Compare commits

...
Author SHA1 Message Date
silveroxides a7c198a439 Bump version to 0.9.7 2026-11-10 21:49:40 +01:00
silveroxides ee3d998a5d Support w6a8_int8 quantization decoding and shape unpacking 2026-11-10 17:29:49 +01:00
silveroxides f13eeda11b Preserve isolated low-bit tensors across all generic merger model types 2026-10-31 23:23:19 +01:00
silveroxides 9ebe34565c Remove arbitrary low-bit tensor count cap and honor discard before preservation 2026-10-31 23:16:02 +01:00
silveroxides 6bbc566a52 Support bare LoRA A/B suffixes without .weight in MiniMax H3 converter 2026-10-31 19:57:42 +01:00
silveroxides 9cda62199c Add include_mode and standard UEL work-unit streaming to PruneKeys nodes 2026-10-22 15:40:40 +02:00
silveroxides cc0ac00128 Move CWBModelGraft to standalone module and add full documentation 2026-10-22 03:40:36 +02:00
silveroxides b2d9a7bd00 Bump version to 0.9.6 2026-10-22 03:29:39 +02:00
silveroxides 8286c2d5d9 Add standalone CWBModelGraft node 2026-10-21 08:34:24 +02:00
silveroxides 88e0e3cdcb Add LoRAModifyAlpha and LoRADenormalizeAlpha nodes 2026-10-19 03:33:42 +02:00
silveroxides 248fdff2a4 Add CWBEmbeddingStack and CWBEmbeddingStackMerger nodes 2026-10-18 18:30:50 +02:00
silveroxides 576c393342 Clarify CWB Custom Configuration tooltips with layman explanations 2026-10-07 07:20:46 +02:00
silveroxides eea3496002 Bump version to 0.9.5 2026-10-06 18:37:43 +02:00
silveroxides f9aaa718c0 Exclude downloader modules via .comfyignore with conditional startup registration 2026-10-06 18:36:34 +02:00
silveroxides eab8b20ef1 Clarify merge mode documentation and parameter strength 2026-10-05 23:02:43 +02:00
silveroxides a08cb23b6c Bump version to 0.9.4 2026-10-05 17:00:22 +02:00
silveroxides 2ee3f30c32 Add MiniMax H3 AdaLN unfold node 2026-10-05 14:54:00 +02:00
silveroxides f383a501a5 Simplify MiniMax H3 fold tooltips and description 2026-10-03 02:25:37 +02:00
silveroxides 14b60d3d9f Bump version to 0.9.3 2026-10-03 02:15:23 +02:00
silveroxides 4f600761dd Make force_clear_cache default False across nodes 2026-10-03 02:14:14 +02:00
silveroxides ae2858bf6b Add MiniMax H3 AdaLN fold node for diffusion models 2026-10-03 02:07:32 +02:00
silveroxides fcc6e5cc4a Support lora.down and lora.up factor suffixes 2026-10-02 21:08:16 +02:00
silveroxides 513d4c60a1 Make test runner quiet with failure-only ticker 2026-10-02 02:02:03 +02:00
silveroxides 2974f394c7 Bump version to 0.9.2 2026-10-02 01:54:08 +02:00
silveroxides 906ef9bf00 Add MiniMax H3 Diffusers to ComfyUI LoRA converter node 2026-10-02 00:02:54 +02:00
silveroxides 5a3b18276c Bump version to 0.9.1 2026-09-28 13:20:43 +02:00
silveroxides 2de5354527 fix: ignore LoRA dtype when saving merged model 2026-09-25 00:45:45 +02:00
silveroxides bbb8dad79d Bump version to 0.9.0 2026-09-23 03:13:54 +02:00
silveroxides 66f795de97 Support mixed-prefix quantized Consensus merges 2026-09-23 02:42:27 +02:00
silveroxides 1861d75aab Match diffusion model layers across prefixes 2026-09-23 02:36:18 +02:00
silveroxides 57bd63e1eb Support quantized diffusion model inputs 2026-09-23 01:53:25 +02:00
silveroxides 7a606f8a02 feat: Add Markdown and CSV analysis reports 2026-09-14 11:10:59 +02:00
silveroxides e8eb2a0868 feat: Analyze LoRA effects on a base model 2026-09-13 14:56:17 +02:00
silveroxides 6339939726 feat: Add per-layer parameter configuration 2026-09-13 06:55:32 +02:00
silveroxides 9e054e6105 docs: Record per-layer parameter plan 2026-09-13 06:39:19 +02:00
silveroxides f691ddc584 feat: Add include mode to layer filters 2026-09-11 03:50:49 +02:00
silveroxides 864ed422d8 chore(.gitignore) complete .gitignore local untracked files specifications 2026-09-08 08:51:10 +02:00
silveroxides 73cefa409d lora(resize) add exclusion and discard filters to resize nodes 2026-09-08 08:39:03 +02:00
silveroxides 0a808b6865 Bump version from 0.7.2 to 0.8.0 2026-09-04 23:31:52 +02:00
silveroxides ae7f742e08 feat(lora): add factorized Delta CWB output 2026-09-04 23:20:46 +02:00
silveroxides af32ebea45 feat(lora): add full-delta CWB merge 2026-09-04 20:20:50 +02:00
silveroxides e421b49575 feat(lora): map alpha normalization keys 2026-09-04 16:34:51 +02:00
silveroxides 7bed3fdbba fix(lora): clarify Lodestone controls 2026-09-04 15:59:25 +02:00
silveroxides ebb29393aa feat(lora): add Lodestone delta mergers 2026-09-04 15:55:42 +02:00
silveroxides 9ef1b49b41 Bump version to 0.7.2 2026-08-25 00:48:59 +02:00
silveroxides 663109aa5c feat(cwb): add focused merge presets 2026-08-24 22:55:06 +02:00
silveroxides 016437d844 Bump version to 0.7.1 2026-08-24 17:48:20 +02:00
silveroxides 5a90fa8a04 feat(cwb): save merge metadata 2026-08-24 17:45:22 +02:00
silveroxides e7e49a5f7e feat(dtype): map dtypes from reference 2026-08-24 16:24:09 +02:00
silveroxides 066c7a34db feat(embeddings): add CWB coalescing 2026-08-20 19:34:21 +02:00
silveroxides d9467c4bf2 chore(release): bump version to 0.7.0 2026-08-18 12:16:33 +02:00
silveroxides 6270b8c7a2 feat(lora): add alpha normalization node 2026-08-18 12:16:21 +02:00
silveroxides 216080cb1a test: add repository test selector 2026-08-18 12:16:11 +02:00
silveroxides f0aee7ec0c Bump version from 0.5.4 to 0.6.0 2026-08-14 20:57:43 +02:00
silveroxides c25d6c4990 feat(cwb): add LoRA presets and diagnostics 2026-08-14 20:53:42 +02:00
silveroxides ece52e651c feat(analysis): add filtered API workflow runner 2026-08-14 19:04:13 +02:00
silveroxides 349e2155cd fix(paths): standardize model artifact outputs 2026-08-14 19:04:05 +02:00
silveroxides f5839f706c docs: define UEL and LoRA merge invariants 2026-08-14 08:44:01 +02:00
silveroxides f93ca54c81 fix(io): stream UEL and normalize LoRA alpha 2026-08-14 08:43:53 +02:00
silveroxides ed8c331ce2 perf(extract): stream adaptive SVD inputs 2026-08-14 07:49:00 +02:00
silveroxides 0d79c8d4f4 docs(audit): record UEL loading gaps 2026-08-13 16:36:18 +02:00
silveroxides 16e493cfdf perf(lora): bound knee SVD probes 2026-08-13 16:36:13 +02:00
silveroxides f76195f249 feat(models): add dtype conversion node 2026-08-13 07:37:36 +02:00
silveroxides e2c11f2290 test(nodes): require static tooltips 2026-08-13 07:27:41 +02:00
silveroxides c84ff102cb fix(lora): preserve per-role dtypes 2026-08-13 07:27:27 +02:00
silveroxides 23fa2ad485 docs(nodes): add schema tooltips 2026-08-13 07:26:37 +02:00
silveroxides 6457c41a5a fix(lora): stream mixed-rank multi-merge 2026-08-09 16:49:31 +02:00
silveroxides 91a38ca4e5 Bump version from 0.5.3 to 0.5.4 2026-08-08 02:34:20 +02:00
silveroxides 725f4a0e68 perf(lora): rebuild resize pipeline 2026-08-08 02:34:08 +02:00
Silver 450adbd9c9 Bump version from 0.5.2 to 0.5.3 2026-08-04 15:55:42 +02:00
Silver 130e6adede Merge pull request #20 from silveroxides/feature/model-similarity-analysis
Feature/model similarity analysis
2026-08-04 15:55:22 +02:00
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
94 changed files with 34839 additions and 1023 deletions
+4
View File
@@ -0,0 +1,4 @@
tests/
nodes/downloader_nodes.py
nodes/downloader_utils.py
+11
View File
@@ -176,3 +176,14 @@ cython_debug/
# PyPI configuration file
.pypirc
# local untracked directories
tools_local/
references_local/
plans_local/
old_files_local/
audit_local/
old_sessions/
# local untracked files
/*.jsonl
+81 -1
View File
@@ -1,2 +1,82 @@
# 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)
- **Diffusion Model Dtype Conversion** – Stream models to fp32, fp16, or bf16 while preserving excluded tensor dtypes
## Layer filters
Nodes with `exclude_patterns` or `skip_patterns` have an appended `include_mode` toggle
(off by default). Turn it on to process only layers matching that same field,
using the node's existing regex or glob syntax. An empty include filter selects
nothing. Nonmatches follow the node's usual exclusion behavior: preserve the
source/anchor layer for pattern-exclusion mergers, resize, or conversion; omit
it from analysis or extraction. LoRA Merge to Model retains its existing skip
semantics: nonmatching base tensors are omitted from the saved model, while
guarded low-bit base tensors are always preserved. Its filter uses regex only.
`discard_patterns`, where available, still takes precedence.
## Example Workflows
### Per-layer parameters
Connect **Layer Parameter Configuration** to `layer_parameters` on standard
two/three-input mergers, extraction nodes, or LoRA resize nodes. For example,
fixed extraction accepts `(blocks\.4[589]\.attn\.qkv_proj) a:64 b:32 c:0.99 d:0`.
Short aliases and full names, with whitespace or comma separators, are supported.
The configurator's documentation output lists every method's bindings and bounds.
Only explicitly assigned values override the receiving node; other settings and
existing filters still apply. Overlapping rules and rules matching no target
layer fail before streaming or writing. LoRA patterns target the existing
normalized logical layer names, not individual factor/alpha suffixes.
<p align="center">
<img src="assets/GetMetaAndkeys.png" width="400" alt="Get Meta and Keys">
<br>
<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.
+104 -4
View File
@@ -1,6 +1,7 @@
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced
from .nodes.metakeys import (
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys
@@ -19,15 +20,68 @@ from .nodes.merger import (
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger
)
from .nodes.lora_extract import (
LoRASaveFromFile, LoRACheckpointSaveFromFile
from .nodes.consensus_merger import CWB_MERGER_NODES
from .nodes.cwb_model_graft import CWB_GRAFT_NODES
from .nodes.cwb_delta_lora_merger import DELTA_CWB_LORA_NODES
from .nodes.lodestone_merger import LODESTONE_MERGER_NODES
from .nodes.model_analysis import MODEL_ANALYSIS_NODES
from .nodes.lora_extract_svd import (
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius
)
from .nodes.dora_extract_wd import (
DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
DoRAExtractKnee, DoRAExtractFrobenius
)
from .nodes.dora_learned_wd import (
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius
)
from .nodes.text_encoder_extract import (
TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio,
TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee,
TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed,
TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile,
TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius,
)
from .nodes.lora_resize import (
LoRANormalizeAlpha,
LoRAModifyAlpha,
LoRADenormalizeAlpha,
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
LoRAMergeToModel
)
try:
from .nodes.downloader_nodes import (
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader
)
DOWNLOADER_NODES: list[type[io.ComfyNode]] = [
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
]
except ImportError:
DOWNLOADER_NODES = []
from .nodes.model_info_nodes import (
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader
)
from .nodes.lora_rename import AnimaLoraRename
from .nodes.minimax_h3_lora_convert import MiniMaxH3DiffusersLoRAConvert
from .nodes.minimax_h3_fold import MiniMaxH3FoldAdaLN, MiniMaxH3UnfoldAdaLN
from .nodes.dtype_conversion import DiffusionModelDtypeConversion
from .nodes.layer_parameters import LayerParameterConfiguration
from .nodes.lora_model_analysis import LoRAOnModelAnalysis
class ModelUtilsExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
LayerParameterConfiguration,
# MetaKeys
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys,
@@ -43,10 +97,56 @@ class ModelUtilsExtension(ComfyExtension):
# Three-Model Mergers
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger,
# Dedicated CWB Mergers
*CWB_MERGER_NODES,
*CWB_GRAFT_NODES,
# Dedicated Delta CWB LoRA Mergers
*DELTA_CWB_LORA_NODES,
# Dedicated Lodestone LoRA Mergers
*LODESTONE_MERGER_NODES,
# Two-Model Analysis
*MODEL_ANALYSIS_NODES,
LoRAOnModelAnalysis,
# LoRA Extraction
LoRASaveFromFile, LoRACheckpointSaveFromFile,
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius,
# DoRA Extraction
DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
DoRAExtractKnee, DoRAExtractFrobenius,
# Learned DoRA Extraction
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius,
# Text Encoder LoRA/DoRA Extraction
TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio,
TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee,
TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed,
TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile,
TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius,
# LoRA Resize
LoRANormalizeAlpha,
LoRAModifyAlpha,
LoRADenormalizeAlpha,
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
# LoRA Multi-Merge
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
# LoRA Merge To Model
LoRAMergeToModel,
# LoRA Utilities
AnimaLoraRename,
MiniMaxH3DiffusersLoRAConvert,
# Dtype Conversion & Folding
DiffusionModelDtypeConversion,
MiniMaxH3FoldAdaLN,
MiniMaxH3UnfoldAdaLN,
# Downloaders (available in manual git/source installations)
*DOWNLOADER_NODES,
# Info Loaders
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader,
]
async def comfy_entrypoint() -> ModelUtilsExtension:
return ModelUtilsExtension()
return ModelUtilsExtension()
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

+323
View File
@@ -0,0 +1,323 @@
# ComfyUI-ModelUtils - Complete Project History
> [Implementation Invariants](IMPLEMENTATION_INVARIANTS.md) is the authoritative contract for UEL tensor I/O, lifecycle management, atomic saving, and LoRA alpha normalization. All new node implementations and changes must follow it. This document is historical and does not override that contract.
## Version Timeline
| Version | Commits | Key Features |
|---------|---------|--------------|
| 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
+93
View File
@@ -0,0 +1,93 @@
# Implementation Invariants
This file is the authoritative engineering contract for model and LoRA tensor I/O and merge behavior in ComfyUI-ModelUtils. New implementations and changes to existing implementations MUST follow these rules. Historical code, audit notes, and examples do not override this contract.
## UnifiedEfficientLoader I/O
1. Production node code MUST use `unifiedefficientloader` for safetensors tensor I/O. Direct safetensors imports and synchronous `get_tensor()` calls are prohibited.
2. Header-only operations MUST use UEL header, key, shape, dtype, and metadata APIs without materializing tensors.
3. Tensor processing MUST use `async_stream` with bounded buffering. The repository default is `batch_size=1` and `prefetch_batches=1` unless a measured benchmark and a documented memory bound justify another value.
4. Multi-input operations MUST maintain an async cursor for each input and construct logical work units. A work unit contains only the tensors required to process one output tensor, LoRA factor pair, or other indivisible operation.
5. Pinned-memory transfer MUST be enabled when the selected processing device benefits from it. CPU-only processing MUST avoid unnecessary pinned copies.
6. Implementations MUST NOT accumulate a complete model, LoRA, or unbounded output batch in RAM. The next work unit may be prefetched while the current unit is processed, but the prefetch bound MUST remain explicit.
7. Every yielded source tensor occurrence MUST be marked processed after its logical work unit completes or aborts. All strong references to source, normalized, padded, device, and output intermediates MUST be removed before advancing to the next work unit.
8. Every async stream and handler MUST close in `finally` or through an equivalent deterministic context manager. Failure paths are subject to the same release rules as successful paths.
9. Saving MUST use UEL incremental writing. Outputs MUST be submitted as they are completed rather than retained in an in-memory state dictionary.
10. File replacement MUST be atomic: write to a temporary sibling file, finalize successfully, then replace the destination. A failure MUST preserve any existing destination and remove the incomplete temporary file.
11. Unsupported low-bit tensors MUST use the shared quantization-preservation path when byte-exact copying is valid. Code MUST NOT silently cast, decode, merge, or relabel unsupported low-bit storage.
The shared implementations in `nodes/uel_io.py`, `nodes/extraction_stream.py`, and `nodes/quantization_guard.py` are the required starting points. New node-local stream or writer abstractions require a demonstrated capability that the shared implementation cannot provide.
## Persistent Artifact Paths and Outputs
1. Persistent model artifacts MUST use `nodes/artifact_paths.py` to construct
their write path from `folder_paths.models_dir` and an explicit canonical
category: `checkpoints`, `diffusion_models`, `text_encoders`, `loras`, or
`embeddings`.
2. Registered search paths, additional paths, and legacy aliases such as
`unet` and `clip` MUST NOT become write roots based on their list position.
3. Nested output names MUST be preserved, exactly one `.safetensors` extension
MUST be applied, and path traversal outside the category root MUST fail.
4. Public filename outputs MUST return the category-relative name with forward
slashes. They MUST NOT return an absolute filesystem path.
5. Loader-compatible filename outputs MUST use `io.AnyType`. Reports and
documentation MUST remain separate string outputs.
6. Persistent model artifacts MUST NOT be written beneath ComfyUI's general
output directory.
## LoRA Alpha Normalization for Merge Operations
1. Every operation that combines a LoRA with another LoRA or applies a LoRA to a model MUST account for per-layer alpha before weighting, padding, concatenation, consensus processing, DARE/TIES processing, or delta application.
2. For a valid factor pair with rank `r`, normalization is performed mathematically as:
`normalized_up = up * (alpha / r)`
`normalized_down = down`
Scaling one factor preserves the represented delta while avoiding unnecessary work and memory.
3. A layer without alpha has an implicit scale of `1.0`. Inputs with and without alpha may be merged only after alpha-bearing pairs have been normalized to this common representation.
4. Normalized merge outputs MUST NOT contain alpha tensors. Alpha MUST NOT be invented for inputs or output formats that do not provide it.
5. Alpha is a scalar control value. Its dtype MUST NOT determine factor computation dtype, merged tensor dtype, or save dtype.
6. Alpha association MUST come from the repository's LoRA pair parser, not from an assumed key spelling. PEFT, Diffusers, ComfyUI, and supported mixed naming forms are governed by the parsed logical pair.
7. Normalization MUST occur before rank padding or other shape alignment. Padding an unnormalized pair and applying its original alpha against the padded rank changes the represented delta.
8. The unscaled factor reference MUST be replaced immediately when normalization creates a new tensor. Both source factors, the alpha scalar, normalized factors, padded factors, and device copies MUST be released when the logical layer completes.
9. A low-bit alpha scalar is valid when both factors use supported floating storage. If an alpha-bearing down or up factor uses unsupported low-bit storage, the operation MUST fail with an actionable error because byte-exact preservation cannot also produce an alpha-free normalized result.
10. Format metadata may describe that normalization occurred, but metadata MUST NOT be used as a substitute for applying the mathematical transformation.
### CWB Mixed-Rank LoRA Alignment
1. Genuine LoRA rank validity MUST come from each factor pair's original shape
or explicit source metadata. Tensor values MUST NOT be inspected to decide
whether a component is padding because a genuine component may be zero.
2. Structural rank padding MUST NOT participate in CWB alignment, consensus,
similarity weighting, fallback behavior, or norm rescaling.
3. Similarity alignment between rank `R` and rank `r` genuine pairs MUST compare
the full rectangular `R x r` component space. A smaller source component may
match any component in the maximum-rank reference.
4. Index alignment MUST pair only shared genuine indices. Higher reference
indices remain singleton groups unless another genuine source contributes.
5. An explicit zero inserted by `mismatch_mode=zeros` is a semantic contributor
and MUST remain distinct from structural rank absence.
6. Down rows and corresponding up columns MUST share validity, alignment, sign
correction, contributor membership, and output position.
7. Output rank MUST remain the largest genuine input rank. Equal maximum-rank
ties retain source order.
The shared implementation in `nodes/lora_alpha.py` is mandatory for merge paths. A format-specific implementation may only replace it when the represented delta differs mathematically and that difference is documented and tested.
## Required Verification
Any implementation that reads, writes, merges, resizes, extracts, or applies model or LoRA tensors MUST include tests appropriate to its scope. Merge changes involving LoRA factors MUST cover:
- alpha present and absent in the same operation;
- non-neutral alpha where `alpha / rank` is not `1.0`;
- normalization before padding or concatenation;
- alpha-free output;
- mixed factor ranks and dtypes;
- full rectangular CWB matching without structural-padding contributors;
- genuine zero-valued components versus explicit mismatch-zero contributors;
- low-bit scalar alpha acceptance and unsupported low-bit factor rejection;
- stream closure, `mark_processed`, and reference release on success and failure;
- atomic destination preservation on writer failure.
Repository guards MUST continue to reject production `get_tensor()` calls and direct safetensors imports. The complete test suite, changed-file lint, compilation, and `git diff --check` are required before commit.
+348
View File
@@ -0,0 +1,348 @@
# Consensus-Weighted Blending Mergers
The dedicated CWB nodes merge two or three safetensors files using
Consensus-Weighted Blending. They calculate an element-wise mean or median
consensus, compare each aligned source vector to that consensus, and normalize
similarity-derived weights before producing the output vector.
The broad mathematical and integration reference is maintained in
[`references/Spec-CWB`](../references/Spec-CWB/README.md). This document
describes the behavior actually exposed by this repository's offline merger
nodes. The specification covers additional contexts, including online
conditioning and spatial fusion, which are not controls of these nodes.
## Node families
Two-input and three-input variants are provided for checkpoints, standalone
diffusion models, text encoders, LoRAs, and embeddings. LoRA also provides a
Multi-Merge variant for 2 to 8 equal-prior inputs.
## Embedding coalescing
`CWB Embedding Self-Coalesce` reduces one embedding by replacing mutually
nearest body-vector pairs that meet both cosine and normalized-position limits.
`CWB Embedding Multi-Merge` accepts 2 to 8 embeddings, performs normal CWB
alignment first, then applies the same reduction. A target row count of zero
continues until no eligible pair remains.
Vision Boundary Embeddings preserves encoded vision-start and vision-end rows
unchanged and applies CWB only to their interior. Normal vision inputs validate
their outer rows before writing. Legacy Boundary Search trims accidental outer
template rows by matching a selected known-clean visual embedding; an absent,
ambiguous, or low-similarity boundary pair is an error rather than a guessed
crop. Disable Vision Boundary Embeddings for open textual-inversion tensors.
Model A supplies output metadata and anchors shared-layer shape, naming, and
mismatch behavior. The output key set is the union of all inputs. A valid
tensor found only in a secondary input is copied unchanged. When Model A lacks
a tensor that is present in multiple secondary inputs, those available sources
are CWB-merged without inserting a zero contribution for Model A. Embedding
mergers additionally use the longest compatible first dimension as their
alignment reference.
## Presets and custom controls
Presets are separated by merge type. Dense tensors, embeddings, and LoRA
factor pairs never share one menu. Every named preset supplies every CWB
setting and cannot inherit hidden widget state.
Manual controls live in the separate CWB Custom Configuration node. Connecting
its `CWB_CONFIG` output completely overrides the selected preset. Leaving it
unconnected uses the named preset.
Preset names expose categorical and boolean behavior. `_idx` and `_sim`
identify alignment, `_medn` and `_mean` identify consensus, and `_rn`, `_dsc`,
`_softcb`, and `_pcp` identify enabled boolean transformations. Dense presets
omit alignment tokens because dense model coordinates are fixed.
Similarity alignment greedily matches source rows to the reference by cosine
similarity. Index alignment keeps absolute row positions. Position weight can
bias similarity matches toward nearby normalized positions. Preserve Common
Prefix copies the longest numerically identical leading row span directly from
the first input.
Inputs have equal prior weight. The similarity calculation supplies the
effective contribution for each aligned vector. If every candidate is rejected
or receives zero weight, CWB falls back to equal weighting.
### Weighting sequence
For each aligned vector group, the merger performs these operations in order:
1. Compute an element-wise mean or median consensus.
2. Measure each vector's cosine similarity to that consensus.
3. Reject similarities below `similarity_threshold`.
4. Optionally remap unequal accepted similarities with Dynamic Similarity
Contrast.
5. Apply `similarity^power_alpha` and, when enabled, the diversity bandpass.
6. Normalize the resulting contributor weights and calculate their weighted
sum. If no usable weight remains, use equal contributor weights.
7. Optionally apply `rescale_norm`.
8. Apply `global_scale`.
### Complete control reference
- **Execution Mode**: `MERGE` writes an output. `DOCUMENTATION ONLY` returns
this document without opening model inputs.
- **Model A**: primary contributor and preservation anchor. It supplies output
metadata and anchors shared names and shapes. The output key set remains the
union of all inputs.
- **Model B / Model C**: additional equal-prior contributors. CWB derives their
effective per-vector influence from consensus similarity; these nodes do not
expose model-level weights.
- **CWB Preset**: complete use-case settings selected from the node's own dense,
embedding, or LoRA registry. A connected CWB Config overrides it completely.
- **Consensus Type**: `mean` uses the arithmetic center; `median` uses the
element-wise median and is less sensitive to coordinate outliers.
- **Alignment Method**: `index` pairs first-axis vectors by position;
`similarity` uses greedy one-to-one cosine matching. Row reordering is
enabled only for embeddings and paired LoRA rank components. Ordinary
checkpoint, diffusion-model, and text-encoder tensors retain fixed
coordinates.
- **Alignment Threshold**: minimum original cosine score for accepting a
greedy row match. It does not filter already-paired contributors during CWB
weighting; that is `similarity_threshold`.
- **Similarity Threshold**: minimum cosine similarity between a participating
vector and the group consensus. Rejected vectors receive zero weight. If all
are rejected, the implementation falls back to equal weights.
- **Power Alpha**: exponent applied to accepted non-negative similarities.
Values above 1 increasingly favor contributors closest to consensus. Zero
makes every accepted contributor's similarity term equal to one.
- **Diversity Beta**: exponent for the diversity bandpass. Values above zero
suppress vectors extremely close to consensus relative to more varied
accepted vectors. Zero disables this term.
- **Rescale Norm**: after blending, preserve the merged direction but replace
its L2 magnitude with the mean L2 magnitude of the participating vectors:
`merged = merged / ||merged|| × mean(||source_i||)`
A zero merged vector remains unchanged. This is performed per aligned row,
not once per complete tensor. For LoRA, it is performed separately for each
paired A row and B column, so enabling it can materially change the effective
low-rank delta magnitude.
- **Global Scale**: multiplies the merged result after CWB. For LoRA it is
applied once through B/up so the represented delta is scaled once rather
than once per factor.
- **Dynamic Similarity Contrast**: when consensus similarities differ, remaps
them into the range 0.7–1.0 before alpha/beta weighting. It does not alter
row-alignment scores.
- **Soft Comfort Bandpass**: when beta is positive, uses
`(1.5 - similarity)^beta` instead of the narrower
`(1.001 - similarity)^beta`. It has no effect when beta is zero.
- **Position Weight**: blends normalized positional affinity into greedy row
selection. Zero uses cosine only and one uses positional affinity only. It
affects only similarity-aligned embeddings and LoRA rank components; the
original cosine score is still checked against the alignment threshold.
- **Preserve Common Prefix**: copies Model A's longest numerically identical
leading first-axis span before CWB. The current implementation applies this
to all tensors with at least two dimensions, including fixed-coordinate
model tensors.
- **Mismatch Mode**: for anchored tensors that are missing or incompatible,
`skip` preserves the anchor, `zeros` supplies a compatible zero contribution
where possible, and `error` aborts. A valid tensor supplied by only one
secondary input is copied unchanged.
- **Output Filename**: filename without extension, written atomically beneath
the matching category in ComfyUI's models directory.
- **Save Dtype**: requested dtype for generated floating tensors. Unless
Override Dtype is enabled, any participating FP32 source keeps the generated
logical result FP32.
- **Process Device**: device for per-layer FP32 arithmetic. A CUDA OOM releases
the failed working set and retries only that layer on CPU.
- **Exclude Patterns**: matching layers are preserved from their anchor instead
of merged.
- **Discard Patterns**: matching tensors or logical LoRA groups are absent from
the output.
- **Glob Patterns**: switches exclude/discard syntax from regular expressions
to shell-style globs. Pattern fields accept one entry per line.
- **Lazy Load**: enables UEL low-memory loading so tensors are streamed by work
unit and explicitly released.
- **Force Clear Cache**: runs garbage collection and clears the CUDA allocator
cache before each layer. This can reduce retained memory at a substantial
speed cost.
- **Override Dtype**: forces generated floating tensors to Save Dtype. Guarded
low-bit tensors and enabled 1D direct differences are exempt.
- **Include 1D Diffs** (LoRA nodes): enables CWB merging of recognized 1D
direct-difference tensors in FP32. When disabled, shared tensors preserve
Model A and secondary-only tensors preserve their earliest provider.
Scalars and 1D generic tensors use direct mean/median handling. Alignment,
similarity weighting, alpha/beta, DSC, and norm rescaling apply to vector groups
in tensors with at least two dimensions, not to those scalar/1D paths.
### Presets implemented by these nodes
Dense nodes offer `balanced_mean`, `robust_medn`, `selective_mean`,
`varied_mean_rn_softcb`, `diverse_medn_rn_dsc_softcb`, and
`strongdiv_medn_rn_dsc_softcb`.
Embedding nodes offer `balanced_idx_mean`, `balanced_sim_mean`,
`robust_idx_medn`, `robust_sim_medn`, `varied_sim_mean_rn_softcb`, and
`diverse_sim_medn_rn_dsc_softcb`, `focused_strong_sim_medn`,
`focused_balance_sim_medn`, `focused_soft_sim_medn`, and
`focused_weak_sim_medn`.
LoRA nodes offer `broad_sim_medn_rn_softcb`,
`moderate_sim_medn_rn_softcb`, `conservative_sim_medn_rn_softcb`,
`direct_idx_medn_rn_softcb`, `broad_sim_mean_rn_softcb`,
`neutral_sim_medn_rn`, `focused_sim_medn_rn_dsc_softcb`, and
`strongfocus_sim_medn_rn_dsc_softcb`, `focused_strong_sim_medn`,
`focused_balance_sim_medn`, `focused_soft_sim_medn`, and
`focused_weak_sim_medn`.
The LoRA default is `broad_sim_medn_rn_softcb`: median consensus,
similarity alignment, zero similarity threshold, alpha 2, beta 4, norm
rescaling, scale 1, DSC disabled, soft comfort bandpass enabled, position
weight 0.05, and common-prefix preservation disabled. Broad, moderate, and
conservative use alignment thresholds 0, 0.0005, and 0.0025. They represent
increasing anchor preservation in automated mixed-rank LoRA tests.
Every LoRA preset fixes all 12 controls explicitly:
| preset | consensus | alignment | align threshold | sim threshold | alpha | beta | RN | scale | DSC | soft CB | position | prefix | intent |
|---|---|---|---:|---:|---:|---:|---|---:|---|---|---:|---|---|
| `broad_sim_medn_rn_softcb` | median | similarity | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Broad row matching; default |
| `moderate_sim_medn_rn_softcb` | median | similarity | 0.0005 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Moderate anchor preservation |
| `conservative_sim_medn_rn_softcb` | median | similarity | 0.0025 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Strong anchor preservation |
| `direct_idx_medn_rn_softcb` | median | index | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Fixed rank-row correspondence |
| `broad_sim_mean_rn_softcb` | mean | similarity | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Mean-consensus broad merge |
| `neutral_sim_medn_rn` | median | similarity | 0 | 0 | 2 | 0 | on | 1 | off | off | 0.05 | off | Similarity weighting without diversity modulation |
| `focused_sim_medn_rn_dsc_softcb` | median | similarity | 0 | 0 | 2 | 4 | on | 1 | on | on | 0.05 | off | Moderate DSC concentration |
| `strongfocus_sim_medn_rn_dsc_softcb` | median | similarity | 0 | 0 | 2 | 7 | on | 1 | on | on | 0.05 | off | Strong DSC concentration |
| `focused_strong_sim_medn` | median | similarity | 0.85 | 0.60 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Strongly focused matching |
| `focused_balance_sim_medn` | median | similarity | 0.75 | 0.55 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Balanced focused matching |
| `focused_soft_sim_medn` | median | similarity | 0.55 | 0.50 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Soft focused matching |
| `focused_weak_sim_medn` | median | similarity | 0.25 | 0.35 | 1.25 | 0 | off | 1 | off | off | 0.20 | off | Weak focused matching |
The four focused median similarity presets have identical settings in the
embedding and LoRA registries.
The LoRA Multi-Merge counterfactual weight sweep reuses each consensus vector
already computed by the active merge. It reports mean and median consensus
results across configured alpha, beta, similarity-threshold, DSC, and comfort
bandpass values without additional model loads or output files.
## Missing tensors and filtering
- `skip` preserves Model A when an anchored input is missing or incompatible.
In embedding-union mode it omits that source for the affected key.
- `zeros` supplies a compatible zero tensor for the missing source.
- `error` stops the merge.
- Exclude patterns preserve the anchored source tensor.
- Discard patterns remove the matching output tensor or logical LoRA layer.
The missing-input modes above apply to Model-A-owned layers. Model A being
absent from a secondary-only layer is not itself a mismatch. A shape conflict
between multiple available secondary providers does use the selected mismatch
mode, with the earliest available input as its preservation anchor.
Patterns use regular expressions by default. Enable Glob Patterns to use glob
syntax instead.
## LoRA behavior
LoRA inputs use the same parser as the resize and multi-merge nodes. It accepts
all seven low-rank A/B spellings recognized by ComfyUI's `LoRAAdapter`, plus
direct `.diff`, `.diff_b`, `.w_norm`, `.b_norm`, and `.set_weight` layers.
Recognized diffusion-model outputs use the canonical
`diffusion_model.<layer>.lora_A.weight` / `.lora_B.weight` convention.
Secondary formats are matched by normalized logical layer name. Alpha is
absorbed mathematically into each B/up factor using its original rank before
alignment, weighting, or merge. Genuine factor ranks remain variable during
CWB. The first maximum-rank input supplies the output component space, so the
output rank remains the largest genuine input rank. Merged outputs never
contain or invent alpha tensors.
Every CWB merge node writes a `cwb.merge` safetensors metadata record. It
contains selected source model names, fully resolved CWB settings, and merge
options. Presets are stored as effective settings rather than by preset name.
Non-empty exclude/discard filters and their glob mode are included. CWB reports
and diagnostics are not saved in metadata.
Similarity alignment operates on paired latent-rank components. A rows and B
columns share one mapping; fixed input/output feature axes are never reordered.
For a rank-R reference and rank-r source it evaluates an R-by-r similarity
matrix, reduced only by a preserved common prefix. Every genuine source
component can therefore match any genuine reference component. Index alignment
maps source component i to reference component i. In either mode, reference
components without a genuine contributor remain singleton groups and are not
weighted or norm-rescaled against nonexistent rank padding.
Structural rank absence is different from `mismatch_mode=zeros`. The latter is
an explicit request for a semantic zero contributor and continues to apply to
every output-rank component where the mismatch policy is active. Numerical zero
values inside a genuine factor do not identify padding; validity comes from the
factor's original rank.
Global scale is applied once to the resulting pair rather than once per factor.
Incomplete pairs and unrecognized tensors are copied from the earliest
available source unchanged.
Companion-bearing groups (`lora_mid`, reshape, DoRA, or set-weight) are kept
atomically from their earliest available source rather than partially
transformed. Secondary-only logical groups are included automatically: a group
from one source is copied unchanged, while compatible groups from multiple
secondary sources are merged using only those sources. Recognized roles are
written with the same canonical output convention as shared layers.
`include_1d_diffs` is disabled by default. Disabled shared 1D direct layers
preserve Model A, while secondary-only layers preserve their earliest source.
Enabled 1D direct layers participate as complete tensors and are always saved
in FP32.
## Precision and quantized inputs
CWB computes in FP32. Unless Override Dtype is enabled, any participating FP32
source keeps the generated logical layer in FP32. Guarded low-bit tensors and
enabled 1D direct diffs are never downcast by the override.
ComfyUI/comfy-kitchen quantized inputs are rejected before an output writer is
opened. Files with three or more bare INT8, UINT8, FP8, or FP4 tensors are also
rejected. One or two isolated low-bit tensors produce one warning per input and
are preserved without CWB arithmetic; an affected LoRA causes its complete
logical layer to be preserved.
## CWB report output
Every CWB merge node provides a separate `cwb_report` string output after the
existing filename and documentation outputs. Documentation-only execution
returns an explicit message that no report exists because no tensors were
processed.
The filename output uses ComfyUI's wildcard socket type and returns the path
relative to the matching ComfyUI model category. A filename such as
`groupfolder/modelname` therefore returns
`groupfolder/modelname.safetensors`. It can connect directly to compatible
loader Combo inputs without exposing the absolute models-directory path.
The report begins with model-wide counters. LoRA arithmetic then adds grouped
mathematical evidence without retaining tensors or dumping complete similarity
matrices. Layers are grouped only when they share:
- down non-rank dimensions;
- up non-rank dimensions;
- ordered genuine input ranks;
- selected reference input and rank;
- alignment method.
Each dimension group reports layer count, effective matrix dimensions, genuine
source-component match coverage, reference-only components, structural rank
slots excluded from arithmetic, explicit mismatch-zero contributors, candidate
similarity min/mean/max, accepted-match similarity percentiles, and rank-one
delta norm ratios. A compact exception section identifies the lowest match
coverage and largest norm changes by layer name.
For a rank-384 reference and rank-128 source, the report identifies a `384x128`
matrix, at most 128 genuine matches, and the 256 structural rank slots that were
excluded rather than treated as zero-valued contributors.
## Streaming and failure safety
Model parameters and direct patches use fixed-coordinate blending; similarity
row matching is limited to embeddings and LoRA rank components. Each logical
layer is loaded, transferred to the processing device, completed, copied back
to CPU, queued for writing, and released before the next layer. Outputs are
written to a temporary sibling and atomically published only after every layer
and asynchronous write succeeds. Failed merges do not replace or finalize the
requested output. A CUDA OOM retries only the affected layer on CPU after
releasing the failed GPU working set; failure of that retry aborts atomically.
Generated files are published under the matching category inside
`folder_paths.models_dir`, never ComfyUI's general output directory.
+23
View File
@@ -0,0 +1,23 @@
# Delta CWB LoRA Mergers
Delta CWB expands each selected LoRA into full per-layer weight changes and
applies CWB to corresponding dense rows. It preserves model coordinates; it
does not align, reorder, or compress rank components.
LoRA 1 supplies metadata and defines the output layer set. The preset controls
how agreement, disagreement, and row magnitude affect the saved changes. A
connected CWB configuration completely replaces the preset. Inputs have equal
prior weight.
Missing-layer handling can preserve LoRA 1, include explicit zero contributors,
or abort. Excluded layers preserve LoRA 1 and discarded layers are omitted.
Outputs contain only canonical `.diff` tensors with alpha already normalized.
`factorized_lora` output factorizes each completed merged layer immediately and
writes canonical `lora_A` and `lora_B` tensors at the selected maximum rank. It
never writes or retains a complete full-difference model. One-dimensional
differences remain `.diff` tensors.
Processing streams one logical layer from each input at a time and writes each
completed layer incrementally. CUDA out-of-memory retries only the failed layer
on CPU. The destination is replaced only after the complete output is written.
+135
View File
@@ -0,0 +1,135 @@
# CWB Model Graft
`CWB Model Graft` (`CWBModelGraft`) is a standalone training-free diffusion model grafting node that transfers selected blocks, operator families, projection sub-slices, or explicitly mapped layers from a donor diffusion model (`model_b`) into a base diffusion model (`model_a`) using Consensus-Weighted Blending (CWB).
`model_a` defines the output architecture, tensor key set, output shapes, and safetensors metadata. Tensors that are protected, excluded by block or sub-part rules, or absent from `model_b` are preserved from `model_a` without modification.
---
## End-to-End Execution Pipeline
For every logical tensor key in `model_a`, `CWB Model Graft` executes the following sequence:
1. **Quantization & Low-Bit Inspection**
- Both `model_a` and `model_b` are opened with `MemoryEfficientSafeOpen` and inspected via `DiffusionQuantization` and `diffusion_key_map`.
- Full-model ComfyUI/comfy-kitchen quantized files or files containing 3 or more bare low-bit (`INT8`, `UINT8`, `FP8`, `FP4`) tensors are rejected before opening an output writer.
- Any layer with 1–2 isolated low-bit tensors in either `model_a` or the resolved `model_b` source keys is preserved byte-for-byte from `model_a` (`force_raw=True`).
2. **Discard & Exclude / Include Pattern Filtering**
- `discard_patterns`: If a `model_a` key (or its physical key) matches `discard_patterns`, the tensor is omitted from the output file.
- `exclude_patterns` with `include_mode = False` (default): Matching keys are preserved unchanged from `model_a`.
- `exclude_patterns` with `include_mode = True`: Acts as a whitelist. Non-matching keys are preserved unchanged from `model_a`. Explicitly matched whitelist keys are eligible for grafting even if they lie outside standard block numbering (such as `text_embedder` or modulation layers).
3. **Donor Key Resolution & Automatic QKV Packing / Splitting**
- **Custom Key Map (`custom_key_map`)**: Checked first. If `base_key` has an explicit mapping `donor_key -> base_key` and `donor_key` exists in `model_b`, `donor_key` is used as the donor source.
- **Direct Match**: If `base_key` exists directly in `model_b`, it is paired 1-to-1.
- **Split-to-Fused QKV (`fuse_qkv`)**: If `model_a` has a fused `<prefix>.proj.qkv.weight` tensor and `model_b` instead stores split `<prefix>.proj.q.weight`, `<prefix>.proj.k.weight`, and `<prefix>.proj.v.weight`, all three donor tensors are streamed in the same work unit and concatenated along `dim=0` (`torch.cat([q, k, v], dim=0)`) before sub-part selection and CWB blending.
- **Fused-to-Split QKV (`split_qkv`)**: If `model_a` has split `<prefix>.proj.{q,k,v}.weight` and `model_b` stores fused `<prefix>.proj.qkv.weight`, the fused donor tensor is streamed and split into 3 equal chunks along `dim=0` (`chunk(3, dim=0)`), selecting index `0` for `q`, `1` for `k`, or `2` for `v`.
- **Unmatched Keys**: If no donor source exists for `base_key`, `model_a`'s tensor is preserved unchanged.
4. **Block Index Detection & Early-Block Protection**
- Main transformer block indices are parsed from keys matching `(^|\.)(layers|blocks|double_blocks|single_blocks|transformer_blocks)\.(\d+)\.`.
- Auxiliary sub-networks (`text_embedder.*`, `token_refiner.*`) do not count toward main transformer block numbering and are preserved from `model_a` unless targeted via `custom_key_map` or `include_mode=True` with `exclude_patterns`.
- `max_block_index` is computed as the highest main block index in `model_a`.
- If a key's `block_idx < protect_early_blocks`, the tensor is preserved 100% from `model_a`.
5. **Operator Classification & Block Scheduling**
- Each main-block key is classified by its token path:
- **Attention (`attention`)**: `.qkv.`, `.qkv_proj.`, `.proj.q.`, `.proj.k.`, `.proj.v.`, `.proj.o.`, `.attn.`, `.self_attn.`, `.q_norm.`, `.k_norm.`, `.to_q.`, `.to_k.`, `.to_v.`, `.to_out.`.
- **Feed-Forward (`ffn`)**: `.ffn_up_gate.`, `.ffn_up.`, `.ffn_gate.`, `.ffn_down.`, `.mlp.`, `.ff.`, `.fc1.`, `.fc2.`, `.w1.`, `.w2.`, `.w3.`.
- Keys inside main blocks that are neither attention nor feed-forward (such as block-level `norm1` and `norm2`) are preserved from `model_a` unless explicitly targeted via `custom_key_map` or `include_mode=True`.
- For classified keys with `block_idx >= protect_early_blocks`, the corresponding rule (`attention_block_rule` or `ffn_block_rule`) determines whether the block is grafted:
- `all_after_protected`: Every block with `block_idx >= protect_early_blocks`.
- `interleaved_even`: Even-numbered blocks (`block_idx % 2 == 0`) with `block_idx >= protect_early_blocks`.
- `interleaved_odd`: Odd-numbered blocks (`block_idx % 2 == 1`) with `block_idx >= protect_early_blocks`.
- `back_half_only`: Blocks in the second half of the network (`block_idx >= max(protect_early_blocks, (max_block_index + 1) // 2)`).
- `none`: No blocks of this operator family are grafted.
6. **Sub-Part Slicing (`attention_subpart` and `ffn_subpart`)**
- **Attention Sub-Parts (`attention_subpart`)**:
- `all`: All attention tensors in selected blocks (`qkv`, `q`, `k`, `v`, `o`, `q_norm`, `k_norm`) are grafted.
- `value_and_output`:
- Split query/key projections and norms (`.proj.q.`, `.proj.k.`, `.to_q.`, `.to_k.`, `.q_norm.`, `.k_norm.`) are preserved 100% from `model_a`.
- Fused QKV weights (`.qkv.`, `.qkv_proj.`) whose first dimension is divisible by 3 are split into `q_a, k_a, v_a = base.chunk(3, dim=0)` and `q_b, k_b, v_b = donor.chunk(3, dim=0)`. `q_a` and `k_a` are kept byte-exact from `model_a`, `v_a` and `v_b` are CWB-blended into `v_merged`, and the tensor is reassembled as `torch.cat([q_a, k_a, v_merged], dim=0)`.
- Split value projections (`.proj.v.`, `.to_v.`) and output projections (`.proj.o.`, `.out_proj.`, `.to_out.`) are CWB-blended in full.
- `output_only`: Only attention output projections (`.proj.o.`, `.out_proj.`, `.to_out.`) are CWB-blended; all Q/K/V and norm tensors are preserved 100% from `model_a`.
- **Feed-Forward Sub-Parts (`ffn_subpart`)**:
- `all`: All feed-forward tensors in selected blocks (`ffn_up_gate`, `ffn_up`, `ffn_gate`, `ffn_down`, `fc1`, `fc2`, `w1`, `w2`, `w3`) are grafted.
- `up_only`:
- Down and gate projections (`.ffn_down.`, `.fc2.`, `.w2.`, `.ff.net.2.`, `.ffn_gate.`, `.w3.`) are preserved 100% from `model_a`.
- Fused up+gate weights (`.ffn_up_gate.`) whose first dimension is divisible by 2 are split into `up_a, gate_a = base.chunk(2, dim=0)` and `up_b, gate_b = donor.chunk(2, dim=0)`. Only `up_a` and `up_b` are CWB-blended into `up_merged`, `gate_a` is preserved 100% from `model_a`, and the tensor is reassembled as `torch.cat([up_merged, gate_a], dim=0)`.
- Split up projections (`.ffn_up.`, `.fc1.`, `.w1.`) are CWB-blended in full.
- `down_only`: Only down projections (`.ffn_down.`, `.fc2.`, `.w2.`, `.ff.net.2.`) are CWB-blended; all up/gate tensors are preserved 100% from `model_a`.
7. **Cross-Shape Alignment, Row Similarity Matching & Norm Locking**
- **Trailing-Dimension Alignment**: If `donor_tensor` and `base_tensor` differ in shape:
- Same element count with different rank: `donor_tensor` is reshaped to `base_tensor.shape`.
- Different rank and different element count: `donor_tensor` is linearly interpolated in FP32 to `base_tensor.numel()` and reshaped to `base_tensor.shape`.
- Same rank (`ndim >= 2`) with different trailing dimensions (`donor_tensor.shape[1:] != base_tensor.shape[1:]`): Each donor row is linearly interpolated in FP32 along its flattened trailing features to match `prod(base_tensor.shape[1:])` and reshaped to `(donor_rows, *base_tensor.shape[1:])`.
- **Row Matching & CWB Blending**: `merge_cwb_tensors` runs in FP32 with `allow_similarity_alignment=True` and `reference_index=0` (`model_a` as reference). When the active CWB settings use `alignment_method="similarity"` (or when row counts differ), donor rows are greedily paired to base rows by cosine similarity (subject to `alignment_threshold` and `position_weight`). Unmatched base rows remain anchored to `model_a`.
- **Base Magnitude Locking**: Whenever `donor_tensor` required shape alignment or `rescale_norm` is enabled in the active CWB settings, the blended tensor's overall Frobenius norm is locked to `model_a`'s tensor norm (`merged * (||base|| / ||merged||)`), preventing activation scale drift across downstream blocks.
8. **Atomic Streaming Output & Metadata**
- Tensors are streamed one work unit at a time through UEL (`batch_size=1`, `prefetch_batches=1`), written incrementally via `atomic_uel_writer`, and released before advancing to the next layer.
- If a CUDA out-of-memory error occurs on a layer, the failed GPU working set is released and that layer retries on CPU.
- Output metadata preserves `model_a`'s metadata and records a JSON summary under `cwb.graft` containing `base_model`, `donor_model`, `attention_preset`, `ffn_preset`, `protect_early_blocks`, `attention_block_rule`, `ffn_block_rule`, `attention_subpart`, and `ffn_subpart`.
---
## Complete Input Reference
- **`execution_mode`** (`MERGE` | `DOCUMENTATION ONLY`, default `MERGE`):
- `MERGE`: Executes the streaming graft and writes the output `.safetensors` file under `models/diffusion_models/`.
- `DOCUMENTATION ONLY`: Returns this documentation string without opening any model files.
- **`model_a`** (Combo, `diffusion_models`):
- Base diffusion model that anchors architecture, output tensor keys, shapes, metadata, and untouched/protected layers.
- **`model_b`** (Combo, `diffusion_models`):
- Donor diffusion model from which selected blocks, sub-projections, or mapped layers are grafted into `model_a`.
- **`attention_preset`** (Combo, default `robust_medn`):
- Dense CWB preset applied to grafted attention tensors (`balanced_mean`, `robust_medn`, `selective_mean`, `varied_mean_rn_softcb`, `diverse_medn_rn_dsc_softcb`, `strongdiv_medn_rn_dsc_softcb`). Overridden completely if `attention_config` is connected.
- **`attention_config`** (`CWB_CONFIG`, optional):
- Optional custom CWB configuration from `CWB Custom Configuration` for attention layers. Enabling `alignment_method = "similarity"` on this config activates greedy cosine row matching across attention rows.
- **`ffn_preset`** (Combo, default `balanced_mean`):
- Dense CWB preset applied to grafted feed-forward (MLP) tensors and custom-mapped non-attention tensors. Overridden completely if `ffn_config` is connected.
- **`ffn_config`** (`CWB_CONFIG`, optional):
- Optional custom CWB configuration from `CWB Custom Configuration` for feed-forward (MLP) layers.
- **`protect_early_blocks`** (Int, default `2`, range `0..256`):
- Number of initial main transformer blocks (`0 .. protect_early_blocks - 1`) kept 100% identical to `model_a` so global image composition and layout remain stable.
- **`attention_block_rule`** (`interleaved_even` | `interleaved_odd` | `back_half_only` | `all_after_protected` | `none`, default `interleaved_even`):
- Block selection schedule for attention layers at index `>= protect_early_blocks`.
- **`ffn_block_rule`** (`all_after_protected` | `interleaved_even` | `interleaved_odd` | `back_half_only` | `none`, default `all_after_protected`):
- Block selection schedule for feed-forward (MLP) layers at index `>= protect_early_blocks`.
- **`attention_subpart`** (`all` | `value_and_output` | `output_only`, default `all`):
- Restricts attention grafting to all attention weights, only the Value (`V`) slice/projection and Output (`O`) projection, or only the Output (`O`) projection.
- **`ffn_subpart`** (`all` | `up_only` | `down_only`, default `all`):
- Restricts feed-forward grafting to all FFN weights, only the Up projection (including the `up` half of packed `ffn_up_gate`), or only the Down projection.
- **`custom_key_map`** (String, multiline, default `""`):
- Optional explicit key mappings, one per line, formatted as `donor_key -> base_key` or `donor_key = base_key` (lines starting with `#` are ignored). Mapped keys bypass block-number filtering so arbitrary donor tensors can be grafted into specific base tensors.
- **`output_filename`** (String, default `cwb_grafted_model`):
- Output filename relative to `models/diffusion_models/` (`.safetensors` extension added automatically).
- **`save_dtype`** (`fp32` | `fp16` | `bf16`, default `bf16`):
- Target floating-point precision when `override_dtype` is enabled or when decoding quantized tensors.
- **`process_device`** (`cuda` | `cpu`, default `cuda`):
- Execution device for per-tensor FP32 CWB math; automatically falls back to `cpu` per tensor on CUDA OOM.
- **`exclude_patterns`** (String, multiline, default `""`):
- Regex (or glob when `glob_patterns=True`) patterns. Preserves matching tensors from `model_a` when `include_mode=False`, or whitelists matching tensors for grafting when `include_mode=True`.
- **`discard_patterns`** (String, multiline, default `""`):
- Regex (or glob) patterns of tensors to omit completely from the saved model.
- **`glob_patterns`** (Boolean, default `False`):
- Interpret `exclude_patterns` and `discard_patterns` as shell-style globs (`*`, `?`) instead of regular expressions.
- **`lazy_load`** (Boolean, default `True`):
- Enable low-memory UEL streaming so only the current work unit resides in memory.
- **`force_clear_cache`** (Boolean, default `False`):
- Run `gc.collect()` and `torch.cuda.empty_cache()` before each tensor work unit.
- **`override_dtype`** (Boolean, default `False`):
- Force grafted floating tensors to `save_dtype`; when `False`, each grafted tensor preserves `model_a`'s floating dtype.
- **`include_mode`** (Boolean, default `False`):
- Invert `exclude_patterns` into a whitelist so only matching tensors are eligible for grafting.
---
## Outputs
1. **`output_filename`** (`*` / `io.AnyType`): Category-relative filename (`<output_filename>.safetensors`) ready to connect directly to diffusion model loader nodes.
2. **`documentation`** (`STRING`): Complete markdown documentation for `CWB Model Graft`.
3. **`cwb_report`** (`STRING`): Diagnostic report detailing active attention/FFN presets, candidate/accepted alignment matches, consensus similarity statistics, CPU fallbacks, total grafted tensor count, and total preserved base tensor count.
+37
View File
@@ -0,0 +1,37 @@
# Lodestone LoRA Mergers
The dedicated Lodestone nodes merge full logical LoRA deltas rather than
blending down and up factors independently. Stored alpha is normalized before
each factor pair is expanded. Results are saved as canonical full direct-difference
tensors and contain no factor or alpha tensors.
The 2-, 3-, and Multi-Merge nodes treat every selected LoRA as one equal-prior
slot, matching the supplied Lodestone calculation.
## Calculation Modes
- `sum`: sum the source deltas.
- `mean`: divide the summed delta by the number of active contributors.
- `slotnorm`: scale each source delta to the median source Frobenius norm, then
average.
- `normmatch`: retain the summed direction and set its Frobenius norm to the
median source norm.
- `slotnorm-normmatch`: equalize source norms, sum, then set the final norm to
the original median source norm.
A zero target or merged norm remains zero.
## Memory Behavior
UEL streams one logical layer at a time and writes each completed full-delta
tensor incrementally through the atomic writer. Dense tensors used for the
merge and Frobenius norms exist only within the active layer. A CUDA
out-of-memory error releases CUDA state and retries that logical layer on CPU.
## Fixed-Input Controls
LoRA 1 anchors output layers and filtering behavior. `skip` preserves LoRA 1
when another input lacks a required layer, `zeros` supplies an explicit zero
slot, and `error` aborts. Source factors may have different ranks when their
expanded dense shapes agree. Exclude and discard patterns apply to complete
logical layers.
+30
View File
@@ -0,0 +1,30 @@
# LoRA-on-model analysis
**Analyze LoRA Effect on Diffusion Model** takes a LoRA as `model_a` and a base
diffusion model as `model_b`. It applies the LoRA transiently and compares the
original base weights against the patched weights. The measured update is
`patched − original`; no merged checkpoint, delta file, or inference is produced.
The comparison and CWB reports use original base as A and patched base as B.
The CSV outputs contain per-tensor difference metrics and row cosine summaries.
Unchanged floating base tensors remain in the analysis, so global metrics reflect
the impact on the selected base model, not just the layers touched by the LoRA.
Application and statistics use FP32, matching the existing analysis arithmetic;
this does not simulate rounding to a separately saved model dtype.
LoRA strength defaults to 1. Alpha/rank normalization, DoRA, LoCon, and reshape
handling use the existing application path and ComfyUI adapter. Direct deltas,
including bias deltas, are applied; explicit `set_weight` patches replace weights
regardless of strength, matching the existing application behavior.
Exclude Patterns and Include Mode filter **base tensor keys** before loading.
Unmapped LoRA groups and unused source tensors are listed separately. Unsupported
low-bit work units and non-floating base tensors are reported and omitted from
numeric metrics; they are not decoded or represented as successfully analyzed.
Application errors fail the operation rather than report a failed patch as zero
impact. A CUDA out-of-memory error retries the complete affected unit on CPU.
Loading uses the shared UEL work-unit streamer: one base tensor and its required
adapter tensors, batch size 1 and prefetch depth 1. Sources are marked processed
and streams closed on success or failure. Only scalar statistics and bounded top
differences survive each unit. No whole model or LoRA tensor dictionary is built.
+155 -47
View File
@@ -2,86 +2,194 @@
---
### Core Convention Across All 2-Model Modes
- **Model A is ALWAYS the base / anchor model.**
- **Model B is ALWAYS the donor / incoming model being merged into Model A.**
- When interpolating or blending between weights:
$$\text{Output} = A \cdot (1 - \text{mask}) + B \cdot \text{mask}$$
- A blend factor / weight of `0.0` yields **100% Model A** (pure Model A).
- A blend factor / weight of `1.0` yields **100% Model B** (pure Model B).
- Intermediate values control how much of Model B replaces or alters Model A.
---
## Weight-Sum
> `A * (1 - α) + B * α`. A simple linear interpolation between two models.
> $\text{Output} = A \cdot (1 - \alpha) + B \cdot \alpha$. Direct linear interpolation between Model A and Model B.
**Models Used:** A (Base), B (Incoming)
**Models Used:** A, B
**Parameters:**
- **Alpha:** Interpolation weight. `0.0` is 100% Model A, `1.0` is 100% Model B.
- **Alpha:** The blend ratio determining what fraction of **Model B** replaces **Model A**.
- `0.0`: **100% Model A, 0% Model B** (exact Model A; Model B is completely ignored).
- `0.25`: **75% Model A, 25% Model B** (mostly Model A, with 25% of Model B blended in).
- `0.50`: **50% Model A, 50% Model B** (equal 50/50 mix of both models).
- `0.75`: **25% Model A, 75% Model B** (mostly Model B, retaining 25% of Model A).
- `1.0`: **0% Model A, 100% Model B** (exact Model B; Model A is completely replaced).
---
## Comparative-Interpolation
> Interpolates between A and B based on the relative differences in their tensor values. Creates a unique blend by deciding how much of each model to use for each individual weight.
> Interpolates between A and B based on the relative differences in their tensor values, creating a selective blend rather than a flat uniform mix.
**Models Used:** A (Base), B (Incoming)
**Models Used:** A, B
**Parameters:**
- **Alpha:** Controls the interpolation curve's strength and shape.
- **Beta:** Switches the interpolation style. `0.0` focuses on similarity, `1.0` focuses on difference.
- **Gamma:** Mixes between a randomized (binomial) interpolation at `0.0` and a smooth (linear) one at `1.0`.
- **Alpha (Blend Intensity):** Controls the overall pull of **Model B into Model A**.
- `0.0`: Heavily favors Model A (virtually no Model B incorporated).
- `0.5`: Balanced curve where weights with moderate differences blend equally.
- `1.0`: Strongly pulls in Model B wherever weights differ.
- **Beta (Feature Focus):** Controls whether the merge targets similar or divergent weights.
- `0.0`: Focuses blending on weights where Model A and Model B are already similar.
- `1.0`: Focuses blending on weights where Model A and Model B differ most.
- **Gamma (Transition Style):** Mixes between discrete selection and smooth gradation.
- `0.0`: Stochastic / binomial selection (individual weights are chosen probabilistically as either pure A or pure B).
- `1.0`: Smooth linear interpolation across all weights.
---
## Power-Up (DARE)
> Adds the unique capabilities of Model B to Model A using the Drop and Rescale (DARE) technique, which often preserves the knowledge of the base model better than simple additions.
> Adds the unique capabilities of Model B onto Model A using the Drop and Rescale (DARE) technique. Computes the task delta $(B - A)$, randomly drops a fraction of the delta parameters, rescales the survivors to maintain expected magnitude, and adds them to Model A.
**Models Used:** A (Base), B (Donor)
**Models Used:** A, B
**Parameters:**
- **Alpha:** The dropout rate. This is the proportion of unique weights from Model B that are *dropped* before merging. A higher value means less of B is merged.
- **Beta:** A final multiplier for the rescaled difference before it's added to Model A.
- **Alpha (Dropout Rate $p$):** Proportion of Model B's delta parameters to randomly zero out.
- Range: `0.0` to `1.0`.
- Typical: `0.7` to `0.9`. For example, `0.85` means 85% of Model B's delta values are discarded, keeping only 15% of the most distinct changes.
- **Beta (Scale Multiplier):** Final strength multiplier for Model B's rescaled delta before adding to Model A.
- `1.0`: Standard 100% strength addition of Model B's surviving features.
- `0.5`: Half-strength addition of Model B's features.
- `1.2`: Amplified addition of Model B's features.
- **Rescaling Logic:** Surviving parameters are automatically scaled by $1 / (1 - p)$ per the DARE paper to preserve the expected total update magnitude.
---
## Power-Up (DARE+TIES)
> Combines DARE (Drop and Rescale) with TIES (Trim, Elect Sign) to merge fine-tuned models cleanly. DARE sparsifies the delta; TIES trims low-magnitude background noise and enforces sign consensus.
**Models Used:** A (Base), B (Donor)
**Parameters:**
- **Alpha (DARE Drop Rate):** Fraction of Model B's delta parameters randomly dropped (`0.5`–`0.9`). Higher values produce a sparser capability transfer.
- **Beta (TIES Trim Quantile):** Fraction of the smallest-magnitude delta parameters eliminated after DARE to remove noise. `0.0` disables trimming; `0.2` trims the lowest 20% by absolute magnitude. Typical: `0.1`–`0.3`.
- **Gamma (Lambda Scale / Strength):** Multiplier for the filtered delta added to Model A.
- `1.0`: Standard 100% strength of Model B's filtered capability added to Model A.
- `0.25`: Gentle 25% addition of Model B's capabilities.
- `1.5`: Amplified capability transfer.
- **Seed:** Random seed for reproducible DARE dropout masks.
**Algorithm:**
1. Compute task vector $\delta = B - A$
2. **DARE:** Apply random binary mask with keep probability $1 - \alpha$; rescale survivors by $1 / (1 - \alpha)$
3. **TIES trim:** Zero out parameters below the $\beta$-quantile magnitude threshold
4. **TIES elect:** Determine dominant sign direction per position; zero out disagreeing parameters
5. Return $A + \gamma \cdot \delta_{\text{filtered}}$
---
## Enhanced Man Interp (Enhanced Manual Interpolation)
> Selective interpolation between Model A and Model B based on normalized pairwise differences, constrained within a user-defined threshold band.
**Models Used:** A (Base), B (Incoming)
**Formula:**
$$\text{diff} = \frac{\max(|A - B|) - |A - B|}{\max(|A - B|)}$$
$$\text{mask}_{\text{threshold}} = (\beta < \text{mean}(\text{diff}) < \gamma)$$
$$\text{blend\_strength} = \text{diff}^{(1/\alpha - 1)} \cdot \text{mask}_{\text{threshold}}$$
$$\text{Output} = A \cdot (1 - \text{interpolated\_mask}) + B \cdot \text{interpolated\_mask}$$
**Parameters:**
- **Alpha (Model B Pull Strength):** Controls how strongly Model B is pulled into Model A.
- Mathematically, the exponent applied to similarity is $(1/\alpha - 1)$.
- **`Alpha = 0.25`**: Exponent is $(1/0.25 - 1) = 3$. This aggressively suppresses Model B's contribution toward zero. Only weights that are virtually identical between both models get even a minor blend of Model B; the final result remains overwhelmingly **Model A** (>90% Model A).
- **`Alpha = 0.50`**: Exponent is $(1/0.50 - 1) = 1$ (linear). Weights with high similarity incorporate Model B proportionally up to 50%.
- **`Alpha = 0.75`**: Exponent is $0.33$. Flattens similarity, pulling in substantially more of Model B wherever differences meet the threshold.
- **`Alpha = 1.0`**: Exponent is $0$. Full maximum pull of Model B across all qualifying weights.
- **Beta (Lower Mean Threshold):** Minimum normalized similarity required before Model B can be blended into Model A. Layers/weights with similarity below $\beta$ remain 100% Model A.
- **Gamma (Upper Mean Threshold):** Maximum normalized similarity allowed for blending. Layers/weights with similarity above $\gamma$ remain 100% Model A. Together, $[\beta, \gamma]$ acts as a bandpass filter targeting only weights with moderate differences.
- **Delta (Smoothness Factor):** Blends between stochastic selection (`0.0`, a Bernoulli coin-flip choosing either pure A or pure B per weight) and smooth continuous blending (`1.0`, smooth linear weighting between A and B).
---
## Enhanced Auto Interp (Enhanced Automatic Interpolation)
> Automated version of Enhanced Interpolation that dynamically computes threshold boundaries centered on the layer's mean difference.
**Models Used:** A (Base), B (Incoming)
**Formula:**
$$\text{threshold\_band} = [\text{mean}(\text{diff}) \cdot (1 - \beta), \;\; \text{mean}(\text{diff}) \cdot (1 + \beta)]$$
$$\text{Output} = A \cdot (1 - \text{interpolated\_mask}) + B \cdot \text{interpolated\_mask}$$
**Parameters:**
- **Alpha (Model B Pull Strength):** Controls how much of Model B replaces Model A inside the active band.
- `0.25`: Low Model B contribution; output remains mostly Model A.
- `0.50`: Moderate, balanced blend of Model B into Model A.
- `1.0`: Maximum Model B influence within the threshold band.
- **Beta (Threshold Width Around Mean):** Defines how wide the window of affected weights is around the layer's average difference.
- `0.1`: Narrow window ($\pm 10\%$ around mean difference).
- `0.5`: Broad window ($\pm 50\%$ around mean difference).
- **Gamma (Smoothness Factor):** Blends between stochastic Bernoulli selection (`0.0`) and smooth continuous linear interpolation (`1.0`).
---
## Weight-Sum Cutoff
> Linear interpolation that only blends Model B into Model A for weights whose relative difference falls within an explicit cutoff window.
**Models Used:** A (Base), B (Incoming)
**Formula:**
$$\text{Output} = A \cdot (1 - \alpha \cdot \text{mask}) + B \cdot (\alpha \cdot \text{mask})$$
**Parameters:**
- **Alpha (Model B Mix Ratio):** The percentage of Model B blended into Model A for weights inside the window.
- `0.0`: 0% Model B (100% Model A everywhere).
- `0.25`: 25% Model B, 75% Model A for weights inside the cutoff window; weights outside remain 100% Model A.
- `0.50`: 50% Model B, 50% Model A inside the window.
- `1.0`: 100% Model B inside the window (complete replacement of qualifying weights).
- **Beta (Upper Difference Threshold):** Upper cutoff for normalized difference.
- **Gamma (Lower Difference Threshold):** Lower cutoff for normalized difference. Weights with difference between $\gamma$ and $\beta$ receive the blend; all other weights stay 100% Model A.
---
## SVD LoRA Extraction
> Extracts the difference between Model B and Model A and saves it as a new LoRA file. **This mode outputs a LoRA, not a full model.** The output file will be saved to your `loras` directory.
> Computes the difference between Model A and Model B and extracts it as a new LoRA file using Singular Value Decomposition (SVD). **Outputs a LoRA safetensors file into your `loras` directory, not a full model.**
**Models Used:** A (Tuned / Modified Model), B (Base Model)
- Computes delta: $\Delta W = A - B$
- Low-rank approximation: $\Delta W \approx U S V^T \implies \text{lora\_B} \cdot \text{lora\_A}$
**Models Used:** A (Tuned Model), B (Base Model)
**Parameters:**
- **Alpha:** The Rank (dimension) for the LoRA's standard layers.
- **Beta:** The Rank (dimension) specifically for 3x3 convolution layers.
- **Gamma:** The clamp quantile for weight values. `0.99` is a typical value to prevent outlier weights from dominating.
- **Alpha (Linear Rank):** Target rank (dimension) for 2D dense/linear layer LoRA factors (e.g. `32`, `64`, `128`).
- **Beta (Conv Rank):** Target rank specifically for 3x3 convolution layer LoRA factors (e.g. `16`, `32`).
- **Gamma (Clamp Quantile):** Outlier clamping quantile (e.g. `0.99` clips the top 1% extreme values before SVD to prevent outlier distortion).
---
## Layer Mismatch Handling
## Layer Mismatch Handling (`mismatch_mode`)
When merging models with different layer structures (e.g., fine-tuned models with added/removed layers, or LoRAs with partial layer coverage), the `mismatch_mode` parameter controls behavior:
When Model A and Model B have differing layer structures or when merging partial LoRAs:
| Mode | Behavior |
|------|----------|
| `skip` | Missing layers in B use A's values **(default)** |
| `zeros` | Missing layers in B are treated as zeros |
| `error` | Fail if any layer is missing |
| `skip` | Missing layers in Model B use Model A's values unchanged **(default)** |
| `zeros` | Missing layers in Model B are treated as zeros |
| `error` | Raise an error and stop execution if any layer is missing in Model B |
**Common scenarios:**
- Merging a fine-tuned model with added/removed layers
- Combining LoRA-extracted differences with base models
- Cross-architecture experiments
**Note:** Extra layers in Model B that don't exist in Model A are currently ignored. The output will always have the same layer structure as Model A.
**Note:** Extra layers in Model B that do not exist in Model A are ignored. Output topology always mirrors Model A.
---
## Layer Filtering (Regex Patterns)
## Dtype Preservation & Override
Use regex patterns to control which layers are merged:
By default, the merger preserves the highest precision `dtype` for each key found across active source models and the requested `save_dtype`:
- **Preserve Higher Precision (default, `override_dtype = False`)**: If a key is `float32` in Model A or Model B, but `save_dtype="bf16"` is selected, that key remains `float32` to protect sensitive normalization layers.
- **Explicit Override (`override_dtype = True`)**: Forces all floating tensors strictly to `save_dtype`.
### `exclude_patterns`
Layers matching any pattern will **keep Model A's values only** (no merge).
---
### `discard_patterns`
Layers matching any pattern will be **removed entirely** from the output.
## Layer Filtering
**Pattern format:**
- Whitespace-separated regex patterns
- Patterns use **substring matching** (not full match)
- Example: `text_model lora` matches any key containing "text_model" OR "lora"
- Example: `layer\.[0-5]\.` matches layers 0-5 using regex syntax
**Pattern Examples:**
| Pattern | Matches |
|---------|---------|
| `text_model` | All text encoder layers |
| `\.norm` | All normalization layers |
| `attn\.(q\|k\|v)` | Query, key, value attention weights |
| `block\.[0-9]\.` | Blocks 0-9 |
- **`exclude_patterns`**: Newline-separated patterns. Matching layers are **not merged** and keep Model A's values unchanged (unless `include_mode` is enabled).
- **`include_mode`**: Inverts `exclude_patterns` into a whitelist: only matching layers are merged; all non-matching layers keep Model A's values unchanged.
- **`discard_patterns`**: Newline-separated patterns. Matching layers are **removed completely** from the saved output.
- **`glob_patterns`**: When enabled, uses shell wildcards (`*`, `?`). When disabled, uses regex substring matching.
+72 -43
View File
@@ -2,82 +2,111 @@
---
### Core Convention Across All 3-Model Modes
- **Model A is ALWAYS the base / receiving model** that receives the merge modifications.
- **Model B is the donor / target fine-tune model.**
- **Model C is the reference base model of Model B.**
- In difference merging, $(B - C)$ calculates the isolated delta (the specific training / fine-tune changes introduced in Model B relative to its base Model C).
- That delta is then added onto Model A with scaling $\alpha$:
$$\text{Output} = A + (B - C) \cdot \alpha$$
---
## Add-Difference
> `A + (B - C) * α`. This classic technique applies the "difference" between models B and C to model A. For example, it can be used to add a style (B) to a base model (A) while subtracting the style's original base (C).
> $\text{Output} = A + (B - C) \cdot \alpha$. Extracts the isolated fine-tune / style delta from $(B - C)$ and applies it directly to Model A.
**Models Used:** A (Receiving Base), B (Donor Fine-Tune), C (Donor's Base Reference)
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** A multiplier for the `(B - C)` difference vector before it's added to A.
- **Beta:** A smoothing toggle. `0` for Off, `1` for On. When On, applies a median and gaussian filter to the difference, which can result in a smoother merge.
- **Alpha (Delta Multiplier / Strength):** Scaling factor for how strongly the $(B - C)$ difference is added to Model A.
- `0.0`: **0% addition** (Output is pure Model A; no changes added).
- `0.25`: **25% strength addition** (gentle, subtle application of Model B's style/features onto Model A).
- `0.50`: **50% strength addition** (moderate application).
- `1.0`: **100% full-strength addition** (applies the exact full delta of B relative to C onto Model A).
- `1.5`: **150% over-amplified addition** (exaggerates the difference; can introduce visual artifacts if pushed too high).
- **Beta (Smoothing Toggle):**
- `0`: Off (raw mathematical delta added directly).
- `1`: On (applies median and Gaussian smoothing to the delta before adding, reducing harsh pixel noise and high-frequency artifacts).
---
## Train-Difference
> A variation of Add-Difference from the SuperMerger extension that uses a more complex scaling factor based on the relative distances between the tensors of A, B, and C.
> A variation of Add-Difference that calculates an adaptive scaling factor based on the relative vector distances between Model A, Model B, and Model C.
**Models Used:** A (Receiving Base), B (Donor Fine-Tune), C (Donor's Base Reference)
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** A multiplier for the final calculated "training" difference.
- **Alpha (Difference Multiplier):** Multiplier for the calculated distance-weighted difference added to Model A.
- `0.0`: Pure Model A.
- `1.0`: Standard 100% strength addition of the adaptive training delta.
- `< 1.0`: Softened application.
- `> 1.0`: Amplified application.
---
## Extract-Features
> A powerful mode that identifies features present in both `(B - A)` and `(C - A)` and adds them to A. Allows for fine-grained control over combining aspects based on their similarity.
> Analyzes feature vectors present in both $(B - A)$ and $(C - A)$, using cosine similarity to isolate shared features and blend them into Model A.
**Models Used:** A (Base), B (Donor 1), C (Donor 2)
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** Weights the merge between Model B (`0.0`) and Model C (`1.0`).
- **Beta:** Controls the focus on similarity (`0.0`) versus dissimilarity (`1.0`).
- **Gamma:** A bias exponent for similarity. Higher values increase the bias.
- **Delta:** A final multiplier for the extracted features before they are added to Model A.
- **Alpha (B vs C Contribution Balance):** Weights the source contribution between Model B and Model C.
- `0.0`: 100% Model B's delta features.
- `0.5`: Equal mix of features from both Model B and Model C.
- `1.0`: 100% Model C's delta features.
- **Beta (Similarity vs Dissimilarity):**
- `0.0`: Focuses exclusively on features where B and C agree / are similar relative to A.
- `1.0`: Focuses on features where B and C diverge / are dissimiliar relative to A.
- **Gamma (Similarity Bias Exponent):** Exponential bias curve shaping how sharply similar features are separated from dissimilar features.
- **Delta (Feature Multiplier / Strength):** Final scaling factor for the extracted features before adding them onto Model A.
- `0.0`: No features added (pure Model A).
- `1.0`: Standard strength addition of extracted features onto Model A.
---
## Add-Dissimilarities
> Identifies features that are dissimilar between Model B and Model C and adds them to Model A. Useful for combining unique aspects of two different models.
> Identifies features that are dissimilar between Model B and Model C (relative to Model A) and transfers them into Model A. Ideal for combining non-overlapping, unique aspects from two different fine-tunes without conflicting duplicate features.
**Models Used:** A (Base), B (Donor 1), C (Donor 2)
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** Weights the comparison between Model B (`0.0`) and Model C (`1.0`).
- **Beta:** A multiplier for the final extracted dissimilar features.
- **Gamma:** A bias exponent for the similarity calculation.
- **Alpha (B vs C Balance):** Balance weighting between Model B and Model C when extracting differences.
- `0.0`: 100% weight to Model B's dissimilarities.
- `0.5`: Balanced extraction from both models.
- `1.0`: 100% weight to Model C's dissimilarities.
- **Beta (Feature Multiplier / Strength):** Final scaling multiplier for the dissimilar features added to Model A.
- `0.0`: No features added (pure Model A).
- `1.0`: Standard strength addition onto Model A.
- **Gamma (Similarity Bias Exponent):** Exponent controlling the cutoff sensitivity between similar and dissimilar directions.
---
## Layer Mismatch Handling
## Layer Mismatch Handling (`mismatch_mode`)
When merging models with different layer structures, the `mismatch_mode` parameter controls behavior:
When merging models with different layer topologies:
| Mode | Behavior |
|------|----------|
| `skip` | Missing layers in B/C use A's values **(default)** |
| `zeros` | Missing layers in B/C are treated as zeros |
| `error` | Fail if any layer is missing |
| `skip` | Missing layers in B or C use Model A's values unchanged **(default)** |
| `zeros` | Missing layers in B or C are treated as zero tensors |
| `error` | Raise an error and abort if any layer is missing across models |
**Special case for B - C subtraction (Add-Difference, Train-Difference):**
- If **both** B and C have the layer: Normal subtraction is performed
- If **either** B or C is missing the layer: The difference is treated as `None` (no contribution from this layer), and Model A's value is used
**Common scenarios:**
- Merging a fine-tuned model with added/removed layers
- Extracting style differences between models with slightly different architectures
- Combining partial LoRAs or adapters
**Note:** Extra layers in Models B/C that don't exist in Model A are currently ignored. The output will always have the same layer structure as Model A.
**Special rule for $(B - C)$ difference operations (Add-Difference, Train-Difference):**
- If **both** B and C have the layer: Normal difference $(B - C)$ is computed and applied to A.
- If **either** B or C is missing the layer: The difference is treated as zero, and Model A's value is preserved untouched.
---
## Layer Filtering (Regex Patterns)
## Dtype Preservation & Override
Use regex patterns to control which layers are merged:
- **Preserve Higher Precision (default, `override_dtype = False`)**: If a key is `float32` in any source model (A, B, or C), it is saved as `float32` regardless of `save_dtype`.
- **Explicit Override (`override_dtype = True`)**: Forces all floating tensors strictly to `save_dtype`.
### `exclude_patterns`
Layers matching any pattern will **keep Model A's values only** (no merge).
---
### `discard_patterns`
Layers matching any pattern will be **removed entirely** from the output.
## Layer Filtering
**Pattern format:**
- Whitespace-separated regex patterns
- Patterns use **substring matching** (not full match)
- Example: `text_model lora` matches any key containing "text_model" OR "lora"
- **`exclude_patterns`**: Newline-separated patterns. Matching layers keep Model A's values untouched.
- **`include_mode`**: Inverts `exclude_patterns` into a whitelist: only matching layers receive the merge; all other layers keep Model A's values.
- **`discard_patterns`**: Newline-separated patterns. Matching layers are omitted completely from the output model.
- **`glob_patterns`**: Enables shell wildcards (`*`, `?`) when checked; uses regex substring matching when unchecked.
+89
View File
@@ -0,0 +1,89 @@
# Two-Model Similarity Analysis
These nodes compare two safetensors models without merging them or writing an
output file. Separate nodes are provided for checkpoints, diffusion models,
text encoders, LoRAs, and embeddings.
`exclude_patterns` accepts one regex per line, or glob patterns when
`glob_patterns` is enabled. By default, matches are excluded from metrics and
topology counts. Enable `include_mode` to analyze only matches instead; an empty
include filter selects no tensors. The same field and matching syntax are used
in both modes.
## Outputs
`comparison_report` contains conventional numerical differences and topology
findings. Both reports use Markdown headings and tables. `cwb_report` separately contains diagnostics derived from the
similarity, alignment, and consensus stages of Consensus-Weighted Blending.
`documentation` contains this reference.
Two CSV-text outputs follow the existing three outputs, preserving their socket
positions: `layerwise_metrics_csv` and `layerwise_cwb_csv`. Connect them to a
text-saving node and use a `.csv` filename; analysis itself writes no files.
The CSV rows follow the same key order and displayed numeric precision as the
layerwise Markdown tables. Dtypes A/B, cosine mean/min/max, and A/B affinities
have separate columns. Undefined values remain `N/A`; an empty analysis produces
headers without data rows. Documentation-only mode returns empty CSV strings.
The reports begin with a global summary, followed by inferred blockwise values
and one entry for every comparable tensor. Block names are approximated from
the tensor-key hierarchy through its first numeric path component; they are not
an architecture-specific interpretation.
## Standard metrics
- MAE, MSE, and RMSE describe absolute error at the original weight scale.
- Maximum absolute difference identifies the largest individual change.
- Relative L2 is `||A-B|| / ((||A||+||B||)/2)` and is scale-aware.
- Cosine similarity compares direction; Pearson correlation compares centered
variation.
- Exact equality, sign agreement, norm ratio, and finite-value coverage expose
differences hidden by averages.
The global section reports both parameter-weighted metrics and equal-layer
averages. No arbitrary combined similarity score is assigned. Undefined
zero-norm or constant-tensor statistics are shown as `N/A`.
Only shared, same-shape floating tensors enter numerical aggregates. Missing
keys and shape mismatches are listed separately and are never padded, cropped,
or replaced with zeros. Non-floating tensors receive an exact-equality check.
Quantized models carrying ComfyUI quantization metadata are rejected because
encoded storage values do not represent directly comparable model weights.
## CWB diagnostics
The CWB report provides pairwise row-cosine statistics and each input's affinity
to both mean and median consensus. It stops before CWB's merge-weighting stage:
no thresholds, contribution weights, diversity weighting, vector
reconstitution, norm rescaling, or output writer are used.
Checkpoint, diffusion-model, and text-encoder tensors always retain their fixed
coordinate alignment. Embedding and LoRA nodes expose one optional
`cwb_similarity_alignment` control. When enabled, it performs
threshold-independent greedy one-to-one matching derived from CWB and reports
coverage, matched-score statistics, and improvement over index alignment. This
operation is quadratic and may be slow for large embeddings.
LoRA down/up tensors are treated as one logical layer. Rank-component matching
uses the product of down and up cosine similarity, including CWB's paired-sign
equivalence for comparison. No source tensor is modified.
## Memory and failure handling
Inputs always use Unified Efficient Loader's low-memory streaming mode. Matching
A/B tensors are loaded concurrently through bounded `async_stream()` pipelines
with `batch_size=1` and `prefetch_batches=1`, but no subsequent logical layer is
requested early. A complete LoRA down/up pair is loaded once and used for both
reports. CUDA inputs use pinned transfer while retaining the current CPU source
tensors for fallback.
Calculations use FP32. When CUDA processing runs out of memory, CUDA memory is
cleared and the current logical layer is retried once on CPU. After success or
failure, all tensor references for that layer are released and every loaded key
is marked processed exactly once before advancing. Both reports list every CPU
fallback. A failed CPU retry or any non-OOM error stops analysis with the layer
name in the error.
`top_weight_differences` bounds the scalar-coordinate list; every comparable
tensor is still present in the layerwise report. `DOCUMENTATION ONLY` loads no
model data.
+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
}
+57
View File
@@ -0,0 +1,57 @@
"""Bounded progressive partial SVD for adaptive extraction modes."""
import torch
ADAPTIVE_PARTIAL_MODES = {"ratio", "quantile", "sv_cumulative", "sv_fro"}
def _bounded_rank(singular_values, mode, parameter, output_cap, matrix_fro_sq):
capped = singular_values[:output_cap]
if capped.numel() == 0 or capped[0] < 1e-8:
return 1, True
if mode == "ratio":
rank = int(torch.sum(capped > capped[0] / parameter).item()) if parameter > 0 else output_cap
return max(1, min(rank, output_cap)), rank < output_cap
if mode in {"quantile", "sv_cumulative"}:
total = capped.sum()
if total < 1e-8:
return 1, True
rank = int(torch.searchsorted(torch.cumsum(capped, 0) / total, parameter).item()) + 1
return max(1, min(rank, output_cap)), rank < output_cap
target_sq = float(parameter) ** 2 * matrix_fro_sq
cumulative = torch.cumsum(capped.square(), 0)
reached = bool(cumulative[-1] >= target_sq)
if reached:
rank = int(torch.searchsorted(cumulative, target_sq).item()) + 1
return max(1, min(rank, output_cap)), rank < output_cap
return output_cap, False
def adaptive_partial_svd(weight, mode, parameter, max_rank, probe_offset, niter):
"""Return a bounded top spectrum and selected output rank."""
max_possible = min(weight.shape)
output_cap = max(1, min(int(max_rank), max_possible))
offset = max(1, int(probe_offset))
initial_probe = min(max_possible, output_cap + offset)
probe_limit = min(max_possible, max(initial_probe, output_cap * 2))
matrix_fro_sq = float(torch.sum(weight.square()).item()) if mode == "sv_fro" else 0.0
probe_rank = initial_probe
while True:
if probe_rank >= max_possible:
U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
else:
U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter)
Vh = V.T
rank, resolved_before_cap = _bounded_rank(
S, mode, parameter, output_cap, matrix_fro_sq
)
near_probe_tail = rank >= max(1, probe_rank - offset)
if resolved_before_cap or not near_probe_tail or probe_rank >= probe_limit:
return U[:, :rank], S[:rank], Vh[:rank, :], rank
probe_rank = probe_limit
+36
View File
@@ -0,0 +1,36 @@
import os
import folder_paths
CANONICAL_MODEL_CATEGORIES = frozenset(
{"checkpoints", "diffusion_models", "text_encoders", "loras", "embeddings"}
)
def canonical_model_artifact_path(category: str, output_filename: str) -> tuple[str, str]:
"""Return the canonical write path and category-relative loader value."""
if category not in CANONICAL_MODEL_CATEGORIES:
raise ValueError(f"Unsupported model artifact category: {category}")
relative_name = output_filename.strip().replace("\\", "/")
if not relative_name:
raise ValueError("Output filename must not be empty.")
if relative_name.startswith("/") or os.path.isabs(relative_name):
raise ValueError("Output filename must be relative to its model category.")
while relative_name.lower().endswith(".safetensors"):
relative_name = relative_name[:-12]
if not relative_name:
raise ValueError("Output filename must include a name before its extension.")
relative_name = f"{relative_name}.safetensors"
category_root = os.path.abspath(os.path.join(folder_paths.models_dir, category))
output_path = os.path.abspath(
os.path.join(category_root, *relative_name.split("/"))
)
if os.path.commonpath((category_root, output_path)) != category_root:
raise ValueError("Output filename must stay inside its model category directory.")
os.makedirs(os.path.dirname(output_path), exist_ok=True)
output_name = os.path.relpath(output_path, category_root).replace(os.sep, "/")
return output_path, output_name
File diff suppressed because it is too large Load Diff
+768
View File
@@ -0,0 +1,768 @@
import fnmatch
import gc
import re
from dataclasses import dataclass, field
import comfy.utils
import folder_paths
import torch
from comfy_api.latest import io
from tqdm import tqdm
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .artifact_paths import canonical_model_artifact_path
from .consensus_merger import (
CWB_CONFIG,
CWBSettings,
CWBDiagnostics,
DENSE_CWB_PRESETS,
merge_cwb_tensors,
resolve_cwb_settings,
)
from .device_utils import (
cleanup_after_operation,
estimate_model_size,
prepare_for_large_operation,
)
from .lora_alpha import normalize_lora_pair
from .lora_resize import (
build_lora_layer_map,
canonical_lora_key,
parse_lora_layers,
select_output_dtype,
validate_canonical_blocks,
)
from .quantization_guard import inspect_low_bit_input
from .uel_io import AsyncTensorCursor, atomic_uel_writer
_OPERATION = "Delta CWB LoRA Merge"
_ALLOWED_ROLES = {
"down", "up", "alpha", "diff", "format", "down_suffix", "up_suffix",
"alpha_suffix",
}
_FACTOR_SUFFIXES = (
".lora_A.default.weight", ".lora_B.default.weight",
".lora_down.weight", ".lora_up.weight",
"_lora.down.weight", "_lora.up.weight",
".lora_A.weight", ".lora_B.weight",
".lora.down.weight", ".lora.up.weight",
".lora_A", ".lora_B",
".lora_linear_layer.down.weight", ".lora_linear_layer.up.weight",
)
_UNSUPPORTED_SUFFIXES = (
".lora_mid.weight", ".reshape_weight", ".dora_scale", ".w_norm",
".b_norm", ".set_weight",
)
@dataclass(frozen=True)
class DeltaCWBLoRASource:
source_index: int
block_name: str
down_key: str | None
up_key: str | None
alpha_key: str | None
diff_key: str | None
source_dtype: torch.dtype
@property
def keys(self):
return tuple(
key for key in (self.down_key, self.up_key, self.alpha_key, self.diff_key)
if key is not None
)
@dataclass(frozen=True)
class DeltaCWBLayerUnit:
core: str
output_key: str
sources: tuple[DeltaCWBLoRASource, ...]
zero_contributors: int
preserve_anchor: bool
@dataclass
class DeltaCWBReport:
merged_layers: int = 0
anchor_preserved_layers: int = 0
zero_filled_contributors: int = 0
rejected_unsupported_layers: int = 0
cuda_cpu_fallbacks: int = 0
output_dtypes: dict[str, int] = field(default_factory=dict)
def record_dtype(self, dtype):
name = str(dtype).removeprefix("torch.")
self.output_dtypes[name] = self.output_dtypes.get(name, 0) + 1
def render(self, diagnostics: CWBDiagnostics, preset_name: str, input_count: int, custom: bool):
dtype_summary = ", ".join(
f"{name}={count}" for name, count in sorted(self.output_dtypes.items())
) or "none"
source = "connected custom configuration" if custom else preset_name
return "\n".join([
"DELTA CWB LORA MERGE SUMMARY",
f"Settings: {source}",
f"Input contributors: {input_count}",
f"Merged layers: {self.merged_layers}",
f"Anchor-preserved layers: {self.anchor_preserved_layers}",
f"Explicit zero contributors: {self.zero_filled_contributors}",
f"Rejected/unsupported layers: {self.rejected_unsupported_layers}",
f"Accepted weighting contributors: {diagnostics.accepted_weighting_contributors}/{diagnostics.weighting_contributors}",
f"All-rejected equal fallbacks: {diagnostics.all_rejected_fallbacks}",
f"Zero-weight equal fallbacks: {diagnostics.zero_weight_fallbacks}",
f"CUDA OOM CPU fallbacks: {self.cuda_cpu_fallbacks}",
f"Output dtypes: {dtype_summary}",
])
def _compile_patterns(patterns, glob_patterns):
entries = [entry.strip() for entry in patterns.splitlines() if entry.strip()]
if glob_patterns:
return tuple(entries)
return tuple(re.compile(entry) for entry in entries)
def _matches_pattern(key, patterns, glob_patterns):
if glob_patterns:
return any(fnmatch.fnmatchcase(key, pattern) for pattern in patterns)
return any(pattern.search(key) is not None for pattern in patterns)
def _factor_delta(down, up, layer):
if down.ndim == 2 and up.ndim == 2:
if down.shape[0] != up.shape[1]:
raise ValueError(f"Delta CWB rank mismatch for '{layer}'.")
return up.float() @ down.float()
if down.ndim == 4 and up.ndim == 4 and tuple(up.shape[2:]) == (1, 1):
rank, channels, height, width = down.shape
if up.shape[1] != rank:
raise ValueError(f"Delta CWB rank mismatch for '{layer}'.")
delta = up.reshape(up.shape[0], rank).float() @ down.reshape(rank, -1).float()
return delta.reshape(up.shape[0], channels, height, width)
raise ValueError(
f"Delta CWB layer '{layer}' requires linear factors or convolutional "
"factors with a 1x1 up kernel."
)
def _factorize_delta(delta, rank):
if delta.ndim not in (2, 4):
return None
shape = delta.shape
matrix = delta if delta.ndim == 2 else delta.reshape(shape[0], -1)
target_rank = max(1, min(int(rank), min(matrix.shape)))
if target_rank == min(matrix.shape):
if matrix.shape[0] <= matrix.shape[1]:
up = torch.eye(matrix.shape[0], device=matrix.device, dtype=matrix.dtype)
down = matrix
else:
up = matrix
down = torch.eye(matrix.shape[1], device=matrix.device, dtype=matrix.dtype)
else:
up, singular_values, right = torch.svd_lowrank(matrix, q=target_rank, niter=2)
up = up * singular_values.unsqueeze(0)
down = right.t()
if delta.ndim == 4:
down = down.reshape(down.shape[0], *shape[1:])
up = up.unsqueeze(-1).unsqueeze(-1)
return down, up
def _inspect_inputs(handlers):
infos = []
for index, handler in enumerate(handlers):
low_bit = inspect_low_bit_input(handler, f"LoRA {index + 1}", _OPERATION)
keys = list(handler.keys())
pairs, passthrough = parse_lora_layers(keys)
unsupported_passthrough = [
key for key in passthrough if key.endswith(_UNSUPPORTED_SUFFIXES)
]
if unsupported_passthrough:
raise ValueError(
f"[{_OPERATION}] LoRA {index + 1} contains unsupported companion "
f"tensor '{unsupported_passthrough[0]}'."
)
incomplete = [key for key in passthrough if key.endswith(_FACTOR_SUFFIXES)]
if incomplete:
raise ValueError(
f"[{_OPERATION}] LoRA {index + 1} contains an incomplete factor "
f"pair at '{incomplete[0]}'."
)
infos.append({"pairs": pairs, "passthrough": passthrough, "low_bit": low_bit})
return infos
def _build_layer_units(
layer_map,
infos,
handlers,
mismatch_mode,
exclude_patterns,
discard_patterns,
glob_patterns,
include_mode=False,
):
units = []
stream_keys = [[] for _ in handlers]
rejected = 0
for core, mapped in layer_map.items():
if not any(index == 0 for index, _ in mapped):
continue
output_key = canonical_lora_key(core, "diff")
if _matches_pattern(output_key, discard_patterns, glob_patterns):
rejected += 1
continue
matched = _matches_pattern(output_key, exclude_patterns, glob_patterns)
preserve = not matched if include_mode else matched
candidates = []
shapes = {}
for index, block in sorted(mapped):
roles = infos[index]["pairs"][block]
unsupported = set(roles) - _ALLOWED_ROLES
if unsupported:
raise ValueError(
f"[{_OPERATION}] Layer '{block}' has unsupported roles: "
f"{', '.join(sorted(unsupported))}."
)
if "down" in roles or "up" in roles:
if not all(role in roles for role in ("down", "up")):
raise ValueError(f"[{_OPERATION}] Layer '{block}' has an incomplete factor pair.")
down_key, up_key = roles["down"], roles["up"]
if down_key in infos[index]["low_bit"] or up_key in infos[index]["low_bit"]:
raise ValueError(f"[{_OPERATION}] Cannot expand low-bit factors for '{block}'.")
down_shape = tuple(handlers[index].get_shape(down_key))
up_shape = tuple(handlers[index].get_shape(up_key))
if len(down_shape) == 2 and len(up_shape) == 2:
if down_shape[0] != up_shape[1]:
raise ValueError(f"[{_OPERATION}] Rank mismatch for '{block}'.")
dense_shape = (up_shape[0], down_shape[1])
elif len(down_shape) == 4 and len(up_shape) == 4 and up_shape[2:] == (1, 1):
if down_shape[0] != up_shape[1]:
raise ValueError(f"[{_OPERATION}] Rank mismatch for '{block}'.")
dense_shape = (up_shape[0], *down_shape[1:])
else:
raise ValueError(
f"[{_OPERATION}] Layer '{block}' requires linear factors or "
"convolutional factors with a 1x1 up kernel."
)
factor_dtypes = (
handlers[index].get_dtype(down_key),
handlers[index].get_dtype(up_key),
)
source_dtype = torch.float32 if torch.float32 in factor_dtypes else factor_dtypes[1]
source = DeltaCWBLoRASource(
index, block, down_key, up_key, roles.get("alpha"), None, source_dtype
)
elif "diff" in roles:
diff_key = roles["diff"]
if diff_key in infos[index]["low_bit"]:
raise ValueError(f"[{_OPERATION}] Cannot merge low-bit delta '{block}'.")
dense_shape = tuple(handlers[index].get_shape(diff_key))
source = DeltaCWBLoRASource(
index, block, None, None, None, diff_key,
handlers[index].get_dtype(diff_key),
)
else:
raise ValueError(f"[{_OPERATION}] Layer '{block}' has no supported representation.")
candidates.append(source)
shapes[index] = dense_shape
anchor = next((source for source in candidates if source.source_index == 0), None)
if anchor is None:
continue
anchor_shape = shapes[0]
compatible = [
source for source in candidates if shapes[source.source_index] == anchor_shape
]
present = {source.source_index for source in compatible}
mismatch_count = len(handlers) - len(present)
if preserve:
sources = (anchor,)
zero_count = 0
elif mismatch_count and mismatch_mode == "error":
raise ValueError(
f"[{_OPERATION}] Layer '{core}' is missing or incompatible in one or more inputs."
)
elif mismatch_count and mismatch_mode == "skip":
sources = (anchor,)
zero_count = 0
preserve = True
else:
sources = tuple(compatible)
zero_count = mismatch_count if mismatch_mode == "zeros" else 0
for source in sources:
stream_keys[source.source_index].extend(source.keys)
units.append(DeltaCWBLayerUnit(core, output_key, sources, zero_count, preserve))
return units, stream_keys, rejected
def _to_processing_device(tensor, device):
if device.type == "cuda":
return transfer_to_gpu_pinned(tensor, device, torch.float32)
return tensor.to(device=device, dtype=torch.float32)
def _process_layer_on_device(
unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics,
output_representation, output_rank,
):
deltas = []
try:
for source in unit.sources:
if source.diff_key is not None:
delta = _to_processing_device(
loaded[(source.source_index, source.diff_key)], device
)
else:
down = loaded[(source.source_index, source.down_key)]
up = loaded[(source.source_index, source.up_key)]
alpha = (
loaded[(source.source_index, source.alpha_key)]
if source.alpha_key is not None else None
)
down, up = normalize_lora_pair(down, up, alpha, layer=source.block_name)
delta = _factor_delta(
_to_processing_device(down, device),
_to_processing_device(up, device),
source.block_name,
)
deltas.append(delta)
if unit.zero_contributors:
deltas.extend(torch.zeros_like(deltas[0]) for _ in range(unit.zero_contributors))
if unit.preserve_anchor:
merged = deltas[0].clone()
else:
merged = merge_cwb_tensors(
deltas,
settings,
reference_index=0,
allow_similarity_alignment=False,
diagnostics=diagnostics,
)
dtype = select_output_dtype(
[source.source_dtype for source in unit.sources],
save_dtype,
is_1d_diff=merged.ndim == 1,
)
if output_representation == "factorized_lora":
factors = _factorize_delta(merged, output_rank)
if factors is not None:
down, up = factors
return [
(canonical_lora_key(unit.core, "down"), down.to(dtype=dtype, device="cpu").contiguous()),
(canonical_lora_key(unit.core, "up"), up.to(dtype=dtype, device="cpu").contiguous()),
], dtype
return [(unit.output_key, merged.to(dtype=dtype, device="cpu").contiguous())], dtype
finally:
deltas.clear()
def _is_cuda_oom(error):
return isinstance(error, torch.cuda.OutOfMemoryError) or (
isinstance(error, RuntimeError) and "out of memory" in str(error).lower()
)
def _release_cuda_oom():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def _process_layer(
unit, loaded, settings, save_dtype, requested_device, diagnostics,
output_representation="full_difference", output_rank=384,
):
device = torch.device(requested_device)
attempt_diagnostics = CWBDiagnostics()
try:
output, dtype = _process_layer_on_device(
unit, loaded, settings, save_dtype, device, attempt_diagnostics,
output_representation, output_rank,
)
diagnostics.absorb_runtime(attempt_diagnostics)
return output, dtype, False
except Exception as error:
if device.type != "cuda" or not _is_cuda_oom(error):
raise
_release_cuda_oom()
retry_diagnostics = CWBDiagnostics()
output, dtype = _process_layer_on_device(
unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics,
output_representation, output_rank,
)
diagnostics.absorb_runtime(retry_diagnostics)
return output, dtype, True
class DeltaCWBLoRAMergerLogic:
@staticmethod
def execute(
lora_paths,
source_names,
settings,
preset_name,
custom_settings,
mismatch_mode,
output_filename,
save_dtype,
output_representation,
output_rank,
process_device,
exclude_patterns,
discard_patterns,
glob_patterns,
force_clear_cache,
include_mode=False,
):
prepare_for_large_operation(
sum(estimate_model_size(path) for path in lora_paths) * 2.5,
torch.device(process_device),
)
handlers = []
cursors = []
report = DeltaCWBReport()
diagnostics = CWBDiagnostics()
try:
for path in lora_paths:
handlers.append(MemoryEfficientSafeOpen(path, low_memory=True))
infos = _inspect_inputs(handlers)
layer_map = build_lora_layer_map(infos, None)
validate_canonical_blocks(layer_map, _OPERATION)
compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
units, stream_keys, rejected = _build_layer_units(
layer_map,
infos,
handlers,
mismatch_mode,
compiled_excludes,
compiled_discards,
glob_patterns,
include_mode,
)
report.rejected_unsupported_layers = rejected
cursors = [
AsyncTensorCursor(
handler,
keys,
pin_memory=torch.device(process_device).type == "cuda",
)
for handler, keys in zip(handlers, stream_keys)
]
output_path, output_name = canonical_model_artifact_path("loras", output_filename)
metadata = (handlers[0].metadata() or {}).copy()
metadata.update({
"merge_method": "delta_cwb",
"source_models": ",".join(source_names),
"alpha_normalized": "true",
"output_representation": (
"factorized_lora" if output_representation == "factorized_lora"
else "full_direct_difference"
),
"cwb_consensus_type": settings.consensus_type,
"cwb_alignment_method": settings.alignment_method,
"cwb_alignment_threshold": str(settings.alignment_threshold),
"cwb_similarity_threshold": str(settings.similarity_threshold),
"cwb_power_alpha": str(settings.power_alpha),
"cwb_diversity_beta": str(settings.diversity_beta),
"cwb_rescale_norm": str(settings.rescale_norm).lower(),
"cwb_global_scale": str(settings.global_scale),
"cwb_dynamic_similarity_contrast": str(settings.dynamic_similarity_contrast).lower(),
"cwb_soft_comfort_bandpass": str(settings.soft_comfort_bandpass).lower(),
"cwb_position_weight": str(settings.position_weight),
"cwb_preserve_common_prefix": str(settings.preserve_common_prefix).lower(),
})
pbar = comfy.utils.ProgressBar(len(units))
with atomic_uel_writer(output_path, metadata) as writer, torch.no_grad():
for unit in tqdm(units, desc="Delta CWB LoRA layers", unit="layers"):
if force_clear_cache:
_release_cuda_oom()
loaded = {}
try:
for source in unit.sources:
for key in source.keys:
loaded[(source.source_index, key)] = cursors[
source.source_index
].take(key)
output_items, dtype, used_cpu = _process_layer(
unit,
loaded,
settings,
save_dtype,
process_device,
diagnostics,
output_representation,
output_rank,
)
writer.write_batch(output_items)
report.merged_layers += int(not unit.preserve_anchor)
report.anchor_preserved_layers += int(unit.preserve_anchor)
report.zero_filled_contributors += unit.zero_contributors
report.cuda_cpu_fallbacks += int(used_cpu)
report.record_dtype(dtype)
del output_items
finally:
for source_index, key in list(loaded):
del loaded[(source_index, key)]
cursors[source_index].release(key)
loaded.clear()
pbar.update(1)
for cursor in cursors:
cursor.finish()
return output_name, report.render(
diagnostics, preset_name, len(lora_paths), custom_settings
)
finally:
for cursor in cursors:
cursor.close()
for handler in handlers:
handler.__exit__(None, None, None)
cleanup_after_operation()
@staticmethod
def execute_node(kwargs, count):
names = [kwargs[f"lora_{index}"] for index in range(1, count + 1)]
paths = [folder_paths.get_full_path_or_raise("loras", name) for name in names]
settings = resolve_cwb_settings(
kwargs["cwb_preset"], DENSE_CWB_PRESETS, kwargs.get("cwb_config")
)
requested_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}[kwargs["save_dtype"]]
filename, report = DeltaCWBLoRAMergerLogic.execute(
paths,
names,
settings,
kwargs["cwb_preset"],
kwargs.get("cwb_config") is not None,
kwargs["mismatch_mode"],
kwargs["output_filename"],
requested_dtype,
kwargs["output_representation"],
kwargs["output_rank"],
kwargs["process_device"],
kwargs["exclude_patterns"],
kwargs["discard_patterns"],
kwargs["glob_patterns"],
kwargs["force_clear_cache"],
kwargs.get("include_mode", False),
)
return io.NodeOutput(filename, report)
def _fixed_inputs(count, default_filename):
options = folder_paths.get_filename_list("loras")
inputs = []
if count:
inputs.append(io.Combo.Input(
"lora_1",
options=options,
tooltip="Metadata and layer anchor. Defines saved layers and preserved values.",
))
inputs.extend(
io.Combo.Input(
f"lora_{index}",
options=options,
tooltip="Equal-prior LoRA contributor. CWB determines its effective row influence.",
)
for index in range(2, count + 1)
)
inputs.extend([
io.Combo.Input(
"cwb_preset",
options=list(DENSE_CWB_PRESETS),
default="balanced_mean",
tooltip="Controls how strongly agreement, disagreement, and row magnitude affect the saved full weight changes. A connected CWB Config overrides it completely.",
),
CWB_CONFIG.Input(
"cwb_config",
optional=True,
tooltip="Optional complete settings override from CWB Custom Configuration.",
),
io.Combo.Input(
"mismatch_mode",
options=["skip", "zeros", "error"],
default="skip",
tooltip="For a missing or incompatible anchored layer: preserve LoRA 1, include an explicit zero contributor, or abort.",
),
io.String.Input(
"output_filename",
default=default_filename,
tooltip="Relative filename for the atomically saved full-difference LoRA.",
),
io.Combo.Input(
"save_dtype",
options=["fp32", "fp16", "bf16"],
default="bf16",
tooltip="Requested saved precision. Participating FP32 inputs preserve FP32.",
),
io.Combo.Input(
"output_representation",
options=["full_difference", "factorized_lora"],
default="full_difference",
tooltip="Save exact full differences or factorize each merged layer directly into compact LoRA factors.",
),
io.Int.Input(
"output_rank",
default=384,
min=1,
max=16384,
tooltip="Maximum rank used when Output Representation is factorized_lora.",
),
io.Combo.Input(
"process_device",
options=["cuda", "cpu"],
default="cuda",
tooltip="Processing device. Only a failed CUDA layer retries on CPU.",
),
io.String.Input(
"exclude_patterns",
default="",
multiline=True,
tooltip="Preserve matching canonical full-difference layers from LoRA 1.",
),
io.String.Input(
"discard_patterns",
default="",
multiline=True,
tooltip="Omit matching canonical full-difference layers from the output.",
),
io.Boolean.Input(
"glob_patterns",
default=False,
tooltip="Use shell-style glob patterns instead of regular expressions.",
),
io.Boolean.Input(
"force_clear_cache",
default=False,
tooltip="Clear Python and CUDA caches before each layer to reduce retained memory.",
),
io.Boolean.Input(
"include_mode",
default=False,
tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from LoRA 1.",
),
])
return inputs
def _multi_inputs(default_filename):
options = folder_paths.get_filename_list("loras")
optional = ["None", *options]
inputs = [
io.Combo.Input(
"lora_count",
options=[str(index) for index in range(2, 9)],
default="2",
tooltip="Number of consecutive selected LoRAs to merge.",
),
io.Combo.Input(
"lora_1",
options=options,
tooltip="Metadata and layer anchor. Defines saved layers and preserved values.",
),
io.Combo.Input(
"lora_2",
options=options,
tooltip="Second equal-prior LoRA contributor.",
),
]
inputs.extend(
io.Combo.Input(
f"lora_{index}",
options=optional,
default="None",
tooltip="Additional equal-prior contributor; required when included by LoRA Count.",
)
for index in range(3, 9)
)
inputs.extend(_fixed_inputs(0, default_filename))
return inputs
class DeltaCWBLoRATwoMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DeltaCWBLoRATwoMerger",
display_name="Delta CWB Merge LoRAs (2)",
category="ModelUtils/LoRA/Merge/Delta CWB",
description="Expand two LoRAs to full weight changes and merge corresponding rows with CWB.",
inputs=_fixed_inputs(2, "delta_cwb_merged_2_lora"),
outputs=[
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="cwb_report"),
],
is_output_node=True,
is_experimental=True,
)
@classmethod
def execute(cls, **kwargs):
return DeltaCWBLoRAMergerLogic.execute_node(kwargs, 2)
class DeltaCWBLoRAThreeMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DeltaCWBLoRAThreeMerger",
display_name="Delta CWB Merge LoRAs (3)",
category="ModelUtils/LoRA/Merge/Delta CWB",
description="Expand three LoRAs to full weight changes and merge corresponding rows with CWB.",
inputs=_fixed_inputs(3, "delta_cwb_merged_3_lora"),
outputs=[
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="cwb_report"),
],
is_output_node=True,
is_experimental=True,
)
@classmethod
def execute(cls, **kwargs):
return DeltaCWBLoRAMergerLogic.execute_node(kwargs, 3)
class DeltaCWBLoRAMultiMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DeltaCWBLoRAMultiMerger",
display_name="Delta CWB LoRA Multi-Merge",
category="ModelUtils/LoRA/Merge/Delta CWB",
description="Expand 2-8 LoRAs to full weight changes and merge corresponding rows with CWB.",
inputs=_multi_inputs("delta_cwb_merged_multi_lora"),
outputs=[
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="cwb_report"),
],
is_output_node=True,
is_experimental=True,
)
@classmethod
def execute(cls, **kwargs):
count = int(kwargs["lora_count"])
missing = [
index for index in range(1, count + 1)
if not kwargs[f"lora_{index}"] or kwargs[f"lora_{index}"] == "None"
]
if missing:
raise ValueError(
f"LoRA Count includes unselected input(s): {', '.join(map(str, missing))}."
)
return DeltaCWBLoRAMergerLogic.execute_node(kwargs, count)
DELTA_CWB_LORA_NODES = [
DeltaCWBLoRATwoMerger,
DeltaCWBLoRAThreeMerger,
DeltaCWBLoRAMultiMerger,
]
+790
View File
@@ -0,0 +1,790 @@
"""Standalone training-free Consensus-Weighted Blending model grafting node."""
from __future__ import annotations
from contextlib import closing
import json
import re
import comfy.utils
import folder_paths
import torch
import torch.nn.functional as F
from comfy_api.latest import io
from tqdm import tqdm
from unifiedefficientloader import MemoryEfficientSafeOpen
from .artifact_paths import canonical_model_artifact_path
from .consensus_merger import (
CWB_CONFIG,
CWBDiagnostics,
CWBSettings,
DENSE_CWB_PRESETS,
_clear_previous_layer,
_is_float_dtype,
_merge_tensors_to_cpu,
_requested_dtype,
resolve_cwb_settings,
)
from .device_utils import (
cleanup_after_operation,
estimate_model_size,
prepare_for_large_operation,
)
from .merger import (
_compile_patterns,
_matches_any_pattern,
load_documentation_from_file,
)
from .quantization_guard import (
DiffusionQuantization,
diffusion_key_map,
write_preserved_tensor,
)
from .uel_io import atomic_uel_writer, stream_work_units
_MAIN_BLOCK_INDEX_RE = re.compile(
r"(?:^|\.)(?:layers|blocks|double_blocks|single_blocks|transformer_blocks)\.(\d+)\."
)
_ATTN_KEY_TOKENS = (
".qkv.",
".qkv_proj.",
".proj.q.",
".proj.k.",
".proj.v.",
".proj.o.",
".attn.",
".self_attn.",
".q_norm.",
".k_norm.",
".to_q.",
".to_k.",
".to_v.",
".to_out.",
)
_FFN_KEY_TOKENS = (
".ffn_up_gate.",
".ffn_up.",
".ffn_gate.",
".ffn_down.",
".mlp.",
".ff.",
".fc1.",
".fc2.",
".w1.",
".w2.",
".w3.",
)
def _extract_main_block_index(key: str) -> int | None:
if "text_embedder." in key or "token_refiner." in key:
return None
match = _MAIN_BLOCK_INDEX_RE.search(key)
return int(match.group(1)) if match else None
def _classify_graft_operator(key: str) -> str | None:
padded = f".{key}."
if any(token in padded for token in _ATTN_KEY_TOKENS):
return "attention"
if any(token in padded for token in _FFN_KEY_TOKENS):
return "ffn"
return None
def _block_rule_allows(
block_idx: int | None,
rule: str,
protect_early_blocks: int,
max_block_index: int,
) -> bool:
if rule == "none":
return False
if block_idx is None:
return False
if block_idx < protect_early_blocks:
return False
if rule == "all_after_protected":
return True
if rule == "interleaved_even":
return (block_idx % 2) == 0
if rule == "interleaved_odd":
return (block_idx % 2) == 1
if rule == "back_half_only":
midpoint = (max_block_index + 1) // 2
return block_idx >= max(protect_early_blocks, midpoint)
return False
def _parse_custom_key_map(raw_map: str) -> dict[str, str]:
"""Parse lines of 'donor_key -> base_key' or 'donor_key = base_key' into {base_key: donor_key}."""
mapping: dict[str, str] = {}
if not raw_map or not raw_map.strip():
return mapping
for line_num, raw_line in enumerate(raw_map.splitlines(), start=1):
line = raw_line.strip()
if not line or line.startswith("#"):
continue
if "->" in line:
donor_part, base_part = line.split("->", 1)
elif "=" in line:
donor_part, base_part = line.split("=", 1)
else:
raise ValueError(
f"Custom key map line {line_num}: expected 'donor_key -> base_key'."
)
donor_key = donor_part.strip()
base_key = base_part.strip()
if not donor_key or not base_key:
raise ValueError(
f"Custom key map line {line_num}: both donor_key and base_key must be non-empty."
)
mapping[base_key] = donor_key
return mapping
def _align_donor_trailing_shape(
base_tensor: torch.Tensor, donor_tensor: torch.Tensor
) -> tuple[torch.Tensor, bool]:
"""Align donor trailing dimensions to base_tensor so CWB similarity alignment can match rows."""
if donor_tensor.ndim != base_tensor.ndim:
if donor_tensor.numel() == base_tensor.numel():
return donor_tensor.reshape(base_tensor.shape), True
flat_donor = donor_tensor.reshape(1, 1, -1).to(dtype=torch.float32)
resized = F.interpolate(
flat_donor, size=base_tensor.numel(), mode="linear", align_corners=False
)
return resized.reshape(base_tensor.shape).to(dtype=donor_tensor.dtype), True
if donor_tensor.shape == base_tensor.shape:
return donor_tensor, False
if base_tensor.ndim <= 1:
return donor_tensor, True
if donor_tensor.shape[1:] == base_tensor.shape[1:]:
return donor_tensor, True
target_trailing = int(torch.tensor(base_tensor.shape[1:]).prod().item())
donor_2d = donor_tensor.reshape(donor_tensor.shape[0], 1, -1).to(dtype=torch.float32)
resized_2d = F.interpolate(
donor_2d, size=target_trailing, mode="linear", align_corners=False
)
reshaped = resized_2d.reshape(donor_tensor.shape[0], *base_tensor.shape[1:]).to(
dtype=donor_tensor.dtype
)
return reshaped, True
def _graft_pair_to_cpu(
base_tensor: torch.Tensor,
donor_tensor: torch.Tensor,
settings: CWBSettings,
device: str,
target_dtype: torch.dtype,
*,
operation_label: str,
diagnostics: CWBDiagnostics,
) -> torch.Tensor:
aligned_donor, shape_changed = _align_donor_trailing_shape(base_tensor, donor_tensor)
merged = _merge_tensors_to_cpu(
[base_tensor, aligned_donor],
settings,
device,
torch.float32,
reference_index=0,
allow_similarity_alignment=True,
operation_label=operation_label,
diagnostics=diagnostics,
)
if shape_changed or settings.rescale_norm:
base_norm = torch.linalg.vector_norm(base_tensor.to(dtype=torch.float32))
merged_norm = torch.linalg.vector_norm(merged)
if merged_norm.item() > 1e-8 and base_norm.item() > 0.0:
merged = merged * (base_norm / merged_norm)
return merged.to(dtype=target_dtype).cpu().contiguous()
def _graft_with_subpart(
key: str,
op_type: str,
subpart: str,
base_tensor: torch.Tensor,
donor_tensor: torch.Tensor,
settings: CWBSettings,
device: str,
target_dtype: torch.dtype,
diagnostics: CWBDiagnostics,
) -> torch.Tensor | None:
"""Return grafted tensor or None if the subpart rule keeps the base tensor untouched."""
padded = f".{key}."
if op_type == "attention":
is_qkv = any(tok in padded for tok in (".qkv.", ".qkv_proj."))
is_q_or_k = any(
tok in padded for tok in (".proj.q.", ".proj.k.", ".to_q.", ".to_k.", ".q_norm.", ".k_norm.")
)
is_v = any(tok in padded for tok in (".proj.v.", ".to_v."))
is_out = any(tok in padded for tok in (".proj.o.", ".out_proj.", ".to_out."))
has_known_attn_part = is_qkv or is_q_or_k or is_v or is_out
if has_known_attn_part and subpart == "output_only":
if not is_out:
return None
elif has_known_attn_part and subpart == "value_and_output":
if is_q_or_k:
return None
if is_qkv and base_tensor.ndim >= 1 and base_tensor.shape[0] % 3 == 0:
aligned_donor, _ = _align_donor_trailing_shape(base_tensor, donor_tensor)
if aligned_donor.shape[0] % 3 == 0:
q_a, k_a, v_a = base_tensor.chunk(3, dim=0)
_, _, v_b = aligned_donor.chunk(3, dim=0)
v_merged = _graft_pair_to_cpu(
v_a,
v_b,
settings,
device,
target_dtype,
operation_label=f"{key}[v]",
diagnostics=diagnostics,
)
return torch.cat(
[
q_a.to(dtype=target_dtype).cpu(),
k_a.to(dtype=target_dtype).cpu(),
v_merged,
],
dim=0,
).contiguous()
elif not (is_v or is_out):
return None
elif op_type == "ffn":
is_up_gate = ".ffn_up_gate." in padded
is_up = any(tok in padded for tok in (".ffn_up.", ".fc1.", ".w1."))
is_gate = any(tok in padded for tok in (".ffn_gate.", ".w3."))
is_down = any(tok in padded for tok in (".ffn_down.", ".fc2.", ".w2.", ".ff.net.2."))
has_known_ffn_part = is_up_gate or is_up or is_gate or is_down
if has_known_ffn_part and subpart == "down_only":
if not is_down:
return None
elif has_known_ffn_part and subpart == "up_only":
if is_down or is_gate:
return None
if is_up_gate and base_tensor.ndim >= 1 and base_tensor.shape[0] % 2 == 0:
aligned_donor, _ = _align_donor_trailing_shape(base_tensor, donor_tensor)
if aligned_donor.shape[0] % 2 == 0:
up_a, gate_a = base_tensor.chunk(2, dim=0)
up_b, _ = aligned_donor.chunk(2, dim=0)
up_merged = _graft_pair_to_cpu(
up_a,
up_b,
settings,
device,
target_dtype,
operation_label=f"{key}[up]",
diagnostics=diagnostics,
)
return torch.cat(
[
up_merged,
gate_a.to(dtype=target_dtype).cpu(),
],
dim=0,
).contiguous()
elif not is_up:
return None
return _graft_pair_to_cpu(
base_tensor,
donor_tensor,
settings,
device,
target_dtype,
operation_label=key,
diagnostics=diagnostics,
)
def _resolve_donor_qkv_plan(
base_key: str,
donor_keys_set: set[str],
custom_map: dict[str, str],
) -> tuple[str, list[str], int | None] | None:
"""Resolve donor key(s) for base_key.
Returns:
(mode, donor_logical_keys, slice_index)
mode is 'direct', 'fuse_qkv' (cat q,k,v), or 'split_qkv' (chunk 3 and pick slice_index).
"""
if base_key in custom_map:
mapped = custom_map[base_key]
if mapped in donor_keys_set:
return ("direct", [mapped], None)
return None
if base_key in donor_keys_set:
return ("direct", [base_key], None)
if base_key.endswith(".proj.qkv.weight"):
prefix = base_key[: -len("qkv.weight")]
parts = [f"{prefix}{p}.weight" for p in ("q", "k", "v")]
if all(k in donor_keys_set for k in parts):
return ("fuse_qkv", parts, None)
for idx, part in enumerate(("q", "k", "v")):
suffix = f".proj.{part}.weight"
if base_key.endswith(suffix):
fused = f"{base_key[:-len(suffix)]}.proj.qkv.weight"
if fused in donor_keys_set:
return ("split_qkv", [fused], idx)
return None
class CWBModelGraft(io.ComfyNode):
@classmethod
def define_schema(cls) -> io.Schema:
models = folder_paths.get_filename_list("diffusion_models")
presets = list(DENSE_CWB_PRESETS)
return io.Schema(
node_id="CWBModelGraft",
display_name="CWB Model Graft",
category="ModelUtils/Merging",
description="Graft selected blocks, sub-projections, or mapped layers from a donor diffusion model into a base diffusion model using Consensus-Weighted Blending.",
inputs=[
io.Combo.Input(
"execution_mode",
options=["MERGE", "DOCUMENTATION ONLY"],
default="MERGE",
tooltip="MERGE writes a grafted model; DOCUMENTATION ONLY opens no model files.",
),
io.Combo.Input(
"model_a",
options=models,
tooltip="Base diffusion model that anchors structure, metadata, and protected blocks.",
),
io.Combo.Input(
"model_b",
options=models,
tooltip="Donor diffusion model supplying parts to graft into Model A.",
),
io.Combo.Input(
"attention_preset",
options=presets,
default="robust_medn",
tooltip="CWB preset applied to grafted attention layers.",
),
CWB_CONFIG.Input(
"attention_config",
optional=True,
tooltip="Optional custom CWB configuration override for attention layers.",
),
io.Combo.Input(
"ffn_preset",
options=presets,
default="balanced_mean",
tooltip="CWB preset applied to grafted feed-forward (MLP) layers.",
),
CWB_CONFIG.Input(
"ffn_config",
optional=True,
tooltip="Optional custom CWB configuration override for feed-forward (MLP) layers.",
),
io.Int.Input(
"protect_early_blocks",
default=2,
min=0,
max=256,
step=1,
tooltip="Number of initial transformer blocks (starting from 0) kept 100% from Model A to preserve global layout.",
),
io.Combo.Input(
"attention_block_rule",
options=[
"interleaved_even",
"interleaved_odd",
"back_half_only",
"all_after_protected",
"none",
],
default="interleaved_even",
tooltip="Which transformer blocks after the protected start have their attention layers grafted.",
),
io.Combo.Input(
"ffn_block_rule",
options=[
"all_after_protected",
"interleaved_even",
"interleaved_odd",
"back_half_only",
"none",
],
default="all_after_protected",
tooltip="Which transformer blocks after the protected start have their feed-forward layers grafted.",
),
io.Combo.Input(
"attention_subpart",
options=["all", "value_and_output", "output_only"],
default="all",
tooltip="Select whether to graft all attention weights, only value and output projections, or only the output projection.",
),
io.Combo.Input(
"ffn_subpart",
options=["all", "up_only", "down_only"],
default="all",
tooltip="Select whether to graft all feed-forward weights, only the up projection, or only the down projection.",
),
io.String.Input(
"custom_key_map",
default="",
multiline=True,
tooltip="Optional explicit key mappings (one per line as 'donor_key -> base_key') to graft across differently named layers.",
),
io.String.Input(
"output_filename",
default="cwb_grafted_model",
tooltip="Filename without extension under ComfyUI's diffusion_models directory.",
),
io.Combo.Input(
"save_dtype",
options=["fp32", "fp16", "bf16"],
default="bf16",
tooltip="Requested dtype for grafted floating tensors.",
),
io.Combo.Input(
"process_device",
options=["cuda", "cpu"],
default="cuda",
tooltip="Per-tensor FP32 CWB computation device; CUDA OOM retries on CPU.",
),
io.String.Input(
"exclude_patterns",
default="",
multiline=True,
tooltip="Newline or whitespace separated patterns to preserve from Model A (or whitelist when include_mode is enabled).",
),
io.String.Input(
"discard_patterns",
default="",
multiline=True,
tooltip="Newline or whitespace separated patterns to omit from the output file.",
),
io.Boolean.Input(
"glob_patterns",
default=False,
tooltip="Interpret exclude and discard entries as shell-style globs instead of regular expressions.",
),
io.Boolean.Input(
"lazy_load",
default=True,
tooltip="Stream only one logical work unit at a time through UEL and release completed tensors immediately.",
),
io.Boolean.Input(
"force_clear_cache",
default=False,
tooltip="Run Python garbage collection and empty CUDA cache before each tensor.",
),
io.Boolean.Input(
"override_dtype",
default=False,
tooltip="Force grafted floating tensors to Save Dtype instead of preserving Model A's float precision.",
),
io.Boolean.Input(
"include_mode",
default=False,
tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are eligible for grafting; nonmatching layers are preserved from Model A.",
),
],
outputs=[
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
io.String.Output(display_name="cwb_report"),
],
is_experimental=True,
)
@classmethod
def execute(cls, **kwargs) -> io.NodeOutput:
documentation = load_documentation_from_file("cwb_model_graft.md")
if kwargs["execution_mode"] == "DOCUMENTATION ONLY":
return io.NodeOutput(
"Documentation mode active. No merge performed.",
documentation,
"CWB report unavailable because no merge was performed.",
)
filename, report = cls._run_graft(kwargs)
return io.NodeOutput(filename, documentation, report)
@classmethod
def _run_graft(cls, params: dict) -> tuple[str, str]:
model_names = [params["model_a"], params["model_b"]]
paths = []
for name in model_names:
path = folder_paths.get_full_path("diffusion_models", name)
if not path:
raise FileNotFoundError(f"diffusion_models input '{name}' was not found.")
paths.append(path)
total_size = sum(estimate_model_size(path) for path in paths)
device = params["process_device"]
if total_size:
prepare_for_large_operation(total_size * 1.2, torch.device(device))
custom_map = _parse_custom_key_map(params.get("custom_key_map", ""))
attn_settings = resolve_cwb_settings(
params["attention_preset"],
DENSE_CWB_PRESETS,
params.get("attention_config"),
)
ffn_settings = resolve_cwb_settings(
params["ffn_preset"],
DENSE_CWB_PRESETS,
params.get("ffn_config"),
)
diagnostics = CWBDiagnostics()
protect_early = int(params.get("protect_early_blocks", 2))
attn_rule = params.get("attention_block_rule", "interleaved_even")
ffn_rule = params.get("ffn_block_rule", "all_after_protected")
attn_subpart = params.get("attention_subpart", "all")
ffn_subpart = params.get("ffn_subpart", "all")
glob_mode = bool(params.get("glob_patterns", False))
include_mode = bool(params.get("include_mode", False))
exclude = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=glob_mode)
discard = _compile_patterns(params.get("discard_patterns", ""), glob_mode=glob_mode)
requested_dtype = _requested_dtype(params["save_dtype"])
handlers = [
MemoryEfficientSafeOpen(path, low_memory=params["lazy_load"])
for path in paths
]
try:
quantizers = [
DiffusionQuantization(handler, path, f"CWB Graft input {idx + 1}")
for idx, (handler, path) in enumerate(zip(handlers, paths))
]
key_maps = [
diffusion_key_map(q.data_keys, f"CWB Graft input {idx + 1}")
for idx, q in enumerate(quantizers)
]
low_bit_sets = [
{
logical
for logical, physical in mapping.items()
if physical in q.isolated_low_bit_keys
}
for mapping, q in zip(key_maps, quantizers)
]
base_keys = sorted(key_maps[0])
donor_keys_set = set(key_maps[1])
block_indices = [
idx
for k in base_keys
if (idx := _extract_main_block_index(k)) is not None
]
max_block_index = max(block_indices) if block_indices else 0
def matches_filter(key: str, patterns: list) -> bool:
aliases = [key, key_maps[0][key]]
return any(
_matches_any_pattern(alias, patterns, glob_mode=glob_mode)
for alias in aliases
)
def should_graft_key(key: str) -> tuple[bool, str | None, tuple[str, list[str], int | None] | None]:
if matches_filter(key, discard):
return (False, None, None)
has_filter_rules = bool(exclude)
matched_filter = matches_filter(key, exclude)
if include_mode and not matched_filter:
return (False, None, None)
if not include_mode and matched_filter:
return (False, None, None)
donor_plan = _resolve_donor_qkv_plan(key, donor_keys_set, custom_map)
if donor_plan is None:
return (False, None, None)
_, donor_logical_keys, _ = donor_plan
if key in low_bit_sets[0] or any(dk in low_bit_sets[1] for dk in donor_logical_keys):
return (False, None, None)
op_type = _classify_graft_operator(key)
if key in custom_map or (include_mode and has_filter_rules and matched_filter):
return (True, op_type or "ffn", donor_plan)
block_idx = _extract_main_block_index(key)
if op_type == "attention":
if _block_rule_allows(block_idx, attn_rule, protect_early, max_block_index):
return (True, "attention", donor_plan)
elif op_type == "ffn":
if _block_rule_allows(block_idx, ffn_rule, protect_early, max_block_index):
return (True, "ffn", donor_plan)
return (False, None, None)
work_units = []
graft_plans: dict[str, tuple[str, tuple[str, list[str], int | None]]] = {}
for key in base_keys:
if matches_filter(key, discard):
work_units.append((key, {}))
continue
graft_ok, op_type, donor_plan = should_graft_key(key)
base_phys = key_maps[0][key]
if not graft_ok:
entries = (
{0: quantizers[0].required_keys(base_phys)}
if base_phys in quantizers[0].quantized_keys
else {}
)
work_units.append((key, entries))
continue
assert op_type is not None and donor_plan is not None
graft_plans[key] = (op_type, donor_plan)
donor_req: list[str] = []
for dk in donor_plan[1]:
d_phys = key_maps[1][dk]
for rk in quantizers[1].required_keys(d_phys):
if rk not in donor_req:
donor_req.append(rk)
work_units.append(
(
key,
{
0: quantizers[0].required_keys(base_phys),
1: donor_req,
},
)
)
output_path, output_rel = canonical_model_artifact_path(
"diffusion_models", params["output_filename"]
)
output_metadata = quantizers[0].output_metadata()
output_metadata["cwb.graft"] = json.dumps(
{
"base_model": params["model_a"],
"donor_model": params["model_b"],
"attention_preset": params["attention_preset"],
"ffn_preset": params["ffn_preset"],
"protect_early_blocks": protect_early,
"attention_block_rule": attn_rule,
"ffn_block_rule": ffn_rule,
"attention_subpart": attn_subpart,
"ffn_subpart": ffn_subpart,
},
sort_keys=True,
)
def decode_logical(index: int, logical_key: str, loaded: dict) -> torch.Tensor:
phys = key_maps[index][logical_key]
if phys in quantizers[index].quantized_keys:
parts = {
part: loaded[(index, part)]
for part in quantizers[index].required_keys(phys)
}
return quantizers[index].decode(phys, parts, torch.float32, device="cpu")
return loaded[(index, phys)]
def preserve_base(writer, key: str, loaded: dict, *, force_raw: bool = False):
phys = key_maps[0][key]
if phys in quantizers[0].quantized_keys and (0, phys) in loaded:
decoded = decode_logical(0, key, loaded)
writer.write(key, decoded.to(requested_dtype).cpu().contiguous())
else:
write_preserved_tensor(
writer,
phys,
handlers[0],
output_key=key,
force_raw=force_raw,
tensor=loaded.get((0, phys)),
)
handler_map = {0: handlers[0], 1: handlers[1]}
grafted_count = 0
preserved_count = 0
pbar = comfy.utils.ProgressBar(len(work_units))
with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing(
stream_work_units(
handler_map,
work_units,
pin_memory=str(device).startswith("cuda"),
)
) as streamed:
for key, loaded in tqdm(streamed, total=len(work_units), desc="CWB grafting tensors", unit="tensors"):
_clear_previous_layer(params)
if matches_filter(key, discard):
pbar.update(1)
continue
if key not in graft_plans:
preserved_count += 1
preserve_base(writer, key, loaded, force_raw=key in low_bit_sets[0])
pbar.update(1)
continue
op_type, (mode, donor_logical_keys, slice_idx) = graft_plans[key]
base_tensor = decode_logical(0, key, loaded)
if mode == "direct":
donor_tensor = decode_logical(1, donor_logical_keys[0], loaded)
elif mode == "fuse_qkv":
donor_tensor = torch.cat(
[decode_logical(1, dk, loaded) for dk in donor_logical_keys],
dim=0,
)
else:
fused_donor = decode_logical(1, donor_logical_keys[0], loaded)
donor_tensor = fused_donor.chunk(3, dim=0)[slice_idx or 0]
base_phys = key_maps[0][key]
base_dtype = quantizers[0].logical_dtype(base_phys, requested_dtype)
if not _is_float_dtype(base_dtype) or not _is_float_dtype(donor_tensor.dtype):
preserved_count += 1
preserve_base(writer, key, loaded)
pbar.update(1)
continue
target_dtype = (
requested_dtype
if params.get("override_dtype", False)
else base_dtype
)
active_settings = attn_settings if op_type == "attention" else ffn_settings
active_subpart = attn_subpart if op_type == "attention" else ffn_subpart
grafted = _graft_with_subpart(
key,
op_type,
active_subpart,
base_tensor,
donor_tensor,
active_settings,
device,
target_dtype,
diagnostics,
)
if grafted is None:
preserved_count += 1
preserve_base(writer, key, loaded)
else:
grafted_count += 1
writer.write(key, grafted)
del base_tensor, donor_tensor, grafted
pbar.update(1)
report = diagnostics.render(
f"attn={params['attention_preset']}, ffn={params['ffn_preset']}",
2,
params.get("attention_config") is not None or params.get("ffn_config") is not None,
)
report += f"\nGrafted tensors: {grafted_count}\nPreserved base tensors: {preserved_count}\n"
return output_rel, report
finally:
for handler in handlers:
handler.__exit__(None, None, None)
cleanup_after_operation()
CWB_GRAFT_NODES = [
CWBModelGraft,
]
+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")
File diff suppressed because it is too large Load Diff
+874
View File
@@ -0,0 +1,874 @@
"""
Learned DoRA Extraction via Optimization.
Extracts Weight-Decomposed Low-Rank Adapters using gradient descent optimization
to minimize the reconstruction MSE loss, matching the true non-linear DoRA formulation.
Uses analytical SVD purely as the initialization seed.
"""
import math
import torch
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
from .layer_parameters import parameter_input, resolve_layer_parameters
from .quantization_guard import DiffusionQuantization, diffusion_key_map
# Import the existing SV ranking and chunking logic so we don't have to duplicate it
from .dora_extract_wd import (
_svd_extract_linear,
_svd_extract_conv, _detect_fused_layer, _compile_patterns,
_matches_any_pattern, _build_lora_output_path, _get_model_inputs,
_format_lora_key, _get_common_inputs
)
def _adaptive_lr_update_cosine(curr_lr: float, improved: bool, worse_loss_counter: int, iteration: int, tensor_shape: tuple, min_lr: float, early_stop_stall: int, lr_factor: float = 0.95, lr_cooldown: int = 1) -> tuple:
M, N = tensor_shape
shape_ratio = abs(M - N) / max(M, N) if max(M, N) > 0 else 0
t = min(worse_loss_counter / max(early_stop_stall, 1), 1.0)
u_factor = (1 + math.cos(2 * math.pi * t)) / 2
if improved:
base_boost = 1.25
distance = base_boost - 1.0
scaled_distance = distance * (1.0 + 0.5 * shape_ratio)
boost_mult = 1.0 + scaled_distance * (1.0 - u_factor)
return min(curr_lr * boost_mult, 100.0), True
else:
cooldown = max(lr_cooldown, 1)
if iteration % cooldown != 0:
return curr_lr, False
min_decay = lr_factor - 0.03 * shape_ratio
max_decay = 0.995
decay_mult = min_decay + (max_decay - min_decay) * u_factor
return max(curr_lr * decay_mult, min_lr), True
def _compute_shape_aware_plateau_params(M: int, N: int, lr_patience: int, lr_factor: float, lr_cooldown: int, lr_shape_influence: float = 1.0) -> tuple:
if lr_shape_influence > 0:
aspect_ratio = max(M, N) / min(M, N) if min(M, N) > 0 else 1.0
ar_factor = math.sqrt(aspect_ratio)
blend = lr_shape_influence
effective_patience = lr_patience
raw_factor = lr_factor
aggressive_factor = raw_factor**ar_factor
effective_factor = raw_factor + (aggressive_factor - raw_factor) * blend
effective_cooldown = lr_cooldown
else:
effective_patience = lr_patience
effective_factor = lr_factor
effective_cooldown = lr_cooldown
return effective_patience, effective_factor, effective_cooldown
def _optimize_dora(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
init_down: torch.Tensor,
init_up: torch.Tensor,
optimize_iters: int,
learning_rate: float,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""
Run gradient descent to optimize lora_down and lora_up against the actual DoRA formula.
Args:
tensor_base: Base model weight
tensor_ft: Finetuned model weight
init_down: SVD-initialized lora_down matrix
init_up: SVD-initialized lora_up matrix
optimize_iters: Number of iterations for AdamW
learning_rate: Learning rate for AdamW
optimizer_type: Optimizer choice (adamw, radam, prodigy)
lr_schedule: LR Schedule choice (exponential, plateau, adaptive)
early_stop_loss: Target loss to stop optimizing
early_stop_stall: Number of non-improving iterations before stopping
Returns:
(optimized_down, optimized_up)
"""
if optimize_iters <= 0:
return init_down, init_up
is_conv = tensor_base.ndim == 4
if is_conv:
out_ch, in_ch, kh, kw = tensor_base.shape
flat_base = tensor_base.reshape(out_ch, -1)
flat_ft = tensor_ft.reshape(out_ch, -1)
# Flatten init matrices for matmul: (out, rank) @ (rank, in*k*k)
rank = init_down.shape[0]
flat_up = init_up.reshape(out_ch, rank)
flat_down = init_down.reshape(rank, -1)
else:
flat_base = tensor_base
flat_ft = tensor_ft
flat_up = init_up
flat_down = init_down
# Magnitude of the finetuned model (fixed DoRA scale)
m = torch.linalg.norm(flat_ft, dim=1, keepdim=True).detach()
eps = 1e-8
# Set up learnable parameters
with torch.inference_mode(False), torch.enable_grad():
# Clone inference tensors to normal tensors for autograd
flat_base_opt = flat_base.clone()
flat_ft_opt = flat_ft.clone()
m_opt = m.clone()
# Free original tensors to save memory
del flat_base, flat_ft, m
B = torch.nn.Parameter(flat_up.clone())
A = torch.nn.Parameter(flat_down.clone())
opt_type = optimizer_type.lower()
if opt_type == "radam":
optimizer = torch.optim.RAdam([B, A], lr=learning_rate)
elif opt_type == "prodigy":
try:
from prodigyplus.prodigy_plus_schedulefree import ProdigyPlusScheduleFree
optimizer = ProdigyPlusScheduleFree([B, A], lr=learning_rate, use_schedulefree=False)
except ImportError:
print("ProdigyPlusScheduleFree optimizer not found. Falling back to AdamW. To install: pip install prodigyplus")
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
else:
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
# Learning rate scheduling and early stopping state
best_loss = float('inf')
worse_loss_counter = 0
plateau_counter = 0
cooldown_counter = 0
current_lr = learning_rate
schedule_type = lr_schedule.lower()
M_dim, N_dim = flat_base_opt.shape[0], flat_base_opt.shape[1]
eff_patience, eff_factor, eff_cooldown = _compute_shape_aware_plateau_params(M_dim, N_dim, lr_patience, lr_factor, lr_cooldown)
# Optimization loop
opt_pbar = tqdm(range(optimize_iters), desc=f" └─ Optimizing {layer_name}", leave=False, dynamic_ncols=True, mininterval=0.5)
for i in opt_pbar:
optimizer.zero_grad()
# 1. Reconstruct directional weight: W_dir = W_base + B @ A
w_approx = flat_base_opt + (B @ A)
# 2. Normalize directional weight
approx_norm = torch.linalg.norm(w_approx, dim=1, keepdim=True) + eps
# 3. Scale by target magnitude
w_reconstructed = m_opt * (w_approx / approx_norm)
# 4. Compute MSE loss against actual finetuned weight
loss = torch.nn.functional.mse_loss(w_reconstructed, flat_ft_opt)
loss_val = loss.item()
if i == 0:
initial_loss = loss_val
best_loss = loss_val
loss.backward()
optimizer.step()
# --- Early Stopping & Schedulers ---
improved = False
if loss_val < best_loss * 0.9999: # 0.01% improvement threshold
best_loss = loss_val
worse_loss_counter = 0
plateau_counter = 0
improved = True
else:
worse_loss_counter += 1
plateau_counter += 1
if loss_val <= early_stop_loss:
break
if worse_loss_counter > early_stop_stall:
break
if current_lr <= early_stop_lr:
break
# LR Schedulers
if schedule_type == "adaptive":
current_lr, updated = _adaptive_lr_update_cosine(
current_lr, improved, worse_loss_counter, i, (M_dim, N_dim), early_stop_lr, early_stop_stall, lr_factor, lr_cooldown
)
if updated:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
elif schedule_type == "plateau":
if cooldown_counter > 0:
cooldown_counter -= 1
elif plateau_counter >= max(eff_patience, 1): # Decay if patience threshold reached
current_lr = max(current_lr * eff_factor, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
plateau_counter = 0 # reset counter after decay
cooldown_counter = eff_cooldown
elif schedule_type == "exponential":
current_lr = max(current_lr * 0.99, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
if i % 10 == 0 or i == optimize_iters - 1:
if schedule_type == "plateau":
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "plat": f"{plateau_counter}/{eff_patience}"}, refresh=False)
else:
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "stall": f"{worse_loss_counter}"}, refresh=False)
final_loss = loss.item()
opt_pbar.close()
print(f" └─ {layer_name} | Loss: {initial_loss:.6f} -> {final_loss:.6f} (stopped at iter {i+1})")
final_up = B.detach()
final_down = A.detach()
# Reshape back to original dimensions
if is_conv:
final_up = final_up.reshape(out_ch, rank, 1, 1)
final_down = final_down.reshape(rank, in_ch, kh, kw)
return final_down, final_up
def _extract_chunked_learned_layer(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
knee_probe_offset: int = 32,
clamp_quantile: float = 0.99,
) -> tuple:
"""Extract Learned DoRA from fused layer by chunking."""
out_dim, in_dim = tensor_base.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk_base = tensor_base[i * chunk_size:(i + 1) * chunk_size, :]
chunk_ft = tensor_ft[i * chunk_size:(i + 1) * chunk_size, :]
norm_base = torch.linalg.norm(chunk_base, dim=1, keepdim=True)
norm_ft = torch.linalg.norm(chunk_ft, dim=1, keepdim=True)
eps = 1e-8
w_dir = chunk_ft * (norm_base / (norm_ft + eps))
weight_diff = w_dir - chunk_base
try:
result, mode_str = _svd_extract_linear(
weight_diff, mode, mode_param, device, max_rank,
clamp_quantile,
knee_probe_offset=knee_probe_offset,
)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
# Apply learned optimization
lora_down, lora_up = _optimize_dora(
chunk_base, chunk_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=f"{layer_name} (Chunk {i+1}/{num_chunks})"
)
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
def extract_dora_learned_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = False,
glob_skip_patterns: bool = False,
knee_probe_offset: int = 32,
include_mode: bool = False,
layer_parameters=None,
) -> None:
"""
Extract Learned DoRA from difference between two models.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[Learned DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
quantization_a = DiffusionQuantization(handler_a, model_a_path, "Learned DoRA Extract model A")
quantization_b = DiffusionQuantization(handler_b, model_b_path, "Learned DoRA Extract model B")
keys_a = set(quantization_a.data_keys)
keys_b = set(quantization_b.data_keys)
logical_a = diffusion_key_map(keys_a, "Learned DoRA Extract model A")
logical_b = diffusion_key_map(keys_b, "Learned DoRA Extract model B")
paired_keys_b = {
key_a: logical_b.get(logical) for logical, key_a in logical_a.items()
}
weight_keys = sorted(
key for key in keys_a
if key.endswith(".weight")
and key not in quantization_a.isolated_low_bit_keys
and paired_keys_b[key] not in quantization_b.isolated_low_bit_keys
)
if mode == "fixed":
profile = "extract:fixed"
defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
"clamp_quantile": clamp_quantile, "min_diff": min_diff}
elif mode == "ratio":
profile = "extract:ratio"
defaults = {"linear_ratio": linear_param, "conv_ratio": conv_param,
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
elif mode == "quantile":
profile = "extract:quantile"
defaults = {"linear_quantile": linear_param, "conv_quantile": conv_param,
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
elif mode == "sv_fro":
profile = "extract:frobenius"
defaults = {"linear_target": linear_param, "conv_target": conv_param,
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
else:
profile = "extract:knee"
defaults = {"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
"clamp_quantile": clamp_quantile, "min_diff": min_diff}
resolved_parameters = resolve_layer_parameters(
layer_parameters, profile, weight_keys, defaults,
node_name="Learned DoRA Extract",
)
work_units = [
(key, key, paired_keys_b[key])
for key in weight_keys
if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
and not (paired_keys_b[key] is None and mismatch_mode == "skip")
]
pbar = comfy.utils.ProgressBar(len(work_units))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key, cpu_a, cpu_b):
key_b = paired_keys_b[key]
values = resolved_parameters.get(key, defaults)
if mode == "fixed":
layer_linear_param, layer_conv_param = values["linear_dim"], values["conv_dim"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
elif mode == "ratio":
layer_linear_param, layer_conv_param = values["linear_ratio"], values["conv_ratio"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
elif mode == "quantile":
layer_linear_param, layer_conv_param = values["linear_quantile"], values["conv_quantile"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
elif mode == "sv_fro":
layer_linear_param, layer_conv_param = values["linear_target"], values["conv_target"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
else:
layer_linear_param, layer_conv_param = linear_param, conv_param
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
layer_clamp_quantile = values["clamp_quantile"]
layer_min_diff = values["min_diff"]
lora_name = _format_lora_key(key)
weight_diff = None
tensor_base = None
tensor_ft = None
layer_device = device
if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)):
return "skipped", None
if key_b is None:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
dora_difference(cpu_a, None, device, keep_sources=True)
)
else:
if cpu_a.shape != cpu_b.shape:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
dora_difference(cpu_a, None, device, keep_sources=True)
)
else:
if cpu_a.ndim < 2:
return "skipped", None
weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
dora_difference(cpu_a, cpu_b, device, keep_sources=True)
)
if weight_diff is None or tensor_base is None or tensor_ft is None:
raise RuntimeError(f"No extraction tensors were produced for {key}")
if layer_min_diff > 0 and weight_diff.abs().max() < layer_min_diff:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if weight_diff.ndim < 2:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
svd_weight = weight_diff
optimize_base = tensor_base
optimize_ft = tensor_ft
try:
# 1. Run Analytical SVD (Initialization)
if is_conv:
result, mode_str = retry_cuda_oom_on_cpu(
lambda: _svd_extract_conv(svd_weight, mode, layer_conv_param, layer_device, layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
lambda: _svd_extract_conv(svd_weight.cpu(), mode, layer_conv_param, "cpu", layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
layer_device,
)
else:
result, mode_str = retry_cuda_oom_on_cpu(
lambda: _svd_extract_linear(svd_weight, mode, layer_linear_param, layer_device, layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
lambda: _svd_extract_linear(svd_weight.cpu(), mode, layer_linear_param, "cpu", layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
layer_device,
)
# 2. Run Optimization (Learned Rounding)
if mode_str != "full" and dora_scale is not None:
lora_down, lora_up, _ = result
optimize_args = (
optimize_iters, learning_rate, optimizer_type, lr_schedule,
lr_patience, lr_factor, lr_cooldown, early_stop_loss,
early_stop_stall, early_stop_lr,
)
lora_down, lora_up = retry_cuda_oom_on_cpu(
lambda: _optimize_dora(optimize_base, optimize_ft, lora_down, lora_up, *optimize_args, layer_name=lora_name),
lambda: _optimize_dora(optimize_base.cpu(), optimize_ft.cpu(), lora_down.cpu(), lora_up.cpu(), *optimize_args, layer_name=lora_name),
layer_device,
)
result = (lora_down, lora_up, _)
except Exception as e:
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[Learned DoRA] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_learned_layer(
tensor_base, tensor_ft, num_chunks, mode, layer_linear_param, layer_device, layer_linear_max_rank, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name,
knee_probe_offset=knee_probe_offset,
clamp_quantile=layer_clamp_quantile,
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff, tensor_base, tensor_ft
return "chunked", layer_results
print(f"[Learned DoRA] Failed: {key}: {e}")
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
rank = lora_down.shape[0]
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff, tensor_base, tensor_ft
return status, layer_results
writer_context = atomic_uel_writer(output_path)
writer = writer_context.__enter__()
stream = None
try:
stream = paired_async_tensors(
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda"),
quantization_a=quantization_a, quantization_b=quantization_b,
compute_dtype=torch.float32,
)
for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting Learned DoRA", unit="layers"):
status, layer_sd = _process_layer(key, cpu_a, cpu_b)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
stats["skipped"] += len(weight_keys) - len(work_units)
except BaseException as exc:
if stream is not None:
stream.close()
writer_context.__exit__(type(exc), exc, exc.__traceback__)
raise
else:
if stream is not None:
stream.close()
writer_context.__exit__(None, None, None)
print(f"[Learned DoRA] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
# =============================================================================
# Node Definitions
# =============================================================================
def _get_learned_inputs():
return [
io.Int.Input("optimize_iters", default=500, min=0, max=10000,
tooltip="Number of gradient descent iterations (0 = skip optimization)"),
io.Float.Input("learning_rate", default=1.0, min=0.0001, max=100.0, step=0.001,
tooltip="Base learning rate. AdamW prefers ~0.01, Prodigy prefers ~1.0"),
io.Combo.Input("optimizer", options=["prodigy", "adamw", "radam"], default="prodigy",
tooltip="Optimization algorithm"),
io.Combo.Input("lr_schedule", options=["plateau", "adaptive", "exponential", "constant"], default="plateau",
tooltip="Learning rate scheduling strategy"),
io.Int.Input("lr_patience", default=2, min=0, max=1000, tooltip="Steps to wait before decaying LR (Plateau)"),
io.Float.Input("lr_factor", default=0.9, min=0.01, max=1.0, step=0.01, tooltip="Factor to decay LR by (Plateau)"),
io.Int.Input("lr_cooldown", default=2, min=0, max=1000, tooltip="Steps to wait after decay before checking again (Plateau)"),
io.Float.Input("early_stop_loss", default=1e-6, min=0.0, max=1.0, step=1e-9,
tooltip="Stop early if MSE loss drops below this value"),
io.Int.Input("early_stop_stall", default=2000, min=0, max=10000,
tooltip="Stop early if loss doesn't improve for this many iterations"),
io.Float.Input("early_stop_lr", default=9.01e-9, min=0.0, max=1.0, step=1e-10,
tooltip="Stop early if learning rate drops below this value"),
]
class DoRALearnedExtractFixed(io.ComfyNode):
"""Extract Learned DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFixed",
display_name="DoRA Learned Extract (Fixed Rank) (WIP)",
category="ModelUtils/DoRA",
description="Extract DoRA via MSE Loss minimization with fixed rank.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
*_get_learned_inputs(),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD iterations for initialization"),
*_get_common_inputs(),
parameter_input("extract:fixed"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns, include_mode=include_mode,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class DoRALearnedExtractRatio(io.ComfyNode):
"""Extract Learned DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractRatio",
display_name="DoRA Learned Extract (Ratio) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping SVs > max(S) / ratio.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_learned_inputs(),
*_get_common_inputs(),
parameter_input("extract:ratio"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset, layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class DoRALearnedExtractQuantile(io.ComfyNode):
"""Extract Learned DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractQuantile",
display_name="DoRA Learned Extract (Quantile) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping enough SVs to reach target percentage.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_learned_inputs(),
*_get_common_inputs(),
parameter_input("extract:quantile"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset, layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class DoRALearnedExtractKnee(io.ComfyNode):
"""Extract Learned DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractKnee",
display_name="DoRA Learned Extract (Knee Detection) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction automatically finding optimal rank using knee detection.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_learned_inputs(),
*_get_common_inputs(),
parameter_input("extract:knee"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class DoRALearnedExtractFrobenius(io.ComfyNode):
"""Extract Learned DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFrobenius",
display_name="DoRA Learned Extract (Frobenius) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction preserving target fraction of Frobenius norm.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_learned_inputs(),
*_get_common_inputs(),
parameter_input("extract:frobenius"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
+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="/", tooltip="Checkpoint subdirectory to scan; Recursive also includes its descendants."))
return io.Schema(
node_id=cls.__name__,
display_name="Checkpoint Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("checkpoints")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class DiffusionModelInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("diffusion_models"), default="/", tooltip="Diffusion-model subdirectory to scan; Recursive also includes its descendants."))
return io.Schema(
node_id=cls.__name__,
display_name="Diffusion Model Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("diffusion_models")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class LoRAInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("loras"), default="/", tooltip="LoRA subdirectory to scan; Recursive also includes its descendants."))
return io.Schema(
node_id=cls.__name__,
display_name="LoRA Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("loras")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class EmbeddingInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("embeddings"), default="/", tooltip="Embedding subdirectory to scan; Recursive also includes its descendants."))
return io.Schema(
node_id=cls.__name__,
display_name="Embedding Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("embeddings")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class VAEInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("vae"), default="/", tooltip="VAE subdirectory to scan; Recursive also includes its descendants."))
return io.Schema(
node_id=cls.__name__,
display_name="VAE Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("vae")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class ControlNetInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("controlnet"), default="/", tooltip="ControlNet subdirectory to scan; Recursive also includes its descendants."))
return io.Schema(
node_id=cls.__name__,
display_name="ControlNet Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("controlnet")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class ManualPathInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.append(io.String.Input("scan_dir", default="", tooltip="Absolute directory to scan for supported model files."))
return io.Schema(
node_id=cls.__name__,
display_name="Manual Path Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, scan_dir, **kwargs):
return super().execute(
recursive=recursive,
nsfw_level=nsfw_level,
max_examples=max_examples,
threads=threads,
api_key=api_key,
scan_dir=scan_dir,
**kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
scan_dir = kwargs.get("scan_dir", "")
if scan_dir.strip():
return [scan_dir.strip()]
return []
+638
View File
@@ -0,0 +1,638 @@
import os
import sys
import hashlib
import json
import sqlite3
import requests
import glob
from pathlib import Path
import threading
import concurrent.futures
import codecs
import shutil
try:
from PIL import Image
HAS_PIL = True
except ImportError:
HAS_PIL = False
print("Warning: Pillow not found. Image workflow extraction disabled.")
try:
from mutagen.mp4 import MP4
HAS_MUTAGEN = True
except ImportError:
HAS_MUTAGEN = False
print("Warning: mutagen not found. MP4 tag extraction will rely entirely on av fallback.")
try:
import av
HAS_AV = True
except ImportError:
HAS_AV = False
print("Warning: av (PyAV) not found. MP4 metadata extraction fallback disabled.")
# Assuming downloader_utils.py is inside nodes/, the root is one level up
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
from unifiedefficientloader import UnifiedSafetensorsLoader
from .model_info_nodes import (
get_potential_preview_files,
load_image_tensor,
get_model_workflows,
get_model_metadata_file,
)
# --- Database Path Logic ---
try:
import folder_paths
base_user_dir = folder_paths.get_public_user_directory("default")
if not base_user_dir:
base_user_dir = root_dir
except (ImportError, AttributeError):
base_user_dir = root_dir
PLUGIN_DATA_DIR = os.path.join(base_user_dir, "ComfyUI-ModelUtils")
os.makedirs(PLUGIN_DATA_DIR, exist_ok=True)
DEFAULT_DB_PATH = os.path.join(PLUGIN_DATA_DIR, "cache.db")
# --- Workflow Extraction Logic ---
def unescape_and_parse_nested_json(json_string):
if not isinstance(json_string, str): return None
stripped = json_string.strip()
if not (stripped.startswith('{') or stripped.startswith('[')): return None
try: return json.loads(stripped)
except json.JSONDecodeError: pass
try:
first_parse = json.loads(stripped)
if isinstance(first_parse, str): return json.loads(first_parse)
except json.JSONDecodeError: pass
try:
decoded = codecs.decode(stripped, "unicode_escape")
return json.loads(decoded)
except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass
try:
unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\")
unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t")
return json.loads(unescaped)
except json.JSONDecodeError: pass
return None
def get_workflow_blob(metadata):
for key in ['workflow', 'prompt']:
val = metadata.get(key)
if val:
parsed = unescape_and_parse_nested_json(val)
if parsed: return parsed
for val in metadata.values():
parsed = unescape_and_parse_nested_json(val)
if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
return parsed
return None
def process_image_workflow_and_compress(filepath_str):
if not HAS_PIL: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
p = Path(filepath_str)
if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
workflow = None
actual_format = None
stats = {"images_processed": 1, "workflows_extracted": 0, "space_saved": 0}
try:
with Image.open(p) as img:
actual_format = img.format
# CivitAI sometimes serves PNGs with .jpeg extension in URL
if actual_format == 'PNG':
workflow = get_workflow_blob(img.info)
if workflow:
# Save workflow
json_path = p.with_suffix('.json')
if not json_path.exists():
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(workflow, f, indent=2, ensure_ascii=False)
print(f"Extracted workflow to {json_path}")
stats["workflows_extracted"] = 1
# Compress image to WebP
webp_path = p.with_suffix('.webp')
if not webp_path.exists():
orig_size = os.path.getsize(p)
img.save(webp_path, 'WEBP', quality=85, method=4)
new_size = os.path.getsize(webp_path)
stats["space_saved"] = max(0, orig_size - new_size)
print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}")
# Delete original if we created a webp replacement
if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG':
os.remove(p)
print(f"Deleted original uncompressed image: {p.name}")
except Exception as e:
print(f"Failed to process image {p.name}: {e}")
return stats
# --- Video Workflow Logic ---
def get_tag_value(mp4, keys):
for key in keys:
if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None
for key in keys:
search_suffix = key.split(':')[-1].lower()
for tag_key, val in mp4.tags.items():
if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None
return None
def extract_metadata_av(mp4_path):
if not HAS_AV: return None
try:
with av.open(str(mp4_path)) as container:
return {k.lower(): v for k, v in container.metadata.items()}
except Exception as e:
print(f"av extraction error for {mp4_path}: {e}")
return None
def process_video_workflow(filepath_str):
p = Path(filepath_str)
if p.suffix.lower() != '.mp4': return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
json_path = p.with_suffix('.json')
if json_path.exists(): return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
workflow_data = None
mp4 = None
stats = {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
if HAS_MUTAGEN:
try: mp4 = MP4(p)
except: pass
if mp4:
for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]:
val = get_tag_value(mp4, tag_keys)
if val:
if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore')
try:
parsed = json.loads(val)
if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
if 'prompt' in parsed and isinstance(parsed['prompt'], str):
try: workflow_data = json.loads(parsed['prompt'])
except: workflow_data = parsed
else:
workflow_data = parsed
break
except json.JSONDecodeError:
# Try more aggressive parsing if it's a nested string
parsed = unescape_and_parse_nested_json(val)
if parsed:
workflow_data = parsed
break
if not workflow_data:
av_data = extract_metadata_av(p)
if av_data:
for key in ['workflow', 'prompt', 'comment', 'description']:
if key in av_data:
val = av_data[key]
try:
parsed = json.loads(val)
if isinstance(parsed, dict):
if 'prompt' in parsed and isinstance(parsed['prompt'], str):
try: workflow_data = json.loads(parsed['prompt'])
except: workflow_data = parsed
else:
workflow_data = parsed
break
except json.JSONDecodeError:
parsed = unescape_and_parse_nested_json(val)
if parsed:
workflow_data = parsed
break
if workflow_data:
try:
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(workflow_data, f, indent=2, ensure_ascii=False)
print(f"Extracted workflow to {json_path.name}")
stats["workflows_extracted"] = 1
except Exception as e:
print(f"Error saving workflow for {p.name}: {e}")
return stats
# --- Constants for Civitai NSFW levels ---
NSFW_LEVELS = {
"None": 1,
"Soft": 2,
"Mature": 4,
"X": 8,
"XXX": 16,
"All": 32 # Anything
}
class HashCache:
def __init__(self, db_path=None):
self.db_path = db_path if db_path else DEFAULT_DB_PATH
self.is_sqlite = self.db_path.endswith('.db') or self.db_path.endswith('.sqlite')
self.lock = threading.Lock()
if self.is_sqlite:
self._init_sqlite()
else:
self.cache = self._load_json()
def _init_sqlite(self):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS file_hashes (
filepath TEXT PRIMARY KEY,
mtime REAL,
size INTEGER,
hash TEXT
)
''')
conn.commit()
def _load_json(self):
if os.path.exists(self.db_path):
with open(self.db_path, "r", encoding="utf-8") as f:
return json.load(f)
return {}
def _save_json(self):
with open(self.db_path, "w", encoding="utf-8") as f:
json.dump(self.cache, f, indent=4)
def _get_sqlite_hash(self, filepath, mtime, size):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,))
row = cursor.fetchone()
if row:
db_mtime, db_size, db_hash = row
if db_mtime == mtime and db_size == size:
return db_hash
return None
def _set_sqlite_hash(self, filepath, mtime, size, file_hash):
with sqlite3.connect(self.db_path, timeout=10.0) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash)
VALUES (?, ?, ?, ?)
''', (filepath, mtime, size, file_hash))
conn.commit()
def get_hash(self, filepath):
if not os.path.exists(filepath):
return None
mtime = os.path.getmtime(filepath)
size = os.path.getsize(filepath)
key = filepath
with self.lock:
if self.is_sqlite:
cached_hash = self._get_sqlite_hash(key, mtime, size)
if cached_hash: return cached_hash
else:
if key in self.cache:
entry = self.cache[key]
if entry["mtime"] == mtime and entry["size"] == size:
return entry["hash"]
# Calculate new hash (outside the lock to allow parallel hashing)
print(f"Calculating SHA256 for: {filepath}")
sha256_hash = self.calculate_sha256(filepath)
with self.lock:
if self.is_sqlite:
self._set_sqlite_hash(key, mtime, size, sha256_hash)
else:
self.cache[key] = {
"mtime": mtime,
"size": size,
"hash": sha256_hash
}
self._save_json()
return sha256_hash
@staticmethod
def calculate_sha256(filename):
sha256_hash = hashlib.sha256()
with open(filename, "rb") as f:
# Read and update hash string value in blocks of 4K
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
class CivitaiAPI:
def __init__(self, api_key=None):
self.api_key = api_key
self.base_url = "https://civitai.com/api/v1"
def get_headers(self):
headers = {}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
return headers
def get_model_info_by_hash(self, file_hash):
url = f"{self.base_url}/model-versions/by-hash/{file_hash}"
response = requests.get(url, headers=self.get_headers())
if response.status_code == 200:
return response.json()
elif response.status_code == 404:
return None
else:
response.raise_for_status()
class Downloader:
@staticmethod
def download_file(url, target_path, headers=None):
if headers is None:
headers = {}
temp_path = target_path + ".downloading"
downloaded_size = 0
if os.path.exists(temp_path):
downloaded_size = os.path.getsize(temp_path)
if downloaded_size > 0:
headers["Range"] = f"bytes={downloaded_size}-"
response = requests.get(url, headers=headers, stream=True)
if response.status_code == 416: # Range not satisfiable, file probably changed or completed
os.remove(temp_path)
downloaded_size = 0
if "Range" in headers:
del headers["Range"]
response = requests.get(url, headers=headers, stream=True)
if response.status_code not in [200, 206]:
print(f"Failed to download from {url}, status code: {response.status_code}")
return False
mode = "ab" if response.status_code == 206 else "wb"
total_size = int(response.headers.get('content-length', 0)) + downloaded_size
print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)")
with open(temp_path, mode) as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
os.rename(temp_path, target_path)
return True
def preview_exists(base_path):
previews = get_potential_preview_files(base_path)
for p in previews:
if os.path.exists(p):
return True
return False
def next_example_image_path(base_path, ext):
i = 0
while len(glob.glob(f"{base_path}.example.{i}.*")) > 0:
i += 1
return f"{base_path}.example.{i}.{ext}"
def filter_images_by_nsfw(images, threshold_level):
valid_images = []
for img in images:
rating = img.get("nsfwLevel", 1)
if rating <= threshold_level:
valid_images.append(img)
return valid_images
def extract_safetensors_metadata(filepath):
try:
# Use low_memory=True to only read the header and metadata without loading tensors
loader = UnifiedSafetensorsLoader(filepath, low_memory=True)
metadata = loader.metadata()
loader.close()
return metadata if metadata else None
except Exception as e:
print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}")
return None
def process_file(filepath, cache, api, nsfw_threshold, max_examples):
filepath_str = str(filepath)
base_path = os.path.splitext(filepath_str)[0]
is_safetensors = filepath_str.lower().endswith(".safetensors")
stats = {"new_hashes": 0, "images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
# Check if all targets already exist
info_path = f"{base_path}.civitai.info"
st_metadata_path = f"{base_path}.metadata.json"
has_info = os.path.exists(info_path)
has_preview = preview_exists(base_path)
has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path)
# Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file
for p in get_potential_preview_files(base_path):
if os.path.exists(p):
ext_lower = p.lower()
if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
res = process_image_workflow_and_compress(p)
for k in stats: stats[k] += res.get(k, 0)
elif ext_lower.endswith('.mp4'):
res = process_video_workflow(p)
for k in stats: stats[k] += res.get(k, 0)
for ex in glob.glob(f"{base_path}.example.*"):
ext_lower = ex.lower()
if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
res = process_image_workflow_and_compress(ex)
for k in stats: stats[k] += res.get(k, 0)
elif ext_lower.endswith('.mp4'):
res = process_video_workflow(ex)
for k in stats: stats[k] += res.get(k, 0)
# If everything exists and we don't want examples or already have enough examples, we might skip API logic
if has_info and has_preview and has_st_metadata and max_examples == 0:
print(f"Skipping fully populated model: {filepath_str}")
return stats
print(f"Processing model: {filepath_str}")
# Safetensors built-in metadata extraction
if is_safetensors and not has_st_metadata:
metadata = extract_safetensors_metadata(filepath_str)
if metadata:
with open(st_metadata_path, "w", encoding="utf-8") as f:
json.dump(metadata, f, indent=4)
print(f"Saved safetensors metadata to {st_metadata_path}")
mtime = os.path.getmtime(filepath_str)
size = os.path.getsize(filepath_str)
# We want to know if it's a NEW hash or cached
is_cached = False
with cache.lock:
if cache.is_sqlite:
cached_hash = cache._get_sqlite_hash(filepath_str, mtime, size)
if cached_hash: is_cached = True
else:
if filepath_str in cache.cache:
entry = cache.cache[filepath_str]
if entry["mtime"] == mtime and entry["size"] == size:
is_cached = True
file_hash = cache.get_hash(filepath_str)
if not file_hash:
return stats
if not is_cached:
stats["new_hashes"] += 1
try:
model_info = api.get_model_info_by_hash(file_hash)
except Exception as e:
print(f"API Error for {filepath_str}: {e}")
return stats
if not model_info:
print(f"Model not found on Civitai: {filepath_str}")
return stats
# 1. Save Info
if not has_info:
with open(info_path, "w", encoding="utf-8") as f:
json.dump(model_info, f, indent=4)
print(f"Saved metadata to {info_path}")
images = model_info.get("images", [])
valid_images = filter_images_by_nsfw(images, nsfw_threshold)
if not valid_images:
print(f"No valid images found matching NSFW criteria for {filepath_str}")
return stats
# 2. Save Preview
if not has_preview and len(valid_images) > 0:
preview_img = valid_images[0]
url = preview_img.get("url")
ext = "jpeg" # Default fallback
if ".png" in url.lower(): ext = "png"
elif ".mp4" in url.lower(): ext = "mp4"
elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
preview_path = f"{base_path}.preview.{ext}"
if Downloader.download_file(url, preview_path):
if ext in ("png", "jpg", "jpeg"):
res = process_image_workflow_and_compress(preview_path)
for k in stats: stats[k] += res.get(k, 0)
elif ext == "mp4":
res = process_video_workflow(preview_path)
for k in stats: stats[k] += res.get(k, 0)
# 3. Save Examples
if max_examples > 0:
existing_examples_count = len(glob.glob(f"{base_path}.example.*"))
# We divide by 2 roughly if there are json files, let's just count unique example indices
existing_indices = set()
for f in glob.glob(f"{base_path}.example.*"):
parts = f.split('.example.')
if len(parts) > 1:
idx = parts[1].split('.')[0]
if idx.isdigit(): existing_indices.add(idx)
needed_examples = max_examples - len(existing_indices)
if needed_examples > 0:
example_images = valid_images[1:max_examples+1]
# Skip images we probably already downloaded by skipping the first N
example_images_to_dl = example_images[len(existing_indices):]
for img in example_images_to_dl:
url = img.get("url")
ext = "jpeg"
if ".png" in url.lower(): ext = "png"
elif ".mp4" in url.lower(): ext = "mp4"
elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
example_path = next_example_image_path(base_path, ext)
if Downloader.download_file(url, example_path):
if ext in ("png", "jpg", "jpeg"):
res = process_image_workflow_and_compress(example_path)
for k in stats: stats[k] += res.get(k, 0)
elif ext == "mp4":
res = process_video_workflow(example_path)
for k in stats: stats[k] += res.get(k, 0)
return stats
def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4):
"""
Scans directories for models and fetches their metadata/previews.
Adapted for direct Python API calls rather than argparse CLI.
"""
db_path = db_path if db_path else DEFAULT_DB_PATH
cache = HashCache(db_path)
api = CivitaiAPI(api_key)
model_exts = [".safetensors", ".ckpt", ".pt"]
files_to_process = []
if isinstance(scan_dirs, str):
scan_dirs = [scan_dirs]
for scan_dir in scan_dirs:
path = Path(scan_dir)
if recursive:
for ext in model_exts:
# rglob includes all directories, so we must filter out hidden ones
for f in path.rglob(f"*{ext}"):
if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
files_to_process.append(f)
else:
for ext in model_exts:
for f in path.glob(f"*{ext}"):
if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
files_to_process.append(f)
nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"])
print(f"Found {len(files_to_process)} models to process.")
global_stats = {
"total_processed": len(files_to_process),
"new_hashes": 0,
"images_processed": 0,
"workflows_extracted": 0,
"space_saved": 0
}
# Import and use ComfyUI's ProgressBar
try:
import comfy.utils
pbar = comfy.utils.ProgressBar(len(files_to_process))
except ImportError:
pbar = None
if threads > 1:
print(f"Starting {threads} threads...")
with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process]
for future in concurrent.futures.as_completed(futures):
try:
res = future.result()
if res:
for k in global_stats: global_stats[k] += res.get(k, 0)
except Exception as e:
print(f"Error in thread: {e}")
if pbar:
pbar.update(1)
else:
for filepath in files_to_process:
res = process_file(filepath, cache, api, nsfw_threshold, max_examples)
if res:
for k in global_stats: global_stats[k] += res.get(k, 0)
if pbar:
pbar.update(1)
return global_stats
+242
View File
@@ -0,0 +1,242 @@
"""Asynchronous dtype conversion for diffusion-model safetensors files."""
from __future__ import annotations
import os
import re
from collections import Counter
from contextlib import nullcontext
import comfy.utils
import folder_paths
import torch
from comfy_api.latest import io
from tqdm import tqdm
from unifiedefficientloader import MemoryEfficientSafeOpen
from .device_utils import cleanup_after_operation
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
TARGET_DTYPES = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}
def _compile_exclusions(pattern_text: str) -> list[re.Pattern]:
patterns = []
for line_number, value in enumerate(pattern_text.splitlines(), start=1):
pattern = value.strip()
if not pattern:
continue
try:
patterns.append(re.compile(pattern))
except re.error as exc:
raise ValueError(
f"Invalid exclusion regex on line {line_number}: {exc}"
) from exc
return patterns
def _is_excluded(key: str, patterns: list[re.Pattern]) -> bool:
return any(pattern.search(key) for pattern in patterns)
def _output_path(output_filename: str) -> tuple[str, str]:
return canonical_model_artifact_path("diffusion_models", output_filename)
def convert_diffusion_model_dtype(
model_name: str,
target_dtype: str,
exclude_patterns: str,
output_filename: str,
reference_model: str | None = None,
include_mode: bool = False,
) -> tuple[str, str]:
try:
destination_dtype = TARGET_DTYPES[target_dtype]
except KeyError as exc:
raise ValueError(f"Unsupported target dtype: {target_dtype}") from exc
source_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_name
)
output_path, output_name = _output_path(output_filename)
if os.path.normcase(os.path.abspath(source_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the input model path.")
exclusions = _compile_exclusions(exclude_patterns)
counts = Counter()
original_dtypes = Counter()
reference_path = None
if reference_model and reference_model != "None":
reference_path = folder_paths.get_full_path_or_raise(
"diffusion_models", reference_model
)
try:
with MemoryEfficientSafeOpen(source_path, low_memory=True) as loader, (
MemoryEfficientSafeOpen(reference_path, low_memory=True)
if reference_path else nullcontext()
) as reference_loader:
metadata = (loader.metadata() or {}).copy()
reference_keys = set(reference_loader.keys()) if reference_loader else set()
with atomic_uel_writer(output_path, metadata) as writer:
keys = list(loader.keys())
progress = comfy.utils.ProgressBar(len(keys))
stream = loader.async_stream(
keys,
batch_size=1,
prefetch_batches=1,
pin_memory=False,
)
for batch in tqdm(
stream,
total=len(keys),
desc="Converting diffusion-model dtype",
unit="tensors",
):
key, source = batch[0]
original_dtypes[str(source.dtype).removeprefix("torch.")] += 1
output = source
if _is_excluded(key, exclusions) != include_mode:
counts["excluded"] += 1
elif not source.is_floating_point():
counts["non_floating"] += 1
elif reference_loader is not None and key in reference_keys:
reference_dtype = reference_loader.get_dtype(key)
if not reference_dtype.is_floating_point:
raise ValueError(
f"Reference tensor '{key}' is {reference_dtype}, but the input tensor is floating."
)
if source.dtype == reference_dtype:
counts["reference_already_matched"] += 1
else:
output = source.to(dtype=reference_dtype).contiguous()
counts["reference_mapped"] += 1
elif reference_loader is not None:
output = source if source.dtype == destination_dtype else source.to(
dtype=destination_dtype
).contiguous()
counts["reference_missing_fallback"] += 1
elif source.dtype == destination_dtype:
counts["already_target"] += 1
else:
output = source.to(dtype=destination_dtype).contiguous()
counts["converted"] += 1
writer.write_batch([(key, output)])
loader.mark_processed(key)
progress.update(1)
del output
del source
batch.clear()
del batch
finally:
cleanup_after_operation()
total = sum(counts.values())
dtype_summary = ", ".join(
f"{name}={count}" for name, count in sorted(original_dtypes.items())
)
report = "\n".join(
(
"DIFFUSION MODEL DTYPE CONVERSION",
f"Input: {model_name}",
f"Reference model: {reference_model or 'None'}",
f"Output: {output_name}",
f"Target dtype: {target_dtype}",
f"Total tensors: {total}",
f"Converted floating tensors: {counts['converted']}",
f"Excluded tensors preserved: {counts['excluded']}",
f"Already target dtype: {counts['already_target']}",
f"Reference-mapped floating tensors: {counts['reference_mapped']}",
f"Reference dtype already matched: {counts['reference_already_matched']}",
f"Reference-missing fallback tensors: {counts['reference_missing_fallback']}",
f"Non-floating tensors preserved: {counts['non_floating']}",
f"Original dtype counts: {dtype_summary or 'None'}",
)
)
return output_name, report
class DiffusionModelDtypeConversion(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DiffusionModelDtypeConversion",
display_name="Diffusion Model Dtype Conversion",
category="ModelUtils/Conversion",
description=(
"Streams a diffusion-model safetensors file into fp32, fp16, "
"or bf16. Floating tensors matching an exclusion regex retain "
"their original dtype, or enable Include Mode to convert only matches; "
"non-floating tensors are always preserved."
),
inputs=[
io.Combo.Input(
"model_name",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Diffusion-model file whose floating tensors will be converted.",
),
io.Combo.Input(
"target_dtype", options=["fp32", "fp16", "bf16"],
tooltip="Target dtype for floating tensors selected by the layer filter.",
),
io.Combo.Input(
"reference_model",
options=["None", *folder_paths.get_filename_list("diffusion_models")],
default="None",
tooltip="Optional authoritative dtype map. Matching floating tensor keys use this model's header dtype without loading its tensors; missing keys use Target Dtype.",
),
io.String.Input(
"exclude_patterns",
default="",
multiline=True,
tooltip=(
"Optional Python regex patterns, one per line. Matching "
"tensor keys retain their original dtype, or are the only keys "
"converted when Include Mode is enabled."
),
),
io.String.Input(
"output_filename", default="converted_model",
tooltip="Output filename without extension, written under ComfyUI's diffusion-model directory.",
),
io.Boolean.Input(
"include_mode", default=False,
tooltip="Use Exclude Patterns as an include-only filter. Only matching tensors are converted; an empty filter converts nothing. Nonmatches retain their original dtype.",
),
],
outputs=[
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="conversion_report"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_name: str,
target_dtype: str,
reference_model: str,
exclude_patterns: str,
output_filename: str,
include_mode: bool = False,
) -> io.NodeOutput:
output_path, report = convert_diffusion_model_dtype(
model_name,
target_dtype,
exclude_patterns,
output_filename,
reference_model,
include_mode,
)
return io.NodeOutput(output_path, report)
+179
View File
@@ -0,0 +1,179 @@
"""Bounded paired UEL streaming helpers for model extraction nodes."""
import torch
from .uel_io import AsyncTensorCursor
def is_cuda_oom(error):
return isinstance(error, torch.cuda.OutOfMemoryError) or (
isinstance(error, RuntimeError)
and "cuda" in str(error).lower()
and "out of memory" in str(error).lower()
)
def release_cuda_oom():
if torch.cuda.is_available():
torch.cuda.empty_cache()
def move_source_tensors(cpu_a, cpu_b, device, dtype=torch.float32):
"""Move one source pair, retrying a confirmed CUDA OOM on CPU."""
use_cuda = str(device).startswith("cuda")
try:
if use_cuda:
tensor_a = cpu_a.to(device=device, dtype=dtype, non_blocking=True)
tensor_b = cpu_b.to(device=device, dtype=dtype, non_blocking=True) if cpu_b is not None else None
else:
tensor_a = cpu_a.to(device=device, dtype=dtype)
tensor_b = cpu_b.to(device=device, dtype=dtype) if cpu_b is not None else None
return tensor_a, tensor_b, device
except BaseException as error:
if not use_cuda or not is_cuda_oom(error):
raise
release_cuda_oom()
return (
cpu_a.to(device="cpu", dtype=dtype),
cpu_b.to(device="cpu", dtype=dtype) if cpu_b is not None else None,
"cpu",
)
def retry_cuda_oom_on_cpu(cuda_call, cpu_call, device):
"""Retry one compute operation on CPU only for a confirmed CUDA OOM."""
try:
return cuda_call()
except BaseException as error:
if not str(device).startswith("cuda") or not is_cuda_oom(error):
raise
release_cuda_oom()
return cpu_call()
def lora_difference(cpu_a, cpu_b, device):
"""Construct A minus B with transfer and arithmetic CUDA OOM fallback."""
tensor_a, tensor_b, actual_device = move_source_tensors(
cpu_a, cpu_b, device
)
if tensor_b is None:
return tensor_a, actual_device
try:
return tensor_a - tensor_b, actual_device
except BaseException as error:
if not str(actual_device).startswith("cuda") or not is_cuda_oom(error):
raise
del tensor_a, tensor_b
release_cuda_oom()
return cpu_a.float() - cpu_b.float(), "cpu"
def dora_difference(cpu_a, cpu_b, device, *, keep_sources=False):
"""Construct a directional DoRA delta with complete CUDA OOM fallback."""
tensor_a, tensor_b, actual_device = move_source_tensors(
cpu_a, cpu_b, device
)
def compute(a, b):
if b is None:
base = torch.zeros_like(a)
result = (a, None, base, a)
return result if keep_sources else result[:2]
if a.ndim < 2:
result = (None, None, b, a)
return result if keep_sources else result[:2]
is_conv = a.ndim == 4
if is_conv:
out_channels = a.shape[0]
norm_a = torch.linalg.norm(a.reshape(out_channels, -1), dim=1, keepdim=True)
norm_b = torch.linalg.norm(b.reshape(out_channels, -1), dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(b, dim=1, keepdim=True)
ratio = norm_b / (norm_a + 1e-8)
if is_conv:
ratio = ratio.view(-1, 1, 1, 1)
scale = norm_a.view(-1, 1, 1, 1)
else:
scale = norm_a
result = (a * ratio - b, scale, b, a)
return result if keep_sources else result[:2]
try:
return (*compute(tensor_a, tensor_b), actual_device)
except BaseException as error:
if not str(actual_device).startswith("cuda") or not is_cuda_oom(error):
raise
del tensor_a, tensor_b
release_cuda_oom()
cpu_b_float = cpu_b.float() if cpu_b is not None else None
return (*compute(cpu_a.float(), cpu_b_float), "cpu")
def paired_async_tensors(
handler_a,
handler_b,
work_units,
*,
pin_memory: bool,
quantization_a=None,
quantization_b=None,
compute_dtype=torch.float32,
):
"""Yield logical A/B tensors and release their source keys per work unit."""
def source_keys(quantization, logical_key):
if logical_key is None:
return ()
return tuple(
quantization.required_keys(logical_key)
if quantization is not None
else (logical_key,)
)
required_a = [source_keys(quantization_a, unit[1]) for unit in work_units]
required_b = [source_keys(quantization_b, unit[2]) for unit in work_units]
keys_a = [key for keys in required_a for key in keys]
keys_b = [key for keys in required_b for key in keys]
cursor_a = AsyncTensorCursor(handler_a, keys_a, pin_memory=pin_memory)
cursor_b = AsyncTensorCursor(handler_b, keys_b, pin_memory=pin_memory)
try:
for (logical_key, key_a, key_b), keys_for_a, keys_for_b in zip(
work_units, required_a, required_b
):
tensors_a = {}
tensors_b = {}
consumed_a = []
consumed_b = []
tensor_a = tensor_b = None
try:
for source_key in keys_for_a:
tensors_a[source_key] = cursor_a.take(source_key)
consumed_a.append(source_key)
for source_key in keys_for_b:
tensors_b[source_key] = cursor_b.take(source_key)
consumed_b.append(source_key)
tensor_a = (
quantization_a.decode(
key_a, tensors_a, compute_dtype, device="cpu"
)
if quantization_a is not None and key_a in quantization_a.quantized_keys
else tensors_a.get(key_a)
)
tensor_b = (
quantization_b.decode(
key_b, tensors_b, compute_dtype, device="cpu"
)
if quantization_b is not None and key_b in quantization_b.quantized_keys
else tensors_b.get(key_b)
)
yield logical_key, tensor_a, tensor_b
finally:
del tensor_a, tensor_b, tensors_a, tensors_b
for source_key in consumed_a:
cursor_a.release(source_key)
for source_key in consumed_b:
cursor_b.release(source_key)
cursor_a.finish()
cursor_b.finish()
finally:
cursor_a.close()
cursor_b.close()
+284
View File
@@ -0,0 +1,284 @@
"""Parsed per-layer numeric overrides shared by merge, extraction and resize."""
from dataclasses import dataclass
import fnmatch
import math
import re
from comfy_api.latest import io
LAYER_PARAMETERS = io.Custom("MODELUTILS_LAYER_PARAMETERS")
@dataclass(frozen=True)
class LayerParameterRule:
line: int
pattern: str
matcher: re.Pattern
assignments: tuple[tuple[str, int | float], ...]
@dataclass(frozen=True)
class LayerParameterRules:
rules: tuple[LayerParameterRule, ...]
glob_patterns: bool = False
@dataclass(frozen=True)
class Parameter:
name: str
aliases: tuple[str, ...]
minimum: int | float
maximum: int | float
integral: bool = False
MERGE_COEFFICIENTS = ("alpha", "beta", "gamma", "delta", "epsilon", "zeta")
MERGE_MODES = {
"Weight-Sum": ("alpha",),
"Train-Difference": ("alpha",),
"Add-Difference": ("alpha", "beta"),
"Power-Up (DARE)": ("alpha", "beta"),
"Comparative-Interpolation": ("alpha", "beta", "gamma"),
"Add-Dissimilarities": ("alpha", "beta", "gamma"),
"SVD LoRA Extraction": ("alpha", "beta", "gamma"),
"Enhanced Auto Interp": ("alpha", "beta", "gamma"),
"Weight-Sum Cutoff": ("alpha", "beta", "gamma"),
"Power-Up (DARE+TIES)": ("alpha", "beta", "gamma"),
"Extract-Features": ("alpha", "beta", "gamma", "delta"),
"Enhanced Man Interp": ("alpha", "beta", "gamma", "delta"),
"Power-Up Enhanced (DARE)": ("alpha", "beta", "gamma", "delta"),
"Power-Up Enhanced (DARE+TIES)": ("beta", "gamma", "delta", "epsilon", "zeta"),
}
EXTRACTION_BINDINGS = {
"fixed": ("linear_dim", "conv_dim", "clamp_quantile", "min_diff"),
"ratio": ("linear_ratio", "conv_ratio", "clamp_quantile", "min_diff", "linear_max_rank", "conv_max_rank"),
"quantile": ("linear_quantile", "conv_quantile", "clamp_quantile", "min_diff", "linear_max_rank", "conv_max_rank"),
"frobenius": ("linear_target", "conv_target", "clamp_quantile", "min_diff", "linear_max_rank", "conv_max_rank"),
"knee": ("linear_max_rank", "conv_max_rank", "clamp_quantile", "min_diff"),
}
RESIZE_BINDINGS = {
"fixed": ("new_rank",),
"ratio": ("max_rank", "ratio"),
"frobenius": ("max_rank", "min_rank", "target"),
"cumulative": ("max_rank", "target"),
}
_ASSIGNMENT = re.compile(r"([a-z][a-z0-9_]*)\s*:\s*([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)")
_SEPARATOR = re.compile(r"(?:\s*,\s*|\s+)")
def parse_rules(text: str, glob_patterns: bool = False) -> LayerParameterRules:
rules = []
for line_number, raw in enumerate(text.splitlines(), 1):
line = raw.strip()
if not line or line.startswith("#"):
continue
row = re.fullmatch(r"\((.*)\)\s+(.+)", line)
if row is None or not row[1].strip():
raise ValueError(f"Layer parameters line {line_number}: expected (pattern) name:value assignments.")
pattern, values = row.groups()
try:
matcher = re.compile(fnmatch.translate(f"*{pattern}*") if glob_patterns else pattern)
except re.error as exc:
raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: invalid regex: {exc}") from exc
assignments = []
seen = set()
position = 0
while position < len(values):
match = _ASSIGNMENT.match(values, position)
if match is None:
raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: invalid assignment near {values[position:]!r}.")
name, literal = match.groups()
if name in seen:
raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: duplicate parameter {name!r}.")
try:
value = float(literal) if any(c in literal for c in ".eE") else int(literal)
finite = math.isfinite(value)
except (ValueError, OverflowError):
finite = False
if not finite:
raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: {name} must be finite.")
assignments.append((name, value))
seen.add(name)
position = match.end()
if position == len(values):
break
separator = _SEPARATOR.match(values, position)
if separator is None or separator.end() == len(values):
raise ValueError(f"Layer parameters line {line_number}, pattern {pattern!r}: expected another assignment after a comma or whitespace.")
position = separator.end()
rules.append(LayerParameterRule(line_number, pattern, matcher, tuple(assignments)))
return LayerParameterRules(tuple(rules), glob_patterns)
def parameter_bindings(profile: str) -> tuple[Parameter, ...]:
family, _, mode = profile.partition(":")
if family == "merge" and mode in MERGE_MODES:
parameters = []
probabilities = {
"Power-Up (DARE)": {"alpha"},
"Power-Up (DARE+TIES)": {"alpha", "beta"},
"Power-Up Enhanced (DARE)": {"beta", "gamma"},
"Power-Up Enhanced (DARE+TIES)": {"beta", "epsilon", "zeta"},
"SVD LoRA Extraction": {"gamma"},
}.get(mode, set())
for name in MERGE_MODES[mode]:
aliases = ("abcdef"[MERGE_COEFFICIENTS.index(name)],)
if name == "gamma":
aliases += ("g",)
if name == "zeta":
aliases += ("z",)
integral = mode == "SVD LoRA Extraction" and name in {"alpha", "beta"}
minimum, maximum = (0, 1) if name in probabilities else (-10, 10)
if integral:
minimum = 1
parameters.append(Parameter(name, aliases, minimum, maximum, integral))
return tuple(parameters)
if family == "extract" and mode in EXTRACTION_BINDINGS:
names = EXTRACTION_BINDINGS[mode]
elif family == "resize" and mode in RESIZE_BINDINGS:
names = RESIZE_BINDINGS[mode]
else:
raise ValueError(f"Unsupported layer parameter bindings: {profile!r}")
parameters = []
for alias, name in zip("abcdef", names):
integral = "rank" in name or name.endswith("_dim")
if integral:
minimum, maximum = 1, 16384 if family == "extract" else 3072
elif "ratio" in name:
minimum, maximum = 1, 100
elif name == "clamp_quantile":
minimum, maximum = 0.5, 1
else:
minimum, maximum = (0.1 if family == "resize" else 0), 1
parameters.append(Parameter(name, (alias,), minimum, maximum, integral))
return tuple(parameters)
def _validated_value(parameter, value, context):
if not isinstance(value, (int, float)) or isinstance(value, bool) or not math.isfinite(value):
raise ValueError(f"{context}: {parameter.name} must be a finite number, got {value!r}.")
if parameter.integral and int(value) != value:
raise ValueError(f"{context}: {parameter.name} must be integral, got {value}.")
if not parameter.minimum <= value <= parameter.maximum:
raise ValueError(f"{context}: {parameter.name} must be in [{parameter.minimum}, {parameter.maximum}], got {value}.")
return int(value) if parameter.integral else value
def resolve_layer_parameters(layer_parameters, profile, layers, defaults, *, node_name):
"""Validate against header-only target names and return matched scalar overrides."""
if layer_parameters is None:
return {}
if not isinstance(layer_parameters, LayerParameterRules):
raise TypeError(f"{node_name}: layer_parameters must come from Layer Parameter Configuration.")
if not layer_parameters.rules:
return {}
bindings = parameter_bindings(profile)
aliases = {alias: parameter for parameter in bindings for alias in (*parameter.aliases, parameter.name)}
assignments = []
for rule in layer_parameters.rules:
resolved = {}
for name, value in rule.assignments:
parameter = aliases.get(name)
context = f"{node_name}: line {rule.line}, pattern {rule.pattern!r}"
if parameter is None:
allowed = ", ".join(p.name for p in bindings)
raise ValueError(f"{context}: unsupported parameter {name!r} for {profile}; allowed: {allowed}.")
if parameter.name in resolved:
raise ValueError(f"{context}: duplicate aliases for {parameter.name}.")
resolved[parameter.name] = _validated_value(parameter, value, context)
assignments.append(resolved)
hits = set()
result = {}
for layer in sorted(set(layers)):
matched = [i for i, rule in enumerate(layer_parameters.rules) if rule.matcher.search(layer)]
if len(matched) > 1:
lines = "; ".join(f"line {layer_parameters.rules[i].line} ({layer_parameters.rules[i].pattern!r})" for i in matched)
raise ValueError(f"{node_name}: layer {layer!r} matches conflicting rules: {lines}.")
if not matched:
continue
index = matched[0]
values = dict(defaults)
values.update(assignments[index])
if profile == "extract:fixed":
for dimension, cap in (("linear_dim", "linear_max_rank"), ("conv_dim", "conv_max_rank")):
if dimension in assignments[index] and cap in defaults:
values[cap] = values[dimension]
rule = layer_parameters.rules[index]
context = f"{node_name}: layer {layer!r}, line {rule.line}, pattern {rule.pattern!r}"
for parameter in bindings:
if parameter.name in values:
values[parameter.name] = _validated_value(parameter, values[parameter.name], context)
if profile == "resize:frobenius" and values["min_rank"] > values["max_rank"]:
raise ValueError(f"{node_name}: layer {layer!r}, line {rule.line}, pattern {rule.pattern!r}: min_rank={values['min_rank']} exceeds max_rank={values['max_rank']}.")
result[layer] = values
hits.add(index)
unmatched = [rule for i, rule in enumerate(layer_parameters.rules) if i not in hits]
if unmatched:
details = "; ".join(f"line {rule.line}, pattern {rule.pattern!r}" for rule in unmatched)
raise ValueError(f"{node_name}: no target layer matches {details}.")
return result
def _binding_description(profile):
return "; ".join(f"{'/'.join(p.aliases)}={p.name}" for p in parameter_bindings(profile))
def parameter_input(profile):
if profile == "merge":
parameters = {p.name: p for mode in MERGE_MODES for p in parameter_bindings(f"merge:{mode}")}
description = "; ".join(f"{'/'.join(parameters[name].aliases)}={name}" for name in MERGE_COEFFICIENTS)
description += ". Only coefficients used by Calc Mode are accepted."
else:
description = _binding_description(profile)
return LAYER_PARAMETERS.Input(
"layer_parameters", optional=True,
tooltip="Optional Layer Parameter Configuration. " + description + " Full names are also accepted. Unassigned values use this node's settings; existing filters still apply.",
)
def documentation():
lines = [
"# Per-layer parameter rules", "",
"Connect Layer Parameter Configuration to the optional layer_parameters socket on standard two/three-input mergers, all 25 extraction variants, or the four LoRA resizers.", "",
"Write (pattern) a:value b:value, or use full parameter names. Commas and whitespace may separate assignments; whitespace around colons is allowed. Signed decimal/scientific numbers are accepted. Blank lines and whole-line # comments are ignored. Inline comments are not supported.", "",
"Patterns preserve backslashes, nested groups, spaces and character classes. Regex uses substring search; Glob Patterns uses case-sensitive glob substring matching. These settings do not change the main node's filters.", "",
"Unconnected or empty configuration changes nothing. Unmatched layers and omitted parameters use main-node settings. Existing exclusions, inclusion filters, discards, mismatch rules and safety guards still apply.", "",
"Invalid syntax/values, duplicate assignments (including aliases), unsupported parameters, overlapping rule lines and rules matching no target layer are errors. Matching and validation happen before tensor streaming or output writing. Name matches are counted before exclusions.", "",
"LoRA rules target normalized logical block names, not factor/alpha suffixes. Existing dotted normalization and reference mapping are retained, including underscores inside names such as qkv_proj. Unresolved flattened names retain the existing fallback; patterns are not silently rewritten.", "",
"## Bindings and bounds", "", "| Receiver/method | Short aliases and full names | Bounds |", "|---|---|---|",
]
profiles = [*(f"merge:{name}" for name in MERGE_MODES), *(f"extract:{name}" for name in EXTRACTION_BINDINGS), *(f"resize:{name}" for name in RESIZE_BINDINGS)]
for profile in profiles:
bounds = "; ".join(f"{p.name}: {'integer ' if p.integral else ''}[{p.minimum}, {p.maximum}]" for p in parameter_bindings(profile))
lines.append(f"| {profile} | {_binding_description(profile)} | {bounds} |")
lines.extend([
"", "Frobenius resize requires resolved min_rank <= max_rank. Existing tensor-dimension rank caps remain. Fixed extraction dimension overrides also update their corresponding rank caps. Probe offsets, modes, iterations, optimizer settings, seeds, files, dtype and device remain main-node settings.",
"", "## Examples", "", "Standard Extract-Features merge:", "", "```text",
r"(blocks\.4[589]\.attn\.qkv_proj) a:0.5 b:0.25 g:0.75 d:1.0", "```", "",
"Fixed extraction:", "", "```text",
r"(blocks\.4[589]\.attn\.qkv_proj) a:64, b:32, c:0.99, d:0.0", "```", "",
"Frobenius resize:", "", "```text",
r"(blocks\.4[589]\.attn\.qkv_proj) max_rank:128 min_rank:1 target:0.9", "```",
])
return "\n".join(lines)
class LayerParameterConfiguration(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LayerParameterConfiguration", display_name="Layer Parameter Configuration",
category="ModelUtils/Configuration",
description="Configure per-layer numeric overrides for standard merging, extraction and LoRA resizing.",
inputs=[
io.String.Input("rules", default="", multiline=True, tooltip="One (pattern) name:value rule per line. Short aliases or full names are accepted; separate assignments with whitespace or commas. See Documentation for receiver-specific bindings."),
io.Boolean.Input("glob_patterns", default=False, tooltip="Use case-sensitive glob substring matching instead of regex substring search. Independent of receiving nodes' exclusion-filter syntax."),
],
outputs=[LAYER_PARAMETERS.Output(display_name="layer_parameters"), io.String.Output(display_name="documentation")],
)
@classmethod
def execute(cls, rules="", glob_patterns=False):
return io.NodeOutput(parse_rules(rules, glob_patterns), documentation())
+313
View File
@@ -0,0 +1,313 @@
import fnmatch
import gc
import re
import comfy.utils
import folder_paths
import torch
from comfy_api.latest import io
from tqdm import tqdm
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .artifact_paths import canonical_model_artifact_path
from .device_utils import cleanup_after_operation, estimate_model_size, prepare_for_large_operation
from .lora_alpha import normalize_lora_pair
from .lora_resize import build_lora_layer_map, canonical_lora_key, parse_lora_layers, select_output_dtype, validate_canonical_blocks
from .quantization_guard import inspect_low_bit_input
from .uel_io import AsyncTensorCursor, atomic_uel_writer
LODESTONE_METHODS = ("sum", "mean", "slotnorm", "normmatch", "slotnorm-normmatch")
_ALLOWED_ROLES = {
"down", "up", "alpha", "diff", "format", "down_suffix", "up_suffix",
"alpha_suffix",
}
def _matches(key, patterns, glob_mode):
entries = [entry.strip() for entry in patterns.splitlines() if entry.strip()]
if glob_mode:
return any(fnmatch.fnmatchcase(key, entry) for entry in entries)
return any(re.search(entry, key) is not None for entry in entries)
def factor_delta(down, up, layer):
if down.ndim == 2 and up.ndim == 2:
if down.shape[0] != up.shape[1]:
raise ValueError(f"Lodestone rank mismatch for '{layer}'.")
return up.float() @ down.float()
if down.ndim == 4 and up.ndim == 4 and tuple(up.shape[2:]) == (1, 1):
rank, channels, height, width = down.shape
if up.shape[1] != rank:
raise ValueError(f"Lodestone rank mismatch for '{layer}'.")
delta = up.reshape(up.shape[0], rank).float() @ down.reshape(rank, -1).float()
return delta.reshape(up.shape[0], channels, height, width)
raise ValueError(
f"Lodestone layer '{layer}' requires linear factors or convolutional "
"factors with a 1x1 up kernel."
)
def merge_lodestone_deltas(deltas, method):
if method not in LODESTONE_METHODS:
raise ValueError(f"Unknown Lodestone method '{method}'.")
if not deltas:
raise ValueError("Lodestone requires at least one delta.")
shape = tuple(deltas[0].shape)
if any(tuple(delta.shape) != shape for delta in deltas[1:]):
raise ValueError("Lodestone source deltas must have matching shapes.")
norms = torch.stack([torch.linalg.vector_norm(delta.float()) for delta in deltas])
target = torch.median(norms)
values = deltas
if method in {"slotnorm", "slotnorm-normmatch"}:
values = [
delta * (target / norm.clamp_min(1e-12))
for delta, norm in zip(deltas, norms)
]
merged = torch.stack(values).sum(dim=0)
if method in {"mean", "slotnorm"}:
merged /= len(values)
elif method in {"normmatch", "slotnorm-normmatch"}:
merged_norm = torch.linalg.vector_norm(merged.float())
if merged_norm > 0:
merged *= target / merged_norm.clamp_min(1e-12)
return merged
def _to_device(tensor, device):
if device.type == "cuda":
return transfer_to_gpu_pinned(tensor, device, torch.float32)
return tensor.to(device=device, dtype=torch.float32)
def _release_cuda_oom():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def _is_cuda_oom(error):
return isinstance(error, torch.cuda.OutOfMemoryError) or (
isinstance(error, RuntimeError) and "out of memory" in str(error).lower()
)
def _build_units(layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns, include_mode=False):
units = []
stream_keys = [[] for _ in handlers]
for core, mapped in layer_map.items():
if not any(index == 0 for index, _ in mapped):
continue
output_key = canonical_lora_key(core, "diff")
if _matches(output_key, discard_patterns, glob_patterns):
continue
matched = _matches(output_key, exclude_patterns, glob_patterns)
preserve = not matched if include_mode else matched
sources = []
for index, block in sorted(mapped):
roles = infos[index]["pairs"][block]
unsupported = set(roles) - _ALLOWED_ROLES
if unsupported:
raise ValueError(f"Lodestone layer '{block}' has unsupported roles: {', '.join(sorted(unsupported))}.")
if "down" in roles or "up" in roles:
if not all(role in roles for role in ("down", "up")):
raise ValueError(f"Lodestone layer '{block}' has an incomplete factor pair.")
keys = [roles["down"], roles["up"]]
if any(key in infos[index]["low_bit"] for key in keys):
raise ValueError(f"Lodestone cannot expand low-bit factors for '{block}'.")
if "alpha" in roles:
keys.append(roles["alpha"])
sources.append((index, block, "pair", keys))
elif "diff" in roles:
key = roles["diff"]
if key in infos[index]["low_bit"]:
raise ValueError(f"Lodestone cannot merge low-bit delta '{block}'.")
sources.append((index, block, "diff", [key]))
if preserve:
sources = [source for source in sources if source[0] == 0]
present = {source[0] for source in sources}
if len(present) < len(handlers) and not preserve:
if mismatch_mode == "error":
raise ValueError(f"Lodestone layer '{core}' is missing from one or more inputs.")
if mismatch_mode == "skip":
sources = [source for source in sources if source[0] == 0]
if not sources:
continue
dtypes = []
for index, _block, kind, keys in sources:
stream_keys[index].extend(keys)
dtypes.append(handlers[index].get_dtype(keys[0] if kind == "diff" else keys[1]))
units.append({
"core": core,
"output_key": output_key,
"sources": sources,
"missing": 0 if preserve or mismatch_mode != "zeros" else len(handlers) - len(present),
"dtypes": dtypes,
})
return units, stream_keys
def _process_on_device(unit, loaded, method, save_dtype, device):
deltas = []
for index, block, kind, keys in unit["sources"]:
if kind == "diff":
delta = _to_device(loaded[(index, keys[0])], device)
else:
down = loaded[(index, keys[0])]
up = loaded[(index, keys[1])]
alpha = loaded[(index, keys[2])] if len(keys) == 3 else None
down, up = normalize_lora_pair(down, up, alpha, layer=block)
delta = factor_delta(_to_device(down, device), _to_device(up, device), block)
deltas.append(delta)
if unit["missing"]:
deltas.extend(torch.zeros_like(deltas[0]) for _ in range(unit["missing"]))
merged = merge_lodestone_deltas(deltas, method)
dtype = select_output_dtype(unit["dtypes"], save_dtype, is_1d_diff=merged.ndim == 1)
return merged.to(dtype).cpu().contiguous()
def _process(unit, loaded, method, save_dtype, requested_device):
device = torch.device(requested_device)
try:
return _process_on_device(unit, loaded, method, save_dtype, device), False
except Exception as error:
if device.type != "cuda" or not _is_cuda_oom(error):
raise
_release_cuda_oom()
return _process_on_device(unit, loaded, method, save_dtype, torch.device("cpu")), True
def merge_lodestone_loras(lora_paths, method, device, save_dtype, output_filename,
mismatch_mode="skip", exclude_patterns="", discard_patterns="",
glob_patterns=False, verbose=True, include_mode=False):
operation = "Lodestone LoRA Merge"
prepare_for_large_operation(sum(estimate_model_size(path) for path in lora_paths) * 2.5, torch.device(device))
handlers = [MemoryEfficientSafeOpen(path, low_memory=True) for path in lora_paths]
cursors = []
try:
infos = []
for index, handler in enumerate(handlers):
low_bit = inspect_low_bit_input(handler, f"LoRA {index + 1}", operation)
pairs, passthrough = parse_lora_layers(handler.keys())
infos.append({"pairs": pairs, "passthrough": passthrough, "low_bit": low_bit})
layer_map = build_lora_layer_map(infos, None)
validate_canonical_blocks(layer_map, operation)
units, stream_keys = _build_units(
layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns, include_mode
)
cursors = [
AsyncTensorCursor(handler, keys, pin_memory=torch.device(device).type == "cuda")
for handler, keys in zip(handlers, stream_keys)
]
output_path, _ = canonical_model_artifact_path("loras", output_filename)
metadata = (handlers[0].metadata() or {}).copy()
metadata.update({"merge_method": f"lodestone:{method}", "alpha_normalized": "true", "output_representation": "full_direct_difference"})
fallback_count = 0
pbar = comfy.utils.ProgressBar(len(units))
with atomic_uel_writer(output_path, metadata) as writer, torch.no_grad():
for unit in tqdm(units, desc=f"Lodestone layers ({method})", unit="layers"):
loaded = {}
try:
for index, _block, _kind, keys in unit["sources"]:
for key in keys:
loaded[(index, key)] = cursors[index].take(key)
output, fallback = _process(unit, loaded, method, save_dtype, device)
writer.write_batch([(unit["output_key"], output)])
fallback_count += int(fallback)
del output
finally:
for index, key in list(loaded):
del loaded[(index, key)]
cursors[index].release(key)
loaded.clear()
pbar.update(1)
for cursor in cursors:
cursor.finish()
if verbose:
print(f"[{operation}] CUDA OOM CPU fallbacks: {fallback_count}")
print(f"[{operation}] Saved {len(units)} full-delta layers to {output_path}")
return output_path
finally:
for cursor in cursors:
cursor.close()
for handler in handlers:
handler.__exit__(None, None, None)
cleanup_after_operation()
def _controls(default_filename):
return [
io.Combo.Input("calc_mode", options=list(LODESTONE_METHODS), default="sum", tooltip="How full LoRA weight changes are combined per layer. sum adds them; mean averages them; slotnorm equalizes each input's Frobenius magnitude before averaging; normmatch adds them and scales the result to a typical input magnitude; slotnorm-normmatch performs both normalization steps."),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", tooltip="What to do when a layer exists in LoRA 1 but is absent from another input. skip copies LoRA 1's layer unchanged; zeros treats each missing input as a zero update; error stops without replacing an existing output file."),
io.String.Input("output_filename", default=default_filename, tooltip="Name of the new .safetensors file. It is saved in ComfyUI's loras folder; omit the extension."),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"], default="bf16", tooltip="Precision used to save the expanded full-weight difference tensors. FP32 is largest and most precise; FP16 is smallest; BF16 has wider numeric range than FP16."),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", tooltip="Where each full layer is expanded and merged. CUDA is faster; if a layer runs out of VRAM, that layer is automatically retried on CPU."),
io.String.Input("exclude_patterns", default="", multiline=True, tooltip="Optional layer-name patterns to copy from LoRA 1 instead of merging. Enter one pattern per line; useful for protecting specific blocks."),
io.String.Input("discard_patterns", default="", multiline=True, tooltip="Optional layer-name patterns to leave out of the saved LoRA completely. Enter one pattern per line; discarded layers cannot affect the model."),
io.Boolean.Input("glob_patterns", default=False, tooltip="Choose filter syntax. Off uses regular expressions; on uses shell-style globs where * matches any text and ? matches one character."),
io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from LoRA 1."),
]
def _execute(kwargs, count):
paths = [folder_paths.get_full_path_or_raise("loras", kwargs[f"lora_{index}"]) for index in range(1, count + 1)]
dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[kwargs["save_dtype"]]
merge_lodestone_loras(paths, kwargs["calc_mode"], kwargs["process_device"], dtype,
kwargs["output_filename"], kwargs["mismatch_mode"],
kwargs["exclude_patterns"], kwargs["discard_patterns"], kwargs["glob_patterns"], include_mode=kwargs.get("include_mode", False))
return canonical_model_artifact_path("loras", kwargs["output_filename"])[1]
class LodestoneLoRATwoMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
options = folder_paths.get_filename_list("loras")
return io.Schema(node_id="LodestoneLoRATwoMerger", display_name="Lodestone Merge LoRAs (2)",
category="ModelUtils/LoRA/Merge/Lodestone",
description="Expand two LoRAs into full per-layer weight changes, combine them with the selected Frobenius-norm method, and save a full-difference LoRA.",
inputs=[io.Combo.Input("lora_1", options=options, tooltip="First LoRA to merge. It also supplies metadata and is the layer kept when Missing Layer Handling is set to skip."), io.Combo.Input("lora_2", options=options, tooltip="Second LoRA to merge. It contributes equally with LoRA 1 wherever both contain the same layer."), *_controls("lodestone_merged_2_lora")],
outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True)
@classmethod
def execute(cls, **kwargs):
return io.NodeOutput(_execute(kwargs, 2))
class LodestoneLoRAThreeMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
options = folder_paths.get_filename_list("loras")
return io.Schema(node_id="LodestoneLoRAThreeMerger", display_name="Lodestone Merge LoRAs (3)",
category="ModelUtils/LoRA/Merge/Lodestone",
description="Expand three LoRAs into full per-layer weight changes, combine them with the selected Frobenius-norm method, and save a full-difference LoRA.",
inputs=[io.Combo.Input("lora_1", options=options, tooltip="First LoRA to merge. It supplies metadata, defines the output layer set, and is kept when Missing Layer Handling is skip."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-strength LoRA contribution for every matching layer."), io.Combo.Input("lora_3", options=options, tooltip="Third equal-strength LoRA contribution for every matching layer."), *_controls("lodestone_merged_3_lora")],
outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True)
@classmethod
def execute(cls, **kwargs):
return io.NodeOutput(_execute(kwargs, 3))
class LodestoneLoRAMultiMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
options = folder_paths.get_filename_list("loras")
optional = ["None", *options]
inputs = [io.Combo.Input("lora_count", options=[str(index) for index in range(2, 9)], default="2", tooltip="Number of consecutive LoRA selectors to merge, starting at LoRA 1. Every included LoRA has equal influence before the selected normalization method."),
io.Combo.Input("lora_1", options=options, tooltip="First LoRA to merge. It supplies metadata, defines the output layer set, and anchors missing-layer handling."), io.Combo.Input("lora_2", options=options, tooltip="Second equal-strength LoRA contribution. LoRA Count always includes this input.")]
inputs.extend(io.Combo.Input(f"lora_{index}", options=optional, default="None", tooltip="Additional equal-strength LoRA contribution. Select a file for every input included by LoRA Count.") for index in range(3, 9))
inputs.extend(_controls("lodestone_merged_multi_lora"))
return io.Schema(node_id="LodestoneLoRAMultiMerger", display_name="Lodestone LoRA Multi-Merge",
category="ModelUtils/LoRA/Merge/Lodestone", description="Expand 2-8 equal-strength LoRAs into full per-layer weight changes, combine them with Frobenius-norm-aware math, and save a full-difference LoRA.", inputs=inputs,
outputs=[io.AnyType.Output(display_name="output_filename")], is_output_node=True)
@classmethod
def execute(cls, **kwargs):
count = int(kwargs["lora_count"])
missing = [index for index in range(1, count + 1) if not kwargs[f"lora_{index}"] or kwargs[f"lora_{index}"] == "None"]
if missing:
raise ValueError(f"LoRA Count includes unselected input(s): {', '.join(map(str, missing))}.")
return io.NodeOutput(_execute(kwargs, count))
LODESTONE_MERGER_NODES = [LodestoneLoRATwoMerger, LodestoneLoRAThreeMerger, LodestoneLoRAMultiMerger]
+31
View File
@@ -0,0 +1,31 @@
import torch
def lora_alpha_scale(alpha: torch.Tensor | None, rank: int, *, layer: str) -> float:
"""Return the effective LoRA multiplier represented by alpha and rank."""
if rank <= 0:
raise ValueError(f"LoRA rank for '{layer}' must be positive.")
if alpha is None:
return 1.0
if alpha.numel() != 1:
raise ValueError(f"LoRA alpha for '{layer}' must be scalar.")
return float(alpha.reshape(-1)[0].item()) / rank
def normalize_lora_pair(
down: torch.Tensor,
up: torch.Tensor,
alpha: torch.Tensor | None,
*,
layer: str,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Materialize an alpha-free factor pair before rank alignment or merging."""
if down.ndim < 2 or up.ndim < 2:
raise ValueError(f"LoRA factors for '{layer}' must have at least two dimensions.")
rank = int(down.shape[0])
if int(up.shape[1]) != rank:
raise ValueError(f"LoRA factor rank mismatch for '{layer}'.")
scale = lora_alpha_scale(alpha, rank, layer=layer)
return down, up if scale == 1.0 else up * scale
-350
View File
@@ -1,350 +0,0 @@
"""
Memory-efficient LoRA extraction node.
This node loads two models from files and extracts a LoRA by computing the difference,
without requiring the full models to be loaded into GPU memory via ComfyUI's model management.
"""
import os
import torch
import folder_paths
import comfy.utils
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_utils import MemoryEfficientSafeOpen
CLAMP_QUANTILE = 0.99
def extract_lora(diff, rank):
"""Extract LoRA weights from a weight difference tensor using SVD."""
conv2d = (len(diff.shape) == 4)
kernel_size = None if not conv2d else diff.size()[2:4]
conv2d_3x3 = conv2d and kernel_size != (1, 1)
out_dim, in_dim = diff.size()[0:2]
rank = min(rank, in_dim, out_dim)
if conv2d:
if conv2d_3x3:
diff = diff.flatten(start_dim=1)
else:
diff = diff.squeeze()
U, S, Vh = torch.linalg.svd(diff.float())
U = U[:, :rank]
S = S[:rank]
U = U @ torch.diag(S)
Vh = Vh[:rank, :]
dist = torch.cat([U.flatten(), Vh.flatten()])
hi_val = torch.quantile(dist, CLAMP_QUANTILE)
low_val = -hi_val
U = U.clamp(low_val, hi_val)
Vh = Vh.clamp(low_val, hi_val)
if conv2d:
U = U.reshape(out_dim, rank, 1, 1)
Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
return (U, Vh)
def extract_lora_from_files(
model_a_path: str,
model_b_path: str,
rank: int,
lora_type: str,
bias_diff: bool,
prefix_lora: str,
process_device: str,
save_dtype: str,
) -> dict[str, torch.Tensor]:
"""
Memory-efficient LoRA extraction by streaming tensors from safetensors files.
Args:
model_a_path: Path to the finetuned/modified model
model_b_path: Path to the base/original model
rank: LoRA rank for SVD decomposition
lora_type: "standard" for LoRA, "full_diff" for full difference
bias_diff: Whether to include bias differences
prefix_lora: Prefix for LoRA keys (e.g., "diffusion_model.")
process_device: Device for processing ("cuda" or "cpu")
save_dtype: Output dtype ("fp16", "bf16", "fp32")
Returns:
Dictionary of LoRA tensors
"""
output_sd = {}
# Get dtype for saving
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16
}.get(save_dtype, torch.float16)
# Open both files with memory-efficient handler
handler_a = MemoryEfficientSafeOpen(model_a_path)
handler_b = MemoryEfficientSafeOpen(model_b_path)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
# Only process keys that exist in both models
common_keys = keys_a & keys_b
weight_keys = [k for k in common_keys if k.endswith(".weight")]
bias_keys = [k for k in common_keys if k.endswith(".bias")] if bias_diff else []
# Log key differences
missing_in_b = keys_a - keys_b
extra_in_b = keys_b - keys_a
if missing_in_b:
print(f"[LoRA Extract] {len(missing_in_b)} keys in model A not in model B (skipped)")
if extra_in_b:
print(f"[LoRA Extract] {len(extra_in_b)} keys in model B not in model A (ignored)")
pbar = comfy.utils.ProgressBar(len(weight_keys) + len(bias_keys))
for key in weight_keys:
# Load tensors one at a time
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
# Compute difference: finetuned - base
weight_diff = tensor_a - tensor_b
# Clean up immediately
del tensor_a, tensor_b
# Extract LoRA key name
lora_key = key[:-7] # Remove ".weight"
if lora_type == "standard":
if weight_diff.ndim < 2:
# 1D weights (e.g., LayerNorm) - store as diff if bias_diff enabled
if bias_diff:
output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu()
else:
try:
lora_up, lora_down = extract_lora(weight_diff, rank)
output_sd[f"{prefix_lora}{lora_key}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu()
output_sd[f"{prefix_lora}{lora_key}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu()
except Exception as e:
print(f"[LoRA Extract] Could not extract LoRA for {key}: {e}")
elif lora_type == "full_diff":
output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu()
del weight_diff
pbar.update(1)
# Process bias keys
for key in bias_keys:
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
bias_diff_tensor = tensor_a - tensor_b
del tensor_a, tensor_b
lora_key = key[:-5] # Remove ".bias"
output_sd[f"{prefix_lora}{lora_key}.diff_b"] = bias_diff_tensor.to(save_torch_dtype).cpu()
del bias_diff_tensor
pbar.update(1)
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
return output_sd
class LoRASaveFromFile(io.ComfyNode):
"""
Memory-efficient LoRA extraction from safetensors files.
Unlike the standard LoraSave node which requires pre-loaded and subtracted models,
this node loads models directly from files one tensor at a time, making it much
more memory efficient for extracting LoRAs from large models.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRASaveFromFile",
display_name="Extract LoRA from Files (Memory Efficient)",
category="ModelUtils/LoRA",
description="Extracts a LoRA by computing the difference between two model files. "
"Much more memory efficient than using ModelSubtract + LoraSave.",
inputs=[
io.Combo.Input(
"model_a",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="The finetuned/modified model (A - B = LoRA)"
),
io.Combo.Input(
"model_b",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="The base/original model (A - B = LoRA)"
),
io.Int.Input("rank", default=64, min=1, max=4096, step=1,
tooltip="LoRA rank - higher preserves more detail but increases size"),
io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard",
tooltip="standard: SVD decomposition to lora_up/lora_down, full_diff: raw weight differences"),
io.Boolean.Input("bias_diff", default=True,
tooltip="Include bias differences in the LoRA"),
io.String.Input("output_filename", default="extracted_lora",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_a: str,
model_b: str,
rank: int,
lora_type: str,
bias_diff: bool,
output_filename: str,
save_dtype: str,
process_device: str,
) -> io.NodeOutput:
# Get full paths
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)
# Extract LoRA
output_sd = extract_lora_from_files(
model_a_path=model_a_path,
model_b_path=model_b_path,
rank=rank,
lora_type=lora_type,
bias_diff=bias_diff,
prefix_lora="diffusion_model.",
process_device=process_device,
save_dtype=save_dtype,
)
if not output_sd:
raise ValueError("No LoRA weights extracted - models may be identical or incompatible")
# Save to loras folder
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path)
print(f"[LoRA Extract] Saved {len(output_sd)} tensors to {output_path}")
return io.NodeOutput(output_path)
class LoRACheckpointSaveFromFile(io.ComfyNode):
"""
Memory-efficient LoRA extraction from checkpoint files (which may contain both model and text encoder).
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRACheckpointSaveFromFile",
display_name="Extract LoRA from Checkpoints (Memory Efficient)",
category="ModelUtils/LoRA",
description="Extracts a LoRA by computing the difference between two checkpoint files. "
"Handles both diffusion model and text encoder weights.",
inputs=[
io.Combo.Input(
"checkpoint_a",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="The finetuned/modified checkpoint (A - B = LoRA)"
),
io.Combo.Input(
"checkpoint_b",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="The base/original checkpoint (A - B = LoRA)"
),
io.Int.Input("rank", default=64, min=1, max=4096, step=1,
tooltip="LoRA rank - higher preserves more detail but increases size"),
io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard",
tooltip="standard: SVD decomposition, full_diff: raw weight differences"),
io.Boolean.Input("bias_diff", default=True,
tooltip="Include bias differences in the LoRA"),
io.String.Input("output_filename", default="extracted_lora",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
checkpoint_a: str,
checkpoint_b: str,
rank: int,
lora_type: str,
bias_diff: bool,
output_filename: str,
save_dtype: str,
process_device: str,
) -> io.NodeOutput:
# Get full paths
checkpoint_a_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_a)
checkpoint_b_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_b)
# Open handlers to detect prefixes
handler_a = MemoryEfficientSafeOpen(checkpoint_a_path)
keys = handler_a.keys()
handler_a.__exit__(None, None, None)
# Auto-detect key prefixes
has_model = any(k.startswith("model.") for k in keys)
has_diffusion = any(k.startswith("diffusion_model.") for k in keys)
has_text_encoders = any(k.startswith("text_encoders.") or k.startswith("cond_stage_model.") for k in keys)
output_sd = {}
# Extract diffusion model LoRA
if has_model or has_diffusion:
prefix = "model.diffusion_model." if has_model else "diffusion_model."
model_lora = extract_lora_from_files(
model_a_path=checkpoint_a_path,
model_b_path=checkpoint_b_path,
rank=rank,
lora_type=lora_type,
bias_diff=bias_diff,
prefix_lora="diffusion_model.",
process_device=process_device,
save_dtype=save_dtype,
)
output_sd.update(model_lora)
# Note: Text encoder extraction would need additional prefix handling
# For full checkpoint support, this would need to be expanded
if not output_sd:
raise ValueError("No LoRA weights extracted - checkpoints may be identical or incompatible")
# Save to loras folder
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path)
print(f"[LoRA Extract] Saved {len(output_sd)} tensors to {output_path}")
return io.NodeOutput(output_path)
File diff suppressed because it is too large Load Diff
+1160
View File
File diff suppressed because it is too large Load Diff
+248
View File
@@ -0,0 +1,248 @@
"""Read-only, bounded analysis of a LoRA's effect on a diffusion model."""
from contextlib import closing
import csv
import gc
import heapq
from io import StringIO
import logging
import folder_paths
import torch
from comfy.weight_adapter import LoRAAdapter
from comfy_api.latest import io
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import cleanup_after_operation
from .lora_alpha import normalize_lora_pair
from .lora_resize import (
AdapterErrorCapture, build_lora_layer_map, canonical_lora_block_name,
layer_tensor_keys, parse_lora_layers, validate_canonical_blocks,
)
from .merger import _compile_patterns, _matches_any_pattern, load_documentation_from_file
from .model_analysis import (
CWBStats, RawStats, TensorRecord, _block_name, _common_inputs,
_compute_pair_on_device, _fmt, _is_cuda_oom, _is_float_dtype, _layer_name,
_metrics, _release_cuda_oom, _render_comparison, _render_cwb,
)
from .quantization_guard import DiffusionQuantization, inspect_low_bit_input
from .uel_io import stream_work_units
def _plan_application(base, lora, base_path, params, base_quant):
pairs, passthrough = parse_lora_layers(lora.keys())
validate_canonical_blocks(pairs, "LoRA-on-model analysis")
layer_map = build_lora_layer_map([{"pairs": pairs}], base_path)
targets = {canonical_lora_block_name(key): key for key in base_quant.data_keys}
patches = {}
mapped = set()
unmapped_targets = []
for core, sources in layer_map.items():
for _, block in sources:
roles = pairs[block]
weight_role = next((role for role in ("set_weight", "diff", "w_norm") if role in roles), None)
specs = []
if weight_role is not None:
target = targets.get(core) if weight_role == "diff" else None
specs.append((target or targets.get(f"{core}.weight"), weight_role))
elif "down" in roles and "up" in roles:
specs.append((targets.get(f"{core}.weight"), "pair"))
bias_role = next((role for role in ("diff_b", "b_norm") if role in roles), None)
if bias_role:
specs.append((targets.get(f"{core}.bias"), bias_role))
for target, role in specs:
if target is not None:
patches.setdefault(target, []).append((block, roles, role))
mapped.add(block)
else:
unmapped_targets.append(f"{block} ({role}: no base target)")
inventory = {name: [] for name in ("a_only", "b_only", "shape", "nonfloat", "low_bit", "excluded_a", "excluded_b")}
unmatched = sorted(set(pairs) - mapped) + unmapped_targets
patterns = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=params.get("glob_patterns", False))
low_base = base_quant.isolated_low_bit_keys
low_lora = inspect_low_bit_input(lora, "LoRA", "LoRA-on-model analysis")
units = []
for key in sorted(base_quant.data_keys):
if _matches_any_pattern(key, patterns, glob_mode=params.get("glob_patterns", False)) != params.get("include_mode", False):
inventory["excluded_a"].append(key)
inventory["excluded_b"].append(key)
continue
specs = patches.get(key, [])
source_keys = list(dict.fromkeys(
source_key for _, roles, role in specs
for source_key in (layer_tensor_keys(roles).values() if role == "pair" else (roles[role],))
))
alpha_keys = {roles["alpha"] for _, roles, role in specs if role == "pair" and "alpha" in roles}
if key in low_base or any(source in low_lora and source not in alpha_keys for source in source_keys):
inventory["low_bit"].append(key)
continue
if not _is_float_dtype(base_quant.logical_dtype(key, torch.float32)):
inventory["nonfloat"].append(key)
continue
entries = {0: base_quant.required_keys(key)}
if source_keys:
entries[1] = source_keys
units.append((key, entries))
return units, patches, inventory, unmatched, passthrough
def _apply_pair(key, weight, block, roles, loaded, strength):
tensors = {source_key: loaded[(1, source_key)] for source_key in layer_tensor_keys(roles).values()}
down = up = adapter = None
try:
down, up = normalize_lora_pair(tensors[roles["down"]], tensors[roles["up"]], tensors.get(roles.get("alpha")), layer=block)
tensors[roles["down"]], tensors[roles["up"]] = down, up
tensors.pop(roles.get("alpha"), None)
adapter = LoRAAdapter.load(block, tensors, None, tensors.get(roles.get("dora_scale")), set())
if adapter is None:
raise ValueError(f"LoRA-on-model analysis: could not load adapter {block!r} for {key!r}.")
capture = AdapterErrorCapture(key)
logger = logging.getLogger()
logger.addHandler(capture)
try:
weight = adapter.calculate_weight(weight, key, strength, 1.0, None, lambda value: value, torch.float32, None)
finally:
logger.removeHandler(capture)
if capture.messages:
raise RuntimeError("; ".join(capture.messages))
return weight
finally:
tensors.clear()
del adapter, down, up, weight
def _apply_and_analyze(key, loaded, specs, strength, device, top_count):
source = loaded[(0, key)]
original = patched = delta = None
try:
original = (
transfer_to_gpu_pinned(source, device=device, dtype=torch.float32)
if str(device).startswith("cuda") else source.to(device=device, dtype=torch.float32)
)
patched = original.clone()
for block, roles, role in specs:
if role == "pair":
patched = _apply_pair(key, patched, block, roles, loaded, strength)
continue
delta = loaded[(1, roles[role])].to(device=device, dtype=torch.float32)
if delta.shape != patched.shape:
raise ValueError(f"LoRA-on-model analysis: {roles[role]} shape {tuple(delta.shape)} does not match {key} {tuple(patched.shape)}.")
if role == "set_weight":
patched.copy_(delta)
else:
patched.add_(delta, alpha=strength)
delta = None
if patched.shape != original.shape:
raise ValueError(f"LoRA-on-model analysis: applying {key!r} changed its shape; it cannot be compared with the base tensor.")
return _compute_pair_on_device(original, patched, device, False, top_count)
finally:
del source, original, patched, delta
def _analyze_unit(key, loaded, specs, strength, device, top_count):
try:
return (*_apply_and_analyze(key, loaded, specs, strength, device, top_count), False)
except Exception as error:
if not str(device).startswith("cuda") or not _is_cuda_oom(error):
raise
error.__traceback__ = None
_release_cuda_oom()
try:
return (*_apply_and_analyze(key, loaded, specs, strength, "cpu", top_count), True)
except Exception as cpu_error:
raise RuntimeError(f"LoRA-on-model analysis CPU fallback failed for {key!r}: {cpu_error}") from cpu_error
def _csv(headers, rows):
output = StringIO(newline="")
writer = csv.writer(output)
writer.writerow(headers)
writer.writerows(rows)
return output.getvalue()
def analyze_lora_on_model(model_a, model_b, params):
lora_path = folder_paths.get_full_path_or_raise("loras", model_a)
base_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
strength = params.get("strength", 1.0)
records, fallback_keys, top_heap = [], [], []
global_raw = RawStats()
layer_stats, block_stats, block_cwb = {}, {}, {}
try:
with MemoryEfficientSafeOpen(base_path, low_memory=True) as base, MemoryEfficientSafeOpen(lora_path, low_memory=True) as lora:
base_quant = DiffusionQuantization(base, base_path, "Base diffusion model")
units, patches, inventory, unmatched, passthrough = _plan_application(base, lora, base_path, params, base_quant)
device = params["process_device"]
top_limit = params["top_weight_differences"]
with closing(stream_work_units({0: base, 1: lora}, units, pin_memory=str(device).startswith("cuda"))) as stream:
for key, loaded in stream:
try:
if key in base_quant.quantized_keys:
loaded[(0, key)] = base_quant.decode(
key,
{source_key: loaded[(0, source_key)] for source_key in base_quant.required_keys(key)},
torch.float32,
)
raw, cwb, top, fallback = _analyze_unit(key, loaded, patches.get(key, []), strength, device, top_limit)
shape = base_quant.logical_shape(key)
records.append(TensorRecord(key, shape, str(base_quant.logical_dtype(key, torch.float32)), str(torch.float32), raw, cwb, fallback))
global_raw.add(raw)
layer_stats.setdefault(_layer_name(key), RawStats()).add(raw)
block_stats.setdefault(_block_name(key), RawStats()).add(raw)
block_cwb.setdefault(_block_name(key), CWBStats()).add(cwb)
if fallback:
fallback_keys.append(key)
for value, index, a, b in top:
entry = (value, key, index, a, b, shape)
if len(top_heap) < top_limit:
heapq.heappush(top_heap, entry)
elif top_limit and value > top_heap[0][0]:
heapq.heapreplace(top_heap, entry)
finally:
loaded.clear()
if params["force_clear_cache"]:
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
label = f"{model_b} + {model_a} (strength={strength})"
comparison = _render_comparison(model_b, label, "LoRA applied to diffusion model", records, global_raw, layer_stats, block_stats, inventory, top_heap, fallback_keys, base.metadata(), base.metadata())
comparison = (
"# LoRA-on-model analysis\n\nModel A input is the LoRA; Model B input is the base diffusion model. "
"The comparison below uses original base as A and patched base as B; delta = patched − original. "
"Unmodified base tensors are included unless filtered. No model is saved.\n\n"
+ comparison
+ "\n## Unmapped LoRA groups\n\n" + ("\n".join(f"- `{name}`" for name in unmatched) or "None.")
+ "\n\n## Unused LoRA tensors\n\n" + ("\n".join(f"- `{name}`" for name in passthrough) or "None.") + "\n"
)
cwb_report = _render_cwb(model_b, label, records, block_cwb, fallback_keys, "original versus LoRA-patched base tensor rows (index only)")
metric_names = ("mae", "mse", "rmse", "max", "relative_l2", "cosine", "pearson", "exact", "sign", "norm_ratio", "coverage")
metrics_csv = _csv(["key", "parameters", *metric_names], ([record.key, record.raw.total, *(_fmt(_metrics(record.raw)[name]) for name in metric_names)] for record in records))
cwb_names = ("pair_mean", "pair_min", "pair_max")
cwb_csv = _csv(["key", *cwb_names], ([record.key, *(_fmt(record.cwb.values()[name]) for name in cwb_names)] for record in records))
return comparison, cwb_report, metrics_csv, cwb_csv
finally:
cleanup_after_operation()
class LoRAOnModelAnalysis(io.ComfyNode):
@classmethod
def define_schema(cls):
inputs = _common_inputs("diffusion_models", alignment_control=False)
inputs[1] = io.Combo.Input("model_a", options=folder_paths.get_filename_list("loras"), tooltip="LoRA to apply transiently to Model B before comparing the original and patched base tensors.")
inputs[2] = io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"), tooltip="Base diffusion model. Its original weights are compared against a transient LoRA-patched copy, one tensor at a time.")
inputs.append(io.Float.Input("strength", default=1.0, min=-10.0, max=10.0, step=0.01, tooltip="LoRA application strength. Alpha scaling and DoRA reconstruction use ComfyUI's adapter; explicit set_weight patches retain their existing replacement semantics."))
return io.Schema(
node_id="LoRAOnModelAnalysis", display_name="Analyze LoRA Effect on Diffusion Model",
category="ModelUtils/Analysis",
description="Compare the base diffusion model with its LoRA-patched weights using bounded streaming. No model is written or inferred.",
inputs=inputs,
outputs=[io.String.Output(display_name="comparison_report"), io.String.Output(display_name="cwb_report"), io.String.Output(display_name="documentation"), io.String.Output(display_name="layerwise_metrics_csv"), io.String.Output(display_name="layerwise_cwb_csv")],
)
@classmethod
def execute(cls, **kwargs):
docs = load_documentation_from_file("lora_model_analysis.md")
if kwargs["execution_mode"] == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No comparison performed.", "No CWB analysis performed.", docs, "", "")
comparison, cwb, metrics_csv, cwb_csv = analyze_lora_on_model(kwargs["model_a"], kwargs["model_b"], kwargs)
return io.NodeOutput(comparison, cwb, docs, metrics_csv, cwb_csv)
+118
View File
@@ -0,0 +1,118 @@
import os
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
REVERSE_RENAME_DICT = {
"transformer_blocks.": "blocks.",
"norm1.linear_1": "adaln_modulation_self_attn.1",
"norm1.linear_2": "adaln_modulation_self_attn.2",
"norm2.linear_1": "adaln_modulation_cross_attn.1",
"norm2.linear_2": "adaln_modulation_cross_attn.2",
"norm3.linear_1": "adaln_modulation_mlp.1",
"norm3.linear_2": "adaln_modulation_mlp.2",
"attn1.to_q": "self_attn.q_proj",
"attn1.to_k": "self_attn.k_proj",
"attn1.to_v": "self_attn.v_proj",
"attn1.to_out.0": "self_attn.output_proj",
"attn2.to_q": "cross_attn.q_proj",
"attn2.to_k": "cross_attn.k_proj",
"attn2.to_v": "cross_attn.v_proj",
"attn2.to_out.0": "cross_attn.output_proj",
"ff.net.0.proj": "mlp.layer1",
"ff.net.2": "mlp.layer2",
"norm_out.linear_1": "final_layer.adaln_modulation.1",
"norm_out.linear_2": "final_layer.adaln_modulation.2",
"proj_out": "final_layer.linear",
"time_embed.t_embedder": "t_embedder.1",
"time_embed.norm": "t_embedding_norm",
"patch_embed.proj": "x_embedder.proj.1",
}
def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filename: str) -> str:
"""
Converts a Diffusers Anima LoRA back to its non-diffusers format.
"""
model_type = "loras"
model_path = folder_paths.get_full_path_or_raise(model_type, lora_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
# Stream tensors, rename on the fly, write immediately
with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
metadata = handler.metadata() or {}
with atomic_uel_writer(output_path, metadata) as writer:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
stream = handler.async_stream(
original_keys, batch_size=1, prefetch_batches=1, pin_memory=False
)
for batch in tqdm(stream, total=len(original_keys), desc="Renaming Anima LoRA keys", unit="keys"):
key, tensor = batch[0]
new_key = key
if key.startswith("text_conditioner."):
new_key = f"diffusion_model.llm_adapter.{key.removeprefix('text_conditioner.')}"
elif key.startswith("transformer."):
new_key = key.removeprefix("transformer.")
for old_part, new_part in REVERSE_RENAME_DICT.items():
new_key = new_key.replace(old_part, new_part)
new_key = f"diffusion_model.{new_key}"
writer.write_batch([(new_key, tensor.contiguous())])
handler.mark_processed(key)
pbar.update(1)
# Cleanup after operation
cleanup_after_operation()
return output_name
class AnimaLoraRename(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="AnimaLoraRename",
display_name="Anima LoRA Rename (Diffusers to Non-Diffusers)",
category="ModelUtils/LoRA",
description="Loads a Diffusers Anima LoRA, reverses its key names to the non-diffusers format, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
tooltip="Diffusers-format Anima LoRA whose tensor keys will be renamed.",
),
io.String.Input("output_filename", default="anima_lora_non_diffusers", tooltip="Output filename without extension, written under ComfyUI's LoRA directory."),
],
outputs=[
io.AnyType.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name: str, output_filename: str) -> io.NodeOutput:
path = _convert_diffusers_to_non_diffusers_anima_lora(lora_name, output_filename)
return io.NodeOutput(path)
+2989
View File
File diff suppressed because it is too large Load Diff
+799 -301
View File
File diff suppressed because it is too large Load Diff
+588 -45
View File
@@ -4,6 +4,8 @@ import torch.nn.functional as F
from enum import Enum
from unifiedefficientloader import transfer_to_gpu_pinned
class MissingTensorBehavior(Enum):
"""Controls behavior when a tensor key is missing from a model."""
ERROR = "error" # Raise exception (strict mode)
@@ -15,6 +17,92 @@ 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
# ################# #
@@ -23,6 +111,7 @@ class Operation:
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
@@ -40,20 +129,31 @@ class Operation:
raise NotImplementedError
def recurse(self, operation):
source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*source_tensors)
self._source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*self._source_tensors)
def merge(self):
return self.merge_func(self)
def clean(self):
if hasattr(self, '_source_tensors'):
for t in self._source_tensors:
del t
del self._source_tensors
for source in self.sources:
if hasattr(source, 'clean'):
source.clean()
class LoadTensor(Operation):
def __init__(self, key, model_name, handlers, device, dtype,
preloaded_tensors=None,
on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None):
super().__init__(key)
self.model_name = model_name
self.handlers = handlers
self.device = device
self.dtype = dtype
self.preloaded_tensors = preloaded_tensors or {}
self.on_missing = on_missing
self.fallback_shape = fallback_shape
self.fallback_dtype = fallback_dtype
@@ -61,9 +161,9 @@ class LoadTensor(Operation):
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():
source = (self.model_name, self.key)
if source not in self.preloaded_tensors:
if self.on_missing == MissingTensorBehavior.ERROR:
raise MissingTensorError(
f"Key '{self.key}' not found in model '{self.model_name}'"
@@ -77,8 +177,18 @@ class LoadTensor(Operation):
)
dtype = self.fallback_dtype if self.fallback_dtype else self.dtype
return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype)
return handler.get_tensor(self.key).to(device=self.device, dtype=self.dtype)
cpu_tensor = self.preloaded_tensors[source]
if self.device == 'cuda':
self._tensor = transfer_to_gpu_pinned(cpu_tensor, self.device, self.dtype)
else:
self._tensor = cpu_tensor.to(device=self.device, dtype=self.dtype)
del cpu_tensor
return self._tensor
def clean(self):
if hasattr(self, '_tensor'):
del self._tensor
class Multiply(Operation):
def __init__(self, key, alpha, *sources):
@@ -96,6 +206,7 @@ class Add(Operation):
return b
if b is None:
return a
a, b = resize_tensors(a, b, mode=self.alignment_mode)
return a + b
@@ -103,16 +214,34 @@ 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):
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
diff, weights, conv2d = a - b, {}, len(a.size()) == 4
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, 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
@@ -141,6 +270,10 @@ 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())
@@ -153,11 +286,17 @@ class ExtractOp(Operation):
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.flatten(), b_f.flatten(), 0).clamp(-1, 1)
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(a.dtype)
return (torch.lerp(a_f, b_f, self.alpha) * torch.lerp(d, 1 - d, self.beta)).to(dtype)
class Similarities(ExtractOp):
@@ -172,10 +311,61 @@ class PowerUpOp(Operation):
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.bernoulli(torch.full_like(delta, 1 - self.alpha), generator=rng)
return (m * delta) / (1 - self.alpha) if (1 - self.alpha) > 0 else torch.zeros_like(delta)
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):
@@ -185,6 +375,7 @@ class InterpolateDifference(Operation):
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
@@ -194,18 +385,78 @@ class InterpolateDifference(Operation):
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'],
"preloaded_tensors": kwargs.get('preloaded_tensors', {}),
"device": kwargs['device'],
"dtype": kwargs['dtype'],
"on_missing": mismatch_mode,
@@ -219,16 +470,20 @@ class CalcMode:
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.'
@@ -236,22 +491,25 @@ class WeightSum(CalcMode):
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)
return Add(key, c, d)
res = Add(key, c, d)
res.alignment_mode = alignment_mode
return res
class AddDifference(CalcMode):
@@ -261,20 +519,24 @@ class AddDifference(CalcMode):
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)
return Add(key, a, diffm)
res = Add(key, a, diffm)
res.alignment_mode = alignment_mode
return res
class TrainDifference(CalcMode):
@@ -284,19 +546,23 @@ class TrainDifference(CalcMode):
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)
return Add(key, a, diffm)
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.'
@@ -304,20 +570,23 @@ class InterpDifference(CalcMode):
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)
return InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
res = InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
class Extract(CalcMode):
@@ -327,19 +596,23 @@ class Extract(CalcMode):
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)
return Add(key, a, multiplied)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class AddDisimilarity(CalcMode):
@@ -349,19 +622,23 @@ class AddDisimilarity(CalcMode):
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)
return Add(key, a, multiplied)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class PowerUp(CalcMode):
@@ -371,18 +648,22 @@ class PowerUp(CalcMode):
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)
res = Multiply(key, kwargs['beta'], deltahat)
return Add(key, a, res)
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):
@@ -392,17 +673,279 @@ class SVDMode(CalcMode):
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)
return SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, 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
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode()]
THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
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': 'Model B pull strength (exponent 1/alpha - 1). 0.25 strongly suppresses B; 0.50 is linear; 1.0 is full B.',
'beta': 'Lower mean threshold. Only weights with normalized difference > beta blend B.',
'gamma': 'Upper mean threshold. Only weights with normalized difference < gamma blend B.',
'delta': 'Smoothness factor. 0.0 is stochastic Bernoulli choice; 1.0 is smooth linear blend.',
}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
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': 'Model B pull strength (exponent 1/alpha - 1). Lower keeps output close to Model A; higher pulls in Model B.',
'beta': 'Threshold width around the layer mean difference (+/- beta).',
'gamma': 'Smoothness factor. 0.0 is stochastic Bernoulli choice; 1.0 is smooth linear blend.',
}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
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': 'Model B blend fraction inside the threshold window. 0.0 is pure Model A; 0.25 is 25% B / 75% A; 1.0 is 100% B.',
'beta': 'Upper difference threshold cutoff.',
'gamma': 'Lower difference threshold cutoff.',
}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
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()]
+9 -92
View File
@@ -1,90 +1,10 @@
import json
import struct
import torch
"""
Utilities for model merging operations.
"""
from typing import Dict, List
class MemoryEfficientSafeOpen:
def __init__(self, filename, device='cpu'):
self.filename = filename
self.device = device
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.file.close()
def keys(self):
return [k for k in self.header.keys() if k != "__metadata__"]
def get_tensor(self, key):
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
return self._deserialize_tensor(tensor_bytes, metadata)
def _read_header(self):
with open(self.filename, "rb") as f:
header_size = struct.unpack("<Q", f.read(8))[0]
header_json = f.read(header_size).decode("utf-8")
return json.loads(header_json), header_size
def _deserialize_tensor(self, tensor_bytes, metadata):
dtype_str = metadata["dtype"]
shape = metadata["shape"]
dtype = self._get_torch_dtype(dtype_str)
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
return self._convert_float8(byte_tensor, dtype_str, shape)
return byte_tensor.view(dtype).reshape(shape)
@staticmethod
def _get_torch_dtype(dtype_str):
dtype_map = {
"F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
"I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
"U8": torch.uint8, "BOOL": torch.bool,
}
if hasattr(torch, "float8_e5m2"):
dtype_map["F8_E5M2"] = torch.float8_e5m2
if hasattr(torch, "float8_e4m3fn"):
dtype_map["F8_E4M3"] = torch.float8_e4m3fn
dtype = dtype_map.get(dtype_str)
if dtype is None:
raise ValueError(f"Unsupported dtype: {dtype_str}")
return dtype
@staticmethod
def _convert_float8(byte_tensor, dtype_str, shape):
if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
return byte_tensor.view(torch.float8_e5m2).reshape(shape)
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
else:
raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")
def get_union_keys(*handlers):
"""Returns the union of all tensor keys across multiple model handlers.
Useful when you want to iterate over all unique keys present in any model.
"""
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:
@@ -92,15 +12,12 @@ def get_union_keys(*handlers):
return sorted(all_keys)
def get_intersection_keys(*handlers):
"""Returns only keys present in ALL model handlers.
Useful when you want to iterate only over keys that exist in every model.
"""
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)
return sorted(all_keys)
+19 -11
View File
@@ -1,15 +1,15 @@
import os
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
from unifiedefficientloader import MemoryEfficientSafeOpen
from .utils import convert_pt_to_safetensors
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)
model_weights = {}
metadata = {}
layer_shapes = ""
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
@@ -19,13 +19,16 @@ def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]:
return f"Conversion failed: {error_message}", ""
model_path = temp_safe_path
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
with MemoryEfficientSafeOpen(model_path, low_memory=True) as handler:
metadata_str = str(handler.metadata() or {})
keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(keys))
lines = []
for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
lines.append(f"{layer_name}, {list(handler.get_shape(layer_name))}\n")
pbar.update(1)
layer_shapes = "".join(lines)
return metadata_str, layer_shapes
@@ -41,6 +44,7 @@ class ModelMetaKeys(io.ComfyNode):
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Diffusion model whose safetensors metadata and tensor keys will be listed.",
),
],
outputs=[
@@ -67,6 +71,7 @@ class TextEncoderMetaKeys(io.ComfyNode):
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
tooltip="Text encoder whose safetensors metadata and tensor keys will be listed.",
),
],
outputs=[
@@ -93,6 +98,7 @@ class LoRAMetaKeys(io.ComfyNode):
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
tooltip="LoRA whose safetensors metadata and tensor keys will be listed.",
),
],
outputs=[
@@ -119,6 +125,7 @@ class CheckpointMetaKeys(io.ComfyNode):
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="Checkpoint whose safetensors metadata and tensor keys will be listed.",
),
],
outputs=[
@@ -145,6 +152,7 @@ class EmbeddingMetaKeys(io.ComfyNode):
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
tooltip="Embedding whose safetensors metadata and tensor keys will be listed.",
),
],
outputs=[
+687
View File
@@ -0,0 +1,687 @@
from __future__ import annotations
import fnmatch
import json
import math
import os
import re
import comfy.utils
from comfy_api.latest import io
import folder_paths
import torch
from unifiedefficientloader import MemoryEfficientSafeOpen
from .artifact_paths import canonical_model_artifact_path
from .device_utils import cleanup_after_operation
from .quantization_guard import inspect_low_bit_input
from .uel_io import atomic_uel_writer
def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
"""Compile newline or whitespace separated patterns."""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for pattern in pattern_string.split():
if not pattern:
continue
if glob_mode:
patterns.append(pattern)
else:
try:
patterns.append(re.compile(pattern))
except re.error as e:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}") from e
return patterns
def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
if not patterns:
return False
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
return any(p.search(key) is not None for p in patterns)
def time_embedding_curve(
w1: torch.Tensor,
b1: torch.Tensor,
w2: torch.Tensor,
b2: torch.Tensor,
grid_size: int,
device: torch.device,
) -> torch.Tensor:
"""Sample silu(time_embedder(t)) on the curve grid."""
t = torch.linspace(0.0, 1.0, grid_size, dtype=torch.float32, device=device)
half = w1.shape[1] // 2
freqs = torch.exp(
-math.log(10000.0)
* torch.arange(half, dtype=torch.float32, device=device)
/ half
)
args = t[:, None] * freqs[None]
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
w1_d, b1_d, w2_d, b2_d = (x.to(device, torch.float32) for x in (w1, b1, w2, b2))
hidden = torch.nn.functional.silu(emb @ w1_d.T + b1_d)
return torch.nn.functional.silu(hidden @ w2_d.T + b2_d)
def _fold_curve_basis(
w1: torch.Tensor,
b1: torch.Tensor,
w2: torch.Tensor,
b2: torch.Tensor,
grid_size: int,
rank: int,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, float]:
"""Calculate SVD curve basis and adaln_t_table."""
compute_device = device
try:
curve = time_embedding_curve(w1, b1, w2, b2, grid_size, compute_device)
u, s, vh = torch.linalg.svd(curve, full_matrices=False)
except (torch.cuda.OutOfMemoryError, RuntimeError):
cleanup_after_operation()
compute_device = torch.device("cpu")
curve = time_embedding_curve(w1, b1, w2, b2, grid_size, compute_device)
u, s, vh = torch.linalg.svd(curve, full_matrices=False)
target_rank = min(rank, vh.shape[0])
table = u[:, :target_rank] * s[:target_rank]
basis = vh[:target_rank].T
error = ((table @ basis.T - curve).norm() / curve.norm()).item()
return table.float().cpu().contiguous(), basis.float().cpu().contiguous(), error
def _detect_time_embedder_keys(all_keys: list[str]) -> tuple[str, dict[str, str]]:
for prefix in ("", "diffusion_model.", "model.diffusion_model."):
candidates = {
"w1": f"{prefix}time_embedder.proj_in.weight",
"b1": f"{prefix}time_embedder.proj_in.bias",
"w2": f"{prefix}time_embedder.proj_out.weight",
"b2": f"{prefix}time_embedder.proj_out.bias",
}
if all(k in all_keys for k in candidates.values()):
return prefix, candidates
raise KeyError(
"Model is missing required time_embedder projection tensors (proj_in/proj_out)."
)
def _is_adaln(key: str) -> bool:
return ".adaln_proj.linear." in key or key.startswith("adaln_proj.linear.")
def fold_minimax_h3_diffusion_model(
input_path: str,
output_path: str,
*,
exclude_patterns: str = "",
include_mode: bool = False,
discard_patterns: str = "",
glob_patterns: bool = False,
curve_rank: int = 8,
curve_grid: int = 1025,
process_device: str = "cuda",
) -> str:
"""Stream and fold MiniMax H3 AdaLN linears in a diffusion model."""
if os.path.normcase(os.path.abspath(input_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the source model path.")
device = torch.device(
process_device if torch.cuda.is_available() and process_device == "cuda" else "cpu"
)
compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
handler = MemoryEfficientSafeOpen(input_path, low_memory=True)
try:
inspect_low_bit_input(handler, f"Model ({input_path})", "MiniMax H3 Fold AdaLN")
all_keys = list(handler.keys())
prefix, time_keys = _detect_time_embedder_keys(all_keys)
# Step 1: Load time embedder and calculate basis
stream = handler.async_stream(
list(time_keys.values()),
batch_size=1,
prefetch_batches=1,
pin_memory=False,
)
time_tensors = {}
try:
for batch in stream:
for k, t in batch:
time_tensors[k] = t
table, basis, error = _fold_curve_basis(
time_tensors[time_keys["w1"]],
time_tensors[time_keys["b1"]],
time_tensors[time_keys["w2"]],
time_tensors[time_keys["b2"]],
curve_grid,
curve_rank,
device,
)
finally:
time_tensors.clear()
for k in time_keys.values():
handler.mark_processed(k)
stream.close()
# Determine which AdaLN layers are folded vs kept full width
any_adaln_kept_unfolded = False
for k in all_keys:
if _is_adaln(k) and not _matches_any_pattern(k, compiled_discards, glob_patterns):
matched = _matches_any_pattern(k, compiled_excludes, glob_patterns)
should_fold = matched if include_mode else not matched
if not should_fold:
any_adaln_kept_unfolded = True
break
metadata = dict(handler.metadata() or {})
if "config" in metadata:
try:
cfg = json.loads(metadata["config"])
if isinstance(cfg, dict):
tf = cfg.setdefault("transformer", {})
tf["adaln_curve_grid"] = table.shape[0]
tf["time_embed_dim"] = table.shape[1]
metadata["config"] = json.dumps(cfg, separators=(",", ":"))
except Exception:
pass
metadata["modelutils_adaln_fold"] = f"rank{table.shape[1]}_curve_grid_{table.shape[0]}"
table_key = f"{prefix}adaln_t_table"
time_embedder_prefix = f"{prefix}time_embedder."
counts = {
"folded": 0,
"preserved": 0,
"discarded": 0,
"time_embedder_dropped": 0,
}
basis_device = basis.to(device)
with atomic_uel_writer(output_path, metadata) as writer:
# Emit adaln_t_table
writer.write_batch([(table_key, table)])
# Filter remaining keys
remaining_keys = [k for k in all_keys if k != table_key]
progress = comfy.utils.ProgressBar(len(remaining_keys))
for key in remaining_keys:
if _matches_any_pattern(key, compiled_discards, glob_patterns):
counts["discarded"] += 1
handler.mark_processed(key)
progress.update(1)
continue
if key.startswith(time_embedder_prefix) and not any_adaln_kept_unfolded:
counts["time_embedder_dropped"] += 1
handler.mark_processed(key)
progress.update(1)
continue
# Stream this tensor
tstream = handler.async_stream(
[key], batch_size=1, prefetch_batches=1, pin_memory=False
)
try:
batch = next(tstream)
tensor = batch[0][1]
if _is_adaln(key):
matched = _matches_any_pattern(key, compiled_excludes, glob_patterns)
should_fold = matched if include_mode else not matched
if should_fold:
if key.endswith(".weight"):
if tensor.ndim != 2 or tensor.shape[1] != basis.shape[0]:
raise ValueError(
f"{key} has shape {tuple(tensor.shape)}, expected [*, {basis.shape[0]}]"
)
folded_w = (
tensor.to(device, torch.float32) @ basis_device
).to(torch.float16).cpu().contiguous()
writer.write_batch([(key, folded_w)])
else:
writer.write_batch([(key, tensor.to(torch.float16).contiguous())])
counts["folded"] += 1
else:
writer.write_batch([(key, tensor.contiguous())])
counts["preserved"] += 1
else:
writer.write_batch([(key, tensor.contiguous())])
counts["preserved"] += 1
finally:
handler.mark_processed(key)
tstream.close()
progress.update(1)
finally:
handler.__exit__(None, None, None)
cleanup_after_operation()
report_lines = [
"MiniMax H3 AdaLN Fold complete",
f"Curve rank: {curve_rank}, grid points: {curve_grid}",
f"Reconstruction relative error: {error:.4e}",
f"AdaLN layers folded: {counts['folded']}",
f"Tensors preserved: {counts['preserved']}",
f"Tensors discarded: {counts['discarded']}",
f"Time embedder dropped: {counts['time_embedder_dropped']}",
f"Output size: {os.path.getsize(output_path)} bytes",
f"Output: {output_path}",
]
return "\n".join(report_lines)
class MiniMaxH3FoldAdaLN(io.ComfyNode):
"""Folds MiniMax H3 AdaLN conditioning linears onto a low-rank curve basis.
Technical details:
Full-width AdaLN linear weights [*, 2688] are projected onto a rank-8 SVD basis
of the silu(time_embedder) curve sampled across a 1025-point grid of t in [0, 1]:
X = silu(t_emb(t)) [1025, 2688] = U S V^T
adaln_t_table = U_8 S_8 (FP32), W' = W V_8 (FP16)
The time embedder is dropped and replaced with adaln_t_table [1025, 8], reducing
model size substantially with negligible reconstruction error (< 1e-4).
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MiniMaxH3FoldAdaLN",
display_name="MiniMax H3 Fold AdaLN",
category="ModelUtils/DiffusionModels",
description="Compresses large MiniMax H3 models by folding time-conditioning layers into a small lookup table, matching official ComfyUI checkpoint formats.",
inputs=[
io.Combo.Input(
"model_name",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Select the original/uncompressed MiniMax H3 model file to fold.",
),
io.String.Input(
"output_filename",
default="minimax_h3_folded",
tooltip="Name for the new folded model file, saved in your diffusion_models folder.",
),
io.String.Input(
"discard_patterns",
default="",
multiline=True,
tooltip="Layers to completely remove from the final file (one name or pattern per line). Useful to strip parts you do not need.",
),
io.Boolean.Input(
"glob_patterns",
default=False,
tooltip="Enable simple wildcard matching (* to match anything). If unchecked, uses regular expressions.",
),
io.Int.Input(
"curve_rank",
default=8,
min=1,
max=64,
step=1,
tooltip="Compression size for the time schedule table. Leave at 8 to match standard ComfyUI MiniMax H3 models.",
),
io.Int.Input(
"curve_grid",
default=1025,
min=65,
max=4097,
step=64,
tooltip="Sampling resolution for the time schedule curve. Leave at 1025 to match standard ComfyUI MiniMax H3 models.",
),
io.Combo.Input(
"process_device",
options=["cuda", "cpu"],
default="cuda",
tooltip="Hardware device used to perform the math. Use cuda (GPU) for speed, or cpu if you run low on VRAM.",
),
io.String.Input(
"exclude_patterns",
default="",
multiline=True,
tooltip="Specific layers to skip folding, keeping them at their original full size (one name or pattern per line).",
),
io.Boolean.Input(
"include_mode",
default=False,
tooltip="Invert the exclude filter: only fold layers matching the pattern, and keep everything else untouched.",
),
],
outputs=[
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="report"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_name: str,
output_filename: str,
discard_patterns: str = "",
glob_patterns: bool = False,
curve_rank: int = 8,
curve_grid: int = 1025,
process_device: str = "cuda",
exclude_patterns: str = "",
include_mode: bool = False,
) -> io.NodeOutput:
input_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_name
)
output_path, output_name = canonical_model_artifact_path(
"diffusion_models", output_filename
)
report = fold_minimax_h3_diffusion_model(
input_path,
output_path,
exclude_patterns=exclude_patterns,
include_mode=include_mode,
discard_patterns=discard_patterns,
glob_patterns=glob_patterns,
curve_rank=curve_rank,
curve_grid=curve_grid,
process_device=process_device,
)
return io.NodeOutput(output_name, report)
def unfold_minimax_h3_diffusion_model(
model_a_path: str,
model_b_path: str,
output_path: str,
*,
exclude_patterns: str = "",
include_mode: bool = False,
discard_patterns: str = "",
glob_patterns: bool = False,
curve_rank: int = 8,
curve_grid: int = 1025,
process_device: str = "cuda",
) -> str:
"""Stream and reconstruct unfolded MiniMax H3 AdaLN linears using a reference base model."""
if os.path.normcase(os.path.abspath(model_a_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the folded model (Model A) path.")
if os.path.normcase(os.path.abspath(model_b_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the base model (Model B) path.")
device = torch.device(
process_device if torch.cuda.is_available() and process_device == "cuda" else "cpu"
)
compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns)
compiled_discards = _compile_patterns(discard_patterns, glob_patterns)
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=True)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=True)
try:
inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "MiniMax H3 Unfold AdaLN")
inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "MiniMax H3 Unfold AdaLN")
all_keys_a = list(handler_a.keys())
all_keys_b = list(handler_b.keys())
prefix_b, time_keys_b = _detect_time_embedder_keys(all_keys_b)
prefix_a = "diffusion_model." if any(k.startswith("diffusion_model.") for k in all_keys_a) else ""
# Step 1: Compute basis and table from base model (Model B)
tstream_b = handler_b.async_stream(
list(time_keys_b.values()),
batch_size=1,
prefetch_batches=1,
pin_memory=False,
)
time_tensors_b = {}
try:
for batch in tstream_b:
for k, t in batch:
time_tensors_b[k] = t
_, basis, error = _fold_curve_basis(
time_tensors_b[time_keys_b["w1"]],
time_tensors_b[time_keys_b["b1"]],
time_tensors_b[time_keys_b["w2"]],
time_tensors_b[time_keys_b["b2"]],
curve_grid,
curve_rank,
device,
)
finally:
for k in time_keys_b.values():
handler_b.mark_processed(k)
tstream_b.close()
basis_device = basis.to(device)
metadata = dict(handler_a.metadata() or {})
if "config" in metadata:
try:
cfg = json.loads(metadata["config"])
if isinstance(cfg, dict):
tf = cfg.setdefault("transformer", {})
tf.pop("adaln_curve_grid", None)
tf["time_embed_dim"] = basis.shape[0]
metadata["config"] = json.dumps(cfg, separators=(",", ":"))
except Exception:
pass
metadata["modelutils_adaln_unfold"] = f"reconstructed_rank{basis.shape[1]}_with_base"
table_key_a = f"{prefix_a}adaln_t_table"
counts = {
"unfolded": 0,
"preserved": 0,
"discarded": 0,
"time_embedder_restored": 0,
}
with atomic_uel_writer(output_path, metadata) as writer:
# Emit the 4 time_embedder tensors restored from Model B
for _, base_k in time_keys_b.items():
target_k = f"{prefix_a}{base_k.removeprefix(prefix_b)}"
if not _matches_any_pattern(target_k, compiled_discards, glob_patterns):
writer.write_batch([(target_k, time_tensors_b[base_k])])
counts["time_embedder_restored"] += 1
time_tensors_b.clear()
remaining_keys_a = [k for k in all_keys_a if k != table_key_a]
progress = comfy.utils.ProgressBar(len(remaining_keys_a))
for key in remaining_keys_a:
if _matches_any_pattern(key, compiled_discards, glob_patterns):
counts["discarded"] += 1
handler_a.mark_processed(key)
progress.update(1)
continue
astream = handler_a.async_stream(
[key], batch_size=1, prefetch_batches=1, pin_memory=False
)
try:
tensor_a = next(astream)[0][1]
if _is_adaln(key):
matched = _matches_any_pattern(key, compiled_excludes, glob_patterns)
should_unfold = matched if include_mode else not matched
if should_unfold and key.endswith(".weight"):
base_key = f"{prefix_b}{key.removeprefix(prefix_a)}"
if base_key not in all_keys_b:
raise KeyError(
f"Matching base layer '{base_key}' not found in Model B."
)
bstream = handler_b.async_stream(
[base_key], batch_size=1, prefetch_batches=1, pin_memory=False
)
try:
base_w = next(bstream)[0][1]
if base_w.ndim != 2 or base_w.shape[1] != basis.shape[0]:
raise ValueError(
f"{base_key} in Model B has shape {tuple(base_w.shape)}, expected [*, {basis.shape[0]}]"
)
bw_f32 = base_w.to(device, torch.float32)
w_folded = tensor_a.to(device, torch.float32)
w_diff = w_folded - bw_f32 @ basis_device
w_unfolded = (bw_f32 + w_diff @ basis_device.T).to(base_w.dtype).cpu().contiguous()
writer.write_batch([(key, w_unfolded)])
counts["unfolded"] += 1
finally:
handler_b.mark_processed(base_key)
bstream.close()
else:
writer.write_batch([(key, tensor_a.contiguous())])
counts["preserved"] += 1
else:
writer.write_batch([(key, tensor_a.contiguous())])
counts["preserved"] += 1
finally:
handler_a.mark_processed(key)
astream.close()
progress.update(1)
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
report_lines = [
"MiniMax H3 AdaLN Unfold complete",
f"Curve rank: {curve_rank}, grid points: {curve_grid}",
f"Base basis reconstruction error: {error:.4e}",
f"AdaLN layers unfolded: {counts['unfolded']}",
f"Time embedder tensors restored: {counts['time_embedder_restored']}",
f"Tensors preserved: {counts['preserved']}",
f"Tensors discarded: {counts['discarded']}",
f"Output size: {os.path.getsize(output_path)} bytes",
f"Output: {output_path}",
]
return "\n".join(report_lines)
class MiniMaxH3UnfoldAdaLN(io.ComfyNode):
"""Reconstruct full-width AdaLN linears [*, 2688] from a folded model and reference base model."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MiniMaxH3UnfoldAdaLN",
display_name="MiniMax H3 Unfold AdaLN",
category="ModelUtils/DiffusionModels",
description="Reconstructs full-width MiniMax H3 AdaLN linears [*, 2688] from a folded model (Model A) using a reference non-folded base model (Model B) to supply the time embedder and orthogonal dimensions.",
inputs=[
io.Combo.Input(
"model_a",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Select the folded MiniMax H3 model (Model A) to unfold.",
),
io.Combo.Input(
"model_b",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Select the reference original/unfolded MiniMax H3 base model (Model B) to supply the time embedder and missing dimensions.",
),
io.String.Input(
"output_filename",
default="minimax_h3_unfolded",
tooltip="Name for the reconstructed unfolded model file, saved in your diffusion_models folder.",
),
io.String.Input(
"discard_patterns",
default="",
multiline=True,
tooltip="Layers to completely remove from the final file (one name or pattern per line). Useful to strip parts you do not need.",
),
io.Boolean.Input(
"glob_patterns",
default=False,
tooltip="Enable simple wildcard matching (* to match anything). If unchecked, uses regular expressions.",
),
io.Int.Input(
"curve_rank",
default=8,
min=1,
max=64,
step=1,
tooltip="Basis rank used when the model was folded (default 8, matching standard ComfyUI H3 checkpoints).",
),
io.Int.Input(
"curve_grid",
default=1025,
min=65,
max=4097,
step=64,
tooltip="Sampling resolution used when the model was folded (default 1025, matching standard ComfyUI H3 checkpoints).",
),
io.Combo.Input(
"process_device",
options=["cuda", "cpu"],
default="cuda",
tooltip="Hardware device used to perform the math. Use cuda (GPU) for speed, or cpu if you run low on VRAM.",
),
io.String.Input(
"exclude_patterns",
default="",
multiline=True,
tooltip="Specific layers to skip unfolding, keeping them at their folded size (one name or pattern per line).",
),
io.Boolean.Input(
"include_mode",
default=False,
tooltip="Invert the exclude filter: only unfold layers matching the pattern, and keep everything else untouched.",
),
],
outputs=[
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="report"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_a: str,
model_b: str,
output_filename: str,
discard_patterns: str = "",
glob_patterns: bool = False,
curve_rank: int = 8,
curve_grid: int = 1025,
process_device: str = "cuda",
exclude_patterns: str = "",
include_mode: bool = False,
) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_a
)
model_b_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_b
)
output_path, output_name = canonical_model_artifact_path(
"diffusion_models", output_filename
)
report = unfold_minimax_h3_diffusion_model(
model_a_path,
model_b_path,
output_path,
exclude_patterns=exclude_patterns,
include_mode=include_mode,
discard_patterns=discard_patterns,
glob_patterns=glob_patterns,
curve_rank=curve_rank,
curve_grid=curve_grid,
process_device=process_device,
)
return io.NodeOutput(output_name, report)
+335
View File
@@ -0,0 +1,335 @@
from __future__ import annotations
from dataclasses import dataclass
import os
import re
import comfy.utils
from comfy_api.latest import io
import folder_paths
import torch
from unifiedefficientloader import MemoryEfficientSafeOpen
from .artifact_paths import canonical_model_artifact_path
from .quantization_guard import inspect_low_bit_input
from .uel_io import atomic_uel_writer
_FACTOR_RE = re.compile(
r"^(?P<module>.+)\.(?:lora_(?P<factor_ab>[AB])|lora[._](?P<factor_word>down|up))(?:\.default)?(?:\.weight)?$"
)
_BLOCK_RE = re.compile(r"^transformer_blocks\.(?P<index>\d+)\.(?P<tail>.+)$")
_REFINER_RE = re.compile(r"^token_refiner\.refiner_blocks\.(?P<index>\d+)\.(?P<tail>.+)$")
_QKV_RE = re.compile(r"^attn\.to_(?P<part>[qkv])$")
_DIRECT_TAILS = {
"attn.to_out.0": "attn.out_proj",
"ff.net.0.proj": "mlp.fc1",
"ff.net.2": "mlp.fc2",
}
_EXTRA_PREFIX_MAP = {
"context_embedder": "condition_proj",
"proj_in": "video_patch_proj",
"proj_out": "final_layer.video_out",
"audio_proj_in": "audio_patch_proj",
"audio_proj_out": "final_layer.audio_out",
"norm_out.norm": "final_layer.norm",
"norm_out.linear": "final_layer.adaln_proj.linear",
"time_embedder.linear_1": "time_embedder.proj_in",
"time_embedder.linear_2": "time_embedder.proj_out",
}
@dataclass(frozen=True)
class _ConversionUnit:
target: str
sources: tuple[str, ...]
qkv: bool
is_swiglu: bool = False
def _target_module(module: str) -> tuple[str, str | None, bool]:
"""Map Diffusers module name to ComfyUI native target module.
Returns:
(target_module, qkv_part, is_swiglu_fc1)
"""
clean_mod = module.removeprefix("base_model.model.").removeprefix("transformer.")
match = _BLOCK_RE.fullmatch(clean_mod)
if match:
prefix = f"diffusion_model.blocks.{match.group('index')}"
tail = match.group("tail")
else:
match = _REFINER_RE.fullmatch(clean_mod)
if match:
prefix = f"diffusion_model.token_refiner.blocks.{match.group('index')}"
tail = match.group("tail")
else:
for old_pfx, new_pfx in _EXTRA_PREFIX_MAP.items():
if clean_mod.startswith(old_pfx):
target = f"diffusion_model.{new_pfx}{clean_mod[len(old_pfx):]}"
return target, None, False
raise ValueError(f"Unsupported MiniMax H3 Diffusers module: {module}")
qkv = _QKV_RE.fullmatch(tail)
if qkv:
return f"{prefix}.attn.qkv_proj", qkv.group("part"), False
if tail in _DIRECT_TAILS:
target = f"{prefix}.{_DIRECT_TAILS[tail]}"
is_swiglu = (tail == "ff.net.0.proj")
return target, None, is_swiglu
raise ValueError(f"Unsupported MiniMax H3 Diffusers module tail: {module}")
def _build_plan(keys: list[str]) -> list[_ConversionUnit]:
"""Build conversion work units from input keys."""
factors: dict[str, dict[str, str]] = {}
for key in keys:
match = _FACTOR_RE.fullmatch(key)
if not match:
raise ValueError(
f"Input is not a compatible MiniMax H3 Diffusers PEFT LoRA; unsupported key: {key}"
)
mod = match.group("module")
factor = match.group("factor_ab") or (
"A" if match.group("factor_word") == "down" else "B"
)
pair = factors.setdefault(mod, {})
if factor in pair:
raise ValueError(f"Duplicate LoRA factor for: {mod}")
pair[factor] = key
missing = [mod for mod, pair in factors.items() if set(pair.keys()) != {"A", "B"}]
if missing:
raise ValueError(f"Incomplete LoRA A/B pair for: {missing[0]}")
direct_units: list[_ConversionUnit] = []
qkv_groups: dict[str, dict[str, tuple[str, str]]] = {}
for mod, pair in factors.items():
target, qkv_part, is_swiglu = _target_module(mod)
source_pair = (pair["A"], pair["B"])
if qkv_part is None:
direct_units.append(
_ConversionUnit(
target=target,
sources=source_pair,
qkv=False,
is_swiglu=is_swiglu,
)
)
else:
qkv_groups.setdefault(target, {})[qkv_part] = source_pair
qkv_units: list[_ConversionUnit] = []
for target, parts in qkv_groups.items():
if set(parts.keys()) != {"q", "k", "v"}:
raise ValueError(f"Incomplete Q/K/V LoRA group for: {target}")
sources = tuple(k for part in ("q", "k", "v") for k in parts[part])
qkv_units.append(_ConversionUnit(target=target, sources=sources, qkv=True))
units = sorted(direct_units + qkv_units, key=lambda u: u.target)
targets = [u.target for u in units]
if len(targets) != len(set(targets)):
raise ValueError("Conversion would produce duplicate ComfyUI module keys")
return units
def _validate_pair(down: torch.Tensor, up: torch.Tensor, label: str) -> None:
if down.ndim != 2 or up.ndim != 2:
raise ValueError(f"LoRA factors for '{label}' must be 2D matrices")
if down.shape[0] != up.shape[1]:
raise ValueError(
f"LoRA factor rank mismatch for '{label}': down={list(down.shape)}, up={list(up.shape)}"
)
def _convert_qkv(
tensors: dict[str, torch.Tensor],
sources: tuple[str, ...],
) -> tuple[torch.Tensor, torch.Tensor]:
"""Fuse separate Q, K, V LoRA factors into unified block-diagonal qkv_proj factors."""
pairs = [
(tensors[sources[i]], tensors[sources[i + 1]])
for i in range(0, 6, 2)
]
for part, (down, up) in zip(("Q", "K", "V"), pairs):
_validate_pair(down, up, part)
in_dims = {down.shape[1] for down, _ in pairs}
if len(in_dims) != 1:
raise ValueError("Q/K/V LoRA input dimensions do not match")
dtypes = {t.dtype for pair in pairs for t in pair}
if len(dtypes) != 1:
raise ValueError("Q/K/V LoRA dtypes do not match")
source_dtype = pairs[0][0].dtype
ranks = [down.shape[0] for down, _ in pairs]
output_dims = [up.shape[0] for _, up in pairs]
total_rank = sum(ranks)
total_out = sum(output_dims)
# Concatenate down matrices vertically: [total_rank, in_dim]
exact_down = torch.cat([down for down, _ in pairs], dim=0).contiguous()
# Place up matrices block-diagonally: [total_out, total_rank]
exact_up = torch.zeros(
(total_out, total_rank),
dtype=source_dtype,
device=pairs[0][1].device,
)
row = 0
col = 0
for (_, part_up), out_dim, part_rank in zip(pairs, output_dims, ranks):
exact_up[row : row + out_dim, col : col + part_rank].copy_(part_up)
row += out_dim
col += part_rank
return exact_down, exact_up.contiguous()
def _convert_swiglu(
down: torch.Tensor, up: torch.Tensor, label: str
) -> tuple[torch.Tensor, torch.Tensor]:
"""Swap SwiGLU output halves from Diffusers [value; gate] to ComfyUI [gate; value]."""
_validate_pair(down, up, label)
if up.shape[0] % 2 != 0:
raise ValueError(
f"SwiGLU LoRA B factor for '{label}' has odd output dimension {up.shape[0]}"
)
val_b, gate_b = up.chunk(2, dim=0)
swapped_up = torch.cat([gate_b, val_b], dim=0).contiguous()
return down.contiguous(), swapped_up
def convert_minimax_h3_diffusers_lora(
input_path: str,
output_path: str,
) -> str:
"""Stream and convert Diffusers MiniMax H3 LoRA into ComfyUI native format."""
if os.path.normcase(os.path.abspath(input_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the source LoRA path")
handler = MemoryEfficientSafeOpen(input_path, low_memory=True)
try:
inspect_low_bit_input(handler, f"LoRA ({input_path})", "MiniMax H3 LoRA Convert")
keys = list(handler.keys())
source_count = len(keys)
units = _build_plan(keys)
metadata = {str(k): str(v) for k, v in (handler.metadata() or {}).items()}
metadata.update(
{
"modelutils_conversion": "diffusers_minimax_h3_to_comfyui",
"modelutils_qkv_fusion": "lossless_block_diagonal",
"modelutils_mlp_swiglu": "swapped_gate_value",
}
)
progress = comfy.utils.ProgressBar(len(units))
output_count = 0
fused_qkv_count = 0
with atomic_uel_writer(output_path, metadata) as writer:
for unit in units:
stream = handler.async_stream(
list(unit.sources),
batch_size=1,
prefetch_batches=1,
pin_memory=False,
)
tensors = {}
try:
for batch in stream:
for k, t in batch:
tensors[k] = t
if unit.qkv:
down, up = _convert_qkv(tensors, unit.sources)
fused_qkv_count += 1
elif unit.is_swiglu:
down_key, up_key = unit.sources
down, up = _convert_swiglu(
tensors[down_key], tensors[up_key], unit.target
)
else:
down_key, up_key = unit.sources
down = tensors[down_key].contiguous()
up = tensors[up_key].contiguous()
_validate_pair(down, up, unit.target)
writer.write_batch(
[
(f"{unit.target}.lora_A.weight", down),
(f"{unit.target}.lora_B.weight", up),
]
)
output_count += 2
finally:
tensors.clear()
for key in unit.sources:
handler.mark_processed(key)
stream.close()
progress.update(1)
finally:
handler.__exit__(None, None, None)
lines = [
"MiniMax H3 Diffusers LoRA conversion complete",
f"Source tensors: {source_count}",
f"Output tensors: {output_count}",
f"Fused QKV modules: {fused_qkv_count}",
f"Output size: {os.path.getsize(output_path)} bytes",
f"Output: {output_path}",
]
return "\n".join(lines)
class MiniMaxH3DiffusersLoRAConvert(io.ComfyNode):
"""Convert Diffusers MiniMax H3 LoRAs to ComfyUI's native layout."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MiniMaxH3DiffusersLoRAConvert",
display_name="MiniMax H3 Diffusers LoRA Convert",
category="ModelUtils/LoRA",
description="Converts Diffusers/PEFT MiniMax H3 LoRAs to ComfyUI's native fused layout (lossless block-diagonal QKV and SwiGLU-aligned MLP).",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
tooltip="Diffusers-format MiniMax H3 LoRA to convert.",
),
io.String.Input(
"output_filename",
default="minimax_h3_comfyui",
tooltip="Output filename without extension, written under ComfyUI's LoRA directory.",
),
],
outputs=[
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="conversion_report"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
lora_name: str,
output_filename: str,
) -> io.NodeOutput:
input_path = folder_paths.get_full_path_or_raise("loras", lora_name)
output_path, output_name = canonical_model_artifact_path(
"loras", output_filename
)
report = convert_minimax_h3_diffusers_lora(input_path, output_path)
return io.NodeOutput(output_name, report)
File diff suppressed because it is too large Load Diff
+241
View File
@@ -0,0 +1,241 @@
import glob
import os
import numpy as np
from PIL import Image, ImageOps
import torch
import folder_paths
from comfy_api.latest import io
def get_potential_preview_files(base_path: str) -> list[str]:
exts = ["png", "jpg", "jpeg", "webp", "mp4"]
files = []
for ext in exts:
files.append(f"{base_path}.preview.{ext}")
files.append(f"{base_path}.{ext}")
return files
def load_image_tensor(path: str):
if not os.path.exists(path):
return None
try:
img = Image.open(path)
img = ImageOps.exif_transpose(img)
image = img.convert("RGB")
image_np = np.array(image).astype(np.float32) / 255.0
return torch.from_numpy(image_np)[None,]
except Exception as e:
print(f"Error loading image {path}: {e}")
return None
def get_model_workflows(base_path: str) -> list[str]:
workflows = []
main_json = f"{base_path}.json"
if os.path.exists(main_json):
workflows.append(main_json)
preview_json = f"{base_path}.preview.json"
if os.path.exists(preview_json) and preview_json not in workflows:
workflows.append(preview_json)
example_jsons = glob.glob(f"{base_path}.example.*.json")
example_jsons.sort()
for ej in example_jsons:
if ej not in workflows:
workflows.append(ej)
return workflows
def get_model_metadata_file(base_path: str) -> str | None:
civitai_info = f"{base_path}.civitai.info"
if os.path.exists(civitai_info):
return civitai_info
metadata_json = f"{base_path}.metadata.json"
if os.path.exists(metadata_json):
return metadata_json
return None
class BaseModelInfoLoader(io.ComfyNode):
CATEGORY = "ModelUtils/Info"
@classmethod
def get_info(cls, category, name, workflow_index):
full_path = folder_paths.get_full_path(category, name)
if not full_path or not os.path.exists(full_path):
return (None, "{}", "{}")
base_path = os.path.splitext(full_path)[0]
# 1. Image Preview
image_tensor = None
previews = get_potential_preview_files(base_path)
for p in previews:
if os.path.exists(p) and not p.lower().endswith('.mp4'):
image_tensor = load_image_tensor(p)
if image_tensor is not None:
break
# 2. Workflow
workflow_json = "{}"
workflows = get_model_workflows(base_path)
if workflows and 0 <= workflow_index < len(workflows):
try:
with open(workflows[workflow_index], 'r', encoding='utf-8') as f:
workflow_json = f.read()
except Exception as e:
print(f"Error reading workflow {workflows[workflow_index]}: {e}")
# 3. Metadata
metadata_json = "{}"
metadata_file = get_model_metadata_file(base_path)
if metadata_file:
try:
with open(metadata_file, 'r', encoding='utf-8') as f:
metadata_json = f.read()
except Exception as e:
print(f"Error reading metadata {metadata_file}: {e}")
return (image_tensor, workflow_json, metadata_json)
class CheckpointInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Checkpoint Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("checkpoint", options=folder_paths.get_filename_list("checkpoints"), tooltip="Checkpoint whose downloaded information and example workflow will be returned."),
io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, checkpoint, workflow_index):
res = cls.get_info("checkpoints", checkpoint, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class LoRAInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="LoRA Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("lora", options=folder_paths.get_filename_list("loras"), tooltip="LoRA whose downloaded information and example workflow will be returned."),
io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, lora, workflow_index):
res = cls.get_info("loras", lora, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class EmbeddingInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Embedding Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings"), tooltip="Embedding whose downloaded information and example workflow will be returned."),
io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, embedding, workflow_index):
res = cls.get_info("embeddings", embedding, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class VAEInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="VAE Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("vae", options=folder_paths.get_filename_list("vae"), tooltip="VAE whose downloaded information and example workflow will be returned."),
io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, vae, workflow_index):
res = cls.get_info("vae", vae, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class ControlNetInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="ControlNet Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("controlnet", options=folder_paths.get_filename_list("controlnet"), tooltip="ControlNet whose downloaded information and example workflow will be returned."),
io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, controlnet, workflow_index):
res = cls.get_info("controlnet", controlnet, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class DiffusionModelInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Diffusion Model Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("diffusion_model", options=folder_paths.get_filename_list("diffusion_models"), tooltip="Diffusion model whose downloaded information and example workflow will be returned."),
io.Int.Input("workflow_index", default=0, min=0, max=100, tooltip="Zero-based index of the saved example workflow to return when multiple workflows are available.")
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, diffusion_model, workflow_index):
res = cls.get_info("diffusion_models", diffusion_model, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
+197 -65
View File
@@ -1,13 +1,32 @@
from contextlib import closing
import os
import re
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer, stream_work_units
from .artifact_paths import canonical_model_artifact_path
from .quantization_guard import (
UPSTREAM_WRITER_STORAGE_CODES,
_is_low_bit,
storage_dtype_code,
write_preserved_tensor,
)
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
use_regex: bool, output_filename: str) -> str:
def _prune_keys(
model_name: str,
model_type: str,
keys_to_prune_str: str,
use_regex: bool,
output_filename: str,
include_mode: bool = False,
) -> str:
"""Shared logic for all PruneKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
@@ -21,32 +40,75 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
else:
model_path_to_load = model_path
model_weights, metadata = load_metadata_from_safetensors(model_path_to_load)
patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
if not patterns:
raw_patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
if not raw_patterns:
raise ValueError("No keys/patterns provided to prune.")
pruned_tensors = {}
for key, tensor in model_weights.items():
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
compiled_patterns = []
if use_regex:
for pat in raw_patterns:
try:
compiled_patterns.append(re.compile(pat))
except re.error as exc:
raise ValueError(f"Invalid regex pattern '{pat}': {exc}") from exc
if not is_match:
pruned_tensors[key] = tensor
model_size_gb = estimate_model_size(model_path_to_load)
if model_size_gb:
prepare_for_large_operation(model_size_gb * 1.2)
model_dir = folder_paths.get_folder_paths(model_type)[0]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
save_file(pruned_tensors, output_path, metadata)
handler = MemoryEfficientSafeOpen(model_path_to_load, low_memory=True)
try:
metadata = (handler.metadata() or {}).copy()
all_keys = list(handler.keys())
kept_keys = []
for key in all_keys:
matched = (
any(pat.search(key) is not None for pat in compiled_patterns)
if use_regex
else any(pat in key for pat in raw_patterns)
)
keep = matched if include_mode else not matched
if keep:
kept_keys.append(key)
return output_path
work_units = []
raw_preserve_keys = set()
for key in kept_keys:
code = storage_dtype_code(handler, key)
if _is_low_bit(handler, key) or (
code is not None and code not in UPSTREAM_WRITER_STORAGE_CODES
):
raw_preserve_keys.add(key)
work_units.append((key, {}))
else:
work_units.append((key, {0: [key]}))
pbar = comfy.utils.ProgressBar(len(all_keys))
with atomic_uel_writer(output_path, metadata) as writer, closing(
stream_work_units({0: handler}, work_units, pin_memory=False)
) as streamed:
for key, loaded in tqdm(streamed, total=len(work_units), desc="Pruning keys", unit="keys"):
if key in raw_preserve_keys:
write_preserved_tensor(
writer, key, handler, output_key=key, force_raw=True
)
else:
tensor = loaded[(0, key)]
try:
write_preserved_tensor(
writer, key, handler, output_key=key, tensor=tensor
)
finally:
del tensor
pbar.update(1)
pbar.update(len(all_keys) - len(work_units))
finally:
handler.__exit__(None, None, None)
cleanup_after_operation()
return output_name
class ModelPruneKeys(io.ComfyNode):
@@ -61,21 +123,35 @@ class ModelPruneKeys(io.ComfyNode):
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Diffusion model from which matching tensor keys will be removed.",
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_model"),
io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
io.String.Input("output_filename", default="pruned_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, diffusionmodel_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(diffusionmodel_name, "diffusion_models",
keys_to_prune, use_regex, output_filename)
def execute(
cls,
diffusionmodel_name: str,
keys_to_prune: str,
use_regex: bool,
output_filename: str,
include_mode: bool = False,
) -> io.NodeOutput:
path = _prune_keys(
diffusionmodel_name,
"diffusion_models",
keys_to_prune,
use_regex,
output_filename,
include_mode=include_mode,
)
return io.NodeOutput(path)
@@ -91,21 +167,35 @@ class TextEncoderPruneKeys(io.ComfyNode):
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
tooltip="Text encoder from which matching tensor keys will be removed.",
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_textencoder"),
io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
io.String.Input("output_filename", default="pruned_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, textencoder_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(textencoder_name, "text_encoders",
keys_to_prune, use_regex, output_filename)
def execute(
cls,
textencoder_name: str,
keys_to_prune: str,
use_regex: bool,
output_filename: str,
include_mode: bool = False,
) -> io.NodeOutput:
path = _prune_keys(
textencoder_name,
"text_encoders",
keys_to_prune,
use_regex,
output_filename,
include_mode=include_mode,
)
return io.NodeOutput(path)
@@ -121,21 +211,35 @@ class LoRAPruneKeys(io.ComfyNode):
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
tooltip="LoRA from which matching tensor keys will be removed.",
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_lora"),
io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
io.String.Input("output_filename", default="pruned_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, lora_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(lora_name, "loras",
keys_to_prune, use_regex, output_filename)
def execute(
cls,
lora_name: str,
keys_to_prune: str,
use_regex: bool,
output_filename: str,
include_mode: bool = False,
) -> io.NodeOutput:
path = _prune_keys(
lora_name,
"loras",
keys_to_prune,
use_regex,
output_filename,
include_mode=include_mode,
)
return io.NodeOutput(path)
@@ -151,21 +255,35 @@ class CheckpointPruneKeys(io.ComfyNode):
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="Checkpoint from which matching tensor keys will be removed.",
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_checkpoint"),
io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
io.String.Input("output_filename", default="pruned_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, ckpt_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(ckpt_name, "checkpoints",
keys_to_prune, use_regex, output_filename)
def execute(
cls,
ckpt_name: str,
keys_to_prune: str,
use_regex: bool,
output_filename: str,
include_mode: bool = False,
) -> io.NodeOutput:
path = _prune_keys(
ckpt_name,
"checkpoints",
keys_to_prune,
use_regex,
output_filename,
include_mode=include_mode,
)
return io.NodeOutput(path)
@@ -181,19 +299,33 @@ class EmbeddingPruneKeys(io.ComfyNode):
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
tooltip="Embedding from which matching tensor keys will be removed.",
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_embedding"),
io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."),
io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."),
io.String.Input("output_filename", default="pruned_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, embedding: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(embedding, "embeddings",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
def execute(
cls,
embedding: str,
keys_to_prune: str,
use_regex: bool,
output_filename: str,
include_mode: bool = False,
) -> io.NodeOutput:
path = _prune_keys(
embedding,
"embeddings",
keys_to_prune,
use_regex,
output_filename,
include_mode=include_mode,
)
return io.NodeOutput(path)
+508
View File
@@ -0,0 +1,508 @@
"""Fail-safe detection for quantized or isolated low-bit safetensors inputs."""
import logging
import json
import struct
import torch
def _low_bit_dtypes():
names = (
"int8",
"uint8",
"float4_e2m1fn_x2",
"float8_e4m3fn",
"float8_e4m3fnuz",
"float8_e5m2",
"float8_e5m2fnuz",
"float8_e8m0fnu",
)
return {getattr(torch, name) for name in names if hasattr(torch, name)}
LOW_BIT_DTYPES = _low_bit_dtypes()
LOW_BIT_STORAGE_CODES = {"I8", "U8", "F4"}
UPSTREAM_WRITER_STORAGE_CODES = {
"F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8",
"U64", "U32", "U16", "BOOL", "C64", "F8_E5M2", "F8_E4M3",
}
_QUANT_SIDECARS = (
"comfy_quant", "weight_scale", "weight_scale_2", "input_scale",
"pre_quant_scale", "weight_s_rel", "weight_s_channel",
"weight_correction", "weight_codebook", "scale_weight", "scale_input",
)
def diffusion_key_map(keys, label):
"""Map optional diffusion wrappers to each file's physical tensor keys."""
mapped = {}
for key in sorted(keys):
logical = key
for prefix in ("model.diffusion_model.", "diffusion_model."):
if key.startswith(prefix):
logical = key[len(prefix):]
break
previous = mapped.get(logical)
if previous is not None:
raise ValueError(
f"{label} has duplicate diffusion layer {logical!r}: "
f"{previous!r} and {key!r}"
)
mapped[logical] = key
return mapped
class DiffusionQuantization:
"""Index and decode Core quantized diffusion weights one layer at a time.
The handler is a UEL low-memory handler. Header inspection uses only its
keys, shapes, dtypes, and metadata; tensor values are supplied to ``decode``
by the caller's bounded UEL work unit.
"""
def __init__(self, handler, model_path: str, label: str):
self.handler = handler
self.model_path = model_path
self.label = label
self._keys = tuple(handler.keys())
self._key_set = set(self._keys)
self._metadata = dict(handler.metadata() or {})
self._layer_configs = {}
self._architecture_shapes = None
self._legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in self._keys)
encoded = self._metadata.get("_quantization_metadata")
if encoded:
try:
layers = json.loads(encoded).get("layers", {})
except (TypeError, ValueError) as exc:
raise ValueError(
f"[{label}] Invalid _quantization_metadata in {model_path}: {exc}"
) from exc
if not isinstance(layers, dict):
raise ValueError(f"[{label}] Invalid quantization layers metadata in {model_path}")
self._layer_configs.update(layers)
quantized = set()
for sidecar in self._keys:
if not sidecar.endswith(".comfy_quant"):
continue
stem = sidecar[:-len(".comfy_quant")]
weight_key = f"{stem}.weight"
if weight_key not in self._key_set:
raise ValueError(f"[{label}] Quantization sidecar has no weight: {sidecar}")
quantized.add(weight_key)
for stem in self._layer_configs:
weight_key = f"{stem}.weight"
if weight_key in self._key_set:
quantized.add(weight_key)
if self._legacy_scaled_fp8:
for sidecar in self._keys:
if sidecar.endswith(".scale_weight"):
stem = sidecar[:-len(".scale_weight")]
weight_key = f"{stem}.weight"
if weight_key in self._key_set:
quantized.add(weight_key)
self._layer_configs.setdefault(stem, {"format": "float8_e4m3fn"})
self.quantized_keys = frozenset(quantized)
sidecars = set()
for key in self.quantized_keys:
stem = key[:-len(".weight")]
for name in _QUANT_SIDECARS:
candidate = f"{stem}.{name}"
if candidate in self._key_set:
sidecars.add(candidate)
for key in self._keys:
if key.endswith(".comfy_quant") and any(
key == f"{weight[:-len('.weight')]}.comfy_quant" for weight in self.quantized_keys
):
sidecars.add(key)
if self._legacy_scaled_fp8:
sidecars.update(key for key in self._keys if key.endswith("scaled_fp8"))
self._sidecars = frozenset(sidecars)
self.data_keys = tuple(key for key in self._keys if key not in self._sidecars)
self.isolated_low_bit_keys = {
key for key in self.data_keys
if key not in self.quantized_keys and _is_low_bit(handler, key)
}
if self.isolated_low_bit_keys:
logging.warning(
"[%s] %s contains %d isolated low-bit tensor(s); preserving/excluding unchanged: %s",
label, model_path, len(self.isolated_low_bit_keys),
", ".join(sorted(self.isolated_low_bit_keys)),
)
def required_keys(self, key: str) -> list[str]:
if key not in self.quantized_keys:
return [key]
stem = key[:-len(".weight")]
return [key] + [
f"{stem}.{name}" for name in _QUANT_SIDECARS
if f"{stem}.{name}" in self._key_set
]
def _architecture_shape(self, key: str) -> tuple[int, ...]:
if self._architecture_shapes is None:
from comfy import model_detection
# Core detection examines keys, shapes and dtypes. Meta tensors supply
# that inventory without materializing any checkpoint tensor data.
state = {
name: torch.empty(
tuple(self.handler.get_shape(name)),
dtype=self.handler.get_dtype(name), device="meta",
)
for name in self._keys
}
prefix = model_detection.unet_prefix_from_state_dict(state)
if self.quantized_keys and all(
name.startswith("model.diffusion_model.") for name in self.quantized_keys
):
prefix = "model.diffusion_model."
elif not any(name.startswith(prefix) for name in state):
prefix = ""
detected = {
name[len(prefix):]: value
for name, value in state.items() if name.startswith(prefix)
}
try:
config = model_detection.model_config_from_unet(
detected, "", metadata=self._metadata,
)
if config is None:
raise ValueError("Core could not detect this diffusion architecture")
with torch.device("meta"):
model = config.get_model(detected, "", device=torch.device("meta"))
modules = dict(getattr(model, "diffusion_model", model).named_modules())
shapes = {}
for weight_key in self.quantized_keys:
module_name = weight_key[len(prefix):].removesuffix(".weight")
module = modules.get(module_name)
if module is None:
continue
shape = getattr(module, "_orig_shape", None)
if shape is None:
weight = getattr(module, "weight", None)
shape = getattr(weight, "shape", None)
if shape is not None:
shapes[weight_key] = tuple(shape)
self._architecture_shapes = shapes
except Exception as exc:
raise ValueError(
f"[{self.label}] Cannot determine exact quantized layer shapes from "
f"Core architecture for {self.model_path}: {exc}"
) from exc
finally:
state.clear()
detected.clear()
if "model" in locals():
del model
shape = self._architecture_shapes.get(key)
if shape is None:
raise ValueError(
f"[{self.label}] Core architecture has no logical shape for quantized layer {key}"
)
stored = tuple(self.handler.get_shape(key))
if len(shape) != 2 or len(stored) != 2:
raise ValueError(f"[{self.label}] Invalid padded quantized shape for {key}: {shape}, {stored}")
stem = key[:-len(".weight")]
config = self._layer_configs.get(stem)
nvfp4 = (isinstance(config, dict) and config.get("format") == "nvfp4") or (
f"{stem}.weight_scale_2" in self._key_set
)
if nvfp4:
expected = ((shape[0] + 15) // 16 * 16, (shape[1] + 15) // 16 * 8)
else:
expected = ((shape[0] + 31) // 32 * 32, (shape[1] + 31) // 32 * 32)
if stored != expected:
raise ValueError(
f"[{self.label}] Core architecture shape {shape} disagrees with packed storage {stored} for {key}"
)
return shape
def _format(self, key: str, tensors=None):
stem = key[:-len(".weight")]
config = self._layer_configs.get(stem)
sidecar = f"{stem}.comfy_quant"
if tensors is None and config is None and sidecar in self._key_set:
stream = self.handler.async_stream(
[sidecar], batch_size=1, prefetch_batches=1, pin_memory=False,
)
yielded_keys = []
try:
batch = next(stream)
yielded_keys = [name for name, _ in batch]
if len(batch) != 1 or batch[0][0] != sidecar:
raise RuntimeError(f"[{self.label}] Unexpected quantization sidecar for {key}")
tensors = {sidecar: batch[0][1]}
finally:
try:
for yielded_key in yielded_keys:
self.handler.mark_processed(yielded_key)
finally:
stream.close()
if tensors is not None and sidecar in tensors:
raw = tensors[sidecar].detach().to(device="cpu").contiguous().reshape(-1).tolist()
try:
config = json.loads(bytes(raw).decode("utf-8"))
except (UnicodeDecodeError, TypeError, ValueError) as exc:
raise ValueError(f"[{self.label}] Invalid quantization sidecar {sidecar}: {exc}") from exc
self._layer_configs[stem] = config
if config is None:
raise ValueError(f"[{self.label}] Missing quantization config for {key}")
if not isinstance(config, dict):
raise ValueError(f"[{self.label}] Invalid quantization config for {key}: expected object")
fmt = config.get("format")
try:
from comfy.quant_ops import QUANT_ALGOS
except ImportError as exc:
raise RuntimeError("ComfyUI quantization support is unavailable") from exc
if fmt not in QUANT_ALGOS:
raise ValueError(f"[{self.label}] Unsupported ComfyUI quantization format {fmt!r} for {key}")
return fmt, config
def logical_shape(self, key: str) -> tuple[int, ...]:
if key not in self.quantized_keys:
return tuple(self.handler.get_shape(key))
fmt, _ = self._format(key)
if fmt == "nvfp4" or fmt == "mxfp8":
return self._architecture_shape(key)
shape = tuple(self.handler.get_shape(key))
if fmt in {"convrot_w4a4", "asym_w4a8_int8"}:
if len(shape) != 2:
raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}")
return shape[0], shape[1] * 2
if fmt == "w6a8_int8":
if len(shape) != 2:
raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}")
if (shape[1] * 8) % 6 != 0:
raise ValueError(f"[{self.label}] Invalid 6-bit packed width {shape[1]} for {key}")
return shape[0], shape[1] * 8 // 6
return shape
def logical_dtype(self, key: str, compute_dtype: torch.dtype) -> torch.dtype:
if key not in self.quantized_keys:
return self.handler.get_dtype(key)
return compute_dtype
def decode(self, key: str, tensors, compute_dtype: torch.dtype, device="cpu") -> torch.Tensor:
if key not in self.quantized_keys:
raise ValueError(f"[{self.label}] {key} is not a recognized quantized diffusion weight")
fmt, config = self._format(key, tensors)
try:
from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
except ImportError as exc:
raise RuntimeError("ComfyUI quantization support is unavailable") from exc
stem = key[:-len(".weight")]
device = torch.device(device)
shape = self.logical_shape(key)
def tensor(name, *, dtype=None, required=False):
tensor_key = f"{stem}.{name}"
if tensor_key not in tensors and name == "weight_scale":
legacy_key = f"{stem}.scale_weight"
if legacy_key in tensors:
tensor_key = legacy_key
value = tensors.get(tensor_key)
if value is None:
if required:
raise ValueError(f"[{self.label}] Missing quantization tensor {tensor_key} for {key}")
return None
value = value.to(device=device)
if dtype is not None:
value = value.view(dtype=dtype)
return value
params_conf = config.get("params", {})
if not isinstance(params_conf, dict):
params_conf = {}
if fmt in ("float8_e4m3fn", "float8_e5m2"):
values = {"scale": tensor("weight_scale", required=True)}
elif fmt == "mxfp8":
values = {"scale": tensor("weight_scale", dtype=torch.float8_e8m0fnu, required=True)}
elif fmt == "nvfp4":
values = {
"scale": tensor("weight_scale_2", required=True),
"block_scale": tensor("weight_scale", dtype=torch.float8_e4m3fn, required=True),
}
elif fmt == "int8_tensorwise":
values = {"scale": tensor("weight_scale", required=True)}
if config.get("convrot", params_conf.get("convrot", False)):
values.update({
"convrot": True,
"convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))),
})
elif fmt == "convrot_w4a4":
values = {
"scale": tensor("weight_scale", required=True),
"convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))),
"quant_group_size": 64,
"linear_dtype": config.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
}
elif fmt in ("asym_w4a8_int8", "w6a8_int8"):
if f"{stem}.weight_correction" in tensors:
raise ValueError(
f"[{self.label}] Core's W4A8 loader does not apply weight_correction for {key}"
)
scale_value = tensors.get(f"{stem}.weight_s_rel")
values = {
"scale": tensor("weight_s_rel", dtype=torch.float8_e4m3fn if scale_value is not None and scale_value.dtype == torch.uint8 else None, required=True),
"s_channel": tensor("weight_s_channel", required=True),
"codebook": tensor("weight_codebook"),
"group_size": int(config.get("group_size", params_conf.get("group_size", 16))),
"convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))),
}
else:
raise ValueError(f"[{self.label}] Unsupported ComfyUI quantization format {fmt!r} for {key}")
layout_type = QUANT_ALGOS[fmt]["comfy_tensor_layout"]
layout = get_layout_class(layout_type)
params = layout.Params(**values, orig_dtype=compute_dtype, orig_shape=shape)
qdata = tensors[key].to(device=device, dtype=QUANT_ALGOS[fmt]["storage_t"])
return QuantizedTensor(qdata, layout_type, params).dequantize()
def output_metadata(self) -> dict:
metadata = dict(self._metadata)
metadata.pop("_quantization_metadata", None)
return metadata
def _header(handler):
if getattr(handler, "_header", None) is not None:
return handler._header
with open(handler.filename, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
return json.loads(source.read(header_size).decode("utf-8"))
def storage_dtype_code(handler, key: str) -> str | None:
entry = _header(handler).get(key)
return entry.get("dtype") if entry else None
def _is_low_bit(handler, key: str) -> bool:
storage_code = storage_dtype_code(handler, key)
if storage_code in LOW_BIT_STORAGE_CODES or (
storage_code is not None and storage_code.startswith("F8_")
):
return True
return handler.get_dtype(key) in LOW_BIT_DTYPES
def inspect_low_bit_input(handler, label: str, operation: str) -> set[str]:
"""Reject quantized models; return isolated low-bit keys that must not be transformed."""
keys = list(handler.keys())
metadata = handler.metadata() or {}
quant_keys = [key for key in keys if key.endswith(".comfy_quant")]
has_quant_metadata = "_quantization_metadata" in metadata
has_legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in keys)
if quant_keys or has_quant_metadata or has_legacy_scaled_fp8:
raise ValueError(
f"[{operation}] {label} uses ComfyUI quantization metadata; quantized inputs are unsupported"
)
low_bit_keys = {key for key in keys if _is_low_bit(handler, key)}
if low_bit_keys:
logging.warning(
"[%s] %s contains %d isolated low-bit tensor(s); "
"preserving/excluding unchanged: %s",
operation,
label,
len(low_bit_keys),
", ".join(sorted(low_bit_keys)),
)
return low_bit_keys
def layer_has_low_bit(block_keys: dict, low_bit_keys: set[str]) -> bool:
return any(
block_keys.get(name) in low_bit_keys
for name in (
"down", "up", "mid", "reshape", "dora_scale",
"diff", "diff_b", "w_norm", "b_norm", "set_weight",
)
)
def write_preserved_tensor(
writer,
key: str,
handler,
output_key: str | None = None,
*,
force_raw: bool = False,
tensor=None,
) -> None:
"""Copy one tensor's original safetensors bytes and dtype into an active writer."""
destination = output_key or key
header = _header(handler)
entry = header[key]
if not force_raw and entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES:
if tensor is not None:
writer.write(destination, tensor.cpu().contiguous())
return
stream = handler.async_stream(
[key], batch_size=1, prefetch_batches=1, pin_memory=False
)
yielded = False
try:
batch = next(stream)
if len(batch) != 1 or batch[0][0] != key:
raise RuntimeError(f"UEL yielded an unexpected preserved tensor for {key}")
yielded = True
writer.write(destination, batch[0][1].cpu().contiguous())
finally:
if yielded:
handler.mark_processed(key)
stream.close()
return
source_start, source_end = entry["data_offsets"]
byte_size = source_end - source_start
with writer._lock:
if destination in writer._manifest:
raise ValueError(f"Tensor '{destination}' has already been written.")
target_start = writer._current_data_offset
target_end = target_start + byte_size
writer._manifest[destination] = {
"dtype": entry["dtype"],
"shape": entry["shape"],
"data_offsets": [target_start, target_end],
}
writer._current_data_offset = target_end
source_path = handler.filename
source_absolute = 8 + getattr(handler, "_header_size", 0) + source_start
if not getattr(handler, "_header_size", None):
with open(source_path, "rb") as source:
source_absolute = 8 + struct.unpack("<Q", source.read(8))[0] + source_start
target_absolute = 8 + writer.max_header_bytes + target_start
def copy_raw():
try:
with open(source_path, "rb") as source:
source.seek(source_absolute)
remaining = byte_size
with writer._lock:
writer._file.seek(target_absolute)
while remaining:
chunk = source.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise EOFError(f"Unexpected EOF while preserving tensor '{key}'")
writer._file.write(chunk)
remaining -= len(chunk)
finally:
writer._semaphore.release()
writer._semaphore.acquire()
writer._futures.append(writer._executor.submit(copy_raw))
+68 -44
View File
@@ -1,11 +1,18 @@
import os
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
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)
@@ -20,7 +27,9 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
else:
model_path_to_load = model_path
model_weights, metadata = load_metadata_from_safetensors(model_path_to_load)
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
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()]
@@ -30,18 +39,28 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
key_map = dict(zip(old_keys, new_keys))
original_keys = list(model_weights.keys())
renamed_tensors = {}
for key in original_keys:
new_key = key_map.get(key, key)
renamed_tensors[new_key] = model_weights[key]
output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
model_dir = folder_paths.get_folder_paths(model_type)[0]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Stream tensors, rename on the fly, write immediately
with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
metadata = handler.metadata() or {}
with atomic_uel_writer(output_path, metadata) as writer:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
stream = handler.async_stream(
original_keys, batch_size=1, prefetch_batches=1, pin_memory=False
)
for batch in tqdm(stream, total=len(original_keys), desc="Renaming keys", unit="keys"):
key, tensor = batch[0]
new_key = key_map.get(key, key)
writer.write_batch([(new_key, tensor.contiguous())])
handler.mark_processed(key)
pbar.update(1)
save_file(renamed_tensors, output_path, metadata)
# Cleanup after operation
cleanup_after_operation()
return output_path
return output_name
class ModelRenameKeys(io.ComfyNode):
@@ -56,20 +75,21 @@ class ModelRenameKeys(io.ComfyNode):
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Diffusion model whose tensor keys will be renamed.",
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_model"),
io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
io.String.Input("output_filename", default="renamed_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, diffusionmodel_name: str, old_keys: str,
def execute(cls, diffusionmodel_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(diffusionmodel_name, "diffusion_models",
path = _rename_keys(diffusionmodel_name, "diffusion_models",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
@@ -86,20 +106,21 @@ class TextEncoderRenameKeys(io.ComfyNode):
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
tooltip="Text encoder whose tensor keys will be renamed.",
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_textencoder"),
io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
io.String.Input("output_filename", default="renamed_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, textencoder_name: str, old_keys: str,
def execute(cls, textencoder_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(textencoder_name, "text_encoders",
path = _rename_keys(textencoder_name, "text_encoders",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
@@ -116,20 +137,21 @@ class LoRARenameKeys(io.ComfyNode):
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
tooltip="LoRA whose tensor keys will be renamed.",
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_lora"),
io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
io.String.Input("output_filename", default="renamed_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, lora_name: str, old_keys: str,
def execute(cls, lora_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(lora_name, "loras",
path = _rename_keys(lora_name, "loras",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
@@ -146,20 +168,21 @@ class CheckpointRenameKeys(io.ComfyNode):
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="Checkpoint whose tensor keys will be renamed.",
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_checkpoint"),
io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
io.String.Input("output_filename", default="renamed_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, ckpt_name: str, old_keys: str,
def execute(cls, ckpt_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(ckpt_name, "checkpoints",
path = _rename_keys(ckpt_name, "checkpoints",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
@@ -176,19 +199,20 @@ class EmbeddingRenameKeys(io.ComfyNode):
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
tooltip="Embedding whose tensor keys will be renamed.",
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_embedding"),
io.String.Input("old_keys", multiline=True, default="", tooltip="Existing key text or pattern, one entry per line, paired by line with New Keys."),
io.String.Input("new_keys", multiline=True, default="", tooltip="Replacement key text, one entry per line, paired by line with Old Keys."),
io.String.Input("output_filename", default="renamed_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, embedding: str, old_keys: str,
def execute(cls, embedding: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(embedding, "embeddings",
path = _rename_keys(embedding, "embeddings",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
return io.NodeOutput(path)
+794
View File
@@ -0,0 +1,794 @@
"""
Text Encoder LoRA and DoRA Extraction.
Specifically targets extracting text encoder deltas as LoRA/DoRA with custom mapping rules
for layer name prefixes and PEFT weight formatting.
"""
import torch
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
from .extraction_stream import dora_difference, lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
from .layer_parameters import parameter_input, resolve_layer_parameters
# Import SVD rank utilities from standard lora extract file to avoid duplication
from .lora_extract_svd import (
_svd_extract_linear,
_svd_extract_conv,
_extract_chunked_layer,
_detect_fused_layer,
_compile_patterns,
_matches_any_pattern
)
def _format_te_lora_key(key: str) -> str:
"""
Format key specifically for Text Encoder LoRA/DoRA saving.
Corrects layer name prefixing based on specific target mappings.
"""
if key.endswith(".weight"):
key = key[:-7]
rules = [
("model.language_model.layers", "text_encoders.transformer.model.layers"),
("model.visual.blocks", "text_encoders.transformer.visual.blocks"),
("model.visual.deepstack_merger_list", "text_encoders.transformer.visual.deepstack_merger_list"),
("model.visual.merger", "text_encoders.transformer.visual.merger"),
("model.visual.pos_embed", "text_encoders.transformer.visual.pos_embed"),
("model.layers", "text_encoders.transformer.model.layers"),
("vision_model.embeddings.position_embedding", "text_encoders.transformer.vision_model.embeddings.position_embedding"),
("vision_model.encoder.layers", "text_encoders.transformer.vision_model.encoder.layers"),
]
for prefix, target in rules:
if key.startswith(prefix):
return target + key[len(prefix):]
# Prepend text_encoders. as default generic fallback
if key.startswith("text_encoders."):
return key
return f"text_encoders.{key}"
def extract_te_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
is_dora: bool = False,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = False,
glob_skip_patterns: bool = False,
knee_probe_offset: int = 32,
include_mode: bool = False,
layer_parameters=None,
) -> None:
"""
Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[TE Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = sorted(k for k in keys_a if k.endswith(".weight"))
if mode == "fixed":
profile = "extract:fixed"
defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
"clamp_quantile": clamp_quantile, "min_diff": min_diff}
elif mode == "ratio":
profile = "extract:ratio"
defaults = {"linear_ratio": linear_param, "conv_ratio": conv_param,
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
elif mode == "quantile":
profile = "extract:quantile"
defaults = {"linear_quantile": linear_param, "conv_quantile": conv_param,
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
elif mode == "sv_fro":
profile = "extract:frobenius"
defaults = {"linear_target": linear_param, "conv_target": conv_param,
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
else:
profile = "extract:knee"
defaults = {"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
"clamp_quantile": clamp_quantile, "min_diff": min_diff}
resolved_parameters = resolve_layer_parameters(
layer_parameters, profile, weight_keys, defaults, node_name="Text Encoder Extract"
)
work_units = [
(key, key, key if key in keys_b else None)
for key in weight_keys
if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
and not (key not in keys_b and mismatch_mode == "skip")
]
pbar = comfy.utils.ProgressBar(len(work_units))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key, cpu_a, cpu_b):
values = resolved_parameters.get(key, defaults)
if mode == "fixed":
layer_linear_param, layer_conv_param = values["linear_dim"], values["conv_dim"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
elif mode == "ratio":
layer_linear_param, layer_conv_param = values["linear_ratio"], values["conv_ratio"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
elif mode == "quantile":
layer_linear_param, layer_conv_param = values["linear_quantile"], values["conv_quantile"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
elif mode == "sv_fro":
layer_linear_param, layer_conv_param = values["linear_target"], values["conv_target"]
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
else:
layer_linear_param, layer_conv_param = linear_param, conv_param
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
layer_clamp_quantile = values["clamp_quantile"]
layer_min_diff = values["min_diff"]
lora_name = _format_te_lora_key(key)
weight_diff = None
layer_device = device
if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)):
return "skipped", None
# Load tensors with pinned memory for CUDA
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
if is_dora:
weight_diff, dora_scale, layer_device = dora_difference(
cpu_a, None, device
)
else:
weight_diff, layer_device = lora_difference(cpu_a, None, device)
dora_scale = None
else:
if cpu_a.shape != cpu_b.shape:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
if is_dora:
weight_diff, dora_scale, layer_device = dora_difference(
cpu_a, None, device
)
else:
weight_diff, layer_device = lora_difference(cpu_a, None, device)
dora_scale = None
else:
if cpu_a.ndim < 2:
return "skipped", None
if is_dora:
weight_diff, dora_scale, layer_device = dora_difference(
cpu_a, cpu_b, device
)
else:
weight_diff, layer_device = lora_difference(
cpu_a, cpu_b, device
)
dora_scale = None
if weight_diff is None:
raise RuntimeError(f"No extraction tensor was produced for {key}")
# Skip small differences
if layer_min_diff > 0 and weight_diff.abs().max() < layer_min_diff:
del weight_diff
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
svd_weight = weight_diff
try:
if is_conv:
result, mode_str = retry_cuda_oom_on_cpu(
lambda: _svd_extract_conv(svd_weight, mode, layer_conv_param, layer_device, layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
lambda: _svd_extract_conv(svd_weight.cpu(), mode, layer_conv_param, "cpu", layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
layer_device,
)
else:
result, mode_str = retry_cuda_oom_on_cpu(
lambda: _svd_extract_linear(svd_weight, mode, layer_linear_param, layer_device, layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
lambda: _svd_extract_linear(svd_weight.cpu(), mode, layer_linear_param, "cpu", layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
layer_device,
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[TE Extract] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_layer(
weight_diff, num_chunks, mode, layer_linear_param, layer_device,
layer_linear_max_rank, knee_probe_offset, layer_clamp_quantile,
)
if lora_up is not None:
# PEFT suffixes
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
if is_dora and dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
print(f"[TE Extract] Failed: {key}: {e}")
del weight_diff
return "skipped", None
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
# Write to PEFT weight suffixes
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
if is_dora and dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
return status, layer_results
writer_context = atomic_uel_writer(output_path)
writer = writer_context.__enter__()
try:
stream = paired_async_tensors(
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
)
for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting TE Layers", unit="layers"):
status, layer_sd = _process_layer(key, cpu_a, cpu_b)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
stats["skipped"] += len(weight_keys) - len(work_units)
except BaseException as exc:
writer_context.__exit__(type(exc), exc, exc.__traceback__)
raise
else:
writer_context.__exit__(None, None, None)
print(f"[TE Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> tuple[str, str]:
"""Build output path for LoRA file."""
return canonical_model_artifact_path("loras", output_filename)
def _get_te_model_inputs():
return [
io.Combo.Input("model_a", options=folder_paths.get_filename_list("text_encoders"),
tooltip="Finetuned Text Encoder model (A - B = LoRA)"),
io.Combo.Input("model_b", options=folder_paths.get_filename_list("text_encoders"),
tooltip="Base Text Encoder model (A - B = LoRA)"),
]
def _get_common_inputs():
return [
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."),
io.Boolean.Input("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
tooltip="Clamp outlier singular values"),
io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", tooltip="Handle missing or incompatible text-encoder tensors by skipping them, substituting zeros where supported, or aborting."),
io.String.Input("output_filename", default="extracted_te_lora", tooltip="Output filename without extension, written under the ComfyUI LoRA directory."),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save extracted text-encoder LoRA factors."),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer extraction arithmetic; CUDA out-of-memory retries the affected layer on CPU where supported."),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
io.Boolean.Input("include_mode", default=False, tooltip="Use Skip Patterns as a whitelist instead. Only matching layers are extracted; an empty whitelist extracts nothing."),
]
# =============================================================================
# LoRA Text Encoder Extract Nodes
# =============================================================================
class TextEncoderLoRAExtractFixed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractFixed",
display_name="TE LoRA Extract (Fixed Rank)",
category="ModelUtils/LoRA Extract (TE)",
description="Extract Text Encoder LoRA with specified fixed rank for each layer type.",
inputs=[
*_get_te_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"),
*_get_common_inputs(),
parameter_input("extract:fixed"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, is_dora=False,
linear_max_rank=linear_dim, conv_max_rank=conv_dim,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractRatio(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractRatio",
display_name="TE LoRA Extract (Ratio)",
category="ModelUtils/LoRA Extract (TE)",
description="Keep singular values > max(S) / ratio.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:ratio"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, is_dora=False,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractQuantile(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractQuantile",
display_name="TE LoRA Extract (Quantile)",
category="ModelUtils/LoRA Extract (TE)",
description="Keep enough singular values to reach target percentage.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:quantile"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, is_dora=False,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractKnee(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractKnee",
display_name="TE LoRA Extract (Knee)",
category="ModelUtils/LoRA Extract (TE)",
description="Automatically find optimal rank using knee detection on singular value curve.",
inputs=[
*_get_te_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:knee"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, is_dora=False,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractFrobenius",
display_name="TE LoRA Extract (Frobenius)",
category="ModelUtils/LoRA Extract (TE)",
description="Preserve target fraction of Frobenius norm.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:frobenius"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, is_dora=False,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
# =============================================================================
# DoRA Text Encoder Extract Nodes
# =============================================================================
class TextEncoderDoRAExtractFixed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractFixed",
display_name="TE DoRA Extract (Fixed Rank)",
category="ModelUtils/DoRA Extract (TE)",
description="Extract Text Encoder DoRA with specified fixed rank for each layer type.",
inputs=[
*_get_te_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"),
*_get_common_inputs(),
parameter_input("extract:fixed"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, is_dora=True,
linear_max_rank=linear_dim, conv_max_rank=conv_dim,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractRatio(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractRatio",
display_name="TE DoRA Extract (Ratio)",
category="ModelUtils/DoRA Extract (TE)",
description="Keep singular values > max(S) / ratio.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:ratio"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, is_dora=True,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractQuantile(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractQuantile",
display_name="TE DoRA Extract (Quantile)",
category="ModelUtils/DoRA Extract (TE)",
description="Keep enough singular values to reach target percentage.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:quantile"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, is_dora=True,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractKnee(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractKnee",
display_name="TE DoRA Extract (Knee)",
category="ModelUtils/DoRA Extract (TE)",
description="Automatically find optimal rank using knee detection on singular value curve.",
inputs=[
*_get_te_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:knee"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, is_dora=True,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractFrobenius",
display_name="TE DoRA Extract (Frobenius)",
category="ModelUtils/DoRA Extract (TE)",
description="Preserve target fraction of Frobenius norm.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
parameter_input("extract:frobenius"),
],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False, layer_parameters=None) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, is_dora=True,
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode,
knee_probe_offset=probe_offset,
layer_parameters=layer_parameters,
)
return io.NodeOutput(output_name)
+115
View File
@@ -0,0 +1,115 @@
import os
import uuid
from collections import deque
from contextlib import contextmanager
from unifiedefficientloader import IncrementalSafetensorsWriter
class AsyncTensorCursor:
"""Ordered, bounded tensor access over one UEL asynchronous stream."""
def __init__(self, handler, keys, *, pin_memory: bool):
self.handler = handler
self.expected = tuple(keys)
self.position = 0
self.pending = deque()
self.stream = None
if self.expected:
self.stream = handler.async_stream(
list(self.expected),
batch_size=1,
prefetch_batches=1,
pin_memory=pin_memory,
)
def take(self, expected_key):
if self.position >= len(self.expected):
raise RuntimeError(f"Unexpected tensor request: {expected_key}")
planned = self.expected[self.position]
if planned != expected_key:
raise RuntimeError(
f"UEL stream order mismatch: expected {planned}, requested {expected_key}"
)
if not self.pending:
self.pending.extend(next(self.stream))
key, tensor = self.pending.popleft()
if key != expected_key:
raise RuntimeError(f"UEL yielded {key} while {expected_key} was expected")
self.position += 1
return tensor
def release(self, key):
self.handler.mark_processed(key)
def finish(self):
if self.position != len(self.expected) or self.pending:
raise RuntimeError("UEL stream ended before all planned tensors were consumed")
if self.stream is not None:
try:
next(self.stream)
except StopIteration:
return
raise RuntimeError("UEL stream yielded unplanned tensors")
def close(self):
if self.stream is not None:
close = getattr(self.stream, "close", None)
if close is not None:
close()
self.stream = None
self.pending.clear()
@contextmanager
def atomic_uel_writer(output_path: str, metadata: dict | None = None):
"""Write incrementally and expose the destination only after finalization."""
directory = os.path.dirname(output_path) or "."
os.makedirs(directory, exist_ok=True)
basename = os.path.basename(output_path)
temporary = os.path.join(directory, f".{basename}.{uuid.uuid4().hex}.tmp")
try:
with IncrementalSafetensorsWriter(
temporary,
metadata=metadata or {},
max_workers=1,
) as writer:
yield writer
os.replace(temporary, output_path)
finally:
if os.path.exists(temporary):
os.remove(temporary)
def stream_work_units(handlers, work_units, *, pin_memory: bool):
"""Yield bounded logical units from one ordered stream per input handler."""
planned = {source: [] for source in handlers}
for _, entries in work_units:
for source, keys in entries.items():
planned[source].extend(keys)
cursors = {
source: AsyncTensorCursor(handler, planned[source], pin_memory=pin_memory)
for source, handler in handlers.items()
}
try:
for logical_key, entries in work_units:
loaded = {}
consumed = []
try:
for source, keys in entries.items():
cursor = cursors[source]
for key in keys:
loaded[(source, key)] = cursor.take(key)
consumed.append((source, key))
yield logical_key, loaded
finally:
loaded.clear()
for source, key in consumed:
cursors[source].release(key)
for cursor in cursors.values():
cursor.finish()
finally:
for cursor in cursors.values():
cursor.close()
+8 -19
View File
@@ -1,7 +1,6 @@
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from comfy.model_management import get_torch_device
from .uel_io import atomic_uel_writer
def convert_pt_to_safetensors(pt_path, safe_path):
try:
@@ -12,23 +11,13 @@ def convert_pt_to_safetensors(pt_path, safe_path):
state_dict = model["model"]
else:
state_dict = model
state_dict = {k.replace("state_dict.", ""): v for k, v in state_dict.items()}
metadata = {"format": "pt"}
save_file(state_dict, safe_path, metadata)
with atomic_uel_writer(safe_path, metadata) as writer:
for key in list(state_dict.keys()):
tensor = state_dict.pop(key)
writer.write_batch(
[(key.replace("state_dict.", ""), tensor.contiguous())]
)
return True, ""
except Exception as e:
return False, str(e)
def load_metadata_from_safetensors(safetensors_file: str) -> dict:
try:
with safe_open(safetensors_file, framework="pt", device="cpu") as f:
model_weights = {}
for key in f.keys():
model_weights[key] = f.get_tensor(key)
metadata = f.metadata()
if metadata is None:
metadata = {}
return model_weights, metadata
except Exception as e:
print(f"Error loading safetensors file: {e}")
return {}, {}
+237
View File
@@ -0,0 +1,237 @@
# Per-layer parameter rules for merge, extraction, and resize
## 1. Summary and initial scope
Add one multiline configurator node that supplies optional per-layer numeric overrides to existing nodes. Matching layers receive the explicitly assigned values; unassigned parameters and unmatched layers retain the main node’s settings. Existing exclusions, inclusion filters, discard rules, mismatch handling, and safety guards remain authoritative.
The initial implementation covers these 39 nodes:
| Family | Included variants | Count |
|---|---|---:|
| Standard mergers | Two- and three-input checkpoint, diffusion-model, text-encoder, LoRA, and embedding mergers | 10 |
| Extraction | Fixed, Ratio, Quantile, Knee, and Frobenius for LoRA, DoRA, learned DoRA, text-encoder LoRA, and text-encoder DoRA | 25 |
| LoRA resize | Fixed, Ratio, Frobenius, and Cumulative | 4 |
Do not expand this implementation to CWB configuration, weighted multi-LoRA merging, LoRA application, Lodestone, optimizer settings, iteration counts, seeds, device selection, dtype selection, or new normalization behavior. These remain post-implementation considerations.
No changes are made during planning. After approval and execution is enabled, the first implementation action is to save this entire approved plan verbatim to `plans/per-layer-parameter-rules.md`. Commit that planning document before editing implementation files. Commit the completed feature after verification.
Preserve the pre-existing analysis/report edits currently outside the filter commit.
## 2. Public interface, syntax, and bindings
### Configurator and connection
Implement `LayerParameterConfiguration`, displayed as **Layer Parameter Configuration**, in a new `nodes/layer_parameters.py` module.
Its inputs are:
| Input | Type | Default | Meaning |
|---|---|---|---|
| `rules` | Multiline string | Empty | One pattern-and-assignment rule per line |
| `glob_patterns` | Boolean | False | False uses regex substring search; True uses glob substring matching |
Its outputs are:
| Output | Type | Meaning |
|---|---|---|
| `layer_parameters` | `io.Custom("MODELUTILS_LAYER_PARAMETERS")` | Immutable, parsed rules |
| `documentation` | String | Syntax, bindings, supported nodes, fallback behavior, errors, and examples |
This follows the existing `io.Custom(...)` configurator/socket mechanism already used by CWB. No frontend extension or dynamic socket-type changes are needed.
Append an optional `layer_parameters` input to each supported receiver. Append `layer_parameters=None` to its execution interface and forward it by keyword. Preserve every existing input identifier, widget order, output, and default.
An unconnected input or an empty rule set leaves the operation unchanged.
### Rule grammar
Accept the following forms:
```text
(blocks\.4[589]\.attn\.qkv_proj) a:0.5 b:0.25 c:0.75 d:1.0
(blocks\.4[589]\.attn\.qkv_proj) alpha:0.5, beta:0.25, gamma:0.75, delta:1.0
(blocks\.4[589]\.attn\.qkv_proj) a:64, b:32, c:0.99, d:0.0
(blocks\.4[589]\.attn\.qkv_proj) max_rank:128 min_rank:1 target:0.9
```
Rules:
- Require a parenthesized pattern followed by at least one `name:value` assignment.
- Accept whitespace, commas, or a mixture between assignments, including whitespace around colons.
- Preserve the pattern’s backslashes, internal whitespace, nested groups, and character classes.
- Accept signed decimal and scientific-notation numeric values. Reject nonfinite values.
- Ignore blank lines and whole-line comments whose first nonwhitespace character is `#`. Do not support inline comments.
- Reject empty patterns, missing values, trailing garbage, duplicate assignments, and invalid regex.
- Use `re.search` for regex matching. For glob matching, use case-sensitive `fnmatchcase` with surrounding wildcards to provide substring matching.
- Keep configurator pattern syntax independent of the main node’s existing exclusion-filter syntax.
Short aliases and full parameter names may be mixed. Two aliases assigning the same parameter on one line are an error, even when their values agree.
### Binding tables
The shared binding definitions must generate both receiver tooltips and configurator documentation, preventing the two from drifting.
#### Standard merge coefficients
Bindings remain stable across calculation modes:
| Short name | Full name |
|---|---|
| `a` | `alpha` |
| `b` | `beta` |
| `c` or `g` | `gamma` |
| `d` | `delta` |
| `e` | `epsilon` |
| `f` or `z` | `zeta` |
Accept only coefficients actually used by the selected calculation mode:
| Calculation mode | Supported coefficients |
|---|---|
| Weight-Sum; Train-Difference | alpha |
| Add-Difference; Power-Up (DARE) | alpha, beta |
| Comparative-Interpolation; Add-Dissimilarities; SVD LoRA Extraction; Enhanced Auto Interp; Weight-Sum Cutoff; Power-Up (DARE+TIES) | alpha, beta, gamma |
| Extract-Features; Enhanced Man Interp; Power-Up Enhanced (DARE) | alpha, beta, gamma, delta |
| Power-Up Enhanced (DARE+TIES) | beta, gamma, delta, epsilon, zeta |
Reject assignments to ignored coefficients rather than pretend they affect the selected mode. Seeds remain main-node settings.
#### Extraction
Use the same bindings across LoRA, DoRA, learned DoRA, and both text-encoder extraction families:
| Method | `a` | `b` | `c` | `d` | `e` | `f` |
|---|---|---|---|---|---|---|
| Fixed | linear_dim | conv_dim | clamp_quantile | min_diff | — | — |
| Ratio | linear_ratio | conv_ratio | clamp_quantile | min_diff | linear_max_rank | conv_max_rank |
| Quantile | linear_quantile | conv_quantile | clamp_quantile | min_diff | linear_max_rank | conv_max_rank |
| Frobenius | linear_target | conv_target | clamp_quantile | min_diff | linear_max_rank | conv_max_rank |
| Knee | linear_max_rank | conv_max_rank | clamp_quantile | min_diff | — | — |
Knee method selection, probe offsets, SVD iterations, chunking controls, and learned-DoRA optimization settings remain on the receiving node.
#### LoRA resize
| Method | `a` | `b` | `c` |
|---|---|---|---|
| Fixed | new_rank | — | — |
| Ratio | max_rank | ratio | — |
| Frobenius | max_rank | min_rank | target |
| Cumulative | max_rank | target | — |
### Value validation
Inherit values from the receiving node—not hardcoded configurator defaults.
For supplied overrides:
- Extraction dimensions and rank caps are integral values in `[1, 16384]`; ratios are `[1, 100]`; quantile/Frobenius targets are `[0, 1]`; `clamp_quantile` is `[0.5, 1]`; `min_diff` is `[0, 1]`.
- Resize ranks are integral values in `[1, 3072]`; ratios are `[1, 100]`; targets are `[0.1, 1]`.
- Validate resolved Frobenius resize `min_rank <= max_rank`, including values inherited from the main node.
- Standard merge coefficients retain existing widget bounds. Where a coefficient represents a quantile or probability, validate `[0, 1]`; generic SVD rank coefficients must also be positive integral values.
- Do not silently truncate fractional ranks or repair invalid assignments. Retain existing mathematical rank caps imposed by tensor dimensions.
## 3. Matching and implementation data flow
### Matching contract
For ordinary model merging and extraction, match the receiving operation’s target tensor names.
For LoRA inputs, match the normalized logical layer name after the existing parsing and normalization path. Do not match factor or alpha suffixes independently.
Use the established dotted naming produced by extraction and existing reference mapping. Preserve internal identifiers such as `qkv_proj`. Do not introduce underscore normalization as the new public convention, rewrite regex punctuation, or redesign model-format conversion.
Retain the existing unresolved flattened-name fallback when no reference mapping exists. Documentation must distinguish this fallback from the normal dotted path. A dotted rule that cannot match the resulting normalized input names produces the ordinary unmatched-rule error; do not silently reinterpret it.
A single rule matching multiple keys belonging to one logical LoRA layer counts as one match. That layer receives one parameter set shared by its factors.
### Strict errors and fallback
Validate all rules before opening the output writer or beginning tensor streaming:
1. Parse syntax and compile patterns once in the configurator.
2. At each receiver, resolve aliases and reject unsupported parameters.
3. Use header/key information to identify target layers and their existing normalized names.
4. Detect matching rules per logical layer.
5. Reject multiple matching rule lines for any one layer—even if their assignments concern different parameters.
6. Reject any rule line matching no target layer.
7. Combine the single matching rule with main-node values and validate the resulting parameters.
Errors identify the receiving node, configuration line number, pattern, and affected layer or parameter. Overlap errors include every conflicting line number.
Determine name matches before existing exclusions so an explicitly excluded layer is not falsely reported as nonexistent. Overrides never reactivate an excluded, discarded, guarded, or otherwise unprocessable layer.
### Shared implementation
Keep parsing, immutable rule records, binding definitions, and matching/error handling in the new shared module.
Store only lightweight per-layer assignments and resolved scalar parameters. Do not store tensors, mutate the shared configuration payload, or modify global/main-node parameter dictionaries between layers.
Register the configurator through the existing extension node list. Do not add dependencies.
### Standard merger integration
In `MergerLogic.execute_merge`:
- Prepare the rule-to-layer mapping alongside existing header-based work planning.
- For LoRA mode, use parsed logical groups to assign the same override to their factor keys; retain existing merge and alpha-normalization behavior.
- Immediately before `calc_mode_class.create_recipe`, construct a layer-local parameter dictionary containing the base settings plus that layer’s overrides.
- Keep tensor references and preloaded-source data local to the current work unit.
- Do not mutate `recipe_params` with overrides that could affect later layers.
Do not change recipe formulas, calculation-mode selection, or source-loading behavior.
### Extraction integration
Thread the optional rules through all 25 node entry points into the four extraction backends.
Prepare resolved parameters before streaming, then use them in each backend’s `_process_layer` before `min_diff` testing and SVD selection.
Specific requirements:
- Fixed-rank overrides update both the requested dimension and corresponding maximum-rank argument; otherwise the original cap would negate an increased rank.
- Adaptive overrides update the existing linear/conv method parameter and rank cap without altering method selection or probe settings.
- Normal execution, CPU retry, and chunked fallback use the same resolved layer values.
- Extend the existing chunked extraction helpers with a trailing `clamp_quantile` argument and pass the effective value through. These helpers currently omit it when calling `_svd_extract_linear`; text-encoder extraction shares the SVD helper.
- Perform configuration validation outside broad extraction exception handlers so a configuration error cannot become a silently skipped layer.
Do not alter SVD, chunk recombination, learned optimization, or tensor-lifecycle algorithms.
### Resize integration
Attach resolved rank/target values to the existing logical-layer work units.
When calling `_resize_lora_factors`, translate the public bindings into the existing arguments:
- Fixed `new_rank` → `new_rank`.
- Adaptive `max_rank` → `new_rank`.
- Ratio/target → `dynamic_param`.
- Frobenius `min_rank` → `min_rank`.
Leave `dynamic_method` on the main node. Preserve whole-layer factor handling, alpha behavior, passthrough data, low-bit handling, dtype selection, and CPU fallback.
## 4. Verification and delivery
Extend the repository-owned test selector with focused rule-parser and integration coverage. Map the new production module explicitly; do not rely solely on the broad repository-contract group.
Required tests:
- Both delimiter styles, mixed short/full names, scientific notation, nested regex groups, CRLF, blank lines, and comments.
- Invalid syntax, invalid regex, unsupported parameters, duplicate aliases, nonfinite values, fractional ranks, and invalid resolved rank bounds.
- Overlapping rules fail with layer and line numbers; unmatched rules fail before streaming or writing.
- Partial assignments inherit remaining main-node values. Unmatched layers retain main-node behavior.
- Existing include/exclude/discard and safety paths remain authoritative.
- LoRA matching uses normalized logical names, preserves internal underscores, and assigns one configuration to paired factors.
- All 39 receiver schemas append the optional socket without moving existing controls; every execution wrapper forwards it correctly.
- Standard merge outputs reflect different coefficients on different layers, without leaking overrides to later layers.
- Every extraction family uses the requested per-layer rank, clamp, and `min_diff`; cover adaptive methods, CPU retry, and forced chunk fallback.
- All four resize variants use their resolved rank/target values, including inherited `min_rank` validation.
- Absent or empty configuration preserves existing behavior.
- Configuration failures preserve any existing destination file.
Use small deterministic tensors and existing I/O fixtures. Do not run model inference or add unrelated performance, architecture, or normalization test projects.
Run focused groups during development, `--changed` after accumulated changes, and the required full gate before the feature commit. Run changed-file lint, compilation, and `git diff --check`; distinguish pre-existing lint findings from new ones.
Deliver the saved approved plan, implementation, generated binding documentation/tooltips, regression tests, and a verified feature commit. Leave unrelated analysis/report edits untouched.
+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.7"
version = "0.9.7"
license = { file = "LICENSE" }
classifiers = [
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
[project.urls]
Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils"
# Used by Comfy Registry https://comfyregistry.org
+5
View File
@@ -0,0 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader>=0.5.2
av
+259
View File
@@ -0,0 +1,259 @@
"""Queue the CWB six-LoRA API workflow and print its merge report."""
from __future__ import annotations
import argparse
import csv
import json
import re
import sys
import time
from pathlib import Path
from urllib import error, request
DEFAULT_WORKFLOW = Path(__file__).with_name("CWB6LORATEST_API.json")
MERGER_CLASS = "CWBLoRAMultiMerger"
CUSTOM_CLASS = "CWBCustomConfiguration"
PREVIEW_CLASS = "PreviewAny"
CALIBRATION_EXCLUDE_PATTERN = (
r"^diffusion_model\.blocks\.(?:[12]|[4-9]|1[0-3]|1[5-9]|2[0-3]|2[56])\."
"\n"
r"^diffusion_model\.txtfusion\.layerwise_blocks\.1\."
"\n"
r"^diffusion_model\.txtfusion\.refiner_blocks\.0\."
)
def parse_counterfactual_weight_sweep(report: str) -> list[dict]:
pattern = re.compile(
r"^(?P<consensus_type>mean|median) \| "
r"(?P<similarity_threshold>[-+0-9.e]+) \| "
r"(?P<power_alpha>[-+0-9.e]+) \| (?P<diversity_beta>[-+0-9.e]+) \| "
r"(?P<dsc>true|false) \| (?P<softcb>true|false) \| "
r"(?P<accepted>\d+)/(?P<contributors>\d+) \| (?P<fallbacks>\d+) \| "
r"(?P<dominant_mean>[-+0-9.e]+)/(?P<dominant_max>[-+0-9.e]+) \| "
r"(?P<effective_mean>[-+0-9.e]+)/(?P<effective_min>[-+0-9.e]+)$",
re.MULTILINE,
)
rows = []
for match in pattern.finditer(report):
values = match.groupdict()
rows.append({
"consensus_type": values["consensus_type"],
"similarity_threshold": float(values["similarity_threshold"]),
"power_alpha": float(values["power_alpha"]),
"diversity_beta": float(values["diversity_beta"]),
"dynamic_similarity_contrast": values["dsc"] == "true",
"soft_comfort_bandpass": values["softcb"] == "true",
"accepted_contributors": int(values["accepted"]),
"contributors": int(values["contributors"]),
"fallbacks": int(values["fallbacks"]),
"dominant_weight_mean": float(values["dominant_mean"]),
"dominant_weight_max": float(values["dominant_max"]),
"effective_contributors_mean": float(values["effective_mean"]),
"effective_contributors_min": float(values["effective_min"]),
})
return rows
def write_counterfactual_csv(path: Path, rows: list[dict]) -> None:
if not rows:
raise RuntimeError("CWB report does not contain a counterfactual weight sweep")
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", newline="", encoding="utf-8") as handle:
writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
writer.writeheader()
writer.writerows(rows)
def _json_request(url: str, payload=None):
data = None if payload is None else json.dumps(payload).encode("utf-8")
req = request.Request(url, data=data)
if data is not None:
req.add_header("Content-Type", "application/json")
try:
with request.urlopen(req) as response:
return json.load(response)
except error.HTTPError as exc:
detail = exc.read().decode("utf-8", errors="replace")
raise RuntimeError(f"ComfyUI request failed ({exc.code}): {detail}") from exc
except error.URLError as exc:
raise RuntimeError(f"Could not reach ComfyUI at {url}: {exc.reason}") from exc
def _single_node(workflow: dict, class_type: str) -> tuple[str, dict]:
matches = [
(node_id, node)
for node_id, node in workflow.items()
if node.get("class_type") == class_type
]
if len(matches) != 1:
raise RuntimeError(
f"Expected exactly one {class_type} node, found {len(matches)}."
)
return matches[0]
def _linked_preview_nodes(workflow: dict, merger_id: str) -> dict[int, str]:
linked = {}
for node_id, node in workflow.items():
if node.get("class_type") != PREVIEW_CLASS:
continue
source = node.get("inputs", {}).get("source")
if (
isinstance(source, list)
and len(source) == 2
and str(source[0]) == merger_id
and isinstance(source[1], int)
):
linked[source[1]] = node_id
missing = {0, 2} - linked.keys()
if missing:
raise RuntimeError(
"Workflow must connect merger output 0 (filename) and output 2 "
"(CWB report) to PreviewAny nodes."
)
return linked
def _prepare_workflow(
workflow: dict,
preset: str | None,
output_filename: str | None,
full_blocks: bool = False,
) -> tuple[str, dict[int, str]]:
merger_id, merger = _single_node(workflow, MERGER_CLASS)
_single_node(workflow, CUSTOM_CLASS)
inputs = merger.setdefault("inputs", {})
if preset is not None:
inputs["cwb_preset"] = preset
inputs.pop("cwb_config", None)
if output_filename is not None:
inputs["output_filename"] = output_filename
if not full_blocks:
existing = inputs.get("exclude_patterns", "").strip()
inputs["exclude_patterns"] = "\n".join(
value for value in (existing, CALIBRATION_EXCLUDE_PATTERN) if value
)
return merger_id, _linked_preview_nodes(workflow, merger_id)
def _prompt_payload(workflow: dict) -> dict:
return {"prompt": workflow}
def _wait_for_history(server: str, prompt_id: str, poll_seconds: float):
while True:
history = _json_request(f"{server}/history/{prompt_id}")
if prompt_id in history:
return history[prompt_id]
time.sleep(poll_seconds)
def _text_output(history_outputs: dict, node_id: str) -> str:
output = history_outputs.get(node_id, {})
candidates = (
output.get("text"),
output.get("ui", {}).get("text") if isinstance(output.get("ui"), dict) else None,
)
for candidate in candidates:
if isinstance(candidate, list) and candidate:
return str(candidate[0])
if isinstance(candidate, str):
return candidate
raise RuntimeError(f"PreviewAny node {node_id} did not return text in API history.")
def main() -> int:
parser = argparse.ArgumentParser(
description="Run the CWB six-LoRA API workflow and collect its report."
)
parser.add_argument("--workflow", type=Path, default=DEFAULT_WORKFLOW)
parser.add_argument("--server", default="http://127.0.0.1:8188")
parser.add_argument("--poll-seconds", type=float, default=2.0)
parser.add_argument(
"--preset",
help="Disconnect the custom configuration and run this CWB preset.",
)
parser.add_argument(
"--output-filename",
help="Override the merger output filename without its extension.",
)
parser.add_argument(
"--result-json",
type=Path,
help="Optionally write prompt ID, filename, and report as JSON.",
)
parser.add_argument(
"--full-blocks",
action="store_true",
help="Process all blocks for final verification instead of the calibration subset.",
)
parser.add_argument("--summary-only", action="store_true")
parser.add_argument("--counterfactual-csv", type=Path)
parser.add_argument("--alignment-threshold", type=float)
parser.add_argument("--alignment-method", choices=["index", "similarity"])
args = parser.parse_args()
if args.poll_seconds <= 0:
parser.error("--poll-seconds must be greater than zero")
if not args.workflow.is_file():
parser.error(f"workflow not found: {args.workflow}")
workflow = json.loads(args.workflow.read_text(encoding="utf-8"))
if args.alignment_threshold is not None:
_, custom = _single_node(workflow, CUSTOM_CLASS)
custom.setdefault("inputs", {})["alignment_threshold"] = (
args.alignment_threshold
)
if args.alignment_method is not None:
_, custom = _single_node(workflow, CUSTOM_CLASS)
custom.setdefault("inputs", {})["alignment_method"] = args.alignment_method
_, linked_outputs = _prepare_workflow(
workflow,
args.preset,
args.output_filename,
args.full_blocks,
)
server = args.server.rstrip("/")
queued = _json_request(f"{server}/prompt", _prompt_payload(workflow))
prompt_id = queued["prompt_id"]
sys.stdout.write(f"Queued prompt {prompt_id}\n")
history = _wait_for_history(server, prompt_id, args.poll_seconds)
status = history.get("status", {})
if status.get("status_str") == "error":
sys.stderr.write(json.dumps(status, indent=2) + "\n")
return 1
outputs = history.get("outputs", {})
output_filename = _text_output(outputs, linked_outputs[0])
cwb_report = _text_output(outputs, linked_outputs[2])
result = {
"prompt_id": prompt_id,
"output_filename": output_filename,
"cwb_report": cwb_report,
}
if args.result_json is not None:
args.result_json.parent.mkdir(parents=True, exist_ok=True)
args.result_json.write_text(
json.dumps(result, indent=2) + "\n",
encoding="utf-8",
)
if args.counterfactual_csv is not None:
write_counterfactual_csv(
args.counterfactual_csv,
parse_counterfactual_weight_sweep(cwb_report),
)
sys.stdout.write(f"Output: {output_filename}\n")
if not args.summary_only:
sys.stdout.write(f"\n{cwb_report}\n")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+418
View File
@@ -0,0 +1,418 @@
"""Run a resumable one-factor CWB LoRA calibration sweep."""
from __future__ import annotations
import argparse
import csv
import json
import os
import re
import sys
import time
from copy import deepcopy
from pathlib import Path
import run_cwb_api_workflow as api
import run_model_analysis_api_workflow as analysis_api
CONFIG_KEYS = (
"consensus_type",
"alignment_method",
"alignment_threshold",
"similarity_threshold",
"power_alpha",
"diversity_beta",
"rescale_norm",
"global_scale",
"dynamic_similarity_contrast",
"soft_comfort_bandpass",
"position_weight",
"preserve_common_prefix",
)
INTEGER_METRICS = {
"alignment_candidates": "Alignment candidates",
"alignment_matches": "Alignment matches",
"anchor_only_groups": "Anchor-only groups",
"weighted_vector_groups": "Weighted vector groups",
"all_rejected_fallbacks": "All-rejected equal fallbacks",
"zero_weight_fallbacks": "Zero-weight equal fallbacks",
"cuda_oom_fallbacks": "CUDA OOM CPU fallbacks",
"lora_groups": "LoRA groups",
}
TRIPLE_METRICS = {
"consensus_similarity": "Consensus similarity min/mean/max",
"dominant_weight": "Dominant normalized weight min/mean/max",
"effective_contributors": "Effective contributors min/mean/max",
}
DELTA_METRICS = (
"alignment_match_fraction",
"anchor_only_groups",
"accepted_weight_fraction",
"consensus_similarity_mean",
"dominant_weight_mean",
"dominant_weight_max",
"effective_contributors_mean",
"effective_contributors_min",
"norm_ratio_mean",
"norm_ratio_min",
"norm_ratio_max",
"all_rejected_fallbacks",
"zero_weight_fallbacks",
"wall_seconds",
"analysis_relative_l2",
"analysis_cosine",
"analysis_exact",
"analysis_norm_ratio",
"analysis_pair_cosine",
"analysis_mean_affinity_a",
"analysis_mean_affinity_b",
"analysis_median_affinity_a",
"analysis_median_affinity_b",
)
def _single_custom_node(workflow: dict) -> tuple[str, dict]:
return api._single_node(workflow, api.CUSTOM_CLASS)
def _baseline_config(workflow: dict) -> dict:
_, node = _single_custom_node(workflow)
inputs = node.get("inputs", {})
missing = [key for key in CONFIG_KEYS if key not in inputs]
if missing:
raise RuntimeError(
"Custom configuration is missing required inputs: " + ", ".join(missing)
)
return {key: inputs[key] for key in CONFIG_KEYS}
def build_sweep_cases(baseline: dict) -> list[dict]:
variations = [
("alignment_method", "index"),
*[("alignment_threshold", value) for value in (0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01)],
*[("position_weight", value) for value in (0.0, 0.0001, 0.0005, 0.001, 0.005, 0.01)],
("consensus_type", "mean"),
*[("similarity_threshold", value) for value in (-0.25, 0.1, 0.25, 0.5, 0.75)],
*[("power_alpha", value) for value in (0.5, 1.0, 3.0, 4.0)],
*[("diversity_beta", value) for value in (0.0, 1.0, 2.0, 4.0, 7.0)],
("rescale_norm", False),
("dynamic_similarity_contrast", True),
("soft_comfort_bandpass", False),
]
cases = [{"id": "baseline", "factor": "baseline", "value": None, "config": baseline.copy()}]
for factor, value in variations:
if baseline[factor] == value:
continue
config = baseline.copy()
config[factor] = value
value_text = str(value).lower().replace(".", "p").replace("-", "neg")
cases.append({
"id": f"{factor}_{value_text}",
"factor": factor,
"value": value,
"config": config,
})
return cases
def parse_cwb_report(report: str) -> dict:
metrics = {}
for metric, label in INTEGER_METRICS.items():
match = re.search(rf"^{re.escape(label)}: (\d+)$", report, re.MULTILINE)
if match:
metrics[metric] = int(match.group(1))
accepted = re.search(
r"^Accepted weighting contributors: (\d+)/(\d+)$",
report,
re.MULTILINE,
)
if not accepted:
raise RuntimeError(
"CWB report lacks contributor-weight diagnostics. Restart ComfyUI "
"after loading the current consensus_merger.py before running the sweep."
)
metrics["accepted_weighting_contributors"] = int(accepted.group(1))
metrics["weighting_contributors"] = int(accepted.group(2))
number = r"(?:n/a|[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[-+]?\d+)?)"
for prefix, label in TRIPLE_METRICS.items():
match = re.search(
rf"^{re.escape(label)}: ({number}) / ({number}) / ({number})$",
report,
re.MULTILINE | re.IGNORECASE,
)
if not match:
raise RuntimeError(f"CWB report is missing metric: {label}")
for suffix, value in zip(("min", "mean", "max"), match.groups()):
metrics[f"{prefix}_{suffix}"] = None if value == "n/a" else float(value)
norm_groups = re.findall(
rf"^Rank-one delta norm ratio count/min/mean/max: (\d+) / "
rf"({number}) / ({number}) / ({number})$",
report,
re.MULTILINE | re.IGNORECASE,
)
numeric_norms = [
(int(count), *(float(value) for value in (minimum, mean, maximum)))
for count, minimum, mean, maximum in norm_groups
if count != "0" and "n/a" not in (minimum, mean, maximum)
]
if numeric_norms:
norm_count = sum(group[0] for group in numeric_norms)
metrics["norm_ratio_count"] = norm_count
metrics["norm_ratio_min"] = min(group[1] for group in numeric_norms)
metrics["norm_ratio_mean"] = (
sum(group[0] * group[2] for group in numeric_norms) / norm_count
)
metrics["norm_ratio_max"] = max(group[3] for group in numeric_norms)
candidates = metrics.get("alignment_candidates", 0)
contributors = metrics["weighting_contributors"]
metrics["alignment_match_fraction"] = (
metrics.get("alignment_matches", 0) / candidates if candidates else 0.0
)
metrics["accepted_weight_fraction"] = (
metrics["accepted_weighting_contributors"] / contributors
if contributors else 0.0
)
return metrics
def parse_analysis_reports(comparison: str, cwb_report: str) -> dict:
global_section = comparison.split("PARAMETER-WEIGHTED GLOBAL METRICS", 1)[1].split(
"EQUAL-LAYER AVERAGES", 1
)[0]
def single(label: str, text: str) -> float:
match = re.search(rf"^{re.escape(label)}: ([-+0-9.e]+)$", text, re.MULTILINE)
if not match:
raise RuntimeError(f"Analysis report is missing metric: {label}")
return float(match.group(1))
metrics = {
"analysis_relative_l2": single("relative_l2", global_section),
"analysis_cosine": single("cosine", global_section),
"analysis_exact": single("exact", global_section),
"analysis_norm_ratio": single("norm_ratio", global_section),
}
pairwise = re.search(
r"^Pairwise row cosine mean/min/max: ([-+0-9.e]+) /",
cwb_report,
re.MULTILINE,
)
mean_affinity = re.search(
r"^Mean-consensus affinity A/B: ([-+0-9.e]+) / ([-+0-9.e]+)$",
cwb_report,
re.MULTILINE,
)
median_affinity = re.search(
r"^Median-consensus affinity A/B: ([-+0-9.e]+) / ([-+0-9.e]+)$",
cwb_report,
re.MULTILINE,
)
if not pairwise or not mean_affinity or not median_affinity:
raise RuntimeError("CWB analysis report is missing global affinity metrics")
metrics.update({
"analysis_pair_cosine": float(pairwise.group(1)),
"analysis_mean_affinity_a": float(mean_affinity.group(1)),
"analysis_mean_affinity_b": float(mean_affinity.group(2)),
"analysis_median_affinity_a": float(median_affinity.group(1)),
"analysis_median_affinity_b": float(median_affinity.group(2)),
})
return metrics
def _run_case(
workflow_template: dict,
case: dict,
server: str,
poll_seconds: float,
full_blocks: bool,
analysis_workflow: dict | None = None,
) -> dict:
workflow = deepcopy(workflow_template)
_, custom = _single_custom_node(workflow)
custom["inputs"].update(case["config"])
_, linked = api._prepare_workflow(
workflow,
preset=None,
output_filename=None,
full_blocks=full_blocks,
)
started = time.perf_counter()
queued = api._json_request(f"{server}/prompt", api._prompt_payload(workflow))
prompt_id = queued["prompt_id"]
history = api._wait_for_history(server, prompt_id, poll_seconds)
wall_seconds = time.perf_counter() - started
status = history.get("status", {})
if status.get("status_str") == "error":
raise RuntimeError(
f"Sweep case {case['id']} failed: {json.dumps(status, indent=2)}"
)
outputs = history.get("outputs", {})
output_filename = api._text_output(outputs, linked[0])
report = api._text_output(outputs, linked[2])
metrics = parse_cwb_report(report)
metrics["wall_seconds"] = wall_seconds
result = {
**case,
"prompt_id": prompt_id,
"output_filename": output_filename,
"metrics": metrics,
"cwb_report": report,
}
if analysis_workflow is not None:
analysis_result = analysis_api.run_analysis(
deepcopy(analysis_workflow), server, poll_seconds
)
metrics.update(parse_analysis_reports(
analysis_result["comparison_report"],
analysis_result["cwb_report"],
))
result["analysis"] = analysis_result
return result
def _load_completed(path: Path) -> list[dict]:
if not path.exists():
return []
completed = []
with path.open("r", encoding="utf-8") as handle:
for line_number, line in enumerate(handle, start=1):
try:
completed.append(json.loads(line))
except json.JSONDecodeError as exc:
raise RuntimeError(
f"Invalid sweep JSONL at line {line_number}: {exc}"
) from exc
return completed
def _append_result(path: Path, result: dict) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(result, ensure_ascii=False) + "\n")
handle.flush()
os.fsync(handle.fileno())
def _write_influence_csv(path: Path, results: list[dict]) -> None:
baseline = next((result for result in results if result["id"] == "baseline"), None)
if baseline is None:
return
fields = ["id", "factor", "value", *DELTA_METRICS]
delta_fields = [f"delta_{metric}" for metric in DELTA_METRICS]
with path.open("w", newline="", encoding="utf-8") as handle:
writer = csv.DictWriter(handle, fieldnames=[*fields, *delta_fields])
writer.writeheader()
base_metrics = baseline["metrics"]
for result in results:
row = {key: result.get(key) for key in ("id", "factor", "value")}
for metric in DELTA_METRICS:
value = result["metrics"].get(metric)
base = base_metrics.get(metric)
row[metric] = value
row[f"delta_{metric}"] = (
value - base if value is not None and base is not None else None
)
writer.writerow(row)
def main() -> int:
parser = argparse.ArgumentParser(
description="Run a resumable one-factor CWB LoRA parameter sweep."
)
parser.add_argument("--workflow", type=Path, default=api.DEFAULT_WORKFLOW)
parser.add_argument("--server", default="http://127.0.0.1:8188")
parser.add_argument("--poll-seconds", type=float, default=2.0)
parser.add_argument("--output-dir", type=Path, default=Path("audit/cwb_sweep"))
parser.add_argument("--full-blocks", action="store_true")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--limit", type=int)
parser.add_argument(
"--case",
action="append",
dest="case_ids",
help="Run only the named case. Repeat to select multiple cases.",
)
parser.add_argument(
"--analysis-workflow",
type=Path,
default=analysis_api.DEFAULT_WORKFLOW,
)
parser.add_argument(
"--analysis-every",
type=int,
default=5,
help="Run candidate-versus-anchor analysis for baseline and every Nth case.",
)
args = parser.parse_args()
workflow = json.loads(args.workflow.read_text(encoding="utf-8"))
analysis_workflow = json.loads(args.analysis_workflow.read_text(encoding="utf-8"))
if args.full_blocks:
_, analysis_node = api._single_node(
analysis_workflow,
analysis_api.ANALYSIS_CLASS,
)
analysis_node.setdefault("inputs", {})["exclude_patterns"] = ""
cases = build_sweep_cases(_baseline_config(workflow))
if args.case_ids:
known_ids = {case["id"] for case in cases}
unknown = [case_id for case_id in args.case_ids if case_id not in known_ids]
if unknown:
raise RuntimeError("Unknown sweep case(s): " + ", ".join(unknown))
selected = set(args.case_ids)
cases = [case for case in cases if case["id"] in selected]
if args.limit is not None:
cases = cases[:args.limit]
if args.dry_run:
for case in cases:
sys.stdout.write(f"{case['id']}: {case['factor']}={case['value']}\n")
sys.stdout.write(f"Cases: {len(cases)}\n")
return 0
results_path = args.output_dir / "runs.jsonl"
completed = _load_completed(results_path)
completed_ids = {result["id"] for result in completed}
server = args.server.rstrip("/")
if args.analysis_every < 1:
raise RuntimeError("--analysis-every must be at least 1")
for index, case in enumerate(cases, start=1):
if case["id"] in completed_ids:
sys.stdout.write(f"[{index}/{len(cases)}] skip {case['id']}\n")
continue
sys.stdout.write(f"[{index}/{len(cases)}] run {case['id']}\n")
result = _run_case(
workflow,
case,
server,
args.poll_seconds,
args.full_blocks,
analysis_workflow=(
analysis_workflow
if index == 1 or index % args.analysis_every == 0
else None
),
)
_append_result(results_path, result)
completed.append(result)
completed_ids.add(case["id"])
sys.stdout.write(
f" {result['metrics']['wall_seconds']:.1f}s, "
f"match={result['metrics']['alignment_match_fraction']:.4f}, "
f"effective={result['metrics']['effective_contributors_mean']:.4f}\n"
)
_write_influence_csv(args.output_dir / "influence.csv", completed)
_write_influence_csv(args.output_dir / "influence.csv", completed)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+73
View File
@@ -0,0 +1,73 @@
"""Queue a model-analysis API workflow and capture both report outputs."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import run_cwb_api_workflow as api
DEFAULT_WORKFLOW = Path("ANALYZETESTCANDIDATE_API.json")
ANALYSIS_CLASS = "LoRAModelAnalysis"
def _linked_previews(workflow: dict, analysis_id: str) -> dict[int, str]:
linked = {}
for node_id, node in workflow.items():
if node.get("class_type") != "PreviewAny":
continue
source = node.get("inputs", {}).get("source")
if isinstance(source, list) and len(source) == 2 and str(source[0]) == analysis_id:
linked[int(source[1])] = node_id
missing = [index for index in (0, 1) if index not in linked]
if missing:
raise RuntimeError(f"Missing PreviewAny capture for analysis output(s): {missing}")
return linked
def run_analysis(workflow: dict, server: str, poll_seconds: float) -> dict:
analysis_id, _ = api._single_node(workflow, ANALYSIS_CLASS)
linked = _linked_previews(workflow, analysis_id)
queued = api._json_request(
f"{server.rstrip('/')}/prompt",
api._prompt_payload(workflow),
)
prompt_id = queued["prompt_id"]
history = api._wait_for_history(server.rstrip("/"), prompt_id, poll_seconds)
status = history.get("status", {})
if status.get("status_str") == "error":
raise RuntimeError(f"Analysis failed: {json.dumps(status, indent=2)}")
outputs = history.get("outputs", {})
return {
"prompt_id": prompt_id,
"comparison_report": api._text_output(outputs, linked[0]),
"cwb_report": api._text_output(outputs, linked[1]),
}
def main() -> int:
parser = argparse.ArgumentParser(description="Run a model-analysis API workflow.")
parser.add_argument("--workflow", type=Path, default=DEFAULT_WORKFLOW)
parser.add_argument("--server", default="http://127.0.0.1:8188")
parser.add_argument("--poll-seconds", type=float, default=2.0)
parser.add_argument("--result-json", type=Path)
args = parser.parse_args()
workflow = json.loads(args.workflow.read_text(encoding="utf-8"))
result = run_analysis(workflow, args.server, args.poll_seconds)
if args.result_json:
args.result_json.parent.mkdir(parents=True, exist_ok=True)
args.result_json.write_text(
json.dumps(result, indent=2, ensure_ascii=False) + "\n",
encoding="utf-8",
)
sys.stdout.write(f"Prompt: {result['prompt_id']}\n")
sys.stdout.write("Analysis reports captured\n")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+8
View File
@@ -0,0 +1,8 @@
def pytest_report_teststatus(report, config):
"""Pytest hook: emit only 'F' for failed tests and nothing for passing/skipped tests."""
if report.when in ("setup", "call") and report.failed:
return "failed", "F", "FAILED"
if report.passed:
return "passed", "", ""
if report.skipped:
return "skipped", "", ""
+249
View File
@@ -0,0 +1,249 @@
"""Select repository tests from changed paths or named subsystem groups."""
# ruff: noqa: T201
from __future__ import annotations
import argparse
import fnmatch
import os
import shutil
import subprocess
import sys
import tomllib
import uuid
from dataclasses import dataclass, field
from pathlib import Path
REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
COMFYUI_ROOT = REPOSITORY_ROOT.parent.parent
MANIFEST_PATH = Path(__file__).with_name("test_groups.toml")
TEMP_ROOT = REPOSITORY_ROOT / ".pytest-tmp"
@dataclass(frozen=True)
class TestGroup:
name: str
paths: tuple[str, ...]
python_tests: tuple[str, ...]
@dataclass
class Selection:
groups: set[str] = field(default_factory=set)
python_tests: set[str] = field(default_factory=set)
reasons: dict[str, set[str]] = field(default_factory=dict)
unmapped: set[str] = field(default_factory=set)
def _parse_groups(text: str) -> dict[str, TestGroup]:
data = tomllib.loads(text)
return {
name: TestGroup(
name=name,
paths=tuple(values.get("paths", ())),
python_tests=tuple(values.get("python_tests", ())),
)
for name, values in data.get("groups", {}).items()
}
def load_groups(path: Path = MANIFEST_PATH) -> dict[str, TestGroup]:
return _parse_groups(path.read_text(encoding="utf-8"))
def load_groups_from_revision(revision: str) -> dict[str, TestGroup]:
manifest = MANIFEST_PATH.relative_to(REPOSITORY_ROOT).as_posix()
result = subprocess.run(
["git", "show", f"{revision}:{manifest}"], cwd=REPOSITORY_ROOT,
check=False, capture_output=True, text=True,
)
return _parse_groups(result.stdout) if result.returncode == 0 else {}
def git_lines(*args: str) -> set[str]:
result = subprocess.run(
["git", *args], cwd=REPOSITORY_ROOT, check=True,
capture_output=True, text=True,
)
return {
line.strip().replace("\\", "/")
for line in result.stdout.splitlines() if line.strip()
}
def changed_paths(base: str | None = None) -> set[str]:
paths = git_lines("diff", "--name-only", "--relative", "HEAD")
if base:
paths.update(git_lines("diff", "--name-only", "--relative", f"{base}...HEAD"))
return paths
def is_production_source(path: str) -> bool:
normalized = path.replace("\\", "/")
return normalized == "__init__.py" or (
normalized.startswith("nodes/") and normalized.endswith(".py")
)
def _add_group(selection: Selection, group: TestGroup, reason: str) -> None:
selection.groups.add(group.name)
selection.python_tests.update(group.python_tests)
selection.reasons.setdefault(group.name, set()).add(reason)
def _path_matches_group(path: str, group: TestGroup) -> bool:
return any(fnmatch.fnmatchcase(path, pattern) for pattern in group.paths)
def _existing_group(group: TestGroup) -> TestGroup:
return TestGroup(
name=group.name, paths=group.paths,
python_tests=tuple(
path for path in group.python_tests if (REPOSITORY_ROOT / path).is_file()
),
)
def select_tests(
paths: set[str], groups: dict[str, TestGroup],
explicit_groups: tuple[str, ...] = (),
historical_groups: tuple[dict[str, TestGroup], ...] = (),
) -> Selection:
selection = Selection()
for name in explicit_groups:
if name not in groups:
raise ValueError(f"Unknown test group: {name}")
_add_group(selection, groups[name], "explicit selection")
for raw_path in sorted(paths):
path = raw_path.replace("\\", "/")
if path.startswith("tests/test_") and path.endswith(".py"):
if (REPOSITORY_ROOT / path).is_file():
selection.python_tests.add(path)
selection.reasons.setdefault("direct test", set()).add(path)
continue
matched = False
for group in groups.values():
if _path_matches_group(path, group):
_add_group(selection, group, path)
matched = True
if not matched and not (REPOSITORY_ROOT / path).exists():
for historical in historical_groups:
for historical_group in historical.values():
if _path_matches_group(path, historical_group):
active = groups.get(historical_group.name)
_add_group(
selection, active or _existing_group(historical_group),
f"{path} (deleted)",
)
matched = True
if not matched and is_production_source(path) and (REPOSITORY_ROOT / path).exists():
selection.unmapped.add(path)
return selection
def tracked_final_tests() -> set[str]:
tracked = git_lines("ls-files", "tests/test_*.py")
return {path for path in tracked if (REPOSITORY_ROOT / path).is_file()}
def print_selection(selection: Selection) -> None:
for group in sorted(selection.groups):
reasons = ", ".join(sorted(selection.reasons.get(group, ())))
print(f"group {group}: {reasons}")
for reason in sorted(selection.reasons.get("direct test", ())):
print(f"direct test: {reason}")
print("python tests:")
for path in sorted(selection.python_tests):
print(f" {path}")
def run_selection(selection: Selection) -> int:
TEMP_ROOT.mkdir(exist_ok=True)
basetemp = (TEMP_ROOT / f"run-{uuid.uuid4().hex}").resolve()
if TEMP_ROOT.resolve() not in basetemp.parents:
raise RuntimeError(f"Unsafe pytest temporary path: {basetemp}")
try:
environment = os.environ.copy()
environment["PYTHONPATH"] = os.pathsep.join(
filter(None, (str(REPOSITORY_ROOT), environment.get("PYTHONPATH", "")))
)
result = subprocess.run(
[
sys.executable,
"-m",
"pytest",
"-q",
"--import-mode=importlib",
"-o",
"addopts=",
"--capture=sys",
"--basetemp",
str(basetemp),
*[
str((REPOSITORY_ROOT / path).resolve())
for path in sorted(selection.python_tests)
],
],
cwd=COMFYUI_ROOT,
env=environment,
check=False,
)
return result.returncode
finally:
if basetemp.exists():
shutil.rmtree(basetemp)
if TEMP_ROOT.exists() and not any(TEMP_ROOT.iterdir()):
TEMP_ROOT.rmdir()
def parser() -> argparse.ArgumentParser:
result = argparse.ArgumentParser(description=__doc__)
mode = result.add_mutually_exclusive_group()
mode.add_argument("--final", action="store_true", help="run every tracked test")
mode.add_argument("--list-groups", action="store_true", help="list configured groups")
result.add_argument("--changed", action="store_true", help="select changed paths (default)")
result.add_argument("--base", help="include committed changes since BASE")
result.add_argument("--group", action="append", default=[], help="add a named group")
result.add_argument("--dry-run", action="store_true", help="explain without executing")
return result
def main(argv: list[str] | None = None) -> int:
args = parser().parse_args(argv)
groups = load_groups()
if args.list_groups:
for name in sorted(groups):
print(name)
return 0
if args.final:
selection = Selection(
groups={"final"}, python_tests=tracked_final_tests(),
reasons={"final": {"every tracked test"}},
)
else:
paths = changed_paths(args.base) if args.changed or args.base or not args.group else set()
revisions = ["HEAD", *([args.base] if args.base else [])]
historical = tuple(load_groups_from_revision(rev) for rev in revisions) if paths else ()
selection = select_tests(paths, groups, tuple(args.group), historical)
if selection.unmapped:
print("Unmapped production source files:", file=sys.stderr)
for path in sorted(selection.unmapped):
print(f" {path}", file=sys.stderr)
print("Update tests/test_groups.toml or use --final.", file=sys.stderr)
return 2
if args.dry_run:
print_selection(selection)
return 0
if not selection.python_tests:
print("No tests selected.")
return 0
return run_selection(selection)
if __name__ == "__main__":
raise SystemExit(main())
+381
View File
@@ -0,0 +1,381 @@
import importlib
import json
import sys
import types
from collections import Counter
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_adaptive_extraction_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def adaptive():
return _load_module("nodes.adaptive_svd")
def _exact_partial(weight, q, niter):
del niter
U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
return U[:, :q], S[:q], Vh[:q, :].T
@pytest.mark.parametrize(
("mode", "parameter", "expected"),
[("ratio", 2.0, 2), ("quantile", 0.7, 2), ("sv_fro", 0.8, 1)],
)
def test_adaptive_partial_rank_matches_reference(monkeypatch, adaptive, mode, parameter, expected):
monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial)
weight = torch.diag(torch.tensor([8.0, 5.0, 2.0, 1.0, 0.5, 0.25]))
_, _, _, rank = adaptive.adaptive_partial_svd(
weight, mode, parameter, max_rank=4, probe_offset=1, niter=2
)
assert rank == expected
def test_unresolved_ratio_expands_and_caps(monkeypatch, adaptive):
probes = []
def tracked(weight, q, niter):
probes.append(q)
return _exact_partial(weight, q, niter)
monkeypatch.setattr(adaptive.torch, "svd_lowrank", tracked)
weight = torch.diag(torch.linspace(10.0, 8.0, 12))
_, _, _, rank = adaptive.adaptive_partial_svd(
weight, "ratio", 100.0, max_rank=4, probe_offset=2, niter=2
)
assert probes == [6, 8]
assert rank == 4
def test_frobenius_uses_energy_outside_output_cap(monkeypatch, adaptive):
monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial)
weight = torch.eye(10)
_, _, _, rank = adaptive.adaptive_partial_svd(
weight, "sv_fro", 0.9, max_rank=3, probe_offset=2, niter=2
)
assert rank == 3
class _FakeStream:
def __init__(self, batches):
self._batches = iter(batches)
self.closed = False
def __iter__(self):
return self
def __next__(self):
return next(self._batches)
def close(self):
self.closed = True
class _FakeHandler:
def __init__(self, values):
self.values = values
self.calls = []
self.marked = []
self.streams = []
def async_stream(self, keys, **kwargs):
self.calls.append((tuple(keys), kwargs))
stream = _FakeStream([[(key, self.values[key])] for key in keys])
self.streams.append(stream)
return stream
def mark_processed(self, key):
self.marked.append(key)
def test_paired_stream_is_bounded_ordered_and_releases_once():
streaming = _load_module("nodes.extraction_stream")
handler_a = _FakeHandler({"a1": torch.tensor(1), "a2": torch.tensor(2)})
handler_b = _FakeHandler({"b1": torch.tensor(3)})
units = [("one", "a1", "b1"), ("two", "a2", None)]
observed = []
for key, tensor_a, tensor_b in streaming.paired_async_tensors(
handler_a, handler_b, units, pin_memory=True
):
observed.append((key, tensor_a.item(), None if tensor_b is None else tensor_b.item()))
assert observed == [("one", 1, 3), ("two", 2, None)]
assert handler_a.calls == [(('a1', 'a2'), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})]
assert handler_b.calls == [(('b1',), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})]
assert Counter(handler_a.marked) == Counter({"a1": 1, "a2": 1})
assert Counter(handler_b.marked) == Counter({"b1": 1})
assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
def test_paired_stream_decodes_quantized_logical_keys_per_work_unit():
streaming = _load_module("nodes.extraction_stream")
handler_a = _FakeHandler({
"layer1.code": torch.tensor(2),
"layer1.scale": torch.tensor(3),
"layer2.code": torch.tensor(4),
"layer2.scale": torch.tensor(5),
})
handler_b = _FakeHandler({"layer1": torch.tensor(1), "layer2": torch.tensor(2)})
class _Quantization:
quantized_keys = {"layer1", "layer2"}
def required_keys(self, key):
return (f"{key}.code", f"{key}.scale")
def decode(self, key, tensors, compute_dtype, device="cpu"):
assert set(tensors) == {f"{key}.code", f"{key}.scale"}
return tensors[f"{key}.code"] * tensors[f"{key}.scale"]
quantization = _Quantization()
observed = list(streaming.paired_async_tensors(
handler_a,
handler_b,
[("layer1", "layer1", "layer1"), ("layer2", "layer2", "layer2")],
pin_memory=False,
quantization_a=quantization,
))
assert [(key, tensor_a.item(), tensor_b.item()) for key, tensor_a, tensor_b in observed] == [
("layer1", 6, 1),
("layer2", 20, 2),
]
assert handler_a.calls[0][0] == (
"layer1.code", "layer1.scale", "layer2.code", "layer2.scale"
)
assert Counter(handler_a.marked) == Counter({
"layer1.code": 1, "layer1.scale": 1, "layer2.code": 1, "layer2.scale": 1,
})
assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
def test_cuda_oom_compute_retries_on_cpu(monkeypatch):
streaming = _load_module("nodes.extraction_stream")
monkeypatch.setattr(streaming, "release_cuda_oom", lambda: None)
calls = []
def fail_cuda():
calls.append("cuda")
raise RuntimeError("CUDA out of memory")
def run_cpu():
calls.append("cpu")
return 7
assert streaming.retry_cuda_oom_on_cpu(fail_cuda, run_cpu, "cuda:0") == 7
assert calls == ["cuda", "cpu"]
def test_adaptive_controls_cover_all_public_families():
modules_and_nodes = [
(_load_module("nodes.lora_extract_svd"), prefix)
for prefix in ("LoRAExtractRatio", "LoRAExtractQuantile", "LoRAExtractFrobenius")
]
modules_and_nodes += [
(_load_module("nodes.dora_extract_wd"), prefix)
for prefix in ("DoRAExtractRatio", "DoRAExtractQuantile", "DoRAExtractFrobenius")
]
modules_and_nodes += [
(_load_module("nodes.dora_learned_wd"), prefix)
for prefix in ("DoRALearnedExtractRatio", "DoRALearnedExtractQuantile", "DoRALearnedExtractFrobenius")
]
modules_and_nodes += [
(_load_module("nodes.text_encoder_extract"), prefix)
for prefix in (
"TextEncoderLoRAExtractRatio", "TextEncoderLoRAExtractQuantile",
"TextEncoderLoRAExtractFrobenius", "TextEncoderDoRAExtractRatio",
"TextEncoderDoRAExtractQuantile", "TextEncoderDoRAExtractFrobenius",
)
]
for module, node_name in modules_and_nodes:
names = [item.id for item in getattr(module, node_name).define_schema().inputs]
probe_index = names.index("probe_offset")
assert names[probe_index + 1] == "linear_max_rank"
def _write_uel(path, tensors):
with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
writer.write_dict(tensors)
def test_lora_extraction_uses_only_async_uel_and_releases_sources(monkeypatch, tmp_path):
extraction = _load_module("nodes.lora_extract_svd")
model_a = tmp_path / "a.safetensors"
model_b = tmp_path / "b.safetensors"
output = tmp_path / "out.safetensors"
key = "model.diffusion_model.layer.weight"
_write_uel(model_a, {key: torch.eye(3)})
_write_uel(model_b, {key: torch.zeros(3, 3)})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
original_async = extraction.MemoryEfficientSafeOpen.async_stream
original_mark = extraction.MemoryEfficientSafeOpen.mark_processed
async_calls = []
marked = []
def track_async(self, keys, **kwargs):
async_calls.append((self.filename, tuple(keys), kwargs))
yield from original_async(self, keys, **kwargs)
def track_mark(self, source_key):
marked.append((self.filename, source_key))
return original_mark(self, source_key)
monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "async_stream", track_async)
monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "mark_processed", track_mark)
monkeypatch.setattr(
extraction.MemoryEfficientSafeOpen,
"get_tensor",
lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("synchronous tensor load")),
)
extraction.extract_lora_from_files(
str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
)
assert output.is_file()
assert len(async_calls) == 2
assert all(options == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, _, options in async_calls)
assert Counter(marked) == Counter((filename, key) for filename, keys, _ in async_calls for key in keys)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
assert sorted(result.keys()) == [
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
]
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_lora_extraction_matches_prefixed_and_bare_models(monkeypatch, tmp_path, prefixed_input):
extraction = _load_module("nodes.lora_extract_svd")
paths = {name: tmp_path / f"{name}.safetensors" for name in ("a", "b")}
output = tmp_path / "mixed_prefix.safetensors"
for name in ("a", "b"):
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
value = torch.tensor([[2.0, 0.0], [0.0, 0.0]]) if name == "a" else torch.zeros((2, 2))
if name == prefixed_input:
stem = key[:-len(".weight")]
_write_uel(paths[name], {
key: value.to(torch.int8),
f"{stem}.weight_scale": torch.tensor(1.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
})
else:
_write_uel(paths[name], {key: value})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
extraction.extract_lora_from_files(
str(paths["a"]), str(paths["b"]), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
keys = sorted(result.keys())
assert keys == [
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
]
stream = result.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
tensors = {}
try:
for key in keys:
batch = next(stream)
assert len(batch) == 1 and batch[0][0] == key
tensors[key] = batch[0][1]
finally:
for key in tensors:
result.mark_processed(key)
stream.close()
delta = tensors[keys[1]] @ tensors[keys[0]]
torch.testing.assert_close(delta, torch.tensor([[2.0, 0.0], [0.0, 0.0]]), atol=1e-4, rtol=0)
@pytest.mark.parametrize("module_name,function_name,extra", [
("nodes.dora_extract_wd", "extract_dora_from_files", {}),
("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
])
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_dora_extraction_matches_prefixed_and_bare_models(
monkeypatch, tmp_path, module_name, function_name, extra, prefixed_input,
):
module = _load_module(module_name)
paths = {name: tmp_path / f"{name}.safetensors" for name in ("a", "b")}
output = tmp_path / "mixed_prefix_dora.safetensors"
for name in ("a", "b"):
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
value = torch.diag(torch.tensor([3.0, 1.0])) if name == "a" else torch.eye(2)
if name == prefixed_input:
stem = key[:-len(".weight")]
_write_uel(paths[name], {
key: value.to(torch.int8),
f"{stem}.weight_scale": torch.tensor(1.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
})
else:
_write_uel(paths[name], {key: value})
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
getattr(module, function_name)(
str(paths["a"]), str(paths["b"]), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False, **extra,
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
keys = set(result.keys())
assert "diffusion_model.layer.lora_down.weight" in keys
assert "diffusion_model.layer.lora_up.weight" in keys
assert not any("comfy_quant" in key or "model.diffusion_model." in key for key in keys)
@pytest.mark.parametrize("module_name, function_name, extra", [
("nodes.lora_extract_svd", "extract_lora_from_files", {"include_1d_diffs": False}),
("nodes.dora_extract_wd", "extract_dora_from_files", {}),
("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": False}),
("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": True}),
])
def test_extraction_include_mode_selects_only_matching_layers(monkeypatch, tmp_path, module_name, function_name, extra):
module = _load_module(module_name)
first, second = tmp_path / "first.safetensors", tmp_path / "second.safetensors"
match, other = "model.encoder.match.weight", "model.encoder.other.weight"
_write_uel(first, {match: torch.eye(2), other: torch.eye(2)})
_write_uel(second, {match: torch.zeros(2, 2), other: torch.zeros(2, 2)})
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
common = dict(linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
skip_patterns_str="match", include_mode=True, **extra)
selected = tmp_path / "selected.safetensors"
getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(selected), **common)
with MemoryEfficientSafeOpen(str(selected), low_memory=True) as result:
keys = list(result.keys())
assert keys
assert all("match" in key for key in keys)
empty = tmp_path / "empty.safetensors"
common["skip_patterns_str"] = ""
getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(empty), **common)
with MemoryEfficientSafeOpen(str(empty), low_memory=True) as result:
assert list(result.keys()) == []
globbed = tmp_path / "globbed.safetensors"
common.update(skip_patterns_str="*.match.weight", glob_skip_patterns=True)
getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(globbed), **common)
with MemoryEfficientSafeOpen(str(globbed), low_memory=True) as result:
assert list(result.keys()) == keys
+120
View File
@@ -0,0 +1,120 @@
import importlib
import os
import sys
import types
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_artifact_path_tests"
MIGRATED_SAVE_MODULES = (
"consensus_merger.py",
"cwb_model_graft.py",
"cwb_delta_lora_merger.py",
"merger.py",
"lora_merger.py",
"lodestone_merger.py",
"lora_resize.py",
"dtype_conversion.py",
"lora_extract_svd.py",
"dora_extract_wd.py",
"dora_learned_wd.py",
"text_encoder_extract.py",
"renamekeys.py",
"prunekeys.py",
"lora_rename.py",
"minimax_h3_lora_convert.py",
"minimax_h3_fold.py",
)
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def artifact_paths():
return _load_module("nodes.artifact_paths")
@pytest.mark.parametrize(
"category",
["checkpoints", "diffusion_models", "text_encoders", "loras", "embeddings"],
)
def test_canonical_categories_use_models_dir_and_return_relative_name(
monkeypatch, tmp_path, artifact_paths, category
):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
output_path, output_name = artifact_paths.canonical_model_artifact_path(
category, "group/subgroup/model"
)
expected = tmp_path / category / "group" / "subgroup" / "model.safetensors"
assert Path(output_path) == expected
assert output_name == "group/subgroup/model.safetensors"
assert expected.parent.is_dir()
assert not os.path.isabs(output_name)
def test_registered_search_paths_do_not_influence_write_root(
monkeypatch, tmp_path, artifact_paths
):
canonical = tmp_path / "cli-model-root"
legacy = tmp_path / "legacy-unet"
additional = tmp_path / "additional-models"
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(canonical))
monkeypatch.setattr(
artifact_paths.folder_paths,
"get_folder_paths",
lambda _: [str(legacy), str(additional)],
)
output_path, output_name = artifact_paths.canonical_model_artifact_path(
"diffusion_models", "nested/model.safetensors"
)
assert Path(output_path) == canonical / "diffusion_models" / "nested" / "model.safetensors"
assert output_name == "nested/model.safetensors"
@pytest.mark.parametrize(
"name",
["model.safetensors", "model.safetensors.safetensors"],
)
def test_appends_exactly_one_extension(monkeypatch, tmp_path, artifact_paths, name):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
_, output_name = artifact_paths.canonical_model_artifact_path("loras", name)
assert output_name == "model.safetensors"
@pytest.mark.parametrize(
"name",
["../escape", "folder/../../escape", "", " ", "/absolute", ".safetensors"],
)
def test_rejects_empty_or_escaping_names(monkeypatch, tmp_path, artifact_paths, name):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
with pytest.raises(ValueError):
artifact_paths.canonical_model_artifact_path("loras", name)
def test_rejects_unrecognized_category(monkeypatch, tmp_path, artifact_paths):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
with pytest.raises(ValueError, match="Unsupported model artifact category"):
artifact_paths.canonical_model_artifact_path("unet", "model")
def test_persistent_save_nodes_use_shared_path_contract_and_wildcard_outputs():
for filename in MIGRATED_SAVE_MODULES:
source = (REPO_ROOT / "nodes" / filename).read_text(encoding="utf-8")
assert (
"canonical_model_artifact_path" in source
or "_build_lora_output_path" in source
), filename
assert 'io.String.Output(display_name="output_path")' not in source, filename
assert 'io.String.Output(display_name="output_filename")' not in source, filename
+49
View File
@@ -0,0 +1,49 @@
import asyncio
import importlib.util
import sys
from pathlib import Path
from comfy_api.latest import io
REPO_ROOT = Path(__file__).resolve().parents[1]
def _load_extension_module():
package_name = "ComfyUI_ModelUtils_init_test"
if package_name not in sys.modules:
spec = importlib.util.spec_from_file_location(
package_name,
REPO_ROOT / "__init__.py",
submodule_search_locations=[str(REPO_ROOT)],
)
module = importlib.util.module_from_spec(spec)
sys.modules[package_name] = module
spec.loader.exec_module(module)
return sys.modules[package_name]
def test_startup_node_registration_when_downloaders_absent(monkeypatch):
ext_module = _load_extension_module()
ext = ext_module.ModelUtilsExtension()
nodes_with_downloader = asyncio.run(ext.get_node_list())
assert len(nodes_with_downloader) > 0
# Simulate registry install by setting DOWNLOADER_NODES to empty
monkeypatch.setattr(ext_module, "DOWNLOADER_NODES", [])
nodes_registry = asyncio.run(ext.get_node_list())
assert len(nodes_registry) < len(nodes_with_downloader)
assert all(issubclass(n, io.ComfyNode) for n in nodes_registry)
def test_model_info_nodes_importable_standalone():
spec = importlib.util.spec_from_file_location(
"modelutils_model_info_nodes_test",
REPO_ROOT / "nodes" / "model_info_nodes.py",
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
assert hasattr(module, "get_potential_preview_files")
assert hasattr(module, "load_image_tensor")
assert hasattr(module, "get_model_workflows")
assert hasattr(module, "get_model_metadata_file")
File diff suppressed because it is too large Load Diff
+141
View File
@@ -0,0 +1,141 @@
import importlib.util
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
def _load_runner():
spec = importlib.util.spec_from_file_location(
"cwb_api_runner_test", REPO_ROOT / "run_cwb_api_workflow.py"
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
@pytest.fixture(scope="module")
def runner():
return _load_runner()
@pytest.fixture
def workflow():
return {
"10": {
"class_type": "CWBCustomConfiguration",
"inputs": {"diversity_beta": 10.0},
},
"11": {
"class_type": "CWBLoRAMultiMerger",
"inputs": {
"cwb_config": ["10", 0],
"cwb_preset": "default",
"output_filename": "original",
},
},
"14": {"class_type": "PreviewAny", "inputs": {"source": ["11", 0]}},
"15": {"class_type": "PreviewAny", "inputs": {"source": ["11", 2]}},
}
def test_custom_export_remains_connected_by_default(runner, workflow):
merger_id, linked = runner._prepare_workflow(workflow, None, None)
assert merger_id == "11"
assert workflow["11"]["inputs"]["cwb_config"] == ["10", 0]
assert workflow["11"]["inputs"]["exclude_patterns"] == (
runner.CALIBRATION_EXCLUDE_PATTERN
)
assert linked == {0: "14", 2: "15"}
def test_preset_run_disconnects_custom_config_and_overrides_filename(
runner, workflow
):
runner._prepare_workflow(workflow, "candidate_sim_mean", "candidate/output")
inputs = workflow["11"]["inputs"]
assert "cwb_config" not in inputs
assert inputs["cwb_preset"] == "candidate_sim_mean"
assert inputs["output_filename"] == "candidate/output"
def test_full_block_verification_does_not_add_calibration_exclusions(
runner, workflow
):
workflow["11"]["inputs"]["exclude_patterns"] = "existing"
runner._prepare_workflow(workflow, None, None, full_blocks=True)
assert workflow["11"]["inputs"]["exclude_patterns"] == "existing"
@pytest.mark.parametrize(
("block", "excluded"),
[(block, block not in {0, 3, 14, 24, 27}) for block in range(28)],
)
def test_calibration_block_ranges_are_exact(runner, block, excluded):
import re
key = f"diffusion_model.blocks.{block}.attn.wq.lora_A.weight"
patterns = runner.CALIBRATION_EXCLUDE_PATTERN.splitlines()
assert any(re.search(pattern, key) for pattern in patterns) is excluded
@pytest.mark.parametrize(
("key", "excluded"),
[
("diffusion_model.txtfusion.layerwise_blocks.1.attn.lora_A.weight", True),
("diffusion_model.txtfusion.refiner_blocks.0.attn.lora_A.weight", True),
("diffusion_model.txtfusion.layerwise_blocks.0.attn.lora_A.weight", False),
],
)
def test_calibration_txtfusion_exclusions_are_exact(runner, key, excluded):
import re
patterns = runner.CALIBRATION_EXCLUDE_PATTERN.splitlines()
assert any(re.search(pattern, key) for pattern in patterns) is excluded
def test_requires_filename_and_report_capture(runner, workflow):
del workflow["15"]
with pytest.raises(RuntimeError, match="output 2"):
runner._prepare_workflow(workflow, None, None)
def test_parses_counterfactual_weight_sweep(runner):
report = """CWB COUNTERFACTUAL WEIGHT SWEEP
median | 0.1 | 2 | 4 | true | false | 12/18 | 3 | 0.6/1 | 2.4/1
"""
rows = runner.parse_counterfactual_weight_sweep(report)
assert rows == [{
"consensus_type": "median",
"similarity_threshold": 0.1,
"power_alpha": 2.0,
"diversity_beta": 4.0,
"dynamic_similarity_contrast": True,
"soft_comfort_bandpass": False,
"accepted_contributors": 12,
"contributors": 18,
"fallbacks": 3,
"dominant_weight_mean": 0.6,
"dominant_weight_max": 1.0,
"effective_contributors_mean": 2.4,
"effective_contributors_min": 1.0,
}]
def test_prompt_payload_contains_only_api_workflow(runner, workflow):
payload = runner._prompt_payload(workflow)
assert payload == {"prompt": workflow}
@pytest.mark.parametrize(
("output", "expected"),
[
({"text": ["value"]}, "value"),
({"text": "value"}, "value"),
({"ui": {"text": ["value"]}}, "value"),
],
)
def test_extracts_supported_preview_history_shapes(runner, output, expected):
assert runner._text_output({"14": output}, "14") == expected
+405
View File
@@ -0,0 +1,405 @@
import importlib
import re
import sys
import types
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_delta_cwb_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def delta_cwb():
return _load_module("nodes.cwb_delta_lora_merger")
def _settings(delta_cwb):
return delta_cwb.CWBSettings(
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.0,
similarity_threshold=-1.0,
power_alpha=0.0,
diversity_beta=0.0,
rescale_norm=False,
global_scale=1.0,
dynamic_similarity_contrast=False,
soft_comfort_bandpass=False,
position_weight=0.0,
preserve_common_prefix=False,
)
def _write(path, tensors, metadata=None):
with IncrementalSafetensorsWriter(str(path), metadata=metadata or {}, max_workers=1) as writer:
writer.write_batch(list(tensors.items()))
def _read(path):
tensors = {}
with MemoryEfficientSafeOpen(str(path), low_memory=True) as handler:
stream = handler.async_stream(
handler.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
)
for batch in stream:
for key, tensor in batch:
tensors[key] = tensor.clone()
handler.mark_processed(key)
return tensors
def test_linear_and_convolutional_factor_expansion(delta_cwb):
down = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
up = torch.tensor([[2.0, -1.0], [0.5, 3.0]])
torch.testing.assert_close(
delta_cwb._factor_delta(down, up, "linear"), up @ down
)
conv_down = torch.arange(18, dtype=torch.float32).reshape(2, 1, 3, 3)
conv_up = torch.tensor([[[[2.0]], [[-1.0]]]])
expected = (conv_up.reshape(1, 2) @ conv_down.reshape(2, -1)).reshape(1, 1, 3, 3)
torch.testing.assert_close(
delta_cwb._factor_delta(conv_down, conv_up, "conv"), expected
)
def test_dense_cwb_keeps_corresponding_rows_fixed(delta_cwb):
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
result = delta_cwb.merge_cwb_tensors(
[reference, reversed_source],
_settings(delta_cwb),
reference_index=0,
allow_similarity_alignment=False,
)
torch.testing.assert_close(result, torch.full((2, 2), 0.5))
def test_layer_units_apply_missing_modes_filters_and_shape_checks(delta_cwb):
class Handler:
def __init__(self, shapes):
self.shapes = shapes
def get_shape(self, key):
return self.shapes[key]
def get_dtype(self, _key):
return torch.float32
roles = {"down": "layer.down", "up": "layer.up"}
infos = [
{"pairs": {"layer": roles}, "low_bit": set()},
{"pairs": {}, "low_bit": set()},
]
handlers = [Handler({"layer.down": (1, 2), "layer.up": (2, 1)}), Handler({})]
layer_map = {"layer": [(0, "layer")]}
empty = delta_cwb._compile_patterns("", False)
with pytest.raises(ValueError, match="missing or incompatible"):
delta_cwb._build_layer_units(
layer_map, infos, handlers, "error", empty, empty, False
)
units, _, _ = delta_cwb._build_layer_units(
layer_map, infos, handlers, "zeros", empty, empty, False
)
assert units[0].zero_contributors == 1
assert units[0].preserve_anchor is False
exclude = delta_cwb._compile_patterns(r"\.diff$", False)
units, _, _ = delta_cwb._build_layer_units(
layer_map, infos, handlers, "zeros", exclude, empty, False
)
assert units[0].preserve_anchor is True
assert units[0].zero_contributors == 0
discard = delta_cwb._compile_patterns("*.diff", True)
units, _, rejected = delta_cwb._build_layer_units(
layer_map, infos, handlers, "skip", empty, discard, True
)
assert units == []
assert rejected == 1
def test_streamed_merge_normalizes_alpha_once_and_supports_mixed_rank(
monkeypatch, tmp_path, delta_cwb
):
first = tmp_path / "first.safetensors"
second = tmp_path / "second.safetensors"
_write(first, {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [0.0]]),
"diffusion_model.layer.alpha": torch.tensor(0.5),
}, {"anchor": "yes"})
_write(second, {
"diffusion_model.layer.lora_A.weight": torch.eye(2),
"diffusion_model.layer.lora_B.weight": torch.diag(torch.tensor([3.0, 4.0])),
})
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
filename, report = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)],
[first.name, second.name],
_settings(delta_cwb),
"balanced_mean",
False,
"skip",
"dense/mixed_rank",
torch.float32,
"full_difference",
384,
"cpu",
"",
"",
False,
False,
)
output = tmp_path / "loras" / filename
tensors = _read(output)
assert filename == "dense/mixed_rank.safetensors"
assert list(tensors) == ["diffusion_model.layer.diff"]
torch.testing.assert_close(
tensors["diffusion_model.layer.diff"],
torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
)
assert "Merged layers: 1" in report
with MemoryEfficientSafeOpen(str(output), low_memory=True) as handler:
metadata = handler.metadata()
assert metadata["anchor"] == "yes"
assert metadata["merge_method"] == "delta_cwb"
assert metadata["alpha_normalized"] == "true"
assert metadata["output_representation"] == "full_direct_difference"
@pytest.mark.parametrize("patterns, glob_patterns, selected", [
(r"\.selected", False, True),
("*.selected*", True, True),
("", False, False),
("no_match", False, False),
])
def test_include_filter_preserves_nonmatches_and_discards_first(
monkeypatch, tmp_path, delta_cwb, patterns, glob_patterns, selected,
):
paths = [tmp_path / f"{name}.safetensors" for name in ("a", "b")]
for path, value in zip(paths, (2.0, 6.0)):
_write(path, {
f"diffusion_model.{layer}.diff": torch.full((2, 2), value)
for layer in ("selected", "other", "selected_discard")
})
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(path) for path in paths], [path.name for path in paths], _settings(delta_cwb),
"balanced_mean", False, "skip", "include_filter", torch.float32,
"full_difference", 384, "cpu", patterns,
"*discard*" if glob_patterns else "discard", glob_patterns, False, include_mode=True,
)
tensors = _read(tmp_path / "loras" / filename)
assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"}
torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 4.0 if selected else 2.0))
torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0))
def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cwb):
first = tmp_path / "diff_first.safetensors"
second = tmp_path / "diff_second.safetensors"
key = "diffusion_model.bias.diff"
_write(first, {key: torch.tensor([1.0, 3.0])})
_write(second, {key: torch.tensor([3.0, 5.0])})
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
"balanced_mean", False, "skip", "one_dim", torch.float16,
"full_difference", 384, "cpu",
"", "", False, False,
)
output = _read(tmp_path / "loras" / filename)[key]
assert output.dtype == torch.float32
torch.testing.assert_close(output, torch.tensor([2.0, 4.0]))
def test_factorized_output_writes_compact_factors_without_dense_tensor(
monkeypatch, tmp_path, delta_cwb
):
first = tmp_path / "factor_first.safetensors"
second = tmp_path / "factor_second.safetensors"
tensors = {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [6.0]]),
}
_write(first, tensors)
_write(second, tensors)
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
"balanced_mean", False, "skip", "factorized", torch.float32,
"factorized_lora", 1, "cpu", "", "", False, False,
)
output_path = tmp_path / "loras" / filename
output = _read(output_path)
assert set(output) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
}
assert "diffusion_model.layer.diff" not in output
reconstructed = (
output["diffusion_model.layer.lora_B.weight"]
@ output["diffusion_model.layer.lora_A.weight"]
)
torch.testing.assert_close(reconstructed, torch.tensor([[3.0, 6.0], [6.0, 12.0]]))
with MemoryEfficientSafeOpen(str(output_path), low_memory=True) as handler:
assert handler.metadata()["output_representation"] == "factorized_lora"
def test_atomic_writer_preserves_destination_on_cwb_failure(
monkeypatch, tmp_path, delta_cwb
):
first = tmp_path / "atomic_first.safetensors"
second = tmp_path / "atomic_second.safetensors"
tensors = {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [0.0]]),
}
_write(first, tensors)
_write(second, tensors)
destination = tmp_path / "loras" / "atomic.safetensors"
destination.parent.mkdir(parents=True)
destination.write_bytes(b"existing destination")
monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None)
released = []
closed = []
opener = delta_cwb.MemoryEfficientSafeOpen
class TrackedHandler:
def __init__(self, path, low_memory=True):
self.handler = opener(path, low_memory=low_memory)
def __getattr__(self, name):
return getattr(self.handler, name)
def mark_processed(self, key):
released.append(key)
return self.handler.mark_processed(key)
def __exit__(self, exc_type, exc, traceback):
closed.append(True)
return self.handler.__exit__(exc_type, exc, traceback)
monkeypatch.setattr(delta_cwb, "MemoryEfficientSafeOpen", TrackedHandler)
monkeypatch.setattr(
delta_cwb,
"merge_cwb_tensors",
lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("planned failure")),
)
with pytest.raises(RuntimeError, match="planned failure"):
delta_cwb.DeltaCWBLoRAMergerLogic.execute(
[str(first), str(second)], [first.name, second.name], _settings(delta_cwb),
"balanced_mean", False, "skip", "atomic", torch.float32,
"full_difference", 384, "cpu",
"", "", False, False,
)
assert destination.read_bytes() == b"existing destination"
assert not list(destination.parent.glob(".atomic.safetensors.*.tmp"))
assert sorted(released) == sorted([*tensors, *tensors])
assert len(closed) == 2
def test_complete_layer_cuda_oom_retries_on_cpu(monkeypatch, delta_cwb):
source = delta_cwb.DeltaCWBLoRASource(
0, "layer", None, None, None, "layer.diff", torch.float32
)
unit = delta_cwb.DeltaCWBLayerUnit(
"layer", "layer.diff", (source,), 0, False
)
calls = []
def process(
_unit, _loaded, _settings, _dtype, device, _diagnostics,
_output_representation, _output_rank,
):
calls.append(device.type)
if device.type == "cuda":
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
return torch.ones(1), torch.float32
monkeypatch.setattr(delta_cwb, "_process_layer_on_device", process)
monkeypatch.setattr(delta_cwb, "_release_cuda_oom", lambda: None)
output, dtype, fallback = delta_cwb._process_layer(
unit, {}, _settings(delta_cwb), torch.float32, "cuda", delta_cwb.CWBDiagnostics()
)
assert calls == ["cuda", "cpu"]
assert fallback is True
assert dtype == torch.float32
assert torch.equal(output, torch.ones(1))
def test_exact_node_schema_contract(delta_cwb):
expected_ids = [
"DeltaCWBLoRATwoMerger",
"DeltaCWBLoRAThreeMerger",
"DeltaCWBLoRAMultiMerger",
]
schemas = [node.define_schema() for node in delta_cwb.DELTA_CWB_LORA_NODES]
assert [schema.node_id for schema in schemas] == expected_ids
assert [schema.display_name for schema in schemas] == [
"Delta CWB Merge LoRAs (2)",
"Delta CWB Merge LoRAs (3)",
"Delta CWB LoRA Multi-Merge",
]
assert all(schema.category == "ModelUtils/LoRA/Merge/Delta CWB" for schema in schemas)
assert all([output.display_name for output in schema.outputs] == ["output_filename", "cwb_report"] for schema in schemas)
assert all(schema.outputs[0].get_io_type() == "*" for schema in schemas)
fixed_ids = [value.id for value in schemas[0].inputs]
assert fixed_ids == [
"lora_1", "lora_2", "cwb_preset", "cwb_config", "mismatch_mode",
"output_filename", "save_dtype", "output_representation", "output_rank",
"process_device", "exclude_patterns",
"discard_patterns", "glob_patterns", "force_clear_cache", "include_mode",
]
multi_ids = [value.id for value in schemas[2].inputs]
assert multi_ids[:9] == ["lora_count", *[f"lora_{index}" for index in range(1, 9)]]
assert schemas[0].inputs[2].default == "balanced_mean"
assert schemas[0].inputs[4].default == "skip"
assert schemas[0].inputs[6].default == "bf16"
assert schemas[0].inputs[7].default == "full_difference"
assert schemas[0].inputs[8].default == 384
assert schemas[0].inputs[9].default == "cuda"
assert schemas[0].inputs[-2].default is False
assert schemas[0].inputs[-1].default is False
assert "agreement" in schemas[0].inputs[2].tooltip
assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[9].tooltip
def test_invalid_regex_and_unselected_multi_input_fail(delta_cwb):
with pytest.raises(re.error):
delta_cwb._compile_patterns("[", False)
with pytest.raises(ValueError, match="unselected input.*3"):
delta_cwb.DeltaCWBLoRAMultiMerger.execute(
lora_count="3", lora_1="a", lora_2="b", lora_3="None"
)
+182
View File
@@ -0,0 +1,182 @@
import importlib
from pathlib import Path
import sys
import types
import pytest
from safetensors.torch import load_file, save_file
import torch
REPO_ROOT = Path(__file__).resolve().parents[1]
@pytest.fixture(scope="module")
def graft_mod():
package_name = "modelutils_cwb_graft_tests"
if package_name not in sys.modules:
package = types.ModuleType(package_name)
package.__path__ = [str(REPO_ROOT)]
sys.modules[package_name] = package
return importlib.import_module(f"{package_name}.nodes.cwb_model_graft")
def _patch_io(monkeypatch, graft_mod, tmp_path, paths):
monkeypatch.setattr(
graft_mod.folder_paths,
"get_full_path",
lambda _kind, name: paths.get(name),
)
monkeypatch.setattr(
graft_mod.folder_paths,
"get_folder_paths",
lambda _kind: [str(tmp_path)],
)
monkeypatch.setattr(graft_mod.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(graft_mod, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(graft_mod, "cleanup_after_operation", lambda: None)
def test_cwb_model_graft_schema_and_documentation(graft_mod):
schema = graft_mod.CWBModelGraft.define_schema()
assert schema.node_id == "CWBModelGraft"
assert schema.display_name == "CWB Model Graft"
ids = [inp.id for inp in schema.inputs]
assert "attention_preset" in ids
assert "ffn_preset" in ids
assert "protect_early_blocks" in ids
assert "attention_block_rule" in ids
assert "ffn_block_rule" in ids
assert "attention_subpart" in ids
assert "ffn_subpart" in ids
assert "custom_key_map" in ids
assert ids[-1] == "include_mode"
doc_out = graft_mod.CWBModelGraft.execute(execution_mode="DOCUMENTATION ONLY")
filename, docs, report = doc_out.result
assert filename == "Documentation mode active. No merge performed."
assert "# CWB Model Graft" in docs
assert "Complete Input Reference" in docs
assert "no merge was performed" in report
def test_cwb_model_graft_protected_blocks_subparts_and_mismatched_shapes(
monkeypatch, tmp_path, graft_mod
):
base_path = tmp_path / "base.safetensors"
donor_path = tmp_path / "donor.safetensors"
base_qkv_0 = torch.ones((6, 4), dtype=torch.float32) * 2.0
base_qkv_2 = torch.cat(
[
torch.ones((2, 4), dtype=torch.float32) * 1.0,
torch.ones((2, 4), dtype=torch.float32) * 2.0,
torch.ones((2, 4), dtype=torch.float32) * 3.0,
],
dim=0,
)
base_ffn_2 = torch.cat(
[
torch.ones((2, 4), dtype=torch.float32) * 4.0,
torch.ones((2, 4), dtype=torch.float32) * 5.0,
],
dim=0,
)
base_custom = torch.ones((2, 4), dtype=torch.float32) * 2.0
save_file(
{
"layers.0.proj.qkv.weight": base_qkv_0,
"layers.2.proj.qkv.weight": base_qkv_2,
"layers.2.proj.ffn_up_gate.weight": base_ffn_2,
"layers.3.proj.qkv.weight": base_qkv_0.clone(),
"custom_target.weight": base_custom,
},
str(base_path),
)
save_file(
{
"layers.0.proj.qkv.weight": torch.ones((6, 4), dtype=torch.float32) * 99.0,
"layers.2.proj.q.weight": torch.ones((2, 4), dtype=torch.float32) * 50.0,
"layers.2.proj.k.weight": torch.ones((2, 4), dtype=torch.float32) * 60.0,
"layers.2.proj.v.weight": torch.tensor(
[[9.0, 1.0, 1.0, 1.0], [1.0, 9.0, 1.0, 1.0]], dtype=torch.float32
),
"layers.2.proj.ffn_up_gate.weight": torch.cat(
[
torch.tensor(
[[8.0, 1.0, 1.0, 1.0], [1.0, 8.0, 1.0, 1.0]],
dtype=torch.float32,
),
torch.ones((2, 4), dtype=torch.float32) * 77.0,
],
dim=0,
),
"layers.3.proj.qkv.weight": torch.ones((6, 4), dtype=torch.float32) * 88.0,
"custom_donor.weight": torch.tensor(
[
[10.0, 2.0, 1.0, 1.0, 1.0, 1.0],
[1.0, 10.0, 2.0, 1.0, 1.0, 1.0],
[1.0, 1.0, 10.0, 2.0, 1.0, 1.0],
],
dtype=torch.float32,
),
},
str(donor_path),
)
paths = {"base": str(base_path), "donor": str(donor_path)}
_patch_io(monkeypatch, graft_mod, tmp_path, paths)
out = graft_mod.CWBModelGraft.execute(
execution_mode="MERGE",
model_a="base",
model_b="donor",
attention_preset="balanced_mean",
attention_config=None,
ffn_preset="balanced_mean",
ffn_config=None,
protect_early_blocks=2,
attention_block_rule="interleaved_even",
ffn_block_rule="all_after_protected",
attention_subpart="value_and_output",
ffn_subpart="up_only",
custom_key_map="custom_donor.weight -> custom_target.weight",
output_filename="grafted",
save_dtype="fp32",
process_device="cpu",
exclude_patterns="",
discard_patterns="",
glob_patterns=False,
lazy_load=True,
force_clear_cache=False,
override_dtype=False,
include_mode=False,
)
rel_name, docs, report = out.result
assert "# CWB Model Graft" in docs
assert "attn=balanced_mean, ffn=balanced_mean" in report
assert "Grafted tensors: 3" in report
tensors = load_file(str(tmp_path / "diffusion_models" / rel_name))
# 1. Protected block 0 is byte-exact identical to Base
assert torch.equal(tensors["layers.0.proj.qkv.weight"], base_qkv_0)
# 2. Odd block 3 attention is untouched under interleaved_even
assert torch.equal(tensors["layers.3.proj.qkv.weight"], base_qkv_0)
# 3. Even block 2 qkv: Q and K slices stay Base, only V slice is grafted from split donor
q_out, k_out, v_out = tensors["layers.2.proj.qkv.weight"].chunk(3, dim=0)
assert torch.equal(q_out, base_qkv_2[:2])
assert torch.equal(k_out, base_qkv_2[2:4])
assert not torch.equal(v_out, base_qkv_2[4:6])
# 4. Block 2 ffn_up_gate: up half is grafted, gate half stays Base
up_out, gate_out = tensors["layers.2.proj.ffn_up_gate.weight"].chunk(2, dim=0)
assert not torch.equal(up_out, base_ffn_2[:2])
assert torch.equal(gate_out, base_ffn_2[2:])
# 5. Mismatched-shape custom mapped tensor blends without changing shape or norm
assert tensors["custom_target.weight"].shape == base_custom.shape
assert not torch.equal(tensors["custom_target.weight"], base_custom)
assert torch.linalg.vector_norm(tensors["custom_target.weight"]).item() == pytest.approx(
torch.linalg.vector_norm(base_custom).item(), rel=1e-5
)
+120
View File
@@ -0,0 +1,120 @@
import importlib.util
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
def _load_sweep():
spec = importlib.util.spec_from_file_location(
"cwb_preset_sweep_test", REPO_ROOT / "run_cwb_preset_sweep.py"
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
@pytest.fixture(scope="module")
def sweep():
return _load_sweep()
@pytest.fixture
def baseline():
return {
"consensus_type": "median",
"alignment_method": "similarity",
"alignment_threshold": 0.0,
"similarity_threshold": 0.0,
"power_alpha": 2.0,
"diversity_beta": 10.0,
"rescale_norm": True,
"global_scale": 1.0,
"dynamic_similarity_contrast": False,
"soft_comfort_bandpass": True,
"position_weight": 0.05,
"preserve_common_prefix": False,
}
def test_cases_are_complete_and_change_exactly_one_factor(sweep, baseline):
cases = sweep.build_sweep_cases(baseline)
assert cases[0]["id"] == "baseline"
assert len({case["id"] for case in cases}) == len(cases)
for case in cases:
assert tuple(case["config"]) == sweep.CONFIG_KEYS
changed = [
key for key in sweep.CONFIG_KEYS
if case["config"][key] != baseline[key]
]
assert changed == ([] if case["id"] == "baseline" else [case["factor"]])
def test_sweep_exposes_expected_full_block_finalist_ids(sweep, baseline):
ids = {case["id"] for case in sweep.build_sweep_cases(baseline)}
assert {
"baseline",
"alignment_method_index",
"alignment_threshold_0p0005",
"alignment_threshold_0p0025",
"consensus_type_mean",
"diversity_beta_4p0",
"rescale_norm_false",
"dynamic_similarity_contrast_true",
"soft_comfort_bandpass_false",
} <= ids
def test_report_parser_extracts_influence_metrics(sweep):
report = """CWB MERGE SUMMARY
Alignment candidates: 100
Alignment matches: 75
Anchor-only groups: 5
Weighted vector groups: 20
All-rejected equal fallbacks: 1
Zero-weight equal fallbacks: 2
Accepted weighting contributors: 50/80
Consensus similarity min/mean/max: -0.5 / 0.25 / 1
Dominant normalized weight min/mean/max: 0.2 / 0.4 / 0.8
Effective contributors min/mean/max: 1.2 / 2.5 / 5
CUDA OOM CPU fallbacks: 0
LoRA groups: 10
Rank-one delta norm ratio count/min/mean/max: 3 / 0.7 / 1.1 / 1.4
Rank-one delta norm ratio count/min/mean/max: 1 / 0.8 / 0.9 / 1.2
"""
metrics = sweep.parse_cwb_report(report)
assert metrics["alignment_match_fraction"] == pytest.approx(0.75)
assert metrics["accepted_weight_fraction"] == pytest.approx(0.625)
assert metrics["dominant_weight_mean"] == pytest.approx(0.4)
assert metrics["effective_contributors_mean"] == pytest.approx(2.5)
assert metrics["norm_ratio_min"] == pytest.approx(0.7)
assert metrics["norm_ratio_count"] == 4
assert metrics["norm_ratio_mean"] == pytest.approx(1.05)
assert metrics["norm_ratio_max"] == pytest.approx(1.4)
def test_report_parser_rejects_pre_diagnostic_server(sweep):
with pytest.raises(RuntimeError, match="Restart ComfyUI"):
sweep.parse_cwb_report("CWB MERGE SUMMARY\nAlignment matches: 1")
def test_analysis_parser_extracts_global_metrics(sweep):
comparison = """MODEL COMPARISON SUMMARY
PARAMETER-WEIGHTED GLOBAL METRICS
relative_l2: 1.04
cosine: 0.46
exact: 0.01
norm_ratio: 0.89
EQUAL-LAYER AVERAGES
"""
cwb_report = """GLOBAL CWB-DERIVED DIAGNOSTICS
Pairwise row cosine mean/min/max: 0.4 / -0.1 / 1
Mean-consensus affinity A/B: 0.2 / 0.3
Median-consensus affinity A/B: 0.25 / 0.35
"""
metrics = sweep.parse_analysis_reports(comparison, cwb_report)
assert metrics["analysis_relative_l2"] == pytest.approx(1.04)
assert metrics["analysis_pair_cosine"] == pytest.approx(0.4)
assert metrics["analysis_median_affinity_b"] == pytest.approx(0.35)
+269
View File
@@ -0,0 +1,269 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import (
IncrementalSafetensorsWriter,
MemoryEfficientSafeOpen,
)
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_dtype_conversion_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def conversion():
return _load_module("nodes.dtype_conversion")
def _write_uel(path, tensors, metadata=None):
with IncrementalSafetensorsWriter(
str(path), metadata=metadata or {}, max_workers=1
) as writer:
writer.write_batch(list(tensors.items()))
def _read_uel(path):
result = {}
with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
metadata = loader.metadata()
for batch in loader.async_stream(
loader.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
):
key, tensor = batch[0]
result[key] = tensor.clone()
loader.mark_processed(key)
return result, metadata
def _patch_runtime(monkeypatch, tmp_path, conversion, source):
monkeypatch.setattr(conversion.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(
conversion.folder_paths,
"get_full_path_or_raise",
lambda model_type, name: str(source),
)
monkeypatch.setattr(conversion, "cleanup_after_operation", lambda: None)
def test_schema_is_diffusion_model_scoped(conversion):
schema = conversion.DiffusionModelDtypeConversion.define_schema()
assert schema.node_id == "DiffusionModelDtypeConversion"
assert schema.category == "ModelUtils/Conversion"
assert [item.id for item in schema.inputs] == [
"model_name",
"target_dtype",
"reference_model",
"exclude_patterns",
"output_filename",
"include_mode",
]
assert schema.inputs[-1].default is False
assert schema.inputs[1].options == ["fp32", "fp16", "bf16"]
assert [item.display_name for item in schema.outputs] == [
"output_path",
"conversion_report",
]
def test_reference_model_dtypes_override_target_without_tensor_loading(
monkeypatch, tmp_path, conversion
):
source = tmp_path / "source.safetensors"
reference = tmp_path / "reference.safetensors"
_write_uel(source, {
"mapped_bf16": torch.ones(2, dtype=torch.float32),
"mapped_fp32": torch.ones(2, dtype=torch.float16),
"fallback": torch.ones(2, dtype=torch.float32),
"nonfloating": torch.tensor([1], dtype=torch.int64),
})
_write_uel(reference, {
"mapped_bf16": torch.ones(2, dtype=torch.bfloat16),
"mapped_fp32": torch.ones(2, dtype=torch.float32),
"nonfloating": torch.tensor([1], dtype=torch.int64),
})
monkeypatch.setattr(conversion.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(
conversion.folder_paths,
"get_full_path_or_raise",
lambda _, name: {"source": str(source), "reference": str(reference)}[name],
)
monkeypatch.setattr(conversion, "cleanup_after_operation", lambda: None)
monkeypatch.setattr(
conversion.MemoryEfficientSafeOpen,
"get_tensor",
lambda *args, **kwargs: pytest.fail("synchronous reference loading was used"),
)
output_path, report = conversion.convert_diffusion_model_dtype(
"source", "fp16", "", "reference_mapped", "reference"
)
result, _ = _read_uel(tmp_path / "diffusion_models" / output_path)
assert result["mapped_bf16"].dtype == torch.bfloat16
assert result["mapped_fp32"].dtype == torch.float32
assert result["fallback"].dtype == torch.float16
assert result["nonfloating"].dtype == torch.int64
assert "Reference-mapped floating tensors: 2" in report
assert "Reference-missing fallback tensors: 1" in report
@pytest.mark.parametrize(
("target_name", "expected_dtype"),
[("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
)
def test_converts_floating_and_preserves_excluded_and_nonfloating(
monkeypatch, tmp_path, conversion, target_name, expected_dtype
):
source = tmp_path / f"source_{target_name}.safetensors"
tensors = {
"blocks.0.weight": torch.arange(4, dtype=torch.float32),
"blocks.0.norm.weight": torch.ones(3, dtype=torch.float32),
"position_ids": torch.tensor([1, 2], dtype=torch.int64),
}
_write_uel(source, tensors, metadata={"source": "test"})
_patch_runtime(monkeypatch, tmp_path, conversion, source)
output_path, report = conversion.convert_diffusion_model_dtype(
source.name,
target_name,
r"\.norm\.weight$",
f"converted_{target_name}",
)
assert output_path == f"converted_{target_name}.safetensors"
result, metadata = _read_uel(
tmp_path / "diffusion_models" / output_path
)
assert result["blocks.0.weight"].dtype == expected_dtype
assert result["blocks.0.norm.weight"].dtype == torch.float32
assert result["position_ids"].dtype == torch.int64
assert metadata == {"source": "test"}
assert "Excluded tensors preserved: 1" in report
assert "Non-floating tensors preserved: 1" in report
@pytest.mark.parametrize("patterns, selected", [
(r"^blocks\.0\.weight$", {"blocks.0.weight"}),
(r"^blocks\.0\.weight$" + "\n" + r"^blocks\.0\.norm\.weight$", {"blocks.0.weight", "blocks.0.norm.weight"}),
("", set()),
("no_such_layer", set()),
])
@pytest.mark.parametrize("use_reference", [False, True])
def test_include_filter_only_converts_matches(
monkeypatch, tmp_path, conversion, patterns, selected, use_reference,
):
source = tmp_path / "source.safetensors"
reference = tmp_path / "reference.safetensors"
tensors = {
"blocks.0.weight": torch.arange(4, dtype=torch.float32),
"blocks.0.norm.weight": torch.ones(3, dtype=torch.float32),
"position_ids": torch.tensor([1, 2], dtype=torch.int64),
}
_write_uel(source, tensors)
_write_uel(reference, {"blocks.0.weight": tensors["blocks.0.weight"].to(torch.bfloat16)})
_patch_runtime(monkeypatch, tmp_path, conversion, source)
monkeypatch.setattr(
conversion.folder_paths, "get_full_path_or_raise",
lambda _, name: str(reference if name == "reference" else source),
)
output_name, _ = conversion.DiffusionModelDtypeConversion.execute(
source.name, "fp16", "reference" if use_reference else "None",
patterns, "include_converted", include_mode=True,
)
result, _ = _read_uel(tmp_path / "diffusion_models" / output_name)
for key, tensor in tensors.items():
expected_dtype = tensor.dtype
if key in selected:
expected_dtype = torch.bfloat16 if use_reference and key == "blocks.0.weight" else torch.float16
assert result[key].dtype == expected_dtype
torch.testing.assert_close(result[key].to(tensor.dtype), tensor)
def test_uses_bounded_async_stream_and_releases_every_tensor(
monkeypatch, tmp_path, conversion
):
source = tmp_path / "source_async.safetensors"
_write_uel(
source,
{
"a": torch.ones(2, dtype=torch.float32),
"b": torch.ones(2, dtype=torch.float16),
"c": torch.ones(2, dtype=torch.int32),
},
)
_patch_runtime(monkeypatch, tmp_path, conversion, source)
async_calls = []
processed = []
original_async = conversion.MemoryEfficientSafeOpen.async_stream
original_mark = conversion.MemoryEfficientSafeOpen.mark_processed
def track_async(loader, keys, **kwargs):
async_calls.append((tuple(keys), kwargs))
yield from original_async(loader, keys, **kwargs)
def track_mark(loader, key):
processed.append(key)
return original_mark(loader, key)
monkeypatch.setattr(
conversion.MemoryEfficientSafeOpen, "async_stream", track_async
)
monkeypatch.setattr(
conversion.MemoryEfficientSafeOpen, "mark_processed", track_mark
)
monkeypatch.setattr(
conversion.MemoryEfficientSafeOpen,
"get_tensor",
lambda *args, **kwargs: pytest.fail("synchronous loading was used"),
)
conversion.convert_diffusion_model_dtype(
source.name, "bf16", "^b$", "converted_async"
)
assert len(async_calls) == 1
assert async_calls[0][1] == {
"batch_size": 1,
"prefetch_batches": 1,
"pin_memory": False,
}
assert processed == ["a", "b", "c"]
def test_invalid_regex_does_not_create_output(monkeypatch, tmp_path, conversion):
source = tmp_path / "source_invalid.safetensors"
_write_uel(source, {"weight": torch.ones(1)})
_patch_runtime(monkeypatch, tmp_path, conversion, source)
with pytest.raises(ValueError, match="Invalid exclusion regex"):
conversion.convert_diffusion_model_dtype(
source.name, "fp16", "[", "must_not_exist"
)
assert not (tmp_path / "diffusion_models" / "must_not_exist.safetensors").exists()
def test_rejects_input_output_collision(monkeypatch, tmp_path, conversion):
model_dir = tmp_path / "diffusion_models"
model_dir.mkdir()
source = model_dir / "same.safetensors"
_write_uel(source, {"weight": torch.ones(1)})
_patch_runtime(monkeypatch, tmp_path, conversion, source)
with pytest.raises(ValueError, match="must differ"):
conversion.convert_diffusion_model_dtype(
source.name, "fp16", "", "same"
)
+80
View File
@@ -0,0 +1,80 @@
[groups.repository_contract]
paths = ["__init__.py", "nodes/*.py"]
python_tests = [
"tests/test_conditional_registration.py",
"tests/test_schema_tooltips.py",
"tests/test_uel_repository_contract.py",
]
[groups.artifact_paths]
paths = [
"nodes/artifact_paths.py", "nodes/dora_extract_wd.py",
"nodes/dora_learned_wd.py", "nodes/dtype_conversion.py",
"nodes/cwb_delta_lora_merger.py", "nodes/cwb_model_graft.py", "nodes/lodestone_merger.py", "nodes/lora_extract_svd.py", "nodes/lora_merger.py",
"nodes/lora_rename.py", "nodes/lora_resize.py", "nodes/merger.py",
"nodes/minimax_h3_lora_convert.py",
"nodes/minimax_h3_fold.py",
"nodes/prunekeys.py", "nodes/renamekeys.py", "nodes/text_encoder_extract.py",
]
python_tests = ["tests/test_artifact_paths.py"]
[groups.cwb]
paths = ["nodes/consensus_merger.py", "nodes/cwb_model_graft.py", "nodes/cwb_delta_lora_merger.py", "run_cwb_api_workflow.py", "run_cwb_preset_sweep.py"]
python_tests = [
"tests/test_consensus_merger.py", "tests/test_cwb_model_graft.py", "tests/test_cwb_api_runner.py",
"tests/test_cwb_delta_lora_merger.py", "tests/test_cwb_preset_sweep.py",
]
[groups.model_analysis]
paths = ["nodes/model_analysis.py", "run_model_analysis_api_workflow.py"]
python_tests = ["tests/test_model_analysis.py", "tests/test_model_analysis_api_runner.py"]
[groups.lora_model_analysis]
paths = ["nodes/lora_model_analysis.py", "nodes/lora_resize.py", "nodes/model_analysis.py"]
python_tests = ["tests/test_lora_model_analysis.py"]
[groups.adaptive_extraction]
paths = [
"nodes/adaptive_svd.py", "nodes/dora_extract_wd.py", "nodes/dora_learned_wd.py",
"nodes/extraction_stream.py", "nodes/lora_extract_svd.py",
"nodes/text_encoder_extract.py",
]
python_tests = ["tests/test_adaptive_extraction_pipeline.py", "tests/test_knee_partial_svd.py"]
[groups.lora_merge]
paths = [
"nodes/cwb_delta_lora_merger.py", "nodes/device_utils.py", "nodes/lodestone_merger.py", "nodes/lora_alpha.py",
"nodes/lora_merger.py", "nodes/quantization_guard.py", "nodes/uel_io.py",
]
python_tests = [
"tests/test_lora_diff_dtype_contract.py", "tests/test_lora_merger_async.py",
"tests/test_cwb_delta_lora_merger.py", "tests/test_lodestone_merger.py", "tests/test_lora_quantization_guard.py",
]
[groups.lora_resize]
paths = ["nodes/lora_resize.py"]
python_tests = ["tests/test_lora_resize_factor_space.py", "tests/test_lora_resize_min_rank.py"]
[groups.merger_filters]
paths = ["nodes/merger.py", "nodes/prunekeys.py"]
python_tests = ["tests/test_merger_filters.py"]
[groups.layer_parameters]
paths = ["nodes/layer_parameters.py", "nodes/merger.py", "nodes/lora_extract_svd.py", "nodes/dora_extract_wd.py", "nodes/dora_learned_wd.py", "nodes/text_encoder_extract.py", "nodes/lora_resize.py", "__init__.py"]
python_tests = ["tests/test_layer_parameters.py", "tests/test_layer_parameter_receivers.py", "tests/test_layer_parameter_extraction.py", "tests/test_layer_parameter_resize.py", "tests/test_merger_filters.py"]
[groups.layer_parameter_syntax]
paths = ["nodes/layer_parameters.py"]
python_tests = ["tests/test_layer_parameters.py"]
[groups.dtype_conversion]
paths = ["nodes/dtype_conversion.py"]
python_tests = ["tests/test_dtype_conversion.py"]
[groups.minimax_h3_lora_convert]
paths = ["nodes/minimax_h3_lora_convert.py"]
python_tests = ["tests/test_minimax_h3_lora_convert.py"]
[groups.minimax_h3_fold]
paths = ["nodes/minimax_h3_fold.py"]
python_tests = ["tests/test_minimax_h3_fold.py"]
+101
View File
@@ -0,0 +1,101 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_knee_partial_svd_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module", params=["nodes.lora_extract_svd", "nodes.dora_extract_wd"])
def extraction(request):
return _load_module(request.param)
def _fake_factors(weight, q, niter):
del niter
return (
torch.eye(weight.shape[0], q),
torch.linspace(float(q), 1.0, q),
torch.eye(weight.shape[1], q),
)
def test_knee_probe_uses_offset_and_expands_at_boundary(monkeypatch, extraction):
probes = []
detected = iter([9, 5])
def fake_svd(weight, q, niter):
probes.append(q)
return _fake_factors(weight, q, niter)
monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd)
monkeypatch.setattr(
extraction, "_compute_rank", lambda singular_values, mode, param, cap: next(detected)
)
_, _, _, rank = extraction._svd_extract_knee_lowrank(
torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2
)
assert probes == [10, 12]
assert rank == 5
def test_knee_probe_stops_when_knee_is_clear(monkeypatch, extraction):
probes = []
def fake_svd(weight, q, niter):
probes.append(q)
return _fake_factors(weight, q, niter)
monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd)
monkeypatch.setattr(extraction, "_compute_rank", lambda *args: 3)
_, _, _, rank = extraction._svd_extract_knee_lowrank(
torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2
)
assert probes == [10]
assert rank == 3
def test_lowrank_factor_construction_does_not_build_residual(monkeypatch, extraction):
monkeypatch.setattr(extraction.torch, "svd_lowrank", _fake_factors)
result, mode = extraction._svd_extract_linear_lowrank(
torch.eye(6), rank=3, device="cpu", clamp_quantile=1.0, niter=2
)
down, up, residual = result
assert mode == "low rank"
assert residual is None
assert torch.equal(up, torch.eye(6, 3) * torch.tensor([3.0, 2.0, 1.0]))
assert torch.equal(down, torch.eye(3, 6))
def test_knee_control_is_adjacent_to_method():
modules_and_nodes = [
(_load_module("nodes.lora_extract_svd"), "LoRAExtractKnee"),
(_load_module("nodes.dora_extract_wd"), "DoRAExtractKnee"),
(_load_module("nodes.dora_learned_wd"), "DoRALearnedExtractKnee"),
(_load_module("nodes.text_encoder_extract"), "TextEncoderLoRAExtractKnee"),
(_load_module("nodes.text_encoder_extract"), "TextEncoderDoRAExtractKnee"),
]
for module, node_name in modules_and_nodes:
names = [item.id for item in getattr(module, node_name).define_schema().inputs]
method_index = names.index("knee_method")
assert names[method_index + 1] == "knee_probe_offset"
+194
View File
@@ -0,0 +1,194 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_layer_parameter_extraction_tests"
def _load(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
def _write(path, tensors):
with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
writer.write_dict(tensors)
def test_fixed_extraction_rule_overrides_only_matching_layer(monkeypatch, tmp_path):
extraction = _load("nodes.lora_extract_svd")
parameters = _load("nodes.layer_parameters")
first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
match = "blocks.45.attn.qkv_proj.weight"
other = "blocks.46.attn.qkv_proj.weight"
_write(first, {match: torch.eye(3), other: torch.eye(3)})
_write(second, {match: torch.zeros(3, 3), other: torch.zeros(3, 3)})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
rules = parameters.parse_rules(r"(blocks\.45\.attn\.qkv_proj) linear_dim:2 clamp_quantile:1")
extraction.extract_lora_from_files(
str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
layer_parameters=rules,
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
assert result.get_shape("diffusion_model.blocks.45.attn.qkv_proj.lora_A.weight")[0] == 2
assert result.get_shape("diffusion_model.blocks.46.attn.qkv_proj.lora_A.weight")[0] == 1
def test_unmatched_fixed_rule_preserves_existing_rank_cap(monkeypatch, tmp_path):
extraction = _load("nodes.lora_extract_svd")
first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
key = "blocks.45.attn.qkv_proj.weight"
_write(first, {key: torch.eye(3)})
_write(second, {key: torch.zeros(3, 3)})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
caps = []
def extract_linear(weight, mode, parameter, device, max_rank, clamp, niter, probe):
caps.append((parameter, max_rank))
return (torch.ones((1, weight.shape[1])), torch.ones((weight.shape[0], 1)), 1), "lora"
monkeypatch.setattr(extraction, "_svd_extract_linear", extract_linear)
extraction.extract_lora_from_files(
str(first), str(second), "fixed", 2, 2, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
)
assert caps == [(2, 1)]
def test_adaptive_rule_and_cpu_retry_use_resolved_values(monkeypatch, tmp_path):
extraction = _load("nodes.lora_extract_svd")
parameters = _load("nodes.layer_parameters")
first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
key = "blocks.45.attn.qkv_proj.weight"
_write(first, {key: torch.eye(3)})
_write(second, {key: torch.zeros(3, 3)})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
calls = []
def retry(cuda_call, cpu_call, _device):
return cpu_call()
def extract_linear(weight, mode, parameter, device, max_rank, clamp, niter, probe):
calls.append((mode, parameter, device, max_rank, clamp, niter, probe))
return (torch.ones((1, weight.shape[1])), torch.ones((weight.shape[0], 1)), 1), "lora"
monkeypatch.setattr(extraction, "retry_cuda_oom_on_cpu", retry)
monkeypatch.setattr(extraction, "_svd_extract_linear", extract_linear)
rules = parameters.parse_rules(r"(blocks\.45) linear_ratio:3 linear_max_rank:2 clamp_quantile:1")
extraction.extract_lora_from_files(
str(first), str(second), "ratio", 2, 2, "cuda", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
layer_parameters=rules,
)
assert calls == [("ratio", 3, "cpu", 2, 1, 2, 32)]
def test_rule_min_diff_gates_matching_layer_before_svd(monkeypatch, tmp_path):
extraction = _load("nodes.lora_extract_svd")
parameters = _load("nodes.layer_parameters")
first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
key = "blocks.45.attn.qkv_proj.weight"
_write(first, {key: torch.eye(2) * 0.5})
_write(second, {key: torch.zeros(2, 2)})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
calls = []
monkeypatch.setattr(extraction, "_svd_extract_linear", lambda *args: calls.append(args))
rules = parameters.parse_rules(r"(blocks\.45) min_diff:1")
extraction.extract_lora_from_files(
str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
layer_parameters=rules,
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
assert list(result.keys()) == []
assert calls == []
def test_chunked_fallback_receives_resolved_clamp(monkeypatch, tmp_path):
extraction = _load("nodes.lora_extract_svd")
parameters = _load("nodes.layer_parameters")
first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
key = "blocks.45.attn.qkv_proj.weight"
_write(first, {key: torch.eye(6, 2)})
_write(second, {key: torch.zeros(6, 2)})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
monkeypatch.setattr(extraction, "_svd_extract_linear", lambda *args: (_ for _ in ()).throw(RuntimeError("force chunk")))
monkeypatch.setattr(extraction, "_detect_fused_layer", lambda *_: 3)
observed = []
def chunk(weight, chunks, mode, parameter, device, rank, probe, clamp):
observed.append((mode, parameter, rank, probe, clamp))
return torch.ones((weight.shape[0], 1)), torch.ones((1, weight.shape[1])), 1
monkeypatch.setattr(extraction, "_extract_chunked_layer", chunk)
rules = parameters.parse_rules(r"(blocks\.45) linear_dim:2 clamp_quantile:1")
extraction.extract_lora_from_files(
str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
layer_parameters=rules,
)
assert observed == [("fixed", 2, 2, 32, 1)]
@pytest.mark.parametrize(
("module_name", "function_name", "extra"),
[
("nodes.dora_extract_wd", "extract_dora_from_files", {}),
("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": False}),
("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": True}),
],
)
def test_fixed_rule_changes_effective_rank_in_every_other_extraction_family(
monkeypatch, tmp_path, module_name, function_name, extra
):
module = _load(module_name)
parameters = _load("nodes.layer_parameters")
first, second, output = (tmp_path / name for name in ("first.safetensors", "second.safetensors", "out.safetensors"))
key = "blocks.45.attn.qkv_proj.weight"
_write(first, {key: torch.eye(3)})
_write(second, {key: torch.zeros(3, 3)})
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
rules = parameters.parse_rules(r"(blocks\.45\.attn\.qkv_proj) a:2")
kwargs = dict(linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
layer_parameters=rules, **extra)
getattr(module, function_name)(
str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(output), **kwargs
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
factor_key = next(key for key in result.keys() if key.endswith((".lora_A.weight", ".lora_down.weight")))
assert result.get_shape(factor_key)[0] == 2
@pytest.mark.parametrize(
("module_name", "nodes"),
[
("nodes.lora_extract_svd", ("LoRAExtractFixed", "LoRAExtractRatio", "LoRAExtractQuantile", "LoRAExtractKnee", "LoRAExtractFrobenius")),
("nodes.dora_extract_wd", ("DoRAExtractFixed", "DoRAExtractRatio", "DoRAExtractQuantile", "DoRAExtractKnee", "DoRAExtractFrobenius")),
("nodes.dora_learned_wd", ("DoRALearnedExtractFixed", "DoRALearnedExtractRatio", "DoRALearnedExtractQuantile", "DoRALearnedExtractKnee", "DoRALearnedExtractFrobenius")),
("nodes.text_encoder_extract", ("TextEncoderLoRAExtractFixed", "TextEncoderLoRAExtractRatio", "TextEncoderLoRAExtractQuantile", "TextEncoderLoRAExtractKnee", "TextEncoderLoRAExtractFrobenius", "TextEncoderDoRAExtractFixed", "TextEncoderDoRAExtractRatio", "TextEncoderDoRAExtractQuantile", "TextEncoderDoRAExtractKnee", "TextEncoderDoRAExtractFrobenius")),
],
)
def test_all_extraction_schemas_append_layer_parameter_socket(module_name, nodes):
module = _load(module_name)
for node_name in nodes:
assert module.__dict__[node_name].define_schema().inputs[-1].id == "layer_parameters"
+112
View File
@@ -0,0 +1,112 @@
import importlib
import inspect
import sys
import types
from pathlib import Path
import pytest
from comfy_api.latest import io
RECEIVERS = [
("merger", 10), ("lora_extract_svd", 5), ("dora_extract_wd", 5),
("dora_learned_wd", 5), ("text_encoder_extract", 10), ("lora_resize", 4),
]
def load(name):
package_name = "modelutils_layer_receiver_tests"
if package_name not in sys.modules:
package = types.ModuleType(package_name)
package.__path__ = [str(Path(__file__).resolve().parents[1])]
sys.modules[package_name] = package
return importlib.import_module(f"{package_name}.nodes.{name}")
def nodes(module):
return [
node for name, node in vars(module).items()
if isinstance(node, type) and not name.startswith("_")
and node.__module__ == module.__name__ and issubclass(node, io.ComfyNode)
and (not module.__name__.endswith(".lora_resize") or name.startswith("LoRAResize"))
]
@pytest.mark.parametrize("module_name,count", RECEIVERS)
def test_all_39_receivers_append_optional_typed_socket(module_name, count):
module = load(module_name)
classes = nodes(module)
assert len(classes) == count
config_output = load("layer_parameters").LayerParameterConfiguration.define_schema().outputs[0]
for node in classes:
schema = node.define_schema()
input_ids = [item.id for item in schema.inputs]
assert input_ids[-2:] == ["include_mode", "layer_parameters"], schema.node_id
assert input_ids.count("layer_parameters") == 1
socket = schema.inputs[-1]
assert socket.optional
assert socket.Parent.io_type == config_output.Parent.io_type == "MODELUTILS_LAYER_PARAMETERS"
assert "a=" in socket.tooltip
assert inspect.signature(node.execute).parameters["layer_parameters"].default is None
@pytest.mark.parametrize("module_name,count", RECEIVERS)
def test_all_39_receivers_forward_configuration(monkeypatch, module_name, count):
module = load(module_name)
rules = load("layer_parameters").parse_rules("(layer) a:2")
captured = []
helper_names = {
"lora_extract_svd": "extract_lora_from_files",
"dora_extract_wd": "extract_dora_from_files",
"dora_learned_wd": "extract_dora_learned_from_files",
"text_encoder_extract": "extract_te_from_files",
"lora_resize": "resize_lora_file",
}
if module_name == "merger":
def merge(model_names, calc_mode, all_modes, params, model_type):
captured.append(params["layer_parameters"])
return "result.safetensors"
monkeypatch.setattr(module.MergerLogic, "execute_merge", merge)
else:
name = helper_names[module_name]
signature = inspect.signature(getattr(module, name))
def process(*args, **kwargs):
captured.append(signature.bind(*args, **kwargs).arguments["layer_parameters"])
return "result.safetensors"
monkeypatch.setattr(module, name, process)
monkeypatch.setattr(module.folder_paths, "get_full_path_or_raise", lambda category, name: name)
if module_name in {"lora_extract_svd", "dora_extract_wd", "dora_learned_wd", "text_encoder_extract"}:
monkeypatch.setattr(module, "_build_lora_output_path", lambda name: (name, name))
if module_name == "lora_resize":
monkeypatch.setattr(module, "canonical_model_artifact_path", lambda category, name: (name, name))
for node in nodes(module):
values = {}
for item in node.define_schema().inputs:
value = getattr(item, "default", None)
if value is None:
options = getattr(item, "options", [])
value = options[0] if options else item.id
values[item.id] = value
values["layer_parameters"] = rules
if "execution_mode" in values:
values["execution_mode"] = "MERGE"
node.execute(**values)
assert len(captured) == count
assert all(value is rules for value in captured)
def test_no_expansion_to_other_node_families():
for module_name in ("consensus_merger", "cwb_model_graft", "cwb_delta_lora_merger", "lodestone_merger", "lora_merger", "model_analysis", "dtype_conversion"):
module = load(module_name)
for node in nodes(module):
assert "layer_parameters" not in {item.id for item in node.define_schema().inputs}
resize = load("lora_resize")
for node in (
resize.LoRAMergeToModel,
resize.LoRANormalizeAlpha,
resize.LoRAModifyAlpha,
resize.LoRADenormalizeAlpha,
):
assert "layer_parameters" not in {item.id for item in node.define_schema().inputs}
+258
View File
@@ -0,0 +1,258 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_layer_parameter_resize_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def resize():
return _load_module("nodes.lora_resize")
@pytest.fixture(scope="module")
def layer_parameters():
return _load_module("nodes.layer_parameters")
def _write_uel(path, tensors):
with IncrementalSafetensorsWriter(str(path), metadata={}) as writer:
writer.write_batch(
[(key, tensor.cpu().contiguous()) for key, tensor in tensors.items()]
)
def _read_uel(path):
result = {}
with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
stream = loader.async_stream(
loader.keys(), batch_size=1, prefetch_batches=1, pin_memory=False
)
for batch in stream:
for key, tensor in batch:
result[key] = tensor.clone()
loader.mark_processed(key)
return result
@pytest.mark.parametrize(
("node_name", "extra"),
[
("LoRAResizeFixed", {"new_rank": 1}),
("LoRAResizeRatio", {"max_rank": 1, "ratio": 2.0}),
(
"LoRAResizeFrobenius",
{"max_rank": 1, "min_rank": 1, "target": 0.9},
),
(
"LoRAResizeCumulative",
{"max_rank": 1, "target": 0.9},
),
],
)
def test_resize_nodes_append_and_forward_layer_parameters(
monkeypatch, tmp_path, resize, node_name, extra
):
node = getattr(resize, node_name)
schema = node.define_schema()
assert schema.inputs[-1].id == "layer_parameters"
calls = []
layer_rules = object()
monkeypatch.setattr(
resize.folder_paths, "get_full_path_or_raise", lambda *args: "input"
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(
resize, "resize_lora_file", lambda *args, **kwargs: calls.append(kwargs)
)
node.execute(
lora_name="input",
output_filename="out",
save_dtype="fp16",
device="cpu",
force_clear_cache=False,
layer_parameters=layer_rules,
**extra,
)
assert calls[0]["layer_parameters"] is layer_rules
def test_fixed_resize_rule_applies_only_to_matched_canonical_layer(
monkeypatch, tmp_path, resize, layer_parameters
):
source = tmp_path / "source.safetensors"
tensors = {}
for name in ("matched", "unmatched"):
tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3)
tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3)
_write_uel(source, tensors)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.resize_lora_file(
str(source),
1,
None,
None,
"cpu",
torch.float32,
"resized",
verbose=False,
layer_parameters=layer_parameters.parse_rules(r"(diffusion_model\.matched) a:2"),
)
actual = _read_uel(output)
assert actual["diffusion_model.matched.lora_A.weight"].shape[0] == 2
assert actual["diffusion_model.unmatched.lora_A.weight"].shape[0] == 1
@pytest.mark.parametrize(
("method", "parameter", "max_rank", "min_rank", "rule", "expected_rank"),
[
(None, None, 1, 1, "a:2", 2),
("sv_ratio", 2.0, 1, 1, "a:3 b:100", 3),
("sv_fro", 0.1, 1, 1, "a:3 b:2 c:0.1", 2),
("sv_cumulative", 0.1, 1, 1, "a:3 b:0.9", 2),
],
)
def test_resize_rule_overrides_each_profile_numerically(
monkeypatch,
tmp_path,
resize,
layer_parameters,
method,
parameter,
max_rank,
min_rank,
rule,
expected_rank,
):
source = tmp_path / f"{method or 'fixed'}.safetensors"
_write_uel(
source,
{
"diffusion_model.target.lora_A.weight": torch.eye(3),
"diffusion_model.target.lora_B.weight": torch.diag(
torch.tensor([10.0, 3.0, 1.0])
),
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.resize_lora_file(
str(source),
max_rank,
method,
parameter,
"cpu",
torch.float32,
f"{method or 'fixed'}_output",
verbose=False,
min_rank=min_rank,
layer_parameters=layer_parameters.parse_rules(
rf"(diffusion_model\.target) {rule}"
),
)
assert _read_uel(output)["diffusion_model.target.lora_A.weight"].shape[0] == expected_rank
def test_resize_rules_respect_include_discard_and_low_bit_precedence(
monkeypatch, tmp_path, resize, layer_parameters
):
source = tmp_path / "precedence.safetensors"
low_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
_write_uel(
source,
{
"diffusion_model.low.lora_A.weight": low_down,
"diffusion_model.low.lora_B.weight": torch.eye(2),
"diffusion_model.normal.lora_A.weight": torch.eye(3),
"diffusion_model.normal.lora_B.weight": torch.eye(3),
"diffusion_model.discard.lora_A.weight": torch.eye(3),
"diffusion_model.discard.lora_B.weight": torch.eye(3),
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.resize_lora_file(
str(source),
2,
None,
None,
"cpu",
torch.float32,
"precedence_output",
verbose=False,
exclude_patterns=r"diffusion_model\.(low|normal|discard)",
discard_patterns=r"diffusion_model\.discard",
include_mode=True,
layer_parameters=layer_parameters.parse_rules(
"\n".join(
[
r"(diffusion_model\.low) a:1",
r"(diffusion_model\.normal) a:1",
r"(diffusion_model\.discard) a:1",
]
)
),
)
actual = _read_uel(output)
assert torch.equal(actual["diffusion_model.low.lora_A.weight"], low_down)
assert actual["diffusion_model.normal.lora_A.weight"].shape[0] == 1
assert not any(".discard." in key for key in actual)
def test_frobenius_rule_validation_precedes_writer(
monkeypatch, tmp_path, resize, layer_parameters
):
source = tmp_path / "source.safetensors"
_write_uel(
source,
{
"diffusion_model.target.lora_A.weight": torch.eye(3),
"diffusion_model.target.lora_B.weight": torch.eye(3),
},
)
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
monkeypatch.setattr(
resize,
"atomic_uel_writer",
lambda *args, **kwargs: pytest.fail("Writer must not open for invalid rule"),
)
with pytest.raises(ValueError, match="min_rank=3 exceeds max_rank=2"):
resize.resize_lora_file(
str(source),
4,
"sv_fro",
0.9,
"cpu",
torch.float32,
"unused",
verbose=False,
layer_parameters=layer_parameters.parse_rules(r"(diffusion_model\.target) a:2 b:3"),
)
+165
View File
@@ -0,0 +1,165 @@
import importlib
import sys
import types
from dataclasses import FrozenInstanceError
from pathlib import Path
import pytest
@pytest.fixture(scope="module")
def rules_module():
name = "modelutils_layer_rule_tests"
package = types.ModuleType(name)
package.__path__ = [str(Path(__file__).resolve().parents[1])]
sys.modules[name] = package
return importlib.import_module(f"{name}.nodes.layer_parameters")
def resolve(module, text, *, profile="extract:fixed", names=("blocks.45.attn.qkv_proj.weight",), defaults=None, glob=False):
return module.resolve_layer_parameters(
module.parse_rules(text, glob), profile, names,
defaults or {"linear_dim": 8, "conv_dim": 4, "clamp_quantile": 0.99, "min_diff": 0},
node_name="Example Node",
)
@pytest.mark.parametrize("assignments", [
"a:64 b:32 c:0.99 d:0.0",
"a : 64, b:32 c:0.99, d : 0.0",
"linear_dim:64 b:3.2e1 clamp_quantile:+.99 d:0e0",
])
def test_aliases_separators_and_inheritance(rules_module, assignments):
module = rules_module
result = resolve(module, rf"(blocks\.4[589]\.attn\.qkv_proj) {assignments}")
assert result == {"blocks.45.attn.qkv_proj.weight": {
"linear_dim": 64, "conv_dim": 32, "clamp_quantile": 0.99, "min_diff": 0,
}}
partial = resolve(module, "(qkv_proj) a:64", names=("other", "blocks.45.attn.qkv_proj.weight"))
assert "other" not in partial
assert partial["blocks.45.attn.qkv_proj.weight"]["conv_dim"] == 4
def test_patterns_preserve_groups_spaces_escapes_and_lines(rules_module):
text = " # explanation\r\n\r\n" + r"((?:blocks\.45|blocks\.49)\.my layer\.(?:qkv_proj)) a:16"
parsed = rules_module.parse_rules(text)
assert parsed.rules[0].line == 3
assert parsed.rules[0].pattern == r"(?:blocks\.45|blocks\.49)\.my layer\.(?:qkv_proj)"
result = rules_module.resolve_layer_parameters(parsed, "extract:fixed", ["blocks.49.my layer.qkv_proj"], {}, node_name="Fixed")
assert result["blocks.49.my layer.qkv_proj"]["linear_dim"] == 16
@pytest.mark.parametrize("text", [
"() a:1", "( ) a:1", "(layer)", "layer a:1", "(layer)a:1",
"([invalid) a:1", "(layer) a:", "(layer) a:nan", "(layer) a:inf",
"(layer) a:1e309", "(layer) a:1a:2", "(layer) a:1,", "(layer) a:1,,b:2",
"(layer) a:1 # comment", "(layer) a:1 a:2", "(layer) a:1 rubbish",
])
def test_invalid_syntax_errors_at_configurator(rules_module, text):
with pytest.raises(ValueError, match="line 1"):
rules_module.LayerParameterConfiguration.execute(text)
@pytest.mark.parametrize("assignments, message", [
("a:1 alpha:2", "unsupported"),
("a:1 linear_dim:1", "duplicate aliases"),
("a:1.5", "integral"),
("a:0", "linear_dim"),
("a:16385", "linear_dim"),
("c:0.49", "clamp_quantile"),
("d:1.01", "min_diff"),
("seed:1", "unsupported"),
])
def test_receiver_validates_parameter_domains(rules_module, assignments, message):
with pytest.raises(ValueError, match=message) as error:
resolve(rules_module, f"(qkv_proj) {assignments}")
assert "Example Node" in str(error.value)
assert "line 1" in str(error.value)
def test_strict_overlap_even_for_disjoint_assignments(rules_module):
with pytest.raises(ValueError, match="conflicting") as error:
resolve(rules_module, "(blocks) a:32\n\n(qkv_proj) b:16")
assert "blocks.45.attn.qkv_proj.weight" in str(error.value)
assert "line 1" in str(error.value) and "line 3" in str(error.value)
def test_zero_match_and_repeated_logical_names(rules_module):
with pytest.raises(ValueError, match="line 2"):
resolve(rules_module, "(blocks) a:32\n(missing) b:16")
assert resolve(rules_module, "(qkv_proj) a:32", names=["qkv_proj", "qkv_proj"]) == {
"qkv_proj": {"linear_dim": 32, "conv_dim": 4, "clamp_quantile": 0.99, "min_diff": 0},
}
def test_glob_is_case_sensitive_substring(rules_module):
result = resolve(rules_module, "(blocks.4[589].attn.qkv_*) a:32", glob=True)
assert len(result) == 1
with pytest.raises(ValueError, match="no target layer"):
resolve(rules_module, "(BLOCKS*) a:32", glob=True)
def test_frobenius_validates_inherited_bounds(rules_module):
defaults = {"max_rank": 128, "min_rank": 32, "target": 0.9}
with pytest.raises(ValueError, match="min_rank=32 exceeds max_rank=16"):
resolve(rules_module, "(qkv_proj) a:16", profile="resize:frobenius", defaults=defaults)
result = resolve(rules_module, "(qkv_proj) a:16 b:1", profile="resize:frobenius", defaults=defaults)
assert next(iter(result.values())) == {"max_rank": 16, "min_rank": 1, "target": 0.9}
assert defaults["max_rank"] == 128
def test_shared_rules_do_not_leak_receiver_state(rules_module):
rules = rules_module.parse_rules("(layer) a:2")
with pytest.raises(FrozenInstanceError):
rules.glob_patterns = True
first = rules_module.resolve_layer_parameters(rules, "extract:fixed", ["layer"], {"conv_dim": 4}, node_name="A")
second = rules_module.resolve_layer_parameters(rules, "resize:fixed", ["layer"], {"new_rank": 8}, node_name="B")
first["layer"]["linear_dim"] = 64
assert second == {"layer": {"new_rank": 2}}
assert rules.rules[0].assignments == (("a", 2),)
assert rules_module.resolve_layer_parameters(None, "anything", [], {}, node_name="A") == {}
assert rules_module.resolve_layer_parameters(rules_module.parse_rules("\n#empty"), "anything", [], {}, node_name="A") == {}
with pytest.raises(TypeError, match="Layer Parameter Configuration"):
rules_module.resolve_layer_parameters("not a config", "extract:fixed", [], {}, node_name="A")
def test_fixed_dimensions_update_only_explicit_rank_caps(rules_module):
defaults = {"linear_dim": 8, "conv_dim": 4, "linear_max_rank": 2, "conv_max_rank": 1, "clamp_quantile": 0.99, "min_diff": 0}
result = resolve(rules_module, "(qkv_proj) a:8", defaults=defaults)
values = next(iter(result.values()))
assert values["linear_max_rank"] == 8
assert values["conv_max_rank"] == 1
values = next(iter(resolve(rules_module, "(qkv_proj) c:.9", defaults=defaults).values()))
assert values["linear_max_rank"] == 2
assert values["conv_max_rank"] == 1
def test_merge_aliases_mode_restrictions_and_quantiles(rules_module):
result = resolve(rules_module, "(qkv_proj) a:.5 b:.25 g:.75 d:1", profile="merge:Extract-Features", defaults={})
assert next(iter(result.values()))["gamma"] == 0.75
with pytest.raises(ValueError, match="duplicate aliases"):
resolve(rules_module, "(qkv_proj) c:.5 g:.5", profile="merge:Extract-Features")
with pytest.raises(ValueError, match="unsupported parameter"):
resolve(rules_module, "(qkv_proj) b:.5", profile="merge:Weight-Sum")
with pytest.raises(ValueError, match="unsupported parameter"):
resolve(rules_module, "(qkv_proj) a:.5", profile="merge:Power-Up Enhanced (DARE+TIES)")
with pytest.raises(ValueError, match="must be in"):
resolve(rules_module, "(qkv_proj) b:1.1", profile="merge:Power-Up (DARE+TIES)")
def test_configurator_documentation_and_socket_types(rules_module):
module = rules_module
schema = module.LayerParameterConfiguration.define_schema()
assert [item.id for item in schema.inputs] == ["rules", "glob_patterns"]
config, docs = module.LayerParameterConfiguration.execute("(layer) a:32")
assert isinstance(config, module.LayerParameterRules)
assert docs.startswith("# Per-layer parameter rules")
assert len(schema.outputs) == 2
profiles = [*(f"merge:{mode}" for mode in module.MERGE_MODES), *(f"extract:{mode}" for mode in module.EXTRACTION_BINDINGS), *(f"resize:{mode}" for mode in module.RESIZE_BINDINGS)]
for profile in profiles:
assert profile in docs
widget = module.parameter_input("merge" if profile.startswith("merge:") else profile)
assert widget.optional
for parameter in module.parameter_bindings(profile):
assert parameter.name in widget.tooltip
assert parameter.name in docs
+186
View File
@@ -0,0 +1,186 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_lodestone_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def lodestone():
return _load_module("nodes.lodestone_merger")
def _write(path, tensors):
with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
writer.write_batch(list(tensors.items()))
def _read(path):
result = {}
with MemoryEfficientSafeOpen(str(path), low_memory=True) as handler:
stream = handler.async_stream(handler.keys(), batch_size=1, prefetch_batches=1, pin_memory=False)
for batch in stream:
for key, tensor in batch:
result[key] = tensor.clone()
handler.mark_processed(key)
return result
@pytest.mark.parametrize(
("method", "expected"),
[
("sum", torch.diag(torch.tensor([2.0, 4.0]))),
("mean", torch.diag(torch.tensor([1.0, 2.0]))),
("slotnorm", torch.eye(2)),
("normmatch", torch.diag(torch.tensor([2.0, 4.0])) * (2.0 / torch.sqrt(torch.tensor(20.0)))),
("slotnorm-normmatch", torch.eye(2) * torch.sqrt(torch.tensor(2.0))),
],
)
def test_methods_match_supplied_dense_math(lodestone, method, expected):
deltas = [torch.diag(torch.tensor([2.0, 0.0])), torch.diag(torch.tensor([0.0, 4.0]))]
torch.testing.assert_close(lodestone.merge_lodestone_deltas(deltas, method), expected)
def test_convolutional_factor_expansion(lodestone):
down = torch.arange(18, dtype=torch.float32).reshape(2, 1, 3, 3)
up = torch.tensor([[[[2.0]], [[-1.0]]]])
expected = (up.reshape(1, 2) @ down.reshape(2, -1)).reshape(1, 1, 3, 3)
torch.testing.assert_close(lodestone.factor_delta(down, up, "conv"), expected)
def test_zero_cancellation_remains_zero(lodestone):
delta = torch.tensor([[1.0, -2.0]])
result = lodestone.merge_lodestone_deltas([delta, -delta], "normmatch")
assert torch.count_nonzero(result) == 0
@pytest.mark.parametrize("count", [3, 8])
def test_three_and_multi_input_dense_means(lodestone, count):
deltas = [torch.full((2, 2), float(index)) for index in range(1, count + 1)]
expected = torch.full((2, 2), (count + 1) / 2)
torch.testing.assert_close(
lodestone.merge_lodestone_deltas(deltas, "mean"), expected
)
def test_layer_plan_enforces_missing_filters_and_low_bit(lodestone):
class Handler:
def get_dtype(self, _key):
return torch.float32
handlers = [Handler(), Handler()]
roles = {"down": "layer.down", "up": "layer.up"}
infos = [
{"pairs": {"layer": roles}, "low_bit": set()},
{"pairs": {}, "low_bit": set()},
]
layer_map = {"layer": [(0, "layer")]}
with pytest.raises(ValueError, match="missing"):
lodestone._build_units(layer_map, infos, handlers, "error", "", "", False)
units, _ = lodestone._build_units(layer_map, infos, handlers, "zeros", "", "", False)
assert units[0]["missing"] == 1
excluded, _ = lodestone._build_units(layer_map, infos, handlers, "zeros", "*.diff", "", True)
assert excluded[0]["missing"] == 0
discarded, _ = lodestone._build_units(layer_map, infos, handlers, "skip", "", "*.diff", True)
assert discarded == []
infos[0]["low_bit"] = {"layer.down"}
with pytest.raises(ValueError, match="low-bit"):
lodestone._build_units(layer_map, infos, handlers, "skip", "", "", False)
def test_streamed_merge_writes_only_alpha_normalized_full_delta(monkeypatch, tmp_path, lodestone):
first = tmp_path / "first.safetensors"
second = tmp_path / "second.safetensors"
_write(first, {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [0.0]]),
"diffusion_model.layer.alpha": torch.tensor(0.5),
})
_write(second, {
"diffusion_model.layer.lora_A.weight": torch.tensor([[0.0, 1.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[0.0], [4.0]]),
})
monkeypatch.setattr(lodestone.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(lodestone, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(lodestone, "cleanup_after_operation", lambda: None)
output = lodestone.merge_lodestone_loras(
[str(first), str(second)], "sum", "cpu", torch.float32, "dense", verbose=False
)
tensors = _read(output)
assert list(tensors) == ["diffusion_model.layer.diff"]
torch.testing.assert_close(tensors["diffusion_model.layer.diff"], torch.diag(torch.tensor([1.0, 4.0])))
@pytest.mark.parametrize("patterns, glob_patterns, selected", [
(r"\.selected", False, True),
("*.selected*", True, True),
("", False, False),
("no_match", False, False),
])
def test_include_filter_preserves_nonmatches_and_discards_first(
monkeypatch, tmp_path, lodestone, patterns, glob_patterns, selected,
):
paths = [tmp_path / f"{name}.safetensors" for name in ("a", "b")]
for path, value in zip(paths, (2.0, 6.0)):
_write(path, {
f"diffusion_model.{layer}.diff": torch.full((2, 2), value)
for layer in ("selected", "other", "selected_discard")
})
monkeypatch.setattr(lodestone.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(lodestone, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(lodestone, "cleanup_after_operation", lambda: None)
output = lodestone.merge_lodestone_loras(
[str(path) for path in paths], "sum", "cpu", torch.float32, "include_filter",
exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=True,
discard_patterns="*discard*" if glob_patterns else "discard", verbose=False,
)
tensors = _read(output)
assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"}
torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 8.0 if selected else 2.0))
torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0))
def test_dedicated_nodes_do_not_modify_factor_merger_schema(lodestone):
standard = _load_module("nodes.lora_merger")
assert [node.define_schema().node_id for node in lodestone.LODESTONE_MERGER_NODES] == [
"LodestoneLoRATwoMerger", "LodestoneLoRAThreeMerger", "LodestoneLoRAMultiMerger"
]
standard_inputs = {item.id: item for item in standard.LoRAMultiMerge.define_schema().inputs}
assert standard_inputs["merge_mode"].options == ["concatenate", "weighted_sum"]
assert "calc_mode" not in standard_inputs
def test_dense_cuda_oom_retries_layer_on_cpu(monkeypatch, lodestone):
unit = {"dtypes": [torch.float32], "missing": 0, "sources": []}
calls = []
def process(_unit, _loaded, _method, _dtype, device):
calls.append(device.type)
if device.type == "cuda":
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
return torch.ones(1)
monkeypatch.setattr(lodestone, "_process_on_device", process)
monkeypatch.setattr(lodestone, "_release_cuda_oom", lambda: None)
output, fallback = lodestone._process(unit, {}, "sum", torch.float32, "cuda")
assert fallback is True
assert calls == ["cuda", "cpu"]
assert torch.equal(output, torch.ones(1))
File diff suppressed because it is too large Load Diff
+379
View File
@@ -0,0 +1,379 @@
import importlib
import sys
import types
from collections import Counter
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import (
IncrementalSafetensorsWriter,
MemoryEfficientSafeOpen,
)
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_lora_merger_async_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def merger():
return _load_module("nodes.lora_merger")
def _write_uel(path, tensors):
with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
writer.write_batch(
[(key, tensor.cpu().contiguous()) for key, tensor in tensors.items()]
)
def _read_uel(path, async_method=None, mark_method=None):
result = {}
with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
stream = (async_method or MemoryEfficientSafeOpen.async_stream)(
loader,
loader.keys(),
batch_size=1,
prefetch_batches=1,
pin_memory=False,
)
for batch in stream:
for key, tensor in batch:
result[key] = tensor.clone()
(mark_method or MemoryEfficientSafeOpen.mark_processed)(loader, key)
return result
def _patch_runtime(monkeypatch, tmp_path, merger):
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
def _mixed_sources(tmp_path):
first = tmp_path / "mixed_first.safetensors"
second = tmp_path / "mixed_second.safetensors"
first_tensors = {
"diffusion_model.small.lora_A.weight": torch.tensor([[1.0, 0.0]]),
"diffusion_model.small.lora_B.weight": torch.tensor([[1.0], [2.0]]),
"diffusion_model.small.alpha": torch.tensor(0.5),
"diffusion_model.large.lora_A.weight": torch.eye(3),
"diffusion_model.large.lora_B.weight": torch.eye(3),
}
second_tensors = {
"diffusion_model.small.lora_A.weight": torch.eye(2),
"diffusion_model.small.lora_B.weight": torch.eye(2),
"diffusion_model.large.lora_A.weight": torch.tensor([[1.0, 1.0, 1.0]]),
"diffusion_model.large.lora_B.weight": torch.tensor([[1.0], [2.0], [3.0]]),
}
_write_uel(first, first_tensors)
_write_uel(second, second_tensors)
return first, second, first_tensors, second_tensors
def test_concatenate_uses_per_layer_rank_and_alpha(monkeypatch, tmp_path, merger):
_patch_runtime(monkeypatch, tmp_path, merger)
first, second, first_tensors, second_tensors = _mixed_sources(tmp_path)
output = merger.merge_multi_loras(
[str(first), str(second)],
[2.0, -0.5],
"concatenate",
"cpu",
torch.float32,
"mixed_concat",
verbose=False,
)
tensors = _read_uel(output)
assert tensors["diffusion_model.small.lora_A.weight"].shape[0] == 3
assert tensors["diffusion_model.large.lora_A.weight"].shape[0] == 4
assert not any(key.endswith(".alpha") for key in tensors)
small_expected = (
2.0
* (0.5 / 1.0)
* first_tensors["diffusion_model.small.lora_B.weight"]
@ first_tensors["diffusion_model.small.lora_A.weight"]
- 0.5
* second_tensors["diffusion_model.small.lora_B.weight"]
@ second_tensors["diffusion_model.small.lora_A.weight"]
)
small_actual = (
tensors["diffusion_model.small.lora_B.weight"]
@ tensors["diffusion_model.small.lora_A.weight"]
)
torch.testing.assert_close(small_actual, small_expected)
@pytest.mark.parametrize("strategy", ["weighted_sum", "dare", "enhanced_dare"])
def test_nonconcatenate_variants_use_each_layers_max_rank(
monkeypatch, tmp_path, merger, strategy
):
_patch_runtime(monkeypatch, tmp_path, merger)
first, second, _, _ = _mixed_sources(tmp_path)
common = [str(first), str(second)], [1.0, 1.0]
if strategy == "weighted_sum":
output = merger.merge_multi_loras(
*common,
"weighted_sum",
"cpu",
torch.float32,
"mixed_weighted",
verbose=False,
)
elif strategy == "dare":
output = merger.merge_multi_loras_dare(
*common,
0.0,
0.0,
7,
"cpu",
torch.float32,
"mixed_dare",
verbose=False,
)
else:
output = merger.merge_multi_loras_dare_enhanced(
*common,
1.0,
1.0,
0.0,
0.0,
7,
"cpu",
torch.float32,
"mixed_enhanced",
verbose=False,
)
tensors = _read_uel(output)
assert tensors["diffusion_model.small.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.large.lora_A.weight"].shape[0] == 3
assert not any(key.endswith(".alpha") for key in tensors)
@pytest.mark.parametrize("strategy", ["weighted_sum", "dare", "enhanced_dare"])
def test_nonconcatenate_variants_apply_source_alpha_once(
monkeypatch, tmp_path, merger, strategy
):
_patch_runtime(monkeypatch, tmp_path, merger)
first = tmp_path / f"alpha_first_{strategy}.safetensors"
second = tmp_path / f"alpha_second_{strategy}.safetensors"
_write_uel(
first,
{
"diffusion_model.layer.lora_A.weight": torch.ones((1, 1)),
"diffusion_model.layer.lora_B.weight": torch.full((1, 1), 2.0),
"diffusion_model.layer.alpha": torch.tensor(0.5),
},
)
_write_uel(
second,
{
"diffusion_model.layer.lora_A.weight": torch.ones((1, 1)),
"diffusion_model.layer.lora_B.weight": torch.full((1, 1), 2.0),
},
)
common = [str(first), str(second)], [1.0, 1.0]
if strategy == "weighted_sum":
output = merger.merge_multi_loras(
*common,
"weighted_sum",
"cpu",
torch.float32,
"alpha_weighted",
verbose=False,
)
expected_down = 2.0
expected_up = 3.0
elif strategy == "dare":
output = merger.merge_multi_loras_dare(
*common,
0.0,
0.0,
9,
"cpu",
torch.float32,
"alpha_dare",
verbose=False,
)
expected_down = 1.0
expected_up = 1.5
else:
output = merger.merge_multi_loras_dare_enhanced(
*common,
1.0,
1.0,
0.0,
0.0,
9,
"cpu",
torch.float32,
"alpha_enhanced",
verbose=False,
)
expected_down = 1.0
expected_up = 1.5
tensors = _read_uel(output)
assert tensors["diffusion_model.layer.lora_A.weight"].item() == pytest.approx(
expected_down
)
assert tensors["diffusion_model.layer.lora_B.weight"].item() == pytest.approx(
expected_up
)
@pytest.mark.parametrize(
"strategy", ["concatenate", "weighted_sum", "dare", "enhanced_dare"]
)
def test_fp32_alpha_does_not_promote_bf16_factors(
monkeypatch, tmp_path, merger, strategy
):
_patch_runtime(monkeypatch, tmp_path, merger)
first = tmp_path / f"bf16_alpha_first_{strategy}.safetensors"
second = tmp_path / f"bf16_alpha_second_{strategy}.safetensors"
tensors = {
"diffusion_model.layer.lora_A.weight": torch.ones(
(1, 2), dtype=torch.bfloat16
),
"diffusion_model.layer.lora_B.weight": torch.ones(
(2, 1), dtype=torch.bfloat16
),
"diffusion_model.layer.alpha": torch.tensor(1.0, dtype=torch.float32),
}
_write_uel(first, tensors)
_write_uel(second, tensors)
common = [str(first), str(second)], [1.0, 1.0]
if strategy in {"concatenate", "weighted_sum"}:
output = merger.merge_multi_loras(
*common,
strategy,
"cpu",
torch.bfloat16,
f"bf16_alpha_{strategy}",
verbose=False,
)
elif strategy == "dare":
output = merger.merge_multi_loras_dare(
*common,
0.0,
0.0,
17,
"cpu",
torch.bfloat16,
"bf16_alpha_dare",
verbose=False,
)
else:
output = merger.merge_multi_loras_dare_enhanced(
*common,
1.0,
1.0,
0.0,
0.0,
17,
"cpu",
torch.bfloat16,
"bf16_alpha_enhanced",
verbose=False,
)
merged = _read_uel(output)
assert merged["diffusion_model.layer.lora_A.weight"].dtype == torch.bfloat16
assert merged["diffusion_model.layer.lora_B.weight"].dtype == torch.bfloat16
def test_all_normal_inputs_use_async_uel_and_release_once(
monkeypatch, tmp_path, merger
):
_patch_runtime(monkeypatch, tmp_path, merger)
first, second, _, _ = _mixed_sources(tmp_path)
original_async = merger.MemoryEfficientSafeOpen.async_stream
original_mark = merger.MemoryEfficientSafeOpen.mark_processed
async_calls = []
marked = []
def track_async(self, keys, **kwargs):
async_calls.append((self.filename, tuple(keys), kwargs, self.low_memory))
yield from original_async(self, keys, **kwargs)
def track_mark(self, key):
marked.append((self.filename, key))
return original_mark(self, key)
def reject_get_tensor(*args, **kwargs):
raise AssertionError("multi-merge used synchronous get_tensor")
monkeypatch.setattr(merger.MemoryEfficientSafeOpen, "async_stream", track_async)
monkeypatch.setattr(merger.MemoryEfficientSafeOpen, "mark_processed", track_mark)
monkeypatch.setattr(merger.MemoryEfficientSafeOpen, "get_tensor", reject_get_tensor)
output = merger.merge_multi_loras(
[str(first), str(second)],
[1.0, 1.0],
"concatenate",
"cpu",
torch.float32,
"async_only",
verbose=False,
)
assert Path(output).is_file()
assert len(async_calls) == 2
planned = Counter(
(filename, key)
for filename, keys, _, _ in async_calls
for key in keys
)
assert Counter(marked) == planned
assert all(count == 1 for count in planned.values())
for _, _, options, low_memory in async_calls:
assert low_memory is True
assert options == {
"batch_size": 1,
"prefetch_batches": 1,
"pin_memory": False,
}
def test_cuda_oom_retries_only_current_unit_on_cpu(
monkeypatch, tmp_path, merger
):
_patch_runtime(monkeypatch, tmp_path, merger)
first, second, _, _ = _mixed_sources(tmp_path)
original = merger._process_merge_unit_on_device
calls = []
def fail_cuda_then_cpu(*args):
process_device = args[-2]
calls.append(process_device.type)
if process_device.type == "cuda":
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
return original(*args)
monkeypatch.setattr(merger, "_process_merge_unit_on_device", fail_cuda_then_cpu)
monkeypatch.setattr(merger, "_release_cuda_oom_state", lambda: None)
output = merger.merge_multi_loras(
[str(first), str(second)],
[1.0, 1.0],
"concatenate",
"cuda",
torch.float32,
"oom_fallback",
verbose=False,
)
assert Path(output).is_file()
assert calls == ["cuda", "cpu", "cuda", "cpu"]
+250
View File
@@ -0,0 +1,250 @@
import csv
import importlib
from io import StringIO
import json
import logging
from pathlib import Path
import sys
import types
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter
@pytest.fixture(scope="module")
def analysis():
name = "modelutils_lora_on_model_tests"
package = types.ModuleType(name)
package.__path__ = [str(Path(__file__).resolve().parents[1])]
sys.modules[name] = package
return importlib.import_module(f"{name}.nodes.lora_model_analysis")
def write(path, tensors):
with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
writer.write_batch(list(tensors.items()))
def setup(monkeypatch, tmp_path, analysis, adapter, base=None):
paths = {"loras": tmp_path / "adapter.safetensors", "diffusion_models": tmp_path / "base.safetensors"}
write(paths["loras"], adapter)
write(paths["diffusion_models"], base if base is not None else {
"model.diffusion_model.layer.weight": torch.ones((2, 2)),
"model.diffusion_model.untouched.weight": torch.ones((2, 2)),
})
monkeypatch.setattr(analysis.folder_paths, "get_full_path_or_raise", lambda category, name: str(paths[category]))
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
return paths
def params(**overrides):
return {"execution_mode": "ANALYZE", "process_device": "cpu", "top_weight_differences": 3, "force_clear_cache": False, "exclude_patterns": "", "glob_patterns": False, "include_mode": False, "strength": 1.0, **overrides}
def pair(alpha=True):
result = {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [4.0]]),
}
if alpha:
result["diffusion_model.layer.alpha"] = torch.tensor(0.5)
return result
@pytest.mark.parametrize("alpha,strength,mae,maximum", [(True, 1.0, 2.25, 4), (False, 1.0, 4.5, 8), (True, 0.5, 1.125, 2), (True, 0.0, 0, 0)])
def test_reports_actual_applied_delta_without_saving(monkeypatch, tmp_path, analysis, alpha, strength, mae, maximum):
paths = setup(monkeypatch, tmp_path, analysis, pair(alpha))
before = {path: path.read_bytes() for path in paths.values()}
comparison, cwb, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params(strength=strength))
rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))}
row = rows["model.diffusion_model.layer.weight"]
assert float(row["mae"]) == pytest.approx(mae)
assert float(row["max"]) == maximum
assert float(rows["model.diffusion_model.untouched.weight"]["mae"]) == 0
assert "patched − original" in comparison
assert "original versus LoRA-patched" in cwb
assert {path: path.read_bytes() for path in paths.values()} == before
assert set(tmp_path.iterdir()) == set(paths.values())
def test_quantized_base_is_dequantized_and_sidecars_are_not_analyzed(monkeypatch, tmp_path, analysis):
prefix = "model.diffusion_model.layer"
quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
setup(monkeypatch, tmp_path, analysis, pair(), {
f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8),
f"{prefix}.weight_scale": torch.tensor(2.0),
f"{prefix}.comfy_quant": quant,
})
comparison, _, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params())
rows = list(csv.DictReader(StringIO(metrics_csv)))
assert [row["key"] for row in rows] == [f"{prefix}.weight"]
assert float(rows[0]["mae"]) == pytest.approx(2.25)
assert "weight_scale" not in comparison
@pytest.mark.parametrize("suffixes", [(".lora_down.weight", ".lora_up.weight"), (".lora_A.default.weight", ".lora_B.default.weight"), (".lora_A", ".lora_B")])
def test_pair_formats_and_flattened_reference_mapping(monkeypatch, tmp_path, analysis, suffixes):
setup(monkeypatch, tmp_path, analysis, {
"lora_unet_layer" + suffixes[0]: torch.ones((1, 2)),
"lora_unet_layer" + suffixes[1]: torch.ones((2, 1)),
})
_, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
rows = list(csv.DictReader(StringIO(text)))
assert float(rows[0]["mae"]) == 1
def test_direct_weight_bias_replacement_and_unmapped_are_explicit(monkeypatch, tmp_path, analysis):
setup(monkeypatch, tmp_path, analysis, {
"diffusion_model.layer.diff": torch.full((2, 2), 2.0),
"diffusion_model.layer.diff_b": torch.full((2,), 3.0),
"diffusion_model.replaced.set_weight": torch.full((2, 2), 5.0),
"diffusion_model.missing.diff": torch.ones((2, 2)),
}, {
"diffusion_model.layer.weight": torch.ones((2, 2)),
"diffusion_model.layer.bias": torch.ones(2),
"diffusion_model.replaced.weight": torch.ones((2, 2)),
})
report, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params(strength=0.5))
rows = {row["key"]: row for row in csv.DictReader(StringIO(text))}
assert float(rows["diffusion_model.layer.weight"]["mae"]) == 1
assert float(rows["diffusion_model.layer.bias"]["mae"]) == 1.5
assert float(rows["diffusion_model.replaced.weight"]["mae"]) == 4
assert "diffusion_model.missing" in report
def test_bounded_stream_release_and_filtering(monkeypatch, tmp_path, analysis):
setup(monkeypatch, tmp_path, analysis, pair())
calls, released = [], []
original_stream = analysis.MemoryEfficientSafeOpen.async_stream
original_mark = analysis.MemoryEfficientSafeOpen.mark_processed
def stream(handler, keys, **kwargs):
calls.append((list(keys), kwargs))
return original_stream(handler, keys, **kwargs)
def mark(handler, key):
released.append(key)
return original_mark(handler, key)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", mark)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "get_tensor", lambda *args: pytest.fail("synchronous tensor load"))
_, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params(exclude_patterns="*layer.weight", glob_patterns=True, include_mode=True))
assert len(list(csv.DictReader(StringIO(text)))) == 1
assert len(calls) == 2
assert all(kwargs == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, kwargs in calls)
assert sorted(released) == sorted(key for keys, _ in calls for key in keys)
assert len(released) == 4
def test_logged_adapter_failure_is_not_reported_as_zero_delta(monkeypatch, tmp_path, analysis):
setup(monkeypatch, tmp_path, analysis, pair())
marks = []
original_mark = analysis.MemoryEfficientSafeOpen.mark_processed
def mark(handler, key):
marks.append(key)
return original_mark(handler, key)
def fail(self, weight, key, *args):
logging.error("ERROR LoRA %s deliberately failed", key)
return weight
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", mark)
monkeypatch.setattr(analysis.LoRAAdapter, "calculate_weight", fail)
with pytest.raises(RuntimeError, match="deliberately failed"):
analysis.analyze_lora_on_model("adapter", "base", params())
assert len(marks) == 4
assert not any(isinstance(handler, analysis.AdapterErrorCapture) for handler in logging.getLogger().handlers)
def test_complete_unit_cuda_retry(monkeypatch, analysis):
devices = []
def calculate(key, loaded, specs, strength, device, top_count):
devices.append(device)
if device == "cuda":
raise torch.cuda.OutOfMemoryError("test")
return analysis.RawStats(), analysis.CWBStats(), []
monkeypatch.setattr(analysis, "_apply_and_analyze", calculate)
monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
assert analysis._analyze_unit("layer", {}, [], 1, "cuda", 0)[-1] is True
assert devices == ["cuda", "cpu"]
def test_dora_uses_reconstructed_weight_not_raw_factor_delta(monkeypatch, tmp_path, analysis):
setup(monkeypatch, tmp_path, analysis, {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [0.0]]),
"diffusion_model.layer.alpha": torch.tensor(0.5),
"diffusion_model.layer.dora_scale": torch.tensor([3.0, 4.0]),
}, {"diffusion_model.layer.weight": torch.eye(2)})
_, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
row = next(csv.DictReader(StringIO(text)))
# Installed Comfy output-axis DoRA scales by the original base row norms.
# The resulting diagonal is (4.5, 4), not the raw-factor result (1.5, 1).
assert float(row["mae"]) == pytest.approx(1.625)
assert float(row["max"]) == pytest.approx(3.5)
def test_locon_mid_and_scalar_low_bit_alpha(monkeypatch, tmp_path, analysis):
setup(monkeypatch, tmp_path, analysis, {
"diffusion_model.layer.lora_down.weight": torch.ones((1, 1, 1, 1)),
"diffusion_model.layer.lora_up.weight": torch.ones((1, 1, 1, 1)),
"diffusion_model.layer.lora_mid.weight": torch.ones((1, 1, 3, 3)),
"diffusion_model.layer.alpha": torch.tensor(1, dtype=torch.uint8),
}, {"diffusion_model.layer.weight": torch.zeros((1, 1, 3, 3))})
_, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
row = next(csv.DictReader(StringIO(text)))
assert float(row["mae"]) == 1
def test_low_bit_factors_reported_without_materializing(monkeypatch, tmp_path, analysis):
tensors = pair()
tensors["diffusion_model.layer.lora_A.weight"] = torch.ones((1, 2), dtype=torch.uint8)
setup(monkeypatch, tmp_path, analysis, tensors)
requested = []
original = analysis.MemoryEfficientSafeOpen.async_stream
def stream(handler, keys, **kwargs):
requested.extend(keys)
return original(handler, keys, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
report, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params())
assert requested == ["model.diffusion_model.untouched.weight"]
assert "model.diffusion_model.layer.weight" in report.replace(r"\_", "_")
assert [row["key"] for row in csv.DictReader(StringIO(text))] == ["model.diffusion_model.untouched.weight"]
def test_streams_close_after_failure(monkeypatch, tmp_path, analysis):
setup(monkeypatch, tmp_path, analysis, pair())
streams_closed = []
original = analysis.MemoryEfficientSafeOpen.async_stream
def stream(handler, keys, **kwargs):
iterator = original(handler, keys, **kwargs)
try:
yield from iterator
finally:
iterator.close()
streams_closed.append(True)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
monkeypatch.setattr(analysis.LoRAAdapter, "calculate_weight", lambda *args: (_ for _ in ()).throw(ValueError("test failure")))
with pytest.raises(ValueError, match="test failure"):
analysis.analyze_lora_on_model("adapter", "base", params())
assert len(streams_closed) == 2
def test_schema_and_documentation_mode(monkeypatch, analysis):
monkeypatch.setattr(analysis.folder_paths, "get_filename_list", lambda category: [category])
schema = analysis.LoRAOnModelAnalysis.define_schema()
assert schema.inputs[1].id == "model_a" and schema.inputs[1].options == ["loras"]
assert schema.inputs[2].id == "model_b" and schema.inputs[2].options == ["diffusion_models"]
assert len(schema.outputs) == 5
assert schema.inputs[-1].id == "strength" and schema.inputs[-1].default == 1
monkeypatch.setattr(analysis, "analyze_lora_on_model", lambda *args: pytest.fail("documentation loaded tensors"))
assert analysis.LoRAOnModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")[2].startswith("# LoRA-on-model analysis")
+645
View File
@@ -0,0 +1,645 @@
import importlib
import json
import logging
import struct
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
from unifiedefficientloader import IncrementalSafetensorsWriter
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_quant_guard_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def modules():
return {
"guard": _load_module("nodes.quantization_guard"),
"extract": _load_module("nodes.lora_extract_svd"),
"resize": _load_module("nodes.lora_resize"),
"multi": _load_module("nodes.lora_merger"),
"generic": _load_module("nodes.merger"),
"operations": _load_module("nodes.merger_ops"),
}
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _inspect_file(path, modules, label="input"):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
return modules["guard"].inspect_low_bit_input(handler, label, "Guard Test")
finally:
handler.__exit__(None, None, None)
def _round_trip_preserved(path, output, modules):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
with IncrementalSafetensorsWriter(str(output)) as writer:
modules["guard"].write_preserved_tensor(writer, "low", handler)
finally:
handler.__exit__(None, None, None)
return load_file(str(output))["low"]
def _tensor_payload(path, key="low"):
with open(path, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
header = json.loads(source.read(header_size).decode("utf-8"))
start, end = header[key]["data_offsets"]
source.seek(8 + header_size + start)
return source.read(end - start)
@pytest.mark.parametrize("dtype", [torch.int8, torch.uint8])
def test_int8_uint8_are_detected_and_warn_once(tmp_path, modules, caplog, dtype):
path = tmp_path / f"{dtype}.safetensors"
save_file({"low": torch.arange(4, dtype=dtype), "normal": torch.ones(1)}, str(path))
with caplog.at_level(logging.WARNING):
detected = _inspect_file(path, modules, "threshold input")
assert detected == {"low"}
records = [record for record in caplog.records if "threshold input" in record.message]
assert len(records) == 1
assert "1 isolated low-bit tensor" in records[0].message
assert "low" in records[0].message
def test_multiple_low_bit_config_tensors_allowed(tmp_path, modules):
for count in (0, 1, 2, 52):
path = tmp_path / f"threshold_{count}.safetensors"
tensors = {
f"diffusion_model.blocks.{i}.attn.comfy_attention.config": torch.ones(2, dtype=torch.uint8)
for i in range(count)
}
tensors["normal"] = torch.ones(1)
save_file(tensors, str(path))
assert len(_inspect_file(path, modules)) == count
@pytest.mark.parametrize(
("tensors", "metadata"),
[
({"layer.comfy_quant": torch.tensor([123], dtype=torch.uint8)}, None),
({"normal": torch.ones(1)}, {"_quantization_metadata": "{}"}),
({"scaled_fp8": torch.ones(2)}, None),
],
)
def test_recognized_comfy_quantization_hard_errors(tmp_path, modules, tensors, metadata):
path = tmp_path / f"quant_{len(list(tmp_path.iterdir()))}.safetensors"
save_file(tensors, str(path), metadata=metadata)
with pytest.raises(ValueError, match="ComfyUI quantization metadata.*unsupported"):
_inspect_file(path, modules, "recognized input")
def test_supported_fp8_storage_round_trip(tmp_path, modules):
candidates = [
dtype
for name in (
"float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz",
"float8_e8m0fnu",
)
if (dtype := getattr(torch, name, None)) is not None
]
supported = []
for index, dtype in enumerate(candidates):
path = tmp_path / f"low_float_{index}.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
continue
supported.append(dtype)
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / f"preserved_{index}.safetensors", modules)
assert preserved.dtype == dtype
torch.testing.assert_close(preserved.float(), loaded.float())
if candidates and not supported:
pytest.skip("Installed safetensors cannot write any installed torch FP8 dtype")
def test_fp4_storage_when_supported(tmp_path, modules):
dtype = getattr(torch, "float4_e2m1fn_x2", None)
if dtype is None:
pytest.skip("Installed PyTorch has no torch.float4_e2m1fn_x2")
path = tmp_path / "fp4.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
pytest.skip("Installed safetensors cannot write torch.float4_e2m1fn_x2")
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / "preserved_fp4.safetensors", modules)
assert preserved.dtype == dtype
assert _tensor_payload(path) == _tensor_payload(tmp_path / "preserved_fp4.safetensors")
def test_realistic_packed_uint8_comfy_sidecar_is_rejected(tmp_path, modules):
path = tmp_path / "packed_comfy.safetensors"
save_file({
"layer.weight": torch.arange(8, dtype=torch.uint8).reshape(2, 4),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
"layer.weight_scale": torch.ones(2),
}, str(path))
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
_inspect_file(path, modules)
def test_resize_preserves_complete_low_bit_layer_and_resizes_normal(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "resize_quant.safetensors"
guarded_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"diffusion_model.guarded.lora_down.weight": guarded_down,
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.guarded.alpha": torch.tensor(2.0),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.normal.lora_up.weight": torch.ones((2, 2)),
"isolated_aux": torch.tensor([7], dtype=torch.int8),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resize_guarded", verbose=False
)
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], guarded_down)
assert result["diffusion_model.guarded.lora_A.weight"].dtype == torch.uint8
assert result["diffusion_model.guarded.lora_B.weight"].dtype == torch.float32
assert result["diffusion_model.guarded.dora_scale"].dtype == torch.float32
assert result["diffusion_model.normal.lora_A.weight"].shape[0] == 1
assert result["isolated_aux"].dtype == torch.int8
def test_extract_excludes_low_bit_candidate_and_extracts_float_layer(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
model_a = tmp_path / "extract_a.safetensors"
model_b = tmp_path / "extract_b.safetensors"
output = tmp_path / "extracted.safetensors"
save_file({
"model.diffusion_model.normal.weight": torch.eye(2),
"model.diffusion_model.quant.weight": torch.ones((2, 2), dtype=torch.uint8),
}, str(model_a))
save_file({
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
"model.diffusion_model.quant.weight": torch.zeros((2, 2), dtype=torch.uint8),
}, str(model_b))
extract.extract_lora_from_files(
str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
force_clear_cache=False,
)
result = load_file(str(output))
assert any("normal" in key for key in result)
assert not any("quant" in key for key in result)
def test_extract_bias_diff_round_trips_through_merge_to_model(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
resize = modules["resize"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
_patch_output(monkeypatch, resize, tmp_path)
finetuned = tmp_path / "bias_finetuned.safetensors"
base = tmp_path / "bias_base.safetensors"
adapter = tmp_path / "bias_adapter.safetensors"
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.25, -0.5]),
}, str(finetuned))
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.0, 0.5]),
}, str(base))
extract.extract_lora_from_files(
str(finetuned), str(base), "fixed", 1, 1, "cpu", "fp16", str(adapter),
force_clear_cache=False, include_1d_diffs=True,
)
extracted = load_file(str(adapter))
assert set(extracted) == {"diffusion_model.layer.diff_b"}
assert extracted["diffusion_model.layer.diff_b"].dtype == torch.float32
merged_path = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"bias_round_trip", verbose=False, include_1d_diffs=True,
)
merged = load_file(merged_path)
torch.testing.assert_close(
merged["model.diffusion_model.layer.bias"],
load_file(str(finetuned))["model.diffusion_model.layer.bias"],
)
def _run_multi_variant(module, variant, paths, output_name):
common = (paths, [1.0, 1.0])
if variant == "standard":
return module.merge_multi_loras(
*common, "weighted_sum", "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
if variant == "dare":
return module.merge_multi_loras_dare(
*common, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
return module.merge_multi_loras_dare_enhanced(
*common, 1.0, 1.0, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("variant", ["standard", "dare", "enhanced"])
def test_multi_variants_preserve_earliest_complete_affected_layer(
monkeypatch, tmp_path, modules, variant
):
multi = modules["multi"]
_patch_output(monkeypatch, multi, tmp_path)
first = tmp_path / f"{variant}_first.safetensors"
second = tmp_path / f"{variant}_second.safetensors"
first_down = torch.tensor([[1.0, 2.0]])
save_file({
"diffusion_model.guarded.lora_down.weight": first_down,
"diffusion_model.guarded.lora_up.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(first))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 2]], dtype=torch.uint8),
"diffusion_model.guarded.lora_up.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([9.0, 9.0]),
"diffusion_model.normal.diff": torch.full((2, 2), 3.0),
}, str(second))
output = _run_multi_variant(multi, variant, [str(first), str(second)], f"guard_{variant}")
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], first_down)
torch.testing.assert_close(
result["diffusion_model.guarded.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
torch.testing.assert_close(
result["diffusion_model.guarded.dora_scale"], torch.tensor([0.125, 0.25])
)
assert "diffusion_model.normal.diff" in result
def test_resize_preserves_multiple_low_bit_config_tensors(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "three_low.safetensors"
tensors = {
f"diffusion_model.blocks.{i}.attn.comfy_attention.config": torch.ones(1, dtype=torch.uint8)
for i in range(3)
}
tensors["diffusion_model.normal.lora_down.weight"] = torch.ones((2, 2))
tensors["diffusion_model.normal.lora_up.weight"] = torch.ones((2, 2))
save_file(tensors, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "three_low_out", verbose=False
)
result = load_file(output)
for i in range(3):
assert result[f"diffusion_model.blocks.{i}.attn.comfy_attention.config"].dtype == torch.uint8
def test_recognized_quantization_creates_no_resize_output(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "recognized_quant.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
}, str(source))
output = tmp_path / "recognized_must_not_exist.safetensors"
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def _generic_params(output_filename):
return {
"mismatch_mode": "skip", "alignment_mode": "pad/crop", "alpha": 0.5,
"beta": 0.0, "gamma": 0.5, "delta": 2.0, "epsilon": 0.01, "zeta": 0.0,
"seed": 0, "output_filename": output_filename, "save_dtype": "bf16",
"override_dtype": True, "device": "cpu", "dtype": torch.float32,
"exclude_patterns": "", "discard_patterns": "", "glob_patterns": False,
"lazy_load": True, "force_clear_cache": False, "include_1d_diffs": True,
}
def test_generic_merge_preserves_model_a_when_either_source_key_is_low_bit(
monkeypatch, tmp_path, modules
):
generic = modules["generic"]
operations = modules["operations"]
model_a = tmp_path / "generic_guard_a.safetensors"
model_b = tmp_path / "generic_guard_b.safetensors"
a_tensors = {
"a_low": torch.tensor([[1, 2]], dtype=torch.uint8),
"b_low": torch.tensor([[0.125, 0.25]], dtype=torch.float32),
"normal": torch.ones((1, 2)),
}
save_file(a_tensors, str(model_a))
save_file({
"a_low": torch.tensor([[9.0, 9.0]]),
"b_low": torch.tensor([[3, 4]], dtype=torch.int8),
"normal": torch.full((1, 2), 3.0),
}, str(model_b))
paths = {"a": str(model_a), "b": str(model_b)}
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
params = _generic_params("generic_guarded")
params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "loras",
)
result = load_file(
str(tmp_path / "loras" / "generic_guarded.safetensors")
)
torch.testing.assert_close(result["a_low"], a_tensors["a_low"])
torch.testing.assert_close(result["b_low"], a_tensors["b_low"])
assert result["a_low"].dtype == torch.uint8
assert result["b_low"].dtype == torch.float32
torch.testing.assert_close(result["normal"].float(), torch.full((1, 2), 2.0))
def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "merge_guard_base.safetensors"
adapter = tmp_path / "merge_guard_adapter.safetensors"
base_quant = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"model.diffusion_model.base_quant.weight": base_quant,
"model.diffusion_model.guarded.weight": torch.zeros((2, 2)),
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 1]], dtype=torch.int8),
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 1)),
"diffusion_model.guarded.diff": torch.full((2, 2), 5.0),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(adapter))
output, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16, "merge_guarded",
verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(output)
torch.testing.assert_close(result["model.diffusion_model.base_quant.weight"], base_quant)
assert result["model.diffusion_model.base_quant.weight"].dtype == torch.uint8
torch.testing.assert_close(
result["model.diffusion_model.guarded.weight"].float(), torch.zeros((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.normal.weight"].float(), torch.ones((2, 2))
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 1/2")
assert "[GUARDED] (1)" in report
assert "diffusion_model.guarded.lora_down.weight" in report
assert "torch.int8" in report
assert "model.diffusion_model.base_quant.weight" in report
def test_merge_to_model_dequantizes_base_and_omits_sidecars(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "quantized_base.safetensors"
adapter = tmp_path / "quantized_adapter.safetensors"
quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
prefix = "model.diffusion_model.layer"
save_file({
f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8),
f"{prefix}.weight_scale": torch.tensor(2.0),
f"{prefix}.comfy_quant": quant,
}, str(base))
save_file({"diffusion_model.layer.diff": torch.ones((2, 2))}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
"dequantized_merge", verbose=False,
)
result = load_file(output)
assert list(result) == [f"{prefix}.weight"]
torch.testing.assert_close(result[f"{prefix}.weight"].float(), torch.full((2, 2), 3.0))
def test_diffusion_decoder_matches_core_quantized_tensor(tmp_path, modules):
from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
path = tmp_path / "core_int8.safetensors"
key = "layer.weight"
qdata = torch.tensor([[1, -2], [3, 4]], dtype=torch.int8)
scale = torch.tensor(0.5)
save_file({
key: qdata,
"layer.weight_scale": scale,
"layer.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
}, str(path))
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
decoder = modules["guard"].DiffusionQuantization(handler, str(path), "Core comparison")
keys = decoder.required_keys(key)
assert set(keys) == {key, "layer.weight_scale", "layer.comfy_quant"}
assert decoder.data_keys == (key,)
stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
tensors = {}
try:
for expected in keys:
batch = next(stream)
assert len(batch) == 1 and batch[0][0] == expected
tensors[expected] = batch[0][1]
actual = decoder.decode(key, tensors, torch.float32)
finally:
for source_key in tensors:
handler.mark_processed(source_key)
stream.close()
layout_name = QUANT_ALGOS["int8_tensorwise"]["comfy_tensor_layout"]
layout = get_layout_class(layout_name)
expected = QuantizedTensor(
qdata, layout_name,
layout.Params(scale=scale, orig_dtype=torch.float32, orig_shape=(2, 2)),
).dequantize()
torch.testing.assert_close(actual, expected)
torch.testing.assert_close(actual, qdata.float() * scale)
finally:
handler.__exit__(None, None, None)
@pytest.mark.parametrize("prefix", ["model.", "model.diffusion_model."])
def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules, prefix):
import comfy.model_detection
stem = f"{prefix}layer"
class HeaderOnly:
def keys(self):
return [f"{stem}.weight", f"{stem}.weight_scale", f"{stem}.weight_scale_2"]
def metadata(self):
return {"_quantization_metadata": json.dumps({"layers": {stem: {"format": "nvfp4"}}})}
def get_shape(self, key):
return {f"{stem}.weight": (16, 8), f"{stem}.weight_scale": (16, 1), f"{stem}.weight_scale_2": ()}[key]
def get_dtype(self, key):
return torch.uint8 if key == f"{stem}.weight" else torch.float32
class Config:
def get_model(self, state, prefix, device):
assert device == torch.device("meta")
assert all(tensor.device.type == "meta" for tensor in state.values())
model = torch.nn.Module()
model.diffusion_model = torch.nn.Module()
model.diffusion_model.layer = torch.nn.Module()
model.diffusion_model.layer._orig_shape = (3, 5)
return model
monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
decoder = modules["guard"].DiffusionQuantization(HeaderOnly(), "header-only", "Shape test")
assert decoder.logical_shape(f"{stem}.weight") == (3, 5)
def test_nvfp4_stream_decode_matches_core_full_dequantization(monkeypatch, tmp_path, modules):
import comfy.model_detection
from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
layout_name = QUANT_ALGOS["nvfp4"]["comfy_tensor_layout"]
layout = get_layout_class(layout_name)
original = torch.arange(15, dtype=torch.float16).reshape(3, 5) / 10
qdata, params = layout.quantize(original)
path = tmp_path / "nvfp4.safetensors"
metadata = {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"format": "nvfp4"}}})}
save_file({
"model.layer.weight": qdata,
"model.layer.weight_scale": params.block_scale,
"model.layer.weight_scale_2": params.scale,
}, str(path), metadata=metadata)
class Config:
def get_model(self, state, prefix, device):
model = torch.nn.Module()
model.diffusion_model = torch.nn.Module()
model.diffusion_model.layer = torch.nn.Module()
model.diffusion_model.layer._orig_shape = tuple(original.shape)
return model
monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
decoder = modules["guard"].DiffusionQuantization(handler, str(path), "NVFP4 test")
key = "model.layer.weight"
assert decoder.logical_shape(key) == tuple(original.shape)
keys = decoder.required_keys(key)
stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
tensors = {}
try:
for source_key in keys:
batch = next(stream)
assert len(batch) == 1 and batch[0][0] == source_key
tensors[source_key] = batch[0][1]
actual = decoder.decode(key, tensors, torch.float16)
finally:
for source_key in tensors:
handler.mark_processed(source_key)
stream.close()
expected = QuantizedTensor(qdata, layout_name, params).dequantize()
torch.testing.assert_close(actual, expected)
assert actual.shape == original.shape
finally:
handler.__exit__(None, None, None)
def test_w6a8_int8_stream_decode_matches_core_dequantization(tmp_path, modules):
from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
layout_name = QUANT_ALGOS["w6a8_int8"]["comfy_tensor_layout"]
layout = get_layout_class(layout_name)
original = torch.randn(4, 256, dtype=torch.float32)
qdata, params = layout.quantize(original, bits=6, codebook=False)
path = tmp_path / "w6a8.safetensors"
key = "model.layer.weight"
metadata = {
"_quantization_metadata": json.dumps({
"layers": {
"model.layer": {
"format": "w6a8_int8",
"group_size": params.group_size,
"convrot_groupsize": params.convrot_groupsize,
}
}
})
}
save_file({
key: qdata,
"model.layer.weight_s_rel": params.scale,
"model.layer.weight_s_channel": params.s_channel,
}, str(path), metadata=metadata)
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
decoder = modules["guard"].DiffusionQuantization(handler, str(path), "W6A8 test")
assert decoder.logical_shape(key) == tuple(original.shape)
keys = decoder.required_keys(key)
stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
tensors = {}
try:
for source_key in keys:
batch = next(stream)
assert len(batch) == 1 and batch[0][0] == source_key
tensors[source_key] = batch[0][1]
actual = decoder.decode(key, tensors, torch.float32)
finally:
for source_key in tensors:
handler.mark_processed(source_key)
stream.close()
expected = QuantizedTensor(qdata, layout_name, params).dequantize()
torch.testing.assert_close(actual, expected)
assert actual.shape == original.shape
finally:
handler.__exit__(None, None, None)
+650
View File
@@ -0,0 +1,650 @@
import importlib
import sys
import types
from collections import Counter
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import (
IncrementalSafetensorsWriter,
MemoryEfficientSafeOpen,
)
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_resize_factor_space_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def resize():
return _load_module("nodes.lora_resize")
def _write_uel(path, tensors, metadata=None):
with IncrementalSafetensorsWriter(str(path), metadata=metadata or {}) as writer:
writer.write_batch(
[(key, tensor.cpu().contiguous()) for key, tensor in tensors.items()]
)
def _read_uel(path, async_method=None, mark_method=None):
result = {}
with MemoryEfficientSafeOpen(str(path), low_memory=True) as loader:
keys = loader.keys()
stream = (async_method or MemoryEfficientSafeOpen.async_stream)(
loader, keys, batch_size=1, prefetch_batches=1, pin_memory=False
)
for batch in stream:
for key, tensor in batch:
result[key] = tensor.clone()
(mark_method or MemoryEfficientSafeOpen.mark_processed)(loader, key)
return result
@pytest.mark.parametrize("glob", [False, True])
def test_resize_filters_preserve_and_discard_whole_layers(monkeypatch, tmp_path, resize, glob):
source = tmp_path / "filters.safetensors"
tensors = {}
for name in ("keep", "drop", "resize"):
tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3, dtype=torch.float16)
tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3, dtype=torch.float16)
tensors[f"diffusion_model.{name}.alpha"] = torch.tensor(1.5)
tensors["diffusion_model.drop.extra"] = torch.ones(1)
tensors["standalone"] = torch.ones(1)
_write_uel(source, tensors)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
original_async = MemoryEfficientSafeOpen.async_stream
streamed = []
def track(self, keys, **kwargs):
streamed.extend(keys)
yield from original_async(self, keys, **kwargs)
monkeypatch.setattr(MemoryEfficientSafeOpen, "async_stream", track)
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.bfloat16, "filtered", verbose=False,
exclude_patterns="*.keep.lora_A.weight\n*.drop.*" if glob else r"\.keep\.lora_A\.weight$" + "\n" + r"\.drop\.",
discard_patterns="*.drop.alpha\nstandalone" if glob else r"\.drop\.alpha$" + "\n^standalone$",
glob_patterns=glob,
)
assert not any(".drop." in key or key == "standalone" for key in streamed)
actual = _read_uel(output, original_async)
assert not any(".drop." in key or key == "standalone" for key in actual)
for key in tensors:
if ".keep." in key:
assert actual[key].dtype == tensors[key].dtype
assert torch.equal(actual[key], tensors[key])
assert actual["diffusion_model.resize.lora_A.weight"].shape[0] == 1
@pytest.mark.parametrize("glob", [False, True])
def test_resize_include_mode_only_resizes_matching_layers(monkeypatch, tmp_path, resize, glob):
source = tmp_path / "include_filters.safetensors"
tensors = {}
for name in ("include", "discard", "skip"):
tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3, dtype=torch.float16)
tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3, dtype=torch.float16)
tensors[f"diffusion_model.{name}.alpha"] = torch.tensor(1.5)
tensors["standalone.include"] = torch.ones(1)
tensors["standalone.skip"] = torch.ones(1)
_write_uel(source, tensors)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.bfloat16, "included", verbose=False,
exclude_patterns="*include*\n*discard*\nstandalone.include" if glob else r"\.(include|discard)\." + "\n^standalone\\.include$",
discard_patterns="*discard*" if glob else r"\.discard\.",
glob_patterns=glob,
include_mode=True,
)
actual = _read_uel(output)
assert actual["diffusion_model.include.lora_A.weight"].shape[0] == 1
assert torch.equal(actual["standalone.include"], tensors["standalone.include"])
assert not any(".discard." in key for key in actual)
for key in tensors:
if ".skip." in key or key == "standalone.skip":
assert torch.equal(actual[key], tensors[key])
def test_resize_include_mode_empty_patterns_preserves_all_tensors(monkeypatch, tmp_path, resize):
source = tmp_path / "empty_include.safetensors"
tensors = {"standalone": torch.ones(2)}
_write_uel(source, tensors)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float32, "empty", verbose=False,
include_mode=True,
)
actual = _read_uel(output)
assert torch.equal(actual["standalone"], tensors["standalone"])
@pytest.mark.parametrize("node_name,extra", [
("LoRAResizeFixed", {"new_rank": 1}),
("LoRAResizeRatio", {"max_rank": 1, "ratio": 2.0}),
("LoRAResizeFrobenius", {"max_rank": 1, "min_rank": 1, "target": 0.9}),
("LoRAResizeCumulative", {"max_rank": 1, "target": 0.9}),
])
def test_all_resize_nodes_forward_filters(monkeypatch, tmp_path, resize, node_name, extra):
node = getattr(resize, node_name)
controls = {item.id: item for item in node.define_schema().inputs}
assert controls["exclude_patterns"].default == ""
assert controls["discard_patterns"].default == ""
assert controls["glob_patterns"].default is False
assert controls["include_mode"].default is False
assert [item.id for item in node.define_schema().inputs if item.id != "layer_parameters"][-1] == "include_mode"
calls = []
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize.folder_paths, "get_full_path_or_raise", lambda *args: "input")
monkeypatch.setattr(resize, "resize_lora_file", lambda *args, **kwargs: calls.append(kwargs))
node.execute(lora_name="input", output_filename="out", save_dtype="fp16", device="cpu",
force_clear_cache=False, exclude_patterns="keep*", discard_patterns="drop*",
glob_patterns=True, include_mode=True, **extra)
assert calls[0]["exclude_patterns"] == "keep*"
assert calls[0]["discard_patterns"] == "drop*"
assert calls[0]["glob_patterns"] is True
assert calls[0]["include_mode"] is True
def test_resize_invalid_filter_fails_before_writer(monkeypatch, tmp_path, resize):
import re
source = tmp_path / "invalid_filter.safetensors"
_write_uel(source, {"layer.diff": torch.ones(2)})
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
def unexpected_writer(*args, **kwargs):
pytest.fail("Writer must not open for invalid patterns")
monkeypatch.setattr(resize, "atomic_uel_writer", unexpected_writer)
with pytest.raises(re.error):
resize.resize_lora_file(str(source), 1, None, None, "cpu", torch.float32,
"unused", exclude_patterns="[")
def _dense_truncated(up, down, rank):
dense = up.float() @ down.float()
u, s, vh = torch.linalg.svd(dense, full_matrices=False)
return (u[:, :rank] * s[:rank]) @ vh[:rank, :]
def test_linear_factor_space_matches_dense_reference_and_bounds_svd(
monkeypatch, resize
):
generator = torch.Generator().manual_seed(123)
down = torch.randn((4, 17), generator=generator)
up = torch.randn((13, 4), generator=generator)
expected = _dense_truncated(up, down, 3)
observed_shapes = []
original_svd = torch.linalg.svd
def track_svd(tensor, *args, **kwargs):
observed_shapes.append(tuple(tensor.shape))
return original_svd(tensor, *args, **kwargs)
monkeypatch.setattr(torch.linalg, "svd", track_svd)
result = resize._resize_lora_factors(down, up, 3, None, None, 1.0)
actual = result["lora_up"] @ result["lora_down"]
assert torch.allclose(actual, expected, atol=2e-5, rtol=2e-5)
assert observed_shapes == [(4, 4)]
def test_convolutional_factor_space_matches_dense_reference(resize):
generator = torch.Generator().manual_seed(456)
down = torch.randn((3, 2, 3, 2), generator=generator)
up = torch.randn((5, 3, 1, 1), generator=generator)
result = resize._resize_lora_factors(down, up, 2, None, None, 1.0)
actual = result["lora_up"].reshape(5, 2) @ result["lora_down"].reshape(2, -1)
expected = _dense_truncated(up.reshape(5, 3), down.reshape(3, -1), 2)
assert result["lora_up"].shape == (5, 2, 1, 1)
assert result["lora_down"].shape == (2, 2, 3, 2)
assert torch.allclose(actual, expected, atol=2e-5, rtol=2e-5)
@pytest.mark.parametrize(
("method", "parameter", "maximum", "expected_rank"),
[
(None, None, 3, 3),
("sv_ratio", 4.0, 4, 2),
("sv_fro", 0.95, 4, 1),
("sv_cumulative", 0.8, 4, 2),
],
)
def test_all_resize_modes_use_the_shared_spectrum(
resize, method, parameter, maximum, expected_rank
):
singular_values = torch.tensor([10.0, 3.0, 1.0, 0.1])
result = resize._resize_lora_factors(
torch.eye(4),
torch.diag(singular_values),
maximum,
method,
parameter,
1.0,
)
assert result["new_rank"] == expected_rank
def test_mixed_source_ranks_are_capped_per_layer(resize):
rank_128 = resize._resize_lora_factors(
torch.eye(128), torch.eye(128), 384, None, None, 1.0
)
rank_384 = resize._resize_lora_factors(
torch.eye(384), torch.eye(384), 384, None, None, 1.0
)
assert rank_128["new_rank"] == 128
assert rank_384["new_rank"] == 384
def test_resize_file_uses_only_async_uel_and_preserves_per_layer_scale(
monkeypatch, tmp_path, resize
):
source = tmp_path / "mixed_rank.safetensors"
_write_uel(
source,
{
"diffusion_model.r2.lora_A.weight": torch.eye(2),
"diffusion_model.r2.lora_B.weight": torch.eye(2),
"diffusion_model.r2.alpha": torch.tensor(1.0),
"diffusion_model.r3.lora_A.weight": torch.eye(3),
"diffusion_model.r3.lora_B.weight": torch.eye(3),
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
original_async = resize.MemoryEfficientSafeOpen.async_stream
original_mark = resize.MemoryEfficientSafeOpen.mark_processed
async_calls = []
marked = []
def track_async(self, keys, **kwargs):
async_calls.append((tuple(keys), kwargs, self.low_memory))
yield from original_async(self, keys, **kwargs)
def track_mark(self, key):
marked.append(key)
return original_mark(self, key)
def reject_sync_get(*args, **kwargs):
raise AssertionError("resize used synchronous get_tensor")
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async)
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "mark_processed", track_mark)
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get)
output = resize.resize_lora_file(
str(source), 3, None, None, "cpu", torch.float32, "mixed_rank_output",
verbose=False,
)
assert len(async_calls) == 1
streamed_keys, options, low_memory = async_calls[0]
assert low_memory is True
assert options == {
"batch_size": 1,
"prefetch_batches": 1,
"pin_memory": False,
}
assert Counter(marked) == Counter(streamed_keys)
assert all(count == 1 for count in Counter(marked).values())
output_tensors = _read_uel(output, original_async, original_mark)
assert output_tensors["diffusion_model.r2.lora_A.weight"].shape[0] == 2
assert output_tensors["diffusion_model.r3.lora_A.weight"].shape[0] == 3
assert output_tensors["diffusion_model.r2.alpha"].item() == pytest.approx(1.0)
assert "diffusion_model.r3.alpha" not in output_tensors
def test_normalize_alpha_uses_bounded_uel_and_removes_alpha(
monkeypatch, tmp_path, resize
):
source = tmp_path / "alpha_input.safetensors"
down = torch.eye(2, dtype=torch.float16)
up = torch.tensor([[2.0, 4.0], [6.0, 8.0]], dtype=torch.float16)
_write_uel(
source,
{
"diffusion_model.block.lora_A.weight": down,
"diffusion_model.block.lora_B.weight": up,
"diffusion_model.block.alpha": torch.tensor(1.0),
"other.weight": torch.tensor([3.0]),
},
metadata={"source": "test"},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
original_async = resize.MemoryEfficientSafeOpen.async_stream
original_mark = resize.MemoryEfficientSafeOpen.mark_processed
async_calls = []
marked = []
def track_async(self, keys, **kwargs):
async_calls.append((tuple(keys), kwargs, self.low_memory))
yield from original_async(self, keys, **kwargs)
def track_mark(self, key):
marked.append(key)
return original_mark(self, key)
def reject_sync_get(*args, **kwargs):
raise AssertionError("alpha normalization used synchronous get_tensor")
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async)
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "mark_processed", track_mark)
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get)
output = resize.normalize_lora_alpha_file(
str(source), "alpha_normalized", verbose=False
)
assert len(async_calls) == 1
streamed_keys, options, low_memory = async_calls[0]
assert streamed_keys == (
"diffusion_model.block.lora_B.weight",
"diffusion_model.block.alpha",
)
assert options == {
"batch_size": 1,
"prefetch_batches": 1,
"pin_memory": False,
}
assert low_memory is True
assert Counter(marked) == Counter(streamed_keys)
output_tensors = _read_uel(output, original_async, original_mark)
assert "diffusion_model.block.alpha" not in output_tensors
assert torch.equal(output_tensors["diffusion_model.block.lora_A.weight"], down)
assert torch.equal(
output_tensors["diffusion_model.block.lora_B.weight"], up * 0.5
)
assert torch.equal(output_tensors["other.weight"], torch.tensor([3.0]))
with MemoryEfficientSafeOpen(output, low_memory=True) as loader:
assert loader.metadata()["alpha_normalized"] == "true"
def test_normalize_alpha_rejects_alpha_free_input(monkeypatch, tmp_path, resize):
source = tmp_path / "no_alpha.safetensors"
_write_uel(
source,
{
"diffusion_model.block.lora_A.weight": torch.eye(2),
"diffusion_model.block.lora_B.weight": torch.eye(2),
},
)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
with pytest.raises(ValueError, match="contains no alpha tensors"):
resize.normalize_lora_alpha_file(
str(source), "unused", verbose=False
)
def test_normalize_alpha_maps_flattened_names_with_reference(
monkeypatch, tmp_path, resize
):
source = tmp_path / "flattened_alpha.safetensors"
reference = tmp_path / "reference.safetensors"
down = torch.eye(2, dtype=torch.float16)
up = torch.tensor([[2.0, 4.0], [6.0, 8.0]], dtype=torch.float16)
_write_uel(
source,
{
"lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight": down,
"lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight": up,
"lora_unet_double_blocks_0_img_attn_qkv.alpha": torch.tensor(1.0),
},
)
_write_uel(
reference,
{"diffusion_model.double_blocks.0.img_attn.qkv.weight": torch.eye(2)},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.normalize_lora_alpha_file(
str(source),
"reference_mapped_alpha",
verbose=False,
reference_model_path=str(reference),
)
output_tensors = _read_uel(output)
assert set(output_tensors) == {
"diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight",
"diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight",
}
assert torch.equal(
output_tensors["diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight"],
down,
)
assert torch.equal(
output_tensors["diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight"],
up * 0.5,
)
def test_resize_does_not_invent_alpha_for_peft_pair(monkeypatch, tmp_path, resize):
source = tmp_path / "peft_no_alpha.safetensors"
_write_uel(
source,
{
"base_model.model.transformer.block.to_q.lora_A.default.weight": (
torch.eye(2)
),
"base_model.model.transformer.block.to_q.lora_B.default.weight": (
torch.eye(2)
),
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
output = resize.resize_lora_file(
str(source), 2, None, None, "cpu", torch.float32, "peft_output",
verbose=False,
)
output_tensors = _read_uel(output)
assert len(output_tensors) == 2
assert not any(key.endswith(".alpha") for key in output_tensors)
def test_cuda_oom_retries_only_current_pair_on_cpu(monkeypatch, resize):
calls = []
original = resize._resize_factors_on_device
def fail_cuda_then_run_cpu(*args):
process_device = args[-1]
calls.append(process_device)
if process_device == "cuda":
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
return original(*args)
monkeypatch.setattr(resize, "_resize_factors_on_device", fail_cuda_then_run_cpu)
monkeypatch.setattr(resize, "_release_cuda_after_oom", lambda: None)
result = resize._resize_lora_factors(
torch.eye(2), torch.eye(2), 1, None, None, 1.0,
process_device="cuda",
)
assert calls == ["cuda", "cpu"]
assert result["cpu_fallback"] is True
def test_low_bit_pair_uses_raw_copy_without_tensor_loading(
monkeypatch, tmp_path, resize
):
source = tmp_path / "low_bit_pair.safetensors"
down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
up = torch.eye(2)
_write_uel(
source,
{
"diffusion_model.low.lora_A.weight": down,
"diffusion_model.low.lora_B.weight": up,
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
async_calls = []
original_async = resize.MemoryEfficientSafeOpen.async_stream
original_mark = resize.MemoryEfficientSafeOpen.mark_processed
def track_async(self, keys, **kwargs):
async_calls.append(tuple(keys))
yield from original_async(self, keys, **kwargs)
def reject_sync_get(*args, **kwargs):
raise AssertionError("low-bit raw copy decoded a tensor")
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "async_stream", track_async)
monkeypatch.setattr(resize.MemoryEfficientSafeOpen, "get_tensor", reject_sync_get)
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "low_bit_output",
verbose=False,
)
assert async_calls == []
output_tensors = _read_uel(output, original_async, original_mark)
assert torch.equal(output_tensors["diffusion_model.low.lora_A.weight"], down)
assert torch.equal(output_tensors["diffusion_model.low.lora_B.weight"], up)
def test_resize_node_interfaces_remove_obsolete_loader_and_iteration_controls(resize):
schemas = [
resize.LoRAResizeFixed.define_schema(),
resize.LoRAResizeRatio.define_schema(),
resize.LoRAResizeFrobenius.define_schema(),
resize.LoRAResizeCumulative.define_schema(),
]
input_ids = [[item.id for item in schema.inputs] for schema in schemas]
assert "svd_niter" not in input_ids[0]
assert all("lazy_load" not in ids for ids in input_ids)
assert input_ids[2][:4] == ["lora_name", "max_rank", "min_rank", "target"]
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
def test_cuda_and_cpu_factor_space_results_match(resize):
generator = torch.Generator().manual_seed(789)
down = torch.randn((4, 11), generator=generator)
up = torch.randn((9, 4), generator=generator)
cpu = resize._resize_lora_factors(down, up, 3, None, None, 1.0, process_device="cpu")
cuda = resize._resize_lora_factors(down, up, 3, None, None, 1.0, process_device="cuda")
cpu_update = cpu["lora_up"] @ cpu["lora_down"]
cuda_update = cuda["lora_up"] @ cuda["lora_down"]
assert torch.allclose(cpu_update, cuda_update, atol=2e-4, rtol=2e-4)
def test_modify_lora_alpha_handles_present_and_absent_alpha(monkeypatch, tmp_path, resize):
source = tmp_path / "mixed_alpha.safetensors"
down_a = torch.eye(4, dtype=torch.float16)
up_a = torch.ones((4, 4), dtype=torch.float16)
down_b = torch.eye(2, dtype=torch.float16)
up_b = torch.ones((2, 2), dtype=torch.float16) * 3.0
_write_uel(
source,
{
"diffusion_model.block_a.lora_A.weight": down_a,
"diffusion_model.block_a.lora_B.weight": up_a,
"diffusion_model.block_a.alpha": torch.tensor(2.0, dtype=torch.float16),
"diffusion_model.block_b.lora_A.weight": down_b,
"diffusion_model.block_b.lora_B.weight": up_b,
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
scaled = resize.modify_lora_alpha_file(
str(source), "scaled_alpha", alpha_mode="scale", alpha_value=1.5, verbose=False
)
scaled_tensors = _read_uel(scaled)
# Factor weights are byte-exact unchanged
assert torch.equal(scaled_tensors["diffusion_model.block_a.lora_A.weight"], down_a)
assert torch.equal(scaled_tensors["diffusion_model.block_a.lora_B.weight"], up_a)
assert torch.equal(scaled_tensors["diffusion_model.block_b.lora_A.weight"], down_b)
assert torch.equal(scaled_tensors["diffusion_model.block_b.lora_B.weight"], up_b)
# block_a had alpha=2.0 -> 2.0 * 1.5 = 3.0; block_b had no alpha (rank=2) -> 2.0 * 1.5 = 3.0
assert scaled_tensors["diffusion_model.block_a.alpha"].item() == pytest.approx(3.0)
assert scaled_tensors["diffusion_model.block_b.alpha"].item() == pytest.approx(3.0)
# Shifted with filter excluding block_a
shifted = resize.modify_lora_alpha_file(
str(source),
"shifted_alpha",
alpha_mode="shift",
alpha_value=2.0,
exclude_patterns="block_a",
verbose=False,
)
shifted_tensors = _read_uel(shifted)
assert shifted_tensors["diffusion_model.block_a.alpha"].item() == pytest.approx(2.0)
assert shifted_tensors["diffusion_model.block_b.alpha"].item() == pytest.approx(4.0)
def test_denormalize_lora_alpha_is_inverse_of_normalize(monkeypatch, tmp_path, resize):
source = tmp_path / "alpha_lora.safetensors"
down = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
up = torch.tensor([[2.0, 0.5], [1.0, 3.0]], dtype=torch.float32)
_write_uel(
source,
{
"diffusion_model.block.lora_A.weight": down,
"diffusion_model.block.lora_B.weight": up,
"diffusion_model.block.alpha": torch.tensor(1.0, dtype=torch.float32),
},
)
monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None)
normalized_path = resize.normalize_lora_alpha_file(
str(source), "norm_step", verbose=False
)
denormalized_path = resize.denormalize_lora_alpha_file(
str(normalized_path),
"denorm_step",
alpha_mode="set_alpha",
alpha_value=1.0,
verbose=False,
)
denorm_tensors = _read_uel(denormalized_path)
assert denorm_tensors["diffusion_model.block.alpha"].item() == pytest.approx(1.0)
assert torch.allclose(denorm_tensors["diffusion_model.block.lora_A.weight"], down)
assert torch.allclose(denorm_tensors["diffusion_model.block.lora_B.weight"], up)
with pytest.raises(ValueError, match="non-zero"):
resize.denormalize_lora_alpha_file(
str(normalized_path), "bad_zero", alpha_mode="set_alpha", alpha_value=0.0, verbose=False
)
+89
View File
@@ -0,0 +1,89 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_resize_min_rank_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def resize():
return _load_module("nodes.lora_resize")
def test_frobenius_minimum_rank_is_enforced(resize):
singular_values = torch.tensor([10.0, 1.0, 0.1, 0.01, 0.001])
rank, _, _ = resize._compute_resize(
singular_values, 4, "sv_fro", 0.5, 1.0, min_rank=3
)
assert rank == 3
def test_default_minimum_rank_remains_one(resize):
singular_values = torch.tensor([10.0, 1.0, 0.1, 0.01, 0.001])
rank, _, _ = resize._compute_resize(
singular_values, 4, "sv_fro", 0.5, 1.0
)
assert rank == 1
def test_factor_space_resize_passes_minimum_rank_to_dynamic_selection(resize):
down = torch.eye(5)
up = torch.diag(torch.tensor([10.0, 1.0, 0.1, 0.01, 0.001]))
result = resize._resize_lora_factors(
down, up, 4, "sv_fro", 0.5, 1.0, min_rank=3
)
assert result["new_rank"] == 3
assert result["lora_down"].shape[0] == 3
def test_minimum_rank_is_capped_by_layer_dimensions(resize):
singular_values = torch.tensor([10.0, 1.0, 0.1])
rank, _, _ = resize._compute_resize(
singular_values, 4, "sv_fro", 0.5, 1.0, min_rank=4
)
assert rank == 3
def test_frobenius_schema_places_minimum_beside_maximum(resize):
schema = resize.LoRAResizeFrobenius.define_schema()
assert [item.id for item in schema.inputs[:4]] == [
"lora_name", "max_rank", "min_rank", "target"
]
assert schema.inputs[2].default == 1
def test_resize_rejects_minimum_above_maximum(tmp_path, resize):
with pytest.raises(ValueError, match="min_rank cannot exceed max rank"):
resize.resize_lora_file(
str(tmp_path / "unused.safetensors"),
4,
"sv_fro",
0.999,
"cpu",
torch.float32,
"unused",
min_rank=5,
)
+284
View File
@@ -0,0 +1,284 @@
import importlib
import json
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_merger_filter_tests"
def _module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def modules():
return _module("nodes.merger"), _module("nodes.merger_ops")
def _params(name, *, patterns="", include=False, discard="", glob=False):
return {
"mismatch_mode": "skip", "alignment_mode": "pad/crop", "alpha": 0.5,
"beta": 0.0, "gamma": 0.5, "delta": 2.0, "epsilon": 0.01,
"zeta": 0.0, "seed": 0, "output_filename": name, "save_dtype": "fp32",
"override_dtype": True, "device": "cpu", "dtype": torch.float32,
"exclude_patterns": patterns, "include_mode": include,
"discard_patterns": discard, "glob_patterns": glob, "lazy_load": True,
"force_clear_cache": False,
}
def _run(monkeypatch, tmp_path, merger, operations, params):
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
save_file({"match.weight": torch.tensor([2.0]), "other.weight": torch.tensor([4.0])}, paths["a"])
save_file({"match.weight": torch.tensor([6.0]), "other.weight": torch.tensor([8.0])}, paths["b"])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params.update({"model_a": "a", "model_b": "b"})
merger.MergerLogic.execute_merge({"model_a": "a", "model_b": "b"}, "Weight-Sum", operations.TWO_MODEL_MODES, params, "checkpoints")
return load_file(str(tmp_path / "checkpoints" / f"{params['output_filename']}.safetensors"))
def test_include_filter_controls_preflight_and_streaming(monkeypatch, tmp_path, modules):
merger, operations = modules
selected = _run(monkeypatch, tmp_path, merger, operations, _params("selected", patterns="match", include=True))
assert selected["match.weight"].item() == 4.0
assert selected["other.weight"].item() == 4.0
empty = _run(monkeypatch, tmp_path, merger, operations, _params("empty", include=True))
assert empty["match.weight"].item() == 2.0
assert empty["other.weight"].item() == 4.0
discarded = _run(monkeypatch, tmp_path, merger, operations, _params("discarded", patterns="match", include=True, discard="match"))
assert "match.weight" not in discarded
assert discarded["other.weight"].item() == 4.0
globbed = _run(monkeypatch, tmp_path, merger, operations, _params("globbed", patterns="*.weight", include=True, glob=True))
assert globbed["match.weight"].item() == 4.0
assert globbed["other.weight"].item() == 6.0
def test_layer_parameters_apply_per_layer_without_leakage(monkeypatch, tmp_path, modules):
merger, operations = modules
def resolve(_payload, operation, layers, defaults, *, node_name):
assert operation == "merge:Weight-Sum"
assert node_name == "checkpoints Weight-Sum merger"
assert set(layers) == {"match.weight", "other.weight"}
assert defaults["alpha"] == 0.5
return {"match.weight": {"alpha": 0.0}, "other.weight": {"alpha": 1.0}}
monkeypatch.setattr(merger, "resolve_layer_parameters", resolve)
result = _run(monkeypatch, tmp_path, merger, operations, _params("per_layer"))
assert result["match.weight"].item() == 2.0
assert result["other.weight"].item() == 8.0
def test_parsed_layer_parameters_change_only_matching_output(monkeypatch, tmp_path, modules):
merger, operations = modules
layer_parameters = _module("nodes.layer_parameters")
params = _params("parsed")
params["layer_parameters"] = layer_parameters.parse_rules("(match\\.weight) alpha:0\n")
result = _run(monkeypatch, tmp_path, merger, operations, params)
assert result["match.weight"].item() == 2.0
assert result["other.weight"].item() == 6.0
def test_layer_parameter_failure_precedes_writer_and_streaming(monkeypatch, tmp_path, modules):
merger, operations = modules
monkeypatch.setattr(
merger,
"resolve_layer_parameters",
lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("unmatched rule line 1")),
)
params = _params("must_not_exist")
with pytest.raises(ValueError, match="unmatched rule"):
_run(monkeypatch, tmp_path, merger, operations, params)
assert not (tmp_path / "checkpoints" / "must_not_exist.safetensors").exists()
def test_lora_factor_group_uses_one_normalized_layer_name(monkeypatch, tmp_path, modules):
merger, operations = modules
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
tensors = {
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [2.0]]),
}
save_file({**tensors, "diffusion_model.layer.alpha": torch.tensor(1.0)}, paths["a"])
save_file(tensors, paths["b"])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params = _params("lora_group")
params.update({"model_a": "a", "model_b": "b", "layer_parameters": _module("nodes.layer_parameters").parse_rules("(diffusion_model\\.layer) alpha:0\n")})
merger.MergerLogic.execute_merge({"model_a": "a", "model_b": "b"}, "Weight-Sum", operations.TWO_MODEL_MODES, params, "loras")
result = load_file(str(tmp_path / "loras" / "lora_group.safetensors"))
assert "diffusion_model.layer.lora_A.weight" in result
assert "diffusion_model.layer.lora_B.weight" in result
assert not any(key.endswith(".alpha") for key in result)
def test_diffusion_merge_decodes_tensorwise_quantization_and_drops_sidecars(
monkeypatch, tmp_path, modules,
):
merger, operations = modules
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
save_file({
"model.diffusion_model.quant.weight": torch.tensor([1, 2], dtype=torch.int8),
"model.diffusion_model.quant.weight_scale": torch.tensor(2.0),
"model.diffusion_model.quant.comfy_quant": quant,
"dense.weight": torch.tensor([2.0, 4.0]),
}, paths["a"])
save_file({
"model.diffusion_model.quant.weight": torch.tensor([3, 6], dtype=torch.int8),
"model.diffusion_model.quant.weight_scale": torch.tensor(2.0),
"model.diffusion_model.quant.comfy_quant": quant,
"dense.weight": torch.tensor([6.0, 8.0]),
}, paths["b"])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params = _params("quantized_merge")
params.update({"model_a": "a", "model_b": "b"})
merger.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "diffusion_models",
)
output_path = tmp_path / "diffusion_models" / "quantized_merge.safetensors"
result = load_file(str(output_path))
key = "quant.weight"
assert set(result) == {key, "dense.weight"}
assert result[key].tolist() == pytest.approx([4.0, 8.0])
assert result[key].dtype == torch.float32
assert result["dense.weight"].tolist() == pytest.approx([4.0, 6.0])
with safe_open(str(output_path), framework="pt", device="cpu") as output:
metadata = output.metadata()
assert metadata in (None, {})
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_diffusion_merge_matches_prefixed_and_bare_models(monkeypatch, tmp_path, modules, prefixed_input):
merger, operations = modules
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
for name in ("a", "b"):
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
if name == prefixed_input:
stem = key[:-len(".weight")]
save_file({
key: torch.tensor([1, 2], dtype=torch.int8),
f"{stem}.weight_scale": torch.tensor(2.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
}, paths[name])
else:
save_file({key: torch.tensor([4.0, 8.0])}, paths[name])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params = _params("mixed_prefix_merge")
params.update({"model_a": "a", "model_b": "b"})
merger.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "diffusion_models",
)
output = load_file(str(tmp_path / "diffusion_models" / "mixed_prefix_merge.safetensors"))
assert list(output) == ["layer.weight"]
torch.testing.assert_close(output["layer.weight"], torch.tensor([3.0, 6.0]))
def test_three_diffusion_models_match_across_prefixes(monkeypatch, tmp_path, modules):
merger, operations = modules
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b", "c")}
save_file({"layer.weight": torch.tensor([1.0, 1.0])}, paths["a"])
save_file({"model.diffusion_model.layer.weight": torch.tensor([4.0, 4.0])}, paths["b"])
save_file({"layer.weight": torch.tensor([2.0, 2.0])}, paths["c"])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params = _params("three_prefix_merge")
params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
merger.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"}, "Add-Difference",
operations.THREE_MODEL_MODES, params, "diffusion_models",
)
output = load_file(str(tmp_path / "diffusion_models" / "three_prefix_merge.safetensors"))
assert list(output) == ["layer.weight"]
torch.testing.assert_close(output["layer.weight"], torch.tensor([2.0, 2.0]))
def test_prunekeys_include_mode_and_uel_streaming(monkeypatch, tmp_path):
prunekeys = _module("nodes.prunekeys")
for cls in (
prunekeys.ModelPruneKeys,
prunekeys.TextEncoderPruneKeys,
prunekeys.LoRAPruneKeys,
prunekeys.CheckpointPruneKeys,
prunekeys.EmbeddingPruneKeys,
):
schema = cls.define_schema()
assert schema.inputs[-1].id == "include_mode"
assert schema.inputs[-1].default is False
src = tmp_path / "source.safetensors"
save_file(
{
"layers.0.attn.weight": torch.tensor([1.0, 2.0]),
"layers.0.ffn.weight": torch.tensor([3.0, 4.0]),
"layers.1.attn.weight": torch.tensor([5.0, 6.0]),
},
str(src),
metadata={"author": "test"},
)
monkeypatch.setattr(prunekeys.folder_paths, "get_full_path_or_raise", lambda _kind, _name: str(src))
monkeypatch.setattr(prunekeys.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(prunekeys, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(prunekeys, "cleanup_after_operation", lambda: None)
# Default include_mode=False removes matching keys
rel_removed = prunekeys.ModelPruneKeys.execute(
diffusionmodel_name="source.safetensors",
keys_to_prune="attn",
use_regex=False,
output_filename="removed_attn",
include_mode=False,
).result[0]
removed_tensors = load_file(str(tmp_path / "diffusion_models" / rel_removed))
assert set(removed_tensors) == {"layers.0.ffn.weight"}
# include_mode=True keeps only matching keys
rel_kept = prunekeys.ModelPruneKeys.execute(
diffusionmodel_name="source.safetensors",
keys_to_prune=r"^layers\.\d+\.attn\.",
use_regex=True,
output_filename="kept_attn",
include_mode=True,
).result[0]
kept_tensors = load_file(str(tmp_path / "diffusion_models" / rel_kept))
assert set(kept_tensors) == {"layers.0.attn.weight", "layers.1.attn.weight"}
+246
View File
@@ -0,0 +1,246 @@
import importlib
from pathlib import Path
import sys
import types
import pytest
from safetensors.torch import load_file, save_file
import torch
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_minimax_fold_tests"
def _load_module(name: str):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def folder_module():
return _load_module("nodes.minimax_h3_fold")
def _make_h3_checkpoint(prefix: str = "") -> dict[str, torch.Tensor]:
return {
f"{prefix}time_embedder.proj_in.weight": torch.eye(16, dtype=torch.float32),
f"{prefix}time_embedder.proj_in.bias": torch.zeros(16, dtype=torch.float32),
f"{prefix}time_embedder.proj_out.weight": torch.ones((2688, 16), dtype=torch.float32) * 0.01,
f"{prefix}time_embedder.proj_out.bias": torch.zeros(2688, dtype=torch.float32),
f"{prefix}blocks.0.adaln_proj.linear.weight": torch.ones((128, 2688), dtype=torch.float32),
f"{prefix}blocks.0.adaln_proj.linear.bias": torch.zeros(128, dtype=torch.float32),
f"{prefix}blocks.1.adaln_proj.linear.weight": torch.ones((128, 2688), dtype=torch.float32) * 2.0,
f"{prefix}blocks.1.adaln_proj.linear.bias": torch.zeros(128, dtype=torch.float32),
f"{prefix}blocks.0.attn.qkv_proj.weight": torch.eye(64, dtype=torch.float32),
}
def test_fold_h3_diffusion_model_exact_basis(folder_module, tmp_path):
source = tmp_path / "model_full.safetensors"
output = tmp_path / "model_folded.safetensors"
original = _make_h3_checkpoint()
save_file(original, str(source))
report = folder_module.fold_minimax_h3_diffusion_model(
str(source), str(output), process_device="cpu"
)
assert "MiniMax H3 AdaLN Fold complete" in report
assert "AdaLN layers folded: 4" in report
converted = load_file(str(output))
assert "adaln_t_table" in converted
assert converted["adaln_t_table"].shape == (1025, 8)
assert converted["adaln_t_table"].dtype == torch.float32
# Folded AdaLN weight and bias
assert converted["blocks.0.adaln_proj.linear.weight"].shape == (128, 8)
assert converted["blocks.0.adaln_proj.linear.weight"].dtype == torch.float16
assert converted["blocks.0.adaln_proj.linear.bias"].dtype == torch.float16
# time_embedder should be dropped
assert "time_embedder.proj_in.weight" not in converted
assert "time_embedder.proj_out.weight" not in converted
# Unrelated layers preserved
assert "blocks.0.attn.qkv_proj.weight" in converted
assert torch.equal(converted["blocks.0.attn.qkv_proj.weight"], original["blocks.0.attn.qkv_proj.weight"])
def test_fold_h3_discard_patterns(folder_module, tmp_path):
source = tmp_path / "model_discard.safetensors"
output = tmp_path / "model_discarded.safetensors"
original = _make_h3_checkpoint()
save_file(original, str(source))
folder_module.fold_minimax_h3_diffusion_model(
str(source),
str(output),
discard_patterns="blocks.0.attn",
process_device="cpu",
)
converted = load_file(str(output))
assert "blocks.0.attn.qkv_proj.weight" not in converted
assert "blocks.0.adaln_proj.linear.weight" in converted
def test_fold_h3_exclude_patterns_preserves_unfolded_and_retains_time_embedder(folder_module, tmp_path):
source = tmp_path / "model_partial.safetensors"
output = tmp_path / "model_partial_folded.safetensors"
original = _make_h3_checkpoint()
save_file(original, str(source))
folder_module.fold_minimax_h3_diffusion_model(
str(source),
str(output),
exclude_patterns="blocks.1",
include_mode=False,
process_device="cpu",
)
converted = load_file(str(output))
# Block 0 folded
assert converted["blocks.0.adaln_proj.linear.weight"].shape == (128, 8)
# Block 1 kept at full width
assert converted["blocks.1.adaln_proj.linear.weight"].shape == (128, 2688)
# time_embedder MUST be retained because block 1 needs it
assert "time_embedder.proj_in.weight" in converted
assert "time_embedder.proj_out.weight" in converted
def test_node_schema_and_execution_contract(folder_module, monkeypatch, tmp_path):
schema = folder_module.MiniMaxH3FoldAdaLN.define_schema()
assert schema.node_id == "MiniMaxH3FoldAdaLN"
assert schema.is_output_node is True
input_ids = [item.id for item in schema.inputs]
assert "model_name" in input_ids
assert "output_filename" in input_ids
assert "exclude_patterns" in input_ids
assert "include_mode" in input_ids
assert "discard_patterns" in input_ids
assert "glob_patterns" in input_ids
# Every input must have a non-empty static tooltip
for item in schema.inputs:
assert hasattr(item, "tooltip") and item.tooltip.strip()
output_displays = [item.display_name for item in schema.outputs]
assert output_displays == ["output_path", "report"]
# Test execution
models_dir = tmp_path / "models"
diff_dir = models_dir / "diffusion_models"
diff_dir.mkdir(parents=True)
source_file = diff_dir / "test_h3_model.safetensors"
save_file(_make_h3_checkpoint("diffusion_model."), str(source_file))
monkeypatch.setattr(folder_module.folder_paths, "models_dir", str(models_dir))
monkeypatch.setattr(
folder_module.folder_paths,
"get_full_path_or_raise",
lambda folder, name: str(diff_dir / name),
)
result = folder_module.MiniMaxH3FoldAdaLN.execute(
model_name="test_h3_model.safetensors",
output_filename="sub/folded_h3",
process_device="cpu",
)
assert result.args[0] == "sub/folded_h3.safetensors"
assert "MiniMax H3 AdaLN Fold complete" in result.args[1]
assert (diff_dir / "sub" / "folded_h3.safetensors").is_file()
def test_unfold_h3_diffusion_model_reconstructs_with_base(folder_module, tmp_path):
base_file = tmp_path / "model_base.safetensors"
folded_file = tmp_path / "model_folded.safetensors"
unfolded_file = tmp_path / "model_unfolded.safetensors"
original_base = _make_h3_checkpoint()
save_file(original_base, str(base_file))
# Fold base model to create folded model A
folder_module.fold_minimax_h3_diffusion_model(
str(base_file), str(folded_file), process_device="cpu"
)
# Unfold model A using base model B
report = folder_module.unfold_minimax_h3_diffusion_model(
str(folded_file), str(base_file), str(unfolded_file), process_device="cpu"
)
assert "MiniMax H3 AdaLN Unfold complete" in report
assert "AdaLN layers unfolded: 2" in report
reconstructed = load_file(str(unfolded_file))
# adaln_t_table should be dropped in unfolded output
assert "adaln_t_table" not in reconstructed
# time_embedder restored from base model B
assert "time_embedder.proj_in.weight" in reconstructed
assert "time_embedder.proj_out.weight" in reconstructed
assert torch.equal(reconstructed["time_embedder.proj_in.weight"], original_base["time_embedder.proj_in.weight"])
assert torch.equal(reconstructed["time_embedder.proj_out.weight"], original_base["time_embedder.proj_out.weight"])
# AdaLN linear weights reconstructed to full width [128, 2688]
assert reconstructed["blocks.0.adaln_proj.linear.weight"].shape == (128, 2688)
assert reconstructed["blocks.1.adaln_proj.linear.weight"].shape == (128, 2688)
# Other tensors preserved
assert torch.equal(reconstructed["blocks.0.attn.qkv_proj.weight"], original_base["blocks.0.attn.qkv_proj.weight"])
def test_unfold_h3_node_schema_and_execution(folder_module, monkeypatch, tmp_path):
schema = folder_module.MiniMaxH3UnfoldAdaLN.define_schema()
assert schema.node_id == "MiniMaxH3UnfoldAdaLN"
assert schema.is_output_node is True
input_ids = [item.id for item in schema.inputs]
assert input_ids == [
"model_a", "model_b", "output_filename", "discard_patterns",
"glob_patterns", "curve_rank", "curve_grid", "process_device",
"exclude_patterns", "include_mode",
]
assert input_ids[-1] == "include_mode"
for item in schema.inputs:
assert hasattr(item, "tooltip") and item.tooltip.strip()
models_dir = tmp_path / "models"
diff_dir = models_dir / "diffusion_models"
diff_dir.mkdir(parents=True)
base_file = diff_dir / "base.safetensors"
folded_file = diff_dir / "folded.safetensors"
save_file(_make_h3_checkpoint(), str(base_file))
folder_module.fold_minimax_h3_diffusion_model(str(base_file), str(folded_file), process_device="cpu")
monkeypatch.setattr(folder_module.folder_paths, "models_dir", str(models_dir))
monkeypatch.setattr(
folder_module.folder_paths,
"get_full_path_or_raise",
lambda folder, name: str(diff_dir / name),
)
result = folder_module.MiniMaxH3UnfoldAdaLN.execute(
model_a="folded.safetensors",
model_b="base.safetensors",
output_filename="sub/unfolded_h3",
process_device="cpu",
)
assert result.args[0] == "sub/unfolded_h3.safetensors"
assert "MiniMax H3 AdaLN Unfold complete" in result.args[1]
assert (diff_dir / "sub" / "unfolded_h3.safetensors").is_file()
+226
View File
@@ -0,0 +1,226 @@
import importlib
from pathlib import Path
import sys
import types
import pytest
from safetensors.torch import load_file, save_file
import torch
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_minimax_convert_tests"
def _load_module(name: str):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def converter():
return _load_module("nodes.minimax_h3_lora_convert")
def _make_source_tensors(prefix: str = "transformer_blocks.0") -> dict[str, torch.Tensor]:
tensors = {}
for offset, part in enumerate(("q", "k", "v"), start=1):
# Down: rank 2, in_dim 3
tensors[f"{prefix}.attn.to_{part}.lora_A.default.weight"] = torch.tensor(
[[1.0, float(offset), 0.0], [0.0, 1.0, float(offset)]],
dtype=torch.float32,
)
# Up: out_dim 2, rank 2
tensors[f"{prefix}.attn.to_{part}.lora_B.default.weight"] = torch.tensor(
[[float(offset), 0.0], [0.0, float(offset) + 1.0]],
dtype=torch.float32,
)
# out_proj: rank 2, in_dim 3, out_dim 2
tensors[f"{prefix}.attn.to_out.0.lora_A.default.weight"] = torch.arange(
6, dtype=torch.float32
).reshape(2, 3)
tensors[f"{prefix}.attn.to_out.0.lora_B.default.weight"] = torch.arange(
4, dtype=torch.float32
).reshape(2, 2)
# mlp.fc1 (SwiGLU): rank 2, in_dim 3, out_dim 4 (halves of size 2)
tensors[f"{prefix}.ff.net.0.proj.lora_A.default.weight"] = torch.ones(
(2, 3), dtype=torch.float32
)
# Val half: rows 0..1 (all 1s), Gate half: rows 2..3 (all 2s)
tensors[f"{prefix}.ff.net.0.proj.lora_B.default.weight"] = torch.cat(
[torch.full((2, 2), 1.0), torch.full((2, 2), 2.0)], dim=0
)
# mlp.fc2: rank 2, in_dim 4, out_dim 3
tensors[f"{prefix}.ff.net.2.lora_A.default.weight"] = torch.ones(
(2, 4), dtype=torch.float32
)
tensors[f"{prefix}.ff.net.2.lora_B.default.weight"] = torch.ones(
(3, 2), dtype=torch.float32
)
return tensors
def test_full_rank_qkv_conversion_is_exact(converter, tmp_path):
source = tmp_path / "source_exact.safetensors"
output = tmp_path / "output_exact.safetensors"
original = _make_source_tensors("token_refiner.refiner_blocks.0")
save_file(original, str(source))
report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
assert "MiniMax H3 Diffusers LoRA conversion complete" in report
assert "Fused QKV modules: 1" in report
converted = load_file(str(output))
down = converted["diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_A.weight"]
up = converted["diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_B.weight"]
# Ranks: 2 + 2 + 2 = 6, in_dim: 3, out_dim: 2 + 2 + 2 = 6
assert down.shape == (6, 3)
assert up.shape == (6, 6)
# Mathematical delta equality: up @ down must equal stacked (up_q @ down_q, up_k @ down_k, up_v @ down_v)
fused_delta = up @ down
expected_delta = torch.cat(
[
original[f"token_refiner.refiner_blocks.0.attn.to_{p}.lora_B.default.weight"]
@ original[f"token_refiner.refiner_blocks.0.attn.to_{p}.lora_A.default.weight"]
for p in ("q", "k", "v")
],
dim=0,
)
assert torch.allclose(fused_delta, expected_delta, atol=1e-5, rtol=1e-5)
def test_swiglu_mlp_fc1_halves_swapped(converter, tmp_path):
source = tmp_path / "source_swiglu.safetensors"
output = tmp_path / "output_swiglu.safetensors"
original = _make_source_tensors("transformer_blocks.0")
save_file(original, str(source))
converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
converted = load_file(str(output))
fc1_b = converted["diffusion_model.blocks.0.mlp.fc1.lora_B.weight"]
orig_b = original["transformer_blocks.0.ff.net.0.proj.lora_B.default.weight"]
# In source: top half (rows 0..1) is 1.0, bottom half (rows 2..3) is 2.0
# In converted: top half must be 2.0 (gate), bottom half must be 1.0 (value)
assert torch.equal(fc1_b[:2], orig_b[2:])
assert torch.equal(fc1_b[2:], orig_b[:2])
def test_rejects_incomplete_qkv_group(converter, tmp_path):
source = _make_source_tensors()
del source["transformer_blocks.0.attn.to_v.lora_B.default.weight"]
path = tmp_path / "incomplete.safetensors"
save_file(source, str(path))
with pytest.raises(ValueError, match="Incomplete LoRA A/B pair"):
converter.convert_minimax_h3_diffusers_lora(
str(path), str(tmp_path / "unused.safetensors")
)
def test_rejects_unsupported_module(converter, tmp_path):
source = {"unknown_layer.lora_A.weight": torch.ones(2, 2), "unknown_layer.lora_B.weight": torch.ones(2, 2)}
path = tmp_path / "unknown.safetensors"
save_file(source, str(path))
with pytest.raises(ValueError, match="Unsupported MiniMax H3 Diffusers module"):
converter.convert_minimax_h3_diffusers_lora(
str(path), str(tmp_path / "unused.safetensors")
)
def test_node_schema_and_execution_contract(converter, monkeypatch, tmp_path):
schema = converter.MiniMaxH3DiffusersLoRAConvert.define_schema()
assert schema.node_id == "MiniMaxH3DiffusersLoRAConvert"
assert schema.is_output_node is True
input_ids = [item.id for item in schema.inputs]
assert input_ids == ["lora_name", "output_filename"]
output_displays = [item.display_name for item in schema.outputs]
assert output_displays == ["output_path", "conversion_report"]
# Test execution through node
models_dir = tmp_path / "models"
loras_dir = models_dir / "loras"
loras_dir.mkdir(parents=True)
source_file = loras_dir / "test_diffusers_lora.safetensors"
save_file(_make_source_tensors(), str(source_file))
monkeypatch.setattr(converter.folder_paths, "models_dir", str(models_dir))
monkeypatch.setattr(
converter.folder_paths,
"get_full_path_or_raise",
lambda folder, name: str(loras_dir / name),
)
result = converter.MiniMaxH3DiffusersLoRAConvert.execute(
lora_name="test_diffusers_lora.safetensors",
output_filename="sub/converted_h3",
)
assert result.args[0] == "sub/converted_h3.safetensors"
assert "MiniMax H3 Diffusers LoRA conversion complete" in result.args[1]
assert (loras_dir / "sub" / "converted_h3.safetensors").is_file()
def test_lora_down_up_suffixes_supported(converter, tmp_path):
source = tmp_path / "source_down_up.safetensors"
output = tmp_path / "output_down_up.safetensors"
original = _make_source_tensors("transformer_blocks.0")
# Replace .lora_A/.lora_B with .lora.down/.lora.up on attention layers (DMAD format)
modified = {}
for k, v in original.items():
if ".attn." in k:
k = k.replace(".lora_A.default.weight", ".lora.down.weight")
k = k.replace(".lora_B.default.weight", ".lora.up.weight")
else:
k = k.replace(".default.weight", ".weight")
modified[k] = v
save_file(modified, str(source))
report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
assert "MiniMax H3 Diffusers LoRA conversion complete" in report
converted = load_file(str(output))
assert "diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight" in converted
assert "diffusion_model.blocks.0.attn.qkv_proj.lora_B.weight" in converted
assert "diffusion_model.blocks.0.mlp.fc1.lora_A.weight" in converted
assert "diffusion_model.blocks.0.mlp.fc1.lora_B.weight" in converted
def test_bare_lora_ab_without_weight_suffix_supported(converter, tmp_path):
source = tmp_path / "source_bare_ab.safetensors"
output = tmp_path / "output_bare_ab.safetensors"
original = _make_source_tensors("token_refiner.refiner_blocks.0")
modified = {
k.replace(".default.weight", ""): v
for k, v in original.items()
}
save_file(modified, str(source))
report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output))
assert "MiniMax H3 Diffusers LoRA conversion complete" in report
converted = load_file(str(output))
assert "diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_A.weight" in converted
assert "diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_B.weight" in converted
assert "diffusion_model.token_refiner.blocks.0.mlp.fc1.lora_A.weight" in converted
assert "diffusion_model.token_refiner.blocks.0.mlp.fc1.lora_B.weight" in converted
+481
View File
@@ -0,0 +1,481 @@
import importlib
import csv
from io import StringIO
import json
import math
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_analysis_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def analysis():
return _load_module("nodes.model_analysis")
def _params(**overrides):
params = {
"cwb_similarity_alignment": False,
"top_weight_differences": 3,
"process_device": "cpu",
"force_clear_cache": False,
"exclude_patterns": "",
"glob_patterns": False,
}
params.update(overrides)
return params
def test_five_node_schema_contract(analysis):
assert len(analysis.MODEL_ANALYSIS_NODES) == 5
assert len({node.NODE_ID for node in analysis.MODEL_ANALYSIS_NODES}) == 5
for node in analysis.MODEL_ANALYSIS_NODES:
schema = node.define_schema()
assert schema.category == "ModelUtils/Analysis"
assert schema.description
assert all(item.tooltip for item in schema.inputs)
assert [item.id for item in schema.inputs[:3]] == [
"execution_mode", "model_a", "model_b",
]
assert schema.inputs[0].options == ["ANALYZE", "DOCUMENTATION ONLY"]
assert "lazy_load" not in [item.id for item in schema.inputs]
input_ids = [item.id for item in schema.inputs]
assert input_ids[-3:] == ["exclude_patterns", "glob_patterns", "include_mode"]
assert schema.inputs[-1].default is False
if node in (analysis.LoRAModelAnalysis, analysis.EmbeddingModelAnalysis):
assert input_ids[3] == "cwb_similarity_alignment"
else:
assert "cwb_similarity_alignment" not in input_ids
assert not {
"cwb_preset", "consensus_type", "alignment_threshold",
"similarity_threshold", "power_alpha", "diversity_beta",
"dynamic_similarity_contrast", "soft_comfort_bandpass",
"position_weight", "preserve_common_prefix",
} & set(input_ids)
assert [item.display_name for item in schema.outputs] == [
"comparison_report", "cwb_report", "documentation",
"layerwise_metrics_csv", "layerwise_cwb_csv",
]
def test_documentation_mode_preserves_reports_and_returns_empty_csv(analysis):
result = analysis.DiffusionModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")
assert len(result.result) == 5
assert "No comparison performed" in result[0]
assert "No CWB analysis performed" in result[1]
assert result[2].startswith("# Two-Model Similarity Analysis")
assert result[3] == result[4] == ""
def test_raw_metrics_have_independent_expected_values(analysis):
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[1.0, 0.0], [5.0, 4.0]])
raw, top = analysis._compute_raw(a, b, 2)
metrics = analysis._metrics(raw)
assert metrics["mae"] == pytest.approx(1.0)
assert metrics["mse"] == pytest.approx(2.0)
assert metrics["rmse"] == pytest.approx(math.sqrt(2.0))
assert metrics["max"] == pytest.approx(2.0)
assert metrics["exact"] == pytest.approx(0.5)
assert len(top) == 2
assert {entry[1] for entry in top} == {1, 2}
def test_cwb_derived_metrics_have_independent_expected_values(analysis):
a = torch.tensor([[1.0, 0.0], [1.0, 0.0]])
b = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
values = analysis._compute_cwb(a, b, False).values()
assert values["pair_mean"] == pytest.approx(0.5)
assert values["pair_min"] == pytest.approx(0.0)
assert values["pair_max"] == pytest.approx(1.0)
assert values["mean_aff_a"] == pytest.approx((1.0 + 2 ** -0.5) / 2)
assert values["mean_aff_b"] == pytest.approx((1.0 + 2 ** -0.5) / 2)
assert values["median_aff_a"] == pytest.approx(0.5)
assert values["median_aff_b"] == pytest.approx(0.5)
assert values["alignment"] is None
def test_zero_constant_and_nonfinite_metrics_are_explicit(analysis):
raw, _ = analysis._compute_raw(torch.zeros(3), torch.zeros(3), 0)
metrics = analysis._metrics(raw)
assert metrics["cosine"] == pytest.approx(1.0)
assert metrics["pearson"] is None
assert metrics["norm_ratio"] == pytest.approx(1.0)
raw, _ = analysis._compute_raw(
torch.tensor([1.0, float("nan"), float("inf")]),
torch.tensor([2.0, 4.0, 5.0]),
0,
)
metrics = analysis._metrics(raw)
assert raw.finite == 1
assert raw.nonfinite_a == 2
assert metrics["coverage"] == pytest.approx(1 / 3)
def test_end_to_end_reports_and_topology_classification(monkeypatch, tmp_path, analysis):
tensor_key = 'model.blocks.0.weight|"quoted",\nname'
path_a = tmp_path / "a.safetensors"
path_b = tmp_path / "b.safetensors"
save_file({
tensor_key: torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
"only_a": torch.ones(1),
"shape": torch.ones(2),
"integer": torch.tensor([1, 2], dtype=torch.int64),
}, str(path_a), metadata={"source": "a"})
save_file({
tensor_key: torch.tensor([[1.0, 0.0], [5.0, 4.0]]),
"only_b": torch.ones(1),
"shape": torch.ones(3),
"integer": torch.tensor([1, 3], dtype=torch.int64),
}, str(path_b), metadata={"source": "b"})
paths = {"a": str(path_a), "b": str(path_b)}
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
processed = []
async_calls = []
original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed
original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream
def track_processed(loader, key):
processed.append((loader.filename, key))
return original_mark_processed(loader, key)
def track_async_stream(loader, keys, **kwargs):
async_calls.append((loader.filename, tuple(keys), kwargs))
yield from original_async_stream(loader, keys, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream)
result = analysis.DiffusionModelAnalysis.execute(
model_a="a", model_b="b", execution_mode="ANALYZE", **_params()
)
comparison, cwb_report, documentation, metrics_csv, cwb_csv = result.result
assert documentation.startswith("# Two-Model Similarity Analysis")
assert comparison.startswith("# MODEL COMPARISON SUMMARY\n\n")
assert "| Comparable floating tensors | 1 |" in comparison
assert "| mae | 1 |" in comparison
assert "| mse | 2 |" in comparison
assert '| model.blocks.0.weight\\|"quoted",<br>name |' in comparison
assert "| --- | --- |" in comparison
metrics_rows = list(csv.DictReader(StringIO(metrics_csv)))
cwb_rows = list(csv.DictReader(StringIO(cwb_csv)))
assert len(metrics_rows) == len(cwb_rows) == 1
assert metrics_rows[0]["key"] == cwb_rows[0]["key"] == tensor_key
assert float(metrics_rows[0]["MAE"]) == 1
assert float(metrics_rows[0]["RMSE"]) == pytest.approx(math.sqrt(2), rel=1e-5)
assert metrics_rows[0]["dtype_A"] == metrics_rows[0]["dtype_B"] == "torch.float32"
assert cwb_rows[0]["matched"] == "N/A"
assert float(cwb_rows[0]["pair_cos_min"]) <= float(cwb_rows[0]["pair_cos_mean"]) <= float(cwb_rows[0]["pair_cos_max"])
top_rows = comparison.split("## TOP SCALAR WEIGHT DIFFERENCES\n\n", 1)[1].split("\n\n", 1)[0].splitlines()[2:]
assert len(top_rows) == 3
assert [float(row.rsplit(" | ", 4)[2]) for row in top_rows] == [2, 2, 0]
assert "[A ONLY] (1)" in comparison
assert "[B ONLY] (1)" in comparison
assert "[SHAPE MISMATCH] (1)" in comparison
assert "[NON-FLOATING] (1)" in comparison
assert "model.blocks.0 |" in comparison
assert "CWB-DERIVED SIMILARITY SUMMARY" in cwb_report
assert "merge weighting" in cwb_report
assert "PRESET-IMPLIED CONTRIBUTIONS" not in cwb_report
assert "Average contribution" not in cwb_report
assert len(processed) == 4
assert {key for _, key in processed} == {tensor_key, "integer"}
assert len(async_calls) == 4
assert all(len(call[1]) == 1 for call in async_calls)
assert all(call[2]["batch_size"] == 1 for call in async_calls)
assert all(call[2]["prefetch_batches"] == 1 for call in async_calls)
assert not list(tmp_path.glob("**/*analysis*"))
def test_lora_cross_format_keys_compare_as_one_logical_pair(monkeypatch, tmp_path, analysis):
path_a = tmp_path / "lora_a.safetensors"
path_b = tmp_path / "lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
}, str(path_a))
save_file({
"base_model.model.diffusion_model.foo.lora_down.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
"base_model.model.diffusion_model.foo.lora_up.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
}, str(path_b))
paths = {"a": str(path_a), "b": str(path_b)}
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
processed = []
async_calls = []
original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed
original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream
def track_processed(loader, key):
processed.append((loader.filename, key))
return original_mark_processed(loader, key)
def track_async_stream(loader, keys, **kwargs):
async_calls.append((loader.filename, tuple(keys), kwargs))
yield from original_async_stream(loader, keys, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream)
comparison, cwb_report, metrics_csv, cwb_csv = analysis.ModelAnalysisLogic.execute(
"a", "b", "loras", _params(), lora_mode=True,
)
assert "| Comparable floating tensors | 2 |" in comparison
assert "[A ONLY] (0)" in comparison
assert "[B ONLY] (0)" in comparison
assert "foo.down |" in comparison
assert "foo.up |" in comparison
assert "foo.\\[down+up\\] |" in cwb_report
assert [row["key"] for row in csv.DictReader(StringIO(metrics_csv))] == ["foo.down", "foo.up"]
assert [row["key"] for row in csv.DictReader(StringIO(cwb_csv))] == ["foo.[down+up]"]
assert len(processed) == 4
assert len({entry for entry in processed}) == 4
assert len(async_calls) == 2
assert all(len(call[1]) == 2 for call in async_calls)
@pytest.mark.parametrize("patterns, glob_patterns, include_mode", [
(r"^diffusion_model\.blocks\.4\.", False, False),
(r"^diffusion_model\.blocks\.3\.", False, True),
("diffusion_model.blocks.3.*", True, True),
])
def test_excluded_lora_blocks_do_not_enter_metrics(
monkeypatch, tmp_path, analysis, patterns, glob_patterns, include_mode,
):
path_a = tmp_path / "lora_a.safetensors"
path_b = tmp_path / "lora_b.safetensors"
tensors_a = {}
tensors_b = {}
for block, value_a, value_b in ((3, 1.0, 2.0), (4, 10.0, 20.0)):
prefix = f"diffusion_model.blocks.{block}.attn"
tensors_a[f"{prefix}.lora_A.weight"] = torch.full((1, 2), value_a)
tensors_a[f"{prefix}.lora_B.weight"] = torch.full((2, 1), value_a)
tensors_b[f"{prefix}.lora_A.weight"] = torch.full((1, 2), value_b)
tensors_b[f"{prefix}.lora_B.weight"] = torch.full((2, 1), value_b)
save_file(tensors_a, str(path_a))
save_file(tensors_b, str(path_b))
paths = {"a": str(path_a), "b": str(path_b)}
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
comparison, cwb_report, metrics_csv, cwb_csv = analysis.ModelAnalysisLogic.execute(
"a", "b", "loras",
_params(exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=include_mode),
lora_mode=True,
)
assert "| Comparable floating tensors | 2 |" in comparison
assert "| Excluded A tensors | 2 |" in comparison
assert "| Excluded B tensors | 2 |" in comparison
layerwise = comparison.split("LAYERWISE TENSOR METRICS", 1)[1].split(
"TOP SCALAR WEIGHT DIFFERENCES", 1
)[0]
assert "blocks\\_3\\_" in layerwise
assert "blocks\\_4\\_" not in layerwise
assert "blocks\\_4\\_" not in cwb_report
assert all("blocks_3_" in row["key"] for row in csv.DictReader(StringIO(metrics_csv)))
assert all("blocks_3_" in row["key"] for row in csv.DictReader(StringIO(cwb_csv)))
@pytest.mark.parametrize("patterns, glob_patterns, expected_keys", [
(r"^blocks\.0\.", False, ["blocks.0.weight"]),
("blocks.0.*", True, ["blocks.0.weight"]),
("", False, []),
("no_such_layer", True, []),
])
def test_include_filter_limits_analysis_inventory(
monkeypatch, tmp_path, analysis, patterns, glob_patterns, expected_keys,
):
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
for name, value in (("a", 1.0), ("b", 2.0)):
save_file({
"blocks.0.weight": torch.full((2, 2), value),
"blocks.1.weight": torch.full((2, 2), value),
f"only_{name}.weight": torch.full((2, 2), value),
}, paths[name])
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
streamed = []
original_stream = analysis.MemoryEfficientSafeOpen.async_stream
def stream(handler, keys, *args, **kwargs):
streamed.extend(keys)
return original_stream(handler, keys, *args, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream)
analysis.ModelAnalysisLogic.execute(
"a", "b", "diffusion_models",
_params(exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=True),
)
assert sorted(streamed) == sorted(expected_keys * 2)
def test_diffusion_analysis_decodes_int8_tensorwise_and_skips_sidecars(
monkeypatch, tmp_path, analysis,
):
prefix = "model.diffusion_model.layer"
quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
save_file({
f"{prefix}.weight": torch.tensor([1, 2], dtype=torch.int8),
f"{prefix}.weight_scale": torch.tensor(2.0),
f"{prefix}.comfy_quant": quant,
"dense.weight": torch.tensor([2.0, 4.0]),
}, paths["a"])
save_file({
f"{prefix}.weight": torch.tensor([3, 6], dtype=torch.int8),
f"{prefix}.weight_scale": torch.tensor(2.0),
f"{prefix}.comfy_quant": quant,
"dense.weight": torch.tensor([4.0, 8.0]),
}, paths["b"])
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
streamed = []
original_stream = analysis.MemoryEfficientSafeOpen.async_stream
def track_stream(handler, keys, *args, **kwargs):
streamed.append((handler.filename, tuple(keys)))
return original_stream(handler, keys, *args, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_stream)
_, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute(
"a", "b", "diffusion_models", _params(),
)
rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))}
assert set(rows) == {"layer.weight", "dense.weight"}
assert float(rows["layer.weight"]["MAE"]) == pytest.approx(6.0)
assert float(rows["dense.weight"]["MAE"]) == pytest.approx(3.0)
for path, keys in streamed:
if path.endswith("a.safetensors") or path.endswith("b.safetensors"):
if f"{prefix}.weight" in keys:
assert f"{prefix}.weight_scale" in keys
assert f"{prefix}.comfy_quant" in keys
assert sum(f"{prefix}.weight" in keys or "dense.weight" in keys for _, keys in streamed) == 4
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_diffusion_analysis_matches_prefixed_and_bare_layers(
monkeypatch, tmp_path, analysis, prefixed_input,
):
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
prefix = "model.diffusion_model."
for name in ("a", "b"):
key = f"{prefix if name == prefixed_input else ''}layer.weight"
if name == prefixed_input:
stem = key[:-len(".weight")]
save_file({
key: torch.tensor([1, 2], dtype=torch.int8),
f"{stem}.weight_scale": torch.tensor(2.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
}, paths[name])
else:
save_file({key: torch.tensor([4.0, 8.0])}, paths[name])
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
_, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute(
"a", "b", "diffusion_models", _params(),
)
rows = list(csv.DictReader(StringIO(metrics_csv)))
assert [row["key"] for row in rows] == ["layer.weight"]
assert float(rows[0]["MAE"]) == pytest.approx(3.0)
def test_diffusion_analysis_rejects_duplicate_prefix_variants(monkeypatch, tmp_path, analysis):
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
save_file({
"layer.weight": torch.ones(2),
"model.diffusion_model.layer.weight": torch.ones(2),
}, paths["a"])
save_file({"layer.weight": torch.ones(2)}, paths["b"])
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
with pytest.raises(ValueError, match="duplicate diffusion layer 'layer.weight'"):
analysis.ModelAnalysisLogic.execute("a", "b", "diffusion_models", _params())
def test_embedding_similarity_alignment_is_explicit_opt_in(analysis):
a = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
b = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
aligned = analysis._compute_cwb(a, b, True).values()
indexed = analysis._compute_cwb(a, b, False).values()
assert aligned["alignment"] == pytest.approx(1.0)
assert aligned["alignment_gain"] == pytest.approx(1.0)
assert indexed["pair_mean"] == pytest.approx(0.0)
assert indexed["alignment"] is None
def test_cuda_oom_retries_once_on_cpu(monkeypatch, analysis):
calls = []
def fake_compute(a, b, device, allow_alignment, top_count):
calls.append(device)
if device == "cuda":
raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory")
return analysis.RawStats(finite=1, total=1), analysis.CWBStats(), []
monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
result = analysis._analyze_with_fallback(
"layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
)
assert calls == ["cuda", "cpu"]
assert result[-1] is True
def test_non_oom_error_is_not_retried(monkeypatch, analysis):
calls = []
def fake_compute(*args):
calls.append(args[2])
raise RuntimeError("unrelated failure")
monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
with pytest.raises(RuntimeError, match="unrelated failure"):
analysis._analyze_with_fallback(
"layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
)
assert calls == ["cuda"]
def test_failed_cpu_retry_names_layer_and_chains_oom(monkeypatch, analysis):
def fake_compute(a, b, device, allow_alignment, top_count):
if device == "cuda":
raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory")
raise ValueError("CPU calculation failed")
monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
with pytest.raises(RuntimeError, match="CPU fallback failed while analyzing 'named.layer'") as caught:
analysis._analyze_with_fallback(
"named.layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
)
assert isinstance(caught.value.__cause__, torch.cuda.OutOfMemoryError)
+56
View File
@@ -0,0 +1,56 @@
import importlib.util
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
@pytest.fixture(scope="module")
def runner():
spec = importlib.util.spec_from_file_location(
"model_analysis_api_runner_test",
REPO_ROOT / "run_model_analysis_api_workflow.py",
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def test_linked_previews_require_both_reports(runner):
workflow = {
"1": {"class_type": "LoRAModelAnalysis", "inputs": {}},
"2": {"class_type": "PreviewAny", "inputs": {"source": ["1", 0]}},
"3": {"class_type": "PreviewAny", "inputs": {"source": ["1", 1]}},
}
assert runner._linked_previews(workflow, "1") == {0: "2", 1: "3"}
def test_linked_previews_reject_missing_report(runner):
workflow = {
"1": {"class_type": "LoRAModelAnalysis", "inputs": {}},
"2": {"class_type": "PreviewAny", "inputs": {"source": ["1", 0]}},
}
with pytest.raises(RuntimeError, match="output"):
runner._linked_previews(workflow, "1")
def test_run_analysis_captures_both_reports(monkeypatch, runner):
workflow = {
"1": {"class_type": "LoRAModelAnalysis", "inputs": {}},
"2": {"class_type": "PreviewAny", "inputs": {"source": ["1", 0]}},
"3": {"class_type": "PreviewAny", "inputs": {"source": ["1", 1]}},
}
monkeypatch.setattr(runner.api, "_json_request", lambda *args: {"prompt_id": "p"})
monkeypatch.setattr(
runner.api,
"_wait_for_history",
lambda *args: {"outputs": {"2": {"text": ["standard"]}, "3": {"text": ["cwb"]}}},
)
result = runner.run_analysis(workflow, "http://server", 0.1)
assert result == {
"prompt_id": "p",
"comparison_report": "standard",
"cwb_report": "cwb",
}
+153
View File
@@ -0,0 +1,153 @@
import ast
import importlib
import inspect
import sys
import types
from pathlib import Path
import pytest
from comfy_api.latest import io
NODE_DIR = Path(__file__).resolve().parents[1] / "nodes"
def test_every_schema_input_has_a_direct_static_tooltip():
missing = []
non_static = []
for path in sorted(NODE_DIR.glob("*.py")):
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
for call in ast.walk(tree):
if not (
isinstance(call, ast.Call)
and isinstance(call.func, ast.Attribute)
and call.func.attr == "Input"
):
continue
name = (
call.args[0].value
if call.args and isinstance(call.args[0], ast.Constant)
else "<dynamic input name>"
)
tooltip = next(
(keyword.value for keyword in call.keywords if keyword.arg == "tooltip"),
None,
)
location = f"{path.name}:{call.lineno}:{name}"
if tooltip is None:
missing.append(location)
elif not (
isinstance(tooltip, ast.Constant)
and isinstance(tooltip.value, str)
and tooltip.value.strip()
):
if not (path.name == "layer_parameters.py" and name == "layer_parameters"):
non_static.append(location)
assert not missing, "Inputs without tooltips:\n" + "\n".join(missing)
assert not non_static, "Tooltips must be direct non-empty string literals:\n" + "\n".join(non_static)
FILTER_MODULES = [
"merger", "consensus_merger", "cwb_model_graft", "cwb_delta_lora_merger", "lodestone_merger",
"lora_resize", "dtype_conversion", "model_analysis",
"lora_extract_svd", "dora_extract_wd", "dora_learned_wd", "text_encoder_extract",
"minimax_h3_fold",
]
def test_filter_schema_audit_covers_all_direct_filter_definitions():
discovered = set()
for path in NODE_DIR.glob("*.py"):
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
for call in ast.walk(tree):
if (
isinstance(call, ast.Call)
and isinstance(call.func, ast.Attribute)
and call.func.attr == "Input"
and call.args
and isinstance(call.args[0], ast.Constant)
and call.args[0].value in {"exclude_patterns", "skip_patterns"}
):
discovered.add(path.stem)
assert discovered <= set(FILTER_MODULES), f"Unaudited filter modules: {discovered - set(FILTER_MODULES)}"
def _load_filter_module(module_name):
package_name = "modelutils_filter_schema_tests"
if package_name not in sys.modules:
package = types.ModuleType(package_name)
package.__path__ = [str(NODE_DIR.parent)]
sys.modules[package_name] = package
return importlib.import_module(f"{package_name}.nodes.{module_name}")
@pytest.mark.parametrize("module_name", FILTER_MODULES)
def test_layer_filter_schemas_append_disabled_include_toggle(module_name):
module = _load_filter_module(module_name)
checked = []
for name, node in vars(module).items():
if name.startswith("_") or not isinstance(node, type):
continue
if node.__module__ != module.__name__ or not issubclass(node, io.ComfyNode):
continue
schema = node.define_schema()
previous_inputs = [item for item in schema.inputs if item.id != "layer_parameters"]
ids = [item.id for item in previous_inputs]
if not {"exclude_patterns", "skip_patterns"}.intersection(ids):
continue
assert ids.count("include_mode") == 1, schema.node_id
assert ids[-1] == "include_mode", schema.node_id
assert isinstance(previous_inputs[-1], io.Boolean.Input), schema.node_id
assert previous_inputs[-1].default is False, schema.node_id
checked.append(schema.node_id)
assert checked, f"No filter nodes checked in {module_name}"
@pytest.mark.parametrize("module_name, helper_name, expected_count", [
("lora_extract_svd", "extract_lora_from_files", 5),
("dora_extract_wd", "extract_dora_from_files", 5),
("dora_learned_wd", "extract_dora_learned_from_files", 5),
("text_encoder_extract", "extract_te_from_files", 10),
])
@pytest.mark.parametrize("include_mode", [False, True])
def test_all_extraction_nodes_forward_include_mode_without_shifting_parameters(
monkeypatch, module_name, helper_name, expected_count, include_mode,
):
module = _load_filter_module(module_name)
signature = inspect.signature(getattr(module, helper_name))
assert [name for name in signature.parameters if name != "layer_parameters"][-1] == "include_mode"
captured = []
def extract(*args, **kwargs):
bound = signature.bind(*args, **kwargs)
bound.apply_defaults()
captured.append(bound.arguments)
monkeypatch.setattr(module, helper_name, extract)
monkeypatch.setattr(module.folder_paths, "get_full_path_or_raise", lambda category, name: name)
monkeypatch.setattr(module, "_build_lora_output_path", lambda name: (name, name))
for name, node in vars(module).items():
if name.startswith("_") or not isinstance(node, type):
continue
if node.__module__ != module.__name__ or not issubclass(node, io.ComfyNode):
continue
schema = node.define_schema()
values = {}
for widget in schema.inputs:
if widget.id == "layer_parameters":
continue
value = getattr(widget, "default", None)
if value is None:
options = getattr(widget, "options", [])
value = options[0] if options else widget.id
values[widget.id] = value
values["include_mode"] = include_mode
values["skip_patterns"] = "selected_layer"
node.execute(**values)
assert captured[-1]["include_mode"] is include_mode, schema.node_id
assert captured[-1]["skip_patterns_str"] == "selected_layer", schema.node_id
assert captured[-1]["knee_probe_offset"] == values.get("knee_probe_offset", 32), schema.node_id
assert len(captured) == expected_count
+114
View File
@@ -0,0 +1,114 @@
import importlib.util
from pathlib import Path
import pytest
import torch
from unifiedefficientloader import IncrementalSafetensorsWriter
REPO_ROOT = Path(__file__).resolve().parents[1]
def _load_uel_io():
spec = importlib.util.spec_from_file_location(
"modelutils_uel_io_test", REPO_ROOT / "nodes" / "uel_io.py"
)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
class _Handler:
def __init__(self, values):
self.values = values
self.calls = []
self.processed = []
self.closed = 0
def async_stream(self, keys, **kwargs):
self.calls.append((list(keys), kwargs))
def generate():
try:
for key in keys:
yield [(key, self.values[key])]
finally:
self.closed += 1
return generate()
def mark_processed(self, key):
self.processed.append(key)
def test_production_nodes_have_no_synchronous_safetensors_io():
violations = []
for path in sorted((REPO_ROOT / "nodes").glob("*.py")):
source = path.read_text(encoding="utf-8")
if ".get_tensor(" in source:
violations.append(f"{path.name}: get_tensor")
if "from safetensors" in source or "import safetensors" in source:
violations.append(f"{path.name}: direct safetensors import")
assert violations == []
def test_work_units_use_bounded_streams_and_release_every_yield():
uel = _load_uel_io()
first = _Handler({"a": torch.tensor([1.0]), "b": torch.tensor([2.0])})
second = _Handler({"x": torch.tensor([3.0])})
units = [("one", {0: ["a"], 1: ["x"]}), ("two", {0: ["b"]})]
observed = []
for logical, loaded in uel.stream_work_units(
{0: first, 1: second}, units, pin_memory=True
):
observed.append((logical, sorted(loaded)))
assert observed == [
("one", [(0, "a"), (1, "x")]),
("two", [(0, "b")]),
]
assert first.processed == ["a", "b"]
assert second.processed == ["x"]
assert first.closed == second.closed == 1
for handler in (first, second):
assert handler.calls[0][1] == {
"batch_size": 1,
"prefetch_batches": 1,
"pin_memory": True,
}
def test_work_unit_failure_releases_consumed_sources_and_closes_streams():
uel = _load_uel_io()
handler = _Handler({"a": torch.tensor([1.0]), "b": torch.tensor([2.0])})
stream = uel.stream_work_units(
{0: handler}, [("one", {0: ["a"]}), ("two", {0: ["b"]})],
pin_memory=False,
)
with pytest.raises(RuntimeError, match="consumer failed"):
try:
for _, _ in stream:
raise RuntimeError("consumer failed")
finally:
stream.close()
assert handler.processed == ["a"]
assert handler.closed == 1
def test_atomic_writer_preserves_destination_on_failure(tmp_path):
uel = _load_uel_io()
output = tmp_path / "output.safetensors"
with IncrementalSafetensorsWriter(str(output), max_workers=1) as writer:
writer.write("old", torch.tensor([1.0]))
original = output.read_bytes()
with pytest.raises(RuntimeError, match="planned failure"):
with uel.atomic_uel_writer(str(output)) as writer:
writer.write("new", torch.tensor([2.0]))
raise RuntimeError("planned failure")
assert output.read_bytes() == original
assert list(tmp_path.glob(".output.safetensors.*.tmp")) == []