Function _get_checkpoint_shard_files

src/diffusers/utils/hub_utils.py:359–467

For a given model: - download and cache all the shards of a sharded checkpoint if `pretrained_model_name_or_path` is a model ID on the Hub - returns the list of paths to all the shards, as well as some metadata. For the description of each arg, see [`PreTrainedModel.from_pre

(
    pretrained_model_name_or_path,
    index_filename,
    cache_dir=None,
    proxies=None,
    local_files_only=False,
    token=None,
    user_agent=None,
    revision=None,
    subfolder="",
    dduf_entries: dict[str, DDUFEntry] | None = None,
)

def _get_checkpoint_shard_files(
    pretrained_model_name_or_path,
    index_filename,
    cache_dir=None,
    proxies=None,
    local_files_only=False,
    token=None,
    user_agent=None,
    revision=None,
    subfolder="",
    dduf_entries: dict[str, DDUFEntry] | None = None,
):
    """
    For a given model:

    - download and cache all the shards of a sharded checkpoint if `pretrained_model_name_or_path` is a model ID on the
      Hub
    - returns the list of paths to all the shards, as well as some metadata.

    For the description of each arg, see [`PreTrainedModel.from_pretrained`]. `index_filename` is the full path to the
    index (downloaded and cached if `pretrained_model_name_or_path` is a model ID on the Hub).
    """
    if dduf_entries:
        if index_filename not in dduf_entries:
            raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.")
    else:
        if not os.path.isfile(index_filename):
            raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.")

    if dduf_entries:
        index = json.loads(dduf_entries[index_filename].read_text())
    else:
        with open(index_filename, "r") as f:
            index = json.loads(f.read())

    original_shard_filenames = sorted(set(index["weight_map"].values()))
    sharded_metadata = index["metadata"]
    sharded_metadata["all_checkpoint_keys"] = list(index["weight_map"].keys())
    sharded_metadata["weight_map"] = index["weight_map"].copy()
    shards_path = os.path.join(pretrained_model_name_or_path, subfolder)

    # First, let's deal with local folder.
    if os.path.isdir(pretrained_model_name_or_path) or dduf_entries:
        shard_filenames = [os.path.join(shards_path, f) for f in original_shard_filenames]
        for shard_file in shard_filenames:
            if dduf_entries:
                if shard_file not in dduf_entries:
                    raise FileNotFoundError(
                        f"{shards_path} does not appear to have a file named {shard_file} which is "
                        "required according to the checkpoint index."
                    )
            else:
                if not os.path.exists(shard_file):
                    raise FileNotFoundError(
                        f"{shards_path} does not appear to have a file named {shard_file} which is "
                        "required according to the checkpoint index."
                    )
        return shard_filenames, sharded_metadata