Compare commits

..
Author SHA1 Message Date
silveroxides 0a82e071b8 fix NodeOutput issue 2026-04-08 17:50:10 +02:00
silveroxides dc43040562 Rename downloader nodes and add nodes for displaying previews, metadata of model and workflows extracted from images/videos 2026-04-08 16:46:35 +02:00
silveroxides 39380f533d Remove subprocess, use av. More info in status output from Downloader nodes 2026-04-08 16:28:45 +02:00
silveroxides 2ac0d46935 Add progress bar 2026-04-08 15:43:05 +02:00
silveroxides 378c327406 Add subdirectory selector 2026-04-08 15:31:42 +02:00
silveroxides 4076c14cb5 Add requirements.txt entries 2026-04-08 15:16:16 +02:00
silveroxides c4c967172b Add Model preview/example image/video downloading, model hashing, preview/example workflow extraction from images and videos followed by compression of images to save storage, saving of hash database and much more 2026-04-08 15:04:24 +02:00
Silver 095ac89164 Bump version to 0.2.9 2026-03-07 19:10:13 +01:00
Silver 2d477767d6 Merge pull request #9 from silveroxides/fix/loading-fixes
Fix models getting stuck in memory when merging
2026-03-07 19:09:32 +01:00
silveroxides 4d49477f58 Fix models getting stuck in memory when merging 2026-03-07 16:32:43 +01:00
Silver 6c4d505d84 Bump version from 0.2.7 to 0.2.8 2026-02-24 17:41:25 +01:00
Silver 1e53a87427 Merge pull request #8 from silveroxides/feature/expand-lora-functionality
expand on functionality and reduce restrictions when extracting LoRA
2026-02-24 17:40:47 +01:00
silveroxides b497022c3c expand on functionality and reduce restrictions when extracting LoRA 2026-02-24 17:39:48 +01:00
Silver 5df8c6711a Merge pull request #7 from silveroxides/fix/loading-imporovements
improve memory efficiency
2026-02-12 08:07:39 +01:00
Silver f4cecdc49b Merge branch 'main' into fix/loading-imporovements 2026-02-12 08:07:31 +01:00
Silver 8b34d57003 Merge pull request #6 from silveroxides/patch/lora-changes
changes to LoRAs
2026-02-12 08:05:19 +01:00
silveroxides 9f8b337584 changes to LoRAs 2026-02-07 16:02:40 +01:00
silveroxides 0d011d1232 improve memory efficiency 2026-02-07 16:01:00 +01:00
silveroxides 6a119e63f5 Bump version from 0.2.6 to 0.2.7 2026-01-23 19:21:02 +01:00
silveroxides 4097326a82 Add pad/crop for mismatched tensors and interpolation mode for weight redistribution across rank 2026-01-23 19:04:51 +01:00
silveroxides 31eee604de Bump version from 0.2.5 to 0.2.6 2026-01-23 18:54:09 +01:00
silveroxides 2bfecf9c2d Fix resize_tensors to handle cases where model B is larger than model A 2026-01-23 18:53:40 +01:00
silveroxides 1ba0932d6b Bump version from 0.2.4 to 0.2.5 2026-01-23 18:42:23 +01:00
silveroxides b1ec55ff42 improve and correct merger operations based on papers 2026-01-23 18:41:35 +01:00
silveroxides 8deb4fd739 perf: Add low_memory mode to reduce RAM usage in LoRA merge
- MemoryEfficientSafeOpen now supports low_memory=True which disables mmap
- Uses direct file reads with readinto() to avoid OS page caching
- _deserialize_tensor avoids bytearray copy when input is already bytearray
- merge_loras_to_model now uses low_memory=True for base model

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

After

Width:  |  Height:  |  Size: 352 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 600 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 592 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 427 KiB

+321
View File
@@ -0,0 +1,321 @@
# ComfyUI-ModelUtils - Complete Project History
## Version Timeline
| Version | Commits | Key Features |
|---------|---------|--------------|
| 0.1.1 | `70d6c13`-`d67cbca` | Initial nodes (MetaKeys) |
| 0.1.2 | `6b662a4` | Icon assets |
| 0.1.4 | `cbe3502` | Key renaming/pruning nodes |
| 0.1.5 | `78b3198` | Model merging nodes |
| 0.1.6 | `31989be` | Embedding nodes, namespace fixes |
| 0.1.7 | `b6fae8d` | Mismatch handling, layer filtering, V3 migration |
| **0.2.0** | `3df073f` | LoRA extraction, memory optimization, pinned memory |
---
## Phase 1: Initial Release (v0.1.1)
**Commit:** `70d6c13` - "Add first nodes, toml and publish"
### Files Created
- `nodes/metakeys.py` - Get model metadata and layer keys
- `nodes/utils.py` - Shared utilities
- `__init__.py` - Node registration
- `pyproject.toml` - Package configuration
### Nodes
- `ModelMetaKeys` - Get diffusion model metadata
- `CheckpointMetaKeys` - Get checkpoint metadata
- `TextEncoderMetaKeys` - Get text encoder metadata
- `LoRAMetaKeys` - Get LoRA metadata
- `EmbeddingMetaKeys` - Get embedding metadata
---
## Phase 2: Key Manipulation (v0.1.4)
**Commit:** `4f08fe6` - "add key renaming and key pruning nodes"
### Files Created
- `nodes/renamekeys.py` - Rename tensor keys
- `nodes/prunekeys.py` - Remove tensor keys by pattern
### Nodes (5 variants each for model types)
- `ModelRenameKeys`, `CheckpointRenameKeys`, `TextEncoderRenameKeys`, `LoRARenameKeys`, `EmbeddingRenameKeys`
- `ModelPruneKeys`, `CheckpointPruneKeys`, `TextEncoderPruneKeys`, `LoRAPruneKeys`, `EmbeddingPruneKeys`
---
## Phase 3: Model Merging (v0.1.5)
**Commit:** `7959370` - "Add model merging nodes"
### Files Created
- `nodes/merger.py` - Main merging logic
- `nodes/merger_ops.py` - Merge operation modes
- `nodes/merger_utils.py` - `MemoryEfficientSafeOpen` class
- `docs/merger_2_model_modes.md` - Documentation
- `docs/merger_3_model_modes.md` - Documentation
### Nodes
- `CheckpointTwoMerger`, `DiffusionModelTwoMerger`, `TextEncoderTwoMerger`, `LoRATwoMerger`
- `CheckpointThreeMerger`, `DiffusionModelThreeMerger`, `TextEncoderThreeMerger`, `LoRAThreeMerger`
### Merge Modes (TWO_MODEL_MODES)
- Weighted Sum, Add Difference, Smooth Add Difference
- Train Difference, Sum Twice, Triple Sum
- Euclidean Add Difference, Distribution Crossover, Ties
- Slerp, Lerp, SVD Supermerger, SVD LoRA Extraction
### Merge Modes (THREE_MODEL_MODES)
- Add Difference (A + (B - C) × α)
### MemoryEfficientSafeOpen (initial)
```python
class MemoryEfficientSafeOpen:
def __init__(self, filename, device='cpu', mmap_mode=True)
def keys() -> List[str]
def get_tensor(key) -> Tensor
# mmap for zero-copy reads
```
---
## Phase 4: Embedding Nodes (v0.1.6)
**Commit:** `0d035ef` - "Add embedding nodes"
### Additions
- Added embedding variants to all node types
- Namespace conflict fixes (`68f22fa`, `5ea63ab`)
- Refactored names for universality (`232dc76`)
---
## Phase 5: Mismatch Handling (v0.1.7)
**Commits:** `19bc047` → `de27186`
### Features Added
- `mismatch_mode` parameter: skip/zeros/error
- Layer filtering with regex patterns
- `exclude_patterns` - Keep model A only (no merge)
- `discard_patterns` - Remove from output entirely
### Files Modified
- `merger.py` - Added mismatch handling logic
- `merger_ops.py` - Added `MissingTensorBehavior` enum, `MissingTensorError`
---
## Phase 6: V3 API Migration (v0.1.7)
**Commit:** `a3116cb` - "migrate to v3 schema and add LoRA extraction node"
### Changes
- Migrated all nodes from legacy API to `comfy_api.latest.io`
- `io.Schema`, `io.ComfyNode`, `io.NodeOutput`
- Added initial LoRA extraction node (diff-based)
---
## Phase 7: LyCORIS Extraction (pre-0.2.0)
**Commits:** `92d3c55` → `18679cc`
### Files Created
- `nodes/lycoris_extract.py` - LyCORIS-based extraction
- `nodes/lora_extract.py` - Simple diff extraction
### Initial Implementation
- Used `lycoris-lora` external dependency
- `lycoris.utils.extract_linear`, `extract_conv`
- Added `requirements.txt` with `lycoris-lora`
### Nodes Split (`18679cc`)
- `LyCORISExtractFixed`, `LyCORISExtractThreshold`
- `LyCORISExtractRatio`, `LyCORISExtractQuantile`
---
## Phase 8: Native SVD Port (0.2.0)
**Commits:** `6d7b2d1` → `eef7377`
### Removed
- `lycoris-lora` dependency
- `lycoris_extract.py`
- `lora_extract.py`
### NEW: [lora_extract_svd.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/lora_extract_svd.py)
**8 Rank Selection Functions:**
```python
_index_sv_fixed(S, dim) # Fixed rank
_index_sv_threshold(S, threshold) # Absolute threshold
_index_sv_ratio(S, ratio) # Relative to max
_index_sv_cumulative(S, target) # Cumulative sum %
_index_sv_fro(S, target) # Frobenius norm
_index_sv_knee(S) # Knee detection
_index_sv_cumulative_knee(S) # Cumsum knee
_index_sv_rel_decrease(S, tau) # Ratio drop
```
**SVD Functions:**
```python
_svd_extract_linear_lowrank(weight, rank, device, clamp_quantile)
_svd_extract_linear(weight, mode, mode_param, device, max_rank, clamp_quantile)
_svd_extract_conv(weight, mode, mode_param, device, max_rank, clamp_quantile)
_detect_fused_layer(key, shape) # QKV/MLP detection
_extract_chunked_layer(...) # Chunked extraction
```
**Nodes:**
- `LoRAExtractFixed` - Uses `svd_lowrank` (10x faster)
- `LoRAExtractRatio` - Relative threshold
- `LoRAExtractQuantile` - Cumulative target
- `LoRAExtractKnee` - Auto knee detection
- `LoRAExtractFrobenius` - Norm preservation
---
## Phase 9: Memory Management (0.2.0)
**Commit:** `6d7b2d1` - "refactor nodes and implement memory management"
### NEW: [device_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/device_utils.py)
Integrates with `comfy.model_management`:
```python
@dataclass
class DeviceCapabilities:
device_type, total_vram_gb, free_vram_gb, total_ram_gb, free_ram_gb
cpu_cores, recommended_workers, max_tensor_size_gb
supports_pinned_memory, supports_async_streams
get_device_capabilities(device) -> DeviceCapabilities
estimate_model_size(path) -> float # GB
can_fit_in_memory(paths, device, safety_factor) -> (bool, str)
prepare_for_large_operation(gb, device) -> bool # Calls free_memory()
cleanup_after_operation() # gc.collect() + soft_empty_cache()
get_optimal_workers(size_gb) -> int
```
### Updates to Existing Nodes
- `metakeys.py` - Header-only reading (no tensor loading)
- `renamekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
- `prunekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
- `merger.py` - Memory prep/cleanup
---
## Phase 10: Pinned Memory (0.2.0)
**Commit:** `b770786` - "Add pinned memory for faster merging and extraction"
### [merger_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/merger_utils.py) Enhancements
```python
# Context manager
class PinnedTensor:
def __enter__(self) -> Tensor # Pins memory
def __exit__(...) # Unpins
# Transfer helper (2-3x faster)
def transfer_to_gpu_pinned(tensor, device='cuda', dtype=None) -> Tensor
# MemoryEfficientSafeOpen additions
def get_tensor_to_gpu(key, device, dtype) # Pinned transfer
def keys_sorted_by_offset() -> List[str] # Sequential I/O
def get_tensors_parallel(keys, workers, sort_by_offset) -> Dict # Parallel
def get_tensor_as_dict(key) -> Dict # JSON config parsing
```
---
## Phase 11: SVD Optimization (0.2.0)
**Commit:** `65c1ae3` - "add svd_lowrank for fixed rank extraction"
### Changes
- Added `_svd_extract_linear_lowrank()` using `torch.svd_lowrank`
- Routes fixed mode to lowrank for **10x speedup**
- Changed `chunk_large_layers` default to `False`
---
## Phase 12: Bug Fixes (0.2.0)
**Commit:** `eef7377` - "Added .contiguous() to saves"
### Fix
Non-contiguous tensor error from SVD slices:
```python
tensor.to(dtype).cpu().contiguous() # Added .contiguous()
```
---
## Current File Structure
```
ComfyUI-ModelUtils/
├── __init__.py
├── pyproject.toml
├── requirements.txt (empty - no external deps)
├── README.md
├── docs/
│ ├── merger_2_model_modes.md
│ └── merger_3_model_modes.md
├── nodes/
│ ├── utils.py
│ ├── metakeys.py
│ ├── renamekeys.py
│ ├── prunekeys.py
│ ├── merger.py
│ ├── merger_ops.py
│ ├── merger_utils.py
│ ├── device_utils.py
│ └── lora_extract_svd.py
└── assets/
└── icon.png
```
---
## Performance Summary
| Feature | Improvement |
|---------|-------------|
| Fixed-rank SVD | **10x** (svd_lowrank) |
| CPU→GPU transfer | **2-3x** (pinned memory) |
| Metadata reading | **100x+** (header only) |
| Parallel I/O | **2-4x** (ThreadPoolExecutor) |
| Memory management | Auto-unload models before heavy ops |
| **LoRA Resize Fast** | **~91x** (merge-to-base extraction) |
---
## Phase 13: Fast LoRA Resize (2026-01-11)
### Summary
Added fast LoRA resize via base model extraction path, achieving ~91x speedup.
### Changes
- Added `resize_lora_via_base()` function using extraction code path
- Added 4 new fast resize nodes requiring base model input
### Files Modified
| File | Changes |
|------|---------|
| `nodes/lora_resize.py` | Added `resize_lora_via_base()`, 4 new node classes |
| `__init__.py` | Registered new nodes |
### New Nodes
- `LoRAResizeViaBaseFixed` - Fast fixed rank resize
- `LoRAResizeViaBaseRatio` - Fast SV ratio resize
- `LoRAResizeViaBaseFrobenius` - Fast Frobenius resize
- `LoRAResizeViaBaseCumulative` - Fast cumulative resize
+83 -3
View File
@@ -23,12 +23,47 @@
---
## Power-Up (DARE)
> Adds the unique capabilities of Model B to Model A using the Drop and Rescale (DARE) technique, which often preserves the knowledge of the base model better than simple additions.
> Adds the unique capabilities of Model B to Model A using the Drop and Rescale (DARE) technique. This implementation handles shape mismatches between models by padding and uses a randomized dropout mask.
**Models Used:** A, B
**Parameters:**
- **Alpha:** The dropout rate. This is the proportion of unique weights from Model B that are *dropped* before merging. A higher value means less of B is merged.
- **Alpha:** The dropout rate ($p$). This is the proportion of delta parameters from Model B that are randomly set to zero.
- **Beta:** A final multiplier for the rescaled difference before it's added to Model A.
- **Rescaling Logic:** Remaining weights are automatically rescaled by $1/(1-p)$ as per the DARE paper to approximate the original embeddings.
---
## Enhanced Man Interp
> Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control.
**Models Used:** A, B
**Parameters:**
- **Alpha:** Interpolation strength.
- **Beta:** Lower mean threshold for filtering differences.
- **Gamma:** Upper mean threshold for filtering differences.
- **Delta:** Smoothness factor (mix between randomized mask and powered differences).
---
## Enhanced Auto Interp
> Automated version of the enhanced interpolation mode that dynamically calculates thresholds based on mean differences.
**Models Used:** A, B
**Parameters:**
- **Alpha:** Interpolation strength.
- **Beta:** Threshold adjustment factor.
- **Gamma:** Smoothness factor.
---
## Weight-Sum Cutoff
> A linear interpolation mode that only applies the merge to weights whose differences fall within a specific threshold range.
**Models Used:** A, B
**Parameters:**
- **Alpha:** Interpolation weight (multiplier for the difference).
- **Beta:** Upper threshold for the difference cutoff.
- **Gamma:** Lower threshold for the difference cutoff.
---
@@ -39,4 +74,49 @@
**Parameters:**
- **Alpha:** The Rank (dimension) for the LoRA's standard layers.
- **Beta:** The Rank (dimension) specifically for 3x3 convolution layers.
- **Gamma:** The clamp quantile for weight values. `0.99` is a typical value to prevent outlier weights from dominating.
- **Gamma:** The clamp quantile for weight values. `0.99` is a typical value to prevent outlier weights from dominating.
---
## Layer Mismatch Handling
When merging models with different layer structures (e.g., fine-tuned models with added/removed layers, or LoRAs with partial layer coverage), the `mismatch_mode` parameter controls behavior:
| Mode | Behavior |
|------|----------|
| `skip` | Missing layers in B use A's values **(default)** |
| `zeros` | Missing layers in B are treated as zeros |
| `error` | Fail if any layer is missing |
**Common scenarios:**
- Merging a fine-tuned model with added/removed layers
- Combining LoRA-extracted differences with base models
- Cross-architecture experiments
**Note:** Extra layers in Model B that don't exist in Model A are currently ignored. The output will always have the same layer structure as Model A.
---
## Layer Filtering (Regex Patterns)
Use regex patterns to control which layers are merged:
### `exclude_patterns`
Layers matching any pattern will **keep Model A's values only** (no merge).
### `discard_patterns`
Layers matching any pattern will be **removed entirely** from the output.
**Pattern format:**
- Whitespace-separated regex patterns
- Patterns use **substring matching** (not full match)
- Example: `text_model lora` matches any key containing "text_model" OR "lora"
- Example: `layer\.[0-5]\.` matches layers 0-5 using regex syntax
**Pattern Examples:**
| Pattern | Matches |
|---------|---------|
| `text_model` | All text encoder layers |
| `\.norm` | All normalization layers |
| `attn\.(q\|k\|v)` | Query, key, value attention weights |
| `block\.[0-9]\.` | Blocks 0-9 |
+44 -4
View File
@@ -22,22 +22,62 @@
---
## Extract-Features
> A powerful mode that identifies features present in both `(B - A)` and `(C - A)` and adds them to A. Allows for fine-grained control over combining aspects based on their similarity.
> A powerful mode that identifies features present in both `(B - A)` and `(C - A)` and adds them to A. It uses per-vector cosine similarity to decide how much of each feature to keep, allowing for fine-grained control over combining aspects.
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** Weights the merge between Model B (`0.0`) and Model C (`1.0`).
- **Beta:** Controls the focus on similarity (`0.0`) versus dissimilarity (`1.0`).
- **Gamma:** A bias exponent for similarity. Higher values increase the bias.
- **Gamma:** A bias exponent for similarity calculation.
- **Delta:** A final multiplier for the extracted features before they are added to Model A.
---
## Add-Dissimilarities
> Identifies features that are dissimilar between Model B and Model C and adds them to Model A. Useful for combining unique aspects of two different models.
> Identifies features that are dissimilar between Model B and Model C (relative to A) and adds them to Model A. Useful for combining unique aspects of two different models.
**Models Used:** A, B, C
**Parameters:**
- **Alpha:** Weights the comparison between Model B (`0.0`) and Model C (`1.0`).
- **Beta:** A multiplier for the final extracted dissimilar features.
- **Gamma:** A bias exponent for the similarity calculation.
- **Gamma:** A bias exponent for the similarity calculation.
---
## Layer Mismatch Handling
When merging models with different layer structures, the `mismatch_mode` parameter controls behavior:
| Mode | Behavior |
|------|----------|
| `skip` | Missing layers in B/C use A's values **(default)** |
| `zeros` | Missing layers in B/C are treated as zeros |
| `error` | Fail if any layer is missing |
**Special case for B - C subtraction (Add-Difference, Train-Difference):**
- If **both** B and C have the layer: Normal subtraction is performed
- If **either** B or C is missing the layer: The difference is treated as `None` (no contribution from this layer), and Model A's value is used
**Common scenarios:**
- Merging a fine-tuned model with added/removed layers
- Extracting style differences between models with slightly different architectures
- Combining partial LoRAs or adapters
**Note:** Extra layers in Models B/C that don't exist in Model A are currently ignored. The output will always have the same layer structure as Model A.
---
## Layer Filtering (Regex Patterns)
Use regex patterns to control which layers are merged:
### `exclude_patterns`
Layers matching any pattern will **keep Model A's values only** (no merge).
### `discard_patterns`
Layers matching any pattern will be **removed entirely** from the output.
**Pattern format:**
- Whitespace-separated regex patterns
- Patterns use **substring matching** (not full match)
- Example: `text_model lora` matches any key containing "text_model" OR "lora"
+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
}
-12
View File
@@ -1,12 +0,0 @@
from .metakeys import *
from .renamekeys import *
from .prunekeys import *
from .merger import *
__all__ = [
"ModelMetaKeys", "TextEncoderMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", "EmbeddingMetaKeys",
"ModelRenameKeys", "TextEncoderRenameKeys", "LoRARenameKeys", "CheckpointRenameKeys", "EmbeddingRenameKeys",
"ModelPruneKeys", "TextEncoderPruneKeys", "LoRAPruneKeys", "CheckpointPruneKeys", "EmbeddingPruneKeys",
"CheckpointTwoMerger", "ModelTwoMerger", "TextEncoderTwoMerger", "LoRATwoMerger", "EmbeddingTwoMerger",
"CheckpointThreeMerger", "ModelThreeMerger", "TextEncoderThreeMerger", "LoRAThreeMerger", "EmbeddingThreeMerger",
]
+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")
+341
View File
@@ -0,0 +1,341 @@
import os
import folder_paths
from comfy_api.latest import io
from .downloader_utils import scan_and_process
def get_subdirectories(folder_name: str) -> list[str]:
"""Helper to get a list of subdirectories for a given model folder name."""
paths = folder_paths.get_folder_paths(folder_name)
if not paths:
return ["/"]
subdirs = set()
for base_path in paths:
if not os.path.exists(base_path):
continue
for root, dirs, _ in os.walk(base_path):
# Ignore hidden directories directly during walk
dirs[:] = [d for d in dirs if not d.startswith('.')]
rel_path = os.path.relpath(root, base_path)
if rel_path == ".":
subdirs.add("/")
else:
# Store with forward slashes for consistency
subdirs.add("/" + rel_path.replace("\\", "/"))
return sorted(list(subdirs))
class BaseInfoMetaDownloaderNode(io.ComfyNode):
CATEGORY = "ModelUtils/Downloader"
OUTPUT_NODE = True
@classmethod
def get_common_inputs(cls):
return {
"recursive": io.Boolean.Input("recursive", default=True, tooltip="Whether to scan subdirectories recursively."),
"nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All", tooltip="Filter models based on NSFW level in metadata."),
"max_examples": io.Int.Input("max_examples", default=0, min=0, max=100, tooltip="Maximum number of examples to process."),
"threads": io.Int.Input("threads", default=4, min=1, max=32, tooltip="Number of threads to use for processing."),
"api_key": io.String.Input("api_key", default="", tooltip="API key for accessing the model repository.")
}
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name=cls.__name__.replace("Downloader", " Downloader"),
category=cls.CATEGORY,
inputs=list(cls.get_common_inputs().values()),
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
description=("Scans specified directories for models, extracts metadata, generates previews, and optionally compresses files.")
)
@classmethod
def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
scan_dirs = cls.get_scan_dirs(**kwargs)
if not scan_dirs:
return io.NodeOutput("No directories to scan.")
# Call the scan_and_process function
stats = scan_and_process(
scan_dirs=scan_dirs,
recursive=recursive,
nsfw_level=nsfw_level,
max_examples=max_examples,
api_key=api_key if api_key else None,
threads=threads
)
mb_saved = stats.get("space_saved", 0) / (1024 * 1024)
output_msg = (
f"Scan Complete for {len(scan_dirs)} directory/directories.\n"
f"Total Models Processed: {stats.get('total_processed', 0)}\n"
f"New Models Hashed: {stats.get('new_hashes', 0)}\n"
f"Images Processed: {stats.get('images_processed', 0)}\n"
f"Workflows Extracted: {stats.get('workflows_extracted', 0)}\n"
f"Space Saved via Compression: {mb_saved:.2f} MB"
)
return io.NodeOutput(output_msg)
@classmethod
def get_scan_dirs(cls, **kwargs) -> list[str]:
return []
class CheckpointInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("checkpoints"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="Checkpoint Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("checkpoints")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class DiffusionModelInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("diffusion_models"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="Diffusion Model Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("diffusion_models")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class LoRAInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("loras"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="LoRA Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("loras")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class EmbeddingInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("embeddings"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="Embedding Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("embeddings")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class VAEInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("vae"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="VAE Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("vae")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class ControlNetInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("controlnet"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="ControlNet Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("controlnet")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class ManualPathInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.append(io.String.Input("scan_dir", default=""))
return io.Schema(
node_id=cls.__name__,
display_name="Manual Path Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, scan_dir, **kwargs):
return super().execute(
recursive=recursive,
nsfw_level=nsfw_level,
max_examples=max_examples,
threads=threads,
api_key=api_key,
scan_dir=scan_dir,
**kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
scan_dir = kwargs.get("scan_dir", "")
if scan_dir.strip():
return [scan_dir.strip()]
return []
+702
View File
@@ -0,0 +1,702 @@
import os
import sys
import hashlib
import json
import sqlite3
import requests
import glob
from pathlib import Path
import threading
import concurrent.futures
import codecs
import shutil
try:
from PIL import Image
HAS_PIL = True
except ImportError:
HAS_PIL = False
print("Warning: Pillow not found. Image workflow extraction disabled.")
try:
from mutagen.mp4 import MP4
HAS_MUTAGEN = True
except ImportError:
HAS_MUTAGEN = False
print("Warning: mutagen not found. MP4 tag extraction will rely entirely on av fallback.")
try:
import av
HAS_AV = True
except ImportError:
HAS_AV = False
print("Warning: av (PyAV) not found. MP4 metadata extraction fallback disabled.")
# Add the reference module to the Python path
# Assuming downloader_utils.py is inside nodes/, the root is one level up
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(root_dir, 'unifiedefficientloader-reference'))
try:
from unifiedefficientloader import UnifiedSafetensorsLoader
HAS_UNIFIED_LOADER = True
except ImportError:
HAS_UNIFIED_LOADER = False
print("Warning: unifiedefficientloader not found. Safetensors metadata extraction disabled.")
# --- Database Path Logic ---
try:
import folder_paths
base_user_dir = folder_paths.get_public_user_directory("default")
if not base_user_dir:
base_user_dir = root_dir
except (ImportError, AttributeError):
base_user_dir = root_dir
PLUGIN_DATA_DIR = os.path.join(base_user_dir, "ComfyUI-ModelUtils")
os.makedirs(PLUGIN_DATA_DIR, exist_ok=True)
DEFAULT_DB_PATH = os.path.join(PLUGIN_DATA_DIR, "cache.db")
# --- Workflow Extraction Logic ---
def unescape_and_parse_nested_json(json_string):
if not isinstance(json_string, str): return None
stripped = json_string.strip()
if not (stripped.startswith('{') or stripped.startswith('[')): return None
try: return json.loads(stripped)
except json.JSONDecodeError: pass
try:
first_parse = json.loads(stripped)
if isinstance(first_parse, str): return json.loads(first_parse)
except json.JSONDecodeError: pass
try:
decoded = codecs.decode(stripped, "unicode_escape")
return json.loads(decoded)
except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass
try:
unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\")
unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t")
return json.loads(unescaped)
except json.JSONDecodeError: pass
return None
def get_workflow_blob(metadata):
for key in ['workflow', 'prompt']:
val = metadata.get(key)
if val:
parsed = unescape_and_parse_nested_json(val)
if parsed: return parsed
for val in metadata.values():
parsed = unescape_and_parse_nested_json(val)
if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
return parsed
return None
def process_image_workflow_and_compress(filepath_str):
if not HAS_PIL: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
p = Path(filepath_str)
if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
workflow = None
actual_format = None
stats = {"images_processed": 1, "workflows_extracted": 0, "space_saved": 0}
try:
with Image.open(p) as img:
actual_format = img.format
# CivitAI sometimes serves PNGs with .jpeg extension in URL
if actual_format == 'PNG':
workflow = get_workflow_blob(img.info)
if workflow:
# Save workflow
json_path = p.with_suffix('.json')
if not json_path.exists():
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(workflow, f, indent=2, ensure_ascii=False)
print(f"Extracted workflow to {json_path}")
stats["workflows_extracted"] = 1
# Compress image to WebP
webp_path = p.with_suffix('.webp')
if not webp_path.exists():
orig_size = os.path.getsize(p)
img.save(webp_path, 'WEBP', quality=85, method=4)
new_size = os.path.getsize(webp_path)
stats["space_saved"] = max(0, orig_size - new_size)
print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}")
# Delete original if we created a webp replacement
if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG':
os.remove(p)
print(f"Deleted original uncompressed image: {p.name}")
except Exception as e:
print(f"Failed to process image {p.name}: {e}")
return stats
# --- Video Workflow Logic ---
def get_tag_value(mp4, keys):
for key in keys:
if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None
for key in keys:
search_suffix = key.split(':')[-1].lower()
for tag_key, val in mp4.tags.items():
if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None
return None
def extract_metadata_av(mp4_path):
if not HAS_AV: return None
try:
with av.open(str(mp4_path)) as container:
return {k.lower(): v for k, v in container.metadata.items()}
except Exception as e:
print(f"av extraction error for {mp4_path}: {e}")
return None
def process_video_workflow(filepath_str):
p = Path(filepath_str)
if p.suffix.lower() != '.mp4': return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
json_path = p.with_suffix('.json')
if json_path.exists(): return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
workflow_data = None
mp4 = None
stats = {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
if HAS_MUTAGEN:
try: mp4 = MP4(p)
except: pass
if mp4:
for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]:
val = get_tag_value(mp4, tag_keys)
if val:
if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore')
try:
parsed = json.loads(val)
if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
if 'prompt' in parsed and isinstance(parsed['prompt'], str):
try: workflow_data = json.loads(parsed['prompt'])
except: workflow_data = parsed
else:
workflow_data = parsed
break
except json.JSONDecodeError:
# Try more aggressive parsing if it's a nested string
parsed = unescape_and_parse_nested_json(val)
if parsed:
workflow_data = parsed
break
if not workflow_data:
av_data = extract_metadata_av(p)
if av_data:
for key in ['workflow', 'prompt', 'comment', 'description']:
if key in av_data:
val = av_data[key]
try:
parsed = json.loads(val)
if isinstance(parsed, dict):
if 'prompt' in parsed and isinstance(parsed['prompt'], str):
try: workflow_data = json.loads(parsed['prompt'])
except: workflow_data = parsed
else:
workflow_data = parsed
break
except json.JSONDecodeError:
parsed = unescape_and_parse_nested_json(val)
if parsed:
workflow_data = parsed
break
if workflow_data:
try:
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(workflow_data, f, indent=2, ensure_ascii=False)
print(f"Extracted workflow to {json_path.name}")
stats["workflows_extracted"] = 1
except Exception as e:
print(f"Error saving workflow for {p.name}: {e}")
return stats
# --- Constants for Civitai NSFW levels ---
NSFW_LEVELS = {
"None": 1,
"Soft": 2,
"Mature": 4,
"X": 8,
"XXX": 16,
"All": 32 # Anything
}
class HashCache:
def __init__(self, db_path=None):
self.db_path = db_path if db_path else DEFAULT_DB_PATH
self.is_sqlite = self.db_path.endswith('.db') or self.db_path.endswith('.sqlite')
self.lock = threading.Lock()
if self.is_sqlite:
self._init_sqlite()
else:
self.cache = self._load_json()
def _init_sqlite(self):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS file_hashes (
filepath TEXT PRIMARY KEY,
mtime REAL,
size INTEGER,
hash TEXT
)
''')
conn.commit()
def _load_json(self):
if os.path.exists(self.db_path):
with open(self.db_path, "r", encoding="utf-8") as f:
return json.load(f)
return {}
def _save_json(self):
with open(self.db_path, "w", encoding="utf-8") as f:
json.dump(self.cache, f, indent=4)
def _get_sqlite_hash(self, filepath, mtime, size):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,))
row = cursor.fetchone()
if row:
db_mtime, db_size, db_hash = row
if db_mtime == mtime and db_size == size:
return db_hash
return None
def _set_sqlite_hash(self, filepath, mtime, size, file_hash):
with sqlite3.connect(self.db_path, timeout=10.0) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash)
VALUES (?, ?, ?, ?)
''', (filepath, mtime, size, file_hash))
conn.commit()
def get_hash(self, filepath):
if not os.path.exists(filepath):
return None
mtime = os.path.getmtime(filepath)
size = os.path.getsize(filepath)
key = filepath
with self.lock:
if self.is_sqlite:
cached_hash = self._get_sqlite_hash(key, mtime, size)
if cached_hash: return cached_hash
else:
if key in self.cache:
entry = self.cache[key]
if entry["mtime"] == mtime and entry["size"] == size:
return entry["hash"]
# Calculate new hash (outside the lock to allow parallel hashing)
print(f"Calculating SHA256 for: {filepath}")
sha256_hash = self.calculate_sha256(filepath)
with self.lock:
if self.is_sqlite:
self._set_sqlite_hash(key, mtime, size, sha256_hash)
else:
self.cache[key] = {
"mtime": mtime,
"size": size,
"hash": sha256_hash
}
self._save_json()
return sha256_hash
@staticmethod
def calculate_sha256(filename):
sha256_hash = hashlib.sha256()
with open(filename, "rb") as f:
# Read and update hash string value in blocks of 4K
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
class CivitaiAPI:
def __init__(self, api_key=None):
self.api_key = api_key
self.base_url = "https://civitai.com/api/v1"
def get_headers(self):
headers = {}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
return headers
def get_model_info_by_hash(self, file_hash):
url = f"{self.base_url}/model-versions/by-hash/{file_hash}"
response = requests.get(url, headers=self.get_headers())
if response.status_code == 200:
return response.json()
elif response.status_code == 404:
return None
else:
response.raise_for_status()
class Downloader:
@staticmethod
def download_file(url, target_path, headers=None):
if headers is None:
headers = {}
temp_path = target_path + ".downloading"
downloaded_size = 0
if os.path.exists(temp_path):
downloaded_size = os.path.getsize(temp_path)
if downloaded_size > 0:
headers["Range"] = f"bytes={downloaded_size}-"
response = requests.get(url, headers=headers, stream=True)
if response.status_code == 416: # Range not satisfiable, file probably changed or completed
os.remove(temp_path)
downloaded_size = 0
if "Range" in headers:
del headers["Range"]
response = requests.get(url, headers=headers, stream=True)
if response.status_code not in [200, 206]:
print(f"Failed to download from {url}, status code: {response.status_code}")
return False
mode = "ab" if response.status_code == 206 else "wb"
total_size = int(response.headers.get('content-length', 0)) + downloaded_size
print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)")
with open(temp_path, mode) as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
os.rename(temp_path, target_path)
return True
def get_potential_preview_files(base_path):
exts = ["png", "jpg", "jpeg", "webp", "mp4"]
files = []
for ext in exts:
files.append(f"{base_path}.preview.{ext}")
files.append(f"{base_path}.{ext}")
return files
def preview_exists(base_path):
previews = get_potential_preview_files(base_path)
for p in previews:
if os.path.exists(p):
return True
return False
def next_example_image_path(base_path, ext):
i = 0
while len(glob.glob(f"{base_path}.example.{i}.*")) > 0:
i += 1
return f"{base_path}.example.{i}.{ext}"
def filter_images_by_nsfw(images, threshold_level):
valid_images = []
for img in images:
rating = img.get("nsfwLevel", 1)
if rating <= threshold_level:
valid_images.append(img)
return valid_images
def extract_safetensors_metadata(filepath):
if not HAS_UNIFIED_LOADER:
return None
try:
# Use low_memory=True to only read the header and metadata without loading tensors
loader = UnifiedSafetensorsLoader(filepath, low_memory=True)
metadata = loader.metadata()
loader.close()
return metadata if metadata else None
except Exception as e:
print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}")
return None
def process_file(filepath, cache, api, nsfw_threshold, max_examples):
filepath_str = str(filepath)
base_path = os.path.splitext(filepath_str)[0]
is_safetensors = filepath_str.lower().endswith(".safetensors")
stats = {"new_hashes": 0, "images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
# Check if all targets already exist
info_path = f"{base_path}.civitai.info"
st_metadata_path = f"{base_path}.metadata.json"
has_info = os.path.exists(info_path)
has_preview = preview_exists(base_path)
has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path)
# Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file
for p in get_potential_preview_files(base_path):
if os.path.exists(p):
ext_lower = p.lower()
if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
res = process_image_workflow_and_compress(p)
for k in stats: stats[k] += res.get(k, 0)
elif ext_lower.endswith('.mp4'):
res = process_video_workflow(p)
for k in stats: stats[k] += res.get(k, 0)
for ex in glob.glob(f"{base_path}.example.*"):
ext_lower = ex.lower()
if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
res = process_image_workflow_and_compress(ex)
for k in stats: stats[k] += res.get(k, 0)
elif ext_lower.endswith('.mp4'):
res = process_video_workflow(ex)
for k in stats: stats[k] += res.get(k, 0)
# If everything exists and we don't want examples or already have enough examples, we might skip API logic
if has_info and has_preview and has_st_metadata and max_examples == 0:
print(f"Skipping fully populated model: {filepath_str}")
return stats
print(f"Processing model: {filepath_str}")
# Safetensors built-in metadata extraction
if is_safetensors and not has_st_metadata:
metadata = extract_safetensors_metadata(filepath_str)
if metadata:
with open(st_metadata_path, "w", encoding="utf-8") as f:
json.dump(metadata, f, indent=4)
print(f"Saved safetensors metadata to {st_metadata_path}")
mtime = os.path.getmtime(filepath_str)
size = os.path.getsize(filepath_str)
# We want to know if it's a NEW hash or cached
is_cached = False
with cache.lock:
if cache.is_sqlite:
cached_hash = cache._get_sqlite_hash(filepath_str, mtime, size)
if cached_hash: is_cached = True
else:
if filepath_str in cache.cache:
entry = cache.cache[filepath_str]
if entry["mtime"] == mtime and entry["size"] == size:
is_cached = True
file_hash = cache.get_hash(filepath_str)
if not file_hash:
return stats
if not is_cached:
stats["new_hashes"] += 1
try:
model_info = api.get_model_info_by_hash(file_hash)
except Exception as e:
print(f"API Error for {filepath_str}: {e}")
return stats
if not model_info:
print(f"Model not found on Civitai: {filepath_str}")
return stats
# 1. Save Info
if not has_info:
with open(info_path, "w", encoding="utf-8") as f:
json.dump(model_info, f, indent=4)
print(f"Saved metadata to {info_path}")
images = model_info.get("images", [])
valid_images = filter_images_by_nsfw(images, nsfw_threshold)
if not valid_images:
print(f"No valid images found matching NSFW criteria for {filepath_str}")
return stats
# 2. Save Preview
if not has_preview and len(valid_images) > 0:
preview_img = valid_images[0]
url = preview_img.get("url")
ext = "jpeg" # Default fallback
if ".png" in url.lower(): ext = "png"
elif ".mp4" in url.lower(): ext = "mp4"
elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
preview_path = f"{base_path}.preview.{ext}"
if Downloader.download_file(url, preview_path):
if ext in ("png", "jpg", "jpeg"):
res = process_image_workflow_and_compress(preview_path)
for k in stats: stats[k] += res.get(k, 0)
elif ext == "mp4":
res = process_video_workflow(preview_path)
for k in stats: stats[k] += res.get(k, 0)
# 3. Save Examples
if max_examples > 0:
existing_examples_count = len(glob.glob(f"{base_path}.example.*"))
# We divide by 2 roughly if there are json files, let's just count unique example indices
existing_indices = set()
for f in glob.glob(f"{base_path}.example.*"):
parts = f.split('.example.')
if len(parts) > 1:
idx = parts[1].split('.')[0]
if idx.isdigit(): existing_indices.add(idx)
needed_examples = max_examples - len(existing_indices)
if needed_examples > 0:
example_images = valid_images[1:max_examples+1]
# Skip images we probably already downloaded by skipping the first N
example_images_to_dl = example_images[len(existing_indices):]
for img in example_images_to_dl:
url = img.get("url")
ext = "jpeg"
if ".png" in url.lower(): ext = "png"
elif ".mp4" in url.lower(): ext = "mp4"
elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
example_path = next_example_image_path(base_path, ext)
if Downloader.download_file(url, example_path):
if ext in ("png", "jpg", "jpeg"):
res = process_image_workflow_and_compress(example_path)
for k in stats: stats[k] += res.get(k, 0)
elif ext == "mp4":
res = process_video_workflow(example_path)
for k in stats: stats[k] += res.get(k, 0)
return stats
import numpy as np
from PIL import Image, ImageOps
import torch
def load_image_tensor(path):
if not os.path.exists(path):
return None
try:
img = Image.open(path)
img = ImageOps.exif_transpose(img)
image = img.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return image
except Exception as e:
print(f"Error loading image {path}: {e}")
return None
def get_model_workflows(base_path):
workflows = []
# Main workflow
main_json = f"{base_path}.json"
if os.path.exists(main_json):
workflows.append(main_json)
# Example workflows
# They are named like base_path.example.N.json
# or base_path.preview.json (if extracted from preview)
preview_json = f"{base_path}.preview.json"
if os.path.exists(preview_json) and preview_json not in workflows:
workflows.append(preview_json)
example_jsons = glob.glob(f"{base_path}.example.*.json")
# Sort them by index
example_jsons.sort()
for ej in example_jsons:
if ej not in workflows:
workflows.append(ej)
return workflows
def get_model_metadata_file(base_path):
# Prefer civitai.info if available
civitai_info = f"{base_path}.civitai.info"
if os.path.exists(civitai_info):
return civitai_info
# Fallback to metadata.json (safetensors metadata)
metadata_json = f"{base_path}.metadata.json"
if os.path.exists(metadata_json):
return metadata_json
return None
def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4):
"""
Scans directories for models and fetches their metadata/previews.
Adapted for direct Python API calls rather than argparse CLI.
"""
db_path = db_path if db_path else DEFAULT_DB_PATH
cache = HashCache(db_path)
api = CivitaiAPI(api_key)
model_exts = [".safetensors", ".ckpt", ".pt"]
files_to_process = []
if isinstance(scan_dirs, str):
scan_dirs = [scan_dirs]
for scan_dir in scan_dirs:
path = Path(scan_dir)
if recursive:
for ext in model_exts:
# rglob includes all directories, so we must filter out hidden ones
for f in path.rglob(f"*{ext}"):
if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
files_to_process.append(f)
else:
for ext in model_exts:
for f in path.glob(f"*{ext}"):
if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
files_to_process.append(f)
nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"])
print(f"Found {len(files_to_process)} models to process.")
global_stats = {
"total_processed": len(files_to_process),
"new_hashes": 0,
"images_processed": 0,
"workflows_extracted": 0,
"space_saved": 0
}
# Import and use ComfyUI's ProgressBar
try:
import comfy.utils
pbar = comfy.utils.ProgressBar(len(files_to_process))
except ImportError:
pbar = None
if threads > 1:
print(f"Starting {threads} threads...")
with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process]
for future in concurrent.futures.as_completed(futures):
try:
res = future.result()
if res:
for k in global_stats: global_stats[k] += res.get(k, 0)
except Exception as e:
print(f"Error in thread: {e}")
if pbar:
pbar.update(1)
else:
for filepath in files_to_process:
res = process_file(filepath, cache, api, nsfw_threshold, max_examples)
if res:
for k in global_stats: global_stats[k] += res.get(k, 0)
if pbar:
pbar.update(1)
return global_stats
+957
View File
@@ -0,0 +1,957 @@
"""
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, max_rank: int = None) -> int:
"""Cumulative mode - keep enough SVs to reach target % of total.
Calculates relative to max_rank if provided, otherwise relative to full.
"""
if max_rank is not None and max_rank < len(S):
total = torch.sum(S[:max_rank])
else:
total = torch.sum(S)
if total < 1e-8:
return 1
cumsum = torch.cumsum(S, dim=0) / total
rank = int(torch.searchsorted(cumsum, target).item()) + 1
return max(1, min(rank, len(S) - 1))
def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Frobenius norm mode - preserve target fraction of Frobenius norm.
Calculates relative to max_rank if provided, otherwise relative to full.
This means "retain target% of what's achievable within max_rank".
"""
if max_rank is not None and max_rank < len(S):
# Calculate relative to what's achievable within max_rank
S_capped = S[:max_rank]
S_sq = S_capped.pow(2)
total_sq = torch.sum(S_sq)
else:
S_sq = S.pow(2)
total_sq = torch.sum(S_sq)
if total_sq < 1e-8:
return 1
# Cumsum of all S (not capped) to find where we reach target
cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq
rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1
return max(1, min(rank, len(S) - 1))
def _index_sv_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
"""Knee detection - find elbow point on singular value curve."""
n = len(S)
if n < 3:
return 1
s_max, s_min = S[0].item(), S[-1].item()
if s_max - s_min < min_sv:
return 1
# Normalize to [0, 1]
s_norm = (S - s_min) / (s_max - s_min)
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
# Distance from diagonal
distances = (x_norm + s_norm - 1).abs()
rank = torch.argmax(distances).item() + 1
return max(1, min(rank, n - 1))
def _index_sv_cumulative_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int:
"""Knee detection on cumulative singular value curve."""
n = len(S)
if n < 3:
return 1
total = torch.sum(S)
if total < min_sv:
return 1
cumsum = torch.cumsum(S, dim=0) / total
y_min, y_max = cumsum[0].item(), cumsum[-1].item()
if y_max - y_min < min_sv:
return 1
y_norm = (cumsum - y_min) / (y_max - y_min)
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
distances = (y_norm - x_norm).abs()
rank = torch.argmax(distances).item() + 1
return max(1, min(rank, n - 1))
def _index_sv_rel_decrease(S: torch.Tensor, tau: float = 0.1) -> int:
"""Relative decrease mode - stop when consecutive SV ratio drops below tau."""
if len(S) < 2:
return 1
ratios = S[1:] / (S[:-1] + 1e-8)
for k in range(len(ratios)):
if ratios[k] < tau:
return k + 1
return len(S)
def _compute_rank(S: torch.Tensor, mode: str, mode_param: float,
max_rank: int = None) -> int:
"""Compute rank based on mode and parameters."""
if mode == "fixed":
rank = _index_sv_fixed(S, int(mode_param))
elif mode == "threshold":
rank = _index_sv_threshold(S, mode_param)
elif mode == "ratio":
rank = _index_sv_ratio(S, mode_param)
elif mode == "quantile" or mode == "sv_cumulative":
rank = _index_sv_cumulative(S, mode_param, max_rank)
elif mode == "sv_fro":
rank = _index_sv_fro(S, mode_param, max_rank)
elif mode == "sv_knee":
rank = _index_sv_knee(S)
elif mode == "sv_cumulative_knee":
rank = _index_sv_cumulative_knee(S)
elif mode == "sv_rel_decrease":
rank = _index_sv_rel_decrease(S, mode_param)
else:
rank = _index_sv_fixed(S, int(mode_param) if mode_param else 64)
if max_rank is not None:
rank = min(rank, max_rank)
return rank
# =============================================================================
# SVD Extraction Functions
# =============================================================================
def _svd_extract_linear_lowrank(
weight: torch.Tensor,
rank: int,
device: str,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms.
Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# Clamp rank to valid range
max_possible_rank = min(out_dim, in_dim)
rank = max(1, min(rank, max_possible_rank))
# Check if decomposition is worthwhile
if rank >= max_possible_rank:
return weight, "full"
# Use svd_lowrank for faster low-rank approximation
try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
# U: [out_dim, rank], S: [rank], V: [in_dim, rank]
Vh = V.T # [rank, in_dim]
except Exception as e:
raise RuntimeError(f"svd_lowrank failed: {e}")
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_linear(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
SVD decomposition for linear layers.
For fixed rank mode, uses svd_lowrank which is much faster.
For adaptive modes, uses full SVD to analyze all singular values.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# For fixed rank, use the faster svd_lowrank
if mode == "fixed" and mode_param > 0:
target_rank = int(mode_param)
if max_rank is not None:
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
# Compute rank based on mode
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim):
return weight, "full"
# Truncate to rank
U = U[:, :rank]
S = S[:rank]
Vh = Vh[:rank, :]
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
max_val = torch.quantile(S, clamp_quantile)
S = S.clamp(max=max_val)
except RuntimeError:
pass # Skip clamping if quantile fails
# Construct LoRA matrices
# lora_up = U @ diag(S), lora_down = Vh
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff for optional sparse bias
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_conv(
weight: torch.Tensor,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
) -> tuple:
"""
SVD decomposition for conv2d layers.
Returns:
(lora_down, lora_up, rank), "low rank" or (weight, "full")
"""
out_ch, in_ch, kh, kw = weight.shape
is_1x1 = (kh == 1 and kw == 1)
# Flatten for SVD
if is_1x1:
mat = weight.view(out_ch, in_ch) # [out_ch, in_ch]
else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile)
if mode_str == "full":
return weight, "full"
lora_down, lora_up, diff = result
rank = lora_down.shape[0]
# Reshape for conv
if is_1x1:
lora_down = lora_down.unsqueeze(-1).unsqueeze(-1) # [rank, in_ch, 1, 1]
lora_up = lora_up.unsqueeze(-1).unsqueeze(-1) # [out_ch, rank, 1, 1]
else:
lora_down = lora_down.reshape(rank, in_ch, kh, kw)
lora_up = lora_up.reshape(out_ch, rank, 1, 1)
return (lora_down, lora_up, diff), "low rank"
# =============================================================================
# Chunked Extraction for Large Layers
# =============================================================================
def _detect_fused_layer(key: str, shape: tuple) -> int:
"""Detect fused QKV/MLP layers and return number of chunks."""
if len(shape) != 2:
return 1
out_dim, in_dim = shape
key_lower = key.lower()
# QKV layers (3 chunks)
if "qkv" in key_lower and out_dim % 3 == 0 and out_dim // 3 >= in_dim // 4:
return 3
# KV layers (2 chunks)
if "kv" in key_lower and "qkv" not in key_lower and out_dim % 2 == 0:
return 2
# Fused MLP (2 chunks)
if ("linear1" in key_lower or "gate_up" in key_lower) and out_dim % 2 == 0 and out_dim >= in_dim * 2:
return 2
return 1
def _extract_chunked_layer(
weight_diff: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
) -> tuple:
"""Extract LoRA from fused layer by chunking."""
out_dim, in_dim = weight_diff.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :]
try:
result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
# Combine chunks
# Note: Using the first chunk's down matrix as the basis for all chunks
# is a better approximation than a simple mean if the chunks share the same input space (like QKV).
# Ideally, we would use a more sophisticated joint SVD, but this is a reasonable fallback.
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
# =============================================================================
# Pattern Matching
# =============================================================================
def _compile_patterns(pattern_string: str) -> 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,
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,
lazy_load: bool = True,
force_clear_cache: bool = True,
) -> dict[str, torch.Tensor]:
"""
Extract LoRA from difference between two models.
Args:
model_a_path: Finetuned model path
model_b_path: Base model path
mode: Rank selection mode
linear_param: Mode parameter for linear layers
conv_param: Mode parameter for conv layers
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, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns):
return "skipped", None
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key)
if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
if mismatch_mode == "zeros":
# For zeros mode, we treat the missing base as a zeroed tensor of same shape
# which is already captured by weight_diff = tensor_a
pass
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
# For zeros/fallback, use tensor_a as the difference
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:
del weight_diff
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[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:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
del weight_diff
return "chunked", layer_results
print(f"[LoRA Extract] Failed: {key}: {e}")
del weight_diff
return "skipped", None
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype)
status = "extracted"
del weight_diff
return status, layer_results
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
output_sd.update(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
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 _reconstruct_dots(key: str) -> str:
"""Reconstruct dot structure from underscored key using heuristics."""
# Block indices: input_blocks_1 -> input_blocks.1
blocks = [
"input_blocks", "output_blocks", "middle_block",
"transformer_blocks", "single_transformer_blocks",
"double_blocks", "single_blocks",
"down_blocks", "up_blocks", "mid_block",
"attentions", "resnets", "upsamplers", "downsamplers"
]
for b in blocks:
key = re.sub(f"{b}_(\\d+)", f"{b}.\\1", key)
# Sequential numbers: .1_1 -> .1.1
key = re.sub(r"(\.\d+)_(\d+)", r"\1.\2", key)
# Separators that should be preceded by dot
separators = [
"transformer_blocks", "single_transformer_blocks",
"attn1", "attn2", "attn",
"img_attn", "txt_attn",
"to_q", "to_k", "to_v", "to_out",
"q_proj", "k_proj", "v_proj", "out_proj",
"qkv", "proj", "ff", "net",
"norm1", "norm2", "norm3", "norm",
"time_emb_proj"
]
for s in separators:
key = key.replace(f"_{s}", f".{s}")
# Special case: to_out_0 -> to_out.0
key = re.sub(r"to_out_(\d+)", r"to_out.\1", key)
return key
def _format_lora_key(key: str) -> str:
"""
Format the key for LoRA saving.
Standardizes to 'diffusion_model.' or 'transformer.' prefix.
"""
if key.endswith(".weight"):
key = key[:-7]
# Handle ComfyUI Checkpoint wrapper
if key.startswith("model.diffusion_model."):
# Check if it's a Diffusers-format transformer inside a checkpoint
inner_key = key[22:] # len("model.diffusion_model.")
if inner_key.startswith("transformer_blocks") or inner_key.startswith("single_transformer_blocks"):
return f"transformer.{inner_key}"
# Standard checkpoint or Original Flux
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys
if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"):
return f"transformer.{key}"
# Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"):
core = key[10:]
dotted = _reconstruct_dots(core)
# Check if reconstructed key implies transformer format
if dotted.startswith("transformer_blocks") or dotted.startswith("single_transformer_blocks"):
return f"transformer.{dotted}"
return f"diffusion_model.{dotted}"
# Handle already correct prefixes
if key.startswith("diffusion_model.") or key.startswith("transformer."):
return key
# Handle known Unet blocks (Diffusers format UNet)
if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]):
return f"diffusion_model.{key}"
# Default fallback to diffusion_model
return f"diffusion_model.{key}"
def _get_common_inputs():
return [
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
io.Boolean.Input("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
tooltip="Clamp outlier singular values"),
io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.String.Input("output_filename", default="extracted_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="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=16384,
tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384,
tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache
)
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=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
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=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
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=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_sd = extract_lora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
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=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
+565
View File
@@ -0,0 +1,565 @@
"""
LoRA Multi-Merge - Merge multiple LoRAs into a single LoRA file.
Resolves different naming conventions and ranks via concatenation.
"""
import os
import torch
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from typing import List, Dict, Tuple, Optional
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from .lora_resize import (
detect_lora_format,
extract_lora_pairs,
detect_lora_rank,
_compile_patterns,
_matches_any_pattern,
_format_lora_key
)
def merge_multi_loras(
lora_paths: List[str],
lora_weights: List[float],
merge_mode: str,
device: str,
save_dtype: torch.dtype,
output_filename: str,
verbose: bool = True,
) -> str:
"""
Merge multiple LoRAs into a single LoRA file.
Modes:
- concatenate: Mathematically sound merge by stacking ranks.
New rank = sum(input ranks). Alpha set to new rank.
- weighted_sum: Weighted sum of A and B weights (Kohya style).
If ranks differ, smaller ones are zero-padded to match the largest rank.
New rank = max(input ranks).
"""
# Estimate memory
total_size_gb = sum(estimate_model_size(p) for p in lora_paths)
if verbose:
print(f"[LoRA Multi-Merge] Preparing memory for {total_size_gb:.2f}GB operation...")
print(f"[LoRA Multi-Merge] Mode: {merge_mode}")
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
try:
# 1. Analyze all LoRAs
layer_map = {}
lora_infos = []
for i, handler in enumerate(handlers):
keys = handler.keys()
fmt = detect_lora_format(keys)
pairs = extract_lora_pairs(keys, fmt)
rank, alpha = detect_lora_rank(handler, pairs)
info = {
"name": os.path.basename(lora_paths[i]),
"format": fmt,
"pairs": pairs,
"rank": rank,
"alpha": alpha,
"weight": lora_weights[i],
"scale": alpha / rank if rank > 0 else 1.0
}
lora_infos.append(info)
if verbose:
print(f"[LoRA Multi-Merge] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}, format={fmt['format']}")
for block_name in pairs.keys():
core = _format_lora_key(block_name)
if core not in layer_map:
layer_map[core] = []
layer_map[core].append(i)
if verbose:
print(f"[LoRA Multi-Merge] Total unique layers after resolving naming: {len(layer_map)}")
output_sd = {}
pbar = comfy.utils.ProgressBar(len(layer_map))
# 2. Merge each core layer
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"):
downs = []
ups = []
is_conv = False
max_rank = 0
# First pass: load and determine max rank
for idx in indices:
info = lora_infos[idx]
handler = handlers[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys:
continue
t_down = handler.get_tensor(block_keys["down"])
t_up = handler.get_tensor(block_keys["up"])
if device == 'cuda':
t_down = transfer_to_gpu_pinned(t_down, device, torch.float32)
t_up = transfer_to_gpu_pinned(t_up, device, torch.float32)
else:
t_down = t_down.to(device=device, dtype=torch.float32)
t_up = t_up.to(device=device, dtype=torch.float32)
# Store with original info for padding/weighting
is_conv = len(t_down.shape) == 4
current_rank = t_down.shape[0]
max_rank = max(max_rank, current_rank)
# Apply scale immediately for concatenate mode, or keep for later
if merge_mode == "concatenate":
effective_weight = info["weight"] * info["scale"]
t_up = t_up * effective_weight
downs.append((t_down, info))
ups.append((t_up, info))
if not downs:
pbar.update(1)
continue
if merge_mode == "concatenate":
# Stack ranks
merged_down = torch.cat([d[0] for d in downs], dim=0)
merged_up = torch.cat([u[0] for u in ups], dim=1)
new_rank = merged_down.shape[0]
new_alpha = float(new_rank)
else:
# weighted_sum (Kohya style)
# Result = sum( weight_i * Padded(B_i) ), sum( weight_i * Padded(A_i) )
# Note: We apply scale to weights here to normalize different alphas
merged_down = torch.zeros_like(downs[0][0])
# Need to handle different ranks via padding
target_down_shape = list(downs[0][0].shape)
target_down_shape[0] = max_rank
target_up_shape = list(ups[0][0].shape)
target_up_shape[1] = max_rank
merged_down = torch.zeros(target_down_shape, device=device, dtype=torch.float32)
merged_up = torch.zeros(target_up_shape, device=device, dtype=torch.float32)
for (t_d, info_d), (t_u, info_u) in zip(downs, ups):
r = t_d.shape[0]
w = info_d["weight"]
# We also apply sqrt(scale) to both to distribute the scale factor?
# Kohya just uses weights. But to be safe with different alphas,
# we apply scale to the final delta.
# For direct weight merge, we'll just use weights.
if r < max_rank:
# Pad down: [r, in...] -> [max_rank, in...]
pad_d = [0] * (len(t_d.shape) * 2)
pad_d[-1] = max_rank - r # last dim in pad is first dim in tensor (reversed)
# Wait, F.pad uses reverse order of dims.
# For [r, in], padding is (0,0, 0, max_rank-r)
padding_d = [0, 0] * (len(t_d.shape) - 1) + [0, max_rank - r]
t_d = torch.nn.functional.pad(t_d, tuple(padding_d))
# Pad up: [out, r...] -> [out, max_rank...]
# For [out, r], padding is (0, max_rank-r, 0, 0)
padding_u = [0, max_rank - r] + [0, 0] * (len(t_u.shape) - 1)
t_u = torch.nn.functional.pad(t_u, tuple(padding_u))
merged_down += w * t_d
merged_up += w * t_u
new_rank = max_rank
# Alpha is usually max of alphas or same as rank
new_alpha = float(max_rank)
# Store
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(new_alpha, dtype=save_dtype)
# Cleanup
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
# 3. Final Summary
if verbose:
matched_stats = {i: 0 for i in range(len(lora_paths))}
for indices in layer_map.values():
for idx in indices:
matched_stats[idx] += 1
print(f"[LoRA Multi-Merge] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
print(f"[LoRA Multi-Merge] Output state dict has {len(output_sd)} tensors")
# 4. Metadata
metadata = {
"ss_training_comment": f"Merged {len(lora_paths)} LoRAs via concatenation",
"ss_network_module": "networks.lora",
}
# 4. 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)
if verbose:
print(f"[LoRA Multi-Merge] Saved merged LoRA to {output_path}")
return output_path
finally:
for h in handlers:
h.__exit__(None, None, None)
cleanup_after_operation()
class LoRAMultiMerge(io.ComfyNode):
"""Merge multiple LoRAs into a single LoRA file."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAMultiMerge",
display_name="LoRA Multi-Merge",
category="ModelUtils/LoRA/Merge",
description="Merge 1-8 LoRAs into a single LoRA file. This resolves different ranks and naming conventions.",
inputs=[
io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2",
tooltip="Number of LoRAs to merge"),
# LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"),
io.Float.Input("weight_1", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 2
io.Combo.Input("lora_2", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_2", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 3
io.Combo.Input("lora_3", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_3", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 4
io.Combo.Input("lora_4", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 5
io.Combo.Input("lora_5", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_5", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 6
io.Combo.Input("lora_6", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_6", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 7
io.Combo.Input("lora_7", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_7", default=1.0, min=-10.0, max=10.0, step=0.01),
# LoRA 8
io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("merge_mode", options=["concatenate", "weighted_sum"], default="concatenate",
tooltip="concatenate: safe, increases rank. weighted_sum: fixed rank (to max input rank), mathematically lossy but standard in Kohya."),
io.String.Input("output_filename", default="merged_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_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
merge_mode, output_filename, save_dtype, device) -> io.NodeOutput:
count = int(lora_count)
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8]
valid_paths = []
valid_weights = []
for i in range(count):
if names[i] and names[i] != "None":
path = folder_paths.get_full_path_or_raise("loras", names[i])
valid_paths.append(path)
valid_weights.append(weights[i])
if not valid_paths:
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = merge_multi_loras(
lora_paths=valid_paths,
lora_weights=valid_weights,
merge_mode=merge_mode,
device=device,
save_dtype=dtype,
output_filename=output_filename
)
return io.NodeOutput(path)
class LoRAMultiMergeDARE(io.ComfyNode):
"""Merge multiple LoRAs into a single LoRA file using DARE-Ties method."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAMultiMergeDARE",
display_name="LoRA Multi-Merge (DARE-Ties)",
category="ModelUtils/LoRA/Merge",
description="Merge 1-8 LoRAs into a single LoRA file using DARE-Ties. Drops small values and resolves sign conflicts.",
inputs=[
io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2"),
# LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"),
io.Float.Input("weight_1", default=1.0, min=-10.0, max=10.0, step=0.01),
# ...
io.Combo.Input("lora_2", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_2", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("lora_3", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_3", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("lora_4", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("lora_5", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_5", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("lora_6", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_6", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("lora_7", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_7", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("drop_rate", default=0.1, min=0.0, max=1.0, step=0.01, tooltip="DARE drop rate"),
io.Float.Input("trim_quantile", default=0.2, min=0.0, max=1.0, step=0.01, tooltip="TIES trim quantile (drops smallest values)"),
io.Int.Input("seed", default=42, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_lora_dare"),
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_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput:
count = int(lora_count)
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8]
valid_paths = []
valid_weights = []
for i in range(count):
if names[i] and names[i] != "None":
path = folder_paths.get_full_path_or_raise("loras", names[i])
valid_paths.append(path)
valid_weights.append(weights[i])
if not valid_paths:
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = merge_multi_loras_dare(
lora_paths=valid_paths,
lora_weights=valid_weights,
drop_rate=drop_rate,
trim_quantile=trim_quantile,
seed=seed,
device=device,
save_dtype=dtype,
output_filename=output_filename
)
return io.NodeOutput(path)
def merge_multi_loras_dare(
lora_paths: List[str],
lora_weights: List[float],
drop_rate: float,
trim_quantile: float,
seed: int,
device: str,
save_dtype: torch.dtype,
output_filename: str,
verbose: bool = True,
) -> str:
"""
Merge multiple LoRAs using DARE-Ties method applied to weights.
Different ranks are handled via zero-padding to the max rank.
"""
total_size_gb = sum(estimate_model_size(p) for p in lora_paths)
if verbose:
print(f"[LoRA Multi-Merge DARE] Preparing memory for {total_size_gb:.2f}GB operation...")
print(f"[LoRA Multi-Merge DARE] Drop rate: {drop_rate}, Trim quantile: {trim_quantile}")
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
rng = torch.Generator(device=device).manual_seed(seed)
try:
# 1. Analyze
layer_map = {}
lora_infos = []
for i, handler in enumerate(handlers):
keys = handler.keys()
fmt = detect_lora_format(keys)
pairs = extract_lora_pairs(keys, fmt)
rank, alpha = detect_lora_rank(handler, pairs)
info = {
"name": os.path.basename(lora_paths[i]),
"pairs": pairs,
"rank": rank,
"alpha": alpha,
"weight": lora_weights[i]
}
lora_infos.append(info)
if verbose:
print(f"[LoRA Multi-Merge DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}")
for block_name in pairs.keys():
core = _format_lora_key(block_name)
if core not in layer_map: layer_map[core] = []
layer_map[core].append(i)
if verbose:
print(f"[LoRA Multi-Merge DARE] Total unique layers after resolving naming: {len(layer_map)}")
output_sd = {}
pbar = comfy.utils.ProgressBar(len(layer_map))
# 2. Merge
with torch.no_grad():
for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"):
downs = []
ups = []
max_rank = 0
for idx in indices:
info = lora_infos[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
max_rank = max(max_rank, t_d.shape[0])
downs.append((t_d, info["weight"]))
ups.append((t_u, info["weight"]))
if not downs:
pbar.update(1)
continue
def process_ties_dare(tensors, weights, dim_to_pad):
# Pad all to max_rank
padded = []
for t, w in tensors:
if t.shape[dim_to_pad] < max_rank:
padding = [0] * (len(t.shape) * 2)
# dim_to_pad 0 (down) -> last pair in pad
# dim_to_pad 1 (up) -> second to last pair in pad
rev_dim = len(t.shape) - 1 - dim_to_pad
padding[rev_dim*2 + 1] = max_rank - t.shape[dim_to_pad]
t = torch.nn.functional.pad(t, tuple(padding))
padded.append(t * w)
# DARE
if drop_rate > 0:
for i in range(len(padded)):
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float()
padded[i] = (padded[i] * mask) / (1 - drop_rate)
# TIES
# 1. Trim
if trim_quantile > 0:
for i in range(len(padded)):
flat = padded[i].abs().flatten()
k = int(len(flat) * trim_quantile)
if k > 0:
threshold = torch.kthvalue(flat, k).values
padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i])
# 2. Elect & Merge
stacked = torch.stack(padded) # [N, ...]
signs = torch.sign(stacked)
sum_signs = signs.sum(dim=0)
dominant_sign = torch.sign(sum_signs)
# Filter those matching dominant sign
mask = (signs == dominant_sign) & (dominant_sign != 0)
filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
# Average matching signs
count = mask.sum(dim=0)
result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
return result
merged_down = process_ties_dare(downs, [1.0]*len(downs), 0)
merged_up = process_ties_dare(ups, [1.0]*len(ups), 1)
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype)
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
# Final Summary
if verbose:
matched_stats = {i: 0 for i in range(len(lora_paths))}
for indices in layer_map.values():
for idx in indices:
matched_stats[idx] += 1
print(f"[LoRA Multi-Merge DARE] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
print(f"[LoRA Multi-Merge DARE] Output state dict has {len(output_sd)} tensors")
# 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, {"ss_training_comment": "Merged via DARE-Ties"})
return output_path
finally:
for h in handlers: h.__exit__(None, None, None)
cleanup_after_operation()
+986
View File
@@ -0,0 +1,986 @@
"""
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,
_svd_extract_linear,
_svd_extract_conv,
_compile_patterns,
_matches_any_pattern,
_format_lora_key,
)
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
- HuggingFace: *.lora.down.weight, *.lora.up.weight (dots instead of underscores)
- Full diff: *.diff, *.diff_b (full weight differences, not low-rank)
- 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,
"is_full_diff": False, # New flag for full diff format
}
# 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[:50] # Check more keys for full_diff detection
# Full diff format: *.diff (NOT .lora_down, just straight weight diff)
if any(k.endswith('.diff') for k in sample_keys):
format_info["format"] = "full_diff"
format_info["down_suffix"] = ".diff" # Used for pairing, but no up
format_info["up_suffix"] = None # No up weight in full diff
format_info["alpha_suffix"] = None
format_info["is_full_diff"] = True
# Kohya/A1111 format: lora_unet_down_blocks_0_*.lora_down.weight
elif 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"
# HuggingFace format: transformer.*.lora.down.weight (dots instead of underscores)
elif any('.lora.down.weight' in k for k in sample_keys):
format_info["format"] = "huggingface"
format_info["down_suffix"] = ".lora.down.weight"
format_info["up_suffix"] = ".lora.up.weight"
format_info["alpha_suffix"] = ".alpha"
# 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.
For full_diff format, groups .diff and .diff_b keys.
Returns:
Dict[block_name, {"down": key, "up": key, "alpha": key}]
For full_diff: {"diff": key, "diff_b": key}
"""
down_suffix = format_info["down_suffix"]
up_suffix = format_info["up_suffix"]
alpha_suffix = format_info["alpha_suffix"]
is_full_diff = format_info.get("is_full_diff", False)
pairs = {}
if is_full_diff:
# Full diff format: group .diff and .diff_b
for key in keys:
if key.endswith('.diff'):
block_name = key[:-5] # Remove .diff
if block_name not in pairs:
pairs[block_name] = {}
pairs[block_name]["diff"] = key
elif key.endswith('.diff_b'):
block_name = key[:-7] # Remove .diff_b
if block_name not in pairs:
pairs[block_name] = {}
pairs[block_name]["diff_b"] = key
else:
# Standard LoRA format
for key in keys:
if down_suffix and 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 and 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 and (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:
# Use kohya convention: S[0]/ratio
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, max_rank)
elif dynamic_method == "sv_fro" and dynamic_param is not None:
new_rank = _index_sv_fro(S, dynamic_param, max_rank)
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"]
# Standardize prefix using heuristics if possible
new_block_name = _format_lora_key(block_name)
output_sd[f"{new_block_name}{down_suffix}"] = result["lora_down"].to(save_dtype)
output_sd[f"{new_block_name}{up_suffix}"] = result["lora_up"].to(save_dtype)
if alpha_suffix:
output_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype)
del result
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
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=3072,
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=3072,
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=3072,
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=3072,
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)
# =============================================================================
# LoRA Merge To Model (Save merged model, skip extraction)
# =============================================================================
def merge_loras_to_model(
lora_paths: List[str],
lora_weights: List[float],
base_model_path: str,
device: str,
save_dtype: torch.dtype,
output_filename: str,
skip_patterns_str: str = "",
verbose: bool = True,
) -> str:
"""
Merge multiple LoRAs into a base model and save the result directly.
Unlike merge_multi_loras_via_base, this function:
- Does NOT extract the result back to a LoRA
- Saves the merged full model to the base model's directory
Args:
lora_paths: List of paths to LoRA files
lora_weights: List of weight strengths (0.0-2.0) for each LoRA
base_model_path: Path to base model
device: Processing device
save_dtype: Output dtype
output_filename: Output filename (without extension)
skip_patterns_str: Regex patterns for layers to skip
verbose: Print progress info
Returns:
Path to saved merged model
"""
# Estimate memory and prepare
total_size_gb = estimate_model_size(base_model_path)
for lp in lora_paths:
total_size_gb += estimate_model_size(lp)
if verbose:
print(f"[LoRA Merge To Model] Preparing memory for {total_size_gb:.2f}GB operation...")
print(f"[LoRA Merge To Model] Merging {len(lora_paths)} LoRAs with weights: {lora_weights}")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
# Open all files - use low_memory for base model to avoid OS page caching
base_handler = MemoryEfficientSafeOpen(base_model_path, low_memory=True)
lora_handlers = [MemoryEfficientSafeOpen(lp) for lp in lora_paths]
try:
# Detect format and extract pairs for each LoRA
lora_infos = []
for i, handler in enumerate(lora_handlers):
keys = handler.keys()
format_info = detect_lora_format(keys)
pairs = extract_lora_pairs(keys, format_info)
network_dim, network_alpha = detect_lora_rank(handler, pairs)
lora_infos.append({
"handler": handler,
"format_info": format_info,
"pairs": pairs,
"network_dim": network_dim,
"network_alpha": network_alpha,
"weight": lora_weights[i],
})
if verbose:
print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}")
# Common prefixes
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."]
LORA_PREFIXES = [
"lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_",
"lycoris_", "diffusion_model.", "transformer.", "unet."
]
def extract_core_layer_base(key: str) -> str:
result = key
if result.endswith(".weight"):
result = result[:-7]
elif result.endswith(".bias"):
result = result[:-5]
for prefix in BASE_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
return result
def extract_core_layer_lora(block_name: str) -> str:
result = block_name
for prefix in LORA_PREFIXES:
if result.startswith(prefix):
result = result[len(prefix):]
break
# Strip trailing .lora suffix (HuggingFace format leaves this after suffix removal)
if result.endswith(".lora"):
result = result[:-5]
return result.replace(".", "_")
# Build LoRA lookup: core layer name (underscored) -> list of (info, block_keys)
lora_lookup = {}
for info in lora_infos:
for block_name, block_keys in info["pairs"].items():
core = extract_core_layer_lora(block_name)
if core not in lora_lookup:
lora_lookup[core] = []
lora_lookup[core].append((info, block_keys))
# Compile skip patterns
skip_patterns = _compile_patterns(skip_patterns_str)
# Preserve metadata from base model
base_metadata = base_handler.metadata().copy() if base_handler.metadata() else {}
base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}"
output_sd = {}
stats = {"merged": 0, "copied": 0, "skipped": 0}
base_keys = list(base_handler.keys())
pbar = comfy.utils.ProgressBar(len(base_keys))
if verbose:
print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...")
with torch.no_grad():
for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"):
# Check skip patterns
if _matches_any_pattern(base_key, skip_patterns):
stats["skipped"] += 1
pbar.update(1)
continue
# Load base weight
cpu_base = base_handler.get_tensor(base_key)
# Only process weight tensors for LoRA merging
if base_key.endswith(".weight"):
core = extract_core_layer_base(base_key)
core_underscored = core.replace(".", "_")
# Check if any LoRA contributes to this layer
if core_underscored in lora_lookup:
# Transfer to GPU for computation
if device == 'cuda':
base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32)
else:
base_weight = cpu_base.to(device=device, dtype=torch.float32)
del cpu_base
# Accumulate deltas from all contributing LoRAs
for info, block_keys in lora_lookup[core_underscored]:
is_full_diff = info["format_info"].get("is_full_diff", False)
if is_full_diff:
# Full diff format
if "diff" not in block_keys:
continue
cpu_diff = info["handler"].get_tensor(block_keys["diff"])
if device == 'cuda':
delta = transfer_to_gpu_pinned(cpu_diff, device, torch.float32)
else:
delta = cpu_diff.to(device=device, dtype=torch.float32)
del cpu_diff
effective_scale = info["weight"]
else:
# Standard LoRA format
if "down" not in block_keys or "up" not in block_keys:
continue
cpu_down = info["handler"].get_tensor(block_keys["down"])
cpu_up = info["handler"].get_tensor(block_keys["up"])
if device == 'cuda':
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
else:
lora_down = cpu_down.to(device=device, dtype=torch.float32)
lora_up = cpu_up.to(device=device, dtype=torch.float32)
del cpu_down, cpu_up
# Get alpha
if "alpha" in block_keys:
alpha_tensor = info["handler"].get_tensor(block_keys["alpha"])
layer_alpha = float(alpha_tensor.item())
else:
layer_alpha = float(info["network_dim"])
layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0
effective_scale = layer_scale * info["weight"]
# Compute delta
is_conv = len(lora_down.shape) == 4
if is_conv:
in_rank, in_size, kernel_size, k_ = lora_down.shape
out_size, out_rank, _, _ = lora_up.shape
delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
delta = delta.reshape(out_size, in_size, kernel_size, kernel_size)
else:
delta = lora_up @ lora_down
del lora_down, lora_up
# Apply delta to base weight
base_weight = base_weight + effective_scale * delta
del delta
# Store merged weight
output_sd[base_key] = base_weight.to(save_dtype).cpu().contiguous()
del base_weight
stats["merged"] += 1
else:
# No LoRA contribution, copy as-is
output_sd[base_key] = cpu_base.to(save_dtype).contiguous()
del cpu_base
stats["copied"] += 1
else:
# Non-weight tensor (bias, norm, etc.), copy as-is
output_sd[base_key] = cpu_base.to(save_dtype).contiguous()
del cpu_base
stats["copied"] += 1
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
if verbose:
print(f"[LoRA Merge To Model] Done: {stats['merged']} merged, {stats['copied']} copied, {stats['skipped']} skipped")
# Save to base model directory
base_dir = os.path.dirname(base_model_path)
os.makedirs(base_dir, exist_ok=True)
output_path = os.path.join(base_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path, base_metadata)
print(f"[LoRA Merge To Model] Saved to {output_path}")
return output_path
finally:
base_handler.__exit__(None, None, None)
for handler in lora_handlers:
handler.__exit__(None, None, None)
cleanup_after_operation()
class LoRAMergeToModel(io.ComfyNode):
"""Merge multiple LoRAs into base model and save as full model."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAMergeToModel",
display_name="LoRA Merge To Model",
category="ModelUtils/LoRA/Merge",
description="Merge 1-8 LoRAs into a base model and save the result. Saves to base model directory.",
inputs=[
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Base model the LoRAs were trained on"),
io.Combo.Input("lora_count", options=["1", "2", "3", "4", "5", "6", "7", "8"], default="2",
tooltip="Number of LoRAs to merge"),
# LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"),
tooltip="First LoRA"),
io.Float.Input("weight_1", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 1"),
# LoRA 2
io.Combo.Input("lora_2", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Second LoRA"),
io.Float.Input("weight_2", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 2"),
# LoRA 3
io.Combo.Input("lora_3", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Third LoRA"),
io.Float.Input("weight_3", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 3"),
# LoRA 4
io.Combo.Input("lora_4", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Fourth LoRA"),
io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 4"),
# LoRA 5
io.Combo.Input("lora_5", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Fifth LoRA"),
io.Float.Input("weight_5", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 5"),
# LoRA 6
io.Combo.Input("lora_6", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Sixth LoRA"),
io.Float.Input("weight_6", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 6"),
# LoRA 7
io.Combo.Input("lora_7", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Seventh LoRA"),
io.Float.Input("weight_7", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 7"),
# LoRA 8
io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Eighth LoRA"),
io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 8"),
# Settings
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip"),
io.String.Input("output_filename", default="merged_model"),
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, base_model, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
skip_patterns, output_filename, save_dtype, device) -> io.NodeOutput:
# Build LoRA list based on count
count = int(lora_count)
lora_names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8][:count]
lora_weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8][:count]
# Filter out "None" entries
valid_loras = [(name, weight) for name, weight in zip(lora_names, lora_weights) if name != "None"]
if not valid_loras:
raise ValueError("At least one LoRA must be selected")
lora_names, lora_weights = zip(*valid_loras)
lora_paths = [folder_paths.get_full_path_or_raise("loras", name) for name in lora_names]
base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = merge_loras_to_model(
lora_paths=list(lora_paths),
lora_weights=list(lora_weights),
base_model_path=base_path,
device=device,
save_dtype=dtype,
output_filename=output_filename,
skip_patterns_str=skip_patterns,
)
return io.NodeOutput(path)
+764 -113
View File
@@ -1,10 +1,17 @@
import os
import re
import torch
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES
from .merger_utils import MemoryEfficientSafeOpen
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
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):
"""Loads documentation from a markdown file in the ../docs/ directory."""
@@ -16,158 +23,802 @@ def load_documentation_from_file(filename):
except FileNotFoundError:
return f"# Documentation File Not Found\n\nPlease ensure `{filename}` exists in the `docs` directory of the custom node."
# --- UI Dictionary Helper Functions ---
def _create_two_model_inputs(model_folder_name):
return {"required": {
"execution_mode": (["MERGE", "DOCUMENTATION ONLY"],),
"model_a": (["None"] + folder_paths.get_filename_list(model_folder_name),),
"model_b": (["None"] + folder_paths.get_filename_list(model_folder_name),),
"calc_mode": ([m.name for m in TWO_MODEL_MODES],),
"alpha": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}),
"beta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}),
"gamma": ("FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.001}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"output_filename": ("STRING", {"default": "merged_2_model"}),
"save_dtype": (["fp32", "fp16", "bf16"],),
"process_device": (["cuda", "cpu"],),
}}
def _create_three_model_inputs(model_folder_name):
return {"required": {
"execution_mode": (["MERGE", "DOCUMENTATION ONLY"],),
"model_a": (["None"] + folder_paths.get_filename_list(model_folder_name),),
"model_b": (["None"] + folder_paths.get_filename_list(model_folder_name),),
"model_c": (["None"] + folder_paths.get_filename_list(model_folder_name),),
"calc_mode": ([m.name for m in THREE_MODEL_MODES],),
"alpha": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}),
"beta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}),
"gamma": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}),
"delta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"output_filename": ("STRING", {"default": "merged_3_model"}),
"save_dtype": (["fp32", "fp16", "bf16"],),
"process_device": (["cuda", "cpu"],),
}}
def _compile_patterns(pattern_string):
"""Compiles whitespace-separated regex patterns into a list of compiled patterns.
Returns empty list if pattern_string is empty or whitespace-only.
Raises ValueError if any pattern is invalid regex.
"""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for pattern in pattern_string.split():
try:
patterns.append(re.compile(pattern))
except re.error as e:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}")
return patterns
def _matches_any_pattern(key, patterns):
"""Returns True if key matches any of the compiled regex patterns (substring search)."""
for pattern in patterns:
if pattern.search(key):
return True
return False
# --- Base Classes for Backend Logic ONLY ---
class MergerLogic:
def _execute_merge(self, model_names, calc_mode, all_modes, recipe_params, model_type):
"""Shared merge execution logic for all merger nodes."""
@staticmethod
def execute_merge(model_names, calc_mode, all_modes, recipe_params, model_type):
calc_mode_class = next((m for m in all_modes if m.name == calc_mode), None)
if not calc_mode_class: raise ValueError(f"Calc mode '{calc_mode}' not found.")
if not calc_mode_class:
raise ValueError(f"Calc mode '{calc_mode}' not found.")
primary_model_name = model_names.get('model_a')
if not primary_model_name or primary_model_name == "None": raise ValueError("Model A is required to run the merge.")
if not primary_model_name or primary_model_name == "None":
raise ValueError("Model A is required to run the merge.")
# Prepare memory before heavy operation
total_size_gb = 0
model_paths = []
for name in model_names.values():
if name and name != "None":
path = folder_paths.get_full_path(model_type, name)
if path:
total_size_gb += estimate_model_size(path)
model_paths.append(path)
if total_size_gb > 0:
process_device = recipe_params.get('device', 'cpu')
print(f"[Merger] Preparing memory for {total_size_gb:.2f}GB merge operation...")
prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device))
handlers = {}
for name in model_names.values():
if name and name != "None":
path = folder_paths.get_full_path(model_type, name)
if not path: raise FileNotFoundError(f"Model '{name}' not found.")
handlers[name] = MemoryEfficientSafeOpen(path)
all_keys = handlers[primary_model_name].keys()
metadata = handlers[primary_model_name].header.get("__metadata__", {})
if not path:
raise FileNotFoundError(f"Model '{name}' not found.")
handlers[name] = MemoryEfficientSafeOpen(path, low_memory=True)
primary_handler = handlers[primary_model_name]
all_keys = primary_handler.keys()
metadata = primary_handler.metadata()
# Convert mismatch_mode string to enum
mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip')
mismatch_mode = MissingTensorBehavior(mismatch_mode_str)
recipe_params['mismatch_mode'] = mismatch_mode
alignment_mode = recipe_params.get('alignment_mode', 'pad/crop')
# Compile filter patterns
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''))
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''))
# Pre-compute key differences for logging
primary_keys = set(all_keys)
for name, handler in handlers.items():
if name != primary_model_name:
secondary_keys = set(handler.keys())
missing = primary_keys - secondary_keys
extra = secondary_keys - primary_keys
if missing:
print(f"[Merger] {name} is missing {len(missing)} keys present in Model A")
if extra:
print(f"[Merger] {name} has {len(extra)} extra keys not in Model A (ignored)")
merged_state_dict = {}
pbar = comfy.utils.ProgressBar(len(all_keys))
save_dtype = recipe_params.pop('save_dtype')
save_torch_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(save_dtype)
recipe_params.update({"handlers": handlers}) # Add handlers, device and dtype are already in
for key in all_keys:
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
result = recipe.merge()
if isinstance(result, dict):
for r_key, r_tensor in result.items(): merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone()
else: merged_state_dict[key] = result.to(save_torch_dtype).cpu().clone()
pbar.update(1)
for handler in handlers.values(): handler.__exit__(None, None, None)
recipe_params.update({"handlers": handlers})
process_device = recipe_params.get('device', 'cpu')
process_dtype = recipe_params.get('dtype', torch.float32)
skipped_keys = 0
excluded_keys = 0
discarded_keys = 0
error_keys = []
with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
pbar.update(1)
continue
# Pre-load Model A's tensor with pinned memory for CUDA
cpu_tensor = primary_handler.get_tensor(key)
if process_device == 'cuda':
tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
else:
tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype)
del cpu_tensor
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
merged_state_dict[key] = tensor_a.detach().to(save_torch_dtype).cpu().clone()
excluded_keys += 1
pbar.update(1)
continue
# Pass tensor_a metadata to recipes for zeros mode and fallback
recipe_params['_tensor_a'] = tensor_a
recipe_params['_tensor_a_shape'] = tensor_a.shape
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
try:
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
result = recipe.merge()
except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR:
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
result = None
error_keys.append(key)
# Handle None result (mismatch occurred with skip mode)
if result is None:
result = tensor_a
skipped_keys += 1
if isinstance(result, dict):
for r_key, r_tensor in result.items():
merged_state_dict[r_key] = r_tensor.detach().to(save_torch_dtype).cpu().clone()
else:
# Ensure compatibility with Model A's architecture.
# If alignment_mode is 'pad/crop', we crop results that were padded.
# If alignment_mode is 'interpolate', resizing happened during operators.
if alignment_mode == 'pad/crop':
target_shape = recipe_params['_tensor_a_shape']
if result.shape != target_shape:
slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape))
result = result[slices]
merged_state_dict[key] = result.detach().to(save_torch_dtype).cpu().clone()
# Clean up references to allow GC immediately.
# Local loop variables must be explicitly deleted to prevent PyTorch from keeping tensors in VRAM.
recipe.clean()
del recipe_params['_tensor_a']
del tensor_a
del recipe
del result
if recipe_params.get('force_clear_cache', True):
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
# Log summary
if excluded_keys > 0:
print(f"[Merger] Excluded {excluded_keys} keys from merge (kept Model A only)")
if discarded_keys > 0:
print(f"[Merger] Discarded {discarded_keys} keys from output")
if skipped_keys > 0:
print(f"[Merger] Used Model A's values for {skipped_keys} keys due to mismatches")
if error_keys:
print(f"[Merger] Unexpected errors on {len(error_keys)} keys: {error_keys[:5]}{'...' if len(error_keys) > 5 else ''}")
for handler in handlers.values():
handler.__exit__(None, None, None)
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
output_dir = folder_paths.get_folder_paths(output_folder)[0]
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
output_dir = folder_paths.get_folder_paths(output_folder)[-1]
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
save_file(merged_state_dict, output_path, metadata=metadata)
return (f"{output_filename}.safetensors",)
class BaseTwoMerger(MergerLogic):
FUNCTION = "merge"; CATEGORY = "ModelUtils/Merging"
RETURN_TYPES = ("STRING", "STRING",); RETURN_NAMES = ("output_filename", "documentation",)
# Cleanup after heavy operation
del merged_state_dict
cleanup_after_operation()
def merge(self, execution_mode, model_a, model_b, calc_mode, alpha, beta, gamma, seed, output_filename, save_dtype, process_device):
return f"{output_filename}.safetensors"
# --- Two-Model Merger Nodes ---
class CheckpointTwoMerger(io.ComfyNode):
MODEL_TYPE = "checkpoints"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointTwoMerger",
display_name="Merge Checkpoints (2 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("checkpoints")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("checkpoints")),
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return ("Documentation mode active. No merge performed.", doc)
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename, = self._execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, self.MODEL_TYPE)
return (filename, doc)
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class BaseThreeMerger(MergerLogic):
FUNCTION = "merge"; CATEGORY = "ModelUtils/Merging"
RETURN_TYPES = ("STRING", "STRING",); RETURN_NAMES = ("output_filename", "documentation",)
def merge(self, execution_mode, model_a, model_b, model_c, calc_mode, alpha, beta, gamma, delta, seed, output_filename, save_dtype, process_device):
class ModelTwoMerger(io.ComfyNode):
MODEL_TYPE = "diffusion_models"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="ModelTwoMerger",
display_name="Merge Models (2 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("diffusion_models")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("diffusion_models")),
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class TextEncoderTwoMerger(io.ComfyNode):
MODEL_TYPE = "text_encoders"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderTwoMerger",
display_name="Merge Text Encoders (2 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("text_encoders")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("text_encoders")),
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class LoRATwoMerger(io.ComfyNode):
MODEL_TYPE = "loras"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRATwoMerger",
display_name="Merge LoRAs (2 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("loras")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("loras")),
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class EmbeddingTwoMerger(io.ComfyNode):
MODEL_TYPE = "embeddings"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingTwoMerger",
display_name="Merge Embeddings (2 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("embeddings")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("embeddings")),
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
# --- Three-Model Merger Nodes ---
class CheckpointThreeMerger(io.ComfyNode):
MODEL_TYPE = "checkpoints"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointThreeMerger",
display_name="Merge Checkpoints (3 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("checkpoints")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("checkpoints")),
io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("checkpoints")),
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return ("Documentation mode active. No merge performed.", doc)
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename, = self._execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, self.MODEL_TYPE)
return (filename, doc)
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
# --- Concrete 2-Model Nodes ---
class CheckpointTwoMerger(BaseTwoMerger):
MODEL_TYPE = "checkpoints"
@classmethod
def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE)
class ModelTwoMerger(BaseTwoMerger):
class ModelThreeMerger(io.ComfyNode):
MODEL_TYPE = "diffusion_models"
@classmethod
def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE)
class TextEncoderTwoMerger(BaseTwoMerger):
MODEL_TYPE = "text_encoders"
@classmethod
def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE)
class LoRATwoMerger(BaseTwoMerger):
MODEL_TYPE = "loras"
@classmethod
def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE)
class EmbeddingTwoMerger(BaseTwoMerger):
MODEL_TYPE = "embeddings"
@classmethod
def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE)
# --- Concrete 3-Model Nodes ---
class CheckpointThreeMerger(BaseThreeMerger):
MODEL_TYPE = "checkpoints"
@classmethod
def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE)
class ModelThreeMerger(BaseThreeMerger):
MODEL_TYPE = "diffusion_models"
@classmethod
def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE)
class TextEncoderThreeMerger(BaseThreeMerger):
MODEL_TYPE = "text_encoders"
@classmethod
def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE)
class LoRAThreeMerger(BaseThreeMerger):
MODEL_TYPE = "loras"
@classmethod
def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE)
class EmbeddingThreeMerger(BaseThreeMerger):
MODEL_TYPE = "embeddings"
@classmethod
def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE)
def define_schema(cls):
return io.Schema(
node_id="ModelThreeMerger",
display_name="Merge Models (3 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("diffusion_models")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("diffusion_models")),
io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("diffusion_models")),
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
__all__ = [
"CheckpointTwoMerger", "ModelTwoMerger", "TextEncoderTwoMerger", "LoRATwoMerger", "EmbeddingTwoMerger",
"CheckpointThreeMerger", "ModelThreeMerger", "TextEncoderThreeMerger", "LoRAThreeMerger", "EmbeddingThreeMerger"
]
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class TextEncoderThreeMerger(io.ComfyNode):
MODEL_TYPE = "text_encoders"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderThreeMerger",
display_name="Merge Text Encoders (3 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("text_encoders")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("text_encoders")),
io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("text_encoders")),
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class LoRAThreeMerger(io.ComfyNode):
MODEL_TYPE = "loras"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoRAThreeMerger",
display_name="Merge LoRAs (3 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("loras")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("loras")),
io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("loras")),
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
class EmbeddingThreeMerger(io.ComfyNode):
MODEL_TYPE = "embeddings"
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingThreeMerger",
display_name="Merge Embeddings (3 Models)",
category="ModelUtils/Merging",
inputs=[
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]),
io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("embeddings")),
io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("embeddings")),
io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("embeddings")),
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
+523 -50
View File
@@ -1,6 +1,107 @@
import torch
import scipy.ndimage
import torch.nn.functional as F
from enum import Enum
from .merger_utils import transfer_to_gpu_pinned
class MissingTensorBehavior(Enum):
"""Controls behavior when a tensor key is missing from a model."""
ERROR = "error" # Raise exception (strict mode)
SKIP = "skip" # Return None, let caller handle
ZEROS = "zeros" # Return zeros tensor matching shape from primary model
class MissingTensorError(Exception):
"""Raised when a tensor is missing and behavior is ERROR."""
pass
# ################# #
# UTILITIES
# ################# #
def resize_by_interpolation(tensor1, tensor2):
"""Resizes tensor2 to match tensor1's shape using interpolation."""
if tensor1 is None or tensor2 is None:
return tensor1, tensor2
if tensor1.shape == tensor2.shape:
return tensor1, tensor2
dims1 = len(tensor1.shape)
dims2 = len(tensor2.shape)
if dims1 != dims2 or dims1 > 4 or dims1 == 0:
return tensor1, tensor2
orig_dtype = tensor2.dtype
t2_f = tensor2.float()
# F.interpolate expects [B, C, ...]
# We treat the tensor as [1, 1, ...dims]
view_shape = [1, 1] + list(tensor2.shape)
t2_f = t2_f.view(view_shape)
# Select mode based on dimensions
if dims1 == 4:
# [Out, In, H, W] -> [1, 1, Out, In, H*W] (5D)
view_shape = [1, 1, tensor2.shape[0], tensor2.shape[1], tensor2.shape[2] * tensor2.shape[3]]
target_size = [tensor1.shape[0], tensor1.shape[1], tensor1.shape[2] * tensor1.shape[3]]
t2_f = t2_f.view(view_shape)
t2_res = F.interpolate(t2_f, size=target_size, mode='trilinear', align_corners=False)
return tensor1, t2_res.view(tensor1.shape).to(orig_dtype)
if dims1 == 1:
mode = 'linear'
elif dims1 == 2:
mode = 'bilinear'
else:
mode = 'trilinear'
t2_res = F.interpolate(t2_f, size=list(tensor1.shape), mode=mode, align_corners=False)
return tensor1, t2_res.view(tensor1.shape).to(orig_dtype)
def resize_tensors(tensor1, tensor2, mode='pad/crop'):
"""Resizes tensors to match shapes using either padding or interpolation."""
if tensor1 is None or tensor2 is None:
return tensor1, tensor2
if tensor1.shape == tensor2.shape:
return tensor1, tensor2
if mode == 'interpolate':
return resize_by_interpolation(tensor1, tensor2)
dims1 = len(tensor1.shape)
dims2 = len(tensor2.shape)
if dims1 != dims2 or dims1 > 4 or dims1 == 0:
return tensor1, tensor2
padding1 = []
padding2 = []
# F.pad expects: (last_dim_left, last_dim_right, prev_dim_top, prev_dim_bottom, ...)
for i in range(dims1 - 1, -1, -1):
s1 = tensor1.shape[i]
s2 = tensor2.shape[i]
if s1 < s2:
padding1.extend([0, s2 - s1])
padding2.extend([0, 0])
elif s2 < s1:
padding1.extend([0, 0])
padding2.extend([0, s1 - s2])
else:
padding1.extend([0, 0])
padding2.extend([0, 0])
if any(p > 0 for p in padding1):
tensor1 = F.pad(tensor1, tuple(padding1))
if any(p > 0 for p in padding2):
tensor2 = F.pad(tensor2, tuple(padding2))
return tensor1, tensor2
# ################# #
# MERGE OPERATORS
@@ -10,6 +111,7 @@ class Operation:
self.key = key
self.sources = tuple(sources)
self.merge_func = self.recurse
self.alignment_mode = 'pad/crop'
self.alpha = None
self.beta = None
self.gamma = None
@@ -27,43 +129,116 @@ class Operation:
raise NotImplementedError
def recurse(self, operation):
source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*source_tensors)
self._source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*self._source_tensors)
def merge(self):
return self.merge_func(self)
def clean(self):
if hasattr(self, '_source_tensors'):
for t in self._source_tensors:
del t
del self._source_tensors
for source in self.sources:
if hasattr(source, 'clean'):
source.clean()
class LoadTensor(Operation):
def __init__(self, key, model_name, handlers, device, dtype):
def __init__(self, key, model_name, handlers, device, dtype,
on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None):
super().__init__(key)
self.model_name = model_name
self.handlers = handlers
self.device = device
self.dtype = dtype
self.on_missing = on_missing
self.fallback_shape = fallback_shape
self.fallback_dtype = fallback_dtype
def merge(self) -> torch.Tensor:
if self.model_name not in self.handlers:
raise ValueError(f"Model '{self.model_name}' is required for this mode but was not provided.")
return self.handlers[self.model_name].get_tensor(self.key).to(device=self.device, dtype=self.dtype)
handler = self.handlers[self.model_name]
if self.key not in handler.keys():
if self.on_missing == MissingTensorBehavior.ERROR:
raise MissingTensorError(
f"Key '{self.key}' not found in model '{self.model_name}'"
)
elif self.on_missing == MissingTensorBehavior.SKIP:
return None
elif self.on_missing == MissingTensorBehavior.ZEROS:
if self.fallback_shape is None:
raise ValueError(
f"Cannot create zeros for '{self.key}' without fallback_shape"
)
dtype = self.fallback_dtype if self.fallback_dtype else self.dtype
return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype)
cpu_tensor = handler.get_tensor(self.key)
if self.device == 'cuda':
self._tensor = transfer_to_gpu_pinned(cpu_tensor, self.device, self.dtype)
else:
self._tensor = cpu_tensor.to(device=self.device, dtype=self.dtype)
del cpu_tensor
return self._tensor
def clean(self):
if hasattr(self, '_tensor'):
del self._tensor
class Multiply(Operation):
def __init__(self, key, alpha, *sources):
super().__init__(key, *sources)
self.alpha = alpha
def oper(self, a) -> torch.Tensor:
if a is None:
return None
return a * self.alpha
class Add(Operation):
def oper(self, a, b) -> torch.Tensor:
if a is None:
return b
if b is None:
return a
a, b = resize_tensors(a, b, mode=self.alignment_mode)
return a + b
class Sub(Operation):
def oper(self, a, b) -> torch.Tensor:
if a is None or b is None:
return None # Cannot compute difference with missing operand
a, b = resize_tensors(a, b, mode=self.alignment_mode)
return a - b
class SVD(Operation):
def __init__(self, key, alpha, beta, gamma, *sources):
def __init__(self, key, alpha, beta, gamma, *sources, target_shape=None):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma = int(alpha), int(beta), gamma
self.target_shape = target_shape
def oper(self, a, b) -> torch.Tensor:
if a is None or b is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
# For SVD extraction, we must ensure the diff matches the original Model A shape
# so the resulting LoRA weights are compatible with the target architecture.
if self.target_shape and a.shape != self.target_shape:
if self.alignment_mode == 'interpolate':
# Re-align both to match target_shape via interpolation
dummy = torch.empty(self.target_shape, device=a.device)
_, a = resize_by_interpolation(dummy, a)
_, b = resize_by_interpolation(dummy, b)
else:
slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(a.shape, self.target_shape))
a = a[slices]
b = b[slices]
diff, weights, conv2d = a - b, {}, len(a.size()) == 4
kernel_size = None if not conv2d else a.size()[2:4]
conv2d_3x3 = conv2d and kernel_size != (1, 1)
@@ -84,11 +259,19 @@ class SVD(Operation):
class Smooth(Operation):
def oper(self, a) -> torch.Tensor:
if a is None:
return None
filtered = scipy.ndimage.median_filter(a.detach().cpu().to(torch.float32).numpy(), size=3)
filtered = scipy.ndimage.gaussian_filter(filtered, sigma=1)
return torch.tensor(filtered, dtype=a.dtype, device=a.device)
class TrainDiff(Operation):
def oper(self, a, b, c) -> torch.Tensor:
if a is None or b is None or c is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
a, c = resize_tensors(a, c, mode=self.alignment_mode)
b, c = resize_tensors(b, c, mode=self.alignment_mode)
if torch.allclose(b.float(), c.float()): return torch.zeros_like(a)
diff_bc = b.float() - c.float()
dist_bc, dist_ab = torch.abs(diff_bc), torch.abs(b.float() - a.float())
@@ -99,27 +282,51 @@ class ExtractOp(Operation):
super().__init__(key, *args)
self.alpha, self.beta, self.gamma = alpha, beta, gamma
def oper(self, base: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
if a is None or b is None:
return None
if base is not None:
base, a = resize_tensors(base, a, mode=self.alignment_mode)
base, b = resize_tensors(base, b, mode=self.alignment_mode)
a, b = resize_tensors(a, b, mode=self.alignment_mode)
dtype = base.dtype if base is not None else a.dtype
base_f = base.float() if base is not None else 0
a_f, b_f = a.float() - base_f, b.float() - base_f
c = torch.cosine_similarity(a_f.flatten(), b_f.flatten(), 0).clamp(-1, 1)
c = torch.cosine_similarity(a_f, b_f, -1).clamp(-1, 1).unsqueeze(-1)
d = ((c + 1) / 2)**self.gamma
return (torch.lerp(a_f, b_f, self.alpha) * torch.lerp(d, 1 - d, self.beta)).to(a.dtype)
return (torch.lerp(a_f, b_f, self.alpha) * torch.lerp(d, 1 - d, self.beta)).to(dtype)
class Similarities(ExtractOp):
def oper(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return super().oper(None, a, b)
def oper(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
if a is None or b is None:
return None
return super().oper(None, a, b)
class PowerUpOp(Operation):
def __init__(self, key, alpha, seed, *sources):
super().__init__(key, *sources)
self.alpha, self.seed = alpha, seed
def oper(self, a, b):
if a is None or b is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta = b - a
keep_prob = 1 - self.alpha
if keep_prob <= 0:
return torch.zeros_like(delta)
rng = torch.Generator(device=delta.device).manual_seed(self.seed)
m = torch.bernoulli(torch.full_like(delta, 1 - self.alpha), generator=rng)
return (m * delta) / (1 - self.alpha) if (1 - self.alpha) > 0 else torch.zeros_like(delta)
m = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).to(delta.dtype)
return (m * delta) / keep_prob
class InterpolateDifference(Operation):
def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma, self.seed = alpha, beta, gamma, seed
def oper(self, a, b):
if a is None or b is None:
return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
alpha, delta = max(self.alpha, 0.001), torch.abs(a - b)
max_delta = torch.max(delta)
if max_delta == 0: return a
@@ -129,6 +336,65 @@ class InterpolateDifference(Operation):
mask = torch.lerp(torch.bernoulli(torch.clamp(diff, 0, 1), generator=rng), diff, self.gamma)
return a * (1 - mask) + b * mask
class ManualEnhancedInterpolateDifferenceOp(Operation):
def __init__(self, key, alpha, beta, gamma, delta, seed, *sources):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma, self.delta, self.seed = alpha, beta, gamma, delta, seed
def oper(self, a, b):
if a is None or b is None: return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta = torch.abs(a - b)
diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta))
mean_diff = torch.mean(diff, 0, keepdim=True)
mask = torch.logical_and(self.beta < mean_diff, mean_diff < self.gamma)
powered_diff = torch.nan_to_num(diff ** (1 / max(self.alpha, 0.001) - 1))
masked_diff = powered_diff * mask.float()
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(torch.clamp(masked_diff, 0, 1), generator=rng)
interpolated_mask = torch.lerp(random_mask, masked_diff, self.delta)
return a * (1 - interpolated_mask) + b * interpolated_mask
class AutoEnhancedInterpolateDifferenceOp(Operation):
def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma, self.seed = alpha, beta, gamma, seed
def oper(self, a, b):
if a is None or b is None: return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta = torch.abs(a - b)
max_delta = torch.max(delta)
diff = torch.nan_to_num((max_delta - delta) / max_delta)
mean_diff = torch.mean(diff)
lower_threshold = mean_diff * (1 - self.beta)
upper_threshold = mean_diff * (1 + self.beta)
mask = torch.logical_and(lower_threshold < diff, diff < upper_threshold)
powered_diff = torch.nan_to_num(diff ** (1 / max(self.alpha, 0.001) - 1))
masked_diff = powered_diff * mask.float()
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(torch.clamp(masked_diff, 0, 1), generator=rng)
interpolated_mask = torch.lerp(random_mask, masked_diff, self.gamma)
return a * (1 - interpolated_mask) + b * interpolated_mask
class WeightSumCutoffOp(Operation):
def __init__(self, key, alpha, beta, gamma, *sources):
super().__init__(key, *sources)
self.alpha, self.beta, self.gamma = alpha, beta, gamma
def oper(self, a, b):
if a is None or b is None: return None
a, b = resize_tensors(a, b, mode=self.alignment_mode)
delta = torch.abs(a - b)
diff = torch.nan_to_num((torch.max(delta) - delta) / torch.max(delta))
mean = torch.mean(diff, 0, True)
mask = torch.logical_and(mean < self.beta, self.gamma < mean)
mul = self.alpha * mask
return a * (1 - mul) + b * mul
# ################# #
# CALCULATION MODES
# ################# #
@@ -136,106 +402,313 @@ class CalcMode:
name = 'calcmode'; description = 'description'; models_used = []; param_docs = {}
def create_recipe(self, key, **kwargs) -> Operation: raise NotImplementedError
@staticmethod
def _get_secondary_loader_args(kwargs, mismatch_mode):
"""Helper to build loader args for secondary models (B, C) with fallback support."""
base_args = {
"handlers": kwargs['handlers'],
"device": kwargs['device'],
"dtype": kwargs['dtype'],
"on_missing": mismatch_mode,
}
# Add fallback info for zeros mode
if '_tensor_a_shape' in kwargs:
base_args['fallback_shape'] = kwargs['_tensor_a_shape']
base_args['fallback_dtype'] = kwargs.get('_tensor_a_dtype', kwargs['dtype'])
return base_args
class PreloadedTensor(Operation):
"""Returns a pre-loaded tensor instead of loading from file.
Used when Model A's tensor is pre-loaded in the merge loop.
"""
def __init__(self, key, tensor):
super().__init__(key)
self.tensor = tensor
def merge(self) -> torch.Tensor:
return self.tensor
def clean(self):
if hasattr(self, 'tensor'):
del self.tensor
class WeightSum(CalcMode):
name = 'Weight-Sum'
description = 'A * (1 - α) + B * α. Simple linear interpolation.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Interpolation weight. 0.0 is 100% Model A, 1.0 is 100% Model B.'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
# Use pre-loaded tensor_a if available
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
if kwargs['alpha'] >= 1: return b
if kwargs['alpha'] <= 0: return a
c = Multiply(key, 1 - kwargs['alpha'], a)
d = Multiply(key, kwargs['alpha'], b)
return Add(key, c, d)
res = Add(key, c, d)
res.alignment_mode = alignment_mode
return res
class AddDifference(CalcMode):
name = 'Add-Difference'
description = 'A + (B - C) * α. Applies the difference between B and C to A.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Multiplier for the (B - C) difference.', 'beta': 'Smoothing (0=Off, 1=On).'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
c = LoadTensor(key, kwargs['model_c'], **loader_args)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
diff = Sub(key, b, c)
diff.alignment_mode = alignment_mode
if kwargs['beta'] == 1: diff = Smooth(key, diff)
diffm = Multiply(key, kwargs['alpha'], diff)
return Add(key, a, diffm)
res = Add(key, a, diffm)
res.alignment_mode = alignment_mode
return res
class TrainDifference(CalcMode):
name = 'Train-Difference'
description = 'A + (B - C) * α, using tensor distances.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Multiplier for the calculated difference.'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
c = LoadTensor(key, kwargs['model_c'], **loader_args)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
diff = TrainDiff(key, a, b, c)
diff.alignment_mode = alignment_mode
diffm = Multiply(key, kwargs['alpha'], diff)
return Add(key, a, diffm)
res = Add(key, a, diffm)
res.alignment_mode = alignment_mode
return res
class InterpDifference(CalcMode):
name = 'Comparative-Interpolation'
description = 'Interpolates A and B based on relative tensor value differences.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Curve strength.', 'beta': 'Style (0=Similarity, 1=Difference).', 'gamma': 'Mix (0=Random, 1=Linear).'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
if key.startswith('cond_stage_model.transformer.text_model.embeddings'): return LoadTensor(key, kwargs['model_a'], **loader_args)
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
return InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
# Skip certain keys entirely
if key.startswith('cond_stage_model.transformer.text_model.embeddings'):
return a
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = InterpolateDifference(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
class Extract(CalcMode):
name = 'Extract-Features'
description = 'Adds features from B and C to A, based on their similarity.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Similarity vs dissimilarity.', 'gamma': 'Similarity bias.', 'delta': 'Multiplier for final features.'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
c = LoadTensor(key, kwargs['model_c'], **loader_args)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
extracted = ExtractOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'] * 15, a, b, c)
extracted.alignment_mode = alignment_mode
multiplied = Multiply(key, kwargs['delta'], extracted)
return Add(key, a, multiplied)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class AddDisimilarity(CalcMode):
name = 'Add-Dissimilarities'
description = 'Adds dissimilar features between B and C to model A.'
models_used = ['A', 'B', 'C']
param_docs = {'alpha': 'Weighting B vs C.', 'beta': 'Multiplier for final features.', 'gamma': 'Similarity bias.'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
c = LoadTensor(key, kwargs['model_c'], **loader_args)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_bc = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_bc)
c = LoadTensor(key, kwargs['model_c'], **loader_args_bc)
extracted = Similarities(key, kwargs['alpha'], 1, kwargs['gamma'] * 15, b, c)
extracted.alignment_mode = alignment_mode
multiplied = Multiply(key, kwargs['beta'], extracted)
return Add(key, a, multiplied)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class PowerUp(CalcMode):
name = 'Power-Up (DARE)'
description = 'Adds capabilities of B to A using Drop and Rescale (DARE).'
models_used = ['A', 'B']
param_docs = {'alpha': 'Dropout rate (proportion of weights from B to *drop*).', 'beta': 'Multiplier for the final difference.'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
deltahat = PowerUpOp(key, kwargs['alpha'], kwargs['seed'], a, b)
res = Multiply(key, kwargs['beta'], deltahat)
return Add(key, a, res)
class SVD(CalcMode):
deltahat.alignment_mode = alignment_mode # Important!
multiplied = Multiply(key, kwargs['beta'], deltahat)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class SVDMode(CalcMode):
name = 'SVD LoRA Extraction'
description = 'Extracts the difference between B and A into a LoRA file.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Rank for standard layers.', 'beta': 'Rank for 3x3 conv layers.', 'gamma': 'Clamp quantile for weights.'}
def create_recipe(self, key, **kwargs):
loader_args = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args)
b = LoadTensor(key, kwargs['model_b'], **loader_args)
return SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b)
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVD()]
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = SVD(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b, target_shape=kwargs.get('_tensor_a_shape'))
res.alignment_mode = alignment_mode
return res
class ManEnhInterpDifference(CalcMode):
name = 'Enhanced Man Interp'
description = 'Enhanced interpolation between each pair of values from A and B depending on their difference relative to other values.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Interpolation strength.', 'beta': 'Lower mean threshold.', 'gamma': 'Upper mean threshold.', 'delta': 'Smoothness factor.'}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
if key.startswith('cond_stage_model.transformer.text_model.embeddings'):
return a
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = ManualEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
class AutoEnhInterpDifference(CalcMode):
name = 'Enhanced Auto Interp'
description = 'Interpolates between each pair of values from A and B depending on their difference relative to other values.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Interpolation strength.', 'beta': 'Threshold adjustment factor.', 'gamma': 'Smoothness factor.'}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
if key.startswith('cond_stage_model.transformer.text_model.embeddings'):
return a
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = AutoEnhancedInterpolateDifferenceOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['seed'], a, b)
res.alignment_mode = alignment_mode
return res
class WeightSumCutoffMode(CalcMode):
name = 'Weight-Sum Cutoff'
description = 'Weight-sum with cutoff based on value differences.'
models_used = ['A', 'B']
param_docs = {'alpha': 'Interpolation weight.', 'beta': 'Upper threshold.', 'gamma': 'Lower threshold.'}
def create_recipe(self, key, **kwargs):
mismatch_mode = kwargs.get('mismatch_mode', MissingTensorBehavior.SKIP)
alignment_mode = kwargs.get('alignment_mode', 'pad/crop')
if '_tensor_a' in kwargs:
a = PreloadedTensor(key, kwargs['_tensor_a'])
else:
loader_args_a = {"handlers": kwargs['handlers'], "device": kwargs['device'], "dtype": kwargs['dtype']}
a = LoadTensor(key, kwargs['model_a'], **loader_args_a)
loader_args_b = self._get_secondary_loader_args(kwargs, mismatch_mode)
b = LoadTensor(key, kwargs['model_b'], **loader_args_b)
res = WeightSumCutoffOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], a, b)
res.alignment_mode = alignment_mode
return res
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()]
THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
+235 -7
View File
@@ -1,9 +1,78 @@
"""
Memory-efficient safetensors utilities with pinned memory support.
Forces direct file reads to minimize OS page cache usage and avoids VRAM accumulation.
"""
import os
import gc
import json
import struct
import torch
from safetensors import safe_open
from typing import Dict, List, Optional, Any
# Module-level configuration for pinned transfer
_verbose = False
_pinned_transfer_stats = {"pinned": 0, "fallback": 0}
def set_verbose(enabled: bool):
"""Enable/disable verbose output for pinned transfers."""
global _verbose
_verbose = enabled
def get_pinned_transfer_stats():
"""Return pinned transfer statistics."""
return _pinned_transfer_stats.copy()
def reset_pinned_transfer_stats():
"""Reset transfer statistics."""
global _pinned_transfer_stats
_pinned_transfer_stats = {"pinned": 0, "fallback": 0}
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."""
global _pinned_transfer_stats
if tensor.device.type != 'cpu' or not torch.cuda.is_available():
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
if not str(device).startswith('cuda'):
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
try:
pinned = tensor.pin_memory()
if dtype is not None:
result = pinned.to(device=device, dtype=dtype, non_blocking=True)
else:
result = pinned.to(device=device, non_blocking=True)
torch.cuda.current_stream().synchronize()
_pinned_transfer_stats["pinned"] += 1
return result
except Exception as e:
_pinned_transfer_stats["fallback"] += 1
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
class MemoryEfficientSafeOpen:
def __init__(self, filename, device='cpu'):
"""Memory-efficient safetensors file reader with streaming and pinned memory support.
Operates strictly in streaming mode (low memory). Tensors are read byte-by-byte from
their exact offsets without using safe_open or memory mapping, ensuring tensors do not
remain persistently in OS page caches or RAM.
"""
def __init__(self, filename: str, device: str = 'cpu', **kwargs):
# Ignore mmap_mode, low_memory, or any other legacy kwargs. We always stream.
self.filename = filename
self.device = device
self.header, self.header_size = self._read_header()
@@ -15,23 +84,118 @@ class MemoryEfficientSafeOpen:
def __exit__(self, exc_type, exc_val, exc_tb):
self.file.close()
def keys(self):
def keys(self) -> List[str]:
"""Return all tensor keys (excluding metadata)."""
return [k for k in self.header.keys() if k != "__metadata__"]
def get_tensor(self, key):
def metadata(self) -> Dict[str, str]:
"""Return file metadata."""
return self.header.get("__metadata__", {})
def get_shape(self, key: str) -> tuple:
"""Get tensor shape without loading tensor data."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in file")
return tuple(self.header[key]["shape"])
def get_ndim(self, key: str) -> int:
"""Get tensor ndim without loading tensor data."""
return len(self.get_shape(key))
def keys_sorted_by_offset(self) -> List[str]:
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 directly from file byte offset."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
tensor_bytes = bytearray(offset_end - offset_start)
self.file.readinto(tensor_bytes)
else:
tensor_bytes = None
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:
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]:
if keys is None:
keys = self.keys()
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:
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]:
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 {}
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
return json.loads(tensor_bytes.decode("utf-8"))
def _read_header(self):
with open(self.filename, "rb") as f:
header_size = struct.unpack("<Q", f.read(8))[0]
@@ -46,7 +210,7 @@ class MemoryEfficientSafeOpen:
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
return self._convert_float8(byte_tensor, dtype_str, shape)
@@ -77,4 +241,68 @@ class MemoryEfficientSafeOpen:
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
else:
raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")
raise ValueError(f"Unsupported float8 type: {dtype_str}")
def _read_tensor_from_file(
filepath: str,
key: str,
metadata: Dict,
header_size: int
) -> Optional[torch.Tensor]:
"""Helper function to read a single tensor from file (for parallel execution).
Each worker opens its own file handle for thread safety.
"""
offset_start, offset_end = metadata["data_offsets"]
dtype_str = metadata["dtype"]
shape = metadata["shape"]
if offset_start == offset_end:
return None
dtype_map = {
"F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
"I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
"U8": torch.uint8, "BOOL": torch.bool,
}
if hasattr(torch, "float8_e5m2"):
dtype_map["F8_E5M2"] = torch.float8_e5m2
if hasattr(torch, "float8_e4m3fn"):
dtype_map["F8_E4M3"] = torch.float8_e4m3fn
dtype = dtype_map.get(dtype_str, torch.float32)
with open(filepath, "rb") as f:
f.seek(header_size + 8 + offset_start)
tensor_bytes = f.read(offset_end - offset_start)
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
return byte_tensor.view(torch.float8_e5m2).reshape(shape)
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
return byte_tensor.view(dtype).reshape(shape)
def get_union_keys(*handlers) -> List[str]:
"""Returns the union of all tensor keys across multiple model handlers."""
all_keys = set()
for handler in handlers:
if handler is not None:
all_keys.update(handler.keys())
return sorted(all_keys)
def get_intersection_keys(*handlers) -> List[str]:
"""Returns only keys present in ALL model handlers."""
handlers = [h for h in handlers if h is not None]
if not handlers:
return []
all_keys = set(handlers[0].keys())
for handler in handlers[1:]:
all_keys &= set(handler.keys())
return sorted(all_keys)
+151 -75
View File
@@ -1,98 +1,174 @@
import os
import json
import struct
import folder_paths
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors import safe_open
from .utils import convert_pt_to_safetensors
def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]:
"""Shared logic for all MetaKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
class BaseMetaKeys:
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("metadata", "keys")
FUNCTION = "get_metakeys"
CATEGORY = "ModelUtils/Keys"
DESCRIPTION = "Loads a model and returns its metadata and layer names."
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
return f"Conversion failed: {error_message}", ""
model_path = temp_safe_path
def _get_metakeys(self, model_name, model_type):
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
model_weights = {}
metadata = {}
layer_shapes = ""
# 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)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
return f"Conversion failed: {error_message}", ""
model_path = temp_safe_path
# Extract metadata and layer shapes from header
metadata = header.pop("__metadata__", {})
metadata_str = str(metadata)
model_weights, metadata = load_metadata_from_safetensors(model_path)
metadata_str = str(metadata)
for layer_name, tensor in model_weights.items():
shape = str(tensor.shape)
shape = shape.replace("torch.Size(", "").replace(")", "")
name_shape = f"{layer_name}, {shape}\n"
layer_shapes += name_shape
return metadata_str, layer_shapes
layer_shapes = ""
keys = [k for k in header.keys()]
pbar = comfy.utils.ProgressBar(len(keys))
for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
info = header[layer_name]
shape = str(info.get("shape", []))
name_shape = f"{layer_name}, {shape}\n"
layer_shapes += name_shape
pbar.update(1)
return metadata_str, layer_shapes
class ModelMetaKeys(BaseMetaKeys):
class ModelMetaKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"diffusionmodel_name": (folder_paths.get_filename_list("diffusion_models"), ),
}
}
def define_schema(cls):
return io.Schema(
node_id="ModelMetaKeys",
display_name="Get Diffusion Model Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a diffusion model and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
def get_metakeys(self, diffusionmodel_name):
return self._get_metakeys(diffusionmodel_name, "diffusion_models")
class TextEncoderMetaKeys(BaseMetaKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"textencoder_name": (folder_paths.get_filename_list("text_encoders"), ),
}
}
def execute(cls, diffusionmodel_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(diffusionmodel_name, "diffusion_models")
return io.NodeOutput(metadata, keys)
def get_metakeys(self, textencoder_name):
return self._get_metakeys(textencoder_name, "text_encoders")
class LoRAMetaKeys(BaseMetaKeys):
class TextEncoderMetaKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"lora_name": (folder_paths.get_filename_list("loras"), ),
}
}
def get_metakeys(self, lora_name):
return self._get_metakeys(lora_name, "loras")
def define_schema(cls):
return io.Schema(
node_id="TextEncoderMetaKeys",
display_name="Get Text Encoder Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a text encoder and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
class CheckpointMetaKeys(BaseMetaKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
}
}
def execute(cls, textencoder_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(textencoder_name, "text_encoders")
return io.NodeOutput(metadata, keys)
def get_metakeys(self, ckpt_name):
return self._get_metakeys(ckpt_name, "checkpoints")
class EmbeddingMetaKeys(BaseMetaKeys):
class LoRAMetaKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"embedding": (folder_paths.get_filename_list("embeddings"), ),
}
}
def define_schema(cls):
return io.Schema(
node_id="LoRAMetaKeys",
display_name="Get LoRA Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a LoRA and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
def get_metakeys(self, embedding):
return self._get_metakeys(embedding, "embeddings")
@classmethod
def execute(cls, lora_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(lora_name, "loras")
return io.NodeOutput(metadata, keys)
__all__ = ("ModelMetaKeys", "TextEncoderMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", "EmbeddingMetaKeys")
class CheckpointMetaKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointMetaKeys",
display_name="Get Checkpoint Metadata & Keys",
category="ModelUtils/Keys",
description="Loads a checkpoint and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
@classmethod
def execute(cls, ckpt_name: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(ckpt_name, "checkpoints")
return io.NodeOutput(metadata, keys)
class EmbeddingMetaKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingMetaKeys",
display_name="Get Embedding Metadata & Keys",
category="ModelUtils/Keys",
description="Loads an embedding and returns its metadata and layer names.",
inputs=[
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
),
],
outputs=[
io.String.Output(display_name="metadata"),
io.String.Output(display_name="keys"),
],
)
@classmethod
def execute(cls, embedding: str) -> io.NodeOutput:
metadata, keys = _get_metakeys(embedding, "embeddings")
return io.NodeOutput(metadata, keys)
+184
View File
@@ -0,0 +1,184 @@
import os
import folder_paths
from comfy_api.latest import io
from .downloader_utils import get_potential_preview_files, load_image_tensor, get_model_workflows, get_model_metadata_file
class BaseModelInfoLoader(io.ComfyNode):
CATEGORY = "ModelUtils/Info"
@classmethod
def get_info(cls, category, name, workflow_index):
full_path = folder_paths.get_full_path(category, name)
if not full_path or not os.path.exists(full_path):
return (None, "{}", "{}")
base_path = os.path.splitext(full_path)[0]
# 1. Image Preview
image_tensor = None
previews = get_potential_preview_files(base_path)
for p in previews:
if os.path.exists(p) and not p.lower().endswith('.mp4'):
image_tensor = load_image_tensor(p)
if image_tensor is not None:
break
# 2. Workflow
workflow_json = "{}"
workflows = get_model_workflows(base_path)
if workflows and 0 <= workflow_index < len(workflows):
try:
with open(workflows[workflow_index], 'r', encoding='utf-8') as f:
workflow_json = f.read()
except Exception as e:
print(f"Error reading workflow {workflows[workflow_index]}: {e}")
# 3. Metadata
metadata_json = "{}"
metadata_file = get_model_metadata_file(base_path)
if metadata_file:
try:
with open(metadata_file, 'r', encoding='utf-8') as f:
metadata_json = f.read()
except Exception as e:
print(f"Error reading metadata {metadata_file}: {e}")
return (image_tensor, workflow_json, metadata_json)
class CheckpointInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Checkpoint Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("checkpoint", options=folder_paths.get_filename_list("checkpoints")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, checkpoint, workflow_index):
res = cls.get_info("checkpoints", checkpoint, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class LoRAInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="LoRA Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("lora", options=folder_paths.get_filename_list("loras")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, lora, workflow_index):
res = cls.get_info("loras", lora, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class EmbeddingInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Embedding Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, embedding, workflow_index):
res = cls.get_info("embeddings", embedding, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class VAEInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="VAE Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("vae", options=folder_paths.get_filename_list("vae")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, vae, workflow_index):
res = cls.get_info("vae", vae, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class ControlNetInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="ControlNet Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("controlnet", options=folder_paths.get_filename_list("controlnet")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, controlnet, workflow_index):
res = cls.get_info("controlnet", controlnet, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class DiffusionModelInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Diffusion Model Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("diffusion_model", options=folder_paths.get_filename_list("diffusion_models")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, diffusion_model, workflow_index):
res = cls.get_info("diffusion_models", diffusion_model, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
+184 -91
View File
@@ -1,38 +1,50 @@
import os
import re
import folder_paths
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
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
class BasePruneKeys:
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("output_path",)
FUNCTION = "prune_keys"
CATEGORY = "ModelUtils/Keys"
DESCRIPTION = "Loads a model, removes specified keys, and saves it as a new safetensors file."
def _prune_keys(self, model_name, model_type, keys_to_prune_str, use_regex, output_filename):
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
use_regex: bool, output_filename: str) -> str:
"""Shared logic for all PruneKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
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)
patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
# 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 {}
if not patterns:
raise ValueError("No keys/patterns provided to prune.")
patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
pruned_tensors = {}
for key, tensor in model_weights.items():
if not patterns:
raise ValueError("No keys/patterns provided to prune.")
# Stream tensors and filter on the fly
pruned_tensors = {}
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
all_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(all_keys))
for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = False
if use_regex:
if any(re.search(pattern, key) for pattern in patterns):
@@ -42,86 +54,167 @@ class BasePruneKeys:
is_match = True
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]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(pruned_tensors, output_path, metadata)
save_file(pruned_tensors, output_path, metadata)
return (output_path,)
# Cleanup after operation
del pruned_tensors
cleanup_after_operation()
class ModelPruneKeys(BasePruneKeys):
return output_path
class ModelPruneKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"diffusionmodel_name": (folder_paths.get_filename_list("diffusion_models"), ),
"keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
"use_regex": ("BOOLEAN", {"default": False}),
"output_filename": ("STRING", {"default": "pruned_model"}),
}
}
def define_schema(cls):
return io.Schema(
node_id="ModelPruneKeys",
display_name="Prune Diffusion Model Keys",
category="ModelUtils/Keys",
description="Loads a diffusion model, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_model"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
def prune_keys(self, diffusionmodel_name, keys_to_prune, use_regex, output_filename):
return self._prune_keys(diffusionmodel_name, "diffusion_models", keys_to_prune, use_regex, output_filename)
class TextEncoderPruneKeys(BasePruneKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"textencoder_name": (folder_paths.get_filename_list("text_encoders"), ),
"keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
"use_regex": ("BOOLEAN", {"default": False}),
"output_filename": ("STRING", {"default": "pruned_textencoder"}),
}
}
def execute(cls, diffusionmodel_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(diffusionmodel_name, "diffusion_models",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
def prune_keys(self, textencoder_name, keys_to_prune, use_regex, output_filename):
return self._prune_keys(textencoder_name, "text_encoders", keys_to_prune, use_regex, output_filename)
class LoRAPruneKeys(BasePruneKeys):
class TextEncoderPruneKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"lora_name": (folder_paths.get_filename_list("loras"), ),
"keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
"use_regex": ("BOOLEAN", {"default": False}),
"output_filename": ("STRING", {"default": "pruned_lora"}),
}
}
def define_schema(cls):
return io.Schema(
node_id="TextEncoderPruneKeys",
display_name="Prune Text Encoder Keys",
category="ModelUtils/Keys",
description="Loads a text encoder, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_textencoder"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
def prune_keys(self, lora_name, keys_to_prune, use_regex, output_filename):
return self._prune_keys(lora_name, "loras", keys_to_prune, use_regex, output_filename)
class CheckpointPruneKeys(BasePruneKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
"keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
"use_regex": ("BOOLEAN", {"default": False}),
"output_filename": ("STRING", {"default": "pruned_checkpoint"}),
}
}
def execute(cls, textencoder_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(textencoder_name, "text_encoders",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
def prune_keys(self, ckpt_name, keys_to_prune, use_regex, output_filename):
return self._prune_keys(ckpt_name, "checkpoints", keys_to_prune, use_regex, output_filename)
class EmbeddingPruneKeys(BasePruneKeys):
class LoRAPruneKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"embedding": (folder_paths.get_filename_list("embeddings"), ),
"keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
"use_regex": ("BOOLEAN", {"default": False}),
"output_filename": ("STRING", {"default": "pruned_embedding"}),
}
}
def define_schema(cls):
return io.Schema(
node_id="LoRAPruneKeys",
display_name="Prune LoRA Keys",
category="ModelUtils/Keys",
description="Loads a LoRA, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_lora"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
def prune_keys(self, embedding, keys_to_prune, use_regex, output_filename):
return self._prune_keys(embedding, "embeddings", keys_to_prune, use_regex, output_filename)
@classmethod
def execute(cls, lora_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(lora_name, "loras",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
class CheckpointPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointPruneKeys",
display_name="Prune Checkpoint Keys",
category="ModelUtils/Keys",
description="Loads a checkpoint, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_checkpoint"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, ckpt_name: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(ckpt_name, "checkpoints",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
class EmbeddingPruneKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingPruneKeys",
display_name="Prune Embedding Keys",
category="ModelUtils/Keys",
description="Loads an embedding, removes specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
),
io.String.Input("keys_to_prune", multiline=True, default=""),
io.Boolean.Input("use_regex", default=False),
io.String.Input("output_filename", default="pruned_embedding"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, embedding: str, keys_to_prune: str,
use_regex: bool, output_filename: str) -> io.NodeOutput:
path = _prune_keys(embedding, "embeddings",
keys_to_prune, use_regex, output_filename)
return io.NodeOutput(path)
+186 -94
View File
@@ -1,122 +1,214 @@
import os
import folder_paths
from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
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
class BaseRenameKeys:
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("output_path",)
FUNCTION = "rename_keys"
CATEGORY = "ModelUtils/Keys"
DESCRIPTION = "Loads a model, renames specified keys, and saves it as a new safetensors file."
def _rename_keys(self, model_name, model_type, old_keys_str, new_keys_str, output_filename):
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
new_keys_str: str, output_filename: str) -> str:
"""Shared logic for all RenameKeys nodes."""
model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
temp_safe_path = model_path + ".safetensors"
if not os.path.exists(temp_safe_path):
conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
if not conversion_successful:
raise Exception(f"Conversion failed: {error_message}")
model_path_to_load = temp_safe_path
else:
model_path_to_load = model_path
model_weights, metadata = load_metadata_from_safetensors(model_path_to_load)
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
old_keys = [key.strip() for key in old_keys_str.strip().split('\n') if key.strip()]
new_keys = [key.strip() for key in new_keys_str.strip().split('\n') if key.strip()]
# 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 {}
if len(old_keys) != len(new_keys):
raise ValueError("The number of old keys must match the number of new keys.")
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()]
key_map = dict(zip(old_keys, new_keys))
if len(old_keys) != len(new_keys):
raise ValueError("The number of old keys must match the number of new keys.")
original_keys = list(model_weights.keys())
renamed_tensors = {}
for key in original_keys:
key_map = dict(zip(old_keys, new_keys))
# Stream tensors and rename on the fly
renamed_tensors = {}
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
original_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(original_keys))
for key in tqdm(original_keys, desc="Renaming keys", unit="keys"):
new_key = key_map.get(key, key)
renamed_tensors[new_key] = model_weights[key]
renamed_tensors[new_key] = handler.get_tensor(key)
pbar.update(1)
model_dir = folder_paths.get_folder_paths(model_type)[0]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(renamed_tensors, output_path, metadata)
save_file(renamed_tensors, output_path, metadata)
return (output_path,)
# Cleanup after operation
del renamed_tensors
cleanup_after_operation()
class ModelRenameKeys(BaseRenameKeys):
return output_path
class ModelRenameKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"diffusionmodel_name": (folder_paths.get_filename_list("diffusion_models"), ),
"old_keys": ("STRING", {"multiline": True, "default": ""}),
"new_keys": ("STRING", {"multiline": True, "default": ""}),
"output_filename": ("STRING", {"default": "renamed_model"}),
}
}
def define_schema(cls):
return io.Schema(
node_id="ModelRenameKeys",
display_name="Rename Diffusion Model Keys",
category="ModelUtils/Keys",
description="Loads a diffusion model, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"diffusionmodel_name",
options=folder_paths.get_filename_list("diffusion_models"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_model"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
def rename_keys(self, diffusionmodel_name, old_keys, new_keys, output_filename):
return self._rename_keys(diffusionmodel_name, "diffusion_models", old_keys, new_keys, output_filename)
class TextEncoderRenameKeys(BaseRenameKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"textencoder_name": (folder_paths.get_filename_list("text_encoders"), ),
"old_keys": ("STRING", {"multiline": True, "default": ""}),
"new_keys": ("STRING", {"multiline": True, "default": ""}),
"output_filename": ("STRING", {"default": "renamed_textencoder"}),
}
}
def execute(cls, diffusionmodel_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(diffusionmodel_name, "diffusion_models",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
def rename_keys(self, textencoder_name, old_keys, new_keys, output_filename):
return self._rename_keys(textencoder_name, "text_encoders", old_keys, new_keys, output_filename)
class LoRARenameKeys(BaseRenameKeys):
class TextEncoderRenameKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"lora_name": (folder_paths.get_filename_list("loras"), ),
"old_keys": ("STRING", {"multiline": True, "default": ""}),
"new_keys": ("STRING", {"multiline": True, "default": ""}),
"output_filename": ("STRING", {"default": "renamed_lora"}),
}
}
def define_schema(cls):
return io.Schema(
node_id="TextEncoderRenameKeys",
display_name="Rename Text Encoder Keys",
category="ModelUtils/Keys",
description="Loads a text encoder, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"textencoder_name",
options=folder_paths.get_filename_list("text_encoders"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_textencoder"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
def rename_keys(self, lora_name, old_keys, new_keys, output_filename):
return self._rename_keys(lora_name, "loras", old_keys, new_keys, output_filename)
class CheckpointRenameKeys(BaseRenameKeys):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
"old_keys": ("STRING", {"multiline": True, "default": ""}),
"new_keys": ("STRING", {"multiline": True, "default": ""}),
"output_filename": ("STRING", {"default": "renamed_checkpoint"}),
}
}
def execute(cls, textencoder_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(textencoder_name, "text_encoders",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
def rename_keys(self, ckpt_name, old_keys, new_keys, output_filename):
return self._rename_keys(ckpt_name, "checkpoints", old_keys, new_keys, output_filename)
class EmbeddingRenameKeys(BaseRenameKeys):
class LoRARenameKeys(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"embedding": (folder_paths.get_filename_list("embeddings"), ),
"old_keys": ("STRING", {"multiline": True, "default": ""}),
"new_keys": ("STRING", {"multiline": True, "default": ""}),
"output_filename": ("STRING", {"default": "renamed_embedding"}),
}
}
def define_schema(cls):
return io.Schema(
node_id="LoRARenameKeys",
display_name="Rename LoRA Keys",
category="ModelUtils/Keys",
description="Loads a LoRA, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_lora"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
def rename_keys(self, embedding, old_keys, new_keys, output_filename):
return self._rename_keys(embedding, "embeddings", old_keys, new_keys, output_filename)
@classmethod
def execute(cls, lora_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(lora_name, "loras",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
class CheckpointRenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CheckpointRenameKeys",
display_name="Rename Checkpoint Keys",
category="ModelUtils/Keys",
description="Loads a checkpoint, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"ckpt_name",
options=folder_paths.get_filename_list("checkpoints"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_checkpoint"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, ckpt_name: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(ckpt_name, "checkpoints",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
class EmbeddingRenameKeys(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="EmbeddingRenameKeys",
display_name="Rename Embedding Keys",
category="ModelUtils/Keys",
description="Loads an embedding, renames specified keys, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"embedding",
options=folder_paths.get_filename_list("embeddings"),
),
io.String.Input("old_keys", multiline=True, default=""),
io.String.Input("new_keys", multiline=True, default=""),
io.String.Input("output_filename", default="renamed_embedding"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
)
@classmethod
def execute(cls, embedding: str, old_keys: str,
new_keys: str, output_filename: str) -> io.NodeOutput:
path = _rename_keys(embedding, "embeddings",
old_keys, new_keys, output_filename)
return io.NodeOutput(path)
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.1.6"
version = "0.2.9"
license = { file = "LICENSE" }
[project.urls]
+5
View File
@@ -0,0 +1,5 @@
requests
Pillow
mutagen
unifiedefficientloader
av