diff --git a/whisper_live/client.py b/whisper_live/client.py index 6364f21..fd53131 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -260,6 +260,7 @@ class Client: while time.time() - self.last_response_received < self.disconnect_if_no_response_for: continue + class TranscriptionTeeClient: """ Client for handling audio recording, streaming, and transcription tasks via one or more @@ -476,6 +477,42 @@ class TranscriptionTeeClient: return process + def save_chunk(self, n_audio_file): + """ + Saves the current audio frames to a WAV file in a separate thread. + + Args: + n_audio_file (int): The index of the audio file which determines the filename. + This helps in maintaining the order and uniqueness of each chunk. + """ + t = threading.Thread( + target=self.write_audio_frames_to_file, + args=(self.frames[:], f"chunks/{n_audio_file}.wav",), + ) + t.start() + + def finalize_recording(self, n_audio_file): + """ + Finalizes the recording process by saving any remaining audio frames, + closing the audio stream, and terminating the process. + + Args: + n_audio_file (int): The file index to be used if there are remaining audio frames to be saved. + This index is incremented before use if the last chunk is saved. + """ + if self.save_output_recording and len(self.frames): + self.write_audio_frames_to_file( + self.frames[:], f"chunks/{n_audio_file}.wav" + ) + n_audio_file += 1 + self.stream.stop_stream() + self.stream.close() + self.p.terminate() + self.close_all_clients() + if self.save_output_recording: + self.write_output_recording(n_audio_file) + self.write_all_clients_srt() + def record(self): """ Record audio data from the input stream and save it to a WAV file. @@ -508,31 +545,13 @@ class TranscriptionTeeClient: # save frames if more than a minute if len(self.frames) > 60 * self.rate: if self.save_output_recording: - t = threading.Thread( - target=self.write_audio_frames_to_file, - args=( - self.frames[:], - f"chunks/{n_audio_file}.wav", - ), - ) - t.start() + self.save_chunk(n_audio_file) n_audio_file += 1 self.frames = b"" self.write_all_clients_srt() except KeyboardInterrupt: - if self.save_output_recording and len(self.frames): - self.write_audio_frames_to_file( - self.frames[:], f"chunks/{n_audio_file}.wav" - ) - n_audio_file += 1 - self.stream.stop_stream() - self.stream.close() - self.p.terminate() - self.close_all_clients() - if self.save_output_recording: - self.write_output_recording(n_audio_file) - self.write_all_clients_srt() + self.finalize_recording(n_audio_file) def write_audio_frames_to_file(self, frames, file_name): """ @@ -608,6 +627,7 @@ class TranscriptionTeeClient: raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16) return raw_data.astype(np.float32) / 32768.0 + class TranscriptionClient(TranscriptionTeeClient): """ Client for handling audio transcription tasks via a single WebSocket connection. @@ -633,8 +653,23 @@ class TranscriptionClient(TranscriptionTeeClient): transcription_client() ``` """ - def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True, save_output_recording=False, output_recording_filename="./output_recording.wav"): + def __init__( + self, + host, + port, + lang=None, + translate=False, + model="small", + use_vad=True, + save_output_recording=False, + output_recording_filename="./output_recording.wav" + ): self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad) if save_output_recording and not output_recording_filename.endswith(".wav"): raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}") - TranscriptionTeeClient.__init__(self, [self.client], save_output_recording=save_output_recording, output_recording_filename=output_recording_filename) + TranscriptionTeeClient.__init__( + self, + [self.client], + save_output_recording=save_output_recording, + output_recording_filename=output_recording_filename + )