Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
61b2f506d0 | ||
|
|
75d1f0f5c9 | ||
|
|
3a71ad9952 | ||
|
|
37090dd7ba | ||
|
|
f58289d042 | ||
|
|
1796f1bea9 |
@@ -23,6 +23,10 @@ from .nodes.lora_extract_svd import (
|
||||
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius
|
||||
)
|
||||
from .nodes.lora_resize import (
|
||||
LoRAResizeFixed, LoRAResizeRatio,
|
||||
LoRAResizeFrobenius, LoRAResizeCumulative
|
||||
)
|
||||
|
||||
|
||||
class ModelUtilsExtension(ComfyExtension):
|
||||
@@ -47,6 +51,9 @@ class ModelUtilsExtension(ComfyExtension):
|
||||
# LoRA Extraction
|
||||
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
|
||||
LoRAExtractKnee, LoRAExtractFrobenius,
|
||||
# LoRA Resize
|
||||
LoRAResizeFixed, LoRAResizeRatio,
|
||||
LoRAResizeFrobenius, LoRAResizeCumulative,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
# ComfyUI-ModelUtils - Complete Project History
|
||||
|
||||
## Version Timeline
|
||||
|
||||
| Version | Commits | Key Features |
|
||||
|---------|---------|--------------|
|
||||
| 0.1.1 | `70d6c13`-`d67cbca` | Initial nodes (MetaKeys) |
|
||||
| 0.1.2 | `6b662a4` | Icon assets |
|
||||
| 0.1.4 | `cbe3502` | Key renaming/pruning nodes |
|
||||
| 0.1.5 | `78b3198` | Model merging nodes |
|
||||
| 0.1.6 | `31989be` | Embedding nodes, namespace fixes |
|
||||
| 0.1.7 | `b6fae8d` | Mismatch handling, layer filtering, V3 migration |
|
||||
| **0.2.0** | `3df073f` | LoRA extraction, memory optimization, pinned memory |
|
||||
|
||||
---
|
||||
|
||||
## Phase 1: Initial Release (v0.1.1)
|
||||
|
||||
**Commit:** `70d6c13` - "Add first nodes, toml and publish"
|
||||
|
||||
### Files Created
|
||||
- `nodes/metakeys.py` - Get model metadata and layer keys
|
||||
- `nodes/utils.py` - Shared utilities
|
||||
- `__init__.py` - Node registration
|
||||
- `pyproject.toml` - Package configuration
|
||||
|
||||
### Nodes
|
||||
- `ModelMetaKeys` - Get diffusion model metadata
|
||||
- `CheckpointMetaKeys` - Get checkpoint metadata
|
||||
- `TextEncoderMetaKeys` - Get text encoder metadata
|
||||
- `LoRAMetaKeys` - Get LoRA metadata
|
||||
- `EmbeddingMetaKeys` - Get embedding metadata
|
||||
|
||||
---
|
||||
|
||||
## Phase 2: Key Manipulation (v0.1.4)
|
||||
|
||||
**Commit:** `4f08fe6` - "add key renaming and key pruning nodes"
|
||||
|
||||
### Files Created
|
||||
- `nodes/renamekeys.py` - Rename tensor keys
|
||||
- `nodes/prunekeys.py` - Remove tensor keys by pattern
|
||||
|
||||
### Nodes (5 variants each for model types)
|
||||
- `ModelRenameKeys`, `CheckpointRenameKeys`, `TextEncoderRenameKeys`, `LoRARenameKeys`, `EmbeddingRenameKeys`
|
||||
- `ModelPruneKeys`, `CheckpointPruneKeys`, `TextEncoderPruneKeys`, `LoRAPruneKeys`, `EmbeddingPruneKeys`
|
||||
|
||||
---
|
||||
|
||||
## Phase 3: Model Merging (v0.1.5)
|
||||
|
||||
**Commit:** `7959370` - "Add model merging nodes"
|
||||
|
||||
### Files Created
|
||||
- `nodes/merger.py` - Main merging logic
|
||||
- `nodes/merger_ops.py` - Merge operation modes
|
||||
- `nodes/merger_utils.py` - `MemoryEfficientSafeOpen` class
|
||||
- `docs/merger_2_model_modes.md` - Documentation
|
||||
- `docs/merger_3_model_modes.md` - Documentation
|
||||
|
||||
### Nodes
|
||||
- `CheckpointTwoMerger`, `DiffusionModelTwoMerger`, `TextEncoderTwoMerger`, `LoRATwoMerger`
|
||||
- `CheckpointThreeMerger`, `DiffusionModelThreeMerger`, `TextEncoderThreeMerger`, `LoRAThreeMerger`
|
||||
|
||||
### Merge Modes (TWO_MODEL_MODES)
|
||||
- Weighted Sum, Add Difference, Smooth Add Difference
|
||||
- Train Difference, Sum Twice, Triple Sum
|
||||
- Euclidean Add Difference, Distribution Crossover, Ties
|
||||
- Slerp, Lerp, SVD Supermerger, SVD LoRA Extraction
|
||||
|
||||
### Merge Modes (THREE_MODEL_MODES)
|
||||
- Add Difference (A + (B - C) × α)
|
||||
|
||||
### MemoryEfficientSafeOpen (initial)
|
||||
```python
|
||||
class MemoryEfficientSafeOpen:
|
||||
def __init__(self, filename, device='cpu', mmap_mode=True)
|
||||
def keys() -> List[str]
|
||||
def get_tensor(key) -> Tensor
|
||||
# mmap for zero-copy reads
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 4: Embedding Nodes (v0.1.6)
|
||||
|
||||
**Commit:** `0d035ef` - "Add embedding nodes"
|
||||
|
||||
### Additions
|
||||
- Added embedding variants to all node types
|
||||
- Namespace conflict fixes (`68f22fa`, `5ea63ab`)
|
||||
- Refactored names for universality (`232dc76`)
|
||||
|
||||
---
|
||||
|
||||
## Phase 5: Mismatch Handling (v0.1.7)
|
||||
|
||||
**Commits:** `19bc047` → `de27186`
|
||||
|
||||
### Features Added
|
||||
- `mismatch_mode` parameter: skip/zeros/error
|
||||
- Layer filtering with regex patterns
|
||||
- `exclude_patterns` - Keep model A only (no merge)
|
||||
- `discard_patterns` - Remove from output entirely
|
||||
|
||||
### Files Modified
|
||||
- `merger.py` - Added mismatch handling logic
|
||||
- `merger_ops.py` - Added `MissingTensorBehavior` enum, `MissingTensorError`
|
||||
|
||||
---
|
||||
|
||||
## Phase 6: V3 API Migration (v0.1.7)
|
||||
|
||||
**Commit:** `a3116cb` - "migrate to v3 schema and add LoRA extraction node"
|
||||
|
||||
### Changes
|
||||
- Migrated all nodes from legacy API to `comfy_api.latest.io`
|
||||
- `io.Schema`, `io.ComfyNode`, `io.NodeOutput`
|
||||
- Added initial LoRA extraction node (diff-based)
|
||||
|
||||
---
|
||||
|
||||
## Phase 7: LyCORIS Extraction (pre-0.2.0)
|
||||
|
||||
**Commits:** `92d3c55` → `18679cc`
|
||||
|
||||
### Files Created
|
||||
- `nodes/lycoris_extract.py` - LyCORIS-based extraction
|
||||
- `nodes/lora_extract.py` - Simple diff extraction
|
||||
|
||||
### Initial Implementation
|
||||
- Used `lycoris-lora` external dependency
|
||||
- `lycoris.utils.extract_linear`, `extract_conv`
|
||||
- Added `requirements.txt` with `lycoris-lora`
|
||||
|
||||
### Nodes Split (`18679cc`)
|
||||
- `LyCORISExtractFixed`, `LyCORISExtractThreshold`
|
||||
- `LyCORISExtractRatio`, `LyCORISExtractQuantile`
|
||||
|
||||
---
|
||||
|
||||
## Phase 8: Native SVD Port (0.2.0)
|
||||
|
||||
**Commits:** `6d7b2d1` → `eef7377`
|
||||
|
||||
### Removed
|
||||
- `lycoris-lora` dependency
|
||||
- `lycoris_extract.py`
|
||||
- `lora_extract.py`
|
||||
|
||||
### NEW: [lora_extract_svd.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/lora_extract_svd.py)
|
||||
|
||||
**8 Rank Selection Functions:**
|
||||
```python
|
||||
_index_sv_fixed(S, dim) # Fixed rank
|
||||
_index_sv_threshold(S, threshold) # Absolute threshold
|
||||
_index_sv_ratio(S, ratio) # Relative to max
|
||||
_index_sv_cumulative(S, target) # Cumulative sum %
|
||||
_index_sv_fro(S, target) # Frobenius norm
|
||||
_index_sv_knee(S) # Knee detection
|
||||
_index_sv_cumulative_knee(S) # Cumsum knee
|
||||
_index_sv_rel_decrease(S, tau) # Ratio drop
|
||||
```
|
||||
|
||||
**SVD Functions:**
|
||||
```python
|
||||
_svd_extract_linear_lowrank(weight, rank, device, clamp_quantile)
|
||||
_svd_extract_linear(weight, mode, mode_param, device, max_rank, clamp_quantile)
|
||||
_svd_extract_conv(weight, mode, mode_param, device, max_rank, clamp_quantile)
|
||||
_detect_fused_layer(key, shape) # QKV/MLP detection
|
||||
_extract_chunked_layer(...) # Chunked extraction
|
||||
```
|
||||
|
||||
**Nodes:**
|
||||
- `LoRAExtractFixed` - Uses `svd_lowrank` (10x faster)
|
||||
- `LoRAExtractRatio` - Relative threshold
|
||||
- `LoRAExtractQuantile` - Cumulative target
|
||||
- `LoRAExtractKnee` - Auto knee detection
|
||||
- `LoRAExtractFrobenius` - Norm preservation
|
||||
|
||||
---
|
||||
|
||||
## Phase 9: Memory Management (0.2.0)
|
||||
|
||||
**Commit:** `6d7b2d1` - "refactor nodes and implement memory management"
|
||||
|
||||
### NEW: [device_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/device_utils.py)
|
||||
|
||||
Integrates with `comfy.model_management`:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class DeviceCapabilities:
|
||||
device_type, total_vram_gb, free_vram_gb, total_ram_gb, free_ram_gb
|
||||
cpu_cores, recommended_workers, max_tensor_size_gb
|
||||
supports_pinned_memory, supports_async_streams
|
||||
|
||||
get_device_capabilities(device) -> DeviceCapabilities
|
||||
estimate_model_size(path) -> float # GB
|
||||
can_fit_in_memory(paths, device, safety_factor) -> (bool, str)
|
||||
prepare_for_large_operation(gb, device) -> bool # Calls free_memory()
|
||||
cleanup_after_operation() # gc.collect() + soft_empty_cache()
|
||||
get_optimal_workers(size_gb) -> int
|
||||
```
|
||||
|
||||
### Updates to Existing Nodes
|
||||
- `metakeys.py` - Header-only reading (no tensor loading)
|
||||
- `renamekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
|
||||
- `prunekeys.py` - `MemoryEfficientSafeOpen` + memory mgmt
|
||||
- `merger.py` - Memory prep/cleanup
|
||||
|
||||
---
|
||||
|
||||
## Phase 10: Pinned Memory (0.2.0)
|
||||
|
||||
**Commit:** `b770786` - "Add pinned memory for faster merging and extraction"
|
||||
|
||||
### [merger_utils.py](file:///f:/ComfyUI/custom_nodes/ComfyUI-ModelUtils/nodes/merger_utils.py) Enhancements
|
||||
|
||||
```python
|
||||
# Context manager
|
||||
class PinnedTensor:
|
||||
def __enter__(self) -> Tensor # Pins memory
|
||||
def __exit__(...) # Unpins
|
||||
|
||||
# Transfer helper (2-3x faster)
|
||||
def transfer_to_gpu_pinned(tensor, device='cuda', dtype=None) -> Tensor
|
||||
|
||||
# MemoryEfficientSafeOpen additions
|
||||
def get_tensor_to_gpu(key, device, dtype) # Pinned transfer
|
||||
def keys_sorted_by_offset() -> List[str] # Sequential I/O
|
||||
def get_tensors_parallel(keys, workers, sort_by_offset) -> Dict # Parallel
|
||||
def get_tensor_as_dict(key) -> Dict # JSON config parsing
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 11: SVD Optimization (0.2.0)
|
||||
|
||||
**Commit:** `65c1ae3` - "add svd_lowrank for fixed rank extraction"
|
||||
|
||||
### Changes
|
||||
- Added `_svd_extract_linear_lowrank()` using `torch.svd_lowrank`
|
||||
- Routes fixed mode to lowrank for **10x speedup**
|
||||
- Changed `chunk_large_layers` default to `False`
|
||||
|
||||
---
|
||||
|
||||
## Phase 12: Bug Fixes (0.2.0)
|
||||
|
||||
**Commit:** `eef7377` - "Added .contiguous() to saves"
|
||||
|
||||
### Fix
|
||||
Non-contiguous tensor error from SVD slices:
|
||||
```python
|
||||
tensor.to(dtype).cpu().contiguous() # Added .contiguous()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Current File Structure
|
||||
|
||||
```
|
||||
ComfyUI-ModelUtils/
|
||||
├── __init__.py
|
||||
├── pyproject.toml
|
||||
├── requirements.txt (empty - no external deps)
|
||||
├── README.md
|
||||
├── docs/
|
||||
│ ├── merger_2_model_modes.md
|
||||
│ └── merger_3_model_modes.md
|
||||
├── nodes/
|
||||
│ ├── utils.py
|
||||
│ ├── metakeys.py
|
||||
│ ├── renamekeys.py
|
||||
│ ├── prunekeys.py
|
||||
│ ├── merger.py
|
||||
│ ├── merger_ops.py
|
||||
│ ├── merger_utils.py
|
||||
│ ├── device_utils.py
|
||||
│ └── lora_extract_svd.py
|
||||
└── assets/
|
||||
└── icon.png
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Performance Summary
|
||||
|
||||
| Feature | Improvement |
|
||||
|---------|-------------|
|
||||
| Fixed-rank SVD | **10x** (svd_lowrank) |
|
||||
| CPU→GPU transfer | **2-3x** (pinned memory) |
|
||||
| Metadata reading | **100x+** (header only) |
|
||||
| Parallel I/O | **2-4x** (ThreadPoolExecutor) |
|
||||
| Memory management | Auto-unload models before heavy ops |
|
||||
+48
-25
@@ -36,10 +36,10 @@ def _index_sv_threshold(S: torch.Tensor, threshold: float) -> int:
|
||||
|
||||
|
||||
def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int:
|
||||
"""Ratio mode - keep singular values > max(S) * ratio."""
|
||||
"""Ratio mode - keep singular values > max(S) / ratio (kohya convention)."""
|
||||
if ratio <= 0:
|
||||
return len(S) - 1
|
||||
min_sv = S[0] * ratio
|
||||
min_sv = S[0] / ratio
|
||||
rank = int(torch.sum(S > min_sv).item())
|
||||
return max(1, min(rank, len(S) - 1))
|
||||
|
||||
@@ -148,6 +148,7 @@ def _svd_extract_linear_lowrank(
|
||||
rank: int,
|
||||
device: str,
|
||||
clamp_quantile: float = 0.99,
|
||||
niter: int = 2,
|
||||
) -> tuple:
|
||||
"""
|
||||
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
|
||||
@@ -155,6 +156,9 @@ def _svd_extract_linear_lowrank(
|
||||
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")
|
||||
"""
|
||||
@@ -170,9 +174,8 @@ def _svd_extract_linear_lowrank(
|
||||
return weight, "full"
|
||||
|
||||
# Use svd_lowrank for faster low-rank approximation
|
||||
# niter controls accuracy (higher = more accurate but slower)
|
||||
try:
|
||||
U, S, V = torch.svd_lowrank(weight, q=rank, niter=2)
|
||||
U, 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:
|
||||
@@ -203,6 +206,7 @@ def _svd_extract_linear(
|
||||
device: str,
|
||||
max_rank: int = None,
|
||||
clamp_quantile: float = 0.99,
|
||||
niter: int = 2,
|
||||
) -> tuple:
|
||||
"""
|
||||
SVD decomposition for linear layers.
|
||||
@@ -221,7 +225,7 @@ def _svd_extract_linear(
|
||||
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)
|
||||
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
|
||||
|
||||
# For adaptive modes, use full SVD to analyze singular values
|
||||
try:
|
||||
@@ -409,6 +413,7 @@ def extract_lora_from_files(
|
||||
skip_patterns_str: str = "",
|
||||
mismatch_mode: str = "skip",
|
||||
chunk_large_layers: bool = True,
|
||||
svd_niter: int = 2,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""
|
||||
Extract LoRA from difference between two models.
|
||||
@@ -528,7 +533,7 @@ def extract_lora_from_files(
|
||||
)
|
||||
else:
|
||||
result, mode_str = _svd_extract_linear(
|
||||
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile
|
||||
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
|
||||
)
|
||||
except Exception as e:
|
||||
# Try chunked extraction for large tensors
|
||||
@@ -640,6 +645,8 @@ class LoRAExtractFixed(io.ComfyNode):
|
||||
tooltip="Rank for linear/attention layers"),
|
||||
io.Int.Input("conv_dim", default=32, min=1, max=4096,
|
||||
tooltip="Rank for conv layers"),
|
||||
io.Int.Input("svd_niter", default=2, min=0, max=10,
|
||||
tooltip="SVD power iterations (higher = more accurate but slower)"),
|
||||
*_get_common_inputs(),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
@@ -647,7 +654,7 @@ class LoRAExtractFixed(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_a, model_b, linear_dim, conv_dim, chunk_large_layers,
|
||||
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
|
||||
clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
save_dtype, device, skip_patterns) -> io.NodeOutput:
|
||||
|
||||
@@ -657,7 +664,7 @@ class LoRAExtractFixed(io.ComfyNode):
|
||||
output_sd = extract_lora_from_files(
|
||||
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
|
||||
"lora_unet_", device, save_dtype, linear_dim, conv_dim,
|
||||
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
|
||||
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter
|
||||
)
|
||||
|
||||
return io.NodeOutput(_save_lora(output_sd, output_filename))
|
||||
@@ -672,13 +679,17 @@ class LoRAExtractRatio(io.ComfyNode):
|
||||
node_id="LoRAExtractRatio",
|
||||
display_name="LoRA Extract (Ratio)",
|
||||
category="ModelUtils/LoRA",
|
||||
description="Keep singular values > max(S) × ratio.",
|
||||
description="Keep singular values > max(S) / ratio.",
|
||||
inputs=[
|
||||
*_get_model_inputs(),
|
||||
io.Float.Input("linear_ratio", default=0.5, min=0.0, max=1.0, step=0.01,
|
||||
tooltip="Ratio threshold for linear layers"),
|
||||
io.Float.Input("conv_ratio", default=0.5, min=0.0, max=1.0, step=0.01,
|
||||
tooltip="Ratio threshold for conv layers"),
|
||||
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
|
||||
tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
|
||||
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
|
||||
tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
|
||||
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for linear layers"),
|
||||
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for conv layers"),
|
||||
*_get_common_inputs(),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
@@ -686,8 +697,8 @@ class LoRAExtractRatio(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, chunk_large_layers,
|
||||
clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
|
||||
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
save_dtype, device, skip_patterns) -> io.NodeOutput:
|
||||
|
||||
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
|
||||
@@ -695,7 +706,7 @@ class LoRAExtractRatio(io.ComfyNode):
|
||||
|
||||
output_sd = extract_lora_from_files(
|
||||
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
|
||||
"lora_unet_", device, save_dtype, None, None,
|
||||
"lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank,
|
||||
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
|
||||
)
|
||||
|
||||
@@ -718,6 +729,10 @@ class LoRAExtractQuantile(io.ComfyNode):
|
||||
tooltip="Target cumulative % for linear layers"),
|
||||
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
|
||||
tooltip="Target cumulative % for conv layers"),
|
||||
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for linear layers"),
|
||||
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for conv layers"),
|
||||
*_get_common_inputs(),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
@@ -725,8 +740,8 @@ class LoRAExtractQuantile(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, chunk_large_layers,
|
||||
clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
|
||||
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
save_dtype, device, skip_patterns) -> io.NodeOutput:
|
||||
|
||||
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
|
||||
@@ -734,7 +749,7 @@ class LoRAExtractQuantile(io.ComfyNode):
|
||||
|
||||
output_sd = extract_lora_from_files(
|
||||
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
|
||||
"lora_unet_", device, save_dtype, None, None,
|
||||
"lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank,
|
||||
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
|
||||
)
|
||||
|
||||
@@ -755,6 +770,10 @@ class LoRAExtractKnee(io.ComfyNode):
|
||||
*_get_model_inputs(),
|
||||
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
|
||||
default="sv_knee", tooltip="Knee detection method"),
|
||||
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for linear layers"),
|
||||
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for conv layers"),
|
||||
*_get_common_inputs(),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
@@ -762,8 +781,8 @@ class LoRAExtractKnee(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_a, model_b, knee_method, chunk_large_layers,
|
||||
clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
|
||||
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
save_dtype, device, skip_patterns) -> io.NodeOutput:
|
||||
|
||||
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
|
||||
@@ -771,7 +790,7 @@ class LoRAExtractKnee(io.ComfyNode):
|
||||
|
||||
output_sd = extract_lora_from_files(
|
||||
model_a_path, model_b_path, knee_method, 0, 0,
|
||||
"lora_unet_", device, save_dtype, None, None,
|
||||
"lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank,
|
||||
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
|
||||
)
|
||||
|
||||
@@ -794,6 +813,10 @@ class LoRAExtractFrobenius(io.ComfyNode):
|
||||
tooltip="Target Frobenius norm fraction for linear"),
|
||||
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
|
||||
tooltip="Target Frobenius norm fraction for conv"),
|
||||
io.Int.Input("linear_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for linear layers"),
|
||||
io.Int.Input("conv_max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum rank for conv layers"),
|
||||
*_get_common_inputs(),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
@@ -801,8 +824,8 @@ class LoRAExtractFrobenius(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_a, model_b, linear_target, conv_target, chunk_large_layers,
|
||||
clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
|
||||
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
|
||||
save_dtype, device, skip_patterns) -> io.NodeOutput:
|
||||
|
||||
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
|
||||
@@ -810,7 +833,7 @@ class LoRAExtractFrobenius(io.ComfyNode):
|
||||
|
||||
output_sd = extract_lora_from_files(
|
||||
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
|
||||
"lora_unet_", device, save_dtype, None, None,
|
||||
"lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank,
|
||||
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,588 @@
|
||||
"""
|
||||
LoRA Resize - Resize existing LoRAs to different ranks.
|
||||
|
||||
Based on kohya_ss resize_lora.py. Merges LoRA weights and re-extracts via SVD.
|
||||
Supports fixed rank and dynamic methods (sv_ratio, sv_fro, sv_cumulative).
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
import torch.linalg as linalg
|
||||
import folder_paths
|
||||
import comfy.utils
|
||||
from tqdm import tqdm
|
||||
from comfy_api.latest import io
|
||||
from safetensors.torch import save_file
|
||||
from typing import Optional, Dict, Tuple, List
|
||||
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
|
||||
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
|
||||
|
||||
# Reuse rank functions from extraction module
|
||||
from .lora_extract_svd import (
|
||||
_index_sv_ratio,
|
||||
_index_sv_cumulative,
|
||||
_index_sv_fro,
|
||||
)
|
||||
|
||||
|
||||
MIN_SV = 1e-6
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# LoRA Format Detection - Support various LoRA sources
|
||||
# =============================================================================
|
||||
|
||||
def detect_lora_format(keys: List[str]) -> Dict:
|
||||
"""
|
||||
Detect LoRA format and key patterns from various sources.
|
||||
|
||||
Supports:
|
||||
- Kohya/A1111: lora_unet_*, lora_te_*
|
||||
- Diffusers: *.lora_A.weight, *.lora_B.weight
|
||||
- PEFT: base_model.model.*.lora_A/B
|
||||
- ComfyUI native: *.lora_down.weight, *.lora_up.weight
|
||||
- Flux/SDXL: Various transformer key patterns
|
||||
|
||||
Returns:
|
||||
Dict with format info and key patterns
|
||||
"""
|
||||
format_info = {
|
||||
"format": "unknown",
|
||||
"down_suffix": ".lora_down.weight",
|
||||
"up_suffix": ".lora_up.weight",
|
||||
"alpha_suffix": ".alpha",
|
||||
"has_alpha": False,
|
||||
"key_count": 0,
|
||||
}
|
||||
|
||||
# Check for alpha keys
|
||||
alpha_keys = [k for k in keys if '.alpha' in k or k.endswith('alpha')]
|
||||
format_info["has_alpha"] = len(alpha_keys) > 0
|
||||
|
||||
# Detect format by key patterns
|
||||
sample_keys = keys[:20] # Check first 20 keys
|
||||
|
||||
# Kohya/A1111 format: lora_unet_down_blocks_0_*.lora_down.weight
|
||||
if any('lora_unet_' in k or 'lora_te' in k for k in sample_keys):
|
||||
format_info["format"] = "kohya"
|
||||
format_info["down_suffix"] = ".lora_down.weight"
|
||||
format_info["up_suffix"] = ".lora_up.weight"
|
||||
|
||||
# Diffusers format: transformer.*.lora_A.weight
|
||||
elif any('.lora_A.weight' in k for k in sample_keys):
|
||||
format_info["format"] = "diffusers"
|
||||
format_info["down_suffix"] = ".lora_A.weight"
|
||||
format_info["up_suffix"] = ".lora_B.weight"
|
||||
format_info["alpha_suffix"] = ".alpha"
|
||||
|
||||
# PEFT format: base_model.model.*.lora_A.default.weight
|
||||
elif any('base_model.model.' in k and 'lora_' in k for k in sample_keys):
|
||||
format_info["format"] = "peft"
|
||||
# Handle nested structure
|
||||
format_info["down_suffix"] = ".lora_A.default.weight"
|
||||
format_info["up_suffix"] = ".lora_B.default.weight"
|
||||
|
||||
# ComfyUI native / standard
|
||||
elif any('.lora_down.weight' in k for k in sample_keys):
|
||||
format_info["format"] = "comfy"
|
||||
format_info["down_suffix"] = ".lora_down.weight"
|
||||
format_info["up_suffix"] = ".lora_up.weight"
|
||||
|
||||
# Count LoRA layers
|
||||
down_keys = [k for k in keys if format_info["down_suffix"] in k]
|
||||
format_info["key_count"] = len(down_keys)
|
||||
|
||||
return format_info
|
||||
|
||||
|
||||
def extract_lora_pairs(keys: List[str], format_info: Dict) -> Dict[str, Dict[str, str]]:
|
||||
"""
|
||||
Group LoRA keys into down/up/alpha pairs.
|
||||
|
||||
Returns:
|
||||
Dict[block_name, {"down": key, "up": key, "alpha": key}]
|
||||
"""
|
||||
down_suffix = format_info["down_suffix"]
|
||||
up_suffix = format_info["up_suffix"]
|
||||
alpha_suffix = format_info["alpha_suffix"]
|
||||
|
||||
pairs = {}
|
||||
|
||||
for key in keys:
|
||||
if down_suffix in key:
|
||||
block_name = key.replace(down_suffix, "")
|
||||
if block_name not in pairs:
|
||||
pairs[block_name] = {}
|
||||
pairs[block_name]["down"] = key
|
||||
elif up_suffix in key:
|
||||
block_name = key.replace(up_suffix, "")
|
||||
if block_name not in pairs:
|
||||
pairs[block_name] = {}
|
||||
pairs[block_name]["up"] = key
|
||||
elif alpha_suffix in key or key.endswith('.alpha'):
|
||||
# Handle various alpha key formats
|
||||
block_name = key.replace(alpha_suffix, "").replace(".alpha", "")
|
||||
if block_name not in pairs:
|
||||
pairs[block_name] = {}
|
||||
pairs[block_name]["alpha"] = key
|
||||
|
||||
return pairs
|
||||
|
||||
|
||||
def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int, float]:
|
||||
"""
|
||||
Detect the rank and alpha of an existing LoRA.
|
||||
|
||||
Returns:
|
||||
(network_dim, network_alpha)
|
||||
"""
|
||||
network_dim = None
|
||||
network_alpha = None
|
||||
|
||||
for block_name, block_keys in pairs.items():
|
||||
if "down" not in block_keys:
|
||||
continue
|
||||
|
||||
# Get dim from down weight shape
|
||||
if network_dim is None:
|
||||
down_key = block_keys["down"]
|
||||
shape = handler.header[down_key]["shape"]
|
||||
# Linear: [rank, in_features] or Conv: [rank, in_ch, k, k]
|
||||
network_dim = shape[0]
|
||||
|
||||
# Get alpha if present
|
||||
if network_alpha is None and "alpha" in block_keys:
|
||||
alpha_tensor = handler.get_tensor(block_keys["alpha"])
|
||||
network_alpha = float(alpha_tensor.item())
|
||||
|
||||
if network_dim is not None and network_alpha is not None:
|
||||
break
|
||||
|
||||
# Default alpha to dim if not found
|
||||
if network_alpha is None:
|
||||
network_alpha = float(network_dim) if network_dim else 1.0
|
||||
if network_dim is None:
|
||||
network_dim = 1
|
||||
|
||||
return network_dim, network_alpha
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Merge Functions
|
||||
# =============================================================================
|
||||
|
||||
def _merge_linear(lora_down: torch.Tensor, lora_up: torch.Tensor, device: str) -> torch.Tensor:
|
||||
"""Merge linear LoRA weights: lora_up @ lora_down."""
|
||||
lora_down = lora_down.to(device=device, dtype=torch.float32)
|
||||
lora_up = lora_up.to(device=device, dtype=torch.float32)
|
||||
return lora_up @ lora_down
|
||||
|
||||
|
||||
def _merge_conv(lora_down: torch.Tensor, lora_up: torch.Tensor, device: str) -> torch.Tensor:
|
||||
"""Merge conv LoRA weights."""
|
||||
in_rank, in_size, kernel_size, k_ = lora_down.shape
|
||||
out_size, out_rank, _, _ = lora_up.shape
|
||||
assert in_rank == out_rank, f"rank mismatch: {in_rank} vs {out_rank}"
|
||||
|
||||
lora_down = lora_down.to(device=device, dtype=torch.float32)
|
||||
lora_up = lora_up.to(device=device, dtype=torch.float32)
|
||||
|
||||
merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
|
||||
return merged.reshape(out_size, in_size, kernel_size, kernel_size)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Extract Functions (after merge)
|
||||
# =============================================================================
|
||||
|
||||
def _extract_linear(
|
||||
weight: torch.Tensor,
|
||||
max_rank: int,
|
||||
dynamic_method: Optional[str],
|
||||
dynamic_param: Optional[float],
|
||||
scale: float,
|
||||
niter: int = 2,
|
||||
) -> Dict:
|
||||
"""Extract LoRA from merged linear weight."""
|
||||
out_size, in_size = weight.shape
|
||||
|
||||
if dynamic_method is None:
|
||||
# Fixed rank - use svd_lowrank for 10x speedup
|
||||
rank = min(max_rank, min(out_size, in_size) - 1)
|
||||
U, S, Vh = torch.svd_lowrank(weight, q=rank, niter=niter)
|
||||
Vh = Vh.T # svd_lowrank returns V, not Vh
|
||||
new_rank = rank
|
||||
new_alpha = float(scale * new_rank)
|
||||
stats = {"sum_retained": 1.0, "fro_retained": 1.0} # Not computed for lowrank
|
||||
else:
|
||||
# Dynamic methods need full SVD to compute rank from all singular values
|
||||
U, S, Vh = linalg.svd(weight, full_matrices=False)
|
||||
new_rank, new_alpha, stats = _compute_resize(S, max_rank, dynamic_method, dynamic_param, scale)
|
||||
U = U[:, :new_rank]
|
||||
S = S[:new_rank]
|
||||
Vh = Vh[:new_rank, :]
|
||||
|
||||
lora_up = U @ torch.diag(S)
|
||||
lora_down = Vh
|
||||
|
||||
return {
|
||||
"lora_down": lora_down.cpu().contiguous(),
|
||||
"lora_up": lora_up.cpu().contiguous(),
|
||||
"new_rank": new_rank,
|
||||
"new_alpha": new_alpha,
|
||||
**stats
|
||||
}
|
||||
|
||||
|
||||
def _extract_conv(
|
||||
weight: torch.Tensor,
|
||||
max_rank: int,
|
||||
dynamic_method: Optional[str],
|
||||
dynamic_param: Optional[float],
|
||||
scale: float,
|
||||
niter: int = 2,
|
||||
) -> Dict:
|
||||
"""Extract LoRA from merged conv weight."""
|
||||
out_ch, in_ch, kh, kw = weight.shape
|
||||
mat = weight.reshape(out_ch, -1)
|
||||
|
||||
if dynamic_method is None:
|
||||
# Fixed rank - use svd_lowrank for 10x speedup
|
||||
rank = min(max_rank, min(mat.shape) - 1)
|
||||
U, S, Vh = torch.svd_lowrank(mat, q=rank, niter=niter)
|
||||
Vh = Vh.T # svd_lowrank returns V, not Vh
|
||||
new_rank = rank
|
||||
new_alpha = float(scale * new_rank)
|
||||
stats = {"sum_retained": 1.0, "fro_retained": 1.0} # Not computed for lowrank
|
||||
else:
|
||||
# Dynamic methods need full SVD to compute rank from all singular values
|
||||
U, S, Vh = linalg.svd(mat, full_matrices=False)
|
||||
new_rank, new_alpha, stats = _compute_resize(S, max_rank, dynamic_method, dynamic_param, scale)
|
||||
U = U[:, :new_rank]
|
||||
S = S[:new_rank]
|
||||
Vh = Vh[:new_rank, :]
|
||||
|
||||
lora_up = (U @ torch.diag(S)).reshape(out_ch, new_rank, 1, 1)
|
||||
lora_down = Vh.reshape(new_rank, in_ch, kh, kw)
|
||||
|
||||
return {
|
||||
"lora_down": lora_down.cpu().contiguous(),
|
||||
"lora_up": lora_up.cpu().contiguous(),
|
||||
"new_rank": new_rank,
|
||||
"new_alpha": new_alpha,
|
||||
**stats
|
||||
}
|
||||
|
||||
|
||||
def _compute_resize(
|
||||
S: torch.Tensor,
|
||||
max_rank: int,
|
||||
dynamic_method: Optional[str],
|
||||
dynamic_param: Optional[float],
|
||||
scale: float
|
||||
) -> Tuple[int, float, Dict]:
|
||||
"""Compute new rank and alpha based on resize method."""
|
||||
|
||||
if dynamic_method == "sv_ratio" and dynamic_param is not None:
|
||||
# Note: _index_sv_ratio uses S[0]*ratio, kohya uses S[0]/ratio
|
||||
# Use kohya convention here for consistency with their tool
|
||||
min_sv = S[0] / dynamic_param
|
||||
new_rank = max(1, int(torch.sum(S > min_sv).item()))
|
||||
elif dynamic_method == "sv_cumulative" and dynamic_param is not None:
|
||||
new_rank = _index_sv_cumulative(S, dynamic_param)
|
||||
elif dynamic_method == "sv_fro" and dynamic_param is not None:
|
||||
new_rank = _index_sv_fro(S, dynamic_param)
|
||||
else:
|
||||
new_rank = max_rank
|
||||
|
||||
# Clamp rank
|
||||
if S[0] < MIN_SV:
|
||||
new_rank = 1
|
||||
else:
|
||||
new_rank = max(1, min(new_rank, max_rank, len(S) - 1))
|
||||
|
||||
new_alpha = float(scale * new_rank)
|
||||
|
||||
# Compute retention stats
|
||||
s_sum = float(torch.sum(torch.abs(S)))
|
||||
s_rank = float(torch.sum(torch.abs(S[:new_rank]))) if new_rank <= len(S) else s_sum
|
||||
|
||||
S_sq = S.pow(2)
|
||||
s_fro = float(torch.sqrt(torch.sum(S_sq)))
|
||||
s_red_fro = float(torch.sqrt(torch.sum(S_sq[:new_rank]))) if new_rank <= len(S) else s_fro
|
||||
|
||||
stats = {
|
||||
"sum_retained": s_rank / s_sum if s_sum > 0 else 1.0,
|
||||
"fro_retained": s_red_fro / s_fro if s_fro > 0 else 1.0,
|
||||
}
|
||||
|
||||
return new_rank, new_alpha, stats
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Main Resize Function
|
||||
# =============================================================================
|
||||
|
||||
def resize_lora_file(
|
||||
lora_path: str,
|
||||
new_rank: int,
|
||||
dynamic_method: Optional[str],
|
||||
dynamic_param: Optional[float],
|
||||
device: str,
|
||||
save_dtype: torch.dtype,
|
||||
output_filename: str,
|
||||
verbose: bool = True,
|
||||
svd_niter: int = 2,
|
||||
) -> str:
|
||||
"""
|
||||
Resize a LoRA file to a new rank.
|
||||
|
||||
Args:
|
||||
lora_path: Path to input LoRA
|
||||
new_rank: Target rank (max rank for dynamic methods)
|
||||
dynamic_method: None, "sv_ratio", "sv_fro", "sv_cumulative"
|
||||
dynamic_param: Parameter for dynamic method
|
||||
device: Processing device
|
||||
save_dtype: Output dtype
|
||||
output_filename: Output filename (without extension)
|
||||
verbose: Print progress info
|
||||
svd_niter: Power iterations for SVD accuracy
|
||||
|
||||
Returns:
|
||||
Path to saved resized LoRA
|
||||
"""
|
||||
# Prepare memory
|
||||
lora_size_gb = estimate_model_size(lora_path)
|
||||
prepare_for_large_operation(lora_size_gb * 2, torch.device(device))
|
||||
|
||||
handler = MemoryEfficientSafeOpen(lora_path)
|
||||
|
||||
try:
|
||||
metadata = handler.metadata().copy()
|
||||
all_keys = handler.keys()
|
||||
|
||||
# Detect format and extract pairs
|
||||
format_info = detect_lora_format(all_keys)
|
||||
pairs = extract_lora_pairs(all_keys, format_info)
|
||||
network_dim, network_alpha = detect_lora_rank(handler, pairs)
|
||||
|
||||
scale = network_alpha / network_dim if network_dim > 0 else 1.0
|
||||
|
||||
if verbose:
|
||||
method_str = f"{dynamic_method}: {dynamic_param}" if dynamic_method else "fixed"
|
||||
print(f"[LoRA Resize] Format: {format_info['format']}, layers: {format_info['key_count']}")
|
||||
print(f"[LoRA Resize] Original dim={network_dim}, alpha={network_alpha:.1f}, scale={scale:.3f}")
|
||||
print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}")
|
||||
|
||||
output_sd = {}
|
||||
fro_list = []
|
||||
pbar = comfy.utils.ProgressBar(len(pairs))
|
||||
|
||||
with torch.no_grad():
|
||||
for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"):
|
||||
if "down" not in block_keys or "up" not in block_keys:
|
||||
pbar.update(1)
|
||||
continue
|
||||
|
||||
lora_down = handler.get_tensor(block_keys["down"])
|
||||
lora_up = handler.get_tensor(block_keys["up"])
|
||||
|
||||
is_conv = len(lora_down.shape) == 4
|
||||
|
||||
# Transfer to GPU with pinned memory if available
|
||||
if device == 'cuda':
|
||||
lora_down = transfer_to_gpu_pinned(lora_down, device, torch.float32)
|
||||
lora_up = transfer_to_gpu_pinned(lora_up, device, torch.float32)
|
||||
|
||||
# Merge and re-extract
|
||||
if is_conv:
|
||||
weight = _merge_conv(lora_down, lora_up, device)
|
||||
result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
|
||||
else:
|
||||
weight = _merge_linear(lora_down, lora_up, device)
|
||||
result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
|
||||
|
||||
del weight, lora_down, lora_up
|
||||
|
||||
fro_list.append(result['fro_retained'])
|
||||
|
||||
# Store using same format suffixes as input
|
||||
down_suffix = format_info["down_suffix"]
|
||||
up_suffix = format_info["up_suffix"]
|
||||
alpha_suffix = format_info["alpha_suffix"]
|
||||
|
||||
output_sd[f"{block_name}{down_suffix}"] = result["lora_down"].to(save_dtype)
|
||||
output_sd[f"{block_name}{up_suffix}"] = result["lora_up"].to(save_dtype)
|
||||
output_sd[f"{block_name}{alpha_suffix}"] = torch.tensor([result["new_alpha"]]).to(save_dtype)
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
if verbose and fro_list:
|
||||
import numpy as np
|
||||
avg_fro = np.mean(fro_list)
|
||||
std_fro = np.std(fro_list)
|
||||
print(f"[LoRA Resize] Average Frobenius retention: {avg_fro:.1%} ± {std_fro:.3f}")
|
||||
|
||||
# Update metadata
|
||||
if dynamic_method:
|
||||
metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}"
|
||||
metadata["ss_network_dim"] = "Dynamic"
|
||||
metadata["ss_network_alpha"] = "Dynamic"
|
||||
else:
|
||||
metadata["ss_training_comment"] = f"Resized from dim {network_dim} to {new_rank}"
|
||||
metadata["ss_network_dim"] = str(new_rank)
|
||||
metadata["ss_network_alpha"] = str(scale * new_rank)
|
||||
|
||||
# Save
|
||||
output_dir = folder_paths.get_folder_paths("loras")[0]
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
|
||||
|
||||
save_file(output_sd, output_path, metadata)
|
||||
print(f"[LoRA Resize] Saved to {output_path}")
|
||||
|
||||
return output_path
|
||||
|
||||
finally:
|
||||
handler.__exit__(None, None, None)
|
||||
cleanup_after_operation()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Node Definitions
|
||||
# =============================================================================
|
||||
|
||||
class LoRAResizeFixed(io.ComfyNode):
|
||||
"""Resize LoRA to a fixed rank."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LoRAResizeFixed",
|
||||
display_name="LoRA Resize (Fixed Rank)",
|
||||
category="ModelUtils/LoRA/Resize",
|
||||
description="Resize existing LoRA to a specific rank by merging and re-extracting via SVD.",
|
||||
inputs=[
|
||||
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
||||
tooltip="LoRA to resize"),
|
||||
io.Int.Input("new_rank", default=64, min=1, max=1024,
|
||||
tooltip="Target rank"),
|
||||
io.Int.Input("svd_niter", default=2, min=0, max=10,
|
||||
tooltip="SVD power iterations (higher = more accurate but slower)"),
|
||||
io.String.Input("output_filename", default="resized_lora"),
|
||||
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
||||
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, lora_name, new_rank, svd_niter, output_filename, save_dtype, device) -> io.NodeOutput:
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
||||
|
||||
path = resize_lora_file(lora_path, new_rank, None, None, device, dtype, output_filename, svd_niter=svd_niter)
|
||||
return io.NodeOutput(path)
|
||||
|
||||
|
||||
class LoRAResizeRatio(io.ComfyNode):
|
||||
"""Resize LoRA keeping SVs above ratio threshold."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LoRAResizeRatio",
|
||||
display_name="LoRA Resize (SV Ratio)",
|
||||
category="ModelUtils/LoRA/Resize",
|
||||
description="Dynamically resize LoRA, keeping singular values where S[i] > S[0]/ratio.",
|
||||
inputs=[
|
||||
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
||||
tooltip="LoRA to resize"),
|
||||
io.Int.Input("max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum allowed rank"),
|
||||
io.Float.Input("ratio", default=2.0, min=1.0, max=100.0, step=0.1,
|
||||
tooltip="Keep SVs where S[i] > S[0]/ratio"),
|
||||
io.String.Input("output_filename", default="resized_lora_ratio"),
|
||||
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
||||
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, lora_name, max_rank, ratio, output_filename, save_dtype, device) -> io.NodeOutput:
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
||||
|
||||
path = resize_lora_file(lora_path, max_rank, "sv_ratio", ratio, device, dtype, output_filename)
|
||||
return io.NodeOutput(path)
|
||||
|
||||
|
||||
class LoRAResizeFrobenius(io.ComfyNode):
|
||||
"""Resize LoRA to preserve Frobenius norm target."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LoRAResizeFrobenius",
|
||||
display_name="LoRA Resize (Frobenius)",
|
||||
category="ModelUtils/LoRA/Resize",
|
||||
description="Dynamically resize LoRA to preserve target fraction of Frobenius norm.",
|
||||
inputs=[
|
||||
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
||||
tooltip="LoRA to resize"),
|
||||
io.Int.Input("max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum allowed rank"),
|
||||
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
|
||||
tooltip="Target Frobenius norm retention (0.9 = 90%)"),
|
||||
io.String.Input("output_filename", default="resized_lora_fro"),
|
||||
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
||||
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device) -> io.NodeOutput:
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
||||
|
||||
path = resize_lora_file(lora_path, max_rank, "sv_fro", target, device, dtype, output_filename)
|
||||
return io.NodeOutput(path)
|
||||
|
||||
|
||||
class LoRAResizeCumulative(io.ComfyNode):
|
||||
"""Resize LoRA to preserve cumulative SV target."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LoRAResizeCumulative",
|
||||
display_name="LoRA Resize (Cumulative)",
|
||||
category="ModelUtils/LoRA/Resize",
|
||||
description="Dynamically resize LoRA to preserve target fraction of cumulative singular values.",
|
||||
inputs=[
|
||||
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"),
|
||||
tooltip="LoRA to resize"),
|
||||
io.Int.Input("max_rank", default=128, min=1, max=1024,
|
||||
tooltip="Maximum allowed rank"),
|
||||
io.Float.Input("target", default=0.9, min=0.1, max=1.0, step=0.01,
|
||||
tooltip="Target cumulative SV retention (0.9 = 90%)"),
|
||||
io.String.Input("output_filename", default="resized_lora_cumulative"),
|
||||
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
|
||||
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="output_path")],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device) -> io.NodeOutput:
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
|
||||
|
||||
path = resize_lora_file(lora_path, max_rank, "sv_cumulative", target, device, dtype, output_filename)
|
||||
return io.NodeOutput(path)
|
||||
Reference in New Issue
Block a user