Update app.py

app.py

@@ -11,6 +11,7 @@ import torch
 import tempfile
 from gradio_imageslider import ImageSlider
 from huggingface_hub import hf_hub_download
+import safetensors
 
 from depth_anything_v2.dpt import DepthAnythingV2
 
@@ -66,9 +67,19 @@ def load_file(filename):
     offset = n + 8
     return {name: create_tensor(storage, info, offset) for name, info in metadata.items() if name != "__metadata__"}
 
+tensor_data = safetensors.load(filepath)
+
+# Convert to PyTorch tensor
+if isinstance(tensor_data, np.ndarray):
+    pytorch_tensor = torch.tensor(tensor_data)
+elif isinstance(tensor_data, safetensors.Tensor):
+    pytorch_tensor = torch.tensor(tensor_data.numpy())  # Assuming safetensors Tensor has a .numpy() method
+else:
+    raise TypeError("Unsupported data type from safetensors")
 
 #state_dict = torch.load(filepath, map_location="cpu", weights_only=True)
-state_dict = load_file(filepath)
+#state_dict = load_file(filepath)
+state_dict = pytorch_tensor
 
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()