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

#3
by kcz358 - opened
Microsoft org

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.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment