Stability AI が新たに公開した「Stable Cascade」を簡単にローカル環境で試す方法

はじめに

Stability AIは以下のような画像生成AIを公開してきました。
「Stable Diffusion 1.x」→「Stable Diffusion 2.x」→「SDXL」→「SDXL Turbo」

今回新たに「Stable Cascade」というモデルを公開しました。

「SDXL」より高速かつ高品質に画像生成ができるようです。

Web上にデモを公開してくれている人もいますが、今回は自身のPC環境で実行するための記事です。

方法

1. Pythonの環境を構築する
2. Pythonスクリプトを実行する

以上の2ステップだけです。モデルのダウンロードなど事前準備は不要です。

Python環境構築

pip install torch==2.2.1+cu118 --index-url https://download.pytorch.org/whl/cu118
pip install diffusers[torch]
pip install transformers

Pythonスクリプトの実行

たったこれだけのスクリプトです。

import torch
from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline

prior = StableCascadePriorPipeline.from_pretrained(
    "stabilityai/stable-cascade-prior",
    torch_dtype=torch.bfloat16,
    variant="bf16",
).to("cuda")

decoder = StableCascadeDecoderPipeline.from_pretrained(
    "stabilityai/stable-cascade",
    torch_dtype=torch.bfloat16,
    variant="bf16",
).to("cuda")

num_images_per_prompt = 1
prompt = "a cat playing with a white ball in the forest"
negative_prompt = "poor quality, worst quality"

prior_output = prior(
    prompt=prompt,
    height=1024,
    width=1024,
    negative_prompt=negative_prompt,
    guidance_scale=4.0,
    num_images_per_prompt=num_images_per_prompt,
    num_inference_steps=20
)
decoder_output = decoder(
    image_embeddings=prior_output.image_embeddings,
    prompt=prompt,
    negative_prompt=negative_prompt,
    guidance_scale=0.0,
    output_type="pil",
    num_inference_steps=10
).images[0]

decoder_output.save("result.png")

結果

PC環境

Windows 11
CUDA 11.8
Python 3.11