From 4e5ba653c69522affca4b74885a3ec6878de60a9 Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Wed, 24 Jan 2024 11:03:36 -0800 Subject: [PATCH] i --- modules/sd_models.py | 6 ++++++ modules/sd_samplers_common.py | 3 +-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index 9910af3d..42729bf5 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -653,6 +653,12 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None): timer.record("create unet patcher") del state_dict_for_forge + def patched_decode_first_stage(sample): + sample = unet_patcher.model.model_config.latent_format.process_out(sample) + return vae_patcher.decode(sample).movedim(-1, 1) * 2.0 - 1.0 + + sd_model.decode_first_stage = patched_decode_first_stage + load_model_weights(sd_model, checkpoint_info, state_dict_for_a1111, timer) del state_dict_for_a1111 timer.record("load weights from state dict") diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index c45db906..7d5b360b 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -52,8 +52,7 @@ def samples_to_images_tensor(sample, approximation=None, model=None): else: if model is None: model = shared.sd_model - sample = model.unet_patcher.model.model_config.latent_format.process_out(sample) - x_sample = model.vae_patcher.decode(sample).movedim(-1, 1) * 2.0 - 1.0 + x_sample = model.decode_first_stage(sample) return x_sample