Update app.py

app.py

@@ -11,7 +11,7 @@ import torch
 import tempfile
 from gradio_imageslider import ImageSlider
 from huggingface_hub import hf_hub_download
-import safetensors
+from safetensors.torch import load_file
 
 from depth_anything_v2.dpt import DepthAnythingV2
 
@@ -47,39 +47,9 @@ model_name = encoder2name[encoder]
 model = DepthAnythingV2(**model_configs[encoder])
 filepath = hf_hub_download(repo_id="depth-anything/Depth-Anything-V2-Metric-Indoor-Large-hf", filename="model.safetensors", repo_type="model")
 
-def create_tensor(storage, info, offset):
-    DTYPES = {"F32": torch.float32}
-    dtype = DTYPES[info["dtype"]]
-    shape = info["shape"]
-    start, stop = info["data_offsets"]
-    return torch.asarray(storage[start + offset : stop + offset], dtype=torch.uint8).view(dtype=dtype).reshape(shape)
-
-def load_file(filename):
-    with open(filename, mode="r", encoding="utf8") as file_obj:
-        with mmap.mmap(file_obj.fileno(), length=0, access=mmap.ACCESS_READ) as m:
-            header = m.read(8)
-            n = int.from_bytes(header, "little")
-            metadata_bytes = m.read(n)
-            metadata = json.loads(metadata_bytes)
-
-    size = os.stat(filename).st_size
-    storage = torch.ByteStorage.from_file(filename, shared=False, size=size).untyped()
-    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 = pytorch_tensor
+state_dict = load_file(filepath)
 
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()