| 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 | |