FLUX、Python、ディフューザーを使用した AI を活用した画像生成 API サービスの作成

FLUX (Black Forest Labs による) は、ここ数か月で AI 画像生成の世界を席巻しました。多くのベンチマークで Stable Diffusion (以前のオープンソースの王様) を上回っただけでなく、いくつかの指標では Dall-E や Midjourney などの独自モデルも上回りました。

しかし、あなたのアプリの 1 つで FLUX を使用するにはどうすればよいでしょうか? Replicate などのサーバーレス ホストの使用を考える人もいるかもしれませんが、これらはすぐに非常に高価になる可能性があり、必要な柔軟性を提供できない可能性があります。そこで、独自のカスタム FLUX サーバーを作成することが役立ちます。

この記事では、Python を使用して独自の FLUX サーバーを作成する手順を説明します。このサーバーを使用すると、単純な API を介してテキスト プロンプトに基づいて画像を生成できます。このサーバーを個人使用で実行する場合でも、運用アプリケーションの一部として展開する場合でも、このガイドは開始に役立ちます。

前提条件

コードに入る前に、必要なツールとライブラリが設定されていることを確認してください。

• Python: マシンに Python 3 がインストールされている必要があります。バージョン 3.10 が望ましいです。
• torch: FLUX を実行するために使用する深層学習フレームワーク。
• ディフューザー: FLUX モデルへのアクセスを提供します。
• トランスフォーマー: ディフューザーの必須の依存関係。
• 文: FLUX トークナイザーを実行するために必要です
• protobuf: FLUX
を実行するために必要です
• 加速: 場合によっては、FLUX モデルをより効率的にロードできるようになります。
• fastapi: 画像生成リクエストを受け入れることができる Web サーバーを作成するフレームワーク。
• uvicorn: FastAPI サーバーを実行するために必要です。
• psutil: マシン上の RAM の量を確認できます。

次のコマンドを実行すると、すべてのライブラリをインストールできます: pip install torch diffusersTransformers Sentencepiece protobuf Accelerate fastapi uvicorn.

M1 または M2 チップを搭載した Mac を使用している場合は、最適なパフォーマンスを得るために Metal で PyTorch をセットアップする必要があります。続行する前に、公式 PyTorch with Metal ガイドに従ってください。

GPU デバイスで FLUX を実行する予定の場合は、少なくとも 12 GB の VRAM があることを確認する必要もあります。または、CPU/MPS で実行する場合は少なくとも 12 GB の RAM (速度は遅くなります)。

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

使用しているハードウェアに基づいて推論を実行する適切なデバイスを選択して、スクリプトを開始しましょう。

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

cpu、cuda (NVIDIA GPU の場合)、または mps (Apple の Metal Performance Shaders の場合) を指定できます。次にスクリプトは、選択したデバイスが利用可能かどうかを確認し、利用できない場合は例外を発生させます。

ステップ 2: FLUX モデルのロード

次に、FLUX モデルを読み込みます。モデルを fp16 精度でロードします。これにより、品質をあまり損なうことなくメモリを節約できます。

この時点で、FLUX モデルはゲートされているため、HuggingFace での認証を求められる場合があります。認証を成功させるには、HuggingFace アカウントを作成し、モデル ページに移動して規約に同意し、アカウント設定から HuggingFace トークンを作成し、それを HF_TOKEN 環境変数としてマシンに追加する必要があります。

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

ここでは、ディフューザー ライブラリを使用して FLUX モデルをロードしています。私たちが使用しているモデルは black-forest-labs/FLUX.1-dev で、fp16 精度でロードされています。

クローズドソースの FLUX Pro モデルと同様に、より高速な推論を備えていますが、詳細度の低い画像を出力する FLUX Schnell という名前のタイムステップ蒸留モデルもあります。
ここではオイラー スケジューラを使用しますが、これを試してみることもできます。スケジューラーの詳細については、こちらをご覧ください。
画像の生成はリソースを大量に消費する可能性があるため、特にメモリが限られた CPU またはデバイスで実行する場合は、メモリ使用量を最適化することが重要です。

from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline
import psutil

model_name = "black-forest-labs/FLUX.1-dev"

print(f"Loading {model_name} on {device}")

pipeline = FluxPipeline.from_pretrained(
      model_name,

      # Diffusion models are generally trained on fp32, but fp16
      # gets us 99% there in terms of quality, with just half the (V)RAM
      torch_dtype=torch.float16,

      # Ensure we don't load any dangerous binary code
      use_safetensors=True

      # We are using Euler here, but you can also use other samplers
      scheduler=FlowMatchEulerDiscreteScheduler()
).to(device)

このコードは、使用可能な合計メモリをチェックし、システムの RAM が 64 GB 未満の場合にアテンション スライスを有効にします。アテンション スライスにより、画像生成中のメモリ使用量が削減されます。これは、リソースが限られているデバイスにとって不可欠です。

ステップ 3: FastAPI を使用して API を作成する

次に、画像を生成するための API を提供する FastAPI サーバーをセットアップします。

# Recommended if running on MPS or CPU with < 64 GB of RAM
total_memory = psutil.virtual_memory().total
total_memory_gb = total_memory / (1024 ** 3)
if (device == 'cpu' or device == 'mps') and total_memory_gb < 64:
      print("Enabling attention slicing")
      pipeline.enable_attention_slicing()

FastAPI は、Python で Web API を構築するための人気のあるフレームワークです。この場合、これを使用して、画像生成のリクエストを受け入れることができるサーバーを作成します。また、GZip ミドルウェアを使用して応答を圧縮しています。これは、画像を Base64 形式で送り返すときに特に便利です。

実稼働環境では、CDN やその他の最適化を利用するために、生成されたイメージを S3 バケットまたは他のクラウド ストレージに保存し、base64 でエンコードされた文字列の代わりに URL を返すことができます。

ステップ 4: リクエスト モデルの定義

次に、API が受け入れるリクエストのモデルを定義する必要があります。

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, conint, confloat
from fastapi.middleware.gzip import GZipMiddleware
from io import BytesIO
import base64

app = FastAPI()

# We will be returning the image as a base64 encoded string
# which we will want compressed
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

この GenerateRequest モデルは、画像の生成に必要なパラメーターを定義します。プロンプト フィールドは、作成するイメージの説明テキストです。他のフィールドには、画像の寸法、推論ステップの数、バッチ サイズが含まれます。

ステップ 5: イメージ生成エンドポイントの作成

次に、イメージ生成リクエストを処理するエンドポイントを作成しましょう。

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

このエンドポイントは画像生成プロセスを処理します。まず、FLUX の要求に従って、高さと幅が 8 の倍数であることを検証します。次に、提供されたプロンプトに基づいて画像を生成し、base64 でエンコードされた文字列として返します。

ステップ 6: サーバーの起動

最後に、スクリプトの実行時にサーバーを起動するコードを追加しましょう。

from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline
import psutil

model_name = "black-forest-labs/FLUX.1-dev"

print(f"Loading {model_name} on {device}")

pipeline = FluxPipeline.from_pretrained(
      model_name,

      # Diffusion models are generally trained on fp32, but fp16
      # gets us 99% there in terms of quality, with just half the (V)RAM
      torch_dtype=torch.float16,

      # Ensure we don't load any dangerous binary code
      use_safetensors=True

      # We are using Euler here, but you can also use other samplers
      scheduler=FlowMatchEulerDiscreteScheduler()
).to(device)

このコードは、ポート 8000 で FastAPI サーバーを起動し、http://localhost:8000 からだけでなく、0.0.0.0 バインディングのおかげで、ホスト マシンの IP アドレスを使用して同じネットワーク上の他のデバイスからもアクセスできるようにします。

ステップ 7: サーバーをローカルでテストする

FLUX サーバーが起動して実行されているので、テストしてみましょう。 HTTP リクエストを行うためのコマンドライン ツールであるcurlを使用して、サーバーと対話できます。

# Recommended if running on MPS or CPU with < 64 GB of RAM
total_memory = psutil.virtual_memory().total
total_memory_gb = total_memory / (1024 ** 3)
if (device == 'cpu' or device == 'mps') and total_memory_gb < 64:
      print("Enabling attention slicing")
      pipeline.enable_attention_slicing()

このコマンドは、curl、jq、base64 ユーティリティがインストールされている UNIX ベースのシステムでのみ機能します。 FLUX サーバーをホストしているハードウェアによっては、完了までに最大で数分かかる場合もあります。

結論

おめでとうございます! Python を使用して独自の FLUX サーバーを正常に作成しました。この設定により、単純な API を介してテキスト プロンプトに基づいて画像を生成できます。基本の FLUX モデルの結果に満足できない場合は、特定のユースケースでのパフォーマンスをさらに向上させるためにモデルを微調整することを検討してください。

完全なコード

このガイドで使用されている完全なコードは以下にあります:

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, conint, confloat
from fastapi.middleware.gzip import GZipMiddleware
from io import BytesIO
import base64

app = FastAPI()

# We will be returning the image as a base64 encoded string
# which we will want compressed
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

以上がFLUX、Python、ディフューザーを使用した AI を活用した画像生成 API サービスの作成の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。