Update print statement from 'Hello' to 'Goodbye'

2026-02-21 18:49:38 -06:00
parent b0938a4543
commit 85b20bcd0c
1 changed files with 197 additions and 121 deletions
@@ -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 = {