feat(modes): t2i handler (base + turbo) with transformer swap and lora ctx

modes.py

@@ -0,0 +1,63 @@
+"""Mode handlers — pure functions over a ZImagePipeline + params dict."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, TypedDict
+
+from PIL import Image
+
+import lora
+
+
+class T2IParams(TypedDict, total=False):
+    prompt: str
+    negative_prompt: str
+    model: str          # "Base" | "Turbo"
+    steps: int
+    cfg: float
+    width: int
+    height: int
+    seed: int
+    lora_path: Path | None
+    lora_strength: float
+
+
+def _swap_transformer(pipe: Any, model_name: str) -> None:
+    """Swap the active transformer in the pipeline's model pool."""
+    variant = "z_image" if model_name == "Base" else "z_image_turbo"
+    pipe.dit = pipe.model_pool.fetch_model("z_image_dit", variant=variant)
+    try:
+        pipe.dit._zis_variant = variant
+    except (AttributeError, RuntimeError):
+        pass
+
+
+def call_t2i(pipe: Any, params: T2IParams) -> tuple[Image.Image, dict[str, Any]]:
+    """Text-to-image. Routes to base (cfg=4, 25 steps) or turbo (cfg=1, 8 steps)."""
+    model_name = params.get("model", "Turbo")
+    is_base = model_name == "Base"
+    _swap_transformer(pipe, model_name)
+
+    kwargs: dict[str, Any] = dict(
+        prompt=params["prompt"],
+        cfg_scale=float(params.get("cfg", 4.0 if is_base else 1.0)),
+        num_inference_steps=int(params.get("steps", 25 if is_base else 8)),
+        sigma_shift=3.0,
+        height=int(params.get("height", 1024)),
+        width=int(params.get("width", 1024)),
+        seed=int(params.get("seed", 0)),
+    )
+    if is_base and params.get("negative_prompt"):
+        kwargs["negative_prompt"] = params["negative_prompt"]
+
+    with lora.applied_lora(pipe, params.get("lora_path"), params.get("lora_strength", 0.0)):
+        image = pipe(**kwargs)
+
+    meta = dict(
+        mode="t2i", model=model_name,
+        steps=kwargs["num_inference_steps"], cfg=kwargs["cfg_scale"],
+        seed=kwargs["seed"], width=kwargs["width"], height=kwargs["height"],
+        lora=str(params.get("lora_path")) if params.get("lora_path") else None,
+        lora_strength=params.get("lora_strength", 0.0),
+    )
+    return image, meta

tests/test_modes.py

@@ -0,0 +1,70 @@
+from unittest.mock import MagicMock
+
+import pytest
+from PIL import Image
+
+import modes
+
+
+@pytest.fixture
+def fake_pipe():
+    """Stand-in pipeline that records its __call__ args and returns a dummy image."""
+    pipe = MagicMock()
+    pipe.dit = MagicMock()
+    pipe.model_pool = MagicMock()
+    pipe.return_value = Image.new("RGB", (64, 64), color=(255, 176, 46))
+    return pipe
+
+
+def test_t2i_turbo_builds_minimal_call(fake_pipe):
+    out, meta = modes.call_t2i(
+        fake_pipe,
+        params=dict(
+            prompt="a cat",
+            negative_prompt="",
+            model="Turbo",
+            steps=8, cfg=1.0,
+            width=1024, height=1024,
+            seed=42,
+            lora_path=None, lora_strength=0.0,
+        ),
+    )
+    fake_pipe.assert_called_once()
+    kwargs = fake_pipe.call_args.kwargs
+    assert kwargs["prompt"] == "a cat"
+    assert kwargs["cfg_scale"] == 1.0
+    assert kwargs["num_inference_steps"] == 8
+    assert kwargs["width"] == 1024
+    assert kwargs["seed"] == 42
+    assert kwargs["sigma_shift"] == 3.0
+    assert "negative_prompt" not in kwargs or not kwargs.get("negative_prompt")
+    assert meta["model"] == "Turbo"
+    assert meta["steps"] == 8
+    assert isinstance(out, Image.Image)
+
+
+def test_t2i_base_passes_negative_prompt_and_cfg4(fake_pipe):
+    modes.call_t2i(
+        fake_pipe,
+        params=dict(
+            prompt="a cat", negative_prompt="blurry, lowres",
+            model="Base", steps=25, cfg=4.0,
+            width=1024, height=1024, seed=42,
+            lora_path=None, lora_strength=0.0,
+        ),
+    )
+    kwargs = fake_pipe.call_args.kwargs
+    assert kwargs["negative_prompt"] == "blurry, lowres"
+    assert kwargs["cfg_scale"] == 4.0
+    assert kwargs["num_inference_steps"] == 25
+
+
+def test_t2i_swaps_transformer_via_model_pool(fake_pipe):
+    modes.call_t2i(
+        fake_pipe,
+        params=dict(prompt="x", negative_prompt="", model="Base", steps=25, cfg=4.0,
+                    width=1024, height=1024, seed=0, lora_path=None, lora_strength=0.0),
+    )
+    fake_pipe.model_pool.fetch_model.assert_called()
+    call = fake_pipe.model_pool.fetch_model.call_args
+    assert call.args[0] == "z_image_dit"