reobustify

app.py

@@ -29,6 +29,8 @@ For further clarification on this topic, feel free to open a [discussion](https:
 `controlnet_id` or `t2i_adapter_id`.
 """
 
+ALLOWED_VARIANTS = ["fp32", "fp16", "bf16"]
+
 
 def format_size(num: int) -> str:
     """Format size in bytes into a human-readable string.
@@ -69,13 +71,12 @@ def load_model_index(pipeline_id, token=None, revision=None):
 
 def get_individual_model_memory(id, token, variant, extension):
     files_in_repo = model_info(id, token=token, files_metadata=True).siblings
-    for x in files_in_repo:
-        if extension in x.rfilename:
-            if variant:
-                if variant in x.rfilename:
-                    return x.size
-            else:
-                return x.size
+    candidates = [x for x in files_in_repo if extension in x.rfilename]
+    if variant:
+        candidate = list(filter(lambda x: variant in x.rfilename, candidates))[0]
+    else:
+        candidate = list(filter(lambda x: all(var not in x.rfilename for var in ALLOWED_VARIANTS[1:]), candidates))[0]
+    return candidate.size
 
 
 def get_component_wise_memory(
@@ -211,7 +212,7 @@ with gr.Interface(
         gr.components.Textbox(lines=1, label="t2i_adapter_id", info="Example: TencentARC/t2iadapter_color_sd14v1"),
         gr.components.Textbox(lines=1, label="hf_token", info="Pass this in case of private repositories."),
         gr.components.Radio(
-            ["fp32", "fp16", "bf16"],
+            ALLOWED_VARIANTS,
             label="variant",
             info="Precision to use for calculation.",
         ),
@@ -232,7 +233,7 @@ with gr.Interface(
             None,
             "TencentARC/t2i-adapter-lineart-sdxl-1.0,TencentARC/t2i-adapter-canny-sdxl-1.0",
             None,
-            "fp32",
+            "fp16",
             None,
             ".safetensors",
         ],