import gradio as gr
from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
from PIL import Image
import torch

# Load model and processor
processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-lightweight-x4-64")
model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-lightweight-x4-64")

def upscale_image(image):
    # Preprocess the input image
    inputs = processor(images=image, return_tensors="pt")
    
    # Perform super-resolution
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Post-process the output to get a high-resolution image
    output_image = processor.postprocess(outputs.logits, target_sizes=[(image.size[1]*4, image.size[0]*4)])[0]
    return output_image

# Gradio interface
iface = gr.Interface(
    fn=upscale_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Image Super Resolution with Swin2SR",
    description="Upload an image and enhance its resolution using the Swin2SR model (4x resolution)."
)

if __name__ == "__main__":
    iface.launch()