from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import torch
from PIL import Image
import io
import base64
from diffusers import StableDiffusionInpaintPipeline
import gc
from fastapi.responses import JSONResponse
import logging

app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variable for the model
pipe = None

def load_model():
    global pipe
    if pipe is None:
        # Use the pre-uploaded model from Hugging Face
        model_id = "Uminosachi/realisticVisionV51_v51VAE-inpainting"
        pipe = StableDiffusionInpaintPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            safety_checker=None
        ).to("cuda")
        pipe.enable_attention_slicing(slice_size="max")
        pipe.enable_sequential_cpu_offload()
    return pipe

@app.on_event("startup")
async def startup_event():
    if torch.cuda.is_available():
        load_model()

def image_to_base64(image: Image.Image) -> str:
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

@app.post("/inpaint")
async def inpaint(
    image: UploadFile = File(...),
    mask: UploadFile = File(...),
    prompt: str = "add some flowers and a fountain",
    negative_prompt: str = "blurry, low quality, distorted"
):
    try:
        # Add file size check (10MB limit)
        max_size = 10 * 1024 * 1024  # 10MB
        if len(await image.read()) > max_size or len(await mask.read()) > max_size:
            return JSONResponse(
                status_code=400,
                content={"error": "File size too large. Maximum size is 10MB"}
            )
        
        # Reset file positions
        await image.seek(0)
        await mask.seek(0)
        
        # Read and process input image
        image_data = await image.read()
        mask_data = await mask.read()
        
        original_image = Image.open(io.BytesIO(image_data))
        mask_image = Image.open(io.BytesIO(mask_data))

        # Resize to multiple of 8
        width, height = (dim - dim % 8 for dim in original_image.size)
        original_image = original_image.resize((width, height))
        mask_image = mask_image.resize((width, height))
        mask_image = mask_image.convert("L")

        # Perform inpainting
        with torch.cuda.amp.autocast():
            output_image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image=original_image,
                mask_image=mask_image,
                num_inference_steps=20,
                guidance_scale=7.5,
            ).images[0]

        # Convert output image to base64
        output_base64 = image_to_base64(output_image)

        # Clean up
        torch.cuda.empty_cache()
        gc.collect()

        return {"status": "success", "image": output_base64}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "healthy", "cuda_available": torch.cuda.is_available()}