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

add lycoris diff extraction

Browse Source
This commit is contained in:
silveroxides
2026-01-01 07:06:19 +01:00
parent 92d3c5582c
commit da13f9f51d
5 changed files with 495 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+4 -1
View File
@@ -22,6 +22,7 @@ from .nodes.merger import (
from .nodes.lora_extract import (
LoRASaveFromFile, LoRACheckpointSaveFromFile
)
from .nodes.lycoris_extract import LyCORISExtractFromFile
class ModelUtilsExtension(ComfyExtension):
@@ -45,8 +46,10 @@ class ModelUtilsExtension(ComfyExtension):
CheckpointThreeMerger, EmbeddingThreeMerger,
# LoRA Extraction
LoRASaveFromFile, LoRACheckpointSaveFromFile,
# LyCORIS Extraction
LyCORISExtractFromFile,
]
async def comfy_entrypoint() -> ModelUtilsExtension:
return ModelUtilsExtension()
return ModelUtilsExtension()
nodes/lora_extract.py
+85 -97
View File
@@ -51,12 +51,15 @@ def extract_lora(diff, rank):
def _compile_patterns(pattern_string: str) -> list:
"""Compiles whitespace-separated regex patterns into a list of compiled patterns."""
"""Compiles newline/whitespace-separated regex patterns into a list of compiled patterns."""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for pattern in pattern_string.split():
pattern = pattern.strip()
if not pattern:
continue
try:
patterns.append(re.compile(pattern))
except re.error as e:
@@ -77,13 +80,12 @@ def extract_lora_from_files(
model_b_path: str,
rank: int,
lora_type: str,
bias_diff: bool,
prefix_lora: str,
process_device: str,
save_dtype: str,
exclude_patterns_str: str = "",
discard_patterns_str: str = "",
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
include_biases: bool = False,
) -> dict[str, torch.Tensor]:
"""
Memory-efficient LoRA extraction by streaming tensors from safetensors files.
@@ -92,14 +94,13 @@ def extract_lora_from_files(
model_a_path: Path to the finetuned/modified model
model_b_path: Path to the base/original model
rank: LoRA rank for SVD decomposition
lora_type: "standard" for LoRA, "full_diff" for full difference
bias_diff: Whether to include bias differences
lora_type: "standard" for LoRA (only 2D+ weights), "full_diff" for full difference
prefix_lora: Prefix for LoRA keys (e.g., "diffusion_model.")
process_device: Device for processing ("cuda" or "cpu")
save_dtype: Output dtype ("fp16", "bf16", "fp32")
exclude_patterns_str: Regex patterns to exclude from extraction (use model A only)
discard_patterns_str: Regex patterns to discard entirely (no output)
mismatch_mode: "skip" (use model A), "zeros" (use zeros), "error" (raise error)
skip_patterns_str: Regex patterns for layers to skip (no LoRA extracted)
mismatch_mode: "skip" (ignore), "zeros" (use model A), "error" (raise error)
include_biases: Whether to include bias differences (only for full_diff mode)
Returns:
Dictionary of LoRA tensors
@@ -114,8 +115,7 @@ def extract_lora_from_files(
}.get(save_dtype, torch.float16)
# Compile patterns
exclude_patterns = _compile_patterns(exclude_patterns_str)
discard_patterns = _compile_patterns(discard_patterns_str)
skip_patterns = _compile_patterns(skip_patterns_str)
# Open both files with memory-efficient handler
handler_a = MemoryEfficientSafeOpen(model_a_path)
@@ -127,7 +127,10 @@ def extract_lora_from_files(
# Get all weight keys from model A (primary)
all_weight_keys = [k for k in keys_a if k.endswith(".weight")]
all_bias_keys = [k for k in keys_a if k.endswith(".bias")] if bias_diff else []
# Only process biases in full_diff mode with include_biases=True
all_bias_keys = []
if lora_type == "full_diff" and include_biases:
all_bias_keys = [k for k in keys_a if k.endswith(".bias")]
# Log key differences
missing_in_b = keys_a - keys_b
@@ -140,27 +143,17 @@ def extract_lora_from_files(
pbar = comfy.utils.ProgressBar(len(all_weight_keys) + len(all_bias_keys))
# Stats
skipped_keys = 0
excluded_keys = 0
discarded_keys = 0
skipped_by_pattern = 0
skipped_by_mismatch = 0
skipped_1d_weights = 0
shape_mismatch_keys = 0
for key in tqdm(all_weight_keys, desc="Extracting LoRA weights", unit="layers"):
lora_key = key[:-7] # Remove ".weight"
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
pbar.update(1)
continue
# Load model A tensor
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
# Check exclude patterns - skip extraction, no output for this key
if _matches_any_pattern(key, exclude_patterns):
excluded_keys += 1
del tensor_a
# Check skip patterns - skip entirely, no LoRA for this layer
if _matches_any_pattern(key, skip_patterns):
skipped_by_pattern += 1
pbar.update(1)
continue
@@ -169,111 +162,112 @@ def extract_lora_from_files(
if mismatch_mode == "error":
raise ValueError(f"Key '{key}' not found in model B (mismatch_mode='error')")
elif mismatch_mode == "skip":
skipped_keys += 1
del tensor_a
skipped_by_mismatch += 1
pbar.update(1)
continue
# zeros mode: use tensor_a as-is (diff with zeros = tensor_a)
# zeros mode: load tensor A (diff with zeros = tensor_a)
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
weight_diff = tensor_a
del tensor_a
else:
# Load both tensors
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
# Check shape mismatch
if tensor_a.shape != tensor_b.shape:
del tensor_b
del tensor_a, tensor_b
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for '{key}': {tensor_a.shape} vs {tensor_b.shape}")
elif mismatch_mode == "skip":
shape_mismatch_keys += 1
del tensor_a
pbar.update(1)
continue
# zeros mode: use tensor_a as-is
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
weight_diff = tensor_a
del tensor_a
else:
# Compute difference: finetuned - base
weight_diff = tensor_a - tensor_b
del tensor_b
del tensor_a
del tensor_a, tensor_b
if lora_type == "standard":
# Standard LoRA: Only decompose 2D+ weights, skip 1D entirely
if weight_diff.ndim < 2:
# 1D weights (e.g., LayerNorm) - store as diff if bias_diff enabled
if bias_diff:
output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu()
else:
try:
lora_up, lora_down = extract_lora(weight_diff, rank)
output_sd[f"{prefix_lora}{lora_key}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu()
output_sd[f"{prefix_lora}{lora_key}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu()
except Exception as e:
print(f"[LoRA Extract] Could not extract LoRA for {key}: {e}")
skipped_1d_weights += 1
del weight_diff
pbar.update(1)
continue
try:
lora_up, lora_down = extract_lora(weight_diff, rank)
output_sd[f"{prefix_lora}{lora_key}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu()
output_sd[f"{prefix_lora}{lora_key}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu()
except Exception as e:
print(f"[LoRA Extract] Could not extract LoRA for {key}: {e}")
elif lora_type == "full_diff":
# Full diff: Store all weight differences
output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu()
del weight_diff
pbar.update(1)
# Process bias keys
for key in tqdm(all_bias_keys, desc="Extracting LoRA biases", unit="layers"):
# Process bias keys (only in full_diff mode with include_biases=True)
for key in tqdm(all_bias_keys, desc="Extracting bias differences", unit="layers"):
lora_key = key[:-5] # Remove ".bias"
# Check discard patterns
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
# Check skip patterns
if _matches_any_pattern(key, skip_patterns):
skipped_by_pattern += 1
pbar.update(1)
continue
# Check exclude patterns
if _matches_any_pattern(key, exclude_patterns):
excluded_keys += 1
pbar.update(1)
continue
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
if key not in keys_b:
if mismatch_mode == "error":
raise ValueError(f"Key '{key}' not found in model B (mismatch_mode='error')")
elif mismatch_mode == "skip":
skipped_keys += 1
del tensor_a
skipped_by_mismatch += 1
pbar.update(1)
continue
# zeros mode
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
bias_diff_tensor = tensor_a
del tensor_a
else:
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
if tensor_a.shape != tensor_b.shape:
del tensor_b
del tensor_a, tensor_b
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for '{key}': {tensor_a.shape} vs {tensor_b.shape}")
elif mismatch_mode == "skip":
shape_mismatch_keys += 1
del tensor_a
pbar.update(1)
continue
# zeros mode
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
bias_diff_tensor = tensor_a
del tensor_a
else:
bias_diff_tensor = tensor_a - tensor_b
del tensor_b
del tensor_a, tensor_b
del tensor_a
output_sd[f"{prefix_lora}{lora_key}.diff_b"] = bias_diff_tensor.to(save_torch_dtype).cpu()
del bias_diff_tensor
pbar.update(1)
# Log summary
if excluded_keys > 0:
print(f"[LoRA Extract] Excluded {excluded_keys} keys from extraction")
if discarded_keys > 0:
print(f"[LoRA Extract] Discarded {discarded_keys} keys from output")
if skipped_keys > 0:
print(f"[LoRA Extract] Skipped {skipped_keys} keys (missing in model B)")
if skipped_by_pattern > 0:
print(f"[LoRA Extract] Skipped {skipped_by_pattern} keys (matched skip patterns)")
if skipped_by_mismatch > 0:
print(f"[LoRA Extract] Skipped {skipped_by_mismatch} keys (missing in model B)")
if shape_mismatch_keys > 0:
print(f"[LoRA Extract] Skipped {shape_mismatch_keys} keys (shape mismatch)")
if skipped_1d_weights > 0:
print(f"[LoRA Extract] Skipped {skipped_1d_weights} 1D weights (cannot SVD decompose)")
finally:
handler_a.__exit__(None, None, None)
@@ -313,20 +307,19 @@ class LoRASaveFromFile(io.ComfyNode):
io.Int.Input("rank", default=64, min=1, max=4096, step=1,
tooltip="LoRA rank - higher preserves more detail but increases size"),
io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard",
tooltip="standard: SVD decomposition to lora_up/lora_down, full_diff: raw weight differences"),
io.Boolean.Input("bias_diff", default=True,
tooltip="Include bias differences in the LoRA"),
tooltip="standard: SVD to lora_up/down (proper LoRA format), "
"full_diff: raw weight differences (.diff format)"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip",
tooltip="How to handle missing keys or shape mismatches between models"),
io.Boolean.Input("include_biases", default=False,
tooltip="Include bias differences (only applies to full_diff mode)"),
io.String.Input("output_filename", default="extracted_lora",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
io.String.Input("exclude_patterns", default="", multiline=True,
tooltip="Regex patterns to exclude from extraction (one per line, whitespace-separated)"),
io.String.Input("discard_patterns", default="", multiline=True,
tooltip="Regex patterns to discard entirely (no output for matching keys)"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip (no LoRA extracted for matching keys)"),
],
outputs=[
io.String.Output(display_name="output_path"),
@@ -341,13 +334,12 @@ class LoRASaveFromFile(io.ComfyNode):
model_b: str,
rank: int,
lora_type: str,
bias_diff: bool,
mismatch_mode: str,
include_biases: bool,
output_filename: str,
save_dtype: str,
process_device: str,
exclude_patterns: str,
discard_patterns: str,
skip_patterns: str,
) -> io.NodeOutput:
# Get full paths
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
@@ -359,13 +351,12 @@ class LoRASaveFromFile(io.ComfyNode):
model_b_path=model_b_path,
rank=rank,
lora_type=lora_type,
bias_diff=bias_diff,
prefix_lora="diffusion_model.",
process_device=process_device,
save_dtype=save_dtype,
exclude_patterns_str=exclude_patterns,
discard_patterns_str=discard_patterns,
skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode,
include_biases=include_biases,
)
if not output_sd:
@@ -409,20 +400,19 @@ class LoRACheckpointSaveFromFile(io.ComfyNode):
io.Int.Input("rank", default=64, min=1, max=4096, step=1,
tooltip="LoRA rank - higher preserves more detail but increases size"),
io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard",
tooltip="standard: SVD decomposition, full_diff: raw weight differences"),
io.Boolean.Input("bias_diff", default=True,
tooltip="Include bias differences in the LoRA"),
tooltip="standard: SVD to lora_up/down (proper LoRA format), "
"full_diff: raw weight differences (.diff format)"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip",
tooltip="How to handle missing keys or shape mismatches between models"),
io.Boolean.Input("include_biases", default=False,
tooltip="Include bias differences (only applies to full_diff mode)"),
io.String.Input("output_filename", default="extracted_lora",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
io.String.Input("exclude_patterns", default="", multiline=True,
tooltip="Regex patterns to exclude from extraction"),
io.String.Input("discard_patterns", default="", multiline=True,
tooltip="Regex patterns to discard entirely"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip (no LoRA extracted for matching keys)"),
],
outputs=[
io.String.Output(display_name="output_path"),
@@ -437,13 +427,12 @@ class LoRACheckpointSaveFromFile(io.ComfyNode):
checkpoint_b: str,
rank: int,
lora_type: str,
bias_diff: bool,
mismatch_mode: str,
include_biases: bool,
output_filename: str,
save_dtype: str,
process_device: str,
exclude_patterns: str,
discard_patterns: str,
skip_patterns: str,
) -> io.NodeOutput:
# Get full paths
checkpoint_a_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_a)
@@ -469,13 +458,12 @@ class LoRACheckpointSaveFromFile(io.ComfyNode):
model_b_path=checkpoint_b_path,
rank=rank,
lora_type=lora_type,
bias_diff=bias_diff,
prefix_lora="diffusion_model.",
process_device=process_device,
save_dtype=save_dtype,
exclude_patterns_str=exclude_patterns,
discard_patterns_str=discard_patterns,
skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode,
include_biases=include_biases,
)
output_sd.update(model_lora)
nodes/lycoris_extract.py
+374
View File
@@ -0,0 +1,374 @@
"""
LyCORIS-based LoRA extraction node.
Uses the LyCORIS library for advanced extraction modes (threshold, ratio, quantile)
with our memory-efficient tensor streaming.
"""
import os
import re
import torch
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from safetensors.torch import save_file
from .merger_utils import MemoryEfficientSafeOpen
try:
from lycoris.utils import extract_linear, extract_conv
LYCORIS_AVAILABLE = True
except ImportError:
LYCORIS_AVAILABLE = False
print("[LyCORIS Extract] Warning: lycoris not installed. Install with: pip install lycoris")
def _compile_patterns(pattern_string: str) -> list:
"""Compiles newline/whitespace-separated regex patterns into a list of compiled patterns."""
if not pattern_string or not pattern_string.strip():
return []
patterns = []
for pattern in pattern_string.split():
pattern = pattern.strip()
if not pattern:
continue
try:
patterns.append(re.compile(pattern))
except re.error as e:
raise ValueError(f"Invalid regex pattern '{pattern}': {e}")
return patterns
def _matches_any_pattern(key: str, patterns: list) -> bool:
"""Returns True if key matches any of the compiled regex patterns (substring search)."""
for pattern in patterns:
if pattern.search(key):
return True
return False
def lycoris_extract_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
prefix_lora: str,
process_device: str,
save_dtype: str,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
use_cp: bool = True,
) -> dict[str, torch.Tensor]:
"""
Memory-efficient LyCORIS extraction by streaming tensors from safetensors files.
Uses LyCORIS's extract_linear and extract_conv functions for proper decomposition.
Args:
model_a_path: Path to the finetuned/modified model
model_b_path: Path to the base/original model
mode: Extraction mode - "fixed", "threshold", "ratio", "quantile", "full"
linear_param: Mode parameter for linear layers (rank/threshold/ratio/quantile)
conv_param: Mode parameter for conv layers (rank/threshold/ratio/quantile)
prefix_lora: Prefix for LoRA keys (e.g., "lora_unet_")
process_device: Device for processing ("cuda" or "cpu")
save_dtype: Output dtype ("fp16", "bf16", "fp32")
skip_patterns_str: Regex patterns for layers to skip
mismatch_mode: "skip" (ignore), "zeros" (use model A), "error" (raise error)
use_cp: Use CP decomposition for 3x3 convolutions
Returns:
Dictionary of LyCORIS tensors
"""
if not LYCORIS_AVAILABLE:
raise ImportError("lycoris library not installed. Install with: pip install lycoris")
output_sd = {}
# Get dtype for saving
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16
}.get(save_dtype, torch.float16)
# Compile patterns
skip_patterns = _compile_patterns(skip_patterns_str)
# Open both files with memory-efficient handler
handler_a = MemoryEfficientSafeOpen(model_a_path)
handler_b = MemoryEfficientSafeOpen(model_b_path)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
# Get all weight keys from model A (primary)
all_weight_keys = [k for k in keys_a if k.endswith(".weight")]
# Log key differences
missing_in_b = keys_a - keys_b
extra_in_b = keys_b - keys_a
if missing_in_b:
print(f"[LyCORIS Extract] {len(missing_in_b)} keys in model A not in model B")
if extra_in_b:
print(f"[LyCORIS Extract] {len(extra_in_b)} keys in model B not in model A (ignored)")
pbar = comfy.utils.ProgressBar(len(all_weight_keys))
# Stats
skipped_by_pattern = 0
skipped_by_mismatch = 0
skipped_1d_weights = 0
shape_mismatch_keys = 0
full_weights = 0
low_rank_weights = 0
for key in tqdm(all_weight_keys, desc="Extracting LyCORIS weights", unit="layers"):
lora_key = key[:-7] # Remove ".weight"
# Convert key format: double_blocks.0.img_attn.proj -> lora_unet_double_blocks_0_img_attn_proj
lora_name = prefix_lora + lora_key.replace(".", "_")
# Check skip patterns
if _matches_any_pattern(key, skip_patterns):
skipped_by_pattern += 1
pbar.update(1)
continue
# Check if key exists in model B
if key not in keys_b:
if mismatch_mode == "error":
raise ValueError(f"Key '{key}' not found in model B (mismatch_mode='error')")
elif mismatch_mode == "skip":
skipped_by_mismatch += 1
pbar.update(1)
continue
# zeros mode: diff = tensor_a - 0 = tensor_a
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
weight_diff = tensor_a
del tensor_a
else:
# Load both tensors
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32)
# Check shape mismatch
if tensor_a.shape != tensor_b.shape:
del tensor_a, tensor_b
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for '{key}': {tensor_a.shape} vs {tensor_b.shape}")
elif mismatch_mode == "skip":
shape_mismatch_keys += 1
pbar.update(1)
continue
# zeros mode
tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32)
weight_diff = tensor_a
del tensor_a
else:
# Compute difference: finetuned - base
weight_diff = tensor_a - tensor_b
del tensor_a, tensor_b
# Skip 1D weights (norms, biases can't be decomposed)
if weight_diff.ndim < 2:
skipped_1d_weights += 1
del weight_diff
pbar.update(1)
continue
# Determine if this is a conv or linear layer
is_conv = weight_diff.ndim == 4
is_linear_conv = is_conv and weight_diff.shape[2] == 1 and weight_diff.shape[3] == 1
try:
if is_conv:
# Use conv extraction
mode_param = linear_param if is_linear_conv else conv_param
result, decompose_mode = extract_conv(
weight_diff,
mode=mode,
mode_param=int(mode_param) if mode == "fixed" else mode_param,
device=process_device,
is_cp=use_cp and not is_linear_conv,
)
else:
# Use linear extraction
result, decompose_mode = extract_linear(
weight_diff,
mode=mode,
mode_param=int(linear_param) if mode == "fixed" else linear_param,
device=process_device,
)
if decompose_mode == "full":
# Store as full diff
output_sd[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu()
full_weights += 1
elif decompose_mode == "low rank":
extract_a, extract_b, diff = result
# Check if CP decomposition was applied (has 3 components for conv)
if is_conv and not is_linear_conv and use_cp and hasattr(result, '__len__') and len(result) == 3:
# Check if extract_a has been further decomposed
if extract_a.ndim == 4 and extract_a.shape[2] > 1:
# Re-do with CP decomposition
dim = extract_a.size(0)
(extract_c, extract_a_new, _), _ = extract_conv(
extract_a.transpose(0, 1),
"fixed",
dim,
process_device,
True,
)
extract_a_new = extract_a_new.transpose(0, 1)
extract_c = extract_c.transpose(0, 1)
output_sd[f"{lora_name}.lora_mid.weight"] = extract_c.to(save_torch_dtype).cpu()
extract_a = extract_a_new
output_sd[f"{lora_name}.lora_down.weight"] = extract_a.to(save_torch_dtype).cpu()
output_sd[f"{lora_name}.lora_up.weight"] = extract_b.to(save_torch_dtype).cpu()
output_sd[f"{lora_name}.alpha"] = torch.tensor([extract_a.shape[0]]).to(save_torch_dtype)
low_rank_weights += 1
del extract_a, extract_b, diff
except Exception as e:
print(f"[LyCORIS Extract] Could not extract for {key}: {e}")
del weight_diff
pbar.update(1)
# Log summary
print(f"[LyCORIS Extract] Extracted {low_rank_weights} low-rank, {full_weights} full weights")
if skipped_by_pattern > 0:
print(f"[LyCORIS Extract] Skipped {skipped_by_pattern} keys (matched skip patterns)")
if skipped_by_mismatch > 0:
print(f"[LyCORIS Extract] Skipped {skipped_by_mismatch} keys (missing in model B)")
if shape_mismatch_keys > 0:
print(f"[LyCORIS Extract] Skipped {shape_mismatch_keys} keys (shape mismatch)")
if skipped_1d_weights > 0:
print(f"[LyCORIS Extract] Skipped {skipped_1d_weights} 1D weights")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
return output_sd
class LyCORISExtractFromFile(io.ComfyNode):
"""
Memory-efficient LyCORIS extraction from safetensors files.
Uses the LyCORIS library for advanced extraction modes while streaming tensors
one at a time for memory efficiency.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LyCORISExtractFromFile",
display_name="Extract LyCORIS from Files (Memory Efficient)",
category="ModelUtils/LoRA",
description="Extracts a LyCORIS/LoCoN by computing the difference between two model files. "
"Uses LyCORIS library for advanced extraction modes (threshold, ratio, quantile).",
inputs=[
io.Combo.Input(
"model_a",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="The finetuned/modified model (A - B = LoRA)"
),
io.Combo.Input(
"model_b",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="The base/original model (A - B = LoRA)"
),
io.Combo.Input("mode", options=["fixed", "threshold", "ratio", "quantile", "full"], default="fixed",
tooltip="Extraction mode: fixed=specific rank, threshold/ratio/quantile=dynamic rank, full=no decomposition"),
io.Int.Input("linear_dim", default=64, min=1, max=4096, step=1,
tooltip="Rank for linear layers (fixed mode) or ignored in other modes"),
io.Int.Input("conv_dim", default=32, min=1, max=4096, step=1,
tooltip="Rank for conv layers (fixed mode) or ignored in other modes"),
io.Float.Input("linear_threshold", default=0.0, min=0.0, max=1.0, step=0.01,
tooltip="Threshold/ratio/quantile for linear layers (non-fixed modes)"),
io.Float.Input("conv_threshold", default=0.0, min=0.0, max=1.0, step=0.01,
tooltip="Threshold/ratio/quantile for conv layers (non-fixed modes)"),
io.Boolean.Input("use_cp", default=True,
tooltip="Use CP decomposition for 3x3 convolutions (reduces size)"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip",
tooltip="How to handle missing keys or shape mismatches between models"),
io.String.Input("output_filename", default="extracted_lycoris",
tooltip="Output filename (without extension)"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda",
tooltip="Device for SVD computation"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip"),
],
outputs=[
io.String.Output(display_name="output_path"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_a: str,
model_b: str,
mode: str,
linear_dim: int,
conv_dim: int,
linear_threshold: float,
conv_threshold: float,
use_cp: bool,
mismatch_mode: str,
output_filename: str,
save_dtype: str,
process_device: str,
skip_patterns: str,
) -> io.NodeOutput:
if not LYCORIS_AVAILABLE:
raise ImportError("lycoris library not installed. Install with: pip install lycoris")
# Get full paths
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
# Determine parameters based on mode
if mode == "fixed":
linear_param = linear_dim
conv_param = conv_dim
else:
linear_param = linear_threshold
conv_param = conv_threshold
# Extract LyCORIS
output_sd = lycoris_extract_from_files(
model_a_path=model_a_path,
model_b_path=model_b_path,
mode=mode,
linear_param=linear_param,
conv_param=conv_param,
prefix_lora="lora_unet_",
process_device=process_device,
save_dtype=save_dtype,
skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode,
use_cp=use_cp,
)
if not output_sd:
raise ValueError("No LyCORIS weights extracted - models may be identical or incompatible")
# Save to loras folder
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path)
print(f"[LyCORIS Extract] Saved {len(output_sd)} tensors to {output_path}")
return io.NodeOutput(output_path)
nodes/merger_utils.py
+31 -5
View File
@@ -1,18 +1,33 @@
import mmap
import json
import struct
import torch
class MemoryEfficientSafeOpen:
def __init__(self, filename, device='cpu'):
"""Memory-efficient safetensors file reader with optional mmap zero-copy support.
When mmap_mode=True, tensor data is read directly from disk via memory-mapped file,
avoiding intermediate copying. This is especially efficient for large models.
"""
def __init__(self, filename, device='cpu', mmap_mode=True):
self.filename = filename
self.device = device
self.mmap_mode = mmap_mode
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
self.mmap_obj = None
if self.mmap_mode:
self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.mmap_obj:
self.mmap_obj.close()
self.file.close()
def keys(self):
@@ -25,10 +40,21 @@ class MemoryEfficientSafeOpen:
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
if self.mmap_mode and self.mmap_obj:
if offset_start != offset_end:
# Calculate absolute offset in file
start = self.header_size + 8 + offset_start
end = self.header_size + 8 + offset_end
# Create tensor from memory view (zero-copy)
# Note: This tensor is valid only while the file/mmap is open
tensor_bytes = memoryview(self.mmap_obj)[start:end]
else:
tensor_bytes = None
else:
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
return self._deserialize_tensor(tensor_bytes, metadata)
requirements.txt
+1
View File
@@ -0,0 +1 @@
lycoris