diff --git a/src/maxtext/multimodal/processor_qwen3_omni.py b/src/maxtext/multimodal/processor_qwen3_omni.py index 787dd82011..ab50497868 100644 --- a/src/maxtext/multimodal/processor_qwen3_omni.py +++ b/src/maxtext/multimodal/processor_qwen3_omni.py @@ -199,7 +199,13 @@ def maybe_pad_video_values_to_max_grid( def smart_resize( - height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280 + height: int, + width: int, + factor: int = 28, + min_pixels: int = 56 * 56, + max_pixels: int = 14 * 14 * 4 * 1280, + max_h_pixels: int | None = None, + max_w_pixels: int | None = None, ): """Rescales the image so that the following conditions are met: @@ -209,6 +215,10 @@ def smart_resize( 3. The aspect ratio of the image is maintained as closely as possible. + 4. (Optional) Neither dimension exceeds max_h_pixels / max_w_pixels. When either + cap is exceeded, both dimensions are scaled down proportionally and re-aligned + to 'factor'. + """ if max(height, width) / min(height, width) > MAX_RATIO: raise ValueError( @@ -224,6 +234,13 @@ def smart_resize( beta = math.sqrt(min_pixels / (height * width)) h_bar = math.ceil(height * beta / factor) * factor w_bar = math.ceil(width * beta / factor) * factor + # Apply per-dimension pixel caps, scaling down proportionally if either is exceeded. + if (max_h_pixels is not None and h_bar > max_h_pixels) or (max_w_pixels is not None and w_bar > max_w_pixels): + cap_h = max_h_pixels if max_h_pixels is not None else h_bar + cap_w = max_w_pixels if max_w_pixels is not None else w_bar + scale = min(cap_h / h_bar, cap_w / w_bar) + h_bar = min(max(factor, round(h_bar * scale / factor) * factor), (cap_h // factor) * factor) + w_bar = min(max(factor, round(w_bar * scale / factor) * factor), (cap_w // factor) * factor) return h_bar, w_bar @@ -461,6 +478,10 @@ def preprocess_video(video, config): nframes, channel, height, width = video.shape max_pixels = max(min(VIDEO_MAX_PIXELS, VIDEO_TOTAL_PIXELS / nframes * FRAME_FACTOR), int(VIDEO_MIN_PIXELS * 1.05)) + max_grid_h = getattr(config, "video_max_grid_h", None) + max_grid_w = getattr(config, "video_max_grid_w", None) + max_h_pixels = int(max_grid_h) * patch_size if max_grid_h is not None else None + max_w_pixels = int(max_grid_w) * patch_size if max_grid_w is not None else None resized_height_1, resized_width_1 = smart_resize( height, width, @@ -488,6 +509,8 @@ def preprocess_video(video, config): factor=patch_size * merge_size, min_pixels=VIDEO_MIN_PIXELS, max_pixels=VIDEO_MAX_PIXELS, + max_h_pixels=max_h_pixels, + max_w_pixels=max_w_pixels, ) # Second resize - process each channel separately to preserve float values @@ -780,7 +803,7 @@ def add_extra_tokens_for_qwen3_omni(tokens, config, processor_output): new_tokens.append(qwen_tokens.audio_pad) audio_data_idx += 1 - new_tokens.append(qwen_tokens.audio_pad) + new_tokens.append(qwen_tokens.audio_end) new_tokens.append(qwen_tokens.vision_end) video_idx += 1 diff --git a/tests/unit/qwen3_omni_layers_test.py b/tests/unit/qwen3_omni_layers_test.py index 4b3227ed7c..d27190dbd1 100644 --- a/tests/unit/qwen3_omni_layers_test.py +++ b/tests/unit/qwen3_omni_layers_test.py @@ -37,7 +37,10 @@ ) from maxtext.layers.decoders import deepstack_process from maxtext.layers.encoders import AudioEncoder -from maxtext.multimodal.processor_qwen3_omni import maybe_pad_video_values_to_max_grid +from maxtext.multimodal.processor_qwen3_omni import ( + maybe_pad_video_values_to_max_grid, + preprocess_video, +) from maxtext.models.qwen3 import ( Qwen3OmniAudioEncoder, Qwen3OmniAudioEncoderLayer, @@ -435,6 +438,48 @@ def test_patch_embed_padded_video_valid_outputs_match_unpadded(self): atol=5e-3, ) + def test_scale_to_fit_video_before_padding(self): + """Test that preprocess_video respects video_max_grid_h/w for wide/tall videos.""" + cfg = pyconfig.initialize( + ["", base_config_path], + model_name="qwen3-omni-30b-a3b", + video_max_grid_t=2, + video_max_grid_h=32, + video_max_grid_w=32, + patch_size_for_vit=16, + num_channels_for_vit=3, + temporal_patch_size_for_vit=2, + ) + + # Wide video (2 frames, 3 channels, 220×896): grid_w >> max_grid_w. + dummy_video_wide = np.ones((2, 3, 220, 896), dtype=np.float32) + dummy_video_wide_ratio = dummy_video_wide.shape[3] / dummy_video_wide.shape[2] + video_processed, video_grid_thw = preprocess_video(dummy_video_wide, cfg) + self.assertLessEqual(video_grid_thw[0, 1], cfg.video_max_grid_h) + self.assertLessEqual(video_grid_thw[0, 2], cfg.video_max_grid_w) + # Check the aspect ratio is mostly preserved after smart_resize scaling to fit max grid. + self.assertAlmostEqual(video_grid_thw[0, 2] / video_grid_thw[0, 1], dummy_video_wide_ratio, delta=0.5) + + # Verify maybe_pad_video_values_to_max_grid succeeds without ValueError. + video_values = np.reshape( + video_processed, + ( + 1, + cfg.num_channels_for_vit, + cfg.temporal_patch_size_for_vit * video_grid_thw[0, 0], + cfg.patch_size_for_vit * video_grid_thw[0, 1], + cfg.patch_size_for_vit * video_grid_thw[0, 2], + ), + ) + padded_values, padded_grid, _ = maybe_pad_video_values_to_max_grid( + video_values, + video_grid_thw, + cfg, + ) + self.assertEqual(padded_values.shape[3], cfg.video_max_grid_h * cfg.patch_size_for_vit) + self.assertEqual(padded_values.shape[4], cfg.video_max_grid_w * cfg.patch_size_for_vit) + np.testing.assert_array_equal(padded_grid, video_grid_thw) # Grid should not change + def test_patch_embed_is_jittable(self): """Test that patch embed is JIT-compilable."""