Update app.py

app.py

@@ -1,300 +1,339 @@
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.responses import JSONResponse, StreamingResponse
-from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
-import torch
-import io
-import base64
-import tempfile
-import os
-from diffusers import AutoPipelineForInpainting, AutoencoderKL
-from typing import Optional
-import time
-
-# Initialize FastAPI app
-app = FastAPI(
-    title="Virtual Try-On API",
-    description="API for virtual clothing try-on using Stable Diffusion XL",
-    version="1.0.0"
-)
-
-# Add CORS middleware for React Native
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Adjust for production
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Global variables for models
-pipeline = None
-segment_body = None
-
-def load_models():
-    """Load all required models"""
-    global pipeline, segment_body
-    
-    print("🔄 Loading VAE...")
-    vae = AutoencoderKL.from_pretrained(
-        "madebyollin/sdxl-vae-fp16-fix",
-        torch_dtype=torch.float16
-    )
-    
-    print("🔄 Loading inpainting pipeline...")
-    pipeline = AutoPipelineForInpainting.from_pretrained(
-        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
-        vae=vae,
-        torch_dtype=torch.float16,
-        variant="fp16",
-        use_safetensors=True
-    )
-    
-    # Use CPU for Hugging Face Spaces
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipeline = pipeline.to(device)
-    
-    print("🔄 Loading IP-Adapter...")
-    pipeline.load_ip_adapter(
-        "h94/IP-Adapter",
-        subfolder="sdxl_models",
-        weight_name="ip-adapter_sdxl.bin",
-        low_cpu_mem_usage=True
-    )
-    
-    print("🔄 Loading body segmentation...")
-    try:
-        from SegBody import segment_body as seg_func
-        segment_body = seg_func
-        print("✅ Body segmentation loaded!")
-    except ImportError:
-        print("⚠️ SegBody module not found, segmentation will be disabled")
-    
-    print("✅ All models loaded successfully!")
-
-@app.on_event("startup")
-async def startup_event():
-    """Load models on startup"""
-    load_models()
-
-@app.get("/")
-async def root():
-    """Health check endpoint"""
-    return {
-        "status": "running",
-        "message": "Virtual Try-On API is running!",
-        "cuda_available": torch.cuda.is_available(),
-        "device": "cuda" if torch.cuda.is_available() else "cpu"
-    }
-
-@app.get("/health")
-async def health():
-    """Health check endpoint"""
-    return {"status": "healthy", "models_loaded": pipeline is not None}
-
-def image_to_base64(image: Image.Image) -> str:
-    """Convert PIL Image to base64 string"""
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return img_str
-
-def base64_to_image(base64_str: str) -> Image.Image:
-    """Convert base64 string to PIL Image"""
-    img_data = base64.b64decode(base64_str)
-    return Image.open(io.BytesIO(img_data)).convert('RGB')
-
-@app.post("/tryon")
-async def virtual_tryon(
-    person_image: UploadFile = File(..., description="Image of the person"),
-    clothing_image: UploadFile = File(..., description="Image of the clothing"),
-    prompt: str = Form("photorealistic, perfect body, beautiful skin, realistic skin, natural skin"),
-    negative_prompt: str = Form("ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings"),
-    ip_scale: float = Form(0.8),
-    strength: float = Form(0.99),
-    guidance_scale: float = Form(7.5),
-    num_steps: int = Form(50),
-    return_format: str = Form("base64", description="base64 or image")
-):
-    """
-    Virtual Try-On endpoint
-    
-    Args:
-        person_image: Image file of the person
-        clothing_image: Image file of the clothing
-        prompt: Generation prompt
-        negative_prompt: Negative prompt
-        ip_scale: IP-Adapter influence (0.0-1.0)
-        strength: Inpainting strength (0.0-1.0)
-        guidance_scale: CFG scale
-        num_steps: Number of inference steps
-        return_format: Response format (base64 or image)
-    
-    Returns:
-        Generated image in specified format
-    """
-    try:
-        if pipeline is None:
-            raise HTTPException(status_code=503, detail="Models not loaded yet")
-        
-        start_time = time.time()
-        
-        # Load and resize images
-        print("📥 Loading images...")
-        person_img = Image.open(person_image.file).convert('RGB').resize((512, 512))
-        clothing_img = Image.open(clothing_image.file).convert('RGB').resize((512, 512))
-        
-        # Generate body segmentation mask
-        print("🎭 Generating segmentation mask...")
-        if segment_body is None:
-            # Create a simple fallback mask (full body) if segmentation not available
-            mask_img = Image.new('L', (512, 512), 255)
-        else:
-            try:
-                # Save person image to temp file for segmentation
-                with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
-                    temp_path = tmp_file.name
-                    person_img.save(temp_path)
-                
-                seg_image, mask_img = segment_body(temp_path, face=False)
-                mask_img = mask_img.resize((512, 512))
-                
-                # Clean up
-                os.unlink(temp_path)
-            except Exception as e:
-                print(f"⚠️ Segmentation failed: {e}, using full mask")
-                mask_img = Image.new('L', (512, 512), 255)
-        
-        # Set IP-Adapter scale
-        pipeline.set_ip_adapter_scale(ip_scale)
-        
-        # Generate virtual try-on
-        print("🎨 Generating virtual try-on...")
-        result = pipeline(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=person_img,
-            mask_image=mask_img,
-            ip_adapter_image=clothing_img,
-            strength=strength,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_steps,
-        )
-        
-        generated_image = result.images[0]
-        
-        processing_time = time.time() - start_time
-        print(f"✅ Generation completed in {processing_time:.2f}s")
-        
-        # Return based on format
-        if return_format == "image":
-            # Return as image file
-            img_byte_arr = io.BytesIO()
-            generated_image.save(img_byte_arr, format='PNG')
-            img_byte_arr.seek(0)
-            return StreamingResponse(img_byte_arr, media_type="image/png")
-        else:
-            # Return as base64 JSON
-            img_base64 = image_to_base64(generated_image)
-            return JSONResponse({
-                "success": True,
-                "image": img_base64,
-                "processing_time": processing_time,
-                "parameters": {
-                    "prompt": prompt,
-                    "ip_scale": ip_scale,
-                    "strength": strength,
-                    "guidance_scale": guidance_scale,
-                    "num_steps": num_steps
-                }
-            })
-    
-    except Exception as e:
-        print(f"❌ Error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/tryon-base64")
-async def virtual_tryon_base64(
-    person_image_base64: str = Form(..., description="Base64 encoded person image"),
-    clothing_image_base64: str = Form(..., description="Base64 encoded clothing image"),
-    prompt: str = Form("photorealistic, perfect body, beautiful skin, realistic skin, natural skin"),
-    negative_prompt: str = Form("ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings"),
-    ip_scale: float = Form(0.8),
-    strength: float = Form(0.99),
-    guidance_scale: float = Form(7.5),
-    num_steps: int = Form(50)
-):
-    """
-    Virtual Try-On endpoint accepting base64 encoded images
-    (Alternative endpoint for easier React Native integration)
-    """
-    try:
-        if pipeline is None:
-            raise HTTPException(status_code=503, detail="Models not loaded yet")
-        
-        start_time = time.time()
-        
-        # Decode base64 images
-        print("📥 Decoding base64 images...")
-        person_img = base64_to_image(person_image_base64).resize((512, 512))
-        clothing_img = base64_to_image(clothing_image_base64).resize((512, 512))
-        
-        # Generate body segmentation mask
-        print("🎭 Generating segmentation mask...")
-        if segment_body is None:
-            mask_img = Image.new('L', (512, 512), 255)
-        else:
-            try:
-                with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
-                    temp_path = tmp_file.name
-                    person_img.save(temp_path)
-                
-                seg_image, mask_img = segment_body(temp_path, face=False)
-                mask_img = mask_img.resize((512, 512))
-                os.unlink(temp_path)
-            except Exception as e:
-                print(f"⚠️ Segmentation failed: {e}, using full mask")
-                mask_img = Image.new('L', (512, 512), 255)
-        
-        # Set IP-Adapter scale
-        pipeline.set_ip_adapter_scale(ip_scale)
-        
-        # Generate virtual try-on
-        print("🎨 Generating virtual try-on...")
-        result = pipeline(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=person_img,
-            mask_image=mask_img,
-            ip_adapter_image=clothing_img,
-            strength=strength,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_steps,
-        )
-        
-        generated_image = result.images[0]
-        processing_time = time.time() - start_time
-        print(f"✅ Generation completed in {processing_time:.2f}s")
-        
-        # Return as base64
-        img_base64 = image_to_base64(generated_image)
-        return JSONResponse({
-            "success": True,
-            "image": img_base64,
-            "processing_time": processing_time,
-            "parameters": {
-                "prompt": prompt,
-                "ip_scale": ip_scale,
-                "strength": strength,
-                "guidance_scale": guidance_scale,
-                "num_steps": num_steps
-            }
-        })
-    
-    except Exception as e:
-        print(f"❌ Error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
+import torch
+import io
+import base64
+import tempfile
+import os
+from diffusers import AutoPipelineForInpainting, AutoencoderKL
+from typing import Optional
+import time
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Virtual Try-On API",
+    description="API for virtual clothing try-on using Stable Diffusion XL",
+    version="1.0.0"
+)
+
+# Add CORS middleware for React Native
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Adjust for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global variables for models
+pipeline = None
+segment_body = None
+
+def load_models():
+    """Load all required models"""
+    global pipeline, segment_body
+    
+    # Determine device and dtype
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    
+    print(f"🔄 Using device: {device}, dtype: {dtype}")
+    
+    print("🔄 Loading VAE...")
+    if torch.cuda.is_available():
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix",
+            torch_dtype=dtype
+        )
+    else:
+        # Use standard VAE for CPU (fp32)
+        vae = AutoencoderKL.from_pretrained(
+            "stabilityai/sdxl-vae",
+            torch_dtype=dtype
+        )
+    
+    print("🔄 Loading inpainting pipeline...")
+    if torch.cuda.is_available():
+        pipeline = AutoPipelineForInpainting.from_pretrained(
+            "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+            vae=vae,
+            torch_dtype=dtype,
+            variant="fp16",
+            use_safetensors=True
+        )
+    else:
+        # Load without fp16 variant for CPU
+        pipeline = AutoPipelineForInpainting.from_pretrained(
+            "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+            vae=vae,
+            torch_dtype=dtype,
+            use_safetensors=True
+        )
+    
+    pipeline = pipeline.to(device)
+    
+    print("🔄 Loading IP-Adapter...")
+    pipeline.load_ip_adapter(
+        "h94/IP-Adapter",
+        subfolder="sdxl_models",
+        weight_name="ip-adapter_sdxl.bin",
+        low_cpu_mem_usage=True
+    )
+    
+    print("🔄 Loading body segmentation...")
+    try:
+        from SegBody import segment_body as seg_func
+        segment_body = seg_func
+        print("✅ Body segmentation loaded!")
+    except ImportError:
+        print("⚠️ SegBody module not found, segmentation will be disabled")
+    
+    print("✅ All models loaded successfully!")
+
+@app.on_event("startup")
+async def startup_event():
+    """Load models on startup"""
+    load_models()
+
+@app.get("/")
+async def root():
+    """Health check endpoint"""
+    return {
+        "status": "running",
+        "message": "Virtual Try-On API is running!",
+        "cuda_available": torch.cuda.is_available(),
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }
+
+@app.get("/health")
+async def health():
+    """Health check endpoint"""
+    return {"status": "healthy", "models_loaded": pipeline is not None}
+
+def image_to_base64(image: Image.Image) -> str:
+    """Convert PIL Image to base64 string"""
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return img_str
+
+def base64_to_image(base64_str: str) -> Image.Image:
+    """Convert base64 string to PIL Image"""
+    img_data = base64.b64decode(base64_str)
+    return Image.open(io.BytesIO(img_data)).convert('RGB')
+
+@app.post("/tryon")
+async def virtual_tryon(
+    person_image: UploadFile = File(..., description="Image of the person"),
+    clothing_image: UploadFile = File(..., description="Image of the clothing"),
+    prompt: str = Form("photorealistic, perfect body, beautiful skin, realistic skin, natural skin"),
+    negative_prompt: str = Form("ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings"),
+    ip_scale: float = Form(0.8),
+    strength: float = Form(0.99),
+    guidance_scale: float = Form(7.5),
+    num_steps: int = Form(50),
+    return_format: str = Form("base64", description="base64 or image")
+):
+    """
+    Virtual Try-On endpoint
+    
+    Args:
+        person_image: Image file of the person
+        clothing_image: Image file of the clothing
+        prompt: Generation prompt
+        negative_prompt: Negative prompt
+        ip_scale: IP-Adapter influence (0.0-1.0)
+        strength: Inpainting strength (0.0-1.0)
+        guidance_scale: CFG scale
+        num_steps: Number of inference steps
+        return_format: Response format (base64 or image)
+    
+    Returns:
+        Generated image in specified format
+    """
+    try:
+        if pipeline is None:
+            raise HTTPException(status_code=503, detail="Models not loaded yet")
+        
+        start_time = time.time()
+        
+        # Load and resize images
+        print("📥 Loading images...")
+        person_img = Image.open(person_image.file).convert('RGB').resize((512, 512))
+        clothing_img = Image.open(clothing_image.file).convert('RGB').resize((512, 512))
+        
+        # Generate body segmentation mask
+        print("🎭 Generating segmentation mask...")
+        if segment_body is None:
+            # Create a simple fallback mask (full body) if segmentation not available
+            mask_img = Image.new('L', (512, 512), 255)
+        else:
+            try:
+                # Try calling segment_body - it might expect a file path or PIL Image
+                try:
+                    # First try with PIL Image directly
+                    seg_image, mask_img = segment_body(person_img, face=False)
+                except (TypeError, AttributeError):
+                    # If that fails, try with file path
+                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
+                        temp_path = tmp_file.name
+                        person_img.save(temp_path)
+                    
+                    seg_image, mask_img = segment_body(temp_path, face=False)
+                    os.unlink(temp_path)
+                
+                # Ensure mask is PIL Image and resize
+                if isinstance(mask_img, str):
+                    mask_img = Image.open(mask_img).convert('L')
+                mask_img = mask_img.resize((512, 512))
+                
+            except Exception as e:
+                print(f"⚠️ Segmentation failed: {e}, using full mask")
+                mask_img = Image.new('L', (512, 512), 255)
+        
+        # Set IP-Adapter scale
+        pipeline.set_ip_adapter_scale(ip_scale)
+        
+        # Generate virtual try-on
+        print("🎨 Generating virtual try-on...")
+        result = pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=person_img,
+            mask_image=mask_img,
+            ip_adapter_image=clothing_img,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_steps,
+        )
+        
+        generated_image = result.images[0]
+        
+        processing_time = time.time() - start_time
+        print(f"✅ Generation completed in {processing_time:.2f}s")
+        
+        # Return based on format
+        if return_format == "image":
+            # Return as image file
+            img_byte_arr = io.BytesIO()
+            generated_image.save(img_byte_arr, format='PNG')
+            img_byte_arr.seek(0)
+            return StreamingResponse(img_byte_arr, media_type="image/png")
+        else:
+            # Return as base64 JSON
+            img_base64 = image_to_base64(generated_image)
+            return JSONResponse({
+                "success": True,
+                "image": img_base64,
+                "processing_time": processing_time,
+                "parameters": {
+                    "prompt": prompt,
+                    "ip_scale": ip_scale,
+                    "strength": strength,
+                    "guidance_scale": guidance_scale,
+                    "num_steps": num_steps
+                }
+            })
+    
+    except Exception as e:
+        print(f"❌ Error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/tryon-base64")
+async def virtual_tryon_base64(
+    person_image_base64: str = Form(..., description="Base64 encoded person image"),
+    clothing_image_base64: str = Form(..., description="Base64 encoded clothing image"),
+    prompt: str = Form("photorealistic, perfect body, beautiful skin, realistic skin, natural skin"),
+    negative_prompt: str = Form("ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings"),
+    ip_scale: float = Form(0.8),
+    strength: float = Form(0.99),
+    guidance_scale: float = Form(7.5),
+    num_steps: int = Form(50)
+):
+    """
+    Virtual Try-On endpoint accepting base64 encoded images
+    (Alternative endpoint for easier React Native integration)
+    """
+    try:
+        if pipeline is None:
+            raise HTTPException(status_code=503, detail="Models not loaded yet")
+        
+        start_time = time.time()
+        
+        # Decode base64 images
+        print("📥 Decoding base64 images...")
+        person_img = base64_to_image(person_image_base64).resize((512, 512))
+        clothing_img = base64_to_image(clothing_image_base64).resize((512, 512))
+        
+        # Generate body segmentation mask
+        print("🎭 Generating segmentation mask...")
+        if segment_body is None:
+            mask_img = Image.new('L', (512, 512), 255)
+        else:
+            try:
+                # Try calling segment_body - it might expect a file path or PIL Image
+                try:
+                    # First try with PIL Image directly
+                    seg_image, mask_img = segment_body(person_img, face=False)
+                except (TypeError, AttributeError):
+                    # If that fails, try with file path
+                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
+                        temp_path = tmp_file.name
+                        person_img.save(temp_path)
+                    
+                    seg_image, mask_img = segment_body(temp_path, face=False)
+                    os.unlink(temp_path)
+                
+                # Ensure mask is PIL Image and resize
+                if isinstance(mask_img, str):
+                    mask_img = Image.open(mask_img).convert('L')
+                mask_img = mask_img.resize((512, 512))
+                
+            except Exception as e:
+                print(f"⚠️ Segmentation failed: {e}, using full mask")
+                mask_img = Image.new('L', (512, 512), 255)
+        
+        # Set IP-Adapter scale
+        pipeline.set_ip_adapter_scale(ip_scale)
+        
+        # Generate virtual try-on
+        print("🎨 Generating virtual try-on...")
+        result = pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=person_img,
+            mask_image=mask_img,
+            ip_adapter_image=clothing_img,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_steps,
+        )
+        
+        generated_image = result.images[0]
+        processing_time = time.time() - start_time
+        print(f"✅ Generation completed in {processing_time:.2f}s")
+        
+        # Return as base64
+        img_base64 = image_to_base64(generated_image)
+        return JSONResponse({
+            "success": True,
+            "image": img_base64,
+            "processing_time": processing_time,
+            "parameters": {
+                "prompt": prompt,
+                "ip_scale": ip_scale,
+                "strength": strength,
+                "guidance_scale": guidance_scale,
+                "num_steps": num_steps
+            }
+        })
+    
+    except Exception as e:
+        print(f"❌ Error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)