Add block swap debug prints and cleanup unused code

Author: kijai

nodes.py

@@ -39,7 +39,7 @@ class WanVideoBlockSwap:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "blocks_to_swap": ("INT", {"default": 20, "min": 0, "max": 40, "step": 1, "tooltip": "Number of double blocks to swap"}),
+                "blocks_to_swap": ("INT", {"default": 20, "min": 0, "max": 40, "step": 1, "tooltip": "Number of transformer blocks to swap, the 14B model has 40, while the 1.3B model has 30 blocks"}),
                 "offload_img_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload img_emb to offload_device"}),
                 "offload_txt_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload time_emb to offload_device"}),
             },

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-WanVideoWrapper"
 description = "ComfyUI diffusers wrapper nodes for WanVideo"
-version = "1.0.3"
+version = "1.0.4"
 license = {file = "LICENSE"}
 dependencies = ["accelerate >= 1.2.1", "diffusers >= 0.31.0", "ftfy"]
 

utils.py

@@ -21,4 +21,11 @@ def print_memory(device):
     log.info(f"Max allocated memory: {max_memory=:.3f} GB")
     log.info(f"Max reserved memory: {max_reserved=:.3f} GB")
     #memory_summary = torch.cuda.memory_summary(device=device, abbreviated=False)
-    #log.info(f"Memory Summary:\n{memory_summary}")
+    #log.info(f"Memory Summary:\n{memory_summary}")
+
+def get_module_memory_mb(module):
+    memory = 0
+    for param in module.parameters():
+        if param.data is not None:
+            memory += param.nelement() * param.element_size()
+    return memory / (1024 * 1024)  # Convert to MB

wanvideo/modules/clip.py

@@ -9,8 +9,6 @@
 import torchvision.transforms as T
 
 from .attention import attention
-from .tokenizers import HuggingfaceTokenizer
-from .xlm_roberta import XLMRoberta
 
 __all__ = [
     'XLMRobertaCLIP',
@@ -155,60 +153,6 @@ def forward(self, x):
             x = x + self.mlp(self.norm2(x))
         return x
 
-
-class AttentionPool(nn.Module):
-
-    def __init__(self,
-                 dim,
-                 mlp_ratio,
-                 num_heads,
-                 activation='gelu',
-                 proj_dropout=0.0,
-                 norm_eps=1e-5):
-        assert dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.mlp_ratio = mlp_ratio
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.proj_dropout = proj_dropout
-        self.norm_eps = norm_eps
-
-        # layers
-        gain = 1.0 / math.sqrt(dim)
-        self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
-        self.to_q = nn.Linear(dim, dim)
-        self.to_kv = nn.Linear(dim, dim * 2)
-        self.proj = nn.Linear(dim, dim)
-        self.norm = LayerNorm(dim, eps=norm_eps)
-        self.mlp = nn.Sequential(
-            nn.Linear(dim, int(dim * mlp_ratio)),
-            QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
-            nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
-
-    def forward(self, x):
-        """
-        x:  [B, L, C].
-        """
-        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
-
-        # compute query, key, value
-        q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)
-        k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)
-
-        # compute attention
-        x = flash_attention(q, k, v, version=2)
-        x = x.reshape(b, 1, c)
-
-        # output
-        x = self.proj(x)
-        x = F.dropout(x, self.proj_dropout, self.training)
-
-        # mlp
-        x = x + self.mlp(self.norm(x))
-        return x[:, 0]
-
-
 class VisionTransformer(nn.Module):
 
     def __init__(self,
@@ -275,9 +219,6 @@ def __init__(self,
             self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
         elif pool_type == 'token_fc':
             self.head = nn.Linear(dim, out_dim)
-        elif pool_type == 'attn_pool':
-            self.head = AttentionPool(dim, mlp_ratio, num_heads, activation,
-                                      proj_dropout, norm_eps)
 
     def forward(self, x, interpolation=False, use_31_block=False):
         b = x.size(0)
@@ -303,31 +244,6 @@ def forward(self, x, interpolation=False, use_31_block=False):
             return x
 
 
-class XLMRobertaWithHead(XLMRoberta):
-
-    def __init__(self, **kwargs):
-        self.out_dim = kwargs.pop('out_dim')
-        super().__init__(**kwargs)
-
-        # head
-        mid_dim = (self.dim + self.out_dim) // 2
-        self.head = nn.Sequential(
-            nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(),
-            nn.Linear(mid_dim, self.out_dim, bias=False))
-
-    def forward(self, ids):
-        # xlm-roberta
-        x = super().forward(ids)
-
-        # average pooling
-        mask = ids.ne(self.pad_id).unsqueeze(-1).to(x)
-        x = (x * mask).sum(dim=1) / mask.sum(dim=1)
-
-        # head
-        x = self.head(x)
-        return x
-
-
 class XLMRobertaCLIP(nn.Module):
 
     def __init__(self,

wanvideo/modules/model.py

@@ -15,7 +15,7 @@
 
 from tqdm import tqdm
 
-from ...utils import log
+from ...utils import log, get_module_memory_mb
 
 def poly1d(coefficients, x):
     result = torch.zeros_like(x)
@@ -547,12 +547,27 @@ def block_swap(self, blocks_to_swap, offload_txt_emb=False, offload_img_emb=Fals
         self.blocks_to_swap = blocks_to_swap
         self.offload_img_emb = offload_img_emb
         self.offload_txt_emb = offload_txt_emb
+
+        total_offload_memory = 0
+        total_main_memory = 0
 
         for b, block in tqdm(enumerate(self.blocks), total=len(self.blocks), desc="Initializing block swap"):
+            block_memory = get_module_memory_mb(block)
+            
             if b > self.blocks_to_swap:
                 block.to(self.main_device)
+                total_main_memory += block_memory
             else:
                 block.to(self.offload_device)
+                total_offload_memory += block_memory
+                
+            #print(f"Block {b}: {block_memory:.2f}MB on {block.parameters().__next__().device}")
+        log.info("----------------------")
+        log.info(f"Block swap memory summary:")
+        log.info(f"Transformer blocks on {self.offload_device}: {total_offload_memory:.2f}MB")
+        log.info(f"Transformer blocks on {self.main_device}: {total_main_memory:.2f}MB")
+        log.info(f"Total Memory: {(total_offload_memory + total_main_memory):.2f}MB")
+        log.info("----------------------")
 
     def forward(
         self,