From b6fefe686beba2e550e8556bb86113a5ddef9318 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 26 Feb 2025 07:51:22 -0500
Subject: [PATCH] Better wan memory estimation.

---
 comfy/supported_models.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index e28bd138..a8212c1f 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -916,6 +916,10 @@ class WAN21_T2V(supported_models_base.BASE):
     vae_key_prefix = ["vae."]
     text_encoder_key_prefix = ["text_encoders."]
 
+    def __init__(self, unet_config):
+        super().__init__(unet_config)
+        self.memory_usage_factor = self.unet_config.get("dim", 2000) / 2000
+
     def get_model(self, state_dict, prefix="", device=None):
         out = model_base.WAN21(self, device=device)
         return out