Update app.py

app.py

@@ -16,7 +16,10 @@ from utils import init_weight_dtype, resize_and_crop, resize_and_padding
 
 class CatVTONService:
     def __init__(self):
-        self.device = "cpu"  # Will use CPU on free tier
+        # Auto-detect device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"🖥️  Using device: {self.device}")
+        
         self.pipeline = None
         self.automasker = None
         self.models_loaded = False
@@ -28,33 +31,46 @@ class CatVTONService:
             
         print("🔄 Loading CatVTON models (this happens once)...")
         
-        # Download model weights from HuggingFace Hub - CACHED automatically
-        repo_path = snapshot_download(
-            repo_id="zhengchong/CatVTON",
-            cache_dir="./model_cache"
-        )
-        
-        print("✅ Models downloaded and cached!")
-        
-        # Initialize pipeline
-        self.pipeline = CatVTONPipeline(
-            base_ckpt="booksforcharlie/stable-diffusion-inpainting",
-            attn_ckpt=repo_path,
-            attn_ckpt_version="mix",
-            weight_dtype=init_weight_dtype("fp16" if self.device == "cuda" else "fp32"),
-            use_tf32=False,  # CPU doesn't support TF32
-            device=self.device
-        )
-        
-        # Initialize automasker
-        self.automasker = AutoMasker(
-            densepose_ckpt=os.path.join(repo_path, "DensePose"),
-            schp_ckpt=os.path.join(repo_path, "SCHP"),
-            device=self.device
-        )
-        
-        self.models_loaded = True
-        print("✅ CatVTON ready!")
+        try:
+            # Download model weights from HuggingFace Hub - CACHED automatically
+            repo_path = snapshot_download(
+                repo_id="zhengchong/CatVTON",
+                cache_dir="./model_cache",
+                resume_download=True,  # Resume if interrupted
+                local_files_only=False  # Allow downloading
+            )
+            
+            print(f"✅ Models downloaded to: {repo_path}")
+            
+            # Determine weight dtype based on device
+            weight_dtype = init_weight_dtype("fp16" if self.device == "cuda" else "fp32")
+            use_tf32 = self.device == "cuda"  # Only use TF32 on CUDA
+            
+            print(f"⚙️  Weight dtype: {weight_dtype}, TF32: {use_tf32}")
+            
+            # Initialize pipeline
+            self.pipeline = CatVTONPipeline(
+                base_ckpt="booksforcharlie/stable-diffusion-inpainting",
+                attn_ckpt=repo_path,
+                attn_ckpt_version="mix",
+                weight_dtype=weight_dtype,
+                use_tf32=use_tf32,
+                device=self.device
+            )
+            
+            # Initialize automasker
+            self.automasker = AutoMasker(
+                densepose_ckpt=os.path.join(repo_path, "DensePose"),
+                schp_ckpt=os.path.join(repo_path, "SCHP"),
+                device=self.device
+            )
+            
+            self.models_loaded = True
+            print("✅ CatVTON ready!")
+            
+        except Exception as e:
+            print(f"❌ Error loading models: {e}")
+            raise
         
     def generate_tryon(self, person_image, garment_image, progress=gr.Progress()):
         """Generate virtual try-on result"""
@@ -96,7 +112,8 @@ class CatVTONService:
             if self.device == "cuda":
                 torch.cuda.empty_cache()
             
-            progress(0.6, desc="Running virtual try-on (this may take 2-5 minutes on CPU)...")
+            device_msg = "GPU - ~30-60 seconds" if self.device == "cuda" else "CPU - ~2-5 minutes"
+            progress(0.6, desc=f"Running virtual try-on on {device_msg}...")
             
             # Run inference
             result = self.pipeline(
@@ -110,19 +127,38 @@ class CatVTONService:
                 width=target_width
             )[0]
             
+            # Clear memory after inference
+            gc.collect()
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+            
             progress(1.0, desc="Complete!")
             
-            return result, "✅ Virtual try-on generated successfully!"
+            return result, f"✅ Virtual try-on generated successfully on {self.device.upper()}!"
             
         except Exception as e:
             import traceback
             error_msg = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
             print(error_msg)
+            
+            # Clear memory on error
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                
             return None, error_msg
 
 # Initialize service
+print("🚀 Initializing CatVTON Service...")
 service = CatVTONService()
 
+# Preload models on startup (optional - comment out if you want lazy loading)
+# try:
+#     service.load_models()
+# except Exception as e:
+#     print(f"⚠️  Could not preload models: {e}")
+#     print("Models will be loaded on first request")
+
 # Create Gradio Interface
 def generate_tryon_interface(person_img, garment_img, progress=gr.Progress()):
     """Wrapper for Gradio"""
@@ -131,7 +167,7 @@ def generate_tryon_interface(person_img, garment_img, progress=gr.Progress()):
 
 # Build UI
 with gr.Blocks(
-    title="Try-Space - CatVTON Virtual Try-On",
+    title="CatVTON Virtual Try-On",
     theme=gr.themes.Soft(),
     css="""
         .gradio-container {max-width: 1200px !important}
@@ -140,13 +176,17 @@ with gr.Blocks(
     """
 ) as demo:
     
-    gr.HTML("""
+    device_info = "🖥️ GPU" if torch.cuda.is_available() else "💻 CPU"
+    processing_time = "30-60 seconds" if torch.cuda.is_available() else "2-5 minutes"
+    
+    gr.HTML(f"""
         <div id="title">
-            <h1>👗 Try-Space - Virtual Try-On</h1>
+            <h1>👗 CatVTON - Virtual Try-On</h1>
         </div>
         <div id="subtitle">
             <p>Upload a person image and a garment to see how it looks on them!</p>
-            <p><strong>⚠️ Note:</strong> Processing takes 2-5 minutes on CPU. First run downloads models (~5GB).</p>
+            <p><strong>Device:</strong> {device_info} | <strong>Processing Time:</strong> ~{processing_time}</p>
+            <p><em>First run downloads models (~5GB) - subsequent runs are faster!</em></p>
         </div>
     """)
     
@@ -176,6 +216,7 @@ with gr.Blocks(
                 - Person should face camera directly
                 - Garment on plain/white background
                 - Works best with shirts, jackets, tops
+                - Avoid images with multiple people
             """)
             
         with gr.Column():
@@ -186,61 +227,38 @@ with gr.Blocks(
             )
             status_output = gr.Textbox(
                 label="Status",
-                lines=3
+                lines=3,
+                show_label=True
             )
     
-    # Examples
-    gr.Markdown("### 📋 Example Images")
-    gr.Examples(
-        examples=[
-            ["examples/person1.jpg", "examples/garment1.jpg"],
-            ["examples/person2.jpg", "examples/garment2.jpg"],
-        ],
-        inputs=[person_input, garment_input],
-        label="Try these examples (if available)"
-    )
-    
-    # API Documentation
-    with gr.Accordion("📡 API Usage", open=False):
-        gr.Markdown(f"""
-        ### Using the Try-Space API
+    # Examples (only show if examples directory exists)
+    if os.path.exists("examples"):
+        gr.Markdown("### 📋 Example Images")
+        example_files = []
+        if os.path.exists("examples/person1.jpg") and os.path.exists("examples/garment1.jpg"):
+            example_files.append(["examples/person1.jpg", "examples/garment1.jpg"])
+        if os.path.exists("examples/person2.jpg") and os.path.exists("examples/garment2.jpg"):
+            example_files.append(["examples/person2.jpg", "examples/garment2.jpg"])
         
-        **Endpoint:** `https://YOUR_USERNAME-try-space.hf.space/api/predict`
-        
-        **Python Example:**
-```python
-        import requests
-        import base64
-        from io import BytesIO
-        from PIL import Image
-        
-        def encode_image(image_path):
-            with open(image_path, "rb") as f:
-                return base64.b64encode(f.read()).decode()
-        
-        url = "https://YOUR_USERNAME-try-space.hf.space/api/predict"
-        
-        response = requests.post(
-            url,
-            json={{
-                "data": [
-                    encode_image("person.jpg"),
-                    encode_image("garment.jpg")
-                ]
-            }}
-        )
+        if example_files:
+            gr.Examples(
+                examples=example_files,
+                inputs=[person_input, garment_input],
+                label="Try these examples"
+            )
+    
+    # Footer
+    gr.Markdown("""
+        ---
+        ### ℹ️ About
+        This app uses **CatVTON** (Concatenation-based Attention Virtual Try-On) for realistic garment transfer.
         
-        result = response.json()
-        print(result)
-```
+        - Model: [zhengchong/CatVTON](https://huggingface.co/zhengchong/CatVTON)
+        - Based on Stable Diffusion Inpainting
+        - Supports upper body garments (shirts, jackets, tops)
         
-        **cURL Example:**
-```bash
-        curl -X POST https://YOUR_USERNAME-try-space.hf.space/api/predict \\
-          -H "Content-Type: application/json" \\
-          -d '{{"data": ["<person_base64>", "<garment_base64>"]}}'
-```
-        """)
+        **Note:** Processing time depends on hardware. GPU is recommended for faster results.
+    """)
     
     # Connect button
     generate_btn.click(
@@ -251,9 +269,21 @@ with gr.Blocks(
 
 # Launch app
 if __name__ == "__main__":
-    demo.queue(max_size=20)  # Enable queue for long-running tasks
+    print("\n" + "="*60)
+    print("🌐 Starting CatVTON Virtual Try-On Server")
+    print("="*60)
+    print(f"Device: {service.device}")
+    print(f"Server: http://0.0.0.0:7860")
+    print("="*60 + "\n")
+    
+    demo.queue(
+        max_size=20,  # Max queue size
+        default_concurrency_limit=2  # Limit concurrent requests
+    )
+    
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        show_error=True
+        show_error=True,
+        share=False  # Don't create public link on HF Spaces
     )