from huggingface_hub import hf_hub_download, model_info
import gradio as gr
import json

COMPONENT_FILTER = [
    "scheduler",
    "feature_extractor",
    "tokenizer",
    "tokenizer_2",
    "_class_name",
    "_diffusers_version",
]

ARTICLE = """
## Notes on how to use the `controlnet_id` and `t2i_adapter_id` fields

Both `controlnet_id` and `t2i_adapter_id` fields support passing multiple checkpoint ids,
e.g., "thibaud/controlnet-openpose-sdxl-1.0,diffusers/controlnet-canny-sdxl-1.0". For 
`t2i_adapter_id`, this could be like - "TencentARC/t2iadapter_keypose_sd14v1,TencentARC/t2iadapter_depth_sd14v1".

Users should take care of passing the underlying base `pipeline_id` appropriately. For example,
passing `pipeline_id` as "runwayml/stable-diffusion-v1-5" and `controlnet_id` as "thibaud/controlnet-openpose-sdxl-1.0"
won't result in an error but these two things aren't meant to compatible. You should pass
a `controlnet_id` that is compatible with "runwayml/stable-diffusion-v1-5".

For further clarification on this topic, feel free to open a [discussion](https://huggingface.co/spaces/diffusers/compute-pipeline-size/discussions).

📔 Also, note that `revision` field is only reserved for `pipeline_id`. It won't have any effect on the
`controlnet_id` or `t2i_adapter_id`.
"""

ALLOWED_VARIANTS = ["fp32", "fp16", "bf16"]


def format_size(num: int) -> str:
    """Format size in bytes into a human-readable string.
    Taken from https://stackoverflow.com/a/1094933
    """
    num_f = float(num)
    for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]:
        if abs(num_f) < 1000.0:
            return f"{num_f:3.1f}{unit}"
        num_f /= 1000.0
    return f"{num_f:.1f}Y"


def format_output(pipeline_id, memory_mapping, variant=None, controlnet_mapping=None, t2i_adapter_mapping=None):
    if variant is None:
        variant = "fp32"

    markdown_str = f"## {pipeline_id} ({variant})\n"

    if memory_mapping:
        for component, memory in memory_mapping.items():
            markdown_str += f"* {component}: {format_size(memory)}\n"
    if controlnet_mapping:
        markdown_str += f"\n## ControlNet(s) ({variant})\n"
        for controlnet_id, memory in controlnet_mapping.items():
            markdown_str += f"* {controlnet_id}: {format_size(memory)}\n"
    if t2i_adapter_mapping:
        markdown_str += f"\n## T2I-Adapters(s) ({variant})\n"
        for t2_adapter_id, memory in t2i_adapter_mapping.items():
            markdown_str += f"* {t2_adapter_id}: {format_size(memory)}\n"

    return markdown_str


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_individual_model_memory(id, token, variant, extension):
    # Retrieve all files in the repository.
    files_in_repo = model_info(id, token=token, files_metadata=True).siblings

    # Filter files by extension and variant (if provided).
    if variant:
        candidates = [x for x in files_in_repo if (extension in x.rfilename) and (variant in x.rfilename)]
        if not candidates:
            raise ValueError(f"Requested variant ({variant}) for {id} couldn't be found with {extension} extension.")
    else:
        candidates = [
            x
            for x in files_in_repo
            if (extension in x.rfilename) and all(var not in x.rfilename for var in ALLOWED_VARIANTS[1:])
        ]
        if not candidates:
            raise ValueError(f"No file for {id} could be found with {extension} extension without specified variants.")

    # Return the size of the first matching file.
    return candidates[0].size


def get_component_wise_memory(
    pipeline_id,
    controlnet_id=None,
    t2i_adapter_id=None,
    token=None,
    variant=None,
    revision=None,
    extension=".safetensors",
):
    if controlnet_id == "":
        controlnet_id = None

    if t2i_adapter_id == "":
        t2i_adapter_id = None

    if controlnet_id and t2i_adapter_id:
        raise ValueError("Both `controlnet_id` and `t2i_adapter_id` cannot be provided.")

    if token == "":
        token = None

    if revision == "":
        revision = None

    if variant == "fp32":
        variant = None

    # Handle ControlNet and T2I-Adapter.
    controlnet_mapping = t2_adapter_mapping = None
    if controlnet_id is not None:
        controlnet_id = controlnet_id.split(",")
        controlnet_sizes = [
            get_individual_model_memory(id_, token=token, variant=variant, extension=extension)
            for id_ in controlnet_id
        ]
        controlnet_mapping = dict(zip(controlnet_id, controlnet_sizes))
    elif t2i_adapter_id is not None:
        t2i_adapter_id = t2i_adapter_id.split(",")
        t2i_adapter_sizes = [
            get_individual_model_memory(id_, token=token, variant=variant, extension=extension)
            for id_ in t2i_adapter_id
        ]
        t2_adapter_mapping = dict(zip(t2i_adapter_id, t2i_adapter_sizes))

    print(f"pipeline_id: {pipeline_id}, variant: {variant}, revision: {revision}, extension: {extension}")

    # Load pipeline metadata.
    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)

    # Check if all the concerned components have the checkpoints in
    # the requested "variant" and "extension".
    print(f"Index dict: {index_dict}")
    for current_component in index_dict:
        if (
            current_component not in COMPONENT_FILTER
            and isinstance(index_dict[current_component], list)
            and len(index_dict[current_component]) == 2
        ):
            current_component_fileobjs = list(filter(lambda x: current_component in x.rfilename, files_in_repo))

            if current_component_fileobjs:
                current_component_filenames = [fileobj.rfilename for fileobj in current_component_fileobjs]
                condition = (  # noqa: E731
                    lambda filename: extension in filename and variant in filename
                    if variant is not None
                    else lambda filename: extension in filename
                )
                variant_present_with_extension = any(condition(filename) for filename in current_component_filenames)
                if not variant_present_with_extension:
                    formatted_filenames = ", ".join(current_component_filenames)
                    raise ValueError(
                        f"Requested extension ({extension}) and variant ({variant}) not present for {current_component}."
                        f" Available files for this component: {formatted_filenames}."
                    )
            else:
                raise ValueError(f"Problem with {current_component}.")

    # Handle text encoder separately when it's sharded.
    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

    # Handle pipeline components.
    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:
                        component_wise_memory[component] = selected_file.size

    return format_output(pipeline_id, component_wise_memory, variant, controlnet_mapping, t2_adapter_mapping)


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    with gr.Column():
        gr.Markdown(
            """<img src="https://huggingface.co/spaces/hf-accelerate/model-memory-usage/resolve/main/measure_model_size.png" style="float: left;" width="150" height="175"><h1>🧨 Diffusers Pipeline Memory Calculator</h1>
    This tool will help you to gauge the memory requirements of a Diffusers pipeline. Pipelines containing text encoders with sharded checkpoints are also supported
    (PixArt-Alpha, for example) 🤗 See instructions below the form on how to pass `controlnet_id` or `t2_adapter_id`. When performing inference, expect to add up to an
    additional 20% to this as found by [EleutherAI](https://blog.eleuther.ai/transformer-math/). The final memory requirement will also depend on the requested resolution. You can click on one
    of the examples below the "Calculate Memory Usage" button to get started. Design adapted from [this Space](https://huggingface.co/spaces/hf-accelerate/model-memory-usage).
    """
        )
        out_text = gr.Markdown()
        with gr.Row():
            pipeline_id = gr.Textbox(lines=1, label="pipeline_id", info="Example: runwayml/stable-diffusion-v1-5")
        with gr.Row():
            controlnet_id = gr.Textbox(lines=1, label="controlnet_id", info="Example: lllyasviel/sd-controlnet-canny")
            t2i_adapter_id = gr.Textbox(
                lines=1, label="t2i_adapter_id", info="Example: TencentARC/t2iadapter_color_sd14v1"
            )
        with gr.Row():
            token = gr.Textbox(lines=1, label="hf_token", info="Pass this in case of private/gated repositories.")
            variant = gr.Radio(
                ALLOWED_VARIANTS,
                label="variant",
                info="Precision to use for calculation.",
            )
            revision = gr.Textbox(lines=1, label="revision", info="Repository revision to use.")
            extension = gr.Radio(
                [".bin", ".safetensors"],
                label="extension",
                info="Extension to use.",
            )
        with gr.Row():
            btn = gr.Button("Calculate Memory Usage")

        gr.Markdown("## Examples")
        gr.Examples(
            [
                ["runwayml/stable-diffusion-v1-5", None, None, None, "fp32", None, ".safetensors"],
                ["PixArt-alpha/PixArt-XL-2-1024-MS", None, None, None, "fp32", None, ".safetensors"],
                [
                    "runwayml/stable-diffusion-v1-5",
                    "lllyasviel/sd-controlnet-canny",
                    None,
                    None,
                    "fp32",
                    None,
                    ".safetensors",
                ],
                [
                    "stabilityai/stable-diffusion-xl-base-1.0",
                    None,
                    "TencentARC/t2i-adapter-lineart-sdxl-1.0,TencentARC/t2i-adapter-canny-sdxl-1.0",
                    None,
                    "fp16",
                    None,
                    ".safetensors",
                ],
                ["stabilityai/stable-cascade", None, None, None, "bf16", None, ".safetensors"],
                ["Deci/DeciDiffusion-v2-0", None, None, None, "fp32", None, ".safetensors"],
            ],
            [pipeline_id, controlnet_id, t2i_adapter_id, token, variant, revision, extension],
            out_text,
            get_component_wise_memory,
            cache_examples=False,
        )

        gr.Markdown(ARTICLE)

    btn.click(
        get_component_wise_memory,
        inputs=[pipeline_id, controlnet_id, t2i_adapter_id, token, variant, revision, extension],
        outputs=[out_text],
        api_name=False,
    )

demo.launch(show_error=True)