diff --git a/README.md b/README.md index 2b26a60..7a19993 100644 --- a/README.md +++ b/README.md @@ -1300,6 +1300,13 @@ phase artifact rather than general evidence for better recent semantics. Run a 24-second L20 smoke test first: +The sliding 2-second recap otherwise invokes Whisper's model-maximum audio +padding 120 times. Run log-Mel extraction on the model GPU and set explicit +right padding to zero. Every EasyCom 0.2-second increment contains 3200 samples +at 16 kHz and is already divisible by Whisper's 160-sample hop, so no external +silence is needed. This preserves the feature extractor's own STFT boundary +handling while avoiding repeated processing of tens of seconds of zeros. + ```bash mkdir -p features/easycom_overlap_working_8s_smoke @@ -1317,6 +1324,8 @@ python streaming_turn_model_directly.py \ --working-memory-seconds 2 \ --video-fps 2 \ --frame-cache-size 128 \ + --audio-feature-device model \ + --audio-right-padding-samples 0 \ --decoder decord \ --decoder-threads 1 \ --max-open-readers 1 \ @@ -1366,6 +1375,8 @@ python streaming_turn_model_directly.py \ --working-memory-seconds 2 \ --video-fps 2 \ --frame-cache-size 128 \ + --audio-feature-device model \ + --audio-right-padding-samples 0 \ --decoder decord \ --decoder-threads 1 \ --max-open-readers 1 \ diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 5e7df2d..9a92431 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -1672,6 +1672,8 @@ class OverlappingWorkingMemoryExtractor: working_memory_seconds: float, fixed_grid_av: FixedGridAV, working_grid_av: FixedGridAV, + audio_feature_device: str | torch.device | None = None, + audio_right_padding_samples: int | None = None, ) -> None: config.validate() if working_memory_seconds <= 0: @@ -1700,6 +1702,13 @@ class OverlappingWorkingMemoryExtractor: self.processor = processor self.config = config self.working_memory_seconds = float(working_memory_seconds) + if ( + audio_right_padding_samples is not None + and audio_right_padding_samples < 0 + ): + raise ValueError("audio_right_padding_samples must be non-negative") + self.audio_feature_device = audio_feature_device + self.audio_right_padding_samples = audio_right_padding_samples # The working recap consumes one block from the matched token budget. self.fixed_block_budget = total_blocks - 1 self.fixed_extractor = OpenStreamReplayExtractor( @@ -1822,6 +1831,10 @@ class OverlappingWorkingMemoryExtractor: audio_path=audio_path, interval=(block_start, block_end), use_audio_in_video=use_audio_in_video, + audio_feature_device=self.audio_feature_device, + audio_right_padding_samples=( + self.audio_right_padding_samples + ), ) completed_blocks.put( CachedOpenMediaBlock( @@ -1874,6 +1887,8 @@ class OverlappingWorkingMemoryExtractor: audio_path=audio_path, interval=(working_start, current_time), use_audio_in_video=use_audio_in_video, + audio_feature_device=self.audio_feature_device, + audio_right_padding_samples=self.audio_right_padding_samples, ) state = self._state_from_prepared( video_path=video_path, diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index db4ee8e..44342fa 100644 --- a/streaming_turn_model_directly.py +++ b/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