Stable DiffusionとFastAPIで画像生成APIを自作する:LinuxサーバーでAIをマスターしよう

Artificial Intelligence tutorial - IT technology blog
Artificial Intelligence tutorial - IT technology blog

なぜ独自の画像生成APIを構築すべきなのか?

AI画像生成機能をアプリケーションに統合しようとする際、まずDALL-E 3やMidjourneyが頭に浮かぶでしょう。確かに便利ですが、すぐに財布の紐が悲鳴を上げることになります。DALL-E 3の標準的なリクエストは1回あたり約$0.04から$0.08かかります。月に10,000リクエストあれば、数万円単位のサービス利用料が簡単に飛んでいってしまいます。

もしNVIDIA製GPUを搭載したLinuxサーバーが手元にあるなら、独自のシステムを構築するのが非常に経済的な選択です。コストを抑えられるだけでなく、モデルを自由にカスタマイズすることも可能です。Civitaiなどで公開されている特定のLoRAやチェックポイントを自由に取り入れ、独自の画像スタイルを作り出すことができます。

最大の挑戦は、オフラインで動作する Pythonスクリプトを、安定したWebサービスへと変換することです。システムはキュー(待ち行列)を適切に処理し、GPUメモリ(VRAM)をオーバーフローさせないように設計する必要があります。私はこの手法を、1日2,000枚の画像を生成する自動コンテンツ作成プロジェクトに適用しましたが、結果は非常にスムーズでした。

コアコンセプト:Stable DiffusionとFastAPI

始める前に、このスタックの3つの主要コンポーネントを確認しておきましょう。

  • Stable Diffusion (SD): オープンソースの拡散モデル。今回は、よりプロフェッショナルにモデルを制御するためにHugging Faceのdiffusersライブラリを使用します。
  • FastAPI: 非常に高速でasync(非同期)処理に優れたPythonフレームワーク。Swagger UIが標準で提供されており、数クリックでAPIのテストが可能です。
  • CUDA/PyTorch: PythonがNVIDIAグラフィックカードの強力な計算能力を活用するための基盤となるレイヤーです。

ハードウェア要件:ボトルネックを避けるために

実務的な経験から、Stable Diffusion v1.5を3〜5秒程度でレンダリングするには、最低でも以下の構成が必要です。

  • OS: Ubuntu 22.04 LTS(AIドライバーにとって最も安定しています)。
  • GPU: 最低8GBのVRAMを搭載したNVIDIA製。中古のRTX 3060などは、現在非常にコストパフォーマンスの良い選択肢です。
  • RAM: システムの遅延を防ぐために16GB以上。
  • ディスク: 20GBの空き容量(各SDモデルは通常2GBから5GBのサイズがあります)。

実践:ゼロから完全なAPIのデプロイまで

ステップ1:Linux環境のセットアップ

NVIDIAドライバーとCUDA Toolkitがインストールされていることを確認してください。その後、システムを汚さないように仮想環境へライブラリを隔離します。

# システムの更新
sudo apt update && sudo apt upgrade -y

# Pythonとvenvのインストール
sudo apt install python3-venv python3-pip -y

# 仮想環境の作成
python3 -m venv sd_env
source sd_env/bin/activate

# コアライブラリのインストール
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install diffusers transformers accelerate fastapi uvicorn python-multipart

ステップ2:画像処理スクリプト(Worker)の作成

モデルを管理するためのworker.pyファイルを作成します。ここでのコツはfp16(16ビット浮動小数点数)を使用することです。これにより、生成される画像の品質をほぼ維持したまま、VRAMの消費量を50%削減できます。

import torch
from diffusers import StableDiffusionPipeline
import io
from PIL import Image

class ImageGenerator:
    def __init__(self):
        self.model_id = "runwayml/stable-diffusion-v1-5"
        # モデルをGPUに直接ロード
        self.pipe = StableDiffusionPipeline.from_pretrained(
            self.model_id, 
            torch_dtype=torch.float16
        )
        self.pipe = self.pipe.to("cuda")
        # VRAMの少ないカード向けにメモリを最適化
        self.pipe.enable_attention_slicing()

    def generate(self, prompt: str):
        with torch.autocast("cuda"):
            image = self.pipe(prompt).images[0]
        
        # ネットワーク経由で送信するために画像をバイト形式で出力
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        return img_byte_arr.getvalue()

ステップ3:FastAPIによるREST APIの構築

次に、上記のロジックをエンドポイントとしてラップします。main.pyはユーザーからのリクエストを受け取るゲートウェイの役割を果たします。

Share: