import torch import torch.nn as nn import types import copy import gc import weakref import threading from types import SimpleNamespace from dataclasses import fields, is_dataclass from concurrent.futures import ThreadPoolExecutor, as_completed import comfy.model_management # ---------------------------------------------------------------------------- # Thread-local state for pipeline mode # ---------------------------------------------------------------------------- _pipeline_state = threading.local() def get_pipeline_mode(): return getattr(_pipeline_state, 'active', False) def set_pipeline_mode(active): _pipeline_state.active = active def as_torch_device(device): if isinstance(device, torch.device): return device return torch.device(device) def move_tree_to_device(x, device, non_blocking=False): """Recursively move tensors inside nested Python containers to `device`.""" device = as_torch_device(device) if isinstance(x, torch.Tensor): if x.device != device: x = x.to(device, non_blocking=non_blocking) if is_float8_dtype(x.dtype) and not device_supports_float8(device): x = x.half() return x if isinstance(x, (list, tuple)): return type(x)(move_tree_to_device(item, device, non_blocking) for item in x) if isinstance(x, dict): return {k: move_tree_to_device(v, device, non_blocking) for k, v in x.items()} if is_dataclass(x) and not isinstance(x, type): new_obj = copy.copy(x) for field in fields(x): val = getattr(x, field.name) setattr(new_obj, field.name, move_tree_to_device(val, device, non_blocking)) return new_obj if isinstance(x, SimpleNamespace): return SimpleNamespace(**{ k: move_tree_to_device(v, device, non_blocking) for k, v in vars(x).items() }) return x BLOCK_LIST_NAMES = ('double_blocks', 'single_blocks', 'transformer_blocks', 'layers') _PATCH_KEY_PREFIXES = ( "diffusion_model.", "model.diffusion_model.", "model.", ) def resolve_home_device(model_wrapper, target_model, device_names): """Device ComfyUI will load the live model onto. After unload_all_models() parameters look like they live on CPU, so we must remember the real home GPU *before* using that as a clone target. Cloning onto the home device puts two full copies on one card and OOMs (e.g. Lumina2 11.7GB x2 on a 24GB GPU). """ candidates = [] if model_wrapper is not None: candidates.append(getattr(model_wrapper, "load_device", None)) inner = getattr(model_wrapper, "model", None) if inner is not None: candidates.append(getattr(inner, "load_device", None)) try: candidates.append(next(target_model.parameters()).device) except (StopIteration, AttributeError): pass for candidate in candidates: if candidate is None: continue name = str(candidate) if name != "cpu": return name for name in device_names or (): if name != "cpu": return name if device_names: return device_names[0] return "cpu" def should_reuse_live_model(dev_name, home_device_str): return str(dev_name) == str(home_device_str) def find_model_patcher(model): if model is None: return None if hasattr(model, "patches") and hasattr(model, "add_patches"): return model for attr in ("model_patcher", "patcher"): inner = getattr(model, attr, None) if inner is not None and hasattr(inner, "patches") and hasattr(inner, "add_patches"): return inner return None def collect_weight_patches(model): """Return ModelPatcher.patches without baking them into the live model.""" patcher = find_model_patcher(model) if patcher is None: return {} patches = getattr(patcher, "patches", None) return patches if patches else {} def snapshot_weight_patches(patches): """Copy the patch mapping so later disarm/clear cannot empty our bake list.""" if not patches: return {} return {key: list(value) for key, value in patches.items()} def disarm_patcher_weight_patches(patcher): """Detach patches so ComfyUI dynamic-VRAM load does not apply LoRA a second time.""" if patcher is None or not hasattr(patcher, "patches"): return None saved = patcher.patches if not saved: return None patcher.patches = {} return saved def restore_patcher_weight_patches(patcher, saved): if patcher is None or saved is None: return patcher.patches = saved def bake_lora_with_patcher(patcher, device="cpu"): """Apply LoRA once through ComfyUI's own loader. Writing patched values into the live model with param.data.copy_ can segfault on quantized / dynamic-VRAM parameters. patch_model() knows how to handle those tensors. """ if patcher is None or not getattr(patcher, "patches", None): return False if not hasattr(patcher, "patch_model"): return False patcher.patch_model(device_to=torch.device(device), force_patch_weights=True) return True def resolve_module_param_key(patch_key, named_params): if patch_key in named_params: return patch_key for prefix in _PATCH_KEY_PREFIXES: if patch_key.startswith(prefix): stripped = patch_key[len(prefix):] if stripped in named_params: return stripped return None def _patch_key_candidates(patch_key): keys = [patch_key] for prefix in _PATCH_KEY_PREFIXES: if patch_key.startswith(prefix): keys.append(patch_key[len(prefix):]) return keys def _unwrap_parallel(obj): while isinstance(obj, ParallelBlock): obj = obj.local_block return obj def _lookup_module_tensor(module, dotted_key): obj = module try: for part in dotted_key.split("."): obj = _unwrap_parallel(getattr(_unwrap_parallel(obj), part)) except AttributeError: return None obj = _unwrap_parallel(obj) if isinstance(obj, (nn.Parameter, torch.Tensor)): return obj return None def _lookup_module(module, dotted_key): if not dotted_key: return _unwrap_parallel(module) obj = module try: for part in dotted_key.split("."): obj = _unwrap_parallel(getattr(_unwrap_parallel(obj), part)) except AttributeError: return None obj = _unwrap_parallel(obj) return obj if isinstance(obj, nn.Module) else None def _module_and_hook_attr(root, tensor_key): for candidate in _patch_key_candidates(tensor_key): if candidate.endswith(".weight"): mod = _lookup_module(root, candidate[: -len(".weight")]) if mod is not None: return mod, "weight_function" if candidate.endswith(".bias"): mod = _lookup_module(root, candidate[: -len(".bias")]) if mod is not None: return mod, "bias_function" return None, None def copy_lora_cast_hooks(source, replica): """Copy ComfyUI lazy-LoRA hooks that state_dict cloning drops.""" if source is None or replica is None or source is replica: return 0 src_map = dict(source.named_modules()) copied = 0 for name, dst in replica.named_modules(): src = src_map.get(name) if src is None: continue for attr in ("weight_function", "bias_function"): val = getattr(src, attr, None) if val: setattr(dst, attr, list(val)) copied += 1 for attr in ("comfy_cast_weights", "comfy_force_cast_weights"): if hasattr(src, attr): setattr(dst, attr, getattr(src, attr)) return copied def _attach_lora_hook(module, hook_attr, patch_list, key): def apply_lora(weight): return _calculate_patched_weight(patch_list, weight, key) apply_lora._parallel_lora_key = key hooks = list(getattr(module, hook_attr, None) or []) if any(getattr(fn, "_parallel_lora_key", None) == key for fn in hooks): setattr(module, hook_attr, hooks) module.comfy_cast_weights = True return False hooks.append(apply_lora) setattr(module, hook_attr, hooks) module.comfy_cast_weights = True return True def attach_lora_hooks_to_replica(replica, patches): """Install ComfyUI-style weight_function hooks. No matmuls — safe at setup and every step.""" if replica is None or not patches: return {"attached": 0, "skipped": 0} attached = 0 skipped = 0 for key, patch_list in patches.items(): mod, hook_attr = _module_and_hook_attr(replica, key) if mod is None: skipped += 1 continue if _attach_lora_hook(mod, hook_attr, patch_list, key): attached += 1 else: attached += 1 return {"attached": attached, "skipped": skipped} def apply_patches_to_replica(replica, patches, progress_label=None): """Bake ModelPatcher.patches into a cloned replica's param.data. Must run after incremental clone. Dynamic VRAM often leaves LoRA only in hooks on the live model; state_dict copies are the base checkpoint. Always bake the replica (never the live model). Stay on the replica device. """ if replica is None or not patches: return {"baked": 0, "attached": 0, "failed": 0, "skipped": 0} baked = 0 attached = 0 failed = 0 skipped = 0 total = len(patches) for index, (key, patch_list) in enumerate(patches.items(), start=1): if progress_label and (index == 1 or index == total or index % 30 == 0): print(f"[ParallelAnything] {progress_label} {index}/{total}", flush=True) param = None for candidate in _patch_key_candidates(key): param = _lookup_module_tensor(replica, candidate) if param is not None: break if param is None: skipped += 1 continue try: with torch.no_grad(): work = param.detach() if work.dtype != torch.float32: work = work.float() else: work = work.clone() patched = _calculate_patched_weight(patch_list, work, key) if not isinstance(patched, torch.Tensor) or patched.shape != param.shape: raise RuntimeError( f"bad LoRA result for {key}: {type(patched)} {getattr(patched, 'shape', None)}" ) param.data.copy_(patched.to(device=param.device, dtype=param.dtype)) baked += 1 except Exception as e: mod, hook_attr = _module_and_hook_attr(replica, key) if mod is not None and ( hasattr(mod, "comfy_cast_weights") or hasattr(mod, "weight_function") ): try: _attach_lora_hook(mod, hook_attr, patch_list, key) attached += 1 continue except Exception: pass failed += 1 print(f"[ParallelAnything] Warning: could not apply LoRA {key}: {e}") return {"baked": baked, "attached": attached, "failed": failed, "skipped": skipped} def sync_lora_to_replica(source, replica, patches, weights_already_baked=False): """Compatibility wrapper: always bake patches onto the replica.""" del source, weights_already_baked return apply_patches_to_replica(replica, patches) def _fallback_calculate_weight(patch_list, weight, key): """Subset of comfy.lora.calculate_weight for tests and simple diff patches.""" out = weight for patch in patch_list: if not isinstance(patch, (tuple, list)) or len(patch) < 2: continue strength = patch[0] value = patch[1] strength_model = patch[2] if len(patch) > 2 else 1.0 if not isinstance(strength, (int, float)): continue if strength_model != 1.0: out = out * strength_model if hasattr(value, "calculate_weight"): computed = value.calculate_weight( out, key, strength, strength_model, None, lambda a: a, torch.float32, None ) if computed is None: raise RuntimeError(f"adapter calculate_weight returned None for {key}") out = computed continue if isinstance(value, torch.Tensor): if value.shape == out.shape: out = out + strength * value.to(device=out.device, dtype=out.dtype) continue if not isinstance(value, (tuple, list)) or len(value) == 0: continue patch_type = "diff" data = value if len(value) == 2 and isinstance(value[0], str): patch_type = value[0] data = value[1] payload = data[0] if isinstance(data, (tuple, list)) else data if not isinstance(payload, torch.Tensor): continue if patch_type == "set": out = payload.to(device=out.device, dtype=out.dtype) elif patch_type == "diff" and payload.shape == out.shape: out = out + strength * payload.to(device=out.device, dtype=out.dtype) return out def _comfy_calculate_weight(): try: import comfy.lora as comfy_lora candidate = getattr(comfy_lora, "calculate_weight", None) if callable(candidate) and getattr(candidate, "__module__", "") == "comfy.lora": return candidate except ImportError: pass return None def _calculate_patched_weight(patch_list, weight, key): calc = _comfy_calculate_weight() if calc is not None: return calc(patch_list, weight, key) return _fallback_calculate_weight(patch_list, weight, key) def bake_patches_into_module(module, patches, progress_label=None): """Bake ModelPatcher weight patches onto a module. Looks up tensors by key path so we never call named_parameters() on a dynamically staged model (that pages the full 10GB+ checkpoint onto GPU and looks like a hang). """ if module is None or not patches: return 0 applied = 0 total = len(patches) for index, (key, patch_list) in enumerate(patches.items(), start=1): if progress_label and (index == 1 or index == total or index % 20 == 0): print(f"[ParallelAnything] {progress_label} {index}/{total}", flush=True) param = None for candidate in _patch_key_candidates(key): param = _lookup_module_tensor(module, candidate) if param is not None: break if param is None: continue try: with torch.no_grad(): work_dtype = torch.float32 if param.dtype != torch.float32 else param.dtype base = param.detach() if base.device.type != "cpu" or base.dtype != work_dtype: base = base.to(device="cpu", dtype=work_dtype) else: base = base.clone() patched = _calculate_patched_weight(patch_list, base, key) if not isinstance(patched, torch.Tensor): raise RuntimeError(f"calculate_weight returned {type(patched).__name__}") if patched.shape != param.data.shape: print( f"[ParallelAnything] Warning: LoRA shape mismatch for {key}: " f"{tuple(patched.shape)} != {tuple(param.data.shape)}" ) continue param.data.copy_(patched.to(device=param.device, dtype=param.dtype)) applied += 1 except Exception as e: print(f"[ParallelAnything] Warning: failed to bake LoRA patch {key}: {e}") if applied: print(f"[ParallelAnything] Baked {applied} LoRA/weight patches into replica", flush=True) elif patches: print(f"[ParallelAnything] Warning: {len(patches)} patches found but none matched replica keys", flush=True) return applied def unwrap_parallel_blocks(module): """Replace ParallelBlock wrappers with the original lead-model blocks. Re-running setup without this nests wrappers and leaves stale peer references that point at deleted replicas. """ if module is None or isinstance(module, ParallelBlock): return for list_name in BLOCK_LIST_NAMES: block_list = getattr(module, list_name, None) if not isinstance(block_list, nn.ModuleList): continue for idx, block in enumerate(block_list): current = block while isinstance(current, ParallelBlock): current = current.local_block if current is not block: block_list[idx] = current for child in module.children(): unwrap_parallel_blocks(child) class ParallelBlock(nn.Module): """ A wrapper that handles dual-mode execution: 1. Data Parallel (Batch > 1): Runs the local block immediately. 2. Pipeline Parallel (Batch = 1): Moves data to assigned device/replica. """ def __init__(self, local_block, block_idx, owner_device, peers, is_last_block, lead_device): super().__init__() self.local_block = local_block self.block_idx = block_idx self.owner_device = as_torch_device(owner_device) self.is_last_block = is_last_block self.lead_device = as_torch_device(lead_device) # Peer replica modules must NOT be registered as children. # ComfyUI later does model.to(lead_device); if owner_block is a # submodule, that silently moves cuda:1 weights onto cuda:0 and # produces wrapper_CUDA_addmm device mismatches. owner_key = str(self.owner_device) if owner_key in peers: owner_block = peers[owner_key] elif self.owner_device == self.lead_device: owner_block = local_block else: raise RuntimeError( f"[ParallelAnything] ParallelBlock {block_idx}: no peer for {owner_key}; " f"have {list(peers)}" ) object.__setattr__(self, "peers", peers) object.__setattr__(self, "_owner_block_holder", [owner_block]) @property def owner_block(self): return self._owner_block_holder[0] def _move_tensor(self, x, device): return move_tree_to_device(x, device, non_blocking=True) def _move_args(self, args, device): return tuple(self._move_tensor(a, device) for a in args) def _move_kwargs(self, kwargs, device): return {k: self._move_tensor(v, device) for k, v in kwargs.items()} def forward(self, *args, **kwargs): if not get_pipeline_mode(): # Standard execution for Data Parallel (Batch > 1) return self.local_block(*args, **kwargs) target_device = self.owner_device # --- CRITICAL FIX --- # Do NOT check args[0].device to skip moving. # Main input might be on target_device (from prev block), # but auxiliary inputs (timesteps, context) might still be on Lead Device. # We must ensure ALL inputs are on target_device. args_moved = self._move_args(args, target_device) kwargs_moved = self._move_kwargs(kwargs, target_device) # Execute on target replica out = self.owner_block(*args_moved, **kwargs_moved) if self.is_last_block: # Return data to Lead Device for final layers/output processing out = self._move_tensor(out, self.lead_device) return out # ---------------------------------------------------------------------------- # Utilities # ---------------------------------------------------------------------------- def is_float8_dtype(dtype): if dtype is None: return False dtype_name = str(dtype) fp8_types = ['float8_e4m3fn', 'float8_e5m2', 'float8_e4m3fnuz', 'float8_e5m2fnuz', 'float8'] return any(t in dtype_name for t in fp8_types) def check_sm80_support(device_name): if not device_name.startswith("cuda"): return True try: idx = int(device_name.split(":")[-1]) capability = torch.cuda.get_device_capability(idx) major, _ = capability return major >= 8 except Exception as e: print(f"[ParallelAnything] Warning: Could not check compute capability for {device_name}: {e}") return False def device_supports_float8(device): if isinstance(device, str): if not device.startswith("cuda"): return False device = torch.device(device) if device.type != 'cuda': return False try: cap = torch.cuda.get_device_capability(device) major, minor = cap return (major, minor) >= (9, 0) except: return False def disable_flash_xformers(model): disable_configs = [ ('set_use_memory_efficient_attention_xformers', False), ('set_use_flash_attention_2', False), ('disable_xformers_memory_efficient_attention', None), ('use_xformers', False), ('use_flash_attention', False), ('use_flash_attention_2', False), ('_use_memory_efficient_attention', False), ('_flash_attention_enabled', False), ] for attr_or_method, value in disable_configs: if hasattr(model, attr_or_method): try: method = getattr(model, attr_or_method) if callable(method) and value is None: method() elif callable(method): method(value) else: setattr(model, attr_or_method, value) except Exception: pass for name, module in model.named_modules(): if any(x in name.lower() for x in ['attn', 'attention', 'transformer']): for attr in ['use_xformers', 'use_flash_attention', 'use_flash_attention_2', '_use_memory_efficient_attention', 'enable_flash', 'enable_xformers']: if hasattr(module, attr): try: setattr(module, attr, False) except AttributeError: pass if hasattr(module, 'set_processor'): try: from diffusers.models.attention_processor import Attention module.set_processor(Attention()) except Exception: pass def clear_flux_caches(model): cache_attrs = [ 'img_ids', 'txt_ids', '_img_ids', '_txt_ids', 'cached_img_ids', 'cached_txt_ids', 'pos_emb', '_pos_emb', 'pos_embed', '_pos_embed', 'cached_pos_emb', 'rope', '_rope', 'freqs_cis', '_freqs_cis', 'freqs', '_freqs', 'cache', '_cache', 'kv_cache', '_kv_cache', 'attn_bias', '_attn_bias', 'rope_cache', '_rope_cache', 'freqs_cis_cache', '_freqs_cis_cache', 'temporal_ids', 'frame_ids', 'video_ids', 'temp_pos_emb', ] cleared_count = 0 def clear_attrs(obj, name_prefix=""): nonlocal cleared_count for attr in cache_attrs: if hasattr(obj, attr): try: val = getattr(obj, attr) if val is not None: setattr(obj, attr, None) cleared_count += 1 except (AttributeError, TypeError): pass clear_attrs(model) for name, module in model.named_modules(): clear_attrs(module, name) if cleared_count > 0: print(f"[ParallelAnything] Cleared {cleared_count} cached tensors") def aggressive_cleanup(): gc.collect() if torch.cuda.is_available(): torch.cuda.synchronize() torch.cuda.empty_cache() for i in range(torch.cuda.device_count()): try: with torch.cuda.device(f'cuda:{i}'): torch.cuda.synchronize() torch.cuda.empty_cache() except: pass comfy.model_management.soft_empty_cache() def cleanup_parallel_model(model_ref): model = model_ref() if isinstance(model_ref, weakref.ref) else model_ref if model is None: return if not getattr(model, '_true_parallel_active', False): return print("[ParallelAnything] Cleaning up parallel model...") should_purge_cache = getattr(model, '_parallel_purge_cache', True) should_purge_models = getattr(model, '_parallel_purge_models', False) saved_patches = getattr(model, '_parallel_saved_patches', None) if saved_patches: try: patcher = saved_patches[0] patches = saved_patches[1] backup = saved_patches[2] if len(saved_patches) > 2 else None if backup is not None and hasattr(patcher, "backup"): patcher.backup = backup if hasattr(patcher, "unpatch_model"): try: patcher.unpatch_model(unpatch_weights=True) except Exception: pass restore_patcher_weight_patches(patcher, patches) except Exception as e: print(f"[ParallelAnything] Warning: Could not restore LoRA patches: {e}") try: delattr(model, '_parallel_saved_patches') except Exception: pass # Restore original forward if hasattr(model, '_original_forward'): try: model.forward = model._original_forward delattr(model, '_original_forward') except Exception: pass # Drop pipeline wrappers before replicas go away so a later setup # does not nest ParallelBlocks around stale peer references. try: unwrap_parallel_blocks(model) except Exception as e: print(f"[ParallelAnything] Warning: Could not unwrap parallel blocks: {e}") # Cleanup replicas if hasattr(model, '_parallel_replicas'): replicas = model._parallel_replicas for dev_name, replica in list(replicas.items()): try: unwrap_parallel_blocks(replica) if replica is not model and hasattr(replica, 'cpu'): replica.cpu() clear_flux_caches(replica) if hasattr(replica, 'gradient_checkpointing'): replica.gradient_checkpointing = False if hasattr(replica, '_gradient_checkpointing_func'): replica._gradient_checkpointing_func = None except Exception as e: print(f"[ParallelAnything] Warning: Error cleaning up replica on {dev_name}: {e}") try: delattr(model, '_parallel_replicas') except Exception: pass # Clear other attributes for attr in ['_true_parallel_active', '_parallel_devices', '_parallel_streams', '_parallel_weights', '_auto_vram_balance', '_parallel_purge_cache', '_parallel_purge_models']: if hasattr(model, attr): try: delattr(model, attr) except Exception: pass if should_purge_models: try: print("[ParallelAnything] Purging models from VRAM...") comfy.model_management.unload_all_models() except Exception as e: print(f"[ParallelAnything] Warning: Could not unload models: {e}") gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() if should_purge_cache: try: comfy.model_management.soft_empty_cache() except: pass if torch.cuda.is_available(): for i in range(torch.cuda.device_count()): try: with torch.cuda.device(f'cuda:{i}'): torch.cuda.empty_cache() except: pass def extract_model_config(model): config = {} possible_attrs = [ 'in_channels', 'out_channels', 'vec_in_dim', 'context_in_dim', 'hidden_size', 'mlp_ratio', 'num_heads', 'depth', 'depth_single_blocks', 'depth_single', 'axes_dim', 'theta', 'patch_size', 'qkv_bias', 'guidance_embed', 'txt_ids_dim', 'img_ids_dim', 'num_res_blocks', 'attention_resolutions', 'dropout', 'channel_mult', 'num_classes', 'use_checkpoint', 'num_heads_upsample', 'use_scale_shift_norm', 'resblock_updown', 'use_new_attention_order', 'adm_in_channels', 'num_noises', 'context_dim', 'n_heads', 'd_head', 'transformer_depth', 'model_channels', 'max_depth', 'num_frames', 'temporal_compression', 'temporal_dim', 'video_length', 'operations', 'ops', ] for attr in possible_attrs: if hasattr(model, attr): try: val = getattr(model, attr) if not callable(val) and not isinstance(val, (nn.Module, torch.nn.Parameter, torch.Tensor)): config[attr] = val except Exception: pass if hasattr(model, 'params'): try: params = model.params if isinstance(params, dict): config.update(params) elif is_dataclass(params): for field_info in fields(params): try: val = getattr(params, field_info.name) if not isinstance(val, (torch.Tensor, nn.Module)): config[field_info.name] = val except Exception: pass else: params_dict = vars(params) if params_dict: config.update({k: v for k, v in params_dict.items() if not isinstance(v, (torch.Tensor, nn.Module))}) except Exception: pass if hasattr(model, 'config'): try: cfg = model.config if isinstance(cfg, dict): config.update({k: v for k, v in cfg.items() if not isinstance(v, (torch.Tensor, nn.Module))}) else: cfg_dict = {k: v for k, v in vars(cfg).items() if not k.startswith('_') and not callable(v) and not isinstance(v, (torch.Tensor, nn.Module))} config.update(cfg_dict) except Exception: pass if hasattr(model, 'unet_config') and isinstance(model.unet_config, dict): config.update({k: v for k, v in model.unet_config.items() if not isinstance(v, (torch.Tensor, nn.Module))}) # ComfyUI models (Lumina2/Flux) need the ops factory. Dropping it makes # model_class(**config) do None.Linear and fall back to the slow recursive clone. for attr in ("operations", "ops"): if hasattr(model, attr): config[attr] = getattr(model, attr) clean_config = {} keep_refs = {"operations", "ops"} for k, v in config.items(): if k in keep_refs: clean_config[k] = v continue if v is not None and not isinstance(v, (torch.Tensor, nn.Module)): try: copy.deepcopy(v) clean_config[k] = v except Exception: pass return clean_config def clone_dataclass_or_object(obj, target_device=None): if is_dataclass(obj): try: field_values = {} for field_info in fields(obj): try: val = getattr(obj, field_info.name) if isinstance(val, torch.Tensor): field_values[field_info.name] = val.clone().detach() elif is_dataclass(val): field_values[field_info.name] = clone_dataclass_or_object(val) elif isinstance(val, (list, tuple)): new_val = [] for v in val: if is_dataclass(v): new_val.append(clone_dataclass_or_object(v)) elif isinstance(v, torch.Tensor): new_val.append(v.clone()) else: new_val.append(copy.deepcopy(v)) field_values[field_info.name] = type(val)(new_val) else: field_values[field_info.name] = copy.deepcopy(val) except Exception: # Fallback to direct reference if deep copy fails field_values[field_info.name] = getattr(obj, field_info.name) return obj.__class__(**field_values) except Exception: pass try: return copy.deepcopy(obj) except Exception: return obj def safe_getattr(obj, attr, default=None): return getattr(obj, attr, default) def clone_module_simple(module, target_device): if module is None: return None module_class = module.__class__ if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)): try: weight = safe_getattr(module, 'weight', None) weight_dtype = weight.dtype if weight is not None else torch.float32 has_bias = safe_getattr(module, 'bias', None) is not None supports_fp8 = device_supports_float8(target_device) if not supports_fp8 and is_float8_dtype(weight_dtype): weight_dtype = torch.float16 if isinstance(module, nn.Linear): new_mod = nn.Linear( safe_getattr(module, 'in_features'), safe_getattr(module, 'out_features'), bias=has_bias, device=target_device, dtype=weight_dtype ) elif isinstance(module, nn.Conv2d): new_mod = nn.Conv2d( safe_getattr(module, 'in_channels'), safe_getattr(module, 'out_channels'), safe_getattr(module, 'kernel_size', 1), stride=safe_getattr(module, 'stride', 1), padding=safe_getattr(module, 'padding', 0), dilation=safe_getattr(module, 'dilation', 1), groups=safe_getattr(module, 'groups', 1), bias=has_bias, padding_mode=safe_getattr(module, 'padding_mode', 'zeros'), device=target_device, dtype=weight_dtype ) elif isinstance(module, nn.Conv1d): new_mod = nn.Conv1d( safe_getattr(module, 'in_channels'), safe_getattr(module, 'out_channels'), safe_getattr(module, 'kernel_size', 1), stride=safe_getattr(module, 'stride', 1), padding=safe_getattr(module, 'padding', 0), dilation=safe_getattr(module, 'dilation', 1), groups=safe_getattr(module, 'groups', 1), bias=has_bias, padding_mode=safe_getattr(module, 'padding_mode', 'zeros'), device=target_device, dtype=weight_dtype ) with torch.no_grad(): if weight is not None: if not supports_fp8 and is_float8_dtype(weight.dtype): weight = weight.half() new_mod.weight.copy_(weight) if has_bias and safe_getattr(module, 'bias') is not None: new_mod.bias.copy_(module.bias) return new_mod except Exception as e: print(f"[ParallelAnything] Warning: Failed to reconstruct {module_class}: {e}") if isinstance(module, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm)): try: weight = safe_getattr(module, 'weight', None) weight_dtype = weight.dtype if weight is not None else torch.float32 supports_fp8 = device_supports_float8(target_device) if not supports_fp8 and is_float8_dtype(weight_dtype): weight_dtype = torch.float16 if isinstance(module, nn.LayerNorm): new_mod = nn.LayerNorm( safe_getattr(module, 'normalized_shape'), eps=safe_getattr(module, 'eps', 1e-5), elementwise_affine=safe_getattr(module, 'elementwise_affine', True), device=target_device, dtype=weight_dtype ) elif isinstance(module, nn.BatchNorm2d): new_mod = nn.BatchNorm2d( safe_getattr(module, 'num_features'), eps=safe_getattr(module, 'eps', 1e-5), momentum=safe_getattr(module, 'momentum', 0.1), affine=safe_getattr(module, 'affine', True), track_running_stats=safe_getattr(module, 'track_running_stats', True), device=target_device, dtype=weight_dtype ) elif isinstance(module, nn.GroupNorm): new_mod = nn.GroupNorm( safe_getattr(module, 'num_groups'), safe_getattr(module, 'num_channels'), eps=safe_getattr(module, 'eps', 1e-5), affine=safe_getattr(module, 'affine', True), device=target_device, dtype=weight_dtype ) with torch.no_grad(): if weight is not None: if not supports_fp8 and is_float8_dtype(weight.dtype): weight = weight.half() new_mod.weight.copy_(weight) if safe_getattr(module, 'bias') is not None: new_mod.bias.copy_(module.bias) if isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): if safe_getattr(module, 'running_mean') is not None: new_mod.running_mean.copy_(module.running_mean) if safe_getattr(module, 'running_var') is not None: new_mod.running_var.copy_(module.running_var) return new_mod except Exception as e: print(f"[ParallelAnything] Warning: Failed to reconstruct norm layer {module_class}: {e}") try: new_mod = module_class.__new__(module_class) nn.Module.__init__(new_mod) supports_fp8 = device_supports_float8(target_device) # Copy parameters for name, param in module.named_parameters(recurse=False): if param is not None: new_param_data = param.clone().detach() if not supports_fp8 and is_float8_dtype(new_param_data.dtype): new_param_data = new_param_data.half() new_param = nn.Parameter(new_param_data.to(device=target_device), requires_grad=False) new_mod.register_parameter(name, new_param) else: new_mod.register_parameter(name, None) # Copy buffers for name, buffer in module.named_buffers(recurse=False): if buffer is not None: new_buffer_data = buffer.clone().detach() if not supports_fp8 and is_float8_dtype(new_buffer_data.dtype): new_buffer_data = new_buffer_data.half() new_buffer = new_buffer_data.to(device=target_device) new_mod.register_buffer(name, new_buffer) else: new_mod.register_buffer(name, None) # Recursively clone children for name, child in module.named_children(): if child is not None: cloned_child = clone_module_simple(child, target_device) new_mod.add_module(name, cloned_child) # Copy other attributes excluded_attrs = { '_parameters', '_buffers', '_modules', '_non_persistent_buffers_set', '_backward_pre_hooks', '_backward_hooks', '_forward_pre_hooks', '_forward_hooks', '_state_dict_hooks', '_load_state_dict_pre_hooks', '_extra_state', '_modules_to_load' } cache_attrs = { 'img_ids', 'txt_ids', '_img_ids', '_txt_ids', 'cached_img_ids', 'cached_txt_ids', 'pos_emb', '_pos_emb', 'pos_embed', '_pos_embed', 'freqs_cis', '_freqs_cis', 'freqs', '_freqs', 'cache', '_cache', 'kv_cache', '_kv_cache', 'attn_bias', '_attn_bias' } for key, value in module.__dict__.items(): if key in excluded_attrs: continue try: if is_dataclass(value): setattr(new_mod, key, clone_dataclass_or_object(value)) elif isinstance(value, torch.Tensor): if key in cache_attrs: setattr(new_mod, key, None) else: tensor_val = value.clone().detach() if not supports_fp8 and is_float8_dtype(tensor_val.dtype): tensor_val = tensor_val.half() setattr(new_mod, key, tensor_val.to(target_device)) elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], torch.Tensor): if key in cache_attrs: setattr(new_mod, key, None) else: converted = [] for t in value: t_conv = t.clone().detach() if not supports_fp8 and is_float8_dtype(t_conv.dtype): t_conv = t_conv.half() converted.append(t_conv.to(target_device)) setattr(new_mod, key, type(value)(converted)) else: setattr(new_mod, key, copy.deepcopy(value)) except Exception: pass return new_mod except Exception as e: raise RuntimeError(f"Failed to clone module {module_class}: {e}") def safe_model_clone(source_model, target_device, disable_flash=False): clear_flux_caches(source_model) try: src_device = next(source_model.parameters()).device except StopIteration: src_device = torch.device('cpu') if str(src_device) == str(target_device): print(f"[ParallelAnything] Model already on {target_device}, using reference...") clear_flux_caches(source_model) return source_model # Move to CPU first if on GPU if src_device.type == 'cuda': model_cpu = source_model.cpu() if torch.cuda.is_available(): torch.cuda.synchronize() torch.cuda.empty_cache() comfy.model_management.soft_empty_cache() else: model_cpu = source_model model_class = model_cpu.__class__ replica = None method_used = "unknown" try: print(f"[ParallelAnything] Attempting efficient clone to {target_device}...") config = extract_model_config(model_cpu) state_dict = model_cpu.state_dict() if not config: raise RuntimeError("Could not extract model config") try: replica = model_class(**config) replica = replica.to(target_device, non_blocking=False) except TypeError as e: if "missing" in str(e) and "required" in str(e): try: replica = model_class(config) replica = replica.to(target_device, non_blocking=False) except Exception: config_obj = SimpleNamespace(**config) replica = model_class(config_obj) replica = replica.to(target_device, non_blocking=False) else: raise print(f"[ParallelAnything] Loading parameters incrementally...") supports_fp8 = device_supports_float8(target_device) with torch.no_grad(): for key in list(state_dict.keys()): try: param = state_dict[key] if '.' in key: parts = key.split('.') module = replica for part in parts[:-1]: module = getattr(module, part) target_param = getattr(module, parts[-1]) else: target_param = getattr(replica, key) if isinstance(target_param, (torch.nn.Parameter, torch.Tensor)): param_data = param.to(target_device, non_blocking=False) if is_float8_dtype(param_data.dtype) and not supports_fp8: param_data = param_data.half() target_param.data.copy_(param_data) del state_dict[key] del param if len(state_dict) % 50 == 0 and torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e: print(f"[ParallelAnything] Warning: Could not load {key}: {e}") del state_dict if src_device.type != 'cuda': del model_cpu gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() method_used = "incremental_gpu" except Exception as recon_error: print(f"[ParallelAnything] Config reconstruction failed: {recon_error}") print(f"[ParallelAnything] Attempting manual recursive cloning...") replica = clone_module_simple(model_cpu, target_device) method_used = "manual_recursive_clone" if src_device.type != 'cuda': del model_cpu gc.collect() if replica is None: raise RuntimeError("All clone methods failed") clear_flux_caches(replica) if not device_supports_float8(target_device): print(f"[ParallelAnything] Converting FP8 weights to FP16 for {target_device}...") def convert_fp8_recursive(m): for name, param in m.named_parameters(): if is_float8_dtype(param.dtype): param.data = param.data.half() for name, buffer in m.named_buffers(): if is_float8_dtype(buffer.dtype): buffer.data = buffer.data.half() for child in m.children(): convert_fp8_recursive(child) convert_fp8_recursive(replica) if hasattr(replica, 'gradient_checkpointing'): replica.gradient_checkpointing = False if hasattr(replica, '_gradient_checkpointing_func'): replica._gradient_checkpointing_func = None if hasattr(replica, 'hooks'): replica.hooks = [] if hasattr(replica, '_hf_hook'): replica._hf_hook = None replica.eval() for param in replica.parameters(): param.requires_grad = False for buffer in replica.buffers(): if buffer.device != target_device: buffer.data = buffer.data.to(target_device) if disable_flash: disable_flash_xformers(replica) print(f"[ParallelAnything] Cloned via {method_used} to {target_device}") return replica def get_free_vram(device_name): try: if device_name.startswith("cuda"): idx = int(device_name.split(":")[-1]) torch.cuda.set_device(idx) free_memory = ( torch.cuda.get_device_properties(idx).total_memory - torch.cuda.memory_allocated(idx) ) return free_memory / (1024 ** 2) except Exception: pass return 0 def auto_split_batch(batch_size, devices, weights): if not any(d.startswith("cuda") for d in devices): return [max(1, int(batch_size * w)) for w in weights] vram_avail = [] for d in devices: if d.startswith("cuda"): vram_avail.append(get_free_vram(d)) else: vram_avail.append(0) total_vram = sum(vram_avail) if total_vram == 0: return [max(1, int(batch_size * w)) for w in weights] adjusted_weights = [] for i, (w, vram) in enumerate(zip(weights, vram_avail)): if vram > 0: vram_weight = vram / total_vram adjusted = 0.7 * w + 0.3 * vram_weight else: adjusted = w adjusted_weights.append(adjusted) total = sum(adjusted_weights) adjusted_weights = [w/total for w in adjusted_weights] split_sizes = [max(1, int(batch_size * w)) for w in adjusted_weights] split_sizes[-1] = batch_size - sum(split_sizes[:-1]) return split_sizes class ParallelDevice: @classmethod def get_available_devices(cls): 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}") try: import torch_directml for i in range(torch_directml.device_count()): devices.append(f"privateuseone:{i}") except ImportError: pass return devices @classmethod def INPUT_TYPES(s): available = s.get_available_devices() 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 (or layers for batch=1) 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 = [] config = { "device": device_id, "percentage": float(percentage), "weight": float(percentage) / 100.0 } new_chain = previous_devices.copy() new_chain.append(config) return (new_chain,) class ParallelDeviceList: @classmethod def get_available_devices(cls): 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: @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" }), "auto_vram_balance": ("BOOLEAN", { "default": True, "tooltip": "Automatically adjust batch split based on available VRAM" }), "purge_cache": ("BOOLEAN", { "default": True, "tooltip": "Purge CUDA cache when cleaning up parallel resources" }), "purge_models": ("BOOLEAN", { "default": False, "tooltip": "Unload all models from VRAM when cleaning up (aggressive memory clearing)" }), } } RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) FUNCTION = "setup_parallel" CATEGORY = "utils/hardware" def setup_parallel(self, model, device_chain, workload_split=True, auto_vram_balance=False, purge_cache=True, purge_models=False): if model is None or not device_chain: return (model,) # Extract the diffusion model from ComfyUI model wrapper if hasattr(model, 'model') and hasattr(model.model, 'diffusion_model'): target_model = model.model.diffusion_model model_wrapper = model elif hasattr(model, 'diffusion_model'): target_model = model.diffusion_model model_wrapper = None else: target_model = model model_wrapper = None # CRITICAL FIX: Check if model is stranded on CPU BEFORE determining original_device try: current_device = next(target_model.parameters()).device if current_device.type == 'cpu': expected_device = None if hasattr(model, 'load_device'): expected_device = model.load_device elif hasattr(model, 'model') and hasattr(model.model, 'load_device'): expected_device = model.model.load_device if expected_device is not None and expected_device.type == 'cuda': print(f"[ParallelAnything] Restoring stranded model from CPU to {expected_device}") target_model = target_model.to(str(expected_device)) if torch.cuda.is_available(): torch.cuda.synchronize() else: mgmt_device = comfy.model_management.get_torch_device() if mgmt_device.type == 'cuda': print(f"[ParallelAnything] Restoring stranded model from CPU to {mgmt_device}") target_model = target_model.to(mgmt_device) if torch.cuda.is_available(): torch.cuda.synchronize() except StopIteration: pass except Exception as e: print(f"[ParallelAnything] Warning during CPU restoration: {e}") try: original_device = next(target_model.parameters()).device original_device_str = str(original_device) except StopIteration: original_device = torch.device('cpu') original_device_str = "cpu" # Original LoRA path: patch_model, then clone. Incremental clone only # copies state_dict, so we also sync weight_function hooks afterward. patch_source = model_wrapper if model_wrapper is not None else model weight_patches = snapshot_weight_patches(collect_weight_patches(patch_source)) has_lora_patches = bool(weight_patches) lora_weights_baked = False if has_lora_patches: print(f"[ParallelAnything] Found {len(weight_patches)} LoRA/weight patches") try: print("[ParallelAnything] Applying LoRA/Custom patches before cloning...") if hasattr(model_wrapper, 'patch_model'): model_wrapper.patch_model( device_to=comfy.model_management.get_torch_device(), force_patch_weights=True, ) print("[ParallelAnything] Patches applied successfully") elif hasattr(model_wrapper, 'model_patcher') and hasattr(model_wrapper.model_patcher, 'patch_model'): model_wrapper.model_patcher.patch_model( device_to=comfy.model_management.get_torch_device(), force_patch_weights=True, ) print("[ParallelAnything] Patches applied via model_patcher") patcher = find_model_patcher(patch_source) lora_weights_baked = bool(getattr(patcher, "backup", None)) except Exception as e: print(f"[ParallelAnything] Warning: Could not apply patches: {e}") # Cleanup previous parallel setup if exists if hasattr(target_model, "_true_parallel_active"): print("[ParallelAnything] Cleaning up previous parallel setup...") cleanup_parallel_model(target_model) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() comfy.model_management.soft_empty_cache() # Do not unload_all_models() here: that unbakes LoRA on current ComfyUI. # Original principle is patch_model then clone the baked weights. aggressive_cleanup() 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]))}") devices_needing_safe_attention = {} for dev_name in device_names: if not check_sm80_support(dev_name): devices_needing_safe_attention[dev_name] = True print(f"[ParallelAnything] {dev_name} < SM_80, disabling Flash/xFormers") for dev in device_names: try: torch.device(dev) except Exception: print(f"[ParallelAnything] Invalid device: {dev}") return (model,) try: original_device = next(target_model.parameters()).device original_device_str = str(original_device) except StopIteration: original_device = torch.device('cpu') original_device_str = "cpu" replicas = {} streams = {} successful_devices = [] successful_weights = [] try: print(f"[ParallelAnything] Cloning to {len(device_names)} devices...") clear_flux_caches(target_model) needs_cpu_transition = original_device.type == 'cuda' and original_device_str not in device_names if needs_cpu_transition: print(f"[ParallelAnything] Moving model to CPU for safe cloning...") target_model = target_model.cpu() aggressive_cleanup() for i, dev_name in enumerate(device_names): dev = torch.device(dev_name) need_safe = dev_name in devices_needing_safe_attention is_original_device = (dev_name == original_device_str and not needs_cpu_transition) # Original rule: with LoRA, clone every device including home # so each replica has independent baked weights. if is_original_device and not has_lora_patches: print(f"[ParallelAnything] Using original model for {dev_name} (skipping clone)") replicas[dev_name] = target_model if dev.type == 'cuda': streams[dev_name] = torch.cuda.Stream(dev) successful_devices.append(dev_name) successful_weights.append(weights[i]) continue try: if is_original_device and has_lora_patches: print(f"[ParallelAnything] Cloning to {dev_name} (preserving LoRA weights)...") else: print(f"[ParallelAnything] Cloning to {dev_name} ({i+1}/{len(device_names)})...") if dev.type == 'cuda': idx = dev.index if dev.index is not None else int(dev_name.split(":")[-1]) torch.cuda.synchronize(idx) free_mem = torch.cuda.mem_get_info(idx)[0] / 1024**3 print(f"[ParallelAnything] VRAM available on {dev_name}: {free_mem:.2f}GB") replica = safe_model_clone(target_model, dev, disable_flash=need_safe) clear_flux_caches(replica) replicas[dev_name] = replica if dev.type == 'cuda': streams[dev_name] = torch.cuda.Stream(dev) status_msg = f"[ParallelAnything] ✓ {dev_name}" if need_safe: status_msg += " (Safe mode)" if has_lora_patches: status_msg += " (LoRA active)" print(status_msg) successful_devices.append(dev_name) successful_weights.append(weights[i]) aggressive_cleanup() except RuntimeError as e: if "out of memory" in str(e).lower(): print(f"[ParallelAnything] OOM cloning to {dev_name}, skipping...") aggressive_cleanup() continue else: raise if len(successful_devices) < len(device_names): if len(successful_devices) == 0: raise RuntimeError("No devices available for parallel processing - all OOM") print(f"[ParallelAnything] Reducing to {len(successful_devices)} devices due to OOM") device_names = successful_devices total_weight = sum(successful_weights) weights = [w/total_weight for w in successful_weights] if needs_cpu_transition and original_device_str not in replicas: try: print(f"[ParallelAnything] Restoring original model to {original_device}") target_model = target_model.to(original_device) except Exception as e: print(f"[ParallelAnything] Warning: Could not restore original model: {e}") except Exception as e: print(f"[ParallelAnything] Error during cloning: {e}") import traceback traceback.print_exc() for dev_name, r in replicas.items(): try: if r is not target_model: r.cpu() del r except Exception: pass aggressive_cleanup() return (model,) print("[ParallelAnything] Configuring Model/Pipeline Parallelism for Batch=1...") lead_device_name = device_names[0] lead_replica = replicas[lead_device_name] unwrap_parallel_blocks(lead_replica) # Never wrap the live ComfyUI model. Wrapping hides Linear keys from # ModelPatcher and our hooks fight dynamic-VRAM/sage (blurry "paint" # images on the 70% home split). wrap_lead = lead_replica is not target_model if not wrap_lead: print("[ParallelAnything] Skipping pipeline wrap on live home model (keeps ComfyUI LoRA/sage)") block_lists = list(BLOCK_LIST_NAMES) for list_name in block_lists: if not wrap_lead: break if hasattr(lead_replica, list_name): local_blocks = getattr(lead_replica, list_name) if not isinstance(local_blocks, nn.ModuleList): continue num_blocks = len(local_blocks) if num_blocks == 0: continue print(f"[ParallelAnything] Configuring {list_name} ({num_blocks} blocks) for pipeline execution...") device_assignments = [] current_block = 0 for i, weight in enumerate(weights): count = int(round(weight * num_blocks)) if i == len(weights) - 1: count = num_blocks - current_block target_dev = torch.device(device_names[i]) for _ in range(count): if current_block < num_blocks: device_assignments.append(target_dev) current_block += 1 for idx in range(num_blocks): assigned_dev = device_assignments[idx] peers = {} for d_name, r in replicas.items(): if hasattr(r, list_name): peers[d_name] = getattr(r, list_name)[idx] original_block = local_blocks[idx] is_last = (idx == num_blocks - 1) wrapper = ParallelBlock( local_block=original_block, block_idx=idx, owner_device=assigned_dev, peers=peers, is_last_block=is_last, lead_device=torch.device(lead_device_name) ) local_blocks[idx] = wrapper # Clones only: bake LoRA into param.data. Hooks do not run on many # cloned Linears, which is why the 30% V100 split had no LoRA while # the live 3090 (ComfyUI path) did. Never bake the live model. if has_lora_patches: seen = set() for dev_name, replica in replicas.items(): if replica is target_model or id(replica) in seen: if replica is target_model: print(f"[ParallelAnything] {dev_name} is live model; ComfyUI applies LoRA") continue seen.add(id(replica)) print(f"[ParallelAnything] Baking LoRA into clone on {dev_name}...", flush=True) baked = apply_patches_to_replica( replica, weight_patches, progress_label=f"{dev_name} LoRA" ) print( f"[ParallelAnything] LoRA on {dev_name}: " f"baked={baked['baked']} attached={baked['attached']} " f"failed={baked['failed']} skipped={baked['skipped']}", flush=True, ) if baked["baked"] + baked["attached"] == 0: raise RuntimeError( f"LoRA patches were found ({len(weight_patches)}) but none " f"applied on clone {dev_name}." ) target_model._parallel_purge_cache = purge_cache target_model._parallel_purge_models = purge_models replicas_ref = replicas devices_ref = tuple(device_names) weights_ref = tuple(weights) lead_device = torch.device(device_names[0]) streams_ref = streams auto_balance_ref = auto_vram_balance def get_batch_size(x): if isinstance(x, torch.Tensor): return x.shape[0] elif isinstance(x, (list, tuple)) and len(x) > 0: if isinstance(x[0], torch.Tensor): batch_sizes = [t.shape[0] for t in x if isinstance(t, torch.Tensor)] if batch_sizes: return batch_sizes[0] return len(x) else: return 1 def split_batch(x, split_sizes): if isinstance(x, torch.Tensor): return torch.split(x, split_sizes, dim=0) elif isinstance(x, (list, tuple)): split_lists = [] for t in x: if isinstance(t, torch.Tensor): split_lists.append(torch.split(t, split_sizes, dim=0)) else: split_lists.append([t] * len(split_sizes)) result = [] for i in range(len(split_sizes)): result.append(type(x)(sl[i] for sl in split_lists)) return result else: return [x] * len(split_sizes) def move_to_device(x, device, non_blocking=False): return move_tree_to_device(x, device, non_blocking=non_blocking) def split_kwargs(kwargs, split_sizes, total_batch): split_kwargs_list = [{} for _ in range(len(split_sizes))] for key, value in kwargs.items(): if isinstance(value, torch.Tensor) and value.shape[0] == total_batch: chunks = torch.split(value, split_sizes, dim=0) for i, chunk in enumerate(chunks): split_kwargs_list[i][key] = chunk elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], torch.Tensor): if all(isinstance(t, torch.Tensor) and t.shape[0] == total_batch for t in value): split_lists = [torch.split(t, split_sizes, dim=0) for t in value] for i in range(len(split_sizes)): split_kwargs_list[i][key] = type(value)(sl[i] for sl in split_lists) else: for i in range(len(split_sizes)): split_kwargs_list[i][key] = value return split_kwargs_list def concatenate_results(results, dim=0): if len(results) == 0: return results first = results[0] if isinstance(first, torch.Tensor): return torch.cat(results, dim=dim) elif isinstance(first, (list, tuple)): concatenated = [] for i in range(len(first)): if isinstance(first[i], torch.Tensor): to_concat = [r[i] for r in results] concatenated.append(torch.cat(to_concat, dim=dim)) else: concatenated.append(first[i]) return type(first)(concatenated) else: return results def parallel_forward(self, x, timesteps, context=None, **kwargs): try: try: batch_size = get_batch_size(x) except ValueError as e: print(f"[ParallelAnything] Error: {e}") raise # Pipeline mode for batch_size == 1 if batch_size == 1 and workload_split: set_pipeline_mode(True) try: lead_replica = replicas_ref[devices_ref[0]] if hasattr(lead_replica, '_original_forward'): return lead_replica._original_forward(x, timesteps, context=context, **kwargs) else: return lead_replica(x, timesteps, context=context, **kwargs) finally: set_pipeline_mode(False) # Single device fallback if batch too small or workload_split disabled if batch_size < len(devices_ref) or not workload_split: set_pipeline_mode(False) with torch.no_grad(): first_replica = replicas_ref[devices_ref[0]] if hasattr(first_replica, '_original_forward'): return first_replica._original_forward(x, timesteps, context=context, **kwargs) else: return first_replica(x, timesteps, context=context, **kwargs) # Data parallelism for batch_size > 1 if auto_balance_ref: split_sizes = auto_split_batch(batch_size, list(devices_ref), list(weights_ref)) else: 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], 'stream': streams_ref.get(dev_name), 'size': size, }) if len(active) == 0: raise RuntimeError("No active devices available") if len(active) == 1: set_pipeline_mode(False) with torch.no_grad(): replica = active[0]['replica'] if hasattr(replica, '_original_forward'): return replica._original_forward(x, timesteps, context=context, **kwargs) else: return replica(x, timesteps, context=context, **kwargs) x_chunks = split_batch(x, [a['size'] for a in active]) t_chunks = split_batch(timesteps, [a['size'] for a in active]) if context is not None: c_chunks = split_batch(context, [a['size'] for a in active]) else: c_chunks = [None] * len(active) kwargs_chunks = split_kwargs(kwargs, [a['size'] for a in active], batch_size) results = [None] * len(active) exceptions = [] def worker(task_idx): set_pipeline_mode(False) task = active[task_idx] dev = task['device'] replica = task['replica'] stream = task['stream'] try: if dev.type == 'cuda': torch.cuda.set_device(dev) x_in = move_to_device(x_chunks[task_idx], dev) t_in = move_to_device(t_chunks[task_idx], dev) c_in = move_to_device(c_chunks[task_idx], dev) if c_chunks[task_idx] is not None else None k_in = {} for k, v in kwargs_chunks[task_idx].items(): k_in[k] = move_to_device(v, dev, non_blocking=False) if hasattr(replica, '_original_forward'): forward_fn = replica._original_forward else: forward_fn = replica.forward if dev.type == 'cuda' and stream is not None: with torch.cuda.device(dev): with torch.cuda.stream(stream): torch.cuda.synchronize(dev) with torch.no_grad(): out = forward_fn(x_in, t_in, context=c_in, **k_in) torch.cuda.synchronize(dev) elif dev.type == 'cuda': with torch.cuda.device(dev): torch.cuda.synchronize(dev) with torch.no_grad(): out = forward_fn(x_in, t_in, context=c_in, **k_in) torch.cuda.synchronize(dev) elif dev.type == 'xpu': with torch.xpu.device(dev): torch.xpu.synchronize(dev) with torch.no_grad(): out = forward_fn(x_in, t_in, context=c_in, **k_in) torch.xpu.synchronize(dev) else: with torch.no_grad(): out = forward_fn(x_in, t_in, context=c_in, **k_in) out = move_to_device(out, lead_device, non_blocking=False) return task_idx, out except Exception as e: return task_idx, e 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): exceptions.append((active[idx]['dev_name'], result)) results[idx] = None else: results[idx] = result if exceptions: for dev_name, exc in exceptions: print(f"[ParallelAnything] Error on {dev_name}: {exc}") raise exceptions[0][1] if any(r is None for r in results): missing = [active[i]['dev_name'] for i, r in enumerate(results) if r is None] raise RuntimeError(f"Missing results from devices: {missing}") return concatenate_results(results, dim=0) except RuntimeError as e: if "out of memory" in str(e).lower(): print("[ParallelAnything] OOM in parallel forward, attempting emergency cleanup...") lead_name = devices_ref[0] lead_replica = replicas_ref[lead_name] for name, replica in list(replicas_ref.items()): if replica is lead_replica: continue try: if hasattr(replica, "cpu"): replica.cpu() except Exception: pass aggressive_cleanup() print("[ParallelAnything] Falling back to single device processing...") set_pipeline_mode(False) with torch.no_grad(): if hasattr(lead_replica, '_original_forward'): return lead_replica._original_forward(x, timesteps, context=context, **kwargs) else: return lead_replica(x, timesteps, context=context, **kwargs) else: raise 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 target_model._parallel_devices = device_names target_model._parallel_streams = streams target_model._parallel_weights = weights target_model._auto_vram_balance = auto_vram_balance weakref.finalize(model, cleanup_parallel_model, weakref.ref(target_model)) # Update ComfyUI's model device tracking to point to lead device if hasattr(model, 'load_device'): model.load_device = lead_device elif hasattr(model, 'model') and hasattr(model.model, 'load_device'): model.model.load_device = lead_device print(f"[ParallelAnything] Parallel setup complete. Devices: {device_names}") if has_lora_patches: print("[ParallelAnything] LoRA weights synchronized across all devices") 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)", }