olivv-cs-ComfyUI-FunPack/h3_repr_steering.py

"""REINS-style representation steering for MiniMax H3.

Everything else this project has built for ratings acts OUTSIDE the model: an attention
bias on the prompt's own tokens (h3_token_weights), or a push on the latent between denoiser
calls (output_guidance, trajectory_guidance). Neither changes what the model itself computes.

This does. It captures the video rows of one block's hidden state per generation, keeps a
running direction from how each generation was rated, and adds that direction back into the
same block's output on the next run -- reaching inside the forward pass instead of around it.
Training-free: no gradient, no MLP, a mean difference of vectors is the whole "model".
("Pulling The REINS", arXiv:2606.17257 -- same idea, applied here via a two-group mean
difference instead of their supervised-PCA.)

Liked vs disliked by the rating's SIGN only, not weighted by magnitude. A magnitude-weighted
version (a 9/10 pulling harder than a 6/10) was tried first and is the statistically better
estimator in the limit -- but it needs enough data for a weakly-rated row's noise to average
out, and at the handful of ratings a key actually has, a mediocre "6" just gets a vote and
dilutes the direction instead of being excluded by it. Binary spends what little data exists
on the clearest examples only; see direction()'s docstring and the 2026-09-04 session before
reintroducing weighting. The weight is still the rating's own `reward` -- already
admissibility-filtered by the caller (conditioning.py only commits inside the same
`_v2_reward_admissible` gate the value functions use), so a Wrong-* identity-mismatch rating
never reaches here at all, not because its number was excluded but because it was never a
quality verdict to begin with.

Masked to VIDEO rows only, using the model's own `mod_segments` tags (0=video, 1=text,
2=audio) -- never audio or text rows. The one documented failure of this class of
intervention (`project_self_consistency_failed.md`) happened on a DIFFERENT model (LTXAV)
by overwriting rows fed back through joint audio-video attention; H3 has no separate
audio-attention path to spare, so touching audio rows here risks the same corruption for a
worse reason. This is the mitigation, not a guarantee -- unvalidated until it's actually run.
"""
from __future__ import annotations

import os

import torch

# Roughly half of H3's 50 blocks -- REINS reports this FRACTION as their sweet spot (their
# own notation: block l in {0..L-1}, effect peaking ~50% of L), on whatever model they tested
# -- not H3, and their own paper picks the actual best layer per model by sweep rather than
# trusting 50% universally. This is that same sweep, run on H3 with real rated data instead
# of guessed from someone else's percentage. Only DEFAULT_BLOCK ever STEERS; the rest of
# CANDIDATE_BLOCKS are captured read-only alongside it so block_sweep() has something to rank.
DEFAULT_BLOCK = 25
# Every block, 0-49: the coarse grid + 41-49 individually was a compromise when only the
# tail's back-loaded weight was known to be worth resolving finely (48+49 alone carrying ~a
# third of all residual movement). 2026-09-04: block 0/1 turned out to be structurally
# distinct too (raw input projection / patch-mixing), not just quiet early blocks -- so
# "coarse in the middle" was an assumption, not a finding. Capturing every block read-only
# costs one extra mean-pool per block per run; steering still only ever touches one block at
# a time (whichever `steer_block` names).
CANDIDATE_BLOCKS = list(range(50))

# Each rating's vote shrinks by this factor for every rating that came after it (read-time,
# nothing stored changes). Same constant as Studio's V2_PATH_OUTCOME_DECAY so the two forget
# at the same speed: half the vote sits in the last ~7 ratings. Without it a plain mean let
# twenty early ratings of one style outvote the new style until it had been out-rated
# one-for-one (the "first videos stick" observation, 2026-09-07).
RECENCY_DECAY = 0.9
MIN_PER_GROUP = 2  # need 2+ POSITIVE-weight and 2+ NEGATIVE-weight rows. Was 3 (parity with
# absolute/taste steering's floor, not derived from anything specific to this mechanism).
# 1 would mean the "direction" is literally one liked descriptor minus one disliked one --
# no averaging, as vulnerable to a single noisy capture as anything else this project has
# learned not to trust from one data point. 2 still averages a pair on each side while
# halving the wait (4 rated runs instead of 6) -- lower this only if you want to see raw
# per-pair noise, not a real direction. A weight near zero still gets logged and still
# contributes, just barely -- it does not count toward this floor.


def _log():
    try:
        from . import funpack_log as fl
    except ImportError:
        import funpack_log as fl
    return fl


def state_path(refinement_key, kind="repr_steer"):
    try:
        from .conditioning import refinement_state_path
    except ImportError:
        from conditioning import refinement_state_path
    return refinement_state_path(refinement_key, kind, prefix="refine_v2", extension="pt")


def _mask_from_mod_segments(mod_segments, seq_len, device, tag):
    """Shared by video_mask_from_mod_segments/audio_mask_from_mod_segments.

    The TARGET streams only: upstream packs ``[text | cond/ref | audio | video]`` and emits
    one mod_segments entry per target stream, so target video is the last entry and target
    audio the one before it. Matching by modality tag alone is wrong — tag 0 also marks the
    i2v anchor pin, guide keyframes, reference images/videos and the vision tokens inside the
    text span, none of which a "video rows" feature may steer or learn from.

    Returns None (not an all-False mask) when the layout cannot be proven — wrong tag at that
    position, or a sequence it does not fit (the 2-block token refiner's text-only calls) —
    so callers no-op rather than guess.
    """
    segs = list(mod_segments or ())
    if len(segs) < 2 or tag not in (0, 2):
        return None
    a, b, row = segs[-1] if tag == 0 else segs[-2]
    try:
        ok = bool(((row % 3) == tag).all()) if torch.is_tensor(row) else int(row) % 3 == tag
    except Exception:  # noqa: BLE001
        return None
    if not ok or not (0 <= a < b <= seq_len):
        return None
    mask = torch.zeros(seq_len, dtype=torch.bool, device=device)
    mask[a:b] = True
    return mask


def video_mask_from_mod_segments(mod_segments, seq_len, device):
    """mod_segments (from the H3 block hook's own args) -> a [seq_len] bool mask, True on the
    TARGET video rows (the last segment). None when that cannot be proven."""
    return _mask_from_mod_segments(mod_segments, seq_len, device, 0)


def audio_mask_from_mod_segments(mod_segments, seq_len, device):
    """Same as video_mask_from_mod_segments but for the TARGET audio rows (second to last)."""
    return _mask_from_mod_segments(mod_segments, seq_len, device, 2)


#: Query steering's store. Its rows are captured BEFORE H3's rotary position encoding; the
#: old "q_steer" rows were taken after it, mean something else, and must not mix with these.
Q_STEER_KIND = "q_steer_prerope"


def rope_table_rows(rope_freqs, seq_len):
    """H3's per-row rotation table ``[1, S, 1, half, 2, 2]`` (the block hook's ``rope_freqs``)
    -> ``[S, half, 2, 2]``, or None when it is not that shape for this sequence."""
    if not torch.is_tensor(rope_freqs) or rope_freqs.ndim != 6 \
            or rope_freqs.shape[1] != seq_len or rope_freqs.shape[-2:] != (2, 2):
        return None
    return rope_freqs[0, :, 0]


def rotate_rows(x, table, inverse=False):
    """Apply (or undo) H3's split-half rotary encoding to ``x`` ``[..., n, D]``.

    Mirrors comfy_kitchen's ``apply_rope_split_half1`` on the first ``2 * half`` dims — dim
    ``i`` pairs with ``i + half``, each pair turned by its row's 2x2 matrix — and passes the
    rest through. The inverse is the transpose, so undoing it is exact.
    """
    half = int(table.shape[-3])
    rot = 2 * half
    if x.shape[-1] < rot:
        raise ValueError(f"head dim {x.shape[-1]} < rotary dim {rot}")
    pairs = x[..., :rot].reshape(*x.shape[:-1], 2, half).movedim(-2, -1).float()
    mat = table.float()
    if inverse:
        mat = mat.transpose(-1, -2)
    out = (mat @ pairs.unsqueeze(-1)).squeeze(-1)
    out = out.movedim(-1, -2).reshape(*x.shape[:-1], rot).to(x.dtype)
    return torch.cat([out, x[..., rot:]], dim=-1)


def capture(hidden_state, video_mask, has_rows=None):
    """[seq_len, hidden] + a video-row mask -> a [hidden] descriptor, or None.

    Mean over rows, nothing more -- hidden_size is architecturally fixed (H3 is one width
    throughout), so there is no varying dimension to pool away the way a latent's spatial
    size needs adaptive_avg_pool1d for.

    `has_rows`: pass True/False when the caller already knows (from a mask it cached and
    already ran .any() on ONCE per seq_len, e.g. h3_repr_steering's/h3_q_steering's own
    hooks) whether `video_mask` has at least one True row -- this skips re-deriving it here
    with a SECOND host sync (bool() of a CUDA tensor) on every one of the 50 candidate
    blocks, every step, which is exactly the kind of stall this hot path needs to avoid (see
    project_reward_model_rework memory: implicated in a silent segfault under comfy_aimdo's
    dynamic-VRAM block prefetch). Leave it None (the default, used by direct/test callers
    with no cache of their own) to get the fully self-contained, host-synced check."""
    if video_mask is None:
        return None
    if has_rows is None:
        has_rows = bool(video_mask.any())
    if not has_rows:
        return None
    return hidden_state[video_mask].detach().float().mean(dim=0)


# --- persistence -------------------------------------------------------------------------
#
# One row per generation: {"desc": {block: tensor}, "weight": float}.
# Kept as a log and the direction recomputed from it on demand -- same reasoning
# trajectory_guidance's train_from_rows gives: this is cheap enough that being able to
# re-derive it (a log carried from another box, a fixed weighting formula) is worth more than
# the saved cycles of an online update. `desc` captures EVERY candidate block, not just the
# one that steers, so block_sweep() can be answered from data already being collected instead
# of needing its own separate collection pass.

def _load(refinement_key, kind="repr_steer"):
    path = state_path(refinement_key, kind)
    if not os.path.exists(path):
        return {"rows": [], "pending": None, "slots": {}}
    try:
        data = torch.load(path, map_location="cpu", weights_only=True)
    except Exception:  # noqa: BLE001
        return {"rows": [], "pending": None, "slots": {}}
    if not isinstance(data, dict):
        return {"rows": [], "pending": None, "slots": {}}
    data.setdefault("rows", [])
    data.setdefault("pending", None)
    data.setdefault("slots", {})
    return data


def _save(refinement_key, data, kind="repr_steer"):
    path = state_path(refinement_key, kind)
    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, descriptors, kind="repr_steer"):
    """Written on EVERY captured run, win or lose -- same fix trajectory_probe needed: a
    pending capture must not survive to be paired with the NEXT run's rating instead of its
    own, and the only way to guarantee that is to overwrite it every time, not just when a
    rating happens to follow. `descriptors` is {block: tensor}, whatever candidate blocks the
    sampler actually captured this run -- not required to cover every entry in
    CANDIDATE_BLOCKS (a shape the hook has never seen just means that block is missing).

    `kind` picks which sidecar file this is read/written under (see state_path) -- a second
    mechanism capturing a DIFFERENT vector space at the same block index (h3_q_steering
    captures Q, not hidden state) uses its own kind so the two never share rows under the
    same block key."""
    if not refinement_key or not descriptors:
        return
    data = _load(refinement_key, kind)
    data["pending"] = {int(b): d.detach().float().cpu() for b, d in descriptors.items()}
    try:
        _save(refinement_key, data, kind)
    except OSError as e:
        _log().failed("H3 representation steering", "save pending", e,
                       "this run's capture is not kept")


def commit(refinement_key, reward, kind="repr_steer"):
    """Pairs the pending capture with a rating's WEIGHT and appends it to the log. `reward`
    is the caller's already-resolved, already-admissibility-filtered reward -- this function
    trusts it rather than re-deriving a verdict, same as the value functions do with the same
    number.

    -> "recorded" | "no_pending" (nothing generated since the last rating) | "no_key". The
    caller prints this, not just the reward it computed -- a reward can be correct and still
    never reach a row if there was nothing to pair it with, and a log line that only ever
    reports the reward it is ABOUT to commit cannot tell that apart from a real commit."""
    if not refinement_key:
        return "no_key"
    data = _load(refinement_key, kind)
    pending = data.get("pending")
    data["pending"] = None
    if pending is None:
        try:
            _save(refinement_key, data, kind)
        except OSError:
            pass
        return "no_pending"
    try:
        data["rows"].append({"desc": pending, "weight": float(reward)})
    except (TypeError, ValueError):
        try:
            _save(refinement_key, data, kind)
        except OSError:
            pass
        return "no_pending"
    try:
        _save(refinement_key, data, kind)
    except OSError as e:
        _log().failed("H3 representation steering", "commit", e, "rating not recorded")
        return "no_pending"
    return "recorded"


def save_capture_slot(refinement_key, slot_id, descriptors):
    """Like save_pending, but keyed by an explicit slot instead of the single overwritable
    `pending` -- for a batch of generations that each need to stay independently rateable, in
    any order, any subset, instead of each new run silently evicting the last one waiting to
    be rated. `pending` is untouched by this (a slotted run does not also become "the"
    pending capture), so the ordinary one-at-a-time rating flow keeps working exactly as
    before for anything that doesn't pass a slot id."""
    if not refinement_key or not slot_id or not descriptors:
        return
    data = _load(refinement_key)
    data["slots"][str(slot_id)] = {int(b): d.detach().float().cpu() for b, d in descriptors.items()}
    try:
        _save(refinement_key, data)
    except OSError as e:
        _log().failed("H3 representation steering", "save capture slot", e,
                       f"slot {slot_id}'s capture is not kept")


def commit_slot(refinement_key, slot_id, reward):
    """Slot equivalent of commit(): pairs ONE named slot's capture with a rating and appends
    it to the same rows log direction() reads, then frees that slot. Slots are independent of
    each other and of `pending` -- rating one does not consume or disturb any other slot, so
    a whole batch can be rated in any order, any subset, whenever -- the actual point of
    a batch. Same return contract as commit(): "recorded" | "no_pending" | "no_key"."""
    if not refinement_key:
        return "no_key"
    data = _load(refinement_key)
    captured = data["slots"].pop(str(slot_id), None)
    if captured is None:
        try:
            _save(refinement_key, data)
        except OSError:
            pass
        return "no_pending"
    try:
        data["rows"].append({"desc": captured, "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 representation steering", "commit slot", e, "rating not recorded")
        return "no_pending"
    return "recorded"


def discard_slot(refinement_key, slot_id):
    """Drop a slot's capture without rating it (Skip) -- costs nothing, frees the state."""
    if not refinement_key:
        return
    data = _load(refinement_key)
    if data["slots"].pop(str(slot_id), None) is not None:
        try:
            _save(refinement_key, data)
        except OSError:
            pass


def clear_all(refinement_key, kind="repr_steer"):
    if not refinement_key:
        return
    try:
        os.remove(state_path(refinement_key, kind))
    except FileNotFoundError:
        pass


def direction(refinement_key, block=None, kind="repr_steer"):
    """-> (unit_vector, n_positive, n_negative), or (None, n_positive, n_negative) if either
    side is under MIN_PER_GROUP. One Awful and one Perfect is not a direction, it is two
    points -- doesn't matter how many barely-positive rows sit between them.

    Recency-weighted mean(liked) - mean(disliked): row i of n carries RECENCY_DECAY**(n-1-i),
    so the newest row votes 1.0 and older rows fade. Liked/disliked by weight SIGN only --
    not weighted by magnitude. Weighting by magnitude (a 9/10 pulling harder than a 6/10) is the statistically
    better estimator, but only once there is enough data for a weakly-rated row's noise to
    average out; at the handful of ratings this actually runs on, a mediocre "6" still gets a
    vote and dilutes the direction with an ambiguous case instead of being excluded by it. A
    binary split spends what little data exists on the clearest examples only. Revisit
    weighting once a key has real volume, not before -- see the 2026-09-04 dev session.

    Raw descriptors, not per-vector normalised: a hidden-state row's magnitude carries real
    content, and normalising it away is the trajectory_probe failure this deliberately avoids.

    `block` defaults to DEFAULT_BLOCK -- the one that actually steers. Rows captured before
    that block was in CANDIDATE_BLOCKS (or missing it for any other reason) are skipped, not
    treated as zero."""
    block = DEFAULT_BLOCK if block is None else int(block)
    data = _load(refinement_key, kind)
    rows = [r for r in data["rows"]
            if isinstance(r, dict) and "weight" in r and isinstance(r.get("desc"), dict)
            and block in r["desc"]]
    # Age counts every recorded rating, liked or not: the decay is "per rating since".
    aged = [(r["desc"][block], r["weight"], RECENCY_DECAY ** (len(rows) - 1 - i))
            for i, r in enumerate(rows)]
    liked = [(d, a) for d, w, a in aged if w > 0]
    disliked = [(d, a) for d, w, a in aged if w < 0]
    n_pos, n_neg = len(liked), len(disliked)
    if n_pos < MIN_PER_GROUP or n_neg < MIN_PER_GROUP:
        return None, n_pos, n_neg

    def _wmean(group):
        ds = torch.stack([d for d, _ in group]).float()
        a = torch.tensor([a for _, a in group], dtype=ds.dtype, device=ds.device)
        return (ds * a.view(-1, *([1] * (ds.dim() - 1)))).sum(dim=0) / a.sum()

    diff = _wmean(liked) - _wmean(disliked)
    norm = diff.norm()
    if not torch.isfinite(norm) or norm <= 1e-8:
        return None, n_pos, n_neg
    return diff / norm, n_pos, n_neg


def _trajectory_probe():
    try:
        from . import trajectory_probe as tp
    except ImportError:
        import trajectory_probe as tp
    return tp


def block_sweep(refinement_key, trials=2000):
    """-> {block: {"separation":..., "p_value":..., "noise_floor":..., "n":...}}, ranked by
    which of CANDIDATE_BLOCKS actually separates liked from disliked in YOUR rated data --
    reusing trajectory_probe's own permutation test (the one that correctly found real signal
    in the early-schedule work this session, after the centring fix) instead of trusting
    REINS' 50%-depth heuristic or porting a diagnostic built for a different question (LTXAV's
    identity-block search measured cross-scene consistency, not liked-vs-disliged separation).

    No cross-prompt grouping -- it was tried (prompt_hash on each row) and dropped: the
    R2V-wired conditioning this runs on never produced a real hash, so the guard was
    permanently inert, and fixing it was declined as not worth the continued attention. Rows
    with weight == 0 (a rating exactly at the scale's neutral midpoint) are dropped -- neither
    liked nor disliked."""
    tp = _trajectory_probe()
    data = _load(refinement_key)
    rows = [r for r in data["rows"]
            if isinstance(r, dict) and isinstance(r.get("desc"), dict)
            and r.get("weight", 0) != 0]
    out = {}
    for block in CANDIDATE_BLOCKS:
        have = [r for r in rows if block in r["desc"]]
        if len(have) < 2 * MIN_PER_GROUP:
            continue
        descriptors = [r["desc"][block] for r in have]
        labels = [r["weight"] > 0 for r in have]
        result = tp.permutation_test(descriptors, labels, trials=trials)
        if result is not None:
            out[block] = result
    return out