Skip to content

Commit 2c764e9

Browse files
authored
Model loading speed optimization (huggingface#5635)
Move unchanging operation out of loop for speed benefit.
1 parent 9c299f1 commit 2c764e9

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

models/modeling_utils.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,8 @@ def load_model_dict_into_meta(model, state_dict, device=None, dtype=None, model_
134134
device = device or torch.device("cpu")
135135
dtype = dtype or torch.float32
136136

137+
accepts_dtype = "dtype" in set(inspect.signature(set_module_tensor_to_device).parameters.keys())
138+
137139
unexpected_keys = []
138140
empty_state_dict = model.state_dict()
139141
for param_name, param in state_dict.items():
@@ -147,7 +149,6 @@ def load_model_dict_into_meta(model, state_dict, device=None, dtype=None, model_
147149
f"Cannot load {model_name_or_path_str}because {param_name} expected shape {empty_state_dict[param_name]}, but got {param.shape}. If you want to instead overwrite randomly initialized weights, please make sure to pass both `low_cpu_mem_usage=False` and `ignore_mismatched_sizes=True`. For more information, see also: https://github.com/huggingface/diffusers/issues/1619#issuecomment-1345604389 as an example."
148150
)
149151

150-
accepts_dtype = "dtype" in set(inspect.signature(set_module_tensor_to_device).parameters.keys())
151152
if accepts_dtype:
152153
set_module_tensor_to_device(model, param_name, device, value=param, dtype=dtype)
153154
else:

0 commit comments

Comments
 (0)