Update app.py

app.py

@@ -3,9 +3,24 @@ import gradio as gr
 import json
 
 
-def bytes_to_giga_bytes(bytes):
-    return f"{(bytes / 1024 / 1024 / 1024):.3f}"
-
+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)
+    for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]:
+        if abs(num_f) < 1000.0:
+            return f"{num_f:3.1f}{unit}"
+        num_f /= 1000.0
+    return f"{num_f:.1f}Y"
+
+def format_output(memory_mapping):
+    markdown_str = ""
+    if memory_mapping:
+        for component, memory in 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)
@@ -73,12 +88,12 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
                         print(selected_file.rfilename)
                         component_wise_memory[component] = bytes_to_giga_bytes(selected_file.size)
 
-    return component_wise_memory
+    return format_output(component_wise_memory)
 
 
 gr.Interface(
     title="Compute component-wise memory of a 🧨 Diffusers pipeline.",
-    description="Sizes will be reported in GB. It can handle pipelines where the text encoder checkpoints might be shared (PixArt, for example).",
+    description="Sizes will be reported in GB.",
     fn=get_component_wise_memory,
     inputs=[
         gr.components.Textbox(lines=1, label="pipeline_id", info="Example: runwayml/stable-diffusion-v1-5"),
@@ -98,11 +113,11 @@ gr.Interface(
             info="Extension to use.",
         ),
     ],
-    outputs="text",
+    outputs="markdown",
     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,