899 lines
39 KiB
Python
899 lines
39 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
|
||
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',
|
||
]
|
||
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):
|
||
"""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:
|
||
delattr(model, '_original_forward')
|
||
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:
|
||
delattr(model, attr)
|
||
except:
|
||
pass
|
||
|
||
# Force garbage collection
|
||
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',
|
||
]
|
||
|
||
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):
|
||
"""Safely get attribute, returning default if it doesn't exist."""
|
||
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__
|
||
|
||
# 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)
|
||
out_channels = safe_getattr(module, 'out_features') or safe_getattr(module, 'out_channels', 0)
|
||
has_bias = safe_getattr(module, 'bias') is not None
|
||
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
|
||
)
|
||
elif isinstance(module, nn.Conv2d):
|
||
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):
|
||
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:
|
||
normalized_shape = safe_getattr(module, 'normalized_shape')
|
||
num_features = safe_getattr(module, 'num_features', 0)
|
||
num_groups = safe_getattr(module, 'num_groups', 1)
|
||
num_channels = safe_getattr(module, 'num_channels', 0)
|
||
eps = safe_getattr(module, 'eps', 1e-5)
|
||
elementwise_affine = safe_getattr(module, 'elementwise_affine', True)
|
||
momentum = safe_getattr(module, 'momentum', 0.1)
|
||
affine = safe_getattr(module, 'affine', True)
|
||
track_running_stats = safe_getattr(module, 'track_running_stats', True)
|
||
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):
|
||
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
|
||
)
|
||
elif isinstance(module, nn.GroupNorm):
|
||
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}")
|
||
|
||
# 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_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)
|
||
new_mod.register_buffer(name, new_buffer)
|
||
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',
|
||
'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']:
|
||
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
|
||
|
||
# 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)
|
||
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}")
|
||
|
||
# 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()
|
||
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:
|
||
"""Alternative: Parallel layout with dropdowns for 1-4 devices."""
|
||
|
||
@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:
|
||
"""Main execution node with TRUE parallelism via model replication."""
|
||
|
||
@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"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "setup_parallel"
|
||
CATEGORY = "utils/hardware"
|
||
|
||
def setup_parallel(self, model, device_chain, workload_split=True):
|
||
if model is None or not device_chain:
|
||
return (model,)
|
||
|
||
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()
|
||
if torch.cuda.is_available():
|
||
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,)
|
||
|
||
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])
|
||
|
||
# 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)
|
||
except:
|
||
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()
|
||
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
|
||
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}")
|
||
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 to restore original device on error
|
||
try:
|
||
if original_device.type == 'cuda':
|
||
target_model = target_model.to(original_device)
|
||
except:
|
||
pass
|
||
return (model,)
|
||
|
||
# Store references for closure
|
||
replicas_ref = replicas
|
||
devices_ref = device_names
|
||
weights_ref = weights
|
||
lead_ref = lead_device
|
||
|
||
def split_kwargs(kwargs, split_sizes, device):
|
||
"""FLUX FIX: Split any tensor kwargs that have batch dimension."""
|
||
split_kwargs_list = [{} for _ in range(len(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
|
||
chunks = torch.split(value, split_sizes, dim=0)
|
||
for i, chunk in enumerate(chunks):
|
||
split_kwargs_list[i][key] = chunk.to(device)
|
||
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 parallel_forward(self, x, timesteps, context=None, **kwargs):
|
||
batch_size = x.shape[0]
|
||
|
||
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)
|
||
|
||
# Calculate splits
|
||
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],
|
||
'size': size,
|
||
})
|
||
|
||
if len(active) == 1:
|
||
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)
|
||
|
||
# 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:
|
||
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)
|
||
|
||
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):
|
||
idx, result = future.result()
|
||
if isinstance(result, Exception):
|
||
raise result
|
||
results[idx] = result
|
||
|
||
return torch.cat(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,)
|
||
|
||
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)",
|
||
}
|