Optimize overlap working-memory audio preprocessing

streaming_turn_model_directly.py

@@ -1044,6 +1044,23 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
     parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
     parser.add_argument("--frame-cache-size", type=int, default=8)
+    parser.add_argument(
+        "--audio-feature-device",
+        default=None,
+        help=(
+            "Whisper log-Mel device for overlap-working-memory, for example "
+            "model, cuda:0, or cpu; omitted preserves processor defaults"
+        ),
+    )
+    parser.add_argument(
+        "--audio-right-padding-samples",
+        type=int,
+        default=None,
+        help=(
+            "For overlap-working-memory, bound processor audio padding to "
+            "actual samples plus this many zeros; omitted uses model maximum"
+        ),
+    )
     parser.add_argument(
         "--decoder",
         choices=("decord", "torchcodec", "pynvcodec"),
@@ -1183,6 +1200,19 @@ def main(
         )
     if args.frame_cache_size < 0:
         raise SystemExit("--frame-cache-size must be >= 0")
+    if (
+        args.audio_right_padding_samples is not None
+        and args.audio_right_padding_samples < 0
+    ):
+        raise SystemExit("--audio-right-padding-samples must be non-negative")
+    if args.mode != "overlap-working-memory" and (
+        args.audio_feature_device is not None
+        or args.audio_right_padding_samples is not None
+    ):
+        raise SystemExit(
+            "Audio preprocessing flags are currently implemented only for "
+            "--mode overlap-working-memory"
+        )
     if args.output is None and not args.media_only:
         raise SystemExit("--output is required unless --media-only is used")
     if args.audio is not None and args.video is None:
@@ -1306,6 +1336,24 @@ def main(
         model.eval()
     model.requires_grad_(False)
     media_path_for_api = reader.clips[0].path
+    audio_feature_device: str | torch.device | None = args.audio_feature_device
+    if audio_feature_device == "model":
+        audio_feature_device = model.device
+    if args.mode == "overlap-working-memory":
+        print(
+            "Overlap audio preprocessing: "
+            f"feature_device={audio_feature_device or 'processor default'}, "
+            "padding="
+            + (
+                "model maximum"
+                if args.audio_right_padding_samples is None
+                else (
+                    "actual samples + >= "
+                    f"{args.audio_right_padding_samples} right-padding samples"
+                )
+            ),
+            flush=True,
+        )
 
     overlap_extractor = None
     working_grid_av = None
@@ -1347,6 +1395,8 @@ def main(
             working_memory_seconds=args.working_memory_seconds,
             fixed_grid_av=reader,
             working_grid_av=working_grid_av,
+            audio_feature_device=audio_feature_device,
+            audio_right_padding_samples=args.audio_right_padding_samples,
         )
         result = overlap_extractor.extract(
             video_path=media_path_for_api,
@@ -1440,6 +1490,12 @@ def main(
                     if working_grid_av is not None
                     else None
                 ),
+                "audio_feature_device": (
+                    str(audio_feature_device)
+                    if audio_feature_device is not None
+                    else "processor_default"
+                ),
+                "audio_right_padding_samples": args.audio_right_padding_samples,
                 "layout": (
                     overlap_extractor.layout_records
                     if overlap_extractor is not None