From f81cff61dfdd340533cc210fbc4d286ed957b915 Mon Sep 17 00:00:00 2001 From: Daniel Martinek Date: Wed, 1 Apr 2026 14:52:42 +0200 Subject: [PATCH] AI: multiple changes --- more_math/SelectiveGuiderMathNode.py | 182 ++++++++++++++++++++------- 1 file changed, 136 insertions(+), 46 deletions(-) diff --git a/more_math/SelectiveGuiderMathNode.py b/more_math/SelectiveGuiderMathNode.py index 8484b87..c859968 100644 --- a/more_math/SelectiveGuiderMathNode.py +++ b/more_math/SelectiveGuiderMathNode.py @@ -30,19 +30,33 @@ class SelectiveGuiderMathNode(io.ComfyNode): ), io.MultiType.Input( io.String.Input("Expression", default="push(0, inp); inp", multiline=False), - types=[io.String, MrmthParseTree] + types=[io.String, MrmthParseTree], ), - io.Combo.Input( - id="hook_target", - options=["all", "dit_block", "attn1", "attn2"], - default="all", - ), - io.Combo.Input( - id="hook_when", - options=["before_x", "just_before_x", "just_after_x", "after_x"], - default="after_x", - ), - io.Int.Input(id="layer_x", default=0, min=-999, max=999), + io.DynamicCombo.Input("hook_mode_cfg", options=[ + io.DynamicCombo.Option("single_target", [ + io.Combo.Input("hook_target", options=["dit_block", "attn1", "attn2"], default="attn1"), + io.DynamicCombo.Input("hook_pos", options=[ + io.DynamicCombo.Option("before_all", []), + io.DynamicCombo.Option("after_all", []), + io.DynamicCombo.Option("relative", [ + io.Combo.Input( + "hook_when", + options=["before_x", "just_before_x", "just_after_x", "after_x"], + default="after_x", + ), + io.Int.Input("layer_x", default=0, min=-999, max=999), + ]), + ]), + ]), + io.DynamicCombo.Option("all_targets", [ + io.Combo.Input( + "hook_when", + options=["before_x", "just_before_x", "just_after_x", "after_x"], + default="after_x", + ), + io.Int.Input("layer_x", default=0, min=-999, max=999), + ]), + ]), io.Boolean.Input(id="remember_stack", default=False), MrmthStack.Input(id="stack", optional=True), ], @@ -55,9 +69,7 @@ class SelectiveGuiderMathNode(io.ComfyNode): V, F, Expression, - hook_target="all", - hook_when="after_x", - layer_x=0, + hook_mode_cfg={}, remember_stack=False, stack={}, ): @@ -69,48 +81,89 @@ class SelectiveGuiderMathNode(io.ComfyNode): V, F, Expression, - hook_target="all", - hook_when="after_x", - layer_x=0, + hook_mode_cfg, remember_stack=False, stack=None, ): if V is None or not hasattr(V, "model_options"): raise ValueError("Vstupní guider musí mít model_options.") + # DynamicCombo -> interní config + mode = hook_mode_cfg.get("hook_mode_cfg", "all_targets") + if mode == "single_target": + hook_target = hook_mode_cfg.get("hook_target", "attn1") + hook_pos = hook_mode_cfg.get("hook_pos", "relative") + if hook_pos in ("before_all", "after_all"): + hook_when = hook_pos + layer_x = 0 + else: + hook_when = hook_mode_cfg.get("hook_when", "after_x") + layer_x = int(hook_mode_cfg.get("layer_x", 0)) + else: + hook_target = "all" + hook_when = hook_mode_cfg.get("hook_when", "after_x") + layer_x = int(hook_mode_cfg.get("layer_x", 0)) + stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {}) tree = parse_expr(Expression) if isinstance(Expression, str) else Expression - def match_when(idx: int, total_blocks: int | None = None) -> bool: - x = int(layer_x) + MAX_HOOK_INDEX = 999 - # záporné indexy od konce + def is_after_phase() -> bool: + return hook_when in ("after_all", "after_x", "just_after_x") + + def resolve_x(total_blocks: int | None) -> int: + x = int(layer_x) if x < 0 and total_blocks is not None: x = total_blocks + x # -1 => last - - # clamp if total_blocks is not None: x = max(0, min(x, total_blocks - 1)) + return x + + def match_index(idx: int, total_blocks: int | None = None) -> bool: + if hook_when == "before_all": + return idx == 0 + if hook_when == "after_all": + return total_blocks is not None and idx == (total_blocks - 1) + + x = resolve_x(total_blocks) if hook_when == "before_x": return idx < x if hook_when == "just_before_x": - # pro x==0 zůstane 0 (užitečné chování místo -1) return idx == (x - 1 if x > 0 else 0) if hook_when == "just_after_x": return idx == x - return idx >= x # after_x + return idx > x # after_x - def run_expr(inp: torch.Tensor) -> torch.Tensor: - stack["hook_last_inp"] = inp.detach() + def target_indices(): + # full range režimy + if hook_when in ("before_all", "after_all", "before_x", "after_x"): + return range(0, MAX_HOOK_INDEX + 1) + # záporné indexy bez total_blocks při registraci + if int(layer_x) < 0: + return range(0, MAX_HOOK_INDEX + 1) + + # just_* optimalizace + if hook_when == "just_before_x": + return [max(0, int(layer_x) - 1)] + if hook_when == "just_after_x": + return [int(layer_x)] + + return [int(layer_x)] + + def run_expr(inp: torch.Tensor, extra_vars: dict | None = None) -> torch.Tensor: variables = { "inp": inp, "sample": inp, } | generate_dim_variables(inp) - for k, v in F.items(): - variables[k] = v if v is not None else 0.0 + if extra_vars: + variables.update(extra_vars) + + for kf, vf in F.items(): + variables[kf] = vf if vf is not None else 0.0 visitor = UnifiedMathVisitor(variables, inp.shape, inp.device, state_storage=stack) out = visitor.visit(tree) @@ -125,41 +178,78 @@ class SelectiveGuiderMathNode(io.ComfyNode): def dit_wrap(args, extra_args): idx = int(args.get("transformer_options", {}).get("block_index", -1)) total_blocks = args.get("transformer_options", {}).get("total_blocks", None) - if not match_when(idx, total_blocks): + + if not match_index(idx, total_blocks): return extra_args["original_block"](args) - args2 = dict(args) - args2["img"] = run_expr(args["img"]) - return extra_args["original_block"](args2) + btype = args.get("transformer_options", {}).get("block_type", "dit") + meta = { + "hook_kind": "dit_block", + "layer_id": float(idx), + "block_name": btype, + "has_qkv": 0.0, + } + + # BEFORE + if not is_after_phase(): + args2 = dict(args) + args2["img"] = run_expr(args["img"], meta) + return extra_args["original_block"](args2) + + # AFTER + out = extra_args["original_block"](args) + if isinstance(out, dict) and "img" in out: + out2 = dict(out) + out2["img"] = run_expr(out["img"], meta) + return out2 + return out patches_replace = topts.get("patches_replace", {}).copy() dit_map = patches_replace.get("dit", {}).copy() - dit_map[("double_block", int(layer_x))] = dit_wrap - dit_map[("single_block", int(layer_x))] = dit_wrap + for li in target_indices(): + dit_map[("double_block", li)] = dit_wrap + dit_map[("single_block", li)] = dit_wrap patches_replace["dit"] = dit_map topts["patches_replace"] = patches_replace patched["transformer_options"] = topts def register_attn(attn_name: str): - def attn_replace(q, k, v, extra_options): idx = int(extra_options.get("block_index", -1)) total_blocks = extra_options.get("total_blocks", None) - if not match_when(idx, total_blocks): + + if not match_index(idx, total_blocks): return optimized_attention(q, k, v, extra_options["n_heads"], transformer_options=extra_options) - q2 = run_expr(q) - return optimized_attention(q2, k, v, extra_options["n_heads"], transformer_options=extra_options) + meta = { + "hook_kind": attn_name, + "i": float(idx), + "heads": float(extra_options.get("n_heads", 0)), + "dim_head": float(extra_options.get("dim_head", 0)), + "q": q, + "k": k, + "v": v, + } + + # BEFORE -> uprav q + if not is_after_phase(): + q2 = run_expr(q, meta) + return optimized_attention(q2, k, v, extra_options["n_heads"], transformer_options=extra_options) + + # AFTER -> uprav attention output + out = optimized_attention(q, k, v, extra_options["n_heads"], transformer_options=extra_options) + return run_expr(out, meta) nonlocal patched for blk in ("input", "middle", "output"): - patched = comfy.model_patcher.set_model_options_patch_replace( - patched, - attn_replace, - attn_name, - blk, - int(layer_x), - ) + for li in target_indices(): + patched = comfy.model_patcher.set_model_options_patch_replace( + patched, + attn_replace, + attn_name, + blk, + li, + ) if hook_target in ("all", "dit_block"): register_dit()