import torch
import torch.nn as nn
import types
import copy
import io
import pickle
import gc
import weakref
from types import SimpleNamespace
from dataclasses import dataclass, fields, is_dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed, TimeoutError
import comfy.model_management

def check_sm80_support(device_name):
    """Check if CUDA device supports SM_80 (Ampere) or higher."""
    if not device_name.startswith("cuda"):
        return True
    try:
        idx = int(device_name.split(":")[-1])
        capability = torch.cuda.get_device_capability(idx)
        major, minor = capability
        return major >= 8
    except Exception as e:
        print(f"[ParallelAnything] Warning: Could not check compute capability for {device_name}: {e}")
        return False

def disable_flash_xformers(model):
    """Aggressively disable Flash Attention and xFormers on the model."""
    disable_methods = [
        ('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_methods:
        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:
                        pass
            if hasattr(module, 'set_processor'):
                try:
                    from diffusers.models.attention_processor import Attention
                    module.set_processor(Attention())
                except:
                    pass

def clear_flux_caches(model):
    """Clear FLUX-specific cached tensors that depend on batch size or input dimensions."""
    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',
        # LTX/Video specific
        'temporal_ids', 'frame_ids', 'video_ids', 'temp_pos_emb',
    ]
    cleared = []
    
    for name, module in model.named_modules():
        for attr in cache_attrs:
            if hasattr(module, attr):
                try:
                    val = getattr(module, attr)
                    if val is not None:
                        setattr(module, attr, None)
                        cleared.append(f"{name}.{attr}" if name else attr)
                except Exception:
                    pass
    
    for attr in cache_attrs:
        if hasattr(model, attr):
            try:
                val = getattr(model, attr)
                if val is not None:
                    setattr(model, attr, None)
                    cleared.append(attr)
            except Exception:
                pass
    
    if cleared:
        print(f"[ParallelAnything] Cleared {len(cleared)} cached tensors")

def cleanup_parallel_model(model_ref):
    """Cleanup function to remove parallel replicas and restore original model."""
    # FIX: Handle weakref properly
    model = model_ref() if isinstance(model_ref, weakref.ref) else model_ref
    if model is None:
        return
    
    if not hasattr(model, '_true_parallel_active'):
        return
    
    print("[ParallelAnything] Cleaning up parallel model...")
    
    if hasattr(model, '_original_forward'):
        model.forward = model._original_forward
        try:
            delattr(model, '_original_forward')
        except:
            pass
    
    if hasattr(model, '_parallel_replicas'):
        replicas = model._parallel_replicas
        for dev_name, replica in replicas.items():
            try:
                replica.cpu()
                for name, module in replica.named_modules():
                    clear_flux_caches(module)
                # FIX: Disable gradient checkpointing cleanup
                if hasattr(replica, 'gradient_checkpointing'):
                    replica.gradient_checkpointing = False
            except:
                pass
        try:
            delattr(model, '_parallel_replicas')
        except:
            pass
    
    for attr in ['_true_parallel_active', '_parallel_devices', 'forward_orig_backup', '_parallel_streams']:
        if hasattr(model, attr):
            try:
                delattr(model, attr)
            except:
                pass
    
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    comfy.model_management.soft_empty_cache()

def extract_model_config(model):
    """Extract initialization config from model instance with FLUX-specific handling."""
    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',
    ]
    
    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:
                        pass
            else:
                try:
                    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:
                    pass
        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:
            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))})
    
    clean_config = {}
    for k, v in config.items():
        if v is not None and not isinstance(v, (torch.Tensor, nn.Module)):
            clean_config[k] = v
    return clean_config

def clone_dataclass_or_object(obj, target_device=None):
    """Deep copy a dataclass or simple object, handling nested structures."""
    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)):
                        field_values[field_info.name] = type(val)(
                            clone_dataclass_or_object(v) if is_dataclass(v) else (v.clone() if isinstance(v, torch.Tensor) else v)
                            for v in val
                        )
                    else:
                        field_values[field_info.name] = copy.deepcopy(val)
                except:
                    field_values[field_info.name] = getattr(obj, field_info.name)
            return obj.__class__(**field_values)
        except:
            pass
    try:
        return copy.deepcopy(obj)
    except:
        return obj

def safe_getattr(obj, attr, default=None):
    return getattr(obj, attr, default)

def clone_module_simple(module, target_device):
    """Simple module cloning that handles Parameters and FLUX-specific attributes."""
    if module is None:
        return None
    
    module_class = module.__class__
    
    # FIX: Cleaner attribute handling for Conv/Linear layers
    if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)):
        try:
            weight_dtype = safe_getattr(safe_getattr(module, 'weight'), 'dtype', torch.float32)
            has_bias = safe_getattr(module, 'bias') is not None
            
            if isinstance(module, nn.Linear):
                in_features = safe_getattr(module, 'in_features')
                out_features = safe_getattr(module, 'out_features')
                new_mod = nn.Linear(in_features, out_features, bias=has_bias, device=target_device, dtype=weight_dtype)
            elif isinstance(module, nn.Conv2d):
                in_channels = safe_getattr(module, 'in_channels')
                out_channels = safe_getattr(module, 'out_channels')
                kernel_size = 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)
                padding_mode = safe_getattr(module, 'padding_mode', 'zeros')
                new_mod = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, 
                                   dilation=dilation, groups=groups, bias=has_bias, padding_mode=padding_mode, 
                                   device=target_device, dtype=weight_dtype)
            elif isinstance(module, nn.Conv1d):
                in_channels = safe_getattr(module, 'in_channels')
                out_channels = safe_getattr(module, 'out_channels')
                kernel_size = 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)
                padding_mode = safe_getattr(module, 'padding_mode', 'zeros')
                new_mod = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, 
                                   dilation=dilation, groups=groups, bias=has_bias, padding_mode=padding_mode, 
                                   device=target_device, dtype=weight_dtype)
            
            with torch.no_grad():
                if hasattr(module, 'weight') and module.weight is not None:
                    new_mod.weight.copy_(module.weight)
                if has_bias and hasattr(module, 'bias') and 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_dtype = safe_getattr(safe_getattr(module, 'weight'), 'dtype', torch.float32)
            if weight_dtype is None:
                weight_dtype = torch.float32
            
            if isinstance(module, nn.LayerNorm):
                normalized_shape = safe_getattr(module, 'normalized_shape')
                eps = safe_getattr(module, 'eps', 1e-5)
                elementwise_affine = safe_getattr(module, 'elementwise_affine', True)
                new_mod = nn.LayerNorm(normalized_shape, eps=eps, elementwise_affine=elementwise_affine, 
                                      device=target_device, dtype=weight_dtype)
            elif isinstance(module, nn.BatchNorm2d):
                num_features = 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)
                new_mod = nn.BatchNorm2d(num_features, eps=eps, momentum=momentum, affine=affine, 
                                        track_running_stats=track_running_stats, device=target_device, dtype=weight_dtype)
            elif isinstance(module, nn.GroupNorm):
                num_groups = safe_getattr(module, 'num_groups')
                num_channels = safe_getattr(module, 'num_channels')
                eps = safe_getattr(module, 'eps', 1e-5)
                affine = safe_getattr(module, 'affine', True)
                new_mod = nn.GroupNorm(num_groups, num_channels, eps=eps, affine=affine, 
                                      device=target_device, dtype=weight_dtype)
            
            with torch.no_grad():
                if hasattr(module, 'weight') and module.weight is not None:
                    new_mod.weight.copy_(module.weight)
                if hasattr(module, 'bias') and module.bias is not None:
                    new_mod.bias.copy_(module.bias)
                if hasattr(module, 'running_mean') and safe_getattr(module, 'running_mean') is not None:
                    new_mod.running_mean.copy_(module.running_mean)
                if hasattr(module, 'running_var') and 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)
        
        for name, param in module.named_parameters(recurse=False):
            if param is not None:
                new_param = nn.Parameter(param.clone().detach().to(device=target_device), requires_grad=False)
                new_mod.register_parameter(name, new_param)
            else:
                new_mod.register_parameter(name, None)
        
        for name, buffer in module.named_buffers(recurse=False):
            if buffer is not None:
                new_buffer = buffer.clone().detach().to(device=target_device)
                new_mod.register_buffer(name, new_buffer)
            else:
                new_mod.register_buffer(name, None)
        
        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)
        
        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',
            'rope_cache', '_rope_cache', 'freqs_cis_cache', '_freqs_cis_cache'
        }
        
        for key, value in module.__dict__.items():
            if key not in ['_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']:
                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:
                            setattr(new_mod, key, value.clone().detach().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:
                            setattr(new_mod, key, type(value)(t.clone().detach().to(target_device) for t in value))
                    else:
                        try:
                            setattr(new_mod, key, copy.copy(value))
                        except:
                            setattr(new_mod, key, copy.deepcopy(value))
                except Exception:
                    pass
        
        new_mod = new_mod.to(target_device)
        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):
    """FLUX-safe model cloning with explicit memory cleanup."""
    clear_flux_caches(source_model)
    
    src_device = next(source_model.parameters()).device if next(source_model.parameters(), None) is not None else torch.device('cpu')
    
    if src_device.type != 'cpu':
        model_cpu = source_model.cpu()
    else:
        model_cpu = source_model
    
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        comfy.model_management.soft_empty_cache()
    
    model_class = model_cpu.__class__
    
    try:
        replica = copy.deepcopy(model_cpu)
        replica = replica.to(target_device)
        method_used = "deepcopy"
    except (pickle.PicklingError, AttributeError, TypeError, RuntimeError) as e:
        if any(x in str(e).lower() for x in ["pickle", "can't pickle", "local object", "copy_"]):
            print(f"[ParallelAnything] Deepcopy failed ({str(e)[:80]}), trying config reconstruction...")
        else:
            raise e
        
        try:
            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)
            except TypeError as e:
                if "missing" in str(e) and "required" in str(e):
                    try:
                        replica = model_class(config)
                    except:
                        config_obj = SimpleNamespace(**config)
                        replica = model_class(config_obj)
                else:
                    raise e
            
            replica.load_state_dict(state_dict, strict=False)
            replica = replica.to(target_device)
            method_used = "config_reconstruction"
        except Exception as recon_error:
            print(f"[ParallelAnything] Config reconstruction failed: {recon_error}")
            print(f"[ParallelAnything] Attempting manual recursive cloning...")
            try:
                replica = clone_module_simple(model_cpu, target_device)
                method_used = "manual_recursive_clone"
            except Exception as manual_error:
                raise RuntimeError(f"All clone methods failed. Last error: {manual_error}")
    
    del model_cpu
    gc.collect()
    clear_flux_caches(replica)
    
    # FIX: Disable gradient checkpointing to save VRAM
    if hasattr(replica, 'gradient_checkpointing'):
        replica.gradient_checkpointing = False
    if hasattr(replica, '_gradient_checkpointing_func'):
        replica._gradient_checkpointing_func = None
    
    # FIX: Handle accelerate hooks/offloading
    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):
    """Get available VRAM in MB for a device."""
    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)  # Convert to MB
    except:
        pass
    return 0

def auto_split_batch(batch_size, devices, weights):
    """Adjust split based on available VRAM."""
    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)  # CPU gets 0 weight in VRAM calculation
    
    total_vram = sum(vram_avail)
    if total_vram == 0:
        return [max(1, int(batch_size * w)) for w in weights]
    
    # Mix percentage weights with VRAM availability (70% user pref, 30% VRAM)
    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)
    
    # Normalize
    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])  # Adjust last to ensure sum = batch_size
    return split_sizes

class ParallelDevice:
    @classmethod
    def get_available_devices(cls):
        devices = ["cpu"]
        if torch.cuda.is_available():
            count = torch.cuda.device_count()
            for i in range(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():
            count = torch.xpu.device_count()
            for i in range(count):
                devices.append(f"xpu:{i}")
        try:
            import torch_directml
            dml_count = torch_directml.device_count()
            for i in range(dml_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 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": False,
                    "tooltip": "Automatically adjust batch split based on available VRAM"
                }),
            }
        }

    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):
        if model is None or not device_chain:
            return (model,)
        
        target_model = model.model.diffusion_model
        
        if hasattr(target_model, "_true_parallel_active"):
            print("[ParallelAnything] Cleaning up previous parallel setup...")
            cleanup_parallel_model(target_model)
            if hasattr(target_model, '_parallel_replicas'):
                delattr(target_model, '_parallel_replicas')
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            comfy.model_management.soft_empty_cache()
        
        total_pct = sum(item["percentage"] for item in device_chain)
        if total_pct <= 0:
            return (model,)
        
        device_names = []
        weights = []
        for item in device_chain:
            weights.append(item["percentage"] / total_pct)
            device_names.append(item["device"])
        
        print(f"[ParallelAnything] Setup: {list(zip(device_names, [f'{w*100:.1f}%' for w in weights]))}")
        
        lead_device = torch.device(device_names[0])
        
        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:
                print(f"[ParallelAnything] Invalid device: {dev}")
                return (model,)
        
        try:
            original_device = next(target_model.parameters()).device
        except StopIteration:
            original_device = torch.device('cpu')
        
        replicas = {}
        streams = {}  # IMPROVEMENT: CUDA streams for parallel execution
        
        try:
            print(f"[ParallelAnything] Cloning to {len(device_names)} devices... (safe mode)")
            clear_flux_caches(target_model)
            
            if original_device.type == 'cuda':
                print(f"[ParallelAnything] Moving model to CPU for safe cloning...")
                target_model = target_model.cpu()
                torch.cuda.empty_cache()
                comfy.model_management.soft_empty_cache()
            
            for dev_name in device_names:
                dev = torch.device(dev_name)
                need_safe = dev_name in devices_needing_safe_attention
                
                try:
                    replica = safe_model_clone(target_model, dev, disable_flash=need_safe)
                    clear_flux_caches(replica)
                    replicas[dev_name] = replica
                    
                    # IMPROVEMENT: Create CUDA stream for this device
                    if dev.type == 'cuda':
                        streams[dev_name] = torch.cuda.Stream(dev)
                    else:
                        streams[dev_name] = None
                    
                    print(f"[ParallelAnything] ✓ {dev_name}" + (" (Safe mode)" if need_safe else ""))
                    
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                        comfy.model_management.soft_empty_cache()
                        
                except Exception as e:
                    print(f"[ParallelAnything] Error cloning to {dev_name}: {e}")
                    import traceback
                    traceback.print_exc()
                    raise
            
            try:
                if original_device.type == 'cuda':
                    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 device: {e}")
                
        except RuntimeError as e:
            print(f"[ParallelAnything] VRAM Error: {e}")
            try:
                if original_device.type == 'cuda':
                    target_model = target_model.to(original_device)
            except:
                pass
            return (model,)
        
        replicas_ref = replicas
        devices_ref = device_names
        weights_ref = weights
        lead_ref = lead_device
        streams_ref = streams
        auto_balance_ref = auto_vram_balance
        
        # FIX: Improved batch size detection with validation
        def get_batch_size(x):
            """Handle both tensor and list inputs with validation."""
            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:
                        # Validate all tensors have same batch size
                        if not all(b == batch_sizes[0] for b in batch_sizes):
                            raise ValueError(f"Inconsistent batch sizes in input list: {batch_sizes}")
                        return batch_sizes[0]
                return len(x)
            else:
                return 1
        
        def split_batch(x, split_sizes):
            """Split tensors or lists of tensors."""
            if isinstance(x, torch.Tensor):
                return torch.split(x, split_sizes, dim=0)
            elif isinstance(x, (list, tuple)):
                # Split each tensor in the list
                split_lists = [torch.split(t, split_sizes, dim=0) if isinstance(t, torch.Tensor) else [t] * len(split_sizes) for t in x]
                # Transpose: list of splits -> splits of lists
                return [type(x)(t[i] for t in split_lists) for i in range(len(split_sizes))]
            else:
                # Non-tensor, broadcast to all
                return [x] * len(split_sizes)
        
        def move_to_device(x, device, non_blocking=True):
            """Move tensor or list/tuple of tensors to device."""
            if isinstance(x, torch.Tensor):
                return x.to(device, non_blocking=non_blocking)
            elif isinstance(x, (list, tuple)):
                moved = [move_to_device(t, device, non_blocking) for t in x]
                return type(x)(moved)
            else:
                return x
        
        # FIX: Don't pre-move tensors in split_kwargs (avoid double transfer)
        def split_kwargs(kwargs, split_sizes):
            """Handle list-of-tensors in kwargs (common in video models)."""
            split_kwargs_list = [{} for _ in range(len(split_sizes))]
            total_batch = sum(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  # Don't move to device yet
                elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], torch.Tensor):
                    # Handle list of tensors (e.g., frame conditioning)
                    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:
                        # Broadcast whole list
                        for i in range(len(split_sizes)):
                            split_kwargs_list[i][key] = value
                else:
                    for i in range(len(split_sizes)):
                        split_kwargs_list[i][key] = value
            return split_kwargs_list
        
        def concatenate_results(results, dim=0):
            """Concatenate tensor or list/tuple of tensors."""
            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)):
                # Concatenate each position in the tuple/list
                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:
                        # Non-tensor, just take from first (assume same across batch)
                        concatenated.append(first[i])
                return type(first)(concatenated)
            else:
                return results
        
        def parallel_forward(self, x, timesteps, context=None, **kwargs):
            # FIX: Validate batch size consistency
            try:
                batch_size = get_batch_size(x)
            except ValueError as e:
                print(f"[ParallelAnything] Error: {e}")
                raise
            
            if batch_size < len(devices_ref) or not workload_split:
                with torch.no_grad():
                    out = replicas_ref[devices_ref[0]](x, timesteps, context=context, **kwargs)
                return out
            
            # IMPROVEMENT: Use auto VRAM balancing if enabled
            if auto_balance_ref:
                split_sizes = auto_split_batch(batch_size, devices_ref, 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) == 1:
                with torch.no_grad():
                    return active[0]['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)
            
            # FIX: Don't pre-move kwargs to device
            kwargs_chunks = split_kwargs(kwargs, [a['size'] for a in active])
            
            results = [None] * len(active)
            exceptions = []
            
            # FIX: Improved error handling with proper cleanup
            def worker(task_idx):
                task = active[task_idx]
                dev = task['device']
                replica = task['replica']
                stream = task['stream']
                
                try:
                    # Move inputs to device (now happens here, not in split_kwargs)
                    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
                    
                    # Handle kwargs
                    k_in = {}
                    for k, v in kwargs_chunks[task_idx].items():
                        k_in[k] = move_to_device(v, dev, non_blocking=False)
                    
                    # IMPROVEMENT: Use CUDA streams for parallel execution
                    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 = replica(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 = replica(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 = replica(x_in, t_in, context=c_in, **k_in)
                            torch.xpu.synchronize(dev)
                    else:
                        with torch.no_grad():
                            out = replica(x_in, t_in, context=c_in, **k_in)
                    
                    # Move result back handling list/tuple
                    out = move_to_device(out, lead_ref, non_blocking=False)
                    return task_idx, out
                    
                except Exception as e:
                    # FIX: Return exception instead of raising to allow cleanup
                    return task_idx, e
            
            # FIX: Use context manager for proper thread cleanup
            with ThreadPoolExecutor(max_workers=len(active)) as executor:
                futures = [executor.submit(worker, i) for i in range(len(active))]
                
                # FIX: Collect all results and check for exceptions
                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
            
            # FIX: Handle exceptions after all threads complete
            if exceptions:
                for dev_name, exc in exceptions:
                    print(f"[ParallelAnything] Error on {dev_name}: {exc}")
                raise exceptions[0][1]  # Raise first exception
            
            # Verify all results present
            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)
        
        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  # Store for cleanup
        
        # FIX: Use weakref properly to avoid strong reference
        weakref.finalize(model, cleanup_parallel_model, weakref.ref(target_model))
        
        model.load_device = lead_device
        return (model,)

NODE_CLASS_MAPPINGS = {
    "ParallelAnything": ParallelAnything,
    "ParallelDevice": ParallelDevice,
    "ParallelDeviceList": ParallelDeviceList,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ParallelAnything": "Parallel Anything (True Multi-GPU)",
    "ParallelDevice": "Parallel Device Config",
    "ParallelDeviceList": "Parallel Device List (1-4x)",
}