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(