import torch
import gradio as gr
from inference_codeformer import inference
from fastapi import FastAPI, UploadFile, File
import uvicorn
from PIL import Image
import io

app = FastAPI()

# xyz
# vishal singh
# 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):
    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)