Optimize Omni audio preprocessing

streaming_turn_model.py

@@ -518,6 +518,8 @@ def _prepare_conversation(
     manual_audios: list[Any] | None = None,
     video_fps: float | None = None,
     text_override: str | None = None,
+    audio_feature_device: str | torch.device | None = None,
+    audio_right_padding_samples: int | None = None,
 ) -> dict[str, torch.Tensor]:
     """Decode/crop media and produce tensors for media turns only."""
 
@@ -557,6 +559,38 @@ def _prepare_conversation(
         "use_audio_in_video": use_audio_in_video,
         **processor_kwargs,
     }
+    audio_kwargs: dict[str, Any] = {}
+    if audio_feature_device is not None:
+        audio_kwargs["device"] = str(audio_feature_device)
+    if audio_right_padding_samples is not None:
+        if audio_right_padding_samples < 0:
+            raise ValueError("audio_right_padding_samples must be non-negative")
+        if audios is None or len(audios) != 1:
+            raise ValueError(
+                "Bounded audio padding currently requires exactly one audio input"
+            )
+        actual_samples = len(audios[0])
+        feature_extractor = processor.feature_extractor
+        hop_length = int(getattr(feature_extractor, "hop_length", 1))
+        if hop_length <= 0:
+            raise ValueError(f"Invalid audio hop_length: {hop_length}")
+        requested_length = actual_samples + audio_right_padding_samples
+        max_length = math.ceil(requested_length / hop_length) * hop_length
+        model_max_length = getattr(feature_extractor, "n_samples", None)
+        if model_max_length is not None:
+            max_length = min(max_length, int(model_max_length))
+        if max_length < actual_samples:
+            raise ValueError(
+                f"Bounded audio max_length {max_length} truncates "
+                f"{actual_samples} input samples"
+            )
+        audio_kwargs["max_length"] = max_length
+    if audio_kwargs:
+        # ProcessorMixin forwards nested modality kwargs to the underlying
+        # WhisperFeatureExtractor. This preserves Qwen's token expansion and
+        # feature_attention_mask construction while allowing CUDA STFT and a
+        # bounded amount of right-side silence.
+        processor_inputs["audio_kwargs"] = audio_kwargs
     # Passing images=None and return_tensors through Qwen's 4.52.3 processor
     # produces a warning on every streaming step. Omit the absent modality and
     # perform the same BatchFeature tensor conversion explicitly afterward.
@@ -1217,6 +1251,8 @@ class OpenStreamReplayExtractor:
         audio_path: str | Path | None,
         interval: tuple[float, float],
         use_audio_in_video: bool,
+        audio_feature_device: str | torch.device | None = None,
+        audio_right_padding_samples: int | None = None,
     ) -> OpenMediaPieces:
         conversation, manual_audios = build_chunk_conversation(
             video_path=video_path,
@@ -1233,6 +1269,8 @@ class OpenStreamReplayExtractor:
             model=self.model,
             manual_audios=manual_audios,
             video_fps=self.config.video_fps,
+            audio_feature_device=audio_feature_device,
+            audio_right_padding_samples=audio_right_padding_samples,
         )
         pieces = _open_media_pieces(
             self.model,