File size: 2,861 Bytes
bf03ba4 4b23311 983e4d7 4b23311 bf03ba4 4b23311 bf03ba4 4b23311 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from huggingface_hub import model_info
import gradio as gr
def bytes_to_giga_bytes(bytes):
return f"{(bytes / 1024 / 1024 / 1024):.3f}"
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 = {}
files_in_repo = model_info(pipeline_id, revision=revision, token=token, files_metadata=True).siblings
for current_file in files_in_repo:
if all(
substring not in current_file.rfilename
for substring in ["scheduler", "feature_extractor", "safety_checker", "tokenizer"]
):
is_folder = len(current_file.rfilename.split("/")) == 2
if is_folder:
filename = 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
):
filename = current_file.rfilename
elif "ema" not in current_file.rfilename:
filename = current_file.rfilename
if filename is not None:
component_wise_memory[component] = bytes_to_giga_bytes(current_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()
|