Stable Cascade + Gradio を自作 PC で動かす

背景

Stable Diffusion を自作 PC で動かしていたりしたのですが、新しく発表のありました Stable Cascade も試してみました。公式ページのコードを参考にしながら私の自作 PC のスペックや前提に合わせて多少補正したコードを共有します。

環境

自作 PC の環境です。関係しそうなパーツのみ記載します。

• CPU: インテル® Core™ i3-12100F
• メインメモリ: DDR4 64GB
• GPU: Zotac GeForce RTX 3060 (12GB VRAM)

それにしても、1年と少し前に自作 PC を組み立てたときは、生成 AI 関連で使うとは全く考えていなかった(息子と私の PC ゲーム目的)のに、パーツ一式の相談に乗ってくれた秋葉原の某 PC ショップの方の選択(特に GPU の VRAM 12GB！)に感謝してます。

試し方

前提

• 家庭内 LAN の他の PC/スマホからはアクセスできるようにする(グローバルには公開しなくて良い)
• 生成速度に影響出るので、VRAM からはみ出ないようにする

コード

最終的にできたコードがこちらです。必要なライブラリは適宜インストールしてください。

import torch
import gradio as gr
from diffusers import StableCascadeCombinedPipeline

pipe = StableCascadeCombinedPipeline.from_pretrained(
    "stabilityai/stable-cascade",
    variant="bf16",
    torch_dtype=torch.bfloat16
)
pipe.enable_model_cpu_offload()

def txt2img(prompt, negative_prompt, width, height):
    return pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_images_per_prompt=1,
        prior_num_inference_steps=30
    ).images[0]

gr.Interface(
    fn=txt2img,
    inputs=[
        "text",
        "text",
        gr.Slider(
            minimum=256,
            maximum=1920,
            value=1024
        ),
        gr.Slider(
            minimum=256,
            maximum=1920,
            value=1024
        )
    ],
    outputs="image"
).launch(
    server_name="0.0.0.0",
    server_port=7860
)

• pip install torch で入れた場合、cuda がうまく使えない旨のエラーが出る場合があります。回避策はこちらを参考にさせていただきました。感謝です。

• huggingface の解説ページにはいくつかサンプルコードがありますが、prior と decoder が分離しているコードを流用しようとすると、Gradio で関数呼び出しを行った時にエラーになるので、StableCascadeCombinedPipeline を使ったサンプルコードを採用しました。
• pipe.to("cuda") としてしまうと、12GB の GPU のメモリでは足りず、メインメモリにはみ出すようになってしまい、こうなると如実に生成スピードに悪影響が出ます。pipe.enable_model_cpu_offload() とする事で、使用する VRAM を減らし 12GB 以内に収まり、画像生成をスムーズに行えるようになりました。
• pipe の引数はこちらを参照し、使うもの・明示的に変更しても良さそうなものを追加しました。

• launch() への引数として、servername=0.0.0.0 / server_port=7860 を指定した事で、同じネットワーク内のパソコンやスマホから http://192.168.0.2:7860 とかでもアクセスできるようにしています。

参考: メモリ使用量の違い

pipe.to("cuda") の場合

共有 GPU メモリが使用されてしまっています。

pipe.enable_model_cpu_offload()

こちらでは無事専用 GPU メモリにおさまりました。

実行結果例

1024 x 1024 の画像生成で1枚概ね 25〜30 秒、768 x 768 だと 20秒程度で生成できています。

Two owls perching on a tree branch in a deep mountain forest.

generative ai

個人的には好きな画風です。

おわりに

そこまでハイスペックというわけでもない自宅の PC でも、悪くない生成スピードで動いてくれて、このクオリティはすごいなと素直に驚きました。

2024/4/16 に確認した時点では非商用での利用のみということで、普段の仕事で使ったりできないのは残念ですが、Stable Diffusion 3 も控えているようですし、引き続き楽しみな分野です。