Optimize overlap working-memory audio preprocessing

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,