refactor: streamline MultiGPU node initialization
Code Changes (__init__.py): - Switches from deepcopy to standard copy for better efficiency - Initializes current_device from model_management instead of hardcoded "cuda:0" - Simplifies node class mapping by removing intermediate TARGET_NODE_CLASS_MAPPINGS - Updates node naming convention: removes underscore from MultiGPU suffix - Adds new supported nodes: "CheckpointLoaderSimple", "ControlNetLoader", "LoadFluxControlNet" - Improves code organization with better comment clarity COMPATIBILITY NOTE: This version restores backward compatibility with workflows using the previous node naming scheme. Both old and new node names will work. Documentation Changes (README.md): - Updates node list to reflect automatic detection system - Adds proper attribution links for required dependencies (ComfyUI-GGUF, x-flux-comfy) - Links to example quantized models like flux1-dev-gguf - Reorganizes loader sections with clear dependency requirements - Updates support links to new maintainer - Removes business/commercial references - Updates credits section to reflect current project status
This commit is contained in:
John Pollock
@@ -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
|
||||
|
||||
|
||||
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).
|
||||
+11
-17
@@ -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)
|
||||
NODE_CLASS_MAPPINGS[f"{name}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[name])
|
||||
|
||||
Reference in New Issue
Block a user