Update app.py

app.py

@@ -1,3 +1,4 @@
+import patch_gradio
 import os
 import sys
 import torch
@@ -5,8 +6,6 @@ import gradio as gr
 from PIL import Image
 import gc
 import traceback
-import base64
-import io
 from huggingface_hub import snapshot_download
 
 sys.path.insert(0, '/app/CatVTON')
@@ -15,17 +14,13 @@ from model.pipeline import CatVTONPipeline
 from model.cloth_masker import AutoMasker
 from utils import init_weight_dtype, resize_and_crop, resize_and_padding
 
-# Global model variables
 pipeline = None
 automasker = None
-models_loaded = False
 
 def load_models():
-    """Load CatVTON models if not already loaded"""
-    global pipeline, automasker, models_loaded
+    global pipeline, automasker
     
-    if models_loaded and pipeline is not None and automasker is not None:
-        print("✅ Models already loaded", file=sys.stderr)
+    if pipeline is not None and automasker is not None:
         return
     
     print("🔄 Loading models...", file=sys.stderr)
@@ -36,139 +31,60 @@ def load_models():
             cache_dir="/tmp/models"
         )
         
-        # Create NSFW placeholder if needed
         nsfw_path = "/tmp/NSFW.jpg"
         if not os.path.exists(nsfw_path):
             Image.new('RGB', (512, 512), color='black').save(nsfw_path)
         
-        # Initialize pipeline
         pipeline = CatVTONPipeline(
             base_ckpt="booksforcharlie/stable-diffusion-inpainting",
             attn_ckpt=repo_path,
             attn_ckpt_version="mix",
             weight_dtype=torch.float16,
             use_tf32=True,
-            device='cuda' if torch.cuda.is_available() else 'cpu'
+            device='cuda'
         )
         
-        # Initialize automasker
         automasker = AutoMasker(
             densepose_ckpt=os.path.join(repo_path, "DensePose"),
             schp_ckpt=os.path.join(repo_path, "SCHP"),
             device='cpu'
         )
         
-        models_loaded = True
-        print("✅ Models loaded successfully!", file=sys.stderr)
-        
-        # Force garbage collection after loading
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        print("✅ Models loaded!", file=sys.stderr)
         
     except Exception as e:
-        print(f"❌ Error loading models: {e}", file=sys.stderr)
-        traceback.print_exc(file=sys.stderr)
-        models_loaded = False
+        print(f"❌ Error: {e}", file=sys.stderr)
+        traceback.print_exc()
         raise
 
-def _convert_to_pil_image(image_input):
-    """Convert various input types to PIL Image"""
-    if image_input is None:
-        return None
-    
-    # Already a PIL Image
-    if isinstance(image_input, Image.Image):
-        return image_input.convert('RGB')
-    
-    # File path (string)
-    if isinstance(image_input, str):
-        # Check if it's a base64 string
-        if image_input.startswith('data:image') or (len(image_input) > 100 and not os.path.exists(image_input)):
-            # Try to decode as base64
-            try:
-                if ',' in image_input:
-                    # Remove data URI prefix
-                    base64_data = image_input.split(',')[1]
-                else:
-                    base64_data = image_input
-                
-                image_bytes = base64.b64decode(base64_data)
-                return Image.open(io.BytesIO(image_bytes)).convert('RGB')
-            except Exception as e:
-                print(f"⚠️ Failed to decode base64, trying as file path: {e}", file=sys.stderr)
-        
-        # Try as file path
-        if os.path.exists(image_input):
-            return Image.open(image_input).convert('RGB')
-        else:
-            raise ValueError(f"Image path does not exist: {image_input}")
-    
-    # Bytes or bytearray
-    if isinstance(image_input, (bytes, bytearray)):
-        return Image.open(io.BytesIO(image_input)).convert('RGB')
-    
-    # Try to convert using PIL's open
-    try:
-        return Image.open(image_input).convert('RGB')
-    except Exception as e:
-        raise ValueError(f"Unable to convert input to PIL Image: {type(image_input)}, error: {e}")
-
 def generate_tryon(person_img, cloth_img):
-    """
-    Generate virtual try-on result from person and garment images.
-    
-    Args:
-        person_img: Person image (file path, PIL Image, base64 string, or bytes)
-        cloth_img: Garment image (file path, PIL Image, base64 string, or bytes)
+    """CRITICAL: Removed progress parameter - causes API issues"""
     
-    Returns:
-        PIL Image of the try-on result
-    """
     print("="*50, file=sys.stderr)
-    print(f"📥 Received inputs - Person: {type(person_img)}, Cloth: {type(cloth_img)}", 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 = "Both person and garment images are required!"
-        print(f"❌ {error_msg}", file=sys.stderr)
-        raise gr.Error(error_msg)
+        raise gr.Error("Both images required!")
     
     try:
-        # Convert inputs to PIL Images
-        print("🔄 Converting inputs to PIL Images...", file=sys.stderr)
-        person_img = _convert_to_pil_image(person_img)
-        cloth_img = _convert_to_pil_image(cloth_img)
-        
-        if person_img is None or cloth_img is None:
-            error_msg = "Failed to convert images to PIL format!"
-            print(f"❌ {error_msg}", file=sys.stderr)
-            raise gr.Error(error_msg)
+        # Convert filepaths to PIL
+        if isinstance(person_img, str):
+            person_img = Image.open(person_img).convert('RGB')
+        if isinstance(cloth_img, str):
+            cloth_img = Image.open(cloth_img).convert('RGB')
         
-        print(f"✅ Images converted - Person: {person_img.size}, Cloth: {cloth_img.size}", file=sys.stderr)
+        print("Images converted to PIL", file=sys.stderr)
         
-        # Load models if not already loaded
         load_models()
         
-        if pipeline is None or automasker is None:
-            error_msg = "Failed to load models. Please try again."
-            print(f"❌ {error_msg}", file=sys.stderr)
-            raise gr.Error(error_msg)
-        
-        # Resize images to target dimensions
         target_height = 1024
         target_width = 768
-        print(f"🔄 Resizing images to {target_width}x{target_height}...", file=sys.stderr)
-        
         person_img = resize_and_crop(person_img, (target_width, target_height))
         cloth_img = resize_and_padding(cloth_img, (target_width, target_height))
         
-        # Generate mask
-        print("🔄 Generating mask...", file=sys.stderr)
         mask = automasker(person_img, "upper")['mask']
         gc.collect()
         
-        # Generate try-on result
-        print("🔄 Generating try-on result (this may take 2-3 minutes)...", file=sys.stderr)
         result = pipeline(
             image=person_img,
             condition_image=cloth_img,
@@ -180,77 +96,32 @@ def generate_tryon(person_img, cloth_img):
             width=target_width
         )[0]
         
-        # Clean up
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        print("✅ Try-on generation completed successfully!", file=sys.stderr)
+        print("✅ Success!", file=sys.stderr)
         return result
         
-    except gr.Error:
-        # Re-raise Gradio errors as-is
-        raise
     except Exception as e:
-        error_msg = f"Error during try-on generation: {str(e)}"
-        print(f"❌ {error_msg}", file=sys.stderr)
-        traceback.print_exc(file=sys.stderr)
-        raise gr.Error(error_msg)
+        print(f"❌ Error: {e}", file=sys.stderr)
+        traceback.print_exc()
+        raise gr.Error(str(e))
 
-# Create Gradio Interface with proper API configuration
+# 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",  # Accepts file paths, but we handle other types in the function
-            sources=["upload", "webcam"],
-        ),
-        gr.Image(
-            label="Garment Image",
-            type="filepath",
-            sources=["upload", "webcam"],
-        )
+        gr.Image(label="Person Image", type="filepath"),
+        gr.Image(label="Garment Image", type="filepath")
     ],
-    outputs=gr.Image(
-        label="Try-On Result",
-        type="pil"  # Return PIL Image for better API compatibility
-    ),
+    outputs=gr.Image(label="Result", type="filepath"),
     title="Try-Space Virtual Try-On",
-    description="""
-    Upload person and garment images to generate a virtual try-on result.
-    
-    **Processing Time:** 2-3 minutes on GPU T4
-    
-    **Tips:**
-    - Use clear, well-lit images
-    - Person should be facing forward
-    - Garment should be on a plain background
-    """,
-    api_name="generate_tryon",  # Named endpoint for API access
-    allow_flagging="never",
-    examples=None,  # Can add examples later if needed
+    description="Upload person and garment images. Processing takes 2-3 minutes on GPU T4.",
+    api_name="generate_tryon",
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    print("🚀 Starting Try-Space Virtual Try-On Space...", file=sys.stderr)
-    
-    # Try to load models at startup (non-blocking)
+    print("🚀 Starting...", file=sys.stderr)
     try:
-        print("🔄 Pre-loading models...", file=sys.stderr)
         load_models()
-    except Exception as e:
-        print(f"⚠️ Failed to pre-load models: {e}", file=sys.stderr)
-        print("⚠️ Models will be loaded on first request", file=sys.stderr)
-    
-    # Launch with queue for better API handling
-    demo.queue(
-        max_size=10,  # Limit queue size
-        default_concurrency_limit=1  # Process one request at a time
-    ).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        share=False,
-        enable_queue=True
-    )
+    except:
+        pass
+    demo.queue().launch(show_error=True)