Scale-wise Distillation 3.5 Medium

Using with Diffusers

Upgrade to the latest version of the 🧨 diffusers library

pip install -U diffusers

and then you can run

import torch
from diffusers import StableDiffusion3Pipeline
from peft import LoraConfig, get_peft_model, PeftModel

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")

lora_path = 'yresearch/swd-medium-6-steps '
pipe.transformer = PeftModel.from_pretrained(
    pipe.transformer,
    distill_check,
)

generator = torch.Generator().manual_seed(0)
prompt = 'A castle made of tortilla chips, in a river made of salsa. There are tiny burritos walking around the castle'
sigmas = [1.0000, 0.9454, 0.8959, 0.7904, 0.7371, 0.6022]
scales = [32, 48, 64, 80, 96, 128]

images = run(
        pipe,
        prompt,
        sigmas=sigmas,
        scales=scales,
        num_inference_steps=6,
    
        guidance_scale=0.0,
        height=int(scales[0] * 8),
        width=int(scales[0] * 8),
        generator=generator,
    ).images