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

## 背景

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

https://ja.stability.ai/blog/stable-cascade

https://huggingface.co/stabilityai/stable-cascade

## 環境

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

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

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

## 試し方

### 前提

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

### コード

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

```python
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 がうまく使えない旨のエラーが出る場合があります。回避策はこちらを参考にさせていただきました。感謝です。

https://qiita.com/sakusaku3939/items/1a133729c7f38e8403ce

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

https://huggingface.co/docs/diffusers/api/pipelines/stable_cascade#diffusers.StableCascadeCombinedPipeline

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



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

pipe.to("cuda") の場合
![stablecascade-cuda.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3535779/4aa7ab19-d183-c83c-cbdb-af8e7f9253d0.png)
共有 GPU メモリが使用されてしまっています。

pipe.enable_model_cpu_offload() 
![stablecascade-cpuoffload.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3535779/5a0e9c90-ddee-63c7-1689-842886337a7e.png)
こちらでは無事専用 GPU メモリにおさまりました。

## 実行結果例

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

```
Two owls perching on a tree branch in a deep mountain forest.
```
![generated-image01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3535779/c4a78e26-a5c2-c3a5-f95b-c920e4312c96.png)

```
generative ai
```
![generated-image02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3535779/81e38f3c-6040-fa15-1f5f-a1c0d8aef4ff.png)
![generated-image03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3535779/03b9500a-7b47-b2a5-b37f-41ad4d00a3ed.png)
個人的には好きな画風です。

## おわりに

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

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