Fix: switch to PaintByExamplePipeline (CatVTON had no model_index.json)

app.py

@@ -1,52 +1,30 @@
 """
-Virtual Try-On — CatVTON + Hugging Face ZeroGPU
-No local GPU or model storage needed. Generated images download to your device.
+Virtual Try-On — Paint-by-Example + Hugging Face ZeroGPU
+Uses an exemplar garment image to guide inpainting on a person photo.
+No local GPU or model storage needed.
 """
 
 import datetime
 import os
-import sys
 
 import gradio as gr
-import numpy as np
 import spaces
 import torch
-from huggingface_hub import snapshot_download
 from PIL import Image, ImageDraw
 
 # ---------------------------------------------------------------------------
 # Persistent storage (/data on ZeroGPU Spaces, /tmp fallback)
 # ---------------------------------------------------------------------------
 DATA_DIR   = "/data" if os.path.exists("/data") else "/tmp"
-MODELS_DIR = os.path.join(DATA_DIR, "catvton_models")
 OUTPUT_DIR = os.path.join(DATA_DIR, "outputs")
-os.makedirs(MODELS_DIR, exist_ok=True)
 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")
 
 # ---------------------------------------------------------------------------
-# Model download — runs once at Space startup on HF servers (not locally)
-# ---------------------------------------------------------------------------
-CATVTON_REPO  = "zhengchong/CatVTON"
-CATVTON_LOCAL = os.path.join(MODELS_DIR, "CatVTON")
-
-def download_models():
-    if os.path.exists(os.path.join(CATVTON_LOCAL, "model_index.json")):
-        print("CatVTON already cached.")
-        return
-    print("Downloading CatVTON (~4 GB) to HF persistent storage…")
-    snapshot_download(
-        repo_id=CATVTON_REPO,
-        local_dir=CATVTON_LOCAL,
-        local_dir_use_symlinks=False,
-        ignore_patterns=["*.md", "*.txt", "*.py"],
-    )
-    print("CatVTON ready.")
-
-# ---------------------------------------------------------------------------
-# Pipeline (loaded lazily inside @spaces.GPU)
+# Pipeline (loaded lazily inside @spaces.GPU to avoid wasting GPU quota)
 # ---------------------------------------------------------------------------
 _pipe = None
 
@@ -54,15 +32,14 @@ def _get_pipe():
     global _pipe
     if _pipe is not None:
         return _pipe
-    from diffusers import StableDiffusionInpaintPipeline
-    _pipe = StableDiffusionInpaintPipeline.from_pretrained(
-        CATVTON_LOCAL,
+    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,
-        safety_checker=None,
-        requires_safety_checker=False,
     ).to("cuda")
     _pipe.set_progress_bar_config(disable=True)
-    print("Pipeline loaded on CUDA.")
+    print("Pipeline ready on CUDA.")
     return _pipe
 
 # ---------------------------------------------------------------------------
@@ -112,17 +89,10 @@ def run_tryon(
     rng = torch.Generator(device="cuda")
     rng.manual_seed(int(seed) if seed != -1 else torch.randint(0, 2**32, (1,)).item())
 
-    prompt = (
-        "a person wearing the garment in the reference image, "
-        "photorealistic, high quality, natural lighting"
-    )
-    negative = "blurry, distorted, deformed, low quality, artifacts"
-
     result = pipe(
-        prompt=prompt,
-        negative_prompt=negative,
         image=person,
         mask_image=mask,
+        example_image=garment,
         num_inference_steps=num_steps,
         guidance_scale=guidance_scale,
         generator=rng,
@@ -146,7 +116,8 @@ with gr.Blocks(title="Virtual Try-On", theme=gr.themes.Soft()) as demo:
         "# 👗 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."
+        "> 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."
     )
 
     with gr.Row():
@@ -182,12 +153,9 @@ with gr.Blocks(title="Virtual Try-On", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         "---\n"
         "**Tips:** front-facing photo · garment on white/neutral background · upper body for shirts\n\n"
-        "First run: ~2-5 min (model download). Subsequent runs: ~15-30s.\n\n"
-        "Built with [CatVTON](https://github.com/zhengchong/CatVTON) · "
+        "Built with [Paint-by-Example](https://github.com/Fantasy-Studio/Paint-by-Example) · "
         "[Gradio](https://gradio.app) · [ZeroGPU](https://huggingface.co/docs/hub/spaces-zerogpu)"
     )
 
-download_models()
-
 if __name__ == "__main__":
     demo.launch()