From 20397b69c2eca4992f0397fe6e223adaffe62b8c Mon Sep 17 00:00:00 2001 From: mcDandy Date: Fri, 28 Aug 2026 21:53:43 +0200 Subject: [PATCH] AI: quick lora math node. TODO: make human walk through --- more_math/LoraMathNode.py | 63 +++++++++++++++ more_math/loraDictCommon.py | 155 ++++++++++++++++++++++++++++++++++++ more_math/nodes.py | 2 + 3 files changed, 220 insertions(+) create mode 100644 more_math/LoraMathNode.py create mode 100644 more_math/loraDictCommon.py diff --git a/more_math/LoraMathNode.py b/more_math/LoraMathNode.py new file mode 100644 index 0000000..5bcde3e --- /dev/null +++ b/more_math/LoraMathNode.py @@ -0,0 +1,63 @@ +from inspect import cleandoc +from comfy_api.latest import io +import copy +from .helper_functions import checkLazyNew +from .Stack import MrmthStack +from .ParseTree import MrmthParseTree +from .loraDictCommon import calculate_lora_dict_autogrow + + +class LoraMathNode(io.ComfyNode): + """ + This node enables the use of math expressions directly on raw LoRA + tensors (LORA_MODEL: dict of lora_up.weight / lora_down.weight / alpha / + diff tensors) using Autogrow inputs. + + Unlike Model/VAE/CLIP Math, these tensors are not attached to a + ModelPatcher, so there is no base-vs-patched distinction - the tensors + already *are* the LoRA's deltas - and therefore no V{n}_d. + """ + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="mrmth_ag_LoraMathNode", + display_name="Lora Math", + category="More math", + inputs=[ + io.Autogrow.Input(id="V", template=io.Autogrow.TemplatePrefix(io.LoraModel.Input("values"), prefix="V", min=1, max=50)), + io.Autogrow.Input(id="F", template=io.Autogrow.TemplatePrefix(io.Float.Input("float", default=0.0, optional=True, lazy=True, force_input=True), prefix="F", min=1, max=50)), + io.MultiType.Input( + io.String.Input("Expression", default="V0", multiline=False), + types=[io.String, MrmthParseTree], + tooltip="Expression to apply on LoRA tensors", + ), + io.Boolean.Input( + id="use_compute_device", + default=True, + display_name="Move tensors to GPU", + tooltip="Temporarily copies LoRA tensors to the compute device for math and moves the result back afterwards.", + ), + MrmthStack.Input(id="stack", tooltip="Access stack between nodes", optional=True) + ], + outputs=[ + io.LoraModel.Output(), + MrmthStack.Output(), + ], + ) + + tooltip = cleandoc(__doc__) + + @classmethod + def check_lazy_status(cls, Expression, V, F, use_compute_device=True, stack={}): + return checkLazyNew(Expression, V, F) + + @classmethod + def execute(cls, V, F, Expression, use_compute_device=True, stack={}) -> io.NodeOutput: + if not any(v is not None for v in V.values()): + raise ValueError("At least one input LoRA is required.") + stack = copy.deepcopy(stack) if stack is not None else {} + + aliases = {"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"} + result = calculate_lora_dict_autogrow(Expression, V=V, F=F, mapping=aliases, stack=stack, use_compute_device=use_compute_device) + return (result, stack) diff --git a/more_math/loraDictCommon.py b/more_math/loraDictCommon.py new file mode 100644 index 0000000..647c8f3 --- /dev/null +++ b/more_math/loraDictCommon.py @@ -0,0 +1,155 @@ +from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable, get_tensor_device, move_to_device, LazyVariableDict, checkLazyNew +from .Parser.UnifiedMathVisitor import UnifiedMathVisitor +import comfy.model_management +import torch + + +def calculate_lora_dict_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], use_compute_device=False): + """ + Calculate a new LORA_MODEL dict[str, Tensor] from math expressions applied + over the UNION of keys from all input LoRA dicts, using Autogrow inputs. + + Unlike calculate_patches_autogrow (Model/VAE/CLIP), there is no + ModelPatcher/backup involved here - a LORA_MODEL is just raw tensors + (lora_up.weight, lora_down.weight, alpha, diff, ...). So V{n} is simply + that tensor as-is (zero-filled when a given key is missing from that + input), and there is no V{n}_d - these tensors already *are* deltas. + """ + if V is None: V = {} + if F is None: F = {} + if mapping is None: mapping = {} + + needed_vars = checkLazyNew(Expr, V, F) + needed_v_names = [v_name for v_name in V.keys() if v_name in needed_vars] + if not needed_v_names: + needed_v_names = [v_name for v_name in V.keys() if V.get(v_name) is not None] + + if not any(V.get(v_name) is not None for v_name in needed_v_names): + return {} + + # Collect all unique keys from all inputs, preserving order from first input + all_keys_list = [] + seen_keys = set() + for d in V.values(): + if d is None: + continue + for key in d.keys(): + if key not in seen_keys: + seen_keys.add(key) + all_keys_list.append(key) + + tree = parse_expr(Expr) if isinstance(Expr, str) else Expr + + original_device = None + compute_device = comfy.model_management.get_torch_device() if use_compute_device else None + + const_f_vars = {k: val if val is not None else 0.0 for k, val in F.items()} + const_f_alias_vars = { + alias: (F[target] if F[target] is not None else 0.0) + for alias, target in mapping.items() + if target in F + } + mapping_items = tuple(mapping.items()) + + input_dicts = {} + for v_name in needed_v_names: + d = V.get(v_name) + if d is None: + continue + if original_device is None: + original_device = get_tensor_device(d) + if compute_device is not None and original_device is not None and compute_device != original_device: + d = move_to_device(d, compute_device) + input_dicts[v_name] = d + + layer_count = len(all_keys_list) + + def make_lazy(value_fn): + def lazy_eval(): + return value_fn() + lazy_eval.is_lazy_var = True + return lazy_eval + + result = {} + + for layer_idx, key in enumerate(all_keys_list): + variables = LazyVariableDict(const_f_vars) + variables.update(const_f_alias_vars) + + variables["L"] = float(layer_idx) + variables["layer"] = float(layer_idx) + variables["LC"] = float(layer_count) + variables["layer_count"] = float(layer_count) + variables["K"] = key + variables["key"] = key + + ref_tensor = None + present = {} + for v_name in V.keys(): + d = input_dicts.get(v_name) + if d is None: + continue + t = d.get(key, None) + if t is not None: + present[v_name] = t + if ref_tensor is None: + ref_tensor = t + + if ref_tensor is None: + continue + + for v_name in V.keys(): + if v_name in present: + t = present[v_name] + variables[v_name] = make_lazy(lambda t=t: t) + else: + variables[v_name] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt)) + + # Populate aliases for V (a, b, c, d) + for alias, target in mapping_items: + if target in variables: + variables[alias] = variables[target] + elif target.startswith("V") and target[1:].isdigit(): + variables[alias] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt)) + + v_cache = [None] + + def lazy_v_stack(): + if v_cache[0] is None: + v_cache[0] = get_v_variable(variables) + return v_cache[0][0] + + lazy_v_stack.is_lazy_var = True + + def lazy_v_count(): + if v_cache[0] is None: + v_cache[0] = get_v_variable(variables) + return float(v_cache[0][1]) + + lazy_v_count.is_lazy_var = True + + variables["V"] = lazy_v_stack + variables["Vcnt"] = lazy_v_count + variables["V_count"] = lazy_v_count + + f_stacked, f_cnt = get_f_variable(F) + if f_stacked is not None: + variables["F"] = f_stacked + variables["Fcnt"] = float(f_cnt) + variables["F_count"] = float(f_cnt) + + variables = variables | generate_dim_variables(ref_tensor) + + visitor = UnifiedMathVisitor(variables, ref_tensor.shape, state_storage=stack) + res = visitor.visit(tree) + res = as_tensor(res, ref_tensor.shape) + + if compute_device is not None and original_device is not None and original_device != compute_device: + res = res.to(device=original_device) + + result[key] = res + + if pbar is not None: + pbar.update(1) + + return result diff --git a/more_math/nodes.py b/more_math/nodes.py index 1a8ed80..92a15e0 100644 --- a/more_math/nodes.py +++ b/more_math/nodes.py @@ -21,6 +21,7 @@ from .VideoMathNode import VideoMathNode from .ModelMathNode import ModelMathNode from .VaeMathNode import VAEMathNode from .ClipMathNode import CLIPMathNode +from .LoraMathNode import LoraMathNode from .SpectrogramToAudioNode import SpectrogramToAudio from .AudioToSpectrogramNode import AudioToSpectrogram @@ -151,6 +152,7 @@ class MoreMathExtension(ComfyExtension): CLIPMathNodeOLD, VAEMathNode, VAEMathNodeOLD, + LoraMathNode, LatentMathNode, LatentMathNodeOLD, ImageMathNode,