Update app.py

app.py

@@ -4,6 +4,8 @@ import matplotlib
 import numpy as np
 import os
 from PIL import Image
+import mmap
+import json
 import spaces
 import torch
 import tempfile
@@ -43,7 +45,31 @@ encoder = 'vitl'
 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")
-state_dict = torch.load(filepath, map_location="cpu", weights_only=True)
+
+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__"}
+
+
+#state_dict = torch.load(filepath, map_location="cpu", weights_only=True)
+state_dict = load_file(filepath)
+
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()