From 4acd51abab786afea793e0cce0e597493866c9c3 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 27 Nov 2025 19:27:53 +0200
Subject: [PATCH] Handle latent dim difference for image model in the VAE
 instead

---
 comfy/model_base.py | 3 ---
 comfy/sd.py         | 2 ++
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index 5ae9ca340..50872601e 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -1690,6 +1690,3 @@ class Kandinsky5_image(Kandinsky5):
 
     def concat_cond(self, **kwargs):
         return None
-
-    def process_latent_out(self, latent): # input is still 5D, return single frame to decode with Flux VAE
-        return self.latent_format.process_out(latent)[:, :, 0]
diff --git a/comfy/sd.py b/comfy/sd.py
index 0bda30fdb..381c20289 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -742,6 +742,8 @@ class VAE:
         self.throw_exception_if_invalid()
         pixel_samples = None
         do_tile = False
+        if self.latent_dim == 2 and samples_in.ndim == 5:
+            samples_in = samples_in[:, :, 0]
         try:
             memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
             model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)