nodes/mcDandy-more_math
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

AI: multiple changes

Browse Source
This commit is contained in:
Daniel Martinek
2026-04-01 14:52:42 +02:00
parent d98b65bcc2
commit f81cff61df
1 changed files with 136 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
more_math/SelectiveGuiderMathNode.py
+136 -46
View File
@@ -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()