diff --git a/tests/test_base_backend.py b/tests/test_base_backend.py index bdb6f95..576d455 100644 --- a/tests/test_base_backend.py +++ b/tests/test_base_backend.py @@ -24,6 +24,24 @@ class ConcreteServeClient(ServeClientBase): pass +class WaitTrackingEvent: + """Threading event that records when wait() is entered.""" + + def __init__(self): + self._event = threading.Event() + self.wait_started = threading.Event() + + def wait(self, timeout=None): + self.wait_started.set() + return self._event.wait(timeout) + + def set(self): + self._event.set() + + def __getattr__(self, name): + return getattr(self._event, name) + + class TestServeClientBaseInit(unittest.TestCase): def test_default_values(self): ws = MagicMock() @@ -91,7 +109,6 @@ class TestAddFrames(unittest.TestCase): # timestamp_offset should be bumped to at least frames_offset self.assertGreaterEqual(self.client.timestamp_offset, self.client.frames_offset) - class TestAddFramesThreadSafety(unittest.TestCase): def test_concurrent_add_frames(self): ws = MagicMock() @@ -114,6 +131,18 @@ class TestAddFramesThreadSafety(unittest.TestCase): self.assertEqual(errors, []) self.assertIsNotNone(client.frames_np) + def test_exception_releases_lock_without_signaling_frames_ready(self): + ws = MagicMock() + client = ConcreteServeClient(client_uid="test", websocket=ws) + client.frames_np = np.array([0.1], dtype=np.float32) + + with patch("whisper_live.backend.base.np.concatenate", side_effect=RuntimeError("boom")): + with self.assertRaisesRegex(RuntimeError, "boom"): + client.add_frames(np.array([0.2], dtype=np.float32)) + + self.assertFalse(client.lock.locked()) + self.assertFalse(client.frames_ready.is_set()) + class TestGetAudioChunkForProcessing(unittest.TestCase): def setUp(self): @@ -266,6 +295,102 @@ class TestCleanup(unittest.TestCase): self.assertTrue(client.exit) +def _supports_thread_time(): + thread_time = getattr(time, "thread_time", None) + if thread_time is None: + return False + try: + thread_time() + except NotImplementedError: + return False + return True + +class TestSpeechToTextWaitingBehavior(unittest.TestCase): + """Tests the first-frame wait behavior in speech_to_text().""" + + def setUp(self): + self.ws = MagicMock() + self.client = ConcreteServeClient(client_uid="test", websocket=self.ws) + self.client.frames_ready = WaitTrackingEvent() + self.transcribe_called = threading.Event() + self.thread_started = threading.Event() + self.cpu_used = None + self.thread = None + + def tearDown(self): + if self.thread is not None and self.thread.is_alive(): + self.client.exit = True + # release wait() directly so a broken cleanup() cannot hang the test process + self.client.frames_ready.set() + self.thread.join(timeout=1.0) + + def _start_speech_thread(self, target=None): + self.thread = threading.Thread(target=target or self.client.speech_to_text) + self.thread.start() + return self.thread + + def _join_speech_thread(self): + self.thread.join(timeout=1.0) + return not self.thread.is_alive() + + def _transcribe_once(self, input_sample): + # mark the first processing step after wait and stop the loop + self.transcribe_called.set() + self.client.exit = True + return [] + + def _measure_waiting_cpu(self): + # measure CPU consumed by speech_to_text loop while it waits for the first frame + self.thread_started.set() + start_cpu = time.thread_time() + self.client.speech_to_text() + self.cpu_used = time.thread_time() - start_cpu + + def test_waits_for_first_frame_before_transcribing(self): + self.client.transcribe_audio = MagicMock(side_effect=self._transcribe_once) + self._start_speech_thread() + self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0)) + self.assertFalse(self.transcribe_called.is_set()) + + self.client.add_frames(np.zeros(self.client.RATE, dtype=np.float32)) + + self.assertTrue(self.transcribe_called.wait(timeout=1.0)) + self.assertTrue(self._join_speech_thread()) + + def test_cleanup_unblocks_waiting_thread_without_audio(self): + self.client.transcribe_audio = MagicMock() + self._start_speech_thread() + self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0)) + + self.client.cleanup() + + self.assertTrue(self._join_speech_thread()) + self.assertTrue(self.client.exit) + self.client.transcribe_audio.assert_not_called() + + def test_exit_flag_unblocks_waiting_thread_without_signal(self): + self.client.transcribe_audio = MagicMock() + self._start_speech_thread() + self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0)) + + self.client.exit = True + + self.assertTrue(self._join_speech_thread()) + self.client.transcribe_audio.assert_not_called() + + @unittest.skipUnless(_supports_thread_time(), "time.thread_time() not supported") + def test_waiting_for_first_frame_uses_negligible_thread_cpu(self): + self._start_speech_thread(target=self._measure_waiting_cpu) + self.assertTrue(self.thread_started.wait(timeout=1.0)) + + time.sleep(0.25) + self.client.cleanup() + + self.assertTrue(self._join_speech_thread()) + self.assertIsNotNone(self.cpu_used) + self.assertLess(self.cpu_used, 0.1) + + class TestTrimTranscript(unittest.TestCase): def setUp(self): self.ws = MagicMock() diff --git a/whisper_live/backend/base.py b/whisper_live/backend/base.py index f2ee52b..afa2eca 100644 --- a/whisper_live/backend/base.py +++ b/whisper_live/backend/base.py @@ -21,6 +21,8 @@ class ServeClientBase(object): """Duration threshold in seconds for clipping audio with no valid segments.""" CLIP_TAIL_DURATION_S = 5 """Duration in seconds of audio to keep after clipping.""" + FIRST_FRAME_WAIT_TIMEOUT_S = 0.1 + """Interval in seconds for re-checking exit while waiting for the first audio frame.""" client_uid: str """A unique identifier for the client.""" @@ -81,13 +83,15 @@ class ServeClientBase(object): # threading self.lock = threading.Lock() + self.frames_ready = threading.Event() def speech_to_text(self): """ Process an audio stream in an infinite loop, continuously transcribing the speech. This method continuously receives audio frames, performs real-time transcription, and sends - transcribed segments to the client via a WebSocket connection. + transcribed segments to the client via a WebSocket connection. The loop blocks until the first + audio frame arrives when a client is connected but still idle. If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction. It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments @@ -103,6 +107,8 @@ class ServeClientBase(object): break if self.frames_np is None: + while self.frames_np is None and not self.exit: + self.frames_ready.wait(timeout=self.FIRST_FRAME_WAIT_TIMEOUT_S) continue if self.clip_audio: @@ -170,7 +176,8 @@ class ServeClientBase(object): This method is responsible for maintaining the audio stream buffer, allowing the continuous addition of audio frames as they are received. It also ensures that the buffer does not exceed a specified size - to prevent excessive memory usage. + to prevent excessive memory usage. When the first frame arrives, it also wakes the transcription + thread so processing can begin. If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds of audio data to maintain a reasonable buffer size. If the buffer is empty, it initializes it with the provided @@ -180,20 +187,20 @@ class ServeClientBase(object): frame_np (numpy.ndarray): The audio frame data as a NumPy array. """ - self.lock.acquire() - if self.frames_np is not None and self.frames_np.shape[0] > self.MAX_BUFFER_DURATION_S*self.RATE: - self.frames_offset += float(self.BUFFER_TRIM_DURATION_S) - self.frames_np = self.frames_np[int(self.BUFFER_TRIM_DURATION_S*self.RATE):] - # check timestamp offset(should be >= self.frame_offset) - # this basically means that there is no speech as timestamp offset hasnt updated - # and is less than frame_offset - if self.timestamp_offset < self.frames_offset: - self.timestamp_offset = self.frames_offset - if self.frames_np is None: - self.frames_np = frame_np.copy() - else: - self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0) - self.lock.release() + with self.lock: + if self.frames_np is not None and self.frames_np.shape[0] > self.MAX_BUFFER_DURATION_S*self.RATE: + self.frames_offset += float(self.BUFFER_TRIM_DURATION_S) + self.frames_np = self.frames_np[int(self.BUFFER_TRIM_DURATION_S*self.RATE):] + # check timestamp offset(should be >= self.frame_offset) + # this basically means that there is no speech as timestamp offset hasnt updated + # and is less than frame_offset + if self.timestamp_offset < self.frames_offset: + self.timestamp_offset = self.frames_offset + if self.frames_np is None: + self.frames_np = frame_np.copy() + else: + self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0) + self.frames_ready.set() def clip_audio_if_no_valid_segment(self): """ @@ -323,6 +330,7 @@ class ServeClientBase(object): """ logging.info("Cleaning up.") self.exit = True + self.frames_ready.set() def get_segment_no_speech_prob(self, segment): return getattr(segment, "no_speech_prob", 0)