Update app.py

app.py

@@ -1,4 +1,4 @@
-import patch_gradio  # Add this as first import
+import patch_gradio
 import os
 import sys
 import torch
@@ -8,41 +8,33 @@ import gc
 import traceback
 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 and automasker is not None:
         return
     
-    print("🔄 Downloading/Loading CatVTON models (first time may take 5-10 mins)...")
+    print("🔄 Loading models...", file=sys.stderr)
     
     try:
-        # Download and cache models
         repo_path = snapshot_download(
             repo_id="zhengchong/CatVTON",
             cache_dir="/tmp/models"
         )
         
-        print(f"✅ Models downloaded to: {repo_path}")
-        
-        # Create NSFW placeholder in writable directory
         nsfw_path = "/tmp/NSFW.jpg"
         if not os.path.exists(nsfw_path):
             Image.new('RGB', (512, 512), color='black').save(nsfw_path)
         
-        print("Initializing pipeline...")
         pipeline = CatVTONPipeline(
             base_ckpt="booksforcharlie/stable-diffusion-inpainting",
             attn_ckpt=repo_path,
@@ -51,76 +43,48 @@ def load_models():
             use_tf32=True,
             device='cuda'
         )
-        print("✅ Pipeline loaded!")
         
-        print("Initializing automasker...")
         automasker = AutoMasker(
             densepose_ckpt=os.path.join(repo_path, "DensePose"),
             schp_ckpt=os.path.join(repo_path, "SCHP"),
             device='cpu'
         )
-        print("✅ Automasker loaded!")
+        
+        print("✅ Models loaded!", file=sys.stderr)
         
     except Exception as e:
-        print(f"❌ Error loading models: {e}", file=sys.stderr)
+        print(f"❌ Error: {e}", file=sys.stderr)
         traceback.print_exc()
         raise
 
-def generate_tryon(person_img, cloth_img, progress=gr.Progress()):
-    """Generate virtual try-on"""
+def generate_tryon(person_img, cloth_img):
+    """CRITICAL: Removed progress parameter - causes API issues"""
     
-    # ADD EXTENSIVE LOGGING FOR API DEBUGGING
-    print("=" * 50, file=sys.stderr)
-    print(f"API CALL RECEIVED", file=sys.stderr)
-    print(f"Person image type: {type(person_img)}", file=sys.stderr)
-    print(f"Cloth image type: {type(cloth_img)}", file=sys.stderr)
+    print("="*50, file=sys.stderr)
+    print(f"Received - Person: {type(person_img)}, Cloth: {type(cloth_img)}", file=sys.stderr)
     
     if person_img is None or cloth_img is None:
-        error_msg = "Please upload both person and garment images!"
-        print(f"ERROR: {error_msg}", file=sys.stderr)
-        raise gr.Error(error_msg)
+        raise gr.Error("Both images required!")
     
     try:
-        # HANDLE DIFFERENT INPUT TYPES (filepath or PIL)
+        # Convert filepaths to PIL
         if isinstance(person_img, str):
-            print(f"Converting person_img from filepath: {person_img}", file=sys.stderr)
             person_img = Image.open(person_img).convert('RGB')
-        elif not isinstance(person_img, Image.Image):
-            print(f"Converting person_img from array", file=sys.stderr)
-            person_img = Image.fromarray(person_img).convert('RGB')
-            
         if isinstance(cloth_img, str):
-            print(f"Converting cloth_img from filepath: {cloth_img}", file=sys.stderr)
             cloth_img = Image.open(cloth_img).convert('RGB')
-        elif not isinstance(cloth_img, Image.Image):
-            print(f"Converting cloth_img from array", file=sys.stderr)
-            cloth_img = Image.fromarray(cloth_img).convert('RGB')
         
-        print(f"Images loaded successfully", file=sys.stderr)
+        print("Images converted to PIL", file=sys.stderr)
         
-        # 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 on GPU T4 - processing takes 2-3 minutes per image.")
-        
-        # Run inference
         result = pipeline(
             image=person_img,
             condition_image=cloth_img,
@@ -132,92 +96,32 @@ def generate_tryon(person_img, cloth_img, progress=gr.Progress()):
             width=target_width
         )[0]
         
-        progress(1.0, desc="Complete! ✨")
-        
-        print("SUCCESS: Try-on generated successfully", file=sys.stderr)
+        print("✅ Success!", file=sys.stderr)
         return result
         
     except Exception as e:
-        error_msg = f"Error during try-on: {str(e)}"
-        print(f"ERROR: {error_msg}", file=sys.stderr)
+        print(f"❌ Error: {e}", file=sys.stderr)
         traceback.print_exc()
-        raise gr.Error(error_msg)
+        raise gr.Error(str(e))
 
-# Create Gradio UI
-with gr.Blocks(
+# CRITICAL: Use gr.Interface instead of gr.Blocks for better API support
+demo = gr.Interface(
+    fn=generate_tryon,
+    inputs=[
+        gr.Image(label="Person Image", type="filepath"),
+        gr.Image(label="Garment Image", type="filepath")
+    ],
+    outputs=gr.Image(label="Result", type="filepath"),
     title="Try-Space Virtual Try-On",
-    theme=gr.themes.Soft()
-) as demo:
-    
-    gr.Markdown("""
-    # 🎨 Try-Space Virtual Try-On
-    ### Upload a person image and garment to see the magic! ✨
-    
-    ⚠️ **Note:** Running on GPU T4 - processing takes 2-3 minutes per image.
-    """)
-    
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 📸 Inputs")
-            person_input = gr.Image(
-                label="👤 Person Image (full body, front-facing)", 
-                type="filepath",  # CHANGED FROM "pil" TO "filepath"
-                height=350
-            )
-            cloth_input = gr.Image(
-                label="👕 Garment Image (flat, white background)", 
-                type="filepath",  # CHANGED FROM "pil" TO "filepath"
-                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",
-                type="filepath",  # ADDED type
-                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:
-    - **GPU T4:** ~2-3 minutes per generation
-    """)
-    
-    # Event handler with API name
-    submit_btn.click(
-        fn=generate_tryon,
-        inputs=[person_input, cloth_input],
-        outputs=output_img,
-        api_name="generate_tryon"  # ADDED THIS - CRITICAL FOR API ACCESS
-    )
+    description="Upload person and garment images. Processing takes 2-3 minutes on GPU T4.",
+    api_name="generate_tryon",
+    allow_flagging="never"
+)
 
-# Launch app
 if __name__ == "__main__":
-    print("🚀 Starting Try-Space Virtual Try-On...")    
-    # Pre-load models at startup
+    print("🚀 Starting...", file=sys.stderr)
     try:
         load_models()
-    except Exception as e:
-        print(f"⚠️ Model loading will happen on first inference: {e}")
-    
-    # Launch with queue and show_error for better debugging
-    demo.queue().launch(show_error=True)  # ADDED show_error=True
+    except:
+        pass
+    demo.queue().launch(show_error=True)