MindSpore实现扩散模型——Palette

原创于 2025-08-15 10:02:23 发布 · 686 阅读

26 ·

CC 4.0 BY-SA版权

文章标签：

#计算机视觉 #人工智能 #深度学习

Palette

论文地址：Palette: Image-to-Image Diffusion Models

代码地址：https://github.com/Janspiry/Palette-Image-to-Image-Diffusion-Models

一、概述

Palette 是一种基于扩散模型的通用图像转换模型，可用于图像着色，图像修复，图像补全和 JPEG 图像恢复等多种转换任务，它通过输入的条件x来构建分布p(y|x)，并通过图像级联的方式引入条件图像。

Palette 的主要特点包括：

统一框架：Palette 模型采用了统一的框架，通过条件扩散模型解决多个图像到图像翻译任务，而无需针对每个任务进行特定的超参数调整、架构定制或辅助损失函数的设计。
样本多样性：通过采用 L2 损失函数，Palette 模型在多个任务上均展现出了较高的样本多样性。这对于需要生成多种可能结果的图像到图像翻译任务尤为重要。
条件 UNet 主干：采用 256×256 分辨率的条件 UNet，通过跳跃连接融合多尺度特征，增强跨层信息传递。

二、主要步骤

1. 预训练基础扩散模型

Palette 的预训练阶段建立在扩散模型的核心框架上，通过学习图像的噪声分布来捕获数据的潜在结构。

扩散过程：逐步向图像添加高斯噪声，将原始图像 x转换为噪声样本 x_T。
去噪过程：训练 UNet 模型学习反向过程，从噪声中恢复原始图像。
关键技术：使用时间嵌入编码当前噪声强度，帮助模型理解去噪阶段。

2. 任务条件化适配

针对不同图像转换任务，Palette 通过以下方式进行条件化：

任务特定提示：为每个任务设计专用的输入提示（如掩码区域、低分辨率图像、灰度图等）。
条件嵌入：将任务提示通过额外的网络层转换为条件嵌入，与 UNet 的中间特征融合。
多任务学习：通过共享主干网络、任务特定分支的方式，使模型同时掌握多个任务。

3. 自适应噪声调度

Palette 引入自适应噪声调度提升不同任务的性能：

任务感知噪声强度：根据任务难度动态调整噪声水平，例如修复任务使用更高噪声强度。
渐进式细化：在去噪后期增加额外的细化步骤，提升细节质量。

4. 后处理与质量优化

生成结果经过以下处理提升最终质量：

细化网络
- 局部增强：对生成的图像进行局部细化，例如增强边缘、平滑过渡区域。
- 对抗训练：引入判别器网络，提升生成图像的真实性。
一致性检查
- 区域匹配：确保生成区域与周围环境自然融合，避免人工痕迹。
- 语义一致性：在多对象场景中，保持对象间的语义关系（如比例、遮挡）。
质量控制
- 置信度估计：预测每个像素的生成置信度，对低置信度区域进行重采样。
- 混合输出：结合多个采样结果，通过加权平均生成最终图像。

三、模型架构

从论文中可以看出，Palette 模型的核心架构就是采用了条件扩散模型作为基础框架，并在此基础上做成了一些改进。

UNet 主干：使用的是 256×256 的条件 UNet，这是一种对称结构。它包含上下采样路径和瓶颈层。同时在上下采样路径中，都包含了每个残差块，它们使用了跳跃连接来整合不同尺度的特征信息，达到捕捉图像的多尺度细节的效果。
条件嵌入：模型融合了多种条件信息来改变生成图像的结果，这其中包括了任务类型嵌入、时间步嵌入和条件图像特征。任务类型嵌入是用于区分图像的转换，例如着色、填充等；而时间步嵌入则是用于帮助模型理解生成过程中的时间信息；条件图像特征则是通过对输入的条件图像进行特征提取，将其作为额外的条件信息输入到模型中，使模型能够根据给定的条件图像生成相应的输出。
多尺度特征融合：在模型生成图像过程中，有一个步骤为跳跃连接，它是将下采样过程中不同维度的图像特征与上采样过程中的对应维度图像特征进行融合，从而丰富图像在不同尺度上的表现。
L2 损失函数：Palette 模型的损失函数是基于去噪目标的，即给定一个训练输出图像y，生成其噪声版本\tilde{y}，并训练神经网络f_θ在条件x和噪声水平指标\gamma下对\tilde{y}进行去噪。而模型所使用的损失函数一般为 L2 范数，其数学公式为：

$E_{(x,y)}E_{\epsilon\sim N(0,I)}E_{\gamma}||f_{\theta}(x,\sqrt{\gamma}y+\sqrt{1-\gamma}\epsilon,\gamma)-\epsilon||_2^2$

四、代码实现

import mindspore
from mindspore import nn, ops
from mindspore.common.initializer import Normal

class Swish(nn.Cell):
    def construct(self, x):
        return x * ops.sigmoid(x)

class TimeEmbedding(nn.Cell):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.half_dim = dim // 2
        self.emb = 9.210340371976184 / (self.half_dim - 1)
        self.linear_1 = nn.Dense(dim, dim * 4)
        self.linear_2 = nn.Dense(dim * 4, dim * 4)
        self.swish = Swish()

    def construct(self, time):
        emb = time * ops.exp(ops.arange(self.half_dim, dtype=mindspore.float32) * -self.emb)
        emb = ops.concat([ops.sin(emb), ops.cos(emb)], axis=1)
        emb = self.linear_1(emb)
        emb = self.swish(emb)
        emb = self.linear_2(emb)
        return emb

class ResidualBlock(nn.Cell):
    def __init__(self, in_channels, out_channels, time_channels, dropout):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, pad_mode='pad')
        self.norm1 = nn.GroupNorm(32, in_channels)
        self.swish = Swish()
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, pad_mode='pad')
        self.norm2 = nn.GroupNorm(32, out_channels)
        self.time_emb_proj = nn.Dense(time_channels, out_channels)
        self.dropout = nn.Dropout(1 - dropout)

        if in_channels != out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1)