Match the reference's directed mask and sentence register

Checked against github.com/Danzer1xxxxChan/H3-World (code/abot/action_script.py and code/diffsynth_h3_action.patch).

- caption_for: port annotate_from_keys9. The camera clause is never omitted — `camera follows him` while moving, `camera holds steady` while still — and the pan adverb attaches to the pan clause rather than after a join. Previously 11 of 17 presets produced a sentence shape absent from training; the default example gave 30 of 37 latent frames such a shape.
- directed_attention: implement mask_mod's leak_out and leak_in. A sentence row as a key is now readable only by itself and its own frame's video rows; as a query it reads its own frame's video rows only. Previously every sentence read every other sentence in all 50 blocks, breaking the cut-vertex property the reference documents as required.
- token refiner: attend block-diagonally, one segment per sentence, matching the reference's split refiner_cu. The LoRA adapts refiner_blocks.{0,1}, which were trained under that segmentation.
- _masked_lse: zero all-masked rows. Once the mask bites, most rows see no sentence key, and softmax over an all -inf row is NaN.
- version the processor patch guard so a reload cannot keep a stale closure.

Remaining divergence, documented in the README: sentences are still encoded jointly through the shared conditioner rather than independently.

README.md

@@ -53,16 +53,20 @@ So a 124-frame request (37 latent frames) is conditioned on a prompt that looks
 
 ```
 A third-person view of a man walking through a city intersection...
-the man walks forward
-the man walks forward
+the man walks forward, camera follows him
+the man walks forward, camera follows him
 ...
-the man walks forward and strafes right
+the man walks forward and strafes right, camera follows him
 ...
 the man stands still, camera pans right sharply
 ```
 
-This Space builds that text from the key script for you — `caption_for()` in `app.py` — and shows you the resolved
-per-frame captions under **Conditioning** after every run.
+Every sentence ends in a camera clause — there is no "no camera" shape in the training data. With no camera key
+held the reference still states what the camera is doing: `camera follows him` while the character is moving,
+`camera holds steady` while he is not. `caption_for()` in `app.py` is a port of the reference's
+`annotate_from_keys9` (`code/abot/action_script.py`), sentence-for-sentence across all 130 reachable key
+combinations, so the strings the model receives here are the strings it was trained on. The resolved per-frame
+captions are shown under **Conditioning** after every run.
 
 ## The directed attention mask
 
@@ -83,19 +87,46 @@ off its flash kernel for the whole 21k-row sequence). The keys are split into th
 | **B** | everything after the text block (~99% of keys) | flash, unmasked |
 
 Each returns its output and its log-sum-exp; the three are recombined with the online-softmax identity, which is
-numerically identical to one masked softmax over the full row. Only **video** queries are restricted (frame *i*'s rows
-see only sentence *i*); the captions themselves still see everything — that is the "directed" part.
+numerically identical to one masked softmax over the full row.
+
+**What "directed" means.** The reference's `mask_mod` cuts the annotations' *outgoing* edges, and only those:
+
+* a sentence row, **as a key**, is readable by itself and by the video rows of its own latent frame — the static
+  prompt, the keyframe anchors, the audio and **every other sentence** are blocked from reading it;
+* a sentence row, **as a query**, reads the video rows of its own latent frame only. It still reads the static
+  prompt, the anchors and the audio freely — that is information flowing *into* the sentence, giving it scene
+  grounding, not a bypass around frame *i*.
+
+Blocking `A_j → A_k` is the load-bearing half. With it, `V_k` is a cut vertex between `A_k` and the rest of the
+sequence; without it `A_j`'s content reaches `A_k` first, and `V_k` reading `A_k` is reading `A_j` too.
 
 The mask is toggleable in **Advanced** so you can see the difference. With it off, the same script produces a video
 that drifts through a blur of every action at once.
 
-Two guards, both in `app.py`:
+Three guards, all in `app.py`:
 
 * The token spans are located by re-tokenizing prefixes of the prompt. If a BPE merge straddles a sentence boundary
   the spans would be wrong, so `build_conditioning_text()` verifies `cuts[-1] == total` and **refuses to mask** rather
   than mask the wrong rows.
-* `MiniMaxH3TokenRefinerBlock` runs the same attention module over the *text stream alone*. The processor detects that
-  (the sequence is too short to contain the video block) and falls straight through to the stock path.
+* Once the mask actually bites, most rows see *no* sentence key at all, so the masked region's log-sum-exp is `-inf`
+  for them. The merge treats that as weight zero, but `softmax` over an all-`-inf` row is `NaN` — so `_masked_lse`
+  zeroes those rows explicitly.
+* `MiniMaxH3TokenRefinerBlock` runs the same attention module over the *text stream alone*, which the processor
+  detects by its length (`cap_end` is the whole text length by construction). That stream is attended
+  **block-diagonally** — one segment for the scene prompt, one per sentence — matching the reference's split
+  `refiner_cu`. Running it as a single document lets sentences mix *before* the DiT backbone, and no mask on the DiT
+  side can close an upstream leak. The LoRA adapts `token_refiner.refiner_blocks.{0,1}`, so those deltas were trained
+  under this segmentation.
+
+### Known divergence from the reference
+
+One difference remains, and it is upstream of everything above. The reference encodes **each sentence
+independently** through the text encoder and concatenates the row blocks after the head, caching by string so a
+repeated sentence is bit-identical. This Space builds one prompt string and sends it to the shared
+`multimodalart/qwen3vl-conditioner`, so sentence *k* is contextualised against the image and every preceding
+sentence before the DiT ever sees it. Closing it needs a raw text-encode endpoint on the conditioner, which it does
+not currently expose. Related: the reference also gives annotation rows a mirrored positional offset
+(`origin = text_len - s_k[-1] - 1`), which is not reachable from this side of the split either.
 
 ## Architecture — why two Spaces
 

app.py

@@ -14,10 +14,16 @@ Two things beyond a plain ``load_lora_weights`` are needed, and this Space imple
 
 2. **A directed attention mask.**  MiniMax-H3 denoises one packed sequence
    ``[text | keyframe anchors | audio | video]`` under *full* self-attention, so by default every video
-   row sees every sentence and the per-frame binding is lost.  The patch below restricts each video
-   row of latent frame ``f`` to sentence ``f`` — and only among the sentences; the scene prompt, the
-   keyframe anchors, the audio rows and the whole video block stay fully visible.  It is *directed*:
-   the sentences themselves are not restricted.
+   row sees every sentence and the per-frame binding is lost.  The patch below cuts the sentences'
+   *outgoing* edges, exactly as the reference's ``mask_mod`` does: sentence ``f`` is readable only by
+   itself and by the video rows of latent frame ``f`` — the scene prompt, the keyframe anchors, the
+   audio and **every other sentence** are blocked from reading it — and sentence ``f`` as a query reads
+   frame ``f``'s video rows only, while still reading the prompt, anchors and audio freely.  That last
+   asymmetry is what "directed" means: information may flow *into* a sentence for scene grounding, but
+   nothing may route *around* frame ``f`` to reach it.  Blocking ``A_j -> A_k`` is the load-bearing
+   half; without it ``A_j``'s content reaches ``A_k`` first and ``V_k`` reading ``A_k`` reads ``A_j``
+   too.  The token refiner runs the same rule as one segment per sentence, since it sees the text
+   stream before the backbone does and no mask here could close a leak that happened there.
 
    Implementing that as a dense ``[S, S]`` mask would force SDPA off its flash kernel.  Instead the
    attention is split by key region and recombined with an online-softmax (log-sum-exp) merge, which
@@ -134,41 +140,76 @@ PRESETS: dict[str, str] = {
 }
 
 
+# The annotation vocabulary, ported verbatim from the author's own
+# ``code/abot/action_script.py`` (github.com/Danzer1xxxxChan/H3-World). Training and
+# inference in the reference run through the same `annotate_from_keys9`, so any drift here
+# is a drift away from the only sentence shapes the LoRA ever saw.
+MOTION = {"W": "walks forward", "S": "walks backward",
+          "A": "strafes left", "D": "strafes right"}
+MOTION_ORDER = ("W", "S", "A", "D")
+MOTION_IDLE = "stands still"
+PAN_KEY = {"J": "left", "L": "right"}
+TILT_KEY = {"I": "tilts down", "K": "tilts up"}
+CAMERA_IDLE = "holds steady"
+CAMERA_FOLLOW = "follows him"
+FAST_KEY = "F"
+
+
+def _purify(held: dict[str, bool], pairs) -> None:
+    """Cancel a pair of mutually exclusive keys that ended up both set.
+
+    The training pooling takes amax over a 4-frame window, so W-then-S inside one bin sets
+    both bits; the reference cancels them rather than emitting "walks forward and walks
+    backward". Reachable here through raw key entry (``WS``).
+    """
+    for a, b in pairs:
+        if held[a] and held[b]:
+            held[a] = held[b] = False
+
+
+def _motion_clause(held: dict[str, bool]) -> str:
+    _purify(held, (("W", "S"), ("A", "D")))
+    words = [MOTION[name] for name in MOTION_ORDER if held[name]]
+    return " and ".join(words) if words else MOTION_IDLE
+
+
+def _camera_clause(held: dict[str, bool], moving: bool) -> str:
+    """The camera clause is **never** omitted.
+
+    When no camera key is held the reference still says what the camera is doing —
+    ``follows him`` while the character moves, ``holds steady`` while he does not — so
+    every sentence in training ends in a camera clause. Dropping it (which this Space used
+    to do) puts the majority of latent frames in a sentence shape the model never saw.
+    """
+    _purify(held, (("J", "L"), ("I", "K")))
+    parts = []
+    for key, side in PAN_KEY.items():
+        if held[key]:
+            parts.append(f"pans {side} {'sharply' if held[FAST_KEY] else 'slowly'}")
+    for key, word in TILT_KEY.items():
+        if held[key]:
+            parts.append(word)
+    if parts:
+        return " and ".join(parts)
+    return CAMERA_FOLLOW if moving else CAMERA_IDLE
+
+
 def caption_for(keys: str, subject: str = DEFAULT_SUBJECT) -> str:
     """Render one latent frame's key state as the sentence H3-World was trained to read.
 
-    The register is the model card's: ``"the man walks forward, camera pans left sharply"`` — a
-    locomotion clause from ``W/A/S/D``, an optional camera clause from ``I/J/K/L``, and ``F`` turning
-    the camera adverb from ``slowly`` to ``sharply``.
+    ``"the man walks forward and strafes right, camera pans left sharply"`` — a locomotion
+    clause from ``W/A/S/D``, a camera clause from ``I/J/K/L`` with ``F`` turning its adverb
+    from ``slowly`` to ``sharply``, and the two joined by ``", camera "``.
+
+    ``subject`` stays configurable, but the camera clause's ``follows him`` is left exactly
+    as trained rather than agreed with it — fidelity to the register beats grammar here.
     """
-    held = set(keys.upper())
-
-    motion = []
-    if "W" in held:
-        motion.append("walks forward")
-    elif "S" in held:
-        motion.append("walks backward")
-    if "A" in held:
-        motion.append("strafes left")
-    elif "D" in held:
-        motion.append("strafes right")
-    body = " and ".join(motion) if motion else "stands still"
-
-    camera = []
-    if "J" in held:
-        camera.append("camera pans left")
-    elif "L" in held:
-        camera.append("camera pans right")
-    if "K" in held:
-        camera.append("camera tilts up")
-    elif "I" in held:
-        camera.append("camera tilts down")
-
-    sentence = f"{subject.strip() or DEFAULT_SUBJECT} {body}"
-    if camera:
-        adverb = "sharply" if "F" in held else "slowly"
-        sentence += ", " + ", ".join(camera) + f" {adverb}"
-    return sentence
+    upper = keys.upper()
+    held = {name: (name in upper) for name in KEYS}
+    # each clause purifies its own copy: cancelling J/L must not disturb W/S
+    motion = _motion_clause(dict(held))
+    camera = _camera_clause(dict(held), motion != MOTION_IDLE)
+    return f"{subject.strip() or DEFAULT_SUBJECT} {motion}, camera {camera}"
 
 
 def parse_script(script: str, num_slots: int) -> list[str]:
@@ -315,6 +356,8 @@ def directed_plan(num_text_tokens: int, num_prompt_tokens: int, cuts, num_latent
 # ── The directed-mask attention patch ────────────────────────────────────────
 
 DIRECTED: dict = {"plan": None}
+# Bumped whenever the patched processor's behaviour changes; see `install_directed_processor`.
+DIRECTED_PATCH_VERSION = "leak_out+leak_in+refiner_segments"
 
 
 def _lse_flash(q, k, v, scale):
@@ -370,7 +413,11 @@ def _masked_lse(query, key, value, mask, scale, chunk: int = 2048):
         if mask is not None:
             scores.masked_fill_(~mask[start:stop].unsqueeze(0).unsqueeze(0), float("-inf"))
         lse[:, :, start:stop] = torch.logsumexp(scores, dim=-1)
-        probs = torch.softmax(scores, dim=-1).to(v.dtype)
+        # A row with NO visible key in this region is normal once the mask actually bites
+        # (static text, the keyframe cond and the audio rows see no sentence at all). Its
+        # log-sum-exp is -inf, which the merge handles as weight zero — but softmax over
+        # an all -inf row is NaN, and NaN * 0 stays NaN, so it has to be zeroed here.
+        probs = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0).to(v.dtype)
         out[:, start:stop] = torch.matmul(probs, v).transpose(1, 2).to(out.dtype)
         del scores, probs, q
     return out, lse
@@ -396,35 +443,115 @@ def _merge_lse(out_a, lse_a, out_b, lse_b, chunk: int = 4096):
     return out_a, lse_a
 
 
-def directed_attention(query, key, value, plan):
-    """Full self-attention with each video row's view of the *sentences* narrowed to its own.
+def _refiner_segmented(query, key, value, plan, scale):
+    """Block-diagonal attention over the text stream: the head, then one sentence each.
 
-    Returns ``None`` when ``plan`` does not describe this call's sequence — the token refiner runs the
-    same attention module over the text stream alone, and it must keep the unmasked path.
+    MiniMax-H3's token refiner runs before the DiT backbone, over the text rows alone. The
+    reference splits its ``refiner_cu`` into one segment per sentence, and is explicit
+    about why: run as a single document, *"by the time A_k reaches the DiT backbone, its
+    representation already contains information from other action sentences, and no mask
+    on the DiT side can close that upstream leak."* The LoRA adapts
+    ``token_refiner.refiner_blocks.{0,1}``, so those deltas were trained under this
+    segmentation.
+    """
+    cap_start = plan["cap_start"]
+    segments = ([(0, cap_start)] if cap_start > 0 else []) + list(plan["spans"])
+    out = torch.zeros_like(query)
+    for start, stop in segments:
+        if stop <= start:
+            continue
+        q = query[:, start:stop]
+        piece, _ = _flash_lse(q, key[:, start:stop], value[:, start:stop], scale)
+        out[:, start:stop] = piece.to(out.dtype)
+    return out
+
+
+def directed_attention(query, key, value, plan):
+    """Self-attention under H3-World's directed mask.
+
+    Two shapes reach this function. The **token refiner**'s stream is the text rows alone
+    (``length == cap_end``) and gets segmented block-diagonally. The **packed sequence**
+    ``[text | keyframe anchors | audio | video]`` gets the directed mask proper, which is
+    two rules, both from the reference's ``mask_mod``:
+
+    * ``leak_out`` — an annotation row as a *key* is readable only by itself and by the
+      video rows of its own latent frame. Static text, the keyframe cond, the audio and
+      **every other annotation** are blocked from reading it. This is the load-bearing one:
+      the reference's own comment says *"'A_k never reads another A_j' is required for this
+      to hold — otherwise A_j's content would flow into A_k first, and V_k reading A_k
+      would be reading A_j too, and the cut-vertex property would break immediately."*
+    * ``leak_in`` — an annotation row as a *query* reads the video rows of its own latent
+      frame only. It still reads static text, the cond and the audio freely: that is
+      information flowing *into* the sentence, giving it scene grounding, not a bypass.
+
+    "Directed" therefore means only the annotations' outgoing edges are cut — not, as an
+    earlier version of this file had it, that the annotations are unrestricted as keys.
+
+    Returns ``None`` when ``plan`` does not describe this call's sequence at all.
     """
     length = query.shape[1]
     cap_start, cap_end = plan["cap_start"], plan["cap_end"]
     rows_per_frame = plan["rows_per_frame"]
     num_latent_frames = plan["num_latent_frames"]
+    scale = query.shape[-1] ** -0.5
+
+    # `cap_end` is the whole text length by construction (see `directed_plan`), so the
+    # refiner's own stream identifies itself by its length.
+    if length == cap_end:
+        return _refiner_segmented(query, key, value, plan, scale)
+
     # The video block is the tail of the packed sequence, so its start follows from the geometry
     # alone — no need to re-derive the keyframe or audio row counts here.
     video_start = length - num_latent_frames * rows_per_frame
     if video_start <= cap_end:
         return None
 
-    scale = query.shape[-1] ** -0.5
+    spans = plan["spans"]
+
+    def video_rows(index):
+        return slice(video_start + index * rows_per_frame, video_start + (index + 1) * rows_per_frame)
+
     out, lse = _flash_lse(query, key[:, cap_end:], value[:, cap_end:], scale)  # region B
     if cap_start > 0:  # region A
         out, lse = _merge_lse(out, lse, *_flash_lse(query, key[:, :cap_start], value[:, :cap_start], scale))
 
-    # region C — the sentences, the only place the mask bites.
-    visible = torch.ones(length, cap_end - cap_start, dtype=torch.bool, device=query.device)
-    for index, (start, stop) in enumerate(plan["spans"]):
-        rows = slice(video_start + index * rows_per_frame, video_start + (index + 1) * rows_per_frame)
-        visible[rows] = False
-        visible[rows, start - cap_start : stop - cap_start] = True
+    # region C — the sentence keys. Default DENIED: only the sentence's own rows and the
+    # video rows of its own latent frame may read it.
+    visible = torch.zeros(length, cap_end - cap_start, dtype=torch.bool, device=query.device)
+    for index, (start, stop) in enumerate(spans):
+        column = slice(start - cap_start, stop - cap_start)
+        visible[start:stop, column] = True          # same_ann:    A_k reads A_k
+        visible[video_rows(index), column] = True   # v_reads_own: V_k reads A_k
     out_c, lse_c = _masked_lse(query, key[:, cap_start:cap_end], value[:, cap_start:cap_end], visible, scale)
     out, _ = _merge_lse(out, lse, out_c, lse_c)
+
+    # `leak_in`: the annotation rows' own output is rebuilt from scratch, because region B
+    # above gave them every video row and they may only have their own frame's. They are a
+    # contiguous slice and there are only a few hundred of them, so recomputing is cheaper
+    # than trying to subtract the wrong contribution back out of a merged softmax.
+    q_ann = query[:, cap_start:cap_end]
+    out_ann, lse_ann = _flash_lse(                                   # cond + audio, unrestricted
+        q_ann, key[:, cap_end:video_start], value[:, cap_end:video_start], scale)
+    if cap_start > 0:                                                # static text, unrestricted
+        out_ann, lse_ann = _merge_lse(
+            out_ann, lse_ann, *_flash_lse(q_ann, key[:, :cap_start], value[:, :cap_start], scale))
+    own = torch.zeros(cap_end - cap_start, cap_end - cap_start, dtype=torch.bool, device=query.device)
+    for start, stop in spans:                                        # its own sentence only
+        local = slice(start - cap_start, stop - cap_start)
+        own[local, local] = True
+    out_ann, lse_ann = _merge_lse(
+        out_ann, lse_ann,
+        *_masked_lse(q_ann, key[:, cap_start:cap_end], value[:, cap_start:cap_end], own, scale))
+    # its own latent frame's video rows, one small attention per frame
+    for index, (start, stop) in enumerate(spans):
+        local = slice(start - cap_start, stop - cap_start)
+        rows = video_rows(index)
+        piece, piece_lse = _flash_lse(q_ann[:, local], key[:, rows], value[:, rows], scale)
+        merged, merged_lse = _merge_lse(
+            out_ann[:, local].clone(), lse_ann[:, :, local].clone(), piece, piece_lse)
+        out_ann[:, local] = merged
+        lse_ann[:, :, local] = merged_lse
+    out[:, cap_start:cap_end] = out_ann.to(out.dtype)
     return out
 
 
@@ -437,9 +564,13 @@ def install_directed_processor() -> None:
     """
     from diffusers.models.transformers import transformer_minimax_h3 as h3
 
-    if getattr(h3.MiniMaxH3AttnProcessor, "_h3world_directed", False):
+    # Version the guard, not just its presence: a reload re-executes this module but the
+    # patched class object survives it, so a bare boolean would keep the PREVIOUS closure
+    # — still pointing at the previous `directed_attention` — installed forever.
+    if getattr(h3.MiniMaxH3AttnProcessor, "_h3world_directed", None) == DIRECTED_PATCH_VERSION:
         return
-    base_call = h3.MiniMaxH3AttnProcessor.__call__
+    base_call = getattr(h3.MiniMaxH3AttnProcessor, "_h3world_base_call",
+                        h3.MiniMaxH3AttnProcessor.__call__)
 
     def __call__(self, attn, hidden_states, rotary_emb=None, attention_mask=None):
         plan = DIRECTED.get("plan")
@@ -463,8 +594,9 @@ def install_directed_processor() -> None:
         attended = attended.flatten(2, 3).type_as(query)
         return attn.to_out[1](attn.to_out[0](attended))
 
+    h3.MiniMaxH3AttnProcessor._h3world_base_call = base_call
     h3.MiniMaxH3AttnProcessor.__call__ = __call__
-    h3.MiniMaxH3AttnProcessor._h3world_directed = True
+    h3.MiniMaxH3AttnProcessor._h3world_directed = DIRECTED_PATCH_VERSION
 
 
 # ── LoRA loading (weight folding) ────────────────────────────────────────────