John Pollock
657fdac13a
Problem: WanVideoWrapper caches device at module load time, causing timesteps and tensors to be created on wrong device when looping between models on different GPUs. Solution: WanVideoSamplerMultiGPU wrapper updates module-level device variable to match current model's device before sampling. Changes: - Added comprehensive logging to trace device allocation through pipeline - Identified module-level device caching as root cause - Simplified WanVideoSamplerMultiGPU to only update device variable - Verified fix works for multi-model workflows with looping
1036 lines
46 KiB
Python
1036 lines
46 KiB
Python
import folder_paths
|
|
from pathlib import Path
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
|
|
class UnetLoaderGGUF:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
|
|
return {
|
|
"required": {
|
|
"unet_name": (unet_names,),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_unet"
|
|
CATEGORY = "bootleg"
|
|
TITLE = "Unet Loader (GGUF)"
|
|
|
|
def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]()
|
|
return original_loader.load_unet(unet_name, dequant_dtype, patch_dtype, patch_on_device)
|
|
|
|
class UnetLoaderGGUFAdvanced(UnetLoaderGGUF):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
|
|
return {
|
|
"required": {
|
|
"unet_name": (unet_names,),
|
|
"dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}),
|
|
"patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}),
|
|
"patch_on_device": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
TITLE = "Unet Loader (GGUF/Advanced)"
|
|
|
|
|
|
class CLIPLoaderGGUF:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"clip_name": (s.get_filename_list(),),
|
|
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "wan"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP",)
|
|
FUNCTION = "load_clip"
|
|
CATEGORY = "bootleg"
|
|
TITLE = "CLIPLoader (GGUF)"
|
|
|
|
@classmethod
|
|
def get_filename_list(s):
|
|
files = []
|
|
files += folder_paths.get_filename_list("clip")
|
|
files += folder_paths.get_filename_list("clip_gguf")
|
|
return sorted(files)
|
|
|
|
def load_data(self, ckpt_paths):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]()
|
|
return original_loader.load_data(ckpt_paths)
|
|
|
|
def load_patcher(self, clip_paths, clip_type, clip_data):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]()
|
|
return original_loader.load_patcher(clip_paths, clip_type, clip_data)
|
|
|
|
def load_clip(self, clip_name, type="stable_diffusion"):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]()
|
|
return original_loader.load_clip(clip_name, type)
|
|
|
|
class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
file_options = (s.get_filename_list(), )
|
|
return {
|
|
"required": {
|
|
"clip_name1": file_options,
|
|
"clip_name2": file_options,
|
|
"type": (("sdxl", "sd3", "flux", "hunyuan_video"),),
|
|
}
|
|
}
|
|
|
|
TITLE = "DualCLIPLoader (GGUF)"
|
|
|
|
def load_clip(self, clip_name1, clip_name2, type):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUF"]()
|
|
clip = original_loader.load_clip(clip_name1, clip_name2, type)
|
|
clip[0].patcher.load(force_patch_weights=True)
|
|
return clip
|
|
|
|
|
|
class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
file_options = (s.get_filename_list(), )
|
|
return {
|
|
"required": {
|
|
"clip_name1": file_options,
|
|
"clip_name2": file_options,
|
|
"clip_name3": file_options,
|
|
}
|
|
}
|
|
|
|
TITLE = "TripleCLIPLoader (GGUF)"
|
|
|
|
def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUF"]()
|
|
return original_loader.load_clip(clip_name1, clip_name2, clip_name3, type)
|
|
|
|
class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
file_options = (s.get_filename_list(), )
|
|
return {
|
|
"required": {
|
|
"clip_name1": file_options,
|
|
"clip_name2": file_options,
|
|
"clip_name3": file_options,
|
|
"clip_name4": file_options,
|
|
}
|
|
}
|
|
|
|
TITLE = "QuadrupleCLIPLoader (GGUF)"
|
|
|
|
def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderGGUF"]()
|
|
return original_loader.load_clip(clip_name1, clip_name2, clip_name3, clip_name4, type)
|
|
|
|
|
|
class LTXVLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),
|
|
{"tooltip": "The name of the checkpoint (model) to load."}),
|
|
"dtype": (["bfloat16", "float32"], {"default": "bfloat16"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "VAE")
|
|
RETURN_NAMES = ("model", "vae")
|
|
FUNCTION = "load"
|
|
CATEGORY = "lightricks/LTXV"
|
|
TITLE = "LTXV Loader"
|
|
OUTPUT_NODE = False
|
|
|
|
def load(self, ckpt_name, dtype):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
|
return original_loader.load(ckpt_name, dtype)
|
|
def _load_unet(self, load_device, offload_device, weights, num_latent_channels, dtype, config=None ):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
|
return original_loader._load_unet(load_device, offload_device, weights, num_latent_channels, dtype, config=None )
|
|
def _load_vae(self, weights, config=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
|
return original_loader._load_vae(weights, config=None)
|
|
|
|
class Florence2ModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
|
|
"precision": (['fp16','bf16','fp32'],),
|
|
"attention": (
|
|
[ 'flash_attention_2', 'sdpa', 'eager'],
|
|
{
|
|
"default": 'sdpa'
|
|
}),
|
|
},
|
|
"optional": {
|
|
"lora": ("PEFTLORA",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FL2MODEL",)
|
|
RETURN_NAMES = ("florence2_model",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Florence2"
|
|
|
|
def loadmodel(self, model, precision, attention, lora=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["Florence2ModelLoader"]()
|
|
return original_loader.loadmodel(model, precision, attention, lora)
|
|
|
|
class DownloadAndLoadFlorence2Model:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": (
|
|
[
|
|
'microsoft/Florence-2-base',
|
|
'microsoft/Florence-2-base-ft',
|
|
'microsoft/Florence-2-large',
|
|
'microsoft/Florence-2-large-ft',
|
|
'HuggingFaceM4/Florence-2-DocVQA',
|
|
'thwri/CogFlorence-2.1-Large',
|
|
'thwri/CogFlorence-2.2-Large',
|
|
'gokaygokay/Florence-2-SD3-Captioner',
|
|
'gokaygokay/Florence-2-Flux-Large',
|
|
'MiaoshouAI/Florence-2-base-PromptGen-v1.5',
|
|
'MiaoshouAI/Florence-2-large-PromptGen-v1.5',
|
|
'MiaoshouAI/Florence-2-base-PromptGen-v2.0',
|
|
'MiaoshouAI/Florence-2-large-PromptGen-v2.0'
|
|
],
|
|
{
|
|
"default": 'microsoft/Florence-2-base'
|
|
}),
|
|
"precision": ([ 'fp16','bf16','fp32'],
|
|
{
|
|
"default": 'fp16'
|
|
}),
|
|
"attention": (
|
|
[ 'flash_attention_2', 'sdpa', 'eager'],
|
|
{
|
|
"default": 'sdpa'
|
|
}),
|
|
},
|
|
"optional": {
|
|
"lora": ("PEFTLORA",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FL2MODEL",)
|
|
RETURN_NAMES = ("florence2_model",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Florence2"
|
|
|
|
def loadmodel(self, model, precision, attention, lora=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
|
|
return original_loader.loadmodel(model, precision, attention, lora)
|
|
|
|
class CheckpointLoaderNF4:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
|
}}
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
FUNCTION = "load_checkpoint"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
|
|
def load_checkpoint(self, ckpt_name):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["CheckpointLoaderNF4"]()
|
|
return original_loader.load_checkpoint(ckpt_name)
|
|
|
|
class LoadFluxControlNet:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"model_name": (["flux-dev", "flux-dev-fp8", "flux-schnell"],),
|
|
"controlnet_path": (folder_paths.get_filename_list("xlabs_controlnets"), ),
|
|
}}
|
|
|
|
RETURN_TYPES = ("FluxControlNet",)
|
|
RETURN_NAMES = ("ControlNet",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "XLabsNodes"
|
|
|
|
def loadmodel(self, model_name, controlnet_path):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["LoadFluxControlNet"]()
|
|
return original_loader.loadmodel(model_name, controlnet_path)
|
|
|
|
class MMAudioModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mmaudio_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from the 'ComfyUI/models/mmaudio' -folder",}),
|
|
|
|
"base_precision": (["fp16", "fp32", "bf16"], {"default": "fp16"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MMAUDIO_MODEL",)
|
|
RETURN_NAMES = ("mmaudio_model", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "MMAudio"
|
|
|
|
def loadmodel(self, mmaudio_model, base_precision):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["MMAudioModelLoader"]()
|
|
return original_loader.loadmodel(mmaudio_model, base_precision)
|
|
|
|
class MMAudioFeatureUtilsLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"vae_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
|
"synchformer_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
|
"clip_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
|
},
|
|
"optional": {
|
|
"bigvgan_vocoder_model": ("VOCODER_MODEL", {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
|
"mode": (["16k", "44k"], {"default": "44k"}),
|
|
"precision": (["fp16", "fp32", "bf16"],
|
|
{"default": "fp16"}
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MMAUDIO_FEATUREUTILS",)
|
|
RETURN_NAMES = ("mmaudio_featureutils", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "MMAudio"
|
|
|
|
def loadmodel(self, vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoader"]()
|
|
return original_loader.loadmodel(vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model)
|
|
|
|
class MMAudioSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mmaudio_model": ("MMAUDIO_MODEL",),
|
|
"feature_utils": ("MMAUDIO_FEATUREUTILS",),
|
|
"duration": ("FLOAT", {"default": 8, "step": 0.01, "tooltip": "Duration of the audio in seconds"}),
|
|
"steps": ("INT", {"default": 25, "step": 1, "tooltip": "Number of steps to interpolate"}),
|
|
"cfg": ("FLOAT", {"default": 4.5, "step": 0.1, "tooltip": "Strength of the conditioning"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"prompt": ("STRING", {"default": "", "multiline": True} ),
|
|
"negative_prompt": ("STRING", {"default": "", "multiline": True} ),
|
|
"mask_away_clip": ("BOOLEAN", {"default": False, "tooltip": "If true, the clip video will be masked away"}),
|
|
"force_offload": ("BOOLEAN", {"default": True, "tooltip": "If true, the model will be offloaded to the offload device"}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
RETURN_NAMES = ("audio", )
|
|
FUNCTION = "sample"
|
|
CATEGORY = "MMAudio"
|
|
|
|
def sample(self, mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["MMAudioSampler"]()
|
|
return original_loader.sample(mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images)
|
|
|
|
class PulidModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "pulid_file": (folder_paths.get_filename_list("pulid"), )}}
|
|
|
|
RETURN_TYPES = ("PULID",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "pulid"
|
|
|
|
def load_model(self, pulid_file):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["PulidModelLoader"]()
|
|
return original_loader.load_model(pulid_file)
|
|
|
|
class PulidInsightFaceLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"provider": (["CPU", "CUDA", "ROCM", "CoreML"], ),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FACEANALYSIS",)
|
|
FUNCTION = "load_insightface"
|
|
CATEGORY = "pulid"
|
|
|
|
def load_insightface(self, provider):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["PulidInsightFaceLoader"]()
|
|
return original_loader.load_insightface(provider)
|
|
|
|
class PulidEvaClipLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {},
|
|
}
|
|
|
|
RETURN_TYPES = ("EVA_CLIP",)
|
|
FUNCTION = "load_eva_clip"
|
|
CATEGORY = "pulid"
|
|
|
|
def load_eva_clip(self):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["PulidEvaClipLoader"]()
|
|
return original_loader.load_eva_clip()
|
|
|
|
|
|
class HyVideoModelLoader:
|
|
@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"], {"default": "bf16"}),
|
|
"quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', '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_varlen",
|
|
"sageattn_varlen",
|
|
"comfy",
|
|
], {"default": "flash_attn"}),
|
|
"compile_args": ("COMPILEARGS", ),
|
|
"block_swap_args": ("BLOCKSWAPARGS", ),
|
|
"lora": ("HYVIDLORA", {"default": None}),
|
|
"auto_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "Enable auto offloading for reduced VRAM usage, implementation from DiffSynth-Studio, slightly different from block swapping and uses even less VRAM, but can be slower as you can't define how much VRAM to use"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HYVIDEOMODEL",)
|
|
RETURN_NAMES = ("model", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "HunyuanVideoWrapper"
|
|
|
|
def loadmodel(self, model, base_precision, load_device, quantization, compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None, auto_cpu_offload=False):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["HyVideoModelLoader"]()
|
|
return original_loader.loadmodel(model, base_precision, load_device, quantization, compile_args, attention_mode, block_swap_args, lora, auto_cpu_offload)
|
|
|
|
class HyVideoVAELoader:
|
|
@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"}
|
|
),
|
|
"compile_args":("COMPILEARGS", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE",)
|
|
RETURN_NAMES = ("vae", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "HunyuanVideoWrapper"
|
|
DESCRIPTION = "Loads Hunyuan VAE model from 'ComfyUI/models/vae'"
|
|
|
|
def loadmodel(self, model_name, precision, compile_args=None):
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["HyVideoVAELoader"]()
|
|
return original_loader.loadmodel(model_name, precision, compile_args)
|
|
|
|
class DownloadAndLoadHyVideoTextEncoder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"llm_model": (["Kijai/llava-llama-3-8b-text-encoder-tokenizer","xtuner/llava-llama-3-8b-v1_1-transformers"],),
|
|
"clip_model": (["disabled","openai/clip-vit-large-patch14",],),
|
|
"precision": (["fp16", "fp32", "bf16"],
|
|
{"default": "bf16"}
|
|
),
|
|
},
|
|
"optional": {
|
|
"apply_final_norm": ("BOOLEAN", {"default": False}),
|
|
"hidden_state_skip_layer": ("INT", {"default": 2}),
|
|
"quantization": (['disabled', 'bnb_nf4', "fp8_e4m3fn"], {"default": 'disabled'}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HYVIDTEXTENCODER",)
|
|
RETURN_NAMES = ("hyvid_text_encoder", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "HunyuanVideoWrapper"
|
|
DESCRIPTION = "Loads Hunyuan text_encoder model from 'ComfyUI/models/LLM'"
|
|
|
|
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)
|
|
class WanVideoModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Use the existing get_device_list function
|
|
from . import get_device_list
|
|
devices = get_device_list()
|
|
|
|
return {
|
|
"required": {
|
|
"model": (folder_paths.get_filename_list("unet_gguf") + 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_e4m3fn_fast_no_ffn", "fp8_e4m3fn_scaled", "fp8_e5m2_scaled"],
|
|
{"default": "disabled", "tooltip": "optional quantization method"}
|
|
),
|
|
"device": (devices, {"default": devices[1] if len(devices) > 1 else devices[0], "tooltip": "Device to load the model to"}),
|
|
},
|
|
"optional": {
|
|
"attention_mode": ([
|
|
"sdpa",
|
|
"flash_attn_2",
|
|
"flash_attn_3",
|
|
"sageattn",
|
|
"sageattn_3",
|
|
"flex_attention",
|
|
"radial_sage_attention",
|
|
], {"default": "sdpa"}),
|
|
"compile_args": ("WANCOMPILEARGS", ),
|
|
"block_swap_args": ("BLOCKSWAPARGS", ),
|
|
"lora": ("WANVIDLORA", {"default": None}),
|
|
"vram_management_args": ("VRAM_MANAGEMENTARGS", {"default": None, "tooltip": "Alternative offloading method from DiffSynth-Studio, more aggressive in reducing memory use than block swapping, but can be slower"}),
|
|
"vace_model": ("VACEPATH", {"default": None, "tooltip": "VACE model to use when not using model that has it included"}),
|
|
"fantasytalking_model": ("FANTASYTALKINGMODEL", {"default": None, "tooltip": "FantasyTalking model https://github.com/Fantasy-AMAP"}),
|
|
"multitalk_model": ("MULTITALKMODEL", {"default": None, "tooltip": "Multitalk model"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVIDEOMODEL",)
|
|
RETURN_NAMES = ("model", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "WanVideoWrapper"
|
|
|
|
def loadmodel(self, model, base_precision, device, quantization,
|
|
compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None, vram_management_args=None, vace_model=None, fantasytalking_model=None, multitalk_model=None):
|
|
import logging
|
|
import comfy.model_management as mm
|
|
import torch
|
|
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] ========== CUSTOM IMPLEMENTATION ==========")
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] User selected device: {device}")
|
|
|
|
# Convert device string to torch device
|
|
selected_device = torch.device(device)
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Torch device: {selected_device}")
|
|
|
|
# Determine load_device parameter for original loader
|
|
# If user selected CPU, use "offload_device", otherwise use "main_device"
|
|
load_device = "offload_device" if device == "cpu" else "main_device"
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Mapped to load_device: {load_device}")
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["WanVideoModelLoader"]()
|
|
|
|
# Patch BOTH WanVideo modules with the selected device
|
|
import sys
|
|
import inspect
|
|
loader_module = inspect.getmodule(original_loader)
|
|
|
|
if loader_module:
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Patching WanVideo modules to use {selected_device}")
|
|
|
|
# Save original devices
|
|
original_device = getattr(loader_module, 'device', None)
|
|
original_offload = getattr(loader_module, 'offload_device', None)
|
|
|
|
# Check if there's a model offload device override (from block swap config)
|
|
model_offload_override = getattr(loader_module, '_model_offload_device_override', None)
|
|
|
|
# Patch nodes_model_loading.py module
|
|
setattr(loader_module, 'device', selected_device)
|
|
if model_offload_override:
|
|
# Use the model offload override for offload_device
|
|
setattr(loader_module, 'offload_device', model_offload_override)
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Using model offload override: {model_offload_override}")
|
|
elif device == "cpu":
|
|
setattr(loader_module, 'offload_device', selected_device)
|
|
|
|
# Patch nodes.py module as well
|
|
nodes_module_name = loader_module.__name__.replace('.nodes_model_loading', '.nodes')
|
|
if nodes_module_name in sys.modules:
|
|
nodes_module = sys.modules[nodes_module_name]
|
|
setattr(nodes_module, 'device', selected_device)
|
|
|
|
# Check for model offload override in nodes module too
|
|
nodes_model_offload_override = getattr(nodes_module, '_model_offload_device_override', None)
|
|
if nodes_model_offload_override:
|
|
setattr(nodes_module, 'offload_device', nodes_model_offload_override)
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Using model offload override for nodes.py: {nodes_model_offload_override}")
|
|
elif device == "cpu":
|
|
setattr(nodes_module, 'offload_device', selected_device)
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Both modules patched successfully")
|
|
|
|
# Call original loader with our patches in place
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Calling original loader with patched device")
|
|
result = original_loader.loadmodel(model, base_precision, load_device, quantization,
|
|
compile_args, attention_mode, block_swap_args, lora, vram_management_args, vace_model, fantasytalking_model, multitalk_model)
|
|
|
|
# Leave patches in place for subsequent operations
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] Model loaded on {selected_device}")
|
|
logging.info(f"[MultiGPU WanVideoModelLoader] ========== COMPLETE ==========")
|
|
|
|
return result
|
|
else:
|
|
logging.error(f"[MultiGPU WanVideoModelLoader] Could not patch modules, falling back")
|
|
return original_loader.loadmodel(model, base_precision, load_device, quantization,
|
|
compile_args, attention_mode, block_swap_args, lora, vram_management_args, vace_model, fantasytalking_model, multitalk_model)
|
|
|
|
|
|
class WanVideoVAELoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
from . import get_device_list
|
|
devices = get_device_list()
|
|
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("vae"),
|
|
{"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}),
|
|
"device": (devices, {"default": devices[1] if len(devices) > 1 else devices[0],
|
|
"tooltip": "Device to load the VAE to"}),
|
|
},
|
|
"optional": {
|
|
"precision": (["fp16", "fp32", "bf16"], {"default": "bf16"}),
|
|
"compile_args": ("WANCOMPILEARGS", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVAE",)
|
|
RETURN_NAMES = ("vae", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "Loads Wan VAE model with explicit device selection"
|
|
|
|
def loadmodel(self, model_name, device, precision="bf16", compile_args=None):
|
|
import logging
|
|
import torch
|
|
|
|
logging.info(f"[MultiGPU WanVideoVAELoader] User selected device: {device}")
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["WanVideoVAELoader"]()
|
|
|
|
# Patch BOTH modules with selected device
|
|
import sys
|
|
import inspect
|
|
loader_module = inspect.getmodule(original_loader)
|
|
|
|
if loader_module:
|
|
selected_device = torch.device(device)
|
|
logging.info(f"[MultiGPU WanVideoVAELoader] Patching modules to use {selected_device}")
|
|
|
|
# For VAE, we want to control where it loads initially
|
|
# Set offload_device to our selected device
|
|
setattr(loader_module, 'offload_device', selected_device)
|
|
setattr(loader_module, 'device', selected_device)
|
|
|
|
# Also patch nodes.py
|
|
nodes_module_name = loader_module.__name__.replace('.nodes_model_loading', '.nodes')
|
|
if nodes_module_name in sys.modules:
|
|
nodes_module = sys.modules[nodes_module_name]
|
|
setattr(nodes_module, 'device', selected_device)
|
|
setattr(nodes_module, 'offload_device', selected_device)
|
|
|
|
result = original_loader.loadmodel(model_name, precision, compile_args)
|
|
|
|
logging.info(f"[MultiGPU WanVideoVAELoader] VAE loaded on {selected_device}")
|
|
return result
|
|
else:
|
|
logging.error(f"[MultiGPU WanVideoVAELoader] Could not patch modules")
|
|
return original_loader.loadmodel(model_name, precision, compile_args)
|
|
|
|
|
|
class LoadWanVideoT5TextEncoder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
from . import get_device_list
|
|
devices = get_device_list()
|
|
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("text_encoders"),
|
|
{"tooltip": "These models are loaded from 'ComfyUI/models/text_encoders'"}),
|
|
"precision": (["fp32", "bf16"], {"default": "bf16"}),
|
|
"device": (devices, {"default": devices[1] if len(devices) > 1 else devices[0],
|
|
"tooltip": "Device to load the text encoder to"}),
|
|
},
|
|
"optional": {
|
|
"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/text_encoders'"
|
|
|
|
def loadmodel(self, model_name, precision, device, quantization="disabled"):
|
|
import logging
|
|
import torch
|
|
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] ========== CUSTOM IMPLEMENTATION ==========")
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] User selected device: {device}")
|
|
|
|
selected_device = torch.device(device)
|
|
load_device = "offload_device" if device == "cpu" else "main_device"
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] Mapped to load_device: {load_device}")
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["LoadWanVideoT5TextEncoder"]()
|
|
|
|
# Patch BOTH WanVideo modules
|
|
import sys
|
|
import inspect
|
|
loader_module = inspect.getmodule(original_loader)
|
|
|
|
if loader_module:
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] Patching WanVideo modules to use {selected_device}")
|
|
|
|
# Patch nodes_model_loading.py
|
|
setattr(loader_module, 'device', selected_device)
|
|
if device == "cpu":
|
|
setattr(loader_module, 'offload_device', selected_device)
|
|
|
|
# Patch nodes.py module as well
|
|
nodes_module_name = loader_module.__name__.replace('.nodes_model_loading', '.nodes')
|
|
if nodes_module_name in sys.modules:
|
|
nodes_module = sys.modules[nodes_module_name]
|
|
setattr(nodes_module, 'device', selected_device)
|
|
if device == "cpu":
|
|
setattr(nodes_module, 'offload_device', selected_device)
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] Both modules patched successfully")
|
|
|
|
result = original_loader.loadmodel(model_name, precision, load_device, quantization)
|
|
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] Text encoder loaded on {selected_device}")
|
|
logging.info(f"[MultiGPU LoadWanVideoT5TextEncoder] ========== COMPLETE ==========")
|
|
|
|
return result
|
|
else:
|
|
logging.error(f"[MultiGPU LoadWanVideoT5TextEncoder] Could not patch modules, falling back")
|
|
return original_loader.loadmodel(model_name, precision, load_device, quantization)
|
|
|
|
class WanVideoTextEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
from . import get_device_list
|
|
devices = get_device_list()
|
|
|
|
return {"required": {
|
|
"positive_prompt": ("STRING", {"default": "", "multiline": True} ),
|
|
"negative_prompt": ("STRING", {"default": "", "multiline": True} ),
|
|
"device": (devices, {"default": devices[1] if len(devices) > 1 else devices[0],
|
|
"tooltip": "Device to run the text encoding on"}),
|
|
},
|
|
"optional": {
|
|
"t5": ("WANTEXTENCODER",),
|
|
"force_offload": ("BOOLEAN", {"default": True}),
|
|
"model_to_offload": ("WANVIDEOMODEL", {"tooltip": "Model to move to offload_device before encoding"}),
|
|
"use_disk_cache": ("BOOLEAN", {"default": False, "tooltip": "Cache the text embeddings to disk for faster re-use"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVIDEOTEXTEMBEDS", )
|
|
RETURN_NAMES = ("text_embeds",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "Encodes text prompts with explicit device selection"
|
|
|
|
def process(self, positive_prompt, negative_prompt, device, t5=None, force_offload=True,
|
|
model_to_offload=None, use_disk_cache=False):
|
|
import logging
|
|
import torch
|
|
|
|
logging.info(f"[MultiGPU WanVideoTextEncode] User selected device: {device}")
|
|
|
|
# Map to original device parameter
|
|
original_device = "gpu" if device != "cpu" else "cpu"
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_encoder = NODE_CLASS_MAPPINGS["WanVideoTextEncode"]()
|
|
|
|
# Patch the modules
|
|
import sys
|
|
import inspect
|
|
encoder_module = inspect.getmodule(original_encoder)
|
|
|
|
if encoder_module:
|
|
selected_device = torch.device(device)
|
|
logging.info(f"[MultiGPU WanVideoTextEncode] Patching module to use {selected_device}")
|
|
setattr(encoder_module, 'device', selected_device)
|
|
|
|
# Also patch nodes_model_loading if needed
|
|
model_loading_name = encoder_module.__name__.replace('.nodes', '.nodes_model_loading')
|
|
if model_loading_name in sys.modules:
|
|
model_loading_module = sys.modules[model_loading_name]
|
|
setattr(model_loading_module, 'device', selected_device)
|
|
|
|
result = original_encoder.process(positive_prompt, negative_prompt, t5=t5,
|
|
force_offload=force_offload, model_to_offload=model_to_offload,
|
|
use_disk_cache=use_disk_cache, device=original_device)
|
|
|
|
logging.info(f"[MultiGPU WanVideoTextEncode] Encoding completed on {selected_device}")
|
|
return result
|
|
else:
|
|
return original_encoder.process(positive_prompt, negative_prompt, t5=t5,
|
|
force_offload=force_offload, model_to_offload=model_to_offload,
|
|
use_disk_cache=use_disk_cache, device=original_device)
|
|
|
|
class LoadWanVideoClipTextEncoder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
from . import get_device_list
|
|
devices = get_device_list()
|
|
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("clip_vision") + folder_paths.get_filename_list("text_encoders"),
|
|
{"tooltip": "These models are loaded from 'ComfyUI/models/clip_vision'"}),
|
|
"precision": (["fp16", "fp32", "bf16"], {"default": "fp16"}),
|
|
"device": (devices, {"default": devices[1] if len(devices) > 1 else devices[0],
|
|
"tooltip": "Device to load the CLIP encoder to"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP_VISION",)
|
|
RETURN_NAMES = ("clip_vision", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "Loads Wan CLIP text encoder model from 'ComfyUI/models/clip_vision'"
|
|
|
|
def loadmodel(self, model_name, precision, device):
|
|
import logging
|
|
import torch
|
|
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] ========== CUSTOM IMPLEMENTATION ==========")
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] User selected device: {device}")
|
|
|
|
selected_device = torch.device(device)
|
|
load_device = "offload_device" if device == "cpu" else "main_device"
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] Mapped to load_device: {load_device}")
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["LoadWanVideoClipTextEncoder"]()
|
|
|
|
# Patch BOTH WanVideo modules
|
|
import sys
|
|
import inspect
|
|
loader_module = inspect.getmodule(original_loader)
|
|
|
|
if loader_module:
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] Patching WanVideo modules to use {selected_device}")
|
|
|
|
# Patch nodes_model_loading.py
|
|
setattr(loader_module, 'device', selected_device)
|
|
if device == "cpu":
|
|
setattr(loader_module, 'offload_device', selected_device)
|
|
|
|
# Patch nodes.py module as well
|
|
nodes_module_name = loader_module.__name__.replace('.nodes_model_loading', '.nodes')
|
|
if nodes_module_name in sys.modules:
|
|
nodes_module = sys.modules[nodes_module_name]
|
|
setattr(nodes_module, 'device', selected_device)
|
|
if device == "cpu":
|
|
setattr(nodes_module, 'offload_device', selected_device)
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] Both modules patched successfully")
|
|
|
|
result = original_loader.loadmodel(model_name, precision, load_device)
|
|
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] CLIP encoder loaded on {selected_device}")
|
|
logging.info(f"[MultiGPU LoadWanVideoClipTextEncoder] ========== COMPLETE ==========")
|
|
|
|
return result
|
|
else:
|
|
logging.error(f"[MultiGPU LoadWanVideoClipTextEncoder] Could not patch modules, falling back")
|
|
return original_loader.loadmodel(model_name, precision, load_device)
|
|
|
|
|
|
|
|
class WanVideoModelLoader_2:
|
|
"""Second instance for multi-model workflows to maintain separate device patches"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Delegate to the primary loader
|
|
return WanVideoModelLoader.INPUT_TYPES()
|
|
|
|
RETURN_TYPES = WanVideoModelLoader.RETURN_TYPES
|
|
RETURN_NAMES = WanVideoModelLoader.RETURN_NAMES
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "Second model loader instance for workflows using multiple models on different devices"
|
|
|
|
def loadmodel(self, model, base_precision, device, quantization,
|
|
compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None,
|
|
vram_management_args=None, vace_model=None, fantasytalking_model=None, multitalk_model=None):
|
|
# Just use the first loader's implementation
|
|
loader = WanVideoModelLoader()
|
|
return loader.loadmodel(model, base_precision, device, quantization,
|
|
compile_args, attention_mode, block_swap_args, lora,
|
|
vram_management_args, vace_model, fantasytalking_model, multitalk_model)
|
|
|
|
|
|
class WanVideoSampler:
|
|
"""Wrapper that ensures correct device patching before sampling"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Get original sampler's inputs
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_types = NODE_CLASS_MAPPINGS["WanVideoSampler"].INPUT_TYPES()
|
|
return original_types
|
|
|
|
RETURN_TYPES = ("LATENT", "LATENT",)
|
|
RETURN_NAMES = ("samples", "denoised_samples",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "MultiGPU-aware sampler that ensures correct device for each model"
|
|
|
|
def process(self, model, **kwargs):
|
|
import logging
|
|
import sys
|
|
|
|
# Get the model's device and update WanVideo modules to match
|
|
model_device = model.load_device
|
|
logging.info(f"[MultiGPU WanVideoSampler] Model device: {model_device}")
|
|
|
|
# Update the device variable in WanVideo modules
|
|
for module_name in sys.modules.keys():
|
|
if 'WanVideoWrapper' in module_name and hasattr(sys.modules[module_name], 'device'):
|
|
sys.modules[module_name].device = model_device
|
|
|
|
# Call original sampler
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_sampler = NODE_CLASS_MAPPINGS["WanVideoSampler"]()
|
|
return original_sampler.process(model, **kwargs)
|
|
|
|
|
|
class WanVideoBlockSwap:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
from . import get_device_list
|
|
devices = get_device_list()
|
|
|
|
return {
|
|
"required": {
|
|
"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"}),
|
|
"swap_device": (devices, {"default": "cpu",
|
|
"tooltip": "Device to swap blocks to during sampling (default: cpu for standard behavior)"}),
|
|
"model_offload_device": (devices, {"default": "cpu",
|
|
"tooltip": "Device to offload entire model to when done (default: cpu)"}),
|
|
"offload_img_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload img_emb to swap_device"}),
|
|
"offload_txt_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload time_emb to swap_device"}),
|
|
},
|
|
"optional": {
|
|
"use_non_blocking": ("BOOLEAN", {"default": False,
|
|
"tooltip": "Use non-blocking memory transfer for offloading, reserves more RAM but is faster"}),
|
|
"vace_blocks_to_swap": ("INT", {"default": 0, "min": 0, "max": 15, "step": 1,
|
|
"tooltip": "Number of VACE blocks to swap, the VACE model has 15 blocks"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("BLOCKSWAPARGS",)
|
|
RETURN_NAMES = ("block_swap_args",)
|
|
FUNCTION = "setargs"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "Block swap settings with explicit device selection for memory management across GPUs"
|
|
|
|
def setargs(self, blocks_to_swap, swap_device, model_offload_device, offload_img_emb, offload_txt_emb,
|
|
use_non_blocking=False, vace_blocks_to_swap=0):
|
|
import logging
|
|
import torch
|
|
import comfy.model_management as mm
|
|
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] ========== CONFIGURATION ==========")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] User selected swap device: {swap_device}")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] User selected model offload device: {model_offload_device}")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Blocks to swap: {blocks_to_swap}")
|
|
|
|
# Convert device strings to torch devices
|
|
selected_swap_device = torch.device(swap_device)
|
|
selected_offload_device = torch.device(model_offload_device)
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Torch swap device: {selected_swap_device}")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Torch model offload device: {selected_offload_device}")
|
|
|
|
# Patch the offload_device in WanVideo modules to use our selected swap device
|
|
# This needs to persist through model loading
|
|
import sys
|
|
|
|
# Find the actual module paths (without the custom_nodes prefix)
|
|
for module_name in sys.modules.keys():
|
|
if 'WanVideoWrapper' in module_name and 'nodes_model_loading' in module_name:
|
|
module = sys.modules[module_name]
|
|
original_offload = getattr(module, 'offload_device', None)
|
|
# For model loading, use the model offload device
|
|
setattr(module, 'offload_device', selected_offload_device)
|
|
# Store the block swap device separately
|
|
setattr(module, '_block_swap_device_override', selected_swap_device)
|
|
setattr(module, '_model_offload_device_override', selected_offload_device)
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Patched {module_name}")
|
|
logging.info(f" - offload_device: {original_offload} -> {selected_offload_device}")
|
|
logging.info(f" - _block_swap_device_override: {selected_swap_device}")
|
|
|
|
if 'WanVideoWrapper' in module_name and module_name.endswith('.nodes'):
|
|
module = sys.modules[module_name]
|
|
original_offload = getattr(module, 'offload_device', None)
|
|
# For nodes.py, set the model offload device
|
|
setattr(module, 'offload_device', selected_offload_device)
|
|
setattr(module, '_block_swap_device_override', selected_swap_device)
|
|
setattr(module, '_model_offload_device_override', selected_offload_device)
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Patched {module_name}")
|
|
logging.info(f" - offload_device: {original_offload} -> {selected_offload_device}")
|
|
|
|
# Also store in block_swap_args so it can be used directly
|
|
block_swap_args = {
|
|
"blocks_to_swap": blocks_to_swap,
|
|
"offload_img_emb": offload_img_emb,
|
|
"offload_txt_emb": offload_txt_emb,
|
|
"use_non_blocking": use_non_blocking,
|
|
"vace_blocks_to_swap": vace_blocks_to_swap,
|
|
"swap_device": swap_device, # For block swapping
|
|
"model_offload_device": model_offload_device, # For full model offload
|
|
}
|
|
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Block swap configuration complete")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Stored swap_device in args: {swap_device}")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] Stored model_offload_device in args: {model_offload_device}")
|
|
logging.info(f"[MultiGPU WanVideoBlockSwap] ========== COMPLETE ==========")
|
|
|
|
return (block_swap_args,)
|