import gradio as gr
import torch
from realesrgan import RealESRGAN

# Load the model
model = RealESRGAN('weights/RealESRGAN_x4plus_anime_6B.pth', scale=4)  # Adjust scale if needed

def upscale_image(input_image, scale_factor):
    # Ensure scale factor is an integer
    scale_factor = int(scale_factor)
    
    # Set the model scale
    model.scale = scale_factor
    
    # Perform the upscaling
    upscaled_image = model.predict(input_image)
    
    return upscaled_image

# Create the Gradio interface
iface = gr.Interface(
    fn=upscale_image,
    inputs=[
        gr.inputs.Image(type="numpy", label="Input Image"),
        gr.inputs.Slider(minimum=2, maximum=4, step=1, default=4, label="Scale Factor")
    ],
    outputs=gr.outputs.Image(type="numpy", label="Upscaled Image"),
    title="Image Upscaler using R-ESRGAN Anime 6B",
    description="Upload an image and select a scale factor to upscale it using the R-ESRGAN Anime 6B model."
)

# Launch the app
iface.launch()