SDXL 1.0 (Stable Diffusion XL 1.0) に専用VAEを適用させてみました

はじめに

こちらにSDXL専用と思われるVAEが公開されていたので使ってみました。
huggingface.co
SDXL 1.0の基本的な使い方はこちらを参照して下さい。
touch-sp.hatenablog.com

Pythonスクリプト

from diffusers import DiffusionPipeline, AutoencoderKL
import torch

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix",
    torch_dtype=torch.float16)

base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)
base.to("cuda")

refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16"
)
refiner.to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
seed = 20000
generator = torch.Generator(device="cuda").manual_seed(seed)
image = base(
    prompt=prompt,
    output_type="latent",
    generator=generator
    ).images[0]
image = refiner(
    prompt=prompt,
    image=image[None, :]).images[0]
image.save("result_2_with_vae.png")