fix
Browse files
app.py
CHANGED
@@ -73,13 +73,25 @@ def load_model_index(pipeline_id, token=None, revision=None):
|
|
73 |
|
74 |
|
75 |
def get_individual_model_memory(id, token, variant, extension):
|
|
|
76 |
files_in_repo = model_info(id, token=token, files_metadata=True).siblings
|
77 |
-
|
|
|
78 |
if variant:
|
79 |
-
|
|
|
|
|
80 |
else:
|
81 |
-
|
82 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
83 |
|
84 |
|
85 |
def get_component_wise_memory(
|
|
|
73 |
|
74 |
|
75 |
def get_individual_model_memory(id, token, variant, extension):
|
76 |
+
# Retrieve all files in the repository.
|
77 |
files_in_repo = model_info(id, token=token, files_metadata=True).siblings
|
78 |
+
|
79 |
+
# Filter files by extension and variant (if provided).
|
80 |
if variant:
|
81 |
+
candidates = [x for x in files_in_repo if (extension in x.rfilename) and (variant in x.rfilename)]
|
82 |
+
if not candidates:
|
83 |
+
raise ValueError(f"Requested variant ({variant}) for {id} couldn't be found with {extension} extension.")
|
84 |
else:
|
85 |
+
candidates = [
|
86 |
+
x
|
87 |
+
for x in files_in_repo
|
88 |
+
if (extension in x.rfilename) and all(var not in x.rfilename for var in ALLOWED_VARIANTS[1:])
|
89 |
+
]
|
90 |
+
if not candidates:
|
91 |
+
raise ValueError(f"No file for {id} could be found with {extension} extension without specified variants.")
|
92 |
+
|
93 |
+
# Return the size of the first matching file.
|
94 |
+
return candidates[0].size
|
95 |
|
96 |
|
97 |
def get_component_wise_memory(
|