diff --git a/docs/IMPLEMENTATION_INVARIANTS.md b/docs/IMPLEMENTATION_INVARIANTS.md index 5786207..9af42b8 100644 --- a/docs/IMPLEMENTATION_INVARIANTS.md +++ b/docs/IMPLEMENTATION_INVARIANTS.md @@ -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. diff --git a/nodes/artifact_paths.py b/nodes/artifact_paths.py new file mode 100644 index 0000000..de7ed0d --- /dev/null +++ b/nodes/artifact_paths.py @@ -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 diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py index 518925d..d687e57 100644 --- a/nodes/dora_extract_wd.py +++ b/nodes/dora_extract_wd.py @@ -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) diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py index 82992b2..c66c0a9 100644 --- a/nodes/dora_learned_wd.py +++ b/nodes/dora_learned_wd.py @@ -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) diff --git a/nodes/dtype_conversion.py b/nodes/dtype_conversion.py index a9bd6fc..7206517 100644 --- a/nodes/dtype_conversion.py +++ b/nodes/dtype_conversion.py @@ -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, diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index ba4409d..242cbeb 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -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) diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index ae3503e..7614bbb 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -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( diff --git a/nodes/lora_rename.py b/nodes/lora_rename.py index fb26275..1004839 100644 --- a/nodes/lora_rename.py +++ b/nodes/lora_rename.py @@ -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, ) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index dbf5a12..cbf3620 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -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) diff --git a/nodes/merger.py b/nodes/merger.py index 346fc62..6db5040 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -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"), ], ) diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index f671b5c..a7b50ab 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -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"), ], ) diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index 79f4e44..ccac302 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -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"), ], ) diff --git a/nodes/text_encoder_extract.py b/nodes/text_encoder_extract.py index 053c3b3..6f3b41c 100644 --- a/nodes/text_encoder_extract.py +++ b/nodes/text_encoder_extract.py @@ -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) diff --git a/tests/test_artifact_paths.py b/tests/test_artifact_paths.py new file mode 100644 index 0000000..9a1e957 --- /dev/null +++ b/tests/test_artifact_paths.py @@ -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 diff --git a/tests/test_dtype_conversion.py b/tests/test_dtype_conversion.py index 9933cba..b913e12 100644 --- a/tests/test_dtype_conversion.py +++ b/tests/test_dtype_conversion.py @@ -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 diff --git a/tests/test_lora_diff_dtype_contract.py b/tests/test_lora_diff_dtype_contract.py index a273335..e4ec4ea 100644 --- a/tests/test_lora_diff_dtype_contract.py +++ b/tests/test_lora_diff_dtype_contract.py @@ -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):