From 8c63b6d4127d9c4050241de1d486808fa8ad8c46 Mon Sep 17 00:00:00 2001
From: jiangchengchengark <3306065226@qq.com>
Date: Tue, 9 Dec 2025 13:10:16 +0800
Subject: [PATCH] Delete the qwen-image vae check branch and use shape[2] to
 perform memory calculation

---
 comfy/sd.py | 23 +++++------------------
 1 file changed, 5 insertions(+), 18 deletions(-)

diff --git a/comfy/sd.py b/comfy/sd.py
index 11699f649..bb0ca8923 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -536,22 +536,7 @@ class VAE:
                     self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
                     self.memory_used_encode = lambda shape, dtype: 3300 * shape[3] * shape[4] * model_management.dtype_size(dtype)
                     self.memory_used_decode = lambda shape, dtype: 8000 * shape[3] * shape[4] * (16 * 16) * model_management.dtype_size(dtype)
-                elif "qwen_image" in sd:  # Qwen-Image-Edit/Qwen-Image
-                    logging.info("qwen-image tag detected, loaded vae in qwen-image memory_used.")
-                    self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
-                    self.upscale_index_formula = (4, 8, 8)
-                    self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
-                    self.downscale_index_formula = (4, 8, 8)
-                    self.latent_dim = 3
-                    self.latent_channels = 16
-                    ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
-                    self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
-                    self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
-                    # One Picture Minimum limit value 1496.23
-                    self.memory_used_encode = lambda shape, dtype: 1500 * shape[3] * shape[4] * model_management.dtype_size(dtype)
-                    # One Picture Minimum limit value 2164
-                    self.memory_used_decode = lambda shape, dtype: 2200 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
-                else:  # Wan 2.1 VAE
+                else:  # Wan 2.1 VAE 
                     dim = sd["decoder.head.0.gamma"].shape[0]
                     self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
                     self.upscale_index_formula = (4, 8, 8)
@@ -562,8 +547,10 @@ class VAE:
                     ddconfig = {"dim": dim, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
                     self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
                     self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
-                    self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
-                    self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
+                    self.memory_used_encode = lambda shape, dtype: (1500 if shape[2]<=4 else 6000) * shape[3] * shape[4] * model_management.dtype_size(dtype)
+                    self.memory_used_decode = lambda shape, dtype: (2050 if shape[2]<=4 else 7000) * shape[3] * shape[4] * (8*8) * model_management.dtype_size(dtype)
+
+
             # Hunyuan 3d v2 2.0 & 2.1
             elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd: