FearL0rd
+188
-113
@@ -68,23 +68,22 @@ class ParallelBlock(nn.Module):
|
||||
return self.local_block(*args, **kwargs)
|
||||
|
||||
target_device = self.owner_device
|
||||
|
||||
|
||||
# --- CRITICAL FIX ---
|
||||
# Do NOT check args[0].device to skip moving.
|
||||
# Main input might be on target_device (from prev block),
|
||||
# Do NOT check args[0].device to skip moving.
|
||||
# Main input might be on target_device (from prev block),
|
||||
# but auxiliary inputs (timesteps, context) might still be on Lead Device.
|
||||
# We must ensure ALL inputs are on target_device.
|
||||
|
||||
args_moved = self._move_args(args, target_device)
|
||||
kwargs_moved = self._move_kwargs(kwargs, target_device)
|
||||
|
||||
|
||||
# Execute on target replica
|
||||
out = self.owner_block(*args_moved, **kwargs_moved)
|
||||
|
||||
if self.is_last_block:
|
||||
# Return data to Lead Device for final layers/output processing
|
||||
out = self._move_tensor(out, self.lead_device)
|
||||
|
||||
|
||||
return out
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
@@ -151,8 +150,7 @@ def disable_flash_xformers(model):
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if any(x in name.lower() for x in ['attn', 'attention', 'transformer']):
|
||||
for attr in ['use_xformers', 'use_flash_attention', 'use_flash_attention_2',
|
||||
'_use_memory_efficient_attention', 'enable_flash', 'enable_xformers']:
|
||||
for attr in ['use_xformers', 'use_flash_attention', 'use_flash_attention_2', '_use_memory_efficient_attention', 'enable_flash', 'enable_xformers']:
|
||||
if hasattr(module, attr):
|
||||
try:
|
||||
setattr(module, attr, False)
|
||||
@@ -219,10 +217,10 @@ def cleanup_parallel_model(model_ref):
|
||||
return
|
||||
|
||||
print("[ParallelAnything] Cleaning up parallel model...")
|
||||
|
||||
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
|
||||
@@ -230,6 +228,7 @@ def cleanup_parallel_model(model_ref):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Cleanup replicas
|
||||
if hasattr(model, '_parallel_replicas'):
|
||||
replicas = model._parallel_replicas
|
||||
for dev_name, replica in list(replicas.items()):
|
||||
@@ -243,12 +242,12 @@ def cleanup_parallel_model(model_ref):
|
||||
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
|
||||
|
||||
# Clear other attributes
|
||||
for attr in ['_true_parallel_active', '_parallel_devices', '_parallel_streams',
|
||||
'_parallel_weights', '_auto_vram_balance', '_parallel_purge_cache',
|
||||
'_parallel_purge_models']:
|
||||
@@ -274,14 +273,13 @@ def cleanup_parallel_model(model_ref):
|
||||
comfy.model_management.soft_empty_cache()
|
||||
except:
|
||||
pass
|
||||
|
||||
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
|
||||
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):
|
||||
config = {}
|
||||
@@ -290,11 +288,11 @@ def extract_model_config(model):
|
||||
'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',
|
||||
'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:
|
||||
@@ -322,8 +320,7 @@ def extract_model_config(model):
|
||||
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))})
|
||||
config.update({k: v for k, v in params_dict.items() if not isinstance(v, (torch.Tensor, nn.Module))})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -331,19 +328,15 @@ def extract_model_config(model):
|
||||
try:
|
||||
cfg = model.config
|
||||
if isinstance(cfg, dict):
|
||||
config.update({k: v for k, v in cfg.items()
|
||||
if not isinstance(v, (torch.Tensor, nn.Module))})
|
||||
config.update({k: v for k, v in cfg.items() if not isinstance(v, (torch.Tensor, nn.Module))})
|
||||
else:
|
||||
cfg_dict = {k: v for k, v in vars(cfg).items()
|
||||
if not k.startswith('_') and not callable(v)
|
||||
and not isinstance(v, (torch.Tensor, nn.Module))}
|
||||
cfg_dict = {k: v for k, v in vars(cfg).items() if not k.startswith('_') and not callable(v) and not isinstance(v, (torch.Tensor, nn.Module))}
|
||||
config.update(cfg_dict)
|
||||
except 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))})
|
||||
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():
|
||||
@@ -385,6 +378,7 @@ def clone_dataclass_or_object(obj, target_device=None):
|
||||
return obj.__class__(**field_values)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
return copy.deepcopy(obj)
|
||||
except Exception:
|
||||
@@ -410,22 +404,41 @@ def clone_module_simple(module, target_device):
|
||||
weight_dtype = torch.float16
|
||||
|
||||
if isinstance(module, nn.Linear):
|
||||
new_mod = nn.Linear(safe_getattr(module, 'in_features'), safe_getattr(module, 'out_features'),
|
||||
bias=has_bias, device=target_device, dtype=weight_dtype)
|
||||
new_mod = nn.Linear(
|
||||
safe_getattr(module, 'in_features'),
|
||||
safe_getattr(module, '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)
|
||||
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)
|
||||
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:
|
||||
@@ -442,27 +455,37 @@ def clone_module_simple(module, target_device):
|
||||
try:
|
||||
weight = safe_getattr(module, 'weight', None)
|
||||
weight_dtype = weight.dtype if weight is not None else torch.float32
|
||||
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
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:
|
||||
@@ -485,6 +508,7 @@ def clone_module_simple(module, target_device):
|
||||
nn.Module.__init__(new_mod)
|
||||
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()
|
||||
@@ -495,6 +519,7 @@ def clone_module_simple(module, target_device):
|
||||
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()
|
||||
@@ -505,22 +530,23 @@ def clone_module_simple(module, target_device):
|
||||
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',
|
||||
'_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'
|
||||
'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():
|
||||
@@ -570,7 +596,8 @@ def safe_model_clone(source_model, target_device, disable_flash=False):
|
||||
clear_flux_caches(source_model)
|
||||
return source_model
|
||||
|
||||
if src_device.type != 'cpu':
|
||||
# Move to CPU first if on GPU
|
||||
if src_device.type == 'cuda':
|
||||
model_cpu = source_model.cpu()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
@@ -627,17 +654,16 @@ def safe_model_clone(source_model, target_device, disable_flash=False):
|
||||
if is_float8_dtype(param_data.dtype) and not supports_fp8:
|
||||
param_data = param_data.half()
|
||||
target_param.data.copy_(param_data)
|
||||
|
||||
del state_dict[key]
|
||||
del param
|
||||
|
||||
if len(state_dict) % 50 == 0 and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
del state_dict[key]
|
||||
del param
|
||||
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':
|
||||
if src_device.type != 'cuda':
|
||||
del model_cpu
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
@@ -648,11 +674,9 @@ def safe_model_clone(source_model, target_device, disable_flash=False):
|
||||
except Exception as recon_error:
|
||||
print(f"[ParallelAnything] Config reconstruction failed: {recon_error}")
|
||||
print(f"[ParallelAnything] Attempting manual recursive cloning...")
|
||||
|
||||
replica = clone_module_simple(model_cpu, target_device)
|
||||
method_used = "manual_recursive_clone"
|
||||
|
||||
if src_device.type != 'cpu':
|
||||
if src_device.type != 'cuda':
|
||||
del model_cpu
|
||||
gc.collect()
|
||||
|
||||
@@ -678,7 +702,6 @@ def safe_model_clone(source_model, target_device, disable_flash=False):
|
||||
replica.gradient_checkpointing = False
|
||||
if hasattr(replica, '_gradient_checkpointing_func'):
|
||||
replica._gradient_checkpointing_func = None
|
||||
|
||||
if hasattr(replica, 'hooks'):
|
||||
replica.hooks = []
|
||||
if hasattr(replica, '_hf_hook'):
|
||||
@@ -703,8 +726,9 @@ def get_free_vram(device_name):
|
||||
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))
|
||||
free_memory = (
|
||||
torch.cuda.get_device_properties(idx).total_memory - torch.cuda.memory_allocated(idx)
|
||||
)
|
||||
return free_memory / (1024 ** 2)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -739,7 +763,6 @@ def auto_split_batch(batch_size, devices, weights):
|
||||
|
||||
split_sizes = [max(1, int(batch_size * w)) for w in adjusted_weights]
|
||||
split_sizes[-1] = batch_size - sum(split_sizes[:-1])
|
||||
|
||||
return split_sizes
|
||||
|
||||
class ParallelDevice:
|
||||
@@ -761,7 +784,7 @@ class ParallelDevice:
|
||||
except ImportError:
|
||||
pass
|
||||
return devices
|
||||
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
available = s.get_available_devices()
|
||||
@@ -786,13 +809,13 @@ class ParallelDevice:
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
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 = []
|
||||
@@ -821,7 +844,7 @@ class ParallelDeviceList:
|
||||
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()
|
||||
@@ -840,18 +863,15 @@ class ParallelDeviceList:
|
||||
"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):
|
||||
|
||||
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)]
|
||||
|
||||
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({
|
||||
@@ -859,7 +879,6 @@ class ParallelDeviceList:
|
||||
"percentage": float(pct),
|
||||
"weight": float(pct) / 100.0
|
||||
})
|
||||
|
||||
return (chain,)
|
||||
|
||||
class ParallelAnything:
|
||||
@@ -889,19 +908,56 @@ class ParallelAnything:
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
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):
|
||||
|
||||
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,)
|
||||
|
||||
target_model = model.model.diffusion_model
|
||||
# Extract the diffusion model from ComfyUI model wrapper
|
||||
if hasattr(model, 'model') and hasattr(model.model, 'diffusion_model'):
|
||||
target_model = model.model.diffusion_model
|
||||
model_wrapper = model
|
||||
elif hasattr(model, 'diffusion_model'):
|
||||
target_model = model.diffusion_model
|
||||
model_wrapper = None
|
||||
else:
|
||||
target_model = model
|
||||
model_wrapper = None
|
||||
|
||||
# Check if ComfyUI model has patches (LoRAs) that need to be applied
|
||||
# ComfyUI stores patches in model.patches (list of tuples) or model.model_patcher.patches
|
||||
has_lora_patches = False
|
||||
if model_wrapper is not None:
|
||||
# Check various locations where ComfyUI might store patches
|
||||
if hasattr(model_wrapper, 'patches') and len(model_wrapper.patches) > 0:
|
||||
has_lora_patches = True
|
||||
print(f"[ParallelAnything] Found {len(model_wrapper.patches)} patches in model.patches")
|
||||
elif hasattr(model_wrapper, 'model_patcher') and hasattr(model_wrapper.model_patcher, 'patches'):
|
||||
if len(model_wrapper.model_patcher.patches) > 0:
|
||||
has_lora_patches = True
|
||||
print(f"[ParallelAnything] Found {len(model_wrapper.model_patcher.patches)} patches in model_patcher.patches")
|
||||
|
||||
# Apply patches if found - this modifies the diffusion_model weights in-place
|
||||
if has_lora_patches:
|
||||
try:
|
||||
print("[ParallelAnything] Applying LoRA/Custom patches before cloning...")
|
||||
# Apply patches to the model weights
|
||||
if hasattr(model_wrapper, 'patch_model'):
|
||||
model_wrapper.patch_model(device_to=comfy.model_management.get_torch_device())
|
||||
print("[ParallelAnything] Patches applied successfully")
|
||||
elif hasattr(model_wrapper, 'model_patcher') and hasattr(model_wrapper.model_patcher, 'patch_model'):
|
||||
model_wrapper.model_patcher.patch_model(device_to=comfy.model_management.get_torch_device())
|
||||
print("[ParallelAnything] Patches applied via model_patcher")
|
||||
except Exception as e:
|
||||
print(f"[ParallelAnything] Warning: Could not apply patches: {e}")
|
||||
has_lora_patches = False
|
||||
|
||||
# Cleanup previous parallel setup if exists
|
||||
if hasattr(target_model, "_true_parallel_active"):
|
||||
print("[ParallelAnything] Cleaning up previous parallel setup...")
|
||||
cleanup_parallel_model(target_model)
|
||||
@@ -965,7 +1021,12 @@ class ParallelAnything:
|
||||
dev = torch.device(dev_name)
|
||||
need_safe = dev_name in devices_needing_safe_attention
|
||||
|
||||
if dev_name == original_device_str and not needs_cpu_transition:
|
||||
# Determine if this is the original device
|
||||
is_original_device = (dev_name == original_device_str and not needs_cpu_transition)
|
||||
|
||||
# CRITICAL FIX: If we have LoRA patches, we MUST clone even the original device
|
||||
# because ComfyUI will manage the original model's weights separately
|
||||
if is_original_device and not has_lora_patches:
|
||||
print(f"[ParallelAnything] Using original model for {dev_name} (skipping clone)")
|
||||
replicas[dev_name] = target_model
|
||||
if dev.type == 'cuda':
|
||||
@@ -975,7 +1036,10 @@ class ParallelAnything:
|
||||
continue
|
||||
|
||||
try:
|
||||
print(f"[ParallelAnything] Cloning to {dev_name} ({i+1}/{len(device_names)})...")
|
||||
if is_original_device and has_lora_patches:
|
||||
print(f"[ParallelAnything] Cloning to {dev_name} (preserving LoRA weights)...")
|
||||
else:
|
||||
print(f"[ParallelAnything] Cloning to {dev_name} ({i+1}/{len(device_names)})...")
|
||||
|
||||
if dev.type == 'cuda':
|
||||
idx = dev.index if dev.index is not None else int(dev_name.split(":")[-1])
|
||||
@@ -985,15 +1049,20 @@ class ParallelAnything:
|
||||
|
||||
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 ""))
|
||||
status_msg = f"[ParallelAnything] ✓ {dev_name}"
|
||||
if need_safe:
|
||||
status_msg += " (Safe mode)"
|
||||
if has_lora_patches:
|
||||
status_msg += " (LoRA active)"
|
||||
print(status_msg)
|
||||
|
||||
successful_devices.append(dev_name)
|
||||
successful_weights.append(weights[i])
|
||||
|
||||
aggressive_cleanup()
|
||||
|
||||
except RuntimeError as e:
|
||||
@@ -1012,6 +1081,7 @@ class ParallelAnything:
|
||||
total_weight = sum(successful_weights)
|
||||
weights = [w/total_weight for w in successful_weights]
|
||||
|
||||
# Restore original to its device if we moved it to CPU 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}")
|
||||
@@ -1032,24 +1102,21 @@ class ParallelAnything:
|
||||
pass
|
||||
aggressive_cleanup()
|
||||
return (model,)
|
||||
|
||||
print("[ParallelAnything] Configuring Model/Pipeline Parallelism for Batch=1...")
|
||||
|
||||
print("[ParallelAnything] Configuring Model/Pipeline Parallelism for Batch=1...")
|
||||
lead_device_name = device_names[0]
|
||||
lead_replica = replicas[lead_device_name]
|
||||
|
||||
block_lists = ['double_blocks', 'single_blocks', 'transformer_blocks', 'layers']
|
||||
|
||||
for list_name in block_lists:
|
||||
if hasattr(lead_replica, list_name):
|
||||
local_blocks = getattr(lead_replica, list_name)
|
||||
if not isinstance(local_blocks, nn.ModuleList):
|
||||
continue
|
||||
|
||||
num_blocks = len(local_blocks)
|
||||
if num_blocks == 0:
|
||||
continue
|
||||
|
||||
|
||||
print(f"[ParallelAnything] Configuring {list_name} ({num_blocks} blocks) for pipeline execution...")
|
||||
|
||||
device_assignments = []
|
||||
@@ -1058,7 +1125,6 @@ class ParallelAnything:
|
||||
count = int(round(weight * num_blocks))
|
||||
if i == len(weights) - 1:
|
||||
count = num_blocks - current_block
|
||||
|
||||
target_dev = torch.device(device_names[i])
|
||||
for _ in range(count):
|
||||
if current_block < num_blocks:
|
||||
@@ -1067,7 +1133,6 @@ class ParallelAnything:
|
||||
|
||||
for idx in range(num_blocks):
|
||||
assigned_dev = device_assignments[idx]
|
||||
|
||||
peers = {}
|
||||
for d_name, r in replicas.items():
|
||||
if hasattr(r, list_name):
|
||||
@@ -1084,9 +1149,8 @@ class ParallelAnything:
|
||||
is_last_block=is_last,
|
||||
lead_device=torch.device(lead_device_name)
|
||||
)
|
||||
|
||||
local_blocks[idx] = wrapper
|
||||
|
||||
|
||||
target_model._parallel_purge_cache = purge_cache
|
||||
target_model._parallel_purge_models = purge_models
|
||||
|
||||
@@ -1182,6 +1246,7 @@ class ParallelAnything:
|
||||
print(f"[ParallelAnything] Error: {e}")
|
||||
raise
|
||||
|
||||
# Pipeline mode for batch_size == 1
|
||||
if batch_size == 1 and workload_split:
|
||||
set_pipeline_mode(True)
|
||||
try:
|
||||
@@ -1193,6 +1258,7 @@ class ParallelAnything:
|
||||
finally:
|
||||
set_pipeline_mode(False)
|
||||
|
||||
# Single device fallback if batch too small or workload_split disabled
|
||||
if batch_size < len(devices_ref) or not workload_split:
|
||||
set_pipeline_mode(False)
|
||||
with torch.no_grad():
|
||||
@@ -1202,6 +1268,7 @@ class ParallelAnything:
|
||||
else:
|
||||
return first_replica(x, timesteps, context=context, **kwargs)
|
||||
|
||||
# Data parallelism for batch_size > 1
|
||||
if auto_balance_ref:
|
||||
split_sizes = auto_split_batch(batch_size, list(devices_ref), list(weights_ref))
|
||||
else:
|
||||
@@ -1234,6 +1301,7 @@ class ParallelAnything:
|
||||
|
||||
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:
|
||||
@@ -1293,6 +1361,7 @@ class ParallelAnything:
|
||||
|
||||
out = move_to_device(out, lead_device, non_blocking=False)
|
||||
return task_idx, out
|
||||
|
||||
except Exception as e:
|
||||
return task_idx, e
|
||||
|
||||
@@ -1334,7 +1403,6 @@ class ParallelAnything:
|
||||
|
||||
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
|
||||
@@ -1344,9 +1412,16 @@ class ParallelAnything:
|
||||
|
||||
weakref.finalize(model, cleanup_parallel_model, weakref.ref(target_model))
|
||||
|
||||
model.load_device = lead_device
|
||||
# Update ComfyUI's model device tracking to point to lead device
|
||||
if hasattr(model, 'load_device'):
|
||||
model.load_device = lead_device
|
||||
elif hasattr(model, 'model') and hasattr(model.model, 'load_device'):
|
||||
model.model.load_device = lead_device
|
||||
|
||||
print(f"[ParallelAnything] Parallel setup complete. Devices: {device_names}")
|
||||
if has_lora_patches:
|
||||
print("[ParallelAnything] LoRA weights synchronized across all devices")
|
||||
|
||||
return (model,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
|
||||
Reference in New Issue
Block a user