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

Code before:
            # load_state_dict will manage the case where we pass a dict instead of a file
            # if state dict is not None, it means that we don't need to read the files from resolved_model_file also
            resolved_model_file = [state_dict]

        # prepare the arguments.
        args_list = [
            (
                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,
            )
            for shard_file in resolved_model_file
        ]

        if is_parallel_loading_enabled:
            offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_shard_files_with_threadpool(
                args_list
            )
            error_msgs += _error_msgs
            mismatched_keys += _mismatched_keys
        else:
            if len(args_list) > 1:
                args_list = logging.tqdm(args_list, desc="Loading checkpoint shards")

            for args in args_list:
                offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_shard_file(args)
                error_msgs += _error_msgs
                mismatched_keys += _mismatched_keys

        empty_device_cache()

        if offload_index is not None and len(offload_index) > 0:
            save_offload_index(offload_index, offload_folder)
            offload_index = None

            if offload_state_dict:
                load_offloaded_weights(model, state_dict_index, state_dict_folder)
                shutil.rmtree(state_dict_folder)

Review comment (suggestion):
Since the same arguments are used across the two loading functions, it's a good candidate for `functools.partial`

```python
        load_fn = partial(
            load_shard_files_with_threadpool if is_parallel_loading_enabled else 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,
        )

        if is_parallel_loading_enabled:
            offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(
                resolved_model_file,
            )
            error_msgs += _error_msgs
            mismatched_keys += _mismatched_keys
        else:
            shard_files = resolved_model_file
            if len(resolved_model_file) > 1:
                shard_files = logging.tqdm(resolved_model_file, desc="Loading checkpoint shards")

            for shard_file in resolved_model_file:
                offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(shard_file)
                error_msgs += _error_msgs
                mismatched_keys += _mismatched_keys
```

Code after:

        if state_dict is not None:
            # load_state_dict will manage the case where we pass a dict instead of a file
            # if state dict is not None, it means that we don't need to read the files from resolved_model_file also
            resolved_model_file = [state_dict]

        # Prepare the loading function sharing the attributes shared between them.
        load_fn = functools.partial(
            _load_shard_files_with_threadpool if is_parallel_loading_enabled else _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,
        )

        if is_parallel_loading_enabled:
            offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(resolved_model_file)
            error_msgs += _error_msgs
            mismatched_keys += _mismatched_keys
        else:
            shard_files = resolved_model_file
            if len(resolved_model_file) > 1:
                shard_files = logging.tqdm(resolved_model_file, desc="Loading checkpoint shards")

            for shard_file in shard_files:
                offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(shard_file)
                error_msgs += _error_msgs
                mismatched_keys += _mismatched_keys

        empty_device_cache()

        if offload_index is not None and len(offload_index) > 0:
            save_offload_index(offload_index, offload_folder)
            offload_index = None

            if offload_state_dict:
                load_offloaded_weights(model, state_dict_index, state_dict_folder)
                shutil.rmtree(state_dict_folder)

        if len(error_msgs) > 0: