From d6981036d0056a0c482eec50d07ca99e3488a621 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:01:21 +0300 Subject: [PATCH] feat: per-block influence probe for H3 (measurement only) Records ||f_i(x_i)|| / ||x_i|| on video rows for every block, averaged over each run's steps, and pairs the profile with that run's rating. Answers the question that gates rating-driven block weighting: is H3's depth profile flat (nothing to aim at) or structured? Reports a scale-invariant flatness number, and once a key has 2+ liked and 2+ disliked runs, which blocks run hotter on the liked ones. Steers nothing. Installed BEFORE any steering hook so the delta it records is the block's own, not one already containing an injection. Bounded cost: two norms per block per step on ~512 strided rows, no extra forward pass, results accumulated as device tensors and synced once per run rather than 50x per step. Says so explicitly when it records nothing, since an empty capture and a flat profile are otherwise indistinguishable downstream. Adjacent to per-block residual steering, removed in b7ed4f9 for finding no rating signal above the 4th decimal -- that pass had no per-step resolution, which its own removal note named as the gap. This either reproduces that null definitively or it doesn't. --- block_influence.py | 194 +++++++++++++++ conditioning.py | 24 ++ movie_editor/backend/settings_card.py | 1 + movie_editor/frontend/engine_settings.js | 6 +- samplers.py | 104 +++++++- tests/test_block_influence.py | 297 +++++++++++++++++++++++ 6 files changed, 623 insertions(+), 3 deletions(-) create mode 100644 block_influence.py create mode 100644 tests/test_block_influence.py diff --git a/block_influence.py b/block_influence.py new file mode 100644 index 00000000..dbc25336 --- /dev/null +++ b/block_influence.py @@ -0,0 +1,194 @@ +"""Per-block influence probe for MiniMax H3. + +Answers one thing nothing else here measures: how much does each of H3's blocks actually +move the VIDEO rows of the hidden stream, and does that profile differ between generations +you rated liked and ones you rated disliked. + +Each block of a residual transformer contributes `x_{i+1} = x_i + f_i(x_i)`. This records +||f_i(x_i)|| / ||x_i|| on video rows only, averaged over every sampling step of a run. Free +in the sense that matters: two norms per block per step, no extra forward pass, no second +model call. Row-subsampled (see samplers._install_block_influence) so the cost does not +scale with H3's 17k+ sequence length. + +WHAT THIS PROVES, AND WHAT IT DOES NOT. A large residual delta means the block moved the +stream a lot at that point -- NOT that the movement survived to the final video. The blocks +behind it, and the final layer, can normalise an early push away entirely, and a small late +delta can reach the output untouched. The causal version of this question is ablation: zero +block i, re-run, measure how far the output latent moved. That is ~50 single-step forwards +per probe -- an offline calibration pass, cached per model+LoRA stack (a LoRA changes the +weights, so it changes the map), not something to run per generation. This module is the +free correlational map, and its job is to answer the question that gates the expensive one: +is the depth profile flat, or is it structured? A flat profile means rating-driven per-block +weighting has nothing to grip. + +PRIOR ART IN THIS REPO, read before trusting a positive result: per-block residual steering +from ratings was built and REMOVED (b7ed4f9) -- block activity was found to vary in the 4th +decimal with no extractable rating signal. That attempt had no per-step resolution, which +its own removal note named as the missing piece; this one keeps every step separate before +averaging. So this either reproduces that null definitively, or it doesn't. Both outcomes +are worth having, and a null here is the cheaper answer. + +H3 is 50 dense uniform blocks with no MoE and no router (see the H3 intervention map), so +there is nothing to REROUTE -- the only levers a block exposes are skip, scale, and add. +This measures which of those would be worth aiming at. +""" +from __future__ import annotations + +import os + +import torch + +# Match h3_repr_steering's floor: a liked-vs-disliked profile difference needs 2+ rated runs +# on each side before it is a difference rather than a pair of points. +MIN_PER_GROUP = 2 + + +def _log(): + try: + from . import funpack_log as fl + except ImportError: + import funpack_log as fl + return fl + + +def state_path(refinement_key): + try: + from .conditioning import refinement_state_path + except ImportError: + from conditioning import refinement_state_path + return refinement_state_path(refinement_key, "block_influence", prefix="refine_v2", + extension="pt") + + +def _load(refinement_key): + path = state_path(refinement_key) + if not os.path.exists(path): + return {"rows": [], "pending": None} + try: + data = torch.load(path, map_location="cpu", weights_only=True) + except Exception: # noqa: BLE001 + return {"rows": [], "pending": None} + if not isinstance(data, dict): + return {"rows": [], "pending": None} + data.setdefault("rows", []) + data.setdefault("pending", None) + return data + + +def _save(refinement_key, data): + path = state_path(refinement_key) + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = path + ".tmp" + torch.save(data, tmp) + os.replace(tmp, path) + + +def save_pending(refinement_key, profile): + """Store this run's per-block profile as the pending candidate the next rating scores. + `profile` is {block_index: mean relative delta}. Overwritten every run, same reason + h3_repr_steering.save_pending overwrites: a pending capture must never survive to be + paired with a LATER run's rating.""" + if not refinement_key or not profile: + return + data = _load(refinement_key) + data["pending"] = {int(b): float(v) for b, v in profile.items()} + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 block influence", "save pending", e, "this run's profile is not kept") + + +def commit(refinement_key, reward): + """Pair the pending profile with the rating that scored its run. + -> "recorded" | "no_pending" | "no_key", same contract h3_repr_steering.commit uses so + the caller can tell a real commit from a reward that had nothing to attach to.""" + if not refinement_key: + return "no_key" + data = _load(refinement_key) + pending = data.get("pending") + data["pending"] = None + if pending is None: + try: + _save(refinement_key, data) + except OSError: + pass + return "no_pending" + try: + data["rows"].append({"profile": pending, "weight": float(reward)}) + except (TypeError, ValueError): + try: + _save(refinement_key, data) + except OSError: + pass + return "no_pending" + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 block influence", "commit", e, "rating not recorded") + return "no_pending" + return "recorded" + + +def clear_all(refinement_key): + if not refinement_key: + return + try: + os.remove(state_path(refinement_key)) + except FileNotFoundError: + pass + + +def _blocks_in(rows): + blocks = set() + for r in rows: + prof = r.get("profile") + if isinstance(prof, dict): + blocks.update(int(b) for b in prof) + return sorted(blocks) + + +def profile(refinement_key): + """-> {"overall": {block: mean}, "liked": {...}|None, "disliked": {...}|None, + "difference": {block: liked-disliked}|None, "n_liked": int, "n_disliked": int, + "flatness": float|None} + + `overall` is the depth profile itself -- where this model does its work, regardless of + ratings. `difference` is the rating-driven signal: blocks whose activity runs higher on + runs you liked. `flatness` is the spread of `overall` as a fraction of its own mean + (coefficient of variation): near 0 means every block moves the stream by the same amount + and there is nothing to aim at; a large value means the depth profile is structured. + + Rows missing a block are skipped for that block rather than counted as zero -- a run that + hooked fewer blocks must not drag a block's mean down.""" + data = _load(refinement_key) + rows = [r for r in data["rows"] + if isinstance(r, dict) and isinstance(r.get("profile"), dict) and "weight" in r] + blocks = _blocks_in(rows) + if not rows or not blocks: + return {"overall": {}, "liked": None, "disliked": None, "difference": None, + "n_liked": 0, "n_disliked": 0, "flatness": None} + + def _mean(subset): + out = {} + for b in blocks: + vals = [float(r["profile"][b]) for r in subset if b in r["profile"]] + if vals: + out[b] = sum(vals) / len(vals) + return out + + liked_rows = [r for r in rows if r["weight"] > 0] + disliked_rows = [r for r in rows if r["weight"] < 0] + overall = _mean(rows) + flatness = None + if overall: + vals = torch.tensor(list(overall.values()), dtype=torch.float32) + m = float(vals.mean()) + flatness = float(vals.std(unbiased=False) / m) if abs(m) > 1e-12 else None + + liked = disliked = difference = None + if len(liked_rows) >= MIN_PER_GROUP and len(disliked_rows) >= MIN_PER_GROUP: + liked, disliked = _mean(liked_rows), _mean(disliked_rows) + difference = {b: liked[b] - disliked[b] for b in liked if b in disliked} + return {"overall": overall, "liked": liked, "disliked": disliked, + "difference": difference, "n_liked": len(liked_rows), + "n_disliked": len(disliked_rows), "flatness": flatness} diff --git a/conditioning.py b/conditioning.py index f759442d..1e6aeba0 100644 --- a/conditioning.py +++ b/conditioning.py @@ -10216,6 +10216,30 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): }.get(_rs_outcome, _rs_outcome) print(f"[FunPackRefiner] H3 representation steering: weight " f"{_rs_reward:+.3f} for rating '{rating_label}' — {_rs_note}") + # Block-influence probe rides the same rating. Measurement only -- it + # steers nothing, so it is committed unconditionally alongside REINS + # rather than behind its own admissibility gate, and stays silent unless + # the sampler actually captured a profile this run. + try: + from . import block_influence as _bi + except ImportError: + import block_influence as _bi + if _bi.commit(refinement_key, _rs_reward) == "recorded": + _bi_prof = _bi.profile(refinement_key) + _bi_flat = _bi_prof.get("flatness") + _bi_msg = (f"flatness {_bi_flat:.3f}" if _bi_flat is not None + else "flatness n/a") + _bi_diff = _bi_prof.get("difference") + if _bi_diff: + _bi_top = sorted(_bi_diff.items(), key=lambda kv: -abs(kv[1]))[:3] + _bi_msg += " | liked-vs-disliked top: " + ", ".join( + f"block {b} {v:+.4f}" for b, v in _bi_top) + else: + _bi_msg += (f" | no liked/disliked split yet " + f"({_bi_prof.get('n_liked', 0)} liked / " + f"{_bi_prof.get('n_disliked', 0)} disliked, need " + f"{_bi.MIN_PER_GROUP} of each)") + print(f"[FunPackRefiner] H3 block influence: recorded — {_bi_msg}") # DynaShift latent memory: pair the sampler's pending raw latent with this rating # — promote into the negative bank on a bad outcome (intrusions: awful / # wrong_appearance; or the quality-missing family via reward <= threshold — see diff --git a/movie_editor/backend/settings_card.py b/movie_editor/backend/settings_card.py index 0f246a94..58b938da 100644 --- a/movie_editor/backend/settings_card.py +++ b/movie_editor/backend/settings_card.py @@ -267,6 +267,7 @@ _EDITOR_DEFAULTS = { "frame_overlap": 16, "h3_phrase_emphasis": False, "h3_phrase_variability": 0.0, + "h3_block_influence": False, "h3_repr_steering": False, "h3_repr_steering_block": "25", "h3_repr_steering_strength": 0.05, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 8ba09277..ee6dd6a8 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -24,6 +24,7 @@ "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", + "h3_block_influence", ]); const RATING_GATED_STUDIO = new Set(["reference_injection", "value_guidance", "steer_mode", "absolute_strength"]); @@ -319,6 +320,9 @@ hint: "Fraction of the block's own hidden-state norm added along the learned direction. Start low and push it deliberately high on a same-seed test to confirm it moves the video at all before trusting a small value." }, { name: "h3_repr_steering_block", label: "Steer at block", kind: "combo", choices: ["1", "5", "10", "15", "20", "25", "30", "35", "40"], default: "25", dependsOn: "h3_repr_steering", hint: "TEST-ONLY: which of the block-sweep's candidate blocks actually gets steered, instead of always block 25. Blocks 1-15 are deep-end territory (the anchor/reference rows haven't been read by the network yet at that depth) -- expect this to behave differently, not just 'weaker', than 20-40. The other candidate blocks are still captured read-only either way." }, + { name: "h3_block_influence", label: "Measure block influence (H3)", kind: "bool", default: false, + hint: "Measures only — steers nothing. Records how much each of H3's 50 blocks moves the picture as it forms, and pairs that with the rating you give the run.", + detail: "Answers whether the depth profile is flat (every block contributes about the same, so weighting blocks by rating has nothing to grip) or structured. Reports a flatness number and, once you have 2+ liked and 2+ disliked runs on the key, which blocks run hotter on the ones you liked. A block moving the stream a lot at that depth is not proof it survived to the final video — that needs per-block ablation, an offline pass. Near-free: two norms per block per step, no extra model pass." }, { name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false, hint: "Feeds the identity pin image the way Best-FaceID identity LoRAs expect it. Needs an Identity pin set.", detail: "Replaces Continuity's identity-pin guide with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself in Models. No effect without an identity pin." }, @@ -508,7 +512,7 @@ const CHAIN_VIEW_KNOBS = { chain_continuity: ["carry_i2v_guides", "carry_overlap_through_anchor"], chain_timing: ["frame_overlap", "transition_duration", "use_same_seed", "cut_opening_frames"], - chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block"], + chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", "h3_block_influence"], chain_decode: ["decode_noise_scale", "decode_timestep", "decode_tile_size"], chain_experimental: ["context_windows", "context_window_length", "context_window_overlap", "context_window_schedule", "context_window_fuse", "context_window_freenoise", "context_window_retain_first", "h3_video_detail", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], }; diff --git a/samplers.py b/samplers.py index ddf64bba..583f83dc 100644 --- a/samplers.py +++ b/samplers.py @@ -2970,6 +2970,10 @@ class FunPackLTXAVSceneChainSampler: "default": "25", "tooltip": "TEST-ONLY: which block actually gets steered (the rest of the sweep's candidate blocks are still captured read-only, same as always). direction() is recomputed for whichever block you pick here -- this does not change what block_sweep ranks, only which block's learned direction gets injected into this run.", }), + "h3_block_influence": ("BOOLEAN", { + "default": False, + "tooltip": "MEASUREMENT ONLY -- steers nothing. Records how much each of H3's 50 blocks moves the video rows of the hidden stream (its residual delta relative to the stream itself), averaged over every step, and pairs that profile with the rating you give the run. Answers the question that gates rating-driven block weighting: is the depth profile flat (every block contributes the same, nothing to aim at) or structured? A large residual delta means the block moved the stream a lot THERE -- not that it survived to the final video; the causal version of that needs per-block ablation. Near-free: two norms per block per step on ~512 strided rows, no extra forward pass.", + }), # A connection socket, never a widget — safe at the end, and it must stay after # every widget above (see the widgets_values note at the top of this block). "second_pass_sigmas": ("SIGMAS", { @@ -3678,7 +3682,7 @@ class FunPackLTXAVSceneChainSampler: alg_anchor_sigma_threshold=0.975, bounded_attention_enabled=False, h3_repr_steering=False, h3_repr_steering_strength=0.05, - h3_repr_steering_block="25", + h3_repr_steering_block="25", h3_block_influence=False, refinement_key="",): if sampler is None: raise ValueError("sampler input is required.") @@ -3741,6 +3745,12 @@ class FunPackLTXAVSceneChainSampler: # key, no rating and no Studio state, so it works on a graph with the Refiner absent # entirely. 1.0 on all three does not even clone the model. model = self._install_h3_final_layer(model, positive) + # Measurement first, steering second: whatever is installed later chains this probe + # as its inner call, so the delta it records is the block's OWN, not one that already + # contains another mechanism's injection. + _influence_capture = [{}] + if h3_block_influence and refinement_key: + model = self._install_block_influence(model, _influence_capture) _repr_capture = [{}] if h3_repr_steering and refinement_key: model = self._install_h3_repr_steering( @@ -3759,6 +3769,23 @@ class FunPackLTXAVSceneChainSampler: except ImportError: import h3_repr_steering as _rs _rs.save_pending(refinement_key, _repr_capture[0]) + if _influence_capture[0]: + try: + from . import block_influence as _bi + except ImportError: + import block_influence as _bi + # One host sync for the whole run, here, instead of per block per step. + _bi.save_pending(refinement_key, { + b: float(torch.stack(v).mean().item()) + for b, v in _influence_capture[0].items() if v}) + elif h3_block_influence and refinement_key: + # An empty capture and a genuinely flat profile look identical downstream, + # and this probe exists to tell those apart -- so it says which one happened + # rather than letting a dead hook read as a finding. + _log.failed("FunPackSceneChain", "H3 block influence probe", + "no block recorded a video-row delta", + "this run contributes no profile (the hooks never saw video rows " + "— not a flat depth profile)") finally: self._remove_bounded_attention(_ba_handles) self._assert_finite_sample(sampled) @@ -6732,6 +6759,78 @@ class FunPackLTXAVSceneChainSampler: "not applied and not capturing this run") return model + def _install_block_influence(self, model, capture_holder, n_blocks=50, max_rows=512): + """MEASUREMENT ONLY (H3). See block_influence.py. + + Records, per block, ||f_i(x_i)|| / ||x_i|| on the VIDEO rows -- how much that block + moved the stream, relative to what the stream already was. Nothing is modified: the + hook returns the block's own output untouched, and a failure inside it is swallowed + rather than allowed to break sampling, because a probe must never be able to cost you + a generation. + + Cost control: video rows are STRIDED down to ~`max_rows` before the two norms, so the + reduction is bounded regardless of H3's 17k+ sequence length, and the per-block + results accumulate as 0-dim CUDA tensors (never .item()) so there is no host sync per + block per step. Install this BEFORE any steering hook -- whichever wrapper is + installed later chains this one as its inner call, so measuring first is what keeps + the number the block's OWN delta instead of one that includes another mechanism's + injection (the same capture-before-injection ordering h3_repr_steering needed). + + `n_blocks` over-covers deliberately: registering a patch for an index the model does + not have is never called, so a shorter stack simply reports fewer blocks.""" + try: + try: + from . import h3_repr_steering as _rs + except ImportError: + import h3_repr_steering as _rs + patched = model.clone() + to = patched.model_options.get("transformer_options", {}).copy() + patches_replace = dict(to.get("patches_replace", {})) + dit_patches = dict(patches_replace.get("dit", {})) + + def _make_probe(block): + inner = dit_patches.get(("double_block", block)) + mask_cache = {} + + def _probe(args, extra): + src = args.get("img") + out = extra["original_block"](args)["img"] if inner is None else \ + inner(args, extra)["img"] + try: + seq_len = int(out.shape[0]) + mask = mask_cache.get(seq_len, "MISS") + if mask == "MISS": + mask = _rs.video_mask_from_mod_segments( + args.get("mod_segments"), seq_len, out.device) + mask_cache[seq_len] = mask + if mask is not None and torch.is_tensor(src) and src.shape == out.shape: + idx = mask.nonzero(as_tuple=True)[0] + if idx.numel(): + stride = max(1, int(idx.numel()) // int(max_rows)) + sel = idx[::stride] + before = src[sel].detach().float() + delta = out[sel].detach().float() - before + base = before.norm() + # Kept as a 0-dim tensor: averaged and read once, after + # sampling, instead of syncing the device 50x per step. + capture_holder[0].setdefault(block, []).append( + delta.norm() / base.clamp(min=1e-8)) + except Exception: # noqa: BLE001 + pass # a probe never breaks a generation + return {"img": out} + return _probe + + for _b in range(int(n_blocks)): + dit_patches[("double_block", _b)] = _make_probe(_b) + patches_replace["dit"] = dit_patches + to["patches_replace"] = patches_replace + patched.model_options["transformer_options"] = to + return patched + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackSceneChain", "H3 block influence probe", _e, + "not measuring this run") + return model + def _install_bounded_attention(self, model, latent, positive): """EXPERIMENTAL (arXiv:2403.16990-inspired, see [[project_bounded_attention]]): mask the text cross-attention (attn2 only — never audio_attn2, never self-attention) so the LEFT @@ -6990,7 +7089,7 @@ class FunPackLTXAVSceneChainSampler: trajectory_guidance=False, trajectory_guidance_strength=0.02, dynashift=False, dynashift_strength=0.3, dynashift_threshold=0.6, h3_repr_steering=False, h3_repr_steering_strength=0.05, - h3_repr_steering_block="25", + h3_repr_steering_block="25", h3_block_influence=False, alg_guide_blur_strength=2.0, alg_guide_blur_sigma_threshold=0.975, alg_anchor=False, alg_anchor_strength=2.0, alg_anchor_sigma_threshold=0.975, identity_transfer_enabled=False, identity_projector="None", source_id=2.0, @@ -7847,6 +7946,7 @@ class FunPackLTXAVSceneChainSampler: h3_repr_steering=h3_repr_steering, h3_repr_steering_strength=h3_repr_steering_strength, h3_repr_steering_block=h3_repr_steering_block, + h3_block_influence=h3_block_influence, refinement_key=refinement_key_input, ) # cut_opening_frames: let the real, untouched i2v anchor condition the scene diff --git a/tests/test_block_influence.py b/tests/test_block_influence.py new file mode 100644 index 00000000..23ab13e8 --- /dev/null +++ b/tests/test_block_influence.py @@ -0,0 +1,297 @@ +"""Per-block influence probe: the pending->rating pairing, the profile maths, and the +flatness readout that decides whether rating-driven block weighting is worth pursuing.""" +import sys + +import pytest +import torch + +sys.path.insert(0, ".") +import block_influence as bi # noqa: E402 + + +@pytest.fixture(autouse=True) +def _isolate_state(monkeypatch, tmp_path): + monkeypatch.setattr(bi, "state_path", lambda key: str(tmp_path / f"{key}.pt")) + + +def _run(key, profile, reward): + bi.save_pending(key, profile) + return bi.commit(key, reward) + + +# --- pending / commit contract ------------------------------------------------- + +def test_commit_reports_recorded_on_a_real_pairing(): + assert _run("k", {0: 0.1, 1: 0.2}, 1.0) == "recorded" + + +def test_commit_reports_no_pending_with_nothing_to_pair(): + assert bi.commit("k", 1.0) == "no_pending" + + +def test_commit_reports_no_key(): + assert bi.commit("", 1.0) == "no_key" + + +def test_pending_overwrites_rather_than_accumulating(): + bi.save_pending("k", {0: 0.1}) + bi.save_pending("k", {0: 0.9}) # a second run before any rating + bi.commit("k", 1.0) + rows = bi._load("k")["rows"] + assert len(rows) == 1 and rows[0]["profile"][0] == 0.9 + + +def test_clear_all_removes_the_state(): + _run("k", {0: 0.1}, 1.0) + bi.clear_all("k") + assert bi.profile("k")["overall"] == {} + + +def test_no_key_is_a_safe_no_op(): + bi.save_pending("", {0: 0.1}) + bi.commit("", 1.0) + bi.clear_all("") + + +# --- profile maths ------------------------------------------------------------- + +def test_overall_is_the_mean_across_runs_regardless_of_rating(): + _run("k", {0: 0.2, 1: 0.4}, 1.0) + _run("k", {0: 0.4, 1: 0.6}, -1.0) + got = bi.profile("k")["overall"] + assert got[0] == pytest.approx(0.3) + assert got[1] == pytest.approx(0.5) + + +def test_difference_needs_min_per_group_on_each_side(): + for _ in range(bi.MIN_PER_GROUP): + _run("k", {0: 0.5}, 1.0) + _run("k", {0: 0.1}, -1.0) # only one disliked + p = bi.profile("k") + assert p["difference"] is None + assert p["n_liked"] == bi.MIN_PER_GROUP and p["n_disliked"] == 1 + + +def test_difference_is_liked_mean_minus_disliked_mean(): + for _ in range(bi.MIN_PER_GROUP): + _run("k", {0: 0.8, 1: 0.2}, 1.0) + _run("k", {0: 0.4, 1: 0.2}, -1.0) + diff = bi.profile("k")["difference"] + assert diff[0] == pytest.approx(0.4) # block 0 runs hotter on liked runs + assert diff[1] == pytest.approx(0.0) # block 1 is indifferent to the rating + + +def test_a_block_missing_from_a_row_is_skipped_not_counted_as_zero(): + """A run that hooked fewer blocks must not drag the absent block's mean toward 0 -- + the same rule h3_repr_steering.direction() follows for a missing descriptor.""" + _run("k", {0: 0.4, 1: 0.4}, 1.0) + _run("k", {0: 0.6}, 1.0) # no block 1 this run + got = bi.profile("k")["overall"] + assert got[0] == pytest.approx(0.5) + assert got[1] == pytest.approx(0.4) + + +def test_neutral_weight_rows_count_toward_overall_but_neither_group(): + _run("k", {0: 1.0}, 0.0) + p = bi.profile("k") + assert p["n_liked"] == 0 and p["n_disliked"] == 0 + assert p["overall"][0] == pytest.approx(1.0) + + +# --- flatness: the number the whole probe exists to produce --------------------- + +def test_a_perfectly_flat_profile_reports_zero_flatness(): + """Every block moving the stream by the same amount is the null result that says + rating-driven block weighting has nothing to grip.""" + _run("k", {b: 0.5 for b in range(10)}, 1.0) + assert bi.profile("k")["flatness"] == pytest.approx(0.0) + + +def test_a_structured_profile_reports_nonzero_flatness(): + _run("k", {0: 0.1, 1: 0.5, 2: 1.5}, 1.0) + assert bi.profile("k")["flatness"] > 0.5 + + +def test_flatness_is_scale_invariant(): + """Coefficient of variation, not raw spread: doubling every block's delta describes the + same shape and must not read as twice as structured.""" + _run("a", {0: 0.1, 1: 0.2, 2: 0.4}, 1.0) + _run("b", {0: 0.2, 1: 0.4, 2: 0.8}, 1.0) + assert bi.profile("a")["flatness"] == pytest.approx(bi.profile("b")["flatness"]) + + +def test_flatness_is_none_with_no_data(): + assert bi.profile("k")["flatness"] is None + + +def test_corrupt_state_file_degrades_to_empty(tmp_path, monkeypatch): + path = tmp_path / "junk.pt" + path.write_bytes(b"not a torch file") + monkeypatch.setattr(bi, "state_path", lambda key: str(path)) + assert bi.profile("k")["overall"] == {} + assert bi.commit("k", 1.0) == "no_pending" + + +def test_rows_survive_a_reload_as_plain_floats(): + """weights_only=True load: the state must never depend on unpickling a custom class.""" + _run("k", {0: 0.25}, 1.0) + rows = bi._load("k")["rows"] + assert isinstance(rows[0]["profile"][0], float) + assert isinstance(rows[0]["weight"], float) + assert not isinstance(rows[0]["profile"][0], torch.Tensor) + + +# --- the sampler hook itself --------------------------------------------------- +# +# The module above is only half the probe. The half that can silently record nothing is the +# hook: if the video mask misses, the input tensor is not where it is expected, or the +# swallow-everything guard eats a real error, the profile comes back empty and looks like +# "this model has no block structure" rather than "the probe never ran". These exercise the +# registered hook directly, the way the model calls it. + +import types # noqa: E402 +from pathlib import Path # noqa: E402 + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +for _name in ( + "comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling", + "comfy.model_sampling", "comfy.nested_tensor", "comfy.sample", + "comfy.samplers", "comfy.utils", +): + sys.modules.setdefault(_name, types.ModuleType(_name)) +sys.modules["comfy.nested_tensor"].NestedTensor = object + +import samplers # noqa: E402 + + +class _Model: + def __init__(self): + self.model_options = {} + + def clone(self): + m = _Model() + m.model_options = {k: dict(v) if isinstance(v, dict) else v + for k, v in self.model_options.items()} + return m + + +def _install(capture, n_blocks=2, max_rows=512): + s = samplers.FunPackLTXAVSceneChainSampler() + patched = s._install_block_influence(_Model(), capture, n_blocks=n_blocks, + max_rows=max_rows) + return patched.model_options["transformer_options"]["patches_replace"]["dit"] + + +def _call(hook, src, out, mod_segments): + """Drive the hook the way minimax/model.py does.""" + return hook({"img": src, "mod_segments": mod_segments}, + {"original_block": lambda a: {"img": out}}) + + +def test_hook_records_the_relative_delta_on_video_rows(): + capture = [{}] + dit = _install(capture) + src = torch.zeros(4, 8) + src[:, 0] = 3.0 # ||video rows|| = 6.0 over 4 rows + out = src.clone() + out[:, 1] = 3.0 # delta has the same norm as the input + res = _call(dit[("double_block", 0)], src, out, [(0, 4, 6)]) # tag 6 % 3 == 0 -> video + assert torch.equal(res["img"], out) # measurement must not modify the stream + recorded = torch.stack(capture[0][0]).mean().item() + assert recorded == pytest.approx(1.0, rel=1e-4) + + +def test_hook_averages_across_steps(): + capture = [{}] + dit = _install(capture) + src = torch.zeros(4, 8) + src[:, 0] = 1.0 + for scale in (1.0, 3.0): # two "steps" with different deltas + out = src.clone() + out[:, 1] = scale + _call(dit[("double_block", 0)], src, out, [(0, 4, 6)]) + assert torch.stack(capture[0][0]).mean().item() == pytest.approx(2.0, rel=1e-4) + + +def test_text_and_audio_rows_are_not_measured(): + """Only video rows count -- a huge delta on text rows must not register.""" + capture = [{}] + dit = _install(capture) + src = torch.zeros(4, 8) + src[:, 0] = 1.0 + out = src.clone() + out[2:, 1] = 50.0 # rows 2-3 are text below + _call(dit[("double_block", 0)], src, out, [(0, 2, 6), (2, 4, 7)]) + recorded = torch.stack(capture[0][0]).mean().item() + assert recorded == pytest.approx(0.0, abs=1e-6) + + +def test_no_video_segment_records_nothing_rather_than_zero(): + capture = [{}] + dit = _install(capture) + src = torch.ones(4, 8) + _call(dit[("double_block", 0)], src, src * 2, [(0, 4, 7)]) # text only + assert capture[0] == {} + + +def test_each_block_accumulates_separately(): + capture = [{}] + dit = _install(capture, n_blocks=3) + src = torch.zeros(4, 8) + src[:, 0] = 1.0 + for block, scale in ((0, 1.0), (2, 4.0)): + out = src.clone() + out[:, 1] = scale + _call(dit[("double_block", block)], src, out, [(0, 4, 6)]) + assert set(capture[0]) == {0, 2} # block 1 was never called + assert torch.stack(capture[0][0]).mean().item() == pytest.approx(1.0, rel=1e-4) + assert torch.stack(capture[0][2]).mean().item() == pytest.approx(4.0, rel=1e-4) + + +def test_row_subsampling_bounds_the_reduction(): + """The stride keeps the cost flat as the sequence grows; the estimate must stay close.""" + capture = [{}] + dit = _install(capture, max_rows=16) + n = 4096 + src = torch.zeros(n, 8) + src[:, 0] = 1.0 + out = src.clone() + out[:, 1] = 2.0 + _call(dit[("double_block", 0)], src, out, [(0, n, 6)]) + assert torch.stack(capture[0][0]).mean().item() == pytest.approx(2.0, rel=1e-3) + + +def test_an_existing_patch_at_the_same_block_is_chained_not_replaced(): + """Installing over another mechanism's hook must call it, not drop it.""" + calls = [] + + class _M(_Model): + pass + + s = samplers.FunPackLTXAVSceneChainSampler() + base = _M() + + def _inner(args, extra): + calls.append(True) + return {"img": args["img"] * 2} + + base.model_options["transformer_options"] = { + "patches_replace": {"dit": {("double_block", 0): _inner}}} + capture = [{}] + patched = s._install_block_influence(base, capture, n_blocks=1) + hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][ + ("double_block", 0)] + src = torch.ones(4, 8) + res = _call(hook, src, src, [(0, 4, 6)]) + assert calls == [True] # inner ran + assert torch.equal(res["img"], src * 2) # and its output is what propagates + + +def test_a_broken_hook_never_breaks_sampling(): + capture = [{}] + dit = _install(capture) + src = torch.ones(4, 8) + # mod_segments of the wrong shape entirely: must pass the block output through anyway. + res = _call(dit[("double_block", 0)], src, src * 3, "not-segments") + assert torch.equal(res["img"], src * 3) + assert capture[0] == {}