扩散自编码器的官方实现

CVPR 2022 (口头报告)论文(论文链接，项目网站，5分钟视频)：

@inproceedings{preechakul2021diffusion,
      title={扩散自编码器：迈向有意义且可解码的表示}, 
      author={Preechakul, Konpat and Chatthee, Nattanat and Wizadwongsa, Suttisak and Suwajanakorn, Supasorn},
      booktitle={IEEE计算机视觉与模式识别会议(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（需要pip install dlib requests）
3. 对齐后的图像将保存在imgs_align目录中

<table> <tr> <th width="33%"> <code>imgs</code>目录中的原始图像<br><img src="https://yellow-cdn.veclightyear.com/835a84d5/d66e3759-7103-4d0e-a173-aaef5834b936.JPG" style="width: 100%"> </th> <th width="33%"> 使用<code>align.py</code>对齐后<br><img src="https://yellow-cdn.veclightyear.com/835a84d5/e6fbceda-d0bf-43c4-8f12-1881adadff7e.png" style="width: 100%"> </th> <th width="33%"> 使用<code>manipulate.ipynb</code>的效果<br><img src="https://yellow-cdn.veclightyear.com/835a84d5/910cd0ca-78f3-4f7e-a5f4-5feb8c624d2f.png" style="width: 100%"> </th> </tr> </table>

模型检查点

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

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云端硬盘：https://drive.google.com/drive/folders/1abNP4QKGbNnymjn8607BF0cwxX2L23jh?usp=sharing

训练

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

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

FFHQ128

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

可以使用以下命令训练分类器（用于操作）：

python run_ffhq128_cls.py

FFHQ256

由于计算成本较高，我们只训练了DiffAE。 这需要8个V100 GPU。

sbatch run_ffhq256.py

任务完成后，您需要训练潜在DPM（只需要1个2080Ti GPU）

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 GPU上运行。

# diffae
python run_celeba64.py