Stable Diffusionをクラウド環境(Modal)で動かす

手元の PC で stable-diffusion-webui を動かしましたが、画像の生成速度が遅いため、より高いスペックで実行できるクラウド環境を使ってみることにしました。

また、今回は webui じゃなく、CLI でシュッと画像生成できるよう、Modal のチュートリアルを参考に Python でスクリプトを書きました。

Modal

Modal はモデル推論、バッチ処理、タスクキュー、ウェブアプリケーションなどをクラウド上で実行できるサービス。

利用してみた所管ですが、AWS Lambda のような、serverless アーキテクチャに近いサービスという印象。

料金は月 10$まで無料(2023/05/21時点)。私がアカウントを申請したタイミングでは、30$分の無料クレジットがもらえました。

無料枠を超える場合は従量課金になるみたい。デフォルトで 30$の上限が設定されていますが、念のため定期的にチェックしておいたほうが良さそうです。

非常に多くの画像を生成するヘビーユーザーは、Google Colab や Paperspace など他の定額制のサービスと比較して選ぶのが良いと思います。

事前準備

  1. Modal のアカウント申請
  2. API トークンの発行

Modalでアカウント申請。waitlist に登録されて、半日ほどでアカウントが払い出されます。

API トークンの発行

アカウントが払い出されたあと、Modal のホーム画面に Getting Started があり、API トークンの発行から動作確認までの手順があります。

Stable Diffusion CLI を Modal 上で実行するスクリプト

GitHub のリポジトリを公開しました。下記の手順で画像が生成され、outputs ディレクトリに出力されます。

  1. リポジトリをgit clone
  2. .env を開き huggingface の API トークンとモデルを設定
  3. Makefile を開いてプロンプトを設定
  4. make run

スクリプトの説明

ディレクトリ構成

.
├── .env # 環境変数
├── Dockerfile # Modalで動作するベースのdockerイメージ
├── Makefile # スクリプト実行用
├── outputs/ # 画像ファイル置き場
├── requirements.txt # 必要なパッケージリスト
└── sd_cli.py # スクリプト本体

.env

スクリプト本体で利用する環境変数。利用したいモデルの設定をここで行います。

詳細
HUGGINGFACE_TOKEN=""
MODEL_REPO_ID="runwayml/stable-diffusion-2-1"
MODEL_NAME="stable-diffusion-2-1"

Dockerfile

詳細
FROM python:3.11.3-slim-bullseye
COPY requirements.txt /
RUN apt update \
    && apt install -y wget git \
    && pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 --pre xformers

Makefile

コマンドライン引数が多いので、Makefile で実行するようにしました。

詳細
run:
    modal run sd_cli.py \
    --prompt "A woman with bob hair" \
    --n-prompt "" \
    --height 512 \
    --width 768 \
    --samples 5

requirements.txt

Stable Diffusion を動かすための必須パッケージのリスト。バージョン管理が不十分なので、後ほど整理していきます。

詳細
accelerate
scipy
diffusers[torch]
safetensors
torch==2.0.1+cu117
torchvision
torchmetrics
omegaconf
transformers

sd_cli.py

Stable Diffusion を実行するスクリプト。

スクリプト詳細
from __future__ import annotations
import io
import os
import time
from datetime import date
from pathlib import Path
from modal import Image, Secret, Stub, method, Mount

stub = Stub("stable-diffusion-cli")
BASE_CACHE_PATH = "/vol/cache"


def download_models():
    """
    Downloads the model from Hugging Face and saves it to the cache path using
    diffusers.StableDiffusionPipeline.from_pretrained().
    """
    import diffusers

    hugging_face_token = os.environ["HUGGINGFACE_TOKEN"]
    model_repo_id = os.environ["MODEL_REPO_ID"]
    cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])

    scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
        model_repo_id,
        subfolder="scheduler",
        use_auth_token=hugging_face_token,
        cache_dir=cache_path,
    )
    scheduler.save_pretrained(cache_path, safe_serialization=True)

    pipe = diffusers.StableDiffusionPipeline.from_pretrained(
        model_repo_id,
        use_auth_token=hugging_face_token,
        cache_dir=cache_path,
    )
    pipe.save_pretrained(cache_path, safe_serialization=True)


stub_image = Image.from_dockerfile(
    path="./Dockerfile",
    context_mount=Mount.from_local_file("./requirements.txt"),
).run_function(
    download_models,
    secrets=[Secret.from_dotenv(__file__)],
)
stub.image = stub_image


@stub.cls(gpu="A10G", secrets=[Secret.from_dotenv(__file__)])
class StableDiffusion:
    """
    A class that wraps the Stable Diffusion pipeline and scheduler.
    """

    def __enter__(self):
        import diffusers
        import torch

        cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
        if os.path.exists(cache_path):
            print(f"The directory '{cache_path}' exists.")
        else:
            print(f"The directory '{cache_path}' does not exist. Download models...")
            download_models()

        torch.backends.cuda.matmul.allow_tf32 = True

        scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
            cache_path,
            subfolder="scheduler",
        )

        self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
            cache_path,
            scheduler=scheduler,
            custom_pipeline="lpw_stable_diffusion",
        ).to("cuda")
        self.pipe.enable_xformers_memory_efficient_attention()

    @method()
    def run_inference(
        self,
        prompt: str,
        n_prompt: str,
        steps: int = 30,
        batch_size: int = 1,
        height: int = 512,
        width: int = 512,
        max_embeddings_multiples: int = 1,
    ) -> list[bytes]:
        """
        Runs the Stable Diffusion pipeline on the given prompt and outputs images.
        """
        import torch

        with torch.inference_mode():
            with torch.autocast("cuda"):
                images = self.pipe(
                    [prompt] * batch_size,
                    negative_prompt=[n_prompt] * batch_size,
                    height=height,
                    width=width,
                    num_inference_steps=steps,
                    guidance_scale=7.5,
                    max_embeddings_multiples=max_embeddings_multiples,
                ).images

        image_output = []
        for image in images:
            with io.BytesIO() as buf:
                image.save(buf, format="PNG")
                image_output.append(buf.getvalue())
        return image_output


@stub.local_entrypoint()
def entrypoint(
    prompt: str,
    n_prompt: str,
    samples: int = 5,
    steps: int = 30,
    batch_size: int = 1,
    height: int = 512,
    width: int = 512,
):
    """
    This function is the entrypoint for the Runway CLI.
    The function pass the given prompt to StableDiffusion on Modal,
    gets back a list of images and outputs images to local.

    The function is called with the following arguments:
    - prompt: the prompt to run inference on
    - n_prompt: the negative prompt to run inference on
    - samples: the number of samples to generate
    - steps: the number of steps to run inference for
    - batch_size: the batch size to use
    - height: the height of the output image
    - width: the width of the output image
    """
    print(f"steps => {steps}, sapmles => {samples}, batch_size => {batch_size}")

    max_embeddings_multiples = 1
    token_count = len(prompt.split())
    if token_count > 77:
        max_embeddings_multiples = token_count // 77 + 1

    print(
        f"token_count => {token_count}, max_embeddings_multiples => {max_embeddings_multiples}"
    )

    directory = Path(f"./outputs/{date.today().strftime('%Y-%m-%d')}")
    if not directory.exists():
        directory.mkdir(exist_ok=True, parents=True)

    stable_diffusion = StableDiffusion()
    for i in range(samples):
        start_time = time.time()
        images = stable_diffusion.run_inference.call(
            prompt,
            n_prompt,
            steps,
            batch_size,
            height,
            width,
            max_embeddings_multiples,
        )
        total_time = time.time() - start_time
        print(
            f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image)."
        )
        for j, image_bytes in enumerate(images):
            formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
            output_path = directory / f"{formatted_time}_{i}_{j}.png"
            print(f"Saving it to {output_path}")
            with open(output_path, "wb") as file:
                file.write(image_bytes)
スクリプト解説

Dockerfile からイメージを構築し、そのイメージをベースに Stable Diffusion を実行します。

↓ 処理の流れ

  1. Docker イメージをビルド
    1. requirements.txt のパッケージのインストールとモデルのダウンロードを行います。
  2. 画像の生成と出力
    1. 引数で渡したプロンプトのトークンをカウントし、パラメータを調整
    2. diffusers を使って画像生成
    3. 画像生成日時をファイル名に設定し、ローカルの outputs/ディレクトリに保存

まとめ

512x512 サイズの画像なら 2~3 秒程度で画像が生成されます。すばらしいですね、、

living room

living room

living room

↓ 今後の課題

  • LoRA, Textual inversion の導入。
  • 必須パッケージのバージョン管理。
  • モデルの切り替えをもっといい感じにできるようにする。
  • img2img, upscale など、Stable Diffusion の他のコマンドを使えるようにする。
  • diffusers のコンフィグが設定されていないカスタムモデルも使えるようにする。