Stable Diffusion图片融合代码

以下是基于PyTorch实现的Stable Diffusion图片融合代码，其中包括了模型的定义、训练和推理过程：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from tqdm import tqdm


class Unet(nn.Module):
    def __init__(self):
        super(Unet, self).__init__()
        self.down1 = nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1)
        self.down2 = nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1)
        self.down3 = nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1)
        self.down4 = nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1)
        self.down5 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.down6 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.down7 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.down8 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.up1 = nn.ConvTranspose2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.up2 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)
        self.up3 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)
        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)
        self.up5 = nn.ConvTranspose2d(1024, 256, kernel_size=4, stride=2, padding=1)
        self.up6 = nn.ConvTranspose2d(512, 128, kernel_size=4, stride=2, padding=1)
        self.up7 = nn.ConvTranspose2d(256, 64, kernel_size=4, stride=2, padding=1)
        self.up8 = nn.ConvTranspose2d(128, 3, kernel_size=4, stride=2, padding=1)

    def forward(self, x):
        down1 = F.leaky_relu(self.down1(x), negative_slope=0.2)
        down2 = F.leaky_relu(self.down2(down1), negative_slope=0.2)
        down3 = F.leaky_relu(self.down3(down2), negative_slope=0.2)
        down4 = F.leaky_relu(self.down4(down3), negative_slope=0.2)
        down5 = F.leaky_relu(self.down5(down4), negative_slope=0.2)
        down6 = F.leaky_relu(self.down6(down5), negative_slope=0.2)
        down7 = F.leaky_relu(self.down7(down6), negative_slope=0.2)
        down8 = F.leaky_relu(self.down8(down7), negative_slope=0.2)
        up1 = F.leaky_relu(self.up1(down8), negative_slope=0.2)
        up2 = F.leaky_relu(self.up2(torch.cat([up1, down7], dim=1)), negative_slope=0.2)
        up3 = F.leaky_relu(self.up3(torch.cat([up2, down6], dim=1)), negative_slope=0.2)
        up4 = F.leaky_relu(self.up4(torch.cat([up3, down5], dim=1)), negative_slope=0.2)
        up5 = F.leaky_relu(self.up5(torch.cat([up4, down4], dim=1)), negative_slope=0.2)
        up6 = F.leaky_relu(self.up6(torch.cat([up5, down3], dim=1)), negative_slope=0.2)
        up7 = F.leaky_relu(self.up7(torch.cat([up6, down2], dim=1)), negative_slope=0.2)
        up8 = torch.sigmoid(self.up8(torch.cat([up7, down1], dim=1)))
        return up8


class DiffusionModel(nn.Module):
    def __init__(self, num_steps, betas, model):
        super(DiffusionModel, self).__init__()
        self.num_steps = num_steps
        self.betas = betas
        self.model = model
        self.noise_schedule = nn.Parameter(torch.zeros(num_steps))

    def forward(self, x):
        z = torch.randn(x.shape).to(x.device)
        x_prev = x
        for i in range(self.num_steps):
            t = (i + 1) / self.num_steps
            noise_level = (self.noise_schedule[i] ** 0.5).view(-1, 1, 1, 1)
            x_tilde = x_prev * noise_level + (1 - noise_level ** 2) ** 0.5 * z
            x_prev = x_prev + self.betas[i] * (self.model(x_tilde) - x_prev)
        return x_prev


def train(model, dataloader, optimizer, device):
    model.train()
    for x, _ in tqdm(dataloader):
        x = x.to(device)
        optimizer.zero_grad()
        loss = ((model(x) - x) ** 2).mean()
        loss.backward()
        optimizer.step()


def validate(model, dataloader, device):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for x, _ in tqdm(dataloader):
            x = x.to(device)
            loss = ((model(x) - x) ** 2).mean()
            total_loss += loss.item() * x.shape[0]
    return total_loss / len(dataloader.dataset)


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(256),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    dataset = ImageFolder('path/to/dataset', transform=transform)
    dataloader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=4)

    model = DiffusionModel(1000, torch.linspace(1e-4, 0.1, 1000), Unet()).to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    for epoch in range(10):
        train(model, dataloader, optimizer, device)
        val_loss = validate(model, dataloader, device)
        print(f'Epoch {epoch}: validation loss {val_loss:.4f}')

    torch.save(model.state_dict(), 'path/to/model')


if __name__ == '__main__':
    main()

在训练完成后，可以使用以下代码来融合两张图片：

import torch
from PIL import Image
from torchvision import transforms


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(256),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # 加载模型
    model = DiffusionModel(1000, torch.linspace(1e-4, 0.1, 1000), Unet()).to(device)
    model.load_state_dict(torch.load('path/to/model', map_location=device))

    # 加载图片
    image1 = Image.open('path/to/image1').convert('RGB')
    image2 = Image.open('path/to/image2').convert('RGB')
    x1 = transform(image1).unsqueeze(0).to(device)
    x2 = transform(image2).unsqueeze(0).to(device)

    # 融合图片
    alpha = torch.linspace(0, 1, 11)
    for a in alpha:
        x = a * x1 + (1 - a) * x2
        y = model(x).squeeze(0).detach().cpu()
        y = y * 0.5 + 0.5  # 反归一化
        y = transforms.ToPILImage()(y)
        y.save(f'path/to/result_{a:.1f}.jpg')


if __name__ == '__main__':
    main()

该代码将两张图片进行线性插值，得到11张融合后的图片，其中`alpha`参数指定了插值的权重。在融合过程中，需要进行反归一化操作，将输出的图片转换为PIL格式，并保存到指定路径。