FearL0rd-ComfyUI-ParallelAnything/any_device_parallel.py

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)",
}