distill-any-depth

app.py

@@ -8,7 +8,8 @@ from distillanydepth.midas.transforms import Resize, NormalizeImage, PrepareForN
 from torchvision.transforms import Compose
 import cv2
 from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file  # 导入 safetensors 库
+from safetensors.torch import load_file
+import spaces
 
 # Helper function to load model from Hugging Face
 def load_model_by_name(arch_name, checkpoint_path, device):
@@ -70,7 +71,8 @@ def process_image(image, model, device):
     depth_image = Image.fromarray(depth_colored_hwc)
     return depth_image
 
-# Gradio interface function
+# Gradio interface function with GPU support
+@spaces.GPU
 def gradio_interface(image):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -114,7 +116,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),  # Only image input, no mode selection
     outputs=gr.Image(type="pil"),  # Only depth image output, no debug info
     title="Depth Estimation Demo",
-    description="Upload an image to see the depth estimation results. Our model is running on CPU, so it may take a while like ~80 seconds to process."
+    description="Upload an image to see the depth estimation results. Our model is running on GPU for faster processing."
 )
 
 # Launch the Gradio interface