[Bugfix] Avoid enlarging MiniMax H3 reference images and videos (#8253)

Signed-off-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Co-authored-by: Roger Wang <hey@rogerw.io>

recipes/MiniMaxAI/MiniMax-H3.md

@@ -725,6 +725,9 @@ curl -sS -X POST "${API_URL}" \
 ### 3. Ref2VA: image-only, image/audio, or mixed references
 
 Run these requests against the combined service or a Ref2VA-only service.
+Reference images keep their original resolution, with dimensions aligned to
+32 pixels. No additional request or server setting is needed.
+
 Image-only Ref2VA omits `audio_reference`; adding one or more audio references
 is optional. The typed fields accept one object or an ordered JSON list.
 `audio_reference` accepts an HTTP(S) URL or a `data:` URL. In one terminal,
@@ -782,6 +785,9 @@ Run this request against the combined service. Repeat the
 `input_references` multipart field once per source video. H3 consumes the
 videos in form order and preserves their original soundtracks during
 conditioning.
+Small reference videos retain their original size; larger videos are scaled
+down to the reference canvas. Dimensions are aligned to 32 pixels. No
+additional request or server setting is needed.
 
 ```bash
 export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -138,7 +138,7 @@ def test_local_encoder_modes_preserve_validated_text_and_unified_handoff(task, t
         (image.height // 32) * (image.width // 32), 96
     )
     pipeline.audio_vae = Mock()
-    prompt = {"prompt": "A person waves."}
+    prompt: dict[str, Any] = {"prompt": "A person waves."}
     if task != "t2va":
         prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
     if text_key is not None:
@@ -1487,19 +1487,39 @@ def test_rainfusion_packed_padding_stays_mask_free_on_unaligned_lengths():
     assert metadata.extra["valid_kv_length"] == 5
 
 
-def test_reference_image_resize_contract():
+@pytest.mark.parametrize(
+    ("size", "expected"),
+    [
+        ((1344, 768), (1344, 768)),
+        ((640, 384), (640, 384)),
+        ((2688, 1536), (2688, 1536)),
+        ((1536, 2688), (1536, 2688)),
+        ((3072, 3072), (3072, 3072)),
+        ((1080, 1440), (1088, 1440)),
+        ((648, 392), (640, 384)),
+    ],
+)
+def test_reference_image_resize_contract(size, expected):
     from PIL import Image
 
     from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
         resolve_minimax_h3_reference_image_shape,
     )
 
-    assert resolve_minimax_h3_reference_image_shape(Image.new("RGB", (1080, 1440))) == (
-        2048,
-        2720,
+    assert resolve_minimax_h3_reference_image_shape(Image.new("RGB", size)) == expected
+
+
+def test_reference_image_resize_rejects_invalid_input():
+    from PIL import Image
+
+    from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+        resolve_minimax_h3_reference_image_shape,
     )
+
     with pytest.raises(ValueError, match="aspect ratio"):
         resolve_minimax_h3_reference_image_shape(Image.new("RGB", (100, 501)))
+    with pytest.raises(ValueError, match="dimensions"):
+        resolve_minimax_h3_reference_image_shape(Image.new("RGB", (128, 128)))
 
 
 def test_fl2va_supports_first_last_and_explicit_frame_index_contracts():
@@ -1678,13 +1698,26 @@ def test_encoder_forward_forwards_video_inputs():
     )
 
 
-def test_reference_video_shape_uses_h3_adapt_shape_policy():
+@pytest.mark.parametrize(
+    ("size", "expected"),
+    [
+        ((640, 384), (640, 384)),
+        ((384, 640), (384, 640)),
+        ((1280, 720), (1280, 704)),
+        ((1344, 768), (1344, 768)),
+        ((1920, 1152), (1280, 768)),
+        ((1152, 1920), (768, 1280)),
+        ((3844, 2160), (1344, 768)),
+        ((2160, 3844), (768, 1344)),
+        ((2048, 2048), (768, 768)),
+    ],
+)
+def test_reference_video_shape_caps_large_inputs_without_enlarging_small_inputs(size, expected):
     from vllm_omni.model_executor.models.minimax_h3.reference_video import (
         _reference_video_shape,
     )
 
-    assert _reference_video_shape(1280, 720) == (1344, 768)
-    assert _reference_video_shape(3844, 2160) == (1344, 768)
+    assert _reference_video_shape(*size) == expected
 
 
 def test_text_encoder_stub_constructs_without_group_or_weights():
@@ -2398,20 +2431,22 @@ def test_ref2va_reference_count_validation_preserves_client_error_metadata():
 
 
 @pytest.mark.parametrize(
-    ("case", "start_time", "expected_duration"),
+    ("case", "start_time", "expected_duration", "input_size", "expected_size"),
     [
-        ("R7", None, 10.0),
-        ("R8", 4.0, 6.0),
+        ("R7", None, 10.0, (3844, 2160), (1344, 768)),
+        ("R8", 4.0, 6.0, (640, 384), (640, 384)),
     ],
 )
-def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_time, expected_duration):
+def test_r7_r8_ref2va_video_segment_matrix(
+    monkeypatch, tmp_path, case, start_time, expected_duration, input_size, expected_size
+):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     source = tmp_path / f"{case}.mp4"
     source.touch()
     metadata = {
-        "width": 1280,
-        "height": 720,
+        "width": input_size[0],
+        "height": input_size[1],
         "fps": 24.0,
         "frame_count": 240,
         "duration": 10.0,
@@ -2440,6 +2475,8 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
     assert prepared[0]["start_time_seconds"] == pytest.approx(start_time or 0.0)
     assert transcode_calls[0][1]["duration_seconds"] == pytest.approx(expected_duration)
     assert transcode_calls[0][1]["target_frame_count"] == 209
+    assert (prepared[0]["width"], prepared[0]["height"]) == expected_size
+    assert (transcode_calls[0][1]["target_width"], transcode_calls[0][1]["target_height"]) == expected_size
 
 
 def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path):

tests/model_executor/stage_input_processors/test_minimax_h3.py

@@ -3,6 +3,7 @@
 """Regression tests for MiniMax H3's disaggregated encoder contract."""
 
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import Mock
 
 import numpy as np
@@ -166,6 +167,36 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
     assert len(media.audios) == 1
 
 
+@pytest.mark.parametrize(
+    ("image_size", "output_size"),
+    [
+        ((1344, 768), (448, 256)),
+        ((1344, 768), (1344, 768)),
+        ((640, 384), (1344, 768)),
+    ],
+)
+def test_ref2va_reference_sizing_is_shared_by_qwen_and_vae(image_size, output_size):
+    from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
+    from vllm_omni.model_executor.models.minimax_h3.encoder_processing import prepare_encoder_inputs
+
+    images = [Image.new("RGB", image_size, color=color) for color in ("red", "green", "blue", "white")]
+    prompt = {"prompt": "Use all four references.", "multi_modal_data": {"image": images}}
+    sampling = OmniDiffusionSamplingParams(height=output_size[1], width=output_size[0], extra_args={"task": "ref2va"})
+
+    local = prepare_encoder_inputs(prompt, sampling)
+    transformed = prepare_encoder_prompt(prompt, [sampling])
+    media = MiniMaxH3Encoder._media_input(transformed["additional_information"])
+    qwen_images = transformed["multi_modal_data"]["image"]
+    assert len(qwen_images) == len(media.images) == 4
+    assert (media.width, media.height) == output_size
+    for index, qwen_image in enumerate(qwen_images):
+        assert qwen_image.size == image_size
+        np.testing.assert_array_equal(np.asarray(qwen_image), media.images[index].numpy())
+        np.testing.assert_array_equal(np.asarray(qwen_image), np.asarray(local.images[index]))
+        torch.testing.assert_close(media.images[index], local.media.images[index])
+    assert [image.size for image in images] == [image_size] * 4
+
+
 def _mock_ref2va_video_with_audio(monkeypatch, *, duration_seconds: float) -> None:
     from vllm_omni.model_executor.models.minimax_h3 import encoder_processing
 
@@ -581,7 +612,7 @@ def test_full_payload_hook_skips_only_absent_output():
     ],
 )
 def test_encoder_handoff_rejects_invalid_full_payload(case, message, via_connector):
-    payload = _full_encoder_output()
+    payload: Any = _full_encoder_output()
     layout = payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY]
     if case == "not_mapping":
         payload = []
@@ -621,7 +652,7 @@ def test_encoder_handoff_cleans_media_without_mutating_prompt_or_stripping_outpu
     incoming = _full_encoder_output()
     retained_hidden = torch.ones(1)
     media = torch.zeros(2)
-    prompt = {
+    prompt: dict[str, Any] = {
         "prompt": "hello",
         "negative_prompt": "blur",
         "multi_modal_data": {"image": object(), "audio": object(), "video": object()},

vllm_omni/model_executor/models/minimax_h3/preprocessing.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Shared MiniMax H3 media normalization and Qwen presentation building.
 
 Builds the positive presentation token stream:
@@ -34,7 +35,6 @@ _VIDEO_TAG = 0
 
 MINIMAX_H3_OUTPUT_SHORT_EDGE = 768
 MINIMAX_H3_OUTPUT_MAX_PIXELS = 768 * 1344
-MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048
 MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32
 MINIMAX_H3_SUPPORTED_ASPECT_RATIOS = {
     "21:9": 21.0 / 9.0,
@@ -135,17 +135,16 @@ def resolve_minimax_h3_aspect_ratio(
 
 
 def resolve_minimax_h3_reference_image_shape(image: Image.Image) -> tuple[int, int]:
-    """Resize an H3 reference image to the official 2048-short-edge canvas."""
+    """Preserve reference resolution, rounding each axis to the 32-pixel grid."""
     width, height = image.size
     ratio = width / height
     if not 0.4 <= ratio <= 2.5:
         raise OmniClientError(f"reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}")
     if min(width, height) < 256 or max(width, height) > 5760:
         raise OmniClientError(f"reference image dimensions must be in [256, 5760] pixels, got {width}x{height}")
-    scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
     return (
-        _align_multiple(width * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE),
-        _align_multiple(height * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE),
+        _align_multiple(width, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE),
+        _align_multiple(height, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE),
     )
 
 

vllm_omni/model_executor/models/minimax_h3/reference_video.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 Ref2VA reference-video preparation."""
 
 from __future__ import annotations
@@ -65,9 +66,9 @@ def deserialize_prepared_reference_videos(value: str) -> tuple[str, list[dict[st
 MINIMAX_H3_FPS = 24.0
 MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS = 2.0
 MINIMAX_H3_QWEN_TEMPORAL_PATCH = 2
+MINIMAX_H3_CANVAS_MULTIPLE = 32
 MINIMAX_H3_BASE_SHORT_EDGE = 768
 MINIMAX_H3_MAX_PIXELS = 768 * 1344
-MINIMAX_H3_CANVAS_MULTIPLE = 32
 MINIMAX_H3_MIN_REFERENCE_DIMENSION = 256
 MINIMAX_H3_MAX_REFERENCE_DIMENSION = 5760
 MINIMAX_H3_MIN_REFERENCE_FPS = 23.976
@@ -290,6 +291,7 @@ def validate_reference_audio_waveforms(values: list[tuple[torch.Tensor, int]]) -
 
 
 def _reference_video_shape(width: int, height: int) -> tuple[int, int]:
+    """Fit large references to the video canvas without enlarging small inputs."""
     if (
         min(width, height) < MINIMAX_H3_MIN_REFERENCE_DIMENSION
         or max(width, height) > MINIMAX_H3_MAX_REFERENCE_DIMENSION
@@ -309,10 +311,14 @@ def _reference_video_shape(width: int, height: int) -> tuple[int, int]:
         scale = math.sqrt(MINIMAX_H3_MAX_PIXELS / area)
         target_width *= scale
         target_height *= scale
-    return (
-        _nearest_multiple(target_width, MINIMAX_H3_CANVAS_MULTIPLE),
-        _nearest_multiple(target_height, MINIMAX_H3_CANVAS_MULTIPLE),
-    )
+    canvas_width = _nearest_multiple(target_width, MINIMAX_H3_CANVAS_MULTIPLE)
+    canvas_height = _nearest_multiple(target_height, MINIMAX_H3_CANVAS_MULTIPLE)
+    if width * height < canvas_width * canvas_height:
+        return (
+            _nearest_multiple(width, MINIMAX_H3_CANVAS_MULTIPLE),
+            _nearest_multiple(height, MINIMAX_H3_CANVAS_MULTIPLE),
+        )
+    return canvas_width, canvas_height
 
 
 def _transcode_reference_video(