|
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 revision == "": |
|
revision = None |
|
|
|
if variant == "fp32": |
|
variant = None |
|
|
|
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 = {} |
|
|
|
|
|
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) |
|
|
|
|
|
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 component_filter): |
|
is_folder = len(current_file.rfilename.split("/")) == 2 |
|
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 ( |
|
variant is not None |
|
and variant in current_file.rfilename |
|
and "ema" not in current_file.rfilename |
|
): |
|
selected_file = current_file |
|
elif variant is None and "ema" not in current_file.rfilename: |
|
selected_file = current_file |
|
|
|
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 |
|
|
|
|
|
gr.Interface( |
|
title="Compute component-wise memory of a 🧨 Diffusers pipeline.", |
|
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"), |
|
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="text", |
|
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() |
