summaryrefslogtreecommitdiffstats
path: root/app
diff options
context:
space:
mode:
Diffstat (limited to 'app')
-rw-r--r--app/app_config.py39
-rw-r--r--app/hi.py384
-rw-r--r--app/requirements.txt7
-rw-r--r--app/shared_thread_data.py9
-rw-r--r--app/stt.py581
-rw-r--r--app/vad.py313
6 files changed, 1333 insertions, 0 deletions
diff --git a/app/app_config.py b/app/app_config.py
new file mode 100644
index 0000000..f911456
--- /dev/null
+++ b/app/app_config.py
@@ -0,0 +1,39 @@
+import os
+import sys
+import typing
+
+def getConfig(path: str) -> typing.Dict[str, typing.Union[str, float, int, bool]]:
+ # Helper functions to detect and convert the type
+ def is_int(value: str) -> bool:
+ try:
+ int(value)
+ return True
+ except ValueError:
+ return False
+
+ def is_float(value: str) -> bool:
+ try:
+ float(value)
+ return True
+ except ValueError:
+ return False
+
+ def convert_value(key: str, value: str):
+ if key.startswith(("enable_", "remove_", "use_", "clear_")):
+ return bool(int(value))
+ elif is_int(value):
+ return int(value)
+ elif is_float(value):
+ return float(value)
+ else:
+ return value
+
+ config = {}
+ with open(path, 'r') as file:
+ for line in file:
+ key_value = line.strip().split(": ", maxsplit=1)
+ key = key_value[0]
+ value = key_value[1] if len(key_value) > 1 else ""
+ config[key] = convert_value(key, value.strip())
+ return config
+
diff --git a/app/hi.py b/app/hi.py
new file mode 100644
index 0000000..0129958
--- /dev/null
+++ b/app/hi.py
@@ -0,0 +1,384 @@
+import app_config
+import argparse
+from math import floor, ceil
+import msvcrt
+from pythonosc import udp_client
+import sentencepiece as spm
+from shared_thread_data import SharedThreadData
+import stt
+import sys
+import threading
+import time
+
+TESTS_ENABLED = True
+
+# 0 = quiet, 1 = verbose, 2 = very verbose
+LOG_LEVEL = 0
+
+def get_tokenizer():
+ model_path = "./custom_unigram_tokenizer_65k/unigram.model"
+ print(f"Loading SentencePiece tokenizer from: {model_path}")
+ sp = spm.SentencePieceProcessor()
+ sp.load(model_path)
+ print(f"Successfully loaded SentencePiece model. Vocab size: {sp.get_piece_size()}")
+ return sp
+
+def parse_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--config", type=str, help="Path to config file (YAML).", required=True)
+ return parser.parse_args()
+
+def assert_equal(a, b):
+ err_msg = f"{a} != {b}"
+ assert a == b, err_msg
+
+# Turn a whitespace-delimited string into a list of strings no longer than
+# `cols`.
+# Preferentially breaks strings at whitespace boundaries. Preserves whitespace
+# between words, except if that whitespace comes between lines. Breaks words
+# longer than `cols` with a hyphen.
+def wrap_line(line: str, cols):
+ # First, split line into alternating chunks of words and whitespace.
+ def get_sequences(line):
+ is_space = False
+ sequences = []
+ seq_start = 0
+ seq_end = -1
+ for i in range(0, len(line)):
+ if line[i].isspace():
+ if is_space:
+ seq_end = i
+ continue
+ # We were looking at text, now we see whitespace.
+ seq = line[seq_start:seq_end+1]
+ if len(seq) > 0:
+ sequences.append(seq)
+ seq_start = i
+ seq_end = i
+ is_space = True
+ else:
+ if not is_space:
+ seq_end = i
+ continue
+ # We were looking at whitespace, now we see text.
+ seq = line[seq_start:seq_end+1]
+ if len(seq) > 0:
+ sequences.append(seq)
+ seq_start = i
+ seq_end = i
+ is_space = False
+ sequences.append(line[seq_start:seq_end+1])
+ return sequences
+ if TESTS_ENABLED:
+ assert_equal(get_sequences("foo"), ["foo"])
+ assert_equal(get_sequences("foo bar"), ["foo", " ", "bar"])
+ assert_equal(get_sequences(" foo bar"), [" ", "foo", " ", "bar"])
+ assert_equal(get_sequences(" foo bar"), [" ", "foo", " ", "bar"])
+
+ # Next, greedily construct lines out of those sequences.
+ # Whitespace gets treated specially. If it would push us over the limit, we
+ # end the line and drop the whitespace.
+ sequences = get_sequences(line)
+ def coalesce_sequences(sequences, cols):
+ cur_line = ""
+ lines = []
+ for seq in sequences:
+ if len(cur_line) + len(seq) <= cols:
+ cur_line += seq
+ continue
+ if seq.isspace():
+ lines.append(cur_line)
+ cur_line = ""
+ continue
+ if len(cur_line) > 0:
+ lines.append(cur_line)
+ # Edge case: text sequence is longer than a line.
+ while len(seq) > cols:
+ seq_prefix = seq[0:cols-1] + "-"
+ seq = seq[cols-1:]
+ lines.append(seq_prefix)
+ cur_line = seq
+ if len(cur_line) > 0:
+ lines.append(cur_line)
+ return lines
+ if TESTS_ENABLED:
+ assert_equal(coalesce_sequences(get_sequences("foo bar"), 3), ["foo", "bar"])
+ assert_equal(coalesce_sequences(get_sequences("foo bar"), 4), ["foo ", "bar"])
+ assert_equal(coalesce_sequences(get_sequences("foo bar"), 4), ["foo", "bar"])
+ assert_equal(coalesce_sequences(get_sequences("foobar"), 3), ["fo-", "ob-", "ar"])
+ assert_equal(coalesce_sequences(get_sequences("f obar"), 3), ["f ", "ob-", "ar"])
+
+ lines = coalesce_sequences(sequences, cols)
+
+ # Next, pad each line with whitespace.
+ def pad_lines(lines, cols):
+ for i in range(0, len(lines)):
+ lines[i] += ' ' * (cols - len(lines[i]))
+ return lines
+ if TESTS_ENABLED:
+ assert_equal(pad_lines(["foo", "ba"], 4), ["foo ", "ba "])
+ assert_equal(pad_lines(["foo"], 2), ["foo"])
+
+ return pad_lines(lines, cols)
+
+def get_blocks(lines, tokenizer, block_width, num_blocks):
+ if LOG_LEVEL == 2:
+ print(f"Lines sent to tokenizer: {''.join(lines)}")
+ tokens = tokenizer.encode_as_ids(''.join(lines))
+ if LOG_LEVEL == 2:
+ print(f"Tokens: {tokens}")
+ pieces = []
+ for tok in tokens:
+ piece = tokenizer.id_to_piece(tok)
+ pieces.append(piece)
+ if LOG_LEVEL == 2:
+ print(f"Pieces: {pieces}")
+
+ # Group tokens into blocks and pad with empty characters.
+ # Also get visual pointers - the location where each block will be rendered.
+ def get_blocks():
+ blocks = []
+ visual_pointer = 0
+ visual_pointers = []
+ for i in range(0, ceil(len(tokens) / block_width)):
+ visual_pointers.append(visual_pointer)
+ block = []
+ for j in range(0, block_width):
+ if i*block_width + j >= len(tokens):
+ # Pad block with empty characters. 65535 is a special token.
+ block += [65535] * (block_width - len(block))
+ break
+ block.append(tokens[i*block_width+j])
+ visual_pointer += len(pieces[i*block_width+j])
+ blocks.append(block)
+ return (blocks, visual_pointers)
+ blocks, visual_pointers = get_blocks()
+ if LOG_LEVEL == 2:
+ print(f"Blocks: {blocks}")
+ print(f"Visual pointers: {visual_pointers}")
+
+ # Set all blocks up to the next `num_blocks` boundary to blank tokens.
+ # This handles the edge case where a prior message wrote data there which
+ # is covering up our new data.
+ def pad_blocks(blocks, visual_pointers):
+ cur_num_blocks = len(blocks)
+ num_pad_blocks = num_blocks - cur_num_blocks
+ for i in range(0, num_pad_blocks):
+ blocks.append([65535] * block_width)
+ visual_pointers.append(255)
+ return blocks, visual_pointers
+ blocks, visual_pointers = pad_blocks(blocks, visual_pointers)
+ if LOG_LEVEL == 2:
+ print(f"Blocks (padded): {blocks}")
+ print(f"Visual pointers (padded): {visual_pointers}")
+
+ return blocks, visual_pointers
+
+def calc_diff(prev_blocks, prev_visual_pointers, cur_blocks,
+ cur_visual_pointers):
+ diff_indices = []
+ diff_blocks = []
+ diff_visual_pointers = []
+
+ for i in range(0, len(cur_blocks)):
+ if i >= len(prev_blocks):
+ diff_blocks.append(cur_blocks[i])
+ diff_visual_pointers.append(cur_visual_pointers[i])
+ diff_indices.append(i)
+ continue
+ if prev_blocks[i] != cur_blocks[i] or prev_visual_pointers[i] != cur_visual_pointers[i]:
+ diff_blocks.append(cur_blocks[i])
+ diff_visual_pointers.append(cur_visual_pointers[i])
+ diff_indices.append(i)
+
+ return diff_indices, diff_blocks, diff_visual_pointers
+
+def send_data(osc_client, indices, blocks, visual_pointers):
+ def split_blocks_by_byte(blocks):
+ blocks_byte00 = []
+ blocks_byte01 = []
+ for block in blocks:
+ block_byte00 = []
+ block_byte01 = []
+ for datum in block:
+ block_byte00.append((datum >> 0) & 0xFF)
+ block_byte01.append((datum >> 8) & 0xFF)
+ blocks_byte00.append(block_byte00)
+ blocks_byte01.append(block_byte01)
+ return blocks_byte00, blocks_byte01
+
+ blocks_byte00, blocks_byte01 = split_blocks_by_byte(blocks)
+ if LOG_LEVEL == 2:
+ print(f"Blocks (byte 00): {blocks_byte00}")
+ print(f"Blocks (byte 01): {blocks_byte01}")
+
+ def send_osc(osc_client, addr, data):
+ #print(f"Sending {data} to {addr}")
+ osc_client.send_message(addr, data)
+
+ for i in range(0, len(blocks)):
+ lp_int = indices[i]
+ lp_param = "_Unigram_Letter_Grid_OSC_Pointer"
+ addr = "/avatar/parameters/" + lp_param
+ send_osc(osc_client, addr, lp_int)
+
+ vp_float = (-127.5 + visual_pointers[i]) / 127.5
+ vp_param = f"_Unigram_Letter_Grid_OSC_Visual_Pointer"
+ addr = "/avatar/parameters/" + vp_param
+ send_osc(osc_client, addr, vp_float)
+ if LOG_LEVEL == 2:
+ print(f"Sending block {blocks[i]} at {visual_pointers[i]} index {indices[i]}")
+ for j in range(0, len(blocks[i])):
+ byte00_float = (-127.5 + blocks_byte00[i][j]) / 127.5
+ byte01_float = (-127.5 + blocks_byte01[i][j]) / 127.5
+ byte00_param = f"_Unigram_Letter_Grid_OSC_Datum{j:02}_Byte00"
+ byte01_param = f"_Unigram_Letter_Grid_OSC_Datum{j:02}_Byte01"
+ addr = "/avatar/parameters/" + byte00_param
+ send_osc(osc_client, addr, byte00_float)
+ addr = "/avatar/parameters/" + byte01_param
+ send_osc(osc_client, addr, byte01_float)
+ time.sleep(0.34)
+
+def getOscClient(ip = "127.0.0.1", port = 9000):
+ return udp_client.SimpleUDPClient(ip, port)
+
+class InputState:
+ def __init__(self):
+ self.page = 0
+ # Initialize the known state of the board to empty array. This will cause
+ # our paging logic to re-send everything the first time around.
+ self.blocks = []
+ self.visual_pointers = []
+ pass
+
+def handle_input(state: InputState, line: str, tokenizer, osc_client, cfg):
+ line_wrapped = wrap_line(line, cfg["cols"])
+ if TESTS_ENABLED:
+ for line in line_wrapped:
+ assert_equal(len(line), cfg["cols"])
+ if LOG_LEVEL == 2:
+ print(f"Wrapped lines: {line_wrapped}")
+
+ # Get several blank lines whenever we roll over.
+ # It's better for the reader to have some continuity when the board pages
+ # over. If we simply replaced the entire screen, it would be harder to
+ # understand.
+ line_rollover = cfg["rows"] - 2
+ blank_line = ' ' * cfg["cols"]
+ # We show a full page, then only `line_rollover` additional lines per page.
+ end_ptr = cfg["rows"]
+ which_page = 0
+ while end_ptr < len(line_wrapped):
+ end_ptr += line_rollover
+ which_page += 1
+ if state.page != which_page:
+ state.blocks = []
+ state.visual_pointers = []
+ state.page = which_page
+ line_wrapped = line_wrapped[end_ptr-cfg["rows"]:]
+
+ # Get blocks and visual pointers.
+ blocks, visual_pointers = get_blocks(line_wrapped, tokenizer,
+ cfg["block_width"], cfg["num_blocks"])
+
+ # Note that because we only send one page of data at a time, we don't have
+ # to worry about wrapping visual pointers! We will basically never run out
+ # of space.
+ indices, diff_blocks, diff_visual_pointers = calc_diff(state.blocks, state.visual_pointers, blocks, visual_pointers)
+ indices = [idx % cfg["num_blocks"] for idx in indices]
+ # Send only one block at a time to make things snappier in interactive use
+ # case.
+ # TODO use a continuation (yield) instead of returning. Then we can be a
+ # little lighter on the cpu. Measurements show that this script is
+ # already very light but we're clearly wasting a lot of work by
+ # re-tokenizing the entire input every time we send a block.
+ if len(indices) == 0:
+ return
+ if indices[0] == len(state.blocks):
+ state.blocks.append(diff_blocks[0])
+ state.visual_pointers.append(diff_visual_pointers[0])
+ elif indices[0] > len(state.blocks):
+ print(f"This should never happen!")
+ sys.exit(1)
+ else:
+ state.blocks[indices[0]] = diff_blocks[0]
+ state.visual_pointers[indices[0]] = diff_visual_pointers[0]
+
+ send_data(osc_client, [indices[0]], [diff_blocks[0]], [diff_visual_pointers[0]])
+
+def osc_thread(shared_data: SharedThreadData):
+ tokenizer = get_tokenizer()
+ osc_client = getOscClient()
+
+ # Prime the board
+ print("Priming the board")
+ input_state = InputState()
+ handle_input(input_state, "", tokenizer, osc_client, shared_data.cfg)
+
+ while not shared_data.exit_event.is_set():
+ word_copy = ""
+ with shared_data.word_lock:
+ word_copy = shared_data.word
+ handle_input(input_state, word_copy, tokenizer, osc_client, shared_data.cfg)
+ time.sleep(0.01)
+
+if __name__ == "__main__":
+ cli_args = parse_args()
+ cfg = app_config.getConfig(cli_args.config)
+ shared_data = SharedThreadData(cfg)
+ osc_thread = threading.Thread(
+ target=osc_thread,
+ args=(shared_data,))
+ osc_thread.start()
+
+ transcribe_thread = threading.Thread(
+ target=stt.transcriptionThread,
+ args=(shared_data,))
+ transcribe_thread.start()
+
+ word_is_over = False
+ local_word = ""
+ while True:
+ char_bytes = msvcrt.getch()
+ if char_bytes == b'\x03': # ctrl+C
+ break
+
+ time.sleep(0.1)
+ continue
+
+
+ try:
+ char = char_bytes.decode('utf-8')
+ if char == '\r' or char == '\n':
+ word_is_over = True
+ continue
+ except UnicodeDecodeError:
+ print(f"Unsupported character: {char_bytes}")
+ if char_bytes == b'\x00' or char_bytes == b'\xe0':
+ special_char = msvcrt.getch()
+ continue
+
+ if char_bytes == b'\x03': # ctrl+C
+ break
+ elif char_bytes == b'\x08': # backspace
+ with shared_data.word_lock:
+ shared_data.word = shared_data.word[:-1]
+ local_word = shared_data.word
+ elif char_bytes == b'\x0c': # ctrl+L
+ with shared_data.word_lock:
+ shared_data.word = ""
+ local_word = shared_data.word
+ elif word_is_over:
+ with shared_data.word_lock:
+ shared_data.word = char
+ local_word = shared_data.word
+ word_is_over = False
+ else:
+ with shared_data.word_lock:
+ shared_data.word += char
+ local_word = shared_data.word
+ print(local_word + "_")
+ shared_data.exit_event.set()
+ osc_thread.join()
+ transcribe_thread.join()
+
diff --git a/app/requirements.txt b/app/requirements.txt
new file mode 100644
index 0000000..4e79312
--- /dev/null
+++ b/app/requirements.txt
@@ -0,0 +1,7 @@
+faster-whisper
+langcodes
+pyaudio
+pydub
+python-osc
+sentencepiece
+
diff --git a/app/shared_thread_data.py b/app/shared_thread_data.py
new file mode 100644
index 0000000..ba0a419
--- /dev/null
+++ b/app/shared_thread_data.py
@@ -0,0 +1,9 @@
+import threading
+
+class SharedThreadData:
+ def __init__(self, cfg):
+ self.word = ""
+ self.word_lock = threading.Lock()
+ self.exit_event = threading.Event()
+ self.cfg = cfg
+
diff --git a/app/stt.py b/app/stt.py
new file mode 100644
index 0000000..34ef2e9
--- /dev/null
+++ b/app/stt.py
@@ -0,0 +1,581 @@
+from faster_whisper import WhisperModel
+import langcodes
+import numpy as np
+import os
+import pyaudio
+from pydub import AudioSegment
+from shared_thread_data import SharedThreadData
+import sys
+import time
+import typing
+import vad
+
+class AudioStream():
+ FORMAT = pyaudio.paInt16
+ # Size of each frame (audio sample), in bytes. If you change FORMAT, make
+ # sure this stays up to date!
+ FRAME_SZ = 2
+ # Frames per second.
+ FPS = 16000
+ CHANNELS = 1
+ def __init__(self):
+ pass
+
+ def getSamples(self) -> bytes:
+ raise NotImplementedError("getSamples is not implemented!")
+
+class MicStream(AudioStream):
+ CHUNK_SZ = 1024
+
+ def __init__(self, which_mic: str):
+ self.p = pyaudio.PyAudio()
+ self.stream = None
+ self.sample_rate = None
+ # Each time pyaudio gives us audio data, it's in the form of a chunk of
+ # samples. We keep these in a list to keep the audio callback as light
+ # as possible. Whenever downstream layers want data, we collapse the
+ # list into a single array of data (a bytes object).
+ self.chunks = []
+ # If set, incoming frames are simply discarded.
+ self.paused = False
+
+ print(f"Finding mic {which_mic}", file=sys.stderr)
+ self.dumpMicDevices()
+
+ got_match = False
+ device_index = -1
+ if which_mic == "index":
+ target_str = "Digital Audio Interface"
+ elif which_mic == "focusrite":
+ target_str = "Focusrite"
+ elif which_mic == "motu":
+ target_str = "In 1-2 (MOTU M Series)"
+ elif which_mic == "beyond":
+ target_str = "Microphone (Beyond)"
+ else:
+ print(f"Mic {which_mic} requested, treating it as a numerical " +
+ "device ID", file=sys.stderr)
+ device_index = int(which_mic)
+ got_match = True
+ if not got_match:
+ info = self.p.get_host_api_info_by_index(0)
+ numdevices = info.get('deviceCount')
+ for i in range(0, numdevices):
+ if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
+ device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name')
+ if target_str in device_name:
+ print(f"Got matching mic: {device_name}",
+ file=sys.stderr)
+ device_index = i
+ got_match = True
+ break
+ if not got_match:
+ raise KeyError(f"Mic {which_mic} not found")
+
+ info = self.p.get_device_info_by_host_api_device_index(0, device_index)
+ print(f"Found mic {which_mic}: {info['name']}", file=sys.stderr)
+ self.sample_rate = int(info['defaultSampleRate'])
+ print(f"Mic sample rate: {self.sample_rate}", file=sys.stderr)
+
+ self.stream = self.p.open(
+ rate=self.sample_rate,
+ channels=AudioStream.CHANNELS,
+ format=AudioStream.FORMAT,
+ input=True,
+ frames_per_buffer=MicStream.CHUNK_SZ,
+ input_device_index=device_index,
+ stream_callback=self.onAudioFramesAvailable)
+
+ self.stream.start_stream()
+
+ AudioStream.__init__(self)
+
+ def pause(self, state: bool = True):
+ self.paused = state
+
+ def dumpMicDevices(self):
+ info = self.p.get_host_api_info_by_index(0)
+ numdevices = info.get('deviceCount')
+
+ for i in range(0, numdevices):
+ if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
+ device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name')
+ print("Input Device id ", i, " - ", device_name)
+
+ def onAudioFramesAvailable(self,
+ frames,
+ frame_count,
+ time_info,
+ status_flags):
+ if self.paused:
+ # Don't literally pause, just start returning silence. This allows
+ # the `min_segment_age_s` check to work while paused.
+ n_frames = int(frame_count * AudioStream.FPS /
+ float(self.sample_rate))
+ self.chunks.append(np.zeros(n_frames,
+ dtype=np.int16).tobytes())
+ return (frames, pyaudio.paContinue)
+
+ decimated = b''
+ # In pyaudio, a `frame` is a single sample of audio data.
+ frame_len = AudioStream.FRAME_SZ
+ next_frame = 0.0
+ # The mic probably has a higher sample rate than Whisper wants, so
+ # decrease the sample rate by dropping samples. Note that this
+ # algorithm only works if the mic's rate is higher than whisper's
+ # expected rate.
+ keep_every = float(self.sample_rate) / AudioStream.FPS
+ for i in range(frame_count):
+ if i >= next_frame:
+ decimated += frames[i*frame_len:(i+1)*frame_len]
+ next_frame += keep_every
+ self.chunks.append(decimated)
+
+ return (frames, pyaudio.paContinue)
+
+ # Get audio data and the corresponding timestamp.
+ def getSamples(self) -> bytes:
+ chunks = self.chunks
+ self.chunks = []
+ result = b''.join(chunks)
+ return result
+
+class AudioCollector:
+ def __init__(self, stream: AudioStream):
+ self.stream = stream
+ self.frames = b''
+ # Note: by design, this is the only spot where we anchor our timestamps
+ # against the real world. This is done to make it possible to profile
+ # test cases which read from disk (at much faster than real speed) in
+ # the same way that we profile real-time data.
+ self.wall_ts = time.time()
+
+ def getAudio(self) -> bytes:
+ frames = self.stream.getSamples()
+ if frames:
+ self.frames += frames
+ return self.frames
+
+ def dropAudioPrefix(self, dur_s: float) -> bytes:
+ n_bytes = int(dur_s * AudioStream.FPS) * self.stream.FRAME_SZ
+ n_bytes = min(n_bytes, len(self.frames))
+ cut_portion = self.frames[:n_bytes]
+ self.frames = self.frames[n_bytes:]
+ self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.FPS
+ return cut_portion
+
+ def dropAudioPrefixByFrames(self, dur_frames: int) -> bytes:
+ n_bytes = dur_frames * self.stream.FRAME_SZ
+ n_bytes = min(n_bytes, len(self.frames))
+ cut_portion = self.frames[:n_bytes]
+ self.frames = self.frames[n_bytes:]
+ self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.FPS
+ return cut_portion
+
+ def keepLast(self, dur_s: float) -> bytes:
+ drop_len = max(0, self.duration() - dur_s)
+ return self.dropAudioPrefix(drop_len)
+
+ def dropAudio(self):
+ self.wall_ts += self.duration()
+ cut_portion = self.frames
+ self.frames = b''
+ return cut_portion
+
+ def duration(self):
+ return len(self.frames) / (AudioStream.FPS * self.stream.FRAME_SZ)
+
+ def begin(self):
+ return self.wall_ts
+
+ def now(self):
+ return self.begin() + self.duration()
+
+class AudioCollectorFilter:
+ def __init__(self, parent: AudioCollector):
+ self.parent = parent
+
+ def getAudio(self) -> bytes:
+ return self.parent.getAudio()
+ def dropAudioPrefix(self, dur_s: float):
+ return self.parent.dropAudioPrefix(dur_s)
+ def dropAudioPrefixByFrames(self, dur_frames: int):
+ return self.parent.dropAudioPrefixByFrames(dur_frames)
+ def keepLast(self, dur_s):
+ return self.parent.keepLast(dur_s)
+ def dropAudio(self):
+ return self.parent.dropAudio()
+ def duration(self):
+ return self.parent.duration()
+ def begin(self):
+ return self.parent.begin()
+ def now(self):
+ return self.parent.now()
+
+# Audio collector that enforces a minimum length on its audio data.
+class LengthEnforcingAudioCollector(AudioCollectorFilter):
+ def __init__(self, parent: AudioCollector, min_duration_s: float):
+ AudioCollectorFilter.__init__(self, parent)
+ self.min_duration_s = min_duration_s
+
+ def getAudio(self) -> bytes:
+ audio = self.parent.getAudio()
+ min_duration_frames = int(self.min_duration_s * AudioStream.FPS)
+ pad_len_frames = max(0, min_duration_frames - int(len(audio) /
+ AudioStream.FRAME_SZ))
+ pad = np.zeros(pad_len_frames, dtype=np.int16).tobytes()
+ return pad + audio
+
+class NormalizingAudioCollector(AudioCollectorFilter):
+ def __init__(self, parent: AudioCollector):
+ AudioCollectorFilter.__init__(self, parent)
+
+ def getAudio(self) -> bytes:
+ audio = self.parent.getAudio()
+
+ audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ,
+ frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS)
+ audio = audio.normalize()
+
+ frames = np.array(audio.get_array_of_samples())
+ frames = np.int16(frames).tobytes()
+
+ return frames
+
+class CompressingAudioCollector(AudioCollectorFilter):
+ def __init__(self, parent: AudioCollector):
+ AudioCollectorFilter.__init__(self, parent)
+
+ def getAudio(self) -> bytes:
+ audio = self.parent.getAudio()
+
+ audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ,
+ frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS)
+ # subtle compression has a slight positive effect on my benchmark
+ audio = audio.compress_dynamic_range(threshold=-10, ratio=2.0)
+
+ frames = np.array(audio.get_array_of_samples())
+ frames = np.int16(frames).tobytes()
+
+ return frames
+
+class AudioSegmenter:
+ def __init__(self,
+ min_silence_ms=250,
+ max_speech_s=5):
+ self.vad_options = vad.VadOptions(
+ min_silence_duration_ms=min_silence_ms,
+ max_speech_duration_s=max_speech_s)
+ pass
+
+ def segmentAudio(self, audio: bytes):
+ audio = np.frombuffer(audio,
+ dtype=np.int16).flatten().astype(np.float32) / 32768.0
+ return vad.get_speech_timestamps(audio, vad_options=self.vad_options)
+
+ # Returns the stable cutoff (if any) and whether there are any segments.
+ def getStableCutoff(self, audio: bytes) -> typing.Tuple[int, bool]:
+ min_delta_frames = int((self.vad_options.min_silence_duration_ms *
+ AudioStream.FPS) / 1000.0)
+ cutoff = None
+
+ last_end = None
+ segments = self.segmentAudio(audio)
+
+ for i in range(len(segments)):
+ s = segments[i]
+ #print(f"s: {s}")
+ #print(f"last_end: {last_end}")
+
+ if last_end:
+ delta_frames = s['start'] - last_end
+ #print(f"delta frames: {delta_frames}")
+ if delta_frames > min_delta_frames:
+ cutoff = s['start']
+ else:
+ last_end = s['end']
+
+ if i == len(segments) - 1:
+ now = int(len(audio) / AudioStream.FRAME_SZ)
+ #print(f"now: {now}")
+ #print(f"min d: {min_delta_frames}")
+ delta_frames = now - s['end']
+ if delta_frames > min_delta_frames:
+ cutoff = now - int(min_delta_frames / 2)
+
+ return (cutoff, len(segments) > 0)
+
+# A segment of transcribed audio. `start_ts` and `end_ts` are floating point
+# number of seconds since the beginning of audio data.
+class Segment:
+ def __init__(self,
+ transcript: str,
+ start_ts: float,
+ end_ts: float,
+ wall_ts: float,
+ avg_logprob: float,
+ no_speech_prob: float,
+ compression_ratio: float):
+ self.transcript = transcript
+ # start_ts, end_ts are timestamps in seconds relative to `wall_ts`.
+ self.start_ts = start_ts
+ self.end_ts = end_ts
+ # wall_ts is the time.time() at which the oldest audio sample leading
+ # to this transcript was collected.
+ self.wall_ts = wall_ts
+ self.avg_logprob = avg_logprob
+ self.no_speech_prob = no_speech_prob
+ self.compression_ratio = compression_ratio
+
+ def __str__(self):
+ ts = f"(ts: {self.start_ts}-{self.end_ts}) "
+
+ wall_ts_start = datetime.utcfromtimestamp(self.start_ts + self.wall_ts).strftime('%H:%M:%S')
+ wall_ts_end = datetime.utcfromtimestamp(self.end_ts + self.wall_ts).strftime('%H:%M:%S')
+ wall_ts = f"(wall ts: {wall_ts_start}-{wall_ts_end}) "
+
+ no_speech = f"(no_speech: {self.no_speech_prob}) "
+ avg_logprob = f"(avg_logprob: {self.avg_logprob}) "
+ return f"{self.transcript} " + ts + wall_ts + no_speech + avg_logprob
+
+class Whisper:
+ def __init__(self,
+ collector: AudioCollector,
+ cfg: typing.Dict):
+ self.collector = collector
+ self.model = None
+ self.cfg = cfg
+
+ abspath = os.path.abspath(__file__)
+ my_dir = os.path.dirname(abspath)
+ parent_dir = os.path.dirname(my_dir)
+
+ model_str = cfg["model"]
+ model_root = os.path.join(parent_dir, "Models",
+ os.path.normpath(model_str))
+ print(f"Model {cfg['model']} will be saved to {model_root}",
+ file=sys.stderr)
+
+ model_device = "cuda"
+ if cfg["use_cpu"]:
+ model_device = "cpu"
+
+ already_downloaded = os.path.exists(model_root)
+
+ self.model = WhisperModel(model_str,
+ device = model_device,
+ device_index = cfg["gpu_idx"],
+ compute_type = cfg["compute_type"],
+ download_root = model_root,
+ local_files_only = already_downloaded)
+
+ def transcribe(self, frames: bytes = None) -> typing.List[Segment]:
+ if frames is None:
+ frames = self.collector.getAudio()
+ # Convert from signed 16-bit int [-32768, 32767] to signed 32-bit float on
+ # [-1, 1].
+ audio = np.frombuffer(frames,
+ dtype=np.int16).flatten().astype(np.float32) / 32768.0
+
+ t0 = time.time()
+ segments, info = self.model.transcribe(
+ audio,
+ language = langcodes.find(self.cfg["language"]).language,
+ vad_filter = True,
+ temperature=0.0,
+ without_timestamps = False)
+ res = []
+ for s in segments:
+ # Manual touchup. I see a decent number of hallucinations sneaking
+ # in with high `no_speech_prob` and modest `avg_logprob`.
+ if s.no_speech_prob > 0.6 and s.avg_logprob < -0.5:
+ if self.cfg["enable_debug_mode"]:
+ print(f"Drop probable hallucination (case 1) " +
+ f"(text='{s.text}', " +
+ f"no_speech_prob={s.no_speech_prob}, " +
+ f"avg_logprob={s.avg_logprob})", file=sys.stderr)
+ continue
+ # Another touchup targeted at the vexatious "thanks for watching!"
+ # hallucination. This triggers a lot when listening to
+ # instrumental/electronic music.
+ if s.no_speech_prob > 0.15 and s.avg_logprob < -0.7:
+ if self.cfg["enable_debug_mode"]:
+ print(f"Drop probable hallucination (case 2) " +
+ f"(text='{s.text}', " +
+ f"no_speech_prob={s.no_speech_prob}, " +
+ f"avg_logprob={s.avg_logprob})", file=sys.stderr)
+ continue
+ if self.cfg["enable_debug_mode"]:
+ print(f"s get: {s}")
+ if s.avg_logprob < -1.0:
+ continue
+ if s.compression_ratio > 2.4:
+ continue
+ res.append(Segment(s.text, s.start, s.end,
+ self.collector.begin(),
+ s.avg_logprob, s.no_speech_prob, s.compression_ratio))
+ t1 = time.time()
+ if self.cfg["enable_debug_mode"]:
+ print(f"Transcription latency (s): {t1 - t0}")
+ return res
+
+class TranscriptCommit:
+ def __init__(self,
+ delta: str,
+ preview: str,
+ latency_s: float = None,
+ thresh_at_commit: int = None,
+ audio: bytes = None,
+ duration_s: float = None,
+ start_ts: float = None):
+ self.delta = delta
+ self.preview = preview
+ self.latency_s = latency_s
+ self.thresh_at_commit = thresh_at_commit
+ self.audio = audio
+ # Time at which the commit is generated
+ self.ts = time.time()
+ # Time corresponding to the start of the segment
+ self.start_ts = start_ts
+ # The duration of the audio segment, in seconds.
+ self.duration_s = duration_s
+
+
+class VadCommitter:
+ def __init__(self,
+ cfg: typing.Dict,
+ collector: AudioCollector,
+ whisper: Whisper,
+ segmenter: AudioSegmenter):
+ self.cfg = cfg
+ self.collector = collector
+ self.whisper = whisper
+ self.segmenter = segmenter
+
+ def getDelta(self) -> TranscriptCommit:
+ audio = self.collector.getAudio()
+ stable_cutoff, has_audio = self.segmenter.getStableCutoff(audio)
+
+ delta = ""
+ commit_audio = None
+ latency_s = None
+ duration_s = self.collector.duration()
+ start_ts = self.collector.begin()
+
+ if has_audio and stable_cutoff:
+ #print(f"stable cutoff get: {stable_cutoff}", file=sys.stderr)
+ latency_s = self.collector.now() - self.collector.begin()
+ duration_s = stable_cutoff / AudioStream.FPS
+ start_ts = self.collector.begin()
+ commit_audio = self.collector.dropAudioPrefixByFrames(stable_cutoff)
+
+ segments = self.whisper.transcribe(commit_audio)
+ delta = ''.join(s.transcript for s in segments)
+ audio = self.collector.getAudio()
+ if self.cfg["enable_debug_mode"]:
+ for s in segments:
+ print(f"commit segment: {s}", file=sys.stderr)
+ print(f"delta get: {delta}", file=sys.stderr)
+
+ if False:
+ ts = datetime.fromtimestamp(self.collector.now() - latency_s)
+ filename = str(ts.strftime('%Y_%m_%d__%H-%M-%S')) + ".wav"
+ saveAudio(commit_audio, filename)
+
+ preview = ""
+ if self.cfg["enable_previews"] and has_audio:
+ segments = self.whisper.transcribe(audio)
+ preview = "".join(s.transcript for s in segments)
+
+ if not has_audio:
+ #print("VAD detects no audio, skip transcription", file=sys.stderr)
+ self.collector.keepLast(1.0)
+
+ return TranscriptCommit(
+ delta.strip(),
+ preview.strip(),
+ latency_s,
+ audio=audio,
+ duration_s=duration_s,
+ start_ts=start_ts)
+
+def transcriptionThread(shared_data: SharedThreadData):
+ last_stable_commit = None
+
+ stream = MicStream(shared_data.cfg["microphone"])
+ collector = AudioCollector(stream)
+ collector = NormalizingAudioCollector(collector)
+ collector = CompressingAudioCollector(collector)
+ whisper = Whisper(collector, shared_data.cfg)
+ segmenter = AudioSegmenter(min_silence_ms=shared_data.cfg["min_silence_duration_ms"],
+ max_speech_s=shared_data.cfg["max_speech_duration_s"])
+ committer = VadCommitter(shared_data.cfg, collector, whisper, segmenter)
+
+ transcript = ""
+ preview = ""
+
+ while not shared_data.exit_event.is_set():
+ time.sleep(shared_data.cfg["transcription_loop_delay_ms"] / 1000.0);
+
+ op = None
+
+ commit = committer.getDelta()
+
+ if len(commit.delta) > 0 or len(commit.preview) > 0:
+ # Avoid re-sending text after long pauses. User controls the length
+ # of the pause in the UI.
+ if shared_data.cfg["reset_after_silence_s"] > 0:
+ silence_duration = 0
+ if last_stable_commit:
+ last_commit_end_ts = \
+ last_stable_commit.start_ts + \
+ last_stable_commit.duration_s
+ silence_duration = commit.start_ts - last_commit_end_ts
+ if silence_duration > shared_data.cfg["reset_after_silence_s"]:
+ print(f"Resetting transcript after {silence_duration}-second "
+ "silence", file=sys.stderr)
+ transcript = ""
+ preview = ""
+ if commit.delta:
+ last_stable_commit = commit
+
+ # Hard-cap displayed transcript length at 4k characters to prevent
+ # runaway memory use in UI. Keep the full transcript to avoid
+ # breaking OSC pager.
+ transcript = transcript[-4096:]
+ def join_segments(a, b):
+ if len(a) > 0 and a[-1] != ' ':
+ return a + ' ' + b
+ else:
+ return a + b
+ transcript = join_segments(transcript, commit.delta)
+ preview = commit.preview
+
+ try:
+ print(f"Transcript: {transcript}")
+ except UnicodeEncodeError:
+ print("Failed to encode transcript - discarding delta",
+ file=sys.stderr)
+ continue
+ try:
+ print(f"Preview: {preview}")
+ except UnicodeEncodeError:
+ print("Failed to encode preview - discarding", file=sys.stderr)
+
+ with shared_data.word_lock:
+ shared_data.word = join_segments(transcript, preview)
+
+ if shared_data.cfg["enable_debug_mode"]:
+ print(f"commit latency: {commit.latency_s}", file=sys.stderr)
+ print(f"commit thresh: {commit.thresh_at_commit}",
+ file=sys.stderr)
+
+ if len(transcript) > 0 and \
+ (not transcript.endswith(' ')) and \
+ (not commit.delta.startswith(' ')):
+ commit.delta = ' ' + commit.delta
+ if len(commit.delta) > 0 and \
+ (not commit.delta.endswith(' ')) and \
+ (not commit.preview.startswith(' ')):
+ commit.preview = ' ' + commit.preview
+
diff --git a/app/vad.py b/app/vad.py
new file mode 100644
index 0000000..10a72d3
--- /dev/null
+++ b/app/vad.py
@@ -0,0 +1,313 @@
+# MIT License
+#
+# Copyright (c) 2023 Guillaume Klein
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import bisect
+import functools
+import os
+import warnings
+
+from typing import List, NamedTuple, Optional
+
+import numpy as np
+
+
+# The code below is adapted from https://github.com/snakers4/silero-vad.
+class VadOptions(NamedTuple):
+ """VAD options.
+
+ Attributes:
+ threshold: Speech threshold. Silero VAD outputs speech probabilities for each audio chunk,
+ probabilities ABOVE this value are considered as SPEECH. It is better to tune this
+ parameter for each dataset separately, but "lazy" 0.5 is pretty good for most datasets.
+ min_speech_duration_ms: Final speech chunks shorter min_speech_duration_ms are thrown out.
+ max_speech_duration_s: Maximum duration of speech chunks in seconds. Chunks longer
+ than max_speech_duration_s will be split at the timestamp of the last silence that
+ lasts more than 100ms (if any), to prevent aggressive cutting. Otherwise, they will be
+ split aggressively just before max_speech_duration_s.
+ min_silence_duration_ms: In the end of each speech chunk wait for min_silence_duration_ms
+ before separating it
+ window_size_samples: Audio chunks of window_size_samples size are fed to the silero VAD model.
+ WARNING! Silero VAD models were trained using 512, 1024, 1536 samples for 16000 sample rate.
+ Values other than these may affect model performance!!
+ speech_pad_ms: Final speech chunks are padded by speech_pad_ms each side
+ """
+
+ threshold: float = 0.5
+ min_speech_duration_ms: int = 250
+ max_speech_duration_s: float = float("inf")
+ min_silence_duration_ms: int = 2000
+ window_size_samples: int = 1024
+ speech_pad_ms: int = 400
+
+
+def get_speech_timestamps(
+ audio: np.ndarray,
+ vad_options: Optional[VadOptions] = None,
+ **kwargs,
+) -> List[dict]:
+ """This method is used for splitting long audios into speech chunks using silero VAD.
+
+ Args:
+ audio: One dimensional float array.
+ vad_options: Options for VAD processing.
+ kwargs: VAD options passed as keyword arguments for backward compatibility.
+
+ Returns:
+ List of dicts containing begin and end samples of each speech chunk.
+ """
+ if vad_options is None:
+ vad_options = VadOptions(**kwargs)
+
+ threshold = vad_options.threshold
+ min_speech_duration_ms = vad_options.min_speech_duration_ms
+ max_speech_duration_s = vad_options.max_speech_duration_s
+ min_silence_duration_ms = vad_options.min_silence_duration_ms
+ window_size_samples = vad_options.window_size_samples
+ speech_pad_ms = vad_options.speech_pad_ms
+
+ if window_size_samples not in [512, 1024, 1536]:
+ warnings.warn(
+ "Unusual window_size_samples! Supported window_size_samples:\n"
+ " - [512, 1024, 1536] for 16000 sampling_rate"
+ )
+
+ sampling_rate = 16000
+ min_speech_samples = sampling_rate * min_speech_duration_ms / 1000
+ speech_pad_samples = sampling_rate * speech_pad_ms / 1000
+ max_speech_samples = (
+ sampling_rate * max_speech_duration_s
+ - window_size_samples
+ - 2 * speech_pad_samples
+ )
+ min_silence_samples = sampling_rate * min_silence_duration_ms / 1000
+ min_silence_samples_at_max_speech = sampling_rate * 98 / 1000
+
+ audio_length_samples = len(audio)
+
+ model = get_vad_model()
+ state = model.get_initial_state(batch_size=1)
+
+ speech_probs = []
+ for current_start_sample in range(0, audio_length_samples, window_size_samples):
+ chunk = audio[current_start_sample : current_start_sample + window_size_samples]
+ if len(chunk) < window_size_samples:
+ chunk = np.pad(chunk, (0, int(window_size_samples - len(chunk))))
+ speech_prob, state = model(chunk, state, sampling_rate)
+ speech_probs.append(speech_prob)
+
+ triggered = False
+ speeches = []
+ current_speech = {}
+ neg_threshold = threshold - 0.15
+
+ # to save potential segment end (and tolerate some silence)
+ temp_end = 0
+ # to save potential segment limits in case of maximum segment size reached
+ prev_end = next_start = 0
+
+ for i, speech_prob in enumerate(speech_probs):
+ if (speech_prob >= threshold) and temp_end:
+ temp_end = 0
+ if next_start < prev_end:
+ next_start = window_size_samples * i
+
+ if (speech_prob >= threshold) and not triggered:
+ triggered = True
+ current_speech["start"] = window_size_samples * i
+ continue
+
+ if (
+ triggered
+ and (window_size_samples * i) - current_speech["start"] > max_speech_samples
+ ):
+ if prev_end:
+ current_speech["end"] = prev_end
+ speeches.append(current_speech)
+ current_speech = {}
+ # previously reached silence (< neg_thres) and is still not speech (< thres)
+ if next_start < prev_end:
+ triggered = False
+ else:
+ current_speech["start"] = next_start
+ prev_end = next_start = temp_end = 0
+ else:
+ current_speech["end"] = window_size_samples * i
+ speeches.append(current_speech)
+ current_speech = {}
+ prev_end = next_start = temp_end = 0
+ triggered = False
+ continue
+
+ if (speech_prob < neg_threshold) and triggered:
+ if not temp_end:
+ temp_end = window_size_samples * i
+ # condition to avoid cutting in very short silence
+ if (window_size_samples * i) - temp_end > min_silence_samples_at_max_speech:
+ prev_end = temp_end
+ if (window_size_samples * i) - temp_end < min_silence_samples:
+ continue
+ else:
+ current_speech["end"] = temp_end
+ if (
+ current_speech["end"] - current_speech["start"]
+ ) > min_speech_samples:
+ speeches.append(current_speech)
+ current_speech = {}
+ prev_end = next_start = temp_end = 0
+ triggered = False
+ continue
+
+ if (
+ current_speech
+ and (audio_length_samples - current_speech["start"]) > min_speech_samples
+ ):
+ current_speech["end"] = audio_length_samples
+ speeches.append(current_speech)
+
+ for i, speech in enumerate(speeches):
+ if i == 0:
+ speech["start"] = int(max(0, speech["start"] - speech_pad_samples))
+ if i != len(speeches) - 1:
+ silence_duration = speeches[i + 1]["start"] - speech["end"]
+ if silence_duration < 2 * speech_pad_samples:
+ speech["end"] += int(silence_duration // 2)
+ speeches[i + 1]["start"] = int(
+ max(0, speeches[i + 1]["start"] - silence_duration // 2)
+ )
+ else:
+ speech["end"] = int(
+ min(audio_length_samples, speech["end"] + speech_pad_samples)
+ )
+ speeches[i + 1]["start"] = int(
+ max(0, speeches[i + 1]["start"] - speech_pad_samples)
+ )
+ else:
+ speech["end"] = int(
+ min(audio_length_samples, speech["end"] + speech_pad_samples)
+ )
+
+ return speeches
+
+
+def collect_chunks(audio: np.ndarray, chunks: List[dict]) -> np.ndarray:
+ """Collects and concatenates audio chunks."""
+ if not chunks:
+ return np.array([], dtype=np.float32)
+
+ return np.concatenate([audio[chunk["start"] : chunk["end"]] for chunk in chunks])
+
+
+class SpeechTimestampsMap:
+ """Helper class to restore original speech timestamps."""
+
+ def __init__(self, chunks: List[dict], sampling_rate: int, time_precision: int = 2):
+ self.sampling_rate = sampling_rate
+ self.time_precision = time_precision
+ self.chunk_end_sample = []
+ self.total_silence_before = []
+
+ previous_end = 0
+ silent_samples = 0
+
+ for chunk in chunks:
+ silent_samples += chunk["start"] - previous_end
+ previous_end = chunk["end"]
+
+ self.chunk_end_sample.append(chunk["end"] - silent_samples)
+ self.total_silence_before.append(silent_samples / sampling_rate)
+
+ def get_original_time(
+ self,
+ time: float,
+ chunk_index: Optional[int] = None,
+ ) -> float:
+ if chunk_index is None:
+ chunk_index = self.get_chunk_index(time)
+
+ total_silence_before = self.total_silence_before[chunk_index]
+ return round(total_silence_before + time, self.time_precision)
+
+ def get_chunk_index(self, time: float) -> int:
+ sample = int(time * self.sampling_rate)
+ return min(
+ bisect.bisect(self.chunk_end_sample, sample),
+ len(self.chunk_end_sample) - 1,
+ )
+
+
+@functools.lru_cache
+def get_vad_model():
+ """Returns the VAD model instance."""
+ abspath = os.path.abspath(__file__)
+ my_dir = os.path.dirname(abspath)
+ path = os.path.join(my_dir, "Models/silero_vad.onnx")
+ return SileroVADModel(path)
+
+
+class SileroVADModel:
+ def __init__(self, path):
+ try:
+ import onnxruntime
+ except ImportError as e:
+ raise RuntimeError(
+ "Applying the VAD filter requires the onnxruntime package"
+ ) from e
+
+ opts = onnxruntime.SessionOptions()
+ opts.inter_op_num_threads = 1
+ opts.intra_op_num_threads = 1
+ opts.log_severity_level = 4
+
+ self.session = onnxruntime.InferenceSession(
+ path,
+ providers=["CPUExecutionProvider"],
+ sess_options=opts,
+ )
+
+ def get_initial_state(self, batch_size: int):
+ h = np.zeros((2, batch_size, 64), dtype=np.float32)
+ c = np.zeros((2, batch_size, 64), dtype=np.float32)
+ return h, c
+
+ def __call__(self, x, state, sr: int):
+ if len(x.shape) == 1:
+ x = np.expand_dims(x, 0)
+ if len(x.shape) > 2:
+ raise ValueError(
+ f"Too many dimensions for input audio chunk {len(x.shape)}"
+ )
+ if sr / x.shape[1] > 31.25:
+ raise ValueError("Input audio chunk is too short")
+
+ h, c = state
+
+ ort_inputs = {
+ "input": x,
+ "h": h,
+ "c": c,
+ "sr": np.array(sr, dtype="int64"),
+ }
+
+ out, h, c = self.session.run(None, ort_inputs)
+ state = (h, c)
+
+ return out, state