app.py

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

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

# Initialize model for CPU
model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32)

upsampler = RealESRGANer(
    scale=4,
    model_path='RealESRGAN_x4plus.pth',
    model=model,
    tile=64,
    tile_pad=10,
    pre_pad=0,
    half=False
)

MAX_RES = 1000

def enhance_image(input_image):
    img = input_image.convert("RGB")
    if max(img.size) > MAX_RES:
        img.thumbnail((MAX_RES, MAX_RES))
    
    img_np = np.array(img)

    try:
        output, _ = upsampler.enhance(img_np)
        enhanced = Image.fromarray(output.astype(np.uint8))
        # Save to file so it can be downloadable
        enhanced.save("enhanced.png", format="PNG")
        return "enhanced.png"
    except Exception as e:
        print("❌ Enhancement error:", e)
        return None

# ✅ Single output: the image with built-in download button
demo = gr.Interface(
    fn=enhance_image,
    inputs=gr.Image(type="pil", label="Upload your image"),
    outputs=gr.Image(type="filepath", label="Enhanced Image (click to download)", image_mode="RGB", show_download_button=True),
    title="🖼️ Real-ESRGAN CPU Enhancer",
    description="Upload an image to upscale and enhance it using Real-ESRGAN. Optimized for CPU (tile=64)."
)

demo.launch()