最近在研究diffusion模型，并对目前最火的stable-diffusion模型很感兴趣，又因为stable-diffusion是一种latent-diffusion模型，故尝试复现latent-diffusion model，并训练自己的LDM。
写这篇博客的出发点是因为，当我跟随github页面上的install教程进行安装的时候，遇到了很多问题，有些是因为网络问题，服务器无法直接下载国外服务器的资源；有些则是软件版本其实是错的，照着安装就是不对。
我修改了以下项目中environment.yaml文件里的配置，这里请跟着我的步骤进行安装
我将environment.yaml文件中的配置提取了出来，防止直接安装yaml文件时出现网络错误，导致无法继续安装
把下面的内容复制到requirement.txt文件里
albumentations==0.4.3
opencv-python==4.1.2.30
pudb==2019.2
imageio==2.9.0
imageio-ffmpeg==0.4.2
pytorch-lightning==1.4.2
omegaconf==2.1.1
test-tube>=0.7.5
streamlit>=0.73.1
einops==0.3.0
torch-fidelity==0.3.0
transformers==4.6.0
torchmetrics==0.6
kornia==0.5.10
six
-e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
-e git+https://github.com/openai/CLIP.git@main#egg=clip
-e .
taming-transformers项目地址
clip项目地址
下载好后解压，然后将两个项目的文件夹名分别修改为taming-transformers和clip，然后把这两个文件夹复制到latent-diffusion项目下的src文件夹下(如果没有src文件夹，自己创建一个)，如下图所示：
执行下面命令创建ldm环境
conda create -n ldm python=3.8.5 # 创建新环境
conda activate ldm # 激活新环境
安装requirement.txt文件
pip install -r requirement.txt
Tips:
1.为什么我没有提前安装torch，因为安装好requirement.txt后，会自动安装上torch2.0，而torch2.0在执行代码时会报错，版本是错误的，因此我手动将其删除，再重新安装低版本的torch
2.如果按照LDM项目里的environment.yaml安装的话，会缺一些包，并且包的版本也是错的，进到导致代码执行报错。于是我Google好久找到了解决方法，以及合适的包。我已经将合适的包以及版本添加进了requirement.txt文件。如：
transformers升级为4.6.0版本，添加torchmetrics=0.6 (解决Bug的Answer)
添加kornia=0.5.10 (torch和kornia的对应关系)，
添加six包。
照着安装你们就不用再走一遍弯路了。
删除toch2.0
pip uninstall torch
不要跟着environment.yaml里的torch版本安装，安装完执行代码依然会报错，cudatookit11.0 就是会出错，要用cudatookit11.1。执行下面命令安装
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
因为国内的服务器大多数是无法直接访问huggingface网站的，进而导致无法联网下载模型，因而被迫只能先将模型下载到本地，然后再上传至服务器。
在这里下载bert-base-uncased的权重，只需要下载如图方块所示的文件即可：
将这三个文件放在一个文件夹内，然后你可以将文件夹上传到服务器上LDM项目下的models目录下，想要系统了解的可以参考这篇文章《huggingface transformers预训练模型如何下载至本地，并使用？》，如图：
然后修改LDM项目下ldm/modules/encoders/modules.py这个文件
对应代码在第58行，修改from_pretrained的路径为本地路径，注意是权重文件对应文件夹名
class BERTTokenizer(AbstractEncoder):
""" Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)"""
def __init__(self, device="cuda", vq_interface=True, max_length=77):
super().__init__()
from transformers import BertTokenizerFast # TODO: add to reuquirements
# self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
# 需要提前从huggingface上把模型文件下载到本地，然后进行本地加载，否则会出现联网错误的问题，导致模型无法下载
self.tokenizer = BertTokenizerFast.from_pretrained("models/bert") # 加载本地文件
self.device = device
self.vq_interface = vq_interface
self.max_length = max_length
测试Text-to-Image任务
如果模型文件你无法在服务器上直接下载，请在本地单独下载，然后再上传至服务器
# 下载预训练权重(5.7GB):
mkdir -p models/ldm/text2img-large/
wget -O models/ldm/text2img-large/model.ckpt https://ommer-lab.com/files/latent-diffusion/nitro/txt2img-f8-large/model.ckpt
# 执行脚本
python scripts/txt2img.py --prompt "a virus monster is playing guitar, oil on canvas" --ddim_eta 0.0 --n_samples 4 --n_iter 4 --scale 5.0 --ddim_steps 50
测试Inpainting任务
# 下载预训练权重
wget -O models/ldm/inpainting_big/last.ckpt https://heibox.uni-heidelberg.de/f/4d9ac7ea40c64582b7c9/?dl=1
# 执行脚本
python scripts/inpaint.py --indir data/inpainting_examples/ --outdir outputs/inpainting_results
如果生成了对应图片，说明环境配置成功
训练集 ILSVRC2012_img_train.tar 137.74GB
http://academictorrents.com/download/a306397ccf9c2ead27155983c254227c0fd938e2.torrent
验证集 ILSVRC2012_img_val.tar 6.28 GB
http://academictorrents.com/download/5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5.torrent
原文链接
下载好后，将train和val的tar压缩包上传至Linux服务器，train和val的目录分别为：
~/.cache/autoencoders/data/ILSVRC2012_train
~/.cache/autoencoders/data/ILSVRC2012_validation
如果没有这个目录手动创建一下，如下图：
Tips:
1.上面我的目录里不止有tar压缩包还有别的文件，是因为我用LDM代码运行后解压过tar文件了，自动生成了好多文件，正常情况下ILSVRC2012_train目录下 只会存在ILSVRC2012_img_train.tar这一个文件。
LDM的训练时分两步进行的，首先训练Autoencoder用于将输入的图像Encoder编码到Latent Space，经过DDIM过程后在将Latent Space解码Decoder回图像。
当环境都配置好后，执行下面代码即可训练Autoencoder
CUDA_VISIBLE_DEVICES=0 python main.py --base configs/autoencoder/autoencoder_kl_32x32x4.yaml -t --gpus 0,
"""
CUDA_VISIBLE_DEVICES和--gpus的GPU id号要一致
如果想要多卡训练，命令如下：
CUDA_VISIBLE_DEVICES=0,1 python main.py --base configs/autoencoder/autoencoder_kl_32x32x4.yaml -t --gpus 0,1
"""
LDM仓库中，作者只提供了训练 KL-regularized autoencoder 的config文件，对于训练 VQ-regularized autoencoder请参考taming-transformers
解释一下LDM仓库中符号的含义：
autoencoder_kl_8x8x64(f=32, d=64)，f=32表示下采样因子，d表示什么含义，我暂时还没搞明白，大家可以告送我，或者等我搞明白了我再补充上
这篇文章《Stable Diffusion 图像生成 攻略二》提到，训练AE和LDM的数据集和配置最好都一样，以避免出错
因为LDM仓库中并未提供在ImageNet数据集上使用KL-regularized训练LDM的config文件，因此，我根据理解自己创建了一个config文件imagenet-ldm-kl-8.yaml，如下：
model:
base_learning_rate: 5.0e-5 # set to target_lr by starting main.py with '--scale_lr False'
target: ldm.models.diffusion.ddpm.LatentDiffusion
params:
linear_start: 0.0015
linear_end: 0.0155
num_timesteps_cond: 1
log_every_t: 200
timesteps: 1000
loss_type: l1
first_stage_key: "image"
cond_stage_key: "image"
image_size: 32
channels: 4
cond_stage_trainable: False
concat_mode: False
scale_by_std: True
monitor: 'val/loss_simple_ema'
scheduler_config: # 10000 warmup steps
target: ldm.lr_scheduler.LambdaLinearScheduler
params:
warm_up_steps: [10000]
cycle_lengths: [10000000000000]
f_start: [1.e-6]
f_max: [1.]
f_min: [ 1.]
unet_config:
target: ldm.modules.diffusionmodules.openaimodel.UNetModel
params:
image_size: 32
in_channels: 4
out_channels: 4
model_channels: 192
attention_resolutions: [ 1, 2, 4, 8 ] # 32, 16, 8, 4
num_res_blocks: 2
channel_mult: [ 1,2,2,4,4 ] # 32, 16, 8, 4, 2
num_heads: 8
use_scale_shift_norm: True
resblock_updown: True
first_stage_config:
target: ldm.models.autoencoder.AutoencoderKL
params:
embed_dim: 4
monitor: "val/rec_loss"
ckpt_path: "models/first_stage_models/kl-f8/model.ckpt" # 修改为自己训练好的Autoencoder模型文件所在路径
ddconfig:
double_z: True
z_channels: 4
resolution: 256
in_channels: 3
out_ch: 3
ch: 128
ch_mult: [ 1,2,4,4 ] # num_down = len(ch_mult)-1
num_res_blocks: 2
attn_resolutions: [ ]
dropout: 0.0
lossconfig:
target: torch.nn.Identity
cond_stage_config: "__is_unconditional__"
data:
target: main.DataModuleFromConfig
params:
batch_size: 64
num_workers: 12
wrap: false
train:
target: ldm.data.imagenet.ImageNetTrain
params:
config:
size: 256
validation:
target: ldm.data.imagenet.ImageNetValidation
params:
config:
size: 256
lightning:
callbacks:
image_logger:
target: main.ImageLogger
params:
batch_frequency: 5000
max_images: 8
increase_log_steps: False
trainer:
benchmark: True
创建好config文件后，执行下面代码开始LDM训练：
CUDA_VISIBLE_DEVICES=0 python main.py --base configs/latent-diffusion/imagenet-ldm-kl-8.yaml -t --gpus 0,
应该是鸽了，改做三维重建了~~~