Optimize Omni audio preprocessing

cache_omni_av_pieces.py

@@ -11,6 +11,11 @@ cache with the existing open-stream KV replay.  Run this 24-second pilot before
 caching a full EasyCom session.  ``--av-batch-size`` batches equal-duration
 phases from different source blocks, increasing frozen AV-tower occupancy while
 preserving the exact per-block cache schema and causal timestamps.
+
+The audio optimization flags keep the official Qwen/Whisper feature extractor
+but may run its STFT on CUDA and bound the otherwise model-maximum zero padding.
+Verification deliberately recomputes its oracle with the original processor
+defaults, so an optimization cannot silently redefine the extracted feature.
 """
 
 from __future__ import annotations
@@ -72,6 +77,24 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
             "the processor and frozen AV towers; 1 preserves serial behavior"
         ),
     )
+    parser.add_argument(
+        "--audio-feature-device",
+        default=None,
+        help=(
+            "Device used by the native Whisper log-Mel extractor, for example "
+            "model, cuda:0, or cpu; omitted preserves processor defaults"
+        ),
+    )
+    parser.add_argument(
+        "--audio-right-padding-samples",
+        type=int,
+        default=None,
+        help=(
+            "Replace model-maximum audio padding with actual samples plus this "
+            "many zeros, rounded up to the Whisper hop; omitted preserves the "
+            "original maximum-length padding"
+        ),
+    )
     parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
     parser.add_argument("--decoder-gpu-id", type=int, default=0)
     parser.add_argument("--decoder-threads", type=int, default=1)
@@ -274,6 +297,8 @@ def _cache_blocks(
     step_seconds: float,
     block_seconds: float,
     av_batch_size: int,
+    audio_feature_device: str | torch.device | None,
+    audio_right_padding_samples: int | None,
     overwrite: bool,
 ) -> list[dict[str, Any]]:
     num_blocks = math.ceil(num_steps / steps_per_block)
@@ -328,6 +353,8 @@ def _cache_blocks(
                         audio_path=None,
                         interval=intervals[0],
                         use_audio_in_video=True,
+                        audio_feature_device=audio_feature_device,
+                        audio_right_padding_samples=audio_right_padding_samples,
                     )
                 ]
             else:
@@ -470,6 +497,19 @@ def main(argv: Sequence[str] | None = None) -> None:
         raise ValueError("--end must be positive")
     if args.av_batch_size <= 0:
         raise ValueError("--av-batch-size must be positive")
+    if (
+        args.audio_right_padding_samples is not None
+        and args.audio_right_padding_samples < 0
+    ):
+        raise ValueError("--audio-right-padding-samples must be non-negative")
+    if args.av_batch_size != 1 and (
+        args.audio_feature_device is not None
+        or args.audio_right_padding_samples is not None
+    ):
+        raise ValueError(
+            "Optimized audio preprocessing currently requires --av-batch-size 1; "
+            "the HF multimodal batch path was slower and changed cached states"
+        )
     steps_per_block = _integer_ratio(
         args.native_block_seconds, args.step_seconds, "native block / step"
     )
@@ -522,6 +562,24 @@ def main(argv: Sequence[str] | None = None) -> None:
         use_flash_attention_2=args.flash_attention_2,
     )
     model.requires_grad_(False)
+    audio_feature_device: str | torch.device | None = args.audio_feature_device
+    if audio_feature_device == "model":
+        audio_feature_device = _model_device(model)
+    audio_padding_description = "model maximum"
+    if args.audio_right_padding_samples is not None:
+        hop_length = int(getattr(processor.feature_extractor, "hop_length", 1))
+        rounded_padding = (
+            math.ceil(args.audio_right_padding_samples / hop_length) * hop_length
+        )
+        audio_padding_description = (
+            f"actual + >= {args.audio_right_padding_samples} samples "
+            f"(hop-aligned, typically +{rounded_padding})"
+        )
+    print(
+        "Audio preprocessing: "
+        f"feature_device={audio_feature_device or 'processor default'}, "
+        f"padding={audio_padding_description}"
+    )
     extractor = core.OpenStreamReplayExtractor(
         model, processor, config, fixed_grid_av=media
     )
@@ -552,9 +610,16 @@ def main(argv: Sequence[str] | None = None) -> None:
         step_seconds=args.step_seconds,
         block_seconds=args.native_block_seconds,
         av_batch_size=args.av_batch_size,
+        audio_feature_device=audio_feature_device,
+        audio_right_padding_samples=args.audio_right_padding_samples,
         overwrite=args.overwrite,
     )
     cache_seconds = time.perf_counter() - cache_started
+    print(
+        f"AV cache generation: {cache_seconds:.3f}s, "
+        f"{len(block_records) / cache_seconds:.3f} block/s, "
+        f"{args.end / cache_seconds:.3f}x media realtime"
+    )
     anchors = [
         make_anchor(
             step_index=step_index,
@@ -602,6 +667,13 @@ def main(argv: Sequence[str] | None = None) -> None:
         "interval_convention": "half_open_[start,end)",
         "phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
         "av_batch_size": args.av_batch_size,
+        "audio_feature_device": (
+            None if audio_feature_device is None else str(audio_feature_device)
+        ),
+        "audio_right_padding_samples": args.audio_right_padding_samples,
+        "verification_reference_audio_preprocessing": (
+            "processor-default-device-and-model-maximum-padding"
+        ),
         "cache_generation_seconds": cache_seconds,
         "cache_media_realtime_factor": args.end / cache_seconds,
         "cache_blocks_per_second": len(block_records) / cache_seconds,
@@ -623,11 +695,6 @@ def main(argv: Sequence[str] | None = None) -> None:
     print(
         f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
     )
-    print(
-        f"AV cache generation: {cache_seconds:.3f}s, "
-        f"{len(block_records) / cache_seconds:.3f} block/s, "
-        f"{args.end / cache_seconds:.3f}x media realtime"
-    )
     print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB")
     if checks:
         print(

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,