Add LTX-2 Video and multi-output model support
Fixed "'list' object has no attribute 'shape'" error when using LTX-2 Video and other diffusion models that return tuples/lists instead of single tensors. Key changes: - Added tensor container helpers: get_batch_size(), split_batch(), move_to_device(), concatenate_results() - parallel_forward now handles list/tuple inputs/outputs (common in video models with temporal conditioning) - Enhanced split_kwargs() to split list-of-tensors kwargs (frame-wise embeddings, temporal IDs) - Results are now properly concatenated whether model returns Tensor, Tuple[Tensor, ...], or List[Tensor] Tested working with: FLUX, SD3.5, LTX-2 Video, HunyuanVideo VRAM: Maintains existing N× model replication strategy Breaking: None (backward compatible with single-tensor outputs)
This commit is contained in:
FearL0rd
+134
-164
@@ -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,)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user