Compare commits

..
Author SHA1 Message Date
silveroxides 0a7578ce73 feat(analysis): add model similarity nodes
Separate CWB-derived diagnostics from merge weighting while streaming each logical layer once.
2026-08-04 15:52:42 +02:00
silveroxides 9ad992d080 feat(extract): register text encoder nodes 2026-08-03 18:57:57 +02:00
silveroxides 3bd4977d2a fix(cwb): keep preset scale neutral
Runtime conditioning scale reduction must not be carried into persisted model or LoRA weights. Keep named merger presets at 1.0 and reserve scaling for explicit custom configuration.
2026-08-03 18:48:14 +02:00
silveroxides 10484831d0 fix(package): use requirements manifest
Remove duplicate dependency metadata from pyproject.toml and rely on requirements.txt as the custom-node installation source.

Release as 0.5.2.
2026-08-02 20:43:32 +02:00
silveroxides 71b7936f92 fix(package): source registry requirements
Use requirements.txt as the single dependency source so Registry releases cannot drift from runtime installation constraints.

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

Fixes #15
2026-06-06 18:21:50 +08:00
Silver d2a81c0f9c Bump version from 0.4.2 to 0.4.3 2026-05-29 14:38:38 +02:00
Silver 2cb7f6c55e Merge pull request #14 from silveroxides/feature/dora-extract
Feature/dora extract
2026-05-29 14:37:56 +02:00
silveroxides 3a505e704f Add work in progress experimental learned extraction nodes 2026-05-29 14:36:26 +02:00
silveroxides 1512fa7390 Add DoRA extraction nodes 2026-05-28 03:27:47 +02:00
Silver b7365bfd79 Bump version from 0.4.1 to 0.4.2 2026-05-25 14:25:29 +02:00
Silver 34a2b78a9c Update requirements.txt 2026-05-25 14:25:00 +02:00
Silver 5780122057 Merge pull request #13 from silveroxides/refactor/incremental-writer
Refactor/incremental writer
2026-05-25 14:21:33 +02:00
Silver f6aac8cedf Merge branch 'main' into refactor/incremental-writer 2026-05-25 14:21:24 +02:00
silveroxides 7aa70d8173 Bump version from 0.4.0 to 0.4.1 2026-05-25 14:20:07 +02:00
silveroxides 26c8dffff9 Bump version from 0.3.8 to 0.4.0 2026-05-25 14:18:59 +02:00
silveroxides f3497658fb update unifiedefficientloader requirements 2026-05-25 14:17:58 +02:00
silveroxides 6747f2e8dc Fix block indices error 2026-05-25 14:16:59 +02:00
Silver 54bb6621d8 Bump version from 0.3.9 to 0.4.0 2026-05-21 11:14:38 +02:00
Silver d0d1063964 Update unifiedefficientloader version to 0.5.0 2026-05-21 11:14:19 +02:00
Silver 518e8e0b6d Bump version from 0.3.8 to 0.3.9 2026-05-16 16:33:59 +02:00
Silver fee30e36f5 Merge pull request #12 from silveroxides/refactor/incremental-writer
Refactor/incremental writer
2026-05-16 16:33:37 +02:00
silveroxides 94d2a93b91 fix layer renaming issues 2026-05-16 16:32:44 +02:00
silveroxides bb915f0599 refactor all saves to use IncrementalSafetensorsWriter 2026-05-16 11:07:20 +02:00
silveroxides f9a96d78a2 Bump version from 0.3.7 to 0.3.8 2026-05-10 18:52:38 +02:00
silveroxides 8ee4654c80 Improve layer filters in the merger nodes 2026-05-10 18:52:10 +02:00
silveroxides 73ce12cf9f Bump version from 0.3.6 to 0.3.7 2026-05-10 18:22:23 +02:00
silveroxides 21c0972a5b Refine LoRA extraction to output accurate layer formats across the board and add experimental new DARE+Ties merge operation variant 2026-05-10 18:21:54 +02:00
silveroxides 3bcd2bf3f8 Bump version from 0.3.5 to 0.3.6 2026-05-09 03:09:30 +02:00
silveroxides 29f03fe979 Add new 'enhanced' PowerUp/DARE-Ties merge modes 2026-05-09 03:08:57 +02:00
silveroxides 2e6eec7f63 Bump version from 0.3.4 to 0.3.5 2026-05-02 19:25:55 +02:00
silveroxides 2e0f79a396 handle Anima net. layer prefix 2026-05-02 19:25:20 +02:00
silveroxides 9f5d894320 Bump version from 0.3.3 to 0.3.4 2026-04-29 13:26:32 +02:00
silveroxides 10d0547e05 bump unifiedefficientloader requirement 2026-04-29 12:39:20 +02:00
silveroxides d2adfb3183 Bump version from 0.3.2 to 0.3.3 2026-04-27 11:49:18 +02:00
silveroxides 8fbbe8a5ae refactor to use only unifiedefficientloader and fix header loading 2026-04-27 11:48:53 +02:00
silveroxides 308d839ce3 Bump version from 0.3.1 to 0.3.2 2026-04-26 20:21:39 +02:00
silveroxides 7f6e2c4300 Add DARE+TIES Merge mode 2026-04-26 20:21:13 +02:00
silveroxides a8098b769e Bump version from 0.3.0 to 0.3.1 2026-04-26 13:23:17 +02:00
silveroxides ff1706a320 make merger go brrrrr 2026-04-26 12:41:19 +02:00
silveroxides 6dd89d4bbc Bump version to 0.3.0 2026-04-25 22:11:07 +02:00
silveroxides 4a2faa5a8e Add incremental saving when merging 2026-04-25 22:09:24 +02:00
silveroxides bb787378fe add usage of unifiedefficientloader module 2026-04-25 19:06:50 +02:00
Silver 95393aea13 Merge pull request #10 from silveroxides/feature/preview-meta-wf-loader
Feature/preview meta wf loader
2026-04-25 19:05:51 +02:00
silveroxides 0a82e071b8 fix NodeOutput issue 2026-04-08 17:50:10 +02:00
silveroxides dc43040562 Rename downloader nodes and add nodes for displaying previews, metadata of model and workflows extracted from images/videos 2026-04-08 16:46:35 +02:00
silveroxides 39380f533d Remove subprocess, use av. More info in status output from Downloader nodes 2026-04-08 16:28:45 +02:00
silveroxides 2ac0d46935 Add progress bar 2026-04-08 15:43:05 +02:00
silveroxides 378c327406 Add subdirectory selector 2026-04-08 15:31:42 +02:00
silveroxides 4076c14cb5 Add requirements.txt entries 2026-04-08 15:16:16 +02:00
silveroxides c4c967172b Add Model preview/example image/video downloading, model hashing, preview/example workflow extraction from images and videos followed by compression of images to save storage, saving of hash database and much more 2026-04-08 15:04:24 +02:00
Silver 095ac89164 Bump version to 0.2.9 2026-03-07 19:10:13 +01:00
Silver 2d477767d6 Merge pull request #9 from silveroxides/fix/loading-fixes
Fix models getting stuck in memory when merging
2026-03-07 19:09:32 +01:00
silveroxides 4d49477f58 Fix models getting stuck in memory when merging 2026-03-07 16:32:43 +01:00
Silver 6c4d505d84 Bump version from 0.2.7 to 0.2.8 2026-02-24 17:41:25 +01:00
Silver 1e53a87427 Merge pull request #8 from silveroxides/feature/expand-lora-functionality
expand on functionality and reduce restrictions when extracting LoRA
2026-02-24 17:40:47 +01:00
silveroxides b497022c3c expand on functionality and reduce restrictions when extracting LoRA 2026-02-24 17:39:48 +01:00
Silver 5df8c6711a Merge pull request #7 from silveroxides/fix/loading-imporovements
improve memory efficiency
2026-02-12 08:07:39 +01:00
Silver f4cecdc49b Merge branch 'main' into fix/loading-imporovements 2026-02-12 08:07:31 +01:00
Silver 8b34d57003 Merge pull request #6 from silveroxides/patch/lora-changes
changes to LoRAs
2026-02-12 08:05:19 +01:00
silveroxides 9f8b337584 changes to LoRAs 2026-02-07 16:02:40 +01:00
silveroxides 0d011d1232 improve memory efficiency 2026-02-07 16:01:00 +01:00
silveroxides 6a119e63f5 Bump version from 0.2.6 to 0.2.7 2026-01-23 19:21:02 +01:00
silveroxides 4097326a82 Add pad/crop for mismatched tensors and interpolation mode for weight redistribution across rank 2026-01-23 19:04:51 +01:00
silveroxides 31eee604de Bump version from 0.2.5 to 0.2.6 2026-01-23 18:54:09 +01:00
silveroxides 2bfecf9c2d Fix resize_tensors to handle cases where model B is larger than model A 2026-01-23 18:53:40 +01:00
silveroxides 1ba0932d6b Bump version from 0.2.4 to 0.2.5 2026-01-23 18:42:23 +01:00
silveroxides b1ec55ff42 improve and correct merger operations based on papers 2026-01-23 18:41:35 +01:00
silveroxides 8deb4fd739 perf: Add low_memory mode to reduce RAM usage in LoRA merge
- MemoryEfficientSafeOpen now supports low_memory=True which disables mmap
- Uses direct file reads with readinto() to avoid OS page caching
- _deserialize_tensor avoids bytearray copy when input is already bytearray
- merge_loras_to_model now uses low_memory=True for base model

This should reduce peak RAM from 2x model size to 1x + overhead
2026-01-15 03:13:51 +01:00
silveroxides e9f7cda706 fix: Free cpu_base after storing to avoid duplicate RAM usage 2026-01-15 02:51:17 +01:00
silveroxides cffb92aa06 fix: Add HuggingFace LoRA format support (.lora.down.weight/.lora.up.weight)
- Add detection for HuggingFace format in detect_lora_format()
- Strip trailing .lora suffix in extract_core_layer_lora() for proper key matching
- Applied fix to resize_lora_via_base, merge_multi_loras_via_base, and merge_loras_to_model
2026-01-15 02:43:42 +01:00
silveroxides ce8c6aab97 feat: Add LoRAMergeToModel node for merging LoRAs directly into model
- New merge_loras_to_model() function that merges LoRAs into base model
  weights and saves the result without extraction
- New LoRAMergeToModel node with simplified inputs (no extraction params)
- Outputs merged model to the base model's directory
2026-01-15 02:30:51 +01:00
silveroxides 6f3fd0900f add support for full diff LoRA in multi LoRA merge 2026-01-12 02:36:56 +01:00
Silver eaed2003a0 Bump version from 0.2.3 to 0.2.4 2026-01-12 01:23:09 +01:00
Silver 78192c594a Merge pull request #5 from silveroxides/feature/resize-with-base
add new LoRA features and up to 4 way LoRA resize+merge node.
2026-01-12 01:22:46 +01:00
silveroxides ede6fde5f3 add new feature to resize LoRA by first adding it to base model 2026-01-12 01:21:19 +01:00
Silver 0726ec952b Update pyproject.toml 2026-01-09 23:02:30 +01:00
silveroxides c6b65d6721 fix: frobenius/cumulative norm calculation now relative to max_rank 2026-01-03 20:55:27 +01:00
Silver 32281ad2d5 Bump version from 0.2.0 to 0.2.1 2026-01-03 20:38:38 +01:00
Silver 68fab8d3d0 Merge pull request #4 from silveroxides/feature/lora-resize
Feature/lora resize
2026-01-03 20:38:13 +01:00
silveroxides 61b2f506d0 feat: add svd_niter parameter for SVD accuracy control 2026-01-03 20:27:56 +01:00
silveroxides 75d1f0f5c9 perf: use svd_lowrank for fixed rank resize (10x faster) 2026-01-03 20:14:25 +01:00
silveroxides 3a71ad9952 feat: add max_rank inputs to dynamic extraction nodes 2026-01-03 20:12:58 +01:00
silveroxides 37090dd7ba feat: add LoRA resize nodes and fix sv_ratio convention 2026-01-03 20:04:56 +01:00
silveroxides f58289d042 document development phases 2026-01-03 19:37:39 +01:00
Silver 1796f1bea9 Merge pull request #3 from silveroxides/feature/lora-extract
Feature/lora extract and version bumped to 0.2.0
2026-01-02 17:01:41 +01:00
Silver 3df073f810 Bump version from 0.1.7 to 0.2.0 2026-01-02 17:01:02 +01:00
silveroxides 0e73c5d889 Update README.md, add images and workflows 2026-01-02 16:59:25 +01:00
silveroxides eef7377061 Added .contiguous() to saves 2026-01-01 21:18:08 +01:00
silveroxides 65c1ae37b7 add svd_lowrank for fixed rank extraction and put experimental chunk large layers default to false 2026-01-01 21:08:15 +01:00
silveroxides b7707862b2 Add pinned memory for faster merging and extraction 2026-01-01 20:49:49 +01:00
silveroxides 6d7b2d1403 refactor nodes and implement memory management 2026-01-01 20:29:08 +01:00
silveroxides bfbbcd9515 fix exception handler 2026-01-01 19:04:25 +01:00
silveroxides 18679cca2c Split lycoris extraction modes into different nodes 2026-01-01 07:21:37 +01:00
silveroxides da13f9f51d add lycoris diff extraction 2026-01-01 07:06:19 +01:00
silveroxides 92d3c5582c add lora extraction, progress bars and more 2026-01-01 06:36:59 +01:00
Silver bb923e8474 Merge pull request #2 from silveroxides/feature/v3-migration
migrate to v3 schema and add LoRA extraction node
2026-01-01 06:35:56 +01:00
silveroxides a3116cbb11 migrate to v3 schema and add LoRA extraction node 2026-01-01 05:43:24 +01:00
silveroxides b6fae8d0d9 Bump version to 0.1.7 2025-12-21 20:32:19 +01:00
silveroxides de27186b90 merge feature mismatch-handling 2025-12-21 20:31:17 +01:00
silveroxides 1131811117 add layer filtering 2025-12-18 22:23:33 +01:00
silveroxides c5df75dce8 fix handling of mismatched layers 2025-12-18 22:04:04 +01:00
silveroxides 19bc0472dc add handling of mismatched layers 2025-12-18 21:43:32 +01:00
silveroxides 6c08202ef5 update .gitignore 2025-12-18 21:22:10 +01:00
silveroxides 68f22fa87b fix Extract namespace conflict 2025-11-24 19:16:03 +01:00
silveroxides 5ea63abfbe fix PowerUp namespace conflict 2025-11-24 19:13:45 +01:00
silveroxides 232dc76db2 refactor names to be more universal 2025-11-24 18:37:18 +01:00
silveroxides 31989beefd Bump version to 0.1.6 2025-11-09 06:57:51 +01:00
silveroxides 0d035ef072 Add embedding nodes 2025-11-09 06:57:23 +01:00
silveroxides 78b3198c14 Bump version to 0.1.5 2025-10-19 10:20:26 +02:00
silveroxides 79593701aa Add model merging nodes 2025-10-19 10:19:52 +02:00
silveroxides cbe3502e83 Bump version to 0.1.4 2025-10-12 05:26:04 +02:00
silveroxides 4f08fe69b2 add key renaming and key pruning nodes 2025-10-12 05:25:27 +02:00
47 changed files with 19021 additions and 155 deletions
+1
View File
@@ -0,0 +1 @@
tests/
+4
View File
@@ -1,3 +1,7 @@
# Repo specific
AGENTS.MD
custom_nodes.examples.md
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
+54 -1
View File
@@ -1,2 +1,55 @@
# ComfyUI-ModelUtils
[WIP]Custom nodes for handling, inspecting, modifying and creating various model files.
A collection of ComfyUI custom nodes for inspecting, modifying, merging, and creating model files. Supports Models, TextEncoders, LoRAs, Checkpoints, and Embeddings.
## Features
- **MetaKeys** – Inspect and display metadata and tensor keys from model files
- **RenameKeys** – Batch rename tensor keys using pattern matching
- **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)
## Example Workflows
<p align="center">
<img src="assets/GetMetaAndkeys.png" width="400" alt="Get Meta and Keys">
<br>
<a href="example_workflows/GetMetaAndkeys.json">📥 GetMetaAndkeys.json</a>
</p>
---
<p align="center">
<img src="assets/LoRA_Extract_nodes.png" width="400" alt="LoRA Extraction Nodes">
<br>
<a href="example_workflows/LoRA_Extract_nodes.json">📥 LoRA_Extract_nodes.json</a>
</p>
---
<p align="center">
<img src="assets/Merging_Examples.png" width="400" alt="Merging Examples">
<br>
<a href="example_workflows/Merging_Examples.json">📥 Merging_Examples.json</a>
</p>
---
<p align="center">
<img src="assets/RenameKeysInModel.png" width="400" alt="Rename Keys in Model">
<br>
<a href="example_workflows/RenameKeysInModel.json">📥 RenameKeysInModel.json</a>
</p>
## Acknowledgements
The LoRA extraction functionality was inspired by and references the excellent work from:
- [kohya-ss/sd-scripts](https://github.com/kohya-ss/sd-scripts) – Training scripts for Stable Diffusion
- [KohakuBlueleaf/LyCORIS](https://github.com/KohakuBlueleaf/LyCORIS) – Advanced LoRA techniques
- [bmaltais/kohya_ss](https://github.com/bmaltais/kohya_ss) – Windows-friendly GUI for sd-scripts
## License
See [LICENSE](LICENSE) for details.
+117 -15
View File
@@ -1,17 +1,119 @@
from .nodes import *
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
NODE_CLASS_MAPPINGS = {
"UNetMetaKeys": UNetMetaKeys,
"CLIPMetaKeys": CLIPMetaKeys,
"LoRAMetaKeys": LoRAMetaKeys,
"CheckpointMetaKeys": CheckpointMetaKeys,
"OffloadModels": OffloadModelsNode,
}
from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced
from .nodes.metakeys import (
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys
)
from .nodes.renamekeys import (
ModelRenameKeys, TextEncoderRenameKeys, LoRARenameKeys,
CheckpointRenameKeys, EmbeddingRenameKeys
)
from .nodes.prunekeys import (
ModelPruneKeys, TextEncoderPruneKeys, LoRAPruneKeys,
CheckpointPruneKeys, EmbeddingPruneKeys
)
from .nodes.merger import (
ModelTwoMerger, TextEncoderTwoMerger, LoRATwoMerger,
CheckpointTwoMerger, EmbeddingTwoMerger,
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger
)
from .nodes.consensus_merger import CWB_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 (
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
LoRAMergeToModel
)
from .nodes.downloader_nodes import (
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader
)
from .nodes.model_info_nodes import (
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader
)
from .nodes.lora_rename import AnimaLoraRename
class ModelUtilsExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
# MetaKeys
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys,
# RenameKeys
ModelRenameKeys, TextEncoderRenameKeys, LoRARenameKeys,
CheckpointRenameKeys, EmbeddingRenameKeys,
# PruneKeys
ModelPruneKeys, TextEncoderPruneKeys, LoRAPruneKeys,
CheckpointPruneKeys, EmbeddingPruneKeys,
# Two-Model Mergers
ModelTwoMerger, TextEncoderTwoMerger, LoRATwoMerger,
CheckpointTwoMerger, EmbeddingTwoMerger,
# Three-Model Mergers
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger,
# Dedicated CWB Mergers
*CWB_MERGER_NODES,
# Two-Model Analysis
*MODEL_ANALYSIS_NODES,
# 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
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
# LoRA Multi-Merge
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
# LoRA Merge To Model
LoRAMergeToModel,
# LoRA Utilities
AnimaLoraRename,
# Downloaders
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
# Info Loaders
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader,
]
async def comfy_entrypoint() -> ModelUtilsExtension:
return ModelUtilsExtension()
NODE_DISPLAY_NAME_MAPPINGS = {
"UNetMetaKeys": "Get UNet Metadata and Keys",
"CLIPMetaKeys": "Get CLIP Metadata and Keys",
"LoRAMetaKeys": "Get LoRA Metadata and Keys",
"CheckpointMetaKeys": "Get Checkpoint Metadata and Keys",
"OffloadModelsNode": "Offload Models with Passthrough",
}
Binary file not shown.

After

Width:  |  Height:  |  Size: 352 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 600 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 592 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 427 KiB

+321
View File
@@ -0,0 +1,321 @@
# ComfyUI-ModelUtils - Complete Project History
## Version Timeline
| Version | Commits | Key Features |
|---------|---------|--------------|
| 0.1.1 | `70d6c13`-`d67cbca` | Initial nodes (MetaKeys) |
| 0.1.2 | `6b662a4` | Icon assets |
| 0.1.4 | `cbe3502` | Key renaming/pruning nodes |
| 0.1.5 | `78b3198` | Model merging nodes |
| 0.1.6 | `31989be` | Embedding nodes, namespace fixes |
| 0.1.7 | `b6fae8d` | Mismatch handling, layer filtering, V3 migration |
| **0.2.0** | `3df073f` | LoRA extraction, memory optimization, pinned memory |
---
## Phase 1: Initial Release (v0.1.1)
**Commit:** `70d6c13` - "Add first nodes, toml and publish"
### Files Created
- `nodes/metakeys.py` - Get model metadata and layer keys
- `nodes/utils.py` - Shared utilities
- `__init__.py` - Node registration
- `pyproject.toml` - Package configuration
### Nodes
- `ModelMetaKeys` - Get diffusion model metadata
- `CheckpointMetaKeys` - Get checkpoint metadata
- `TextEncoderMetaKeys` - Get text encoder metadata
- `LoRAMetaKeys` - Get LoRA metadata
- `EmbeddingMetaKeys` - Get embedding metadata
---
## Phase 2: Key Manipulation (v0.1.4)
**Commit:** `4f08fe6` - "add key renaming and key pruning nodes"
### Files Created
- `nodes/renamekeys.py` - Rename tensor keys
- `nodes/prunekeys.py` - Remove tensor keys by pattern
### Nodes (5 variants each for model types)
- `ModelRenameKeys`, `CheckpointRenameKeys`, `TextEncoderRenameKeys`, `LoRARenameKeys`, `EmbeddingRenameKeys`
- `ModelPruneKeys`, `CheckpointPruneKeys`, `TextEncoderPruneKeys`, `LoRAPruneKeys`, `EmbeddingPruneKeys`
---
## Phase 3: Model Merging (v0.1.5)
**Commit:** `7959370` - "Add model merging nodes"
### Files Created
- `nodes/merger.py` - Main merging logic
- `nodes/merger_ops.py` - Merge operation modes
- `nodes/merger_utils.py` - `MemoryEfficientSafeOpen` class
- `docs/merger_2_model_modes.md` - Documentation
- `docs/merger_3_model_modes.md` - Documentation
### Nodes
- `CheckpointTwoMerger`, `DiffusionModelTwoMerger`, `TextEncoderTwoMerger`, `LoRATwoMerger`
- `CheckpointThreeMerger`, `DiffusionModelThreeMerger`, `TextEncoderThreeMerger`, `LoRAThreeMerger`
### Merge Modes (TWO_MODEL_MODES)
- Weighted Sum, Add Difference, Smooth Add Difference
- Train Difference, Sum Twice, Triple Sum
- Euclidean Add Difference, Distribution Crossover, Ties
- Slerp, Lerp, SVD Supermerger, SVD LoRA Extraction
### Merge Modes (THREE_MODEL_MODES)
- Add Difference (A + (B - C) × α)
### MemoryEfficientSafeOpen (initial)
```python
class MemoryEfficientSafeOpen:
def __init__(self, filename, device='cpu', mmap_mode=True)
def keys() -> List[str]
def get_tensor(key) -> Tensor
# mmap for zero-copy reads
```
---
## Phase 4: Embedding Nodes (v0.1.6)
**Commit:** `0d035ef` - "Add embedding nodes"
### Additions
- Added embedding variants to all node types
- Namespace conflict fixes (`68f22fa`, `5ea63ab`)
- Refactored names for universality (`232dc76`)
---
## Phase 5: Mismatch Handling (v0.1.7)
**Commits:** `19bc047` → `de27186`
### Features Added
- `mismatch_mode` parameter: skip/zeros/error
- Layer filtering with regex patterns
- `exclude_patterns` - Keep model A only (no merge)
- `discard_patterns` - Remove from output entirely
### Files Modified
- `merger.py` - Added mismatch handling logic
- `merger_ops.py` - Added `MissingTensorBehavior` enum, `MissingTensorError`
---
## Phase 6: V3 API Migration (v0.1.7)
**Commit:** `a3116cb` - "migrate to v3 schema and add LoRA extraction node"
### Changes
- Migrated all nodes from legacy API to `comfy_api.latest.io`
- `io.Schema`, `io.ComfyNode`, `io.NodeOutput`
- Added initial LoRA extraction node (diff-based)
---
## Phase 7: LyCORIS Extraction (pre-0.2.0)
**Commits:** `92d3c55` → `18679cc`
### Files Created
- `nodes/lycoris_extract.py` - LyCORIS-based extraction
- `nodes/lora_extract.py` - Simple diff extraction
### Initial Implementation
- Used `lycoris-lora` external dependency
- `lycoris.utils.extract_linear`, `extract_conv`
- Added `requirements.txt` with `lycoris-lora`
### Nodes Split (`18679cc`)
- `LyCORISExtractFixed`, `LyCORISExtractThreshold`
- `LyCORISExtractRatio`, `LyCORISExtractQuantile`
---
## Phase 8: Native SVD Port (0.2.0)
**Commits:** `6d7b2d1` → `eef7377`
### Removed
- `lycoris-lora` dependency
- `lycoris_extract.py`
- `lora_extract.py`
### NEW: [lora_extract_svd.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/lora_extract_svd.py)
**8 Rank Selection Functions:**
```python
_index_sv_fixed(S, dim) # Fixed rank
_index_sv_threshold(S, threshold) # Absolute threshold
_index_sv_ratio(S, ratio) # Relative to max
_index_sv_cumulative(S, target) # Cumulative sum %
_index_sv_fro(S, target) # Frobenius norm
_index_sv_knee(S) # Knee detection
_index_sv_cumulative_knee(S) # Cumsum knee
_index_sv_rel_decrease(S, tau) # Ratio drop
```
**SVD Functions:**
```python
_svd_extract_linear_lowrank(weight, rank, device, clamp_quantile)
_svd_extract_linear(weight, mode, mode_param, device, max_rank, clamp_quantile)
_svd_extract_conv(weight, mode, mode_param, device, max_rank, clamp_quantile)
_detect_fused_layer(key, shape) # QKV/MLP detection
_extract_chunked_layer(...) # Chunked extraction
```
**Nodes:**
- `LoRAExtractFixed` - Uses `svd_lowrank` (10x faster)
- `LoRAExtractRatio` - Relative threshold
- `LoRAExtractQuantile` - Cumulative target
- `LoRAExtractKnee` - Auto knee detection
- `LoRAExtractFrobenius` - Norm preservation
---
## Phase 9: Memory Management (0.2.0)
**Commit:** `6d7b2d1` - "refactor nodes and implement memory management"
### NEW: [device_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/device_utils.py)
Integrates with `comfy.model_management`:
```python
@dataclass
class DeviceCapabilities:
device_type, total_vram_gb, free_vram_gb, total_ram_gb, free_ram_gb
cpu_cores, recommended_workers, max_tensor_size_gb
supports_pinned_memory, supports_async_streams
get_device_capabilities(device) -> DeviceCapabilities
estimate_model_size(path) -> float # GB
can_fit_in_memory(paths, device, safety_factor) -> (bool, str)
prepare_for_large_operation(gb, device) -> bool # Calls free_memory()
cleanup_after_operation() # gc.collect() + soft_empty_cache()
get_optimal_workers(size_gb) -> int
```
### Updates to Existing Nodes
- `metakeys.py` - Header-only reading (no tensor loading)
- `renamekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
- `prunekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
- `merger.py` - Memory prep/cleanup
---
## Phase 10: Pinned Memory (0.2.0)
**Commit:** `b770786` - "Add pinned memory for faster merging and extraction"
### [merger_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/merger_utils.py) Enhancements
```python
# Context manager
class PinnedTensor:
def __enter__(self) -> Tensor # Pins memory
def __exit__(...) # Unpins
# Transfer helper (2-3x faster)
def transfer_to_gpu_pinned(tensor, device='cuda', dtype=None) -> Tensor
# MemoryEfficientSafeOpen additions
def get_tensor_to_gpu(key, device, dtype) # Pinned transfer
def keys_sorted_by_offset() -> List[str] # Sequential I/O
def get_tensors_parallel(keys, workers, sort_by_offset) -> Dict # Parallel
def get_tensor_as_dict(key) -> Dict # JSON config parsing
```
---
## Phase 11: SVD Optimization (0.2.0)
**Commit:** `65c1ae3` - "add svd_lowrank for fixed rank extraction"
### Changes
- Added `_svd_extract_linear_lowrank()` using `torch.svd_lowrank`
- Routes fixed mode to lowrank for **10x speedup**
- Changed `chunk_large_layers` default to `False`
---
## Phase 12: Bug Fixes (0.2.0)
**Commit:** `eef7377` - "Added .contiguous() to saves"
### Fix
Non-contiguous tensor error from SVD slices:
```python
tensor.to(dtype).cpu().contiguous() # Added .contiguous()
```
---
## Current File Structure
```
ComfyUI-ModelUtils/
├── __init__.py
├── pyproject.toml
├── requirements.txt (empty - no external deps)
├── README.md
├── docs/
│ ├── merger_2_model_modes.md
│ └── merger_3_model_modes.md
├── nodes/
│ ├── utils.py
│ ├── metakeys.py
│ ├── renamekeys.py
│ ├── prunekeys.py
│ ├── merger.py
│ ├── merger_ops.py
│ ├── merger_utils.py
│ ├── device_utils.py
│ └── lora_extract_svd.py
└── assets/
└── icon.png
```
---
## Performance Summary
| Feature | Improvement |
|---------|-------------|
| Fixed-rank SVD | **10x** (svd_lowrank) |
| CPU→GPU transfer | **2-3x** (pinned memory) |
| Metadata reading | **100x+** (header only) |
| Parallel I/O | **2-4x** (ThreadPoolExecutor) |
| Memory management | Auto-unload models before heavy ops |
| **LoRA Resize Fast** | **~91x** (merge-to-base extraction) |
---
## Phase 13: Fast LoRA Resize (2026-01-11)
### Summary
Added fast LoRA resize via base model extraction path, achieving ~91x speedup.
### Changes
- Added `resize_lora_via_base()` function using extraction code path
- Added 4 new fast resize nodes requiring base model input
### Files Modified
| File | Changes |
|------|---------|
| `nodes/lora_resize.py` | Added `resize_lora_via_base()`, 4 new node classes |
| `__init__.py` | Registered new nodes |
### New Nodes
- `LoRAResizeViaBaseFixed` - Fast fixed rank resize
- `LoRAResizeViaBaseRatio` - Fast SV ratio resize
- `LoRAResizeViaBaseFrobenius` - Fast Frobenius resize
- `LoRAResizeViaBaseCumulative` - Fast cumulative resize
+111
View File
@@ -0,0 +1,111 @@
# Consensus-Weighted Blending Mergers
The dedicated CWB nodes merge two or three safetensors files using
Consensus-Weighted Blending. They calculate an element-wise mean or median
consensus, compare each aligned source vector to that consensus, and normalize
similarity-derived weights before producing the output vector.
## Node families
Two-input and three-input variants are provided for checkpoints, standalone
diffusion models, text encoders, LoRAs, and embeddings.
Model A supplies output metadata and anchors shared-layer shape, naming, and
mismatch behavior. The output key set is the union of all inputs. A valid
tensor found only in a secondary input is copied unchanged. When Model A lacks
a tensor that is present in multiple secondary inputs, those available sources
are CWB-merged without inserting a zero contribution for Model A. Embedding
mergers additionally use the longest compatible first dimension as their
alignment reference.
## Presets and custom controls
`baseline` is the default. The standard clarity, smooth, varied, and diversity
presets are joined by `power_blend` and the DSC presets demonstrated by the CWB
conditioning implementation. Select `custom` to make every manual CWB control
authoritative. Every named merger preset keeps `global_scale` at `1.0`; scaling
persisted model or LoRA weights is only performed when explicitly selected in
the custom controls.
Similarity alignment greedily matches source rows to the reference by cosine
similarity. Index alignment keeps absolute row positions. Position weight can
bias similarity matches toward nearby normalized positions. Preserve Common
Prefix copies the longest numerically identical leading row span directly from
the first input.
Inputs have equal prior weight. The similarity calculation supplies the
effective contribution for each aligned vector. If every candidate is rejected
or receives zero weight, CWB falls back to equal weighting.
## Missing tensors and filtering
- `skip` preserves Model A when an anchored input is missing or incompatible.
In embedding-union mode it omits that source for the affected key.
- `zeros` supplies a compatible zero tensor for the missing source.
- `error` stops the merge.
- Exclude patterns preserve the anchored source tensor.
- Discard patterns remove the matching output tensor or logical LoRA layer.
The missing-input modes above apply to Model-A-owned layers. Model A being
absent from a secondary-only layer is not itself a mismatch. A shape conflict
between multiple available secondary providers does use the selected mismatch
mode, with the earliest available input as its preservation anchor.
Patterns use regular expressions by default. Enable Glob Patterns to use glob
syntax instead.
## LoRA behavior
LoRA inputs use the same parser as the resize and multi-merge nodes. It accepts
all seven low-rank A/B spellings recognized by ComfyUI's `LoRAAdapter`, plus
direct `.diff`, `.diff_b`, `.w_norm`, `.b_norm`, and `.set_weight` layers.
Recognized diffusion-model outputs use the canonical
`diffusion_model.<layer>.lora_A.weight` / `.lora_B.weight` convention.
Secondary formats are matched by normalized logical layer name. Differing
ranks are zero-padded like the DARE/TIES mergers: A/down tensors on dimension 0
and B/up tensors on dimension 1. An existing anchor-source alpha key is updated
to the resulting maximum rank; no alpha key is invented when the anchor has
none.
Similarity alignment operates on paired latent-rank components. A rows and B
columns share one mapping; fixed input/output feature axes are never reordered.
Input alpha scaling is absorbed into the B factor before blending, and global
scale is applied once to the resulting pair rather than once per factor.
Incomplete pairs and unrecognized tensors are copied from the earliest
available source unchanged.
Companion-bearing groups (`lora_mid`, reshape, DoRA, or set-weight) are kept
atomically from their earliest available source rather than partially
transformed. Secondary-only logical groups are included automatically: a group
from one source is copied unchanged, while compatible groups from multiple
secondary sources are merged using only those sources. Recognized roles are
written with the same canonical output convention as shared layers.
`include_1d_diffs` is disabled by default. Disabled shared 1D direct layers
preserve Model A, while secondary-only layers preserve their earliest source.
Enabled 1D direct layers participate as complete tensors and are always saved
in FP32.
## Precision and quantized inputs
CWB computes in FP32. Unless Override Dtype is enabled, any participating FP32
source keeps the generated logical layer in FP32. Guarded low-bit tensors and
enabled 1D direct diffs are never downcast by the override.
ComfyUI/comfy-kitchen quantized inputs are rejected before an output writer is
opened. Files with three or more bare INT8, UINT8, FP8, or FP4 tensors are also
rejected. One or two isolated low-bit tensors produce one warning per input and
are preserved without CWB arithmetic; an affected LoRA causes its complete
logical layer to be preserved.
## Streaming and failure safety
Model parameters and direct patches use fixed-coordinate blending; similarity
row matching is limited to embeddings and LoRA rank components. Each logical
layer is loaded, transferred to the processing device, completed, copied back
to CPU, queued for writing, and released before the next layer. Outputs are
written to a temporary sibling and atomically published only after every layer
and asynchronous write succeeds. Failed merges do not replace or finalize the
requested output. A CUDA OOM retries only the affected layer on CPU after
releasing the failed GPU working set; failure of that retry aborts atomically.
Generated files are published under the matching category inside
`folder_paths.models_dir`, never ComfyUI's general output directory.
+155
View File
@@ -0,0 +1,155 @@
# 2-Model Merger Documentation
---
## Weight-Sum
> `A * (1 - α) + B * α`. A simple linear interpolation between two models.
**Models Used:** A, B
**Parameters:**
- **Alpha:** Interpolation weight. `0.0` is 100% Model A, `1.0` is 100% Model B.
---
## 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.
**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`.
---
## 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.
**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.
---
## Power-Up (DARE+TIES)
> Combines DARE (Drop and Rescale) with TIES (Trim, Elect Sign) — the current best-practice algorithm for merging fine-tuned models, particularly LoRA-tuned ones. DARE sparsifies the task vector via random dropout and rescaling; TIES then removes residual low-magnitude noise and enforces sign consistency. The result is a cleaner, more targeted capability transfer from B to A with less interference.
**Models Used:** A, B
**Parameters:**
- **Alpha:** DARE drop rate. Fraction of delta parameters randomly zeroed. Higher values produce a sparser, more targeted delta. Typical range: `0.5`–`0.9`.
- **Beta:** TIES trim quantile. Fraction of the smallest-magnitude delta parameters zeroed after DARE, removing residual noise. `0.0` disables trimming. Typical range: `0.1`–`0.3`.
- **Gamma:** Lambda scale. Final multiplier applied to the merged delta before adding to Model A. Equivalent to the task arithmetic scaling coefficient. Typical range: `0.5`–`1.5`.
- **Seed:** Random seed for the DARE dropout mask. Controls reproducibility.
**Algorithm:**
1. Compute task vector `δ = B − A`
2. **DARE:** apply random binary mask with keep probability `1 − alpha`; rescale survivors by `1/(1−alpha)`
3. **TIES trim:** zero out parameters below the `beta`-quantile magnitude threshold
4. **TIES elect:** determine dominant sign per position; zero out disagreeing parameters
5. Return `A + gamma × δ_filtered`
**Guidance:**
- Start with `alpha=0.9, beta=0.2, gamma=1.0` for LoRA-derived models
- Lower `alpha` (e.g. `0.5`) for full fine-tunes where more delta mass should be retained
- Increase `gamma` beyond `1.0` to amplify B's influence; decrease below `1.0` to soften it
---
## Enhanced Man Interp
> Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control.
**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).
---
## Enhanced Auto Interp
> Automated version of the enhanced interpolation mode that dynamically calculates thresholds based on mean differences.
**Models Used:** A, B
**Parameters:**
- **Alpha:** Interpolation strength.
- **Beta:** Threshold adjustment factor.
- **Gamma:** Smoothness factor.
---
## Weight-Sum Cutoff
> A linear interpolation mode that only applies the merge to weights whose differences fall within a specific threshold range.
**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.
---
## 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.
**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.
---
## Layer Mismatch Handling
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:
| 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 |
**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.
---
## Dtype Preservation & Override
By default, the merger preserves the highest precision `dtype` for each key found across all active source models (Model A and Model B) and the requested `save_dtype`.
- **Preserve Higher Precision**: If a key is `float32` in any of the source models, but you requested `save_dtype="bf16"` or `"fp16"`, that key will be saved as `float32` anyway. This prevents degrading critical high-precision layers (such as normalization layers or layers with restored outliers).
- **Explicit Override**: If you prefer to force the entire model to be saved strictly as the selected `save_dtype` (e.g. for complete downcasting to `bf16`), toggle `override_dtype` to `True`.
---
## Layer Filtering (Regex Patterns)
Use regex patterns to control which layers are merged:
### `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.
**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 |
+92
View File
@@ -0,0 +1,92 @@
# 3-Model Merger Documentation
---
## 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).
**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.
---
## 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.
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** A multiplier for the final calculated "training" difference.
---
## 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.
**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.
---
## 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.
**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.
---
## Layer Mismatch Handling
When merging models with different layer structures, the `mismatch_mode` parameter controls behavior:
| 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 |
**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.
---
## Dtype Preservation & Override
By default, the merger preserves the highest precision `dtype` for each key found across all active source models (Model A, B, and C) and the requested `save_dtype`.
- **Preserve Higher Precision**: If a key is `float32` in any of the active source models, but you requested `save_dtype="bf16"` or `"fp16"`, that key will be saved as `float32` anyway. This prevents degrading critical high-precision layers (such as normalization layers or layers with restored outliers).
- **Explicit Override**: If you prefer to force the entire model to be saved strictly as the selected `save_dtype` (e.g. for complete downcasting to `bf16`), toggle `override_dtype` to `True`.
---
## Layer Filtering (Regex Patterns)
Use regex patterns to control which layers are merged:
### `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.
**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"
+75
View File
@@ -0,0 +1,75 @@
# 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.
## Outputs
`comparison_report` contains conventional numerical differences and topology
findings. `cwb_report` separately contains diagnostics derived from the
similarity, alignment, and consensus stages of Consensus-Weighted Blending.
`documentation` contains this reference.
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.
+600
View File
@@ -0,0 +1,600 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 31,
"last_link_id": 19,
"nodes": [
{
"id": 16,
"type": "PreviewAny",
"pos": [
-3608,
200
],
"size": [
832,
232
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 18
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 17,
"type": "ModelMetaKeys",
"pos": [
-2800,
504
],
"size": [
346.253515625,
78
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
13
]
},
{
"name": "keys",
"type": "STRING",
"links": [
17
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "ModelMetaKeys"
},
"widgets_values": [
"flux2-dev.safetensors"
]
},
{
"id": 18,
"type": "TextEncoderMetaKeys",
"pos": [
-3176,
680
],
"size": [
329.4234375,
78
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
15
]
},
{
"name": "keys",
"type": "STRING",
"links": [
16
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "TextEncoderMetaKeys"
},
"widgets_values": [
"clip_l.safetensors"
]
},
{
"id": 19,
"type": "CheckpointMetaKeys",
"pos": [
-3168,
512
],
"size": [
317.6724609375,
78
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
12
]
},
{
"name": "keys",
"type": "STRING",
"links": [
18
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "CheckpointMetaKeys"
},
"widgets_values": [
"unsorted\\v1-5-pruned-emaonly.safetensors"
]
},
{
"id": 21,
"type": "LoRAMetaKeys",
"pos": [
-2800,
680
],
"size": [
277.0943359375,
78
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
14
]
},
{
"name": "keys",
"type": "STRING",
"links": [
19
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAMetaKeys"
},
"widgets_values": [
"Chroma\\FlatAnimeStyle_Chroma.safetensors"
]
},
{
"id": 24,
"type": "PreviewAny",
"pos": [
-2728,
216
],
"size": [
512,
200
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 17
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 25,
"type": "PreviewAny",
"pos": [
-2816,
832
],
"size": [
736,
248
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 19
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 27,
"type": "PreviewAny",
"pos": [
-3496,
848
],
"size": [
632,
224
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 16
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 28,
"type": "PreviewAny",
"pos": [
-3600,
664
],
"size": [
352,
136
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 15
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 29,
"type": "PreviewAny",
"pos": [
-2408,
648
],
"size": [
360,
120
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 14
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 30,
"type": "PreviewAny",
"pos": [
-3576,
480
],
"size": [
336,
128
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 12
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 31,
"type": "PreviewAny",
"pos": [
-2368,
472
],
"size": [
328,
120
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 13
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
}
],
"links": [
[
12,
19,
0,
30,
0,
"STRING"
],
[
13,
17,
0,
31,
0,
"STRING"
],
[
14,
21,
0,
29,
0,
"STRING"
],
[
15,
18,
0,
28,
0,
"STRING"
],
[
16,
18,
1,
27,
0,
"STRING"
],
[
17,
17,
1,
24,
0,
"STRING"
],
[
18,
19,
1,
16,
0,
"STRING"
],
[
19,
21,
1,
25,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.9070294784580497,
"offset": [
3855.323424397858,
-56.106573378634465
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+522
View File
@@ -0,0 +1,522 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 10,
"last_link_id": 5,
"nodes": [
{
"id": 1,
"type": "LoRAExtractFixed",
"pos": [
-1848,
16
],
"size": [
296,
352
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
3
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractFixed"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
64,
32,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 2,
"type": "LoRAExtractFrobenius",
"pos": [
-2432,
24
],
"size": [
296,
352
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
1
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractFrobenius"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
0.9,
0.9,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 3,
"type": "LoRAExtractKnee",
"pos": [
-2416,
440
],
"size": [
280,
328
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
2
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractKnee"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
"sv_knee",
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 4,
"type": "LoRAExtractQuantile",
"pos": [
-1248,
128
],
"size": [
288,
352
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
5
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractQuantile"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
0.9,
0.9,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 5,
"type": "LoRAExtractRatio",
"pos": [
-1824,
424
],
"size": [
296,
352
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
4
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractRatio"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
0.5,
0.5,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 6,
"type": "PreviewAny",
"pos": [
-2096,
136
],
"size": [
210,
168
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 1
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 7,
"type": "PreviewAny",
"pos": [
-2096,
448
],
"size": [
210,
168
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 2
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 8,
"type": "PreviewAny",
"pos": [
-1504,
96
],
"size": [
210,
168
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 3
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 9,
"type": "PreviewAny",
"pos": [
-1472,
432
],
"size": [
210,
168
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 4
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 10,
"type": "PreviewAny",
"pos": [
-904,
184
],
"size": [
210,
168
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 5
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
}
],
"links": [
[
1,
2,
0,
6,
0,
"STRING"
],
[
2,
3,
0,
7,
0,
"STRING"
],
[
3,
1,
0,
8,
0,
"STRING"
],
[
4,
5,
0,
9,
0,
"STRING"
],
[
5,
4,
0,
10,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.05,
"offset": [
2532,
84
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+460
View File
@@ -0,0 +1,460 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 16,
"last_link_id": 11,
"nodes": [
{
"id": 6,
"type": "PreviewAny",
"pos": [
-3552,
136
],
"size": [
336,
160
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 6
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 7,
"type": "PreviewAny",
"pos": [
-2792,
144
],
"size": [
456,
464
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 9
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
false
]
},
{
"id": 11,
"type": "ModelThreeMerger",
"pos": [
-3120,
112
],
"size": [
296,
522
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": null
},
{
"name": "documentation",
"type": "STRING",
"links": [
9
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "ModelThreeMerger"
},
"widgets_values": [
"DOCUMENTATION ONLY",
"Chroma1-HD.safetensors",
"Chroma-DC-2K.safetensors",
"None",
"Add-Difference",
"skip",
0.5,
0.5,
0.5,
0.5,
996811747788113,
"randomize",
"merged_3_model",
"fp32",
"cuda",
"",
""
]
},
{
"id": 12,
"type": "ModelTwoMerger",
"pos": [
-3120,
704
],
"size": [
304,
474
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": [
11
]
},
{
"name": "documentation",
"type": "STRING",
"links": null
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "ModelTwoMerger"
},
"widgets_values": [
"MERGE",
"None",
"None",
"Weight-Sum",
"skip",
0.5,
0.5,
0.99,
388546580517523,
"randomize",
"merged_2_model",
"fp32",
"cuda",
"",
""
]
},
{
"id": 13,
"type": "LoRAThreeMerger",
"pos": [
-3872,
128
],
"size": [
296,
522
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": [
6
]
},
{
"name": "documentation",
"type": "STRING",
"links": []
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAThreeMerger"
},
"widgets_values": [
"MERGE",
"None",
"None",
"None",
"Add-Difference",
"skip",
0.5,
0.5,
0.5,
0.5,
1120102695669269,
"randomize",
"merged_3_lora",
"fp32",
"cuda",
"",
""
]
},
{
"id": 14,
"type": "LoRATwoMerger",
"pos": [
-3872,
720
],
"size": [
296,
474
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": []
},
{
"name": "documentation",
"type": "STRING",
"links": [
10
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRATwoMerger"
},
"widgets_values": [
"DOCUMENTATION ONLY",
"Hyper-FLUX.1-dev-16steps-lora-rescale_z1.safetensors",
"Hyper-FLUX.1-dev-16steps-lora-rescale_z2.safetensors",
"Weight-Sum",
"skip",
0.5,
0.5,
0.99,
1020143314396572,
"randomize",
"merged_2_lora",
"fp32",
"cuda",
"",
""
]
},
{
"id": 15,
"type": "PreviewAny",
"pos": [
-2768,
720
],
"size": [
432,
160
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 11
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 16,
"type": "PreviewAny",
"pos": [
-3544,
712
],
"size": [
400,
448
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 10
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
false
]
}
],
"links": [
[
6,
13,
0,
6,
0,
"STRING"
],
[
9,
11,
1,
7,
0,
"STRING"
],
[
10,
14,
1,
16,
0,
"STRING"
],
[
11,
12,
0,
15,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.9070294784580494,
"offset": [
3972,
-12
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+583
View File
@@ -0,0 +1,583 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 46,
"last_link_id": 31,
"nodes": [
{
"id": 35,
"type": "LoRAMetaKeys",
"pos": [
-3136,
512
],
"size": [
376,
78
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": null
},
{
"name": "keys",
"type": "STRING",
"links": [
21,
22
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAMetaKeys"
},
"widgets_values": [
"qwen\\gmqwen_k8_piid_4step.safetensors"
]
},
{
"id": 36,
"type": "LoRARenameKeys",
"pos": [
-2856,
864
],
"size": [
320,
190
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "old_keys",
"type": "STRING",
"widget": {
"name": "old_keys"
},
"link": 27
},
{
"name": "new_keys",
"type": "STRING",
"widget": {
"name": "new_keys"
},
"link": 28
}
],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
31
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRARenameKeys"
},
"widgets_values": [
"qwen\\gmqwen_k8_piid_4step.safetensors",
"",
"",
"gmqwen_k8_piid_4step"
]
},
{
"id": 37,
"type": "PreviewAny",
"pos": [
-3792,
920
],
"size": [
592,
384
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 21
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 40,
"type": "RegexReplace",
"pos": [
-3400,
568
],
"size": [
210,
292
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "string",
"type": "STRING",
"widget": {
"name": "string"
},
"link": 22
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": [
23,
25
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "RegexReplace"
},
"widgets_values": [
"",
",\\s\\[.+$",
"",
true,
true,
false,
0
]
},
{
"id": 41,
"type": "PreviewAny",
"pos": [
-3784,
528
],
"size": [
360,
328
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 23
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 42,
"type": "RegexReplace",
"pos": [
-3176,
1016
],
"size": [
288,
292
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "string",
"type": "STRING",
"widget": {
"name": "string"
},
"link": 26
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": [
28,
30
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "RegexReplace"
},
"widgets_values": [
"",
"^(\\w+)",
"diffusion_model.\\1",
true,
true,
false,
0
]
},
{
"id": 43,
"type": "RegexReplace",
"pos": [
-3152,
648
],
"size": [
288,
292
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "string",
"type": "STRING",
"widget": {
"name": "string"
},
"link": 25
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": [
26,
27,
29
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "RegexReplace"
},
"widgets_values": [
"",
"^(proj_out.+|norm_out.+)$\\n",
"",
true,
true,
false,
0
]
},
{
"id": 44,
"type": "PreviewAny",
"pos": [
-2832,
632
],
"size": [
640,
192
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 29
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 45,
"type": "PreviewAny",
"pos": [
-2848,
1096
],
"size": [
640,
192
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 30
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 46,
"type": "PreviewAny",
"pos": [
-2520,
872
],
"size": [
336,
168
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 31
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
}
],
"links": [
[
21,
35,
1,
37,
0,
"STRING"
],
[
22,
35,
1,
40,
0,
"STRING"
],
[
23,
40,
0,
41,
0,
"STRING"
],
[
25,
40,
0,
43,
0,
"STRING"
],
[
26,
43,
0,
42,
0,
"STRING"
],
[
27,
43,
0,
36,
0,
"STRING"
],
[
28,
42,
0,
36,
1,
"STRING"
],
[
29,
43,
0,
44,
0,
"STRING"
],
[
30,
42,
0,
45,
0,
"STRING"
],
[
31,
36,
0,
46,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1,
"offset": [
4191.258750150775,
-356.1917593715161
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
-10
View File
@@ -1,10 +0,0 @@
from .metakeys import *
from .modelunloader import *
__all__ = (
"UNetMetaKeys",
"CLIPMetaKeys",
"LoRAMetaKeys",
"CheckpointMetaKeys",
"OffloadModels",
)
File diff suppressed because it is too large Load Diff
+268
View File
@@ -0,0 +1,268 @@
"""
Device capabilities and memory management utilities.
Integrates with ComfyUI's model_management for safe memory operations.
"""
import os
import gc
import psutil
import torch
from typing import Optional, Tuple
from dataclasses import dataclass
# Import ComfyUI's model management functions
try:
from comfy.model_management import (
get_torch_device,
get_free_memory,
get_total_memory,
free_memory,
soft_empty_cache,
unload_all_models,
is_device_cuda,
)
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
@dataclass
class DeviceCapabilities:
"""Device capabilities for I/O and compute optimization."""
device_type: str # 'cuda', 'cpu', 'mps', 'xpu'
total_vram_gb: float
free_vram_gb: float
total_ram_gb: float
free_ram_gb: float
cpu_cores: int
recommended_workers: int
max_tensor_size_gb: float # Largest tensor we can safely load
supports_pinned_memory: bool
supports_async_streams: bool
def get_device_capabilities(device: Optional[torch.device] = None) -> DeviceCapabilities:
"""
Analyze device capabilities for optimal I/O configuration.
Args:
device: Target device (default: ComfyUI's torch device)
Returns:
DeviceCapabilities dataclass with recommended settings
"""
if device is None:
if COMFY_AVAILABLE:
device = get_torch_device()
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device_type = device.type if hasattr(device, 'type') else 'cpu'
# Get RAM info
ram_info = psutil.virtual_memory()
total_ram_gb = ram_info.total / (1024**3)
free_ram_gb = ram_info.available / (1024**3)
# Get VRAM info
if device_type == 'cuda' and torch.cuda.is_available():
if COMFY_AVAILABLE:
total_vram = get_total_memory(device)
free_vram = get_free_memory(device)
else:
free_vram, total_vram = torch.cuda.mem_get_info(device)
total_vram_gb = total_vram / (1024**3)
free_vram_gb = free_vram / (1024**3)
supports_pinned = True
supports_async = True
elif device_type == 'mps':
total_vram_gb = total_ram_gb * 0.5 # MPS shares RAM
free_vram_gb = free_ram_gb * 0.5
supports_pinned = False
supports_async = False
else:
total_vram_gb = 0
free_vram_gb = 0
supports_pinned = False
supports_async = False
# CPU cores
cpu_cores = os.cpu_count() or 4
# Recommended workers based on hardware
if device_type == 'cuda':
# GPU: fewer workers since I/O is not the bottleneck
recommended_workers = min(4, cpu_cores // 2)
else:
# CPU: more workers for parallel I/O
recommended_workers = min(8, cpu_cores)
# Max tensor size: leave headroom for operations
if device_type == 'cuda':
max_tensor_gb = min(free_vram_gb * 0.8, free_ram_gb * 0.5)
else:
max_tensor_gb = free_ram_gb * 0.6
return DeviceCapabilities(
device_type=device_type,
total_vram_gb=total_vram_gb,
free_vram_gb=free_vram_gb,
total_ram_gb=total_ram_gb,
free_ram_gb=free_ram_gb,
cpu_cores=cpu_cores,
recommended_workers=recommended_workers,
max_tensor_size_gb=max_tensor_gb,
supports_pinned_memory=supports_pinned,
supports_async_streams=supports_async,
)
def estimate_model_size(file_path: str) -> float:
"""Estimate model size in GB from file size."""
try:
return os.path.getsize(file_path) / (1024**3)
except OSError:
return 0.0
def can_fit_in_memory(
model_paths: list[str],
device: Optional[torch.device] = None,
safety_factor: float = 0.8
) -> Tuple[bool, str]:
"""
Check if models can fit in available memory.
Args:
model_paths: List of model file paths
device: Target device
safety_factor: Fraction of free memory to use (default 0.8)
Returns:
(can_fit, reason_message)
"""
caps = get_device_capabilities(device)
total_size = sum(estimate_model_size(p) for p in model_paths)
if caps.device_type == 'cuda':
available = caps.free_vram_gb * safety_factor
mem_type = "VRAM"
else:
available = caps.free_ram_gb * safety_factor
mem_type = "RAM"
if total_size <= available:
return True, f"OK: {total_size:.2f}GB needed, {available:.2f}GB {mem_type} available"
else:
return False, f"Insufficient memory: {total_size:.2f}GB needed but only {available:.2f}GB {mem_type} available"
def prepare_for_large_operation(
estimated_memory_gb: float,
device: Optional[torch.device] = None
) -> bool:
"""
Prepare device for a large memory operation by freeing resources.
Args:
estimated_memory_gb: Estimated memory needed in GB
device: Target device
Returns:
True if preparation succeeded
"""
if not COMFY_AVAILABLE:
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return True
if device is None:
device = get_torch_device()
memory_bytes = int(estimated_memory_gb * 1024**3)
# Use ComfyUI's free_memory which unloads models intelligently
try:
free_memory(memory_bytes, device)
return True
except Exception as e:
print(f"[DeviceUtils] Warning: Could not free memory: {e}")
return False
def cleanup_after_operation():
"""Clean up memory after a large operation."""
gc.collect()
if COMFY_AVAILABLE:
soft_empty_cache()
elif torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def get_optimal_workers(
model_size_gb: float,
caps: Optional[DeviceCapabilities] = None
) -> int:
"""
Calculate optimal number of I/O workers based on model size and device.
Args:
model_size_gb: Model size in GB
caps: Pre-calculated device capabilities (optional)
Returns:
Recommended number of workers
"""
if caps is None:
caps = get_device_capabilities()
# For very large models, fewer workers to avoid memory pressure
if model_size_gb > caps.free_ram_gb * 0.5:
return max(1, caps.recommended_workers // 2)
# For small models, use more workers
if model_size_gb < 2.0:
return min(8, caps.cpu_cores)
return caps.recommended_workers
def get_processing_device_string(prefer_cuda: bool = True) -> str:
"""
Get the best available processing device as a string.
Args:
prefer_cuda: Prefer CUDA over CPU if available
Returns:
Device string ('cuda', 'mps', 'cpu')
"""
if prefer_cuda:
if torch.cuda.is_available():
return "cuda"
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
return "mps"
return "cpu"
def print_device_info():
"""Print device capabilities for debugging."""
caps = get_device_capabilities()
print(f"\n{'='*50}")
print(" Device Capabilities")
print(f"{'='*50}")
print(f" Device Type: {caps.device_type}")
print(f" Total VRAM: {caps.total_vram_gb:.2f} GB")
print(f" Free VRAM: {caps.free_vram_gb:.2f} GB")
print(f" Total RAM: {caps.total_ram_gb:.2f} GB")
print(f" Free RAM: {caps.free_ram_gb:.2f} GB")
print(f" CPU Cores: {caps.cpu_cores}")
print(f" Recommended Workers: {caps.recommended_workers}")
print(f" Max Tensor Size: {caps.max_tensor_size_gb:.2f} GB")
print(f" Pinned Memory: {caps.supports_pinned_memory}")
print(f" Async Streams: {caps.supports_async_streams}")
print(f"{'='*50}\n")
+989
View File
@@ -0,0 +1,989 @@
"""
DoRA Extraction via SVD and Weight Decomposition.
Native implementation of Weight-Decomposed Low-Rank Adapter extraction
with multiple rank selection modes. Extracted models use the DoRA structure.
"""
import fnmatch
import os
import re
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# =============================================================================
# Rank Index Functions - Determine LoRA rank from singular values
# =============================================================================
def _index_sv_fixed(S: torch.Tensor, dim: int) -> int:
"""Fixed rank mode - use specified dimension."""
return max(1, min(dim, len(S) - 1))
def _index_sv_threshold(S: torch.Tensor, threshold: float) -> int:
"""Threshold mode - keep singular values above absolute threshold."""
rank = int(torch.sum(S > threshold).item())
return max(1, min(rank, len(S) - 1))
def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int:
"""Ratio mode - keep singular values > max(S) / ratio (kohya convention)."""
if ratio <= 0:
return len(S) - 1
min_sv = S[0] / ratio
rank = int(torch.sum(S > min_sv).item())
return max(1, min(rank, len(S) - 1))
def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Cumulative mode - keep enough SVs to reach target % of total.
Calculates relative to max_rank if provided, otherwise relative to full.
"""
if max_rank is not None and max_rank < len(S):
total = torch.sum(S[:max_rank])
else:
total = torch.sum(S)
if total < 1e-8:
return 1
cumsum = torch.cumsum(S, dim=0) / total
rank = int(torch.searchsorted(cumsum, target).item()) + 1
return max(1, min(rank, len(S) - 1))
def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Frobenius norm mode - preserve target fraction of Frobenius norm.
Calculates relative to max_rank if provided, otherwise relative to full.
This means "retain target% of what's achievable within max_rank".
"""
if max_rank is not None and max_rank < len(S):
# Calculate relative to what's achievable within max_rank
S_capped = S[:max_rank]
S_sq = S_capped.pow(2)
total_sq = torch.sum(S_sq)
else:
S_sq = S.pow(2)
total_sq = torch.sum(S_sq)
if total_sq < 1e-8:
return 1
# Cumsum of all S (not capped) to find where we reach target
cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq
rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1
return max(1, min(rank, len(S) - 1))
def _index_sv_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
"""Knee detection - find elbow point on singular value curve."""
n = len(S)
if n < 3:
return 1
s_max, s_min = S[0].item(), S[-1].item()
if s_max - s_min < min_sv:
return 1
# Normalize to [0, 1]
s_norm = (S - s_min) / (s_max - s_min)
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
# Distance from diagonal
distances = (x_norm + s_norm - 1).abs()
rank = torch.argmax(distances).item() + 1
return max(1, min(rank, n - 1))
def _index_sv_cumulative_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
"""Knee detection on cumulative singular value curve."""
n = len(S)
if n < 3:
return 1
total = torch.sum(S)
if total < min_sv:
return 1
cumsum = torch.cumsum(S, dim=0) / total
y_min, y_max = cumsum[0].item(), cumsum[-1].item()
if y_max - y_min < min_sv:
return 1
y_norm = (cumsum - y_min) / (y_max - y_min)
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
distances = (y_norm - x_norm).abs()
rank = torch.argmax(distances).item() + 1
return max(1, min(rank, n - 1))
def _index_sv_rel_decrease(S: torch.Tensor, tau: float = 0.1) -> int:
"""Relative decrease mode - stop when consecutive SV ratio drops below tau."""
if len(S) < 2:
return 1
ratios = S[1:] / (S[:-1] + 1e-8)
for k in range(len(ratios)):
if ratios[k] < tau:
return k + 1
return len(S)
def _compute_rank(S: torch.Tensor, mode: str, mode_param: float,
max_rank: int = None) -> int:
"""Compute rank based on mode and parameters."""
if mode == "fixed":
rank = _index_sv_fixed(S, int(mode_param))
elif mode == "threshold":
rank = _index_sv_threshold(S, mode_param)
elif mode == "ratio":
rank = _index_sv_ratio(S, mode_param)
elif mode == "quantile" or mode == "sv_cumulative":
rank = _index_sv_cumulative(S, mode_param, max_rank)
elif mode == "sv_fro":
rank = _index_sv_fro(S, mode_param, max_rank)
elif mode == "sv_knee":
rank = _index_sv_knee(S)
elif mode == "sv_cumulative_knee":
rank = _index_sv_cumulative_knee(S)
elif mode == "sv_rel_decrease":
rank = _index_sv_rel_decrease(S, mode_param)
else:
rank = _index_sv_fixed(S, int(mode_param) if mode_param else 64)
if max_rank is not None:
rank = min(rank, max_rank)
return rank
# =============================================================================
# SVD Extraction Functions
# =============================================================================
def _svd_extract_linear_lowrank(
weight: torch.Tensor,
rank: int,
device: str,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms.
Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# Clamp rank to valid range
max_possible_rank = min(out_dim, in_dim)
rank = max(1, min(rank, max_possible_rank))
# Check if decomposition is worthwhile
if rank >= max_possible_rank:
return weight, "full"
# Use svd_lowrank for faster low-rank approximation
try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
# U: [out_dim, rank], S: [rank], V: [in_dim, rank]
Vh = V.T # [rank, in_dim]
except Exception as e:
raise RuntimeError(f"svd_lowrank failed: {e}")
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_linear(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
SVD decomposition for linear layers.
For fixed rank mode, uses svd_lowrank which is much faster.
For adaptive modes, uses full SVD to analyze all singular values.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# For fixed rank, use the faster svd_lowrank
if mode == "fixed" and mode_param > 0:
target_rank = int(mode_param)
if max_rank is not None:
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
# Compute rank based on mode
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim):
return weight, "full"
# Truncate to rank
U = U[:, :rank]
S = S[:rank]
Vh = Vh[:rank, :]
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass # Skip clamping if quantile fails
# Construct LoRA matrices
# lora_up = U @ diag(S), lora_down = Vh
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff for optional sparse bias
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_conv(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
) -> tuple:
"""
SVD decomposition for conv2d layers.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
out_ch, in_ch, kh, kw = weight.shape
is_1x1 = (kh == 1 and kw == 1)
# Flatten for SVD
if is_1x1:
mat = weight.view(out_ch, in_ch) # [out_ch, in_ch]
else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile)
if mode_str == "full":
return weight, "full"
lora_down, lora_up, diff = result
rank = lora_down.shape[0]
# Reshape for conv
if is_1x1:
lora_down = lora_down.unsqueeze(-1).unsqueeze(-1) # [rank, in_ch, 1, 1]
lora_up = lora_up.unsqueeze(-1).unsqueeze(-1) # [out_ch, rank, 1, 1]
else:
lora_down = lora_down.reshape(rank, in_ch, kh, kw)
lora_up = lora_up.reshape(out_ch, rank, 1, 1)
return (lora_down, lora_up, diff), "low rank"
# =============================================================================
# Chunked Extraction for Large Layers
# =============================================================================
def _detect_fused_layer(key: str, shape: tuple) -> int:
"""Detect fused QKV/MLP layers and return number of chunks."""
if len(shape) != 2:
return 1
out_dim, in_dim = shape
key_lower = key.lower()
# QKV layers (3 chunks)
if "qkv" in key_lower and out_dim % 3 == 0 and out_dim // 3 >= in_dim // 4:
return 3
# KV layers (2 chunks)
if "kv" in key_lower and "qkv" not in key_lower and out_dim % 2 == 0:
return 2
# Fused MLP (2 chunks)
if ("linear1" in key_lower or "gate_up" in key_lower) and out_dim % 2 == 0 and out_dim >= in_dim * 2:
return 2
return 1
def _extract_chunked_layer(
weight_diff: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
) -> tuple:
"""Extract LoRA from fused layer by chunking."""
out_dim, in_dim = weight_diff.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :]
try:
result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
# Combine chunks
# Note: Using the first chunk's down matrix as the basis for all chunks
# is a better approximation than a simple mean if the chunks share the same input space (like QKV).
# Ideally, we would use a more sophisticated joint SVD, but this is a reasonable fallback.
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
# =============================================================================
# Pattern Matching
# =============================================================================
def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
"""Compile patterns from whitespace-separated string.
In regex mode (default) returns compiled re.Pattern objects.
In glob mode returns plain strings; fnmatch handles matching.
"""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for p in pattern_string.split():
p = p.strip()
if not p:
continue
if glob_mode:
patterns.append(p)
else:
try:
patterns.append(re.compile(p))
except re.error:
pass
return patterns
def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
"""Check if key matches any pattern.
Glob mode: each pattern is matched as a substring glob using fnmatch
(dots are literal, * matches any sequence of characters).
Regex mode: each pattern is a compiled re.Pattern, matched via search().
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
return any(p.search(key) for p in patterns)
# =============================================================================
# Main Extraction Function
# =============================================================================
def extract_dora_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
Args:
model_a_path: Finetuned model path
model_b_path: Base model path
mode: Rank selection mode
linear_param: Mode parameter for linear layers
conv_param: Mode parameter for conv layers
device: Computation device
save_dtype: Output dtype
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers
clamp_quantile: Quantile for weight clamping
min_diff: Minimum difference threshold
skip_patterns_str: Patterns for layers to skip
mismatch_mode: How to handle mismatches
chunk_large_layers: Enable chunked extraction
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) # 1.5x for SVD headroom
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key)
if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
# Without base, DoRA fallback just treats tensor_a as diff
dora_scale = None
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
weight_diff = tensor_a
dora_scale = None
del tensor_b
else:
# 1-D weights (e.g. x_pad_token, norm .weight) have no
# weight-decomposition direction and are skipped below
# anyway; bail out before the norm, which would index
# dim=1 and raise on a 1-D tensor.
if tensor_a.ndim < 2:
del tensor_a, tensor_b
return "skipped", None
# Weight Decomposition calculation
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
# Compute directional diff: W_dir = W_ft * (norm_base / norm_ft)
# Use a small epsilon to avoid division by zero
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del tensor_a, tensor_b, w_dir
# Skip small differences
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[DoRA Extract] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_layer(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
print(f"[DoRA Extract] Failed: {key}: {e}")
del weight_diff
return "skipped", None
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, rank_or_diff = result
# rank is the size of lora_down's first dimension
rank = lora_down.shape[0]
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting DoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[DoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> str:
"""Build output path for LoRA file."""
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
# =============================================================================
# Node Definitions
# =============================================================================
def _get_model_inputs():
return [
io.Combo.Input("model_a", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Finetuned model (A - B = LoRA)"),
io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Base model (A - B = LoRA)"),
]
def _format_lora_key(key: str) -> str:
"""
Format the key for LoRA saving.
Always standardizes to 'diffusion_model.' prefix for maximum compatibility.
ComfyUI maps 'diffusion_model.<layer>' generically for all model types, making
it universally compatible. The 'transformer.' prefix only works for a subset of
model types via model-specific mapping code and is intentionally not used here.
"""
if key.endswith(".weight"):
key = key[:-7]
# Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*)
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # strip len("model.diffusion_model.")
return f"diffusion_model.{inner_key}"
# Handle net.* wrapper
if key.startswith("net."):
inner_key = key[4:] # strip len("net.")
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys (transformer_blocks.* / single_transformer_blocks.* without prefix)
if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"):
return f"diffusion_model.{key}"
# Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"):
return key
# Handle already-prefixed keys
if key.startswith("diffusion_model."):
return key
# Absorb now-invalid transformer. prefix into diffusion_model. for compatibility
if key.startswith("transformer."):
return f"diffusion_model.{key[12:]}"
# Handle known Diffusers UNet blocks
if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]):
return f"diffusion_model.{key}"
# Default fallback
return f"diffusion_model.{key}"
def _get_common_inputs():
return [
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
io.Boolean.Input("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
tooltip="Clamp outlier singular values"),
io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.String.Input("output_filename", default="extracted_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
]
class DoRAExtractFixed(io.ComfyNode):
"""Extract DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractFixed",
display_name="DoRA Extract (Fixed Rank)",
category="ModelUtils/DoRA",
description="Extract DoRA with specified fixed rank for each layer type.",
inputs=[
*_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384,
tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384,
tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractRatio(io.ComfyNode):
"""Extract DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractRatio",
display_name="DoRA Extract (Ratio)",
category="ModelUtils/DoRA",
description="Keep singular values > max(S) / ratio.",
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractQuantile(io.ComfyNode):
"""Extract DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractQuantile",
display_name="DoRA Extract (Quantile)",
category="ModelUtils/DoRA",
description="Keep enough singular values to reach target percentage.",
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target cumulative % for linear layers"),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target cumulative % for conv layers"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractKnee(io.ComfyNode):
"""Extract DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractKnee",
display_name="DoRA Extract (Knee Detection)",
category="ModelUtils/DoRA",
description="Automatically find optimal rank using knee detection on singular value curve.",
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
default="sv_knee", tooltip="Knee detection method"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractFrobenius(io.ComfyNode):
"""Extract DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractFrobenius",
display_name="DoRA Extract (Frobenius)",
category="ModelUtils/DoRA",
description="Preserve target fraction of Frobenius norm.",
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target Frobenius norm fraction for linear"),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target Frobenius norm fraction for conv"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
+790
View File
@@ -0,0 +1,790 @@
"""
Learned DoRA Extraction via Optimization.
Extracts Weight-Decomposed Low-Rank Adapters using gradient descent optimization
to minimize the reconstruction MSE loss, matching the true non-linear DoRA formulation.
Uses analytical SVD purely as the initialization seed.
"""
import fnmatch
import os
import re
import math
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# Import the existing SV ranking and chunking logic so we don't have to duplicate it
from .dora_extract_wd import (
_compute_rank, _svd_extract_linear_lowrank, _svd_extract_linear,
_svd_extract_conv, _detect_fused_layer, _compile_patterns,
_matches_any_pattern, _build_lora_output_path, _get_model_inputs,
_format_lora_key, _get_common_inputs
)
def _adaptive_lr_update_cosine(curr_lr: float, improved: bool, worse_loss_counter: int, iteration: int, tensor_shape: tuple, min_lr: float, early_stop_stall: int, lr_factor: float = 0.95, lr_cooldown: int = 1) -> tuple:
M, N = tensor_shape
shape_ratio = abs(M - N) / max(M, N) if max(M, N) > 0 else 0
t = min(worse_loss_counter / max(early_stop_stall, 1), 1.0)
u_factor = (1 + math.cos(2 * math.pi * t)) / 2
if improved:
base_boost = 1.25
distance = base_boost - 1.0
scaled_distance = distance * (1.0 + 0.5 * shape_ratio)
boost_mult = 1.0 + scaled_distance * (1.0 - u_factor)
return min(curr_lr * boost_mult, 100.0), True
else:
cooldown = max(lr_cooldown, 1)
if iteration % cooldown != 0:
return curr_lr, False
min_decay = lr_factor - 0.03 * shape_ratio
max_decay = 0.995
decay_mult = min_decay + (max_decay - min_decay) * u_factor
return max(curr_lr * decay_mult, min_lr), True
def _compute_shape_aware_plateau_params(M: int, N: int, lr_patience: int, lr_factor: float, lr_cooldown: int, lr_shape_influence: float = 1.0) -> tuple:
if lr_shape_influence > 0:
aspect_ratio = max(M, N) / min(M, N) if min(M, N) > 0 else 1.0
ar_factor = math.sqrt(aspect_ratio)
blend = lr_shape_influence
effective_patience = lr_patience
raw_factor = lr_factor
aggressive_factor = raw_factor**ar_factor
effective_factor = raw_factor + (aggressive_factor - raw_factor) * blend
effective_cooldown = lr_cooldown
else:
effective_patience = lr_patience
effective_factor = lr_factor
effective_cooldown = lr_cooldown
return effective_patience, effective_factor, effective_cooldown
def _optimize_dora(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
init_down: torch.Tensor,
init_up: torch.Tensor,
optimize_iters: int,
learning_rate: float,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""
Run gradient descent to optimize lora_down and lora_up against the actual DoRA formula.
Args:
tensor_base: Base model weight
tensor_ft: Finetuned model weight
init_down: SVD-initialized lora_down matrix
init_up: SVD-initialized lora_up matrix
optimize_iters: Number of iterations for AdamW
learning_rate: Learning rate for AdamW
optimizer_type: Optimizer choice (adamw, radam, prodigy)
lr_schedule: LR Schedule choice (exponential, plateau, adaptive)
early_stop_loss: Target loss to stop optimizing
early_stop_stall: Number of non-improving iterations before stopping
Returns:
(optimized_down, optimized_up)
"""
if optimize_iters <= 0:
return init_down, init_up
is_conv = tensor_base.ndim == 4
if is_conv:
out_ch, in_ch, kh, kw = tensor_base.shape
flat_base = tensor_base.reshape(out_ch, -1)
flat_ft = tensor_ft.reshape(out_ch, -1)
# Flatten init matrices for matmul: (out, rank) @ (rank, in*k*k)
rank = init_down.shape[0]
flat_up = init_up.reshape(out_ch, rank)
flat_down = init_down.reshape(rank, -1)
else:
flat_base = tensor_base
flat_ft = tensor_ft
flat_up = init_up
flat_down = init_down
# Magnitude of the finetuned model (fixed DoRA scale)
m = torch.linalg.norm(flat_ft, dim=1, keepdim=True).detach()
eps = 1e-8
# Set up learnable parameters
with torch.inference_mode(False), torch.enable_grad():
# Clone inference tensors to normal tensors for autograd
flat_base_opt = flat_base.clone()
flat_ft_opt = flat_ft.clone()
m_opt = m.clone()
# Free original tensors to save memory
del flat_base, flat_ft, m
B = torch.nn.Parameter(flat_up.clone())
A = torch.nn.Parameter(flat_down.clone())
opt_type = optimizer_type.lower()
if opt_type == "radam":
optimizer = torch.optim.RAdam([B, A], lr=learning_rate)
elif opt_type == "prodigy":
try:
from prodigyplus.prodigy_plus_schedulefree import ProdigyPlusScheduleFree
optimizer = ProdigyPlusScheduleFree([B, A], lr=learning_rate, use_schedulefree=False)
except ImportError:
print("ProdigyPlusScheduleFree optimizer not found. Falling back to AdamW. To install: pip install prodigyplus")
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
else:
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
# Learning rate scheduling and early stopping state
best_loss = float('inf')
worse_loss_counter = 0
plateau_counter = 0
cooldown_counter = 0
current_lr = learning_rate
schedule_type = lr_schedule.lower()
M_dim, N_dim = flat_base_opt.shape[0], flat_base_opt.shape[1]
eff_patience, eff_factor, eff_cooldown = _compute_shape_aware_plateau_params(M_dim, N_dim, lr_patience, lr_factor, lr_cooldown)
# Optimization loop
opt_pbar = tqdm(range(optimize_iters), desc=f" └─ Optimizing {layer_name}", leave=False, dynamic_ncols=True, mininterval=0.5)
for i in opt_pbar:
optimizer.zero_grad()
# 1. Reconstruct directional weight: W_dir = W_base + B @ A
w_approx = flat_base_opt + (B @ A)
# 2. Normalize directional weight
approx_norm = torch.linalg.norm(w_approx, dim=1, keepdim=True) + eps
# 3. Scale by target magnitude
w_reconstructed = m_opt * (w_approx / approx_norm)
# 4. Compute MSE loss against actual finetuned weight
loss = torch.nn.functional.mse_loss(w_reconstructed, flat_ft_opt)
loss_val = loss.item()
if i == 0:
initial_loss = loss_val
best_loss = loss_val
loss.backward()
optimizer.step()
# --- Early Stopping & Schedulers ---
improved = False
if loss_val < best_loss * 0.9999: # 0.01% improvement threshold
best_loss = loss_val
worse_loss_counter = 0
plateau_counter = 0
improved = True
else:
worse_loss_counter += 1
plateau_counter += 1
if loss_val <= early_stop_loss:
break
if worse_loss_counter > early_stop_stall:
break
if current_lr <= early_stop_lr:
break
# LR Schedulers
if schedule_type == "adaptive":
current_lr, updated = _adaptive_lr_update_cosine(
current_lr, improved, worse_loss_counter, i, (M_dim, N_dim), early_stop_lr, early_stop_stall, lr_factor, lr_cooldown
)
if updated:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
elif schedule_type == "plateau":
if cooldown_counter > 0:
cooldown_counter -= 1
elif plateau_counter >= max(eff_patience, 1): # Decay if patience threshold reached
current_lr = max(current_lr * eff_factor, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
plateau_counter = 0 # reset counter after decay
cooldown_counter = eff_cooldown
elif schedule_type == "exponential":
current_lr = max(current_lr * 0.99, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
if i % 10 == 0 or i == optimize_iters - 1:
if schedule_type == "plateau":
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "plat": f"{plateau_counter}/{eff_patience}"}, refresh=False)
else:
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "stall": f"{worse_loss_counter}"}, refresh=False)
final_loss = loss.item()
opt_pbar.close()
print(f" └─ {layer_name} | Loss: {initial_loss:.6f} -> {final_loss:.6f} (stopped at iter {i+1})")
final_up = B.detach()
final_down = A.detach()
# Reshape back to original dimensions
if is_conv:
final_up = final_up.reshape(out_ch, rank, 1, 1)
final_down = final_down.reshape(rank, in_ch, kh, kw)
return final_down, final_up
def _extract_chunked_learned_layer(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""Extract Learned DoRA from fused layer by chunking."""
out_dim, in_dim = tensor_base.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk_base = tensor_base[i * chunk_size:(i + 1) * chunk_size, :]
chunk_ft = tensor_ft[i * chunk_size:(i + 1) * chunk_size, :]
norm_base = torch.linalg.norm(chunk_base, dim=1, keepdim=True)
norm_ft = torch.linalg.norm(chunk_ft, dim=1, keepdim=True)
eps = 1e-8
w_dir = chunk_ft * (norm_base / (norm_ft + eps))
weight_diff = w_dir - chunk_base
try:
result, mode_str = _svd_extract_linear(weight_diff, mode, mode_param, device, max_rank)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
# Apply learned optimization
lora_down, lora_up = _optimize_dora(
chunk_base, chunk_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=f"{layer_name} (Chunk {i+1}/{num_chunks})"
)
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
def extract_dora_learned_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract Learned DoRA from difference between two models.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[Learned DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
# Cannot extract true DoRA without base. Fallback to diff.
cpu_a = handler_a.get_tensor(key)
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) if use_pinned else cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
weight_diff = tensor_a
tensor_base = torch.zeros_like(tensor_a)
tensor_ft = tensor_a
dora_scale = None
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
weight_diff = tensor_a
tensor_base = torch.zeros_like(tensor_a)
tensor_ft = tensor_a
dora_scale = None
del tensor_b
else:
tensor_base = tensor_b
tensor_ft = tensor_a
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del w_dir
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if weight_diff.ndim < 2:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
# 1. Run Analytical SVD (Initialization)
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
# 2. Run Optimization (Learned Rounding)
if mode_str != "full" and dora_scale is not None:
lora_down, lora_up, _ = result
lora_down, lora_up = _optimize_dora(
tensor_base, tensor_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
)
result = (lora_down, lora_up, _)
except Exception as e:
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[Learned DoRA] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_learned_layer(
tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff, tensor_base, tensor_ft
return "chunked", layer_results
print(f"[Learned DoRA] Failed: {key}: {e}")
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
rank = lora_down.shape[0]
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff, tensor_base, tensor_ft
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting Learned DoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[Learned DoRA] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
# =============================================================================
# Node Definitions
# =============================================================================
def _get_learned_inputs():
return [
io.Int.Input("optimize_iters", default=500, min=0, max=10000,
tooltip="Number of gradient descent iterations (0 = skip optimization)"),
io.Float.Input("learning_rate", default=1.0, min=0.0001, max=100.0, step=0.001,
tooltip="Base learning rate. AdamW prefers ~0.01, Prodigy prefers ~1.0"),
io.Combo.Input("optimizer", options=["prodigy", "adamw", "radam"], default="prodigy",
tooltip="Optimization algorithm"),
io.Combo.Input("lr_schedule", options=["plateau", "adaptive", "exponential", "constant"], default="plateau",
tooltip="Learning rate scheduling strategy"),
io.Int.Input("lr_patience", default=2, min=0, max=1000, tooltip="Steps to wait before decaying LR (Plateau)"),
io.Float.Input("lr_factor", default=0.9, min=0.01, max=1.0, step=0.01, tooltip="Factor to decay LR by (Plateau)"),
io.Int.Input("lr_cooldown", default=2, min=0, max=1000, tooltip="Steps to wait after decay before checking again (Plateau)"),
io.Float.Input("early_stop_loss", default=1e-6, min=0.0, max=1.0, step=1e-9,
tooltip="Stop early if MSE loss drops below this value"),
io.Int.Input("early_stop_stall", default=2000, min=0, max=10000,
tooltip="Stop early if loss doesn't improve for this many iterations"),
io.Float.Input("early_stop_lr", default=9.01e-9, min=0.0, max=1.0, step=1e-10,
tooltip="Stop early if learning rate drops below this value"),
]
class DoRALearnedExtractFixed(io.ComfyNode):
"""Extract Learned DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFixed",
display_name="DoRA Learned Extract (Fixed Rank) (WIP)",
category="ModelUtils/DoRA",
description="Extract DoRA via MSE Loss minimization with fixed rank.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
*_get_learned_inputs(),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD iterations for initialization"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractRatio(io.ComfyNode):
"""Extract Learned DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractRatio",
display_name="DoRA Learned Extract (Ratio) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping SVs > max(S) / ratio.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractQuantile(io.ComfyNode):
"""Extract Learned DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractQuantile",
display_name="DoRA Learned Extract (Quantile) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping enough SVs to reach target percentage.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractKnee(io.ComfyNode):
"""Extract Learned DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractKnee",
display_name="DoRA Learned Extract (Knee Detection) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction automatically finding optimal rank using knee detection.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractFrobenius(io.ComfyNode):
"""Extract Learned DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFrobenius",
display_name="DoRA Learned Extract (Frobenius) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction preserving target fraction of Frobenius norm.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
+341
View File
@@ -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="/"))
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="/"))
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="/"))
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="/"))
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="/"))
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="/"))
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=""))
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 []
+695
View File
@@ -0,0 +1,695 @@
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.")
# Add the reference module to the Python path
# Assuming downloader_utils.py is inside nodes/, the root is one level up
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(root_dir, 'unifiedefficientloader-reference'))
from unifiedefficientloader import UnifiedSafetensorsLoader
# --- 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 get_potential_preview_files(base_path):
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 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
import numpy as np
from PIL import Image, ImageOps
import torch
def load_image_tensor(path):
if not os.path.exists(path):
return None
try:
img = Image.open(path)
img = ImageOps.exif_transpose(img)
image = img.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return image
except Exception as e:
print(f"Error loading image {path}: {e}")
return None
def get_model_workflows(base_path):
workflows = []
# Main workflow
main_json = f"{base_path}.json"
if os.path.exists(main_json):
workflows.append(main_json)
# Example workflows
# They are named like base_path.example.N.json
# or base_path.preview.json (if extracted from preview)
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")
# Sort them by index
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):
# Prefer civitai.info if available
civitai_info = f"{base_path}.civitai.info"
if os.path.exists(civitai_info):
return civitai_info
# Fallback to metadata.json (safetensors metadata)
metadata_json = f"{base_path}.metadata.json"
if os.path.exists(metadata_json):
return metadata_json
return None
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
View File
File diff suppressed because it is too large Load Diff
+1249
View File
File diff suppressed because it is too large Load Diff
+119
View File
@@ -0,0 +1,119 @@
import os
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
REVERSE_RENAME_DICT = {
"transformer_blocks.": "blocks.",
"norm1.linear_1": "adaln_modulation_self_attn.1",
"norm1.linear_2": "adaln_modulation_self_attn.2",
"norm2.linear_1": "adaln_modulation_cross_attn.1",
"norm2.linear_2": "adaln_modulation_cross_attn.2",
"norm3.linear_1": "adaln_modulation_mlp.1",
"norm3.linear_2": "adaln_modulation_mlp.2",
"attn1.to_q": "self_attn.q_proj",
"attn1.to_k": "self_attn.k_proj",
"attn1.to_v": "self_attn.v_proj",
"attn1.to_out.0": "self_attn.output_proj",
"attn2.to_q": "cross_attn.q_proj",
"attn2.to_k": "cross_attn.k_proj",
"attn2.to_v": "cross_attn.v_proj",
"attn2.to_out.0": "cross_attn.output_proj",
"ff.net.0.proj": "mlp.layer1",
"ff.net.2": "mlp.layer2",
"norm_out.linear_1": "final_layer.adaln_modulation.1",
"norm_out.linear_2": "final_layer.adaln_modulation.2",
"proj_out": "final_layer.linear",
"time_embed.t_embedder": "t_embedder.1",
"time_embed.norm": "t_embedding_norm",
"patch_embed.proj": "x_embedder.proj.1",
}
def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filename: str) -> str:
"""
Converts a Diffusers Anima LoRA back to its non-diffusers format.
"""
model_type = "loras"
model_path = folder_paths.get_full_path_or_raise(model_type, lora_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
# Get metadata from safe_open (header only, no tensor loading)
with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
metadata = f.metadata() or {}
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Stream tensors, rename on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
for key in tqdm(original_keys, desc="Renaming Anima LoRA keys", unit="keys"):
new_key = key
if key.startswith("text_conditioner."):
new_key = f"diffusion_model.llm_adapter.{key.removeprefix('text_conditioner.')}"
elif key.startswith("transformer."):
new_key = key.removeprefix("transformer.")
for old_part, new_part in REVERSE_RENAME_DICT.items():
new_key = new_key.replace(old_part, new_part)
new_key = f"diffusion_model.{new_key}"
writer.write(new_key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
cleanup_after_operation()
return output_path
class AnimaLoraRename(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="AnimaLoraRename",
display_name="Anima LoRA Rename (Diffusers to Non-Diffusers)",
category="ModelUtils/LoRA",
description="Loads a Diffusers Anima LoRA, reverses its key names to the non-diffusers format, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("output_filename", default="anima_lora_non_diffusers"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name: str, output_filename: str) -> io.NodeOutput:
path = _convert_diffusers_to_non_diffusers_anima_lora(lora_name, output_filename)
return io.NodeOutput(path)
+1718
View File
File diff suppressed because it is too large Load Diff
+1022
View File
File diff suppressed because it is too large Load Diff
+935
View File
@@ -0,0 +1,935 @@
import torch
import scipy.ndimage
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)
SKIP = "skip" # Return None, let caller handle
ZEROS = "zeros" # Return zeros tensor matching shape from primary model
class MissingTensorError(Exception):
"""Raised when a tensor is missing and behavior is ERROR."""
pass
# ################# #
# UTILITIES
# ################# #
def resize_by_interpolation(tensor1, tensor2):
"""Resizes tensor2 to match tensor1's shape using interpolation."""
if tensor1 is None or tensor2 is None:
return tensor1, tensor2
if tensor1.shape == tensor2.shape:
return tensor1, tensor2
dims1 = len(tensor1.shape)
dims2 = len(tensor2.shape)
if dims1 != dims2 or dims1 > 4 or dims1 == 0:
return tensor1, tensor2
orig_dtype = tensor2.dtype
t2_f = tensor2.float()
# F.interpolate expects [B, C, ...]
# We treat the tensor as [1, 1, ...dims]
view_shape = [1, 1] + list(tensor2.shape)
t2_f = t2_f.view(view_shape)
# Select mode based on dimensions
if dims1 == 4:
# [Out, In, H, W] -> [1, 1, Out, In, H*W] (5D)
view_shape = [1, 1, tensor2.shape[0], tensor2.shape[1], tensor2.shape[2] * tensor2.shape[3]]
target_size = [tensor1.shape[0], tensor1.shape[1], tensor1.shape[2] * tensor1.shape[3]]
t2_f = t2_f.view(view_shape)
t2_res = F.interpolate(t2_f, size=target_size, mode='trilinear', align_corners=False)
return tensor1, t2_res.view(tensor1.shape).to(orig_dtype)
if dims1 == 1:
mode = 'linear'
elif dims1 == 2:
mode = 'bilinear'
else:
mode = 'trilinear'
t2_res = F.interpolate(t2_f, size=list(tensor1.shape), mode=mode, align_corners=False)
return tensor1, t2_res.view(tensor1.shape).to(orig_dtype)
def resize_tensors(tensor1, tensor2, mode='pad/crop'):
"""Resizes tensors to match shapes using either padding or interpolation."""
if tensor1 is None or tensor2 is None:
return tensor1, tensor2
if tensor1.shape == tensor2.shape:
return tensor1, tensor2
if mode == 'interpolate':
return resize_by_interpolation(tensor1, tensor2)
dims1 = len(tensor1.shape)
dims2 = len(tensor2.shape)
if dims1 != dims2 or dims1 > 4 or dims1 == 0:
return tensor1, tensor2
padding1 = []
padding2 = []
# F.pad expects: (last_dim_left, last_dim_right, prev_dim_top, prev_dim_bottom, ...)
for i in range(dims1 - 1, -1, -1):
s1 = tensor1.shape[i]
s2 = tensor2.shape[i]
if s1 < s2:
padding1.extend([0, s2 - s1])
padding2.extend([0, 0])
elif s2 < s1:
padding1.extend([0, 0])
padding2.extend([0, s1 - s2])
else:
padding1.extend([0, 0])
padding2.extend([0, 0])
if any(p > 0 for p in padding1):
tensor1 = F.pad(tensor1, tuple(padding1))
if any(p > 0 for p in padding2):
tensor2 = F.pad(tensor2, tuple(padding2))
return tensor1, tensor2
# ################# #
# MERGE OPERATORS
# ################# #
class Operation:
def __init__(self, key, *sources):
self.key = key
self.sources = tuple(sources)
self.merge_func = self.recurse
self.alignment_mode = 'pad/crop'
self.alpha = None
self.beta = None
self.gamma = None
self.delta = None
self.seed = None
def __eq__(self, other):
return (type(self), self.key, self.alpha, self.beta, self.gamma, self.delta, self.seed, self.sources) == \
(type(other), other.key, other.alpha, other.beta, other.gamma, other.delta, other.seed, other.sources)
def __hash__(self):
return hash((type(self), self.key, self.alpha, self.beta, self.gamma, self.delta, self.seed, self.sources))
def oper(self, *args) -> torch.Tensor:
raise NotImplementedError
def recurse(self, operation):
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,
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.on_missing = on_missing
self.fallback_shape = fallback_shape
self.fallback_dtype = fallback_dtype
def merge(self) -> torch.Tensor:
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():
if self.on_missing == MissingTensorBehavior.ERROR:
raise MissingTensorError(
f"Key '{self.key}' not found in model '{self.model_name}'"
)
elif self.on_missing == MissingTensorBehavior.SKIP:
return None
elif self.on_missing == MissingTensorBehavior.ZEROS:
if self.fallback_shape is None:
raise ValueError(
f"Cannot create zeros for '{self.key}' without fallback_shape"
)
dtype = self.fallback_dtype if self.fallback_dtype else self.dtype
return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype)
cpu_tensor = handler.get_tensor(self.key)
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):
super().__init__(key, *sources)
self.alpha = alpha
def oper(self, a) -> torch.Tensor:
if a is None:
return None
return a * self.alpha
class Add(Operation):
def oper(self, a, b) -> torch.Tensor:
if a is None:
return b
if b is None:
return a
a, b = resize_tensors(a, b, mode=self.alignment_mode)
return a + b
class Sub(Operation):
def oper(self, a, b) -> torch.Tensor:
if a is None or b is None:
return None # Cannot compute difference with missing operand
a, b = resize_tensors(a, b, mode=self.alignment_mode)
return a - b
class SVD(Operation):
def __init__(self, key, alpha, beta, gamma, *sources, target_shape=None):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma = int(alpha), int(beta), gamma
self.target_shape = target_shape
def oper(self, a, b) -> torch.Tensor:
if a is None or b is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
# For SVD extraction, we must ensure the diff matches the original Model A shape
# so the resulting LoRA weights are compatible with the target architecture.
if self.target_shape and a.shape != self.target_shape:
if self.alignment_mode == 'interpolate':
# Re-align both to match target_shape via interpolation
dummy = torch.empty(self.target_shape, device=a.device)
_, a = resize_by_interpolation(dummy, a)
_, b = resize_by_interpolation(dummy, b)
else:
slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(a.shape, self.target_shape))
a = a[slices]
b = b[slices]
diff, weights, 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
rank = min(rank, a.size()[1], a.size()[0])
matrix = diff.flatten(start_dim=1) if conv2d_3x3 else diff.squeeze() if conv2d else diff
U, S, Vh = torch.linalg.svd(matrix)
U, S, Vh = U[:, :rank], S[:rank], Vh[:rank, :]
U = U @ torch.diag(S)
dist = torch.cat([U.flatten(), Vh.flatten()])
hi_val = torch.quantile(dist, self.gamma)
U, Vh = U.clamp(-hi_val, hi_val), Vh.clamp(-hi_val, hi_val)
if conv2d:
Vh = Vh.reshape(rank, a.size()[1], *kernel_size) if conv2d_3x3 else Vh.reshape(rank, a.size()[1], 1, 1)
U = U.reshape(a.size()[0], rank, 1, 1)
weights_sd = {f"{self.key}.lora_up.weight": U, f"{self.key}.lora_down.weight": Vh, f"{self.key}.alpha": torch.tensor(float(rank))}
return weights_sd
class Smooth(Operation):
def oper(self, a) -> torch.Tensor:
if a is None:
return None
filtered = scipy.ndimage.median_filter(a.detach().cpu().to(torch.float32).numpy(), size=3)
filtered = scipy.ndimage.gaussian_filter(filtered, sigma=1)
return torch.tensor(filtered, dtype=a.dtype, device=a.device)
class TrainDiff(Operation):
def oper(self, a, b, c) -> torch.Tensor:
if a is None or b is None or c is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
a, c = resize_tensors(a, c, mode=self.alignment_mode)
b, c = resize_tensors(b, c, mode=self.alignment_mode)
if torch.allclose(b.float(), c.float()): return torch.zeros_like(a)
diff_bc = b.float() - c.float()
dist_bc, dist_ab = torch.abs(diff_bc), torch.abs(b.float() - a.float())
scale = torch.where(dist_bc + dist_ab != 0, dist_ab / (dist_bc + dist_ab), 0.0)
return (torch.sign(diff_bc) * torch.abs(scale) * torch.abs(diff_bc)).to(a.dtype) * 1.8
class ExtractOp(Operation):
def __init__(self, key, alpha, beta, gamma, *args):
super().__init__(key, *args)
self.alpha, self.beta, self.gamma = alpha, beta, gamma
def oper(self, base: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
if a is None or b is None:
return None
if base is not None:
base, a = resize_tensors(base, a, mode=self.alignment_mode)
base, b = resize_tensors(base, b, mode=self.alignment_mode)
a, b = resize_tensors(a, b, mode=self.alignment_mode)
dtype = base.dtype if base is not None else a.dtype
base_f = base.float() if base is not None else 0
a_f, b_f = a.float() - base_f, b.float() - base_f
c = torch.cosine_similarity(a_f, b_f, -1).clamp(-1, 1).unsqueeze(-1)
d = ((c + 1) / 2)**self.gamma
return (torch.lerp(a_f, b_f, self.alpha) * torch.lerp(d, 1 - d, self.beta)).to(dtype)
class Similarities(ExtractOp):
def oper(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
if a is None or b is None:
return None
return super().oper(None, a, b)
class PowerUpOp(Operation):
def __init__(self, key, alpha, seed, *sources):
super().__init__(key, *sources)
self.alpha, self.seed = alpha, 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
keep_prob = 1 - self.alpha
if keep_prob <= 0:
return torch.zeros_like(delta)
rng = torch.Generator(device=delta.device).manual_seed(self.seed)
m = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).to(delta.dtype)
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)
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)
alpha, delta = max(self.alpha, 0.001), torch.abs(a - b)
max_delta = torch.max(delta)
if max_delta == 0: return a
diff = ((max_delta - delta) / max_delta if self.beta != 1 else delta / max_delta) ** (1 / alpha - 1)
diff = torch.nan_to_num(diff)
rng = torch.Generator(device=diff.device).manual_seed(self.seed)
mask = torch.lerp(torch.bernoulli(torch.clamp(diff, 0, 1), generator=rng), diff, self.gamma)
return a * (1 - mask) + b * mask
class ManualEnhancedInterpolateDifferenceOp(Operation):
def __init__(self, key, alpha, beta, gamma, delta, seed, *sources):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma, self.delta, self.seed = alpha, beta, gamma, delta, 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 = torch.abs(a - b)
diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta))
mean_diff = torch.mean(diff, 0, keepdim=True)
mask = torch.logical_and(self.beta < mean_diff, mean_diff < self.gamma)
powered_diff = torch.nan_to_num(diff ** (1 / max(self.alpha, 0.001) - 1))
masked_diff = powered_diff * mask.float()
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(torch.clamp(masked_diff, 0, 1), generator=rng)
interpolated_mask = torch.lerp(random_mask, masked_diff, self.delta)
return a * (1 - interpolated_mask) + b * interpolated_mask
class AutoEnhancedInterpolateDifferenceOp(Operation):
def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources)
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 = torch.abs(a - b)
max_delta = torch.max(delta)
diff = torch.nan_to_num((max_delta - delta) / max_delta)
mean_diff = torch.mean(diff)
lower_threshold = mean_diff * (1 - self.beta)
upper_threshold = mean_diff * (1 + self.beta)
mask = torch.logical_and(lower_threshold < diff, diff < upper_threshold)
powered_diff = torch.nan_to_num(diff ** (1 / max(self.alpha, 0.001) - 1))
masked_diff = powered_diff * mask.float()
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(torch.clamp(masked_diff, 0, 1), generator=rng)
interpolated_mask = torch.lerp(random_mask, masked_diff, self.gamma)
return a * (1 - interpolated_mask) + b * interpolated_mask
class WeightSumCutoffOp(Operation):
def __init__(self, key, alpha, beta, gamma, *sources):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma = alpha, beta, gamma
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 = torch.abs(a - b)
diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta))
mean = torch.mean(diff, 0, True)
mask = torch.logical_and(mean < self.beta, self.gamma < mean)
mul = self.alpha * mask
return a * (1 - mul) + b * mul
# ################# #
# CALCULATION MODES
# ################# #
class CalcMode:
name = 'calcmode'; description = 'description'; models_used = []; param_docs = {}
def create_recipe(self, key, **kwargs) -> Operation: raise NotImplementedError
@staticmethod
def _get_secondary_loader_args(kwargs, mismatch_mode):
"""Helper to build loader args for secondary models (B, C) with fallback support."""
base_args = {
"handlers": kwargs['handlers'],
"device": kwargs['device'],
"dtype": kwargs['dtype'],
"on_missing": mismatch_mode,
}
# Add fallback info for zeros mode
if '_tensor_a_shape' in kwargs:
base_args['fallback_shape'] = kwargs['_tensor_a_shape']
base_args['fallback_dtype'] = kwargs.get('_tensor_a_dtype', kwargs['dtype'])
return base_args
class PreloadedTensor(Operation):
"""Returns a pre-loaded tensor instead of loading from file.
Used when Model A's tensor is pre-loaded in the merge loop.
"""
def __init__(self, key, tensor):
super().__init__(key)
self.tensor = tensor
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.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Interpolation weight. 0.0 is 100% Model A, 1.0 is 100% Model B.'}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
# Use pre-loaded tensor_a if available
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)
if kwargs['alpha'] >= 1: return b
if kwargs['alpha'] <= 0: return a
c = Multiply(key, 1 - kwargs['alpha'], a)
d = Multiply(key, kwargs['alpha'], b)
res = Add(key, c, d)
res.alignment_mode = alignment_mode
return res
class AddDifference(CalcMode):
name = 'Add-Difference'
description = 'A + (B - C) * α. Applies the difference between B and C to A.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Multiplier for the (B - C) difference.', 'beta': 'Smoothing (0=Off, 1=On).'}
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_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
diff = Sub(key, b, c)
diff.alignment_mode = alignment_mode
if kwargs['beta'] == 1: diff = Smooth(key, diff)
diffm = Multiply(key, kwargs['alpha'], diff)
res = Add(key, a, diffm)
res.alignment_mode = alignment_mode
return res
class TrainDifference(CalcMode):
name = 'Train-Difference'
description = 'A + (B - C) * α, using tensor distances.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Multiplier for the calculated 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_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
diff = TrainDiff(key, a, b, c)
diff.alignment_mode = alignment_mode
diffm = Multiply(key, kwargs['alpha'], diff)
res = Add(key, a, diffm)
res.alignment_mode = alignment_mode
return res
class InterpDifference(CalcMode):
name = 'Comparative-Interpolation'
description = 'Interpolates A and B based on relative tensor value differences.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Curve strength.', 'beta': 'Style (0=Similarity, 1=Difference).', 'gamma': 'Mix (0=Random, 1=Linear).'}
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)
# Skip certain keys entirely
if key.startswith('cond_stage_model.transformer.text_model.embeddings'):
return a
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
class Extract(CalcMode):
name = 'Extract-Features'
description = 'Adds features from B and C to A, based on their similarity.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Similarity vs dissimilarity.', 'gamma': 'Similarity bias.', 'delta': 'Multiplier for final features.'}
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_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
extracted = ExtractOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'] * 15, a, b, c)
extracted.alignment_mode = alignment_mode
multiplied = Multiply(key, kwargs['delta'], extracted)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class AddDisimilarity(CalcMode):
name = 'Add-Dissimilarities'
description = 'Adds dissimilar features between B and C to model A.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Multiplier for final features.', 'gamma': 'Similarity bias.'}
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_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
extracted = Similarities(key, kwargs['alpha'], 1, kwargs['gamma'] * 15, b, c)
extracted.alignment_mode = alignment_mode
multiplied = Multiply(key, kwargs['beta'], extracted)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class PowerUp(CalcMode):
name = 'Power-Up (DARE)'
description = 'Adds capabilities of B to A using Drop and Rescale (DARE).'
models_used = ['A', 'B']
param_docs = {'alpha': 'Dropout rate (proportion of weights from B to *drop*).', 'beta': '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 = PowerUpOp(key, kwargs['alpha'], kwargs['seed'], a, b)
deltahat.alignment_mode = alignment_mode # Important!
multiplied = Multiply(key, kwargs['beta'], deltahat)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class SVDMode(CalcMode):
name = 'SVD LoRA Extraction'
description = 'Extracts the difference between B and A into a LoRA file.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Rank for standard layers.', 'beta': 'Rank for 3x3 conv layers.', 'gamma': 'Clamp quantile for weights.'}
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)
res = SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b, target_shape=kwargs.get('_tensor_a_shape'))
res.alignment_mode = alignment_mode
return res
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.'}
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)
if key.startswith('cond_stage_model.transformer.text_model.embeddings'):
return a
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = ManualEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
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.'}
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)
if key.startswith('cond_stage_model.transformer.text_model.embeddings'):
return a
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = AutoEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
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.'}
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)
res = WeightSumCutoffOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b)
res.alignment_mode = alignment_mode
return res
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()]
+23
View File
@@ -0,0 +1,23 @@
"""
Utilities for model merging operations.
"""
from typing import Dict, List
def get_union_keys(*handlers) -> List[str]:
"""Returns the union of all tensor keys across multiple model handlers."""
all_keys = set()
for handler in handlers:
if handler is not None:
all_keys.update(handler.keys())
return sorted(all_keys)
def get_intersection_keys(*handlers) -> List[str]:
"""Returns only keys present in ALL model handlers."""
handlers = [h for h in handlers if h is not None]
if not handlers:
return []
all_keys = set(handlers[0].keys())
for handler in handlers[1:]:
all_keys &= set(handler.keys())
return sorted(all_keys)
+153 -97
View File
@@ -1,118 +1,174 @@
import torch
from safetensors import safe_open
from safetensors.torch import save_file
import os
import json
import struct
import folder_paths
from comfy.model_management import get_torch_device
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
def convert_pt_to_safetensors(pt_path, safe_path):
try:
model = torch.load(pt_path, map_location="cpu")
if "state_dict" in model:
state_dict = model["state_dict"]
elif "model" in model:
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)
return True, ""
except Exception as e:
return False, str(e)
def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]:
"""Shared logic for all MetaKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
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 {}, {}
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:
return f"Conversion failed: {error_message}", ""
model_path = temp_safe_path
class BaseMetaKeys:
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("metadata", "keys")
FUNCTION = "get_metakeys"
CATEGORY = "utils"
DESCRIPTION = "Loads a model and returns its metadata and layer names."
# 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)
def _get_metakeys(self, model_name, model_type):
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
model_weights = {}
metadata = {}
layer_shapes = ""
# Extract metadata and layer shapes from header
metadata = header.pop("__metadata__", {})
metadata_str = str(metadata)
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:
return f"Conversion failed: {error_message}", ""
model_path = temp_safe_path
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)
model_weights, metadata = load_metadata_from_safetensors(model_path)
metadata_str = str(metadata)
for layer_name, tensor in model_weights.items():
shape = str(tensor.shape)
shape = shape.replace("torch.Size(", "").replace(")", "")
name_shape = f"{layer_name}, {shape}\n"
layer_shapes += name_shape
return metadata_str, layer_shapes
return metadata_str, layer_shapes
class UNetMetaKeys(BaseMetaKeys):
class ModelMetaKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"unet_name": (folder_paths.get_filename_list("diffusion_models"), ),
}
}
def define_schema(cls):
return io.Schema(
node_id="ModelMetaKeys",
display_name="Get Diffusion Model Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a diffusion model and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
def get_metakeys(self, unet_name):
return self._get_metakeys(unet_name, "diffusion_models")
class CLIPMetaKeys(BaseMetaKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip_name": (folder_paths.get_filename_list("text_encoders"), ),
}
}
def execute(cls, diffusionmodel_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(diffusionmodel_name, "diffusion_models")
return io.NodeOutput(metadata, keys)
def get_metakeys(self, clip_name):
return self._get_metakeys(clip_name, "text_encoders")
class LoRAMetaKeys(BaseMetaKeys):
class TextEncoderMetaKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"lora_name": (folder_paths.get_filename_list("loras"), ),
}
}
def get_metakeys(self, lora_name):
return self._get_metakeys(lora_name, "loras")
def define_schema(cls):
return io.Schema(
node_id="TextEncoderMetaKeys",
display_name="Get Text Encoder Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a text encoder and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
class CheckpointMetaKeys(BaseMetaKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
}
}
def execute(cls, textencoder_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(textencoder_name, "text_encoders")
return io.NodeOutput(metadata, keys)
def get_metakeys(self, ckpt_name):
return self._get_metakeys(ckpt_name, "checkpoints")
__all__ = ("UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys")
class LoRAMetaKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAMetaKeys",
display_name="Get LoRA Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a LoRA and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
@classmethod
def execute(cls, lora_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(lora_name, "loras")
return io.NodeOutput(metadata, keys)
class CheckpointMetaKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointMetaKeys",
display_name="Get Checkpoint Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a checkpoint and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
@classmethod
def execute(cls, ckpt_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(ckpt_name, "checkpoints")
return io.NodeOutput(metadata, keys)
class EmbeddingMetaKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingMetaKeys",
display_name="Get Embedding Metadata & Keys",
category="ModelUtils/Keys",
description="Loads an embedding and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
@classmethod
def execute(cls, embedding: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(embedding, "embeddings")
return io.NodeOutput(metadata, keys)
File diff suppressed because it is too large Load Diff
+184
View File
@@ -0,0 +1,184 @@
import os
import folder_paths
from comfy_api.latest import io
from .downloader_utils import get_potential_preview_files, load_image_tensor, get_model_workflows, get_model_metadata_file
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")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
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")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
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")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
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")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
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")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
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")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
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])
-31
View File
@@ -1,31 +0,0 @@
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
import comfy.model_management as mm
class OffloadModelsNode(ComfyNodeABC):
@classmethod
def INPUT_TYPES(s) -> InputTypeDict:
return {
"required": {
"any": (IO.ANY, {}),
"force": (IO.BOOLEAN, {"default": False, "tooltip": "If True, forcibly unloads models from GPU."}),
"enable_offload": (IO.BOOLEAN, {"default": False, "tooltip": "If True, calls ComfyUI's `unload_all_models()` function."}),
}
}
RETURN_TYPES = (IO.ANY, )
RETURN_NAMES = ("any", )
FUNCTION = 'function'
CATEGORY = 'advanced'
def function(self, any, force, enable_offload):
if enable_offload == True:
print("OffloadModels: `enable_offload` is True. Calling `mm.unload_all_models()`.")
mm.unload_all_models()
if Force == True:
mm.soft_empty_cache(Force=True)
else:
mm.soft_empty_cache()
return (any,)
__all__ = ("OffloadModelsNode")
+221
View File
@@ -0,0 +1,221 @@
import os
import re
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
use_regex: bool, output_filename: str) -> str:
"""Shared logic for all PruneKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
# Get metadata from safe_open (header only, no tensor loading)
with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
metadata = f.metadata() or {}
patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
if not patterns:
raise ValueError("No keys/patterns provided to prune.")
# 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")
# Stream tensors, filter on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
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
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match:
writer.write(key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
cleanup_after_operation()
return output_path
class ModelPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ModelPruneKeys",
display_name="Prune Diffusion Model Keys",
category="ModelUtils/Keys",
description="Loads a diffusion model, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_model"),
],
outputs=[
io.String.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)
return io.NodeOutput(path)
class TextEncoderPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderPruneKeys",
display_name="Prune Text Encoder Keys",
category="ModelUtils/Keys",
description="Loads a text encoder, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_textencoder"),
],
outputs=[
io.String.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)
return io.NodeOutput(path)
class LoRAPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAPruneKeys",
display_name="Prune LoRA Keys",
category="ModelUtils/Keys",
description="Loads a LoRA, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_lora"),
],
outputs=[
io.String.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)
return io.NodeOutput(path)
class CheckpointPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointPruneKeys",
display_name="Prune Checkpoint Keys",
category="ModelUtils/Keys",
description="Loads a checkpoint, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_checkpoint"),
],
outputs=[
io.String.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)
return io.NodeOutput(path)
class EmbeddingPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingPruneKeys",
display_name="Prune Embedding Keys",
category="ModelUtils/Keys",
description="Loads an embedding, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_embedding"),
],
outputs=[
io.String.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)
+145
View File
@@ -0,0 +1,145 @@
"""Fail-safe detection for quantized or isolated low-bit safetensors inputs."""
import logging
import json
import struct
import torch
MAX_ISOLATED_LOW_BIT_TENSORS = 2
def _low_bit_dtypes():
names = (
"int8",
"uint8",
"float4_e2m1fn_x2",
"float8_e4m3fn",
"float8_e4m3fnuz",
"float8_e5m2",
"float8_e5m2fnuz",
"float8_e8m0fnu",
)
return {getattr(torch, name) for name in names if hasattr(torch, name)}
LOW_BIT_DTYPES = _low_bit_dtypes()
LOW_BIT_STORAGE_CODES = {"I8", "U8", "F4"}
UPSTREAM_WRITER_STORAGE_CODES = {
"F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8",
"U64", "U32", "U16", "BOOL", "C64", "F8_E5M2", "F8_E4M3",
}
def _header(handler):
if getattr(handler, "_header", None) is not None:
return handler._header
with open(handler.filename, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
return json.loads(source.read(header_size).decode("utf-8"))
def storage_dtype_code(handler, key: str) -> str | None:
entry = _header(handler).get(key)
return entry.get("dtype") if entry else None
def _is_low_bit(handler, key: str) -> bool:
storage_code = storage_dtype_code(handler, key)
if storage_code in LOW_BIT_STORAGE_CODES or (
storage_code is not None and storage_code.startswith("F8_")
):
return True
return handler.get_dtype(key) in LOW_BIT_DTYPES
def inspect_low_bit_input(handler, label: str, operation: str) -> set[str]:
"""Reject quantized models; return isolated low-bit keys that must not be transformed."""
keys = list(handler.keys())
metadata = handler.metadata() or {}
quant_keys = [key for key in keys if key.endswith(".comfy_quant")]
has_quant_metadata = "_quantization_metadata" in metadata
has_legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in keys)
if quant_keys or has_quant_metadata or has_legacy_scaled_fp8:
raise ValueError(
f"[{operation}] {label} uses ComfyUI quantization metadata; quantized inputs are unsupported"
)
low_bit_keys = {key for key in keys if _is_low_bit(handler, key)}
if len(low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS:
raise ValueError(
f"[{operation}] {label} contains {len(low_bit_keys)} low-bit tensors; "
"casted or quantized inputs are unsupported"
)
if low_bit_keys:
logging.warning(
"[%s] %s contains %d isolated low-bit tensor(s); "
"preserving/excluding unchanged: %s",
operation,
label,
len(low_bit_keys),
", ".join(sorted(low_bit_keys)),
)
return low_bit_keys
def layer_has_low_bit(block_keys: dict, low_bit_keys: set[str]) -> bool:
return any(
block_keys.get(name) in low_bit_keys
for name in (
"down", "up", "alpha", "mid", "reshape", "dora_scale",
"diff", "diff_b", "w_norm", "b_norm", "set_weight",
)
)
def write_preserved_tensor(writer, key: str, handler, output_key: str | None = None) -> None:
"""Copy one tensor's original safetensors bytes and dtype into an active writer."""
destination = output_key or key
header = _header(handler)
entry = header[key]
if entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES:
writer.write(destination, handler.get_tensor(key).cpu().contiguous())
return
source_start, source_end = entry["data_offsets"]
byte_size = source_end - source_start
with writer._lock:
if destination in writer._manifest:
raise ValueError(f"Tensor '{destination}' has already been written.")
target_start = writer._current_data_offset
target_end = target_start + byte_size
writer._manifest[destination] = {
"dtype": entry["dtype"],
"shape": entry["shape"],
"data_offsets": [target_start, target_end],
}
writer._current_data_offset = target_end
source_path = handler.filename
source_absolute = 8 + getattr(handler, "_header_size", 0) + source_start
if not getattr(handler, "_header_size", None):
with open(source_path, "rb") as source:
source_absolute = 8 + struct.unpack("<Q", source.read(8))[0] + source_start
target_absolute = 8 + writer.max_header_bytes + target_start
def copy_raw():
try:
with open(source_path, "rb") as source:
source.seek(source_absolute)
remaining = byte_size
with writer._lock:
writer._file.seek(target_absolute)
while remaining:
chunk = source.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise EOFError(f"Unexpected EOF while preserving tensor '{key}'")
writer._file.write(chunk)
remaining -= len(chunk)
finally:
writer._semaphore.release()
writer._semaphore.acquire()
writer._futures.append(writer._executor.submit(copy_raw))
+216
View File
@@ -0,0 +1,216 @@
import os
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
new_keys_str: str, output_filename: str) -> str:
"""Shared logic for all RenameKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
# Get metadata from safe_open (header only, no tensor loading)
with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
metadata = f.metadata() or {}
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()]
if len(old_keys) != len(new_keys):
raise ValueError("The number of old keys must match the number of new keys.")
key_map = dict(zip(old_keys, new_keys))
# 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")
# Stream tensors, rename on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
for key in tqdm(original_keys, desc="Renaming keys", unit="keys"):
new_key = key_map.get(key, key)
writer.write(new_key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
cleanup_after_operation()
return output_path
class ModelRenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ModelRenameKeys",
display_name="Rename Diffusion Model Keys",
category="ModelUtils/Keys",
description="Loads a diffusion model, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_model"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, diffusionmodel_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(diffusionmodel_name, "diffusion_models",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
class TextEncoderRenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderRenameKeys",
display_name="Rename Text Encoder Keys",
category="ModelUtils/Keys",
description="Loads a text encoder, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_textencoder"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, textencoder_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(textencoder_name, "text_encoders",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
class LoRARenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRARenameKeys",
display_name="Rename LoRA Keys",
category="ModelUtils/Keys",
description="Loads a LoRA, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_lora"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, lora_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(lora_name, "loras",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
class CheckpointRenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointRenameKeys",
display_name="Rename Checkpoint Keys",
category="ModelUtils/Keys",
description="Loads a checkpoint, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_checkpoint"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, ckpt_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(ckpt_name, "checkpoints",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
class EmbeddingRenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingRenameKeys",
display_name="Rename Embedding Keys",
category="ModelUtils/Keys",
description="Loads an embedding, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_embedding"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, embedding: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(embedding, "embeddings",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
+721
View File
@@ -0,0 +1,721 @@
"""
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 fnmatch
import os
import re
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# Import SVD rank utilities from standard lora extract file to avoid duplication
from .lora_extract_svd import (
_compute_rank,
_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 = True,
glob_skip_patterns: bool = False,
) -> 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 = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_te_lora_key(key)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key)
if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
dora_scale = None
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
weight_diff = tensor_a
dora_scale = None
del tensor_b
else:
# Skip 1D weights
if tensor_a.ndim < 2:
del tensor_a, tensor_b
return "skipped", None
if is_dora:
# Weight Decomposition calculation
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
# Compute directional diff
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del tensor_a, tensor_b, w_dir
else:
weight_diff = tensor_a - tensor_b
dora_scale = None
del tensor_a, tensor_b
# Skip small differences
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[TE Extract] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_layer(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
# 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 = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting TE Layers", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[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) -> str:
"""Build output path for LoRA file."""
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
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=True, tooltip="Clear CUDA cache after each layer"),
io.Boolean.Input("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
tooltip="Clamp outlier singular values"),
io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.String.Input("output_filename", default="extracted_te_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
]
# =============================================================================
# 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(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
# =============================================================================
# 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(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
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),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _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
)
return io.NodeOutput(output_path)
+34
View File
@@ -0,0 +1,34 @@
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from comfy.model_management import get_torch_device
def convert_pt_to_safetensors(pt_path, safe_path):
try:
model = torch.load(pt_path, map_location="cpu")
if "state_dict" in model:
state_dict = model["state_dict"]
elif "model" in model:
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)
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 {}, {}
+12 -1
View File
@@ -1,9 +1,20 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.1.3"
version = "0.5.2"
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
+5
View File
@@ -0,0 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader>=0.5.2
av
+845
View File
@@ -0,0 +1,845 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_cwb_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def cwb():
return _load_module("nodes.consensus_merger")
def _patch_io(monkeypatch, module, tmp_path, paths):
monkeypatch.setattr(module.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _result_path(tmp_path, category, result):
return tmp_path / category / result
def _params(output_filename, **overrides):
params = {
"cwb_preset": "custom",
"consensus_type": "mean",
"alignment_method": "index",
"alignment_threshold": 0.4,
"similarity_threshold": 0.0,
"power_alpha": 2.0,
"diversity_beta": 0.0,
"rescale_norm": False,
"global_scale": 1.0,
"dynamic_similarity_contrast": False,
"soft_comfort_bandpass": False,
"position_weight": 0.0,
"preserve_common_prefix": False,
"mismatch_mode": "skip",
"output_filename": output_filename,
"save_dtype": "fp16",
"process_device": "cpu",
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"override_dtype": False,
"include_1d_diffs": False,
}
params.update(overrides)
return params
def test_preset_resolution_matches_extended_contract(cwb):
assert all(
preset["global_scale"] == pytest.approx(1.0)
for preset in cwb.CWB_PRESETS.values()
)
power = cwb.resolve_cwb_settings(cwb_preset="power_blend")
assert power.consensus_type == "median"
assert power.alignment_threshold == pytest.approx(0.9)
assert power.power_alpha == pytest.approx(8.0)
assert power.dynamic_similarity_contrast is True
assert power.soft_comfort_bandpass is False
baseline = cwb.resolve_cwb_settings(
cwb_preset="baseline",
alignment_method="index",
dynamic_similarity_contrast=True,
soft_comfort_bandpass=True,
)
assert baseline.alignment_method == "similarity"
assert baseline.dynamic_similarity_contrast is False
assert baseline.soft_comfort_bandpass is False
varied = cwb.resolve_cwb_settings(
cwb_preset="varied_merge",
global_scale=1.25,
position_weight=0.2,
preserve_common_prefix=True,
)
assert varied.global_scale == pytest.approx(1.25)
assert varied.position_weight == pytest.approx(0.2)
assert varied.preserve_common_prefix is True
def test_cwb_math_index_prefix_and_scale(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=1.0,
)
first = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
second = torch.tensor([[3.0, 0.0], [0.0, 3.0]])
torch.testing.assert_close(
cwb.merge_cwb_tensors([first, second], settings),
torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
)
prefix_settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=2.0,
preserve_common_prefix=True,
)
prefixed = cwb.merge_cwb_tensors(
[torch.tensor([[9.0, 9.0], [1.0, 0.0]]),
torch.tensor([[9.0, 9.0], [3.0, 0.0]])],
prefix_settings,
)
torch.testing.assert_close(prefixed[0], torch.tensor([9.0, 9.0]))
torch.testing.assert_close(prefixed[1], torch.tensor([4.0, 0.0]))
def test_similarity_alignment_reorders_matching_rows(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
torch.testing.assert_close(
cwb.merge_cwb_tensors([reference, reversed_source], settings),
torch.tensor([[1.5, 0.0], [0.0, 1.5]]),
)
scores = torch.ones((2, 2))
biased = cwb._position_biased_scores(scores, 1.0)
assert biased[0, 0] > biased[0, 1]
def test_fixed_coordinate_merge_never_reorders_model_rows(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
expected = torch.stack([
cwb.merge_consensus_group(
torch.stack([reference[row], reversed_source[row]]), settings
)
for row in range(reference.shape[0])
])
torch.testing.assert_close(
cwb.merge_cwb_tensors(
[reference, reversed_source],
settings,
allow_similarity_alignment=False,
),
expected,
)
def test_lora_similarity_alignment_pairs_a_rows_with_b_columns(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference_down = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reference_up = torch.tensor([[2.0, 0.0], [0.0, 3.0]])
source_down = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
source_up = torch.tensor([[0.0, -2.0], [3.0, 0.0]])
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[reference_down, source_down],
[reference_up, source_up],
settings,
reference_index=0,
)
torch.testing.assert_close(merged_down, reference_down)
torch.testing.assert_close(merged_up, reference_up)
def test_lora_krea_shape_scores_only_rank_components(monkeypatch, cwb):
rank = 256
output_features = 36864
down = torch.ones((rank, 2))
up = torch.ones((output_features, rank))
matrix_shapes = []
def bounded_mm(left, right):
matrix_shapes.append((left.shape, right.shape))
return torch.eye(left.shape[0], right.shape[1], dtype=left.dtype)
monkeypatch.setattr(cwb.torch, "mm", bounded_mm)
monkeypatch.setattr(
cwb,
"merge_consensus_group",
lambda stacked, settings, **kwargs: stacked[0].clone(),
)
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
alignment_method="similarity",
alignment_threshold=0.5,
)
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[down, down], [up, up], settings, reference_index=0
)
assert merged_down.shape == down.shape
assert merged_up.shape == up.shape
assert matrix_shapes
assert all(left[0] == rank and right[1] == rank for left, right in matrix_shapes)
def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=3.0,
)
down = torch.ones((1, 1))
up = torch.ones((1, 1))
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
[(0, down, up, 2.0), (1, down, up, 2.0)],
settings,
"cpu",
torch.float32,
)
assert rank == 1
torch.testing.assert_close(merged_down, torch.ones((1, 1)))
torch.testing.assert_close(merged_up, torch.full((1, 1), 6.0))
def test_lora_pair_alignment_supports_convolution_factors(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
)
down = torch.arange(6, dtype=torch.float32).reshape(2, 3, 1, 1)
up = torch.arange(8, dtype=torch.float32).reshape(4, 2, 1, 1)
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[down, down], [up, up], settings, reference_index=0
)
torch.testing.assert_close(merged_down, down)
torch.testing.assert_close(merged_up, up)
def test_lora_cuda_oom_retries_current_pair_on_cpu(monkeypatch, caplog, cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
)
original_to_compute = cwb._to_compute
attempted_devices = []
def fail_cuda(tensor, device):
attempted_devices.append(device)
if str(device).startswith("cuda"):
raise torch.OutOfMemoryError("injected allocation failure")
return original_to_compute(tensor, device)
monkeypatch.setattr(cwb, "_to_compute", fail_cuda)
monkeypatch.setattr(cwb, "_release_failed_cuda_operation", lambda: None)
down = torch.ones((1, 2))
up = torch.ones((2, 1))
with caplog.at_level("WARNING"):
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
[(0, down, up, 1.0), (1, down, up, 1.0)],
settings,
"cuda",
torch.float32,
operation_label="diffusion_model.foo",
)
assert rank == 1
torch.testing.assert_close(merged_down, down)
torch.testing.assert_close(merged_up, up)
assert attempted_devices[0] == "cuda"
assert "cpu" in attempted_devices
assert "retrying this layer on CPU" in caplog.text
assert "diffusion_model.foo" in caplog.text
def test_norm_rescale_and_dsc_bandpass_are_finite(cwb):
rescale = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
rescale_norm=True,
)
merged = cwb.merge_consensus_group(
torch.tensor([[2.0, 0.0], [0.0, 2.0]]), rescale
)
assert torch.linalg.vector_norm(merged).item() == pytest.approx(2.0)
dsc = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="median",
diversity_beta=1.5,
dynamic_similarity_contrast=True,
soft_comfort_bandpass=True,
)
result = cwb.merge_consensus_group(
torch.tensor([[1.0, 0.0], [0.8, 0.2], [0.0, 1.0]]), dsc
)
assert torch.isfinite(result).all()
def test_all_ten_schemas_and_lora_switch_contract(cwb):
assert len(cwb.CWB_MERGER_NODES) == 10
assert len({node.NODE_ID for node in cwb.CWB_MERGER_NODES}) == 10
for node in cwb.CWB_MERGER_NODES:
schema = node.define_schema()
ids = [value.id for value in schema.inputs]
assert ids[:3] == ["execution_mode", "model_a", "model_b"]
if node.INPUT_COUNT == 3:
assert ids[3] == "model_c"
if node.LORA_MODE:
assert schema.inputs[-1].id == "include_1d_diffs"
assert schema.inputs[-1].default is False
else:
assert "include_1d_diffs" not in ids
def test_streaming_dtype_nonfloat_and_secondary_preservation(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "a.safetensors"
b = tmp_path / "b.safetensors"
save_file({
"shared": torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32),
"only_a": torch.tensor([5.0, 6.0], dtype=torch.float16),
"metadata_tensor": torch.tensor([1, 2], dtype=torch.int64),
}, str(a), metadata={"source": "A"})
save_file({
"shared": torch.tensor([[3.0, 0.0], [0.0, 3.0]], dtype=torch.float16),
"extra": torch.ones((2, 2)),
"metadata_tensor": torch.tensor([9, 9], dtype=torch.int64),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("anchored")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert set(tensors) == {"shared", "only_a", "extra", "metadata_tensor"}
torch.testing.assert_close(tensors["shared"], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
assert tensors["shared"].dtype == torch.float32
torch.testing.assert_close(tensors["only_a"], torch.tensor([5.0, 6.0], dtype=torch.float16))
torch.testing.assert_close(tensors["metadata_tensor"], torch.tensor([1, 2]))
torch.testing.assert_close(tensors["extra"], torch.ones((2, 2)))
overridden = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("anchored_override", override_dtype=True),
)
overridden_tensors = load_file(
str(_result_path(tmp_path, "diffusion_models", overridden))
)
assert overridden_tensors["shared"].dtype == torch.float16
assert overridden_tensors["metadata_tensor"].dtype == torch.int64
def test_failed_merge_preserves_existing_output_and_removes_temporary_file(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "atomic_a.safetensors"
b = tmp_path / "atomic_b.safetensors"
output = tmp_path / "diffusion_models" / "atomic_output.safetensors"
output.parent.mkdir(parents=True)
save_file({"layer": torch.ones((2, 2))}, str(a))
save_file({"layer": torch.full((2, 2), 2.0)}, str(b))
save_file({"existing": torch.tensor([7.0])}, str(output))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
def fail_merge(*args, **kwargs):
raise RuntimeError("injected CWB failure")
monkeypatch.setattr(cwb, "merge_cwb_tensors", fail_merge)
with pytest.raises(RuntimeError, match="injected CWB failure"):
cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("atomic_output")
)
existing = load_file(str(output))
assert set(existing) == {"existing"}
torch.testing.assert_close(existing["existing"], torch.tensor([7.0]))
assert not list(output.parent.glob(".atomic_output.safetensors.*.tmp"))
def test_embedding_union_uses_longest_first_dimension(monkeypatch, tmp_path, cwb):
a = tmp_path / "embed_a.safetensors"
b = tmp_path / "embed_b.safetensors"
save_file({"emb": torch.tensor([[1.0, 0.0], [0.0, 1.0]])}, str(a))
save_file({
"emb": torch.tensor([[3.0, 0.0], [0.0, 3.0], [4.0, 4.0]]),
"secondary_only": torch.tensor([[7.0, 0.0]]),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"embeddings",
_params("embedding_union"),
embedding_union=True,
)
tensors = load_file(str(_result_path(tmp_path, "embeddings", result)))
assert tensors["emb"].shape == (3, 2)
torch.testing.assert_close(tensors["emb"][:2], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
torch.testing.assert_close(tensors["emb"][2], torch.tensor([4.0, 4.0]))
torch.testing.assert_close(tensors["secondary_only"], torch.tensor([[7.0, 0.0]]))
def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb):
paths = {}
for name, value in (("a", 1.0), ("b", 3.0), ("c", 5.0)):
path = tmp_path / f"three_{name}.safetensors"
save_file({"layer": torch.tensor([[value, 0.0]])}, str(path))
paths[name] = str(path)
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("three_inputs")
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
torch.testing.assert_close(tensor, torch.tensor([[3.0, 0.0]]))
def test_two_input_preserves_secondary_only_tensor_exactly(monkeypatch, tmp_path, cwb):
a = tmp_path / "secondary_a.safetensors"
b = tmp_path / "secondary_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary": torch.tensor([2.0, 3.0], dtype=torch.float32),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params(
"secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
assert tensor.dtype == torch.float32
def test_three_input_merges_available_secondary_sources(monkeypatch, tmp_path, cwb):
a = tmp_path / "union_a.safetensors"
b = tmp_path / "union_b.safetensors"
c = tmp_path / "union_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([2.0]),
"b_only": torch.tensor([7.0]),
}, str(b))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([4.0]),
"c_only": torch.tensor([9.0]),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("secondary_union")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
torch.testing.assert_close(tensors["secondary_shared"], torch.tensor([3.0]))
torch.testing.assert_close(tensors["b_only"], torch.tensor([7.0]))
torch.testing.assert_close(tensors["c_only"], torch.tensor([9.0]))
def test_secondary_shape_conflict_skip_preserves_earliest_source(monkeypatch, tmp_path, cwb):
a = tmp_path / "shape_a.safetensors"
b = tmp_path / "shape_b.safetensors"
c = tmp_path / "shape_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({"secondary": torch.tensor([2.0, 3.0])}, str(b))
save_file({"secondary": torch.tensor([[8.0, 9.0]])}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"],
"diffusion_models",
_params("secondary_shape_skip", mismatch_mode="skip"),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
def test_secondary_only_filters_and_low_bit_preservation(monkeypatch, tmp_path, cwb):
a = tmp_path / "filter_a.safetensors"
b = tmp_path / "filter_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"secondary.keep": torch.tensor([4, 5], dtype=torch.uint8),
"secondary.drop": torch.tensor([6.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("secondary_filters", discard_patterns=r"secondary\.drop"),
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert "secondary.drop" not in tensors
torch.testing.assert_close(
tensors["secondary.keep"], torch.tensor([4, 5], dtype=torch.uint8)
)
def test_lora_cross_format_companion_group_preserves_model_a_canonically(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_a.safetensors"
b = tmp_path / "lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]], dtype=torch.float32),
"diffusion_model.foo.lora_B.weight": torch.tensor([[1.0], [2.0]], dtype=torch.float32),
"diffusion_model.foo.dora_scale": torch.tensor([1.0, 1.0]),
}, str(a))
save_file({
"lora_unet_foo.lora_down.weight": torch.tensor([[3.0, 4.0], [5.0, 6.0]]),
"lora_unet_foo.lora_up.weight": torch.tensor([[3.0, 5.0], [4.0, 6.0]]),
"lora_unet_secondary.lora_down.weight": torch.ones((1, 2)),
"lora_unet_secondary.lora_up.weight": torch.ones((2, 1)),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("lora_cross_format"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.dora_scale",
"lora_unet_secondary.lora_A.weight",
"lora_unet_secondary.lora_B.weight",
}
assert tensors["diffusion_model.foo.lora_A.weight"].shape == (1, 2)
assert tensors["diffusion_model.foo.lora_B.weight"].shape == (2, 1)
assert tensors["diffusion_model.foo.lora_A.weight"].dtype == torch.float32
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([1.0, 1.0])
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_A.weight"], torch.ones((1, 2))
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_B.weight"], torch.ones((2, 1))
)
def test_lora_secondary_only_pair_is_preserved_without_scaling(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_secondary_a.safetensors"
b = tmp_path / "lora_secondary_b.safetensors"
save_file({
"diffusion_model.base.diff": torch.tensor([[1.0]]),
}, str(a))
save_file({
"transformer.extra.lora_down.weight": torch.tensor([[2.0, 3.0]]),
"transformer.extra.lora_up.weight": torch.tensor([[4.0], [5.0]]),
"transformer.extra.alpha": torch.tensor(1.0),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params(
"lora_secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"],
torch.tensor([[2.0, 3.0]]),
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"],
torch.tensor([[4.0], [5.0]]),
)
assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.extra.alpha"].item() == pytest.approx(1.0)
def test_lora_three_input_merges_b_and_c_group_without_a(monkeypatch, tmp_path, cwb):
a = tmp_path / "lora_union_a.safetensors"
b = tmp_path / "lora_union_b.safetensors"
c = tmp_path / "lora_union_c.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.extra.lora_A.weight": torch.full((1, 2), 1.0),
"diffusion_model.extra.lora_B.weight": torch.full((2, 1), 1.0),
}, str(b))
save_file({
"transformer.extra.lora_down.weight": torch.full((1, 2), 3.0),
"transformer.extra.lora_up.weight": torch.full((2, 1), 3.0),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "loras", _params("lora_secondary_union"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.full((1, 2), 2.0)
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.full((2, 1), 2.0)
)
def test_lora_secondary_direct_and_companion_groups_are_atomic(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_atomic_a.safetensors"
b = tmp_path / "lora_atomic_b.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.norm.diff": torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
"diffusion_model.extra.lora_A.weight": torch.tensor([[4.0, 5.0]]),
"diffusion_model.extra.lora_B.weight": torch.tensor([[6.0], [7.0]]),
"diffusion_model.extra.dora_scale": torch.tensor([8.0, 9.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("lora_secondary_atomic", global_scale=3.0),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.norm.diff"],
torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
)
assert tensors["diffusion_model.norm.diff"].dtype == torch.bfloat16
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.tensor([[4.0, 5.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.tensor([[6.0], [7.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.dora_scale"], torch.tensor([8.0, 9.0])
)
def test_lora_rejects_duplicate_normalized_cores_in_secondary_input(cwb):
with pytest.raises(ValueError, match="same logical layer 'foo'.*input 2"):
cwb._secondary_lora_map(
{
"diffusion_model.foo": {},
"transformer.foo": {},
},
input_index=1,
)
def test_lora_peft_prefix_matching_and_existing_alpha(monkeypatch, tmp_path, cwb):
a = tmp_path / "peft_a.safetensors"
b = tmp_path / "peft_b.safetensors"
save_file({
"base_model.model.diffusion_model.foo.lora_A.weight": torch.ones((1, 2)),
"base_model.model.diffusion_model.foo.lora_B.weight": torch.ones((2, 1)),
"base_model.model.diffusion_model.foo.alpha": torch.tensor(1.0),
}, str(a))
save_file({
"diffusion_model.foo.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.foo.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.foo.alpha": torch.tensor(2.0),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("peft_prefix"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.alpha",
}
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(2.0)
def test_lora_mochi_inputs_emit_preferred_canonical_output(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "mochi_a.safetensors"
b = tmp_path / "mochi_b.safetensors"
for path, value in ((a, 1.0), (b, 3.0)):
save_file({
"diffusion_model.foo.lora_A": torch.full((1, 2), value),
"diffusion_model.foo.lora_B": torch.full((2, 1), value),
"diffusion_model.foo.alpha": torch.tensor(1.0),
}, str(path))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("mochi_canonical"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.alpha",
}
def test_lora_1d_direct_switch_is_default_fallback_and_fp32(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "direct_a.safetensors"
b = tmp_path / "direct_b.safetensors"
save_file({"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16)}, str(a))
save_file({"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16)}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
default = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("direct_default"), lora_mode=True
)
default_tensor = load_file(str(_result_path(tmp_path, "loras", default)))[
"diffusion_model.norm.diff"
]
torch.testing.assert_close(default_tensor, torch.tensor([1.0, 2.0], dtype=torch.bfloat16))
assert default_tensor.dtype == torch.bfloat16
enabled = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("direct_enabled", include_1d_diffs=True, override_dtype=True),
lora_mode=True,
)
enabled_tensor = load_file(str(_result_path(tmp_path, "loras", enabled)))[
"diffusion_model.norm.diff"
]
torch.testing.assert_close(enabled_tensor, torch.tensor([2.0, 3.0]))
assert enabled_tensor.dtype == torch.float32
def test_quantized_marker_errors_before_output(monkeypatch, tmp_path, cwb):
a = tmp_path / "quant_a.safetensors"
b = tmp_path / "quant_b.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor([1], dtype=torch.uint8),
}, str(a))
save_file({"layer.weight": torch.ones((2, 2))}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("must_not_exist")
)
assert not (tmp_path / "diffusion_models" / "must_not_exist.safetensors").exists()
def test_isolated_low_bit_preserves_anchored_tensor(monkeypatch, tmp_path, cwb):
a = tmp_path / "guard_a.safetensors"
b = tmp_path / "guard_b.safetensors"
save_file({"layer": torch.tensor([1.0, 2.0])}, str(a))
save_file({"layer": torch.tensor([8, 9], dtype=torch.uint8)}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("guarded_generic")
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
torch.testing.assert_close(tensor, torch.tensor([1.0, 2.0]))
assert tensor.dtype == torch.float32
def test_isolated_low_bit_preserves_complete_lora_layer(monkeypatch, tmp_path, cwb):
a = tmp_path / "guard_lora_a.safetensors"
b = tmp_path / "guard_lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.foo.alpha": torch.tensor(1.0),
"diffusion_model.foo.dora_scale": torch.tensor([5.0, 6.0]),
}, str(a))
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[9.0, 9.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.foo.alpha": torch.tensor(1, dtype=torch.uint8),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("guarded_lora"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.foo.lora_A.weight"], torch.tensor([[1.0, 2.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.foo.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(1.0)
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([5.0, 6.0])
)
+899
View File
@@ -0,0 +1,899 @@
import importlib
import logging
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_contract_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def modules():
return _load_module("nodes.lora_resize"), _load_module("nodes.lora_merger")
@pytest.fixture(scope="module")
def generic_modules():
return _load_module("nodes.merger"), _load_module("nodes.merger_ops")
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
COMFY_LORA_PAIR_SUFFIXES = [
(".lora_down.weight", ".lora_up.weight"),
("_lora.down.weight", "_lora.up.weight"),
(".lora_A.weight", ".lora_B.weight"),
(".lora.down.weight", ".lora.up.weight"),
(".lora_A", ".lora_B"),
(".lora_linear_layer.down.weight", ".lora_linear_layer.up.weight"),
(".lora_A.default.weight", ".lora_B.default.weight"),
]
@pytest.mark.parametrize("down_suffix,up_suffix", COMFY_LORA_PAIR_SUFFIXES)
def test_all_comfy_lora_pair_spellings_parse_to_canonical_keys(
modules, down_suffix, up_suffix
):
resize, _ = modules
block = "diffusion_model.layer"
layers, passthrough = resize.parse_lora_layers([
f"{block}{down_suffix}",
f"{block}{up_suffix}",
f"{block}.alpha",
])
assert not passthrough
assert layers[block]["down"] == f"{block}{down_suffix}"
assert layers[block]["up"] == f"{block}{up_suffix}"
assert resize.canonical_lora_key(block, "down") == f"{block}.lora_A.weight"
assert resize.canonical_lora_key(block, "up") == f"{block}.lora_B.weight"
@pytest.mark.parametrize("down_suffix,up_suffix", COMFY_LORA_PAIR_SUFFIXES)
def test_merge_to_model_applies_every_comfy_lora_pair_spelling(
monkeypatch, tmp_path, modules, down_suffix, up_suffix
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
variant = COMFY_LORA_PAIR_SUFFIXES.index((down_suffix, up_suffix))
base = tmp_path / f"base_{variant}.safetensors"
adapter = tmp_path / f"adapter_{variant}.safetensors"
block = "diffusion_model.layer"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
}, str(base))
save_file({
f"{block}{down_suffix}": torch.tensor([[1.0, 2.0]]),
f"{block}{up_suffix}": torch.tensor([[3.0], [4.0]]),
f"{block}.alpha": torch.tensor(0.5),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
f"merged_variant_{variant}", verbose=False,
)
result = load_file(output)["model.diffusion_model.layer.weight"]
torch.testing.assert_close(
result,
torch.tensor([[1.5, 3.0], [2.0, 4.0]]),
)
assert result.dtype == torch.float32
def test_merge_to_model_applies_comfy_direct_norm_and_set_forms(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "direct_base.safetensors"
adapter = tmp_path / "direct_adapter.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.norm.bias": torch.zeros(2),
"model.diffusion_model.replace.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.w_norm": torch.tensor([1.0, 2.0]),
"diffusion_model.norm.b_norm": torch.tensor([3.0, 4.0]),
"diffusion_model.replace.set_weight": torch.full((2, 2), 7.0),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16,
"direct_forms", verbose=False, include_1d_diffs=True,
)
result = load_file(output)
torch.testing.assert_close(
result["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0])
)
torch.testing.assert_close(
result["model.diffusion_model.norm.bias"], torch.tensor([1.5, 2.0])
)
torch.testing.assert_close(
result["model.diffusion_model.replace.weight"], torch.full((2, 2), 7.0)
)
def test_merge_to_model_uses_comfy_reshape_companion(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "reshape_base.safetensors"
adapter = tmp_path / "reshape_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((1, 2)),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.layer.reshape_weight": torch.tensor([2, 2]),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
"reshape_companion", verbose=False,
)
result = load_file(output)["model.diffusion_model.layer.weight"]
assert result.shape == (2, 2)
torch.testing.assert_close(result, torch.tensor([[3.0, 6.0], [4.0, 8.0]]))
def test_merge_to_model_reports_success_and_separates_base_only_biases(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "summary_base.safetensors"
adapter = tmp_path / "summary_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"summary", verbose=True, return_report=True,
)
output = capsys.readouterr().out
assert Path(path).is_file()
assert report.startswith("SUCCESS: 1/1 adapter groups fully applied")
assert "[PATCHED]" not in report
assert "[APPLIED]" not in report
assert "L1: summary_adapter.safetensors, strength=1.0" in report
assert "[BASE-ONLY BIASES RETAINED] (1)" in report
assert "model.diffusion_model.layer.bias" in report
assert "SUCCESS: 1/1 groups applied" in output
assert "1 tensors patched, 1 base tensors not patched, 0 adapter issues" in output
assert "model.diffusion_model.layer.bias" not in output
def test_success_report_stays_compact_and_console_is_summary_only(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "many_success_base.safetensors"
adapter = tmp_path / "many_success_adapter.safetensors"
save_file({
f"model.diffusion_model.layer_{index}.weight": torch.zeros((2, 2))
for index in range(40)
}, str(base))
save_file({
f"diffusion_model.layer_{index}.diff": torch.ones((2, 2))
for index in range(40)
}, str(adapter))
_, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"many_success", verbose=True, return_report=True,
)
console = capsys.readouterr().out
assert report.startswith("SUCCESS: 40/40")
assert "layer_0.weight" not in report
assert "ADAPTER GROUP ISSUES\n- None" in report
assert len(report) < 1000
assert max(map(len, report.splitlines())) < 160
assert "layer_0.weight" not in console
assert len(console.splitlines()) <= 9
def test_merge_to_model_report_identifies_shape_mismatch_and_unresolved_group(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "diagnostic_base.safetensors"
adapter = tmp_path / "diagnostic_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.layer.diff": torch.ones(3),
"diffusion_model.orphan.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.orphan.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
_, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"diagnostic", verbose=False, include_1d_diffs=True, return_report=True,
)
assert report.startswith("COMPLETED WITH ERRORS: 0/2")
assert "[MAPPED BUT UNCHANGED] (1)" in report
assert "[SHAPE MISMATCH] (1)" in report
assert "source_shape=(3,), target_shape=(2, 2)" in report
assert "[UNRESOLVED] (1)" in report
assert "diffusion_model.orphan" in report
assert "no matching base tensor" in report
def test_merge_to_model_detects_comfy_adapter_logged_failure_atomically(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "logged_failure_base.safetensors"
adapter = tmp_path / "logged_failure_adapter.safetensors"
base_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
save_file({"model.diffusion_model.layer.weight": base_weight}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
class LoggedFailureAdapter:
def calculate_weight(self, weight, key, *args, **kwargs):
weight.add_(99.0)
logging.error("ERROR lora %s synthetic internal failure", key)
return weight
class FakeLoRAAdapter:
@staticmethod
def load(*args, **kwargs):
return LoggedFailureAdapter()
monkeypatch.setattr(resize, "LoRAAdapter", FakeLoRAAdapter)
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"logged_failure", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.layer.weight"], base_weight
)
assert report.startswith("COMPLETED WITH ERRORS: 0/1")
assert "[CALCULATION FAILURE] (1)" in report
assert "synthetic internal failure" in report
assert "[MAPPED BUT UNCHANGED] (1)" in report
def test_merge_to_model_report_classifies_disabled_1d_as_exclusion(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "disabled_1d_base.safetensors"
adapter = tmp_path / "disabled_1d_adapter.safetensors"
save_file({"model.diffusion_model.norm.weight": torch.zeros(2)}, str(base))
save_file({"diffusion_model.norm.diff": torch.ones(2)}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"disabled_1d", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.norm.weight"], torch.zeros(2)
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 0/1")
assert "[1D DISABLED] (1)" in report
assert "source=diffusion_model.norm.diff" in report
assert "include_1d_diffs is disabled" in report
def test_merge_to_model_report_marks_partially_applied_group(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "partial_base.safetensors"
adapter = tmp_path / "partial_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
"diffusion_model.layer.diff_b": torch.ones(3),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"partial", verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(path)
torch.testing.assert_close(
result["model.diffusion_model.layer.weight"], torch.ones((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.layer.bias"], torch.zeros(2)
)
assert report.startswith("COMPLETED WITH ERRORS: 0/1")
assert "[PARTIALLY APPLIED] (1)" in report
assert "source_shape=(3,), target_shape=(2,)" in report
def test_merge_to_model_detects_raised_adapter_failure_atomically(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "raised_failure_base.safetensors"
adapter = tmp_path / "raised_failure_adapter.safetensors"
base_weight = torch.eye(2)
save_file({"model.diffusion_model.layer.weight": base_weight}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
class RaisedFailureAdapter:
def calculate_weight(self, weight, key, *args, **kwargs):
weight.zero_()
raise RuntimeError("synthetic raised failure")
class FakeLoRAAdapter:
@staticmethod
def load(*args, **kwargs):
return RaisedFailureAdapter()
monkeypatch.setattr(resize, "LoRAAdapter", FakeLoRAAdapter)
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"raised_failure", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.layer.weight"], base_weight
)
assert "[CALCULATION FAILURE] (1)" in report
assert "RuntimeError: synthetic raised failure" in report
def test_merge_to_model_report_lists_skipped_tensor_and_pattern(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "skip_base.safetensors"
adapter = tmp_path / "skip_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"skip", skip_patterns_str=r"layer\.weight$", verbose=False,
return_report=True,
)
assert "model.diffusion_model.layer.weight" not in load_file(path)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 0/1")
assert "[OMITTED BY SKIP PATTERN] (1)" in report
assert r"pattern=layer\.weight$" in report
assert "[SKIPPED BY PATTERN] (1)" in report
def test_merge_to_model_node_returns_path_blank_line_report(
monkeypatch, modules
):
resize, _ = modules
monkeypatch.setattr(
resize.folder_paths, "get_full_path_or_raise",
lambda category, name: f"{category}/{name}",
)
monkeypatch.setattr(
resize, "merge_loras_to_model",
lambda **kwargs: ("saved/model.safetensors", "SUCCESS\nDETAILS"),
)
result = resize.LoRAMergeToModel.execute(
"base.safetensors", "1",
"one.safetensors", 1.0,
"None", 1.0, "None", 1.0, "None", 1.0,
"None", 1.0, "None", 1.0, "None", 1.0, "None", 1.0,
"", "merged", "fp32", "cpu", True, False, False,
)
assert result.result == ("saved/model.safetensors\n\nSUCCESS\nDETAILS",)
def test_mixed_parser_scans_every_key_and_retains_unpaired(modules):
resize, _ = modules
filler = [f"metadata.tensor_{index}" for index in range(60)]
keys = filler + [
"diffusion_model.linear.lora_down.weight",
"diffusion_model.linear.lora_up.weight",
"diffusion_model.linear.alpha",
"diffusion_model.norm.diff",
"diffusion_model.block.diff_b",
"diffusion_model.orphan.lora_down.weight",
"diffusion_model.linear.dora_scale",
]
info = resize.detect_lora_format(keys)
layers, passthrough = resize.parse_lora_layers(keys)
assert info["format"] == "mixed"
assert layers["diffusion_model.linear"]["down"].endswith("lora_down.weight")
assert layers["diffusion_model.norm"]["diff"].endswith(".diff")
assert layers["diffusion_model.block"]["diff_b"].endswith(".diff_b")
assert "diffusion_model.orphan.lora_down.weight" in passthrough
assert layers["diffusion_model.linear"]["dora_scale"].endswith(".dora_scale")
def test_dtype_contract(modules):
resize, _ = modules
assert resize.select_output_dtype([torch.bfloat16], torch.float16, is_1d_diff=True) == torch.float32
assert resize.select_output_dtype([torch.float32, torch.bfloat16], torch.float16) == torch.float32
assert resize.select_output_dtype([torch.float32], torch.bfloat16, force=True) == torch.bfloat16
def test_include_1d_switches_are_appended_and_default_off(modules, generic_modules):
resize, multi = modules
generic, _ = generic_modules
node_classes = [
generic.LoRATwoMerger,
generic.LoRAThreeMerger,
multi.LoRAMultiMerge,
multi.LoRAMultiMergeDARE,
multi.LoRAMultiMergeDAREEnhanced,
resize.LoRAMergeToModel,
]
for node_class in node_classes:
final_input = node_class.define_schema().inputs[-1]
assert final_input.id == "include_1d_diffs"
assert final_input.default is False
assert len(resize.LoRAMergeToModel.define_schema().outputs) == 1
def test_resize_preserves_mixed_direct_auxiliary_and_fp32(monkeypatch, tmp_path, modules):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mixed.safetensors"
save_file({
"diffusion_model.linear.lora_down.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4),
"diffusion_model.linear.lora_up.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2),
"diffusion_model.linear.alpha": torch.tensor(2.0, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([0.125, -0.25], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.linear.dora_scale": torch.ones(4, dtype=torch.float32),
"diffusion_model.orphan.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resized", verbose=False
)
tensors = load_file(output)
assert tensors["diffusion_model.linear.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.linear.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.lora_B.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.alpha"].dtype == torch.float32
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.scale.diff"].dtype == torch.float16
assert tensors["diffusion_model.lin.diff"].dtype == torch.float32
assert tensors["diffusion_model.linear.dora_scale"].dtype == torch.float32
assert tensors["diffusion_model.orphan.lora_down.weight"].dtype == torch.float16
pure_source = tmp_path / "pure_diff.safetensors"
save_file({
"diffusion_model.pure.diff": torch.ones((2, 2), dtype=torch.bfloat16),
"diffusion_model.pure_norm.diff": torch.ones(2, dtype=torch.bfloat16),
}, str(pure_source))
pure_output = resize.resize_lora_file(
str(pure_source), 1, None, None, "cpu", torch.float16, "pure_resized", verbose=False
)
pure_tensors = load_file(pure_output)
assert set(pure_tensors) == {"diffusion_model.pure.diff", "diffusion_model.pure_norm.diff"}
assert pure_tensors["diffusion_model.pure.diff"].dtype == torch.float16
assert pure_tensors["diffusion_model.pure_norm.diff"].dtype == torch.float32
def test_resize_normalizes_mochi_pair_to_preferred_output(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mochi.safetensors"
save_file({
"diffusion_model.layer.lora_A": torch.ones((2, 2)),
"diffusion_model.layer.lora_B": torch.ones((2, 2)),
"diffusion_model.layer.alpha": torch.tensor(2.0),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float32,
"mochi_normalized", verbose=False,
)
tensors = load_file(output)
assert set(tensors) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
"diffusion_model.layer.alpha",
}
def test_standard_multi_merge_normalizes_mochi_pair_to_preferred_output(
monkeypatch, tmp_path, modules
):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "multi_mochi_a.safetensors"
second = tmp_path / "multi_mochi_b.safetensors"
for path, value in ((first, 1.0), (second, 3.0)):
save_file({
"diffusion_model.layer.lora_A": torch.full((1, 2), value),
"diffusion_model.layer.lora_B": torch.full((2, 1), value),
"diffusion_model.layer.alpha": torch.tensor(1.0),
}, str(path))
output = merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float32, "multi_mochi_normalized", verbose=False,
)
tensors = load_file(output)
assert set(tensors) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
"diffusion_model.layer.alpha",
}
def test_standard_multi_merge_emits_weighted_direct_union(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "first.safetensors"
second = tmp_path / "second.safetensors"
base = tmp_path / "multi_base.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.matrix.weight": torch.zeros((2, 2)),
"model.diffusion_model.only_first.weight": torch.zeros(()),
"model.diffusion_model.scale": torch.zeros(()),
"model.diffusion_model.lowrank.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.only_first.diff": torch.tensor(4.0, dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(1.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.full((1, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.full((2, 1), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.bfloat16),
}, str(second))
output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged", base_model_path=str(base), verbose=False,
include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([-1.0, 0.0]))
torch.testing.assert_close(tensors["diffusion_model.matrix.diff"], torch.zeros((2, 2)))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.matrix.diff"].dtype == torch.float32
assert tensors["diffusion_model.only_first.diff"].item() == pytest.approx(8.0)
assert tensors["diffusion_model.only_first.diff"].dtype == torch.float16
assert tensors["diffusion_model.scale.diff"].item() == pytest.approx(2.0)
assert tensors["diffusion_model.lowrank.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.lowrank.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_B.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.alpha"].dtype == torch.float32
default_output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged_default", base_model_path=str(base), verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.matrix.diff" in default_tensors
assert "diffusion_model.scale.diff" in default_tensors
assert "diffusion_model.lowrank.lora_A.weight" in default_tensors
def test_multi_merge_rejects_direct_shape_mismatch(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "shape_first.safetensors"
second = tmp_path / "shape_second.safetensors"
save_file({"diffusion_model.bad.diff": torch.ones(2)}, str(first))
save_file({"diffusion_model.bad.diff": torch.ones(3)}, str(second))
with pytest.raises(ValueError, match="Direct LoRA shape mismatch"):
merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float16, "bad_shapes", verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("enhanced", [False, True])
def test_dare_variants_emit_direct_layers(monkeypatch, tmp_path, modules, enhanced):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / f"dare_first_{enhanced}.safetensors"
second = tmp_path / f"dare_second_{enhanced}.safetensors"
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.float32),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.bfloat16),
}, str(second))
if enhanced:
output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced", verbose=False, include_1d_diffs=True,
)
else:
output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare", verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0]))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_B.weight"].dtype == torch.float32
if enhanced:
default_output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced_default", verbose=False,
)
else:
default_output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare_default", verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.lowrank.lora_A.weight" in default_tensors
def test_merge_to_model_applies_generic_direct_keys_and_preserves_dtype(
monkeypatch, tmp_path, modules, caplog
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "base.safetensors"
adapter = tmp_path / "adapter.safetensors"
save_file({
"model.diffusion_model.block.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.norm.weight": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.block.bias": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.scale": torch.tensor(1.0, dtype=torch.bfloat16),
"model.diffusion_model.lin": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.bad.weight": torch.ones((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.untouched": torch.tensor(1e-8, dtype=torch.float32),
}, str(base))
save_file({
"diffusion_model.block.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.block.diff_b": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.bad.diff": torch.ones(3, dtype=torch.float32),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model",
verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0]))
torch.testing.assert_close(tensors["model.diffusion_model.block.bias"], torch.tensor([1.5, 2.0]))
assert tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5))
torch.testing.assert_close(tensors["model.diffusion_model.bad.weight"].float(), torch.ones((2, 2)))
assert "Shape mismatch" in caplog.text
assert tensors["model.diffusion_model.block.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.norm.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.block.bias"].dtype == torch.float32
assert tensors["model.diffusion_model.scale"].dtype == torch.float16
assert tensors["model.diffusion_model.lin"].dtype == torch.float32
assert tensors["model.diffusion_model.untouched"].dtype == torch.float32
default_output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model_default",
verbose=False,
)
default_tensors = load_file(default_output)
torch.testing.assert_close(
default_tensors["model.diffusion_model.norm.weight"].float(), torch.zeros(2)
)
torch.testing.assert_close(
default_tensors["model.diffusion_model.block.bias"].float(), torch.zeros(2)
)
assert default_tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(
default_tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5)
)
def _generic_params(output_filename, override_dtype, include_1d_diffs=False):
return {
"mismatch_mode": "skip",
"alignment_mode": "pad/crop",
"alpha": 0.5,
"beta": 0.0,
"gamma": 0.5,
"delta": 2.0,
"epsilon": 0.01,
"zeta": 0.0,
"seed": 0,
"output_filename": output_filename,
"save_dtype": "bf16",
"override_dtype": override_dtype,
"device": "cpu",
"dtype": torch.float32,
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"include_1d_diffs": include_1d_diffs,
}
def test_generic_two_and_three_mergers_enforce_direct_dtype_contract(
monkeypatch, tmp_path, generic_modules
):
generic, operations = generic_modules
paths = {}
for name, tensors in {
"a": {
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
},
"b": {
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 3.0, dtype=torch.bfloat16),
"diffusion_model.secondary_only.diff": torch.ones(1, dtype=torch.float32),
},
"c": {
"diffusion_model.norm.diff": torch.tensor([0.5, 1.0], dtype=torch.float32),
"diffusion_model.matrix.diff": torch.full((2, 2), 0.5, dtype=torch.bfloat16),
},
}.items():
path = tmp_path / f"{name}.safetensors"
save_file(tensors, str(path))
paths[name] = str(path)
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
two_params = _generic_params("generic_two", override_dtype=True)
two_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_params,
"loras",
)
two = load_file(str(tmp_path / "loras" / "generic_two.safetensors"))
torch.testing.assert_close(
two["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert two["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert two["diffusion_model.matrix.diff"].dtype == torch.bfloat16
assert "diffusion_model.secondary_only.diff" not in two
two_enabled_params = _generic_params(
"generic_two_enabled", override_dtype=True, include_1d_diffs=True
)
two_enabled_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_enabled_params,
"loras",
)
two_enabled = load_file(
str(tmp_path / "loras" / "generic_two_enabled.safetensors")
)
torch.testing.assert_close(
two_enabled["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0])
)
assert two_enabled["diffusion_model.norm.diff"].dtype == torch.float32
three_params = _generic_params("generic_three", override_dtype=False)
three_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_params,
"loras",
)
three = load_file(str(tmp_path / "loras" / "generic_three.safetensors"))
torch.testing.assert_close(
three["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert three["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert three["diffusion_model.matrix.diff"].dtype == torch.float32
three_enabled_params = _generic_params(
"generic_three_enabled", override_dtype=False, include_1d_diffs=True
)
three_enabled_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_enabled_params,
"loras",
)
three_enabled = load_file(
str(tmp_path / "loras" / "generic_three_enabled.safetensors")
)
torch.testing.assert_close(
three_enabled["diffusion_model.norm.diff"], torch.tensor([2.25, 3.5])
)
assert three_enabled["diffusion_model.norm.diff"].dtype == torch.float32
+437
View File
@@ -0,0 +1,437 @@
import importlib
import json
import logging
import struct
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_quant_guard_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def modules():
return {
"guard": _load_module("nodes.quantization_guard"),
"extract": _load_module("nodes.lora_extract_svd"),
"resize": _load_module("nodes.lora_resize"),
"multi": _load_module("nodes.lora_merger"),
"generic": _load_module("nodes.merger"),
"operations": _load_module("nodes.merger_ops"),
}
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _inspect_file(path, modules, label="input"):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
return modules["guard"].inspect_low_bit_input(handler, label, "Guard Test")
finally:
handler.__exit__(None, None, None)
def _round_trip_preserved(path, output, modules):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
with modules["resize"].IncrementalSafetensorsWriter(str(output)) as writer:
modules["guard"].write_preserved_tensor(writer, "low", handler)
finally:
handler.__exit__(None, None, None)
return load_file(str(output))["low"]
def _tensor_payload(path, key="low"):
with open(path, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
header = json.loads(source.read(header_size).decode("utf-8"))
start, end = header[key]["data_offsets"]
source.seek(8 + header_size + start)
return source.read(end - start)
@pytest.mark.parametrize("dtype", [torch.int8, torch.uint8])
def test_int8_uint8_are_detected_and_warn_once(tmp_path, modules, caplog, dtype):
path = tmp_path / f"{dtype}.safetensors"
save_file({"low": torch.arange(4, dtype=dtype), "normal": torch.ones(1)}, str(path))
with caplog.at_level(logging.WARNING):
detected = _inspect_file(path, modules, "threshold input")
assert detected == {"low"}
records = [record for record in caplog.records if "threshold input" in record.message]
assert len(records) == 1
assert "1 isolated low-bit tensor" in records[0].message
assert "low" in records[0].message
def test_exact_two_tensor_threshold(tmp_path, modules):
for count in range(4):
path = tmp_path / f"threshold_{count}.safetensors"
tensors = {f"low_{i}": torch.ones(2, dtype=torch.uint8) for i in range(count)}
tensors["normal"] = torch.ones(1)
save_file(tensors, str(path))
if count <= 2:
assert len(_inspect_file(path, modules)) == count
else:
with pytest.raises(ValueError, match="3 low-bit tensors.*casted or quantized"):
_inspect_file(path, modules)
@pytest.mark.parametrize(
("tensors", "metadata"),
[
({"layer.comfy_quant": torch.tensor([123], dtype=torch.uint8)}, None),
({"normal": torch.ones(1)}, {"_quantization_metadata": "{}"}),
({"scaled_fp8": torch.ones(2)}, None),
],
)
def test_recognized_comfy_quantization_hard_errors(tmp_path, modules, tensors, metadata):
path = tmp_path / f"quant_{len(list(tmp_path.iterdir()))}.safetensors"
save_file(tensors, str(path), metadata=metadata)
with pytest.raises(ValueError, match="ComfyUI quantization metadata.*unsupported"):
_inspect_file(path, modules, "recognized input")
def test_supported_fp8_storage_round_trip(tmp_path, modules):
candidates = [
dtype
for name in (
"float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz",
"float8_e8m0fnu",
)
if (dtype := getattr(torch, name, None)) is not None
]
supported = []
for index, dtype in enumerate(candidates):
path = tmp_path / f"low_float_{index}.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
continue
supported.append(dtype)
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / f"preserved_{index}.safetensors", modules)
assert preserved.dtype == dtype
torch.testing.assert_close(preserved.float(), loaded.float())
if candidates and not supported:
pytest.skip("Installed safetensors cannot write any installed torch FP8 dtype")
def test_fp4_storage_when_supported(tmp_path, modules):
dtype = getattr(torch, "float4_e2m1fn_x2", None)
if dtype is None:
pytest.skip("Installed PyTorch has no torch.float4_e2m1fn_x2")
path = tmp_path / "fp4.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
pytest.skip("Installed safetensors cannot write torch.float4_e2m1fn_x2")
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / "preserved_fp4.safetensors", modules)
assert preserved.dtype == dtype
assert _tensor_payload(path) == _tensor_payload(tmp_path / "preserved_fp4.safetensors")
def test_realistic_packed_uint8_comfy_sidecar_is_rejected(tmp_path, modules):
path = tmp_path / "packed_comfy.safetensors"
save_file({
"layer.weight": torch.arange(8, dtype=torch.uint8).reshape(2, 4),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
"layer.weight_scale": torch.ones(2),
}, str(path))
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
_inspect_file(path, modules)
def test_resize_preserves_complete_low_bit_layer_and_resizes_normal(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "resize_quant.safetensors"
guarded_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"diffusion_model.guarded.lora_down.weight": guarded_down,
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.guarded.alpha": torch.tensor(2.0),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.normal.lora_up.weight": torch.ones((2, 2)),
"isolated_aux": torch.tensor([7], dtype=torch.int8),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resize_guarded", verbose=False
)
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], guarded_down)
assert result["diffusion_model.guarded.lora_A.weight"].dtype == torch.uint8
assert result["diffusion_model.guarded.lora_B.weight"].dtype == torch.float32
assert result["diffusion_model.guarded.dora_scale"].dtype == torch.float32
assert result["diffusion_model.normal.lora_A.weight"].shape[0] == 1
assert result["isolated_aux"].dtype == torch.int8
def test_extract_excludes_low_bit_candidate_and_extracts_float_layer(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
model_a = tmp_path / "extract_a.safetensors"
model_b = tmp_path / "extract_b.safetensors"
output = tmp_path / "extracted.safetensors"
save_file({
"model.diffusion_model.normal.weight": torch.eye(2),
"model.diffusion_model.quant.weight": torch.ones((2, 2), dtype=torch.uint8),
}, str(model_a))
save_file({
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
"model.diffusion_model.quant.weight": torch.zeros((2, 2), dtype=torch.uint8),
}, str(model_b))
extract.extract_lora_from_files(
str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
force_clear_cache=False,
)
result = load_file(str(output))
assert any("normal" in key for key in result)
assert not any("quant" in key for key in result)
def test_extract_bias_diff_round_trips_through_merge_to_model(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
resize = modules["resize"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
_patch_output(monkeypatch, resize, tmp_path)
finetuned = tmp_path / "bias_finetuned.safetensors"
base = tmp_path / "bias_base.safetensors"
adapter = tmp_path / "bias_adapter.safetensors"
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.25, -0.5]),
}, str(finetuned))
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.0, 0.5]),
}, str(base))
extract.extract_lora_from_files(
str(finetuned), str(base), "fixed", 1, 1, "cpu", "fp16", str(adapter),
force_clear_cache=False, include_1d_diffs=True,
)
extracted = load_file(str(adapter))
assert set(extracted) == {"diffusion_model.layer.diff_b"}
assert extracted["diffusion_model.layer.diff_b"].dtype == torch.float32
merged_path = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"bias_round_trip", verbose=False, include_1d_diffs=True,
)
merged = load_file(merged_path)
torch.testing.assert_close(
merged["model.diffusion_model.layer.bias"],
load_file(str(finetuned))["model.diffusion_model.layer.bias"],
)
def _run_multi_variant(module, variant, paths, output_name):
common = (paths, [1.0, 1.0])
if variant == "standard":
return module.merge_multi_loras(
*common, "weighted_sum", "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
if variant == "dare":
return module.merge_multi_loras_dare(
*common, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
return module.merge_multi_loras_dare_enhanced(
*common, 1.0, 1.0, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("variant", ["standard", "dare", "enhanced"])
def test_multi_variants_preserve_earliest_complete_affected_layer(
monkeypatch, tmp_path, modules, variant
):
multi = modules["multi"]
_patch_output(monkeypatch, multi, tmp_path)
first = tmp_path / f"{variant}_first.safetensors"
second = tmp_path / f"{variant}_second.safetensors"
first_down = torch.tensor([[1.0, 2.0]])
save_file({
"diffusion_model.guarded.lora_down.weight": first_down,
"diffusion_model.guarded.lora_up.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(first))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 2]], dtype=torch.uint8),
"diffusion_model.guarded.lora_up.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([9.0, 9.0]),
"diffusion_model.normal.diff": torch.full((2, 2), 3.0),
}, str(second))
output = _run_multi_variant(multi, variant, [str(first), str(second)], f"guard_{variant}")
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], first_down)
torch.testing.assert_close(
result["diffusion_model.guarded.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
torch.testing.assert_close(
result["diffusion_model.guarded.dora_scale"], torch.tensor([0.125, 0.25])
)
assert "diffusion_model.normal.diff" in result
def test_hard_error_happens_before_resize_output_creation(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "three_low.safetensors"
save_file({f"low_{i}": torch.ones(1, dtype=torch.uint8) for i in range(3)}, str(source))
output = tmp_path / "must_not_exist.safetensors"
with pytest.raises(ValueError, match="3 low-bit tensors"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def test_recognized_quantization_creates_no_resize_output(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "recognized_quant.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
}, str(source))
output = tmp_path / "recognized_must_not_exist.safetensors"
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def _generic_params(output_filename):
return {
"mismatch_mode": "skip", "alignment_mode": "pad/crop", "alpha": 0.5,
"beta": 0.0, "gamma": 0.5, "delta": 2.0, "epsilon": 0.01, "zeta": 0.0,
"seed": 0, "output_filename": output_filename, "save_dtype": "bf16",
"override_dtype": True, "device": "cpu", "dtype": torch.float32,
"exclude_patterns": "", "discard_patterns": "", "glob_patterns": False,
"lazy_load": True, "force_clear_cache": False, "include_1d_diffs": True,
}
def test_generic_merge_preserves_model_a_when_either_source_key_is_low_bit(
monkeypatch, tmp_path, modules
):
generic = modules["generic"]
operations = modules["operations"]
model_a = tmp_path / "generic_guard_a.safetensors"
model_b = tmp_path / "generic_guard_b.safetensors"
a_tensors = {
"a_low": torch.tensor([[1, 2]], dtype=torch.uint8),
"b_low": torch.tensor([[0.125, 0.25]], dtype=torch.float32),
"normal": torch.ones((1, 2)),
}
save_file(a_tensors, str(model_a))
save_file({
"a_low": torch.tensor([[9.0, 9.0]]),
"b_low": torch.tensor([[3, 4]], dtype=torch.int8),
"normal": torch.full((1, 2), 3.0),
}, str(model_b))
paths = {"a": str(model_a), "b": str(model_b)}
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
params = _generic_params("generic_guarded")
params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "loras",
)
result = load_file(
str(tmp_path / "loras" / "generic_guarded.safetensors")
)
torch.testing.assert_close(result["a_low"], a_tensors["a_low"])
torch.testing.assert_close(result["b_low"], a_tensors["b_low"])
assert result["a_low"].dtype == torch.uint8
assert result["b_low"].dtype == torch.float32
torch.testing.assert_close(result["normal"].float(), torch.full((1, 2), 2.0))
def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "merge_guard_base.safetensors"
adapter = tmp_path / "merge_guard_adapter.safetensors"
base_quant = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"model.diffusion_model.base_quant.weight": base_quant,
"model.diffusion_model.guarded.weight": torch.zeros((2, 2)),
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 1]], dtype=torch.int8),
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 1)),
"diffusion_model.guarded.diff": torch.full((2, 2), 5.0),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(adapter))
output, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16, "merge_guarded",
verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(output)
torch.testing.assert_close(result["model.diffusion_model.base_quant.weight"], base_quant)
assert result["model.diffusion_model.base_quant.weight"].dtype == torch.uint8
torch.testing.assert_close(
result["model.diffusion_model.guarded.weight"].float(), torch.zeros((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.normal.weight"].float(), torch.ones((2, 2))
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 1/2")
assert "[GUARDED] (1)" in report
assert "diffusion_model.guarded.lora_down.weight" in report
assert "torch.int8" in report
assert "model.diffusion_model.base_quant.weight" in report
+282
View File
@@ -0,0 +1,282 @@
import importlib
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,
}
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 [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]
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",
]
def test_documentation_mode_has_three_separate_outputs(analysis):
result = analysis.DiffusionModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")
assert len(result.result) == 3
assert "No comparison performed" in result[0]
assert "No CWB analysis performed" in result[1]
assert result[2].startswith("# Two-Model Similarity Analysis")
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):
path_a = tmp_path / "a.safetensors"
path_b = tmp_path / "b.safetensors"
save_file({
"model.blocks.0.weight": 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({
"model.blocks.0.weight": 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)
comparison, cwb_report = analysis.ModelAnalysisLogic.execute(
"a", "b", "diffusion_models", _params()
)
assert "Comparable floating tensors: 1" in comparison
assert "mae: 1" in comparison
assert "mse: 2" in comparison
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} == {"model.blocks.0.weight", "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 = 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 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)
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)