Run Iris-3B in-process on ZeroGPU via @spaces.GPU fn nodes

app.py

@@ -1,17 +1,77 @@
-"""Gradio Workflow app for the Iris-3B Space (speridlabs/iris-3b).
+"""Gradio Workflow app that runs Iris-3B in-process on ZeroGPU.
 
-Two pipelines on one canvas, editable in the browser:
+The canvas's Iris-3B nodes are `kind: "fn"` operators backed by the
+`illustrate` function below: the 3B pixel-space DiT and its text encoder are
+loaded once at module scope, and each call books its own GPU slot via
+`@spaces.GPU` — instead of calling the remote speridlabs/iris-3b Space through
+gradio_client.
 
-1. "Iris Text-to-Image": Prompt -> enhance_prompt (Python fn) -> Iris-3B /run
-   -> Generated Image + Seed Used.
-2. "Aspect Ratio Comparison": one Prompt fans out to two Iris-3B nodes at
-   1344x768 and 768x1344 (same seed) -> Wide + Tall subjects.
-
-The graph lives in workflow.json; the canvas autosaves edits back to it.
+The `iris3b` package is not on PyPI, so it is pulled from the
+speridlabs/iris-3b Space repo at startup; the weights come from the
+speridlabs/iris-3b model repo. Requires the Space hardware to be set to
+ZeroGPU (Settings -> Hardware).
 """
 
+import os
+import random
+import sys
+import tempfile
+
+# `spaces` has to be imported before anything initializes CUDA. It only exists
+# on Hugging Face Spaces, so shim it for local runs.
+try:
+    import spaces
+except ImportError:
+
+    class _SpacesShim:
+        @staticmethod
+        def GPU(duration=60):
+            return lambda fn: fn
+
+    spaces = _SpacesShim()
+
+import torch
+from gradio_client import utils as client_utils
+from huggingface_hub import snapshot_download
+from PIL import Image
+
 import gradio as gr
 
+MODEL_REPO = os.environ.get("IRIS_MODEL_REPO", "speridlabs/iris-3b")
+CODE_REPO = "speridlabs/iris-3b"  # Space repo holding the iris3b package
+
+_src = snapshot_download(CODE_REPO, repo_type="space", allow_patterns=["iris3b/**"])
+sys.path.insert(0, _src)
+
+import iris3b.text  # noqa: F401, E402  (registers encoders)
+from iris3b.config import inference_config  # noqa: E402
+from iris3b.models.dit import IrisDiT  # noqa: E402
+from iris3b.registry import TEXT_ENCODERS  # noqa: E402
+from iris3b.sampling import generate, load_for_inference  # noqa: E402
+
+# On ZeroGPU, `spaces` forces torch.cuda.is_available() to True (CUDA emulation
+# at import, real CUDA inside @spaces.GPU), so this resolves to cuda there and
+# to cpu when run locally.
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# (height, width) buckets with real training mass, from the original demo.
+SIZES = {f"{w}×{h}": (h, w) for h, w in [(768, 1344), (832, 1280), (896, 1152), (1024, 1024),
+                                          (1152, 896), (1280, 832), (1344, 768)]}
+
+# Load once at import and move to the device here, not inside the function: on
+# ZeroGPU the move is intercepted and the weights are transferred when a GPU is
+# actually allocated, so module scope costs no GPU time.
+source = MODEL_REPO if os.path.isdir(MODEL_REPO) else snapshot_download(
+    MODEL_REPO, token=os.environ.get("HF_TOKEN")
+)
+raw, weights = load_for_inference(source)
+cfg = inference_config(raw, [])
+with torch.device("meta"):
+    model = IrisDiT(cfg.model)
+model.load_state_dict(weights, strict=True, assign=True)
+model = model.eval().to(device=DEVICE, dtype=torch.float32)  # FP32 weights, BF16 autocast at inference
+text_encoder = TEXT_ENCODERS.build(cfg.text_encoder.name, cfg.text_encoder, device=DEVICE)
+
 QUALITY_SUFFIX = "highly detailed, sharp focus, professional quality"
 
 
@@ -25,7 +85,46 @@ def enhance_prompt(prompt: str) -> str:
     return f"{prompt}, {QUALITY_SUFFIX}"
 
 
-demo = gr.Workflow(graph="workflow.json", bind={"enhance_prompt": enhance_prompt})
+# duration is per node, not per pipeline: each illustrate call books and
+# releases its own GPU slot. ZeroGPU multiplies this by the hardware factor
+# (1.5 on RTX PRO 6000 Blackwell), so 120s books 180s against the caller.
+@spaces.GPU(duration=120)
+def illustrate(prompt: str, negative: str, size: str, steps: float,
+               cfg_scale: float, seed: float, randomize: bool):
+    """Generate one image with Iris-3B. Returns (image file dict, seed)."""
+    if not (prompt or "").strip():
+        raise gr.Error("Connect a prompt to this node.")
+    if randomize:
+        seed = random.randint(0, 2**31 - 1)
+    if size not in SIZES:
+        size = "1024×1024"
+    height, width = SIZES[size]
+    with torch.autocast(DEVICE.type, dtype=torch.bfloat16, enabled=DEVICE.type == "cuda"):
+        image = generate(
+            model, text_encoder, [prompt], height=height, width=width, steps=int(steps),
+            order=cfg.sample.order, cfg_scale=cfg_scale,
+            cfg_interval=tuple(cfg.sample.cfg_interval), shift=cfg.flow.shift,
+            negative_prompt=negative or "",
+            generator=torch.Generator(device=DEVICE).manual_seed(int(seed)), device=DEVICE,
+            num_train_timesteps=cfg.flow.num_train_timesteps, prediction=cfg.flow.prediction,
+        )[0]
+    pixels = ((image.float().clamp(-1, 1) + 1) * 127.5).round().byte().permute(1, 2, 0).cpu().numpy()
+    path = os.path.join(tempfile.gettempdir(), f"iris_{os.urandom(8).hex()}.png")
+    Image.fromarray(pixels).save(path)
+    # The canvas renders media only from a dict carrying a url; launch()
+    # already allowlists the system tempdir.
+    media = {
+        "path": path,
+        "url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
+        "is_file": True,
+    }
+    return media, int(seed)
+
+
+demo = gr.Workflow(
+    graph="workflow.json",
+    bind={"enhance_prompt": enhance_prompt, "illustrate": illustrate},
+)
 
 if __name__ == "__main__":
     demo.launch()