From d0bbcb51fc211292310794da144e60500a732248 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Sun, 18 Oct 2026 00:58:12 +0200 Subject: [PATCH] AI: audit 1 --- more_math/AudioMathNode.py | 6 ++-- more_math/ClipMathNode.py | 4 +-- more_math/ConditioningMathNode.py | 6 ++-- more_math/FloatMathNode.py | 6 ++-- more_math/GuiderMathNode.py | 4 +-- more_math/ImageMathNode.py | 6 ++-- more_math/LatentMathNode.py | 6 ++-- more_math/LoraMathNode.py | 4 +-- more_math/MaskMathNode.py | 6 ++-- more_math/ModelMathNode.py | 4 +-- more_math/NoiseMathNode.py | 6 ++-- more_math/Parser/UnifiedMathVisitor.py | 12 ++++--- more_math/SelectiveGuiderMathNode.py | 4 +-- more_math/SigmasMathNode.py | 4 +-- more_math/StringMathNode.py | 5 +-- more_math/VaeMathNode.py | 4 +-- more_math/VideoMathNode.py | 4 +-- more_math/helper_functions.py | 37 ++++++--------------- more_math/loraDictCommon.py | 4 +-- more_math/modelLikeCommon.py | 23 ++++--------- tests/test_assignment.py | 39 ++++++++++++----------- tests/test_unified_visitor_regressions.py | 31 +++++++++++++++++- 22 files changed, 116 insertions(+), 109 deletions(-) diff --git a/more_math/AudioMathNode.py b/more_math/AudioMathNode.py index e2846b3..beea704 100644 --- a/more_math/AudioMathNode.py +++ b/more_math/AudioMathNode.py @@ -74,17 +74,17 @@ class AudioMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0, remember_stack=False,use_compute_device=True,stack={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0, remember_stack=False,use_compute_device=True,stack=None): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",batching=0, remember_stack=False,use_compute_device=True, stack={}): + def execute(cls, V, F, Expression, length_mismatch="tile",batching=0, remember_stack=False,use_compute_device=True, stack=None): # Identify all present audio inputs and their keys tensor_keys = [k for k, v in V.items() if v is not None and isinstance(v, dict) and "waveform" in v] if not tensor_keys: raise ValueError("At least one audio input is required.") - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) original_waveform = V[tensor_keys[0]]["waveform"] original_device = get_tensor_device(original_waveform) compute_device = comfy.model_management.get_torch_device() if use_compute_device else original_device diff --git a/more_math/ClipMathNode.py b/more_math/ClipMathNode.py index 445455a..77cdfab 100644 --- a/more_math/ClipMathNode.py +++ b/more_math/ClipMathNode.py @@ -58,12 +58,12 @@ class CLIPMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack=None,dict_mode=False): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False) -> io.NodeOutput: + def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack=None,dict_mode=False) -> io.NodeOutput: # Determine reference CLIP a = V.get("V0") stack = copy.deepcopy(stack) if stack is not None else {} diff --git a/more_math/ConditioningMathNode.py b/more_math/ConditioningMathNode.py index b4b10da..7cd2480 100644 --- a/more_math/ConditioningMathNode.py +++ b/more_math/ConditioningMathNode.py @@ -67,19 +67,19 @@ class ConditioningMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression,Expression_pi, V, F, length_mismatch="tile", batching=0,remember_stack=False,use_compute_device=True, stack={}): + def check_lazy_status(cls, Expression,Expression_pi, V, F, length_mismatch="tile", batching=0,remember_stack=False,use_compute_device=True, stack=None): d = checkLazyNew(Expression,V,F) b = checkLazyNew(Expression_pi,V,F) return d|b @classmethod - def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile", batching=0,remember_stack=False,use_compute_device=True, stack={}): + def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile", batching=0,remember_stack=False,use_compute_device=True, stack=None): # Identify all present conditioning inputs tensor_keys = [k for k, v in V.items() if v is not None and isinstance(v, list) and len(v) > 0] if not tensor_keys: raise ValueError("At least one input is required.") - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) needed_vars = checkLazyNew(Expression, V, F) | checkLazyNew(Expression_pi, V, F) needed_tensor_keys = [k for k in tensor_keys if k in needed_vars] diff --git a/more_math/FloatMathNode.py b/more_math/FloatMathNode.py index 1ccaad9..3d7b670 100644 --- a/more_math/FloatMathNode.py +++ b/more_math/FloatMathNode.py @@ -58,13 +58,13 @@ class FloatMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, FloatFunc, V, remember_stack=False, stack={}): + def check_lazy_status(cls, FloatFunc, V, remember_stack=False, stack=None): # remember_stack ani stack nemění lazy logiku return checkLazyNew(FloatFunc, V, V) @classmethod - def execute(cls, FloatFunc, V, remember_stack=False, stack={}): - work_stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + def execute(cls, FloatFunc, V, remember_stack=False, stack=None): + work_stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) variables = {} # Populate aliases diff --git a/more_math/GuiderMathNode.py b/more_math/GuiderMathNode.py index 1ba8a6b..70e9d63 100644 --- a/more_math/GuiderMathNode.py +++ b/more_math/GuiderMathNode.py @@ -63,7 +63,7 @@ class GuiderMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression,Expression1, V, F,remember_stack=True,stack={}): + def check_lazy_status(cls, Expression,Expression1, V, F,remember_stack=True,stack=None): d = checkLazyNew(Expression,V,F) b = checkLazyNew(Expression1,V,F) return d|b @@ -71,7 +71,7 @@ class GuiderMathNode(io.ComfyNode): @classmethod def execute(cls, V, F, Expression, Expression1,remember_stack = True, stack=None): - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) guider = MathGuider(V, F, Expression, Expression1, stack); stack = stack if remember_stack else copy.deepcopy(stack) return (guider, stack) diff --git a/more_math/ImageMathNode.py b/more_math/ImageMathNode.py index 136f0a1..3820b0f 100644 --- a/more_math/ImageMathNode.py +++ b/more_math/ImageMathNode.py @@ -63,11 +63,11 @@ class ImageMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0,remember_stack=False,use_compute_device=True,stack={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0,remember_stack=False,use_compute_device=True,stack=None): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="error",batching=0,remember_stack=False,use_compute_device=True,stack={}): + def execute(cls, V, F, Expression, length_mismatch="error",batching=0,remember_stack=False,use_compute_device=True,stack=None): # I and F are Autogrow.Type which is dict[str, Any] # Identify all present tensors and their keys @@ -80,7 +80,7 @@ class ImageMathNode(io.ComfyNode): compute_device = comfy.model_management.get_torch_device() if use_compute_device else original_device working_V = move_to_device(V, compute_device) if compute_device is not None and compute_device != original_device else V tensors = [working_V[k] for k in tensor_keys] - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) # Normalize all tensors together to find the common target shape normalized_tensors = normalize_to_common_shape(*tensors, mode=length_mismatch) diff --git a/more_math/LatentMathNode.py b/more_math/LatentMathNode.py index 265bafc..8e42e70 100644 --- a/more_math/LatentMathNode.py +++ b/more_math/LatentMathNode.py @@ -78,11 +78,11 @@ class LatentMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, Expression, V, F,batching, length_mismatch="tile",remember_stack=False,use_compute_device=True,stack={}): + def check_lazy_status(cls, Expression, V, F,batching, length_mismatch="tile",remember_stack=False,use_compute_device=True,stack=None): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression,batching, length_mismatch="tile",remember_stack=False,use_compute_device=True,stack={}) -> io.NodeOutput: + def execute(cls, V, F, Expression,batching, length_mismatch="tile",remember_stack=False,use_compute_device=True,stack=None) -> io.NodeOutput: # Determine reference latent ref_latent = None for lat in V.values(): @@ -92,7 +92,7 @@ class LatentMathNode(io.ComfyNode): if ref_latent is None: raise ValueError("At least one input is required.") - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) original_device = get_tensor_device(ref_latent.get("samples")) compute_device = comfy.model_management.get_torch_device() if use_compute_device else original_device working_V = move_to_device(V, compute_device) if compute_device is not None and compute_device != original_device else V diff --git a/more_math/LoraMathNode.py b/more_math/LoraMathNode.py index 7fcc18e..90ae94c 100644 --- a/more_math/LoraMathNode.py +++ b/more_math/LoraMathNode.py @@ -56,11 +56,11 @@ class LoraMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, Expression, V, F, use_compute_device=True, stack={}, dict_mode=False): + def check_lazy_status(cls, Expression, V, F, use_compute_device=True, stack=None, dict_mode=False): return checkLazyNew(Expression, V, F) @classmethod - def execute(cls, V, F, Expression, use_compute_device=True, stack={}, dict_mode=False) -> io.NodeOutput: + def execute(cls, V, F, Expression, use_compute_device=True, stack=None, dict_mode=False) -> 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 {} diff --git a/more_math/MaskMathNode.py b/more_math/MaskMathNode.py index 2bd718b..f6414db 100644 --- a/more_math/MaskMathNode.py +++ b/more_math/MaskMathNode.py @@ -56,18 +56,18 @@ class MaskMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0,remember_stack=False,stack={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0,remember_stack=False,stack=None): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",batching=0,remember_stack=False,stack={}): + def execute(cls, V, F, Expression, length_mismatch="tile",batching=0,remember_stack=False,stack=None): # Identify all present tensors and their keys tensor_keys = [k for k, v in V.items() if v is not None] if not tensor_keys: raise ValueError("At least one input is required.") tensors = [V[k] for k in tensor_keys] - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) # Normalize all tensors together normalized_tensors = normalize_to_common_shape(*tensors, mode=length_mismatch) V_norm = dict(zip(tensor_keys, normalized_tensors)) diff --git a/more_math/ModelMathNode.py b/more_math/ModelMathNode.py index 3f3bf7d..068ba08 100644 --- a/more_math/ModelMathNode.py +++ b/more_math/ModelMathNode.py @@ -59,11 +59,11 @@ class ModelMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack=None,dict_mode=False): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False) -> io.NodeOutput: + def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack=None,dict_mode=False) -> io.NodeOutput: # Determine reference model for cloning a = V.get("V0") stack = copy.deepcopy(stack) if stack is not None else {} diff --git a/more_math/NoiseMathNode.py b/more_math/NoiseMathNode.py index c1ec1ec..3d7d94a 100644 --- a/more_math/NoiseMathNode.py +++ b/more_math/NoiseMathNode.py @@ -53,12 +53,12 @@ class NoiseMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Noise, V, F,remember_stack=False,stack={}): + def check_lazy_status(cls, Noise, V, F,remember_stack=False,stack=None): return checkLazyNew(Noise,V,F) @classmethod - def execute(cls, Noise, V,F,remember_stack=False,stack={}): - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + def execute(cls, Noise, V,F,remember_stack=False,stack=None): + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) executer = NoiseExecutor(V,F, Noise,stack) stack = stack if remember_stack else copy.deepcopy(stack) return (executer,stack) diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index ff5bff3..6f39949 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -3,6 +3,7 @@ import os import torch import math import inspect +import weakref import torch.nn.functional as F from comfy.nested_tensor import NestedTensor from comfy.model_management import throw_exception_if_processing_interrupted @@ -142,6 +143,7 @@ class UnifiedMathVisitor(MathExprVisitor): self.depth = depth self._scope_stack = [] self._state_storage = state_storage if state_storage is not None else {} + self._owned_tensors = weakref.WeakValueDictionary() def visit(self, tree): if tree is None: @@ -3062,6 +3064,12 @@ class UnifiedMathVisitor(MathExprVisitor): target = self.variables[var_name] if self._is_tensor(target): + # Inputs can be cached outputs of upstream nodes, so copy before the first in-place write. + if self._is_plain_tensor(target) and self._owned_tensors.get(id(target)) is not target: + target = target.clone() + self._owned_tensors[id(target)] = target + self.variables[var_name] = target + # Process indices for PyTorch torch_indices = [] for idx in indices: @@ -3180,10 +3188,6 @@ class UnifiedMathVisitor(MathExprVisitor): throw_exception_if_processing_interrupted() func_name = ctx.VARIABLE().getText() - builtin_visit = f"visit{func_name[:1].upper()}{func_name[1:]}Func" - if hasattr(self, builtin_visit): - return (yield getattr(self, builtin_visit)(ctx)) - # 1) Pokud proměnná existuje a je to lambda uložená v variables, aplikuj ji if func_name in self.variables and isinstance(self.variables[func_name], LambdaFunction): lam = self.variables[func_name] diff --git a/more_math/SelectiveGuiderMathNode.py b/more_math/SelectiveGuiderMathNode.py index 1b1369a..a5f6340 100644 --- a/more_math/SelectiveGuiderMathNode.py +++ b/more_math/SelectiveGuiderMathNode.py @@ -57,7 +57,7 @@ class SelectiveGuiderMathNode(io.ComfyNode): hook_target="all", remember_stack=False, debug_hooks=False, - stack={}, + stack=None, ): return checkLazyNew(Expression, {}, F) @@ -76,7 +76,7 @@ class SelectiveGuiderMathNode(io.ComfyNode): if V is None or not hasattr(V, "model_options"): raise ValueError("Vstupní guider musí mít model_options.") - stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) tree = parse_expr(Expression) if isinstance(Expression, str) else Expression hit_flags = { diff --git a/more_math/SigmasMathNode.py b/more_math/SigmasMathNode.py index 36abfb1..653d31a 100644 --- a/more_math/SigmasMathNode.py +++ b/more_math/SigmasMathNode.py @@ -45,11 +45,11 @@ class SigmasMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack=None): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",stack={}): + def execute(cls, V, F, Expression, length_mismatch="tile",stack=None): # I and F are Autogrow.Type which is dict[str, Any] # Determine reference image for zero-initialization (fallback for a,b,c,d) diff --git a/more_math/StringMathNode.py b/more_math/StringMathNode.py index 82ab33f..9fed137 100644 --- a/more_math/StringMathNode.py +++ b/more_math/StringMathNode.py @@ -53,12 +53,13 @@ class StringMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression, V, F,batching=0, remember_stack=False,stack={}): + def check_lazy_status(cls, Expression, V, F,batching=0, remember_stack=False,stack=None): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, batching=0, remember_stack=False, stack={}): + def execute(cls, V, F, Expression, batching=0, remember_stack=False, stack=None): + stack = {} if stack is None else (stack if remember_stack else copy.deepcopy(stack)) variables = { "a": V.get("V0",""), "b": V.get("V1",""), "c": V.get("V2",""), "d": V.get("V3",""), diff --git a/more_math/VaeMathNode.py b/more_math/VaeMathNode.py index 1062460..6fd0bb2 100644 --- a/more_math/VaeMathNode.py +++ b/more_math/VaeMathNode.py @@ -55,11 +55,11 @@ class VAEMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack=None,dict_mode=False): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False) -> io.NodeOutput: + def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack=None,dict_mode=False) -> io.NodeOutput: # Determine reference VAE a = V.get("V0") if a is None: diff --git a/more_math/VideoMathNode.py b/more_math/VideoMathNode.py index bce29d7..f912841 100644 --- a/more_math/VideoMathNode.py +++ b/more_math/VideoMathNode.py @@ -96,13 +96,13 @@ class VideoMathNode(io.ComfyNode): ) @classmethod - def check_lazy_status(cls, Expression,Expression_pi, V, F, length_mismatch="tile",use_compute_device=True,stack={}): + def check_lazy_status(cls, Expression,Expression_pi, V, F, length_mismatch="tile",use_compute_device=True,stack=None): d = checkLazyNew(Expression,V,F) b = checkLazyNew(Expression_pi,V,F) return d|b @classmethod - def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile",use_compute_device=True,stack={}): + def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile",use_compute_device=True,stack=None): tensor_keys = [k for k, v in V.items() if v is not None] if not tensor_keys: raise ValueError("At least one input is required.") diff --git a/more_math/helper_functions.py b/more_math/helper_functions.py index 4b5e7c2..ef77264 100644 --- a/more_math/helper_functions.py +++ b/more_math/helper_functions.py @@ -337,31 +337,6 @@ def get_f_variable(f_dict): except Exception: return None, len(ordered_values) -def _collect_reads_only( node, needed_vars, assigned_vars, shadowed_vars): - if node is None: - return - - node_type = type(node).__name__ - - if node_type == 'VariableExpContext': - var_name = node.VARIABLE().getText() - if var_name in shadowed_vars: - return - if var_name not in assigned_vars: - needed_vars.add(var_name) - return - - if node_type == 'FunctionDefContext': - return - - if node_type == 'VarDefContext': - for expr in node.expr(): - _collect_reads_only(expr, needed_vars, assigned_vars, shadowed_vars) - return - - for i in range(node.getChildCount()): - _collect_reads_only(node.getChild(i), needed_vars, assigned_vars, shadowed_vars) - def _collect_vars_from_node( node, needed_vars, assigned_vars, shadowed_vars): """Recursively collect variable reads from an AST node""" if node is None: @@ -379,10 +354,16 @@ def _collect_vars_from_node( node, needed_vars, assigned_vars, shadowed_vars): needed_vars.add(var_name) return - # Skip + # Function bodies are scanned by checkLazyNew with their parameters shadowed if node_type == 'FunctionDefContext': return + # Compound (x += ...) and indexed (x[i] = ...) assignments read their target + if node_type == 'VarDefContext' and (node.LBRACKET() or node.EQUEALS() is None): + var_name = node.VARIABLE().getText() + if var_name not in shadowed_vars and var_name not in assigned_vars: + needed_vars.add(var_name) + # Recursively visit children for i in range(node.getChildCount()): _collect_vars_from_node(node.getChild(i), needed_vars, assigned_vars, shadowed_vars) @@ -411,14 +392,14 @@ def checkLazyNew(Expression, V, F): rule_name = parser.ruleNames[child.getRuleIndex()] if child.getRuleIndex() < len(parser.ruleNames) else None - # Process function definitions: scan for reads but ignore writes + # Process function definitions: scan the body for reads, ignoring parameters if rule_name == 'funcDef': func_params = set() if child.paramList(): for param in child.paramList().VARIABLE(): func_params.add(param.getText()) - _collect_reads_only(child, needed_vars, assigned_vars, func_params) + _collect_vars_from_node(child.block() or child.expr(), needed_vars, assigned_vars, func_params) # Top-level assignments elif rule_name == 'varDef': diff --git a/more_math/loraDictCommon.py b/more_math/loraDictCommon.py index b96afc6..88aaf80 100644 --- a/more_math/loraDictCommon.py +++ b/more_math/loraDictCommon.py @@ -5,7 +5,7 @@ import torch from collections.abc import Mapping -def calculate_lora_dict_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], use_compute_device=False): +def calculate_lora_dict_autogrow(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False): """ Calculate a new LORA_MODEL dict[str, Tensor] from math expressions applied @@ -157,7 +157,7 @@ def calculate_lora_dict_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], return result -def calculate_lora_dict_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], use_compute_device=False): +def calculate_lora_dict_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False): if V is None: V = {} if F is None: diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 33f14ff..656c001 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, LazyVariableDict, checkLazyNew +from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable, get_tensor_device, LazyVariableDict, checkLazyNew from .Parser.UnifiedMathVisitor import MathDict, ChangeTrackingDict, UnifiedMathVisitor import comfy import comfy.model_management @@ -109,7 +109,7 @@ def calculate_patches(Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0 """Legacy calculate_patches for backward compatibility.""" return calculate_patches_autogrow(Model, V={"V0": a, "V1": b, "V2": c, "V3": d}, F={"F0": w, "F1": x, "F2": y, "F3": z}, pbar=None, mapping={"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"}, use_compute_device=use_compute_device) -def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], use_compute_device=False): +def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False): """ Calculate patches for model-like objects (Model, VAE, CLIP) using Autogrow inputs. Iterates over the UNION of keys from all input models to support merging disjoint architectures/patches. @@ -192,8 +192,6 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us continue if original_device is None: original_device = get_tensor_device(current_sd) - if compute_device is not None and original_device is not None and compute_device != original_device: - current_sd = move_to_device(current_sd, compute_device) current_state_dicts[v_name] = current_sd # Progress bar if possible (comfy.utils.ProgressBar might assume unthreaded?) @@ -308,15 +306,8 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us variables = variables | generate_dim_variables(ref_tensor) - # Execute math - # Store original V0 before visiting, since visitor may modify it + # Read V0 before visiting, since the expression may rebind it original_v0 = variables.get("V0") - if original_v0 is not None: - # If it's a tensor, clone it; if it's a dict, copy its items - if isinstance(original_v0, torch.Tensor): - original_v0 = original_v0.clone() - elif isinstance(original_v0, dict): - original_v0 = MathDict({k: (v.clone() if isinstance(v, torch.Tensor) else v) for k, v in original_v0.items()}) visitor = UnifiedMathVisitor(variables, ref_tensor.shape,state_storage=stack) res = visitor.visit(tree) @@ -358,7 +349,7 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us return patches -def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], use_compute_device=False): +def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False): """ Calculate patches for model-like objects (Model, VAE, CLIP) using Dictionary mode. Evaluates the expression once for all keys simultaneously. @@ -423,8 +414,6 @@ def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], u continue if original_device is None: original_device = get_tensor_device(state_dict) - if compute_device is not None and original_device is not None and compute_device != original_device: - state_dict = move_to_device(state_dict, compute_device) state_dicts[name] = state_dict if not state_dicts: @@ -495,9 +484,9 @@ def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], u variables[f"{name}_d"] = make_delta_dict visitor_device = compute_device if compute_device is not None else next(iter(references.values())).device - # Store a copy of V0 BEFORE visiting, since the visitor may modify it + # Snapshot V0's entries BEFORE visiting, since the expression may reassign them v0_dict = variables.get("V0", MathDict({key: torch.zeros_like(reference) for key, reference in references.items()})) - originals = MathDict({k: (v.clone() if isinstance(v, torch.Tensor) else v) for k, v in v0_dict.items()}) + originals = dict(v0_dict) result = UnifiedMathVisitor(variables, device=visitor_device, state_storage=stack).visit(tree) if not isinstance(result, (MathDict, dict)): if isinstance(result, (float, int)): diff --git a/tests/test_assignment.py b/tests/test_assignment.py index 8d784e4..45a33b7 100644 --- a/tests/test_assignment.py +++ b/tests/test_assignment.py @@ -21,16 +21,19 @@ def test_tensor_assignment(): vars = {"T": T} # 1. Scalar to position - parse_and_visit("T[0, 0] = 5.0;", vars) - assert T[0, 0] == 5.0 + res = parse_and_visit("T[0, 0] = 5.0; T", vars) + assert res[0, 0] == 5.0 # 2. Scalar to slice (broadcast) - parse_and_visit("T[1] = 3.0;", vars) - assert torch.allclose(T[1], torch.tensor([3.0, 3.0])) + res = parse_and_visit("T[1] = 3.0; T", vars) + assert torch.allclose(res[1], torch.tensor([3.0, 3.0])) # 3. Tensor to slice - parse_and_visit("T[0] = [1, 2];", vars) - assert torch.allclose(T[0], torch.tensor([1.0, 2.0])) + res = parse_and_visit("T[0] = [1, 2]; T", vars) + assert torch.allclose(res[0], torch.tensor([1.0, 2.0])) + + # The input tensor itself is never written to + assert torch.equal(T, torch.zeros((2, 2))) def test_4d_assignment(): # val[0] expect 4d if dim0=1 or 3d @@ -39,17 +42,17 @@ def test_4d_assignment(): # Slice assignment val_3d = torch.ones((3, 4, 4)) - parse_and_visit("T[0] = V1;", {"T": T, "V1": val_3d}) - assert torch.allclose(T[0], val_3d) + res = parse_and_visit("T[0] = V1; T", {"T": T, "V1": val_3d}) + assert torch.allclose(res[0], val_3d) # 4D with leading 1 assignment val_4d_1 = torch.ones((1, 3, 4, 4)) * 2.0 - parse_and_visit("T[1] = V2;", {"T": T, "V2": val_4d_1}) - assert torch.allclose(T[1], val_4d_1[0]) + res = parse_and_visit("T[1] = V2; T", {"T": T, "V2": val_4d_1}) + assert torch.allclose(res[1], val_4d_1[0]) # Position assignment - parse_and_visit("T[0, 1, 2, 3] = 9.0;", vars) - assert T[0, 1, 2, 3] == 9.0 + res = parse_and_visit("T[0, 1, 2, 3] = 9.0; T", vars) + assert res[0, 1, 2, 3] == 9.0 def test_list_assignment(): L = [[1, 2], [3, 4]] @@ -67,19 +70,19 @@ def test_enhanced_assignment_logic(): # 1. Scalar filling: val[0] = scalar T = torch.zeros((2, 2, 2)) vars = {"T": T} - parse_and_visit("T[0] = 5.0;", vars) - assert torch.all(T[0] == 5.0) + res = parse_and_visit("T[0] = 5.0; T", vars) + assert torch.all(res[0] == 5.0) # 2. 1-element tensor/list filling - parse_and_visit("T[1] = [7.0];", vars) - assert torch.all(T[1] == 7.0) + res = parse_and_visit("T[1] = [7.0]; T", vars) + assert torch.all(res[1] == 7.0) # 3. Rank matching (squeeze leading 1s) # Target slice T[0, 0] is 1D (shape [2]) # Value is 3D (shape [1, 1, 2]) val_3d = torch.tensor([[[9.0, 9.0]]]) - parse_and_visit("T[0, 0] = V1;", {"T": T, "V1": val_3d}) - assert torch.allclose(T[0, 0], torch.tensor([9.0, 9.0])) + res = parse_and_visit("T[0, 0] = V1; T", {"T": T, "V1": val_3d}) + assert torch.allclose(res[0, 0], torch.tensor([9.0, 9.0])) if __name__ == "__main__": test_multidimensional_indexing() diff --git a/tests/test_unified_visitor_regressions.py b/tests/test_unified_visitor_regressions.py index 1c89645..3282959 100644 --- a/tests/test_unified_visitor_regressions.py +++ b/tests/test_unified_visitor_regressions.py @@ -7,7 +7,8 @@ import torch from PIL import ImageFont from more_math.Parser.UnifiedMathVisitor import ChangeTrackingDict, MathDict, UnifiedMathVisitor -from more_math.helper_functions import parse_expr +from more_math.helper_functions import checkLazyNew, parse_expr +from more_math.FloatMathNode import FloatMathNode def _visit(expression, variables): @@ -220,3 +221,31 @@ def test_perlin_uses_full_unit_range(): result = _visit("perlin(1, 8, 1, 0, [4096])", {}) assert result.min() < 0.05 and result.max() > 0.95 + + +def test_indexed_assignment_does_not_modify_input_tensor(): + a = torch.zeros(3) + result = _visit("a[0] = 5; a[1] = 6; a", {"a": a}) + + assert torch.equal(a, torch.zeros(3)) + assert torch.equal(result, torch.tensor([5.0, 6.0, 0.0])) + + +def test_user_function_named_like_builtin_visitor(): + assert _visit("sigmoid(x) -> x + 1; sigmoid(1)", {}) == 2 + + +@pytest.mark.parametrize("expression, needed", [ + ("w += 1; w", True), + ("f() -> w * 2; f()", True), + ("f(w) -> w * 2; f(1)", False), +]) +def test_lazy_check_finds_hidden_reads(expression, needed): + assert ("F0" in checkLazyNew(expression, {"V0": 1.0}, {"F0": None})) == needed + + +def test_unconnected_stack_is_not_shared_between_runs(): + FloatMathNode.execute("stack_push(0, 1)", {"V0": 0.0}, remember_stack=True) + _, stack = FloatMathNode.execute("stack_push(0, 1)", {"V0": 0.0}, remember_stack=True) + + assert stack == {0: [1.0]}