feat(lora): add full-delta CWB merge
This commit is contained in:
silveroxides
@@ -21,6 +21,7 @@ from .nodes.merger import (
|
||||
CheckpointThreeMerger, EmbeddingThreeMerger
|
||||
)
|
||||
from .nodes.consensus_merger import CWB_MERGER_NODES
|
||||
from .nodes.cwb_delta_lora_merger import DELTA_CWB_LORA_NODES
|
||||
from .nodes.lodestone_merger import LODESTONE_MERGER_NODES
|
||||
from .nodes.model_analysis import MODEL_ANALYSIS_NODES
|
||||
from .nodes.lora_extract_svd import (
|
||||
@@ -82,6 +83,8 @@ class ModelUtilsExtension(ComfyExtension):
|
||||
CheckpointThreeMerger, EmbeddingThreeMerger,
|
||||
# Dedicated CWB Mergers
|
||||
*CWB_MERGER_NODES,
|
||||
# Dedicated Delta CWB LoRA Mergers
|
||||
*DELTA_CWB_LORA_NODES,
|
||||
# Dedicated Lodestone LoRA Mergers
|
||||
*LODESTONE_MERGER_NODES,
|
||||
# Two-Model Analysis
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# Delta CWB LoRA Mergers
|
||||
|
||||
Delta CWB expands each selected LoRA into full per-layer weight changes and
|
||||
applies CWB to corresponding dense rows. It preserves model coordinates; it
|
||||
does not align, reorder, or compress rank components.
|
||||
|
||||
LoRA 1 supplies metadata and defines the output layer set. The preset controls
|
||||
how agreement, disagreement, and row magnitude affect the saved changes. A
|
||||
connected CWB configuration completely replaces the preset. Inputs have equal
|
||||
prior weight.
|
||||
|
||||
Missing-layer handling can preserve LoRA 1, include explicit zero contributors,
|
||||
or abort. Excluded layers preserve LoRA 1 and discarded layers are omitted.
|
||||
Outputs contain only canonical `.diff` tensors with alpha already normalized.
|
||||
|
||||
Processing streams one logical layer from each input at a time and writes each
|
||||
completed layer incrementally. CUDA out-of-memory retries only the failed layer
|
||||
on CPU. The destination is replaced only after the complete output is written.
|
||||
@@ -0,0 +1,699 @@
|
||||
import fnmatch
|
||||
import gc
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
from tqdm import tqdm
|
||||
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
|
||||
|
||||
from .artifact_paths import canonical_model_artifact_path
|
||||
from .consensus_merger import (
|
||||
CWB_CONFIG,
|
||||
CWBSettings,
|
||||
CWBDiagnostics,
|
||||
DENSE_CWB_PRESETS,
|
||||
merge_cwb_tensors,
|
||||
resolve_cwb_settings,
|
||||
)
|
||||
from .device_utils import (
|
||||
cleanup_after_operation,
|
||||
estimate_model_size,
|
||||
prepare_for_large_operation,
|
||||
)
|
||||
from .lora_alpha import normalize_lora_pair
|
||||
from .lora_resize import (
|
||||
build_lora_layer_map,
|
||||
canonical_lora_key,
|
||||
parse_lora_layers,
|
||||
select_output_dtype,
|
||||
validate_canonical_blocks,
|
||||
)
|
||||
from .quantization_guard import inspect_low_bit_input
|
||||
from .uel_io import AsyncTensorCursor, atomic_uel_writer
|
||||
|
||||
|
||||
_OPERATION = "Delta CWB LoRA Merge"
|
||||
_ALLOWED_ROLES = {
|
||||
"down", "up", "alpha", "diff", "format", "down_suffix", "up_suffix",
|
||||
"alpha_suffix",
|
||||
}
|
||||
_FACTOR_SUFFIXES = (
|
||||
".lora_A.default.weight", ".lora_B.default.weight",
|
||||
".lora_down.weight", ".lora_up.weight",
|
||||
"_lora.down.weight", "_lora.up.weight",
|
||||
".lora_A.weight", ".lora_B.weight",
|
||||
".lora.down.weight", ".lora.up.weight",
|
||||
".lora_A", ".lora_B",
|
||||
".lora_linear_layer.down.weight", ".lora_linear_layer.up.weight",
|
||||
)
|
||||
_UNSUPPORTED_SUFFIXES = (
|
||||
".lora_mid.weight", ".reshape_weight", ".dora_scale", ".w_norm",
|
||||
".b_norm", ".set_weight",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DeltaCWBLoRASource:
|
||||
source_index: int
|
||||
block_name: str
|
||||
down_key: str | None
|
||||
up_key: str | None
|
||||
alpha_key: str | None
|
||||
diff_key: str | None
|
||||
source_dtype: torch.dtype
|
||||
|
||||
@property
|
||||
def keys(self):
|
||||
return tuple(
|
||||
key for key in (self.down_key, self.up_key, self.alpha_key, self.diff_key)
|
||||
if key is not None
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DeltaCWBLayerUnit:
|
||||
core: str
|
||||
output_key: str
|
||||
sources: tuple[DeltaCWBLoRASource, ...]
|
||||
zero_contributors: int
|
||||
preserve_anchor: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeltaCWBReport:
|
||||
merged_layers: int = 0
|
||||
anchor_preserved_layers: int = 0
|
||||
zero_filled_contributors: int = 0
|
||||
rejected_unsupported_layers: int = 0
|
||||
cuda_cpu_fallbacks: int = 0
|
||||
output_dtypes: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
def record_dtype(self, dtype):
|
||||
name = str(dtype).removeprefix("torch.")
|
||||
self.output_dtypes[name] = self.output_dtypes.get(name, 0) + 1
|
||||
|
||||
def render(self, diagnostics: CWBDiagnostics, preset_name: str, input_count: int, custom: bool):
|
||||
dtype_summary = ", ".join(
|
||||
f"{name}={count}" for name, count in sorted(self.output_dtypes.items())
|
||||
) or "none"
|
||||
source = "connected custom configuration" if custom else preset_name
|
||||
return "\n".join([
|
||||
"DELTA CWB LORA MERGE SUMMARY",
|
||||
f"Settings: {source}",
|
||||
f"Input contributors: {input_count}",
|
||||
f"Merged layers: {self.merged_layers}",
|
||||
f"Anchor-preserved layers: {self.anchor_preserved_layers}",
|
||||
f"Explicit zero contributors: {self.zero_filled_contributors}",
|
||||
f"Rejected/unsupported layers: {self.rejected_unsupported_layers}",
|
||||
f"Accepted weighting contributors: {diagnostics.accepted_weighting_contributors}/{diagnostics.weighting_contributors}",
|
||||
f"All-rejected equal fallbacks: {diagnostics.all_rejected_fallbacks}",
|
||||
f"Zero-weight equal fallbacks: {diagnostics.zero_weight_fallbacks}",
|
||||
f"CUDA OOM CPU fallbacks: {self.cuda_cpu_fallbacks}",
|
||||
f"Output dtypes: {dtype_summary}",
|
||||
])
|
||||
|
||||
|
||||
def _compile_patterns(patterns, glob_patterns):
|
||||
entries = [entry.strip() for entry in patterns.splitlines() if entry.strip()]
|
||||
if glob_patterns:
|
||||
return tuple(entries)
|
||||
return tuple(re.compile(entry) for entry in entries)
|
||||
|
||||
|
||||
def _matches_pattern(key, patterns, glob_patterns):
|
||||
if glob_patterns:
|
||||
return any(fnmatch.fnmatchcase(key, pattern) for pattern in patterns)
|
||||
return any(pattern.search(key) is not None for pattern in patterns)
|
||||
|
||||
|
||||
def _factor_delta(down, up, layer):
|
||||
if down.ndim == 2 and up.ndim == 2:
|
||||
if down.shape[0] != up.shape[1]:
|
||||
raise ValueError(f"Delta CWB rank mismatch for '{layer}'.")
|
||||
return up.float() @ down.float()
|
||||
if down.ndim == 4 and up.ndim == 4 and tuple(up.shape[2:]) == (1, 1):
|
||||
rank, channels, height, width = down.shape
|
||||
if up.shape[1] != rank:
|
||||
raise ValueError(f"Delta CWB rank mismatch for '{layer}'.")
|
||||
delta = up.reshape(up.shape[0], rank).float() @ down.reshape(rank, -1).float()
|
||||
return delta.reshape(up.shape[0], channels, height, width)
|
||||
raise ValueError(
|
||||
f"Delta CWB layer '{layer}' requires linear factors or convolutional "
|
||||
"factors with a 1x1 up kernel."
|
||||
)
|
||||
|
||||
|
||||
def _inspect_inputs(handlers):
|
||||
infos = []
|
||||
for index, handler in enumerate(handlers):
|
||||
low_bit = inspect_low_bit_input(handler, f"LoRA {index + 1}", _OPERATION)
|
||||
keys = list(handler.keys())
|
||||
pairs, passthrough = parse_lora_layers(keys)
|
||||
unsupported_passthrough = [
|
||||
key for key in passthrough if key.endswith(_UNSUPPORTED_SUFFIXES)
|
||||
]
|
||||
if unsupported_passthrough:
|
||||
raise ValueError(
|
||||
f"[{_OPERATION}] LoRA {index + 1} contains unsupported companion "
|
||||
f"tensor '{unsupported_passthrough[0]}'."
|
||||
)
|
||||
incomplete = [key for key in passthrough if key.endswith(_FACTOR_SUFFIXES)]
|
||||
if incomplete:
|
||||
raise ValueError(
|
||||
f"[{_OPERATION}] LoRA {index + 1} contains an incomplete factor "
|
||||
f"pair at '{incomplete[0]}'."
|
||||
)
|
||||
infos.append({"pairs": pairs, "passthrough": passthrough, "low_bit": low_bit})
|
||||
return infos
|
||||
|
||||
|
||||
def _build_layer_units(
|
||||
layer_map,
|
||||
infos,
|
||||
handlers,
|
||||
mismatch_mode,
|
||||
exclude_patterns,
|
||||
discard_patterns,
|
||||
glob_patterns,
|
||||
):
|
||||
units = []
|
||||
stream_keys = [[] for _ in handlers]
|
||||
rejected = 0
|
||||
for core, mapped in layer_map.items():
|
||||
if not any(index == 0 for index, _ in mapped):
|
||||
continue
|
||||
output_key = canonical_lora_key(core, "diff")
|
||||
if _matches_pattern(output_key, discard_patterns, glob_patterns):
|
||||
rejected += 1
|
||||
continue
|
||||
preserve = _matches_pattern(output_key, exclude_patterns, glob_patterns)
|
||||
candidates = []
|
||||
shapes = {}
|
||||
for index, block in sorted(mapped):
|
||||
roles = infos[index]["pairs"][block]
|
||||
unsupported = set(roles) - _ALLOWED_ROLES
|
||||
if unsupported:
|
||||
raise ValueError(
|
||||
f"[{_OPERATION}] Layer '{block}' has unsupported roles: "
|
||||
f"{', '.join(sorted(unsupported))}."
|
||||
)
|
||||
if "down" in roles or "up" in roles:
|
||||
if not all(role in roles for role in ("down", "up")):
|
||||
raise ValueError(f"[{_OPERATION}] Layer '{block}' has an incomplete factor pair.")
|
||||
down_key, up_key = roles["down"], roles["up"]
|
||||
if down_key in infos[index]["low_bit"] or up_key in infos[index]["low_bit"]:
|
||||
raise ValueError(f"[{_OPERATION}] Cannot expand low-bit factors for '{block}'.")
|
||||
down_shape = tuple(handlers[index].get_shape(down_key))
|
||||
up_shape = tuple(handlers[index].get_shape(up_key))
|
||||
if len(down_shape) == 2 and len(up_shape) == 2:
|
||||
if down_shape[0] != up_shape[1]:
|
||||
raise ValueError(f"[{_OPERATION}] Rank mismatch for '{block}'.")
|
||||
dense_shape = (up_shape[0], down_shape[1])
|
||||
elif len(down_shape) == 4 and len(up_shape) == 4 and up_shape[2:] == (1, 1):
|
||||
if down_shape[0] != up_shape[1]:
|
||||
raise ValueError(f"[{_OPERATION}] Rank mismatch for '{block}'.")
|
||||
dense_shape = (up_shape[0], *down_shape[1:])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"[{_OPERATION}] Layer '{block}' requires linear factors or "
|
||||
"convolutional factors with a 1x1 up kernel."
|
||||
)
|
||||
factor_dtypes = (
|
||||
handlers[index].get_dtype(down_key),
|
||||
handlers[index].get_dtype(up_key),
|
||||
)
|
||||
source_dtype = torch.float32 if torch.float32 in factor_dtypes else factor_dtypes[1]
|
||||
source = DeltaCWBLoRASource(
|
||||
index, block, down_key, up_key, roles.get("alpha"), None, source_dtype
|
||||
)
|
||||
elif "diff" in roles:
|
||||
diff_key = roles["diff"]
|
||||
if diff_key in infos[index]["low_bit"]:
|
||||
raise ValueError(f"[{_OPERATION}] Cannot merge low-bit delta '{block}'.")
|
||||
dense_shape = tuple(handlers[index].get_shape(diff_key))
|
||||
source = DeltaCWBLoRASource(
|
||||
index, block, None, None, None, diff_key,
|
||||
handlers[index].get_dtype(diff_key),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"[{_OPERATION}] Layer '{block}' has no supported representation.")
|
||||
candidates.append(source)
|
||||
shapes[index] = dense_shape
|
||||
|
||||
anchor = next((source for source in candidates if source.source_index == 0), None)
|
||||
if anchor is None:
|
||||
continue
|
||||
anchor_shape = shapes[0]
|
||||
compatible = [
|
||||
source for source in candidates if shapes[source.source_index] == anchor_shape
|
||||
]
|
||||
present = {source.source_index for source in compatible}
|
||||
mismatch_count = len(handlers) - len(present)
|
||||
if preserve:
|
||||
sources = (anchor,)
|
||||
zero_count = 0
|
||||
elif mismatch_count and mismatch_mode == "error":
|
||||
raise ValueError(
|
||||
f"[{_OPERATION}] Layer '{core}' is missing or incompatible in one or more inputs."
|
||||
)
|
||||
elif mismatch_count and mismatch_mode == "skip":
|
||||
sources = (anchor,)
|
||||
zero_count = 0
|
||||
preserve = True
|
||||
else:
|
||||
sources = tuple(compatible)
|
||||
zero_count = mismatch_count if mismatch_mode == "zeros" else 0
|
||||
|
||||
for source in sources:
|
||||
stream_keys[source.source_index].extend(source.keys)
|
||||
units.append(DeltaCWBLayerUnit(core, output_key, sources, zero_count, preserve))
|
||||
return units, stream_keys, rejected
|
||||
|
||||
|
||||
def _to_processing_device(tensor, device):
|
||||
if device.type == "cuda":
|
||||
return transfer_to_gpu_pinned(tensor, device, torch.float32)
|
||||
return tensor.to(device=device, dtype=torch.float32)
|
||||
|
||||
|
||||
def _process_layer_on_device(
|
||||
unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics
|
||||
):
|
||||
deltas = []
|
||||
try:
|
||||
for source in unit.sources:
|
||||
if source.diff_key is not None:
|
||||
delta = _to_processing_device(
|
||||
loaded[(source.source_index, source.diff_key)], device
|
||||
)
|
||||
else:
|
||||
down = loaded[(source.source_index, source.down_key)]
|
||||
up = loaded[(source.source_index, source.up_key)]
|
||||
alpha = (
|
||||
loaded[(source.source_index, source.alpha_key)]
|
||||
if source.alpha_key is not None else None
|
||||
)
|
||||
down, up = normalize_lora_pair(down, up, alpha, layer=source.block_name)
|
||||
delta = _factor_delta(
|
||||
_to_processing_device(down, device),
|
||||
_to_processing_device(up, device),
|
||||
source.block_name,
|
||||
)
|
||||
deltas.append(delta)
|
||||
if unit.zero_contributors:
|
||||
deltas.extend(torch.zeros_like(deltas[0]) for _ in range(unit.zero_contributors))
|
||||
if unit.preserve_anchor:
|
||||
merged = deltas[0].clone()
|
||||
else:
|
||||
merged = merge_cwb_tensors(
|
||||
deltas,
|
||||
settings,
|
||||
reference_index=0,
|
||||
allow_similarity_alignment=False,
|
||||
diagnostics=diagnostics,
|
||||
)
|
||||
dtype = select_output_dtype(
|
||||
[source.source_dtype for source in unit.sources],
|
||||
save_dtype,
|
||||
is_1d_diff=merged.ndim == 1,
|
||||
)
|
||||
return merged.to(dtype=dtype, device="cpu").contiguous(), dtype
|
||||
finally:
|
||||
deltas.clear()
|
||||
|
||||
|
||||
def _is_cuda_oom(error):
|
||||
return isinstance(error, torch.cuda.OutOfMemoryError) or (
|
||||
isinstance(error, RuntimeError) and "out of memory" in str(error).lower()
|
||||
)
|
||||
|
||||
|
||||
def _release_cuda_oom():
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _process_layer(unit, loaded, settings, save_dtype, requested_device, diagnostics):
|
||||
device = torch.device(requested_device)
|
||||
attempt_diagnostics = CWBDiagnostics()
|
||||
try:
|
||||
output, dtype = _process_layer_on_device(
|
||||
unit, loaded, settings, save_dtype, device, attempt_diagnostics
|
||||
)
|
||||
diagnostics.absorb_runtime(attempt_diagnostics)
|
||||
return output, dtype, False
|
||||
except Exception as error:
|
||||
if device.type != "cuda" or not _is_cuda_oom(error):
|
||||
raise
|
||||
_release_cuda_oom()
|
||||
retry_diagnostics = CWBDiagnostics()
|
||||
output, dtype = _process_layer_on_device(
|
||||
unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics
|
||||
)
|
||||
diagnostics.absorb_runtime(retry_diagnostics)
|
||||
return output, dtype, True
|
||||
|
||||
|
||||
class DeltaCWBLoRAMergerLogic:
|
||||
@staticmethod
|
||||
def execute(
|
||||
lora_paths,
|
||||
source_names,
|
||||
settings,
|
||||
preset_name,
|
||||
custom_settings,
|
||||
mismatch_mode,
|
||||
output_filename,
|
||||
save_dtype,
|
||||
process_device,
|
||||
exclude_patterns,
|
||||
discard_patterns,
|
||||
glob_patterns,
|
||||
force_clear_cache,
|
||||
):
|
||||
prepare_for_large_operation(
|
||||
sum(estimate_model_size(path) for path in lora_paths) * 2.5,
|
||||
torch.device(process_device),
|
||||
)
|
||||
handlers = []
|
||||
cursors = []
|
||||
report = DeltaCWBReport()
|
||||
diagnostics = CWBDiagnostics()
|
||||
try:
|
||||
for path in lora_paths:
|
||||
handlers.append(MemoryEfficientSafeOpen(path, low_memory=True))
|
||||
infos = _inspect_inputs(handlers)
|
||||
layer_map = build_lora_layer_map(infos, None)
|
||||
validate_canonical_blocks(layer_map, _OPERATION)
|
||||
compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
|
||||
compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
|
||||
units, stream_keys, rejected = _build_layer_units(
|
||||
layer_map,
|
||||
infos,
|
||||
handlers,
|
||||
mismatch_mode,
|
||||
compiled_excludes,
|
||||
compiled_discards,
|
||||
glob_patterns,
|
||||
)
|
||||
report.rejected_unsupported_layers = rejected
|
||||
cursors = [
|
||||
AsyncTensorCursor(
|
||||
handler,
|
||||
keys,
|
||||
pin_memory=torch.device(process_device).type == "cuda",
|
||||
)
|
||||
for handler, keys in zip(handlers, stream_keys)
|
||||
]
|
||||
output_path, output_name = canonical_model_artifact_path("loras", output_filename)
|
||||
metadata = (handlers[0].metadata() or {}).copy()
|
||||
metadata.update({
|
||||
"merge_method": "delta_cwb",
|
||||
"source_models": ",".join(source_names),
|
||||
"alpha_normalized": "true",
|
||||
"output_representation": "full_direct_difference",
|
||||
"cwb_consensus_type": settings.consensus_type,
|
||||
"cwb_alignment_method": settings.alignment_method,
|
||||
"cwb_alignment_threshold": str(settings.alignment_threshold),
|
||||
"cwb_similarity_threshold": str(settings.similarity_threshold),
|
||||
"cwb_power_alpha": str(settings.power_alpha),
|
||||
"cwb_diversity_beta": str(settings.diversity_beta),
|
||||
"cwb_rescale_norm": str(settings.rescale_norm).lower(),
|
||||
"cwb_global_scale": str(settings.global_scale),
|
||||
"cwb_dynamic_similarity_contrast": str(settings.dynamic_similarity_contrast).lower(),
|
||||
"cwb_soft_comfort_bandpass": str(settings.soft_comfort_bandpass).lower(),
|
||||
"cwb_position_weight": str(settings.position_weight),
|
||||
"cwb_preserve_common_prefix": str(settings.preserve_common_prefix).lower(),
|
||||
})
|
||||
pbar = comfy.utils.ProgressBar(len(units))
|
||||
with atomic_uel_writer(output_path, metadata) as writer, torch.no_grad():
|
||||
for unit in tqdm(units, desc="Delta CWB LoRA layers", unit="layers"):
|
||||
if force_clear_cache:
|
||||
_release_cuda_oom()
|
||||
loaded = {}
|
||||
try:
|
||||
for source in unit.sources:
|
||||
for key in source.keys:
|
||||
loaded[(source.source_index, key)] = cursors[
|
||||
source.source_index
|
||||
].take(key)
|
||||
output, dtype, used_cpu = _process_layer(
|
||||
unit,
|
||||
loaded,
|
||||
settings,
|
||||
save_dtype,
|
||||
process_device,
|
||||
diagnostics,
|
||||
)
|
||||
writer.write_batch([(unit.output_key, output)])
|
||||
report.merged_layers += int(not unit.preserve_anchor)
|
||||
report.anchor_preserved_layers += int(unit.preserve_anchor)
|
||||
report.zero_filled_contributors += unit.zero_contributors
|
||||
report.cuda_cpu_fallbacks += int(used_cpu)
|
||||
report.record_dtype(dtype)
|
||||
del output
|
||||
finally:
|
||||
for source_index, key in list(loaded):
|
||||
del loaded[(source_index, key)]
|
||||
cursors[source_index].release(key)
|
||||
loaded.clear()
|
||||
pbar.update(1)
|
||||
for cursor in cursors:
|
||||
cursor.finish()
|
||||
return output_name, report.render(
|
||||
diagnostics, preset_name, len(lora_paths), custom_settings
|
||||
)
|
||||
finally:
|
||||
for cursor in cursors:
|
||||
cursor.close()
|
||||
for handler in handlers:
|
||||
handler.__exit__(None, None, None)
|
||||
cleanup_after_operation()
|
||||
|
||||
@staticmethod
|
||||
def execute_node(kwargs, count):
|
||||
names = [kwargs[f"lora_{index}"] for index in range(1, count + 1)]
|
||||
paths = [folder_paths.get_full_path_or_raise("loras", name) for name in names]
|
||||
settings = resolve_cwb_settings(
|
||||
kwargs["cwb_preset"], DENSE_CWB_PRESETS, kwargs.get("cwb_config")
|
||||
)
|
||||
requested_dtype = {
|
||||
"fp32": torch.float32,
|
||||
"fp16": torch.float16,
|
||||
"bf16": torch.bfloat16,
|
||||
}[kwargs["save_dtype"]]
|
||||
filename, report = DeltaCWBLoRAMergerLogic.execute(
|
||||
paths,
|
||||
names,
|
||||
settings,
|
||||
kwargs["cwb_preset"],
|
||||
kwargs.get("cwb_config") is not None,
|
||||
kwargs["mismatch_mode"],
|
||||
kwargs["output_filename"],
|
||||
requested_dtype,
|
||||
kwargs["process_device"],
|
||||
kwargs["exclude_patterns"],
|
||||
kwargs["discard_patterns"],
|
||||
kwargs["glob_patterns"],
|
||||
kwargs["force_clear_cache"],
|
||||
)
|
||||
return io.NodeOutput(filename, report)
|
||||
|
||||
|
||||
def _fixed_inputs(count, default_filename):
|
||||
options = folder_paths.get_filename_list("loras")
|
||||
inputs = []
|
||||
if count:
|
||||
inputs.append(io.Combo.Input(
|
||||
"lora_1",
|
||||
options=options,
|
||||
tooltip="Metadata and layer anchor. Defines saved layers and preserved values.",
|
||||
))
|
||||
inputs.extend(
|
||||
io.Combo.Input(
|
||||
f"lora_{index}",
|
||||
options=options,
|
||||
tooltip="Equal-prior LoRA contributor. CWB determines its effective row influence.",
|
||||
)
|
||||
for index in range(2, count + 1)
|
||||
)
|
||||
inputs.extend([
|
||||
io.Combo.Input(
|
||||
"cwb_preset",
|
||||
options=list(DENSE_CWB_PRESETS),
|
||||
default="balanced_mean",
|
||||
tooltip="Controls how strongly agreement, disagreement, and row magnitude affect the saved full weight changes. A connected CWB Config overrides it completely.",
|
||||
),
|
||||
CWB_CONFIG.Input(
|
||||
"cwb_config",
|
||||
optional=True,
|
||||
tooltip="Optional complete settings override from CWB Custom Configuration.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"mismatch_mode",
|
||||
options=["skip", "zeros", "error"],
|
||||
default="skip",
|
||||
tooltip="For a missing or incompatible anchored layer: preserve LoRA 1, include an explicit zero contributor, or abort.",
|
||||
),
|
||||
io.String.Input(
|
||||
"output_filename",
|
||||
default=default_filename,
|
||||
tooltip="Relative filename for the atomically saved full-difference LoRA.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"save_dtype",
|
||||
options=["fp32", "fp16", "bf16"],
|
||||
default="bf16",
|
||||
tooltip="Requested saved precision. Participating FP32 inputs preserve FP32.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"process_device",
|
||||
options=["cuda", "cpu"],
|
||||
default="cuda",
|
||||
tooltip="Processing device. Only a failed CUDA layer retries on CPU.",
|
||||
),
|
||||
io.String.Input(
|
||||
"exclude_patterns",
|
||||
default="",
|
||||
multiline=True,
|
||||
tooltip="Preserve matching canonical full-difference layers from LoRA 1.",
|
||||
),
|
||||
io.String.Input(
|
||||
"discard_patterns",
|
||||
default="",
|
||||
multiline=True,
|
||||
tooltip="Omit matching canonical full-difference layers from the output.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"glob_patterns",
|
||||
default=False,
|
||||
tooltip="Use shell-style glob patterns instead of regular expressions.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"force_clear_cache",
|
||||
default=True,
|
||||
tooltip="Clear Python and CUDA caches before each layer to reduce retained memory.",
|
||||
),
|
||||
])
|
||||
return inputs
|
||||
|
||||
|
||||
def _multi_inputs(default_filename):
|
||||
options = folder_paths.get_filename_list("loras")
|
||||
optional = ["None", *options]
|
||||
inputs = [
|
||||
io.Combo.Input(
|
||||
"lora_count",
|
||||
options=[str(index) for index in range(2, 9)],
|
||||
default="2",
|
||||
tooltip="Number of consecutive selected LoRAs to merge.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"lora_1",
|
||||
options=options,
|
||||
tooltip="Metadata and layer anchor. Defines saved layers and preserved values.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"lora_2",
|
||||
options=options,
|
||||
tooltip="Second equal-prior LoRA contributor.",
|
||||
),
|
||||
]
|
||||
inputs.extend(
|
||||
io.Combo.Input(
|
||||
f"lora_{index}",
|
||||
options=optional,
|
||||
default="None",
|
||||
tooltip="Additional equal-prior contributor; required when included by LoRA Count.",
|
||||
)
|
||||
for index in range(3, 9)
|
||||
)
|
||||
inputs.extend(_fixed_inputs(0, default_filename))
|
||||
return inputs
|
||||
|
||||
|
||||
class DeltaCWBLoRATwoMerger(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DeltaCWBLoRATwoMerger",
|
||||
display_name="Delta CWB Merge LoRAs (2)",
|
||||
category="ModelUtils/LoRA/Merge/Delta CWB",
|
||||
description="Expand two LoRAs to full weight changes and merge corresponding rows with CWB.",
|
||||
inputs=_fixed_inputs(2, "delta_cwb_merged_2_lora"),
|
||||
outputs=[
|
||||
io.AnyType.Output(display_name="output_filename"),
|
||||
io.String.Output(display_name="cwb_report"),
|
||||
],
|
||||
is_output_node=True,
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, **kwargs):
|
||||
return DeltaCWBLoRAMergerLogic.execute_node(kwargs, 2)
|
||||
|
||||
|
||||
class DeltaCWBLoRAThreeMerger(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DeltaCWBLoRAThreeMerger",
|
||||
display_name="Delta CWB Merge LoRAs (3)",
|
||||
category="ModelUtils/LoRA/Merge/Delta CWB",
|
||||
description="Expand three LoRAs to full weight changes and merge corresponding rows with CWB.",
|
||||
inputs=_fixed_inputs(3, "delta_cwb_merged_3_lora"),
|
||||
outputs=[
|
||||
io.AnyType.Output(display_name="output_filename"),
|
||||
io.String.Output(display_name="cwb_report"),
|
||||
],
|
||||
is_output_node=True,
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, **kwargs):
|
||||
return DeltaCWBLoRAMergerLogic.execute_node(kwargs, 3)
|
||||
|
||||
|
||||
class DeltaCWBLoRAMultiMerger(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DeltaCWBLoRAMultiMerger",
|
||||
display_name="Delta CWB LoRA Multi-Merge",
|
||||
category="ModelUtils/LoRA/Merge/Delta CWB",
|
||||
description="Expand 2-8 LoRAs to full weight changes and merge corresponding rows with CWB.",
|
||||
inputs=_multi_inputs("delta_cwb_merged_multi_lora"),
|
||||
outputs=[
|
||||
io.AnyType.Output(display_name="output_filename"),
|
||||
io.String.Output(display_name="cwb_report"),
|
||||
],
|
||||
is_output_node=True,
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, **kwargs):
|
||||
count = int(kwargs["lora_count"])
|
||||
missing = [
|
||||
index for index in range(1, count + 1)
|
||||
if not kwargs[f"lora_{index}"] or kwargs[f"lora_{index}"] == "None"
|
||||
]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"LoRA Count includes unselected input(s): {', '.join(map(str, missing))}."
|
||||
)
|
||||
return DeltaCWBLoRAMergerLogic.execute_node(kwargs, count)
|
||||
|
||||
|
||||
DELTA_CWB_LORA_NODES = [
|
||||
DeltaCWBLoRATwoMerger,
|
||||
DeltaCWBLoRAThreeMerger,
|
||||
DeltaCWBLoRAMultiMerger,
|
||||
]
|
||||
@@ -11,6 +11,7 @@ REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE_NAME = "modelutils_artifact_path_tests"
|
||||
MIGRATED_SAVE_MODULES = (
|
||||
"consensus_merger.py",
|
||||
"cwb_delta_lora_merger.py",
|
||||
"merger.py",
|
||||
"lora_merger.py",
|
||||
"lodestone_merger.py",
|
||||
|
||||
@@ -0,0 +1,329 @@
|
||||
import importlib
|
||||
import re
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE_NAME = "modelutils_delta_cwb_tests"
|
||||
|
||||
|
||||
def _load_module(name):
|
||||
if PACKAGE_NAME not in sys.modules:
|
||||
package = types.ModuleType(PACKAGE_NAME)
|
||||
package.__path__ = [str(REPO_ROOT)]
|
||||
sys.modules[PACKAGE_NAME] = package
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def delta_cwb():
|
||||
return _load_module("nodes.cwb_delta_lora_merger")
|
||||
|
||||
|
||||
def _settings(delta_cwb):
|
||||
return delta_cwb.CWBSettings(
|
||||
consensus_type="mean",
|
||||
alignment_method="similarity",
|
||||
alignment_threshold=0.0,
|
||||
similarity_threshold=-1.0,
|
||||
power_alpha=0.0,
|
||||
diversity_beta=0.0,
|
||||
rescale_norm=False,
|
||||
global_scale=1.0,
|
||||
dynamic_similarity_contrast=False,
|
||||
soft_comfort_bandpass=False,
|
||||
position_weight=0.0,
|
||||
preserve_common_prefix=False,
|
||||
)
|
||||
|
||||
|
||||
def _write(path, tensors, metadata=None):
|
||||
with IncrementalSafetensorsWriter(str(path), metadata=metadata or {}, max_workers=1) as writer:
|
||||
writer.write_batch(list(tensors.items()))
|
||||
|
||||
|
||||
def _read(path):
|
||||
tensors = {}
|
||||
with MemoryEfficientSafeOpen(str(path), low_memory=True) as handler:
|
||||
stream = handler.async_stream(
|
||||
handler.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
|
||||
)
|
||||
for batch in stream:
|
||||
for key, tensor in batch:
|
||||
tensors[key] = tensor.clone()
|
||||
handler.mark_processed(key)
|
||||
return tensors
|
||||
|
||||
|
||||
def test_linear_and_convolutional_factor_expansion(delta_cwb):
|
||||
down = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
||||
up = torch.tensor([[2.0, -1.0], [0.5, 3.0]])
|
||||
torch.testing.assert_close(
|
||||
delta_cwb._factor_delta(down, up, "linear"), up @ down
|
||||
)
|
||||
|
||||
conv_down = torch.arange(18, dtype=torch.float32).reshape(2, 1, 3, 3)
|
||||
conv_up = torch.tensor([[[[2.0]], [[-1.0]]]])
|
||||
expected = (conv_up.reshape(1, 2) @ conv_down.reshape(2, -1)).reshape(1, 1, 3, 3)
|
||||
torch.testing.assert_close(
|
||||
delta_cwb._factor_delta(conv_down, conv_up, "conv"), expected
|
||||
)
|
||||
|
||||
|
||||
def test_dense_cwb_keeps_corresponding_rows_fixed(delta_cwb):
|
||||
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
||||
reversed_source = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
|
||||
result = delta_cwb.merge_cwb_tensors(
|
||||
[reference, reversed_source],
|
||||
_settings(delta_cwb),
|
||||
reference_index=0,
|
||||
allow_similarity_alignment=False,
|
||||
)
|
||||
torch.testing.assert_close(result, torch.full((2, 2), 0.5))
|
||||
|
||||
|
||||
def test_layer_units_apply_missing_modes_filters_and_shape_checks(delta_cwb):
|
||||
class Handler:
|
||||
def __init__(self, shapes):
|
||||
self.shapes = shapes
|
||||
|
||||
def get_shape(self, key):
|
||||
return self.shapes[key]
|
||||
|
||||
def get_dtype(self, _key):
|
||||
return torch.float32
|
||||
|
||||
roles = {"down": "layer.down", "up": "layer.up"}
|
||||
infos = [
|
||||
{"pairs": {"layer": roles}, "low_bit": set()},
|
||||
{"pairs": {}, "low_bit": set()},
|
||||
]
|
||||
handlers = [Handler({"layer.down": (1, 2), "layer.up": (2, 1)}), Handler({})]
|
||||
layer_map = {"layer": [(0, "layer")]}
|
||||
empty = delta_cwb._compile_patterns("", False)
|
||||
|
||||
with pytest.raises(ValueError, match="missing or incompatible"):
|
||||
delta_cwb._build_layer_units(
|
||||
layer_map, infos, handlers, "error", empty, empty, False
|
||||
)
|
||||
units, _, _ = delta_cwb._build_layer_units(
|
||||
layer_map, infos, handlers, "zeros", empty, empty, False
|
||||
)
|
||||
assert units[0].zero_contributors == 1
|
||||
assert units[0].preserve_anchor is False
|
||||
|
||||
exclude = delta_cwb._compile_patterns(r"\.diff$", False)
|
||||
units, _, _ = delta_cwb._build_layer_units(
|
||||
layer_map, infos, handlers, "zeros", exclude, empty, False
|
||||
)
|
||||
assert units[0].preserve_anchor is True
|
||||
assert units[0].zero_contributors == 0
|
||||
|
||||
discard = delta_cwb._compile_patterns("*.diff", True)
|
||||
units, _, rejected = delta_cwb._build_layer_units(
|
||||
layer_map, infos, handlers, "skip", empty, discard, True
|
||||
)
|
||||
assert units == []
|
||||
assert rejected == 1
|
||||
|
||||
|
||||
def test_streamed_merge_normalizes_alpha_once_and_supports_mixed_rank(
|
||||
monkeypatch, tmp_path, delta_cwb
|
||||
):
|
||||
first = tmp_path / "first.safetensors"
|
||||
second = tmp_path / "second.safetensors"
|
||||
_write(first, {
|
||||
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
|
||||
"diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [0.0]]),
|
||||
"diffusion_model.layer.alpha": torch.tensor(0.5),
|
||||
}, {"anchor": "yes"})
|
||||
_write(second, {
|
||||
"diffusion_model.layer.lora_A.weight": torch.eye(2),
|
||||
"diffusion_model.layer.lora_B.weight": torch.diag(torch.tensor([3.0, 4.0])),
|
||||
})
|
||||
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
|
||||
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
|
||||
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
|
||||
|
||||
filename, report = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
|
||||
[str(first), str(second)],
|
||||
[first.name, second.name],
|
||||
_settings(delta_cwb),
|
||||
"balanced_mean",
|
||||
False,
|
||||
"skip",
|
||||
"dense/mixed_rank",
|
||||
torch.float32,
|
||||
"cpu",
|
||||
"",
|
||||
"",
|
||||
False,
|
||||
False,
|
||||
)
|
||||
output = tmp_path / "loras" / filename
|
||||
tensors = _read(output)
|
||||
|
||||
assert filename == "dense/mixed_rank.safetensors"
|
||||
assert list(tensors) == ["diffusion_model.layer.diff"]
|
||||
torch.testing.assert_close(
|
||||
tensors["diffusion_model.layer.diff"],
|
||||
torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
|
||||
)
|
||||
assert "Merged layers: 1" in report
|
||||
with MemoryEfficientSafeOpen(str(output), low_memory=True) as handler:
|
||||
metadata = handler.metadata()
|
||||
assert metadata["anchor"] == "yes"
|
||||
assert metadata["merge_method"] == "delta_cwb"
|
||||
assert metadata["alpha_normalized"] == "true"
|
||||
assert metadata["output_representation"] == "full_direct_difference"
|
||||
|
||||
|
||||
def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cwb):
|
||||
first = tmp_path / "diff_first.safetensors"
|
||||
second = tmp_path / "diff_second.safetensors"
|
||||
key = "diffusion_model.bias.diff"
|
||||
_write(first, {key: torch.tensor([1.0, 3.0])})
|
||||
_write(second, {key: torch.tensor([3.0, 5.0])})
|
||||
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
|
||||
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
|
||||
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
|
||||
|
||||
filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
|
||||
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
|
||||
"balanced_mean", False, "skip", "one_dim", torch.float16, "cpu",
|
||||
"", "", False, False,
|
||||
)
|
||||
output = _read(tmp_path / "loras" / filename)[key]
|
||||
assert output.dtype == torch.float32
|
||||
torch.testing.assert_close(output, torch.tensor([2.0, 4.0]))
|
||||
|
||||
|
||||
def test_atomic_writer_preserves_destination_on_cwb_failure(
|
||||
monkeypatch, tmp_path, delta_cwb
|
||||
):
|
||||
first = tmp_path / "atomic_first.safetensors"
|
||||
second = tmp_path / "atomic_second.safetensors"
|
||||
tensors = {
|
||||
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
|
||||
"diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [0.0]]),
|
||||
}
|
||||
_write(first, tensors)
|
||||
_write(second, tensors)
|
||||
destination = tmp_path / "loras" / "atomic.safetensors"
|
||||
destination.parent.mkdir(parents=True)
|
||||
destination.write_bytes(b"existing destination")
|
||||
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
|
||||
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
|
||||
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
|
||||
released = []
|
||||
closed = []
|
||||
opener = delta_cwb.MemoryEfficientSafeOpen
|
||||
|
||||
class TrackedHandler:
|
||||
def __init__(self, path, low_memory=True):
|
||||
self.handler = opener(path, low_memory=low_memory)
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.handler, name)
|
||||
|
||||
def mark_processed(self, key):
|
||||
released.append(key)
|
||||
return self.handler.mark_processed(key)
|
||||
|
||||
def __exit__(self, exc_type, exc, traceback):
|
||||
closed.append(True)
|
||||
return self.handler.__exit__(exc_type, exc, traceback)
|
||||
|
||||
monkeypatch.setattr(delta_cwb, "MemoryEfficientSafeOpen", TrackedHandler)
|
||||
monkeypatch.setattr(
|
||||
delta_cwb,
|
||||
"merge_cwb_tensors",
|
||||
lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("planned failure")),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="planned failure"):
|
||||
delta_cwb.DeltaCWBLoRAMergerLogic.execute(
|
||||
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
|
||||
"balanced_mean", False, "skip", "atomic", torch.float32, "cpu",
|
||||
"", "", False, False,
|
||||
)
|
||||
assert destination.read_bytes() == b"existing destination"
|
||||
assert not list(destination.parent.glob(".atomic.safetensors.*.tmp"))
|
||||
assert sorted(released) == sorted([*tensors, *tensors])
|
||||
assert len(closed) == 2
|
||||
|
||||
|
||||
def test_complete_layer_cuda_oom_retries_on_cpu(monkeypatch, delta_cwb):
|
||||
source = delta_cwb.DeltaCWBLoRASource(
|
||||
0, "layer", None, None, None, "layer.diff", torch.float32
|
||||
)
|
||||
unit = delta_cwb.DeltaCWBLayerUnit(
|
||||
"layer", "layer.diff", (source,), 0, False
|
||||
)
|
||||
calls = []
|
||||
|
||||
def process(_unit, _loaded, _settings, _dtype, device, _diagnostics):
|
||||
calls.append(device.type)
|
||||
if device.type == "cuda":
|
||||
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
|
||||
return torch.ones(1), torch.float32
|
||||
|
||||
monkeypatch.setattr(delta_cwb, "_process_layer_on_device", process)
|
||||
monkeypatch.setattr(delta_cwb, "_release_cuda_oom", lambda: None)
|
||||
output, dtype, fallback = delta_cwb._process_layer(
|
||||
unit, {}, _settings(delta_cwb), torch.float32, "cuda", delta_cwb.CWBDiagnostics()
|
||||
)
|
||||
assert calls == ["cuda", "cpu"]
|
||||
assert fallback is True
|
||||
assert dtype == torch.float32
|
||||
assert torch.equal(output, torch.ones(1))
|
||||
|
||||
|
||||
def test_exact_node_schema_contract(delta_cwb):
|
||||
expected_ids = [
|
||||
"DeltaCWBLoRATwoMerger",
|
||||
"DeltaCWBLoRAThreeMerger",
|
||||
"DeltaCWBLoRAMultiMerger",
|
||||
]
|
||||
schemas = [node.define_schema() for node in delta_cwb.DELTA_CWB_LORA_NODES]
|
||||
assert [schema.node_id for schema in schemas] == expected_ids
|
||||
assert [schema.display_name for schema in schemas] == [
|
||||
"Delta CWB Merge LoRAs (2)",
|
||||
"Delta CWB Merge LoRAs (3)",
|
||||
"Delta CWB LoRA Multi-Merge",
|
||||
]
|
||||
assert all(schema.category == "ModelUtils/LoRA/Merge/Delta CWB" for schema in schemas)
|
||||
assert all([output.display_name for output in schema.outputs] == ["output_filename", "cwb_report"] for schema in schemas)
|
||||
assert all(schema.outputs[0].get_io_type() == "*" for schema in schemas)
|
||||
|
||||
fixed_ids = [value.id for value in schemas[0].inputs]
|
||||
assert fixed_ids == [
|
||||
"lora_1", "lora_2", "cwb_preset", "cwb_config", "mismatch_mode",
|
||||
"output_filename", "save_dtype", "process_device", "exclude_patterns",
|
||||
"discard_patterns", "glob_patterns", "force_clear_cache",
|
||||
]
|
||||
multi_ids = [value.id for value in schemas[2].inputs]
|
||||
assert multi_ids[:9] == ["lora_count", *[f"lora_{index}" for index in range(1, 9)]]
|
||||
assert schemas[0].inputs[2].default == "balanced_mean"
|
||||
assert schemas[0].inputs[4].default == "skip"
|
||||
assert schemas[0].inputs[6].default == "bf16"
|
||||
assert schemas[0].inputs[7].default == "cuda"
|
||||
assert schemas[0].inputs[-2].default is False
|
||||
assert schemas[0].inputs[-1].default is True
|
||||
assert "agreement" in schemas[0].inputs[2].tooltip
|
||||
assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[7].tooltip
|
||||
|
||||
|
||||
def test_invalid_regex_and_unselected_multi_input_fail(delta_cwb):
|
||||
with pytest.raises(re.error):
|
||||
delta_cwb._compile_patterns("[", False)
|
||||
with pytest.raises(ValueError, match="unselected input.*3"):
|
||||
delta_cwb.DeltaCWBLoRAMultiMerger.execute(
|
||||
lora_count="3", lora_1="a", lora_2="b", lora_3="None"
|
||||
)
|
||||
@@ -9,17 +9,17 @@ python_tests = [
|
||||
paths = [
|
||||
"nodes/artifact_paths.py", "nodes/dora_extract_wd.py",
|
||||
"nodes/dora_learned_wd.py", "nodes/dtype_conversion.py",
|
||||
"nodes/lodestone_merger.py", "nodes/lora_extract_svd.py", "nodes/lora_merger.py",
|
||||
"nodes/cwb_delta_lora_merger.py", "nodes/lodestone_merger.py", "nodes/lora_extract_svd.py", "nodes/lora_merger.py",
|
||||
"nodes/lora_rename.py", "nodes/lora_resize.py", "nodes/merger.py",
|
||||
"nodes/prunekeys.py", "nodes/renamekeys.py", "nodes/text_encoder_extract.py",
|
||||
]
|
||||
python_tests = ["tests/test_artifact_paths.py"]
|
||||
|
||||
[groups.cwb]
|
||||
paths = ["nodes/consensus_merger.py", "run_cwb_api_workflow.py", "run_cwb_preset_sweep.py"]
|
||||
paths = ["nodes/consensus_merger.py", "nodes/cwb_delta_lora_merger.py", "run_cwb_api_workflow.py", "run_cwb_preset_sweep.py"]
|
||||
python_tests = [
|
||||
"tests/test_consensus_merger.py", "tests/test_cwb_api_runner.py",
|
||||
"tests/test_cwb_preset_sweep.py",
|
||||
"tests/test_cwb_delta_lora_merger.py", "tests/test_cwb_preset_sweep.py",
|
||||
]
|
||||
|
||||
[groups.model_analysis]
|
||||
@@ -36,12 +36,12 @@ python_tests = ["tests/test_adaptive_extraction_pipeline.py", "tests/test_knee_p
|
||||
|
||||
[groups.lora_merge]
|
||||
paths = [
|
||||
"nodes/device_utils.py", "nodes/lodestone_merger.py", "nodes/lora_alpha.py",
|
||||
"nodes/cwb_delta_lora_merger.py", "nodes/device_utils.py", "nodes/lodestone_merger.py", "nodes/lora_alpha.py",
|
||||
"nodes/lora_merger.py", "nodes/quantization_guard.py", "nodes/uel_io.py",
|
||||
]
|
||||
python_tests = [
|
||||
"tests/test_lora_diff_dtype_contract.py", "tests/test_lora_merger_async.py",
|
||||
"tests/test_lodestone_merger.py", "tests/test_lora_quantization_guard.py",
|
||||
"tests/test_cwb_delta_lora_merger.py", "tests/test_lodestone_merger.py", "tests/test_lora_quantization_guard.py",
|
||||
]
|
||||
|
||||
[groups.lora_resize]
|
||||
|
||||
Reference in New Issue
Block a user