import gradio as gr
import torch
from PIL import Image
import numpy as np
import os
import requests

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

# Download model if not present
model_path = 'RealESRGAN_x4plus.pth'
model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
if not os.path.exists(model_path):
    print("🔽 Downloading model...")
    r = requests.get(model_url)
    with open(model_path, 'wb') as f:
        f.write(r.content)
    print("✅ Model downloaded!")

# Check if GPU is available
use_half = torch.cuda.is_available()

# Load model
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
                num_block=23, num_grow_ch=32, scale=4)

upscaler = RealESRGANer(
    scale=4,
    model_path=model_path,
    model=model,
    tile=256,         # speeds up processing for large images
    tile_pad=10,
    pre_pad=0,
    half=use_half
)

def upscale_image(input_img, dpi):
    img = np.array(input_img)
    try:
        output, _ = upscaler.enhance(img, outscale=1)
        output_pil = Image.fromarray(output)

        from io import BytesIO
        buffer = BytesIO()
        output_pil.save(buffer, format="JPEG", dpi=(dpi, dpi))
        buffer.seek(0)
        return Image.open(buffer)
    except Exception as e:
        return f"❌ Error during upscaling: {str(e)}"

iface = gr.Interface(
    fn=upscale_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(72, 300, value=100, step=1, label="Set DPI (for stock upload)")
    ],
    outputs=gr.Image(type="pil", label="Upscaled Image"),
    title="🔍 Free AI Image Upscaler (Real-ESRGAN 4x)",
    description="Fast and free image upscaler using Real-ESRGAN with DPI control (Adobe Stock compatible). Best performance on GPU."
)

iface.launch()