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

AI sgmn runs

Browse Source
This commit is contained in:
Daniel Martinek
2026-04-01 21:13:30 +02:00
parent e6fc81d1a3
commit cff12d66e0
1 changed files with 179 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
more_math/SelectiveGuiderMathNode.py
+179 -72
View File
@@ -2,6 +2,9 @@ import copy
import torch
import comfy.model_patcher
from comfy_api.latest import io
import comfy.hooks
import logging
import comfy.patcher_extension
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
from .helper_functions import checkLazyNew, generate_dim_variables, parse_expr, as_tensor
@@ -88,7 +91,6 @@ class SelectiveGuiderMathNode(io.ComfyNode):
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")
@@ -115,65 +117,75 @@ class SelectiveGuiderMathNode(io.ComfyNode):
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
x = total_blocks + x
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:
def match_index(idx: int, total_blocks: int | None) -> bool:
# fallback: když neznáme total_blocks, "all" režimy nefiltruj
if total_blocks is None:
if hook_when in ("before_all", "after_all"):
return True
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":
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
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 side_from_cond_or_uncond(cond_or_uncond):
if isinstance(cond_or_uncond, (list, tuple)) and len(cond_or_uncond) > 0:
vals = list(cond_or_uncond)
uniq = set(vals)
if len(uniq) == 1:
idx = int(vals[0])
if idx == 0:
return ("positive", 0)
if idx == 1:
return ("negative", 1)
return ("mixed", -1)
return ("unknown", -1)
def run_expr(inp: torch.Tensor, extra_vars: dict | None = None) -> torch.Tensor:
# universal defaults: always defined
variables = {
"inp": inp,
"sample": inp,
# hook identity
"hook_kind": "unknown",
"block_name": "",
"layer_key": "",
"layer_id": -1.0,
"i": -1.0,
"total_blocks": -1.0,
# attention-related (always present, valid flag says if real)
"has_qkv": 0.0,
"q": inp,
"k": inp,
"v": inp,
"heads": 0.0,
"dim_head": 0.0,
"cond_side": "unknown",
"cond_index": -1.0,
"is_positive": 0.0,
"is_negative": 0.0,
} | generate_dim_variables(inp)
if extra_vars:
@@ -186,11 +198,8 @@ class SelectiveGuiderMathNode(io.ComfyNode):
out = visitor.visit(tree)
return as_tensor(out, inp.shape).to(inp.device)
patched = comfy.model_patcher.create_model_options_clone(V.model_options)
patched.setdefault("transformer_options", {})
def register_dit():
topts = patched["transformer_options"]
def register_dit(patched_dict, forced_side=None):
topts = patched_dict["transformer_options"]
def dit_wrap(args, extra_args):
idx = int(args.get("transformer_options", {}).get("block_index", -1))
@@ -198,8 +207,15 @@ class SelectiveGuiderMathNode(io.ComfyNode):
if not match_index(idx, total_blocks):
return extra_args["original_block"](args)
if not hit_flags["dit"]:
logging.warning(f"[SelectiveGuiderMathNode] DIT HIT idx={idx} total={total_blocks}")
hit_flags["dit"] = True
btype = args.get("transformer_options", {}).get("block_type", "dit")
co = args.get("transformer_options", {}).get("cond_or_uncond", None)
auto_side, auto_idx = side_from_cond_or_uncond(co)
side = forced_side if forced_side is not None else auto_side
idx_side = 0 if side == "positive" else (1 if side == "negative" else auto_idx)
meta = {
"hook_kind": "dit_block",
"layer_id": float(idx),
@@ -208,15 +224,17 @@ class SelectiveGuiderMathNode(io.ComfyNode):
"block_name": btype,
"layer_key": f"dit.{btype}.{idx}",
"has_qkv": 0.0,
"cond_side": side,
"cond_index": float(idx_side),
"is_positive": 1.0 if side == "positive" else 0.0,
"is_negative": 1.0 if side == "negative" else 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)
@@ -231,63 +249,152 @@ class SelectiveGuiderMathNode(io.ComfyNode):
dit_map[("single_block", li)] = dit_wrap
patches_replace["dit"] = dit_map
topts["patches_replace"] = patches_replace
patched["transformer_options"] = topts
patched_dict["transformer_options"] = topts
def register_attn(attn_name: str):
def make_attn_replace(block_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)
def register_attn(patched_dict, attn_name: str, forced_side=None):
topts = patched_dict["transformer_options"].copy()
if not match_index(idx, total_blocks):
return optimized_attention(q, k, v, extra_options["n_heads"], transformer_options=extra_options)
def attn_override(original_attn, q, k, v, heads, **kwargs):
transformer_options = kwargs.get("transformer_options", {})
meta = {
"hook_kind": attn_name,
"block_name": block_name,
"layer_id": float(idx),
"i": float(idx),
"total_blocks": float(total_blocks) if total_blocks is not None else -1.0,
"has_qkv": 1.0,
blk = transformer_options.get("block", None) # ("input"/"middle"/"output", idx)
if isinstance(blk, (tuple, list)) and len(blk) >= 2:
stage = str(blk[0])
layer_id = int(blk[1])
else:
stage = "unknown"
layer_id = int(transformer_options.get("block_index", -1))
"q": q,
"k": k,
"v": v,
"heads": float(extra_options.get("n_heads", 0)),
"dim_head": float(extra_options.get("dim_head", 0)),
"layer_key": f"{block_name}.{attn_name}.{idx}",
}
# filtr podle cíle attn1/attn2
# v SD1 je spolehlivé rozlišení přes aktuální index uvnitř transformeru menší;
# proto zatím pouštíme obě a filtrujeme až expression logikou přes hook_kind/layer_key.
idx = layer_id
total_blocks = None
# BEFORE -> modify 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)
if not match_index(idx, total_blocks):
return original_attn(q, k, v, heads, **kwargs)
# AFTER -> modify attn output
out = optimized_attention(q, k, v, extra_options["n_heads"], transformer_options=extra_options)
return run_expr(out, meta)
return attn_replace
co = transformer_options.get("cond_or_uncond", None)
auto_side, auto_idx = side_from_cond_or_uncond(co)
side = forced_side if forced_side is not None else auto_side
idx_side = 0 if side == "positive" else (1 if side == "negative" else auto_idx)
nonlocal patched
for blk in ("input", "middle", "output"):
attn_replace = make_attn_replace(blk)
for li in target_indices():
patched = comfy.model_patcher.set_model_options_patch_replace(
patched,
attn_replace,
attn_name,
blk,
li,
)
if not hit_flags["attn"]:
logging.warning(f"[SelectiveGuiderMathNode] ATTN HIT stage={stage} layer_id={layer_id} side={side}")
hit_flags["attn"] = True
if hook_target in ("all", "dit_block"):
register_dit()
if hook_target in ("all", "attn1"):
register_attn("attn1")
if hook_target in ("all", "attn2"):
register_attn("attn2")
meta = {
"hook_kind": attn_name,
"block_name": stage,
"layer_id": float(layer_id),
"i": float(layer_id),
"total_blocks": -1.0,
"has_qkv": 1.0,
"q": q,
"k": k,
"v": v,
"heads": float(heads),
"dim_head": float(q.shape[-1] // heads) if heads > 0 else 0.0,
"layer_key": f"{stage}.{layer_id}.{attn_name}",
"cond_side": side,
"cond_index": float(idx_side),
"is_positive": 1.0 if side == "positive" else 0.0,
"is_negative": 1.0 if side == "negative" else 0.0,
}
V.model_options = patched
if not is_after_phase():
q2 = run_expr(q, meta)
return original_attn(q2, k, v, heads, **kwargs)
out = original_attn(q, k, v, heads, **kwargs)
return run_expr(out, meta)
topts["optimized_attention_override"] = attn_override
patched_dict["transformer_options"] = topts
return patched_dict
def build_transformers_for_side(side_label=None):
p = comfy.model_patcher.create_model_options_clone(V.model_options)
p.setdefault("transformer_options", {})
if hook_target in ("all", "dit_block"):
register_dit(p, forced_side=side_label)
if hook_target in ("all", "attn1"):
p = register_attn(p, "attn1", forced_side=side_label)
if hook_target in ("all", "attn2"):
p = register_attn(p, "attn2", forced_side=side_label)
return p["transformer_options"]
def attach_side_hook(side_name: str):
if not hasattr(V, "original_conds") or side_name not in V.original_conds:
return
tdict = build_transformers_for_side(side_name)
hook = comfy.hooks.TransformerOptionsHook(
transformers_dict=tdict,
hook_scope=comfy.hooks.EnumHookScope.HookedOnly,
)
hg = comfy.hooks.HookGroup()
hg.add(hook)
new_conds = []
for c in V.original_conds[side_name]:
cc = c.copy()
old = cc.get("hooks", None)
cc["hooks"] = comfy.hooks.HookGroup.combine_all_hooks([old, hg])
new_conds.append(cc)
V.original_conds[side_name] = new_conds
if hasattr(V, "original_conds"):
attach_side_hook("positive")
attach_side_hook("negative")
logging.warning(f"[SelectiveGuiderMathNode] has original_conds: {list(V.original_conds.keys())}")
for side in ("positive", "negative"):
if side in V.original_conds:
cnt = sum(1 for c in V.original_conds[side] if c.get("hooks", None) is not None)
logging.warning(f"[SelectiveGuiderMathNode] side={side} conds_with_hooks={cnt}/{len(V.original_conds[side])}")
else:
patched = comfy.model_patcher.create_model_options_clone(V.model_options)
patched.setdefault("transformer_options", {})
if hook_target in ("all", "dit_block"):
register_dit(patched, forced_side=None)
if hook_target in ("all", "attn1"):
patched = register_attn(patched, "attn1", forced_side=None)
if hook_target in ("all", "attn2"):
patched = register_attn(patched, "attn2", forced_side=None)
V.model_options = patched
if hasattr(V, "model_patcher") and V.model_patcher is not None and not getattr(V, "_mrmth_debug_callbacks_added", False):
hit_flags = {"attn": False, "dit": False}
def _on_register_all_hook_patches(model_patcher, hooks, target_dict, model_options, registered):
try:
reg_len = len(registered) if registered is not None else 0
in_len = len(hooks) if hooks is not None else 0
tr_count = len(registered.get_type(comfy.hooks.EnumHookType.TransformerOptions)) if registered is not None else 0
logging.warning(f"[SelectiveGuiderMathNode] REGISTER hooks_in={in_len} registered={reg_len} transformer_registered={tr_count}")
if model_options is not None:
rh = model_options.get("registered_hooks", None)
logging.warning(f"[SelectiveGuiderMathNode] model_options.registered_hooks={'set' if rh is not None else 'None'}")
except Exception as e:
logging.warning(f"[SelectiveGuiderMathNode] REGISTER debug error: {e}")
def _on_apply_hooks(model_patcher, hooks):
try:
hl = len(hooks) if hooks is not None else 0
logging.warning(f"[SelectiveGuiderMathNode] APPLY hooks_len={hl}")
except Exception as e:
logging.warning(f"[SelectiveGuiderMathNode] APPLY debug error: {e}")
V.model_patcher.add_callback_with_key(
comfy.patcher_extension.CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES,
"mrmth_selective_debug",
_on_register_all_hook_patches,
)
V.model_patcher.add_callback_with_key(
comfy.patcher_extension.CallbacksMP.ON_APPLY_HOOKS,
"mrmth_selective_debug",
_on_apply_hooks,
)
V._mrmth_debug_callbacks_added = True
stack = stack if remember_stack else copy.deepcopy(stack)
# do execute(), před register_*:
return (V, stack)