From 9ad992d080ea6c4b36053a7eed594c74d2843534 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 3 Aug 2026 18:57:57 +0200 Subject: [PATCH 1/2] feat(extract): register text encoder nodes --- __init__.py | 13 + nodes/text_encoder_extract.py | 721 ++++++++++++++++++++++++++++++++++ 2 files changed, 734 insertions(+) create mode 100644 nodes/text_encoder_extract.py diff --git a/__init__.py b/__init__.py index 2e1e381..516ba14 100644 --- a/__init__.py +++ b/__init__.py @@ -33,6 +33,13 @@ from .nodes.dora_learned_wd import ( DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile, DoRALearnedExtractKnee, DoRALearnedExtractFrobenius ) +from .nodes.text_encoder_extract import ( + TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio, + TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee, + TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed, + TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile, + TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius, +) from .nodes.lora_resize import ( LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, @@ -80,6 +87,12 @@ class ModelUtilsExtension(ComfyExtension): # Learned DoRA Extraction DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile, DoRALearnedExtractKnee, DoRALearnedExtractFrobenius, + # Text Encoder LoRA/DoRA Extraction + TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio, + TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee, + TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed, + TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile, + TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius, # LoRA Resize LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, diff --git a/nodes/text_encoder_extract.py b/nodes/text_encoder_extract.py new file mode 100644 index 0000000..562cc35 --- /dev/null +++ b/nodes/text_encoder_extract.py @@ -0,0 +1,721 @@ +""" +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 fnmatch +import os +import re +import torch +import torch.linalg as linalg +import folder_paths +import comfy.utils +from tqdm import tqdm +from comfy_api.latest import io +from .device_utils import ( + estimate_model_size, prepare_for_large_operation, + cleanup_after_operation +) + +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter + +# Import SVD rank utilities from standard lora extract file to avoid duplication +from .lora_extract_svd import ( + _compute_rank, + _svd_extract_linear, + _svd_extract_conv, + _extract_chunked_layer, + _detect_fused_layer, + _compile_patterns, + _matches_any_pattern +) + + +def _format_te_lora_key(key: str) -> str: + """ + Format key specifically for Text Encoder LoRA/DoRA saving. + Corrects layer name prefixing based on specific target mappings. + """ + if key.endswith(".weight"): + key = key[:-7] + + rules = [ + ("model.language_model.layers", "text_encoders.transformer.model.layers"), + ("model.visual.blocks", "text_encoders.transformer.visual.blocks"), + ("model.visual.deepstack_merger_list", "text_encoders.transformer.visual.deepstack_merger_list"), + ("model.visual.merger", "text_encoders.transformer.visual.merger"), + ("model.visual.pos_embed", "text_encoders.transformer.visual.pos_embed"), + ("model.layers", "text_encoders.transformer.model.layers"), + ("vision_model.embeddings.position_embedding", "text_encoders.transformer.vision_model.embeddings.position_embedding"), + ("vision_model.encoder.layers", "text_encoders.transformer.vision_model.encoder.layers"), + ] + + for prefix, target in rules: + if key.startswith(prefix): + return target + key[len(prefix):] + + # Prepend text_encoders. as default generic fallback + if key.startswith("text_encoders."): + return key + return f"text_encoders.{key}" + + +def extract_te_from_files( + model_a_path: str, + model_b_path: str, + mode: str, + linear_param: float, + conv_param: float, + device: str, + save_dtype: str, + output_path: str, + is_dora: bool = False, + linear_max_rank: int = None, + conv_max_rank: int = None, + clamp_quantile: float = 0.99, + min_diff: float = 0.0, + skip_patterns_str: str = "", + mismatch_mode: str = "skip", + chunk_large_layers: bool = True, + svd_niter: int = 2, + lazy_load: bool = True, + force_clear_cache: bool = True, + glob_skip_patterns: bool = False, +) -> None: + """ + Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk. + """ + save_torch_dtype = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, + }.get(save_dtype, torch.float16) + + skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns) + + # Prepare memory before heavy operation + total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path) + print(f"[TE Extract] Preparing memory for {total_size_gb:.2f}GB operation...") + prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) + + handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load) + handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load) + + try: + keys_a = set(handler_a.keys()) + keys_b = set(handler_b.keys()) + weight_keys = [k for k in keys_a if k.endswith(".weight")] + pbar = comfy.utils.ProgressBar(len(weight_keys)) + stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} + + def _process_layer(key): + lora_name = _format_te_lora_key(key) + + if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): + return "skipped", None + + # Load tensors with pinned memory for CUDA + use_pinned = device == 'cuda' + + if key not in keys_b: + if mismatch_mode == "skip": + return "skipped", None + if mismatch_mode == "error": + raise ValueError(f"Key {key} not found in model B") + + cpu_a = handler_a.get_tensor(key) + if use_pinned: + weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32) + else: + weight_diff = cpu_a.to(device=device, dtype=torch.float32) + del cpu_a + + dora_scale = None + else: + cpu_a = handler_a.get_tensor(key) + cpu_b = handler_b.get_tensor(key) + + if use_pinned: + tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) + tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32) + else: + tensor_a = cpu_a.to(device=device, dtype=torch.float32) + tensor_b = cpu_b.to(device=device, dtype=torch.float32) + del cpu_a, cpu_b + + if tensor_a.shape != tensor_b.shape: + if mismatch_mode == "skip": + del tensor_a, tensor_b + return "skipped", None + if mismatch_mode == "error": + raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}") + + weight_diff = tensor_a + dora_scale = None + del tensor_b + else: + # Skip 1D weights + if tensor_a.ndim < 2: + del tensor_a, tensor_b + return "skipped", None + + if is_dora: + # Weight Decomposition calculation + is_conv_dim = tensor_a.ndim == 4 + + if is_conv_dim: + out_ch = tensor_a.shape[0] + flat_a = tensor_a.reshape(out_ch, -1) + flat_b = tensor_b.reshape(out_ch, -1) + norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True) + norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True) + else: + norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True) + norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True) + + # Compute directional diff + eps = 1e-8 + norm_ratio = norm_b / (norm_a + eps) + if is_conv_dim: + norm_ratio = norm_ratio.view(-1, 1, 1, 1) + dora_scale = norm_a.view(-1, 1, 1, 1) + else: + dora_scale = norm_a + + w_dir = tensor_a * norm_ratio + weight_diff = w_dir - tensor_b + del tensor_a, tensor_b, w_dir + else: + weight_diff = tensor_a - tensor_b + dora_scale = None + del tensor_a, tensor_b + + # Skip small differences + if min_diff > 0 and weight_diff.abs().max() < min_diff: + del weight_diff + return "skipped", None + + # Skip 1D tensors + if weight_diff.ndim < 2: + del weight_diff + return "skipped", None + + is_conv = weight_diff.ndim == 4 + layer_results = {} + + try: + if is_conv: + result, mode_str = _svd_extract_conv( + weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile + ) + else: + result, mode_str = _svd_extract_linear( + weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + ) + except Exception as e: + # Try chunked extraction for large tensors + if chunk_large_layers and not is_conv: + num_chunks = _detect_fused_layer(key, weight_diff.shape) + if num_chunks > 1: + print(f"[TE Extract] Chunked: {key} ({num_chunks} chunks)") + lora_up, lora_down, rank = _extract_chunked_layer( + weight_diff, num_chunks, mode, linear_param, device, linear_max_rank + ) + if lora_up is not None: + # PEFT suffixes + layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() + if is_dora and dora_scale is not None: + layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous() + del weight_diff + return "chunked", layer_results + + print(f"[TE Extract] Failed: {key}: {e}") + del weight_diff + return "skipped", None + + # Store result + if mode_str == "full": + layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous() + status = "full" + else: + lora_down, lora_up, _ = result + # Write to PEFT weight suffixes + layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() + if is_dora and dora_scale is not None: + layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous() + status = "extracted" + + del weight_diff + return status, layer_results + + writer = IncrementalSafetensorsWriter(output_path) + writer.__enter__() + try: + for key in tqdm(weight_keys, desc="Extracting TE Layers", unit="layers"): + status, layer_sd = _process_layer(key) + stats[status] += 1 + if layer_sd: + writer.write_dict(layer_sd) + + if force_clear_cache: + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + pbar.update(1) + finally: + writer.__exit__(None, None, None) + + print(f"[TE Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " + f"{stats['full']} full, {stats['skipped']} skipped") + + finally: + handler_a.__exit__(None, None, None) + handler_b.__exit__(None, None, None) + cleanup_after_operation() + + +def _build_lora_output_path(output_filename: str) -> str: + """Build output path for LoRA file.""" + output_dir = folder_paths.get_folder_paths("loras")[0] + os.makedirs(output_dir, exist_ok=True) + return os.path.join(output_dir, f"{output_filename.strip()}.safetensors") + + +def _get_te_model_inputs(): + return [ + io.Combo.Input("model_a", options=folder_paths.get_filename_list("text_encoders"), + tooltip="Finetuned Text Encoder model (A - B = LoRA)"), + io.Combo.Input("model_b", options=folder_paths.get_filename_list("text_encoders"), + tooltip="Base Text Encoder model (A - B = LoRA)"), + ] + + +def _get_common_inputs(): + return [ + io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), + io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), + io.Boolean.Input("chunk_large_layers", default=False, + tooltip="Split large fused layers (QKV, MLP) into chunks"), + io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01, + tooltip="Clamp outlier singular values"), + io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001, + tooltip="Skip layers with max difference below this"), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.String.Input("output_filename", default="extracted_te_lora"), + io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), + io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), + io.String.Input("skip_patterns", default="", multiline=True, + tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"), + io.Boolean.Input("glob_skip_patterns", default=False, + tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). " + "When False (default), patterns are Python regex matched as substrings."), + ] + + +# ============================================================================= +# LoRA Text Encoder Extract Nodes +# ============================================================================= + +class TextEncoderLoRAExtractFixed(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderLoRAExtractFixed", + display_name="TE LoRA Extract (Fixed Rank)", + category="ModelUtils/LoRA Extract (TE)", + description="Extract Text Encoder LoRA with specified fixed rank for each layer type.", + inputs=[ + *_get_te_model_inputs(), + io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"), + io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"), + io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, + clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "fixed", linear_dim, conv_dim, + device, save_dtype, output_path, is_dora=False, + linear_max_rank=linear_dim, conv_max_rank=conv_dim, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + 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) + + +class TextEncoderLoRAExtractRatio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderLoRAExtractRatio", + display_name="TE LoRA Extract (Ratio)", + category="ModelUtils/LoRA Extract (TE)", + description="Keep singular values > max(S) / ratio.", + inputs=[ + *_get_te_model_inputs(), + io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1), + io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, + device, save_dtype, output_path, is_dora=False, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +class TextEncoderLoRAExtractQuantile(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderLoRAExtractQuantile", + display_name="TE LoRA Extract (Quantile)", + category="ModelUtils/LoRA Extract (TE)", + description="Keep enough singular values to reach target percentage.", + inputs=[ + *_get_te_model_inputs(), + io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01), + io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, + device, save_dtype, output_path, is_dora=False, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +class TextEncoderLoRAExtractKnee(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderLoRAExtractKnee", + display_name="TE LoRA Extract (Knee)", + category="ModelUtils/LoRA Extract (TE)", + description="Automatically find optimal rank using knee detection on singular value curve.", + inputs=[ + *_get_te_model_inputs(), + io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, knee_method, 0, 0, + device, save_dtype, output_path, is_dora=False, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +class TextEncoderLoRAExtractFrobenius(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderLoRAExtractFrobenius", + display_name="TE LoRA Extract (Frobenius)", + category="ModelUtils/LoRA Extract (TE)", + description="Preserve target fraction of Frobenius norm.", + inputs=[ + *_get_te_model_inputs(), + io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01), + io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "sv_fro", linear_target, conv_target, + device, save_dtype, output_path, is_dora=False, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +# ============================================================================= +# DoRA Text Encoder Extract Nodes +# ============================================================================= + +class TextEncoderDoRAExtractFixed(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderDoRAExtractFixed", + display_name="TE DoRA Extract (Fixed Rank)", + category="ModelUtils/DoRA Extract (TE)", + description="Extract Text Encoder DoRA with specified fixed rank for each layer type.", + inputs=[ + *_get_te_model_inputs(), + io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"), + io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"), + io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, + clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "fixed", linear_dim, conv_dim, + device, save_dtype, output_path, is_dora=True, + linear_max_rank=linear_dim, conv_max_rank=conv_dim, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + 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) + + +class TextEncoderDoRAExtractRatio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderDoRAExtractRatio", + display_name="TE DoRA Extract (Ratio)", + category="ModelUtils/DoRA Extract (TE)", + description="Keep singular values > max(S) / ratio.", + inputs=[ + *_get_te_model_inputs(), + io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1), + io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, + device, save_dtype, output_path, is_dora=True, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +class TextEncoderDoRAExtractQuantile(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderDoRAExtractQuantile", + display_name="TE DoRA Extract (Quantile)", + category="ModelUtils/DoRA Extract (TE)", + description="Keep enough singular values to reach target percentage.", + inputs=[ + *_get_te_model_inputs(), + io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01), + io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, + device, save_dtype, output_path, is_dora=True, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +class TextEncoderDoRAExtractKnee(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderDoRAExtractKnee", + display_name="TE DoRA Extract (Knee)", + category="ModelUtils/DoRA Extract (TE)", + description="Automatically find optimal rank using knee detection on singular value curve.", + inputs=[ + *_get_te_model_inputs(), + io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, knee_method, 0, 0, + device, save_dtype, output_path, is_dora=True, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) + + +class TextEncoderDoRAExtractFrobenius(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderDoRAExtractFrobenius", + display_name="TE DoRA Extract (Frobenius)", + category="ModelUtils/DoRA Extract (TE)", + description="Preserve target fraction of Frobenius norm.", + inputs=[ + *_get_te_model_inputs(), + io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01), + io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("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) + + extract_te_from_files( + model_a_path, model_b_path, "sv_fro", linear_target, conv_target, + device, save_dtype, output_path, is_dora=True, + linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, + clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + ) + return io.NodeOutput(output_path) From 0a7578ce738e7db6f2408feb58009ecab9d4a9dc Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 4 Aug 2026 15:52:42 +0200 Subject: [PATCH 2/2] feat(analysis): add model similarity nodes Separate CWB-derived diagnostics from merge weighting while streaming each logical layer once. --- __init__.py | 3 + docs/model_analysis.md | 75 +++ nodes/model_analysis.py | 1128 ++++++++++++++++++++++++++++++++++ tests/test_model_analysis.py | 282 +++++++++ 4 files changed, 1488 insertions(+) create mode 100644 docs/model_analysis.md create mode 100644 nodes/model_analysis.py create mode 100644 tests/test_model_analysis.py diff --git a/__init__.py b/__init__.py index 516ba14..772beb8 100644 --- a/__init__.py +++ b/__init__.py @@ -21,6 +21,7 @@ from .nodes.merger import ( CheckpointThreeMerger, EmbeddingThreeMerger ) from .nodes.consensus_merger import CWB_MERGER_NODES +from .nodes.model_analysis import MODEL_ANALYSIS_NODES from .nodes.lora_extract_svd import ( LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractKnee, LoRAExtractFrobenius @@ -78,6 +79,8 @@ class ModelUtilsExtension(ComfyExtension): CheckpointThreeMerger, EmbeddingThreeMerger, # Dedicated CWB Mergers *CWB_MERGER_NODES, + # Two-Model Analysis + *MODEL_ANALYSIS_NODES, # LoRA Extraction LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractKnee, LoRAExtractFrobenius, diff --git a/docs/model_analysis.md b/docs/model_analysis.md new file mode 100644 index 0000000..d428696 --- /dev/null +++ b/docs/model_analysis.md @@ -0,0 +1,75 @@ +# Two-Model Similarity Analysis + +These nodes compare two safetensors models without merging them or writing an +output file. Separate nodes are provided for checkpoints, diffusion models, +text encoders, LoRAs, and embeddings. + +## Outputs + +`comparison_report` contains conventional numerical differences and topology +findings. `cwb_report` separately contains diagnostics derived from the +similarity, alignment, and consensus stages of Consensus-Weighted Blending. +`documentation` contains this reference. + +The reports begin with a global summary, followed by inferred blockwise values +and one entry for every comparable tensor. Block names are approximated from +the tensor-key hierarchy through its first numeric path component; they are not +an architecture-specific interpretation. + +## Standard metrics + +- MAE, MSE, and RMSE describe absolute error at the original weight scale. +- Maximum absolute difference identifies the largest individual change. +- Relative L2 is `||A-B|| / ((||A||+||B||)/2)` and is scale-aware. +- Cosine similarity compares direction; Pearson correlation compares centered + variation. +- Exact equality, sign agreement, norm ratio, and finite-value coverage expose + differences hidden by averages. + +The global section reports both parameter-weighted metrics and equal-layer +averages. No arbitrary combined similarity score is assigned. Undefined +zero-norm or constant-tensor statistics are shown as `N/A`. + +Only shared, same-shape floating tensors enter numerical aggregates. Missing +keys and shape mismatches are listed separately and are never padded, cropped, +or replaced with zeros. Non-floating tensors receive an exact-equality check. +Quantized models carrying ComfyUI quantization metadata are rejected because +encoded storage values do not represent directly comparable model weights. + +## CWB diagnostics + +The CWB report provides pairwise row-cosine statistics and each input's affinity +to both mean and median consensus. It stops before CWB's merge-weighting stage: +no thresholds, contribution weights, diversity weighting, vector +reconstitution, norm rescaling, or output writer are used. + +Checkpoint, diffusion-model, and text-encoder tensors always retain their fixed +coordinate alignment. Embedding and LoRA nodes expose one optional +`cwb_similarity_alignment` control. When enabled, it performs +threshold-independent greedy one-to-one matching derived from CWB and reports +coverage, matched-score statistics, and improvement over index alignment. This +operation is quadratic and may be slow for large embeddings. + +LoRA down/up tensors are treated as one logical layer. Rank-component matching +uses the product of down and up cosine similarity, including CWB's paired-sign +equivalence for comparison. No source tensor is modified. + +## Memory and failure handling + +Inputs always use Unified Efficient Loader's low-memory streaming mode. Matching +A/B tensors are loaded concurrently through bounded `async_stream()` pipelines +with `batch_size=1` and `prefetch_batches=1`, but no subsequent logical layer is +requested early. A complete LoRA down/up pair is loaded once and used for both +reports. CUDA inputs use pinned transfer while retaining the current CPU source +tensors for fallback. + +Calculations use FP32. When CUDA processing runs out of memory, CUDA memory is +cleared and the current logical layer is retried once on CPU. After success or +failure, all tensor references for that layer are released and every loaded key +is marked processed exactly once before advancing. Both reports list every CPU +fallback. A failed CPU retry or any non-OOM error stops analysis with the layer +name in the error. + +`top_weight_differences` bounds the scalar-coordinate list; every comparable +tensor is still present in the layerwise report. `DOCUMENTATION ONLY` loads no +model data. diff --git a/nodes/model_analysis.py b/nodes/model_analysis.py new file mode 100644 index 0000000..a6c6b5e --- /dev/null +++ b/nodes/model_analysis.py @@ -0,0 +1,1128 @@ +"""Streaming two-model similarity analysis nodes.""" + +from __future__ import annotations + +import gc +import heapq +import math +import re +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import comfy.utils +import folder_paths +import torch +import torch.nn.functional as F +from comfy_api.latest import io +from tqdm import tqdm +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + +from .consensus_merger import ( + _secondary_lora_map, +) +from .device_utils import ( + cleanup_after_operation, + estimate_model_size, + prepare_for_large_operation, +) +from .merger import load_documentation_from_file +from .lora_resize import layer_tensor_keys, parse_lora_layers, validate_canonical_blocks +from .quantization_guard import inspect_low_bit_input + + +FLOAT_DTYPES = {torch.float16, torch.bfloat16, torch.float32, torch.float64} +PARAMETER_SUFFIXES = ( + ".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight", + ".lora.down.weight", ".lora.up.weight", ".reshape_weight", ".dora_scale", + ".set_weight", ".diff_b", ".w_norm", ".b_norm", ".alpha", ".weight", + ".bias", ".scale", ".diff", ".lin", +) + + +def _is_float_dtype(dtype: torch.dtype) -> bool: + return dtype in FLOAT_DTYPES + + +def _fmt(value: float | None) -> str: + if value is None or not math.isfinite(value): + return "N/A" + return f"{value:.6g}" + + +def _shape_text(shape: tuple[int, ...]) -> str: + return "scalar" if not shape else "x".join(str(value) for value in shape) + + +def _layer_name(key: str) -> str: + for suffix in PARAMETER_SUFFIXES: + if key.endswith(suffix): + return key[:-len(suffix)] or key + return key + + +def _block_name(key: str) -> str: + layer = _layer_name(key) + numeric = re.search(r"(?:^|[._])\d+(?=$|[._])", layer) + if numeric: + return layer[:numeric.end()] + split_at = max(layer.rfind("."), layer.rfind("_")) + return layer[:split_at] if split_at > 0 else layer + + +@dataclass +class RawStats: + finite: int = 0 + total: int = 0 + sum_abs: float = 0.0 + sum_sq_diff: float = 0.0 + dot: float = 0.0 + sum_a2: float = 0.0 + sum_b2: float = 0.0 + sum_a: float = 0.0 + sum_b: float = 0.0 + exact: int = 0 + sign_equal: int = 0 + max_abs: float = 0.0 + nonfinite_a: int = 0 + nonfinite_b: int = 0 + + def add(self, other: "RawStats") -> None: + for name in ( + "finite", "total", "sum_abs", "sum_sq_diff", "dot", "sum_a2", + "sum_b2", "sum_a", "sum_b", "exact", "sign_equal", + "nonfinite_a", "nonfinite_b", + ): + setattr(self, name, getattr(self, name) + getattr(other, name)) + self.max_abs = max(self.max_abs, other.max_abs) + + +def _metrics(raw: RawStats) -> dict[str, float | None]: + if raw.finite == 0: + return {name: None for name in ( + "mae", "mse", "rmse", "max", "relative_l2", "cosine", "pearson", + "exact", "sign", "norm_ratio", "coverage", + )} + n = raw.finite + mse = raw.sum_sq_diff / n + norm_a = math.sqrt(max(raw.sum_a2, 0.0)) + norm_b = math.sqrt(max(raw.sum_b2, 0.0)) + if norm_a == 0.0 and norm_b == 0.0: + cosine = 1.0 if raw.sum_sq_diff == 0.0 else None + norm_ratio = 1.0 + elif norm_a == 0.0 or norm_b == 0.0: + cosine = None + norm_ratio = None if norm_a == 0.0 else 0.0 + else: + cosine = raw.dot / (norm_a * norm_b) + norm_ratio = norm_b / norm_a + centered_a = max(raw.sum_a2 - raw.sum_a * raw.sum_a / n, 0.0) + centered_b = max(raw.sum_b2 - raw.sum_b * raw.sum_b / n, 0.0) + pearson_denom = math.sqrt(centered_a * centered_b) + pearson = ( + (raw.dot - raw.sum_a * raw.sum_b / n) / pearson_denom + if pearson_denom > 0.0 else None + ) + average_norm = (norm_a + norm_b) / 2.0 + relative_l2 = math.sqrt(raw.sum_sq_diff) / average_norm if average_norm else 0.0 + return { + "mae": raw.sum_abs / n, + "mse": mse, + "rmse": math.sqrt(max(mse, 0.0)), + "max": raw.max_abs, + "relative_l2": relative_l2, + "cosine": cosine, + "pearson": pearson, + "exact": raw.exact / n, + "sign": raw.sign_equal / n, + "norm_ratio": norm_ratio, + "coverage": raw.finite / raw.total if raw.total else None, + } + + +@dataclass +class CWBStats: + rows: int = 0 + affinity_rows: int = 0 + pair_sum: float = 0.0 + pair_min: float = math.inf + pair_max: float = -math.inf + mean_aff_a: float = 0.0 + mean_aff_b: float = 0.0 + median_aff_a: float = 0.0 + median_aff_b: float = 0.0 + reference_rows: int = 0 + matched_rows: int = 0 + alignment_rows: int = 0 + alignment_sum: float = 0.0 + alignment_min: float = math.inf + alignment_max: float = -math.inf + index_alignment_sum: float = 0.0 + + def add(self, other: "CWBStats") -> None: + if other.rows: + self.pair_min = min(self.pair_min, other.pair_min) + self.pair_max = max(self.pair_max, other.pair_max) + if other.alignment_rows: + self.alignment_min = min(self.alignment_min, other.alignment_min) + self.alignment_max = max(self.alignment_max, other.alignment_max) + for name in ( + "rows", "affinity_rows", "pair_sum", "mean_aff_a", "mean_aff_b", + "median_aff_a", "median_aff_b", "reference_rows", "matched_rows", + "alignment_rows", "alignment_sum", "index_alignment_sum", + ): + setattr(self, name, getattr(self, name) + getattr(other, name)) + + def values(self) -> dict[str, float | None]: + n = self.rows + affinity_n = self.affinity_rows + alignment_n = self.alignment_rows + alignment = self.alignment_sum / alignment_n if alignment_n else None + return { + "pair_mean": self.pair_sum / n if n else None, + "pair_min": self.pair_min if n else None, + "pair_max": self.pair_max if n else None, + "mean_aff_a": self.mean_aff_a / affinity_n if affinity_n else None, + "mean_aff_b": self.mean_aff_b / affinity_n if affinity_n else None, + "median_aff_a": self.median_aff_a / affinity_n if affinity_n else None, + "median_aff_b": self.median_aff_b / affinity_n if affinity_n else None, + "matched": self.matched_rows / self.reference_rows if self.reference_rows else None, + "alignment": alignment, + "alignment_min": self.alignment_min if alignment_n else None, + "alignment_max": self.alignment_max if alignment_n else None, + "alignment_gain": ( + alignment - self.index_alignment_sum / alignment_n + if alignment_n else None + ), + } + + +@dataclass +class TensorRecord: + key: str + shape: tuple[int, ...] + dtype_a: str + dtype_b: str + raw: RawStats + cwb: CWBStats + cpu_fallback: bool = False + + +def _row_cosines(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (F.normalize(a, dim=1, eps=1e-8) * F.normalize(b, dim=1, eps=1e-8)).sum(1) + + +def _compute_raw( + a: torch.Tensor, + b: torch.Tensor, + top_count: int, +) -> tuple[RawStats, list[tuple[float, int, float, float]]]: + flat_a = a.reshape(-1) + flat_b = b.reshape(-1) + finite_a = torch.isfinite(flat_a) + finite_b = torch.isfinite(flat_b) + finite = finite_a & finite_b + raw = RawStats( + total=flat_a.numel(), + nonfinite_a=int((~finite_a).sum().item()), + nonfinite_b=int((~finite_b).sum().item()), + ) + if not finite.any(): + return raw, [] + va = flat_a[finite] + vb = flat_b[finite] + diff = va - vb + abs_diff = diff.abs() + raw.finite = va.numel() + raw.sum_abs = float(abs_diff.sum().item()) + raw.sum_sq_diff = float((diff * diff).sum().item()) + raw.dot = float((va * vb).sum().item()) + raw.sum_a2 = float((va * va).sum().item()) + raw.sum_b2 = float((vb * vb).sum().item()) + raw.sum_a = float(va.sum().item()) + raw.sum_b = float(vb.sum().item()) + raw.exact = int((va == vb).sum().item()) + raw.sign_equal = int((torch.sign(va) == torch.sign(vb)).sum().item()) + raw.max_abs = float(abs_diff.max().item()) + if top_count <= 0: + return raw, [] + masked = torch.where(finite, (flat_a - flat_b).abs(), torch.full_like(flat_a, -1.0)) + count = min(top_count, masked.numel()) + values, indices = torch.topk(masked, count) + entries = [] + for value, index in zip(values.tolist(), indices.tolist()): + if value < 0.0: + continue + entries.append((float(value), int(index), float(flat_a[index].item()), float(flat_b[index].item()))) + return raw, entries + + +def _greedy_alignment(similarities: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Return threshold-independent one-to-one matches using CWB's greedy rule.""" + if not similarities.numel(): + empty = torch.empty(0, device=similarities.device, dtype=torch.long) + return empty, empty + scores = similarities.clone() + reference_rows = [] + source_rows = [] + for _ in range(min(scores.shape)): + flat_index = torch.argmax(scores) + reference_row = flat_index // scores.shape[1] + source_row = flat_index % scores.shape[1] + reference_rows.append(reference_row) + source_rows.append(source_row) + scores[reference_row, :] = -torch.inf + scores[:, source_row] = -torch.inf + return torch.stack(reference_rows), torch.stack(source_rows) + + +def _consensus_affinities(left: torch.Tensor, right: torch.Tensor) -> tuple[torch.Tensor, ...]: + mean_consensus = (left + right) / 2.0 + median_consensus = torch.median(torch.stack((left, right)), dim=0).values + return ( + _row_cosines(left, mean_consensus), + _row_cosines(right, mean_consensus), + _row_cosines(left, median_consensus), + _row_cosines(right, median_consensus), + ) + + +def _compute_cwb(a: torch.Tensor, b: torch.Tensor, allow_alignment: bool) -> CWBStats: + if a.ndim < 2: + rows_a = a.reshape(1, -1) + rows_b = b.reshape(1, -1) + else: + rows_a = a.reshape(a.shape[0], -1) + rows_b = b.reshape(b.shape[0], -1) + valid_a = torch.isfinite(rows_a).all(dim=1) + valid_b = torch.isfinite(rows_b).all(dim=1) + paired_valid = valid_a & valid_b + indexed_a = rows_a[paired_valid] + indexed_b = rows_b[paired_valid] + if not indexed_a.shape[0]: + return CWBStats() + + index_scores = _row_cosines(indexed_a, indexed_b) + stats = CWBStats() + aligned = allow_alignment and a.ndim >= 2 + if aligned: + finite_a = rows_a[valid_a] + finite_b = rows_b[valid_b] + similarities = torch.mm( + F.normalize(finite_a, dim=1, eps=1e-8), + F.normalize(finite_b, dim=1, eps=1e-8).T, + ) + reference_rows, source_rows = _greedy_alignment(similarities) + if not reference_rows.numel(): + return stats + compared_a = finite_a[reference_rows] + compared_b = finite_b[source_rows] + pair = similarities[reference_rows, source_rows] + stats.reference_rows = max(finite_a.shape[0], finite_b.shape[0]) + stats.matched_rows = pair.numel() + stats.alignment_rows = pair.numel() + stats.alignment_sum = float(pair.sum().item()) + stats.alignment_min = float(pair.min().item()) + stats.alignment_max = float(pair.max().item()) + stats.index_alignment_sum = float(index_scores.mean().item()) * pair.numel() + else: + compared_a = indexed_a + compared_b = indexed_b + pair = index_scores + + count = pair.numel() + if count == 0: + return stats + mean_a, mean_b, median_a, median_b = _consensus_affinities(compared_a, compared_b) + stats.pair_sum = float(pair.sum().item()) + stats.pair_min = float(pair.min().item()) + stats.pair_max = float(pair.max().item()) + stats.mean_aff_a = float(mean_a.sum().item()) + stats.mean_aff_b = float(mean_b.sum().item()) + stats.median_aff_a = float(median_a.sum().item()) + stats.median_aff_b = float(median_b.sum().item()) + stats.rows = count + stats.affinity_rows = count + return stats + + +def _compute_pair_on_device( + tensor_a: torch.Tensor, + tensor_b: torch.Tensor, + device: str, + allow_alignment: bool, + top_count: int, +) -> tuple[RawStats, CWBStats, list[tuple[float, int, float, float]]]: + if str(device).startswith("cuda"): + a = transfer_to_gpu_pinned(tensor_a, device=device, dtype=torch.float32) + b = transfer_to_gpu_pinned(tensor_b, device=device, dtype=torch.float32) + else: + a = tensor_a.to(device=device, dtype=torch.float32) + b = tensor_b.to(device=device, dtype=torch.float32) + raw, top = _compute_raw(a, b, top_count) + cwb = _compute_cwb(a, b, allow_alignment) + return raw, cwb, top + + +def _is_cuda_oom(exc: BaseException) -> bool: + oom_type = getattr(torch.cuda, "OutOfMemoryError", RuntimeError) + return isinstance(exc, oom_type) or ( + isinstance(exc, RuntimeError) and "out of memory" in str(exc).lower() + ) + + +def _release_cuda_oom() -> None: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _analyze_with_fallback( + key: str, + tensor_a: torch.Tensor, + tensor_b: torch.Tensor, + device: str, + allow_alignment: bool, + top_count: int, +): + try: + return (*_compute_pair_on_device( + tensor_a, tensor_b, device, allow_alignment, top_count + ), False) + except BaseException as oom: + if not str(device).startswith("cuda") or not _is_cuda_oom(oom): + raise + oom.__traceback__ = None + _release_cuda_oom() + try: + result = _compute_pair_on_device( + tensor_a, tensor_b, "cpu", allow_alignment, top_count + ) + except BaseException as cpu_error: + raise RuntimeError( + f"CPU fallback failed while analyzing '{key}': {cpu_error}" + ) from oom + return (*result, True) + + +def _compute_lora_pair_on_device( + down_a: torch.Tensor, + up_a: torch.Tensor, + down_b: torch.Tensor, + up_b: torch.Tensor, + device: str, + allow_alignment: bool, + top_count: int, +) -> tuple[ + RawStats, list[tuple[float, int, float, float]], + RawStats, list[tuple[float, int, float, float]], CWBStats, +]: + if str(device).startswith("cuda"): + da = transfer_to_gpu_pinned(down_a, device=device, dtype=torch.float32) + ua = transfer_to_gpu_pinned(up_a, device=device, dtype=torch.float32) + db = transfer_to_gpu_pinned(down_b, device=device, dtype=torch.float32) + ub = transfer_to_gpu_pinned(up_b, device=device, dtype=torch.float32) + else: + da = down_a.to(device=device, dtype=torch.float32) + ua = up_a.to(device=device, dtype=torch.float32) + db = down_b.to(device=device, dtype=torch.float32) + ub = up_b.to(device=device, dtype=torch.float32) + rank = da.shape[0] + if db.shape[0] != rank or ua.shape[1] != rank or ub.shape[1] != rank: + raise ValueError("LoRA down/up tensors do not share one rank.") + raw_down, top_down = _compute_raw(da, db, top_count) + raw_up, top_up = _compute_raw(ua, ub, top_count) + da_rows = da.reshape(rank, -1) + db_rows = db.reshape(rank, -1) + ua_rows = ua.movedim(1, 0).reshape(rank, -1) + ub_rows = ub.movedim(1, 0).reshape(rank, -1) + valid_a = torch.isfinite(da_rows).all(1) & torch.isfinite(ua_rows).all(1) + valid_b = torch.isfinite(db_rows).all(1) & torch.isfinite(ub_rows).all(1) + paired_valid = valid_a & valid_b + indexed_da = da_rows[paired_valid] + indexed_db = db_rows[paired_valid] + indexed_ua = ua_rows[paired_valid] + indexed_ub = ub_rows[paired_valid] + stats = CWBStats() + if not indexed_da.shape[0]: + return raw_down, top_down, raw_up, top_up, stats + + index_scores = _row_cosines(indexed_da, indexed_db) * _row_cosines(indexed_ua, indexed_ub) + if allow_alignment: + finite_da, finite_ua = da_rows[valid_a], ua_rows[valid_a] + finite_db, finite_ub = db_rows[valid_b], ub_rows[valid_b] + down_matrix = torch.mm( + F.normalize(finite_da, dim=1, eps=1e-8), + F.normalize(finite_db, dim=1, eps=1e-8).T, + ) + up_matrix = torch.mm( + F.normalize(finite_ua, dim=1, eps=1e-8), + F.normalize(finite_ub, dim=1, eps=1e-8).T, + ) + contribution_matrix = down_matrix * up_matrix + reference_rows, source_rows = _greedy_alignment(contribution_matrix) + compared_da = finite_da[reference_rows] + compared_ua = finite_ua[reference_rows] + compared_db = finite_db[source_rows] + compared_ub = finite_ub[source_rows] + selected_down = down_matrix[reference_rows, source_rows] + selected_up = up_matrix[reference_rows, source_rows] + sign_flip = (selected_down < 0) & (selected_up < 0) + compared_db = torch.where(sign_flip[:, None], -compared_db, compared_db) + compared_ub = torch.where(sign_flip[:, None], -compared_ub, compared_ub) + pair = contribution_matrix[reference_rows, source_rows] + stats.reference_rows = max(finite_da.shape[0], finite_db.shape[0]) + stats.matched_rows = pair.numel() + stats.alignment_rows = pair.numel() + stats.alignment_sum = float(pair.sum().item()) + stats.alignment_min = float(pair.min().item()) + stats.alignment_max = float(pair.max().item()) + stats.index_alignment_sum = float(index_scores.mean().item()) * pair.numel() + else: + compared_da, compared_db = indexed_da, indexed_db + compared_ua, compared_ub = indexed_ua, indexed_ub + pair = index_scores + + stats.rows = pair.numel() + stats.pair_sum = float(pair.sum().item()) + stats.pair_min = float(pair.min().item()) + stats.pair_max = float(pair.max().item()) + for left, right in ((compared_da, compared_db), (compared_ua, compared_ub)): + mean_a, mean_b, median_a, median_b = _consensus_affinities(left, right) + stats.mean_aff_a += float(mean_a.sum().item()) + stats.mean_aff_b += float(mean_b.sum().item()) + stats.median_aff_a += float(median_a.sum().item()) + stats.median_aff_b += float(median_b.sum().item()) + stats.affinity_rows += left.shape[0] + return raw_down, top_down, raw_up, top_up, stats + + +def _analyze_lora_with_fallback( + key: str, + tensors: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + device: str, + allow_alignment: bool, + top_count: int, +): + try: + return (*_compute_lora_pair_on_device( + *tensors, device, allow_alignment, top_count + ), False) + except BaseException as oom: + if not str(device).startswith("cuda") or not _is_cuda_oom(oom): + raise + oom.__traceback__ = None + _release_cuda_oom() + try: + result = _compute_lora_pair_on_device( + *tensors, "cpu", allow_alignment, top_count + ) + except BaseException as cpu_error: + raise RuntimeError( + f"CPU fallback failed while analyzing LoRA layer '{key}': {cpu_error}" + ) from oom + return (*result, True) + + +def _coordinate(shape: tuple[int, ...], flat_index: int) -> str: + if not shape: + return "()" + result = [] + remaining = flat_index + for size in reversed(shape): + result.append(remaining % size) + remaining //= size + return str(tuple(reversed(result))) + + +def _load_paired_unit(handlers, units, *, pin_memory: bool): + """Load one logical A/B analysis unit concurrently, with no cross-unit prefetch.""" + stream_a = handlers[0].async_stream( + [unit[1] for unit in units], + batch_size=1, + prefetch_batches=1, + pin_memory=pin_memory, + ) + stream_b = handlers[1].async_stream( + [unit[2] for unit in units], + batch_size=1, + prefetch_batches=1, + pin_memory=pin_memory, + ) + loaded = [] + try: + with ThreadPoolExecutor(max_workers=2) as executor: + for unit in units: + future_a = executor.submit(next, stream_a) + future_b = executor.submit(next, stream_b) + batch_a = future_a.result() + batch_b = future_b.result() + if len(batch_a) != 1 or len(batch_b) != 1: + raise RuntimeError("Asynchronous model analysis expected one tensor per batch.") + key_a, tensor_a = batch_a[0] + key_b, tensor_b = batch_b[0] + if key_a != unit[1] or key_b != unit[2]: + raise RuntimeError( + "Asynchronous model analysis returned tensors out of order." + ) + loaded.append((*unit, tensor_a, tensor_b)) + finally: + stream_a.close() + stream_b.close() + return loaded + + +def _average_layer_metrics(layer_stats: dict[str, RawStats]) -> dict[str, float | None]: + metrics = [_metrics(raw) for raw in layer_stats.values()] + result = {} + for name in ("mae", "mse", "rmse", "relative_l2", "cosine", "pearson", "exact", "sign", "norm_ratio"): + values = [entry[name] for entry in metrics if entry[name] is not None and math.isfinite(entry[name])] + result[name] = sum(values) / len(values) if values else None + return result + + +def _metadata_lines(metadata_a: dict | None, metadata_b: dict | None) -> list[str]: + a = metadata_a or {} + b = metadata_b or {} + keys = sorted(set(a) | set(b)) + differences = [key for key in keys if a.get(key) != b.get(key)] + lines = [f"Metadata keys: A={len(a)}, B={len(b)}, differing={len(differences)}"] + lines.extend(f"- {key}" for key in differences) + return lines + + +def _render_comparison( + model_a: str, + model_b: str, + model_type: str, + records: list[TensorRecord], + global_raw: RawStats, + layer_stats: dict[str, RawStats], + block_stats: dict[str, RawStats], + inventory: dict[str, list[str]], + top_weights: list[tuple], + fallback_keys: list[str], + metadata_a: dict | None, + metadata_b: dict | None, +) -> str: + weighted = _metrics(global_raw) + layer_average = _average_layer_metrics(layer_stats) + lines = [ + "MODEL COMPARISON SUMMARY", + f"Type: {model_type}", + f"Model A: {model_a}", + f"Model B: {model_b}", + f"Comparable floating tensors: {len(records)}", + f"Comparable parameters: {global_raw.total}", + f"A-only keys: {len(inventory['a_only'])}", + f"B-only keys: {len(inventory['b_only'])}", + f"Shape mismatches: {len(inventory['shape'])}", + f"Non-floating shared tensors: {len(inventory['nonfloat'])}", + f"Unsupported low-bit tensors: {len(inventory['low_bit'])}", + f"CUDA OOM CPU fallbacks: {len(fallback_keys)}", + "", + "PARAMETER-WEIGHTED GLOBAL METRICS", + ] + for name in ("mae", "mse", "rmse", "max", "relative_l2", "cosine", "pearson", "exact", "sign", "norm_ratio", "coverage"): + lines.append(f"{name}: {_fmt(weighted[name])}") + lines.extend(["", "EQUAL-LAYER AVERAGES"]) + for name, value in layer_average.items(): + lines.append(f"{name}: {_fmt(value)}") + + metrics_by_record = [(record, _metrics(record.raw)) for record in records] + lines.extend(["", "MOST DIVERGENT TENSORS", "By relative L2:"]) + for record, metric in sorted(metrics_by_record, key=lambda item: item[1]["relative_l2"] or -1.0, reverse=True)[:10]: + lines.append(f"- {record.key}: {_fmt(metric['relative_l2'])}") + lines.append("By MAE:") + for record, metric in sorted(metrics_by_record, key=lambda item: item[1]["mae"] or -1.0, reverse=True)[:10]: + lines.append(f"- {record.key}: {_fmt(metric['mae'])}") + lines.append("By lowest cosine:") + cosine_records = [(record, metric) for record, metric in metrics_by_record if metric["cosine"] is not None] + for record, metric in sorted(cosine_records, key=lambda item: item[1]["cosine"])[:10]: + lines.append(f"- {record.key}: {_fmt(metric['cosine'])}") + + lines.extend(["", "INFERRED BLOCKWISE METRICS", "block | params | MAE | RMSE | relative_L2 | cosine"]) + for block in sorted(block_stats): + raw = block_stats[block] + metric = _metrics(raw) + lines.append( + f"{block} | {raw.total} | {_fmt(metric['mae'])} | {_fmt(metric['rmse'])} | " + f"{_fmt(metric['relative_l2'])} | {_fmt(metric['cosine'])}" + ) + + lines.extend(["", "LAYERWISE TENSOR METRICS", "key | shape | dtype A/B | MAE | RMSE | max_abs | relative_L2 | cosine | pearson | exact | sign | norm_B/A"]) + for record, metric in sorted(metrics_by_record, key=lambda item: item[0].key): + lines.append( + f"{record.key} | {_shape_text(record.shape)} | {record.dtype_a}/{record.dtype_b} | " + f"{_fmt(metric['mae'])} | {_fmt(metric['rmse'])} | {_fmt(metric['max'])} | " + f"{_fmt(metric['relative_l2'])} | {_fmt(metric['cosine'])} | {_fmt(metric['pearson'])} | " + f"{_fmt(metric['exact'])} | {_fmt(metric['sign'])} | {_fmt(metric['norm_ratio'])}" + ) + + lines.extend(["", "TOP SCALAR WEIGHT DIFFERENCES"]) + for value, key, flat_index, a_value, b_value, shape in sorted(top_weights, reverse=True): + lines.append( + f"- {key}{_coordinate(shape, flat_index)}: |delta|={_fmt(value)}, " + f"A={_fmt(a_value)}, B={_fmt(b_value)}" + ) + if not top_weights: + lines.append("- None") + + lines.extend(["", "TOPOLOGY AND DATA FINDINGS"]) + for label, name in ( + ("A ONLY", "a_only"), ("B ONLY", "b_only"), ("SHAPE MISMATCH", "shape"), + ("NON-FLOATING", "nonfloat"), ("LOW-BIT UNSUPPORTED", "low_bit"), + ): + lines.append(f"[{label}] ({len(inventory[name])})") + lines.extend(f"- {value}" for value in inventory[name]) + if not inventory[name]: + lines.append("- None") + lines.append(f"[CUDA OOM -> CPU] ({len(fallback_keys)})") + lines.extend(f"- {key}" for key in fallback_keys) + if not fallback_keys: + lines.append("- None") + lines.extend(["", "METADATA"]) + lines.extend(_metadata_lines(metadata_a, metadata_b)) + return "\n".join(lines) + + +def _render_cwb( + model_a: str, + model_b: str, + records: list[TensorRecord], + block_cwb: dict[str, CWBStats], + fallback_keys: list[str], + alignment_scope: str, +) -> str: + global_cwb = CWBStats() + for record in records: + global_cwb.add(record.cwb) + values = global_cwb.values() + lines = [ + "CWB-DERIVED SIMILARITY SUMMARY", + f"Model A: {model_a}", + f"Model B: {model_b}", + f"Alignment scope: {alignment_scope}", + "Boundary: CWB similarity, alignment, and consensus diagnostics only.", + "Excluded: merge weighting, contribution calculation, vector reconstitution, norm rescaling, and output writing.", + "", + "GLOBAL CWB-DERIVED DIAGNOSTICS", + f"Pairwise row cosine mean/min/max: {_fmt(values['pair_mean'])} / {_fmt(values['pair_min'])} / {_fmt(values['pair_max'])}", + f"Mean-consensus affinity A/B: {_fmt(values['mean_aff_a'])} / {_fmt(values['mean_aff_b'])}", + f"Median-consensus affinity A/B: {_fmt(values['median_aff_a'])} / {_fmt(values['median_aff_b'])}", + f"Greedy alignment coverage: {_fmt(values['matched'])}", + f"Greedy matched cosine mean/min/max: {_fmt(values['alignment'])} / {_fmt(values['alignment_min'])} / {_fmt(values['alignment_max'])}", + f"Greedy improvement over index alignment: {_fmt(values['alignment_gain'])}", + "", + "INFERRED BLOCKWISE CWB DIAGNOSTICS", + "block | pair_cos mean/min/max | mean_aff_A/B | median_aff_A/B | matched | alignment | gain", + ] + for block in sorted(block_cwb): + entry = block_cwb[block].values() + lines.append( + f"{block} | {_fmt(entry['pair_mean'])}/{_fmt(entry['pair_min'])}/{_fmt(entry['pair_max'])} | " + f"{_fmt(entry['mean_aff_a'])}/{_fmt(entry['mean_aff_b'])} | " + f"{_fmt(entry['median_aff_a'])}/{_fmt(entry['median_aff_b'])} | " + f"{_fmt(entry['matched'])} | {_fmt(entry['alignment'])} | {_fmt(entry['alignment_gain'])}" + ) + lines.extend([ + "", "LAYERWISE CWB DIAGNOSTICS", + "key | pair_cos mean/min/max | mean_aff_A/B | median_aff_A/B | matched | alignment | gain", + ]) + for record in sorted(records, key=lambda value: value.key): + entry = record.cwb.values() + lines.append( + f"{record.key} | {_fmt(entry['pair_mean'])}/{_fmt(entry['pair_min'])}/{_fmt(entry['pair_max'])} | " + f"{_fmt(entry['mean_aff_a'])}/{_fmt(entry['mean_aff_b'])} | " + f"{_fmt(entry['median_aff_a'])}/{_fmt(entry['median_aff_b'])} | " + f"{_fmt(entry['matched'])} | {_fmt(entry['alignment'])} | {_fmt(entry['alignment_gain'])}" + ) + lines.append("") + lines.append(f"CUDA OOM CPU fallbacks: {len(fallback_keys)}") + lines.extend(f"- {key}" for key in fallback_keys) + return "\n".join(lines) + + +class ModelAnalysisLogic: + @classmethod + def execute( + cls, + model_a: str, + model_b: str, + model_type: str, + params: dict, + *, + embedding_alignment: bool = False, + lora_mode: bool = False, + ) -> tuple[str, str]: + paths = [] + for name in (model_a, model_b): + path = folder_paths.get_full_path(model_type, name) + if not path: + raise FileNotFoundError(f"{model_type} input '{name}' was not found.") + paths.append(path) + total_size = sum(estimate_model_size(path) for path in paths) + if total_size: + prepare_for_large_operation(total_size * 1.2, torch.device(params["process_device"])) + + handlers = [MemoryEfficientSafeOpen(path, low_memory=True) for path in paths] + try: + low_bit = [ + inspect_low_bit_input(handler, f"Input {index + 1} ({path})", "Model Analysis") + for index, (handler, path) in enumerate(zip(handlers, paths)) + ] + similarity_alignment = bool(params.get("cwb_similarity_alignment", False)) + keys_a = set(handlers[0].keys()) + keys_b = set(handlers[1].keys()) + inventory = { + "a_only": [], "b_only": [], + "shape": [], "nonfloat": [], "low_bit": [], + } + units: list[tuple[str, str, str]] = [] + lora_pairs: list[tuple[str, dict[str, str], dict[str, str]]] = [] + if lora_mode: + parsed_a, parsed_b = parse_lora_layers(list(keys_a)), parse_lora_layers(list(keys_b)) + pairs_a, passthrough_a = parsed_a + pairs_b, passthrough_b = parsed_b + pairs_a = {block: layer_tensor_keys(roles) for block, roles in pairs_a.items()} + pairs_b = {block: layer_tensor_keys(roles) for block, roles in pairs_b.items()} + validate_canonical_blocks(pairs_a, "Model Analysis input A") + validate_canonical_blocks(pairs_b, "Model Analysis input B") + map_a = _secondary_lora_map(pairs_a, input_index=0) + map_b = _secondary_lora_map(pairs_b, input_index=1) + for core in sorted(set(map_a) | set(map_b)): + if core not in map_a: + inventory["b_only"].extend( + f"{core}.{role}: {key}" for role, key in sorted(pairs_b[map_b[core]].items()) + ) + continue + if core not in map_b: + inventory["a_only"].extend( + f"{core}.{role}: {key}" for role, key in sorted(pairs_a[map_a[core]].items()) + ) + continue + roles_a = pairs_a[map_a[core]] + roles_b = pairs_b[map_b[core]] + complete_pair = all( + role in roles_a and role in roles_b for role in ("down", "up") + ) + if complete_pair: + down_shape_a = tuple(handlers[0].get_shape(roles_a["down"])) + up_shape_a = tuple(handlers[0].get_shape(roles_a["up"])) + down_shape_b = tuple(handlers[1].get_shape(roles_b["down"])) + up_shape_b = tuple(handlers[1].get_shape(roles_b["up"])) + complete_pair = ( + down_shape_a == down_shape_b + and up_shape_a == up_shape_b + and len(down_shape_a) >= 2 + and len(up_shape_a) >= 2 + and down_shape_a[0] == up_shape_a[1] + and down_shape_b[0] == up_shape_b[1] + and all(_is_float_dtype(dtype) for dtype in ( + handlers[0].get_dtype(roles_a["down"]), + handlers[0].get_dtype(roles_a["up"]), + handlers[1].get_dtype(roles_b["down"]), + handlers[1].get_dtype(roles_b["up"]), + )) + ) + for role in sorted(set(roles_a) | set(roles_b)): + display = f"{core}.{role}" + if role not in roles_a: + inventory["b_only"].append(f"{display}: {roles_b[role]}") + elif role not in roles_b: + inventory["a_only"].append(f"{display}: {roles_a[role]}") + elif not (complete_pair and role in {"down", "up"}): + units.append((display, roles_a[role], roles_b[role])) + if complete_pair: + lora_pairs.append((core, roles_a, roles_b)) + pass_a = set(passthrough_a) + pass_b = set(passthrough_b) + units.extend((key, key, key) for key in sorted(pass_a & pass_b)) + inventory["a_only"].extend(sorted(pass_a - pass_b)) + inventory["b_only"].extend(sorted(pass_b - pass_a)) + else: + units = [(key, key, key) for key in sorted(keys_a & keys_b)] + inventory["a_only"] = sorted(keys_a - keys_b) + inventory["b_only"] = sorted(keys_b - keys_a) + global_raw = RawStats() + layer_stats: dict[str, RawStats] = {} + block_stats: dict[str, RawStats] = {} + block_cwb: dict[str, CWBStats] = {} + records: list[TensorRecord] = [] + cwb_records: list[TensorRecord] = [] + fallback_keys: list[str] = [] + top_heap: list[tuple] = [] + top_limit = params["top_weight_differences"] + + def add_standard_record( + key, shape, dtype_a, dtype_b, raw, cwb, top, fallback, + *, include_cwb=True, + ): + record = TensorRecord( + key=key, shape=shape, dtype_a=str(dtype_a), dtype_b=str(dtype_b), + raw=raw, cwb=cwb, cpu_fallback=fallback, + ) + records.append(record) + global_raw.add(raw) + layer = _layer_name(key) + block = _block_name(key) + layer_stats.setdefault(layer, RawStats()).add(raw) + block_stats.setdefault(block, RawStats()).add(raw) + if include_cwb and cwb.rows: + cwb_records.append(record) + block_cwb.setdefault(block, CWBStats()).add(cwb) + for value, flat_index, a_value, b_value in top: + entry = (value, key, flat_index, a_value, b_value, shape) + if len(top_heap) < top_limit: + heapq.heappush(top_heap, entry) + elif value > top_heap[0][0]: + heapq.heapreplace(top_heap, entry) + + def clear_after_unit(): + if params["force_clear_cache"]: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + ready_units = [] + skipped_units = 0 + for key, key_a, key_b in units: + if key_a in low_bit[0] or key_b in low_bit[1]: + inventory["low_bit"].append(f"{key}: A={key_a}, B={key_b}") + skipped_units += 1 + continue + shape_a = tuple(handlers[0].get_shape(key_a)) + shape_b = tuple(handlers[1].get_shape(key_b)) + if shape_a != shape_b: + inventory["shape"].append(f"{key}: A={shape_a}, B={shape_b}") + skipped_units += 1 + continue + ready_units.append(( + key, key_a, key_b, shape_a, + handlers[0].get_dtype(key_a), handlers[1].get_dtype(key_b), + )) + progress = comfy.utils.ProgressBar(len(units) + len(lora_pairs)) + if skipped_units: + progress.update(skipped_units) + + with torch.no_grad(): + iterator = tqdm(ready_units, desc="Analyzing model tensors", unit="tensors") + for key, key_a, key_b, shape_a, dtype_a, dtype_b in iterator: + loaded = _load_paired_unit( + handlers, [(key, key_a, key_b)], + pin_memory=str(params["process_device"]).startswith("cuda"), + ) + _, _, _, tensor_a, tensor_b = loaded[0] + try: + if not _is_float_dtype(dtype_a) or not _is_float_dtype(dtype_b): + equal = torch.equal(tensor_a, tensor_b) + inventory["nonfloat"].append( + f"{key}: A={dtype_a}, B={dtype_b}, exact_equal={equal}" + ) + continue + raw, cwb, top, fallback = _analyze_with_fallback( + key, tensor_a, tensor_b, params["process_device"], + embedding_alignment and similarity_alignment, top_limit, + ) + add_standard_record( + key, shape_a, dtype_a, dtype_b, raw, cwb, top, fallback + ) + if fallback: + fallback_keys.append(key) + finally: + del tensor_a, tensor_b + loaded.clear() + handlers[0].mark_processed(key_a) + handlers[1].mark_processed(key_b) + clear_after_unit() + progress.update(1) + + for core, roles_a, roles_b in tqdm(lora_pairs, desc="Analyzing LoRA pairs", unit="layers"): + pair_keys = (roles_a["down"], roles_a["up"], roles_b["down"], roles_b["up"]) + if pair_keys[0] in low_bit[0] or pair_keys[1] in low_bit[0] or pair_keys[2] in low_bit[1] or pair_keys[3] in low_bit[1]: + inventory["low_bit"].append(f"{core}.[down+up]") + progress.update(1) + continue + down_shape_a = tuple(handlers[0].get_shape(pair_keys[0])) + up_shape_a = tuple(handlers[0].get_shape(pair_keys[1])) + down_shape_b = tuple(handlers[1].get_shape(pair_keys[2])) + up_shape_b = tuple(handlers[1].get_shape(pair_keys[3])) + pair_units = [ + ("down", pair_keys[0], pair_keys[2]), + ("up", pair_keys[1], pair_keys[3]), + ] + loaded_items = _load_paired_unit( + handlers, + pair_units, + pin_memory=str(params["process_device"]).startswith("cuda"), + ) + loaded = { + role: (tensor_a, tensor_b) + for role, _, _, tensor_a, tensor_b in loaded_items + } + tensors = ( + loaded["down"][0], loaded["up"][0], + loaded["down"][1], loaded["up"][1], + ) + try: + raw_down, top_down, raw_up, top_up, cwb, fallback = ( + _analyze_lora_with_fallback( + core, tensors, params["process_device"], + similarity_alignment, top_limit, + ) + ) + add_standard_record( + f"{core}.down", down_shape_a, tensors[0].dtype, tensors[2].dtype, + raw_down, CWBStats(), top_down, fallback, include_cwb=False, + ) + add_standard_record( + f"{core}.up", up_shape_a, tensors[1].dtype, tensors[3].dtype, + raw_up, CWBStats(), top_up, fallback, include_cwb=False, + ) + synthetic = TensorRecord( + key=f"{core}.[down+up]", shape=down_shape_a, + dtype_a=str(tensors[0].dtype), dtype_b=str(tensors[2].dtype), + raw=RawStats(), cwb=cwb, cpu_fallback=fallback, + ) + cwb_records.append(synthetic) + block_cwb.setdefault(_block_name(core), CWBStats()).add(cwb) + if fallback: + fallback_keys.append(synthetic.key) + finally: + del tensors + loaded.clear() + loaded_items.clear() + handlers[0].mark_processed(pair_keys[0]) + handlers[0].mark_processed(pair_keys[1]) + handlers[1].mark_processed(pair_keys[2]) + handlers[1].mark_processed(pair_keys[3]) + clear_after_unit() + progress.update(1) + + comparison = _render_comparison( + model_a, model_b, model_type, records, global_raw, layer_stats, + block_stats, inventory, top_heap, fallback_keys, + handlers[0].metadata(), handlers[1].metadata(), + ) + cwb_report = _render_cwb( + model_a, model_b, cwb_records, block_cwb, fallback_keys, + ( + "paired LoRA rank components with greedy alignment" + if lora_mode and similarity_alignment + else "paired LoRA rank components (index only)" + if lora_mode + else "embedding rows with greedy alignment" + if embedding_alignment and similarity_alignment + else "embedding rows (index only)" + if embedding_alignment else "fixed tensor rows (index only)" + ), + ) + return comparison, cwb_report + finally: + for handler in handlers: + handler.__exit__(None, None, None) + cleanup_after_operation() + + +def _common_inputs(model_type: str, *, alignment_control: bool): + inputs = [ + io.Combo.Input("execution_mode", options=["ANALYZE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=folder_paths.get_filename_list(model_type)), + io.Combo.Input("model_b", options=folder_paths.get_filename_list(model_type)), + ] + if alignment_control: + inputs.append(io.Boolean.Input( + "cwb_similarity_alignment", + default=False, + tooltip=( + "Run threshold-independent CWB-style greedy one-to-one alignment. " + "This is quadratic and can be slow for large embeddings." + ), + )) + inputs.extend([ + io.Int.Input("top_weight_differences", default=20, min=0, max=1000), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.Boolean.Input("force_clear_cache", default=True), + ]) + return inputs + + +class _ModelAnalysisNode(io.ComfyNode): + MODEL_TYPE = "diffusion_models" + NODE_ID = "" + DISPLAY_NAME = "" + EMBEDDING_ALIGNMENT = False + LORA_MODE = False + + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.NODE_ID, + display_name=cls.DISPLAY_NAME, + category="ModelUtils/Analysis", + inputs=_common_inputs( + cls.MODEL_TYPE, + alignment_control=cls.EMBEDDING_ALIGNMENT or cls.LORA_MODE, + ), + outputs=[ + io.String.Output(display_name="comparison_report"), + io.String.Output(display_name="cwb_report"), + io.String.Output(display_name="documentation"), + ], + is_experimental=True, + ) + + @classmethod + def execute(cls, **kwargs): + documentation = load_documentation_from_file("model_analysis.md") + if kwargs["execution_mode"] == "DOCUMENTATION ONLY": + return io.NodeOutput( + "Documentation mode active. No comparison performed.", + "Documentation mode active. No CWB analysis performed.", + documentation, + ) + comparison, cwb_report = ModelAnalysisLogic.execute( + kwargs["model_a"], kwargs["model_b"], cls.MODEL_TYPE, kwargs, + embedding_alignment=cls.EMBEDDING_ALIGNMENT, + lora_mode=cls.LORA_MODE, + ) + return io.NodeOutput(comparison, cwb_report, documentation) + + +class CheckpointModelAnalysis(_ModelAnalysisNode): + MODEL_TYPE = "checkpoints" + NODE_ID = "CheckpointModelAnalysis" + DISPLAY_NAME = "Analyze Checkpoint Similarity (2 Models)" + + +class DiffusionModelAnalysis(_ModelAnalysisNode): + NODE_ID = "DiffusionModelAnalysis" + DISPLAY_NAME = "Analyze Diffusion Model Similarity (2 Models)" + + +class TextEncoderModelAnalysis(_ModelAnalysisNode): + MODEL_TYPE = "text_encoders" + NODE_ID = "TextEncoderModelAnalysis" + DISPLAY_NAME = "Analyze Text Encoder Similarity (2 Models)" + + +class LoRAModelAnalysis(_ModelAnalysisNode): + MODEL_TYPE = "loras" + NODE_ID = "LoRAModelAnalysis" + DISPLAY_NAME = "Analyze LoRA Similarity (2 Models)" + LORA_MODE = True + + +class EmbeddingModelAnalysis(_ModelAnalysisNode): + MODEL_TYPE = "embeddings" + NODE_ID = "EmbeddingModelAnalysis" + DISPLAY_NAME = "Analyze Embedding Similarity (2 Models)" + EMBEDDING_ALIGNMENT = True + + +MODEL_ANALYSIS_NODES = [ + CheckpointModelAnalysis, + DiffusionModelAnalysis, + TextEncoderModelAnalysis, + LoRAModelAnalysis, + EmbeddingModelAnalysis, +] diff --git a/tests/test_model_analysis.py b/tests/test_model_analysis.py new file mode 100644 index 0000000..cf386b1 --- /dev/null +++ b/tests/test_model_analysis.py @@ -0,0 +1,282 @@ +import importlib +import math +import sys +import types +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_NAME = "modelutils_analysis_tests" + + +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 analysis(): + return _load_module("nodes.model_analysis") + + +def _params(**overrides): + params = { + "cwb_similarity_alignment": False, + "top_weight_differences": 3, + "process_device": "cpu", + "force_clear_cache": False, + } + params.update(overrides) + return params + + +def test_five_node_schema_contract(analysis): + assert len(analysis.MODEL_ANALYSIS_NODES) == 5 + assert len({node.NODE_ID for node in analysis.MODEL_ANALYSIS_NODES}) == 5 + for node in analysis.MODEL_ANALYSIS_NODES: + schema = node.define_schema() + assert schema.category == "ModelUtils/Analysis" + assert [item.id for item in schema.inputs[:3]] == [ + "execution_mode", "model_a", "model_b", + ] + assert schema.inputs[0].options == ["ANALYZE", "DOCUMENTATION ONLY"] + assert "lazy_load" not in [item.id for item in schema.inputs] + input_ids = [item.id for item in schema.inputs] + if node in (analysis.LoRAModelAnalysis, analysis.EmbeddingModelAnalysis): + assert input_ids[3] == "cwb_similarity_alignment" + else: + assert "cwb_similarity_alignment" not in input_ids + assert not { + "cwb_preset", "consensus_type", "alignment_threshold", + "similarity_threshold", "power_alpha", "diversity_beta", + "dynamic_similarity_contrast", "soft_comfort_bandpass", + "position_weight", "preserve_common_prefix", + } & set(input_ids) + assert [item.display_name for item in schema.outputs] == [ + "comparison_report", "cwb_report", "documentation", + ] + + +def test_documentation_mode_has_three_separate_outputs(analysis): + result = analysis.DiffusionModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY") + assert len(result.result) == 3 + assert "No comparison performed" in result[0] + assert "No CWB analysis performed" in result[1] + assert result[2].startswith("# Two-Model Similarity Analysis") + + +def test_raw_metrics_have_independent_expected_values(analysis): + a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + b = torch.tensor([[1.0, 0.0], [5.0, 4.0]]) + raw, top = analysis._compute_raw(a, b, 2) + metrics = analysis._metrics(raw) + assert metrics["mae"] == pytest.approx(1.0) + assert metrics["mse"] == pytest.approx(2.0) + assert metrics["rmse"] == pytest.approx(math.sqrt(2.0)) + assert metrics["max"] == pytest.approx(2.0) + assert metrics["exact"] == pytest.approx(0.5) + assert len(top) == 2 + assert {entry[1] for entry in top} == {1, 2} + + +def test_cwb_derived_metrics_have_independent_expected_values(analysis): + a = torch.tensor([[1.0, 0.0], [1.0, 0.0]]) + b = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + values = analysis._compute_cwb(a, b, False).values() + assert values["pair_mean"] == pytest.approx(0.5) + assert values["pair_min"] == pytest.approx(0.0) + assert values["pair_max"] == pytest.approx(1.0) + assert values["mean_aff_a"] == pytest.approx((1.0 + 2 ** -0.5) / 2) + assert values["mean_aff_b"] == pytest.approx((1.0 + 2 ** -0.5) / 2) + assert values["median_aff_a"] == pytest.approx(0.5) + assert values["median_aff_b"] == pytest.approx(0.5) + assert values["alignment"] is None + + +def test_zero_constant_and_nonfinite_metrics_are_explicit(analysis): + raw, _ = analysis._compute_raw(torch.zeros(3), torch.zeros(3), 0) + metrics = analysis._metrics(raw) + assert metrics["cosine"] == pytest.approx(1.0) + assert metrics["pearson"] is None + assert metrics["norm_ratio"] == pytest.approx(1.0) + + raw, _ = analysis._compute_raw( + torch.tensor([1.0, float("nan"), float("inf")]), + torch.tensor([2.0, 4.0, 5.0]), + 0, + ) + metrics = analysis._metrics(raw) + assert raw.finite == 1 + assert raw.nonfinite_a == 2 + assert metrics["coverage"] == pytest.approx(1 / 3) + + +def test_end_to_end_reports_and_topology_classification(monkeypatch, tmp_path, analysis): + path_a = tmp_path / "a.safetensors" + path_b = tmp_path / "b.safetensors" + save_file({ + "model.blocks.0.weight": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "only_a": torch.ones(1), + "shape": torch.ones(2), + "integer": torch.tensor([1, 2], dtype=torch.int64), + }, str(path_a), metadata={"source": "a"}) + save_file({ + "model.blocks.0.weight": torch.tensor([[1.0, 0.0], [5.0, 4.0]]), + "only_b": torch.ones(1), + "shape": torch.ones(3), + "integer": torch.tensor([1, 3], dtype=torch.int64), + }, str(path_b), metadata={"source": "b"}) + paths = {"a": str(path_a), "b": str(path_b)} + monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name]) + monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None) + monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None) + processed = [] + async_calls = [] + original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed + original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream + + def track_processed(loader, key): + processed.append((loader.filename, key)) + return original_mark_processed(loader, key) + + def track_async_stream(loader, keys, **kwargs): + async_calls.append((loader.filename, tuple(keys), kwargs)) + yield from original_async_stream(loader, keys, **kwargs) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed) + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream) + + comparison, cwb_report = analysis.ModelAnalysisLogic.execute( + "a", "b", "diffusion_models", _params() + ) + assert "Comparable floating tensors: 1" in comparison + assert "mae: 1" in comparison + assert "mse: 2" in comparison + assert "[A ONLY] (1)" in comparison + assert "[B ONLY] (1)" in comparison + assert "[SHAPE MISMATCH] (1)" in comparison + assert "[NON-FLOATING] (1)" in comparison + assert "model.blocks.0 |" in comparison + assert "CWB-DERIVED SIMILARITY SUMMARY" in cwb_report + assert "merge weighting" in cwb_report + assert "PRESET-IMPLIED CONTRIBUTIONS" not in cwb_report + assert "Average contribution" not in cwb_report + assert len(processed) == 4 + assert {key for _, key in processed} == {"model.blocks.0.weight", "integer"} + assert len(async_calls) == 4 + assert all(len(call[1]) == 1 for call in async_calls) + assert all(call[2]["batch_size"] == 1 for call in async_calls) + assert all(call[2]["prefetch_batches"] == 1 for call in async_calls) + assert not list(tmp_path.glob("**/*analysis*")) + + +def test_lora_cross_format_keys_compare_as_one_logical_pair(monkeypatch, tmp_path, analysis): + path_a = tmp_path / "lora_a.safetensors" + path_b = tmp_path / "lora_b.safetensors" + save_file({ + "diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + "diffusion_model.foo.lora_B.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + }, str(path_a)) + save_file({ + "base_model.model.diffusion_model.foo.lora_down.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + "base_model.model.diffusion_model.foo.lora_up.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + }, str(path_b)) + paths = {"a": str(path_a), "b": str(path_b)} + monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name]) + monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None) + monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None) + processed = [] + async_calls = [] + original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed + original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream + + def track_processed(loader, key): + processed.append((loader.filename, key)) + return original_mark_processed(loader, key) + + def track_async_stream(loader, keys, **kwargs): + async_calls.append((loader.filename, tuple(keys), kwargs)) + yield from original_async_stream(loader, keys, **kwargs) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed) + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream) + + comparison, cwb_report = analysis.ModelAnalysisLogic.execute( + "a", "b", "loras", _params(), lora_mode=True, + ) + assert "Comparable floating tensors: 2" in comparison + assert "[A ONLY] (0)" in comparison + assert "[B ONLY] (0)" in comparison + assert "foo.down |" in comparison + assert "foo.up |" in comparison + assert "foo.[down+up] |" in cwb_report + assert len(processed) == 4 + assert len({entry for entry in processed}) == 4 + assert len(async_calls) == 2 + assert all(len(call[1]) == 2 for call in async_calls) + + +def test_embedding_similarity_alignment_is_explicit_opt_in(analysis): + a = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + b = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) + aligned = analysis._compute_cwb(a, b, True).values() + indexed = analysis._compute_cwb(a, b, False).values() + assert aligned["alignment"] == pytest.approx(1.0) + assert aligned["alignment_gain"] == pytest.approx(1.0) + assert indexed["pair_mean"] == pytest.approx(0.0) + assert indexed["alignment"] is None + + +def test_cuda_oom_retries_once_on_cpu(monkeypatch, analysis): + calls = [] + + def fake_compute(a, b, device, allow_alignment, top_count): + calls.append(device) + if device == "cuda": + raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory") + return analysis.RawStats(finite=1, total=1), analysis.CWBStats(), [] + + monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute) + monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None) + result = analysis._analyze_with_fallback( + "layer", torch.ones(1), torch.ones(1), "cuda", False, 1, + ) + assert calls == ["cuda", "cpu"] + assert result[-1] is True + + +def test_non_oom_error_is_not_retried(monkeypatch, analysis): + calls = [] + + def fake_compute(*args): + calls.append(args[2]) + raise RuntimeError("unrelated failure") + + monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute) + with pytest.raises(RuntimeError, match="unrelated failure"): + analysis._analyze_with_fallback( + "layer", torch.ones(1), torch.ones(1), "cuda", False, 1, + ) + assert calls == ["cuda"] + + +def test_failed_cpu_retry_names_layer_and_chains_oom(monkeypatch, analysis): + def fake_compute(a, b, device, allow_alignment, top_count): + if device == "cuda": + raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory") + raise ValueError("CPU calculation failed") + + monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute) + monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None) + with pytest.raises(RuntimeError, match="CPU fallback failed while analyzing 'named.layer'") as caught: + analysis._analyze_with_fallback( + "named.layer", torch.ones(1), torch.ones(1), "cuda", False, 1, + ) + assert isinstance(caught.value.__cause__, torch.cuda.OutOfMemoryError)