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

Code before:
        ignore_mismatched_sizes,
    )
    error_msgs = []
    if low_cpu_mem_usage:
        offload_index, state_dict_index = load_model_dict_into_meta(
            model,
            state_dict,
            device_map=device_map,
            dtype=dtype,
            hf_quantizer=hf_quantizer,
            keep_in_fp32_modules=keep_in_fp32_modules,
            unexpected_keys=unexpected_keys,
            offload_folder=offload_folder,
            offload_index=offload_index,
            state_dict_index=state_dict_index,
            state_dict_folder=state_dict_folder,
        )
    else:
        if assign_to_params_buffers is None:
            assign_to_params_buffers = check_support_param_buffer_assignment(model, state_dict)

        error_msgs += _load_state_dict_into_model(model, state_dict, assign_to_params_buffers)
    return offload_index, state_dict_index, mismatched_keys, error_msgs


def load_shard_files_with_threadpool(args_list):
    num_workers = int(os.environ.get("HF_PARALLEL_LOADING_WORKERS", "8"))

    # Do not spawn anymore workers than you need
    num_workers = min(len(args_list), num_workers)

    logger.info(f"Loading model weights in parallel with {num_workers} workers...")

    error_msgs = []
    mismatched_keys = []

    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        with logging.tqdm(total=len(args_list), desc="Loading checkpoint shards") as pbar:
            futures = [executor.submit(load_shard_file, arg) for arg in args_list]
            for future in as_completed(futures):
                result = future.result()
                offload_index, state_dict_index, _mismatched_keys, _error_msgs = result
                error_msgs += _error_msgs
                mismatched_keys += _mismatched_keys
                pbar.update(1)

    return offload_index, state_dict_index, mismatched_keys, error_msgs


def _find_mismatched_keys(
    state_dict,

Review comment (refactor):
Would have this also be an internal utility and take named kwargs, it just makes loading requirements more clear. 

```suggestion
def _load_shard_files_with_threadpool(
    shard_files,
    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:
        ignore_mismatched_sizes,
    )
    error_msgs = []
    if low_cpu_mem_usage:
        offload_index, state_dict_index = load_model_dict_into_meta(
            model,
            state_dict,
            device_map=device_map,
            dtype=dtype,
            hf_quantizer=hf_quantizer,
            keep_in_fp32_modules=keep_in_fp32_modules,
            unexpected_keys=unexpected_keys,
            offload_folder=offload_folder,
            offload_index=offload_index,
            state_dict_index=state_dict_index,
            state_dict_folder=state_dict_folder,
        )
    else:
        assign_to_params_buffers = check_support_param_buffer_assignment(model, state_dict)

        error_msgs += _load_state_dict_into_model(model, state_dict, assign_to_params_buffers)
    return offload_index, state_dict_index, mismatched_keys, error_msgs


def _load_shard_files_with_threadpool(
    shard_files,
    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,
):
    # Do not spawn anymore workers than you need
    num_workers = min(len(shard_files), DEFAULT_HF_PARALLEL_LOADING_WORKERS)

    logger.info(f"Loading model weights in parallel with {num_workers} workers...")

    error_msgs = []
    mismatched_keys = []

    load_one = functools.partial(