Function load_state_dict

src/diffusers/models/model_loading_utils.py:155–210

Reads a checkpoint file, returning properly formatted errors if they arise.

(
    checkpoint_file: str | os.PathLike,
    dduf_entries: dict[str, DDUFEntry] | None = None,
    disable_mmap: bool = False,
    map_location: str | torch.device = "cpu",
)

def load_state_dict(
    checkpoint_file: str | os.PathLike,
    dduf_entries: dict[str, DDUFEntry] | None = None,
    disable_mmap: bool = False,
    map_location: str | torch.device = "cpu",
):
    """
    Reads a checkpoint file, returning properly formatted errors if they arise.
    """
    # TODO: maybe refactor a bit this part where we pass a dict here
    if isinstance(checkpoint_file, dict):
        return checkpoint_file
    try:
        file_extension = os.path.basename(checkpoint_file).split(".")[-1]
        if file_extension == SAFETENSORS_FILE_EXTENSION:
            if dduf_entries:
                # tensors are loaded on cpu
                with dduf_entries[checkpoint_file].as_mmap() as mm:
                    return safetensors.torch.load(mm)
            if disable_mmap:
                return safetensors.torch.load(open(checkpoint_file, "rb").read())
            else:
                return safetensors.torch.load_file(checkpoint_file, device=map_location)
        elif file_extension == GGUF_FILE_EXTENSION:
            return load_gguf_checkpoint(checkpoint_file)
        else:
            extra_args = {}
            weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {}
            # mmap can only be used with files serialized with zipfile-based format.
            if (
                isinstance(checkpoint_file, str)
                and map_location != "meta"
                and is_torch_version(">=", "2.1.0")
                and is_zipfile(checkpoint_file)
                and not disable_mmap
            ):
                extra_args = {"mmap": True}
            return torch.load(checkpoint_file, map_location=map_location, **weights_only_kwarg, **extra_args)
    except Exception as e:
        try:
            with open(checkpoint_file) as f:
                if f.read().startswith("version"):
                    raise OSError(
                        "You seem to have cloned a repository without having git-lfs installed. Please install "
                        "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder "
                        "you cloned."
                    )
                else:
                    raise ValueError(
                        f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained "
                        "model. Make sure you have saved the model properly."
                    ) from e
        except (UnicodeDecodeError, ValueError):
            raise OSError(
                f"Unable to load weights from checkpoint file for '{checkpoint_file}' at '{checkpoint_file}'. "
            )