Update forge_loader.py

lllyasviel · Feb 29, 2024 · 28ccd3d · 28ccd3d
1 parent 0799e8b
commit 28ccd3d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/modules_forge/forge_loader.py b/modules_forge/forge_loader.py
@@ -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()