SuperUltimateVaceTools - fixed tilling, MultiGPU - fixed torch
This commit is contained in:
Cristian Predut
@@ -0,0 +1,3 @@
|
||||
https://arxiv.org/html/2504.14899v1
|
||||
|
||||
https://github.com/alibaba-damo-academy/Uni3C
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
https://github.com/bbaudio-2025/ComfyUI-SuperUltimateVaceTools
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,774 @@
|
||||
"""
|
||||
DisTorch Safetensor Memory Management Module
|
||||
Contains all safetensor related code for distributed memory management
|
||||
"""
|
||||
|
||||
import sys
|
||||
import torch
|
||||
import logging
|
||||
import hashlib
|
||||
import re
|
||||
import gc
|
||||
|
||||
logger = logging.getLogger("MultiGPU")
|
||||
import copy
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
import comfy.model_management as mm
|
||||
import comfy.model_patcher
|
||||
from .device_utils import get_device_list, soft_empty_cache_multigpu
|
||||
from .model_management_mgpu import multigpu_memory_log, force_full_system_cleanup
|
||||
|
||||
|
||||
safetensor_allocation_store = {}
|
||||
safetensor_settings_store = {}
|
||||
|
||||
|
||||
def _block_fields(block):
|
||||
"""Safely unpack a block tuple regardless of extra fields."""
|
||||
module_size = block[0] if len(block) > 0 else 0
|
||||
module_name = block[1] if len(block) > 1 else ""
|
||||
module_object = block[2] if len(block) > 2 else None
|
||||
params = block[3] if len(block) > 3 else None
|
||||
return module_size, module_name, module_object, params
|
||||
|
||||
|
||||
def create_safetensor_model_hash(model, caller):
|
||||
"""Create a unique hash for a safetensor model to track allocations"""
|
||||
if hasattr(model, 'model'):
|
||||
# For ModelPatcher objects
|
||||
actual_model = model.model
|
||||
model_type = type(actual_model).__name__
|
||||
# Use ComfyUI's model_size if available
|
||||
if hasattr(model, 'model_size'):
|
||||
model_size = model.model_size()
|
||||
else:
|
||||
model_size = sum(p.numel() * p.element_size() for p in actual_model.parameters())
|
||||
if hasattr(model, 'model_state_dict'):
|
||||
first_layers = str(list(model.model_state_dict().keys())[:3])
|
||||
else:
|
||||
first_layers = str(list(actual_model.state_dict().keys())[:3])
|
||||
else:
|
||||
# Direct model
|
||||
model_type = type(model).__name__
|
||||
model_size = sum(p.numel() * p.element_size() for p in model.parameters())
|
||||
first_layers = str(list(model.state_dict().keys())[:3])
|
||||
|
||||
identifier = f"{model_type}_{model_size}_{first_layers}"
|
||||
final_hash = hashlib.sha256(identifier.encode()).hexdigest()
|
||||
|
||||
# DEBUG STATEMENT - ALWAYS LOG THE HASH
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Created hash for {caller}: {final_hash[:8]}...")
|
||||
return final_hash
|
||||
|
||||
def register_patched_safetensor_modelpatcher():
|
||||
"""Register and patch the ModelPatcher for distributed safetensor loading"""
|
||||
from comfy.model_patcher import wipe_lowvram_weight, move_weight_functions
|
||||
# Patch ComfyUI's ModelPatcher
|
||||
if not hasattr(comfy.model_patcher.ModelPatcher, '_distorch_patched'):
|
||||
|
||||
|
||||
# PATCH load_models_gpu with correct memory calculations per model flags
|
||||
original_load_models_gpu = mm.load_models_gpu
|
||||
|
||||
def patched_load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False):
|
||||
from comfy.model_management import cleanup_models_gc, get_free_memory, free_memory, current_loaded_models
|
||||
from comfy.model_management import VRAMState, vram_state, lowvram_available, MIN_WEIGHT_MEMORY_RATIO
|
||||
from comfy.model_management import minimum_inference_memory, extra_reserved_memory, is_device_cpu
|
||||
|
||||
multigpu_memory_log("load_models_gpu_top_level", "start")
|
||||
|
||||
cleanup_models_gc()
|
||||
|
||||
inference_memory = minimum_inference_memory()
|
||||
extra_reserved_mem = extra_reserved_memory()
|
||||
memory_required_total = memory_required + extra_reserved_mem
|
||||
extra_mem = max(inference_memory, memory_required_total)
|
||||
if minimum_memory_required is None:
|
||||
minimum_memory_required = extra_mem
|
||||
else:
|
||||
minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_mem)
|
||||
|
||||
models_temp = set()
|
||||
for m in models:
|
||||
models_temp.add(m)
|
||||
for mm_patch in m.model_patches_models():
|
||||
models_temp.add(mm_patch)
|
||||
|
||||
models = models_temp
|
||||
|
||||
models_to_load = []
|
||||
|
||||
for x in models:
|
||||
loaded_model = mm.LoadedModel(x)
|
||||
try:
|
||||
loaded_model_index = current_loaded_models.index(loaded_model)
|
||||
except:
|
||||
loaded_model_index = None
|
||||
|
||||
if loaded_model_index is not None:
|
||||
loaded = current_loaded_models[loaded_model_index]
|
||||
loaded.currently_used = True
|
||||
models_to_load.append(loaded)
|
||||
else:
|
||||
if hasattr(x, "model"):
|
||||
logging.info(f"Requested to load {x.model.__class__.__name__}")
|
||||
models_to_load.append(loaded_model)
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
to_unload = []
|
||||
for i in range(len(current_loaded_models)):
|
||||
if loaded_model.model.is_clone(current_loaded_models[i].model):
|
||||
to_unload = [i] + to_unload
|
||||
for i in to_unload:
|
||||
model_to_unload = current_loaded_models.pop(i)
|
||||
model_to_unload.model.detach(unpatch_all=False)
|
||||
model_to_unload.model_finalizer.detach()
|
||||
|
||||
# DisTorch Processing
|
||||
total_memory_required = {}
|
||||
eject_device = None
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
device = loaded_model.device
|
||||
base_memory = loaded_model.model_memory_required(device)
|
||||
|
||||
# Check DisTorch flags
|
||||
is_distorch = hasattr(loaded_model.model.model, '_mgpu_virtual_vram_gb')
|
||||
has_eject = hasattr(loaded_model.model.model, '_mgpu_eject_models')
|
||||
|
||||
if has_eject:
|
||||
eject_device = device
|
||||
logger.mgpu_mm_log("DisTorch eject_models=True, is_distorch=True - MAX memory eviction")
|
||||
|
||||
if is_distorch:
|
||||
# is_distorch=True: use compute device allocation size
|
||||
virtual_vram_gb = loaded_model.model.model._mgpu_virtual_vram_gb
|
||||
virtual_vram_bytes = virtual_vram_gb * (1024**3)
|
||||
adjusted_memory = max(0, base_memory - virtual_vram_bytes)
|
||||
total_memory_required[device] = total_memory_required.get(device, 0) + adjusted_memory
|
||||
logger.mgpu_mm_log(f"DisTorch is_distorch=True, model adjusted {(base_memory - virtual_vram_bytes)/(1024**3):.2f}GB for device {device}")
|
||||
else:
|
||||
# is_distorch=False: use full model size
|
||||
total_memory_required[device] = total_memory_required.get(device, 0) + base_memory
|
||||
logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] Standard model {(base_memory)/(1024**3):.2f}GB for device {device}")
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
requested_mem = total_memory_required[device] * 1.1 + extra_mem
|
||||
logger.mgpu_mm_log(f"[FREE_MEMORY_CALL] Device {device}: requesting {requested_mem/(1024**3):.2f}GB = {total_memory_required[device]/(1024**3):.2f}GB * 1.1 + {extra_mem/(1024**3):.2f}GB inference")
|
||||
|
||||
|
||||
multigpu_memory_log("free_memory", "pre")
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
if device == eject_device:
|
||||
total_device_memory = mm.get_total_memory(device)
|
||||
logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] eject_models=1, is_distorch=1 → using MAX memory ({total_device_memory/(1024**3):.2f}GB) for eviction")
|
||||
free_memory(total_device_memory,device)
|
||||
else:
|
||||
logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] eject_models=0, using Comfy Core Computed memory ({(total_memory_required[device] * 1.1 + extra_mem)/(1024**3):.2f}GB) for eviction")
|
||||
free_memory(total_memory_required[device] * 1.1 + extra_mem, device)
|
||||
|
||||
multigpu_memory_log("free_memory/minimum_memory_required", "post/pre")
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
free_mem = get_free_memory(device)
|
||||
free_mem_gb = free_mem / (1024**3)
|
||||
min_required_gb = minimum_memory_required / (1024**3)
|
||||
logger.mgpu_mm_log(f"[MIN_MEMORY_CHECK] Device {device}: free={free_mem_gb:.2f}GB, required={min_required_gb:.2f}GB, will_evict={free_mem < minimum_memory_required}")
|
||||
|
||||
if free_mem < minimum_memory_required:
|
||||
models_l = free_memory(minimum_memory_required, device)
|
||||
logger.mgpu_mm_log(f"[EVICTION] Device {device}: unloaded {len(models_l)} models due to insufficient memory")
|
||||
logging.info("{} models unloaded.".format(len(models_l)))
|
||||
|
||||
multigpu_memory_log("minimum_memory_required", "post")
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
model = loaded_model.model
|
||||
torch_dev = model.load_device
|
||||
if is_device_cpu(torch_dev):
|
||||
vram_set_state = VRAMState.DISABLED
|
||||
else:
|
||||
vram_set_state = vram_state
|
||||
lowvram_model_memory = 0
|
||||
if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load:
|
||||
loaded_memory = loaded_model.model_loaded_memory()
|
||||
current_free_mem = get_free_memory(torch_dev) + loaded_memory
|
||||
|
||||
lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory()))
|
||||
lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory)
|
||||
|
||||
if vram_set_state == VRAMState.NO_VRAM:
|
||||
lowvram_model_memory = 0.1
|
||||
|
||||
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
current_loaded_models.insert(0, loaded_model)
|
||||
|
||||
# Replace the module function
|
||||
mm.load_models_gpu = patched_load_models_gpu
|
||||
|
||||
original_partially_load = comfy.model_patcher.ModelPatcher.partially_load
|
||||
|
||||
def new_partially_load(self, device_to, extra_memory=0, full_load=False, force_patch_weights=False, **kwargs):
|
||||
"""Override to use our static device assignments"""
|
||||
global safetensor_allocation_store
|
||||
|
||||
debug_hash = create_safetensor_model_hash(self, "partial_load")
|
||||
multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "pre-load")
|
||||
allocations = safetensor_allocation_store.get(debug_hash)
|
||||
|
||||
# Set default precision flag before checking
|
||||
if not hasattr(self.model, '_distorch_high_precision_loras'):
|
||||
self.model._distorch_high_precision_loras = True
|
||||
|
||||
if not allocations:
|
||||
result = original_partially_load(self, device_to, extra_memory, force_patch_weights)
|
||||
multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "post-load")
|
||||
if hasattr(self, '_distorch_block_assignments'):
|
||||
del self._distorch_block_assignments
|
||||
return result
|
||||
|
||||
if not hasattr(self.model, 'current_weight_patches_uuid'):
|
||||
self.model.current_weight_patches_uuid = None
|
||||
|
||||
unpatch_weights = self.model.current_weight_patches_uuid is not None and (self.model.current_weight_patches_uuid != self.patches_uuid or force_patch_weights)
|
||||
|
||||
if unpatch_weights:
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Patches changed or forced. Unpatching model.")
|
||||
self.unpatch_model(self.offload_device, unpatch_weights=True)
|
||||
|
||||
self.patch_model(load_weights=False)
|
||||
|
||||
mem_counter = 0
|
||||
|
||||
is_clip_model = getattr(self, 'is_clip', False)
|
||||
device_assignments = analyze_safetensor_loading(self, allocations, is_clip=is_clip_model)
|
||||
|
||||
model_original_dtype = comfy.utils.weight_dtype(self.model.state_dict())
|
||||
high_precision_loras = getattr(self.model, "_distorch_high_precision_loras", True)
|
||||
loading = self._load_list()
|
||||
loading.sort(reverse=True)
|
||||
for block in loading:
|
||||
module_size, module_name, module_object, params = _block_fields(block)
|
||||
if not unpatch_weights and hasattr(module_object, "comfy_patched_weights") and module_object.comfy_patched_weights == True:
|
||||
block_target_device = device_assignments['block_assignments'].get(module_name, device_to)
|
||||
current_module_device = None
|
||||
try:
|
||||
if any(p.numel() > 0 for p in module_object.parameters(recurse=False)):
|
||||
current_module_device = next(module_object.parameters(recurse=False)).device
|
||||
except StopIteration:
|
||||
pass
|
||||
|
||||
if current_module_device is not None and str(current_module_device) != str(block_target_device):
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Moving already patched {module_name} to {block_target_device}")
|
||||
module_object.to(block_target_device)
|
||||
|
||||
mem_counter += module_size
|
||||
continue
|
||||
|
||||
# Step 1: Write block/tensor to compute device first
|
||||
module_object.to(device_to)
|
||||
|
||||
# Step 2: Apply LoRa patches while on compute device
|
||||
weight_key = "{}.weight".format(module_name)
|
||||
bias_key = "{}.bias".format(module_name)
|
||||
|
||||
if weight_key in self.patches:
|
||||
self.patch_weight_to_device(weight_key, device_to=device_to)
|
||||
if weight_key in self.weight_wrapper_patches:
|
||||
module_object.weight_function.extend(self.weight_wrapper_patches[weight_key])
|
||||
|
||||
if bias_key in self.patches:
|
||||
self.patch_weight_to_device(bias_key, device_to=device_to)
|
||||
if bias_key in self.weight_wrapper_patches:
|
||||
module_object.bias_function.extend(self.weight_wrapper_patches[bias_key])
|
||||
|
||||
# Step 3: FP8 casting for CPU storage (if enabled)
|
||||
block_target_device = device_assignments['block_assignments'].get(module_name, device_to)
|
||||
has_patches = weight_key in self.patches or bias_key in self.patches
|
||||
|
||||
if not high_precision_loras and block_target_device == "cpu" and has_patches and model_original_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
|
||||
for param_name, param in module_object.named_parameters():
|
||||
if param.dtype.is_floating_point:
|
||||
cast_data = comfy.float.stochastic_rounding(param.data, torch.float8_e4m3fn)
|
||||
new_param = torch.nn.Parameter(cast_data.to(torch.float8_e4m3fn))
|
||||
new_param.requires_grad = param.requires_grad
|
||||
setattr(module_object, param_name, new_param)
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Cast {module_name}.{param_name} to FP8 for CPU storage")
|
||||
|
||||
# Step 4: Move to ultimate destination based on DisTorch assignment
|
||||
if block_target_device != device_to:
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Moving {module_name} from {device_to} to {block_target_device}")
|
||||
module_object.to(block_target_device)
|
||||
module_object.comfy_cast_weights = True
|
||||
|
||||
# Mark as patched and update memory counter
|
||||
module_object.comfy_patched_weights = True
|
||||
mem_counter += module_size
|
||||
|
||||
self.model.current_weight_patches_uuid = self.patches_uuid
|
||||
|
||||
logger.info("[MultiGPU DisTorch V2] DisTorch loading completed.")
|
||||
logger.info(f"[MultiGPU DisTorch V2] Total memory: {mem_counter / (1024 * 1024):.2f}MB")
|
||||
multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "post-load")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
comfy.model_patcher.ModelPatcher.partially_load = new_partially_load
|
||||
comfy.model_patcher.ModelPatcher._distorch_patched = True
|
||||
logger.info("[MultiGPU Core Patching] Successfully patched ModelPatcher.partially_load")
|
||||
|
||||
def _extract_clip_head_blocks(raw_block_list, compute_device):
|
||||
"""Identify and pre-assign CLIP head blocks to compute device returning head_blocks, distributable_blocks, block_assignments, and head_memory."""
|
||||
head_keywords = ['embed', 'wte', 'wpe', 'token_embedding', 'position_embedding']
|
||||
head_blocks = []
|
||||
distributable_blocks = []
|
||||
head_memory = 0
|
||||
block_assignments = {}
|
||||
|
||||
for block in raw_block_list:
|
||||
module_size, module_name, module_object, params = _block_fields(block)
|
||||
if any(kw in module_name.lower() for kw in head_keywords):
|
||||
head_blocks.append((module_size, module_name, module_object, params))
|
||||
block_assignments[module_name] = compute_device
|
||||
head_memory += module_size
|
||||
else:
|
||||
distributable_blocks.append((module_size, module_name, module_object, params))
|
||||
|
||||
return head_blocks, distributable_blocks, block_assignments, head_memory
|
||||
|
||||
def analyze_safetensor_loading(model_patcher, allocations_string, is_clip=False):
|
||||
"""
|
||||
Analyze and distribute safetensor model blocks across devices.
|
||||
Supports CLIP head preservation when is_clip=True.
|
||||
"""
|
||||
DEVICE_RATIOS_DISTORCH = {}
|
||||
device_table = {}
|
||||
distorch_alloc = allocations_string
|
||||
virtual_vram_gb = 0.0
|
||||
|
||||
distorch_alloc, virtual_vram_str = allocations_string.split('#')
|
||||
|
||||
compute_device = virtual_vram_str.split(';')[0]
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Compute Device: {compute_device}")
|
||||
|
||||
if not distorch_alloc:
|
||||
mode = "fraction"
|
||||
distorch_alloc = calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str)
|
||||
|
||||
elif any(c in distorch_alloc.lower() for c in ['g', 'm', 'k', 'b']):
|
||||
mode = "byte"
|
||||
distorch_alloc = calculate_fraction_from_byte_expert_string(model_patcher, distorch_alloc)
|
||||
elif "%" in distorch_alloc:
|
||||
mode = "ratio"
|
||||
distorch_alloc = calculate_fraction_from_ratio_expert_string(model_patcher, distorch_alloc)
|
||||
|
||||
all_devices = get_device_list()
|
||||
present_devices = {item.split(',')[0] for item in distorch_alloc.split(';') if ',' in item}
|
||||
for device in all_devices:
|
||||
if device not in present_devices:
|
||||
distorch_alloc += f";{device},0.0"
|
||||
|
||||
eq_line = "=" * 50
|
||||
dash_line = "-" * 50
|
||||
fmt_assign = "{:<18}{:>7}{:>14}{:>10}"
|
||||
|
||||
logger.info(eq_line)
|
||||
logger.info(f"[MultiGPU DisTorch V2] Final Allocation String:\n{distorch_alloc}")
|
||||
|
||||
for allocation in distorch_alloc.split(';'):
|
||||
if ',' not in allocation:
|
||||
continue
|
||||
dev_name, fraction = allocation.split(',')
|
||||
fraction = float(fraction)
|
||||
total_mem_bytes = mm.get_total_memory(torch.device(dev_name))
|
||||
alloc_gb = (total_mem_bytes * fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[dev_name] = alloc_gb
|
||||
device_table[dev_name] = {
|
||||
"fraction": fraction,
|
||||
"total_gb": total_mem_bytes / (1024**3),
|
||||
"alloc_gb": alloc_gb
|
||||
}
|
||||
|
||||
logger.info(eq_line)
|
||||
logger.info(" DisTorch2 Model Device Allocations")
|
||||
logger.info(eq_line)
|
||||
|
||||
fmt_rosetta = "{:<8}{:>9}{:>9}{:>11}{:>10}"
|
||||
logger.info(fmt_rosetta.format("Device", "VRAM GB", "Dev %", "Model GB", "Dist %"))
|
||||
logger.info(dash_line)
|
||||
|
||||
sorted_devices = sorted(device_table.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
total_allocated_model_bytes = sum(d["alloc_gb"] * (1024**3) for d in device_table.values())
|
||||
|
||||
for dev in sorted_devices:
|
||||
total_dev_gb = device_table[dev]["total_gb"]
|
||||
alloc_fraction = device_table[dev]["fraction"]
|
||||
alloc_gb = device_table[dev]["alloc_gb"]
|
||||
|
||||
dist_ratio_percent = (alloc_gb * (1024**3) / total_allocated_model_bytes) * 100 if total_allocated_model_bytes > 0 else 0
|
||||
|
||||
logger.info(fmt_rosetta.format(
|
||||
dev,
|
||||
f"{total_dev_gb:.2f}",
|
||||
f"{alloc_fraction*100:.1f}%",
|
||||
f"{alloc_gb:.2f}",
|
||||
f"{dist_ratio_percent:.1f}%"
|
||||
))
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
block_summary = {}
|
||||
block_list = []
|
||||
memory_by_type = defaultdict(int)
|
||||
total_memory = 0
|
||||
|
||||
raw_block_list = model_patcher._load_list()
|
||||
total_memory = sum(_block_fields(block)[0] for block in raw_block_list)
|
||||
|
||||
MIN_BLOCK_THRESHOLD = total_memory * 0.0001
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Total model memory: {total_memory} bytes")
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Tiny block threshold (0.01%): {MIN_BLOCK_THRESHOLD} bytes")
|
||||
|
||||
# CLIP-specific: Extract head blocks and get pre-assignments
|
||||
head_memory = 0
|
||||
block_assignments = {}
|
||||
if is_clip:
|
||||
head_blocks, distributable_raw, block_assignments, head_memory = \
|
||||
_extract_clip_head_blocks(raw_block_list, compute_device)
|
||||
logger.info(f"[MultiGPU DisTorch V2 CLIP] Preserving {len(head_blocks)} head layer(s) ({head_memory/(1024**2):.2f} MB) on compute device: {compute_device}")
|
||||
else:
|
||||
distributable_raw = raw_block_list
|
||||
|
||||
# Build all_blocks list for summary (using full raw_block_list)
|
||||
all_blocks = []
|
||||
for block in raw_block_list:
|
||||
module_size, module_name, module_object, params = _block_fields(block)
|
||||
block_type = type(module_object).__name__
|
||||
# Populate summary dictionaries
|
||||
block_summary[block_type] = block_summary.get(block_type, 0) + 1
|
||||
memory_by_type[block_type] += module_size
|
||||
all_blocks.append((module_name, module_object, block_type, module_size))
|
||||
|
||||
# Use distributable blocks for actual allocation (for CLIP, this excludes heads)
|
||||
distributable_all_blocks = []
|
||||
for block in distributable_raw:
|
||||
module_size, module_name, module_object, params = _block_fields(block)
|
||||
distributable_all_blocks.append((module_name, module_object, type(module_object).__name__, module_size))
|
||||
|
||||
block_list = [b for b in distributable_all_blocks if b[3] >= MIN_BLOCK_THRESHOLD]
|
||||
tiny_block_list = [b for b in distributable_all_blocks if b[3] < MIN_BLOCK_THRESHOLD]
|
||||
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Total blocks: {len(all_blocks)}")
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Distributable blocks: {len(block_list)}")
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Tiny blocks (<0.01%): {len(tiny_block_list)}")
|
||||
|
||||
logger.info(" DisTorch2 Model Layer Distribution")
|
||||
logger.info(dash_line)
|
||||
fmt_layer = "{:<18}{:>7}{:>14}{:>10}"
|
||||
logger.info(fmt_layer.format("Layer Type", "Layers", "Memory (MB)", "% Total"))
|
||||
logger.info(dash_line)
|
||||
|
||||
for layer_type, count in block_summary.items():
|
||||
mem_mb = memory_by_type[layer_type] / (1024 * 1024)
|
||||
mem_percent = (memory_by_type[layer_type] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logger.info(fmt_layer.format(layer_type[:18], str(count), f"{mem_mb:.2f}", f"{mem_percent:.1f}%"))
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
# Distribute blocks sequentially from the tail of the model
|
||||
|
||||
device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()}
|
||||
block_assignments = {}
|
||||
|
||||
# Create a memory quota for each donor device based on its calculated allocation.
|
||||
donor_devices = [d for d in sorted_devices]
|
||||
donor_quotas = {
|
||||
dev: device_table[dev]["alloc_gb"] * (1024**3)
|
||||
for dev in donor_devices
|
||||
}
|
||||
|
||||
# CLIP-specific: Adjust compute_device quota to account for locked head blocks
|
||||
if is_clip and compute_device in donor_quotas and head_memory > 0:
|
||||
donor_quotas[compute_device] = max(0, donor_quotas[compute_device] - head_memory)
|
||||
logger.debug(f"[MultiGPU DisTorch V2 CLIP] Adjusted {compute_device} quota by -{head_memory/(1024**2):.2f} MB for head preservation")
|
||||
|
||||
# Iterate from the TAIL of the model, assigning blocks to donors until their quotas are filled.
|
||||
for block_name, module, block_type, block_memory in reversed(block_list):
|
||||
assigned_to_donor = False
|
||||
for donor in donor_devices:
|
||||
if donor_quotas[donor] >= block_memory:
|
||||
block_assignments[block_name] = donor
|
||||
donor_quotas[donor] -= block_memory
|
||||
assigned_to_donor = True
|
||||
break # Move to the next block
|
||||
|
||||
if not assigned_to_donor: #Note - small rounding errors and tensor-fitting on devices make a block occasionally an orphan. We treat orphans the same as tiny_block_list as they are generally small rounding errors
|
||||
block_assignments[block_name] = compute_device
|
||||
|
||||
if tiny_block_list:
|
||||
for block_name, module, block_type, block_memory in tiny_block_list:
|
||||
block_assignments[block_name] = compute_device
|
||||
|
||||
# Populate device_assignments from the final block_assignments
|
||||
for block_name, device in block_assignments.items():
|
||||
# Find the block in the original list to get all its info
|
||||
for b_name, b_module, b_type, b_mem in all_blocks:
|
||||
if b_name == block_name:
|
||||
device_assignments[device].append((b_name, b_module, b_type, b_mem))
|
||||
break
|
||||
|
||||
logger.info("DisTorch2 Model Final Device/Layer Assignments")
|
||||
logger.info(dash_line)
|
||||
logger.info(fmt_assign.format("Device", "Layers", "Memory (MB)", "% Total"))
|
||||
logger.info(dash_line)
|
||||
|
||||
if tiny_block_list:
|
||||
tiny_block_memory = sum(b[3] for b in tiny_block_list)
|
||||
tiny_mem_mb = tiny_block_memory / (1024 * 1024)
|
||||
tiny_mem_percent = (tiny_block_memory / total_memory) * 100 if total_memory > 0 else 0
|
||||
device_label = f"{compute_device} (<0.01%)"
|
||||
logger.info(fmt_assign.format(device_label, str(len(tiny_block_list)), f"{tiny_mem_mb:.2f}", f"{tiny_mem_percent:.1f}%"))
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Tiny block memory breakdown: {tiny_block_memory} bytes ({tiny_mem_mb:.2f} MB), which is {tiny_mem_percent:.4f}% of total model memory.")
|
||||
|
||||
total_assigned_memory = 0
|
||||
device_memories = {}
|
||||
|
||||
for device, blocks in device_assignments.items():
|
||||
dist_blocks = [b for b in blocks if b[3] >= MIN_BLOCK_THRESHOLD]
|
||||
if not dist_blocks:
|
||||
continue
|
||||
|
||||
device_memory = sum(b[3] for b in dist_blocks)
|
||||
device_memories[device] = device_memory
|
||||
total_assigned_memory += device_memory
|
||||
|
||||
sorted_assignments = sorted(device_memories.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
for dev in sorted_assignments:
|
||||
# Get only the distributed blocks for the count
|
||||
dist_blocks = [b for b in device_assignments[dev] if b[3] >= MIN_BLOCK_THRESHOLD]
|
||||
if not dist_blocks:
|
||||
continue
|
||||
|
||||
mem_mb = device_memories[dev] / (1024 * 1024)
|
||||
mem_percent = (device_memories[dev] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logger.info(fmt_assign.format(dev, str(len(dist_blocks)), f"{mem_mb:.2f}", f"{mem_percent:.1f}%"))
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
return {
|
||||
"device_assignments": device_assignments,
|
||||
"block_assignments": block_assignments
|
||||
}
|
||||
|
||||
def parse_memory_string(mem_str):
|
||||
"""Parses a memory string (e.g., '4.0g', '512M') and returns bytes."""
|
||||
mem_str = mem_str.strip().lower()
|
||||
match = re.match(r'(\d+\.?\d*)\s*([gmkb]?)', mem_str)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid memory string format: {mem_str}")
|
||||
|
||||
val, unit = match.groups()
|
||||
val = float(val)
|
||||
|
||||
if unit == 'g':
|
||||
return val * (1024**3)
|
||||
elif unit == 'm':
|
||||
return val * (1024**2)
|
||||
elif unit == 'k':
|
||||
return val * 1024
|
||||
else: # b or no unit
|
||||
return val
|
||||
|
||||
def calculate_fraction_from_byte_expert_string(model_patcher, byte_str):
|
||||
"""Convert byte allocation string (e.g. 'cuda:1,4gb;cpu,*') to fractional VRAM allocation string respecting device order and byte quotas."""
|
||||
raw_block_list = model_patcher._load_list()
|
||||
total_model_memory = sum(_block_fields(block)[0] for block in raw_block_list)
|
||||
remaining_model_bytes = total_model_memory
|
||||
|
||||
# Use a list of tuples to preserve the user-defined order
|
||||
parsed_allocations = []
|
||||
wildcard_device = "cpu" # Default wildcard device
|
||||
|
||||
for allocation in byte_str.split(';'):
|
||||
if ',' not in allocation:
|
||||
continue
|
||||
dev_name, val_str = allocation.split(',', 1)
|
||||
is_wildcard = '*' in val_str
|
||||
|
||||
if is_wildcard:
|
||||
wildcard_device = dev_name
|
||||
# Don't add wildcard to the priority list yet
|
||||
else:
|
||||
byte_val = parse_memory_string(val_str)
|
||||
parsed_allocations.append({'device': dev_name, 'bytes': byte_val})
|
||||
|
||||
final_byte_allocations = defaultdict(int)
|
||||
|
||||
# Process devices with specific byte allocations first, in order
|
||||
for alloc in parsed_allocations:
|
||||
dev = alloc['device']
|
||||
requested_bytes = alloc['bytes']
|
||||
|
||||
# Determine the actual bytes to allocate to this device
|
||||
bytes_to_assign = min(requested_bytes, remaining_model_bytes)
|
||||
|
||||
if bytes_to_assign > 0:
|
||||
final_byte_allocations[dev] = bytes_to_assign
|
||||
remaining_model_bytes -= bytes_to_assign
|
||||
logger.info(f"[MultiGPU DisTorch V2] Assigning {bytes_to_assign / (1024**2):.2f}MB of model to {dev} (requested {requested_bytes / (1024**2):.2f}MB).")
|
||||
|
||||
if remaining_model_bytes <= 0:
|
||||
logger.info("[MultiGPU DisTorch V2] All model blocks have been allocated. Subsequent devices in the string will receive no assignment.")
|
||||
break
|
||||
|
||||
# Assign any leftover model bytes to the wildcard device
|
||||
if remaining_model_bytes > 0:
|
||||
final_byte_allocations[wildcard_device] += remaining_model_bytes
|
||||
logger.info(f"[MultiGPU DisTorch V2] Assigning remaining {remaining_model_bytes / (1024**2):.2f}MB of model to wildcard device '{wildcard_device}'.")
|
||||
|
||||
# Convert the final byte allocations to VRAM fractions
|
||||
allocation_parts = []
|
||||
for dev, bytes_alloc in final_byte_allocations.items():
|
||||
total_device_vram = mm.get_total_memory(torch.device(dev))
|
||||
if total_device_vram > 0:
|
||||
fraction = bytes_alloc / total_device_vram
|
||||
allocation_parts.append(f"{dev},{fraction:.4f}")
|
||||
|
||||
allocations_string = ";".join(allocation_parts)
|
||||
|
||||
return allocations_string
|
||||
|
||||
def calculate_fraction_from_ratio_expert_string(model_patcher, ratio_str):
|
||||
"""Convert ratio allocation string (e.g. 'cuda:0,25%;cpu,75%') describing model split to fractional VRAM allocation string."""
|
||||
raw_block_list = model_patcher._load_list()
|
||||
total_model_memory = sum(_block_fields(block)[0] for block in raw_block_list)
|
||||
|
||||
raw_ratios = {}
|
||||
for allocation in ratio_str.split(';'):
|
||||
if ',' not in allocation: continue
|
||||
dev_name, val_str = allocation.split(',', 1)
|
||||
# Assumes the value is a unitless ratio number, ignores '%' for simplicity.
|
||||
value = float(val_str.replace('%','').strip())
|
||||
raw_ratios[dev_name] = value
|
||||
|
||||
total_ratio_parts = sum(raw_ratios.values())
|
||||
allocation_parts = []
|
||||
|
||||
for dev, ratio_val in raw_ratios.items():
|
||||
bytes_of_model_for_device = (ratio_val / total_ratio_parts) * total_model_memory
|
||||
|
||||
total_vram_of_device = mm.get_total_memory(torch.device(dev))
|
||||
|
||||
if total_vram_of_device > 0:
|
||||
required_fraction = bytes_of_model_for_device / total_vram_of_device
|
||||
allocation_parts.append(f"{dev},{required_fraction:.4f}")
|
||||
|
||||
ratio_values = [str(v) for v in raw_ratios.values()]
|
||||
ratio_string = ":".join(ratio_values)
|
||||
|
||||
normalized_pcts = [(v / total_ratio_parts) * 100 for v in raw_ratios.values()]
|
||||
|
||||
put_parts = []
|
||||
for i, dev_name in enumerate(raw_ratios.keys()):
|
||||
put_parts.append(f"{int(normalized_pcts[i])}% on {dev_name}")
|
||||
|
||||
if len(put_parts) == 1:
|
||||
put_part = put_parts[0]
|
||||
elif len(put_parts) == 2:
|
||||
put_part = f"{put_parts[0]} and {put_parts[1]}"
|
||||
else:
|
||||
put_part = ", ".join(put_parts[:-1]) + f", and {put_parts[-1]}"
|
||||
|
||||
logger.info(f"[MultiGPU DisTorch V2] Ratio(%) Mode - {ratio_str} -> {ratio_string} ratio, put {put_part}")
|
||||
|
||||
allocations_string = ";".join(allocation_parts)
|
||||
|
||||
return allocations_string
|
||||
|
||||
def calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str):
|
||||
"""Calculate virtual VRAM allocation string for distributed safetensor loading"""
|
||||
recipient_device, vram_amount, donors = virtual_vram_str.split(';')
|
||||
virtual_vram_gb = float(vram_amount)
|
||||
|
||||
eq_line = "=" * 47
|
||||
dash_line = "-" * 47
|
||||
fmt_assign = "{:<8} {:<6} {:>11} {:>9} {:>9}"
|
||||
|
||||
logger.info(eq_line)
|
||||
logger.info(" DisTorch2 Model Virtual VRAM Analysis")
|
||||
logger.info(eq_line)
|
||||
logger.info(fmt_assign.format("Object", "Role", "Original(GB)", "Total(GB)", "Virt(GB)"))
|
||||
logger.info(dash_line)
|
||||
|
||||
# Calculate recipient VRAM
|
||||
recipient_vram = mm.get_total_memory(torch.device(recipient_device)) / (1024**3)
|
||||
recipient_virtual = recipient_vram + virtual_vram_gb
|
||||
|
||||
logger.info(fmt_assign.format(recipient_device, 'recip', f"{recipient_vram:.2f}GB",f"{recipient_virtual:.2f}GB", f"+{virtual_vram_gb:.2f}GB"))
|
||||
|
||||
# Handle donor devices
|
||||
ram_donors = [d for d in donors.split(',')]
|
||||
remaining_vram_needed = virtual_vram_gb
|
||||
|
||||
donor_device_info = {}
|
||||
donor_allocations = {}
|
||||
|
||||
for donor in ram_donors:
|
||||
donor_vram = mm.get_total_memory(torch.device(donor)) / (1024**3)
|
||||
max_donor_capacity = donor_vram
|
||||
|
||||
donation = min(remaining_vram_needed, max_donor_capacity)
|
||||
donor_virtual = donor_vram - donation
|
||||
remaining_vram_needed -= donation
|
||||
donor_allocations[donor] = donation
|
||||
|
||||
donor_device_info[donor] = (donor_vram, donor_virtual)
|
||||
logger.info(fmt_assign.format(donor, 'donor', f"{donor_vram:.2f}GB", f"{donor_virtual:.2f}GB", f"-{donation:.2f}GB"))
|
||||
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
# Calculate model size
|
||||
model = model_patcher.model if hasattr(model_patcher, 'model') else model_patcher
|
||||
total_memory = 0
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if hasattr(module, "weight"):
|
||||
if module.weight is not None:
|
||||
total_memory += module.weight.numel() * module.weight.element_size()
|
||||
if hasattr(module, "bias") and module.bias is not None:
|
||||
total_memory += module.bias.numel() * module.bias.element_size()
|
||||
|
||||
model_size_gb = total_memory / (1024**3)
|
||||
new_model_size_gb = max(0, model_size_gb - virtual_vram_gb)
|
||||
|
||||
logger.info(fmt_assign.format('model', 'model', f"{model_size_gb:.2f}GB",f"{new_model_size_gb:.2f}GB", f"-{virtual_vram_gb:.2f}GB"))
|
||||
|
||||
# Warning if model too large
|
||||
if model_size_gb > (recipient_vram * 0.9):
|
||||
required_offload_gb = model_size_gb - (recipient_vram * 0.9)
|
||||
logger.warning(f"\n\n[MultiGPU DisTorch V2] Model size ({model_size_gb:.2f}GB) is larger than 90% of available VRAM on: {recipient_device} ({recipient_vram * 0.9:.2f}GB).")
|
||||
logger.warning(f"[MultiGPU DisTorch V2] To prevent an OOM error, set 'virtual_vram_gb' to at least {required_offload_gb:.2f}.\n\n")
|
||||
|
||||
new_on_recipient = max(0, model_size_gb - virtual_vram_gb)
|
||||
|
||||
# Build allocation string
|
||||
allocation_parts = []
|
||||
recipient_percent = new_on_recipient / recipient_vram
|
||||
allocation_parts.append(f"{recipient_device},{recipient_percent:.4f}")
|
||||
|
||||
for donor in ram_donors:
|
||||
donor_vram = donor_device_info[donor][0]
|
||||
donor_percent = donor_allocations[donor] / donor_vram
|
||||
allocation_parts.append(f"{donor},{donor_percent:.4f}")
|
||||
|
||||
allocations_string = ";".join(allocation_parts)
|
||||
return allocations_string
|
||||
@@ -0,0 +1,761 @@
|
||||
"""
|
||||
DisTorch Safetensor Memory Management Module
|
||||
Contains all safetensor related code for distributed memory management
|
||||
"""
|
||||
|
||||
import sys
|
||||
import torch
|
||||
import logging
|
||||
import hashlib
|
||||
import re
|
||||
import gc
|
||||
|
||||
logger = logging.getLogger("MultiGPU")
|
||||
import copy
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
import comfy.model_management as mm
|
||||
import comfy.model_patcher
|
||||
from .device_utils import get_device_list, soft_empty_cache_multigpu
|
||||
from .model_management_mgpu import multigpu_memory_log, force_full_system_cleanup
|
||||
|
||||
|
||||
safetensor_allocation_store = {}
|
||||
safetensor_settings_store = {}
|
||||
|
||||
|
||||
def create_safetensor_model_hash(model, caller):
|
||||
"""Create a unique hash for a safetensor model to track allocations"""
|
||||
if hasattr(model, 'model'):
|
||||
# For ModelPatcher objects
|
||||
actual_model = model.model
|
||||
model_type = type(actual_model).__name__
|
||||
# Use ComfyUI's model_size if available
|
||||
if hasattr(model, 'model_size'):
|
||||
model_size = model.model_size()
|
||||
else:
|
||||
model_size = sum(p.numel() * p.element_size() for p in actual_model.parameters())
|
||||
if hasattr(model, 'model_state_dict'):
|
||||
first_layers = str(list(model.model_state_dict().keys())[:3])
|
||||
else:
|
||||
first_layers = str(list(actual_model.state_dict().keys())[:3])
|
||||
else:
|
||||
# Direct model
|
||||
model_type = type(model).__name__
|
||||
model_size = sum(p.numel() * p.element_size() for p in model.parameters())
|
||||
first_layers = str(list(model.state_dict().keys())[:3])
|
||||
|
||||
identifier = f"{model_type}_{model_size}_{first_layers}"
|
||||
final_hash = hashlib.sha256(identifier.encode()).hexdigest()
|
||||
|
||||
# DEBUG STATEMENT - ALWAYS LOG THE HASH
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Created hash for {caller}: {final_hash[:8]}...")
|
||||
return final_hash
|
||||
|
||||
def register_patched_safetensor_modelpatcher():
|
||||
"""Register and patch the ModelPatcher for distributed safetensor loading"""
|
||||
from comfy.model_patcher import wipe_lowvram_weight, move_weight_functions
|
||||
# Patch ComfyUI's ModelPatcher
|
||||
if not hasattr(comfy.model_patcher.ModelPatcher, '_distorch_patched'):
|
||||
|
||||
|
||||
# PATCH load_models_gpu with correct memory calculations per model flags
|
||||
original_load_models_gpu = mm.load_models_gpu
|
||||
|
||||
def patched_load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False):
|
||||
from comfy.model_management import cleanup_models_gc, get_free_memory, free_memory, current_loaded_models
|
||||
from comfy.model_management import VRAMState, vram_state, lowvram_available, MIN_WEIGHT_MEMORY_RATIO
|
||||
from comfy.model_management import minimum_inference_memory, extra_reserved_memory, is_device_cpu
|
||||
|
||||
multigpu_memory_log("load_models_gpu_top_level", "start")
|
||||
|
||||
cleanup_models_gc()
|
||||
|
||||
inference_memory = minimum_inference_memory()
|
||||
extra_reserved_mem = extra_reserved_memory()
|
||||
memory_required_total = memory_required + extra_reserved_mem
|
||||
extra_mem = max(inference_memory, memory_required_total)
|
||||
if minimum_memory_required is None:
|
||||
minimum_memory_required = extra_mem
|
||||
else:
|
||||
minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_mem)
|
||||
|
||||
models_temp = set()
|
||||
for m in models:
|
||||
models_temp.add(m)
|
||||
for mm_patch in m.model_patches_models():
|
||||
models_temp.add(mm_patch)
|
||||
|
||||
models = models_temp
|
||||
|
||||
models_to_load = []
|
||||
|
||||
for x in models:
|
||||
loaded_model = mm.LoadedModel(x)
|
||||
try:
|
||||
loaded_model_index = current_loaded_models.index(loaded_model)
|
||||
except:
|
||||
loaded_model_index = None
|
||||
|
||||
if loaded_model_index is not None:
|
||||
loaded = current_loaded_models[loaded_model_index]
|
||||
loaded.currently_used = True
|
||||
models_to_load.append(loaded)
|
||||
else:
|
||||
if hasattr(x, "model"):
|
||||
logging.info(f"Requested to load {x.model.__class__.__name__}")
|
||||
models_to_load.append(loaded_model)
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
to_unload = []
|
||||
for i in range(len(current_loaded_models)):
|
||||
if loaded_model.model.is_clone(current_loaded_models[i].model):
|
||||
to_unload = [i] + to_unload
|
||||
for i in to_unload:
|
||||
model_to_unload = current_loaded_models.pop(i)
|
||||
model_to_unload.model.detach(unpatch_all=False)
|
||||
model_to_unload.model_finalizer.detach()
|
||||
|
||||
# DisTorch Processing
|
||||
total_memory_required = {}
|
||||
eject_device = None
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
device = loaded_model.device
|
||||
base_memory = loaded_model.model_memory_required(device)
|
||||
|
||||
# Check DisTorch flags
|
||||
is_distorch = hasattr(loaded_model.model.model, '_mgpu_virtual_vram_gb')
|
||||
has_eject = hasattr(loaded_model.model.model, '_mgpu_eject_models')
|
||||
|
||||
if has_eject:
|
||||
eject_device = device
|
||||
logger.mgpu_mm_log("DisTorch eject_models=True, is_distorch=True - MAX memory eviction")
|
||||
|
||||
if is_distorch:
|
||||
# is_distorch=True: use compute device allocation size
|
||||
virtual_vram_gb = loaded_model.model.model._mgpu_virtual_vram_gb
|
||||
virtual_vram_bytes = virtual_vram_gb * (1024**3)
|
||||
adjusted_memory = max(0, base_memory - virtual_vram_bytes)
|
||||
total_memory_required[device] = total_memory_required.get(device, 0) + adjusted_memory
|
||||
logger.mgpu_mm_log(f"DisTorch is_distorch=True, model adjusted {(base_memory - virtual_vram_bytes)/(1024**3):.2f}GB for device {device}")
|
||||
else:
|
||||
# is_distorch=False: use full model size
|
||||
total_memory_required[device] = total_memory_required.get(device, 0) + base_memory
|
||||
logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] Standard model {(base_memory)/(1024**3):.2f}GB for device {device}")
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
requested_mem = total_memory_required[device] * 1.1 + extra_mem
|
||||
logger.mgpu_mm_log(f"[FREE_MEMORY_CALL] Device {device}: requesting {requested_mem/(1024**3):.2f}GB = {total_memory_required[device]/(1024**3):.2f}GB * 1.1 + {extra_mem/(1024**3):.2f}GB inference")
|
||||
|
||||
|
||||
multigpu_memory_log("free_memory", "pre")
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
if device == eject_device:
|
||||
total_device_memory = mm.get_total_memory(device)
|
||||
logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] eject_models=1, is_distorch=1 → using MAX memory ({total_device_memory/(1024**3):.2f}GB) for eviction")
|
||||
free_memory(total_device_memory,device)
|
||||
else:
|
||||
logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] eject_models=0, using Comfy Core Computed memory ({(total_memory_required[device] * 1.1 + extra_mem)/(1024**3):.2f}GB) for eviction")
|
||||
free_memory(total_memory_required[device] * 1.1 + extra_mem, device)
|
||||
|
||||
multigpu_memory_log("free_memory/minimum_memory_required", "post/pre")
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
free_mem = get_free_memory(device)
|
||||
free_mem_gb = free_mem / (1024**3)
|
||||
min_required_gb = minimum_memory_required / (1024**3)
|
||||
logger.mgpu_mm_log(f"[MIN_MEMORY_CHECK] Device {device}: free={free_mem_gb:.2f}GB, required={min_required_gb:.2f}GB, will_evict={free_mem < minimum_memory_required}")
|
||||
|
||||
if free_mem < minimum_memory_required:
|
||||
models_l = free_memory(minimum_memory_required, device)
|
||||
logger.mgpu_mm_log(f"[EVICTION] Device {device}: unloaded {len(models_l)} models due to insufficient memory")
|
||||
logging.info("{} models unloaded.".format(len(models_l)))
|
||||
|
||||
multigpu_memory_log("minimum_memory_required", "post")
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
model = loaded_model.model
|
||||
torch_dev = model.load_device
|
||||
if is_device_cpu(torch_dev):
|
||||
vram_set_state = VRAMState.DISABLED
|
||||
else:
|
||||
vram_set_state = vram_state
|
||||
lowvram_model_memory = 0
|
||||
if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load:
|
||||
loaded_memory = loaded_model.model_loaded_memory()
|
||||
current_free_mem = get_free_memory(torch_dev) + loaded_memory
|
||||
|
||||
lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory()))
|
||||
lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory)
|
||||
|
||||
if vram_set_state == VRAMState.NO_VRAM:
|
||||
lowvram_model_memory = 0.1
|
||||
|
||||
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
current_loaded_models.insert(0, loaded_model)
|
||||
|
||||
# Replace the module function
|
||||
mm.load_models_gpu = patched_load_models_gpu
|
||||
|
||||
original_partially_load = comfy.model_patcher.ModelPatcher.partially_load
|
||||
|
||||
def new_partially_load(self, device_to, extra_memory=0, full_load=False, force_patch_weights=False, **kwargs):
|
||||
"""Override to use our static device assignments"""
|
||||
global safetensor_allocation_store
|
||||
|
||||
debug_hash = create_safetensor_model_hash(self, "partial_load")
|
||||
multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "pre-load")
|
||||
allocations = safetensor_allocation_store.get(debug_hash)
|
||||
|
||||
# Set default precision flag before checking
|
||||
if not hasattr(self.model, '_distorch_high_precision_loras'):
|
||||
self.model._distorch_high_precision_loras = True
|
||||
|
||||
if not allocations:
|
||||
result = original_partially_load(self, device_to, extra_memory, force_patch_weights)
|
||||
multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "post-load")
|
||||
if hasattr(self, '_distorch_block_assignments'):
|
||||
del self._distorch_block_assignments
|
||||
return result
|
||||
|
||||
if not hasattr(self.model, 'current_weight_patches_uuid'):
|
||||
self.model.current_weight_patches_uuid = None
|
||||
|
||||
unpatch_weights = self.model.current_weight_patches_uuid is not None and (self.model.current_weight_patches_uuid != self.patches_uuid or force_patch_weights)
|
||||
|
||||
if unpatch_weights:
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Patches changed or forced. Unpatching model.")
|
||||
self.unpatch_model(self.offload_device, unpatch_weights=True)
|
||||
|
||||
self.patch_model(load_weights=False)
|
||||
|
||||
mem_counter = 0
|
||||
|
||||
is_clip_model = getattr(self, 'is_clip', False)
|
||||
device_assignments = analyze_safetensor_loading(self, allocations, is_clip=is_clip_model)
|
||||
|
||||
model_original_dtype = comfy.utils.weight_dtype(self.model.state_dict())
|
||||
high_precision_loras = getattr(self.model, "_distorch_high_precision_loras", True)
|
||||
loading = self._load_list()
|
||||
loading.sort(reverse=True)
|
||||
for module_size, module_name, module_object, params in loading:
|
||||
if not unpatch_weights and hasattr(module_object, "comfy_patched_weights") and module_object.comfy_patched_weights == True:
|
||||
block_target_device = device_assignments['block_assignments'].get(module_name, device_to)
|
||||
current_module_device = None
|
||||
try:
|
||||
if any(p.numel() > 0 for p in module_object.parameters(recurse=False)):
|
||||
current_module_device = next(module_object.parameters(recurse=False)).device
|
||||
except StopIteration:
|
||||
pass
|
||||
|
||||
if current_module_device is not None and str(current_module_device) != str(block_target_device):
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Moving already patched {module_name} to {block_target_device}")
|
||||
module_object.to(block_target_device)
|
||||
|
||||
mem_counter += module_size
|
||||
continue
|
||||
|
||||
# Step 1: Write block/tensor to compute device first
|
||||
module_object.to(device_to)
|
||||
|
||||
# Step 2: Apply LoRa patches while on compute device
|
||||
weight_key = "{}.weight".format(module_name)
|
||||
bias_key = "{}.bias".format(module_name)
|
||||
|
||||
if weight_key in self.patches:
|
||||
self.patch_weight_to_device(weight_key, device_to=device_to)
|
||||
if weight_key in self.weight_wrapper_patches:
|
||||
module_object.weight_function.extend(self.weight_wrapper_patches[weight_key])
|
||||
|
||||
if bias_key in self.patches:
|
||||
self.patch_weight_to_device(bias_key, device_to=device_to)
|
||||
if bias_key in self.weight_wrapper_patches:
|
||||
module_object.bias_function.extend(self.weight_wrapper_patches[bias_key])
|
||||
|
||||
# Step 3: FP8 casting for CPU storage (if enabled)
|
||||
block_target_device = device_assignments['block_assignments'].get(module_name, device_to)
|
||||
has_patches = weight_key in self.patches or bias_key in self.patches
|
||||
|
||||
if not high_precision_loras and block_target_device == "cpu" and has_patches and model_original_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
|
||||
for param_name, param in module_object.named_parameters():
|
||||
if param.dtype.is_floating_point:
|
||||
cast_data = comfy.float.stochastic_rounding(param.data, torch.float8_e4m3fn)
|
||||
new_param = torch.nn.Parameter(cast_data.to(torch.float8_e4m3fn))
|
||||
new_param.requires_grad = param.requires_grad
|
||||
setattr(module_object, param_name, new_param)
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Cast {module_name}.{param_name} to FP8 for CPU storage")
|
||||
|
||||
# Step 4: Move to ultimate destination based on DisTorch assignment
|
||||
if block_target_device != device_to:
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Moving {module_name} from {device_to} to {block_target_device}")
|
||||
module_object.to(block_target_device)
|
||||
module_object.comfy_cast_weights = True
|
||||
|
||||
# Mark as patched and update memory counter
|
||||
module_object.comfy_patched_weights = True
|
||||
mem_counter += module_size
|
||||
|
||||
self.model.current_weight_patches_uuid = self.patches_uuid
|
||||
|
||||
logger.info("[MultiGPU DisTorch V2] DisTorch loading completed.")
|
||||
logger.info(f"[MultiGPU DisTorch V2] Total memory: {mem_counter / (1024 * 1024):.2f}MB")
|
||||
multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "post-load")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
comfy.model_patcher.ModelPatcher.partially_load = new_partially_load
|
||||
comfy.model_patcher.ModelPatcher._distorch_patched = True
|
||||
logger.info("[MultiGPU Core Patching] Successfully patched ModelPatcher.partially_load")
|
||||
|
||||
def _extract_clip_head_blocks(raw_block_list, compute_device):
|
||||
"""Identify and pre-assign CLIP head blocks to compute device returning head_blocks, distributable_blocks, block_assignments, and head_memory."""
|
||||
head_keywords = ['embed', 'wte', 'wpe', 'token_embedding', 'position_embedding']
|
||||
head_blocks = []
|
||||
distributable_blocks = []
|
||||
head_memory = 0
|
||||
block_assignments = {}
|
||||
|
||||
for module_size, module_name, module_object, params in raw_block_list:
|
||||
if any(kw in module_name.lower() for kw in head_keywords):
|
||||
head_blocks.append((module_size, module_name, module_object, params))
|
||||
block_assignments[module_name] = compute_device
|
||||
head_memory += module_size
|
||||
else:
|
||||
distributable_blocks.append((module_size, module_name, module_object, params))
|
||||
|
||||
return head_blocks, distributable_blocks, block_assignments, head_memory
|
||||
|
||||
def analyze_safetensor_loading(model_patcher, allocations_string, is_clip=False):
|
||||
"""
|
||||
Analyze and distribute safetensor model blocks across devices.
|
||||
Supports CLIP head preservation when is_clip=True.
|
||||
"""
|
||||
DEVICE_RATIOS_DISTORCH = {}
|
||||
device_table = {}
|
||||
distorch_alloc = allocations_string
|
||||
virtual_vram_gb = 0.0
|
||||
|
||||
distorch_alloc, virtual_vram_str = allocations_string.split('#')
|
||||
|
||||
compute_device = virtual_vram_str.split(';')[0]
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Compute Device: {compute_device}")
|
||||
|
||||
if not distorch_alloc:
|
||||
mode = "fraction"
|
||||
distorch_alloc = calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str)
|
||||
|
||||
elif any(c in distorch_alloc.lower() for c in ['g', 'm', 'k', 'b']):
|
||||
mode = "byte"
|
||||
distorch_alloc = calculate_fraction_from_byte_expert_string(model_patcher, distorch_alloc)
|
||||
elif "%" in distorch_alloc:
|
||||
mode = "ratio"
|
||||
distorch_alloc = calculate_fraction_from_ratio_expert_string(model_patcher, distorch_alloc)
|
||||
|
||||
all_devices = get_device_list()
|
||||
present_devices = {item.split(',')[0] for item in distorch_alloc.split(';') if ',' in item}
|
||||
for device in all_devices:
|
||||
if device not in present_devices:
|
||||
distorch_alloc += f";{device},0.0"
|
||||
|
||||
eq_line = "=" * 50
|
||||
dash_line = "-" * 50
|
||||
fmt_assign = "{:<18}{:>7}{:>14}{:>10}"
|
||||
|
||||
logger.info(eq_line)
|
||||
logger.info(f"[MultiGPU DisTorch V2] Final Allocation String:\n{distorch_alloc}")
|
||||
|
||||
for allocation in distorch_alloc.split(';'):
|
||||
if ',' not in allocation:
|
||||
continue
|
||||
dev_name, fraction = allocation.split(',')
|
||||
fraction = float(fraction)
|
||||
total_mem_bytes = mm.get_total_memory(torch.device(dev_name))
|
||||
alloc_gb = (total_mem_bytes * fraction) / (1024**3)
|
||||
DEVICE_RATIOS_DISTORCH[dev_name] = alloc_gb
|
||||
device_table[dev_name] = {
|
||||
"fraction": fraction,
|
||||
"total_gb": total_mem_bytes / (1024**3),
|
||||
"alloc_gb": alloc_gb
|
||||
}
|
||||
|
||||
logger.info(eq_line)
|
||||
logger.info(" DisTorch2 Model Device Allocations")
|
||||
logger.info(eq_line)
|
||||
|
||||
fmt_rosetta = "{:<8}{:>9}{:>9}{:>11}{:>10}"
|
||||
logger.info(fmt_rosetta.format("Device", "VRAM GB", "Dev %", "Model GB", "Dist %"))
|
||||
logger.info(dash_line)
|
||||
|
||||
sorted_devices = sorted(device_table.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
total_allocated_model_bytes = sum(d["alloc_gb"] * (1024**3) for d in device_table.values())
|
||||
|
||||
for dev in sorted_devices:
|
||||
total_dev_gb = device_table[dev]["total_gb"]
|
||||
alloc_fraction = device_table[dev]["fraction"]
|
||||
alloc_gb = device_table[dev]["alloc_gb"]
|
||||
|
||||
dist_ratio_percent = (alloc_gb * (1024**3) / total_allocated_model_bytes) * 100 if total_allocated_model_bytes > 0 else 0
|
||||
|
||||
logger.info(fmt_rosetta.format(
|
||||
dev,
|
||||
f"{total_dev_gb:.2f}",
|
||||
f"{alloc_fraction*100:.1f}%",
|
||||
f"{alloc_gb:.2f}",
|
||||
f"{dist_ratio_percent:.1f}%"
|
||||
))
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
block_summary = {}
|
||||
block_list = []
|
||||
memory_by_type = defaultdict(int)
|
||||
total_memory = 0
|
||||
|
||||
raw_block_list = model_patcher._load_list()
|
||||
total_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
|
||||
MIN_BLOCK_THRESHOLD = total_memory * 0.0001
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Total model memory: {total_memory} bytes")
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Tiny block threshold (0.01%): {MIN_BLOCK_THRESHOLD} bytes")
|
||||
|
||||
# CLIP-specific: Extract head blocks and get pre-assignments
|
||||
head_memory = 0
|
||||
block_assignments = {}
|
||||
if is_clip:
|
||||
head_blocks, distributable_raw, block_assignments, head_memory = \
|
||||
_extract_clip_head_blocks(raw_block_list, compute_device)
|
||||
logger.info(f"[MultiGPU DisTorch V2 CLIP] Preserving {len(head_blocks)} head layer(s) ({head_memory/(1024**2):.2f} MB) on compute device: {compute_device}")
|
||||
else:
|
||||
distributable_raw = raw_block_list
|
||||
|
||||
# Build all_blocks list for summary (using full raw_block_list)
|
||||
all_blocks = []
|
||||
for module_size, module_name, module_object, params in raw_block_list:
|
||||
block_type = type(module_object).__name__
|
||||
# Populate summary dictionaries
|
||||
block_summary[block_type] = block_summary.get(block_type, 0) + 1
|
||||
memory_by_type[block_type] += module_size
|
||||
all_blocks.append((module_name, module_object, block_type, module_size))
|
||||
|
||||
# Use distributable blocks for actual allocation (for CLIP, this excludes heads)
|
||||
distributable_all_blocks = []
|
||||
for module_size, module_name, module_object, params in distributable_raw:
|
||||
distributable_all_blocks.append((module_name, module_object, type(module_object).__name__, module_size))
|
||||
|
||||
block_list = [b for b in distributable_all_blocks if b[3] >= MIN_BLOCK_THRESHOLD]
|
||||
tiny_block_list = [b for b in distributable_all_blocks if b[3] < MIN_BLOCK_THRESHOLD]
|
||||
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Total blocks: {len(all_blocks)}")
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Distributable blocks: {len(block_list)}")
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Tiny blocks (<0.01%): {len(tiny_block_list)}")
|
||||
|
||||
logger.info(" DisTorch2 Model Layer Distribution")
|
||||
logger.info(dash_line)
|
||||
fmt_layer = "{:<18}{:>7}{:>14}{:>10}"
|
||||
logger.info(fmt_layer.format("Layer Type", "Layers", "Memory (MB)", "% Total"))
|
||||
logger.info(dash_line)
|
||||
|
||||
for layer_type, count in block_summary.items():
|
||||
mem_mb = memory_by_type[layer_type] / (1024 * 1024)
|
||||
mem_percent = (memory_by_type[layer_type] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logger.info(fmt_layer.format(layer_type[:18], str(count), f"{mem_mb:.2f}", f"{mem_percent:.1f}%"))
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
# Distribute blocks sequentially from the tail of the model
|
||||
|
||||
device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()}
|
||||
block_assignments = {}
|
||||
|
||||
# Create a memory quota for each donor device based on its calculated allocation.
|
||||
donor_devices = [d for d in sorted_devices]
|
||||
donor_quotas = {
|
||||
dev: device_table[dev]["alloc_gb"] * (1024**3)
|
||||
for dev in donor_devices
|
||||
}
|
||||
|
||||
# CLIP-specific: Adjust compute_device quota to account for locked head blocks
|
||||
if is_clip and compute_device in donor_quotas and head_memory > 0:
|
||||
donor_quotas[compute_device] = max(0, donor_quotas[compute_device] - head_memory)
|
||||
logger.debug(f"[MultiGPU DisTorch V2 CLIP] Adjusted {compute_device} quota by -{head_memory/(1024**2):.2f} MB for head preservation")
|
||||
|
||||
# Iterate from the TAIL of the model, assigning blocks to donors until their quotas are filled.
|
||||
for block_name, module, block_type, block_memory in reversed(block_list):
|
||||
assigned_to_donor = False
|
||||
for donor in donor_devices:
|
||||
if donor_quotas[donor] >= block_memory:
|
||||
block_assignments[block_name] = donor
|
||||
donor_quotas[donor] -= block_memory
|
||||
assigned_to_donor = True
|
||||
break # Move to the next block
|
||||
|
||||
if not assigned_to_donor: #Note - small rounding errors and tensor-fitting on devices make a block occasionally an orphan. We treat orphans the same as tiny_block_list as they are generally small rounding errors
|
||||
block_assignments[block_name] = compute_device
|
||||
|
||||
if tiny_block_list:
|
||||
for block_name, module, block_type, block_memory in tiny_block_list:
|
||||
block_assignments[block_name] = compute_device
|
||||
|
||||
# Populate device_assignments from the final block_assignments
|
||||
for block_name, device in block_assignments.items():
|
||||
# Find the block in the original list to get all its info
|
||||
for b_name, b_module, b_type, b_mem in all_blocks:
|
||||
if b_name == block_name:
|
||||
device_assignments[device].append((b_name, b_module, b_type, b_mem))
|
||||
break
|
||||
|
||||
logger.info("DisTorch2 Model Final Device/Layer Assignments")
|
||||
logger.info(dash_line)
|
||||
logger.info(fmt_assign.format("Device", "Layers", "Memory (MB)", "% Total"))
|
||||
logger.info(dash_line)
|
||||
|
||||
if tiny_block_list:
|
||||
tiny_block_memory = sum(b[3] for b in tiny_block_list)
|
||||
tiny_mem_mb = tiny_block_memory / (1024 * 1024)
|
||||
tiny_mem_percent = (tiny_block_memory / total_memory) * 100 if total_memory > 0 else 0
|
||||
device_label = f"{compute_device} (<0.01%)"
|
||||
logger.info(fmt_assign.format(device_label, str(len(tiny_block_list)), f"{tiny_mem_mb:.2f}", f"{tiny_mem_percent:.1f}%"))
|
||||
logger.debug(f"[MultiGPU DisTorch V2] Tiny block memory breakdown: {tiny_block_memory} bytes ({tiny_mem_mb:.2f} MB), which is {tiny_mem_percent:.4f}% of total model memory.")
|
||||
|
||||
total_assigned_memory = 0
|
||||
device_memories = {}
|
||||
|
||||
for device, blocks in device_assignments.items():
|
||||
dist_blocks = [b for b in blocks if b[3] >= MIN_BLOCK_THRESHOLD]
|
||||
if not dist_blocks:
|
||||
continue
|
||||
|
||||
device_memory = sum(b[3] for b in dist_blocks)
|
||||
device_memories[device] = device_memory
|
||||
total_assigned_memory += device_memory
|
||||
|
||||
sorted_assignments = sorted(device_memories.keys(), key=lambda d: (d == "cpu", d))
|
||||
|
||||
for dev in sorted_assignments:
|
||||
# Get only the distributed blocks for the count
|
||||
dist_blocks = [b for b in device_assignments[dev] if b[3] >= MIN_BLOCK_THRESHOLD]
|
||||
if not dist_blocks:
|
||||
continue
|
||||
|
||||
mem_mb = device_memories[dev] / (1024 * 1024)
|
||||
mem_percent = (device_memories[dev] / total_memory) * 100 if total_memory > 0 else 0
|
||||
logger.info(fmt_assign.format(dev, str(len(dist_blocks)), f"{mem_mb:.2f}", f"{mem_percent:.1f}%"))
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
return {
|
||||
"device_assignments": device_assignments,
|
||||
"block_assignments": block_assignments
|
||||
}
|
||||
|
||||
def parse_memory_string(mem_str):
|
||||
"""Parses a memory string (e.g., '4.0g', '512M') and returns bytes."""
|
||||
mem_str = mem_str.strip().lower()
|
||||
match = re.match(r'(\d+\.?\d*)\s*([gmkb]?)', mem_str)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid memory string format: {mem_str}")
|
||||
|
||||
val, unit = match.groups()
|
||||
val = float(val)
|
||||
|
||||
if unit == 'g':
|
||||
return val * (1024**3)
|
||||
elif unit == 'm':
|
||||
return val * (1024**2)
|
||||
elif unit == 'k':
|
||||
return val * 1024
|
||||
else: # b or no unit
|
||||
return val
|
||||
|
||||
def calculate_fraction_from_byte_expert_string(model_patcher, byte_str):
|
||||
"""Convert byte allocation string (e.g. 'cuda:1,4gb;cpu,*') to fractional VRAM allocation string respecting device order and byte quotas."""
|
||||
raw_block_list = model_patcher._load_list()
|
||||
total_model_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
remaining_model_bytes = total_model_memory
|
||||
|
||||
# Use a list of tuples to preserve the user-defined order
|
||||
parsed_allocations = []
|
||||
wildcard_device = "cpu" # Default wildcard device
|
||||
|
||||
for allocation in byte_str.split(';'):
|
||||
if ',' not in allocation:
|
||||
continue
|
||||
dev_name, val_str = allocation.split(',', 1)
|
||||
is_wildcard = '*' in val_str
|
||||
|
||||
if is_wildcard:
|
||||
wildcard_device = dev_name
|
||||
# Don't add wildcard to the priority list yet
|
||||
else:
|
||||
byte_val = parse_memory_string(val_str)
|
||||
parsed_allocations.append({'device': dev_name, 'bytes': byte_val})
|
||||
|
||||
final_byte_allocations = defaultdict(int)
|
||||
|
||||
# Process devices with specific byte allocations first, in order
|
||||
for alloc in parsed_allocations:
|
||||
dev = alloc['device']
|
||||
requested_bytes = alloc['bytes']
|
||||
|
||||
# Determine the actual bytes to allocate to this device
|
||||
bytes_to_assign = min(requested_bytes, remaining_model_bytes)
|
||||
|
||||
if bytes_to_assign > 0:
|
||||
final_byte_allocations[dev] = bytes_to_assign
|
||||
remaining_model_bytes -= bytes_to_assign
|
||||
logger.info(f"[MultiGPU DisTorch V2] Assigning {bytes_to_assign / (1024**2):.2f}MB of model to {dev} (requested {requested_bytes / (1024**2):.2f}MB).")
|
||||
|
||||
if remaining_model_bytes <= 0:
|
||||
logger.info("[MultiGPU DisTorch V2] All model blocks have been allocated. Subsequent devices in the string will receive no assignment.")
|
||||
break
|
||||
|
||||
# Assign any leftover model bytes to the wildcard device
|
||||
if remaining_model_bytes > 0:
|
||||
final_byte_allocations[wildcard_device] += remaining_model_bytes
|
||||
logger.info(f"[MultiGPU DisTorch V2] Assigning remaining {remaining_model_bytes / (1024**2):.2f}MB of model to wildcard device '{wildcard_device}'.")
|
||||
|
||||
# Convert the final byte allocations to VRAM fractions
|
||||
allocation_parts = []
|
||||
for dev, bytes_alloc in final_byte_allocations.items():
|
||||
total_device_vram = mm.get_total_memory(torch.device(dev))
|
||||
if total_device_vram > 0:
|
||||
fraction = bytes_alloc / total_device_vram
|
||||
allocation_parts.append(f"{dev},{fraction:.4f}")
|
||||
|
||||
allocations_string = ";".join(allocation_parts)
|
||||
|
||||
return allocations_string
|
||||
|
||||
def calculate_fraction_from_ratio_expert_string(model_patcher, ratio_str):
|
||||
"""Convert ratio allocation string (e.g. 'cuda:0,25%;cpu,75%') describing model split to fractional VRAM allocation string."""
|
||||
raw_block_list = model_patcher._load_list()
|
||||
total_model_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
|
||||
raw_ratios = {}
|
||||
for allocation in ratio_str.split(';'):
|
||||
if ',' not in allocation: continue
|
||||
dev_name, val_str = allocation.split(',', 1)
|
||||
# Assumes the value is a unitless ratio number, ignores '%' for simplicity.
|
||||
value = float(val_str.replace('%','').strip())
|
||||
raw_ratios[dev_name] = value
|
||||
|
||||
total_ratio_parts = sum(raw_ratios.values())
|
||||
allocation_parts = []
|
||||
|
||||
for dev, ratio_val in raw_ratios.items():
|
||||
bytes_of_model_for_device = (ratio_val / total_ratio_parts) * total_model_memory
|
||||
|
||||
total_vram_of_device = mm.get_total_memory(torch.device(dev))
|
||||
|
||||
if total_vram_of_device > 0:
|
||||
required_fraction = bytes_of_model_for_device / total_vram_of_device
|
||||
allocation_parts.append(f"{dev},{required_fraction:.4f}")
|
||||
|
||||
ratio_values = [str(v) for v in raw_ratios.values()]
|
||||
ratio_string = ":".join(ratio_values)
|
||||
|
||||
normalized_pcts = [(v / total_ratio_parts) * 100 for v in raw_ratios.values()]
|
||||
|
||||
put_parts = []
|
||||
for i, dev_name in enumerate(raw_ratios.keys()):
|
||||
put_parts.append(f"{int(normalized_pcts[i])}% on {dev_name}")
|
||||
|
||||
if len(put_parts) == 1:
|
||||
put_part = put_parts[0]
|
||||
elif len(put_parts) == 2:
|
||||
put_part = f"{put_parts[0]} and {put_parts[1]}"
|
||||
else:
|
||||
put_part = ", ".join(put_parts[:-1]) + f", and {put_parts[-1]}"
|
||||
|
||||
logger.info(f"[MultiGPU DisTorch V2] Ratio(%) Mode - {ratio_str} -> {ratio_string} ratio, put {put_part}")
|
||||
|
||||
allocations_string = ";".join(allocation_parts)
|
||||
|
||||
return allocations_string
|
||||
|
||||
def calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str):
|
||||
"""Calculate virtual VRAM allocation string for distributed safetensor loading"""
|
||||
recipient_device, vram_amount, donors = virtual_vram_str.split(';')
|
||||
virtual_vram_gb = float(vram_amount)
|
||||
|
||||
eq_line = "=" * 47
|
||||
dash_line = "-" * 47
|
||||
fmt_assign = "{:<8} {:<6} {:>11} {:>9} {:>9}"
|
||||
|
||||
logger.info(eq_line)
|
||||
logger.info(" DisTorch2 Model Virtual VRAM Analysis")
|
||||
logger.info(eq_line)
|
||||
logger.info(fmt_assign.format("Object", "Role", "Original(GB)", "Total(GB)", "Virt(GB)"))
|
||||
logger.info(dash_line)
|
||||
|
||||
# Calculate recipient VRAM
|
||||
recipient_vram = mm.get_total_memory(torch.device(recipient_device)) / (1024**3)
|
||||
recipient_virtual = recipient_vram + virtual_vram_gb
|
||||
|
||||
logger.info(fmt_assign.format(recipient_device, 'recip', f"{recipient_vram:.2f}GB",f"{recipient_virtual:.2f}GB", f"+{virtual_vram_gb:.2f}GB"))
|
||||
|
||||
# Handle donor devices
|
||||
ram_donors = [d for d in donors.split(',')]
|
||||
remaining_vram_needed = virtual_vram_gb
|
||||
|
||||
donor_device_info = {}
|
||||
donor_allocations = {}
|
||||
|
||||
for donor in ram_donors:
|
||||
donor_vram = mm.get_total_memory(torch.device(donor)) / (1024**3)
|
||||
max_donor_capacity = donor_vram
|
||||
|
||||
donation = min(remaining_vram_needed, max_donor_capacity)
|
||||
donor_virtual = donor_vram - donation
|
||||
remaining_vram_needed -= donation
|
||||
donor_allocations[donor] = donation
|
||||
|
||||
donor_device_info[donor] = (donor_vram, donor_virtual)
|
||||
logger.info(fmt_assign.format(donor, 'donor', f"{donor_vram:.2f}GB", f"{donor_virtual:.2f}GB", f"-{donation:.2f}GB"))
|
||||
|
||||
|
||||
logger.info(dash_line)
|
||||
|
||||
# Calculate model size
|
||||
model = model_patcher.model if hasattr(model_patcher, 'model') else model_patcher
|
||||
total_memory = 0
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if hasattr(module, "weight"):
|
||||
if module.weight is not None:
|
||||
total_memory += module.weight.numel() * module.weight.element_size()
|
||||
if hasattr(module, "bias") and module.bias is not None:
|
||||
total_memory += module.bias.numel() * module.bias.element_size()
|
||||
|
||||
model_size_gb = total_memory / (1024**3)
|
||||
new_model_size_gb = max(0, model_size_gb - virtual_vram_gb)
|
||||
|
||||
logger.info(fmt_assign.format('model', 'model', f"{model_size_gb:.2f}GB",f"{new_model_size_gb:.2f}GB", f"-{virtual_vram_gb:.2f}GB"))
|
||||
|
||||
# Warning if model too large
|
||||
if model_size_gb > (recipient_vram * 0.9):
|
||||
required_offload_gb = model_size_gb - (recipient_vram * 0.9)
|
||||
logger.warning(f"\n\n[MultiGPU DisTorch V2] Model size ({model_size_gb:.2f}GB) is larger than 90% of available VRAM on: {recipient_device} ({recipient_vram * 0.9:.2f}GB).")
|
||||
logger.warning(f"[MultiGPU DisTorch V2] To prevent an OOM error, set 'virtual_vram_gb' to at least {required_offload_gb:.2f}.\n\n")
|
||||
|
||||
new_on_recipient = max(0, model_size_gb - virtual_vram_gb)
|
||||
|
||||
# Build allocation string
|
||||
allocation_parts = []
|
||||
recipient_percent = new_on_recipient / recipient_vram
|
||||
allocation_parts.append(f"{recipient_device},{recipient_percent:.4f}")
|
||||
|
||||
for donor in ram_donors:
|
||||
donor_vram = donor_device_info[donor][0]
|
||||
donor_percent = donor_allocations[donor] / donor_vram
|
||||
allocation_parts.append(f"{donor},{donor_percent:.4f}")
|
||||
|
||||
allocations_string = ";".join(allocation_parts)
|
||||
return allocations_string
|
||||
@@ -0,0 +1,439 @@
|
||||
import torch
|
||||
import folder_paths
|
||||
from pathlib import Path
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
from .device_utils import get_device_list
|
||||
from .model_management_mgpu import force_full_system_cleanup
|
||||
|
||||
class UnetLoaderGGUF:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (unet_names,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_unet"
|
||||
CATEGORY = "bootleg"
|
||||
TITLE = "Unet Loader (GGUF)"
|
||||
|
||||
def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None):
|
||||
"""Load GGUF format UNet model."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]()
|
||||
return original_loader.load_unet(unet_name, dequant_dtype, patch_dtype, patch_on_device)
|
||||
|
||||
class UnetLoaderGGUFAdvanced(UnetLoaderGGUF):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (unet_names,),
|
||||
"dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}),
|
||||
"patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}),
|
||||
"patch_on_device": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
TITLE = "Unet Loader (GGUF/Advanced)"
|
||||
|
||||
|
||||
class CLIPLoaderGGUF:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
import nodes
|
||||
base = nodes.CLIPLoader.INPUT_TYPES()
|
||||
return {
|
||||
"required": {
|
||||
"clip_name": (s.get_filename_list(),),
|
||||
"type": base["required"]["type"],
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "load_clip"
|
||||
CATEGORY = "bootleg"
|
||||
TITLE = "CLIPLoader (GGUF)"
|
||||
|
||||
@classmethod
|
||||
def get_filename_list(s):
|
||||
"""Get combined list of CLIP and CLIP_GGUF model files."""
|
||||
files = []
|
||||
files += folder_paths.get_filename_list("clip")
|
||||
files += folder_paths.get_filename_list("clip_gguf")
|
||||
return sorted(files)
|
||||
|
||||
def load_data(self, ckpt_paths):
|
||||
"""Load CLIP model data from checkpoint paths."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]()
|
||||
return original_loader.load_data(ckpt_paths)
|
||||
|
||||
def load_patcher(self, clip_paths, clip_type, clip_data):
|
||||
"""Create ModelPatcher for CLIP model."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]()
|
||||
return original_loader.load_patcher(clip_paths, clip_type, clip_data)
|
||||
|
||||
def load_clip(self, clip_name, type="stable_diffusion", device=None):
|
||||
"""Load CLIP model from GGUF or standard format."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]()
|
||||
return original_loader.load_clip(clip_name, type)
|
||||
|
||||
class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
import nodes
|
||||
base = nodes.DualCLIPLoader.INPUT_TYPES()
|
||||
file_options = (s.get_filename_list(), )
|
||||
return {
|
||||
"required": {
|
||||
"clip_name1": file_options,
|
||||
"clip_name2": file_options,
|
||||
"type": base["required"]["type"],
|
||||
}
|
||||
}
|
||||
|
||||
TITLE = "DualCLIPLoader (GGUF)"
|
||||
|
||||
def load_clip(self, clip_name1, clip_name2, type, device=None):
|
||||
"""Load dual CLIP model configuration."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUF"]()
|
||||
clip = original_loader.load_clip(clip_name1, clip_name2, type)
|
||||
clip[0].patcher.load(force_patch_weights=True)
|
||||
return clip
|
||||
|
||||
|
||||
class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
file_options = (s.get_filename_list(), )
|
||||
return {
|
||||
"required": {
|
||||
"clip_name1": file_options,
|
||||
"clip_name2": file_options,
|
||||
"clip_name3": file_options,
|
||||
}
|
||||
}
|
||||
|
||||
TITLE = "TripleCLIPLoader (GGUF)"
|
||||
|
||||
def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"):
|
||||
"""Load triple CLIP model configuration for SD3."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUF"]()
|
||||
return original_loader.load_clip(clip_name1, clip_name2, clip_name3, type)
|
||||
|
||||
class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
file_options = (s.get_filename_list(), )
|
||||
return {
|
||||
"required": {
|
||||
"clip_name1": file_options,
|
||||
"clip_name2": file_options,
|
||||
"clip_name3": file_options,
|
||||
"clip_name4": file_options,
|
||||
}
|
||||
}
|
||||
|
||||
TITLE = "QuadrupleCLIPLoader (GGUF)"
|
||||
|
||||
def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"):
|
||||
"""Load quadruple CLIP model configuration."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderGGUF"]()
|
||||
return original_loader.load_clip(clip_name1, clip_name2, clip_name3, clip_name4, type)
|
||||
|
||||
|
||||
class LTXVLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),
|
||||
{"tooltip": "The name of the checkpoint (model) to load."}),
|
||||
"dtype": (["bfloat16", "float32"], {"default": "bfloat16"})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "VAE")
|
||||
RETURN_NAMES = ("model", "vae")
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def load(self, ckpt_name, dtype):
|
||||
"""Load LTXV model and VAE with specified precision."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
||||
return original_loader.load(ckpt_name, dtype)
|
||||
def _load_unet(self, load_device, offload_device, weights, num_latent_channels, dtype, config=None ):
|
||||
"""Load LTXV UNet with device-specific configuration."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
||||
return original_loader._load_unet(load_device, offload_device, weights, num_latent_channels, dtype, config=None )
|
||||
def _load_vae(self, weights, config=None):
|
||||
"""Load LTXV VAE from weights."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]()
|
||||
return original_loader._load_vae(weights, config=None)
|
||||
|
||||
class Florence2ModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
|
||||
"precision": (['fp16','bf16','fp32'],),
|
||||
"attention": (
|
||||
[ 'flash_attention_2', 'sdpa', 'eager'],
|
||||
{
|
||||
"default": 'sdpa'
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"lora": ("PEFTLORA",),
|
||||
"convert_to_safetensors": ("BOOLEAN", {"default": False, "tooltip": "Some of the older model weights are not saved in .safetensors format, which seem to cause longer loading times, this option converts the .bin weights to .safetensors"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FL2MODEL",)
|
||||
RETURN_NAMES = ("florence2_model",)
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "Florence2"
|
||||
|
||||
def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
|
||||
"""Load Florence2 vision model with specified precision and attention mode."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["Florence2ModelLoader"]()
|
||||
return original_loader.loadmodel(model, precision, attention, lora, convert_to_safetensors)
|
||||
|
||||
class DownloadAndLoadFlorence2Model:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"model": (
|
||||
[
|
||||
'microsoft/Florence-2-base',
|
||||
'microsoft/Florence-2-base-ft',
|
||||
'microsoft/Florence-2-large',
|
||||
'microsoft/Florence-2-large-ft',
|
||||
'HuggingFaceM4/Florence-2-DocVQA',
|
||||
'thwri/CogFlorence-2.1-Large',
|
||||
'thwri/CogFlorence-2.2-Large',
|
||||
'gokaygokay/Florence-2-SD3-Captioner',
|
||||
'gokaygokay/Florence-2-Flux-Large',
|
||||
'MiaoshouAI/Florence-2-base-PromptGen-v1.5',
|
||||
'MiaoshouAI/Florence-2-large-PromptGen-v1.5',
|
||||
'MiaoshouAI/Florence-2-base-PromptGen-v2.0',
|
||||
'MiaoshouAI/Florence-2-large-PromptGen-v2.0',
|
||||
'PJMixers-Images/Florence-2-base-Castollux-v0.5'
|
||||
],
|
||||
{
|
||||
"default": 'microsoft/Florence-2-base'
|
||||
}),
|
||||
"precision": ([ 'fp16','bf16','fp32'],
|
||||
{
|
||||
"default": 'fp16'
|
||||
}),
|
||||
"attention": (
|
||||
[ 'flash_attention_2', 'sdpa', 'eager'],
|
||||
{
|
||||
"default": 'sdpa'
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"lora": ("PEFTLORA",),
|
||||
"convert_to_safetensors": ("BOOLEAN", {"default": False, "tooltip": "Some of the older model weights are not saved in .safetensors format, which seem to cause longer loading times, this option converts the .bin weights to .safetensors"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FL2MODEL",)
|
||||
RETURN_NAMES = ("florence2_model",)
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "Florence2"
|
||||
|
||||
def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
|
||||
"""Download and load Florence2 model from HuggingFace."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]()
|
||||
return original_loader.loadmodel(model, precision, attention, lora, convert_to_safetensors)
|
||||
|
||||
class CheckpointLoaderNF4:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
||||
}}
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
FUNCTION = "load_checkpoint"
|
||||
|
||||
CATEGORY = "loaders"
|
||||
|
||||
|
||||
def load_checkpoint(self, ckpt_name):
|
||||
"""Load checkpoint in NF4 quantized format."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["CheckpointLoaderNF4"]()
|
||||
return original_loader.load_checkpoint(ckpt_name)
|
||||
|
||||
class LoadFluxControlNet:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"model_name": (["flux-dev", "flux-dev-fp8", "flux-schnell"],),
|
||||
"controlnet_path": (folder_paths.get_filename_list("xlabs_controlnets"), ),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("FluxControlNet",)
|
||||
RETURN_NAMES = ("ControlNet",)
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "XLabsNodes"
|
||||
|
||||
def loadmodel(self, model_name, controlnet_path):
|
||||
"""Load Flux ControlNet model."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["LoadFluxControlNet"]()
|
||||
return original_loader.loadmodel(model_name, controlnet_path)
|
||||
|
||||
class MMAudioModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mmaudio_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from the 'ComfyUI/models/mmaudio' -folder",}),
|
||||
|
||||
"base_precision": (["fp16", "fp32", "bf16"], {"default": "fp16"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MMAUDIO_MODEL",)
|
||||
RETURN_NAMES = ("mmaudio_model", )
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "MMAudio"
|
||||
|
||||
def loadmodel(self, mmaudio_model, base_precision):
|
||||
"""Load MMAudio model with specified precision."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["MMAudioModelLoader"]()
|
||||
return original_loader.loadmodel(mmaudio_model, base_precision)
|
||||
|
||||
class MMAudioFeatureUtilsLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
||||
"synchformer_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
||||
"clip_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
||||
},
|
||||
"optional": {
|
||||
"bigvgan_vocoder_model": ("VOCODER_MODEL", {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}),
|
||||
"mode": (["16k", "44k"], {"default": "44k"}),
|
||||
"precision": (["fp16", "fp32", "bf16"],
|
||||
{"default": "fp16"}
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MMAUDIO_FEATUREUTILS",)
|
||||
RETURN_NAMES = ("mmaudio_featureutils", )
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "MMAudio"
|
||||
|
||||
def loadmodel(self, vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model=None):
|
||||
"""Load MMAudio feature extraction utilities including VAE, Synchformer, and CLIP."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoader"]()
|
||||
return original_loader.loadmodel(vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model)
|
||||
|
||||
class MMAudioSampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mmaudio_model": ("MMAUDIO_MODEL",),
|
||||
"feature_utils": ("MMAUDIO_FEATUREUTILS",),
|
||||
"duration": ("FLOAT", {"default": 8, "step": 0.01, "tooltip": "Duration of the audio in seconds"}),
|
||||
"steps": ("INT", {"default": 25, "step": 1, "tooltip": "Number of steps to interpolate"}),
|
||||
"cfg": ("FLOAT", {"default": 4.5, "step": 0.1, "tooltip": "Strength of the conditioning"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True} ),
|
||||
"negative_prompt": ("STRING", {"default": "", "multiline": True} ),
|
||||
"mask_away_clip": ("BOOLEAN", {"default": False, "tooltip": "If true, the clip video will be masked away"}),
|
||||
"force_offload": ("BOOLEAN", {"default": True, "tooltip": "If true, the model will be offloaded to the offload device"}),
|
||||
},
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AUDIO",)
|
||||
RETURN_NAMES = ("audio", )
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "MMAudio"
|
||||
|
||||
def sample(self, mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images=None):
|
||||
"""Sample audio from MMAudio model with conditioning."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["MMAudioSampler"]()
|
||||
return original_loader.sample(mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images)
|
||||
|
||||
class PulidModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "pulid_file": (folder_paths.get_filename_list("pulid"), )}}
|
||||
|
||||
RETURN_TYPES = ("PULID",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "pulid"
|
||||
|
||||
def load_model(self, pulid_file):
|
||||
"""Load PuLID identity preservation model."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["PulidModelLoader"]()
|
||||
return original_loader.load_model(pulid_file)
|
||||
|
||||
class PulidInsightFaceLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"provider": (["CPU", "CUDA", "ROCM", "CoreML"], ),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FACEANALYSIS",)
|
||||
FUNCTION = "load_insightface"
|
||||
CATEGORY = "pulid"
|
||||
|
||||
def load_insightface(self, provider):
|
||||
"""Load InsightFace face analysis model for PuLID."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["PulidInsightFaceLoader"]()
|
||||
return original_loader.load_insightface(provider)
|
||||
|
||||
class PulidEvaClipLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("EVA_CLIP",)
|
||||
FUNCTION = "load_eva_clip"
|
||||
CATEGORY = "pulid"
|
||||
|
||||
def load_eva_clip(self):
|
||||
"""Load EVA CLIP model for PuLID."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["PulidEvaClipLoader"]()
|
||||
return original_loader.load_eva_clip()
|
||||
|
||||
class UNetLoaderLP:
|
||||
"""UNet Loader (Low Precision) - sets LoRA precision to False for CPU storage optimization"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ),
|
||||
}}
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_unet"
|
||||
CATEGORY = "loaders"
|
||||
TITLE = "UNet Loader (LP)"
|
||||
|
||||
def load_unet(self, unet_name):
|
||||
"""Load UNet with low-precision LoRA flag for CPU storage optimization."""
|
||||
original_loader = NODE_CLASS_MAPPINGS["UNETLoader"]()
|
||||
out = original_loader.load_unet(unet_name)
|
||||
|
||||
# Set the low-precision LoRA flag on the loaded model
|
||||
if hasattr(out[0], 'model'):
|
||||
out[0].model._distorch_high_precision_loras = False
|
||||
elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
|
||||
out[0].patcher.model._distorch_high_precision_loras = False
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,105 @@
|
||||
got prompt
|
||||
loading images: 100%|██████████| 241/241 [00:20<00:00, 11.71it/s]
|
||||
Using pytorch attention in VAE
|
||||
Using pytorch attention in VAE
|
||||
VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16
|
||||
Found quantization metadata version 1
|
||||
[MultiGPU Core Patching] text_encoder_device_patched returning device: cuda:0 (current_text_encoder_device=cuda:0)
|
||||
Using MixedPrecisionOps for text encoder
|
||||
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
|
||||
Requested to load WanTEModel
|
||||
loaded completely; 21474.80 MB usable, 6419.49 MB loaded, full load: True
|
||||
[MultiGPU Core Patching] Successfully patched ModelPatcher.partially_load
|
||||
gguf qtypes: F32 (836), Q4_K (345), Q6_K (144), F16 (6)
|
||||
model weight dtype torch.float16, manual cast: None
|
||||
model_type FLOW
|
||||
Requested to load WanVAE
|
||||
loaded completely; 13567.55 MB usable, 242.03 MB loaded, full load: True
|
||||
Requested to load WAN21_Vace
|
||||
===============================================
|
||||
DisTorch2 Model Virtual VRAM Analysis
|
||||
===============================================
|
||||
Object Role Original(GB) Total(GB) Virt(GB)
|
||||
-----------------------------------------------
|
||||
cuda:0 recip 23.99GB 27.99GB +4.00GB
|
||||
cpu donor 127.12GB 123.12GB -4.00GB
|
||||
-----------------------------------------------
|
||||
model model 10.83GB 6.83GB -4.00GB
|
||||
==================================================
|
||||
[MultiGPU DisTorch V2] Final Allocation String:
|
||||
cuda:0,0.2848;cpu,0.0315
|
||||
==================================================
|
||||
DisTorch2 Model Device Allocations
|
||||
==================================================
|
||||
Device VRAM GB Dev % Model GB Dist %
|
||||
--------------------------------------------------
|
||||
cuda:0 23.99 28.5% 6.83 63.0%
|
||||
cpu 127.12 3.1% 4.00 37.0%
|
||||
--------------------------------------------------
|
||||
!!! Exception during processing !!! too many values to unpack (expected 4)
|
||||
Traceback (most recent call last):
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 515, in execute
|
||||
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 329, in get_output_data
|
||||
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 303, in _async_map_node_over_list
|
||||
await process_inputs(input_dict, i)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 291, in process_inputs
|
||||
result = f(**inputs)
|
||||
^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\ComfyUI-SuperUltimateVaceTools\nodes.py", line 468, in upscale_video
|
||||
sample = nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, conditions['positive'], conditions['negative'], conditions['out_latent'],
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\nodes.py", line 1505, in common_ksampler
|
||||
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\sample.py", line 60, in sample
|
||||
samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\ComfyUI-TiledDiffusion\utils.py", line 51, in KSampler_sample
|
||||
return orig_fn(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 1163, in sample
|
||||
return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 1053, in sample
|
||||
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 1035, in sample
|
||||
output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\patcher_extension.py", line 112, in execute
|
||||
return self.original(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 984, in outer_sample
|
||||
self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds, self.model_options)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\sampler_helpers.py", line 130, in prepare_sampling
|
||||
return executor.execute(model, noise_shape, conds, model_options=model_options)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\patcher_extension.py", line 112, in execute
|
||||
return self.original(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\sampler_helpers.py", line 138, in _prepare_sampling
|
||||
comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required + inference_memory, minimum_memory_required=minimum_memory_required + inference_memory)
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 199, in patched_load_models_gpu
|
||||
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\model_management.py", line 506, in model_load
|
||||
self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\model_management.py", line 536, in model_use_more_vram
|
||||
return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 240, in new_partially_load
|
||||
device_assignments = analyze_safetensor_loading(self, allocations, is_clip=is_clip_model)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 422, in analyze_safetensor_loading
|
||||
total_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 422, in <genexpr>
|
||||
total_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
ValueError: too many values to unpack (expected 4)
|
||||
|
||||
Prompt executed in 61.79 seconds
|
||||
@@ -0,0 +1,145 @@
|
||||
got prompt
|
||||
loading images: 100%|██████████| 81/81 [00:12<00:00, 6.29it/s]
|
||||
loading images: 100%|██████████| 81/81 [00:06<00:00, 12.09it/s]
|
||||
Using pytorch attention in VAE
|
||||
Using pytorch attention in VAE
|
||||
VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16
|
||||
[MultiGPU Core Patching] text_encoder_device_patched returning device: cuda:0 (current_text_encoder_device=cuda:0)
|
||||
Using scaled fp8: fp8 matrix mult: False, scale input: False
|
||||
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
|
||||
Requested to load WanTEModel
|
||||
loaded completely; 21368.67 MB usable, 6419.48 MB loaded, full load: True
|
||||
gguf qtypes: F32 (836), Q4_K (345), Q6_K (144), F16 (6)
|
||||
model weight dtype torch.float16, manual cast: None
|
||||
model_type FLOW
|
||||
Requested to load WanVAE
|
||||
loaded completely; 5713.71 MB usable, 242.03 MB loaded, full load: True
|
||||
Requested to load WAN21_Vace
|
||||
===============================================
|
||||
DisTorch2 Model Virtual VRAM Analysis
|
||||
===============================================
|
||||
Object Role Original(GB) Total(GB) Virt(GB)
|
||||
-----------------------------------------------
|
||||
cuda:0 recip 23.99GB 27.99GB +4.00GB
|
||||
cpu donor 127.12GB 123.12GB -4.00GB
|
||||
-----------------------------------------------
|
||||
model model 10.83GB 6.83GB -4.00GB
|
||||
==================================================
|
||||
[MultiGPU DisTorch V2] Final Allocation String:
|
||||
cuda:0,0.2848;cpu,0.0315
|
||||
==================================================
|
||||
DisTorch2 Model Device Allocations
|
||||
==================================================
|
||||
Device VRAM GB Dev % Model GB Dist %
|
||||
--------------------------------------------------
|
||||
cuda:0 23.99 28.5% 6.83 63.0%
|
||||
cpu 127.12 3.1% 4.00 37.0%
|
||||
--------------------------------------------------
|
||||
!!! Exception during processing !!! too many values to unpack (expected 4)
|
||||
Traceback (most recent call last):
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 510, in execute
|
||||
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 324, in get_output_data
|
||||
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 298, in _async_map_node_over_list
|
||||
await process_inputs(input_dict, i)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 286, in process_inputs
|
||||
result = f(**inputs)
|
||||
^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\ComfyUI-SuperUltimateVaceTools\nodes.py", line 461, in upscale_video
|
||||
sample = nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, conditions['positive'], conditions['negative'], conditions['out_latent'],
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\nodes.py", line 1502, in common_ksampler
|
||||
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\sample.py", line 60, in sample
|
||||
samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\ComfyUI-TiledDiffusion\utils.py", line 51, in KSampler_sample
|
||||
return orig_fn(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 1163, in sample
|
||||
return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 1053, in sample
|
||||
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 1035, in sample
|
||||
output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\patcher_extension.py", line 112, in execute
|
||||
return self.original(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\samplers.py", line 984, in outer_sample
|
||||
self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds, self.model_options)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\sampler_helpers.py", line 130, in prepare_sampling
|
||||
return executor.execute(model, noise_shape, conds, model_options=model_options)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\patcher_extension.py", line 112, in execute
|
||||
return self.original(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\sampler_helpers.py", line 138, in _prepare_sampling
|
||||
comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required + inference_memory, minimum_memory_required=minimum_memory_required + inference_memory)
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 199, in patched_load_models_gpu
|
||||
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\model_management.py", line 506, in model_load
|
||||
self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\comfy\model_management.py", line 536, in model_use_more_vram
|
||||
return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 240, in new_partially_load
|
||||
device_assignments = analyze_safetensor_loading(self, allocations, is_clip=is_clip_model)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 422, in analyze_safetensor_loading
|
||||
total_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\comfyui-multigpu\distorch_2.py", line 422, in <genexpr>
|
||||
total_memory = sum(module_size for module_size, _, _, _ in raw_block_list)
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
ValueError: too many values to unpack (expected 4)
|
||||
|
||||
Prompt executed in 64.32 seconds
|
||||
got prompt
|
||||
Prompt executed in 4.87 seconds
|
||||
got prompt
|
||||
loading images: 100%|██████████| 241/241 [00:31<00:00, 7.54it/s]
|
||||
Prompt executed in 36.81 seconds
|
||||
got prompt
|
||||
Using pytorch attention in VAE
|
||||
Using pytorch attention in VAE
|
||||
VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16
|
||||
[MultiGPU Core Patching] text_encoder_device_patched returning device: cuda:0 (current_text_encoder_device=cuda:0)
|
||||
Using scaled fp8: fp8 matrix mult: False, scale input: False
|
||||
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
|
||||
Requested to load WanTEModel
|
||||
loaded completely; 21368.67 MB usable, 6419.48 MB loaded, full load: True
|
||||
gguf qtypes: F32 (836), Q4_K (345), Q6_K (144), F16 (6)
|
||||
model weight dtype torch.float16, manual cast: None
|
||||
model_type FLOW
|
||||
Requested to load WanVAE
|
||||
loaded completely; 13522.79 MB usable, 242.03 MB loaded, full load: True
|
||||
!!! Exception during processing !!! shape '[241, 42, 8, 73, 8]' is invalid for input of size 47613888
|
||||
Traceback (most recent call last):
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 510, in execute
|
||||
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 324, in get_output_data
|
||||
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 298, in _async_map_node_over_list
|
||||
await process_inputs(input_dict, i)
|
||||
File "C:\Users\ai.machine\AppData\Local\Programs\@comfyorgcomfyui-electron\resources\ComfyUI\execution.py", line 286, in process_inputs
|
||||
result = f(**inputs)
|
||||
^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\ComfyUI-SuperUltimateVaceTools\nodes.py", line 458, in upscale_video
|
||||
conditions = vace_cond_execute(positive, negative, vae, width, height, length_n, 1, strength, input_video=video, control_video=controls, control_masks=mask_ctl, reference_image=refimg, latent_strength_list=None)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "C:\Users\ai.machine\Documents\ComfyUI\custom_nodes\ComfyUI-SuperUltimateVaceTools\nodes.py", line 256, in vace_cond_execute
|
||||
mask = mask.view(length, height_mask, vae_stride, width_mask, vae_stride)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
RuntimeError: shape '[241, 42, 8, 73, 8]' is invalid for input of size 47613888
|
||||
|
||||
Prompt executed in 32.76 seconds
|
||||
Reference in New Issue
Block a user