diffusersでスケジューラを読み込んで画像生成する方法

はじめに

これはhuggingface/diffusersモジュールを利用して画像生成する際に、スケジューラ(Euler等)をを読み込む方法です。

Diffusersではすでに何種類か用意されているので、それを使ってパイプラインを生成する

今回はそれの備忘録

コード

import datetime
import torch
from diffusers import (
    DiffusionPipeline,
    DDIMScheduler,
    DDPMScheduler,
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    HeunDiscreteScheduler,
    KDPM2AncestralDiscreteScheduler,
    KDPM2DiscreteScheduler,
    UniPCMultistepScheduler,
)

MODEL_ID=''
SCHEDULER = 'EulerAncestralDiscrete'
CACHE_DIR = './cache'
DEVICE = 'cuda'

SCHEDULERS_LIST = {
    'DDIM' : DDIMScheduler,
    'DDPM' : DDPMScheduler,
    'DEISMultistep' : DEISMultistepScheduler,
    'DPMSolverMultistep' : DPMSolverMultistepScheduler,
    'DPMSolverSinglestep' : DPMSolverSinglestepScheduler,
    'EulerAncestralDiscrete' : EulerAncestralDiscreteScheduler,
    'EulerDiscrete' : EulerDiscreteScheduler,
    'HeunDiscrete' : HeunDiscreteScheduler,
    'KDPM2AncestralDiscrete' : KDPM2AncestralDiscreteScheduler,
    'KDPM2Discrete' : KDPM2DiscreteScheduler,
    'UniPCMultistep' : UniPCMultistepScheduler
}

# Create Pipeline
pipe = DiffusionPipeline.from_pretrained(
    pretrained_model_name_or_path=MODEL_ID,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR,
)

# Add Scheduler
pipe.scheduler = SCHEDULERS_LIST[SCHEDULER].from_pretrained(
    pretrained_model_name_or_path=MODEL_ID,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR,
    subfolder='scheduler'
)

pipe = pipe.to(DEVICE)

# Generate Image
image = pipe(
    prompt='',
    height=000,
    width=000,
    num_inference_steps=000,
    guidance_scale=000,
    negative_prompt='',
).images[0]

# Save Image
image.save("images/" + str(datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')) + ".png")

上記の下記がスケジューラ読み込み部分です

pipe.scheduler = SCHEDULERS_LIST[SCHEDULER].from_pretrained(
    pretrained_model_name_or_path=MODEL_ID,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR,
    subfolder='scheduler'
)

SCHEDULERS_LIST[SCHEDULER] の部分がクラス名の選択になります。