扩散自编码器的官方实现

一篇 CVPR 2022（口头报告）论文（论文、网站、5分钟视频）：

@inproceedings{preechakul2021diffusion,
      title={Diffusion Autoencoders: Toward a Meaningful and Decodable Representation}, 
      author={Preechakul, Konpat and Chatthee, Nattanat and Wizadwongsa, Suttisak and Suwajanakorn, Supasorn},
      booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 
      year={2022},
}

使用方法

⚙️ 尝试 Colab 教程：

🤗 尝试网页演示：

注意：由于代码库预计会有大量改动，请在使用前先 fork 该仓库。

前置条件

请参阅 requirements.txt 文件。

pip install -r requirements.txt

快速入门

提供了一个 Jupyter Notebook。

无条件生成：sample.ipynb

操控：manipulate.ipynb

插值：interpolate.ipynb

自编码：autoencoding.ipynb

对齐您自己的图像：

1. 将图像放入 imgs 目录。
2. 运行 align.py（需要安装 dlib 和 requests）。
3. 结果图像将保存在 imgs_align 目录中。

`imgs` 目录中的原始图像	使用 `align.py` 对齐后的图像	使用 `manipulate.ipynb` 处理后的图像

检查点

我们为以下模型提供了检查点：

1. DDIM：FFHQ128 (72M, 130M)，Bedroom128，Horse128
2. DiffAE（仅自编码）：FFHQ256，FFHQ128 (72M, 130M)，Bedroom128，Horse128
3. DiffAE（带有潜在 DPM，可采样）：FFHQ256，FFHQ128，Bedroom128，Horse128
4. DiffAE 的分类器（用于操控）：FFHQ256 在 CelebAHQ 上的潜在表示，FFHQ128 在 CelebAHQ 上的潜在表示

检查点应放置在一个单独的目录 checkpoints 中。下载检查点并将其放入 checkpoints 目录中，目录结构应如下所示：

checkpoints/
- bedroom128_autoenc
    - last.ckpt # diffae 检查点
    - latent.ckpt # 数据集上的预测 z_sem
- bedroom128_autoenc_latent
    - last.ckpt # diffae + 潜在 DPM 检查点
- bedroom128_ddpm
- ...

LMDB 数据集

我们不拥有以下任何数据集。为了方便起见，我们提供了可以直接使用的 LMDB 格式数据集。

• FFHQ
• CelebAHQ

链接已失效

注：我正在尝试恢复以下链接。

• CelebA
• LSUN Bedroom
• LSUN Horse

目录结构应如下所示：

datasets/
- bedroom256.lmdb
- celebahq256.lmdb
- celeba.lmdb
- ffhq256.lmdb
- horse256.lmdb

您也可以从原始来源下载数据，并使用我们提供的代码将其打包成 LMDB 文件。各数据集的原始来源如下：

• FFHQ（https://github.com/NVlabs/ffhq-dataset）
• CelebAHQ（https://github.com/switchablenorms/CelebAMask-HQ）
• CelebA（https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html）
• LSUN（https://github.com/fyu/lsun）

转换代码如下：

data_resize_bedroom.py
data_resize_celebhq.py
data_resize_celeba.py
data_resize_ffhq.py
data_resize_horse.py

Google Drive：https://drive.google.com/drive/folders/1abNP4QKGbNnymjn8607BF0cwxX2L23jh?usp=sharing

训练

我们提供了用于训练和评估 DDIM 以及 DiffAE（包括潜在 DPM）的脚本，适用于以下数据集：FFHQ128、FFHQ256、Bedroom128、Horse128、Celeba64（D2C 裁剪版本）。通常，评估结果（FID 分数）将保存在 eval 目录中。

注意：大多数实验在训练 DPM 模型时至少需要 4 张 V100 显卡，而在训练配套的潜在 DPM 时则只需要 1 张 2080Ti 显卡。

FFHQ128

# diffae
python run_ffhq128.py
# ddim
python run_ffhq128_ddim.py

可用于训练分类器（用于操控）的脚本如下：

python run_ffhq128_cls.py

FFHQ256

由于计算成本较高，我们仅训练了 DiffAE。这需要 8 张 V100 显卡。

sbatch run_ffhq256.py

任务完成后，您还需要训练潜在 DPM（只需 1 张 2080Ti 显卡）：

python run_ffhq256_latent.py

可用于训练分类器（用于操控）的脚本如下：

python run_ffhq256_cls.py

Bedroom128

# diffae
python run_bedroom128.py
# ddim
python run_bedroom128_ddim.py

Horse128

# diffae
python run_horse128.py
# ddim
python run_horse128_ddim.py

Celeba64

此实验可以在 2080Ti 显卡上运行。

# diffae
python run_celeba64.py

DiffAE 快速上手指南

DiffAE (Diffusion Autoencoders) 是一个基于扩散模型的自编码器，旨在生成具有语义意义且可解码的潜在表示。它支持无条件图像生成、属性编辑、插值以及图像重建。

环境准备

• 系统要求：Linux 或 macOS（Windows 需使用 WSL），建议配备 NVIDIA GPU。
• 硬件建议：
  • 推理/微调：至少 1 张 2080Ti 或同等算力显卡。
  • 从头训练 DPM 模型：建议 4 张 V100 及以上。
• 前置依赖：
  • Python 3.8+
  • PyTorch
  • Git
  • dlib 和 requests（用于人脸对齐功能）

注意：由于代码库可能频繁更新，建议在使用前先 Fork 该仓库。

安装步骤

1. 克隆仓库

   git clone https://github.com/preechakul/diffae.git
   cd diffae
   

2. 安装 Python 依赖 推荐使用国内镜像源加速安装（如清华源）：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装人脸对齐额外依赖 如果需要处理自定义人脸图像进行属性编辑，需安装以下库：

   pip install dlib requests -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：dlib 可能需要系统级编译工具，如 cmake 和 build-essential)

4. 下载预训练模型 (Checkpoints) 将下载的模型文件放入项目根目录下的 checkpoints 文件夹。目录结构应如下所示：

   checkpoints/
   - bedroom128_autoenc
       - last.ckpt
       - latent.ckpt
   - bedroom128_autoenc_latent
       - last.ckpt
   - ... (其他模型文件夹)
   

   模型下载地址请参考原仓库 README 中的 Google Drive 链接。

5. 准备数据集 (可选) 如需使用自带脚本转换数据集，请将原始数据下载后放入 datasets 目录，或使用提供的 LMDB 文件。目录结构示例：

   datasets/
   - ffhq256.lmdb
   - bedroom256.lmdb
   - ...
   

基本使用

DiffAE 提供了多个 Jupyter Notebook 以便快速体验不同功能。请确保已下载对应的 Checkpoints 并放置在正确目录。

1. 启动示例

直接在浏览器中打开对应的 .ipynb 文件运行即可：

• 无条件图像生成: sample.ipynb
• 图像属性编辑 (Manipulation): manipulate.ipynb
• 图像插值 (Interpolation): interpolate.ipynb
• 图像自编码重建: autoencoding.ipynb

2. 处理自定义图像进行编辑

若要对本地人脸图片进行属性编辑（如改变发型、表情等），需先进行人脸对齐预处理：

1. 将原始图片放入 imgs 目录。
2. 运行对齐脚本：

   python align.py
   

3. 对齐后的图片将生成在 imgs_align 目录。
4. 使用 manipulate.ipynb 加载 imgs_align 中的图片进行编辑。

3. 在线体验 (无需本地部署)

如果暂时无法配置本地环境，可尝试以下在线方案：

• Google Colab: 点击此处打开 walkthrough
• Replicate Web Demo: 点击此处体验