Update app.py

app.py

@@ -2,10 +2,10 @@ from huggingface_hub import model_info, hf_hub_download
 import gradio as gr
 import json
 
+component_filter = ["scheduler", "safety_checker", "tokenizer"]
 
 def format_size(num: int) -> str:
     """Format size in bytes into a human-readable string.
-
     Taken from https://stackoverflow.com/a/1094933
     """
     num_f = float(num)
@@ -43,10 +43,25 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
     files_in_repo = model_info(pipeline_id, revision=revision, token=token, files_metadata=True).siblings
     index_dict = load_model_index(pipeline_id, token=token, revision=revision)
 
-    is_text_encoder_shared = any(".index.json" in file_obj.rfilename for file_obj in files_in_repo)
-    component_wise_memory = {}
+    # 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"])
+    for current_component in index_dict:
+        if current_component not in index_filter:
+            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 = lambda filename: extension in filename and variant in filename if variant is not None else lambda filename: extension in filename
+                variant_present_with_extension = any(condition(filename) for filename in current_component_filenames)
+                if not variant_present_with_extension:
+                    raise ValueError(f"Requested extension ({extension}) and variant ({variant}) not present for {current_component}. Available files for this component:\n{current_component_filenames}.")
+            else:
+                raise ValueError(f"Problem with {current_component}.")
+
 
     # Handle text encoder separately when it's sharded.
+    is_text_encoder_shared = any(".index.json" in file_obj.rfilename for file_obj in files_in_repo)
+    component_wise_memory = {}
     if is_text_encoder_shared:
         for current_file in files_in_repo:
             if "text_encoder" in current_file.rfilename:
@@ -60,10 +75,7 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
                     else:
                         component_wise_memory["text_encoder"] += selected_file.size
 
-    print(component_wise_memory)
-
     # Handle pipeline components.
-    component_filter = ["scheduler", "feature_extractor", "safety_checker", "tokenizer"]
     if is_text_encoder_shared:
         component_filter.append("text_encoder")
 
@@ -87,37 +99,4 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
                         print(selected_file.rfilename)
                         component_wise_memory[component] = selected_file.size
 
-    return format_output(pipeline_id, component_wise_memory)
-
-
-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) 🤗",
-    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",
-            ],
-            label="variant",
-            info="Precision to use for calculation.",
-        ),
-        gr.components.Textbox(lines=1, label="revision", info="Repository revision to use."),
-        gr.components.Dropdown(
-            [".bin", ".safetensors"],
-            label="extension",
-            info="Extension to use.",
-        ),
-    ],
-    outputs=[gr.Markdown(label="Output")],
-    examples=[
-        ["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"],
-    ],
-    theme=gr.themes.Soft(),
-    allow_flagging=False,
-).launch(show_error=True)
+    return format_output(pipeline_id, component_wise_memory)