digital-garbage
3148a5f62a
"<Picture n>: " and "<Audio n>: " are emitted by the H3 tokenizer, not written by anyone. An AUDIO reference is the awkward one: it contributes its label and NO vision block, so it lands inside the same run of text tags as the prompt. Placement is now anchored at the END of that run rather than requiring the run to nearly equal the prompt, so any number of label tokens in front no longer shifts the bias — or refuses it outright, which is what the slack cap did. Protection follows the same boundary: the whole tensor is stashed and everything before the verified prompt start is restored, so a label is no more steerable than the picture is. Without a tokenizer it falls back to the end of the vision block — a label may stay steerable, the reference never does.
475 lines
22 KiB
Python
475 lines
22 KiB
Python
"""Real token weighting for MiniMax H3.
|
|
|
|
Qwen ignores `(word:1.2)`. ComfyUI's `MiniMaxH3Tokenizer.tokenize_with_weights` hardcodes
|
|
every weight to 1.0 and runs the raw HF tokenizer over the literal string, so the syntax
|
|
does not weight anything — it reaches the encoder as punctuation inside your sentence, and
|
|
`sd1_clip`'s weighting arm never runs because no weight is ever != 1.0.
|
|
|
|
Scaling the embeddings would not have worked anyway. Qwen's conditioning rows are
|
|
CONTEXTUAL hidden states — row i is "having read the prompt this far", not "the vector for
|
|
this word" — so there is no per-word magnitude to turn up.
|
|
|
|
What H3 does have is a single packed self-attention stream with the prompt's tokens IN it.
|
|
Every video token attends over those positions, so the weighting CLIP only ever
|
|
approximated can be done exactly: add `log(w)` to the attention logits at that phrase's key
|
|
positions. softmax(logits + log w) is softmax(logits) with that key's share of the
|
|
attention mass multiplied by w, which is what a weight is supposed to mean.
|
|
|
|
Costs one broadcast add per attention call. The catch is stated where it is installed: a
|
|
biased call cannot run on a block-sparse kernel, so SLA goes dense for the run.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import re
|
|
from typing import Optional
|
|
|
|
# (phrase:1.3) / (phrase:-0.4). Escaped \( is left alone so a literal bracket still works.
|
|
_WEIGHTED = re.compile(r"(?<!\\)\(([^():]*?):\s*(-?\d+(?:\.\d+)?)\s*\)")
|
|
|
|
# Below this a weight is treated as "remove entirely" rather than log(0) = -inf, which
|
|
# produces NaN the moment a query attends to nothing else.
|
|
MIN_WEIGHT = 1e-3
|
|
MASKED_BIAS = -30.0
|
|
# A guard, not a preference: bias is added to raw logits, and past this the phrase wins
|
|
# every query regardless of content, which reads as the prompt collapsing onto one word.
|
|
MAX_ABS_BIAS = 6.0
|
|
# How far the text run may exceed the measured prompt before placement is refused. A couple
|
|
# of tokens is a separator or a trimmed tag; more than that means a different string.
|
|
_BASE_SLACK = 4
|
|
# How far overlapping spans may reinforce each other, as a multiple of the STRONGEST single
|
|
# span at that position. A lone span is never touched — an explicit (word:2.0) means 2.0 —
|
|
# but three overlapping x1.5 spans stop at x1.75 instead of compounding to x3.4.
|
|
OVERLAP_HEADROOM = 1.5
|
|
|
|
|
|
def parse(text: str):
|
|
"""Strip weight syntax. -> (clean_text, [(start_char, end_char, weight), ...]).
|
|
|
|
Spans index into `clean_text`, so they survive the removal of the brackets. Returns the
|
|
text unchanged and an empty list when there is nothing to parse, which is the common
|
|
case and must cost nothing.
|
|
"""
|
|
if not text or "(" not in text:
|
|
return text, []
|
|
out = []
|
|
spans = []
|
|
pos = 0
|
|
for m in _WEIGHTED.finditer(text):
|
|
out.append(text[pos:m.start()])
|
|
phrase = m.group(1)
|
|
start = sum(len(p) for p in out)
|
|
out.append(phrase)
|
|
spans.append((start, start + len(phrase), float(m.group(2))))
|
|
pos = m.end()
|
|
out.append(text[pos:])
|
|
return "".join(out), spans
|
|
|
|
|
|
def token_spans(tokenizer, clean_text: str, char_spans):
|
|
"""Map character spans to token index ranges. -> [(start_tok, end_tok, weight), ...].
|
|
|
|
Uses the tokenizer's own offset mapping rather than re-tokenizing each phrase on its
|
|
own: "cat" tokenized alone and "cat" inside a sentence are not always the same ids, and
|
|
a span that is off by one token weights the wrong word.
|
|
"""
|
|
if not char_spans:
|
|
return []
|
|
try:
|
|
enc = tokenizer(clean_text, add_special_tokens=False, return_offsets_mapping=True)
|
|
offsets = enc["offset_mapping"]
|
|
except Exception: # noqa: BLE001 — a slow tokenizer has no offsets; weighting is off
|
|
return []
|
|
return token_spans_from_offsets(offsets, char_spans)
|
|
|
|
|
|
def bias_value(weight: float) -> float:
|
|
"""The additive logit bias a multiplicative weight corresponds to."""
|
|
if weight <= MIN_WEIGHT:
|
|
return MASKED_BIAS
|
|
return max(-MAX_ABS_BIAS, min(MAX_ABS_BIAS, math.log(weight)))
|
|
|
|
|
|
def build_bias(spans, prompt_tokens: int, cond_len: int, seq_len: int, device, dtype,
|
|
base: Optional[int] = None):
|
|
"""An additive [1, 1, 1, seq_len] attention bias, or None if nothing applies.
|
|
|
|
The prompt is the LAST thing the H3 tokenizer appends — reference pictures, audio and
|
|
video labels all go in front of it — so the prompt's tokens are the final
|
|
`prompt_tokens` rows of the conditioning block, and the conditioning block leads the
|
|
packed sequence. That places prompt token j at packed position
|
|
`cond_len - prompt_tokens + j`, without needing to know what the references added.
|
|
"""
|
|
import torch
|
|
|
|
if not spans or prompt_tokens <= 0 or cond_len <= 0 or seq_len < cond_len:
|
|
return None
|
|
if base is None:
|
|
base = cond_len - prompt_tokens
|
|
if base < 0 or base + prompt_tokens > cond_len:
|
|
return None
|
|
bias = None
|
|
strongest = 0.0
|
|
for start, end, weight in spans:
|
|
b = bias_value(weight)
|
|
if b == 0.0:
|
|
continue
|
|
lo, hi = base + max(0, start), base + min(prompt_tokens, end)
|
|
if hi <= lo or hi > cond_len:
|
|
continue
|
|
if bias is None:
|
|
bias = torch.zeros(1, 1, 1, seq_len, device=device, dtype=dtype)
|
|
bias[..., lo:hi] += b
|
|
strongest = max(strongest, abs(b))
|
|
if bias is not None and strongest > 0.0:
|
|
# Spans OVERLAP — a phrase and one of its own words both match the same tokens, and
|
|
# their biases ADD. That is intended, so the word reinforces the phrase. Unbounded it
|
|
# compounds: three overlapping x1.5 spans is an effective x3.4 on one token, which is
|
|
# far past anything weights_from_memory would ever return and reads as the prompt
|
|
# collapsing onto a few words.
|
|
#
|
|
# Bounded against the strongest SINGLE span rather than a constant, so a lone span is
|
|
# never touched: an explicit (word:2.0) still means 2.0.
|
|
limit = min(MAX_ABS_BIAS, strongest * OVERLAP_HEADROOM)
|
|
bias.clamp_(min=-limit, max=limit)
|
|
return bias
|
|
|
|
|
|
def make_override(spans, prompt_tokens: int, cond_len: int, inner=None, on_apply=None,
|
|
base: Optional[int] = None):
|
|
"""An `optimized_attention_override` that adds the bias, then delegates.
|
|
|
|
`inner` is the override this one displaces (SLA's, or the mask-safe one) so installing
|
|
weighting never silently discards the backend the user chose. The bias is cached per
|
|
(sequence length, device, dtype) — it is the same tensor on every call and every block.
|
|
"""
|
|
cache: dict = {}
|
|
|
|
def override(func, q, k, v, heads, mask=None, attn_precision=None,
|
|
skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
def run(m):
|
|
if inner is not None:
|
|
return inner(func, q, k, v, heads, mask=m, attn_precision=attn_precision,
|
|
skip_reshape=skip_reshape,
|
|
skip_output_reshape=skip_output_reshape, **kwargs)
|
|
return func(q, k, v, heads, mask=m, attn_precision=attn_precision,
|
|
skip_reshape=skip_reshape,
|
|
skip_output_reshape=skip_output_reshape, **kwargs)
|
|
|
|
try:
|
|
# Only the packed self-attention: k and q are the same sequence there. The
|
|
# 2-block token refiner and any cross-shaped call are left alone.
|
|
if not skip_reshape or k.ndim != 4 or q.shape[2] != k.shape[2]:
|
|
return run(mask)
|
|
seq_len = int(k.shape[2])
|
|
if seq_len < cond_len:
|
|
return run(mask)
|
|
key = (seq_len, k.device, k.dtype)
|
|
if key not in cache:
|
|
cache[key] = build_bias(spans, prompt_tokens, cond_len, seq_len,
|
|
k.device, k.dtype, base=base)
|
|
if on_apply is not None and cache[key] is not None:
|
|
on_apply(seq_len)
|
|
bias = cache[key]
|
|
if bias is None:
|
|
return run(mask)
|
|
return run(bias if mask is None else mask + bias)
|
|
except Exception: # noqa: BLE001 — weighting must never cost the step
|
|
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.
|
|
#
|
|
# The ceiling is Studio's OWN emphasis constant. `_v2_build_attn2_patch` has boosted an
|
|
# emphasised phrase's K/V by exactly 1.25 on LTX since long before this existed, and that
|
|
# path works. There is no reason a different mechanism for the same intent should push
|
|
# harder, so 1.25 is the most any phrase can reach here too.
|
|
EMPHASIS_CEILING = 0.25
|
|
MIN_LEARNED_WEIGHT = 0.25
|
|
MAX_LEARNED_WEIGHT = 1.0 + EMPHASIS_CEILING
|
|
# The top of _v2_auto_strength's own clamp. Used to read that strength as a 0..1 fraction of
|
|
# how much evidence there is, rather than re-deriving a confidence scale here.
|
|
MAX_AUTO_STRENGTH = 0.085
|
|
|
|
|
|
def strength_from_auto(auto_strength):
|
|
"""Studio's `_v2_auto_strength` -> how hard to emphasise, in 0..EMPHASIS_CEILING.
|
|
|
|
That function already answers "how much evidence is there": it ramps with bad_streak and
|
|
the reward EMA off a 0.030 base and clamps to [0.008, 0.085], so a SINGLE bad rating
|
|
lands near the middle rather than at the top. Reusing it means one Awful nudges a phrase
|
|
to about 1.13 and a sustained bad streak approaches 1.19 — never the full 1.25, and never
|
|
a jump to maximum off one data point.
|
|
"""
|
|
try:
|
|
frac = float(auto_strength) / MAX_AUTO_STRENGTH
|
|
except (TypeError, ValueError, ZeroDivisionError):
|
|
return EMPHASIS_CEILING * 0.5
|
|
return EMPHASIS_CEILING * max(0.0, min(1.0, frac))
|
|
|
|
_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=EMPHASIS_CEILING * 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
|
|
out.append((text, delta * scale(entry.get("kind", unit.get("kind")))))
|
|
if not out:
|
|
return []
|
|
# RELATIVE, not absolute. Category scores are confidences in the 0..1 range and a phrase
|
|
# rarely carries much of any one axis, so `1 + strength * score` put a whole prompt at
|
|
# x1.03 — measurably applied and doing nothing. Normalising against the strongest
|
|
# candidate makes `strength` mean "how much the best phrase gets", and keeps the ordering
|
|
# the memory learned. How big `strength` itself may be is decided by `strength_from_auto`
|
|
# off Studio's existing evidence ramp, NOT by the size of one rating.
|
|
top = max(score for _, score in out)
|
|
if top <= 0.0:
|
|
return []
|
|
scored = []
|
|
for text, score in out:
|
|
weight = 1.0 + float(strength) * (score / top)
|
|
weight = max(MIN_LEARNED_WEIGHT, min(MAX_LEARNED_WEIGHT, weight))
|
|
if abs(weight - 1.0) > 1e-3:
|
|
scored.append((text, weight))
|
|
# Strongest first, so a cap downstream keeps the phrases the rating cared MOST about —
|
|
# taking the longest instead meant an arbitrary eight out of ninety-four.
|
|
scored.sort(key=lambda p: -p[1])
|
|
return scored
|
|
|
|
|
|
def order_for_application(weighted):
|
|
"""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.
|
|
Separate from the ranking above, which decides WHICH phrases survive a cap."""
|
|
return sorted(weighted or [], key=lambda p: -len(p[0]))
|
|
|
|
|
|
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 prompt_region(token_tags, cond_len):
|
|
"""(start, end) of the PROMPT inside an H3 conditioning, or None.
|
|
|
|
An r2v conditioning is laid out reference-first: the `<Picture n>: ` label, the vision
|
|
block, then the encoded prompt. `minimax_token_tags` marks a vision block 0 and text 1,
|
|
so the prompt is the trailing unbroken run of 1s and everything before it belongs to the
|
|
reference.
|
|
|
|
That boundary is the only thing Studio needs. Rows before it are the picture and are not
|
|
Studio's to touch; rows from it on are the prompt and are.
|
|
"""
|
|
tags = _as_list(token_tags)
|
|
if not tags or cond_len <= 0:
|
|
return None
|
|
end = min(len(tags), cond_len)
|
|
while end > 0 and int(tags[end - 1]) != 1:
|
|
end -= 1 # trailing non-text; should not happen, cheap to allow
|
|
start = end
|
|
while start > 0 and int(tags[start - 1]) == 1:
|
|
start -= 1
|
|
return (start, end) if end > start else None
|
|
|
|
|
|
def prompt_base(token_tags, cond_len, prompt_tokens):
|
|
"""Index of the prompt's first token inside the conditioning block, or None.
|
|
|
|
The fallback is arithmetic: the tokenizer appends the prompt LAST, so it is the final
|
|
`prompt_tokens` rows. That holds only if the text tokenized here is exactly the text that
|
|
was encoded, and a run whose conditioning is 367 rows against 368 modality tags shows how
|
|
easily that bookkeeping drifts.
|
|
|
|
`minimax_token_tags` removes the assumption. Tags are 1 for text and 0 for a vision
|
|
block, and references are `<Picture N>: ` + vision BEFORE the prompt, so the prompt is
|
|
the last unbroken run of 1s. Using it means a reference or a trimmed tag shifts the
|
|
prompt and the spans follow.
|
|
"""
|
|
if prompt_tokens <= 0 or cond_len <= 0:
|
|
return None
|
|
region = prompt_region(token_tags, cond_len)
|
|
if region:
|
|
run_start, run_end = region
|
|
run = run_end - run_start
|
|
# The prompt is the tokenizer's LAST segment, so the tail of the text run is the
|
|
# prompt — but only if the text measured here is the text that was encoded. When a
|
|
# reference node owns the conditioning it encoded the whole combined prompt while
|
|
# this may be one scene's, and taking the tail then would bias a window of the wrong
|
|
# words. A run that does not match the measurement is not something to place by.
|
|
if run == prompt_tokens:
|
|
return run_start
|
|
if run > prompt_tokens and run - prompt_tokens <= _BASE_SLACK:
|
|
return run_end - prompt_tokens
|
|
return None
|
|
base = cond_len - prompt_tokens
|
|
return base if base >= 0 else None
|
|
|
|
|
|
def _as_list(tags):
|
|
if tags is None:
|
|
return []
|
|
try:
|
|
return tags.reshape(-1).tolist()
|
|
except AttributeError:
|
|
return list(tags) if isinstance(tags, (list, tuple)) else []
|
|
|
|
|
|
def choose_encoded_text(tokenizer, candidates, token_tags, cond_len):
|
|
"""Pick the candidate string whose token count matches the conditioning's text run.
|
|
|
|
A wired conditioning was encoded by ANOTHER node, from a string the editor expanded for
|
|
it — Studio keeps the raw prompt and expands per scene, so its copy can differ. Guessing
|
|
which is not necessary: `minimax_token_tags` says how many text positions the tensor
|
|
actually holds, so each candidate can be CHECKED against it and the one that fits wins.
|
|
|
|
Returns (text, token_count, base) or (None, 0, None) when nothing fits. `base` is
|
|
measured back from the END of the text run, so labels sitting in front of the prompt —
|
|
an audio reference's "<Audio n>: " has no vision block to separate it — do not shift it.
|
|
"""
|
|
region = prompt_region(token_tags, cond_len)
|
|
if not region:
|
|
return None, 0, None
|
|
run_start, run_end = region
|
|
run = run_end - run_start
|
|
best = None
|
|
for text in candidates:
|
|
text = str(text or "").strip()
|
|
if not text:
|
|
continue
|
|
try:
|
|
n = len(tokenizer(text, add_special_tokens=False,
|
|
return_offsets_mapping=True)["offset_mapping"])
|
|
except Exception: # noqa: BLE001
|
|
continue
|
|
# Anchored at the END, and deliberately not capped: the prompt is the LAST thing the
|
|
# tokenizer appends, but the text in front of it is not always a vision block. An
|
|
# AUDIO reference contributes "<Audio n>: " and no vision block at all, so it lands
|
|
# inside the same run of text tags as the prompt. Requiring the run to nearly equal
|
|
# the prompt would refuse those runs outright; measuring back from the end places the
|
|
# prompt correctly however many label tokens precede it.
|
|
gap = run - n
|
|
if gap >= 0 and (best is None or gap < best[0]):
|
|
best = (gap, text, n)
|
|
if best is None:
|
|
return None, 0, None
|
|
_gap, text, n = best
|
|
return text, n, run_end - n
|
|
|
|
|
|
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
|