⚡ 闪电扩散 ⚡

本仓库是论文《闪电扩散：加速任何条件扩散模型实现少步图像生成》的官方实现。

<a href="http://arxiv.org/abs/2406.02347"> <img src='https://img.shields.io/badge/论文-2406.02347-green' /> </a> <a href="https://gojasper.github.io/flash-diffusion-project/"> <img src='https://img.shields.io/badge/项目-页面-blue' /> </a> <a href='https://creativecommons.org/licenses/by-nd/4.0/legalcode'> <img src="https://img.shields.io/badge/许可证-CC.BY.NC-purple" /> </a> <a href="https://huggingface.co/spaces/jasperai/flash-diffusion"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20演示-FlashPixart-orange' /> </a> <a href="https://huggingface.co/spaces/jasperai/flash-sd3"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20演示-FlashSD3-orange' /> </a> <a href="https://huggingface.co/spaces/jasperai/flash-lora"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20演示-FlashLoRAs-orange' /> </a> <a href="https://huggingface.co/jasperai/flash-sd"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20检查点-FlashSD-yellow' /> </a> <a href="https://huggingface.co/jasperai/flash-sdxl"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20检查点-FlashSDXL-yellow' /> </a> <a href="https://huggingface.co/jasperai/flash-pixart"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20检查点-FlashPixart-yellow' /> </a> <a href="https://huggingface.co/jasperai/flash-sd3"> <img src='https://img.shields.io/badge/%F0%9F%A4%97%20检查点-FlashSD3-yellow' /> </a> </a> <a href="https://huggingface.co/jasperai/flash-sdxl/tree/main/comfy"> <img src='https://img.shields.io/badge/Comfy-FlashSDXL-black' /> </a> <figure> <img style="width:600px;" src="https://yellow-cdn.veclightyear.com/ab5030c0/e2bf9ed2-ef0b-4d57-b956-112718eb942c.jpg"> <figcaption> 使用4个NFE生成的图像 </figcaption> </figure>

在本论文中，我们提出了一种高效、快速、多功能且与LoRA兼容的蒸馏方法来加速预训练扩散模型的生成：闪电扩散。该方法在COCO 2014和COCO 2017数据集上的少步图像生成方面，在FID和CLIP-Score方面达到了最先进的性能，同时仅需要几个GPU小时的训练，且可训练参数比现有方法更少。除了其高效性外，该方法的多功能性还体现在多个任务中，如文本到图像、图像修复、换脸、超分辨率，以及使用不同的扩散模型骨干网络，无论是使用基于UNet的去噪器（SD1.5、SDXL）还是DiT（Pixart-α），以及适配器。在所有情况下，该方法都能在保持高质量图像生成的同时，大幅减少采样步骤的数量。

方法

我们的方法旨在创建一种快速、可靠且适应性强的方法，适用于各种用途。我们提出训练一个学生模型，在单步中预测教师模型对损坏输入样本的多步预测去噪结果。此外，我们从一个可调整的分布中采样时间步，该分布在训练过程中会发生变化，以帮助学生模型针对特定时间步。

结果

<it>闪电扩散</it>与各种骨干网络兼容，例如

Flash Stable Diffusion 3，从Stable Diffusion 3教师模型蒸馏而来
Flash SDXL，从SDXL教师模型蒸馏而来
Flash Pixart (DiT)，从Pixart-α教师模型蒸馏而来
Flash SD，从SD1.5教师模型蒸馏而来

它还可以用于以免训练的方式加速现有的LoRA。查看此部分了解更多信息。

为文本到图像变换骨干网络

Flash Stable Diffusion 3 (MMDiT)

<figure> <img style="width:600px;" src="https://yellow-cdn.veclightyear.com/ab5030c0/5126d5c1-0eea-41f9-9786-46e7ae7c2ebc.jpg"> <figcaption> 使用4个NFEs生成的图像 </figcaption> </figure>

Flash SDXL (UNet)

<figure> <img style="width:600px;" src="https://yellow-cdn.veclightyear.com/ab5030c0/9dae8841-eaca-4ebe-9c4b-fdb5ff4a651e.jpg"> <figcaption> 使用4个NFEs生成的图像 </figcaption> </figure>

Flash Pixart (DiT)

<figure> <img style="width:600px;" src="https://yellow-cdn.veclightyear.com/ab5030c0/4c9dc010-9b33-4966-b136-462a18912baf.jpg"> <figcaption> 使用4个NFEs生成的图像 </figcaption> </figure>

Flash SD

<figure> <img style="width:600px;" src="https://yellow-cdn.veclightyear.com/ab5030c0/b244cb4a-860d-4dba-8498-a04f007fc16f.jpg"> <figcaption> 使用4个NFEs生成的图像 </figcaption> </figure> </details>

各种用例

无需训练的LoRA加速

SDXL LoRAs

<figure> <img style="width:600px;" src="https://yellow-cdn.veclightyear.com/ab5030c0/4b2e93bb-7c57-4bd6-8f78-b3e39926fb5b.jpg"> <figcaption> 使用4个NFEs生成的图像 </figcaption> </figure> </details>

设置

要开始使用，你首先需要创建一个至少安装了python3.10的虚拟环境并激活它

使用`venv`

python3.10 -m venv envs/flash_diffusion
source envs/flash_diffusion/bin/activate

使用`conda`

conda create -n flash_diffusion python=3.10
conda activate flash_diffusion

然后安装所需的依赖项（如果在GPU上）并以可编辑模式安装该仓库

pip install --upgrade pip
pip install -r requirements.txt
pip install -e .

蒸馏现有的T2I模型

论文主要实验的复现脚本位于examples目录中。我们提供了4个不同的脚本：

train_flash_sd3.py：蒸馏SD3模型
train_flash_sdxl.py：蒸馏SDXL模型
train_flash_pixart：蒸馏Pixart-α模型
train_flash_canny_adapter.py：蒸馏T2I Canny适配器
train_flash_sd.py：蒸馏SD1.5模型

在examples\configs中，你会找到与每个脚本相关的配置yaml文件。你唯一需要做的就是修改yaml中的SHARDS_PATH_OR_URLS部分，以便模型在你自己的数据上进行训练。请注意，这个包使用webdataset来处理数据流，因此你使用的URL应该按照webdataset格式进行格式化。特别是对于这4个示例，每个样本需要由一个包含图像的jpg文件和一个包含标题（在caption键下）和图像美学评分（aesthetic_score）的json文件组成：

sample = {
    "jpg": dummy_image,
    "json": {
        "caption": "dummy caption",
        "aesthetic_score": 6.0
    }
}

然后可以通过简单运行以下命令来启动这些脚本

# 设置要使用的 GPU 和节点数量
export SLURM_NPROCS=1
export SLURM_NNODES=1

# 蒸馏 SD1.5
python3.10 examples/train_flash_sd.py

# 蒸馏 SDXL1.0
python3.10 examples/train_flash_sdxl.py

# 蒸馏 Pixart-α
python3.10 examples/train_flash_pixart.py

# 蒸馏 T2I Canny 适配器
python3.10 examples/train_flash_canny_adapter.py

自定义条件扩散模型蒸馏训练示例

该包还旨在支持自定义模型蒸馏。

from copy import deepcopy
from flash.models.unets import DiffusersUNet2DCondWrapper
from flash.models.vae import AutoencoderKLDiffusers, AutoencoderKLDiffusersConfig
from flash.models.embedders import (
    ClipEmbedder,
    ClipEmbedderConfig,
    ClipEmbedderWithProjection,
    ConditionerWrapper,
)

# 创建 VAE
vae_config = AutoencoderKLDiffusersConfig(
	"stabilityai/sdxl-vae" # 来自 HF Hub 的 VAE
) 
vae = AutoencoderKLDiffusers(config=vae_config)

## 创建条件嵌入器 ##
# 返回两种条件类型的 Clip 条件嵌入器
embedder_1_config = ClipEmbedderConfig(
    version="stabilityai/stable-diffusion-xl-base-1.0", # 来自 HF Hub
    text_embedder_subfolder="text_encoder_2",
    tokenizer_subfolder="tokenizer_2",
    input_key="text",
    always_return_pooled=True, # 返回一个一维张量
)
embeddder_1 = ClipEmbedder(config=embedder_1_config)

# 作用于通过拼接注入 UNET 的低分辨率图像的嵌入器
embedder_2_config = TorchNNEmbedderConfig(
    nn_modules=["torch.nn.Conv2d"],
    nn_modules_kwargs=[
       dict(
          in_channels=3,
	  out_channels=6,
          kernel_size=3,
          padding=1,
          stride=2,
       ),
    ],
    input_key="downsampled_image",
    unconditional_conditioning_rate=request.param,
)
embedder_2 = TorchNNEmbedder(config=embedder_2_config)

conditioner_wrapper = ConditionerWrapper(
    conditioners=[embedder1, embedder2]
)

# 创建教师去噪器
unet = DiffusersUNet2DCondWrapper(
    in_channels=4 + 6,  # VAE 通道 + 拼接条件
    out_channels=4,  # VAE 通道
    cross_attention_dim=1280,  # 交叉注意力条件
    projection_class_embeddings_input_dim=1280,  # 添加条件
    class_embed_type="projection",
)

# 加载教师权重
...

# 创建学生去噪器
student_denoiser = deepcopy(teacher_denoiser)

使用 Hugging Face 管道进行推理 🤗

import torch
from diffusers import PixArtAlphaPipeline, Transformer2DModel, LCMScheduler
from peft import PeftModel

# 加载 LoRA
transformer = Transformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    subfolder="transformer",
    torch_dtype=torch.float16
)
transformer = PeftModel.from_pretrained(
    transformer,
    "jasperai/flash-pixart"
)

# 管道
pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    transformer=transformer,
    torch_dtype=torch.float16
)

# 调度器
pipe.scheduler = LCMScheduler.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    subfolder="scheduler",
    timestep_spacing="trailing",
)

pipe.to("cuda")

prompt = "在茂密的森林里，一只浣熊正在读书。"

image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]

在 ComfyUI 中使用 Flash

要在 ComfyUI 中本地使用 FlashSDXL，你需要：

确保你的 ComfyUI 安装是最新的
从 Hugging Face 下载检查点。如果你不知道怎么做，前往"Files and Version"，进入 comfy 文件夹，点击 FlashSDXL.safetensors 旁边的下载按钮
将新的检查点文件移动到你本地的 comfyUI/models/loras/ 文件夹在 sd_xl_base_1.0_0.9vae.safetensors 之上作为 LoRA 使用，本仓库的 examples/comfy 中提供了一个简单的 ComfyUI 工作流程 JSON 文件

免责声明：模型已经被训练为使用 cfg scale 为 1 和 lcm 调度器，但参数可以稍微调整。

将 Flash Diffusion 与现有 LoRA 结合 🎨

Flash Diffusion 模型还可以与现有的 LoRA 结合，以无需训练的方式实现少步生成。它们可以直接集成到 Hugging Face 管道中。示例如下。

from diffusers import DiffusionPipeline, LCMScheduler
import torch

user_lora_id = "TheLastBen/Papercut_SDXL"
trigger_word = "papercut"

flash_lora_id = "jasperai/flash-sdxl"

# 加载管道
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    variant="fp16"
)

# 设置调度器
pipe.scheduler = LCMScheduler.from_config(
    pipe.scheduler.config
)

# 加载 LoRA
pipe.load_lora_weights(flash_lora_id, adapter_name="flash")
pipe.load_lora_weights(user_lora_id, adapter_name="lora")

pipe.set_adapters(["flash", "lora"], adapter_weights=[1.0, 1.0])
pipe.to(device="cuda", dtype=torch.float16)

prompt = f"{trigger_word} 一只可爱的柯基"

image = pipe(
    prompt,
    num_inference_steps=4,
    guidance_scale=0
).images[0]

许可证

本代码根据 Creative Commons BY-NC 4.0 许可证发布。

引用

如果你觉得这项工作有用或在研究中使用了它，请考虑引用我们

@misc{chadebec2024flash,
      title={Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation}, 
      author={Clement Chadebec and Onur Tasar and Eyal Benaroche and Benjamin Aubin},
      year={2024},
      eprint={2406.02347},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}