WIP

Dockerfile

@@ -1,3 +1,5 @@
+# TODO: This doesn't work, copied over from M4T but needs an update
+
 FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && \

app.py

@@ -6,101 +6,150 @@ import gradio as gr
 import numpy as np
 import torch
 import torchaudio
-from seamless_communication.models.inference.translator import Translator
 
-
-from m4t_app import *
 from simuleval_transcoder import *
-# from simuleval_transcoder import *
 
 from pydub import AudioSegment
 import time
 from time import sleep
 
-# m4t_demo()
+from seamless_communication.cli.streaming.agents.tt_waitk_unity_s2t_m4t import (
+    TestTimeWaitKUnityS2TM4T,
+)
+
+language_code_to_name = {
+    "cmn": "Mandarin Chinese",
+    "deu": "German",
+    "eng": "English",
+    "fra": "French",
+    "spa": "Spanish",
+}
+S2ST_TARGET_LANGUAGE_NAMES = language_code_to_name.values()
+LANGUAGE_NAME_TO_CODE = {v: k for k, v in language_code_to_name.items()}
+
+DEFAULT_TARGET_LANGUAGE = "English"
+
+# TODO: Update this so it takes in target langs from input, refactor sample rate
+transcoder = SimulevalTranscoder(
+    sample_rate=48_000,
+    debug=False,
+    buffer_limit=1,
+)
+
+def start_recording():
+    logger.debug(f"start_recording: starting transcoder")
+    transcoder.start()
+
+
+def translate_audio_segment(audio):
+    logger.debug(f"translate_audio_segment: incoming audio")
+    sample_rate, data = audio
+
+    transcoder.process_incoming_bytes(data.tobytes(), 'eng', sample_rate)
 
-USE_M4T = True
+    speech_and_text_output =  transcoder.get_buffered_output()
+    if speech_and_text_output is None:
+        logger.debug("No output from transcoder.get_buffered_output()")
+        return None, None
 
-Transcoder = SimulevalTranscoder()
+    logger.debug(f"We DID get output from the transcoder! {speech_and_text_output}")
 
-def translate_audio_file_segment(audio_file):
-    print("translate_m4t state")
+    text = None
+    speech = None
 
-    return predict(
-        task_name="S2ST",
-        audio_source="microphone",
-        input_audio_mic=audio_file,
-        input_audio_file=None,
-        input_text="",
-        source_language="English",
-        target_language="Portuguese",
-    )
+    if speech_and_text_output.speech_samples:
+        speech = (speech_and_text_output.speech_samples, speech_and_text_output.speech_sample_rate)
 
+    if speech_and_text_output.text:
+        text = speech_and_text_output.text
+        if speech_and_text_output.final:
+            text += "\n"
 
-def translate_m4t_callback(
+    return speech, text
+
+def streaming_input_callback(
     audio_file, translated_audio_bytes_state, translated_text_state
 ):
-    translated_wav_segment, translated_text = translate_audio_file_segment(audio_file)
-    print('translated_audio_bytes_state', translated_audio_bytes_state)
-    print('translated_wav_segment', translated_wav_segment)
+    translated_wav_segment, translated_text = translate_audio_segment(audio_file)
+    logger.debug(f'translated_audio_bytes_state {translated_audio_bytes_state}')
+    logger.debug(f'translated_wav_segment {translated_wav_segment}')
+
+    # TODO: accumulate each segment to provide a continuous audio segment
 
-    # combine translated wav into larger..
-    if type(translated_audio_bytes_state) is not tuple:
-        translated_audio_bytes_state = translated_wav_segment
-    else:
+    if translated_wav_segment is not None:
+        sample_rate, audio_bytes = translated_wav_segment
+        audio_np_array = np.frombuffer(audio_bytes, dtype=np.float32, count=3)
 
-        translated_audio_bytes_state = (translated_audio_bytes_state[0], np.append(translated_audio_bytes_state[1], translated_wav_segment[1]))
 
-    # translated_wav_segment[1]
+        # combine translated wav
+        if type(translated_audio_bytes_state) is not tuple:
+            translated_audio_bytes_state = (sample_rate, audio_np_array)
+            # translated_audio_bytes_state = np.array([])
+        else:
+
+            translated_audio_bytes_state = (translated_audio_bytes_state[0], np.append(translated_audio_bytes_state[1], translated_wav_segment[1]))
 
+    if translated_text is not None:
+        translated_text_state += " | " + str(translated_text)
 
-    translated_text_state += " | " + str(translated_text)
+    # most_recent_input_audio_segment = (most_recent_input_audio_segment[0], np.append(most_recent_input_audio_segment[1], audio_file[1]))
+
+    # Not necessary but for readability.
+    most_recent_input_audio_segment = audio_file
+    translated_wav_segment = translated_wav_segment
+    output_translation_combined = translated_audio_bytes_state
+    stream_output_text = translated_text_state
     return [
-        audio_file,
+        most_recent_input_audio_segment,
         translated_wav_segment,
-        translated_audio_bytes_state,
-        translated_text_state,
+        output_translation_combined,
+        stream_output_text,
         translated_audio_bytes_state,
         translated_text_state,
     ]
 
 
 def clear():
-    print("Clearing State")
+    logger.debug(f"Clearing State")
     return [bytes(), ""]
 
 
 def blocks():
     with gr.Blocks() as demo:
+
+        with gr.Row():
+            # Hook this up once supported
+            target_language = gr.Dropdown(
+                label="Target language",
+                choices=S2ST_TARGET_LANGUAGE_NAMES,
+                value=DEFAULT_TARGET_LANGUAGE,
+            )
+
         translated_audio_bytes_state = gr.State(None)
         translated_text_state = gr.State("")
 
-        # input_audio = gr.Audio(label="Input Audio", type="filepath", format="mp3")
-        if USE_M4T:
-            input_audio = gr.Audio(
-                label="Input Audio",
-                type="filepath",
-                source="microphone",
-                streaming=True,
-            )
-        else:
-            input_audio = gr.Audio(
-                label="Input Audio",
-                type="filepath",
-                format="mp3",
-                source="microphone",
-                streaming=True,
-            )
+        input_audio = gr.Audio(
+            label="Input Audio",
+            # source="microphone", # gradio==3.41.0
+            sources=["microphone"], # new gradio seems to call this less often...
+            streaming=True,
+        )
+
+        # input_audio = gr.Audio(
+        #     label="Input Audio",
+        #     type="filepath",
+        #     source="microphone",
+        #     streaming=True,
+        # )
 
         most_recent_input_audio_segment = gr.Audio(
             label="Recent Input Audio Segment segments",
-            format="bytes",
+            # format="bytes",
             streaming=True
         )
-        # TODO: Should add combined input audio segments...
-
-        stream_as_bytes_btn = gr.Button("Translate most recent recording segment")
 
+        # Force translate
+        stream_as_bytes_btn = gr.Button("Force translate most recent recording segment (ask for model output)")
         output_translation_segment = gr.Audio(
             label="Translated audio segment",
             autoplay=False,
@@ -119,7 +168,7 @@ def blocks():
         stream_output_text = gr.Textbox(label="Translated text")
 
         stream_as_bytes_btn.click(
-            translate_m4t_callback,
+            streaming_input_callback,
             [input_audio, translated_audio_bytes_state, translated_text_state],
             [
                 most_recent_input_audio_segment,
@@ -131,8 +180,21 @@ def blocks():
             ],
         )
 
-        input_audio.change(
-            translate_m4t_callback,
+        # input_audio.change(
+        #     streaming_input_callback,
+        #     [input_audio, translated_audio_bytes_state, translated_text_state],
+        #     [
+        #         most_recent_input_audio_segment,
+        #         output_translation_segment,
+        #         output_translation_combined,
+        #         stream_output_text,
+        #         translated_audio_bytes_state,
+        #         translated_text_state,
+        #     ],
+        # )
+
+        input_audio.stream(
+            streaming_input_callback,
             [input_audio, translated_audio_bytes_state, translated_text_state],
             [
                 most_recent_input_audio_segment,
@@ -143,8 +205,11 @@ def blocks():
                 translated_text_state,
             ],
         )
-        # input_audio.change(stream_bytes, [input_audio, translated_audio_bytes_state, translated_text_state], [most_recent_input_audio_segment, stream_output_text, translated_audio_bytes_state, translated_text_state])
-        # input_audio.change(lambda input_audio: recorded_audio, [input_audio], [recorded_audio])
+
+        input_audio.start_recording(
+            start_recording,
+        )
+
         input_audio.clear(
             clear, None, [translated_audio_bytes_state, translated_text_state]
         )
@@ -154,6 +219,4 @@ def blocks():
 
     demo.queue().launch()
 
-
-# if __name__ == "__main__":
 blocks()

internal_demo_simuleval_transcoder.py

@@ -1,272 +0,0 @@
-from simuleval.utils.agent import build_system_from_dir
-from typing import Any, Tuple
-import numpy as np
-import soundfile
-from fairseq.data.audio.audio_utils import convert_waveform
-import io
-import asyncio
-from simuleval.data.segments import SpeechSegment, EmptySegment
-import threading
-import math
-import logging
-import sys
-from pathlib import Path
-import time
-from g2p_en import G2p
-import torch
-import traceback
-import time
-import random
-
-from .speech_and_text_output import SpeechAndTextOutput
-
-MODEL_SAMPLE_RATE = 16_000
-
-logger = logging.getLogger()
-logger.addHandler(logging.StreamHandler(sys.stdout))
-
-
-class SimulevalTranscoder:
-    def __init__(self, agent, sample_rate, debug, buffer_limit):
-        self.agent = agent
-        self.input_queue = asyncio.Queue()
-        self.output_queue = asyncio.Queue()
-        self.states = self.agent.build_states()
-        if debug:
-            self.states[0].debug = True
-        self.incoming_sample_rate = sample_rate
-        self.close = False
-        self.g2p = G2p()
-
-        # buffer all outgoing translations within this amount of time
-        self.output_buffer_idle_ms = 5000
-        self.output_buffer_size_limit = (
-            buffer_limit  # phonemes for text, seconds for speech
-        )
-        self.output_buffer_cur_size = 0
-        self.output_buffer = []
-        self.speech_output_sample_rate = None
-
-        self.last_output_ts = time.time() * 1000
-        self.timeout_ms = (
-            30000  # close the transcoder thread after this amount of silence
-        )
-        self.first_input_ts = None
-        self.first_output_ts = None
-        self.output_data_type = None  # speech or text
-        self.debug = debug
-        self.debug_ts = f"{time.time()}_{random.randint(1000, 9999)}"
-        if self.debug:
-            debug_folder = Path(__file__).resolve().parent.parent / "debug"
-            self.test_incoming_wav = soundfile.SoundFile(
-                debug_folder / f"{self.debug_ts}_test_incoming.wav",
-                mode="w+",
-                format="WAV",
-                subtype="PCM_16",
-                samplerate=self.incoming_sample_rate,
-                channels=1,
-            )
-            self.states[0].test_input_segments_wav = soundfile.SoundFile(
-                debug_folder / f"{self.debug_ts}_test_input_segments.wav",
-                mode="w+",
-                format="WAV",
-                samplerate=MODEL_SAMPLE_RATE,
-                channels=1,
-            )
-
-    def debug_log(self, *args):
-        if self.debug:
-            logger.info(*args)
-
-    @classmethod
-    def build_agent(cls, model_path):
-        logger.info(f"Building simuleval agent: {model_path}")
-        agent = build_system_from_dir(
-            Path(__file__).resolve().parent.parent / f"models/{model_path}",
-            config_name="vad_main.yaml",
-        )
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        agent.to(device, fp16=True)
-        logger.info(
-            f"Successfully built simuleval agent {model_path} on device {device}"
-        )
-
-        return agent
-
-    def process_incoming_bytes(self, incoming_bytes):
-        segment, _sr = self._preprocess_wav(incoming_bytes)
-        # # segment is array([0, 0, 0, ..., 0, 0, 0], dtype=int16)
-        self.input_queue.put_nowait(segment)
-
-    def get_input_segment(self):
-        if self.input_queue.empty():
-            return None
-        chunk = self.input_queue.get_nowait()
-        self.input_queue.task_done()
-        return chunk
-
-    def _preprocess_wav(self, data: Any) -> Tuple[np.ndarray, int]:
-        segment, sample_rate = soundfile.read(
-            io.BytesIO(data),
-            dtype="float32",
-            always_2d=True,
-            frames=-1,
-            start=0,
-            format="RAW",
-            subtype="PCM_16",
-            samplerate=self.incoming_sample_rate,
-            channels=1,
-        )
-        if self.debug:
-            self.test_incoming_wav.seek(0, soundfile.SEEK_END)
-            self.test_incoming_wav.write(segment)
-
-        segment = segment.T
-        segment, new_sample_rate = convert_waveform(
-            segment,
-            sample_rate,
-            normalize_volume=False,
-            to_mono=True,
-            to_sample_rate=MODEL_SAMPLE_RATE,
-        )
-
-        assert MODEL_SAMPLE_RATE == new_sample_rate
-        segment = segment.squeeze(axis=0)
-        return segment, new_sample_rate
-
-    def process_pipeline_impl(self, input_segment):
-        try:
-            output_segment = self.agent.pushpop(input_segment, self.states)
-            if (
-                self.states[0].first_input_ts is not None
-                and self.first_input_ts is None
-            ):
-                # TODO: this is hacky
-                self.first_input_ts = self.states[0].first_input_ts
-
-            if not output_segment.is_empty:
-                self.output_queue.put_nowait(output_segment)
-
-            if output_segment.finished:
-                self.debug_log("OUTPUT SEGMENT IS FINISHED. Resetting states.")
-
-                for state in self.states:
-                    state.reset()
-
-                if self.debug:
-                    # when we rebuild states, this value is reset to whatever
-                    # is in the system dir config, which defaults debug=False.
-                    self.states[0].debug = True
-        except Exception as e:
-            logger.error(f"Got exception while processing pipeline: {e}")
-            traceback.print_exc()
-        return input_segment
-
-    def process_pipeline_loop(self):
-        if self.close:
-            return  # closes the thread
-
-        self.debug_log("processing_pipeline")
-        while not self.close:
-            input_segment = self.get_input_segment()
-            if input_segment is None:
-                if self.states[0].is_fresh_state:  # TODO: this is hacky
-                    time.sleep(0.3)
-                else:
-                    time.sleep(0.03)
-                continue
-            self.process_pipeline_impl(input_segment)
-        self.debug_log("finished processing_pipeline")
-
-    def process_pipeline_once(self):
-        if self.close:
-            return
-
-        self.debug_log("processing pipeline once")
-        input_segment = self.get_input_segment()
-        if input_segment is None:
-            return
-        self.process_pipeline_impl(input_segment)
-        self.debug_log("finished processing_pipeline_once")
-
-    def get_output_segment(self):
-        if self.output_queue.empty():
-            return None
-
-        output_chunk = self.output_queue.get_nowait()
-        self.output_queue.task_done()
-        return output_chunk
-
-    def start(self):
-        self.debug_log("starting transcoder in a thread")
-        threading.Thread(target=self.process_pipeline_loop).start()
-
-    def first_translation_time(self):
-        return round((self.first_output_ts - self.first_input_ts) / 1000, 2)
-
-    def get_buffered_output(self) -> SpeechAndTextOutput:
-        now = time.time() * 1000
-        self.debug_log(f"get_buffered_output queue size: {self.output_queue.qsize()}")
-        while not self.output_queue.empty():
-            tmp_out = self.get_output_segment()
-            if tmp_out and len(tmp_out.content) > 0:
-                if not self.output_data_type:
-                    self.output_data_type = tmp_out.data_type
-                if len(self.output_buffer) == 0:
-                    self.last_output_ts = now
-                self._populate_output_buffer(tmp_out)
-                self._increment_output_buffer_size(tmp_out)
-
-                if tmp_out.finished:
-                    res = self._gather_output_buffer_data(final=True)
-                    self.output_buffer = []
-                    self.increment_output_buffer_size = 0
-                    self.last_output_ts = now
-                    self.first_output_ts = now
-                    return res
-
-        if len(self.output_buffer) > 0 and (
-            now - self.last_output_ts >= self.output_buffer_idle_ms
-            or self.output_buffer_cur_size >= self.output_buffer_size_limit
-        ):
-            self.last_output_ts = now
-            res = self._gather_output_buffer_data(final=False)
-            self.output_buffer = []
-            self.output_buffer_phoneme_count = 0
-            self.first_output_ts = now
-            return res
-        else:
-            return None
-
-    def _gather_output_buffer_data(self, final):
-        if self.output_data_type == "text":
-            return SpeechAndTextOutput(text=" ".join(self.output_buffer), final=final)
-        elif self.output_data_type == "speech":
-            return SpeechAndTextOutput(
-                speech_samples=self.output_buffer,
-                speech_sample_rate=MODEL_SAMPLE_RATE,
-                final=final,
-            )
-        else:
-            raise ValueError(
-                f"Invalid output buffer data type: {self.output_data_type}"
-            )
-
-    def _increment_output_buffer_size(self, segment):
-        if segment.data_type == "text":
-            self.output_buffer_cur_size += self._compute_phoneme_count(segment.content)
-        elif segment.data_type == "speech":
-            self.output_buffer_cur_size += (
-                len(segment.content) / MODEL_SAMPLE_RATE
-            )  # seconds
-
-    def _populate_output_buffer(self, segment):
-        if segment.data_type == "text":
-            self.output_buffer.append(segment.content)
-        elif segment.data_type == "speech":
-            self.output_buffer += segment.content
-        else:
-            raise ValueError(f"Invalid segment data type: {segment.data_type}")
-
-    def _compute_phoneme_count(self, string: str) -> int:
-        return len([x for x in self.g2p(string) if x != " "])

requirements.txt

@@ -1,18 +1,21 @@
+# TODO: fairseq2 install is complicated so currently done outside
+
 # fairseq2==0.1.0
 
-# Temp to skip
-# git+https://github.com/mduppes/fairseq2.git@93420c86ba01349ee8f90d7adda439b666b50557
 # git+https://github.com/facebookresearch/seamless_communication
-./seamless_communication
+# ./fairseq2
+# ./seamless_communication
 # comment this out to test fairseq1 first
 # git+https://github.com/facebookresearch/SimulEval.git
 gradio==3.41.0
 huggingface_hub==0.16.4
-torch==2.0.1
-torchaudio==2.0.2
-transformers==4.32.1
+# torch==2.1.0
+# torchaudio==2.0.2
+# transformers==4.32.1
 pydub
-
+g2p_en
+colorlog
+# git+ssh://git@github.com/facebookresearch/SimulEval.git
 
 # Can't import fairseq1 together.. causes conflict:
 #The conflict is caused by:

seamless_communication

@@ -1 +0,0 @@
-Subproject commit 02405dfd0c187d625aa66255ff8c39f98031a091

simuleval_transcoder.py

@@ -1,225 +1,455 @@
 
+from typing import Any, List, Tuple, Union, Optional
+import numpy as np
+import soundfile
+import io
+import asyncio
+from simuleval.agents.pipeline import TreeAgentPipeline
+from simuleval.agents.states import AgentStates
+from simuleval.data.segments import Segment, EmptySegment, SpeechSegment
+import threading
+import math
+import logging
+import sys
 from pathlib import Path
-from typing import Callable, Dict, List, Optional, Tuple, Union
-
+import time
+from g2p_en import G2p
 import torch
-import torch.nn as nn
-from fairseq2.assets.card import AssetCard
-from fairseq2.data import Collater
-from fairseq2.data.audio import AudioDecoder, WaveformToFbankConverter
-from fairseq2.data.text.text_tokenizer import TextTokenizer
-from fairseq2.data.typing import StringLike
-from fairseq2.generation import SequenceToTextOutput, SequenceGeneratorOptions
-from fairseq2.memory import MemoryBlock
-from fairseq2.typing import DataType, Device
-from torch import Tensor
-from enum import Enum, auto
-from seamless_communication.models.inference.ngram_repeat_block_processor import (
-    NGramRepeatBlockProcessor,
+import traceback
+import time
+import random
+import colorlog
+
+# Sanity check that pipeline is loadable
+from seamless_communication.cli.streaming.agents.tt_waitk_unity_s2t_m4t import (
+    # TestTimeWaitKUnityS2TM4T,
+    TestTimeWaitKUnityS2TM4TVAD
 )
 
-from seamless_communication.models.unity import (
-    UnitTokenizer,
-    UnitYGenerator,
-    UnitYModel,
-    load_unity_model,
-    load_unity_text_tokenizer,
-    load_unity_unit_tokenizer,
+from simuleval.utils.agent import build_system_args
+
+MODEL_SAMPLE_RATE = 16_000
+
+logger = logging.getLogger(__name__)
+logger.propagate = False
+handler = colorlog.StreamHandler(stream=sys.stdout)
+formatter = colorlog.ColoredFormatter(
+    "%(log_color)s[%(asctime)s][%(levelname)s][%(module)s]:%(reset)s %(message)s",
+    reset=True,
+    log_colors={
+        "DEBUG": "cyan",
+        "INFO": "green",
+        "WARNING": "yellow",
+        "ERROR": "red",
+        "CRITICAL": "red,bg_white",
+    },
 )
-from seamless_communication.models.unity.generator import SequenceToUnitOutput
-from seamless_communication.models.vocoder import load_vocoder_model, Vocoder
+handler.setFormatter(formatter)
+logger.addHandler(handler)
+logger.setLevel(logging.DEBUG)
 
 
+# TODO: Integrate this better so target lang and others can be changed. Also currently dependent on devserver internals
+def build_agent():
+    config = {
+        'dataloader': 'fairseq2_s2t',
+        'data_file': '/large_experiments/seamless/ust/abinesh/data/s2st50_manifests/50-10/simuleval/dev_mtedx_filt_50-10_debug.tsv',
+        'model_name': 'seamlessM4T_v2_large',
+        'device': 'cuda:0',
+        'source_segment_size': 320,
+        'waitk_lagging': 7,
+        'fixed_pre_decision_ratio': 2,
+        'init_target_tokens': '</s> __eng__',
+        'max_len_a': 0,
+        'max_len_b': 200,
+        'agent_class': 'seamless_communication.cli.streaming.agents.tt_waitk_unity_s2t_m4t.TestTimeWaitKUnityS2TM4TVAD',
+        'task': 's2st',
+        'tgt_lang': 'eng',
+        'latency_metrics': 'StartOffset EndOffset AL',
+        'output': 'TestTimeWaitKUnityS2TM4TVAD-wait7-debug'
+    }
 
-# from seamless_communication.models.streaming.agents import (
-#     SileroVADAgent,
-#     TestTimeWaitKS2TVAD,
-#     TestTimeWaitKUnityV1M4T
-# )
+    agent , _ = build_system_args(config)
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # agent.to(device, fp16=True)
+    logger.info(
+        f"Successfully built simuleval agent"
+    )
 
-from seamless_communication.cli.streaming.agents.tt_waitk_unity_s2t_m4t import (
-    TestTimeWaitKUnityS2TM4T,
-)
+    return agent
 
-from seamless_communication.cli.streaming.dataloader import Fairseq2SpeechToTextDataloader
+class SpeechAndTextOutput:
+    def __init__(
+        self,
+        text: str = None,
+        speech_samples: list = None,
+        speech_sample_rate: float = None,
+        final: bool = False,
+    ):
+        self.text = text
+        self.speech_samples = speech_samples
+        self.speech_sample_rate = speech_sample_rate
+        self.final = final
 
-### From test_pipeline
-import math
-import soundfile
-from argparse import Namespace, ArgumentParser
-from simuleval.data.segments import SpeechSegment, EmptySegment
-from simuleval.utils import build_system_from_dir
-from pathlib import Path
-import numpy as np
+class OutputSegments:
+    def __init__(self, segments: Union[List[Segment], Segment]):
+        if isinstance(segments, Segment):
+            segments = [segments]
+        self.segments: List[Segment] = [s for s in segments]
 
-class AudioFrontEnd:
-    def __init__(self, wav_file, segment_size) -> None:
-        self.samples, self.sample_rate = soundfile.read(wav_file)
-        # print(len(self.samples), self.samples[:100])
-        self.samples = self.samples.tolist()
-        self.segment_size = segment_size
-        self.step = 0
-    def send_segment(self):
-        """
-        This is the front-end logic in simuleval instance.py
-        """
-        num_samples = math.ceil(self.segment_size / 1000 * self.sample_rate)
-        print("self.segment_size", self.segment_size)
-        print('num_samples is', num_samples)
-        print('self.sample_rate is', self.sample_rate)
-        if self.step < len(self.samples):
-            if self.step + num_samples >= len(self.samples):
-                samples = self.samples[self.step :]
-                is_finished = True
-            else:
-                samples = self.samples[self.step : self.step + num_samples]
-                is_finished = False
-            self.step = min(self.step + num_samples, len(self.samples))
-            # print("len(samples) is", len(samples))
-            # import pdb
-            # pdb.set_trace()
-            segment = SpeechSegment(
-                index=self.step / self.sample_rate * 1000,
-                content=samples,
-                sample_rate=self.sample_rate,
-                finished=is_finished,
-            )
-        else:
-            # Finish reading this audio
-            segment = EmptySegment(
-                index=self.step / self.sample_rate * 1000,
-                finished=True,
-            )
-        return segment
+    @property
+    def is_empty(self):
+        return all(segment.is_empty for segment in self.segments)
 
+    @property
+    def finished(self):
+        return all(segment.finished for segment in self.segments)
 
+    def compute_length(self, g2p):
+        lengths = []
+        for segment in self.segments:
+            if segment.data_type == "text":
+                lengths.append(len([x for x in g2p(segment.content) if x != " "]))
+            elif segment.data_type == "speech":
+                lengths.append(len(segment.content) / MODEL_SAMPLE_RATE)
+            elif isinstance(segment, EmptySegment):
+                continue
+            else:
+                logger.warning(
+                    f"Unexpected data_type: {segment.data_type} not in 'speech', 'text'"
+                )
+        return max(lengths)
 
-def load_model_for_inference(
-    load_model_fn: Callable[..., nn.Module],
-    model_name_or_card: Union[str, AssetCard],
-    device: Device,
-    dtype: DataType,
-) -> nn.Module:
-    model = load_model_fn(model_name_or_card, device=device, dtype=dtype)
-    model.eval()
-    return model
+    @classmethod
+    def join_output_buffer(
+        cls, buffer: List[List[Segment]], output: SpeechAndTextOutput
+    ):
+        num_segments = len(buffer[0])
+        for i in range(num_segments):
+            segment_list = [
+                buffer[j][i]
+                for j in range(len(buffer))
+                if buffer[j][i].data_type is not None
+            ]
+            if len(segment_list) == 0:
+                continue
+            if len(set(segment.data_type for segment in segment_list)) != 1:
+                logger.warning(
+                    f"Data type mismatch at {i}: {set(segment.data_type for segment in segment_list)}"
+                )
+                continue
+            data_type = segment_list[0].data_type
+            if data_type == "text":
+                if output.text is not None:
+                    logger.warning("Multiple text outputs, overwriting!")
+                output.text = " ".join([segment.content for segment in segment_list])
+            elif data_type == "speech":
+                if output.speech_samples is not None:
+                    logger.warning("Multiple speech outputs, overwriting!")
+                speech_out = []
+                for segment in segment_list:
+                    speech_out += segment.content
+                output.speech_samples = speech_out
+                output.speech_sample_rate = MODEL_SAMPLE_RATE
+            elif isinstance(segment_list[0], EmptySegment):
+                continue
+            else:
+                logger.warning(
+                    f"Invalid output buffer data type: {data_type}, expected 'speech' or 'text"
+                )
 
-def load_model_fairseq2():
-    data_configs = dict(
-        dataloader="fairseq2_s2t",
-        data_file="/large_experiments/seamless/ust/abinesh/data/s2st50_manifests/50-10/simuleval/dev_mtedx_filt_50-10_debug.tsv",
-    )
+        return output
 
-    model_configs = dict(
-        model_name="seamlessM4T_v2_large",
-        device="cuda:0",
-        source_segment_size=320,
-        waitk_lagging=7,
-        fixed_pre_decision_ratio=2,
-        init_target_tokens="</s> __eng__",
-        max_len_a=0,
-        max_len_b=200,
-        agent_class="seamless_communication.cli.streaming.agents.tt_waitk_unity_s2t_m4t.TestTimeWaitKUnityS2TM4T",
-        task="s2st",
-        tgt_lang="eng",
-    )
+    def __repr__(self) -> str:
+        repr_str = str(self.segments)
+        return f"{self.__class__.__name__}(\n\t{repr_str}\n)"
 
-    eval_configs = dict(
-        latency_metrics="StartOffset EndOffset AL",
-        output=f"{TestTimeWaitKUnityS2TM4T.__name__}-wait{model_configs['waitk_lagging']}-debug",
-    )
 
-    model = TestTimeWaitKUnityS2TM4T({**data_configs, **model_configs, **eval_configs})
-    print("model", model)
+def convert_waveform(
+    waveform: Union[np.ndarray, torch.Tensor],
+    sample_rate: int,
+    normalize_volume: bool = False,
+    to_mono: bool = False,
+    to_sample_rate: Optional[int] = None,
+) -> Tuple[Union[np.ndarray, torch.Tensor], int]:
+    """convert a waveform:
+    - to a target sample rate
+    - from multi-channel to mono channel
+    - volume normalization
 
-    evaluate(
-        TestTimeWaitKUnityS2TM4T, {**data_configs, **model_configs, **eval_configs}
-    )
+    Args:
+        waveform (numpy.ndarray or torch.Tensor): 2D original waveform
+            (channels x length)
+        sample_rate (int): original sample rate
+        normalize_volume (bool): perform volume normalization
+        to_mono (bool): convert to mono channel if having multiple channels
+        to_sample_rate (Optional[int]): target sample rate
+    Returns:
+        waveform (numpy.ndarray): converted 2D waveform (channels x length)
+        sample_rate (float): target sample rate
+    """
+    try:
+        import torchaudio.sox_effects as ta_sox
+    except ImportError:
+        raise ImportError("Please install torchaudio: pip install torchaudio")
+
+    effects = []
+    if normalize_volume:
+        effects.append(["gain", "-n"])
+    if to_sample_rate is not None and to_sample_rate != sample_rate:
+        effects.append(["rate", f"{to_sample_rate}"])
+    if to_mono and waveform.shape[0] > 1:
+        effects.append(["channels", "1"])
+    if len(effects) > 0:
+        is_np_input = isinstance(waveform, np.ndarray)
+        _waveform = torch.from_numpy(waveform) if is_np_input else waveform
+        converted, converted_sample_rate = ta_sox.apply_effects_tensor(
+            _waveform, sample_rate, effects
+        )
+        if is_np_input:
+            converted = converted.numpy()
+        return converted, converted_sample_rate
+    return waveform, sample_rate
 
 class SimulevalTranscoder:
-    # def __init__(self, agent, sample_rate, debug, buffer_limit):
-    def __init__(self):
-        # print("MDUPPES in here", SileroVADAgent, TestTimeWaitKS2TVAD)
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    def __init__(self, sample_rate, debug, buffer_limit):
+        self.agent = build_agent()
+        self.input_queue = asyncio.Queue()
+        self.output_queue = asyncio.Queue()
+        self.states = self.agent.build_states()
+        if debug:
+            self.get_states_root().debug = True
+        self.incoming_sample_rate = sample_rate
+        self.close = False
+        self.g2p = G2p()
+
+        # buffer all outgoing translations within this amount of time
+        self.output_buffer_idle_ms = 5000
+        self.output_buffer_size_limit = (
+            buffer_limit  # phonemes for text, seconds for speech
+        )
+        self.output_buffer_cur_size = 0
+        self.output_buffer: List[List[Segment]] = []
+        self.speech_output_sample_rate = None
+
+        self.last_output_ts = time.time() * 1000
+        self.timeout_ms = (
+            30000  # close the transcoder thread after this amount of silence
+        )
+        self.first_input_ts = None
+        self.first_output_ts = None
+        self.debug = debug
+        self.debug_ts = f"{time.time()}_{random.randint(1000, 9999)}"
+        if self.debug:
+            debug_folder = Path(__file__).resolve().parent.parent / "debug"
+            self.test_incoming_wav = soundfile.SoundFile(
+                debug_folder / f"{self.debug_ts}_test_incoming.wav",
+                mode="w+",
+                format="WAV",
+                subtype="PCM_16",
+                samplerate=self.incoming_sample_rate,
+                channels=1,
+            )
+            self.get_states_root().test_input_segments_wav = soundfile.SoundFile(
+                debug_folder / f"{self.debug_ts}_test_input_segments.wav",
+                mode="w+",
+                format="WAV",
+                samplerate=MODEL_SAMPLE_RATE,
+                channels=1,
+            )
 
+    def get_states_root(self) -> AgentStates:
+        if isinstance(self.agent, TreeAgentPipeline):
+            # self.states is a dict
+            return self.states[self.agent.source_module]
+        else:
+            # self.states is a list
+            return self.states[0]
 
+    def reset_states(self):
+        if isinstance(self.agent, TreeAgentPipeline):
+            states_iter = self.states.values()
+        else:
+            states_iter = self.states
+        for state in states_iter:
+            state.reset()
 
-        load_model_fairseq2()
+    def debug_log(self, *args):
+        if self.debug:
+            logger.info(*args)
 
-        device = "cpu"
-        print("DEVICE", device)
-        model_name_or_card="seamlessM4T_medium"
-        vocoder_name_or_card="vocoder_36langs"
-        # dtype=torch.float16,
-        # For CPU Mode need to use 32, float16 causes errors downstream
-        dtype=dtype=torch.float32
+    def process_incoming_bytes(self, incoming_bytes, target_language, sample_rate):
+        # TODO: currently just taking sample rate here, refactor sample rate
+        # bytes is 16bit signed int
+        self.incoming_sample_rate = sample_rate
+        segment, sr = self._preprocess_wav(incoming_bytes)
 
-        model: UnitYModel = load_model_for_inference(
-            load_unity_model, model_name_or_card, device, dtype
+        segment = SpeechSegment(
+            content=segment, sample_rate=sr, tgt_lang=target_language
         )
+        # # segment is array([0, 0, 0, ..., 0, 0, 0], dtype=int16)
+        self.input_queue.put_nowait(segment)
 
+    def get_input_segment(self):
+        if self.input_queue.empty():
+            return None
+        chunk = self.input_queue.get_nowait()
+        self.input_queue.task_done()
+        return chunk
+
+    def _preprocess_wav(self, data: Any) -> Tuple[np.ndarray, int]:
+        segment, sample_rate = soundfile.read(
+            io.BytesIO(data),
+            dtype="float32",
+            always_2d=True,
+            frames=-1,
+            start=0,
+            format="RAW",
+            subtype="PCM_16",
+            samplerate=self.incoming_sample_rate,
+            channels=1,
+        )
+        if self.debug:
+            self.test_incoming_wav.seek(0, soundfile.SEEK_END)
+            self.test_incoming_wav.write(segment)
 
-        print(model, type(model))
-        parser = ArgumentParser()
-        source_segment_size = 320  # milliseconds
-        audio_frontend = AudioFrontEnd(
-            wav_file="/checkpoint/mduppes/samples/marta.wav",
-            segment_size=source_segment_size,
+        segment = segment.T
+        segment, new_sample_rate = convert_waveform(
+            segment,
+            sample_rate,
+            normalize_volume=False,
+            to_mono=True,
+            to_sample_rate=MODEL_SAMPLE_RATE,
         )
 
-        # mostly taken from S2S first agent: OnlineFeatureExtractorAgent defaults
-        SHIFT_SIZE = 10
-        WINDOW_SIZE = 25
-        SAMPLE_RATE = 16000
-        FEATURE_DIM = 80
-
-        # args and convert to namespace so it can be accesed via .
-        args = {
-            "shift_size": SHIFT_SIZE,
-            "window_size": WINDOW_SIZE,
-            "sample_rate": audio_frontend.sample_rate,
-            "feature_dim": 160, # from Wav2Vec2Frontend
-            "denormalize": False, # not sure..
-            "global_stats": None, # default file path containing cmvn stats..
-        }
-        print(args)
-        args = Namespace(**args)
-
-        pipeline = TestTimeWaitKUnityV1M4T(model, args)
-        system_states = pipeline.build_states()
-        print('system states:')
-        for state in system_states:
-            print(state, vars(state))
-
-        input_segment = np.empty(0, dtype=np.int16)
-        segments = []
-        while True:
-            speech_segment = audio_frontend.send_segment()
-            input_segment = np.concatenate((input_segment, np.array(speech_segment.content)))
-            # Translation happens here
-            output_segment = pipeline.pushpop(speech_segment, system_states)
-            print('pushpop result')
-            print(output_segment)
-            print('system states after pushpop:')
-            for state in system_states:
-                print(state, vars(state))
+        assert MODEL_SAMPLE_RATE == new_sample_rate
+        segment = segment.squeeze(axis=0)
+        return segment, new_sample_rate
+
+    def process_pipeline_impl(self, input_segment):
+        try:
+            with torch.no_grad():
+                output_segment = OutputSegments(
+                    self.agent.pushpop(input_segment, self.states)
+                )
+            if (
+                self.get_states_root().first_input_ts is not None
+                and self.first_input_ts is None
+            ):
+                # TODO: this is hacky
+                self.first_input_ts = self.get_states_root().first_input_ts
+
+            if not output_segment.is_empty:
+                self.output_queue.put_nowait(output_segment)
+
             if output_segment.finished:
-                segments.append(input_segment)
-                input_segment = np.empty(0, dtype=np.int16)
-                print("Resetting states")
-                for state in system_states:
-                    state.reset()
-            if speech_segment.finished:
-                break
-        # The VAD-segmented samples from the full input audio
-        for i, seg in enumerate(segments):
-            with soundfile.SoundFile(
-                Path("/checkpoint/mduppes/samples") / f"marta_{i}.wav",
-                mode="w+",
-                format="WAV",
-                samplerate=16000,
-                channels=1,
-            ) as f:
-                f.seek(0, soundfile.SEEK_END)
-                f.write(seg)
+                self.debug_log("OUTPUT SEGMENT IS FINISHED. Resetting states.")
+
+                self.reset_states()
+
+                if self.debug:
+                    # when we rebuild states, this value is reset to whatever
+                    # is in the system dir config, which defaults debug=False.
+                    self.get_states_root().debug = True
+        except Exception as e:
+            logger.error(f"Got exception while processing pipeline: {e}")
+            traceback.print_exc()
+        return input_segment
+
+    def process_pipeline_loop(self):
+        if self.close:
+            return  # closes the thread
+
+        self.debug_log("processing_pipeline")
+        while not self.close:
+            input_segment = self.get_input_segment()
+            if input_segment is None:
+                # if self.get_states_root().is_fresh_state:  # TODO: this is hacky
+                #     time.sleep(0.3)
+                # else:
+                time.sleep(0.03)
+                continue
+            self.process_pipeline_impl(input_segment)
+        self.debug_log("finished processing_pipeline")
+
+    def process_pipeline_once(self):
+        if self.close:
+            return
+
+        self.debug_log("processing pipeline once")
+        input_segment = self.get_input_segment()
+        if input_segment is None:
+            return
+        self.process_pipeline_impl(input_segment)
+        self.debug_log("finished processing_pipeline_once")
+
+    def get_output_segment(self):
+        if self.output_queue.empty():
+            return None
+
+        output_chunk = self.output_queue.get_nowait()
+        self.output_queue.task_done()
+        return output_chunk
+
+    def start(self):
+        self.debug_log("starting transcoder in a thread")
+        threading.Thread(target=self.process_pipeline_loop).start()
+
+    def first_translation_time(self):
+        return round((self.first_output_ts - self.first_input_ts) / 1000, 2)
+
+    def get_buffered_output(self) -> SpeechAndTextOutput:
+        now = time.time() * 1000
+        self.debug_log(f"get_buffered_output queue size: {self.output_queue.qsize()}")
+        while not self.output_queue.empty():
+            tmp_out = self.get_output_segment()
+            if tmp_out and tmp_out.compute_length(self.g2p) > 0:
+                if len(self.output_buffer) == 0:
+                    self.last_output_ts = now
+                self._populate_output_buffer(tmp_out)
+                self._increment_output_buffer_size(tmp_out)
+
+                if tmp_out.finished:
+                    self.debug_log("tmp_out.finished")
+                    res = self._gather_output_buffer_data(final=True)
+                    self.debug_log(f"gathered output data: {res}")
+                    self.output_buffer = []
+                    self.increment_output_buffer_size = 0
+                    self.last_output_ts = now
+                    self.first_output_ts = now
+                    return res
+            else:
+                self.debug_log("tmp_out.compute_length is not > 0")
+
+        if len(self.output_buffer) > 0 and (
+            now - self.last_output_ts >= self.output_buffer_idle_ms
+            or self.output_buffer_cur_size >= self.output_buffer_size_limit
+        ):
+            self.debug_log(
+                "[get_buffered_output] output_buffer is not empty. getting res to return."
+            )
+            self.last_output_ts = now
+            res = self._gather_output_buffer_data(final=False)
+            self.debug_log(f"gathered output data: {res}")
+            self.output_buffer = []
+            self.output_buffer_phoneme_count = 0
+            self.first_output_ts = now
+            return res
+        else:
+            self.debug_log("[get_buffered_output] output_buffer is empty...")
+            return None
+
+    def _gather_output_buffer_data(self, final):
+        output = SpeechAndTextOutput()
+        output.final = final
+        output = OutputSegments.join_output_buffer(self.output_buffer, output)
+        return output
+
+    def _increment_output_buffer_size(self, segment: OutputSegments):
+        self.output_buffer_cur_size += segment.compute_length(self.g2p)
+
+    def _populate_output_buffer(self, segment: OutputSegments):
+        self.output_buffer.append(segment.segments)
 
+    def _compute_phoneme_count(self, string: str) -> int:
+        return len([x for x in self.g2p(string) if x != " "])