diff --git a/__init__.py b/__init__.py index 556630c..da7e686 100644 --- a/__init__.py +++ b/__init__.py @@ -226,11 +226,11 @@ from .wanvideo import ( WanVideoEncode, LoadWanVideoClipTextEncoder, WanVideoClipVisionEncode, - WanVideoControlnetLoaderMultiGPU, - FantasyTalkingModelLoaderMultiGPU, - Wav2VecModelLoaderMultiGPU, - WanVideoUni3C_ControlnetLoaderMultiGPU, - DownloadAndLoadWav2VecModelMultiGPU, + WanVideoControlnetLoader, + FantasyTalkingModelLoader, + Wav2VecModelLoader, + WanVideoUni3C_ControlnetLoader, + DownloadAndLoadWav2VecModel, ) from .wrappers import ( @@ -382,11 +382,11 @@ wanvideo_nodes = { "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) diff --git a/wanvideo.py b/wanvideo.py index 69f8214..dde7f1d 100644 --- a/wanvideo.py +++ b/wanvideo.py @@ -714,116 +714,6 @@ class WanVideoClipVisionEncode: 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 @@ class WanVideoControlnetLoaderMultiGPU: 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 @@ class FantasyTalkingModelLoaderMultiGPU: 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 @@ class Wav2VecModelLoaderMultiGPU: 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()