Update app.py

app.py

@@ -6,6 +6,7 @@ import matplotlib.pyplot as plt
 import subprocess
 import spaces
 import torch
+import io
 
 # Run the script to get pretrained models
 subprocess.run(["bash", "get_pretrained_models.sh"])
@@ -14,11 +15,30 @@ subprocess.run(["bash", "get_pretrained_models.sh"])
 model, transform = depth_pro.create_model_and_transforms()
 model.eval()
 
+def resize_image(image_path, max_size=1024):
+    with Image.open(image_path) as img:
+        # Calculate the new size while maintaining aspect ratio
+        ratio = max_size / max(img.size)
+        new_size = tuple([int(x * ratio) for x in img.size])
+        
+        # Resize the image
+        img = img.resize(new_size, Image.LANCZOS)
+        
+        # Save the resized image to a bytes buffer
+        buffer = io.BytesIO()
+        img.save(buffer, format="PNG")
+        buffer.seek(0)
+        
+        return buffer
+
 @spaces.GPU(duration=120)
 def predict_depth(input_image):
     try:
+        # Resize the input image
+        resized_image = resize_image(input_image)
+        
         # Preprocess the image
-        result = depth_pro.load_rgb(input_image)
+        result = depth_pro.load_rgb(resized_image)
         image = result[0]
         f_px = result[-1]  # Assuming f_px is the last item in the returned tuple
         image = transform(image)
@@ -63,7 +83,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="filepath"),
     outputs=[gr.Image(type="filepath", label="Depth Map"), gr.Textbox(label="Focal Length or Error Message")],
     title="Depth Prediction Demo",
-    description="Upload an image to predict its depth map and focal length."
+    description="Upload an image to predict its depth map and focal length. Large images will be automatically resized."
 )
 
 # Launch the interface