feat(backend): ZImageStudioBackend with @spaces.GPU and mode dispatch

backend.py

@@ -10,6 +10,8 @@ try:
 except ImportError:
     spaces = None  # type: ignore[assignment]
 
+import modes
+
 
 _BASE_DURATION_S: dict[str, int] = {
     "t2i":        20,   # fixed setup + decode
@@ -42,3 +44,62 @@ def duration_for(
 
     est = (base + per_step * steps + cold_buffer) * size_factor * multiplier
     return max(60, min(int(est), 180))
+
+
+def _identity(fn):
+    return fn
+
+
+_ON_SPACES = bool(os.environ.get("SPACES_ZERO_GPU"))
+_GPU = spaces.GPU(duration=lambda *a, **kw: duration_for(*a[1:3], **kw)) \
+       if (spaces is not None and _ON_SPACES) else _identity
+
+
+def _build_pipeline() -> Any:
+    """Construct the DiffSynth ZImagePipeline. Imported lazily to keep tests fast."""
+    import torch
+    from diffsynth.pipelines.z_image import ZImagePipeline
+
+    import models
+
+    device = models.auto_device()
+    vram_cfg: dict[str, Any] = {}
+    if device != "cpu":
+        vram_cfg = dict(
+            offload_dtype=torch.bfloat16, offload_device="cpu",
+            onload_dtype=torch.bfloat16,  onload_device="cpu",
+            preparing_dtype=torch.bfloat16, preparing_device=device,
+            computation_dtype=torch.bfloat16, computation_device=device,
+        )
+
+    pipe = ZImagePipeline.from_pretrained(
+        torch_dtype=torch.bfloat16,
+        device=device,
+        model_configs=models.build_diffsynth_configs(vram_cfg=vram_cfg),
+        tokenizer_config=models.build_diffsynth_configs(
+            (models.TOKENIZER_CONFIG,), vram_cfg=None,
+        )[0],
+        vram_limit=models.vram_limit_for(device),
+    )
+    return pipe
+
+
+_DISPATCH = {
+    "t2i":        modes.call_t2i,
+    "controlnet": modes.call_controlnet,
+    "upscale":    modes.call_upscale,
+}
+
+
+class ZImageStudioBackend:
+    """One-process backend wrapping the DiffSynth ZImagePipeline."""
+
+    def __init__(self) -> None:
+        self.pipeline = _build_pipeline()
+
+    @_GPU
+    def generate(self, mode: str, params: dict[str, Any]) -> tuple[Any, dict[str, Any]]:
+        handler = _DISPATCH.get(mode)
+        if handler is None:
+            raise ValueError(f"unknown mode: {mode!r}; expected one of {list(_DISPATCH)}")
+        return handler(self.pipeline, params)

tests/test_backend.py

@@ -32,3 +32,49 @@ def test_duration_upscale_has_realesrgan_overhead():
     t2i = backend.duration_for(mode="t2i", params=dict(model="Turbo", steps=8, width=1024, height=1024))
     upsc = backend.duration_for(mode="upscale", params=dict(refine_steps=5, width=1024, height=1024))
     assert upsc > t2i
+
+
+from unittest.mock import MagicMock
+
+import pytest
+from PIL import Image
+
+
+@pytest.fixture
+def fake_backend(monkeypatch):
+    """A ZImageStudioBackend whose constructor doesn't actually build a pipeline."""
+    monkeypatch.setattr(backend, "_build_pipeline", lambda *a, **kw: MagicMock())
+    b = backend.ZImageStudioBackend()
+    b.pipeline.return_value = Image.new("RGB", (32, 32))
+    b.pipeline.dit = MagicMock()
+    b.pipeline.model_pool = MagicMock()
+    return b
+
+
+def test_backend_generate_routes_t2i(fake_backend):
+    img, meta = fake_backend.generate(
+        mode="t2i",
+        params=dict(prompt="cat", negative_prompt="", model="Turbo",
+                    steps=8, cfg=1.0, width=1024, height=1024, seed=42,
+                    lora_path=None, lora_strength=0.0),
+    )
+    assert isinstance(img, Image.Image)
+    assert meta["mode"] == "t2i"
+    assert meta["model"] == "Turbo"
+
+
+def test_backend_generate_routes_controlnet(fake_backend, monkeypatch):
+    monkeypatch.setattr(backend.modes, "preprocessors",
+                        type("P", (), {"run": staticmethod(lambda m, i: i)}))
+    img, meta = fake_backend.generate(
+        mode="controlnet",
+        params=dict(prompt="cat", input_image=Image.new("RGB", (64, 64)),
+                    preprocessor="Canny", controlnet_scale=1.0,
+                    steps=9, seed=0, lora_path=None, lora_strength=0.0),
+    )
+    assert meta["mode"] == "controlnet"
+
+
+def test_backend_generate_unknown_mode_raises(fake_backend):
+    with pytest.raises(ValueError):
+        fake_backend.generate(mode="dance", params={})