import gradio as gr
from transformers import pipeline
from PIL import Image
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipeline(task="image-to-image", model="caidas/swin2SR-lightweight-x2-64", device=device)

def upscale_image(image):
    upscaled = pipe(image)
    return upscaled

iface = gr.Interface(
    fn=upscale_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Image Super Resolution",
    description="Upscale an image using Swin2SR."
)

iface.launch()