audio: add support for raw pcm input via server flag
fixes #
This commit is contained in:
@@ -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,
|
||||||
)
|
)
|
||||||
@@ -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):
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user