| 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() | |