Create app.py

app.py

@@ -0,0 +1,185 @@
+import os
+import sys
+import torch
+import gradio as gr
+from PIL import Image
+import gc
+from huggingface_hub import snapshot_download
+
+# Add CatVTON to path
+sys.path.insert(0, '/app/CatVTON')
+
+from model.pipeline import CatVTONPipeline
+from model.cloth_masker import AutoMasker
+from utils import init_weight_dtype, resize_and_crop, resize_and_padding
+
+# Global variables
+pipeline = None
+automasker = None
+
+def load_models():
+    """Load models once at startup"""
+    global pipeline, automasker
+    
+    if pipeline is not None:
+        return
+    
+    print("🔄 Downloading/Loading CatVTON models (first time may take 5-10 mins)...")
+    
+    # Download and cache models
+    repo_path = snapshot_download(
+        repo_id="zhengchong/CatVTON",
+        cache_dir="/tmp/models"  # For CPU basic tier
+    )
+    
+    print("✅ Models downloaded! Initializing pipeline...")
+    
+    # Initialize pipeline for CPU
+    pipeline = CatVTONPipeline(
+        base_ckpt="booksforcharlie/stable-diffusion-inpainting",
+        attn_ckpt=repo_path,
+        attn_ckpt_version="mix",
+        weight_dtype=init_weight_dtype("fp32"),  # CPU uses fp32
+        use_tf32=False,  # CPU doesn't support TF32
+        device='cpu'
+    )
+    
+    automasker = AutoMasker(
+        densepose_ckpt=os.path.join(repo_path, "DensePose"),
+        schp_ckpt=os.path.join(repo_path, "SCHP"),
+        device='cpu'
+    )
+    
+    print("✅ Models loaded successfully!")
+
+def generate_tryon(person_img, cloth_img, progress=gr.Progress()):
+    """Generate virtual try-on"""
+    
+    if person_img is None or cloth_img is None:
+        raise gr.Error("Please upload both person and garment images!")
+    
+    try:
+        # Load models
+        progress(0.05, desc="Loading models...")
+        load_models()
+        
+        progress(0.15, desc="Processing images...")
+        
+        # Resize images
+        target_height = 1024
+        target_width = 768
+        person_img = resize_and_crop(person_img, (target_width, target_height))
+        cloth_img = resize_and_padding(cloth_img, (target_width, target_height))
+        
+        progress(0.35, desc="Generating body mask...")
+        
+        # Generate mask
+        mask = automasker(person_img, "upper")['mask']
+        
+        # Clear memory
+        gc.collect()
+        
+        progress(0.50, desc="Running virtual try-on (this may take 5-10 mins on CPU)...")
+        
+        # Run inference
+        result = pipeline(
+            image=person_img,
+            condition_image=cloth_img,
+            mask=mask,
+            num_inference_steps=50,
+            guidance_scale=2.5,
+            seed=42,
+            height=target_height,
+            width=target_width
+        )[0]
+        
+        progress(1.0, desc="Complete! ✨")
+        
+        return result
+        
+    except Exception as e:
+        raise gr.Error(f"Error: {str(e)}")
+
+# Create Gradio UI
+with gr.Blocks(
+    title="CatVTON Virtual Try-On",
+    theme=gr.themes.Soft()
+) as demo:
+    
+    gr.Markdown("""
+    # 🎨 CatVTON Virtual Try-On
+    ### Upload a person image and garment to see the magic! ✨
+    
+    ⚠️ **Note:** Running on CPU - processing takes 5-10 minutes per image.
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 📸 Inputs")
+            person_input = gr.Image(
+                label="👤 Person Image (full body, front-facing)", 
+                type="pil",
+                height=350
+            )
+            cloth_input = gr.Image(
+                label="👕 Garment Image (flat, white background)", 
+                type="pil",
+                height=350
+            )
+            
+            with gr.Row():
+                clear_btn = gr.ClearButton(
+                    [person_input, cloth_input],
+                    value="🗑️ Clear"
+                )
+                submit_btn = gr.Button(
+                    "🚀 Generate Try-On",
+                    variant="primary",
+                    size="lg"
+                )
+        
+        with gr.Column():
+            gr.Markdown("### ✨ Result")
+            output_img = gr.Image(
+                label="Virtual Try-On Result",
+                height=700
+            )
+    
+    gr.Markdown("""
+    ---
+    ### 💡 Tips for Best Results:
+    - ✅ Use well-lit, clear images
+    - ✅ Person should face the camera directly
+    - ✅ Garment should be flat or on white background
+    - ✅ Works best with shirts, jackets, or tops
+    - ✅ Avoid extreme angles or poses
+    
+    ### ⏱️ Processing Time:
+    - **CPU Basic:** ~5-10 minutes per generation
+    - **Upgrade to GPU:** Reduces to ~2-3 minutes
+    """)
+    
+    # Event handler
+    submit_btn.click(
+        fn=generate_tryon,
+        inputs=[person_input, cloth_input],
+        outputs=output_img
+    )
+
+# Launch app
+if __name__ == "__main__":
+    print("🚀 Starting CatVTON Virtual Try-On...")
+    print("⏳ First run will download models (~2-3 GB)...")
+    
+    # Pre-load models at startup
+    try:
+        load_models()
+    except Exception as e:
+        print(f"⚠️ Model loading will happen on first inference: {e}")
+    
+    # Launch
+    demo.queue(max_size=5).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )