digital-garbageandClaude Sonnet 5
8a6ab12ce6
Default 0.5 shift left phraseA dominating the video. [phraseA|phraseB:N] overrides strength per blend. Fixed a shared regex gap along the way: every weight/strength markup in this file (not just blend) silently rejected leading-dot decimals like ".8" -- fell through to literal text with no error, reproducing the exact "override did nothing" bug this feature exists to fix. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
659 lines
31 KiB
Python
659 lines
31 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
|
|
|
|
# A signed number a weight/strength may be written as, including a leading-dot shorthand
|
|
# (".8") that `-?\d+(?:\.\d+)?` alone would silently refuse -- refused text falls through to
|
|
# being read as literal phrase content instead of a weight, with nothing to say the number
|
|
# was never applied.
|
|
_NUM = r"-?(?:\d+\.\d*|\.\d+|\d+)"
|
|
|
|
# (phrase:1.3) / (phrase:-0.4). Escaped \( is left alone so a literal bracket still works.
|
|
_WEIGHTED = re.compile(r"(?<!\\)\(([^():]*?):\s*(" + _NUM + r")\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
|
|
|
|
|
|
# [phrase@2.0-3.5] / [phrase:1.5@2.0-3.5]: a TIMED phrase. Seconds from the scene start.
|
|
# Square brackets so a shortcut replacement's parentheses and colons can sit inside; only a
|
|
# `:number` directly before the `@t0-t1]` tail is a weight.
|
|
_TIMED = re.compile(
|
|
r"(?<!\\)\[([^\[\]]*?)(?::\s*(" + _NUM + r"))?\s*@\s*(\d+(?:\.\d+)?)\s*-\s*(\d+(?:\.\d+)?)\s*\]")
|
|
# H3's latent time grid: latent frame k spans FRAME_PER_TOKEN[k % 5] pixel frames at 24 fps.
|
|
FPS = 24
|
|
FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
|
|
|
|
|
|
# [phrase:1.5] - a weight with no window, in brackets so a shortcut replacement's own
|
|
# parentheses cannot break it (the reason timed phrases use brackets at all).
|
|
_BRACKET_WEIGHTED = re.compile(r"(?<!\\)\[([^\[\]@]*?):\s*(" + _NUM + r")\s*\]")
|
|
|
|
# [phraseA|phraseB]: BLEND. phraseA stays in the prompt Studio encodes (prefix/postfix and
|
|
# every other phrase are exactly what was typed); phraseB is a second sentence re-encoded
|
|
# with phraseA swapped out for it, so its own average hidden state can be pulled toward
|
|
# phraseA's. See `_v2_apply_h3_phrase_blend` (conditioning.py) for why this is a MEAN-VECTOR
|
|
# shift and not a per-token blend: H3's conditioning rows are Qwen's CONTEXTUAL hidden
|
|
# states (the same reason the module docstring above gives for why weighting cannot scale a
|
|
# token's embedding), so two differently-tokenized phrases have no per-position
|
|
# correspondence to blend token-by-token — only their own average position in context does.
|
|
#
|
|
# Colons, `@`, digits — anything except brackets and the `|` separator itself — are allowed
|
|
# freely in either phrase. `dog:1.5` or `dog@1-2` as an alt phrase is NOT stripped before its
|
|
# separate re-encode (there is no enclosing bracket left, once substituted into the sentence,
|
|
# for a second `parse_markup` pass to catch it), so Qwen reads it as literal punctuation in
|
|
# that phrase's own sentence — exactly the "reaches the encoder as punctuation" outcome the
|
|
# module docstring already describes for `(word:1.2)` failing to weight anything on H3.
|
|
# That is a quality footnote, not a leak: it degrades one phrase's own encoding, the same
|
|
# already-accepted trade-off `[cat (sitting):1.5]` gets elsewhere in this file (parens kept
|
|
# as literal pass-through text, not re-parsed). Earlier attempts to regex-blacklist "looks
|
|
# like leftover markup" shapes here (excluding `:`/`@` outright, then a trailing-weight
|
|
# check, then a trailing-timed-window check) each broke ordinary use of a colon or `@` in
|
|
# natural prose ("wide shot: character walks", "call me@work") while a DIFFERENT shape
|
|
# always found a way through anyway (`dog:5@2`, `dog@1—2` with an em dash, `@` followed by
|
|
# an unrelated dash many words later). Blacklisting a look is not a fixable boundary; not
|
|
# blacklisting at all is consistent with how the rest of this file already treats it.
|
|
#
|
|
# An OPTIONAL `:strength` on phraseB overrides `BLEND_STRENGTH_DEFAULT` for that one blend
|
|
# — `[cat|dog:0.8]`. This deliberately claims the same trailing "colon then a bare number"
|
|
# shape that would otherwise be ambiguous with literal phraseB text, the same way
|
|
# `_BRACKET_WEIGHTED` already claims that shape for its own phrase. Non-greedy so an
|
|
# ordinary internal colon ("shot: walks") is still just text — only a colon-number
|
|
# immediately before the closing `]` is read as a strength override.
|
|
_BLENDED = re.compile(r"(?<!\\)\[([^\[\]|]+)\|([^\[\]|]+?)(?::\s*(" + _NUM + r"))?\]")
|
|
# Default blend amount when no `:strength` is given: half the distance from phraseA's
|
|
# average hidden state toward phraseB's. Applies as a bare fraction, same scale as a typed
|
|
# strength (0 = no shift, 1 = fully phraseB's mean, room to overshoot past 1 if asked). Not
|
|
# clamped: 0 is a legitimate no-op, and a negative value pushes phraseA's rows AWAY from
|
|
# phraseB's mean instead of toward it. Deliberately left open, not validated -- same policy
|
|
# as every other weight/strength value in this file.
|
|
BLEND_STRENGTH_DEFAULT = 0.5
|
|
|
|
_MARKUP = re.compile(_TIMED.pattern + "|" + _BLENDED.pattern + "|" + _WEIGHTED.pattern
|
|
+ "|" + _BRACKET_WEIGHTED.pattern)
|
|
|
|
|
|
def parse_markup(text: str):
|
|
"""Strip ALL THREE prompt markups in one left-to-right pass.
|
|
|
|
-> (clean_text, weighted [(start, end, weight)], timed [(start, end, weight, t0, t1)],
|
|
blended [(start, end, alt_phrase, strength)]).
|
|
One pass, so every span list indexes the same clean text; stripping them one after the
|
|
other would leave the first list's offsets pointing into a string that no longer exists.
|
|
"""
|
|
if not text or ("(" not in text and "[" not in text):
|
|
return text, [], [], []
|
|
out, weighted, timed, blended, pos = [], [], [], [], 0
|
|
for m in _MARKUP.finditer(text):
|
|
out.append(text[pos:m.start()])
|
|
start = sum(len(p) for p in out)
|
|
if m.group(1) is not None: # [phrase(:w)?@t0-t1]
|
|
phrase = m.group(1)
|
|
t0, t1 = float(m.group(3)), float(m.group(4))
|
|
if t1 > t0 and phrase.strip():
|
|
timed.append((start, start + len(phrase),
|
|
float(m.group(2)) if m.group(2) else 1.0, t0, t1))
|
|
elif m.group(5) is not None: # [phraseA|phraseB(:strength)?]
|
|
phrase, alt = m.group(5), m.group(6)
|
|
strength = float(m.group(7)) if m.group(7) is not None else BLEND_STRENGTH_DEFAULT
|
|
if phrase.strip() and alt.strip():
|
|
blended.append((start, start + len(phrase), alt.strip(), strength))
|
|
else: # (phrase:w) or [phrase:w]
|
|
phrase = m.group(8) if m.group(8) is not None else m.group(10)
|
|
weight = m.group(9) if m.group(9) is not None else m.group(11)
|
|
weighted.append((start, start + len(phrase), float(weight)))
|
|
out.append(phrase)
|
|
pos = m.end()
|
|
out.append(text[pos:])
|
|
return "".join(out), weighted, timed, blended
|
|
|
|
|
|
def parse_timed(text: str):
|
|
"""`parse_markup` for callers that only want the windows; the clean text has ALL
|
|
markup removed (the model must never see any of it)."""
|
|
clean, _weighted, timed, _blended = parse_markup(text)
|
|
return clean, timed
|
|
|
|
|
|
def video_row_window(latent_t: int, frame_rows: int, t0: float, t1: float):
|
|
"""Seconds -> [row_start, row_end) inside the target video segment, or None.
|
|
|
|
Latent frame k covers pixel frames [cum[k], cum[k+1]); a window covers every latent
|
|
frame it touches. Rows are frame-major, `frame_rows` per latent frame.
|
|
"""
|
|
if latent_t <= 0 or frame_rows <= 0:
|
|
return None
|
|
cum = [0]
|
|
for k in range(latent_t):
|
|
cum.append(cum[-1] + FRAME_PER_TOKEN[k % 5])
|
|
f0, f1 = int(round(t0 * FPS)), int(round(t1 * FPS))
|
|
if f1 <= f0 or f0 >= cum[-1]:
|
|
return None
|
|
k0 = max(k for k in range(latent_t) if cum[k] <= f0)
|
|
k1 = min(latent_t, max(k for k in range(latent_t) if cum[k] < f1) + 1)
|
|
if k1 <= k0:
|
|
return None
|
|
return k0 * frame_rows, k1 * frame_rows
|
|
|
|
|
|
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, timed=None, latent_t: int = 0,
|
|
frame_rows: int = 0):
|
|
"""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.
|
|
|
|
`timed` = [(start_tok, end_tok, weight, t0, t1)] phrases that apply to a time window
|
|
of the target video only. A key bias is broadcast over every query, so a per-window
|
|
bias needs the queries split: softmax is per query row, so attending the video rows in
|
|
chunks (one per distinct window region) with a chunk-specific key bias is EXACT, and
|
|
costs only extra kernel launches. Inside its window a phrase gets `log(weight)`;
|
|
outside it the phrase is masked from the video rows (H3-World's single-egress rule:
|
|
an action span only influences its own interval). Text, cond/ref and audio rows keep
|
|
the untimed bias, so audio never sees a window.
|
|
"""
|
|
cache: dict = {}
|
|
timed = list(timed or [])
|
|
|
|
def _plan(seq_len, device, dtype):
|
|
base_bias = build_bias(spans, prompt_tokens, cond_len, seq_len, device, dtype, base=base)
|
|
if not timed:
|
|
return base_bias
|
|
import torch
|
|
n_video = latent_t * frame_rows
|
|
video_start = seq_len - n_video
|
|
if n_video <= 0 or video_start < cond_len:
|
|
return base_bias
|
|
origin = cond_len - prompt_tokens if base is None else base
|
|
keyed = []
|
|
for start, end, weight, t0, t1 in timed:
|
|
lo, hi = origin + max(0, start), origin + min(prompt_tokens, end)
|
|
win = video_row_window(latent_t, frame_rows, t0, t1)
|
|
if hi <= lo or hi > cond_len or win is None:
|
|
continue
|
|
keyed.append((lo, hi, bias_value(weight), video_start + win[0], video_start + win[1]))
|
|
if not keyed:
|
|
return base_bias
|
|
edges = {0, video_start, seq_len}
|
|
for _lo, _hi, _b, r0, r1 in keyed:
|
|
edges.update((r0, r1))
|
|
edges = sorted(e for e in edges if 0 <= e <= seq_len)
|
|
chunks = []
|
|
for a, b in zip(edges, edges[1:]):
|
|
if b <= a:
|
|
continue
|
|
bias = (base_bias.clone() if base_bias is not None
|
|
else torch.zeros(1, 1, 1, seq_len, device=device, dtype=dtype))
|
|
if a >= video_start:
|
|
for lo, hi, b_in, r0, r1 in keyed:
|
|
if r0 <= a and b <= r1:
|
|
bias[..., lo:hi] += b_in
|
|
else:
|
|
bias[..., lo:hi] = MASKED_BIAS
|
|
chunks.append((a, b, bias))
|
|
return chunks
|
|
|
|
def override(func, q, k, v, heads, mask=None, attn_precision=None,
|
|
skip_reshape=False, skip_output_reshape=False, **kwargs):
|
|
def run(m, qq=None):
|
|
qq = q if qq is None else qq
|
|
if inner is not None:
|
|
return inner(func, qq, k, v, heads, mask=m, attn_precision=attn_precision,
|
|
skip_reshape=skip_reshape,
|
|
skip_output_reshape=skip_output_reshape, **kwargs)
|
|
return func(qq, 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] = _plan(seq_len, k.device, k.dtype)
|
|
if on_apply is not None and cache[key] is not None:
|
|
on_apply(seq_len)
|
|
plan = cache[key]
|
|
if plan is None:
|
|
return run(mask)
|
|
if not isinstance(plan, list):
|
|
return run(plan if mask is None else mask + plan)
|
|
import torch
|
|
outs = []
|
|
for a, b, bias in plan:
|
|
outs.append(run(bias if mask is None else mask + bias,
|
|
q[:, :, a:b].contiguous()))
|
|
# [B, S, H*D] unless the caller kept heads split ([B, H, S, D]).
|
|
return torch.cat(outs, dim=2 if outs[0].ndim == 4 else 1)
|
|
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
|