From 18679cca2cfa1220546754d7e1f975ef93805a19 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Thu, 1 Jan 2026 07:21:37 +0100 Subject: [PATCH] Split lycoris extraction modes into different nodes --- __init__.py | 9 +- nodes/lycoris_extract.py | 437 ++++++++++++++++++++++++++++----------- requirements.txt | 2 +- 3 files changed, 320 insertions(+), 128 deletions(-) diff --git a/__init__.py b/__init__.py index 6d3e6d7..9258e61 100644 --- a/__init__.py +++ b/__init__.py @@ -22,7 +22,10 @@ from .nodes.merger import ( from .nodes.lora_extract import ( LoRASaveFromFile, LoRACheckpointSaveFromFile ) -from .nodes.lycoris_extract import LyCORISExtractFromFile +from .nodes.lycoris_extract import ( + LyCORISExtractFixed, LyCORISExtractThreshold, + LyCORISExtractRatio, LyCORISExtractQuantile, LyCORISExtractFull +) class ModelUtilsExtension(ComfyExtension): @@ -47,9 +50,11 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Extraction LoRASaveFromFile, LoRACheckpointSaveFromFile, # LyCORIS Extraction - LyCORISExtractFromFile, + LyCORISExtractFixed, LyCORISExtractThreshold, + LyCORISExtractRatio, LyCORISExtractQuantile, LyCORISExtractFull, ] async def comfy_entrypoint() -> ModelUtilsExtension: return ModelUtilsExtension() + diff --git a/nodes/lycoris_extract.py b/nodes/lycoris_extract.py index cd9a7db..4cb1946 100644 --- a/nodes/lycoris_extract.py +++ b/nodes/lycoris_extract.py @@ -1,8 +1,8 @@ """ -LyCORIS-based LoRA extraction node. +LyCORIS-based LoRA extraction nodes. -Uses the LyCORIS library for advanced extraction modes (threshold, ratio, quantile) -with our memory-efficient tensor streaming. +Uses the LyCORIS library for advanced extraction modes with our memory-efficient tensor streaming. +Provides separate nodes for each extraction mode with mode-specific parameters. """ import os import re @@ -19,7 +19,7 @@ try: LYCORIS_AVAILABLE = True except ImportError: LYCORIS_AVAILABLE = False - print("[LyCORIS Extract] Warning: lycoris not installed. Install with: pip install lycoris") + print("[LyCORIS Extract] Warning: lycoris not installed. Install with: pip install lycoris-lora") def _compile_patterns(pattern_string: str) -> list: @@ -40,7 +40,7 @@ def _compile_patterns(pattern_string: str) -> list: def _matches_any_pattern(key: str, patterns: list) -> bool: - """Returns True if key matches any of the compiled regex patterns (substring search).""" + """Returns True if key matches any of the compiled regex patterns.""" for pattern in patterns: if pattern.search(key): return True @@ -59,44 +59,40 @@ def lycoris_extract_from_files( skip_patterns_str: str = "", mismatch_mode: str = "skip", use_cp: bool = True, + use_sparse_bias: bool = False, + sparsity: float = 0.98, ) -> 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") + linear_param: Mode parameter for linear layers + conv_param: Mode parameter for conv layers + prefix_lora: Prefix for LoRA keys + process_device: Device for processing + save_dtype: Output dtype skip_patterns_str: Regex patterns for layers to skip - mismatch_mode: "skip" (ignore), "zeros" (use model A), "error" (raise error) + mismatch_mode: How to handle mismatches use_cp: Use CP decomposition for 3x3 convolutions - - Returns: - Dictionary of LyCORIS tensors + use_sparse_bias: Enable sparse bias storage + sparsity: Sparsity threshold for sparse bias (0-1) """ if not LYCORIS_AVAILABLE: - raise ImportError("lycoris library not installed. Install with: pip install lycoris") + raise ImportError("lycoris library not installed. Install with: pip install lycoris-lora") 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) @@ -104,10 +100,8 @@ def lycoris_extract_from_files( 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: @@ -117,7 +111,6 @@ def lycoris_extract_from_files( pbar = comfy.utils.ProgressBar(len(all_weight_keys)) - # Stats skipped_by_pattern = 0 skipped_by_mismatch = 0 skipped_1d_weights = 0 @@ -126,65 +119,54 @@ def lycoris_extract_from_files( 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_key = key[:-7] 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')") + raise ValueError(f"Key '{key}' not found in model B") 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}") + raise ValueError(f"Shape mismatch for '{key}'") 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, @@ -194,7 +176,6 @@ def lycoris_extract_from_files( is_cp=use_cp and not is_linear_conv, ) else: - # Use linear extraction result, decompose_mode = extract_linear( weight_diff, mode=mode, @@ -203,17 +184,14 @@ def lycoris_extract_from_files( ) 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 + # CP decomposition for 3x3 convs + if is_conv and not is_linear_conv and use_cp: 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), @@ -231,6 +209,18 @@ def lycoris_extract_from_files( 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) + # Sparse bias storage + if use_sparse_bias and diff is not None: + diff_flat = diff.detach().cpu().reshape(extract_b.size(0), -1) + abs_diff = torch.abs(diff_flat) + threshold = float(torch.quantile(abs_diff, sparsity)) + sparse_diff = diff_flat.masked_fill(abs_diff < threshold, 0).to_sparse().coalesce() + + if sparse_diff._nnz() > 0: + output_sd[f"{lora_name}.bias_indices"] = sparse_diff.indices().to(torch.int16) + output_sd[f"{lora_name}.bias_values"] = sparse_diff.values().to(save_torch_dtype) + output_sd[f"{lora_name}.bias_size"] = torch.tensor(diff_flat.shape).to(torch.int16) + low_rank_weights += 1 del extract_a, extract_b, diff @@ -240,7 +230,6 @@ def lycoris_extract_from_files( 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)") @@ -258,117 +247,315 @@ def lycoris_extract_from_files( return output_sd -class LyCORISExtractFromFile(io.ComfyNode): - """ - Memory-efficient LyCORIS extraction from safetensors files. +def _save_lycoris(output_sd: dict, output_filename: str) -> str: + """Save LyCORIS state dict to loras folder.""" + if not output_sd: + raise ValueError("No LyCORIS weights extracted - models may be identical or incompatible") - Uses the LyCORIS library for advanced extraction modes while streaming tensors - one at a time for memory efficiency. - """ + 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 output_path + + +# Common inputs shared by all extraction nodes +def _get_common_inputs(): + return [ + 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)" + ), + ] + + +def _get_common_optional_inputs(): + return [ + io.Boolean.Input("use_cp", default=True, + tooltip="Use CP decomposition for 3x3 convolutions (reduces file size)"), + io.Boolean.Input("use_sparse_bias", default=False, + tooltip="Enable sparse bias storage for residual error compensation"), + io.Float.Input("sparsity", default=0.98, min=0.0, max=1.0, step=0.01, + tooltip="Sparsity threshold for sparse bias (higher = more sparse, smaller file)"), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", + tooltip="How to handle missing keys or shape mismatches"), + 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 (whitespace-separated)"), + ] + + +class LyCORISExtractFixed(io.ComfyNode): + """Extract LyCORIS with fixed rank for linear and conv layers.""" @classmethod def define_schema(cls): return io.Schema( - node_id="LyCORISExtractFromFile", - display_name="Extract LyCORIS from Files (Memory Efficient)", + node_id="LyCORISExtractFixed", + display_name="LyCORIS Extract (Fixed Rank)", 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).", + description="Extracts LyCORIS with fixed LoRA rank. " + "Specify separate ranks for linear and conv layers.", 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"), + *_get_common_inputs(), 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"), + tooltip="LoRA rank for linear/attention layers"), 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"), + tooltip="LoRA rank for conv layers"), + *_get_common_optional_inputs(), ], + 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: + def execute(cls, model_a, model_b, linear_dim, conv_dim, use_cp, use_sparse_bias, + sparsity, mismatch_mode, output_filename, save_dtype, process_device, + skip_patterns) -> io.NodeOutput: if not LYCORIS_AVAILABLE: - raise ImportError("lycoris library not installed. Install with: pip install lycoris") + raise ImportError("lycoris library not installed. Install with: pip install lycoris-lora") - # 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, + mode="fixed", + linear_param=linear_dim, + conv_param=conv_dim, 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, + use_sparse_bias=use_sparse_bias, + sparsity=sparsity, ) - if not output_sd: - raise ValueError("No LyCORIS weights extracted - models may be identical or incompatible") + return io.NodeOutput(_save_lycoris(output_sd, output_filename)) + + +class LyCORISExtractThreshold(io.ComfyNode): + """Extract LyCORIS with singular value threshold.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LyCORISExtractThreshold", + display_name="LyCORIS Extract (Threshold)", + category="ModelUtils/LoRA", + description="Extracts LyCORIS using singular value threshold. " + "Ranks are determined dynamically based on threshold.", + inputs=[ + *_get_common_inputs(), + io.Float.Input("linear_threshold", default=0.01, min=0.0, max=100.0, step=0.001, + tooltip="Singular value threshold for linear layers (absolute value)"), + io.Float.Input("conv_threshold", default=0.01, min=0.0, max=100.0, step=0.001, + tooltip="Singular value threshold for conv layers (absolute value)"), + *_get_common_optional_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_threshold, conv_threshold, use_cp, + use_sparse_bias, sparsity, mismatch_mode, output_filename, save_dtype, + process_device, skip_patterns) -> io.NodeOutput: + if not LYCORIS_AVAILABLE: + raise ImportError("lycoris library not installed. Install with: pip install lycoris-lora") - # 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") + 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) - save_file(output_sd, output_path) - print(f"[LyCORIS Extract] Saved {len(output_sd)} tensors to {output_path}") + output_sd = lycoris_extract_from_files( + model_a_path=model_a_path, + model_b_path=model_b_path, + mode="threshold", + linear_param=linear_threshold, + conv_param=conv_threshold, + 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, + use_sparse_bias=use_sparse_bias, + sparsity=sparsity, + ) - return io.NodeOutput(output_path) + return io.NodeOutput(_save_lycoris(output_sd, output_filename)) + + +class LyCORISExtractRatio(io.ComfyNode): + """Extract LyCORIS with singular value ratio.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LyCORISExtractRatio", + display_name="LyCORIS Extract (Ratio)", + category="ModelUtils/LoRA", + description="Extracts LyCORIS using singular value ratio. " + "Includes singular values >= ratio * max_singular_value.", + inputs=[ + *_get_common_inputs(), + io.Float.Input("linear_ratio", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Ratio of max singular value for linear layers (0-1)"), + io.Float.Input("conv_ratio", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Ratio of max singular value for conv layers (0-1)"), + *_get_common_optional_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, use_cp, use_sparse_bias, + sparsity, mismatch_mode, output_filename, save_dtype, process_device, + skip_patterns) -> io.NodeOutput: + if not LYCORIS_AVAILABLE: + raise ImportError("lycoris library not installed. Install with: pip install lycoris-lora") + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + + output_sd = lycoris_extract_from_files( + model_a_path=model_a_path, + model_b_path=model_b_path, + mode="ratio", + linear_param=linear_ratio, + conv_param=conv_ratio, + 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, + use_sparse_bias=use_sparse_bias, + sparsity=sparsity, + ) + + return io.NodeOutput(_save_lycoris(output_sd, output_filename)) + + +class LyCORISExtractQuantile(io.ComfyNode): + """Extract LyCORIS with cumulative singular value quantile.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LyCORISExtractQuantile", + display_name="LyCORIS Extract (Quantile)", + category="ModelUtils/LoRA", + description="Extracts LyCORIS using cumulative singular value quantile. " + "Includes enough singular values to reach the specified percentage of total.", + inputs=[ + *_get_common_inputs(), + io.Float.Input("linear_quantile", default=0.75, min=0.0, max=1.0, step=0.01, + tooltip="Cumulative quantile for linear layers (0-1, e.g. 0.75 = 75% of variance)"), + io.Float.Input("conv_quantile", default=0.75, min=0.0, max=1.0, step=0.01, + tooltip="Cumulative quantile for conv layers (0-1, e.g. 0.75 = 75% of variance)"), + *_get_common_optional_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, use_cp, + use_sparse_bias, sparsity, mismatch_mode, output_filename, save_dtype, + process_device, skip_patterns) -> io.NodeOutput: + if not LYCORIS_AVAILABLE: + raise ImportError("lycoris library not installed. Install with: pip install lycoris-lora") + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + + output_sd = lycoris_extract_from_files( + model_a_path=model_a_path, + model_b_path=model_b_path, + mode="quantile", + linear_param=linear_quantile, + conv_param=conv_quantile, + 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, + use_sparse_bias=use_sparse_bias, + sparsity=sparsity, + ) + + return io.NodeOutput(_save_lycoris(output_sd, output_filename)) + + +class LyCORISExtractFull(io.ComfyNode): + """Extract full weight differences (no decomposition).""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LyCORISExtractFull", + display_name="LyCORIS Extract (Full Diff)", + category="ModelUtils/LoRA", + description="Extracts full weight differences without SVD decomposition. " + "Larger file size but perfect reconstruction.", + inputs=[ + *_get_common_inputs(), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", + tooltip="How to handle missing keys or shape mismatches"), + io.String.Input("output_filename", default="extracted_full_diff", + 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 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, model_b, mismatch_mode, output_filename, save_dtype, + process_device, skip_patterns) -> io.NodeOutput: + if not LYCORIS_AVAILABLE: + raise ImportError("lycoris library not installed. Install with: pip install lycoris-lora") + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + + output_sd = lycoris_extract_from_files( + model_a_path=model_a_path, + model_b_path=model_b_path, + mode="full", + linear_param=0, + conv_param=0, + prefix_lora="lora_unet_", + process_device=process_device, + save_dtype=save_dtype, + skip_patterns_str=skip_patterns, + mismatch_mode=mismatch_mode, + use_cp=False, + use_sparse_bias=False, + sparsity=0.0, + ) + + return io.NodeOutput(_save_lycoris(output_sd, output_filename)) diff --git a/requirements.txt b/requirements.txt index aea0844..ced22fc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1 @@ -lycoris +lycoris-lora