Update app.py

app.py

@@ -1,37 +1,63 @@
-from huggingface_hub import model_info
+from huggingface_hub import model_info, hf_hub_download
 import gradio as gr
+import json
 
 
 def bytes_to_giga_bytes(bytes):
     return f"{(bytes / 1024 / 1024 / 1024):.3f}"
 
 
+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
 
-    if variant == "":
-        variant = None
-
     if revision == "":
         revision = None
 
     if variant == "fp32":
         variant = None
 
-    print(pipeline_id, variant, revision, extension)
-    component_wise_memory = {}
+    print(f"pipeline_id: {pipeline_id}, variant: {variant}, revision: {revision}, extension: {extension}")
 
     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 = {}
+
+    # Handle text encoder separately when it's sharded.
+    if is_text_encoder_shared:
+        for current_file in files_in_repo:
+            if "text_encoder" in current_file.rfilename:
+                if not current_file.rfilename.endswith(".json") and current_file.rfilename.endswith(extension):
+                    if variant is not None and variant in current_file.rfilename:
+                        selected_file = current_file
+                    else:
+                        selected_file = current_file
+                    if "text_encoder" not in component_wise_memory:
+                        component_wise_memory["text_encoder"] = selected_file.size
+                    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")
 
     for current_file in files_in_repo:
-        if all(
-            substring not in current_file.rfilename
-            for substring in ["scheduler", "feature_extractor", "safety_checker", "tokenizer"]
-        ):
+        if all(substring not in current_file.rfilename for substring in component_filter):
             is_folder = len(current_file.rfilename.split("/")) == 2
-            if is_folder:
-                filename = None
+            if is_folder and current_file.rfilename.split("/")[0] in index_dict:
+                selected_file = None
                 if not current_file.rfilename.endswith(".json") and current_file.rfilename.endswith(extension):
                     component = current_file.rfilename.split("/")[0]
                     if (
@@ -39,12 +65,13 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
                         and variant in current_file.rfilename
                         and "ema" not in current_file.rfilename
                     ):
-                        filename = current_file.rfilename
-                    elif "ema" not in current_file.rfilename:
-                        filename = current_file.rfilename
+                        selected_file = current_file
+                    elif variant is None and "ema" not in current_file.rfilename:
+                        selected_file = current_file
 
-                    if filename is not None:
-                        component_wise_memory[component] = bytes_to_giga_bytes(current_file.size)
+                    if selected_file is not None:
+                        print(selected_file.rfilename)
+                        component_wise_memory[component] = bytes_to_giga_bytes(selected_file.size)
 
     return component_wise_memory
 
@@ -75,7 +102,8 @@ gr.Interface(
     examples=[
         ["runwayml/stable-diffusion-v1-5", None, "fp32", None, ".safetensors"],
         ["stabilityai/stable-diffusion-xl-base-1.0", None, "fp16", None, ".safetensors"],
+        [""],
     ],
     theme=gr.themes.Soft(),
     allow_flagging=False,
-).launch()
+).launch()