api_server.py

"""
FastAPI server for Virtual Try-On using multiple model providers.

This server provides a simple REST API endpoint for virtual try-on generation
using various image generation models:
- Nano Banana and Nano Banana Pro (Google Gemini)
- FLUX 2 Pro and FLUX 2 Flex (Black Forest Labs)
"""

import os
import io
import base64
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Tuple
from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Import adapters
from tryon.api.nano_banana import NanoBananaAdapter, NanoBananaProAdapter
from tryon.api.flux2 import Flux2ProAdapter, Flux2FlexAdapter

# Create output directory for generated images
OUTPUT_DIR = Path("outputs/virtual_tryon")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Supported aspect ratios for both adapters
SUPPORTED_ASPECT_RATIOS = [
    "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"
]


def calculate_aspect_ratio(image: Image.Image) -> str:
    """
    Calculate the aspect ratio from an image and return the closest supported ratio.
    
    Args:
        image: PIL Image object
        
    Returns:
        str: Aspect ratio string in format "W:H" (e.g., "16:9")
    """
    width, height = image.size
    ratio = width / height
    
    # Map of supported ratios to their decimal values
    ratio_map = {
        "1:1": 1.0,
        "2:3": 2/3,
        "3:2": 3/2,
        "3:4": 3/4,
        "4:3": 4/3,
        "4:5": 4/5,
        "5:4": 5/4,
        "9:16": 9/16,
        "16:9": 16/9,
        "21:9": 21/9,
    }
    
    # Find the closest matching aspect ratio
    closest_ratio = "1:1"  # Default
    min_diff = float('inf')
    
    for ratio_str, ratio_value in ratio_map.items():
        diff = abs(ratio - ratio_value)
        if diff < min_diff:
            min_diff = diff
            closest_ratio = ratio_str
    
    return closest_ratio


def get_image_dimensions(image: Image.Image) -> Tuple[int, int]:
    """
    Get image dimensions (width, height).
    
    Args:
        image: PIL Image object
        
    Returns:
        tuple: (width, height)
    """
    return image.size


def calculate_resolution(image: Image.Image) -> str:
    """
    Calculate resolution from image dimensions in "widthxheight" format.
    
    Args:
        image: PIL Image object
        
    Returns:
        str: Resolution string in format "widthxheight" (e.g., "1024x1024")
    """
    width, height = image.size
    return f"{width}x{height}"


def map_resolution_to_pro_format(image: Image.Image) -> str:
    """
    Map image resolution to Nano Banana Pro format ("1K", "2K", or "4K").
    
    The mapping is based on the maximum dimension:
    - max_dimension <= 1500: "1K"
    - max_dimension <= 3000: "2K"
    - max_dimension > 3000: "4K"
    
    Args:
        image: PIL Image object
        
    Returns:
        str: Resolution in format "1K", "2K", or "4K"
    """
    width, height = image.size
    max_dimension = max(width, height)
    
    if max_dimension <= 1500:
        return "1K"
    elif max_dimension <= 3000:
        return "2K"
    else:
        return "4K"

app = FastAPI(
    title="TryOn AI Virtual Try-On API",
    description="Virtual try-on API using multiple model providers (Nano Banana, FLUX 2 Pro, FLUX 2 Flex)",
    version="1.0.0"
)

# CORS middleware to allow requests from frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000", 
        "http://127.0.0.1:3000",
        "http://localhost:5173",
        "https://fyp-frontend-sandy.vercel.app",
        "*"  # Allow all origins for Hugging Face deployment (you can restrict this later)
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def root():
    """Root endpoint."""
    return {
        "message": "TryOn AI Virtual Try-On API",
        "version": "1.0.0",
        "endpoints": {
            "POST /api/v1/virtual-tryon": "Generate virtual try-on image"
        },
        "providers": [
            "nano-banana",
            "nano-banana-pro",
            "flux-2-pro",
            "flux-2-flex"
        ]
    }


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {"status": "healthy"}


@app.post("/api/v1/virtual-tryon")
async def virtual_tryon(
    model_image: UploadFile = File(..., description="Model/person image"),
    garment_images: List[UploadFile] = File(..., description="Garment images"),
    provider: str = Form(default="nano-banana", description="Provider: 'nano-banana', 'nano-banana-pro', 'flux-2-pro', or 'flux-2-flex'"),
    prompt: Optional[str] = Form(default=None, description="Optional custom prompt"),
    resolution: Optional[str] = Form(default="1K", description="Resolution for nano-banana-pro: '1K', '2K', or '4K'"),
    aspect_ratio: Optional[str] = Form(default=None, description="Optional aspect ratio (e.g., '16:9')"),
    width: Optional[int] = Form(default=None, description="Output image width (for FLUX 2 models)"),
    height: Optional[int] = Form(default=None, description="Output image height (for FLUX 2 models)"),
    seed: Optional[int] = Form(default=None, description="Random seed for reproducibility (for FLUX 2 models)"),
    guidance: Optional[float] = Form(default=None, description="Guidance scale 1.5-10 (for FLUX 2 Flex, default: 3.5)"),
    steps: Optional[int] = Form(default=None, description="Number of generation steps (for FLUX 2 Flex, default: 28)"),
    safety_tolerance: Optional[int] = Form(default=2, description="Safety tolerance 0-5 (for FLUX 2 models, default: 2)")
):
    """
    Generate virtual try-on image from model image and garment images.
    
    Uses multi-image composition feature of various models to combine
    the model image with multiple garment images.
    
    Supported providers:
    - nano-banana: Google Gemini Nano Banana (basic)
    - nano-banana-pro: Google Gemini Nano Banana Pro (supports resolution)
    - flux-2-pro: Black Forest Labs FLUX 2 Pro (high quality)
    - flux-2-flex: Black Forest Labs FLUX 2 Flex (advanced controls)
    
    Args:
        model_image: Single model/person image
        garment_images: List of garment images (top, jeans, scarf, hat, etc.)
        provider: Model provider ('nano-banana', 'nano-banana-pro', 'flux-2-pro', or 'flux-2-flex')
        prompt: Optional custom prompt for generation
        resolution: Resolution for nano-banana-pro ('1K', '2K', or '4K')
        aspect_ratio: Optional aspect ratio (for Nano Banana models)
        width: Output image width in pixels (for FLUX 2 models, minimum: 64)
        height: Output image height in pixels (for FLUX 2 models, minimum: 64)
        seed: Random seed for reproducibility (for FLUX 2 models)
        guidance: Guidance scale 1.5-10 (for FLUX 2 Flex only, default: 3.5)
        steps: Number of generation steps (for FLUX 2 Flex only, default: 28)
        safety_tolerance: Safety tolerance 0-5 (for FLUX 2 models, default: 2)
        
    Returns:
        JSON response with base64-encoded result image
    """
    try:
        # Validate provider
        valid_providers = ["nano-banana", "nano-banana-pro", "flux-2-pro", "flux-2-flex"]
        if provider not in valid_providers:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid provider '{provider}'. Must be one of: {', '.join(valid_providers)}"
            )
    
        # Validate inputs
        if not model_image:
            raise HTTPException(status_code=400, detail="Model image is required")
    
        if not garment_images or len(garment_images) == 0:
            raise HTTPException(status_code=400, detail="At least one garment image is required")
    
        # Read model image
        model_image_bytes = await model_image.read()
        model_pil = Image.open(io.BytesIO(model_image_bytes))
    
        # Calculate aspect ratio and resolution from model image
        calculated_aspect_ratio = calculate_aspect_ratio(model_pil)
        calculated_resolution = calculate_resolution(model_pil)
        model_width, model_height = get_image_dimensions(model_pil)
    
        # Use calculated aspect ratio if not provided, otherwise use the provided one
        final_aspect_ratio = aspect_ratio if aspect_ratio else calculated_aspect_ratio
    
        # Map resolution to appropriate format based on provider
        # For nano-banana-pro, use "1K", "2K", or "4K" format
        # For nano-banana, resolution is not used (only aspect ratio)
        # For FLUX 2 models, use width/height parameters instead
        if provider == "nano-banana-pro":
            # Use provided resolution if valid, otherwise map from image dimensions
            if resolution and resolution in ["1K", "2K", "4K"]:
                final_resolution = resolution
            else:
                final_resolution = map_resolution_to_pro_format(model_pil)
        elif provider in ["flux-2-pro", "flux-2-flex"]:
            # For FLUX 2, use model dimensions as default width/height if not provided
            if width is None:
                width = model_width
            if height is None:
                height = model_height
            final_resolution = f"{width}x{height}"
        else:
            # For nano-banana, resolution is not used, but keep calculated for reference
            final_resolution = calculated_resolution
    
        # Read garment images and combine with model image
        # First image should be model, followed by garments
        images_list = [model_pil]
        for garment_file in garment_images:
            garment_bytes = await garment_file.read()
            garment_pil = Image.open(io.BytesIO(garment_bytes))
            images_list.append(garment_pil)
    
        # Prepare prompt
        if not prompt:
            prompt = (
                "Create a realistic virtual try-on image showing the person wearing the provided garments. "
                "CRITICAL REQUIREMENTS - Preserve all details exactly:\n"
                "1. GARMENT EXTRACTION: The garment images may contain people wearing the garments. "
                "IGNORE and EXTRACT ONLY the garment itself - do not use any person, model, or human figure "
                "from the garment images. Focus solely on the garment: its shape, design, patterns, colors, "
                "textures, and all visual details. Remove or ignore any human elements from garment images.\n"
                "2. GARMENT PRESERVATION: Keep ALL garment details completely intact - patterns, colors, textures, "
                "designs, prints, logos, text, embroidery, sequins, and any decorative elements must remain "
                "identical to the original garment images. Do not alter, fade, or modify any garment features.\n"
                "3. PERSON PRESERVATION: Keep the person's face, body shape, skin tone, hair, and physical "
                "characteristics exactly as shown in the FIRST image (model image). Only apply the extracted "
                "garments from the subsequent images to this person. Do not use any person from garment images.\n"
                "4. PARTIAL GARMENT HANDLING: If the person in the model image is wearing a full-body outfit "
                "(dress, jumpsuit, etc.) but the provided garment is only upper-body (top, shirt, blouse) or "
                "lower-body (pants, jeans, skirt), place the provided garment correctly over the corresponding "
                "body part. For the remaining uncovered body parts, generate an appropriate complementary garment "
                "that matches: (a) the person's physical characteristics and body type, (b) the person's style "
                "and personality traits visible in the model image, (c) the style, color scheme, and design "
                "aesthetic of the provided garment. The complementary garment should create a cohesive, "
                "harmonious outfit that looks natural and well-coordinated.\n"
                "5. FITTING: The extracted garments should fit naturally on the person's body from the first image, "
                "following their body contours and proportions realistically, while maintaining all original "
                "garment details from the garment images.\n"
                "6. COMPOSITION: The first image is the model/person to dress. The following images contain "
                "garments (top, bottom, accessories, etc.) - extract ONLY the garments from these images, "
                "ignoring any people shown. Combine the extracted garments to create a cohesive outfit where "
                "each garment maintains its original appearance and fits the person naturally.\n"
                "7. REALISM: The final image should look like a professional photograph of the person from the "
                "first image wearing the exact extracted garments (and complementary garments if needed), with "
                "realistic lighting, shadows, and fabric draping."
            )
    
        # Initialize adapter and generate
        if provider == "nano-banana":
            adapter = NanoBananaAdapter()
            # Generate with basic adapter using calculated aspect ratio
            result_images = adapter.generate_multi_image(
                images=images_list,
                prompt=prompt,
                aspect_ratio=final_aspect_ratio
            )
        elif provider == "nano-banana-pro":
            adapter = NanoBananaProAdapter()
            # Generate with Pro adapter (supports resolution) using calculated resolution and aspect ratio
            result_images = adapter.generate_multi_image(
                images=images_list,
                prompt=prompt,
                resolution=final_resolution,
                aspect_ratio=final_aspect_ratio
            )
        elif provider == "flux-2-pro":
            adapter = Flux2ProAdapter()
            # Generate with FLUX 2 Pro adapter
            # Build kwargs for FLUX 2 Pro
            flux_kwargs = {
                "safety_tolerance": safety_tolerance if safety_tolerance is not None else 2,
                "output_format": "png"
            }
            if width is not None:
                flux_kwargs["width"] = width
            if height is not None:
                flux_kwargs["height"] = height
            if seed is not None:
                flux_kwargs["seed"] = seed
        
            result_images = adapter.generate_multi_image(
                prompt=prompt,
                images=images_list,
                **flux_kwargs
            )
        else:  # flux-2-flex
            adapter = Flux2FlexAdapter()
            # Generate with FLUX 2 Flex adapter (supports guidance and steps)
            # Build kwargs for FLUX 2 Flex
            flux_kwargs = {
                "safety_tolerance": safety_tolerance if safety_tolerance is not None else 2,
                "output_format": "png",
                "guidance": guidance if guidance is not None else 5,
                "steps": steps if steps is not None else 50,
                "prompt_upsampling": True
            }
            if width is not None:
                flux_kwargs["width"] = width
            if height is not None:
                flux_kwargs["height"] = height
            if seed is not None:
                flux_kwargs["seed"] = seed
        
            result_images = adapter.generate_multi_image(
                prompt=prompt,
                images=images_list,
                **flux_kwargs
            )
    
        if not result_images:
            raise HTTPException(status_code=500, detail="No images generated")
    
        # Get first result image and convert to PIL Image
        result_image = result_images[0]
    
        # Convert image to PIL Image if needed
        # FLUX 2 adapters return PIL Images directly
        # Nano Banana adapters return Google GenAI image types that need conversion
        if not isinstance(result_image, Image.Image):
            # Google GenAI image type has image_bytes attribute
            if hasattr(result_image, 'image_bytes'):
                # Convert bytes to PIL Image
                result_image = Image.open(io.BytesIO(result_image.image_bytes))
            elif hasattr(result_image, 'to_pil'):
                # If it has a to_pil method, use it
                result_image = result_image.to_pil()
            else:
                # Try to get bytes from the image object
                try:
                    # Some GenAI image types expose bytes directly
                    image_bytes = bytes(result_image)
                    result_image = Image.open(io.BytesIO(image_bytes))
                except (TypeError, AttributeError):
                    raise HTTPException(
                        status_code=500,
                        detail=f"Unable to convert image type {type(result_image)} to PIL Image. "
                               f"Image attributes: {dir(result_image)}"
                    )
    
        # Generate filename with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = f"tryon_{provider}_{timestamp}.png"
        filepath = OUTPUT_DIR / filename
    
        # Save image to disk
        try:
            # Ensure image is in RGB mode for saving
            if result_image.mode != 'RGB':
                result_image = result_image.convert('RGB')
        
            # Save to file
            result_image.save(str(filepath), 'PNG')
        
            # Also save to BytesIO for base64 encoding
            img_buffer = io.BytesIO()
            result_image.save(img_buffer, 'PNG')
            img_buffer.seek(0)
            img_base64 = base64.b64encode(img_buffer.read()).decode('utf-8')
        
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Error saving image: {str(e)}"
            )
    
        # Build response with provider-specific metadata
        response_data = {
            "success": True,
            "image": f"data:image/png;base64,{img_base64}",
            "provider": provider,
            "num_garments": len(garment_images),
            "saved_path": str(filepath),
            "filename": filename,
            "model_dimensions": {"width": model_width, "height": model_height},
        }
    
        # Add provider-specific metadata
        if provider in ["nano-banana", "nano-banana-pro"]:
            response_data.update({
                "aspect_ratio": final_aspect_ratio,
                "calculated_aspect_ratio": calculated_aspect_ratio,
                "resolution": final_resolution,
                "calculated_resolution": calculated_resolution
            })
        elif provider in ["flux-2-pro", "flux-2-flex"]:
            response_data.update({
                "output_dimensions": {"width": width or model_width, "height": height or model_height},
                "safety_tolerance": safety_tolerance if safety_tolerance is not None else 2,
            })
            if seed is not None:
                response_data["seed"] = seed
            if provider == "flux-2-flex":
                response_data.update({
                    "guidance": guidance if guidance is not None else 3.5,
                    "steps": steps if steps is not None else 28,
                })
    
        return JSONResponse(response_data)
    
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        import traceback
        error_details = f"Error generating try-on: {str(e)}\n{traceback.format_exc()}"
        print(error_details)  # Log to console
        raise HTTPException(status_code=500, detail=f"Error generating try-on: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)