diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 91727be..981e071 100644 --- a/train_omni_turn_lora.py +++ b/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}, "