サーバーの作成

Diffusersのpipelinesは、サーバーの推論エンジンとして使用できます。同時に複数のユーザーからリクエストされる画像の生成に対応するため、並行およびマルチスレッドのリクエストをサポートしています。

このガイドでは、サーバーでStableDiffusion3Pipelineを使用する方法を示します。ただし、お好みのパイプラインを使用しても構いません。

examples/serverフォルダに移動し、すべての依存関係をインストールすることから始めます。

pip install .
pip install -f requirements.txt

次のコマンドでサーバーを起動します。

python server.py

サーバーは http://localhost:8000 でアクセスできます。このモデルには、以下のコマンドで curl できます。

curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations

一部の依存関係のアップグレードが必要な場合は、pip-toolsまたはuvのいずれかを使用できます。例えば、以下のコマンドを使用してuvで依存関係をアップグレードします。

uv pip compile requirements.in -o requirements.txt

サーバーはFastAPIで構築されています。v1/images/generationsのエンドポイントは以下の通りです。

@app.post("/v1/images/generations")
async def generate_image(image_input: TextToImageInput):
    try:
        loop = asyncio.get_event_loop()
        scheduler = shared_pipeline.pipeline.scheduler.from_config(shared_pipeline.pipeline.scheduler.config)
        pipeline = StableDiffusion3Pipeline.from_pipe(shared_pipeline.pipeline, scheduler=scheduler)
        generator = torch.Generator(device="cuda")  # or "mps", "xpu", "cpu"
        generator.manual_seed(random.randint(0, 10000000))
        output = await loop.run_in_executor(None, lambda: pipeline(image_input.prompt, generator = generator))
        logger.info(f"output: {output}")
        image_url = save_image(output.images[0])
        return {"data": [{"url": image_url}]}
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        elif hasattr(e, 'message'):
            raise HTTPException(status_code=500, detail=e.message + traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e) + traceback.format_exc())

generate_image 関数は async キーワードで非同期として定義されており、この関数内で発生する処理が必ずしも即座に結果を返さないことを FastAPI に伝えるためです。関数内の処理が他の Task を await する必要がある地点に達すると、メインスレッドは他の HTTP リクエストの応答処理に戻ります。この動作は、以下のコードで示されています。

output = await loop.run_in_executor(None, lambda: pipeline(image_input.prompt, generator = generator))

この時点で、pipeline関数の実行は新しいスレッドに移動され、メインスレッドはpipelineから結果が返ってくるまで他の処理を行います。

この実装のもう一つの重要な側面は、shared_pipeline から pipeline を作成することです。これにより、新しいリクエストごとに独立したスレッドで動作する独自のジェネレーターとスケジューラーを確保しつつ、基盤となるモデルを GPU に複数回ロードすることを回避できます。特にスケジューラーはスレッドセーフではなく、複数のスレッドで同じスケジューラーを使用しようとすると、IndexError: index 21 is out of bounds for dimension 0 with size 21 のようなエラーが発生します。

Update on GitHub