nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(paths): standardize model artifact outputs

Browse Source
This commit is contained in:
silveroxides
2026-08-14 19:04:05 +02:00
parent f5839f706c
commit 349e2155cd
16 changed files with 362 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/IMPLEMENTATION_INVARIANTS.md
+38
View File
@@ -18,6 +18,23 @@ This file is the authoritative engineering contract for model and LoRA tensor I/
The shared implementations in `nodes/uel_io.py`, `nodes/extraction_stream.py`, and `nodes/quantization_guard.py` are the required starting points. New node-local stream or writer abstractions require a demonstrated capability that the shared implementation cannot provide.
## Persistent Artifact Paths and Outputs
1. Persistent model artifacts MUST use `nodes/artifact_paths.py` to construct
their write path from `folder_paths.models_dir` and an explicit canonical
category: `checkpoints`, `diffusion_models`, `text_encoders`, `loras`, or
`embeddings`.
2. Registered search paths, additional paths, and legacy aliases such as
`unet` and `clip` MUST NOT become write roots based on their list position.
3. Nested output names MUST be preserved, exactly one `.safetensors` extension
MUST be applied, and path traversal outside the category root MUST fail.
4. Public filename outputs MUST return the category-relative name with forward
slashes. They MUST NOT return an absolute filesystem path.
5. Loader-compatible filename outputs MUST use `io.AnyType`. Reports and
documentation MUST remain separate string outputs.
6. Persistent model artifacts MUST NOT be written beneath ComfyUI's general
output directory.
## LoRA Alpha Normalization for Merge Operations
1. Every operation that combines a LoRA with another LoRA or applies a LoRA to a model MUST account for per-layer alpha before weighting, padding, concatenation, consensus processing, DARE/TIES processing, or delta application.
@@ -37,6 +54,25 @@ The shared implementations in `nodes/uel_io.py`, `nodes/extraction_stream.py`, a
9. A low-bit alpha scalar is valid when both factors use supported floating storage. If an alpha-bearing down or up factor uses unsupported low-bit storage, the operation MUST fail with an actionable error because byte-exact preservation cannot also produce an alpha-free normalized result.
10. Format metadata may describe that normalization occurred, but metadata MUST NOT be used as a substitute for applying the mathematical transformation.
### CWB Mixed-Rank LoRA Alignment
1. Genuine LoRA rank validity MUST come from each factor pair's original shape
or explicit source metadata. Tensor values MUST NOT be inspected to decide
whether a component is padding because a genuine component may be zero.
2. Structural rank padding MUST NOT participate in CWB alignment, consensus,
similarity weighting, fallback behavior, or norm rescaling.
3. Similarity alignment between rank `R` and rank `r` genuine pairs MUST compare
the full rectangular `R x r` component space. A smaller source component may
match any component in the maximum-rank reference.
4. Index alignment MUST pair only shared genuine indices. Higher reference
indices remain singleton groups unless another genuine source contributes.
5. An explicit zero inserted by `mismatch_mode=zeros` is a semantic contributor
and MUST remain distinct from structural rank absence.
6. Down rows and corresponding up columns MUST share validity, alignment, sign
correction, contributor membership, and output position.
7. Output rank MUST remain the largest genuine input rank. Equal maximum-rank
ties retain source order.
The shared implementation in `nodes/lora_alpha.py` is mandatory for merge paths. A format-specific implementation may only replace it when the represented delta differs mathematically and that difference is documented and tested.
## Required Verification
@@ -48,6 +84,8 @@ Any implementation that reads, writes, merges, resizes, extracts, or applies mod
- normalization before padding or concatenation;
- alpha-free output;
- mixed factor ranks and dtypes;
- full rectangular CWB matching without structural-padding contributors;
- genuine zero-valued components versus explicit mismatch-zero contributors;
- low-bit scalar alpha acceptance and unsupported low-bit factor rejection;
- stream closure, `mark_processed`, and reference release on success and failure;
- atomic destination preservation on writer failure.
nodes/artifact_paths.py
+36
View File
@@ -0,0 +1,36 @@
import os
import folder_paths
CANONICAL_MODEL_CATEGORIES = frozenset(
{"checkpoints", "diffusion_models", "text_encoders", "loras", "embeddings"}
)
def canonical_model_artifact_path(category: str, output_filename: str) -> tuple[str, str]:
"""Return the canonical write path and category-relative loader value."""
if category not in CANONICAL_MODEL_CATEGORIES:
raise ValueError(f"Unsupported model artifact category: {category}")
relative_name = output_filename.strip().replace("\\", "/")
if not relative_name:
raise ValueError("Output filename must not be empty.")
if relative_name.startswith("/") or os.path.isabs(relative_name):
raise ValueError("Output filename must be relative to its model category.")
while relative_name.lower().endswith(".safetensors"):
relative_name = relative_name[:-12]
if not relative_name:
raise ValueError("Output filename must include a name before its extension.")
relative_name = f"{relative_name}.safetensors"
category_root = os.path.abspath(os.path.join(folder_paths.models_dir, category))
output_path = os.path.abspath(
os.path.join(category_root, *relative_name.split("/"))
)
if os.path.commonpath((category_root, output_path)) != category_root:
raise ValueError("Output filename must stay inside its model category directory.")
os.makedirs(os.path.dirname(output_path), exist_ok=True)
output_name = os.path.relpath(output_path, category_root).replace(os.sep, "/")
return output_path, output_name
nodes/dora_extract_wd.py
+18 -20
View File
@@ -5,7 +5,6 @@ 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
@@ -20,6 +19,7 @@ from .device_utils import (
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd
from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
@@ -690,11 +690,9 @@ def extract_dora_from_files(
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> str:
def _build_lora_output_path(output_filename: str) -> tuple[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")
return canonical_model_artifact_path("loras", output_filename)
# =============================================================================
@@ -798,7 +796,7 @@ class DoRAExtractFixed(io.ComfyNode):
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -809,7 +807,7 @@ class DoRAExtractFixed(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
@@ -818,7 +816,7 @@ class DoRAExtractFixed(io.ComfyNode):
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRAExtractRatio(io.ComfyNode):
@@ -845,7 +843,7 @@ class DoRAExtractRatio(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -856,7 +854,7 @@ class DoRAExtractRatio(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
@@ -866,7 +864,7 @@ class DoRAExtractRatio(io.ComfyNode):
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRAExtractQuantile(io.ComfyNode):
@@ -893,7 +891,7 @@ class DoRAExtractQuantile(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -904,7 +902,7 @@ class DoRAExtractQuantile(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
@@ -914,7 +912,7 @@ class DoRAExtractQuantile(io.ComfyNode):
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRAExtractKnee(io.ComfyNode):
@@ -939,7 +937,7 @@ class DoRAExtractKnee(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -951,7 +949,7 @@ class DoRAExtractKnee(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
@@ -962,7 +960,7 @@ class DoRAExtractKnee(io.ComfyNode):
knee_probe_offset=knee_probe_offset,
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRAExtractFrobenius(io.ComfyNode):
@@ -989,7 +987,7 @@ class DoRAExtractFrobenius(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -1000,7 +998,7 @@ class DoRAExtractFrobenius(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
@@ -1010,4 +1008,4 @@ class DoRAExtractFrobenius(io.ComfyNode):
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
nodes/dora_learned_wd.py
+15 -15
View File
@@ -590,7 +590,7 @@ class DoRALearnedExtractFixed(io.ComfyNode):
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")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -601,7 +601,7 @@ class DoRALearnedExtractFixed(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
@@ -610,7 +610,7 @@ class DoRALearnedExtractFixed(io.ComfyNode):
lazy_load, force_clear_cache, glob_skip_patterns
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRALearnedExtractRatio(io.ComfyNode):
@@ -634,7 +634,7 @@ class DoRALearnedExtractRatio(io.ComfyNode):
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -645,7 +645,7 @@ class DoRALearnedExtractRatio(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
@@ -655,7 +655,7 @@ class DoRALearnedExtractRatio(io.ComfyNode):
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRALearnedExtractQuantile(io.ComfyNode):
@@ -679,7 +679,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode):
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -690,7 +690,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
@@ -700,7 +700,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode):
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRALearnedExtractKnee(io.ComfyNode):
@@ -723,7 +723,7 @@ class DoRALearnedExtractKnee(io.ComfyNode):
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -735,7 +735,7 @@ class DoRALearnedExtractKnee(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
@@ -745,7 +745,7 @@ class DoRALearnedExtractKnee(io.ComfyNode):
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class DoRALearnedExtractFrobenius(io.ComfyNode):
@@ -769,7 +769,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode):
*_get_learned_inputs(),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -780,7 +780,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_dora_learned_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
@@ -790,4 +790,4 @@ class DoRALearnedExtractFrobenius(io.ComfyNode):
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
nodes/dtype_conversion.py
+7 -13
View File
@@ -15,6 +15,7 @@ from unifiedefficientloader import MemoryEfficientSafeOpen
from .device_utils import cleanup_after_operation
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
TARGET_DTYPES = {
@@ -43,15 +44,8 @@ def _is_excluded(key: str, patterns: list[re.Pattern]) -> bool:
return any(pattern.search(key) for pattern in patterns)
def _output_path(output_filename: str) -> str:
name = output_filename.strip()
if not name:
raise ValueError("Output filename must not be empty.")
if name.lower().endswith(".safetensors"):
name = name[:-12]
output_dir = os.path.join(folder_paths.models_dir, "diffusion_models")
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, f"{name}.safetensors")
def _output_path(output_filename: str) -> tuple[str, str]:
return canonical_model_artifact_path("diffusion_models", output_filename)
def convert_diffusion_model_dtype(
@@ -68,7 +62,7 @@ def convert_diffusion_model_dtype(
source_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_name
)
output_path = _output_path(output_filename)
output_path, output_name = _output_path(output_filename)
if os.path.normcase(os.path.abspath(source_path)) == os.path.normcase(
os.path.abspath(output_path)
):
@@ -127,7 +121,7 @@ def convert_diffusion_model_dtype(
(
"DIFFUSION MODEL DTYPE CONVERSION",
f"Input: {model_name}",
f"Output: {output_path}",
f"Output: {output_name}",
f"Target dtype: {target_dtype}",
f"Total tensors: {total}",
f"Converted floating tensors: {counts['converted']}",
@@ -137,7 +131,7 @@ def convert_diffusion_model_dtype(
f"Original dtype counts: {dtype_summary or 'None'}",
)
)
return output_path, report
return output_name, report
class DiffusionModelDtypeConversion(io.ComfyNode):
@@ -177,7 +171,7 @@ class DiffusionModelDtypeConversion(io.ComfyNode):
),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="conversion_report"),
],
is_output_node=True,
nodes/lora_extract_svd.py
+18 -20
View File
@@ -5,7 +5,6 @@ Native implementation of Low-Rank Adapter extraction with multiple rank selectio
No external dependencies required.
"""
import fnmatch
import os
import re
import torch
import torch.linalg as linalg
@@ -20,6 +19,7 @@ from .device_utils import (
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
from .quantization_guard import inspect_low_bit_input
from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd
from .extraction_stream import lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
@@ -715,11 +715,9 @@ def extract_lora_from_files(
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> str:
def _build_lora_output_path(output_filename: str) -> tuple[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")
return canonical_model_artifact_path("loras", output_filename)
# =============================================================================
@@ -833,7 +831,7 @@ class LoRAExtractFixed(io.ComfyNode):
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -845,7 +843,7 @@ class LoRAExtractFixed(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
@@ -854,7 +852,7 @@ class LoRAExtractFixed(io.ComfyNode):
lazy_load, force_clear_cache, glob_skip_patterns, include_1d_diffs
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class LoRAExtractRatio(io.ComfyNode):
@@ -881,7 +879,7 @@ class LoRAExtractRatio(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -893,7 +891,7 @@ class LoRAExtractRatio(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_lora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
@@ -904,7 +902,7 @@ class LoRAExtractRatio(io.ComfyNode):
knee_probe_offset=probe_offset,
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class LoRAExtractQuantile(io.ComfyNode):
@@ -931,7 +929,7 @@ class LoRAExtractQuantile(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -943,7 +941,7 @@ class LoRAExtractQuantile(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_lora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
@@ -954,7 +952,7 @@ class LoRAExtractQuantile(io.ComfyNode):
knee_probe_offset=probe_offset,
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class LoRAExtractKnee(io.ComfyNode):
@@ -979,7 +977,7 @@ class LoRAExtractKnee(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -992,7 +990,7 @@ class LoRAExtractKnee(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_lora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
@@ -1003,7 +1001,7 @@ class LoRAExtractKnee(io.ComfyNode):
knee_probe_offset=knee_probe_offset,
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class LoRAExtractFrobenius(io.ComfyNode):
@@ -1030,7 +1028,7 @@ class LoRAExtractFrobenius(io.ComfyNode):
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -1042,7 +1040,7 @@ class LoRAExtractFrobenius(io.ComfyNode):
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)
output_path, output_name = _build_lora_output_path(output_filename)
extract_lora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
@@ -1053,4 +1051,4 @@ class LoRAExtractFrobenius(io.ComfyNode):
knee_probe_offset=probe_offset,
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
nodes/lora_merger.py
+11 -11
View File
@@ -24,6 +24,7 @@ from .lora_resize import (
from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor
from .uel_io import AsyncTensorCursor as _AsyncTensorCursor, atomic_uel_writer
from .lora_alpha import normalize_lora_pair
from .artifact_paths import canonical_model_artifact_path
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
@@ -752,11 +753,7 @@ def _run_multi_lora_merge(
for handler, keys in zip(handlers, stream_keys)
]
output_dir = os.path.join(folder_paths.models_dir, "loras")
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(
output_dir, f"{output_filename.strip()}.safetensors"
)
output_path, _ = canonical_model_artifact_path("loras", output_filename)
metadata = {
"ss_training_comment": f"Merged {len(lora_paths)} LoRAs via {strategy}",
"ss_network_module": "networks.lora",
@@ -947,7 +944,7 @@ class LoRAMultiMerge(io.ComfyNode):
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Include and merge 1D direct-diff tensors as FP32."),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -989,7 +986,8 @@ class LoRAMultiMerge(io.ComfyNode):
include_1d_diffs=include_1d_diffs,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
class LoRAMultiMergeDARE(io.ComfyNode):
@@ -1036,7 +1034,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Include and merge 1D direct-diff tensors as FP32."),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -1081,7 +1079,8 @@ class LoRAMultiMergeDARE(io.ComfyNode):
include_1d_diffs=include_1d_diffs,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
def merge_multi_loras_dare(
lora_paths: List[str],
@@ -1156,7 +1155,7 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Include and merge 1D direct-diff tensors as FP32."),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -1203,7 +1202,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode):
include_1d_diffs=include_1d_diffs,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
def merge_multi_loras_dare_enhanced(
nodes/lora_rename.py
+4 -4
View File
@@ -8,6 +8,7 @@ from .device_utils import estimate_model_size, prepare_for_large_operation, clea
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
REVERSE_RENAME_DICT = {
"transformer_blocks.": "blocks.",
@@ -56,8 +57,7 @@ def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filena
model_size_gb = estimate_model_size(model_path_to_load)
prepare_for_large_operation(model_size_gb * 1.2)
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
# Stream tensors, rename on the fly, write immediately
with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
@@ -87,7 +87,7 @@ def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filena
# Cleanup after operation
cleanup_after_operation()
return output_path
return output_name
class AnimaLoraRename(io.ComfyNode):
@@ -107,7 +107,7 @@ class AnimaLoraRename(io.ComfyNode):
io.String.Input("output_filename", default="anima_lora_non_diffusers", tooltip="Output filename without extension, written under ComfyUI's LoRA directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
is_output_node=True,
)
nodes/lora_resize.py
+27 -20
View File
@@ -18,6 +18,7 @@ from .device_utils import estimate_model_size, prepare_for_large_operation, clea
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .uel_io import atomic_uel_writer, stream_work_units
from .artifact_paths import canonical_model_artifact_path
from .lora_alpha import normalize_lora_pair
from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor
from typing import Optional, Dict, Tuple, List
@@ -689,9 +690,7 @@ def resize_lora_file(
metadata["ss_network_dim"] = "Dynamic"
metadata["ss_network_alpha"] = "Dynamic"
output_dir = os.path.join(folder_paths.models_dir, "loras")
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
output_path, _ = canonical_model_artifact_path("loras", output_filename)
fro_list = []
cpu_fallbacks = 0
@@ -918,7 +917,7 @@ class LoRAResizeFixed(io.ComfyNode):
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."),
io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -927,11 +926,12 @@ class LoRAResizeFixed(io.ComfyNode):
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(
resize_lora_file(
lora_path, new_rank, None, None, device, dtype, output_filename,
force_clear_cache=force_clear_cache,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
class LoRAResizeRatio(io.ComfyNode):
@@ -956,7 +956,7 @@ class LoRAResizeRatio(io.ComfyNode):
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."),
io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -965,11 +965,12 @@ class LoRAResizeRatio(io.ComfyNode):
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(
resize_lora_file(
lora_path, max_rank, "sv_ratio", ratio, device, dtype, output_filename,
force_clear_cache=force_clear_cache,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
class LoRAResizeFrobenius(io.ComfyNode):
@@ -996,7 +997,7 @@ class LoRAResizeFrobenius(io.ComfyNode):
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."),
io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -1005,12 +1006,13 @@ class LoRAResizeFrobenius(io.ComfyNode):
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(
resize_lora_file(
lora_path, max_rank, "sv_fro", target, device, dtype, output_filename,
force_clear_cache=force_clear_cache,
min_rank=min_rank,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
class LoRAResizeCumulative(io.ComfyNode):
@@ -1035,7 +1037,7 @@ class LoRAResizeCumulative(io.ComfyNode):
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."),
io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -1044,11 +1046,12 @@ class LoRAResizeCumulative(io.ComfyNode):
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(
resize_lora_file(
lora_path, max_rank, "sv_cumulative", target, device, dtype, output_filename,
force_clear_cache=force_clear_cache,
)
return io.NodeOutput(path)
_, output_name = canonical_model_artifact_path("loras", output_filename)
return io.NodeOutput(output_name)
@@ -1254,9 +1257,9 @@ def merge_loras_to_model(
base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}"
# 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")
output_path, _ = canonical_model_artifact_path(
"diffusion_models", output_filename
)
base_outcomes = {
"PATCHED": [],
@@ -1903,7 +1906,10 @@ class LoRAMergeToModel(io.ComfyNode):
io.Boolean.Input("include_1d_diffs", default=False,
tooltip="Apply 1D direct-diff tensors as FP32. Disabled preserves prior behavior."),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="merge_report"),
],
is_output_node=True,
)
@@ -1942,4 +1948,5 @@ class LoRAMergeToModel(io.ComfyNode):
include_1d_diffs=include_1d_diffs,
return_report=True,
)
return io.NodeOutput(f"{path}\n\n{report}")
_, output_name = canonical_model_artifact_path("diffusion_models", output_filename)
return io.NodeOutput(output_name, report)
nodes/merger.py
+14 -17
View File
@@ -17,6 +17,7 @@ from .quantization_guard import inspect_low_bit_input, write_preserved_tensor
from .lora_resize import is_direct_diff_key, parse_lora_layers
from .uel_io import atomic_uel_writer, stream_work_units
from .lora_alpha import lora_alpha_scale
from .artifact_paths import canonical_model_artifact_path
def load_documentation_from_file(filename):
@@ -175,10 +176,8 @@ class MergerLogic:
error_keys = []
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
output_dir = os.path.join(folder_paths.models_dir, output_folder)
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
output_path, output_name = canonical_model_artifact_path(output_folder, output_filename)
alpha_normalization = {}
alpha_keys = set()
@@ -373,16 +372,14 @@ class MergerLogic:
handler.__exit__(None, None, None)
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
output_dir = os.path.join(folder_paths.models_dir, output_folder)
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
output_path, output_name = canonical_model_artifact_path(output_folder, output_filename)
if os.path.exists(output_path):
print(f"[Merger] Output saved to {output_path}")
# Cleanup after heavy operation
cleanup_after_operation()
return f"{output_filename}.safetensors"
return output_name
# --- Two-Model Merger Nodes ---
@@ -423,7 +420,7 @@ class CheckpointTwoMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -489,7 +486,7 @@ class ModelTwoMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -555,7 +552,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -623,7 +620,7 @@ class LoRATwoMerger(io.ComfyNode):
tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -691,7 +688,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -760,7 +757,7 @@ class CheckpointThreeMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -827,7 +824,7 @@ class ModelThreeMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -894,7 +891,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -963,7 +960,7 @@ class LoRAThreeMerger(io.ComfyNode):
tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
@@ -1032,7 +1029,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."),
],
outputs=[
io.String.Output(display_name="output_filename"),
io.AnyType.Output(display_name="output_filename"),
io.String.Output(display_name="documentation"),
],
)
nodes/prunekeys.py
+8 -8
View File
@@ -9,6 +9,7 @@ from .device_utils import estimate_model_size, prepare_for_large_operation, clea
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
@@ -36,8 +37,7 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
raise ValueError("No keys/patterns provided to prune.")
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
# Stream tensors, filter on the fly, write immediately
with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
@@ -65,7 +65,7 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
# Cleanup after operation
cleanup_after_operation()
return output_path
return output_name
class ModelPruneKeys(io.ComfyNode):
@@ -87,7 +87,7 @@ class ModelPruneKeys(io.ComfyNode):
io.String.Input("output_filename", default="pruned_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -118,7 +118,7 @@ class TextEncoderPruneKeys(io.ComfyNode):
io.String.Input("output_filename", default="pruned_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -149,7 +149,7 @@ class LoRAPruneKeys(io.ComfyNode):
io.String.Input("output_filename", default="pruned_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -180,7 +180,7 @@ class CheckpointPruneKeys(io.ComfyNode):
io.String.Input("output_filename", default="pruned_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -211,7 +211,7 @@ class EmbeddingPruneKeys(io.ComfyNode):
io.String.Input("output_filename", default="pruned_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
nodes/renamekeys.py
+8 -9
View File
@@ -8,6 +8,7 @@ from .device_utils import estimate_model_size, prepare_for_large_operation, clea
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
@@ -38,9 +39,7 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
key_map = dict(zip(old_keys, new_keys))
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
model_dir = folder_paths.get_folder_paths(model_type)[-1]
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
output_path, output_name = canonical_model_artifact_path(model_type, output_filename)
# Stream tensors, rename on the fly, write immediately
with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler:
@@ -61,7 +60,7 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str,
# Cleanup after operation
cleanup_after_operation()
return output_path
return output_name
class ModelRenameKeys(io.ComfyNode):
@@ -83,7 +82,7 @@ class ModelRenameKeys(io.ComfyNode):
io.String.Input("output_filename", default="renamed_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -114,7 +113,7 @@ class TextEncoderRenameKeys(io.ComfyNode):
io.String.Input("output_filename", default="renamed_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -145,7 +144,7 @@ class LoRARenameKeys(io.ComfyNode):
io.String.Input("output_filename", default="renamed_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -176,7 +175,7 @@ class CheckpointRenameKeys(io.ComfyNode):
io.String.Input("output_filename", default="renamed_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
@@ -207,7 +206,7 @@ class EmbeddingRenameKeys(io.ComfyNode):
io.String.Input("output_filename", default="renamed_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."),
],
outputs=[
io.String.Output(display_name="output_path"),
io.AnyType.Output(display_name="output_path"),
],
)
nodes/text_encoder_extract.py
+33 -35
View File
@@ -4,7 +4,6 @@ Text Encoder LoRA and DoRA Extraction.
Specifically targets extracting text encoder deltas as LoRA/DoRA with custom mapping rules
for layer name prefixes and PEFT weight formatting.
"""
import os
import torch
import folder_paths
import comfy.utils
@@ -17,6 +16,7 @@ from .device_utils import (
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
from .extraction_stream import dora_difference, lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
# Import SVD rank utilities from standard lora extract file to avoid duplication
@@ -267,11 +267,9 @@ def extract_te_from_files(
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> str:
def _build_lora_output_path(output_filename: str) -> tuple[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")
return canonical_model_artifact_path("loras", output_filename)
def _get_te_model_inputs():
@@ -324,7 +322,7 @@ class TextEncoderLoRAExtractFixed(io.ComfyNode):
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -335,7 +333,7 @@ class TextEncoderLoRAExtractFixed(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
@@ -345,7 +343,7 @@ class TextEncoderLoRAExtractFixed(io.ComfyNode):
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractRatio(io.ComfyNode):
@@ -365,7 +363,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -376,7 +374,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
@@ -387,7 +385,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractQuantile(io.ComfyNode):
@@ -407,7 +405,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -418,7 +416,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
@@ -429,7 +427,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractKnee(io.ComfyNode):
@@ -448,7 +446,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -460,7 +458,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
@@ -471,7 +469,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
@@ -491,7 +489,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -502,7 +500,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
@@ -513,7 +511,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
# =============================================================================
@@ -535,7 +533,7 @@ class TextEncoderDoRAExtractFixed(io.ComfyNode):
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -546,7 +544,7 @@ class TextEncoderDoRAExtractFixed(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
@@ -556,7 +554,7 @@ class TextEncoderDoRAExtractFixed(io.ComfyNode):
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractRatio(io.ComfyNode):
@@ -576,7 +574,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -587,7 +585,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
@@ -598,7 +596,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractQuantile(io.ComfyNode):
@@ -618,7 +616,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -629,7 +627,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
@@ -640,7 +638,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractKnee(io.ComfyNode):
@@ -659,7 +657,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -671,7 +669,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
@@ -682,7 +680,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
@@ -702,7 +700,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
outputs=[io.AnyType.Output(display_name="output_path")],
is_output_node=True,
)
@@ -713,7 +711,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a)
model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b)
output_path = _build_lora_output_path(output_filename)
output_path, output_name = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
@@ -724,4 +722,4 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=probe_offset
)
return io.NodeOutput(output_path)
return io.NodeOutput(output_name)
tests/test_artifact_paths.py
+115
View File
@@ -0,0 +1,115 @@
import importlib
import os
import sys
import types
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_artifact_path_tests"
MIGRATED_SAVE_MODULES = (
"consensus_merger.py",
"merger.py",
"lora_merger.py",
"lora_resize.py",
"dtype_conversion.py",
"lora_extract_svd.py",
"dora_extract_wd.py",
"dora_learned_wd.py",
"text_encoder_extract.py",
"renamekeys.py",
"prunekeys.py",
"lora_rename.py",
)
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def artifact_paths():
return _load_module("nodes.artifact_paths")
@pytest.mark.parametrize(
"category",
["checkpoints", "diffusion_models", "text_encoders", "loras", "embeddings"],
)
def test_canonical_categories_use_models_dir_and_return_relative_name(
monkeypatch, tmp_path, artifact_paths, category
):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
output_path, output_name = artifact_paths.canonical_model_artifact_path(
category, "group/subgroup/model"
)
expected = tmp_path / category / "group" / "subgroup" / "model.safetensors"
assert Path(output_path) == expected
assert output_name == "group/subgroup/model.safetensors"
assert expected.parent.is_dir()
assert not os.path.isabs(output_name)
def test_registered_search_paths_do_not_influence_write_root(
monkeypatch, tmp_path, artifact_paths
):
canonical = tmp_path / "cli-model-root"
legacy = tmp_path / "legacy-unet"
additional = tmp_path / "additional-models"
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(canonical))
monkeypatch.setattr(
artifact_paths.folder_paths,
"get_folder_paths",
lambda _: [str(legacy), str(additional)],
)
output_path, output_name = artifact_paths.canonical_model_artifact_path(
"diffusion_models", "nested/model.safetensors"
)
assert Path(output_path) == canonical / "diffusion_models" / "nested" / "model.safetensors"
assert output_name == "nested/model.safetensors"
@pytest.mark.parametrize(
"name",
["model.safetensors", "model.safetensors.safetensors"],
)
def test_appends_exactly_one_extension(monkeypatch, tmp_path, artifact_paths, name):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
_, output_name = artifact_paths.canonical_model_artifact_path("loras", name)
assert output_name == "model.safetensors"
@pytest.mark.parametrize(
"name",
["../escape", "folder/../../escape", "", " ", "/absolute", ".safetensors"],
)
def test_rejects_empty_or_escaping_names(monkeypatch, tmp_path, artifact_paths, name):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
with pytest.raises(ValueError):
artifact_paths.canonical_model_artifact_path("loras", name)
def test_rejects_unrecognized_category(monkeypatch, tmp_path, artifact_paths):
monkeypatch.setattr(artifact_paths.folder_paths, "models_dir", str(tmp_path))
with pytest.raises(ValueError, match="Unsupported model artifact category"):
artifact_paths.canonical_model_artifact_path("unet", "model")
def test_persistent_save_nodes_use_shared_path_contract_and_wildcard_outputs():
for filename in MIGRATED_SAVE_MODULES:
source = (REPO_ROOT / "nodes" / filename).read_text(encoding="utf-8")
assert (
"canonical_model_artifact_path" in source
or "_build_lora_output_path" in source
), filename
assert 'io.String.Output(display_name="output_path")' not in source, filename
assert 'io.String.Output(display_name="output_filename")' not in source, filename
tests/test_dtype_conversion.py
+4 -1
View File
@@ -97,7 +97,10 @@ def test_converts_floating_and_preserves_excluded_and_nonfloating(
r"\.norm\.weight$",
f"converted_{target_name}",
)
result, metadata = _read_uel(output_path)
assert output_path == f"converted_{target_name}.safetensors"
result, metadata = _read_uel(
tmp_path / "diffusion_models" / output_path
)
assert result["blocks.0.weight"].dtype == expected_dtype
assert result["blocks.0.norm.weight"].dtype == torch.float32
tests/test_lora_diff_dtype_contract.py
+6 -2
View File
@@ -438,7 +438,7 @@ def test_merge_to_model_node_returns_path_blank_line_report(
"", "merged", "fp32", "cpu", True, False, False,
)
assert result.result == ("saved/model.safetensors\n\nSUCCESS\nDETAILS",)
assert result.result == ("merged.safetensors", "SUCCESS\nDETAILS")
def test_mixed_parser_scans_every_key_and_retains_unpaired(modules):
@@ -489,7 +489,11 @@ def test_include_1d_switches_are_appended_and_default_off(modules, generic_modul
assert final_input.id == "include_1d_diffs"
assert final_input.default is False
assert len(resize.LoRAMergeToModel.define_schema().outputs) == 1
outputs = resize.LoRAMergeToModel.define_schema().outputs
assert [output.display_name for output in outputs] == [
"output_path",
"merge_report",
]
def test_resize_preserves_mixed_direct_auxiliary_and_fp32(monkeypatch, tmp_path, modules):