diff --git a/__init__.py b/__init__.py index 7f7c9df..3974ba1 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,3 @@ -# __init__.py import copy import torch import sys @@ -255,27 +254,22 @@ class HunyuanVideoEmbeddingsAdapter: CATEGORY = "multigpu" def adapt_embeddings(self, hyvid_embeds): - # Create main conditioning tensor cond = hyvid_embeds["prompt_embeds"] - # Create pooled dict with all our extra information pooled_dict = { "pooled_output": hyvid_embeds["prompt_embeds_2"], "cross_attn": hyvid_embeds["prompt_embeds"], "attention_mask": hyvid_embeds["attention_mask"], } - # Add CLIP's attention mask if present if hyvid_embeds["attention_mask_2"] is not None: pooled_dict["attention_mask_controlnet"] = hyvid_embeds["attention_mask_2"] - # Add guidance if present - typically these and negative_xxxx are empty for HunyuanVideo if hyvid_embeds["cfg"] is not None: pooled_dict["guidance"] = float(hyvid_embeds["cfg"]) pooled_dict["start_percent"] = float(hyvid_embeds["start_percent"]) if hyvid_embeds["start_percent"] is not None else 0.0 pooled_dict["end_percent"] = float(hyvid_embeds["end_percent"]) if hyvid_embeds["end_percent"] is not None else 1.0 - # Finally create the conditioning list in the exact format that encode_from_tokens_scheduled returns return ([[cond, pooled_dict]],) def override_class(cls): @@ -365,690 +359,78 @@ def override_class_with_distorch(cls): return NodeOverrideDisTorch - -NODE_CLASS_MAPPINGS = { - "DeviceSelectorMultiGPU": DeviceSelectorMultiGPU, - "HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter, -} - def check_module_exists(module_path): full_path = os.path.join(folder_paths.get_folder_paths("custom_nodes")[0], module_path) logging.info(f"MultiGPU: Checking for module at {full_path}") - if not os.path.exists(full_path): logging.info(f"MultiGPU: Module {module_path} not found - skipping") return False - logging.info(f"MultiGPU: Found {module_path}, creating compatible MultiGPU nodes") return True -def register_module(target_nodes): - from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS +from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS +from .nodes import ( + UnetLoaderGGUF, UnetLoaderGGUFAdvanced, + CLIPLoaderGGUF, DualCLIPLoaderGGUF, TripleCLIPLoaderGGUF, + LTXVLoader, + Florence2ModelLoader, DownloadAndLoadFlorence2Model, + CheckpointLoaderNF4, + LoadFluxControlNet, + MMAudioModelLoader, MMAudioFeatureUtilsLoader, MMAudioSampler, + PulidModelLoader, PulidInsightFaceLoader, PulidEvaClipLoader, + HyVideoModelLoader, HyVideoModelLoaderDiffSynth, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder +) - for node in target_nodes: - NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) - return +NODE_CLASS_MAPPINGS = { + "DeviceSelectorMultiGPU": DeviceSelectorMultiGPU, + "HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter +} +NODE_CLASS_MAPPINGS["UNETLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["UNETLoader"]) +NODE_CLASS_MAPPINGS["VAELoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["VAELoader"]) +NODE_CLASS_MAPPINGS["CLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"]) +NODE_CLASS_MAPPINGS["DualCLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["DualCLIPLoader"]) +NODE_CLASS_MAPPINGS["TripleCLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["TripleCLIPLoader"]) +NODE_CLASS_MAPPINGS["CheckpointLoaderSimpleMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]) +NODE_CLASS_MAPPINGS["ControlNetLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["ControlNetLoader"]) -def register_UnetLoaderGGUFMultiGPU(): - global 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) - - # Create both MultiGPU versions of the base class - UnetLoaderGGUFMultiGPU = override_class(UnetLoaderGGUF) - NODE_CLASS_MAPPINGS["UnetLoaderGGUFMultiGPU"] = UnetLoaderGGUFMultiGPU - logging.info(f"MultiGPU: Registered UnetLoaderGGUFMultiGPU") - - UnetLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(UnetLoaderGGUF) - NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = UnetLoaderGGUFDisTorchMultiGPU - logging.info(f"MultiGPU: Registered UnetLoaderGGUFDisTorchMultiGPU") - - 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)" - - # Create both MultiGPU versions of the advanced class - UnetLoaderGGUFAdvancedMultiGPU = override_class(UnetLoaderGGUFAdvanced) - NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedMultiGPU"] = UnetLoaderGGUFAdvancedMultiGPU - logging.info(f"MultiGPU: Registered UnetLoaderGGUFAdvancedMultiGPU") - - UnetLoaderGGUFAdvancedDisTorchMultiGPU = override_class_with_distorch(UnetLoaderGGUFAdvanced) - NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedDisTorchMultiGPU"] = UnetLoaderGGUFAdvancedDisTorchMultiGPU - logging.info(f"MultiGPU: Registered UnetLoaderGGUFAdvancedDisTorchMultiGPU") - -def register_CLIPLoaderGGUFMultiGPU(): - global NODE_CLASS_MAPPINGS - - 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"],), - } - } - - 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) - - # Create the MultiGPU version of the base class - CLIPLoaderGGUFMultiGPU = override_class(CLIPLoaderGGUF) - NODE_CLASS_MAPPINGS["CLIPLoaderGGUFMultiGPU"] = CLIPLoaderGGUFMultiGPU - logging.info(f"MultiGPU: Registered CLIPLoaderGGUFMultiGPU") - - CLIPLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(CLIPLoaderGGUF) - NODE_CLASS_MAPPINGS["CLIPLoaderGGUFDisTorchMultiGPU"] = CLIPLoaderGGUFDisTorchMultiGPU - logging.info(f"MultiGPU: Registered CLIPLoaderGGUFDisTorchMultiGPU") - - # Now create the advanced version that inherits from the MultiGPU base class - - 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 - # Create the MultiGPU version of the advanced class - DualCLIPLoaderGGUFMultiGPU = override_class(DualCLIPLoaderGGUF) - NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFMultiGPU"] = DualCLIPLoaderGGUFMultiGPU - logging.info(f"MultiGPU: Registered DualCLIPLoaderGGUFMultiGPU") - - DualCLIPLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(DualCLIPLoaderGGUF) - NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFDisTorchMultiGPU"] = DualCLIPLoaderGGUFDisTorchMultiGPU - logging.info(f"MultiGPU: Registered DualCLIPLoaderGGUFDisTorchMultiGPU") - - 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) - # Create the MultiGPU version of the advanced class - TripleCLIPLoaderGGUFMultiGPU = override_class(TripleCLIPLoaderGGUF) - NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFMultiGPU"] = TripleCLIPLoaderGGUFMultiGPU - logging.info(f"MultiGPU: Registered TripleCLIPLoaderGGUFMultiGPU") - - TripleCLIPLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(TripleCLIPLoaderGGUF) - NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFDisTorchMultiGPU"] = TripleCLIPLoaderGGUFDisTorchMultiGPU - logging.info(f"MultiGPU: Registered TripleCLIPLoaderGGUFDisTorchMultiGPU") - -def register_LTXVLoaderMultiGPU(): - global NODE_CLASS_MAPPINGS - - 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) +if check_module_exists("ComfyUI-LTXVideo") or check_module_exists("comfyui-ltxvideo"): NODE_CLASS_MAPPINGS["LTXVLoaderMultiGPU"] = override_class(LTXVLoader) - logging.info(f"MultiGPU: Registered LTXVLoaderMultiGPU") - -def register_Florence2ModelLoaderMultiGPU(): - global NODE_CLASS_MAPPINGS - - 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) - NODE_CLASS_MAPPINGS["Florence2ModelLoaderMultiGPU"] = override_class(Florence2ModelLoader) - logging.info(f"MultiGPU: Registered Florence2ModelLoaderMultiGPU") - -def register_DownloadAndLoadFlorence2ModelMultiGPU(): - global NODE_CLASS_MAPPINGS - - 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) - NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2ModelMultiGPU"] = override_class(DownloadAndLoadFlorence2Model) - logging.info(f"MultiGPU: Registered DownloadAndLoadFlorence2ModelMultiGPU") - -def register_CheckpointLoaderNF4(): - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["CheckpointLoaderNF4MultiGPU"] = override_class(CheckpointLoaderNF4) - logging.info(f"MultiGPU: Registered CheckpointLoaderNF4MultiGPU") - -def register_LoadFluxControlNetMultiGPU(): - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["LoadFluxControlNetMultiGPU"] = override_class(LoadFluxControlNet) - logging.info(f"MultiGPU: Registered LoadFluxControlNetMultiGPU") - -def register_MMAudioModelLoaderMultiGPU(): - - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["MMAudioModelLoaderMultiGPU"] = override_class(MMAudioModelLoader) - logging.info(f"MultiGPU: Registered MMAudioModelLoaderMultiGPU") - -def register_MMAudioFeatureUtilsLoaderMultiGPU(): - - global NODE_CLASS_MAPPINGS - - - 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) - - NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoaderMultiGPU"] = override_class(MMAudioFeatureUtilsLoader) - logging.info(f"MultiGPU: Registered MMAudioFeatureUtilsLoaderMultiGPU") - -def register_MMAudioSamplerMultiGPU(): - - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["MMAudioSamplerMultiGPU"] = override_class(MMAudioSampler) - logging.info(f"MultiGPU: Registered MMAudioSamplerMultiGPU") - -def register_PulidModelLoader(): - - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["PulidModelLoaderMultiGPU"] = override_class(PulidModelLoader) - logging.info(f"MultiGPU: Registered PulidModelLoaderMultiGPU") - -def register_PulidInsightFaceLoader(): - - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["PulidInsightFaceLoaderMultiGPU"] = override_class(PulidInsightFaceLoader) - logging.info(f"MultiGPU: Registered PulidInsightFaceLoaderMultiGPU") - -def register_PulidEvaClipLoader(): - - global NODE_CLASS_MAPPINGS - - 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() - - NODE_CLASS_MAPPINGS["PulidEvaClipLoaderMultiGPU"] = override_class(PulidEvaClipLoader) - logging.info(f"MultiGPU: Registered PulidEvaClipLoaderMultiGPU") - -def register_HyVideoModelLoader(): - global NODE_CLASS_MAPPINGS - - # Keep original MultiGPU wrapper unchanged - 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) - - # Add new DiffSynth-style node - class HyVideoModelLoaderDiffSynth: - @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"}), - }, - "optional": { - "attention_mode": ([ - "sdpa", - "flash_attn_varlen", - "sageattn_varlen", - "comfy", - ], {"default": "flash_attn"}), - "compile_args": ("COMPILEARGS", ), - "block_swap_args": ("BLOCKSWAPARGS", ), - "lora": ("HYVIDLORA", {"default": None}), - } - } - - RETURN_TYPES = ("HYVIDEOMODEL",) - RETURN_NAMES = ("model", ) - FUNCTION = "loadmodel" - CATEGORY = "HunyuanVideoWrapper" - - def loadmodel(self, model, base_precision, quantization, compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None): - from nodes import NODE_CLASS_MAPPINGS - original_loader = NODE_CLASS_MAPPINGS["HyVideoModelLoader"]() - # Use DiffSynth's auto offloading approach - return original_loader.loadmodel(model, base_precision, "main_device", quantization, - compile_args, attention_mode, block_swap_args, lora, - auto_cpu_offload=True) - - # Register both with MultiGPU wrapper - NODE_CLASS_MAPPINGS["HyVideoModelLoaderMultiGPU"] = override_class(HyVideoModelLoader) - NODE_CLASS_MAPPINGS["HyVideoModelLoaderDiffSynthMultiGPU"] = override_class_with_offload(HyVideoModelLoaderDiffSynth) - - logging.info(f"MultiGPU: Registered HyVideoModelLoader nodes") - -def register_HyVideoVAELoader(): - - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["HyVideoVAELoaderMultiGPU"] = override_class(HyVideoVAELoader) - logging.info(f"MultiGPU: Registered HyVideoVAELoaderMultiGPU") - -def register_DownloadAndLoadHyVideoTextEncoder(): - - global NODE_CLASS_MAPPINGS - - 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) - - NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoderMultiGPU"] = override_class(DownloadAndLoadHyVideoTextEncoder) - logging.info(f"MultiGPU: Registered DownloadAndLoadHyVideoTextEncoderMultiGPU") - -# Register desired nodes -register_module(["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"]) -if check_module_exists("ComfyUI-LTXVideo") or check_module_exists("comfyui-ltxvideo"): - register_LTXVLoaderMultiGPU() if check_module_exists("ComfyUI-Florence2") or check_module_exists("comfyui-florence2"): - register_Florence2ModelLoaderMultiGPU() - register_DownloadAndLoadFlorence2ModelMultiGPU() -if check_module_exists("ComfyUI_bitsandbytes_NF4") or check_module_exists("comfyui_bitsandbytes_nf4"): - register_CheckpointLoaderNF4() -if check_module_exists("x-flux-comfyui") or check_module_exists("x-flux-comfyui"): - register_LoadFluxControlNetMultiGPU() -if check_module_exists("ComfyUI-MMAudio") or check_module_exists("comfyui-mmaudio"): - register_MMAudioModelLoaderMultiGPU() - register_MMAudioFeatureUtilsLoaderMultiGPU() - register_MMAudioSamplerMultiGPU() -if check_module_exists("ComfyUI-GGUF") or check_module_exists("comfyui-gguf"): - register_UnetLoaderGGUFMultiGPU() - register_CLIPLoaderGGUFMultiGPU() -if check_module_exists("PuLID_ComfyUI") or check_module_exists("pulid_comfyui"): - register_PulidModelLoader() - register_PulidInsightFaceLoader() - register_PulidEvaClipLoader() -if check_module_exists("ComfyUI-HunyuanVideoWrapper") or check_module_exists("comfyui-hunyuanvideowrapper"): - register_HyVideoModelLoader() - register_HyVideoVAELoader() - register_DownloadAndLoadHyVideoTextEncoder() + NODE_CLASS_MAPPINGS["Florence2ModelLoaderMultiGPU"] = override_class(Florence2ModelLoader) + NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2ModelMultiGPU"] = override_class(DownloadAndLoadFlorence2Model) +if check_module_exists("ComfyUI_bitsandbytes_NF4") or check_module_exists("comfyui_bitsandbytes_nf4"): + NODE_CLASS_MAPPINGS["CheckpointLoaderNF4MultiGPU"] = override_class(CheckpointLoaderNF4) + +if check_module_exists("x-flux-comfyui") or check_module_exists("x-flux-comfyui"): + NODE_CLASS_MAPPINGS["LoadFluxControlNetMultiGPU"] = override_class(LoadFluxControlNet) + +if check_module_exists("ComfyUI-MMAudio") or check_module_exists("comfyui-mmaudio"): + NODE_CLASS_MAPPINGS["MMAudioModelLoaderMultiGPU"] = override_class(MMAudioModelLoader) + NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoaderMultiGPU"] = override_class(MMAudioFeatureUtilsLoader) + NODE_CLASS_MAPPINGS["MMAudioSamplerMultiGPU"] = override_class(MMAudioSampler) + +if check_module_exists("ComfyUI-GGUF") or check_module_exists("comfyui-gguf"): + NODE_CLASS_MAPPINGS["UnetLoaderGGUFMultiGPU"] = override_class(UnetLoaderGGUF) + NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(UnetLoaderGGUF) + NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedMultiGPU"] = override_class(UnetLoaderGGUFAdvanced) + NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedDisTorchMultiGPU"] = override_class_with_distorch(UnetLoaderGGUFAdvanced) + NODE_CLASS_MAPPINGS["CLIPLoaderGGUFMultiGPU"] = override_class(CLIPLoaderGGUF) + NODE_CLASS_MAPPINGS["CLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(CLIPLoaderGGUF) + NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFMultiGPU"] = override_class(DualCLIPLoaderGGUF) + NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(DualCLIPLoaderGGUF) + NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFMultiGPU"] = override_class(TripleCLIPLoaderGGUF) + NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(TripleCLIPLoaderGGUF) + +if check_module_exists("PuLID_ComfyUI") or check_module_exists("pulid_comfyui"): + NODE_CLASS_MAPPINGS["PulidModelLoaderMultiGPU"] = override_class(PulidModelLoader) + NODE_CLASS_MAPPINGS["PulidInsightFaceLoaderMultiGPU"] = override_class(PulidInsightFaceLoader) + NODE_CLASS_MAPPINGS["PulidEvaClipLoaderMultiGPU"] = override_class(PulidEvaClipLoader) + +if check_module_exists("ComfyUI-HunyuanVideoWrapper") or check_module_exists("comfyui-hunyuanvideowrapper"): + NODE_CLASS_MAPPINGS["HyVideoModelLoaderMultiGPU"] = override_class(HyVideoModelLoader) + NODE_CLASS_MAPPINGS["HyVideoVAELoaderMultiGPU"] = override_class(HyVideoVAELoader) + NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoderMultiGPU"] = override_class(DownloadAndLoadHyVideoTextEncoder) logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..a824d14 --- /dev/null +++ b/nodes.py @@ -0,0 +1,496 @@ +import os +import torch +import folder_paths +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"],), + } + } + + 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 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) + +# Add new DiffSynth-style node +class HyVideoModelLoaderDiffSynth: + @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"}), + }, + "optional": { + "attention_mode": ([ + "sdpa", + "flash_attn_varlen", + "sageattn_varlen", + "comfy", + ], {"default": "flash_attn"}), + "compile_args": ("COMPILEARGS", ), + "block_swap_args": ("BLOCKSWAPARGS", ), + "lora": ("HYVIDLORA", {"default": None}), + } + } + + RETURN_TYPES = ("HYVIDEOMODEL",) + RETURN_NAMES = ("model", ) + FUNCTION = "loadmodel" + CATEGORY = "HunyuanVideoWrapper" + + def loadmodel(self, model, base_precision, quantization, compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None): + from nodes import NODE_CLASS_MAPPINGS + original_loader = NODE_CLASS_MAPPINGS["HyVideoModelLoader"]() + # Use DiffSynth's auto offloading approach + return original_loader.loadmodel(model, base_precision, "main_device", quantization, + compile_args, attention_mode, block_swap_args, lora, + auto_cpu_offload=True) + +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) \ No newline at end of file