Repository: huggingface/diffusers
File: src/diffusers/models/model_loading_utils.py (Python)

Code before:
    checking if the model explicitly disables it, then by ensuring that the state dict keys are a subset of the model's
    parameters.

    """
    if model_to_load.device.type == "meta":
        return False

    if len([key for key in state_dict if key.startswith(start_prefix)]) == 0:
        return False

    # Some models explicitly do not support param buffer assignment
    if not getattr(model_to_load, "_supports_param_buffer_assignment", True):
        logger.debug(
            f"{model_to_load.__class__.__name__} does not support param buffer assignment, loading will be slower"
        )
        return False

    # If the model does, the incoming `state_dict` and the `model_to_load` must be the same dtype
    first_key = next(iter(model_to_load.state_dict().keys()))
    if start_prefix + first_key in state_dict:
        return state_dict[start_prefix + first_key].dtype == model_to_load.state_dict()[first_key].dtype

    return False


def load_shard_file(args):
    (
        model,
        model_state_dict,
        shard_file,
        device_map,
        dtype,
        hf_quantizer,
        keep_in_fp32_modules,
        dduf_entries,
        loaded_keys,
        unexpected_keys,
        offload_index,
        offload_folder,
        state_dict_index,
        state_dict_folder,
        ignore_mismatched_sizes,
        low_cpu_mem_usage,
    ) = args
    assign_to_params_buffers = None
    state_dict = load_state_dict(shard_file, dduf_entries=dduf_entries)
    mismatched_keys = _find_mismatched_keys(
        state_dict,
        model_state_dict,
        loaded_keys,
        ignore_mismatched_sizes,

Review comment (refactor):
Let's make this an internal utility function and use named kwargs. It makes it clear what this function needs to load the shard
```suggestion
def load_shard_file(
    shard_file,
    model,
    model_state_dict,
    device_map=None,
    dtype=None,
    hf_quantizer=None,
    keep_in_fp32_modules=None,
    dduf_entries=None,
    loaded_keys=None,
    unexpected_keys=None,
    offload_index=None,
    offload_folder=None,
    state_dict_index=None,
    state_dict_folder=None,
    ignore_mismatched_sizes=False,
    low_cpu_mem_usage=False,
):

```



Code after:
    """
    Checks if `model_to_load` supports param buffer assignment (such as when loading in empty weights) by first
    checking if the model explicitly disables it, then by ensuring that the state dict keys are a subset of the model's
    parameters.

    """
    if model_to_load.device.type == "meta":
        return False

    if len([key for key in state_dict if key.startswith(start_prefix)]) == 0:
        return False

    # Some models explicitly do not support param buffer assignment
    if not getattr(model_to_load, "_supports_param_buffer_assignment", True):
        logger.debug(
            f"{model_to_load.__class__.__name__} does not support param buffer assignment, loading will be slower"
        )
        return False

    # If the model does, the incoming `state_dict` and the `model_to_load` must be the same dtype
    first_key = next(iter(model_to_load.state_dict().keys()))
    if start_prefix + first_key in state_dict:
        return state_dict[start_prefix + first_key].dtype == model_to_load.state_dict()[first_key].dtype

    return False


def _load_shard_file(
    shard_file,
    model,
    model_state_dict,
    device_map=None,
    dtype=None,
    hf_quantizer=None,
    keep_in_fp32_modules=None,
    dduf_entries=None,
    loaded_keys=None,
    unexpected_keys=None,
    offload_index=None,
    offload_folder=None,
    state_dict_index=None,
    state_dict_folder=None,
    ignore_mismatched_sizes=False,
    low_cpu_mem_usage=False,
):
    state_dict = load_state_dict(shard_file, dduf_entries=dduf_entries)
    mismatched_keys = _find_mismatched_keys(
        state_dict,
        model_state_dict,
        loaded_keys,
        ignore_mismatched_sizes,