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

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