digital-garbage
b3c2dd3315
Studio has known which words to weight since 2.0 — _v2_concept_units_for_run splits every run into phrases, single words and n-grams; _v2_update_phrase_memory trains each against the rating's missing/wrong axes. The only consumer was _v2_build_attn2_patch, and H3 skips it: one packed self-attention stream, no cross-attention. So the learning ran on every generation and reached nothing. It now becomes an attention-logit bias on the prompt's token positions in the packed stream, which is the same statement the K/V x1.25 emphasis was making, at a place H3 has. Boost only by default. Located with the tokenizer's offset mapping (verified against ComfyUI's H3 vocabulary), not by re-encoding each phrase through Qwen3-VL-32B the way _v2_find_phrase_token_ranges does. The bias is a mask and SLA routes masked calls to dense, so a weighted run is a dense run. Said once, at install.
415 lines
15 KiB
Python
415 lines
15 KiB
Python
"""Attention-logit token weighting for MiniMax H3.
|
|
|
|
Qwen never applies (word:1.2) — ComfyUI's H3 tokenizer hardcodes every weight to 1.0 — so
|
|
the syntax has to be stripped before encoding and re-applied where the model can act on it:
|
|
the prompt's key positions in the packed self-attention stream.
|
|
"""
|
|
import math
|
|
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
sys.path.insert(0, ".")
|
|
import h3_token_weights as tw # noqa: E402
|
|
|
|
|
|
# --- parsing ---------------------------------------------------------------
|
|
|
|
def test_plain_text_is_returned_untouched():
|
|
assert tw.parse("a cat on a mat") == ("a cat on a mat", [])
|
|
|
|
|
|
def test_the_syntax_is_removed_from_the_text():
|
|
"""Left in, Qwen reads the brackets and digits as part of the sentence — worse than no
|
|
weighting at all."""
|
|
clean, spans = tw.parse("a (fluffy:1.4) cat")
|
|
assert clean == "a fluffy cat"
|
|
assert [(s, e) for s, e, _ in spans] == [(2, 8)]
|
|
assert clean[2:8] == "fluffy"
|
|
|
|
|
|
def test_spans_index_the_cleaned_text_not_the_original():
|
|
clean, spans = tw.parse("(red:1.2) car and a (blue:0.5) boat")
|
|
assert clean == "red car and a blue boat"
|
|
assert [clean[s:e] for s, e, _ in spans] == ["red", "blue"]
|
|
assert [w for _, _, w in spans] == [1.2, 0.5]
|
|
|
|
|
|
def test_a_negative_weight_is_parsed():
|
|
_, spans = tw.parse("(blurry:-0.4) shot")
|
|
assert spans[0][2] == -0.4
|
|
|
|
|
|
def test_an_escaped_bracket_is_left_alone():
|
|
clean, spans = tw.parse(r"a \(not a weight:1.2\) thing")
|
|
assert spans == []
|
|
assert clean == r"a \(not a weight:1.2\) thing"
|
|
|
|
|
|
# --- weight -> bias --------------------------------------------------------
|
|
|
|
def test_a_weight_of_one_is_no_bias_at_all():
|
|
assert tw.bias_value(1.0) == 0.0
|
|
|
|
|
|
def test_the_bias_is_the_log_of_the_weight():
|
|
"""softmax(logits + log w) multiplies that key's attention share by w. That identity is
|
|
the whole reason this is the right place to weight."""
|
|
assert tw.bias_value(2.0) == pytest.approx(math.log(2.0))
|
|
assert tw.bias_value(0.5) == pytest.approx(math.log(0.5))
|
|
|
|
|
|
def test_zero_masks_instead_of_producing_negative_infinity():
|
|
"""log(0) is -inf, and a query whose every key is -inf comes out NaN."""
|
|
assert tw.bias_value(0.0) == tw.MASKED_BIAS
|
|
assert math.isfinite(tw.bias_value(0.0))
|
|
|
|
|
|
def test_an_extreme_weight_is_clamped():
|
|
assert tw.bias_value(10_000.0) == tw.MAX_ABS_BIAS
|
|
|
|
|
|
# --- placing the bias in the packed sequence -------------------------------
|
|
|
|
def test_the_bias_lands_on_the_prompt_tail_of_the_conditioning():
|
|
"""References are prepended by the tokenizer, so the prompt occupies the END of the
|
|
conditioning block, and the conditioning block leads the packed sequence."""
|
|
bias = tw.build_bias([(0, 2, 2.0)], prompt_tokens=5, cond_len=9, seq_len=100,
|
|
device="cpu", dtype=torch.float32)
|
|
|
|
assert bias.shape == (1, 1, 1, 100)
|
|
assert torch.allclose(bias[0, 0, 0, 4:6], torch.full((2,), math.log(2.0)))
|
|
assert bias[0, 0, 0, :4].abs().sum() == 0 # the reference tokens are untouched
|
|
assert bias[0, 0, 0, 6:].abs().sum() == 0 # so is audio and video
|
|
|
|
|
|
def test_no_spans_means_no_tensor():
|
|
assert tw.build_bias([], 5, 9, 100, "cpu", torch.float32) is None
|
|
|
|
|
|
def test_a_span_reaching_past_the_prompt_is_dropped_not_wrapped():
|
|
assert tw.build_bias([(4, 99, 2.0)], prompt_tokens=5, cond_len=9, seq_len=100,
|
|
device="cpu", dtype=torch.float32) is not None
|
|
assert tw.build_bias([(50, 60, 2.0)], prompt_tokens=5, cond_len=9, seq_len=100,
|
|
device="cpu", dtype=torch.float32) is None
|
|
|
|
|
|
def test_a_sequence_shorter_than_the_conditioning_is_refused():
|
|
assert tw.build_bias([(0, 2, 2.0)], 5, 9, 4, "cpu", torch.float32) is None
|
|
|
|
|
|
# --- the override ----------------------------------------------------------
|
|
|
|
def _q(seq_len, heads=2, dim=8):
|
|
return torch.zeros(1, heads, seq_len, dim)
|
|
|
|
|
|
def test_the_override_biases_the_packed_self_attention():
|
|
seen = {}
|
|
|
|
def func(q, k, v, heads, mask=None, **kw):
|
|
seen["mask"] = mask
|
|
return q
|
|
|
|
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
|
|
ov(func, _q(100), _q(100), _q(100), 2, skip_reshape=True)
|
|
|
|
assert seen["mask"] is not None
|
|
assert seen["mask"][0, 0, 0, 4] == pytest.approx(math.log(2.0))
|
|
|
|
|
|
def test_an_existing_mask_is_added_to_not_replaced():
|
|
"""Stomping a real padding mask would corrupt attention rather than weight it."""
|
|
seen = {}
|
|
|
|
def func(q, k, v, heads, mask=None, **kw):
|
|
seen["mask"] = mask
|
|
return q
|
|
|
|
existing = torch.full((1, 1, 1, 100), -1.0)
|
|
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
|
|
ov(func, _q(100), _q(100), _q(100), 2, mask=existing, skip_reshape=True)
|
|
|
|
assert seen["mask"][0, 0, 0, 4] == pytest.approx(-1.0 + math.log(2.0))
|
|
assert seen["mask"][0, 0, 0, 0] == pytest.approx(-1.0)
|
|
|
|
|
|
def test_the_token_refiner_and_cross_shaped_calls_are_left_alone():
|
|
"""q and k are the same sequence only in the packed self-attention."""
|
|
seen = {}
|
|
|
|
def func(q, k, v, heads, mask=None, **kw):
|
|
seen["mask"] = mask
|
|
return q
|
|
|
|
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
|
|
ov(func, _q(100), _q(40), _q(40), 2, skip_reshape=True)
|
|
assert seen["mask"] is None
|
|
|
|
ov(func, _q(100), _q(100), _q(100), 2, skip_reshape=False)
|
|
assert seen["mask"] is None
|
|
|
|
|
|
def test_the_displaced_override_still_runs():
|
|
"""There is one override slot. Taking it without chaining would silently discard the
|
|
attention backend the user selected."""
|
|
calls = []
|
|
|
|
def inner(func, q, k, v, heads, mask=None, **kw):
|
|
calls.append(mask)
|
|
return q
|
|
|
|
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9, inner=inner)
|
|
ov(lambda *a, **kw: None, _q(100), _q(100), _q(100), 2, skip_reshape=True)
|
|
|
|
assert len(calls) == 1 and calls[0] is not None
|
|
|
|
|
|
def test_a_failure_inside_the_override_keeps_the_step():
|
|
"""Weighting is a refinement. It must never be the reason a render dies."""
|
|
def func(q, k, v, heads, mask=None, **kw):
|
|
return "ran"
|
|
|
|
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
|
|
broken = torch.zeros(1) # no .shape[2]; build_bias will raise
|
|
assert ov(func, broken, broken, broken, 2, skip_reshape=True) == "ran"
|
|
|
|
|
|
def test_the_bias_is_built_once_per_sequence_shape():
|
|
built = []
|
|
|
|
def func(q, k, v, heads, mask=None, **kw):
|
|
return q
|
|
|
|
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9,
|
|
on_apply=built.append)
|
|
for _ in range(4):
|
|
ov(func, _q(100), _q(100), _q(100), 2, skip_reshape=True)
|
|
|
|
assert built == [100]
|
|
|
|
|
|
# --- token spans -----------------------------------------------------------
|
|
|
|
class _FakeTokenizer:
|
|
"""Character-per-token, which makes the offset arithmetic checkable by hand."""
|
|
def __call__(self, text, add_special_tokens=False, return_offsets_mapping=False):
|
|
return {"offset_mapping": [(i, i + 1) for i in range(len(text))]}
|
|
|
|
|
|
def test_char_spans_become_token_ranges():
|
|
clean, spans = tw.parse("a (bc:1.5) d")
|
|
assert clean == "a bc d"
|
|
assert tw.token_spans(_FakeTokenizer(), clean, spans) == [(2, 4, 1.5)]
|
|
|
|
|
|
def test_a_tokenizer_without_offsets_turns_weighting_off_rather_than_guessing():
|
|
class _NoOffsets:
|
|
def __call__(self, *a, **kw):
|
|
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
|