【SD3的Turbo也来了】Jasper AI用Flash Diffusion的蒸馏技术为SD3提速

Flash Diffusion 是一种新颖的图像生成方法，旨在显著提高现有条件扩散模型（例如教师模型）的速度，而无需牺牲性能。它通过引入以下技术来实现这一点：

• 蒸馏损失：在教师模型和学生模型之间引入蒸馏损失，以鼓励学生模型学习教师模型的输出分布。
• GAN 损失：使用生成对抗网络（GAN）损失来提高生成图像的质量。
• 高效网络架构：使用高效的网络架构来减少计算成本。
• 步骤蒸馏：使用步骤蒸馏来进一步提高效率。

通过结合这些技术，Flash Diffusion 可以实现比现有方法更快的速度，同时生成高质量的图像。事实上，它已经证明能够在移动设备上以不到 2 秒的速度生成图像。

安装依赖

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

Demo

官方

import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

image = pipe(
    "A cat holding a sign that says hello world",
    negative_prompt="",
    num_inference_steps=28,
    guidance_scale=7.0,
).images[0]
image

改进后

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]

原本的sd3 medium需要28步推理， 而它只需要4步推理，推理速度显著。

未来

Flash Diffusion 是一项具有巨大潛力的新技术。它有可能彻底改变我们生成和使用图像的方式。随着该技术的不断发展，我们可以期待看到更加令人印象深刻的应用程序出现。

以下是一些关于 Flash Diffusion 未来发展的潜在方向：

• 进一步提高速度和质量：研究人员将继续致力于提高 Flash Diffusion 的速度和质量。
• 开发新的应用程序：新的 Flash Diffusion 应用程序将被开发出来，例如视频生成和 3D 图像生成。
• 使 Flash Diffusion 更易于使用：Flash Diffusion 将变得更加易于使用，以便更多人可以利用它。

总而言之，Flash Diffusion 是一项开创性的技术，有望在未来几年对图像生成领域产生重大影响。

我会定期在CSDN分享我的学习心得，项目经验和行业动态。如果你对某个领域感兴趣，或者想要了解更多技术干货，请关注我的账号，一起成长！