license: apache-2.0 library_name: diffusers pipeline_tag: text-to-image

层级蒸馏3.5大模型

层级蒸馏（SwD）是一种通过逐步增加生成过程中的空间分辨率来加速扩散模型（DMs）的新框架。
与全分辨率模型相比，SwD实现了显著的加速（2.5倍至10倍），同时保持甚至提高了图像质量。

项目页面: https://yandex-research.github.io/swd
GitHub: https://github.com/yandex-research/swd
演示: https://huggingface.co/spaces/dbaranchuk/Scale-wise-Distillation

使用方法

升级至最新版本的🧨 diffusers库

pip install -U diffusers

然后您可以运行
（可能需要指定可见设备：%env CUDA_VISIBLE_DEVICES=0，以正确加载LoRAs。）

import torch
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large",
                                                torch_dtype=torch.float16,
                                                custom_pipeline='quickjkee/swd_pipeline')
pipe = pipe.to("cuda")
lora_path = 'yresearch/swd-large-4-steps'
pipe.transformer = PeftModel.from_pretrained(
    pipe.transformer,
    lora_path,
)
generator = torch.Generator().manual_seed(1)
prompt = '一只猫举着标牌，上面写着“更快采样”'
sigmas = [1.0000, 0.8959, 0.7371, 0.6022, 0.0000]
scales = [64, 80, 96, 128]
images = pipe(
    prompt,
    sigmas=torch.tensor(sigmas).to('cuda'),
    timesteps=torch.tensor(sigmas[:-1]).to('cuda') * 1000,
    scales=scales,
    guidance_scale=0.0,
    height=int(scales[0] * 8),
    width=int(scales[0] * 8),
    generator=generator,
).images

引用

@article{starodubcev2025swd,
  title={扩散模型的层级蒸馏},
  author={Nikita Starodubcev and Denis Kuznedelev and Artem Babenko and Dmitry Baranchuk},
  journal={arXiv预印本 arXiv:2503.16397},
  year={2025}
}