ByteDance が公開した SDXL-Lightning を使ってみる

huggingface.co
Diffusersから簡単に使えます。

サンプルスクリプトをみるとSDXLのUnet部分を変更しているだけに見えます。

Pythonスクリプト

import torch
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
from safetensors.torch import load_file

# from https://huggingface.co/ByteDance/SDXL-Lightning
ckpt = "sdxl_lightning_8step_unet.safetensors" 

pipe = StableDiffusionXLPipeline.from_pretrained(
    "model/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda")

unet = UNet2DConditionModel.from_config(pipe.unet.config).to("cuda", torch.float16)
unet.load_state_dict(load_file(ckpt, device="cuda"))
pipe.unet = unet
pipe.scheduler = EulerDiscreteScheduler.from_config(
    pipe.scheduler.config,
    timestep_spacing="trailing"
)

for i in range(10):
    seed = 2024 + i * 224
    generator = torch.manual_seed(seed)
    image = pipe(
        "A girl smiling",
        num_inference_steps=8,
        guidance_scale=0,
        generator=generator
    ).images[0]
    image.save(f"result_{seed}.png")