From fcafa25a31047ce0cf83ca79eff2eabcedde2e6f Mon Sep 17 00:00:00 2001
From: fluffydiveX <tia21sb008f4skosul1afp@gmail.com>
Date: Sun, 30 Mar 2025 12:30:12 +0900
Subject: [PATCH] Revert "add wan blockswap node"

This reverts commit 3022d1b6c59097213ee06c8070ecbc7de5c63f2e.
---
 nodes.py | 46 +++-------------------------------------------
 1 file changed, 3 insertions(+), 43 deletions(-)

diff --git a/nodes.py b/nodes.py
index 2b7f89e..29ac464 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2,7 +2,7 @@ import comfy.model_management
 import gc
 from comfy.patcher_extension import CallbacksMP
 from comfy.model_patcher import ModelPatcher
-from comfy.model_base import HunyuanVideo, WAN21
+from comfy.model_base import HunyuanVideo
 
 #Based on https://github.com/kijai/ComfyUI-HunyuanVideoWrapper
 class HunyuanVideoBlockSwap:
@@ -48,50 +48,10 @@ class HunyuanVideoBlockSwap:
 
         return (model, )
 
-class WanVideoBlockSwap:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "blocks_to_swap": ("INT", {"default": 40, "min": 0, "max": 40, "step": 1, "tooltip": "Number of blocks to swap"}),
-                "offload_txt_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload text_embedding layer"}),
-                "offload_img_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload img_emb layer"}),
-            },
-        }
-    RETURN_TYPES = ("MODEL",)
-    CATEGORY = "ComfyUI-hvBlockswap"
-    FUNCTION = "set_callback"
-
-    def set_callback(self, model: ModelPatcher, blocks_to_swap, offload_txt_emb, offload_img_emb):
-        
-        def swap_blocks(model: ModelPatcher, device_to, lowvram_model_memory, force_patch_weights, full_load):
-            base_model = model.model
-            if isinstance(base_model, WAN21):
-                unet = base_model.diffusion_model
-                for b, block in enumerate(unet.blocks):
-                    if b < blocks_to_swap:
-                        block.to(model.offload_device)
-
-                if offload_txt_emb:
-                    unet.text_embedding.to(model.offload_device)
-                if offload_img_emb:
-                    unet.img_emb.to(model.offload_device)
-
-            comfy.model_management.soft_empty_cache()
-            gc.collect()
-        
-        model = model.clone()
-        model.add_callback(CallbacksMP.ON_LOAD,swap_blocks)
-
-        return (model, )
-    
 NODE_CLASS_MAPPINGS = {
-    "hvBlockSwap": HunyuanVideoBlockSwap,
-    "WanBlockSwap": WanVideoBlockSwap
+    "hvBlockSwap": HunyuanVideoBlockSwap
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "hvBlockSwap": "HunyuanVideoBlockSwap",
-    "WanBlockSwap": "WanVideoBlockSwap"
+    "hvBlockSwap": "HunyuanVideoBlockSwap"
 }
\ No newline at end of file