Fix: Corrected StableVITONInference initialization and aligned logs

main.py

@@ -3,7 +3,9 @@ FastAPI Backend for StableVITON Virtual Try-On
 Provides REST API endpoint for virtual try-on inference
 """
 
-from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
+from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Form, Request
+from fastapi.exceptions import RequestValidationError
+from contextlib import asynccontextmanager
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
@@ -27,11 +29,34 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan context manager for startup and shutdown events"""
+    global model
+    try:
+        logger.info("Initializing StableVITON Inference Engine (Lifespan Startup)...")
+        # Initialize the local inference engine
+        model = StableVITONInference(
+            model_path=os.getenv("MODEL_PATH", "yisol/IDM-VTON"),
+            device="auto"
+        )
+        logger.info("StableVITON Inference Engine initialized successfully")
+    except Exception:
+        logger.exception("Failed to initialize StableVITON Inference Engine during startup")
+        model = None
+        
+    yield
+    
+    if model:
+        model.cleanup()
+        logger.info("Model cleaned up (Lifespan Shutdown)")
+
 # Initialize FastAPI app
 app = FastAPI(
     title="StableVITON Virtual Try-On API",
     description="AI-powered virtual try-on service using StableVITON",
-    version="1.0.0"
+    version="1.0.0",
+    lifespan=lifespan
 )
 
 # CORS configuration - allow all origins for demo (restrict in production)
@@ -43,6 +68,28 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    # Log the detailed error to terminal for the user
+    # Convert errors to a string-safe list to avoid JSON serialization issues with ctx/ValueError
+    safe_errors = []
+    for err in exc.errors():
+        safe_err = {k: (str(v) if k in ['ctx', 'input'] else v) for k, v in err.items()}
+        safe_errors.append(safe_err)
+        
+    print(f"\n[!] VALIDATION ERROR: {safe_errors}")
+    logger.error(f"Validation error details: {safe_errors}")
+    
+    return JSONResponse(
+        status_code=422,
+        content={
+            "success": False,
+            "error": "Request validation failed",
+            "details": safe_errors,
+            "error_code": "VALIDATION_ERROR"
+        },
+    )
+
 # Global model instance (loaded once at startup)
 model: Optional[StableVITONInference] = None
 
@@ -66,63 +113,17 @@ class TryOnResponse(BaseModel):
     error_code: Optional[str] = None
 
 
-@app.on_event("startup")
-async def startup_event():
-    """Load model on startup"""
-    global model
-    try:
-        logger.info("Loading StableVITON model...")
-        
-        model = StableVITONInference(
-            model_path=os.getenv("MODEL_PATH", "yisol/IDM-VTON"),
-            device=os.getenv("DEVICE", "auto"),
-            resolution=int(os.getenv("RESOLUTION", "768")),
-            num_inference_steps=int(os.getenv("NUM_INFERENCE_STEPS", "30")),
-            guidance_scale=float(os.getenv("GUIDANCE_SCALE", "2.5"))
-        )
-
-        logger.info("Model loaded successfully")
-        
-        # Warmup inference
-        logger.info("Running warmup inference...")
-        dummy_person = Image.new("RGB", (512, 768), color=(200, 200, 200))
-        dummy_garment = Image.new("RGB", (512, 512), color=(100, 150, 200))
-        model.tryon(dummy_person, dummy_garment)
-        model.cleanup()
-        logger.info("Warmup completed")
-        
-    except Exception as e:
-        logger.error(f"Failed to load model: {e}")
-        raise
-
-
-@app.on_event("shutdown")
-async def shutdown_event():
-    """Cleanup on shutdown"""
-    global model
-    if model:
-        model.cleanup()
-        logger.info("Model cleaned up")
+# Note: Startup and shutdown are now handled by lifespan context manager
 
 
 def validate_image(file: UploadFile) -> None:
-    """
-    Validate uploaded image file.
-    
-    Args:
-        file: Uploaded file
-        
-    Raises:
-        HTTPException: If validation fails
-    """
-    # Check content type
+    """Validate uploaded image file."""
     if file.content_type not in ALLOWED_EXTENSIONS:
         raise HTTPException(
             status_code=400,
             detail=f"Invalid file type. Allowed: {ALLOWED_EXTENSIONS}"
         )
     
-    # Check file size (this is approximate, actual size checked during read)
     if hasattr(file, 'size') and file.size > MAX_IMAGE_SIZE:
         raise HTTPException(
             status_code=400,
@@ -131,15 +132,7 @@ def validate_image(file: UploadFile) -> None:
 
 
 def image_to_base64(image: Image.Image) -> str:
-    """
-    Convert PIL Image to base64 string.
-    
-    Args:
-        image: PIL Image
-        
-    Returns:
-        Base64 encoded string with data URI prefix
-    """
+    """Convert PIL Image to base64 string."""
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
@@ -148,22 +141,9 @@ def image_to_base64(image: Image.Image) -> str:
 
 
 async def read_image_from_upload(file: UploadFile) -> Image.Image:
-    """
-    Read PIL Image from uploaded file.
-    
-    Args:
-        file: Uploaded file
-        
-    Returns:
-        PIL Image
-        
-    Raises:
-        HTTPException: If image cannot be read
-    """
+    """Read PIL Image from uploaded file."""
     try:
         contents = await file.read()
-        
-        # Check actual size
         if len(contents) > MAX_IMAGE_SIZE:
             raise HTTPException(
                 status_code=400,
@@ -171,20 +151,6 @@ async def read_image_from_upload(file: UploadFile) -> Image.Image:
             )
         
         image = Image.open(io.BytesIO(contents))
-        
-        # Validate dimensions
-        width, height = image.size
-        if width < 256 or height < 256:
-            raise HTTPException(
-                status_code=400,
-                detail="Image too small. Minimum dimensions: 256x256"
-            )
-        if width > 2048 or height > 2048:
-            raise HTTPException(
-                status_code=400,
-                detail="Image too large. Maximum dimensions: 2048x2048"
-            )
-        
         return image
         
     except HTTPException:
@@ -201,90 +167,113 @@ async def read_image_from_upload(file: UploadFile) -> Image.Image:
 async def root():
     """Root endpoint"""
     return {
-        "message": "StableVITON Virtual Try-On API",
-        "version": "1.0.0",
+        "message": "StableVITON Virtual Try-On API (Remote Version)",
+        "version": "1.2.0",
         "endpoints": {
             "/tryon": "POST - Virtual try-on inference",
+            "/demo": "GET - Web browser demo interface",
             "/health": "GET - Health check"
         }
     }
 
 
+@app.get("/demo")
+async def get_demo():
+    """Serve the static demo.html file"""
+    from fastapi.responses import FileResponse
+    demo_path = os.path.join(os.path.dirname(__file__), "demo.html")
+    if os.path.exists(demo_path):
+        return FileResponse(demo_path)
+    raise HTTPException(status_code=404, detail="demo.html not found")
+
+
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
     return {
         "status": "healthy",
-        "model_loaded": model is not None,
+        "model_connected": model is not None,
         "timestamp": time.time()
     }
 
 
 @app.post("/tryon", response_model=TryOnResponse)
 async def virtual_tryon(
+    request: Request,
     person_image: UploadFile = File(..., description="Full-body photo of person"),
-    garment_image: UploadFile = File(..., description="Garment image")
+    garment_image: UploadFile = File(..., description="Garment image"),
+    category: str = Form("tops"),
+    garment_photo_type: str = Form("model"),
+    num_timesteps: Optional[int] = Form(50),
+    guidance_scale: Optional[float] = Form(1.5),
+    seed: Optional[int] = Form(42),
+    segmentation_free: Optional[bool] = Form(True)
 ):
     """
-    Perform virtual try-on inference.
-    
-    Args:
-        person_image: Full-body photo of the person
-        garment_image: Image of the garment to try on
-        
-    Returns:
-        TryOnResponse with result image or error
+    Perform virtual try-on inference via remote API.
     """
     start_time = time.time()
     
     try:
-        # Check if model is loaded
+        # Check if model client is initialized
         if model is None:
             raise HTTPException(
                 status_code=503,
-                detail="Model not loaded. Please try again later."
+                detail="Remote API client not initialized. Please check server logs."
             )
         
-        # Validate files
+        # Step 1: Request Received
+        print("\n>>> Step 1: Request received at /tryon")
+        print(f"Headers: {request.headers}")
+        logger.info("Step 1: Request received and logging started")
+        
+        # Step 2: Validate Files
         validate_image(person_image)
         validate_image(garment_image)
+        print(">>> Step 2: Input images validated successfully")
+        logger.info("Step 2: Input images validated successfully")
         
-        logger.info(f"Processing try-on request: {person_image.filename}, {garment_image.filename}")
-        
-        # Read images
+        # Step 3: Read Images
         person_img = await read_image_from_upload(person_image)
         garment_img = await read_image_from_upload(garment_image)
+        print(">>> Step 3: Images loaded into memory")
+        logger.info("Step 3: Images loaded into memory")
         
-        # Acquire processing lock (single request at a time)
-        acquired = processing_lock.acquire(blocking=False)
-        if not acquired:
-            raise HTTPException(
-                status_code=503,
-                detail="Server is busy processing another request. Please try again in a moment."
-            )
+        # Step 4: Call Remote Inference
+        print(f">>> Step 4: Calling remote Gradio API (Category: {category})")
+        logger.info("Step 4: Calling remote Gradio API inference...")
         
-        try:
-            # Run inference
-            logger.info("Running inference...")
-            result_image = model.tryon(person_img, garment_img)
-            
-            # Convert to base64
-            result_base64 = image_to_base64(result_image)
-            
-            # Cleanup
-            model.cleanup()
-            
-            processing_time = time.time() - start_time
-            logger.info(f"Inference completed in {processing_time:.2f}s")
-            
-            return TryOnResponse(
-                success=True,
-                result_image=result_base64,
-                processing_time=processing_time
-            )
-            
-        finally:
-            processing_lock.release()
+        # Ensure correct types for the model wrapper
+        n_steps = int(num_timesteps) if num_timesteps is not None else 50
+        g_scale = float(guidance_scale) if guidance_scale is not None else 1.5
+        s_seed = int(seed) if seed is not None else 42
+        seg_free = bool(segmentation_free) if segmentation_free is not None else True
+
+        result_image = model.tryon(
+            person_image=person_img,
+            garment_image=garment_img,
+            category=category,
+            garment_photo_type=garment_photo_type,
+            num_timesteps=n_steps,
+            guidance_scale=g_scale,
+            seed=s_seed,
+            segmentation_free=seg_free
+        )
+        print(">>> Step 5: Remote inference completed successfully")
+        logger.info("Step 5: Remote inference completed")
+        
+        # Step 6: Convert to base64
+        result_base64 = image_to_base64(result_image)
+        print(">>> Step 6: Result processed and sending to frontend")
+        
+        processing_time = time.time() - start_time
+        logger.info(f"Remote inference completed in {processing_time:.2f}s")
+        
+        return TryOnResponse(
+            success=True,
+            result_image=result_base64,
+            processing_time=processing_time
+        )
         
     except HTTPException:
         raise