diff --git a/more_math/GuiderMathNode.py b/more_math/GuiderMathNode.py index c7ff33d..ee5213d 100644 --- a/more_math/GuiderMathNode.py +++ b/more_math/GuiderMathNode.py @@ -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 diff --git a/more_math/VideoMathNode.py b/more_math/VideoMathNode.py index ca97a32..3db2c27 100644 --- a/more_math/VideoMathNode.py +++ b/more_math/VideoMathNode.py @@ -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: diff --git a/more_math/helper_functions.py b/more_math/helper_functions.py index 77636bd..2c67ba6 100644 --- a/more_math/helper_functions.py +++ b/more_math/helper_functions.py @@ -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.""" diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 64411e0..1c65cdf 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -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