multiple improvements.

app.py

@@ -2,7 +2,14 @@ from huggingface_hub import model_info, hf_hub_download
 import gradio as gr
 import json
 
-COMPONENT_FILTER = ["scheduler", "safety_checker", "feature_extractor", "tokenizer", "tokenizer_2"]
+COMPONENT_FILTER = [
+    "scheduler",
+    "feature_extractor",
+    "tokenizer",
+    "tokenizer_2",
+    "_class_name",
+    "_diffusers_version",
+]
 
 
 def format_size(num: int) -> str:
@@ -48,14 +55,17 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
     index_dict = load_model_index(pipeline_id, token=token, revision=revision)
 
     # Check if all the concerned components have the checkpoints in the requested "variant" and "extension".
-    index_filter = COMPONENT_FILTER.copy()
-    index_filter.extend(["_class_name", "_diffusers_version", "force_zeros_for_empty_prompt", "add_watermarker"])
+    print(f"Index dict: {index_dict}")
     for current_component in index_dict:
-        if current_component not in index_filter:
+        if (
+            current_component not in COMPONENT_FILTER
+            and isinstance(index_dict[current_component], list)
+            and len(index_dict[current_component]) == 2
+        ):
             current_component_fileobjs = list(filter(lambda x: current_component in x.rfilename, files_in_repo))
             if current_component_fileobjs:
                 current_component_filenames = [fileobj.rfilename for fileobj in current_component_fileobjs]
-                condition = (
+                condition = (  # noqa: E731
                     lambda filename: extension in filename and variant in filename
                     if variant is not None
                     else lambda filename: extension in filename
@@ -105,29 +115,26 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
                         selected_file = current_file
 
                     if selected_file is not None:
-                        print(selected_file.rfilename)
                         component_wise_memory[component] = selected_file.size
 
     return format_output(pipeline_id, component_wise_memory)
 
 
-gr.Interface(
+with gr.Interface(
     title="Compute component-wise memory of a 🧨 Diffusers pipeline.",
-    description="Sizes will be reported in GB. Pipelines containing text encoders with sharded checkpoints are also supported (PixArt-Alpha, for example) 🤗",
+    description="Pipelines containing text encoders with sharded checkpoints are also supported"
+    " (PixArt-Alpha, for example) 🤗",
     fn=get_component_wise_memory,
     inputs=[
         gr.components.Textbox(lines=1, label="pipeline_id", info="Example: runwayml/stable-diffusion-v1-5"),
         gr.components.Textbox(lines=1, label="hf_token", info="Pass this in case of private repositories."),
-        gr.components.Dropdown(
-            [
-                "fp32",
-                "fp16",
-            ],
+        gr.components.Radio(
+            ["fp32", "fp16", "bf16"],
             label="variant",
             info="Precision to use for calculation.",
         ),
         gr.components.Textbox(lines=1, label="revision", info="Repository revision to use."),
-        gr.components.Dropdown(
+        gr.components.Radio(
             [".bin", ".safetensors"],
             label="extension",
             info="Extension to use.",
@@ -138,7 +145,10 @@ gr.Interface(
         ["runwayml/stable-diffusion-v1-5", None, "fp32", None, ".safetensors"],
         ["stabilityai/stable-diffusion-xl-base-1.0", None, "fp16", None, ".safetensors"],
         ["PixArt-alpha/PixArt-XL-2-1024-MS", None, "fp32", None, ".safetensors"],
+        ["stabilityai/stable-cascade", None, "bf16", None, ".safetensors"],
+        ["Deci/DeciDiffusion-v2-0", None, "fp32", None, ".safetensors"],
     ],
     theme=gr.themes.Soft(),
-    allow_flagging=False,
-).launch(show_error=True)
+    allow_flagging="never",
+) as demo:
+    demo.launch(show_error=True)