Execute model_loaded_callback after moving to target device

This commit is contained in:
Nuullll 2024-01-06 20:03:33 +08:00
parent b00b429477
commit a183de04e3
2 changed files with 5 additions and 4 deletions

View File

@ -842,13 +842,13 @@ def reload_model_weights(sd_model=None, info=None, forced_reload=False):
sd_hijack.model_hijack.hijack(sd_model)
timer.record("hijack")
script_callbacks.model_loaded_callback(sd_model)
timer.record("script callbacks")
if not sd_model.lowvram:
sd_model.to(devices.device)
timer.record("move model to device")
script_callbacks.model_loaded_callback(sd_model)
timer.record("script callbacks")
print(f"Weights loaded in {timer.summary()}.")
model_data.set_sd_model(sd_model)

View File

@ -273,10 +273,11 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
load_vae(sd_model, vae_file, vae_source)
sd_hijack.model_hijack.hijack(sd_model)
script_callbacks.model_loaded_callback(sd_model)
if not sd_model.lowvram:
sd_model.to(devices.device)
script_callbacks.model_loaded_callback(sd_model)
print("VAE weights loaded.")
return sd_model