Add video-scaling function for SFT padding

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

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."""