Update app.py

app.py

@@ -4,6 +4,7 @@ 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
@@ -15,6 +16,7 @@ def format_size(num: int) -> str:
         num_f /= 1000.0
     return f"{num_f:.1f}Y"
 
+
 def format_output(pipeline_id, memory_mapping):
     markdown_str = f"## {pipeline_id}\n"
     if memory_mapping:
@@ -22,12 +24,14 @@ def format_output(pipeline_id, memory_mapping):
             markdown_str += f"* {component}: {format_size(memory)}\n"
     return markdown_str
 
+
 def load_model_index(pipeline_id, token=None, revision=None):
     index_path = hf_hub_download(repo_id=pipeline_id, filename="model_index.json", revision=revision, token=token)
     with open(index_path, "r") as f:
         index_dict = json.load(f)
     return index_dict
 
+
 def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=None, extension=".safetensors"):
     if token == "":
         token = None
@@ -48,17 +52,22 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
     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))
+            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
+                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}.")
+                    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 = {}
@@ -99,4 +108,37 @@ 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)
+    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)