theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-19 23:08:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

accelerate load model

Browse Source
This commit is contained in:
tc2000731
2024-10-18 15:29:50 +08:00
parent 7d7d72dcfe
commit dfbf43e463
3 changed files with 58 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
diffsynth/models/model_manager.py
View File

@@ -50,7 +50,7 @@ from ..extensions.RIFE import IFNet
from ..extensions.ESRGAN import RRDBNet
from ..configs.model_config import model_loader_configs, huggingface_model_loader_configs, patch_model_loader_configs
from .utils import load_state_dict
from .utils import load_state_dict, init_weights_on_device
@@ -106,8 +106,10 @@ def load_model_from_single_file(state_dict, model_names, model_classes, model_re
else:
model_state_dict, extra_kwargs = state_dict_results, {}
torch_dtype = torch.float32 if extra_kwargs.get("upcast_to_float32", False) else torch_dtype
model = model_class(**extra_kwargs).to(dtype=torch_dtype, device=device)
model.load_state_dict(model_state_dict)
with init_weights_on_device():
model= model_class(**extra_kwargs)
model.load_state_dict(model_state_dict, assign=True)
model = model.to(dtype=torch_dtype, device=device)
loaded_model_names.append(model_name)
loaded_models.append(model)
return loaded_model_names, loaded_models