import os
import gradio as gr
import torch
from PIL import Image
from torchvision.transforms import ToTensor, ToPILImage
from torchvision import models

# Load the R-ESRGAN Anime model (you need to set up this model properly)
model = torch.hub.load('xinntao/Real-ESRGAN', 'restorer', model='R-ESRGAN_Anime_X6')

def upscale_image(image, scale_factor):
    # Ensure image is in PIL format
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    # Upscale using the model
    with torch.no_grad():
        upscaled_image = model(image, scale=scale_factor)

    return upscaled_image

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

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()