FearL0rd
8a3085da78
Critical Bug Fixes: Memory Leak: weakref.finalize now uses weakref.ref(target_model) instead of direct reference Race Condition: Worker threads now return exceptions instead of raising them, allowing proper aggregation and cleanup Batch Validation: get_batch_size() now validates consistent batch dimensions across tensor lists Double Device Transfer: Removed .to(device) from split_kwargs(), now only happens in worker() Attribute Handling: Cleaner separation of in_features vs in_channels in layer cloning Improvements Added: CUDA Streams: Each GPU now uses its own CUDA stream for true parallel execution Auto VRAM Balance: New auto_vram_balance parameter adjusts splits based on available VRAM (70% user preference / 30% availability) Gradient Checkpointing: Automatically disabled on replicas to save VRAM Accelerate Hooks: Clears _hf_hook and hooks attributes for compatibility with accelerate offloading Enhanced Cache Clearing: Added rope_cache and freqs_cis_cache to FLUX cleanup Thread Safety: Proper executor context management and result validation Performance Optimizations: Non-blocking transfers where safe Synchronization only on CUDA/XPU devices Batch size auto-adjustment prevents OOM on uneven GPU setups
1031 lines
45 KiB
Python
1031 lines
45 KiB
Python
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)",
|
|
}
|