File size: 19,242 Bytes
e488ce7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 | """The ComfyUI surface of concept_lab: three thin nodes, no rules.
There is exactly one reason these exist this early (DEC-CL-0013 said nodes
wait for T12, DEC-CL-0018 amends its TIMING): a capture has to happen inside
the ComfyUI process, so the seam needs a node before anything can be
measured. They hold no policy. The Arm node builds a TapSpec, reads a stack
fingerprint off the live ModelPatcher, opens a TapSession and installs it on
a CLONE of the model; the Flush node calls finalize(). Every rule is in
concept_lab/, and every refusal comes from there too.
[MODEL] -> MAI Concept Capture Arm -> [MODEL] -> KSampler -> ... -> decode
| |
+--> CONCEPT_CAPTURE -> MAI Concept Capture Flush <+
(after)
The `after` input is what orders the flush behind the decode. ComfyUI runs
the graph by data dependency, and a flush that is not downstream of anything
can execute before the sampler has taken a single step.
The third node is the mirror of the first (DEC-CL-0020): where the tap READS
block states, `MAI Concept Inject Delta` ADDS a stored one back, so the same
seam serves the capture and the premise test that spends it. It holds no
rules either; the pack, the geometry guard and the controls live in
`concept_lab/backends/h3_inject.py`.
ALPHA. The pass-through gate (E0: tap-installed-disabled renders must be
pixel-identical to the same seed unpatched) has NOT been run live yet. If
these nodes fail it, they get pulled rather than patched around.
"""
from __future__ import annotations
import json
import os
import sys
# ComfyUI loads a pack by file location and does NOT put the pack directory
# on sys.path (nodes.py:2243-2265), while every module under concept_lab/
# imports itself absolutely (`from concept_lab import ...`) so that the CLI,
# the tests and the MCP server all address one module identity. Same idiom
# as cli.py:33-34.
_PACK = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _PACK not in sys.path:
sys.path.insert(0, _PACK)
from concept_lab import types as T # noqa: E402
from concept_lab.backends import h3_inject # noqa: E402
from concept_lab.backends import h3_tap # noqa: E402
CATEGORY = "MAI/concept (alpha)"
class _AnyType(str):
def __ne__(self, other):
return False
ANY = _AnyType("*")
def _ints(s, field, who="MAI Concept Capture Arm"):
out = []
for part in str(s or "").split(","):
part = part.strip()
if not part:
continue
try:
out.append(int(part))
except ValueError:
raise ValueError(f"{who}: {field} wants "
f"comma-separated integers, got {part!r}")
return tuple(out)
def _floats(s, field):
out = []
for part in str(s or "").split(","):
part = part.strip()
if not part:
continue
try:
out.append(float(part))
except ValueError:
raise ValueError(f"MAI Concept Capture Arm: {field} wants "
f"comma-separated numbers, got {part!r}")
return tuple(out)
def _words(s):
return tuple(p.strip() for p in str(s or "").split(",") if p.strip())
class MAIConceptCaptureArm:
"""Install a capture tap on one arm of a comparison."""
DESCRIPTION = (
"EXPERIMENTAL (alpha), new 2026-08-15. Installs the concept_lab "
"block tap on a CLONE of the model and hands the clone on, so the "
"sampler downstream renders exactly as it would have while selected "
"DiT block states are written to the concept space.\n\n"
"One node is ONE ARM of one comparison (R = the actual reference, "
"N = a shape-matched control, 0 = no reference). All arms of a "
"comparison have to run in ONE ComfyUI process: this pipeline is "
"bit-identical within a launch and only close across launches, so a "
"restart makes a new replicate rather than the same condition "
"(DEC-CL-0015).\n\n"
"enabled=False is the pass-through control: the tap is installed and "
"records nothing, which is how you prove the instrumentation is free."
)
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"model": ("MODEL", {"tooltip": "the model this arm renders with; "
"a patched CLONE comes out, the input is "
"untouched"}),
"arm": (list(T.ARMS), {"default": "R",
"tooltip": "R actual reference / N shape-matched control "
"/ 0 no reference / free (uncontrolled)"}),
"variant_name": ("STRING", {"default": "",
"tooltip": "a label; identity comes from what "
"the model actually saw (text "
"context, reference latents), "
"measured by the tap"}),
"seed": ("INT", {"default": 0, "min": 0,
"max": 0xffffffffffffffff,
"tooltip": "MUST equal the sampler's noise seed. This "
"node cannot read it (the seed lives in the "
"sampler, not in the model), so it is copied "
"by hand and a mismatch silently mislabels "
"the capture; the tap now refuses on "
"mismatch"}),
"replicate": ("INT", {"default": 0, "min": 0, "max": 999,
"tooltip": "which repeat of this condition, inside "
"this launch"}),
"blocks": ("STRING", {"default": "0,5,10,20,30,40,49",
"tooltip": "DiT block indices to tap. An experiment "
"config, not a format fact"}),
"step_quantiles": ("STRING", {"default": "0.15,0.5,0.85",
"tooltip": "schedule-relative positions; 0 is "
"the first step, 1 the last"}),
"steps": ("STRING", {"default": "",
"tooltip": "explicit step indices instead; giving both "
"this and step_quantiles is refused"}),
"segments": ("STRING", {"default": "video,ref_img",
"tooltip": "packed-layout segment names: text, cond, "
"cond_audio, ref_img, ref_audio, audio, "
"video"}),
"store": ("STRING", {"default": "stats,frame_mean,sketch",
"tooltip": "comma-joined subset of stats, frame_mean, "
"sketch, full. 'full' is the whole state and "
"costs accordingly"}),
"sketch_dim": ("INT", {"default": 256, "min": 8, "max": 4096,
"tooltip": "width of the seeded random projection"}),
"sketch_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff,
"tooltip": "recorded in the manifest; an "
"unrecorded projection is an "
"unreproducible capture"}),
"stack_label": ("STRING", {"default": "",
"tooltip": "the checkpoint name, from you. A "
"ModelPatcher does not carry the UNET "
"filename and this node will not guess "
"one"}),
}, "optional": {
"sampler_label": ("STRING", {"default": "",
"tooltip": "sampler/scheduler/shift as you ran "
"them; it joins the capture's "
"identity, so two arms must agree"}),
"corpus_row_id": ("STRING", {"default": "",
"tooltip": "the corpus row this arm renders, if "
"there is one"}),
"enabled": ("BOOLEAN", {"default": True,
"tooltip": "False installs the tap and records "
"NOTHING: the pass-through control"}),
"cond_only": ("BOOLEAN", {"default": True,
"tooltip": "skip the unconditional forward; CFG "
"makes two forwards per step and they "
"are two different conditions"}),
"index_root": ("STRING", {"default": "",
"tooltip": "manifest root; empty = the space "
"default (H3_CONCEPT_INDEX_DIR)"}),
"tensor_root": ("STRING", {"default": "",
"tooltip": "tensor root; empty = the space "
"default (H3_CONCEPT_TENSOR_DIR)"}),
"notes": ("STRING", {"default": "", "multiline": True}),
}}
RETURN_TYPES = ("MODEL", "CONCEPT_CAPTURE")
RETURN_NAMES = ("model", "capture")
FUNCTION = "arm"
CATEGORY = CATEGORY
@classmethod
def IS_CHANGED(cls, **kwargs):
# A capture is a measurement, and a cached measurement is not one:
# re-running the same graph must build a NEW session and record
# again, so this node never reports itself unchanged.
return float("nan")
def arm(self, model, arm, variant_name, seed, replicate, blocks,
step_quantiles, steps, segments, store, sketch_dim, sketch_seed,
stack_label, sampler_label="", corpus_row_id="", enabled=True,
cond_only=True, index_root="", tensor_root="", notes=""):
step_idx = _ints(steps, "steps")
tap = T.TapSpec(
blocks=_ints(blocks, "blocks"),
steps=step_idx,
step_quantiles=() if step_idx else _floats(step_quantiles,
"step_quantiles"),
segments=_words(segments), store=store,
sketch_dim=int(sketch_dim), sketch_seed=int(sketch_seed))
from concept_lab.backends import get_backend
get_backend("h3").validate_tap(tap)
stack = h3_tap.fingerprint_from_patcher(
model, stack_label or "UNNAMED (stack_label was left empty)")
session = h3_tap.TapSession(
tap=tap, arm=arm, variant_name=variant_name or "unnamed",
seed=int(seed), replicate=int(replicate), stack=stack,
index_root=index_root or None, tensor_root=tensor_root or None,
corpus_row_id=corpus_row_id or None, enabled=bool(enabled),
cond_only=bool(cond_only), notes=notes,
sampler_label=sampler_label)
return (session.install(model), session)
class MAIConceptCaptureFlush:
"""File the capture. Wire `after` to something the decode produced."""
DESCRIPTION = (
"EXPERIMENTAL (alpha), new 2026-08-15. Finalizes a capture: writes "
"the manifest into the concept space and reports what landed.\n\n"
"`after` accepts anything and is used ONLY to order execution. "
"ComfyUI runs a graph by data dependency, so a flush that is not "
"downstream of the render can fire before the sampler has taken a "
"step; wire the decoded images (or whatever the last node emits) "
"into `after` and the flush happens when the render is done.\n\n"
"The tensors are already on disk by then: the tap writes as it "
"captures rather than holding a run's worth of state in VRAM.")
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"capture": ("CONCEPT_CAPTURE", {}),
"after": (ANY, {"tooltip": "execution ordering only; wire the "
"decode's output so the flush runs last"}),
}}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("report",)
FUNCTION = "flush"
CATEGORY = CATEGORY
OUTPUT_NODE = True
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("nan")
def flush(self, capture, after=None):
# `after` is the decode's IMAGE batch in every graph that wires this
# correctly, so it is also free evidence about what the render came
# out as: hashed, recorded in the notes, never part of identity
# (DEC-CL-0006). Anything that is not an image tensor records null.
return (json.dumps(
capture.finalize(
render_fingerprint=h3_tap.render_fingerprint(after)),
indent=2, sort_keys=True),)
class MAIConceptInjectDelta:
"""Add a stored delta to the target-video rows at chosen blocks/steps."""
DESCRIPTION = (
"EXPERIMENTAL (alpha), new 2026-08-15. The mirror of the capture tap: "
"loads a delta pack (a .safetensors of per-(block, step) frame_mean "
"maps, optional patch_mean maps and optional full block states, plus "
"its sidecar JSON) and ADDS it, times alpha, to the target-video rows "
"of the block output, on a CLONE of the model.\n\n"
"mode='full' adds the whole stored state instead of a reduction, so "
"the thing injected is the thing a capture measured. Those tensors "
"are big (~160 MB per position) and are read from the pack one at a "
"time, on first use, straight into the model's device and dtype.\n\n"
"It answers one question and holds no theory: does a delta measured "
"at the tap do anything when it is put back? Run it against its own "
"controls. control='zero' runs every line of the patch with zeroes "
"(if that render differs from the unpatched one, the plumbing is the "
"effect); control='time_shuffle' keeps the delta's energy and "
"destroys its timing (if that renders the same as 'none', the timing "
"is not what is doing the work); control='space_shuffle' does the "
"same to WHERE the delta sits inside a frame.\n\n"
"The pack's (latent_t, frame_rows) must match the run's: the node "
"refuses before the first block otherwise. A different prompt or a "
"different audio length is fine, the delta lives on video rows.")
@classmethod
def INPUT_TYPES(cls):
return {"required": {
"model": ("MODEL", {"tooltip": "a patched CLONE comes out; the "
"input is untouched"}),
"pack_path": ("STRING", {"default": "",
"tooltip": "absolute path to the pack "
".safetensors; its sidecar <pack>.json "
"(or <pack>.safetensors.json) must sit "
"beside it"}),
"blocks": ("STRING", {"default": "",
"tooltip": "comma-separated DiT blocks to inject at; "
"empty = every block the pack carries. Must "
"not collide with a tap's blocks on the "
"same model"}),
"alpha": ("FLOAT", {"default": 1.0, "min": -8.0, "max": 8.0,
"step": 0.05,
"tooltip": "scale on the delta; negative subtracts"}),
"step_mode": (list(h3_inject.STEP_MODES),
{"default": "captured_only",
"tooltip": "captured_only fires only on the steps "
"the pack was captured at; nearest_all "
"fires every step using the nearest "
"captured entry"}),
"mode": (list(h3_inject.MODES), {"default": "frame",
"tooltip": "frame = one vector per latent frame; "
"separable = frame + patch - grand mean, the "
"rank-2 map (needs patch_mean in the pack); "
"full = the whole stored block state, row for "
"row (needs full states in the pack, ~160 MB "
"each, read one at a time)"}),
"control": (list(h3_inject.CONTROLS), {"default": "none",
"tooltip": "the arm: none / zero (plumbing control) / "
"time_shuffle (same energy, wrong "
"timing) / "
"space_shuffle (same energy, wrong place; "
"refused in mode 'frame', which has no "
"space axis) / sign_flip"}),
"control_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff,
"tooltip": "seeds the shuffle permutations, on "
"the CPU, so an arm is re-runnable "
"from its label"}),
}, "optional": {
"cond_only": ("BOOLEAN", {"default": True,
"tooltip": "inject only on the conditional forward; "
"CFG makes two per step and they are two "
"different conditions"}),
"enabled": ("BOOLEAN", {"default": True,
"tooltip": "False installs the patches and adds "
"NOTHING: the pass-through control"}),
}}
RETURN_TYPES = ("MODEL",)
RETURN_NAMES = ("model",)
FUNCTION = "inject"
CATEGORY = CATEGORY
@classmethod
def IS_CHANGED(cls, **kwargs):
# same reason as the Arm node: an injection is an experiment, and a
# cached experiment is not one
return float("nan")
def inject(self, model, pack_path, blocks, alpha, step_mode, mode,
control, control_seed, cond_only=True, enabled=True):
pack = h3_inject.DeltaPack.load(pack_path, control=control,
control_seed=int(control_seed))
session = h3_inject.InjectSession(
pack=pack,
blocks=_ints(blocks, "blocks", "MAI Concept Inject Delta"),
alpha=float(alpha),
step_mode=step_mode, mode=mode, cond_only=bool(cond_only),
enabled=bool(enabled))
return (session.install(model),)
NODE_CLASS_MAPPINGS = {
"MAIConceptCaptureArm": MAIConceptCaptureArm,
"MAIConceptCaptureFlush": MAIConceptCaptureFlush,
"MAIConceptInjectDelta": MAIConceptInjectDelta,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MAIConceptCaptureArm": "MAI Concept Capture Arm (alpha)",
"MAIConceptCaptureFlush": "MAI Concept Capture Flush (alpha)",
"MAIConceptInjectDelta": "MAI Concept Inject Delta (alpha)",
}
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|