Update app.py

app.py

@@ -6,7 +6,8 @@ import matplotlib.pyplot as plt
 import subprocess
 import spaces
 import torch
-import io
+import tempfile
+import os
 
 # Run the script to get pretrained models
 subprocess.run(["bash", "get_pretrained_models.sh"])
@@ -24,21 +25,20 @@ def resize_image(image_path, max_size=1024):
         # 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
+        # Create a temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
+            img.save(temp_file, format="PNG")
+            return temp_file.name
 
 @spaces.GPU(duration=120)
 def predict_depth(input_image):
+    temp_file = None
     try:
         # Resize the input image
-        resized_image = resize_image(input_image)
+        temp_file = resize_image(input_image)
         
         # Preprocess the image
-        result = depth_pro.load_rgb(resized_image)
+        result = depth_pro.load_rgb(temp_file)
         image = result[0]
         f_px = result[-1]  # Assuming f_px is the last item in the returned tuple
         image = transform(image)
@@ -76,6 +76,10 @@ def predict_depth(input_image):
         return output_path, f"Focal length: {focallength_px:.2f} pixels"
     except Exception as e:
         return None, f"An error occurred: {str(e)}"
+    finally:
+        # Clean up the temporary file
+        if temp_file and os.path.exists(temp_file):
+            os.remove(temp_file)
 
 # Create Gradio interface
 iface = gr.Interface(