fix: bound streaming VAD frame buffers

funasr/frontends/wav_frontend.py

@@ -262,6 +262,8 @@ def forward_lfr_cmvn(
 class WavFrontendOnline(nn.Module):
     """Conventional frontend structure for streaming ASR/VAD."""
 
+    supports_aligned_waveforms = True
+
     def __init__(
         self,
         cmvn_file: str = None,
@@ -514,17 +516,30 @@ def forward(self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs):
         cache = kwargs.get("cache", {})
         if len(cache) == 0:
             self.init_cache(cache)
+        cache.setdefault("waveform_buffer", None)
+        cache.setdefault("waveform_buffer_start_sample", 0)
+        cache.setdefault("emitted_lfr_frames", 0)
+        cache["aligned_waveforms"] = torch.empty(0)
+        return_waveform = kwargs.get("return_waveform", False)
 
         batch_size = input.shape[0]
         assert (
             batch_size == 1
         ), "we support to extract feature online only when the batch size is equal to 1 now"
+        # VAD opts in to exact score-aligned waveform spans.
+        if return_waveform and input.shape[1] > 0:
+            cache["waveform_buffer"] = (
+                input.clone()
+                if cache["waveform_buffer"] is None
+                else torch.cat((cache["waveform_buffer"], input), dim=1)
+            )
 
         waveforms, feats, feats_lengths = self.forward_fbank(
             input, input_lengths, cache=cache
         )  # input shape: B T D
+        has_fbank_frames = bool(feats.shape[0])
 
-        if feats.shape[0]:
+        if has_fbank_frames:
 
             cache["waveforms"] = torch.cat((cache["reserve_waveforms"], waveforms), dim=1)
 
@@ -587,6 +602,41 @@ def forward(self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs):
                 feats, feats_lengths, _ = self.forward_lfr_cmvn(
                     feats, feats_lengths, is_final, cache=cache
                 )
+
+        if return_waveform and feats.ndim == 3 and feats.shape[1] > 0:
+            score_frame_shift = self.lfr_n * self.frame_shift_sample_length
+            aligned_start = cache["emitted_lfr_frames"] * score_frame_shift
+            aligned_sample_count = (
+                (feats.shape[1] - 1) * score_frame_shift + self.frame_sample_length
+            )
+            aligned_end = aligned_start + aligned_sample_count
+            waveform_buffer = cache["waveform_buffer"]
+            buffer_start = cache["waveform_buffer_start_sample"]
+            buffer_end = buffer_start + (
+                0 if waveform_buffer is None else waveform_buffer.shape[1]
+            )
+            if (
+                waveform_buffer is None
+                or aligned_start < buffer_start
+                or aligned_end > buffer_end
+            ):
+                raise RuntimeError(
+                    "Frontend emitted LFR frames without enough waveform samples: "
+                    f"need [{aligned_start}, {aligned_end}), "
+                    f"have [{buffer_start}, {buffer_end})"
+                )
+            local_start = aligned_start - buffer_start
+            local_end = aligned_end - buffer_start
+            cache["aligned_waveforms"] = waveform_buffer[
+                :, local_start:local_end
+            ].clone()
+            cache["emitted_lfr_frames"] += feats.shape[1]
+
+            next_frame_start = cache["emitted_lfr_frames"] * score_frame_shift
+            next_buffer_start = min(next_frame_start, buffer_end)
+            drop_samples = next_buffer_start - buffer_start
+            cache["waveform_buffer"] = waveform_buffer[:, drop_samples:].clone()
+            cache["waveform_buffer_start_sample"] = next_buffer_start
         # if is_final:
         #     self.init_cache(cache)
         return feats, feats_lengths
@@ -603,6 +653,10 @@ def init_cache(self, cache: dict = None):
         cache["input_cache"] = torch.empty(0)
         cache["lfr_splice_cache"] = []
         cache["waveforms"] = None
+        cache["aligned_waveforms"] = torch.empty(0)
+        cache["waveform_buffer"] = None
+        cache["waveform_buffer_start_sample"] = 0
+        cache["emitted_lfr_frames"] = 0
         cache["fbanks"] = None
         cache["fbanks_lens"] = None
         return cache

funasr/models/fsmn_vad_streaming/model.py

@@ -403,6 +403,35 @@ def __init__(
         self.encoder = encoder
         self.encoder_conf = encoder_conf
 
+    def DropCachedFrames(self, drop_frames: int, cache: dict = None) -> None:
+        """Release score, decibel, and waveform history before an absolute frame."""
+        if cache is None:
+            cache = {}
+        stats = cache["stats"]
+        if drop_frames > stats.frm_cnt:
+            raise RuntimeError(
+                f"Cannot drop through frame {drop_frames}; only {stats.frm_cnt} frames exist"
+            )
+
+        real_drop_frames = drop_frames - stats.last_drop_frames
+        if real_drop_frames <= 0:
+            return
+
+        frame_shift_length = int(
+            self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000
+        )
+        drop_samples = real_drop_frames * frame_shift_length
+        stats.data_buf_all = stats.data_buf_all[drop_samples:].clone()
+        stats.decibel = stats.decibel[real_drop_frames:]
+        stats.scores = stats.scores[:, real_drop_frames:, :].clone()
+        stats.last_drop_frames = drop_frames
+
+        data_buf_offset = max(
+            stats.data_buf_start_frame - stats.last_drop_frames,
+            0,
+        ) * frame_shift_length
+        stats.data_buf = stats.data_buf_all[data_buf_offset:]
+
     def ResetDetection(self, cache: dict = None):
         """Resetdetection.
 
@@ -424,44 +453,80 @@ def ResetDetection(self, cache: dict = None):
         if cache["stats"].output_data_buf:
             assert cache["stats"].output_data_buf[-1].contain_seg_end_point == True
             drop_frames = int(cache["stats"].output_data_buf[-1].end_ms / self.vad_opts.frame_in_ms)
-            real_drop_frames = drop_frames - cache["stats"].last_drop_frames
-            cache["stats"].last_drop_frames = drop_frames
-            cache["stats"].data_buf_all = cache["stats"].data_buf_all[
-                real_drop_frames
-                * int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) :
-            ]
-            cache["stats"].decibel = cache["stats"].decibel[real_drop_frames:]
-            cache["stats"].scores = cache["stats"].scores[:, real_drop_frames:, :]
-
-    def ComputeDecibel(self, cache: dict = None) -> None:
+            self.DropCachedFrames(drop_frames, cache=cache)
+
+    def ComputeDecibel(
+        self, cache: dict = None, frame_count: Optional[int] = None
+    ) -> None:
         """Computedecibel.
 
             Args:
+                frame_count: Number of VAD score frames emitted for this waveform.
                 cache: State cache dict for streaming inference.
             """
         if cache is None:
             cache = {}
         frame_sample_length = int(self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000)
         frame_shift_length = int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000)
-        if cache["stats"].data_buf_all is None:
-            cache["stats"].data_buf_all = cache["stats"].waveform[
-                0
-            ]  # cache["stats"].data_buf is pointed to cache["stats"].waveform[0]
-            cache["stats"].data_buf = cache["stats"].data_buf_all
+        if frame_count is not None and int(frame_count) <= 0:
+            return
+        stats = cache["stats"]
+        waveform = stats.waveform[0]
+
+        frame_count_was_provided = frame_count is not None
+        available_frames = max(
+            (waveform.numel() - frame_sample_length) // frame_shift_length + 1,
+            0,
+        )
+        frame_count = available_frames if frame_count is None else int(frame_count)
+        if frame_count <= 0:
+            return
+
+        # The frontend provides the exact waveform span for these VAD score frames.
+        aligned_sample_count = (
+            (frame_count - 1) * frame_shift_length + frame_sample_length
+        )
+        if waveform.numel() != aligned_sample_count:
+            if frame_count_was_provided or waveform.numel() < aligned_sample_count:
+                raise RuntimeError(
+                    "VAD score frames and waveform samples are not aligned: "
+                    f"expected {aligned_sample_count}, got {waveform.numel()}"
+                )
+            waveform = waveform[:aligned_sample_count]
+
+        aligned_waveform = waveform
+
+        if stats.data_buf_all is None:
+            stats.data_buf_all = aligned_waveform.clone()
         else:
-            cache["stats"].data_buf_all = torch.cat(
-                (cache["stats"].data_buf_all, cache["stats"].waveform[0])
+            new_sample_count = frame_count * frame_shift_length
+            stats.data_buf_all = torch.cat(
+                (stats.data_buf_all, aligned_waveform[-new_sample_count:])
             )
-
-        waveform_numpy = cache["stats"].waveform.numpy()
 
-        offsets = np.arange(0, waveform_numpy.shape[1] - frame_sample_length + 1, frame_shift_length)
-        frames = waveform_numpy[0, offsets[:, np.newaxis] + np.arange(frame_sample_length)]
+        data_buf_offset = max(
+            stats.data_buf_start_frame - stats.last_drop_frames,
+            0,
+        ) * frame_shift_length
+        stats.data_buf = stats.data_buf_all[data_buf_offset:]
 
-        decibel_numpy = 10 * np.log10(np.sum(np.square(frames), axis=1) + 0.000001)
+        waveform_numpy = aligned_waveform.detach().cpu().numpy()
+
+        offsets = np.arange(
+            0,
+            waveform_numpy.shape[0] - frame_sample_length + 1,
+            frame_shift_length,
+        )
+        frames = waveform_numpy[
+            offsets[:, np.newaxis] + np.arange(frame_sample_length)
+        ]
+
+        decibel_numpy = 10 * np.log10(
+            np.sum(np.square(frames), axis=1) + 0.000001
+        )
         decibel_numpy = decibel_numpy.tolist()
 
-        cache["stats"].decibel.extend(decibel_numpy)
+        stats.decibel.extend(decibel_numpy)
 
 
     def ComputeScores(self, feats: torch.Tensor, cache: dict = None) -> None:
@@ -710,7 +775,6 @@ def GetFrameState(self, t: int, cache: dict = None):
         # for each frame, calc log posterior probability of each state
         if cur_decibel < self.vad_opts.decibel_thres:
             frame_state = FrameState.kFrameStateSil
-            self.DetectOneFrame(frame_state, t, False, cache=cache)
             return frame_state
 
         sum_score = 0.0
@@ -777,17 +841,20 @@ def forward(
             """
         if cache is None:
             cache = {}
+        if feats.numel() == 0:
+            return []
         # if len(cache) == 0:
         #     self.AllResetDetection()
         # self.waveform = waveform  # compute decibel for each frame
         cache["stats"].waveform = waveform
         is_streaming_input = kwargs.get("is_streaming_input", True)
-        self.ComputeDecibel(cache=cache)
+        self.ComputeDecibel(frame_count=feats.shape[1], cache=cache)
         self.ComputeScores(feats, cache=cache)
         if not is_final:
             self.DetectCommonFrames(cache=cache)
         else:
             self.DetectLastFrames(cache=cache)
+        self.DropCachedFrames(cache["stats"].data_buf_start_frame, cache=cache)
         segments = []
         for batch_num in range(0, feats.shape[0]):  # only support batch_size = 1 now
             segment_batch = []
@@ -946,6 +1013,9 @@ def inference(
         speech_to_sil_ms = self.vad_opts.speech_to_sil_time_thres
         accumulated_ms = cache.get("_dynamic_accumulated_ms", 0)
         in_speech = cache.get("_dynamic_in_speech", False)
+        frontend_inference_kwargs = {}
+        if getattr(frontend, "supports_aligned_waveforms", False):
+            frontend_inference_kwargs["return_waveform"] = True
 
         for i in range(n):
             kwargs["is_final"] = _is_final and i == n - 1
@@ -973,6 +1043,7 @@ def inference(
                 frontend=frontend,
                 cache=cache["frontend"],
                 is_final=kwargs["is_final"],
+                **frontend_inference_kwargs,
             )
             time3 = time.perf_counter()
             meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
@@ -982,9 +1053,18 @@ def inference(
             speech = speech.to(device=kwargs["device"])
             speech_lengths = speech_lengths.to(device=kwargs["device"])
 
+            waveform = cache["frontend"].get("aligned_waveforms")
+            if waveform is None or waveform.numel() == 0:
+                # Custom frontends may expose an exact span under the legacy key.
+                waveform = cache["frontend"].get("waveforms")
+            if speech.numel() > 0 and (waveform is None or waveform.numel() == 0):
+                raise RuntimeError(
+                    "Streaming VAD frontend must provide aligned_waveforms or waveforms "
+                    "when emitting feature frames"
+                )
             batch = {
                 "feats": speech,
-                "waveform": cache["frontend"]["waveforms"],
+                "waveform": waveform,
                 "is_final": kwargs["is_final"],
                 "cache": cache,
                 "is_streaming_input": is_streaming_input,