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

news2024/10/25 4:14:52

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]