""" FastAPI Backend for StableVITON Virtual Try-On Provides REST API endpoint for virtual try-on inference """ 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 from PIL import Image import io import base64 import os import time import asyncio from typing import Optional import logging from queue import Queue from threading import Lock from inference_wrapper import StableVITONInference # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) 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 API inference engine model = StableVITONInference( model_path=os.getenv("MODEL_PATH", "merve/fashn-vton-1.5"), hf_token=os.getenv("HF_TOKEN") ) 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", lifespan=lifespan ) # CORS configuration - allow all origins for demo (restrict in production) app.add_middleware( CORSMiddleware, allow_origins=["*"], # Change to specific domains in production allow_credentials=True, allow_methods=["*"], 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 # Request queue for single-request processing request_queue = Queue() processing_lock = Lock() # Configuration MAX_IMAGE_SIZE = 10 * 1024 * 1024 # 10MB ALLOWED_EXTENSIONS = {"image/jpeg", "image/png", "image/jpg"} REQUEST_TIMEOUT = 90 # seconds class TryOnResponse(BaseModel): """Response model for try-on endpoint""" success: bool result_image: Optional[str] = None processing_time: Optional[float] = None model_version: str = "stablevton-v1" error: Optional[str] = None error_code: Optional[str] = None # Note: Startup and shutdown are now handled by lifespan context manager def validate_image(file: UploadFile) -> None: """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}" ) if hasattr(file, 'size') and file.size > MAX_IMAGE_SIZE: raise HTTPException( status_code=400, detail=f"File too large. Maximum size: {MAX_IMAGE_SIZE / (1024*1024)}MB" ) def image_to_base64(image: Image.Image) -> str: """Convert PIL Image to base64 string.""" buffered = io.BytesIO() image.save(buffered, format="PNG") img_bytes = buffered.getvalue() img_base64 = base64.b64encode(img_bytes).decode('utf-8') return f"data:image/png;base64,{img_base64}" async def read_image_from_upload(file: UploadFile) -> Image.Image: """Read PIL Image from uploaded file.""" try: contents = await file.read() if len(contents) > MAX_IMAGE_SIZE: raise HTTPException( status_code=400, detail=f"File too large. Maximum size: {MAX_IMAGE_SIZE / (1024*1024)}MB" ) image = Image.open(io.BytesIO(contents)) return image except HTTPException: raise except Exception as e: logger.error(f"Failed to read image: {e}") raise HTTPException( status_code=400, detail=f"Invalid image file: {str(e)}" ) @app.get("/") async def root(): """Root endpoint""" return { "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_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"), 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 via remote API. """ start_time = time.time() try: # Check if model client is initialized if model is None: raise HTTPException( status_code=503, detail="Remote API client not initialized. Please check server logs." ) # 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") # 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") # Step 4: Call Remote Inference print(f">>> Step 4: Calling remote Gradio API (Category: {category})") logger.info("Step 4: Calling remote Gradio API inference...") # 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 except Exception as e: logger.error(f"Inference failed: {e}", exc_info=True) return TryOnResponse( success=False, error=str(e), error_code="INFERENCE_FAILED" ) if __name__ == "__main__": import uvicorn uvicorn.run( "main:app", host="0.0.0.0", port=7860, reload=False, # Set to True for development log_level="info" )