Update app.py

app.py

@@ -6,6 +6,20 @@ from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
 import gradio as gr
 import spaces
 
+# Function to resize image to max 2048x2048 while maintaining aspect ratio
+def resize_image(image, max_size=2048):
+    width, height = image.size
+    if width > max_size or height > max_size:
+        aspect_ratio = width / height
+        if width > height:
+            new_width = max_size
+            new_height = int(new_width / aspect_ratio)
+        else:
+            new_height = max_size
+            new_width = int(new_height * aspect_ratio)
+        image = image.resize((new_width, new_height), Image.LANCZOS)
+    return image
+
 # Function to upscale an image using Swin2SR
 def upscale_image(image, model, processor, device):
     # Convert the image to RGB format
@@ -30,6 +44,9 @@ def main(image, model_choice, save_as_jpg=True):
     # Check if GPU is available and set the device accordingly
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
+    # Resize the input image
+    image = resize_image(image)
+    
     # Define model paths
     model_paths = {
         "Pixel Perfect": "caidas/swin2SR-classical-sr-x4-64",
@@ -72,7 +89,7 @@ interface = gr.Interface(
     ],
     outputs=gr.File(label="Download Upscaled Image"),
     title="Image Upscaler",
-    description="Upload an image, select a model, upscale it, and download the new image.",
+    description="Upload an image, select a model, upscale it, and download the new image. Images larger than 2048x2048 will be resized while maintaining aspect ratio.",
 )
 
 # Launch the interface