Compare commits

...
Author SHA1 Message Date
silveroxides 89f4def83e feat(cwb): preserve secondary-only layers
- build generic and LoRA worklists from the union of inputs
- copy single-source layers exactly and merge available providers
- keep logical LoRA groups canonical, atomic, and collision-safe
- retain mismatch, dtype, quantization, streaming, and OOM contracts
2026-08-02 20:26:41 +02:00
silveroxides aca6147f35 chore(package): exclude tests 2026-08-02 18:12:01 +02:00
silveroxides f390c12776 fix(merge): enforce LoRA tensor contracts
- preserve direct diffs, precision, and guarded low-bit tensors
- stream paired CWB layers with atomic CPU OOM fallback
- save generated merge artifacts under models_dir
2026-08-02 18:11:19 +02:00
silveroxides bed9c6e8ff feat(merge): add dedicated CWB mergers 2026-08-01 13:14:14 +02:00
silveroxides 49b66ea64e fix(lora): enforce adapter tensor contracts 2026-08-01 12:15:39 +02:00
silveroxides f828ca7dfc chore: bump version to 0.4.6 2026-07-31 05:55:30 +02:00
silveroxides 47be3fa1bd fix(lora): preserve FP32 source layers 2026-07-31 05:55:21 +02:00
silveroxides e972b53cf0 chore: bump version to 0.4.5 2026-07-31 05:49:02 +02:00
silveroxides 4cb065b9d2 feat(lora): add opt-in full 1D diffs 2026-07-31 05:48:53 +02:00
Silver febab8ffcd Merge pull request #16 from Booyaka101/fix/dora-extract-1d-tensor-crash
Fix DoRA extraction crash on 1-D weights (x_pad_token, norm weights)
2026-07-11 11:26:54 +02:00
Silver 06365311e2 Merge pull request #18 from silveroxides/fix/loading-refactor
Refactor loading strategies in mergers and LoRA resize nodes.
2026-07-11 11:23:07 +02:00
silveroxides 9375ced5cf Refactor loading strategies in mergers and LoRA resize nodes. 2026-07-11 11:21:40 +02:00
Silver 314f5ef67b Bump version from 0.4.3 to 0.4.4 2026-06-06 18:58:20 +02:00
Silver 841b70c384 Merge pull request #17 from silveroxides/feature/dora-extract
Feature/dora extract
2026-06-06 18:58:00 +02:00
Booyaka101 22e7701411 fix: skip 1-D weights before the norm in DoRA extraction
The matched-shape branch computed torch.linalg.norm(tensor_a, dim=1)
for every non-4-D tensor, including 1-D weights (e.g. x_pad_token,
norm .weight). A 1-D tensor has no dim=1, so this raised
'Dimension out of range ... but got 1' and aborted the whole
extraction. The code already skips 1-D tensors a few lines later, so
bail out before the norm instead.

Fixes #15
2026-06-06 18:21:50 +08:00
16 changed files with 5567 additions and 528 deletions
+1
View File
@@ -0,0 +1 @@
tests/
+3
View File
@@ -20,6 +20,7 @@ from .nodes.merger import (
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger
)
from .nodes.consensus_merger import CWB_MERGER_NODES
from .nodes.lora_extract_svd import (
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius
@@ -68,6 +69,8 @@ class ModelUtilsExtension(ComfyExtension):
# Three-Model Mergers
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger,
# Dedicated CWB Mergers
*CWB_MERGER_NODES,
# LoRA Extraction
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius,
+109
View File
@@ -0,0 +1,109 @@
# Consensus-Weighted Blending Mergers
The dedicated CWB nodes merge two or three safetensors files using
Consensus-Weighted Blending. They calculate an element-wise mean or median
consensus, compare each aligned source vector to that consensus, and normalize
similarity-derived weights before producing the output vector.
## Node families
Two-input and three-input variants are provided for checkpoints, standalone
diffusion models, text encoders, LoRAs, and embeddings.
Model A supplies output metadata and anchors shared-layer shape, naming, and
mismatch behavior. The output key set is the union of all inputs. A valid
tensor found only in a secondary input is copied unchanged. When Model A lacks
a tensor that is present in multiple secondary inputs, those available sources
are CWB-merged without inserting a zero contribution for Model A. Embedding
mergers additionally use the longest compatible first dimension as their
alignment reference.
## Presets and custom controls
`baseline` is the default. The standard clarity, smooth, varied, and diversity
presets are joined by `power_blend` and the DSC presets demonstrated by the CWB
conditioning implementation. Select `custom` to make every manual CWB control
authoritative.
Similarity alignment greedily matches source rows to the reference by cosine
similarity. Index alignment keeps absolute row positions. Position weight can
bias similarity matches toward nearby normalized positions. Preserve Common
Prefix copies the longest numerically identical leading row span directly from
the first input.
Inputs have equal prior weight. The similarity calculation supplies the
effective contribution for each aligned vector. If every candidate is rejected
or receives zero weight, CWB falls back to equal weighting.
## Missing tensors and filtering
- `skip` preserves Model A when an anchored input is missing or incompatible.
In embedding-union mode it omits that source for the affected key.
- `zeros` supplies a compatible zero tensor for the missing source.
- `error` stops the merge.
- Exclude patterns preserve the anchored source tensor.
- Discard patterns remove the matching output tensor or logical LoRA layer.
The missing-input modes above apply to Model-A-owned layers. Model A being
absent from a secondary-only layer is not itself a mismatch. A shape conflict
between multiple available secondary providers does use the selected mismatch
mode, with the earliest available input as its preservation anchor.
Patterns use regular expressions by default. Enable Glob Patterns to use glob
syntax instead.
## LoRA behavior
LoRA inputs use the same parser as the resize and multi-merge nodes. It accepts
all seven low-rank A/B spellings recognized by ComfyUI's `LoRAAdapter`, plus
direct `.diff`, `.diff_b`, `.w_norm`, `.b_norm`, and `.set_weight` layers.
Recognized diffusion-model outputs use the canonical
`diffusion_model.<layer>.lora_A.weight` / `.lora_B.weight` convention.
Secondary formats are matched by normalized logical layer name. Differing
ranks are zero-padded like the DARE/TIES mergers: A/down tensors on dimension 0
and B/up tensors on dimension 1. An existing anchor-source alpha key is updated
to the resulting maximum rank; no alpha key is invented when the anchor has
none.
Similarity alignment operates on paired latent-rank components. A rows and B
columns share one mapping; fixed input/output feature axes are never reordered.
Input alpha scaling is absorbed into the B factor before blending, and global
scale is applied once to the resulting pair rather than once per factor.
Incomplete pairs and unrecognized tensors are copied from the earliest
available source unchanged.
Companion-bearing groups (`lora_mid`, reshape, DoRA, or set-weight) are kept
atomically from their earliest available source rather than partially
transformed. Secondary-only logical groups are included automatically: a group
from one source is copied unchanged, while compatible groups from multiple
secondary sources are merged using only those sources. Recognized roles are
written with the same canonical output convention as shared layers.
`include_1d_diffs` is disabled by default. Disabled shared 1D direct layers
preserve Model A, while secondary-only layers preserve their earliest source.
Enabled 1D direct layers participate as complete tensors and are always saved
in FP32.
## Precision and quantized inputs
CWB computes in FP32. Unless Override Dtype is enabled, any participating FP32
source keeps the generated logical layer in FP32. Guarded low-bit tensors and
enabled 1D direct diffs are never downcast by the override.
ComfyUI/comfy-kitchen quantized inputs are rejected before an output writer is
opened. Files with three or more bare INT8, UINT8, FP8, or FP4 tensors are also
rejected. One or two isolated low-bit tensors produce one warning per input and
are preserved without CWB arithmetic; an affected LoRA causes its complete
logical layer to be preserved.
## Streaming and failure safety
Model parameters and direct patches use fixed-coordinate blending; similarity
row matching is limited to embeddings and LoRA rank components. Each logical
layer is loaded, transferred to the processing device, completed, copied back
to CPU, queued for writing, and released before the next layer. Outputs are
written to a temporary sibling and atomically published only after every layer
and asynchronous write succeeds. Failed merges do not replace or finalize the
requested output. A CUDA OOM retries only the affected layer on CPU after
releasing the failed GPU working set; failure of that retry aborts atomically.
Generated files are published under the matching category inside
`folder_paths.models_dir`, never ComfyUI's general output directory.
+9
View File
@@ -121,6 +121,15 @@ When merging models with different layer structures (e.g., fine-tuned models wit
---
## Dtype Preservation & Override
By default, the merger preserves the highest precision `dtype` for each key found across all active source models (Model A and Model B) and the requested `save_dtype`.
- **Preserve Higher Precision**: If a key is `float32` in any of the source models, but you requested `save_dtype="bf16"` or `"fp16"`, that key will be saved as `float32` anyway. This prevents degrading critical high-precision layers (such as normalization layers or layers with restored outliers).
- **Explicit Override**: If you prefer to force the entire model to be saved strictly as the selected `save_dtype` (e.g. for complete downcasting to `bf16`), toggle `override_dtype` to `True`.
---
## Layer Filtering (Regex Patterns)
Use regex patterns to control which layers are merged:
+9
View File
@@ -67,6 +67,15 @@ When merging models with different layer structures, the `mismatch_mode` paramet
---
## Dtype Preservation & Override
By default, the merger preserves the highest precision `dtype` for each key found across all active source models (Model A, B, and C) and the requested `save_dtype`.
- **Preserve Higher Precision**: If a key is `float32` in any of the active source models, but you requested `save_dtype="bf16"` or `"fp16"`, that key will be saved as `float32` anyway. This prevents degrading critical high-precision layers (such as normalization layers or layers with restored outliers).
- **Explicit Override**: If you prefer to force the entire model to be saved strictly as the selected `save_dtype` (e.g. for complete downcasting to `bf16`), toggle `override_dtype` to `True`.
---
## Layer Filtering (Regex Patterns)
Use regex patterns to control which layers are merged:
File diff suppressed because it is too large Load Diff
+8
View File
@@ -551,6 +551,14 @@ def extract_dora_from_files(
dora_scale = None
del tensor_b
else:
# 1-D weights (e.g. x_pad_token, norm .weight) have no
# weight-decomposition direction and are skipped below
# anyway; bail out before the norm, which would index
# dim=1 and raise on a 1-D tensor.
if tensor_a.ndim < 2:
del tensor_a, tensor_b
return "skipped", None
# Weight Decomposition calculation
is_conv_dim = tensor_a.ndim == 4
+70 -21
View File
@@ -19,6 +19,7 @@ from .device_utils import (
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
from .quantization_guard import inspect_low_bit_input
@@ -458,6 +459,7 @@ def extract_lora_from_files(
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
include_1d_diffs: bool = False,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
@@ -469,7 +471,7 @@ def extract_lora_from_files(
linear_param: Mode parameter for linear layers
conv_param: Mode parameter for conv layers
device: Computation device
save_dtype: Output dtype
save_dtype: Output dtype for layers whose sources are not FP32
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers
@@ -480,6 +482,8 @@ def extract_lora_from_files(
chunk_large_layers: Enable chunked extraction
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
include_1d_diffs: Save 1D tensors in FP32 and keys ending in .scale or
.lin as full .diff patches instead of skipping them.
"""
save_torch_dtype = {
"fp32": torch.float32,
@@ -498,10 +502,22 @@ def extract_lora_from_files(
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
low_bit_a = inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "LoRA Extract")
low_bit_b = inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "LoRA Extract")
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
direct_diff_suffixes = (".scale", ".lin")
extraction_keys = [
k for k in keys_a
if k not in low_bit_a
and k not in low_bit_b
and (k.endswith(".weight")
or (
include_1d_diffs
and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes))
))
]
pbar = comfy.utils.ProgressBar(len(extraction_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
@@ -510,6 +526,13 @@ def extract_lora_from_files(
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
source_dtypes = [handler_a.get_dtype(key)]
if key in keys_b:
source_dtypes.append(handler_b.get_dtype(key))
layer_save_dtype = (
torch.float32 if torch.float32 in source_dtypes else save_torch_dtype
)
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
@@ -561,7 +584,18 @@ def extract_lora_from_files(
del weight_diff
return "skipped", None
# Skip 1D tensors
# SVD requires matrix-shaped tensors. Optionally preserve 1D weights
# and scale/lin parameters as exact runtime-applicable diff patches.
is_direct_diff = weight_diff.ndim == 1 or key.endswith(direct_diff_suffixes)
if include_1d_diffs and is_direct_diff:
direct_diff_dtype = torch.float32 if weight_diff.ndim == 1 else layer_save_dtype
layer_results = {
_format_direct_diff_key(key, lora_name):
weight_diff.to(direct_diff_dtype).cpu().contiguous()
}
del weight_diff
return "full", layer_results
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
@@ -588,8 +622,8 @@ def extract_lora_from_files(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(layer_save_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(layer_save_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
@@ -599,12 +633,12 @@ def extract_lora_from_files(
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.diff"] = weight_diff.to(layer_save_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(layer_save_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(layer_save_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
@@ -613,7 +647,7 @@ def extract_lora_from_files(
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
for key in tqdm(extraction_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
@@ -704,6 +738,13 @@ def _format_lora_key(key: str) -> str:
return f"diffusion_model.{key}"
def _format_direct_diff_key(source_key: str, lora_name: str) -> str:
"""Map a full delta to ComfyUI's canonical direct-patch key."""
if source_key.endswith(".bias"):
return f"{lora_name[:-5]}.diff_b"
return f"{lora_name}.diff"
def _get_common_inputs():
return [
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
@@ -716,13 +757,16 @@ def _get_common_inputs():
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.String.Input("output_filename", default="extracted_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16",
tooltip="Output dtype; layers stored as FP32 in either source remain FP32"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Include 1D tensors as FP32 and keys ending in .scale or .lin as full .diff patches"),
]
@@ -753,7 +797,8 @@ class LoRAExtractFixed(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
include_1d_diffs=False) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -763,7 +808,7 @@ class LoRAExtractFixed(io.ComfyNode):
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
lazy_load, force_clear_cache, glob_skip_patterns, include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -798,7 +843,8 @@ class LoRAExtractRatio(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
include_1d_diffs=False) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -809,7 +855,7 @@ class LoRAExtractRatio(io.ComfyNode):
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -844,7 +890,8 @@ class LoRAExtractQuantile(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
include_1d_diffs=False) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -855,7 +902,7 @@ class LoRAExtractQuantile(io.ComfyNode):
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -888,7 +935,8 @@ class LoRAExtractKnee(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
include_1d_diffs=False) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -899,7 +947,7 @@ class LoRAExtractKnee(io.ComfyNode):
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -934,7 +982,8 @@ class LoRAExtractFrobenius(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
include_1d_diffs=False) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -945,7 +994,7 @@ class LoRAExtractFrobenius(io.ComfyNode):
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
+484 -247
View File
File diff suppressed because it is too large Load Diff
+952 -229
View File
File diff suppressed because it is too large Load Diff
+105 -30
View File
@@ -12,6 +12,8 @@ from .device_utils import (
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
from .quantization_guard import inspect_low_bit_input, write_preserved_tensor
from .lora_resize import is_direct_diff_key
def load_documentation_from_file(filename):
@@ -105,6 +107,18 @@ class MergerLogic:
all_keys = primary_handler.keys()
metadata = primary_handler.metadata()
low_bit_keys_by_name = {}
if model_type == "loras":
try:
for name, handler in handlers.items():
low_bit_keys_by_name[name] = inspect_low_bit_input(
handler, f"LoRA ({name})", "Generic LoRA Merge"
)
except Exception:
for handler in handlers.values():
handler.__exit__(None, None, None)
raise
# Convert mismatch_mode string to enum
mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip')
mismatch_mode = MissingTensorBehavior(mismatch_mode_str)
@@ -132,6 +146,26 @@ class MergerLogic:
pbar = comfy.utils.ProgressBar(len(all_keys))
save_dtype = recipe_params.pop('save_dtype')
save_torch_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(save_dtype)
override_dtype = recipe_params.pop('override_dtype', False)
include_1d_diffs = recipe_params.pop('include_1d_diffs', False)
def determine_target_dtype(original_dtypes, target_dtype, override):
if override:
return target_dtype
priority_map = {
torch.float32: 3,
torch.bfloat16: 2,
torch.float16: 2,
}
highest_dtype = target_dtype
highest_priority = priority_map.get(highest_dtype, 1)
for dtype in original_dtypes:
p = priority_map.get(dtype, 1)
if p > highest_priority:
highest_priority = p
highest_dtype = dtype
return highest_dtype
recipe_params.update({"handlers": handlers})
process_device = recipe_params.get('device', 'cpu')
process_dtype = recipe_params.get('dtype', torch.float32)
@@ -142,8 +176,7 @@ class MergerLogic:
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]
output_dir = os.path.join(folder_paths.models_dir, output_folder)
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")
@@ -174,6 +207,14 @@ class MergerLogic:
with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
preserve_low_bit = model_type == "loras" and any(
key in low_bit_keys for low_bit_keys in low_bit_keys_by_name.values()
)
if preserve_low_bit:
_flush_batch()
write_preserved_tensor(writer, key, primary_handler)
pbar.update(1)
continue
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode):
discarded_keys += 1
@@ -181,6 +222,22 @@ class MergerLogic:
# Pre-load Model A's tensor with pinned memory for CUDA
cpu_tensor = primary_handler.get_tensor(key)
# Determine original dtypes across all active models for this key
original_dtypes = [primary_handler.get_dtype(key)]
for name, handler in handlers.items():
if name != primary_model_name and key in handler.keys():
original_dtypes.append(handler.get_dtype(key))
target_dtype = determine_target_dtype(original_dtypes, save_torch_dtype, override_dtype)
preserve_model_a_1d = model_type == "loras" and cpu_tensor.ndim == 1 and not include_1d_diffs
if preserve_model_a_1d:
target_dtype = determine_target_dtype(
[primary_handler.get_dtype(key)], save_torch_dtype, override_dtype
)
elif model_type == "loras" and is_direct_diff_key(key) and cpu_tensor.ndim == 1:
target_dtype = torch.float32
if process_device == 'cuda':
tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
else:
@@ -188,13 +245,14 @@ class MergerLogic:
del cpu_tensor
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode):
t = tensor_a.detach().to(save_torch_dtype).cpu()
if preserve_model_a_1d or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode):
t = tensor_a.detach().to(target_dtype).cpu()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
excluded_keys += 1
if not preserve_model_a_1d:
excluded_keys += 1
del tensor_a
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
@@ -224,7 +282,7 @@ class MergerLogic:
if isinstance(result, dict):
for r_key, r_tensor in result.items():
t = r_tensor.detach().to(save_torch_dtype).cpu()
t = r_tensor.detach().to(target_dtype).cpu()
batch_buffer[r_key] = t
batch_bytes += t.numel() * t.element_size()
@@ -238,7 +296,7 @@ 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]
t = result.detach().to(save_torch_dtype).cpu()
t = result.detach().to(target_dtype).cpu()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
@@ -275,8 +333,7 @@ class MergerLogic:
handler.__exit__(None, None, None)
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]
output_dir = os.path.join(folder_paths.models_dir, output_folder)
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")
@@ -323,6 +380,7 @@ class CheckpointTwoMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -335,7 +393,7 @@ class CheckpointTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -344,7 +402,7 @@ class CheckpointTwoMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -388,6 +446,7 @@ class ModelTwoMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -400,7 +459,7 @@ class ModelTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -409,7 +468,7 @@ class ModelTwoMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -453,6 +512,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -465,7 +525,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -474,7 +534,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -518,6 +578,9 @@ class LoRATwoMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force merged non-1D tensors to save_dtype. Enabled 1D direct diffs remain FP32."),
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -530,7 +593,8 @@ class LoRATwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool,
include_1d_diffs: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -539,11 +603,12 @@ class LoRATwoMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
"include_1d_diffs": include_1d_diffs,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -583,6 +648,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -595,7 +661,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -604,7 +670,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -651,6 +717,7 @@ class CheckpointThreeMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -663,7 +730,7 @@ class CheckpointThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -672,7 +739,7 @@ class CheckpointThreeMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -717,6 +784,7 @@ class ModelThreeMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -729,7 +797,7 @@ class ModelThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -738,7 +806,7 @@ class ModelThreeMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -783,6 +851,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -795,7 +864,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -804,7 +873,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -849,6 +918,9 @@ class LoRAThreeMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force merged non-1D tensors to save_dtype. Enabled 1D direct diffs remain FP32."),
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -861,7 +933,8 @@ class LoRAThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool,
include_1d_diffs: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -870,11 +943,12 @@ class LoRAThreeMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
"include_1d_diffs": include_1d_diffs,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -915,6 +989,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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"),
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -927,7 +1002,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -936,7 +1011,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"output_filename": output_filename, "save_dtype": save_dtype, "override_dtype": override_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
@@ -944,4 +1019,4 @@ class EmbeddingThreeMerger(io.ComfyNode):
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
return io.NodeOutput(filename, doc)
+145
View File
@@ -0,0 +1,145 @@
"""Fail-safe detection for quantized or isolated low-bit safetensors inputs."""
import logging
import json
import struct
import torch
MAX_ISOLATED_LOW_BIT_TENSORS = 2
def _low_bit_dtypes():
names = (
"int8",
"uint8",
"float4_e2m1fn_x2",
"float8_e4m3fn",
"float8_e4m3fnuz",
"float8_e5m2",
"float8_e5m2fnuz",
"float8_e8m0fnu",
)
return {getattr(torch, name) for name in names if hasattr(torch, name)}
LOW_BIT_DTYPES = _low_bit_dtypes()
LOW_BIT_STORAGE_CODES = {"I8", "U8", "F4"}
UPSTREAM_WRITER_STORAGE_CODES = {
"F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8",
"U64", "U32", "U16", "BOOL", "C64", "F8_E5M2", "F8_E4M3",
}
def _header(handler):
if getattr(handler, "_header", None) is not None:
return handler._header
with open(handler.filename, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
return json.loads(source.read(header_size).decode("utf-8"))
def storage_dtype_code(handler, key: str) -> str | None:
entry = _header(handler).get(key)
return entry.get("dtype") if entry else None
def _is_low_bit(handler, key: str) -> bool:
storage_code = storage_dtype_code(handler, key)
if storage_code in LOW_BIT_STORAGE_CODES or (
storage_code is not None and storage_code.startswith("F8_")
):
return True
return handler.get_dtype(key) in LOW_BIT_DTYPES
def inspect_low_bit_input(handler, label: str, operation: str) -> set[str]:
"""Reject quantized models; return isolated low-bit keys that must not be transformed."""
keys = list(handler.keys())
metadata = handler.metadata() or {}
quant_keys = [key for key in keys if key.endswith(".comfy_quant")]
has_quant_metadata = "_quantization_metadata" in metadata
has_legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in keys)
if quant_keys or has_quant_metadata or has_legacy_scaled_fp8:
raise ValueError(
f"[{operation}] {label} uses ComfyUI quantization metadata; quantized inputs are unsupported"
)
low_bit_keys = {key for key in keys if _is_low_bit(handler, key)}
if len(low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS:
raise ValueError(
f"[{operation}] {label} contains {len(low_bit_keys)} low-bit tensors; "
"casted or quantized inputs are unsupported"
)
if low_bit_keys:
logging.warning(
"[%s] %s contains %d isolated low-bit tensor(s); "
"preserving/excluding unchanged: %s",
operation,
label,
len(low_bit_keys),
", ".join(sorted(low_bit_keys)),
)
return low_bit_keys
def layer_has_low_bit(block_keys: dict, low_bit_keys: set[str]) -> bool:
return any(
block_keys.get(name) in low_bit_keys
for name in (
"down", "up", "alpha", "mid", "reshape", "dora_scale",
"diff", "diff_b", "w_norm", "b_norm", "set_weight",
)
)
def write_preserved_tensor(writer, key: str, handler, output_key: str | None = None) -> None:
"""Copy one tensor's original safetensors bytes and dtype into an active writer."""
destination = output_key or key
header = _header(handler)
entry = header[key]
if entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES:
writer.write(destination, handler.get_tensor(key).cpu().contiguous())
return
source_start, source_end = entry["data_offsets"]
byte_size = source_end - source_start
with writer._lock:
if destination in writer._manifest:
raise ValueError(f"Tensor '{destination}' has already been written.")
target_start = writer._current_data_offset
target_end = target_start + byte_size
writer._manifest[destination] = {
"dtype": entry["dtype"],
"shape": entry["shape"],
"data_offsets": [target_start, target_end],
}
writer._current_data_offset = target_end
source_path = handler.filename
source_absolute = 8 + getattr(handler, "_header_size", 0) + source_start
if not getattr(handler, "_header_size", None):
with open(source_path, "rb") as source:
source_absolute = 8 + struct.unpack("<Q", source.read(8))[0] + source_start
target_absolute = 8 + writer.max_header_bytes + target_start
def copy_raw():
try:
with open(source_path, "rb") as source:
source.seek(source_absolute)
remaining = byte_size
with writer._lock:
writer._file.seek(target_absolute)
while remaining:
chunk = source.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise EOFError(f"Unexpected EOF while preserving tensor '{key}'")
writer._file.write(chunk)
remaining -= len(chunk)
finally:
writer._semaphore.release()
writer._semaphore.acquire()
writer._futures.append(writer._executor.submit(copy_raw))
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.4.3"
version = "0.4.6"
license = { file = "LICENSE" }
classifiers = [
+840
View File
@@ -0,0 +1,840 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_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 cwb():
return _load_module("nodes.consensus_merger")
def _patch_io(monkeypatch, module, tmp_path, paths):
monkeypatch.setattr(module.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _result_path(tmp_path, category, result):
return tmp_path / category / result
def _params(output_filename, **overrides):
params = {
"cwb_preset": "custom",
"consensus_type": "mean",
"alignment_method": "index",
"alignment_threshold": 0.4,
"similarity_threshold": 0.0,
"power_alpha": 2.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,
"mismatch_mode": "skip",
"output_filename": output_filename,
"save_dtype": "fp16",
"process_device": "cpu",
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"override_dtype": False,
"include_1d_diffs": False,
}
params.update(overrides)
return params
def test_preset_resolution_matches_extended_contract(cwb):
power = cwb.resolve_cwb_settings(cwb_preset="power_blend")
assert power.consensus_type == "median"
assert power.alignment_threshold == pytest.approx(0.9)
assert power.power_alpha == pytest.approx(8.0)
assert power.dynamic_similarity_contrast is True
assert power.soft_comfort_bandpass is False
baseline = cwb.resolve_cwb_settings(
cwb_preset="baseline",
alignment_method="index",
dynamic_similarity_contrast=True,
soft_comfort_bandpass=True,
)
assert baseline.alignment_method == "similarity"
assert baseline.dynamic_similarity_contrast is False
assert baseline.soft_comfort_bandpass is False
varied = cwb.resolve_cwb_settings(
cwb_preset="varied_merge",
global_scale=1.25,
position_weight=0.2,
preserve_common_prefix=True,
)
assert varied.global_scale == pytest.approx(1.25)
assert varied.position_weight == pytest.approx(0.2)
assert varied.preserve_common_prefix is True
def test_cwb_math_index_prefix_and_scale(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=1.0,
)
first = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
second = torch.tensor([[3.0, 0.0], [0.0, 3.0]])
torch.testing.assert_close(
cwb.merge_cwb_tensors([first, second], settings),
torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
)
prefix_settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=2.0,
preserve_common_prefix=True,
)
prefixed = cwb.merge_cwb_tensors(
[torch.tensor([[9.0, 9.0], [1.0, 0.0]]),
torch.tensor([[9.0, 9.0], [3.0, 0.0]])],
prefix_settings,
)
torch.testing.assert_close(prefixed[0], torch.tensor([9.0, 9.0]))
torch.testing.assert_close(prefixed[1], torch.tensor([4.0, 0.0]))
def test_similarity_alignment_reorders_matching_rows(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
torch.testing.assert_close(
cwb.merge_cwb_tensors([reference, reversed_source], settings),
torch.tensor([[1.5, 0.0], [0.0, 1.5]]),
)
scores = torch.ones((2, 2))
biased = cwb._position_biased_scores(scores, 1.0)
assert biased[0, 0] > biased[0, 1]
def test_fixed_coordinate_merge_never_reorders_model_rows(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
expected = torch.stack([
cwb.merge_consensus_group(
torch.stack([reference[row], reversed_source[row]]), settings
)
for row in range(reference.shape[0])
])
torch.testing.assert_close(
cwb.merge_cwb_tensors(
[reference, reversed_source],
settings,
allow_similarity_alignment=False,
),
expected,
)
def test_lora_similarity_alignment_pairs_a_rows_with_b_columns(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference_down = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reference_up = torch.tensor([[2.0, 0.0], [0.0, 3.0]])
source_down = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
source_up = torch.tensor([[0.0, -2.0], [3.0, 0.0]])
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[reference_down, source_down],
[reference_up, source_up],
settings,
reference_index=0,
)
torch.testing.assert_close(merged_down, reference_down)
torch.testing.assert_close(merged_up, reference_up)
def test_lora_krea_shape_scores_only_rank_components(monkeypatch, cwb):
rank = 256
output_features = 36864
down = torch.ones((rank, 2))
up = torch.ones((output_features, rank))
matrix_shapes = []
def bounded_mm(left, right):
matrix_shapes.append((left.shape, right.shape))
return torch.eye(left.shape[0], right.shape[1], dtype=left.dtype)
monkeypatch.setattr(cwb.torch, "mm", bounded_mm)
monkeypatch.setattr(
cwb,
"merge_consensus_group",
lambda stacked, settings, **kwargs: stacked[0].clone(),
)
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
alignment_method="similarity",
alignment_threshold=0.5,
)
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[down, down], [up, up], settings, reference_index=0
)
assert merged_down.shape == down.shape
assert merged_up.shape == up.shape
assert matrix_shapes
assert all(left[0] == rank and right[1] == rank for left, right in matrix_shapes)
def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=3.0,
)
down = torch.ones((1, 1))
up = torch.ones((1, 1))
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
[(0, down, up, 2.0), (1, down, up, 2.0)],
settings,
"cpu",
torch.float32,
)
assert rank == 1
torch.testing.assert_close(merged_down, torch.ones((1, 1)))
torch.testing.assert_close(merged_up, torch.full((1, 1), 6.0))
def test_lora_pair_alignment_supports_convolution_factors(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
)
down = torch.arange(6, dtype=torch.float32).reshape(2, 3, 1, 1)
up = torch.arange(8, dtype=torch.float32).reshape(4, 2, 1, 1)
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[down, down], [up, up], settings, reference_index=0
)
torch.testing.assert_close(merged_down, down)
torch.testing.assert_close(merged_up, up)
def test_lora_cuda_oom_retries_current_pair_on_cpu(monkeypatch, caplog, cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
)
original_to_compute = cwb._to_compute
attempted_devices = []
def fail_cuda(tensor, device):
attempted_devices.append(device)
if str(device).startswith("cuda"):
raise torch.OutOfMemoryError("injected allocation failure")
return original_to_compute(tensor, device)
monkeypatch.setattr(cwb, "_to_compute", fail_cuda)
monkeypatch.setattr(cwb, "_release_failed_cuda_operation", lambda: None)
down = torch.ones((1, 2))
up = torch.ones((2, 1))
with caplog.at_level("WARNING"):
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
[(0, down, up, 1.0), (1, down, up, 1.0)],
settings,
"cuda",
torch.float32,
operation_label="diffusion_model.foo",
)
assert rank == 1
torch.testing.assert_close(merged_down, down)
torch.testing.assert_close(merged_up, up)
assert attempted_devices[0] == "cuda"
assert "cpu" in attempted_devices
assert "retrying this layer on CPU" in caplog.text
assert "diffusion_model.foo" in caplog.text
def test_norm_rescale_and_dsc_bandpass_are_finite(cwb):
rescale = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
rescale_norm=True,
)
merged = cwb.merge_consensus_group(
torch.tensor([[2.0, 0.0], [0.0, 2.0]]), rescale
)
assert torch.linalg.vector_norm(merged).item() == pytest.approx(2.0)
dsc = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="median",
diversity_beta=1.5,
dynamic_similarity_contrast=True,
soft_comfort_bandpass=True,
)
result = cwb.merge_consensus_group(
torch.tensor([[1.0, 0.0], [0.8, 0.2], [0.0, 1.0]]), dsc
)
assert torch.isfinite(result).all()
def test_all_ten_schemas_and_lora_switch_contract(cwb):
assert len(cwb.CWB_MERGER_NODES) == 10
assert len({node.NODE_ID for node in cwb.CWB_MERGER_NODES}) == 10
for node in cwb.CWB_MERGER_NODES:
schema = node.define_schema()
ids = [value.id for value in schema.inputs]
assert ids[:3] == ["execution_mode", "model_a", "model_b"]
if node.INPUT_COUNT == 3:
assert ids[3] == "model_c"
if node.LORA_MODE:
assert schema.inputs[-1].id == "include_1d_diffs"
assert schema.inputs[-1].default is False
else:
assert "include_1d_diffs" not in ids
def test_streaming_dtype_nonfloat_and_secondary_preservation(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "a.safetensors"
b = tmp_path / "b.safetensors"
save_file({
"shared": torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32),
"only_a": torch.tensor([5.0, 6.0], dtype=torch.float16),
"metadata_tensor": torch.tensor([1, 2], dtype=torch.int64),
}, str(a), metadata={"source": "A"})
save_file({
"shared": torch.tensor([[3.0, 0.0], [0.0, 3.0]], dtype=torch.float16),
"extra": torch.ones((2, 2)),
"metadata_tensor": torch.tensor([9, 9], dtype=torch.int64),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("anchored")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert set(tensors) == {"shared", "only_a", "extra", "metadata_tensor"}
torch.testing.assert_close(tensors["shared"], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
assert tensors["shared"].dtype == torch.float32
torch.testing.assert_close(tensors["only_a"], torch.tensor([5.0, 6.0], dtype=torch.float16))
torch.testing.assert_close(tensors["metadata_tensor"], torch.tensor([1, 2]))
torch.testing.assert_close(tensors["extra"], torch.ones((2, 2)))
overridden = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("anchored_override", override_dtype=True),
)
overridden_tensors = load_file(
str(_result_path(tmp_path, "diffusion_models", overridden))
)
assert overridden_tensors["shared"].dtype == torch.float16
assert overridden_tensors["metadata_tensor"].dtype == torch.int64
def test_failed_merge_preserves_existing_output_and_removes_temporary_file(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "atomic_a.safetensors"
b = tmp_path / "atomic_b.safetensors"
output = tmp_path / "diffusion_models" / "atomic_output.safetensors"
output.parent.mkdir(parents=True)
save_file({"layer": torch.ones((2, 2))}, str(a))
save_file({"layer": torch.full((2, 2), 2.0)}, str(b))
save_file({"existing": torch.tensor([7.0])}, str(output))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
def fail_merge(*args, **kwargs):
raise RuntimeError("injected CWB failure")
monkeypatch.setattr(cwb, "merge_cwb_tensors", fail_merge)
with pytest.raises(RuntimeError, match="injected CWB failure"):
cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("atomic_output")
)
existing = load_file(str(output))
assert set(existing) == {"existing"}
torch.testing.assert_close(existing["existing"], torch.tensor([7.0]))
assert not list(output.parent.glob(".atomic_output.safetensors.*.tmp"))
def test_embedding_union_uses_longest_first_dimension(monkeypatch, tmp_path, cwb):
a = tmp_path / "embed_a.safetensors"
b = tmp_path / "embed_b.safetensors"
save_file({"emb": torch.tensor([[1.0, 0.0], [0.0, 1.0]])}, str(a))
save_file({
"emb": torch.tensor([[3.0, 0.0], [0.0, 3.0], [4.0, 4.0]]),
"secondary_only": torch.tensor([[7.0, 0.0]]),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"embeddings",
_params("embedding_union"),
embedding_union=True,
)
tensors = load_file(str(_result_path(tmp_path, "embeddings", result)))
assert tensors["emb"].shape == (3, 2)
torch.testing.assert_close(tensors["emb"][:2], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
torch.testing.assert_close(tensors["emb"][2], torch.tensor([4.0, 4.0]))
torch.testing.assert_close(tensors["secondary_only"], torch.tensor([[7.0, 0.0]]))
def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb):
paths = {}
for name, value in (("a", 1.0), ("b", 3.0), ("c", 5.0)):
path = tmp_path / f"three_{name}.safetensors"
save_file({"layer": torch.tensor([[value, 0.0]])}, str(path))
paths[name] = str(path)
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("three_inputs")
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
torch.testing.assert_close(tensor, torch.tensor([[3.0, 0.0]]))
def test_two_input_preserves_secondary_only_tensor_exactly(monkeypatch, tmp_path, cwb):
a = tmp_path / "secondary_a.safetensors"
b = tmp_path / "secondary_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary": torch.tensor([2.0, 3.0], dtype=torch.float32),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params(
"secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
assert tensor.dtype == torch.float32
def test_three_input_merges_available_secondary_sources(monkeypatch, tmp_path, cwb):
a = tmp_path / "union_a.safetensors"
b = tmp_path / "union_b.safetensors"
c = tmp_path / "union_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([2.0]),
"b_only": torch.tensor([7.0]),
}, str(b))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([4.0]),
"c_only": torch.tensor([9.0]),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("secondary_union")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
torch.testing.assert_close(tensors["secondary_shared"], torch.tensor([3.0]))
torch.testing.assert_close(tensors["b_only"], torch.tensor([7.0]))
torch.testing.assert_close(tensors["c_only"], torch.tensor([9.0]))
def test_secondary_shape_conflict_skip_preserves_earliest_source(monkeypatch, tmp_path, cwb):
a = tmp_path / "shape_a.safetensors"
b = tmp_path / "shape_b.safetensors"
c = tmp_path / "shape_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({"secondary": torch.tensor([2.0, 3.0])}, str(b))
save_file({"secondary": torch.tensor([[8.0, 9.0]])}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"],
"diffusion_models",
_params("secondary_shape_skip", mismatch_mode="skip"),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
def test_secondary_only_filters_and_low_bit_preservation(monkeypatch, tmp_path, cwb):
a = tmp_path / "filter_a.safetensors"
b = tmp_path / "filter_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"secondary.keep": torch.tensor([4, 5], dtype=torch.uint8),
"secondary.drop": torch.tensor([6.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("secondary_filters", discard_patterns=r"secondary\.drop"),
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert "secondary.drop" not in tensors
torch.testing.assert_close(
tensors["secondary.keep"], torch.tensor([4, 5], dtype=torch.uint8)
)
def test_lora_cross_format_companion_group_preserves_model_a_canonically(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_a.safetensors"
b = tmp_path / "lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]], dtype=torch.float32),
"diffusion_model.foo.lora_B.weight": torch.tensor([[1.0], [2.0]], dtype=torch.float32),
"diffusion_model.foo.dora_scale": torch.tensor([1.0, 1.0]),
}, str(a))
save_file({
"lora_unet_foo.lora_down.weight": torch.tensor([[3.0, 4.0], [5.0, 6.0]]),
"lora_unet_foo.lora_up.weight": torch.tensor([[3.0, 5.0], [4.0, 6.0]]),
"lora_unet_secondary.lora_down.weight": torch.ones((1, 2)),
"lora_unet_secondary.lora_up.weight": torch.ones((2, 1)),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("lora_cross_format"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.dora_scale",
"lora_unet_secondary.lora_A.weight",
"lora_unet_secondary.lora_B.weight",
}
assert tensors["diffusion_model.foo.lora_A.weight"].shape == (1, 2)
assert tensors["diffusion_model.foo.lora_B.weight"].shape == (2, 1)
assert tensors["diffusion_model.foo.lora_A.weight"].dtype == torch.float32
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([1.0, 1.0])
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_A.weight"], torch.ones((1, 2))
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_B.weight"], torch.ones((2, 1))
)
def test_lora_secondary_only_pair_is_preserved_without_scaling(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_secondary_a.safetensors"
b = tmp_path / "lora_secondary_b.safetensors"
save_file({
"diffusion_model.base.diff": torch.tensor([[1.0]]),
}, str(a))
save_file({
"transformer.extra.lora_down.weight": torch.tensor([[2.0, 3.0]]),
"transformer.extra.lora_up.weight": torch.tensor([[4.0], [5.0]]),
"transformer.extra.alpha": torch.tensor(1.0),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params(
"lora_secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"],
torch.tensor([[2.0, 3.0]]),
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"],
torch.tensor([[4.0], [5.0]]),
)
assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.extra.alpha"].item() == pytest.approx(1.0)
def test_lora_three_input_merges_b_and_c_group_without_a(monkeypatch, tmp_path, cwb):
a = tmp_path / "lora_union_a.safetensors"
b = tmp_path / "lora_union_b.safetensors"
c = tmp_path / "lora_union_c.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.extra.lora_A.weight": torch.full((1, 2), 1.0),
"diffusion_model.extra.lora_B.weight": torch.full((2, 1), 1.0),
}, str(b))
save_file({
"transformer.extra.lora_down.weight": torch.full((1, 2), 3.0),
"transformer.extra.lora_up.weight": torch.full((2, 1), 3.0),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "loras", _params("lora_secondary_union"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.full((1, 2), 2.0)
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.full((2, 1), 2.0)
)
def test_lora_secondary_direct_and_companion_groups_are_atomic(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_atomic_a.safetensors"
b = tmp_path / "lora_atomic_b.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.norm.diff": torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
"diffusion_model.extra.lora_A.weight": torch.tensor([[4.0, 5.0]]),
"diffusion_model.extra.lora_B.weight": torch.tensor([[6.0], [7.0]]),
"diffusion_model.extra.dora_scale": torch.tensor([8.0, 9.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("lora_secondary_atomic", global_scale=3.0),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.norm.diff"],
torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
)
assert tensors["diffusion_model.norm.diff"].dtype == torch.bfloat16
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.tensor([[4.0, 5.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.tensor([[6.0], [7.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.dora_scale"], torch.tensor([8.0, 9.0])
)
def test_lora_rejects_duplicate_normalized_cores_in_secondary_input(cwb):
with pytest.raises(ValueError, match="same logical layer 'foo'.*input 2"):
cwb._secondary_lora_map(
{
"diffusion_model.foo": {},
"transformer.foo": {},
},
input_index=1,
)
def test_lora_peft_prefix_matching_and_existing_alpha(monkeypatch, tmp_path, cwb):
a = tmp_path / "peft_a.safetensors"
b = tmp_path / "peft_b.safetensors"
save_file({
"base_model.model.diffusion_model.foo.lora_A.weight": torch.ones((1, 2)),
"base_model.model.diffusion_model.foo.lora_B.weight": torch.ones((2, 1)),
"base_model.model.diffusion_model.foo.alpha": torch.tensor(1.0),
}, str(a))
save_file({
"diffusion_model.foo.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.foo.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.foo.alpha": torch.tensor(2.0),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("peft_prefix"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.alpha",
}
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(2.0)
def test_lora_mochi_inputs_emit_preferred_canonical_output(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "mochi_a.safetensors"
b = tmp_path / "mochi_b.safetensors"
for path, value in ((a, 1.0), (b, 3.0)):
save_file({
"diffusion_model.foo.lora_A": torch.full((1, 2), value),
"diffusion_model.foo.lora_B": torch.full((2, 1), value),
"diffusion_model.foo.alpha": torch.tensor(1.0),
}, str(path))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("mochi_canonical"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.alpha",
}
def test_lora_1d_direct_switch_is_default_fallback_and_fp32(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "direct_a.safetensors"
b = tmp_path / "direct_b.safetensors"
save_file({"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16)}, str(a))
save_file({"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16)}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
default = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("direct_default"), lora_mode=True
)
default_tensor = load_file(str(_result_path(tmp_path, "loras", default)))[
"diffusion_model.norm.diff"
]
torch.testing.assert_close(default_tensor, torch.tensor([1.0, 2.0], dtype=torch.bfloat16))
assert default_tensor.dtype == torch.bfloat16
enabled = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("direct_enabled", include_1d_diffs=True, override_dtype=True),
lora_mode=True,
)
enabled_tensor = load_file(str(_result_path(tmp_path, "loras", enabled)))[
"diffusion_model.norm.diff"
]
torch.testing.assert_close(enabled_tensor, torch.tensor([2.0, 3.0]))
assert enabled_tensor.dtype == torch.float32
def test_quantized_marker_errors_before_output(monkeypatch, tmp_path, cwb):
a = tmp_path / "quant_a.safetensors"
b = tmp_path / "quant_b.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor([1], dtype=torch.uint8),
}, str(a))
save_file({"layer.weight": torch.ones((2, 2))}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("must_not_exist")
)
assert not (tmp_path / "diffusion_models" / "must_not_exist.safetensors").exists()
def test_isolated_low_bit_preserves_anchored_tensor(monkeypatch, tmp_path, cwb):
a = tmp_path / "guard_a.safetensors"
b = tmp_path / "guard_b.safetensors"
save_file({"layer": torch.tensor([1.0, 2.0])}, str(a))
save_file({"layer": torch.tensor([8, 9], dtype=torch.uint8)}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("guarded_generic")
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
torch.testing.assert_close(tensor, torch.tensor([1.0, 2.0]))
assert tensor.dtype == torch.float32
def test_isolated_low_bit_preserves_complete_lora_layer(monkeypatch, tmp_path, cwb):
a = tmp_path / "guard_lora_a.safetensors"
b = tmp_path / "guard_lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.foo.alpha": torch.tensor(1.0),
"diffusion_model.foo.dora_scale": torch.tensor([5.0, 6.0]),
}, str(a))
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[9.0, 9.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.foo.alpha": torch.tensor(1, dtype=torch.uint8),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("guarded_lora"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.foo.lora_A.weight"], torch.tensor([[1.0, 2.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.foo.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(1.0)
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([5.0, 6.0])
)
+899
View File
@@ -0,0 +1,899 @@
import importlib
import logging
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_contract_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 modules():
return _load_module("nodes.lora_resize"), _load_module("nodes.lora_merger")
@pytest.fixture(scope="module")
def generic_modules():
return _load_module("nodes.merger"), _load_module("nodes.merger_ops")
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
COMFY_LORA_PAIR_SUFFIXES = [
(".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"),
(".lora_A.default.weight", ".lora_B.default.weight"),
]
@pytest.mark.parametrize("down_suffix,up_suffix", COMFY_LORA_PAIR_SUFFIXES)
def test_all_comfy_lora_pair_spellings_parse_to_canonical_keys(
modules, down_suffix, up_suffix
):
resize, _ = modules
block = "diffusion_model.layer"
layers, passthrough = resize.parse_lora_layers([
f"{block}{down_suffix}",
f"{block}{up_suffix}",
f"{block}.alpha",
])
assert not passthrough
assert layers[block]["down"] == f"{block}{down_suffix}"
assert layers[block]["up"] == f"{block}{up_suffix}"
assert resize.canonical_lora_key(block, "down") == f"{block}.lora_A.weight"
assert resize.canonical_lora_key(block, "up") == f"{block}.lora_B.weight"
@pytest.mark.parametrize("down_suffix,up_suffix", COMFY_LORA_PAIR_SUFFIXES)
def test_merge_to_model_applies_every_comfy_lora_pair_spelling(
monkeypatch, tmp_path, modules, down_suffix, up_suffix
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
variant = COMFY_LORA_PAIR_SUFFIXES.index((down_suffix, up_suffix))
base = tmp_path / f"base_{variant}.safetensors"
adapter = tmp_path / f"adapter_{variant}.safetensors"
block = "diffusion_model.layer"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
}, str(base))
save_file({
f"{block}{down_suffix}": torch.tensor([[1.0, 2.0]]),
f"{block}{up_suffix}": torch.tensor([[3.0], [4.0]]),
f"{block}.alpha": torch.tensor(0.5),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
f"merged_variant_{variant}", verbose=False,
)
result = load_file(output)["model.diffusion_model.layer.weight"]
torch.testing.assert_close(
result,
torch.tensor([[1.5, 3.0], [2.0, 4.0]]),
)
assert result.dtype == torch.float32
def test_merge_to_model_applies_comfy_direct_norm_and_set_forms(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "direct_base.safetensors"
adapter = tmp_path / "direct_adapter.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.norm.bias": torch.zeros(2),
"model.diffusion_model.replace.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.w_norm": torch.tensor([1.0, 2.0]),
"diffusion_model.norm.b_norm": torch.tensor([3.0, 4.0]),
"diffusion_model.replace.set_weight": torch.full((2, 2), 7.0),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16,
"direct_forms", verbose=False, include_1d_diffs=True,
)
result = load_file(output)
torch.testing.assert_close(
result["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0])
)
torch.testing.assert_close(
result["model.diffusion_model.norm.bias"], torch.tensor([1.5, 2.0])
)
torch.testing.assert_close(
result["model.diffusion_model.replace.weight"], torch.full((2, 2), 7.0)
)
def test_merge_to_model_uses_comfy_reshape_companion(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "reshape_base.safetensors"
adapter = tmp_path / "reshape_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((1, 2)),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.layer.reshape_weight": torch.tensor([2, 2]),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
"reshape_companion", verbose=False,
)
result = load_file(output)["model.diffusion_model.layer.weight"]
assert result.shape == (2, 2)
torch.testing.assert_close(result, torch.tensor([[3.0, 6.0], [4.0, 8.0]]))
def test_merge_to_model_reports_success_and_separates_base_only_biases(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "summary_base.safetensors"
adapter = tmp_path / "summary_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"summary", verbose=True, return_report=True,
)
output = capsys.readouterr().out
assert Path(path).is_file()
assert report.startswith("SUCCESS: 1/1 adapter groups fully applied")
assert "[PATCHED]" not in report
assert "[APPLIED]" not in report
assert "L1: summary_adapter.safetensors, strength=1.0" in report
assert "[BASE-ONLY BIASES RETAINED] (1)" in report
assert "model.diffusion_model.layer.bias" in report
assert "SUCCESS: 1/1 groups applied" in output
assert "1 tensors patched, 1 base tensors not patched, 0 adapter issues" in output
assert "model.diffusion_model.layer.bias" not in output
def test_success_report_stays_compact_and_console_is_summary_only(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "many_success_base.safetensors"
adapter = tmp_path / "many_success_adapter.safetensors"
save_file({
f"model.diffusion_model.layer_{index}.weight": torch.zeros((2, 2))
for index in range(40)
}, str(base))
save_file({
f"diffusion_model.layer_{index}.diff": torch.ones((2, 2))
for index in range(40)
}, str(adapter))
_, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"many_success", verbose=True, return_report=True,
)
console = capsys.readouterr().out
assert report.startswith("SUCCESS: 40/40")
assert "layer_0.weight" not in report
assert "ADAPTER GROUP ISSUES\n- None" in report
assert len(report) < 1000
assert max(map(len, report.splitlines())) < 160
assert "layer_0.weight" not in console
assert len(console.splitlines()) <= 9
def test_merge_to_model_report_identifies_shape_mismatch_and_unresolved_group(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "diagnostic_base.safetensors"
adapter = tmp_path / "diagnostic_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.layer.diff": torch.ones(3),
"diffusion_model.orphan.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.orphan.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
_, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"diagnostic", verbose=False, include_1d_diffs=True, return_report=True,
)
assert report.startswith("COMPLETED WITH ERRORS: 0/2")
assert "[MAPPED BUT UNCHANGED] (1)" in report
assert "[SHAPE MISMATCH] (1)" in report
assert "source_shape=(3,), target_shape=(2, 2)" in report
assert "[UNRESOLVED] (1)" in report
assert "diffusion_model.orphan" in report
assert "no matching base tensor" in report
def test_merge_to_model_detects_comfy_adapter_logged_failure_atomically(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "logged_failure_base.safetensors"
adapter = tmp_path / "logged_failure_adapter.safetensors"
base_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
save_file({"model.diffusion_model.layer.weight": base_weight}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
class LoggedFailureAdapter:
def calculate_weight(self, weight, key, *args, **kwargs):
weight.add_(99.0)
logging.error("ERROR lora %s synthetic internal failure", key)
return weight
class FakeLoRAAdapter:
@staticmethod
def load(*args, **kwargs):
return LoggedFailureAdapter()
monkeypatch.setattr(resize, "LoRAAdapter", FakeLoRAAdapter)
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"logged_failure", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.layer.weight"], base_weight
)
assert report.startswith("COMPLETED WITH ERRORS: 0/1")
assert "[CALCULATION FAILURE] (1)" in report
assert "synthetic internal failure" in report
assert "[MAPPED BUT UNCHANGED] (1)" in report
def test_merge_to_model_report_classifies_disabled_1d_as_exclusion(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "disabled_1d_base.safetensors"
adapter = tmp_path / "disabled_1d_adapter.safetensors"
save_file({"model.diffusion_model.norm.weight": torch.zeros(2)}, str(base))
save_file({"diffusion_model.norm.diff": torch.ones(2)}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"disabled_1d", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.norm.weight"], torch.zeros(2)
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 0/1")
assert "[1D DISABLED] (1)" in report
assert "source=diffusion_model.norm.diff" in report
assert "include_1d_diffs is disabled" in report
def test_merge_to_model_report_marks_partially_applied_group(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "partial_base.safetensors"
adapter = tmp_path / "partial_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
"diffusion_model.layer.diff_b": torch.ones(3),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"partial", verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(path)
torch.testing.assert_close(
result["model.diffusion_model.layer.weight"], torch.ones((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.layer.bias"], torch.zeros(2)
)
assert report.startswith("COMPLETED WITH ERRORS: 0/1")
assert "[PARTIALLY APPLIED] (1)" in report
assert "source_shape=(3,), target_shape=(2,)" in report
def test_merge_to_model_detects_raised_adapter_failure_atomically(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "raised_failure_base.safetensors"
adapter = tmp_path / "raised_failure_adapter.safetensors"
base_weight = torch.eye(2)
save_file({"model.diffusion_model.layer.weight": base_weight}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
class RaisedFailureAdapter:
def calculate_weight(self, weight, key, *args, **kwargs):
weight.zero_()
raise RuntimeError("synthetic raised failure")
class FakeLoRAAdapter:
@staticmethod
def load(*args, **kwargs):
return RaisedFailureAdapter()
monkeypatch.setattr(resize, "LoRAAdapter", FakeLoRAAdapter)
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"raised_failure", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.layer.weight"], base_weight
)
assert "[CALCULATION FAILURE] (1)" in report
assert "RuntimeError: synthetic raised failure" in report
def test_merge_to_model_report_lists_skipped_tensor_and_pattern(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "skip_base.safetensors"
adapter = tmp_path / "skip_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"skip", skip_patterns_str=r"layer\.weight$", verbose=False,
return_report=True,
)
assert "model.diffusion_model.layer.weight" not in load_file(path)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 0/1")
assert "[OMITTED BY SKIP PATTERN] (1)" in report
assert r"pattern=layer\.weight$" in report
assert "[SKIPPED BY PATTERN] (1)" in report
def test_merge_to_model_node_returns_path_blank_line_report(
monkeypatch, modules
):
resize, _ = modules
monkeypatch.setattr(
resize.folder_paths, "get_full_path_or_raise",
lambda category, name: f"{category}/{name}",
)
monkeypatch.setattr(
resize, "merge_loras_to_model",
lambda **kwargs: ("saved/model.safetensors", "SUCCESS\nDETAILS"),
)
result = resize.LoRAMergeToModel.execute(
"base.safetensors", "1",
"one.safetensors", 1.0,
"None", 1.0, "None", 1.0, "None", 1.0,
"None", 1.0, "None", 1.0, "None", 1.0, "None", 1.0,
"", "merged", "fp32", "cpu", True, False, False,
)
assert result.result == ("saved/model.safetensors\n\nSUCCESS\nDETAILS",)
def test_mixed_parser_scans_every_key_and_retains_unpaired(modules):
resize, _ = modules
filler = [f"metadata.tensor_{index}" for index in range(60)]
keys = filler + [
"diffusion_model.linear.lora_down.weight",
"diffusion_model.linear.lora_up.weight",
"diffusion_model.linear.alpha",
"diffusion_model.norm.diff",
"diffusion_model.block.diff_b",
"diffusion_model.orphan.lora_down.weight",
"diffusion_model.linear.dora_scale",
]
info = resize.detect_lora_format(keys)
layers, passthrough = resize.parse_lora_layers(keys)
assert info["format"] == "mixed"
assert layers["diffusion_model.linear"]["down"].endswith("lora_down.weight")
assert layers["diffusion_model.norm"]["diff"].endswith(".diff")
assert layers["diffusion_model.block"]["diff_b"].endswith(".diff_b")
assert "diffusion_model.orphan.lora_down.weight" in passthrough
assert layers["diffusion_model.linear"]["dora_scale"].endswith(".dora_scale")
def test_dtype_contract(modules):
resize, _ = modules
assert resize.select_output_dtype([torch.bfloat16], torch.float16, is_1d_diff=True) == torch.float32
assert resize.select_output_dtype([torch.float32, torch.bfloat16], torch.float16) == torch.float32
assert resize.select_output_dtype([torch.float32], torch.bfloat16, force=True) == torch.bfloat16
def test_include_1d_switches_are_appended_and_default_off(modules, generic_modules):
resize, multi = modules
generic, _ = generic_modules
node_classes = [
generic.LoRATwoMerger,
generic.LoRAThreeMerger,
multi.LoRAMultiMerge,
multi.LoRAMultiMergeDARE,
multi.LoRAMultiMergeDAREEnhanced,
resize.LoRAMergeToModel,
]
for node_class in node_classes:
final_input = node_class.define_schema().inputs[-1]
assert final_input.id == "include_1d_diffs"
assert final_input.default is False
assert len(resize.LoRAMergeToModel.define_schema().outputs) == 1
def test_resize_preserves_mixed_direct_auxiliary_and_fp32(monkeypatch, tmp_path, modules):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mixed.safetensors"
save_file({
"diffusion_model.linear.lora_down.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4),
"diffusion_model.linear.lora_up.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2),
"diffusion_model.linear.alpha": torch.tensor(2.0, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([0.125, -0.25], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.linear.dora_scale": torch.ones(4, dtype=torch.float32),
"diffusion_model.orphan.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resized", verbose=False
)
tensors = load_file(output)
assert tensors["diffusion_model.linear.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.linear.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.lora_B.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.alpha"].dtype == torch.float32
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.scale.diff"].dtype == torch.float16
assert tensors["diffusion_model.lin.diff"].dtype == torch.float32
assert tensors["diffusion_model.linear.dora_scale"].dtype == torch.float32
assert tensors["diffusion_model.orphan.lora_down.weight"].dtype == torch.float16
pure_source = tmp_path / "pure_diff.safetensors"
save_file({
"diffusion_model.pure.diff": torch.ones((2, 2), dtype=torch.bfloat16),
"diffusion_model.pure_norm.diff": torch.ones(2, dtype=torch.bfloat16),
}, str(pure_source))
pure_output = resize.resize_lora_file(
str(pure_source), 1, None, None, "cpu", torch.float16, "pure_resized", verbose=False
)
pure_tensors = load_file(pure_output)
assert set(pure_tensors) == {"diffusion_model.pure.diff", "diffusion_model.pure_norm.diff"}
assert pure_tensors["diffusion_model.pure.diff"].dtype == torch.float16
assert pure_tensors["diffusion_model.pure_norm.diff"].dtype == torch.float32
def test_resize_normalizes_mochi_pair_to_preferred_output(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mochi.safetensors"
save_file({
"diffusion_model.layer.lora_A": torch.ones((2, 2)),
"diffusion_model.layer.lora_B": torch.ones((2, 2)),
"diffusion_model.layer.alpha": torch.tensor(2.0),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float32,
"mochi_normalized", verbose=False,
)
tensors = load_file(output)
assert set(tensors) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
"diffusion_model.layer.alpha",
}
def test_standard_multi_merge_normalizes_mochi_pair_to_preferred_output(
monkeypatch, tmp_path, modules
):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "multi_mochi_a.safetensors"
second = tmp_path / "multi_mochi_b.safetensors"
for path, value in ((first, 1.0), (second, 3.0)):
save_file({
"diffusion_model.layer.lora_A": torch.full((1, 2), value),
"diffusion_model.layer.lora_B": torch.full((2, 1), value),
"diffusion_model.layer.alpha": torch.tensor(1.0),
}, str(path))
output = merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float32, "multi_mochi_normalized", verbose=False,
)
tensors = load_file(output)
assert set(tensors) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
"diffusion_model.layer.alpha",
}
def test_standard_multi_merge_emits_weighted_direct_union(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "first.safetensors"
second = tmp_path / "second.safetensors"
base = tmp_path / "multi_base.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.matrix.weight": torch.zeros((2, 2)),
"model.diffusion_model.only_first.weight": torch.zeros(()),
"model.diffusion_model.scale": torch.zeros(()),
"model.diffusion_model.lowrank.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.only_first.diff": torch.tensor(4.0, dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(1.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.full((1, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.full((2, 1), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.bfloat16),
}, str(second))
output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged", base_model_path=str(base), verbose=False,
include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([-1.0, 0.0]))
torch.testing.assert_close(tensors["diffusion_model.matrix.diff"], torch.zeros((2, 2)))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.matrix.diff"].dtype == torch.float32
assert tensors["diffusion_model.only_first.diff"].item() == pytest.approx(8.0)
assert tensors["diffusion_model.only_first.diff"].dtype == torch.float16
assert tensors["diffusion_model.scale.diff"].item() == pytest.approx(2.0)
assert tensors["diffusion_model.lowrank.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.lowrank.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_B.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.alpha"].dtype == torch.float32
default_output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged_default", base_model_path=str(base), verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.matrix.diff" in default_tensors
assert "diffusion_model.scale.diff" in default_tensors
assert "diffusion_model.lowrank.lora_A.weight" in default_tensors
def test_multi_merge_rejects_direct_shape_mismatch(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "shape_first.safetensors"
second = tmp_path / "shape_second.safetensors"
save_file({"diffusion_model.bad.diff": torch.ones(2)}, str(first))
save_file({"diffusion_model.bad.diff": torch.ones(3)}, str(second))
with pytest.raises(ValueError, match="Direct LoRA shape mismatch"):
merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float16, "bad_shapes", verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("enhanced", [False, True])
def test_dare_variants_emit_direct_layers(monkeypatch, tmp_path, modules, enhanced):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / f"dare_first_{enhanced}.safetensors"
second = tmp_path / f"dare_second_{enhanced}.safetensors"
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.float32),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.bfloat16),
}, str(second))
if enhanced:
output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced", verbose=False, include_1d_diffs=True,
)
else:
output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare", verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0]))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_B.weight"].dtype == torch.float32
if enhanced:
default_output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced_default", verbose=False,
)
else:
default_output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare_default", verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.lowrank.lora_A.weight" in default_tensors
def test_merge_to_model_applies_generic_direct_keys_and_preserves_dtype(
monkeypatch, tmp_path, modules, caplog
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "base.safetensors"
adapter = tmp_path / "adapter.safetensors"
save_file({
"model.diffusion_model.block.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.norm.weight": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.block.bias": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.scale": torch.tensor(1.0, dtype=torch.bfloat16),
"model.diffusion_model.lin": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.bad.weight": torch.ones((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.untouched": torch.tensor(1e-8, dtype=torch.float32),
}, str(base))
save_file({
"diffusion_model.block.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.block.diff_b": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.bad.diff": torch.ones(3, dtype=torch.float32),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model",
verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0]))
torch.testing.assert_close(tensors["model.diffusion_model.block.bias"], torch.tensor([1.5, 2.0]))
assert tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5))
torch.testing.assert_close(tensors["model.diffusion_model.bad.weight"].float(), torch.ones((2, 2)))
assert "Shape mismatch" in caplog.text
assert tensors["model.diffusion_model.block.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.norm.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.block.bias"].dtype == torch.float32
assert tensors["model.diffusion_model.scale"].dtype == torch.float16
assert tensors["model.diffusion_model.lin"].dtype == torch.float32
assert tensors["model.diffusion_model.untouched"].dtype == torch.float32
default_output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model_default",
verbose=False,
)
default_tensors = load_file(default_output)
torch.testing.assert_close(
default_tensors["model.diffusion_model.norm.weight"].float(), torch.zeros(2)
)
torch.testing.assert_close(
default_tensors["model.diffusion_model.block.bias"].float(), torch.zeros(2)
)
assert default_tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(
default_tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5)
)
def _generic_params(output_filename, override_dtype, include_1d_diffs=False):
return {
"mismatch_mode": "skip",
"alignment_mode": "pad/crop",
"alpha": 0.5,
"beta": 0.0,
"gamma": 0.5,
"delta": 2.0,
"epsilon": 0.01,
"zeta": 0.0,
"seed": 0,
"output_filename": output_filename,
"save_dtype": "bf16",
"override_dtype": override_dtype,
"device": "cpu",
"dtype": torch.float32,
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"include_1d_diffs": include_1d_diffs,
}
def test_generic_two_and_three_mergers_enforce_direct_dtype_contract(
monkeypatch, tmp_path, generic_modules
):
generic, operations = generic_modules
paths = {}
for name, tensors in {
"a": {
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
},
"b": {
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 3.0, dtype=torch.bfloat16),
"diffusion_model.secondary_only.diff": torch.ones(1, dtype=torch.float32),
},
"c": {
"diffusion_model.norm.diff": torch.tensor([0.5, 1.0], dtype=torch.float32),
"diffusion_model.matrix.diff": torch.full((2, 2), 0.5, dtype=torch.bfloat16),
},
}.items():
path = tmp_path / f"{name}.safetensors"
save_file(tensors, str(path))
paths[name] = str(path)
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
two_params = _generic_params("generic_two", override_dtype=True)
two_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_params,
"loras",
)
two = load_file(str(tmp_path / "loras" / "generic_two.safetensors"))
torch.testing.assert_close(
two["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert two["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert two["diffusion_model.matrix.diff"].dtype == torch.bfloat16
assert "diffusion_model.secondary_only.diff" not in two
two_enabled_params = _generic_params(
"generic_two_enabled", override_dtype=True, include_1d_diffs=True
)
two_enabled_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_enabled_params,
"loras",
)
two_enabled = load_file(
str(tmp_path / "loras" / "generic_two_enabled.safetensors")
)
torch.testing.assert_close(
two_enabled["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0])
)
assert two_enabled["diffusion_model.norm.diff"].dtype == torch.float32
three_params = _generic_params("generic_three", override_dtype=False)
three_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_params,
"loras",
)
three = load_file(str(tmp_path / "loras" / "generic_three.safetensors"))
torch.testing.assert_close(
three["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert three["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert three["diffusion_model.matrix.diff"].dtype == torch.float32
three_enabled_params = _generic_params(
"generic_three_enabled", override_dtype=False, include_1d_diffs=True
)
three_enabled_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_enabled_params,
"loras",
)
three_enabled = load_file(
str(tmp_path / "loras" / "generic_three_enabled.safetensors")
)
torch.testing.assert_close(
three_enabled["diffusion_model.norm.diff"], torch.tensor([2.25, 3.5])
)
assert three_enabled["diffusion_model.norm.diff"].dtype == torch.float32
+437
View File
@@ -0,0 +1,437 @@
import importlib
import json
import logging
import struct
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_quant_guard_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 modules():
return {
"guard": _load_module("nodes.quantization_guard"),
"extract": _load_module("nodes.lora_extract_svd"),
"resize": _load_module("nodes.lora_resize"),
"multi": _load_module("nodes.lora_merger"),
"generic": _load_module("nodes.merger"),
"operations": _load_module("nodes.merger_ops"),
}
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _inspect_file(path, modules, label="input"):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
return modules["guard"].inspect_low_bit_input(handler, label, "Guard Test")
finally:
handler.__exit__(None, None, None)
def _round_trip_preserved(path, output, modules):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
with modules["resize"].IncrementalSafetensorsWriter(str(output)) as writer:
modules["guard"].write_preserved_tensor(writer, "low", handler)
finally:
handler.__exit__(None, None, None)
return load_file(str(output))["low"]
def _tensor_payload(path, key="low"):
with open(path, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
header = json.loads(source.read(header_size).decode("utf-8"))
start, end = header[key]["data_offsets"]
source.seek(8 + header_size + start)
return source.read(end - start)
@pytest.mark.parametrize("dtype", [torch.int8, torch.uint8])
def test_int8_uint8_are_detected_and_warn_once(tmp_path, modules, caplog, dtype):
path = tmp_path / f"{dtype}.safetensors"
save_file({"low": torch.arange(4, dtype=dtype), "normal": torch.ones(1)}, str(path))
with caplog.at_level(logging.WARNING):
detected = _inspect_file(path, modules, "threshold input")
assert detected == {"low"}
records = [record for record in caplog.records if "threshold input" in record.message]
assert len(records) == 1
assert "1 isolated low-bit tensor" in records[0].message
assert "low" in records[0].message
def test_exact_two_tensor_threshold(tmp_path, modules):
for count in range(4):
path = tmp_path / f"threshold_{count}.safetensors"
tensors = {f"low_{i}": torch.ones(2, dtype=torch.uint8) for i in range(count)}
tensors["normal"] = torch.ones(1)
save_file(tensors, str(path))
if count <= 2:
assert len(_inspect_file(path, modules)) == count
else:
with pytest.raises(ValueError, match="3 low-bit tensors.*casted or quantized"):
_inspect_file(path, modules)
@pytest.mark.parametrize(
("tensors", "metadata"),
[
({"layer.comfy_quant": torch.tensor([123], dtype=torch.uint8)}, None),
({"normal": torch.ones(1)}, {"_quantization_metadata": "{}"}),
({"scaled_fp8": torch.ones(2)}, None),
],
)
def test_recognized_comfy_quantization_hard_errors(tmp_path, modules, tensors, metadata):
path = tmp_path / f"quant_{len(list(tmp_path.iterdir()))}.safetensors"
save_file(tensors, str(path), metadata=metadata)
with pytest.raises(ValueError, match="ComfyUI quantization metadata.*unsupported"):
_inspect_file(path, modules, "recognized input")
def test_supported_fp8_storage_round_trip(tmp_path, modules):
candidates = [
dtype
for name in (
"float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz",
"float8_e8m0fnu",
)
if (dtype := getattr(torch, name, None)) is not None
]
supported = []
for index, dtype in enumerate(candidates):
path = tmp_path / f"low_float_{index}.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
continue
supported.append(dtype)
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / f"preserved_{index}.safetensors", modules)
assert preserved.dtype == dtype
torch.testing.assert_close(preserved.float(), loaded.float())
if candidates and not supported:
pytest.skip("Installed safetensors cannot write any installed torch FP8 dtype")
def test_fp4_storage_when_supported(tmp_path, modules):
dtype = getattr(torch, "float4_e2m1fn_x2", None)
if dtype is None:
pytest.skip("Installed PyTorch has no torch.float4_e2m1fn_x2")
path = tmp_path / "fp4.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
pytest.skip("Installed safetensors cannot write torch.float4_e2m1fn_x2")
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / "preserved_fp4.safetensors", modules)
assert preserved.dtype == dtype
assert _tensor_payload(path) == _tensor_payload(tmp_path / "preserved_fp4.safetensors")
def test_realistic_packed_uint8_comfy_sidecar_is_rejected(tmp_path, modules):
path = tmp_path / "packed_comfy.safetensors"
save_file({
"layer.weight": torch.arange(8, dtype=torch.uint8).reshape(2, 4),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
"layer.weight_scale": torch.ones(2),
}, str(path))
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
_inspect_file(path, modules)
def test_resize_preserves_complete_low_bit_layer_and_resizes_normal(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "resize_quant.safetensors"
guarded_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"diffusion_model.guarded.lora_down.weight": guarded_down,
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.guarded.alpha": torch.tensor(2.0),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.normal.lora_up.weight": torch.ones((2, 2)),
"isolated_aux": torch.tensor([7], dtype=torch.int8),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resize_guarded", verbose=False
)
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], guarded_down)
assert result["diffusion_model.guarded.lora_A.weight"].dtype == torch.uint8
assert result["diffusion_model.guarded.lora_B.weight"].dtype == torch.float32
assert result["diffusion_model.guarded.dora_scale"].dtype == torch.float32
assert result["diffusion_model.normal.lora_A.weight"].shape[0] == 1
assert result["isolated_aux"].dtype == torch.int8
def test_extract_excludes_low_bit_candidate_and_extracts_float_layer(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
model_a = tmp_path / "extract_a.safetensors"
model_b = tmp_path / "extract_b.safetensors"
output = tmp_path / "extracted.safetensors"
save_file({
"model.diffusion_model.normal.weight": torch.eye(2),
"model.diffusion_model.quant.weight": torch.ones((2, 2), dtype=torch.uint8),
}, str(model_a))
save_file({
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
"model.diffusion_model.quant.weight": torch.zeros((2, 2), dtype=torch.uint8),
}, str(model_b))
extract.extract_lora_from_files(
str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
force_clear_cache=False,
)
result = load_file(str(output))
assert any("normal" in key for key in result)
assert not any("quant" in key for key in result)
def test_extract_bias_diff_round_trips_through_merge_to_model(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
resize = modules["resize"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
_patch_output(monkeypatch, resize, tmp_path)
finetuned = tmp_path / "bias_finetuned.safetensors"
base = tmp_path / "bias_base.safetensors"
adapter = tmp_path / "bias_adapter.safetensors"
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.25, -0.5]),
}, str(finetuned))
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.0, 0.5]),
}, str(base))
extract.extract_lora_from_files(
str(finetuned), str(base), "fixed", 1, 1, "cpu", "fp16", str(adapter),
force_clear_cache=False, include_1d_diffs=True,
)
extracted = load_file(str(adapter))
assert set(extracted) == {"diffusion_model.layer.diff_b"}
assert extracted["diffusion_model.layer.diff_b"].dtype == torch.float32
merged_path = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"bias_round_trip", verbose=False, include_1d_diffs=True,
)
merged = load_file(merged_path)
torch.testing.assert_close(
merged["model.diffusion_model.layer.bias"],
load_file(str(finetuned))["model.diffusion_model.layer.bias"],
)
def _run_multi_variant(module, variant, paths, output_name):
common = (paths, [1.0, 1.0])
if variant == "standard":
return module.merge_multi_loras(
*common, "weighted_sum", "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
if variant == "dare":
return module.merge_multi_loras_dare(
*common, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
return module.merge_multi_loras_dare_enhanced(
*common, 1.0, 1.0, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("variant", ["standard", "dare", "enhanced"])
def test_multi_variants_preserve_earliest_complete_affected_layer(
monkeypatch, tmp_path, modules, variant
):
multi = modules["multi"]
_patch_output(monkeypatch, multi, tmp_path)
first = tmp_path / f"{variant}_first.safetensors"
second = tmp_path / f"{variant}_second.safetensors"
first_down = torch.tensor([[1.0, 2.0]])
save_file({
"diffusion_model.guarded.lora_down.weight": first_down,
"diffusion_model.guarded.lora_up.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(first))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 2]], dtype=torch.uint8),
"diffusion_model.guarded.lora_up.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([9.0, 9.0]),
"diffusion_model.normal.diff": torch.full((2, 2), 3.0),
}, str(second))
output = _run_multi_variant(multi, variant, [str(first), str(second)], f"guard_{variant}")
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], first_down)
torch.testing.assert_close(
result["diffusion_model.guarded.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
torch.testing.assert_close(
result["diffusion_model.guarded.dora_scale"], torch.tensor([0.125, 0.25])
)
assert "diffusion_model.normal.diff" in result
def test_hard_error_happens_before_resize_output_creation(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "three_low.safetensors"
save_file({f"low_{i}": torch.ones(1, dtype=torch.uint8) for i in range(3)}, str(source))
output = tmp_path / "must_not_exist.safetensors"
with pytest.raises(ValueError, match="3 low-bit tensors"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def test_recognized_quantization_creates_no_resize_output(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "recognized_quant.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
}, str(source))
output = tmp_path / "recognized_must_not_exist.safetensors"
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def _generic_params(output_filename):
return {
"mismatch_mode": "skip", "alignment_mode": "pad/crop", "alpha": 0.5,
"beta": 0.0, "gamma": 0.5, "delta": 2.0, "epsilon": 0.01, "zeta": 0.0,
"seed": 0, "output_filename": output_filename, "save_dtype": "bf16",
"override_dtype": True, "device": "cpu", "dtype": torch.float32,
"exclude_patterns": "", "discard_patterns": "", "glob_patterns": False,
"lazy_load": True, "force_clear_cache": False, "include_1d_diffs": True,
}
def test_generic_merge_preserves_model_a_when_either_source_key_is_low_bit(
monkeypatch, tmp_path, modules
):
generic = modules["generic"]
operations = modules["operations"]
model_a = tmp_path / "generic_guard_a.safetensors"
model_b = tmp_path / "generic_guard_b.safetensors"
a_tensors = {
"a_low": torch.tensor([[1, 2]], dtype=torch.uint8),
"b_low": torch.tensor([[0.125, 0.25]], dtype=torch.float32),
"normal": torch.ones((1, 2)),
}
save_file(a_tensors, str(model_a))
save_file({
"a_low": torch.tensor([[9.0, 9.0]]),
"b_low": torch.tensor([[3, 4]], dtype=torch.int8),
"normal": torch.full((1, 2), 3.0),
}, str(model_b))
paths = {"a": str(model_a), "b": str(model_b)}
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
params = _generic_params("generic_guarded")
params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "loras",
)
result = load_file(
str(tmp_path / "loras" / "generic_guarded.safetensors")
)
torch.testing.assert_close(result["a_low"], a_tensors["a_low"])
torch.testing.assert_close(result["b_low"], a_tensors["b_low"])
assert result["a_low"].dtype == torch.uint8
assert result["b_low"].dtype == torch.float32
torch.testing.assert_close(result["normal"].float(), torch.full((1, 2), 2.0))
def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "merge_guard_base.safetensors"
adapter = tmp_path / "merge_guard_adapter.safetensors"
base_quant = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"model.diffusion_model.base_quant.weight": base_quant,
"model.diffusion_model.guarded.weight": torch.zeros((2, 2)),
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 1]], dtype=torch.int8),
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 1)),
"diffusion_model.guarded.diff": torch.full((2, 2), 5.0),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(adapter))
output, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16, "merge_guarded",
verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(output)
torch.testing.assert_close(result["model.diffusion_model.base_quant.weight"], base_quant)
assert result["model.diffusion_model.base_quant.weight"].dtype == torch.uint8
torch.testing.assert_close(
result["model.diffusion_model.guarded.weight"].float(), torch.zeros((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.normal.weight"].float(), torch.ones((2, 2))
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 1/2")
assert "[GUARDED] (1)" in report
assert "diffusion_model.guarded.lora_down.weight" in report
assert "torch.int8" in report
assert "model.diffusion_model.base_quant.weight" in report