Flash Sd3

🚀 ⚡ 闪电扩散：FlashSD3 ⚡

闪电扩散（Flash Diffusion）是一种扩散蒸馏方法，由Jasper Research的Clément Chadebec、Onur Tasar、Eyal Benaroche和Benjamin Aubin在论文Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation中提出。此模型是SD3模型的一个9040万参数的LoRA蒸馏版本，能够在4步内生成1024x1024的图像。查看我们的实时演示和官方Github仓库。

🚀 快速开始

该模型可以直接使用diffusers库中的StableDiffusion3Pipeline。它可以将所需的采样步骤减少到4步。

📦 安装指南

⚠️ 重要提示

首先，你需要通过运行以下命令安装特定版本的diffusers。

pip install git+https://github.com/initml/diffusers.git@clement/feature/flash_sd3

💻 使用示例

基础用法

import torch
from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlowMatchEulerDiscreteScheduler
from peft import PeftModel

# Load LoRA
transformer = SD3Transformer2DModel.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    subfolder="transformer",
    torch_dtype=torch.float16,
)
transformer = PeftModel.from_pretrained(transformer, "jasperai/flash-sd3")


# Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    transformer=transformer,
    torch_dtype=torch.float16,
    text_encoder_3=None,
    tokenizer_3=None
)

# Scheduler
pipe.scheduler = FlashFlowMatchEulerDiscreteScheduler.from_pretrained(
  "stabilityai/stable-diffusion-3-medium-diffusers",
  subfolder="scheduler",
)

pipe.to("cuda")

prompt = "A raccoon trapped inside a glass jar full of colorful candies, the background is steamy with vivid colors."

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

🔧 技术细节

该模型在2个H100 GPU上训练了约50小时。

💡 使用建议

如果在包含文本的图像数据集上进行蒸馏，模型在处理文本方面的表现会更好，你可以自行尝试。

📄 引用

如果您发现这项工作有用或在您的研究中使用了它，请考虑引用我们：

@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}
}

📄 许可证

该模型遵循知识共享署名-非商业性使用许可协议（Creative Commons BY - NC license）。