From 583a56f1ac3d9be9ce2f8c1715d77423347d4a4e Mon Sep 17 00:00:00 2001
From: tanaka takashi <isazaka@gmail.com>
Date: Wed, 14 Oct 2026 02:21:50 +0900
Subject: fix: 画像生成モデルのロードを初回生成まで遅延する
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

stdio MCPのツール一覧取得でも、受付開始前にModelManagerがSD本体・
外部VAE・LoRAをロードし、CUDAへ転送していた。画像生成しない通常の
質問でもこの起動コストが発生するため、有効な生成要求の実行時まで
モデル初期化を遅らせ、一覧取得時の不要な遅延とVRAM消費を除去する。

実装:
- ModelManagerの構築時は未ロード状態とし、冪等なload()を追加する。
- 初期化ロックで二重ロードを防ぎ、成功時だけモデルを保持する。
  失敗時は未ロード状態を維持し、次の要求で再試行できるようにする。
- 3種類の生成メソッドからload()を呼び、既存サービスの入力検証・
  生成ロックを経由したワーカースレッド内で初期化する。
- ControlNetと前処理モデルの既存の遅延ロード・モデル共有を維持する。
- 起動・設定読み込みと実際のモデルロードのログを分離し、ロードの
  開始・完了・失敗をstderrへ出力する。
- benchmarkの計測区間で明示的にload()を呼び、ロード時間とVRAMの
  計測の意味を維持する。READMEにロードタイミングと確認手順を追記する。

変更範囲:
- stdioのプロセス起動・終了方式を維持し、常駐化やHTTP切り替え、
  チャット側の接続管理・一覧キャッシュ・アイドル解放は追加しない。
- MCPスキーマ、設定形式、CLI引数、プロンプト・seed処理、scheduler、
  生成サイズ、PNGメタデータ・保存仕様は変更しない。
- 設定とモデルファイルの存在確認は起動時に残し、モデル互換性・
  CUDA・推論依存関係のエラー検出は初回生成時へ移す。
- 同一プロセス内ではモデルを再利用し、終了時に解放する。
  別の生成プロセスでの再ロードは引き続き許容する。

回帰テスト:
- 修正前の既存69件に対し、修正後は追加18件を含む87件すべて成功。
  既存テストと期待値は変更していない。
- CLI/stdioでの一覧取得時にモデルロードと推論ライブラリのimportを
  禁止しても成功することを確認する。
- 全生成方式の初回ロード、ワーカースレッド実行、同時要求、冪等性、
  不正入力・無効なControlNet指定、失敗後の再試行、ログを検証する。
- ruff check .は成功、pyrightは0 errors / 0 warnings。

実GPU・stdioでの確認:
- 現在の設定（外部VAE、2件のLoRA、LCM、6 steps）を維持して計測した。
- 起動から一覧取得までの時間は8.233秒から0.671秒に短縮した。
- 一覧取得前後のGPU全体の使用量は修正前808→3233 MiBに対し、
  修正後682→682 MiBで増加しなかった。
- 同じ接続で2枚生成してもベースモデルのロードは1回で、一覧取得では
  0回だった。ロード時間は7.812秒で、初回生成要求に移った。
- 同じ設定・プロンプト・seedによる修正前後の画像は画素が完全一致。
  PNG形式・寸法・メタデータと、返却PNGおよび保存PNGの一致を確認した。
- 最終版スクリプトでの再確認でも一覧取得0.655秒、初回生成14.089秒、
  2枚目1.221秒で、画素・保存結果の一致と生成プロセス終了を確認した。
- 通常質問の親アプリの処理経路でも、LLMのみをスタブ化し、実際の
  設定済みuvコマンドで一覧取得・3ツールの登録を確認した。
- 計測スクリプトと詳細な検証報告を追加する。生のJSON・stderr・PNGは
  gitignore対象のdata/verification-lazy-loadに保存する。

確認の限界:
- Windowsのプロセス別VRAMはN/Aのため、GPU全体の使用量・プロセス一覧・
  tasklistによる終了確認で評価した。全体値には他プロセスの変動がある。
- 計測値は単一セッションの観測であり、統計的な速度保証ではない。
- ブラウザ操作と実LLMを含むチャット全体、および現在の設定で無効な
  ControlNetの実GPU生成は未実施。誘導生成・Inpaintのロード順は
  スタブを用いた自動テストで確認した。
---
 tests/test_lazy_loading.py | 245 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 245 insertions(+)
 create mode 100644 tests/test_lazy_loading.py

(limited to 'tests/test_lazy_loading.py')

diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py
new file mode 100644
index 0000000..bd4b341
--- /dev/null
+++ b/tests/test_lazy_loading.py
@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import replace
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+from mcp import Client, ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from PIL import Image
+
+from mcp_sd15.errors import ConfigError
+from mcp_sd15.model_manager import ModelManager
+from mcp_sd15.pipeline_factory import PipelineFactory
+from mcp_sd15.schema import TextToImageRequest
+from mcp_sd15.server import create_server
+from mcp_sd15.service import ImageGenerationService
+
+
+@pytest.fixture
+def lazy_backend(app_config, monkeypatch):
+    loads: list[int] = []
+    calls: list[dict[str, object]] = []
+
+    def pipeline(**kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(images=[Image.new("RGB", (4, 4), "red")])
+
+    def create_base(self):
+        loads.append(threading.get_ident())
+        return pipeline
+
+    monkeypatch.setattr(PipelineFactory, "create_base", create_base)
+    backend = ModelManager(app_config)
+    monkeypatch.setattr(backend, "_embed_prompts", lambda *args: ("positive", "negative"))
+    monkeypatch.setattr(backend, "_generator", lambda seed: seed)
+
+    def preprocess(image, **kwargs):
+        assert backend.base is pipeline  # Base must load before any GPU/preprocessor work.
+        return image
+
+    monkeypatch.setattr(backend, "_depth", preprocess)
+    monkeypatch.setattr(backend, "_openpose", preprocess)
+    monkeypatch.setattr(backend, "_inpaint_condition", lambda image, mask: preprocess(image))
+    monkeypatch.setattr(backend, "_pipeline_for", lambda *args, **kwargs: pipeline)
+    return backend, loads, calls
+
+
+async def test_discovery_does_not_load_model(app_config, lazy_backend) -> None:
+    backend, loads, _ = lazy_backend
+    assert backend.base is None
+    server = create_server(ImageGenerationService(app_config, backend))
+    assert loads == []
+    async with Client(server, mode="legacy") as client:
+        tools = (await client.list_tools()).tools
+        assert len(tools) == 3
+        assert loads == []
+    assert backend.base is None
+
+
+@pytest.mark.parametrize(
+    ("tool", "arguments"),
+    [
+        ("generate_image", {}),
+        ("generate_guided_image", {"control_mode": "depth"}),
+        ("generate_guided_image", {"control_mode": "openpose"}),
+        ("generate_guided_image", {"control_mode": "openpose_full"}),
+        ("inpaint_image", {}),
+    ],
+)
+async def test_first_generation_loads_once_in_worker(
+    app_config, lazy_backend, image_base64, tool, arguments, caplog
+) -> None:
+    caplog.set_level(logging.INFO, logger="mcp_sd15.model_manager")
+    backend, loads, calls = lazy_backend
+    arguments = {"positive_prompt": "cat", **arguments}
+    if tool != "generate_image":
+        arguments["input_image"] = image_base64
+    if tool == "inpaint_image":
+        arguments["mask_image"] = image_base64
+    server = create_server(ImageGenerationService(app_config, backend))
+    async with Client(server) as client:
+        for _ in range(2):
+            result = await client.call_tool(tool, arguments)
+            assert result.is_error is False
+            assert result.content[0].type == "image"
+    assert len(loads) == 1
+    assert loads[0] != threading.get_ident()
+    assert len(calls) == 2
+    assert calls[0]["num_inference_steps"] == app_config.generation.steps
+    assert calls[0]["guidance_scale"] == app_config.generation.cfg_scale
+    assert caplog.text.count("model load started") == 1
+    assert caplog.text.count("model load completed") == 1
+
+
+async def test_simultaneous_requests_serialize_load_and_generation(
+    app_config, lazy_backend
+) -> None:
+    backend, loads, calls = lazy_backend
+    service = ImageGenerationService(app_config, backend)
+    await asyncio.gather(
+        *(service.generate_image(TextToImageRequest.create("cat")) for _ in range(4))
+    )
+    assert len(loads) == 1
+    assert len(calls) == 4
+
+
+def test_load_is_thread_safe_and_idempotent(app_config, monkeypatch) -> None:
+    loads: list[int] = []
+    pipeline = object()
+
+    def create_base(self):
+        loads.append(threading.get_ident())
+        time.sleep(0.03)
+        return pipeline
+
+    monkeypatch.setattr(PipelineFactory, "create_base", create_base)
+    backend = ModelManager(app_config)
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        list(executor.map(lambda _: backend.load(), range(8)))
+    backend.load()
+    assert len(loads) == 1
+    assert backend.base is pipeline
+
+
+@pytest.mark.parametrize(
+    ("tool", "arguments"),
+    [
+        ("generate_image", {"positive_prompt": ""}),
+        ("generate_image", {"size": "invalid"}),
+        ("generate_guided_image", {"control_mode": "invalid", "input_image": "invalid"}),
+        ("generate_guided_image", {"control_mode": "depth", "input_image": "invalid"}),
+        ("inpaint_image", {"input_image": "invalid", "mask_image": "invalid"}),
+    ],
+)
+async def test_invalid_requests_do_not_load(app_config, lazy_backend, tool, arguments) -> None:
+    backend, loads, _ = lazy_backend
+    server = create_server(ImageGenerationService(app_config, backend))
+    async with Client(server) as client:
+        result = await client.call_tool(tool, {"positive_prompt": "cat", **arguments})
+    assert result.is_error is True
+    assert loads == []
+
+
+async def test_invalid_mask_does_not_load(app_config, lazy_backend, image_base64) -> None:
+    backend, loads, _ = lazy_backend
+    async with Client(create_server(ImageGenerationService(app_config, backend))) as client:
+        result = await client.call_tool(
+            "inpaint_image",
+            {"positive_prompt": "cat", "input_image": image_base64, "mask_image": "invalid"},
+        )
+    assert result.is_error is True
+    assert loads == []
+
+
+@pytest.mark.parametrize("tool", ["generate_guided_image", "inpaint_image"])
+async def test_disabled_controlnet_does_not_load(app_config, lazy_backend, image_base64, tool):
+    backend, loads, _ = lazy_backend
+    config = replace(
+        app_config,
+        controlnet=replace(
+            app_config.controlnet,
+            depth=replace(app_config.controlnet.depth, enabled=False),
+            inpaint=replace(app_config.controlnet.inpaint, enabled=False),
+        ),
+    )
+    arguments = {"positive_prompt": "cat", "input_image": image_base64}
+    if tool == "generate_guided_image":
+        arguments["control_mode"] = "depth"
+    else:
+        arguments["mask_image"] = image_base64
+    async with Client(create_server(ImageGenerationService(config, backend))) as client:
+        result = await client.call_tool(tool, arguments)
+    assert result.is_error is True
+    assert loads == []
+
+
+async def test_load_failure_is_tool_error_and_next_call_retries(
+    app_config, lazy_backend, monkeypatch, caplog
+):
+    backend, loads, _ = lazy_backend
+    create_base = backend.factory.create_base
+    attempts = 0
+
+    def fail_once():
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            raise ConfigError("test model could not load")
+        return create_base()
+
+    monkeypatch.setattr(backend.factory, "create_base", fail_once)
+    async with Client(create_server(ImageGenerationService(app_config, backend))) as client:
+        result = await client.call_tool("generate_image", {"positive_prompt": "cat"})
+        assert result.is_error is True
+        assert "test model could not load" in caplog.text
+        assert backend.base is None
+        assert len((await client.list_tools()).tools) == 3
+        result = await client.call_tool("generate_image", {"positive_prompt": "cat"})
+        assert result.is_error is False
+    assert attempts == 2
+    assert len(loads) == 1
+
+
+async def test_cli_stdio_discovery_forbids_inference_runtime(tmp_path: Path) -> None:
+    checkpoint = tmp_path / "checkpoints" / "placeholder.safetensors"
+    checkpoint.parent.mkdir()
+    checkpoint.touch()
+    config = tmp_path / "config.toml"
+    config.write_text('[model]\nroot = "."\ncheckpoint = "placeholder.safetensors"\n')
+    # Exercise the real CLI while making any inference import or model load fatal.
+    bootstrap = """
+import importlib.abc
+import sys
+class NoInference(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname.split('.')[0] in {'torch', 'diffusers', 'controlnet_aux', 'sd_embed'}:
+            raise AssertionError('inference imported during discovery: ' + fullname)
+sys.meta_path.insert(0, NoInference())
+from mcp_sd15.pipeline_factory import PipelineFactory
+def forbidden(self):
+    raise AssertionError('model loaded during discovery')
+PipelineFactory.create_base = forbidden
+from mcp_sd15.cli import main
+raise SystemExit(main())
+"""
+    params = StdioServerParameters(
+        command=sys.executable,
+        args=["-c", bootstrap, "--config", str(config)],
+    )
+    async with (
+        stdio_client(params) as (read, write),
+        ClientSession(read, write, read_timeout_seconds=15) as client,
+    ):
+        await client.initialize()
+        tools = (await client.list_tools()).tools
+        assert {tool.name for tool in tools} == {
+            "generate_image", "generate_guided_image", "inpaint_image"
+        }
-- 
cgit v1.2.3