Add video-scaling function for SFT padding

Author: hengtaoguo

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

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