diff --git a/README.md b/README.md
index 585e6cb..736d4bb 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,55 @@
# ComfyUI-ModelUtils
-[WIP]Custom nodes for handling, inspecting, modifying and creating various model files.
+
+A collection of ComfyUI custom nodes for inspecting, modifying, merging, and creating model files. Supports Models, TextEncoders, LoRAs, Checkpoints, and Embeddings.
+
+## Features
+
+- **MetaKeys** – Inspect and display metadata and tensor keys from model files
+- **RenameKeys** – Batch rename tensor keys using pattern matching
+- **PruneKeys** – Remove unwanted layers/keys from models
+- **Mergers** – Combine 2 or 3 models with configurable blend modes and ratios
+- **LoRA Extraction** – Extract LoRA adapters from model pairs using various SVD rank selection methods (Fixed, Ratio, Quantile, Knee-detection, Frobenius-norm)
+
+## Example Workflows
+
+
+ 
+
+ 📥 GetMetaAndkeys.json
+
+
+---
+
+
+ 
+
+ 📥 LoRA_Extract_nodes.json
+
+
+---
+
+
+ 
+
+ 📥 Merging_Examples.json
+
+
+---
+
+
+ 
+
+ 📥 RenameKeysInModel.json
+
+
+## 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.
diff --git a/__init__.py b/__init__.py
index 82c8b43..10e45ce 100644
--- a/__init__.py
+++ b/__init__.py
@@ -19,8 +19,9 @@ from .nodes.merger import (
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger
)
-from .nodes.lora_extract import (
- LoRASaveFromFile, LoRACheckpointSaveFromFile
+from .nodes.lora_extract_svd import (
+ LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
+ LoRAExtractKnee, LoRAExtractFrobenius
)
@@ -44,9 +45,11 @@ class ModelUtilsExtension(ComfyExtension):
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger,
# LoRA Extraction
- LoRASaveFromFile, LoRACheckpointSaveFromFile,
+ LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
+ LoRAExtractKnee, LoRAExtractFrobenius,
]
async def comfy_entrypoint() -> ModelUtilsExtension:
- return ModelUtilsExtension()
\ No newline at end of file
+ return ModelUtilsExtension()
+
diff --git a/assets/GetMetaAndkeys.png b/assets/GetMetaAndkeys.png
new file mode 100644
index 0000000..9adf418
Binary files /dev/null and b/assets/GetMetaAndkeys.png differ
diff --git a/assets/LoRA_Extract_nodes.png b/assets/LoRA_Extract_nodes.png
new file mode 100644
index 0000000..6a18e68
Binary files /dev/null and b/assets/LoRA_Extract_nodes.png differ
diff --git a/assets/Merging_Examples.png b/assets/Merging_Examples.png
new file mode 100644
index 0000000..565f64e
Binary files /dev/null and b/assets/Merging_Examples.png differ
diff --git a/assets/RenameKeysInModel.png b/assets/RenameKeysInModel.png
new file mode 100644
index 0000000..60f5858
Binary files /dev/null and b/assets/RenameKeysInModel.png differ
diff --git a/example_workflows/GetMetaAndkeys.json b/example_workflows/GetMetaAndkeys.json
new file mode 100644
index 0000000..b7395a3
--- /dev/null
+++ b/example_workflows/GetMetaAndkeys.json
@@ -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
+}
\ No newline at end of file
diff --git a/example_workflows/LoRA_Extract_nodes.json b/example_workflows/LoRA_Extract_nodes.json
new file mode 100644
index 0000000..62c4668
--- /dev/null
+++ b/example_workflows/LoRA_Extract_nodes.json
@@ -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
+}
\ No newline at end of file
diff --git a/example_workflows/Merging_Examples.json b/example_workflows/Merging_Examples.json
new file mode 100644
index 0000000..5f92f11
--- /dev/null
+++ b/example_workflows/Merging_Examples.json
@@ -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
+}
\ No newline at end of file
diff --git a/example_workflows/RenameKeysInModel.json b/example_workflows/RenameKeysInModel.json
new file mode 100644
index 0000000..4fbabcc
--- /dev/null
+++ b/example_workflows/RenameKeysInModel.json
@@ -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
+}
\ No newline at end of file
diff --git a/nodes/device_utils.py b/nodes/device_utils.py
new file mode 100644
index 0000000..6ad7588
--- /dev/null
+++ b/nodes/device_utils.py
@@ -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")
diff --git a/nodes/lora_extract.py b/nodes/lora_extract.py
deleted file mode 100644
index fe56431..0000000
--- a/nodes/lora_extract.py
+++ /dev/null
@@ -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)
diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py
new file mode 100644
index 0000000..d831a8d
--- /dev/null
+++ b/nodes/lora_extract_svd.py
@@ -0,0 +1,817 @@
+"""
+LoRA Extraction via SVD decomposition.
+
+Native implementation of Low-Rank Adapter extraction with multiple rank selection modes.
+No external dependencies required.
+"""
+import os
+import re
+import torch
+import torch.linalg as linalg
+import folder_paths
+import comfy.utils
+from tqdm import tqdm
+from comfy_api.latest import io
+from safetensors.torch import save_file
+from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
+from .device_utils import (
+ estimate_model_size, prepare_for_large_operation,
+ cleanup_after_operation, get_device_capabilities
+)
+
+
+# =============================================================================
+# Rank Index Functions - Determine LoRA rank from singular values
+# =============================================================================
+
+def _index_sv_fixed(S: torch.Tensor, dim: int) -> int:
+ """Fixed rank mode - use specified dimension."""
+ return max(1, min(dim, len(S) - 1))
+
+
+def _index_sv_threshold(S: torch.Tensor, threshold: float) -> int:
+ """Threshold mode - keep singular values above absolute threshold."""
+ rank = int(torch.sum(S > threshold).item())
+ return max(1, min(rank, len(S) - 1))
+
+
+def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int:
+ """Ratio mode - keep singular values > max(S) * ratio."""
+ if ratio <= 0:
+ return len(S) - 1
+ min_sv = S[0] * ratio
+ rank = int(torch.sum(S > min_sv).item())
+ return max(1, min(rank, len(S) - 1))
+
+
+def _index_sv_cumulative(S: torch.Tensor, target: float) -> int:
+ """Cumulative mode - keep enough SVs to reach target % of total."""
+ total = torch.sum(S)
+ if total < 1e-8:
+ return 1
+ cumsum = torch.cumsum(S, dim=0) / total
+ rank = int(torch.searchsorted(cumsum, target).item()) + 1
+ return max(1, min(rank, len(S) - 1))
+
+
+def _index_sv_fro(S: torch.Tensor, target: float) -> int:
+ """Frobenius norm mode - preserve target fraction of Frobenius norm."""
+ S_sq = S.pow(2)
+ total_sq = torch.sum(S_sq)
+ if total_sq < 1e-8:
+ return 1
+ cumsum = torch.cumsum(S_sq, dim=0) / total_sq
+ rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1
+ return max(1, min(rank, len(S) - 1))
+
+
+def _index_sv_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
+ """Knee detection - find elbow point on singular value curve."""
+ n = len(S)
+ if n < 3:
+ return 1
+ s_max, s_min = S[0].item(), S[-1].item()
+ if s_max - s_min < min_sv:
+ return 1
+ # Normalize to [0, 1]
+ s_norm = (S - s_min) / (s_max - s_min)
+ x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
+ # Distance from diagonal
+ distances = (x_norm + s_norm - 1).abs()
+ rank = torch.argmax(distances).item() + 1
+ return max(1, min(rank, n - 1))
+
+
+def _index_sv_cumulative_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
+ """Knee detection on cumulative singular value curve."""
+ n = len(S)
+ if n < 3:
+ return 1
+ total = torch.sum(S)
+ if total < min_sv:
+ return 1
+ cumsum = torch.cumsum(S, dim=0) / total
+ y_min, y_max = cumsum[0].item(), cumsum[-1].item()
+ if y_max - y_min < min_sv:
+ return 1
+ y_norm = (cumsum - y_min) / (y_max - y_min)
+ x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
+ distances = (y_norm - x_norm).abs()
+ rank = torch.argmax(distances).item() + 1
+ return max(1, min(rank, n - 1))
+
+
+def _index_sv_rel_decrease(S: torch.Tensor, tau: float = 0.1) -> int:
+ """Relative decrease mode - stop when consecutive SV ratio drops below tau."""
+ if len(S) < 2:
+ return 1
+ ratios = S[1:] / (S[:-1] + 1e-8)
+ for k in range(len(ratios)):
+ if ratios[k] < tau:
+ return k + 1
+ return len(S)
+
+
+def _compute_rank(S: torch.Tensor, mode: str, mode_param: float,
+ max_rank: int = None) -> int:
+ """Compute rank based on mode and parameters."""
+ if mode == "fixed":
+ rank = _index_sv_fixed(S, int(mode_param))
+ elif mode == "threshold":
+ rank = _index_sv_threshold(S, mode_param)
+ elif mode == "ratio":
+ rank = _index_sv_ratio(S, mode_param)
+ elif mode == "quantile" or mode == "sv_cumulative":
+ rank = _index_sv_cumulative(S, mode_param)
+ elif mode == "sv_fro":
+ rank = _index_sv_fro(S, mode_param)
+ elif mode == "sv_knee":
+ rank = _index_sv_knee(S)
+ elif mode == "sv_cumulative_knee":
+ rank = _index_sv_cumulative_knee(S)
+ elif mode == "sv_rel_decrease":
+ rank = _index_sv_rel_decrease(S, mode_param)
+ else:
+ rank = _index_sv_fixed(S, int(mode_param) if mode_param else 64)
+
+ if max_rank is not None:
+ rank = min(rank, max_rank)
+ return rank
+
+
+# =============================================================================
+# SVD Extraction Functions
+# =============================================================================
+
+def _svd_extract_linear_lowrank(
+ weight: torch.Tensor,
+ rank: int,
+ device: str,
+ clamp_quantile: float = 0.99,
+) -> tuple:
+ """
+ Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
+
+ Much faster than full SVD when rank << min(m, n) because it only computes
+ the top-k singular values using randomized algorithms.
+
+ Returns:
+ (lora_down, lora_up, diff), "low rank" or (weight, "full")
+ """
+ weight = weight.to(device=device, dtype=torch.float32)
+ out_dim, in_dim = weight.shape
+
+ # Clamp rank to valid range
+ max_possible_rank = min(out_dim, in_dim)
+ rank = max(1, min(rank, max_possible_rank - 1))
+
+ # Check if decomposition is worthwhile
+ if rank >= max_possible_rank // 2:
+ return weight, "full"
+
+ # Use svd_lowrank for faster low-rank approximation
+ # niter controls accuracy (higher = more accurate but slower)
+ try:
+ U, S, V = torch.svd_lowrank(weight, q=rank, niter=2)
+ # U: [out_dim, rank], S: [rank], V: [in_dim, rank]
+ Vh = V.T # [rank, in_dim]
+ except Exception as e:
+ raise RuntimeError(f"svd_lowrank failed: {e}")
+
+ # Clamp outliers
+ if clamp_quantile < 1.0 and len(S) > 0:
+ try:
+ max_val = torch.quantile(S, clamp_quantile)
+ S = S.clamp(max=max_val)
+ except RuntimeError:
+ pass
+
+ # Construct LoRA matrices
+ lora_up = U @ torch.diag(S) # [out_dim, rank]
+ lora_down = Vh # [rank, in_dim]
+
+ # Compute reconstruction diff
+ diff = weight - (lora_up @ lora_down)
+
+ return (lora_down, lora_up, diff), "low rank"
+
+
+def _svd_extract_linear(
+ weight: torch.Tensor,
+ mode: str,
+ mode_param: float,
+ device: str,
+ max_rank: int = None,
+ clamp_quantile: float = 0.99,
+) -> tuple:
+ """
+ SVD decomposition for linear layers.
+
+ For fixed rank mode, uses svd_lowrank which is much faster.
+ For adaptive modes, uses full SVD to analyze all singular values.
+
+ Returns:
+ (lora_down, lora_up, rank), "low rank" or (weight, "full")
+ """
+ weight = weight.to(device=device, dtype=torch.float32)
+ out_dim, in_dim = weight.shape
+
+ # For fixed rank, use the faster svd_lowrank
+ if mode == "fixed" and mode_param > 0:
+ target_rank = int(mode_param)
+ if max_rank is not None:
+ target_rank = min(target_rank, max_rank)
+ return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile)
+
+ # For adaptive modes, use full SVD to analyze singular values
+ try:
+ U, S, Vh = linalg.svd(weight, full_matrices=False)
+ except Exception as e:
+ raise RuntimeError(f"SVD failed: {e}")
+
+ # Compute rank based on mode
+ rank = _compute_rank(S, mode, mode_param, max_rank)
+
+ # Check if decomposition is worthwhile
+ if rank >= min(out_dim, in_dim) // 2:
+ return weight, "full"
+
+ # Truncate to rank
+ U = U[:, :rank]
+ S = S[:rank]
+ Vh = Vh[:rank, :]
+
+ # Clamp outliers
+ if clamp_quantile < 1.0 and len(S) > 0:
+ try:
+ max_val = torch.quantile(S, clamp_quantile)
+ S = S.clamp(max=max_val)
+ except RuntimeError:
+ pass # Skip clamping if quantile fails
+
+ # Construct LoRA matrices
+ # lora_up = U @ diag(S), lora_down = Vh
+ lora_up = U @ torch.diag(S) # [out_dim, rank]
+ lora_down = Vh # [rank, in_dim]
+
+ # Compute reconstruction diff for optional sparse bias
+ diff = weight - (lora_up @ lora_down)
+
+ return (lora_down, lora_up, diff), "low rank"
+
+
+def _svd_extract_conv(
+ weight: torch.Tensor,
+ mode: str,
+ mode_param: float,
+ device: str,
+ max_rank: int = None,
+ clamp_quantile: float = 0.99,
+) -> tuple:
+ """
+ SVD decomposition for conv2d layers.
+
+ Returns:
+ (lora_down, lora_up, rank), "low rank" or (weight, "full")
+ """
+ out_ch, in_ch, kh, kw = weight.shape
+ is_1x1 = (kh == 1 and kw == 1)
+
+ # Flatten for SVD
+ if is_1x1:
+ mat = weight.squeeze() # [out_ch, in_ch]
+ else:
+ mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
+
+ result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile)
+
+ if mode_str == "full":
+ return weight, "full"
+
+ lora_down, lora_up, diff = result
+ rank = lora_down.shape[0]
+
+ # Reshape for conv
+ if is_1x1:
+ lora_down = lora_down.unsqueeze(-1).unsqueeze(-1) # [rank, in_ch, 1, 1]
+ lora_up = lora_up.unsqueeze(-1).unsqueeze(-1) # [out_ch, rank, 1, 1]
+ else:
+ lora_down = lora_down.reshape(rank, in_ch, kh, kw)
+ lora_up = lora_up.reshape(out_ch, rank, 1, 1)
+
+ return (lora_down, lora_up, diff), "low rank"
+
+
+# =============================================================================
+# Chunked Extraction for Large Layers
+# =============================================================================
+
+def _detect_fused_layer(key: str, shape: tuple) -> int:
+ """Detect fused QKV/MLP layers and return number of chunks."""
+ if len(shape) != 2:
+ return 1
+
+ out_dim, in_dim = shape
+ key_lower = key.lower()
+
+ # QKV layers (3 chunks)
+ if "qkv" in key_lower and out_dim % 3 == 0 and out_dim // 3 >= in_dim // 4:
+ return 3
+
+ # KV layers (2 chunks)
+ if "kv" in key_lower and "qkv" not in key_lower and out_dim % 2 == 0:
+ return 2
+
+ # Fused MLP (2 chunks)
+ if ("linear1" in key_lower or "gate_up" in key_lower) and out_dim % 2 == 0 and out_dim >= in_dim * 2:
+ return 2
+
+ return 1
+
+
+def _extract_chunked_layer(
+ weight_diff: torch.Tensor,
+ num_chunks: int,
+ mode: str,
+ mode_param: float,
+ device: str,
+ max_rank: int = None,
+) -> tuple:
+ """Extract LoRA from fused layer by chunking."""
+ out_dim, in_dim = weight_diff.shape
+ chunk_size = out_dim // num_chunks
+
+ all_lora_up = []
+ all_lora_down = []
+
+ for i in range(num_chunks):
+ chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :]
+
+ try:
+ result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank)
+ if mode_str == "full":
+ return None, None, 0
+ lora_down, lora_up, _ = result
+ all_lora_down.append(lora_down)
+ all_lora_up.append(lora_up)
+ except Exception:
+ return None, None, 0
+
+ # Combine chunks
+ combined_lora_up = torch.cat(all_lora_up, dim=0)
+ combined_lora_down = torch.stack(all_lora_down, dim=0).mean(dim=0)
+ rank = combined_lora_down.shape[0]
+
+ return combined_lora_up, combined_lora_down, rank
+
+
+# =============================================================================
+# Pattern Matching
+# =============================================================================
+
+def _compile_patterns(pattern_string: str) -> list:
+ """Compile regex patterns from whitespace-separated string."""
+ if not pattern_string or not pattern_string.strip():
+ return []
+ patterns = []
+ for p in pattern_string.split():
+ p = p.strip()
+ if p:
+ try:
+ patterns.append(re.compile(p))
+ except re.error:
+ pass
+ return patterns
+
+
+def _matches_any_pattern(key: str, patterns: list) -> bool:
+ """Check if key matches any pattern."""
+ return any(p.search(key) for p in patterns)
+
+
+# =============================================================================
+# Main Extraction Function
+# =============================================================================
+
+def extract_lora_from_files(
+ model_a_path: str,
+ model_b_path: str,
+ mode: str,
+ linear_param: float,
+ conv_param: float,
+ prefix: str,
+ device: str,
+ save_dtype: str,
+ linear_max_rank: int = None,
+ conv_max_rank: int = None,
+ clamp_quantile: float = 0.99,
+ min_diff: float = 0.0,
+ skip_patterns_str: str = "",
+ mismatch_mode: str = "skip",
+ chunk_large_layers: bool = True,
+) -> dict[str, torch.Tensor]:
+ """
+ Extract LoRA from difference between two models.
+
+ Args:
+ model_a_path: Finetuned model path
+ model_b_path: Base model path
+ mode: Rank selection mode
+ linear_param: Mode parameter for linear layers
+ conv_param: Mode parameter for conv layers
+ prefix: LoRA key prefix (e.g. "lora_unet_")
+ device: Computation device
+ save_dtype: Output dtype
+ linear_max_rank: Max rank for linear layers
+ conv_max_rank: Max rank for conv layers
+ clamp_quantile: Quantile for weight clamping
+ min_diff: Minimum difference threshold
+ skip_patterns_str: Patterns for layers to skip
+ mismatch_mode: How to handle mismatches
+ chunk_large_layers: Enable chunked extraction
+ """
+ output_sd = {}
+
+ save_torch_dtype = {
+ "fp32": torch.float32,
+ "fp16": torch.float16,
+ "bf16": torch.bfloat16,
+ }.get(save_dtype, torch.float16)
+
+ skip_patterns = _compile_patterns(skip_patterns_str)
+
+ # Prepare memory before heavy operation
+ total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
+ print(f"[LoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
+ prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) # 1.5x for SVD headroom
+
+ handler_a = MemoryEfficientSafeOpen(model_a_path)
+ handler_b = MemoryEfficientSafeOpen(model_b_path)
+
+ try:
+ keys_a = set(handler_a.keys())
+ keys_b = set(handler_b.keys())
+
+ weight_keys = [k for k in keys_a if k.endswith(".weight")]
+
+ pbar = comfy.utils.ProgressBar(len(weight_keys))
+
+ stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
+
+ for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
+ lora_key = key[:-7]
+ lora_name = prefix + lora_key.replace(".", "_")
+
+ if _matches_any_pattern(key, skip_patterns):
+ stats["skipped"] += 1
+ pbar.update(1)
+ continue
+
+ # Load tensors with pinned memory for CUDA
+ use_pinned = device == 'cuda'
+
+ if key not in keys_b:
+ if mismatch_mode == "skip":
+ stats["skipped"] += 1
+ pbar.update(1)
+ continue
+ cpu_a = handler_a.get_tensor(key)
+ if use_pinned:
+ weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
+ else:
+ weight_diff = cpu_a.to(device=device, dtype=torch.float32)
+ del cpu_a
+ else:
+ cpu_a = handler_a.get_tensor(key)
+ cpu_b = handler_b.get_tensor(key)
+
+ if use_pinned:
+ tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
+ tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
+ else:
+ tensor_a = cpu_a.to(device=device, dtype=torch.float32)
+ tensor_b = cpu_b.to(device=device, dtype=torch.float32)
+ del cpu_a, cpu_b
+
+ if tensor_a.shape != tensor_b.shape:
+ if mismatch_mode == "skip":
+ stats["skipped"] += 1
+ pbar.update(1)
+ del tensor_a, tensor_b
+ continue
+ weight_diff = tensor_a
+ del tensor_b
+ else:
+ weight_diff = tensor_a - tensor_b
+ del tensor_a, tensor_b
+
+ # Skip small differences
+ if min_diff > 0 and weight_diff.abs().max() < min_diff:
+ stats["skipped"] += 1
+ del weight_diff
+ pbar.update(1)
+ continue
+
+ # Skip 1D tensors
+ if weight_diff.ndim < 2:
+ stats["skipped"] += 1
+ del weight_diff
+ pbar.update(1)
+ continue
+
+ is_conv = weight_diff.ndim == 4
+
+ try:
+ if is_conv:
+ result, mode_str = _svd_extract_conv(
+ weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
+ )
+ else:
+ result, mode_str = _svd_extract_linear(
+ weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile
+ )
+ 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"[LoRA Extract] Chunked: {key} ({num_chunks} chunks)")
+ lora_up, lora_down, rank = _extract_chunked_layer(
+ weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
+ )
+ if lora_up is not None:
+ output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
+ output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
+ output_sd[f"{lora_name}.alpha"] = torch.tensor([rank]).to(save_torch_dtype)
+ stats["chunked"] += 1
+ del weight_diff
+ pbar.update(1)
+ continue
+
+ print(f"[LoRA Extract] Failed: {key}: {e}")
+ stats["skipped"] += 1
+ del weight_diff
+ pbar.update(1)
+ continue
+
+ # Store result
+ if mode_str == "full":
+ output_sd[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
+ stats["full"] += 1
+ else:
+ lora_down, lora_up, _ = result
+ output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
+ output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
+ output_sd[f"{lora_name}.alpha"] = torch.tensor([lora_down.shape[0]]).to(save_torch_dtype)
+ stats["extracted"] += 1
+
+ del weight_diff
+ pbar.update(1)
+
+ print(f"[LoRA 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()
+
+ return output_sd
+
+
+def _save_lora(output_sd: dict, output_filename: str) -> str:
+ """Save LoRA to file."""
+ if not output_sd:
+ raise ValueError("No LoRA weights extracted")
+
+ 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 to {output_path}")
+
+ return output_path
+
+
+# =============================================================================
+# Node Definitions
+# =============================================================================
+
+def _get_model_inputs():
+ return [
+ io.Combo.Input("model_a", options=folder_paths.get_filename_list("diffusion_models"),
+ tooltip="Finetuned model (A - B = LoRA)"),
+ io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"),
+ tooltip="Base model (A - B = LoRA)"),
+ ]
+
+
+def _get_common_inputs():
+ return [
+ io.Boolean.Input("chunk_large_layers", default=False,
+ tooltip="Split large fused layers (QKV, MLP) into chunks"),
+ io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
+ tooltip="Clamp outlier singular values"),
+ io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
+ tooltip="Skip layers with max difference below this"),
+ io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
+ io.String.Input("output_filename", default="extracted_lora"),
+ io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
+ io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
+ io.String.Input("skip_patterns", default="", multiline=True,
+ tooltip="Regex patterns for layers to skip"),
+ ]
+
+
+class LoRAExtractFixed(io.ComfyNode):
+ """Extract LoRA with fixed rank."""
+
+ @classmethod
+ def define_schema(cls):
+ return io.Schema(
+ node_id="LoRAExtractFixed",
+ display_name="LoRA Extract (Fixed Rank)",
+ category="ModelUtils/LoRA",
+ description="Extract LoRA with specified fixed rank for each layer type.",
+ inputs=[
+ *_get_model_inputs(),
+ io.Int.Input("linear_dim", default=64, min=1, max=4096,
+ tooltip="Rank for linear/attention layers"),
+ io.Int.Input("conv_dim", default=32, min=1, max=4096,
+ tooltip="Rank for conv layers"),
+ *_get_common_inputs(),
+ ],
+ outputs=[io.String.Output(display_name="output_path")],
+ is_output_node=True,
+ )
+
+ @classmethod
+ def execute(cls, model_a, model_b, linear_dim, conv_dim, chunk_large_layers,
+ clamp_quantile, min_diff, mismatch_mode, output_filename,
+ save_dtype, device, skip_patterns) -> 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_sd = extract_lora_from_files(
+ model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
+ "lora_unet_", device, save_dtype, linear_dim, conv_dim,
+ clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
+ )
+
+ return io.NodeOutput(_save_lora(output_sd, output_filename))
+
+
+class LoRAExtractRatio(io.ComfyNode):
+ """Extract LoRA with singular value ratio threshold."""
+
+ @classmethod
+ def define_schema(cls):
+ return io.Schema(
+ node_id="LoRAExtractRatio",
+ display_name="LoRA Extract (Ratio)",
+ category="ModelUtils/LoRA",
+ description="Keep singular values > max(S) × ratio.",
+ inputs=[
+ *_get_model_inputs(),
+ io.Float.Input("linear_ratio", default=0.5, min=0.0, max=1.0, step=0.01,
+ tooltip="Ratio threshold for linear layers"),
+ io.Float.Input("conv_ratio", default=0.5, min=0.0, max=1.0, step=0.01,
+ tooltip="Ratio threshold for conv layers"),
+ *_get_common_inputs(),
+ ],
+ outputs=[io.String.Output(display_name="output_path")],
+ is_output_node=True,
+ )
+
+ @classmethod
+ def execute(cls, model_a, model_b, linear_ratio, conv_ratio, chunk_large_layers,
+ clamp_quantile, min_diff, mismatch_mode, output_filename,
+ save_dtype, device, skip_patterns) -> 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_sd = extract_lora_from_files(
+ model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
+ "lora_unet_", device, save_dtype, None, None,
+ clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
+ )
+
+ return io.NodeOutput(_save_lora(output_sd, output_filename))
+
+
+class LoRAExtractQuantile(io.ComfyNode):
+ """Extract LoRA with cumulative singular value quantile."""
+
+ @classmethod
+ def define_schema(cls):
+ return io.Schema(
+ node_id="LoRAExtractQuantile",
+ display_name="LoRA Extract (Quantile)",
+ category="ModelUtils/LoRA",
+ description="Keep enough singular values to reach target percentage.",
+ inputs=[
+ *_get_model_inputs(),
+ io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
+ tooltip="Target cumulative % for linear layers"),
+ io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
+ tooltip="Target cumulative % for conv layers"),
+ *_get_common_inputs(),
+ ],
+ outputs=[io.String.Output(display_name="output_path")],
+ is_output_node=True,
+ )
+
+ @classmethod
+ def execute(cls, model_a, model_b, linear_quantile, conv_quantile, chunk_large_layers,
+ clamp_quantile, min_diff, mismatch_mode, output_filename,
+ save_dtype, device, skip_patterns) -> 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_sd = extract_lora_from_files(
+ model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
+ "lora_unet_", device, save_dtype, None, None,
+ clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
+ )
+
+ return io.NodeOutput(_save_lora(output_sd, output_filename))
+
+
+class LoRAExtractKnee(io.ComfyNode):
+ """Extract LoRA with automatic knee detection."""
+
+ @classmethod
+ def define_schema(cls):
+ return io.Schema(
+ node_id="LoRAExtractKnee",
+ display_name="LoRA Extract (Knee Detection)",
+ category="ModelUtils/LoRA",
+ description="Automatically find optimal rank using knee detection on singular value curve.",
+ inputs=[
+ *_get_model_inputs(),
+ io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
+ default="sv_knee", tooltip="Knee detection method"),
+ *_get_common_inputs(),
+ ],
+ outputs=[io.String.Output(display_name="output_path")],
+ is_output_node=True,
+ )
+
+ @classmethod
+ def execute(cls, model_a, model_b, knee_method, chunk_large_layers,
+ clamp_quantile, min_diff, mismatch_mode, output_filename,
+ save_dtype, device, skip_patterns) -> 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_sd = extract_lora_from_files(
+ model_a_path, model_b_path, knee_method, 0, 0,
+ "lora_unet_", device, save_dtype, None, None,
+ clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
+ )
+
+ return io.NodeOutput(_save_lora(output_sd, output_filename))
+
+
+class LoRAExtractFrobenius(io.ComfyNode):
+ """Extract LoRA preserving target fraction of Frobenius norm."""
+
+ @classmethod
+ def define_schema(cls):
+ return io.Schema(
+ node_id="LoRAExtractFrobenius",
+ display_name="LoRA Extract (Frobenius)",
+ category="ModelUtils/LoRA",
+ description="Preserve target fraction of Frobenius norm.",
+ inputs=[
+ *_get_model_inputs(),
+ io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01,
+ tooltip="Target Frobenius norm fraction for linear"),
+ io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
+ tooltip="Target Frobenius norm fraction for conv"),
+ *_get_common_inputs(),
+ ],
+ outputs=[io.String.Output(display_name="output_path")],
+ is_output_node=True,
+ )
+
+ @classmethod
+ def execute(cls, model_a, model_b, linear_target, conv_target, chunk_large_layers,
+ clamp_quantile, min_diff, mismatch_mode, output_filename,
+ save_dtype, device, skip_patterns) -> 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_sd = extract_lora_from_files(
+ model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
+ "lora_unet_", device, save_dtype, None, None,
+ clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
+ )
+
+ return io.NodeOutput(_save_lora(output_sd, output_filename))
diff --git a/nodes/merger.py b/nodes/merger.py
index 3679565..da4d625 100644
--- a/nodes/merger.py
+++ b/nodes/merger.py
@@ -3,10 +3,14 @@ import re
import torch
import folder_paths
import comfy.utils
+from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
-from .merger_utils import MemoryEfficientSafeOpen
+from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
+from .device_utils import (
+ estimate_model_size, prepare_for_large_operation, cleanup_after_operation
+)
def load_documentation_from_file(filename):
@@ -22,13 +26,13 @@ def load_documentation_from_file(filename):
def _compile_patterns(pattern_string):
"""Compiles whitespace-separated regex patterns into a list of compiled patterns.
-
+
Returns empty list if pattern_string is empty or whitespace-only.
Raises ValueError if any pattern is invalid regex.
"""
if not pattern_string or not pattern_string.strip():
return []
-
+
patterns = []
for pattern in pattern_string.split():
try:
@@ -48,7 +52,7 @@ def _matches_any_pattern(key, patterns):
class MergerLogic:
"""Shared merge execution logic for all merger nodes."""
-
+
@staticmethod
def execute_merge(model_names, calc_mode, all_modes, recipe_params, model_type):
calc_mode_class = next((m for m in all_modes if m.name == calc_mode), None)
@@ -57,7 +61,22 @@ class MergerLogic:
primary_model_name = model_names.get('model_a')
if not primary_model_name or primary_model_name == "None":
raise ValueError("Model A is required to run the merge.")
-
+
+ # Prepare memory before heavy operation
+ total_size_gb = 0
+ model_paths = []
+ for name in model_names.values():
+ if name and name != "None":
+ path = folder_paths.get_full_path(model_type, name)
+ if path:
+ total_size_gb += estimate_model_size(path)
+ model_paths.append(path)
+
+ if total_size_gb > 0:
+ process_device = recipe_params.get('device', 'cpu')
+ print(f"[Merger] Preparing memory for {total_size_gb:.2f}GB merge operation...")
+ prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device))
+
handlers = {}
for name in model_names.values():
if name and name != "None":
@@ -65,20 +84,20 @@ class MergerLogic:
if not path:
raise FileNotFoundError(f"Model '{name}' not found.")
handlers[name] = MemoryEfficientSafeOpen(path)
-
+
primary_handler = handlers[primary_model_name]
all_keys = primary_handler.keys()
metadata = primary_handler.header.get("__metadata__", {})
-
+
# Convert mismatch_mode string to enum
mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip')
mismatch_mode = MissingTensorBehavior(mismatch_mode_str)
recipe_params['mismatch_mode'] = mismatch_mode
-
+
# Compile filter patterns
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''))
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''))
-
+
# Pre-compute key differences for logging
primary_keys = set(all_keys)
for name, handler in handlers.items():
@@ -90,7 +109,7 @@ class MergerLogic:
print(f"[Merger] {name} is missing {len(missing)} keys present in Model A")
if extra:
print(f"[Merger] {name} has {len(extra)} extra keys not in Model A (ignored)")
-
+
merged_state_dict = {}
pbar = comfy.utils.ProgressBar(len(all_keys))
save_dtype = recipe_params.pop('save_dtype')
@@ -98,34 +117,39 @@ class MergerLogic:
recipe_params.update({"handlers": handlers})
process_device = recipe_params.get('device', 'cpu')
process_dtype = recipe_params.get('dtype', torch.float32)
-
+
skipped_keys = 0
excluded_keys = 0
discarded_keys = 0
error_keys = []
-
- for key in all_keys:
+
+ for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
pbar.update(1)
continue
-
- # Pre-load Model A's tensor
- tensor_a = primary_handler.get_tensor(key).to(device=process_device, dtype=process_dtype)
-
+
+ # Pre-load Model A's tensor with pinned memory for CUDA
+ cpu_tensor = primary_handler.get_tensor(key)
+ if process_device == 'cuda':
+ tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
+ else:
+ tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype)
+ del cpu_tensor
+
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
merged_state_dict[key] = tensor_a.to(save_torch_dtype).cpu().clone()
excluded_keys += 1
pbar.update(1)
continue
-
+
# Pass tensor_a metadata to recipes for zeros mode and fallback
recipe_params['_tensor_a'] = tensor_a
recipe_params['_tensor_a_shape'] = tensor_a.shape
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
-
+
try:
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
result = recipe.merge()
@@ -134,22 +158,22 @@ class MergerLogic:
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
result = None
error_keys.append(key)
-
+
# Handle None result (mismatch occurred with skip mode)
if result is None:
result = tensor_a
skipped_keys += 1
-
+
if isinstance(result, dict):
for r_key, r_tensor in result.items():
merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone()
else:
merged_state_dict[key] = result.to(save_torch_dtype).cpu().clone()
-
+
# Clean up tensor_a reference to allow GC
del recipe_params['_tensor_a']
pbar.update(1)
-
+
# Log summary
if excluded_keys > 0:
print(f"[Merger] Excluded {excluded_keys} keys from merge (kept Model A only)")
@@ -159,15 +183,22 @@ class MergerLogic:
print(f"[Merger] Used Model A's values for {skipped_keys} keys due to mismatches")
if error_keys:
print(f"[Merger] Unexpected errors on {len(error_keys)} keys: {error_keys[:5]}{'...' if len(error_keys) > 5 else ''}")
-
+
for handler in handlers.values():
handler.__exit__(None, None, None)
+
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
- output_dir = folder_paths.get_folder_paths(output_folder)[0]
+ # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
+ output_dir = folder_paths.get_folder_paths(output_folder)[-1]
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
save_file(merged_state_dict, output_path, metadata=metadata)
+
+ # Cleanup after heavy operation
+ del merged_state_dict
+ cleanup_after_operation()
+
return f"{output_filename}.safetensors"
diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py
index be93a92..02b5f79 100644
--- a/nodes/merger_ops.py
+++ b/nodes/merger_ops.py
@@ -61,7 +61,7 @@ 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():
if self.on_missing == MissingTensorBehavior.ERROR:
@@ -77,7 +77,7 @@ 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)
class Multiply(Operation):
@@ -200,7 +200,7 @@ class InterpolateDifference(Operation):
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."""
@@ -219,13 +219,13 @@ 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
@@ -236,17 +236,17 @@ 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)
-
+
# 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)
@@ -261,13 +261,13 @@ 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)
-
+
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)
@@ -284,13 +284,13 @@ 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)
-
+
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)
@@ -304,17 +304,17 @@ 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)
-
+
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)
@@ -327,13 +327,13 @@ 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)
-
+
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)
@@ -349,13 +349,13 @@ 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)
-
+
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)
@@ -371,13 +371,13 @@ 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)
-
+
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)
@@ -392,13 +392,13 @@ 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)
-
+
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)
diff --git a/nodes/merger_utils.py b/nodes/merger_utils.py
index 59a0c4a..aadb5a4 100644
--- a/nodes/merger_utils.py
+++ b/nodes/merger_utils.py
@@ -1,36 +1,266 @@
+"""
+Memory-efficient safetensors utilities with parallel I/O and pinned memory support.
+
+Based on benchmarking, parallel reading with ThreadPoolExecutor provides 2-4x speedup
+for large models compared to sequential reads.
+
+Worker count is automatically optimized based on device capabilities when workers=None.
+Pinned memory enables faster CPU→GPU transfers (2-3x improvement).
+"""
+import os
+import mmap
import json
import struct
import torch
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Dict, List, Optional, Any, Tuple
+
+# Import ComfyUI's memory pinning functions if available
+try:
+ from comfy.model_management import pin_memory, unpin_memory, is_device_cuda
+ PINNING_AVAILABLE = True
+except ImportError:
+ PINNING_AVAILABLE = False
+ def pin_memory(tensor): return False
+ def unpin_memory(tensor): return False
+ def is_device_cuda(device): return str(device).startswith('cuda')
+
+
+class PinnedTensor:
+ """Context manager for automatic pinned memory management.
+
+ Usage:
+ with PinnedTensor(cpu_tensor) as pinned:
+ gpu_tensor = pinned.to('cuda')
+ # Memory automatically unpinned after transfer
+ """
+
+ def __init__(self, tensor: torch.Tensor, auto_pin: bool = True):
+ self.tensor = tensor
+ self.is_pinned = False
+ self.auto_pin = auto_pin and PINNING_AVAILABLE
+
+ def __enter__(self) -> torch.Tensor:
+ if self.auto_pin and self.tensor.device.type == 'cpu':
+ self.is_pinned = pin_memory(self.tensor)
+ return self.tensor
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ if self.is_pinned:
+ unpin_memory(self.tensor)
+ self.is_pinned = False
+
+
+def transfer_to_gpu_pinned(
+ tensor: torch.Tensor,
+ device: str = 'cuda',
+ dtype: Optional[torch.dtype] = None
+) -> torch.Tensor:
+ """Transfer tensor to GPU using pinned memory for faster transfer.
+
+ Args:
+ tensor: CPU tensor to transfer
+ device: Target GPU device
+ dtype: Optional dtype conversion
+
+ Returns:
+ Tensor on GPU device
+ """
+ if not PINNING_AVAILABLE or not is_device_cuda(torch.device(device)):
+ # Fall back to regular transfer
+ if dtype:
+ return tensor.to(device=device, dtype=dtype)
+ return tensor.to(device=device)
+
+ # Pin memory for faster transfer
+ was_pinned = pin_memory(tensor)
+ try:
+ if dtype:
+ result = tensor.to(device=device, dtype=dtype, non_blocking=True)
+ else:
+ result = tensor.to(device=device, non_blocking=True)
+ # Sync to ensure transfer is complete before unpinning
+ if was_pinned:
+ torch.cuda.current_stream().synchronize()
+ return result
+ finally:
+ if was_pinned:
+ unpin_memory(tensor)
+
class MemoryEfficientSafeOpen:
- def __init__(self, filename, device='cpu'):
+ """Memory-efficient safetensors file reader with mmap, parallel I/O, and pinned memory support.
+
+ Features:
+ - mmap mode: Zero-copy tensor access via memory-mapped file
+ - Parallel loading: Multi-threaded tensor reads for 2-4x speedup
+ - Sorted batch reads: Keys sorted by file offset for sequential I/O
+ - Auto-optimized workers: Adjusts parallelism based on device capabilities
+ - Pinned memory: Faster CPU→GPU transfers when transferring to CUDA
+
+ Args:
+ filename: Path to safetensors file
+ device: Target device (default 'cpu')
+ mmap_mode: Use memory-mapped file for zero-copy (default True)
+ """
+
+ def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True):
self.filename = filename
self.device = device
+ self.mmap_mode = mmap_mode
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
+ self.mmap_obj = None
+
+ if self.mmap_mode:
+ self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
+ if self.mmap_obj:
+ self.mmap_obj.close()
self.file.close()
- def keys(self):
+ def keys(self) -> List[str]:
+ """Return all tensor keys (excluding metadata)."""
return [k for k in self.header.keys() if k != "__metadata__"]
- def get_tensor(self, key):
+ def metadata(self) -> Dict[str, str]:
+ """Return file metadata."""
+ return self.header.get("__metadata__", {})
+
+ def keys_sorted_by_offset(self) -> List[str]:
+ """Return keys sorted by file offset for optimal sequential I/O."""
+ keys_with_offsets = []
+ for key in self.keys():
+ offset = self.header[key]["data_offsets"][0]
+ keys_with_offsets.append((key, offset))
+ keys_with_offsets.sort(key=lambda x: x[1])
+ return [k for k, _ in keys_with_offsets]
+
+ def get_tensor(self, key: str) -> torch.Tensor:
+ """Load a single tensor by key."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
- tensor_bytes = None
- if offset_start != offset_end:
+ if self.mmap_mode and self.mmap_obj:
+ if offset_start != offset_end:
+ start = self.header_size + 8 + offset_start
+ end = self.header_size + 8 + offset_end
+ tensor_bytes = memoryview(self.mmap_obj)[start:end]
+ else:
+ tensor_bytes = None
+ else:
+ 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 get_tensor_to_gpu(
+ self,
+ key: str,
+ device: str = 'cuda',
+ dtype: Optional[torch.dtype] = None
+ ) -> torch.Tensor:
+ """Load a tensor and transfer to GPU using pinned memory for faster transfer.
+
+ Args:
+ key: Tensor key to load
+ device: Target GPU device (default 'cuda')
+ dtype: Optional dtype conversion
+
+ Returns:
+ Tensor on GPU device
+ """
+ cpu_tensor = self.get_tensor(key)
+ return transfer_to_gpu_pinned(cpu_tensor, device, dtype)
+
+ def get_tensors_parallel(
+ self,
+ keys: Optional[List[str]] = None,
+ workers: Optional[int] = None,
+ sort_by_offset: bool = True
+ ) -> Dict[str, torch.Tensor]:
+ """Load multiple tensors in parallel using ThreadPoolExecutor.
+
+ Args:
+ keys: List of keys to load (default: all keys)
+ workers: Number of worker threads (default: auto-calculated based on device)
+ sort_by_offset: Sort keys by file offset before loading (default: True)
+
+ Returns:
+ Dict mapping key -> tensor
+ """
+ if keys is None:
+ keys = self.keys()
+
+ # Auto-calculate workers based on device capabilities
+ if workers is None:
+ try:
+ from .device_utils import get_optimal_workers, estimate_model_size
+ model_size = estimate_model_size(self.filename)
+ workers = get_optimal_workers(model_size)
+ except ImportError:
+ workers = min(4, os.cpu_count() or 4)
+
+ if sort_by_offset:
+ # Sort by offset for better I/O patterns
+ keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys]
+ keys_with_offsets.sort(key=lambda x: x[1])
+ keys = [k for k, _ in keys_with_offsets]
+
+ result = {}
+
+ with ThreadPoolExecutor(max_workers=workers) as executor:
+ futures = {
+ executor.submit(
+ _read_tensor_from_file,
+ self.filename,
+ key,
+ self.header[key],
+ self.header_size
+ ): key
+ for key in keys
+ }
+
+ for future in as_completed(futures):
+ key = futures[future]
+ try:
+ tensor = future.result()
+ if tensor is not None:
+ result[key] = tensor
+ except Exception as e:
+ print(f"[MemoryEfficientSafeOpen] Failed to load {key}: {e}")
+
+ return result
+
+ def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
+ """Load a tensor and decode as JSON dict (for config tensors)."""
+ if key not in self.header:
+ raise KeyError(f"Tensor '{key}' not found")
+
+ metadata = self.header[key]
+ offset_start, offset_end = metadata["data_offsets"]
+
+ if offset_start == offset_end:
+ return {}
+
+ if self.mmap_mode and self.mmap_obj:
+ start = self.header_size + 8 + offset_start
+ end = self.header_size + 8 + offset_end
+ tensor_bytes = bytes(self.mmap_obj[start:end])
+ else:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
- return self._deserialize_tensor(tensor_bytes, metadata)
+ return json.loads(tensor_bytes.decode("utf-8"))
def _read_header(self):
with open(self.filename, "rb") as f:
@@ -80,11 +310,52 @@ class MemoryEfficientSafeOpen:
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 _read_tensor_from_file(
+ filepath: str,
+ key: str,
+ metadata: Dict,
+ header_size: int
+) -> Optional[torch.Tensor]:
+ """Helper function to read a single tensor from file (for parallel execution).
+
+ Each worker opens its own file handle for thread safety.
"""
+ offset_start, offset_end = metadata["data_offsets"]
+ dtype_str = metadata["dtype"]
+ shape = metadata["shape"]
+
+ if offset_start == offset_end:
+ return None
+
+ dtype_map = {
+ "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
+ "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
+ "U8": torch.uint8, "BOOL": torch.bool,
+ }
+ if hasattr(torch, "float8_e5m2"):
+ dtype_map["F8_E5M2"] = torch.float8_e5m2
+ if hasattr(torch, "float8_e4m3fn"):
+ dtype_map["F8_E4M3"] = torch.float8_e4m3fn
+
+ dtype = dtype_map.get(dtype_str, torch.float32)
+
+ with open(filepath, "rb") as f:
+ f.seek(header_size + 8 + offset_start)
+ tensor_bytes = f.read(offset_end - offset_start)
+
+ byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
+
+ if dtype_str in ["F8_E5M2", "F8_E4M3"]:
+ if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
+ return byte_tensor.view(torch.float8_e5m2).reshape(shape)
+ elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
+ return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
+
+ return byte_tensor.view(dtype).reshape(shape)
+
+
+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,11 +363,8 @@ 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 []
diff --git a/nodes/metakeys.py b/nodes/metakeys.py
index 626a383..4a869f2 100644
--- a/nodes/metakeys.py
+++ b/nodes/metakeys.py
@@ -1,15 +1,17 @@
import os
+import json
+import struct
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 safetensors import safe_open
+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 +21,26 @@ 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)
+ # Read header only - no tensor loading needed
+ with open(model_path, "rb") as f:
+ header_size = struct.unpack(" str:
"""Shared logic for all PruneKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
@@ -21,31 +26,47 @@ 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)
+ # Prepare memory before operation
+ model_size_gb = estimate_model_size(model_path_to_load)
+ prepare_for_large_operation(model_size_gb * 1.2)
+
+ # Get metadata from safe_open (header only, no tensor loading)
+ with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
+ metadata = f.metadata() or {}
patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
if not patterns:
raise ValueError("No keys/patterns provided to prune.")
+ # Stream tensors and filter on the fly
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
+ with MemoryEfficientSafeOpen(model_path_to_load) as handler:
+ all_keys = handler.keys()
+ pbar = comfy.utils.ProgressBar(len(all_keys))
+ for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
+ is_match = False
+ if use_regex:
+ if any(re.search(pattern, key) for pattern in patterns):
+ is_match = True
+ else:
+ if any(pattern in key for pattern in patterns):
+ is_match = True
- if not is_match:
- pruned_tensors[key] = tensor
+ if not is_match:
+ pruned_tensors[key] = handler.get_tensor(key)
+ pbar.update(1)
- model_dir = folder_paths.get_folder_paths(model_type)[0]
+ # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
+ model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(pruned_tensors, output_path, metadata)
+ # Cleanup after operation
+ del pruned_tensors
+ cleanup_after_operation()
+
return output_path
@@ -72,9 +93,9 @@ class ModelPruneKeys(io.ComfyNode):
)
@classmethod
- def execute(cls, diffusionmodel_name: str, keys_to_prune: str,
+ 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",
+ path = _prune_keys(diffusionmodel_name, "diffusion_models",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
@@ -102,9 +123,9 @@ class TextEncoderPruneKeys(io.ComfyNode):
)
@classmethod
- def execute(cls, textencoder_name: str, keys_to_prune: str,
+ 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",
+ path = _prune_keys(textencoder_name, "text_encoders",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
@@ -132,9 +153,9 @@ class LoRAPruneKeys(io.ComfyNode):
)
@classmethod
- def execute(cls, lora_name: str, keys_to_prune: str,
+ def execute(cls, lora_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
- path = _prune_keys(lora_name, "loras",
+ path = _prune_keys(lora_name, "loras",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
@@ -162,9 +183,9 @@ class CheckpointPruneKeys(io.ComfyNode):
)
@classmethod
- def execute(cls, ckpt_name: str, keys_to_prune: str,
+ def execute(cls, ckpt_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
- path = _prune_keys(ckpt_name, "checkpoints",
+ path = _prune_keys(ckpt_name, "checkpoints",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
@@ -192,8 +213,8 @@ class EmbeddingPruneKeys(io.ComfyNode):
)
@classmethod
- def execute(cls, embedding: str, keys_to_prune: str,
+ def execute(cls, embedding: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
- path = _prune_keys(embedding, "embeddings",
+ path = _prune_keys(embedding, "embeddings",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
\ No newline at end of file
diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py
index d856a70..625f753 100644
--- a/nodes/renamekeys.py
+++ b/nodes/renamekeys.py
@@ -1,11 +1,16 @@
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 safetensors import safe_open
+from .utils import convert_pt_to_safetensors
+from .merger_utils import MemoryEfficientSafeOpen
+from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
-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 +25,13 @@ 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)
+
+ # Get metadata from safe_open (header only, no tensor loading)
+ with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
+ metadata = f.metadata() or {}
old_keys = [key.strip() for key in old_keys_str.strip().split('\n') if key.strip()]
new_keys = [key.strip() for key in new_keys_str.strip().split('\n') if key.strip()]
@@ -30,17 +41,26 @@ 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())
+ # Stream tensors and rename on the fly
renamed_tensors = {}
- for key in original_keys:
- new_key = key_map.get(key, key)
- renamed_tensors[new_key] = model_weights[key]
+ with MemoryEfficientSafeOpen(model_path_to_load) as handler:
+ original_keys = handler.keys()
+ pbar = comfy.utils.ProgressBar(len(original_keys))
+ for key in tqdm(original_keys, desc="Renaming keys", unit="keys"):
+ new_key = key_map.get(key, key)
+ renamed_tensors[new_key] = handler.get_tensor(key)
+ pbar.update(1)
- model_dir = folder_paths.get_folder_paths(model_type)[0]
+ # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
+ model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(renamed_tensors, output_path, metadata)
+ # Cleanup after operation
+ del renamed_tensors
+ cleanup_after_operation()
+
return output_path
@@ -67,9 +87,9 @@ class ModelRenameKeys(io.ComfyNode):
)
@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)
@@ -97,9 +117,9 @@ class TextEncoderRenameKeys(io.ComfyNode):
)
@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)
@@ -127,9 +147,9 @@ class LoRARenameKeys(io.ComfyNode):
)
@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)
@@ -157,9 +177,9 @@ class CheckpointRenameKeys(io.ComfyNode):
)
@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)
@@ -187,8 +207,8 @@ class EmbeddingRenameKeys(io.ComfyNode):
)
@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)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 4db0a3f..4d3baad 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
-version = "0.1.7"
+version = "0.2.0"
license = { file = "LICENSE" }
[project.urls]
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..d2ce485
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+# No external dependencies required