改进的分布匹配蒸馏用于快速图像合成 [Huggingface 仓库][ComfyUI][Colab]

少步文本到图像生成。

image/jpeg

改进的分布匹配蒸馏用于快速图像合成,
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Frédo Durand, William T. Freeman
arXiv 技术报告 (arXiv 2405.14867)

联系方式

如果您对论文有任何问题，请随时与我们联系！

Tianwei Yin tianweiy@mit.edu

摘要

近期的方法在将扩散模型蒸馏成高效的单步生成器方面显示出了前景。其中，分布匹配蒸馏(DMD)生成的单步生成器在分布上与其教师模型匹配，而无需与教师模型的采样轨迹强制一一对应。然而，为了确保训练稳定，DMD需要一个额外的回归损失，该损失使用教师模型通过确定性采样器多步生成的大量噪声-图像对来计算。这对大规模文本到图像合成来说成本高昂，并限制了学生模型的质量，将其过于紧密地与教师模型的原始采样路径绑定。我们引入了DMD2，一套技术来消除这一限制并改进DMD训练。首先，我们消除了回归损失和昂贵数据集构建的需求。我们表明由此产生的不稳定性是由于假评论家无法准确估计生成样本的分布，并提出了两个时间尺度更新规则作为补救措施。其次，我们将GAN损失整合到蒸馏过程中，区分生成的样本和真实图像。这使我们能够在真实数据上训练学生模型，缓解了来自教师模型的不完美真实分数估计，并提高了质量。最后，我们修改了训练程序以实现多步采样。我们识别并解决了这种情况下的训练-推理输入不匹配问题，通过在训练时模拟推理时的生成器样本。综合起来，我们的改进在单步图像生成方面创造了新的基准，在ImageNet-64x64上的FID分数为1.28，在零样本COCO 2014上为8.35，尽管推理成本降低了500倍，但仍超过了原始教师模型。此外，我们通过蒸馏SDXL展示了我们的方法可以生成百万像素级图像，在少步方法中展现出卓越的视觉质量。

环境设置

# 在conda环境中
conda create -n dmd2 python=3.8 -y 
conda activate dmd2 

pip install --upgrade anyio
pip install -r requirements.txt
python setup.py  develop

推理示例

ImageNet

python -m demo.imagenet_example  --checkpoint_path IMAGENET_CKPT_PATH

文本到图像

# 注意：在演示页面上，点击"使用Tiny VAE进行更快解码"以启用使用来自[madebyollin](https://huggingface.co/madebyollin/taesdxl)的Tiny VAE，实现更快的速度和更低的内存消耗

# 4步（质量远高于1步）
python -m demo.text_to_image_sdxl --checkpoint_path SDXL_CKPT_PATH --precision float16

# 1步 
python -m demo.text_to_image_sdxl --num_step 1 --checkpoint_path SDXL_CKPT_PATH --precision float16 --conditioning_timestep 399

我们也可以使用标准的diffuser管道：

4步UNet生成

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo_name = "tianweiy/DMD2"
ckpt_name = "dmd2_sdxl_4step_unet_fp16.bin"
# 加载模型
unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location="cuda"))
pipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
prompt="一张猫的照片"

# LCMScheduler的默认时间步与我们用于训练的不同
image=pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0, timesteps=[999, 749, 499, 249]).images[0]

4步LoRA生成

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo_name = "tianweiy/DMD2"
ckpt_name = "dmd2_sdxl_4step_lora_fp16.safetensors"
# 加载模型
pipe = DiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
pipe.fuse_lora(lora_scale=1.0)  # 对于社区模型，我们可能想要将比例设置得更小

pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
prompt="一张猫的照片"

# LCMScheduler的默认时间步与我们用于训练的不同
image=pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0, timesteps=[999, 749, 499, 249]).images[0]

1步UNet生成

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo_name = "tianweiy/DMD2"
ckpt_name = "dmd2_sdxl_1step_unet_fp16.bin"
# 加载模型
unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location="cuda"))
pipe = DiffusionPipeline.from_pretrained(base_model_id, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
prompt="一张猫的照片"
image=pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[399]).images[0]

4步T2I适配器

from diffusers import StableDiffusionXLAdapterPipeline, T2IAdapter, AutoencoderKL, UNet2DConditionModel, LCMScheduler
from diffusers.utils import load_image, make_image_grid
from controlnet_aux.canny import CannyDetector
from huggingface_hub import hf_hub_download
import torch

# 加载适配器
adapter = T2IAdapter.from_pretrained("TencentARC/t2i-adapter-canny-sdxl-1.0", torch_dtype=torch.float16, varient="fp16").to("cuda")

vae=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)

base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo_name = "tianweiy/DMD2"
ckpt_name = "dmd2_sdxl_4step_unet_fp16.bin"
# 加载模型
unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name), map_location="cuda"))

pipe = StableDiffusionXLAdapterPipeline.from_pretrained(
    base_model_id, unet=unet, vae=vae, adapter=adapter, torch_dtype=torch.float16, variant="fp16", 
).to("cuda")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()

canny_detector = CannyDetector()

url = "https://huggingface.co/Adapter/t2iadapter/resolve/main/figs_SDXLV1.0/org_canny.jpg"
image = load_image(url)

# 以低分辨率检测Canny边缘图以避免高频细节
image = canny_detector(image, detect_resolution=384, image_resolution=1024)#.resize((1024, 1024))

prompt = "真实的神秘仙子，魔法，4K图片，高质量"

gen_images = pipe(
  prompt=prompt,
  image=image,
  num_inference_steps=4,
  guidance_scale=0, 
  adapter_conditioning_scale=0.8, 
  adapter_conditioning_factor=0.5,
  timesteps=[999, 749, 499, 249]
).images[0]
gen_images.save('out_canny.png')

预训练模型可以在ImageNet和SDXL中找到。

训练和评估

ImageNet-64x64

详细信息请参考ImageNet-64x64。

SDXL

详细信息请参考SDXL。

SDv1.5

详细信息请参考SDv1.5。

许可证

改进的分布匹配蒸馏（Improved Distribution Matching Distillation）根据知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议发布。

已知问题

目前SDXL训练的FSDP非常慢；非常感谢任何帮助！
目前LORA训练实际上比完全微调更慢，并且占用相同的内存；非常感谢任何帮助！

引用

如果您发现DMD2对您的研究有用或相关，请引用我们的论文：

@article{yin2024improved,
    title={改进的分布匹配蒸馏用于快速图像生成},
    author={Yin, Tianwei and Gharbi, Micha{\"e}l and Park, Taesung and Zhang, Richard and Shechtman, Eli and Durand, Fredo and Freeman, William T},
    journal={arXiv 2405.14867},
    year={2024}
}

@inproceedings{yin2024onestep,
    title={基于分布匹配蒸馏的单步扩散},
    author={Yin, Tianwei and Gharbi, Micha{\"e}l and Zhang, Richard and Shechtman, Eli and Durand, Fr{\'e}do and Freeman, William T and Park, Taesung},
    booktitle={CVPR},
    year={2024}
}

第三方代码

EDM用于dnnlib、torch_utils和edm文件夹。

致谢

这项工作是在Tianwei Yin作为麻省理工学院全日制学生期间完成的。它是基于我们对原始DMD论文的重新实现而开发的。这项工作得到了国家科学基金会在合作协议PHY-2019786（NSF人工智能和基础交互研究所，http://iaifi.org/）下的支持，以及NSF Grant 2105819、NSF CISE award 1955864，以及来自Google、GIST、亚马逊和广达电脑的资金支持。