利用扩散模型实现高质量图像生成【原理解析与代码实践】

import torch
import numpy as np

def forward_diffusion_process(x_0, timesteps, noise_schedule):
    """
    正向扩散过程：给图像x_0逐步加噪，生成一系列x_t。
    :param x_0: 原始图像 (batch_size, C, H, W)
    :param timesteps: 扩散的时间步数
    :param noise_schedule: 噪声系数 α_t 列表
    :return: x_t 向量
    """
    noise_schedule = torch.tensor(noise_schedule).float()
    batch_size, C, H, W = x_0.shape
    x_t = x_0.clone()
    x_t_list = [x_0]
    
    for t in range(timesteps):
        noise = torch.randn_like(x_0)
        alpha_t = noise_schedule[t]
        x_t = np.sqrt(alpha_t) * x_t + np.sqrt(1 - alpha_t) * noise
        x_t_list.append(x_t)
    
    return x_t_list

import torch.nn as nn

class DenoiseNetwork(nn.Module):
    def __init__(self, channels):
        super(DenoiseNetwork, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(channels, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, channels, kernel_size=3, padding=1)
        )
    
    def forward(self, x, t):
        return self.net(x)
    
def denoise_step(denoise_net, x_t, t):
    """
    逆向去噪步骤：使用去噪网络从 x_t 生成去噪后的 x_{t-1}
    :param denoise_net: 去噪网络
    :param x_t: 在时间步 t 的加噪图像
    :param t: 时间步
    :return: 去噪后的图像 x_{t-1}
    """
    predicted_x_t_minus_1 = denoise_net(x_t, t)
    return predicted_x_t_minus_1

def train_diffusion_model(denoise_net, optimizer, data_loader, timesteps, noise_schedule):
    """
    训练扩散模型
    :param denoise_net: 去噪网络
    :param optimizer: 优化器
    :param data_loader: 数据加载器
    :param timesteps: 扩散时间步
    :param noise_schedule: 噪声系数 α_t 列表
    """
    denoise_net.train()
    
    for epoch in range(epochs):
        for x_0 in data_loader:
            optimizer.zero_grad()
            
            # 生成随机时间步
            t = torch.randint(0, timesteps, (x_0.shape[0],)).long()
            
            # 获取正向扩散的 x_t
            x_t_list = forward_diffusion_process(x_0, timesteps, noise_schedule)
            x_t = x_t_list[t]
            
            # 去噪网络预测
            predicted_x_t_minus_1 = denoise_step(denoise_net, x_t, t)
            
            # 计算损失并反向传播
            loss = nn.MSELoss()(predicted_x_t_minus_1, x_0)
            loss.backward()
            optimizer.step()

def generate_image(denoise_net, timesteps, noise_schedule):
    """
    使用训练好的去噪网络生成高质量图像
    :param denoise_net: 训练好的去噪网络
    :param timesteps: 扩散时间步
    :param noise_schedule: 噪声系数 α_t 列表
    :return: 生成的图像
    """
    # 从随机噪声开始
    x_t = torch.randn((1, 3, 64, 64))
    
    for t in reversed(range(timesteps)):
        # 使用去噪网络生成 x_{t-1}
        x_t

 = denoise_step(denoise_net, x_t, t)
    
    return x_t

class DAE(nn.Module):
    def __init__(self):
        super(DAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),
            nn.Tanh(),
        )

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z)
    
# 应用在扩散模型的逆向去噪过程
def denoise_with_dae(dae, x_t, t):
    return dae(x_t)

def adaptive_noise_schedule(timesteps, method='cosine'):
    """
    自适应噪声调度函数
    :param timesteps: 扩散时间步数
    :param method: 调度方式 ('cosine' 或 'linear')
    :return: 自适应噪声调度 α_t
    """
    if method == 'cosine':
        return np.cos(np.linspace(0, np.pi / 2, timesteps))**2
    elif method == 'linear':
        return np.linspace(1, 0, timesteps)
    else:
        raise ValueError("Unknown method")
    
# 使用优化后的噪声调度
noise_schedule = adaptive_noise_schedule(timesteps=1000, method='cosine')

class AttentionBlock(nn.Module):
    def __init__(self, channels):
        super(AttentionBlock, self).__init__()
        self.attn = nn.MultiheadAttention(embed_dim=channels, num_heads=4)

    def forward(self, x):
        batch_size, C, H, W = x.size()
        x = x.view(batch_size, C, H * W).permute(2, 0, 1)
        x, _ = self.attn(x, x, x)
        x = x.permute(1, 2, 0).view(batch_size, C, H, W)
        return x

class DenoiseNetWithAttention(nn.Module):
    def __init__(self, channels):
        super(DenoiseNetWithAttention, self).__init__()
        self.conv1 = nn.Conv2d(channels, 64, kernel_size=3, padding=1)
        self.attn_block = AttentionBlock(64)
        self.conv2 = nn.Conv2d(64, channels, kernel_size=3, padding=1)

    def forward(self, x, t):
        x = torch.relu(self.conv1(x))
        x = self.attn_block(x)
        return torch.tanh(self.conv2(x))

class VariationalDiffusion(nn.Module):
    def __init__(self, timesteps, latent_dim):
        super(VariationalDiffusion, self).__init__()
        self.encoder = nn.Linear(latent_dim, timesteps)
        self.decoder = nn.Linear(timesteps, latent_dim)
    
    def forward(self, x_t):
        # 编码当前图像
        z = self.encoder(x_t)
        # 解码生成下一个时间步
        x_t_minus_1 = self.decoder(z)
        return x_t_minus_1

def reduce_steps_diffusion(denoise_net, x_T, reduced_timesteps):
    """
    步数减少的扩散生成过程
    :param denoise_net: 去噪网络
    :param x_T: 初始噪声图像
    :param reduced_timesteps: 减少后的时间步数
    :return: 最终生成的图像
    """
    x_t = x_T
    steps = np.linspace(0, 1000, reduced_timesteps).astype(int)
    
    for t in reversed(steps):
        x_t = denoise_step(denoise_net, x_t, t)
    
    return x_t