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
Silver c3902c235c Merge branch 'main' into feature/dora-extract 2026-06-06 18:57:51 +02:00
silveroxides 2b496bdd8c Bump version to 0.4.2 2026-06-06 18:50:37 +02:00
silveroxides 34223abb25 Add anima lora renaming node 2026-06-06 18:50:17 +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
Silver d2a81c0f9c Bump version from 0.4.2 to 0.4.3 2026-05-29 14:38:38 +02:00
Silver 2cb7f6c55e Merge pull request #14 from silveroxides/feature/dora-extract
Feature/dora extract
2026-05-29 14:37:56 +02:00
silveroxides 3a505e704f Add work in progress experimental learned extraction nodes 2026-05-29 14:36:26 +02:00
silveroxides 1512fa7390 Add DoRA extraction nodes 2026-05-28 03:27:47 +02:00
Silver b7365bfd79 Bump version from 0.4.1 to 0.4.2 2026-05-25 14:25:29 +02:00
Silver 34a2b78a9c Update requirements.txt 2026-05-25 14:25:00 +02:00
Silver 5780122057 Merge pull request #13 from silveroxides/refactor/incremental-writer
Refactor/incremental writer
2026-05-25 14:21:33 +02:00
19 changed files with 7476 additions and 529 deletions
+1
View File
@@ -0,0 +1 @@
tests/
+21
View File
@@ -20,10 +20,19 @@ 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
)
from .nodes.dora_extract_wd import (
DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
DoRAExtractKnee, DoRAExtractFrobenius
)
from .nodes.dora_learned_wd import (
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius
)
from .nodes.lora_resize import (
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
@@ -39,6 +48,8 @@ from .nodes.model_info_nodes import (
)
from .nodes.lora_rename import AnimaLoraRename
class ModelUtilsExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
@@ -58,9 +69,17 @@ class ModelUtilsExtension(ComfyExtension):
# Three-Model Mergers
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger,
# Dedicated CWB Mergers
*CWB_MERGER_NODES,
# LoRA Extraction
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius,
# DoRA Extraction
DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
DoRAExtractKnee, DoRAExtractFrobenius,
# Learned DoRA Extraction
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius,
# LoRA Resize
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
@@ -68,6 +87,8 @@ class ModelUtilsExtension(ComfyExtension):
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
# LoRA Merge To Model
LoRAMergeToModel,
# LoRA Utilities
AnimaLoraRename,
# Downloaders
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
+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
+989
View File
@@ -0,0 +1,989 @@
"""
DoRA Extraction via SVD and Weight Decomposition.
Native implementation of Weight-Decomposed Low-Rank Adapter extraction
with multiple rank selection modes. Extracted models use the DoRA structure.
"""
import fnmatch
import os
import re
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# =============================================================================
# Rank Index Functions - Determine LoRA rank from singular values
# =============================================================================
def _index_sv_fixed(S: torch.Tensor, dim: int) -> int:
"""Fixed rank mode - use specified dimension."""
return max(1, min(dim, len(S) - 1))
def _index_sv_threshold(S: torch.Tensor, threshold: float) -> int:
"""Threshold mode - keep singular values above absolute threshold."""
rank = int(torch.sum(S > threshold).item())
return max(1, min(rank, len(S) - 1))
def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int:
"""Ratio mode - keep singular values > max(S) / ratio (kohya convention)."""
if ratio <= 0:
return len(S) - 1
min_sv = S[0] / ratio
rank = int(torch.sum(S > min_sv).item())
return max(1, min(rank, len(S) - 1))
def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Cumulative mode - keep enough SVs to reach target % of total.
Calculates relative to max_rank if provided, otherwise relative to full.
"""
if max_rank is not None and max_rank < len(S):
total = torch.sum(S[:max_rank])
else:
total = torch.sum(S)
if total < 1e-8:
return 1
cumsum = torch.cumsum(S, dim=0) / total
rank = int(torch.searchsorted(cumsum, target).item()) + 1
return max(1, min(rank, len(S) - 1))
def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Frobenius norm mode - preserve target fraction of Frobenius norm.
Calculates relative to max_rank if provided, otherwise relative to full.
This means "retain target% of what's achievable within max_rank".
"""
if max_rank is not None and max_rank < len(S):
# Calculate relative to what's achievable within max_rank
S_capped = S[:max_rank]
S_sq = S_capped.pow(2)
total_sq = torch.sum(S_sq)
else:
S_sq = S.pow(2)
total_sq = torch.sum(S_sq)
if total_sq < 1e-8:
return 1
# Cumsum of all S (not capped) to find where we reach target
cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq
rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1
return max(1, min(rank, len(S) - 1))
def _index_sv_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
"""Knee detection - find elbow point on singular value curve."""
n = len(S)
if n < 3:
return 1
s_max, s_min = S[0].item(), S[-1].item()
if s_max - s_min < min_sv:
return 1
# Normalize to [0, 1]
s_norm = (S - s_min) / (s_max - s_min)
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
# Distance from diagonal
distances = (x_norm + s_norm - 1).abs()
rank = torch.argmax(distances).item() + 1
return max(1, min(rank, n - 1))
def _index_sv_cumulative_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
"""Knee detection on cumulative singular value curve."""
n = len(S)
if n < 3:
return 1
total = torch.sum(S)
if total < min_sv:
return 1
cumsum = torch.cumsum(S, dim=0) / total
y_min, y_max = cumsum[0].item(), cumsum[-1].item()
if y_max - y_min < min_sv:
return 1
y_norm = (cumsum - y_min) / (y_max - y_min)
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
distances = (y_norm - x_norm).abs()
rank = torch.argmax(distances).item() + 1
return max(1, min(rank, n - 1))
def _index_sv_rel_decrease(S: torch.Tensor, tau: float = 0.1) -> int:
"""Relative decrease mode - stop when consecutive SV ratio drops below tau."""
if len(S) < 2:
return 1
ratios = S[1:] / (S[:-1] + 1e-8)
for k in range(len(ratios)):
if ratios[k] < tau:
return k + 1
return len(S)
def _compute_rank(S: torch.Tensor, mode: str, mode_param: float,
max_rank: int = None) -> int:
"""Compute rank based on mode and parameters."""
if mode == "fixed":
rank = _index_sv_fixed(S, int(mode_param))
elif mode == "threshold":
rank = _index_sv_threshold(S, mode_param)
elif mode == "ratio":
rank = _index_sv_ratio(S, mode_param)
elif mode == "quantile" or mode == "sv_cumulative":
rank = _index_sv_cumulative(S, mode_param, max_rank)
elif mode == "sv_fro":
rank = _index_sv_fro(S, mode_param, max_rank)
elif mode == "sv_knee":
rank = _index_sv_knee(S)
elif mode == "sv_cumulative_knee":
rank = _index_sv_cumulative_knee(S)
elif mode == "sv_rel_decrease":
rank = _index_sv_rel_decrease(S, mode_param)
else:
rank = _index_sv_fixed(S, int(mode_param) if mode_param else 64)
if max_rank is not None:
rank = min(rank, max_rank)
return rank
# =============================================================================
# SVD Extraction Functions
# =============================================================================
def _svd_extract_linear_lowrank(
weight: torch.Tensor,
rank: int,
device: str,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms.
Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# Clamp rank to valid range
max_possible_rank = min(out_dim, in_dim)
rank = max(1, min(rank, max_possible_rank))
# Check if decomposition is worthwhile
if rank >= max_possible_rank:
return weight, "full"
# Use svd_lowrank for faster low-rank approximation
try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
# U: [out_dim, rank], S: [rank], V: [in_dim, rank]
Vh = V.T # [rank, in_dim]
except Exception as e:
raise RuntimeError(f"svd_lowrank failed: {e}")
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_linear(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
SVD decomposition for linear layers.
For fixed rank mode, uses svd_lowrank which is much faster.
For adaptive modes, uses full SVD to analyze all singular values.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# For fixed rank, use the faster svd_lowrank
if mode == "fixed" and mode_param > 0:
target_rank = int(mode_param)
if max_rank is not None:
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
# Compute rank based on mode
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim):
return weight, "full"
# Truncate to rank
U = U[:, :rank]
S = S[:rank]
Vh = Vh[:rank, :]
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass # Skip clamping if quantile fails
# Construct LoRA matrices
# lora_up = U @ diag(S), lora_down = Vh
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff for optional sparse bias
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_conv(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
) -> tuple:
"""
SVD decomposition for conv2d layers.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
out_ch, in_ch, kh, kw = weight.shape
is_1x1 = (kh == 1 and kw == 1)
# Flatten for SVD
if is_1x1:
mat = weight.view(out_ch, in_ch) # [out_ch, in_ch]
else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile)
if mode_str == "full":
return weight, "full"
lora_down, lora_up, diff = result
rank = lora_down.shape[0]
# Reshape for conv
if is_1x1:
lora_down = lora_down.unsqueeze(-1).unsqueeze(-1) # [rank, in_ch, 1, 1]
lora_up = lora_up.unsqueeze(-1).unsqueeze(-1) # [out_ch, rank, 1, 1]
else:
lora_down = lora_down.reshape(rank, in_ch, kh, kw)
lora_up = lora_up.reshape(out_ch, rank, 1, 1)
return (lora_down, lora_up, diff), "low rank"
# =============================================================================
# Chunked Extraction for Large Layers
# =============================================================================
def _detect_fused_layer(key: str, shape: tuple) -> int:
"""Detect fused QKV/MLP layers and return number of chunks."""
if len(shape) != 2:
return 1
out_dim, in_dim = shape
key_lower = key.lower()
# QKV layers (3 chunks)
if "qkv" in key_lower and out_dim % 3 == 0 and out_dim // 3 >= in_dim // 4:
return 3
# KV layers (2 chunks)
if "kv" in key_lower and "qkv" not in key_lower and out_dim % 2 == 0:
return 2
# Fused MLP (2 chunks)
if ("linear1" in key_lower or "gate_up" in key_lower) and out_dim % 2 == 0 and out_dim >= in_dim * 2:
return 2
return 1
def _extract_chunked_layer(
weight_diff: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
) -> tuple:
"""Extract LoRA from fused layer by chunking."""
out_dim, in_dim = weight_diff.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :]
try:
result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
# Combine chunks
# Note: Using the first chunk's down matrix as the basis for all chunks
# is a better approximation than a simple mean if the chunks share the same input space (like QKV).
# Ideally, we would use a more sophisticated joint SVD, but this is a reasonable fallback.
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
# =============================================================================
# Pattern Matching
# =============================================================================
def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
"""Compile patterns from whitespace-separated string.
In regex mode (default) returns compiled re.Pattern objects.
In glob mode returns plain strings; fnmatch handles matching.
"""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for p in pattern_string.split():
p = p.strip()
if not p:
continue
if glob_mode:
patterns.append(p)
else:
try:
patterns.append(re.compile(p))
except re.error:
pass
return patterns
def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
"""Check if key matches any pattern.
Glob mode: each pattern is matched as a substring glob using fnmatch
(dots are literal, * matches any sequence of characters).
Regex mode: each pattern is a compiled re.Pattern, matched via search().
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
return any(p.search(key) for p in patterns)
# =============================================================================
# Main Extraction Function
# =============================================================================
def extract_dora_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
Args:
model_a_path: Finetuned model path
model_b_path: Base model path
mode: Rank selection mode
linear_param: Mode parameter for linear layers
conv_param: Mode parameter for conv layers
device: Computation device
save_dtype: Output dtype
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers
clamp_quantile: Quantile for weight clamping
min_diff: Minimum difference threshold
skip_patterns_str: Patterns for layers to skip
mismatch_mode: How to handle mismatches
chunk_large_layers: Enable chunked extraction
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) # 1.5x for SVD headroom
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
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))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key)
if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
# Without base, DoRA fallback just treats tensor_a as diff
dora_scale = None
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
weight_diff = tensor_a
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
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
# Compute directional diff: W_dir = W_ft * (norm_base / norm_ft)
# Use a small epsilon to avoid division by zero
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del tensor_a, tensor_b, w_dir
# Skip small differences
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[DoRA Extract] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_layer(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
print(f"[DoRA Extract] Failed: {key}: {e}")
del weight_diff
return "skipped", None
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, rank_or_diff = result
# rank is the size of lora_down's first dimension
rank = lora_down.shape[0]
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting DoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[DoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> str:
"""Build output path for LoRA file."""
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
# =============================================================================
# Node Definitions
# =============================================================================
def _get_model_inputs():
return [
io.Combo.Input("model_a", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Finetuned model (A - B = LoRA)"),
io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Base model (A - B = LoRA)"),
]
def _format_lora_key(key: str) -> str:
"""
Format the key for LoRA saving.
Always standardizes to 'diffusion_model.' prefix for maximum compatibility.
ComfyUI maps 'diffusion_model.<layer>' generically for all model types, making
it universally compatible. The 'transformer.' prefix only works for a subset of
model types via model-specific mapping code and is intentionally not used here.
"""
if key.endswith(".weight"):
key = key[:-7]
# Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*)
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # strip len("model.diffusion_model.")
return f"diffusion_model.{inner_key}"
# Handle net.* wrapper
if key.startswith("net."):
inner_key = key[4:] # strip len("net.")
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys (transformer_blocks.* / single_transformer_blocks.* without prefix)
if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"):
return f"diffusion_model.{key}"
# Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"):
return key
# Handle already-prefixed keys
if key.startswith("diffusion_model."):
return key
# Absorb now-invalid transformer. prefix into diffusion_model. for compatibility
if key.startswith("transformer."):
return f"diffusion_model.{key[12:]}"
# Handle known Diffusers UNet blocks
if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]):
return f"diffusion_model.{key}"
# Default fallback
return f"diffusion_model.{key}"
def _get_common_inputs():
return [
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("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
tooltip="Clamp outlier singular values"),
io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
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("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."),
]
class DoRAExtractFixed(io.ComfyNode):
"""Extract DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractFixed",
display_name="DoRA Extract (Fixed Rank)",
category="ModelUtils/DoRA",
description="Extract DoRA with specified fixed rank for each layer type.",
inputs=[
*_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384,
tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384,
tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
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
)
return io.NodeOutput(output_path)
class DoRAExtractRatio(io.ComfyNode):
"""Extract DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractRatio",
display_name="DoRA Extract (Ratio)",
category="ModelUtils/DoRA",
description="Keep singular values > max(S) / ratio.",
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
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
)
return io.NodeOutput(output_path)
class DoRAExtractQuantile(io.ComfyNode):
"""Extract DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractQuantile",
display_name="DoRA Extract (Quantile)",
category="ModelUtils/DoRA",
description="Keep enough singular values to reach target percentage.",
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target cumulative % for linear layers"),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target cumulative % for conv layers"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
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
)
return io.NodeOutput(output_path)
class DoRAExtractKnee(io.ComfyNode):
"""Extract DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractKnee",
display_name="DoRA Extract (Knee Detection)",
category="ModelUtils/DoRA",
description="Automatically find optimal rank using knee detection on singular value curve.",
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
default="sv_knee", tooltip="Knee detection method"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
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
)
return io.NodeOutput(output_path)
class DoRAExtractFrobenius(io.ComfyNode):
"""Extract DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractFrobenius",
display_name="DoRA Extract (Frobenius)",
category="ModelUtils/DoRA",
description="Preserve target fraction of Frobenius norm.",
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target Frobenius norm fraction for linear"),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target Frobenius norm fraction for conv"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
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
)
return io.NodeOutput(output_path)
+790
View File
@@ -0,0 +1,790 @@
"""
Learned DoRA Extraction via Optimization.
Extracts Weight-Decomposed Low-Rank Adapters using gradient descent optimization
to minimize the reconstruction MSE loss, matching the true non-linear DoRA formulation.
Uses analytical SVD purely as the initialization seed.
"""
import fnmatch
import os
import re
import math
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# Import the existing SV ranking and chunking logic so we don't have to duplicate it
from .dora_extract_wd import (
_compute_rank, _svd_extract_linear_lowrank, _svd_extract_linear,
_svd_extract_conv, _detect_fused_layer, _compile_patterns,
_matches_any_pattern, _build_lora_output_path, _get_model_inputs,
_format_lora_key, _get_common_inputs
)
def _adaptive_lr_update_cosine(curr_lr: float, improved: bool, worse_loss_counter: int, iteration: int, tensor_shape: tuple, min_lr: float, early_stop_stall: int, lr_factor: float = 0.95, lr_cooldown: int = 1) -> tuple:
M, N = tensor_shape
shape_ratio = abs(M - N) / max(M, N) if max(M, N) > 0 else 0
t = min(worse_loss_counter / max(early_stop_stall, 1), 1.0)
u_factor = (1 + math.cos(2 * math.pi * t)) / 2
if improved:
base_boost = 1.25
distance = base_boost - 1.0
scaled_distance = distance * (1.0 + 0.5 * shape_ratio)
boost_mult = 1.0 + scaled_distance * (1.0 - u_factor)
return min(curr_lr * boost_mult, 100.0), True
else:
cooldown = max(lr_cooldown, 1)
if iteration % cooldown != 0:
return curr_lr, False
min_decay = lr_factor - 0.03 * shape_ratio
max_decay = 0.995
decay_mult = min_decay + (max_decay - min_decay) * u_factor
return max(curr_lr * decay_mult, min_lr), True
def _compute_shape_aware_plateau_params(M: int, N: int, lr_patience: int, lr_factor: float, lr_cooldown: int, lr_shape_influence: float = 1.0) -> tuple:
if lr_shape_influence > 0:
aspect_ratio = max(M, N) / min(M, N) if min(M, N) > 0 else 1.0
ar_factor = math.sqrt(aspect_ratio)
blend = lr_shape_influence
effective_patience = lr_patience
raw_factor = lr_factor
aggressive_factor = raw_factor**ar_factor
effective_factor = raw_factor + (aggressive_factor - raw_factor) * blend
effective_cooldown = lr_cooldown
else:
effective_patience = lr_patience
effective_factor = lr_factor
effective_cooldown = lr_cooldown
return effective_patience, effective_factor, effective_cooldown
def _optimize_dora(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
init_down: torch.Tensor,
init_up: torch.Tensor,
optimize_iters: int,
learning_rate: float,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""
Run gradient descent to optimize lora_down and lora_up against the actual DoRA formula.
Args:
tensor_base: Base model weight
tensor_ft: Finetuned model weight
init_down: SVD-initialized lora_down matrix
init_up: SVD-initialized lora_up matrix
optimize_iters: Number of iterations for AdamW
learning_rate: Learning rate for AdamW
optimizer_type: Optimizer choice (adamw, radam, prodigy)
lr_schedule: LR Schedule choice (exponential, plateau, adaptive)
early_stop_loss: Target loss to stop optimizing
early_stop_stall: Number of non-improving iterations before stopping
Returns:
(optimized_down, optimized_up)
"""
if optimize_iters <= 0:
return init_down, init_up
is_conv = tensor_base.ndim == 4
if is_conv:
out_ch, in_ch, kh, kw = tensor_base.shape
flat_base = tensor_base.reshape(out_ch, -1)
flat_ft = tensor_ft.reshape(out_ch, -1)
# Flatten init matrices for matmul: (out, rank) @ (rank, in*k*k)
rank = init_down.shape[0]
flat_up = init_up.reshape(out_ch, rank)
flat_down = init_down.reshape(rank, -1)
else:
flat_base = tensor_base
flat_ft = tensor_ft
flat_up = init_up
flat_down = init_down
# Magnitude of the finetuned model (fixed DoRA scale)
m = torch.linalg.norm(flat_ft, dim=1, keepdim=True).detach()
eps = 1e-8
# Set up learnable parameters
with torch.inference_mode(False), torch.enable_grad():
# Clone inference tensors to normal tensors for autograd
flat_base_opt = flat_base.clone()
flat_ft_opt = flat_ft.clone()
m_opt = m.clone()
# Free original tensors to save memory
del flat_base, flat_ft, m
B = torch.nn.Parameter(flat_up.clone())
A = torch.nn.Parameter(flat_down.clone())
opt_type = optimizer_type.lower()
if opt_type == "radam":
optimizer = torch.optim.RAdam([B, A], lr=learning_rate)
elif opt_type == "prodigy":
try:
from prodigyplus.prodigy_plus_schedulefree import ProdigyPlusScheduleFree
optimizer = ProdigyPlusScheduleFree([B, A], lr=learning_rate, use_schedulefree=False)
except ImportError:
print("ProdigyPlusScheduleFree optimizer not found. Falling back to AdamW. To install: pip install prodigyplus")
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
else:
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
# Learning rate scheduling and early stopping state
best_loss = float('inf')
worse_loss_counter = 0
plateau_counter = 0
cooldown_counter = 0
current_lr = learning_rate
schedule_type = lr_schedule.lower()
M_dim, N_dim = flat_base_opt.shape[0], flat_base_opt.shape[1]
eff_patience, eff_factor, eff_cooldown = _compute_shape_aware_plateau_params(M_dim, N_dim, lr_patience, lr_factor, lr_cooldown)
# Optimization loop
opt_pbar = tqdm(range(optimize_iters), desc=f" └─ Optimizing {layer_name}", leave=False, dynamic_ncols=True, mininterval=0.5)
for i in opt_pbar:
optimizer.zero_grad()
# 1. Reconstruct directional weight: W_dir = W_base + B @ A
w_approx = flat_base_opt + (B @ A)
# 2. Normalize directional weight
approx_norm = torch.linalg.norm(w_approx, dim=1, keepdim=True) + eps
# 3. Scale by target magnitude
w_reconstructed = m_opt * (w_approx / approx_norm)
# 4. Compute MSE loss against actual finetuned weight
loss = torch.nn.functional.mse_loss(w_reconstructed, flat_ft_opt)
loss_val = loss.item()
if i == 0:
initial_loss = loss_val
best_loss = loss_val
loss.backward()
optimizer.step()
# --- Early Stopping & Schedulers ---
improved = False
if loss_val < best_loss * 0.9999: # 0.01% improvement threshold
best_loss = loss_val
worse_loss_counter = 0
plateau_counter = 0
improved = True
else:
worse_loss_counter += 1
plateau_counter += 1
if loss_val <= early_stop_loss:
break
if worse_loss_counter > early_stop_stall:
break
if current_lr <= early_stop_lr:
break
# LR Schedulers
if schedule_type == "adaptive":
current_lr, updated = _adaptive_lr_update_cosine(
current_lr, improved, worse_loss_counter, i, (M_dim, N_dim), early_stop_lr, early_stop_stall, lr_factor, lr_cooldown
)
if updated:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
elif schedule_type == "plateau":
if cooldown_counter > 0:
cooldown_counter -= 1
elif plateau_counter >= max(eff_patience, 1): # Decay if patience threshold reached
current_lr = max(current_lr * eff_factor, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
plateau_counter = 0 # reset counter after decay
cooldown_counter = eff_cooldown
elif schedule_type == "exponential":
current_lr = max(current_lr * 0.99, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
if i % 10 == 0 or i == optimize_iters - 1:
if schedule_type == "plateau":
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "plat": f"{plateau_counter}/{eff_patience}"}, refresh=False)
else:
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "stall": f"{worse_loss_counter}"}, refresh=False)
final_loss = loss.item()
opt_pbar.close()
print(f" └─ {layer_name} | Loss: {initial_loss:.6f} -> {final_loss:.6f} (stopped at iter {i+1})")
final_up = B.detach()
final_down = A.detach()
# Reshape back to original dimensions
if is_conv:
final_up = final_up.reshape(out_ch, rank, 1, 1)
final_down = final_down.reshape(rank, in_ch, kh, kw)
return final_down, final_up
def _extract_chunked_learned_layer(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""Extract Learned DoRA from fused layer by chunking."""
out_dim, in_dim = tensor_base.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk_base = tensor_base[i * chunk_size:(i + 1) * chunk_size, :]
chunk_ft = tensor_ft[i * chunk_size:(i + 1) * chunk_size, :]
norm_base = torch.linalg.norm(chunk_base, dim=1, keepdim=True)
norm_ft = torch.linalg.norm(chunk_ft, dim=1, keepdim=True)
eps = 1e-8
w_dir = chunk_ft * (norm_base / (norm_ft + eps))
weight_diff = w_dir - chunk_base
try:
result, mode_str = _svd_extract_linear(weight_diff, mode, mode_param, device, max_rank)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
# Apply learned optimization
lora_down, lora_up = _optimize_dora(
chunk_base, chunk_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=f"{layer_name} (Chunk {i+1}/{num_chunks})"
)
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
def extract_dora_learned_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract Learned DoRA from difference between two models.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[Learned DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
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))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
# Cannot extract true DoRA without base. Fallback to diff.
cpu_a = handler_a.get_tensor(key)
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) if use_pinned else cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
weight_diff = tensor_a
tensor_base = torch.zeros_like(tensor_a)
tensor_ft = tensor_a
dora_scale = None
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
weight_diff = tensor_a
tensor_base = torch.zeros_like(tensor_a)
tensor_ft = tensor_a
dora_scale = None
del tensor_b
else:
tensor_base = tensor_b
tensor_ft = tensor_a
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del w_dir
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if weight_diff.ndim < 2:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
# 1. Run Analytical SVD (Initialization)
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
# 2. Run Optimization (Learned Rounding)
if mode_str != "full" and dora_scale is not None:
lora_down, lora_up, _ = result
lora_down, lora_up = _optimize_dora(
tensor_base, tensor_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
)
result = (lora_down, lora_up, _)
except Exception as e:
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[Learned DoRA] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_learned_layer(
tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff, tensor_base, tensor_ft
return "chunked", layer_results
print(f"[Learned DoRA] Failed: {key}: {e}")
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
rank = lora_down.shape[0]
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff, tensor_base, tensor_ft
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting Learned DoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[Learned DoRA] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
# =============================================================================
# Node Definitions
# =============================================================================
def _get_learned_inputs():
return [
io.Int.Input("optimize_iters", default=500, min=0, max=10000,
tooltip="Number of gradient descent iterations (0 = skip optimization)"),
io.Float.Input("learning_rate", default=1.0, min=0.0001, max=100.0, step=0.001,
tooltip="Base learning rate. AdamW prefers ~0.01, Prodigy prefers ~1.0"),
io.Combo.Input("optimizer", options=["prodigy", "adamw", "radam"], default="prodigy",
tooltip="Optimization algorithm"),
io.Combo.Input("lr_schedule", options=["plateau", "adaptive", "exponential", "constant"], default="plateau",
tooltip="Learning rate scheduling strategy"),
io.Int.Input("lr_patience", default=2, min=0, max=1000, tooltip="Steps to wait before decaying LR (Plateau)"),
io.Float.Input("lr_factor", default=0.9, min=0.01, max=1.0, step=0.01, tooltip="Factor to decay LR by (Plateau)"),
io.Int.Input("lr_cooldown", default=2, min=0, max=1000, tooltip="Steps to wait after decay before checking again (Plateau)"),
io.Float.Input("early_stop_loss", default=1e-6, min=0.0, max=1.0, step=1e-9,
tooltip="Stop early if MSE loss drops below this value"),
io.Int.Input("early_stop_stall", default=2000, min=0, max=10000,
tooltip="Stop early if loss doesn't improve for this many iterations"),
io.Float.Input("early_stop_lr", default=9.01e-9, min=0.0, max=1.0, step=1e-10,
tooltip="Stop early if learning rate drops below this value"),
]
class DoRALearnedExtractFixed(io.ComfyNode):
"""Extract Learned DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFixed",
display_name="DoRA Learned Extract (Fixed Rank) (WIP)",
category="ModelUtils/DoRA",
description="Extract DoRA via MSE Loss minimization with fixed rank.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
*_get_learned_inputs(),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD iterations for initialization"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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
)
return io.NodeOutput(output_path)
class DoRALearnedExtractRatio(io.ComfyNode):
"""Extract Learned DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractRatio",
display_name="DoRA Learned Extract (Ratio) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping SVs > max(S) / ratio.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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
)
return io.NodeOutput(output_path)
class DoRALearnedExtractQuantile(io.ComfyNode):
"""Extract Learned DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractQuantile",
display_name="DoRA Learned Extract (Quantile) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping enough SVs to reach target percentage.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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
)
return io.NodeOutput(output_path)
class DoRALearnedExtractKnee(io.ComfyNode):
"""Extract Learned DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractKnee",
display_name="DoRA Learned Extract (Knee Detection) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction automatically finding optimal rank using knee detection.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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
)
return io.NodeOutput(output_path)
class DoRALearnedExtractFrobenius(io.ComfyNode):
"""Extract Learned DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFrobenius",
display_name="DoRA Learned Extract (Frobenius) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction preserving target fraction of Frobenius norm.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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
)
return io.NodeOutput(output_path)
+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
+119
View File
@@ -0,0 +1,119 @@
import os
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
REVERSE_RENAME_DICT = {
"transformer_blocks.": "blocks.",
"norm1.linear_1": "adaln_modulation_self_attn.1",
"norm1.linear_2": "adaln_modulation_self_attn.2",
"norm2.linear_1": "adaln_modulation_cross_attn.1",
"norm2.linear_2": "adaln_modulation_cross_attn.2",
"norm3.linear_1": "adaln_modulation_mlp.1",
"norm3.linear_2": "adaln_modulation_mlp.2",
"attn1.to_q": "self_attn.q_proj",
"attn1.to_k": "self_attn.k_proj",
"attn1.to_v": "self_attn.v_proj",
"attn1.to_out.0": "self_attn.output_proj",
"attn2.to_q": "cross_attn.q_proj",
"attn2.to_k": "cross_attn.k_proj",
"attn2.to_v": "cross_attn.v_proj",
"attn2.to_out.0": "cross_attn.output_proj",
"ff.net.0.proj": "mlp.layer1",
"ff.net.2": "mlp.layer2",
"norm_out.linear_1": "final_layer.adaln_modulation.1",
"norm_out.linear_2": "final_layer.adaln_modulation.2",
"proj_out": "final_layer.linear",
"time_embed.t_embedder": "t_embedder.1",
"time_embed.norm": "t_embedding_norm",
"patch_embed.proj": "x_embedder.proj.1",
}
def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filename: str) -> str:
"""
Converts a Diffusers Anima LoRA back to its non-diffusers format.
"""
model_type = "loras"
model_path = folder_paths.get_full_path_or_raise(model_type, lora_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
# Get metadata from safe_open (header only, no tensor loading)
with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
metadata = f.metadata() or {}
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Stream tensors, rename on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
for key in tqdm(original_keys, desc="Renaming Anima LoRA keys", unit="keys"):
new_key = key
if key.startswith("text_conditioner."):
new_key = f"diffusion_model.llm_adapter.{key.removeprefix('text_conditioner.')}"
elif key.startswith("transformer."):
new_key = key.removeprefix("transformer.")
for old_part, new_part in REVERSE_RENAME_DICT.items():
new_key = new_key.replace(old_part, new_part)
new_key = f"diffusion_model.{new_key}"
writer.write(new_key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
cleanup_after_operation()
return output_path
class AnimaLoraRename(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="AnimaLoraRename",
display_name="Anima LoRA Rename (Diffusers to Non-Diffusers)",
category="ModelUtils/LoRA",
description="Loads a Diffusers Anima LoRA, reverses its key names to the non-diffusers format, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("output_filename", default="anima_lora_non_diffusers"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name: str, output_filename: str) -> io.NodeOutput:
path = _convert_diffusers_to_non_diffusers_anima_lora(lora_name, output_filename)
return io.NodeOutput(path)
+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.1"
version = "0.4.6"
license = { file = "LICENSE" }
classifiers = [
+1 -1
View File
@@ -1,5 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader>=0.5.1
unifiedefficientloader>=0.5.2
av
+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