diff --git a/src/maxtext/multimodal/processor_qwen3_omni.py b/src/maxtext/multimodal/processor_qwen3_omni.py index 787dd82011..a73e770cf5 100644 --- a/src/maxtext/multimodal/processor_qwen3_omni.py +++ b/src/maxtext/multimodal/processor_qwen3_omni.py @@ -227,6 +227,45 @@ def smart_resize( return h_bar, w_bar +def scale_to_fit_video_grid( + height: int, + width: int, + max_grid_h: int | None, + max_grid_w: int | None, + patch_size: int, + merge_size: int, +) -> tuple[int, int]: + """Rescales height and width proportionally if they exceed the maximum grid dimensions. + + Args: + height: Image height in pixels. + width: Image width in pixels. + max_grid_h: Maximum allowed height in grid units (patches), or None. + max_grid_w: Maximum allowed width in grid units (patches), or None. + patch_size: ViT patch size in pixels. + merge_size: Spatial merge size in patches. + + Returns: + Tuple of (scaled_height, scaled_width) in pixels, divisible by factor = patch_size * merge_size. + """ + if max_grid_h is None and max_grid_w is None: + return height, width + if max_grid_h is None or max_grid_w is None: + raise ValueError("video_max_grid_h and video_max_grid_w must be set together or both None.") + + factor = patch_size * merge_size + max_h_px = int(max_grid_h) * patch_size + max_w_px = int(max_grid_w) * patch_size + + if height <= max_h_px and width <= max_w_px: + return height, width + + scale = min(max_h_px / height, max_w_px / width) + scaled_h = max(factor, math.floor(height * scale / factor) * factor) + scaled_w = max(factor, math.floor(width * scale / factor) * factor) + return scaled_h, scaled_w + + def pre_process_qwen3_image(image: np.ndarray | list[np.ndarray], config, force_resize=None): """Performs a bi-linear resize (with anti-aliasing) and normalizes the image.""" patch_size = config.patch_size_for_vit @@ -468,6 +507,14 @@ def preprocess_video(video, config): min_pixels=VIDEO_MIN_PIXELS, max_pixels=max_pixels, ) + resized_height_1, resized_width_1 = scale_to_fit_video_grid( + resized_height_1, + resized_width_1, + max_grid_h=getattr(config, "video_max_grid_h", None), + max_grid_w=getattr(config, "video_max_grid_w", None), + patch_size=patch_size, + merge_size=merge_size, + ) # First resize - using PIL to match HuggingFace behavior resized_frames = [] @@ -489,6 +536,14 @@ def preprocess_video(video, config): min_pixels=VIDEO_MIN_PIXELS, max_pixels=VIDEO_MAX_PIXELS, ) + resized_height_2, resized_width_2 = scale_to_fit_video_grid( + resized_height_2, + resized_width_2, + max_grid_h=getattr(config, "video_max_grid_h", None), + max_grid_w=getattr(config, "video_max_grid_w", None), + patch_size=patch_size, + merge_size=merge_size, + ) # Second resize - process each channel separately to preserve float values final_frames = [] @@ -780,7 +835,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..c41cef5f2f 100644 --- a/tests/unit/qwen3_omni_layers_test.py +++ b/tests/unit/qwen3_omni_layers_test.py @@ -37,7 +37,11 @@ ) 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, + scale_to_fit_video_grid, +) from maxtext.models.qwen3 import ( Qwen3OmniAudioEncoder, Qwen3OmniAudioEncoderLayer, @@ -435,6 +439,69 @@ def test_patch_embed_padded_video_valid_outputs_match_unpadded(self): atol=5e-3, ) + def test_scale_to_fit_video_before_padding(self): + """Test scale-to-fit before padding for wide/tall videos exceeding video_max_grid_h/w.""" + patch_size = self.config.patch_size_for_vit # 16 + merge_size = self.config.spatial_merge_size_for_vit # 2 + max_grid_h = 32 + max_grid_w = 32 + # Maximum pixel dimensions allowed: max_grid * patch_size = 32 * 16 = 512 + max_h_px = max_grid_h * patch_size # 512 + max_w_px = max_grid_w * patch_size # 512 + + # 1. Wide video (224×896): grid_w=56 > 32; scale=512/896 → scaled to 128×512 px (grid 8×32). + scaled_h, scaled_w = scale_to_fit_video_grid( + 224, + 896, + max_grid_h=max_grid_h, + max_grid_w=max_grid_w, + patch_size=patch_size, + merge_size=merge_size, + ) + self.assertEqual(scaled_h, 128) # 224 → 128 px (grid_h 14 → 8) + self.assertEqual(scaled_w, 512) # 896 → 512 px (grid_w 56 → 32) + self.assertLessEqual(scaled_h, max_h_px) # 128 ≤ 512 + self.assertLessEqual(scaled_w, max_w_px) # 512 ≤ 512 + self.assertLessEqual(scaled_h // patch_size, max_grid_h) # 8 ≤ 32 + self.assertLessEqual(scaled_w // patch_size, max_grid_w) # 32 ≤ 32 + self.assertEqual(scaled_h % (patch_size * merge_size), 0) # divisible by 32 + self.assertEqual(scaled_w % (patch_size * merge_size), 0) # divisible by 32 + + # 2. Test preprocess_video with a wide video array and verify grid_thw fits within max_grid + config_video_scale = pyconfig.initialize( + ["", base_config_path], + model_name="qwen3-omni-30b-a3b", + video_max_grid_t=2, + video_max_grid_h=max_grid_h, + video_max_grid_w=max_grid_w, + ) + # Dummy video: 2 frames, 3 channels, height 224, width 896 + dummy_video = np.ones((2, 3, 224, 896), dtype=np.float32) + video_processed, video_grid_thw = preprocess_video(dummy_video, config_video_scale) + + self.assertLessEqual(video_grid_thw[0, 1], max_grid_h) + self.assertLessEqual(video_grid_thw[0, 2], max_grid_w) + + # 3. Verify that passing the scaled video to maybe_pad_video_values_to_max_grid succeeds without ValueError + video_values = np.reshape( + video_processed, + ( + 1, + config_video_scale.num_channels_for_vit, + config_video_scale.temporal_patch_size_for_vit * video_grid_thw[0, 0], + config_video_scale.patch_size_for_vit * video_grid_thw[0, 1], + config_video_scale.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, + config_video_scale, + ) + self.assertEqual(padded_values.shape[3], max_grid_h * config_video_scale.patch_size_for_vit) + self.assertEqual(padded_values.shape[4], max_grid_w * config_video_scale.patch_size_for_vit) + np.testing.assert_array_equal(padded_grid, video_grid_thw) # padding must not change grid_thw + def test_patch_embed_is_jittable(self): """Test that patch embed is JIT-compilable."""