mirrors/stable-diffusion-webui
1
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2026-01-08 02:02:36 -08:00
Code Issues Projects Releases Packages Wiki Activity

send weights to target device instead of CPU memory

Browse source
This commit is contained in:
AUTOMATIC1111 2023-08-16 12:11:01 +03:00
parent 57e59c14c8
commit eaba3d7349
2 changed files with 31 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

17
modules/sd_models.py
View file

@ -518,6 +518,13 @@ def send_model_to_cpu(m):
devices.torch_gc()
def model_target_device():
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
return devices.cpu
else:
return devices.device
def send_model_to_device(m):
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.setup_for_low_vram(m, shared.cmd_opts.medvram)
@ -579,7 +586,15 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
timer.record("create model")
with sd_disable_initialization.LoadStateDictOnMeta(state_dict, devices.cpu):
if shared.cmd_opts.no_half:
weight_dtype_conversion = None
else:
weight_dtype_conversion = {
'first_stage_model': None,
'': torch.float16,
}
with sd_disable_initialization.LoadStateDictOnMeta(state_dict, device=model_target_device(), weight_dtype_conversion=weight_dtype_conversion):
load_model_weights(sd_model, checkpoint_info, state_dict, timer)
timer.record("load weights from state dict")