feat(modes): upscale handler (realesrgan + z-image-turbo refinement)

modes.py

@@ -8,6 +8,7 @@ from PIL import Image
 
 import lora
 import preprocessors
+import upscale
 
 try:
     from diffsynth.diffusion.base_pipeline import ControlNetInput
@@ -111,3 +112,37 @@ def call_controlnet(pipe: Any, params: dict[str, Any]) -> tuple[Image.Image, dic
         lora_strength=params.get("lora_strength", 0.0),
     )
     return image, meta
+
+
+def call_upscale(pipe: Any, params: dict[str, Any]) -> tuple[Image.Image, dict[str, Any]]:
+    """Upscale — RealESRGAN x4 → 0.5 resize → Z-Image-Turbo img2img refinement."""
+    input_image: Image.Image | None = params.get("input_image")
+    if input_image is None:
+        raise ValueError("Upscale mode requires an input image")
+
+    upscaled = upscale.realesrgan_2x(input_image, model_path=params["esrgan_model_path"])
+
+    _swap_transformer(pipe, "Turbo")
+
+    kwargs: dict[str, Any] = dict(
+        prompt=params.get("prompt", "masterpiece, 8k"),
+        cfg_scale=1.0,
+        num_inference_steps=int(params.get("refine_steps", 5)),
+        sigma_shift=3.0,
+        input_image=upscaled,
+        denoising_strength=float(params.get("refine_denoise", 0.33)),
+        seed=int(params.get("seed", 0)),
+    )
+
+    with lora.applied_lora(pipe, params.get("lora_path"), params.get("lora_strength", 0.0)):
+        image = pipe(**kwargs)
+
+    meta = dict(
+        mode="upscale", model="Turbo",
+        refine_steps=kwargs["num_inference_steps"],
+        refine_denoise=kwargs["denoising_strength"],
+        seed=kwargs["seed"], width=upscaled.size[0], height=upscaled.size[1],
+        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

@@ -110,3 +110,42 @@ def test_controlnet_rejects_missing_input_image(fake_pipe):
                         controlnet_scale=1.0, steps=9, seed=0,
                         lora_path=None, lora_strength=0.0),
         )
+
+
+def test_upscale_runs_realesrgan_then_pipeline(fake_pipe, monkeypatch):
+    calls = {"upscale": None}
+    def fake_2x(img, model_path):
+        calls["upscale"] = (img.size, str(model_path))
+        w, h = img.size
+        return img.resize((w * 2, h * 2), Image.LANCZOS)
+    monkeypatch.setattr(modes, "upscale", type("U", (), {"realesrgan_2x": staticmethod(fake_2x)}))
+
+    input_image = Image.new("RGB", (512, 512))
+    out, meta = modes.call_upscale(
+        fake_pipe,
+        params=dict(
+            prompt="masterpiece, 8k",
+            input_image=input_image,
+            refine_steps=5,
+            refine_denoise=0.33,
+            seed=42,
+            lora_path=None, lora_strength=0.0,
+            esrgan_model_path="/fake/path/RealESRGAN_x4plus.pth",
+        ),
+    )
+
+    assert calls["upscale"] == ((512, 512), "/fake/path/RealESRGAN_x4plus.pth")
+    kwargs = fake_pipe.call_args.kwargs
+    assert kwargs["input_image"].size == (1024, 1024)  # 2x via fake_2x
+    assert kwargs["denoising_strength"] == 0.33
+    assert kwargs["num_inference_steps"] == 5
+    assert kwargs["cfg_scale"] == 1.0
+    assert meta["mode"] == "upscale"
+
+
+def test_upscale_rejects_missing_image(fake_pipe):
+    with pytest.raises(ValueError):
+        modes.call_upscale(fake_pipe, params=dict(prompt="x", input_image=None,
+                                                   refine_steps=5, refine_denoise=0.33, seed=0,
+                                                   lora_path=None, lora_strength=0.0,
+                                                   esrgan_model_path="/fake.pth"))