Protect alphas_cumprod from downcasting

This commit is contained in:
drhead 2023-11-29 17:38:53 -05:00 committed by GitHub
parent f0f100e67b
commit b25c126ccd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions

View File

@ -387,7 +387,11 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
if shared.cmd_opts.upcast_sampling and depth_model: if shared.cmd_opts.upcast_sampling and depth_model:
model.depth_model = None model.depth_model = None
alphas_cumprod = model.alphas_cumprod
model.alphas_cumprod = None
model.half() model.half()
model.alphas_cumprod = alphas_cumprod
model.alphas_cumprod_original = alphas_cumprod
model.first_stage_model = vae model.first_stage_model = vae
if depth_model: if depth_model:
model.depth_model = depth_model model.depth_model = depth_model
@ -642,6 +646,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None):
else: else:
weight_dtype_conversion = { weight_dtype_conversion = {
'first_stage_model': None, 'first_stage_model': None,
'alphas_cumprod': None,
'': torch.float16, '': torch.float16,
} }