冷扩散：无噪声地逆向任意图像变换

这是 冷扩散 的官方 PyTorch 实现。由马里兰大学的 Arpit Bansal、Eitan Borgnia、Hong-Min Chu、Jie Li 和 Hamid Kazemi 共同开发。该仓库包含用于训练和测试多种冷扩散模型的代码，这些模型针对以下图像退化类型：高斯模糊、变形、高斯掩码、分辨率下采样、图像雪景以及色彩去饱和。我们的实现基于 lucidrains 提供的去噪扩散仓库，后者是 DDPM 的 PyTorch 实现。

更新

2022年10月12日：使用模糊和变形处理生成 CelebA 数据集的预训练模型，以及使用模糊处理生成 AFHQ 数据集的预训练模型，已添加至 我们的云盘。

引用我们的工作

如需引用我们的论文，请使用以下 BibTeX 条目：

@misc{bansal2022cold,
      title={Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise}, 
      author={Arpit Bansal and Eitan Borgnia and Hong-Min Chu and Jie S. Li and Hamid Kazemi and Furong Huang and Micah Goldblum and Jonas Geiping and Tom Goldstein},
      year={2022},
      eprint={2208.09392},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

数据集准备

我们使用 create_data.py 文件将数据拆分为用于训练和测试的独立文件夹。MNIST 和 CIFAR-10 数据集可以直接通过 create_data.py 处理，但需要提供 CelebA 数据集文件夹的路径。AFHQ 数据集已经按训练和测试数据分好文件夹，因此无需额外预处理。每个数据文件夹的保存目录可以修改，并将在各种热/冷扩散模型的训练脚本中使用。

去噪（热）扩散模型

为完整起见，我们提供了用于生成可去除 Gaussian 噪声模型的代码，以便与我们在无条件生成任务上的冷扩散模型直接比较。这些模型的代码位于 denoising-diffusion-pytorch 文件夹中。

训练

我们为每个数据集的模型训练准备了单独的脚本，即 <dataset>_noise_<resolution>.py。根据论文第 5 节的内容，我们包含了针对 128×128 分辨率的 CelebA 和 AFHQ 数据集的脚本。--time_steps 参数可用于调整达到最终各向同性 Gaussian 噪声分布所需的步数。

--sampling_routine 参数允许在不同采样算法之间切换。选择 default 将使用我们论文中的算法 1 进行采样；x0_step_down 是我们论文中算法 2 的固定噪声版本；而 ddim 则是 Song 等人 提出的、采用估计噪声的算法 2 版本。

--save_folder 参数指定了保存训练好的模型以及用于跟踪进度的训练数据样本的路径。保存和进度跟踪的频率可以在 denoising_diffusion_pytorch.py 中定义的 Trainer 类中进行修改。data_path 参数则指定了数据集准备步骤中生成的训练数据文件夹路径。

以下是训练去噪扩散模型的示例脚本：

python <dataset>_noise_<resolution>.py --time_steps 200 --sampling_routine x0_step_down --save_folder <模型保存路径> --data_path <训练数据路径>

测试

测试模型时，务必输入与训练时相同的步数。采样方法可以不同，此外还可以使用 test_type 参数指定要进行的测试类型。有关不同测试方法的说明，请参阅 denoising_diffusion_pytorch.py 中定义的 Trainer 类。

下面提供两个测试脚本示例。第一个脚本使用 DDIM 采样方法并采用估计噪声，第二个脚本则使用我们在论文算法 2 中描述的固定噪声方法。

python <dataset>_noise_<resolution>_test.py --time_steps 200 --sampling_routine ddim --save_folder <图片保存路径> --load_path <模型加载路径> --data_path <数据路径> --test_type test_sample_and_save_for_fid

python <dataset>_noise_<resolution>_test.py --time_steps 200 --sampling_routine x0_step_down --save_folder <图片保存路径> --load_path <模型加载路径> --data_path <数据路径> --test_type test_sample_and_save_for_fid

冷扩散模型

我们介绍了论文中探讨的各种冷扩散模型的训练和测试流程。对于每种冷扩散，都提供了两种不同退化方案的示例脚本。有条件生成的方案可用于复现论文第 4 节的结果，而无条件生成的方案则可用于复现第 5 节的结果。除非另有说明，否则参数和流程与热扩散模型的详细描述完全相同。

高斯模糊

这些模型的代码位于 deblurring-diffusion-pytorch 文件夹中。对于高斯模糊，用于指定退化时间表的参数包括 --time_steps、--blur_size、--blur_std 和 blur_routine。有关这些参数的具体含义，请参阅训练脚本中的帮助信息。

条件生成

训练

python mnist_train.py --time_steps 20 --blur_size 11 --blur_std 7.0 --blur_routine 'Constant' --sampling_routine x0_step_down --data_path <数据文件夹路径> --save_folder <保存模型路径> 
python cifar10_train.py --time_steps 50 --blur_routine 'Special_6_routine' --sampling_routine x0_step_down --data_path <数据文件夹路径> --save_folder <保存模型路径> 
python celebA_128.py --time_steps 200 --blur_size 15 --blur_std 0.01 --blur_routine Exponential_reflect --sampling_routine x0_step_down --data_path <数据文件夹路径> --save_folder <保存模型路径> 

测试

python mnist_test.py --time_steps 20 --blur_size 11 --blur_std 7.0 --blur_routine 'Constant' --sampling_routine 'x0_step_down' --save_folder <保存结果路径> --data_path <数据文件夹路径> --test_type test_data
python cifar10_test.py --time_steps 50 --blur_routine 'Special_6_routine' --sampling_routine 'x0_step_down' --save_folder <保存结果路径> --data_path <数据文件夹路径> --test_type test_data
python celebA_128_test.py --time_steps 200 --blur_size 15 --blur_std 0.01 --blur_routine Exponential_reflect --sampling_routine x0_step_down --save_folder <保存结果路径> --data_path <数据文件夹路径> --test_type test_data

无条件生成

训练

python celebA_128.py --discrete --time_steps 300 --blur_size 27 --blur_std 0.01 --blur_routine Exponential --sampling_routine x0_step_down --data_path <数据文件夹路径> --save_folder <保存模型路径>
python AFHQ_128.py --discrete --time_steps 300 --blur_size 27 --blur_std 0.01 --blur_routine Exponential --sampling_routine x0_step_down --data_path <数据文件夹路径> --save_folder <保存模型路径>

测试

以下是论文中所展示的模型/数据集的两组测试脚本。第一组脚本对应于在完全对称条件下进行的采样，第二组脚本则对应于在对称性被破坏的情况下进行的采样。--gmm_cluster 参数指定了拟合到退化图像分布上的高斯混合模型（GMM）的模式数。--noise 参数用于控制从 GMM 采样后对称性被破坏的程度。

python celebA_128_test.py --gmm_cluster 1 --noise 0.000 --discrete --time_steps 300 --blur_size 27 --blur_std 0.01 --blur_routine Exponential --sampling_routine x0_step_down --save_folder <保存结果路径> --load_path <加载模型路径> --data_path <数据文件夹路径> --test_type train_distribution_mean_blur_torch_gmm_ablation
python AFHQ_128_test.py --gmm_cluster 1 --noise 0.000 --discrete --time_steps 300 --blur_size 27 --blur_std 0.01 --blur_routine Exponential --sampling_routine x0_step_down --save_folder <保存结果路径> --load_path <加载模型路径> --data_path <数据文件夹路径> --test_type train_distribution_mean_blur_torch_gmm_ablation

python celebA_128_test.py --gmm_cluster 1 --noise 0.002 --discrete --time_steps 300 --blur_size 27 --blur_std 0.01 --blur_routine Exponential --sampling_routine x0_step_down --save_folder <保存结果路径> --load_path <加载模型路径> --data_path <数据文件夹路径> --test_type train_distribution_mean_blur_torch_gmm_ablation
python AFHQ_128_test.py --gmm_cluster 1 --noise 0.002 --discrete --time_steps 300 --blur_size 27 --blur_std 0.01 --blur_routine Exponential --sampling_routine x0_step_down --save_folder <保存结果路径> --load_path <加载模型路径> --data_path <数据文件夹路径> --test_type train_distribution_mean_blur_torch_gmm_ablation

动物变形术

这些模型的代码位于 demixing-diffusion-pytorch 文件夹中。

无条件生成

训练

python AFHQ_128_to_celebA_128.py --time_steps 200 --sampling_routine x0_step_down --save_folder <保存模型路径> --data_path_start <起始数据流形路径> --data_path_end <终止数据流形路径>

测试

python AFHQ_128_to_celebA_128_test.py --time_steps 200 --sampling_routine x0_step_down --save_folder <保存图像路径> --load_path <加载模型路径> --data_path_start <起始数据流形路径> --data_path_end <终止数据流形路径> --test_type test_sample_and_save_for_fid

图像修复

条件生成的代码位于 defading-diffusion-pytorch 文件夹中，而非条件生成的代码则位于 defading-generation-diffusion-pytorch 文件夹中。对于高斯掩码，用于指定退化过程的参数包括 --time_steps、--kernel_std 和 fade_routine。有关这些参数的具体含义，请参阅训练脚本中的帮助信息。

条件生成

训练

python mnist_train.py --time_steps 50 --save_folder <保存模型的路径> --discrete --sampling_routine x0_step_down --train_steps 700000 --kernel_std 0.1 --fade_routine Random_Incremental --data_path <数据文件夹路径>
python cifar10_train.py --time_steps 50 --save_folder <保存模型的路径> --discrete --sampling_routine x0_step_down --train_steps 700000 --kernel_std 0.1 --fade_routine Random_Incremental --data_path <数据文件夹路径>
python celebA_train.py --time_steps 100 --fade_routine Incremental --save_folder <保存模型的路径> --sampling_routine x0_step_down --train_steps 350000 --kernel_std 0.2 --image_size 128 --dataset celebA --data_path <数据文件夹路径>

测试

python mnist_test.py --time_steps 50 --save_folder test_mnist --discrete --sampling_routine x0_step_down --kernel_std 0.1 --initial_mask 1 --image_size 28 --fade_routine Random_Incremental --load_path <加载模型的路径> --data_path <数据文件夹路径> --test_type test_data 
python cifar10_test.py --time_steps 50 --save_folder test_cifar10 --discrete --sampling_routine x0_step_down --kernel_std 0.1 --initial_mask 1 --image_size 32 --fade_routine Random_Incremental --load_path <加载模型的路径> --data_path < 数据文件夹路径 > --test_type test_data
python celebA_test.py --time_steps 100 --fade_routine Incremental --save_folder test_celebA --sampling_routine x0_step_down --kernel_std 0.2 --initial_mask 1 --image_size 128 --dataset celebA --load_path <加载模型的路径> --data_path < 数据文件夹路径 > --test_type test_data

非条件生成

训练

python celebA_128.py --reverse --kernel_std 0.05 --initial_mask 1 --time_steps 750 --sampling_routine x0_step_down --save_folder <保存模型的路径> --data_path < 数据文件夹路径 >

测试

python celebA_constant_128_test.py --reverse --kernel_std 0.05 --time_steps 750 --sampling_routine x0_step_down --save_folder <保存图像的路径> --data_path < 数据文件夹路径 > --load_path < 加载模型的路径 > --test_type test_sample_and_save_for_fid

超分辨率

这些模型的代码位于 resolution-diffusion-pytorch 文件夹中。对于下采样，用于指定退化过程的参数包括 --time_steps 和 resolution_routine。有关这些参数的具体含义，请参阅训练脚本中的帮助信息。

条件生成

训练

python mnist_train.py --time_steps 3 --resolution_routine 'Incremental_factor_2' --save_folder <保存模型的路径>
python cifar10_train.py --time_steps 3 --resolution_routine 'Incremental_factor_2' --save_folder <保存模型的路径>
python celebA_128.py --time_steps 4 --resolution_routine 'Incremental_factor_2' --save_folder <保存模型的路径>

测试

python mnist_test.py --time_steps 3 --train_routine 'Final' --sampling_routine 'x0_step_down' --resolution_routine 'Incremental_factor_2' --save_folder <保存图像的路径> --load_path <加载模型的路径> --test_type test_data
python cifar10_test.py --time_steps 3 --train_routine 'Final' --sampling_routine 'x0_step_down' --resolution_routine 'Incremental_factor_2' --save_folder <保存图像的路径> --load_path < 加载模型的路径 > --test_type test_data
python celebA_test.py --time_steps 4 --train_routine 'Final' --sampling_routine 'x0_step_down' --resolution_routine 'Incremental_factor_2' --save_folder <保存图像的路径> --load_path < 加载模型的路径 > --test_type test_data

非条件生成

训练

python celebA_test.py --time_steps 4 --train_routine 'Final' --sampling_routine 'x0_step_down' --resolution_routine 'Incremental_factor_2' --save_folder <保存图像的路径> --load_path < 加载模型的路径 > --test_type test_data

测试

python celebA_test.py --time_steps 4 --train_routine 'Final' --sampling_routine 'x0_step_down' --resolution_routine 'Incremental_factor_2' --save_folder <保存图像的路径> --load_path < 加载模型的路径 > --test_type test_data

降雪效果

这些模型的代码位于 snowification 文件夹中。该代码源自现有的 ImageNet C 仓库。用于指定退化过程的参数包括 --time_steps、--snow_level、--random_snow 和 --fix_brightness。有关这些参数的具体含义，请参阅训练脚本中的帮助信息。

条件生成

训练

python train.py --dataset cifar10 --time_steps 200 --forward_process_type ‘Snow’ --snow_level 3 --exp_name <实验名称>  --dataset_folder <数据集路径> --random_snow --fix_brightness  --sampling_routine x0_step_down
python train.py --dataset celebA --time_steps 200 --forward_process_type ‘Snow’ --snow_level 4 --exp_name <实验名称> --dataset_folder <数据集路径> --random_snow --fix_brightness  --sampling_routine x0_step_down

测试

python test.py --dataset cifar10 --time_steps 200 --forward_process_type ‘Snow’ --snow_level 3 --exp_name <实验名称> --dataset_folder <数据集路径> --random_snow --fix_brightness --resume_training --sampling_routine x0_step_down --test_type test_data --order_seed 1
python test.py --dataset celebA --time_steps 200 --forward_process_type ‘Snow’ --snow_level 4 --exp_name <实验名称> --dataset_folder <数据集路径> --random_snow --fix_brightness --resume_training --sampling_routine x0_step_down --test_type test_data --order_seed 1

彩色化

这些模型的代码位于 decolor-diffusion 文件夹中。对于颜色去饱和处理，用于指定退化时间表的参数是 --time_steps、decolor_total_remove 和 decolor_routine。有关这些参数的具体含义，请参阅训练脚本中的帮助信息。

条件生成

训练

python train.py --dataset cifar10 --time_steps 20 --forward_process_type ‘Decolorization’ --exp_name <exp_name> --decolor_total_remove --decolor_routine ‘Linear’ --dataset_folder <path-to-dataset>
python train.py --dataset celebA --time_steps 20 --forward_process_type ‘Decolorization’ --exp_name <exp_name> --decolor_total_remove --decolor_routine ‘Linear’ --dataset_folder <path-to-dataset>

测试

python test.py --dataset cifar10 --time_steps 20 --forward_process_type ‘Decolorization’ --exp_name <exp-name>  --decolor_total_remove --decolor_routine ‘Linear’ --dataset_folder <path-to-dataset> --sampling_routine x0_step_down --test_type test_data --order_seed 1
python test.py --dataset celebA --time_steps 20 --forward_process_type ‘Decolorization’ --exp_name <exp-name>  --decolor_total_remove --decolor_routine ‘Linear’ --dataset_folder <path-to-dataset> --sampling_routine x0_step_down --test_type test_data --order_seed 1

Cold Diffusion Models 快速上手指南

Cold Diffusion 是一种无需添加噪声即可逆转任意图像变换（如高斯模糊、掩膜、降采样等）的扩散模型。本指南基于官方 PyTorch 实现，帮助开发者快速完成环境配置、安装及基础运行。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS。Windows 用户建议使用 WSL2。
• Python: 3.8 或更高版本。
• GPU: 推荐使用 NVIDIA GPU 以加速训练和推理（需安装对应的 CUDA 驱动）。
• 依赖库:
  • PyTorch (最新版)
  • torchvision
  • tqdm, PIL, numpy, scipy 等常用科学计算库。

国内加速建议： 安装 PyTorch 时，推荐使用清华或中科大镜像源以提升下载速度：

pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库 从 GitHub 获取源代码：

   git clone https://github.com/arpitbansal97/Cold-Diffusion-Models.git
   cd Cold-Diffusion-Models
   

2. 安装 Python 依赖 项目根目录下通常包含 requirements.txt（若不存在，请根据代码导入手动安装核心库）。执行以下命令安装依赖：

   pip install -r requirements.txt
   

   注：若缺少 requirements.txt，请至少确保安装了 torch, torchvision, tqdm, pillow, numpy。

3. 数据集准备 使用提供的脚本处理数据集。对于 MNIST 和 CIFAR-10 可直接处理，CelebA 需指定路径，AFHQ 无需预处理。

   # 示例：处理 MNIST 或 CIFAR-10
   python create_data.py --dataset mnist --data_path ./data/mnist
   # 对于 CelebA，需提前下载并解压，然后指定路径
   

基本使用

Cold Diffusion 支持多种退化类型（如高斯模糊、Animorphosis、图像修复等）。以下以**高斯模糊去模糊（Conditional Generation）**为例，展示最简训练与测试流程。

1. 训练模型 (Training)

以 MNIST 数据集为例，训练一个逆转高斯模糊的 Cold Diffusion 模型：

python mnist_train.py \
  --time_steps 20 \
  --blur_size 11 \
  --blur_std 7.0 \
  --blur_routine 'Constant' \
  --sampling_routine x0_step_down \
  --data_path ./data/mnist \
  --save_folder ./models/mnist_blur

参数说明：

• --time_steps: 扩散步数。
• --blur_size / --blur_std: 定义高斯模糊核的大小和标准差。
• --blur_routine: 退化调度策略（如 'Constant', 'Exponential' 等）。
• --sampling_routine: 采样算法，推荐使用 x0_step_down (论文算法 2)。

2. 测试/推理 (Testing)

使用训练好的模型对测试数据进行去模糊还原：

python mnist_test.py \
  --time_steps 20 \
  --blur_size 11 \
  --blur_std 7.0 \
  --blur_routine 'Constant' \
  --sampling_routine 'x0_step_down' \
  --save_folder ./results/mnist_blur \
  --data_path ./data/mnist \
  --load_path ./models/mnist_blur \
  --test_type test_data

注意：

• 测试时的 --time_steps、--blur_size、--blur_std 和 --blur_routine 必须与训练时保持一致。
• --load_path 指向训练保存的模型文件夹。
• 生成结果将保存在 --save_folder 指定的目录中。

其他任务简述

• 无条件生成 (Unconditional Generation): 参考 README 中 celebA_128.py 脚本，需增加 --discrete 参数并调整步数（如 300 步）。
• 图像形态转换 (Animorphosis): 使用 demixing-diffusion-pytorch 文件夹下的脚本，需指定起始和结束数据流形路径 (--data_path_start, --data_path_end)。
• 图像修复 (Inpainting): 使用 defading-diffusion-pytorch 文件夹，通过 --kernel_std 和 --fade_routine 控制掩膜退化过程。