main.py

"""
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"
    )