Fix VAEEncode with taesd3.

NeoWorldTeam · Jun 16, 2024 · 746a041 · 746a041
1 parent 04e8798
commit 746a041
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/comfy/sd.py b/comfy/sd.py
@@ -189,7 +189,8 @@ def __init__(self, sd=None, device=None, config=None, dtype=None):
                                                             encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config},
                                                             decoder_config={'target': "comfy.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config})
             elif "taesd_decoder.1.weight" in sd:
-                self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=sd["taesd_decoder.1.weight"].shape[1])
+                self.latent_channels = sd["taesd_decoder.1.weight"].shape[1]
+                self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=self.latent_channels)
             elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade
                 self.first_stage_model = StageA()
                 self.downscale_ratio = 4