Optimize online-tail audio preprocessing

train_omni_turn_lora.py

@@ -113,6 +113,25 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         help="Maximum producer-ready CPU batches queued ahead of training",
     )
     parser.add_argument("--frame-cache-size", type=int, default=64)
+    parser.add_argument(
+        "--audio-feature-device",
+        default=None,
+        help=(
+            "Device used by the native Whisper log-Mel extractor for online "
+            "current tails, for example model, cuda:0, or cpu; omitted "
+            "preserves processor defaults"
+        ),
+    )
+    parser.add_argument(
+        "--audio-right-padding-samples",
+        type=int,
+        default=None,
+        help=(
+            "For online current tails, replace model-maximum audio padding "
+            "with actual samples plus this many zeros (hop-aligned); omitted "
+            "preserves the original maximum-length padding"
+        ),
+    )
     parser.add_argument(
         "--decoder",
         choices=("decord", "torchcodec", "pynvcodec"),
@@ -318,6 +337,8 @@ class OnlineCurrentTailMaterializer:
         decoder_gpu_id: int,
         decoder_threads: int,
         max_open_readers: int,
+        audio_feature_device: str | torch.device | None,
+        audio_right_padding_samples: int | None,
     ) -> None:
         if len(dataset.sources) != 1:
             raise ValueError("Online-tail benchmark supports exactly one cache directory")
@@ -354,6 +375,13 @@ class OnlineCurrentTailMaterializer:
         self.step_seconds = config.step_seconds
         self.block_seconds = config.native_block_seconds
         self.steps_per_block = int(source["static"]["steps_per_block"])
+        if audio_feature_device == "model":
+            try:
+                audio_feature_device = next(model.parameters()).device
+            except StopIteration as exc:
+                raise RuntimeError("Online-tail AV model has no parameters") from exc
+        self.audio_feature_device = audio_feature_device
+        self.audio_right_padding_samples = audio_right_padding_samples
         self.calls = 0
         self.online_av_calls = 0
         self.total_seconds = 0.0
@@ -382,6 +410,8 @@ class OnlineCurrentTailMaterializer:
                 audio_path=None,
                 interval=(block_start, interval_end),
                 use_audio_in_video=True,
+                audio_feature_device=self.audio_feature_device,
+                audio_right_padding_samples=self.audio_right_padding_samples,
             )
             closure_embeddings = self.model.get_input_embeddings()(
                 pieces.closure_ids
@@ -659,6 +689,18 @@ def main(argv: Sequence[str] | None = None) -> None:
             raise ValueError("--producer-device must be a CUDA device such as cuda:1")
     if args.prefetch_depth <= 0:
         raise ValueError("--prefetch-depth 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 not args.online_current_tail and (
+        args.audio_feature_device is not None
+        or args.audio_right_padding_samples is not None
+    ):
+        raise ValueError(
+            "Audio preprocessing overrides apply only with --online-current-tail"
+        )
     if args.num_workers != 0:
         print(
             "Warning: each worker owns an independent block-file LRU; use 0 for "
@@ -763,11 +805,23 @@ def main(argv: Sequence[str] | None = None) -> None:
             decoder_gpu_id=args.decoder_gpu_id,
             decoder_threads=args.decoder_threads,
             max_open_readers=args.max_open_readers,
+            audio_feature_device=args.audio_feature_device,
+            audio_right_padding_samples=args.audio_right_padding_samples,
         )
         print(
             "Cache policy benchmark: complete 2s history cached; "
             "current 0.2--1.8s tail computed online"
         )
+        print(
+            "Online-tail audio preprocessing: "
+            f"feature_device={args.audio_feature_device or 'processor default'}, "
+            "padding="
+            + (
+                "model maximum"
+                if args.audio_right_padding_samples is None
+                else f"actual + >= {args.audio_right_padding_samples} samples"
+            )
+        )
         if args.producer_device is not None:
             print(
                 f"Producer-consumer enabled: frozen AV={args.producer_device}, "