File size: 587 Bytes
2051791 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | import json
import os
def get_model_storage_size(
base_dir,
index_file="model.safetensors.index.json",
model_file="model.safetensors",
):
size = 0
index_file = os.path.join(base_dir, index_file)
if os.path.exists(index_file):
# model is split into shards
with open(index_file, "r") as f:
index = json.load(f)
for shard in set(index["weight_map"].values()):
size += os.path.getsize(os.path.join(base_dir, shard))
else:
size = os.path.getsize(os.path.join(base_dir, model_file))
return size
|