Compare commits

..
Author SHA1 Message Date
Silver c3902c235c Merge branch 'main' into feature/dora-extract 2026-06-06 18:57:51 +02:00
silveroxides 2b496bdd8c Bump version to 0.4.2 2026-06-06 18:50:37 +02:00
silveroxides 34223abb25 Add anima lora renaming node 2026-06-06 18:50:17 +02:00
Silver d2a81c0f9c Bump version from 0.4.2 to 0.4.3 2026-05-29 14:38:38 +02:00
Silver 2cb7f6c55e Merge pull request #14 from silveroxides/feature/dora-extract
Feature/dora extract
2026-05-29 14:37:56 +02:00
silveroxides 3a505e704f Add work in progress experimental learned extraction nodes 2026-05-29 14:36:26 +02:00
silveroxides 1512fa7390 Add DoRA extraction nodes 2026-05-28 03:27:47 +02:00
Silver b7365bfd79 Bump version from 0.4.1 to 0.4.2 2026-05-25 14:25:29 +02:00
Silver 34a2b78a9c Update requirements.txt 2026-05-25 14:25:00 +02:00
Silver 5780122057 Merge pull request #13 from silveroxides/refactor/incremental-writer
Refactor/incremental writer
2026-05-25 14:21:33 +02:00
Silver f6aac8cedf Merge branch 'main' into refactor/incremental-writer 2026-05-25 14:21:24 +02:00
silveroxides 7aa70d8173 Bump version from 0.4.0 to 0.4.1 2026-05-25 14:20:07 +02:00
silveroxides 26c8dffff9 Bump version from 0.3.8 to 0.4.0 2026-05-25 14:18:59 +02:00
silveroxides f3497658fb update unifiedefficientloader requirements 2026-05-25 14:17:58 +02:00
silveroxides 6747f2e8dc Fix block indices error 2026-05-25 14:16:59 +02:00
Silver 54bb6621d8 Bump version from 0.3.9 to 0.4.0 2026-05-21 11:14:38 +02:00
Silver d0d1063964 Update unifiedefficientloader version to 0.5.0 2026-05-21 11:14:19 +02:00
Silver 518e8e0b6d Bump version from 0.3.8 to 0.3.9 2026-05-16 16:33:59 +02:00
Silver fee30e36f5 Merge pull request #12 from silveroxides/refactor/incremental-writer
Refactor/incremental writer
2026-05-16 16:33:37 +02:00
silveroxides 94d2a93b91 fix layer renaming issues 2026-05-16 16:32:44 +02:00
silveroxides bb915f0599 refactor all saves to use IncrementalSafetensorsWriter 2026-05-16 11:07:20 +02:00
silveroxides f9a96d78a2 Bump version from 0.3.7 to 0.3.8 2026-05-10 18:52:38 +02:00
silveroxides 8ee4654c80 Improve layer filters in the merger nodes 2026-05-10 18:52:10 +02:00
silveroxides 73ce12cf9f Bump version from 0.3.6 to 0.3.7 2026-05-10 18:22:23 +02:00
silveroxides 21c0972a5b Refine LoRA extraction to output accurate layer formats across the board and add experimental new DARE+Ties merge operation variant 2026-05-10 18:21:54 +02:00
silveroxides 3bcd2bf3f8 Bump version from 0.3.5 to 0.3.6 2026-05-09 03:09:30 +02:00
silveroxides 29f03fe979 Add new 'enhanced' PowerUp/DARE-Ties merge modes 2026-05-09 03:08:57 +02:00
silveroxides 2e6eec7f63 Bump version from 0.3.4 to 0.3.5 2026-05-02 19:25:55 +02:00
silveroxides 2e0f79a396 handle Anima net. layer prefix 2026-05-02 19:25:20 +02:00
silveroxides 9f5d894320 Bump version from 0.3.3 to 0.3.4 2026-04-29 13:26:32 +02:00
silveroxides 10d0547e05 bump unifiedefficientloader requirement 2026-04-29 12:39:20 +02:00
silveroxides d2adfb3183 Bump version from 0.3.2 to 0.3.3 2026-04-27 11:49:18 +02:00
silveroxides 8fbbe8a5ae refactor to use only unifiedefficientloader and fix header loading 2026-04-27 11:48:53 +02:00
silveroxides 308d839ce3 Bump version from 0.3.1 to 0.3.2 2026-04-26 20:21:39 +02:00
silveroxides 7f6e2c4300 Add DARE+TIES Merge mode 2026-04-26 20:21:13 +02:00
silveroxides a8098b769e Bump version from 0.3.0 to 0.3.1 2026-04-26 13:23:17 +02:00
silveroxides ff1706a320 make merger go brrrrr 2026-04-26 12:41:19 +02:00
silveroxides 6dd89d4bbc Bump version to 0.3.0 2026-04-25 22:11:07 +02:00
silveroxides 4a2faa5a8e Add incremental saving when merging 2026-04-25 22:09:24 +02:00
silveroxides bb787378fe add usage of unifiedefficientloader module 2026-04-25 19:06:50 +02:00
Silver 95393aea13 Merge pull request #10 from silveroxides/feature/preview-meta-wf-loader
Feature/preview meta wf loader
2026-04-25 19:05:51 +02:00
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 0d011d1232 improve memory efficiency 2026-02-07 16:01:00 +01:00
18 changed files with 4896 additions and 1133 deletions
+34 -2
View File
@@ -1,7 +1,7 @@
from typing_extensions import override from typing_extensions import override
from comfy_api.latest import ComfyExtension, io from comfy_api.latest import ComfyExtension, io
from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE from .nodes.lora_merger import LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced
from .nodes.metakeys import ( from .nodes.metakeys import (
ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys, ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys,
CheckpointMetaKeys, EmbeddingMetaKeys CheckpointMetaKeys, EmbeddingMetaKeys
@@ -24,13 +24,31 @@ from .nodes.lora_extract_svd import (
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius LoRAExtractKnee, LoRAExtractFrobenius
) )
from .nodes.dora_extract_wd import (
DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
DoRAExtractKnee, DoRAExtractFrobenius
)
from .nodes.dora_learned_wd import (
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius
)
from .nodes.lora_resize import ( from .nodes.lora_resize import (
LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative, LoRAResizeFrobenius, LoRAResizeCumulative,
LoRAMergeToModel 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
)
from .nodes.lora_rename import AnimaLoraRename
class ModelUtilsExtension(ComfyExtension): class ModelUtilsExtension(ComfyExtension):
@override @override
async def get_node_list(self) -> list[type[io.ComfyNode]]: async def get_node_list(self) -> list[type[io.ComfyNode]]:
@@ -53,13 +71,27 @@ class ModelUtilsExtension(ComfyExtension):
# LoRA Extraction # LoRA Extraction
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius, LoRAExtractKnee, LoRAExtractFrobenius,
# DoRA Extraction
DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile,
DoRAExtractKnee, DoRAExtractFrobenius,
# Learned DoRA Extraction
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius,
# LoRA Resize # LoRA Resize
LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative, LoRAResizeFrobenius, LoRAResizeCumulative,
# LoRA Multi-Merge # LoRA Multi-Merge
LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMerge, LoRAMultiMergeDARE, LoRAMultiMergeDAREEnhanced,
# LoRA Merge To Model # LoRA Merge To Model
LoRAMergeToModel, LoRAMergeToModel,
# LoRA Utilities
AnimaLoraRename,
# Downloaders
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
# Info Loaders
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader,
] ]
+24
View File
@@ -33,6 +33,30 @@
--- ---
## Power-Up (DARE+TIES)
> Combines DARE (Drop and Rescale) with TIES (Trim, Elect Sign) — the current best-practice algorithm for merging fine-tuned models, particularly LoRA-tuned ones. DARE sparsifies the task vector via random dropout and rescaling; TIES then removes residual low-magnitude noise and enforces sign consistency. The result is a cleaner, more targeted capability transfer from B to A with less interference.
**Models Used:** A, B
**Parameters:**
- **Alpha:** DARE drop rate. Fraction of delta parameters randomly zeroed. Higher values produce a sparser, more targeted delta. Typical range: `0.5`–`0.9`.
- **Beta:** TIES trim quantile. Fraction of the smallest-magnitude delta parameters zeroed after DARE, removing residual noise. `0.0` disables trimming. Typical range: `0.1`–`0.3`.
- **Gamma:** Lambda scale. Final multiplier applied to the merged delta before adding to Model A. Equivalent to the task arithmetic scaling coefficient. Typical range: `0.5`–`1.5`.
- **Seed:** Random seed for the DARE dropout mask. Controls reproducibility.
**Algorithm:**
1. Compute task vector `δ = B − A`
2. **DARE:** apply random binary mask with keep probability `1 − alpha`; rescale survivors by `1/(1−alpha)`
3. **TIES trim:** zero out parameters below the `beta`-quantile magnitude threshold
4. **TIES elect:** determine dominant sign per position; zero out disagreeing parameters
5. Return `A + gamma × δ_filtered`
**Guidance:**
- Start with `alpha=0.9, beta=0.2, gamma=1.0` for LoRA-derived models
- Lower `alpha` (e.g. `0.5`) for full fine-tunes where more delta mass should be retained
- Increase `gamma` beyond `1.0` to amplify B's influence; decrease below `1.0` to soften it
---
## Enhanced Man Interp ## Enhanced Man Interp
> Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control. > Sophisticated interpolation between values from A and B depending on their difference relative to other values, with manual threshold control.
+981
View File
@@ -0,0 +1,981 @@
"""
DoRA Extraction via SVD and Weight Decomposition.
Native implementation of Weight-Decomposed Low-Rank Adapter extraction
with multiple rank selection modes. Extracted models use the DoRA structure.
"""
import fnmatch
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 .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# =============================================================================
# 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, glob_mode: bool = False) -> list:
"""Compile patterns from whitespace-separated string.
In regex mode (default) returns compiled re.Pattern objects.
In glob mode returns plain strings; fnmatch handles matching.
"""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for p in pattern_string.split():
p = p.strip()
if not p:
continue
if glob_mode:
patterns.append(p)
else:
try:
patterns.append(re.compile(p))
except re.error:
pass
return patterns
def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
"""Check if key matches any pattern.
Glob mode: each pattern is matched as a substring glob using fnmatch
(dots are literal, * matches any sequence of characters).
Regex mode: each pattern is a compiled re.Pattern, matched via search().
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
return any(p.search(key) for p in patterns)
# =============================================================================
# Main Extraction Function
# =============================================================================
def extract_dora_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: 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,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
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
output_path: Full path to output safetensors file
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
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[DoRA 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, glob_mode=glob_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
# Without base, DoRA fallback just treats tensor_a as diff
dora_scale = None
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}")
weight_diff = tensor_a
dora_scale = None
del tensor_b
else:
# Weight Decomposition calculation
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
# Compute directional diff: W_dir = W_ft * (norm_base / norm_ft)
# Use a small epsilon to avoid division by zero
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del tensor_a, tensor_b, w_dir
# 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"[DoRA 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, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
print(f"[DoRA 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, rank_or_diff = result
# rank is the size of lora_down's first dimension
rank = lora_down.shape[0]
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(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting DoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[DoRA 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()
def _build_lora_output_path(output_filename: str) -> str:
"""Build output path for LoRA file."""
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
# =============================================================================
# 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 _format_lora_key(key: str) -> str:
"""
Format the key for LoRA saving.
Always standardizes to 'diffusion_model.' prefix for maximum compatibility.
ComfyUI maps 'diffusion_model.<layer>' generically for all model types, making
it universally compatible. The 'transformer.' prefix only works for a subset of
model types via model-specific mapping code and is intentionally not used here.
"""
if key.endswith(".weight"):
key = key[:-7]
# Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*)
if key.startswith("model.diffusion_model."):
inner_key = key[22:] # strip len("model.diffusion_model.")
return f"diffusion_model.{inner_key}"
# Handle net.* wrapper
if key.startswith("net."):
inner_key = key[4:] # strip len("net.")
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys (transformer_blocks.* / single_transformer_blocks.* without prefix)
if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"):
return f"diffusion_model.{key}"
# Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"):
return key
# Handle already-prefixed keys
if key.startswith("diffusion_model."):
return key
# Absorb now-invalid transformer. prefix into diffusion_model. for compatibility
if key.startswith("transformer."):
return f"diffusion_model.{key[12:]}"
# Handle known Diffusers UNet blocks
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
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="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
]
class DoRAExtractFixed(io.ComfyNode):
"""Extract DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractFixed",
display_name="DoRA Extract (Fixed Rank)",
category="ModelUtils/DoRA",
description="Extract DoRA 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, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractRatio(io.ComfyNode):
"""Extract DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractRatio",
display_name="DoRA Extract (Ratio)",
category="ModelUtils/DoRA",
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, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, 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,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractQuantile(io.ComfyNode):
"""Extract DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractQuantile",
display_name="DoRA Extract (Quantile)",
category="ModelUtils/DoRA",
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, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, 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,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractKnee(io.ComfyNode):
"""Extract DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractKnee",
display_name="DoRA Extract (Knee Detection)",
category="ModelUtils/DoRA",
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, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, 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,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRAExtractFrobenius(io.ComfyNode):
"""Extract DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRAExtractFrobenius",
display_name="DoRA Extract (Frobenius)",
category="ModelUtils/DoRA",
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, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, 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,
glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
+790
View File
@@ -0,0 +1,790 @@
"""
Learned DoRA Extraction via Optimization.
Extracts Weight-Decomposed Low-Rank Adapters using gradient descent optimization
to minimize the reconstruction MSE loss, matching the true non-linear DoRA formulation.
Uses analytical SVD purely as the initialization seed.
"""
import fnmatch
import os
import re
import math
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# Import the existing SV ranking and chunking logic so we don't have to duplicate it
from .dora_extract_wd import (
_compute_rank, _svd_extract_linear_lowrank, _svd_extract_linear,
_svd_extract_conv, _detect_fused_layer, _compile_patterns,
_matches_any_pattern, _build_lora_output_path, _get_model_inputs,
_format_lora_key, _get_common_inputs
)
def _adaptive_lr_update_cosine(curr_lr: float, improved: bool, worse_loss_counter: int, iteration: int, tensor_shape: tuple, min_lr: float, early_stop_stall: int, lr_factor: float = 0.95, lr_cooldown: int = 1) -> tuple:
M, N = tensor_shape
shape_ratio = abs(M - N) / max(M, N) if max(M, N) > 0 else 0
t = min(worse_loss_counter / max(early_stop_stall, 1), 1.0)
u_factor = (1 + math.cos(2 * math.pi * t)) / 2
if improved:
base_boost = 1.25
distance = base_boost - 1.0
scaled_distance = distance * (1.0 + 0.5 * shape_ratio)
boost_mult = 1.0 + scaled_distance * (1.0 - u_factor)
return min(curr_lr * boost_mult, 100.0), True
else:
cooldown = max(lr_cooldown, 1)
if iteration % cooldown != 0:
return curr_lr, False
min_decay = lr_factor - 0.03 * shape_ratio
max_decay = 0.995
decay_mult = min_decay + (max_decay - min_decay) * u_factor
return max(curr_lr * decay_mult, min_lr), True
def _compute_shape_aware_plateau_params(M: int, N: int, lr_patience: int, lr_factor: float, lr_cooldown: int, lr_shape_influence: float = 1.0) -> tuple:
if lr_shape_influence > 0:
aspect_ratio = max(M, N) / min(M, N) if min(M, N) > 0 else 1.0
ar_factor = math.sqrt(aspect_ratio)
blend = lr_shape_influence
effective_patience = lr_patience
raw_factor = lr_factor
aggressive_factor = raw_factor**ar_factor
effective_factor = raw_factor + (aggressive_factor - raw_factor) * blend
effective_cooldown = lr_cooldown
else:
effective_patience = lr_patience
effective_factor = lr_factor
effective_cooldown = lr_cooldown
return effective_patience, effective_factor, effective_cooldown
def _optimize_dora(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
init_down: torch.Tensor,
init_up: torch.Tensor,
optimize_iters: int,
learning_rate: float,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""
Run gradient descent to optimize lora_down and lora_up against the actual DoRA formula.
Args:
tensor_base: Base model weight
tensor_ft: Finetuned model weight
init_down: SVD-initialized lora_down matrix
init_up: SVD-initialized lora_up matrix
optimize_iters: Number of iterations for AdamW
learning_rate: Learning rate for AdamW
optimizer_type: Optimizer choice (adamw, radam, prodigy)
lr_schedule: LR Schedule choice (exponential, plateau, adaptive)
early_stop_loss: Target loss to stop optimizing
early_stop_stall: Number of non-improving iterations before stopping
Returns:
(optimized_down, optimized_up)
"""
if optimize_iters <= 0:
return init_down, init_up
is_conv = tensor_base.ndim == 4
if is_conv:
out_ch, in_ch, kh, kw = tensor_base.shape
flat_base = tensor_base.reshape(out_ch, -1)
flat_ft = tensor_ft.reshape(out_ch, -1)
# Flatten init matrices for matmul: (out, rank) @ (rank, in*k*k)
rank = init_down.shape[0]
flat_up = init_up.reshape(out_ch, rank)
flat_down = init_down.reshape(rank, -1)
else:
flat_base = tensor_base
flat_ft = tensor_ft
flat_up = init_up
flat_down = init_down
# Magnitude of the finetuned model (fixed DoRA scale)
m = torch.linalg.norm(flat_ft, dim=1, keepdim=True).detach()
eps = 1e-8
# Set up learnable parameters
with torch.inference_mode(False), torch.enable_grad():
# Clone inference tensors to normal tensors for autograd
flat_base_opt = flat_base.clone()
flat_ft_opt = flat_ft.clone()
m_opt = m.clone()
# Free original tensors to save memory
del flat_base, flat_ft, m
B = torch.nn.Parameter(flat_up.clone())
A = torch.nn.Parameter(flat_down.clone())
opt_type = optimizer_type.lower()
if opt_type == "radam":
optimizer = torch.optim.RAdam([B, A], lr=learning_rate)
elif opt_type == "prodigy":
try:
from prodigyplus.prodigy_plus_schedulefree import ProdigyPlusScheduleFree
optimizer = ProdigyPlusScheduleFree([B, A], lr=learning_rate, use_schedulefree=False)
except ImportError:
print("ProdigyPlusScheduleFree optimizer not found. Falling back to AdamW. To install: pip install prodigyplus")
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
else:
optimizer = torch.optim.AdamW([B, A], lr=learning_rate)
# Learning rate scheduling and early stopping state
best_loss = float('inf')
worse_loss_counter = 0
plateau_counter = 0
cooldown_counter = 0
current_lr = learning_rate
schedule_type = lr_schedule.lower()
M_dim, N_dim = flat_base_opt.shape[0], flat_base_opt.shape[1]
eff_patience, eff_factor, eff_cooldown = _compute_shape_aware_plateau_params(M_dim, N_dim, lr_patience, lr_factor, lr_cooldown)
# Optimization loop
opt_pbar = tqdm(range(optimize_iters), desc=f" └─ Optimizing {layer_name}", leave=False, dynamic_ncols=True, mininterval=0.5)
for i in opt_pbar:
optimizer.zero_grad()
# 1. Reconstruct directional weight: W_dir = W_base + B @ A
w_approx = flat_base_opt + (B @ A)
# 2. Normalize directional weight
approx_norm = torch.linalg.norm(w_approx, dim=1, keepdim=True) + eps
# 3. Scale by target magnitude
w_reconstructed = m_opt * (w_approx / approx_norm)
# 4. Compute MSE loss against actual finetuned weight
loss = torch.nn.functional.mse_loss(w_reconstructed, flat_ft_opt)
loss_val = loss.item()
if i == 0:
initial_loss = loss_val
best_loss = loss_val
loss.backward()
optimizer.step()
# --- Early Stopping & Schedulers ---
improved = False
if loss_val < best_loss * 0.9999: # 0.01% improvement threshold
best_loss = loss_val
worse_loss_counter = 0
plateau_counter = 0
improved = True
else:
worse_loss_counter += 1
plateau_counter += 1
if loss_val <= early_stop_loss:
break
if worse_loss_counter > early_stop_stall:
break
if current_lr <= early_stop_lr:
break
# LR Schedulers
if schedule_type == "adaptive":
current_lr, updated = _adaptive_lr_update_cosine(
current_lr, improved, worse_loss_counter, i, (M_dim, N_dim), early_stop_lr, early_stop_stall, lr_factor, lr_cooldown
)
if updated:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
elif schedule_type == "plateau":
if cooldown_counter > 0:
cooldown_counter -= 1
elif plateau_counter >= max(eff_patience, 1): # Decay if patience threshold reached
current_lr = max(current_lr * eff_factor, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
plateau_counter = 0 # reset counter after decay
cooldown_counter = eff_cooldown
elif schedule_type == "exponential":
current_lr = max(current_lr * 0.99, early_stop_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
if i % 10 == 0 or i == optimize_iters - 1:
if schedule_type == "plateau":
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "plat": f"{plateau_counter}/{eff_patience}"}, refresh=False)
else:
opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "stall": f"{worse_loss_counter}"}, refresh=False)
final_loss = loss.item()
opt_pbar.close()
print(f" └─ {layer_name} | Loss: {initial_loss:.6f} -> {final_loss:.6f} (stopped at iter {i+1})")
final_up = B.detach()
final_down = A.detach()
# Reshape back to original dimensions
if is_conv:
final_up = final_up.reshape(out_ch, rank, 1, 1)
final_down = final_down.reshape(rank, in_ch, kh, kw)
return final_down, final_up
def _extract_chunked_learned_layer(
tensor_base: torch.Tensor,
tensor_ft: torch.Tensor,
num_chunks: int,
mode: str,
mode_param: float,
device: str,
max_rank: int = None,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
) -> tuple:
"""Extract Learned DoRA from fused layer by chunking."""
out_dim, in_dim = tensor_base.shape
chunk_size = out_dim // num_chunks
all_lora_up = []
all_lora_down = []
for i in range(num_chunks):
chunk_base = tensor_base[i * chunk_size:(i + 1) * chunk_size, :]
chunk_ft = tensor_ft[i * chunk_size:(i + 1) * chunk_size, :]
norm_base = torch.linalg.norm(chunk_base, dim=1, keepdim=True)
norm_ft = torch.linalg.norm(chunk_ft, dim=1, keepdim=True)
eps = 1e-8
w_dir = chunk_ft * (norm_base / (norm_ft + eps))
weight_diff = w_dir - chunk_base
try:
result, mode_str = _svd_extract_linear(weight_diff, mode, mode_param, device, max_rank)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
# Apply learned optimization
lora_down, lora_up = _optimize_dora(
chunk_base, chunk_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=f"{layer_name} (Chunk {i+1}/{num_chunks})"
)
all_lora_down.append(lora_down)
all_lora_up.append(lora_up)
except Exception:
return None, None, 0
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
def extract_dora_learned_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
optimize_iters: int = 500,
learning_rate: float = 1.0,
optimizer_type: str = "prodigy",
lr_schedule: str = "plateau",
lr_patience: int = 2,
lr_factor: float = 0.9,
lr_cooldown: int = 2,
early_stop_loss: float = 1e-6,
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
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,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract Learned DoRA from difference between two models.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[Learned DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
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, glob_mode=glob_skip_patterns):
return "skipped", None
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")
# Cannot extract true DoRA without base. Fallback to diff.
cpu_a = handler_a.get_tensor(key)
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) if use_pinned else cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
weight_diff = tensor_a
tensor_base = torch.zeros_like(tensor_a)
tensor_ft = tensor_a
dora_scale = None
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}")
weight_diff = tensor_a
tensor_base = torch.zeros_like(tensor_a)
tensor_ft = tensor_a
dora_scale = None
del tensor_b
else:
tensor_base = tensor_b
tensor_ft = tensor_a
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del w_dir
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
if weight_diff.ndim < 2:
del weight_diff, tensor_base, tensor_ft
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
# 1. Run Analytical SVD (Initialization)
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
)
# 2. Run Optimization (Learned Rounding)
if mode_str != "full" and dora_scale is not None:
lora_down, lora_up, _ = result
lora_down, lora_up = _optimize_dora(
tensor_base, tensor_ft, lora_down, lora_up, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
)
result = (lora_down, lora_up, _)
except Exception as e:
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[Learned DoRA] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_learned_layer(
tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
)
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, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff, tensor_base, tensor_ft
return "chunked", layer_results
print(f"[Learned DoRA] Failed: {key}: {e}")
del weight_diff, tensor_base, tensor_ft
return "skipped", None
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
rank = lora_down.shape[0]
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(rank, dtype=save_torch_dtype)
if dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff, tensor_base, tensor_ft
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting Learned DoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[Learned DoRA] 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()
# =============================================================================
# Node Definitions
# =============================================================================
def _get_learned_inputs():
return [
io.Int.Input("optimize_iters", default=500, min=0, max=10000,
tooltip="Number of gradient descent iterations (0 = skip optimization)"),
io.Float.Input("learning_rate", default=1.0, min=0.0001, max=100.0, step=0.001,
tooltip="Base learning rate. AdamW prefers ~0.01, Prodigy prefers ~1.0"),
io.Combo.Input("optimizer", options=["prodigy", "adamw", "radam"], default="prodigy",
tooltip="Optimization algorithm"),
io.Combo.Input("lr_schedule", options=["plateau", "adaptive", "exponential", "constant"], default="plateau",
tooltip="Learning rate scheduling strategy"),
io.Int.Input("lr_patience", default=2, min=0, max=1000, tooltip="Steps to wait before decaying LR (Plateau)"),
io.Float.Input("lr_factor", default=0.9, min=0.01, max=1.0, step=0.01, tooltip="Factor to decay LR by (Plateau)"),
io.Int.Input("lr_cooldown", default=2, min=0, max=1000, tooltip="Steps to wait after decay before checking again (Plateau)"),
io.Float.Input("early_stop_loss", default=1e-6, min=0.0, max=1.0, step=1e-9,
tooltip="Stop early if MSE loss drops below this value"),
io.Int.Input("early_stop_stall", default=2000, min=0, max=10000,
tooltip="Stop early if loss doesn't improve for this many iterations"),
io.Float.Input("early_stop_lr", default=9.01e-9, min=0.0, max=1.0, step=1e-10,
tooltip="Stop early if learning rate drops below this value"),
]
class DoRALearnedExtractFixed(io.ComfyNode):
"""Extract Learned DoRA with fixed rank."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFixed",
display_name="DoRA Learned Extract (Fixed Rank) (WIP)",
category="ModelUtils/DoRA",
description="Extract DoRA via MSE Loss minimization with fixed rank.",
is_experimental=True,
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"),
*_get_learned_inputs(),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD iterations for initialization"),
*_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, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractRatio(io.ComfyNode):
"""Extract Learned DoRA with singular value ratio threshold."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractRatio",
display_name="DoRA Learned Extract (Ratio) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping SVs > max(S) / ratio.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_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,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractQuantile(io.ComfyNode):
"""Extract Learned DoRA with cumulative singular value quantile."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractQuantile",
display_name="DoRA Learned Extract (Quantile) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction keeping enough SVs to reach target percentage.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_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,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractKnee(io.ComfyNode):
"""Extract Learned DoRA with automatic knee detection."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractKnee",
display_name="DoRA Learned Extract (Knee Detection) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction automatically finding optimal rank using knee detection.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_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,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
class DoRALearnedExtractFrobenius(io.ComfyNode):
"""Extract Learned DoRA preserving target fraction of Frobenius norm."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DoRALearnedExtractFrobenius",
display_name="DoRA Learned Extract (Frobenius) (WIP)",
category="ModelUtils/DoRA",
description="Learned extraction preserving target fraction of Frobenius norm.",
is_experimental=True,
inputs=[
*_get_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_learned_inputs(),
*_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,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_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_path = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, 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, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
+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 []
+695
View File
@@ -0,0 +1,695 @@
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'))
from unifiedefficientloader import UnifiedSafetensorsLoader
# --- 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):
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
+195 -179
View File
@@ -4,6 +4,7 @@ LoRA Extraction via SVD decomposition.
Native implementation of Low-Rank Adapter extraction with multiple rank selection modes. Native implementation of Low-Rank Adapter extraction with multiple rank selection modes.
No external dependencies required. No external dependencies required.
""" """
import fnmatch
import os import os
import re import re
import torch import torch
@@ -12,13 +13,15 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io 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 ( from .device_utils import (
estimate_model_size, prepare_for_large_operation, estimate_model_size, prepare_for_large_operation,
cleanup_after_operation, get_device_capabilities cleanup_after_operation, get_device_capabilities
) )
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# ============================================================================= # =============================================================================
# Rank Index Functions - Determine LoRA rank from singular values # Rank Index Functions - Determine LoRA rank from singular values
@@ -46,14 +49,14 @@ def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int:
def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -> int: def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Cumulative mode - keep enough SVs to reach target % of total. """Cumulative mode - keep enough SVs to reach target % of total.
Calculates relative to max_rank if provided, otherwise relative to full. Calculates relative to max_rank if provided, otherwise relative to full.
""" """
if max_rank is not None and max_rank < len(S): if max_rank is not None and max_rank < len(S):
total = torch.sum(S[:max_rank]) total = torch.sum(S[:max_rank])
else: else:
total = torch.sum(S) total = torch.sum(S)
if total < 1e-8: if total < 1e-8:
return 1 return 1
cumsum = torch.cumsum(S, dim=0) / total cumsum = torch.cumsum(S, dim=0) / total
@@ -63,7 +66,7 @@ def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -
def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int: def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Frobenius norm mode - preserve target fraction of Frobenius norm. """Frobenius norm mode - preserve target fraction of Frobenius norm.
Calculates relative to max_rank if provided, otherwise relative to full. Calculates relative to max_rank if provided, otherwise relative to full.
This means "retain target% of what's achievable within max_rank". This means "retain target% of what's achievable within max_rank".
""" """
@@ -75,10 +78,10 @@ def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
else: else:
S_sq = S.pow(2) S_sq = S.pow(2)
total_sq = torch.sum(S_sq) total_sq = torch.sum(S_sq)
if total_sq < 1e-8: if total_sq < 1e-8:
return 1 return 1
# Cumsum of all S (not capped) to find where we reach target # Cumsum of all S (not capped) to find where we reach target
cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq
rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1 rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1
@@ -172,27 +175,27 @@ def _svd_extract_linear_lowrank(
) -> tuple: ) -> tuple:
""" """
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction. 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 Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms. the top-k singular values using randomized algorithms.
Args: Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower) niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns: Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full") (lora_down, lora_up, diff), "low rank" or (weight, "full")
""" """
weight = weight.to(device=device, dtype=torch.float32) weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape out_dim, in_dim = weight.shape
# Clamp rank to valid range # Clamp rank to valid range
max_possible_rank = min(out_dim, in_dim) max_possible_rank = min(out_dim, in_dim)
rank = max(1, min(rank, max_possible_rank - 1)) rank = max(1, min(rank, max_possible_rank))
# Check if decomposition is worthwhile # Check if decomposition is worthwhile
if rank >= max_possible_rank // 2: if rank >= max_possible_rank:
return weight, "full" return weight, "full"
# Use svd_lowrank for faster low-rank approximation # Use svd_lowrank for faster low-rank approximation
try: try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter) U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
@@ -200,7 +203,7 @@ def _svd_extract_linear_lowrank(
Vh = V.T # [rank, in_dim] Vh = V.T # [rank, in_dim]
except Exception as e: except Exception as e:
raise RuntimeError(f"svd_lowrank failed: {e}") raise RuntimeError(f"svd_lowrank failed: {e}")
# Clamp outliers # Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0: if clamp_quantile < 1.0 and len(S) > 0:
try: try:
@@ -208,14 +211,14 @@ def _svd_extract_linear_lowrank(
S = S.clamp(max=max_val) S = S.clamp(max=max_val)
except RuntimeError: except RuntimeError:
pass pass
# Construct LoRA matrices # Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank] lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim] lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff # Compute reconstruction diff
diff = weight - (lora_up @ lora_down) diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank" return (lora_down, lora_up, diff), "low rank"
@@ -239,7 +242,7 @@ def _svd_extract_linear(
""" """
weight = weight.to(device=device, dtype=torch.float32) weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape out_dim, in_dim = weight.shape
# For fixed rank, use the faster svd_lowrank # For fixed rank, use the faster svd_lowrank
if mode == "fixed" and mode_param > 0: if mode == "fixed" and mode_param > 0:
target_rank = int(mode_param) target_rank = int(mode_param)
@@ -257,7 +260,7 @@ def _svd_extract_linear(
rank = _compute_rank(S, mode, mode_param, max_rank) rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile # Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim) // 2: if rank >= min(out_dim, in_dim):
return weight, "full" return weight, "full"
# Truncate to rank # Truncate to rank
@@ -303,7 +306,7 @@ def _svd_extract_conv(
# Flatten for SVD # Flatten for SVD
if is_1x1: if is_1x1:
mat = weight.squeeze() # [out_ch, in_ch] mat = weight.view(out_ch, in_ch) # [out_ch, in_ch]
else: else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k] mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
@@ -382,8 +385,11 @@ def _extract_chunked_layer(
return None, None, 0 return None, None, 0
# Combine chunks # 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_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = torch.stack(all_lora_down, dim=0).mean(dim=0) combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0] rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank return combined_lora_up, combined_lora_down, rank
@@ -393,14 +399,22 @@ def _extract_chunked_layer(
# Pattern Matching # Pattern Matching
# ============================================================================= # =============================================================================
def _compile_patterns(pattern_string: str) -> list: def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list:
"""Compile regex patterns from whitespace-separated string.""" """Compile patterns from whitespace-separated string.
In regex mode (default) returns compiled re.Pattern objects.
In glob mode returns plain strings; fnmatch handles matching.
"""
if not pattern_string or not pattern_string.strip(): if not pattern_string or not pattern_string.strip():
return [] return []
patterns = [] patterns = []
for p in pattern_string.split(): for p in pattern_string.split():
p = p.strip() p = p.strip()
if p: if not p:
continue
if glob_mode:
patterns.append(p)
else:
try: try:
patterns.append(re.compile(p)) patterns.append(re.compile(p))
except re.error: except re.error:
@@ -408,8 +422,15 @@ def _compile_patterns(pattern_string: str) -> list:
return patterns return patterns
def _matches_any_pattern(key: str, patterns: list) -> bool: def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool:
"""Check if key matches any pattern.""" """Check if key matches any pattern.
Glob mode: each pattern is matched as a substring glob using fnmatch
(dots are literal, * matches any sequence of characters).
Regex mode: each pattern is a compiled re.Pattern, matched via search().
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
return any(p.search(key) for p in patterns) return any(p.search(key) for p in patterns)
@@ -425,6 +446,7 @@ def extract_lora_from_files(
conv_param: float, conv_param: float,
device: str, device: str,
save_dtype: str, save_dtype: str,
output_path: str,
linear_max_rank: int = None, linear_max_rank: int = None,
conv_max_rank: int = None, conv_max_rank: int = None,
clamp_quantile: float = 0.99, clamp_quantile: float = 0.99,
@@ -433,9 +455,12 @@ def extract_lora_from_files(
mismatch_mode: str = "skip", mismatch_mode: str = "skip",
chunk_large_layers: bool = True, chunk_large_layers: bool = True,
svd_niter: int = 2, svd_niter: int = 2,
) -> dict[str, torch.Tensor]: lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> None:
""" """
Extract LoRA from difference between two models. Extract LoRA from difference between two models, writing incrementally to disk.
Args: Args:
model_a_path: Finetuned model path model_a_path: Finetuned model path
@@ -445,6 +470,7 @@ def extract_lora_from_files(
conv_param: Mode parameter for conv layers conv_param: Mode parameter for conv layers
device: Computation device device: Computation device
save_dtype: Output dtype save_dtype: Output dtype
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers conv_max_rank: Max rank for conv layers
clamp_quantile: Quantile for weight clamping clamp_quantile: Quantile for weight clamping
@@ -452,61 +478,62 @@ def extract_lora_from_files(
skip_patterns_str: Patterns for layers to skip skip_patterns_str: Patterns for layers to skip
mismatch_mode: How to handle mismatches mismatch_mode: How to handle mismatches
chunk_large_layers: Enable chunked extraction chunk_large_layers: Enable chunked extraction
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
""" """
output_sd = {}
save_torch_dtype = { save_torch_dtype = {
"fp32": torch.float32, "fp32": torch.float32,
"fp16": torch.float16, "fp16": torch.float16,
"bf16": torch.bfloat16, "bf16": torch.bfloat16,
}.get(save_dtype, torch.float16) }.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str) skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation # Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path) 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...") 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 prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) # 1.5x for SVD headroom
handler_a = MemoryEfficientSafeOpen(model_a_path) handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path) handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try: try:
keys_a = set(handler_a.keys()) keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys()) keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")] weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys)) pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"): def _process_layer(key):
lora_name = _format_lora_key(key) lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns): if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
stats["skipped"] += 1 return "skipped", None
pbar.update(1)
continue
# Load tensors with pinned memory for CUDA # Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda' use_pinned = device == 'cuda'
if key not in keys_b: if key not in keys_b:
if mismatch_mode == "skip": if mismatch_mode == "skip":
stats["skipped"] += 1 return "skipped", None
pbar.update(1) if mismatch_mode == "error":
continue raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key) cpu_a = handler_a.get_tensor(key)
if use_pinned: if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32) weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else: else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32) weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a 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: else:
cpu_a = handler_a.get_tensor(key) cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key) cpu_b = handler_b.get_tensor(key)
if use_pinned: if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32) tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
@@ -517,10 +544,12 @@ def extract_lora_from_files(
if tensor_a.shape != tensor_b.shape: if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip": if mismatch_mode == "skip":
stats["skipped"] += 1
pbar.update(1)
del tensor_a, tensor_b del tensor_a, tensor_b
continue 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 weight_diff = tensor_a
del tensor_b del tensor_b
else: else:
@@ -529,19 +558,16 @@ def extract_lora_from_files(
# Skip small differences # Skip small differences
if min_diff > 0 and weight_diff.abs().max() < min_diff: if min_diff > 0 and weight_diff.abs().max() < min_diff:
stats["skipped"] += 1
del weight_diff del weight_diff
pbar.update(1) return "skipped", None
continue
# Skip 1D tensors # Skip 1D tensors
if weight_diff.ndim < 2: if weight_diff.ndim < 2:
stats["skipped"] += 1
del weight_diff del weight_diff
pbar.update(1) return "skipped", None
continue
is_conv = weight_diff.ndim == 4 is_conv = weight_diff.ndim == 4
layer_results = {}
try: try:
if is_conv: if is_conv:
@@ -562,33 +588,46 @@ def extract_lora_from_files(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
) )
if lora_up is not None: if lora_up is not None:
output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
stats["chunked"] += 1
del weight_diff del weight_diff
pbar.update(1) return "chunked", layer_results
continue
print(f"[LoRA Extract] Failed: {key}: {e}") print(f"[LoRA Extract] Failed: {key}: {e}")
stats["skipped"] += 1
del weight_diff del weight_diff
pbar.update(1) return "skipped", None
continue
# Store result # Store result
if mode_str == "full": if mode_str == "full":
output_sd[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
stats["full"] += 1 status = "full"
else: else:
lora_down, lora_up, _ = result lora_down, lora_up, _ = result
output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype) status = "extracted"
stats["extracted"] += 1
del weight_diff del weight_diff
pbar.update(1) return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped") f"{stats['full']} full, {stats['skipped']} skipped")
@@ -598,22 +637,12 @@ def extract_lora_from_files(
handler_b.__exit__(None, None, None) handler_b.__exit__(None, None, None)
cleanup_after_operation() 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")
def _build_lora_output_path(output_filename: str) -> str:
"""Build output path for LoRA file."""
output_dir = folder_paths.get_folder_paths("loras")[0] output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") return 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
# ============================================================================= # =============================================================================
@@ -629,86 +658,56 @@ def _get_model_inputs():
] ]
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: def _format_lora_key(key: str) -> str:
""" """
Format the key for LoRA saving. Format the key for LoRA saving.
Standardizes to 'diffusion_model.' or 'transformer.' prefix. Always standardizes to 'diffusion_model.' prefix for maximum compatibility.
ComfyUI maps 'diffusion_model.<layer>' generically for all model types, making
it universally compatible. The 'transformer.' prefix only works for a subset of
model types via model-specific mapping code and is intentionally not used here.
""" """
if key.endswith(".weight"): if key.endswith(".weight"):
key = key[:-7] key = key[:-7]
# Handle ComfyUI Checkpoint wrapper # Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*)
if key.startswith("model.diffusion_model."): if key.startswith("model.diffusion_model."):
# Check if it's a Diffusers-format transformer inside a checkpoint inner_key = key[22:] # strip len("model.diffusion_model.")
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}" return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys # Handle net.* wrapper
if key.startswith("net."):
inner_key = key[4:] # strip len("net.")
return f"diffusion_model.{inner_key}"
# Handle direct Diffusers keys (transformer_blocks.* / single_transformer_blocks.* without prefix)
if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"): if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"):
return f"transformer.{key}" return f"diffusion_model.{key}"
# Handle legacy lora_unet_ prefix (for resizing without base) # Handle legacy lora_unet_ prefix (for resizing without base)
if key.startswith("lora_unet_"): 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 return key
# Handle known Unet blocks (Diffusers format UNet) # Handle already-prefixed keys
if key.startswith("diffusion_model."):
return key
# Absorb now-invalid transformer. prefix into diffusion_model. for compatibility
if key.startswith("transformer."):
return f"diffusion_model.{key[12:]}"
# Handle known Diffusers UNet blocks
if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]): 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}" return f"diffusion_model.{key}"
# Default fallback to diffusion_model # Default fallback
return f"diffusion_model.{key}" return f"diffusion_model.{key}"
def _get_common_inputs(): def _get_common_inputs():
return [ 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, io.Boolean.Input("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"), 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, io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
@@ -720,7 +719,10 @@ def _get_common_inputs():
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True, io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip"), tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
io.Boolean.Input("glob_skip_patterns", default=False,
tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
] ]
@@ -736,9 +738,9 @@ class LoRAExtractFixed(io.ComfyNode):
description="Extract LoRA with specified fixed rank for each layer type.", description="Extract LoRA with specified fixed rank for each layer type.",
inputs=[ inputs=[
*_get_model_inputs(), *_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=4096, io.Int.Input("linear_dim", default=64, min=1, max=16384,
tooltip="Rank for linear/attention layers"), tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=4096, io.Int.Input("conv_dim", default=32, min=1, max=16384,
tooltip="Rank for conv layers"), tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10, io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"), tooltip="SVD power iterations (higher = more accurate but slower)"),
@@ -751,18 +753,20 @@ class LoRAExtractFixed(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, 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, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput: save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) 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) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files( extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim, model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, linear_dim, conv_dim, device, save_dtype, output_path, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache, glob_skip_patterns
) )
return io.NodeOutput(_save_lora(output_sd, output_filename)) return io.NodeOutput(output_path)
class LoRAExtractRatio(io.ComfyNode): class LoRAExtractRatio(io.ComfyNode):
@@ -781,9 +785,9 @@ class LoRAExtractRatio(io.ComfyNode):
tooltip="Ratio threshold for linear layers (higher = more SVs kept)"), 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, 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)"), tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072, io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"), tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072, io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"), tooltip="Maximum rank for conv layers"),
*_get_common_inputs(), *_get_common_inputs(),
], ],
@@ -794,18 +798,21 @@ class LoRAExtractRatio(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, 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, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput: save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) 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) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files( extract_lora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, linear_max_rank, conv_max_rank, device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
) )
return io.NodeOutput(_save_lora(output_sd, output_filename)) return io.NodeOutput(output_path)
class LoRAExtractQuantile(io.ComfyNode): class LoRAExtractQuantile(io.ComfyNode):
@@ -824,9 +831,9 @@ class LoRAExtractQuantile(io.ComfyNode):
tooltip="Target cumulative % for linear layers"), tooltip="Target cumulative % for linear layers"),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target cumulative % for conv layers"), tooltip="Target cumulative % for conv layers"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072, io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"), tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072, io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"), tooltip="Maximum rank for conv layers"),
*_get_common_inputs(), *_get_common_inputs(),
], ],
@@ -837,18 +844,21 @@ class LoRAExtractQuantile(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, 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, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput: save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) 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) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files( extract_lora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, linear_max_rank, conv_max_rank, device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
) )
return io.NodeOutput(_save_lora(output_sd, output_filename)) return io.NodeOutput(output_path)
class LoRAExtractKnee(io.ComfyNode): class LoRAExtractKnee(io.ComfyNode):
@@ -865,9 +875,9 @@ class LoRAExtractKnee(io.ComfyNode):
*_get_model_inputs(), *_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
default="sv_knee", tooltip="Knee detection method"), default="sv_knee", tooltip="Knee detection method"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072, io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"), tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072, io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"), tooltip="Maximum rank for conv layers"),
*_get_common_inputs(), *_get_common_inputs(),
], ],
@@ -878,18 +888,21 @@ class LoRAExtractKnee(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, 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, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput: save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) 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) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files( extract_lora_from_files(
model_a_path, model_b_path, knee_method, 0, 0, model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, linear_max_rank, conv_max_rank, device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
) )
return io.NodeOutput(_save_lora(output_sd, output_filename)) return io.NodeOutput(output_path)
class LoRAExtractFrobenius(io.ComfyNode): class LoRAExtractFrobenius(io.ComfyNode):
@@ -908,9 +921,9 @@ class LoRAExtractFrobenius(io.ComfyNode):
tooltip="Target Frobenius norm fraction for linear"), tooltip="Target Frobenius norm fraction for linear"),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target Frobenius norm fraction for conv"), tooltip="Target Frobenius norm fraction for conv"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072, io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"), tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072, io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"), tooltip="Maximum rank for conv layers"),
*_get_common_inputs(), *_get_common_inputs(),
], ],
@@ -921,15 +934,18 @@ class LoRAExtractFrobenius(io.ComfyNode):
@classmethod @classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, 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, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput: save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) 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) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
output_path = _build_lora_output_path(output_filename)
output_sd = extract_lora_from_files( extract_lora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target, model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, linear_max_rank, conv_max_rank, device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
) )
return io.NodeOutput(_save_lora(output_sd, output_filename)) return io.NodeOutput(output_path)
+709 -246
View File
File diff suppressed because it is too large Load Diff
+119
View File
@@ -0,0 +1,119 @@
import os
import folder_paths
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
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
REVERSE_RENAME_DICT = {
"transformer_blocks.": "blocks.",
"norm1.linear_1": "adaln_modulation_self_attn.1",
"norm1.linear_2": "adaln_modulation_self_attn.2",
"norm2.linear_1": "adaln_modulation_cross_attn.1",
"norm2.linear_2": "adaln_modulation_cross_attn.2",
"norm3.linear_1": "adaln_modulation_mlp.1",
"norm3.linear_2": "adaln_modulation_mlp.2",
"attn1.to_q": "self_attn.q_proj",
"attn1.to_k": "self_attn.k_proj",
"attn1.to_v": "self_attn.v_proj",
"attn1.to_out.0": "self_attn.output_proj",
"attn2.to_q": "cross_attn.q_proj",
"attn2.to_k": "cross_attn.k_proj",
"attn2.to_v": "cross_attn.v_proj",
"attn2.to_out.0": "cross_attn.output_proj",
"ff.net.0.proj": "mlp.layer1",
"ff.net.2": "mlp.layer2",
"norm_out.linear_1": "final_layer.adaln_modulation.1",
"norm_out.linear_2": "final_layer.adaln_modulation.2",
"proj_out": "final_layer.linear",
"time_embed.t_embedder": "t_embedder.1",
"time_embed.norm": "t_embedding_norm",
"patch_embed.proj": "x_embedder.proj.1",
}
def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filename: str) -> str:
"""
Converts a Diffusers Anima LoRA back to its non-diffusers format.
"""
model_type = "loras"
model_path = folder_paths.get_full_path_or_raise(model_type, lora_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
# Prepare memory before operation
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
# Get metadata from safe_open (header only, no tensor loading)
with safe_open(model_path_to_load, framework="pt", device="cpu") as f:
metadata = f.metadata() or {}
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Stream tensors, rename on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
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 Anima LoRA keys", unit="keys"):
new_key = key
if key.startswith("text_conditioner."):
new_key = f"diffusion_model.llm_adapter.{key.removeprefix('text_conditioner.')}"
elif key.startswith("transformer."):
new_key = key.removeprefix("transformer.")
for old_part, new_part in REVERSE_RENAME_DICT.items():
new_key = new_key.replace(old_part, new_part)
new_key = f"diffusion_model.{new_key}"
writer.write(new_key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation
cleanup_after_operation()
return output_path
class AnimaLoraRename(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="AnimaLoraRename",
display_name="Anima LoRA Rename (Diffusers to Non-Diffusers)",
category="ModelUtils/LoRA",
description="Loads a Diffusers Anima LoRA, reverses its key names to the non-diffusers format, and saves it as a new safetensors file.",
inputs=[
io.Combo.Input(
"lora_name",
options=folder_paths.get_filename_list("loras"),
),
io.String.Input("output_filename", default="anima_lora_non_diffusers"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(cls, lora_name: str, output_filename: str) -> io.NodeOutput:
path = _convert_diffusers_to_non_diffusers_anima_lora(lora_name, output_filename)
return io.NodeOutput(path)
+195 -155
View File
@@ -12,11 +12,13 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io 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 from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
from typing import Optional, Dict, Tuple, List
# Reuse rank functions from extraction module # Reuse rank functions from extraction module
from .lora_extract_svd import ( from .lora_extract_svd import (
_index_sv_ratio, _index_sv_ratio,
@@ -189,7 +191,7 @@ def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int
# Get dim from down weight shape # Get dim from down weight shape
if network_dim is None: if network_dim is None:
down_key = block_keys["down"] down_key = block_keys["down"]
shape = handler.header[down_key]["shape"] shape = handler.get_shape(down_key)
# Linear: [rank, in_features] or Conv: [rank, in_ch, k, k] # Linear: [rank, in_features] or Conv: [rank, in_ch, k, k]
network_dim = shape[0] network_dim = shape[0]
@@ -416,60 +418,7 @@ def resize_lora_file(
print(f"[LoRA Resize] Original dim={network_dim}, alpha={network_alpha:.1f}, scale={scale:.3f}") 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}") print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}")
output_sd = {} # Build metadata and output path before loop (all values known at this point)
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)
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: if dynamic_method:
metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}" metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}"
metadata["ss_network_dim"] = "Dynamic" metadata["ss_network_dim"] = "Dynamic"
@@ -479,14 +428,78 @@ def resize_lora_file(
metadata["ss_network_dim"] = str(new_rank) metadata["ss_network_dim"] = str(new_rank)
metadata["ss_network_alpha"] = str(scale * new_rank) metadata["ss_network_alpha"] = str(scale * new_rank)
# Save
output_dir = folder_paths.get_folder_paths("loras")[0] output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path, metadata) fro_list = []
print(f"[LoRA Resize] Saved to {output_path}") pbar = comfy.utils.ProgressBar(len(pairs))
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
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"]
# Preserve original key format from input LoRA verbatim
new_block_name = block_name
block_sd = {
f"{new_block_name}{down_suffix}": result["lora_down"].to(save_dtype),
f"{new_block_name}{up_suffix}": result["lora_up"].to(save_dtype),
}
if alpha_suffix:
block_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype)
writer.write_dict(block_sd)
del result
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
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}")
print(f"[LoRA Resize] Saved to {output_path}")
return output_path return output_path
finally: finally:
@@ -706,7 +719,7 @@ def merge_loras_to_model(
print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}") print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}")
# Common prefixes # Common prefixes
BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."] BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."]
LORA_PREFIXES = [ LORA_PREFIXES = [
"lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_",
"lycoris_", "diffusion_model.", "transformer.", "unet." "lycoris_", "diffusion_model.", "transformer.", "unet."
@@ -752,7 +765,11 @@ def merge_loras_to_model(
base_metadata = base_handler.metadata().copy() if base_handler.metadata() else {} 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}" base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}"
output_sd = {} # Build output path before loop
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")
stats = {"merged": 0, "copied": 0, "skipped": 0} stats = {"merged": 0, "copied": 0, "skipped": 0}
base_keys = list(base_handler.keys()) base_keys = list(base_handler.keys())
pbar = comfy.utils.ProgressBar(len(base_keys)) pbar = comfy.utils.ProgressBar(len(base_keys))
@@ -760,114 +777,116 @@ def merge_loras_to_model(
if verbose: if verbose:
print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...") print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...")
with torch.no_grad(): writer = IncrementalSafetensorsWriter(output_path, metadata=base_metadata)
for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): writer.__enter__()
# Check skip patterns try:
if _matches_any_pattern(base_key, skip_patterns): with torch.no_grad():
stats["skipped"] += 1 for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"):
pbar.update(1) # Check skip patterns
continue if _matches_any_pattern(base_key, skip_patterns):
stats["skipped"] += 1
pbar.update(1)
continue
# Load base weight # Load base weight
cpu_base = base_handler.get_tensor(base_key) cpu_base = base_handler.get_tensor(base_key)
# Only process weight tensors for LoRA merging # Only process weight tensors for LoRA merging
if base_key.endswith(".weight"): if base_key.endswith(".weight"):
core = extract_core_layer_base(base_key) core = extract_core_layer_base(base_key)
core_underscored = core.replace(".", "_") core_underscored = core.replace(".", "_")
# Check if any LoRA contributes to this layer # Check if any LoRA contributes to this layer
if core_underscored in lora_lookup: if core_underscored in lora_lookup:
# Transfer to GPU for computation # Transfer to GPU for computation
if device == 'cuda': if device == 'cuda':
base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32) 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: else:
# Standard LoRA format base_weight = cpu_base.to(device=device, dtype=torch.float32)
if "down" not in block_keys or "up" not in block_keys: del cpu_base
continue
cpu_down = info["handler"].get_tensor(block_keys["down"]) # Accumulate deltas from all contributing LoRAs
cpu_up = info["handler"].get_tensor(block_keys["up"]) for info, block_keys in lora_lookup[core_underscored]:
if device == 'cuda': is_full_diff = info["format_info"].get("is_full_diff", False)
lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32) 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: else:
lora_down = cpu_down.to(device=device, dtype=torch.float32) # Standard LoRA format
lora_up = cpu_up.to(device=device, dtype=torch.float32) if "down" not in block_keys or "up" not in block_keys:
del cpu_down, cpu_up continue
# Get alpha cpu_down = info["handler"].get_tensor(block_keys["down"])
if "alpha" in block_keys: cpu_up = info["handler"].get_tensor(block_keys["up"])
alpha_tensor = info["handler"].get_tensor(block_keys["alpha"]) if device == 'cuda':
layer_alpha = float(alpha_tensor.item()) lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32)
else: lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32)
layer_alpha = float(info["network_dim"]) else:
layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0 lora_down = cpu_down.to(device=device, dtype=torch.float32)
effective_scale = layer_scale * info["weight"] lora_up = cpu_up.to(device=device, dtype=torch.float32)
del cpu_down, cpu_up
# Compute delta # Get alpha
is_conv = len(lora_down.shape) == 4 if "alpha" in block_keys:
if is_conv: alpha_tensor = info["handler"].get_tensor(block_keys["alpha"])
in_rank, in_size, kernel_size, k_ = lora_down.shape layer_alpha = float(alpha_tensor.item())
out_size, out_rank, _, _ = lora_up.shape else:
delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1) layer_alpha = float(info["network_dim"])
delta = delta.reshape(out_size, in_size, kernel_size, kernel_size) layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0
else: effective_scale = layer_scale * info["weight"]
delta = lora_up @ lora_down
del lora_down, lora_up
# Apply delta to base weight # Compute delta
base_weight = base_weight + effective_scale * delta is_conv = len(lora_down.shape) == 4
del delta 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
# Store merged weight # Apply delta to base weight
output_sd[base_key] = base_weight.to(save_dtype).cpu().contiguous() base_weight = base_weight + effective_scale * delta
del base_weight del delta
stats["merged"] += 1
# Write merged weight
writer.write(base_key, base_weight.to(save_dtype).cpu().contiguous())
del base_weight
stats["merged"] += 1
else:
# No LoRA contribution, write as-is
writer.write(base_key, cpu_base.to(save_dtype).contiguous())
del cpu_base
stats["copied"] += 1
else: else:
# No LoRA contribution, copy as-is # Non-weight tensor (bias, norm, etc.), write as-is
output_sd[base_key] = cpu_base.to(save_dtype).contiguous() writer.write(base_key, cpu_base.to(save_dtype).contiguous())
del cpu_base del cpu_base
stats["copied"] += 1 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) pbar.update(1)
finally:
writer.__exit__(None, None, None)
if verbose: if verbose:
print(f"[LoRA Merge To Model] Done: {stats['merged']} merged, {stats['copied']} copied, {stats['skipped']} skipped") 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}") print(f"[LoRA Merge To Model] Saved to {output_path}")
return output_path return output_path
finally: finally:
@@ -886,11 +905,11 @@ class LoRAMergeToModel(io.ComfyNode):
node_id="LoRAMergeToModel", node_id="LoRAMergeToModel",
display_name="LoRA Merge To Model", display_name="LoRA Merge To Model",
category="ModelUtils/LoRA/Merge", category="ModelUtils/LoRA/Merge",
description="Merge 1-4 LoRAs into a base model and save the result. Saves to base model directory.", description="Merge 1-8 LoRAs into a base model and save the result. Saves to base model directory.",
inputs=[ inputs=[
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"), io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Base model the LoRAs were trained on"), tooltip="Base model the LoRAs were trained on"),
io.Combo.Input("lora_count", options=["1", "2", "3", "4"], default="2", io.Combo.Input("lora_count", options=["1", "2", "3", "4", "5", "6", "7", "8"], default="2",
tooltip="Number of LoRAs to merge"), tooltip="Number of LoRAs to merge"),
# LoRA 1 # LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"),
@@ -912,6 +931,26 @@ class LoRAMergeToModel(io.ComfyNode):
default="None", tooltip="Fourth LoRA"), default="None", tooltip="Fourth LoRA"),
io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01, io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 4"), 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 # Settings
io.String.Input("skip_patterns", default="", multiline=True, io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip"), tooltip="Regex patterns for layers to skip"),
@@ -926,12 +965,13 @@ class LoRAMergeToModel(io.ComfyNode):
@classmethod @classmethod
def execute(cls, base_model, lora_count, def execute(cls, base_model, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4, 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: skip_patterns, output_filename, save_dtype, device) -> io.NodeOutput:
# Build LoRA list based on count # Build LoRA list based on count
count = int(lora_count) count = int(lora_count)
lora_names = [lora_1, lora_2, lora_3, lora_4][: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][:count] lora_weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8][:count]
# Filter out "None" entries # Filter out "None" entries
valid_loras = [(name, weight) for name, weight in zip(lora_names, lora_weights) if name != "None"] valid_loras = [(name, weight) for name, weight in zip(lora_names, lora_weights) if name != "None"]
+316 -140
View File
@@ -1,3 +1,4 @@
import fnmatch
import os import os
import re import re
import torch import torch
@@ -5,13 +6,13 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import ( from .device_utils import (
estimate_model_size, prepare_for_large_operation, cleanup_after_operation estimate_model_size, prepare_for_large_operation, cleanup_after_operation
) )
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
def load_documentation_from_file(filename): def load_documentation_from_file(filename):
"""Loads documentation from a markdown file in the ../docs/ directory.""" """Loads documentation from a markdown file in the ../docs/ directory."""
@@ -24,26 +25,38 @@ def load_documentation_from_file(filename):
return f"# Documentation File Not Found\n\nPlease ensure `{filename}` exists in the `docs` directory of the custom node." return f"# Documentation File Not Found\n\nPlease ensure `{filename}` exists in the `docs` directory of the custom node."
def _compile_patterns(pattern_string): def _compile_patterns(pattern_string, glob_mode=False):
"""Compiles whitespace-separated regex patterns into a list of compiled patterns. """Compiles whitespace-separated patterns into a list.
Returns empty list if pattern_string is empty or whitespace-only. In regex mode (default): returns compiled re.Pattern objects.
Raises ValueError if any pattern is invalid regex. Raises ValueError if any pattern is invalid regex.
In glob mode: returns plain strings; fnmatch handles matching.
""" """
if not pattern_string or not pattern_string.strip(): if not pattern_string or not pattern_string.strip():
return [] return []
patterns = [] patterns = []
for pattern in pattern_string.split(): for pattern in pattern_string.split():
try: if not pattern:
patterns.append(re.compile(pattern)) continue
except re.error as e: if glob_mode:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}") patterns.append(pattern)
else:
try:
patterns.append(re.compile(pattern))
except re.error as e:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}")
return patterns return patterns
def _matches_any_pattern(key, patterns): def _matches_any_pattern(key, patterns, glob_mode=False):
"""Returns True if key matches any of the compiled regex patterns (substring search).""" """Returns True if key matches any pattern.
Glob mode: fnmatch substring match — dots are literal, * matches any sequence.
Regex mode: compiled re.Pattern substring search.
"""
if glob_mode:
return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns)
for pattern in patterns: for pattern in patterns:
if pattern.search(key): if pattern.search(key):
return True return True
@@ -76,6 +89,9 @@ class MergerLogic:
process_device = recipe_params.get('device', 'cpu') process_device = recipe_params.get('device', 'cpu')
print(f"[Merger] Preparing memory for {total_size_gb:.2f}GB merge operation...") 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)) prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device))
flush_threshold_bytes = int(total_size_gb * 1024**3 / 16)
else:
flush_threshold_bytes = 256 * 1024 * 1024 # 256MB fallback (min 8GB RAM assumed)
handlers = {} handlers = {}
for name in model_names.values(): for name in model_names.values():
@@ -83,11 +99,11 @@ class MergerLogic:
path = folder_paths.get_full_path(model_type, name) path = folder_paths.get_full_path(model_type, name)
if not path: if not path:
raise FileNotFoundError(f"Model '{name}' not found.") raise FileNotFoundError(f"Model '{name}' not found.")
handlers[name] = MemoryEfficientSafeOpen(path) handlers[name] = MemoryEfficientSafeOpen(path, low_memory=True)
primary_handler = handlers[primary_model_name] primary_handler = handlers[primary_model_name]
all_keys = primary_handler.keys() all_keys = primary_handler.keys()
metadata = primary_handler.header.get("__metadata__", {}) metadata = primary_handler.metadata()
# Convert mismatch_mode string to enum # Convert mismatch_mode string to enum
mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip') mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip')
@@ -96,8 +112,9 @@ class MergerLogic:
alignment_mode = recipe_params.get('alignment_mode', 'pad/crop') alignment_mode = recipe_params.get('alignment_mode', 'pad/crop')
# Compile filter patterns # Compile filter patterns
exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', '')) glob_mode = recipe_params.get('glob_patterns', False)
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', '')) exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''), glob_mode=glob_mode)
discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''), glob_mode=glob_mode)
# Pre-compute key differences for logging # Pre-compute key differences for logging
primary_keys = set(all_keys) primary_keys = set(all_keys)
@@ -124,65 +141,125 @@ class MergerLogic:
discarded_keys = 0 discarded_keys = 0
error_keys = [] error_keys = []
for key in tqdm(all_keys, desc="Merging layers", unit="layers"): output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
# Check discard patterns first - skip entirely # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
if _matches_any_pattern(key, discard_patterns): output_dir = folder_paths.get_folder_paths(output_folder)[-1]
discarded_keys += 1 os.makedirs(output_dir, exist_ok=True)
pbar.update(1) output_filename = recipe_params.get("output_filename")
continue output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
# Pre-load Model A's tensor with pinned memory for CUDA writer = None
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 writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
if _matches_any_pattern(key, exclude_patterns): writer.__enter__()
merged_state_dict[key] = tensor_a.to(save_torch_dtype).cpu().clone()
excluded_keys += 1
pbar.update(1)
continue
# Pass tensor_a metadata to recipes for zeros mode and fallback batch_buffer = {}
recipe_params['_tensor_a'] = tensor_a batch_bytes = 0
recipe_params['_tensor_a_shape'] = tensor_a.shape batch_key_count = 0
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
try: def _flush_batch():
recipe = calc_mode_class.create_recipe(key=key, **recipe_params) nonlocal batch_buffer, batch_bytes, batch_key_count
result = recipe.merge() if not batch_buffer:
except MissingTensorError as e: return
if mismatch_mode == MissingTensorBehavior.ERROR: writer.write_dict(batch_buffer)
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}") pbar.update(batch_key_count)
result = None batch_buffer = {}
error_keys.append(key) batch_bytes = 0
batch_key_count = 0
if recipe_params.get('force_clear_cache', True):
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Handle None result (mismatch occurred with skip mode) with torch.no_grad():
if result is None: for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
result = tensor_a # Check discard patterns first - skip entirely
skipped_keys += 1 if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode):
discarded_keys += 1
continue
if isinstance(result, dict): # Pre-load Model A's tensor with pinned memory for CUDA
for r_key, r_tensor in result.items(): cpu_tensor = primary_handler.get_tensor(key)
merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone() if process_device == 'cuda':
else: tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
# Ensure compatibility with Model A's architecture. else:
# If alignment_mode is 'pad/crop', we crop results that were padded. tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype)
# If alignment_mode is 'interpolate', resizing happened during operators. del cpu_tensor
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.to(save_torch_dtype).cpu().clone() # Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode):
t = tensor_a.detach().to(save_torch_dtype).cpu()
# Clean up tensor_a reference to allow GC batch_buffer[key] = t
del recipe_params['_tensor_a'] batch_bytes += t.numel() * t.element_size()
pbar.update(1) batch_key_count += 1
excluded_keys += 1
del tensor_a
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
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:
_flush_batch()
writer.__exit__(None, None, None)
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():
t = r_tensor.detach().to(save_torch_dtype).cpu()
batch_buffer[r_key] = t
batch_bytes += t.numel() * t.element_size()
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]
t = result.detach().to(save_torch_dtype).cpu()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
# 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 (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
_flush_batch()
writer.__exit__(None, None, None)
# Log summary # Log summary
if excluded_keys > 0: if excluded_keys > 0:
@@ -203,10 +280,9 @@ class MergerLogic:
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename") output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors") output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
save_file(merged_state_dict, output_path, metadata=metadata) if os.path.exists(output_path):
print(f"[Merger] Output saved to {output_path}")
# Cleanup after heavy operation # Cleanup after heavy operation
del merged_state_dict
cleanup_after_operation() cleanup_after_operation()
return f"{output_filename}.safetensors" return f"{output_filename}.safetensors"
@@ -230,16 +306,23 @@ class CheckpointTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_checkpoint"), io.String.Input("output_filename", default="merged_2_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -250,8 +333,9 @@ class CheckpointTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -259,10 +343,12 @@ class CheckpointTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -285,16 +371,23 @@ class ModelTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_model"), io.String.Input("output_filename", default="merged_2_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -305,8 +398,9 @@ class ModelTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -314,10 +408,12 @@ class ModelTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -340,16 +436,23 @@ class TextEncoderTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_textencoder"), io.String.Input("output_filename", default="merged_2_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -360,8 +463,9 @@ class TextEncoderTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -369,10 +473,12 @@ class TextEncoderTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -395,16 +501,23 @@ class LoRATwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_lora"), io.String.Input("output_filename", default="merged_2_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -415,8 +528,9 @@ class LoRATwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -424,10 +538,12 @@ class LoRATwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -450,16 +566,23 @@ class EmbeddingTwoMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_2_embedding"), io.String.Input("output_filename", default="merged_2_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -470,8 +593,9 @@ class EmbeddingTwoMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md') doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -479,10 +603,12 @@ class EmbeddingTwoMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
} }
model_names = {"model_a": model_a, "model_b": model_b} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -508,16 +634,23 @@ class CheckpointThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_checkpoint"), io.String.Input("output_filename", default="merged_3_checkpoint"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -528,8 +661,9 @@ class CheckpointThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -537,10 +671,12 @@ class CheckpointThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_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} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -564,16 +700,23 @@ class ModelThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_model"), io.String.Input("output_filename", default="merged_3_model"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -584,8 +727,9 @@ class ModelThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -593,10 +737,12 @@ class ModelThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_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} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -620,16 +766,23 @@ class TextEncoderThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.99, min=-10.0, max=10.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_textencoder"), io.String.Input("output_filename", default="merged_3_textencoder"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -640,8 +793,9 @@ class TextEncoderThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -649,10 +803,12 @@ class TextEncoderThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_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} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -676,16 +832,23 @@ class LoRAThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_lora"), io.String.Input("output_filename", default="merged_3_lora"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -696,8 +859,9 @@ class LoRAThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -705,10 +869,12 @@ class LoRAThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_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} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -732,16 +898,23 @@ class EmbeddingThreeMerger(io.ComfyNode):
io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), 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("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"), 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("alpha", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("beta", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("gamma", default=0.5, min=-10.0, max=10.0, step=0.01),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), io.Float.Input("delta", default=2.0, min=-10.0, max=10.0, step=0.01),
io.Float.Input("epsilon", default=0.01, min=-10.0, max=10.0, step=0.01),
io.Float.Input("zeta", default=0.0, min=-10.0, max=10.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_embedding"), io.String.Input("output_filename", default="merged_3_embedding"),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]),
io.Combo.Input("process_device", options=["cuda", "cpu"]), io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True), io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True), io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("glob_patterns", default=False,
tooltip="When True, exclude/discard patterns use glob syntax (* = any sequence, dots are literal). "
"When False (default), patterns are Python regex matched as substrings."),
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=[ outputs=[
io.String.Output(display_name="output_filename"), io.String.Output(display_name="output_filename"),
@@ -752,8 +925,9 @@ class EmbeddingThreeMerger(io.ComfyNode):
@classmethod @classmethod
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, 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, 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, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: process_device: str, exclude_patterns: str, discard_patterns: str,
glob_patterns: bool, lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md') doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY": if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc) return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -761,10 +935,12 @@ class EmbeddingThreeMerger(io.ComfyNode):
recipe_params = { recipe_params = {
"model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode,
"mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode, "mismatch_mode": mismatch_mode, "alignment_mode": alignment_mode,
"alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "epsilon": epsilon, "zeta": zeta, "seed": seed,
"output_filename": output_filename, "save_dtype": save_dtype, "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, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_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} 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) filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
+250 -4
View File
@@ -4,6 +4,8 @@ import torch.nn.functional as F
from enum import Enum from enum import Enum
from unifiedefficientloader import transfer_to_gpu_pinned
class MissingTensorBehavior(Enum): class MissingTensorBehavior(Enum):
"""Controls behavior when a tensor key is missing from a model.""" """Controls behavior when a tensor key is missing from a model."""
ERROR = "error" # Raise exception (strict mode) ERROR = "error" # Raise exception (strict mode)
@@ -127,12 +129,21 @@ class Operation:
raise NotImplementedError raise NotImplementedError
def recurse(self, operation): def recurse(self, operation):
source_tensors = [source_oper.merge() for source_oper in operation.sources] self._source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*source_tensors) return operation.oper(*self._source_tensors)
def merge(self): def merge(self):
return self.merge_func(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): 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): on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None):
@@ -165,7 +176,17 @@ class LoadTensor(Operation):
dtype = self.fallback_dtype if self.fallback_dtype else self.dtype dtype = self.fallback_dtype if self.fallback_dtype else self.dtype
return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype) return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype)
return handler.get_tensor(self.key).to(device=self.device, dtype=self.dtype) 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): class Multiply(Operation):
def __init__(self, key, alpha, *sources): def __init__(self, key, alpha, *sources):
@@ -298,6 +319,53 @@ class PowerUpOp(Operation):
return (m * delta) / keep_prob return (m * delta) / keep_prob
class DARETIESOp(Operation):
"""DARE + TIES combined operation on a single task vector (delta = b - a).
Step 1 — DARE: randomly drop `drop_rate` fraction of delta parameters,
rescale survivors by 1/(1-drop_rate).
Step 2 — TIES trim: zero out parameters below the `trim_quantile` magnitude threshold.
Step 3 — TIES elect: determine dominant sign per parameter position.
Step 4 — TIES merge: zero out parameters that disagree with the dominant sign.
For a two-model merge the TIES elect/merge step degenerates to a no-op (single vector,
sign is always self-consistent) but the trim step still reduces interference from
near-zero noisy deltas. The dominant sign filter is included for correctness and
forward compatibility with n-model extensions.
"""
def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources)
# alpha = drop_rate (DARE), beta = trim_quantile (TIES), gamma = lambda scale
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 = (b - a).float()
# --- DARE: drop and rescale ---
keep_prob = 1.0 - self.alpha
if keep_prob <= 0:
return torch.zeros_like(a)
rng = torch.Generator(device=delta.device).manual_seed(self.seed)
dare_mask = (torch.empty_like(delta).uniform_(0, 1, generator=rng) < keep_prob).float()
delta = (delta * dare_mask) / keep_prob
# --- TIES trim: zero parameters below magnitude threshold ---
if self.beta > 0:
flat_abs = delta.abs().flatten()
k = max(1, int(flat_abs.numel() * self.beta))
threshold = torch.kthvalue(flat_abs, k).values
delta = torch.where(delta.abs() < threshold, torch.zeros_like(delta), delta)
# --- TIES elect: dominant sign (trivially self-consistent for single vector) ---
dominant_sign = torch.sign(delta)
delta = torch.where(torch.sign(delta) == dominant_sign, delta, torch.zeros_like(delta))
return (a + self.gamma * delta).to(a.dtype)
class InterpolateDifference(Operation): class InterpolateDifference(Operation):
def __init__(self, key, alpha, beta, gamma, seed, *sources): def __init__(self, key, alpha, beta, gamma, seed, *sources):
super().__init__(key, *sources) super().__init__(key, *sources)
@@ -409,6 +477,10 @@ class PreloadedTensor(Operation):
def merge(self) -> torch.Tensor: def merge(self) -> torch.Tensor:
return self.tensor return self.tensor
def clean(self):
if hasattr(self, 'tensor'):
del self.tensor
class WeightSum(CalcMode): class WeightSum(CalcMode):
name = 'Weight-Sum' name = 'Weight-Sum'
description = 'A * (1 - α) + B * α. Simple linear interpolation.' description = 'A * (1 - α) + B * α. Simple linear interpolation.'
@@ -685,5 +757,179 @@ class WeightSumCutoffMode(CalcMode):
return res return res
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode()] class DARETIESMode(CalcMode):
name = 'Power-Up (DARE+TIES)'
description = 'Adds capabilities of B to A using DARE (Drop and Rescale) followed by TIES (Trim, Elect Sign). Best practice for LoRA-tuned model merging.'
models_used = ['A', 'B']
param_docs = {
'alpha': 'DARE drop rate — fraction of delta parameters randomly zeroed. Higher = more aggressive sparsification.',
'beta': 'TIES trim quantile — fraction of smallest-magnitude delta parameters zeroed after DARE. Removes residual noise.',
'gamma': 'Lambda — scale multiplier applied to the final merged delta before adding to Model A.',
'seed': 'Random seed for the DARE dropout mask.',
}
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)
op = DARETIESOp(
key,
kwargs['alpha'], # drop_rate
kwargs['beta'], # trim_quantile
kwargs['gamma'], # lambda scale
kwargs['seed'],
a, b,
)
op.alignment_mode = alignment_mode
return op
class EnhancedPowerUpOp(Operation):
def __init__(self, key, alpha, beta, gamma, delta, seed, *sources):
super().__init__(key, *sources)
self.mask_power = alpha # Mask power
self.min_keep_prob = beta # Min keep prob
self.mask_smooth = gamma # Smoothness (0 or 1)
self.delta = delta # Final multiplier
self.seed = 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_val = (b - a).float()
abs_delta = torch.abs(delta_val)
max_delta = torch.max(abs_delta)
if max_delta == 0: return torch.zeros_like(delta_val)
prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(prob, generator=rng)
interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
rescaled_delta = delta_val * interpolated_mask
return rescaled_delta
class EnhancedDARETIESOp(Operation):
def __init__(self, key, alpha, beta, gamma, delta, epsilon, zeta, seed, *sources):
super().__init__(key, *sources)
self.alpha = alpha # Unused directly for drop scale
self.beta = beta # TIES trim quantile
self.gamma = gamma # TIES lambda scale
self.mask_power = delta
self.min_keep_prob = epsilon
self.mask_smooth = zeta
self.seed = 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_val = (b - a).float()
abs_delta = torch.abs(delta_val)
max_delta = torch.max(abs_delta)
if max_delta > 0:
prob = torch.clamp((abs_delta / max_delta) ** max(self.mask_power, 0.001), min=self.min_keep_prob, max=1.0)
rng = torch.Generator(device=a.device).manual_seed(self.seed)
random_mask = torch.bernoulli(prob, generator=rng)
interpolated_mask = torch.lerp(random_mask, prob, self.mask_smooth)
delta_val = delta_val * interpolated_mask
if self.beta > 0:
flat_abs = delta_val.abs().flatten()
k = max(1, int(flat_abs.numel() * self.beta))
threshold = torch.kthvalue(flat_abs, k).values
delta_val = torch.where(delta_val.abs() < threshold, torch.zeros_like(delta_val), delta_val)
dominant_sign = torch.sign(delta_val)
delta_val = torch.where(torch.sign(delta_val) == dominant_sign, delta_val, torch.zeros_like(delta_val))
return (a + self.gamma * delta_val).to(a.dtype)
class EnhancedPowerUp(CalcMode):
name = 'Power-Up Enhanced (DARE)'
description = 'Adds capabilities of B to A using Enhanced Drop and Rescale based on relative tensor differences.'
models_used = ['A', 'B']
param_docs = {
'alpha': 'Mask power (Curve. 2.0 = quadratic).',
'beta': 'Minimum keep probability (Floor to prevent explosion).',
'gamma': 'Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale).',
'delta': 'Multiplier for the final difference.',
}
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)
deltahat = EnhancedPowerUpOp(key, kwargs['alpha'], kwargs['beta'], kwargs['gamma'], kwargs['delta'], kwargs['seed'], a, b)
deltahat.alignment_mode = alignment_mode
multiplied = Multiply(key, kwargs['delta'], deltahat)
res = Add(key, a, multiplied)
res.alignment_mode = alignment_mode
return res
class EnhancedDARETIESMode(CalcMode):
name = 'Power-Up Enhanced (DARE+TIES)'
description = 'Adds capabilities of B to A using Enhanced DARE mask followed by TIES.'
models_used = ['A', 'B']
param_docs = {
'alpha': 'Unused (Drop scale handled via mask interpolation).',
'beta': 'TIES trim quantile (fraction of smallest-magnitude delta zeroed).',
'gamma': 'TIES Lambda scale.',
'delta': 'Mask power (Curve. 2.0 = quadratic).',
'epsilon': 'Minimum keep probability (Floor to prevent explosion).',
'zeta': 'Mask smoothness factor (0.0 = pure dropout, 1.0 = soft scale).',
}
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)
op = EnhancedDARETIESOp(
key,
kwargs.get('alpha', 0.5),
kwargs.get('beta', 0.2),
kwargs.get('gamma', 1.0),
kwargs.get('delta', 0.5),
kwargs.get('epsilon', 0.5),
kwargs.get('zeta', 0.5),
kwargs['seed'],
a, b,
)
op.alignment_mode = alignment_mode
return op
TWO_MODEL_MODES = [WeightSum(), InterpDifference(), PowerUp(), DARETIESMode(), SVDMode(), ManEnhInterpDifference(), AutoEnhInterpDifference(), WeightSumCutoffMode(), EnhancedPowerUp(), EnhancedDARETIESMode()]
THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()] THREE_MODEL_MODES = [AddDifference(), TrainDifference(), Extract(), AddDisimilarity()]
+3 -369
View File
@@ -1,373 +1,7 @@
""" """
Memory-efficient safetensors utilities with parallel I/O and pinned memory support. Utilities for model merging operations.
Based on benchmarking, parallel reading with ThreadPoolExecutor provides 2-4x speedup
for large models compared to sequential reads.
Worker count is automatically optimized based on device capabilities when workers=None.
Pinned memory enables faster CPU→GPU transfers (2-3x improvement).
""" """
import os from typing import Dict, List
import mmap
import json
import struct
import torch
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Optional, Any, Tuple
# Import ComfyUI's memory pinning functions if available
try:
from comfy.model_management import pin_memory, unpin_memory, is_device_cuda
PINNING_AVAILABLE = True
except ImportError:
PINNING_AVAILABLE = False
def pin_memory(tensor): return False
def unpin_memory(tensor): return False
def is_device_cuda(device): return str(device).startswith('cuda')
class PinnedTensor:
"""Context manager for automatic pinned memory management.
Usage:
with PinnedTensor(cpu_tensor) as pinned:
gpu_tensor = pinned.to('cuda')
# Memory automatically unpinned after transfer
"""
def __init__(self, tensor: torch.Tensor, auto_pin: bool = True):
self.tensor = tensor
self.is_pinned = False
self.auto_pin = auto_pin and PINNING_AVAILABLE
def __enter__(self) -> torch.Tensor:
if self.auto_pin and self.tensor.device.type == 'cpu':
self.is_pinned = pin_memory(self.tensor)
return self.tensor
def __exit__(self, exc_type, exc_val, exc_tb):
if self.is_pinned:
unpin_memory(self.tensor)
self.is_pinned = False
def transfer_to_gpu_pinned(
tensor: torch.Tensor,
device: str = 'cuda',
dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
"""Transfer tensor to GPU using pinned memory for faster transfer.
Args:
tensor: CPU tensor to transfer
device: Target GPU device
dtype: Optional dtype conversion
Returns:
Tensor on GPU device
"""
if not PINNING_AVAILABLE or not is_device_cuda(torch.device(device)):
# Fall back to regular transfer
if dtype:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
# Pin memory for faster transfer
was_pinned = pin_memory(tensor)
try:
if dtype:
result = tensor.to(device=device, dtype=dtype, non_blocking=True)
else:
result = tensor.to(device=device, non_blocking=True)
# Sync to ensure transfer is complete before unpinning
if was_pinned:
torch.cuda.current_stream().synchronize()
return result
finally:
if was_pinned:
unpin_memory(tensor)
class MemoryEfficientSafeOpen:
"""Memory-efficient safetensors file reader with mmap, parallel I/O, and pinned memory support.
Features:
- mmap mode: Zero-copy tensor access via memory-mapped file
- low_memory mode: Direct file reads to minimize OS page cache usage
- Parallel loading: Multi-threaded tensor reads for 2-4x speedup
- Sorted batch reads: Keys sorted by file offset for sequential I/O
- Auto-optimized workers: Adjusts parallelism based on device capabilities
- Pinned memory: Faster CPU→GPU transfers when transferring to CUDA
Args:
filename: Path to safetensors file
device: Target device (default 'cpu')
mmap_mode: Use memory-mapped file for zero-copy (default True)
low_memory: Use direct file reads to minimize memory (overrides mmap_mode)
"""
def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True, low_memory: bool = False):
self.filename = filename
self.device = device
# low_memory mode forces mmap off to avoid OS page caching
self.low_memory = low_memory
self.mmap_mode = mmap_mode and not low_memory
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
self.mmap_obj = None
if self.mmap_mode:
self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.mmap_obj:
self.mmap_obj.close()
self.file.close()
def keys(self) -> List[str]:
"""Return all tensor keys (excluding metadata)."""
return [k for k in self.header.keys() if k != "__metadata__"]
def metadata(self) -> Dict[str, str]:
"""Return file metadata."""
return self.header.get("__metadata__", {})
def keys_sorted_by_offset(self) -> List[str]:
"""Return keys sorted by file offset for optimal sequential I/O."""
keys_with_offsets = []
for key in self.keys():
offset = self.header[key]["data_offsets"][0]
keys_with_offsets.append((key, offset))
keys_with_offsets.sort(key=lambda x: x[1])
return [k for k, _ in keys_with_offsets]
def get_tensor(self, key: str) -> torch.Tensor:
"""Load a single tensor by key."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
if self.mmap_mode and self.mmap_obj:
if offset_start != offset_end:
start = self.header_size + 8 + offset_start
end = self.header_size + 8 + offset_end
tensor_bytes = memoryview(self.mmap_obj)[start:end]
else:
tensor_bytes = None
else:
# Non-mmap mode: use pre-allocated bytearray with readinto for minimal copies
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
# Pre-allocate writable buffer and read directly into it
tensor_bytes = bytearray(offset_end - offset_start)
self.file.readinto(tensor_bytes)
return self._deserialize_tensor(tensor_bytes, metadata)
def get_tensor_to_gpu(
self,
key: str,
device: str = 'cuda',
dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
"""Load a tensor and transfer to GPU using pinned memory for faster transfer.
Args:
key: Tensor key to load
device: Target GPU device (default 'cuda')
dtype: Optional dtype conversion
Returns:
Tensor on GPU device
"""
cpu_tensor = self.get_tensor(key)
return transfer_to_gpu_pinned(cpu_tensor, device, dtype)
def get_tensors_parallel(
self,
keys: Optional[List[str]] = None,
workers: Optional[int] = None,
sort_by_offset: bool = True
) -> Dict[str, torch.Tensor]:
"""Load multiple tensors in parallel using ThreadPoolExecutor.
Args:
keys: List of keys to load (default: all keys)
workers: Number of worker threads (default: auto-calculated based on device)
sort_by_offset: Sort keys by file offset before loading (default: True)
Returns:
Dict mapping key -> tensor
"""
if keys is None:
keys = self.keys()
# Auto-calculate workers based on device capabilities
if workers is None:
try:
from .device_utils import get_optimal_workers, estimate_model_size
model_size = estimate_model_size(self.filename)
workers = get_optimal_workers(model_size)
except ImportError:
workers = min(4, os.cpu_count() or 4)
if sort_by_offset:
# Sort by offset for better I/O patterns
keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys]
keys_with_offsets.sort(key=lambda x: x[1])
keys = [k for k, _ in keys_with_offsets]
result = {}
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(
_read_tensor_from_file,
self.filename,
key,
self.header[key],
self.header_size
): key
for key in keys
}
for future in as_completed(futures):
key = futures[future]
try:
tensor = future.result()
if tensor is not None:
result[key] = tensor
except Exception as e:
print(f"[MemoryEfficientSafeOpen] Failed to load {key}: {e}")
return result
def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
"""Load a tensor and decode as JSON dict (for config tensors)."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
if offset_start == offset_end:
return {}
if self.mmap_mode and self.mmap_obj:
start = self.header_size + 8 + offset_start
end = self.header_size + 8 + offset_end
tensor_bytes = bytes(self.mmap_obj[start:end])
else:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
return 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]
header_json = f.read(header_size).decode("utf-8")
return json.loads(header_json), header_size
def _deserialize_tensor(self, tensor_bytes, metadata):
dtype_str = metadata["dtype"]
shape = metadata["shape"]
dtype = self._get_torch_dtype(dtype_str)
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
# Avoid extra copy if already a bytearray (low_memory mode)
if isinstance(tensor_bytes, bytearray):
byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
else:
# mmap memoryview needs copy to create writable tensor
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
return self._convert_float8(byte_tensor, dtype_str, shape)
return byte_tensor.view(dtype).reshape(shape)
@staticmethod
def _get_torch_dtype(dtype_str):
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)
if dtype is None:
raise ValueError(f"Unsupported dtype: {dtype_str}")
return dtype
@staticmethod
def _convert_float8(byte_tensor, dtype_str, shape):
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)
else:
raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")
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]: def get_union_keys(*handlers) -> List[str]:
"""Returns the union of all tensor keys across multiple model handlers.""" """Returns the union of all tensor keys across multiple model handlers."""
@@ -386,4 +20,4 @@ def get_intersection_keys(*handlers) -> List[str]:
all_keys = set(handlers[0].keys()) all_keys = set(handlers[0].keys())
for handler in handlers[1:]: for handler in handlers[1:]:
all_keys &= set(handler.keys()) all_keys &= set(handler.keys())
return sorted(all_keys) return sorted(all_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])
+23 -22
View File
@@ -4,12 +4,12 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open from safetensors import safe_open
from .utils import convert_pt_to_safetensors 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 from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
use_regex: bool, output_filename: str) -> str: use_regex: bool, output_filename: str) -> str:
@@ -39,32 +39,33 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
if not patterns: if not patterns:
raise ValueError("No keys/patterns provided to prune.") raise ValueError("No keys/patterns provided to prune.")
# Stream tensors and filter on the fly
pruned_tensors = {}
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
all_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(all_keys))
for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = False
if use_regex:
if any(re.search(pattern, key) for pattern in patterns):
is_match = True
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match:
pruned_tensors[key] = handler.get_tensor(key)
pbar.update(1)
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet # 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] model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(pruned_tensors, output_path, metadata) # Stream tensors, filter on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
with MemoryEfficientSafeOpen(model_path_to_load) as handler:
all_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(all_keys))
for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = False
if use_regex:
if any(re.search(pattern, key) for pattern in patterns):
is_match = True
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match:
writer.write(key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation # Cleanup after operation
del pruned_tensors
cleanup_after_operation() cleanup_after_operation()
return output_path return output_path
+16 -14
View File
@@ -3,12 +3,13 @@ import folder_paths
import comfy.utils import comfy.utils
from tqdm import tqdm from tqdm import tqdm
from comfy_api.latest import io from comfy_api.latest import io
from safetensors.torch import save_file
from safetensors import safe_open from safetensors import safe_open
from .utils import convert_pt_to_safetensors 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 from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
def _rename_keys(model_name: str, model_type: str, old_keys_str: str, def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
new_keys_str: str, output_filename: str) -> str: new_keys_str: str, output_filename: str) -> str:
@@ -41,24 +42,25 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
key_map = dict(zip(old_keys, new_keys)) key_map = dict(zip(old_keys, new_keys))
# 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] = handler.get_tensor(key)
pbar.update(1)
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet # 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] model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
save_file(renamed_tensors, output_path, metadata) # Stream tensors, rename on the fly, write immediately
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
try:
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)
writer.write(new_key, handler.get_tensor(key).contiguous())
pbar.update(1)
finally:
writer.__exit__(None, None, None)
# Cleanup after operation # Cleanup after operation
del renamed_tensors
cleanup_after_operation() cleanup_after_operation()
return output_path return output_path
+16 -1
View File
@@ -1,9 +1,24 @@
[project] [project]
name = "ComfyUI-ModelUtils" name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files." description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.2.7" version = "0.4.3"
license = { file = "LICENSE" } license = { file = "LICENSE" }
classifiers = [
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"unifiedefficientloader==0.4.4"
]
[project.urls] [project.urls]
Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils" Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils"
# Used by Comfy Registry https://comfyregistry.org # Used by Comfy Registry https://comfyregistry.org
+5 -1
View File
@@ -1 +1,5 @@
# No external dependencies required requests
Pillow
mutagen
unifiedefficientloader>=0.5.2
av