import gradio as gr
import torch
import os
import requests
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from PIL import Image
import numpy as np

# Ensure model weights exist
MODEL_URL = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
MODEL_PATH = 'weights/RealESRGAN_x4plus.pth'

os.makedirs('weights', exist_ok=True)
if not os.path.exists(MODEL_PATH):
    print('Downloading model...')
    response = requests.get(MODEL_URL, stream=True)
    with open(MODEL_PATH, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

# Load Real-ESRGAN model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
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=400, tile_pad=10, pre_pad=0, half=True)

# Image upscale function
def upscale_image(image):
    img = np.array(image)
    output, _ = upscaler.enhance(img, outscale=4)
    return image, Image.fromarray(output)

# Gradio UI
iface = gr.Interface(
    fn=upscale_image,
    inputs=gr.Image(type='pil'),
    outputs=[gr.Image(type='pil', label='Original Image'), gr.Image(type='pil', label='Upscaled Image')],
    title='Image Upscaler',
    description='Upload an image to see the before and after upscaling using Real-ESRGAN.'
)

if __name__ == '__main__':
    iface.launch()