From caee8716b67a4615023e462740aa63362dc54893 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 30 Dec 2024 07:54:57 -0600 Subject: [PATCH 1/5] Refactor - intermediate step of Implementing MultiGPU node registration and class definition retrieval for custom nodes --- __init__.py | 233 ++++++++++++++++++++++++++++++---------------------- 1 file changed, 136 insertions(+), 97 deletions(-) diff --git a/__init__.py b/__init__.py index 072f070..583c4a5 100644 --- a/__init__.py +++ b/__init__.py @@ -1,29 +1,103 @@ -import time +import os +import ast import copy import torch -import comfy.model_management -import os -import importlib.util import logging +from typing import Dict, Type -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logging.info("MultiGPU: Initialization started") -current_device = comfy.model_management.get_torch_device() -logging.info(f"MultiGPU: Initial device {current_device}") +def find_node_definition(module_path: str, node_name: str) -> Dict: + """ + Finds a specific node class definition by searching Python files in the given path. + """ + search_dir = os.path.join("custom_nodes", module_path) + if not os.path.exists(search_dir): + logging.info(f"MultiGPU: No custom_nodes directory {module_path}, skipping") + return None + + py_files = [] + for root, _, files in os.walk(search_dir): + for file in files: + if file.endswith('.py'): + py_files.append(os.path.join(root, file)) + + if not py_files: + return None + + try: + for file_path in py_files: + with open(file_path, 'r', encoding='utf-8') as f: + tree = ast.parse(f.read()) -def get_torch_device_patched(): - if ( - not torch.cuda.is_available() - or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU - or "cpu" in str(current_device).lower() - ): - return torch.device("cpu") - return torch.device(current_device) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == node_name: + class_info = { + 'input_types': None, + 'return_types': None, + 'function': None, + 'category': None + } -comfy.model_management.get_torch_device = get_torch_device_patched + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == 'INPUT_TYPES': + if any(d.id == 'classmethod' for d in item.decorator_list if isinstance(d, ast.Name)): + for stmt in item.body: + if isinstance(stmt, ast.Return): + try: + class_info['input_types'] = ast.literal_eval(stmt.value) + except: + pass + + elif isinstance(item, ast.Assign): + for target in item.targets: + if isinstance(target, ast.Name): + try: + if target.id == 'RETURN_TYPES': + class_info['return_types'] = ast.literal_eval(item.value) + elif target.id == 'FUNCTION': + class_info['function'] = ast.literal_eval(item.value) + elif target.id == 'CATEGORY': + class_info['category'] = ast.literal_eval(item.value) + except: + pass + + return class_info + + except Exception as e: + logging.error(f"MultiGPU: Error scanning for {node_name}: {str(e)}") + + return None + +def create_multigpu_node(node_name: str, class_info: Dict) -> Type: + """Creates a MultiGPU version of the node""" + class MultiGPUNode: + @classmethod + def INPUT_TYPES(cls): + inputs = copy.deepcopy(class_info['input_types']) + if inputs is None: + inputs = {"required": {}, "optional": {}} + elif "required" not in inputs: + inputs["required"] = {} + + devices = ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())] + inputs["required"]["device"] = (devices,) + return inputs + + RETURN_TYPES = class_info['return_types'] if class_info['return_types'] is not None else tuple() + FUNCTION = "override" + CATEGORY = "multigpu" + + def override(self, *args, device="cpu", **kwargs): + global current_device + current_device = device + return args if isinstance(args, tuple) else (args,) + + MultiGPUNode.__name__ = f"{node_name}MultiGPU" + return MultiGPUNode def override_class(cls): + """Creates a MultiGPU version of a node class that preserves original functionality.""" class NodeOverride(cls): @classmethod def INPUT_TYPES(s): @@ -43,96 +117,61 @@ def override_class(cls): return NodeOverride -def register_module(module_path, target_nodes): +def register_core_nodes(target_nodes: list): + """Register MultiGPU versions of core ComfyUI nodes""" try: - # For core nodes, skip module loading and just register from global mappings - if not module_path: - logging.info("MultiGPU: Starting core node registration") - from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS - for node in target_nodes: - if node in GLOBAL_NODE_CLASS_MAPPINGS: - NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) - logging.info(f"MultiGPU: Registered core node {node}") - else: - logging.info(f"MultiGPU: Core node {node} not found - this shouldn't happen!") - return - - # For custom nodes, try to load module first - full_path = os.path.join("custom_nodes", module_path, "__init__.py") - 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 - - logging.info(f"MultiGPU: Found {module_path}, attempting to load") - spec = importlib.util.spec_from_file_location(module_path, full_path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - logging.info(f"MultiGPU: Executed {module_path} initialization") - from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS - logging.info(f"MultiGPU: Looking for {module_path} nodes in global mappings") + logging.info("MultiGPU: Processing core nodes") for node in target_nodes: if node in GLOBAL_NODE_CLASS_MAPPINGS: NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) - logging.info(f"MultiGPU: Successfully wrapped {node}") + logging.info(f"MultiGPU: Registered {node}MultiGPU") else: - logging.info(f"MultiGPU: Node {node} from {module_path} not found in global mappings") + logging.info(f"MultiGPU: Core node {node} not found") + except Exception as e: + logging.error(f"MultiGPU: Error processing core nodes: {str(e)}") + +def register_module(module_path: str, target_nodes: list): + """Register MultiGPU versions of custom nodes""" + try: + # Handle custom nodes + logging.info(f"MultiGPU: Processing module {module_path}") + search_dir = os.path.join("custom_nodes", module_path) + if not os.path.exists(search_dir): + logging.info(f"MultiGPU: Module directory {module_path} not found, skipping") + return + + # List all Python files in the module once + py_files = [] + for root, _, files in os.walk(search_dir): + for file in files: + if file.endswith('.py'): + py_files.append(os.path.basename(file)) + + if py_files: + logging.info(f"MultiGPU: Searching in {module_path}: {', '.join(sorted(py_files))}") + + for node in target_nodes: + class_info = find_node_definition(module_path, node) + if class_info: + NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = create_multigpu_node(node, class_info) + logging.info(f"MultiGPU: Registered {node}MultiGPU") except Exception as e: - logging.info(f"MultiGPU: Error processing {module_path}: {str(e)}") + logging.error(f"MultiGPU: Error in {module_path}: {str(e)}") +# Initialize NODE_CLASS_MAPPINGS = {} +current_device = None -# Let's test just one new module at a time, starting with GGUF -logging.info("MultiGPU: Starting Core ComfyUI registration") -register_module("", [ - "UNETLoader", - "VAELoader", - "CLIPLoader", - "DualCLIPLoader", - "TripleCLIPLoader", - "CheckpointLoaderSimple", - "ControlNetLoader" -]) +# Register all modules +register_module("", ["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader","TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"]) +register_module("ComfyUI-GGUF", ["UnetLoaderGGUF", "UnetLoaderGGUFAdvanced", "CLIPLoaderGGUF","DualCLIPLoaderGGUF", "TripleCLIPLoaderGGUF"]) +register_module("x-flux-comfyui", ["LoadFluxControlNet"]) +register_module("ComfyUI-Florence2", ["Florence2ModelLoader", "DownloadAndLoadFlorence2Model"]) +register_module("ComfyUI-LTXVideo", ["LTXVLoader"]) +register_module("ComfyUI-MMAudio", [ "MMAudioFeatureUtilsLoader", "MMAudioModelLoader", "MMAudioSampler"]) +register_module("ComfyUI_bitsandbytes_NF4", ["CheckpointLoaderNF4"]) -logging.info("MultiGPU: Starting GGUF registration") -register_module("ComfyUI-GGUF", [ - "UnetLoaderGGUF", - "UnetLoaderGGUFAdvanced", - "CLIPLoaderGGUF", - "DualCLIPLoaderGGUF", - "TripleCLIPLoaderGGUF" -]) -logging.info("MultiGPU: Starting X-Flux ControlNet registration") -register_module("x-flux-comfyui", [ - "LoadFluxControlNet" -]) - -logging.info("MultiGPU: Starting Florence2 registration") -register_module("ComfyUI-Florence2", [ - "Florence2ModelLoader", - "DownloadAndLoadFlorence2Model" -]) - -logging.info("MultiGPU: Starting LTXVideo registration") -register_module("ComfyUI-LTXVideo", [ - "LTXVLoader" -]) - -logging.info("MultiGPU: Starting MMAudio registration") -register_module("ComfyUI-MMAudio", [ - "MMAudioFeatureUtilsLoader", - "MMAudioModelLoader", - "MMAudioSampler" -]) - -logging.info("MultiGPU: Starting NF4 registration") -register_module("ComfyUI_bnb_nf4_fp4_Loaders", [ - "CheckpointLoaderNF4", - "UNETLoaderNF4" -]) - - -logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") \ No newline at end of file +logging.info("MultiGPU: Registration complete") +logging.info(f"MultiGPU: Registered nodes: {', '.join(sorted(NODE_CLASS_MAPPINGS.keys()))}") \ No newline at end of file From 14758b1985b134be39337e9ca6e097d884724317 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 30 Dec 2024 10:12:28 -0600 Subject: [PATCH 2/5] changed methodology (again) to run their code and then scan their LOCAL dict, not the global dict. This will work and be robust I strongly believe Still too much new code, but we'll slim it down later --- __init__.py | 259 ++++++++++++++++++++++++---------------------------- 1 file changed, 121 insertions(+), 138 deletions(-) diff --git a/__init__.py b/__init__.py index 583c4a5..6ab8f88 100644 --- a/__init__.py +++ b/__init__.py @@ -1,107 +1,42 @@ -import os -import ast +import time import copy import torch +import comfy.model_management +import os +import importlib.util import logging -from typing import Dict, Type +############################################################################## +# INITIAL SETUP +############################################################################## +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logging.info("MultiGPU: Initialization started") -def find_node_definition(module_path: str, node_name: str) -> Dict: - """ - Finds a specific node class definition by searching Python files in the given path. - """ - search_dir = os.path.join("custom_nodes", module_path) - if not os.path.exists(search_dir): - logging.info(f"MultiGPU: No custom_nodes directory {module_path}, skipping") - return None - - py_files = [] - for root, _, files in os.walk(search_dir): - for file in files: - if file.endswith('.py'): - py_files.append(os.path.join(root, file)) - - if not py_files: - return None - - try: - for file_path in py_files: - with open(file_path, 'r', encoding='utf-8') as f: - tree = ast.parse(f.read()) +current_device = comfy.model_management.get_torch_device() +logging.info(f"MultiGPU: Initial device {current_device}") - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and node.name == node_name: - class_info = { - 'input_types': None, - 'return_types': None, - 'function': None, - 'category': None - } +def get_torch_device_patched(): + if ( + not torch.cuda.is_available() + or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU + or "cpu" in str(current_device).lower() + ): + return torch.device("cpu") + return torch.device(current_device) - for item in node.body: - if isinstance(item, ast.FunctionDef) and item.name == 'INPUT_TYPES': - if any(d.id == 'classmethod' for d in item.decorator_list if isinstance(d, ast.Name)): - for stmt in item.body: - if isinstance(stmt, ast.Return): - try: - class_info['input_types'] = ast.literal_eval(stmt.value) - except: - pass - - elif isinstance(item, ast.Assign): - for target in item.targets: - if isinstance(target, ast.Name): - try: - if target.id == 'RETURN_TYPES': - class_info['return_types'] = ast.literal_eval(item.value) - elif target.id == 'FUNCTION': - class_info['function'] = ast.literal_eval(item.value) - elif target.id == 'CATEGORY': - class_info['category'] = ast.literal_eval(item.value) - except: - pass - - return class_info - - except Exception as e: - logging.error(f"MultiGPU: Error scanning for {node_name}: {str(e)}") - - return None - -def create_multigpu_node(node_name: str, class_info: Dict) -> Type: - """Creates a MultiGPU version of the node""" - class MultiGPUNode: - @classmethod - def INPUT_TYPES(cls): - inputs = copy.deepcopy(class_info['input_types']) - if inputs is None: - inputs = {"required": {}, "optional": {}} - elif "required" not in inputs: - inputs["required"] = {} - - devices = ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())] - inputs["required"]["device"] = (devices,) - return inputs - - RETURN_TYPES = class_info['return_types'] if class_info['return_types'] is not None else tuple() - FUNCTION = "override" - CATEGORY = "multigpu" - - def override(self, *args, device="cpu", **kwargs): - global current_device - current_device = device - return args if isinstance(args, tuple) else (args,) - - MultiGPUNode.__name__ = f"{node_name}MultiGPU" - return MultiGPUNode +comfy.model_management.get_torch_device = get_torch_device_patched +############################################################################## +# OVERRIDE CLASS +############################################################################## def override_class(cls): - """Creates a MultiGPU version of a node class that preserves original functionality.""" class NodeOverride(cls): @classmethod def INPUT_TYPES(s): inputs = copy.deepcopy(cls.INPUT_TYPES()) + # In case some node forgot "required" + if "required" not in inputs: + inputs["required"] = {} devices = ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())] inputs["required"]["device"] = (devices,) return inputs @@ -109,69 +44,117 @@ def override_class(cls): CATEGORY = "multigpu" FUNCTION = "override" - def override(self, *args, device, **kwargs): + def override(self, *args, device="cpu", **kwargs): global current_device current_device = device fn = getattr(super(), cls.FUNCTION) return fn(*args, **kwargs) - return NodeOverride -def register_core_nodes(target_nodes: list): - """Register MultiGPU versions of core ComfyUI nodes""" +############################################################################## +# OUR LOCAL MAPPING OF MULTIGPU OVERRIDE NODES +# (No strict reason it has to be before the function defs, but it's typical.) +############################################################################## +NODE_CLASS_MAPPINGS = {} + +############################################################################## +# PART 1: CORE NODES +############################################################################## +def register_core_nodes(core_node_names): + """ + Uses ComfyUI's GLOBAL_NODE_CLASS_MAPPINGS to wrap core nodes. + """ + logging.info("MultiGPU: Starting core node registration") try: from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS - logging.info("MultiGPU: Processing core nodes") - for node in target_nodes: - if node in GLOBAL_NODE_CLASS_MAPPINGS: - NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) - logging.info(f"MultiGPU: Registered {node}MultiGPU") - else: - logging.info(f"MultiGPU: Core node {node} not found") - except Exception as e: - logging.error(f"MultiGPU: Error processing core nodes: {str(e)}") + except ImportError as e: + logging.error(f"MultiGPU: Could not import ComfyUI global node mappings: {e}") + return + + for node in core_node_names: + if node in GLOBAL_NODE_CLASS_MAPPINGS: + NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) + logging.info(f"MultiGPU: Registered core node {node}") + else: + logging.info(f"MultiGPU: Core node '{node}' not found in ComfyUI global mappings") + +############################################################################## +# PART 2: CUSTOM NODES (LOCAL DICTIONARY APPROACH, with extra debug logs) +############################################################################## +def register_module(module_path: str, target_nodes: list, local_map_name="NODE_CLASS_MAPPINGS"): + """ + 1) Load custom_nodes//__init__.py via importlib + 2) Grab the dictionary named `local_map_name` (default: 'NODE_CLASS_MAPPINGS') from that module + 3) For each node in target_nodes, if found in that local dictionary, wrap it + 4) No fallback to ComfyUI's global mappings + """ + base_dir = os.path.join("custom_nodes", module_path) + init_file = os.path.join(base_dir, "__init__.py") + + logging.info(f"MultiGPU: Checking custom node module at: {init_file}") + if not os.path.exists(init_file): + logging.info(f"MultiGPU: Module {module_path} not found or missing __init__.py, skipping.") + return -def register_module(module_path: str, target_nodes: list): - """Register MultiGPU versions of custom nodes""" try: - # Handle custom nodes - logging.info(f"MultiGPU: Processing module {module_path}") - search_dir = os.path.join("custom_nodes", module_path) - if not os.path.exists(search_dir): - logging.info(f"MultiGPU: Module directory {module_path} not found, skipping") - return - - # List all Python files in the module once - py_files = [] - for root, _, files in os.walk(search_dir): - for file in files: - if file.endswith('.py'): - py_files.append(os.path.basename(file)) - - if py_files: - logging.info(f"MultiGPU: Searching in {module_path}: {', '.join(sorted(py_files))}") - - for node in target_nodes: - class_info = find_node_definition(module_path, node) - if class_info: - NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = create_multigpu_node(node, class_info) - logging.info(f"MultiGPU: Registered {node}MultiGPU") - + logging.info(f"MultiGPU: Found {module_path}, loading local dictionary from {init_file}") + spec = importlib.util.spec_from_file_location(module_path, init_file) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + logging.info(f"MultiGPU: Executed {module_path} initialization") except Exception as e: - logging.error(f"MultiGPU: Error in {module_path}: {str(e)}") + logging.error(f"MultiGPU: Error loading {module_path}: {e}") + return -# Initialize -NODE_CLASS_MAPPINGS = {} -current_device = None + # Grab that module's local dictionary (e.g. NODE_CLASS_MAPPINGS) + local_map = getattr(module, local_map_name, None) + if not local_map: + logging.info(f"MultiGPU: {module_path} has no '{local_map_name}' dictionary, skipping override.") + return + + # DEBUG: Show everything this node map provides + all_defined_nodes = list(local_map.keys()) + logging.info(f"MultiGPU: {module_path} local dict keys: {all_defined_nodes}") + + # Wrap each node we want + for node in target_nodes: + if node in local_map: + mgpu_class = override_class(local_map[node]) + NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = mgpu_class + logging.info(f"MultiGPU: Successfully wrapped {node} from {module_path}") + else: + logging.info(f"MultiGPU: Node '{node}' not found in {module_path}'s local dictionary") + + +############################################################################## +# EXAMPLE USAGE +############################################################################## + +# 1) CORE NODES +register_core_nodes([ + "UNETLoader", + "VAELoader", + "CLIPLoader", + "DualCLIPLoader", + "TripleCLIPLoader", + "CheckpointLoaderSimple", + "ControlNetLoader", +]) + +# 2) CUSTOM NODES +register_module("ComfyUI-GGUF", [ + "UnetLoaderGGUF", + "UnetLoaderGGUFAdvanced", + "CLIPLoaderGGUF", + "DualCLIPLoaderGGUF", + "TripleCLIPLoaderGGUF", +]) -# Register all modules -register_module("", ["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader","TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"]) -register_module("ComfyUI-GGUF", ["UnetLoaderGGUF", "UnetLoaderGGUFAdvanced", "CLIPLoaderGGUF","DualCLIPLoaderGGUF", "TripleCLIPLoaderGGUF"]) register_module("x-flux-comfyui", ["LoadFluxControlNet"]) register_module("ComfyUI-Florence2", ["Florence2ModelLoader", "DownloadAndLoadFlorence2Model"]) register_module("ComfyUI-LTXVideo", ["LTXVLoader"]) -register_module("ComfyUI-MMAudio", [ "MMAudioFeatureUtilsLoader", "MMAudioModelLoader", "MMAudioSampler"]) +register_module("ComfyUI-MMAudio", ["MMAudioFeatureUtilsLoader", "MMAudioModelLoader", "MMAudioSampler"]) register_module("ComfyUI_bitsandbytes_NF4", ["CheckpointLoaderNF4"]) -logging.info("MultiGPU: Registration complete") -logging.info(f"MultiGPU: Registered nodes: {', '.join(sorted(NODE_CLASS_MAPPINGS.keys()))}") \ No newline at end of file +logging.info("MultiGPU: Registration complete.") +logging.info("MultiGPU: Final mappings: " + ", ".join(sorted(NODE_CLASS_MAPPINGS.keys()))) From fcf054006d2f346ee1176cf75ebf57e293071547 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 30 Dec 2024 11:03:55 -0600 Subject: [PATCH 3/5] Mew method merged in with old code base so less of a shock. --- __init__.py | 156 +++++++++++++++++----------------------------------- 1 file changed, 51 insertions(+), 105 deletions(-) diff --git a/__init__.py b/__init__.py index 6ab8f88..fd278ed 100644 --- a/__init__.py +++ b/__init__.py @@ -6,9 +6,6 @@ import os import importlib.util import logging -############################################################################## -# INITIAL SETUP -############################################################################## logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logging.info("MultiGPU: Initialization started") @@ -26,17 +23,11 @@ def get_torch_device_patched(): comfy.model_management.get_torch_device = get_torch_device_patched -############################################################################## -# OVERRIDE CLASS -############################################################################## def override_class(cls): class NodeOverride(cls): @classmethod def INPUT_TYPES(s): inputs = copy.deepcopy(cls.INPUT_TYPES()) - # In case some node forgot "required" - if "required" not in inputs: - inputs["required"] = {} devices = ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())] inputs["required"]["device"] = (devices,) return inputs @@ -44,117 +35,72 @@ def override_class(cls): CATEGORY = "multigpu" FUNCTION = "override" - def override(self, *args, device="cpu", **kwargs): + def override(self, *args, device, **kwargs): global current_device current_device = device fn = getattr(super(), cls.FUNCTION) return fn(*args, **kwargs) + return NodeOverride -############################################################################## -# OUR LOCAL MAPPING OF MULTIGPU OVERRIDE NODES -# (No strict reason it has to be before the function defs, but it's typical.) -############################################################################## NODE_CLASS_MAPPINGS = {} -############################################################################## -# PART 1: CORE NODES -############################################################################## -def register_core_nodes(core_node_names): - """ - Uses ComfyUI's GLOBAL_NODE_CLASS_MAPPINGS to wrap core nodes. - """ - logging.info("MultiGPU: Starting core node registration") +def register_module(module_path, target_nodes): try: - from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS - except ImportError as e: - logging.error(f"MultiGPU: Could not import ComfyUI global node mappings: {e}") - return + # For core nodes, skip module loading and just register from the global mappings + if not module_path: + logging.info("MultiGPU: Starting core node registration") + from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS + for node in target_nodes: + if node in GLOBAL_NODE_CLASS_MAPPINGS: + NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) + logging.info(f"MultiGPU: Registered core node {node}") + else: + logging.info(f"MultiGPU: Core node {node} not found - this shouldn't happen!") + return - for node in core_node_names: - if node in GLOBAL_NODE_CLASS_MAPPINGS: - NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) - logging.info(f"MultiGPU: Registered core node {node}") - else: - logging.info(f"MultiGPU: Core node '{node}' not found in ComfyUI global mappings") + # For custom nodes, try to load the module first + full_path = os.path.join("custom_nodes", module_path, "__init__.py") + logging.info(f"MultiGPU: Checking for module at {full_path}") -############################################################################## -# PART 2: CUSTOM NODES (LOCAL DICTIONARY APPROACH, with extra debug logs) -############################################################################## -def register_module(module_path: str, target_nodes: list, local_map_name="NODE_CLASS_MAPPINGS"): - """ - 1) Load custom_nodes//__init__.py via importlib - 2) Grab the dictionary named `local_map_name` (default: 'NODE_CLASS_MAPPINGS') from that module - 3) For each node in target_nodes, if found in that local dictionary, wrap it - 4) No fallback to ComfyUI's global mappings - """ - base_dir = os.path.join("custom_nodes", module_path) - init_file = os.path.join(base_dir, "__init__.py") + if not os.path.exists(full_path): + logging.info(f"MultiGPU: Module {module_path} not found - skipping") + return - logging.info(f"MultiGPU: Checking custom node module at: {init_file}") - if not os.path.exists(init_file): - logging.info(f"MultiGPU: Module {module_path} not found or missing __init__.py, skipping.") - return - - try: - logging.info(f"MultiGPU: Found {module_path}, loading local dictionary from {init_file}") - spec = importlib.util.spec_from_file_location(module_path, init_file) + logging.info(f"MultiGPU: Found {module_path}, attempting to load") + spec = importlib.util.spec_from_file_location(module_path, full_path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) logging.info(f"MultiGPU: Executed {module_path} initialization") + + # Use the module's local dictionary instead of the global one + local_map_name = "NODE_CLASS_MAPPINGS" + local_map = getattr(module, local_map_name, None) + if not local_map: + logging.info(f"MultiGPU: {module_path} has no '{local_map_name}' dictionary, skipping override.") + return + + all_defined_nodes = list(local_map.keys()) + logging.info(f"MultiGPU: {module_path} local dict keys: {all_defined_nodes}") + + for node in target_nodes: + if node in local_map: + mgpu_class = override_class(local_map[node]) + NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = mgpu_class + logging.info(f"MultiGPU: Successfully wrapped {node} from {module_path}") + else: + logging.info(f"MultiGPU: Node '{node}' not found in {module_path}'s local dictionary") + except Exception as e: - logging.error(f"MultiGPU: Error loading {module_path}: {e}") - return + logging.info(f"MultiGPU: Error processing {module_path}: {str(e)}") - # Grab that module's local dictionary (e.g. NODE_CLASS_MAPPINGS) - local_map = getattr(module, local_map_name, None) - if not local_map: - logging.info(f"MultiGPU: {module_path} has no '{local_map_name}' dictionary, skipping override.") - return +# Register desired nodes +register_module("", ["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"]) +register_module("ComfyUI-GGUF", ["UnetLoaderGGUF","UnetLoaderGGUFAdvanced","CLIPLoaderGGUF","DualCLIPLoaderGGUF","TripleCLIPLoaderGGUF"]) +register_module("x-flux-comfyui", ["LoadFluxControlNet"]) +register_module("ComfyUI-Florence2", ["Florence2ModelLoader","DownloadAndLoadFlorence2Model"]) +register_module("ComfyUI-LTXVideo", ["LTXVLoader"]) +register_module("ComfyUI-MMAudio", ["MMAudioFeatureUtilsLoader","MMAudioModelLoader","MMAudioSampler"]) +register_module("ComfyUI_bitsandbytes_NF4", ["CheckpointLoaderNF4",]) - # DEBUG: Show everything this node map provides - all_defined_nodes = list(local_map.keys()) - logging.info(f"MultiGPU: {module_path} local dict keys: {all_defined_nodes}") - - # Wrap each node we want - for node in target_nodes: - if node in local_map: - mgpu_class = override_class(local_map[node]) - NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = mgpu_class - logging.info(f"MultiGPU: Successfully wrapped {node} from {module_path}") - else: - logging.info(f"MultiGPU: Node '{node}' not found in {module_path}'s local dictionary") - - -############################################################################## -# EXAMPLE USAGE -############################################################################## - -# 1) CORE NODES -register_core_nodes([ - "UNETLoader", - "VAELoader", - "CLIPLoader", - "DualCLIPLoader", - "TripleCLIPLoader", - "CheckpointLoaderSimple", - "ControlNetLoader", -]) - -# 2) CUSTOM NODES -register_module("ComfyUI-GGUF", [ - "UnetLoaderGGUF", - "UnetLoaderGGUFAdvanced", - "CLIPLoaderGGUF", - "DualCLIPLoaderGGUF", - "TripleCLIPLoaderGGUF", -]) - -register_module("x-flux-comfyui", ["LoadFluxControlNet"]) -register_module("ComfyUI-Florence2", ["Florence2ModelLoader", "DownloadAndLoadFlorence2Model"]) -register_module("ComfyUI-LTXVideo", ["LTXVLoader"]) -register_module("ComfyUI-MMAudio", ["MMAudioFeatureUtilsLoader", "MMAudioModelLoader", "MMAudioSampler"]) -register_module("ComfyUI_bitsandbytes_NF4", ["CheckpointLoaderNF4"]) - -logging.info("MultiGPU: Registration complete.") -logging.info("MultiGPU: Final mappings: " + ", ".join(sorted(NODE_CLASS_MAPPINGS.keys()))) +logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") \ No newline at end of file From 198365dfbefa5d587102c9c373fdc505b85352ac Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 30 Dec 2024 20:48:41 -0600 Subject: [PATCH 4/5] Add hard-coded registration for LTX and Florence2 nodes in MultiGPU setup for debug purposes. Actual nodes pick up the underlying structure at runtime now that the global NODE_CLASS_MAPPINGS has been updated with their information, I pull it directly from there. A work-around for the loading sequencing problems, but hopefully one that requrires little upkeep as any changes to the underlying structure is picked-up at runtime. --- __init__.py | 127 +++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 125 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index fd278ed..c7f4f10 100644 --- a/__init__.py +++ b/__init__.py @@ -3,8 +3,10 @@ import copy import torch import comfy.model_management import os +from pathlib import Path # Add this import import importlib.util import logging +import folder_paths logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logging.info("MultiGPU: Initialization started") @@ -94,12 +96,133 @@ def register_module(module_path, target_nodes): except Exception as e: logging.info(f"MultiGPU: Error processing {module_path}: {str(e)}") +def register_LTXmodule(module_path, node_list): + """Hard-coded registration for LTX Video nodes""" + 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): + # Get original node instance + from nodes import NODE_CLASS_MAPPINGS + original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]() + + # Use original node to load model and VAE + model, vae = original_loader.load(ckpt_name, dtype) + + # Return original objects + return (model, vae) + + ltx_nodes = { + "LTXVLoader": LTXVLoader + } + + for node_name in node_list: + if node_name in ltx_nodes: + NODE_CLASS_MAPPINGS[f"{node_name}MultiGPU"] = override_class(ltx_nodes[node_name]) + logging.info(f"MultiGPU: Registered hard-coded LTX node {node_name}") + +def register_Florence2module(module_path, node_list): + """Hard-coded registration for Florence2 nodes""" + 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): + # Get original node instance + from nodes import NODE_CLASS_MAPPINGS + original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]() + + # Use original node to load model + return original_loader.loadmodel(model, precision, attention, lora) + + 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): + # Get original node instance + from nodes import NODE_CLASS_MAPPINGS + original_loader = NODE_CLASS_MAPPINGS["Florence2ModelLoader"]() + + # Use original node to load model + return original_loader.loadmodel(model, precision, attention, lora) + + florence2_nodes = { + "Florence2ModelLoader": Florence2ModelLoader, + "DownloadAndLoadFlorence2Model": DownloadAndLoadFlorence2Model + } + + for node_name in node_list: + if node_name in florence2_nodes: + NODE_CLASS_MAPPINGS[f"{node_name}MultiGPU"] = override_class(florence2_nodes[node_name]) + logging.info(f"MultiGPU: Registered hard-coded Florence2 node {node_name}") + # Register desired nodes register_module("", ["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"]) register_module("ComfyUI-GGUF", ["UnetLoaderGGUF","UnetLoaderGGUFAdvanced","CLIPLoaderGGUF","DualCLIPLoaderGGUF","TripleCLIPLoaderGGUF"]) register_module("x-flux-comfyui", ["LoadFluxControlNet"]) -register_module("ComfyUI-Florence2", ["Florence2ModelLoader","DownloadAndLoadFlorence2Model"]) -register_module("ComfyUI-LTXVideo", ["LTXVLoader"]) +register_Florence2module("ComfyUI-Florence2", ["Florence2ModelLoader", "DownloadAndLoadFlorence2Model"]) +register_LTXmodule("ComfyUI-LTXVideo", ["LTXVLoader"]) register_module("ComfyUI-MMAudio", ["MMAudioFeatureUtilsLoader","MMAudioModelLoader","MMAudioSampler"]) register_module("ComfyUI_bitsandbytes_NF4", ["CheckpointLoaderNF4",]) From 1e553cfe224b1e611e3595a6b2ab2fe6e47006b3 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 30 Dec 2024 21:00:48 -0600 Subject: [PATCH 5/5] Add utility function to check module existence before registration for hard-coded MultiGPU nodes --- __init__.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/__init__.py b/__init__.py index c7f4f10..50bf97c 100644 --- a/__init__.py +++ b/__init__.py @@ -47,6 +47,19 @@ def override_class(cls): NODE_CLASS_MAPPINGS = {} + +def check_module_exists(module_path): + """Utility function to check if module exists""" + full_path = os.path.join("custom_nodes", module_path, "__init__.py") + 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}, attempting to load") + return True + def register_module(module_path, target_nodes): try: # For core nodes, skip module loading and just register from the global mappings @@ -100,6 +113,9 @@ def register_LTXmodule(module_path, node_list): """Hard-coded registration for LTX Video nodes""" global NODE_CLASS_MAPPINGS + if not check_module_exists(module_path): + return + class LTXVLoader: @classmethod def INPUT_TYPES(s): @@ -142,6 +158,9 @@ def register_Florence2module(module_path, node_list): """Hard-coded registration for Florence2 nodes""" global NODE_CLASS_MAPPINGS + if not check_module_exists(module_path): + return + class DownloadAndLoadFlorence2Model: @classmethod def INPUT_TYPES(s):