olivv-cs-ComfyUI-FunPack/block_steer.py

"""Block Steering (EXPERIMENTAL): rating-learned per-block residual gains for LTX-AV.

RLHF at the MODEL level instead of the conditioning level. The rating system already
attributes quality to sampling settings and conditioning; this adds the missing
coordinate system - which REGIONS OF THE MODEL were doing the work when a generation
came out good or bad - and then steers per-block application accordingly.

Three stages, all riding existing FunPack infra:

1. MEASURE - forward hooks on the transformer blocks record, per model call, each
   block's relative residual delta on the VIDEO stream: ||vx_out - vx_in|| / ||vx_in||.
   A block that barely moves the hidden state is coasting; one with a large relative
   delta is doing the work for this prompt at this noise level. Hooks are tagged with
   ltx_enhancements' _FUNPACK_HOOK_TAG so the existing leak-strip covers them.

2. ATTRIBUTE - two signals, folded into a per-key `block_profile.pt` sidecar:
   * run-level: the end-of-run mean activity fingerprint is staged (`block_snapshot.pt`,
     same stage-then-pair convention as x0_snapshot) and paired with the user's rating
     reward on the next Refiner cycle - liked / disliked / all-runs EMA profiles.
   * step-level: when output_guidance's value function is trained, a model-function
     wrapper scores each step's x0_hat (one tiny-MLP forward, no extra model pass) and
     credits blocks by (value delta this step) x (their activity this step, centered).
     Converges in a handful of runs - matters because keys are retrained per instance.

3. STEER - per-block gains on the video residual only: vx_out' = vx_in + g_b*(vx_out -
   vx_in), g_b = 1 + strength * score_b, scores zero-mean/unit-max so emphasis is
   REDISTRIBUTED rather than globally inflated, clamped to +/-MAX_GAIN_DELTA. The audio
   stream is never touched (residual-gain on the video half is the intervention family
   that works on LTX-AV; hidden-state injection through joint attention is the one that
   corrupts audio). Exact no-op until the profile has enough paired data.

Everything is per-refinement-key and rebuilt from ratings in a few generations - a
mechanism upgrade, not a data store to migrate.
"""

from __future__ import annotations

import os

import torch

try:
    from .ltx_enhancements import _FUNPACK_HOOK_TAG, _funpack_locate_blocks
except ImportError:
    from ltx_enhancements import _FUNPACK_HOOK_TAG, _funpack_locate_blocks

# Steering never moves a block's residual contribution by more than this, regardless of
# strength widget or how lopsided the learned scores are.
MAX_GAIN_DELTA = 0.10
# Profile readiness: at least this many rating-paired runs (liked+disliked pools)…
MIN_RATED_RUNS = 2
# …or this many value-credit runs, before scores are considered meaningful.
MIN_CREDIT_RUNS = 2
# Blend weights when both attribution signals are available.
RATING_WEIGHT = 0.6
CREDIT_WEIGHT = 0.4
_EPS = 1e-8


def _video_stream(x):
    """The video tensor of a block input/output: (vx, ax) tuple on LTX-AV, plain tensor
    on video-only LTXV."""
    if isinstance(x, (tuple, list)) and len(x) >= 1:
        return x[0]
    return x if isinstance(x, torch.Tensor) else None


# ---------------------------------------------------------------------------
# Stage 1 - measurement
# ---------------------------------------------------------------------------

class BlockActivityRecorder:
    """Accumulates per-block relative residual deltas (video stream) across a run.

    `cur` always holds the most recent model call's per-block activity so the value-
    credit wrapper (which runs after all blocks in the same call) can read it.
    """

    def __init__(self, n_blocks: int):
        self.n_blocks = int(n_blocks)
        self.sums = torch.zeros(self.n_blocks, dtype=torch.float32)
        self.counts = torch.zeros(self.n_blocks, dtype=torch.float32)
        self.cur = torch.zeros(self.n_blocks, dtype=torch.float32)
        # value-credit accumulation (filled by the credit wrapper, if installed)
        self.credit = torch.zeros(self.n_blocks, dtype=torch.float32)
        self.credit_steps = 0
        self._prev_score = None

    def record(self, idx: int, rel_delta: float):
        if 0 <= idx < self.n_blocks:
            self.sums[idx] += rel_delta
            self.counts[idx] += 1.0
            self.cur[idx] = rel_delta

    def fingerprint(self):
        """Mean per-block activity over the run, or None if nothing was recorded."""
        if float(self.counts.max()) <= 0:
            return None
        return self.sums / self.counts.clamp(min=1.0)

    def credit_step(self, score: float):
        """Fold one model call's value score into per-block credit: blocks that were
        MORE active than average during a step where predicted quality rose (fell) get
        positive (negative) credit."""
        if self._prev_score is not None:
            delta = float(score) - self._prev_score
            centered = self.cur - self.cur.mean()
            self.credit += delta * centered
            self.credit_steps += 1
        self._prev_score = float(score)

    def credit_vector(self):
        if self.credit_steps <= 0:
            return None
        return self.credit / float(self.credit_steps)


def install_recorder(model, recorder: BlockActivityRecorder):
    """Tagged forward hooks on every transformer block; returns handles (possibly [])."""
    blocks = _funpack_locate_blocks(model)
    if blocks is None:
        return []
    handles = []
    for i, blk in enumerate(blocks):
        def _hook(_module, args, output, _idx=i, _rec=recorder):
            try:
                v_in = _video_stream(args[0] if args else None)
                v_out = _video_stream(output)
                if v_in is None or v_out is None or v_in.shape != v_out.shape:
                    return output
                with torch.no_grad():
                    denom = float(v_in.norm()) + _EPS
                    _rec.record(_idx, float((v_out - v_in).norm()) / denom)
            except Exception:
                pass
            return output
        setattr(_hook, _FUNPACK_HOOK_TAG, True)
        handles.append(blk.register_forward_hook(_hook))
    return handles


def make_credit_wrapper(old_wrapper, recorder: BlockActivityRecorder, value_fn):
    """model_function_wrapper that scores each call's x0_hat with the (already trained)
    output-space value function and folds the step-to-step delta into block credit.
    Installed INNERMOST so it scores the base model's own prediction, before
    embed_guidance / output_guidance corrections layer on top."""

    def _call(apply_fn, a):
        if old_wrapper is not None:
            return old_wrapper(apply_fn, a)
        return apply_fn(a["input"], a["timestep"], **a.get("c", {}))

    def _wrapper(apply_fn, args, _rec=recorder, _vf=value_fn):
        denoised = _call(apply_fn, args)
        try:
            with torch.inference_mode(False), torch.no_grad():
                c = _vf.compress(denoised.detach().float())
                score = float(_vf(c).mean())
            _rec.credit_step(score)
        except Exception:
            pass
        return denoised

    return _wrapper


# ---------------------------------------------------------------------------
# Stage 3 - steering
# ---------------------------------------------------------------------------

def install_steer(model, gains):
    """Tagged forward hooks scaling each block's VIDEO residual by its gain:
    vx' = vx_in + g*(vx_out - vx_in). Audio passes through untouched. Returns handles;
    blocks whose gain is ~1.0 get no hook at all (byte-identical native path)."""
    blocks = _funpack_locate_blocks(model)
    if blocks is None:
        return []
    handles = []
    for i, blk in enumerate(blocks):
        g = float(gains[i]) if i < len(gains) else 1.0
        if abs(g - 1.0) < 1e-4:
            continue
        def _hook(_module, args, output, _g=g):
            try:
                v_in = _video_stream(args[0] if args else None)
                if isinstance(output, tuple) and len(output) == 2:
                    v_out, a_out = output
                    if v_in is None or v_in.shape != v_out.shape:
                        return output
                    return (v_in + _g * (v_out - v_in), a_out)
                if isinstance(output, torch.Tensor) and v_in is not None and v_in.shape == output.shape:
                    return v_in + _g * (output - v_in)
            except Exception:
                pass
            return output
        setattr(_hook, _FUNPACK_HOOK_TAG, True)
        handles.append(blk.register_forward_hook(_hook))
    return handles


def remove_handles(handles):
    for h in handles or []:
        try:
            h.remove()
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Stage 2 - persistence + attribution
# ---------------------------------------------------------------------------

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


def save_run_snapshot(refinement_key, recorder: BlockActivityRecorder):
    """End of run: stage this run's fingerprint (+ optional value credit) so the NEXT
    rating cycle can pair it with a reward. Atomic write - a rating read never sees a
    partial file. Same convention as the sampler's x0_snapshot."""
    if not refinement_key:
        return False
    fp = recorder.fingerprint()
    if fp is None:
        return False
    try:
        path = _state_path(refinement_key, "block_snapshot")
        os.makedirs(os.path.dirname(path), exist_ok=True)
        tmp = path + ".tmp"
        torch.save({
            "activity": fp.cpu(),
            "credit": (recorder.credit_vector().cpu() if recorder.credit_vector() is not None else None),
            "n_blocks": recorder.n_blocks,
        }, tmp)
        os.replace(tmp, path)
        return True
    except Exception as e:
        print(f"[FunPackBlockSteer] snapshot save failed: {e}")
        return False


def _ema(old, x, n):
    a = max(0.2, 1.0 / (float(n) + 1.0))
    return x if old is None else old * (1.0 - a) + x * a


def update_profile_with_rating(refinement_key, reward):
    """Pair the staged snapshot with a rating reward (called by the Refiner at the same
    site that trains the output value function). Folds:
      * every run into the all-runs mean (the contrast baseline),
      * reward >= 0.3 into the liked profile, reward <= 0.0 into the disliked profile
        (the mild-positive band is too ambiguous to move either pole - same rule as the
        Absolute store),
      * the run's value credit (rating-independent, self-supervised) into the credit EMA.
    Consumes the snapshot so one run is never paired twice. Returns the number of
    rating-paired runs, or None when there was nothing to pair."""
    if not refinement_key:
        return None
    try:
        snap_path = _state_path(refinement_key, "block_snapshot")
        if not os.path.exists(snap_path):
            return None
        snap = torch.load(snap_path, map_location="cpu", weights_only=False)
        activity = snap.get("activity")
        if not isinstance(activity, torch.Tensor):
            return None
        prof_path = _state_path(refinement_key, "block_profile")
        prof = {}
        if os.path.exists(prof_path):
            prof = torch.load(prof_path, map_location="cpu", weights_only=False)
            if int(prof.get("n_blocks", -1)) != activity.numel():
                prof = {}  # different model depth - start fresh
        prof["n_blocks"] = activity.numel()

        prof["all_mean"] = _ema(prof.get("all_mean"), activity, prof.get("all_n", 0))
        prof["all_n"] = int(prof.get("all_n", 0)) + 1
        r = float(reward)
        if r >= 0.3:
            prof["liked"] = _ema(prof.get("liked"), activity, prof.get("liked_n", 0))
            prof["liked_n"] = int(prof.get("liked_n", 0)) + 1
        elif r <= 0.0:
            prof["disliked"] = _ema(prof.get("disliked"), activity, prof.get("disliked_n", 0))
            prof["disliked_n"] = int(prof.get("disliked_n", 0)) + 1
        credit = snap.get("credit")
        if isinstance(credit, torch.Tensor) and credit.numel() == activity.numel():
            prof["credit"] = _ema(prof.get("credit"), credit, prof.get("credit_n", 0))
            prof["credit_n"] = int(prof.get("credit_n", 0)) + 1

        tmp = prof_path + ".tmp"
        torch.save(prof, tmp)
        os.replace(tmp, prof_path)
        os.remove(snap_path)  # consumed - never pair one run with two ratings
        return int(prof.get("liked_n", 0)) + int(prof.get("disliked_n", 0))
    except Exception as e:
        print(f"[FunPackBlockSteer] profile update failed: {e}")
        return None


def _normalize(v):
    v = v - v.mean()
    m = float(v.abs().max())
    return v / m if m > _EPS else None


def load_block_scores(refinement_key):
    """Combined, normalized per-block scores in [-1, 1] (zero-mean, unit max-abs), or
    None while the profile lacks enough paired data. Rating contrast is liked-vs-
    disliked when both poles exist, else the available pole against the all-runs mean."""
    if not refinement_key:
        return None
    try:
        prof_path = _state_path(refinement_key, "block_profile")
        if not os.path.exists(prof_path):
            return None
        prof = torch.load(prof_path, map_location="cpu", weights_only=False)
        liked, disliked = prof.get("liked"), prof.get("disliked")
        liked_n, disliked_n = int(prof.get("liked_n", 0)), int(prof.get("disliked_n", 0))
        all_mean, all_n = prof.get("all_mean"), int(prof.get("all_n", 0))
        credit, credit_n = prof.get("credit"), int(prof.get("credit_n", 0))

        rating_part = None
        if liked_n + disliked_n >= MIN_RATED_RUNS:
            if liked is not None and disliked is not None:
                rating_part = _normalize(liked - disliked)
            elif liked is not None and all_mean is not None and all_n >= 2:
                rating_part = _normalize(liked - all_mean)
            elif disliked is not None and all_mean is not None and all_n >= 2:
                rating_part = _normalize(-(disliked - all_mean))
        credit_part = _normalize(credit) if (credit is not None and credit_n >= MIN_CREDIT_RUNS) else None

        if rating_part is not None and credit_part is not None:
            combined = RATING_WEIGHT * rating_part + CREDIT_WEIGHT * credit_part
        else:
            combined = rating_part if rating_part is not None else credit_part
        if combined is None:
            return None
        return _normalize(combined)
    except Exception as e:
        print(f"[FunPackBlockSteer] score load failed: {e}")
        return None


def gains_from_scores(scores, strength):
    """Per-block gains 1 + strength*score, hard-clamped to 1 +/- MAX_GAIN_DELTA."""
    s = max(0.0, float(strength))
    g = 1.0 + (s * scores).clamp(-MAX_GAIN_DELTA, MAX_GAIN_DELTA)
    return [float(x) for x in g]