from fastapi import FastAPI, File, UploadFile
from fastapi.responses import Response
from io import BytesIO
from PIL import Image
import torch
import uvicorn
import os

# Import the CodeFormer model processing function
from codeformer_model import enhance_image  # Make sure this function is defined

app = FastAPI()

@app.post("/enhance")
async def enhance_image_api(file: UploadFile = File(...)):
    try:
        # Load image
        image = Image.open(file.file).convert("RGB")

        # Process the image using the CodeFormer model
        enhanced_image = enhance_image(image)

        # Convert the processed image to bytes
        img_byte_arr = BytesIO()
        enhanced_image.save(img_byte_arr, format="PNG")
        img_byte_arr = img_byte_arr.getvalue()

        return Response(content=img_byte_arr, media_type="image/png")

    except Exception as e:
        return {"error": str(e)}

# Required to run on Hugging Face Spaces
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))