From 6be16be3fd59e551c591ee2051e0ce458b1cee58 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Fri, 11 Sep 2026 10:39:02 +0200 Subject: [PATCH] AI: add dict mode to model like nodes. TODO build grammer and check --- more_math/ClipMathNode.py | 15 ++- more_math/LoraMathNode.py | 15 ++- more_math/ModelMathNode.py | 15 ++- more_math/Parser/MathExpr.g4 | 7 +- more_math/Parser/UnifiedMathVisitor.py | 47 +++++++++ more_math/Parser/inbuilt_symbols.py | 5 +- more_math/VaeMathNode.py | 15 ++- more_math/loraDictCommon.py | 73 ++++++++++++- more_math/modelLikeCommon.py | 136 ++++++++++++++++++++++++- web/inbuilt_symbols.js | 5 +- 10 files changed, 312 insertions(+), 21 deletions(-) diff --git a/more_math/ClipMathNode.py b/more_math/ClipMathNode.py index ca38eb7..e79d71f 100644 --- a/more_math/ClipMathNode.py +++ b/more_math/ClipMathNode.py @@ -35,6 +35,12 @@ class CLIPMathNode(io.ComfyNode): default="error", tooltip="How to handle mismatched layer counts." ), + io.Boolean.Input( + id="dict_mode", + default=False, + display_name="Dictionary mode", + tooltip="Evaluates CLIP weights as dictionaries. Per-layer and dimension variables are unavailable.", + ), io.Boolean.Input( id="use_compute_device", default=True, @@ -52,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={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={}) -> io.NodeOutput: + def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False) -> io.NodeOutput: # Determine reference CLIP a = V.get("V0") stack = copy.deepcopy(stack) if stack is not None else {} @@ -76,13 +82,14 @@ class CLIPMathNode(io.ComfyNode): patchers_V[k] = v.patcher # Call autogrow patch calculation - from .modelLikeCommon import calculate_patches_autogrow + from .modelLikeCommon import calculate_patches_autogrow, calculate_patches_dict_mode 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,pbar=pbar, mapping=aliases,stack=stack,use_compute_device=use_compute_device) + calculate_patches = calculate_patches_dict_mode if dict_mode else calculate_patches_autogrow + patches = calculate_patches(Expression, V=patchers_V, F=F,pbar=pbar, mapping=aliases,stack=stack,use_compute_device=use_compute_device) out_clip = a.clone() if patches: diff --git a/more_math/LoraMathNode.py b/more_math/LoraMathNode.py index cef0ac8..987f980 100644 --- a/more_math/LoraMathNode.py +++ b/more_math/LoraMathNode.py @@ -4,7 +4,7 @@ import copy from .helper_functions import checkLazyNew from .Stack import MrmthStack from .ParseTree import MrmthParseTree -from .loraDictCommon import calculate_lora_dict_autogrow +from .loraDictCommon import calculate_lora_dict_autogrow, calculate_lora_dict_dict_mode class LoraMathNode(io.ComfyNode): @@ -32,6 +32,12 @@ class LoraMathNode(io.ComfyNode): types=[io.String, MrmthParseTree], tooltip="Expression to apply on LoRA tensors", ), + io.Boolean.Input( + id="dict_mode", + default=False, + display_name="Dictionary mode", + tooltip="Evaluates LoRA tensors as dictionaries. Per-layer and dimension variables are unavailable.", + ), io.Boolean.Input( id="use_compute_device", default=True, @@ -49,15 +55,16 @@ class LoraMathNode(io.ComfyNode): tooltip = cleandoc(__doc__) @classmethod - def check_lazy_status(cls, Expression, V, F, use_compute_device=True, stack={}): + def check_lazy_status(cls, Expression, V, F, use_compute_device=True, stack={}, dict_mode=False): return checkLazyNew(Expression, V, F) @classmethod - def execute(cls, V, F, Expression, use_compute_device=True, stack={}) -> io.NodeOutput: + def execute(cls, V, F, Expression, use_compute_device=True, stack={}, 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 {} 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) + calculate_lora_dict = calculate_lora_dict_dict_mode if dict_mode else calculate_lora_dict_autogrow + result = calculate_lora_dict(Expression, V=V, F=F, mapping=aliases, stack=stack, use_compute_device=use_compute_device) return (result, stack) diff --git a/more_math/ModelMathNode.py b/more_math/ModelMathNode.py index eb18170..3f3bf7d 100644 --- a/more_math/ModelMathNode.py +++ b/more_math/ModelMathNode.py @@ -5,7 +5,7 @@ from .Stack import MrmthStack from .ParseTree import MrmthParseTree import copy import comfy.utils -from .modelLikeCommon import calculate_patches_autogrow +from .modelLikeCommon import calculate_patches_autogrow, calculate_patches_dict_mode class ModelMathNode(io.ComfyNode): @@ -35,6 +35,12 @@ class ModelMathNode(io.ComfyNode): default="error", tooltip="How to handle mismatched layer counts. For models, this usually defaults to broadcast (zero for missing layers)." ), + io.Boolean.Input( + id="dict_mode", + default=False, + display_name="Dictionary mode", + tooltip="Evaluates model weights as dictionaries. Per-layer and dimension variables are unavailable.", + ), io.Boolean.Input( id="use_compute_device", default=True, @@ -53,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={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={}) -> io.NodeOutput: + def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False) -> io.NodeOutput: # Determine reference model for cloning a = V.get("V0") stack = copy.deepcopy(stack) if stack is not None else {} @@ -77,7 +83,8 @@ class ModelMathNode(io.ComfyNode): 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,pbar=pbar, mapping=aliases,stack=stack,use_compute_device=use_compute_device) + calculate_patches = calculate_patches_dict_mode if dict_mode else calculate_patches_autogrow + patches = calculate_patches(Expression, V=V, F=F,pbar=pbar, mapping=aliases,stack=stack,use_compute_device=use_compute_device) out_model = a.clone() if patches: diff --git a/more_math/Parser/MathExpr.g4 b/more_math/Parser/MathExpr.g4 index d6564e6..f712502 100644 --- a/more_math/Parser/MathExpr.g4 +++ b/more_math/Parser/MathExpr.g4 @@ -310,6 +310,10 @@ func1: count(x) - returns the number of elements in x */ | COUNT LPAREN expr RPAREN # CountFunc + /** + keys(x) - returns a list of dictionary keys in insertion order + */ + | KEYS LPAREN expr RPAREN # KeysFunc /** repeat(x, count, [dims]) - repeats tensor elements; count may be scalar or per-dim list */ @@ -977,6 +981,7 @@ DIST: 'dist' | 'distance'; REMAP: 'remap'; COSSIM: 'cossim' | 'cosine_similarity'; COUNT: 'count' | 'cnt' | 'length'; +KEYS: 'keys'; REPEAT: 'repeat'; FLATTEN: 'flatten'; APPEND: 'append'; @@ -1132,4 +1137,4 @@ VARIABLE: [a-zA-Z_] [a-zA-Z_0-9]*; SL_COMMENT: '#' ~[\r\n]* -> skip; ML_COMMENT: '/*' .*? '*/' -> skip; -WS: [ \t\r\n]+ -> skip; \ No newline at end of file +WS: [ \t\r\n]+ -> skip; diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index ccb3fdd..fd39e83 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -114,6 +114,28 @@ class UnifiedMathVisitor(MathExprVisitor): def _is_list(self, val): return isinstance(val, (list, tuple)) + def _is_math_dict(self, val): + return isinstance(val, MathDict) + + def _map_math_dict(self, op, *values): + dictionaries = [value for value in values if self._is_math_dict(value)] + if not dictionaries: + return op(*values) + + keys = tuple(dictionaries[0]) + if any(tuple(value) != keys for value in dictionaries[1:]): + raise ValueError("Dictionary operations require matching keys") + + return MathDict({key: op(*(value[key] if self._is_math_dict(value) else value for value in values)) for key in keys}) + + def _where_value(self, condition, true_value, false_value): + if self._is_tensor(condition) or self._is_tensor(true_value) or self._is_tensor(false_value): + condition_tensor = self._promote_to_tensor(condition) + if condition_tensor.dtype != torch.bool: + condition_tensor = condition_tensor != 0 + return torch.where(condition_tensor, self._promote_to_tensor(true_value), self._promote_to_tensor(false_value)).contiguous() + return true_value if bool(condition) else false_value + def _decode_string_literal(self, text): return text[1:-1].replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r').replace('\\\\', '\\').replace('\\"', '"').replace("\\'", "'") @@ -133,6 +155,8 @@ class UnifiedMathVisitor(MathExprVisitor): Generic binary operation handler. """ try: + if self._is_math_dict(a) or self._is_math_dict(b): + return self._map_math_dict(lambda left, right: self._bin_op(left, right, torch_op, scalar_op, ctx), a, b) if self._is_plain_tensor(a) and a.numel() == 1: a = float(a.flatten()[0].item()) if self._is_plain_tensor(b) and b.numel() == 1: @@ -228,6 +252,8 @@ class UnifiedMathVisitor(MathExprVisitor): def _unary_op(self, a, torch_op, scalar_op): + if self._is_math_dict(a): + return self._map_math_dict(lambda value: self._unary_op(value, torch_op, scalar_op), a) if self._is_plain_tensor(a) and a.numel() == 1: a = float(a.flatten()[0].item()) @@ -1092,6 +1118,8 @@ class UnifiedMathVisitor(MathExprVisitor): b = (yield ctx.expr(1)) w = (yield ctx.expr(2)) + if self._is_math_dict(a) or self._is_math_dict(b) or self._is_math_dict(w): + return self._map_math_dict(self._lerp_helper, a, b, w) if self._is_list(w): return [self._lerp_helper(a[i] if self._is_list(a) else a, b[i] if self._is_list(b) else b, @@ -1594,12 +1622,20 @@ class UnifiedMathVisitor(MathExprVisitor): def visitCountFunc(self, ctx): val = yield ctx.expr() + if isinstance(val, dict): + return float(len(val)) if self._is_list(val): return float(len(val)) if self._is_tensor(val): return val.numel() return 1.0 + def visitKeysFunc(self, ctx): + val = yield ctx.expr() + if not isinstance(val, dict): + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: keys() requires a dictionary") + return list(val) + def visitMeanFunc(self, ctx): return self._reduction_op( (yield ctx.expr()), lambda x: torch.mean(x.float()), lambda x: sum(x) / len(x) if x else 0.0 @@ -2431,6 +2467,15 @@ class UnifiedMathVisitor(MathExprVisitor): def visitCallExp(self, ctx): func_name = ctx.VARIABLE().getText() + if func_name == "keys": + args = ctx.exprList().expr() if ctx.exprList() else [] + if len(args) != 1: + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: keys() expects 1 argument") + value = yield args[0] + if not isinstance(value, dict): + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: keys() requires a dictionary") + return list(value) + # 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] @@ -4703,6 +4748,8 @@ class UnifiedMathVisitor(MathExprVisitor): a = (yield ctx.expr(1)) b = (yield ctx.expr(2)) + if self._is_math_dict(cond) or self._is_math_dict(a) or self._is_math_dict(b): + return self._map_math_dict(self._where_value, cond, a, b) if self._is_tensor(cond) or self._is_tensor(a) or self._is_tensor(b): cond_t = self._promote_to_tensor(cond) if cond_t.dtype != torch.bool: diff --git a/more_math/Parser/inbuilt_symbols.py b/more_math/Parser/inbuilt_symbols.py index 0deb877..d966b9e 100644 --- a/more_math/Parser/inbuilt_symbols.py +++ b/more_math/Parser/inbuilt_symbols.py @@ -9,6 +9,8 @@ INBUILT_FUNCTIONS = { 'abs', 'acos', 'acosh', 'add_key', 'all', 'angle', 'any', 'append', 'argmax', 'argmin', 'argsort', 'as_nested_tensor', 'asin', 'asinh', 'atan', 'atan2', 'atanh', 'band', 'batch_shuffle', 'bitcount', 'bitwise_and', 'bitwise_not', 'bitwise_or', 'bitwise_xor', 'blur', 'bnot', 'bor', 'botk', 'botk_ind', 'botk_indices', 'bxor', 'cat', 'ceil', 'cellular', 'cellular_noise', 'cielab_to_rgb', 'clamp', 'cnt', 'concat', 'concatenate', 'conv', 'convolution', 'coordinates', 'coords', 'corr', 'correlation', 'cos', 'cosh', 'cosine_similarity', 'cossim', 'count', 'cov', 'crop', 'cross', 'cubic', 'cubic_ease', 'cumprod', 'cumsum', 'diag', 'diagonal_matrix', 'dilate', 'dist', 'distance', 'domain_warp', 'domain_warp_noise', 'dot', 'edge', 'elastic', 'elastic_ease', 'endswith', 'entropy', 'erf', 'erfinv', 'erode', 'exp', 'ezconv', 'ezconvolution', 'fft', 'find', 'flatten', 'flip', 'float', 'floor', 'flow_ang', 'flow_angle', 'flow_apply', 'flow_mag', 'flow_magnitude', 'flow_to_image', 'fract', 'gamma', 'gaussian', 'gelu', 'get_value', 'hist', 'histogram', 'hsv_to_rgb', 'ifft', 'int', 'int_to_rgb', 'interpolate_area', 'interpolate_linear', 'interpolate_nearest', 'interpolate_nearest_exact', 'join', 'length', 'lerp', 'linspace', 'ln', 'log', 'logspace', 'lower', 'map', 'matmul', 'mean', 'median', 'mode', 'moment', 'morph_close', 'morph_open', 'motion_mask', 'nan_to_num', 'noise', 'now', 'nvl', 'oklab_to_rgb', 'overlay', 'pad', 'percentile', 'perlin', 'perlin_noise', 'perm', 'permute', 'pinv', 'plasma', 'plasma_noise', 'popcnt', 'popcount', 'pow', 'prcnt', 'print', 'print_shape', 'pshp', 'quantile', 'quartil', 'quartile', 'rand', 'randb', 'randbeta', 'randc', 'rande', 'randg', 'randgumbel', 'randchi2', 'randl', 'randln', 'randn', 'random_bernoulli', 'random_beta', 'random_cauchy', 'random_exponential', 'random_gamma', 'random_gumbel', 'random_chi2', 'random_laplace', 'random_log_normal', 'random_normal', 'random_poisson', 'random_studentt', 'random_uniform', 'random_weibull', 'randp', 'randt', 'randu', 'randw', 'range', 'relu', 'remap', 'remove_key', 'repeat', 'replace', 'reshape', 'rgb_to_cielab', 'rgb_to_hsv', 'rgb_to_int', 'rgb_to_oklab', 'ridged', 'ridged_noise', 'rife', 'rm_kay', 'roll', 'round', 'rshp', 'select', 'shape', 'shuffle', 'sigm', 'sign', 'sin', 'sine', 'sine_ease', 'singular_value_decomposition', 'sinh', 'smax', 'smin', 'smootherstep', 'smoothstep', 'snorm', 'softmax', 'softmin', 'softplus', 'sort', 'split', 'sqrt', 'squeeze', 'stack_clear', 'stack_get', 'stack_has', 'stack_pop', 'stack_push', 'startswith', 'std', 'step', 'substr', 'substring', 'sum', 'svd', 'swap', 'tan', 'tanh', 'tensor', 'text_image', 'timestamp', 'tmax', 'tmin', 'tnorm', 'topk', 'topk_ind', 'topk_indices', 'trim', 'turbulence', 'unique', 'unsqueeze', 'upper', 'var', 'voronoi', 'voronoi_noise', 'where', 'worley' } +INBUILT_FUNCTIONS.add('keys') + INBUILT_FUNCTION_META = { 'abs': {'min_args': 1, 'max_args': 1, 'snippet': 'abs()', 'description': 'abs(x) - applies per element absolute value function. Same as |x| for numbers.'}, 'acos': {'min_args': 1, 'max_args': 1, 'snippet': 'acos()', 'description': 'acos(x) - applies arcus cosinus function to value or each element of value'}, @@ -61,6 +63,7 @@ INBUILT_FUNCTION_META = { 'cosine_similarity': {'min_args': 2, 'max_args': 2, 'snippet': 'cosine_similarity()', 'description': 'cosine_similarity(x, y) - computes the cosine similarity between x and y.'}, 'cossim': {'min_args': 2, 'max_args': 2, 'snippet': 'cossim()', 'description': 'cosine_similarity(x, y) - computes the cosine similarity between x and y.'}, 'count': {'min_args': 1, 'max_args': 1, 'snippet': 'count()', 'description': 'count(x) - returns the number of elements in x'}, + 'keys': {'min_args': 1, 'max_args': 1, 'snippet': 'keys()', 'description': 'keys(x) - returns dictionary keys in insertion order'}, 'cov': {'min_args': 2, 'max_args': 2, 'snippet': 'cov()', 'description': 'cov(x, y) - computes the covariance matrix of x and y'}, 'crop': {'min_args': 3, 'max_args': 3, 'snippet': 'crop()', 'description': 'crop(x, start, size) - crops tensor x. Take size and begin at start. start should be the low corner. The result is at most size (if size is bigger then rest of tensor).'}, 'cross': {'min_args': 2, 'max_args': 2, 'snippet': 'cross()', 'description': 'cross(x, y) - computes the cross product of x and y. Last dimension must be 3'}, @@ -259,4 +262,4 @@ INBUILT_FUNCTION_META = { 'voronoi_noise': {'min_args': 2, 'max_args': 5, 'snippet': 'voronoi_noise()', 'description': 'cellular_noise(seed, scale, [jitter], [offset], [shape]) - generates Cellular/Voronoi noise'}, 'where': {'min_args': 3, 'max_args': 3, 'snippet': 'where()', 'description': 'where(condition, x, y) - returns x if condition is true, else y. Runs per element.'}, 'worley': {'min_args': 2, 'max_args': 5, 'snippet': 'worley()', 'description': 'cellular_noise(seed, scale, [jitter], [offset], [shape]) - generates Cellular/Voronoi noise'}, -} \ No newline at end of file +} diff --git a/more_math/VaeMathNode.py b/more_math/VaeMathNode.py index ef7f5f6..732b178 100644 --- a/more_math/VaeMathNode.py +++ b/more_math/VaeMathNode.py @@ -32,6 +32,12 @@ class VAEMathNode(io.ComfyNode): default="error", tooltip="How to handle mismatched layer counts. For models, this usually defaults to broadcast (zero for missing layers)." ), + io.Boolean.Input( + id="dict_mode", + default=False, + display_name="Dictionary mode", + tooltip="Evaluates VAE weights as dictionaries. Per-layer and dimension variables are unavailable.", + ), io.Boolean.Input( id="use_compute_device", default=True, @@ -49,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={}): + def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False): return checkLazyNew(Expression,V,F) @classmethod - def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={}) -> io.NodeOutput: + def execute(cls, V, F, Expression, length_mismatch="tile",use_compute_device=True,stack={},dict_mode=False) -> io.NodeOutput: # Determine reference VAE a = V.get("V0") if a is None: @@ -74,11 +80,12 @@ class VAEMathNode(io.ComfyNode): patchers_V[k] = v.patcher # Calculate patches using the patchers (weights are in patcher.model.state_dict) - from .modelLikeCommon import calculate_patches_autogrow + from .modelLikeCommon import calculate_patches_autogrow, calculate_patches_dict_mode 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, pbar=pbar, mapping=aliases, stack=stack, use_compute_device=use_compute_device) + calculate_patches = calculate_patches_dict_mode if dict_mode else calculate_patches_autogrow + patches = calculate_patches(Expression, V=patchers_V, F=F, pbar=pbar, mapping=aliases, stack=stack, use_compute_device=use_compute_device) # VAE does not have a clone method, so we shallow copy and clone the patcher out_vae = copy.copy(a) diff --git a/more_math/loraDictCommon.py b/more_math/loraDictCommon.py index 647c8f3..7be8e00 100644 --- a/more_math/loraDictCommon.py +++ b/more_math/loraDictCommon.py @@ -1,5 +1,5 @@ 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 +from .Parser.UnifiedMathVisitor import MathDict, UnifiedMathVisitor import comfy.model_management import torch @@ -153,3 +153,74 @@ def calculate_lora_dict_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], pbar.update(1) return result + + +def calculate_lora_dict_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], use_compute_device=False): + if V is None: + V = {} + if F is None: + F = {} + if mapping is None: + mapping = {} + + needed_vars = checkLazyNew(Expr, V, F) + needed_v_names = [name for name in V if name in needed_vars] + if "V0" in V and "V0" not in needed_v_names: + needed_v_names.append("V0") + if not needed_v_names: + needed_v_names = [name for name, value in V.items() if value is not None] + + input_dicts = {name: V[name] for name in needed_v_names if V.get(name) is not None} + if not input_dicts: + return {} + + all_keys = [] + seen_keys = set() + for value in input_dicts.values(): + for key in value: + if key not in seen_keys: + seen_keys.add(key) + all_keys.append(key) + + compute_device = comfy.model_management.get_torch_device() if use_compute_device else None + original_device = get_tensor_device(next(iter(input_dicts.values()))) + if compute_device is not None and original_device is not None and compute_device != original_device: + input_dicts = {name: move_to_device(value, compute_device) for name, value in input_dicts.items()} + + references = {} + for key in all_keys: + for value in input_dicts.values(): + tensor = value.get(key) + if tensor is not None: + references[key] = tensor + break + + def weight_dict(name): + value = input_dicts.get(name, {}) + return MathDict({key: value.get(key, torch.zeros_like(references[key])) for key in all_keys}) + + variables = LazyVariableDict({name: value if value is not None else 0.0 for name, value in F.items()}) + for alias, target in mapping.items(): + if target in F: + variables[alias] = F[target] if F[target] is not None else 0.0 + for name in needed_v_names: + variables[name] = weight_dict(name) + for alias, target in mapping.items(): + if target in V: + variables[alias] = variables.get(target, weight_dict(target)) + + visitor_device = compute_device if compute_device is not None else next(iter(references.values())).device + result = UnifiedMathVisitor(variables, device=visitor_device, state_storage=stack).visit(parse_expr(Expr) if isinstance(Expr, str) else Expr) + if not isinstance(result, MathDict): + if isinstance(result, (float, int)): + result = MathDict({key: torch.full_like(reference, result) for key, reference in references.items()}) + else: + raise ValueError("Dictionary mode expressions must return a dictionary") + if tuple(result) != tuple(all_keys): + raise ValueError("Dictionary mode expressions must preserve all LoRA keys") + + if compute_device is not None and original_device is not None and original_device != compute_device: + result = {key: value.to(device=original_device) for key, value in result.items()} + if pbar is not None: + pbar.update(len(result)) + return dict(result) diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 822c9c1..4aa4a89 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -1,10 +1,11 @@ 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 +from .Parser.UnifiedMathVisitor import MathDict, UnifiedMathVisitor import comfy import comfy.model_management import comfy.model_patcher import comfy.lora import torch +from collections import defaultdict def _get_model_patcher(obj): @@ -322,3 +323,136 @@ 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): + if V is None: + V = {} + if F is None: + F = {} + if mapping is None: + mapping = {} + + needed_vars = checkLazyNew(Expr, V, F) + needed_v_names = [name for name in V if name in needed_vars or f"{name}_d" in needed_vars] + if "V0" in V and "V0" not in needed_v_names: + needed_v_names.append("V0") + if not needed_v_names: + needed_v_names = [name for name, value in V.items() if value is not None] + + models = [value for value in V.values() if value is not None] + if not models: + return {} + + all_keys = [] + seen_keys = set() + for model in models: + state_dict = get_effective_state_dict(model) + if state_dict is None: + continue + for key in state_dict: + if key not in seen_keys: + seen_keys.add(key) + all_keys.append(key) + + tree = parse_expr(Expr) if isinstance(Expr, str) else Expr + compute_device = comfy.model_management.get_torch_device() if use_compute_device else None + original_device = None + state_dicts = {} + for name in needed_v_names: + value = V.get(name) + if value is None: + continue + state_dict = get_effective_state_dict(value) + if state_dict is None: + 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: + return {} + + all_keys = [key for key in all_keys if any(key in state_dict for state_dict in state_dicts.values())] + + delta_names = {name for name in needed_v_names if f"{name}_d" in needed_vars} + weights = {name: {} for name in needed_v_names} + bases = {name: {} for name in delta_names} + references = {} + + for key in all_keys: + for name in needed_v_names: + state_dict = state_dicts.get(name) + raw_weight = state_dict.get(key) if state_dict is not None else None + if raw_weight is None: + continue + model = V[name] + base_weight = get_base_weight(model, key) + if base_weight is None: + base_weight = raw_weight + base_weight = coerce_like(base_weight, raw_weight, device=compute_device) + patched_weight = get_patched_weight(model, key, base_weight) + if patched_weight is None: + patched_weight = raw_weight + patched_weight = coerce_like(patched_weight, raw_weight, device=compute_device) + weights[name][key] = patched_weight + if name in delta_names: + bases[name][key] = base_weight + if key not in references: + references[key] = patched_weight + + if not references: + return {} + + def weight_dict(name): + return MathDict({key: weights[name].get(key, torch.zeros_like(references[key])) for key in all_keys}) + + variables = LazyVariableDict({name: value if value is not None else 0.0 for name, value in F.items()}) + for alias, target in mapping.items(): + if target in F: + variables[alias] = F[target] if F[target] is not None else 0.0 + + for name in needed_v_names: + variables[name] = weight_dict(name) + + for alias, target in mapping.items(): + if target in V: + variables[alias] = variables.get(target, weight_dict(target)) + + for name in delta_names: + def make_delta_dict(name=name): + return MathDict({ + key: weights[name][key] - bases[name][key] if key in weights[name] else torch.zeros_like(references[key]) + for key in all_keys + }) + + make_delta_dict.is_lazy_var = True + variables[f"{name}_d"] = make_delta_dict + + visitor_device = compute_device if compute_device is not None else next(iter(references.values())).device + result = UnifiedMathVisitor(variables, device=visitor_device, state_storage=stack).visit(tree) + if not isinstance(result, MathDict): + if isinstance(result, (float, int)): + result = MathDict({key: torch.full_like(reference, result) for key, reference in references.items()}) + else: + raise ValueError("Dictionary mode expressions must return a dictionary") + + if tuple(result) != tuple(all_keys): + raise ValueError("Dictionary mode expressions must preserve all model keys") + + originals = variables.get("V0", MathDict({key: torch.zeros_like(reference) for key, reference in references.items()})) + patches = {} + for key, value in result.items(): + if not torch.is_tensor(value): + value = torch.full_like(references[key], value) + diff = value - originals[key] + if not torch.all(diff == 0): + patches[key] = (diff,) + if pbar is not None: + pbar.update(1) + + if compute_device is not None and original_device is not None and original_device != compute_device: + patches = {key: (value[0].to(device=original_device),) for key, value in patches.items()} + return patches diff --git a/web/inbuilt_symbols.js b/web/inbuilt_symbols.js index 4097f80..d0743d0 100644 --- a/web/inbuilt_symbols.js +++ b/web/inbuilt_symbols.js @@ -262,4 +262,7 @@ export const FUNCTION_META = { voronoi_noise: { minArgs: 2, maxArgs: 5, snippet: "voronoi_noise()", description: "cellular_noise(seed, scale, [jitter], [offset], [shape]) - generates Cellular/Voronoi noise" }, where: { minArgs: 3, maxArgs: 3, snippet: "where()", description: "where(condition, x, y) - returns x if condition is true, else y. Runs per element." }, worley: { minArgs: 2, maxArgs: 5, snippet: "worley()", description: "cellular_noise(seed, scale, [jitter], [offset], [shape]) - generates Cellular/Voronoi noise" }, -}; \ No newline at end of file +}; + +FUNCTIONS.add("keys"); +FUNCTION_META.keys = { minArgs: 1, maxArgs: 1, snippet: "keys()", description: "keys(x) - returns dictionary keys in insertion order" };