795 lines
43 KiB
Python
795 lines
43 KiB
Python
"""
|
|
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 torch
|
|
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
|
|
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
|
|
from .layer_parameters import parameter_input, resolve_layer_parameters
|
|
|
|
# Import SVD rank utilities from standard lora extract file to avoid duplication
|
|
from .lora_extract_svd import (
|
|
_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,
|
|
knee_probe_offset: int = 32,
|
|
include_mode: bool = False,
|
|
layer_parameters=None,
|
|
) -> 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 = sorted(k for k in keys_a if k.endswith(".weight"))
|
|
if mode == "fixed":
|
|
profile = "extract:fixed"
|
|
defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
|
|
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
|
|
"clamp_quantile": clamp_quantile, "min_diff": min_diff}
|
|
elif mode == "ratio":
|
|
profile = "extract:ratio"
|
|
defaults = {"linear_ratio": linear_param, "conv_ratio": conv_param,
|
|
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
|
|
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
|
|
elif mode == "quantile":
|
|
profile = "extract:quantile"
|
|
defaults = {"linear_quantile": linear_param, "conv_quantile": conv_param,
|
|
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
|
|
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
|
|
elif mode == "sv_fro":
|
|
profile = "extract:frobenius"
|
|
defaults = {"linear_target": linear_param, "conv_target": conv_param,
|
|
"clamp_quantile": clamp_quantile, "min_diff": min_diff,
|
|
"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank}
|
|
else:
|
|
profile = "extract:knee"
|
|
defaults = {"linear_max_rank": linear_max_rank, "conv_max_rank": conv_max_rank,
|
|
"clamp_quantile": clamp_quantile, "min_diff": min_diff}
|
|
resolved_parameters = resolve_layer_parameters(
|
|
layer_parameters, profile, weight_keys, defaults, node_name="Text Encoder Extract"
|
|
)
|
|
work_units = [
|
|
(key, key, key if key in keys_b else None)
|
|
for key in weight_keys
|
|
if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
|
|
and not (key not in keys_b and mismatch_mode == "skip")
|
|
]
|
|
pbar = comfy.utils.ProgressBar(len(work_units))
|
|
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
|
|
|
|
def _process_layer(key, cpu_a, cpu_b):
|
|
values = resolved_parameters.get(key, defaults)
|
|
if mode == "fixed":
|
|
layer_linear_param, layer_conv_param = values["linear_dim"], values["conv_dim"]
|
|
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
|
|
elif mode == "ratio":
|
|
layer_linear_param, layer_conv_param = values["linear_ratio"], values["conv_ratio"]
|
|
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
|
|
elif mode == "quantile":
|
|
layer_linear_param, layer_conv_param = values["linear_quantile"], values["conv_quantile"]
|
|
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
|
|
elif mode == "sv_fro":
|
|
layer_linear_param, layer_conv_param = values["linear_target"], values["conv_target"]
|
|
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
|
|
else:
|
|
layer_linear_param, layer_conv_param = linear_param, conv_param
|
|
layer_linear_max_rank, layer_conv_max_rank = values["linear_max_rank"], values["conv_max_rank"]
|
|
layer_clamp_quantile = values["clamp_quantile"]
|
|
layer_min_diff = values["min_diff"]
|
|
lora_name = _format_te_lora_key(key)
|
|
weight_diff = None
|
|
layer_device = device
|
|
|
|
if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)):
|
|
return "skipped", None
|
|
|
|
# Load tensors with pinned memory for 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")
|
|
|
|
if is_dora:
|
|
weight_diff, dora_scale, layer_device = dora_difference(
|
|
cpu_a, None, device
|
|
)
|
|
else:
|
|
weight_diff, layer_device = lora_difference(cpu_a, None, device)
|
|
dora_scale = None
|
|
else:
|
|
if cpu_a.shape != cpu_b.shape:
|
|
if mismatch_mode == "skip":
|
|
return "skipped", None
|
|
if mismatch_mode == "error":
|
|
raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
|
|
|
|
if is_dora:
|
|
weight_diff, dora_scale, layer_device = dora_difference(
|
|
cpu_a, None, device
|
|
)
|
|
else:
|
|
weight_diff, layer_device = lora_difference(cpu_a, None, device)
|
|
dora_scale = None
|
|
else:
|
|
if cpu_a.ndim < 2:
|
|
return "skipped", None
|
|
if is_dora:
|
|
weight_diff, dora_scale, layer_device = dora_difference(
|
|
cpu_a, cpu_b, device
|
|
)
|
|
else:
|
|
weight_diff, layer_device = lora_difference(
|
|
cpu_a, cpu_b, device
|
|
)
|
|
dora_scale = None
|
|
|
|
if weight_diff is None:
|
|
raise RuntimeError(f"No extraction tensor was produced for {key}")
|
|
|
|
# Skip small differences
|
|
if layer_min_diff > 0 and weight_diff.abs().max() < layer_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 = {}
|
|
svd_weight = weight_diff
|
|
|
|
try:
|
|
if is_conv:
|
|
result, mode_str = retry_cuda_oom_on_cpu(
|
|
lambda: _svd_extract_conv(svd_weight, mode, layer_conv_param, layer_device, layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
|
|
lambda: _svd_extract_conv(svd_weight.cpu(), mode, layer_conv_param, "cpu", layer_conv_max_rank, layer_clamp_quantile, knee_probe_offset),
|
|
layer_device,
|
|
)
|
|
else:
|
|
result, mode_str = retry_cuda_oom_on_cpu(
|
|
lambda: _svd_extract_linear(svd_weight, mode, layer_linear_param, layer_device, layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
|
|
lambda: _svd_extract_linear(svd_weight.cpu(), mode, layer_linear_param, "cpu", layer_linear_max_rank, layer_clamp_quantile, svd_niter, knee_probe_offset),
|
|
layer_device,
|
|
)
|
|
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, layer_linear_param, layer_device,
|
|
layer_linear_max_rank, knee_probe_offset, layer_clamp_quantile,
|
|
)
|
|
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_context = atomic_uel_writer(output_path)
|
|
writer = writer_context.__enter__()
|
|
try:
|
|
stream = paired_async_tensors(
|
|
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
|
|
)
|
|
for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting TE Layers", unit="layers"):
|
|
status, layer_sd = _process_layer(key, cpu_a, cpu_b)
|
|
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)
|
|
stats["skipped"] += len(weight_keys) - len(work_units)
|
|
except BaseException as exc:
|
|
writer_context.__exit__(type(exc), exc, exc.__traceback__)
|
|
raise
|
|
else:
|
|
writer_context.__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) -> tuple[str, str]:
|
|
"""Build output path for LoRA file."""
|
|
return canonical_model_artifact_path("loras", output_filename)
|
|
|
|
|
|
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=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."),
|
|
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", tooltip="Handle missing or incompatible text-encoder tensors by skipping them, substituting zeros where supported, or aborting."),
|
|
io.String.Input("output_filename", default="extracted_te_lora", tooltip="Output filename without extension, written under the ComfyUI LoRA directory."),
|
|
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save extracted text-encoder LoRA factors."),
|
|
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer extraction arithmetic; CUDA out-of-memory retries the affected layer on CPU where supported."),
|
|
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."),
|
|
io.Boolean.Input("include_mode", default=False, tooltip="Use Skip Patterns as a whitelist instead. Only matching layers are extracted; an empty whitelist extracts nothing."),
|
|
]
|
|
|
|
|
|
# =============================================================================
|
|
# 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(),
|
|
parameter_input("extract:fixed"),
|
|
],
|
|
outputs=[io.AnyType.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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
|
|
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
|
|
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:ratio"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
knee_probe_offset=probe_offset,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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, tooltip="Target cumulative singular-value fraction for linear layers."),
|
|
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
|
|
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:quantile"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
knee_probe_offset=probe_offset,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
|
|
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:knee"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
|
|
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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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, tooltip="Target Frobenius-norm fraction retained for linear layers."),
|
|
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
|
|
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:frobenius"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
knee_probe_offset=probe_offset,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
# =============================================================================
|
|
# 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(),
|
|
parameter_input("extract:fixed"),
|
|
],
|
|
outputs=[io.AnyType.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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
|
|
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
|
|
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:ratio"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
knee_probe_offset=probe_offset,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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, tooltip="Target cumulative singular-value fraction for linear layers."),
|
|
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
|
|
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:quantile"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
knee_probe_offset=probe_offset,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
|
|
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:knee"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
|
|
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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, knee_probe_offset=knee_probe_offset, include_mode=include_mode,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|
|
|
|
|
|
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, tooltip="Target Frobenius-norm fraction retained for linear layers."),
|
|
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
|
|
io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
|
|
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
|
|
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
|
|
*_get_common_inputs(),
|
|
parameter_input("extract:frobenius"),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="output_path")],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, 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, include_mode=False, layer_parameters=None) -> 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, output_name = _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, include_mode=include_mode,
|
|
knee_probe_offset=probe_offset,
|
|
layer_parameters=layer_parameters,
|
|
)
|
|
return io.NodeOutput(output_name)
|