Compare commits

...
Author SHA1 Message Date
silveroxides 61b2f506d0 feat: add svd_niter parameter for SVD accuracy control 2026-01-03 20:27:56 +01:00
silveroxides 75d1f0f5c9 perf: use svd_lowrank for fixed rank resize (10x faster) 2026-01-03 20:14:25 +01:00
silveroxides 3a71ad9952 feat: add max_rank inputs to dynamic extraction nodes 2026-01-03 20:12:58 +01:00
silveroxides 37090dd7ba feat: add LoRA resize nodes and fix sv_ratio convention 2026-01-03 20:04:56 +01:00
silveroxides f58289d042 document development phases 2026-01-03 19:37:39 +01:00
Silver 1796f1bea9 Merge pull request #3 from silveroxides/feature/lora-extract
Feature/lora extract and version bumped to 0.2.0
2026-01-02 17:01:41 +01:00
Silver 3df073f810 Bump version from 0.1.7 to 0.2.0 2026-01-02 17:01:02 +01:00
silveroxides 0e73c5d889 Update README.md, add images and workflows 2026-01-02 16:59:25 +01:00
silveroxides eef7377061 Added .contiguous() to saves 2026-01-01 21:18:08 +01:00
silveroxides 65c1ae37b7 add svd_lowrank for fixed rank extraction and put experimental chunk large layers default to false 2026-01-01 21:08:15 +01:00
silveroxides b7707862b2 Add pinned memory for faster merging and extraction 2026-01-01 20:49:49 +01:00
silveroxides 6d7b2d1403 refactor nodes and implement memory management 2026-01-01 20:29:08 +01:00
silveroxides bfbbcd9515 fix exception handler 2026-01-01 19:04:25 +01:00
silveroxides 18679cca2c Split lycoris extraction modes into different nodes 2026-01-01 07:21:37 +01:00
silveroxides da13f9f51d add lycoris diff extraction 2026-01-01 07:06:19 +01:00
silveroxides 92d3c5582c add lora extraction, progress bars and more 2026-01-01 06:36:59 +01:00
Silver bb923e8474 Merge pull request #2 from silveroxides/feature/v3-migration
migrate to v3 schema and add LoRA extraction node
2026-01-01 06:35:56 +01:00
23 changed files with 4695 additions and 469 deletions
+54 -1
View File
@@ -1,2 +1,55 @@
# ComfyUI-ModelUtils # 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.
+13 -3
View File
@@ -19,8 +19,13 @@ from .nodes.merger import (
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger, ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger CheckpointThreeMerger, EmbeddingThreeMerger
) )
from .nodes.lora_extract import ( from .nodes.lora_extract_svd import (
LoRASaveFromFile, LoRACheckpointSaveFromFile LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius
)
from .nodes.lora_resize import (
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative
) )
@@ -44,9 +49,14 @@ class ModelUtilsExtension(ComfyExtension):
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger, ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger, CheckpointThreeMerger, EmbeddingThreeMerger,
# LoRA Extraction # LoRA Extraction
LoRASaveFromFile, LoRACheckpointSaveFromFile, LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius,
# LoRA Resize
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
] ]
async def comfy_entrypoint() -> ModelUtilsExtension: 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

+296
View File
@@ -0,0 +1,296 @@
# ComfyUI-ModelUtils - Complete Project History
## Version Timeline
| Version | Commits | Key Features |
|---------|---------|--------------|
| 0.1.1 | `70d6c13`-`d67cbca` | Initial nodes (MetaKeys) |
| 0.1.2 | `6b662a4` | Icon assets |
| 0.1.4 | `cbe3502` | Key renaming/pruning nodes |
| 0.1.5 | `78b3198` | Model merging nodes |
| 0.1.6 | `31989be` | Embedding nodes, namespace fixes |
| 0.1.7 | `b6fae8d` | Mismatch handling, layer filtering, V3 migration |
| **0.2.0** | `3df073f` | LoRA extraction, memory optimization, pinned memory |
---
## Phase 1: Initial Release (v0.1.1)
**Commit:** `70d6c13` - "Add first nodes, toml and publish"
### Files Created
- `nodes/metakeys.py` - Get model metadata and layer keys
- `nodes/utils.py` - Shared utilities
- `__init__.py` - Node registration
- `pyproject.toml` - Package configuration
### Nodes
- `ModelMetaKeys` - Get diffusion model metadata
- `CheckpointMetaKeys` - Get checkpoint metadata
- `TextEncoderMetaKeys` - Get text encoder metadata
- `LoRAMetaKeys` - Get LoRA metadata
- `EmbeddingMetaKeys` - Get embedding metadata
---
## Phase 2: Key Manipulation (v0.1.4)
**Commit:** `4f08fe6` - "add key renaming and key pruning nodes"
### Files Created
- `nodes/renamekeys.py` - Rename tensor keys
- `nodes/prunekeys.py` - Remove tensor keys by pattern
### Nodes (5 variants each for model types)
- `ModelRenameKeys`, `CheckpointRenameKeys`, `TextEncoderRenameKeys`, `LoRARenameKeys`, `EmbeddingRenameKeys`
- `ModelPruneKeys`, `CheckpointPruneKeys`, `TextEncoderPruneKeys`, `LoRAPruneKeys`, `EmbeddingPruneKeys`
---
## Phase 3: Model Merging (v0.1.5)
**Commit:** `7959370` - "Add model merging nodes"
### Files Created
- `nodes/merger.py` - Main merging logic
- `nodes/merger_ops.py` - Merge operation modes
- `nodes/merger_utils.py` - `MemoryEfficientSafeOpen` class
- `docs/merger_2_model_modes.md` - Documentation
- `docs/merger_3_model_modes.md` - Documentation
### Nodes
- `CheckpointTwoMerger`, `DiffusionModelTwoMerger`, `TextEncoderTwoMerger`, `LoRATwoMerger`
- `CheckpointThreeMerger`, `DiffusionModelThreeMerger`, `TextEncoderThreeMerger`, `LoRAThreeMerger`
### Merge Modes (TWO_MODEL_MODES)
- Weighted Sum, Add Difference, Smooth Add Difference
- Train Difference, Sum Twice, Triple Sum
- Euclidean Add Difference, Distribution Crossover, Ties
- Slerp, Lerp, SVD Supermerger, SVD LoRA Extraction
### Merge Modes (THREE_MODEL_MODES)
- Add Difference (A + (B - C) × α)
### MemoryEfficientSafeOpen (initial)
```python
class MemoryEfficientSafeOpen:
def __init__(self, filename, device='cpu', mmap_mode=True)
def keys() -> List[str]
def get_tensor(key) -> Tensor
# mmap for zero-copy reads
```
---
## Phase 4: Embedding Nodes (v0.1.6)
**Commit:** `0d035ef` - "Add embedding nodes"
### Additions
- Added embedding variants to all node types
- Namespace conflict fixes (`68f22fa`, `5ea63ab`)
- Refactored names for universality (`232dc76`)
---
## Phase 5: Mismatch Handling (v0.1.7)
**Commits:** `19bc047` → `de27186`
### Features Added
- `mismatch_mode` parameter: skip/zeros/error
- Layer filtering with regex patterns
- `exclude_patterns` - Keep model A only (no merge)
- `discard_patterns` - Remove from output entirely
### Files Modified
- `merger.py` - Added mismatch handling logic
- `merger_ops.py` - Added `MissingTensorBehavior` enum, `MissingTensorError`
---
## Phase 6: V3 API Migration (v0.1.7)
**Commit:** `a3116cb` - "migrate to v3 schema and add LoRA extraction node"
### Changes
- Migrated all nodes from legacy API to `comfy_api.latest.io`
- `io.Schema`, `io.ComfyNode`, `io.NodeOutput`
- Added initial LoRA extraction node (diff-based)
---
## Phase 7: LyCORIS Extraction (pre-0.2.0)
**Commits:** `92d3c55` → `18679cc`
### Files Created
- `nodes/lycoris_extract.py` - LyCORIS-based extraction
- `nodes/lora_extract.py` - Simple diff extraction
### Initial Implementation
- Used `lycoris-lora` external dependency
- `lycoris.utils.extract_linear`, `extract_conv`
- Added `requirements.txt` with `lycoris-lora`
### Nodes Split (`18679cc`)
- `LyCORISExtractFixed`, `LyCORISExtractThreshold`
- `LyCORISExtractRatio`, `LyCORISExtractQuantile`
---
## Phase 8: Native SVD Port (0.2.0)
**Commits:** `6d7b2d1` → `eef7377`
### Removed
- `lycoris-lora` dependency
- `lycoris_extract.py`
- `lora_extract.py`
### NEW: [lora_extract_svd.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/lora_extract_svd.py)
**8 Rank Selection Functions:**
```python
_index_sv_fixed(S, dim) # Fixed rank
_index_sv_threshold(S, threshold) # Absolute threshold
_index_sv_ratio(S, ratio) # Relative to max
_index_sv_cumulative(S, target) # Cumulative sum %
_index_sv_fro(S, target) # Frobenius norm
_index_sv_knee(S) # Knee detection
_index_sv_cumulative_knee(S) # Cumsum knee
_index_sv_rel_decrease(S, tau) # Ratio drop
```
**SVD Functions:**
```python
_svd_extract_linear_lowrank(weight, rank, device, clamp_quantile)
_svd_extract_linear(weight, mode, mode_param, device, max_rank, clamp_quantile)
_svd_extract_conv(weight, mode, mode_param, device, max_rank, clamp_quantile)
_detect_fused_layer(key, shape) # QKV/MLP detection
_extract_chunked_layer(...) # Chunked extraction
```
**Nodes:**
- `LoRAExtractFixed` - Uses `svd_lowrank` (10x faster)
- `LoRAExtractRatio` - Relative threshold
- `LoRAExtractQuantile` - Cumulative target
- `LoRAExtractKnee` - Auto knee detection
- `LoRAExtractFrobenius` - Norm preservation
---
## Phase 9: Memory Management (0.2.0)
**Commit:** `6d7b2d1` - "refactor nodes and implement memory management"
### NEW: [device_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/device_utils.py)
Integrates with `comfy.model_management`:
```python
@dataclass
class DeviceCapabilities:
device_type, total_vram_gb, free_vram_gb, total_ram_gb, free_ram_gb
cpu_cores, recommended_workers, max_tensor_size_gb
supports_pinned_memory, supports_async_streams
get_device_capabilities(device) -> DeviceCapabilities
estimate_model_size(path) -> float # GB
can_fit_in_memory(paths, device, safety_factor) -> (bool, str)
prepare_for_large_operation(gb, device) -> bool # Calls free_memory()
cleanup_after_operation() # gc.collect() + soft_empty_cache()
get_optimal_workers(size_gb) -> int
```
### Updates to Existing Nodes
- `metakeys.py` - Header-only reading (no tensor loading)
- `renamekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
- `prunekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
- `merger.py` - Memory prep/cleanup
---
## Phase 10: Pinned Memory (0.2.0)
**Commit:** `b770786` - "Add pinned memory for faster merging and extraction"
### [merger_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/merger_utils.py) Enhancements
```python
# Context manager
class PinnedTensor:
def __enter__(self) -> Tensor # Pins memory
def __exit__(...) # Unpins
# Transfer helper (2-3x faster)
def transfer_to_gpu_pinned(tensor, device='cuda', dtype=None) -> Tensor
# MemoryEfficientSafeOpen additions
def get_tensor_to_gpu(key, device, dtype) # Pinned transfer
def keys_sorted_by_offset() -> List[str] # Sequential I/O
def get_tensors_parallel(keys, workers, sort_by_offset) -> Dict # Parallel
def get_tensor_as_dict(key) -> Dict # JSON config parsing
```
---
## Phase 11: SVD Optimization (0.2.0)
**Commit:** `65c1ae3` - "add svd_lowrank for fixed rank extraction"
### Changes
- Added `_svd_extract_linear_lowrank()` using `torch.svd_lowrank`
- Routes fixed mode to lowrank for **10x speedup**
- Changed `chunk_large_layers` default to `False`
---
## Phase 12: Bug Fixes (0.2.0)
**Commit:** `eef7377` - "Added .contiguous() to saves"
### Fix
Non-contiguous tensor error from SVD slices:
```python
tensor.to(dtype).cpu().contiguous() # Added .contiguous()
```
---
## Current File Structure
```
ComfyUI-ModelUtils/
├── __init__.py
├── pyproject.toml
├── requirements.txt (empty - no external deps)
├── README.md
├── docs/
│ ├── merger_2_model_modes.md
│ └── merger_3_model_modes.md
├── nodes/
│ ├── utils.py
│ ├── metakeys.py
│ ├── renamekeys.py
│ ├── prunekeys.py
│ ├── merger.py
│ ├── merger_ops.py
│ ├── merger_utils.py
│ ├── device_utils.py
│ └── lora_extract_svd.py
└── assets/
└── icon.png
```
---
## Performance Summary
| Feature | Improvement |
|---------|-------------|
| Fixed-rank SVD | **10x** (svd_lowrank) |
| CPU→GPU transfer | **2-3x** (pinned memory) |
| Metadata reading | **100x+** (header only) |
| Parallel I/O | **2-4x** (ThreadPoolExecutor) |
| Memory management | Auto-unload models before heavy ops |
+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)
+840
View File
@@ -0,0 +1,840 @@
"""
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 (kohya convention)."""
if ratio <= 0:
return len(S) - 1
min_sv = S[0] / ratio
rank = int(torch.sum(S > min_sv).item())
return max(1, min(rank, len(S) - 1))
def _index_sv_cumulative(S: torch.Tensor, target: float) -> 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,
niter: int = 2,
) -> tuple:
"""
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms.
Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# Clamp rank to valid range
max_possible_rank = min(out_dim, in_dim)
rank = max(1, min(rank, max_possible_rank - 1))
# Check if decomposition is worthwhile
if rank >= max_possible_rank // 2:
return weight, "full"
# Use svd_lowrank for faster low-rank approximation
try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
# U: [out_dim, rank], S: [rank], V: [in_dim, rank]
Vh = V.T # [rank, in_dim]
except Exception as e:
raise RuntimeError(f"svd_lowrank failed: {e}")
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_linear(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
SVD decomposition for linear layers.
For fixed rank mode, uses svd_lowrank which is much faster.
For adaptive modes, uses full SVD to analyze all singular values.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# For fixed rank, use the faster svd_lowrank
if mode == "fixed" and mode_param > 0:
target_rank = int(mode_param)
if max_rank is not None:
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
# Compute rank based on mode
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim) // 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,
svd_niter: int = 2,
) -> 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, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[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"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> 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, svd_niter
)
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=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> 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, linear_max_rank, conv_max_rank,
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"),
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> 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, linear_max_rank, conv_max_rank,
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"),
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> 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, linear_max_rank, conv_max_rank,
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"),
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> 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, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
+588
View File
@@ -0,0 +1,588 @@
"""
LoRA Resize - Resize existing LoRAs to different ranks.
Based on kohya_ss resize_lora.py. Merges LoRA weights and re-extracts via SVD.
Supports fixed rank and dynamic methods (sv_ratio, sv_fro, sv_cumulative).
"""
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 typing import Optional, Dict, Tuple, List
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
# Reuse rank functions from extraction module
from .lora_extract_svd import (
_index_sv_ratio,
_index_sv_cumulative,
_index_sv_fro,
)
MIN_SV = 1e-6
# =============================================================================
# LoRA Format Detection - Support various LoRA sources
# =============================================================================
def detect_lora_format(keys: List[str]) -> Dict:
"""
Detect LoRA format and key patterns from various sources.
Supports:
- Kohya/A1111: lora_unet_*, lora_te_*
- Diffusers: *.lora_A.weight, *.lora_B.weight
- PEFT: base_model.model.*.lora_A/B
- ComfyUI native: *.lora_down.weight, *.lora_up.weight
- Flux/SDXL: Various transformer key patterns
Returns:
Dict with format info and key patterns
"""
format_info = {
"format": "unknown",
"down_suffix": ".lora_down.weight",
"up_suffix": ".lora_up.weight",
"alpha_suffix": ".alpha",
"has_alpha": False,
"key_count": 0,
}
# Check for alpha keys
alpha_keys = [k for k in keys if '.alpha' in k or k.endswith('alpha')]
format_info["has_alpha"] = len(alpha_keys) > 0
# Detect format by key patterns
sample_keys = keys[:20] # Check first 20 keys
# Kohya/A1111 format: lora_unet_down_blocks_0_*.lora_down.weight
if any('lora_unet_' in k or 'lora_te' in k for k in sample_keys):
format_info["format"] = "kohya"
format_info["down_suffix"] = ".lora_down.weight"
format_info["up_suffix"] = ".lora_up.weight"
# Diffusers format: transformer.*.lora_A.weight
elif any('.lora_A.weight' in k for k in sample_keys):
format_info["format"] = "diffusers"
format_info["down_suffix"] = ".lora_A.weight"
format_info["up_suffix"] = ".lora_B.weight"
format_info["alpha_suffix"] = ".alpha"
# PEFT format: base_model.model.*.lora_A.default.weight
elif any('base_model.model.' in k and 'lora_' in k for k in sample_keys):
format_info["format"] = "peft"
# Handle nested structure
format_info["down_suffix"] = ".lora_A.default.weight"
format_info["up_suffix"] = ".lora_B.default.weight"
# ComfyUI native / standard
elif any('.lora_down.weight' in k for k in sample_keys):
format_info["format"] = "comfy"
format_info["down_suffix"] = ".lora_down.weight"
format_info["up_suffix"] = ".lora_up.weight"
# Count LoRA layers
down_keys = [k for k in keys if format_info["down_suffix"] in k]
format_info["key_count"] = len(down_keys)
return format_info
def extract_lora_pairs(keys: List[str], format_info: Dict) -> Dict[str, Dict[str, str]]:
"""
Group LoRA keys into down/up/alpha pairs.
Returns:
Dict[block_name, {"down": key, "up": key, "alpha": key}]
"""
down_suffix = format_info["down_suffix"]
up_suffix = format_info["up_suffix"]
alpha_suffix = format_info["alpha_suffix"]
pairs = {}
for key in keys:
if down_suffix in key:
block_name = key.replace(down_suffix, "")
if block_name not in pairs:
pairs[block_name] = {}
pairs[block_name]["down"] = key
elif up_suffix in key:
block_name = key.replace(up_suffix, "")
if block_name not in pairs:
pairs[block_name] = {}
pairs[block_name]["up"] = key
elif alpha_suffix in key or key.endswith('.alpha'):
# Handle various alpha key formats
block_name = key.replace(alpha_suffix, "").replace(".alpha", "")
if block_name not in pairs:
pairs[block_name] = {}
pairs[block_name]["alpha"] = key
return pairs
def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int, float]:
"""
Detect the rank and alpha of an existing LoRA.
Returns:
(network_dim, network_alpha)
"""
network_dim = None
network_alpha = None
for block_name, block_keys in pairs.items():
if "down" not in block_keys:
continue
# Get dim from down weight shape
if network_dim is None:
down_key = block_keys["down"]
shape = handler.header[down_key]["shape"]
# Linear: [rank, in_features] or Conv: [rank, in_ch, k, k]
network_dim = shape[0]
# Get alpha if present
if network_alpha is None and "alpha" in block_keys:
alpha_tensor = handler.get_tensor(block_keys["alpha"])
network_alpha = float(alpha_tensor.item())
if network_dim is not None and network_alpha is not None:
break
# Default alpha to dim if not found
if network_alpha is None:
network_alpha = float(network_dim) if network_dim else 1.0
if network_dim is None:
network_dim = 1
return network_dim, network_alpha
# =============================================================================
# Merge Functions
# =============================================================================
def _merge_linear(lora_down: torch.Tensor, lora_up: torch.Tensor, device: str) -> torch.Tensor:
"""Merge linear LoRA weights: lora_up @ lora_down."""
lora_down = lora_down.to(device=device, dtype=torch.float32)
lora_up = lora_up.to(device=device, dtype=torch.float32)
return lora_up @ lora_down
def _merge_conv(lora_down: torch.Tensor, lora_up: torch.Tensor, device: str) -> torch.Tensor:
"""Merge conv LoRA weights."""
in_rank, in_size, kernel_size, k_ = lora_down.shape
out_size, out_rank, _, _ = lora_up.shape
assert in_rank == out_rank, f"rank mismatch: {in_rank} vs {out_rank}"
lora_down = lora_down.to(device=device, dtype=torch.float32)
lora_up = lora_up.to(device=device, dtype=torch.float32)
merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
return merged.reshape(out_size, in_size, kernel_size, kernel_size)
# =============================================================================
# Extract Functions (after merge)
# =============================================================================
def _extract_linear(
weight: torch.Tensor,
max_rank: int,
dynamic_method: Optional[str],
dynamic_param: Optional[float],
scale: float,
niter: int = 2,
) -> Dict:
"""Extract LoRA from merged linear weight."""
out_size, in_size = weight.shape
if dynamic_method is None:
# Fixed rank - use svd_lowrank for 10x speedup
rank = min(max_rank, min(out_size, in_size) - 1)
U, S, Vh = torch.svd_lowrank(weight, q=rank, niter=niter)
Vh = Vh.T # svd_lowrank returns V, not Vh
new_rank = rank
new_alpha = float(scale * new_rank)
stats = {"sum_retained": 1.0, "fro_retained": 1.0} # Not computed for lowrank
else:
# Dynamic methods need full SVD to compute rank from all singular values
U, S, Vh = linalg.svd(weight, full_matrices=False)
new_rank, new_alpha, stats = _compute_resize(S, max_rank, dynamic_method, dynamic_param, scale)
U = U[:, :new_rank]
S = S[:new_rank]
Vh = Vh[:new_rank, :]
lora_up = U @ torch.diag(S)
lora_down = Vh
return {
"lora_down": lora_down.cpu().contiguous(),
"lora_up": lora_up.cpu().contiguous(),
"new_rank": new_rank,
"new_alpha": new_alpha,
**stats
}
def _extract_conv(
weight: torch.Tensor,
max_rank: int,
dynamic_method: Optional[str],
dynamic_param: Optional[float],
scale: float,
niter: int = 2,
) -> Dict:
"""Extract LoRA from merged conv weight."""
out_ch, in_ch, kh, kw = weight.shape
mat = weight.reshape(out_ch, -1)
if dynamic_method is None:
# Fixed rank - use svd_lowrank for 10x speedup
rank = min(max_rank, min(mat.shape) - 1)
U, S, Vh = torch.svd_lowrank(mat, q=rank, niter=niter)
Vh = Vh.T # svd_lowrank returns V, not Vh
new_rank = rank
new_alpha = float(scale * new_rank)
stats = {"sum_retained": 1.0, "fro_retained": 1.0} # Not computed for lowrank
else:
# Dynamic methods need full SVD to compute rank from all singular values
U, S, Vh = linalg.svd(mat, full_matrices=False)
new_rank, new_alpha, stats = _compute_resize(S, max_rank, dynamic_method, dynamic_param, scale)
U = U[:, :new_rank]
S = S[:new_rank]
Vh = Vh[:new_rank, :]
lora_up = (U @ torch.diag(S)).reshape(out_ch, new_rank, 1, 1)
lora_down = Vh.reshape(new_rank, in_ch, kh, kw)
return {
"lora_down": lora_down.cpu().contiguous(),
"lora_up": lora_up.cpu().contiguous(),
"new_rank": new_rank,
"new_alpha": new_alpha,
**stats
}
def _compute_resize(
S: torch.Tensor,
max_rank: int,
dynamic_method: Optional[str],
dynamic_param: Optional[float],
scale: float
) -> Tuple[int, float, Dict]:
"""Compute new rank and alpha based on resize method."""
if dynamic_method == "sv_ratio" and dynamic_param is not None:
# Note: _index_sv_ratio uses S[0]*ratio, kohya uses S[0]/ratio
# Use kohya convention here for consistency with their tool
min_sv = S[0] / dynamic_param
new_rank = max(1, int(torch.sum(S > min_sv).item()))
elif dynamic_method == "sv_cumulative" and dynamic_param is not None:
new_rank = _index_sv_cumulative(S, dynamic_param)
elif dynamic_method == "sv_fro" and dynamic_param is not None:
new_rank = _index_sv_fro(S, dynamic_param)
else:
new_rank = max_rank
# Clamp rank
if S[0] < MIN_SV:
new_rank = 1
else:
new_rank = max(1, min(new_rank, max_rank, len(S) - 1))
new_alpha = float(scale * new_rank)
# Compute retention stats
s_sum = float(torch.sum(torch.abs(S)))
s_rank = float(torch.sum(torch.abs(S[:new_rank]))) if new_rank <= len(S) else s_sum
S_sq = S.pow(2)
s_fro = float(torch.sqrt(torch.sum(S_sq)))
s_red_fro = float(torch.sqrt(torch.sum(S_sq[:new_rank]))) if new_rank <= len(S) else s_fro
stats = {
"sum_retained": s_rank / s_sum if s_sum > 0 else 1.0,
"fro_retained": s_red_fro / s_fro if s_fro > 0 else 1.0,
}
return new_rank, new_alpha, stats
# =============================================================================
# Main Resize Function
# =============================================================================
def resize_lora_file(
lora_path: str,
new_rank: int,
dynamic_method: Optional[str],
dynamic_param: Optional[float],
device: str,
save_dtype: torch.dtype,
output_filename: str,
verbose: bool = True,
svd_niter: int = 2,
) -> str:
"""
Resize a LoRA file to a new rank.
Args:
lora_path: Path to input LoRA
new_rank: Target rank (max rank for dynamic methods)
dynamic_method: None, "sv_ratio", "sv_fro", "sv_cumulative"
dynamic_param: Parameter for dynamic method
device: Processing device
save_dtype: Output dtype
output_filename: Output filename (without extension)
verbose: Print progress info
svd_niter: Power iterations for SVD accuracy
Returns:
Path to saved resized LoRA
"""
# Prepare memory
lora_size_gb = estimate_model_size(lora_path)
prepare_for_large_operation(lora_size_gb * 2, torch.device(device))
handler = MemoryEfficientSafeOpen(lora_path)
try:
metadata = handler.metadata().copy()
all_keys = handler.keys()
# Detect format and extract pairs
format_info = detect_lora_format(all_keys)
pairs = extract_lora_pairs(all_keys, format_info)
network_dim, network_alpha = detect_lora_rank(handler, pairs)
scale = network_alpha / network_dim if network_dim > 0 else 1.0
if verbose:
method_str = f"{dynamic_method}: {dynamic_param}" if dynamic_method else "fixed"
print(f"[LoRA Resize] Format: {format_info['format']}, layers: {format_info['key_count']}")
print(f"[LoRA Resize] Original dim={network_dim}, alpha={network_alpha:.1f}, scale={scale:.3f}")
print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}")
output_sd = {}
fro_list = []
pbar = comfy.utils.ProgressBar(len(pairs))
with torch.no_grad():
for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"):
if "down" not in block_keys or "up" not in block_keys:
pbar.update(1)
continue
lora_down = handler.get_tensor(block_keys["down"])
lora_up = handler.get_tensor(block_keys["up"])
is_conv = len(lora_down.shape) == 4
# Transfer to GPU with pinned memory if available
if device == 'cuda':
lora_down = transfer_to_gpu_pinned(lora_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(lora_up, device, torch.float32)
# Merge and re-extract
if is_conv:
weight = _merge_conv(lora_down, lora_up, device)
result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
else:
weight = _merge_linear(lora_down, lora_up, device)
result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
del weight, lora_down, lora_up
fro_list.append(result['fro_retained'])
# Store using same format suffixes as input
down_suffix = format_info["down_suffix"]
up_suffix = format_info["up_suffix"]
alpha_suffix = format_info["alpha_suffix"]
output_sd[f"{block_name}{down_suffix}"] = result["lora_down"].to(save_dtype)
output_sd[f"{block_name}{up_suffix}"] = result["lora_up"].to(save_dtype)
output_sd[f"{block_name}{alpha_suffix}"] = torch.tensor([result["new_alpha"]]).to(save_dtype)
pbar.update(1)
if verbose and fro_list:
import numpy as np
avg_fro = np.mean(fro_list)
std_fro = np.std(fro_list)
print(f"[LoRA Resize] Average Frobenius retention: {avg_fro:.1%} ± {std_fro:.3f}")
# Update metadata
if dynamic_method:
metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}"
metadata["ss_network_dim"] = "Dynamic"
metadata["ss_network_alpha"] = "Dynamic"
else:
metadata["ss_training_comment"] = f"Resized from dim {network_dim} to {new_rank}"
metadata["ss_network_dim"] = str(new_rank)
metadata["ss_network_alpha"] = str(scale * new_rank)
# Save
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, metadata)
print(f"[LoRA Resize] Saved to {output_path}")
return output_path
finally:
handler.__exit__(None, None, None)
cleanup_after_operation()
# =============================================================================
# Node Definitions
# =============================================================================
class LoRAResizeFixed(io.ComfyNode):
"""Resize LoRA to a fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAResizeFixed",
display_name="LoRA Resize (Fixed Rank)",
category="ModelUtils/LoRA/Resize",
description="Resize existing LoRA to a specific rank by merging and re-extracting via SVD.",
inputs=[
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
tooltip="LoRA to resize"),
io.Int.Input("new_rank", default=64, min=1, max=1024,
tooltip="Target rank"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
io.String.Input("output_filename", default="resized_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name, new_rank, svd_niter, output_filename, save_dtype, device) -> io.NodeOutput:
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(lora_path, new_rank, None, None, device, dtype, output_filename, svd_niter=svd_niter)
return io.NodeOutput(path)
class LoRAResizeRatio(io.ComfyNode):
"""Resize LoRA keeping SVs above ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAResizeRatio",
display_name="LoRA Resize (SV Ratio)",
category="ModelUtils/LoRA/Resize",
description="Dynamically resize LoRA, keeping singular values where S[i] > S[0]/ratio.",
inputs=[
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
tooltip="LoRA to resize"),
io.Int.Input("max_rank", default=128, min=1, max=1024,
tooltip="Maximum allowed rank"),
io.Float.Input("ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Keep SVs where S[i] > S[0]/ratio"),
io.String.Input("output_filename", default="resized_lora_ratio"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name, max_rank, ratio, output_filename, save_dtype, device) -> io.NodeOutput:
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(lora_path, max_rank, "sv_ratio", ratio, device, dtype, output_filename)
return io.NodeOutput(path)
class LoRAResizeFrobenius(io.ComfyNode):
"""Resize LoRA to preserve Frobenius norm target."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAResizeFrobenius",
display_name="LoRA Resize (Frobenius)",
category="ModelUtils/LoRA/Resize",
description="Dynamically resize LoRA to preserve target fraction of Frobenius norm.",
inputs=[
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
tooltip="LoRA to resize"),
io.Int.Input("max_rank", default=128, min=1, max=1024,
tooltip="Maximum allowed rank"),
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
tooltip="Target Frobenius norm retention (0.9 = 90%)"),
io.String.Input("output_filename", default="resized_lora_fro"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device) -> io.NodeOutput:
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(lora_path, max_rank, "sv_fro", target, device, dtype, output_filename)
return io.NodeOutput(path)
class LoRAResizeCumulative(io.ComfyNode):
"""Resize LoRA to preserve cumulative SV target."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAResizeCumulative",
display_name="LoRA Resize (Cumulative)",
category="ModelUtils/LoRA/Resize",
description="Dynamically resize LoRA to preserve target fraction of cumulative singular values.",
inputs=[
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
tooltip="LoRA to resize"),
io.Int.Input("max_rank", default=128, min=1, max=1024,
tooltip="Maximum allowed rank"),
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
tooltip="Target cumulative SV retention (0.9 = 90%)"),
io.String.Input("output_filename", default="resized_lora_cumulative"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device) -> io.NodeOutput:
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(lora_path, max_rank, "sv_cumulative", target, device, dtype, output_filename)
return io.NodeOutput(path)
+36 -5
View File
@@ -3,10 +3,14 @@ import re
import torch import torch
import folder_paths import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file from safetensors.torch import save_file
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError 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): def load_documentation_from_file(filename):
@@ -58,6 +62,21 @@ class MergerLogic:
if not primary_model_name or primary_model_name == "None": if not primary_model_name or primary_model_name == "None":
raise ValueError("Model A is required to run the merge.") 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 = {} handlers = {}
for name in model_names.values(): for name in model_names.values():
if name and name != "None": if name and name != "None":
@@ -104,15 +123,20 @@ class MergerLogic:
discarded_keys = 0 discarded_keys = 0
error_keys = [] error_keys = []
for key in all_keys: for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely # Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns): if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1 discarded_keys += 1
pbar.update(1) pbar.update(1)
continue continue
# Pre-load Model A's tensor # Pre-load Model A's tensor with pinned memory for CUDA
tensor_a = primary_handler.get_tensor(key).to(device=process_device, dtype=process_dtype) 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 # Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns): if _matches_any_pattern(key, exclude_patterns):
@@ -162,12 +186,19 @@ class MergerLogic:
for handler in handlers.values(): for handler in handlers.values():
handler.__exit__(None, None, None) handler.__exit__(None, None, None)
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type 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) os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename") output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors") output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
save_file(merged_state_dict, output_path, metadata=metadata) 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" return f"{output_filename}.safetensors"
+282 -14
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 json
import struct import struct
import torch 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: 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.filename = filename
self.device = device self.device = device
self.mmap_mode = mmap_mode
self.header, self.header_size = self._read_header() self.header, self.header_size = self._read_header()
self.file = open(filename, "rb") 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): def __enter__(self):
return self return self
def __exit__(self, exc_type, exc_val, exc_tb): def __exit__(self, exc_type, exc_val, exc_tb):
if self.mmap_obj:
self.mmap_obj.close()
self.file.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__"] 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: if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file") raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key] metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"] offset_start, offset_end = metadata["data_offsets"]
tensor_bytes = None if self.mmap_mode and self.mmap_obj:
if offset_start != offset_end: 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) self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - 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): def _read_header(self):
with open(self.filename, "rb") as f: 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.") raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")
def get_union_keys(*handlers): def _read_tensor_from_file(
"""Returns the union of all tensor keys across multiple model handlers. filepath: str,
key: str,
metadata: Dict,
header_size: int
) -> Optional[torch.Tensor]:
"""Helper function to read a single tensor from file (for parallel execution).
Useful when you want to iterate over all unique keys present in any model. 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() all_keys = set()
for handler in handlers: for handler in handlers:
if handler is not None: if handler is not None:
@@ -92,11 +363,8 @@ def get_union_keys(*handlers):
return sorted(all_keys) return sorted(all_keys)
def get_intersection_keys(*handlers): def get_intersection_keys(*handlers) -> List[str]:
"""Returns only keys present in ALL model handlers. """Returns only keys present in ALL model handlers."""
Useful when you want to iterate only over keys that exist in every model.
"""
handlers = [h for h in handlers if h is not None] handlers = [h for h in handlers if h is not None]
if not handlers: if not handlers:
return [] return []
+23 -8
View File
@@ -1,15 +1,17 @@
import os import os
import json
import struct
import folder_paths import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io 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]: def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]:
"""Shared logic for all MetaKeys nodes.""" """Shared logic for all MetaKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name) 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')): if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors" 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}", "" return f"Conversion failed: {error_message}", ""
model_path = temp_safe_path 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) metadata_str = str(metadata)
for layer_name, tensor in model_weights.items():
shape = str(tensor.shape) layer_shapes = ""
shape = shape.replace("torch.Size(", "").replace(")", "") 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" name_shape = f"{layer_name}, {shape}\n"
layer_shapes += name_shape layer_shapes += name_shape
pbar.update(1)
return metadata_str, layer_shapes return metadata_str, layer_shapes
+34 -13
View File
@@ -1,9 +1,14 @@
import os import os
import re import re
import folder_paths import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file 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,
@@ -21,31 +26,47 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
else: else:
model_path_to_load = model_path 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()] patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
if not patterns: if not patterns:
raise ValueError("No keys/patterns provided to prune.") raise ValueError("No keys/patterns provided to prune.")
# Stream tensors and filter on the fly
pruned_tensors = {} pruned_tensors = {}
for key, tensor in model_weights.items(): with MemoryEfficientSafeOpen(model_path_to_load) as handler:
is_match = False all_keys = handler.keys()
if use_regex: pbar = comfy.utils.ProgressBar(len(all_keys))
if any(re.search(pattern, key) for pattern in patterns): for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = True is_match = False
else: if use_regex:
if any(pattern in key for pattern in patterns): if any(re.search(pattern, key) for pattern in patterns):
is_match = True is_match = True
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match: if not is_match:
pruned_tensors[key] = tensor 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") output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(pruned_tensors, output_path, metadata) save_file(pruned_tensors, output_path, metadata)
# Cleanup after operation
del pruned_tensors
cleanup_after_operation()
return output_path return output_path
+27 -7
View File
@@ -1,8 +1,13 @@
import os import os
import folder_paths import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file 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,
@@ -20,7 +25,13 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
else: else:
model_path_to_load = model_path 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()] 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()] 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)) key_map = dict(zip(old_keys, new_keys))
original_keys = list(model_weights.keys()) # Stream tensors and rename on the fly
renamed_tensors = {} renamed_tensors = {}
for key in original_keys: with MemoryEfficientSafeOpen(model_path_to_load) as handler:
new_key = key_map.get(key, key) original_keys = handler.keys()
renamed_tensors[new_key] = model_weights[key] 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") output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(renamed_tensors, output_path, metadata) save_file(renamed_tensors, output_path, metadata)
# Cleanup after operation
del renamed_tensors
cleanup_after_operation()
return output_path return output_path
+1 -1
View File
@@ -1,7 +1,7 @@
[project] [project]
name = "ComfyUI-ModelUtils" name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files." description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.1.7" version = "0.2.0"
license = { file = "LICENSE" } license = { file = "LICENSE" }
[project.urls] [project.urls]
+1
View File
@@ -0,0 +1 @@
# No external dependencies required