commit e39e19fc467f5fb1a48a26af4c65871f6490b4b3 Author: FearL0rd Date: Mon Feb 2 10:48:52 2026 -0600 initial release diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..b36c781 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .any_device_parallel import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/any_device_parallel.py b/any_device_parallel.py new file mode 100644 index 0000000..edc1184 --- /dev/null +++ b/any_device_parallel.py @@ -0,0 +1,334 @@ +import torch +import torch.nn as nn +import types +import copy +from concurrent.futures import ThreadPoolExecutor, as_completed +import comfy.model_management + +class ParallelDevice: + """Individual device configuration node with auto-detected hardware.""" + + @classmethod + def get_available_devices(cls): + """Dynamically detect all available compute devices.""" + devices = ["cpu"] # CPU always available + + # NVIDIA CUDA (also covers AMD ROCm, which uses CUDA interface in PyTorch) + if torch.cuda.is_available(): + count = torch.cuda.device_count() + for i in range(count): + name = torch.cuda.get_device_name(i) + # Format: "cuda:0 (NVIDIA RTX 4090)" truncated for readability + safe_name = name[:25] + "..." if len(name) > 28 else name + devices.append(f"cuda:{i}") + + # Apple Silicon (Metal Performance Shaders) + if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): + devices.append("mps") + + # Intel XPU (Arc GPUs, etc.) + if hasattr(torch, 'xpu') and torch.xpu.is_available(): + count = torch.xpu.device_count() + for i in range(count): + devices.append(f"xpu:{i}") + + # DirectML (Windows ML devices) + try: + import torch_directml + dml_count = torch_directml.device_count() + for i in range(dml_count): + devices.append(f"privateuseone:{i}") # DirectML device naming + except ImportError: + pass + + return devices + + @classmethod + def INPUT_TYPES(s): + available = s.get_available_devices() + # Prefer cuda:0 as default if available + default = "cuda:0" if any(d.startswith("cuda:0") for d in available) else available[0] + + return { + "required": { + "device_id": (available, { + "default": default, + "tooltip": "Select available compute device (CPU/CUDA/MPS/XPU)" + }), + "percentage": ("FLOAT", { + "default": 50.0, + "min": 1.0, + "max": 100.0, + "step": 1.0, + "tooltip": "Percentage of batch to process on this device" + }), + }, + "optional": { + "previous_devices": ("DEVICE_CHAIN", { + "tooltip": "Connect from another ParallelDevice node to chain multiple GPUs" + }), + } + } + + RETURN_TYPES = ("DEVICE_CHAIN",) + RETURN_NAMES = ("device_chain",) + FUNCTION = "add_device" + CATEGORY = "utils/hardware" + DESCRIPTION = "Add a GPU/CPU/MPS/XPU device to the parallel processing chain" + + def add_device(self, device_id, percentage, previous_devices=None): + if previous_devices is None: + previous_devices = [] + + # device_id comes directly from dropdown selection + config = { + "device": device_id, + "percentage": float(percentage), + "weight": float(percentage) / 100.0 + } + + # Append to chain + new_chain = previous_devices.copy() + new_chain.append(config) + + return (new_chain,) + +class ParallelDeviceList: + """Alternative: Parallel layout with dropdowns for 1-4 devices.""" + + @classmethod + def get_available_devices(cls): + """Get devices for dropdown options.""" + devices = ["cpu"] + if torch.cuda.is_available(): + for i in range(torch.cuda.device_count()): + devices.append(f"cuda:{i}") + if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): + devices.append("mps") + if hasattr(torch, 'xpu') and torch.xpu.is_available(): + for i in range(torch.xpu.device_count()): + devices.append(f"xpu:{i}") + return devices + + @classmethod + def INPUT_TYPES(s): + devices = s.get_available_devices() + def_dev = "cuda:0" if "cuda:0" in devices else devices[0] + + return { + "required": { + "device_1": (devices, {"default": def_dev}), + "pct_1": ("FLOAT", {"default": 50.0, "min": 1.0, "max": 100.0, "step": 1.0}), + "device_2": (devices, {"default": devices[1] if len(devices) > 1 else def_dev}), + "pct_2": ("FLOAT", {"default": 50.0, "min": 0.0, "max": 100.0, "step": 1.0}), + }, + "optional": { + "device_3": (devices, {"default": devices[2] if len(devices) > 2 else "cpu"}), + "pct_3": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 1.0}), + "device_4": (devices, {"default": devices[3] if len(devices) > 3 else "cpu"}), + "pct_4": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 1.0}), + } + } + + RETURN_TYPES = ("DEVICE_CHAIN",) + RETURN_NAMES = ("device_chain",) + FUNCTION = "create_list" + CATEGORY = "utils/hardware" + + def create_list(self, device_1, pct_1, device_2, pct_2, device_3="cpu", pct_3=0, device_4="cpu", pct_4=0): + chain = [] + devices = [(device_1, pct_1), (device_2, pct_2), (device_3, pct_3), (device_4, pct_4)] + + for dev_str, pct in devices: + if pct > 0: + chain.append({ + "device": dev_str, + "percentage": float(pct), + "weight": float(pct) / 100.0 + }) + return (chain,) + +class ParallelAnything: + """Main execution node with TRUE parallelism via model replication.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "device_chain": ("DEVICE_CHAIN", {"tooltip": "Connect from ParallelDevice nodes"}), + }, + "optional": { + "workload_split": ("BOOLEAN", {"default": True, "tooltip": "Enable multi-device processing"}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "setup_parallel" + CATEGORY = "utils/hardware" + + def setup_parallel(self, model, device_chain, workload_split=True): + if model is None or not device_chain: + return (model,) + + # Normalize weights + total_pct = sum(item["percentage"] for item in device_chain) + if total_pct <= 0: + return (model,) + + device_names = [] + weights = [] + for item in device_chain: + weights.append(item["percentage"] / total_pct) + device_names.append(item["device"]) + + print(f"[ParallelAnything] Setup: {list(zip(device_names, [f'{w*100:.1f}%' for w in weights]))}") + + lead_device = torch.device(device_names[0]) + target_model = model.model.diffusion_model + + if hasattr(target_model, "_true_parallel_active"): + return (model,) + + # Validate devices + for dev in device_names: + try: + torch.device(dev) + except: + print(f"[ParallelAnything] Invalid device: {dev}") + return (model,) + + # Create replicas (N× VRAM) + replicas = {} + try: + print(f"[ParallelAnything] Cloning to {len(device_names)} devices...") + state_dict = copy.deepcopy(target_model.state_dict()) + + for dev_name in device_names: + dev = torch.device(dev_name) + + try: + model_class = target_model.__class__ + if hasattr(target_model, 'config'): + replica = model_class(**target_model.config) + elif hasattr(target_model, 'unet_config'): + replica = model_class(**target_model.unet_config) + else: + replica = model_class() + replica.load_state_dict(state_dict, strict=False) + except Exception as e: + replica = copy.deepcopy(target_model) + + replica.to(dev).eval() + for param in replica.parameters(): + param.requires_grad = False + replicas[dev_name] = replica + print(f"[ParallelAnything] ✓ {dev_name}") + + except RuntimeError as e: + print(f"[ParallelAnything] VRAM Error: {e}") + return (model,) + + # Store references for closure + replicas_ref = replicas + devices_ref = device_names + weights_ref = weights + lead_ref = lead_device + + def parallel_forward(self, x, timesteps, context=None, **kwargs): + batch_size = x.shape[0] + + if batch_size < len(devices_ref) or not workload_split: + with torch.no_grad(): + return replicas_ref[devices_ref[0]](x, timesteps, context=context, **kwargs) + + # Calculate splits + split_sizes = [max(1, int(batch_size * w)) for w in weights_ref] + split_sizes[-1] = batch_size - sum(split_sizes[:-1]) + + active = [] + for idx, (dev_name, size) in enumerate(zip(devices_ref, split_sizes)): + if size > 0: + active.append({ + 'idx': idx, + 'dev_name': dev_name, + 'device': torch.device(dev_name), + 'replica': replicas_ref[dev_name], + 'size': size + }) + + if len(active) == 1: + with torch.no_grad(): + return active[0]['replica'](x, timesteps, context=context, **kwargs) + + # Split tensors + x_chunks = torch.split(x, [a['size'] for a in active], dim=0) + t_chunks = torch.split(timesteps, [a['size'] for a in active], dim=0) + c_chunks = torch.split(context, [a['size'] for a in active], dim=0) if context is not None else [None] * len(active) + + results = [None] * len(active) + + def worker(task_idx): + task = active[task_idx] + dev = task['device'] + replica = task['replica'] + + try: + # Handle device context + if dev.type == 'cuda' or dev.type == 'xpu': + with torch.cuda.device(dev) if dev.type == 'cuda' else torch.xpu.device(dev): + torch.cuda.synchronize(dev) if dev.type == 'cuda' else torch.xpu.synchronize(dev) + + x_in = x_chunks[task_idx].to(dev, non_blocking=True) + t_in = t_chunks[task_idx].to(dev, non_blocking=True) + c_in = c_chunks[task_idx].to(dev, non_blocking=True) if c_chunks[task_idx] is not None else None + + with torch.no_grad(): + out = replica(x_in, t_in, context=c_in, **kwargs) + + torch.cuda.synchronize(dev) if dev.type == 'cuda' else torch.xpu.synchronize(dev) + return task_idx, out.to(lead_ref, non_blocking=False) + else: + # CPU/MPS don't need explicit context managers + x_in = x_chunks[task_idx].to(dev) + t_in = t_chunks[task_idx].to(dev) + c_in = c_chunks[task_idx].to(dev) if c_chunks[task_idx] is not None else None + + with torch.no_grad(): + out = replica(x_in, t_in, context=c_in, **kwargs) + return task_idx, out.to(lead_ref) + + except Exception as e: + return task_idx, e + + # TRUE PARALLEL EXECUTION + with ThreadPoolExecutor(max_workers=len(active)) as executor: + futures = [executor.submit(worker, i) for i in range(len(active))] + for future in as_completed(futures): + idx, result = future.result() + if isinstance(result, Exception): + raise result + results[idx] = result + + return torch.cat(results, dim=0) + + target_model._original_forward = target_model.forward + target_model.forward = types.MethodType(parallel_forward, target_model) + target_model._true_parallel_active = True + target_model._parallel_replicas = replicas + model.load_device = lead_device + + return (model,) + +NODE_CLASS_MAPPINGS = { + "ParallelAnything": ParallelAnything, + "ParallelDevice": ParallelDevice, + "ParallelDeviceList": ParallelDeviceList +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ParallelAnything": "Parallel Anything (True Multi-GPU)", + "ParallelDevice": "Parallel Device Config", + "ParallelDeviceList": "Parallel Device List (1-4x)" +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ff41503 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +# pyproject.toml +[project] +name = "ComfyUI-ParallelAnything" # Unique identifier for your node. Immutable after creation. +description = "This suite of nodes unlocks high-performance parallel processing in ComfyUI by utilizing **Model Replication**. Unlike standard offloading which moves a single model instance between devices, these nodes create independent replicas of the model on each selected GPU/CPU, allowing for true simultaneous batch processing." +version = "1.0.0" # Custom Node version. Must be semantically versioned. +license = { file = "LICENSE.txt" } +dependencies = [] # Filled in from requirements.txt + +[project.urls] +Repository = "https://github.com/FearL0rd/ComfyUI-ParallelAnything" + +[tool.comfy] +PublisherId = "fearl0rd" # TODO (fill in Publisher ID from Comfy Registry Website). +DisplayName = "ComfyUI-ParallelAnything" # Display name for the Custom Node. Can be changed later. +Icon = "https://raw.githubusercontent.com/FearL0rd/ComfyUI-ParallelAnything/refs/heads/main/icon.png" # SVG, PNG, JPG or GIF (MAX. 800x400px) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b1e3a74 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +torch +types +gc +sys \ No newline at end of file