diff --git a/more_math/ClipMathNode.py b/more_math/ClipMathNode.py index a1a4245..3181d4d 100644 --- a/more_math/ClipMathNode.py +++ b/more_math/ClipMathNode.py @@ -6,6 +6,7 @@ from .Parser.MathExprParser import MathExprParser import re from .Stack import MrmthStack import copy +import comfy.utils @@ -95,12 +96,11 @@ class CLIPMathNode(io.ComfyNode): # Call autogrow patch calculation from .modelLikeCommon import calculate_patches_autogrow - # Populate aliases map for backward compatibility in expressions (users might still use a/b/c/d/w/x/y/z) - # a=V0, b=V1, etc created by us or expected by user? - # The prompt says aliases are supported in check_lazy_status. Variables map in helper handles logic. + layer_count = V.get("V0").model.state_dict().__len__() if hasattr(V.get("V0"), "model") and hasattr(V.get("V0").model, "state_dict") else 0 + pbar = comfy.utils.ProgressBar(layer_count) aliases = {"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"} - patches = calculate_patches_autogrow(Expression, V=patchers_V, F=F, mapping=aliases,stack=stack) + patches = calculate_patches_autogrow(Expression, V=patchers_V, F=F,pbar=pbar, mapping=aliases,stack=stack) out_clip = a.clone() if patches: diff --git a/more_math/ModelMathNode.py b/more_math/ModelMathNode.py index fa42e8a..3e42a70 100644 --- a/more_math/ModelMathNode.py +++ b/more_math/ModelMathNode.py @@ -6,6 +6,9 @@ from .Parser.MathExprParser import MathExprParser import re from .Stack import MrmthStack import copy +import comfy.utils +from .modelLikeCommon import calculate_patches_autogrow + class ModelMathNode(io.ComfyNode): """ @@ -88,12 +91,13 @@ class ModelMathNode(io.ComfyNode): if a is None: raise ValueError("At least one input model is required.") + layer_count = V.get("V0").model.state_dict().__len__() if hasattr(V.get("V0"), "model") and hasattr(V.get("V0").model, "state_dict") else 0 + pbar = comfy.utils.ProgressBar(layer_count) - from .modelLikeCommon import calculate_patches_autogrow aliases = {"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"} - patches = calculate_patches_autogrow(Expression, V=V, F=F, mapping=aliases,stack=stack) + patches = calculate_patches_autogrow(Expression, V=V, F=F,pbar=pbar, mapping=aliases,stack=stack) out_model = a.clone() if patches: diff --git a/more_math/VaeMathNode.py b/more_math/VaeMathNode.py index eb2d158..41e3ee6 100644 --- a/more_math/VaeMathNode.py +++ b/more_math/VaeMathNode.py @@ -7,6 +7,7 @@ from .Parser.MathExprLexer import MathExprLexer from .Parser.MathExprParser import MathExprParser import re from .Stack import MrmthStack +import comfy.utils class VAEMathNode(io.ComfyNode): """ @@ -96,6 +97,8 @@ class VAEMathNode(io.ComfyNode): # Calculate patches using the patchers (weights are in patcher.model.state_dict) from .modelLikeCommon import calculate_patches_autogrow aliases = {"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"} + layer_count = V.get("V0").model.state_dict().__len__() if hasattr(V.get("V0"), "model") and hasattr(V.get("V0").model, "state_dict") else 0 + pbar = comfy.utils.ProgressBar(layer_count) patches = calculate_patches_autogrow(Expression, V=patchers_V, F=F, mapping=aliases,stack=stack) # VAE does not have a clone method, so we shallow copy and clone the patcher diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index c134643..78c13c8 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -1,13 +1,14 @@ from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable from .Parser.UnifiedMathVisitor import UnifiedMathVisitor import torch +import comfy.utils def calculate_patches(Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.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}, mapping={"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"}) -def calculate_patches_autogrow(Expr, V, F, mapping=None,stack = []): +def calculate_patches_autogrow(Expr, V, F,pbar, mapping=None,stack = []): """ 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. @@ -53,6 +54,9 @@ def calculate_patches_autogrow(Expr, V, F, mapping=None,stack = []): # Progress bar if possible (comfy.utils.ProgressBar might assume unthreaded?) # Just skip for utility or use if substantial. + all_keys_list = list(all_keys) + layer_count = len(all_keys_list) + for layer_idx, key in enumerate(all_keys_list): for key in all_keys: variables = {} @@ -66,6 +70,11 @@ def calculate_patches_autogrow(Expr, V, F, mapping=None,stack = []): if target in F: variables[alias] = F[target] if F[target] is not None else 0.0 + variables["L"] = float(layer_idx) + variables["layer"] = float(layer_idx) + variables["LC"] = float(layer_count) + variables["layer_count"] = float(layer_count) + # Inject weights for this key from V models valid_key = False ref_tensor = None @@ -130,5 +139,6 @@ def calculate_patches_autogrow(Expr, V, F, mapping=None,stack = []): if not torch.all(diff == 0): patches[key] = (diff,) + pbar.update(1) return patches