From 6e28f56c9225fa2fa58c7834a47df72f48c331e2 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:50:37 +0300 Subject: [PATCH] feat: REINS-style representation steering for H3 (experimental) Reaches inside the forward pass instead of around it: captures the video rows of one block's hidden state per generation, and once a refinement key has 3+ liked and 3+ disliked runs, adds their mean difference back into that block's output on every later run. No gradient, no MLP -- a mean-difference of vectors is the whole model. Masked to video rows only via the block's own mod_segments tags, never audio or text -- the one documented failure of this class of intervention (project_self_consistency_failed.md) happened on a different model (LTXAV) via joint audio-video attention. Off by default, H3-only, unvalidated. --- conditioning.py | 12 ++ h3_repr_steering.py | 187 +++++++++++++++++++++++ movie_editor/backend/settings_card.py | 3 + movie_editor/frontend/engine_settings.js | 7 +- samplers.py | 92 ++++++++++- tests/test_h3_repr_steering.py | 128 ++++++++++++++++ 6 files changed, 427 insertions(+), 2 deletions(-) create mode 100644 h3_repr_steering.py create mode 100644 tests/test_h3_repr_steering.py diff --git a/conditioning.py b/conditioning.py index b2aef94a..0bd00895 100644 --- a/conditioning.py +++ b/conditioning.py @@ -10141,6 +10141,18 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): n_out = self._v2_train_output_value_function(refinement_key, float(learning_profile.get("reward", 0.0))) if n_out is not None: print(f"[FunPackRefiner] Output value function updated — {n_out} samples") + # H3 representation steering: pairs the sampler's pending hidden-state capture + # with this rating. H3-only (checked here, not left to the sampler, so a + # non-H3 run's pending capture from a stale key is never mistaken for one) — + # same reward already resolved above, no separate admissibility question. + try: + from . import minimax_h3 as _h3 + from . import h3_repr_steering as _rs + except ImportError: + import minimax_h3 as _h3 + import h3_repr_steering as _rs + if _h3.is_h3_clip(clip): + _rs.commit(refinement_key, float(learning_profile.get("reward", 0.0))) # DynaShift negative memory: pair the sampler's pending raw latent with this # rating — promote into the per-key negative bank when the rating marks the run # a bad outcome (intrusions: awful / wrong_appearance; or the quality-missing diff --git a/h3_repr_steering.py b/h3_repr_steering.py new file mode 100644 index 00000000..354436ab --- /dev/null +++ b/h3_repr_steering.py @@ -0,0 +1,187 @@ +"""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 (liked runs' mean minus disliked runs' mean), 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 difference-of- +means instead of their supervised-PCA, because a two-group direction is what a mean +difference already gives you exactts.) + +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 measured this depth as the sweet spot: early +# enough that the rest of the stack can still act on the change, late enough that the +# concept the direction encodes has actually formed. +DEFAULT_BLOCK = 25 + +MIN_PER_GROUP = 3 # "need 3+ liked generations" -- same floor absolute/taste steering uses. + + +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, "repr_steer", prefix="refine_v2", extension="pt") + + +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 + VIDEO rows. Each entry is (a, b, row) where row is either a scalar `t_row*3 + tag` or, + for masked/multi-timestep rows, a tensor of them -- tag is `row % 3`, and 0 is video by + construction (`seg_tag` in model.py: video/cond/ref_img all tag 0). Returns None if + nothing matches, which happens on a shape this has never seen -- callers no-op rather + than guess.""" + mask = torch.zeros(seq_len, dtype=torch.bool, device=device) + found = False + for a, b, row in mod_segments or (): + if torch.is_tensor(row): + tags = (row % 3) == 0 + if tags.any(): + mask[a:b][tags] = True + found = True + elif int(row) % 3 == 0: + mask[a:b] = True + found = True + return mask if found else None + + +def capture(hidden_state, video_mask): + """[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.""" + if video_mask is None or not bool(video_mask.any()): + return None + return hidden_state[video_mask].detach().float().mean(dim=0) + + +# --- persistence ------------------------------------------------------------------------- +# +# One row per generation: {"desc": tensor, "reward": float}. Kept as a log and the direction +# recomputed from it on demand -- same reasoning trajectory_guidance's train_from_rows gives: +# a mean-difference is cheap enough that being able to re-derive it (a log carried from +# another box, a bad reward-sign fix) is worth more than the saved cycles of updating two +# running sums in place. + +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, descriptor): + """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.""" + if not refinement_key or descriptor is None: + return + data = _load(refinement_key) + data["pending"] = descriptor.detach().float().cpu() + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 representation steering", "save pending", e, + "this run's capture is not kept") + + +def commit(refinement_key, reward): + """Pairs the pending capture with a rating's reward and appends it to the log. A neutral + / skipped rating (reward is None) discards the pending capture without logging a row -- + same as trajectory_probe: no reward means no learning signal, not a zero one.""" + if not refinement_key: + return + data = _load(refinement_key) + pending = data.get("pending") + data["pending"] = None + if pending is None or reward is None: + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 representation steering", "commit", e, "rating not recorded") + return + try: + reward = float(reward) + except (TypeError, ValueError): + try: + _save(refinement_key, data) + except OSError: + pass + return + data["rows"].append({"desc": pending, "reward": reward}) + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 representation steering", "commit", e, "rating not recorded") + + +def clear_all(refinement_key): + if not refinement_key: + return + try: + os.remove(state_path(refinement_key)) + except FileNotFoundError: + pass + + +def direction(refinement_key): + """-> (unit_vector, n_liked, n_disliked), or (None, n_liked, n_disliked) if either group + is under MIN_PER_GROUP. One Awful and one Perfect is not a direction, it is two points. + + Raw means, not per-vector normalised: a hidden-state row's magnitude carries real content + (this is NOT the trajectory_probe failure mode). Shared content sits in BOTH groups' + means equally and cancels in the subtraction by construction -- centring only matters + when a vector is compared to itself normalised, which nothing here does.""" + data = _load(refinement_key) + liked = [r["desc"] for r in data["rows"] if isinstance(r, dict) and r.get("reward", 0) > 0] + disliked = [r["desc"] for r in data["rows"] if isinstance(r, dict) and r.get("reward", 0) < 0] + if len(liked) < MIN_PER_GROUP or len(disliked) < MIN_PER_GROUP: + return None, len(liked), len(disliked) + diff = torch.stack(liked).mean(dim=0) - torch.stack(disliked).mean(dim=0) + norm = diff.norm() + if not torch.isfinite(norm) or norm <= 1e-8: + return None, len(liked), len(disliked) + return diff / norm, len(liked), len(disliked) diff --git a/movie_editor/backend/settings_card.py b/movie_editor/backend/settings_card.py index eb3d33da..40101d77 100644 --- a/movie_editor/backend/settings_card.py +++ b/movie_editor/backend/settings_card.py @@ -196,6 +196,7 @@ _OWNED_BY = { "debug_log": ("identity_transfer_enabled", None), "dynashift_strength": ("dynashift", None), "dynashift_threshold": ("dynashift", None), + "h3_repr_steering_strength": ("h3_repr_steering", None), "embed_guidance_source": ("embed_guidance", None), "embed_guidance_strength": ("embed_guidance", None), "id_strength": ("identity_transfer_enabled", None), @@ -265,6 +266,8 @@ _EDITOR_DEFAULTS = { "frame_overlap": 16, "h3_phrase_emphasis": False, "h3_phrase_variability": 0.0, + "h3_repr_steering": False, + "h3_repr_steering_strength": 0.05, "h3_video_detail": 1.0, "id_strength": 1.0, "identity_transfer_enabled": False, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index db817970..e6f01a3c 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -23,6 +23,7 @@ "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", + "h3_repr_steering", "h3_repr_steering_strength", ]); const RATING_GATED_STUDIO = new Set(["reference_injection", "value_guidance", "steer_mode", "absolute_strength"]); @@ -312,6 +313,10 @@ hint: "How much of the matched bad direction is removed per step. 0.3 is a gentle nudge; 1.0 removes it outright each step." }, { name: "dynashift_threshold", label: "DynaShift match threshold", kind: "float", default: 0.6, min: 0.3, max: 0.95, step: 0.05, dependsOn: "dynashift", hint: "How closely a frame must resemble a banked bad one before steering starts. Lower is more aggressive and more likely to push away from content that was actually fine." }, + { name: "h3_repr_steering", label: "Representation steering (H3, experimental)", kind: "bool", default: false, + hint: "EXPERIMENTAL, unvalidated, H3 only. Reaches inside the model instead of around it: captures one block's video-row hidden state each generation, and once you have 3+ liked and 3+ disliked runs on this refinement key, adds their mean difference back into that block's output on every later run. No architectural ceiling like the attention-bias mechanisms have — push the strength too far and coherence can break with no warning." }, + { name: "h3_repr_steering_strength", label: "Representation steering strength", kind: "float", default: 0.05, min: 0.0, max: 2.0, step: 0.01, dependsOn: "h3_repr_steering", + 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: "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." }, @@ -501,7 +506,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"], + 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"], 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 83322a21..d4e0a7b7 100644 --- a/samplers.py +++ b/samplers.py @@ -2798,6 +2798,14 @@ class FunPackLTXAVSceneChainSampler: "default": 0.6, "min": 0.3, "max": 0.95, "step": 0.05, "tooltip": "Frame-similarity gate: a current frame must match a banked negative frame above this cosine similarity before any steering applies. Steering strength ramps from 0 at the threshold to full at similarity 1.0, so it self-releases once the unwanted feature is gone. Lower = more aggressive (risks steering away from legitimately similar content).", }), + "h3_repr_steering": ("BOOLEAN", { + "default": False, + "tooltip": "EXPERIMENTAL, unvalidated (H3 only). Every other rating-driven mechanism acts on the prompt's attention or on the latent between denoiser calls -- this one reaches INSIDE the model's forward pass instead. Captures the video rows of one block's hidden state each generation; once you have 3+ liked and 3+ disliked runs on this refinement key, the mean difference between them is added back into that same block's output on every later run. Training-free (no MLP, no gradient -- a mean of vectors is the whole model). Masked to video rows only, never audio or text.", + }), + "h3_repr_steering_strength": ("FLOAT", { + "default": 0.05, "min": 0.0, "max": 2.0, "step": 0.01, + "tooltip": "How much of the learned direction to add, as a fraction of that block's own hidden-state norm. Unlike the attention-bias channel, there is no architectural ceiling here -- push too far and coherence breaks with no warning. Start low; the way to tell it is doing anything at all is to push this deliberately high on a same-seed test and see if the video visibly moves, then back off.", + }), "alg_guide_blur_strength": ("FLOAT", { "default": 2.0, "min": 1.0, "max": 4.0, "step": 0.1, "tooltip": "Downsample factor for the guide-frame blur (alg_blur_guides). Higher = blurrier guide/JoyAI frames during the affected steps. Independent of the sampler's anchor alg_strength.", @@ -3658,7 +3666,9 @@ class FunPackLTXAVSceneChainSampler: 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, - bounded_attention_enabled=False,): + bounded_attention_enabled=False, + h3_repr_steering=False, h3_repr_steering_strength=0.05, + refinement_key="",): if sampler is None: raise ValueError("sampler input is required.") if not isinstance(sigmas, torch.Tensor): @@ -3720,6 +3730,10 @@ 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) + _repr_capture = [None] + if h3_repr_steering and refinement_key: + model = self._install_h3_repr_steering( + model, refinement_key, h3_repr_steering_strength, _repr_capture) try: sampled = comfy.sample.sample_custom( @@ -3727,6 +3741,12 @@ class FunPackLTXAVSceneChainSampler: noise_mask=latent.get("noise_mask"), seed=int(seed), callback=_progress_cb if pbar is not None else None, ) + if _repr_capture[0] is not None: + try: + from . import h3_repr_steering as _rs + except ImportError: + import h3_repr_steering as _rs + _rs.save_pending(refinement_key, _repr_capture[0]) finally: self._remove_bounded_attention(_ba_handles) self._assert_finite_sample(sampled) @@ -6535,6 +6555,71 @@ class FunPackLTXAVSceneChainSampler: "the rating's per-phrase weighting is NOT being applied") return model + def _install_h3_repr_steering(self, model, refinement_key, strength, capture_holder): + """EXPERIMENTAL, unvalidated (H3 only). See h3_repr_steering.py. + + Reaches inside the forward pass: hooks block DEFAULT_BLOCK, adds the learned liked- + minus-disliked direction to that block's VIDEO rows (never audio or text — derived + fresh from `mod_segments` each call, cached by seq_len), and captures this run's own + descriptor into `capture_holder[0]` so the caller can save it as the next pending row + once sampling finishes. No-ops (clears nothing, costs nothing) when there are not yet + 3+ liked and 3+ disliked runs on this key — direction() says so, this just reports it. + """ + try: + try: + from . import h3_repr_steering as _rs + except ImportError: + import h3_repr_steering as _rs + direction, n_liked, n_disliked = _rs.direction(refinement_key) + if direction is None: + print(f"[FunPackSceneChain] H3 representation steering: not enough data yet " + f"({n_liked} liked / {n_disliked} disliked, need {_rs.MIN_PER_GROUP} of " + f"each) — capturing this run, not steering it.") + + 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", {})) + inner = dit_patches.get(("double_block", _rs.DEFAULT_BLOCK)) + mask_cache = {} + _strength = float(strength) + + def _hook(args, extra): + out = extra["original_block"](args)["img"] if inner is None else \ + inner(args, extra)["img"] + 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: + if direction is not None and _strength > 0.0: + rows = out[mask] + row_norm = rows.detach().float().norm(dim=-1).mean() + if torch.isfinite(row_norm) and row_norm > 0: + out = out.clone() + out[mask] = rows + (direction.to(out.dtype).to(out.device) + * _strength * row_norm).to(rows.dtype) + desc = _rs.capture(out, mask) + if desc is not None: + capture_holder[0] = desc + return {"img": out} + + dit_patches[("double_block", _rs.DEFAULT_BLOCK)] = _hook + patches_replace["dit"] = dit_patches + to["patches_replace"] = patches_replace + patched.model_options["transformer_options"] = to + if direction is not None: + print(f"[FunPackSceneChain] H3 representation steering: applying learned " + f"direction ({n_liked} liked / {n_disliked} disliked) at block " + f"{_rs.DEFAULT_BLOCK}, strength {_strength:g}.") + return patched + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackSceneChain", "H3 representation steering", _e, + "not applied and not capturing 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 @@ -6792,6 +6877,7 @@ class FunPackLTXAVSceneChainSampler: output_guidance=False, output_guidance_strength=0.02, 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, 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, @@ -7617,6 +7703,9 @@ class FunPackLTXAVSceneChainSampler: alg_anchor_strength=alg_anchor_strength, alg_anchor_sigma_threshold=alg_anchor_sigma_threshold, bounded_attention_enabled=bounded_attention_enabled, + h3_repr_steering=h3_repr_steering, + h3_repr_steering_strength=h3_repr_steering_strength, + refinement_key=refinement_key_input, ) # cut_opening_frames: let the real, untouched i2v anchor condition the scene # at full strength, then cut it out of the finished clip instead of weakening @@ -8296,6 +8385,7 @@ class FunPackLTXAVSceneChainSampler: output_guidance=False, output_guidance_strength=0.02, 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, 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, diff --git a/tests/test_h3_repr_steering.py b/tests/test_h3_repr_steering.py new file mode 100644 index 00000000..0355d17c --- /dev/null +++ b/tests/test_h3_repr_steering.py @@ -0,0 +1,128 @@ +import sys + +import pytest +import torch + +sys.path.insert(0, ".") +import h3_repr_steering as rs # noqa: E402 + + +@pytest.fixture(autouse=True) +def _isolate_state(monkeypatch, tmp_path): + monkeypatch.setattr(rs, "state_path", lambda key: str(tmp_path / f"{key}.pt")) + + +# --- video_mask_from_mod_segments ------------------------------------------- + +def test_scalar_rows_tag_video_by_mod_3(): + # row = t_row*3 + tag; tag 0 = video, 1 = text, 2 = audio (seg_tag in model.py) + mod_segments = [(0, 4, 6), (4, 10, 7), (10, 14, 8)] # video, text, audio + mask = rs.video_mask_from_mod_segments(mod_segments, seq_len=14, device="cpu") + assert mask.tolist() == [True] * 4 + [False] * 10 + + +def test_tensor_rows_use_the_same_mod_3_rule(): + rows = torch.tensor([6, 7, 6, 8]) # video, text, video, audio + mod_segments = [(0, 4, rows)] + mask = rs.video_mask_from_mod_segments(mod_segments, seq_len=4, device="cpu") + assert mask.tolist() == [True, False, True, False] + + +def test_no_video_segment_returns_none_rather_than_an_empty_mask(): + mod_segments = [(0, 4, 7), (4, 8, 8)] # text, audio only + assert rs.video_mask_from_mod_segments(mod_segments, seq_len=8, device="cpu") is None + + +def test_empty_mod_segments_returns_none(): + assert rs.video_mask_from_mod_segments([], seq_len=10, device="cpu") is None + + +# --- capture ----------------------------------------------------------------- + +def test_capture_means_over_video_rows_only(): + h = torch.tensor([[1.0, 1.0], [3.0, 3.0], [9.0, 9.0]]) + mask = torch.tensor([True, True, False]) + desc = rs.capture(h, mask) + assert torch.allclose(desc, torch.tensor([2.0, 2.0])) + + +def test_capture_with_no_video_rows_is_none(): + h = torch.zeros(3, 2) + mask = torch.zeros(3, dtype=torch.bool) + assert rs.capture(h, mask) is None + + +def test_capture_with_no_mask_is_none(): + assert rs.capture(torch.zeros(3, 2), None) is None + + +# --- persistence + direction -------------------------------------------------- + +def test_direction_needs_min_per_group_of_each(): + for i in range(rs.MIN_PER_GROUP - 1): + rs.save_pending("k", torch.tensor([1.0, 0.0])) + rs.commit("k", 1.0) + for i in range(rs.MIN_PER_GROUP): + rs.save_pending("k", torch.tensor([-1.0, 0.0])) + rs.commit("k", -1.0) + direction, n_liked, n_disliked = rs.direction("k") + assert direction is None + assert n_liked == rs.MIN_PER_GROUP - 1 + assert n_disliked == rs.MIN_PER_GROUP + + +def test_direction_points_from_disliked_mean_to_liked_mean(): + for _ in range(rs.MIN_PER_GROUP): + rs.save_pending("k", torch.tensor([5.0, 0.0])) + rs.commit("k", 1.0) + for _ in range(rs.MIN_PER_GROUP): + rs.save_pending("k", torch.tensor([-5.0, 0.0])) + rs.commit("k", -1.0) + direction, n_liked, n_disliked = rs.direction("k") + assert direction is not None + assert n_liked == rs.MIN_PER_GROUP and n_disliked == rs.MIN_PER_GROUP + assert torch.allclose(direction, torch.tensor([1.0, 0.0]), atol=1e-5) + + +def test_shared_content_cancels_without_needing_explicit_centring(): + """Both groups carry the same huge shared offset (simulating prompt content); the + direction must still point along the true liked-vs-disliked axis, not the shared one.""" + shared = torch.tensor([1000.0, 1000.0]) + for _ in range(rs.MIN_PER_GROUP): + rs.save_pending("k", shared + torch.tensor([1.0, 0.0])) + rs.commit("k", 1.0) + for _ in range(rs.MIN_PER_GROUP): + rs.save_pending("k", shared + torch.tensor([-1.0, 0.0])) + rs.commit("k", -1.0) + direction, _, _ = rs.direction("k") + assert torch.allclose(direction, torch.tensor([1.0, 0.0]), atol=1e-5) + + +def test_pending_is_discarded_on_a_neutral_rating(): + rs.save_pending("k", torch.tensor([1.0, 0.0])) + rs.commit("k", None) + direction, n_liked, n_disliked = rs.direction("k") + assert n_liked == 0 and n_disliked == 0 + + +def test_pending_overwrites_not_accumulates(): + rs.save_pending("k", torch.tensor([1.0, 0.0])) + rs.save_pending("k", torch.tensor([2.0, 0.0])) # a second capture before any rating + rs.commit("k", 1.0) + data = rs._load("k") + assert len(data["rows"]) == 1 + assert torch.allclose(data["rows"][0]["desc"], torch.tensor([2.0, 0.0])) + + +def test_clear_all_removes_the_state_file(): + rs.save_pending("k", torch.tensor([1.0, 0.0])) + rs.commit("k", 1.0) + rs.clear_all("k") + direction, n_liked, n_disliked = rs.direction("k") + assert direction is None and n_liked == 0 and n_disliked == 0 + + +def test_no_refinement_key_is_a_safe_no_op(): + rs.save_pending("", torch.tensor([1.0])) # must not raise + rs.commit("", 1.0) + rs.clear_all("")