VAE patcher and more types of unet patches

This commit is contained in:
lllyasviel
2024-02-29 22:37:34 -08:00
committed by GitHub
parent b59deaa382
commit ef35383b4a
6 changed files with 81 additions and 10 deletions

View File

@@ -235,14 +235,14 @@ def load_model_for_a1111(timer, checkpoint_info=None, state_dict=None):
@torch.inference_mode()
def patched_decode_first_stage(x):
sample = forge_objects.unet.model.model_config.latent_format.process_out(x)
sample = forge_objects.vae.decode(sample).movedim(-1, 1) * 2.0 - 1.0
sample = sd_model.forge_objects.unet.model.model_config.latent_format.process_out(x)
sample = sd_model.forge_objects.vae.decode(sample).movedim(-1, 1) * 2.0 - 1.0
return sample.to(x)
@torch.inference_mode()
def patched_encode_first_stage(x):
sample = forge_objects.vae.encode(x.movedim(1, -1) * 0.5 + 0.5)
sample = forge_objects.unet.model.model_config.latent_format.process_in(sample)
sample = sd_model.forge_objects.vae.encode(x.movedim(1, -1) * 0.5 + 0.5)
sample = sd_model.forge_objects.unet.model.model_config.latent_format.process_in(sample)
return sample.to(x)
sd_model.ema_scope = lambda *args, **kwargs: contextlib.nullcontext()