少ないstep数で画像生成を可能にする Latent Consistency Model (LCM) を Diffusers から使う

github.com

環境

Windows 11
CUDA 11.7
Python 3.10

pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117
pip install diffusers[torch]
pip install transformers

モデルのダウンロード

あらかじめ以下でモデルをダウンロードしました。

git lfs install
git clone https://huggingface.co/latent-consistency/lcm-lora-sdxl

後述するスクリプトを以下のように変更すると事前準備不要で自動的にダウンロードされます。
変更前

pipe.load_lora_weights("lcm-lora-sdxl/pytorch_lora_weights.safetensors")

変更後

pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")

SDXLモデルもあらかじめダウンロードが済んでいる前提です。
もしダウンロードしていないなら後述するスクリプトを以下のように変更すると事前準備不要で自動的にダウンロードされます。
変更前

model_id = "model/stable-diffusion-xl-base-1.0"

変更後

model_id = "stabilityai/stable-diffusion-xl-base-1.0"

Pythonスクリプト

「guidance_scale」を1（または 0）に設定するのが重要だそうです。

from diffusers import AutoPipelineForText2Image, LCMScheduler
import torch

model_id = "model/stable-diffusion-xl-base-1.0"

pipe = AutoPipelineForText2Image.from_pretrained(
    model_id,
    torch_dtype=torch.float16
).to("cuda")

pipe.load_lora_weights("lcm-lora-sdxl/pytorch_lora_weights.safetensors")

pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux"
seed = 10000
generator = torch.manual_seed(seed)
image = pipe(
    prompt=prompt,
    num_inference_steps=4,
    guidance_scale=1,
    generator=generator
).images[0]

image.save("lcm_result.png")

結果の比較

LCMを使う場合と使わない場合とを比較してみました。

比較した結果

上がLCMなし、下がLCMありです。
左からstep数を4→8→12→20→30と変更しています。

比較するためのPythonスクリプト

from diffusers import AutoPipelineForText2Image, LCMScheduler
import torch
import os

model_id = "model/stable-diffusion-xl-base-1.0"

pipe = AutoPipelineForText2Image.from_pretrained(
    model_id,
    torch_dtype=torch.float16
).to("cuda")

prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux"
seed = 10000

os.makedirs("normal_result", exist_ok=True)
for step in [4, 8, 12, 20]:
    generator = torch.manual_seed(seed)
    image = pipe(
        prompt=prompt,
        num_inference_steps=step,
        generator=generator
    ).images[0]

    image.save(os.path.join("normal_result", f"result_step{step}.png"))

pipe.load_lora_weights("lcm-lora-sdxl/pytorch_lora_weights.safetensors")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

os.makedirs("lcm_result", exist_ok=True)
for step in [4, 8, 12, 20]:
    generator = torch.manual_seed(seed)
    image = pipe(
        prompt=prompt,
        num_inference_steps=step,
        guidance_scale=1,
        generator=generator
    ).images[0]

    image.save(os.path.join("lcm_result", f"result_step{step}.png"))