nodes/mcDandy-more_math
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

AI: make more things lazy

Browse Source
This commit is contained in:
mcDandy
2026-08-27 18:25:03 +02:00
parent 75b9857672
commit 2693d11275
4 changed files with 137 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
more_math/GuiderMathNode.py
+13 -54
View File
@@ -8,7 +8,8 @@ from .helper_functions import (
make_zero_like,
as_tensor,
get_v_variable,
get_f_variable
get_f_variable,
LazyVariableDict,
)
from comfy_api.latest import io
import comfy.sampler_helpers
@@ -21,38 +22,6 @@ from .ParseTree import MrmthParseTree
import copy
class LazyVariableDict(dict):
def __getitem__(self, key):
val = super().__getitem__(key)
if callable(val) and getattr(val, "is_lazy_var", False):
val = val()
self[key] = val
return val
def get(self, key, default=None):
try:
return self[key]
except KeyError:
return default
def copy(self):
return LazyVariableDict(self)
def __or__(self, other):
res = self.copy()
res.update(other)
return res
def __ror__(self, other):
res = LazyVariableDict(other)
res.update(self)
return res
def __ior__(self, other):
self.update(other)
return self
class GuiderMathNode(io.ComfyNode):
"""
Enables math expressions on Guiders (sampler inputs) with autogrow support.
@@ -180,6 +149,10 @@ class MathGuider:
frame_count = eval_samples.shape[time_dim] if time_dim is not None else 1
def make_lazy_var(value_fn):
value_fn.is_lazy_var = True
return value_fn
variables = LazyVariableDict()
variables.update({
"w": self.F.get("F0", 0.0),
@@ -206,37 +179,23 @@ class MathGuider:
for k, func in g_results.items():
variables[k] = func
def lazy_a(): return variables["V0"] if "V0" in g_results else make_zero_like(eval_samples)
lazy_a.is_lazy_var = True
variables["a"] = lazy_a
def lazy_b(): return variables["V1"] if "V1" in g_results else make_zero_like(eval_samples)
lazy_b.is_lazy_var = True
variables["b"] = lazy_b
def lazy_c(): return variables["V2"] if "V2" in g_results else make_zero_like(eval_samples)
lazy_c.is_lazy_var = True
variables["c"] = lazy_c
def lazy_d(): return variables["V3"] if "V3" in g_results else make_zero_like(eval_samples)
lazy_d.is_lazy_var = True
variables["d"] = lazy_d
variables["a"] = make_lazy_var(lambda: variables["V0"] if "V0" in g_results else make_zero_like(eval_samples))
variables["b"] = make_lazy_var(lambda: variables["V1"] if "V1" in g_results else make_zero_like(eval_samples))
variables["c"] = make_lazy_var(lambda: variables["V2"] if "V2" in g_results else make_zero_like(eval_samples))
variables["d"] = make_lazy_var(lambda: variables["V3"] if "V3" in g_results else make_zero_like(eval_samples))
def lazy_stacked_v():
evaluated_g_results = {k: variables[k] for k in g_results.keys()}
v_stacked, v_cnt = get_v_variable(evaluated_g_results)
return v_stacked
lazy_stacked_v.is_lazy_var = True
variables["V"] = make_lazy_var(lazy_stacked_v)
def lazy_v_cnt():
evaluated_g_results = {k: variables[k] for k in g_results.keys()}
v_stacked, v_cnt = get_v_variable(evaluated_g_results)
return float(v_cnt)
lazy_v_cnt.is_lazy_var = True
variables["V"] = lazy_stacked_v
variables["Vcnt"] = lazy_v_cnt
variables["V_count"] = lazy_v_cnt
variables["Vcnt"] = make_lazy_var(lazy_v_cnt)
variables["V_count"] = variables["Vcnt"]
for k, v in self.F.items():
variables[k] = v if v is not None else 0.0
more_math/VideoMathNode.py
+31 -19
View File
@@ -1,4 +1,4 @@
from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable, checkLazyNew
from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable, checkLazyNew, LazyVariableDict
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
from comfy_api.latest import io
from .Stack import MrmthStack
@@ -86,7 +86,7 @@ class VideoMathNode(io.ComfyNode):
if tensor is not None and tensor.shape[0] != common_shape[0]:
raise ValueError(f"Input '{name}' has shape {tensor.shape[0]}, expected {common_shape[0]} to match input.")
variables = {
variables = LazyVariableDict({
"a": ae, "b": be, "c": ce, "d": de,
"w": F.get("F0", 0.0) if F.get("F0") is not None else 0.0,
"x": F.get("F1", 0.0) if F.get("F1") is not None else 0.0,
@@ -106,13 +106,21 @@ class VideoMathNode(io.ComfyNode):
"batch_count": float(ae.shape[0]),
"N": float(ae.shape[3]),
"channel_count": float(ae.shape[3]),
} | generate_dim_variables(ae)
} | generate_dim_variables(ae) | V_norm | sample_rates)
v_stacked, v_cnt = get_v_variable(V_norm, length_mismatch=length_mismatch)
if v_stacked is not None:
variables["V"] = v_stacked
variables["Vcnt"] = float(v_cnt)
variables["V_count"] = float(v_cnt)
for v_name, v_tensor in V_norm.items():
variables[v_name] = (lambda t=v_tensor: t)
variables[v_name].is_lazy_var = True
def lazy_v_stack():
return get_v_variable(V_norm, length_mismatch=length_mismatch)[0]
lazy_v_stack.is_lazy_var = True
def lazy_v_count():
return float(get_v_variable(V_norm, length_mismatch=length_mismatch)[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:
@@ -120,9 +128,6 @@ class VideoMathNode(io.ComfyNode):
variables["Fcnt"] = float(f_cnt)
variables["F_count"] = float(f_cnt)
# Add all dynamic inputs
variables.update(V_norm)
for k, val in F.items():
variables[k] = val if val is not None else 0.0
@@ -163,7 +168,7 @@ class VideoMathNode(io.ComfyNode):
# Ensure legacy are normalized
a_w, b_w, c_w, d_w = normalize_to_common_shape(a_w, b_w, c_w, d_w, mode=length_mismatch)
variables = {
variables = LazyVariableDict({
"a": a_w, "b": b_w, "c": c_w, "d": d_w,
"w": F.get("F0", 0.0) if F.get("F0") is not None else 0.0,
"x": F.get("F1", 0.0) if F.get("F1") is not None else 0.0,
@@ -179,14 +184,21 @@ class VideoMathNode(io.ComfyNode):
"batch": getIndexTensorAlongDim(a_w, 0),
"T": float(a_w.shape[0]),
"batch_count": float(a_w.shape[0]),
} | generate_dim_variables(a_w) | V_norm_waveforms | sample_rates
} | generate_dim_variables(a_w) | sample_rates)
v_stacked, v_cnt = get_v_variable(V_norm_waveforms, length_mismatch=length_mismatch)
if v_stacked is not None:
# This 'variables' is the one for the second eval in VideoMathNode
variables["V"] = v_stacked
variables["Vcnt"] = float(v_cnt)
variables["V_count"] = float(v_cnt)
for v_name, v_tensor in V_norm_waveforms.items():
variables[v_name] = (lambda t=v_tensor: t)
variables[v_name].is_lazy_var = True
def lazy_v_stack():
return get_v_variable(V_norm_waveforms, length_mismatch=length_mismatch)[0]
lazy_v_stack.is_lazy_var = True
def lazy_v_count():
return float(get_v_variable(V_norm_waveforms, length_mismatch=length_mismatch)[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:
more_math/helper_functions.py
+33
View File
@@ -8,6 +8,39 @@ MathExprLexer, MathExprParser, _ = get_antlr_modules()
import re
class LazyVariableDict(dict):
def __getitem__(self, key):
val = super().__getitem__(key)
if callable(val) and getattr(val, "is_lazy_var", False):
val = val()
self[key] = val
return val
def get(self, key, default=None):
try:
return self[key]
except KeyError:
return default
def copy(self):
return LazyVariableDict(self)
def __or__(self, other):
res = self.copy()
res.update(other)
return res
def __ror__(self, other):
res = LazyVariableDict(other)
res.update(self)
return res
def __ior__(self, other):
self.update(other)
return self
class ThrowingErrorListener(ErrorListener):
"""Error listener that raises ValueError on syntax errors."""
more_math/modelLikeCommon.py
+60 -42
View File
@@ -1,4 +1,4 @@
from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable, get_tensor_device, move_to_device
from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable, get_tensor_device, move_to_device, LazyVariableDict
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
import comfy
import comfy.model_management
@@ -134,10 +134,17 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us
patches = {}
original_device = None
compute_device = comfy.model_management.get_torch_device() if use_compute_device else None
needs_deltas = isinstance(Expr, str) and "_d" in Expr or not isinstance(Expr, str)
needs_deltas = not (isinstance(Expr, str) and "_d" not in Expr)
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())
current_state_dicts = {}
base_state_dicts = {}
base_state_dicts = {} if needs_deltas else None
for v_name, v_val in V.items():
if v_val is None:
continue
@@ -158,17 +165,18 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us
# Progress bar if possible (comfy.utils.ProgressBar might assume unthreaded?)
# Just skip for utility or use if substantial.
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
for layer_idx, key in enumerate(all_keys_list):
variables = {}
# Populate F variables (constants for all keys)
for k, val in F.items():
variables[k] = val if val is not None else 0.0
variables = LazyVariableDict(const_f_vars)
# Also populate mapped aliases for F (w, x, y, z)
for alias, target in mapping.items():
if target in F:
variables[alias] = F[target] if F[target] is not None else 0.0
variables.update(const_f_alias_vars)
variables["L"] = float(layer_idx)
variables["layer"] = float(layer_idx)
@@ -180,7 +188,7 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us
# Inject weights for this key from V models
valid_key = False
ref_tensor = None
deltas = {}
present_weights = {}
for v_name in V.keys():
current_sd = current_state_dicts.get(v_name)
@@ -188,22 +196,10 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us
continue
w_tensor = current_sd.get(key, None)
if w_tensor is not None:
variables[v_name] = w_tensor
ref_tensor = w_tensor
present_weights[v_name] = w_tensor
if ref_tensor is None:
ref_tensor = w_tensor
valid_key = True
if needs_deltas:
base_sd = base_state_dicts.get(v_name)
base_tensor = base_sd.get(key, None) if base_sd is not None else None
if base_tensor is None:
base_tensor = w_tensor
base_tensor = coerce_like(base_tensor, w_tensor, device=compute_device)
try:
deltas[v_name] = w_tensor - base_tensor
except Exception:
deltas[v_name] = torch.zeros_like(w_tensor)
else:
# Missing key in this model will be handled later (zero init)
pass
if not valid_key:
continue
@@ -212,32 +208,55 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us
if ref_tensor is None:
continue # Should not happen if valid_key is true
# Fill missing models with zeros
# Fill missing models with lazy values so V[n] keeps working
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
for v_name in V.keys():
if v_name not in variables:
variables[v_name] = torch.zeros_like(ref_tensor)
if v_name not in deltas:
deltas[v_name] = torch.zeros_like(variables[v_name])
variables[f"{v_name}_d"] = deltas[v_name]
if v_name in present_weights:
w_tensor = present_weights[v_name]
variables[v_name] = make_lazy(lambda wt=w_tensor: wt)
if needs_deltas:
base_sd = base_state_dicts.get(v_name)
base_tensor = base_sd.get(key, None) if base_sd is not None else None
if base_tensor is None:
base_tensor = w_tensor
base_tensor = coerce_like(base_tensor, w_tensor, device=compute_device)
variables[f"{v_name}_d"] = make_lazy(lambda wt=w_tensor, bt=base_tensor: wt - bt)
else:
variables[v_name] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
if needs_deltas:
variables[f"{v_name}_d"] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
# Populate aliases for V (a, b, c, d)
for alias, target in mapping.items():
for alias, target in mapping_items:
if target in variables:
variables[alias] = variables[target]
elif target in V:
# V key exists in input dict but this specific layer key is missing
variables[alias] = torch.zeros_like(ref_tensor)
variables[alias] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
else:
# Target doesn't exist at all (e.g., V1 not provided)
# Check if it's a V-key pattern and zero-fill
if target.startswith("V") and target[1:].isdigit():
variables[alias] = torch.zeros_like(ref_tensor)
variables[alias] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
v_stacked, v_cnt = get_v_variable(variables)
if v_stacked is not None:
variables["V"] = v_stacked
variables["Vcnt"] = float(v_cnt)
variables["V_count"] = float(v_cnt)
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:
@@ -246,7 +265,6 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us
variables["F_count"] = float(f_cnt)
variables = variables | generate_dim_variables(ref_tensor)
variables |= {f"{v_name}_d": variables[v_name] - variables["V0"] for v_name in V.keys() if v_name in variables and "V0" in variables}
# Execute math