diff --git a/any_device_parallel.py b/any_device_parallel.py index dbe1cfe..423e329 100644 --- a/any_device_parallel.py +++ b/any_device_parallel.py @@ -36,6 +36,7 @@ def disable_flash_xformers(model): ('_use_memory_efficient_attention', False), ('_flash_attention_enabled', False), ] + for attr_or_method, value in disable_methods: if hasattr(model, attr_or_method): try: @@ -51,8 +52,7 @@ def disable_flash_xformers(model): for name, module in model.named_modules(): if any(x in name.lower() for x in ['attn', 'attention', 'transformer']): - for attr in ['use_xformers', 'use_flash_attention', 'use_flash_attention_2', - '_use_memory_efficient_attention', 'enable_flash', 'enable_xformers']: + for attr in ['use_xformers', 'use_flash_attention', 'use_flash_attention_2', '_use_memory_efficient_attention', 'enable_flash', 'enable_xformers']: if hasattr(module, attr): try: setattr(module, attr, False) @@ -72,6 +72,8 @@ def clear_flux_caches(model): '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', + # LTX/Video specific + 'temporal_ids', 'frame_ids', 'video_ids', 'temp_pos_emb', ] cleared = [] for name, module in model.named_modules(): @@ -100,10 +102,8 @@ def cleanup_parallel_model(model): """Cleanup function to remove parallel replicas and restore original model.""" if not hasattr(model, '_true_parallel_active'): return - print("[ParallelAnything] Cleaning up parallel model...") - # Restore original forward if hasattr(model, '_original_forward'): model.forward = model._original_forward try: @@ -111,25 +111,20 @@ def cleanup_parallel_model(model): except: pass - # Clear replicas explicitly if hasattr(model, '_parallel_replicas'): replicas = model._parallel_replicas for dev_name, replica in replicas.items(): try: - # Move to CPU first to free GPU memory replica.cpu() - # Clear attributes that might hold tensors for name, module in replica.named_modules(): clear_flux_caches(module) except: pass - # Delete the dict try: delattr(model, '_parallel_replicas') except: pass - # Remove other attributes for attr in ['_true_parallel_active', '_parallel_devices', 'forward_orig_backup']: if hasattr(model, attr): try: @@ -137,11 +132,10 @@ def cleanup_parallel_model(model): except: pass - # Force garbage collection gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() - comfy.model_management.soft_empty_cache() + comfy.model_management.soft_empty_cache() def extract_model_config(model): """Extract initialization config from model instance with FLUX-specific handling.""" @@ -155,6 +149,8 @@ def extract_model_config(model): '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', + # Video specific + 'num_frames', 'temporal_compression', 'temporal_dim', 'video_length', ] for attr in possible_attrs: @@ -183,8 +179,7 @@ def extract_model_config(model): 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))}) + config.update({k: v for k, v in params_dict.items() if not isinstance(v, (torch.Tensor, nn.Module))}) except: pass except Exception: @@ -194,19 +189,15 @@ def extract_model_config(model): try: cfg = model.config if isinstance(cfg, dict): - config.update({k: v for k, v in cfg.items() - if not isinstance(v, (torch.Tensor, nn.Module))}) + config.update({k: v for k, v in cfg.items() if not isinstance(v, (torch.Tensor, nn.Module))}) else: - cfg_dict = {k: v for k, v in vars(cfg).items() - if not k.startswith('_') and not callable(v) - and not isinstance(v, (torch.Tensor, nn.Module))} + cfg_dict = {k: v for k, v in vars(cfg).items() if not k.startswith('_') and not callable(v) and not isinstance(v, (torch.Tensor, nn.Module))} config.update(cfg_dict) except: pass if hasattr(model, 'unet_config') and isinstance(model.unet_config, dict): - config.update({k: v for k, v in model.unet_config.items() - if not isinstance(v, (torch.Tensor, nn.Module))}) + config.update({k: v for k, v in model.unet_config.items() if not isinstance(v, (torch.Tensor, nn.Module))}) clean_config = {} for k, v in config.items(): @@ -228,8 +219,7 @@ def clone_dataclass_or_object(obj, target_device=None): 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) + clone_dataclass_or_object(v) if is_dataclass(v) else (v.clone() if isinstance(v, torch.Tensor) else v) for v in val ) else: @@ -245,7 +235,6 @@ def clone_dataclass_or_object(obj, target_device=None): return obj def safe_getattr(obj, attr, default=None): - """Safely get attribute, returning default if it doesn't exist.""" return getattr(obj, attr, default) def clone_module_simple(module, target_device): @@ -254,7 +243,6 @@ def clone_module_simple(module, target_device): return None module_class = module.__class__ - # Try standard reconstruction for basic layers with SAFE attribute access if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)): try: in_channels = safe_getattr(module, 'in_features') or safe_getattr(module, 'in_channels', 0) @@ -263,10 +251,7 @@ def clone_module_simple(module, target_device): weight_dtype = safe_getattr(safe_getattr(module, 'weight'), 'dtype', torch.float32) if isinstance(module, nn.Linear): - new_mod = nn.Linear( - in_channels, out_channels, bias=has_bias, - device=target_device, dtype=weight_dtype - ) + new_mod = nn.Linear(in_channels, out_channels, bias=has_bias, device=target_device, dtype=weight_dtype) elif isinstance(module, nn.Conv2d): kernel_size = safe_getattr(module, 'kernel_size', 1) stride = safe_getattr(module, 'stride', 1) @@ -274,12 +259,9 @@ def clone_module_simple(module, target_device): 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 - ) + 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): kernel_size = safe_getattr(module, 'kernel_size', 1) stride = safe_getattr(module, 'stride', 1) @@ -287,12 +269,9 @@ def clone_module_simple(module, target_device): 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 - ) + 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: @@ -319,21 +298,14 @@ def clone_module_simple(module, target_device): weight_dtype = torch.float32 if isinstance(module, nn.LayerNorm): - new_mod = nn.LayerNorm( - normalized_shape, eps=eps, elementwise_affine=elementwise_affine, - device=target_device, dtype=weight_dtype - ) + new_mod = nn.LayerNorm(normalized_shape, eps=eps, elementwise_affine=elementwise_affine, + device=target_device, dtype=weight_dtype) elif isinstance(module, nn.BatchNorm2d): - new_mod = nn.BatchNorm2d( - num_features, eps=eps, momentum=momentum, affine=affine, - track_running_stats=track_running_stats, - device=target_device, dtype=weight_dtype - ) + 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): - new_mod = nn.GroupNorm( - num_groups, num_channels, eps=eps, affine=affine, - device=target_device, dtype=weight_dtype - ) + 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: @@ -348,23 +320,17 @@ def clone_module_simple(module, target_device): except Exception as e: print(f"[ParallelAnything] Warning: Failed to reconstruct norm layer {module_class}: {e}") - # Fallback: Create empty instance and manually populate try: new_mod = module_class.__new__(module_class) nn.Module.__init__(new_mod) - # Copy parameters 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_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) - # Copy buffers for name, buffer in module.named_buffers(recurse=False): if buffer is not None: new_buffer = buffer.clone().detach().to(device=target_device) @@ -372,23 +338,22 @@ def clone_module_simple(module, target_device): else: new_mod.register_buffer(name, None) - # Recursively clone sub-modules 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) - # FLUX FIX: Skip batch-dependent cached tensors 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', + '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 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']: + 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)) @@ -409,7 +374,6 @@ def clone_module_simple(module, target_device): setattr(new_mod, key, copy.deepcopy(value)) except Exception: pass - new_mod = new_mod.to(target_device) return new_mod except Exception as e: @@ -418,20 +382,17 @@ def clone_module_simple(module, target_device): 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 - # Force cleanup of any temporary CPU allocations 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) @@ -441,13 +402,11 @@ def safe_model_clone(source_model, target_device, disable_flash=False): 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: @@ -459,7 +418,6 @@ def safe_model_clone(source_model, target_device, disable_flash=False): 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" @@ -472,46 +430,35 @@ def safe_model_clone(source_model, target_device, disable_flash=False): except Exception as manual_error: raise RuntimeError(f"All clone methods failed. Last error: {manual_error}") - # Clean up CPU model copy explicitly del model_cpu gc.collect() - clear_flux_caches(replica) 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 class ParallelDevice: - """Individual device configuration node with auto-detected hardware.""" - @classmethod def get_available_devices(cls): - """Dynamically detect all available compute devices.""" 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() @@ -519,7 +466,6 @@ class ParallelDevice: devices.append(f"privateuseone:{i}") except ImportError: pass - return devices @classmethod @@ -533,7 +479,10 @@ class ParallelDevice: "tooltip": "Select available compute device (CPU/CUDA/MPS/XPU)" }), "percentage": ("FLOAT", { - "default": 50.0, "min": 1.0, "max": 100.0, "step": 1.0, + "default": 50.0, + "min": 1.0, + "max": 100.0, + "step": 1.0, "tooltip": "Percentage of batch to process on this device" }), }, @@ -553,20 +502,16 @@ class ParallelDevice: 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: - """Alternative: Parallel layout with dropdowns for 1-4 devices.""" - @classmethod def get_available_devices(cls): devices = ["cpu"] @@ -617,8 +562,6 @@ class ParallelDeviceList: return (chain,) class ParallelAnything: - """Main execution node with TRUE parallelism via model replication.""" - @classmethod def INPUT_TYPES(s): return { @@ -645,11 +588,9 @@ class ParallelAnything: target_model = model.model.diffusion_model - # Check if already parallelized - cleanup old replicas first if hasattr(target_model, "_true_parallel_active"): print("[ParallelAnything] Cleaning up previous parallel setup...") cleanup_parallel_model(target_model) - # Also clear any lingering references if hasattr(target_model, '_parallel_replicas'): delattr(target_model, '_parallel_replicas') gc.collect() @@ -657,7 +598,6 @@ class ParallelAnything: torch.cuda.empty_cache() comfy.model_management.soft_empty_cache() - # Normalize weights total_pct = sum(item["percentage"] for item in device_chain) if total_pct <= 0: return (model,) @@ -671,14 +611,12 @@ class ParallelAnything: print(f"[ParallelAnything] Setup: {list(zip(device_names, [f'{w*100:.1f}%' for w in weights]))}") lead_device = torch.device(device_names[0]) - # Check which devices need flash/xformers disabled 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") - # Validate devices for dev in device_names: try: torch.device(dev) @@ -686,19 +624,15 @@ class ParallelAnything: print(f"[ParallelAnything] Invalid device: {dev}") return (model,) - # Store original device to restore later if needed try: original_device = next(target_model.parameters()).device except StopIteration: original_device = torch.device('cpu') - # Create replicas (N× VRAM) replicas = {} try: print(f"[ParallelAnything] Cloning to {len(device_names)} devices... (safe mode)") clear_flux_caches(target_model) - - # Move original to CPU to free GPU memory during cloning if original_device.type == 'cuda': print(f"[ParallelAnything] Moving model to CPU for safe cloning...") target_model = target_model.cpu() @@ -708,25 +642,20 @@ class ParallelAnything: 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 print(f"[ParallelAnything] ✓ {dev_name}" + (" (Safe mode)" if need_safe else "")) - - # Aggressive cleanup between clones to prevent OOM 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 - # Optionally restore original model to its original device (keep it usable) try: if original_device.type == 'cuda': print(f"[ParallelAnything] Restoring original model to {original_device}") @@ -736,7 +665,6 @@ class ParallelAnything: except RuntimeError as e: print(f"[ParallelAnything] VRAM Error: {e}") - # Try to restore original device on error try: if original_device.type == 'cuda': target_model = target_model.to(original_device) @@ -744,35 +672,102 @@ class ParallelAnything: pass return (model,) - # Store references for closure replicas_ref = replicas devices_ref = device_names weights_ref = weights lead_ref = lead_device + # VIDEO FIX: Helper functions for tensor handling + def get_batch_size(x): + """VIDEO FIX: Handle both tensor and list inputs.""" + if isinstance(x, torch.Tensor): + return x.shape[0] + elif isinstance(x, (list, tuple)) and len(x) > 0: + # Assume list of tensors, get batch from first element + if isinstance(x[0], torch.Tensor): + return x[0].shape[0] + else: + return len(x) + else: + return 1 + + def split_batch(x, split_sizes): + """VIDEO FIX: 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): + """VIDEO FIX: 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 + def split_kwargs(kwargs, split_sizes, device): - """FLUX FIX: Split any tensor kwargs that have batch dimension.""" + """VIDEO FIX: 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] == sum(split_sizes): - # This is a batch tensor, split it + 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.to(device) + elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], torch.Tensor): + # Handle list of tensors (e.g., frame conditioning) + # Check if all have batch dim + 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: - # Not a batch tensor or doesn't match batch size, broadcast to all for i in range(len(split_sizes)): split_kwargs_list[i][key] = value return split_kwargs_list + def concatenate_results(results, dim=0): + """VIDEO FIX: 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): - batch_size = x.shape[0] - + batch_size = get_batch_size(x) if batch_size < len(devices_ref) or not workload_split: with torch.no_grad(): - return replicas_ref[devices_ref[0]](x, timesteps, context=context, **kwargs) + out = replicas_ref[devices_ref[0]](x, timesteps, context=context, **kwargs) + return out - # Calculate splits split_sizes = [max(1, int(batch_size * w)) for w in weights_ref] split_sizes[-1] = batch_size - sum(split_sizes[:-1]) @@ -791,77 +786,54 @@ class ParallelAnything: with torch.no_grad(): return active[0]['replica'](x, timesteps, context=context, **kwargs) - # Split tensors - x_chunks = torch.split(x, [a['size'] for a in active], dim=0) - t_chunks = torch.split(timesteps, [a['size'] for a in active], dim=0) - c_chunks = torch.split(context, [a['size'] for a in active], dim=0) if context is not None else [None] * len(active) + # VIDEO FIX: Use robust splitting + 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) - # FLUX FIX: Split kwargs tensors (guidance, y, etc.) kwargs_chunks = split_kwargs(kwargs, [a['size'] for a in active], lead_ref) - results = [None] * len(active) def worker(task_idx): task = active[task_idx] dev = task['device'] replica = task['replica'] - try: + 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) + if dev.type == 'cuda': with torch.cuda.device(dev): torch.cuda.synchronize(dev) - x_in = x_chunks[task_idx].to(dev, non_blocking=True) - t_in = t_chunks[task_idx].to(dev, non_blocking=True) - c_in = c_chunks[task_idx].to(dev, non_blocking=True) if c_chunks[task_idx] is not None else None - - # Get the split kwargs for this chunk - k_in = kwargs_chunks[task_idx] - # Move any remaining tensors to device - for k, v in k_in.items(): - if isinstance(v, torch.Tensor) and v.device != dev: - k_in[k] = v.to(dev) - with torch.no_grad(): out = replica(x_in, t_in, context=c_in, **k_in) torch.cuda.synchronize(dev) - return task_idx, out.to(lead_ref, non_blocking=False) - elif dev.type == 'xpu': with torch.xpu.device(dev): torch.xpu.synchronize(dev) - x_in = x_chunks[task_idx].to(dev, non_blocking=True) - t_in = t_chunks[task_idx].to(dev, non_blocking=True) - c_in = c_chunks[task_idx].to(dev, non_blocking=True) if c_chunks[task_idx] is not None else None - - k_in = kwargs_chunks[task_idx] - for k, v in k_in.items(): - if isinstance(v, torch.Tensor) and v.device != dev: - k_in[k] = v.to(dev) - with torch.no_grad(): out = replica(x_in, t_in, context=c_in, **k_in) torch.xpu.synchronize(dev) - return task_idx, out.to(lead_ref, non_blocking=False) - else: - # CPU/MPS - x_in = x_chunks[task_idx].to(dev) - t_in = t_chunks[task_idx].to(dev) - c_in = c_chunks[task_idx].to(dev) if c_chunks[task_idx] is not None else None - - k_in = kwargs_chunks[task_idx] - for k, v in k_in.items(): - if isinstance(v, torch.Tensor) and v.device != dev: - k_in[k] = v.to(dev) - with torch.no_grad(): out = replica(x_in, t_in, context=c_in, **k_in) - return task_idx, out.to(lead_ref) - + + # VIDEO FIX: Move result back handling list/tuple + out = move_to_device(out, lead_ref, non_blocking=False) + return task_idx, out except Exception as e: return task_idx, e - # TRUE PARALLEL EXECUTION 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): @@ -870,18 +842,16 @@ class ParallelAnything: raise result results[idx] = result - return torch.cat(results, dim=0) + # VIDEO FIX: Robust concatenation + return concatenate_results(results, dim=0) - # Replace forward method 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 - # Register cleanup on model deletion weakref.finalize(model, cleanup_parallel_model, target_model) - model.load_device = lead_device return (model,)