feat(modes): controlnet handler (turbo + union 2.1 + preprocessor)

modes.py

@@ -7,6 +7,17 @@ from typing import Any, TypedDict
 from PIL import Image
 
 import lora
+import preprocessors
+
+try:
+    from diffsynth.diffusion.base_pipeline import ControlNetInput
+except ImportError:
+    from dataclasses import dataclass
+
+    @dataclass
+    class ControlNetInput:  # type: ignore[no-redef]
+        image: Any
+        scale: float = 1.0
 
 
 class T2IParams(TypedDict, total=False):
@@ -61,3 +72,42 @@ def call_t2i(pipe: Any, params: T2IParams) -> tuple[Image.Image, dict[str, Any]]
         lora_strength=params.get("lora_strength", 0.0),
     )
     return image, meta
+
+
+def call_controlnet(pipe: Any, params: dict[str, Any]) -> tuple[Image.Image, dict[str, Any]]:
+    """ControlNet — Turbo + Z-Image-Turbo-Fun-Controlnet-Union-2.1."""
+    input_image: Image.Image | None = params.get("input_image")
+    if input_image is None:
+        raise ValueError("ControlNet mode requires an input image")
+
+    preproc_mode = params.get("preprocessor", "Canny")
+    control_image = preprocessors.run(preproc_mode, input_image)
+
+    _swap_transformer(pipe, "Turbo")
+
+    cn_input = ControlNetInput(image=control_image, scale=float(params.get("controlnet_scale", 1.0)))
+
+    kwargs: dict[str, Any] = dict(
+        prompt=params["prompt"],
+        cfg_scale=1.0,
+        num_inference_steps=int(params.get("steps", 9)),
+        sigma_shift=3.0,
+        height=control_image.size[1],
+        width=control_image.size[0],
+        seed=int(params.get("seed", 0)),
+        controlnet_inputs=[cn_input],
+    )
+
+    with lora.applied_lora(pipe, params.get("lora_path"), params.get("lora_strength", 0.0)):
+        image = pipe(**kwargs)
+
+    meta = dict(
+        mode="controlnet", model="Turbo",
+        preprocessor=preproc_mode,
+        controlnet_scale=cn_input.scale,
+        steps=kwargs["num_inference_steps"], cfg=1.0,
+        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

@@ -68,3 +68,45 @@ def test_t2i_swaps_transformer_via_model_pool(fake_pipe):
     fake_pipe.model_pool.fetch_model.assert_called()
     call = fake_pipe.model_pool.fetch_model.call_args
     assert call.args[0] == "z_image_dit"
+
+
+def test_controlnet_calls_preprocessor_then_pipeline(fake_pipe, monkeypatch):
+    canny_called = []
+    def fake_run(mode, img):
+        canny_called.append((mode, img.size))
+        return img  # passthrough for test
+    monkeypatch.setattr(modes, "preprocessors", type("P", (), {"run": staticmethod(fake_run)}))
+
+    input_image = Image.new("RGB", (1024, 1024))
+    out, meta = modes.call_controlnet(
+        fake_pipe,
+        params=dict(
+            prompt="cinematic portrait",
+            input_image=input_image,
+            preprocessor="Canny",
+            controlnet_scale=1.0,
+            steps=9,
+            seed=42,
+            lora_path=None, lora_strength=0.0,
+        ),
+    )
+
+    assert canny_called == [("Canny", (1024, 1024))]
+    kwargs = fake_pipe.call_args.kwargs
+    assert "controlnet_inputs" in kwargs
+    cn_in = kwargs["controlnet_inputs"]
+    assert len(cn_in) == 1
+    assert cn_in[0].scale == 1.0
+    assert kwargs["num_inference_steps"] == 9
+    assert kwargs["cfg_scale"] == 1.0
+    assert meta["preprocessor"] == "Canny"
+
+
+def test_controlnet_rejects_missing_input_image(fake_pipe):
+    with pytest.raises(ValueError):
+        modes.call_controlnet(
+            fake_pipe,
+            params=dict(prompt="x", input_image=None, preprocessor="Canny",
+                        controlnet_scale=1.0, steps=9, seed=0,
+                        lora_path=None, lora_strength=0.0),
+        )