Compare commits

...
12 changed files with 425 additions and 375 deletions
+24
View File
@@ -33,6 +33,30 @@
---
## Power-Up (DARE+TIES)
> Combines DARE (Drop and Rescale) with TIES (Trim, Elect Sign) — the current best-practice algorithm for merging fine-tuned models, particularly LoRA-tuned ones. DARE sparsifies the task vector via random dropout and rescaling; TIES then removes residual low-magnitude noise and enforces sign consistency. The result is a cleaner, more targeted capability transfer from B to A with less interference.
**Models Used:** A, B
**Parameters:**
- **Alpha:** DARE drop rate. Fraction of delta parameters randomly zeroed. Higher values produce a sparser, more targeted delta. Typical range: `0.5`–`0.9`.
- **Beta:** TIES trim quantile. Fraction of the smallest-magnitude delta parameters zeroed after DARE, removing residual noise. `0.0` disables trimming. Typical range: `0.1`–`0.3`.
- **Gamma:** Lambda scale. Final multiplier applied to the merged delta before adding to Model A. Equivalent to the task arithmetic scaling coefficient. Typical range: `0.5`–`1.5`.
- **Seed:** Random seed for the DARE dropout mask. Controls reproducibility.
**Algorithm:**
1. Compute task vector `δ = B − A`
2. **DARE:** apply random binary mask with keep probability `1 − alpha`; rescale survivors by `1/(1−alpha)`
3. **TIES trim:** zero out parameters below the `beta`-quantile magnitude threshold
4. **TIES elect:** determine dominant sign per position; zero out disagreeing parameters
5. Return `A + gamma × δ_filtered`
**Guidance:**
- Start with `alpha=0.9, beta=0.2, gamma=1.0` for LoRA-derived models
- Lower `alpha` (e.g. `0.5`) for full fine-tunes where more delta mass should be retained
- Increase `gamma` beyond `1.0` to amplify B's influence; decrease below `1.0` to soften it
---
## Enhanced Man Interp
> Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control.
+1 -8
View File
@@ -36,12 +36,7 @@ except ImportError:
# Assuming downloader_utils.py is inside nodes/, the root is one level up
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(root_dir, 'unifiedefficientloader-reference'))
try:
from unifiedefficientloader import UnifiedSafetensorsLoader
HAS_UNIFIED_LOADER = True
except ImportError:
HAS_UNIFIED_LOADER = False
print("Warning: unifiedefficientloader not found. Safetensors metadata extraction disabled.")
from unifiedefficientloader import UnifiedSafetensorsLoader
# --- Database Path Logic ---
try:
@@ -416,8 +411,6 @@ def filter_images_by_nsfw(images, threshold_level):
return valid_images
def extract_safetensors_metadata(filepath):
if not HAS_UNIFIED_LOADER:
return None
try:
# Use low_memory=True to only read the header and metadata without loading tensors
loader = UnifiedSafetensorsLoader(filepath, low_memory=True)
+10 -3
View File
@@ -13,12 +13,15 @@ import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
# =============================================================================
# Rank Index Functions - Determine LoRA rank from singular values
@@ -689,9 +692,13 @@ def _format_lora_key(key: str) -> str:
key = key[:-7]
# Handle ComfyUI Checkpoint wrapper
if key.startswith("model.diffusion_model."):
if key.startswith("model.diffusion_model.") or key.startswith("net."):
# Check if it's a Diffusers-format transformer inside a checkpoint
inner_key = key[22:] # len("model.diffusion_model.")
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # len("model.diffusion_model.")
else:
inner_key = key[4:] # len("net.")
if inner_key.startswith("transformer_blocks") or inner_key.startswith("single_transformer_blocks"):
return f"transformer.{inner_key}"
# Standard checkpoint or Original Flux
+51 -12
View File
@@ -10,8 +10,8 @@ from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from typing import List, Dict, Tuple, Optional
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from .merger_utils import parse_custom_weights, get_custom_parameters
from .lora_resize import (
detect_lora_format,
extract_lora_pairs,
@@ -21,6 +21,10 @@ from .lora_resize import (
_format_lora_key
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
def merge_multi_loras(
lora_paths: List[str],
lora_weights: List[float],
@@ -28,6 +32,7 @@ def merge_multi_loras(
device: str,
save_dtype: torch.dtype,
output_filename: str,
custom_weights: str = "",
verbose: bool = True,
) -> str:
"""
@@ -48,6 +53,8 @@ def merge_multi_loras(
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
parsed_custom_weights = parse_custom_weights(custom_weights)
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
try:
@@ -90,6 +97,9 @@ def merge_multi_loras(
# 2. Merge each core layer
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"):
# Apply per-layer overrides (regex/substring matching)
layer_overrides = get_custom_parameters(core, parsed_custom_weights, {})
downs = []
ups = []
is_conv = False
@@ -97,7 +107,14 @@ def merge_multi_loras(
# First pass: load and determine max rank
for idx in indices:
info = lora_infos[idx]
info = lora_infos[idx].copy()
# Apply overrides: w1, w2, w3... or 'weight'
lora_key = f'w{idx+1}'
if lora_key in layer_overrides:
info["weight"] = layer_overrides[lora_key]
elif "weight" in layer_overrides:
info["weight"] = layer_overrides["weight"]
handler = handlers[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
@@ -275,6 +292,7 @@ class LoRAMultiMerge(io.ComfyNode):
io.Combo.Input("merge_mode", options=["concatenate", "weighted_sum"], default="concatenate",
tooltip="concatenate: safe, increases rank. weighted_sum: fixed rank (to max input rank), mathematically lossy but standard in Kohya."),
io.String.Input("output_filename", default="merged_lora"),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights for each LoRA. Syntax: 'regex:w1,w2,w3...' or 'regex:w1=0.5,w2=1.0'"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
],
@@ -286,7 +304,7 @@ class LoRAMultiMerge(io.ComfyNode):
def execute(cls, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
merge_mode, output_filename, save_dtype, device) -> io.NodeOutput:
merge_mode, output_filename, custom_weights, save_dtype, device) -> io.NodeOutput:
count = int(lora_count)
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
@@ -312,7 +330,8 @@ class LoRAMultiMerge(io.ComfyNode):
merge_mode=merge_mode,
device=device,
save_dtype=dtype,
output_filename=output_filename
output_filename=output_filename,
custom_weights=custom_weights
)
return io.NodeOutput(path)
@@ -353,6 +372,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
io.Float.Input("trim_quantile", default=0.2, min=0.0, max=1.0, step=0.01, tooltip="TIES trim quantile (drops smallest values)"),
io.Int.Input("seed", default=42, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_lora_dare"),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights for each LoRA. Syntax: 'regex:w1,w2,w3...' or 'regex:w1=0.5,w2=1.0'"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
],
@@ -364,7 +384,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
def execute(cls, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput:
drop_rate, trim_quantile, seed, output_filename, custom_weights, save_dtype, device) -> io.NodeOutput:
count = int(lora_count)
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
@@ -392,7 +412,8 @@ class LoRAMultiMergeDARE(io.ComfyNode):
seed=seed,
device=device,
save_dtype=dtype,
output_filename=output_filename
output_filename=output_filename,
custom_weights=custom_weights
)
return io.NodeOutput(path)
@@ -406,6 +427,7 @@ def merge_multi_loras_dare(
device: str,
save_dtype: torch.dtype,
output_filename: str,
custom_weights: str = "",
verbose: bool = True,
) -> str:
"""
@@ -418,6 +440,9 @@ def merge_multi_loras_dare(
print(f"[LoRA Multi-Merge DARE] Drop rate: {drop_rate}, Trim quantile: {trim_quantile}")
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
parsed_custom_weights = parse_custom_weights(custom_weights)
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
rng = torch.Generator(device=device).manual_seed(seed)
@@ -457,12 +482,22 @@ def merge_multi_loras_dare(
# 2. Merge
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"):
# Apply per-layer overrides (regex/substring matching)
layer_overrides = get_custom_parameters(core, parsed_custom_weights, {})
downs = []
ups = []
max_rank = 0
for idx in indices:
info = lora_infos[idx]
info = lora_infos[idx].copy()
# Apply overrides: w1, w2, w3... or 'weight'
lora_key = f'w{idx+1}'
if lora_key in layer_overrides:
info["weight"] = layer_overrides[lora_key]
elif "weight" in layer_overrides:
info["weight"] = layer_overrides["weight"]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
@@ -491,18 +526,22 @@ def merge_multi_loras_dare(
t = torch.nn.functional.pad(t, tuple(padding))
padded.append(t * w)
# Apply per-layer overrides for DARE/TIES parameters
current_drop_rate = layer_overrides.get('alpha', layer_overrides.get('drop_rate', drop_rate))
current_trim_quantile = layer_overrides.get('beta', layer_overrides.get('trim_quantile', trim_quantile))
# DARE
if drop_rate > 0:
if current_drop_rate > 0:
for i in range(len(padded)):
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float()
padded[i] = (padded[i] * mask) / (1 - drop_rate)
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > current_drop_rate).float()
padded[i] = (padded[i] * mask) / (1 - current_drop_rate)
# TIES
# 1. Trim
if trim_quantile > 0:
if current_trim_quantile > 0:
for i in range(len(padded)):
flat = padded[i].abs().flatten()
k = int(len(flat) * trim_quantile)
k = int(len(flat) * current_trim_quantile)
if k > 0:
threshold = torch.kthvalue(flat, k).values
padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i])
+5 -2
View File
@@ -14,9 +14,12 @@ from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from typing import Optional, Dict, Tuple, List
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
# Reuse rank functions from extraction module
from .lora_extract_svd import (
_index_sv_ratio,
@@ -189,7 +192,7 @@ def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int
# Get dim from down weight shape
if network_dim is None:
down_key = block_keys["down"]
shape = handler.header[down_key]["shape"]
shape = handler.get_shape(down_key)
# Linear: [rank, in_features] or Conv: [rank, in_ch, k, k]
network_dim = shape[0]
+142 -56
View File
@@ -7,11 +7,13 @@ from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .merger_utils import parse_custom_weights, get_custom_parameters
from .device_utils import (
estimate_model_size, prepare_for_large_operation, cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
def load_documentation_from_file(filename):
"""Loads documentation from a markdown file in the ../docs/ directory."""
@@ -76,6 +78,9 @@ class MergerLogic:
process_device = recipe_params.get('device', 'cpu')
print(f"[Merger] Preparing memory for {total_size_gb:.2f}GB merge operation...")
prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device))
flush_threshold_bytes = int(total_size_gb * 1024**3 / 16)
else:
flush_threshold_bytes = 256 * 1024 * 1024 # 256MB fallback (min 8GB RAM assumed)
handlers = {}
for name in model_names.values():
@@ -99,6 +104,9 @@ class MergerLogic:
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''))
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''))
# Parse custom per-layer weights
parsed_custom_weights = parse_custom_weights(recipe_params.get('custom_weights', ''))
# Pre-compute key differences for logging
primary_keys = set(all_keys)
for name, handler in handlers.items():
@@ -124,12 +132,42 @@ class MergerLogic:
discarded_keys = 0
error_keys = []
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
output_dir = folder_paths.get_folder_paths(output_folder)[-1]
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
writer = None
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
batch_buffer = {}
batch_bytes = 0
batch_key_count = 0
def _flush_batch():
nonlocal batch_buffer, batch_bytes, batch_key_count
if not batch_buffer:
return
writer.write_dict(batch_buffer)
pbar.update(batch_key_count)
batch_buffer = {}
batch_bytes = 0
batch_key_count = 0
if recipe_params.get('force_clear_cache', True):
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
pbar.update(1)
continue
# Pre-load Model A's tensor with pinned memory for CUDA
@@ -142,9 +180,15 @@ class MergerLogic:
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
merged_state_dict[key] = tensor_a.detach().to(save_torch_dtype).cpu().clone()
t = tensor_a.detach().to(save_torch_dtype).cpu()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
excluded_keys += 1
pbar.update(1)
del tensor_a
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
continue
# Pass tensor_a metadata to recipes for zeros mode and fallback
@@ -153,10 +197,15 @@ class MergerLogic:
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
try:
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
# Apply per-layer parameter overrides (regex/substring matching)
layer_params = get_custom_parameters(key, parsed_custom_weights, recipe_params)
recipe = calc_mode_class.create_recipe(key=key, **layer_params)
result = recipe.merge()
except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR:
_flush_batch()
writer.__exit__(None, None, None)
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
result = None
error_keys.append(key)
@@ -168,7 +217,10 @@ class MergerLogic:
if isinstance(result, dict):
for r_key, r_tensor in result.items():
merged_state_dict[r_key] = r_tensor.detach().to(save_torch_dtype).cpu().clone()
t = r_tensor.detach().to(save_torch_dtype).cpu()
batch_buffer[r_key] = t
batch_bytes += t.numel() * t.element_size()
else:
# Ensure compatibility with Model A's architecture.
# If alignment_mode is 'pad/crop', we crop results that were padded.
@@ -179,7 +231,13 @@ class MergerLogic:
slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape))
result = result[slices]
merged_state_dict[key] = result.detach().to(save_torch_dtype).cpu().clone()
t = result.detach().to(save_torch_dtype).cpu()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
# Clean up references to allow GC immediately.
# Local loop variables must be explicitly deleted to prevent PyTorch from keeping tensors in VRAM.
@@ -189,13 +247,12 @@ class MergerLogic:
del recipe
del result
if recipe_params.get('force_clear_cache', True):
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
pbar.update(1)
_flush_batch()
writer.__exit__(None, None, None)
# Log summary
if excluded_keys > 0:
@@ -216,10 +273,9 @@ class MergerLogic:
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
save_file(merged_state_dict, output_path, metadata=metadata)
if os.path.exists(output_path):
print(f"[Merger] Output saved to {output_path}")
# Cleanup after heavy operation
del merged_state_dict
cleanup_after_operation()
return f"{output_filename}.safetensors"
@@ -243,16 +299,17 @@ class CheckpointTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -267,6 +324,7 @@ class CheckpointTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -279,6 +337,7 @@ class CheckpointTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -302,16 +361,17 @@ class ModelTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -326,6 +386,7 @@ class ModelTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -338,6 +399,7 @@ class ModelTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -361,16 +423,17 @@ class TextEncoderTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -385,6 +448,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -397,6 +461,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -420,16 +485,17 @@ class LoRATwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -444,6 +510,7 @@ class LoRATwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -456,6 +523,7 @@ class LoRATwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -479,16 +547,17 @@ class EmbeddingTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -503,6 +572,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -515,6 +585,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
@@ -541,16 +612,17 @@ class CheckpointThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -565,6 +637,7 @@ class CheckpointThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -577,6 +650,7 @@ class CheckpointThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -601,16 +675,17 @@ class ModelThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -625,6 +700,7 @@ class ModelThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -637,6 +713,7 @@ class ModelThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -661,16 +738,17 @@ class TextEncoderThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -685,6 +763,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -697,6 +776,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -721,16 +801,17 @@ class LoRAThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -745,6 +826,7 @@ class LoRAThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -757,6 +839,7 @@ class LoRAThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
@@ -781,16 +864,17 @@ class EmbeddingThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.String.Input("custom_weights", default="", multiline=True, tooltip="Per-layer weights/params. Syntax: 'regex:weight' or 'regex:param=val,param2=val'"),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
@@ -805,6 +889,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
custom_weights: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
@@ -817,6 +902,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"custom_weights": custom_weights,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
+85 -2
View File
@@ -4,7 +4,7 @@ import torch.nn.functional as F
from enum import Enum
from .merger_utils import transfer_to_gpu_pinned
from unifiedefficientloader import transfer_to_gpu_pinned
class MissingTensorBehavior(Enum):
"""Controls behavior when a tensor key is missing from a model."""
@@ -319,6 +319,53 @@ class PowerUpOp(Operation):
return (m * delta) / keep_prob
class DARETIESOp(Operation):
"""DARE + TIES combined operation on a single task vector (delta = b - a).
Step 1 — DARE: randomly drop `drop_rate` fraction of delta parameters,
rescale survivors by 1/(1-drop_rate).
Step 2 — TIES trim: zero out parameters below the `trim_quantile` magnitude threshold.
Step 3 — TIES elect: determine dominant sign per parameter position.
Step 4 — TIES merge: zero out parameters that disagree with the dominant sign.
For a two-model merge the TIES elect/merge step degenerates to a no-op (single vector,
sign is always self-consistent) but the trim step still reduces interference from
near-zero noisy deltas. The dominant sign filter is included for correctness and
forward compatibility with n-model extensions.
"""
def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources)
# alpha = drop_rate (DARE), beta = trim_quantile (TIES), gamma = lambda scale
self.alpha, self.beta, self.gamma, self.seed = alpha, beta, gamma, seed
def oper(self, a, b):
if a is None or b is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta = (b - a).float()
# --- DARE: drop and rescale ---
keep_prob = 1.0 - self.alpha
if keep_prob <= 0:
return torch.zeros_like(a)
rng = torch.Generator(device=delta.device).manual_seed(self.seed)
dare_mask = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).float()
delta = (delta * dare_mask) / keep_prob
# --- TIES trim: zero parameters below magnitude threshold ---
if self.beta > 0:
flat_abs = delta.abs().flatten()
k = max(1, int(flat_abs.numel() * self.beta))
threshold = torch.kthvalue(flat_abs, k).values
delta = torch.where(delta.abs() < threshold, torch.zeros_like(delta), delta)
# --- TIES elect: dominant sign (trivially self-consistent for single vector) ---
dominant_sign = torch.sign(delta)
delta = torch.where(torch.sign(delta) == dominant_sign, delta, torch.zeros_like(delta))
return (a + self.gamma * delta).to(a.dtype)
class InterpolateDifference(Operation):
def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources)
@@ -710,5 +757,41 @@ class WeightSumCutoffMode(CalcMode):
return res
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
class DARETIESMode(CalcMode):
name = 'Power-Up (DARE+TIES)'
description = 'Adds capabilities of B to A using DARE (Drop and Rescale) followed by TIES (Trim, Elect Sign). Best practice for LoRA-tuned model merging.'
models_used = ['A', 'B']
param_docs = {
'alpha': 'DARE drop rate — fraction of delta parameters randomly zeroed. Higher = more aggressive sparsification.',
'beta': 'TIES trim quantile — fraction of smallest-magnitude delta parameters zeroed after DARE. Removes residual noise.',
'gamma': 'Lambda — scale multiplier applied to the final merged delta before adding to Model A.',
'seed': 'Random seed for the DARE dropout mask.',
}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
op = DARETIESOp(
key,
kwargs['alpha'], # drop_rate
kwargs['beta'], # trim_quantile
kwargs['gamma'], # lambda scale
kwargs['seed'],
a, b,
)
op.alignment_mode = alignment_mode
return op
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
+85 -288
View File
@@ -1,292 +1,8 @@
"""
Memory-efficient safetensors utilities with pinned memory support.
Forces direct file reads to minimize OS page cache usage and avoids VRAM accumulation.
Utilities for model merging operations.
"""
import os
import gc
import json
import struct
import torch
from safetensors import safe_open
from typing import Dict, List, Optional, Any
# Module-level configuration for pinned transfer
_verbose = False
_pinned_transfer_stats = {"pinned": 0, "fallback": 0}
def set_verbose(enabled: bool):
"""Enable/disable verbose output for pinned transfers."""
global _verbose
_verbose = enabled
def get_pinned_transfer_stats():
"""Return pinned transfer statistics."""
return _pinned_transfer_stats.copy()
def reset_pinned_transfer_stats():
"""Reset transfer statistics."""
global _pinned_transfer_stats
_pinned_transfer_stats = {"pinned": 0, "fallback": 0}
def transfer_to_gpu_pinned(
tensor: torch.Tensor,
device: str = 'cuda',
dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
"""Transfer tensor to GPU using pinned memory for faster transfer."""
global _pinned_transfer_stats
if tensor.device.type != 'cpu' or not torch.cuda.is_available():
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
if not str(device).startswith('cuda'):
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
try:
pinned = tensor.pin_memory()
if dtype is not None:
result = pinned.to(device=device, dtype=dtype, non_blocking=True)
else:
result = pinned.to(device=device, non_blocking=True)
torch.cuda.current_stream().synchronize()
_pinned_transfer_stats["pinned"] += 1
return result
except Exception as e:
_pinned_transfer_stats["fallback"] += 1
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
class MemoryEfficientSafeOpen:
"""Memory-efficient safetensors file reader with streaming and pinned memory support.
Operates strictly in streaming mode (low memory). Tensors are read byte-by-byte from
their exact offsets without using safe_open or memory mapping, ensuring tensors do not
remain persistently in OS page caches or RAM.
"""
def __init__(self, filename: str, device: str = 'cpu', **kwargs):
# Ignore mmap_mode, low_memory, or any other legacy kwargs. We always stream.
self.filename = filename
self.device = device
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.file.close()
def keys(self) -> List[str]:
"""Return all tensor keys (excluding metadata)."""
return [k for k in self.header.keys() if k != "__metadata__"]
def metadata(self) -> Dict[str, str]:
"""Return file metadata."""
return self.header.get("__metadata__", {})
def get_shape(self, key: str) -> tuple:
"""Get tensor shape without loading tensor data."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in file")
return tuple(self.header[key]["shape"])
def get_ndim(self, key: str) -> int:
"""Get tensor ndim without loading tensor data."""
return len(self.get_shape(key))
def keys_sorted_by_offset(self) -> List[str]:
keys_with_offsets = []
for key in self.keys():
offset = self.header[key]["data_offsets"][0]
keys_with_offsets.append((key, offset))
keys_with_offsets.sort(key=lambda x: x[1])
return [k for k, _ in keys_with_offsets]
def get_tensor(self, key: str) -> torch.Tensor:
"""Load a single tensor by key directly from file byte offset."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = bytearray(offset_end - offset_start)
self.file.readinto(tensor_bytes)
else:
tensor_bytes = None
return self._deserialize_tensor(tensor_bytes, metadata)
def get_tensor_to_gpu(
self,
key: str,
device: str = 'cuda',
dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
cpu_tensor = self.get_tensor(key)
return transfer_to_gpu_pinned(cpu_tensor, device, dtype)
def get_tensors_parallel(
self,
keys: Optional[List[str]] = None,
workers: Optional[int] = None,
sort_by_offset: bool = True
) -> Dict[str, torch.Tensor]:
if keys is None:
keys = self.keys()
if workers is None:
try:
from .device_utils import get_optimal_workers, estimate_model_size
model_size = estimate_model_size(self.filename)
workers = get_optimal_workers(model_size)
except ImportError:
workers = min(4, os.cpu_count() or 4)
if sort_by_offset:
keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys]
keys_with_offsets.sort(key=lambda x: x[1])
keys = [k for k, _ in keys_with_offsets]
result = {}
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(
_read_tensor_from_file,
self.filename,
key,
self.header[key],
self.header_size
): key
for key in keys
}
for future in as_completed(futures):
key = futures[future]
try:
tensor = future.result()
if tensor is not None:
result[key] = tensor
except Exception as e:
print(f"[MemoryEfficientSafeOpen] Failed to load {key}: {e}")
return result
def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
if offset_start == offset_end:
return {}
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
return json.loads(tensor_bytes.decode("utf-8"))
def _read_header(self):
with open(self.filename, "rb") as f:
header_size = struct.unpack("<Q", f.read(8))[0]
header_json = f.read(header_size).decode("utf-8")
return json.loads(header_json), header_size
def _deserialize_tensor(self, tensor_bytes, metadata):
dtype_str = metadata["dtype"]
shape = metadata["shape"]
dtype = self._get_torch_dtype(dtype_str)
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
return self._convert_float8(byte_tensor, dtype_str, shape)
return byte_tensor.view(dtype).reshape(shape)
@staticmethod
def _get_torch_dtype(dtype_str):
dtype_map = {
"F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
"I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
"U8": torch.uint8, "BOOL": torch.bool,
}
if hasattr(torch, "float8_e5m2"):
dtype_map["F8_E5M2"] = torch.float8_e5m2
if hasattr(torch, "float8_e4m3fn"):
dtype_map["F8_E4M3"] = torch.float8_e4m3fn
dtype = dtype_map.get(dtype_str)
if dtype is None:
raise ValueError(f"Unsupported dtype: {dtype_str}")
return dtype
@staticmethod
def _convert_float8(byte_tensor, dtype_str, shape):
if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
return byte_tensor.view(torch.float8_e5m2).reshape(shape)
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
else:
raise ValueError(f"Unsupported float8 type: {dtype_str}")
def _read_tensor_from_file(
filepath: str,
key: str,
metadata: Dict,
header_size: int
) -> Optional[torch.Tensor]:
"""Helper function to read a single tensor from file (for parallel execution).
Each worker opens its own file handle for thread safety.
"""
offset_start, offset_end = metadata["data_offsets"]
dtype_str = metadata["dtype"]
shape = metadata["shape"]
if offset_start == offset_end:
return None
dtype_map = {
"F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
"I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
"U8": torch.uint8, "BOOL": torch.bool,
}
if hasattr(torch, "float8_e5m2"):
dtype_map["F8_E5M2"] = torch.float8_e5m2
if hasattr(torch, "float8_e4m3fn"):
dtype_map["F8_E4M3"] = torch.float8_e4m3fn
dtype = dtype_map.get(dtype_str, torch.float32)
with open(filepath, "rb") as f:
f.seek(header_size + 8 + offset_start)
tensor_bytes = f.read(offset_end - offset_start)
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
return byte_tensor.view(torch.float8_e5m2).reshape(shape)
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
return byte_tensor.view(dtype).reshape(shape)
from typing import Dict, List, Tuple, Any
import re
def get_union_keys(*handlers) -> List[str]:
"""Returns the union of all tensor keys across multiple model handlers."""
@@ -305,4 +21,85 @@ def get_intersection_keys(*handlers) -> List[str]:
all_keys = set(handlers[0].keys())
for handler in handlers[1:]:
all_keys &= set(handler.keys())
return sorted(all_keys)
return sorted(all_keys)
def parse_custom_weights(custom_weights_str: str) -> List[Tuple[re.Pattern, Dict[str, float]]]:
"""Parses multiline string of custom weights.
Format: pattern:weight OR pattern:param=value,param2=value
Returns list of (compiled_regex, parameter_dict).
"""
if not custom_weights_str or not custom_weights_str.strip():
return []
results = []
lines = custom_weights_str.strip().split('\n')
for line in lines:
line = line.strip()
if not line or line.startswith('#'):
continue
if ':' not in line:
continue
pattern_str, weight_str = line.rsplit(':', 1)
pattern_str = pattern_str.strip()
weight_str = weight_str.strip()
try:
pattern = re.compile(pattern_str)
except re.error:
# Fallback to literal substring if regex invalid
pattern = re.compile(re.escape(pattern_str))
params = {}
if '=' in weight_str:
# param=value,param2=value
parts = weight_str.split(',')
for part in parts:
if '=' in part:
k, v = part.split('=', 1)
k = k.strip()
# Map common aliases
if k in ['a', 'alpha']: k = 'alpha'
elif k in ['b', 'beta']: k = 'beta'
elif k in ['g', 'gamma']: k = 'gamma'
elif k in ['d', 'delta']: k = 'delta'
elif k in ['w', 'weight']: k = 'weight'
elif re.match(r'^w[1-8]$', k): pass # Keep w1-w8 as is
try:
params[k] = float(v.strip())
except ValueError:
pass
else:
# single float or comma-separated list of floats (for multi-LoRA)
if ',' in weight_str:
parts = weight_str.split(',')
for i, p in enumerate(parts):
try:
params[f'w{i+1}'] = float(p.strip())
except ValueError:
pass
else:
try:
val = float(weight_str)
params['alpha'] = val # default for merger
params['weight'] = val # default for lora_merger
except ValueError:
pass
if params:
results.append((pattern, params))
return results
def get_custom_parameters(key: str, parsed_weights: List[Tuple[re.Pattern, Dict[str, float]]], base_params: Dict[str, Any]) -> Dict[str, Any]:
"""Returns updated parameters if key matches any pattern."""
new_params = base_params.copy()
for pattern, custom_dict in parsed_weights:
if pattern.search(key):
new_params.update(custom_dict)
break # First match wins
return new_params
+2 -1
View File
@@ -7,9 +7,10 @@ from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .merger_utils import MemoryEfficientSafeOpen
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
use_regex: bool, output_filename: str) -> str:
+3 -1
View File
@@ -6,9 +6,11 @@ from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .merger_utils import MemoryEfficientSafeOpen
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
new_keys_str: str, output_filename: str) -> str:
+16 -1
View File
@@ -1,9 +1,24 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.2.9"
version = "0.3.5"
license = { file = "LICENSE" }
classifiers = [
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"unifiedefficientloader==0.4.4"
]
[project.urls]
Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils"
# Used by Comfy Registry https://comfyregistry.org
+1 -1
View File
@@ -1,5 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader
unifiedefficientloader==0.4.4
av