Update app.py
Browse files
app.py
CHANGED
@@ -2,7 +2,7 @@ from huggingface_hub import model_info, hf_hub_download
|
|
2 |
import gradio as gr
|
3 |
import json
|
4 |
|
5 |
-
|
6 |
|
7 |
|
8 |
def format_size(num: int) -> str:
|
@@ -48,8 +48,8 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
|
|
48 |
index_dict = load_model_index(pipeline_id, token=token, revision=revision)
|
49 |
|
50 |
# Check if all the concerned components have the checkpoints in the requested "variant" and "extension".
|
51 |
-
index_filter =
|
52 |
-
index_filter.extend(["_class_name", "_diffusers_version"])
|
53 |
for current_component in index_dict:
|
54 |
if current_component not in index_filter:
|
55 |
current_component_fileobjs = list(filter(lambda x: current_component in x.rfilename, files_in_repo))
|
@@ -86,10 +86,10 @@ def get_component_wise_memory(pipeline_id, token=None, variant=None, revision=No
|
|
86 |
|
87 |
# Handle pipeline components.
|
88 |
if is_text_encoder_shared:
|
89 |
-
|
90 |
|
91 |
for current_file in files_in_repo:
|
92 |
-
if all(substring not in current_file.rfilename for substring in
|
93 |
is_folder = len(current_file.rfilename.split("/")) == 2
|
94 |
if is_folder and current_file.rfilename.split("/")[0] in index_dict:
|
95 |
selected_file = None
|
|
|
2 |
import gradio as gr
|
3 |
import json
|
4 |
|
5 |
+
COMPONENT_FILTER = ["scheduler", "safety_checker", "tokenizer"]
|
6 |
|
7 |
|
8 |
def format_size(num: int) -> str:
|
|
|
48 |
index_dict = load_model_index(pipeline_id, token=token, revision=revision)
|
49 |
|
50 |
# Check if all the concerned components have the checkpoints in the requested "variant" and "extension".
|
51 |
+
index_filter = COMPONENT_FILTER.copy()
|
52 |
+
index_filter.extend(["_class_name", "_diffusers_version", "force_zeros_for_empty_prompt", "add_watermarker"])
|
53 |
for current_component in index_dict:
|
54 |
if current_component not in index_filter:
|
55 |
current_component_fileobjs = list(filter(lambda x: current_component in x.rfilename, files_in_repo))
|
|
|
86 |
|
87 |
# Handle pipeline components.
|
88 |
if is_text_encoder_shared:
|
89 |
+
COMPONENT_FILTER.append("text_encoder")
|
90 |
|
91 |
for current_file in files_in_repo:
|
92 |
+
if all(substring not in current_file.rfilename for substring in COMPONENT_FILTER):
|
93 |
is_folder = len(current_file.rfilename.split("/")) == 2
|
94 |
if is_folder and current_file.rfilename.split("/")[0] in index_dict:
|
95 |
selected_file = None
|