Update app.py

app.py

@@ -15,8 +15,8 @@ def format_size(num: int) -> str:
         num_f /= 1000.0
     return f"{num_f:.1f}Y"
 
-def format_output(memory_mapping):
-    markdown_str = ""
+def format_output(pipeline_id, memory_mapping):
+    markdown_str = f"## {pipeline_id}\n"
     if memory_mapping:
         for component, memory in memory_mapping.items():
             markdown_str += f"* {component}: {format_size(memory)}\n"
@@ -28,7 +28,6 @@ def load_model_index(pipeline_id, token=None, revision=None):
         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
@@ -88,7 +87,7 @@ 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(component_wise_memory)
+    return format_output(pipeline_id, component_wise_memory)
 
 
 gr.Interface(
@@ -113,7 +112,7 @@ gr.Interface(
             info="Extension to use.",
         ),
     ],
-    outputs="markdown",
+    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"],