digital-garbageandClaude Opus 5.5
f99944e5e4
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
47 lines
1.8 KiB
Python
47 lines
1.8 KiB
Python
"""Which rows of H3's joint sequence are the TARGET video and audio.
|
|
|
|
Upstream packs `[text | cond/ref | audio | video]` and emits one mod_segments
|
|
entry per stream, so target video is the LAST entry and target audio the one
|
|
before it. Matching by modality tag alone is wrong: tag 0 also marks the i2v
|
|
anchor pin, guide keyframes, reference media and the vision tokens inside the
|
|
text span, none of which a "video rows" feature may touch (fa04d6c audit).
|
|
|
|
None -- never an all-False mask -- when the layout cannot be proven (the token
|
|
refiner's text-only calls, an upstream layout change), so callers no-op
|
|
instead of guessing.
|
|
"""
|
|
|
|
import torch
|
|
|
|
_TAG = {"video": (0, -1), "audio": (2, -2)}
|
|
|
|
|
|
def target_rows(mod_segments, seq_len, device, stream):
|
|
if stream not in _TAG:
|
|
return None
|
|
tag, at = _TAG[stream]
|
|
segs = list(mod_segments or ())
|
|
if len(segs) < 2:
|
|
return None
|
|
a, b, row = segs[at]
|
|
try:
|
|
ok = bool(((row % 3) == tag).all()) if torch.is_tensor(row) else int(row) % 3 == tag
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
if not ok or not (0 <= a < b <= seq_len):
|
|
return None
|
|
mask = torch.zeros(seq_len, dtype=torch.bool, device=device)
|
|
mask[a:b] = True
|
|
return mask
|
|
|
|
|
|
def text_rows(model_kwargs, length):
|
|
"""[length] bool mask of the WORD rows in the text conditioning the model is
|
|
called with, or None when it can't be told. H3's text span carries picture
|
|
tokens too (reference images, tag 0); words are tag 1."""
|
|
payload = (model_kwargs or {}).get("minimax_payload")
|
|
tags = payload.get("text_token_tags") if isinstance(payload, dict) else None
|
|
if tags is None or not hasattr(tags, "view") or tags.numel() != length:
|
|
return None
|
|
return tags.view(-1) == 1
|