Compare commits

...
21 changed files with 3781 additions and 469 deletions
+54 -1
View File
@@ -1,2 +1,55 @@
# ComfyUI-ModelUtils
[WIP]Custom nodes for handling, inspecting, modifying and creating various model files.
A collection of ComfyUI custom nodes for inspecting, modifying, merging, and creating model files. Supports Models, TextEncoders, LoRAs, Checkpoints, and Embeddings.
## Features
- **MetaKeys** – Inspect and display metadata and tensor keys from model files
- **RenameKeys** – Batch rename tensor keys using pattern matching
- **PruneKeys** – Remove unwanted layers/keys from models
- **Mergers** – Combine 2 or 3 models with configurable blend modes and ratios
- **LoRA Extraction** – Extract LoRA adapters from model pairs using various SVD rank selection methods (Fixed, Ratio, Quantile, Knee-detection, Frobenius-norm)
## Example Workflows
<p align="center">
<img src="assets/GetMetaAndkeys.png" width="400" alt="Get Meta and Keys">
<br>
<a href="example_workflows/GetMetaAndkeys.json">📥 GetMetaAndkeys.json</a>
</p>
---
<p align="center">
<img src="assets/LoRA_Extract_nodes.png" width="400" alt="LoRA Extraction Nodes">
<br>
<a href="example_workflows/LoRA_Extract_nodes.json">📥 LoRA_Extract_nodes.json</a>
</p>
---
<p align="center">
<img src="assets/Merging_Examples.png" width="400" alt="Merging Examples">
<br>
<a href="example_workflows/Merging_Examples.json">📥 Merging_Examples.json</a>
</p>
---
<p align="center">
<img src="assets/RenameKeysInModel.png" width="400" alt="Rename Keys in Model">
<br>
<a href="example_workflows/RenameKeysInModel.json">📥 RenameKeysInModel.json</a>
</p>
## Acknowledgements
The LoRA extraction functionality was inspired by and references the excellent work from:
- [kohya-ss/sd-scripts](https://github.com/kohya-ss/sd-scripts) – Training scripts for Stable Diffusion
- [KohakuBlueleaf/LyCORIS](https://github.com/KohakuBlueleaf/LyCORIS) – Advanced LoRA techniques
- [bmaltais/kohya_ss](https://github.com/bmaltais/kohya_ss) – Windows-friendly GUI for sd-scripts
## License
See [LICENSE](LICENSE) for details.
+7 -4
View File
@@ -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()
return ModelUtilsExtension()
Binary file not shown.

After

Width:  |  Height:  |  Size: 352 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 600 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 592 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 427 KiB

+600
View File
@@ -0,0 +1,600 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 31,
"last_link_id": 19,
"nodes": [
{
"id": 16,
"type": "PreviewAny",
"pos": [
-3608,
200
],
"size": [
832,
232
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 18
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 17,
"type": "ModelMetaKeys",
"pos": [
-2800,
504
],
"size": [
346.253515625,
78
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
13
]
},
{
"name": "keys",
"type": "STRING",
"links": [
17
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "ModelMetaKeys"
},
"widgets_values": [
"flux2-dev.safetensors"
]
},
{
"id": 18,
"type": "TextEncoderMetaKeys",
"pos": [
-3176,
680
],
"size": [
329.4234375,
78
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
15
]
},
{
"name": "keys",
"type": "STRING",
"links": [
16
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "TextEncoderMetaKeys"
},
"widgets_values": [
"clip_l.safetensors"
]
},
{
"id": 19,
"type": "CheckpointMetaKeys",
"pos": [
-3168,
512
],
"size": [
317.6724609375,
78
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
12
]
},
{
"name": "keys",
"type": "STRING",
"links": [
18
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "CheckpointMetaKeys"
},
"widgets_values": [
"unsorted\\v1-5-pruned-emaonly.safetensors"
]
},
{
"id": 21,
"type": "LoRAMetaKeys",
"pos": [
-2800,
680
],
"size": [
277.0943359375,
78
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": [
14
]
},
{
"name": "keys",
"type": "STRING",
"links": [
19
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAMetaKeys"
},
"widgets_values": [
"Chroma\\FlatAnimeStyle_Chroma.safetensors"
]
},
{
"id": 24,
"type": "PreviewAny",
"pos": [
-2728,
216
],
"size": [
512,
200
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 17
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 25,
"type": "PreviewAny",
"pos": [
-2816,
832
],
"size": [
736,
248
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 19
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 27,
"type": "PreviewAny",
"pos": [
-3496,
848
],
"size": [
632,
224
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 16
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 28,
"type": "PreviewAny",
"pos": [
-3600,
664
],
"size": [
352,
136
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 15
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 29,
"type": "PreviewAny",
"pos": [
-2408,
648
],
"size": [
360,
120
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 14
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 30,
"type": "PreviewAny",
"pos": [
-3576,
480
],
"size": [
336,
128
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 12
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 31,
"type": "PreviewAny",
"pos": [
-2368,
472
],
"size": [
328,
120
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 13
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
}
],
"links": [
[
12,
19,
0,
30,
0,
"STRING"
],
[
13,
17,
0,
31,
0,
"STRING"
],
[
14,
21,
0,
29,
0,
"STRING"
],
[
15,
18,
0,
28,
0,
"STRING"
],
[
16,
18,
1,
27,
0,
"STRING"
],
[
17,
17,
1,
24,
0,
"STRING"
],
[
18,
19,
1,
16,
0,
"STRING"
],
[
19,
21,
1,
25,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.9070294784580497,
"offset": [
3855.323424397858,
-56.106573378634465
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+522
View File
@@ -0,0 +1,522 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 10,
"last_link_id": 5,
"nodes": [
{
"id": 1,
"type": "LoRAExtractFixed",
"pos": [
-1848,
16
],
"size": [
296,
352
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
3
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractFixed"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
64,
32,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 2,
"type": "LoRAExtractFrobenius",
"pos": [
-2432,
24
],
"size": [
296,
352
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
1
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractFrobenius"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
0.9,
0.9,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 3,
"type": "LoRAExtractKnee",
"pos": [
-2416,
440
],
"size": [
280,
328
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
2
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractKnee"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
"sv_knee",
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 4,
"type": "LoRAExtractQuantile",
"pos": [
-1248,
128
],
"size": [
288,
352
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
5
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractQuantile"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
0.9,
0.9,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 5,
"type": "LoRAExtractRatio",
"pos": [
-1824,
424
],
"size": [
296,
352
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
4
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAExtractRatio"
},
"widgets_values": [
"13499-bf16.safetensors",
"13499-bf16.safetensors",
0.5,
0.5,
false,
0.99,
0,
"skip",
"extracted_lora",
"fp16",
"cuda",
""
]
},
{
"id": 6,
"type": "PreviewAny",
"pos": [
-2096,
136
],
"size": [
210,
168
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 1
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 7,
"type": "PreviewAny",
"pos": [
-2096,
448
],
"size": [
210,
168
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 2
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 8,
"type": "PreviewAny",
"pos": [
-1504,
96
],
"size": [
210,
168
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 3
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 9,
"type": "PreviewAny",
"pos": [
-1472,
432
],
"size": [
210,
168
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 4
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 10,
"type": "PreviewAny",
"pos": [
-904,
184
],
"size": [
210,
168
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 5
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
}
],
"links": [
[
1,
2,
0,
6,
0,
"STRING"
],
[
2,
3,
0,
7,
0,
"STRING"
],
[
3,
1,
0,
8,
0,
"STRING"
],
[
4,
5,
0,
9,
0,
"STRING"
],
[
5,
4,
0,
10,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.05,
"offset": [
2532,
84
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+460
View File
@@ -0,0 +1,460 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 16,
"last_link_id": 11,
"nodes": [
{
"id": 6,
"type": "PreviewAny",
"pos": [
-3552,
136
],
"size": [
336,
160
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 6
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 7,
"type": "PreviewAny",
"pos": [
-2792,
144
],
"size": [
456,
464
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 9
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
false
]
},
{
"id": 11,
"type": "ModelThreeMerger",
"pos": [
-3120,
112
],
"size": [
296,
522
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": null
},
{
"name": "documentation",
"type": "STRING",
"links": [
9
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "ModelThreeMerger"
},
"widgets_values": [
"DOCUMENTATION ONLY",
"Chroma1-HD.safetensors",
"Chroma-DC-2K.safetensors",
"None",
"Add-Difference",
"skip",
0.5,
0.5,
0.5,
0.5,
996811747788113,
"randomize",
"merged_3_model",
"fp32",
"cuda",
"",
""
]
},
{
"id": 12,
"type": "ModelTwoMerger",
"pos": [
-3120,
704
],
"size": [
304,
474
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": [
11
]
},
{
"name": "documentation",
"type": "STRING",
"links": null
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "ModelTwoMerger"
},
"widgets_values": [
"MERGE",
"None",
"None",
"Weight-Sum",
"skip",
0.5,
0.5,
0.99,
388546580517523,
"randomize",
"merged_2_model",
"fp32",
"cuda",
"",
""
]
},
{
"id": 13,
"type": "LoRAThreeMerger",
"pos": [
-3872,
128
],
"size": [
296,
522
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": [
6
]
},
{
"name": "documentation",
"type": "STRING",
"links": []
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAThreeMerger"
},
"widgets_values": [
"MERGE",
"None",
"None",
"None",
"Add-Difference",
"skip",
0.5,
0.5,
0.5,
0.5,
1120102695669269,
"randomize",
"merged_3_lora",
"fp32",
"cuda",
"",
""
]
},
{
"id": 14,
"type": "LoRATwoMerger",
"pos": [
-3872,
720
],
"size": [
296,
474
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "output_filename",
"type": "STRING",
"links": []
},
{
"name": "documentation",
"type": "STRING",
"links": [
10
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRATwoMerger"
},
"widgets_values": [
"DOCUMENTATION ONLY",
"Hyper-FLUX.1-dev-16steps-lora-rescale_z1.safetensors",
"Hyper-FLUX.1-dev-16steps-lora-rescale_z2.safetensors",
"Weight-Sum",
"skip",
0.5,
0.5,
0.99,
1020143314396572,
"randomize",
"merged_2_lora",
"fp32",
"cuda",
"",
""
]
},
{
"id": 15,
"type": "PreviewAny",
"pos": [
-2768,
720
],
"size": [
432,
160
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 11
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 16,
"type": "PreviewAny",
"pos": [
-3544,
712
],
"size": [
400,
448
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 10
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
false
]
}
],
"links": [
[
6,
13,
0,
6,
0,
"STRING"
],
[
9,
11,
1,
7,
0,
"STRING"
],
[
10,
14,
1,
16,
0,
"STRING"
],
[
11,
12,
0,
15,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.9070294784580494,
"offset": [
3972,
-12
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+583
View File
@@ -0,0 +1,583 @@
{
"id": "e07ecc9b-30b9-45e4-9ed6-5fbb7e510257",
"revision": 0,
"last_node_id": 46,
"last_link_id": 31,
"nodes": [
{
"id": 35,
"type": "LoRAMetaKeys",
"pos": [
-3136,
512
],
"size": [
376,
78
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "metadata",
"type": "STRING",
"links": null
},
{
"name": "keys",
"type": "STRING",
"links": [
21,
22
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRAMetaKeys"
},
"widgets_values": [
"qwen\\gmqwen_k8_piid_4step.safetensors"
]
},
{
"id": 36,
"type": "LoRARenameKeys",
"pos": [
-2856,
864
],
"size": [
320,
190
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "old_keys",
"type": "STRING",
"widget": {
"name": "old_keys"
},
"link": 27
},
{
"name": "new_keys",
"type": "STRING",
"widget": {
"name": "new_keys"
},
"link": 28
}
],
"outputs": [
{
"name": "output_path",
"type": "STRING",
"links": [
31
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfyui_modelutils",
"ver": "eef737706127fe4a2c28172ab706e92318e0dc60",
"Node name for S&R": "LoRARenameKeys"
},
"widgets_values": [
"qwen\\gmqwen_k8_piid_4step.safetensors",
"",
"",
"gmqwen_k8_piid_4step"
]
},
{
"id": 37,
"type": "PreviewAny",
"pos": [
-3792,
920
],
"size": [
592,
384
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 21
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 40,
"type": "RegexReplace",
"pos": [
-3400,
568
],
"size": [
210,
292
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "string",
"type": "STRING",
"widget": {
"name": "string"
},
"link": 22
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": [
23,
25
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "RegexReplace"
},
"widgets_values": [
"",
",\\s\\[.+$",
"",
true,
true,
false,
0
]
},
{
"id": 41,
"type": "PreviewAny",
"pos": [
-3784,
528
],
"size": [
360,
328
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 23
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 42,
"type": "RegexReplace",
"pos": [
-3176,
1016
],
"size": [
288,
292
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "string",
"type": "STRING",
"widget": {
"name": "string"
},
"link": 26
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": [
28,
30
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "RegexReplace"
},
"widgets_values": [
"",
"^(\\w+)",
"diffusion_model.\\1",
true,
true,
false,
0
]
},
{
"id": 43,
"type": "RegexReplace",
"pos": [
-3152,
648
],
"size": [
288,
292
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "string",
"type": "STRING",
"widget": {
"name": "string"
},
"link": 25
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": [
26,
27,
29
]
}
],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "RegexReplace"
},
"widgets_values": [
"",
"^(proj_out.+|norm_out.+)$\\n",
"",
true,
true,
false,
0
]
},
{
"id": 44,
"type": "PreviewAny",
"pos": [
-2832,
632
],
"size": [
640,
192
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 29
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 45,
"type": "PreviewAny",
"pos": [
-2848,
1096
],
"size": [
640,
192
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 30
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
},
{
"id": 46,
"type": "PreviewAny",
"pos": [
-2520,
872
],
"size": [
336,
168
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 31
}
],
"outputs": [],
"properties": {
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65,
"cnr_id": "comfy-core",
"ver": "0.7.0",
"Node name for S&R": "PreviewAny"
},
"widgets_values": [
null,
null,
true
]
}
],
"links": [
[
21,
35,
1,
37,
0,
"STRING"
],
[
22,
35,
1,
40,
0,
"STRING"
],
[
23,
40,
0,
41,
0,
"STRING"
],
[
25,
40,
0,
43,
0,
"STRING"
],
[
26,
43,
0,
42,
0,
"STRING"
],
[
27,
43,
0,
36,
0,
"STRING"
],
[
28,
42,
0,
36,
1,
"STRING"
],
[
29,
43,
0,
44,
0,
"STRING"
],
[
30,
42,
0,
45,
0,
"STRING"
],
[
31,
36,
0,
46,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1,
"offset": [
4191.258750150775,
-356.1917593715161
]
},
"frontendVersion": "1.37.2",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
+268
View File
@@ -0,0 +1,268 @@
"""
Device capabilities and memory management utilities.
Integrates with ComfyUI's model_management for safe memory operations.
"""
import os
import gc
import psutil
import torch
from typing import Optional, Tuple
from dataclasses import dataclass
# Import ComfyUI's model management functions
try:
from comfy.model_management import (
get_torch_device,
get_free_memory,
get_total_memory,
free_memory,
soft_empty_cache,
unload_all_models,
is_device_cuda,
)
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
@dataclass
class DeviceCapabilities:
"""Device capabilities for I/O and compute optimization."""
device_type: str # 'cuda', 'cpu', 'mps', 'xpu'
total_vram_gb: float
free_vram_gb: float
total_ram_gb: float
free_ram_gb: float
cpu_cores: int
recommended_workers: int
max_tensor_size_gb: float # Largest tensor we can safely load
supports_pinned_memory: bool
supports_async_streams: bool
def get_device_capabilities(device: Optional[torch.device] = None) -> DeviceCapabilities:
"""
Analyze device capabilities for optimal I/O configuration.
Args:
device: Target device (default: ComfyUI's torch device)
Returns:
DeviceCapabilities dataclass with recommended settings
"""
if device is None:
if COMFY_AVAILABLE:
device = get_torch_device()
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device_type = device.type if hasattr(device, 'type') else 'cpu'
# Get RAM info
ram_info = psutil.virtual_memory()
total_ram_gb = ram_info.total / (1024**3)
free_ram_gb = ram_info.available / (1024**3)
# Get VRAM info
if device_type == 'cuda' and torch.cuda.is_available():
if COMFY_AVAILABLE:
total_vram = get_total_memory(device)
free_vram = get_free_memory(device)
else:
free_vram, total_vram = torch.cuda.mem_get_info(device)
total_vram_gb = total_vram / (1024**3)
free_vram_gb = free_vram / (1024**3)
supports_pinned = True
supports_async = True
elif device_type == 'mps':
total_vram_gb = total_ram_gb * 0.5 # MPS shares RAM
free_vram_gb = free_ram_gb * 0.5
supports_pinned = False
supports_async = False
else:
total_vram_gb = 0
free_vram_gb = 0
supports_pinned = False
supports_async = False
# CPU cores
cpu_cores = os.cpu_count() or 4
# Recommended workers based on hardware
if device_type == 'cuda':
# GPU: fewer workers since I/O is not the bottleneck
recommended_workers = min(4, cpu_cores // 2)
else:
# CPU: more workers for parallel I/O
recommended_workers = min(8, cpu_cores)
# Max tensor size: leave headroom for operations
if device_type == 'cuda':
max_tensor_gb = min(free_vram_gb * 0.8, free_ram_gb * 0.5)
else:
max_tensor_gb = free_ram_gb * 0.6
return DeviceCapabilities(
device_type=device_type,
total_vram_gb=total_vram_gb,
free_vram_gb=free_vram_gb,
total_ram_gb=total_ram_gb,
free_ram_gb=free_ram_gb,
cpu_cores=cpu_cores,
recommended_workers=recommended_workers,
max_tensor_size_gb=max_tensor_gb,
supports_pinned_memory=supports_pinned,
supports_async_streams=supports_async,
)
def estimate_model_size(file_path: str) -> float:
"""Estimate model size in GB from file size."""
try:
return os.path.getsize(file_path) / (1024**3)
except OSError:
return 0.0
def can_fit_in_memory(
model_paths: list[str],
device: Optional[torch.device] = None,
safety_factor: float = 0.8
) -> Tuple[bool, str]:
"""
Check if models can fit in available memory.
Args:
model_paths: List of model file paths
device: Target device
safety_factor: Fraction of free memory to use (default 0.8)
Returns:
(can_fit, reason_message)
"""
caps = get_device_capabilities(device)
total_size = sum(estimate_model_size(p) for p in model_paths)
if caps.device_type == 'cuda':
available = caps.free_vram_gb * safety_factor
mem_type = "VRAM"
else:
available = caps.free_ram_gb * safety_factor
mem_type = "RAM"
if total_size <= available:
return True, f"OK: {total_size:.2f}GB needed, {available:.2f}GB {mem_type} available"
else:
return False, f"Insufficient memory: {total_size:.2f}GB needed but only {available:.2f}GB {mem_type} available"
def prepare_for_large_operation(
estimated_memory_gb: float,
device: Optional[torch.device] = None
) -> bool:
"""
Prepare device for a large memory operation by freeing resources.
Args:
estimated_memory_gb: Estimated memory needed in GB
device: Target device
Returns:
True if preparation succeeded
"""
if not COMFY_AVAILABLE:
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return True
if device is None:
device = get_torch_device()
memory_bytes = int(estimated_memory_gb * 1024**3)
# Use ComfyUI's free_memory which unloads models intelligently
try:
free_memory(memory_bytes, device)
return True
except Exception as e:
print(f"[DeviceUtils] Warning: Could not free memory: {e}")
return False
def cleanup_after_operation():
"""Clean up memory after a large operation."""
gc.collect()
if COMFY_AVAILABLE:
soft_empty_cache()
elif torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def get_optimal_workers(
model_size_gb: float,
caps: Optional[DeviceCapabilities] = None
) -> int:
"""
Calculate optimal number of I/O workers based on model size and device.
Args:
model_size_gb: Model size in GB
caps: Pre-calculated device capabilities (optional)
Returns:
Recommended number of workers
"""
if caps is None:
caps = get_device_capabilities()
# For very large models, fewer workers to avoid memory pressure
if model_size_gb > caps.free_ram_gb * 0.5:
return max(1, caps.recommended_workers // 2)
# For small models, use more workers
if model_size_gb < 2.0:
return min(8, caps.cpu_cores)
return caps.recommended_workers
def get_processing_device_string(prefer_cuda: bool = True) -> str:
"""
Get the best available processing device as a string.
Args:
prefer_cuda: Prefer CUDA over CPU if available
Returns:
Device string ('cuda', 'mps', 'cpu')
"""
if prefer_cuda:
if torch.cuda.is_available():
return "cuda"
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
return "mps"
return "cpu"
def print_device_info():
"""Print device capabilities for debugging."""
caps = get_device_capabilities()
print(f"\n{'='*50}")
print(" Device Capabilities")
print(f"{'='*50}")
print(f" Device Type: {caps.device_type}")
print(f" Total VRAM: {caps.total_vram_gb:.2f} GB")
print(f" Free VRAM: {caps.free_vram_gb:.2f} GB")
print(f" Total RAM: {caps.total_ram_gb:.2f} GB")
print(f" Free RAM: {caps.free_ram_gb:.2f} GB")
print(f" CPU Cores: {caps.cpu_cores}")
print(f" Recommended Workers: {caps.recommended_workers}")
print(f" Max Tensor Size: {caps.max_tensor_size_gb:.2f} GB")
print(f" Pinned Memory: {caps.supports_pinned_memory}")
print(f" Async Streams: {caps.supports_async_streams}")
print(f"{'='*50}\n")
-350
View File
@@ -1,350 +0,0 @@
"""
Memory-efficient LoRA extraction node.
This node loads two models from files and extracts a LoRA by computing the difference,
without requiring the full models to be loaded into GPU memory via ComfyUI's model management.
"""
import os
import torch
import folder_paths
import comfy.utils
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_utils import MemoryEfficientSafeOpen
CLAMP_QUANTILE = 0.99
def extract_lora(diff, rank):
"""Extract LoRA weights from a weight difference tensor using SVD."""
conv2d = (len(diff.shape) == 4)
kernel_size = None if not conv2d else diff.size()[2:4]
conv2d_3x3 = conv2d and kernel_size != (1, 1)
out_dim, in_dim = diff.size()[0:2]
rank = min(rank, in_dim, out_dim)
if conv2d:
if conv2d_3x3:
diff = diff.flatten(start_dim=1)
else:
diff = diff.squeeze()
U, S, Vh = torch.linalg.svd(diff.float())
U = U[:, :rank]
S = S[:rank]
U = U @ torch.diag(S)
Vh = Vh[:rank, :]
dist = torch.cat([U.flatten(), Vh.flatten()])
hi_val = torch.quantile(dist, CLAMP_QUANTILE)
low_val = -hi_val
U = U.clamp(low_val, hi_val)
Vh = Vh.clamp(low_val, hi_val)
if conv2d:
U = U.reshape(out_dim, rank, 1, 1)
Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
return (U, Vh)
def extract_lora_from_files(
model_a_path: str,
model_b_path: str,
rank: int,
lora_type: str,
bias_diff: bool,
prefix_lora: str,
process_device: str,
save_dtype: str,
) -> dict[str, torch.Tensor]:
"""
Memory-efficient LoRA extraction by streaming tensors from safetensors files.
Args:
model_a_path: Path to the finetuned/modified model
model_b_path: Path to the base/original model
rank: LoRA rank for SVD decomposition
lora_type: "standard" for LoRA, "full_diff" for full difference
bias_diff: Whether to include bias differences
prefix_lora: Prefix for LoRA keys (e.g., "diffusion_model.")
process_device: Device for processing ("cuda" or "cpu")
save_dtype: Output dtype ("fp16", "bf16", "fp32")
Returns:
Dictionary of LoRA tensors
"""
output_sd = {}
# Get dtype for saving
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16
}.get(save_dtype, torch.float16)
# Open both files with memory-efficient handler
handler_a = MemoryEfficientSafeOpen(model_a_path)
handler_b = MemoryEfficientSafeOpen(model_b_path)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
# Only process keys that exist in both models
common_keys = keys_a & keys_b
weight_keys = [k for k in common_keys if k.endswith(".weight")]
bias_keys = [k for k in common_keys if k.endswith(".bias")] if bias_diff else []
# Log key differences
missing_in_b = keys_a - keys_b
extra_in_b = keys_b - keys_a
if missing_in_b:
print(f"[LoRA Extract] {len(missing_in_b)} keys in model A not in model B (skipped)")
if extra_in_b:
print(f"[LoRA Extract] {len(extra_in_b)} keys in model B not in model A (ignored)")
pbar = comfy.utils.ProgressBar(len(weight_keys) + len(bias_keys))
for key in weight_keys:
# Load tensors one at a time
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
# Compute difference: finetuned - base
weight_diff = tensor_a - tensor_b
# Clean up immediately
del tensor_a, tensor_b
# Extract LoRA key name
lora_key = key[:-7] # Remove ".weight"
if lora_type == "standard":
if weight_diff.ndim < 2:
# 1D weights (e.g., LayerNorm) - store as diff if bias_diff enabled
if bias_diff:
output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu()
else:
try:
lora_up, lora_down = extract_lora(weight_diff, rank)
output_sd[f"{prefix_lora}{lora_key}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu()
output_sd[f"{prefix_lora}{lora_key}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu()
except Exception as e:
print(f"[LoRA Extract] Could not extract LoRA for {key}: {e}")
elif lora_type == "full_diff":
output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu()
del weight_diff
pbar.update(1)
# Process bias keys
for key in bias_keys:
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
bias_diff_tensor = tensor_a - tensor_b
del tensor_a, tensor_b
lora_key = key[:-5] # Remove ".bias"
output_sd[f"{prefix_lora}{lora_key}.diff_b"] = bias_diff_tensor.to(save_torch_dtype).cpu()
del bias_diff_tensor
pbar.update(1)
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
return output_sd
class LoRASaveFromFile(io.ComfyNode):
"""
Memory-efficient LoRA extraction from safetensors files.
Unlike the standard LoraSave node which requires pre-loaded and subtracted models,
this node loads models directly from files one tensor at a time, making it much
more memory efficient for extracting LoRAs from large models.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRASaveFromFile",
display_name="Extract LoRA from Files (Memory Efficient)",
category="ModelUtils/LoRA",
description="Extracts a LoRA by computing the difference between two model files. "
"Much more memory efficient than using ModelSubtract + LoraSave.",
inputs=[
io.Combo.Input(
"model_a",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="The finetuned/modified model (A - B = LoRA)"
),
io.Combo.Input(
"model_b",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="The base/original model (A - B = LoRA)"
),
io.Int.Input("rank", default=64, min=1, max=4096, step=1,
tooltip="LoRA rank - higher preserves more detail but increases size"),
io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard",
tooltip="standard: SVD decomposition to lora_up/lora_down, full_diff: raw weight differences"),
io.Boolean.Input("bias_diff", default=True,
tooltip="Include bias differences in the LoRA"),
io.String.Input("output_filename", default="extracted_lora",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_a: str,
model_b: str,
rank: int,
lora_type: str,
bias_diff: bool,
output_filename: str,
save_dtype: str,
process_device: str,
) -> io.NodeOutput:
# Get full paths
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
# Extract LoRA
output_sd = extract_lora_from_files(
model_a_path=model_a_path,
model_b_path=model_b_path,
rank=rank,
lora_type=lora_type,
bias_diff=bias_diff,
prefix_lora="diffusion_model.",
process_device=process_device,
save_dtype=save_dtype,
)
if not output_sd:
raise ValueError("No LoRA weights extracted - models may be identical or incompatible")
# Save to loras folder
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path)
print(f"[LoRA Extract] Saved {len(output_sd)} tensors to {output_path}")
return io.NodeOutput(output_path)
class LoRACheckpointSaveFromFile(io.ComfyNode):
"""
Memory-efficient LoRA extraction from checkpoint files (which may contain both model and text encoder).
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRACheckpointSaveFromFile",
display_name="Extract LoRA from Checkpoints (Memory Efficient)",
category="ModelUtils/LoRA",
description="Extracts a LoRA by computing the difference between two checkpoint files. "
"Handles both diffusion model and text encoder weights.",
inputs=[
io.Combo.Input(
"checkpoint_a",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="The finetuned/modified checkpoint (A - B = LoRA)"
),
io.Combo.Input(
"checkpoint_b",
options=folder_paths.get_filename_list("checkpoints"),
tooltip="The base/original checkpoint (A - B = LoRA)"
),
io.Int.Input("rank", default=64, min=1, max=4096, step=1,
tooltip="LoRA rank - higher preserves more detail but increases size"),
io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard",
tooltip="standard: SVD decomposition, full_diff: raw weight differences"),
io.Boolean.Input("bias_diff", default=True,
tooltip="Include bias differences in the LoRA"),
io.String.Input("output_filename", default="extracted_lora",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
checkpoint_a: str,
checkpoint_b: str,
rank: int,
lora_type: str,
bias_diff: bool,
output_filename: str,
save_dtype: str,
process_device: str,
) -> io.NodeOutput:
# Get full paths
checkpoint_a_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_a)
checkpoint_b_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_b)
# Open handlers to detect prefixes
handler_a = MemoryEfficientSafeOpen(checkpoint_a_path)
keys = handler_a.keys()
handler_a.__exit__(None, None, None)
# Auto-detect key prefixes
has_model = any(k.startswith("model.") for k in keys)
has_diffusion = any(k.startswith("diffusion_model.") for k in keys)
has_text_encoders = any(k.startswith("text_encoders.") or k.startswith("cond_stage_model.") for k in keys)
output_sd = {}
# Extract diffusion model LoRA
if has_model or has_diffusion:
prefix = "model.diffusion_model." if has_model else "diffusion_model."
model_lora = extract_lora_from_files(
model_a_path=checkpoint_a_path,
model_b_path=checkpoint_b_path,
rank=rank,
lora_type=lora_type,
bias_diff=bias_diff,
prefix_lora="diffusion_model.",
process_device=process_device,
save_dtype=save_dtype,
)
output_sd.update(model_lora)
# Note: Text encoder extraction would need additional prefix handling
# For full checkpoint support, this would need to be expanded
if not output_sd:
raise ValueError("No LoRA weights extracted - checkpoints may be identical or incompatible")
# Save to loras folder
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path)
print(f"[LoRA Extract] Saved {len(output_sd)} tensors to {output_path}")
return io.NodeOutput(output_path)
+817
View File
@@ -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))
+56 -25
View File
@@ -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"
+23 -23
View File
@@ -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)
+283 -15
View File
@@ -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 []
+23 -8
View File
@@ -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("<Q", f.read(8))[0]
header_json = f.read(header_size).decode("utf-8")
header = json.loads(header_json)
# Extract metadata and layer shapes from header
metadata = header.pop("__metadata__", {})
metadata_str = str(metadata)
for layer_name, tensor in model_weights.items():
shape = str(tensor.shape)
shape = shape.replace("torch.Size(", "").replace(")", "")
layer_shapes = ""
keys = [k for k in header.keys()]
pbar = comfy.utils.ProgressBar(len(keys))
for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
info = header[layer_name]
shape = str(info.get("shape", []))
name_shape = f"{layer_name}, {shape}\n"
layer_shapes += name_shape
pbar.update(1)
return metadata_str, layer_shapes
+45 -24
View File
@@ -1,12 +1,17 @@
import os
import re
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
from 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 _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
use_regex: bool, output_filename: str) -> str:
"""Shared logic for all PruneKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
@@ -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)
+38 -18
View File
@@ -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)
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.1.7"
version = "0.2.0"
license = { file = "LICENSE" }
[project.urls]
+1
View File
@@ -0,0 +1 @@
# No external dependencies required