diff --git a/README.md b/README.md index 2fc1c3e..036d29a 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ ### Experimental nodes for using multiple GPUs in a single ComfyUI workflow. -This extension adds new nodes for model loading that allow you to specify the GPU to use for each model. It monkey patches the memory management of ComfyUI in a hacky way and is neither a comprehensive solution nor a well-tested one. Use at your own risk. +This extension adds device selection capabilities to model loading nodes in ComfyUI. It monkey patches the memory management of ComfyUI in a hacky way and is neither a comprehensive solution nor a well-tested one. Use at your own risk. Note that this does not add parallelism. The workflow steps are still executed sequentially just on different GPUs. Any potential speedup comes from not having to constantly load and unload models from VRAM. @@ -12,16 +12,29 @@ Clone this repository inside `ComfyUI/custom_nodes/`. ## Nodes -![](examples/nodes.png) +The extension automatically creates MultiGPU versions of loader nodes. Each MultiGPU node has the same functionality as its original counterpart but adds a `device` parameter that allows you to specify the GPU to use. -The extension adds new loader nodes corresponding to the default ones. The new nodes have the same functionality but add a new `device` parameter that allows you to specify the GPU to use. +Currently supported nodes (automatically detected if available): +- Standard ComfyUI loaders: + - CheckpointLoaderSimpleMultiGPU + - CLIPLoaderMultiGPU + - ControlNetLoaderMultiGPU + - DualCLIPLoaderMultiGPU + - TripleCLIPLoaderMultiGPU + - UNETLoaderMultiGPU + - VAELoaderMultiGPU -- `CheckpointLoaderMultiGPU` -- `CLIPLoaderMultiGPU` -- `ControlNetLoaderMultiGPU` -- `DualCLIPLoaderMultiGPU` -- `UNETLoaderMultiGPU` -- `VAELoaderMultiGPU` +- GGUF loaders (requires [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF)): + - UnetLoaderGGUFMultiGPU (supports quantized models like [flux1-dev-gguf](https://huggingface.co/city96/FLUX.1-dev-gguf)) + - UnetLoaderGGUFAdvancedMultiGPU + - CLIPLoaderGGUFMultiGPU + - DualCLIPLoaderGGUFMultiGPU + - TripleCLIPLoaderGGUFMultiGPU + +- Additional supported nodes: + - LoadFluxControlNet (requires [x-flux-comfy](https://github.com/XLabAI/x-flux-comfyui)) + +All MultiGPU nodes can be found in the "multigpu" category in the node menu. ## Example workflows @@ -47,10 +60,10 @@ This workflow loads a FLUX.1-dev model and an SDXL model in the same workflow. T ## Support -If you encounter problems, please [open an issue](https://github.com/neuratech-ai/ComfyUI-MultiGPU/issues/new). Attach the workflow if possible. +If you encounter problems, please [open an issue](https://github.com/pollockjj/ComfyUI-MultiGPU/issues/new). Attach the workflow if possible. ## Credits -Made by [Alexander Dzhoganov](https://github.com/AlexanderDzhoganov). - -For business inquiries, email [sales@neuratech.io](mailto:sales@neuratech.io) or visit [our website](https://neuratech.io/). +Originally created by [Alexander Dzhoganov](https://github.com/AlexanderDzhoganov). +Implementation improved by [City96](https://v100s.net/). +Currently maintained by [pollockjj](https://github.com/pollockjj). \ No newline at end of file diff --git a/__init__.py b/__init__.py index c92b572..f87ba6e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,18 +1,16 @@ import time +import copy import torch - -from copy import deepcopy import comfy.model_management -current_device = "cuda:0" +current_device = comfy.model_management.get_torch_device() + def get_torch_device_patched(): - global current_device if ( not torch.cuda.is_available() or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU ): return torch.device("cpu") - return torch.device(current_device) comfy.model_management.get_torch_device = get_torch_device_patched @@ -21,7 +19,7 @@ def override_class(cls): class NodeOverride(cls): @classmethod def INPUT_TYPES(s): - inputs = deepcopy(cls.INPUT_TYPES()) + inputs = copy.deepcopy(cls.INPUT_TYPES()) inputs["required"]["device"] = ([f"cuda:{i}" for i in range(torch.cuda.device_count())],) return inputs @@ -33,24 +31,20 @@ def override_class(cls): current_device = device fn = getattr(super(), cls.FUNCTION) return fn(*args, **kwargs) - + return NodeOverride -# This is to make sure the other nodes are already loaded -time.sleep(2) +time.sleep(2) # This is to make sure the other nodes are already loaded from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS TARGET_NODE_NAMES = { - "UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoader", - "UnetLoaderGGUF", "UnetLoaderGGUFAdvanced", "CLIPLoaderGGUF", "DualCLIPLoaderGGUF", "TripleCLIPLoaderGGUF", + "UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader", + "UnetLoaderGGUF", "UnetLoaderGGUFAdvanced", "CLIPLoaderGGUF", "DualCLIPLoaderGGUF", "TripleCLIPLoaderGGUF", + "LoadFluxControlNet", } -TARGET_NODE_CLASS_MAPPINGS = {} +NODE_CLASS_MAPPINGS = {} for name in TARGET_NODE_NAMES: if name not in GLOBAL_NODE_CLASS_MAPPINGS: continue - TARGET_NODE_CLASS_MAPPINGS[name] = GLOBAL_NODE_CLASS_MAPPINGS[name] - -NODE_CLASS_MAPPINGS = {} -for name, cls in TARGET_NODE_CLASS_MAPPINGS.items(): - NODE_CLASS_MAPPINGS[f"{name}_MultiGPU"] = override_class(cls) \ No newline at end of file + NODE_CLASS_MAPPINGS[f"{name}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[name])