Update app.py

app.py

@@ -1,4 +1,3 @@
-import patch_gradio
 import os
 import sys
 import torch
@@ -6,6 +5,8 @@ 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')
@@ -14,13 +15,17 @@ 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():
-    global pipeline, automasker
+    """Load CatVTON models if not already loaded"""
+    global pipeline, automasker, models_loaded
     
-    if pipeline is not None and automasker is not None:
+    if models_loaded and pipeline is not None and automasker is not None:
+        print("✅ Models already loaded", file=sys.stderr)
         return
     
     print("🔄 Loading models...", file=sys.stderr)
@@ -31,60 +36,139 @@ 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'
+            device='cuda' if torch.cuda.is_available() else 'cpu'
         )
         
+        # Initialize automasker
         automasker = AutoMasker(
             densepose_ckpt=os.path.join(repo_path, "DensePose"),
             schp_ckpt=os.path.join(repo_path, "SCHP"),
             device='cpu'
         )
         
-        print("✅ Models loaded!", file=sys.stderr)
+        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()
         
     except Exception as e:
-        print(f"❌ Error: {e}", file=sys.stderr)
-        traceback.print_exc()
+        print(f"❌ Error loading models: {e}", file=sys.stderr)
+        traceback.print_exc(file=sys.stderr)
+        models_loaded = False
         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):
-    """CRITICAL: Removed progress parameter - causes API issues"""
+    """
+    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)
     
+    Returns:
+        PIL Image of the try-on result
+    """
     print("="*50, file=sys.stderr)
-    print(f"Received - Person: {type(person_img)}, Cloth: {type(cloth_img)}", file=sys.stderr)
+    print(f"📥 Received inputs - Person: {type(person_img)}, Cloth: {type(cloth_img)}", file=sys.stderr)
     
     if person_img is None or cloth_img is None:
-        raise gr.Error("Both images required!")
+        error_msg = "Both person and garment images are required!"
+        print(f"❌ {error_msg}", file=sys.stderr)
+        raise gr.Error(error_msg)
     
     try:
-        # 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')
+        # 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)
         
-        print("Images converted to PIL", file=sys.stderr)
+        print(f"✅ Images converted - Person: {person_img.size}, Cloth: {cloth_img.size}", 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,
@@ -96,32 +180,77 @@ def generate_tryon(person_img, cloth_img):
             width=target_width
         )[0]
         
-        print("✅ Success!", file=sys.stderr)
+        # Clean up
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        print("✅ Try-on generation completed successfully!", file=sys.stderr)
         return result
         
+    except gr.Error:
+        # Re-raise Gradio errors as-is
+        raise
     except Exception as e:
-        print(f"❌ Error: {e}", file=sys.stderr)
-        traceback.print_exc()
-        raise gr.Error(str(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)
 
-# CRITICAL: Use gr.Interface instead of gr.Blocks for better API support
+# Create Gradio Interface with proper API configuration
 demo = gr.Interface(
     fn=generate_tryon,
     inputs=[
-        gr.Image(label="Person Image", type="filepath"),
-        gr.Image(label="Garment Image", type="filepath")
+        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"],
+        )
     ],
-    outputs=gr.Image(label="Result", type="filepath"),
+    outputs=gr.Image(
+        label="Try-On Result",
+        type="pil"  # Return PIL Image for better API compatibility
+    ),
     title="Try-Space Virtual Try-On",
-    description="Upload person and garment images. Processing takes 2-3 minutes on GPU T4.",
-    api_name="generate_tryon",
-    allow_flagging="never"
+    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
 )
 
 if __name__ == "__main__":
-    print("🚀 Starting...", file=sys.stderr)
+    print("🚀 Starting Try-Space Virtual Try-On Space...", file=sys.stderr)
+    
+    # Try to load models at startup (non-blocking)
     try:
+        print("🔄 Pre-loading models...", file=sys.stderr)
         load_models()
-    except:
-        pass
-    demo.queue().launch(show_error=True)
+    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
+    )