Bump version to 0.9.7

Feature/model similarity analysis

.comfyignore

@@ -0,0 +1,4 @@
+tests/
+nodes/downloader_nodes.py
+nodes/downloader_utils.py
+

.gitignore

@@ -176,3 +176,14 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# local untracked directories
+tools_local/
+references_local/
+plans_local/
+old_files_local/
+audit_local/
+old_sessions/
+
+# local untracked files
+/*.jsonl

README.md

@@ -9,9 +9,36 @@ A collection of ComfyUI custom nodes for inspecting, modifying, merging, and cre
 - **PruneKeys** – Remove unwanted layers/keys from models
 - **Mergers** – Combine 2 or 3 models with configurable blend modes and ratios
 - **LoRA Extraction** – Extract LoRA adapters from model pairs using various SVD rank selection methods (Fixed, Ratio, Quantile, Knee-detection, Frobenius-norm)
+- **Diffusion Model Dtype Conversion** – Stream models to fp32, fp16, or bf16 while preserving excluded tensor dtypes
+
+## Layer filters
+
+Nodes with `exclude_patterns` or `skip_patterns` have an appended `include_mode` toggle
+(off by default). Turn it on to process only layers matching that same field,
+using the node's existing regex or glob syntax. An empty include filter selects
+nothing. Nonmatches follow the node's usual exclusion behavior: preserve the
+source/anchor layer for pattern-exclusion mergers, resize, or conversion; omit
+it from analysis or extraction. LoRA Merge to Model retains its existing skip
+semantics: nonmatching base tensors are omitted from the saved model, while
+guarded low-bit base tensors are always preserved. Its filter uses regex only.
+`discard_patterns`, where available, still takes precedence.
 
 ## Example Workflows
 
+### Per-layer parameters
+
+Connect **Layer Parameter Configuration** to `layer_parameters` on standard
+two/three-input mergers, extraction nodes, or LoRA resize nodes. For example,
+fixed extraction accepts `(blocks\.4[589]\.attn\.qkv_proj) a:64 b:32 c:0.99 d:0`.
+Short aliases and full names, with whitespace or comma separators, are supported.
+The configurator's documentation output lists every method's bindings and bounds.
+
+Only explicitly assigned values override the receiving node; other settings and
+existing filters still apply. Overlapping rules and rules matching no target
+layer fail before streaming or writing. LoRA patterns target the existing
+normalized logical layer names, not individual factor/alpha suffixes.
+
+
 <p align="center">
   <img src="assets/GetMetaAndkeys.png" width="400" alt="Get Meta and Keys">
   <br>

__init__.py

@@ -1,7 +1,7 @@
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, io
 
-from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE
+from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced
 from .nodes.metakeys import (
     ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
     CheckpointMetaKeys, EmbeddingMetaKeys
@@ -20,21 +20,68 @@ from .nodes.merger import (
     ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
     CheckpointThreeMerger, EmbeddingThreeMerger
 )
+from .nodes.consensus_merger import CWB_MERGER_NODES
+from .nodes.cwb_model_graft import CWB_GRAFT_NODES
+from .nodes.cwb_delta_lora_merger import DELTA_CWB_LORA_NODES
+from .nodes.lodestone_merger import LODESTONE_MERGER_NODES
+from .nodes.model_analysis import MODEL_ANALYSIS_NODES
 from .nodes.lora_extract_svd import (
     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.text_encoder_extract import (
+    TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio,
+    TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee,
+    TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed,
+    TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile,
+    TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius,
+)
 from .nodes.lora_resize import (
+    LoRANormalizeAlpha,
+    LoRAModifyAlpha,
+    LoRADenormalizeAlpha,
     LoRAResizeFixed, LoRAResizeRatio,
     LoRAResizeFrobenius, LoRAResizeCumulative,
     LoRAMergeToModel
 )
+try:
+    from .nodes.downloader_nodes import (
+        CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
+        VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader
+    )
+    DOWNLOADER_NODES: list[type[io.ComfyNode]] = [
+        CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
+        VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
+    ]
+except ImportError:
+    DOWNLOADER_NODES = []
 
+from .nodes.model_info_nodes import (
+    CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
+    VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader
+)
+
+
+from .nodes.lora_rename import AnimaLoraRename
+from .nodes.minimax_h3_lora_convert import MiniMaxH3DiffusersLoRAConvert
+from .nodes.minimax_h3_fold import MiniMaxH3FoldAdaLN, MiniMaxH3UnfoldAdaLN
+from .nodes.dtype_conversion import DiffusionModelDtypeConversion
+from .nodes.layer_parameters import LayerParameterConfiguration
+from .nodes.lora_model_analysis import LoRAOnModelAnalysis
 
 class ModelUtilsExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
         return [
+            LayerParameterConfiguration,
             # MetaKeys
             ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
             CheckpointMetaKeys, EmbeddingMetaKeys,
@@ -50,16 +97,53 @@ class ModelUtilsExtension(ComfyExtension):
             # Three-Model Mergers
             ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
             CheckpointThreeMerger, EmbeddingThreeMerger,
+            # Dedicated CWB Mergers
+            *CWB_MERGER_NODES,
+            *CWB_GRAFT_NODES,
+            # Dedicated Delta CWB LoRA Mergers
+            *DELTA_CWB_LORA_NODES,
+            # Dedicated Lodestone LoRA Mergers
+            *LODESTONE_MERGER_NODES,
+            # Two-Model Analysis
+            *MODEL_ANALYSIS_NODES,
+            LoRAOnModelAnalysis,
             # LoRA Extraction
             LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
             LoRAExtractKnee, LoRAExtractFrobenius,
+            # DoRA Extraction
+            DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
+            DoRAExtractKnee, DoRAExtractFrobenius,
+            # Learned DoRA Extraction
+            DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
+            DoRALearnedExtractKnee, DoRALearnedExtractFrobenius,
+            # Text Encoder LoRA/DoRA Extraction
+            TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio,
+            TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee,
+            TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed,
+            TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile,
+            TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius,
             # LoRA Resize
+            LoRANormalizeAlpha,
+            LoRAModifyAlpha,
+            LoRADenormalizeAlpha,
             LoRAResizeFixed, LoRAResizeRatio,
             LoRAResizeFrobenius, LoRAResizeCumulative,
             # LoRA Multi-Merge
-            LoRAMultiMerge, LoRAMultiMergeDARE,
+            LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
             # LoRA Merge To Model
             LoRAMergeToModel,
+            # LoRA Utilities
+            AnimaLoraRename,
+            MiniMaxH3DiffusersLoRAConvert,
+            # Dtype Conversion & Folding
+            DiffusionModelDtypeConversion,
+            MiniMaxH3FoldAdaLN,
+            MiniMaxH3UnfoldAdaLN,
+            # Downloaders (available in manual git/source installations)
+            *DOWNLOADER_NODES,
+            # Info Loaders
+            CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
+            VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader,
         ]
 
 

docs/DEVELOPMENT.md

@@ -1,5 +1,7 @@
 # ComfyUI-ModelUtils - Complete Project History
 
+> [Implementation Invariants](IMPLEMENTATION_INVARIANTS.md) is the authoritative contract for UEL tensor I/O, lifecycle management, atomic saving, and LoRA alpha normalization. All new node implementations and changes must follow it. This document is historical and does not override that contract.
+
 ## Version Timeline
 
 | Version | Commits | Key Features |

docs/IMPLEMENTATION_INVARIANTS.md

@@ -0,0 +1,93 @@
+# Implementation Invariants
+
+This file is the authoritative engineering contract for model and LoRA tensor I/O and merge behavior in ComfyUI-ModelUtils. New implementations and changes to existing implementations MUST follow these rules. Historical code, audit notes, and examples do not override this contract.
+
+## UnifiedEfficientLoader I/O
+
+1. Production node code MUST use `unifiedefficientloader` for safetensors tensor I/O. Direct safetensors imports and synchronous `get_tensor()` calls are prohibited.
+2. Header-only operations MUST use UEL header, key, shape, dtype, and metadata APIs without materializing tensors.
+3. Tensor processing MUST use `async_stream` with bounded buffering. The repository default is `batch_size=1` and `prefetch_batches=1` unless a measured benchmark and a documented memory bound justify another value.
+4. Multi-input operations MUST maintain an async cursor for each input and construct logical work units. A work unit contains only the tensors required to process one output tensor, LoRA factor pair, or other indivisible operation.
+5. Pinned-memory transfer MUST be enabled when the selected processing device benefits from it. CPU-only processing MUST avoid unnecessary pinned copies.
+6. Implementations MUST NOT accumulate a complete model, LoRA, or unbounded output batch in RAM. The next work unit may be prefetched while the current unit is processed, but the prefetch bound MUST remain explicit.
+7. Every yielded source tensor occurrence MUST be marked processed after its logical work unit completes or aborts. All strong references to source, normalized, padded, device, and output intermediates MUST be removed before advancing to the next work unit.
+8. Every async stream and handler MUST close in `finally` or through an equivalent deterministic context manager. Failure paths are subject to the same release rules as successful paths.
+9. Saving MUST use UEL incremental writing. Outputs MUST be submitted as they are completed rather than retained in an in-memory state dictionary.
+10. File replacement MUST be atomic: write to a temporary sibling file, finalize successfully, then replace the destination. A failure MUST preserve any existing destination and remove the incomplete temporary file.
+11. Unsupported low-bit tensors MUST use the shared quantization-preservation path when byte-exact copying is valid. Code MUST NOT silently cast, decode, merge, or relabel unsupported low-bit storage.
+
+The shared implementations in `nodes/uel_io.py`, `nodes/extraction_stream.py`, and `nodes/quantization_guard.py` are the required starting points. New node-local stream or writer abstractions require a demonstrated capability that the shared implementation cannot provide.
+
+## Persistent Artifact Paths and Outputs
+
+1. Persistent model artifacts MUST use `nodes/artifact_paths.py` to construct
+   their write path from `folder_paths.models_dir` and an explicit canonical
+   category: `checkpoints`, `diffusion_models`, `text_encoders`, `loras`, or
+   `embeddings`.
+2. Registered search paths, additional paths, and legacy aliases such as
+   `unet` and `clip` MUST NOT become write roots based on their list position.
+3. Nested output names MUST be preserved, exactly one `.safetensors` extension
+   MUST be applied, and path traversal outside the category root MUST fail.
+4. Public filename outputs MUST return the category-relative name with forward
+   slashes. They MUST NOT return an absolute filesystem path.
+5. Loader-compatible filename outputs MUST use `io.AnyType`. Reports and
+   documentation MUST remain separate string outputs.
+6. Persistent model artifacts MUST NOT be written beneath ComfyUI's general
+   output directory.
+
+## LoRA Alpha Normalization for Merge Operations
+
+1. Every operation that combines a LoRA with another LoRA or applies a LoRA to a model MUST account for per-layer alpha before weighting, padding, concatenation, consensus processing, DARE/TIES processing, or delta application.
+2. For a valid factor pair with rank `r`, normalization is performed mathematically as:
+
+   `normalized_up = up * (alpha / r)`
+
+   `normalized_down = down`
+
+   Scaling one factor preserves the represented delta while avoiding unnecessary work and memory.
+3. A layer without alpha has an implicit scale of `1.0`. Inputs with and without alpha may be merged only after alpha-bearing pairs have been normalized to this common representation.
+4. Normalized merge outputs MUST NOT contain alpha tensors. Alpha MUST NOT be invented for inputs or output formats that do not provide it.
+5. Alpha is a scalar control value. Its dtype MUST NOT determine factor computation dtype, merged tensor dtype, or save dtype.
+6. Alpha association MUST come from the repository's LoRA pair parser, not from an assumed key spelling. PEFT, Diffusers, ComfyUI, and supported mixed naming forms are governed by the parsed logical pair.
+7. Normalization MUST occur before rank padding or other shape alignment. Padding an unnormalized pair and applying its original alpha against the padded rank changes the represented delta.
+8. The unscaled factor reference MUST be replaced immediately when normalization creates a new tensor. Both source factors, the alpha scalar, normalized factors, padded factors, and device copies MUST be released when the logical layer completes.
+9. A low-bit alpha scalar is valid when both factors use supported floating storage. If an alpha-bearing down or up factor uses unsupported low-bit storage, the operation MUST fail with an actionable error because byte-exact preservation cannot also produce an alpha-free normalized result.
+10. Format metadata may describe that normalization occurred, but metadata MUST NOT be used as a substitute for applying the mathematical transformation.
+
+### CWB Mixed-Rank LoRA Alignment
+
+1. Genuine LoRA rank validity MUST come from each factor pair's original shape
+   or explicit source metadata. Tensor values MUST NOT be inspected to decide
+   whether a component is padding because a genuine component may be zero.
+2. Structural rank padding MUST NOT participate in CWB alignment, consensus,
+   similarity weighting, fallback behavior, or norm rescaling.
+3. Similarity alignment between rank `R` and rank `r` genuine pairs MUST compare
+   the full rectangular `R x r` component space. A smaller source component may
+   match any component in the maximum-rank reference.
+4. Index alignment MUST pair only shared genuine indices. Higher reference
+   indices remain singleton groups unless another genuine source contributes.
+5. An explicit zero inserted by `mismatch_mode=zeros` is a semantic contributor
+   and MUST remain distinct from structural rank absence.
+6. Down rows and corresponding up columns MUST share validity, alignment, sign
+   correction, contributor membership, and output position.
+7. Output rank MUST remain the largest genuine input rank. Equal maximum-rank
+   ties retain source order.
+
+The shared implementation in `nodes/lora_alpha.py` is mandatory for merge paths. A format-specific implementation may only replace it when the represented delta differs mathematically and that difference is documented and tested.
+
+## Required Verification
+
+Any implementation that reads, writes, merges, resizes, extracts, or applies model or LoRA tensors MUST include tests appropriate to its scope. Merge changes involving LoRA factors MUST cover:
+
+- alpha present and absent in the same operation;
+- non-neutral alpha where `alpha / rank` is not `1.0`;
+- normalization before padding or concatenation;
+- alpha-free output;
+- mixed factor ranks and dtypes;
+- full rectangular CWB matching without structural-padding contributors;
+- genuine zero-valued components versus explicit mismatch-zero contributors;
+- low-bit scalar alpha acceptance and unsupported low-bit factor rejection;
+- stream closure, `mark_processed`, and reference release on success and failure;
+- atomic destination preservation on writer failure.
+
+Repository guards MUST continue to reject production `get_tensor()` calls and direct safetensors imports. The complete test suite, changed-file lint, compilation, and `git diff --check` are required before commit.

docs/consensus_mergers.md

@@ -0,0 +1,348 @@
+# 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.
+
+The broad mathematical and integration reference is maintained in
+[`references/Spec-CWB`](../references/Spec-CWB/README.md). This document
+describes the behavior actually exposed by this repository's offline merger
+nodes. The specification covers additional contexts, including online
+conditioning and spatial fusion, which are not controls of these nodes.
+
+## Node families
+
+Two-input and three-input variants are provided for checkpoints, standalone
+diffusion models, text encoders, LoRAs, and embeddings. LoRA also provides a
+Multi-Merge variant for 2 to 8 equal-prior inputs.
+
+## Embedding coalescing
+
+`CWB Embedding Self-Coalesce` reduces one embedding by replacing mutually
+nearest body-vector pairs that meet both cosine and normalized-position limits.
+`CWB Embedding Multi-Merge` accepts 2 to 8 embeddings, performs normal CWB
+alignment first, then applies the same reduction. A target row count of zero
+continues until no eligible pair remains.
+
+Vision Boundary Embeddings preserves encoded vision-start and vision-end rows
+unchanged and applies CWB only to their interior. Normal vision inputs validate
+their outer rows before writing. Legacy Boundary Search trims accidental outer
+template rows by matching a selected known-clean visual embedding; an absent,
+ambiguous, or low-similarity boundary pair is an error rather than a guessed
+crop. Disable Vision Boundary Embeddings for open textual-inversion tensors.
+
+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
+
+Presets are separated by merge type. Dense tensors, embeddings, and LoRA
+factor pairs never share one menu. Every named preset supplies every CWB
+setting and cannot inherit hidden widget state.
+
+Manual controls live in the separate CWB Custom Configuration node. Connecting
+its `CWB_CONFIG` output completely overrides the selected preset. Leaving it
+unconnected uses the named preset.
+
+Preset names expose categorical and boolean behavior. `_idx` and `_sim`
+identify alignment, `_medn` and `_mean` identify consensus, and `_rn`, `_dsc`,
+`_softcb`, and `_pcp` identify enabled boolean transformations. Dense presets
+omit alignment tokens because dense model coordinates are fixed.
+
+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.
+
+### Weighting sequence
+
+For each aligned vector group, the merger performs these operations in order:
+
+1. Compute an element-wise mean or median consensus.
+2. Measure each vector's cosine similarity to that consensus.
+3. Reject similarities below `similarity_threshold`.
+4. Optionally remap unequal accepted similarities with Dynamic Similarity
+   Contrast.
+5. Apply `similarity^power_alpha` and, when enabled, the diversity bandpass.
+6. Normalize the resulting contributor weights and calculate their weighted
+   sum. If no usable weight remains, use equal contributor weights.
+7. Optionally apply `rescale_norm`.
+8. Apply `global_scale`.
+
+### Complete control reference
+
+- **Execution Mode**: `MERGE` writes an output. `DOCUMENTATION ONLY` returns
+  this document without opening model inputs.
+- **Model A**: primary contributor and preservation anchor. It supplies output
+  metadata and anchors shared names and shapes. The output key set remains the
+  union of all inputs.
+- **Model B / Model C**: additional equal-prior contributors. CWB derives their
+  effective per-vector influence from consensus similarity; these nodes do not
+  expose model-level weights.
+- **CWB Preset**: complete use-case settings selected from the node's own dense,
+  embedding, or LoRA registry. A connected CWB Config overrides it completely.
+- **Consensus Type**: `mean` uses the arithmetic center; `median` uses the
+  element-wise median and is less sensitive to coordinate outliers.
+- **Alignment Method**: `index` pairs first-axis vectors by position;
+  `similarity` uses greedy one-to-one cosine matching. Row reordering is
+  enabled only for embeddings and paired LoRA rank components. Ordinary
+  checkpoint, diffusion-model, and text-encoder tensors retain fixed
+  coordinates.
+- **Alignment Threshold**: minimum original cosine score for accepting a
+  greedy row match. It does not filter already-paired contributors during CWB
+  weighting; that is `similarity_threshold`.
+- **Similarity Threshold**: minimum cosine similarity between a participating
+  vector and the group consensus. Rejected vectors receive zero weight. If all
+  are rejected, the implementation falls back to equal weights.
+- **Power Alpha**: exponent applied to accepted non-negative similarities.
+  Values above 1 increasingly favor contributors closest to consensus. Zero
+  makes every accepted contributor's similarity term equal to one.
+- **Diversity Beta**: exponent for the diversity bandpass. Values above zero
+  suppress vectors extremely close to consensus relative to more varied
+  accepted vectors. Zero disables this term.
+- **Rescale Norm**: after blending, preserve the merged direction but replace
+  its L2 magnitude with the mean L2 magnitude of the participating vectors:
+
+  `merged = merged / ||merged|| × mean(||source_i||)`
+
+  A zero merged vector remains unchanged. This is performed per aligned row,
+  not once per complete tensor. For LoRA, it is performed separately for each
+  paired A row and B column, so enabling it can materially change the effective
+  low-rank delta magnitude.
+- **Global Scale**: multiplies the merged result after CWB. For LoRA it is
+  applied once through B/up so the represented delta is scaled once rather
+  than once per factor.
+- **Dynamic Similarity Contrast**: when consensus similarities differ, remaps
+  them into the range 0.7–1.0 before alpha/beta weighting. It does not alter
+  row-alignment scores.
+- **Soft Comfort Bandpass**: when beta is positive, uses
+  `(1.5 - similarity)^beta` instead of the narrower
+  `(1.001 - similarity)^beta`. It has no effect when beta is zero.
+- **Position Weight**: blends normalized positional affinity into greedy row
+  selection. Zero uses cosine only and one uses positional affinity only. It
+  affects only similarity-aligned embeddings and LoRA rank components; the
+  original cosine score is still checked against the alignment threshold.
+- **Preserve Common Prefix**: copies Model A's longest numerically identical
+  leading first-axis span before CWB. The current implementation applies this
+  to all tensors with at least two dimensions, including fixed-coordinate
+  model tensors.
+- **Mismatch Mode**: for anchored tensors that are missing or incompatible,
+  `skip` preserves the anchor, `zeros` supplies a compatible zero contribution
+  where possible, and `error` aborts. A valid tensor supplied by only one
+  secondary input is copied unchanged.
+- **Output Filename**: filename without extension, written atomically beneath
+  the matching category in ComfyUI's models directory.
+- **Save Dtype**: requested dtype for generated floating tensors. Unless
+  Override Dtype is enabled, any participating FP32 source keeps the generated
+  logical result FP32.
+- **Process Device**: device for per-layer FP32 arithmetic. A CUDA OOM releases
+  the failed working set and retries only that layer on CPU.
+- **Exclude Patterns**: matching layers are preserved from their anchor instead
+  of merged.
+- **Discard Patterns**: matching tensors or logical LoRA groups are absent from
+  the output.
+- **Glob Patterns**: switches exclude/discard syntax from regular expressions
+  to shell-style globs. Pattern fields accept one entry per line.
+- **Lazy Load**: enables UEL low-memory loading so tensors are streamed by work
+  unit and explicitly released.
+- **Force Clear Cache**: runs garbage collection and clears the CUDA allocator
+  cache before each layer. This can reduce retained memory at a substantial
+  speed cost.
+- **Override Dtype**: forces generated floating tensors to Save Dtype. Guarded
+  low-bit tensors and enabled 1D direct differences are exempt.
+- **Include 1D Diffs** (LoRA nodes): enables CWB merging of recognized 1D
+  direct-difference tensors in FP32. When disabled, shared tensors preserve
+  Model A and secondary-only tensors preserve their earliest provider.
+
+Scalars and 1D generic tensors use direct mean/median handling. Alignment,
+similarity weighting, alpha/beta, DSC, and norm rescaling apply to vector groups
+in tensors with at least two dimensions, not to those scalar/1D paths.
+
+### Presets implemented by these nodes
+
+Dense nodes offer `balanced_mean`, `robust_medn`, `selective_mean`,
+`varied_mean_rn_softcb`, `diverse_medn_rn_dsc_softcb`, and
+`strongdiv_medn_rn_dsc_softcb`.
+
+Embedding nodes offer `balanced_idx_mean`, `balanced_sim_mean`,
+`robust_idx_medn`, `robust_sim_medn`, `varied_sim_mean_rn_softcb`, and
+`diverse_sim_medn_rn_dsc_softcb`, `focused_strong_sim_medn`,
+`focused_balance_sim_medn`, `focused_soft_sim_medn`, and
+`focused_weak_sim_medn`.
+
+LoRA nodes offer `broad_sim_medn_rn_softcb`,
+`moderate_sim_medn_rn_softcb`, `conservative_sim_medn_rn_softcb`,
+`direct_idx_medn_rn_softcb`, `broad_sim_mean_rn_softcb`,
+`neutral_sim_medn_rn`, `focused_sim_medn_rn_dsc_softcb`, and
+`strongfocus_sim_medn_rn_dsc_softcb`, `focused_strong_sim_medn`,
+`focused_balance_sim_medn`, `focused_soft_sim_medn`, and
+`focused_weak_sim_medn`.
+
+The LoRA default is `broad_sim_medn_rn_softcb`: median consensus,
+similarity alignment, zero similarity threshold, alpha 2, beta 4, norm
+rescaling, scale 1, DSC disabled, soft comfort bandpass enabled, position
+weight 0.05, and common-prefix preservation disabled. Broad, moderate, and
+conservative use alignment thresholds 0, 0.0005, and 0.0025. They represent
+increasing anchor preservation in automated mixed-rank LoRA tests.
+
+Every LoRA preset fixes all 12 controls explicitly:
+
+| preset | consensus | alignment | align threshold | sim threshold | alpha | beta | RN | scale | DSC | soft CB | position | prefix | intent |
+|---|---|---|---:|---:|---:|---:|---|---:|---|---|---:|---|---|
+| `broad_sim_medn_rn_softcb` | median | similarity | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Broad row matching; default |
+| `moderate_sim_medn_rn_softcb` | median | similarity | 0.0005 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Moderate anchor preservation |
+| `conservative_sim_medn_rn_softcb` | median | similarity | 0.0025 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Strong anchor preservation |
+| `direct_idx_medn_rn_softcb` | median | index | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Fixed rank-row correspondence |
+| `broad_sim_mean_rn_softcb` | mean | similarity | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Mean-consensus broad merge |
+| `neutral_sim_medn_rn` | median | similarity | 0 | 0 | 2 | 0 | on | 1 | off | off | 0.05 | off | Similarity weighting without diversity modulation |
+| `focused_sim_medn_rn_dsc_softcb` | median | similarity | 0 | 0 | 2 | 4 | on | 1 | on | on | 0.05 | off | Moderate DSC concentration |
+| `strongfocus_sim_medn_rn_dsc_softcb` | median | similarity | 0 | 0 | 2 | 7 | on | 1 | on | on | 0.05 | off | Strong DSC concentration |
+| `focused_strong_sim_medn` | median | similarity | 0.85 | 0.60 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Strongly focused matching |
+| `focused_balance_sim_medn` | median | similarity | 0.75 | 0.55 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Balanced focused matching |
+| `focused_soft_sim_medn` | median | similarity | 0.55 | 0.50 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Soft focused matching |
+| `focused_weak_sim_medn` | median | similarity | 0.25 | 0.35 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Weak focused matching |
+
+The four focused median similarity presets have identical settings in the
+embedding and LoRA registries.
+
+The LoRA Multi-Merge counterfactual weight sweep reuses each consensus vector
+already computed by the active merge. It reports mean and median consensus
+results across configured alpha, beta, similarity-threshold, DSC, and comfort
+bandpass values without additional model loads or output files.
+
+## 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. Alpha is
+absorbed mathematically into each B/up factor using its original rank before
+alignment, weighting, or merge. Genuine factor ranks remain variable during
+CWB. The first maximum-rank input supplies the output component space, so the
+output rank remains the largest genuine input rank. Merged outputs never
+contain or invent alpha tensors.
+
+Every CWB merge node writes a `cwb.merge` safetensors metadata record. It
+contains selected source model names, fully resolved CWB settings, and merge
+options. Presets are stored as effective settings rather than by preset name.
+Non-empty exclude/discard filters and their glob mode are included. CWB reports
+and diagnostics are not saved in metadata.
+Similarity alignment operates on paired latent-rank components. A rows and B
+columns share one mapping; fixed input/output feature axes are never reordered.
+For a rank-R reference and rank-r source it evaluates an R-by-r similarity
+matrix, reduced only by a preserved common prefix. Every genuine source
+component can therefore match any genuine reference component. Index alignment
+maps source component i to reference component i. In either mode, reference
+components without a genuine contributor remain singleton groups and are not
+weighted or norm-rescaled against nonexistent rank padding.
+
+Structural rank absence is different from `mismatch_mode=zeros`. The latter is
+an explicit request for a semantic zero contributor and continues to apply to
+every output-rank component where the mismatch policy is active. Numerical zero
+values inside a genuine factor do not identify padding; validity comes from the
+factor's original rank.
+
+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.
+
+## CWB report output
+
+Every CWB merge node provides a separate `cwb_report` string output after the
+existing filename and documentation outputs. Documentation-only execution
+returns an explicit message that no report exists because no tensors were
+processed.
+
+The filename output uses ComfyUI's wildcard socket type and returns the path
+relative to the matching ComfyUI model category. A filename such as
+`groupfolder/modelname` therefore returns
+`groupfolder/modelname.safetensors`. It can connect directly to compatible
+loader Combo inputs without exposing the absolute models-directory path.
+
+The report begins with model-wide counters. LoRA arithmetic then adds grouped
+mathematical evidence without retaining tensors or dumping complete similarity
+matrices. Layers are grouped only when they share:
+
+- down non-rank dimensions;
+- up non-rank dimensions;
+- ordered genuine input ranks;
+- selected reference input and rank;
+- alignment method.
+
+Each dimension group reports layer count, effective matrix dimensions, genuine
+source-component match coverage, reference-only components, structural rank
+slots excluded from arithmetic, explicit mismatch-zero contributors, candidate
+similarity min/mean/max, accepted-match similarity percentiles, and rank-one
+delta norm ratios. A compact exception section identifies the lowest match
+coverage and largest norm changes by layer name.
+
+For a rank-384 reference and rank-128 source, the report identifies a `384x128`
+matrix, at most 128 genuine matches, and the 256 structural rank slots that were
+excluded rather than treated as zero-valued contributors.
+
+## 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.

docs/cwb_delta_lora_mergers.md

@@ -0,0 +1,23 @@
+# Delta CWB LoRA Mergers
+
+Delta CWB expands each selected LoRA into full per-layer weight changes and
+applies CWB to corresponding dense rows. It preserves model coordinates; it
+does not align, reorder, or compress rank components.
+
+LoRA 1 supplies metadata and defines the output layer set. The preset controls
+how agreement, disagreement, and row magnitude affect the saved changes. A
+connected CWB configuration completely replaces the preset. Inputs have equal
+prior weight.
+
+Missing-layer handling can preserve LoRA 1, include explicit zero contributors,
+or abort. Excluded layers preserve LoRA 1 and discarded layers are omitted.
+Outputs contain only canonical `.diff` tensors with alpha already normalized.
+
+`factorized_lora` output factorizes each completed merged layer immediately and
+writes canonical `lora_A` and `lora_B` tensors at the selected maximum rank. It
+never writes or retains a complete full-difference model. One-dimensional
+differences remain `.diff` tensors.
+
+Processing streams one logical layer from each input at a time and writes each
+completed layer incrementally. CUDA out-of-memory retries only the failed layer
+on CPU. The destination is replaced only after the complete output is written.

docs/cwb_model_graft.md

@@ -0,0 +1,135 @@
+# CWB Model Graft
+
+`CWB Model Graft` (`CWBModelGraft`) is a standalone training-free diffusion model grafting node that transfers selected blocks, operator families, projection sub-slices, or explicitly mapped layers from a donor diffusion model (`model_b`) into a base diffusion model (`model_a`) using Consensus-Weighted Blending (CWB).
+
+`model_a` defines the output architecture, tensor key set, output shapes, and safetensors metadata. Tensors that are protected, excluded by block or sub-part rules, or absent from `model_b` are preserved from `model_a` without modification.
+
+---
+
+## End-to-End Execution Pipeline
+
+For every logical tensor key in `model_a`, `CWB Model Graft` executes the following sequence:
+
+1. **Quantization & Low-Bit Inspection**
+   - Both `model_a` and `model_b` are opened with `MemoryEfficientSafeOpen` and inspected via `DiffusionQuantization` and `diffusion_key_map`.
+   - Full-model ComfyUI/comfy-kitchen quantized files or files containing 3 or more bare low-bit (`INT8`, `UINT8`, `FP8`, `FP4`) tensors are rejected before opening an output writer.
+   - Any layer with 1–2 isolated low-bit tensors in either `model_a` or the resolved `model_b` source keys is preserved byte-for-byte from `model_a` (`force_raw=True`).
+
+2. **Discard & Exclude / Include Pattern Filtering**
+   - `discard_patterns`: If a `model_a` key (or its physical key) matches `discard_patterns`, the tensor is omitted from the output file.
+   - `exclude_patterns` with `include_mode = False` (default): Matching keys are preserved unchanged from `model_a`.
+   - `exclude_patterns` with `include_mode = True`: Acts as a whitelist. Non-matching keys are preserved unchanged from `model_a`. Explicitly matched whitelist keys are eligible for grafting even if they lie outside standard block numbering (such as `text_embedder` or modulation layers).
+
+3. **Donor Key Resolution & Automatic QKV Packing / Splitting**
+   - **Custom Key Map (`custom_key_map`)**: Checked first. If `base_key` has an explicit mapping `donor_key -> base_key` and `donor_key` exists in `model_b`, `donor_key` is used as the donor source.
+   - **Direct Match**: If `base_key` exists directly in `model_b`, it is paired 1-to-1.
+   - **Split-to-Fused QKV (`fuse_qkv`)**: If `model_a` has a fused `<prefix>.proj.qkv.weight` tensor and `model_b` instead stores split `<prefix>.proj.q.weight`, `<prefix>.proj.k.weight`, and `<prefix>.proj.v.weight`, all three donor tensors are streamed in the same work unit and concatenated along `dim=0` (`torch.cat([q, k, v], dim=0)`) before sub-part selection and CWB blending.
+   - **Fused-to-Split QKV (`split_qkv`)**: If `model_a` has split `<prefix>.proj.{q,k,v}.weight` and `model_b` stores fused `<prefix>.proj.qkv.weight`, the fused donor tensor is streamed and split into 3 equal chunks along `dim=0` (`chunk(3, dim=0)`), selecting index `0` for `q`, `1` for `k`, or `2` for `v`.
+   - **Unmatched Keys**: If no donor source exists for `base_key`, `model_a`'s tensor is preserved unchanged.
+
+4. **Block Index Detection & Early-Block Protection**
+   - Main transformer block indices are parsed from keys matching `(^|\.)(layers|blocks|double_blocks|single_blocks|transformer_blocks)\.(\d+)\.`.
+   - Auxiliary sub-networks (`text_embedder.*`, `token_refiner.*`) do not count toward main transformer block numbering and are preserved from `model_a` unless targeted via `custom_key_map` or `include_mode=True` with `exclude_patterns`.
+   - `max_block_index` is computed as the highest main block index in `model_a`.
+   - If a key's `block_idx < protect_early_blocks`, the tensor is preserved 100% from `model_a`.
+
+5. **Operator Classification & Block Scheduling**
+   - Each main-block key is classified by its token path:
+     - **Attention (`attention`)**: `.qkv.`, `.qkv_proj.`, `.proj.q.`, `.proj.k.`, `.proj.v.`, `.proj.o.`, `.attn.`, `.self_attn.`, `.q_norm.`, `.k_norm.`, `.to_q.`, `.to_k.`, `.to_v.`, `.to_out.`.
+     - **Feed-Forward (`ffn`)**: `.ffn_up_gate.`, `.ffn_up.`, `.ffn_gate.`, `.ffn_down.`, `.mlp.`, `.ff.`, `.fc1.`, `.fc2.`, `.w1.`, `.w2.`, `.w3.`.
+   - Keys inside main blocks that are neither attention nor feed-forward (such as block-level `norm1` and `norm2`) are preserved from `model_a` unless explicitly targeted via `custom_key_map` or `include_mode=True`.
+   - For classified keys with `block_idx >= protect_early_blocks`, the corresponding rule (`attention_block_rule` or `ffn_block_rule`) determines whether the block is grafted:
+     - `all_after_protected`: Every block with `block_idx >= protect_early_blocks`.
+     - `interleaved_even`: Even-numbered blocks (`block_idx % 2 == 0`) with `block_idx >= protect_early_blocks`.
+     - `interleaved_odd`: Odd-numbered blocks (`block_idx % 2 == 1`) with `block_idx >= protect_early_blocks`.
+     - `back_half_only`: Blocks in the second half of the network (`block_idx >= max(protect_early_blocks, (max_block_index + 1) // 2)`).
+     - `none`: No blocks of this operator family are grafted.
+
+6. **Sub-Part Slicing (`attention_subpart` and `ffn_subpart`)**
+   - **Attention Sub-Parts (`attention_subpart`)**:
+     - `all`: All attention tensors in selected blocks (`qkv`, `q`, `k`, `v`, `o`, `q_norm`, `k_norm`) are grafted.
+     - `value_and_output`:
+       - Split query/key projections and norms (`.proj.q.`, `.proj.k.`, `.to_q.`, `.to_k.`, `.q_norm.`, `.k_norm.`) are preserved 100% from `model_a`.
+       - Fused QKV weights (`.qkv.`, `.qkv_proj.`) whose first dimension is divisible by 3 are split into `q_a, k_a, v_a = base.chunk(3, dim=0)` and `q_b, k_b, v_b = donor.chunk(3, dim=0)`. `q_a` and `k_a` are kept byte-exact from `model_a`, `v_a` and `v_b` are CWB-blended into `v_merged`, and the tensor is reassembled as `torch.cat([q_a, k_a, v_merged], dim=0)`.
+       - Split value projections (`.proj.v.`, `.to_v.`) and output projections (`.proj.o.`, `.out_proj.`, `.to_out.`) are CWB-blended in full.
+     - `output_only`: Only attention output projections (`.proj.o.`, `.out_proj.`, `.to_out.`) are CWB-blended; all Q/K/V and norm tensors are preserved 100% from `model_a`.
+   - **Feed-Forward Sub-Parts (`ffn_subpart`)**:
+     - `all`: All feed-forward tensors in selected blocks (`ffn_up_gate`, `ffn_up`, `ffn_gate`, `ffn_down`, `fc1`, `fc2`, `w1`, `w2`, `w3`) are grafted.
+     - `up_only`:
+       - Down and gate projections (`.ffn_down.`, `.fc2.`, `.w2.`, `.ff.net.2.`, `.ffn_gate.`, `.w3.`) are preserved 100% from `model_a`.
+       - Fused up+gate weights (`.ffn_up_gate.`) whose first dimension is divisible by 2 are split into `up_a, gate_a = base.chunk(2, dim=0)` and `up_b, gate_b = donor.chunk(2, dim=0)`. Only `up_a` and `up_b` are CWB-blended into `up_merged`, `gate_a` is preserved 100% from `model_a`, and the tensor is reassembled as `torch.cat([up_merged, gate_a], dim=0)`.
+       - Split up projections (`.ffn_up.`, `.fc1.`, `.w1.`) are CWB-blended in full.
+     - `down_only`: Only down projections (`.ffn_down.`, `.fc2.`, `.w2.`, `.ff.net.2.`) are CWB-blended; all up/gate tensors are preserved 100% from `model_a`.
+
+7. **Cross-Shape Alignment, Row Similarity Matching & Norm Locking**
+   - **Trailing-Dimension Alignment**: If `donor_tensor` and `base_tensor` differ in shape:
+     - Same element count with different rank: `donor_tensor` is reshaped to `base_tensor.shape`.
+     - Different rank and different element count: `donor_tensor` is linearly interpolated in FP32 to `base_tensor.numel()` and reshaped to `base_tensor.shape`.
+     - Same rank (`ndim >= 2`) with different trailing dimensions (`donor_tensor.shape[1:] != base_tensor.shape[1:]`): Each donor row is linearly interpolated in FP32 along its flattened trailing features to match `prod(base_tensor.shape[1:])` and reshaped to `(donor_rows, *base_tensor.shape[1:])`.
+   - **Row Matching & CWB Blending**: `merge_cwb_tensors` runs in FP32 with `allow_similarity_alignment=True` and `reference_index=0` (`model_a` as reference). When the active CWB settings use `alignment_method="similarity"` (or when row counts differ), donor rows are greedily paired to base rows by cosine similarity (subject to `alignment_threshold` and `position_weight`). Unmatched base rows remain anchored to `model_a`.
+   - **Base Magnitude Locking**: Whenever `donor_tensor` required shape alignment or `rescale_norm` is enabled in the active CWB settings, the blended tensor's overall Frobenius norm is locked to `model_a`'s tensor norm (`merged * (||base|| / ||merged||)`), preventing activation scale drift across downstream blocks.
+
+8. **Atomic Streaming Output & Metadata**
+   - Tensors are streamed one work unit at a time through UEL (`batch_size=1`, `prefetch_batches=1`), written incrementally via `atomic_uel_writer`, and released before advancing to the next layer.
+   - If a CUDA out-of-memory error occurs on a layer, the failed GPU working set is released and that layer retries on CPU.
+   - Output metadata preserves `model_a`'s metadata and records a JSON summary under `cwb.graft` containing `base_model`, `donor_model`, `attention_preset`, `ffn_preset`, `protect_early_blocks`, `attention_block_rule`, `ffn_block_rule`, `attention_subpart`, and `ffn_subpart`.
+
+---
+
+## Complete Input Reference
+
+- **`execution_mode`** (`MERGE` | `DOCUMENTATION ONLY`, default `MERGE`):
+  - `MERGE`: Executes the streaming graft and writes the output `.safetensors` file under `models/diffusion_models/`.
+  - `DOCUMENTATION ONLY`: Returns this documentation string without opening any model files.
+- **`model_a`** (Combo, `diffusion_models`):
+  - Base diffusion model that anchors architecture, output tensor keys, shapes, metadata, and untouched/protected layers.
+- **`model_b`** (Combo, `diffusion_models`):
+  - Donor diffusion model from which selected blocks, sub-projections, or mapped layers are grafted into `model_a`.
+- **`attention_preset`** (Combo, default `robust_medn`):
+  - Dense CWB preset applied to grafted attention tensors (`balanced_mean`, `robust_medn`, `selective_mean`, `varied_mean_rn_softcb`, `diverse_medn_rn_dsc_softcb`, `strongdiv_medn_rn_dsc_softcb`). Overridden completely if `attention_config` is connected.
+- **`attention_config`** (`CWB_CONFIG`, optional):
+  - Optional custom CWB configuration from `CWB Custom Configuration` for attention layers. Enabling `alignment_method = "similarity"` on this config activates greedy cosine row matching across attention rows.
+- **`ffn_preset`** (Combo, default `balanced_mean`):
+  - Dense CWB preset applied to grafted feed-forward (MLP) tensors and custom-mapped non-attention tensors. Overridden completely if `ffn_config` is connected.
+- **`ffn_config`** (`CWB_CONFIG`, optional):
+  - Optional custom CWB configuration from `CWB Custom Configuration` for feed-forward (MLP) layers.
+- **`protect_early_blocks`** (Int, default `2`, range `0..256`):
+  - Number of initial main transformer blocks (`0 .. protect_early_blocks - 1`) kept 100% identical to `model_a` so global image composition and layout remain stable.
+- **`attention_block_rule`** (`interleaved_even` | `interleaved_odd` | `back_half_only` | `all_after_protected` | `none`, default `interleaved_even`):
+  - Block selection schedule for attention layers at index `>= protect_early_blocks`.
+- **`ffn_block_rule`** (`all_after_protected` | `interleaved_even` | `interleaved_odd` | `back_half_only` | `none`, default `all_after_protected`):
+  - Block selection schedule for feed-forward (MLP) layers at index `>= protect_early_blocks`.
+- **`attention_subpart`** (`all` | `value_and_output` | `output_only`, default `all`):
+  - Restricts attention grafting to all attention weights, only the Value (`V`) slice/projection and Output (`O`) projection, or only the Output (`O`) projection.
+- **`ffn_subpart`** (`all` | `up_only` | `down_only`, default `all`):
+  - Restricts feed-forward grafting to all FFN weights, only the Up projection (including the `up` half of packed `ffn_up_gate`), or only the Down projection.
+- **`custom_key_map`** (String, multiline, default `""`):
+  - Optional explicit key mappings, one per line, formatted as `donor_key -> base_key` or `donor_key = base_key` (lines starting with `#` are ignored). Mapped keys bypass block-number filtering so arbitrary donor tensors can be grafted into specific base tensors.
+- **`output_filename`** (String, default `cwb_grafted_model`):
+  - Output filename relative to `models/diffusion_models/` (`.safetensors` extension added automatically).
+- **`save_dtype`** (`fp32` | `fp16` | `bf16`, default `bf16`):
+  - Target floating-point precision when `override_dtype` is enabled or when decoding quantized tensors.
+- **`process_device`** (`cuda` | `cpu`, default `cuda`):
+  - Execution device for per-tensor FP32 CWB math; automatically falls back to `cpu` per tensor on CUDA OOM.
+- **`exclude_patterns`** (String, multiline, default `""`):
+  - Regex (or glob when `glob_patterns=True`) patterns. Preserves matching tensors from `model_a` when `include_mode=False`, or whitelists matching tensors for grafting when `include_mode=True`.
+- **`discard_patterns`** (String, multiline, default `""`):
+  - Regex (or glob) patterns of tensors to omit completely from the saved model.
+- **`glob_patterns`** (Boolean, default `False`):
+  - Interpret `exclude_patterns` and `discard_patterns` as shell-style globs (`*`, `?`) instead of regular expressions.
+- **`lazy_load`** (Boolean, default `True`):
+  - Enable low-memory UEL streaming so only the current work unit resides in memory.
+- **`force_clear_cache`** (Boolean, default `False`):
+  - Run `gc.collect()` and `torch.cuda.empty_cache()` before each tensor work unit.
+- **`override_dtype`** (Boolean, default `False`):
+  - Force grafted floating tensors to `save_dtype`; when `False`, each grafted tensor preserves `model_a`'s floating dtype.
+- **`include_mode`** (Boolean, default `False`):
+  - Invert `exclude_patterns` into a whitelist so only matching tensors are eligible for grafting.
+
+---
+
+## Outputs
+
+1. **`output_filename`** (`*` / `io.AnyType`): Category-relative filename (`<output_filename>.safetensors`) ready to connect directly to diffusion model loader nodes.
+2. **`documentation`** (`STRING`): Complete markdown documentation for `CWB Model Graft`.
+3. **`cwb_report`** (`STRING`): Diagnostic report detailing active attention/FFN presets, candidate/accepted alignment matches, consensus similarity statistics, CPU fallbacks, total grafted tensor count, and total preserved base tensor count.

docs/lodestone_mergers.md

@@ -0,0 +1,37 @@
+# Lodestone LoRA Mergers
+
+The dedicated Lodestone nodes merge full logical LoRA deltas rather than
+blending down and up factors independently. Stored alpha is normalized before
+each factor pair is expanded. Results are saved as canonical full direct-difference
+tensors and contain no factor or alpha tensors.
+
+The 2-, 3-, and Multi-Merge nodes treat every selected LoRA as one equal-prior
+slot, matching the supplied Lodestone calculation.
+
+## Calculation Modes
+
+- `sum`: sum the source deltas.
+- `mean`: divide the summed delta by the number of active contributors.
+- `slotnorm`: scale each source delta to the median source Frobenius norm, then
+  average.
+- `normmatch`: retain the summed direction and set its Frobenius norm to the
+  median source norm.
+- `slotnorm-normmatch`: equalize source norms, sum, then set the final norm to
+  the original median source norm.
+
+A zero target or merged norm remains zero.
+
+## Memory Behavior
+
+UEL streams one logical layer at a time and writes each completed full-delta
+tensor incrementally through the atomic writer. Dense tensors used for the
+merge and Frobenius norms exist only within the active layer. A CUDA
+out-of-memory error releases CUDA state and retries that logical layer on CPU.
+
+## Fixed-Input Controls
+
+LoRA 1 anchors output layers and filtering behavior. `skip` preserves LoRA 1
+when another input lacks a required layer, `zeros` supplies an explicit zero
+slot, and `error` aborts. Source factors may have different ranks when their
+expanded dense shapes agree. Exclude and discard patterns apply to complete
+logical layers.

docs/lora_model_analysis.md

@@ -0,0 +1,30 @@
+# LoRA-on-model analysis
+
+**Analyze LoRA Effect on Diffusion Model** takes a LoRA as `model_a` and a base
+diffusion model as `model_b`. It applies the LoRA transiently and compares the
+original base weights against the patched weights. The measured update is
+`patched − original`; no merged checkpoint, delta file, or inference is produced.
+
+The comparison and CWB reports use original base as A and patched base as B.
+The CSV outputs contain per-tensor difference metrics and row cosine summaries.
+Unchanged floating base tensors remain in the analysis, so global metrics reflect
+the impact on the selected base model, not just the layers touched by the LoRA.
+Application and statistics use FP32, matching the existing analysis arithmetic;
+this does not simulate rounding to a separately saved model dtype.
+
+LoRA strength defaults to 1. Alpha/rank normalization, DoRA, LoCon, and reshape
+handling use the existing application path and ComfyUI adapter. Direct deltas,
+including bias deltas, are applied; explicit `set_weight` patches replace weights
+regardless of strength, matching the existing application behavior.
+
+Exclude Patterns and Include Mode filter **base tensor keys** before loading.
+Unmapped LoRA groups and unused source tensors are listed separately. Unsupported
+low-bit work units and non-floating base tensors are reported and omitted from
+numeric metrics; they are not decoded or represented as successfully analyzed.
+Application errors fail the operation rather than report a failed patch as zero
+impact. A CUDA out-of-memory error retries the complete affected unit on CPU.
+
+Loading uses the shared UEL work-unit streamer: one base tensor and its required
+adapter tensors, batch size 1 and prefetch depth 1. Sources are marked processed
+and streams closed on success or failure. Only scalar statistics and bounded top
+differences survive each unit. No whole model or LoRA tensor dictionary is built.

docs/merger_2_model_modes.md

@@ -2,121 +2,194 @@
 
 ---
 
+### Core Convention Across All 2-Model Modes
+- **Model A is ALWAYS the base / anchor model.**
+- **Model B is ALWAYS the donor / incoming model being merged into Model A.**
+- When interpolating or blending between weights:
+  $$\text{Output} = A \cdot (1 - \text{mask}) + B \cdot \text{mask}$$
+- A blend factor / weight of `0.0` yields **100% Model A** (pure Model A).
+- A blend factor / weight of `1.0` yields **100% Model B** (pure Model B).
+- Intermediate values control how much of Model B replaces or alters Model A.
+
+---
+
 ## Weight-Sum
-> `A * (1 - α) + B * α`. A simple linear interpolation between two models.
+> $\text{Output} = A \cdot (1 - \alpha) + B \cdot \alpha$. Direct linear interpolation between Model A and Model B.
+
+**Models Used:** A (Base), B (Incoming)
 
-**Models Used:** A, B
 **Parameters:**
-- **Alpha:** Interpolation weight. `0.0` is 100% Model A, `1.0` is 100% Model B.
+- **Alpha:** The blend ratio determining what fraction of **Model B** replaces **Model A**.
+  - `0.0`: **100% Model A, 0% Model B** (exact Model A; Model B is completely ignored).
+  - `0.25`: **75% Model A, 25% Model B** (mostly Model A, with 25% of Model B blended in).
+  - `0.50`: **50% Model A, 50% Model B** (equal 50/50 mix of both models).
+  - `0.75`: **25% Model A, 75% Model B** (mostly Model B, retaining 25% of Model A).
+  - `1.0`: **0% Model A, 100% Model B** (exact Model B; Model A is completely replaced).
 
 ---
 
 ## Comparative-Interpolation
-> Interpolates between A and B based on the relative differences in their tensor values. Creates a unique blend by deciding how much of each model to use for each individual weight.
+> Interpolates between A and B based on the relative differences in their tensor values, creating a selective blend rather than a flat uniform mix.
+
+**Models Used:** A (Base), B (Incoming)
 
-**Models Used:** A, B
 **Parameters:**
-- **Alpha:** Controls the interpolation curve's strength and shape.
-- **Beta:** Switches the interpolation style. `0.0` focuses on similarity, `1.0` focuses on difference.
-- **Gamma:** Mixes between a randomized (binomial) interpolation at `0.0` and a smooth (linear) one at `1.0`.
+- **Alpha (Blend Intensity):** Controls the overall pull of **Model B into Model A**.
+  - `0.0`: Heavily favors Model A (virtually no Model B incorporated).
+  - `0.5`: Balanced curve where weights with moderate differences blend equally.
+  - `1.0`: Strongly pulls in Model B wherever weights differ.
+- **Beta (Feature Focus):** Controls whether the merge targets similar or divergent weights.
+  - `0.0`: Focuses blending on weights where Model A and Model B are already similar.
+  - `1.0`: Focuses blending on weights where Model A and Model B differ most.
+- **Gamma (Transition Style):** Mixes between discrete selection and smooth gradation.
+  - `0.0`: Stochastic / binomial selection (individual weights are chosen probabilistically as either pure A or pure B).
+  - `1.0`: Smooth linear interpolation across all weights.
 
 ---
 
 ## Power-Up (DARE)
-> Adds the unique capabilities of Model B to Model A using the Drop and Rescale (DARE) technique. This implementation handles shape mismatches between models by padding and uses a randomized dropout mask.
+> Adds the unique capabilities of Model B onto Model A using the Drop and Rescale (DARE) technique. Computes the task delta $(B - A)$, randomly drops a fraction of the delta parameters, rescales the survivors to maintain expected magnitude, and adds them to Model A.
+
+**Models Used:** A (Base), B (Donor)
 
-**Models Used:** A, B
 **Parameters:**
-- **Alpha:** The dropout rate ($p$). This is the proportion of delta parameters from Model B that are randomly set to zero.
-- **Beta:** A final multiplier for the rescaled difference before it's added to Model A.
-- **Rescaling Logic:** Remaining weights are automatically rescaled by $1/(1-p)$ as per the DARE paper to approximate the original embeddings.
+- **Alpha (Dropout Rate $p$):** Proportion of Model B's delta parameters to randomly zero out.
+  - Range: `0.0` to `1.0`.
+  - Typical: `0.7` to `0.9`. For example, `0.85` means 85% of Model B's delta values are discarded, keeping only 15% of the most distinct changes.
+- **Beta (Scale Multiplier):** Final strength multiplier for Model B's rescaled delta before adding to Model A.
+  - `1.0`: Standard 100% strength addition of Model B's surviving features.
+  - `0.5`: Half-strength addition of Model B's features.
+  - `1.2`: Amplified addition of Model B's features.
+- **Rescaling Logic:** Surviving parameters are automatically scaled by $1 / (1 - p)$ per the DARE paper to preserve the expected total update magnitude.
 
 ---
 
-## Enhanced Man Interp
-> Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control.
+## Power-Up (DARE+TIES)
+> Combines DARE (Drop and Rescale) with TIES (Trim, Elect Sign) to merge fine-tuned models cleanly. DARE sparsifies the delta; TIES trims low-magnitude background noise and enforces sign consensus.
+
+**Models Used:** A (Base), B (Donor)
 
-**Models Used:** A, B
 **Parameters:**
-- **Alpha:** Interpolation strength.
-- **Beta:** Lower mean threshold for filtering differences.
-- **Gamma:** Upper mean threshold for filtering differences.
-- **Delta:** Smoothness factor (mix between randomized mask and powered differences).
+- **Alpha (DARE Drop Rate):** Fraction of Model B's delta parameters randomly dropped (`0.5`–`0.9`). Higher values produce a sparser capability transfer.
+- **Beta (TIES Trim Quantile):** Fraction of the smallest-magnitude delta parameters eliminated after DARE to remove noise. `0.0` disables trimming; `0.2` trims the lowest 20% by absolute magnitude. Typical: `0.1`–`0.3`.
+- **Gamma (Lambda Scale / Strength):** Multiplier for the filtered delta added to Model A.
+  - `1.0`: Standard 100% strength of Model B's filtered capability added to Model A.
+  - `0.25`: Gentle 25% addition of Model B's capabilities.
+  - `1.5`: Amplified capability transfer.
+- **Seed:** Random seed for reproducible DARE dropout masks.
+
+**Algorithm:**
+1. Compute task vector $\delta = B - A$
+2. **DARE:** Apply random binary mask with keep probability $1 - \alpha$; rescale survivors by $1 / (1 - \alpha)$
+3. **TIES trim:** Zero out parameters below the $\beta$-quantile magnitude threshold
+4. **TIES elect:** Determine dominant sign direction per position; zero out disagreeing parameters
+5. Return $A + \gamma \cdot \delta_{\text{filtered}}$
 
 ---
 
-## Enhanced Auto Interp
-> Automated version of the enhanced interpolation mode that dynamically calculates thresholds based on mean differences.
+## Enhanced Man Interp (Enhanced Manual Interpolation)
+> Selective interpolation between Model A and Model B based on normalized pairwise differences, constrained within a user-defined threshold band.
+
+**Models Used:** A (Base), B (Incoming)
+
+**Formula:**
+$$\text{diff} = \frac{\max(|A - B|) - |A - B|}{\max(|A - B|)}$$
+$$\text{mask}_{\text{threshold}} = (\beta < \text{mean}(\text{diff}) < \gamma)$$
+$$\text{blend\_strength} = \text{diff}^{(1/\alpha - 1)} \cdot \text{mask}_{\text{threshold}}$$
+$$\text{Output} = A \cdot (1 - \text{interpolated\_mask}) + B \cdot \text{interpolated\_mask}$$
 
-**Models Used:** A, B
 **Parameters:**
-- **Alpha:** Interpolation strength.
-- **Beta:** Threshold adjustment factor.
-- **Gamma:** Smoothness factor.
+- **Alpha (Model B Pull Strength):** Controls how strongly Model B is pulled into Model A.
+  - Mathematically, the exponent applied to similarity is $(1/\alpha - 1)$.
+  - **`Alpha = 0.25`**: Exponent is $(1/0.25 - 1) = 3$. This aggressively suppresses Model B's contribution toward zero. Only weights that are virtually identical between both models get even a minor blend of Model B; the final result remains overwhelmingly **Model A** (>90% Model A).
+  - **`Alpha = 0.50`**: Exponent is $(1/0.50 - 1) = 1$ (linear). Weights with high similarity incorporate Model B proportionally up to 50%.
+  - **`Alpha = 0.75`**: Exponent is $0.33$. Flattens similarity, pulling in substantially more of Model B wherever differences meet the threshold.
+  - **`Alpha = 1.0`**: Exponent is $0$. Full maximum pull of Model B across all qualifying weights.
+- **Beta (Lower Mean Threshold):** Minimum normalized similarity required before Model B can be blended into Model A. Layers/weights with similarity below $\beta$ remain 100% Model A.
+- **Gamma (Upper Mean Threshold):** Maximum normalized similarity allowed for blending. Layers/weights with similarity above $\gamma$ remain 100% Model A. Together, $[\beta, \gamma]$ acts as a bandpass filter targeting only weights with moderate differences.
+- **Delta (Smoothness Factor):** Blends between stochastic selection (`0.0`, a Bernoulli coin-flip choosing either pure A or pure B per weight) and smooth continuous blending (`1.0`, smooth linear weighting between A and B).
+
+---
+
+## Enhanced Auto Interp (Enhanced Automatic Interpolation)
+> Automated version of Enhanced Interpolation that dynamically computes threshold boundaries centered on the layer's mean difference.
+
+**Models Used:** A (Base), B (Incoming)
+
+**Formula:**
+$$\text{threshold\_band} = [\text{mean}(\text{diff}) \cdot (1 - \beta), \;\; \text{mean}(\text{diff}) \cdot (1 + \beta)]$$
+$$\text{Output} = A \cdot (1 - \text{interpolated\_mask}) + B \cdot \text{interpolated\_mask}$$
+
+**Parameters:**
+- **Alpha (Model B Pull Strength):** Controls how much of Model B replaces Model A inside the active band.
+  - `0.25`: Low Model B contribution; output remains mostly Model A.
+  - `0.50`: Moderate, balanced blend of Model B into Model A.
+  - `1.0`: Maximum Model B influence within the threshold band.
+- **Beta (Threshold Width Around Mean):** Defines how wide the window of affected weights is around the layer's average difference.
+  - `0.1`: Narrow window ($\pm 10\%$ around mean difference).
+  - `0.5`: Broad window ($\pm 50\%$ around mean difference).
+- **Gamma (Smoothness Factor):** Blends between stochastic Bernoulli selection (`0.0`) and smooth continuous linear interpolation (`1.0`).
 
 ---
 
 ## Weight-Sum Cutoff
-> A linear interpolation mode that only applies the merge to weights whose differences fall within a specific threshold range.
+> Linear interpolation that only blends Model B into Model A for weights whose relative difference falls within an explicit cutoff window.
+
+**Models Used:** A (Base), B (Incoming)
+
+**Formula:**
+$$\text{Output} = A \cdot (1 - \alpha \cdot \text{mask}) + B \cdot (\alpha \cdot \text{mask})$$
 
-**Models Used:** A, B
 **Parameters:**
-- **Alpha:** Interpolation weight (multiplier for the difference).
-- **Beta:** Upper threshold for the difference cutoff.
-- **Gamma:** Lower threshold for the difference cutoff.
+- **Alpha (Model B Mix Ratio):** The percentage of Model B blended into Model A for weights inside the window.
+  - `0.0`: 0% Model B (100% Model A everywhere).
+  - `0.25`: 25% Model B, 75% Model A for weights inside the cutoff window; weights outside remain 100% Model A.
+  - `0.50`: 50% Model B, 50% Model A inside the window.
+  - `1.0`: 100% Model B inside the window (complete replacement of qualifying weights).
+- **Beta (Upper Difference Threshold):** Upper cutoff for normalized difference.
+- **Gamma (Lower Difference Threshold):** Lower cutoff for normalized difference. Weights with difference between $\gamma$ and $\beta$ receive the blend; all other weights stay 100% Model A.
 
 ---
 
 ## SVD LoRA Extraction
-> Extracts the difference between Model B and Model A and saves it as a new LoRA file. **This mode outputs a LoRA, not a full model.** The output file will be saved to your `loras` directory.
+> Computes the difference between Model A and Model B and extracts it as a new LoRA file using Singular Value Decomposition (SVD). **Outputs a LoRA safetensors file into your `loras` directory, not a full model.**
+
+**Models Used:** A (Tuned / Modified Model), B (Base Model)
+- Computes delta: $\Delta W = A - B$
+- Low-rank approximation: $\Delta W \approx U S V^T \implies \text{lora\_B} \cdot \text{lora\_A}$
 
-**Models Used:** A (Tuned Model), B (Base Model)
 **Parameters:**
-- **Alpha:** The Rank (dimension) for the LoRA's standard layers.
-- **Beta:** The Rank (dimension) specifically for 3x3 convolution layers.
-- **Gamma:** The clamp quantile for weight values. `0.99` is a typical value to prevent outlier weights from dominating.
+- **Alpha (Linear Rank):** Target rank (dimension) for 2D dense/linear layer LoRA factors (e.g. `32`, `64`, `128`).
+- **Beta (Conv Rank):** Target rank specifically for 3x3 convolution layer LoRA factors (e.g. `16`, `32`).
+- **Gamma (Clamp Quantile):** Outlier clamping quantile (e.g. `0.99` clips the top 1% extreme values before SVD to prevent outlier distortion).
 
 ---
 
-## Layer Mismatch Handling
+## Layer Mismatch Handling (`mismatch_mode`)
 
-When merging models with different layer structures (e.g., fine-tuned models with added/removed layers, or LoRAs with partial layer coverage), the `mismatch_mode` parameter controls behavior:
+When Model A and Model B have differing layer structures or when merging partial LoRAs:
 
 | Mode | Behavior |
 |------|----------|
-| `skip` | Missing layers in B use A's values **(default)** |
-| `zeros` | Missing layers in B are treated as zeros |
-| `error` | Fail if any layer is missing |
+| `skip` | Missing layers in Model B use Model A's values unchanged **(default)** |
+| `zeros` | Missing layers in Model B are treated as zeros |
+| `error` | Raise an error and stop execution if any layer is missing in Model B |
 
-**Common scenarios:**
-- Merging a fine-tuned model with added/removed layers
-- Combining LoRA-extracted differences with base models
-- Cross-architecture experiments
-
-**Note:** Extra layers in Model B that don't exist in Model A are currently ignored. The output will always have the same layer structure as Model A.
+**Note:** Extra layers in Model B that do not exist in Model A are ignored. Output topology always mirrors Model A.
 
 ---
 
-## Layer Filtering (Regex Patterns)
+## Dtype Preservation & Override
 
-Use regex patterns to control which layers are merged:
+By default, the merger preserves the highest precision `dtype` for each key found across active source models and the requested `save_dtype`:
+- **Preserve Higher Precision (default, `override_dtype = False`)**: If a key is `float32` in Model A or Model B, but `save_dtype="bf16"` is selected, that key remains `float32` to protect sensitive normalization layers.
+- **Explicit Override (`override_dtype = True`)**: Forces all floating tensors strictly to `save_dtype`.
 
-### `exclude_patterns`
-Layers matching any pattern will **keep Model A's values only** (no merge).
+---
 
-### `discard_patterns`
-Layers matching any pattern will be **removed entirely** from the output.
+## Layer Filtering
 
-**Pattern format:**
-- Whitespace-separated regex patterns
-- Patterns use **substring matching** (not full match)
-- Example: `text_model lora` matches any key containing "text_model" OR "lora"
-- Example: `layer\.[0-5]\.` matches layers 0-5 using regex syntax
-
-**Pattern Examples:**
-| Pattern | Matches |
-|---------|---------|
-| `text_model` | All text encoder layers |
-| `\.norm` | All normalization layers |
-| `attn\.(q\|k\|v)` | Query, key, value attention weights |
-| `block\.[0-9]\.` | Blocks 0-9 |
+- **`exclude_patterns`**: Newline-separated patterns. Matching layers are **not merged** and keep Model A's values unchanged (unless `include_mode` is enabled).
+- **`include_mode`**: Inverts `exclude_patterns` into a whitelist: only matching layers are merged; all non-matching layers keep Model A's values unchanged.
+- **`discard_patterns`**: Newline-separated patterns. Matching layers are **removed completely** from the saved output.
+- **`glob_patterns`**: When enabled, uses shell wildcards (`*`, `?`). When disabled, uses regex substring matching.

docs/merger_3_model_modes.md

@@ -2,82 +2,111 @@
 
 ---
 
+### Core Convention Across All 3-Model Modes
+- **Model A is ALWAYS the base / receiving model** that receives the merge modifications.
+- **Model B is the donor / target fine-tune model.**
+- **Model C is the reference base model of Model B.**
+- In difference merging, $(B - C)$ calculates the isolated delta (the specific training / fine-tune changes introduced in Model B relative to its base Model C).
+- That delta is then added onto Model A with scaling $\alpha$:
+  $$\text{Output} = A + (B - C) \cdot \alpha$$
+
+---
+
 ## Add-Difference
-> `A + (B - C) * α`. This classic technique applies the "difference" between models B and C to model A. For example, it can be used to add a style (B) to a base model (A) while subtracting the style's original base (C).
+> $\text{Output} = A + (B - C) \cdot \alpha$. Extracts the isolated fine-tune / style delta from $(B - C)$ and applies it directly to Model A.
+
+**Models Used:** A (Receiving Base), B (Donor Fine-Tune), C (Donor's Base Reference)
 
-**Models Used:** A, B, C
 **Parameters:**
-- **Alpha:** A multiplier for the `(B - C)` difference vector before it's added to A.
-- **Beta:** A smoothing toggle. `0` for Off, `1` for On. When On, applies a median and gaussian filter to the difference, which can result in a smoother merge.
+- **Alpha (Delta Multiplier / Strength):** Scaling factor for how strongly the $(B - C)$ difference is added to Model A.
+  - `0.0`: **0% addition** (Output is pure Model A; no changes added).
+  - `0.25`: **25% strength addition** (gentle, subtle application of Model B's style/features onto Model A).
+  - `0.50`: **50% strength addition** (moderate application).
+  - `1.0`: **100% full-strength addition** (applies the exact full delta of B relative to C onto Model A).
+  - `1.5`: **150% over-amplified addition** (exaggerates the difference; can introduce visual artifacts if pushed too high).
+- **Beta (Smoothing Toggle):**
+  - `0`: Off (raw mathematical delta added directly).
+  - `1`: On (applies median and Gaussian smoothing to the delta before adding, reducing harsh pixel noise and high-frequency artifacts).
 
 ---
 
 ## Train-Difference
-> A variation of Add-Difference from the SuperMerger extension that uses a more complex scaling factor based on the relative distances between the tensors of A, B, and C.
+> A variation of Add-Difference that calculates an adaptive scaling factor based on the relative vector distances between Model A, Model B, and Model C.
+
+**Models Used:** A (Receiving Base), B (Donor Fine-Tune), C (Donor's Base Reference)
 
-**Models Used:** A, B, C
 **Parameters:**
-- **Alpha:** A multiplier for the final calculated "training" difference.
+- **Alpha (Difference Multiplier):** Multiplier for the calculated distance-weighted difference added to Model A.
+  - `0.0`: Pure Model A.
+  - `1.0`: Standard 100% strength addition of the adaptive training delta.
+  - `< 1.0`: Softened application.
+  - `> 1.0`: Amplified application.
 
 ---
 
 ## Extract-Features
-> A powerful mode that identifies features present in both `(B - A)` and `(C - A)` and adds them to A. It uses per-vector cosine similarity to decide how much of each feature to keep, allowing for fine-grained control over combining aspects.
+> Analyzes feature vectors present in both $(B - A)$ and $(C - A)$, using cosine similarity to isolate shared features and blend them into Model A.
+
+**Models Used:** A (Base), B (Donor 1), C (Donor 2)
 
-**Models Used:** A, B, C
 **Parameters:**
-- **Alpha:** Weights the merge between Model B (`0.0`) and Model C (`1.0`).
-- **Beta:** Controls the focus on similarity (`0.0`) versus dissimilarity (`1.0`).
-- **Gamma:** A bias exponent for similarity calculation.
-- **Delta:** A final multiplier for the extracted features before they are added to Model A.
+- **Alpha (B vs C Contribution Balance):** Weights the source contribution between Model B and Model C.
+  - `0.0`: 100% Model B's delta features.
+  - `0.5`: Equal mix of features from both Model B and Model C.
+  - `1.0`: 100% Model C's delta features.
+- **Beta (Similarity vs Dissimilarity):**
+  - `0.0`: Focuses exclusively on features where B and C agree / are similar relative to A.
+  - `1.0`: Focuses on features where B and C diverge / are dissimiliar relative to A.
+- **Gamma (Similarity Bias Exponent):** Exponential bias curve shaping how sharply similar features are separated from dissimilar features.
+- **Delta (Feature Multiplier / Strength):** Final scaling factor for the extracted features before adding them onto Model A.
+  - `0.0`: No features added (pure Model A).
+  - `1.0`: Standard strength addition of extracted features onto Model A.
 
 ---
 
 ## Add-Dissimilarities
-> Identifies features that are dissimilar between Model B and Model C (relative to A) and adds them to Model A. Useful for combining unique aspects of two different models.
+> Identifies features that are dissimilar between Model B and Model C (relative to Model A) and transfers them into Model A. Ideal for combining non-overlapping, unique aspects from two different fine-tunes without conflicting duplicate features.
+
+**Models Used:** A (Base), B (Donor 1), C (Donor 2)
 
-**Models Used:** A, B, C
 **Parameters:**
-- **Alpha:** Weights the comparison between Model B (`0.0`) and Model C (`1.0`).
-- **Beta:** A multiplier for the final extracted dissimilar features.
-- **Gamma:** A bias exponent for the similarity calculation.
+- **Alpha (B vs C Balance):** Balance weighting between Model B and Model C when extracting differences.
+  - `0.0`: 100% weight to Model B's dissimilarities.
+  - `0.5`: Balanced extraction from both models.
+  - `1.0`: 100% weight to Model C's dissimilarities.
+- **Beta (Feature Multiplier / Strength):** Final scaling multiplier for the dissimilar features added to Model A.
+  - `0.0`: No features added (pure Model A).
+  - `1.0`: Standard strength addition onto Model A.
+- **Gamma (Similarity Bias Exponent):** Exponent controlling the cutoff sensitivity between similar and dissimilar directions.
 
 ---
 
-## Layer Mismatch Handling
+## Layer Mismatch Handling (`mismatch_mode`)
 
-When merging models with different layer structures, the `mismatch_mode` parameter controls behavior:
+When merging models with different layer topologies:
 
 | Mode | Behavior |
 |------|----------|
-| `skip` | Missing layers in B/C use A's values **(default)** |
-| `zeros` | Missing layers in B/C are treated as zeros |
-| `error` | Fail if any layer is missing |
+| `skip` | Missing layers in B or C use Model A's values unchanged **(default)** |
+| `zeros` | Missing layers in B or C are treated as zero tensors |
+| `error` | Raise an error and abort if any layer is missing across models |
 
-**Special case for B - C subtraction (Add-Difference, Train-Difference):**
-- If **both** B and C have the layer: Normal subtraction is performed
-- If **either** B or C is missing the layer: The difference is treated as `None` (no contribution from this layer), and Model A's value is used
-
-**Common scenarios:**
-- Merging a fine-tuned model with added/removed layers
-- Extracting style differences between models with slightly different architectures
-- Combining partial LoRAs or adapters
-
-**Note:** Extra layers in Models B/C that don't exist in Model A are currently ignored. The output will always have the same layer structure as Model A.
+**Special rule for $(B - C)$ difference operations (Add-Difference, Train-Difference):**
+- If **both** B and C have the layer: Normal difference $(B - C)$ is computed and applied to A.
+- If **either** B or C is missing the layer: The difference is treated as zero, and Model A's value is preserved untouched.
 
 ---
 
-## Layer Filtering (Regex Patterns)
+## Dtype Preservation & Override
 
-Use regex patterns to control which layers are merged:
+- **Preserve Higher Precision (default, `override_dtype = False`)**: If a key is `float32` in any source model (A, B, or C), it is saved as `float32` regardless of `save_dtype`.
+- **Explicit Override (`override_dtype = True`)**: Forces all floating tensors strictly to `save_dtype`.
 
-### `exclude_patterns`
-Layers matching any pattern will **keep Model A's values only** (no merge).
+---
 
-### `discard_patterns`
-Layers matching any pattern will be **removed entirely** from the output.
+## Layer Filtering
 
-**Pattern format:**
-- Whitespace-separated regex patterns
-- Patterns use **substring matching** (not full match)
-- Example: `text_model lora` matches any key containing "text_model" OR "lora"
+- **`exclude_patterns`**: Newline-separated patterns. Matching layers keep Model A's values untouched.
+- **`include_mode`**: Inverts `exclude_patterns` into a whitelist: only matching layers receive the merge; all other layers keep Model A's values.
+- **`discard_patterns`**: Newline-separated patterns. Matching layers are omitted completely from the output model.
+- **`glob_patterns`**: Enables shell wildcards (`*`, `?`) when checked; uses regex substring matching when unchecked.

docs/model_analysis.md

@@ -0,0 +1,89 @@
+# Two-Model Similarity Analysis
+
+These nodes compare two safetensors models without merging them or writing an
+output file. Separate nodes are provided for checkpoints, diffusion models,
+text encoders, LoRAs, and embeddings.
+
+`exclude_patterns` accepts one regex per line, or glob patterns when
+`glob_patterns` is enabled. By default, matches are excluded from metrics and
+topology counts. Enable `include_mode` to analyze only matches instead; an empty
+include filter selects no tensors. The same field and matching syntax are used
+in both modes.
+
+## Outputs
+
+`comparison_report` contains conventional numerical differences and topology
+findings. Both reports use Markdown headings and tables. `cwb_report` separately contains diagnostics derived from the
+similarity, alignment, and consensus stages of Consensus-Weighted Blending.
+`documentation` contains this reference.
+
+Two CSV-text outputs follow the existing three outputs, preserving their socket
+positions: `layerwise_metrics_csv` and `layerwise_cwb_csv`. Connect them to a
+text-saving node and use a `.csv` filename; analysis itself writes no files.
+The CSV rows follow the same key order and displayed numeric precision as the
+layerwise Markdown tables. Dtypes A/B, cosine mean/min/max, and A/B affinities
+have separate columns. Undefined values remain `N/A`; an empty analysis produces
+headers without data rows. Documentation-only mode returns empty CSV strings.
+
+The reports begin with a global summary, followed by inferred blockwise values
+and one entry for every comparable tensor. Block names are approximated from
+the tensor-key hierarchy through its first numeric path component; they are not
+an architecture-specific interpretation.
+
+## Standard metrics
+
+- MAE, MSE, and RMSE describe absolute error at the original weight scale.
+- Maximum absolute difference identifies the largest individual change.
+- Relative L2 is `||A-B|| / ((||A||+||B||)/2)` and is scale-aware.
+- Cosine similarity compares direction; Pearson correlation compares centered
+  variation.
+- Exact equality, sign agreement, norm ratio, and finite-value coverage expose
+  differences hidden by averages.
+
+The global section reports both parameter-weighted metrics and equal-layer
+averages. No arbitrary combined similarity score is assigned. Undefined
+zero-norm or constant-tensor statistics are shown as `N/A`.
+
+Only shared, same-shape floating tensors enter numerical aggregates. Missing
+keys and shape mismatches are listed separately and are never padded, cropped,
+or replaced with zeros. Non-floating tensors receive an exact-equality check.
+Quantized models carrying ComfyUI quantization metadata are rejected because
+encoded storage values do not represent directly comparable model weights.
+
+## CWB diagnostics
+
+The CWB report provides pairwise row-cosine statistics and each input's affinity
+to both mean and median consensus. It stops before CWB's merge-weighting stage:
+no thresholds, contribution weights, diversity weighting, vector
+reconstitution, norm rescaling, or output writer are used.
+
+Checkpoint, diffusion-model, and text-encoder tensors always retain their fixed
+coordinate alignment. Embedding and LoRA nodes expose one optional
+`cwb_similarity_alignment` control. When enabled, it performs
+threshold-independent greedy one-to-one matching derived from CWB and reports
+coverage, matched-score statistics, and improvement over index alignment. This
+operation is quadratic and may be slow for large embeddings.
+
+LoRA down/up tensors are treated as one logical layer. Rank-component matching
+uses the product of down and up cosine similarity, including CWB's paired-sign
+equivalence for comparison. No source tensor is modified.
+
+## Memory and failure handling
+
+Inputs always use Unified Efficient Loader's low-memory streaming mode. Matching
+A/B tensors are loaded concurrently through bounded `async_stream()` pipelines
+with `batch_size=1` and `prefetch_batches=1`, but no subsequent logical layer is
+requested early. A complete LoRA down/up pair is loaded once and used for both
+reports. CUDA inputs use pinned transfer while retaining the current CPU source
+tensors for fallback.
+
+Calculations use FP32. When CUDA processing runs out of memory, CUDA memory is
+cleared and the current logical layer is retried once on CPU. After success or
+failure, all tensor references for that layer are released and every loaded key
+is marked processed exactly once before advancing. Both reports list every CPU
+fallback. A failed CPU retry or any non-OOM error stops analysis with the layer
+name in the error.
+
+`top_weight_differences` bounds the scalar-coordinate list; every comparable
+tensor is still present in the layerwise report. `DOCUMENTATION ONLY` loads no
+model data.

nodes/adaptive_svd.py

@@ -0,0 +1,57 @@
+"""Bounded progressive partial SVD for adaptive extraction modes."""
+
+import torch
+
+
+ADAPTIVE_PARTIAL_MODES = {"ratio", "quantile", "sv_cumulative", "sv_fro"}
+
+
+def _bounded_rank(singular_values, mode, parameter, output_cap, matrix_fro_sq):
+    capped = singular_values[:output_cap]
+    if capped.numel() == 0 or capped[0] < 1e-8:
+        return 1, True
+
+    if mode == "ratio":
+        rank = int(torch.sum(capped > capped[0] / parameter).item()) if parameter > 0 else output_cap
+        return max(1, min(rank, output_cap)), rank < output_cap
+
+    if mode in {"quantile", "sv_cumulative"}:
+        total = capped.sum()
+        if total < 1e-8:
+            return 1, True
+        rank = int(torch.searchsorted(torch.cumsum(capped, 0) / total, parameter).item()) + 1
+        return max(1, min(rank, output_cap)), rank < output_cap
+
+    target_sq = float(parameter) ** 2 * matrix_fro_sq
+    cumulative = torch.cumsum(capped.square(), 0)
+    reached = bool(cumulative[-1] >= target_sq)
+    if reached:
+        rank = int(torch.searchsorted(cumulative, target_sq).item()) + 1
+        return max(1, min(rank, output_cap)), rank < output_cap
+    return output_cap, False
+
+
+def adaptive_partial_svd(weight, mode, parameter, max_rank, probe_offset, niter):
+    """Return a bounded top spectrum and selected output rank."""
+    max_possible = min(weight.shape)
+    output_cap = max(1, min(int(max_rank), max_possible))
+    offset = max(1, int(probe_offset))
+    initial_probe = min(max_possible, output_cap + offset)
+    probe_limit = min(max_possible, max(initial_probe, output_cap * 2))
+    matrix_fro_sq = float(torch.sum(weight.square()).item()) if mode == "sv_fro" else 0.0
+
+    probe_rank = initial_probe
+    while True:
+        if probe_rank >= max_possible:
+            U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
+        else:
+            U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter)
+            Vh = V.T
+
+        rank, resolved_before_cap = _bounded_rank(
+            S, mode, parameter, output_cap, matrix_fro_sq
+        )
+        near_probe_tail = rank >= max(1, probe_rank - offset)
+        if resolved_before_cap or not near_probe_tail or probe_rank >= probe_limit:
+            return U[:, :rank], S[:rank], Vh[:rank, :], rank
+        probe_rank = probe_limit

nodes/artifact_paths.py

@@ -0,0 +1,36 @@
+import os
+
+import folder_paths
+
+
+CANONICAL_MODEL_CATEGORIES = frozenset(
+    {"checkpoints", "diffusion_models", "text_encoders", "loras", "embeddings"}
+)
+
+
+def canonical_model_artifact_path(category: str, output_filename: str) -> tuple[str, str]:
+    """Return the canonical write path and category-relative loader value."""
+    if category not in CANONICAL_MODEL_CATEGORIES:
+        raise ValueError(f"Unsupported model artifact category: {category}")
+
+    relative_name = output_filename.strip().replace("\\", "/")
+    if not relative_name:
+        raise ValueError("Output filename must not be empty.")
+    if relative_name.startswith("/") or os.path.isabs(relative_name):
+        raise ValueError("Output filename must be relative to its model category.")
+    while relative_name.lower().endswith(".safetensors"):
+        relative_name = relative_name[:-12]
+    if not relative_name:
+        raise ValueError("Output filename must include a name before its extension.")
+    relative_name = f"{relative_name}.safetensors"
+
+    category_root = os.path.abspath(os.path.join(folder_paths.models_dir, category))
+    output_path = os.path.abspath(
+        os.path.join(category_root, *relative_name.split("/"))
+    )
+    if os.path.commonpath((category_root, output_path)) != category_root:
+        raise ValueError("Output filename must stay inside its model category directory.")
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    output_name = os.path.relpath(output_path, category_root).replace(os.sep, "/")
+    return output_path, output_name

nodes/cwb_delta_lora_merger.py

@@ -0,0 +1,768 @@
+import fnmatch
+import gc
+import re
+from dataclasses import dataclass, field
+
+import comfy.utils
+import folder_paths
+import torch
+from comfy_api.latest import io
+from tqdm import tqdm
+from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
+
+from .artifact_paths import canonical_model_artifact_path
+from .consensus_merger import (
+    CWB_CONFIG,
+    CWBSettings,
+    CWBDiagnostics,
+    DENSE_CWB_PRESETS,
+    merge_cwb_tensors,
+    resolve_cwb_settings,
+)
+from .device_utils import (
+    cleanup_after_operation,
+    estimate_model_size,
+    prepare_for_large_operation,
+)
+from .lora_alpha import normalize_lora_pair
+from .lora_resize import (
+    build_lora_layer_map,
+    canonical_lora_key,
+    parse_lora_layers,
+    select_output_dtype,
+    validate_canonical_blocks,
+)
+from .quantization_guard import inspect_low_bit_input
+from .uel_io import AsyncTensorCursor, atomic_uel_writer
+
+
+_OPERATION = "Delta CWB LoRA Merge"
+_ALLOWED_ROLES = {
+    "down", "up", "alpha", "diff", "format", "down_suffix", "up_suffix",
+    "alpha_suffix",
+}
+_FACTOR_SUFFIXES = (
+    ".lora_A.default.weight", ".lora_B.default.weight",
+    ".lora_down.weight", ".lora_up.weight",
+    "_lora.down.weight", "_lora.up.weight",
+    ".lora_A.weight", ".lora_B.weight",
+    ".lora.down.weight", ".lora.up.weight",
+    ".lora_A", ".lora_B",
+    ".lora_linear_layer.down.weight", ".lora_linear_layer.up.weight",
+)
+_UNSUPPORTED_SUFFIXES = (
+    ".lora_mid.weight", ".reshape_weight", ".dora_scale", ".w_norm",
+    ".b_norm", ".set_weight",
+)
+
+
+@dataclass(frozen=True)
+class DeltaCWBLoRASource:
+    source_index: int
+    block_name: str
+    down_key: str | None
+    up_key: str | None
+    alpha_key: str | None
+    diff_key: str | None
+    source_dtype: torch.dtype
+
+    @property
+    def keys(self):
+        return tuple(
+            key for key in (self.down_key, self.up_key, self.alpha_key, self.diff_key)
+            if key is not None
+        )
+
+
+@dataclass(frozen=True)
+class DeltaCWBLayerUnit:
+    core: str
+    output_key: str
+    sources: tuple[DeltaCWBLoRASource, ...]
+    zero_contributors: int
+    preserve_anchor: bool
+
+
+@dataclass
+class DeltaCWBReport:
+    merged_layers: int = 0
+    anchor_preserved_layers: int = 0
+    zero_filled_contributors: int = 0
+    rejected_unsupported_layers: int = 0
+    cuda_cpu_fallbacks: int = 0
+    output_dtypes: dict[str, int] = field(default_factory=dict)
+
+    def record_dtype(self, dtype):
+        name = str(dtype).removeprefix("torch.")
+        self.output_dtypes[name] = self.output_dtypes.get(name, 0) + 1
+
+    def render(self, diagnostics: CWBDiagnostics, preset_name: str, input_count: int, custom: bool):
+        dtype_summary = ", ".join(
+            f"{name}={count}" for name, count in sorted(self.output_dtypes.items())
+        ) or "none"
+        source = "connected custom configuration" if custom else preset_name
+        return "\n".join([
+            "DELTA CWB LORA MERGE SUMMARY",
+            f"Settings: {source}",
+            f"Input contributors: {input_count}",
+            f"Merged layers: {self.merged_layers}",
+            f"Anchor-preserved layers: {self.anchor_preserved_layers}",
+            f"Explicit zero contributors: {self.zero_filled_contributors}",
+            f"Rejected/unsupported layers: {self.rejected_unsupported_layers}",
+            f"Accepted weighting contributors: {diagnostics.accepted_weighting_contributors}/{diagnostics.weighting_contributors}",
+            f"All-rejected equal fallbacks: {diagnostics.all_rejected_fallbacks}",
+            f"Zero-weight equal fallbacks: {diagnostics.zero_weight_fallbacks}",
+            f"CUDA OOM CPU fallbacks: {self.cuda_cpu_fallbacks}",
+            f"Output dtypes: {dtype_summary}",
+        ])
+
+
+def _compile_patterns(patterns, glob_patterns):
+    entries = [entry.strip() for entry in patterns.splitlines() if entry.strip()]
+    if glob_patterns:
+        return tuple(entries)
+    return tuple(re.compile(entry) for entry in entries)
+
+
+def _matches_pattern(key, patterns, glob_patterns):
+    if glob_patterns:
+        return any(fnmatch.fnmatchcase(key, pattern) for pattern in patterns)
+    return any(pattern.search(key) is not None for pattern in patterns)
+
+
+def _factor_delta(down, up, layer):
+    if down.ndim == 2 and up.ndim == 2:
+        if down.shape[0] != up.shape[1]:
+            raise ValueError(f"Delta CWB rank mismatch for '{layer}'.")
+        return up.float() @ down.float()
+    if down.ndim == 4 and up.ndim == 4 and tuple(up.shape[2:]) == (1, 1):
+        rank, channels, height, width = down.shape
+        if up.shape[1] != rank:
+            raise ValueError(f"Delta CWB rank mismatch for '{layer}'.")
+        delta = up.reshape(up.shape[0], rank).float() @ down.reshape(rank, -1).float()
+        return delta.reshape(up.shape[0], channels, height, width)
+    raise ValueError(
+        f"Delta CWB layer '{layer}' requires linear factors or convolutional "
+        "factors with a 1x1 up kernel."
+    )
+
+
+def _factorize_delta(delta, rank):
+    if delta.ndim not in (2, 4):
+        return None
+    shape = delta.shape
+    matrix = delta if delta.ndim == 2 else delta.reshape(shape[0], -1)
+    target_rank = max(1, min(int(rank), min(matrix.shape)))
+    if target_rank == min(matrix.shape):
+        if matrix.shape[0] <= matrix.shape[1]:
+            up = torch.eye(matrix.shape[0], device=matrix.device, dtype=matrix.dtype)
+            down = matrix
+        else:
+            up = matrix
+            down = torch.eye(matrix.shape[1], device=matrix.device, dtype=matrix.dtype)
+    else:
+        up, singular_values, right = torch.svd_lowrank(matrix, q=target_rank, niter=2)
+        up = up * singular_values.unsqueeze(0)
+        down = right.t()
+    if delta.ndim == 4:
+        down = down.reshape(down.shape[0], *shape[1:])
+        up = up.unsqueeze(-1).unsqueeze(-1)
+    return down, up
+
+
+def _inspect_inputs(handlers):
+    infos = []
+    for index, handler in enumerate(handlers):
+        low_bit = inspect_low_bit_input(handler, f"LoRA {index + 1}", _OPERATION)
+        keys = list(handler.keys())
+        pairs, passthrough = parse_lora_layers(keys)
+        unsupported_passthrough = [
+            key for key in passthrough if key.endswith(_UNSUPPORTED_SUFFIXES)
+        ]
+        if unsupported_passthrough:
+            raise ValueError(
+                f"[{_OPERATION}] LoRA {index + 1} contains unsupported companion "
+                f"tensor '{unsupported_passthrough[0]}'."
+            )
+        incomplete = [key for key in passthrough if key.endswith(_FACTOR_SUFFIXES)]
+        if incomplete:
+            raise ValueError(
+                f"[{_OPERATION}] LoRA {index + 1} contains an incomplete factor "
+                f"pair at '{incomplete[0]}'."
+            )
+        infos.append({"pairs": pairs, "passthrough": passthrough, "low_bit": low_bit})
+    return infos
+
+
+def _build_layer_units(
+    layer_map,
+    infos,
+    handlers,
+    mismatch_mode,
+    exclude_patterns,
+    discard_patterns,
+    glob_patterns,
+    include_mode=False,
+):
+    units = []
+    stream_keys = [[] for _ in handlers]
+    rejected = 0
+    for core, mapped in layer_map.items():
+        if not any(index == 0 for index, _ in mapped):
+            continue
+        output_key = canonical_lora_key(core, "diff")
+        if _matches_pattern(output_key, discard_patterns, glob_patterns):
+            rejected += 1
+            continue
+        matched = _matches_pattern(output_key, exclude_patterns, glob_patterns)
+        preserve = not matched if include_mode else matched
+        candidates = []
+        shapes = {}
+        for index, block in sorted(mapped):
+            roles = infos[index]["pairs"][block]
+            unsupported = set(roles) - _ALLOWED_ROLES
+            if unsupported:
+                raise ValueError(
+                    f"[{_OPERATION}] Layer '{block}' has unsupported roles: "
+                    f"{', '.join(sorted(unsupported))}."
+                )
+            if "down" in roles or "up" in roles:
+                if not all(role in roles for role in ("down", "up")):
+                    raise ValueError(f"[{_OPERATION}] Layer '{block}' has an incomplete factor pair.")
+                down_key, up_key = roles["down"], roles["up"]
+                if down_key in infos[index]["low_bit"] or up_key in infos[index]["low_bit"]:
+                    raise ValueError(f"[{_OPERATION}] Cannot expand low-bit factors for '{block}'.")
+                down_shape = tuple(handlers[index].get_shape(down_key))
+                up_shape = tuple(handlers[index].get_shape(up_key))
+                if len(down_shape) == 2 and len(up_shape) == 2:
+                    if down_shape[0] != up_shape[1]:
+                        raise ValueError(f"[{_OPERATION}] Rank mismatch for '{block}'.")
+                    dense_shape = (up_shape[0], down_shape[1])
+                elif len(down_shape) == 4 and len(up_shape) == 4 and up_shape[2:] == (1, 1):
+                    if down_shape[0] != up_shape[1]:
+                        raise ValueError(f"[{_OPERATION}] Rank mismatch for '{block}'.")
+                    dense_shape = (up_shape[0], *down_shape[1:])
+                else:
+                    raise ValueError(
+                        f"[{_OPERATION}] Layer '{block}' requires linear factors or "
+                        "convolutional factors with a 1x1 up kernel."
+                    )
+                factor_dtypes = (
+                    handlers[index].get_dtype(down_key),
+                    handlers[index].get_dtype(up_key),
+                )
+                source_dtype = torch.float32 if torch.float32 in factor_dtypes else factor_dtypes[1]
+                source = DeltaCWBLoRASource(
+                    index, block, down_key, up_key, roles.get("alpha"), None, source_dtype
+                )
+            elif "diff" in roles:
+                diff_key = roles["diff"]
+                if diff_key in infos[index]["low_bit"]:
+                    raise ValueError(f"[{_OPERATION}] Cannot merge low-bit delta '{block}'.")
+                dense_shape = tuple(handlers[index].get_shape(diff_key))
+                source = DeltaCWBLoRASource(
+                    index, block, None, None, None, diff_key,
+                    handlers[index].get_dtype(diff_key),
+                )
+            else:
+                raise ValueError(f"[{_OPERATION}] Layer '{block}' has no supported representation.")
+            candidates.append(source)
+            shapes[index] = dense_shape
+
+        anchor = next((source for source in candidates if source.source_index == 0), None)
+        if anchor is None:
+            continue
+        anchor_shape = shapes[0]
+        compatible = [
+            source for source in candidates if shapes[source.source_index] == anchor_shape
+        ]
+        present = {source.source_index for source in compatible}
+        mismatch_count = len(handlers) - len(present)
+        if preserve:
+            sources = (anchor,)
+            zero_count = 0
+        elif mismatch_count and mismatch_mode == "error":
+            raise ValueError(
+                f"[{_OPERATION}] Layer '{core}' is missing or incompatible in one or more inputs."
+            )
+        elif mismatch_count and mismatch_mode == "skip":
+            sources = (anchor,)
+            zero_count = 0
+            preserve = True
+        else:
+            sources = tuple(compatible)
+            zero_count = mismatch_count if mismatch_mode == "zeros" else 0
+
+        for source in sources:
+            stream_keys[source.source_index].extend(source.keys)
+        units.append(DeltaCWBLayerUnit(core, output_key, sources, zero_count, preserve))
+    return units, stream_keys, rejected
+
+
+def _to_processing_device(tensor, device):
+    if device.type == "cuda":
+        return transfer_to_gpu_pinned(tensor, device, torch.float32)
+    return tensor.to(device=device, dtype=torch.float32)
+
+
+def _process_layer_on_device(
+    unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics,
+    output_representation, output_rank,
+):
+    deltas = []
+    try:
+        for source in unit.sources:
+            if source.diff_key is not None:
+                delta = _to_processing_device(
+                    loaded[(source.source_index, source.diff_key)], device
+                )
+            else:
+                down = loaded[(source.source_index, source.down_key)]
+                up = loaded[(source.source_index, source.up_key)]
+                alpha = (
+                    loaded[(source.source_index, source.alpha_key)]
+                    if source.alpha_key is not None else None
+                )
+                down, up = normalize_lora_pair(down, up, alpha, layer=source.block_name)
+                delta = _factor_delta(
+                    _to_processing_device(down, device),
+                    _to_processing_device(up, device),
+                    source.block_name,
+                )
+            deltas.append(delta)
+        if unit.zero_contributors:
+            deltas.extend(torch.zeros_like(deltas[0]) for _ in range(unit.zero_contributors))
+        if unit.preserve_anchor:
+            merged = deltas[0].clone()
+        else:
+            merged = merge_cwb_tensors(
+                deltas,
+                settings,
+                reference_index=0,
+                allow_similarity_alignment=False,
+                diagnostics=diagnostics,
+            )
+        dtype = select_output_dtype(
+            [source.source_dtype for source in unit.sources],
+            save_dtype,
+            is_1d_diff=merged.ndim == 1,
+        )
+        if output_representation == "factorized_lora":
+            factors = _factorize_delta(merged, output_rank)
+            if factors is not None:
+                down, up = factors
+                return [
+                    (canonical_lora_key(unit.core, "down"), down.to(dtype=dtype, device="cpu").contiguous()),
+                    (canonical_lora_key(unit.core, "up"), up.to(dtype=dtype, device="cpu").contiguous()),
+                ], dtype
+        return [(unit.output_key, merged.to(dtype=dtype, device="cpu").contiguous())], dtype
+    finally:
+        deltas.clear()
+
+
+def _is_cuda_oom(error):
+    return isinstance(error, torch.cuda.OutOfMemoryError) or (
+        isinstance(error, RuntimeError) and "out of memory" in str(error).lower()
+    )
+
+
+def _release_cuda_oom():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _process_layer(
+    unit, loaded, settings, save_dtype, requested_device, diagnostics,
+    output_representation="full_difference", output_rank=384,
+):
+    device = torch.device(requested_device)
+    attempt_diagnostics = CWBDiagnostics()
+    try:
+        output, dtype = _process_layer_on_device(
+            unit, loaded, settings, save_dtype, device, attempt_diagnostics,
+            output_representation, output_rank,
+        )
+        diagnostics.absorb_runtime(attempt_diagnostics)
+        return output, dtype, False
+    except Exception as error:
+        if device.type != "cuda" or not _is_cuda_oom(error):
+            raise
+        _release_cuda_oom()
+        retry_diagnostics = CWBDiagnostics()
+        output, dtype = _process_layer_on_device(
+            unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics,
+            output_representation, output_rank,
+        )
+        diagnostics.absorb_runtime(retry_diagnostics)
+        return output, dtype, True
+
+
+class DeltaCWBLoRAMergerLogic:
+    @staticmethod
+    def execute(
+        lora_paths,
+        source_names,
+        settings,
+        preset_name,
+        custom_settings,
+        mismatch_mode,
+        output_filename,
+        save_dtype,
+        output_representation,
+        output_rank,
+        process_device,
+        exclude_patterns,
+        discard_patterns,
+        glob_patterns,
+        force_clear_cache,
+        include_mode=False,
+    ):
+        prepare_for_large_operation(
+            sum(estimate_model_size(path) for path in lora_paths) * 2.5,
+            torch.device(process_device),
+        )
+        handlers = []
+        cursors = []
+        report = DeltaCWBReport()
+        diagnostics = CWBDiagnostics()
+        try:
+            for path in lora_paths:
+                handlers.append(MemoryEfficientSafeOpen(path, low_memory=True))
+            infos = _inspect_inputs(handlers)
+            layer_map = build_lora_layer_map(infos, None)
+            validate_canonical_blocks(layer_map, _OPERATION)
+            compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
+            compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
+            units, stream_keys, rejected = _build_layer_units(
+                layer_map,
+                infos,
+                handlers,
+                mismatch_mode,
+                compiled_excludes,
+                compiled_discards,
+                glob_patterns,
+                include_mode,
+            )
+            report.rejected_unsupported_layers = rejected
+            cursors = [
+                AsyncTensorCursor(
+                    handler,
+                    keys,
+                    pin_memory=torch.device(process_device).type == "cuda",
+                )
+                for handler, keys in zip(handlers, stream_keys)
+            ]
+            output_path, output_name = canonical_model_artifact_path("loras", output_filename)
+            metadata = (handlers[0].metadata() or {}).copy()
+            metadata.update({
+                "merge_method": "delta_cwb",
+                "source_models": ",".join(source_names),
+                "alpha_normalized": "true",
+                "output_representation": (
+                    "factorized_lora" if output_representation == "factorized_lora"
+                    else "full_direct_difference"
+                ),
+                "cwb_consensus_type": settings.consensus_type,
+                "cwb_alignment_method": settings.alignment_method,
+                "cwb_alignment_threshold": str(settings.alignment_threshold),
+                "cwb_similarity_threshold": str(settings.similarity_threshold),
+                "cwb_power_alpha": str(settings.power_alpha),
+                "cwb_diversity_beta": str(settings.diversity_beta),
+                "cwb_rescale_norm": str(settings.rescale_norm).lower(),
+                "cwb_global_scale": str(settings.global_scale),
+                "cwb_dynamic_similarity_contrast": str(settings.dynamic_similarity_contrast).lower(),
+                "cwb_soft_comfort_bandpass": str(settings.soft_comfort_bandpass).lower(),
+                "cwb_position_weight": str(settings.position_weight),
+                "cwb_preserve_common_prefix": str(settings.preserve_common_prefix).lower(),
+            })
+            pbar = comfy.utils.ProgressBar(len(units))
+            with atomic_uel_writer(output_path, metadata) as writer, torch.no_grad():
+                for unit in tqdm(units, desc="Delta CWB LoRA layers", unit="layers"):
+                    if force_clear_cache:
+                        _release_cuda_oom()
+                    loaded = {}
+                    try:
+                        for source in unit.sources:
+                            for key in source.keys:
+                                loaded[(source.source_index, key)] = cursors[
+                                    source.source_index
+                                ].take(key)
+                        output_items, dtype, used_cpu = _process_layer(
+                            unit,
+                            loaded,
+                            settings,
+                            save_dtype,
+                            process_device,
+                            diagnostics,
+                            output_representation,
+                            output_rank,
+                        )
+                        writer.write_batch(output_items)
+                        report.merged_layers += int(not unit.preserve_anchor)
+                        report.anchor_preserved_layers += int(unit.preserve_anchor)
+                        report.zero_filled_contributors += unit.zero_contributors
+                        report.cuda_cpu_fallbacks += int(used_cpu)
+                        report.record_dtype(dtype)
+                        del output_items
+                    finally:
+                        for source_index, key in list(loaded):
+                            del loaded[(source_index, key)]
+                            cursors[source_index].release(key)
+                        loaded.clear()
+                    pbar.update(1)
+            for cursor in cursors:
+                cursor.finish()
+            return output_name, report.render(
+                diagnostics, preset_name, len(lora_paths), custom_settings
+            )
+        finally:
+            for cursor in cursors:
+                cursor.close()
+            for handler in handlers:
+                handler.__exit__(None, None, None)
+            cleanup_after_operation()
+
+    @staticmethod
+    def execute_node(kwargs, count):
+        names = [kwargs[f"lora_{index}"] for index in range(1, count + 1)]
+        paths = [folder_paths.get_full_path_or_raise("loras", name) for name in names]
+        settings = resolve_cwb_settings(
+            kwargs["cwb_preset"], DENSE_CWB_PRESETS, kwargs.get("cwb_config")
+        )
+        requested_dtype = {
+            "fp32": torch.float32,
+            "fp16": torch.float16,
+            "bf16": torch.bfloat16,
+        }[kwargs["save_dtype"]]
+        filename, report = DeltaCWBLoRAMergerLogic.execute(
+            paths,
+            names,
+            settings,
+            kwargs["cwb_preset"],
+            kwargs.get("cwb_config") is not None,
+            kwargs["mismatch_mode"],
+            kwargs["output_filename"],
+            requested_dtype,
+            kwargs["output_representation"],
+            kwargs["output_rank"],
+            kwargs["process_device"],
+            kwargs["exclude_patterns"],
+            kwargs["discard_patterns"],
+            kwargs["glob_patterns"],
+            kwargs["force_clear_cache"],
+            kwargs.get("include_mode", False),
+        )
+        return io.NodeOutput(filename, report)
+
+
+def _fixed_inputs(count, default_filename):
+    options = folder_paths.get_filename_list("loras")
+    inputs = []
+    if count:
+        inputs.append(io.Combo.Input(
+            "lora_1",
+            options=options,
+            tooltip="Metadata and layer anchor. Defines saved layers and preserved values.",
+        ))
+    inputs.extend(
+        io.Combo.Input(
+            f"lora_{index}",
+            options=options,
+            tooltip="Equal-prior LoRA contributor. CWB determines its effective row influence.",
+        )
+        for index in range(2, count + 1)
+    )
+    inputs.extend([
+        io.Combo.Input(
+            "cwb_preset",
+            options=list(DENSE_CWB_PRESETS),
+            default="balanced_mean",
+            tooltip="Controls how strongly agreement, disagreement, and row magnitude affect the saved full weight changes. A connected CWB Config overrides it completely.",
+        ),
+        CWB_CONFIG.Input(
+            "cwb_config",
+            optional=True,
+            tooltip="Optional complete settings override from CWB Custom Configuration.",
+        ),
+        io.Combo.Input(
+            "mismatch_mode",
+            options=["skip", "zeros", "error"],
+            default="skip",
+            tooltip="For a missing or incompatible anchored layer: preserve LoRA 1, include an explicit zero contributor, or abort.",
+        ),
+        io.String.Input(
+            "output_filename",
+            default=default_filename,
+            tooltip="Relative filename for the atomically saved full-difference LoRA.",
+        ),
+        io.Combo.Input(
+            "save_dtype",
+            options=["fp32", "fp16", "bf16"],
+            default="bf16",
+            tooltip="Requested saved precision. Participating FP32 inputs preserve FP32.",
+        ),
+        io.Combo.Input(
+            "output_representation",
+            options=["full_difference", "factorized_lora"],
+            default="full_difference",
+            tooltip="Save exact full differences or factorize each merged layer directly into compact LoRA factors.",
+        ),
+        io.Int.Input(
+            "output_rank",
+            default=384,
+            min=1,
+            max=16384,
+            tooltip="Maximum rank used when Output Representation is factorized_lora.",
+        ),
+        io.Combo.Input(
+            "process_device",
+            options=["cuda", "cpu"],
+            default="cuda",
+            tooltip="Processing device. Only a failed CUDA layer retries on CPU.",
+        ),
+        io.String.Input(
+            "exclude_patterns",
+            default="",
+            multiline=True,
+            tooltip="Preserve matching canonical full-difference layers from LoRA 1.",
+        ),
+        io.String.Input(
+            "discard_patterns",
+            default="",
+            multiline=True,
+            tooltip="Omit matching canonical full-difference layers from the output.",
+        ),
+        io.Boolean.Input(
+            "glob_patterns",
+            default=False,
+            tooltip="Use shell-style glob patterns instead of regular expressions.",
+        ),
+        io.Boolean.Input(
+            "force_clear_cache",
+            default=False,
+            tooltip="Clear Python and CUDA caches before each layer to reduce retained memory.",
+        ),
+        io.Boolean.Input(
+            "include_mode",
+            default=False,
+            tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from LoRA 1.",
+        ),
+    ])
+    return inputs
+
+
+def _multi_inputs(default_filename):
+    options = folder_paths.get_filename_list("loras")
+    optional = ["None", *options]
+    inputs = [
+        io.Combo.Input(
+            "lora_count",
+            options=[str(index) for index in range(2, 9)],
+            default="2",
+            tooltip="Number of consecutive selected LoRAs to merge.",
+        ),
+        io.Combo.Input(
+            "lora_1",
+            options=options,
+            tooltip="Metadata and layer anchor. Defines saved layers and preserved values.",
+        ),
+        io.Combo.Input(
+            "lora_2",
+            options=options,
+            tooltip="Second equal-prior LoRA contributor.",
+        ),
+    ]
+    inputs.extend(
+        io.Combo.Input(
+            f"lora_{index}",
+            options=optional,
+            default="None",
+            tooltip="Additional equal-prior contributor; required when included by LoRA Count.",
+        )
+        for index in range(3, 9)
+    )
+    inputs.extend(_fixed_inputs(0, default_filename))
+    return inputs
+
+
+class DeltaCWBLoRATwoMerger(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="DeltaCWBLoRATwoMerger",
+            display_name="Delta CWB Merge LoRAs (2)",
+            category="ModelUtils/LoRA/Merge/Delta CWB",
+            description="Expand two LoRAs to full weight changes and merge corresponding rows with CWB.",
+            inputs=_fixed_inputs(2, "delta_cwb_merged_2_lora"),
+            outputs=[
+                io.AnyType.Output(display_name="output_filename"),
+                io.String.Output(display_name="cwb_report"),
+            ],
+            is_output_node=True,
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, **kwargs):
+        return DeltaCWBLoRAMergerLogic.execute_node(kwargs, 2)
+
+
+class DeltaCWBLoRAThreeMerger(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="DeltaCWBLoRAThreeMerger",
+            display_name="Delta CWB Merge LoRAs (3)",
+            category="ModelUtils/LoRA/Merge/Delta CWB",
+            description="Expand three LoRAs to full weight changes and merge corresponding rows with CWB.",
+            inputs=_fixed_inputs(3, "delta_cwb_merged_3_lora"),
+            outputs=[
+                io.AnyType.Output(display_name="output_filename"),
+                io.String.Output(display_name="cwb_report"),
+            ],
+            is_output_node=True,
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, **kwargs):
+        return DeltaCWBLoRAMergerLogic.execute_node(kwargs, 3)
+
+
+class DeltaCWBLoRAMultiMerger(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="DeltaCWBLoRAMultiMerger",
+            display_name="Delta CWB LoRA Multi-Merge",
+            category="ModelUtils/LoRA/Merge/Delta CWB",
+            description="Expand 2-8 LoRAs to full weight changes and merge corresponding rows with CWB.",
+            inputs=_multi_inputs("delta_cwb_merged_multi_lora"),
+            outputs=[
+                io.AnyType.Output(display_name="output_filename"),
+                io.String.Output(display_name="cwb_report"),
+            ],
+            is_output_node=True,
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, **kwargs):
+        count = int(kwargs["lora_count"])
+        missing = [
+            index for index in range(1, count + 1)
+            if not kwargs[f"lora_{index}"] or kwargs[f"lora_{index}"] == "None"
+        ]
+        if missing:
+            raise ValueError(
+                f"LoRA Count includes unselected input(s): {', '.join(map(str, missing))}."
+            )
+        return DeltaCWBLoRAMergerLogic.execute_node(kwargs, count)
+
+
+DELTA_CWB_LORA_NODES = [
+    DeltaCWBLoRATwoMerger,
+    DeltaCWBLoRAThreeMerger,
+    DeltaCWBLoRAMultiMerger,
+]

nodes/cwb_model_graft.py

@@ -0,0 +1,790 @@
+"""Standalone training-free Consensus-Weighted Blending model grafting node."""
+
+from __future__ import annotations
+
+from contextlib import closing
+import json
+import re
+
+import comfy.utils
+import folder_paths
+import torch
+import torch.nn.functional as F
+from comfy_api.latest import io
+from tqdm import tqdm
+from unifiedefficientloader import MemoryEfficientSafeOpen
+
+from .artifact_paths import canonical_model_artifact_path
+from .consensus_merger import (
+    CWB_CONFIG,
+    CWBDiagnostics,
+    CWBSettings,
+    DENSE_CWB_PRESETS,
+    _clear_previous_layer,
+    _is_float_dtype,
+    _merge_tensors_to_cpu,
+    _requested_dtype,
+    resolve_cwb_settings,
+)
+from .device_utils import (
+    cleanup_after_operation,
+    estimate_model_size,
+    prepare_for_large_operation,
+)
+from .merger import (
+    _compile_patterns,
+    _matches_any_pattern,
+    load_documentation_from_file,
+)
+from .quantization_guard import (
+    DiffusionQuantization,
+    diffusion_key_map,
+    write_preserved_tensor,
+)
+from .uel_io import atomic_uel_writer, stream_work_units
+
+
+_MAIN_BLOCK_INDEX_RE = re.compile(
+    r"(?:^|\.)(?:layers|blocks|double_blocks|single_blocks|transformer_blocks)\.(\d+)\."
+)
+_ATTN_KEY_TOKENS = (
+    ".qkv.",
+    ".qkv_proj.",
+    ".proj.q.",
+    ".proj.k.",
+    ".proj.v.",
+    ".proj.o.",
+    ".attn.",
+    ".self_attn.",
+    ".q_norm.",
+    ".k_norm.",
+    ".to_q.",
+    ".to_k.",
+    ".to_v.",
+    ".to_out.",
+)
+_FFN_KEY_TOKENS = (
+    ".ffn_up_gate.",
+    ".ffn_up.",
+    ".ffn_gate.",
+    ".ffn_down.",
+    ".mlp.",
+    ".ff.",
+    ".fc1.",
+    ".fc2.",
+    ".w1.",
+    ".w2.",
+    ".w3.",
+)
+
+
+def _extract_main_block_index(key: str) -> int | None:
+    if "text_embedder." in key or "token_refiner." in key:
+        return None
+    match = _MAIN_BLOCK_INDEX_RE.search(key)
+    return int(match.group(1)) if match else None
+
+
+def _classify_graft_operator(key: str) -> str | None:
+    padded = f".{key}."
+    if any(token in padded for token in _ATTN_KEY_TOKENS):
+        return "attention"
+    if any(token in padded for token in _FFN_KEY_TOKENS):
+        return "ffn"
+    return None
+
+
+def _block_rule_allows(
+    block_idx: int | None,
+    rule: str,
+    protect_early_blocks: int,
+    max_block_index: int,
+) -> bool:
+    if rule == "none":
+        return False
+    if block_idx is None:
+        return False
+    if block_idx < protect_early_blocks:
+        return False
+    if rule == "all_after_protected":
+        return True
+    if rule == "interleaved_even":
+        return (block_idx % 2) == 0
+    if rule == "interleaved_odd":
+        return (block_idx % 2) == 1
+    if rule == "back_half_only":
+        midpoint = (max_block_index + 1) // 2
+        return block_idx >= max(protect_early_blocks, midpoint)
+    return False
+
+
+def _parse_custom_key_map(raw_map: str) -> dict[str, str]:
+    """Parse lines of 'donor_key -> base_key' or 'donor_key = base_key' into {base_key: donor_key}."""
+    mapping: dict[str, str] = {}
+    if not raw_map or not raw_map.strip():
+        return mapping
+    for line_num, raw_line in enumerate(raw_map.splitlines(), start=1):
+        line = raw_line.strip()
+        if not line or line.startswith("#"):
+            continue
+        if "->" in line:
+            donor_part, base_part = line.split("->", 1)
+        elif "=" in line:
+            donor_part, base_part = line.split("=", 1)
+        else:
+            raise ValueError(
+                f"Custom key map line {line_num}: expected 'donor_key -> base_key'."
+            )
+        donor_key = donor_part.strip()
+        base_key = base_part.strip()
+        if not donor_key or not base_key:
+            raise ValueError(
+                f"Custom key map line {line_num}: both donor_key and base_key must be non-empty."
+            )
+        mapping[base_key] = donor_key
+    return mapping
+
+
+def _align_donor_trailing_shape(
+    base_tensor: torch.Tensor, donor_tensor: torch.Tensor
+) -> tuple[torch.Tensor, bool]:
+    """Align donor trailing dimensions to base_tensor so CWB similarity alignment can match rows."""
+    if donor_tensor.ndim != base_tensor.ndim:
+        if donor_tensor.numel() == base_tensor.numel():
+            return donor_tensor.reshape(base_tensor.shape), True
+        flat_donor = donor_tensor.reshape(1, 1, -1).to(dtype=torch.float32)
+        resized = F.interpolate(
+            flat_donor, size=base_tensor.numel(), mode="linear", align_corners=False
+        )
+        return resized.reshape(base_tensor.shape).to(dtype=donor_tensor.dtype), True
+
+    if donor_tensor.shape == base_tensor.shape:
+        return donor_tensor, False
+
+    if base_tensor.ndim <= 1:
+        return donor_tensor, True
+
+    if donor_tensor.shape[1:] == base_tensor.shape[1:]:
+        return donor_tensor, True
+
+    target_trailing = int(torch.tensor(base_tensor.shape[1:]).prod().item())
+    donor_2d = donor_tensor.reshape(donor_tensor.shape[0], 1, -1).to(dtype=torch.float32)
+    resized_2d = F.interpolate(
+        donor_2d, size=target_trailing, mode="linear", align_corners=False
+    )
+    reshaped = resized_2d.reshape(donor_tensor.shape[0], *base_tensor.shape[1:]).to(
+        dtype=donor_tensor.dtype
+    )
+    return reshaped, True
+
+
+def _graft_pair_to_cpu(
+    base_tensor: torch.Tensor,
+    donor_tensor: torch.Tensor,
+    settings: CWBSettings,
+    device: str,
+    target_dtype: torch.dtype,
+    *,
+    operation_label: str,
+    diagnostics: CWBDiagnostics,
+) -> torch.Tensor:
+    aligned_donor, shape_changed = _align_donor_trailing_shape(base_tensor, donor_tensor)
+    merged = _merge_tensors_to_cpu(
+        [base_tensor, aligned_donor],
+        settings,
+        device,
+        torch.float32,
+        reference_index=0,
+        allow_similarity_alignment=True,
+        operation_label=operation_label,
+        diagnostics=diagnostics,
+    )
+    if shape_changed or settings.rescale_norm:
+        base_norm = torch.linalg.vector_norm(base_tensor.to(dtype=torch.float32))
+        merged_norm = torch.linalg.vector_norm(merged)
+        if merged_norm.item() > 1e-8 and base_norm.item() > 0.0:
+            merged = merged * (base_norm / merged_norm)
+    return merged.to(dtype=target_dtype).cpu().contiguous()
+
+
+def _graft_with_subpart(
+    key: str,
+    op_type: str,
+    subpart: str,
+    base_tensor: torch.Tensor,
+    donor_tensor: torch.Tensor,
+    settings: CWBSettings,
+    device: str,
+    target_dtype: torch.dtype,
+    diagnostics: CWBDiagnostics,
+) -> torch.Tensor | None:
+    """Return grafted tensor or None if the subpart rule keeps the base tensor untouched."""
+    padded = f".{key}."
+    if op_type == "attention":
+        is_qkv = any(tok in padded for tok in (".qkv.", ".qkv_proj."))
+        is_q_or_k = any(
+            tok in padded for tok in (".proj.q.", ".proj.k.", ".to_q.", ".to_k.", ".q_norm.", ".k_norm.")
+        )
+        is_v = any(tok in padded for tok in (".proj.v.", ".to_v."))
+        is_out = any(tok in padded for tok in (".proj.o.", ".out_proj.", ".to_out."))
+        has_known_attn_part = is_qkv or is_q_or_k or is_v or is_out
+
+        if has_known_attn_part and subpart == "output_only":
+            if not is_out:
+                return None
+        elif has_known_attn_part and subpart == "value_and_output":
+            if is_q_or_k:
+                return None
+            if is_qkv and base_tensor.ndim >= 1 and base_tensor.shape[0] % 3 == 0:
+                aligned_donor, _ = _align_donor_trailing_shape(base_tensor, donor_tensor)
+                if aligned_donor.shape[0] % 3 == 0:
+                    q_a, k_a, v_a = base_tensor.chunk(3, dim=0)
+                    _, _, v_b = aligned_donor.chunk(3, dim=0)
+                    v_merged = _graft_pair_to_cpu(
+                        v_a,
+                        v_b,
+                        settings,
+                        device,
+                        target_dtype,
+                        operation_label=f"{key}[v]",
+                        diagnostics=diagnostics,
+                    )
+                    return torch.cat(
+                        [
+                            q_a.to(dtype=target_dtype).cpu(),
+                            k_a.to(dtype=target_dtype).cpu(),
+                            v_merged,
+                        ],
+                        dim=0,
+                    ).contiguous()
+            elif not (is_v or is_out):
+                return None
+
+    elif op_type == "ffn":
+        is_up_gate = ".ffn_up_gate." in padded
+        is_up = any(tok in padded for tok in (".ffn_up.", ".fc1.", ".w1."))
+        is_gate = any(tok in padded for tok in (".ffn_gate.", ".w3."))
+        is_down = any(tok in padded for tok in (".ffn_down.", ".fc2.", ".w2.", ".ff.net.2."))
+        has_known_ffn_part = is_up_gate or is_up or is_gate or is_down
+
+        if has_known_ffn_part and subpart == "down_only":
+            if not is_down:
+                return None
+        elif has_known_ffn_part and subpart == "up_only":
+            if is_down or is_gate:
+                return None
+            if is_up_gate and base_tensor.ndim >= 1 and base_tensor.shape[0] % 2 == 0:
+                aligned_donor, _ = _align_donor_trailing_shape(base_tensor, donor_tensor)
+                if aligned_donor.shape[0] % 2 == 0:
+                    up_a, gate_a = base_tensor.chunk(2, dim=0)
+                    up_b, _ = aligned_donor.chunk(2, dim=0)
+                    up_merged = _graft_pair_to_cpu(
+                        up_a,
+                        up_b,
+                        settings,
+                        device,
+                        target_dtype,
+                        operation_label=f"{key}[up]",
+                        diagnostics=diagnostics,
+                    )
+                    return torch.cat(
+                        [
+                            up_merged,
+                            gate_a.to(dtype=target_dtype).cpu(),
+                        ],
+                        dim=0,
+                    ).contiguous()
+            elif not is_up:
+                return None
+
+    return _graft_pair_to_cpu(
+        base_tensor,
+        donor_tensor,
+        settings,
+        device,
+        target_dtype,
+        operation_label=key,
+        diagnostics=diagnostics,
+    )
+
+
+def _resolve_donor_qkv_plan(
+    base_key: str,
+    donor_keys_set: set[str],
+    custom_map: dict[str, str],
+) -> tuple[str, list[str], int | None] | None:
+    """Resolve donor key(s) for base_key.
+
+    Returns:
+        (mode, donor_logical_keys, slice_index)
+        mode is 'direct', 'fuse_qkv' (cat q,k,v), or 'split_qkv' (chunk 3 and pick slice_index).
+    """
+    if base_key in custom_map:
+        mapped = custom_map[base_key]
+        if mapped in donor_keys_set:
+            return ("direct", [mapped], None)
+        return None
+    if base_key in donor_keys_set:
+        return ("direct", [base_key], None)
+
+    if base_key.endswith(".proj.qkv.weight"):
+        prefix = base_key[: -len("qkv.weight")]
+        parts = [f"{prefix}{p}.weight" for p in ("q", "k", "v")]
+        if all(k in donor_keys_set for k in parts):
+            return ("fuse_qkv", parts, None)
+
+    for idx, part in enumerate(("q", "k", "v")):
+        suffix = f".proj.{part}.weight"
+        if base_key.endswith(suffix):
+            fused = f"{base_key[:-len(suffix)]}.proj.qkv.weight"
+            if fused in donor_keys_set:
+                return ("split_qkv", [fused], idx)
+
+    return None
+
+
+class CWBModelGraft(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        models = folder_paths.get_filename_list("diffusion_models")
+        presets = list(DENSE_CWB_PRESETS)
+        return io.Schema(
+            node_id="CWBModelGraft",
+            display_name="CWB Model Graft",
+            category="ModelUtils/Merging",
+            description="Graft selected blocks, sub-projections, or mapped layers from a donor diffusion model into a base diffusion model using Consensus-Weighted Blending.",
+            inputs=[
+                io.Combo.Input(
+                    "execution_mode",
+                    options=["MERGE", "DOCUMENTATION ONLY"],
+                    default="MERGE",
+                    tooltip="MERGE writes a grafted model; DOCUMENTATION ONLY opens no model files.",
+                ),
+                io.Combo.Input(
+                    "model_a",
+                    options=models,
+                    tooltip="Base diffusion model that anchors structure, metadata, and protected blocks.",
+                ),
+                io.Combo.Input(
+                    "model_b",
+                    options=models,
+                    tooltip="Donor diffusion model supplying parts to graft into Model A.",
+                ),
+                io.Combo.Input(
+                    "attention_preset",
+                    options=presets,
+                    default="robust_medn",
+                    tooltip="CWB preset applied to grafted attention layers.",
+                ),
+                CWB_CONFIG.Input(
+                    "attention_config",
+                    optional=True,
+                    tooltip="Optional custom CWB configuration override for attention layers.",
+                ),
+                io.Combo.Input(
+                    "ffn_preset",
+                    options=presets,
+                    default="balanced_mean",
+                    tooltip="CWB preset applied to grafted feed-forward (MLP) layers.",
+                ),
+                CWB_CONFIG.Input(
+                    "ffn_config",
+                    optional=True,
+                    tooltip="Optional custom CWB configuration override for feed-forward (MLP) layers.",
+                ),
+                io.Int.Input(
+                    "protect_early_blocks",
+                    default=2,
+                    min=0,
+                    max=256,
+                    step=1,
+                    tooltip="Number of initial transformer blocks (starting from 0) kept 100% from Model A to preserve global layout.",
+                ),
+                io.Combo.Input(
+                    "attention_block_rule",
+                    options=[
+                        "interleaved_even",
+                        "interleaved_odd",
+                        "back_half_only",
+                        "all_after_protected",
+                        "none",
+                    ],
+                    default="interleaved_even",
+                    tooltip="Which transformer blocks after the protected start have their attention layers grafted.",
+                ),
+                io.Combo.Input(
+                    "ffn_block_rule",
+                    options=[
+                        "all_after_protected",
+                        "interleaved_even",
+                        "interleaved_odd",
+                        "back_half_only",
+                        "none",
+                    ],
+                    default="all_after_protected",
+                    tooltip="Which transformer blocks after the protected start have their feed-forward layers grafted.",
+                ),
+                io.Combo.Input(
+                    "attention_subpart",
+                    options=["all", "value_and_output", "output_only"],
+                    default="all",
+                    tooltip="Select whether to graft all attention weights, only value and output projections, or only the output projection.",
+                ),
+                io.Combo.Input(
+                    "ffn_subpart",
+                    options=["all", "up_only", "down_only"],
+                    default="all",
+                    tooltip="Select whether to graft all feed-forward weights, only the up projection, or only the down projection.",
+                ),
+                io.String.Input(
+                    "custom_key_map",
+                    default="",
+                    multiline=True,
+                    tooltip="Optional explicit key mappings (one per line as 'donor_key -> base_key') to graft across differently named layers.",
+                ),
+                io.String.Input(
+                    "output_filename",
+                    default="cwb_grafted_model",
+                    tooltip="Filename without extension under ComfyUI's diffusion_models directory.",
+                ),
+                io.Combo.Input(
+                    "save_dtype",
+                    options=["fp32", "fp16", "bf16"],
+                    default="bf16",
+                    tooltip="Requested dtype for grafted floating tensors.",
+                ),
+                io.Combo.Input(
+                    "process_device",
+                    options=["cuda", "cpu"],
+                    default="cuda",
+                    tooltip="Per-tensor FP32 CWB computation device; CUDA OOM retries on CPU.",
+                ),
+                io.String.Input(
+                    "exclude_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip="Newline or whitespace separated patterns to preserve from Model A (or whitelist when include_mode is enabled).",
+                ),
+                io.String.Input(
+                    "discard_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip="Newline or whitespace separated patterns to omit from the output file.",
+                ),
+                io.Boolean.Input(
+                    "glob_patterns",
+                    default=False,
+                    tooltip="Interpret exclude and discard entries as shell-style globs instead of regular expressions.",
+                ),
+                io.Boolean.Input(
+                    "lazy_load",
+                    default=True,
+                    tooltip="Stream only one logical work unit at a time through UEL and release completed tensors immediately.",
+                ),
+                io.Boolean.Input(
+                    "force_clear_cache",
+                    default=False,
+                    tooltip="Run Python garbage collection and empty CUDA cache before each tensor.",
+                ),
+                io.Boolean.Input(
+                    "override_dtype",
+                    default=False,
+                    tooltip="Force grafted floating tensors to Save Dtype instead of preserving Model A's float precision.",
+                ),
+                io.Boolean.Input(
+                    "include_mode",
+                    default=False,
+                    tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are eligible for grafting; nonmatching layers are preserved from Model A.",
+                ),
+            ],
+            outputs=[
+                io.AnyType.Output(display_name="output_filename"),
+                io.String.Output(display_name="documentation"),
+                io.String.Output(display_name="cwb_report"),
+            ],
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, **kwargs) -> io.NodeOutput:
+        documentation = load_documentation_from_file("cwb_model_graft.md")
+        if kwargs["execution_mode"] == "DOCUMENTATION ONLY":
+            return io.NodeOutput(
+                "Documentation mode active. No merge performed.",
+                documentation,
+                "CWB report unavailable because no merge was performed.",
+            )
+        filename, report = cls._run_graft(kwargs)
+        return io.NodeOutput(filename, documentation, report)
+
+    @classmethod
+    def _run_graft(cls, params: dict) -> tuple[str, str]:
+        model_names = [params["model_a"], params["model_b"]]
+        paths = []
+        for name in model_names:
+            path = folder_paths.get_full_path("diffusion_models", name)
+            if not path:
+                raise FileNotFoundError(f"diffusion_models input '{name}' was not found.")
+            paths.append(path)
+
+        total_size = sum(estimate_model_size(path) for path in paths)
+        device = params["process_device"]
+        if total_size:
+            prepare_for_large_operation(total_size * 1.2, torch.device(device))
+
+        custom_map = _parse_custom_key_map(params.get("custom_key_map", ""))
+        attn_settings = resolve_cwb_settings(
+            params["attention_preset"],
+            DENSE_CWB_PRESETS,
+            params.get("attention_config"),
+        )
+        ffn_settings = resolve_cwb_settings(
+            params["ffn_preset"],
+            DENSE_CWB_PRESETS,
+            params.get("ffn_config"),
+        )
+        diagnostics = CWBDiagnostics()
+        protect_early = int(params.get("protect_early_blocks", 2))
+        attn_rule = params.get("attention_block_rule", "interleaved_even")
+        ffn_rule = params.get("ffn_block_rule", "all_after_protected")
+        attn_subpart = params.get("attention_subpart", "all")
+        ffn_subpart = params.get("ffn_subpart", "all")
+        glob_mode = bool(params.get("glob_patterns", False))
+        include_mode = bool(params.get("include_mode", False))
+        exclude = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=glob_mode)
+        discard = _compile_patterns(params.get("discard_patterns", ""), glob_mode=glob_mode)
+        requested_dtype = _requested_dtype(params["save_dtype"])
+
+        handlers = [
+            MemoryEfficientSafeOpen(path, low_memory=params["lazy_load"])
+            for path in paths
+        ]
+        try:
+            quantizers = [
+                DiffusionQuantization(handler, path, f"CWB Graft input {idx + 1}")
+                for idx, (handler, path) in enumerate(zip(handlers, paths))
+            ]
+            key_maps = [
+                diffusion_key_map(q.data_keys, f"CWB Graft input {idx + 1}")
+                for idx, q in enumerate(quantizers)
+            ]
+            low_bit_sets = [
+                {
+                    logical
+                    for logical, physical in mapping.items()
+                    if physical in q.isolated_low_bit_keys
+                }
+                for mapping, q in zip(key_maps, quantizers)
+            ]
+
+            base_keys = sorted(key_maps[0])
+            donor_keys_set = set(key_maps[1])
+            block_indices = [
+                idx
+                for k in base_keys
+                if (idx := _extract_main_block_index(k)) is not None
+            ]
+            max_block_index = max(block_indices) if block_indices else 0
+
+            def matches_filter(key: str, patterns: list) -> bool:
+                aliases = [key, key_maps[0][key]]
+                return any(
+                    _matches_any_pattern(alias, patterns, glob_mode=glob_mode)
+                    for alias in aliases
+                )
+
+            def should_graft_key(key: str) -> tuple[bool, str | None, tuple[str, list[str], int | None] | None]:
+                if matches_filter(key, discard):
+                    return (False, None, None)
+                has_filter_rules = bool(exclude)
+                matched_filter = matches_filter(key, exclude)
+                if include_mode and not matched_filter:
+                    return (False, None, None)
+                if not include_mode and matched_filter:
+                    return (False, None, None)
+
+                donor_plan = _resolve_donor_qkv_plan(key, donor_keys_set, custom_map)
+                if donor_plan is None:
+                    return (False, None, None)
+                _, donor_logical_keys, _ = donor_plan
+                if key in low_bit_sets[0] or any(dk in low_bit_sets[1] for dk in donor_logical_keys):
+                    return (False, None, None)
+
+                op_type = _classify_graft_operator(key)
+                if key in custom_map or (include_mode and has_filter_rules and matched_filter):
+                    return (True, op_type or "ffn", donor_plan)
+
+                block_idx = _extract_main_block_index(key)
+                if op_type == "attention":
+                    if _block_rule_allows(block_idx, attn_rule, protect_early, max_block_index):
+                        return (True, "attention", donor_plan)
+                elif op_type == "ffn":
+                    if _block_rule_allows(block_idx, ffn_rule, protect_early, max_block_index):
+                        return (True, "ffn", donor_plan)
+                return (False, None, None)
+
+            work_units = []
+            graft_plans: dict[str, tuple[str, tuple[str, list[str], int | None]]] = {}
+            for key in base_keys:
+                if matches_filter(key, discard):
+                    work_units.append((key, {}))
+                    continue
+                graft_ok, op_type, donor_plan = should_graft_key(key)
+                base_phys = key_maps[0][key]
+                if not graft_ok:
+                    entries = (
+                        {0: quantizers[0].required_keys(base_phys)}
+                        if base_phys in quantizers[0].quantized_keys
+                        else {}
+                    )
+                    work_units.append((key, entries))
+                    continue
+
+                assert op_type is not None and donor_plan is not None
+                graft_plans[key] = (op_type, donor_plan)
+                donor_req: list[str] = []
+                for dk in donor_plan[1]:
+                    d_phys = key_maps[1][dk]
+                    for rk in quantizers[1].required_keys(d_phys):
+                        if rk not in donor_req:
+                            donor_req.append(rk)
+                work_units.append(
+                    (
+                        key,
+                        {
+                            0: quantizers[0].required_keys(base_phys),
+                            1: donor_req,
+                        },
+                    )
+                )
+
+            output_path, output_rel = canonical_model_artifact_path(
+                "diffusion_models", params["output_filename"]
+            )
+            output_metadata = quantizers[0].output_metadata()
+            output_metadata["cwb.graft"] = json.dumps(
+                {
+                    "base_model": params["model_a"],
+                    "donor_model": params["model_b"],
+                    "attention_preset": params["attention_preset"],
+                    "ffn_preset": params["ffn_preset"],
+                    "protect_early_blocks": protect_early,
+                    "attention_block_rule": attn_rule,
+                    "ffn_block_rule": ffn_rule,
+                    "attention_subpart": attn_subpart,
+                    "ffn_subpart": ffn_subpart,
+                },
+                sort_keys=True,
+            )
+
+            def decode_logical(index: int, logical_key: str, loaded: dict) -> torch.Tensor:
+                phys = key_maps[index][logical_key]
+                if phys in quantizers[index].quantized_keys:
+                    parts = {
+                        part: loaded[(index, part)]
+                        for part in quantizers[index].required_keys(phys)
+                    }
+                    return quantizers[index].decode(phys, parts, torch.float32, device="cpu")
+                return loaded[(index, phys)]
+
+            def preserve_base(writer, key: str, loaded: dict, *, force_raw: bool = False):
+                phys = key_maps[0][key]
+                if phys in quantizers[0].quantized_keys and (0, phys) in loaded:
+                    decoded = decode_logical(0, key, loaded)
+                    writer.write(key, decoded.to(requested_dtype).cpu().contiguous())
+                else:
+                    write_preserved_tensor(
+                        writer,
+                        phys,
+                        handlers[0],
+                        output_key=key,
+                        force_raw=force_raw,
+                        tensor=loaded.get((0, phys)),
+                    )
+
+            handler_map = {0: handlers[0], 1: handlers[1]}
+            grafted_count = 0
+            preserved_count = 0
+            pbar = comfy.utils.ProgressBar(len(work_units))
+            with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing(
+                stream_work_units(
+                    handler_map,
+                    work_units,
+                    pin_memory=str(device).startswith("cuda"),
+                )
+            ) as streamed:
+                for key, loaded in tqdm(streamed, total=len(work_units), desc="CWB grafting tensors", unit="tensors"):
+                    _clear_previous_layer(params)
+                    if matches_filter(key, discard):
+                        pbar.update(1)
+                        continue
+                    if key not in graft_plans:
+                        preserved_count += 1
+                        preserve_base(writer, key, loaded, force_raw=key in low_bit_sets[0])
+                        pbar.update(1)
+                        continue
+
+                    op_type, (mode, donor_logical_keys, slice_idx) = graft_plans[key]
+                    base_tensor = decode_logical(0, key, loaded)
+                    if mode == "direct":
+                        donor_tensor = decode_logical(1, donor_logical_keys[0], loaded)
+                    elif mode == "fuse_qkv":
+                        donor_tensor = torch.cat(
+                            [decode_logical(1, dk, loaded) for dk in donor_logical_keys],
+                            dim=0,
+                        )
+                    else:
+                        fused_donor = decode_logical(1, donor_logical_keys[0], loaded)
+                        donor_tensor = fused_donor.chunk(3, dim=0)[slice_idx or 0]
+
+                    base_phys = key_maps[0][key]
+                    base_dtype = quantizers[0].logical_dtype(base_phys, requested_dtype)
+                    if not _is_float_dtype(base_dtype) or not _is_float_dtype(donor_tensor.dtype):
+                        preserved_count += 1
+                        preserve_base(writer, key, loaded)
+                        pbar.update(1)
+                        continue
+
+                    target_dtype = (
+                        requested_dtype
+                        if params.get("override_dtype", False)
+                        else base_dtype
+                    )
+                    active_settings = attn_settings if op_type == "attention" else ffn_settings
+                    active_subpart = attn_subpart if op_type == "attention" else ffn_subpart
+                    grafted = _graft_with_subpart(
+                        key,
+                        op_type,
+                        active_subpart,
+                        base_tensor,
+                        donor_tensor,
+                        active_settings,
+                        device,
+                        target_dtype,
+                        diagnostics,
+                    )
+                    if grafted is None:
+                        preserved_count += 1
+                        preserve_base(writer, key, loaded)
+                    else:
+                        grafted_count += 1
+                        writer.write(key, grafted)
+                    del base_tensor, donor_tensor, grafted
+                    pbar.update(1)
+
+            report = diagnostics.render(
+                f"attn={params['attention_preset']}, ffn={params['ffn_preset']}",
+                2,
+                params.get("attention_config") is not None or params.get("ffn_config") is not None,
+            )
+            report += f"\nGrafted tensors: {grafted_count}\nPreserved base tensors: {preserved_count}\n"
+            return output_rel, report
+        finally:
+            for handler in handlers:
+                handler.__exit__(None, None, None)
+            cleanup_after_operation()
+
+
+CWB_GRAFT_NODES = [
+    CWBModelGraft,
+]

nodes/dora_learned_wd.py

@@ -0,0 +1,874 @@
+"""
+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 math
+import torch
+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
+from .uel_io import atomic_uel_writer
+from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
+from .layer_parameters import parameter_input, resolve_layer_parameters
+from .quantization_guard import DiffusionQuantization, diffusion_key_map
+
+# Import the existing SV ranking and chunking logic so we don't have to duplicate it
+from .dora_extract_wd import (
+    _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",
+    knee_probe_offset: int = 32,
+    clamp_quantile: float = 0.99,
+) -> 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,
+                clamp_quantile,
+                knee_probe_offset=knee_probe_offset,
+            )
+            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 = False,
+    glob_skip_patterns: bool = False,
+    knee_probe_offset: int = 32,
+    include_mode: bool = False,
+    layer_parameters=None,
+) -> 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:
+        quantization_a = DiffusionQuantization(handler_a, model_a_path, "Learned DoRA Extract model A")
+        quantization_b = DiffusionQuantization(handler_b, model_b_path, "Learned DoRA Extract model B")
+        keys_a = set(quantization_a.data_keys)
+        keys_b = set(quantization_b.data_keys)
+        logical_a = diffusion_key_map(keys_a, "Learned DoRA Extract model A")
+        logical_b = diffusion_key_map(keys_b, "Learned DoRA Extract model B")
+        paired_keys_b = {
+            key_a: logical_b.get(logical) for logical, key_a in logical_a.items()
+        }
+        weight_keys = sorted(
+            key for key in keys_a
+            if key.endswith(".weight")
+            and key not in quantization_a.isolated_low_bit_keys
+            and paired_keys_b[key] not in quantization_b.isolated_low_bit_keys
+        )
+        if mode == "fixed":
+            profile = "extract:fixed"
+            defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff}
+        elif mode == "ratio":
+            profile = "extract:ratio"
+            defaults = {"linear_ratio": linear_param, "conv_ratio": conv_param,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
+        elif mode == "quantile":
+            profile = "extract:quantile"
+            defaults = {"linear_quantile": linear_param, "conv_quantile": conv_param,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
+        elif mode == "sv_fro":
+            profile = "extract:frobenius"
+            defaults = {"linear_target": linear_param, "conv_target": conv_param,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
+        else:
+            profile = "extract:knee"
+            defaults = {"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff}
+        resolved_parameters = resolve_layer_parameters(
+            layer_parameters, profile, weight_keys, defaults,
+            node_name="Learned DoRA Extract",
+        )
+        work_units = [
+            (key, key, paired_keys_b[key])
+            for key in weight_keys
+            if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
+            and not (paired_keys_b[key] is None and mismatch_mode == "skip")
+        ]
+        pbar = comfy.utils.ProgressBar(len(work_units))
+        stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
+
+        def _process_layer(key, cpu_a, cpu_b):
+            key_b = paired_keys_b[key]
+            values = resolved_parameters.get(key, defaults)
+            if mode == "fixed":
+                layer_linear_param, layer_conv_param = values["linear_dim"], values["conv_dim"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            elif mode == "ratio":
+                layer_linear_param, layer_conv_param = values["linear_ratio"], values["conv_ratio"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            elif mode == "quantile":
+                layer_linear_param, layer_conv_param = values["linear_quantile"], values["conv_quantile"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            elif mode == "sv_fro":
+                layer_linear_param, layer_conv_param = values["linear_target"], values["conv_target"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            else:
+                layer_linear_param, layer_conv_param = linear_param, conv_param
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            layer_clamp_quantile = values["clamp_quantile"]
+            layer_min_diff = values["min_diff"]
+            lora_name = _format_lora_key(key)
+            weight_diff = None
+            tensor_base = None
+            tensor_ft = None
+            layer_device = device
+
+            if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)):
+                return "skipped", None
+
+            if key_b is None:
+                if mismatch_mode == "skip":
+                    return "skipped", None
+                if mismatch_mode == "error":
+                    raise ValueError(f"Key {key} not found in model B")
+
+                weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
+                    dora_difference(cpu_a, None, device, keep_sources=True)
+                )
+            else:
+                if cpu_a.shape != cpu_b.shape:
+                    if mismatch_mode == "skip":
+                        return "skipped", None
+                    if mismatch_mode == "error":
+                        raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
+
+                    weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
+                        dora_difference(cpu_a, None, device, keep_sources=True)
+                    )
+                else:
+                    if cpu_a.ndim < 2:
+                        return "skipped", None
+                    weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
+                        dora_difference(cpu_a, cpu_b, device, keep_sources=True)
+                    )
+
+            if weight_diff is None or tensor_base is None or tensor_ft is None:
+                raise RuntimeError(f"No extraction tensors were produced for {key}")
+
+            if layer_min_diff > 0 and weight_diff.abs().max() < layer_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 = {}
+            svd_weight = weight_diff
+            optimize_base = tensor_base
+            optimize_ft = tensor_ft
+
+            try:
+                # 1. Run Analytical SVD (Initialization)
+                if is_conv:
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_conv(svd_weight, mode, layer_conv_param, layer_device, layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
+                        lambda: _svd_extract_conv(svd_weight.cpu(), mode, layer_conv_param, "cpu", layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
+                        layer_device,
+                    )
+                else:
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_linear(svd_weight, mode, layer_linear_param, layer_device, layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
+                        lambda: _svd_extract_linear(svd_weight.cpu(), mode, layer_linear_param, "cpu", layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
+                        layer_device,
+                    )
+
+                # 2. Run Optimization (Learned Rounding)
+                if mode_str != "full" and dora_scale is not None:
+                    lora_down, lora_up, _ = result
+                    optimize_args = (
+                        optimize_iters, learning_rate, optimizer_type, lr_schedule,
+                        lr_patience, lr_factor, lr_cooldown, early_stop_loss,
+                        early_stop_stall, early_stop_lr,
+                    )
+                    lora_down, lora_up = retry_cuda_oom_on_cpu(
+                        lambda: _optimize_dora(optimize_base, optimize_ft, lora_down, lora_up, *optimize_args, layer_name=lora_name),
+                        lambda: _optimize_dora(optimize_base.cpu(), optimize_ft.cpu(), lora_down.cpu(), lora_up.cpu(), *optimize_args, layer_name=lora_name),
+                        layer_device,
+                    )
+                    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, layer_linear_param, layer_device, layer_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,
+                            knee_probe_offset=knee_probe_offset,
+                            clamp_quantile=layer_clamp_quantile,
+                        )
+                        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_context = atomic_uel_writer(output_path)
+        writer = writer_context.__enter__()
+        stream = None
+        try:
+            stream = paired_async_tensors(
+                handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda"),
+                quantization_a=quantization_a, quantization_b=quantization_b,
+                compute_dtype=torch.float32,
+            )
+            for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting Learned DoRA", unit="layers"):
+                status, layer_sd = _process_layer(key, cpu_a, cpu_b)
+                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)
+            stats["skipped"] += len(weight_keys) - len(work_units)
+        except BaseException as exc:
+            if stream is not None:
+                stream.close()
+            writer_context.__exit__(type(exc), exc, exc.__traceback__)
+            raise
+        else:
+            if stream is not None:
+                stream.close()
+            writer_context.__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(),
+                parameter_input("extract:fixed"),
+            ],
+            outputs=[io.AnyType.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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
+            layer_parameters=layer_parameters,
+        )
+
+        return io.NodeOutput(output_name)
+
+
+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, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
+                io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_learned_inputs(),
+                *_get_common_inputs(),
+                parameter_input("extract:ratio"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
+            knee_probe_offset=probe_offset, layer_parameters=layer_parameters,
+        )
+
+        return io.NodeOutput(output_name)
+
+
+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, tooltip="Target cumulative singular-value fraction for linear layers."),
+                io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_learned_inputs(),
+                *_get_common_inputs(),
+                parameter_input("extract:quantile"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
+            knee_probe_offset=probe_offset, layer_parameters=layer_parameters,
+        )
+
+        return io.NodeOutput(output_name)
+
+
+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", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
+                io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_learned_inputs(),
+                *_get_common_inputs(),
+                parameter_input("extract:knee"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
+                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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
+            layer_parameters=layer_parameters,
+        )
+
+        return io.NodeOutput(output_name)
+
+
+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, tooltip="Target Frobenius-norm fraction retained for linear layers."),
+                io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_learned_inputs(),
+                *_get_common_inputs(),
+                parameter_input("extract:frobenius"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+
+        return io.NodeOutput(output_name)

nodes/downloader_nodes.py

@@ -0,0 +1,341 @@
+import os
+import folder_paths
+from comfy_api.latest import io
+from .downloader_utils import scan_and_process
+
+def get_subdirectories(folder_name: str) -> list[str]:
+    """Helper to get a list of subdirectories for a given model folder name."""
+    paths = folder_paths.get_folder_paths(folder_name)
+    if not paths:
+        return ["/"]
+
+    subdirs = set()
+    for base_path in paths:
+        if not os.path.exists(base_path):
+            continue
+        for root, dirs, _ in os.walk(base_path):
+            # Ignore hidden directories directly during walk
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+
+            rel_path = os.path.relpath(root, base_path)
+            if rel_path == ".":
+                subdirs.add("/")
+            else:
+                # Store with forward slashes for consistency
+                subdirs.add("/" + rel_path.replace("\\", "/"))
+
+    return sorted(list(subdirs))
+
+class BaseInfoMetaDownloaderNode(io.ComfyNode):
+    CATEGORY = "ModelUtils/Downloader"
+    OUTPUT_NODE = True
+
+    @classmethod
+    def get_common_inputs(cls):
+        return {
+            "recursive": io.Boolean.Input("recursive", default=True, tooltip="Whether to scan subdirectories recursively."),
+            "nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All", tooltip="Filter models based on NSFW level in metadata."),
+            "max_examples": io.Int.Input("max_examples", default=0, min=0, max=100, tooltip="Maximum number of examples to process."),
+            "threads": io.Int.Input("threads", default=4, min=1, max=32, tooltip="Number of threads to use for processing."),
+            "api_key": io.String.Input("api_key", default="", tooltip="API key for accessing the model repository.")
+        }
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name=cls.__name__.replace("Downloader", " Downloader"),
+            category=cls.CATEGORY,
+            inputs=list(cls.get_common_inputs().values()),
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+            description=("Scans specified directories for models, extracts metadata, generates previews, and optionally compresses files.")
+        )
+
+    @classmethod
+    def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        scan_dirs = cls.get_scan_dirs(**kwargs)
+        if not scan_dirs:
+            return io.NodeOutput("No directories to scan.")
+
+        # Call the scan_and_process function
+        stats = scan_and_process(
+            scan_dirs=scan_dirs,
+            recursive=recursive,
+            nsfw_level=nsfw_level,
+            max_examples=max_examples,
+            api_key=api_key if api_key else None,
+            threads=threads
+        )
+
+        mb_saved = stats.get("space_saved", 0) / (1024 * 1024)
+
+        output_msg = (
+            f"Scan Complete for {len(scan_dirs)} directory/directories.\n"
+            f"Total Models Processed: {stats.get('total_processed', 0)}\n"
+            f"New Models Hashed: {stats.get('new_hashes', 0)}\n"
+            f"Images Processed: {stats.get('images_processed', 0)}\n"
+            f"Workflows Extracted: {stats.get('workflows_extracted', 0)}\n"
+            f"Space Saved via Compression: {mb_saved:.2f} MB"
+        )
+
+        return io.NodeOutput(output_msg)
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs) -> list[str]:
+        return []
+
+class CheckpointInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("checkpoints"), default="/", tooltip="Checkpoint subdirectory to scan; Recursive also includes its descendants."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Checkpoint Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        return super().execute(
+            recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
+            threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        base_paths = folder_paths.get_folder_paths("checkpoints")
+        if not base_paths:
+            return []
+        sub = kwargs.get("subdirectory", "/")
+        if sub == "/":
+            return base_paths
+
+        target_paths = []
+        for bp in base_paths:
+            target_path = os.path.join(bp, sub.lstrip("/"))
+            if os.path.exists(target_path):
+                target_paths.append(target_path)
+        return target_paths
+
+class DiffusionModelInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("diffusion_models"), default="/", tooltip="Diffusion-model subdirectory to scan; Recursive also includes its descendants."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Diffusion Model Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        return super().execute(
+            recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
+            threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        base_paths = folder_paths.get_folder_paths("diffusion_models")
+        if not base_paths:
+            return []
+        sub = kwargs.get("subdirectory", "/")
+        if sub == "/":
+            return base_paths
+
+        target_paths = []
+        for bp in base_paths:
+            target_path = os.path.join(bp, sub.lstrip("/"))
+            if os.path.exists(target_path):
+                target_paths.append(target_path)
+        return target_paths
+
+class LoRAInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("loras"), default="/", tooltip="LoRA subdirectory to scan; Recursive also includes its descendants."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="LoRA Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        return super().execute(
+            recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
+            threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        base_paths = folder_paths.get_folder_paths("loras")
+        if not base_paths:
+            return []
+        sub = kwargs.get("subdirectory", "/")
+        if sub == "/":
+            return base_paths
+
+        target_paths = []
+        for bp in base_paths:
+            target_path = os.path.join(bp, sub.lstrip("/"))
+            if os.path.exists(target_path):
+                target_paths.append(target_path)
+        return target_paths
+
+class EmbeddingInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("embeddings"), default="/", tooltip="Embedding subdirectory to scan; Recursive also includes its descendants."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Embedding Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        return super().execute(
+            recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
+            threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        base_paths = folder_paths.get_folder_paths("embeddings")
+        if not base_paths:
+            return []
+        sub = kwargs.get("subdirectory", "/")
+        if sub == "/":
+            return base_paths
+
+        target_paths = []
+        for bp in base_paths:
+            target_path = os.path.join(bp, sub.lstrip("/"))
+            if os.path.exists(target_path):
+                target_paths.append(target_path)
+        return target_paths
+
+class VAEInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("vae"), default="/", tooltip="VAE subdirectory to scan; Recursive also includes its descendants."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="VAE Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        return super().execute(
+            recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
+            threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        base_paths = folder_paths.get_folder_paths("vae")
+        if not base_paths:
+            return []
+        sub = kwargs.get("subdirectory", "/")
+        if sub == "/":
+            return base_paths
+
+        target_paths = []
+        for bp in base_paths:
+            target_path = os.path.join(bp, sub.lstrip("/"))
+            if os.path.exists(target_path):
+                target_paths.append(target_path)
+        return target_paths
+
+class ControlNetInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("controlnet"), default="/", tooltip="ControlNet subdirectory to scan; Recursive also includes its descendants."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="ControlNet Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
+        return super().execute(
+            recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
+            threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        base_paths = folder_paths.get_folder_paths("controlnet")
+        if not base_paths:
+            return []
+        sub = kwargs.get("subdirectory", "/")
+        if sub == "/":
+            return base_paths
+
+        target_paths = []
+        for bp in base_paths:
+            target_path = os.path.join(bp, sub.lstrip("/"))
+            if os.path.exists(target_path):
+                target_paths.append(target_path)
+        return target_paths
+
+class ManualPathInfoMetaDownloader(BaseInfoMetaDownloaderNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = list(cls.get_common_inputs().values())
+        inputs.append(io.String.Input("scan_dir", default="", tooltip="Absolute directory to scan for supported model files."))
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Manual Path Downloader",
+            category=cls.CATEGORY,
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="status")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, scan_dir, **kwargs):
+        return super().execute(
+            recursive=recursive,
+            nsfw_level=nsfw_level,
+            max_examples=max_examples,
+            threads=threads,
+            api_key=api_key,
+            scan_dir=scan_dir,
+            **kwargs
+        )
+
+    @classmethod
+    def get_scan_dirs(cls, **kwargs):
+        scan_dir = kwargs.get("scan_dir", "")
+        if scan_dir.strip():
+            return [scan_dir.strip()]
+        return []

nodes/downloader_utils.py

@@ -0,0 +1,638 @@
+import os
+import sys
+import hashlib
+import json
+import sqlite3
+import requests
+import glob
+from pathlib import Path
+import threading
+import concurrent.futures
+import codecs
+import shutil
+
+try:
+    from PIL import Image
+    HAS_PIL = True
+except ImportError:
+    HAS_PIL = False
+    print("Warning: Pillow not found. Image workflow extraction disabled.")
+
+try:
+    from mutagen.mp4 import MP4
+    HAS_MUTAGEN = True
+except ImportError:
+    HAS_MUTAGEN = False
+    print("Warning: mutagen not found. MP4 tag extraction will rely entirely on av fallback.")
+
+try:
+    import av
+    HAS_AV = True
+except ImportError:
+    HAS_AV = False
+    print("Warning: av (PyAV) not found. MP4 metadata extraction fallback disabled.")
+
+# Assuming downloader_utils.py is inside nodes/, the root is one level up
+root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+from unifiedefficientloader import UnifiedSafetensorsLoader
+from .model_info_nodes import (
+    get_potential_preview_files,
+    load_image_tensor,
+    get_model_workflows,
+    get_model_metadata_file,
+)
+
+# --- Database Path Logic ---
+try:
+    import folder_paths
+    base_user_dir = folder_paths.get_public_user_directory("default")
+    if not base_user_dir:
+        base_user_dir = root_dir
+except (ImportError, AttributeError):
+    base_user_dir = root_dir
+
+PLUGIN_DATA_DIR = os.path.join(base_user_dir, "ComfyUI-ModelUtils")
+os.makedirs(PLUGIN_DATA_DIR, exist_ok=True)
+DEFAULT_DB_PATH = os.path.join(PLUGIN_DATA_DIR, "cache.db")
+
+
+# --- Workflow Extraction Logic ---
+
+def unescape_and_parse_nested_json(json_string):
+    if not isinstance(json_string, str): return None
+    stripped = json_string.strip()
+    if not (stripped.startswith('{') or stripped.startswith('[')): return None
+    try: return json.loads(stripped)
+    except json.JSONDecodeError: pass
+    try:
+        first_parse = json.loads(stripped)
+        if isinstance(first_parse, str): return json.loads(first_parse)
+    except json.JSONDecodeError: pass
+    try:
+        decoded = codecs.decode(stripped, "unicode_escape")
+        return json.loads(decoded)
+    except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass
+    try:
+        unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\")
+        unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t")
+        return json.loads(unescaped)
+    except json.JSONDecodeError: pass
+    return None
+
+def get_workflow_blob(metadata):
+    for key in ['workflow', 'prompt']:
+        val = metadata.get(key)
+        if val:
+            parsed = unescape_and_parse_nested_json(val)
+            if parsed: return parsed
+    for val in metadata.values():
+        parsed = unescape_and_parse_nested_json(val)
+        if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
+            return parsed
+    return None
+
+def process_image_workflow_and_compress(filepath_str):
+    if not HAS_PIL: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
+    p = Path(filepath_str)
+    if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
+
+    workflow = None
+    actual_format = None
+    stats = {"images_processed": 1, "workflows_extracted": 0, "space_saved": 0}
+    try:
+        with Image.open(p) as img:
+            actual_format = img.format
+            # CivitAI sometimes serves PNGs with .jpeg extension in URL
+            if actual_format == 'PNG':
+                workflow = get_workflow_blob(img.info)
+                if workflow:
+                    # Save workflow
+                    json_path = p.with_suffix('.json')
+                    if not json_path.exists():
+                        with open(json_path, 'w', encoding='utf-8') as f:
+                            json.dump(workflow, f, indent=2, ensure_ascii=False)
+                        print(f"Extracted workflow to {json_path}")
+                        stats["workflows_extracted"] = 1
+
+                    # Compress image to WebP
+                    webp_path = p.with_suffix('.webp')
+                    if not webp_path.exists():
+                        orig_size = os.path.getsize(p)
+                        img.save(webp_path, 'WEBP', quality=85, method=4)
+                        new_size = os.path.getsize(webp_path)
+                        stats["space_saved"] = max(0, orig_size - new_size)
+                        print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}")
+
+        # Delete original if we created a webp replacement
+        if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG':
+            os.remove(p)
+            print(f"Deleted original uncompressed image: {p.name}")
+    except Exception as e:
+        print(f"Failed to process image {p.name}: {e}")
+    return stats
+
+# --- Video Workflow Logic ---
+def get_tag_value(mp4, keys):
+    for key in keys:
+        if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None
+    for key in keys:
+        search_suffix = key.split(':')[-1].lower()
+        for tag_key, val in mp4.tags.items():
+            if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None
+    return None
+
+def extract_metadata_av(mp4_path):
+    if not HAS_AV: return None
+    try:
+        with av.open(str(mp4_path)) as container:
+            return {k.lower(): v for k, v in container.metadata.items()}
+    except Exception as e:
+        print(f"av extraction error for {mp4_path}: {e}")
+        return None
+
+def process_video_workflow(filepath_str):
+    p = Path(filepath_str)
+    if p.suffix.lower() != '.mp4': return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
+
+    json_path = p.with_suffix('.json')
+    if json_path.exists(): return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
+
+    workflow_data = None
+    mp4 = None
+    stats = {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
+    if HAS_MUTAGEN:
+        try: mp4 = MP4(p)
+        except: pass
+
+    if mp4:
+        for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]:
+            val = get_tag_value(mp4, tag_keys)
+            if val:
+                if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore')
+                try:
+                    parsed = json.loads(val)
+                    if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
+                        if 'prompt' in parsed and isinstance(parsed['prompt'], str):
+                            try: workflow_data = json.loads(parsed['prompt'])
+                            except: workflow_data = parsed
+                        else:
+                            workflow_data = parsed
+                        break
+                except json.JSONDecodeError:
+                    # Try more aggressive parsing if it's a nested string
+                    parsed = unescape_and_parse_nested_json(val)
+                    if parsed:
+                        workflow_data = parsed
+                        break
+
+    if not workflow_data:
+        av_data = extract_metadata_av(p)
+        if av_data:
+            for key in ['workflow', 'prompt', 'comment', 'description']:
+                if key in av_data:
+                    val = av_data[key]
+                    try:
+                        parsed = json.loads(val)
+                        if isinstance(parsed, dict):
+                            if 'prompt' in parsed and isinstance(parsed['prompt'], str):
+                                try: workflow_data = json.loads(parsed['prompt'])
+                                except: workflow_data = parsed
+                            else:
+                                workflow_data = parsed
+                            break
+                    except json.JSONDecodeError:
+                        parsed = unescape_and_parse_nested_json(val)
+                        if parsed:
+                            workflow_data = parsed
+                            break
+
+    if workflow_data:
+        try:
+            with open(json_path, 'w', encoding='utf-8') as f:
+                json.dump(workflow_data, f, indent=2, ensure_ascii=False)
+            print(f"Extracted workflow to {json_path.name}")
+            stats["workflows_extracted"] = 1
+        except Exception as e:
+            print(f"Error saving workflow for {p.name}: {e}")
+    return stats
+
+# --- Constants for Civitai NSFW levels ---
+NSFW_LEVELS = {
+    "None": 1,
+    "Soft": 2,
+    "Mature": 4,
+    "X": 8,
+    "XXX": 16,
+    "All": 32  # Anything
+}
+
+class HashCache:
+    def __init__(self, db_path=None):
+        self.db_path = db_path if db_path else DEFAULT_DB_PATH
+        self.is_sqlite = self.db_path.endswith('.db') or self.db_path.endswith('.sqlite')
+        self.lock = threading.Lock()
+        if self.is_sqlite:
+            self._init_sqlite()
+        else:
+            self.cache = self._load_json()
+
+    def _init_sqlite(self):
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS file_hashes (
+                    filepath TEXT PRIMARY KEY,
+                    mtime REAL,
+                    size INTEGER,
+                    hash TEXT
+                )
+            ''')
+            conn.commit()
+
+    def _load_json(self):
+        if os.path.exists(self.db_path):
+            with open(self.db_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        return {}
+
+    def _save_json(self):
+        with open(self.db_path, "w", encoding="utf-8") as f:
+            json.dump(self.cache, f, indent=4)
+
+    def _get_sqlite_hash(self, filepath, mtime, size):
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,))
+            row = cursor.fetchone()
+            if row:
+                db_mtime, db_size, db_hash = row
+                if db_mtime == mtime and db_size == size:
+                    return db_hash
+        return None
+
+    def _set_sqlite_hash(self, filepath, mtime, size, file_hash):
+        with sqlite3.connect(self.db_path, timeout=10.0) as conn:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash)
+                VALUES (?, ?, ?, ?)
+            ''', (filepath, mtime, size, file_hash))
+            conn.commit()
+
+    def get_hash(self, filepath):
+        if not os.path.exists(filepath):
+            return None
+        mtime = os.path.getmtime(filepath)
+        size = os.path.getsize(filepath)
+        key = filepath
+
+        with self.lock:
+            if self.is_sqlite:
+                cached_hash = self._get_sqlite_hash(key, mtime, size)
+                if cached_hash: return cached_hash
+            else:
+                if key in self.cache:
+                    entry = self.cache[key]
+                    if entry["mtime"] == mtime and entry["size"] == size:
+                        return entry["hash"]
+
+        # Calculate new hash (outside the lock to allow parallel hashing)
+        print(f"Calculating SHA256 for: {filepath}")
+        sha256_hash = self.calculate_sha256(filepath)
+
+        with self.lock:
+            if self.is_sqlite:
+                self._set_sqlite_hash(key, mtime, size, sha256_hash)
+            else:
+                self.cache[key] = {
+                    "mtime": mtime,
+                    "size": size,
+                    "hash": sha256_hash
+                }
+                self._save_json()
+
+        return sha256_hash
+
+    @staticmethod
+    def calculate_sha256(filename):
+        sha256_hash = hashlib.sha256()
+        with open(filename, "rb") as f:
+            # Read and update hash string value in blocks of 4K
+            for byte_block in iter(lambda: f.read(4096), b""):
+                sha256_hash.update(byte_block)
+        return sha256_hash.hexdigest()
+
+class CivitaiAPI:
+    def __init__(self, api_key=None):
+        self.api_key = api_key
+        self.base_url = "https://civitai.com/api/v1"
+
+    def get_headers(self):
+        headers = {}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        return headers
+
+    def get_model_info_by_hash(self, file_hash):
+        url = f"{self.base_url}/model-versions/by-hash/{file_hash}"
+        response = requests.get(url, headers=self.get_headers())
+        if response.status_code == 200:
+            return response.json()
+        elif response.status_code == 404:
+            return None
+        else:
+            response.raise_for_status()
+
+class Downloader:
+    @staticmethod
+    def download_file(url, target_path, headers=None):
+        if headers is None:
+            headers = {}
+
+        temp_path = target_path + ".downloading"
+
+        downloaded_size = 0
+        if os.path.exists(temp_path):
+            downloaded_size = os.path.getsize(temp_path)
+
+        if downloaded_size > 0:
+            headers["Range"] = f"bytes={downloaded_size}-"
+
+        response = requests.get(url, headers=headers, stream=True)
+
+        if response.status_code == 416: # Range not satisfiable, file probably changed or completed
+            os.remove(temp_path)
+            downloaded_size = 0
+            if "Range" in headers:
+                del headers["Range"]
+            response = requests.get(url, headers=headers, stream=True)
+
+        if response.status_code not in [200, 206]:
+            print(f"Failed to download from {url}, status code: {response.status_code}")
+            return False
+
+        mode = "ab" if response.status_code == 206 else "wb"
+
+        total_size = int(response.headers.get('content-length', 0)) + downloaded_size
+
+        print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)")
+
+        with open(temp_path, mode) as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+
+        os.rename(temp_path, target_path)
+        return True
+
+def preview_exists(base_path):
+    previews = get_potential_preview_files(base_path)
+    for p in previews:
+        if os.path.exists(p):
+            return True
+    return False
+
+def next_example_image_path(base_path, ext):
+    i = 0
+    while len(glob.glob(f"{base_path}.example.{i}.*")) > 0:
+        i += 1
+    return f"{base_path}.example.{i}.{ext}"
+
+def filter_images_by_nsfw(images, threshold_level):
+    valid_images = []
+    for img in images:
+        rating = img.get("nsfwLevel", 1)
+        if rating <= threshold_level:
+            valid_images.append(img)
+    return valid_images
+
+def extract_safetensors_metadata(filepath):
+    try:
+        # Use low_memory=True to only read the header and metadata without loading tensors
+        loader = UnifiedSafetensorsLoader(filepath, low_memory=True)
+        metadata = loader.metadata()
+        loader.close()
+        return metadata if metadata else None
+    except Exception as e:
+        print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}")
+        return None
+
+def process_file(filepath, cache, api, nsfw_threshold, max_examples):
+    filepath_str = str(filepath)
+    base_path = os.path.splitext(filepath_str)[0]
+    is_safetensors = filepath_str.lower().endswith(".safetensors")
+
+    stats = {"new_hashes": 0, "images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
+
+    # Check if all targets already exist
+    info_path = f"{base_path}.civitai.info"
+    st_metadata_path = f"{base_path}.metadata.json"
+
+    has_info = os.path.exists(info_path)
+    has_preview = preview_exists(base_path)
+    has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path)
+
+    # Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file
+    for p in get_potential_preview_files(base_path):
+        if os.path.exists(p):
+            ext_lower = p.lower()
+            if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
+                res = process_image_workflow_and_compress(p)
+                for k in stats: stats[k] += res.get(k, 0)
+            elif ext_lower.endswith('.mp4'):
+                res = process_video_workflow(p)
+                for k in stats: stats[k] += res.get(k, 0)
+
+    for ex in glob.glob(f"{base_path}.example.*"):
+        ext_lower = ex.lower()
+        if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
+            res = process_image_workflow_and_compress(ex)
+            for k in stats: stats[k] += res.get(k, 0)
+        elif ext_lower.endswith('.mp4'):
+            res = process_video_workflow(ex)
+            for k in stats: stats[k] += res.get(k, 0)
+
+    # If everything exists and we don't want examples or already have enough examples, we might skip API logic
+    if has_info and has_preview and has_st_metadata and max_examples == 0:
+        print(f"Skipping fully populated model: {filepath_str}")
+        return stats
+
+    print(f"Processing model: {filepath_str}")
+
+    # Safetensors built-in metadata extraction
+    if is_safetensors and not has_st_metadata:
+        metadata = extract_safetensors_metadata(filepath_str)
+        if metadata:
+            with open(st_metadata_path, "w", encoding="utf-8") as f:
+                json.dump(metadata, f, indent=4)
+            print(f"Saved safetensors metadata to {st_metadata_path}")
+
+    mtime = os.path.getmtime(filepath_str)
+    size = os.path.getsize(filepath_str)
+
+    # We want to know if it's a NEW hash or cached
+    is_cached = False
+    with cache.lock:
+        if cache.is_sqlite:
+            cached_hash = cache._get_sqlite_hash(filepath_str, mtime, size)
+            if cached_hash: is_cached = True
+        else:
+            if filepath_str in cache.cache:
+                entry = cache.cache[filepath_str]
+                if entry["mtime"] == mtime and entry["size"] == size:
+                    is_cached = True
+
+    file_hash = cache.get_hash(filepath_str)
+    if not file_hash:
+        return stats
+
+    if not is_cached:
+        stats["new_hashes"] += 1
+
+    try:
+        model_info = api.get_model_info_by_hash(file_hash)
+    except Exception as e:
+        print(f"API Error for {filepath_str}: {e}")
+        return stats
+
+    if not model_info:
+        print(f"Model not found on Civitai: {filepath_str}")
+        return stats
+
+    # 1. Save Info
+    if not has_info:
+        with open(info_path, "w", encoding="utf-8") as f:
+            json.dump(model_info, f, indent=4)
+        print(f"Saved metadata to {info_path}")
+
+    images = model_info.get("images", [])
+    valid_images = filter_images_by_nsfw(images, nsfw_threshold)
+
+    if not valid_images:
+        print(f"No valid images found matching NSFW criteria for {filepath_str}")
+        return stats
+
+    # 2. Save Preview
+    if not has_preview and len(valid_images) > 0:
+        preview_img = valid_images[0]
+        url = preview_img.get("url")
+        ext = "jpeg" # Default fallback
+        if ".png" in url.lower(): ext = "png"
+        elif ".mp4" in url.lower(): ext = "mp4"
+        elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
+
+        preview_path = f"{base_path}.preview.{ext}"
+        if Downloader.download_file(url, preview_path):
+            if ext in ("png", "jpg", "jpeg"):
+                res = process_image_workflow_and_compress(preview_path)
+                for k in stats: stats[k] += res.get(k, 0)
+            elif ext == "mp4":
+                res = process_video_workflow(preview_path)
+                for k in stats: stats[k] += res.get(k, 0)
+
+    # 3. Save Examples
+    if max_examples > 0:
+        existing_examples_count = len(glob.glob(f"{base_path}.example.*"))
+        # We divide by 2 roughly if there are json files, let's just count unique example indices
+        existing_indices = set()
+        for f in glob.glob(f"{base_path}.example.*"):
+            parts = f.split('.example.')
+            if len(parts) > 1:
+                idx = parts[1].split('.')[0]
+                if idx.isdigit(): existing_indices.add(idx)
+
+        needed_examples = max_examples - len(existing_indices)
+
+        if needed_examples > 0:
+            example_images = valid_images[1:max_examples+1]
+            # Skip images we probably already downloaded by skipping the first N
+            example_images_to_dl = example_images[len(existing_indices):]
+
+            for img in example_images_to_dl:
+                url = img.get("url")
+                ext = "jpeg"
+                if ".png" in url.lower(): ext = "png"
+                elif ".mp4" in url.lower(): ext = "mp4"
+                elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
+
+                example_path = next_example_image_path(base_path, ext)
+                if Downloader.download_file(url, example_path):
+                    if ext in ("png", "jpg", "jpeg"):
+                        res = process_image_workflow_and_compress(example_path)
+                        for k in stats: stats[k] += res.get(k, 0)
+                    elif ext == "mp4":
+                        res = process_video_workflow(example_path)
+                        for k in stats: stats[k] += res.get(k, 0)
+    return stats
+
+
+def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4):
+    """
+    Scans directories for models and fetches their metadata/previews.
+    Adapted for direct Python API calls rather than argparse CLI.
+    """
+    db_path = db_path if db_path else DEFAULT_DB_PATH
+    cache = HashCache(db_path)
+    api = CivitaiAPI(api_key)
+
+    model_exts = [".safetensors", ".ckpt", ".pt"]
+
+    files_to_process = []
+    if isinstance(scan_dirs, str):
+        scan_dirs = [scan_dirs]
+
+    for scan_dir in scan_dirs:
+        path = Path(scan_dir)
+        if recursive:
+            for ext in model_exts:
+                # rglob includes all directories, so we must filter out hidden ones
+                for f in path.rglob(f"*{ext}"):
+                    if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
+                        files_to_process.append(f)
+        else:
+            for ext in model_exts:
+                for f in path.glob(f"*{ext}"):
+                    if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
+                        files_to_process.append(f)
+
+    nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"])
+
+    print(f"Found {len(files_to_process)} models to process.")
+
+    global_stats = {
+        "total_processed": len(files_to_process),
+        "new_hashes": 0,
+        "images_processed": 0,
+        "workflows_extracted": 0,
+        "space_saved": 0
+    }
+
+    # Import and use ComfyUI's ProgressBar
+    try:
+        import comfy.utils
+        pbar = comfy.utils.ProgressBar(len(files_to_process))
+    except ImportError:
+        pbar = None
+
+    if threads > 1:
+        print(f"Starting {threads} threads...")
+        with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
+            futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process]
+            for future in concurrent.futures.as_completed(futures):
+                try:
+                    res = future.result()
+                    if res:
+                        for k in global_stats: global_stats[k] += res.get(k, 0)
+                except Exception as e:
+                    print(f"Error in thread: {e}")
+                if pbar:
+                    pbar.update(1)
+    else:
+        for filepath in files_to_process:
+            res = process_file(filepath, cache, api, nsfw_threshold, max_examples)
+            if res:
+                for k in global_stats: global_stats[k] += res.get(k, 0)
+            if pbar:
+                pbar.update(1)
+
+    return global_stats

nodes/dtype_conversion.py

@@ -0,0 +1,242 @@
+"""Asynchronous dtype conversion for diffusion-model safetensors files."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections import Counter
+from contextlib import nullcontext
+
+import comfy.utils
+import folder_paths
+import torch
+from comfy_api.latest import io
+from tqdm import tqdm
+from unifiedefficientloader import MemoryEfficientSafeOpen
+
+from .device_utils import cleanup_after_operation
+from .uel_io import atomic_uel_writer
+from .artifact_paths import canonical_model_artifact_path
+
+
+TARGET_DTYPES = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+
+def _compile_exclusions(pattern_text: str) -> list[re.Pattern]:
+    patterns = []
+    for line_number, value in enumerate(pattern_text.splitlines(), start=1):
+        pattern = value.strip()
+        if not pattern:
+            continue
+        try:
+            patterns.append(re.compile(pattern))
+        except re.error as exc:
+            raise ValueError(
+                f"Invalid exclusion regex on line {line_number}: {exc}"
+            ) from exc
+    return patterns
+
+
+def _is_excluded(key: str, patterns: list[re.Pattern]) -> bool:
+    return any(pattern.search(key) for pattern in patterns)
+
+
+def _output_path(output_filename: str) -> tuple[str, str]:
+    return canonical_model_artifact_path("diffusion_models", output_filename)
+
+
+def convert_diffusion_model_dtype(
+    model_name: str,
+    target_dtype: str,
+    exclude_patterns: str,
+    output_filename: str,
+    reference_model: str | None = None,
+    include_mode: bool = False,
+) -> tuple[str, str]:
+    try:
+        destination_dtype = TARGET_DTYPES[target_dtype]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported target dtype: {target_dtype}") from exc
+
+    source_path = folder_paths.get_full_path_or_raise(
+        "diffusion_models", model_name
+    )
+    output_path, output_name = _output_path(output_filename)
+    if os.path.normcase(os.path.abspath(source_path)) == os.path.normcase(
+        os.path.abspath(output_path)
+    ):
+        raise ValueError("Output path must differ from the input model path.")
+
+    exclusions = _compile_exclusions(exclude_patterns)
+    counts = Counter()
+    original_dtypes = Counter()
+    reference_path = None
+    if reference_model and reference_model != "None":
+        reference_path = folder_paths.get_full_path_or_raise(
+            "diffusion_models", reference_model
+        )
+
+    try:
+        with MemoryEfficientSafeOpen(source_path, low_memory=True) as loader, (
+            MemoryEfficientSafeOpen(reference_path, low_memory=True)
+            if reference_path else nullcontext()
+        ) as reference_loader:
+            metadata = (loader.metadata() or {}).copy()
+            reference_keys = set(reference_loader.keys()) if reference_loader else set()
+            with atomic_uel_writer(output_path, metadata) as writer:
+                keys = list(loader.keys())
+                progress = comfy.utils.ProgressBar(len(keys))
+                stream = loader.async_stream(
+                    keys,
+                    batch_size=1,
+                    prefetch_batches=1,
+                    pin_memory=False,
+                )
+                for batch in tqdm(
+                    stream,
+                    total=len(keys),
+                    desc="Converting diffusion-model dtype",
+                    unit="tensors",
+                ):
+                    key, source = batch[0]
+                    original_dtypes[str(source.dtype).removeprefix("torch.")] += 1
+                    output = source
+                    if _is_excluded(key, exclusions) != include_mode:
+                        counts["excluded"] += 1
+                    elif not source.is_floating_point():
+                        counts["non_floating"] += 1
+                    elif reference_loader is not None and key in reference_keys:
+                        reference_dtype = reference_loader.get_dtype(key)
+                        if not reference_dtype.is_floating_point:
+                            raise ValueError(
+                                f"Reference tensor '{key}' is {reference_dtype}, but the input tensor is floating."
+                            )
+                        if source.dtype == reference_dtype:
+                            counts["reference_already_matched"] += 1
+                        else:
+                            output = source.to(dtype=reference_dtype).contiguous()
+                            counts["reference_mapped"] += 1
+                    elif reference_loader is not None:
+                        output = source if source.dtype == destination_dtype else source.to(
+                            dtype=destination_dtype
+                        ).contiguous()
+                        counts["reference_missing_fallback"] += 1
+                    elif source.dtype == destination_dtype:
+                        counts["already_target"] += 1
+                    else:
+                        output = source.to(dtype=destination_dtype).contiguous()
+                        counts["converted"] += 1
+
+                    writer.write_batch([(key, output)])
+                    loader.mark_processed(key)
+                    progress.update(1)
+                    del output
+                    del source
+                    batch.clear()
+                    del batch
+    finally:
+        cleanup_after_operation()
+
+    total = sum(counts.values())
+    dtype_summary = ", ".join(
+        f"{name}={count}" for name, count in sorted(original_dtypes.items())
+    )
+    report = "\n".join(
+        (
+            "DIFFUSION MODEL DTYPE CONVERSION",
+            f"Input: {model_name}",
+            f"Reference model: {reference_model or 'None'}",
+            f"Output: {output_name}",
+            f"Target dtype: {target_dtype}",
+            f"Total tensors: {total}",
+            f"Converted floating tensors: {counts['converted']}",
+            f"Excluded tensors preserved: {counts['excluded']}",
+            f"Already target dtype: {counts['already_target']}",
+            f"Reference-mapped floating tensors: {counts['reference_mapped']}",
+            f"Reference dtype already matched: {counts['reference_already_matched']}",
+            f"Reference-missing fallback tensors: {counts['reference_missing_fallback']}",
+            f"Non-floating tensors preserved: {counts['non_floating']}",
+            f"Original dtype counts: {dtype_summary or 'None'}",
+        )
+    )
+    return output_name, report
+
+
+class DiffusionModelDtypeConversion(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="DiffusionModelDtypeConversion",
+            display_name="Diffusion Model Dtype Conversion",
+            category="ModelUtils/Conversion",
+            description=(
+                "Streams a diffusion-model safetensors file into fp32, fp16, "
+                "or bf16. Floating tensors matching an exclusion regex retain "
+                "their original dtype, or enable Include Mode to convert only matches; "
+                "non-floating tensors are always preserved."
+            ),
+            inputs=[
+                io.Combo.Input(
+                    "model_name",
+                    options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Diffusion-model file whose floating tensors will be converted.",
+                ),
+                io.Combo.Input(
+                    "target_dtype", options=["fp32", "fp16", "bf16"],
+                    tooltip="Target dtype for floating tensors selected by the layer filter.",
+                ),
+                io.Combo.Input(
+                    "reference_model",
+                    options=["None", *folder_paths.get_filename_list("diffusion_models")],
+                    default="None",
+                    tooltip="Optional authoritative dtype map. Matching floating tensor keys use this model's header dtype without loading its tensors; missing keys use Target Dtype.",
+                ),
+                io.String.Input(
+                    "exclude_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip=(
+                        "Optional Python regex patterns, one per line. Matching "
+                        "tensor keys retain their original dtype, or are the only keys "
+                        "converted when Include Mode is enabled."
+                    ),
+                ),
+                io.String.Input(
+                    "output_filename", default="converted_model",
+                    tooltip="Output filename without extension, written under ComfyUI's diffusion-model directory.",
+                ),
+                io.Boolean.Input(
+                    "include_mode", default=False,
+                    tooltip="Use Exclude Patterns as an include-only filter. Only matching tensors are converted; an empty filter converts nothing. Nonmatches retain their original dtype.",
+                ),
+            ],
+            outputs=[
+                io.AnyType.Output(display_name="output_path"),
+                io.String.Output(display_name="conversion_report"),
+            ],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        model_name: str,
+        target_dtype: str,
+        reference_model: str,
+        exclude_patterns: str,
+        output_filename: str,
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        output_path, report = convert_diffusion_model_dtype(
+            model_name,
+            target_dtype,
+            exclude_patterns,
+            output_filename,
+            reference_model,
+            include_mode,
+        )
+        return io.NodeOutput(output_path, report)

nodes/extraction_stream.py

@@ -0,0 +1,179 @@
+"""Bounded paired UEL streaming helpers for model extraction nodes."""
+
+import torch
+
+from .uel_io import AsyncTensorCursor
+
+def is_cuda_oom(error):
+    return isinstance(error, torch.cuda.OutOfMemoryError) or (
+        isinstance(error, RuntimeError)
+        and "cuda" in str(error).lower()
+        and "out of memory" in str(error).lower()
+    )
+
+
+def release_cuda_oom():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def move_source_tensors(cpu_a, cpu_b, device, dtype=torch.float32):
+    """Move one source pair, retrying a confirmed CUDA OOM on CPU."""
+    use_cuda = str(device).startswith("cuda")
+    try:
+        if use_cuda:
+            tensor_a = cpu_a.to(device=device, dtype=dtype, non_blocking=True)
+            tensor_b = cpu_b.to(device=device, dtype=dtype, non_blocking=True) if cpu_b is not None else None
+        else:
+            tensor_a = cpu_a.to(device=device, dtype=dtype)
+            tensor_b = cpu_b.to(device=device, dtype=dtype) if cpu_b is not None else None
+        return tensor_a, tensor_b, device
+    except BaseException as error:
+        if not use_cuda or not is_cuda_oom(error):
+            raise
+        release_cuda_oom()
+        return (
+            cpu_a.to(device="cpu", dtype=dtype),
+            cpu_b.to(device="cpu", dtype=dtype) if cpu_b is not None else None,
+            "cpu",
+        )
+
+
+def retry_cuda_oom_on_cpu(cuda_call, cpu_call, device):
+    """Retry one compute operation on CPU only for a confirmed CUDA OOM."""
+    try:
+        return cuda_call()
+    except BaseException as error:
+        if not str(device).startswith("cuda") or not is_cuda_oom(error):
+            raise
+        release_cuda_oom()
+        return cpu_call()
+
+
+def lora_difference(cpu_a, cpu_b, device):
+    """Construct A minus B with transfer and arithmetic CUDA OOM fallback."""
+    tensor_a, tensor_b, actual_device = move_source_tensors(
+        cpu_a, cpu_b, device
+    )
+    if tensor_b is None:
+        return tensor_a, actual_device
+    try:
+        return tensor_a - tensor_b, actual_device
+    except BaseException as error:
+        if not str(actual_device).startswith("cuda") or not is_cuda_oom(error):
+            raise
+        del tensor_a, tensor_b
+        release_cuda_oom()
+        return cpu_a.float() - cpu_b.float(), "cpu"
+
+
+def dora_difference(cpu_a, cpu_b, device, *, keep_sources=False):
+    """Construct a directional DoRA delta with complete CUDA OOM fallback."""
+    tensor_a, tensor_b, actual_device = move_source_tensors(
+        cpu_a, cpu_b, device
+    )
+
+    def compute(a, b):
+        if b is None:
+            base = torch.zeros_like(a)
+            result = (a, None, base, a)
+            return result if keep_sources else result[:2]
+        if a.ndim < 2:
+            result = (None, None, b, a)
+            return result if keep_sources else result[:2]
+        is_conv = a.ndim == 4
+        if is_conv:
+            out_channels = a.shape[0]
+            norm_a = torch.linalg.norm(a.reshape(out_channels, -1), dim=1, keepdim=True)
+            norm_b = torch.linalg.norm(b.reshape(out_channels, -1), dim=1, keepdim=True)
+        else:
+            norm_a = torch.linalg.norm(a, dim=1, keepdim=True)
+            norm_b = torch.linalg.norm(b, dim=1, keepdim=True)
+        ratio = norm_b / (norm_a + 1e-8)
+        if is_conv:
+            ratio = ratio.view(-1, 1, 1, 1)
+            scale = norm_a.view(-1, 1, 1, 1)
+        else:
+            scale = norm_a
+        result = (a * ratio - b, scale, b, a)
+        return result if keep_sources else result[:2]
+
+    try:
+        return (*compute(tensor_a, tensor_b), actual_device)
+    except BaseException as error:
+        if not str(actual_device).startswith("cuda") or not is_cuda_oom(error):
+            raise
+        del tensor_a, tensor_b
+        release_cuda_oom()
+        cpu_b_float = cpu_b.float() if cpu_b is not None else None
+        return (*compute(cpu_a.float(), cpu_b_float), "cpu")
+
+
+def paired_async_tensors(
+    handler_a,
+    handler_b,
+    work_units,
+    *,
+    pin_memory: bool,
+    quantization_a=None,
+    quantization_b=None,
+    compute_dtype=torch.float32,
+):
+    """Yield logical A/B tensors and release their source keys per work unit."""
+    def source_keys(quantization, logical_key):
+        if logical_key is None:
+            return ()
+        return tuple(
+            quantization.required_keys(logical_key)
+            if quantization is not None
+            else (logical_key,)
+        )
+
+    required_a = [source_keys(quantization_a, unit[1]) for unit in work_units]
+    required_b = [source_keys(quantization_b, unit[2]) for unit in work_units]
+    keys_a = [key for keys in required_a for key in keys]
+    keys_b = [key for keys in required_b for key in keys]
+    cursor_a = AsyncTensorCursor(handler_a, keys_a, pin_memory=pin_memory)
+    cursor_b = AsyncTensorCursor(handler_b, keys_b, pin_memory=pin_memory)
+    try:
+        for (logical_key, key_a, key_b), keys_for_a, keys_for_b in zip(
+            work_units, required_a, required_b
+        ):
+            tensors_a = {}
+            tensors_b = {}
+            consumed_a = []
+            consumed_b = []
+            tensor_a = tensor_b = None
+            try:
+                for source_key in keys_for_a:
+                    tensors_a[source_key] = cursor_a.take(source_key)
+                    consumed_a.append(source_key)
+                for source_key in keys_for_b:
+                    tensors_b[source_key] = cursor_b.take(source_key)
+                    consumed_b.append(source_key)
+                tensor_a = (
+                    quantization_a.decode(
+                        key_a, tensors_a, compute_dtype, device="cpu"
+                    )
+                    if quantization_a is not None and key_a in quantization_a.quantized_keys
+                    else tensors_a.get(key_a)
+                )
+                tensor_b = (
+                    quantization_b.decode(
+                        key_b, tensors_b, compute_dtype, device="cpu"
+                    )
+                    if quantization_b is not None and key_b in quantization_b.quantized_keys
+                    else tensors_b.get(key_b)
+                )
+                yield logical_key, tensor_a, tensor_b
+            finally:
+                del tensor_a, tensor_b, tensors_a, tensors_b
+                for source_key in consumed_a:
+                    cursor_a.release(source_key)
+                for source_key in consumed_b:
+                    cursor_b.release(source_key)
+        cursor_a.finish()
+        cursor_b.finish()
+    finally:
+        cursor_a.close()
+        cursor_b.close()

nodes/layer_parameters.py

@@ -0,0 +1,284 @@
+"""Parsed per-layer numeric overrides shared by merge, extraction and resize."""
+
+from dataclasses import dataclass
+import fnmatch
+import math
+import re
+
+from comfy_api.latest import io
+
+
+LAYER_PARAMETERS = io.Custom("MODELUTILS_LAYER_PARAMETERS")
+
+
+@dataclass(frozen=True)
+class LayerParameterRule:
+    line: int
+    pattern: str
+    matcher: re.Pattern
+    assignments: tuple[tuple[str, int | float], ...]
+
+
+@dataclass(frozen=True)
+class LayerParameterRules:
+    rules: tuple[LayerParameterRule, ...]
+    glob_patterns: bool = False
+
+
+@dataclass(frozen=True)
+class Parameter:
+    name: str
+    aliases: tuple[str, ...]
+    minimum: int | float
+    maximum: int | float
+    integral: bool = False
+
+
+MERGE_COEFFICIENTS = ("alpha", "beta", "gamma", "delta", "epsilon", "zeta")
+MERGE_MODES = {
+    "Weight-Sum": ("alpha",),
+    "Train-Difference": ("alpha",),
+    "Add-Difference": ("alpha", "beta"),
+    "Power-Up (DARE)": ("alpha", "beta"),
+    "Comparative-Interpolation": ("alpha", "beta", "gamma"),
+    "Add-Dissimilarities": ("alpha", "beta", "gamma"),
+    "SVD LoRA Extraction": ("alpha", "beta", "gamma"),
+    "Enhanced Auto Interp": ("alpha", "beta", "gamma"),
+    "Weight-Sum Cutoff": ("alpha", "beta", "gamma"),
+    "Power-Up (DARE+TIES)": ("alpha", "beta", "gamma"),
+    "Extract-Features": ("alpha", "beta", "gamma", "delta"),
+    "Enhanced Man Interp": ("alpha", "beta", "gamma", "delta"),
+    "Power-Up Enhanced (DARE)": ("alpha", "beta", "gamma", "delta"),
+    "Power-Up Enhanced (DARE+TIES)": ("beta", "gamma", "delta", "epsilon", "zeta"),
+}
+EXTRACTION_BINDINGS = {
+    "fixed": ("linear_dim", "conv_dim", "clamp_quantile", "min_diff"),
+    "ratio": ("linear_ratio", "conv_ratio", "clamp_quantile", "min_diff", "linear_max_rank", "conv_max_rank"),
+    "quantile": ("linear_quantile", "conv_quantile", "clamp_quantile", "min_diff", "linear_max_rank", "conv_max_rank"),
+    "frobenius": ("linear_target", "conv_target", "clamp_quantile", "min_diff", "linear_max_rank", "conv_max_rank"),
+    "knee": ("linear_max_rank", "conv_max_rank", "clamp_quantile", "min_diff"),
+}
+RESIZE_BINDINGS = {
+    "fixed": ("new_rank",),
+    "ratio": ("max_rank", "ratio"),
+    "frobenius": ("max_rank", "min_rank", "target"),
+    "cumulative": ("max_rank", "target"),
+}
+_ASSIGNMENT = re.compile(r"([a-z][a-z0-9_]*)\s*:\s*([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)")
+_SEPARATOR = re.compile(r"(?:\s*,\s*|\s+)")
+
+
+def parse_rules(text: str, glob_patterns: bool = False) -> LayerParameterRules:
+    rules = []
+    for line_number, raw in enumerate(text.splitlines(), 1):
+        line = raw.strip()
+        if not line or line.startswith("#"):
+            continue
+        row = re.fullmatch(r"\((.*)\)\s+(.+)", line)
+        if row is None or not row[1].strip():
+            raise ValueError(f"Layer parameters line {line_number}: expected (pattern) name:value assignments.")
+        pattern, values = row.groups()
+        try:
+            matcher = re.compile(fnmatch.translate(f"*{pattern}*") if glob_patterns else pattern)
+        except re.error as exc:
+            raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: invalid regex: {exc}") from exc
+        assignments = []
+        seen = set()
+        position = 0
+        while position < len(values):
+            match = _ASSIGNMENT.match(values, position)
+            if match is None:
+                raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: invalid assignment near {values[position:]!r}.")
+            name, literal = match.groups()
+            if name in seen:
+                raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: duplicate parameter {name!r}.")
+            try:
+                value = float(literal) if any(c in literal for c in ".eE") else int(literal)
+                finite = math.isfinite(value)
+            except (ValueError, OverflowError):
+                finite = False
+            if not finite:
+                raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: {name} must be finite.")
+            assignments.append((name, value))
+            seen.add(name)
+            position = match.end()
+            if position == len(values):
+                break
+            separator = _SEPARATOR.match(values, position)
+            if separator is None or separator.end() == len(values):
+                raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: expected another assignment after a comma or whitespace.")
+            position = separator.end()
+        rules.append(LayerParameterRule(line_number, pattern, matcher, tuple(assignments)))
+    return LayerParameterRules(tuple(rules), glob_patterns)
+
+
+def parameter_bindings(profile: str) -> tuple[Parameter, ...]:
+    family, _, mode = profile.partition(":")
+    if family == "merge" and mode in MERGE_MODES:
+        parameters = []
+        probabilities = {
+            "Power-Up (DARE)": {"alpha"},
+            "Power-Up (DARE+TIES)": {"alpha", "beta"},
+            "Power-Up Enhanced (DARE)": {"beta", "gamma"},
+            "Power-Up Enhanced (DARE+TIES)": {"beta", "epsilon", "zeta"},
+            "SVD LoRA Extraction": {"gamma"},
+        }.get(mode, set())
+        for name in MERGE_MODES[mode]:
+            aliases = ("abcdef"[MERGE_COEFFICIENTS.index(name)],)
+            if name == "gamma":
+                aliases += ("g",)
+            if name == "zeta":
+                aliases += ("z",)
+            integral = mode == "SVD LoRA Extraction" and name in {"alpha", "beta"}
+            minimum, maximum = (0, 1) if name in probabilities else (-10, 10)
+            if integral:
+                minimum = 1
+            parameters.append(Parameter(name, aliases, minimum, maximum, integral))
+        return tuple(parameters)
+    if family == "extract" and mode in EXTRACTION_BINDINGS:
+        names = EXTRACTION_BINDINGS[mode]
+    elif family == "resize" and mode in RESIZE_BINDINGS:
+        names = RESIZE_BINDINGS[mode]
+    else:
+        raise ValueError(f"Unsupported layer parameter bindings: {profile!r}")
+    parameters = []
+    for alias, name in zip("abcdef", names):
+        integral = "rank" in name or name.endswith("_dim")
+        if integral:
+            minimum, maximum = 1, 16384 if family == "extract" else 3072
+        elif "ratio" in name:
+            minimum, maximum = 1, 100
+        elif name == "clamp_quantile":
+            minimum, maximum = 0.5, 1
+        else:
+            minimum, maximum = (0.1 if family == "resize" else 0), 1
+        parameters.append(Parameter(name, (alias,), minimum, maximum, integral))
+    return tuple(parameters)
+
+
+def _validated_value(parameter, value, context):
+    if not isinstance(value, (int, float)) or isinstance(value, bool) or not math.isfinite(value):
+        raise ValueError(f"{context}: {parameter.name} must be a finite number, got {value!r}.")
+    if parameter.integral and int(value) != value:
+        raise ValueError(f"{context}: {parameter.name} must be integral, got {value}.")
+    if not parameter.minimum <= value <= parameter.maximum:
+        raise ValueError(f"{context}: {parameter.name} must be in [{parameter.minimum}, {parameter.maximum}], got {value}.")
+    return int(value) if parameter.integral else value
+
+
+def resolve_layer_parameters(layer_parameters, profile, layers, defaults, *, node_name):
+    """Validate against header-only target names and return matched scalar overrides."""
+    if layer_parameters is None:
+        return {}
+    if not isinstance(layer_parameters, LayerParameterRules):
+        raise TypeError(f"{node_name}: layer_parameters must come from Layer Parameter Configuration.")
+    if not layer_parameters.rules:
+        return {}
+    bindings = parameter_bindings(profile)
+    aliases = {alias: parameter for parameter in bindings for alias in (*parameter.aliases, parameter.name)}
+    assignments = []
+    for rule in layer_parameters.rules:
+        resolved = {}
+        for name, value in rule.assignments:
+            parameter = aliases.get(name)
+            context = f"{node_name}: line {rule.line}, pattern {rule.pattern!r}"
+            if parameter is None:
+                allowed = ", ".join(p.name for p in bindings)
+                raise ValueError(f"{context}: unsupported parameter {name!r} for {profile}; allowed: {allowed}.")
+            if parameter.name in resolved:
+                raise ValueError(f"{context}: duplicate aliases for {parameter.name}.")
+            resolved[parameter.name] = _validated_value(parameter, value, context)
+        assignments.append(resolved)
+    hits = set()
+    result = {}
+    for layer in sorted(set(layers)):
+        matched = [i for i, rule in enumerate(layer_parameters.rules) if rule.matcher.search(layer)]
+        if len(matched) > 1:
+            lines = "; ".join(f"line {layer_parameters.rules[i].line} ({layer_parameters.rules[i].pattern!r})" for i in matched)
+            raise ValueError(f"{node_name}: layer {layer!r} matches conflicting rules: {lines}.")
+        if not matched:
+            continue
+        index = matched[0]
+        values = dict(defaults)
+        values.update(assignments[index])
+        if profile == "extract:fixed":
+            for dimension, cap in (("linear_dim", "linear_max_rank"), ("conv_dim", "conv_max_rank")):
+                if dimension in assignments[index] and cap in defaults:
+                    values[cap] = values[dimension]
+        rule = layer_parameters.rules[index]
+        context = f"{node_name}: layer {layer!r}, line {rule.line}, pattern {rule.pattern!r}"
+        for parameter in bindings:
+            if parameter.name in values:
+                values[parameter.name] = _validated_value(parameter, values[parameter.name], context)
+        if profile == "resize:frobenius" and values["min_rank"] > values["max_rank"]:
+            raise ValueError(f"{node_name}: layer {layer!r}, line {rule.line}, pattern {rule.pattern!r}: min_rank={values['min_rank']} exceeds max_rank={values['max_rank']}.")
+        result[layer] = values
+        hits.add(index)
+    unmatched = [rule for i, rule in enumerate(layer_parameters.rules) if i not in hits]
+    if unmatched:
+        details = "; ".join(f"line {rule.line}, pattern {rule.pattern!r}" for rule in unmatched)
+        raise ValueError(f"{node_name}: no target layer matches {details}.")
+    return result
+
+
+def _binding_description(profile):
+    return "; ".join(f"{'/'.join(p.aliases)}={p.name}" for p in parameter_bindings(profile))
+
+
+def parameter_input(profile):
+    if profile == "merge":
+        parameters = {p.name: p for mode in MERGE_MODES for p in parameter_bindings(f"merge:{mode}")}
+        description = "; ".join(f"{'/'.join(parameters[name].aliases)}={name}" for name in MERGE_COEFFICIENTS)
+        description += ". Only coefficients used by Calc Mode are accepted."
+    else:
+        description = _binding_description(profile)
+    return LAYER_PARAMETERS.Input(
+        "layer_parameters", optional=True,
+        tooltip="Optional Layer Parameter Configuration. " + description + " Full names are also accepted. Unassigned values use this node's settings; existing filters still apply.",
+    )
+
+
+def documentation():
+    lines = [
+        "# Per-layer parameter rules", "",
+        "Connect Layer Parameter Configuration to the optional layer_parameters socket on standard two/three-input mergers, all 25 extraction variants, or the four LoRA resizers.", "",
+        "Write (pattern) a:value b:value, or use full parameter names. Commas and whitespace may separate assignments; whitespace around colons is allowed. Signed decimal/scientific numbers are accepted. Blank lines and whole-line # comments are ignored. Inline comments are not supported.", "",
+        "Patterns preserve backslashes, nested groups, spaces and character classes. Regex uses substring search; Glob Patterns uses case-sensitive glob substring matching. These settings do not change the main node's filters.", "",
+        "Unconnected or empty configuration changes nothing. Unmatched layers and omitted parameters use main-node settings. Existing exclusions, inclusion filters, discards, mismatch rules and safety guards still apply.", "",
+        "Invalid syntax/values, duplicate assignments (including aliases), unsupported parameters, overlapping rule lines and rules matching no target layer are errors. Matching and validation happen before tensor streaming or output writing. Name matches are counted before exclusions.", "",
+        "LoRA rules target normalized logical block names, not factor/alpha suffixes. Existing dotted normalization and reference mapping are retained, including underscores inside names such as qkv_proj. Unresolved flattened names retain the existing fallback; patterns are not silently rewritten.", "",
+        "## Bindings and bounds", "", "| Receiver/method | Short aliases and full names | Bounds |", "|---|---|---|",
+    ]
+    profiles = [*(f"merge:{name}" for name in MERGE_MODES), *(f"extract:{name}" for name in EXTRACTION_BINDINGS), *(f"resize:{name}" for name in RESIZE_BINDINGS)]
+    for profile in profiles:
+        bounds = "; ".join(f"{p.name}: {'integer ' if p.integral else ''}[{p.minimum}, {p.maximum}]" for p in parameter_bindings(profile))
+        lines.append(f"| {profile} | {_binding_description(profile)} | {bounds} |")
+    lines.extend([
+        "", "Frobenius resize requires resolved min_rank <= max_rank. Existing tensor-dimension rank caps remain. Fixed extraction dimension overrides also update their corresponding rank caps. Probe offsets, modes, iterations, optimizer settings, seeds, files, dtype and device remain main-node settings.",
+        "", "## Examples", "", "Standard Extract-Features merge:", "", "```text",
+        r"(blocks\.4[589]\.attn\.qkv_proj) a:0.5 b:0.25 g:0.75 d:1.0", "```", "",
+        "Fixed extraction:", "", "```text",
+        r"(blocks\.4[589]\.attn\.qkv_proj) a:64, b:32, c:0.99, d:0.0", "```", "",
+        "Frobenius resize:", "", "```text",
+        r"(blocks\.4[589]\.attn\.qkv_proj) max_rank:128 min_rank:1 target:0.9", "```",
+    ])
+    return "\n".join(lines)
+
+
+class LayerParameterConfiguration(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LayerParameterConfiguration", display_name="Layer Parameter Configuration",
+            category="ModelUtils/Configuration",
+            description="Configure per-layer numeric overrides for standard merging, extraction and LoRA resizing.",
+            inputs=[
+                io.String.Input("rules", default="", multiline=True, tooltip="One (pattern) name:value rule per line. Short aliases or full names are accepted; separate assignments with whitespace or commas. See Documentation for receiver-specific bindings."),
+                io.Boolean.Input("glob_patterns", default=False, tooltip="Use case-sensitive glob substring matching instead of regex substring search. Independent of receiving nodes' exclusion-filter syntax."),
+            ],
+            outputs=[LAYER_PARAMETERS.Output(display_name="layer_parameters"), io.String.Output(display_name="documentation")],
+        )
+
+    @classmethod
+    def execute(cls, rules="", glob_patterns=False):
+        return io.NodeOutput(parse_rules(rules, glob_patterns), documentation())

nodes/lodestone_merger.py

@@ -0,0 +1,313 @@
+import fnmatch
+import gc
+import re
+
+import comfy.utils
+import folder_paths
+import torch
+from comfy_api.latest import io
+from tqdm import tqdm
+from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
+
+from .artifact_paths import canonical_model_artifact_path
+from .device_utils import cleanup_after_operation, estimate_model_size, prepare_for_large_operation
+from .lora_alpha import normalize_lora_pair
+from .lora_resize import build_lora_layer_map, canonical_lora_key, parse_lora_layers, select_output_dtype, validate_canonical_blocks
+from .quantization_guard import inspect_low_bit_input
+from .uel_io import AsyncTensorCursor, atomic_uel_writer
+
+
+LODESTONE_METHODS = ("sum", "mean", "slotnorm", "normmatch", "slotnorm-normmatch")
+_ALLOWED_ROLES = {
+    "down", "up", "alpha", "diff", "format", "down_suffix", "up_suffix",
+    "alpha_suffix",
+}
+
+
+def _matches(key, patterns, glob_mode):
+    entries = [entry.strip() for entry in patterns.splitlines() if entry.strip()]
+    if glob_mode:
+        return any(fnmatch.fnmatchcase(key, entry) for entry in entries)
+    return any(re.search(entry, key) is not None for entry in entries)
+
+
+def factor_delta(down, up, layer):
+    if down.ndim == 2 and up.ndim == 2:
+        if down.shape[0] != up.shape[1]:
+            raise ValueError(f"Lodestone rank mismatch for '{layer}'.")
+        return up.float() @ down.float()
+    if down.ndim == 4 and up.ndim == 4 and tuple(up.shape[2:]) == (1, 1):
+        rank, channels, height, width = down.shape
+        if up.shape[1] != rank:
+            raise ValueError(f"Lodestone rank mismatch for '{layer}'.")
+        delta = up.reshape(up.shape[0], rank).float() @ down.reshape(rank, -1).float()
+        return delta.reshape(up.shape[0], channels, height, width)
+    raise ValueError(
+        f"Lodestone layer '{layer}' requires linear factors or convolutional "
+        "factors with a 1x1 up kernel."
+    )
+
+
+def merge_lodestone_deltas(deltas, method):
+    if method not in LODESTONE_METHODS:
+        raise ValueError(f"Unknown Lodestone method '{method}'.")
+    if not deltas:
+        raise ValueError("Lodestone requires at least one delta.")
+    shape = tuple(deltas[0].shape)
+    if any(tuple(delta.shape) != shape for delta in deltas[1:]):
+        raise ValueError("Lodestone source deltas must have matching shapes.")
+    norms = torch.stack([torch.linalg.vector_norm(delta.float()) for delta in deltas])
+    target = torch.median(norms)
+    values = deltas
+    if method in {"slotnorm", "slotnorm-normmatch"}:
+        values = [
+            delta * (target / norm.clamp_min(1e-12))
+            for delta, norm in zip(deltas, norms)
+        ]
+    merged = torch.stack(values).sum(dim=0)
+    if method in {"mean", "slotnorm"}:
+        merged /= len(values)
+    elif method in {"normmatch", "slotnorm-normmatch"}:
+        merged_norm = torch.linalg.vector_norm(merged.float())
+        if merged_norm > 0:
+            merged *= target / merged_norm.clamp_min(1e-12)
+    return merged
+
+
+def _to_device(tensor, device):
+    if device.type == "cuda":
+        return transfer_to_gpu_pinned(tensor, device, torch.float32)
+    return tensor.to(device=device, dtype=torch.float32)
+
+
+def _release_cuda_oom():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _is_cuda_oom(error):
+    return isinstance(error, torch.cuda.OutOfMemoryError) or (
+        isinstance(error, RuntimeError) and "out of memory" in str(error).lower()
+    )
+
+
+def _build_units(layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns, include_mode=False):
+    units = []
+    stream_keys = [[] for _ in handlers]
+    for core, mapped in layer_map.items():
+        if not any(index == 0 for index, _ in mapped):
+            continue
+        output_key = canonical_lora_key(core, "diff")
+        if _matches(output_key, discard_patterns, glob_patterns):
+            continue
+        matched = _matches(output_key, exclude_patterns, glob_patterns)
+        preserve = not matched if include_mode else matched
+        sources = []
+        for index, block in sorted(mapped):
+            roles = infos[index]["pairs"][block]
+            unsupported = set(roles) - _ALLOWED_ROLES
+            if unsupported:
+                raise ValueError(f"Lodestone layer '{block}' has unsupported roles: {', '.join(sorted(unsupported))}.")
+            if "down" in roles or "up" in roles:
+                if not all(role in roles for role in ("down", "up")):
+                    raise ValueError(f"Lodestone layer '{block}' has an incomplete factor pair.")
+                keys = [roles["down"], roles["up"]]
+                if any(key in infos[index]["low_bit"] for key in keys):
+                    raise ValueError(f"Lodestone cannot expand low-bit factors for '{block}'.")
+                if "alpha" in roles:
+                    keys.append(roles["alpha"])
+                sources.append((index, block, "pair", keys))
+            elif "diff" in roles:
+                key = roles["diff"]
+                if key in infos[index]["low_bit"]:
+                    raise ValueError(f"Lodestone cannot merge low-bit delta '{block}'.")
+                sources.append((index, block, "diff", [key]))
+        if preserve:
+            sources = [source for source in sources if source[0] == 0]
+        present = {source[0] for source in sources}
+        if len(present) < len(handlers) and not preserve:
+            if mismatch_mode == "error":
+                raise ValueError(f"Lodestone layer '{core}' is missing from one or more inputs.")
+            if mismatch_mode == "skip":
+                sources = [source for source in sources if source[0] == 0]
+        if not sources:
+            continue
+        dtypes = []
+        for index, _block, kind, keys in sources:
+            stream_keys[index].extend(keys)
+            dtypes.append(handlers[index].get_dtype(keys[0] if kind == "diff" else keys[1]))
+        units.append({
+            "core": core,
+            "output_key": output_key,
+            "sources": sources,
+            "missing": 0 if preserve or mismatch_mode != "zeros" else len(handlers) - len(present),
+            "dtypes": dtypes,
+        })
+    return units, stream_keys
+
+
+def _process_on_device(unit, loaded, method, save_dtype, device):
+    deltas = []
+    for index, block, kind, keys in unit["sources"]:
+        if kind == "diff":
+            delta = _to_device(loaded[(index, keys[0])], device)
+        else:
+            down = loaded[(index, keys[0])]
+            up = loaded[(index, keys[1])]
+            alpha = loaded[(index, keys[2])] if len(keys) == 3 else None
+            down, up = normalize_lora_pair(down, up, alpha, layer=block)
+            delta = factor_delta(_to_device(down, device), _to_device(up, device), block)
+        deltas.append(delta)
+    if unit["missing"]:
+        deltas.extend(torch.zeros_like(deltas[0]) for _ in range(unit["missing"]))
+    merged = merge_lodestone_deltas(deltas, method)
+    dtype = select_output_dtype(unit["dtypes"], save_dtype, is_1d_diff=merged.ndim == 1)
+    return merged.to(dtype).cpu().contiguous()
+
+
+def _process(unit, loaded, method, save_dtype, requested_device):
+    device = torch.device(requested_device)
+    try:
+        return _process_on_device(unit, loaded, method, save_dtype, device), False
+    except Exception as error:
+        if device.type != "cuda" or not _is_cuda_oom(error):
+            raise
+        _release_cuda_oom()
+        return _process_on_device(unit, loaded, method, save_dtype, torch.device("cpu")), True
+
+
+def merge_lodestone_loras(lora_paths, method, device, save_dtype, output_filename,
+                          mismatch_mode="skip", exclude_patterns="", discard_patterns="",
+                          glob_patterns=False, verbose=True, include_mode=False):
+    operation = "Lodestone LoRA Merge"
+    prepare_for_large_operation(sum(estimate_model_size(path) for path in lora_paths) * 2.5, torch.device(device))
+    handlers = [MemoryEfficientSafeOpen(path, low_memory=True) for path in lora_paths]
+    cursors = []
+    try:
+        infos = []
+        for index, handler in enumerate(handlers):
+            low_bit = inspect_low_bit_input(handler, f"LoRA {index + 1}", operation)
+            pairs, passthrough = parse_lora_layers(handler.keys())
+            infos.append({"pairs": pairs, "passthrough": passthrough, "low_bit": low_bit})
+        layer_map = build_lora_layer_map(infos, None)
+        validate_canonical_blocks(layer_map, operation)
+        units, stream_keys = _build_units(
+            layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns, include_mode
+        )
+        cursors = [
+            AsyncTensorCursor(handler, keys, pin_memory=torch.device(device).type == "cuda")
+            for handler, keys in zip(handlers, stream_keys)
+        ]
+        output_path, _ = canonical_model_artifact_path("loras", output_filename)
+        metadata = (handlers[0].metadata() or {}).copy()
+        metadata.update({"merge_method": f"lodestone:{method}", "alpha_normalized": "true", "output_representation": "full_direct_difference"})
+        fallback_count = 0
+        pbar = comfy.utils.ProgressBar(len(units))
+        with atomic_uel_writer(output_path, metadata) as writer, torch.no_grad():
+            for unit in tqdm(units, desc=f"Lodestone layers ({method})", unit="layers"):
+                loaded = {}
+                try:
+                    for index, _block, _kind, keys in unit["sources"]:
+                        for key in keys:
+                            loaded[(index, key)] = cursors[index].take(key)
+                    output, fallback = _process(unit, loaded, method, save_dtype, device)
+                    writer.write_batch([(unit["output_key"], output)])
+                    fallback_count += int(fallback)
+                    del output
+                finally:
+                    for index, key in list(loaded):
+                        del loaded[(index, key)]
+                        cursors[index].release(key)
+                    loaded.clear()
+                pbar.update(1)
+        for cursor in cursors:
+            cursor.finish()
+        if verbose:
+            print(f"[{operation}] CUDA OOM CPU fallbacks: {fallback_count}")
+            print(f"[{operation}] Saved {len(units)} full-delta layers to {output_path}")
+        return output_path
+    finally:
+        for cursor in cursors:
+            cursor.close()
+        for handler in handlers:
+            handler.__exit__(None, None, None)
+        cleanup_after_operation()
+
+
+def _controls(default_filename):
+    return [
+        io.Combo.Input("calc_mode", options=list(LODESTONE_METHODS), default="sum", tooltip="How full LoRA weight changes are combined per layer. sum adds them; mean averages them; slotnorm equalizes each input's Frobenius magnitude before averaging; normmatch adds them and scales the result to a typical input magnitude; slotnorm-normmatch performs both normalization steps."),
+        io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", tooltip="What to do when a layer exists in LoRA 1 but is absent from another input. skip copies LoRA 1's layer unchanged; zeros treats each missing input as a zero update; error stops without replacing an existing output file."),
+        io.String.Input("output_filename", default=default_filename, tooltip="Name of the new .safetensors file. It is saved in ComfyUI's loras folder; omit the extension."),
+        io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"], default="bf16", tooltip="Precision used to save the expanded full-weight difference tensors. FP32 is largest and most precise; FP16 is smallest; BF16 has wider numeric range than FP16."),
+        io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", tooltip="Where each full layer is expanded and merged. CUDA is faster; if a layer runs out of VRAM, that layer is automatically retried on CPU."),
+        io.String.Input("exclude_patterns", default="", multiline=True, tooltip="Optional layer-name patterns to copy from LoRA 1 instead of merging. Enter one pattern per line; useful for protecting specific blocks."),
+        io.String.Input("discard_patterns", default="", multiline=True, tooltip="Optional layer-name patterns to leave out of the saved LoRA completely. Enter one pattern per line; discarded layers cannot affect the model."),
+        io.Boolean.Input("glob_patterns", default=False, tooltip="Choose filter syntax. Off uses regular expressions; on uses shell-style globs where * matches any text and ? matches one character."),
+        io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from LoRA 1."),
+    ]
+
+
+def _execute(kwargs, count):
+    paths = [folder_paths.get_full_path_or_raise("loras", kwargs[f"lora_{index}"]) for index in range(1, count + 1)]
+    dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[kwargs["save_dtype"]]
+    merge_lodestone_loras(paths, kwargs["calc_mode"], kwargs["process_device"], dtype,
+                          kwargs["output_filename"], kwargs["mismatch_mode"],
+                          kwargs["exclude_patterns"], kwargs["discard_patterns"], kwargs["glob_patterns"], include_mode=kwargs.get("include_mode", False))
+    return canonical_model_artifact_path("loras", kwargs["output_filename"])[1]
+
+
+class LodestoneLoRATwoMerger(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        options = folder_paths.get_filename_list("loras")
+        return io.Schema(node_id="LodestoneLoRATwoMerger", display_name="Lodestone Merge LoRAs (2)",
+                         category="ModelUtils/LoRA/Merge/Lodestone",
+                         description="Expand two LoRAs into full per-layer weight changes, combine them with the selected Frobenius-norm method, and save a full-difference LoRA.",
+                         inputs=[io.Combo.Input("lora_1", options=options, tooltip="First LoRA to merge. It also supplies metadata and is the layer kept when Missing Layer Handling is set to skip."), io.Combo.Input("lora_2", options=options, tooltip="Second LoRA to merge. It contributes equally with LoRA 1 wherever both contain the same layer."), *_controls("lodestone_merged_2_lora")],
+                         outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True)
+
+    @classmethod
+    def execute(cls, **kwargs):
+        return io.NodeOutput(_execute(kwargs, 2))
+
+
+class LodestoneLoRAThreeMerger(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        options = folder_paths.get_filename_list("loras")
+        return io.Schema(node_id="LodestoneLoRAThreeMerger", display_name="Lodestone Merge LoRAs (3)",
+                         category="ModelUtils/LoRA/Merge/Lodestone",
+                         description="Expand three LoRAs into full per-layer weight changes, combine them with the selected Frobenius-norm method, and save a full-difference LoRA.",
+                         inputs=[io.Combo.Input("lora_1", options=options, tooltip="First LoRA to merge. It supplies metadata, defines the output layer set, and is kept when Missing Layer Handling is skip."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-strength LoRA contribution for every matching layer."), io.Combo.Input("lora_3", options=options, tooltip="Third equal-strength LoRA contribution for every matching layer."), *_controls("lodestone_merged_3_lora")],
+                         outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True)
+
+    @classmethod
+    def execute(cls, **kwargs):
+        return io.NodeOutput(_execute(kwargs, 3))
+
+
+class LodestoneLoRAMultiMerger(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        options = folder_paths.get_filename_list("loras")
+        optional = ["None", *options]
+        inputs = [io.Combo.Input("lora_count", options=[str(index) for index in range(2, 9)], default="2", tooltip="Number of consecutive LoRA selectors to merge, starting at LoRA 1. Every included LoRA has equal influence before the selected normalization method."),
+                  io.Combo.Input("lora_1", options=options, tooltip="First LoRA to merge. It supplies metadata, defines the output layer set, and anchors missing-layer handling."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-strength LoRA contribution. LoRA Count always includes this input.")]
+        inputs.extend(io.Combo.Input(f"lora_{index}", options=optional, default="None", tooltip="Additional equal-strength LoRA contribution. Select a file for every input included by LoRA Count.") for index in range(3, 9))
+        inputs.extend(_controls("lodestone_merged_multi_lora"))
+        return io.Schema(node_id="LodestoneLoRAMultiMerger", display_name="Lodestone LoRA Multi-Merge",
+                         category="ModelUtils/LoRA/Merge/Lodestone", description="Expand 2-8 equal-strength LoRAs into full per-layer weight changes, combine them with Frobenius-norm-aware math, and save a full-difference LoRA.", inputs=inputs,
+                         outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True)
+
+    @classmethod
+    def execute(cls, **kwargs):
+        count = int(kwargs["lora_count"])
+        missing = [index for index in range(1, count + 1) if not kwargs[f"lora_{index}"] or kwargs[f"lora_{index}"] == "None"]
+        if missing:
+            raise ValueError(f"LoRA Count includes unselected input(s): {', '.join(map(str, missing))}.")
+        return io.NodeOutput(_execute(kwargs, count))
+
+
+LODESTONE_MERGER_NODES = [LodestoneLoRATwoMerger, LodestoneLoRAThreeMerger, LodestoneLoRAMultiMerger]

nodes/lora_alpha.py

@@ -0,0 +1,31 @@
+import torch
+
+
+def lora_alpha_scale(alpha: torch.Tensor | None, rank: int, *, layer: str) -> float:
+    """Return the effective LoRA multiplier represented by alpha and rank."""
+
+    if rank <= 0:
+        raise ValueError(f"LoRA rank for '{layer}' must be positive.")
+    if alpha is None:
+        return 1.0
+    if alpha.numel() != 1:
+        raise ValueError(f"LoRA alpha for '{layer}' must be scalar.")
+    return float(alpha.reshape(-1)[0].item()) / rank
+
+
+def normalize_lora_pair(
+    down: torch.Tensor,
+    up: torch.Tensor,
+    alpha: torch.Tensor | None,
+    *,
+    layer: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Materialize an alpha-free factor pair before rank alignment or merging."""
+
+    if down.ndim < 2 or up.ndim < 2:
+        raise ValueError(f"LoRA factors for '{layer}' must have at least two dimensions.")
+    rank = int(down.shape[0])
+    if int(up.shape[1]) != rank:
+        raise ValueError(f"LoRA factor rank mismatch for '{layer}'.")
+    scale = lora_alpha_scale(alpha, rank, layer=layer)
+    return down, up if scale == 1.0 else up * scale

nodes/lora_model_analysis.py

@@ -0,0 +1,248 @@
+"""Read-only, bounded analysis of a LoRA's effect on a diffusion model."""
+
+from contextlib import closing
+import csv
+import gc
+import heapq
+from io import StringIO
+import logging
+
+import folder_paths
+import torch
+from comfy.weight_adapter import LoRAAdapter
+from comfy_api.latest import io
+from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
+
+from .device_utils import cleanup_after_operation
+from .lora_alpha import normalize_lora_pair
+from .lora_resize import (
+    AdapterErrorCapture, build_lora_layer_map, canonical_lora_block_name,
+    layer_tensor_keys, parse_lora_layers, validate_canonical_blocks,
+)
+from .merger import _compile_patterns, _matches_any_pattern, load_documentation_from_file
+from .model_analysis import (
+    CWBStats, RawStats, TensorRecord, _block_name, _common_inputs,
+    _compute_pair_on_device, _fmt, _is_cuda_oom, _is_float_dtype, _layer_name,
+    _metrics, _release_cuda_oom, _render_comparison, _render_cwb,
+)
+from .quantization_guard import DiffusionQuantization, inspect_low_bit_input
+from .uel_io import stream_work_units
+
+
+def _plan_application(base, lora, base_path, params, base_quant):
+    pairs, passthrough = parse_lora_layers(lora.keys())
+    validate_canonical_blocks(pairs, "LoRA-on-model analysis")
+    layer_map = build_lora_layer_map([{"pairs": pairs}], base_path)
+    targets = {canonical_lora_block_name(key): key for key in base_quant.data_keys}
+    patches = {}
+    mapped = set()
+    unmapped_targets = []
+    for core, sources in layer_map.items():
+        for _, block in sources:
+            roles = pairs[block]
+            weight_role = next((role for role in ("set_weight", "diff", "w_norm") if role in roles), None)
+            specs = []
+            if weight_role is not None:
+                target = targets.get(core) if weight_role == "diff" else None
+                specs.append((target or targets.get(f"{core}.weight"), weight_role))
+            elif "down" in roles and "up" in roles:
+                specs.append((targets.get(f"{core}.weight"), "pair"))
+            bias_role = next((role for role in ("diff_b", "b_norm") if role in roles), None)
+            if bias_role:
+                specs.append((targets.get(f"{core}.bias"), bias_role))
+            for target, role in specs:
+                if target is not None:
+                    patches.setdefault(target, []).append((block, roles, role))
+                    mapped.add(block)
+                else:
+                    unmapped_targets.append(f"{block} ({role}: no base target)")
+    inventory = {name: [] for name in ("a_only", "b_only", "shape", "nonfloat", "low_bit", "excluded_a", "excluded_b")}
+    unmatched = sorted(set(pairs) - mapped) + unmapped_targets
+    patterns = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=params.get("glob_patterns", False))
+    low_base = base_quant.isolated_low_bit_keys
+    low_lora = inspect_low_bit_input(lora, "LoRA", "LoRA-on-model analysis")
+    units = []
+    for key in sorted(base_quant.data_keys):
+        if _matches_any_pattern(key, patterns, glob_mode=params.get("glob_patterns", False)) != params.get("include_mode", False):
+            inventory["excluded_a"].append(key)
+            inventory["excluded_b"].append(key)
+            continue
+        specs = patches.get(key, [])
+        source_keys = list(dict.fromkeys(
+            source_key for _, roles, role in specs
+            for source_key in (layer_tensor_keys(roles).values() if role == "pair" else (roles[role],))
+        ))
+        alpha_keys = {roles["alpha"] for _, roles, role in specs if role == "pair" and "alpha" in roles}
+        if key in low_base or any(source in low_lora and source not in alpha_keys for source in source_keys):
+            inventory["low_bit"].append(key)
+            continue
+        if not _is_float_dtype(base_quant.logical_dtype(key, torch.float32)):
+            inventory["nonfloat"].append(key)
+            continue
+        entries = {0: base_quant.required_keys(key)}
+        if source_keys:
+            entries[1] = source_keys
+        units.append((key, entries))
+    return units, patches, inventory, unmatched, passthrough
+
+
+def _apply_pair(key, weight, block, roles, loaded, strength):
+    tensors = {source_key: loaded[(1, source_key)] for source_key in layer_tensor_keys(roles).values()}
+    down = up = adapter = None
+    try:
+        down, up = normalize_lora_pair(tensors[roles["down"]], tensors[roles["up"]], tensors.get(roles.get("alpha")), layer=block)
+        tensors[roles["down"]], tensors[roles["up"]] = down, up
+        tensors.pop(roles.get("alpha"), None)
+        adapter = LoRAAdapter.load(block, tensors, None, tensors.get(roles.get("dora_scale")), set())
+        if adapter is None:
+            raise ValueError(f"LoRA-on-model analysis: could not load adapter {block!r} for {key!r}.")
+        capture = AdapterErrorCapture(key)
+        logger = logging.getLogger()
+        logger.addHandler(capture)
+        try:
+            weight = adapter.calculate_weight(weight, key, strength, 1.0, None, lambda value: value, torch.float32, None)
+        finally:
+            logger.removeHandler(capture)
+        if capture.messages:
+            raise RuntimeError("; ".join(capture.messages))
+        return weight
+    finally:
+        tensors.clear()
+        del adapter, down, up, weight
+
+
+def _apply_and_analyze(key, loaded, specs, strength, device, top_count):
+    source = loaded[(0, key)]
+    original = patched = delta = None
+    try:
+        original = (
+            transfer_to_gpu_pinned(source, device=device, dtype=torch.float32)
+            if str(device).startswith("cuda") else source.to(device=device, dtype=torch.float32)
+        )
+        patched = original.clone()
+        for block, roles, role in specs:
+            if role == "pair":
+                patched = _apply_pair(key, patched, block, roles, loaded, strength)
+                continue
+            delta = loaded[(1, roles[role])].to(device=device, dtype=torch.float32)
+            if delta.shape != patched.shape:
+                raise ValueError(f"LoRA-on-model analysis: {roles[role]} shape {tuple(delta.shape)} does not match {key} {tuple(patched.shape)}.")
+            if role == "set_weight":
+                patched.copy_(delta)
+            else:
+                patched.add_(delta, alpha=strength)
+            delta = None
+        if patched.shape != original.shape:
+            raise ValueError(f"LoRA-on-model analysis: applying {key!r} changed its shape; it cannot be compared with the base tensor.")
+        return _compute_pair_on_device(original, patched, device, False, top_count)
+    finally:
+        del source, original, patched, delta
+
+
+def _analyze_unit(key, loaded, specs, strength, device, top_count):
+    try:
+        return (*_apply_and_analyze(key, loaded, specs, strength, device, top_count), False)
+    except Exception as error:
+        if not str(device).startswith("cuda") or not _is_cuda_oom(error):
+            raise
+        error.__traceback__ = None
+        _release_cuda_oom()
+        try:
+            return (*_apply_and_analyze(key, loaded, specs, strength, "cpu", top_count), True)
+        except Exception as cpu_error:
+            raise RuntimeError(f"LoRA-on-model analysis CPU fallback failed for {key!r}: {cpu_error}") from cpu_error
+
+
+def _csv(headers, rows):
+    output = StringIO(newline="")
+    writer = csv.writer(output)
+    writer.writerow(headers)
+    writer.writerows(rows)
+    return output.getvalue()
+
+
+def analyze_lora_on_model(model_a, model_b, params):
+    lora_path = folder_paths.get_full_path_or_raise("loras", model_a)
+    base_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
+    strength = params.get("strength", 1.0)
+    records, fallback_keys, top_heap = [], [], []
+    global_raw = RawStats()
+    layer_stats, block_stats, block_cwb = {}, {}, {}
+    try:
+        with MemoryEfficientSafeOpen(base_path, low_memory=True) as base, MemoryEfficientSafeOpen(lora_path, low_memory=True) as lora:
+            base_quant = DiffusionQuantization(base, base_path, "Base diffusion model")
+            units, patches, inventory, unmatched, passthrough = _plan_application(base, lora, base_path, params, base_quant)
+            device = params["process_device"]
+            top_limit = params["top_weight_differences"]
+            with closing(stream_work_units({0: base, 1: lora}, units, pin_memory=str(device).startswith("cuda"))) as stream:
+                for key, loaded in stream:
+                    try:
+                        if key in base_quant.quantized_keys:
+                            loaded[(0, key)] = base_quant.decode(
+                                key,
+                                {source_key: loaded[(0, source_key)] for source_key in base_quant.required_keys(key)},
+                                torch.float32,
+                            )
+                        raw, cwb, top, fallback = _analyze_unit(key, loaded, patches.get(key, []), strength, device, top_limit)
+                        shape = base_quant.logical_shape(key)
+                        records.append(TensorRecord(key, shape, str(base_quant.logical_dtype(key, torch.float32)), str(torch.float32), raw, cwb, fallback))
+                        global_raw.add(raw)
+                        layer_stats.setdefault(_layer_name(key), RawStats()).add(raw)
+                        block_stats.setdefault(_block_name(key), RawStats()).add(raw)
+                        block_cwb.setdefault(_block_name(key), CWBStats()).add(cwb)
+                        if fallback:
+                            fallback_keys.append(key)
+                        for value, index, a, b in top:
+                            entry = (value, key, index, a, b, shape)
+                            if len(top_heap) < top_limit:
+                                heapq.heappush(top_heap, entry)
+                            elif top_limit and value > top_heap[0][0]:
+                                heapq.heapreplace(top_heap, entry)
+                    finally:
+                        loaded.clear()
+                        if params["force_clear_cache"]:
+                            gc.collect()
+                            if torch.cuda.is_available():
+                                torch.cuda.empty_cache()
+            label = f"{model_b} + {model_a} (strength={strength})"
+            comparison = _render_comparison(model_b, label, "LoRA applied to diffusion model", records, global_raw, layer_stats, block_stats, inventory, top_heap, fallback_keys, base.metadata(), base.metadata())
+            comparison = (
+                "# LoRA-on-model analysis\n\nModel A input is the LoRA; Model B input is the base diffusion model. "
+                "The comparison below uses original base as A and patched base as B; delta = patched − original. "
+                "Unmodified base tensors are included unless filtered. No model is saved.\n\n"
+                + comparison
+                + "\n## Unmapped LoRA groups\n\n" + ("\n".join(f"- `{name}`" for name in unmatched) or "None.")
+                + "\n\n## Unused LoRA tensors\n\n" + ("\n".join(f"- `{name}`" for name in passthrough) or "None.") + "\n"
+            )
+            cwb_report = _render_cwb(model_b, label, records, block_cwb, fallback_keys, "original versus LoRA-patched base tensor rows (index only)")
+            metric_names = ("mae", "mse", "rmse", "max", "relative_l2", "cosine", "pearson", "exact", "sign", "norm_ratio", "coverage")
+            metrics_csv = _csv(["key", "parameters", *metric_names], ([record.key, record.raw.total, *(_fmt(_metrics(record.raw)[name]) for name in metric_names)] for record in records))
+            cwb_names = ("pair_mean", "pair_min", "pair_max")
+            cwb_csv = _csv(["key", *cwb_names], ([record.key, *(_fmt(record.cwb.values()[name]) for name in cwb_names)] for record in records))
+            return comparison, cwb_report, metrics_csv, cwb_csv
+    finally:
+        cleanup_after_operation()
+
+
+class LoRAOnModelAnalysis(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        inputs = _common_inputs("diffusion_models", alignment_control=False)
+        inputs[1] = io.Combo.Input("model_a", options=folder_paths.get_filename_list("loras"), tooltip="LoRA to apply transiently to Model B before comparing the original and patched base tensors.")
+        inputs[2] = io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"), tooltip="Base diffusion model. Its original weights are compared against a transient LoRA-patched copy, one tensor at a time.")
+        inputs.append(io.Float.Input("strength", default=1.0, min=-10.0, max=10.0, step=0.01, tooltip="LoRA application strength. Alpha scaling and DoRA reconstruction use ComfyUI's adapter; explicit set_weight patches retain their existing replacement semantics."))
+        return io.Schema(
+            node_id="LoRAOnModelAnalysis", display_name="Analyze LoRA Effect on Diffusion Model",
+            category="ModelUtils/Analysis",
+            description="Compare the base diffusion model with its LoRA-patched weights using bounded streaming. No model is written or inferred.",
+            inputs=inputs,
+            outputs=[io.String.Output(display_name="comparison_report"), io.String.Output(display_name="cwb_report"), io.String.Output(display_name="documentation"), io.String.Output(display_name="layerwise_metrics_csv"), io.String.Output(display_name="layerwise_cwb_csv")],
+        )
+
+    @classmethod
+    def execute(cls, **kwargs):
+        docs = load_documentation_from_file("lora_model_analysis.md")
+        if kwargs["execution_mode"] == "DOCUMENTATION ONLY":
+            return io.NodeOutput("Documentation mode active. No comparison performed.", "No CWB analysis performed.", docs, "", "")
+        comparison, cwb, metrics_csv, cwb_csv = analyze_lora_on_model(kwargs["model_a"], kwargs["model_b"], kwargs)
+        return io.NodeOutput(comparison, cwb, docs, metrics_csv, cwb_csv)

nodes/lora_rename.py

@@ -0,0 +1,118 @@
+import os
+import folder_paths
+import comfy.utils
+from tqdm import tqdm
+from comfy_api.latest import io
+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
+from .uel_io import atomic_uel_writer
+from .artifact_paths import canonical_model_artifact_path
+
+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)
+
+    output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
+
+    # Stream tensors, rename on the fly, write immediately
+    with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
+        metadata = handler.metadata() or {}
+        with atomic_uel_writer(output_path, metadata) as writer:
+            original_keys = handler.keys()
+            pbar = comfy.utils.ProgressBar(len(original_keys))
+            stream = handler.async_stream(
+                original_keys, batch_size=1, prefetch_batches=1, pin_memory=False
+            )
+            for batch in tqdm(stream, total=len(original_keys), desc="Renaming Anima LoRA keys", unit="keys"):
+                key, tensor = batch[0]
+                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_batch([(new_key, tensor.contiguous())])
+                handler.mark_processed(key)
+                pbar.update(1)
+
+    # Cleanup after operation
+    cleanup_after_operation()
+
+    return output_name
+
+
+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"),
+                    tooltip="Diffusers-format Anima LoRA whose tensor keys will be renamed.",
+                ),
+                io.String.Input("output_filename", default="anima_lora_non_diffusers", tooltip="Output filename without extension, written under ComfyUI's LoRA directory."),
+            ],
+            outputs=[
+                io.AnyType.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)

nodes/merger_ops.py

@@ -4,6 +4,8 @@ import torch.nn.functional as F
 from enum import Enum
 
 
+from unifiedefficientloader import transfer_to_gpu_pinned
+
 class MissingTensorBehavior(Enum):
     """Controls behavior when a tensor key is missing from a model."""
     ERROR = "error"      # Raise exception (strict mode)
@@ -127,20 +129,31 @@ class Operation:
         raise NotImplementedError
 
     def recurse(self, operation):
-        source_tensors = [source_oper.merge() for source_oper in operation.sources]
-        return operation.oper(*source_tensors)
+        self._source_tensors = [source_oper.merge() for source_oper in operation.sources]
+        return operation.oper(*self._source_tensors)
 
     def merge(self):
         return self.merge_func(self)
 
+    def clean(self):
+        if hasattr(self, '_source_tensors'):
+            for t in self._source_tensors:
+                del t
+            del self._source_tensors
+        for source in self.sources:
+            if hasattr(source, 'clean'):
+                source.clean()
+
 class LoadTensor(Operation):
     def __init__(self, key, model_name, handlers, device, dtype,
+                 preloaded_tensors=None,
                  on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None):
         super().__init__(key)
         self.model_name = model_name
         self.handlers = handlers
         self.device = device
         self.dtype = dtype
+        self.preloaded_tensors = preloaded_tensors or {}
         self.on_missing = on_missing
         self.fallback_shape = fallback_shape
         self.fallback_dtype = fallback_dtype
@@ -149,8 +162,8 @@ class LoadTensor(Operation):
         if self.model_name not in self.handlers:
             raise ValueError(f"Model '{self.model_name}' is required for this mode but was not provided.")
 
-        handler = self.handlers[self.model_name]
-        if self.key not in handler.keys():
+        source = (self.model_name, self.key)
+        if source not in self.preloaded_tensors:
             if self.on_missing == MissingTensorBehavior.ERROR:
                 raise MissingTensorError(
                     f"Key '{self.key}' not found in model '{self.model_name}'"
@@ -165,7 +178,17 @@ class LoadTensor(Operation):
                 dtype = self.fallback_dtype if self.fallback_dtype else self.dtype
                 return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype)
 
-        return handler.get_tensor(self.key).to(device=self.device, dtype=self.dtype)
+        cpu_tensor = self.preloaded_tensors[source]
+        if self.device == 'cuda':
+            self._tensor = transfer_to_gpu_pinned(cpu_tensor, self.device, self.dtype)
+        else:
+            self._tensor = cpu_tensor.to(device=self.device, dtype=self.dtype)
+        del cpu_tensor
+        return self._tensor
+
+    def clean(self):
+        if hasattr(self, '_tensor'):
+            del self._tensor
 
 class Multiply(Operation):
     def __init__(self, key, alpha, *sources):
@@ -218,7 +241,7 @@ class SVD(Operation):
                 a = a[slices]
                 b = b[slices]
 
-        diff, weights, conv2d = a - b, {}, len(a.size()) == 4
+        diff, conv2d = a - b, len(a.size()) == 4
         kernel_size = None if not conv2d else a.size()[2:4]
         conv2d_3x3 = conv2d and kernel_size != (1, 1)
         rank = self.alpha if not conv2d_3x3 or self.beta is None else self.beta
@@ -298,6 +321,53 @@ class PowerUpOp(Operation):
         return (m * delta) / keep_prob
 
 
+class DARETIESOp(Operation):
+    """DARE + TIES combined operation on a single task vector (delta = b - a).
+
+    Step 1, DARE: randomly drop `drop_rate` fraction of delta parameters,
+                   rescale survivors by 1/(1-drop_rate).
+    Step 2, TIES trim: zero out parameters below the `trim_quantile` magnitude threshold.
+    Step 3, TIES elect: determine dominant sign per parameter position.
+    Step 4, TIES merge: zero out parameters that disagree with the dominant sign.
+
+    For a two-model merge the TIES elect/merge step degenerates to a no-op (single vector,
+    sign is always self-consistent) but the trim step still reduces interference from
+    near-zero noisy deltas. The dominant sign filter is included for correctness and
+    forward compatibility with n-model extensions.
+    """
+    def __init__(self, key, alpha, beta, gamma, seed, *sources):
+        super().__init__(key, *sources)
+        # alpha = drop_rate (DARE), beta = trim_quantile (TIES), gamma = lambda scale
+        self.alpha, self.beta, self.gamma, self.seed = alpha, beta, gamma, seed
+
+    def oper(self, a, b):
+        if a is None or b is None:
+            return None
+        a, b = resize_tensors(a, b, mode=self.alignment_mode)
+        delta = (b - a).float()
+
+        # --- DARE: drop and rescale ---
+        keep_prob = 1.0 - self.alpha
+        if keep_prob <= 0:
+            return torch.zeros_like(a)
+        rng = torch.Generator(device=delta.device).manual_seed(self.seed)
+        dare_mask = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).float()
+        delta = (delta * dare_mask) / keep_prob
+
+        # --- TIES trim: zero parameters below magnitude threshold ---
+        if self.beta > 0:
+            flat_abs = delta.abs().flatten()
+            k = max(1, int(flat_abs.numel() * self.beta))
+            threshold = torch.kthvalue(flat_abs, k).values
+            delta = torch.where(delta.abs() < threshold, torch.zeros_like(delta), delta)
+
+        # --- TIES elect: dominant sign (trivially self-consistent for single vector) ---
+        dominant_sign = torch.sign(delta)
+        delta = torch.where(torch.sign(delta) == dominant_sign, delta, torch.zeros_like(delta))
+
+        return (a + self.gamma * delta).to(a.dtype)
+
+
 class InterpolateDifference(Operation):
     def __init__(self, key, alpha, beta, gamma, seed, *sources):
         super().__init__(key, *sources)
@@ -386,6 +456,7 @@ class CalcMode:
         """Helper to build loader args for secondary models (B, C) with fallback support."""
         base_args = {
             "handlers": kwargs['handlers'],
+            "preloaded_tensors": kwargs.get('preloaded_tensors', {}),
             "device": kwargs['device'],
             "dtype": kwargs['dtype'],
             "on_missing": mismatch_mode,
@@ -409,6 +480,10 @@ class PreloadedTensor(Operation):
     def merge(self) -> torch.Tensor:
         return self.tensor
 
+    def clean(self):
+        if hasattr(self, 'tensor'):
+            del self.tensor
+
 class WeightSum(CalcMode):
     name = 'Weight-Sum'
     description = 'A * (1 - α) + B * α. Simple linear interpolation.'
@@ -617,7 +692,12 @@ class ManEnhInterpDifference(CalcMode):
     name = 'Enhanced Man Interp'
     description = 'Enhanced interpolation between each pair of values from A and B depending on their difference relative to other values.'
     models_used = ['A', 'B']
-    param_docs = {'alpha': 'Interpolation strength.', 'beta': 'Lower mean threshold.', 'gamma': 'Upper mean threshold.', 'delta': 'Smoothness factor.'}
+    param_docs = {
+        'alpha': 'Model B pull strength (exponent 1/alpha - 1). 0.25 strongly suppresses B; 0.50 is linear; 1.0 is full B.',
+        'beta': 'Lower mean threshold. Only weights with normalized difference > beta blend B.',
+        'gamma': 'Upper mean threshold. Only weights with normalized difference < gamma blend B.',
+        'delta': 'Smoothness factor. 0.0 is stochastic Bernoulli choice; 1.0 is smooth linear blend.',
+    }
     def create_recipe(self, key, **kwargs):
         mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
         alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
@@ -642,7 +722,11 @@ class AutoEnhInterpDifference(CalcMode):
     name = 'Enhanced Auto Interp'
     description = 'Interpolates between each pair of values from A and B depending on their difference relative to other values.'
     models_used = ['A', 'B']
-    param_docs = {'alpha': 'Interpolation strength.', 'beta': 'Threshold adjustment factor.', 'gamma': 'Smoothness factor.'}
+    param_docs = {
+        'alpha': 'Model B pull strength (exponent 1/alpha - 1). Lower keeps output close to Model A; higher pulls in Model B.',
+        'beta': 'Threshold width around the layer mean difference (+/- beta).',
+        'gamma': 'Smoothness factor. 0.0 is stochastic Bernoulli choice; 1.0 is smooth linear blend.',
+    }
     def create_recipe(self, key, **kwargs):
         mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
         alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
@@ -667,7 +751,11 @@ class WeightSumCutoffMode(CalcMode):
     name = 'Weight-Sum Cutoff'
     description = 'Weight-sum with cutoff based on value differences.'
     models_used = ['A', 'B']
-    param_docs = {'alpha': 'Interpolation weight.', 'beta': 'Upper threshold.', 'gamma': 'Lower threshold.'}
+    param_docs = {
+        'alpha': 'Model B blend fraction inside the threshold window. 0.0 is pure Model A; 0.25 is 25% B / 75% A; 1.0 is 100% B.',
+        'beta': 'Upper difference threshold cutoff.',
+        'gamma': 'Lower difference threshold cutoff.',
+    }
     def create_recipe(self, key, **kwargs):
         mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
         alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
@@ -685,5 +773,179 @@ class WeightSumCutoffMode(CalcMode):
         return res
 
 
-TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
-THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
+class DARETIESMode(CalcMode):
+    name = 'Power-Up (DARE+TIES)'
+    description = 'Adds capabilities of B to A using DARE (Drop and Rescale) followed by TIES (Trim, Elect Sign). Best practice for LoRA-tuned model merging.'
+    models_used = ['A', 'B']
+    param_docs = {
+        'alpha': 'DARE drop rate: fraction of delta parameters randomly zeroed. Higher = more aggressive sparsification.',
+        'beta':  'TIES trim quantile: fraction of smallest-magnitude delta parameters zeroed after DARE. Removes residual noise.',
+        'gamma': 'Lambda: scale multiplier applied to the final merged delta before adding to Model A.',
+        'seed':  'Random seed for the DARE dropout mask.',
+    }
+
+    def create_recipe(self, key, **kwargs):
+        mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
+        alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
+
+        if '_tensor_a' in kwargs:
+            a = PreloadedTensor(key, kwargs['_tensor_a'])
+        else:
+            loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
+            a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
+
+        loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
+        b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
+
+        op = DARETIESOp(
+            key,
+            kwargs['alpha'],   # drop_rate
+            kwargs['beta'],    # trim_quantile
+            kwargs['gamma'],   # lambda scale
+            kwargs['seed'],
+            a, b,
+        )
+        op.alignment_mode = alignment_mode
+        return op
+
+
+class EnhancedPowerUpOp(Operation):
+    def __init__(self, key, alpha, beta, gamma, delta, seed, *sources):
+        super().__init__(key, *sources)
+        self.mask_power = alpha       # Mask power
+        self.min_keep_prob = beta     # Min keep prob
+        self.mask_smooth = gamma      # Smoothness (0 or 1)
+        self.delta = delta            # Final multiplier
+        self.seed = seed
+
+    def oper(self, a, b):
+        if a is None or b is None: return None
+        a, b = resize_tensors(a, b, mode=self.alignment_mode)
+        delta_val = (b - a).float()
+        abs_delta = torch.abs(delta_val)
+        max_delta = torch.max(abs_delta)
+        if max_delta == 0: return torch.zeros_like(delta_val)
+
+        prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
+
+        rng = torch.Generator(device=a.device).manual_seed(self.seed)
+        random_mask = torch.bernoulli(prob, generator=rng)
+        interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
+
+        rescaled_delta = delta_val * interpolated_mask
+        return rescaled_delta
+
+
+class EnhancedDARETIESOp(Operation):
+    def __init__(self, key, alpha, beta, gamma, delta, epsilon, zeta, seed, *sources):
+        super().__init__(key, *sources)
+        self.alpha = alpha       # Unused directly for drop scale
+        self.beta = beta         # TIES trim quantile
+        self.gamma = gamma       # TIES lambda scale
+        self.mask_power = delta
+        self.min_keep_prob = epsilon
+        self.mask_smooth = zeta
+        self.seed = seed
+
+    def oper(self, a, b):
+        if a is None or b is None: return None
+        a, b = resize_tensors(a, b, mode=self.alignment_mode)
+        delta_val = (b - a).float()
+        abs_delta = torch.abs(delta_val)
+        max_delta = torch.max(abs_delta)
+
+        if max_delta > 0:
+            prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
+
+            rng = torch.Generator(device=a.device).manual_seed(self.seed)
+            random_mask = torch.bernoulli(prob, generator=rng)
+            interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
+
+            delta_val = delta_val * interpolated_mask
+
+        if self.beta > 0:
+            flat_abs = delta_val.abs().flatten()
+            k = max(1, int(flat_abs.numel() * self.beta))
+            threshold = torch.kthvalue(flat_abs, k).values
+            delta_val = torch.where(delta_val.abs() < threshold, torch.zeros_like(delta_val), delta_val)
+
+        dominant_sign = torch.sign(delta_val)
+        delta_val = torch.where(torch.sign(delta_val) == dominant_sign, delta_val, torch.zeros_like(delta_val))
+
+        return (a + self.gamma * delta_val).to(a.dtype)
+
+
+class EnhancedPowerUp(CalcMode):
+    name = 'Power-Up Enhanced (DARE)'
+    description = 'Adds capabilities of B to A using Enhanced Drop and Rescale based on relative tensor differences.'
+    models_used = ['A', 'B']
+    param_docs = {
+        'alpha': 'Mask power (Curve. 2.0 = quadratic).',
+        'beta': 'Minimum keep probability (Floor to prevent explosion).',
+        'gamma': 'Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale).',
+        'delta': 'Multiplier for the final difference.',
+    }
+    def create_recipe(self, key, **kwargs):
+        mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
+        alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
+
+        if '_tensor_a' in kwargs:
+            a = PreloadedTensor(key, kwargs['_tensor_a'])
+        else:
+            loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
+            a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
+
+        loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
+        b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
+
+        deltahat = EnhancedPowerUpOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b)
+        deltahat.alignment_mode = alignment_mode
+        multiplied = Multiply(key, kwargs['delta'], deltahat)
+        res = Add(key, a, multiplied)
+        res.alignment_mode = alignment_mode
+        return res
+
+
+class EnhancedDARETIESMode(CalcMode):
+    name = 'Power-Up Enhanced (DARE+TIES)'
+    description = 'Adds capabilities of B to A using Enhanced DARE mask followed by TIES.'
+    models_used = ['A', 'B']
+    param_docs = {
+        'alpha': 'Unused (Drop scale handled via mask interpolation).',
+        'beta': 'TIES trim quantile (fraction of smallest-magnitude delta zeroed).',
+        'gamma': 'TIES Lambda scale.',
+        'delta': 'Mask power (Curve. 2.0 = quadratic).',
+        'epsilon': 'Minimum keep probability (Floor to prevent explosion).',
+        'zeta': 'Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale).',
+    }
+
+    def create_recipe(self, key, **kwargs):
+        mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
+        alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
+
+        if '_tensor_a' in kwargs:
+            a = PreloadedTensor(key, kwargs['_tensor_a'])
+        else:
+            loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
+            a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
+
+        loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
+        b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
+
+        op = EnhancedDARETIESOp(
+            key,
+            kwargs.get('alpha', 0.5),
+            kwargs.get('beta', 0.2),
+            kwargs.get('gamma', 1.0),
+            kwargs.get('delta', 0.5),
+            kwargs.get('epsilon', 0.5),
+            kwargs.get('zeta', 0.5),
+            kwargs['seed'],
+            a, b,
+        )
+        op.alignment_mode = alignment_mode
+        return op
+
+
+TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode(), EnhancedPowerUp(), EnhancedDARETIESMode()]
+THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]

nodes/merger_utils.py

@@ -1,373 +1,7 @@
 """
-Memory-efficient safetensors utilities with parallel I/O and pinned memory support.
-
-Based on benchmarking, parallel reading with ThreadPoolExecutor provides 2-4x speedup
-for large models compared to sequential reads.
-
-Worker count is automatically optimized based on device capabilities when workers=None.
-Pinned memory enables faster CPU→GPU transfers (2-3x improvement).
+Utilities for model merging operations.
 """
-import os
-import mmap
-import json
-import struct
-import torch
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Dict, List, Optional, Any, Tuple
-
-# Import ComfyUI's memory pinning functions if available
-try:
-    from comfy.model_management import pin_memory, unpin_memory, is_device_cuda
-    PINNING_AVAILABLE = True
-except ImportError:
-    PINNING_AVAILABLE = False
-    def pin_memory(tensor): return False
-    def unpin_memory(tensor): return False
-    def is_device_cuda(device): return str(device).startswith('cuda')
-
-
-class PinnedTensor:
-    """Context manager for automatic pinned memory management.
-
-    Usage:
-        with PinnedTensor(cpu_tensor) as pinned:
-            gpu_tensor = pinned.to('cuda')
-        # Memory automatically unpinned after transfer
-    """
-
-    def __init__(self, tensor: torch.Tensor, auto_pin: bool = True):
-        self.tensor = tensor
-        self.is_pinned = False
-        self.auto_pin = auto_pin and PINNING_AVAILABLE
-
-    def __enter__(self) -> torch.Tensor:
-        if self.auto_pin and self.tensor.device.type == 'cpu':
-            self.is_pinned = pin_memory(self.tensor)
-        return self.tensor
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.is_pinned:
-            unpin_memory(self.tensor)
-            self.is_pinned = False
-
-
-def transfer_to_gpu_pinned(
-    tensor: torch.Tensor,
-    device: str = 'cuda',
-    dtype: Optional[torch.dtype] = None
-) -> torch.Tensor:
-    """Transfer tensor to GPU using pinned memory for faster transfer.
-
-    Args:
-        tensor: CPU tensor to transfer
-        device: Target GPU device
-        dtype: Optional dtype conversion
-
-    Returns:
-        Tensor on GPU device
-    """
-    if not PINNING_AVAILABLE or not is_device_cuda(torch.device(device)):
-        # Fall back to regular transfer
-        if dtype:
-            return tensor.to(device=device, dtype=dtype)
-        return tensor.to(device=device)
-
-    # Pin memory for faster transfer
-    was_pinned = pin_memory(tensor)
-    try:
-        if dtype:
-            result = tensor.to(device=device, dtype=dtype, non_blocking=True)
-        else:
-            result = tensor.to(device=device, non_blocking=True)
-        # Sync to ensure transfer is complete before unpinning
-        if was_pinned:
-            torch.cuda.current_stream().synchronize()
-        return result
-    finally:
-        if was_pinned:
-            unpin_memory(tensor)
-
-
-class MemoryEfficientSafeOpen:
-    """Memory-efficient safetensors file reader with mmap, parallel I/O, and pinned memory support.
-
-    Features:
-    - mmap mode: Zero-copy tensor access via memory-mapped file
-    - low_memory mode: Direct file reads to minimize OS page cache usage
-    - Parallel loading: Multi-threaded tensor reads for 2-4x speedup
-    - Sorted batch reads: Keys sorted by file offset for sequential I/O
-    - Auto-optimized workers: Adjusts parallelism based on device capabilities
-    - Pinned memory: Faster CPU→GPU transfers when transferring to CUDA
-
-    Args:
-        filename: Path to safetensors file
-        device: Target device (default 'cpu')
-        mmap_mode: Use memory-mapped file for zero-copy (default True)
-        low_memory: Use direct file reads to minimize memory (overrides mmap_mode)
-    """
-
-    def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True, low_memory: bool = False):
-        self.filename = filename
-        self.device = device
-        # low_memory mode forces mmap off to avoid OS page caching
-        self.low_memory = low_memory
-        self.mmap_mode = mmap_mode and not low_memory
-        self.header, self.header_size = self._read_header()
-        self.file = open(filename, "rb")
-        self.mmap_obj = None
-
-        if self.mmap_mode:
-            self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.mmap_obj:
-            self.mmap_obj.close()
-        self.file.close()
-
-
-    def keys(self) -> List[str]:
-        """Return all tensor keys (excluding metadata)."""
-        return [k for k in self.header.keys() if k != "__metadata__"]
-
-    def metadata(self) -> Dict[str, str]:
-        """Return file metadata."""
-        return self.header.get("__metadata__", {})
-
-    def keys_sorted_by_offset(self) -> List[str]:
-        """Return keys sorted by file offset for optimal sequential I/O."""
-        keys_with_offsets = []
-        for key in self.keys():
-            offset = self.header[key]["data_offsets"][0]
-            keys_with_offsets.append((key, offset))
-        keys_with_offsets.sort(key=lambda x: x[1])
-        return [k for k, _ in keys_with_offsets]
-
-    def get_tensor(self, key: str) -> torch.Tensor:
-        """Load a single tensor by key."""
-        if key not in self.header:
-            raise KeyError(f"Tensor '{key}' not found in the file")
-
-        metadata = self.header[key]
-        offset_start, offset_end = metadata["data_offsets"]
-
-        if self.mmap_mode and self.mmap_obj:
-            if offset_start != offset_end:
-                start = self.header_size + 8 + offset_start
-                end = self.header_size + 8 + offset_end
-                tensor_bytes = memoryview(self.mmap_obj)[start:end]
-            else:
-                tensor_bytes = None
-        else:
-            # Non-mmap mode: use pre-allocated bytearray with readinto for minimal copies
-            tensor_bytes = None
-            if offset_start != offset_end:
-                self.file.seek(self.header_size + 8 + offset_start)
-                # Pre-allocate writable buffer and read directly into it
-                tensor_bytes = bytearray(offset_end - offset_start)
-                self.file.readinto(tensor_bytes)
-
-        return self._deserialize_tensor(tensor_bytes, metadata)
-
-
-    def get_tensor_to_gpu(
-        self,
-        key: str,
-        device: str = 'cuda',
-        dtype: Optional[torch.dtype] = None
-    ) -> torch.Tensor:
-        """Load a tensor and transfer to GPU using pinned memory for faster transfer.
-
-        Args:
-            key: Tensor key to load
-            device: Target GPU device (default 'cuda')
-            dtype: Optional dtype conversion
-
-        Returns:
-            Tensor on GPU device
-        """
-        cpu_tensor = self.get_tensor(key)
-        return transfer_to_gpu_pinned(cpu_tensor, device, dtype)
-
-    def get_tensors_parallel(
-        self,
-        keys: Optional[List[str]] = None,
-        workers: Optional[int] = None,
-        sort_by_offset: bool = True
-    ) -> Dict[str, torch.Tensor]:
-        """Load multiple tensors in parallel using ThreadPoolExecutor.
-
-        Args:
-            keys: List of keys to load (default: all keys)
-            workers: Number of worker threads (default: auto-calculated based on device)
-            sort_by_offset: Sort keys by file offset before loading (default: True)
-
-        Returns:
-            Dict mapping key -> tensor
-        """
-        if keys is None:
-            keys = self.keys()
-
-        # Auto-calculate workers based on device capabilities
-        if workers is None:
-            try:
-                from .device_utils import get_optimal_workers, estimate_model_size
-                model_size = estimate_model_size(self.filename)
-                workers = get_optimal_workers(model_size)
-            except ImportError:
-                workers = min(4, os.cpu_count() or 4)
-
-        if sort_by_offset:
-            # Sort by offset for better I/O patterns
-            keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys]
-            keys_with_offsets.sort(key=lambda x: x[1])
-            keys = [k for k, _ in keys_with_offsets]
-
-        result = {}
-
-        with ThreadPoolExecutor(max_workers=workers) as executor:
-            futures = {
-                executor.submit(
-                    _read_tensor_from_file,
-                    self.filename,
-                    key,
-                    self.header[key],
-                    self.header_size
-                ): key
-                for key in keys
-            }
-
-            for future in as_completed(futures):
-                key = futures[future]
-                try:
-                    tensor = future.result()
-                    if tensor is not None:
-                        result[key] = tensor
-                except Exception as e:
-                    print(f"[MemoryEfficientSafeOpen] Failed to load {key}: {e}")
-
-        return result
-
-    def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
-        """Load a tensor and decode as JSON dict (for config tensors)."""
-        if key not in self.header:
-            raise KeyError(f"Tensor '{key}' not found")
-
-        metadata = self.header[key]
-        offset_start, offset_end = metadata["data_offsets"]
-
-        if offset_start == offset_end:
-            return {}
-
-        if self.mmap_mode and self.mmap_obj:
-            start = self.header_size + 8 + offset_start
-            end = self.header_size + 8 + offset_end
-            tensor_bytes = bytes(self.mmap_obj[start:end])
-        else:
-            self.file.seek(self.header_size + 8 + offset_start)
-            tensor_bytes = self.file.read(offset_end - offset_start)
-
-        return json.loads(tensor_bytes.decode("utf-8"))
-
-    def _read_header(self):
-        with open(self.filename, "rb") as f:
-            header_size = struct.unpack("<Q", f.read(8))[0]
-            header_json = f.read(header_size).decode("utf-8")
-        return json.loads(header_json), header_size
-
-    def _deserialize_tensor(self, tensor_bytes, metadata):
-        dtype_str = metadata["dtype"]
-        shape = metadata["shape"]
-        dtype = self._get_torch_dtype(dtype_str)
-
-        if tensor_bytes is None:
-            byte_tensor = torch.empty(0, dtype=torch.uint8)
-        else:
-            # Avoid extra copy if already a bytearray (low_memory mode)
-            if isinstance(tensor_bytes, bytearray):
-                byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
-            else:
-                # mmap memoryview needs copy to create writable tensor
-                byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
-
-        if dtype_str in ["F8_E5M2", "F8_E4M3"]:
-            return self._convert_float8(byte_tensor, dtype_str, shape)
-
-        return byte_tensor.view(dtype).reshape(shape)
-
-
-    @staticmethod
-    def _get_torch_dtype(dtype_str):
-        dtype_map = {
-            "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
-            "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
-            "U8": torch.uint8, "BOOL": torch.bool,
-        }
-        if hasattr(torch, "float8_e5m2"):
-            dtype_map["F8_E5M2"] = torch.float8_e5m2
-        if hasattr(torch, "float8_e4m3fn"):
-            dtype_map["F8_E4M3"] = torch.float8_e4m3fn
-
-        dtype = dtype_map.get(dtype_str)
-        if dtype is None:
-            raise ValueError(f"Unsupported dtype: {dtype_str}")
-        return dtype
-
-    @staticmethod
-    def _convert_float8(byte_tensor, dtype_str, shape):
-        if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
-            return byte_tensor.view(torch.float8_e5m2).reshape(shape)
-        elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
-            return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
-        else:
-            raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")
-
-
-def _read_tensor_from_file(
-    filepath: str,
-    key: str,
-    metadata: Dict,
-    header_size: int
-) -> Optional[torch.Tensor]:
-    """Helper function to read a single tensor from file (for parallel execution).
-
-    Each worker opens its own file handle for thread safety.
-    """
-    offset_start, offset_end = metadata["data_offsets"]
-    dtype_str = metadata["dtype"]
-    shape = metadata["shape"]
-
-    if offset_start == offset_end:
-        return None
-
-    dtype_map = {
-        "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
-        "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
-        "U8": torch.uint8, "BOOL": torch.bool,
-    }
-    if hasattr(torch, "float8_e5m2"):
-        dtype_map["F8_E5M2"] = torch.float8_e5m2
-    if hasattr(torch, "float8_e4m3fn"):
-        dtype_map["F8_E4M3"] = torch.float8_e4m3fn
-
-    dtype = dtype_map.get(dtype_str, torch.float32)
-
-    with open(filepath, "rb") as f:
-        f.seek(header_size + 8 + offset_start)
-        tensor_bytes = f.read(offset_end - offset_start)
-
-    byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
-
-    if dtype_str in ["F8_E5M2", "F8_E4M3"]:
-        if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
-            return byte_tensor.view(torch.float8_e5m2).reshape(shape)
-        elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
-            return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
-
-    return byte_tensor.view(dtype).reshape(shape)
-
+from typing import Dict, List
 
 def get_union_keys(*handlers) -> List[str]:
     """Returns the union of all tensor keys across multiple model handlers."""
@@ -386,4 +20,4 @@ def get_intersection_keys(*handlers) -> List[str]:
     all_keys = set(handlers[0].keys())
     for handler in handlers[1:]:
         all_keys &= set(handler.keys())
-    return sorted(all_keys)
+    return sorted(all_keys)

nodes/metakeys.py

@@ -1,11 +1,9 @@
 import os
-import json
-import struct
 import folder_paths
 import comfy.utils
 from tqdm import tqdm
 from comfy_api.latest import io
-from safetensors import safe_open
+from unifiedefficientloader import MemoryEfficientSafeOpen
 from .utils import convert_pt_to_safetensors
 
 
@@ -21,25 +19,15 @@ def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]:
                 return f"Conversion failed: {error_message}", ""
         model_path = temp_safe_path
 
-    # Read header only - no tensor loading needed
-    with open(model_path, "rb") as f:
-        header_size = struct.unpack("<Q", f.read(8))[0]
-        header_json = f.read(header_size).decode("utf-8")
-    header = json.loads(header_json)
-
-    # Extract metadata and layer shapes from header
-    metadata = header.pop("__metadata__", {})
-    metadata_str = str(metadata)
-
-    layer_shapes = ""
-    keys = [k for k in header.keys()]
-    pbar = comfy.utils.ProgressBar(len(keys))
-    for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
-        info = header[layer_name]
-        shape = str(info.get("shape", []))
-        name_shape = f"{layer_name}, {shape}\n"
-        layer_shapes += name_shape
-        pbar.update(1)
+    with MemoryEfficientSafeOpen(model_path, low_memory=True) as handler:
+        metadata_str = str(handler.metadata() or {})
+        keys = handler.keys()
+        pbar = comfy.utils.ProgressBar(len(keys))
+        lines = []
+        for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
+            lines.append(f"{layer_name}, {list(handler.get_shape(layer_name))}\n")
+            pbar.update(1)
+        layer_shapes = "".join(lines)
 
     return metadata_str, layer_shapes
 
@@ -56,6 +44,7 @@ class ModelMetaKeys(io.ComfyNode):
                 io.Combo.Input(
                     "diffusionmodel_name",
                     options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Diffusion model whose safetensors metadata and tensor keys will be listed.",
                 ),
             ],
             outputs=[
@@ -82,6 +71,7 @@ class TextEncoderMetaKeys(io.ComfyNode):
                 io.Combo.Input(
                     "textencoder_name",
                     options=folder_paths.get_filename_list("text_encoders"),
+                    tooltip="Text encoder whose safetensors metadata and tensor keys will be listed.",
                 ),
             ],
             outputs=[
@@ -108,6 +98,7 @@ class LoRAMetaKeys(io.ComfyNode):
                 io.Combo.Input(
                     "lora_name",
                     options=folder_paths.get_filename_list("loras"),
+                    tooltip="LoRA whose safetensors metadata and tensor keys will be listed.",
                 ),
             ],
             outputs=[
@@ -134,6 +125,7 @@ class CheckpointMetaKeys(io.ComfyNode):
                 io.Combo.Input(
                     "ckpt_name",
                     options=folder_paths.get_filename_list("checkpoints"),
+                    tooltip="Checkpoint whose safetensors metadata and tensor keys will be listed.",
                 ),
             ],
             outputs=[
@@ -160,6 +152,7 @@ class EmbeddingMetaKeys(io.ComfyNode):
                 io.Combo.Input(
                     "embedding",
                     options=folder_paths.get_filename_list("embeddings"),
+                    tooltip="Embedding whose safetensors metadata and tensor keys will be listed.",
                 ),
             ],
             outputs=[

nodes/minimax_h3_fold.py

@@ -0,0 +1,687 @@
+from __future__ import annotations
+
+import fnmatch
+import json
+import math
+import os
+import re
+
+import comfy.utils
+from comfy_api.latest import io
+import folder_paths
+import torch
+from unifiedefficientloader import MemoryEfficientSafeOpen
+
+from .artifact_paths import canonical_model_artifact_path
+from .device_utils import cleanup_after_operation
+from .quantization_guard import inspect_low_bit_input
+from .uel_io import atomic_uel_writer
+
+
+def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
+    """Compile newline or whitespace separated patterns."""
+    if not pattern_string or not pattern_string.strip():
+        return []
+    patterns = []
+    for pattern in pattern_string.split():
+        if not pattern:
+            continue
+        if glob_mode:
+            patterns.append(pattern)
+        else:
+            try:
+                patterns.append(re.compile(pattern))
+            except re.error as e:
+                raise ValueError(f"Invalid regex pattern '{pattern}': {e}") from e
+    return patterns
+
+
+def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
+    if not patterns:
+        return False
+    if glob_mode:
+        return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
+    return any(p.search(key) is not None for p in patterns)
+
+
+def time_embedding_curve(
+    w1: torch.Tensor,
+    b1: torch.Tensor,
+    w2: torch.Tensor,
+    b2: torch.Tensor,
+    grid_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Sample silu(time_embedder(t)) on the curve grid."""
+    t = torch.linspace(0.0, 1.0, grid_size, dtype=torch.float32, device=device)
+    half = w1.shape[1] // 2
+    freqs = torch.exp(
+        -math.log(10000.0)
+        * torch.arange(half, dtype=torch.float32, device=device)
+        / half
+    )
+    args = t[:, None] * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    w1_d, b1_d, w2_d, b2_d = (x.to(device, torch.float32) for x in (w1, b1, w2, b2))
+    hidden = torch.nn.functional.silu(emb @ w1_d.T + b1_d)
+    return torch.nn.functional.silu(hidden @ w2_d.T + b2_d)
+
+
+def _fold_curve_basis(
+    w1: torch.Tensor,
+    b1: torch.Tensor,
+    w2: torch.Tensor,
+    b2: torch.Tensor,
+    grid_size: int,
+    rank: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, float]:
+    """Calculate SVD curve basis and adaln_t_table."""
+    compute_device = device
+    try:
+        curve = time_embedding_curve(w1, b1, w2, b2, grid_size, compute_device)
+        u, s, vh = torch.linalg.svd(curve, full_matrices=False)
+    except (torch.cuda.OutOfMemoryError, RuntimeError):
+        cleanup_after_operation()
+        compute_device = torch.device("cpu")
+        curve = time_embedding_curve(w1, b1, w2, b2, grid_size, compute_device)
+        u, s, vh = torch.linalg.svd(curve, full_matrices=False)
+
+    target_rank = min(rank, vh.shape[0])
+    table = u[:, :target_rank] * s[:target_rank]
+    basis = vh[:target_rank].T
+    error = ((table @ basis.T - curve).norm() / curve.norm()).item()
+    return table.float().cpu().contiguous(), basis.float().cpu().contiguous(), error
+
+
+def _detect_time_embedder_keys(all_keys: list[str]) -> tuple[str, dict[str, str]]:
+    for prefix in ("", "diffusion_model.", "model.diffusion_model."):
+        candidates = {
+            "w1": f"{prefix}time_embedder.proj_in.weight",
+            "b1": f"{prefix}time_embedder.proj_in.bias",
+            "w2": f"{prefix}time_embedder.proj_out.weight",
+            "b2": f"{prefix}time_embedder.proj_out.bias",
+        }
+        if all(k in all_keys for k in candidates.values()):
+            return prefix, candidates
+    raise KeyError(
+        "Model is missing required time_embedder projection tensors (proj_in/proj_out)."
+    )
+
+
+def _is_adaln(key: str) -> bool:
+    return ".adaln_proj.linear." in key or key.startswith("adaln_proj.linear.")
+
+
+def fold_minimax_h3_diffusion_model(
+    input_path: str,
+    output_path: str,
+    *,
+    exclude_patterns: str = "",
+    include_mode: bool = False,
+    discard_patterns: str = "",
+    glob_patterns: bool = False,
+    curve_rank: int = 8,
+    curve_grid: int = 1025,
+    process_device: str = "cuda",
+) -> str:
+    """Stream and fold MiniMax H3 AdaLN linears in a diffusion model."""
+    if os.path.normcase(os.path.abspath(input_path)) == os.path.normcase(
+        os.path.abspath(output_path)
+    ):
+        raise ValueError("Output path must differ from the source model path.")
+
+    device = torch.device(
+        process_device if torch.cuda.is_available() and process_device == "cuda" else "cpu"
+    )
+
+    compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
+    compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
+
+    handler = MemoryEfficientSafeOpen(input_path, low_memory=True)
+    try:
+        inspect_low_bit_input(handler, f"Model ({input_path})", "MiniMax H3 Fold AdaLN")
+        all_keys = list(handler.keys())
+        prefix, time_keys = _detect_time_embedder_keys(all_keys)
+
+        # Step 1: Load time embedder and calculate basis
+        stream = handler.async_stream(
+            list(time_keys.values()),
+            batch_size=1,
+            prefetch_batches=1,
+            pin_memory=False,
+        )
+        time_tensors = {}
+        try:
+            for batch in stream:
+                for k, t in batch:
+                    time_tensors[k] = t
+            table, basis, error = _fold_curve_basis(
+                time_tensors[time_keys["w1"]],
+                time_tensors[time_keys["b1"]],
+                time_tensors[time_keys["w2"]],
+                time_tensors[time_keys["b2"]],
+                curve_grid,
+                curve_rank,
+                device,
+            )
+        finally:
+            time_tensors.clear()
+            for k in time_keys.values():
+                handler.mark_processed(k)
+            stream.close()
+
+        # Determine which AdaLN layers are folded vs kept full width
+        any_adaln_kept_unfolded = False
+        for k in all_keys:
+            if _is_adaln(k) and not _matches_any_pattern(k, compiled_discards, glob_patterns):
+                matched = _matches_any_pattern(k, compiled_excludes, glob_patterns)
+                should_fold = matched if include_mode else not matched
+                if not should_fold:
+                    any_adaln_kept_unfolded = True
+                    break
+
+        metadata = dict(handler.metadata() or {})
+        if "config" in metadata:
+            try:
+                cfg = json.loads(metadata["config"])
+                if isinstance(cfg, dict):
+                    tf = cfg.setdefault("transformer", {})
+                    tf["adaln_curve_grid"] = table.shape[0]
+                    tf["time_embed_dim"] = table.shape[1]
+                    metadata["config"] = json.dumps(cfg, separators=(",", ":"))
+            except Exception:
+                pass
+        metadata["modelutils_adaln_fold"] = f"rank{table.shape[1]}_curve_grid_{table.shape[0]}"
+
+        table_key = f"{prefix}adaln_t_table"
+        time_embedder_prefix = f"{prefix}time_embedder."
+
+        counts = {
+            "folded": 0,
+            "preserved": 0,
+            "discarded": 0,
+            "time_embedder_dropped": 0,
+        }
+
+        basis_device = basis.to(device)
+
+        with atomic_uel_writer(output_path, metadata) as writer:
+            # Emit adaln_t_table
+            writer.write_batch([(table_key, table)])
+
+            # Filter remaining keys
+            remaining_keys = [k for k in all_keys if k != table_key]
+            progress = comfy.utils.ProgressBar(len(remaining_keys))
+
+            for key in remaining_keys:
+                if _matches_any_pattern(key, compiled_discards, glob_patterns):
+                    counts["discarded"] += 1
+                    handler.mark_processed(key)
+                    progress.update(1)
+                    continue
+
+                if key.startswith(time_embedder_prefix) and not any_adaln_kept_unfolded:
+                    counts["time_embedder_dropped"] += 1
+                    handler.mark_processed(key)
+                    progress.update(1)
+                    continue
+
+                # Stream this tensor
+                tstream = handler.async_stream(
+                    [key], batch_size=1, prefetch_batches=1, pin_memory=False
+                )
+                try:
+                    batch = next(tstream)
+                    tensor = batch[0][1]
+
+                    if _is_adaln(key):
+                        matched = _matches_any_pattern(key, compiled_excludes, glob_patterns)
+                        should_fold = matched if include_mode else not matched
+                        if should_fold:
+                            if key.endswith(".weight"):
+                                if tensor.ndim != 2 or tensor.shape[1] != basis.shape[0]:
+                                    raise ValueError(
+                                        f"{key} has shape {tuple(tensor.shape)}, expected [*, {basis.shape[0]}]"
+                                    )
+                                folded_w = (
+                                    tensor.to(device, torch.float32) @ basis_device
+                                ).to(torch.float16).cpu().contiguous()
+                                writer.write_batch([(key, folded_w)])
+                            else:
+                                writer.write_batch([(key, tensor.to(torch.float16).contiguous())])
+                            counts["folded"] += 1
+                        else:
+                            writer.write_batch([(key, tensor.contiguous())])
+                            counts["preserved"] += 1
+                    else:
+                        writer.write_batch([(key, tensor.contiguous())])
+                        counts["preserved"] += 1
+                finally:
+                    handler.mark_processed(key)
+                    tstream.close()
+
+                progress.update(1)
+    finally:
+        handler.__exit__(None, None, None)
+        cleanup_after_operation()
+
+    report_lines = [
+        "MiniMax H3 AdaLN Fold complete",
+        f"Curve rank: {curve_rank}, grid points: {curve_grid}",
+        f"Reconstruction relative error: {error:.4e}",
+        f"AdaLN layers folded: {counts['folded']}",
+        f"Tensors preserved: {counts['preserved']}",
+        f"Tensors discarded: {counts['discarded']}",
+        f"Time embedder dropped: {counts['time_embedder_dropped']}",
+        f"Output size: {os.path.getsize(output_path)} bytes",
+        f"Output: {output_path}",
+    ]
+    return "\n".join(report_lines)
+
+
+class MiniMaxH3FoldAdaLN(io.ComfyNode):
+    """Folds MiniMax H3 AdaLN conditioning linears onto a low-rank curve basis.
+
+    Technical details:
+    Full-width AdaLN linear weights [*, 2688] are projected onto a rank-8 SVD basis
+    of the silu(time_embedder) curve sampled across a 1025-point grid of t in [0, 1]:
+      X = silu(t_emb(t)) [1025, 2688] = U S V^T
+      adaln_t_table = U_8 S_8 (FP32),  W' = W V_8 (FP16)
+    The time embedder is dropped and replaced with adaln_t_table [1025, 8], reducing
+    model size substantially with negligible reconstruction error (< 1e-4).
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3FoldAdaLN",
+            display_name="MiniMax H3 Fold AdaLN",
+            category="ModelUtils/DiffusionModels",
+            description="Compresses large MiniMax H3 models by folding time-conditioning layers into a small lookup table, matching official ComfyUI checkpoint formats.",
+            inputs=[
+                io.Combo.Input(
+                    "model_name",
+                    options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Select the original/uncompressed MiniMax H3 model file to fold.",
+                ),
+                io.String.Input(
+                    "output_filename",
+                    default="minimax_h3_folded",
+                    tooltip="Name for the new folded model file, saved in your diffusion_models folder.",
+                ),
+                io.String.Input(
+                    "discard_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip="Layers to completely remove from the final file (one name or pattern per line). Useful to strip parts you do not need.",
+                ),
+                io.Boolean.Input(
+                    "glob_patterns",
+                    default=False,
+                    tooltip="Enable simple wildcard matching (* to match anything). If unchecked, uses regular expressions.",
+                ),
+                io.Int.Input(
+                    "curve_rank",
+                    default=8,
+                    min=1,
+                    max=64,
+                    step=1,
+                    tooltip="Compression size for the time schedule table. Leave at 8 to match standard ComfyUI MiniMax H3 models.",
+                ),
+                io.Int.Input(
+                    "curve_grid",
+                    default=1025,
+                    min=65,
+                    max=4097,
+                    step=64,
+                    tooltip="Sampling resolution for the time schedule curve. Leave at 1025 to match standard ComfyUI MiniMax H3 models.",
+                ),
+                io.Combo.Input(
+                    "process_device",
+                    options=["cuda", "cpu"],
+                    default="cuda",
+                    tooltip="Hardware device used to perform the math. Use cuda (GPU) for speed, or cpu if you run low on VRAM.",
+                ),
+                io.String.Input(
+                    "exclude_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip="Specific layers to skip folding, keeping them at their original full size (one name or pattern per line).",
+                ),
+                io.Boolean.Input(
+                    "include_mode",
+                    default=False,
+                    tooltip="Invert the exclude filter: only fold layers matching the pattern, and keep everything else untouched.",
+                ),
+            ],
+            outputs=[
+                io.AnyType.Output(display_name="output_path"),
+                io.String.Output(display_name="report"),
+            ],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        model_name: str,
+        output_filename: str,
+        discard_patterns: str = "",
+        glob_patterns: bool = False,
+        curve_rank: int = 8,
+        curve_grid: int = 1025,
+        process_device: str = "cuda",
+        exclude_patterns: str = "",
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        input_path = folder_paths.get_full_path_or_raise(
+            "diffusion_models", model_name
+        )
+        output_path, output_name = canonical_model_artifact_path(
+            "diffusion_models", output_filename
+        )
+        report = fold_minimax_h3_diffusion_model(
+            input_path,
+            output_path,
+            exclude_patterns=exclude_patterns,
+            include_mode=include_mode,
+            discard_patterns=discard_patterns,
+            glob_patterns=glob_patterns,
+            curve_rank=curve_rank,
+            curve_grid=curve_grid,
+            process_device=process_device,
+        )
+        return io.NodeOutput(output_name, report)
+
+
+def unfold_minimax_h3_diffusion_model(
+    model_a_path: str,
+    model_b_path: str,
+    output_path: str,
+    *,
+    exclude_patterns: str = "",
+    include_mode: bool = False,
+    discard_patterns: str = "",
+    glob_patterns: bool = False,
+    curve_rank: int = 8,
+    curve_grid: int = 1025,
+    process_device: str = "cuda",
+) -> str:
+    """Stream and reconstruct unfolded MiniMax H3 AdaLN linears using a reference base model."""
+    if os.path.normcase(os.path.abspath(model_a_path)) == os.path.normcase(
+        os.path.abspath(output_path)
+    ):
+        raise ValueError("Output path must differ from the folded model (Model A) path.")
+    if os.path.normcase(os.path.abspath(model_b_path)) == os.path.normcase(
+        os.path.abspath(output_path)
+    ):
+        raise ValueError("Output path must differ from the base model (Model B) path.")
+
+    device = torch.device(
+        process_device if torch.cuda.is_available() and process_device == "cuda" else "cpu"
+    )
+
+    compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
+    compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
+
+    handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=True)
+    handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=True)
+
+    try:
+        inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "MiniMax H3 Unfold AdaLN")
+        inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "MiniMax H3 Unfold AdaLN")
+
+        all_keys_a = list(handler_a.keys())
+        all_keys_b = list(handler_b.keys())
+
+        prefix_b, time_keys_b = _detect_time_embedder_keys(all_keys_b)
+        prefix_a = "diffusion_model." if any(k.startswith("diffusion_model.") for k in all_keys_a) else ""
+
+        # Step 1: Compute basis and table from base model (Model B)
+        tstream_b = handler_b.async_stream(
+            list(time_keys_b.values()),
+            batch_size=1,
+            prefetch_batches=1,
+            pin_memory=False,
+        )
+        time_tensors_b = {}
+        try:
+            for batch in tstream_b:
+                for k, t in batch:
+                    time_tensors_b[k] = t
+            _, basis, error = _fold_curve_basis(
+                time_tensors_b[time_keys_b["w1"]],
+                time_tensors_b[time_keys_b["b1"]],
+                time_tensors_b[time_keys_b["w2"]],
+                time_tensors_b[time_keys_b["b2"]],
+                curve_grid,
+                curve_rank,
+                device,
+            )
+        finally:
+            for k in time_keys_b.values():
+                handler_b.mark_processed(k)
+            tstream_b.close()
+
+        basis_device = basis.to(device)
+
+        metadata = dict(handler_a.metadata() or {})
+        if "config" in metadata:
+            try:
+                cfg = json.loads(metadata["config"])
+                if isinstance(cfg, dict):
+                    tf = cfg.setdefault("transformer", {})
+                    tf.pop("adaln_curve_grid", None)
+                    tf["time_embed_dim"] = basis.shape[0]
+                    metadata["config"] = json.dumps(cfg, separators=(",", ":"))
+            except Exception:
+                pass
+        metadata["modelutils_adaln_unfold"] = f"reconstructed_rank{basis.shape[1]}_with_base"
+
+        table_key_a = f"{prefix_a}adaln_t_table"
+
+        counts = {
+            "unfolded": 0,
+            "preserved": 0,
+            "discarded": 0,
+            "time_embedder_restored": 0,
+        }
+
+        with atomic_uel_writer(output_path, metadata) as writer:
+            # Emit the 4 time_embedder tensors restored from Model B
+            for _, base_k in time_keys_b.items():
+                target_k = f"{prefix_a}{base_k.removeprefix(prefix_b)}"
+                if not _matches_any_pattern(target_k, compiled_discards, glob_patterns):
+                    writer.write_batch([(target_k, time_tensors_b[base_k])])
+                    counts["time_embedder_restored"] += 1
+
+            time_tensors_b.clear()
+
+            remaining_keys_a = [k for k in all_keys_a if k != table_key_a]
+            progress = comfy.utils.ProgressBar(len(remaining_keys_a))
+
+            for key in remaining_keys_a:
+                if _matches_any_pattern(key, compiled_discards, glob_patterns):
+                    counts["discarded"] += 1
+                    handler_a.mark_processed(key)
+                    progress.update(1)
+                    continue
+
+                astream = handler_a.async_stream(
+                    [key], batch_size=1, prefetch_batches=1, pin_memory=False
+                )
+                try:
+                    tensor_a = next(astream)[0][1]
+
+                    if _is_adaln(key):
+                        matched = _matches_any_pattern(key, compiled_excludes, glob_patterns)
+                        should_unfold = matched if include_mode else not matched
+                        if should_unfold and key.endswith(".weight"):
+                            base_key = f"{prefix_b}{key.removeprefix(prefix_a)}"
+                            if base_key not in all_keys_b:
+                                raise KeyError(
+                                    f"Matching base layer '{base_key}' not found in Model B."
+                                )
+                            bstream = handler_b.async_stream(
+                                [base_key], batch_size=1, prefetch_batches=1, pin_memory=False
+                            )
+                            try:
+                                base_w = next(bstream)[0][1]
+                                if base_w.ndim != 2 or base_w.shape[1] != basis.shape[0]:
+                                    raise ValueError(
+                                        f"{base_key} in Model B has shape {tuple(base_w.shape)}, expected [*, {basis.shape[0]}]"
+                                    )
+                                bw_f32 = base_w.to(device, torch.float32)
+                                w_folded = tensor_a.to(device, torch.float32)
+                                w_diff = w_folded - bw_f32 @ basis_device
+                                w_unfolded = (bw_f32 + w_diff @ basis_device.T).to(base_w.dtype).cpu().contiguous()
+                                writer.write_batch([(key, w_unfolded)])
+                                counts["unfolded"] += 1
+                            finally:
+                                handler_b.mark_processed(base_key)
+                                bstream.close()
+                        else:
+                            writer.write_batch([(key, tensor_a.contiguous())])
+                            counts["preserved"] += 1
+                    else:
+                        writer.write_batch([(key, tensor_a.contiguous())])
+                        counts["preserved"] += 1
+                finally:
+                    handler_a.mark_processed(key)
+                    astream.close()
+
+                progress.update(1)
+    finally:
+        handler_a.__exit__(None, None, None)
+        handler_b.__exit__(None, None, None)
+        cleanup_after_operation()
+
+    report_lines = [
+        "MiniMax H3 AdaLN Unfold complete",
+        f"Curve rank: {curve_rank}, grid points: {curve_grid}",
+        f"Base basis reconstruction error: {error:.4e}",
+        f"AdaLN layers unfolded: {counts['unfolded']}",
+        f"Time embedder tensors restored: {counts['time_embedder_restored']}",
+        f"Tensors preserved: {counts['preserved']}",
+        f"Tensors discarded: {counts['discarded']}",
+        f"Output size: {os.path.getsize(output_path)} bytes",
+        f"Output: {output_path}",
+    ]
+    return "\n".join(report_lines)
+
+
+class MiniMaxH3UnfoldAdaLN(io.ComfyNode):
+    """Reconstruct full-width AdaLN linears [*, 2688] from a folded model and reference base model."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3UnfoldAdaLN",
+            display_name="MiniMax H3 Unfold AdaLN",
+            category="ModelUtils/DiffusionModels",
+            description="Reconstructs full-width MiniMax H3 AdaLN linears [*, 2688] from a folded model (Model A) using a reference non-folded base model (Model B) to supply the time embedder and orthogonal dimensions.",
+            inputs=[
+                io.Combo.Input(
+                    "model_a",
+                    options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Select the folded MiniMax H3 model (Model A) to unfold.",
+                ),
+                io.Combo.Input(
+                    "model_b",
+                    options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Select the reference original/unfolded MiniMax H3 base model (Model B) to supply the time embedder and missing dimensions.",
+                ),
+                io.String.Input(
+                    "output_filename",
+                    default="minimax_h3_unfolded",
+                    tooltip="Name for the reconstructed unfolded model file, saved in your diffusion_models folder.",
+                ),
+                io.String.Input(
+                    "discard_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip="Layers to completely remove from the final file (one name or pattern per line). Useful to strip parts you do not need.",
+                ),
+                io.Boolean.Input(
+                    "glob_patterns",
+                    default=False,
+                    tooltip="Enable simple wildcard matching (* to match anything). If unchecked, uses regular expressions.",
+                ),
+                io.Int.Input(
+                    "curve_rank",
+                    default=8,
+                    min=1,
+                    max=64,
+                    step=1,
+                    tooltip="Basis rank used when the model was folded (default 8, matching standard ComfyUI H3 checkpoints).",
+                ),
+                io.Int.Input(
+                    "curve_grid",
+                    default=1025,
+                    min=65,
+                    max=4097,
+                    step=64,
+                    tooltip="Sampling resolution used when the model was folded (default 1025, matching standard ComfyUI H3 checkpoints).",
+                ),
+                io.Combo.Input(
+                    "process_device",
+                    options=["cuda", "cpu"],
+                    default="cuda",
+                    tooltip="Hardware device used to perform the math. Use cuda (GPU) for speed, or cpu if you run low on VRAM.",
+                ),
+                io.String.Input(
+                    "exclude_patterns",
+                    default="",
+                    multiline=True,
+                    tooltip="Specific layers to skip unfolding, keeping them at their folded size (one name or pattern per line).",
+                ),
+                io.Boolean.Input(
+                    "include_mode",
+                    default=False,
+                    tooltip="Invert the exclude filter: only unfold layers matching the pattern, and keep everything else untouched.",
+                ),
+            ],
+            outputs=[
+                io.AnyType.Output(display_name="output_path"),
+                io.String.Output(display_name="report"),
+            ],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        model_a: str,
+        model_b: str,
+        output_filename: str,
+        discard_patterns: str = "",
+        glob_patterns: bool = False,
+        curve_rank: int = 8,
+        curve_grid: int = 1025,
+        process_device: str = "cuda",
+        exclude_patterns: str = "",
+        include_mode: bool = 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
+        )
+        output_path, output_name = canonical_model_artifact_path(
+            "diffusion_models", output_filename
+        )
+        report = unfold_minimax_h3_diffusion_model(
+            model_a_path,
+            model_b_path,
+            output_path,
+            exclude_patterns=exclude_patterns,
+            include_mode=include_mode,
+            discard_patterns=discard_patterns,
+            glob_patterns=glob_patterns,
+            curve_rank=curve_rank,
+            curve_grid=curve_grid,
+            process_device=process_device,
+        )
+        return io.NodeOutput(output_name, report)

nodes/minimax_h3_lora_convert.py

@@ -0,0 +1,335 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+import re
+
+import comfy.utils
+from comfy_api.latest import io
+import folder_paths
+import torch
+from unifiedefficientloader import MemoryEfficientSafeOpen
+
+from .artifact_paths import canonical_model_artifact_path
+from .quantization_guard import inspect_low_bit_input
+from .uel_io import atomic_uel_writer
+
+_FACTOR_RE = re.compile(
+    r"^(?P<module>.+)\.(?:lora_(?P<factor_ab>[AB])|lora[._](?P<factor_word>down|up))(?:\.default)?(?:\.weight)?$"
+)
+_BLOCK_RE = re.compile(r"^transformer_blocks\.(?P<index>\d+)\.(?P<tail>.+)$")
+_REFINER_RE = re.compile(r"^token_refiner\.refiner_blocks\.(?P<index>\d+)\.(?P<tail>.+)$")
+_QKV_RE = re.compile(r"^attn\.to_(?P<part>[qkv])$")
+
+_DIRECT_TAILS = {
+    "attn.to_out.0": "attn.out_proj",
+    "ff.net.0.proj": "mlp.fc1",
+    "ff.net.2": "mlp.fc2",
+}
+
+_EXTRA_PREFIX_MAP = {
+    "context_embedder": "condition_proj",
+    "proj_in": "video_patch_proj",
+    "proj_out": "final_layer.video_out",
+    "audio_proj_in": "audio_patch_proj",
+    "audio_proj_out": "final_layer.audio_out",
+    "norm_out.norm": "final_layer.norm",
+    "norm_out.linear": "final_layer.adaln_proj.linear",
+    "time_embedder.linear_1": "time_embedder.proj_in",
+    "time_embedder.linear_2": "time_embedder.proj_out",
+}
+
+
+@dataclass(frozen=True)
+class _ConversionUnit:
+    target: str
+    sources: tuple[str, ...]
+    qkv: bool
+    is_swiglu: bool = False
+
+
+def _target_module(module: str) -> tuple[str, str | None, bool]:
+    """Map Diffusers module name to ComfyUI native target module.
+
+    Returns:
+        (target_module, qkv_part, is_swiglu_fc1)
+    """
+    clean_mod = module.removeprefix("base_model.model.").removeprefix("transformer.")
+    match = _BLOCK_RE.fullmatch(clean_mod)
+    if match:
+        prefix = f"diffusion_model.blocks.{match.group('index')}"
+        tail = match.group("tail")
+    else:
+        match = _REFINER_RE.fullmatch(clean_mod)
+        if match:
+            prefix = f"diffusion_model.token_refiner.blocks.{match.group('index')}"
+            tail = match.group("tail")
+        else:
+            for old_pfx, new_pfx in _EXTRA_PREFIX_MAP.items():
+                if clean_mod.startswith(old_pfx):
+                    target = f"diffusion_model.{new_pfx}{clean_mod[len(old_pfx):]}"
+                    return target, None, False
+            raise ValueError(f"Unsupported MiniMax H3 Diffusers module: {module}")
+
+    qkv = _QKV_RE.fullmatch(tail)
+    if qkv:
+        return f"{prefix}.attn.qkv_proj", qkv.group("part"), False
+
+    if tail in _DIRECT_TAILS:
+        target = f"{prefix}.{_DIRECT_TAILS[tail]}"
+        is_swiglu = (tail == "ff.net.0.proj")
+        return target, None, is_swiglu
+
+    raise ValueError(f"Unsupported MiniMax H3 Diffusers module tail: {module}")
+
+
+def _build_plan(keys: list[str]) -> list[_ConversionUnit]:
+    """Build conversion work units from input keys."""
+    factors: dict[str, dict[str, str]] = {}
+    for key in keys:
+        match = _FACTOR_RE.fullmatch(key)
+        if not match:
+            raise ValueError(
+                f"Input is not a compatible MiniMax H3 Diffusers PEFT LoRA; unsupported key: {key}"
+            )
+        mod = match.group("module")
+        factor = match.group("factor_ab") or (
+            "A" if match.group("factor_word") == "down" else "B"
+        )
+        pair = factors.setdefault(mod, {})
+        if factor in pair:
+            raise ValueError(f"Duplicate LoRA factor for: {mod}")
+        pair[factor] = key
+
+    missing = [mod for mod, pair in factors.items() if set(pair.keys()) != {"A", "B"}]
+    if missing:
+        raise ValueError(f"Incomplete LoRA A/B pair for: {missing[0]}")
+
+    direct_units: list[_ConversionUnit] = []
+    qkv_groups: dict[str, dict[str, tuple[str, str]]] = {}
+
+    for mod, pair in factors.items():
+        target, qkv_part, is_swiglu = _target_module(mod)
+        source_pair = (pair["A"], pair["B"])
+        if qkv_part is None:
+            direct_units.append(
+                _ConversionUnit(
+                    target=target,
+                    sources=source_pair,
+                    qkv=False,
+                    is_swiglu=is_swiglu,
+                )
+            )
+        else:
+            qkv_groups.setdefault(target, {})[qkv_part] = source_pair
+
+    qkv_units: list[_ConversionUnit] = []
+    for target, parts in qkv_groups.items():
+        if set(parts.keys()) != {"q", "k", "v"}:
+            raise ValueError(f"Incomplete Q/K/V LoRA group for: {target}")
+        sources = tuple(k for part in ("q", "k", "v") for k in parts[part])
+        qkv_units.append(_ConversionUnit(target=target, sources=sources, qkv=True))
+
+    units = sorted(direct_units + qkv_units, key=lambda u: u.target)
+    targets = [u.target for u in units]
+    if len(targets) != len(set(targets)):
+        raise ValueError("Conversion would produce duplicate ComfyUI module keys")
+
+    return units
+
+
+def _validate_pair(down: torch.Tensor, up: torch.Tensor, label: str) -> None:
+    if down.ndim != 2 or up.ndim != 2:
+        raise ValueError(f"LoRA factors for '{label}' must be 2D matrices")
+    if down.shape[0] != up.shape[1]:
+        raise ValueError(
+            f"LoRA factor rank mismatch for '{label}': down={list(down.shape)}, up={list(up.shape)}"
+        )
+
+
+def _convert_qkv(
+    tensors: dict[str, torch.Tensor],
+    sources: tuple[str, ...],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse separate Q, K, V LoRA factors into unified block-diagonal qkv_proj factors."""
+    pairs = [
+        (tensors[sources[i]], tensors[sources[i + 1]])
+        for i in range(0, 6, 2)
+    ]
+    for part, (down, up) in zip(("Q", "K", "V"), pairs):
+        _validate_pair(down, up, part)
+
+    in_dims = {down.shape[1] for down, _ in pairs}
+    if len(in_dims) != 1:
+        raise ValueError("Q/K/V LoRA input dimensions do not match")
+
+    dtypes = {t.dtype for pair in pairs for t in pair}
+    if len(dtypes) != 1:
+        raise ValueError("Q/K/V LoRA dtypes do not match")
+
+    source_dtype = pairs[0][0].dtype
+    ranks = [down.shape[0] for down, _ in pairs]
+    output_dims = [up.shape[0] for _, up in pairs]
+    total_rank = sum(ranks)
+    total_out = sum(output_dims)
+
+    # Concatenate down matrices vertically: [total_rank, in_dim]
+    exact_down = torch.cat([down for down, _ in pairs], dim=0).contiguous()
+
+    # Place up matrices block-diagonally: [total_out, total_rank]
+    exact_up = torch.zeros(
+        (total_out, total_rank),
+        dtype=source_dtype,
+        device=pairs[0][1].device,
+    )
+    row = 0
+    col = 0
+    for (_, part_up), out_dim, part_rank in zip(pairs, output_dims, ranks):
+        exact_up[row : row + out_dim, col : col + part_rank].copy_(part_up)
+        row += out_dim
+        col += part_rank
+
+    return exact_down, exact_up.contiguous()
+
+
+def _convert_swiglu(
+    down: torch.Tensor, up: torch.Tensor, label: str
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Swap SwiGLU output halves from Diffusers [value; gate] to ComfyUI [gate; value]."""
+    _validate_pair(down, up, label)
+    if up.shape[0] % 2 != 0:
+        raise ValueError(
+            f"SwiGLU LoRA B factor for '{label}' has odd output dimension {up.shape[0]}"
+        )
+    val_b, gate_b = up.chunk(2, dim=0)
+    swapped_up = torch.cat([gate_b, val_b], dim=0).contiguous()
+    return down.contiguous(), swapped_up
+
+
+def convert_minimax_h3_diffusers_lora(
+    input_path: str,
+    output_path: str,
+) -> str:
+    """Stream and convert Diffusers MiniMax H3 LoRA into ComfyUI native format."""
+    if os.path.normcase(os.path.abspath(input_path)) == os.path.normcase(
+        os.path.abspath(output_path)
+    ):
+        raise ValueError("Output path must differ from the source LoRA path")
+
+    handler = MemoryEfficientSafeOpen(input_path, low_memory=True)
+    try:
+        inspect_low_bit_input(handler, f"LoRA ({input_path})", "MiniMax H3 LoRA Convert")
+        keys = list(handler.keys())
+        source_count = len(keys)
+        units = _build_plan(keys)
+
+        metadata = {str(k): str(v) for k, v in (handler.metadata() or {}).items()}
+        metadata.update(
+            {
+                "modelutils_conversion": "diffusers_minimax_h3_to_comfyui",
+                "modelutils_qkv_fusion": "lossless_block_diagonal",
+                "modelutils_mlp_swiglu": "swapped_gate_value",
+            }
+        )
+
+        progress = comfy.utils.ProgressBar(len(units))
+        output_count = 0
+        fused_qkv_count = 0
+
+        with atomic_uel_writer(output_path, metadata) as writer:
+            for unit in units:
+                stream = handler.async_stream(
+                    list(unit.sources),
+                    batch_size=1,
+                    prefetch_batches=1,
+                    pin_memory=False,
+                )
+                tensors = {}
+                try:
+                    for batch in stream:
+                        for k, t in batch:
+                            tensors[k] = t
+
+                    if unit.qkv:
+                        down, up = _convert_qkv(tensors, unit.sources)
+                        fused_qkv_count += 1
+                    elif unit.is_swiglu:
+                        down_key, up_key = unit.sources
+                        down, up = _convert_swiglu(
+                            tensors[down_key], tensors[up_key], unit.target
+                        )
+                    else:
+                        down_key, up_key = unit.sources
+                        down = tensors[down_key].contiguous()
+                        up = tensors[up_key].contiguous()
+                        _validate_pair(down, up, unit.target)
+
+                    writer.write_batch(
+                        [
+                            (f"{unit.target}.lora_A.weight", down),
+                            (f"{unit.target}.lora_B.weight", up),
+                        ]
+                    )
+                    output_count += 2
+                finally:
+                    tensors.clear()
+                    for key in unit.sources:
+                        handler.mark_processed(key)
+                    stream.close()
+
+                progress.update(1)
+    finally:
+        handler.__exit__(None, None, None)
+
+    lines = [
+        "MiniMax H3 Diffusers LoRA conversion complete",
+        f"Source tensors: {source_count}",
+        f"Output tensors: {output_count}",
+        f"Fused QKV modules: {fused_qkv_count}",
+        f"Output size: {os.path.getsize(output_path)} bytes",
+        f"Output: {output_path}",
+    ]
+    return "\n".join(lines)
+
+
+class MiniMaxH3DiffusersLoRAConvert(io.ComfyNode):
+    """Convert Diffusers MiniMax H3 LoRAs to ComfyUI's native layout."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3DiffusersLoRAConvert",
+            display_name="MiniMax H3 Diffusers LoRA Convert",
+            category="ModelUtils/LoRA",
+            description="Converts Diffusers/PEFT MiniMax H3 LoRAs to ComfyUI's native fused layout (lossless block-diagonal QKV and SwiGLU-aligned MLP).",
+            inputs=[
+                io.Combo.Input(
+                    "lora_name",
+                    options=folder_paths.get_filename_list("loras"),
+                    tooltip="Diffusers-format MiniMax H3 LoRA to convert.",
+                ),
+                io.String.Input(
+                    "output_filename",
+                    default="minimax_h3_comfyui",
+                    tooltip="Output filename without extension, written under ComfyUI's LoRA directory.",
+                ),
+            ],
+            outputs=[
+                io.AnyType.Output(display_name="output_path"),
+                io.String.Output(display_name="conversion_report"),
+            ],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        lora_name: str,
+        output_filename: str,
+    ) -> io.NodeOutput:
+        input_path = folder_paths.get_full_path_or_raise("loras", lora_name)
+        output_path, output_name = canonical_model_artifact_path(
+            "loras", output_filename
+        )
+        report = convert_minimax_h3_diffusers_lora(input_path, output_path)
+        return io.NodeOutput(output_name, report)

nodes/model_info_nodes.py

@@ -0,0 +1,241 @@
+import glob
+import os
+import numpy as np
+from PIL import Image, ImageOps
+import torch
+import folder_paths
+from comfy_api.latest import io
+
+
+def get_potential_preview_files(base_path: str) -> list[str]:
+    exts = ["png", "jpg", "jpeg", "webp", "mp4"]
+    files = []
+    for ext in exts:
+        files.append(f"{base_path}.preview.{ext}")
+        files.append(f"{base_path}.{ext}")
+    return files
+
+
+def load_image_tensor(path: str):
+    if not os.path.exists(path):
+        return None
+    try:
+        img = Image.open(path)
+        img = ImageOps.exif_transpose(img)
+        image = img.convert("RGB")
+        image_np = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(image_np)[None,]
+    except Exception as e:
+        print(f"Error loading image {path}: {e}")
+        return None
+
+
+def get_model_workflows(base_path: str) -> list[str]:
+    workflows = []
+    main_json = f"{base_path}.json"
+    if os.path.exists(main_json):
+        workflows.append(main_json)
+
+    preview_json = f"{base_path}.preview.json"
+    if os.path.exists(preview_json) and preview_json not in workflows:
+        workflows.append(preview_json)
+
+    example_jsons = glob.glob(f"{base_path}.example.*.json")
+    example_jsons.sort()
+    for ej in example_jsons:
+        if ej not in workflows:
+            workflows.append(ej)
+
+    return workflows
+
+
+def get_model_metadata_file(base_path: str) -> str | None:
+    civitai_info = f"{base_path}.civitai.info"
+    if os.path.exists(civitai_info):
+        return civitai_info
+
+    metadata_json = f"{base_path}.metadata.json"
+    if os.path.exists(metadata_json):
+        return metadata_json
+
+    return None
+
+class BaseModelInfoLoader(io.ComfyNode):
+    CATEGORY = "ModelUtils/Info"
+
+    @classmethod
+    def get_info(cls, category, name, workflow_index):
+        full_path = folder_paths.get_full_path(category, name)
+        if not full_path or not os.path.exists(full_path):
+            return (None, "{}", "{}")
+
+        base_path = os.path.splitext(full_path)[0]
+
+        # 1. Image Preview
+        image_tensor = None
+        previews = get_potential_preview_files(base_path)
+        for p in previews:
+            if os.path.exists(p) and not p.lower().endswith('.mp4'):
+                image_tensor = load_image_tensor(p)
+                if image_tensor is not None:
+                    break
+
+        # 2. Workflow
+        workflow_json = "{}"
+        workflows = get_model_workflows(base_path)
+        if workflows and 0 <= workflow_index < len(workflows):
+            try:
+                with open(workflows[workflow_index], 'r', encoding='utf-8') as f:
+                    workflow_json = f.read()
+            except Exception as e:
+                print(f"Error reading workflow {workflows[workflow_index]}: {e}")
+
+        # 3. Metadata
+        metadata_json = "{}"
+        metadata_file = get_model_metadata_file(base_path)
+        if metadata_file:
+            try:
+                with open(metadata_file, 'r', encoding='utf-8') as f:
+                    metadata_json = f.read()
+            except Exception as e:
+                print(f"Error reading metadata {metadata_file}: {e}")
+
+        return (image_tensor, workflow_json, metadata_json)
+
+class CheckpointInfoLoader(BaseModelInfoLoader):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Checkpoint Info Loader",
+            category=cls.CATEGORY,
+            inputs=[
+                io.Combo.Input("checkpoint", options=folder_paths.get_filename_list("checkpoints"), tooltip="Checkpoint whose downloaded information and example workflow will be returned."),
+                io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
+            ],
+            outputs=[
+                io.Image.Output(display_name="preview"),
+                io.String.Output(display_name="workflow_json"),
+                io.String.Output(display_name="metadata_json")
+            ]
+        )
+
+    @classmethod
+    def execute(cls, checkpoint, workflow_index):
+        res = cls.get_info("checkpoints", checkpoint, workflow_index)
+        return io.NodeOutput(res[0], res[1], res[2])
+
+class LoRAInfoLoader(BaseModelInfoLoader):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="LoRA Info Loader",
+            category=cls.CATEGORY,
+            inputs=[
+                io.Combo.Input("lora", options=folder_paths.get_filename_list("loras"), tooltip="LoRA whose downloaded information and example workflow will be returned."),
+                io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
+            ],
+            outputs=[
+                io.Image.Output(display_name="preview"),
+                io.String.Output(display_name="workflow_json"),
+                io.String.Output(display_name="metadata_json")
+            ]
+        )
+
+    @classmethod
+    def execute(cls, lora, workflow_index):
+        res = cls.get_info("loras", lora, workflow_index)
+        return io.NodeOutput(res[0], res[1], res[2])
+
+class EmbeddingInfoLoader(BaseModelInfoLoader):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Embedding Info Loader",
+            category=cls.CATEGORY,
+            inputs=[
+                io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings"), tooltip="Embedding whose downloaded information and example workflow will be returned."),
+                io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
+            ],
+            outputs=[
+                io.Image.Output(display_name="preview"),
+                io.String.Output(display_name="workflow_json"),
+                io.String.Output(display_name="metadata_json")
+            ]
+        )
+
+    @classmethod
+    def execute(cls, embedding, workflow_index):
+        res = cls.get_info("embeddings", embedding, workflow_index)
+        return io.NodeOutput(res[0], res[1], res[2])
+
+class VAEInfoLoader(BaseModelInfoLoader):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="VAE Info Loader",
+            category=cls.CATEGORY,
+            inputs=[
+                io.Combo.Input("vae", options=folder_paths.get_filename_list("vae"), tooltip="VAE whose downloaded information and example workflow will be returned."),
+                io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
+            ],
+            outputs=[
+                io.Image.Output(display_name="preview"),
+                io.String.Output(display_name="workflow_json"),
+                io.String.Output(display_name="metadata_json")
+            ]
+        )
+
+    @classmethod
+    def execute(cls, vae, workflow_index):
+        res = cls.get_info("vae", vae, workflow_index)
+        return io.NodeOutput(res[0], res[1], res[2])
+
+class ControlNetInfoLoader(BaseModelInfoLoader):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="ControlNet Info Loader",
+            category=cls.CATEGORY,
+            inputs=[
+                io.Combo.Input("controlnet", options=folder_paths.get_filename_list("controlnet"), tooltip="ControlNet whose downloaded information and example workflow will be returned."),
+                io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
+            ],
+            outputs=[
+                io.Image.Output(display_name="preview"),
+                io.String.Output(display_name="workflow_json"),
+                io.String.Output(display_name="metadata_json")
+            ]
+        )
+
+    @classmethod
+    def execute(cls, controlnet, workflow_index):
+        res = cls.get_info("controlnet", controlnet, workflow_index)
+        return io.NodeOutput(res[0], res[1], res[2])
+
+class DiffusionModelInfoLoader(BaseModelInfoLoader):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id=cls.__name__,
+            display_name="Diffusion Model Info Loader",
+            category=cls.CATEGORY,
+            inputs=[
+                io.Combo.Input("diffusion_model", options=folder_paths.get_filename_list("diffusion_models"), tooltip="Diffusion model whose downloaded information and example workflow will be returned."),
+                io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
+            ],
+            outputs=[
+                io.Image.Output(display_name="preview"),
+                io.String.Output(display_name="workflow_json"),
+                io.String.Output(display_name="metadata_json")
+            ]
+        )
+
+    @classmethod
+    def execute(cls, diffusion_model, workflow_index):
+        res = cls.get_info("diffusion_models", diffusion_model, workflow_index)
+        return io.NodeOutput(res[0], res[1], res[2])

nodes/prunekeys.py

@@ -1,18 +1,32 @@
+from contextlib import closing
 import os
 import re
 import folder_paths
 import comfy.utils
 from tqdm import tqdm
 from comfy_api.latest import io
-from safetensors.torch import save_file
-from safetensors import safe_open
 from .utils import convert_pt_to_safetensors
-from .merger_utils import MemoryEfficientSafeOpen
 from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
 
+from unifiedefficientloader import MemoryEfficientSafeOpen
+from .uel_io import atomic_uel_writer, stream_work_units
+from .artifact_paths import canonical_model_artifact_path
+from .quantization_guard import (
+    UPSTREAM_WRITER_STORAGE_CODES,
+    _is_low_bit,
+    storage_dtype_code,
+    write_preserved_tensor,
+)
 
-def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
-                use_regex: bool, output_filename: str) -> str:
+
+def _prune_keys(
+    model_name: str,
+    model_type: str,
+    keys_to_prune_str: str,
+    use_regex: bool,
+    output_filename: str,
+    include_mode: bool = False,
+) -> str:
     """Shared logic for all PruneKeys nodes."""
     model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
 
@@ -26,48 +40,75 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
     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 {}
-
-    patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
-
-    if not patterns:
+    raw_patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
+    if not raw_patterns:
         raise ValueError("No keys/patterns provided to prune.")
 
-    # Stream tensors and filter on the fly
-    pruned_tensors = {}
-    with MemoryEfficientSafeOpen(model_path_to_load) as handler:
-        all_keys = handler.keys()
-        pbar = comfy.utils.ProgressBar(len(all_keys))
-        for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
-            is_match = False
-            if use_regex:
-                if any(re.search(pattern, key) for pattern in patterns):
-                    is_match = True
+    compiled_patterns = []
+    if use_regex:
+        for pat in raw_patterns:
+            try:
+                compiled_patterns.append(re.compile(pat))
+            except re.error as exc:
+                raise ValueError(f"Invalid regex pattern '{pat}': {exc}") from exc
+
+    model_size_gb = estimate_model_size(model_path_to_load)
+    if model_size_gb:
+        prepare_for_large_operation(model_size_gb * 1.2)
+
+    output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
+
+    handler = MemoryEfficientSafeOpen(model_path_to_load, low_memory=True)
+    try:
+        metadata = (handler.metadata() or {}).copy()
+        all_keys = list(handler.keys())
+        kept_keys = []
+        for key in all_keys:
+            matched = (
+                any(pat.search(key) is not None for pat in compiled_patterns)
+                if use_regex
+                else any(pat in key for pat in raw_patterns)
+            )
+            keep = matched if include_mode else not matched
+            if keep:
+                kept_keys.append(key)
+
+        work_units = []
+        raw_preserve_keys = set()
+        for key in kept_keys:
+            code = storage_dtype_code(handler, key)
+            if _is_low_bit(handler, key) or (
+                code is not None and code not in UPSTREAM_WRITER_STORAGE_CODES
+            ):
+                raw_preserve_keys.add(key)
+                work_units.append((key, {}))
             else:
-                if any(pattern in key for pattern in patterns):
-                    is_match = True
+                work_units.append((key, {0: [key]}))
 
-            if not is_match:
-                pruned_tensors[key] = handler.get_tensor(key)
-            pbar.update(1)
+        pbar = comfy.utils.ProgressBar(len(all_keys))
+        with atomic_uel_writer(output_path, metadata) as writer, closing(
+            stream_work_units({0: handler}, work_units, pin_memory=False)
+        ) as streamed:
+            for key, loaded in tqdm(streamed, total=len(work_units), desc="Pruning keys", unit="keys"):
+                if key in raw_preserve_keys:
+                    write_preserved_tensor(
+                        writer, key, handler, output_key=key, force_raw=True
+                    )
+                else:
+                    tensor = loaded[(0, key)]
+                    try:
+                        write_preserved_tensor(
+                            writer, key, handler, output_key=key, tensor=tensor
+                        )
+                    finally:
+                        del tensor
+                pbar.update(1)
+            pbar.update(len(all_keys) - len(work_units))
+    finally:
+        handler.__exit__(None, None, None)
+        cleanup_after_operation()
 
-    # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
-    model_dir = folder_paths.get_folder_paths(model_type)[-1]
-    output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
-
-    save_file(pruned_tensors, output_path, metadata)
-
-    # Cleanup after operation
-    del pruned_tensors
-    cleanup_after_operation()
-
-    return output_path
+    return output_name
 
 
 class ModelPruneKeys(io.ComfyNode):
@@ -82,21 +123,35 @@ class ModelPruneKeys(io.ComfyNode):
                 io.Combo.Input(
                     "diffusionmodel_name",
                     options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Diffusion model from which matching tensor keys will be removed.",
                 ),
-                io.String.Input("keys_to_prune", multiline=True, default=""),
-                io.Boolean.Input("use_regex", default=False),
-                io.String.Input("output_filename", default="pruned_model"),
+                io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
+                io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
+                io.String.Input("output_filename", default="pruned_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
+                io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
     @classmethod
-    def execute(cls, diffusionmodel_name: str, keys_to_prune: str,
-                use_regex: bool, output_filename: str) -> io.NodeOutput:
-        path = _prune_keys(diffusionmodel_name, "diffusion_models",
-                          keys_to_prune, use_regex, output_filename)
+    def execute(
+        cls,
+        diffusionmodel_name: str,
+        keys_to_prune: str,
+        use_regex: bool,
+        output_filename: str,
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        path = _prune_keys(
+            diffusionmodel_name,
+            "diffusion_models",
+            keys_to_prune,
+            use_regex,
+            output_filename,
+            include_mode=include_mode,
+        )
         return io.NodeOutput(path)
 
 
@@ -112,21 +167,35 @@ class TextEncoderPruneKeys(io.ComfyNode):
                 io.Combo.Input(
                     "textencoder_name",
                     options=folder_paths.get_filename_list("text_encoders"),
+                    tooltip="Text encoder from which matching tensor keys will be removed.",
                 ),
-                io.String.Input("keys_to_prune", multiline=True, default=""),
-                io.Boolean.Input("use_regex", default=False),
-                io.String.Input("output_filename", default="pruned_textencoder"),
+                io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
+                io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
+                io.String.Input("output_filename", default="pruned_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
+                io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
     @classmethod
-    def execute(cls, textencoder_name: str, keys_to_prune: str,
-                use_regex: bool, output_filename: str) -> io.NodeOutput:
-        path = _prune_keys(textencoder_name, "text_encoders",
-                          keys_to_prune, use_regex, output_filename)
+    def execute(
+        cls,
+        textencoder_name: str,
+        keys_to_prune: str,
+        use_regex: bool,
+        output_filename: str,
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        path = _prune_keys(
+            textencoder_name,
+            "text_encoders",
+            keys_to_prune,
+            use_regex,
+            output_filename,
+            include_mode=include_mode,
+        )
         return io.NodeOutput(path)
 
 
@@ -142,21 +211,35 @@ class LoRAPruneKeys(io.ComfyNode):
                 io.Combo.Input(
                     "lora_name",
                     options=folder_paths.get_filename_list("loras"),
+                    tooltip="LoRA from which matching tensor keys will be removed.",
                 ),
-                io.String.Input("keys_to_prune", multiline=True, default=""),
-                io.Boolean.Input("use_regex", default=False),
-                io.String.Input("output_filename", default="pruned_lora"),
+                io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
+                io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
+                io.String.Input("output_filename", default="pruned_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
+                io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
     @classmethod
-    def execute(cls, lora_name: str, keys_to_prune: str,
-                use_regex: bool, output_filename: str) -> io.NodeOutput:
-        path = _prune_keys(lora_name, "loras",
-                          keys_to_prune, use_regex, output_filename)
+    def execute(
+        cls,
+        lora_name: str,
+        keys_to_prune: str,
+        use_regex: bool,
+        output_filename: str,
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        path = _prune_keys(
+            lora_name,
+            "loras",
+            keys_to_prune,
+            use_regex,
+            output_filename,
+            include_mode=include_mode,
+        )
         return io.NodeOutput(path)
 
 
@@ -172,21 +255,35 @@ class CheckpointPruneKeys(io.ComfyNode):
                 io.Combo.Input(
                     "ckpt_name",
                     options=folder_paths.get_filename_list("checkpoints"),
+                    tooltip="Checkpoint from which matching tensor keys will be removed.",
                 ),
-                io.String.Input("keys_to_prune", multiline=True, default=""),
-                io.Boolean.Input("use_regex", default=False),
-                io.String.Input("output_filename", default="pruned_checkpoint"),
+                io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
+                io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
+                io.String.Input("output_filename", default="pruned_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
+                io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
     @classmethod
-    def execute(cls, ckpt_name: str, keys_to_prune: str,
-                use_regex: bool, output_filename: str) -> io.NodeOutput:
-        path = _prune_keys(ckpt_name, "checkpoints",
-                          keys_to_prune, use_regex, output_filename)
+    def execute(
+        cls,
+        ckpt_name: str,
+        keys_to_prune: str,
+        use_regex: bool,
+        output_filename: str,
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        path = _prune_keys(
+            ckpt_name,
+            "checkpoints",
+            keys_to_prune,
+            use_regex,
+            output_filename,
+            include_mode=include_mode,
+        )
         return io.NodeOutput(path)
 
 
@@ -202,19 +299,33 @@ class EmbeddingPruneKeys(io.ComfyNode):
                 io.Combo.Input(
                     "embedding",
                     options=folder_paths.get_filename_list("embeddings"),
+                    tooltip="Embedding from which matching tensor keys will be removed.",
                 ),
-                io.String.Input("keys_to_prune", multiline=True, default=""),
-                io.Boolean.Input("use_regex", default=False),
-                io.String.Input("output_filename", default="pruned_embedding"),
+                io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
+                io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
+                io.String.Input("output_filename", default="pruned_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
+                io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
     @classmethod
-    def execute(cls, embedding: str, keys_to_prune: str,
-                use_regex: bool, output_filename: str) -> io.NodeOutput:
-        path = _prune_keys(embedding, "embeddings",
-                          keys_to_prune, use_regex, output_filename)
-        return io.NodeOutput(path)
+    def execute(
+        cls,
+        embedding: str,
+        keys_to_prune: str,
+        use_regex: bool,
+        output_filename: str,
+        include_mode: bool = False,
+    ) -> io.NodeOutput:
+        path = _prune_keys(
+            embedding,
+            "embeddings",
+            keys_to_prune,
+            use_regex,
+            output_filename,
+            include_mode=include_mode,
+        )
+        return io.NodeOutput(path)

nodes/quantization_guard.py

@@ -0,0 +1,508 @@
+"""Fail-safe detection for quantized or isolated low-bit safetensors inputs."""
+
+import logging
+import json
+import struct
+
+import torch
+
+
+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",
+}
+
+
+_QUANT_SIDECARS = (
+    "comfy_quant", "weight_scale", "weight_scale_2", "input_scale",
+    "pre_quant_scale", "weight_s_rel", "weight_s_channel",
+    "weight_correction", "weight_codebook", "scale_weight", "scale_input",
+)
+
+
+def diffusion_key_map(keys, label):
+    """Map optional diffusion wrappers to each file's physical tensor keys."""
+    mapped = {}
+    for key in sorted(keys):
+        logical = key
+        for prefix in ("model.diffusion_model.", "diffusion_model."):
+            if key.startswith(prefix):
+                logical = key[len(prefix):]
+                break
+        previous = mapped.get(logical)
+        if previous is not None:
+            raise ValueError(
+                f"{label} has duplicate diffusion layer {logical!r}: "
+                f"{previous!r} and {key!r}"
+            )
+        mapped[logical] = key
+    return mapped
+
+
+class DiffusionQuantization:
+    """Index and decode Core quantized diffusion weights one layer at a time.
+
+    The handler is a UEL low-memory handler. Header inspection uses only its
+    keys, shapes, dtypes, and metadata; tensor values are supplied to ``decode``
+    by the caller's bounded UEL work unit.
+    """
+
+    def __init__(self, handler, model_path: str, label: str):
+        self.handler = handler
+        self.model_path = model_path
+        self.label = label
+        self._keys = tuple(handler.keys())
+        self._key_set = set(self._keys)
+        self._metadata = dict(handler.metadata() or {})
+        self._layer_configs = {}
+        self._architecture_shapes = None
+        self._legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in self._keys)
+
+        encoded = self._metadata.get("_quantization_metadata")
+        if encoded:
+            try:
+                layers = json.loads(encoded).get("layers", {})
+            except (TypeError, ValueError) as exc:
+                raise ValueError(
+                    f"[{label}] Invalid _quantization_metadata in {model_path}: {exc}"
+                ) from exc
+            if not isinstance(layers, dict):
+                raise ValueError(f"[{label}] Invalid quantization layers metadata in {model_path}")
+            self._layer_configs.update(layers)
+
+        quantized = set()
+        for sidecar in self._keys:
+            if not sidecar.endswith(".comfy_quant"):
+                continue
+            stem = sidecar[:-len(".comfy_quant")]
+            weight_key = f"{stem}.weight"
+            if weight_key not in self._key_set:
+                raise ValueError(f"[{label}] Quantization sidecar has no weight: {sidecar}")
+            quantized.add(weight_key)
+
+        for stem in self._layer_configs:
+            weight_key = f"{stem}.weight"
+            if weight_key in self._key_set:
+                quantized.add(weight_key)
+
+        if self._legacy_scaled_fp8:
+            for sidecar in self._keys:
+                if sidecar.endswith(".scale_weight"):
+                    stem = sidecar[:-len(".scale_weight")]
+                    weight_key = f"{stem}.weight"
+                    if weight_key in self._key_set:
+                        quantized.add(weight_key)
+                        self._layer_configs.setdefault(stem, {"format": "float8_e4m3fn"})
+
+        self.quantized_keys = frozenset(quantized)
+        sidecars = set()
+        for key in self.quantized_keys:
+            stem = key[:-len(".weight")]
+            for name in _QUANT_SIDECARS:
+                candidate = f"{stem}.{name}"
+                if candidate in self._key_set:
+                    sidecars.add(candidate)
+        for key in self._keys:
+            if key.endswith(".comfy_quant") and any(
+                key == f"{weight[:-len('.weight')]}.comfy_quant" for weight in self.quantized_keys
+            ):
+                sidecars.add(key)
+        if self._legacy_scaled_fp8:
+            sidecars.update(key for key in self._keys if key.endswith("scaled_fp8"))
+        self._sidecars = frozenset(sidecars)
+        self.data_keys = tuple(key for key in self._keys if key not in self._sidecars)
+
+        self.isolated_low_bit_keys = {
+            key for key in self.data_keys
+            if key not in self.quantized_keys and _is_low_bit(handler, key)
+        }
+        if self.isolated_low_bit_keys:
+            logging.warning(
+                "[%s] %s contains %d isolated low-bit tensor(s); preserving/excluding unchanged: %s",
+                label, model_path, len(self.isolated_low_bit_keys),
+                ", ".join(sorted(self.isolated_low_bit_keys)),
+            )
+
+    def required_keys(self, key: str) -> list[str]:
+        if key not in self.quantized_keys:
+            return [key]
+        stem = key[:-len(".weight")]
+        return [key] + [
+            f"{stem}.{name}" for name in _QUANT_SIDECARS
+            if f"{stem}.{name}" in self._key_set
+        ]
+
+    def _architecture_shape(self, key: str) -> tuple[int, ...]:
+        if self._architecture_shapes is None:
+            from comfy import model_detection
+
+            # Core detection examines keys, shapes and dtypes. Meta tensors supply
+            # that inventory without materializing any checkpoint tensor data.
+            state = {
+                name: torch.empty(
+                    tuple(self.handler.get_shape(name)),
+                    dtype=self.handler.get_dtype(name), device="meta",
+                )
+                for name in self._keys
+            }
+            prefix = model_detection.unet_prefix_from_state_dict(state)
+            if self.quantized_keys and all(
+                name.startswith("model.diffusion_model.") for name in self.quantized_keys
+            ):
+                prefix = "model.diffusion_model."
+            elif not any(name.startswith(prefix) for name in state):
+                prefix = ""
+            detected = {
+                name[len(prefix):]: value
+                for name, value in state.items() if name.startswith(prefix)
+            }
+            try:
+                config = model_detection.model_config_from_unet(
+                    detected, "", metadata=self._metadata,
+                )
+                if config is None:
+                    raise ValueError("Core could not detect this diffusion architecture")
+                with torch.device("meta"):
+                    model = config.get_model(detected, "", device=torch.device("meta"))
+                modules = dict(getattr(model, "diffusion_model", model).named_modules())
+                shapes = {}
+                for weight_key in self.quantized_keys:
+                    module_name = weight_key[len(prefix):].removesuffix(".weight")
+                    module = modules.get(module_name)
+                    if module is None:
+                        continue
+                    shape = getattr(module, "_orig_shape", None)
+                    if shape is None:
+                        weight = getattr(module, "weight", None)
+                        shape = getattr(weight, "shape", None)
+                    if shape is not None:
+                        shapes[weight_key] = tuple(shape)
+                self._architecture_shapes = shapes
+            except Exception as exc:
+                raise ValueError(
+                    f"[{self.label}] Cannot determine exact quantized layer shapes from "
+                    f"Core architecture for {self.model_path}: {exc}"
+                ) from exc
+            finally:
+                state.clear()
+                detected.clear()
+                if "model" in locals():
+                    del model
+
+        shape = self._architecture_shapes.get(key)
+        if shape is None:
+            raise ValueError(
+                f"[{self.label}] Core architecture has no logical shape for quantized layer {key}"
+            )
+        stored = tuple(self.handler.get_shape(key))
+        if len(shape) != 2 or len(stored) != 2:
+            raise ValueError(f"[{self.label}] Invalid padded quantized shape for {key}: {shape}, {stored}")
+        stem = key[:-len(".weight")]
+        config = self._layer_configs.get(stem)
+        nvfp4 = (isinstance(config, dict) and config.get("format") == "nvfp4") or (
+            f"{stem}.weight_scale_2" in self._key_set
+        )
+        if nvfp4:
+            expected = ((shape[0] + 15) // 16 * 16, (shape[1] + 15) // 16 * 8)
+        else:
+            expected = ((shape[0] + 31) // 32 * 32, (shape[1] + 31) // 32 * 32)
+        if stored != expected:
+            raise ValueError(
+                f"[{self.label}] Core architecture shape {shape} disagrees with packed storage {stored} for {key}"
+            )
+        return shape
+
+    def _format(self, key: str, tensors=None):
+        stem = key[:-len(".weight")]
+        config = self._layer_configs.get(stem)
+        sidecar = f"{stem}.comfy_quant"
+        if tensors is None and config is None and sidecar in self._key_set:
+            stream = self.handler.async_stream(
+                [sidecar], batch_size=1, prefetch_batches=1, pin_memory=False,
+            )
+            yielded_keys = []
+            try:
+                batch = next(stream)
+                yielded_keys = [name for name, _ in batch]
+                if len(batch) != 1 or batch[0][0] != sidecar:
+                    raise RuntimeError(f"[{self.label}] Unexpected quantization sidecar for {key}")
+                tensors = {sidecar: batch[0][1]}
+            finally:
+                try:
+                    for yielded_key in yielded_keys:
+                        self.handler.mark_processed(yielded_key)
+                finally:
+                    stream.close()
+        if tensors is not None and sidecar in tensors:
+            raw = tensors[sidecar].detach().to(device="cpu").contiguous().reshape(-1).tolist()
+            try:
+                config = json.loads(bytes(raw).decode("utf-8"))
+            except (UnicodeDecodeError, TypeError, ValueError) as exc:
+                raise ValueError(f"[{self.label}] Invalid quantization sidecar {sidecar}: {exc}") from exc
+            self._layer_configs[stem] = config
+        if config is None:
+            raise ValueError(f"[{self.label}] Missing quantization config for {key}")
+        if not isinstance(config, dict):
+            raise ValueError(f"[{self.label}] Invalid quantization config for {key}: expected object")
+        fmt = config.get("format")
+        try:
+            from comfy.quant_ops import QUANT_ALGOS
+        except ImportError as exc:
+            raise RuntimeError("ComfyUI quantization support is unavailable") from exc
+        if fmt not in QUANT_ALGOS:
+            raise ValueError(f"[{self.label}] Unsupported ComfyUI quantization format {fmt!r} for {key}")
+        return fmt, config
+
+    def logical_shape(self, key: str) -> tuple[int, ...]:
+        if key not in self.quantized_keys:
+            return tuple(self.handler.get_shape(key))
+        fmt, _ = self._format(key)
+        if fmt == "nvfp4" or fmt == "mxfp8":
+            return self._architecture_shape(key)
+        shape = tuple(self.handler.get_shape(key))
+        if fmt in {"convrot_w4a4", "asym_w4a8_int8"}:
+            if len(shape) != 2:
+                raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}")
+            return shape[0], shape[1] * 2
+        if fmt == "w6a8_int8":
+            if len(shape) != 2:
+                raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}")
+            if (shape[1] * 8) % 6 != 0:
+                raise ValueError(f"[{self.label}] Invalid 6-bit packed width {shape[1]} for {key}")
+            return shape[0], shape[1] * 8 // 6
+        return shape
+
+    def logical_dtype(self, key: str, compute_dtype: torch.dtype) -> torch.dtype:
+        if key not in self.quantized_keys:
+            return self.handler.get_dtype(key)
+        return compute_dtype
+
+    def decode(self, key: str, tensors, compute_dtype: torch.dtype, device="cpu") -> torch.Tensor:
+        if key not in self.quantized_keys:
+            raise ValueError(f"[{self.label}] {key} is not a recognized quantized diffusion weight")
+        fmt, config = self._format(key, tensors)
+        try:
+            from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+        except ImportError as exc:
+            raise RuntimeError("ComfyUI quantization support is unavailable") from exc
+
+        stem = key[:-len(".weight")]
+        device = torch.device(device)
+        shape = self.logical_shape(key)
+
+        def tensor(name, *, dtype=None, required=False):
+            tensor_key = f"{stem}.{name}"
+            if tensor_key not in tensors and name == "weight_scale":
+                legacy_key = f"{stem}.scale_weight"
+                if legacy_key in tensors:
+                    tensor_key = legacy_key
+            value = tensors.get(tensor_key)
+            if value is None:
+                if required:
+                    raise ValueError(f"[{self.label}] Missing quantization tensor {tensor_key} for {key}")
+                return None
+            value = value.to(device=device)
+            if dtype is not None:
+                value = value.view(dtype=dtype)
+            return value
+
+        params_conf = config.get("params", {})
+        if not isinstance(params_conf, dict):
+            params_conf = {}
+        if fmt in ("float8_e4m3fn", "float8_e5m2"):
+            values = {"scale": tensor("weight_scale", required=True)}
+        elif fmt == "mxfp8":
+            values = {"scale": tensor("weight_scale", dtype=torch.float8_e8m0fnu, required=True)}
+        elif fmt == "nvfp4":
+            values = {
+                "scale": tensor("weight_scale_2", required=True),
+                "block_scale": tensor("weight_scale", dtype=torch.float8_e4m3fn, required=True),
+            }
+        elif fmt == "int8_tensorwise":
+            values = {"scale": tensor("weight_scale", required=True)}
+            if config.get("convrot", params_conf.get("convrot", False)):
+                values.update({
+                    "convrot": True,
+                    "convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))),
+                })
+        elif fmt == "convrot_w4a4":
+            values = {
+                "scale": tensor("weight_scale", required=True),
+                "convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))),
+                "quant_group_size": 64,
+                "linear_dtype": config.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
+            }
+        elif fmt in ("asym_w4a8_int8", "w6a8_int8"):
+            if f"{stem}.weight_correction" in tensors:
+                raise ValueError(
+                    f"[{self.label}] Core's W4A8 loader does not apply weight_correction for {key}"
+                )
+            scale_value = tensors.get(f"{stem}.weight_s_rel")
+            values = {
+                "scale": tensor("weight_s_rel", dtype=torch.float8_e4m3fn if scale_value is not None and scale_value.dtype == torch.uint8 else None, required=True),
+                "s_channel": tensor("weight_s_channel", required=True),
+                "codebook": tensor("weight_codebook"),
+                "group_size": int(config.get("group_size", params_conf.get("group_size", 16))),
+                "convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))),
+            }
+        else:
+            raise ValueError(f"[{self.label}] Unsupported ComfyUI quantization format {fmt!r} for {key}")
+
+        layout_type = QUANT_ALGOS[fmt]["comfy_tensor_layout"]
+        layout = get_layout_class(layout_type)
+        params = layout.Params(**values, orig_dtype=compute_dtype, orig_shape=shape)
+        qdata = tensors[key].to(device=device, dtype=QUANT_ALGOS[fmt]["storage_t"])
+        return QuantizedTensor(qdata, layout_type, params).dequantize()
+
+    def output_metadata(self) -> dict:
+        metadata = dict(self._metadata)
+        metadata.pop("_quantization_metadata", None)
+        return metadata
+
+
+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 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", "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,
+    *,
+    force_raw: bool = False,
+    tensor=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 not force_raw and entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES:
+        if tensor is not None:
+            writer.write(destination, tensor.cpu().contiguous())
+            return
+        stream = handler.async_stream(
+            [key], batch_size=1, prefetch_batches=1, pin_memory=False
+        )
+        yielded = False
+        try:
+            batch = next(stream)
+            if len(batch) != 1 or batch[0][0] != key:
+                raise RuntimeError(f"UEL yielded an unexpected preserved tensor for {key}")
+            yielded = True
+            writer.write(destination, batch[0][1].cpu().contiguous())
+        finally:
+            if yielded:
+                handler.mark_processed(key)
+            stream.close()
+        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))

nodes/renamekeys.py

@@ -3,12 +3,14 @@ import folder_paths
 import comfy.utils
 from tqdm import tqdm
 from comfy_api.latest import io
-from safetensors.torch import save_file
-from safetensors import safe_open
 from .utils import convert_pt_to_safetensors
-from .merger_utils import MemoryEfficientSafeOpen
 from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
 
+from unifiedefficientloader import MemoryEfficientSafeOpen
+from .uel_io import atomic_uel_writer
+from .artifact_paths import canonical_model_artifact_path
+
+
 
 def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
                  new_keys_str: str, output_filename: str) -> str:
@@ -29,10 +31,6 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
     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 {}
-
     old_keys = [key.strip() for key in old_keys_str.strip().split('\n') if key.strip()]
     new_keys = [key.strip() for key in new_keys_str.strip().split('\n') if key.strip()]
 
@@ -41,27 +39,28 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
 
     key_map = dict(zip(old_keys, new_keys))
 
-    # Stream tensors and rename on the fly
-    renamed_tensors = {}
-    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 keys", unit="keys"):
-            new_key = key_map.get(key, key)
-            renamed_tensors[new_key] = handler.get_tensor(key)
-            pbar.update(1)
+    output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
 
-    # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
-    model_dir = folder_paths.get_folder_paths(model_type)[-1]
-    output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
-
-    save_file(renamed_tensors, output_path, metadata)
+    # Stream tensors, rename on the fly, write immediately
+    with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
+        metadata = handler.metadata() or {}
+        with atomic_uel_writer(output_path, metadata) as writer:
+            original_keys = handler.keys()
+            pbar = comfy.utils.ProgressBar(len(original_keys))
+            stream = handler.async_stream(
+                original_keys, batch_size=1, prefetch_batches=1, pin_memory=False
+            )
+            for batch in tqdm(stream, total=len(original_keys), desc="Renaming keys", unit="keys"):
+                key, tensor = batch[0]
+                new_key = key_map.get(key, key)
+                writer.write_batch([(new_key, tensor.contiguous())])
+                handler.mark_processed(key)
+                pbar.update(1)
 
     # Cleanup after operation
-    del renamed_tensors
     cleanup_after_operation()
 
-    return output_path
+    return output_name
 
 
 class ModelRenameKeys(io.ComfyNode):
@@ -76,13 +75,14 @@ class ModelRenameKeys(io.ComfyNode):
                 io.Combo.Input(
                     "diffusionmodel_name",
                     options=folder_paths.get_filename_list("diffusion_models"),
+                    tooltip="Diffusion model whose tensor keys will be renamed.",
                 ),
-                io.String.Input("old_keys", multiline=True, default=""),
-                io.String.Input("new_keys", multiline=True, default=""),
-                io.String.Input("output_filename", default="renamed_model"),
+                io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
+                io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
+                io.String.Input("output_filename", default="renamed_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
@@ -106,13 +106,14 @@ class TextEncoderRenameKeys(io.ComfyNode):
                 io.Combo.Input(
                     "textencoder_name",
                     options=folder_paths.get_filename_list("text_encoders"),
+                    tooltip="Text encoder whose tensor keys will be renamed.",
                 ),
-                io.String.Input("old_keys", multiline=True, default=""),
-                io.String.Input("new_keys", multiline=True, default=""),
-                io.String.Input("output_filename", default="renamed_textencoder"),
+                io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
+                io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
+                io.String.Input("output_filename", default="renamed_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
@@ -136,13 +137,14 @@ class LoRARenameKeys(io.ComfyNode):
                 io.Combo.Input(
                     "lora_name",
                     options=folder_paths.get_filename_list("loras"),
+                    tooltip="LoRA whose tensor keys will be renamed.",
                 ),
-                io.String.Input("old_keys", multiline=True, default=""),
-                io.String.Input("new_keys", multiline=True, default=""),
-                io.String.Input("output_filename", default="renamed_lora"),
+                io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
+                io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
+                io.String.Input("output_filename", default="renamed_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
@@ -166,13 +168,14 @@ class CheckpointRenameKeys(io.ComfyNode):
                 io.Combo.Input(
                     "ckpt_name",
                     options=folder_paths.get_filename_list("checkpoints"),
+                    tooltip="Checkpoint whose tensor keys will be renamed.",
                 ),
-                io.String.Input("old_keys", multiline=True, default=""),
-                io.String.Input("new_keys", multiline=True, default=""),
-                io.String.Input("output_filename", default="renamed_checkpoint"),
+                io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
+                io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
+                io.String.Input("output_filename", default="renamed_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
@@ -196,13 +199,14 @@ class EmbeddingRenameKeys(io.ComfyNode):
                 io.Combo.Input(
                     "embedding",
                     options=folder_paths.get_filename_list("embeddings"),
+                    tooltip="Embedding whose tensor keys will be renamed.",
                 ),
-                io.String.Input("old_keys", multiline=True, default=""),
-                io.String.Input("new_keys", multiline=True, default=""),
-                io.String.Input("output_filename", default="renamed_embedding"),
+                io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
+                io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
+                io.String.Input("output_filename", default="renamed_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
             ],
             outputs=[
-                io.String.Output(display_name="output_path"),
+                io.AnyType.Output(display_name="output_path"),
             ],
         )
 
@@ -211,4 +215,4 @@ class EmbeddingRenameKeys(io.ComfyNode):
                 new_keys: str, output_filename: str) -> io.NodeOutput:
         path = _rename_keys(embedding, "embeddings",
                            old_keys, new_keys, output_filename)
-        return io.NodeOutput(path)
+        return io.NodeOutput(path)

nodes/text_encoder_extract.py

@@ -0,0 +1,794 @@
+"""
+Text Encoder LoRA and DoRA Extraction.
+
+Specifically targets extracting text encoder deltas as LoRA/DoRA with custom mapping rules
+for layer name prefixes and PEFT weight formatting.
+"""
+import torch
+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
+from .uel_io import atomic_uel_writer
+from .artifact_paths import canonical_model_artifact_path
+from .extraction_stream import dora_difference, lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
+from .layer_parameters import parameter_input, resolve_layer_parameters
+
+# Import SVD rank utilities from standard lora extract file to avoid duplication
+from .lora_extract_svd import (
+    _svd_extract_linear,
+    _svd_extract_conv,
+    _extract_chunked_layer,
+    _detect_fused_layer,
+    _compile_patterns,
+    _matches_any_pattern
+)
+
+
+def _format_te_lora_key(key: str) -> str:
+    """
+    Format key specifically for Text Encoder LoRA/DoRA saving.
+    Corrects layer name prefixing based on specific target mappings.
+    """
+    if key.endswith(".weight"):
+        key = key[:-7]
+
+    rules = [
+        ("model.language_model.layers", "text_encoders.transformer.model.layers"),
+        ("model.visual.blocks", "text_encoders.transformer.visual.blocks"),
+        ("model.visual.deepstack_merger_list", "text_encoders.transformer.visual.deepstack_merger_list"),
+        ("model.visual.merger", "text_encoders.transformer.visual.merger"),
+        ("model.visual.pos_embed", "text_encoders.transformer.visual.pos_embed"),
+        ("model.layers", "text_encoders.transformer.model.layers"),
+        ("vision_model.embeddings.position_embedding", "text_encoders.transformer.vision_model.embeddings.position_embedding"),
+        ("vision_model.encoder.layers", "text_encoders.transformer.vision_model.encoder.layers"),
+    ]
+
+    for prefix, target in rules:
+        if key.startswith(prefix):
+            return target + key[len(prefix):]
+
+    # Prepend text_encoders. as default generic fallback
+    if key.startswith("text_encoders."):
+        return key
+    return f"text_encoders.{key}"
+
+
+def extract_te_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,
+    is_dora: bool = False,
+    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 = False,
+    glob_skip_patterns: bool = False,
+    knee_probe_offset: int = 32,
+    include_mode: bool = False,
+    layer_parameters=None,
+) -> None:
+    """
+    Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk.
+    """
+    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"[TE 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 = sorted(k for k in keys_a if k.endswith(".weight"))
+        if mode == "fixed":
+            profile = "extract:fixed"
+            defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff}
+        elif mode == "ratio":
+            profile = "extract:ratio"
+            defaults = {"linear_ratio": linear_param, "conv_ratio": conv_param,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
+        elif mode == "quantile":
+            profile = "extract:quantile"
+            defaults = {"linear_quantile": linear_param, "conv_quantile": conv_param,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
+        elif mode == "sv_fro":
+            profile = "extract:frobenius"
+            defaults = {"linear_target": linear_param, "conv_target": conv_param,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff,
+                        "linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
+        else:
+            profile = "extract:knee"
+            defaults = {"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
+                        "clamp_quantile": clamp_quantile, "min_diff": min_diff}
+        resolved_parameters = resolve_layer_parameters(
+            layer_parameters, profile, weight_keys, defaults, node_name="Text Encoder Extract"
+        )
+        work_units = [
+            (key, key, key if key in keys_b else None)
+            for key in weight_keys
+            if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
+            and not (key not in keys_b and mismatch_mode == "skip")
+        ]
+        pbar = comfy.utils.ProgressBar(len(work_units))
+        stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
+
+        def _process_layer(key, cpu_a, cpu_b):
+            values = resolved_parameters.get(key, defaults)
+            if mode == "fixed":
+                layer_linear_param, layer_conv_param = values["linear_dim"], values["conv_dim"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            elif mode == "ratio":
+                layer_linear_param, layer_conv_param = values["linear_ratio"], values["conv_ratio"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            elif mode == "quantile":
+                layer_linear_param, layer_conv_param = values["linear_quantile"], values["conv_quantile"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            elif mode == "sv_fro":
+                layer_linear_param, layer_conv_param = values["linear_target"], values["conv_target"]
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            else:
+                layer_linear_param, layer_conv_param = linear_param, conv_param
+                layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
+            layer_clamp_quantile = values["clamp_quantile"]
+            layer_min_diff = values["min_diff"]
+            lora_name = _format_te_lora_key(key)
+            weight_diff = None
+            layer_device = device
+
+            if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)):
+                return "skipped", None
+
+            # Load tensors with pinned memory for 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")
+
+                if is_dora:
+                    weight_diff, dora_scale, layer_device = dora_difference(
+                        cpu_a, None, device
+                    )
+                else:
+                    weight_diff, layer_device = lora_difference(cpu_a, None, device)
+                    dora_scale = None
+            else:
+                if cpu_a.shape != cpu_b.shape:
+                    if mismatch_mode == "skip":
+                        return "skipped", None
+                    if mismatch_mode == "error":
+                        raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
+
+                    if is_dora:
+                        weight_diff, dora_scale, layer_device = dora_difference(
+                            cpu_a, None, device
+                        )
+                    else:
+                        weight_diff, layer_device = lora_difference(cpu_a, None, device)
+                        dora_scale = None
+                else:
+                    if cpu_a.ndim < 2:
+                        return "skipped", None
+                    if is_dora:
+                        weight_diff, dora_scale, layer_device = dora_difference(
+                            cpu_a, cpu_b, device
+                        )
+                    else:
+                        weight_diff, layer_device = lora_difference(
+                            cpu_a, cpu_b, device
+                        )
+                        dora_scale = None
+
+            if weight_diff is None:
+                raise RuntimeError(f"No extraction tensor was produced for {key}")
+
+            # Skip small differences
+            if layer_min_diff > 0 and weight_diff.abs().max() < layer_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 = {}
+            svd_weight = weight_diff
+
+            try:
+                if is_conv:
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_conv(svd_weight, mode, layer_conv_param, layer_device, layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
+                        lambda: _svd_extract_conv(svd_weight.cpu(), mode, layer_conv_param, "cpu", layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
+                        layer_device,
+                    )
+                else:
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_linear(svd_weight, mode, layer_linear_param, layer_device, layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
+                        lambda: _svd_extract_linear(svd_weight.cpu(), mode, layer_linear_param, "cpu", layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
+                        layer_device,
+                    )
+            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"[TE Extract] Chunked: {key} ({num_chunks} chunks)")
+                        lora_up, lora_down, rank = _extract_chunked_layer(
+                            weight_diff, num_chunks, mode, layer_linear_param, layer_device,
+                            layer_linear_max_rank, knee_probe_offset, layer_clamp_quantile,
+                        )
+                        if lora_up is not None:
+                            # PEFT suffixes
+                            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()
+                            if is_dora and 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"[TE 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, _ = result
+                # Write to PEFT weight suffixes
+                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()
+                if is_dora and 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_context = atomic_uel_writer(output_path)
+        writer = writer_context.__enter__()
+        try:
+            stream = paired_async_tensors(
+                handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
+            )
+            for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting TE Layers", unit="layers"):
+                status, layer_sd = _process_layer(key, cpu_a, cpu_b)
+                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)
+            stats["skipped"] += len(weight_keys) - len(work_units)
+        except BaseException as exc:
+            writer_context.__exit__(type(exc), exc, exc.__traceback__)
+            raise
+        else:
+            writer_context.__exit__(None, None, None)
+
+        print(f"[TE 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) -> tuple[str, str]:
+    """Build output path for LoRA file."""
+    return canonical_model_artifact_path("loras", output_filename)
+
+
+def _get_te_model_inputs():
+    return [
+        io.Combo.Input("model_a", options=folder_paths.get_filename_list("text_encoders"),
+                      tooltip="Finetuned Text Encoder model (A - B = LoRA)"),
+        io.Combo.Input("model_b", options=folder_paths.get_filename_list("text_encoders"),
+                      tooltip="Base Text Encoder model (A - B = LoRA)"),
+    ]
+
+
+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=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."),
+        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", tooltip="Handle missing or incompatible text-encoder tensors by skipping them, substituting zeros where supported, or aborting."),
+        io.String.Input("output_filename", default="extracted_te_lora", tooltip="Output filename without extension, written under the ComfyUI LoRA directory."),
+        io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save extracted text-encoder LoRA factors."),
+        io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer extraction arithmetic; CUDA out-of-memory retries the affected layer on CPU where supported."),
+        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_mode", default=False, tooltip="Use Skip Patterns as a whitelist instead. Only matching layers are extracted; an empty whitelist extracts nothing."),
+    ]
+
+
+# =============================================================================
+# LoRA Text Encoder Extract Nodes
+# =============================================================================
+
+class TextEncoderLoRAExtractFixed(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderLoRAExtractFixed",
+            display_name="TE LoRA Extract (Fixed Rank)",
+            category="ModelUtils/LoRA Extract (TE)",
+            description="Extract Text Encoder LoRA with specified fixed rank for each layer type.",
+            inputs=[
+                *_get_te_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"),
+                *_get_common_inputs(),
+                parameter_input("extract:fixed"),
+            ],
+            outputs=[io.AnyType.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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
+            device, save_dtype, output_path, is_dora=False,
+            linear_max_rank=linear_dim, conv_max_rank=conv_dim,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderLoRAExtractRatio(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderLoRAExtractRatio",
+            display_name="TE LoRA Extract (Ratio)",
+            category="ModelUtils/LoRA Extract (TE)",
+            description="Keep singular values > max(S) / ratio.",
+            inputs=[
+                *_get_te_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 values retain more singular values."),
+                io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:ratio"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, 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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
+            device, save_dtype, output_path, is_dora=False,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderLoRAExtractQuantile(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderLoRAExtractQuantile",
+            display_name="TE LoRA Extract (Quantile)",
+            category="ModelUtils/LoRA Extract (TE)",
+            description="Keep enough singular values to reach target percentage.",
+            inputs=[
+                *_get_te_model_inputs(),
+                io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
+                io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:quantile"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, 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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
+            device, save_dtype, output_path, is_dora=False,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderLoRAExtractKnee(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderLoRAExtractKnee",
+            display_name="TE LoRA Extract (Knee)",
+            category="ModelUtils/LoRA Extract (TE)",
+            description="Automatically find optimal rank using knee detection on singular value curve.",
+            inputs=[
+                *_get_te_model_inputs(),
+                io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
+                io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:knee"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
+                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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, knee_method, 0, 0,
+            device, save_dtype, output_path, is_dora=False,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache,
+            glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderLoRAExtractFrobenius",
+            display_name="TE LoRA Extract (Frobenius)",
+            category="ModelUtils/LoRA Extract (TE)",
+            description="Preserve target fraction of Frobenius norm.",
+            inputs=[
+                *_get_te_model_inputs(),
+                io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
+                io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:frobenius"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, 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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
+            device, save_dtype, output_path, is_dora=False,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+# =============================================================================
+# DoRA Text Encoder Extract Nodes
+# =============================================================================
+
+class TextEncoderDoRAExtractFixed(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderDoRAExtractFixed",
+            display_name="TE DoRA Extract (Fixed Rank)",
+            category="ModelUtils/DoRA Extract (TE)",
+            description="Extract Text Encoder DoRA with specified fixed rank for each layer type.",
+            inputs=[
+                *_get_te_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"),
+                *_get_common_inputs(),
+                parameter_input("extract:fixed"),
+            ],
+            outputs=[io.AnyType.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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
+            device, save_dtype, output_path, is_dora=True,
+            linear_max_rank=linear_dim, conv_max_rank=conv_dim,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderDoRAExtractRatio(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderDoRAExtractRatio",
+            display_name="TE DoRA Extract (Ratio)",
+            category="ModelUtils/DoRA Extract (TE)",
+            description="Keep singular values > max(S) / ratio.",
+            inputs=[
+                *_get_te_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 values retain more singular values."),
+                io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:ratio"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, 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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
+            device, save_dtype, output_path, is_dora=True,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderDoRAExtractQuantile(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderDoRAExtractQuantile",
+            display_name="TE DoRA Extract (Quantile)",
+            category="ModelUtils/DoRA Extract (TE)",
+            description="Keep enough singular values to reach target percentage.",
+            inputs=[
+                *_get_te_model_inputs(),
+                io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
+                io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:quantile"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, 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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
+            device, save_dtype, output_path, is_dora=True,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderDoRAExtractKnee(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderDoRAExtractKnee",
+            display_name="TE DoRA Extract (Knee)",
+            category="ModelUtils/DoRA Extract (TE)",
+            description="Automatically find optimal rank using knee detection on singular value curve.",
+            inputs=[
+                *_get_te_model_inputs(),
+                io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
+                io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:knee"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
+                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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, knee_method, 0, 0,
+            device, save_dtype, output_path, is_dora=True,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache,
+            glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)
+
+
+class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncoderDoRAExtractFrobenius",
+            display_name="TE DoRA Extract (Frobenius)",
+            category="ModelUtils/DoRA Extract (TE)",
+            description="Preserve target fraction of Frobenius norm.",
+            inputs=[
+                *_get_te_model_inputs(),
+                io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
+                io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
+                io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
+                io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
+                *_get_common_inputs(),
+                parameter_input("extract:frobenius"),
+            ],
+            outputs=[io.AnyType.Output(display_name="output_path")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, 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, include_mode=False, layer_parameters=None) -> io.NodeOutput:
+
+        model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
+        model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
+        output_path, output_name = _build_lora_output_path(output_filename)
+
+        extract_te_from_files(
+            model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
+            device, save_dtype, output_path, is_dora=True,
+            linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
+            clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
+            mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
+            knee_probe_offset=probe_offset,
+            layer_parameters=layer_parameters,
+        )
+        return io.NodeOutput(output_name)

nodes/uel_io.py

@@ -0,0 +1,115 @@
+import os
+import uuid
+from collections import deque
+from contextlib import contextmanager
+
+from unifiedefficientloader import IncrementalSafetensorsWriter
+
+
+class AsyncTensorCursor:
+    """Ordered, bounded tensor access over one UEL asynchronous stream."""
+
+    def __init__(self, handler, keys, *, pin_memory: bool):
+        self.handler = handler
+        self.expected = tuple(keys)
+        self.position = 0
+        self.pending = deque()
+        self.stream = None
+        if self.expected:
+            self.stream = handler.async_stream(
+                list(self.expected),
+                batch_size=1,
+                prefetch_batches=1,
+                pin_memory=pin_memory,
+            )
+
+    def take(self, expected_key):
+        if self.position >= len(self.expected):
+            raise RuntimeError(f"Unexpected tensor request: {expected_key}")
+        planned = self.expected[self.position]
+        if planned != expected_key:
+            raise RuntimeError(
+                f"UEL stream order mismatch: expected {planned}, requested {expected_key}"
+            )
+        if not self.pending:
+            self.pending.extend(next(self.stream))
+        key, tensor = self.pending.popleft()
+        if key != expected_key:
+            raise RuntimeError(f"UEL yielded {key} while {expected_key} was expected")
+        self.position += 1
+        return tensor
+
+    def release(self, key):
+        self.handler.mark_processed(key)
+
+    def finish(self):
+        if self.position != len(self.expected) or self.pending:
+            raise RuntimeError("UEL stream ended before all planned tensors were consumed")
+        if self.stream is not None:
+            try:
+                next(self.stream)
+            except StopIteration:
+                return
+            raise RuntimeError("UEL stream yielded unplanned tensors")
+
+    def close(self):
+        if self.stream is not None:
+            close = getattr(self.stream, "close", None)
+            if close is not None:
+                close()
+            self.stream = None
+        self.pending.clear()
+
+
+@contextmanager
+def atomic_uel_writer(output_path: str, metadata: dict | None = None):
+    """Write incrementally and expose the destination only after finalization."""
+
+    directory = os.path.dirname(output_path) or "."
+    os.makedirs(directory, exist_ok=True)
+    basename = os.path.basename(output_path)
+    temporary = os.path.join(directory, f".{basename}.{uuid.uuid4().hex}.tmp")
+    try:
+        with IncrementalSafetensorsWriter(
+            temporary,
+            metadata=metadata or {},
+            max_workers=1,
+        ) as writer:
+            yield writer
+        os.replace(temporary, output_path)
+    finally:
+        if os.path.exists(temporary):
+            os.remove(temporary)
+
+
+def stream_work_units(handlers, work_units, *, pin_memory: bool):
+    """Yield bounded logical units from one ordered stream per input handler."""
+
+    planned = {source: [] for source in handlers}
+    for _, entries in work_units:
+        for source, keys in entries.items():
+            planned[source].extend(keys)
+    cursors = {
+        source: AsyncTensorCursor(handler, planned[source], pin_memory=pin_memory)
+        for source, handler in handlers.items()
+    }
+    try:
+        for logical_key, entries in work_units:
+            loaded = {}
+            consumed = []
+            try:
+                for source, keys in entries.items():
+                    cursor = cursors[source]
+                    for key in keys:
+                        loaded[(source, key)] = cursor.take(key)
+                        consumed.append((source, key))
+                yield logical_key, loaded
+            finally:
+                loaded.clear()
+                for source, key in consumed:
+                    cursors[source].release(key)
+        for cursor in cursors.values():
+            cursor.finish()
+    finally:
+        for cursor in cursors.values():
+            cursor.close()

nodes/utils.py

@@ -1,7 +1,6 @@
 import torch
-from safetensors import safe_open
-from safetensors.torch import save_file
-from comfy.model_management import get_torch_device
+
+from .uel_io import atomic_uel_writer
 
 def convert_pt_to_safetensors(pt_path, safe_path):
     try:
@@ -12,23 +11,13 @@ def convert_pt_to_safetensors(pt_path, safe_path):
             state_dict = model["model"]
         else:
             state_dict = model
-        state_dict = {k.replace("state_dict.", ""): v for k, v in state_dict.items()}
         metadata = {"format": "pt"}
-        save_file(state_dict, safe_path, metadata)
+        with atomic_uel_writer(safe_path, metadata) as writer:
+            for key in list(state_dict.keys()):
+                tensor = state_dict.pop(key)
+                writer.write_batch(
+                    [(key.replace("state_dict.", ""), tensor.contiguous())]
+                )
         return True, ""
     except Exception as e:
         return False, str(e)
-
-def load_metadata_from_safetensors(safetensors_file: str) -> dict:
-    try:
-        with safe_open(safetensors_file, framework="pt", device="cpu") as f:
-            model_weights = {}
-            for key in f.keys():
-                model_weights[key] = f.get_tensor(key)
-            metadata = f.metadata()
-        if metadata is None:
-            metadata = {}
-        return model_weights, metadata
-    except Exception as e:
-        print(f"Error loading safetensors file: {e}")
-        return {}, {}

plans/per-layer-parameter-rules.md

@@ -0,0 +1,237 @@
+# Per-layer parameter rules for merge, extraction, and resize
+
+## 1. Summary and initial scope
+
+Add one multiline configurator node that supplies optional per-layer numeric overrides to existing nodes. Matching layers receive the explicitly assigned values; unassigned parameters and unmatched layers retain the main node’s settings. Existing exclusions, inclusion filters, discard rules, mismatch handling, and safety guards remain authoritative.
+
+The initial implementation covers these 39 nodes:
+
+| Family | Included variants | Count |
+|---|---|---:|
+| Standard mergers | Two- and three-input checkpoint, diffusion-model, text-encoder, LoRA, and embedding mergers | 10 |
+| Extraction | Fixed, Ratio, Quantile, Knee, and Frobenius for LoRA, DoRA, learned DoRA, text-encoder LoRA, and text-encoder DoRA | 25 |
+| LoRA resize | Fixed, Ratio, Frobenius, and Cumulative | 4 |
+
+Do not expand this implementation to CWB configuration, weighted multi-LoRA merging, LoRA application, Lodestone, optimizer settings, iteration counts, seeds, device selection, dtype selection, or new normalization behavior. These remain post-implementation considerations.
+
+No changes are made during planning. After approval and execution is enabled, the first implementation action is to save this entire approved plan verbatim to `plans/per-layer-parameter-rules.md`. Commit that planning document before editing implementation files. Commit the completed feature after verification.
+
+Preserve the pre-existing analysis/report edits currently outside the filter commit.
+
+## 2. Public interface, syntax, and bindings
+
+### Configurator and connection
+
+Implement `LayerParameterConfiguration`, displayed as **Layer Parameter Configuration**, in a new `nodes/layer_parameters.py` module.
+
+Its inputs are:
+
+| Input | Type | Default | Meaning |
+|---|---|---|---|
+| `rules` | Multiline string | Empty | One pattern-and-assignment rule per line |
+| `glob_patterns` | Boolean | False | False uses regex substring search; True uses glob substring matching |
+
+Its outputs are:
+
+| Output | Type | Meaning |
+|---|---|---|
+| `layer_parameters` | `io.Custom("MODELUTILS_LAYER_PARAMETERS")` | Immutable, parsed rules |
+| `documentation` | String | Syntax, bindings, supported nodes, fallback behavior, errors, and examples |
+
+This follows the existing `io.Custom(...)` configurator/socket mechanism already used by CWB. No frontend extension or dynamic socket-type changes are needed.
+
+Append an optional `layer_parameters` input to each supported receiver. Append `layer_parameters=None` to its execution interface and forward it by keyword. Preserve every existing input identifier, widget order, output, and default.
+
+An unconnected input or an empty rule set leaves the operation unchanged.
+
+### Rule grammar
+
+Accept the following forms:
+
+```text
+(blocks\.4[589]\.attn\.qkv_proj) a:0.5 b:0.25 c:0.75 d:1.0
+(blocks\.4[589]\.attn\.qkv_proj) alpha:0.5, beta:0.25, gamma:0.75, delta:1.0
+(blocks\.4[589]\.attn\.qkv_proj) a:64, b:32, c:0.99, d:0.0
+(blocks\.4[589]\.attn\.qkv_proj) max_rank:128 min_rank:1 target:0.9
+```
+
+Rules:
+
+- Require a parenthesized pattern followed by at least one `name:value` assignment.
+- Accept whitespace, commas, or a mixture between assignments, including whitespace around colons.
+- Preserve the pattern’s backslashes, internal whitespace, nested groups, and character classes.
+- Accept signed decimal and scientific-notation numeric values. Reject nonfinite values.
+- Ignore blank lines and whole-line comments whose first nonwhitespace character is `#`. Do not support inline comments.
+- Reject empty patterns, missing values, trailing garbage, duplicate assignments, and invalid regex.
+- Use `re.search` for regex matching. For glob matching, use case-sensitive `fnmatchcase` with surrounding wildcards to provide substring matching.
+- Keep configurator pattern syntax independent of the main node’s existing exclusion-filter syntax.
+
+Short aliases and full parameter names may be mixed. Two aliases assigning the same parameter on one line are an error, even when their values agree.
+
+### Binding tables
+
+The shared binding definitions must generate both receiver tooltips and configurator documentation, preventing the two from drifting.
+
+#### Standard merge coefficients
+
+Bindings remain stable across calculation modes:
+
+| Short name | Full name |
+|---|---|
+| `a` | `alpha` |
+| `b` | `beta` |
+| `c` or `g` | `gamma` |
+| `d` | `delta` |
+| `e` | `epsilon` |
+| `f` or `z` | `zeta` |
+
+Accept only coefficients actually used by the selected calculation mode:
+
+| Calculation mode | Supported coefficients |
+|---|---|
+| Weight-Sum; Train-Difference | alpha |
+| Add-Difference; Power-Up (DARE) | alpha, beta |
+| Comparative-Interpolation; Add-Dissimilarities; SVD LoRA Extraction; Enhanced Auto Interp; Weight-Sum Cutoff; Power-Up (DARE+TIES) | alpha, beta, gamma |
+| Extract-Features; Enhanced Man Interp; Power-Up Enhanced (DARE) | alpha, beta, gamma, delta |
+| Power-Up Enhanced (DARE+TIES) | beta, gamma, delta, epsilon, zeta |
+
+Reject assignments to ignored coefficients rather than pretend they affect the selected mode. Seeds remain main-node settings.
+
+#### Extraction
+
+Use the same bindings across LoRA, DoRA, learned DoRA, and both text-encoder extraction families:
+
+| Method | `a` | `b` | `c` | `d` | `e` | `f` |
+|---|---|---|---|---|---|---|
+| Fixed | linear_dim | conv_dim | clamp_quantile | min_diff | — | — |
+| Ratio | linear_ratio | conv_ratio | clamp_quantile | min_diff | linear_max_rank | conv_max_rank |
+| Quantile | linear_quantile | conv_quantile | clamp_quantile | min_diff | linear_max_rank | conv_max_rank |
+| Frobenius | linear_target | conv_target | clamp_quantile | min_diff | linear_max_rank | conv_max_rank |
+| Knee | linear_max_rank | conv_max_rank | clamp_quantile | min_diff | — | — |
+
+Knee method selection, probe offsets, SVD iterations, chunking controls, and learned-DoRA optimization settings remain on the receiving node.
+
+#### LoRA resize
+
+| Method | `a` | `b` | `c` |
+|---|---|---|---|
+| Fixed | new_rank | — | — |
+| Ratio | max_rank | ratio | — |
+| Frobenius | max_rank | min_rank | target |
+| Cumulative | max_rank | target | — |
+
+### Value validation
+
+Inherit values from the receiving node—not hardcoded configurator defaults.
+
+For supplied overrides:
+
+- Extraction dimensions and rank caps are integral values in `[1, 16384]`; ratios are `[1, 100]`; quantile/Frobenius targets are `[0, 1]`; `clamp_quantile` is `[0.5, 1]`; `min_diff` is `[0, 1]`.
+- Resize ranks are integral values in `[1, 3072]`; ratios are `[1, 100]`; targets are `[0.1, 1]`.
+- Validate resolved Frobenius resize `min_rank <= max_rank`, including values inherited from the main node.
+- Standard merge coefficients retain existing widget bounds. Where a coefficient represents a quantile or probability, validate `[0, 1]`; generic SVD rank coefficients must also be positive integral values.
+- Do not silently truncate fractional ranks or repair invalid assignments. Retain existing mathematical rank caps imposed by tensor dimensions.
+
+## 3. Matching and implementation data flow
+
+### Matching contract
+
+For ordinary model merging and extraction, match the receiving operation’s target tensor names.
+
+For LoRA inputs, match the normalized logical layer name after the existing parsing and normalization path. Do not match factor or alpha suffixes independently.
+
+Use the established dotted naming produced by extraction and existing reference mapping. Preserve internal identifiers such as `qkv_proj`. Do not introduce underscore normalization as the new public convention, rewrite regex punctuation, or redesign model-format conversion.
+
+Retain the existing unresolved flattened-name fallback when no reference mapping exists. Documentation must distinguish this fallback from the normal dotted path. A dotted rule that cannot match the resulting normalized input names produces the ordinary unmatched-rule error; do not silently reinterpret it.
+
+A single rule matching multiple keys belonging to one logical LoRA layer counts as one match. That layer receives one parameter set shared by its factors.
+
+### Strict errors and fallback
+
+Validate all rules before opening the output writer or beginning tensor streaming:
+
+1. Parse syntax and compile patterns once in the configurator.
+2. At each receiver, resolve aliases and reject unsupported parameters.
+3. Use header/key information to identify target layers and their existing normalized names.
+4. Detect matching rules per logical layer.
+5. Reject multiple matching rule lines for any one layer—even if their assignments concern different parameters.
+6. Reject any rule line matching no target layer.
+7. Combine the single matching rule with main-node values and validate the resulting parameters.
+
+Errors identify the receiving node, configuration line number, pattern, and affected layer or parameter. Overlap errors include every conflicting line number.
+
+Determine name matches before existing exclusions so an explicitly excluded layer is not falsely reported as nonexistent. Overrides never reactivate an excluded, discarded, guarded, or otherwise unprocessable layer.
+
+### Shared implementation
+
+Keep parsing, immutable rule records, binding definitions, and matching/error handling in the new shared module.
+
+Store only lightweight per-layer assignments and resolved scalar parameters. Do not store tensors, mutate the shared configuration payload, or modify global/main-node parameter dictionaries between layers.
+
+Register the configurator through the existing extension node list. Do not add dependencies.
+
+### Standard merger integration
+
+In `MergerLogic.execute_merge`:
+
+- Prepare the rule-to-layer mapping alongside existing header-based work planning.
+- For LoRA mode, use parsed logical groups to assign the same override to their factor keys; retain existing merge and alpha-normalization behavior.
+- Immediately before `calc_mode_class.create_recipe`, construct a layer-local parameter dictionary containing the base settings plus that layer’s overrides.
+- Keep tensor references and preloaded-source data local to the current work unit.
+- Do not mutate `recipe_params` with overrides that could affect later layers.
+
+Do not change recipe formulas, calculation-mode selection, or source-loading behavior.
+
+### Extraction integration
+
+Thread the optional rules through all 25 node entry points into the four extraction backends.
+
+Prepare resolved parameters before streaming, then use them in each backend’s `_process_layer` before `min_diff` testing and SVD selection.
+
+Specific requirements:
+
+- Fixed-rank overrides update both the requested dimension and corresponding maximum-rank argument; otherwise the original cap would negate an increased rank.
+- Adaptive overrides update the existing linear/conv method parameter and rank cap without altering method selection or probe settings.
+- Normal execution, CPU retry, and chunked fallback use the same resolved layer values.
+- Extend the existing chunked extraction helpers with a trailing `clamp_quantile` argument and pass the effective value through. These helpers currently omit it when calling `_svd_extract_linear`; text-encoder extraction shares the SVD helper.
+- Perform configuration validation outside broad extraction exception handlers so a configuration error cannot become a silently skipped layer.
+
+Do not alter SVD, chunk recombination, learned optimization, or tensor-lifecycle algorithms.
+
+### Resize integration
+
+Attach resolved rank/target values to the existing logical-layer work units.
+
+When calling `_resize_lora_factors`, translate the public bindings into the existing arguments:
+
+- Fixed `new_rank` → `new_rank`.
+- Adaptive `max_rank` → `new_rank`.
+- Ratio/target → `dynamic_param`.
+- Frobenius `min_rank` → `min_rank`.
+
+Leave `dynamic_method` on the main node. Preserve whole-layer factor handling, alpha behavior, passthrough data, low-bit handling, dtype selection, and CPU fallback.
+
+## 4. Verification and delivery
+
+Extend the repository-owned test selector with focused rule-parser and integration coverage. Map the new production module explicitly; do not rely solely on the broad repository-contract group.
+
+Required tests:
+
+- Both delimiter styles, mixed short/full names, scientific notation, nested regex groups, CRLF, blank lines, and comments.
+- Invalid syntax, invalid regex, unsupported parameters, duplicate aliases, nonfinite values, fractional ranks, and invalid resolved rank bounds.
+- Overlapping rules fail with layer and line numbers; unmatched rules fail before streaming or writing.
+- Partial assignments inherit remaining main-node values. Unmatched layers retain main-node behavior.
+- Existing include/exclude/discard and safety paths remain authoritative.
+- LoRA matching uses normalized logical names, preserves internal underscores, and assigns one configuration to paired factors.
+- All 39 receiver schemas append the optional socket without moving existing controls; every execution wrapper forwards it correctly.
+- Standard merge outputs reflect different coefficients on different layers, without leaking overrides to later layers.
+- Every extraction family uses the requested per-layer rank, clamp, and `min_diff`; cover adaptive methods, CPU retry, and forced chunk fallback.
+- All four resize variants use their resolved rank/target values, including inherited `min_rank` validation.
+- Absent or empty configuration preserves existing behavior.
+- Configuration failures preserve any existing destination file.
+
+Use small deterministic tensors and existing I/O fixtures. Do not run model inference or add unrelated performance, architecture, or normalization test projects.
+
+Run focused groups during development, `--changed` after accumulated changes, and the required full gate before the feature commit. Run changed-file lint, compilation, and `git diff --check`; distinguish pre-existing lint findings from new ones.
+
+Deliver the saved approved plan, implementation, generated binding documentation/tooltips, regression tests, and a verified feature commit. Leave unrelated analysis/report edits untouched.

pyproject.toml

@@ -1,9 +1,20 @@
 [project]
 name = "ComfyUI-ModelUtils"
 description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
-version = "0.2.7"
+version = "0.9.7"
 license = { file = "LICENSE" }
 
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
 [project.urls]
 Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils"
 #  Used by Comfy Registry https://comfyregistry.org

requirements.txt

@@ -1 +1,5 @@
-# No external dependencies required
+requests
+Pillow
+mutagen
+unifiedefficientloader>=0.5.2
+av

run_cwb_api_workflow.py

@@ -0,0 +1,259 @@
+"""Queue the CWB six-LoRA API workflow and print its merge report."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import re
+import sys
+import time
+from pathlib import Path
+from urllib import error, request
+
+
+DEFAULT_WORKFLOW = Path(__file__).with_name("CWB6LORATEST_API.json")
+MERGER_CLASS = "CWBLoRAMultiMerger"
+CUSTOM_CLASS = "CWBCustomConfiguration"
+PREVIEW_CLASS = "PreviewAny"
+CALIBRATION_EXCLUDE_PATTERN = (
+    r"^diffusion_model\.blocks\.(?:[12]|[4-9]|1[0-3]|1[5-9]|2[0-3]|2[56])\."
+    "\n"
+    r"^diffusion_model\.txtfusion\.layerwise_blocks\.1\."
+    "\n"
+    r"^diffusion_model\.txtfusion\.refiner_blocks\.0\."
+)
+
+
+def parse_counterfactual_weight_sweep(report: str) -> list[dict]:
+    pattern = re.compile(
+        r"^(?P<consensus_type>mean|median) \| "
+        r"(?P<similarity_threshold>[-+0-9.e]+) \| "
+        r"(?P<power_alpha>[-+0-9.e]+) \| (?P<diversity_beta>[-+0-9.e]+) \| "
+        r"(?P<dsc>true|false) \| (?P<softcb>true|false) \| "
+        r"(?P<accepted>\d+)/(?P<contributors>\d+) \| (?P<fallbacks>\d+) \| "
+        r"(?P<dominant_mean>[-+0-9.e]+)/(?P<dominant_max>[-+0-9.e]+) \| "
+        r"(?P<effective_mean>[-+0-9.e]+)/(?P<effective_min>[-+0-9.e]+)$",
+        re.MULTILINE,
+    )
+    rows = []
+    for match in pattern.finditer(report):
+        values = match.groupdict()
+        rows.append({
+            "consensus_type": values["consensus_type"],
+            "similarity_threshold": float(values["similarity_threshold"]),
+            "power_alpha": float(values["power_alpha"]),
+            "diversity_beta": float(values["diversity_beta"]),
+            "dynamic_similarity_contrast": values["dsc"] == "true",
+            "soft_comfort_bandpass": values["softcb"] == "true",
+            "accepted_contributors": int(values["accepted"]),
+            "contributors": int(values["contributors"]),
+            "fallbacks": int(values["fallbacks"]),
+            "dominant_weight_mean": float(values["dominant_mean"]),
+            "dominant_weight_max": float(values["dominant_max"]),
+            "effective_contributors_mean": float(values["effective_mean"]),
+            "effective_contributors_min": float(values["effective_min"]),
+        })
+    return rows
+
+
+def write_counterfactual_csv(path: Path, rows: list[dict]) -> None:
+    if not rows:
+        raise RuntimeError("CWB report does not contain a counterfactual weight sweep")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def _json_request(url: str, payload=None):
+    data = None if payload is None else json.dumps(payload).encode("utf-8")
+    req = request.Request(url, data=data)
+    if data is not None:
+        req.add_header("Content-Type", "application/json")
+    try:
+        with request.urlopen(req) as response:
+            return json.load(response)
+    except error.HTTPError as exc:
+        detail = exc.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"ComfyUI request failed ({exc.code}): {detail}") from exc
+    except error.URLError as exc:
+        raise RuntimeError(f"Could not reach ComfyUI at {url}: {exc.reason}") from exc
+
+
+def _single_node(workflow: dict, class_type: str) -> tuple[str, dict]:
+    matches = [
+        (node_id, node)
+        for node_id, node in workflow.items()
+        if node.get("class_type") == class_type
+    ]
+    if len(matches) != 1:
+        raise RuntimeError(
+            f"Expected exactly one {class_type} node, found {len(matches)}."
+        )
+    return matches[0]
+
+
+def _linked_preview_nodes(workflow: dict, merger_id: str) -> dict[int, str]:
+    linked = {}
+    for node_id, node in workflow.items():
+        if node.get("class_type") != PREVIEW_CLASS:
+            continue
+        source = node.get("inputs", {}).get("source")
+        if (
+            isinstance(source, list)
+            and len(source) == 2
+            and str(source[0]) == merger_id
+            and isinstance(source[1], int)
+        ):
+            linked[source[1]] = node_id
+    missing = {0, 2} - linked.keys()
+    if missing:
+        raise RuntimeError(
+            "Workflow must connect merger output 0 (filename) and output 2 "
+            "(CWB report) to PreviewAny nodes."
+        )
+    return linked
+
+
+def _prepare_workflow(
+    workflow: dict,
+    preset: str | None,
+    output_filename: str | None,
+    full_blocks: bool = False,
+) -> tuple[str, dict[int, str]]:
+    merger_id, merger = _single_node(workflow, MERGER_CLASS)
+    _single_node(workflow, CUSTOM_CLASS)
+    inputs = merger.setdefault("inputs", {})
+
+    if preset is not None:
+        inputs["cwb_preset"] = preset
+        inputs.pop("cwb_config", None)
+    if output_filename is not None:
+        inputs["output_filename"] = output_filename
+    if not full_blocks:
+        existing = inputs.get("exclude_patterns", "").strip()
+        inputs["exclude_patterns"] = "\n".join(
+            value for value in (existing, CALIBRATION_EXCLUDE_PATTERN) if value
+        )
+
+    return merger_id, _linked_preview_nodes(workflow, merger_id)
+
+
+def _prompt_payload(workflow: dict) -> dict:
+    return {"prompt": workflow}
+
+
+def _wait_for_history(server: str, prompt_id: str, poll_seconds: float):
+    while True:
+        history = _json_request(f"{server}/history/{prompt_id}")
+        if prompt_id in history:
+            return history[prompt_id]
+        time.sleep(poll_seconds)
+
+
+def _text_output(history_outputs: dict, node_id: str) -> str:
+    output = history_outputs.get(node_id, {})
+    candidates = (
+        output.get("text"),
+        output.get("ui", {}).get("text") if isinstance(output.get("ui"), dict) else None,
+    )
+    for candidate in candidates:
+        if isinstance(candidate, list) and candidate:
+            return str(candidate[0])
+        if isinstance(candidate, str):
+            return candidate
+    raise RuntimeError(f"PreviewAny node {node_id} did not return text in API history.")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Run the CWB six-LoRA API workflow and collect its report."
+    )
+    parser.add_argument("--workflow", type=Path, default=DEFAULT_WORKFLOW)
+    parser.add_argument("--server", default="http://127.0.0.1:8188")
+    parser.add_argument("--poll-seconds", type=float, default=2.0)
+    parser.add_argument(
+        "--preset",
+        help="Disconnect the custom configuration and run this CWB preset.",
+    )
+    parser.add_argument(
+        "--output-filename",
+        help="Override the merger output filename without its extension.",
+    )
+    parser.add_argument(
+        "--result-json",
+        type=Path,
+        help="Optionally write prompt ID, filename, and report as JSON.",
+    )
+    parser.add_argument(
+        "--full-blocks",
+        action="store_true",
+        help="Process all blocks for final verification instead of the calibration subset.",
+    )
+    parser.add_argument("--summary-only", action="store_true")
+    parser.add_argument("--counterfactual-csv", type=Path)
+    parser.add_argument("--alignment-threshold", type=float)
+    parser.add_argument("--alignment-method", choices=["index", "similarity"])
+    args = parser.parse_args()
+
+    if args.poll_seconds <= 0:
+        parser.error("--poll-seconds must be greater than zero")
+    if not args.workflow.is_file():
+        parser.error(f"workflow not found: {args.workflow}")
+
+    workflow = json.loads(args.workflow.read_text(encoding="utf-8"))
+    if args.alignment_threshold is not None:
+        _, custom = _single_node(workflow, CUSTOM_CLASS)
+        custom.setdefault("inputs", {})["alignment_threshold"] = (
+            args.alignment_threshold
+        )
+    if args.alignment_method is not None:
+        _, custom = _single_node(workflow, CUSTOM_CLASS)
+        custom.setdefault("inputs", {})["alignment_method"] = args.alignment_method
+    _, linked_outputs = _prepare_workflow(
+        workflow,
+        args.preset,
+        args.output_filename,
+        args.full_blocks,
+    )
+    server = args.server.rstrip("/")
+    queued = _json_request(f"{server}/prompt", _prompt_payload(workflow))
+    prompt_id = queued["prompt_id"]
+    sys.stdout.write(f"Queued prompt {prompt_id}\n")
+
+    history = _wait_for_history(server, prompt_id, args.poll_seconds)
+    status = history.get("status", {})
+    if status.get("status_str") == "error":
+        sys.stderr.write(json.dumps(status, indent=2) + "\n")
+        return 1
+
+    outputs = history.get("outputs", {})
+    output_filename = _text_output(outputs, linked_outputs[0])
+    cwb_report = _text_output(outputs, linked_outputs[2])
+    result = {
+        "prompt_id": prompt_id,
+        "output_filename": output_filename,
+        "cwb_report": cwb_report,
+    }
+    if args.result_json is not None:
+        args.result_json.parent.mkdir(parents=True, exist_ok=True)
+        args.result_json.write_text(
+            json.dumps(result, indent=2) + "\n",
+            encoding="utf-8",
+        )
+    if args.counterfactual_csv is not None:
+        write_counterfactual_csv(
+            args.counterfactual_csv,
+            parse_counterfactual_weight_sweep(cwb_report),
+        )
+
+    sys.stdout.write(f"Output: {output_filename}\n")
+    if not args.summary_only:
+        sys.stdout.write(f"\n{cwb_report}\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

run_cwb_preset_sweep.py

@@ -0,0 +1,418 @@
+"""Run a resumable one-factor CWB LoRA calibration sweep."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import os
+import re
+import sys
+import time
+from copy import deepcopy
+from pathlib import Path
+
+import run_cwb_api_workflow as api
+import run_model_analysis_api_workflow as analysis_api
+
+
+CONFIG_KEYS = (
+    "consensus_type",
+    "alignment_method",
+    "alignment_threshold",
+    "similarity_threshold",
+    "power_alpha",
+    "diversity_beta",
+    "rescale_norm",
+    "global_scale",
+    "dynamic_similarity_contrast",
+    "soft_comfort_bandpass",
+    "position_weight",
+    "preserve_common_prefix",
+)
+
+INTEGER_METRICS = {
+    "alignment_candidates": "Alignment candidates",
+    "alignment_matches": "Alignment matches",
+    "anchor_only_groups": "Anchor-only groups",
+    "weighted_vector_groups": "Weighted vector groups",
+    "all_rejected_fallbacks": "All-rejected equal fallbacks",
+    "zero_weight_fallbacks": "Zero-weight equal fallbacks",
+    "cuda_oom_fallbacks": "CUDA OOM CPU fallbacks",
+    "lora_groups": "LoRA groups",
+}
+
+TRIPLE_METRICS = {
+    "consensus_similarity": "Consensus similarity min/mean/max",
+    "dominant_weight": "Dominant normalized weight min/mean/max",
+    "effective_contributors": "Effective contributors min/mean/max",
+}
+
+DELTA_METRICS = (
+    "alignment_match_fraction",
+    "anchor_only_groups",
+    "accepted_weight_fraction",
+    "consensus_similarity_mean",
+    "dominant_weight_mean",
+    "dominant_weight_max",
+    "effective_contributors_mean",
+    "effective_contributors_min",
+    "norm_ratio_mean",
+    "norm_ratio_min",
+    "norm_ratio_max",
+    "all_rejected_fallbacks",
+    "zero_weight_fallbacks",
+    "wall_seconds",
+    "analysis_relative_l2",
+    "analysis_cosine",
+    "analysis_exact",
+    "analysis_norm_ratio",
+    "analysis_pair_cosine",
+    "analysis_mean_affinity_a",
+    "analysis_mean_affinity_b",
+    "analysis_median_affinity_a",
+    "analysis_median_affinity_b",
+)
+
+
+def _single_custom_node(workflow: dict) -> tuple[str, dict]:
+    return api._single_node(workflow, api.CUSTOM_CLASS)
+
+
+def _baseline_config(workflow: dict) -> dict:
+    _, node = _single_custom_node(workflow)
+    inputs = node.get("inputs", {})
+    missing = [key for key in CONFIG_KEYS if key not in inputs]
+    if missing:
+        raise RuntimeError(
+            "Custom configuration is missing required inputs: " + ", ".join(missing)
+        )
+    return {key: inputs[key] for key in CONFIG_KEYS}
+
+
+def build_sweep_cases(baseline: dict) -> list[dict]:
+    variations = [
+        ("alignment_method", "index"),
+        *[("alignment_threshold", value) for value in (0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01)],
+        *[("position_weight", value) for value in (0.0, 0.0001, 0.0005, 0.001, 0.005, 0.01)],
+        ("consensus_type", "mean"),
+        *[("similarity_threshold", value) for value in (-0.25, 0.1, 0.25, 0.5, 0.75)],
+        *[("power_alpha", value) for value in (0.5, 1.0, 3.0, 4.0)],
+        *[("diversity_beta", value) for value in (0.0, 1.0, 2.0, 4.0, 7.0)],
+        ("rescale_norm", False),
+        ("dynamic_similarity_contrast", True),
+        ("soft_comfort_bandpass", False),
+    ]
+    cases = [{"id": "baseline", "factor": "baseline", "value": None, "config": baseline.copy()}]
+    for factor, value in variations:
+        if baseline[factor] == value:
+            continue
+        config = baseline.copy()
+        config[factor] = value
+        value_text = str(value).lower().replace(".", "p").replace("-", "neg")
+        cases.append({
+            "id": f"{factor}_{value_text}",
+            "factor": factor,
+            "value": value,
+            "config": config,
+        })
+    return cases
+
+
+def parse_cwb_report(report: str) -> dict:
+    metrics = {}
+    for metric, label in INTEGER_METRICS.items():
+        match = re.search(rf"^{re.escape(label)}: (\d+)$", report, re.MULTILINE)
+        if match:
+            metrics[metric] = int(match.group(1))
+
+    accepted = re.search(
+        r"^Accepted weighting contributors: (\d+)/(\d+)$",
+        report,
+        re.MULTILINE,
+    )
+    if not accepted:
+        raise RuntimeError(
+            "CWB report lacks contributor-weight diagnostics. Restart ComfyUI "
+            "after loading the current consensus_merger.py before running the sweep."
+        )
+    metrics["accepted_weighting_contributors"] = int(accepted.group(1))
+    metrics["weighting_contributors"] = int(accepted.group(2))
+
+    number = r"(?:n/a|[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[-+]?\d+)?)"
+    for prefix, label in TRIPLE_METRICS.items():
+        match = re.search(
+            rf"^{re.escape(label)}: ({number}) / ({number}) / ({number})$",
+            report,
+            re.MULTILINE | re.IGNORECASE,
+        )
+        if not match:
+            raise RuntimeError(f"CWB report is missing metric: {label}")
+        for suffix, value in zip(("min", "mean", "max"), match.groups()):
+            metrics[f"{prefix}_{suffix}"] = None if value == "n/a" else float(value)
+
+    norm_groups = re.findall(
+        rf"^Rank-one delta norm ratio count/min/mean/max: (\d+) / "
+        rf"({number}) / ({number}) / ({number})$",
+        report,
+        re.MULTILINE | re.IGNORECASE,
+    )
+    numeric_norms = [
+        (int(count), *(float(value) for value in (minimum, mean, maximum)))
+        for count, minimum, mean, maximum in norm_groups
+        if count != "0" and "n/a" not in (minimum, mean, maximum)
+    ]
+    if numeric_norms:
+        norm_count = sum(group[0] for group in numeric_norms)
+        metrics["norm_ratio_count"] = norm_count
+        metrics["norm_ratio_min"] = min(group[1] for group in numeric_norms)
+        metrics["norm_ratio_mean"] = (
+            sum(group[0] * group[2] for group in numeric_norms) / norm_count
+        )
+        metrics["norm_ratio_max"] = max(group[3] for group in numeric_norms)
+
+    candidates = metrics.get("alignment_candidates", 0)
+    contributors = metrics["weighting_contributors"]
+    metrics["alignment_match_fraction"] = (
+        metrics.get("alignment_matches", 0) / candidates if candidates else 0.0
+    )
+    metrics["accepted_weight_fraction"] = (
+        metrics["accepted_weighting_contributors"] / contributors
+        if contributors else 0.0
+    )
+    return metrics
+
+
+def parse_analysis_reports(comparison: str, cwb_report: str) -> dict:
+    global_section = comparison.split("PARAMETER-WEIGHTED GLOBAL METRICS", 1)[1].split(
+        "EQUAL-LAYER AVERAGES", 1
+    )[0]
+
+    def single(label: str, text: str) -> float:
+        match = re.search(rf"^{re.escape(label)}: ([-+0-9.e]+)$", text, re.MULTILINE)
+        if not match:
+            raise RuntimeError(f"Analysis report is missing metric: {label}")
+        return float(match.group(1))
+
+    metrics = {
+        "analysis_relative_l2": single("relative_l2", global_section),
+        "analysis_cosine": single("cosine", global_section),
+        "analysis_exact": single("exact", global_section),
+        "analysis_norm_ratio": single("norm_ratio", global_section),
+    }
+    pairwise = re.search(
+        r"^Pairwise row cosine mean/min/max: ([-+0-9.e]+) /",
+        cwb_report,
+        re.MULTILINE,
+    )
+    mean_affinity = re.search(
+        r"^Mean-consensus affinity A/B: ([-+0-9.e]+) / ([-+0-9.e]+)$",
+        cwb_report,
+        re.MULTILINE,
+    )
+    median_affinity = re.search(
+        r"^Median-consensus affinity A/B: ([-+0-9.e]+) / ([-+0-9.e]+)$",
+        cwb_report,
+        re.MULTILINE,
+    )
+    if not pairwise or not mean_affinity or not median_affinity:
+        raise RuntimeError("CWB analysis report is missing global affinity metrics")
+    metrics.update({
+        "analysis_pair_cosine": float(pairwise.group(1)),
+        "analysis_mean_affinity_a": float(mean_affinity.group(1)),
+        "analysis_mean_affinity_b": float(mean_affinity.group(2)),
+        "analysis_median_affinity_a": float(median_affinity.group(1)),
+        "analysis_median_affinity_b": float(median_affinity.group(2)),
+    })
+    return metrics
+
+
+def _run_case(
+    workflow_template: dict,
+    case: dict,
+    server: str,
+    poll_seconds: float,
+    full_blocks: bool,
+    analysis_workflow: dict | None = None,
+) -> dict:
+    workflow = deepcopy(workflow_template)
+    _, custom = _single_custom_node(workflow)
+    custom["inputs"].update(case["config"])
+    _, linked = api._prepare_workflow(
+        workflow,
+        preset=None,
+        output_filename=None,
+        full_blocks=full_blocks,
+    )
+    started = time.perf_counter()
+    queued = api._json_request(f"{server}/prompt", api._prompt_payload(workflow))
+    prompt_id = queued["prompt_id"]
+    history = api._wait_for_history(server, prompt_id, poll_seconds)
+    wall_seconds = time.perf_counter() - started
+    status = history.get("status", {})
+    if status.get("status_str") == "error":
+        raise RuntimeError(
+            f"Sweep case {case['id']} failed: {json.dumps(status, indent=2)}"
+        )
+    outputs = history.get("outputs", {})
+    output_filename = api._text_output(outputs, linked[0])
+    report = api._text_output(outputs, linked[2])
+    metrics = parse_cwb_report(report)
+    metrics["wall_seconds"] = wall_seconds
+    result = {
+        **case,
+        "prompt_id": prompt_id,
+        "output_filename": output_filename,
+        "metrics": metrics,
+        "cwb_report": report,
+    }
+    if analysis_workflow is not None:
+        analysis_result = analysis_api.run_analysis(
+            deepcopy(analysis_workflow), server, poll_seconds
+        )
+        metrics.update(parse_analysis_reports(
+            analysis_result["comparison_report"],
+            analysis_result["cwb_report"],
+        ))
+        result["analysis"] = analysis_result
+    return result
+
+
+def _load_completed(path: Path) -> list[dict]:
+    if not path.exists():
+        return []
+    completed = []
+    with path.open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            try:
+                completed.append(json.loads(line))
+            except json.JSONDecodeError as exc:
+                raise RuntimeError(
+                    f"Invalid sweep JSONL at line {line_number}: {exc}"
+                ) from exc
+    return completed
+
+
+def _append_result(path: Path, result: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a", encoding="utf-8") as handle:
+        handle.write(json.dumps(result, ensure_ascii=False) + "\n")
+        handle.flush()
+        os.fsync(handle.fileno())
+
+
+def _write_influence_csv(path: Path, results: list[dict]) -> None:
+    baseline = next((result for result in results if result["id"] == "baseline"), None)
+    if baseline is None:
+        return
+    fields = ["id", "factor", "value", *DELTA_METRICS]
+    delta_fields = [f"delta_{metric}" for metric in DELTA_METRICS]
+    with path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=[*fields, *delta_fields])
+        writer.writeheader()
+        base_metrics = baseline["metrics"]
+        for result in results:
+            row = {key: result.get(key) for key in ("id", "factor", "value")}
+            for metric in DELTA_METRICS:
+                value = result["metrics"].get(metric)
+                base = base_metrics.get(metric)
+                row[metric] = value
+                row[f"delta_{metric}"] = (
+                    value - base if value is not None and base is not None else None
+                )
+            writer.writerow(row)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Run a resumable one-factor CWB LoRA parameter sweep."
+    )
+    parser.add_argument("--workflow", type=Path, default=api.DEFAULT_WORKFLOW)
+    parser.add_argument("--server", default="http://127.0.0.1:8188")
+    parser.add_argument("--poll-seconds", type=float, default=2.0)
+    parser.add_argument("--output-dir", type=Path, default=Path("audit/cwb_sweep"))
+    parser.add_argument("--full-blocks", action="store_true")
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--limit", type=int)
+    parser.add_argument(
+        "--case",
+        action="append",
+        dest="case_ids",
+        help="Run only the named case. Repeat to select multiple cases.",
+    )
+    parser.add_argument(
+        "--analysis-workflow",
+        type=Path,
+        default=analysis_api.DEFAULT_WORKFLOW,
+    )
+    parser.add_argument(
+        "--analysis-every",
+        type=int,
+        default=5,
+        help="Run candidate-versus-anchor analysis for baseline and every Nth case.",
+    )
+    args = parser.parse_args()
+
+    workflow = json.loads(args.workflow.read_text(encoding="utf-8"))
+    analysis_workflow = json.loads(args.analysis_workflow.read_text(encoding="utf-8"))
+    if args.full_blocks:
+        _, analysis_node = api._single_node(
+            analysis_workflow,
+            analysis_api.ANALYSIS_CLASS,
+        )
+        analysis_node.setdefault("inputs", {})["exclude_patterns"] = ""
+    cases = build_sweep_cases(_baseline_config(workflow))
+    if args.case_ids:
+        known_ids = {case["id"] for case in cases}
+        unknown = [case_id for case_id in args.case_ids if case_id not in known_ids]
+        if unknown:
+            raise RuntimeError("Unknown sweep case(s): " + ", ".join(unknown))
+        selected = set(args.case_ids)
+        cases = [case for case in cases if case["id"] in selected]
+    if args.limit is not None:
+        cases = cases[:args.limit]
+    if args.dry_run:
+        for case in cases:
+            sys.stdout.write(f"{case['id']}: {case['factor']}={case['value']}\n")
+        sys.stdout.write(f"Cases: {len(cases)}\n")
+        return 0
+
+    results_path = args.output_dir / "runs.jsonl"
+    completed = _load_completed(results_path)
+    completed_ids = {result["id"] for result in completed}
+    server = args.server.rstrip("/")
+    if args.analysis_every < 1:
+        raise RuntimeError("--analysis-every must be at least 1")
+    for index, case in enumerate(cases, start=1):
+        if case["id"] in completed_ids:
+            sys.stdout.write(f"[{index}/{len(cases)}] skip {case['id']}\n")
+            continue
+        sys.stdout.write(f"[{index}/{len(cases)}] run {case['id']}\n")
+        result = _run_case(
+            workflow,
+            case,
+            server,
+            args.poll_seconds,
+            args.full_blocks,
+            analysis_workflow=(
+                analysis_workflow
+                if index == 1 or index % args.analysis_every == 0
+                else None
+            ),
+        )
+        _append_result(results_path, result)
+        completed.append(result)
+        completed_ids.add(case["id"])
+        sys.stdout.write(
+            f"  {result['metrics']['wall_seconds']:.1f}s, "
+            f"match={result['metrics']['alignment_match_fraction']:.4f}, "
+            f"effective={result['metrics']['effective_contributors_mean']:.4f}\n"
+        )
+        _write_influence_csv(args.output_dir / "influence.csv", completed)
+
+    _write_influence_csv(args.output_dir / "influence.csv", completed)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

run_model_analysis_api_workflow.py

@@ -0,0 +1,73 @@
+"""Queue a model-analysis API workflow and capture both report outputs."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import run_cwb_api_workflow as api
+
+
+DEFAULT_WORKFLOW = Path("ANALYZETESTCANDIDATE_API.json")
+ANALYSIS_CLASS = "LoRAModelAnalysis"
+
+
+def _linked_previews(workflow: dict, analysis_id: str) -> dict[int, str]:
+    linked = {}
+    for node_id, node in workflow.items():
+        if node.get("class_type") != "PreviewAny":
+            continue
+        source = node.get("inputs", {}).get("source")
+        if isinstance(source, list) and len(source) == 2 and str(source[0]) == analysis_id:
+            linked[int(source[1])] = node_id
+    missing = [index for index in (0, 1) if index not in linked]
+    if missing:
+        raise RuntimeError(f"Missing PreviewAny capture for analysis output(s): {missing}")
+    return linked
+
+
+def run_analysis(workflow: dict, server: str, poll_seconds: float) -> dict:
+    analysis_id, _ = api._single_node(workflow, ANALYSIS_CLASS)
+    linked = _linked_previews(workflow, analysis_id)
+    queued = api._json_request(
+        f"{server.rstrip('/')}/prompt",
+        api._prompt_payload(workflow),
+    )
+    prompt_id = queued["prompt_id"]
+    history = api._wait_for_history(server.rstrip("/"), prompt_id, poll_seconds)
+    status = history.get("status", {})
+    if status.get("status_str") == "error":
+        raise RuntimeError(f"Analysis failed: {json.dumps(status, indent=2)}")
+    outputs = history.get("outputs", {})
+    return {
+        "prompt_id": prompt_id,
+        "comparison_report": api._text_output(outputs, linked[0]),
+        "cwb_report": api._text_output(outputs, linked[1]),
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Run a model-analysis API workflow.")
+    parser.add_argument("--workflow", type=Path, default=DEFAULT_WORKFLOW)
+    parser.add_argument("--server", default="http://127.0.0.1:8188")
+    parser.add_argument("--poll-seconds", type=float, default=2.0)
+    parser.add_argument("--result-json", type=Path)
+    args = parser.parse_args()
+
+    workflow = json.loads(args.workflow.read_text(encoding="utf-8"))
+    result = run_analysis(workflow, args.server, args.poll_seconds)
+    if args.result_json:
+        args.result_json.parent.mkdir(parents=True, exist_ok=True)
+        args.result_json.write_text(
+            json.dumps(result, indent=2, ensure_ascii=False) + "\n",
+            encoding="utf-8",
+        )
+    sys.stdout.write(f"Prompt: {result['prompt_id']}\n")
+    sys.stdout.write("Analysis reports captured\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/conftest.py

@@ -0,0 +1,8 @@
+def pytest_report_teststatus(report, config):
+    """Pytest hook: emit only 'F' for failed tests and nothing for passing/skipped tests."""
+    if report.when in ("setup", "call") and report.failed:
+        return "failed", "F", "FAILED"
+    if report.passed:
+        return "passed", "", ""
+    if report.skipped:
+        return "skipped", "", ""

tests/run_tests.py

@@ -0,0 +1,249 @@
+"""Select repository tests from changed paths or named subsystem groups."""
+
+# ruff: noqa: T201
+
+from __future__ import annotations
+
+import argparse
+import fnmatch
+import os
+import shutil
+import subprocess
+import sys
+import tomllib
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
+COMFYUI_ROOT = REPOSITORY_ROOT.parent.parent
+MANIFEST_PATH = Path(__file__).with_name("test_groups.toml")
+TEMP_ROOT = REPOSITORY_ROOT / ".pytest-tmp"
+
+
+@dataclass(frozen=True)
+class TestGroup:
+    name: str
+    paths: tuple[str, ...]
+    python_tests: tuple[str, ...]
+
+
+@dataclass
+class Selection:
+    groups: set[str] = field(default_factory=set)
+    python_tests: set[str] = field(default_factory=set)
+    reasons: dict[str, set[str]] = field(default_factory=dict)
+    unmapped: set[str] = field(default_factory=set)
+
+
+def _parse_groups(text: str) -> dict[str, TestGroup]:
+    data = tomllib.loads(text)
+    return {
+        name: TestGroup(
+            name=name,
+            paths=tuple(values.get("paths", ())),
+            python_tests=tuple(values.get("python_tests", ())),
+        )
+        for name, values in data.get("groups", {}).items()
+    }
+
+
+def load_groups(path: Path = MANIFEST_PATH) -> dict[str, TestGroup]:
+    return _parse_groups(path.read_text(encoding="utf-8"))
+
+
+def load_groups_from_revision(revision: str) -> dict[str, TestGroup]:
+    manifest = MANIFEST_PATH.relative_to(REPOSITORY_ROOT).as_posix()
+    result = subprocess.run(
+        ["git", "show", f"{revision}:{manifest}"], cwd=REPOSITORY_ROOT,
+        check=False, capture_output=True, text=True,
+    )
+    return _parse_groups(result.stdout) if result.returncode == 0 else {}
+
+
+def git_lines(*args: str) -> set[str]:
+    result = subprocess.run(
+        ["git", *args], cwd=REPOSITORY_ROOT, check=True,
+        capture_output=True, text=True,
+    )
+    return {
+        line.strip().replace("\\", "/")
+        for line in result.stdout.splitlines() if line.strip()
+    }
+
+
+def changed_paths(base: str | None = None) -> set[str]:
+    paths = git_lines("diff", "--name-only", "--relative", "HEAD")
+    if base:
+        paths.update(git_lines("diff", "--name-only", "--relative", f"{base}...HEAD"))
+    return paths
+
+
+def is_production_source(path: str) -> bool:
+    normalized = path.replace("\\", "/")
+    return normalized == "__init__.py" or (
+        normalized.startswith("nodes/") and normalized.endswith(".py")
+    )
+
+
+def _add_group(selection: Selection, group: TestGroup, reason: str) -> None:
+    selection.groups.add(group.name)
+    selection.python_tests.update(group.python_tests)
+    selection.reasons.setdefault(group.name, set()).add(reason)
+
+
+def _path_matches_group(path: str, group: TestGroup) -> bool:
+    return any(fnmatch.fnmatchcase(path, pattern) for pattern in group.paths)
+
+
+def _existing_group(group: TestGroup) -> TestGroup:
+    return TestGroup(
+        name=group.name, paths=group.paths,
+        python_tests=tuple(
+            path for path in group.python_tests if (REPOSITORY_ROOT / path).is_file()
+        ),
+    )
+
+
+def select_tests(
+    paths: set[str], groups: dict[str, TestGroup],
+    explicit_groups: tuple[str, ...] = (),
+    historical_groups: tuple[dict[str, TestGroup], ...] = (),
+) -> Selection:
+    selection = Selection()
+    for name in explicit_groups:
+        if name not in groups:
+            raise ValueError(f"Unknown test group: {name}")
+        _add_group(selection, groups[name], "explicit selection")
+
+    for raw_path in sorted(paths):
+        path = raw_path.replace("\\", "/")
+        if path.startswith("tests/test_") and path.endswith(".py"):
+            if (REPOSITORY_ROOT / path).is_file():
+                selection.python_tests.add(path)
+                selection.reasons.setdefault("direct test", set()).add(path)
+            continue
+
+        matched = False
+        for group in groups.values():
+            if _path_matches_group(path, group):
+                _add_group(selection, group, path)
+                matched = True
+        if not matched and not (REPOSITORY_ROOT / path).exists():
+            for historical in historical_groups:
+                for historical_group in historical.values():
+                    if _path_matches_group(path, historical_group):
+                        active = groups.get(historical_group.name)
+                        _add_group(
+                            selection, active or _existing_group(historical_group),
+                            f"{path} (deleted)",
+                        )
+                        matched = True
+        if not matched and is_production_source(path) and (REPOSITORY_ROOT / path).exists():
+            selection.unmapped.add(path)
+    return selection
+
+
+def tracked_final_tests() -> set[str]:
+    tracked = git_lines("ls-files", "tests/test_*.py")
+    return {path for path in tracked if (REPOSITORY_ROOT / path).is_file()}
+
+
+def print_selection(selection: Selection) -> None:
+    for group in sorted(selection.groups):
+        reasons = ", ".join(sorted(selection.reasons.get(group, ())))
+        print(f"group {group}: {reasons}")
+    for reason in sorted(selection.reasons.get("direct test", ())):
+        print(f"direct test: {reason}")
+    print("python tests:")
+    for path in sorted(selection.python_tests):
+        print(f"  {path}")
+
+
+def run_selection(selection: Selection) -> int:
+    TEMP_ROOT.mkdir(exist_ok=True)
+    basetemp = (TEMP_ROOT / f"run-{uuid.uuid4().hex}").resolve()
+    if TEMP_ROOT.resolve() not in basetemp.parents:
+        raise RuntimeError(f"Unsafe pytest temporary path: {basetemp}")
+    try:
+        environment = os.environ.copy()
+        environment["PYTHONPATH"] = os.pathsep.join(
+            filter(None, (str(REPOSITORY_ROOT), environment.get("PYTHONPATH", "")))
+        )
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "pytest",
+                "-q",
+                "--import-mode=importlib",
+                "-o",
+                "addopts=",
+                "--capture=sys",
+                "--basetemp",
+                str(basetemp),
+                *[
+                    str((REPOSITORY_ROOT / path).resolve())
+                    for path in sorted(selection.python_tests)
+                ],
+            ],
+            cwd=COMFYUI_ROOT,
+            env=environment,
+            check=False,
+        )
+        return result.returncode
+    finally:
+        if basetemp.exists():
+            shutil.rmtree(basetemp)
+        if TEMP_ROOT.exists() and not any(TEMP_ROOT.iterdir()):
+            TEMP_ROOT.rmdir()
+
+
+def parser() -> argparse.ArgumentParser:
+    result = argparse.ArgumentParser(description=__doc__)
+    mode = result.add_mutually_exclusive_group()
+    mode.add_argument("--final", action="store_true", help="run every tracked test")
+    mode.add_argument("--list-groups", action="store_true", help="list configured groups")
+    result.add_argument("--changed", action="store_true", help="select changed paths (default)")
+    result.add_argument("--base", help="include committed changes since BASE")
+    result.add_argument("--group", action="append", default=[], help="add a named group")
+    result.add_argument("--dry-run", action="store_true", help="explain without executing")
+    return result
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parser().parse_args(argv)
+    groups = load_groups()
+    if args.list_groups:
+        for name in sorted(groups):
+            print(name)
+        return 0
+    if args.final:
+        selection = Selection(
+            groups={"final"}, python_tests=tracked_final_tests(),
+            reasons={"final": {"every tracked test"}},
+        )
+    else:
+        paths = changed_paths(args.base) if args.changed or args.base or not args.group else set()
+        revisions = ["HEAD", *([args.base] if args.base else [])]
+        historical = tuple(load_groups_from_revision(rev) for rev in revisions) if paths else ()
+        selection = select_tests(paths, groups, tuple(args.group), historical)
+
+    if selection.unmapped:
+        print("Unmapped production source files:", file=sys.stderr)
+        for path in sorted(selection.unmapped):
+            print(f"  {path}", file=sys.stderr)
+        print("Update tests/test_groups.toml or use --final.", file=sys.stderr)
+        return 2
+    if args.dry_run:
+        print_selection(selection)
+        return 0
+    if not selection.python_tests:
+        print("No tests selected.")
+        return 0
+    return run_selection(selection)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_adaptive_extraction_pipeline.py

@@ -0,0 +1,381 @@
+import importlib
+import json
+import sys
+import types
+from collections import Counter
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_adaptive_extraction_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 adaptive():
+    return _load_module("nodes.adaptive_svd")
+
+
+def _exact_partial(weight, q, niter):
+    del niter
+    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
+    return U[:, :q], S[:q], Vh[:q, :].T
+
+
+@pytest.mark.parametrize(
+    ("mode", "parameter", "expected"),
+    [("ratio", 2.0, 2), ("quantile", 0.7, 2), ("sv_fro", 0.8, 1)],
+)
+def test_adaptive_partial_rank_matches_reference(monkeypatch, adaptive, mode, parameter, expected):
+    monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial)
+    weight = torch.diag(torch.tensor([8.0, 5.0, 2.0, 1.0, 0.5, 0.25]))
+    _, _, _, rank = adaptive.adaptive_partial_svd(
+        weight, mode, parameter, max_rank=4, probe_offset=1, niter=2
+    )
+    assert rank == expected
+
+
+def test_unresolved_ratio_expands_and_caps(monkeypatch, adaptive):
+    probes = []
+
+    def tracked(weight, q, niter):
+        probes.append(q)
+        return _exact_partial(weight, q, niter)
+
+    monkeypatch.setattr(adaptive.torch, "svd_lowrank", tracked)
+    weight = torch.diag(torch.linspace(10.0, 8.0, 12))
+    _, _, _, rank = adaptive.adaptive_partial_svd(
+        weight, "ratio", 100.0, max_rank=4, probe_offset=2, niter=2
+    )
+    assert probes == [6, 8]
+    assert rank == 4
+
+
+def test_frobenius_uses_energy_outside_output_cap(monkeypatch, adaptive):
+    monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial)
+    weight = torch.eye(10)
+    _, _, _, rank = adaptive.adaptive_partial_svd(
+        weight, "sv_fro", 0.9, max_rank=3, probe_offset=2, niter=2
+    )
+    assert rank == 3
+
+
+class _FakeStream:
+    def __init__(self, batches):
+        self._batches = iter(batches)
+        self.closed = False
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return next(self._batches)
+
+    def close(self):
+        self.closed = True
+
+
+class _FakeHandler:
+    def __init__(self, values):
+        self.values = values
+        self.calls = []
+        self.marked = []
+        self.streams = []
+
+    def async_stream(self, keys, **kwargs):
+        self.calls.append((tuple(keys), kwargs))
+        stream = _FakeStream([[(key, self.values[key])] for key in keys])
+        self.streams.append(stream)
+        return stream
+
+    def mark_processed(self, key):
+        self.marked.append(key)
+
+
+def test_paired_stream_is_bounded_ordered_and_releases_once():
+    streaming = _load_module("nodes.extraction_stream")
+    handler_a = _FakeHandler({"a1": torch.tensor(1), "a2": torch.tensor(2)})
+    handler_b = _FakeHandler({"b1": torch.tensor(3)})
+    units = [("one", "a1", "b1"), ("two", "a2", None)]
+
+    observed = []
+    for key, tensor_a, tensor_b in streaming.paired_async_tensors(
+        handler_a, handler_b, units, pin_memory=True
+    ):
+        observed.append((key, tensor_a.item(), None if tensor_b is None else tensor_b.item()))
+
+    assert observed == [("one", 1, 3), ("two", 2, None)]
+    assert handler_a.calls == [(('a1', 'a2'), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})]
+    assert handler_b.calls == [(('b1',), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})]
+    assert Counter(handler_a.marked) == Counter({"a1": 1, "a2": 1})
+    assert Counter(handler_b.marked) == Counter({"b1": 1})
+    assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
+
+
+def test_paired_stream_decodes_quantized_logical_keys_per_work_unit():
+    streaming = _load_module("nodes.extraction_stream")
+    handler_a = _FakeHandler({
+        "layer1.code": torch.tensor(2),
+        "layer1.scale": torch.tensor(3),
+        "layer2.code": torch.tensor(4),
+        "layer2.scale": torch.tensor(5),
+    })
+    handler_b = _FakeHandler({"layer1": torch.tensor(1), "layer2": torch.tensor(2)})
+
+    class _Quantization:
+        quantized_keys = {"layer1", "layer2"}
+
+        def required_keys(self, key):
+            return (f"{key}.code", f"{key}.scale")
+
+        def decode(self, key, tensors, compute_dtype, device="cpu"):
+            assert set(tensors) == {f"{key}.code", f"{key}.scale"}
+            return tensors[f"{key}.code"] * tensors[f"{key}.scale"]
+
+    quantization = _Quantization()
+    observed = list(streaming.paired_async_tensors(
+        handler_a,
+        handler_b,
+        [("layer1", "layer1", "layer1"), ("layer2", "layer2", "layer2")],
+        pin_memory=False,
+        quantization_a=quantization,
+    ))
+
+    assert [(key, tensor_a.item(), tensor_b.item()) for key, tensor_a, tensor_b in observed] == [
+        ("layer1", 6, 1),
+        ("layer2", 20, 2),
+    ]
+    assert handler_a.calls[0][0] == (
+        "layer1.code", "layer1.scale", "layer2.code", "layer2.scale"
+    )
+    assert Counter(handler_a.marked) == Counter({
+        "layer1.code": 1, "layer1.scale": 1, "layer2.code": 1, "layer2.scale": 1,
+    })
+    assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
+
+
+def test_cuda_oom_compute_retries_on_cpu(monkeypatch):
+    streaming = _load_module("nodes.extraction_stream")
+    monkeypatch.setattr(streaming, "release_cuda_oom", lambda: None)
+    calls = []
+
+    def fail_cuda():
+        calls.append("cuda")
+        raise RuntimeError("CUDA out of memory")
+
+    def run_cpu():
+        calls.append("cpu")
+        return 7
+
+    assert streaming.retry_cuda_oom_on_cpu(fail_cuda, run_cpu, "cuda:0") == 7
+    assert calls == ["cuda", "cpu"]
+
+
+def test_adaptive_controls_cover_all_public_families():
+    modules_and_nodes = [
+        (_load_module("nodes.lora_extract_svd"), prefix)
+        for prefix in ("LoRAExtractRatio", "LoRAExtractQuantile", "LoRAExtractFrobenius")
+    ]
+    modules_and_nodes += [
+        (_load_module("nodes.dora_extract_wd"), prefix)
+        for prefix in ("DoRAExtractRatio", "DoRAExtractQuantile", "DoRAExtractFrobenius")
+    ]
+    modules_and_nodes += [
+        (_load_module("nodes.dora_learned_wd"), prefix)
+        for prefix in ("DoRALearnedExtractRatio", "DoRALearnedExtractQuantile", "DoRALearnedExtractFrobenius")
+    ]
+    modules_and_nodes += [
+        (_load_module("nodes.text_encoder_extract"), prefix)
+        for prefix in (
+            "TextEncoderLoRAExtractRatio", "TextEncoderLoRAExtractQuantile",
+            "TextEncoderLoRAExtractFrobenius", "TextEncoderDoRAExtractRatio",
+            "TextEncoderDoRAExtractQuantile", "TextEncoderDoRAExtractFrobenius",
+        )
+    ]
+
+    for module, node_name in modules_and_nodes:
+        names = [item.id for item in getattr(module, node_name).define_schema().inputs]
+        probe_index = names.index("probe_offset")
+        assert names[probe_index + 1] == "linear_max_rank"
+
+
+def _write_uel(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
+        writer.write_dict(tensors)
+
+
+def test_lora_extraction_uses_only_async_uel_and_releases_sources(monkeypatch, tmp_path):
+    extraction = _load_module("nodes.lora_extract_svd")
+    model_a = tmp_path / "a.safetensors"
+    model_b = tmp_path / "b.safetensors"
+    output = tmp_path / "out.safetensors"
+    key = "model.diffusion_model.layer.weight"
+    _write_uel(model_a, {key: torch.eye(3)})
+    _write_uel(model_b, {key: torch.zeros(3, 3)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    original_async = extraction.MemoryEfficientSafeOpen.async_stream
+    original_mark = extraction.MemoryEfficientSafeOpen.mark_processed
+    async_calls = []
+    marked = []
+
+    def track_async(self, keys, **kwargs):
+        async_calls.append((self.filename, tuple(keys), kwargs))
+        yield from original_async(self, keys, **kwargs)
+
+    def track_mark(self, source_key):
+        marked.append((self.filename, source_key))
+        return original_mark(self, source_key)
+
+    monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "async_stream", track_async)
+    monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "mark_processed", track_mark)
+    monkeypatch.setattr(
+        extraction.MemoryEfficientSafeOpen,
+        "get_tensor",
+        lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("synchronous tensor load")),
+    )
+
+    extraction.extract_lora_from_files(
+        str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+    )
+
+    assert output.is_file()
+    assert len(async_calls) == 2
+    assert all(options == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, _, options in async_calls)
+    assert Counter(marked) == Counter((filename, key) for filename, keys, _ in async_calls for key in keys)
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        assert sorted(result.keys()) == [
+            "diffusion_model.layer.lora_A.weight",
+            "diffusion_model.layer.lora_B.weight",
+        ]
+
+
+@pytest.mark.parametrize("prefixed_input", ["a", "b"])
+def test_lora_extraction_matches_prefixed_and_bare_models(monkeypatch, tmp_path, prefixed_input):
+    extraction = _load_module("nodes.lora_extract_svd")
+    paths = {name: tmp_path / f"{name}.safetensors" for name in ("a", "b")}
+    output = tmp_path / "mixed_prefix.safetensors"
+    for name in ("a", "b"):
+        key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
+        value = torch.tensor([[2.0, 0.0], [0.0, 0.0]]) if name == "a" else torch.zeros((2, 2))
+        if name == prefixed_input:
+            stem = key[:-len(".weight")]
+            _write_uel(paths[name], {
+                key: value.to(torch.int8),
+                f"{stem}.weight_scale": torch.tensor(1.0),
+                f"{stem}.comfy_quant": torch.tensor(
+                    list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
+                ),
+            })
+        else:
+            _write_uel(paths[name], {key: value})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    extraction.extract_lora_from_files(
+        str(paths["a"]), str(paths["b"]), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+    )
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        keys = sorted(result.keys())
+        assert keys == [
+            "diffusion_model.layer.lora_A.weight",
+            "diffusion_model.layer.lora_B.weight",
+        ]
+        stream = result.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for key in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == key
+                tensors[key] = batch[0][1]
+        finally:
+            for key in tensors:
+                result.mark_processed(key)
+            stream.close()
+    delta = tensors[keys[1]] @ tensors[keys[0]]
+    torch.testing.assert_close(delta, torch.tensor([[2.0, 0.0], [0.0, 0.0]]), atol=1e-4, rtol=0)
+
+
+@pytest.mark.parametrize("module_name,function_name,extra", [
+    ("nodes.dora_extract_wd", "extract_dora_from_files", {}),
+    ("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
+])
+@pytest.mark.parametrize("prefixed_input", ["a", "b"])
+def test_dora_extraction_matches_prefixed_and_bare_models(
+    monkeypatch, tmp_path, module_name, function_name, extra, prefixed_input,
+):
+    module = _load_module(module_name)
+    paths = {name: tmp_path / f"{name}.safetensors" for name in ("a", "b")}
+    output = tmp_path / "mixed_prefix_dora.safetensors"
+    for name in ("a", "b"):
+        key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
+        value = torch.diag(torch.tensor([3.0, 1.0])) if name == "a" else torch.eye(2)
+        if name == prefixed_input:
+            stem = key[:-len(".weight")]
+            _write_uel(paths[name], {
+                key: value.to(torch.int8),
+                f"{stem}.weight_scale": torch.tensor(1.0),
+                f"{stem}.comfy_quant": torch.tensor(
+                    list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
+                ),
+            })
+        else:
+            _write_uel(paths[name], {key: value})
+    monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
+    getattr(module, function_name)(
+        str(paths["a"]), str(paths["b"]), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False, **extra,
+    )
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        keys = set(result.keys())
+    assert "diffusion_model.layer.lora_down.weight" in keys
+    assert "diffusion_model.layer.lora_up.weight" in keys
+    assert not any("comfy_quant" in key or "model.diffusion_model." in key for key in keys)
+
+
+@pytest.mark.parametrize("module_name, function_name, extra", [
+    ("nodes.lora_extract_svd", "extract_lora_from_files", {"include_1d_diffs": False}),
+    ("nodes.dora_extract_wd", "extract_dora_from_files", {}),
+    ("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
+    ("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": False}),
+    ("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": True}),
+])
+def test_extraction_include_mode_selects_only_matching_layers(monkeypatch, tmp_path, module_name, function_name, extra):
+    module = _load_module(module_name)
+    first, second = tmp_path / "first.safetensors", tmp_path / "second.safetensors"
+    match, other = "model.encoder.match.weight", "model.encoder.other.weight"
+    _write_uel(first, {match: torch.eye(2), other: torch.eye(2)})
+    _write_uel(second, {match: torch.zeros(2, 2), other: torch.zeros(2, 2)})
+    monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
+    common = dict(linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+                  skip_patterns_str="match", include_mode=True, **extra)
+    selected = tmp_path / "selected.safetensors"
+    getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(selected), **common)
+    with MemoryEfficientSafeOpen(str(selected), low_memory=True) as result:
+        keys = list(result.keys())
+        assert keys
+        assert all("match" in key for key in keys)
+    empty = tmp_path / "empty.safetensors"
+    common["skip_patterns_str"] = ""
+    getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(empty), **common)
+    with MemoryEfficientSafeOpen(str(empty), low_memory=True) as result:
+        assert list(result.keys()) == []
+    globbed = tmp_path / "globbed.safetensors"
+    common.update(skip_patterns_str="*.match.weight", glob_skip_patterns=True)
+    getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(globbed), **common)
+    with MemoryEfficientSafeOpen(str(globbed), low_memory=True) as result:
+        assert list(result.keys()) == keys

tests/test_artifact_paths.py

@@ -0,0 +1,120 @@
+import importlib
+import os
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_artifact_path_tests"
+MIGRATED_SAVE_MODULES = (
+    "consensus_merger.py",
+    "cwb_model_graft.py",
+    "cwb_delta_lora_merger.py",
+    "merger.py",
+    "lora_merger.py",
+    "lodestone_merger.py",
+    "lora_resize.py",
+    "dtype_conversion.py",
+    "lora_extract_svd.py",
+    "dora_extract_wd.py",
+    "dora_learned_wd.py",
+    "text_encoder_extract.py",
+    "renamekeys.py",
+    "prunekeys.py",
+    "lora_rename.py",
+    "minimax_h3_lora_convert.py",
+    "minimax_h3_fold.py",
+)
+
+
+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 artifact_paths():
+    return _load_module("nodes.artifact_paths")
+
+
+@pytest.mark.parametrize(
+    "category",
+    ["checkpoints", "diffusion_models", "text_encoders", "loras", "embeddings"],
+)
+def test_canonical_categories_use_models_dir_and_return_relative_name(
+    monkeypatch, tmp_path, artifact_paths, category
+):
+    monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
+    output_path, output_name = artifact_paths.canonical_model_artifact_path(
+        category, "group/subgroup/model"
+    )
+
+    expected = tmp_path / category / "group" / "subgroup" / "model.safetensors"
+    assert Path(output_path) == expected
+    assert output_name == "group/subgroup/model.safetensors"
+    assert expected.parent.is_dir()
+    assert not os.path.isabs(output_name)
+
+
+def test_registered_search_paths_do_not_influence_write_root(
+    monkeypatch, tmp_path, artifact_paths
+):
+    canonical = tmp_path / "cli-model-root"
+    legacy = tmp_path / "legacy-unet"
+    additional = tmp_path / "additional-models"
+    monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(canonical))
+    monkeypatch.setattr(
+        artifact_paths.folder_paths,
+        "get_folder_paths",
+        lambda _: [str(legacy), str(additional)],
+    )
+
+    output_path, output_name = artifact_paths.canonical_model_artifact_path(
+        "diffusion_models", "nested/model.safetensors"
+    )
+
+    assert Path(output_path) == canonical / "diffusion_models" / "nested" / "model.safetensors"
+    assert output_name == "nested/model.safetensors"
+
+
+@pytest.mark.parametrize(
+    "name",
+    ["model.safetensors", "model.safetensors.safetensors"],
+)
+def test_appends_exactly_one_extension(monkeypatch, tmp_path, artifact_paths, name):
+    monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
+    _, output_name = artifact_paths.canonical_model_artifact_path("loras", name)
+    assert output_name == "model.safetensors"
+
+
+@pytest.mark.parametrize(
+    "name",
+    ["../escape", "folder/../../escape", "", "   ", "/absolute", ".safetensors"],
+)
+def test_rejects_empty_or_escaping_names(monkeypatch, tmp_path, artifact_paths, name):
+    monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
+    with pytest.raises(ValueError):
+        artifact_paths.canonical_model_artifact_path("loras", name)
+
+
+def test_rejects_unrecognized_category(monkeypatch, tmp_path, artifact_paths):
+    monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
+    with pytest.raises(ValueError, match="Unsupported model artifact category"):
+        artifact_paths.canonical_model_artifact_path("unet", "model")
+
+
+def test_persistent_save_nodes_use_shared_path_contract_and_wildcard_outputs():
+    for filename in MIGRATED_SAVE_MODULES:
+        source = (REPO_ROOT / "nodes" / filename).read_text(encoding="utf-8")
+        assert (
+            "canonical_model_artifact_path" in source
+            or "_build_lora_output_path" in source
+        ), filename
+        assert 'io.String.Output(display_name="output_path")' not in source, filename
+        assert 'io.String.Output(display_name="output_filename")' not in source, filename

tests/test_conditional_registration.py

@@ -0,0 +1,49 @@
+import asyncio
+import importlib.util
+import sys
+from pathlib import Path
+from comfy_api.latest import io
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_extension_module():
+    package_name = "ComfyUI_ModelUtils_init_test"
+    if package_name not in sys.modules:
+        spec = importlib.util.spec_from_file_location(
+            package_name,
+            REPO_ROOT / "__init__.py",
+            submodule_search_locations=[str(REPO_ROOT)],
+        )
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[package_name] = module
+        spec.loader.exec_module(module)
+    return sys.modules[package_name]
+
+
+def test_startup_node_registration_when_downloaders_absent(monkeypatch):
+    ext_module = _load_extension_module()
+
+    ext = ext_module.ModelUtilsExtension()
+    nodes_with_downloader = asyncio.run(ext.get_node_list())
+    assert len(nodes_with_downloader) > 0
+
+    # Simulate registry install by setting DOWNLOADER_NODES to empty
+    monkeypatch.setattr(ext_module, "DOWNLOADER_NODES", [])
+    nodes_registry = asyncio.run(ext.get_node_list())
+    assert len(nodes_registry) < len(nodes_with_downloader)
+    assert all(issubclass(n, io.ComfyNode) for n in nodes_registry)
+
+
+def test_model_info_nodes_importable_standalone():
+    spec = importlib.util.spec_from_file_location(
+        "modelutils_model_info_nodes_test",
+        REPO_ROOT / "nodes" / "model_info_nodes.py",
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    assert hasattr(module, "get_potential_preview_files")
+    assert hasattr(module, "load_image_tensor")
+    assert hasattr(module, "get_model_workflows")
+    assert hasattr(module, "get_model_metadata_file")

tests/test_cwb_api_runner.py

@@ -0,0 +1,141 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_runner():
+    spec = importlib.util.spec_from_file_location(
+        "cwb_api_runner_test", REPO_ROOT / "run_cwb_api_workflow.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture(scope="module")
+def runner():
+    return _load_runner()
+
+
+@pytest.fixture
+def workflow():
+    return {
+        "10": {
+            "class_type": "CWBCustomConfiguration",
+            "inputs": {"diversity_beta": 10.0},
+        },
+        "11": {
+            "class_type": "CWBLoRAMultiMerger",
+            "inputs": {
+                "cwb_config": ["10", 0],
+                "cwb_preset": "default",
+                "output_filename": "original",
+            },
+        },
+        "14": {"class_type": "PreviewAny", "inputs": {"source": ["11", 0]}},
+        "15": {"class_type": "PreviewAny", "inputs": {"source": ["11", 2]}},
+    }
+
+
+def test_custom_export_remains_connected_by_default(runner, workflow):
+    merger_id, linked = runner._prepare_workflow(workflow, None, None)
+    assert merger_id == "11"
+    assert workflow["11"]["inputs"]["cwb_config"] == ["10", 0]
+    assert workflow["11"]["inputs"]["exclude_patterns"] == (
+        runner.CALIBRATION_EXCLUDE_PATTERN
+    )
+    assert linked == {0: "14", 2: "15"}
+
+
+def test_preset_run_disconnects_custom_config_and_overrides_filename(
+    runner, workflow
+):
+    runner._prepare_workflow(workflow, "candidate_sim_mean", "candidate/output")
+    inputs = workflow["11"]["inputs"]
+    assert "cwb_config" not in inputs
+    assert inputs["cwb_preset"] == "candidate_sim_mean"
+    assert inputs["output_filename"] == "candidate/output"
+
+
+def test_full_block_verification_does_not_add_calibration_exclusions(
+    runner, workflow
+):
+    workflow["11"]["inputs"]["exclude_patterns"] = "existing"
+    runner._prepare_workflow(workflow, None, None, full_blocks=True)
+    assert workflow["11"]["inputs"]["exclude_patterns"] == "existing"
+
+
+@pytest.mark.parametrize(
+    ("block", "excluded"),
+    [(block, block not in {0, 3, 14, 24, 27}) for block in range(28)],
+)
+def test_calibration_block_ranges_are_exact(runner, block, excluded):
+    import re
+
+    key = f"diffusion_model.blocks.{block}.attn.wq.lora_A.weight"
+    patterns = runner.CALIBRATION_EXCLUDE_PATTERN.splitlines()
+    assert any(re.search(pattern, key) for pattern in patterns) is excluded
+
+
+@pytest.mark.parametrize(
+    ("key", "excluded"),
+    [
+        ("diffusion_model.txtfusion.layerwise_blocks.1.attn.lora_A.weight", True),
+        ("diffusion_model.txtfusion.refiner_blocks.0.attn.lora_A.weight", True),
+        ("diffusion_model.txtfusion.layerwise_blocks.0.attn.lora_A.weight", False),
+    ],
+)
+def test_calibration_txtfusion_exclusions_are_exact(runner, key, excluded):
+    import re
+
+    patterns = runner.CALIBRATION_EXCLUDE_PATTERN.splitlines()
+    assert any(re.search(pattern, key) for pattern in patterns) is excluded
+
+
+def test_requires_filename_and_report_capture(runner, workflow):
+    del workflow["15"]
+    with pytest.raises(RuntimeError, match="output 2"):
+        runner._prepare_workflow(workflow, None, None)
+
+
+def test_parses_counterfactual_weight_sweep(runner):
+    report = """CWB COUNTERFACTUAL WEIGHT SWEEP
+median | 0.1 | 2 | 4 | true | false | 12/18 | 3 | 0.6/1 | 2.4/1
+"""
+    rows = runner.parse_counterfactual_weight_sweep(report)
+    assert rows == [{
+        "consensus_type": "median",
+        "similarity_threshold": 0.1,
+        "power_alpha": 2.0,
+        "diversity_beta": 4.0,
+        "dynamic_similarity_contrast": True,
+        "soft_comfort_bandpass": False,
+        "accepted_contributors": 12,
+        "contributors": 18,
+        "fallbacks": 3,
+        "dominant_weight_mean": 0.6,
+        "dominant_weight_max": 1.0,
+        "effective_contributors_mean": 2.4,
+        "effective_contributors_min": 1.0,
+    }]
+
+
+def test_prompt_payload_contains_only_api_workflow(runner, workflow):
+    payload = runner._prompt_payload(workflow)
+    assert payload == {"prompt": workflow}
+
+
+@pytest.mark.parametrize(
+    ("output", "expected"),
+    [
+        ({"text": ["value"]}, "value"),
+        ({"text": "value"}, "value"),
+        ({"ui": {"text": ["value"]}}, "value"),
+    ],
+)
+def test_extracts_supported_preview_history_shapes(runner, output, expected):
+    assert runner._text_output({"14": output}, "14") == expected

tests/test_cwb_delta_lora_merger.py

@@ -0,0 +1,405 @@
+import importlib
+import re
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_delta_cwb_tests"
+
+
+def _load_module(name):
+    if PACKAGE_NAME not in sys.modules:
+        package = types.ModuleType(PACKAGE_NAME)
+        package.__path__ = [str(REPO_ROOT)]
+        sys.modules[PACKAGE_NAME] = package
+    return importlib.import_module(f"{PACKAGE_NAME}.{name}")
+
+
+@pytest.fixture(scope="module")
+def delta_cwb():
+    return _load_module("nodes.cwb_delta_lora_merger")
+
+
+def _settings(delta_cwb):
+    return delta_cwb.CWBSettings(
+        consensus_type="mean",
+        alignment_method="similarity",
+        alignment_threshold=0.0,
+        similarity_threshold=-1.0,
+        power_alpha=0.0,
+        diversity_beta=0.0,
+        rescale_norm=False,
+        global_scale=1.0,
+        dynamic_similarity_contrast=False,
+        soft_comfort_bandpass=False,
+        position_weight=0.0,
+        preserve_common_prefix=False,
+    )
+
+
+def _write(path, tensors, metadata=None):
+    with IncrementalSafetensorsWriter(str(path), metadata=metadata or {}, max_workers=1) as writer:
+        writer.write_batch(list(tensors.items()))
+
+
+def _read(path):
+    tensors = {}
+    with MemoryEfficientSafeOpen(str(path), low_memory=True) as handler:
+        stream = handler.async_stream(
+            handler.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
+        )
+        for batch in stream:
+            for key, tensor in batch:
+                tensors[key] = tensor.clone()
+                handler.mark_processed(key)
+    return tensors
+
+
+def test_linear_and_convolutional_factor_expansion(delta_cwb):
+    down = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    up = torch.tensor([[2.0, -1.0], [0.5, 3.0]])
+    torch.testing.assert_close(
+        delta_cwb._factor_delta(down, up, "linear"), up @ down
+    )
+
+    conv_down = torch.arange(18, dtype=torch.float32).reshape(2, 1, 3, 3)
+    conv_up = torch.tensor([[[[2.0]], [[-1.0]]]])
+    expected = (conv_up.reshape(1, 2) @ conv_down.reshape(2, -1)).reshape(1, 1, 3, 3)
+    torch.testing.assert_close(
+        delta_cwb._factor_delta(conv_down, conv_up, "conv"), expected
+    )
+
+
+def test_dense_cwb_keeps_corresponding_rows_fixed(delta_cwb):
+    reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+    reversed_source = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
+    result = delta_cwb.merge_cwb_tensors(
+        [reference, reversed_source],
+        _settings(delta_cwb),
+        reference_index=0,
+        allow_similarity_alignment=False,
+    )
+    torch.testing.assert_close(result, torch.full((2, 2), 0.5))
+
+
+def test_layer_units_apply_missing_modes_filters_and_shape_checks(delta_cwb):
+    class Handler:
+        def __init__(self, shapes):
+            self.shapes = shapes
+
+        def get_shape(self, key):
+            return self.shapes[key]
+
+        def get_dtype(self, _key):
+            return torch.float32
+
+    roles = {"down": "layer.down", "up": "layer.up"}
+    infos = [
+        {"pairs": {"layer": roles}, "low_bit": set()},
+        {"pairs": {}, "low_bit": set()},
+    ]
+    handlers = [Handler({"layer.down": (1, 2), "layer.up": (2, 1)}), Handler({})]
+    layer_map = {"layer": [(0, "layer")]}
+    empty = delta_cwb._compile_patterns("", False)
+
+    with pytest.raises(ValueError, match="missing or incompatible"):
+        delta_cwb._build_layer_units(
+            layer_map, infos, handlers, "error", empty, empty, False
+        )
+    units, _, _ = delta_cwb._build_layer_units(
+        layer_map, infos, handlers, "zeros", empty, empty, False
+    )
+    assert units[0].zero_contributors == 1
+    assert units[0].preserve_anchor is False
+
+    exclude = delta_cwb._compile_patterns(r"\.diff$", False)
+    units, _, _ = delta_cwb._build_layer_units(
+        layer_map, infos, handlers, "zeros", exclude, empty, False
+    )
+    assert units[0].preserve_anchor is True
+    assert units[0].zero_contributors == 0
+
+    discard = delta_cwb._compile_patterns("*.diff", True)
+    units, _, rejected = delta_cwb._build_layer_units(
+        layer_map, infos, handlers, "skip", empty, discard, True
+    )
+    assert units == []
+    assert rejected == 1
+
+
+def test_streamed_merge_normalizes_alpha_once_and_supports_mixed_rank(
+    monkeypatch, tmp_path, delta_cwb
+):
+    first = tmp_path / "first.safetensors"
+    second = tmp_path / "second.safetensors"
+    _write(first, {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [0.0]]),
+        "diffusion_model.layer.alpha": torch.tensor(0.5),
+    }, {"anchor": "yes"})
+    _write(second, {
+        "diffusion_model.layer.lora_A.weight": torch.eye(2),
+        "diffusion_model.layer.lora_B.weight": torch.diag(torch.tensor([3.0, 4.0])),
+    })
+    monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
+
+    filename, report = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
+        [str(first), str(second)],
+        [first.name, second.name],
+        _settings(delta_cwb),
+        "balanced_mean",
+        False,
+        "skip",
+        "dense/mixed_rank",
+        torch.float32,
+        "full_difference",
+        384,
+        "cpu",
+        "",
+        "",
+        False,
+        False,
+    )
+    output = tmp_path / "loras" / filename
+    tensors = _read(output)
+
+    assert filename == "dense/mixed_rank.safetensors"
+    assert list(tensors) == ["diffusion_model.layer.diff"]
+    torch.testing.assert_close(
+        tensors["diffusion_model.layer.diff"],
+        torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
+    )
+    assert "Merged layers: 1" in report
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as handler:
+        metadata = handler.metadata()
+    assert metadata["anchor"] == "yes"
+    assert metadata["merge_method"] == "delta_cwb"
+    assert metadata["alpha_normalized"] == "true"
+    assert metadata["output_representation"] == "full_direct_difference"
+
+
+@pytest.mark.parametrize("patterns, glob_patterns, selected", [
+    (r"\.selected", False, True),
+    ("*.selected*", True, True),
+    ("", False, False),
+    ("no_match", False, False),
+])
+def test_include_filter_preserves_nonmatches_and_discards_first(
+    monkeypatch, tmp_path, delta_cwb, patterns, glob_patterns, selected,
+):
+    paths = [tmp_path / f"{name}.safetensors" for name in ("a", "b")]
+    for path, value in zip(paths, (2.0, 6.0)):
+        _write(path, {
+            f"diffusion_model.{layer}.diff": torch.full((2, 2), value)
+            for layer in ("selected", "other", "selected_discard")
+        })
+    monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
+    filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
+        [str(path) for path in paths], [path.name for path in paths], _settings(delta_cwb),
+        "balanced_mean", False, "skip", "include_filter", torch.float32,
+        "full_difference", 384, "cpu", patterns,
+        "*discard*" if glob_patterns else "discard", glob_patterns, False, include_mode=True,
+    )
+    tensors = _read(tmp_path / "loras" / filename)
+    assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"}
+    torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 4.0 if selected else 2.0))
+    torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0))
+
+
+def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cwb):
+    first = tmp_path / "diff_first.safetensors"
+    second = tmp_path / "diff_second.safetensors"
+    key = "diffusion_model.bias.diff"
+    _write(first, {key: torch.tensor([1.0, 3.0])})
+    _write(second, {key: torch.tensor([3.0, 5.0])})
+    monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
+
+    filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
+        [str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
+        "balanced_mean", False, "skip", "one_dim", torch.float16,
+        "full_difference", 384, "cpu",
+        "", "", False, False,
+    )
+    output = _read(tmp_path / "loras" / filename)[key]
+    assert output.dtype == torch.float32
+    torch.testing.assert_close(output, torch.tensor([2.0, 4.0]))
+
+
+def test_factorized_output_writes_compact_factors_without_dense_tensor(
+    monkeypatch, tmp_path, delta_cwb
+):
+    first = tmp_path / "factor_first.safetensors"
+    second = tmp_path / "factor_second.safetensors"
+    tensors = {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [6.0]]),
+    }
+    _write(first, tensors)
+    _write(second, tensors)
+    monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
+
+    filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
+        [str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
+        "balanced_mean", False, "skip", "factorized", torch.float32,
+        "factorized_lora", 1, "cpu", "", "", False, False,
+    )
+    output_path = tmp_path / "loras" / filename
+    output = _read(output_path)
+    assert set(output) == {
+        "diffusion_model.layer.lora_A.weight",
+        "diffusion_model.layer.lora_B.weight",
+    }
+    assert "diffusion_model.layer.diff" not in output
+    reconstructed = (
+        output["diffusion_model.layer.lora_B.weight"]
+        @ output["diffusion_model.layer.lora_A.weight"]
+    )
+    torch.testing.assert_close(reconstructed, torch.tensor([[3.0, 6.0], [6.0, 12.0]]))
+    with MemoryEfficientSafeOpen(str(output_path), low_memory=True) as handler:
+        assert handler.metadata()["output_representation"] == "factorized_lora"
+
+
+def test_atomic_writer_preserves_destination_on_cwb_failure(
+    monkeypatch, tmp_path, delta_cwb
+):
+    first = tmp_path / "atomic_first.safetensors"
+    second = tmp_path / "atomic_second.safetensors"
+    tensors = {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [0.0]]),
+    }
+    _write(first, tensors)
+    _write(second, tensors)
+    destination = tmp_path / "loras" / "atomic.safetensors"
+    destination.parent.mkdir(parents=True)
+    destination.write_bytes(b"existing destination")
+    monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
+    released = []
+    closed = []
+    opener = delta_cwb.MemoryEfficientSafeOpen
+
+    class TrackedHandler:
+        def __init__(self, path, low_memory=True):
+            self.handler = opener(path, low_memory=low_memory)
+
+        def __getattr__(self, name):
+            return getattr(self.handler, name)
+
+        def mark_processed(self, key):
+            released.append(key)
+            return self.handler.mark_processed(key)
+
+        def __exit__(self, exc_type, exc, traceback):
+            closed.append(True)
+            return self.handler.__exit__(exc_type, exc, traceback)
+
+    monkeypatch.setattr(delta_cwb, "MemoryEfficientSafeOpen", TrackedHandler)
+    monkeypatch.setattr(
+        delta_cwb,
+        "merge_cwb_tensors",
+        lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("planned failure")),
+    )
+
+    with pytest.raises(RuntimeError, match="planned failure"):
+        delta_cwb.DeltaCWBLoRAMergerLogic.execute(
+            [str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
+            "balanced_mean", False, "skip", "atomic", torch.float32,
+            "full_difference", 384, "cpu",
+            "", "", False, False,
+        )
+    assert destination.read_bytes() == b"existing destination"
+    assert not list(destination.parent.glob(".atomic.safetensors.*.tmp"))
+    assert sorted(released) == sorted([*tensors, *tensors])
+    assert len(closed) == 2
+
+
+def test_complete_layer_cuda_oom_retries_on_cpu(monkeypatch, delta_cwb):
+    source = delta_cwb.DeltaCWBLoRASource(
+        0, "layer", None, None, None, "layer.diff", torch.float32
+    )
+    unit = delta_cwb.DeltaCWBLayerUnit(
+        "layer", "layer.diff", (source,), 0, False
+    )
+    calls = []
+
+    def process(
+        _unit, _loaded, _settings, _dtype, device, _diagnostics,
+        _output_representation, _output_rank,
+    ):
+        calls.append(device.type)
+        if device.type == "cuda":
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory")
+        return torch.ones(1), torch.float32
+
+    monkeypatch.setattr(delta_cwb, "_process_layer_on_device", process)
+    monkeypatch.setattr(delta_cwb, "_release_cuda_oom", lambda: None)
+    output, dtype, fallback = delta_cwb._process_layer(
+        unit, {}, _settings(delta_cwb), torch.float32, "cuda", delta_cwb.CWBDiagnostics()
+    )
+    assert calls == ["cuda", "cpu"]
+    assert fallback is True
+    assert dtype == torch.float32
+    assert torch.equal(output, torch.ones(1))
+
+
+def test_exact_node_schema_contract(delta_cwb):
+    expected_ids = [
+        "DeltaCWBLoRATwoMerger",
+        "DeltaCWBLoRAThreeMerger",
+        "DeltaCWBLoRAMultiMerger",
+    ]
+    schemas = [node.define_schema() for node in delta_cwb.DELTA_CWB_LORA_NODES]
+    assert [schema.node_id for schema in schemas] == expected_ids
+    assert [schema.display_name for schema in schemas] == [
+        "Delta CWB Merge LoRAs (2)",
+        "Delta CWB Merge LoRAs (3)",
+        "Delta CWB LoRA Multi-Merge",
+    ]
+    assert all(schema.category == "ModelUtils/LoRA/Merge/Delta CWB" for schema in schemas)
+    assert all([output.display_name for output in schema.outputs] == ["output_filename", "cwb_report"] for schema in schemas)
+    assert all(schema.outputs[0].get_io_type() == "*" for schema in schemas)
+
+    fixed_ids = [value.id for value in schemas[0].inputs]
+    assert fixed_ids == [
+        "lora_1", "lora_2", "cwb_preset", "cwb_config", "mismatch_mode",
+        "output_filename", "save_dtype", "output_representation", "output_rank",
+        "process_device", "exclude_patterns",
+        "discard_patterns", "glob_patterns", "force_clear_cache", "include_mode",
+    ]
+    multi_ids = [value.id for value in schemas[2].inputs]
+    assert multi_ids[:9] == ["lora_count", *[f"lora_{index}" for index in range(1, 9)]]
+    assert schemas[0].inputs[2].default == "balanced_mean"
+    assert schemas[0].inputs[4].default == "skip"
+    assert schemas[0].inputs[6].default == "bf16"
+    assert schemas[0].inputs[7].default == "full_difference"
+    assert schemas[0].inputs[8].default == 384
+    assert schemas[0].inputs[9].default == "cuda"
+    assert schemas[0].inputs[-2].default is False
+    assert schemas[0].inputs[-1].default is False
+    assert "agreement" in schemas[0].inputs[2].tooltip
+    assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[9].tooltip
+
+
+def test_invalid_regex_and_unselected_multi_input_fail(delta_cwb):
+    with pytest.raises(re.error):
+        delta_cwb._compile_patterns("[", False)
+    with pytest.raises(ValueError, match="unselected input.*3"):
+        delta_cwb.DeltaCWBLoRAMultiMerger.execute(
+            lora_count="3", lora_1="a", lora_2="b", lora_3="None"
+        )

tests/test_cwb_model_graft.py

@@ -0,0 +1,182 @@
+import importlib
+from pathlib import Path
+import sys
+import types
+
+import pytest
+from safetensors.torch import load_file, save_file
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+@pytest.fixture(scope="module")
+def graft_mod():
+    package_name = "modelutils_cwb_graft_tests"
+    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}.nodes.cwb_model_graft")
+
+
+def _patch_io(monkeypatch, graft_mod, tmp_path, paths):
+    monkeypatch.setattr(
+        graft_mod.folder_paths,
+        "get_full_path",
+        lambda _kind, name: paths.get(name),
+    )
+    monkeypatch.setattr(
+        graft_mod.folder_paths,
+        "get_folder_paths",
+        lambda _kind: [str(tmp_path)],
+    )
+    monkeypatch.setattr(graft_mod.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(graft_mod, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(graft_mod, "cleanup_after_operation", lambda: None)
+
+
+def test_cwb_model_graft_schema_and_documentation(graft_mod):
+    schema = graft_mod.CWBModelGraft.define_schema()
+    assert schema.node_id == "CWBModelGraft"
+    assert schema.display_name == "CWB Model Graft"
+    ids = [inp.id for inp in schema.inputs]
+    assert "attention_preset" in ids
+    assert "ffn_preset" in ids
+    assert "protect_early_blocks" in ids
+    assert "attention_block_rule" in ids
+    assert "ffn_block_rule" in ids
+    assert "attention_subpart" in ids
+    assert "ffn_subpart" in ids
+    assert "custom_key_map" in ids
+    assert ids[-1] == "include_mode"
+
+    doc_out = graft_mod.CWBModelGraft.execute(execution_mode="DOCUMENTATION ONLY")
+    filename, docs, report = doc_out.result
+    assert filename == "Documentation mode active. No merge performed."
+    assert "# CWB Model Graft" in docs
+    assert "Complete Input Reference" in docs
+    assert "no merge was performed" in report
+
+
+def test_cwb_model_graft_protected_blocks_subparts_and_mismatched_shapes(
+    monkeypatch, tmp_path, graft_mod
+):
+    base_path = tmp_path / "base.safetensors"
+    donor_path = tmp_path / "donor.safetensors"
+
+    base_qkv_0 = torch.ones((6, 4), dtype=torch.float32) * 2.0
+    base_qkv_2 = torch.cat(
+        [
+            torch.ones((2, 4), dtype=torch.float32) * 1.0,
+            torch.ones((2, 4), dtype=torch.float32) * 2.0,
+            torch.ones((2, 4), dtype=torch.float32) * 3.0,
+        ],
+        dim=0,
+    )
+    base_ffn_2 = torch.cat(
+        [
+            torch.ones((2, 4), dtype=torch.float32) * 4.0,
+            torch.ones((2, 4), dtype=torch.float32) * 5.0,
+        ],
+        dim=0,
+    )
+    base_custom = torch.ones((2, 4), dtype=torch.float32) * 2.0
+
+    save_file(
+        {
+            "layers.0.proj.qkv.weight": base_qkv_0,
+            "layers.2.proj.qkv.weight": base_qkv_2,
+            "layers.2.proj.ffn_up_gate.weight": base_ffn_2,
+            "layers.3.proj.qkv.weight": base_qkv_0.clone(),
+            "custom_target.weight": base_custom,
+        },
+        str(base_path),
+    )
+
+    save_file(
+        {
+            "layers.0.proj.qkv.weight": torch.ones((6, 4), dtype=torch.float32) * 99.0,
+            "layers.2.proj.q.weight": torch.ones((2, 4), dtype=torch.float32) * 50.0,
+            "layers.2.proj.k.weight": torch.ones((2, 4), dtype=torch.float32) * 60.0,
+            "layers.2.proj.v.weight": torch.tensor(
+                [[9.0, 1.0, 1.0, 1.0], [1.0, 9.0, 1.0, 1.0]], dtype=torch.float32
+            ),
+            "layers.2.proj.ffn_up_gate.weight": torch.cat(
+                [
+                    torch.tensor(
+                        [[8.0, 1.0, 1.0, 1.0], [1.0, 8.0, 1.0, 1.0]],
+                        dtype=torch.float32,
+                    ),
+                    torch.ones((2, 4), dtype=torch.float32) * 77.0,
+                ],
+                dim=0,
+            ),
+            "layers.3.proj.qkv.weight": torch.ones((6, 4), dtype=torch.float32) * 88.0,
+            "custom_donor.weight": torch.tensor(
+                [
+                    [10.0, 2.0, 1.0, 1.0, 1.0, 1.0],
+                    [1.0, 10.0, 2.0, 1.0, 1.0, 1.0],
+                    [1.0, 1.0, 10.0, 2.0, 1.0, 1.0],
+                ],
+                dtype=torch.float32,
+            ),
+        },
+        str(donor_path),
+    )
+
+    paths = {"base": str(base_path), "donor": str(donor_path)}
+    _patch_io(monkeypatch, graft_mod, tmp_path, paths)
+
+    out = graft_mod.CWBModelGraft.execute(
+        execution_mode="MERGE",
+        model_a="base",
+        model_b="donor",
+        attention_preset="balanced_mean",
+        attention_config=None,
+        ffn_preset="balanced_mean",
+        ffn_config=None,
+        protect_early_blocks=2,
+        attention_block_rule="interleaved_even",
+        ffn_block_rule="all_after_protected",
+        attention_subpart="value_and_output",
+        ffn_subpart="up_only",
+        custom_key_map="custom_donor.weight -> custom_target.weight",
+        output_filename="grafted",
+        save_dtype="fp32",
+        process_device="cpu",
+        exclude_patterns="",
+        discard_patterns="",
+        glob_patterns=False,
+        lazy_load=True,
+        force_clear_cache=False,
+        override_dtype=False,
+        include_mode=False,
+    )
+
+    rel_name, docs, report = out.result
+    assert "# CWB Model Graft" in docs
+    assert "attn=balanced_mean, ffn=balanced_mean" in report
+    assert "Grafted tensors: 3" in report
+    tensors = load_file(str(tmp_path / "diffusion_models" / rel_name))
+
+    # 1. Protected block 0 is byte-exact identical to Base
+    assert torch.equal(tensors["layers.0.proj.qkv.weight"], base_qkv_0)
+    # 2. Odd block 3 attention is untouched under interleaved_even
+    assert torch.equal(tensors["layers.3.proj.qkv.weight"], base_qkv_0)
+    # 3. Even block 2 qkv: Q and K slices stay Base, only V slice is grafted from split donor
+    q_out, k_out, v_out = tensors["layers.2.proj.qkv.weight"].chunk(3, dim=0)
+    assert torch.equal(q_out, base_qkv_2[:2])
+    assert torch.equal(k_out, base_qkv_2[2:4])
+    assert not torch.equal(v_out, base_qkv_2[4:6])
+    # 4. Block 2 ffn_up_gate: up half is grafted, gate half stays Base
+    up_out, gate_out = tensors["layers.2.proj.ffn_up_gate.weight"].chunk(2, dim=0)
+    assert not torch.equal(up_out, base_ffn_2[:2])
+    assert torch.equal(gate_out, base_ffn_2[2:])
+    # 5. Mismatched-shape custom mapped tensor blends without changing shape or norm
+    assert tensors["custom_target.weight"].shape == base_custom.shape
+    assert not torch.equal(tensors["custom_target.weight"], base_custom)
+    assert torch.linalg.vector_norm(tensors["custom_target.weight"]).item() == pytest.approx(
+        torch.linalg.vector_norm(base_custom).item(), rel=1e-5
+    )

tests/test_cwb_preset_sweep.py

@@ -0,0 +1,120 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_sweep():
+    spec = importlib.util.spec_from_file_location(
+        "cwb_preset_sweep_test", REPO_ROOT / "run_cwb_preset_sweep.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture(scope="module")
+def sweep():
+    return _load_sweep()
+
+
+@pytest.fixture
+def baseline():
+    return {
+        "consensus_type": "median",
+        "alignment_method": "similarity",
+        "alignment_threshold": 0.0,
+        "similarity_threshold": 0.0,
+        "power_alpha": 2.0,
+        "diversity_beta": 10.0,
+        "rescale_norm": True,
+        "global_scale": 1.0,
+        "dynamic_similarity_contrast": False,
+        "soft_comfort_bandpass": True,
+        "position_weight": 0.05,
+        "preserve_common_prefix": False,
+    }
+
+
+def test_cases_are_complete_and_change_exactly_one_factor(sweep, baseline):
+    cases = sweep.build_sweep_cases(baseline)
+    assert cases[0]["id"] == "baseline"
+    assert len({case["id"] for case in cases}) == len(cases)
+    for case in cases:
+        assert tuple(case["config"]) == sweep.CONFIG_KEYS
+        changed = [
+            key for key in sweep.CONFIG_KEYS
+            if case["config"][key] != baseline[key]
+        ]
+        assert changed == ([] if case["id"] == "baseline" else [case["factor"]])
+
+
+def test_sweep_exposes_expected_full_block_finalist_ids(sweep, baseline):
+    ids = {case["id"] for case in sweep.build_sweep_cases(baseline)}
+    assert {
+        "baseline",
+        "alignment_method_index",
+        "alignment_threshold_0p0005",
+        "alignment_threshold_0p0025",
+        "consensus_type_mean",
+        "diversity_beta_4p0",
+        "rescale_norm_false",
+        "dynamic_similarity_contrast_true",
+        "soft_comfort_bandpass_false",
+    } <= ids
+
+
+def test_report_parser_extracts_influence_metrics(sweep):
+    report = """CWB MERGE SUMMARY
+Alignment candidates: 100
+Alignment matches: 75
+Anchor-only groups: 5
+Weighted vector groups: 20
+All-rejected equal fallbacks: 1
+Zero-weight equal fallbacks: 2
+Accepted weighting contributors: 50/80
+Consensus similarity min/mean/max: -0.5 / 0.25 / 1
+Dominant normalized weight min/mean/max: 0.2 / 0.4 / 0.8
+Effective contributors min/mean/max: 1.2 / 2.5 / 5
+CUDA OOM CPU fallbacks: 0
+LoRA groups: 10
+Rank-one delta norm ratio count/min/mean/max: 3 / 0.7 / 1.1 / 1.4
+Rank-one delta norm ratio count/min/mean/max: 1 / 0.8 / 0.9 / 1.2
+"""
+    metrics = sweep.parse_cwb_report(report)
+    assert metrics["alignment_match_fraction"] == pytest.approx(0.75)
+    assert metrics["accepted_weight_fraction"] == pytest.approx(0.625)
+    assert metrics["dominant_weight_mean"] == pytest.approx(0.4)
+    assert metrics["effective_contributors_mean"] == pytest.approx(2.5)
+    assert metrics["norm_ratio_min"] == pytest.approx(0.7)
+    assert metrics["norm_ratio_count"] == 4
+    assert metrics["norm_ratio_mean"] == pytest.approx(1.05)
+    assert metrics["norm_ratio_max"] == pytest.approx(1.4)
+
+
+def test_report_parser_rejects_pre_diagnostic_server(sweep):
+    with pytest.raises(RuntimeError, match="Restart ComfyUI"):
+        sweep.parse_cwb_report("CWB MERGE SUMMARY\nAlignment matches: 1")
+
+
+def test_analysis_parser_extracts_global_metrics(sweep):
+    comparison = """MODEL COMPARISON SUMMARY
+PARAMETER-WEIGHTED GLOBAL METRICS
+relative_l2: 1.04
+cosine: 0.46
+exact: 0.01
+norm_ratio: 0.89
+EQUAL-LAYER AVERAGES
+"""
+    cwb_report = """GLOBAL CWB-DERIVED DIAGNOSTICS
+Pairwise row cosine mean/min/max: 0.4 / -0.1 / 1
+Mean-consensus affinity A/B: 0.2 / 0.3
+Median-consensus affinity A/B: 0.25 / 0.35
+"""
+    metrics = sweep.parse_analysis_reports(comparison, cwb_report)
+    assert metrics["analysis_relative_l2"] == pytest.approx(1.04)
+    assert metrics["analysis_pair_cosine"] == pytest.approx(0.4)
+    assert metrics["analysis_median_affinity_b"] == pytest.approx(0.35)

tests/test_dtype_conversion.py

@@ -0,0 +1,269 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import (
+    IncrementalSafetensorsWriter,
+    MemoryEfficientSafeOpen,
+)
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_dtype_conversion_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 conversion():
+    return _load_module("nodes.dtype_conversion")
+
+
+def _write_uel(path, tensors, metadata=None):
+    with IncrementalSafetensorsWriter(
+        str(path), metadata=metadata or {}, max_workers=1
+    ) as writer:
+        writer.write_batch(list(tensors.items()))
+
+
+def _read_uel(path):
+    result = {}
+    with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
+        metadata = loader.metadata()
+        for batch in loader.async_stream(
+            loader.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
+        ):
+            key, tensor = batch[0]
+            result[key] = tensor.clone()
+            loader.mark_processed(key)
+    return result, metadata
+
+
+def _patch_runtime(monkeypatch, tmp_path, conversion, source):
+    monkeypatch.setattr(conversion.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(
+        conversion.folder_paths,
+        "get_full_path_or_raise",
+        lambda model_type, name: str(source),
+    )
+    monkeypatch.setattr(conversion, "cleanup_after_operation", lambda: None)
+
+
+def test_schema_is_diffusion_model_scoped(conversion):
+    schema = conversion.DiffusionModelDtypeConversion.define_schema()
+    assert schema.node_id == "DiffusionModelDtypeConversion"
+    assert schema.category == "ModelUtils/Conversion"
+    assert [item.id for item in schema.inputs] == [
+        "model_name",
+        "target_dtype",
+        "reference_model",
+        "exclude_patterns",
+        "output_filename",
+        "include_mode",
+    ]
+    assert schema.inputs[-1].default is False
+    assert schema.inputs[1].options == ["fp32", "fp16", "bf16"]
+    assert [item.display_name for item in schema.outputs] == [
+        "output_path",
+        "conversion_report",
+    ]
+
+
+def test_reference_model_dtypes_override_target_without_tensor_loading(
+    monkeypatch, tmp_path, conversion
+):
+    source = tmp_path / "source.safetensors"
+    reference = tmp_path / "reference.safetensors"
+    _write_uel(source, {
+        "mapped_bf16": torch.ones(2, dtype=torch.float32),
+        "mapped_fp32": torch.ones(2, dtype=torch.float16),
+        "fallback": torch.ones(2, dtype=torch.float32),
+        "nonfloating": torch.tensor([1], dtype=torch.int64),
+    })
+    _write_uel(reference, {
+        "mapped_bf16": torch.ones(2, dtype=torch.bfloat16),
+        "mapped_fp32": torch.ones(2, dtype=torch.float32),
+        "nonfloating": torch.tensor([1], dtype=torch.int64),
+    })
+    monkeypatch.setattr(conversion.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(
+        conversion.folder_paths,
+        "get_full_path_or_raise",
+        lambda _, name: {"source": str(source), "reference": str(reference)}[name],
+    )
+    monkeypatch.setattr(conversion, "cleanup_after_operation", lambda: None)
+    monkeypatch.setattr(
+        conversion.MemoryEfficientSafeOpen,
+        "get_tensor",
+        lambda *args, **kwargs: pytest.fail("synchronous reference loading was used"),
+    )
+
+    output_path, report = conversion.convert_diffusion_model_dtype(
+        "source", "fp16", "", "reference_mapped", "reference"
+    )
+    result, _ = _read_uel(tmp_path / "diffusion_models" / output_path)
+    assert result["mapped_bf16"].dtype == torch.bfloat16
+    assert result["mapped_fp32"].dtype == torch.float32
+    assert result["fallback"].dtype == torch.float16
+    assert result["nonfloating"].dtype == torch.int64
+    assert "Reference-mapped floating tensors: 2" in report
+    assert "Reference-missing fallback tensors: 1" in report
+
+
+@pytest.mark.parametrize(
+    ("target_name", "expected_dtype"),
+    [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
+)
+def test_converts_floating_and_preserves_excluded_and_nonfloating(
+    monkeypatch, tmp_path, conversion, target_name, expected_dtype
+):
+    source = tmp_path / f"source_{target_name}.safetensors"
+    tensors = {
+        "blocks.0.weight": torch.arange(4, dtype=torch.float32),
+        "blocks.0.norm.weight": torch.ones(3, dtype=torch.float32),
+        "position_ids": torch.tensor([1, 2], dtype=torch.int64),
+    }
+    _write_uel(source, tensors, metadata={"source": "test"})
+    _patch_runtime(monkeypatch, tmp_path, conversion, source)
+
+    output_path, report = conversion.convert_diffusion_model_dtype(
+        source.name,
+        target_name,
+        r"\.norm\.weight$",
+        f"converted_{target_name}",
+    )
+    assert output_path == f"converted_{target_name}.safetensors"
+    result, metadata = _read_uel(
+        tmp_path / "diffusion_models" / output_path
+    )
+
+    assert result["blocks.0.weight"].dtype == expected_dtype
+    assert result["blocks.0.norm.weight"].dtype == torch.float32
+    assert result["position_ids"].dtype == torch.int64
+    assert metadata == {"source": "test"}
+    assert "Excluded tensors preserved: 1" in report
+    assert "Non-floating tensors preserved: 1" in report
+
+
+@pytest.mark.parametrize("patterns, selected", [
+    (r"^blocks\.0\.weight$", {"blocks.0.weight"}),
+    (r"^blocks\.0\.weight$" + "\n" + r"^blocks\.0\.norm\.weight$", {"blocks.0.weight", "blocks.0.norm.weight"}),
+    ("", set()),
+    ("no_such_layer", set()),
+])
+@pytest.mark.parametrize("use_reference", [False, True])
+def test_include_filter_only_converts_matches(
+    monkeypatch, tmp_path, conversion, patterns, selected, use_reference,
+):
+    source = tmp_path / "source.safetensors"
+    reference = tmp_path / "reference.safetensors"
+    tensors = {
+        "blocks.0.weight": torch.arange(4, dtype=torch.float32),
+        "blocks.0.norm.weight": torch.ones(3, dtype=torch.float32),
+        "position_ids": torch.tensor([1, 2], dtype=torch.int64),
+    }
+    _write_uel(source, tensors)
+    _write_uel(reference, {"blocks.0.weight": tensors["blocks.0.weight"].to(torch.bfloat16)})
+    _patch_runtime(monkeypatch, tmp_path, conversion, source)
+    monkeypatch.setattr(
+        conversion.folder_paths, "get_full_path_or_raise",
+        lambda _, name: str(reference if name == "reference" else source),
+    )
+
+    output_name, _ = conversion.DiffusionModelDtypeConversion.execute(
+        source.name, "fp16", "reference" if use_reference else "None",
+        patterns, "include_converted", include_mode=True,
+    )
+    result, _ = _read_uel(tmp_path / "diffusion_models" / output_name)
+    for key, tensor in tensors.items():
+        expected_dtype = tensor.dtype
+        if key in selected:
+            expected_dtype = torch.bfloat16 if use_reference and key == "blocks.0.weight" else torch.float16
+        assert result[key].dtype == expected_dtype
+        torch.testing.assert_close(result[key].to(tensor.dtype), tensor)
+
+
+def test_uses_bounded_async_stream_and_releases_every_tensor(
+    monkeypatch, tmp_path, conversion
+):
+    source = tmp_path / "source_async.safetensors"
+    _write_uel(
+        source,
+        {
+            "a": torch.ones(2, dtype=torch.float32),
+            "b": torch.ones(2, dtype=torch.float16),
+            "c": torch.ones(2, dtype=torch.int32),
+        },
+    )
+    _patch_runtime(monkeypatch, tmp_path, conversion, source)
+    async_calls = []
+    processed = []
+    original_async = conversion.MemoryEfficientSafeOpen.async_stream
+    original_mark = conversion.MemoryEfficientSafeOpen.mark_processed
+
+    def track_async(loader, keys, **kwargs):
+        async_calls.append((tuple(keys), kwargs))
+        yield from original_async(loader, keys, **kwargs)
+
+    def track_mark(loader, key):
+        processed.append(key)
+        return original_mark(loader, key)
+
+    monkeypatch.setattr(
+        conversion.MemoryEfficientSafeOpen, "async_stream", track_async
+    )
+    monkeypatch.setattr(
+        conversion.MemoryEfficientSafeOpen, "mark_processed", track_mark
+    )
+    monkeypatch.setattr(
+        conversion.MemoryEfficientSafeOpen,
+        "get_tensor",
+        lambda *args, **kwargs: pytest.fail("synchronous loading was used"),
+    )
+
+    conversion.convert_diffusion_model_dtype(
+        source.name, "bf16", "^b$", "converted_async"
+    )
+
+    assert len(async_calls) == 1
+    assert async_calls[0][1] == {
+        "batch_size": 1,
+        "prefetch_batches": 1,
+        "pin_memory": False,
+    }
+    assert processed == ["a", "b", "c"]
+
+
+def test_invalid_regex_does_not_create_output(monkeypatch, tmp_path, conversion):
+    source = tmp_path / "source_invalid.safetensors"
+    _write_uel(source, {"weight": torch.ones(1)})
+    _patch_runtime(monkeypatch, tmp_path, conversion, source)
+
+    with pytest.raises(ValueError, match="Invalid exclusion regex"):
+        conversion.convert_diffusion_model_dtype(
+            source.name, "fp16", "[", "must_not_exist"
+        )
+
+    assert not (tmp_path / "diffusion_models" / "must_not_exist.safetensors").exists()
+
+
+def test_rejects_input_output_collision(monkeypatch, tmp_path, conversion):
+    model_dir = tmp_path / "diffusion_models"
+    model_dir.mkdir()
+    source = model_dir / "same.safetensors"
+    _write_uel(source, {"weight": torch.ones(1)})
+    _patch_runtime(monkeypatch, tmp_path, conversion, source)
+
+    with pytest.raises(ValueError, match="must differ"):
+        conversion.convert_diffusion_model_dtype(
+            source.name, "fp16", "", "same"
+        )

tests/test_groups.toml

@@ -0,0 +1,80 @@
+[groups.repository_contract]
+paths = ["__init__.py", "nodes/*.py"]
+python_tests = [
+  "tests/test_conditional_registration.py",
+  "tests/test_schema_tooltips.py",
+  "tests/test_uel_repository_contract.py",
+]
+
+[groups.artifact_paths]
+paths = [
+  "nodes/artifact_paths.py", "nodes/dora_extract_wd.py",
+  "nodes/dora_learned_wd.py", "nodes/dtype_conversion.py",
+  "nodes/cwb_delta_lora_merger.py", "nodes/cwb_model_graft.py", "nodes/lodestone_merger.py", "nodes/lora_extract_svd.py", "nodes/lora_merger.py",
+  "nodes/lora_rename.py", "nodes/lora_resize.py", "nodes/merger.py",
+  "nodes/minimax_h3_lora_convert.py",
+  "nodes/minimax_h3_fold.py",
+  "nodes/prunekeys.py", "nodes/renamekeys.py", "nodes/text_encoder_extract.py",
+]
+python_tests = ["tests/test_artifact_paths.py"]
+
+[groups.cwb]
+paths = ["nodes/consensus_merger.py", "nodes/cwb_model_graft.py", "nodes/cwb_delta_lora_merger.py", "run_cwb_api_workflow.py", "run_cwb_preset_sweep.py"]
+python_tests = [
+  "tests/test_consensus_merger.py", "tests/test_cwb_model_graft.py", "tests/test_cwb_api_runner.py",
+  "tests/test_cwb_delta_lora_merger.py", "tests/test_cwb_preset_sweep.py",
+]
+
+[groups.model_analysis]
+paths = ["nodes/model_analysis.py", "run_model_analysis_api_workflow.py"]
+python_tests = ["tests/test_model_analysis.py", "tests/test_model_analysis_api_runner.py"]
+
+[groups.lora_model_analysis]
+paths = ["nodes/lora_model_analysis.py", "nodes/lora_resize.py", "nodes/model_analysis.py"]
+python_tests = ["tests/test_lora_model_analysis.py"]
+
+[groups.adaptive_extraction]
+paths = [
+  "nodes/adaptive_svd.py", "nodes/dora_extract_wd.py", "nodes/dora_learned_wd.py",
+  "nodes/extraction_stream.py", "nodes/lora_extract_svd.py",
+  "nodes/text_encoder_extract.py",
+]
+python_tests = ["tests/test_adaptive_extraction_pipeline.py", "tests/test_knee_partial_svd.py"]
+
+[groups.lora_merge]
+paths = [
+  "nodes/cwb_delta_lora_merger.py", "nodes/device_utils.py", "nodes/lodestone_merger.py", "nodes/lora_alpha.py",
+  "nodes/lora_merger.py", "nodes/quantization_guard.py", "nodes/uel_io.py",
+]
+python_tests = [
+  "tests/test_lora_diff_dtype_contract.py", "tests/test_lora_merger_async.py",
+  "tests/test_cwb_delta_lora_merger.py", "tests/test_lodestone_merger.py", "tests/test_lora_quantization_guard.py",
+]
+
+[groups.lora_resize]
+paths = ["nodes/lora_resize.py"]
+python_tests = ["tests/test_lora_resize_factor_space.py", "tests/test_lora_resize_min_rank.py"]
+
+[groups.merger_filters]
+paths = ["nodes/merger.py", "nodes/prunekeys.py"]
+python_tests = ["tests/test_merger_filters.py"]
+
+[groups.layer_parameters]
+paths = ["nodes/layer_parameters.py", "nodes/merger.py", "nodes/lora_extract_svd.py", "nodes/dora_extract_wd.py", "nodes/dora_learned_wd.py", "nodes/text_encoder_extract.py", "nodes/lora_resize.py", "__init__.py"]
+python_tests = ["tests/test_layer_parameters.py", "tests/test_layer_parameter_receivers.py", "tests/test_layer_parameter_extraction.py", "tests/test_layer_parameter_resize.py", "tests/test_merger_filters.py"]
+
+[groups.layer_parameter_syntax]
+paths = ["nodes/layer_parameters.py"]
+python_tests = ["tests/test_layer_parameters.py"]
+
+[groups.dtype_conversion]
+paths = ["nodes/dtype_conversion.py"]
+python_tests = ["tests/test_dtype_conversion.py"]
+
+[groups.minimax_h3_lora_convert]
+paths = ["nodes/minimax_h3_lora_convert.py"]
+python_tests = ["tests/test_minimax_h3_lora_convert.py"]
+
+[groups.minimax_h3_fold]
+paths = ["nodes/minimax_h3_fold.py"]
+python_tests = ["tests/test_minimax_h3_fold.py"]

tests/test_knee_partial_svd.py

@@ -0,0 +1,101 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_knee_partial_svd_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", params=["nodes.lora_extract_svd", "nodes.dora_extract_wd"])
+def extraction(request):
+    return _load_module(request.param)
+
+
+def _fake_factors(weight, q, niter):
+    del niter
+    return (
+        torch.eye(weight.shape[0], q),
+        torch.linspace(float(q), 1.0, q),
+        torch.eye(weight.shape[1], q),
+    )
+
+
+def test_knee_probe_uses_offset_and_expands_at_boundary(monkeypatch, extraction):
+    probes = []
+    detected = iter([9, 5])
+
+    def fake_svd(weight, q, niter):
+        probes.append(q)
+        return _fake_factors(weight, q, niter)
+
+    monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd)
+    monkeypatch.setattr(
+        extraction, "_compute_rank", lambda singular_values, mode, param, cap: next(detected)
+    )
+
+    _, _, _, rank = extraction._svd_extract_knee_lowrank(
+        torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2
+    )
+
+    assert probes == [10, 12]
+    assert rank == 5
+
+
+def test_knee_probe_stops_when_knee_is_clear(monkeypatch, extraction):
+    probes = []
+
+    def fake_svd(weight, q, niter):
+        probes.append(q)
+        return _fake_factors(weight, q, niter)
+
+    monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd)
+    monkeypatch.setattr(extraction, "_compute_rank", lambda *args: 3)
+
+    _, _, _, rank = extraction._svd_extract_knee_lowrank(
+        torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2
+    )
+
+    assert probes == [10]
+    assert rank == 3
+
+
+def test_lowrank_factor_construction_does_not_build_residual(monkeypatch, extraction):
+    monkeypatch.setattr(extraction.torch, "svd_lowrank", _fake_factors)
+
+    result, mode = extraction._svd_extract_linear_lowrank(
+        torch.eye(6), rank=3, device="cpu", clamp_quantile=1.0, niter=2
+    )
+    down, up, residual = result
+
+    assert mode == "low rank"
+    assert residual is None
+    assert torch.equal(up, torch.eye(6, 3) * torch.tensor([3.0, 2.0, 1.0]))
+    assert torch.equal(down, torch.eye(3, 6))
+
+
+def test_knee_control_is_adjacent_to_method():
+    modules_and_nodes = [
+        (_load_module("nodes.lora_extract_svd"), "LoRAExtractKnee"),
+        (_load_module("nodes.dora_extract_wd"), "DoRAExtractKnee"),
+        (_load_module("nodes.dora_learned_wd"), "DoRALearnedExtractKnee"),
+        (_load_module("nodes.text_encoder_extract"), "TextEncoderLoRAExtractKnee"),
+        (_load_module("nodes.text_encoder_extract"), "TextEncoderDoRAExtractKnee"),
+    ]
+
+    for module, node_name in modules_and_nodes:
+        names = [item.id for item in getattr(module, node_name).define_schema().inputs]
+        method_index = names.index("knee_method")
+        assert names[method_index + 1] == "knee_probe_offset"

tests/test_layer_parameter_extraction.py

@@ -0,0 +1,194 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_layer_parameter_extraction_tests"
+
+
+def _load(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}")
+
+
+def _write(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
+        writer.write_dict(tensors)
+
+
+def test_fixed_extraction_rule_overrides_only_matching_layer(monkeypatch, tmp_path):
+    extraction = _load("nodes.lora_extract_svd")
+    parameters = _load("nodes.layer_parameters")
+    first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
+    match = "blocks.45.attn.qkv_proj.weight"
+    other = "blocks.46.attn.qkv_proj.weight"
+    _write(first, {match: torch.eye(3), other: torch.eye(3)})
+    _write(second, {match: torch.zeros(3, 3), other: torch.zeros(3, 3)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    rules = parameters.parse_rules(r"(blocks\.45\.attn\.qkv_proj) linear_dim:2 clamp_quantile:1")
+
+    extraction.extract_lora_from_files(
+        str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+        layer_parameters=rules,
+    )
+
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        assert result.get_shape("diffusion_model.blocks.45.attn.qkv_proj.lora_A.weight")[0] == 2
+        assert result.get_shape("diffusion_model.blocks.46.attn.qkv_proj.lora_A.weight")[0] == 1
+
+
+def test_unmatched_fixed_rule_preserves_existing_rank_cap(monkeypatch, tmp_path):
+    extraction = _load("nodes.lora_extract_svd")
+    first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
+    key = "blocks.45.attn.qkv_proj.weight"
+    _write(first, {key: torch.eye(3)})
+    _write(second, {key: torch.zeros(3, 3)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    caps = []
+
+    def extract_linear(weight, mode, parameter, device, max_rank, clamp, niter, probe):
+        caps.append((parameter, max_rank))
+        return (torch.ones((1, weight.shape[1])), torch.ones((weight.shape[0], 1)), 1), "lora"
+
+    monkeypatch.setattr(extraction, "_svd_extract_linear", extract_linear)
+    extraction.extract_lora_from_files(
+        str(first), str(second), "fixed", 2, 2, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+    )
+    assert caps == [(2, 1)]
+
+
+def test_adaptive_rule_and_cpu_retry_use_resolved_values(monkeypatch, tmp_path):
+    extraction = _load("nodes.lora_extract_svd")
+    parameters = _load("nodes.layer_parameters")
+    first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
+    key = "blocks.45.attn.qkv_proj.weight"
+    _write(first, {key: torch.eye(3)})
+    _write(second, {key: torch.zeros(3, 3)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    calls = []
+
+    def retry(cuda_call, cpu_call, _device):
+        return cpu_call()
+
+    def extract_linear(weight, mode, parameter, device, max_rank, clamp, niter, probe):
+        calls.append((mode, parameter, device, max_rank, clamp, niter, probe))
+        return (torch.ones((1, weight.shape[1])), torch.ones((weight.shape[0], 1)), 1), "lora"
+
+    monkeypatch.setattr(extraction, "retry_cuda_oom_on_cpu", retry)
+    monkeypatch.setattr(extraction, "_svd_extract_linear", extract_linear)
+    rules = parameters.parse_rules(r"(blocks\.45) linear_ratio:3 linear_max_rank:2 clamp_quantile:1")
+    extraction.extract_lora_from_files(
+        str(first), str(second), "ratio", 2, 2, "cuda", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+        layer_parameters=rules,
+    )
+    assert calls == [("ratio", 3, "cpu", 2, 1, 2, 32)]
+
+
+def test_rule_min_diff_gates_matching_layer_before_svd(monkeypatch, tmp_path):
+    extraction = _load("nodes.lora_extract_svd")
+    parameters = _load("nodes.layer_parameters")
+    first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
+    key = "blocks.45.attn.qkv_proj.weight"
+    _write(first, {key: torch.eye(2) * 0.5})
+    _write(second, {key: torch.zeros(2, 2)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    calls = []
+    monkeypatch.setattr(extraction, "_svd_extract_linear", lambda *args: calls.append(args))
+    rules = parameters.parse_rules(r"(blocks\.45) min_diff:1")
+    extraction.extract_lora_from_files(
+        str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+        layer_parameters=rules,
+    )
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        assert list(result.keys()) == []
+    assert calls == []
+
+
+def test_chunked_fallback_receives_resolved_clamp(monkeypatch, tmp_path):
+    extraction = _load("nodes.lora_extract_svd")
+    parameters = _load("nodes.layer_parameters")
+    first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
+    key = "blocks.45.attn.qkv_proj.weight"
+    _write(first, {key: torch.eye(6, 2)})
+    _write(second, {key: torch.zeros(6, 2)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    monkeypatch.setattr(extraction, "_svd_extract_linear", lambda *args: (_ for _ in ()).throw(RuntimeError("force chunk")))
+    monkeypatch.setattr(extraction, "_detect_fused_layer", lambda *_: 3)
+    observed = []
+
+    def chunk(weight, chunks, mode, parameter, device, rank, probe, clamp):
+        observed.append((mode, parameter, rank, probe, clamp))
+        return torch.ones((weight.shape[0], 1)), torch.ones((1, weight.shape[1])), 1
+
+    monkeypatch.setattr(extraction, "_extract_chunked_layer", chunk)
+    rules = parameters.parse_rules(r"(blocks\.45) linear_dim:2 clamp_quantile:1")
+    extraction.extract_lora_from_files(
+        str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+        layer_parameters=rules,
+    )
+    assert observed == [("fixed", 2, 2, 32, 1)]
+
+
+@pytest.mark.parametrize(
+    ("module_name", "function_name", "extra"),
+    [
+        ("nodes.dora_extract_wd", "extract_dora_from_files", {}),
+        ("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
+        ("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": False}),
+        ("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": True}),
+    ],
+)
+def test_fixed_rule_changes_effective_rank_in_every_other_extraction_family(
+    monkeypatch, tmp_path, module_name, function_name, extra
+):
+    module = _load(module_name)
+    parameters = _load("nodes.layer_parameters")
+    first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
+    key = "blocks.45.attn.qkv_proj.weight"
+    _write(first, {key: torch.eye(3)})
+    _write(second, {key: torch.zeros(3, 3)})
+    monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
+    rules = parameters.parse_rules(r"(blocks\.45\.attn\.qkv_proj) a:2")
+    kwargs = dict(linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+                  layer_parameters=rules, **extra)
+    getattr(module, function_name)(
+        str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output), **kwargs
+    )
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        factor_key = next(key for key in result.keys() if key.endswith((".lora_A.weight", ".lora_down.weight")))
+        assert result.get_shape(factor_key)[0] == 2
+
+
+@pytest.mark.parametrize(
+    ("module_name", "nodes"),
+    [
+        ("nodes.lora_extract_svd", ("LoRAExtractFixed", "LoRAExtractRatio", "LoRAExtractQuantile", "LoRAExtractKnee", "LoRAExtractFrobenius")),
+        ("nodes.dora_extract_wd", ("DoRAExtractFixed", "DoRAExtractRatio", "DoRAExtractQuantile", "DoRAExtractKnee", "DoRAExtractFrobenius")),
+        ("nodes.dora_learned_wd", ("DoRALearnedExtractFixed", "DoRALearnedExtractRatio", "DoRALearnedExtractQuantile", "DoRALearnedExtractKnee", "DoRALearnedExtractFrobenius")),
+        ("nodes.text_encoder_extract", ("TextEncoderLoRAExtractFixed", "TextEncoderLoRAExtractRatio", "TextEncoderLoRAExtractQuantile", "TextEncoderLoRAExtractKnee", "TextEncoderLoRAExtractFrobenius", "TextEncoderDoRAExtractFixed", "TextEncoderDoRAExtractRatio", "TextEncoderDoRAExtractQuantile", "TextEncoderDoRAExtractKnee", "TextEncoderDoRAExtractFrobenius")),
+    ],
+)
+def test_all_extraction_schemas_append_layer_parameter_socket(module_name, nodes):
+    module = _load(module_name)
+    for node_name in nodes:
+        assert module.__dict__[node_name].define_schema().inputs[-1].id == "layer_parameters"

tests/test_layer_parameter_receivers.py

@@ -0,0 +1,112 @@
+import importlib
+import inspect
+import sys
+import types
+from pathlib import Path
+
+import pytest
+from comfy_api.latest import io
+
+
+RECEIVERS = [
+    ("merger", 10), ("lora_extract_svd", 5), ("dora_extract_wd", 5),
+    ("dora_learned_wd", 5), ("text_encoder_extract", 10), ("lora_resize", 4),
+]
+
+
+def load(name):
+    package_name = "modelutils_layer_receiver_tests"
+    if package_name not in sys.modules:
+        package = types.ModuleType(package_name)
+        package.__path__ = [str(Path(__file__).resolve().parents[1])]
+        sys.modules[package_name] = package
+    return importlib.import_module(f"{package_name}.nodes.{name}")
+
+
+def nodes(module):
+    return [
+        node for name, node in vars(module).items()
+        if isinstance(node, type) and not name.startswith("_")
+        and node.__module__ == module.__name__ and issubclass(node, io.ComfyNode)
+        and (not module.__name__.endswith(".lora_resize") or name.startswith("LoRAResize"))
+    ]
+
+
+@pytest.mark.parametrize("module_name,count", RECEIVERS)
+def test_all_39_receivers_append_optional_typed_socket(module_name, count):
+    module = load(module_name)
+    classes = nodes(module)
+    assert len(classes) == count
+    config_output = load("layer_parameters").LayerParameterConfiguration.define_schema().outputs[0]
+    for node in classes:
+        schema = node.define_schema()
+        input_ids = [item.id for item in schema.inputs]
+        assert input_ids[-2:] == ["include_mode", "layer_parameters"], schema.node_id
+        assert input_ids.count("layer_parameters") == 1
+        socket = schema.inputs[-1]
+        assert socket.optional
+        assert socket.Parent.io_type == config_output.Parent.io_type == "MODELUTILS_LAYER_PARAMETERS"
+        assert "a=" in socket.tooltip
+        assert inspect.signature(node.execute).parameters["layer_parameters"].default is None
+
+
+@pytest.mark.parametrize("module_name,count", RECEIVERS)
+def test_all_39_receivers_forward_configuration(monkeypatch, module_name, count):
+    module = load(module_name)
+    rules = load("layer_parameters").parse_rules("(layer) a:2")
+    captured = []
+    helper_names = {
+        "lora_extract_svd": "extract_lora_from_files",
+        "dora_extract_wd": "extract_dora_from_files",
+        "dora_learned_wd": "extract_dora_learned_from_files",
+        "text_encoder_extract": "extract_te_from_files",
+        "lora_resize": "resize_lora_file",
+    }
+    if module_name == "merger":
+        def merge(model_names, calc_mode, all_modes, params, model_type):
+            captured.append(params["layer_parameters"])
+            return "result.safetensors"
+        monkeypatch.setattr(module.MergerLogic, "execute_merge", merge)
+    else:
+        name = helper_names[module_name]
+        signature = inspect.signature(getattr(module, name))
+
+        def process(*args, **kwargs):
+            captured.append(signature.bind(*args, **kwargs).arguments["layer_parameters"])
+            return "result.safetensors"
+
+        monkeypatch.setattr(module, name, process)
+    monkeypatch.setattr(module.folder_paths, "get_full_path_or_raise", lambda category, name: name)
+    if module_name in {"lora_extract_svd", "dora_extract_wd", "dora_learned_wd", "text_encoder_extract"}:
+        monkeypatch.setattr(module, "_build_lora_output_path", lambda name: (name, name))
+    if module_name == "lora_resize":
+        monkeypatch.setattr(module, "canonical_model_artifact_path", lambda category, name: (name, name))
+    for node in nodes(module):
+        values = {}
+        for item in node.define_schema().inputs:
+            value = getattr(item, "default", None)
+            if value is None:
+                options = getattr(item, "options", [])
+                value = options[0] if options else item.id
+            values[item.id] = value
+        values["layer_parameters"] = rules
+        if "execution_mode" in values:
+            values["execution_mode"] = "MERGE"
+        node.execute(**values)
+    assert len(captured) == count
+    assert all(value is rules for value in captured)
+
+
+def test_no_expansion_to_other_node_families():
+    for module_name in ("consensus_merger", "cwb_model_graft", "cwb_delta_lora_merger", "lodestone_merger", "lora_merger", "model_analysis", "dtype_conversion"):
+        module = load(module_name)
+        for node in nodes(module):
+            assert "layer_parameters" not in {item.id for item in node.define_schema().inputs}
+    resize = load("lora_resize")
+    for node in (
+        resize.LoRAMergeToModel,
+        resize.LoRANormalizeAlpha,
+        resize.LoRAModifyAlpha,
+        resize.LoRADenormalizeAlpha,
+    ):
+        assert "layer_parameters" not in {item.id for item in node.define_schema().inputs}

tests/test_layer_parameter_resize.py

@@ -0,0 +1,258 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_layer_parameter_resize_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 resize():
+    return _load_module("nodes.lora_resize")
+
+
+@pytest.fixture(scope="module")
+def layer_parameters():
+    return _load_module("nodes.layer_parameters")
+
+
+def _write_uel(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), metadata={}) as writer:
+        writer.write_batch(
+            [(key, tensor.cpu().contiguous()) for key, tensor in tensors.items()]
+        )
+
+
+def _read_uel(path):
+    result = {}
+    with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
+        stream = loader.async_stream(
+            loader.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
+        )
+        for batch in stream:
+            for key, tensor in batch:
+                result[key] = tensor.clone()
+                loader.mark_processed(key)
+    return result
+
+
+@pytest.mark.parametrize(
+    ("node_name", "extra"),
+    [
+        ("LoRAResizeFixed", {"new_rank": 1}),
+        ("LoRAResizeRatio", {"max_rank": 1, "ratio": 2.0}),
+        (
+            "LoRAResizeFrobenius",
+            {"max_rank": 1, "min_rank": 1, "target": 0.9},
+        ),
+        (
+            "LoRAResizeCumulative",
+            {"max_rank": 1, "target": 0.9},
+        ),
+    ],
+)
+def test_resize_nodes_append_and_forward_layer_parameters(
+    monkeypatch, tmp_path, resize, node_name, extra
+):
+    node = getattr(resize, node_name)
+    schema = node.define_schema()
+    assert schema.inputs[-1].id == "layer_parameters"
+
+    calls = []
+    layer_rules = object()
+    monkeypatch.setattr(
+        resize.folder_paths, "get_full_path_or_raise", lambda *args: "input"
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(
+        resize, "resize_lora_file", lambda *args, **kwargs: calls.append(kwargs)
+    )
+    node.execute(
+        lora_name="input",
+        output_filename="out",
+        save_dtype="fp16",
+        device="cpu",
+        force_clear_cache=False,
+        layer_parameters=layer_rules,
+        **extra,
+    )
+    assert calls[0]["layer_parameters"] is layer_rules
+
+
+def test_fixed_resize_rule_applies_only_to_matched_canonical_layer(
+    monkeypatch, tmp_path, resize, layer_parameters
+):
+    source = tmp_path / "source.safetensors"
+    tensors = {}
+    for name in ("matched", "unmatched"):
+        tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3)
+        tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3)
+    _write_uel(source, tensors)
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.resize_lora_file(
+        str(source),
+        1,
+        None,
+        None,
+        "cpu",
+        torch.float32,
+        "resized",
+        verbose=False,
+        layer_parameters=layer_parameters.parse_rules(r"(diffusion_model\.matched) a:2"),
+    )
+
+    actual = _read_uel(output)
+    assert actual["diffusion_model.matched.lora_A.weight"].shape[0] == 2
+    assert actual["diffusion_model.unmatched.lora_A.weight"].shape[0] == 1
+
+
+@pytest.mark.parametrize(
+    ("method", "parameter", "max_rank", "min_rank", "rule", "expected_rank"),
+    [
+        (None, None, 1, 1, "a:2", 2),
+        ("sv_ratio", 2.0, 1, 1, "a:3 b:100", 3),
+        ("sv_fro", 0.1, 1, 1, "a:3 b:2 c:0.1", 2),
+        ("sv_cumulative", 0.1, 1, 1, "a:3 b:0.9", 2),
+    ],
+)
+def test_resize_rule_overrides_each_profile_numerically(
+    monkeypatch,
+    tmp_path,
+    resize,
+    layer_parameters,
+    method,
+    parameter,
+    max_rank,
+    min_rank,
+    rule,
+    expected_rank,
+):
+    source = tmp_path / f"{method or 'fixed'}.safetensors"
+    _write_uel(
+        source,
+        {
+            "diffusion_model.target.lora_A.weight": torch.eye(3),
+            "diffusion_model.target.lora_B.weight": torch.diag(
+                torch.tensor([10.0, 3.0, 1.0])
+            ),
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.resize_lora_file(
+        str(source),
+        max_rank,
+        method,
+        parameter,
+        "cpu",
+        torch.float32,
+        f"{method or 'fixed'}_output",
+        verbose=False,
+        min_rank=min_rank,
+        layer_parameters=layer_parameters.parse_rules(
+            rf"(diffusion_model\.target) {rule}"
+        ),
+    )
+
+    assert _read_uel(output)["diffusion_model.target.lora_A.weight"].shape[0] == expected_rank
+
+
+def test_resize_rules_respect_include_discard_and_low_bit_precedence(
+    monkeypatch, tmp_path, resize, layer_parameters
+):
+    source = tmp_path / "precedence.safetensors"
+    low_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
+    _write_uel(
+        source,
+        {
+            "diffusion_model.low.lora_A.weight": low_down,
+            "diffusion_model.low.lora_B.weight": torch.eye(2),
+            "diffusion_model.normal.lora_A.weight": torch.eye(3),
+            "diffusion_model.normal.lora_B.weight": torch.eye(3),
+            "diffusion_model.discard.lora_A.weight": torch.eye(3),
+            "diffusion_model.discard.lora_B.weight": torch.eye(3),
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.resize_lora_file(
+        str(source),
+        2,
+        None,
+        None,
+        "cpu",
+        torch.float32,
+        "precedence_output",
+        verbose=False,
+        exclude_patterns=r"diffusion_model\.(low|normal|discard)",
+        discard_patterns=r"diffusion_model\.discard",
+        include_mode=True,
+        layer_parameters=layer_parameters.parse_rules(
+            "\n".join(
+                [
+                    r"(diffusion_model\.low) a:1",
+                    r"(diffusion_model\.normal) a:1",
+                    r"(diffusion_model\.discard) a:1",
+                ]
+            )
+        ),
+    )
+
+    actual = _read_uel(output)
+    assert torch.equal(actual["diffusion_model.low.lora_A.weight"], low_down)
+    assert actual["diffusion_model.normal.lora_A.weight"].shape[0] == 1
+    assert not any(".discard." in key for key in actual)
+
+
+def test_frobenius_rule_validation_precedes_writer(
+    monkeypatch, tmp_path, resize, layer_parameters
+):
+    source = tmp_path / "source.safetensors"
+    _write_uel(
+        source,
+        {
+            "diffusion_model.target.lora_A.weight": torch.eye(3),
+            "diffusion_model.target.lora_B.weight": torch.eye(3),
+        },
+    )
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+    monkeypatch.setattr(
+        resize,
+        "atomic_uel_writer",
+        lambda *args, **kwargs: pytest.fail("Writer must not open for invalid rule"),
+    )
+
+    with pytest.raises(ValueError, match="min_rank=3 exceeds max_rank=2"):
+        resize.resize_lora_file(
+            str(source),
+            4,
+            "sv_fro",
+            0.9,
+            "cpu",
+            torch.float32,
+            "unused",
+            verbose=False,
+            layer_parameters=layer_parameters.parse_rules(r"(diffusion_model\.target) a:2 b:3"),
+        )

tests/test_layer_parameters.py

@@ -0,0 +1,165 @@
+import importlib
+import sys
+import types
+from dataclasses import FrozenInstanceError
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture(scope="module")
+def rules_module():
+    name = "modelutils_layer_rule_tests"
+    package = types.ModuleType(name)
+    package.__path__ = [str(Path(__file__).resolve().parents[1])]
+    sys.modules[name] = package
+    return importlib.import_module(f"{name}.nodes.layer_parameters")
+
+
+def resolve(module, text, *, profile="extract:fixed", names=("blocks.45.attn.qkv_proj.weight",), defaults=None, glob=False):
+    return module.resolve_layer_parameters(
+        module.parse_rules(text, glob), profile, names,
+        defaults or {"linear_dim": 8, "conv_dim": 4, "clamp_quantile": 0.99, "min_diff": 0},
+        node_name="Example Node",
+    )
+
+
+@pytest.mark.parametrize("assignments", [
+    "a:64 b:32 c:0.99 d:0.0",
+    "a : 64, b:32 c:0.99, d : 0.0",
+    "linear_dim:64 b:3.2e1 clamp_quantile:+.99 d:0e0",
+])
+def test_aliases_separators_and_inheritance(rules_module, assignments):
+    module = rules_module
+    result = resolve(module, rf"(blocks\.4[589]\.attn\.qkv_proj) {assignments}")
+    assert result == {"blocks.45.attn.qkv_proj.weight": {
+        "linear_dim": 64, "conv_dim": 32, "clamp_quantile": 0.99, "min_diff": 0,
+    }}
+    partial = resolve(module, "(qkv_proj) a:64", names=("other", "blocks.45.attn.qkv_proj.weight"))
+    assert "other" not in partial
+    assert partial["blocks.45.attn.qkv_proj.weight"]["conv_dim"] == 4
+
+
+def test_patterns_preserve_groups_spaces_escapes_and_lines(rules_module):
+    text = " # explanation\r\n\r\n" + r"((?:blocks\.45|blocks\.49)\.my layer\.(?:qkv_proj)) a:16"
+    parsed = rules_module.parse_rules(text)
+    assert parsed.rules[0].line == 3
+    assert parsed.rules[0].pattern == r"(?:blocks\.45|blocks\.49)\.my layer\.(?:qkv_proj)"
+    result = rules_module.resolve_layer_parameters(parsed, "extract:fixed", ["blocks.49.my layer.qkv_proj"], {}, node_name="Fixed")
+    assert result["blocks.49.my layer.qkv_proj"]["linear_dim"] == 16
+
+
+@pytest.mark.parametrize("text", [
+    "() a:1", "( ) a:1", "(layer)", "layer a:1", "(layer)a:1",
+    "([invalid) a:1", "(layer) a:", "(layer) a:nan", "(layer) a:inf",
+    "(layer) a:1e309", "(layer) a:1a:2", "(layer) a:1,", "(layer) a:1,,b:2",
+    "(layer) a:1 # comment", "(layer) a:1 a:2", "(layer) a:1 rubbish",
+])
+def test_invalid_syntax_errors_at_configurator(rules_module, text):
+    with pytest.raises(ValueError, match="line 1"):
+        rules_module.LayerParameterConfiguration.execute(text)
+
+
+@pytest.mark.parametrize("assignments, message", [
+    ("a:1 alpha:2", "unsupported"),
+    ("a:1 linear_dim:1", "duplicate aliases"),
+    ("a:1.5", "integral"),
+    ("a:0", "linear_dim"),
+    ("a:16385", "linear_dim"),
+    ("c:0.49", "clamp_quantile"),
+    ("d:1.01", "min_diff"),
+    ("seed:1", "unsupported"),
+])
+def test_receiver_validates_parameter_domains(rules_module, assignments, message):
+    with pytest.raises(ValueError, match=message) as error:
+        resolve(rules_module, f"(qkv_proj) {assignments}")
+    assert "Example Node" in str(error.value)
+    assert "line 1" in str(error.value)
+
+
+def test_strict_overlap_even_for_disjoint_assignments(rules_module):
+    with pytest.raises(ValueError, match="conflicting") as error:
+        resolve(rules_module, "(blocks) a:32\n\n(qkv_proj) b:16")
+    assert "blocks.45.attn.qkv_proj.weight" in str(error.value)
+    assert "line 1" in str(error.value) and "line 3" in str(error.value)
+
+
+def test_zero_match_and_repeated_logical_names(rules_module):
+    with pytest.raises(ValueError, match="line 2"):
+        resolve(rules_module, "(blocks) a:32\n(missing) b:16")
+    assert resolve(rules_module, "(qkv_proj) a:32", names=["qkv_proj", "qkv_proj"]) == {
+        "qkv_proj": {"linear_dim": 32, "conv_dim": 4, "clamp_quantile": 0.99, "min_diff": 0},
+    }
+
+
+def test_glob_is_case_sensitive_substring(rules_module):
+    result = resolve(rules_module, "(blocks.4[589].attn.qkv_*) a:32", glob=True)
+    assert len(result) == 1
+    with pytest.raises(ValueError, match="no target layer"):
+        resolve(rules_module, "(BLOCKS*) a:32", glob=True)
+
+
+def test_frobenius_validates_inherited_bounds(rules_module):
+    defaults = {"max_rank": 128, "min_rank": 32, "target": 0.9}
+    with pytest.raises(ValueError, match="min_rank=32 exceeds max_rank=16"):
+        resolve(rules_module, "(qkv_proj) a:16", profile="resize:frobenius", defaults=defaults)
+    result = resolve(rules_module, "(qkv_proj) a:16 b:1", profile="resize:frobenius", defaults=defaults)
+    assert next(iter(result.values())) == {"max_rank": 16, "min_rank": 1, "target": 0.9}
+    assert defaults["max_rank"] == 128
+
+
+def test_shared_rules_do_not_leak_receiver_state(rules_module):
+    rules = rules_module.parse_rules("(layer) a:2")
+    with pytest.raises(FrozenInstanceError):
+        rules.glob_patterns = True
+    first = rules_module.resolve_layer_parameters(rules, "extract:fixed", ["layer"], {"conv_dim": 4}, node_name="A")
+    second = rules_module.resolve_layer_parameters(rules, "resize:fixed", ["layer"], {"new_rank": 8}, node_name="B")
+    first["layer"]["linear_dim"] = 64
+    assert second == {"layer": {"new_rank": 2}}
+    assert rules.rules[0].assignments == (("a", 2),)
+    assert rules_module.resolve_layer_parameters(None, "anything", [], {}, node_name="A") == {}
+    assert rules_module.resolve_layer_parameters(rules_module.parse_rules("\n#empty"), "anything", [], {}, node_name="A") == {}
+    with pytest.raises(TypeError, match="Layer Parameter Configuration"):
+        rules_module.resolve_layer_parameters("not a config", "extract:fixed", [], {}, node_name="A")
+
+
+def test_fixed_dimensions_update_only_explicit_rank_caps(rules_module):
+    defaults = {"linear_dim": 8, "conv_dim": 4, "linear_max_rank": 2, "conv_max_rank": 1, "clamp_quantile": 0.99, "min_diff": 0}
+    result = resolve(rules_module, "(qkv_proj) a:8", defaults=defaults)
+    values = next(iter(result.values()))
+    assert values["linear_max_rank"] == 8
+    assert values["conv_max_rank"] == 1
+    values = next(iter(resolve(rules_module, "(qkv_proj) c:.9", defaults=defaults).values()))
+    assert values["linear_max_rank"] == 2
+    assert values["conv_max_rank"] == 1
+
+
+def test_merge_aliases_mode_restrictions_and_quantiles(rules_module):
+    result = resolve(rules_module, "(qkv_proj) a:.5 b:.25 g:.75 d:1", profile="merge:Extract-Features", defaults={})
+    assert next(iter(result.values()))["gamma"] == 0.75
+    with pytest.raises(ValueError, match="duplicate aliases"):
+        resolve(rules_module, "(qkv_proj) c:.5 g:.5", profile="merge:Extract-Features")
+    with pytest.raises(ValueError, match="unsupported parameter"):
+        resolve(rules_module, "(qkv_proj) b:.5", profile="merge:Weight-Sum")
+    with pytest.raises(ValueError, match="unsupported parameter"):
+        resolve(rules_module, "(qkv_proj) a:.5", profile="merge:Power-Up Enhanced (DARE+TIES)")
+    with pytest.raises(ValueError, match="must be in"):
+        resolve(rules_module, "(qkv_proj) b:1.1", profile="merge:Power-Up (DARE+TIES)")
+
+
+def test_configurator_documentation_and_socket_types(rules_module):
+    module = rules_module
+    schema = module.LayerParameterConfiguration.define_schema()
+    assert [item.id for item in schema.inputs] == ["rules", "glob_patterns"]
+    config, docs = module.LayerParameterConfiguration.execute("(layer) a:32")
+    assert isinstance(config, module.LayerParameterRules)
+    assert docs.startswith("# Per-layer parameter rules")
+    assert len(schema.outputs) == 2
+    profiles = [*(f"merge:{mode}" for mode in module.MERGE_MODES), *(f"extract:{mode}" for mode in module.EXTRACTION_BINDINGS), *(f"resize:{mode}" for mode in module.RESIZE_BINDINGS)]
+    for profile in profiles:
+        assert profile in docs
+        widget = module.parameter_input("merge" if profile.startswith("merge:") else profile)
+        assert widget.optional
+        for parameter in module.parameter_bindings(profile):
+            assert parameter.name in widget.tooltip
+            assert parameter.name in docs

tests/test_lodestone_merger.py

@@ -0,0 +1,186 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_lodestone_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 lodestone():
+    return _load_module("nodes.lodestone_merger")
+
+
+def _write(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
+        writer.write_batch(list(tensors.items()))
+
+
+def _read(path):
+    result = {}
+    with MemoryEfficientSafeOpen(str(path), low_memory=True) as handler:
+        stream = handler.async_stream(handler.keys(), batch_size=1, prefetch_batches=1, pin_memory=False)
+        for batch in stream:
+            for key, tensor in batch:
+                result[key] = tensor.clone()
+                handler.mark_processed(key)
+    return result
+
+
+@pytest.mark.parametrize(
+    ("method", "expected"),
+    [
+        ("sum", torch.diag(torch.tensor([2.0, 4.0]))),
+        ("mean", torch.diag(torch.tensor([1.0, 2.0]))),
+        ("slotnorm", torch.eye(2)),
+        ("normmatch", torch.diag(torch.tensor([2.0, 4.0])) * (2.0 / torch.sqrt(torch.tensor(20.0)))),
+        ("slotnorm-normmatch", torch.eye(2) * torch.sqrt(torch.tensor(2.0))),
+    ],
+)
+def test_methods_match_supplied_dense_math(lodestone, method, expected):
+    deltas = [torch.diag(torch.tensor([2.0, 0.0])), torch.diag(torch.tensor([0.0, 4.0]))]
+    torch.testing.assert_close(lodestone.merge_lodestone_deltas(deltas, method), expected)
+
+
+def test_convolutional_factor_expansion(lodestone):
+    down = torch.arange(18, dtype=torch.float32).reshape(2, 1, 3, 3)
+    up = torch.tensor([[[[2.0]], [[-1.0]]]])
+    expected = (up.reshape(1, 2) @ down.reshape(2, -1)).reshape(1, 1, 3, 3)
+    torch.testing.assert_close(lodestone.factor_delta(down, up, "conv"), expected)
+
+
+def test_zero_cancellation_remains_zero(lodestone):
+    delta = torch.tensor([[1.0, -2.0]])
+    result = lodestone.merge_lodestone_deltas([delta, -delta], "normmatch")
+    assert torch.count_nonzero(result) == 0
+
+
+@pytest.mark.parametrize("count", [3, 8])
+def test_three_and_multi_input_dense_means(lodestone, count):
+    deltas = [torch.full((2, 2), float(index)) for index in range(1, count + 1)]
+    expected = torch.full((2, 2), (count + 1) / 2)
+    torch.testing.assert_close(
+        lodestone.merge_lodestone_deltas(deltas, "mean"), expected
+    )
+
+
+def test_layer_plan_enforces_missing_filters_and_low_bit(lodestone):
+    class Handler:
+        def get_dtype(self, _key):
+            return torch.float32
+
+    handlers = [Handler(), Handler()]
+    roles = {"down": "layer.down", "up": "layer.up"}
+    infos = [
+        {"pairs": {"layer": roles}, "low_bit": set()},
+        {"pairs": {}, "low_bit": set()},
+    ]
+    layer_map = {"layer": [(0, "layer")]}
+    with pytest.raises(ValueError, match="missing"):
+        lodestone._build_units(layer_map, infos, handlers, "error", "", "", False)
+    units, _ = lodestone._build_units(layer_map, infos, handlers, "zeros", "", "", False)
+    assert units[0]["missing"] == 1
+    excluded, _ = lodestone._build_units(layer_map, infos, handlers, "zeros", "*.diff", "", True)
+    assert excluded[0]["missing"] == 0
+    discarded, _ = lodestone._build_units(layer_map, infos, handlers, "skip", "", "*.diff", True)
+    assert discarded == []
+
+    infos[0]["low_bit"] = {"layer.down"}
+    with pytest.raises(ValueError, match="low-bit"):
+        lodestone._build_units(layer_map, infos, handlers, "skip", "", "", False)
+
+
+def test_streamed_merge_writes_only_alpha_normalized_full_delta(monkeypatch, tmp_path, lodestone):
+    first = tmp_path / "first.safetensors"
+    second = tmp_path / "second.safetensors"
+    _write(first, {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [0.0]]),
+        "diffusion_model.layer.alpha": torch.tensor(0.5),
+    })
+    _write(second, {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[0.0, 1.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[0.0], [4.0]]),
+    })
+    monkeypatch.setattr(lodestone.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(lodestone, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(lodestone, "cleanup_after_operation", lambda: None)
+
+    output = lodestone.merge_lodestone_loras(
+        [str(first), str(second)], "sum", "cpu", torch.float32, "dense", verbose=False
+    )
+    tensors = _read(output)
+
+    assert list(tensors) == ["diffusion_model.layer.diff"]
+    torch.testing.assert_close(tensors["diffusion_model.layer.diff"], torch.diag(torch.tensor([1.0, 4.0])))
+
+
+@pytest.mark.parametrize("patterns, glob_patterns, selected", [
+    (r"\.selected", False, True),
+    ("*.selected*", True, True),
+    ("", False, False),
+    ("no_match", False, False),
+])
+def test_include_filter_preserves_nonmatches_and_discards_first(
+    monkeypatch, tmp_path, lodestone, patterns, glob_patterns, selected,
+):
+    paths = [tmp_path / f"{name}.safetensors" for name in ("a", "b")]
+    for path, value in zip(paths, (2.0, 6.0)):
+        _write(path, {
+            f"diffusion_model.{layer}.diff": torch.full((2, 2), value)
+            for layer in ("selected", "other", "selected_discard")
+        })
+    monkeypatch.setattr(lodestone.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(lodestone, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(lodestone, "cleanup_after_operation", lambda: None)
+    output = lodestone.merge_lodestone_loras(
+        [str(path) for path in paths], "sum", "cpu", torch.float32, "include_filter",
+        exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=True,
+        discard_patterns="*discard*" if glob_patterns else "discard", verbose=False,
+    )
+    tensors = _read(output)
+    assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"}
+    torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 8.0 if selected else 2.0))
+    torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0))
+
+
+def test_dedicated_nodes_do_not_modify_factor_merger_schema(lodestone):
+    standard = _load_module("nodes.lora_merger")
+    assert [node.define_schema().node_id for node in lodestone.LODESTONE_MERGER_NODES] == [
+        "LodestoneLoRATwoMerger", "LodestoneLoRAThreeMerger", "LodestoneLoRAMultiMerger"
+    ]
+    standard_inputs = {item.id: item for item in standard.LoRAMultiMerge.define_schema().inputs}
+    assert standard_inputs["merge_mode"].options == ["concatenate", "weighted_sum"]
+    assert "calc_mode" not in standard_inputs
+
+
+def test_dense_cuda_oom_retries_layer_on_cpu(monkeypatch, lodestone):
+    unit = {"dtypes": [torch.float32], "missing": 0, "sources": []}
+    calls = []
+
+    def process(_unit, _loaded, _method, _dtype, device):
+        calls.append(device.type)
+        if device.type == "cuda":
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory")
+        return torch.ones(1)
+
+    monkeypatch.setattr(lodestone, "_process_on_device", process)
+    monkeypatch.setattr(lodestone, "_release_cuda_oom", lambda: None)
+    output, fallback = lodestone._process(unit, {}, "sum", torch.float32, "cuda")
+    assert fallback is True
+    assert calls == ["cuda", "cpu"]
+    assert torch.equal(output, torch.ones(1))

tests/test_lora_merger_async.py

@@ -0,0 +1,379 @@
+import importlib
+import sys
+import types
+from collections import Counter
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import (
+    IncrementalSafetensorsWriter,
+    MemoryEfficientSafeOpen,
+)
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_lora_merger_async_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 merger():
+    return _load_module("nodes.lora_merger")
+
+
+def _write_uel(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
+        writer.write_batch(
+            [(key, tensor.cpu().contiguous()) for key, tensor in tensors.items()]
+        )
+
+
+def _read_uel(path, async_method=None, mark_method=None):
+    result = {}
+    with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
+        stream = (async_method or MemoryEfficientSafeOpen.async_stream)(
+            loader,
+            loader.keys(),
+            batch_size=1,
+            prefetch_batches=1,
+            pin_memory=False,
+        )
+        for batch in stream:
+            for key, tensor in batch:
+                result[key] = tensor.clone()
+                (mark_method or MemoryEfficientSafeOpen.mark_processed)(loader, key)
+    return result
+
+
+def _patch_runtime(monkeypatch, tmp_path, merger):
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+
+
+def _mixed_sources(tmp_path):
+    first = tmp_path / "mixed_first.safetensors"
+    second = tmp_path / "mixed_second.safetensors"
+    first_tensors = {
+        "diffusion_model.small.lora_A.weight": torch.tensor([[1.0, 0.0]]),
+        "diffusion_model.small.lora_B.weight": torch.tensor([[1.0], [2.0]]),
+        "diffusion_model.small.alpha": torch.tensor(0.5),
+        "diffusion_model.large.lora_A.weight": torch.eye(3),
+        "diffusion_model.large.lora_B.weight": torch.eye(3),
+    }
+    second_tensors = {
+        "diffusion_model.small.lora_A.weight": torch.eye(2),
+        "diffusion_model.small.lora_B.weight": torch.eye(2),
+        "diffusion_model.large.lora_A.weight": torch.tensor([[1.0, 1.0, 1.0]]),
+        "diffusion_model.large.lora_B.weight": torch.tensor([[1.0], [2.0], [3.0]]),
+    }
+    _write_uel(first, first_tensors)
+    _write_uel(second, second_tensors)
+    return first, second, first_tensors, second_tensors
+
+
+def test_concatenate_uses_per_layer_rank_and_alpha(monkeypatch, tmp_path, merger):
+    _patch_runtime(monkeypatch, tmp_path, merger)
+    first, second, first_tensors, second_tensors = _mixed_sources(tmp_path)
+
+    output = merger.merge_multi_loras(
+        [str(first), str(second)],
+        [2.0, -0.5],
+        "concatenate",
+        "cpu",
+        torch.float32,
+        "mixed_concat",
+        verbose=False,
+    )
+    tensors = _read_uel(output)
+
+    assert tensors["diffusion_model.small.lora_A.weight"].shape[0] == 3
+    assert tensors["diffusion_model.large.lora_A.weight"].shape[0] == 4
+    assert not any(key.endswith(".alpha") for key in tensors)
+
+    small_expected = (
+        2.0
+        * (0.5 / 1.0)
+        * first_tensors["diffusion_model.small.lora_B.weight"]
+        @ first_tensors["diffusion_model.small.lora_A.weight"]
+        - 0.5
+        * second_tensors["diffusion_model.small.lora_B.weight"]
+        @ second_tensors["diffusion_model.small.lora_A.weight"]
+    )
+    small_actual = (
+        tensors["diffusion_model.small.lora_B.weight"]
+        @ tensors["diffusion_model.small.lora_A.weight"]
+    )
+    torch.testing.assert_close(small_actual, small_expected)
+
+
+@pytest.mark.parametrize("strategy", ["weighted_sum", "dare", "enhanced_dare"])
+def test_nonconcatenate_variants_use_each_layers_max_rank(
+    monkeypatch, tmp_path, merger, strategy
+):
+    _patch_runtime(monkeypatch, tmp_path, merger)
+    first, second, _, _ = _mixed_sources(tmp_path)
+    common = [str(first), str(second)], [1.0, 1.0]
+    if strategy == "weighted_sum":
+        output = merger.merge_multi_loras(
+            *common,
+            "weighted_sum",
+            "cpu",
+            torch.float32,
+            "mixed_weighted",
+            verbose=False,
+        )
+    elif strategy == "dare":
+        output = merger.merge_multi_loras_dare(
+            *common,
+            0.0,
+            0.0,
+            7,
+            "cpu",
+            torch.float32,
+            "mixed_dare",
+            verbose=False,
+        )
+    else:
+        output = merger.merge_multi_loras_dare_enhanced(
+            *common,
+            1.0,
+            1.0,
+            0.0,
+            0.0,
+            7,
+            "cpu",
+            torch.float32,
+            "mixed_enhanced",
+            verbose=False,
+        )
+    tensors = _read_uel(output)
+    assert tensors["diffusion_model.small.lora_A.weight"].shape[0] == 2
+    assert tensors["diffusion_model.large.lora_A.weight"].shape[0] == 3
+    assert not any(key.endswith(".alpha") for key in tensors)
+
+
+@pytest.mark.parametrize("strategy", ["weighted_sum", "dare", "enhanced_dare"])
+def test_nonconcatenate_variants_apply_source_alpha_once(
+    monkeypatch, tmp_path, merger, strategy
+):
+    _patch_runtime(monkeypatch, tmp_path, merger)
+    first = tmp_path / f"alpha_first_{strategy}.safetensors"
+    second = tmp_path / f"alpha_second_{strategy}.safetensors"
+    _write_uel(
+        first,
+        {
+            "diffusion_model.layer.lora_A.weight": torch.ones((1, 1)),
+            "diffusion_model.layer.lora_B.weight": torch.full((1, 1), 2.0),
+            "diffusion_model.layer.alpha": torch.tensor(0.5),
+        },
+    )
+    _write_uel(
+        second,
+        {
+            "diffusion_model.layer.lora_A.weight": torch.ones((1, 1)),
+            "diffusion_model.layer.lora_B.weight": torch.full((1, 1), 2.0),
+        },
+    )
+    common = [str(first), str(second)], [1.0, 1.0]
+    if strategy == "weighted_sum":
+        output = merger.merge_multi_loras(
+            *common,
+            "weighted_sum",
+            "cpu",
+            torch.float32,
+            "alpha_weighted",
+            verbose=False,
+        )
+        expected_down = 2.0
+        expected_up = 3.0
+    elif strategy == "dare":
+        output = merger.merge_multi_loras_dare(
+            *common,
+            0.0,
+            0.0,
+            9,
+            "cpu",
+            torch.float32,
+            "alpha_dare",
+            verbose=False,
+        )
+        expected_down = 1.0
+        expected_up = 1.5
+    else:
+        output = merger.merge_multi_loras_dare_enhanced(
+            *common,
+            1.0,
+            1.0,
+            0.0,
+            0.0,
+            9,
+            "cpu",
+            torch.float32,
+            "alpha_enhanced",
+            verbose=False,
+        )
+        expected_down = 1.0
+        expected_up = 1.5
+
+    tensors = _read_uel(output)
+    assert tensors["diffusion_model.layer.lora_A.weight"].item() == pytest.approx(
+        expected_down
+    )
+    assert tensors["diffusion_model.layer.lora_B.weight"].item() == pytest.approx(
+        expected_up
+    )
+
+
+@pytest.mark.parametrize(
+    "strategy", ["concatenate", "weighted_sum", "dare", "enhanced_dare"]
+)
+def test_fp32_alpha_does_not_promote_bf16_factors(
+    monkeypatch, tmp_path, merger, strategy
+):
+    _patch_runtime(monkeypatch, tmp_path, merger)
+    first = tmp_path / f"bf16_alpha_first_{strategy}.safetensors"
+    second = tmp_path / f"bf16_alpha_second_{strategy}.safetensors"
+    tensors = {
+        "diffusion_model.layer.lora_A.weight": torch.ones(
+            (1, 2), dtype=torch.bfloat16
+        ),
+        "diffusion_model.layer.lora_B.weight": torch.ones(
+            (2, 1), dtype=torch.bfloat16
+        ),
+        "diffusion_model.layer.alpha": torch.tensor(1.0, dtype=torch.float32),
+    }
+    _write_uel(first, tensors)
+    _write_uel(second, tensors)
+    common = [str(first), str(second)], [1.0, 1.0]
+
+    if strategy in {"concatenate", "weighted_sum"}:
+        output = merger.merge_multi_loras(
+            *common,
+            strategy,
+            "cpu",
+            torch.bfloat16,
+            f"bf16_alpha_{strategy}",
+            verbose=False,
+        )
+    elif strategy == "dare":
+        output = merger.merge_multi_loras_dare(
+            *common,
+            0.0,
+            0.0,
+            17,
+            "cpu",
+            torch.bfloat16,
+            "bf16_alpha_dare",
+            verbose=False,
+        )
+    else:
+        output = merger.merge_multi_loras_dare_enhanced(
+            *common,
+            1.0,
+            1.0,
+            0.0,
+            0.0,
+            17,
+            "cpu",
+            torch.bfloat16,
+            "bf16_alpha_enhanced",
+            verbose=False,
+        )
+
+    merged = _read_uel(output)
+    assert merged["diffusion_model.layer.lora_A.weight"].dtype == torch.bfloat16
+    assert merged["diffusion_model.layer.lora_B.weight"].dtype == torch.bfloat16
+
+
+def test_all_normal_inputs_use_async_uel_and_release_once(
+    monkeypatch, tmp_path, merger
+):
+    _patch_runtime(monkeypatch, tmp_path, merger)
+    first, second, _, _ = _mixed_sources(tmp_path)
+    original_async = merger.MemoryEfficientSafeOpen.async_stream
+    original_mark = merger.MemoryEfficientSafeOpen.mark_processed
+    async_calls = []
+    marked = []
+
+    def track_async(self, keys, **kwargs):
+        async_calls.append((self.filename, tuple(keys), kwargs, self.low_memory))
+        yield from original_async(self, keys, **kwargs)
+
+    def track_mark(self, key):
+        marked.append((self.filename, key))
+        return original_mark(self, key)
+
+    def reject_get_tensor(*args, **kwargs):
+        raise AssertionError("multi-merge used synchronous get_tensor")
+
+    monkeypatch.setattr(merger.MemoryEfficientSafeOpen, "async_stream", track_async)
+    monkeypatch.setattr(merger.MemoryEfficientSafeOpen, "mark_processed", track_mark)
+    monkeypatch.setattr(merger.MemoryEfficientSafeOpen, "get_tensor", reject_get_tensor)
+
+    output = merger.merge_multi_loras(
+        [str(first), str(second)],
+        [1.0, 1.0],
+        "concatenate",
+        "cpu",
+        torch.float32,
+        "async_only",
+        verbose=False,
+    )
+
+    assert Path(output).is_file()
+    assert len(async_calls) == 2
+    planned = Counter(
+        (filename, key)
+        for filename, keys, _, _ in async_calls
+        for key in keys
+    )
+    assert Counter(marked) == planned
+    assert all(count == 1 for count in planned.values())
+    for _, _, options, low_memory in async_calls:
+        assert low_memory is True
+        assert options == {
+            "batch_size": 1,
+            "prefetch_batches": 1,
+            "pin_memory": False,
+        }
+
+
+def test_cuda_oom_retries_only_current_unit_on_cpu(
+    monkeypatch, tmp_path, merger
+):
+    _patch_runtime(monkeypatch, tmp_path, merger)
+    first, second, _, _ = _mixed_sources(tmp_path)
+    original = merger._process_merge_unit_on_device
+    calls = []
+
+    def fail_cuda_then_cpu(*args):
+        process_device = args[-2]
+        calls.append(process_device.type)
+        if process_device.type == "cuda":
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory")
+        return original(*args)
+
+    monkeypatch.setattr(merger, "_process_merge_unit_on_device", fail_cuda_then_cpu)
+    monkeypatch.setattr(merger, "_release_cuda_oom_state", lambda: None)
+
+    output = merger.merge_multi_loras(
+        [str(first), str(second)],
+        [1.0, 1.0],
+        "concatenate",
+        "cuda",
+        torch.float32,
+        "oom_fallback",
+        verbose=False,
+    )
+
+    assert Path(output).is_file()
+    assert calls == ["cuda", "cpu", "cuda", "cpu"]

tests/test_lora_model_analysis.py

@@ -0,0 +1,250 @@
+import csv
+import importlib
+from io import StringIO
+import json
+import logging
+from pathlib import Path
+import sys
+import types
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter
+
+
+@pytest.fixture(scope="module")
+def analysis():
+    name = "modelutils_lora_on_model_tests"
+    package = types.ModuleType(name)
+    package.__path__ = [str(Path(__file__).resolve().parents[1])]
+    sys.modules[name] = package
+    return importlib.import_module(f"{name}.nodes.lora_model_analysis")
+
+
+def write(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
+        writer.write_batch(list(tensors.items()))
+
+
+def setup(monkeypatch, tmp_path, analysis, adapter, base=None):
+    paths = {"loras": tmp_path / "adapter.safetensors", "diffusion_models": tmp_path / "base.safetensors"}
+    write(paths["loras"], adapter)
+    write(paths["diffusion_models"], base if base is not None else {
+        "model.diffusion_model.layer.weight": torch.ones((2, 2)),
+        "model.diffusion_model.untouched.weight": torch.ones((2, 2)),
+    })
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path_or_raise", lambda category, name: str(paths[category]))
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    return paths
+
+
+def params(**overrides):
+    return {"execution_mode": "ANALYZE", "process_device": "cpu", "top_weight_differences": 3, "force_clear_cache": False, "exclude_patterns": "", "glob_patterns": False, "include_mode": False, "strength": 1.0, **overrides}
+
+
+def pair(alpha=True):
+    result = {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [4.0]]),
+    }
+    if alpha:
+        result["diffusion_model.layer.alpha"] = torch.tensor(0.5)
+    return result
+
+
+@pytest.mark.parametrize("alpha,strength,mae,maximum", [(True, 1.0, 2.25, 4), (False, 1.0, 4.5, 8), (True, 0.5, 1.125, 2), (True, 0.0, 0, 0)])
+def test_reports_actual_applied_delta_without_saving(monkeypatch, tmp_path, analysis, alpha, strength, mae, maximum):
+    paths = setup(monkeypatch, tmp_path, analysis, pair(alpha))
+    before = {path: path.read_bytes() for path in paths.values()}
+    comparison, cwb, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params(strength=strength))
+    rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))}
+    row = rows["model.diffusion_model.layer.weight"]
+    assert float(row["mae"]) == pytest.approx(mae)
+    assert float(row["max"]) == maximum
+    assert float(rows["model.diffusion_model.untouched.weight"]["mae"]) == 0
+    assert "patched − original" in comparison
+    assert "original versus LoRA-patched" in cwb
+    assert {path: path.read_bytes() for path in paths.values()} == before
+    assert set(tmp_path.iterdir()) == set(paths.values())
+
+
+def test_quantized_base_is_dequantized_and_sidecars_are_not_analyzed(monkeypatch, tmp_path, analysis):
+    prefix = "model.diffusion_model.layer"
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    setup(monkeypatch, tmp_path, analysis, pair(), {
+        f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+    })
+    comparison, _, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params())
+    rows = list(csv.DictReader(StringIO(metrics_csv)))
+    assert [row["key"] for row in rows] == [f"{prefix}.weight"]
+    assert float(rows[0]["mae"]) == pytest.approx(2.25)
+    assert "weight_scale" not in comparison
+
+
+@pytest.mark.parametrize("suffixes", [(".lora_down.weight", ".lora_up.weight"), (".lora_A.default.weight", ".lora_B.default.weight"), (".lora_A", ".lora_B")])
+def test_pair_formats_and_flattened_reference_mapping(monkeypatch, tmp_path, analysis, suffixes):
+    setup(monkeypatch, tmp_path, analysis, {
+        "lora_unet_layer" + suffixes[0]: torch.ones((1, 2)),
+        "lora_unet_layer" + suffixes[1]: torch.ones((2, 1)),
+    })
+    _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
+    rows = list(csv.DictReader(StringIO(text)))
+    assert float(rows[0]["mae"]) == 1
+
+
+def test_direct_weight_bias_replacement_and_unmapped_are_explicit(monkeypatch, tmp_path, analysis):
+    setup(monkeypatch, tmp_path, analysis, {
+        "diffusion_model.layer.diff": torch.full((2, 2), 2.0),
+        "diffusion_model.layer.diff_b": torch.full((2,), 3.0),
+        "diffusion_model.replaced.set_weight": torch.full((2, 2), 5.0),
+        "diffusion_model.missing.diff": torch.ones((2, 2)),
+    }, {
+        "diffusion_model.layer.weight": torch.ones((2, 2)),
+        "diffusion_model.layer.bias": torch.ones(2),
+        "diffusion_model.replaced.weight": torch.ones((2, 2)),
+    })
+    report, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params(strength=0.5))
+    rows = {row["key"]: row for row in csv.DictReader(StringIO(text))}
+    assert float(rows["diffusion_model.layer.weight"]["mae"]) == 1
+    assert float(rows["diffusion_model.layer.bias"]["mae"]) == 1.5
+    assert float(rows["diffusion_model.replaced.weight"]["mae"]) == 4
+    assert "diffusion_model.missing" in report
+
+
+def test_bounded_stream_release_and_filtering(monkeypatch, tmp_path, analysis):
+    setup(monkeypatch, tmp_path, analysis, pair())
+    calls, released = [], []
+    original_stream = analysis.MemoryEfficientSafeOpen.async_stream
+    original_mark = analysis.MemoryEfficientSafeOpen.mark_processed
+
+    def stream(handler, keys, **kwargs):
+        calls.append((list(keys), kwargs))
+        return original_stream(handler, keys, **kwargs)
+
+    def mark(handler, key):
+        released.append(key)
+        return original_mark(handler, key)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", mark)
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "get_tensor", lambda *args: pytest.fail("synchronous tensor load"))
+    _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params(exclude_patterns="*layer.weight", glob_patterns=True, include_mode=True))
+    assert len(list(csv.DictReader(StringIO(text)))) == 1
+    assert len(calls) == 2
+    assert all(kwargs == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, kwargs in calls)
+    assert sorted(released) == sorted(key for keys, _ in calls for key in keys)
+    assert len(released) == 4
+
+
+def test_logged_adapter_failure_is_not_reported_as_zero_delta(monkeypatch, tmp_path, analysis):
+    setup(monkeypatch, tmp_path, analysis, pair())
+    marks = []
+    original_mark = analysis.MemoryEfficientSafeOpen.mark_processed
+
+    def mark(handler, key):
+        marks.append(key)
+        return original_mark(handler, key)
+
+    def fail(self, weight, key, *args):
+        logging.error("ERROR LoRA %s deliberately failed", key)
+        return weight
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", mark)
+    monkeypatch.setattr(analysis.LoRAAdapter, "calculate_weight", fail)
+    with pytest.raises(RuntimeError, match="deliberately failed"):
+        analysis.analyze_lora_on_model("adapter", "base", params())
+    assert len(marks) == 4
+    assert not any(isinstance(handler, analysis.AdapterErrorCapture) for handler in logging.getLogger().handlers)
+
+
+def test_complete_unit_cuda_retry(monkeypatch, analysis):
+    devices = []
+
+    def calculate(key, loaded, specs, strength, device, top_count):
+        devices.append(device)
+        if device == "cuda":
+            raise torch.cuda.OutOfMemoryError("test")
+        return analysis.RawStats(), analysis.CWBStats(), []
+
+    monkeypatch.setattr(analysis, "_apply_and_analyze", calculate)
+    monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
+    assert analysis._analyze_unit("layer", {}, [], 1, "cuda", 0)[-1] is True
+    assert devices == ["cuda", "cpu"]
+
+
+def test_dora_uses_reconstructed_weight_not_raw_factor_delta(monkeypatch, tmp_path, analysis):
+    setup(monkeypatch, tmp_path, analysis, {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [0.0]]),
+        "diffusion_model.layer.alpha": torch.tensor(0.5),
+        "diffusion_model.layer.dora_scale": torch.tensor([3.0, 4.0]),
+    }, {"diffusion_model.layer.weight": torch.eye(2)})
+    _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
+    row = next(csv.DictReader(StringIO(text)))
+    # Installed Comfy output-axis DoRA scales by the original base row norms.
+    # The resulting diagonal is (4.5, 4), not the raw-factor result (1.5, 1).
+    assert float(row["mae"]) == pytest.approx(1.625)
+    assert float(row["max"]) == pytest.approx(3.5)
+
+
+def test_locon_mid_and_scalar_low_bit_alpha(monkeypatch, tmp_path, analysis):
+    setup(monkeypatch, tmp_path, analysis, {
+        "diffusion_model.layer.lora_down.weight": torch.ones((1, 1, 1, 1)),
+        "diffusion_model.layer.lora_up.weight": torch.ones((1, 1, 1, 1)),
+        "diffusion_model.layer.lora_mid.weight": torch.ones((1, 1, 3, 3)),
+        "diffusion_model.layer.alpha": torch.tensor(1, dtype=torch.uint8),
+    }, {"diffusion_model.layer.weight": torch.zeros((1, 1, 3, 3))})
+    _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
+    row = next(csv.DictReader(StringIO(text)))
+    assert float(row["mae"]) == 1
+
+
+def test_low_bit_factors_reported_without_materializing(monkeypatch, tmp_path, analysis):
+    tensors = pair()
+    tensors["diffusion_model.layer.lora_A.weight"] = torch.ones((1, 2), dtype=torch.uint8)
+    setup(monkeypatch, tmp_path, analysis, tensors)
+    requested = []
+    original = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def stream(handler, keys, **kwargs):
+        requested.extend(keys)
+        return original(handler, keys, **kwargs)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
+    report, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
+    assert requested == ["model.diffusion_model.untouched.weight"]
+    assert "model.diffusion_model.layer.weight" in report.replace(r"\_", "_")
+    assert [row["key"] for row in csv.DictReader(StringIO(text))] == ["model.diffusion_model.untouched.weight"]
+
+
+def test_streams_close_after_failure(monkeypatch, tmp_path, analysis):
+    setup(monkeypatch, tmp_path, analysis, pair())
+    streams_closed = []
+    original = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def stream(handler, keys, **kwargs):
+        iterator = original(handler, keys, **kwargs)
+        try:
+            yield from iterator
+        finally:
+            iterator.close()
+            streams_closed.append(True)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
+    monkeypatch.setattr(analysis.LoRAAdapter, "calculate_weight", lambda *args: (_ for _ in ()).throw(ValueError("test failure")))
+    with pytest.raises(ValueError, match="test failure"):
+        analysis.analyze_lora_on_model("adapter", "base", params())
+    assert len(streams_closed) == 2
+
+
+def test_schema_and_documentation_mode(monkeypatch, analysis):
+    monkeypatch.setattr(analysis.folder_paths, "get_filename_list", lambda category: [category])
+    schema = analysis.LoRAOnModelAnalysis.define_schema()
+    assert schema.inputs[1].id == "model_a" and schema.inputs[1].options == ["loras"]
+    assert schema.inputs[2].id == "model_b" and schema.inputs[2].options == ["diffusion_models"]
+    assert len(schema.outputs) == 5
+    assert schema.inputs[-1].id == "strength" and schema.inputs[-1].default == 1
+    monkeypatch.setattr(analysis, "analyze_lora_on_model", lambda *args: pytest.fail("documentation loaded tensors"))
+    assert analysis.LoRAOnModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")[2].startswith("# LoRA-on-model analysis")

tests/test_lora_quantization_guard.py

@@ -0,0 +1,645 @@
+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
+from unifiedefficientloader import IncrementalSafetensorsWriter
+
+
+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 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_multiple_low_bit_config_tensors_allowed(tmp_path, modules):
+    for count in (0, 1, 2, 52):
+        path = tmp_path / f"threshold_{count}.safetensors"
+        tensors = {
+            f"diffusion_model.blocks.{i}.attn.comfy_attention.config": torch.ones(2, dtype=torch.uint8)
+            for i in range(count)
+        }
+        tensors["normal"] = torch.ones(1)
+        save_file(tensors, str(path))
+        assert len(_inspect_file(path, modules)) == count
+
+
+@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_resize_preserves_multiple_low_bit_config_tensors(monkeypatch, tmp_path, modules):
+    resize = modules["resize"]
+    _patch_output(monkeypatch, resize, tmp_path)
+    source = tmp_path / "three_low.safetensors"
+    tensors = {
+        f"diffusion_model.blocks.{i}.attn.comfy_attention.config": torch.ones(1, dtype=torch.uint8)
+        for i in range(3)
+    }
+    tensors["diffusion_model.normal.lora_down.weight"] = torch.ones((2, 2))
+    tensors["diffusion_model.normal.lora_up.weight"] = torch.ones((2, 2))
+    save_file(tensors, str(source))
+
+    output = resize.resize_lora_file(
+        str(source), 1, None, None, "cpu", torch.float16, "three_low_out", verbose=False
+    )
+    result = load_file(output)
+    for i in range(3):
+        assert result[f"diffusion_model.blocks.{i}.attn.comfy_attention.config"].dtype == torch.uint8
+
+
+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
+
+
+def test_merge_to_model_dequantizes_base_and_omits_sidecars(monkeypatch, tmp_path, modules):
+    resize = modules["resize"]
+    _patch_output(monkeypatch, resize, tmp_path)
+    base = tmp_path / "quantized_base.safetensors"
+    adapter = tmp_path / "quantized_adapter.safetensors"
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    prefix = "model.diffusion_model.layer"
+    save_file({
+        f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+    }, str(base))
+    save_file({"diffusion_model.layer.diff": torch.ones((2, 2))}, str(adapter))
+
+    output = resize.merge_loras_to_model(
+        [str(adapter)], [1.0], str(base), "cpu", torch.float16,
+        "dequantized_merge", verbose=False,
+    )
+    result = load_file(output)
+    assert list(result) == [f"{prefix}.weight"]
+    torch.testing.assert_close(result[f"{prefix}.weight"].float(), torch.full((2, 2), 3.0))
+
+
+def test_diffusion_decoder_matches_core_quantized_tensor(tmp_path, modules):
+    from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+
+    path = tmp_path / "core_int8.safetensors"
+    key = "layer.weight"
+    qdata = torch.tensor([[1, -2], [3, 4]], dtype=torch.int8)
+    scale = torch.tensor(0.5)
+    save_file({
+        key: qdata,
+        "layer.weight_scale": scale,
+        "layer.comfy_quant": torch.tensor(
+            list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
+        ),
+    }, str(path))
+    handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
+    try:
+        decoder = modules["guard"].DiffusionQuantization(handler, str(path), "Core comparison")
+        keys = decoder.required_keys(key)
+        assert set(keys) == {key, "layer.weight_scale", "layer.comfy_quant"}
+        assert decoder.data_keys == (key,)
+        stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for expected in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == expected
+                tensors[expected] = batch[0][1]
+            actual = decoder.decode(key, tensors, torch.float32)
+        finally:
+            for source_key in tensors:
+                handler.mark_processed(source_key)
+            stream.close()
+
+        layout_name = QUANT_ALGOS["int8_tensorwise"]["comfy_tensor_layout"]
+        layout = get_layout_class(layout_name)
+        expected = QuantizedTensor(
+            qdata, layout_name,
+            layout.Params(scale=scale, orig_dtype=torch.float32, orig_shape=(2, 2)),
+        ).dequantize()
+        torch.testing.assert_close(actual, expected)
+        torch.testing.assert_close(actual, qdata.float() * scale)
+    finally:
+        handler.__exit__(None, None, None)
+
+
+@pytest.mark.parametrize("prefix", ["model.", "model.diffusion_model."])
+def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules, prefix):
+    import comfy.model_detection
+
+    stem = f"{prefix}layer"
+
+    class HeaderOnly:
+        def keys(self):
+            return [f"{stem}.weight", f"{stem}.weight_scale", f"{stem}.weight_scale_2"]
+
+        def metadata(self):
+            return {"_quantization_metadata": json.dumps({"layers": {stem: {"format": "nvfp4"}}})}
+
+        def get_shape(self, key):
+            return {f"{stem}.weight": (16, 8), f"{stem}.weight_scale": (16, 1), f"{stem}.weight_scale_2": ()}[key]
+
+        def get_dtype(self, key):
+            return torch.uint8 if key == f"{stem}.weight" else torch.float32
+
+    class Config:
+        def get_model(self, state, prefix, device):
+            assert device == torch.device("meta")
+            assert all(tensor.device.type == "meta" for tensor in state.values())
+            model = torch.nn.Module()
+            model.diffusion_model = torch.nn.Module()
+            model.diffusion_model.layer = torch.nn.Module()
+            model.diffusion_model.layer._orig_shape = (3, 5)
+            return model
+
+    monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
+    decoder = modules["guard"].DiffusionQuantization(HeaderOnly(), "header-only", "Shape test")
+    assert decoder.logical_shape(f"{stem}.weight") == (3, 5)
+
+
+def test_nvfp4_stream_decode_matches_core_full_dequantization(monkeypatch, tmp_path, modules):
+    import comfy.model_detection
+    from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+
+    layout_name = QUANT_ALGOS["nvfp4"]["comfy_tensor_layout"]
+    layout = get_layout_class(layout_name)
+    original = torch.arange(15, dtype=torch.float16).reshape(3, 5) / 10
+    qdata, params = layout.quantize(original)
+    path = tmp_path / "nvfp4.safetensors"
+    metadata = {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"format": "nvfp4"}}})}
+    save_file({
+        "model.layer.weight": qdata,
+        "model.layer.weight_scale": params.block_scale,
+        "model.layer.weight_scale_2": params.scale,
+    }, str(path), metadata=metadata)
+
+    class Config:
+        def get_model(self, state, prefix, device):
+            model = torch.nn.Module()
+            model.diffusion_model = torch.nn.Module()
+            model.diffusion_model.layer = torch.nn.Module()
+            model.diffusion_model.layer._orig_shape = tuple(original.shape)
+            return model
+
+    monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
+    handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
+    try:
+        decoder = modules["guard"].DiffusionQuantization(handler, str(path), "NVFP4 test")
+        key = "model.layer.weight"
+        assert decoder.logical_shape(key) == tuple(original.shape)
+        keys = decoder.required_keys(key)
+        stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for source_key in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == source_key
+                tensors[source_key] = batch[0][1]
+            actual = decoder.decode(key, tensors, torch.float16)
+        finally:
+            for source_key in tensors:
+                handler.mark_processed(source_key)
+            stream.close()
+        expected = QuantizedTensor(qdata, layout_name, params).dequantize()
+        torch.testing.assert_close(actual, expected)
+        assert actual.shape == original.shape
+    finally:
+        handler.__exit__(None, None, None)
+
+
+def test_w6a8_int8_stream_decode_matches_core_dequantization(tmp_path, modules):
+    from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+
+    layout_name = QUANT_ALGOS["w6a8_int8"]["comfy_tensor_layout"]
+    layout = get_layout_class(layout_name)
+    original = torch.randn(4, 256, dtype=torch.float32)
+    qdata, params = layout.quantize(original, bits=6, codebook=False)
+    path = tmp_path / "w6a8.safetensors"
+    key = "model.layer.weight"
+    metadata = {
+        "_quantization_metadata": json.dumps({
+            "layers": {
+                "model.layer": {
+                    "format": "w6a8_int8",
+                    "group_size": params.group_size,
+                    "convrot_groupsize": params.convrot_groupsize,
+                }
+            }
+        })
+    }
+    save_file({
+        key: qdata,
+        "model.layer.weight_s_rel": params.scale,
+        "model.layer.weight_s_channel": params.s_channel,
+    }, str(path), metadata=metadata)
+
+    handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
+    try:
+        decoder = modules["guard"].DiffusionQuantization(handler, str(path), "W6A8 test")
+        assert decoder.logical_shape(key) == tuple(original.shape)
+        keys = decoder.required_keys(key)
+        stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for source_key in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == source_key
+                tensors[source_key] = batch[0][1]
+            actual = decoder.decode(key, tensors, torch.float32)
+        finally:
+            for source_key in tensors:
+                handler.mark_processed(source_key)
+            stream.close()
+        expected = QuantizedTensor(qdata, layout_name, params).dequantize()
+        torch.testing.assert_close(actual, expected)
+        assert actual.shape == original.shape
+    finally:
+        handler.__exit__(None, None, None)

tests/test_lora_resize_factor_space.py

@@ -0,0 +1,650 @@
+import importlib
+import sys
+import types
+from collections import Counter
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import (
+    IncrementalSafetensorsWriter,
+    MemoryEfficientSafeOpen,
+)
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_resize_factor_space_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 resize():
+    return _load_module("nodes.lora_resize")
+
+
+def _write_uel(path, tensors, metadata=None):
+    with IncrementalSafetensorsWriter(str(path), metadata=metadata or {}) as writer:
+        writer.write_batch(
+            [(key, tensor.cpu().contiguous()) for key, tensor in tensors.items()]
+        )
+
+
+def _read_uel(path, async_method=None, mark_method=None):
+    result = {}
+    with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
+        keys = loader.keys()
+        stream = (async_method or MemoryEfficientSafeOpen.async_stream)(
+            loader, keys, batch_size=1, prefetch_batches=1, pin_memory=False
+        )
+        for batch in stream:
+            for key, tensor in batch:
+                result[key] = tensor.clone()
+                (mark_method or MemoryEfficientSafeOpen.mark_processed)(loader, key)
+    return result
+
+
+@pytest.mark.parametrize("glob", [False, True])
+def test_resize_filters_preserve_and_discard_whole_layers(monkeypatch, tmp_path, resize, glob):
+    source = tmp_path / "filters.safetensors"
+    tensors = {}
+    for name in ("keep", "drop", "resize"):
+        tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3, dtype=torch.float16)
+        tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3, dtype=torch.float16)
+        tensors[f"diffusion_model.{name}.alpha"] = torch.tensor(1.5)
+    tensors["diffusion_model.drop.extra"] = torch.ones(1)
+    tensors["standalone"] = torch.ones(1)
+    _write_uel(source, tensors)
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+    original_async = MemoryEfficientSafeOpen.async_stream
+    streamed = []
+
+    def track(self, keys, **kwargs):
+        streamed.extend(keys)
+        yield from original_async(self, keys, **kwargs)
+
+    monkeypatch.setattr(MemoryEfficientSafeOpen, "async_stream", track)
+    output = resize.resize_lora_file(
+        str(source), 1, None, None, "cpu", torch.bfloat16, "filtered", verbose=False,
+        exclude_patterns="*.keep.lora_A.weight\n*.drop.*" if glob else r"\.keep\.lora_A\.weight$" + "\n" + r"\.drop\.",
+        discard_patterns="*.drop.alpha\nstandalone" if glob else r"\.drop\.alpha$" + "\n^standalone$",
+        glob_patterns=glob,
+    )
+    assert not any(".drop." in key or key == "standalone" for key in streamed)
+    actual = _read_uel(output, original_async)
+    assert not any(".drop." in key or key == "standalone" for key in actual)
+    for key in tensors:
+        if ".keep." in key:
+            assert actual[key].dtype == tensors[key].dtype
+            assert torch.equal(actual[key], tensors[key])
+    assert actual["diffusion_model.resize.lora_A.weight"].shape[0] == 1
+
+
+@pytest.mark.parametrize("glob", [False, True])
+def test_resize_include_mode_only_resizes_matching_layers(monkeypatch, tmp_path, resize, glob):
+    source = tmp_path / "include_filters.safetensors"
+    tensors = {}
+    for name in ("include", "discard", "skip"):
+        tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3, dtype=torch.float16)
+        tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3, dtype=torch.float16)
+        tensors[f"diffusion_model.{name}.alpha"] = torch.tensor(1.5)
+    tensors["standalone.include"] = torch.ones(1)
+    tensors["standalone.skip"] = torch.ones(1)
+    _write_uel(source, tensors)
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.resize_lora_file(
+        str(source), 1, None, None, "cpu", torch.bfloat16, "included", verbose=False,
+        exclude_patterns="*include*\n*discard*\nstandalone.include" if glob else r"\.(include|discard)\." + "\n^standalone\\.include$",
+        discard_patterns="*discard*" if glob else r"\.discard\.",
+        glob_patterns=glob,
+        include_mode=True,
+    )
+
+    actual = _read_uel(output)
+    assert actual["diffusion_model.include.lora_A.weight"].shape[0] == 1
+    assert torch.equal(actual["standalone.include"], tensors["standalone.include"])
+    assert not any(".discard." in key for key in actual)
+    for key in tensors:
+        if ".skip." in key or key == "standalone.skip":
+            assert torch.equal(actual[key], tensors[key])
+
+
+def test_resize_include_mode_empty_patterns_preserves_all_tensors(monkeypatch, tmp_path, resize):
+    source = tmp_path / "empty_include.safetensors"
+    tensors = {"standalone": torch.ones(2)}
+    _write_uel(source, tensors)
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.resize_lora_file(
+        str(source), 1, None, None, "cpu", torch.float32, "empty", verbose=False,
+        include_mode=True,
+    )
+
+    actual = _read_uel(output)
+    assert torch.equal(actual["standalone"], tensors["standalone"])
+
+
+@pytest.mark.parametrize("node_name,extra", [
+    ("LoRAResizeFixed", {"new_rank": 1}),
+    ("LoRAResizeRatio", {"max_rank": 1, "ratio": 2.0}),
+    ("LoRAResizeFrobenius", {"max_rank": 1, "min_rank": 1, "target": 0.9}),
+    ("LoRAResizeCumulative", {"max_rank": 1, "target": 0.9}),
+])
+def test_all_resize_nodes_forward_filters(monkeypatch, tmp_path, resize, node_name, extra):
+    node = getattr(resize, node_name)
+    controls = {item.id: item for item in node.define_schema().inputs}
+    assert controls["exclude_patterns"].default == ""
+    assert controls["discard_patterns"].default == ""
+    assert controls["glob_patterns"].default is False
+    assert controls["include_mode"].default is False
+    assert [item.id for item in node.define_schema().inputs if item.id != "layer_parameters"][-1] == "include_mode"
+    calls = []
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize.folder_paths, "get_full_path_or_raise", lambda *args: "input")
+    monkeypatch.setattr(resize, "resize_lora_file", lambda *args, **kwargs: calls.append(kwargs))
+    node.execute(lora_name="input", output_filename="out", save_dtype="fp16", device="cpu",
+                 force_clear_cache=False, exclude_patterns="keep*", discard_patterns="drop*",
+                 glob_patterns=True, include_mode=True, **extra)
+    assert calls[0]["exclude_patterns"] == "keep*"
+    assert calls[0]["discard_patterns"] == "drop*"
+    assert calls[0]["glob_patterns"] is True
+    assert calls[0]["include_mode"] is True
+
+
+def test_resize_invalid_filter_fails_before_writer(monkeypatch, tmp_path, resize):
+    import re
+
+    source = tmp_path / "invalid_filter.safetensors"
+    _write_uel(source, {"layer.diff": torch.ones(2)})
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    def unexpected_writer(*args, **kwargs):
+        pytest.fail("Writer must not open for invalid patterns")
+
+    monkeypatch.setattr(resize, "atomic_uel_writer", unexpected_writer)
+    with pytest.raises(re.error):
+        resize.resize_lora_file(str(source), 1, None, None, "cpu", torch.float32,
+                                "unused", exclude_patterns="[")
+
+
+def _dense_truncated(up, down, rank):
+    dense = up.float() @ down.float()
+    u, s, vh = torch.linalg.svd(dense, full_matrices=False)
+    return (u[:, :rank] * s[:rank]) @ vh[:rank, :]
+
+
+def test_linear_factor_space_matches_dense_reference_and_bounds_svd(
+    monkeypatch, resize
+):
+    generator = torch.Generator().manual_seed(123)
+    down = torch.randn((4, 17), generator=generator)
+    up = torch.randn((13, 4), generator=generator)
+    expected = _dense_truncated(up, down, 3)
+    observed_shapes = []
+    original_svd = torch.linalg.svd
+
+    def track_svd(tensor, *args, **kwargs):
+        observed_shapes.append(tuple(tensor.shape))
+        return original_svd(tensor, *args, **kwargs)
+
+    monkeypatch.setattr(torch.linalg, "svd", track_svd)
+    result = resize._resize_lora_factors(down, up, 3, None, None, 1.0)
+
+    actual = result["lora_up"] @ result["lora_down"]
+    assert torch.allclose(actual, expected, atol=2e-5, rtol=2e-5)
+    assert observed_shapes == [(4, 4)]
+
+
+def test_convolutional_factor_space_matches_dense_reference(resize):
+    generator = torch.Generator().manual_seed(456)
+    down = torch.randn((3, 2, 3, 2), generator=generator)
+    up = torch.randn((5, 3, 1, 1), generator=generator)
+
+    result = resize._resize_lora_factors(down, up, 2, None, None, 1.0)
+
+    actual = result["lora_up"].reshape(5, 2) @ result["lora_down"].reshape(2, -1)
+    expected = _dense_truncated(up.reshape(5, 3), down.reshape(3, -1), 2)
+    assert result["lora_up"].shape == (5, 2, 1, 1)
+    assert result["lora_down"].shape == (2, 2, 3, 2)
+    assert torch.allclose(actual, expected, atol=2e-5, rtol=2e-5)
+
+
+@pytest.mark.parametrize(
+    ("method", "parameter", "maximum", "expected_rank"),
+    [
+        (None, None, 3, 3),
+        ("sv_ratio", 4.0, 4, 2),
+        ("sv_fro", 0.95, 4, 1),
+        ("sv_cumulative", 0.8, 4, 2),
+    ],
+)
+def test_all_resize_modes_use_the_shared_spectrum(
+    resize, method, parameter, maximum, expected_rank
+):
+    singular_values = torch.tensor([10.0, 3.0, 1.0, 0.1])
+    result = resize._resize_lora_factors(
+        torch.eye(4),
+        torch.diag(singular_values),
+        maximum,
+        method,
+        parameter,
+        1.0,
+    )
+    assert result["new_rank"] == expected_rank
+
+
+def test_mixed_source_ranks_are_capped_per_layer(resize):
+    rank_128 = resize._resize_lora_factors(
+        torch.eye(128), torch.eye(128), 384, None, None, 1.0
+    )
+    rank_384 = resize._resize_lora_factors(
+        torch.eye(384), torch.eye(384), 384, None, None, 1.0
+    )
+
+    assert rank_128["new_rank"] == 128
+    assert rank_384["new_rank"] == 384
+
+
+def test_resize_file_uses_only_async_uel_and_preserves_per_layer_scale(
+    monkeypatch, tmp_path, resize
+):
+    source = tmp_path / "mixed_rank.safetensors"
+    _write_uel(
+        source,
+        {
+            "diffusion_model.r2.lora_A.weight": torch.eye(2),
+            "diffusion_model.r2.lora_B.weight": torch.eye(2),
+            "diffusion_model.r2.alpha": torch.tensor(1.0),
+            "diffusion_model.r3.lora_A.weight": torch.eye(3),
+            "diffusion_model.r3.lora_B.weight": torch.eye(3),
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    original_async = resize.MemoryEfficientSafeOpen.async_stream
+    original_mark = resize.MemoryEfficientSafeOpen.mark_processed
+    async_calls = []
+    marked = []
+
+    def track_async(self, keys, **kwargs):
+        async_calls.append((tuple(keys), kwargs, self.low_memory))
+        yield from original_async(self, keys, **kwargs)
+
+    def track_mark(self, key):
+        marked.append(key)
+        return original_mark(self, key)
+
+    def reject_sync_get(*args, **kwargs):
+        raise AssertionError("resize used synchronous get_tensor")
+
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async)
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "mark_processed", track_mark)
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get)
+
+    output = resize.resize_lora_file(
+        str(source), 3, None, None, "cpu", torch.float32, "mixed_rank_output",
+        verbose=False,
+    )
+
+    assert len(async_calls) == 1
+    streamed_keys, options, low_memory = async_calls[0]
+    assert low_memory is True
+    assert options == {
+        "batch_size": 1,
+        "prefetch_batches": 1,
+        "pin_memory": False,
+    }
+    assert Counter(marked) == Counter(streamed_keys)
+    assert all(count == 1 for count in Counter(marked).values())
+
+    output_tensors = _read_uel(output, original_async, original_mark)
+    assert output_tensors["diffusion_model.r2.lora_A.weight"].shape[0] == 2
+    assert output_tensors["diffusion_model.r3.lora_A.weight"].shape[0] == 3
+    assert output_tensors["diffusion_model.r2.alpha"].item() == pytest.approx(1.0)
+    assert "diffusion_model.r3.alpha" not in output_tensors
+
+
+def test_normalize_alpha_uses_bounded_uel_and_removes_alpha(
+    monkeypatch, tmp_path, resize
+):
+    source = tmp_path / "alpha_input.safetensors"
+    down = torch.eye(2, dtype=torch.float16)
+    up = torch.tensor([[2.0, 4.0], [6.0, 8.0]], dtype=torch.float16)
+    _write_uel(
+        source,
+        {
+            "diffusion_model.block.lora_A.weight": down,
+            "diffusion_model.block.lora_B.weight": up,
+            "diffusion_model.block.alpha": torch.tensor(1.0),
+            "other.weight": torch.tensor([3.0]),
+        },
+        metadata={"source": "test"},
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    original_async = resize.MemoryEfficientSafeOpen.async_stream
+    original_mark = resize.MemoryEfficientSafeOpen.mark_processed
+    async_calls = []
+    marked = []
+
+    def track_async(self, keys, **kwargs):
+        async_calls.append((tuple(keys), kwargs, self.low_memory))
+        yield from original_async(self, keys, **kwargs)
+
+    def track_mark(self, key):
+        marked.append(key)
+        return original_mark(self, key)
+
+    def reject_sync_get(*args, **kwargs):
+        raise AssertionError("alpha normalization used synchronous get_tensor")
+
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async)
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "mark_processed", track_mark)
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get)
+
+    output = resize.normalize_lora_alpha_file(
+        str(source), "alpha_normalized", verbose=False
+    )
+
+    assert len(async_calls) == 1
+    streamed_keys, options, low_memory = async_calls[0]
+    assert streamed_keys == (
+        "diffusion_model.block.lora_B.weight",
+        "diffusion_model.block.alpha",
+    )
+    assert options == {
+        "batch_size": 1,
+        "prefetch_batches": 1,
+        "pin_memory": False,
+    }
+    assert low_memory is True
+    assert Counter(marked) == Counter(streamed_keys)
+
+    output_tensors = _read_uel(output, original_async, original_mark)
+    assert "diffusion_model.block.alpha" not in output_tensors
+    assert torch.equal(output_tensors["diffusion_model.block.lora_A.weight"], down)
+    assert torch.equal(
+        output_tensors["diffusion_model.block.lora_B.weight"], up * 0.5
+    )
+    assert torch.equal(output_tensors["other.weight"], torch.tensor([3.0]))
+    with MemoryEfficientSafeOpen(output, low_memory=True) as loader:
+        assert loader.metadata()["alpha_normalized"] == "true"
+
+
+def test_normalize_alpha_rejects_alpha_free_input(monkeypatch, tmp_path, resize):
+    source = tmp_path / "no_alpha.safetensors"
+    _write_uel(
+        source,
+        {
+            "diffusion_model.block.lora_A.weight": torch.eye(2),
+            "diffusion_model.block.lora_B.weight": torch.eye(2),
+        },
+    )
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    with pytest.raises(ValueError, match="contains no alpha tensors"):
+        resize.normalize_lora_alpha_file(
+            str(source), "unused", verbose=False
+        )
+
+
+def test_normalize_alpha_maps_flattened_names_with_reference(
+    monkeypatch, tmp_path, resize
+):
+    source = tmp_path / "flattened_alpha.safetensors"
+    reference = tmp_path / "reference.safetensors"
+    down = torch.eye(2, dtype=torch.float16)
+    up = torch.tensor([[2.0, 4.0], [6.0, 8.0]], dtype=torch.float16)
+    _write_uel(
+        source,
+        {
+            "lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight": down,
+            "lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight": up,
+            "lora_unet_double_blocks_0_img_attn_qkv.alpha": torch.tensor(1.0),
+        },
+    )
+    _write_uel(
+        reference,
+        {"diffusion_model.double_blocks.0.img_attn.qkv.weight": torch.eye(2)},
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.normalize_lora_alpha_file(
+        str(source),
+        "reference_mapped_alpha",
+        verbose=False,
+        reference_model_path=str(reference),
+    )
+
+    output_tensors = _read_uel(output)
+    assert set(output_tensors) == {
+        "diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight",
+        "diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight",
+    }
+    assert torch.equal(
+        output_tensors["diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight"],
+        down,
+    )
+    assert torch.equal(
+        output_tensors["diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight"],
+        up * 0.5,
+    )
+
+
+def test_resize_does_not_invent_alpha_for_peft_pair(monkeypatch, tmp_path, resize):
+    source = tmp_path / "peft_no_alpha.safetensors"
+    _write_uel(
+        source,
+        {
+            "base_model.model.transformer.block.to_q.lora_A.default.weight": (
+                torch.eye(2)
+            ),
+            "base_model.model.transformer.block.to_q.lora_B.default.weight": (
+                torch.eye(2)
+            ),
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    output = resize.resize_lora_file(
+        str(source), 2, None, None, "cpu", torch.float32, "peft_output",
+        verbose=False,
+    )
+
+    output_tensors = _read_uel(output)
+    assert len(output_tensors) == 2
+    assert not any(key.endswith(".alpha") for key in output_tensors)
+
+
+def test_cuda_oom_retries_only_current_pair_on_cpu(monkeypatch, resize):
+    calls = []
+    original = resize._resize_factors_on_device
+
+    def fail_cuda_then_run_cpu(*args):
+        process_device = args[-1]
+        calls.append(process_device)
+        if process_device == "cuda":
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory")
+        return original(*args)
+
+    monkeypatch.setattr(resize, "_resize_factors_on_device", fail_cuda_then_run_cpu)
+    monkeypatch.setattr(resize, "_release_cuda_after_oom", lambda: None)
+
+    result = resize._resize_lora_factors(
+        torch.eye(2), torch.eye(2), 1, None, None, 1.0,
+        process_device="cuda",
+    )
+
+    assert calls == ["cuda", "cpu"]
+    assert result["cpu_fallback"] is True
+
+
+def test_low_bit_pair_uses_raw_copy_without_tensor_loading(
+    monkeypatch, tmp_path, resize
+):
+    source = tmp_path / "low_bit_pair.safetensors"
+    down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
+    up = torch.eye(2)
+    _write_uel(
+        source,
+        {
+            "diffusion_model.low.lora_A.weight": down,
+            "diffusion_model.low.lora_B.weight": up,
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+    async_calls = []
+    original_async = resize.MemoryEfficientSafeOpen.async_stream
+    original_mark = resize.MemoryEfficientSafeOpen.mark_processed
+
+    def track_async(self, keys, **kwargs):
+        async_calls.append(tuple(keys))
+        yield from original_async(self, keys, **kwargs)
+
+    def reject_sync_get(*args, **kwargs):
+        raise AssertionError("low-bit raw copy decoded a tensor")
+
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async)
+    monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get)
+
+    output = resize.resize_lora_file(
+        str(source), 1, None, None, "cpu", torch.float16, "low_bit_output",
+        verbose=False,
+    )
+
+    assert async_calls == []
+    output_tensors = _read_uel(output, original_async, original_mark)
+    assert torch.equal(output_tensors["diffusion_model.low.lora_A.weight"], down)
+    assert torch.equal(output_tensors["diffusion_model.low.lora_B.weight"], up)
+
+
+def test_resize_node_interfaces_remove_obsolete_loader_and_iteration_controls(resize):
+    schemas = [
+        resize.LoRAResizeFixed.define_schema(),
+        resize.LoRAResizeRatio.define_schema(),
+        resize.LoRAResizeFrobenius.define_schema(),
+        resize.LoRAResizeCumulative.define_schema(),
+    ]
+    input_ids = [[item.id for item in schema.inputs] for schema in schemas]
+
+    assert "svd_niter" not in input_ids[0]
+    assert all("lazy_load" not in ids for ids in input_ids)
+    assert input_ids[2][:4] == ["lora_name", "max_rank", "min_rank", "target"]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
+def test_cuda_and_cpu_factor_space_results_match(resize):
+    generator = torch.Generator().manual_seed(789)
+    down = torch.randn((4, 11), generator=generator)
+    up = torch.randn((9, 4), generator=generator)
+
+    cpu = resize._resize_lora_factors(down, up, 3, None, None, 1.0, process_device="cpu")
+    cuda = resize._resize_lora_factors(down, up, 3, None, None, 1.0, process_device="cuda")
+
+    cpu_update = cpu["lora_up"] @ cpu["lora_down"]
+    cuda_update = cuda["lora_up"] @ cuda["lora_down"]
+    assert torch.allclose(cpu_update, cuda_update, atol=2e-4, rtol=2e-4)
+
+
+def test_modify_lora_alpha_handles_present_and_absent_alpha(monkeypatch, tmp_path, resize):
+    source = tmp_path / "mixed_alpha.safetensors"
+    down_a = torch.eye(4, dtype=torch.float16)
+    up_a = torch.ones((4, 4), dtype=torch.float16)
+    down_b = torch.eye(2, dtype=torch.float16)
+    up_b = torch.ones((2, 2), dtype=torch.float16) * 3.0
+    _write_uel(
+        source,
+        {
+            "diffusion_model.block_a.lora_A.weight": down_a,
+            "diffusion_model.block_a.lora_B.weight": up_a,
+            "diffusion_model.block_a.alpha": torch.tensor(2.0, dtype=torch.float16),
+            "diffusion_model.block_b.lora_A.weight": down_b,
+            "diffusion_model.block_b.lora_B.weight": up_b,
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    scaled = resize.modify_lora_alpha_file(
+        str(source), "scaled_alpha", alpha_mode="scale", alpha_value=1.5, verbose=False
+    )
+    scaled_tensors = _read_uel(scaled)
+    # Factor weights are byte-exact unchanged
+    assert torch.equal(scaled_tensors["diffusion_model.block_a.lora_A.weight"], down_a)
+    assert torch.equal(scaled_tensors["diffusion_model.block_a.lora_B.weight"], up_a)
+    assert torch.equal(scaled_tensors["diffusion_model.block_b.lora_A.weight"], down_b)
+    assert torch.equal(scaled_tensors["diffusion_model.block_b.lora_B.weight"], up_b)
+    # block_a had alpha=2.0 -> 2.0 * 1.5 = 3.0; block_b had no alpha (rank=2) -> 2.0 * 1.5 = 3.0
+    assert scaled_tensors["diffusion_model.block_a.alpha"].item() == pytest.approx(3.0)
+    assert scaled_tensors["diffusion_model.block_b.alpha"].item() == pytest.approx(3.0)
+
+    # Shifted with filter excluding block_a
+    shifted = resize.modify_lora_alpha_file(
+        str(source),
+        "shifted_alpha",
+        alpha_mode="shift",
+        alpha_value=2.0,
+        exclude_patterns="block_a",
+        verbose=False,
+    )
+    shifted_tensors = _read_uel(shifted)
+    assert shifted_tensors["diffusion_model.block_a.alpha"].item() == pytest.approx(2.0)
+    assert shifted_tensors["diffusion_model.block_b.alpha"].item() == pytest.approx(4.0)
+
+
+def test_denormalize_lora_alpha_is_inverse_of_normalize(monkeypatch, tmp_path, resize):
+    source = tmp_path / "alpha_lora.safetensors"
+    down = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
+    up = torch.tensor([[2.0, 0.5], [1.0, 3.0]], dtype=torch.float32)
+    _write_uel(
+        source,
+        {
+            "diffusion_model.block.lora_A.weight": down,
+            "diffusion_model.block.lora_B.weight": up,
+            "diffusion_model.block.alpha": torch.tensor(1.0, dtype=torch.float32),
+        },
+    )
+    monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
+
+    normalized_path = resize.normalize_lora_alpha_file(
+        str(source), "norm_step", verbose=False
+    )
+    denormalized_path = resize.denormalize_lora_alpha_file(
+        str(normalized_path),
+        "denorm_step",
+        alpha_mode="set_alpha",
+        alpha_value=1.0,
+        verbose=False,
+    )
+    denorm_tensors = _read_uel(denormalized_path)
+    assert denorm_tensors["diffusion_model.block.alpha"].item() == pytest.approx(1.0)
+    assert torch.allclose(denorm_tensors["diffusion_model.block.lora_A.weight"], down)
+    assert torch.allclose(denorm_tensors["diffusion_model.block.lora_B.weight"], up)
+
+    with pytest.raises(ValueError, match="non-zero"):
+        resize.denormalize_lora_alpha_file(
+            str(normalized_path), "bad_zero", alpha_mode="set_alpha", alpha_value=0.0, verbose=False
+        )

tests/test_lora_resize_min_rank.py

@@ -0,0 +1,89 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_resize_min_rank_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 resize():
+    return _load_module("nodes.lora_resize")
+
+
+def test_frobenius_minimum_rank_is_enforced(resize):
+    singular_values = torch.tensor([10.0, 1.0, 0.1, 0.01, 0.001])
+
+    rank, _, _ = resize._compute_resize(
+        singular_values, 4, "sv_fro", 0.5, 1.0, min_rank=3
+    )
+
+    assert rank == 3
+
+
+def test_default_minimum_rank_remains_one(resize):
+    singular_values = torch.tensor([10.0, 1.0, 0.1, 0.01, 0.001])
+
+    rank, _, _ = resize._compute_resize(
+        singular_values, 4, "sv_fro", 0.5, 1.0
+    )
+
+    assert rank == 1
+
+
+def test_factor_space_resize_passes_minimum_rank_to_dynamic_selection(resize):
+    down = torch.eye(5)
+    up = torch.diag(torch.tensor([10.0, 1.0, 0.1, 0.01, 0.001]))
+
+    result = resize._resize_lora_factors(
+        down, up, 4, "sv_fro", 0.5, 1.0, min_rank=3
+    )
+
+    assert result["new_rank"] == 3
+    assert result["lora_down"].shape[0] == 3
+
+
+def test_minimum_rank_is_capped_by_layer_dimensions(resize):
+    singular_values = torch.tensor([10.0, 1.0, 0.1])
+
+    rank, _, _ = resize._compute_resize(
+        singular_values, 4, "sv_fro", 0.5, 1.0, min_rank=4
+    )
+
+    assert rank == 3
+
+
+def test_frobenius_schema_places_minimum_beside_maximum(resize):
+    schema = resize.LoRAResizeFrobenius.define_schema()
+
+    assert [item.id for item in schema.inputs[:4]] == [
+        "lora_name", "max_rank", "min_rank", "target"
+    ]
+    assert schema.inputs[2].default == 1
+
+
+def test_resize_rejects_minimum_above_maximum(tmp_path, resize):
+    with pytest.raises(ValueError, match="min_rank cannot exceed max rank"):
+        resize.resize_lora_file(
+            str(tmp_path / "unused.safetensors"),
+            4,
+            "sv_fro",
+            0.999,
+            "cpu",
+            torch.float32,
+            "unused",
+            min_rank=5,
+        )

tests/test_merger_filters.py

@@ -0,0 +1,284 @@
+import importlib
+import json
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_merger_filter_tests"
+
+
+def _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 _module("nodes.merger"), _module("nodes.merger_ops")
+
+
+def _params(name, *, patterns="", include=False, discard="", glob=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": name, "save_dtype": "fp32",
+        "override_dtype": True, "device": "cpu", "dtype": torch.float32,
+        "exclude_patterns": patterns, "include_mode": include,
+        "discard_patterns": discard, "glob_patterns": glob, "lazy_load": True,
+        "force_clear_cache": False,
+    }
+
+
+def _run(monkeypatch, tmp_path, merger, operations, params):
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    save_file({"match.weight": torch.tensor([2.0]), "other.weight": torch.tensor([4.0])}, paths["a"])
+    save_file({"match.weight": torch.tensor([6.0]), "other.weight": torch.tensor([8.0])}, paths["b"])
+    monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
+    monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+    params.update({"model_a": "a", "model_b": "b"})
+    merger.MergerLogic.execute_merge({"model_a": "a", "model_b": "b"}, "Weight-Sum", operations.TWO_MODEL_MODES, params, "checkpoints")
+    return load_file(str(tmp_path / "checkpoints" / f"{params['output_filename']}.safetensors"))
+
+
+def test_include_filter_controls_preflight_and_streaming(monkeypatch, tmp_path, modules):
+    merger, operations = modules
+    selected = _run(monkeypatch, tmp_path, merger, operations, _params("selected", patterns="match", include=True))
+    assert selected["match.weight"].item() == 4.0
+    assert selected["other.weight"].item() == 4.0
+
+    empty = _run(monkeypatch, tmp_path, merger, operations, _params("empty", include=True))
+    assert empty["match.weight"].item() == 2.0
+    assert empty["other.weight"].item() == 4.0
+
+    discarded = _run(monkeypatch, tmp_path, merger, operations, _params("discarded", patterns="match", include=True, discard="match"))
+    assert "match.weight" not in discarded
+    assert discarded["other.weight"].item() == 4.0
+
+    globbed = _run(monkeypatch, tmp_path, merger, operations, _params("globbed", patterns="*.weight", include=True, glob=True))
+    assert globbed["match.weight"].item() == 4.0
+    assert globbed["other.weight"].item() == 6.0
+
+
+def test_layer_parameters_apply_per_layer_without_leakage(monkeypatch, tmp_path, modules):
+    merger, operations = modules
+
+    def resolve(_payload, operation, layers, defaults, *, node_name):
+        assert operation == "merge:Weight-Sum"
+        assert node_name == "checkpoints Weight-Sum merger"
+        assert set(layers) == {"match.weight", "other.weight"}
+        assert defaults["alpha"] == 0.5
+        return {"match.weight": {"alpha": 0.0}, "other.weight": {"alpha": 1.0}}
+
+    monkeypatch.setattr(merger, "resolve_layer_parameters", resolve)
+    result = _run(monkeypatch, tmp_path, merger, operations, _params("per_layer"))
+    assert result["match.weight"].item() == 2.0
+    assert result["other.weight"].item() == 8.0
+
+
+def test_parsed_layer_parameters_change_only_matching_output(monkeypatch, tmp_path, modules):
+    merger, operations = modules
+    layer_parameters = _module("nodes.layer_parameters")
+    params = _params("parsed")
+    params["layer_parameters"] = layer_parameters.parse_rules("(match\\.weight) alpha:0\n")
+    result = _run(monkeypatch, tmp_path, merger, operations, params)
+    assert result["match.weight"].item() == 2.0
+    assert result["other.weight"].item() == 6.0
+
+
+def test_layer_parameter_failure_precedes_writer_and_streaming(monkeypatch, tmp_path, modules):
+    merger, operations = modules
+    monkeypatch.setattr(
+        merger,
+        "resolve_layer_parameters",
+        lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("unmatched rule line 1")),
+    )
+    params = _params("must_not_exist")
+    with pytest.raises(ValueError, match="unmatched rule"):
+        _run(monkeypatch, tmp_path, merger, operations, params)
+    assert not (tmp_path / "checkpoints" / "must_not_exist.safetensors").exists()
+
+
+def test_lora_factor_group_uses_one_normalized_layer_name(monkeypatch, tmp_path, modules):
+    merger, operations = modules
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    tensors = {
+        "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
+        "diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [2.0]]),
+    }
+    save_file({**tensors, "diffusion_model.layer.alpha": torch.tensor(1.0)}, paths["a"])
+    save_file(tensors, paths["b"])
+    monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
+    monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+    params = _params("lora_group")
+    params.update({"model_a": "a", "model_b": "b", "layer_parameters": _module("nodes.layer_parameters").parse_rules("(diffusion_model\\.layer) alpha:0\n")})
+    merger.MergerLogic.execute_merge({"model_a": "a", "model_b": "b"}, "Weight-Sum", operations.TWO_MODEL_MODES, params, "loras")
+    result = load_file(str(tmp_path / "loras" / "lora_group.safetensors"))
+    assert "diffusion_model.layer.lora_A.weight" in result
+    assert "diffusion_model.layer.lora_B.weight" in result
+    assert not any(key.endswith(".alpha") for key in result)
+
+
+def test_diffusion_merge_decodes_tensorwise_quantization_and_drops_sidecars(
+    monkeypatch, tmp_path, modules,
+):
+    merger, operations = modules
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    save_file({
+        "model.diffusion_model.quant.weight": torch.tensor([1, 2], dtype=torch.int8),
+        "model.diffusion_model.quant.weight_scale": torch.tensor(2.0),
+        "model.diffusion_model.quant.comfy_quant": quant,
+        "dense.weight": torch.tensor([2.0, 4.0]),
+    }, paths["a"])
+    save_file({
+        "model.diffusion_model.quant.weight": torch.tensor([3, 6], dtype=torch.int8),
+        "model.diffusion_model.quant.weight_scale": torch.tensor(2.0),
+        "model.diffusion_model.quant.comfy_quant": quant,
+        "dense.weight": torch.tensor([6.0, 8.0]),
+    }, paths["b"])
+
+    monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
+    monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+    params = _params("quantized_merge")
+    params.update({"model_a": "a", "model_b": "b"})
+
+    merger.MergerLogic.execute_merge(
+        {"model_a": "a", "model_b": "b"}, "Weight-Sum",
+        operations.TWO_MODEL_MODES, params, "diffusion_models",
+    )
+
+    output_path = tmp_path / "diffusion_models" / "quantized_merge.safetensors"
+    result = load_file(str(output_path))
+    key = "quant.weight"
+    assert set(result) == {key, "dense.weight"}
+    assert result[key].tolist() == pytest.approx([4.0, 8.0])
+    assert result[key].dtype == torch.float32
+    assert result["dense.weight"].tolist() == pytest.approx([4.0, 6.0])
+    with safe_open(str(output_path), framework="pt", device="cpu") as output:
+        metadata = output.metadata()
+    assert metadata in (None, {})
+
+
+@pytest.mark.parametrize("prefixed_input", ["a", "b"])
+def test_diffusion_merge_matches_prefixed_and_bare_models(monkeypatch, tmp_path, modules, prefixed_input):
+    merger, operations = modules
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    for name in ("a", "b"):
+        key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
+        if name == prefixed_input:
+            stem = key[:-len(".weight")]
+            save_file({
+                key: torch.tensor([1, 2], dtype=torch.int8),
+                f"{stem}.weight_scale": torch.tensor(2.0),
+                f"{stem}.comfy_quant": torch.tensor(
+                    list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
+                ),
+            }, paths[name])
+        else:
+            save_file({key: torch.tensor([4.0, 8.0])}, paths[name])
+    monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
+    monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+    params = _params("mixed_prefix_merge")
+    params.update({"model_a": "a", "model_b": "b"})
+    merger.MergerLogic.execute_merge(
+        {"model_a": "a", "model_b": "b"}, "Weight-Sum",
+        operations.TWO_MODEL_MODES, params, "diffusion_models",
+    )
+    output = load_file(str(tmp_path / "diffusion_models" / "mixed_prefix_merge.safetensors"))
+    assert list(output) == ["layer.weight"]
+    torch.testing.assert_close(output["layer.weight"], torch.tensor([3.0, 6.0]))
+
+
+def test_three_diffusion_models_match_across_prefixes(monkeypatch, tmp_path, modules):
+    merger, operations = modules
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b", "c")}
+    save_file({"layer.weight": torch.tensor([1.0, 1.0])}, paths["a"])
+    save_file({"model.diffusion_model.layer.weight": torch.tensor([4.0, 4.0])}, paths["b"])
+    save_file({"layer.weight": torch.tensor([2.0, 2.0])}, paths["c"])
+    monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
+    monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+    params = _params("three_prefix_merge")
+    params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
+    merger.MergerLogic.execute_merge(
+        {"model_a": "a", "model_b": "b", "model_c": "c"}, "Add-Difference",
+        operations.THREE_MODEL_MODES, params, "diffusion_models",
+    )
+    output = load_file(str(tmp_path / "diffusion_models" / "three_prefix_merge.safetensors"))
+    assert list(output) == ["layer.weight"]
+    torch.testing.assert_close(output["layer.weight"], torch.tensor([2.0, 2.0]))
+
+
+def test_prunekeys_include_mode_and_uel_streaming(monkeypatch, tmp_path):
+    prunekeys = _module("nodes.prunekeys")
+    for cls in (
+        prunekeys.ModelPruneKeys,
+        prunekeys.TextEncoderPruneKeys,
+        prunekeys.LoRAPruneKeys,
+        prunekeys.CheckpointPruneKeys,
+        prunekeys.EmbeddingPruneKeys,
+    ):
+        schema = cls.define_schema()
+        assert schema.inputs[-1].id == "include_mode"
+        assert schema.inputs[-1].default is False
+
+    src = tmp_path / "source.safetensors"
+    save_file(
+        {
+            "layers.0.attn.weight": torch.tensor([1.0, 2.0]),
+            "layers.0.ffn.weight": torch.tensor([3.0, 4.0]),
+            "layers.1.attn.weight": torch.tensor([5.0, 6.0]),
+        },
+        str(src),
+        metadata={"author": "test"},
+    )
+    monkeypatch.setattr(prunekeys.folder_paths, "get_full_path_or_raise", lambda _kind, _name: str(src))
+    monkeypatch.setattr(prunekeys.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(prunekeys, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(prunekeys, "cleanup_after_operation", lambda: None)
+
+    # Default include_mode=False removes matching keys
+    rel_removed = prunekeys.ModelPruneKeys.execute(
+        diffusionmodel_name="source.safetensors",
+        keys_to_prune="attn",
+        use_regex=False,
+        output_filename="removed_attn",
+        include_mode=False,
+    ).result[0]
+    removed_tensors = load_file(str(tmp_path / "diffusion_models" / rel_removed))
+    assert set(removed_tensors) == {"layers.0.ffn.weight"}
+
+    # include_mode=True keeps only matching keys
+    rel_kept = prunekeys.ModelPruneKeys.execute(
+        diffusionmodel_name="source.safetensors",
+        keys_to_prune=r"^layers\.\d+\.attn\.",
+        use_regex=True,
+        output_filename="kept_attn",
+        include_mode=True,
+    ).result[0]
+    kept_tensors = load_file(str(tmp_path / "diffusion_models" / rel_kept))
+    assert set(kept_tensors) == {"layers.0.attn.weight", "layers.1.attn.weight"}

tests/test_minimax_h3_fold.py

@@ -0,0 +1,246 @@
+import importlib
+from pathlib import Path
+import sys
+import types
+
+import pytest
+from safetensors.torch import load_file, save_file
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_minimax_fold_tests"
+
+
+def _load_module(name: str):
+    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 folder_module():
+    return _load_module("nodes.minimax_h3_fold")
+
+
+def _make_h3_checkpoint(prefix: str = "") -> dict[str, torch.Tensor]:
+    return {
+        f"{prefix}time_embedder.proj_in.weight": torch.eye(16, dtype=torch.float32),
+        f"{prefix}time_embedder.proj_in.bias": torch.zeros(16, dtype=torch.float32),
+        f"{prefix}time_embedder.proj_out.weight": torch.ones((2688, 16), dtype=torch.float32) * 0.01,
+        f"{prefix}time_embedder.proj_out.bias": torch.zeros(2688, dtype=torch.float32),
+        f"{prefix}blocks.0.adaln_proj.linear.weight": torch.ones((128, 2688), dtype=torch.float32),
+        f"{prefix}blocks.0.adaln_proj.linear.bias": torch.zeros(128, dtype=torch.float32),
+        f"{prefix}blocks.1.adaln_proj.linear.weight": torch.ones((128, 2688), dtype=torch.float32) * 2.0,
+        f"{prefix}blocks.1.adaln_proj.linear.bias": torch.zeros(128, dtype=torch.float32),
+        f"{prefix}blocks.0.attn.qkv_proj.weight": torch.eye(64, dtype=torch.float32),
+    }
+
+
+def test_fold_h3_diffusion_model_exact_basis(folder_module, tmp_path):
+    source = tmp_path / "model_full.safetensors"
+    output = tmp_path / "model_folded.safetensors"
+
+    original = _make_h3_checkpoint()
+    save_file(original, str(source))
+
+    report = folder_module.fold_minimax_h3_diffusion_model(
+        str(source), str(output), process_device="cpu"
+    )
+    assert "MiniMax H3 AdaLN Fold complete" in report
+    assert "AdaLN layers folded: 4" in report
+
+    converted = load_file(str(output))
+
+    assert "adaln_t_table" in converted
+    assert converted["adaln_t_table"].shape == (1025, 8)
+    assert converted["adaln_t_table"].dtype == torch.float32
+
+    # Folded AdaLN weight and bias
+    assert converted["blocks.0.adaln_proj.linear.weight"].shape == (128, 8)
+    assert converted["blocks.0.adaln_proj.linear.weight"].dtype == torch.float16
+    assert converted["blocks.0.adaln_proj.linear.bias"].dtype == torch.float16
+
+    # time_embedder should be dropped
+    assert "time_embedder.proj_in.weight" not in converted
+    assert "time_embedder.proj_out.weight" not in converted
+
+    # Unrelated layers preserved
+    assert "blocks.0.attn.qkv_proj.weight" in converted
+    assert torch.equal(converted["blocks.0.attn.qkv_proj.weight"], original["blocks.0.attn.qkv_proj.weight"])
+
+
+def test_fold_h3_discard_patterns(folder_module, tmp_path):
+    source = tmp_path / "model_discard.safetensors"
+    output = tmp_path / "model_discarded.safetensors"
+
+    original = _make_h3_checkpoint()
+    save_file(original, str(source))
+
+    folder_module.fold_minimax_h3_diffusion_model(
+        str(source),
+        str(output),
+        discard_patterns="blocks.0.attn",
+        process_device="cpu",
+    )
+    converted = load_file(str(output))
+
+    assert "blocks.0.attn.qkv_proj.weight" not in converted
+    assert "blocks.0.adaln_proj.linear.weight" in converted
+
+
+def test_fold_h3_exclude_patterns_preserves_unfolded_and_retains_time_embedder(folder_module, tmp_path):
+    source = tmp_path / "model_partial.safetensors"
+    output = tmp_path / "model_partial_folded.safetensors"
+
+    original = _make_h3_checkpoint()
+    save_file(original, str(source))
+
+    folder_module.fold_minimax_h3_diffusion_model(
+        str(source),
+        str(output),
+        exclude_patterns="blocks.1",
+        include_mode=False,
+        process_device="cpu",
+    )
+    converted = load_file(str(output))
+
+    # Block 0 folded
+    assert converted["blocks.0.adaln_proj.linear.weight"].shape == (128, 8)
+    # Block 1 kept at full width
+    assert converted["blocks.1.adaln_proj.linear.weight"].shape == (128, 2688)
+
+    # time_embedder MUST be retained because block 1 needs it
+    assert "time_embedder.proj_in.weight" in converted
+    assert "time_embedder.proj_out.weight" in converted
+
+
+def test_node_schema_and_execution_contract(folder_module, monkeypatch, tmp_path):
+    schema = folder_module.MiniMaxH3FoldAdaLN.define_schema()
+    assert schema.node_id == "MiniMaxH3FoldAdaLN"
+    assert schema.is_output_node is True
+
+    input_ids = [item.id for item in schema.inputs]
+    assert "model_name" in input_ids
+    assert "output_filename" in input_ids
+    assert "exclude_patterns" in input_ids
+    assert "include_mode" in input_ids
+    assert "discard_patterns" in input_ids
+    assert "glob_patterns" in input_ids
+
+    # Every input must have a non-empty static tooltip
+    for item in schema.inputs:
+        assert hasattr(item, "tooltip") and item.tooltip.strip()
+
+    output_displays = [item.display_name for item in schema.outputs]
+    assert output_displays == ["output_path", "report"]
+
+    # Test execution
+    models_dir = tmp_path / "models"
+    diff_dir = models_dir / "diffusion_models"
+    diff_dir.mkdir(parents=True)
+
+    source_file = diff_dir / "test_h3_model.safetensors"
+    save_file(_make_h3_checkpoint("diffusion_model."), str(source_file))
+
+    monkeypatch.setattr(folder_module.folder_paths, "models_dir", str(models_dir))
+    monkeypatch.setattr(
+        folder_module.folder_paths,
+        "get_full_path_or_raise",
+        lambda folder, name: str(diff_dir / name),
+    )
+
+    result = folder_module.MiniMaxH3FoldAdaLN.execute(
+        model_name="test_h3_model.safetensors",
+        output_filename="sub/folded_h3",
+        process_device="cpu",
+    )
+
+    assert result.args[0] == "sub/folded_h3.safetensors"
+    assert "MiniMax H3 AdaLN Fold complete" in result.args[1]
+    assert (diff_dir / "sub" / "folded_h3.safetensors").is_file()
+
+
+def test_unfold_h3_diffusion_model_reconstructs_with_base(folder_module, tmp_path):
+    base_file = tmp_path / "model_base.safetensors"
+    folded_file = tmp_path / "model_folded.safetensors"
+    unfolded_file = tmp_path / "model_unfolded.safetensors"
+
+    original_base = _make_h3_checkpoint()
+    save_file(original_base, str(base_file))
+
+    # Fold base model to create folded model A
+    folder_module.fold_minimax_h3_diffusion_model(
+        str(base_file), str(folded_file), process_device="cpu"
+    )
+
+    # Unfold model A using base model B
+    report = folder_module.unfold_minimax_h3_diffusion_model(
+        str(folded_file), str(base_file), str(unfolded_file), process_device="cpu"
+    )
+    assert "MiniMax H3 AdaLN Unfold complete" in report
+    assert "AdaLN layers unfolded: 2" in report
+
+    reconstructed = load_file(str(unfolded_file))
+
+    # adaln_t_table should be dropped in unfolded output
+    assert "adaln_t_table" not in reconstructed
+
+    # time_embedder restored from base model B
+    assert "time_embedder.proj_in.weight" in reconstructed
+    assert "time_embedder.proj_out.weight" in reconstructed
+    assert torch.equal(reconstructed["time_embedder.proj_in.weight"], original_base["time_embedder.proj_in.weight"])
+    assert torch.equal(reconstructed["time_embedder.proj_out.weight"], original_base["time_embedder.proj_out.weight"])
+
+    # AdaLN linear weights reconstructed to full width [128, 2688]
+    assert reconstructed["blocks.0.adaln_proj.linear.weight"].shape == (128, 2688)
+    assert reconstructed["blocks.1.adaln_proj.linear.weight"].shape == (128, 2688)
+
+    # Other tensors preserved
+    assert torch.equal(reconstructed["blocks.0.attn.qkv_proj.weight"], original_base["blocks.0.attn.qkv_proj.weight"])
+
+
+def test_unfold_h3_node_schema_and_execution(folder_module, monkeypatch, tmp_path):
+    schema = folder_module.MiniMaxH3UnfoldAdaLN.define_schema()
+    assert schema.node_id == "MiniMaxH3UnfoldAdaLN"
+    assert schema.is_output_node is True
+
+    input_ids = [item.id for item in schema.inputs]
+    assert input_ids == [
+        "model_a", "model_b", "output_filename", "discard_patterns",
+        "glob_patterns", "curve_rank", "curve_grid", "process_device",
+        "exclude_patterns", "include_mode",
+    ]
+    assert input_ids[-1] == "include_mode"
+
+    for item in schema.inputs:
+        assert hasattr(item, "tooltip") and item.tooltip.strip()
+
+    models_dir = tmp_path / "models"
+    diff_dir = models_dir / "diffusion_models"
+    diff_dir.mkdir(parents=True)
+
+    base_file = diff_dir / "base.safetensors"
+    folded_file = diff_dir / "folded.safetensors"
+
+    save_file(_make_h3_checkpoint(), str(base_file))
+    folder_module.fold_minimax_h3_diffusion_model(str(base_file), str(folded_file), process_device="cpu")
+
+    monkeypatch.setattr(folder_module.folder_paths, "models_dir", str(models_dir))
+    monkeypatch.setattr(
+        folder_module.folder_paths,
+        "get_full_path_or_raise",
+        lambda folder, name: str(diff_dir / name),
+    )
+
+    result = folder_module.MiniMaxH3UnfoldAdaLN.execute(
+        model_a="folded.safetensors",
+        model_b="base.safetensors",
+        output_filename="sub/unfolded_h3",
+        process_device="cpu",
+    )
+
+    assert result.args[0] == "sub/unfolded_h3.safetensors"
+    assert "MiniMax H3 AdaLN Unfold complete" in result.args[1]
+    assert (diff_dir / "sub" / "unfolded_h3.safetensors").is_file()

tests/test_minimax_h3_lora_convert.py

@@ -0,0 +1,226 @@
+import importlib
+from pathlib import Path
+import sys
+import types
+
+import pytest
+from safetensors.torch import load_file, save_file
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_minimax_convert_tests"
+
+
+def _load_module(name: str):
+    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 converter():
+    return _load_module("nodes.minimax_h3_lora_convert")
+
+
+def _make_source_tensors(prefix: str = "transformer_blocks.0") -> dict[str, torch.Tensor]:
+    tensors = {}
+    for offset, part in enumerate(("q", "k", "v"), start=1):
+        # Down: rank 2, in_dim 3
+        tensors[f"{prefix}.attn.to_{part}.lora_A.default.weight"] = torch.tensor(
+            [[1.0, float(offset), 0.0], [0.0, 1.0, float(offset)]],
+            dtype=torch.float32,
+        )
+        # Up: out_dim 2, rank 2
+        tensors[f"{prefix}.attn.to_{part}.lora_B.default.weight"] = torch.tensor(
+            [[float(offset), 0.0], [0.0, float(offset) + 1.0]],
+            dtype=torch.float32,
+        )
+
+    # out_proj: rank 2, in_dim 3, out_dim 2
+    tensors[f"{prefix}.attn.to_out.0.lora_A.default.weight"] = torch.arange(
+        6, dtype=torch.float32
+    ).reshape(2, 3)
+    tensors[f"{prefix}.attn.to_out.0.lora_B.default.weight"] = torch.arange(
+        4, dtype=torch.float32
+    ).reshape(2, 2)
+
+    # mlp.fc1 (SwiGLU): rank 2, in_dim 3, out_dim 4 (halves of size 2)
+    tensors[f"{prefix}.ff.net.0.proj.lora_A.default.weight"] = torch.ones(
+        (2, 3), dtype=torch.float32
+    )
+    # Val half: rows 0..1 (all 1s), Gate half: rows 2..3 (all 2s)
+    tensors[f"{prefix}.ff.net.0.proj.lora_B.default.weight"] = torch.cat(
+        [torch.full((2, 2), 1.0), torch.full((2, 2), 2.0)], dim=0
+    )
+
+    # mlp.fc2: rank 2, in_dim 4, out_dim 3
+    tensors[f"{prefix}.ff.net.2.lora_A.default.weight"] = torch.ones(
+        (2, 4), dtype=torch.float32
+    )
+    tensors[f"{prefix}.ff.net.2.lora_B.default.weight"] = torch.ones(
+        (3, 2), dtype=torch.float32
+    )
+
+    return tensors
+
+
+def test_full_rank_qkv_conversion_is_exact(converter, tmp_path):
+    source = tmp_path / "source_exact.safetensors"
+    output = tmp_path / "output_exact.safetensors"
+
+    original = _make_source_tensors("token_refiner.refiner_blocks.0")
+    save_file(original, str(source))
+
+    report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
+    assert "MiniMax H3 Diffusers LoRA conversion complete" in report
+    assert "Fused QKV modules: 1" in report
+
+    converted = load_file(str(output))
+
+    down = converted["diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_A.weight"]
+    up = converted["diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_B.weight"]
+
+    # Ranks: 2 + 2 + 2 = 6, in_dim: 3, out_dim: 2 + 2 + 2 = 6
+    assert down.shape == (6, 3)
+    assert up.shape == (6, 6)
+
+    # Mathematical delta equality: up @ down must equal stacked (up_q @ down_q, up_k @ down_k, up_v @ down_v)
+    fused_delta = up @ down
+    expected_delta = torch.cat(
+        [
+            original[f"token_refiner.refiner_blocks.0.attn.to_{p}.lora_B.default.weight"]
+            @ original[f"token_refiner.refiner_blocks.0.attn.to_{p}.lora_A.default.weight"]
+            for p in ("q", "k", "v")
+        ],
+        dim=0,
+    )
+    assert torch.allclose(fused_delta, expected_delta, atol=1e-5, rtol=1e-5)
+
+
+def test_swiglu_mlp_fc1_halves_swapped(converter, tmp_path):
+    source = tmp_path / "source_swiglu.safetensors"
+    output = tmp_path / "output_swiglu.safetensors"
+
+    original = _make_source_tensors("transformer_blocks.0")
+    save_file(original, str(source))
+
+    converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
+    converted = load_file(str(output))
+
+    fc1_b = converted["diffusion_model.blocks.0.mlp.fc1.lora_B.weight"]
+    orig_b = original["transformer_blocks.0.ff.net.0.proj.lora_B.default.weight"]
+
+    # In source: top half (rows 0..1) is 1.0, bottom half (rows 2..3) is 2.0
+    # In converted: top half must be 2.0 (gate), bottom half must be 1.0 (value)
+    assert torch.equal(fc1_b[:2], orig_b[2:])
+    assert torch.equal(fc1_b[2:], orig_b[:2])
+
+
+def test_rejects_incomplete_qkv_group(converter, tmp_path):
+    source = _make_source_tensors()
+    del source["transformer_blocks.0.attn.to_v.lora_B.default.weight"]
+
+    path = tmp_path / "incomplete.safetensors"
+    save_file(source, str(path))
+
+    with pytest.raises(ValueError, match="Incomplete LoRA A/B pair"):
+        converter.convert_minimax_h3_diffusers_lora(
+            str(path), str(tmp_path / "unused.safetensors")
+        )
+
+
+def test_rejects_unsupported_module(converter, tmp_path):
+    source = {"unknown_layer.lora_A.weight": torch.ones(2, 2), "unknown_layer.lora_B.weight": torch.ones(2, 2)}
+    path = tmp_path / "unknown.safetensors"
+    save_file(source, str(path))
+
+    with pytest.raises(ValueError, match="Unsupported MiniMax H3 Diffusers module"):
+        converter.convert_minimax_h3_diffusers_lora(
+            str(path), str(tmp_path / "unused.safetensors")
+        )
+
+
+def test_node_schema_and_execution_contract(converter, monkeypatch, tmp_path):
+    schema = converter.MiniMaxH3DiffusersLoRAConvert.define_schema()
+    assert schema.node_id == "MiniMaxH3DiffusersLoRAConvert"
+    assert schema.is_output_node is True
+
+    input_ids = [item.id for item in schema.inputs]
+    assert input_ids == ["lora_name", "output_filename"]
+
+    output_displays = [item.display_name for item in schema.outputs]
+    assert output_displays == ["output_path", "conversion_report"]
+
+    # Test execution through node
+    models_dir = tmp_path / "models"
+    loras_dir = models_dir / "loras"
+    loras_dir.mkdir(parents=True)
+
+    source_file = loras_dir / "test_diffusers_lora.safetensors"
+    save_file(_make_source_tensors(), str(source_file))
+
+    monkeypatch.setattr(converter.folder_paths, "models_dir", str(models_dir))
+    monkeypatch.setattr(
+        converter.folder_paths,
+        "get_full_path_or_raise",
+        lambda folder, name: str(loras_dir / name),
+    )
+
+    result = converter.MiniMaxH3DiffusersLoRAConvert.execute(
+        lora_name="test_diffusers_lora.safetensors",
+        output_filename="sub/converted_h3",
+    )
+
+    assert result.args[0] == "sub/converted_h3.safetensors"
+    assert "MiniMax H3 Diffusers LoRA conversion complete" in result.args[1]
+    assert (loras_dir / "sub" / "converted_h3.safetensors").is_file()
+
+
+def test_lora_down_up_suffixes_supported(converter, tmp_path):
+    source = tmp_path / "source_down_up.safetensors"
+    output = tmp_path / "output_down_up.safetensors"
+
+    original = _make_source_tensors("transformer_blocks.0")
+    # Replace .lora_A/.lora_B with .lora.down/.lora.up on attention layers (DMAD format)
+    modified = {}
+    for k, v in original.items():
+        if ".attn." in k:
+            k = k.replace(".lora_A.default.weight", ".lora.down.weight")
+            k = k.replace(".lora_B.default.weight", ".lora.up.weight")
+        else:
+            k = k.replace(".default.weight", ".weight")
+        modified[k] = v
+
+    save_file(modified, str(source))
+    report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
+    assert "MiniMax H3 Diffusers LoRA conversion complete" in report
+    converted = load_file(str(output))
+
+    assert "diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight" in converted
+    assert "diffusion_model.blocks.0.attn.qkv_proj.lora_B.weight" in converted
+    assert "diffusion_model.blocks.0.mlp.fc1.lora_A.weight" in converted
+    assert "diffusion_model.blocks.0.mlp.fc1.lora_B.weight" in converted
+
+
+def test_bare_lora_ab_without_weight_suffix_supported(converter, tmp_path):
+    source = tmp_path / "source_bare_ab.safetensors"
+    output = tmp_path / "output_bare_ab.safetensors"
+
+    original = _make_source_tensors("token_refiner.refiner_blocks.0")
+    modified = {
+        k.replace(".default.weight", ""): v
+        for k, v in original.items()
+    }
+
+    save_file(modified, str(source))
+    report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
+    assert "MiniMax H3 Diffusers LoRA conversion complete" in report
+    converted = load_file(str(output))
+
+    assert "diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_A.weight" in converted
+    assert "diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_B.weight" in converted
+    assert "diffusion_model.token_refiner.blocks.0.mlp.fc1.lora_A.weight" in converted
+    assert "diffusion_model.token_refiner.blocks.0.mlp.fc1.lora_B.weight" in converted
+

tests/test_model_analysis.py

@@ -0,0 +1,481 @@
+import importlib
+import csv
+from io import StringIO
+import json
+import math
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_analysis_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 analysis():
+    return _load_module("nodes.model_analysis")
+
+
+def _params(**overrides):
+    params = {
+        "cwb_similarity_alignment": False,
+        "top_weight_differences": 3,
+        "process_device": "cpu",
+        "force_clear_cache": False,
+        "exclude_patterns": "",
+        "glob_patterns": False,
+    }
+    params.update(overrides)
+    return params
+
+
+def test_five_node_schema_contract(analysis):
+    assert len(analysis.MODEL_ANALYSIS_NODES) == 5
+    assert len({node.NODE_ID for node in analysis.MODEL_ANALYSIS_NODES}) == 5
+    for node in analysis.MODEL_ANALYSIS_NODES:
+        schema = node.define_schema()
+        assert schema.category == "ModelUtils/Analysis"
+        assert schema.description
+        assert all(item.tooltip for item in schema.inputs)
+        assert [item.id for item in schema.inputs[:3]] == [
+            "execution_mode", "model_a", "model_b",
+        ]
+        assert schema.inputs[0].options == ["ANALYZE", "DOCUMENTATION ONLY"]
+        assert "lazy_load" not in [item.id for item in schema.inputs]
+        input_ids = [item.id for item in schema.inputs]
+        assert input_ids[-3:] == ["exclude_patterns", "glob_patterns", "include_mode"]
+        assert schema.inputs[-1].default is False
+        if node in (analysis.LoRAModelAnalysis, analysis.EmbeddingModelAnalysis):
+            assert input_ids[3] == "cwb_similarity_alignment"
+        else:
+            assert "cwb_similarity_alignment" not in input_ids
+        assert not {
+            "cwb_preset", "consensus_type", "alignment_threshold",
+            "similarity_threshold", "power_alpha", "diversity_beta",
+            "dynamic_similarity_contrast", "soft_comfort_bandpass",
+            "position_weight", "preserve_common_prefix",
+        } & set(input_ids)
+        assert [item.display_name for item in schema.outputs] == [
+            "comparison_report", "cwb_report", "documentation",
+            "layerwise_metrics_csv", "layerwise_cwb_csv",
+        ]
+
+
+def test_documentation_mode_preserves_reports_and_returns_empty_csv(analysis):
+    result = analysis.DiffusionModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")
+    assert len(result.result) == 5
+    assert "No comparison performed" in result[0]
+    assert "No CWB analysis performed" in result[1]
+    assert result[2].startswith("# Two-Model Similarity Analysis")
+    assert result[3] == result[4] == ""
+
+
+def test_raw_metrics_have_independent_expected_values(analysis):
+    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    b = torch.tensor([[1.0, 0.0], [5.0, 4.0]])
+    raw, top = analysis._compute_raw(a, b, 2)
+    metrics = analysis._metrics(raw)
+    assert metrics["mae"] == pytest.approx(1.0)
+    assert metrics["mse"] == pytest.approx(2.0)
+    assert metrics["rmse"] == pytest.approx(math.sqrt(2.0))
+    assert metrics["max"] == pytest.approx(2.0)
+    assert metrics["exact"] == pytest.approx(0.5)
+    assert len(top) == 2
+    assert {entry[1] for entry in top} == {1, 2}
+
+
+def test_cwb_derived_metrics_have_independent_expected_values(analysis):
+    a = torch.tensor([[1.0, 0.0], [1.0, 0.0]])
+    b = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+    values = analysis._compute_cwb(a, b, False).values()
+    assert values["pair_mean"] == pytest.approx(0.5)
+    assert values["pair_min"] == pytest.approx(0.0)
+    assert values["pair_max"] == pytest.approx(1.0)
+    assert values["mean_aff_a"] == pytest.approx((1.0 + 2 ** -0.5) / 2)
+    assert values["mean_aff_b"] == pytest.approx((1.0 + 2 ** -0.5) / 2)
+    assert values["median_aff_a"] == pytest.approx(0.5)
+    assert values["median_aff_b"] == pytest.approx(0.5)
+    assert values["alignment"] is None
+
+
+def test_zero_constant_and_nonfinite_metrics_are_explicit(analysis):
+    raw, _ = analysis._compute_raw(torch.zeros(3), torch.zeros(3), 0)
+    metrics = analysis._metrics(raw)
+    assert metrics["cosine"] == pytest.approx(1.0)
+    assert metrics["pearson"] is None
+    assert metrics["norm_ratio"] == pytest.approx(1.0)
+
+    raw, _ = analysis._compute_raw(
+        torch.tensor([1.0, float("nan"), float("inf")]),
+        torch.tensor([2.0, 4.0, 5.0]),
+        0,
+    )
+    metrics = analysis._metrics(raw)
+    assert raw.finite == 1
+    assert raw.nonfinite_a == 2
+    assert metrics["coverage"] == pytest.approx(1 / 3)
+
+
+def test_end_to_end_reports_and_topology_classification(monkeypatch, tmp_path, analysis):
+    tensor_key = 'model.blocks.0.weight|"quoted",\nname'
+    path_a = tmp_path / "a.safetensors"
+    path_b = tmp_path / "b.safetensors"
+    save_file({
+        tensor_key: torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+        "only_a": torch.ones(1),
+        "shape": torch.ones(2),
+        "integer": torch.tensor([1, 2], dtype=torch.int64),
+    }, str(path_a), metadata={"source": "a"})
+    save_file({
+        tensor_key: torch.tensor([[1.0, 0.0], [5.0, 4.0]]),
+        "only_b": torch.ones(1),
+        "shape": torch.ones(3),
+        "integer": torch.tensor([1, 3], dtype=torch.int64),
+    }, str(path_b), metadata={"source": "b"})
+    paths = {"a": str(path_a), "b": str(path_b)}
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    processed = []
+    async_calls = []
+    original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed
+    original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def track_processed(loader, key):
+        processed.append((loader.filename, key))
+        return original_mark_processed(loader, key)
+
+    def track_async_stream(loader, keys, **kwargs):
+        async_calls.append((loader.filename, tuple(keys), kwargs))
+        yield from original_async_stream(loader, keys, **kwargs)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed)
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream)
+
+    result = analysis.DiffusionModelAnalysis.execute(
+        model_a="a", model_b="b", execution_mode="ANALYZE", **_params()
+    )
+    comparison, cwb_report, documentation, metrics_csv, cwb_csv = result.result
+    assert documentation.startswith("# Two-Model Similarity Analysis")
+    assert comparison.startswith("# MODEL COMPARISON SUMMARY\n\n")
+    assert "| Comparable floating tensors | 1 |" in comparison
+    assert "| mae | 1 |" in comparison
+    assert "| mse | 2 |" in comparison
+    assert '| model.blocks.0.weight\\|"quoted",<br>name |' in comparison
+    assert "| --- | --- |" in comparison
+    metrics_rows = list(csv.DictReader(StringIO(metrics_csv)))
+    cwb_rows = list(csv.DictReader(StringIO(cwb_csv)))
+    assert len(metrics_rows) == len(cwb_rows) == 1
+    assert metrics_rows[0]["key"] == cwb_rows[0]["key"] == tensor_key
+    assert float(metrics_rows[0]["MAE"]) == 1
+    assert float(metrics_rows[0]["RMSE"]) == pytest.approx(math.sqrt(2), rel=1e-5)
+    assert metrics_rows[0]["dtype_A"] == metrics_rows[0]["dtype_B"] == "torch.float32"
+    assert cwb_rows[0]["matched"] == "N/A"
+    assert float(cwb_rows[0]["pair_cos_min"]) <= float(cwb_rows[0]["pair_cos_mean"]) <= float(cwb_rows[0]["pair_cos_max"])
+    top_rows = comparison.split("## TOP SCALAR WEIGHT DIFFERENCES\n\n", 1)[1].split("\n\n", 1)[0].splitlines()[2:]
+    assert len(top_rows) == 3
+    assert [float(row.rsplit(" | ", 4)[2]) for row in top_rows] == [2, 2, 0]
+    assert "[A ONLY] (1)" in comparison
+    assert "[B ONLY] (1)" in comparison
+    assert "[SHAPE MISMATCH] (1)" in comparison
+    assert "[NON-FLOATING] (1)" in comparison
+    assert "model.blocks.0 |" in comparison
+    assert "CWB-DERIVED SIMILARITY SUMMARY" in cwb_report
+    assert "merge weighting" in cwb_report
+    assert "PRESET-IMPLIED CONTRIBUTIONS" not in cwb_report
+    assert "Average contribution" not in cwb_report
+    assert len(processed) == 4
+    assert {key for _, key in processed} == {tensor_key, "integer"}
+    assert len(async_calls) == 4
+    assert all(len(call[1]) == 1 for call in async_calls)
+    assert all(call[2]["batch_size"] == 1 for call in async_calls)
+    assert all(call[2]["prefetch_batches"] == 1 for call in async_calls)
+    assert not list(tmp_path.glob("**/*analysis*"))
+
+
+def test_lora_cross_format_keys_compare_as_one_logical_pair(monkeypatch, tmp_path, analysis):
+    path_a = tmp_path / "lora_a.safetensors"
+    path_b = tmp_path / "lora_b.safetensors"
+    save_file({
+        "diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+        "diffusion_model.foo.lora_B.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+    }, str(path_a))
+    save_file({
+        "base_model.model.diffusion_model.foo.lora_down.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+        "base_model.model.diffusion_model.foo.lora_up.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+    }, str(path_b))
+    paths = {"a": str(path_a), "b": str(path_b)}
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    processed = []
+    async_calls = []
+    original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed
+    original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def track_processed(loader, key):
+        processed.append((loader.filename, key))
+        return original_mark_processed(loader, key)
+
+    def track_async_stream(loader, keys, **kwargs):
+        async_calls.append((loader.filename, tuple(keys), kwargs))
+        yield from original_async_stream(loader, keys, **kwargs)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed)
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream)
+
+    comparison, cwb_report, metrics_csv, cwb_csv = analysis.ModelAnalysisLogic.execute(
+        "a", "b", "loras", _params(), lora_mode=True,
+    )
+    assert "| Comparable floating tensors | 2 |" in comparison
+    assert "[A ONLY] (0)" in comparison
+    assert "[B ONLY] (0)" in comparison
+    assert "foo.down |" in comparison
+    assert "foo.up |" in comparison
+    assert "foo.\\[down+up\\] |" in cwb_report
+    assert [row["key"] for row in csv.DictReader(StringIO(metrics_csv))] == ["foo.down", "foo.up"]
+    assert [row["key"] for row in csv.DictReader(StringIO(cwb_csv))] == ["foo.[down+up]"]
+    assert len(processed) == 4
+    assert len({entry for entry in processed}) == 4
+    assert len(async_calls) == 2
+    assert all(len(call[1]) == 2 for call in async_calls)
+
+
+@pytest.mark.parametrize("patterns, glob_patterns, include_mode", [
+    (r"^diffusion_model\.blocks\.4\.", False, False),
+    (r"^diffusion_model\.blocks\.3\.", False, True),
+    ("diffusion_model.blocks.3.*", True, True),
+])
+def test_excluded_lora_blocks_do_not_enter_metrics(
+    monkeypatch, tmp_path, analysis, patterns, glob_patterns, include_mode,
+):
+    path_a = tmp_path / "lora_a.safetensors"
+    path_b = tmp_path / "lora_b.safetensors"
+    tensors_a = {}
+    tensors_b = {}
+    for block, value_a, value_b in ((3, 1.0, 2.0), (4, 10.0, 20.0)):
+        prefix = f"diffusion_model.blocks.{block}.attn"
+        tensors_a[f"{prefix}.lora_A.weight"] = torch.full((1, 2), value_a)
+        tensors_a[f"{prefix}.lora_B.weight"] = torch.full((2, 1), value_a)
+        tensors_b[f"{prefix}.lora_A.weight"] = torch.full((1, 2), value_b)
+        tensors_b[f"{prefix}.lora_B.weight"] = torch.full((2, 1), value_b)
+    save_file(tensors_a, str(path_a))
+    save_file(tensors_b, str(path_b))
+    paths = {"a": str(path_a), "b": str(path_b)}
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+
+    comparison, cwb_report, metrics_csv, cwb_csv = analysis.ModelAnalysisLogic.execute(
+        "a", "b", "loras",
+        _params(exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=include_mode),
+        lora_mode=True,
+    )
+
+    assert "| Comparable floating tensors | 2 |" in comparison
+    assert "| Excluded A tensors | 2 |" in comparison
+    assert "| Excluded B tensors | 2 |" in comparison
+    layerwise = comparison.split("LAYERWISE TENSOR METRICS", 1)[1].split(
+        "TOP SCALAR WEIGHT DIFFERENCES", 1
+    )[0]
+    assert "blocks\\_3\\_" in layerwise
+    assert "blocks\\_4\\_" not in layerwise
+    assert "blocks\\_4\\_" not in cwb_report
+    assert all("blocks_3_" in row["key"] for row in csv.DictReader(StringIO(metrics_csv)))
+    assert all("blocks_3_" in row["key"] for row in csv.DictReader(StringIO(cwb_csv)))
+
+
+@pytest.mark.parametrize("patterns, glob_patterns, expected_keys", [
+    (r"^blocks\.0\.", False, ["blocks.0.weight"]),
+    ("blocks.0.*", True, ["blocks.0.weight"]),
+    ("", False, []),
+    ("no_such_layer", True, []),
+])
+def test_include_filter_limits_analysis_inventory(
+    monkeypatch, tmp_path, analysis, patterns, glob_patterns, expected_keys,
+):
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    for name, value in (("a", 1.0), ("b", 2.0)):
+        save_file({
+            "blocks.0.weight": torch.full((2, 2), value),
+            "blocks.1.weight": torch.full((2, 2), value),
+            f"only_{name}.weight": torch.full((2, 2), value),
+        }, paths[name])
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    streamed = []
+    original_stream = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def stream(handler, keys, *args, **kwargs):
+        streamed.extend(keys)
+        return original_stream(handler, keys, *args, **kwargs)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
+    analysis.ModelAnalysisLogic.execute(
+        "a", "b", "diffusion_models",
+        _params(exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=True),
+    )
+    assert sorted(streamed) == sorted(expected_keys * 2)
+
+
+def test_diffusion_analysis_decodes_int8_tensorwise_and_skips_sidecars(
+    monkeypatch, tmp_path, analysis,
+):
+    prefix = "model.diffusion_model.layer"
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    save_file({
+        f"{prefix}.weight": torch.tensor([1, 2], dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+        "dense.weight": torch.tensor([2.0, 4.0]),
+    }, paths["a"])
+    save_file({
+        f"{prefix}.weight": torch.tensor([3, 6], dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+        "dense.weight": torch.tensor([4.0, 8.0]),
+    }, paths["b"])
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    streamed = []
+    original_stream = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def track_stream(handler, keys, *args, **kwargs):
+        streamed.append((handler.filename, tuple(keys)))
+        return original_stream(handler, keys, *args, **kwargs)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_stream)
+    _, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute(
+        "a", "b", "diffusion_models", _params(),
+    )
+
+    rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))}
+    assert set(rows) == {"layer.weight", "dense.weight"}
+    assert float(rows["layer.weight"]["MAE"]) == pytest.approx(6.0)
+    assert float(rows["dense.weight"]["MAE"]) == pytest.approx(3.0)
+    for path, keys in streamed:
+        if path.endswith("a.safetensors") or path.endswith("b.safetensors"):
+            if f"{prefix}.weight" in keys:
+                assert f"{prefix}.weight_scale" in keys
+                assert f"{prefix}.comfy_quant" in keys
+    assert sum(f"{prefix}.weight" in keys or "dense.weight" in keys for _, keys in streamed) == 4
+
+
+@pytest.mark.parametrize("prefixed_input", ["a", "b"])
+def test_diffusion_analysis_matches_prefixed_and_bare_layers(
+    monkeypatch, tmp_path, analysis, prefixed_input,
+):
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    prefix = "model.diffusion_model."
+    for name in ("a", "b"):
+        key = f"{prefix if name == prefixed_input else ''}layer.weight"
+        if name == prefixed_input:
+            stem = key[:-len(".weight")]
+            save_file({
+                key: torch.tensor([1, 2], dtype=torch.int8),
+                f"{stem}.weight_scale": torch.tensor(2.0),
+                f"{stem}.comfy_quant": torch.tensor(
+                    list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
+                ),
+            }, paths[name])
+        else:
+            save_file({key: torch.tensor([4.0, 8.0])}, paths[name])
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+
+    _, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute(
+        "a", "b", "diffusion_models", _params(),
+    )
+    rows = list(csv.DictReader(StringIO(metrics_csv)))
+    assert [row["key"] for row in rows] == ["layer.weight"]
+    assert float(rows[0]["MAE"]) == pytest.approx(3.0)
+
+
+def test_diffusion_analysis_rejects_duplicate_prefix_variants(monkeypatch, tmp_path, analysis):
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    save_file({
+        "layer.weight": torch.ones(2),
+        "model.diffusion_model.layer.weight": torch.ones(2),
+    }, paths["a"])
+    save_file({"layer.weight": torch.ones(2)}, paths["b"])
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    with pytest.raises(ValueError, match="duplicate diffusion layer 'layer.weight'"):
+        analysis.ModelAnalysisLogic.execute("a", "b", "diffusion_models", _params())
+
+
+def test_embedding_similarity_alignment_is_explicit_opt_in(analysis):
+    a = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+    b = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
+    aligned = analysis._compute_cwb(a, b, True).values()
+    indexed = analysis._compute_cwb(a, b, False).values()
+    assert aligned["alignment"] == pytest.approx(1.0)
+    assert aligned["alignment_gain"] == pytest.approx(1.0)
+    assert indexed["pair_mean"] == pytest.approx(0.0)
+    assert indexed["alignment"] is None
+
+
+def test_cuda_oom_retries_once_on_cpu(monkeypatch, analysis):
+    calls = []
+
+    def fake_compute(a, b, device, allow_alignment, top_count):
+        calls.append(device)
+        if device == "cuda":
+            raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory")
+        return analysis.RawStats(finite=1, total=1), analysis.CWBStats(), []
+
+    monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
+    monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
+    result = analysis._analyze_with_fallback(
+        "layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
+    )
+    assert calls == ["cuda", "cpu"]
+    assert result[-1] is True
+
+
+def test_non_oom_error_is_not_retried(monkeypatch, analysis):
+    calls = []
+
+    def fake_compute(*args):
+        calls.append(args[2])
+        raise RuntimeError("unrelated failure")
+
+    monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
+    with pytest.raises(RuntimeError, match="unrelated failure"):
+        analysis._analyze_with_fallback(
+            "layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
+        )
+    assert calls == ["cuda"]
+
+
+def test_failed_cpu_retry_names_layer_and_chains_oom(monkeypatch, analysis):
+    def fake_compute(a, b, device, allow_alignment, top_count):
+        if device == "cuda":
+            raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory")
+        raise ValueError("CPU calculation failed")
+
+    monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
+    monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
+    with pytest.raises(RuntimeError, match="CPU fallback failed while analyzing 'named.layer'") as caught:
+        analysis._analyze_with_fallback(
+            "named.layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
+        )
+    assert isinstance(caught.value.__cause__, torch.cuda.OutOfMemoryError)

tests/test_model_analysis_api_runner.py

@@ -0,0 +1,56 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+@pytest.fixture(scope="module")
+def runner():
+    spec = importlib.util.spec_from_file_location(
+        "model_analysis_api_runner_test",
+        REPO_ROOT / "run_model_analysis_api_workflow.py",
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_linked_previews_require_both_reports(runner):
+    workflow = {
+        "1": {"class_type": "LoRAModelAnalysis", "inputs": {}},
+        "2": {"class_type": "PreviewAny", "inputs": {"source": ["1", 0]}},
+        "3": {"class_type": "PreviewAny", "inputs": {"source": ["1", 1]}},
+    }
+    assert runner._linked_previews(workflow, "1") == {0: "2", 1: "3"}
+
+
+def test_linked_previews_reject_missing_report(runner):
+    workflow = {
+        "1": {"class_type": "LoRAModelAnalysis", "inputs": {}},
+        "2": {"class_type": "PreviewAny", "inputs": {"source": ["1", 0]}},
+    }
+    with pytest.raises(RuntimeError, match="output"):
+        runner._linked_previews(workflow, "1")
+
+
+def test_run_analysis_captures_both_reports(monkeypatch, runner):
+    workflow = {
+        "1": {"class_type": "LoRAModelAnalysis", "inputs": {}},
+        "2": {"class_type": "PreviewAny", "inputs": {"source": ["1", 0]}},
+        "3": {"class_type": "PreviewAny", "inputs": {"source": ["1", 1]}},
+    }
+    monkeypatch.setattr(runner.api, "_json_request", lambda *args: {"prompt_id": "p"})
+    monkeypatch.setattr(
+        runner.api,
+        "_wait_for_history",
+        lambda *args: {"outputs": {"2": {"text": ["standard"]}, "3": {"text": ["cwb"]}}},
+    )
+    result = runner.run_analysis(workflow, "http://server", 0.1)
+    assert result == {
+        "prompt_id": "p",
+        "comparison_report": "standard",
+        "cwb_report": "cwb",
+    }

tests/test_schema_tooltips.py

@@ -0,0 +1,153 @@
+import ast
+import importlib
+import inspect
+import sys
+import types
+from pathlib import Path
+
+import pytest
+from comfy_api.latest import io
+
+
+NODE_DIR = Path(__file__).resolve().parents[1] / "nodes"
+
+
+def test_every_schema_input_has_a_direct_static_tooltip():
+    missing = []
+    non_static = []
+
+    for path in sorted(NODE_DIR.glob("*.py")):
+        tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+        for call in ast.walk(tree):
+            if not (
+                isinstance(call, ast.Call)
+                and isinstance(call.func, ast.Attribute)
+                and call.func.attr == "Input"
+            ):
+                continue
+
+            name = (
+                call.args[0].value
+                if call.args and isinstance(call.args[0], ast.Constant)
+                else "<dynamic input name>"
+            )
+            tooltip = next(
+                (keyword.value for keyword in call.keywords if keyword.arg == "tooltip"),
+                None,
+            )
+            location = f"{path.name}:{call.lineno}:{name}"
+            if tooltip is None:
+                missing.append(location)
+            elif not (
+                isinstance(tooltip, ast.Constant)
+                and isinstance(tooltip.value, str)
+                and tooltip.value.strip()
+            ):
+                if not (path.name == "layer_parameters.py" and name == "layer_parameters"):
+                    non_static.append(location)
+
+    assert not missing, "Inputs without tooltips:\n" + "\n".join(missing)
+    assert not non_static, "Tooltips must be direct non-empty string literals:\n" + "\n".join(non_static)
+
+
+FILTER_MODULES = [
+    "merger", "consensus_merger", "cwb_model_graft", "cwb_delta_lora_merger", "lodestone_merger",
+    "lora_resize", "dtype_conversion", "model_analysis",
+    "lora_extract_svd", "dora_extract_wd", "dora_learned_wd", "text_encoder_extract",
+    "minimax_h3_fold",
+]
+
+
+def test_filter_schema_audit_covers_all_direct_filter_definitions():
+    discovered = set()
+    for path in NODE_DIR.glob("*.py"):
+        tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+        for call in ast.walk(tree):
+            if (
+                isinstance(call, ast.Call)
+                and isinstance(call.func, ast.Attribute)
+                and call.func.attr == "Input"
+                and call.args
+                and isinstance(call.args[0], ast.Constant)
+                and call.args[0].value in {"exclude_patterns", "skip_patterns"}
+            ):
+                discovered.add(path.stem)
+    assert discovered <= set(FILTER_MODULES), f"Unaudited filter modules: {discovered - set(FILTER_MODULES)}"
+
+
+def _load_filter_module(module_name):
+    package_name = "modelutils_filter_schema_tests"
+    if package_name not in sys.modules:
+        package = types.ModuleType(package_name)
+        package.__path__ = [str(NODE_DIR.parent)]
+        sys.modules[package_name] = package
+    return importlib.import_module(f"{package_name}.nodes.{module_name}")
+
+
+@pytest.mark.parametrize("module_name", FILTER_MODULES)
+def test_layer_filter_schemas_append_disabled_include_toggle(module_name):
+    module = _load_filter_module(module_name)
+    checked = []
+    for name, node in vars(module).items():
+        if name.startswith("_") or not isinstance(node, type):
+            continue
+        if node.__module__ != module.__name__ or not issubclass(node, io.ComfyNode):
+            continue
+        schema = node.define_schema()
+        previous_inputs = [item for item in schema.inputs if item.id != "layer_parameters"]
+        ids = [item.id for item in previous_inputs]
+        if not {"exclude_patterns", "skip_patterns"}.intersection(ids):
+            continue
+        assert ids.count("include_mode") == 1, schema.node_id
+        assert ids[-1] == "include_mode", schema.node_id
+        assert isinstance(previous_inputs[-1], io.Boolean.Input), schema.node_id
+        assert previous_inputs[-1].default is False, schema.node_id
+        checked.append(schema.node_id)
+    assert checked, f"No filter nodes checked in {module_name}"
+
+
+@pytest.mark.parametrize("module_name, helper_name, expected_count", [
+    ("lora_extract_svd", "extract_lora_from_files", 5),
+    ("dora_extract_wd", "extract_dora_from_files", 5),
+    ("dora_learned_wd", "extract_dora_learned_from_files", 5),
+    ("text_encoder_extract", "extract_te_from_files", 10),
+])
+@pytest.mark.parametrize("include_mode", [False, True])
+def test_all_extraction_nodes_forward_include_mode_without_shifting_parameters(
+    monkeypatch, module_name, helper_name, expected_count, include_mode,
+):
+    module = _load_filter_module(module_name)
+    signature = inspect.signature(getattr(module, helper_name))
+    assert [name for name in signature.parameters if name != "layer_parameters"][-1] == "include_mode"
+    captured = []
+
+    def extract(*args, **kwargs):
+        bound = signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+        captured.append(bound.arguments)
+
+    monkeypatch.setattr(module, helper_name, extract)
+    monkeypatch.setattr(module.folder_paths, "get_full_path_or_raise", lambda category, name: name)
+    monkeypatch.setattr(module, "_build_lora_output_path", lambda name: (name, name))
+    for name, node in vars(module).items():
+        if name.startswith("_") or not isinstance(node, type):
+            continue
+        if node.__module__ != module.__name__ or not issubclass(node, io.ComfyNode):
+            continue
+        schema = node.define_schema()
+        values = {}
+        for widget in schema.inputs:
+            if widget.id == "layer_parameters":
+                continue
+            value = getattr(widget, "default", None)
+            if value is None:
+                options = getattr(widget, "options", [])
+                value = options[0] if options else widget.id
+            values[widget.id] = value
+        values["include_mode"] = include_mode
+        values["skip_patterns"] = "selected_layer"
+        node.execute(**values)
+        assert captured[-1]["include_mode"] is include_mode, schema.node_id
+        assert captured[-1]["skip_patterns_str"] == "selected_layer", schema.node_id
+        assert captured[-1]["knee_probe_offset"] == values.get("knee_probe_offset", 32), schema.node_id
+    assert len(captured) == expected_count

tests/test_uel_repository_contract.py

@@ -0,0 +1,114 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_uel_io():
+    spec = importlib.util.spec_from_file_location(
+        "modelutils_uel_io_test", REPO_ROOT / "nodes" / "uel_io.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class _Handler:
+    def __init__(self, values):
+        self.values = values
+        self.calls = []
+        self.processed = []
+        self.closed = 0
+
+    def async_stream(self, keys, **kwargs):
+        self.calls.append((list(keys), kwargs))
+
+        def generate():
+            try:
+                for key in keys:
+                    yield [(key, self.values[key])]
+            finally:
+                self.closed += 1
+
+        return generate()
+
+    def mark_processed(self, key):
+        self.processed.append(key)
+
+
+def test_production_nodes_have_no_synchronous_safetensors_io():
+    violations = []
+    for path in sorted((REPO_ROOT / "nodes").glob("*.py")):
+        source = path.read_text(encoding="utf-8")
+        if ".get_tensor(" in source:
+            violations.append(f"{path.name}: get_tensor")
+        if "from safetensors" in source or "import safetensors" in source:
+            violations.append(f"{path.name}: direct safetensors import")
+    assert violations == []
+
+
+def test_work_units_use_bounded_streams_and_release_every_yield():
+    uel = _load_uel_io()
+    first = _Handler({"a": torch.tensor([1.0]), "b": torch.tensor([2.0])})
+    second = _Handler({"x": torch.tensor([3.0])})
+    units = [("one", {0: ["a"], 1: ["x"]}), ("two", {0: ["b"]})]
+
+    observed = []
+    for logical, loaded in uel.stream_work_units(
+        {0: first, 1: second}, units, pin_memory=True
+    ):
+        observed.append((logical, sorted(loaded)))
+
+    assert observed == [
+        ("one", [(0, "a"), (1, "x")]),
+        ("two", [(0, "b")]),
+    ]
+    assert first.processed == ["a", "b"]
+    assert second.processed == ["x"]
+    assert first.closed == second.closed == 1
+    for handler in (first, second):
+        assert handler.calls[0][1] == {
+            "batch_size": 1,
+            "prefetch_batches": 1,
+            "pin_memory": True,
+        }
+
+
+def test_work_unit_failure_releases_consumed_sources_and_closes_streams():
+    uel = _load_uel_io()
+    handler = _Handler({"a": torch.tensor([1.0]), "b": torch.tensor([2.0])})
+    stream = uel.stream_work_units(
+        {0: handler}, [("one", {0: ["a"]}), ("two", {0: ["b"]})],
+        pin_memory=False,
+    )
+
+    with pytest.raises(RuntimeError, match="consumer failed"):
+        try:
+            for _, _ in stream:
+                raise RuntimeError("consumer failed")
+        finally:
+            stream.close()
+
+    assert handler.processed == ["a"]
+    assert handler.closed == 1
+
+
+def test_atomic_writer_preserves_destination_on_failure(tmp_path):
+    uel = _load_uel_io()
+    output = tmp_path / "output.safetensors"
+    with IncrementalSafetensorsWriter(str(output), max_workers=1) as writer:
+        writer.write("old", torch.tensor([1.0]))
+    original = output.read_bytes()
+
+    with pytest.raises(RuntimeError, match="planned failure"):
+        with uel.atomic_uel_writer(str(output)) as writer:
+            writer.write("new", torch.tensor([2.0]))
+            raise RuntimeError("planned failure")
+
+    assert output.read_bytes() == original
+    assert list(tmp_path.glob(".output.safetensors.*.tmp")) == []