Fix status updates: use gr.Progress instead of yield generator

app.py

@@ -19,7 +19,6 @@ DATA_DIR   = "/data" if os.path.exists("/data") else "/tmp"
 OUTPUT_DIR = os.path.join(DATA_DIR, "outputs")
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# Point HF cache to persistent storage so model downloads survive restarts
 os.environ["HF_HOME"]               = os.path.join(DATA_DIR, "hf_cache")
 os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.join(DATA_DIR, "hf_cache", "hub")
 
@@ -42,12 +41,12 @@ def _make_mask(size: int, cloth_type: str) -> Image.Image:
         d.rectangle([int(size*.10), int(size*.18), int(size*.90), int(size*.65)], fill=255)
     elif cloth_type == "lower":
         d.rectangle([int(size*.05), int(size*.55), int(size*.95), int(size*1.0)], fill=255)
-    else:  # overall / dress
+    else:
         d.rectangle([int(size*.05), int(size*.15), int(size*.95), int(size*1.0)], fill=255)
     return mask
 
 # ---------------------------------------------------------------------------
-# GPU-decorated inference — plain return, no yield (spaces.GPU doesn't stream)
+# GPU-decorated inference
 # ---------------------------------------------------------------------------
 _pipe = None
 
@@ -57,13 +56,13 @@ def _run_inference(person: Image.Image, garment: Image.Image, mask: Image.Image,
     global _pipe
     if _pipe is None:
         from diffusers import PaintByExamplePipeline
-        print("Loading Paint-by-Example pipeline (~5 GB, first run only)…")
+        print("Loading Paint-by-Example (~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.")
+        print("Pipeline ready.")
 
     rng = torch.Generator(device="cuda")
     rng.manual_seed(int(seed) if seed != -1 else torch.randint(0, 2**32, (1,)).item())
@@ -79,7 +78,7 @@ def _run_inference(person: Image.Image, garment: Image.Image, mask: Image.Image,
     return result.images
 
 # ---------------------------------------------------------------------------
-# Outer generator — handles status streaming to Gradio UI
+# Gradio inference — uses gr.Progress for live status updates
 # ---------------------------------------------------------------------------
 def run_tryon(
     person_image: Image.Image,
@@ -88,11 +87,12 @@ def run_tryon(
     num_steps: int,
     guidance_scale: float,
     seed: int,
+    progress=gr.Progress(track_tqdm=True),
 ):
     if person_image is None or garment_image is None:
         raise gr.Error("Please upload both a person photo and a garment image.")
 
-    yield [], [], "⏳ Requesting GPU + loading model… (first run ~3 min, then ~30s)"
+    progress(0, desc="⏳ Requesting GPU + loading model (first run ~3 min)…")
 
     person  = _fit_to_square(person_image)
     garment = _fit_to_square(garment_image)
@@ -100,6 +100,8 @@ def run_tryon(
 
     output_images = _run_inference(person, garment, mask, num_steps, guidance_scale, seed)
 
+    progress(0.9, desc="💾 Saving result…")
+
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     saved_paths = []
     for i, img in enumerate(output_images):
@@ -107,7 +109,8 @@ def run_tryon(
         img.save(path, format="PNG")
         saved_paths.append(path)
 
-    yield output_images, saved_paths, "✅ Done! Download your result below."
+    progress(1.0, desc="✅ Done!")
+    return output_images, saved_paths, "✅ Done! Download your result below."
 
 # ---------------------------------------------------------------------------
 # Gradio UI
@@ -116,9 +119,8 @@ with gr.Blocks(title="Virtual Try-On", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         "# 👗 Virtual Try-On\n"
         "Upload a **person photo** and a **garment image**, select the type, then click **Try On**.\n\n"
-        "> Runs entirely on **Hugging Face ZeroGPU** (free A10G) — no local GPU needed.  \n"
-        "> Models download once to HF persistent storage. Images save to your device via the Download button.\n\n"
-        "> **First run:** ~2-3 min (model download). **Subsequent runs:** ~15-30s."
+        "> Runs on **Hugging Face ZeroGPU** (free A10G) — no local GPU needed.  \n"
+        "> **First run:** ~2-3 min (model download ~5 GB). **Subsequent runs:** ~15-30s."
     )
 
     with gr.Row():