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

Code before:

        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 = int(os.environ.get("HF_PARALLEL_LOADING_WORKERS", str(DEFAULT_HF_PARALLEL_LOADING_WORKERS)))
    num_workers = min(len(shard_files), num_workers)

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

    error_msgs = []
    mismatched_keys = []

    load_one = functools.partial(
        _load_shard_file,
        model=model,
        model_state_dict=model_state_dict,
        device_map=device_map,
        dtype=dtype,
        hf_quantizer=hf_quantizer,
        keep_in_fp32_modules=keep_in_fp32_modules,
        dduf_entries=dduf_entries,
        loaded_keys=loaded_keys,
        unexpected_keys=unexpected_keys,
        offload_index=offload_index,
        offload_folder=offload_folder,
        state_dict_index=state_dict_index,
        state_dict_folder=state_dict_folder,
        ignore_mismatched_sizes=ignore_mismatched_sizes,
        low_cpu_mem_usage=low_cpu_mem_usage,
    )


Review comment (suggestion):
Nit. Would move DEFAULT_HF_PARALLEL_LOADING_WORKERS into constants and invoke this way. 

```suggestion
    num_workers = min(len(shard_files), DEFAULT_HF_PARALLEL_LOADING_WORKERS)
```

Code after:

        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(
        _load_shard_file,
        model=model,
        model_state_dict=model_state_dict,
        device_map=device_map,
        dtype=dtype,
        hf_quantizer=hf_quantizer,
        keep_in_fp32_modules=keep_in_fp32_modules,
        dduf_entries=dduf_entries,
        loaded_keys=loaded_keys,
        unexpected_keys=unexpected_keys,
        offload_index=offload_index,
        offload_folder=offload_folder,
        state_dict_index=state_dict_index,
        state_dict_folder=state_dict_folder,
        ignore_mismatched_sizes=ignore_mismatched_sizes,
        low_cpu_mem_usage=low_cpu_mem_usage,
    )

    with ThreadPoolExecutor(max_workers=num_workers) as executor: