diff --git a/README.md b/README.md index 54a4c7f..9ece51f 100644 --- a/README.md +++ b/README.md @@ -257,7 +257,7 @@ Accept raw PCM int16 audio from clients (useful for embedded devices): ```bash python3 run_server.py --port 9090 --backend faster_whisper --raw_pcm_input ``` -Audio is automatically normalized to float32 range [-1.0, 1.0]. +Audio is automatically normalized to float32 range [-1.0, 1.0]. Clients can also set `audio_format` in the initial websocket options to `float32` (default), `int16`, or `uint8`. ## Browser Extensions - Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server). diff --git a/tests/test_server_extended.py b/tests/test_server_extended.py index 611db28..419755e 100644 --- a/tests/test_server_extended.py +++ b/tests/test_server_extended.py @@ -273,6 +273,16 @@ class TestTranscriptionServerGetAudio(unittest.TestCase): self.assertTrue(np.all(result >= -1.0)) self.assertTrue(np.all(result <= 1.0)) + def test_uint8_audio_format_normalizes_unsigned_pcm(self): + import numpy as np + ws = MagicMock() + self.server.audio_formats[ws] = "uint8" + pcm = np.array([0, 128, 255], dtype=np.uint8) + ws.recv.return_value = pcm.tobytes() + result = self.server.get_audio_from_websocket(ws) + expected = (pcm.astype(np.float32) - 128.0) / 128.0 + np.testing.assert_array_almost_equal(result, expected) + def test_raw_pcm_input_off_reads_float32(self): import numpy as np self.server.raw_pcm_input = False diff --git a/whisper_live/server.py b/whisper_live/server.py index 8132119..8504683 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -178,6 +178,7 @@ class TranscriptionServer: self.single_model = False self.batch_config = None self.raw_pcm_input = False + self.audio_formats = {} self.segment_post_processor = None def initialize_client( @@ -361,7 +362,11 @@ class TranscriptionServer: frame_data = websocket.recv() if frame_data == b"END_OF_AUDIO": return False - if self.raw_pcm_input: + audio_format = self.audio_formats.get(websocket) + if audio_format == "uint8": + audio_np = np.frombuffer(frame_data, dtype=np.uint8) + return (audio_np.astype(np.float32) - 128.0) / 128.0 + if self.raw_pcm_input or audio_format == "int16": audio_np = np.frombuffer(frame_data, dtype=np.int16) return audio_np.astype(np.float32) / 32768.0 return np.frombuffer(frame_data, dtype=np.float32) @@ -378,6 +383,10 @@ class TranscriptionServer: wl_metrics.track_connection_rejected(reason="full") websocket.close() return False # Indicates that the connection should not continue + audio_format = options.get("audio_format", "float32") + if audio_format not in {"float32", "int16", "uint8"}: + raise ValueError(f"Unsupported audio_format: {audio_format}") + self.audio_formats[websocket] = audio_format if self.backend.is_tensorrt(): self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE) @@ -844,4 +853,5 @@ class TranscriptionServer: # Wait for translation thread to finish if hasattr(client, 'translation_thread') and client.translation_thread: client.translation_thread.join(timeout=2.0) - self.client_manager.remove_client(websocket) \ No newline at end of file + self.client_manager.remove_client(websocket) + self.audio_formats.pop(websocket, None)