diff --git a/CHANGELOG.md b/CHANGELOG.md index 38f527b7..d426511e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ ## [Unreleased] +### Added +- **The rating's per-phrase emphasis now reaches MiniMax H3.** Studio has decomposed every + run into phrases, words and n-grams and scored each against the rating since 2.0, but the + only thing reading that was the attn2 K/V patch — which H3 skips, having one packed + self-attention stream and no cross-attention. The learning ran every generation and landed + nowhere. It is now applied where H3 can act on it: `log(weight)` added to the attention + logits at that phrase's token positions, so its share of the attention mass is multiplied + by the weight. Boost only; damping a phrase you typed is opt-in. Phrases are located with + the tokenizer's offset mapping rather than re-encoded through a 32B text encoder. + A weighted run is a dense run — the bias is a mask, and SLA cannot carry one. + ### Reverted - **GGUF loading is back to expanding the checkpoint at load.** The quantized path through ComfyUI-GGUF loaded in seconds instead of a minute, and produced a model the sampler could diff --git a/conditioning.py b/conditioning.py index 03adbc18..33aea36a 100644 --- a/conditioning.py +++ b/conditioning.py @@ -13065,7 +13065,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): guess_mode=guess_mode, guess_direction=guess_direction, guess_range=guess_range, guess_freeze_seed=guess_freeze_seed) return ( - self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan, clip=clip, encode_cache=encode_cache), + self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan, clip=clip, encode_cache=encode_cache, phrase_memory=global_state.get("phrase_memory"), axis_feedback=repair_feedback), status, training_info, loss_graph, @@ -13083,7 +13083,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): repair_feedback, current_family_slot, _vf_for_memory, _concept_dir, _concept_strength, _current_final) return ( - self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan, clip=clip, encode_cache=encode_cache), + self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan, clip=clip, encode_cache=encode_cache, phrase_memory=global_state.get("phrase_memory"), axis_feedback=repair_feedback), status + enhancement_status, training_info, loss_graph, @@ -13756,11 +13756,83 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): out.append([cond, meta] if isinstance(entry, list) else (cond, meta)) return out + def _v2_apply_h3_token_weights(self, conditioning_list, clip, phrase_memory=None, + axis_feedback=None, fallback_text=""): + """Tag each scene with the token spans the RATING says deserve more attention. + + Studio has decomposed and scored every phrase, word and n-gram since 2.0 + (_v2_concept_units_for_run, _v2_update_phrase_memory). Until now the only consumer + was the attn2 K/V patch, which H3 skips outright — one packed self-attention stream, + no cross-attention — so on H3 that learning ran every generation and landed nowhere. + + This writes it as (start, end, weight) token spans for the sampler to turn into an + attention-logit bias. Text-only and cheap: phrases are located with the tokenizer's + offset mapping, not by re-encoding each one through a 32B text encoder. No behaviour + change unless the sampler's toggle is on. + """ + if clip is None or not isinstance(phrase_memory, dict) or not phrase_memory: + return conditioning_list + try: + try: # matches how this file reaches its siblings + from . import minimax_h3 as _h3 + from . import h3_token_weights as _tw + except ImportError: + import minimax_h3 as _h3 + import h3_token_weights as _tw + if not _h3.is_h3_clip(clip): + return conditioning_list + tokenizer = _tw.h3_tokenizer(clip) + if tokenizer is None: + return conditioning_list + feedback = axis_feedback or {} + units = [{"text": text, "kind": entry.get("kind"), + "effective_category_scores": entry.get("effective_category_scores")} + for text, entry in phrase_memory.items() if isinstance(entry, dict)] + weighted = _tw.weights_from_memory( + units, phrase_memory, + missing_axes=feedback.get("missing_axes", ()), + wrong_axes=feedback.get("wrong_axes", ())) + if not weighted: + return conditioning_list + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackStudio", "H3 token weighting", _e, + "the rating's per-phrase emphasis is NOT being applied") + return conditioning_list + + out = [] + applied = 0 + for entry in conditioning_list or []: + if not (isinstance(entry, (list, tuple)) and len(entry) >= 2 and isinstance(entry[1], dict)): + out.append(entry) + continue + cond, meta = entry[0], dict(entry[1]) + text = str(meta.get("funpack_scene_text") or "").strip() or str(fallback_text or "") + try: + from . import h3_token_weights as _tw + except ImportError: + import h3_token_weights as _tw + spans, prompt_tokens = _tw.locate(tokenizer, text, weighted[:8]) + if spans and prompt_tokens: + meta["funpack_h3_token_weights"] = { + "spans": spans, "prompt_tokens": int(prompt_tokens)} + applied += 1 + out.append([cond, meta] if isinstance(entry, list) else (cond, meta)) + if applied: + _log.note_on_change( + "h3:token_weights", "FunPackStudio", + f"rating-derived phrase emphasis: {len(weighted)} phrase(s) weighted across " + f"{applied} scene(s) (e.g. " + + ", ".join(f"{t!r} x{w:.2f}" for t, w in weighted[:3]) + + ") — applied as an attention bias on H3's packed stream, where the attn2 " + "K/V emphasis has nothing to hook.") + return out + def _v2_finalize_conditioning(self, conditioning_list, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=None, temporal_style="natural", temporal_fallback_text="", scene_refinement_keys=None, learning_profile=None, - conditioning_plan=None, clip=None, encode_cache=None): + conditioning_plan=None, clip=None, encode_cache=None, + phrase_memory=None, axis_feedback=None): """Single output hook for both steering modes. Relative = per-key VF ascend (current behaviour). Absolute = global taste pull. Both = layer them. Finally, if Interactive Guessing has learned a safe-spread ceiling, clamp the output conditioning's video-channel @@ -13774,6 +13846,9 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): out = self._v2_apply_rapid_temporal(out, temporal_style) out = self._v2_apply_auto_temporal(out, temporal_style, fallback_text=temporal_fallback_text) out = self._v2_apply_bounded_attention(out, clip, fallback_text=temporal_fallback_text, encode_cache=encode_cache) + out = self._v2_apply_h3_token_weights(out, clip, phrase_memory=phrase_memory, + axis_feedback=axis_feedback, + fallback_text=temporal_fallback_text) if mode in ("relative", "both"): out = self._v2_apply_scene_refinement_keys( out, scene_refinement_keys, refinement_key, diff --git a/h3_token_weights.py b/h3_token_weights.py index f8c41246..730d8d69 100644 --- a/h3_token_weights.py +++ b/h3_token_weights.py @@ -73,12 +73,7 @@ def token_spans(tokenizer, clean_text: str, char_spans): offsets = enc["offset_mapping"] except Exception: # noqa: BLE001 — a slow tokenizer has no offsets; weighting is off return [] - out = [] - for c_start, c_end, weight in char_spans: - toks = [i for i, (a, b) in enumerate(offsets) if a < c_end and b > c_start] - if toks: - out.append((toks[0], toks[-1] + 1, weight)) - return out + return token_spans_from_offsets(offsets, char_spans) def bias_value(weight: float) -> float: @@ -160,3 +155,144 @@ def make_override(spans, prompt_tokens: int, cond_len: int, inner=None, on_apply return run(mask) return override + + +# --- rating-derived weights ------------------------------------------------ +# +# Studio already knows which words to weight. `_v2_concept_units_for_run` decomposes every +# run into phrases, single words and n-grams; `_v2_update_phrase_memory` trains each one's +# category scores against the rating's missing / satisfied / regressed / wrong axes. What +# was missing was a CONSUMER: the only thing reading that signal was the attn2 K/V patch, +# which H3 skips outright (one packed self-attention stream, no cross-attention). So the +# learning ran every generation and landed nowhere. +# +# This turns an entry into a multiplicative weight, which `bias_value` turns into the logit +# bias. Nothing here decides what a phrase MEANS — that is already decided upstream. + +# Boost only. A phrase carrying an axis the rating said was MISSING gets more attention; +# nothing is taken away unless damping is asked for explicitly. Damping can suppress a +# phrase the user typed, which is a knob going inert on its own — an opt-in, not a default. +MIN_LEARNED_WEIGHT = 0.25 +MAX_LEARNED_WEIGHT = 3.0 + +_KIND_SCALE = {"prompt_phrase": 1.0, "phrase": 1.0, "auto_phrase": 0.72, + "repair_candidate": 0.64, "ngram": 0.62, "token": 0.24} + + +def _kind_scale(kind): + return _KIND_SCALE.get(str(kind or "phrase"), 0.5) + + +def weights_from_memory(phrases, phrase_memory=None, missing_axes=(), wrong_axes=(), + wrong_appearance=False, strength=0.5, damp=False, + kind_scale=None): + """[(text, weight), ...] for the phrases the rating has something to say about. + + `phrases` are this run's classified units (each with `text`, `kind` and + `effective_category_scores`); `phrase_memory` is the learned store, which overrides a + unit's own scores when it has seen the phrase before. Phrases the rating is neutral + about are omitted rather than returned at 1.0, so nothing is biased by accident. + """ + scale = kind_scale or _kind_scale + missing = tuple(missing_axes or ()) + wrong = tuple(wrong_axes or ()) + memory = phrase_memory if isinstance(phrase_memory, dict) else {} + out = [] + for unit in phrases or []: + if not isinstance(unit, dict): + continue + text = str(unit.get("text", "")).strip().lower() + if not text: + continue + entry = memory.get(text) if isinstance(memory.get(text), dict) else {} + scores = entry.get("effective_category_scores") or unit.get("effective_category_scores") \ + or unit.get("category_scores") or {} + if not isinstance(scores, dict): + continue + try: + up = sum(float(scores.get(a, 0.0)) for a in missing) + down = sum(float(scores.get(a, 0.0)) for a in wrong) + if wrong_appearance: + down += float(scores.get("appearance", 0.0)) + except (TypeError, ValueError): + continue + delta = up - (down if damp else 0.0) + if delta <= 0.0: + continue + weight = 1.0 + float(strength) * scale(entry.get("kind", unit.get("kind"))) * delta + weight = max(MIN_LEARNED_WEIGHT, min(MAX_LEARNED_WEIGHT, weight)) + if abs(weight - 1.0) > 1e-3: + out.append((text, weight)) + # Longest first: a phrase and one of its own words can both be weighted, and applying the + # phrase first means the word's bias lands on top of it rather than instead of it. + out.sort(key=lambda p: -len(p[0])) + return out + + +def spans_from_ranges(ranges, weight): + """`_v2_find_phrase_token_ranges` gives (start, end) pairs; pair them with a weight.""" + return [(int(s), int(e), float(weight)) for s, e in ranges or [] if int(e) > int(s)] + + +def locate(tokenizer, prompt_text, weighted_phrases): + """Weighted phrases -> ([(start_tok, end_tok, weight)], prompt_token_count). + + Deliberately NOT `_v2_find_phrase_token_ranges`, which finds a phrase by cosine + similarity against the encoded sequence and so re-encodes every phrase it is given. On + H3 the text encoder is Qwen3-VL-32B, making that up to eight extra 32B forward passes + per generation for something the tokenizer answers exactly and for free. + + Matching is case-insensitive because phrase memory stores everything lowercased while + the prompt keeps the user's capitalisation. Every occurrence is weighted, not just the + first: a word repeated for emphasis meant it both times. + """ + text = str(prompt_text or "") + if not text: + return [], 0 + try: + enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) + offsets = list(enc["offset_mapping"]) + except Exception: # noqa: BLE001 — no offsets means no weighting, not a broken run + return [], 0 + total = len(offsets) + if not weighted_phrases: + return [], total + haystack = text.lower() + char_spans = [] + for phrase, weight in weighted_phrases: + needle = str(phrase or "").strip().lower() + if not needle: + continue + start = haystack.find(needle) + while start != -1: + char_spans.append((start, start + len(needle), float(weight))) + start = haystack.find(needle, start + len(needle)) + return token_spans_from_offsets(offsets, char_spans), total + + +def token_spans_from_offsets(offsets, char_spans): + """Character spans -> token index ranges, given a tokenizer's offset mapping.""" + out = [] + for c_start, c_end, weight in char_spans: + toks = [i for i, (a, b) in enumerate(offsets) if a < c_end and b > c_start] + if toks: + out.append((toks[0], toks[-1] + 1, float(weight))) + return out + + +def h3_tokenizer(clip): + """The raw HF tokenizer behind an H3 CLIP, or None. + + `clip.tokenizer` is a MiniMaxH3Tokenizer, which holds a Qwen3VLSDTokenizer under the + encoder's name, which holds the transformers tokenizer. Verified to return real offset + mappings for this vocabulary; anything else here returns None and weighting stays off. + """ + try: + inner = getattr(getattr(clip, "tokenizer", None), "qwen3vl_32b", None) + tok = getattr(inner, "tokenizer", None) + if tok is None: + return None + tok("probe", add_special_tokens=False, return_offsets_mapping=True) + return tok + except Exception: # noqa: BLE001 + return None diff --git a/samplers.py b/samplers.py index b966b1b1..aba01e4e 100644 --- a/samplers.py +++ b/samplers.py @@ -3912,6 +3912,7 @@ class FunPackLTXAVSceneChainSampler: # the toggles above which only work on Distilled Flow), so install/remove here rather # than via extra_options. Cheap to attempt (no-ops fast without the right metadata). _ba_handles = self._install_bounded_attention(model, latent, positive) if bounded_attention_enabled else [] + model = self._install_h3_token_weights(model, positive) try: sampled = comfy.sample.sample_custom( @@ -6591,6 +6592,53 @@ class FunPackLTXAVSceneChainSampler: w_idx = idx % w return (w_idx >= (w // 2)).long() + def _install_h3_token_weights(self, model, positive): + """Apply the rating-derived phrase emphasis Studio tagged onto the conditioning. + + Returns `model` untouched when there is no tag, which is every non-H3 run and every + H3 run before the first rating. Never raises: this is a refinement, not a + prerequisite. + + The bias is an attention MASK, and SLA routes masked calls to dense + (sla_attention.py) because a block-sparse kernel cannot carry a per-key bias. So a + weighted run is a dense run — said out loud, once, rather than discovered as a + slowdown. + """ + try: + meta = None + if isinstance(positive, list) and positive and isinstance(positive[0], (list, tuple)) \ + and len(positive[0]) >= 2 and isinstance(positive[0][1], dict): + meta = positive[0][1].get("funpack_h3_token_weights") + if not isinstance(meta, dict): + return model + spans = meta.get("spans") or [] + prompt_tokens = int(meta.get("prompt_tokens") or 0) + if not spans or prompt_tokens <= 0: + return model + cond = positive[0][0] + cond_len = int(cond.shape[1]) if hasattr(cond, "shape") and cond.dim() >= 2 else 0 + if cond_len <= 0: + return model + try: + from . import h3_token_weights as _tw + except ImportError: + import h3_token_weights as _tw + + patched = model.clone() + to = patched.model_options.get("transformer_options", {}).copy() + inner = to.get("optimized_attention_override") + to["optimized_attention_override"] = _tw.make_override( + spans, prompt_tokens, cond_len, inner=inner) + patched.model_options["transformer_options"] = to + print(f"[FunPackStudio] H3 phrase emphasis: {len(spans)} token span(s) biased in " + f"the packed attention stream. This is an attention mask, so SLA runs " + f"DENSE for this generation.") + return patched + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackStudio", "H3 phrase emphasis", _e, + "the rating's per-phrase weighting is NOT being applied") + 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 diff --git a/tests/test_h3_token_weights.py b/tests/test_h3_token_weights.py index 9c59eb73..eb24e8ed 100644 --- a/tests/test_h3_token_weights.py +++ b/tests/test_h3_token_weights.py @@ -210,3 +210,205 @@ def test_a_tokenizer_without_offsets_turns_weighting_off_rather_than_guessing(): raise ValueError("slow tokenizer") assert tw.token_spans(_NoOffsets(), "a bc d", [(2, 4, 1.5)]) == [] + + +# --- rating-derived weights ------------------------------------------------ + +def _unit(text, kind="phrase", **scores): + return {"text": text, "kind": kind, "effective_category_scores": scores} + + +def test_a_phrase_carrying_a_missing_axis_is_boosted(): + """This is the signal that already existed and had nowhere to go on H3.""" + got = tw.weights_from_memory([_unit("red scarf", details=0.8)], missing_axes=("details",)) + assert len(got) == 1 + text, weight = got[0] + assert text == "red scarf" and weight > 1.0 + + +def test_a_phrase_the_rating_is_neutral_about_is_omitted(): + """Returned at 1.0 it would still enter the bias tensor and cost a branch for nothing.""" + assert tw.weights_from_memory([_unit("red scarf", action=0.9)], + missing_axes=("details",)) == [] + + +def test_damping_is_off_unless_asked_for(): + """Suppressing a phrase the user typed is a knob going inert on its own.""" + units = [_unit("red scarf", details=0.1, quality=0.9)] + assert tw.weights_from_memory(units, missing_axes=("details",), wrong_axes=("quality",)) \ + != [] + assert tw.weights_from_memory(units, missing_axes=("details",), wrong_axes=("quality",), + damp=True) == [] + + +def test_wrong_appearance_damps_appearance_phrases_only_when_damping(): + units = [_unit("blonde hair", appearance=0.9, details=0.5)] + assert tw.weights_from_memory(units, missing_axes=("details",), + wrong_appearance=True, damp=True) == [] + assert tw.weights_from_memory(units, missing_axes=("details",), + wrong_appearance=True) != [] + + +def test_a_single_word_carries_less_than_a_whole_phrase(): + """Mirrors _v2_memory_kind_scale: phrase 1.0, ngram 0.62, token 0.24.""" + phrase = tw.weights_from_memory([_unit("a red scarf", "phrase", details=0.8)], + missing_axes=("details",))[0][1] + word = tw.weights_from_memory([_unit("scarf", "token", details=0.8)], + missing_axes=("details",))[0][1] + assert phrase > word > 1.0 + + +def test_the_learned_memory_overrides_the_run_s_own_scores(): + """The point of phrase_memory is that it accumulates across ratings.""" + units = [_unit("red scarf", details=0.0)] + memory = {"red scarf": {"kind": "phrase", "effective_category_scores": {"details": 0.9}}} + assert tw.weights_from_memory(units, memory, missing_axes=("details",)) != [] + + +def test_weights_are_clamped(): + got = tw.weights_from_memory([_unit("x", details=99.0)], missing_axes=("details",), + strength=50.0) + assert got[0][1] == tw.MAX_LEARNED_WEIGHT + + +def test_longer_phrases_are_applied_first(): + """A phrase and one of its own words can both be weighted; the word's bias should land + on top of the phrase's, not instead of it.""" + units = [_unit("scarf", "token", details=0.8), _unit("a red scarf", "phrase", details=0.8)] + assert [t for t, _ in tw.weights_from_memory(units, missing_axes=("details",))] \ + == ["a red scarf", "scarf"] + + +def test_malformed_units_are_skipped_not_fatal(): + units = [None, {}, {"text": "ok", "effective_category_scores": "not a dict"}, + {"text": "fine", "effective_category_scores": {"details": 0.5}}] + assert [t for t, _ in tw.weights_from_memory(units, missing_axes=("details",))] == ["fine"] + + +def test_ranges_become_weighted_spans(): + assert tw.spans_from_ranges([(4, 7), (9, 9)], 1.5) == [(4, 7, 1.5)] + + +# --- locating phrases in the prompt --------------------------------------- + +def test_locate_finds_a_phrase_case_insensitively(): + """Phrase memory stores everything lowercased; the prompt keeps the user's caps.""" + spans, total = tw.locate(_FakeTokenizer(), "A Red Scarf", [("red scarf", 1.5)]) + assert total == 11 + assert spans == [(2, 11, 1.5)] + + +def test_locate_weights_every_occurrence(): + spans, _ = tw.locate(_FakeTokenizer(), "cat and cat", [("cat", 1.5)]) + assert spans == [(0, 3, 1.5), (8, 11, 1.5)] + + +def test_locate_reports_the_prompt_length_even_with_nothing_to_weight(): + """The length is what places the bias in the packed sequence, so it is needed whether or + not any phrase matched.""" + assert tw.locate(_FakeTokenizer(), "a cat", []) == ([], 5) + + +def test_a_phrase_that_is_not_in_the_prompt_is_skipped(): + spans, _ = tw.locate(_FakeTokenizer(), "a dog", [("red scarf", 1.5)]) + assert spans == [] + + +def test_locate_survives_a_tokenizer_without_offsets(): + class _NoOffsets: + def __call__(self, *a, **kw): + raise ValueError("slow tokenizer") + + assert tw.locate(_NoOffsets(), "a cat", [("cat", 1.5)]) == ([], 0) + + +def test_the_real_h3_vocabulary_returns_usable_offsets(): + """The mechanism rests on this. Qwen's is a SLOW tokenizer, where offset mapping is not + guaranteed — checked against the actual vocabulary ComfyUI ships for H3.""" + from transformers import Qwen2Tokenizer + path = ("/Users/dex/Documents/ComfyUI/comfy/text_encoders/qwen25_tokenizer") + try: + tok = Qwen2Tokenizer.from_pretrained(path) + except Exception: + pytest.skip("ComfyUI's H3 tokenizer vocabulary is not present") + + spans, total = tw.locate(tok, "a fluffy cat", [("fluffy", 2.0)]) + + assert total == 3 + assert spans == [(1, 2, 2.0)] # 'fluffy' is exactly token 1 + + +def test_the_tokenizer_is_reached_through_the_clip_wrapper(): + import types + tok = _FakeTokenizer() + clip = types.SimpleNamespace( + tokenizer=types.SimpleNamespace(qwen3vl_32b=types.SimpleNamespace(tokenizer=tok))) + assert tw.h3_tokenizer(clip) is tok + assert tw.h3_tokenizer(types.SimpleNamespace(tokenizer=None)) is None + + +# --- sampler install ------------------------------------------------------- + +@pytest.fixture +def chain(): + import samplers + for name in dir(samplers): + obj = getattr(samplers, name) + if isinstance(obj, type) and hasattr(obj, "_install_h3_token_weights"): + return obj() + pytest.skip("no sampler exposes _install_h3_token_weights") + + +class _Patcher: + def __init__(self): + self.model_options = {} + + def clone(self): + c = _Patcher() + c.model_options = dict(self.model_options) + return c + + +def _tagged(spans=((0, 2, 2.0),), prompt_tokens=5, cond_len=9): + return [[torch.zeros(1, cond_len, 8), + {"funpack_h3_token_weights": {"spans": list(spans), + "prompt_tokens": prompt_tokens}}]] + + +def test_an_untagged_run_is_left_completely_alone(chain): + """Every non-H3 run, and every H3 run before the first rating.""" + model = _Patcher() + assert chain._install_h3_token_weights(model, [[torch.zeros(1, 9, 8), {}]]) is model + assert chain._install_h3_token_weights(model, None) is model + + +def test_a_tagged_run_gets_an_override(chain): + out = chain._install_h3_token_weights(_Patcher(), _tagged()) + assert "optimized_attention_override" in out.model_options["transformer_options"] + + +def test_the_existing_override_is_chained_not_replaced(chain): + """One slot. Taking it outright would silently discard SLA or the chosen backend.""" + model = _Patcher() + sentinel = lambda *a, **kw: "inner ran" + model.model_options["transformer_options"] = {"optimized_attention_override": sentinel} + + out = chain._install_h3_token_weights(model, _tagged()) + ov = out.model_options["transformer_options"]["optimized_attention_override"] + + assert ov is not sentinel + q = torch.zeros(1, 2, 100, 8) + assert ov(lambda *a, **kw: None, q, q, q, 2, skip_reshape=True) == "inner ran" + + +def test_the_original_model_options_are_not_mutated(chain): + model = _Patcher() + model.model_options["transformer_options"] = {} + chain._install_h3_token_weights(model, _tagged()) + assert "optimized_attention_override" not in model.model_options["transformer_options"] + + +def test_a_malformed_tag_does_not_break_the_run(chain): + model = _Patcher() + bad = [[torch.zeros(1, 9, 8), {"funpack_h3_token_weights": {"spans": "nonsense"}}]] + assert chain._install_h3_token_weights(model, bad) is model