Fix <|image_pad|> misalignment when images= and videos= coexist

Fixes microsoft/Mage#28.

Two independent defects made a single processor(images=..., videos=...) call
produce visual tensors that do not line up with the prompt placeholders:

1. codec_video_processing_mage_vl.rewrite_text_with_codec_positions used
text.find(VISION_START) / text.rfind(VISION_END), which spans from the
first vision block to the last. Any image block sitting between them was
wiped out along with the video block, so the images lost their
placeholders entirely and generate() raised
'Image features and image tokens do not match'. Now matches exactly one
<|vision_start|><|video_pad|><|vision_end|> block, mirroring what the
frames backend already did.

2. processing_mage_vl.__call__ ran the IMAGE PATH after the video branches
had already rewritten the video block into literal <|image_pad|> runs.
_expand_image_pads restarts from the start of the string on every
replace, so it consumed the video's placeholders. The tensor
concatenation was also unconditionally video-rows-then-image-rows, which
is only correct when the video precedes every image in the prompt. The
image path now runs first (while video placeholders are still
<|video_pad|>), both video branches emit per-visual slots, and the slots
are concatenated in recorded prompt order.

Reordering rows is safe for the vision tower: _build_cu_seqlens blocks
strictly per image_grid_thw row, so image and video rows never attend to
each other regardless of position.

Verification
------------
Single-modality paths are byte-identical to the stock processor
(torch.equal on input_ids / pixel_values / image_grid_thw / patch_positions):
codec video-only, frames video-only, 1 image, 2 images -> all True

Mixed images + video, stock vs patched:
codec [image, video] 3775/3776 ValueError -> 3776/3776 match
codec [video, image] 3775/3776 ValueError -> 3776/3776 match
codec [image, image, video] 5822/5824 ValueError -> 5824/5824 match
codec [video, image, image] 5822/5824 ValueError -> 5824/5824 match
frames all four orders totals matched by coincidence but the k-th
placeholder run was paired with the wrong
tensor rows; runs == grid rows is now True
element-by-element

Generation with unmodified weights, examples/dog.jpg + soccer-broadcast.mp4:
codec [video, image] 'A dog is in the reference photo, and the video shows
a sports broadcast with four commentators...'
frames [video, image] 'A dog is in the reference photo, and the video shows
a football match between England and Argentina.'

Not addressed: the codec branch still replicates one text per video, so
multiple videos in a single prompt remains unsupported, as before.

codec_video_processing_mage_vl.py

@@ -27,6 +27,7 @@ from __future__ import annotations
 import hashlib
 import json
 import os
+import re
 import shutil
 import subprocess
 import sys
@@ -49,6 +50,13 @@ from PIL import Image
 VISION_START = "<|vision_start|>"
 VISION_END = "<|vision_end|>"
 IMAGE_PAD = "<|image_pad|>"
+VIDEO_PAD = "<|video_pad|>"
+
+# Matches exactly one chat-template video placeholder block, so that image
+# blocks elsewhere in the prompt are left untouched.
+_VIDEO_BLOCK_RE = re.compile(
+    re.escape(VISION_START) + r"\s*" + re.escape(VIDEO_PAD) + r"\s*" + re.escape(VISION_END)
+)
 
 
 # ----------------------------------------------------------------- config
@@ -222,13 +230,13 @@ def rewrite_text_with_codec_positions(
     for timestamp, token_count in _timestamp_runs(patch_positions, fps, decimals):
         parts.extend([timestamp, VISION_START, IMAGE_PAD * token_count, VISION_END, "\n"])
     vision_text = "".join(parts)
-    first_vs, last_ve = text.find(VISION_START), text.rfind(VISION_END)
-    if first_vs == -1 or last_ve == -1:
+    match = _VIDEO_BLOCK_RE.search(text)
+    if match is None:
         return text
-    tail_start = last_ve + len(VISION_END)
+    tail_start = match.end()
     if tail_start < len(text) and text[tail_start] == "\n":
         tail_start += 1
-    return text[:first_vs] + vision_text + text[tail_start:]
+    return text[:match.start()] + vision_text + text[tail_start:]
 
 
 def drop_padding_canvases(

processing_mage_vl.py

@@ -237,6 +237,68 @@ class MageVLProcessor:
 
         out: dict = {}
 
+        # ---------------- VISUAL PLACEHOLDER ORDER ----------------
+        # Record the prompt order of every visual placeholder *before* any
+        # rewriting, so the visual tensors can be concatenated in exactly the
+        # order the model consumes the <|image_pad|> slots.
+        _pad_re = re.compile(re.escape(IMAGE_PAD) + "|" + re.escape(VIDEO_PAD))
+        order_per_text = [
+            ["image" if m.group(0) == IMAGE_PAD else "video" for m in _pad_re.finditer(s)]
+            for s in text
+        ]
+
+        # Per-visual (pixel_values, image_grid_thw, patch_positions) triples,
+        # kept separate until the final prompt-order assembly.
+        image_slots: List[tuple] = []
+        video_slots: List[tuple] = []
+
+        # ---------------- IMAGE PATH ----------------
+        # Runs *before* the video paths: at this point every <|image_pad|> in
+        # `text` still belongs to a real image, because video placeholders are
+        # still <|video_pad|>. Running it afterwards would let the image
+        # expansion consume placeholders the video rewrite had just emitted.
+        if images is not None:
+            if self.image_processor is None:
+                raise ValueError("images passed but no image_processor configured.")
+            image_outputs = self.image_processor(images=images, return_tensors="pt")
+            image_grid_thw = image_outputs["image_grid_thw"]
+
+            # Expand each <|image_pad|> placeholder to the number of merged tokens.
+            sms = self.spatial_merge_size
+            merge_factor = sms * sms
+            image_token_counts = (
+                (image_grid_thw[:, 0] * image_grid_thw[:, 1] * image_grid_thw[:, 2])
+                // merge_factor
+            ).tolist()
+            img_idx = 0
+
+            def _expand_image_pads(s: str) -> str:
+                nonlocal img_idx
+                while IMAGE_PAD in s:
+                    if img_idx >= len(image_token_counts):
+                        break
+                    n = int(image_token_counts[img_idx])
+                    s = s.replace(IMAGE_PAD, "<|placeholder|>" * n, 1)
+                    img_idx += 1
+                return s.replace("<|placeholder|>", IMAGE_PAD)
+
+            text = [_expand_image_pads(s) for s in text]
+
+            try:
+                from .video_processing_mage_vl import build_patch_positions
+            except ImportError:
+                from video_processing_mage_vl import build_patch_positions
+            image_pp = build_patch_positions(image_grid_thw, spatial_merge_size=sms)
+            offset = 0
+            for row in image_grid_thw:
+                n = int(row[0]) * int(row[1]) * int(row[2])
+                image_slots.append((
+                    image_outputs["pixel_values"][offset: offset + n],
+                    row.unsqueeze(0),
+                    image_pp[offset: offset + n],
+                ))
+                offset += n
+
         # ---------------- CODEC VIDEO BACKEND ----------------
         # Codec path: replaces the frame-sampling VideoProcessor entirely.
         # Each video -> N canvases + src_patch_position; we rewrite the
@@ -284,6 +346,7 @@ class MageVLProcessor:
             if len(rewritten_texts) != len(videos_list):
                 if len(rewritten_texts) == 1 and len(videos_list) >= 1:
                     rewritten_texts = rewritten_texts * len(videos_list)
+                    order_per_text = order_per_text * len(videos_list)
                 else:
                     raise ValueError(
                         f"codec video backend: got {len(rewritten_texts)} texts but {len(videos_list)} videos"
@@ -311,9 +374,7 @@ class MageVLProcessor:
                 all_grid_thw.append(image_grid_thw)
                 all_patch_positions.append(patch_positions)
 
-            out["pixel_values"] = torch.cat(all_pixel_values, dim=0)
-            out["image_grid_thw"] = torch.cat(all_grid_thw, dim=0)
-            out["patch_positions"] = torch.cat(all_patch_positions, dim=0)
+            video_slots.extend(zip(all_pixel_values, all_grid_thw, all_patch_positions))
             text = rewritten_texts
             # Codec branch handled the video. Suppress the frame-sampling block below.
             videos = None
@@ -412,68 +473,40 @@ class MageVLProcessor:
             # image_grid_thw is the only adjustment needed for the model's
             # forward to treat each frame as a separate image (matching the
             # multi-image inference path).
-            out["pixel_values"] = video_outputs["pixel_values_videos"]
             vgthw = video_outputs["video_grid_thw"]
-            expanded_rows = []
+            video_pv = video_outputs["pixel_values_videos"]
+            video_pp = video_outputs["patch_positions"]
+            offset = 0
             for row in vgthw:
                 T_v, H_v, W_v = int(row[0]), int(row[1]), int(row[2])
-                expanded_rows.extend([[1, H_v, W_v]] * T_v)
-            out["image_grid_thw"] = torch.tensor(expanded_rows, dtype=vgthw.dtype)
-            out["patch_positions"] = video_outputs["patch_positions"]
-
-        # ---------------- IMAGE PATH ----------------
-        if images is not None:
-            if self.image_processor is None:
-                raise ValueError("images passed but no image_processor configured.")
-            image_outputs = self.image_processor(
-                images=images, return_tensors="pt"
-            )
-            image_grid_thw = image_outputs["image_grid_thw"]
-
-            # Expand each <|image_pad|> placeholder to the number of merged tokens.
-            sms = self.spatial_merge_size
-            merge_factor = sms * sms
-            image_token_counts = (
-                (image_grid_thw[:, 0] * image_grid_thw[:, 1] * image_grid_thw[:, 2])
-                // merge_factor
-            ).tolist()
-            img_idx = 0
-
-            def _expand_image_pads(s: str) -> str:
-                nonlocal img_idx
-                while IMAGE_PAD in s:
-                    if img_idx >= len(image_token_counts):
-                        break
-                    n = int(image_token_counts[img_idx])
-                    s = s.replace(IMAGE_PAD, "<|placeholder|>" * n, 1)
-                    img_idx += 1
-                return s.replace("<|placeholder|>", IMAGE_PAD)
-
-            text = [_expand_image_pads(s) for s in text]
-
-            # If videos and images coexist, prefer concatenation of patch tensors.
-            if "pixel_values" in out:
-                out["pixel_values"] = torch.cat(
-                    [out["pixel_values"], image_outputs["pixel_values"]], dim=0
-                )
-                out["image_grid_thw"] = torch.cat(
-                    [out["image_grid_thw"], image_outputs["image_grid_thw"]], dim=0
-                )
-                # Build image patch_positions and concat.
-                from .video_processing_mage_vl import build_patch_positions
-                image_pp = build_patch_positions(
-                    image_outputs["image_grid_thw"], spatial_merge_size=sms
-                )
-                out["patch_positions"] = torch.cat(
-                    [out["patch_positions"], image_pp], dim=0
-                )
+                n = T_v * H_v * W_v
+                video_slots.append((
+                    video_pv[offset: offset + n],
+                    torch.tensor([[1, H_v, W_v]] * T_v, dtype=vgthw.dtype),
+                    video_pp[offset: offset + n],
+                ))
+                offset += n
+
+        # ---------------- ASSEMBLE VISUALS IN PROMPT ORDER ----------------
+        # The model consumes <|image_pad|> slots in prompt order and
+        # pixel_values in row order, so the visual tensors must be stacked in
+        # the order their placeholders appear in the prompt.
+        if image_slots or video_slots:
+            flat_order = [kind for per_text in order_per_text for kind in per_text]
+            n_img = sum(k == "image" for k in flat_order)
+            n_vid = sum(k == "video" for k in flat_order)
+            if n_img == len(image_slots) and n_vid == len(video_slots):
+                img_it, vid_it = iter(image_slots), iter(video_slots)
+                ordered = [next(img_it) if k == "image" else next(vid_it) for k in flat_order]
             else:
-                out["pixel_values"] = image_outputs["pixel_values"]
-                out["image_grid_thw"] = image_outputs["image_grid_thw"]
-                from .video_processing_mage_vl import build_patch_positions
-                out["patch_positions"] = build_patch_positions(
-                    image_outputs["image_grid_thw"], spatial_merge_size=sms
-                )
+                # Placeholder bookkeeping did not line up (e.g. a caller-supplied
+                # prompt that bypassed the chat template). Fall back to the legacy
+                # videos-then-images concatenation rather than failing hard.
+                ordered = list(video_slots) + list(image_slots)
+            out["pixel_values"] = torch.cat([s[0] for s in ordered], dim=0)
+            out["image_grid_thw"] = torch.cat([s[1] for s in ordered], dim=0)
+            out["patch_positions"] = torch.cat([s[2] for s in ordered], dim=0)
+
 
         # ---------------- VIDEO PATH FINAL EXPANSION ----------------
         # When `videos` was given (and possibly without `images`), the per-frame