import gradio as gr
import requests
from PIL import Image
import os
import torch
import numpy as np
from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution


processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64")

def enhance(image):
    # prepare image for the model
    inputs = processor(image, return_tensors="pt")

    # forward pass
    with torch.no_grad():
        outputs = model(**inputs)

    # postprocess
    output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
    output = np.moveaxis(output, source=0, destination=-1)
    output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
    
    return Image.fromarray(output)

title = "Mojo Solo Image Super-Resolution" 

gr.Interface(
    enhance, 
    gr.inputs.Image(type="pil", label="Input").style(height=260),
    gr.inputs.Image(type="pil", label="Ouput").style(height=240),
    title=title,
    ).launch(enable_queue=True, share= True)