Fix streaming: split into outer generator + inner @spaces.GPU function

app.py

@@ -23,25 +23,6 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 os.environ["HF_HOME"]               = os.path.join(DATA_DIR, "hf_cache")
 os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.join(DATA_DIR, "hf_cache", "hub")
 
-# ---------------------------------------------------------------------------
-# Pipeline (loaded lazily inside @spaces.GPU to avoid wasting GPU quota)
-# ---------------------------------------------------------------------------
-_pipe = None
-
-def _get_pipe():
-    global _pipe
-    if _pipe is not None:
-        return _pipe
-    from diffusers import PaintByExamplePipeline
-    print("Loading Paint-by-Example pipeline (~5 GB, first run only)…")
-    _pipe = PaintByExamplePipeline.from_pretrained(
-        "Fantasy-Studio/Paint-by-Example",
-        torch_dtype=torch.float16,
-    ).to("cuda")
-    _pipe.set_progress_bar_config(disable=True)
-    print("Pipeline ready on CUDA.")
-    return _pipe
-
 # ---------------------------------------------------------------------------
 # Image helpers
 # ---------------------------------------------------------------------------
@@ -66,9 +47,40 @@ def _make_mask(size: int, cloth_type: str) -> Image.Image:
     return mask
 
 # ---------------------------------------------------------------------------
-# ZeroGPU inference
+# GPU-decorated inference — plain return, no yield (spaces.GPU doesn't stream)
 # ---------------------------------------------------------------------------
+_pipe = None
+
 @spaces.GPU(duration=120)
+def _run_inference(person: Image.Image, garment: Image.Image, mask: Image.Image,
+                   num_steps: int, guidance_scale: float, seed: int):
+    global _pipe
+    if _pipe is None:
+        from diffusers import PaintByExamplePipeline
+        print("Loading Paint-by-Example pipeline (~5 GB, first run only)…")
+        _pipe = PaintByExamplePipeline.from_pretrained(
+            "Fantasy-Studio/Paint-by-Example",
+            torch_dtype=torch.float16,
+        ).to("cuda")
+        _pipe.set_progress_bar_config(disable=True)
+        print("Pipeline ready on CUDA.")
+
+    rng = torch.Generator(device="cuda")
+    rng.manual_seed(int(seed) if seed != -1 else torch.randint(0, 2**32, (1,)).item())
+
+    result = _pipe(
+        image=person,
+        mask_image=mask,
+        example_image=garment,
+        num_inference_steps=num_steps,
+        guidance_scale=guidance_scale,
+        generator=rng,
+    )
+    return result.images
+
+# ---------------------------------------------------------------------------
+# Outer generator — handles status streaming to Gradio UI
+# ---------------------------------------------------------------------------
 def run_tryon(
     person_image: Image.Image,
     garment_image: Image.Image,
@@ -80,28 +92,13 @@ def run_tryon(
     if person_image is None or garment_image is None:
         raise gr.Error("Please upload both a person photo and a garment image.")
 
-    yield [], [], "⏳ GPU allocated — loading model (first run: ~3 min download)…"
-
-    pipe = _get_pipe()
+    yield [], [], "⏳ Requesting GPU + loading model… (first run ~3 min, then ~30s)"
 
     person  = _fit_to_square(person_image)
     garment = _fit_to_square(garment_image)
     mask    = _make_mask(TARGET_SIZE, cloth_type)
 
-    rng = torch.Generator(device="cuda")
-    rng.manual_seed(int(seed) if seed != -1 else torch.randint(0, 2**32, (1,)).item())
-
-    yield [], [], f"🎨 Running {num_steps} diffusion steps…"
-
-    result = pipe(
-        image=person,
-        mask_image=mask,
-        example_image=garment,
-        num_inference_steps=num_steps,
-        guidance_scale=guidance_scale,
-        generator=rng,
-    )
-    output_images = result.images
+    output_images = _run_inference(person, garment, mask, num_steps, guidance_scale, seed)
 
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     saved_paths = []