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

少步数文本到图像生成。

改进的分布匹配蒸馏用于快速图像生成，
田伟音、米夏埃尔·加尔比、泰松·朴、理查德·张、伊莱·谢赫特曼、弗雷多·杜兰、威廉·T·弗里曼
NeurIPS 2024 (arXiv 2405.14867)

联系方式

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

田伟音 tianweiy@mit.edu

摘要

最近的一些方法通过将扩散模型蒸馏为高效的单步生成器展现了潜力。其中，分布匹配蒸馏（DMD）能够生成在分布上与其教师模型一致的单步生成器，而无需强制与教师模型的采样轨迹一一对应。然而，为了确保训练的稳定性，DMD需要额外使用由教师模型借助确定性采样器经过多步生成的一大批噪声-图像对来计算回归损失。这对于大规模的文本到图像合成来说成本高昂，并且限制了学生模型的质量，使其过于依赖教师模型的原始采样路径。我们提出了 DMD2，这是一套可以克服这一局限并改进 DMD 训练的技术。首先，我们消除了回归损失以及构建昂贵数据集的需求。我们证明了由此产生的不稳定性是由于虚假判别器未能准确估计生成样本的分布所致，并提出了一种双时间尺度更新规则作为解决方案。其次，我们将在蒸馏过程中引入 GAN 损失，用以区分生成样本和真实图像。这使得我们能够在真实数据上训练学生模型，从而缓解教师模型对真实分数估计不准确的问题，并提升生成质量。最后，我们修改了训练流程以支持多步采样。我们识别并解决了在这种情况下训练与推理输入不匹配的问题，即在训练时模拟推理阶段的生成器样本。综合来看，我们的改进在单步图像生成领域树立了新的标杆，在 ImageNet-64x64 数据集上的 FID 得分为 1.28，在零样本 COCO 2014 数据集上的 FID 得分为 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"

# 加载模型。
with torch.device("meta"):
    unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to(torch.float16)
state_dict_path = hf_hub_download(repo_name, ckpt_name)
unet.load_state_dict(torch.load(state_dict_path), assign=True)
unet.to("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="a photo of a cat"

# 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="a photo of a cat"

# 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 有下载，从安全传感器中加载文件，基础型号是“稳定AI稳定的扩散XL基底1.0”，仓库名称是“天伟义DMD2”，检查名是“dmd2_sdxl_1step_unet_fp16.bin”。# 加载模型。
unet = UNet2DConditionModel from config base model id，subfolder “unet” to “cuda”，torch float16。Unet load state dict from hf hub download repo name，ckpt name，map location “cuda”。Pipe = Diffusion Pipeline from pre-trained base model id，unet = unet，torch dtype = torch float16，variant “fp16”。To “cuda”。Pipe scheduler = LCMScheduler from config pipe scheduler’s config。Prompt = “a photo of a cat”。Image = pipe prompt = prompt，num inference steps = 1，guidance scale = 0，timesteps = [399]。Images[0]。

4 步 T2I Adapter

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 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。

许可证

改进型分布匹配蒸馏技术根据 知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议 发布。

已知问题

• 目前用于 SDXL 训练的 FSDP 效率非常低；非常欢迎帮助！
• 目前的 LORA 训练实际上比全量微调还要慢，并且占用的内存相同；同样非常欢迎帮助！

引用

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

@inproceedings{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},
    booktitle={NeurIPS},
    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 文件夹中。

致谢

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

DMD2 快速上手指南

DMD2 (Improved Distribution Matching Distillation) 是一个用于快速图像生成的开源工具，能够将扩散模型蒸馏为高效的单步或少数几步生成器。它在保持高质量的同时，将推理成本降低了 500 倍，支持 ImageNet 和 SDXL 等模型。

环境准备

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: 3.8
• GPU: 支持 CUDA 的 NVIDIA 显卡（运行 SDXL 建议显存 16GB+）
• 前置依赖: Conda 包管理器

安装步骤

1. 创建并激活 Conda 环境

   conda create -n dmd2 python=3.8 -y 
   conda activate dmd2 
   

2. 安装依赖库 为了获得更快的下载速度，建议配置 pip 使用国内镜像源（如清华源）。

   # 升级 anyio
   pip install --upgrade anyio -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 安装项目依赖
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 以开发模式安装当前包
   python setup.py develop
   

基本使用

以下示例展示如何使用预训练模型进行文生图（Text-to-Image）推理。请确保已下载相应的检查点文件（Checkpoint），或代码会自动从 Hugging Face 拉取。

1. 命令行快速推理 (SDXL)

4 步生成（推荐，质量更高）：

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

提示: 在演示页面中，可点击 "Use Tiny VAE for faster decoding" 以启用更小的 VAE 模型，从而显著提升解码速度并降低显存占用。

2. Python 代码调用 (Diffusers Pipeline)

如果你希望在代码中集成 DMD2，可以使用标准的 diffusers 管道。以下是加载 4 步 UNet 模型的示例：

import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
from huggingface_hub import hf_hub_download

base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
repo_name = "tianweiy/DMD2"
ckpt_name = "dmd2_sdxl_4step_unet_fp16.bin"

# 加载模型配置
with torch.device("meta"):
    unet = UNet2DConditionModel.from_config(base_model_id, subfolder="unet").to(torch.float16)

# 下载并加载权重
state_dict_path = hf_hub_download(repo_name, ckpt_name)
unet.load_state_dict(torch.load(state_dict_path), assign=True)
unet.to("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 = "a photo of a cat"

# 执行推理
# 注意：LCMScheduler 的默认时间步与训练时不同，需手动指定 timesteps
image = pipe(
    prompt=prompt, 
    num_inference_steps=4, 
    guidance_scale=0, 
    timesteps=[999, 749, 499, 249]
).images[0]

image.save("output.png")

3. 进阶功能：LoRA 与 ControlNet

DMD2 同样支持 LoRA 微调模式和 T2I Adapter (如 Canny 边缘检测)。

• LoRA 生成: 使用 pipe.load_lora_weights() 加载 .safetensors 格式的 LoRA 权重，并调用 pipe.fuse_lora()。
• ControlNet/Adapter: 结合 StableDiffusionXLAdapterPipeline 和 T2IAdapter，可在 4 步内完成条件图像生成（参考 README 中的 "4-step T2I Adapter" 代码段）。

更多预训练模型详情及训练脚本，请访问项目的 Hugging Face Repo 或查看 experiments 目录下的文档。