From b875a7512e2229d9870bac448bdfd4d2c852f680 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Sun, 23 Aug 2026 16:26:01 +0800 Subject: [PATCH] Optimize Omni audio preprocessing --- cache_omni_av_pieces.py | 77 ++++++++++++++++++++++++++++++++++++++--- streaming_turn_model.py | 38 ++++++++++++++++++++ 2 files changed, 110 insertions(+), 5 deletions(-) diff --git a/cache_omni_av_pieces.py b/cache_omni_av_pieces.py index 77a4e5e..1f4c130 100644 --- a/cache_omni_av_pieces.py +++ b/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( diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 65a7938..669666e 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -518,6 +518,8 @@ def _prepare_conversation( manual_audios: list[Any] | None = None, video_fps: float | None = None, text_override: str | None = None, + audio_feature_device: str | torch.device | None = None, + audio_right_padding_samples: int | None = None, ) -> dict[str, torch.Tensor]: """Decode/crop media and produce tensors for media turns only.""" @@ -557,6 +559,38 @@ def _prepare_conversation( "use_audio_in_video": use_audio_in_video, **processor_kwargs, } + audio_kwargs: dict[str, Any] = {} + if audio_feature_device is not None: + audio_kwargs["device"] = str(audio_feature_device) + if audio_right_padding_samples is not None: + if audio_right_padding_samples < 0: + raise ValueError("audio_right_padding_samples must be non-negative") + if audios is None or len(audios) != 1: + raise ValueError( + "Bounded audio padding currently requires exactly one audio input" + ) + actual_samples = len(audios[0]) + feature_extractor = processor.feature_extractor + hop_length = int(getattr(feature_extractor, "hop_length", 1)) + if hop_length <= 0: + raise ValueError(f"Invalid audio hop_length: {hop_length}") + requested_length = actual_samples + audio_right_padding_samples + max_length = math.ceil(requested_length / hop_length) * hop_length + model_max_length = getattr(feature_extractor, "n_samples", None) + if model_max_length is not None: + max_length = min(max_length, int(model_max_length)) + if max_length < actual_samples: + raise ValueError( + f"Bounded audio max_length {max_length} truncates " + f"{actual_samples} input samples" + ) + audio_kwargs["max_length"] = max_length + if audio_kwargs: + # ProcessorMixin forwards nested modality kwargs to the underlying + # WhisperFeatureExtractor. This preserves Qwen's token expansion and + # feature_attention_mask construction while allowing CUDA STFT and a + # bounded amount of right-side silence. + processor_inputs["audio_kwargs"] = audio_kwargs # Passing images=None and return_tensors through Qwen's 4.52.3 processor # produces a warning on every streaming step. Omit the absent modality and # perform the same BatchFeature tensor conversion explicitly afterward. @@ -1217,6 +1251,8 @@ class OpenStreamReplayExtractor: audio_path: str | Path | None, interval: tuple[float, float], use_audio_in_video: bool, + audio_feature_device: str | torch.device | None = None, + audio_right_padding_samples: int | None = None, ) -> OpenMediaPieces: conversation, manual_audios = build_chunk_conversation( video_path=video_path, @@ -1233,6 +1269,8 @@ class OpenStreamReplayExtractor: model=self.model, manual_audios=manual_audios, video_fps=self.config.video_fps, + audio_feature_device=audio_feature_device, + audio_right_padding_samples=audio_right_padding_samples, ) pieces = _open_media_pieces( self.model,