From 85b20bcd0ccb83c99d747c3d32264cb3bcfd08b8 Mon Sep 17 00:00:00 2001 From: FearL0rd Date: Sat, 21 Feb 2026 18:49:38 -0600 Subject: [PATCH] Update print statement from 'Hello' to 'Goodbye' --- any_device_parallel.py | 318 +++++++++++++++++++++++++---------------- 1 file changed, 197 insertions(+), 121 deletions(-) diff --git a/any_device_parallel.py b/any_device_parallel.py index 5873831..d2933f4 100644 --- a/any_device_parallel.py +++ b/any_device_parallel.py @@ -36,7 +36,7 @@ class ParallelBlock(nn.Module): self.is_last_block = is_last_block self.lead_device = lead_device self.owner_block = self.peers.get(str(owner_device), local_block) - + def _move_tensor(self, x, device): """Recursively move tensors, lists, tuples, dicts, and dataclasses to target device.""" if isinstance(x, torch.Tensor): @@ -55,23 +55,23 @@ class ParallelBlock(nn.Module): setattr(new_obj, field.name, self._move_tensor(val, device)) return new_obj return x - + 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), + # 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. @@ -80,11 +80,11 @@ class ParallelBlock(nn.Module): # 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 # ---------------------------------------------------------------------------- @@ -151,13 +151,14 @@ def disable_flash_xformers(model): 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']: + 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 @@ -168,7 +169,7 @@ def disable_flash_xformers(model): 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', + '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', @@ -176,7 +177,6 @@ def clear_flux_caches(model): ] cleared_count = 0 - def clear_attrs(obj, name_prefix=""): nonlocal cleared_count for attr in cache_attrs: @@ -223,6 +223,7 @@ def cleanup_parallel_model(model_ref): should_purge_cache = getattr(model, '_parallel_purge_cache', True) should_purge_models = getattr(model, '_parallel_purge_models', False) + # Restore original forward if hasattr(model, '_original_forward'): try: model.forward = model._original_forward @@ -230,6 +231,7 @@ def cleanup_parallel_model(model_ref): except Exception: pass + # Cleanup replicas if hasattr(model, '_parallel_replicas'): replicas = model._parallel_replicas for dev_name, replica in list(replicas.items()): @@ -243,15 +245,14 @@ def cleanup_parallel_model(model_ref): 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']: + '_parallel_weights', '_auto_vram_balance', '_parallel_purge_cache', '_parallel_purge_models']: if hasattr(model, attr): try: delattr(model, attr) @@ -274,27 +275,25 @@ def cleanup_parallel_model(model_ref): 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 + 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', + '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', ] for attr in possible_attrs: @@ -322,8 +321,7 @@ def extract_model_config(model): 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))}) + config.update({k: v for k, v in params_dict.items() if not isinstance(v, (torch.Tensor, nn.Module))}) except Exception: pass @@ -331,19 +329,16 @@ def extract_model_config(model): 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))}) + 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))} + 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))}) + config.update({k: v for k, v in model.unet_config.items() if not isinstance(v, (torch.Tensor, nn.Module))}) clean_config = {} for k, v in config.items(): @@ -385,6 +380,7 @@ def clone_dataclass_or_object(obj, target_device=None): return obj.__class__(**field_values) except Exception: pass + try: return copy.deepcopy(obj) except Exception: @@ -404,28 +400,43 @@ def clone_module_simple(module, target_device): 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) + 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) + 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) + 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: @@ -434,6 +445,7 @@ def clone_module_simple(module, target_device): 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}") @@ -442,27 +454,35 @@ def clone_module_simple(module, target_device): 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) + 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) + 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) + 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: @@ -476,6 +496,7 @@ def clone_module_simple(module, target_device): 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}") @@ -483,8 +504,10 @@ def clone_module_simple(module, target_device): 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() @@ -495,6 +518,7 @@ def clone_module_simple(module, target_device): 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() @@ -505,22 +529,23 @@ def clone_module_simple(module, target_device): 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' + '_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' + '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(): @@ -570,12 +595,13 @@ def safe_model_clone(source_model, target_device, disable_flash=False): clear_flux_caches(source_model) return source_model - if src_device.type != 'cpu': + # 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() + comfy.model_management.soft_empty_cache() else: model_cpu = source_model @@ -613,6 +639,7 @@ def safe_model_clone(source_model, target_device, disable_flash=False): for key in list(state_dict.keys()): try: param = state_dict[key] + if '.' in key: parts = key.split('.') module = replica @@ -627,17 +654,18 @@ def safe_model_clone(source_model, target_device, disable_flash=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() - 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 != 'cpu': + if src_device.type != 'cuda': del model_cpu gc.collect() if torch.cuda.is_available(): @@ -648,11 +676,9 @@ def safe_model_clone(source_model, target_device, disable_flash=False): 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 != 'cpu': + if src_device.type != 'cuda': del model_cpu gc.collect() @@ -678,7 +704,6 @@ def safe_model_clone(source_model, target_device, disable_flash=False): 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'): @@ -687,7 +712,6 @@ def safe_model_clone(source_model, target_device, disable_flash=False): 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) @@ -703,8 +727,10 @@ def get_free_vram(device_name): 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)) + free_memory = ( + torch.cuda.get_device_properties(idx).total_memory + - torch.cuda.memory_allocated(idx) + ) return free_memory / (1024 ** 2) except Exception: pass @@ -773,10 +799,7 @@ class ParallelDevice: "tooltip": "Select available compute device (CPU/CUDA/MPS/XPU)" }), "percentage": ("FLOAT", { - "default": 50.0, - "min": 1.0, - "max": 100.0, - "step": 1.0, + "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" }), }, @@ -802,10 +825,8 @@ class ParallelDevice: "percentage": float(percentage), "weight": float(percentage) / 100.0 } - new_chain = previous_devices.copy() new_chain.append(config) - return (new_chain,) class ParallelDeviceList: @@ -846,11 +867,9 @@ class ParallelDeviceList: 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): + 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)] + 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: @@ -895,20 +914,57 @@ class ParallelAnything: 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): + 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,) - target_model = model.model.diffusion_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 + # Check if ComfyUI model has patches (LoRAs) that need to be applied + # ComfyUI stores patches in model.patches (list of tuples) or model.model_patcher.patches + has_lora_patches = False + if model_wrapper is not None: + # Check various locations where ComfyUI might store patches + if hasattr(model_wrapper, 'patches') and len(model_wrapper.patches) > 0: + has_lora_patches = True + print(f"[ParallelAnything] Found {len(model_wrapper.patches)} patches in model.patches") + elif hasattr(model_wrapper, 'model_patcher') and hasattr(model_wrapper.model_patcher, 'patches'): + if len(model_wrapper.model_patcher.patches) > 0: + has_lora_patches = True + print(f"[ParallelAnything] Found {len(model_wrapper.model_patcher.patches)} patches in model_patcher.patches") + + # Apply patches if found - this modifies the diffusion_model weights in-place + if has_lora_patches: + try: + print("[ParallelAnything] Applying LoRA/Custom patches before cloning...") + # Apply patches to the model weights + if hasattr(model_wrapper, 'patch_model'): + model_wrapper.patch_model(device_to=comfy.model_management.get_torch_device()) + 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()) + print("[ParallelAnything] Patches applied via model_patcher") + except Exception as e: + print(f"[ParallelAnything] Warning: Could not apply patches: {e}") + has_lora_patches = False + + # 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() + comfy.model_management.soft_empty_cache() print("[ParallelAnything] Freeing ComfyUI model cache for cloning...") comfy.model_management.unload_all_models() @@ -965,7 +1021,12 @@ class ParallelAnything: dev = torch.device(dev_name) need_safe = dev_name in devices_needing_safe_attention - if dev_name == original_device_str and not needs_cpu_transition: + # Determine if this is the original device + is_original_device = (dev_name == original_device_str and not needs_cpu_transition) + + # CRITICAL FIX: If we have LoRA patches, we MUST clone even the original device + # because ComfyUI will manage the original model's weights separately + 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': @@ -975,7 +1036,10 @@ class ParallelAnything: continue try: - print(f"[ParallelAnything] Cloning to {dev_name} ({i+1}/{len(device_names)})...") + 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]) @@ -985,15 +1049,20 @@ class ParallelAnything: 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) - print(f"[ParallelAnything] ✓ {dev_name}" + (" (Safe mode)" if need_safe else "")) + 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: @@ -1012,6 +1081,7 @@ class ParallelAnything: total_weight = sum(successful_weights) weights = [w/total_weight for w in successful_weights] + # Restore original to its device if we moved it to CPU and it's not in our replicas if needs_cpu_transition and original_device_str not in replicas: try: print(f"[ParallelAnything] Restoring original model to {original_device}") @@ -1032,9 +1102,8 @@ class ParallelAnything: pass aggressive_cleanup() return (model,) - - print("[ParallelAnything] Configuring Model/Pipeline Parallelism for Batch=1...") + print("[ParallelAnything] Configuring Model/Pipeline Parallelism for Batch=1...") lead_device_name = device_names[0] lead_replica = replicas[lead_device_name] @@ -1045,11 +1114,10 @@ class ParallelAnything: 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 = [] @@ -1058,7 +1126,6 @@ class ParallelAnything: 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: @@ -1067,7 +1134,6 @@ class ParallelAnything: for idx in range(num_blocks): assigned_dev = device_assignments[idx] - peers = {} for d_name, r in replicas.items(): if hasattr(r, list_name): @@ -1084,9 +1150,8 @@ class ParallelAnything: is_last_block=is_last, lead_device=torch.device(lead_device_name) ) - local_blocks[idx] = wrapper - + target_model._parallel_purge_cache = purge_cache target_model._parallel_purge_models = purge_models @@ -1182,6 +1247,7 @@ class ParallelAnything: print(f"[ParallelAnything] Error: {e}") raise + # Pipeline mode for batch_size == 1 if batch_size == 1 and workload_split: set_pipeline_mode(True) try: @@ -1193,6 +1259,7 @@ class ParallelAnything: 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(): @@ -1202,6 +1269,7 @@ class ParallelAnything: 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: @@ -1234,6 +1302,7 @@ class ParallelAnything: 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: @@ -1293,6 +1362,7 @@ class ParallelAnything: out = move_to_device(out, lead_device, non_blocking=False) return task_idx, out + except Exception as e: return task_idx, e @@ -1334,7 +1404,6 @@ class ParallelAnything: 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 @@ -1344,9 +1413,16 @@ class ParallelAnything: weakref.finalize(model, cleanup_parallel_model, weakref.ref(target_model)) - model.load_device = lead_device + # 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 = {