Support for kijai's ComfyUI-WanVideoWrapper

README.md

@@ -122,9 +122,12 @@ Currently supported nodes (automatically detected if available):
   - CheckpointLoaderNF4MultiGPU
 - HunyuanVideoWrapper (requires [ComfyUI-HunyuanVideoWrapper](https://github.com/kijai/ComfyUI-HunyuanVideoWrapper)):
   - HyVideoModelLoaderMultiGPU
-  - HyVideoModelLoaderDiffSynthMultiGPU (**NEW** - MultiGPU-specific node for offloading to an `offload_device` using MultiGPU's device selectors)
   - HyVideoVAELoaderMultiGPU
   - DownloadAndLoadHyVideoTextEncoderMultiGPU
+- WanVideoWrapper (requires [ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper)):
+  - WanVideoModelLoader
+  - WanVideoVAELoader
+  - LoadWanVideoT5TextEncoder
 - **Native to ComfyUI-MultiGPU**
   - DeviceSelectorMultiGPU - Allows user to link loaders together to use the same selected device
   - HunyuanVideoEmbeddingsAdapter - Allows Kijai's excellent IP2V CLIP for HunyuanVideo to be used with Comfy Core sampler.
@@ -146,6 +149,11 @@ This workflow attaches a HunyuanVideo GGUF-quantized model on `cuda:0` for compu
 - [examples/flux1dev_gguf_distorch.json](https://github.com/pollockjj/ComfyUI-MultiGPU/blob/main/examples/flux1dev_gguf_distorch.json)
 This workflow loads a FLUX.1-dev model on `cuda:0` for compute and distrubutes its UNet across multiple CUDA devices using new DisTorch distributed-load methodology. While the text encoders and VAE are loaded on GPU 1 and use `cuda:1` for compute.
 
+### Split Wan Video generation across multiple resources
+
+- [examples/hunyuanvideowrapper_native_vae.json](https://github.com/pollockjj/ComfyUI-MultiGPU/blob/main/examples/wanvideo_T2V_example_MultiGPU.json)
+This workflow is a simple extension of [kijai's T2V example](https://github.com/kijai/ComfyUI-WanVideoWrapper/blob/main/example_workflows/wanvideo_T2V_example_02.json) from his custom_node.
+
 ### Split Hunyuan Video generation across multiple resources
 
 - [examples/hunyuanvideowrapper_native_vae.json](https://github.com/pollockjj/ComfyUI-MultiGPU/blob/main/examples/hunyuanvideowrapper_native_vae.json)

__init__.py

@@ -27,7 +27,8 @@ from .nodes import (
     LoadFluxControlNet,
     MMAudioModelLoader, MMAudioFeatureUtilsLoader, MMAudioSampler,
     PulidModelLoader, PulidInsightFaceLoader, PulidEvaClipLoader,
-    HyVideoModelLoader, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder
+    HyVideoModelLoader, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder,
+    WanVideoModelLoader, WanVideoVAELoader, LoadWanVideoT5TextEncoder
 )
 
 current_device = mm.get_torch_device()
@@ -735,4 +736,10 @@ if check_module_exists("ComfyUI-HunyuanVideoWrapper") or check_module_exists("co
     NODE_CLASS_MAPPINGS["HyVideoVAELoaderMultiGPU"] = override_class(HyVideoVAELoader)
     NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoderMultiGPU"] = override_class(DownloadAndLoadHyVideoTextEncoder)
 
+if check_module_exists("ComfyUI-WanVideoWrapper") or check_module_exists("comfyui-wanvideowrapper"):
+    NODE_CLASS_MAPPINGS["WanVideoModelLoaderMultiGPU"] = override_class(WanVideoModelLoader)
+    NODE_CLASS_MAPPINGS["WanVideoVAELoaderMultiGPU"] = override_class(WanVideoVAELoader)
+    NODE_CLASS_MAPPINGS["LoadWanVideoT5TextEncoderMultiGPU"] = override_class(LoadWanVideoT5TextEncoder)
+
+
 logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")

nodes.py

@@ -473,4 +473,99 @@ class DownloadAndLoadHyVideoTextEncoder:
     def loadmodel(self, llm_model, clip_model, precision,  apply_final_norm=False, hidden_state_skip_layer=2, quantization="disabled"):
         from nodes import NODE_CLASS_MAPPINGS
         original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoder"]()
-        return original_loader.loadmodel(llm_model, clip_model, precision, apply_final_norm, hidden_state_skip_layer, quantization)
+        return original_loader.loadmodel(llm_model, clip_model, precision, apply_final_norm, hidden_state_skip_layer, quantization)
+class WanVideoModelLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": (folder_paths.get_filename_list("diffusion_models"),
+                          {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' folder",}),
+                "base_precision": (["fp32", "bf16", "fp16", "fp16_fast"], {"default": "bf16"}),
+                "quantization": (
+                    ['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2', 'fp8_scaled',
+                     'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6",
+                     "torchao_int4", "torchao_int8"],
+                    {"default": 'disabled', "tooltip": "optional quantization method"}
+                ),
+                "load_device": (["main_device"], {"default": "main_device"}),
+            },
+            "optional": {
+                "attention_mode": ([
+                    "sdpa",
+                    "flash_attn_2",
+                    "flash_attn_3",
+                    "sageattn",
+                ], {"default": "sdpa"}),
+                "compile_args": ("WANCOMPILEARGS", ),
+                "block_swap_args": ("BLOCKSWAPARGS", ),
+                "lora": ("WANVIDLORA", {"default": None}),
+                "vram_management_args": ("VRAM_MANAGEMENTARGS",
+                                          {"default": None, "tooltip": "Alternative offloading method"}),
+            }
+        }
+
+    RETURN_TYPES = ("WANVIDEOMODEL",)
+    RETURN_NAMES = ("model", )
+    FUNCTION = "loadmodel"
+    CATEGORY = "WanVideoWrapper"
+
+    def loadmodel(self, model, base_precision, load_device, quantization,
+                  compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None, vram_management_args=None):
+        from nodes import NODE_CLASS_MAPPINGS
+        original_loader = NODE_CLASS_MAPPINGS["WanVideoModelLoader"]()
+        return original_loader.loadmodel(model, base_precision, load_device, quantization,
+                                         compile_args, attention_mode, block_swap_args, lora, vram_management_args)
+
+
+class WanVideoVAELoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model_name": (folder_paths.get_filename_list("vae"),
+                               {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}),
+            },
+            "optional": {
+                "precision": (["fp16", "fp32", "bf16"], {"default": "bf16"}),
+            }
+        }
+
+    RETURN_TYPES = ("WANVAE",)
+    RETURN_NAMES = ("vae", )
+    FUNCTION = "loadmodel"
+    CATEGORY = "WanVideoWrapper"
+    DESCRIPTION = "Loads Wan VAE model from 'ComfyUI/models/vae'"
+
+    def loadmodel(self, model_name, precision):
+        from nodes import NODE_CLASS_MAPPINGS
+        original_loader = NODE_CLASS_MAPPINGS["WanVideoVAELoader"]()
+        return original_loader.loadmodel(model_name, precision)
+
+
+class LoadWanVideoT5TextEncoder:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model_name": (folder_paths.get_filename_list("text_encoders"),
+                               {"tooltip": "These models are loaded from 'ComfyUI/models/text_encoders'"}),
+                "precision": (["fp16", "fp32", "bf16"], {"default": "bf16"}),
+            },
+            "optional": {
+                "load_device": (["main_device"], {"default": "main_device"}),
+                "quantization": (['disabled', 'fp8_e4m3fn'],
+                                 {"default": 'disabled', "tooltip": "optional quantization method"}),
+            }
+        }
+
+    RETURN_TYPES = ("WANTEXTENCODER",)
+    RETURN_NAMES = ("wan_t5_model", )
+    FUNCTION = "loadmodel"
+    CATEGORY = "WanVideoWrapper"
+    DESCRIPTION = "Loads Wan text_encoder model from 'ComfyUI/models/LLM'"
+
+    def loadmodel(self, model_name, precision, load_device="offload_device", quantization="disabled"):
+        from nodes import NODE_CLASS_MAPPINGS
+        original_loader = NODE_CLASS_MAPPINGS["LoadWanVideoT5TextEncoder"]()
+        return original_loader.loadmodel(model_name, precision, load_device, quantization)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-multigpu"
-description = "This custom_node for ComfyUI adds one-click 'Virtual VRAM' for any GGUF UNet and CLIP loader, managing the offload of layers to DRAM or VRAM to maximize the latent space of your card. Also includes nodes for directly loading entire components (UNet, CLIP, VAE) onto the device you choose. Includes 16 examples covering common use cases."
-version = "1.6.2"
+description = "This custom_node for ComfyUI adds one-click 'Virtual VRAM' for any GGUF UNet and CLIP loader, managing the offload of layers to DRAM or VRAM to maximize the latent space of your card. Also includes nodes for directly loading entire components (UNet, CLIP, VAE) onto the device you choose. Includes 16 examples covering common use cases. Includes support for kijai's ComfyUI-WanVideoWrapper and ComfyUI-HunyuanVideoWrapper, among other popular loaders."
+version = "1.7.0"
 license = {file = "LICENSE"}
 
 [project.urls]