[Gemma4] Pass num_frames_per_video to get_video_features in forward (#50237)

Gemma4: pass num_frames_per_video through forward to get_video_features

Author: yzkkai

src/transformers/models/gemma4/modeling_gemma4.py

@@ -2078,7 +2078,7 @@ def forward(
 
         if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None:
             mm_encoder_outputs["video"] = self.get_video_features(
-                pixel_values_videos, video_position_ids, return_dict=True
+                pixel_values_videos, video_position_ids, num_frames_per_video=num_frames_per_video, return_dict=True
             )
 
         if mm_encoder_outputs.get("video") is not None:
@@ -2271,6 +2271,7 @@ def forward(
             use_cache=use_cache,
             image_position_ids=image_position_ids,
             video_position_ids=video_position_ids,
+            num_frames_per_video=num_frames_per_video,
             mm_encoder_outputs=mm_encoder_outputs,
             return_dict=True,
             **kwargs,

src/transformers/models/gemma4/modular_gemma4.py

@@ -1081,6 +1081,7 @@ def forward(
         use_cache: bool | None = None,
         image_position_ids: torch.LongTensor | None = None,
         video_position_ids: torch.LongTensor | None = None,
+        num_frames_per_video: torch.LongTensor | None = None,
         mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Gemma4UnifiedModelOutputWithPast:
@@ -1093,6 +1094,9 @@ def forward(
         video_position_ids (`torch.LongTensor` of shape `(total_num_frames, max_patches, 2)`, *optional*):
             2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding.
             Passed through to the vision encoder for positional embedding computation.
+        num_frames_per_video (`torch.LongTensor` of shape `(num_videos,)`, *optional*):
+            Number of frames belonging to each video. Required whenever `pixel_values_videos` is passed,
+            since the frames of all videos are concatenated along a single axis.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -1137,7 +1141,7 @@ def forward(
 
         if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None:
             mm_encoder_outputs["video"] = self.get_video_features(
-                pixel_values_videos, video_position_ids, return_dict=True
+                pixel_values_videos, video_position_ids, num_frames_per_video=num_frames_per_video, return_dict=True
             )
 
         if mm_encoder_outputs.get("video") is not None:
@@ -1279,6 +1283,7 @@ def forward(
             use_cache=use_cache,
             image_position_ids=image_position_ids,
             video_position_ids=video_position_ids,
+            num_frames_per_video=num_frames_per_video,
             mm_encoder_outputs=mm_encoder_outputs,
             return_dict=True,
             **kwargs,

src/transformers/models/gemma4_unified/modeling_gemma4_unified.py

@@ -464,6 +464,38 @@ def test_training(self):
         loss = model(**inputs).loss
         loss.backward()
 
+    def test_forward_with_videos(self):
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+        model = Gemma4ForConditionalGeneration(config).to(torch_device).eval()
+
+        # Two videos with 2 and 3 frames
+        num_frames_per_video = torch.tensor([2, 3], device=torch_device)
+        pixel_values_videos = self.model_tester.create_pixel_values(5).to(torch_device)
+        video_position_ids = self.model_tester.create_image_position_ids(5)
+        num_tokens = self.model_tester.num_image_tokens  # per frame
+        input_ids = torch.ones((2, 3 * num_tokens + 2), dtype=torch.long, device=torch_device)
+        input_ids[0, 1 : 1 + 2 * num_tokens] = config.video_token_id
+        input_ids[1, 1 : 1 + 3 * num_tokens] = config.video_token_id
+        mm_token_type_ids = torch.zeros_like(input_ids)
+        mm_token_type_ids[input_ids == config.video_token_id] = 2
+
+        with torch.no_grad():
+            video_features = model.model.get_video_features(
+                pixel_values_videos, video_position_ids, num_frames_per_video=num_frames_per_video
+            )
+            # `generate` encodes the videos first, so check that raw pixels give the same output
+            expected = model(
+                input_ids, mm_token_type_ids=mm_token_type_ids, mm_encoder_outputs={"video": video_features}
+            )
+            output = model(
+                input_ids,
+                mm_token_type_ids=mm_token_type_ids,
+                pixel_values_videos=pixel_values_videos,
+                video_position_ids=video_position_ids,
+                num_frames_per_video=num_frames_per_video,
+            )
+        torch.testing.assert_close(output.logits, expected.logits)
+
     def test_vision_axial_rope(self):
         # override -> model shipped weirdly to from the start, pos IDs have actual batch dim
 

src/transformers/models/gemma4_unified/modular_gemma4_unified.py

@@ -427,6 +427,39 @@ def skip_mm_output_format(self):
     def test_num_layers_is_small(self):
         pass
 
+    def test_forward_with_videos(self):
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+        model = Gemma4UnifiedForConditionalGeneration(config).to(torch_device).eval()
+
+        # Two videos with 2 and 3 frames
+        num_frames_per_video = torch.tensor([2, 3], device=torch_device)
+        num_tokens = 2  # per frame, one token per patch
+        patch_dim = config.vision_config.model_patch_size**2 * config.vision_config.num_channels
+        pixel_values_videos = floats_tensor([5, num_tokens, patch_dim]).to(torch_device)
+        video_position_ids = torch.ones(5, num_tokens, 2, dtype=torch.long, device=torch_device)
+        input_ids = torch.ones((2, 3 * num_tokens + 2), dtype=torch.long, device=torch_device)
+        input_ids[0, 1 : 1 + 2 * num_tokens] = config.video_token_id
+        input_ids[1, 1 : 1 + 3 * num_tokens] = config.video_token_id
+        mm_token_type_ids = torch.zeros_like(input_ids)
+        mm_token_type_ids[input_ids == config.video_token_id] = 2
+
+        with torch.no_grad():
+            video_features = model.model.get_video_features(
+                pixel_values_videos, video_position_ids, num_frames_per_video=num_frames_per_video
+            )
+            # `generate` encodes the videos first, so check that raw pixels give the same output
+            expected = model(
+                input_ids, mm_token_type_ids=mm_token_type_ids, mm_encoder_outputs={"video": video_features}
+            )
+            output = model(
+                input_ids,
+                mm_token_type_ids=mm_token_type_ids,
+                pixel_values_videos=pixel_values_videos,
+                video_position_ids=video_position_ids,
+                num_frames_per_video=num_frames_per_video,
+            )
+        torch.testing.assert_close(output.logits, expected.logits)
+
     @unittest.skip("We use vision based masks that use `block_sequence_ids` which force mask materialization")
     def test_sdpa_can_dispatch_on_flash(self):
         pass