import io
import torch

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from PIL import Image
from RealESRGAN import RealESRGAN

app = FastAPI()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RealESRGAN(device, scale=2)
model.load_weights("weights/RealESRGAN_x2.pth", download=True)


@app.get("/")
async def root():
    return {"message": "API is running"}


@app.post("/enhance-img")
async def enhance_img(file: UploadFile = File(...)):
    try:
        print('Inside the function')
        data = await file.read()
        print('File read')
        img = Image.open(io.BytesIO(data)).convert("RGB")
        print('Passing file into network')
        sr = model.predict(img)
        print("File output")
        buf = io.BytesIO()
        sr.save(buf, format="PNG")
        buf.seek(0)
        print('Done')
        return StreamingResponse(buf, media_type="image/png")

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