correcting implementation of salvagable nodes and obiliterating one abomination

Author: pollockjj

__init__.py

@@ -226,11 +226,11 @@ def text_encoder_device_patched():
     WanVideoEncode,
     LoadWanVideoClipTextEncoder,
     WanVideoClipVisionEncode,
-    WanVideoControlnetLoaderMultiGPU,
-    FantasyTalkingModelLoaderMultiGPU,
-    Wav2VecModelLoaderMultiGPU,
-    WanVideoUni3C_ControlnetLoaderMultiGPU,
-    DownloadAndLoadWav2VecModelMultiGPU,
+    WanVideoControlnetLoader,
+    FantasyTalkingModelLoader,
+    Wav2VecModelLoader,
+    WanVideoUni3C_ControlnetLoader,
+    DownloadAndLoadWav2VecModel,
 )
 
 from .wrappers import (
@@ -382,11 +382,11 @@ def register_and_count(module_names, node_map):
     "WanVideoEncodeMultiGPU": WanVideoEncode,
     "LoadWanVideoClipTextEncoderMultiGPU": LoadWanVideoClipTextEncoder,
     "WanVideoClipVisionEncodeMultiGPU": WanVideoClipVisionEncode,
-    "WanVideoControlnetLoaderMultiGPU": WanVideoControlnetLoaderMultiGPU,
-    "FantasyTalkingModelLoaderMultiGPU": FantasyTalkingModelLoaderMultiGPU,
-    "Wav2VecModelLoaderMultiGPU": Wav2VecModelLoaderMultiGPU,
-    "WanVideoUni3C_ControlnetLoaderMultiGPU": WanVideoUni3C_ControlnetLoaderMultiGPU,
-    "DownloadAndLoadWav2VecModelMultiGPU": DownloadAndLoadWav2VecModelMultiGPU,
+    "WanVideoControlnetLoaderMultiGPU": WanVideoControlnetLoader,
+    "FantasyTalkingModelLoaderMultiGPU": FantasyTalkingModelLoader,
+    "Wav2VecModelLoaderMultiGPU": Wav2VecModelLoader,
+    "WanVideoUni3C_ControlnetLoaderMultiGPU": WanVideoUni3C_ControlnetLoader,
+    "DownloadAndLoadWav2VecModelMultiGPU": DownloadAndLoadWav2VecModel,
 }
 register_and_count(["ComfyUI-WanVideoWrapper", "comfyui-wanvideowrapper"], wanvideo_nodes)
 

wanvideo.py

@@ -714,116 +714,6 @@ def process(self, clip_vision, load_device, image_1, strength_1, strength_2, for
             encode_module.device = original_module_device
 
 class WanVideoControlnetLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": (folder_paths.get_filename_list("controlnet"), {"tooltip": "These models are loaded from the 'ComfyUI/models/controlnet' -folder",}),
-
-            "base_precision": (["fp32", "bf16", "fp16"], {"default": "bf16"}),
-            "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2', 'fp8_e4m3fn_fast_no_ffn'], {"default": 'disabled', "tooltip": "optional quantization method"}),
-            "load_device": (["main_device", "offload_device"], {"default": "main_device", "tooltip": "Initial device to load the model to, NOT recommended with the larger models unless you have 48GB+ VRAM"}),
-            },
-        }
-
-    RETURN_TYPES = ("WANVIDEOCONTROLNET",)
-    RETURN_NAMES = ("controlnet", )
-    FUNCTION = "loadmodel"
-    CATEGORY = "WanVideoWrapper"
-    DESCRIPTION = "Loads ControlNet model from 'https://huggingface.co/collections/TheDenk/wan21-controlnets-68302b430411dafc0d74d2fc'"
-
-    def loadmodel(self, model, base_precision, load_device, quantization):
-
-        device = mm.get_torch_device()
-        offload_device = mm.unet_offload_device()
-
-        transformer_load_device = device if load_device == "main_device" else offload_device
-        
-        base_dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "fp8_e4m3fn_fast": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp16_fast": torch.float16, "fp32": torch.float32}[base_precision]
-        
-        model_path = folder_paths.get_full_path_or_raise("controlnet", model)
-      
-        sd = load_torch_file(model_path, device=transformer_load_device, safe_load=True)
-        
-        num_layers = 8 if "blocks.7.scale_shift_table" in sd else 6
-        out_proj_dim = sd["controlnet_blocks.0.bias"].shape[0]
-        downscale_coef = 16 if out_proj_dim == 3072 else 8
-        vae_channels = 48 if out_proj_dim == 3072 else 16
-
-        if not "control_encoder.0.0.weight" in sd:
-            raise ValueError("Invalid ControlNet model")
-
-        controlnet_cfg = {
-            "added_kv_proj_dim": None,
-            "attention_head_dim": 128,
-            "cross_attn_norm": None,
-            "downscale_coef": downscale_coef,
-            "eps": 1e-06,
-            "ffn_dim": 8960,
-            "freq_dim": 256,
-            "image_dim": None,
-            "in_channels": 3,
-            "num_attention_heads": 12,
-            "num_layers": num_layers,
-            "out_proj_dim": out_proj_dim,
-            "patch_size": [
-                1,
-                2,
-                2
-            ],
-            "qk_norm": "rms_norm_across_heads",
-            "rope_max_seq_len": 1024,
-            "text_dim": 4096,
-            "vae_channels": vae_channels
-            }
-        print(f"Loading WanControlnet with config: {controlnet_cfg}")
-        
-        from .wan_controlnet import WanControlnet
-
-        with init_empty_weights():
-            controlnet = WanControlnet(**controlnet_cfg)
-        controlnet.eval()
-        
-        if quantization == "disabled":
-            for k, v in sd.items():
-                if isinstance(v, torch.Tensor):
-                    if v.dtype == torch.float8_e4m3fn:
-                        quantization = "fp8_e4m3fn"
-                        break
-                    elif v.dtype == torch.float8_e5m2:
-                        quantization = "fp8_e5m2"
-                        break
-
-        if "fp8_e4m3fn" in quantization:
-            dtype = torch.float8_e4m3fn
-        elif quantization == "fp8_e5m2":
-            dtype = torch.float8_e5m2
-        else:
-            dtype = base_dtype
-        params_to_keep = {"norm", "head", "time_in", "vector_in", "controlnet_patch_embedding", "time_", "img_emb", "modulation", "text_embedding", "adapter"}
-    
-        log.info("Using accelerate to load and assign controlnet model weights to device...")
-        param_count = sum(1 for _ in controlnet.named_parameters())
-        for name, param in tqdm(controlnet.named_parameters(), 
-                desc=f"Loading transformer parameters to {transformer_load_device}", 
-                total=param_count,
-                leave=True):
-            dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype
-            if "controlnet_patch_embedding" in name:
-                dtype_to_use = torch.float32
-            set_module_tensor_to_device(controlnet, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name])
-        
-        del sd
-
-        if load_device == "offload_device" and controlnet.device != offload_device:
-            log.info(f"Moving controlnet model from {controlnet.device} to {offload_device}")
-            controlnet.to(offload_device)
-            gc.collect()
-            mm.soft_empty_cache()
-
-        return (controlnet,)
-
-class WanVideoControlnetLoaderMultiGPU:
     @classmethod
     def INPUT_TYPES(s):
         devices = get_device_list()
@@ -852,7 +742,7 @@ def loadmodel(self, model, base_precision, load_device, quantization, device):
         original_loader = NODE_CLASS_MAPPINGS["WanVideoControlnetLoader"]()
         return original_loader.loadmodel(model, base_precision, load_device, quantization)
 
-class FantasyTalkingModelLoaderMultiGPU:
+class FantasyTalkingModelLoader:
     @classmethod
     def INPUT_TYPES(s):
         devices = get_device_list()
@@ -879,7 +769,7 @@ def loadmodel(self, model, base_precision, device):
         original_loader = NODE_CLASS_MAPPINGS["FantasyTalkingModelLoader"]()
         return original_loader.loadmodel(model, base_precision)
 
-class Wav2VecModelLoaderMultiGPU:
+class Wav2VecModelLoader:
     @classmethod
     def INPUT_TYPES(s):
         devices = get_device_list()
@@ -907,7 +797,7 @@ def loadmodel(self, model, base_precision, load_device, device):
         original_loader = NODE_CLASS_MAPPINGS["Wav2VecModelLoader"]()
         return original_loader.loadmodel(model, base_precision, load_device)
 
-class DownloadAndLoadWav2VecModelMultiGPU:
+class DownloadAndLoadWav2VecModel:
     @classmethod
     def INPUT_TYPES(s):
         devices = get_device_list()