Update app.py

app.py

@@ -1,32 +1,43 @@
 import gradio as gr
 from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
 from PIL import Image
+import numpy as np
 import torch
+import requests
 
-# Load model and processor
-processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-lightweight-x4-64")
-model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-lightweight-x4-64")
+# Load the model and processor
+processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr")
+model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr").to('cuda')
 
-def upscale_image(image):
+# Define the function for super-resolution
+def super_resolve(image):
     # Preprocess the input image
-    inputs = processor(images=image, return_tensors="pt")
-    
+    inputs = processor(images=image, return_tensors="pt").to('cuda')
+
     # 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]
+    # Get the reconstructed tensor from the outputs
+    reconstructed_tensor = outputs.reconstruction
+
+    # Move the tensor to the CPU and convert it to a NumPy array
+    image_tensor = reconstructed_tensor.squeeze().cpu()
+    image_np = image_tensor.permute(1, 2, 0).numpy()  # Permute to make it HxWxC
+
+    # Rescale the values from [0, 1] to [0, 255]
+    image_np = np.clip(image_np, 0, 1)
+    image_np = (image_np * 255).astype(np.uint8)
+
+    # Convert the NumPy array back to an image
+    output_image = Image.fromarray(image_np)
+
     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()
+# Create the Gradio interface
+inputs = gr.inputs.Image(type="pil", label="Upload an Image")
+outputs = gr.outputs.Image(type="pil", label="Super-Resolved Image")
+
+gr.Interface(fn=super_resolve, inputs=inputs, outputs=outputs, title="Image Super-Resolution with Swin2SR",
+             description="Upload an image to upscale it using the Swin2SR model for real-world super-resolution."
+).launch()