import sys
import os
import torch
import gradio as gr
from fastapi import FastAPI, UploadFile, File
import uvicorn
from PIL import Image
import io

# Ensure the `basicsr` directory is in the system path
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(current_dir, "basicsr"))

# Import the inference function
from inference_codeformer import inference  

app = FastAPI()

# Load the CodeFormer model
model_path = "weights/CodeFormer.pth"
device = "cuda" if torch.cuda.is_available() else "cpu"

@app.post("/enhance")
async def enhance_image(file: UploadFile = File(...), upscale: int = 2, fidelity: float = 0.5):
    """API Endpoint to enhance images using CodeFormer"""
    image = Image.open(io.BytesIO(await file.read()))
    image.save("input.png")

    output_path = inference(
        input_path="input.png",
        upscale=upscale,
        fidelity=fidelity,
        model_path=model_path,
        device=device
    )

    return {"enhanced_image": f"https://your-space-name.hf.space/{output_path}"}

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