Conference

Diffusing Colors: Image Colorization with Text Guided Diffusion

Nir Zabari, Aharon Azulay, Alexey Gorkor, Tavi Halperin, Ohad Fried

Lightricks; Reichman University

SIGGRAPH Asia 2023Image & VideoImage ColorizationDiffusion ModelLatent DiffusionCLIPCold Diffusion
Paper

一句话总结

在预训练潜空间扩散模型（Stable Diffusion 的 VAE 潜空间）里把”RGB→灰度”定义为退化算子，用 Cold Diffusion 框架微调 U-Net 迭代地为灰度图生成色彩，支持文本细粒度控制，并用 CLIP 排序器自动挑选最合适的饱和度，取得当时最优的自动上色效果。

研究背景

  • 领域现状：灰度图上色是经典的计算机图形学问题。传统做法用大规模彩色数据集训练回归模型，或用 GAN、扩散模型等生成式方法采样多样的颜色。文本引导上色也有一些工作，但多局限于”颜色词—物体”的简单绑定。
  • 核心痛点：上色本质是病态问题（一张灰度图对应无穷多合理彩色图）。朴素回归会”回归到均值”，产出灰暗、饱和度低的结果；GAN 存在模式崩溃与训练不稳定；已有文本上色方法表达力弱、依赖小规模预定义词表，且用户输入一旦必需就损害易用性。此外历史照片缺乏真彩色参考，难以构造训练对。
  • 本文 idea：复用一个强大的预训练潜空间扩散模型，把上色任务重新表述为潜空间里的 Cold Diffusion 过程；既能在无任何用户输入时自动上色，也能用描述性文本 prompt 做细粒度全局风格控制，同时用一个 CLIP 排序器自动决定颜色浓淡。

方法

整体框架：先分析 VAE 潜空间中”颜色”的性质，发现彩色图与其灰度图在潜空间的残差向量（称为 color-latent）的缩放，近似等价于像素空间里改变图像的鲜艳度/饱和度。基于这一线性性质，把”去除色度”当作退化算子 \(D\)，在潜空间套用 Cold Diffusion 训练一个时间条件 U-Net 去预测色彩残差；推理时从灰度潜码出发迭代恢复颜色，最后把输出的亮度通道替换为输入原亮度以保留细节。

flowchart LR
  A["灰度图 x'"] --> B["VAE 编码 z'x"]
  B --> C["U-Net 预测色彩残差 Δ (条件: t, 文本 c)"]
  C --> D["Cold Diffusion 迭代恢复"]
  D --> E["CLIP 排序器选最佳饱和度 scale s"]
  E --> F["VAE 解码"]
  F --> G["替换亮度通道为输入 luma"]
  G --> H["彩色输出"]

关键设计：

  1. 潜空间颜色分析：对彩色图 \(z_x = E(x)\) 与灰度图 \(z'_x = E(x')\) 求残差 \(\Delta = z_x - z'_x\)。实验（FID 与 Δ-Colorfulness 指标）表明在潜空间缩放 \(\Delta\) 与在像素空间缩放色度通道行为高度一致，因此可以放心在潜空间做上色，享受更快的训练/推理。VAE 本身在含灰度图的大数据集上训练过，灰度图也落在其编码范围内。

  2. 潜空间 Cold Diffusion（LCDM）：退化算子 \(D\) 去掉 RGB 的色度，恢复算子 \(R\)（U-Net）重新注入色度。任意凸组合 \(\alpha D(z_x) + (1-\alpha) z_x\) 都对应一张位于 \(E\) 张成空间内的自然图像，因此潜空间的线性退化能为不同时间步生成真实训练样本。定义 \(\Delta_t = t \cdot \Delta\)，训练目标是让 U-Net 预测色彩残差：\(\min \mathbb{E}\left[\lVert z_x - (z^t_x + \hat{\Delta}_t) \rVert\right]\)，其中 \(\hat{\Delta}_t = \text{UNet}(z^t_x, t, c)\)。作者强调预测残差 \(\Delta\) 比直接预测潜码效果更好（即扩散文献里的 \(\epsilon\)-prediction）。文本经 CLIP 编码为条件 \(c\)。

  3. 迭代颜色采样：借鉴 Cold Diffusion 的推理，从灰度潜码 \(z_{x'}\) 出发预测初始残差，再按预设时间步表反复”重退化—再预测”逐步重建彩色图。作者证明 negative prompt 与 classifier-free guidance 在 Cold Diffusion 语境下同样适用；步数越多颜色越鲜艳、色域越广、方差越大。最后替换亮度通道以规避 VAE 解码在亮度上的伪影（色度多为低频，更抗解码伪影）。

  4. 颜色浓淡排序器：推理时把最终残差按 \(\hat{z}_x = z_{x'} + s \cdot \Delta_0\) 缩放（\(s=1\) 为默认），交互模式下用户用一个滑块调节且几乎无额外开销。为支持全自动模式，作者在 CLIP 表征上训练一个线性回归排序器：人工标注 500 对训练 / 100 对测试图像的”哪张颜色更好”，用最大化偏好差的损失训练，使其能对单张图按 scale 打分。该排序器在 100 对测试中与人类偏好一致 92 对。全自动流水线还探索了多种自动构造 prompt 的策略，最终发现”空正向 prompt + 固定的反灰度负向 prompt + CFG”效果最好。

实验结果

在 ImageNet ctest10k（10k）与 COCO-Stuff 验证集（5k）上与多种基线对比，主指标为 FID（感知真实性）与 Absolute Δ-Colorfulness（与真值的鲜艳度差异），二者均越低越好。带 Ranker 的版本用自动排序器选饱和度；带 * 的行使用真值彩图的自动生成 caption，仅供参考。

方法 ImageNet FID↓ ImageNet ΔCLR↓ COCO FID↓ COCO ΔCLR↓
BigColor 3.58 1.41 8.23 2.12
Disco 5.57 9.81 10.59 11.82
ColTran 6.37 2.97 11.65 2.07
UniColor 9.43 2.60 11.16 1.91
本文 (Scale=0.8) 3.52 8.35 8.17 9.59
本文 (Ranker) 3.69 0.21 8.05 1.68
本文 (Ranker)* 2.60 1.62 6.62 0.74

本文方法在 FID 上与最强基线 BigColor 相当或更优，并在 Absolute Δ-Colorfulness 上取得最低（最接近真值）的鲜艳度差异，作者认为这源于人工标注时偏好更自然的颜色浓淡。此外还做了用户研究：在 79 张多来源图像上与五个顶尖方法对比，本文以较大优势被最多用户选为”最偏好”，Elo 评分明显领先（本文 86，第二名 BigColor 65）。步数消融显示 2→100 步会让颜色更浓、色域更广；方法也可接在图像修复流水线之后进一步提升老照片上色质量。

亮点与局限

  • 亮点：
    • 把上色重新表述为潜空间 Cold Diffusion，退化算子就是”RGB→灰度”，思路简洁且能复用强大预训练扩散先验。
    • 用回归损失训练而非对抗训练，稳定、不易模式崩溃，同时借迭代扩散保持表达力与多样性。
    • 文本控制超越”颜色词—物体”的简单绑定，支持”golden hour”“film grain”“in the winter”等全局风格描述。
    • 替换亮度通道 + 色度低频特性，有效规避 VAE 解码伪影，这是潜空间生成方法常见的短板。
    • CLIP 排序器让全自动模式也能挑到自然的饱和度，兼顾自动化与可控性。
  • 局限：
    • 依赖预训练 VAE/Stable Diffusion，方法性能与偏置受其潜空间性质约束。
    • 排序器基于作者自己标注的 500/100 对偏好数据，”更好颜色”带主观性，可能不普适。
    • 历史照片仍无真彩色参考，文本只是提供了一个控制旋钮而非真正解决数据稀缺。
    • 扩散迭代推理相比一次前向的回归/GAN 方法更慢，步数与质量需权衡。

延伸思考

  • 把退化算子从”去色度”推广到其它可逆/近似可逆的图像退化（去噪、去模糊、超分），或许能用同一潜空间 Cold Diffusion 框架统一多种图像恢复任务。
  • CLIP 排序器这种”用少量人工偏好对训练线性打分器再自动选参”的思路，可迁移到扩散生成中其它需要挑选超参（引导强度、步数、种子）的场景。
  • 亮度替换技巧提示：在潜空间生成任务中，把对结构/亮度敏感的通道从输入直接旁路、只让网络负责低频/色度信息，是缓解 VAE 解码伪影的通用手段，值得在视频上色、纹理生成等任务里验证时序/空间一致性。