diffusers/optimized-diffusers-code

1import collections2from pathlib import Path3import functools4import os5import safetensors.torch6from huggingface_hub import model_info7import tempfile8import torch9import functools10import os11import requests12import struct13from huggingface_hub import hf_hub_url14 15DTYPE_MAP = {"F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16}16 17 18# https://huggingface.co/docs/safetensors/v0.3.2/metadata_parsing#python19def _parse_single_file(url):20    print(f"{url=}")21    token = os.getenv("HF_TOKEN")22    assert token, "HF_TOKEN must be set"23    headers = {"Range": "bytes=0-7", "Authorization": f"Bearer {token}"}24    response = requests.get(url, headers=headers)25    length_of_header = struct.unpack("<Q", response.content)[0]26    headers = {"Range": f"bytes=8-{7 + length_of_header}", "Authorization": f"Bearer {token}"}27    response = requests.get(url, headers=headers)28    header = response.json()29    return header30 31 32def _get_dtype_from_safetensor_file(file_path):33    """Inspects a safetensors file and returns the dtype of the first tensor.34 35    If it's not a safetensors file and a URL instead, we query it.36    """37    if "https" in file_path:38        metadata = _parse_single_file(file_path)39        except_format_metadata_keys = sorted({k for k in metadata if k != "__metadata__"})40        string_dtype = metadata[except_format_metadata_keys[0]]["dtype"]41        return DTYPE_MAP[string_dtype]42    try:43        # load_file is simple and sufficient for this info-gathering purpose.44        state_dict = safetensors.torch.load_file(file_path)45        if not state_dict:46            return "N/A (empty)"47 48        # Get the dtype from the first tensor in the state dict49        first_tensor = next(iter(state_dict.values()))50        return first_tensor.dtype51    except Exception as e:52        print(f"Warning: Could not determine dtype from {file_path}. Error: {e}")53        return "N/A (error)"54 55 56def _process_components(component_files, file_accessor_fn, disable_bf16=False):57    """58    Generic function to process components, calculate size, and determine dtype.59 60    Args:61        component_files (dict): A dictionary mapping component names to lists of file objects.62        file_accessor_fn (function): A function that takes a file object and returns63                                     a tuple of (local_path_for_inspection, size_in_bytes, relative_filename).64        disable_bf16 (bool): To disable using `torch.bfloat16`. Use it at your own risk.65 66    Returns:67        dict: A dictionary containing the total memory and detailed component info.68    """69    components_info = {}70    total_size_bytes = 071 72    for name, files in component_files.items():73        # Get dtype by inspecting the first file of the component74        first_file = files[0]75 76        # The accessor function handles how to get the path (download vs local)77        # and its size and relative name.78        inspection_path, _, _ = file_accessor_fn(first_file)79        dtype = _get_dtype_from_safetensor_file(inspection_path)80 81        component_size_bytes = 082        component_file_details = []83        for f in files:84            _, size_bytes, rel_filename = file_accessor_fn(f)85            component_size_bytes += size_bytes86            component_file_details.append({"filename": rel_filename, "size_mb": size_bytes / (1024**2)})87 88        if dtype == torch.float32 and not disable_bf16:89            print(90                f"The `dtype` for component ({name}) is torch.float32. Since bf16 computation is not disabled "91                "we will slash the total size of this component by 2."92            )93            total_size_bytes += component_size_bytes / 294        else:95            total_size_bytes += component_size_bytes96 97        components_info[name] = {98            "size_gb": round(component_size_bytes / (1024**3), 3),99            "dtype": dtype,100            "files": sorted(component_file_details, key=lambda x: x["filename"]),101        }102 103    return {104        "total_loading_memory_gb": round(total_size_bytes / (1024**3), 3),105        "components": components_info,106    }107 108 109@functools.lru_cache()110def _determine_memory_from_hub_ckpt(ckpt_id, variant=None, disable_bf16=False):111    """112    Determines memory and dtypes for a checkpoint on the Hugging Face Hub.113    """114    files_in_repo = model_info(ckpt_id, files_metadata=True, token=os.getenv("HF_TOKEN")).siblings115    all_safetensors_siblings = [116        s for s in files_in_repo if s.rfilename.endswith(".safetensors") and "/" in s.rfilename117    ]118    if variant:119        all_safetensors_siblings = [f for f in all_safetensors_siblings if variant in f.rfilename]120 121    component_files = collections.defaultdict(list)122    for sibling in all_safetensors_siblings:123        component_name = Path(sibling.rfilename).parent.name124        component_files[component_name].append(sibling)125 126    with tempfile.TemporaryDirectory() as temp_dir:127 128        def hub_file_accessor(file_obj):129            """Accessor for Hub files: downloads them and returns path/size."""130            print(f"Querying '{file_obj.rfilename}' for inspection...")131            url = hf_hub_url(ckpt_id, file_obj.rfilename)132            return url, file_obj.size, file_obj.rfilename133 134        # We only need to download one file per component for dtype inspection.135        # To make this efficient, we create a specialized accessor for the processing loop136        # that only downloads the *first* file encountered for a component.137        downloaded_for_inspection = {}138 139        def efficient_hub_accessor(file_obj):140            component_name = Path(file_obj.rfilename).parent.name141            if component_name not in downloaded_for_inspection:142                path, _, _ = hub_file_accessor(file_obj)143                downloaded_for_inspection[component_name] = path144 145            inspection_path = downloaded_for_inspection[component_name]146            return inspection_path, file_obj.size, file_obj.rfilename147 148        return _process_components(component_files, efficient_hub_accessor, disable_bf16)149 150 151@functools.lru_cache()152def _determine_memory_from_local_ckpt(path: str, variant=None, disable_bf16=False):153    """154    Determines memory and dtypes for a local checkpoint.155    """156    ckpt_path = Path(path)157    if not ckpt_path.is_dir():158        return {"error": f"Checkpoint path '{path}' not found or is not a directory."}159 160    all_safetensors_paths = list(ckpt_path.glob("**/*.safetensors"))161    if variant:162        all_safetensors_paths = [p for p in all_safetensors_paths if variant in p.name]163 164    component_files = collections.defaultdict(list)165    for file_path in all_safetensors_paths:166        component_name = file_path.parent.name167        component_files[component_name].append(file_path)168 169    def local_file_accessor(file_path):170        """Accessor for local files: just returns their path and size."""171        return file_path, file_path.stat().st_size, str(file_path.relative_to(ckpt_path))172 173    return _process_components(component_files, local_file_accessor, disable_bf16)174 175 176def determine_pipe_loading_memory(ckpt_id: str, variant=None, disable_bf16=False):177    """178    Determines the memory and dtypes for a pipeline, whether it's local or on the Hub.179    """180    if os.path.isdir(ckpt_id):181        return _determine_memory_from_local_ckpt(ckpt_id, variant, disable_bf16)182    else:183        return _determine_memory_from_hub_ckpt(ckpt_id, variant, disable_bf16)184 185 186if __name__ == "__main__":187    output = _determine_memory_from_hub_ckpt("Wan-AI/Wan2.1-T2V-14B-Diffusers")188    total_size_gb = output["total_loading_memory_gb"]189    safetensor_files = output["components"]190    print(f"{total_size_gb=} GB")191    print(f"{safetensor_files=}")192    print("\n")