From 2f2eb2bd74ac18650106edf2aef1e77150c2be58 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 18 Aug 2025 20:46:54 -0400
Subject: [PATCH] Fix.

---
 comfy/model_base.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index 15bd7abef..6c861b15e 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -1325,6 +1325,7 @@ class Omnigen2(BaseModel):
 class QwenImage(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)
+        self.memory_usage_factor_conds = ("ref_latents",)
 
     def extra_conds(self, **kwargs):
         out = super().extra_conds(**kwargs)
@@ -1342,3 +1343,10 @@ class QwenImage(BaseModel):
             if ref_latents_method is not None:
                 out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
         return out
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
+        return out