FearL0rd
e75de86a40
Fix dtype compatibility for FP8-converted and native Float models Resolves "mat1 and mat2 must have the same dtype" errors for Flux (BFloat16 vs Half) and "expected scalar type Float but found Half" for video models (WAN). Changes: - Implement universal dtype matching: inputs now dynamically convert to match model parameter dtype (Float32, Float16, BFloat16) rather than assuming FP16 - Remove conditional _converted_from_fp8 flag in favor of runtime dtype detection - Add dtype synchronization in both single-device fallback and multi-device worker paths - Preserve non-floating point tensors (timesteps, masks) from unnecessary conversion - Ensure FP8 weights converted to FP16 on pre-Hopper GPUs properly align with inputs This fixes compatibility for: - Flux FP8 models on V100/3090 (Ampere) where weights become FP16 but inputs are BF16 - WAN video models and other native Float32 architectures that require full precision - Mixed precision environments where model replicas may have different dtypes per device
1379 lines
60 KiB
Python
1379 lines
60 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import types
|
|
import copy
|
|
import gc
|
|
import weakref
|
|
from types import SimpleNamespace
|
|
from dataclasses import dataclass, fields, is_dataclass
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import comfy.model_management
|
|
|
|
def is_float8_dtype(dtype):
|
|
"""Check if dtype is any FP8 variant (E4M3, E5M2, etc.)."""
|
|
if dtype is None:
|
|
return False
|
|
dtype_name = str(dtype)
|
|
fp8_types = ['float8_e4m3fn', 'float8_e5m2', 'float8_e4m3fnuz', 'float8_e5m2fnuz', 'float8']
|
|
return any(t in dtype_name for t in fp8_types)
|
|
|
|
def check_sm80_support(device_name):
|
|
"""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, _ = capability
|
|
return major >= 8
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Could not check compute capability for {device_name}: {e}")
|
|
return False
|
|
|
|
def device_supports_float8(device):
|
|
"""Check if CUDA device supports Float8 computation (SM90+ / Hopper/Ada)."""
|
|
if isinstance(device, str):
|
|
if not device.startswith("cuda"):
|
|
return False
|
|
device = torch.device(device)
|
|
if device.type != 'cuda':
|
|
return False
|
|
try:
|
|
cap = torch.cuda.get_device_capability(device)
|
|
major, minor = cap
|
|
# FP8 computation requires Hopper/Ada (SM9.0+) or newer
|
|
return (major, minor) >= (9, 0)
|
|
except:
|
|
return False
|
|
|
|
def disable_flash_xformers(model):
|
|
"""Aggressively disable Flash Attention and xFormers on the model."""
|
|
disable_configs = [
|
|
('set_use_memory_efficient_attention_xformers', False),
|
|
('set_use_flash_attention_2', False),
|
|
('disable_xformers_memory_efficient_attention', None),
|
|
('use_xformers', False),
|
|
('use_flash_attention', False),
|
|
('use_flash_attention_2', False),
|
|
('_use_memory_efficient_attention', False),
|
|
('_flash_attention_enabled', False),
|
|
]
|
|
for attr_or_method, value in disable_configs:
|
|
if hasattr(model, attr_or_method):
|
|
try:
|
|
method = getattr(model, attr_or_method)
|
|
if callable(method) and value is None:
|
|
method()
|
|
elif callable(method):
|
|
method(value)
|
|
else:
|
|
setattr(model, attr_or_method, value)
|
|
except Exception:
|
|
pass
|
|
|
|
# Recursively check modules
|
|
for name, module in model.named_modules():
|
|
if any(x in name.lower() for x in ['attn', 'attention', 'transformer']):
|
|
for attr in ['use_xformers', 'use_flash_attention', 'use_flash_attention_2', '_use_memory_efficient_attention', 'enable_flash', 'enable_xformers']:
|
|
if hasattr(module, attr):
|
|
try:
|
|
setattr(module, attr, False)
|
|
except AttributeError:
|
|
pass
|
|
# Reset attention processor if diffusers-style
|
|
if hasattr(module, 'set_processor'):
|
|
try:
|
|
from diffusers.models.attention_processor import Attention
|
|
module.set_processor(Attention())
|
|
except Exception:
|
|
pass
|
|
|
|
def clear_flux_caches(model):
|
|
"""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',
|
|
'temporal_ids', 'frame_ids', 'video_ids', 'temp_pos_emb',
|
|
]
|
|
cleared_count = 0
|
|
def clear_attrs(obj, name_prefix=""):
|
|
nonlocal cleared_count
|
|
for attr in cache_attrs:
|
|
if hasattr(obj, attr):
|
|
try:
|
|
val = getattr(obj, attr)
|
|
if val is not None:
|
|
setattr(obj, attr, None)
|
|
cleared_count += 1
|
|
except (AttributeError, TypeError):
|
|
pass
|
|
# Clear on model itself
|
|
clear_attrs(model)
|
|
# Clear on all submodules
|
|
for name, module in model.named_modules():
|
|
clear_attrs(module, name)
|
|
|
|
if cleared_count > 0:
|
|
print(f"[ParallelAnything] Cleared {cleared_count} cached tensors")
|
|
|
|
def aggressive_cleanup():
|
|
"""Force aggressive CUDA memory cleanup."""
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
# Force allocator to release cached blocks
|
|
for i in range(torch.cuda.device_count()):
|
|
try:
|
|
with torch.cuda.device(f'cuda:{i}'):
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
except:
|
|
pass
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
def cleanup_parallel_model(model_ref):
|
|
"""Cleanup function to remove parallel replicas and restore original model."""
|
|
model = model_ref() if isinstance(model_ref, weakref.ref) else model_ref
|
|
if model is None:
|
|
return
|
|
if not getattr(model, '_true_parallel_active', False):
|
|
return
|
|
|
|
print("[ParallelAnything] Cleaning up parallel model...")
|
|
|
|
# Check purge preferences
|
|
should_purge_cache = getattr(model, '_parallel_purge_cache', True)
|
|
should_purge_models = getattr(model, '_parallel_purge_models', False)
|
|
|
|
# Restore original forward
|
|
if hasattr(model, '_original_forward'):
|
|
try:
|
|
model.forward = model._original_forward
|
|
delattr(model, '_original_forward')
|
|
except Exception:
|
|
pass
|
|
|
|
# Cleanup replicas
|
|
if hasattr(model, '_parallel_replicas'):
|
|
replicas = model._parallel_replicas
|
|
for dev_name, replica in list(replicas.items()):
|
|
try:
|
|
# Move to CPU to free GPU memory
|
|
if hasattr(replica, 'cpu'):
|
|
replica.cpu()
|
|
# Clear caches
|
|
clear_flux_caches(replica)
|
|
# Disable gradient checkpointing
|
|
if hasattr(replica, 'gradient_checkpointing'):
|
|
replica.gradient_checkpointing = False
|
|
if hasattr(replica, '_gradient_checkpointing_func'):
|
|
replica._gradient_checkpointing_func = None
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Error cleaning up replica on {dev_name}: {e}")
|
|
try:
|
|
delattr(model, '_parallel_replicas')
|
|
except Exception:
|
|
pass
|
|
|
|
# Remove all parallel attributes
|
|
for attr in ['_true_parallel_active', '_parallel_devices', '_parallel_streams',
|
|
'_parallel_weights', '_auto_vram_balance', '_parallel_purge_cache',
|
|
'_parallel_purge_models']:
|
|
if hasattr(model, attr):
|
|
try:
|
|
delattr(model, attr)
|
|
except Exception:
|
|
pass
|
|
|
|
# VRAM Purge Logic
|
|
if should_purge_models:
|
|
try:
|
|
print("[ParallelAnything] Purging models from VRAM...")
|
|
comfy.model_management.unload_all_models()
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Could not unload models: {e}")
|
|
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
if should_purge_cache:
|
|
try:
|
|
comfy.model_management.soft_empty_cache()
|
|
except:
|
|
pass
|
|
|
|
# Multi-GPU cache clearing
|
|
if torch.cuda.is_available():
|
|
for i in range(torch.cuda.device_count()):
|
|
try:
|
|
with torch.cuda.device(f'cuda:{i}'):
|
|
torch.cuda.empty_cache()
|
|
except:
|
|
pass
|
|
|
|
def extract_model_config(model):
|
|
"""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
|
|
|
|
# Handle params attribute (common in FLUX)
|
|
if hasattr(model, 'params'):
|
|
try:
|
|
params = model.params
|
|
if isinstance(params, dict):
|
|
config.update(params)
|
|
elif is_dataclass(params):
|
|
for field_info in fields(params):
|
|
try:
|
|
val = getattr(params, field_info.name)
|
|
if not isinstance(val, (torch.Tensor, nn.Module)):
|
|
config[field_info.name] = val
|
|
except Exception:
|
|
pass
|
|
else:
|
|
params_dict = vars(params)
|
|
if params_dict:
|
|
config.update({k: v for k, v in params_dict.items()
|
|
if not isinstance(v, (torch.Tensor, nn.Module))})
|
|
except Exception:
|
|
pass
|
|
|
|
# Handle config attribute
|
|
if hasattr(model, 'config'):
|
|
try:
|
|
cfg = model.config
|
|
if isinstance(cfg, dict):
|
|
config.update({k: v for k, v in cfg.items()
|
|
if not isinstance(v, (torch.Tensor, nn.Module))})
|
|
else:
|
|
cfg_dict = {k: v for k, v in vars(cfg).items()
|
|
if not k.startswith('_') and not callable(v)
|
|
and not isinstance(v, (torch.Tensor, nn.Module))}
|
|
config.update(cfg_dict)
|
|
except Exception:
|
|
pass
|
|
|
|
if hasattr(model, 'unet_config') and isinstance(model.unet_config, dict):
|
|
config.update({k: v for k, v in model.unet_config.items()
|
|
if not isinstance(v, (torch.Tensor, nn.Module))})
|
|
|
|
# Filter None values and non-serializable objects
|
|
clean_config = {}
|
|
for k, v in config.items():
|
|
if v is not None and not isinstance(v, (torch.Tensor, nn.Module)):
|
|
try:
|
|
copy.deepcopy(v) # Test if serializable
|
|
clean_config[k] = v
|
|
except Exception:
|
|
pass
|
|
|
|
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)):
|
|
new_val = []
|
|
for v in val:
|
|
if is_dataclass(v):
|
|
new_val.append(clone_dataclass_or_object(v))
|
|
elif isinstance(v, torch.Tensor):
|
|
new_val.append(v.clone())
|
|
else:
|
|
new_val.append(copy.deepcopy(v))
|
|
field_values[field_info.name] = type(val)(new_val)
|
|
else:
|
|
field_values[field_info.name] = copy.deepcopy(val)
|
|
except Exception:
|
|
field_values[field_info.name] = getattr(obj, field_info.name)
|
|
return obj.__class__(**field_values)
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
return copy.deepcopy(obj)
|
|
except Exception:
|
|
return obj
|
|
|
|
def safe_getattr(obj, attr, default=None):
|
|
"""Safely get attribute, returning default if not exists."""
|
|
return getattr(obj, attr, default)
|
|
|
|
def clone_module_simple(module, target_device):
|
|
"""Simple module cloning that handles Parameters and FLUX-specific attributes with FP8 support."""
|
|
if module is None:
|
|
return None
|
|
|
|
module_class = module.__class__
|
|
|
|
# Handle specific layer types explicitly to preserve dtypes and devices
|
|
if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)):
|
|
try:
|
|
weight = safe_getattr(module, 'weight', None)
|
|
weight_dtype = weight.dtype if weight is not None else torch.float32
|
|
has_bias = safe_getattr(module, 'bias', None) is not None
|
|
|
|
# Check FP8 support for target device
|
|
supports_fp8 = device_supports_float8(target_device)
|
|
if not supports_fp8 and is_float8_dtype(weight_dtype):
|
|
weight_dtype = torch.float16
|
|
|
|
if isinstance(module, nn.Linear):
|
|
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):
|
|
new_mod = nn.Conv2d(
|
|
safe_getattr(module, 'in_channels'),
|
|
safe_getattr(module, 'out_channels'),
|
|
safe_getattr(module, 'kernel_size', 1),
|
|
stride=safe_getattr(module, 'stride', 1),
|
|
padding=safe_getattr(module, 'padding', 0),
|
|
dilation=safe_getattr(module, 'dilation', 1),
|
|
groups=safe_getattr(module, 'groups', 1),
|
|
bias=has_bias,
|
|
padding_mode=safe_getattr(module, 'padding_mode', 'zeros'),
|
|
device=target_device, dtype=weight_dtype
|
|
)
|
|
elif isinstance(module, nn.Conv1d):
|
|
new_mod = nn.Conv1d(
|
|
safe_getattr(module, 'in_channels'),
|
|
safe_getattr(module, 'out_channels'),
|
|
safe_getattr(module, 'kernel_size', 1),
|
|
stride=safe_getattr(module, 'stride', 1),
|
|
padding=safe_getattr(module, 'padding', 0),
|
|
dilation=safe_getattr(module, 'dilation', 1),
|
|
groups=safe_getattr(module, 'groups', 1),
|
|
bias=has_bias,
|
|
padding_mode=safe_getattr(module, 'padding_mode', 'zeros'),
|
|
device=target_device, dtype=weight_dtype
|
|
)
|
|
|
|
with torch.no_grad():
|
|
if weight is not None:
|
|
# Convert FP8 to FP16 if needed
|
|
if not supports_fp8 and is_float8_dtype(weight.dtype):
|
|
weight = weight.half()
|
|
new_mod.weight.copy_(weight)
|
|
if has_bias and safe_getattr(module, 'bias') is not None:
|
|
new_mod.bias.copy_(module.bias)
|
|
|
|
return new_mod
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Failed to reconstruct {module_class}: {e}")
|
|
|
|
# Handle normalization layers
|
|
if isinstance(module, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm)):
|
|
try:
|
|
weight = safe_getattr(module, 'weight', None)
|
|
weight_dtype = weight.dtype if weight is not None else torch.float32
|
|
|
|
# Check FP8 support for target device
|
|
supports_fp8 = device_supports_float8(target_device)
|
|
if not supports_fp8 and is_float8_dtype(weight_dtype):
|
|
weight_dtype = torch.float16
|
|
|
|
if isinstance(module, nn.LayerNorm):
|
|
new_mod = nn.LayerNorm(
|
|
safe_getattr(module, 'normalized_shape'),
|
|
eps=safe_getattr(module, 'eps', 1e-5),
|
|
elementwise_affine=safe_getattr(module, 'elementwise_affine', True),
|
|
device=target_device, dtype=weight_dtype
|
|
)
|
|
elif isinstance(module, nn.BatchNorm2d):
|
|
new_mod = nn.BatchNorm2d(
|
|
safe_getattr(module, 'num_features'),
|
|
eps=safe_getattr(module, 'eps', 1e-5),
|
|
momentum=safe_getattr(module, 'momentum', 0.1),
|
|
affine=safe_getattr(module, 'affine', True),
|
|
track_running_stats=safe_getattr(module, 'track_running_stats', True),
|
|
device=target_device, dtype=weight_dtype
|
|
)
|
|
elif isinstance(module, nn.GroupNorm):
|
|
new_mod = nn.GroupNorm(
|
|
safe_getattr(module, 'num_groups'),
|
|
safe_getattr(module, 'num_channels'),
|
|
eps=safe_getattr(module, 'eps', 1e-5),
|
|
affine=safe_getattr(module, 'affine', True),
|
|
device=target_device, dtype=weight_dtype
|
|
)
|
|
|
|
with torch.no_grad():
|
|
if weight is not None:
|
|
if not supports_fp8 and is_float8_dtype(weight.dtype):
|
|
weight = weight.half()
|
|
new_mod.weight.copy_(weight)
|
|
if safe_getattr(module, 'bias') is not None:
|
|
new_mod.bias.copy_(module.bias)
|
|
if isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):
|
|
if safe_getattr(module, 'running_mean') is not None:
|
|
new_mod.running_mean.copy_(module.running_mean)
|
|
if safe_getattr(module, 'running_var') is not None:
|
|
new_mod.running_var.copy_(module.running_var)
|
|
|
|
return new_mod
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Failed to reconstruct norm layer {module_class}: {e}")
|
|
|
|
# Generic fallback
|
|
try:
|
|
new_mod = module_class.__new__(module_class)
|
|
nn.Module.__init__(new_mod)
|
|
|
|
# Check if target supports FP8
|
|
supports_fp8 = device_supports_float8(target_device)
|
|
|
|
# Copy parameters
|
|
for name, param in module.named_parameters(recurse=False):
|
|
if param is not None:
|
|
new_param_data = param.clone().detach()
|
|
# Convert FP8 to FP16 for older GPUs
|
|
if not supports_fp8 and is_float8_dtype(new_param_data.dtype):
|
|
new_param_data = new_param_data.half()
|
|
new_param = nn.Parameter(new_param_data.to(device=target_device),
|
|
requires_grad=False)
|
|
new_mod.register_parameter(name, new_param)
|
|
else:
|
|
new_mod.register_parameter(name, None)
|
|
|
|
# Copy buffers
|
|
for name, buffer in module.named_buffers(recurse=False):
|
|
if buffer is not None:
|
|
new_buffer_data = buffer.clone().detach()
|
|
# Convert FP8 to FP16 for older GPUs
|
|
if not supports_fp8 and is_float8_dtype(new_buffer_data.dtype):
|
|
new_buffer_data = new_buffer_data.half()
|
|
new_buffer = new_buffer_data.to(device=target_device)
|
|
new_mod.register_buffer(name, new_buffer)
|
|
else:
|
|
new_mod.register_buffer(name, None)
|
|
|
|
# Recursively clone children
|
|
for name, child in module.named_children():
|
|
if child is not None:
|
|
cloned_child = clone_module_simple(child, target_device)
|
|
new_mod.add_module(name, cloned_child)
|
|
|
|
# Copy other attributes
|
|
excluded_attrs = {
|
|
'_parameters', '_buffers', '_modules', '_non_persistent_buffers_set',
|
|
'_backward_pre_hooks', '_backward_hooks', '_forward_pre_hooks', '_forward_hooks',
|
|
'_state_dict_hooks', '_load_state_dict_pre_hooks', '_extra_state', '_modules_to_load'
|
|
}
|
|
cache_attrs = {
|
|
'img_ids', 'txt_ids', '_img_ids', '_txt_ids', 'cached_img_ids', 'cached_txt_ids',
|
|
'pos_emb', '_pos_emb', 'pos_embed', '_pos_embed', 'freqs_cis', '_freqs_cis',
|
|
'freqs', '_freqs', 'cache', '_cache', 'kv_cache', '_kv_cache', 'attn_bias', '_attn_bias'
|
|
}
|
|
|
|
for key, value in module.__dict__.items():
|
|
if key in excluded_attrs:
|
|
continue
|
|
try:
|
|
if is_dataclass(value):
|
|
setattr(new_mod, key, clone_dataclass_or_object(value))
|
|
elif isinstance(value, torch.Tensor):
|
|
if key in cache_attrs:
|
|
setattr(new_mod, key, None)
|
|
else:
|
|
tensor_val = value.clone().detach()
|
|
# Convert FP8 to FP16 if needed
|
|
if not supports_fp8 and is_float8_dtype(tensor_val.dtype):
|
|
tensor_val = tensor_val.half()
|
|
setattr(new_mod, key, tensor_val.to(target_device))
|
|
elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], torch.Tensor):
|
|
if key in cache_attrs:
|
|
setattr(new_mod, key, None)
|
|
else:
|
|
converted = []
|
|
for t in value:
|
|
t_conv = t.clone().detach()
|
|
if not supports_fp8 and is_float8_dtype(t_conv.dtype):
|
|
t_conv = t_conv.half()
|
|
converted.append(t_conv.to(target_device))
|
|
setattr(new_mod, key, type(value)(converted))
|
|
else:
|
|
setattr(new_mod, key, copy.deepcopy(value))
|
|
except Exception:
|
|
pass
|
|
|
|
return new_mod
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to clone module {module_class}: {e}")
|
|
|
|
def safe_model_clone(source_model, target_device, disable_flash=False):
|
|
"""FLUX-safe model cloning with explicit memory cleanup and FP8 handling."""
|
|
clear_flux_caches(source_model)
|
|
|
|
# Determine source device
|
|
try:
|
|
src_device = next(source_model.parameters()).device
|
|
except StopIteration:
|
|
src_device = torch.device('cpu')
|
|
|
|
# If already on target, return reference to avoid double memory usage
|
|
if str(src_device) == str(target_device):
|
|
print(f"[ParallelAnything] Model already on {target_device}, using reference...")
|
|
clear_flux_caches(source_model)
|
|
return source_model
|
|
|
|
# Move to CPU for safe cloning if on GPU
|
|
if src_device.type != 'cpu':
|
|
model_cpu = source_model.cpu()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
comfy.model_management.soft_empty_cache()
|
|
else:
|
|
model_cpu = source_model
|
|
|
|
model_class = model_cpu.__class__
|
|
replica = None
|
|
method_used = "unknown"
|
|
|
|
try:
|
|
# Try config reconstruction with direct-to-GPU creation
|
|
print(f"[ParallelAnything] Attempting efficient clone to {target_device}...")
|
|
config = extract_model_config(model_cpu)
|
|
state_dict = model_cpu.state_dict()
|
|
|
|
if not config:
|
|
raise RuntimeError("Could not extract model config")
|
|
|
|
# Create model directly on target device if possible
|
|
try:
|
|
replica = model_class(**config)
|
|
replica = replica.to(target_device, non_blocking=False)
|
|
except TypeError as e:
|
|
if "missing" in str(e) and "required" in str(e):
|
|
try:
|
|
replica = model_class(config)
|
|
replica = replica.to(target_device, non_blocking=False)
|
|
except Exception:
|
|
config_obj = SimpleNamespace(**config)
|
|
replica = model_class(config_obj)
|
|
replica = replica.to(target_device, non_blocking=False)
|
|
else:
|
|
raise
|
|
|
|
# Load state dict piece by piece to minimize memory spike
|
|
print(f"[ParallelAnything] Loading parameters incrementally...")
|
|
supports_fp8 = device_supports_float8(target_device)
|
|
|
|
with torch.no_grad():
|
|
for key in list(state_dict.keys()):
|
|
try:
|
|
param = state_dict[key]
|
|
if '.' in key:
|
|
# Navigate to submodule
|
|
parts = key.split('.')
|
|
module = replica
|
|
for part in parts[:-1]:
|
|
module = getattr(module, part)
|
|
target_param = getattr(module, parts[-1])
|
|
else:
|
|
target_param = getattr(replica, key)
|
|
|
|
if isinstance(target_param, (torch.nn.Parameter, torch.Tensor)):
|
|
# Load directly to target device with FP8 handling
|
|
param_data = param.to(target_device, non_blocking=False)
|
|
if is_float8_dtype(param_data.dtype) and not supports_fp8:
|
|
param_data = param_data.half()
|
|
target_param.data.copy_(param_data)
|
|
|
|
# Delete immediately to free memory
|
|
del state_dict[key]
|
|
del param
|
|
|
|
# Periodic cleanup for very large models
|
|
if len(state_dict) % 50 == 0 and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Could not load {key}: {e}")
|
|
|
|
del state_dict
|
|
if src_device.type != 'cpu':
|
|
del model_cpu
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
method_used = "incremental_gpu"
|
|
|
|
except Exception as recon_error:
|
|
print(f"[ParallelAnything] Config reconstruction failed: {recon_error}")
|
|
print(f"[ParallelAnything] Attempting manual recursive cloning...")
|
|
|
|
# Fallback to manual cloning
|
|
replica = clone_module_simple(model_cpu, target_device)
|
|
method_used = "manual_recursive_clone"
|
|
|
|
if src_device.type != 'cpu':
|
|
del model_cpu
|
|
gc.collect()
|
|
|
|
if replica is None:
|
|
raise RuntimeError("All clone methods failed")
|
|
|
|
# Post-processing
|
|
clear_flux_caches(replica)
|
|
|
|
# CRITICAL FIX: Convert FP8 to FP16 for older GPUs after cloning
|
|
if not device_supports_float8(target_device):
|
|
print(f"[ParallelAnything] Converting FP8 weights to FP16 for {target_device}...")
|
|
converted_from_fp8 = False
|
|
|
|
def convert_fp8_recursive(m):
|
|
nonlocal converted_from_fp8
|
|
for name, param in m.named_parameters():
|
|
if is_float8_dtype(param.dtype):
|
|
param.data = param.data.half()
|
|
converted_from_fp8 = True
|
|
for name, buffer in m.named_buffers():
|
|
if is_float8_dtype(buffer.dtype):
|
|
buffer.data = buffer.data.half()
|
|
converted_from_fp8 = True
|
|
for child in m.children():
|
|
convert_fp8_recursive(child)
|
|
|
|
convert_fp8_recursive(replica)
|
|
if converted_from_fp8:
|
|
replica._converted_from_fp8 = True
|
|
print(f"[ParallelAnything] Converted FP8 weights to FP16 for {target_device}")
|
|
|
|
# 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
|
|
|
|
# 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
|
|
|
|
# Ensure buffers are on correct device
|
|
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 Exception:
|
|
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)
|
|
|
|
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]) # Ensure sum = batch_size
|
|
|
|
return split_sizes
|
|
|
|
class ParallelDevice:
|
|
@classmethod
|
|
def get_available_devices(cls):
|
|
devices = ["cpu"]
|
|
if torch.cuda.is_available():
|
|
for i in range(torch.cuda.device_count()):
|
|
devices.append(f"cuda:{i}")
|
|
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
devices.append("mps")
|
|
if hasattr(torch, 'xpu') and torch.xpu.is_available():
|
|
for i in range(torch.xpu.device_count()):
|
|
devices.append(f"xpu:{i}")
|
|
try:
|
|
import torch_directml
|
|
for i in range(torch_directml.device_count()):
|
|
devices.append(f"privateuseone:{i}")
|
|
except ImportError:
|
|
pass
|
|
return devices
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
available = s.get_available_devices()
|
|
default = "cuda:0" if any(d.startswith("cuda:0") for d in available) else available[0]
|
|
return {
|
|
"required": {
|
|
"device_id": (available, {
|
|
"default": default,
|
|
"tooltip": "Select available compute device (CPU/CUDA/MPS/XPU)"
|
|
}),
|
|
"percentage": ("FLOAT", {
|
|
"default": 50.0, "min": 1.0, "max": 100.0, "step": 1.0,
|
|
"tooltip": "Percentage of batch to process on this device"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"previous_devices": ("DEVICE_CHAIN", {
|
|
"tooltip": "Connect from another ParallelDevice node to chain multiple GPUs"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("DEVICE_CHAIN",)
|
|
RETURN_NAMES = ("device_chain",)
|
|
FUNCTION = "add_device"
|
|
CATEGORY = "utils/hardware"
|
|
DESCRIPTION = "Add a GPU/CPU/MPS/XPU device to the parallel processing chain"
|
|
|
|
def add_device(self, device_id, percentage, previous_devices=None):
|
|
if previous_devices is None:
|
|
previous_devices = []
|
|
|
|
config = {
|
|
"device": device_id,
|
|
"percentage": float(percentage),
|
|
"weight": float(percentage) / 100.0
|
|
}
|
|
|
|
new_chain = previous_devices.copy()
|
|
new_chain.append(config)
|
|
|
|
return (new_chain,)
|
|
|
|
class ParallelDeviceList:
|
|
@classmethod
|
|
def get_available_devices(cls):
|
|
devices = ["cpu"]
|
|
if torch.cuda.is_available():
|
|
for i in range(torch.cuda.device_count()):
|
|
devices.append(f"cuda:{i}")
|
|
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
devices.append("mps")
|
|
if hasattr(torch, 'xpu') and torch.xpu.is_available():
|
|
for i in range(torch.xpu.device_count()):
|
|
devices.append(f"xpu:{i}")
|
|
return devices
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
devices = s.get_available_devices()
|
|
def_dev = "cuda:0" if "cuda:0" in devices else devices[0]
|
|
return {
|
|
"required": {
|
|
"device_1": (devices, {"default": def_dev}),
|
|
"pct_1": ("FLOAT", {"default": 50.0, "min": 1.0, "max": 100.0, "step": 1.0}),
|
|
"device_2": (devices, {"default": devices[1] if len(devices) > 1 else def_dev}),
|
|
"pct_2": ("FLOAT", {"default": 50.0, "min": 0.0, "max": 100.0, "step": 1.0}),
|
|
},
|
|
"optional": {
|
|
"device_3": (devices, {"default": devices[2] if len(devices) > 2 else "cpu"}),
|
|
"pct_3": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 1.0}),
|
|
"device_4": (devices, {"default": devices[3] if len(devices) > 3 else "cpu"}),
|
|
"pct_4": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("DEVICE_CHAIN",)
|
|
RETURN_NAMES = ("device_chain",)
|
|
FUNCTION = "create_list"
|
|
CATEGORY = "utils/hardware"
|
|
|
|
def create_list(self, device_1, pct_1, device_2, pct_2, device_3="cpu", pct_3=0, device_4="cpu", pct_4=0):
|
|
chain = []
|
|
devices = [(device_1, pct_1), (device_2, pct_2), (device_3, pct_3), (device_4, pct_4)]
|
|
|
|
for dev_str, pct in devices:
|
|
if pct > 0:
|
|
chain.append({
|
|
"device": dev_str,
|
|
"percentage": float(pct),
|
|
"weight": float(pct) / 100.0
|
|
})
|
|
|
|
return (chain,)
|
|
|
|
class ParallelAnything:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"device_chain": ("DEVICE_CHAIN", {"tooltip": "Connect from ParallelDevice nodes"}),
|
|
},
|
|
"optional": {
|
|
"workload_split": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Enable multi-device processing"
|
|
}),
|
|
"auto_vram_balance": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Automatically adjust batch split based on available VRAM"
|
|
}),
|
|
"purge_cache": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Purge CUDA cache when cleaning up parallel resources"
|
|
}),
|
|
"purge_models": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Unload all models from VRAM when cleaning up (aggressive memory clearing)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "setup_parallel"
|
|
CATEGORY = "utils/hardware"
|
|
|
|
def setup_parallel(self, model, device_chain, workload_split=True, auto_vram_balance=False,
|
|
purge_cache=True, purge_models=False):
|
|
if model is None or not device_chain:
|
|
return (model,)
|
|
|
|
# Get the actual diffusion model
|
|
target_model = model.model.diffusion_model
|
|
|
|
# Cleanup existing parallel setup if present
|
|
if hasattr(target_model, "_true_parallel_active"):
|
|
print("[ParallelAnything] Cleaning up previous parallel setup...")
|
|
cleanup_parallel_model(target_model)
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
# AGGRESSIVE: Unload ALL other models from ComfyUI before cloning to free VRAM
|
|
print("[ParallelAnything] Freeing ComfyUI model cache for cloning...")
|
|
comfy.model_management.unload_all_models()
|
|
aggressive_cleanup()
|
|
|
|
# Validate percentages
|
|
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]))}")
|
|
|
|
# Check for devices needing safe attention (older GPUs)
|
|
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)
|
|
except Exception:
|
|
print(f"[ParallelAnything] Invalid device: {dev}")
|
|
return (model,)
|
|
|
|
# Store original device
|
|
try:
|
|
original_device = next(target_model.parameters()).device
|
|
original_device_str = str(original_device)
|
|
except StopIteration:
|
|
original_device = torch.device('cpu')
|
|
original_device_str = "cpu"
|
|
|
|
replicas = {}
|
|
streams = {}
|
|
successful_devices = []
|
|
successful_weights = []
|
|
|
|
try:
|
|
print(f"[ParallelAnything] Cloning to {len(device_names)} devices...")
|
|
clear_flux_caches(target_model)
|
|
|
|
# Move original to CPU to free VRAM for cloning, unless it's the only device and matches
|
|
needs_cpu_transition = original_device.type == 'cuda' and original_device_str not in device_names
|
|
if needs_cpu_transition:
|
|
print(f"[ParallelAnything] Moving model to CPU for safe cloning...")
|
|
target_model = target_model.cpu()
|
|
aggressive_cleanup()
|
|
|
|
# Clone to each device with OOM handling
|
|
for i, dev_name in enumerate(device_names):
|
|
dev = torch.device(dev_name)
|
|
need_safe = dev_name in devices_needing_safe_attention
|
|
|
|
# CRITICAL: If target is same as original device, use original directly
|
|
if dev_name == original_device_str and not needs_cpu_transition:
|
|
print(f"[ParallelAnything] Using original model for {dev_name} (skipping clone)")
|
|
replicas[dev_name] = target_model
|
|
if dev.type == 'cuda':
|
|
streams[dev_name] = torch.cuda.Stream(dev)
|
|
successful_devices.append(dev_name)
|
|
successful_weights.append(weights[i])
|
|
continue
|
|
|
|
try:
|
|
print(f"[ParallelAnything] Cloning to {dev_name} ({i+1}/{len(device_names)})...")
|
|
|
|
# Check VRAM availability before clone
|
|
if dev.type == 'cuda':
|
|
idx = dev.index if dev.index is not None else int(dev_name.split(":")[-1])
|
|
torch.cuda.synchronize(idx)
|
|
free_mem = torch.cuda.mem_get_info(idx)[0] / 1024**3
|
|
print(f"[ParallelAnything] VRAM available on {dev_name}: {free_mem:.2f}GB")
|
|
|
|
replica = safe_model_clone(target_model, dev, disable_flash=need_safe)
|
|
clear_flux_caches(replica)
|
|
replicas[dev_name] = replica
|
|
|
|
if dev.type == 'cuda':
|
|
streams[dev_name] = torch.cuda.Stream(dev)
|
|
|
|
print(f"[ParallelAnything] ✓ {dev_name}" + (" (Safe mode)" if need_safe else ""))
|
|
successful_devices.append(dev_name)
|
|
successful_weights.append(weights[i])
|
|
|
|
# Aggressive cleanup after clone
|
|
aggressive_cleanup()
|
|
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e).lower():
|
|
print(f"[ParallelAnything] OOM cloning to {dev_name}, skipping...")
|
|
aggressive_cleanup()
|
|
continue
|
|
else:
|
|
raise
|
|
|
|
# Handle case where we lost some devices
|
|
if len(successful_devices) < len(device_names):
|
|
if len(successful_devices) == 0:
|
|
raise RuntimeError("No devices available for parallel processing - all OOM")
|
|
print(f"[ParallelAnything] Reducing to {len(successful_devices)} devices due to OOM")
|
|
device_names = successful_devices
|
|
# Renormalize weights
|
|
total_weight = sum(successful_weights)
|
|
weights = [w/total_weight for w in successful_weights]
|
|
|
|
# Restore original only if we moved it and it's not in our replicas
|
|
if needs_cpu_transition and original_device_str not in replicas:
|
|
try:
|
|
print(f"[ParallelAnything] Restoring original model to {original_device}")
|
|
target_model = target_model.to(original_device)
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Warning: Could not restore original model: {e}")
|
|
|
|
except Exception as e:
|
|
print(f"[ParallelAnything] Error during cloning: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
# Cleanup partial replicas
|
|
for dev_name, r in replicas.items():
|
|
try:
|
|
if r is not target_model:
|
|
r.cpu()
|
|
del r
|
|
except Exception:
|
|
pass
|
|
|
|
aggressive_cleanup()
|
|
return (model,)
|
|
|
|
# Store purge preferences
|
|
target_model._parallel_purge_cache = purge_cache
|
|
target_model._parallel_purge_models = purge_models
|
|
|
|
# Setup parallel forward function
|
|
# Store references in closure - use tuples to prevent modification
|
|
replicas_ref = replicas
|
|
devices_ref = tuple(device_names)
|
|
weights_ref = tuple(weights)
|
|
lead_device = torch.device(device_names[0])
|
|
streams_ref = streams
|
|
auto_balance_ref = auto_vram_balance
|
|
original_device_ref = original_device_str
|
|
|
|
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:
|
|
if not all(b == batch_sizes[0] for b in batch_sizes):
|
|
raise ValueError(f"Inconsistent batch sizes: {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 = []
|
|
for t in x:
|
|
if isinstance(t, torch.Tensor):
|
|
split_lists.append(torch.split(t, split_sizes, dim=0))
|
|
else:
|
|
# Non-tensor, broadcast
|
|
split_lists.append([t] * len(split_sizes))
|
|
|
|
# Transpose: list of splits -> splits of lists
|
|
result = []
|
|
for i in range(len(split_sizes)):
|
|
result.append(type(x)(sl[i] for sl in split_lists))
|
|
return result
|
|
else:
|
|
# Non-tensor, broadcast to all
|
|
return [x] * len(split_sizes)
|
|
|
|
def move_to_device(x, device, non_blocking=False):
|
|
"""Move tensor or list/tuple of tensors to device with FP8 handling."""
|
|
if isinstance(x, torch.Tensor):
|
|
# Only move if not already on target device
|
|
if x.device != torch.device(device):
|
|
x = x.to(device, non_blocking=non_blocking)
|
|
# CRITICAL FIX: Convert FP8 to FP16 for old GPUs
|
|
if is_float8_dtype(x.dtype) and not device_supports_float8(device):
|
|
x = x.half()
|
|
return x
|
|
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, total_batch):
|
|
"""Split keyword arguments that contain batch tensors."""
|
|
split_kwargs_list = [{} for _ in range(len(split_sizes))]
|
|
|
|
for key, value in kwargs.items():
|
|
if isinstance(value, torch.Tensor) and value.shape[0] == total_batch:
|
|
chunks = torch.split(value, split_sizes, dim=0)
|
|
for i, chunk in enumerate(chunks):
|
|
split_kwargs_list[i][key] = chunk
|
|
elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], torch.Tensor):
|
|
if all(isinstance(t, torch.Tensor) and t.shape[0] == total_batch for t in value):
|
|
split_lists = [torch.split(t, split_sizes, dim=0) for t in value]
|
|
for i in range(len(split_sizes)):
|
|
split_kwargs_list[i][key] = type(value)(sl[i] for sl in split_lists)
|
|
else:
|
|
for i in range(len(split_sizes)):
|
|
split_kwargs_list[i][key] = value
|
|
|
|
return split_kwargs_list
|
|
|
|
def concatenate_results(results, dim=0):
|
|
"""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)):
|
|
concatenated = []
|
|
for i in range(len(first)):
|
|
if isinstance(first[i], torch.Tensor):
|
|
to_concat = [r[i] for r in results]
|
|
concatenated.append(torch.cat(to_concat, dim=dim))
|
|
else:
|
|
concatenated.append(first[i])
|
|
return type(first)(concatenated)
|
|
else:
|
|
return results
|
|
|
|
def parallel_forward(self, x, timesteps, context=None, **kwargs):
|
|
try:
|
|
# Validate batch size
|
|
try:
|
|
batch_size = get_batch_size(x)
|
|
except ValueError as e:
|
|
print(f"[ParallelAnything] Error: {e}")
|
|
raise
|
|
|
|
# Fallback to single device if batch too small
|
|
if batch_size < len(devices_ref) or not workload_split:
|
|
with torch.no_grad():
|
|
first_replica = replicas_ref[devices_ref[0]]
|
|
# Use _original_forward to avoid recursion if available
|
|
if hasattr(first_replica, '_original_forward'):
|
|
return first_replica._original_forward(x, timesteps, context=context, **kwargs)
|
|
else:
|
|
return first_replica(x, timesteps, context=context, **kwargs)
|
|
|
|
# Calculate split sizes
|
|
if auto_balance_ref:
|
|
split_sizes = auto_split_batch(batch_size, list(devices_ref), list(weights_ref))
|
|
else:
|
|
split_sizes = [max(1, int(batch_size * w)) for w in weights_ref]
|
|
split_sizes[-1] = batch_size - sum(split_sizes[:-1])
|
|
|
|
# Build active device list
|
|
active = []
|
|
for idx, (dev_name, size) in enumerate(zip(devices_ref, split_sizes)):
|
|
if size > 0:
|
|
active.append({
|
|
'idx': idx,
|
|
'dev_name': dev_name,
|
|
'device': torch.device(dev_name),
|
|
'replica': replicas_ref[dev_name],
|
|
'stream': streams_ref.get(dev_name),
|
|
'size': size,
|
|
})
|
|
|
|
if len(active) == 0:
|
|
raise RuntimeError("No active devices available")
|
|
|
|
if len(active) == 1:
|
|
with torch.no_grad():
|
|
replica = active[0]['replica']
|
|
# Get model dtype - could be Float, Half, BFloat16, etc.
|
|
model_dtype = next(replica.parameters()).dtype
|
|
|
|
# CRITICAL FIX: Always ensure input dtype matches model dtype
|
|
# This fixes both FP8->FP16 conversion and works for video models
|
|
if x.dtype != model_dtype and x.is_floating_point():
|
|
x = x.to(model_dtype)
|
|
if context is not None and context.dtype != model_dtype and context.is_floating_point():
|
|
context = context.to(model_dtype)
|
|
|
|
if hasattr(replica, '_original_forward'):
|
|
return replica._original_forward(x, timesteps, context=context, **kwargs)
|
|
else:
|
|
return replica(x, timesteps, context=context, **kwargs)
|
|
|
|
# Split inputs
|
|
x_chunks = split_batch(x, [a['size'] for a in active])
|
|
t_chunks = split_batch(timesteps, [a['size'] for a in active])
|
|
|
|
if context is not None:
|
|
c_chunks = split_batch(context, [a['size'] for a in active])
|
|
else:
|
|
c_chunks = [None] * len(active)
|
|
|
|
kwargs_chunks = split_kwargs(kwargs, [a['size'] for a in active], batch_size)
|
|
|
|
results = [None] * len(active)
|
|
exceptions = []
|
|
|
|
def worker(task_idx):
|
|
task = active[task_idx]
|
|
dev = task['device']
|
|
replica = task['replica']
|
|
stream = task['stream']
|
|
|
|
try:
|
|
# Move inputs to device
|
|
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
|
|
|
|
# CRITICAL FIX: Get model dtype and ensure inputs match it
|
|
model_dtype = next(replica.parameters()).dtype
|
|
|
|
# Convert floating point tensors to match model dtype
|
|
# This handles: Flux FP16 model + BF16 inputs -> both FP16
|
|
# And: WAN Float model + Float inputs -> both Float (no change)
|
|
if x_in.dtype != model_dtype and x_in.is_floating_point():
|
|
x_in = x_in.to(model_dtype)
|
|
if c_in is not None and c_in.dtype != model_dtype and c_in.is_floating_point():
|
|
c_in = c_in.to(model_dtype)
|
|
|
|
# Move kwargs and ensure dtype compatibility
|
|
k_in = {}
|
|
for k, v in kwargs_chunks[task_idx].items():
|
|
v_moved = move_to_device(v, dev, non_blocking=False)
|
|
if isinstance(v_moved, torch.Tensor) and v_moved.is_floating_point() and v_moved.dtype != model_dtype:
|
|
v_moved = v_moved.to(model_dtype)
|
|
k_in[k] = v_moved
|
|
|
|
# Get forward function
|
|
if hasattr(replica, '_original_forward'):
|
|
forward_fn = replica._original_forward
|
|
else:
|
|
forward_fn = replica.forward
|
|
|
|
# Execute with appropriate synchronization
|
|
if dev.type == 'cuda' and stream is not None:
|
|
with torch.cuda.device(dev):
|
|
with torch.cuda.stream(stream):
|
|
torch.cuda.synchronize(dev)
|
|
with torch.no_grad():
|
|
out = forward_fn(x_in, t_in, context=c_in, **k_in)
|
|
torch.cuda.synchronize(dev)
|
|
elif dev.type == 'cuda':
|
|
with torch.cuda.device(dev):
|
|
torch.cuda.synchronize(dev)
|
|
with torch.no_grad():
|
|
out = forward_fn(x_in, t_in, context=c_in, **k_in)
|
|
torch.cuda.synchronize(dev)
|
|
elif dev.type == 'xpu':
|
|
with torch.xpu.device(dev):
|
|
torch.xpu.synchronize(dev)
|
|
with torch.no_grad():
|
|
out = forward_fn(x_in, t_in, context=c_in, **k_in)
|
|
torch.xpu.synchronize(dev)
|
|
else:
|
|
with torch.no_grad():
|
|
out = forward_fn(x_in, t_in, context=c_in, **k_in)
|
|
|
|
# Move result back to lead device
|
|
out = move_to_device(out, lead_device, non_blocking=False)
|
|
|
|
return task_idx, out
|
|
|
|
except Exception as e:
|
|
return task_idx, e
|
|
|
|
# Execute in thread pool
|
|
with ThreadPoolExecutor(max_workers=len(active)) as executor:
|
|
futures = [executor.submit(worker, i) for i in range(len(active))]
|
|
|
|
for future in as_completed(futures):
|
|
idx, result = future.result()
|
|
if isinstance(result, Exception):
|
|
exceptions.append((active[idx]['dev_name'], result))
|
|
results[idx] = None
|
|
else:
|
|
results[idx] = result
|
|
|
|
# Handle exceptions
|
|
if exceptions:
|
|
for dev_name, exc in exceptions:
|
|
print(f"[ParallelAnything] Error on {dev_name}: {exc}")
|
|
raise exceptions[0][1]
|
|
|
|
# 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)
|
|
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e).lower():
|
|
print("[ParallelAnything] OOM in parallel forward, attempting emergency cleanup...")
|
|
aggressive_cleanup()
|
|
|
|
# Try single device fallback on lead device
|
|
print("[ParallelAnything] Falling back to single device processing...")
|
|
with torch.no_grad():
|
|
first_replica = replicas_ref[devices_ref[0]]
|
|
if hasattr(first_replica, '_original_forward'):
|
|
return first_replica._original_forward(x, timesteps, context=context, **kwargs)
|
|
else:
|
|
return first_replica(x, timesteps, context=context, **kwargs)
|
|
else:
|
|
raise
|
|
|
|
# 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
|
|
target_model._parallel_streams = streams
|
|
target_model._parallel_weights = weights
|
|
target_model._auto_vram_balance = auto_vram_balance
|
|
|
|
# Register cleanup
|
|
weakref.finalize(model, cleanup_parallel_model, weakref.ref(target_model))
|
|
|
|
# Set load device to lead device for ComfyUI memory management
|
|
model.load_device = lead_device
|
|
|
|
print(f"[ParallelAnything] Parallel setup complete. Using devices: {device_names}")
|
|
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)",
|
|
}
|