sayakpaul HF staff commited on
Commit
1632490
1 Parent(s): 0c24640
Files changed (1) hide show
  1. app.py +16 -4
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
- candidates = [x for x in files_in_repo if extension in x.rfilename]
 
78
  if variant:
79
- candidate = list(filter(lambda x: variant in x.rfilename, candidates))[0]
 
 
80
  else:
81
- candidate = list(filter(lambda x: all(var not in x.rfilename for var in ALLOWED_VARIANTS[1:]), candidates))[0]
82
- return candidate.size
 
 
 
 
 
 
 
 
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(