import io
import uvicorn
import numpy as np
from PIL import Image
import cv2
import torch
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import StreamingResponse
from realesrgan import RealESRGAN
import gradio as gr
import threading

# -----------------
# Load Model
# -----------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RealESRGAN(device, scale=4)
model.load_weights('https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5/RealESRGAN_x4plus.pth')

# -----------------
# FastAPI (Custom API)
# -----------------
api = FastAPI(title="Real-ESRGAN Custom API")

@api.post("/upscale")
async def upscale_image(file: UploadFile = File(...)):
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    image_np = np.array(image)

    upscaled = model.predict(image_np)
    upscaled_img = Image.fromarray(upscaled)

    buf = io.BytesIO()
    upscaled_img.save(buf, format="PNG")
    buf.seek(0)

    return StreamingResponse(buf, media_type="image/png")

# -----------------
# Gradio UI
# -----------------
def gradio_upscale(img):
    img_np = np.array(img)
    result = model.predict(img_np)
    return Image.fromarray(result)

ui = gr.Interface(
    fn=gradio_upscale,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(type="pil", label="Upscaled Image (4x)"),
    title="Real-ESRGAN Image Upscaler",
    description="4x AI Image Upscaling using Real-ESRGAN"
)

# -----------------
# Run both API + UI
# -----------------
def start_api():
    uvicorn.run(api, host="0.0.0.0", port=7861)

threading.Thread(target=start_api).start()

ui.launch(server_name="0.0.0.0", server_port=7860)