Update app.py

app.py

@@ -1,10 +1,8 @@
-"""FASHN VTON v1.5 HuggingFace Space Demo."""
+"""FASHN VTON v1.5 HuggingFace Space Demo (CPU Compatible Version)."""
 
 import os
 import platform
-
 import gradio as gr
-import spaces
 import torch
 from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -18,20 +16,17 @@ WEIGHTS_DIR = os.path.join(SCRIPT_DIR, "weights")
 CATEGORIES = ["tops", "bottoms", "one-pieces"]
 GARMENT_PHOTO_TYPES = ["model", "flat-lay"]
 
-# Global pipeline instance (lazy loaded)
 _pipeline = None
 
 
 # ----------------- HELPERS ----------------- #
 
-
 def download_weights():
     """Download model weights from HuggingFace Hub."""
     os.makedirs(WEIGHTS_DIR, exist_ok=True)
     dwpose_dir = os.path.join(WEIGHTS_DIR, "dwpose")
     os.makedirs(dwpose_dir, exist_ok=True)
 
-    # Download TryOnModel weights
     tryon_path = os.path.join(WEIGHTS_DIR, "model.safetensors")
     if not os.path.exists(tryon_path):
         print("Downloading TryOnModel weights...")
@@ -41,7 +36,6 @@ def download_weights():
             local_dir=WEIGHTS_DIR,
         )
 
-    # Download DWPose models
     dwpose_files = ["yolox_l.onnx", "dw-ll_ucoco_384.onnx"]
     for filename in dwpose_files:
         filepath = os.path.join(dwpose_dir, filename)
@@ -53,56 +47,45 @@ def download_weights():
                 local_dir=dwpose_dir,
             )
 
-    print("Weights downloaded successfully!")
+    print("Weights ready!")
 
 
 # ----------------- MODEL LOADING ----------------- #
 
-
 def get_pipeline():
-    """Lazy-load the pipeline on first use (ensures GPU is available on ZeroGPU)."""
+    """Lazy-load pipeline (CPU/GPU auto detect)."""
     global _pipeline
+
     if _pipeline is None:
-        # Check CUDA availability (will be true inside @spaces.GPU context)
-        if not torch.cuda.is_available():
-            raise gr.Error(
-                "CUDA is not available. This demo requires a GPU to run. "
-                "If you're on HuggingFace Spaces, please try again in a moment."
-            )
 
-        # ---------------------------------- Diagnostics ---------------------------------- #
-        print(f"Python           : {platform.python_version()}")
-        print(f"PyTorch          : {torch.__version__}")
-        print(f" • built for CUDA : {torch.version.cuda}")
-        if torch.backends.cudnn.is_available():
-            print(f" • built for cuDNN: {torch.backends.cudnn.version()}")
-        print(f"torch.cuda.is_available(): {torch.cuda.is_available()}")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {device}")
+
+        print(f"Python: {platform.python_version()}")
+        print(f"PyTorch: {torch.__version__}")
+
         if torch.cuda.is_available():
             dev = torch.cuda.current_device()
-            cc = torch.cuda.get_device_capability(dev)
-            print(f"GPU {dev}: {torch.cuda.get_device_name(dev)} (compute {cc[0]}.{cc[1]})")
-
-        # Enable TF32 for faster computation on Ampere+ GPUs
-        if torch.cuda.get_device_properties(0).major >= 8:
-            torch.backends.cuda.matmul.allow_tf32 = True
-            torch.backends.cudnn.allow_tf32 = True
+            print(f"GPU: {torch.cuda.get_device_name(dev)}")
 
-        print("Downloading weights (if needed)...")
+        print("Downloading weights if needed...")
         download_weights()
 
         print("Loading pipeline...")
         from fashn_vton import TryOnPipeline
 
-        _pipeline = TryOnPipeline(weights_dir=WEIGHTS_DIR, device="cuda")
-        print("Pipeline loaded on CUDA!")
+        _pipeline = TryOnPipeline(
+            weights_dir=WEIGHTS_DIR,
+            device=device,
+        )
+
+        print(f"Pipeline loaded successfully on {device}")
 
     return _pipeline
 
 
 # ----------------- INFERENCE ----------------- #
 
-
-@spaces.GPU
 def try_on(
     person_image: Image.Image,
     garment_image: Image.Image,
@@ -113,26 +96,23 @@ def try_on(
     seed: int,
     segmentation_free: bool,
 ) -> Image.Image:
-    """Run virtual try-on inference."""
+
     if person_image is None:
         raise gr.Error("Please upload a person image")
     if garment_image is None:
         raise gr.Error("Please upload a garment image")
 
-    # Handle seed (guard against None or invalid values)
     if seed is None or seed < 0:
         seed = 42
 
-    # Convert to RGB if needed
     if person_image.mode != "RGB":
         person_image = person_image.convert("RGB")
+
     if garment_image.mode != "RGB":
         garment_image = garment_image.convert("RGB")
 
-    # Get pipeline (lazy loads on first call)
     pipeline = get_pipeline()
 
-    # Run inference
     result = pipeline(
         person_image=person_image,
         garment_image=garment_image,
@@ -150,7 +130,6 @@ def try_on(
 
 # ----------------- UI ----------------- #
 
-# Custom CSS
 CUSTOM_CSS = """
 .contain img {
     object-fit: contain !important;
@@ -159,151 +138,40 @@ CUSTOM_CSS = """
 }
 """
 
-# Load HTML content
-with open(os.path.join(SCRIPT_DIR, "banner.html"), "r") as f:
-    banner_html = f.read()
-with open(os.path.join(SCRIPT_DIR, "tips.html"), "r") as f:
-    tips_html = f.read()
-
-# Build example paths
-examples_dir = os.path.join(ASSETS_DIR, "examples")
-
-# Paired examples: [person_path, garment_path, category, garment_photo_type]
-paired_examples = [
-    [os.path.join(examples_dir, "person1.png"), os.path.join(examples_dir, "garment1.jpeg"), "one-pieces", "model"],
-    [os.path.join(examples_dir, "person2.png"), os.path.join(examples_dir, "garment2.webp"), "tops", "model"],
-    [os.path.join(examples_dir, "person3.png"), os.path.join(examples_dir, "garment3.jpeg"), "tops", "flat-lay"],
-    [os.path.join(examples_dir, "person4.png"), os.path.join(examples_dir, "garment4.webp"), "tops", "model"],
-    [os.path.join(examples_dir, "person5.png"), os.path.join(examples_dir, "garment5.jpeg"), "bottoms", "flat-lay"],
-    [os.path.join(examples_dir, "person6.png"), os.path.join(examples_dir, "garment6.webp"), "one-pieces", "model"],
-]
-
-# Individual examples (classic from repo)
-person_only_examples = [os.path.join(examples_dir, "person0.png")]
-
-# Garment examples with their settings: (image_path, category, photo_type)
-# Order matters - index in Gallery corresponds to this list
-garment_examples_data = [
-    (os.path.join(examples_dir, "garment0.png"), "tops", "model"),
-    (os.path.join(examples_dir, "garment7.jpg"), "tops", "flat-lay"),
-]
-garment_gallery_images = [item[0] for item in garment_examples_data]
-
-
-def on_garment_gallery_select(evt: gr.SelectData):
-    """Handle garment gallery selection - load image and update dropdowns."""
-    idx = evt.index
-    if idx < len(garment_examples_data):
-        image_path, cat, photo_type = garment_examples_data[idx]
-        return Image.open(image_path), cat, photo_type
-    return None, "tops", "model"
-
-
-# Build UI
 with gr.Blocks(css=CUSTOM_CSS) as demo:
-    # Header
-    gr.HTML(banner_html)
-    gr.HTML(tips_html)
-
-    with gr.Row(equal_height=False):
-        # Column 1: Person
-        with gr.Column(scale=1):
-            person_image = gr.Image(
-                label="Person Image",
-                type="pil",
-                sources=["upload", "clipboard"],
-                elem_classes=["contain"],
-            )
 
-            # Individual person examples
-            gr.Examples(
-                examples=person_only_examples,
-                inputs=person_image,
-                label="Person Examples",
-            )
+    gr.Markdown("# 👕 FASHN VTON v1.5 (CPU Compatible)")
 
-        # Column 2: Garment
-        with gr.Column(scale=1):
-            garment_image = gr.Image(
-                label="Garment Image",
-                type="pil",
-                sources=["upload", "clipboard"],
-                elem_classes=["contain"],
-            )
+    with gr.Row():
+        with gr.Column():
+            person_image = gr.Image(label="Person Image", type="pil")
 
-            with gr.Row():
-                category = gr.Dropdown(
-                    choices=CATEGORIES,
-                    value="tops",
-                    label="Category",
-                )
-                garment_photo_type = gr.Dropdown(
-                    choices=GARMENT_PHOTO_TYPES,
-                    value="model",
-                    label="Photo Type",
-                )
-
-            # Garment examples as clickable gallery
-            gr.Markdown("**Garment Examples** (click to load with settings)")
-            garment_gallery = gr.Gallery(
-                value=garment_gallery_images,
-                columns=2,
-                rows=1,
-                height="auto",
-                object_fit="contain",
-                show_label=False,
-                allow_preview=False,
+        with gr.Column():
+            garment_image = gr.Image(label="Garment Image", type="pil")
+
+            category = gr.Dropdown(
+                choices=CATEGORIES,
+                value="tops",
+                label="Category",
             )
 
-        # Column 3: Result
-        with gr.Column(scale=1):
-            result_image = gr.Image(
-                label="Try-On Result",
-                type="pil",
-                interactive=False,
-                elem_classes=["contain"],
+            garment_photo_type = gr.Dropdown(
+                choices=GARMENT_PHOTO_TYPES,
+                value="model",
+                label="Photo Type",
             )
 
-            run_button = gr.Button("Try On", variant="primary", size="lg")
+        with gr.Column():
+            result_image = gr.Image(label="Try-On Result", type="pil")
+
+            run_button = gr.Button("Try On", variant="primary")
 
-            # Advanced settings
             with gr.Accordion("Advanced Settings", open=False):
-                num_timesteps = gr.Slider(
-                    minimum=10,
-                    maximum=50,
-                    value=50,
-                    step=5,
-                    label="Sampling Steps",
-                    info="Higher = better quality, slower.",
-                )
-                guidance_scale = gr.Slider(
-                    minimum=1.0,
-                    maximum=3.0,
-                    value=1.5,
-                    step=0.1,
-                    label="Guidance Scale",
-                    info="How closely to follow the garment. 1.5 recommended.",
-                )
-                seed = gr.Number(
-                    value=42,
-                    label="Seed",
-                    info="Random seed for reproducibility.",
-                    precision=0,
-                )
-                segmentation_free = gr.Checkbox(
-                    value=True,
-                    label="Segmentation Free",
-                    info="Preserves body features and allows unconstrained garment volume. Disable for tighter garment fitting.",
-                )
-
-    # Paired examples at the bottom
-    gr.Examples(
-        examples=paired_examples,
-        inputs=[person_image, garment_image, category, garment_photo_type],
-        label="Complete Examples (click to load person + garment + settings)",
-    )
+                num_timesteps = gr.Slider(10, 50, value=50, step=5, label="Sampling Steps")
+                guidance_scale = gr.Slider(1.0, 3.0, value=1.5, step=0.1, label="Guidance Scale")
+                seed = gr.Number(value=42, label="Seed", precision=0)
+                segmentation_free = gr.Checkbox(value=True, label="Segmentation Free")
 
-    # Event handlers
     run_button.click(
         fn=try_on,
         inputs=[
@@ -316,18 +184,10 @@ with gr.Blocks(css=CUSTOM_CSS) as demo:
             seed,
             segmentation_free,
         ],
-        outputs=[result_image],
-    )
-
-    # Garment gallery selection - loads image and updates dropdowns
-    garment_gallery.select(
-        fn=on_garment_gallery_select,
-        inputs=None,
-        outputs=[garment_image, category, garment_photo_type],
+        outputs=result_image,
     )
 
-# Configure queue with concurrency limit to prevent GPU OOM
-demo.queue(default_concurrency_limit=1, max_size=30)
+demo.queue(default_concurrency_limit=1, max_size=20)
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()