audio: add support for raw pcm input via server flag

fixes #
This commit is contained in:
Aaron Boxer
2026-04-17 09:17:12 -04:00
parent b1cd51ac8a
commit f5340ddf1e
3 changed files with 36 additions and 1 deletions
+7
View File
@@ -84,6 +84,12 @@ if __name__ == "__main__":
default=50, default=50,
help='Maximum time in ms to wait for batch to fill (default: 50).' help='Maximum time in ms to wait for batch to fill (default: 50).'
) )
parser.add_argument(
'--raw_pcm_input',
action='store_true',
help='Expect raw PCM int16 audio from clients instead of float32. '
'Audio will be normalized to float32 range [-1.0, 1.0].'
)
args = parser.parse_args() args = parser.parse_args()
if args.backend == "tensorrt": if args.backend == "tensorrt":
@@ -113,4 +119,5 @@ if __name__ == "__main__":
batch_enabled=args.batch_inference, batch_enabled=args.batch_inference,
batch_max_size=args.batch_max_size, batch_max_size=args.batch_max_size,
batch_window_ms=args.batch_window_ms, batch_window_ms=args.batch_window_ms,
raw_pcm_input=args.raw_pcm_input,
) )
+22
View File
@@ -169,6 +169,28 @@ class TestTranscriptionServerGetAudio(unittest.TestCase):
result = self.server.get_audio_from_websocket(ws) result = self.server.get_audio_from_websocket(ws)
np.testing.assert_array_almost_equal(result, audio) np.testing.assert_array_almost_equal(result, audio)
def test_raw_pcm_input_normalizes_int16(self):
import numpy as np
self.server.raw_pcm_input = True
ws = MagicMock()
pcm = np.array([0, 16384, -16384, 32767], dtype=np.int16)
ws.recv.return_value = pcm.tobytes()
result = self.server.get_audio_from_websocket(ws)
expected = pcm.astype(np.float32) / 32768.0
np.testing.assert_array_almost_equal(result, expected)
self.assertTrue(result.dtype == np.float32)
self.assertTrue(np.all(result >= -1.0))
self.assertTrue(np.all(result <= 1.0))
def test_raw_pcm_input_off_reads_float32(self):
import numpy as np
self.server.raw_pcm_input = False
ws = MagicMock()
audio = np.array([0.5, -0.5], dtype=np.float32)
ws.recv.return_value = audio.tobytes()
result = self.server.get_audio_from_websocket(ws)
np.testing.assert_array_almost_equal(result, audio)
class TestTranscriptionServerHandleNewConnection(unittest.TestCase): class TestTranscriptionServerHandleNewConnection(unittest.TestCase):
def setUp(self): def setUp(self):
+7 -1
View File
@@ -160,6 +160,7 @@ class TranscriptionServer:
self.use_vad = True self.use_vad = True
self.single_model = False self.single_model = False
self.batch_config = None self.batch_config = None
self.raw_pcm_input = False
def initialize_client( def initialize_client(
self, websocket, options, faster_whisper_custom_model_path, self, websocket, options, faster_whisper_custom_model_path,
@@ -316,6 +317,9 @@ class TranscriptionServer:
frame_data = websocket.recv() frame_data = websocket.recv()
if frame_data == b"END_OF_AUDIO": if frame_data == b"END_OF_AUDIO":
return False return False
if self.raw_pcm_input:
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) return np.frombuffer(frame_data, dtype=np.float32)
def handle_new_connection(self, websocket, faster_whisper_custom_model_path, def handle_new_connection(self, websocket, faster_whisper_custom_model_path,
@@ -431,7 +435,8 @@ class TranscriptionServer:
cors_origins: Optional[str] = None, cors_origins: Optional[str] = None,
batch_enabled=False, batch_enabled=False,
batch_max_size=8, batch_max_size=8,
batch_window_ms=50): batch_window_ms=50,
raw_pcm_input=False):
""" """
Run the transcription server. Run the transcription server.
@@ -449,6 +454,7 @@ class TranscriptionServer:
to 50. to 50.
""" """
self.cache_path = cache_path self.cache_path = cache_path
self.raw_pcm_input = raw_pcm_input
self.client_manager = ClientManager(max_clients, max_connection_time) self.client_manager = ClientManager(max_clients, max_connection_time)
if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path): if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path):
if "/" not in faster_whisper_custom_model_path: if "/" not in faster_whisper_custom_model_path: