Update app.py

app.py

@@ -5,6 +5,19 @@ from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
 import gradio as gr
 import spaces
 
+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
+
 def split_image(image, chunk_size=512):
     width, height = image.size
     chunks = []
@@ -30,8 +43,11 @@ def upscale_chunk(chunk, model, processor, device):
     output_image = (output * 255.0).round().astype(np.uint8)
     return Image.fromarray(output_image)
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def main(image, model_choice, save_as_jpg=True):
+    # Resize the input image
+    image = resize_image(image)
+    
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     model_paths = {
@@ -87,7 +103,7 @@ interface = gr.Interface(
         gr.Textbox(label="Error Message", visible=True)
     ],
     title="Image Upscaler",
-    description="Upload an image, select a model, and upscale it. The image will be processed in 512x512 pixel chunks to handle large images efficiently.",
+    description="Upload an image, select a model, and upscale it. Images larger than 2048x2048 will be resized while maintaining aspect ratio. The image will be processed in 512x512 pixel chunks for efficient handling.",
 )
 
 interface.launch()