Inference endpoint

LLM/language_model.py

@@ -145,3 +145,4 @@ def process(self, prompt):
 
  # don't forget last sentence
  yield (printable_text, language_code)
+ yield b"DONE"

TTS/melo_handler.py

@@ -56,6 +56,11 @@ def warmup(self):
  _ = self.model.tts_to_file("text", self.speaker_id, quiet=True)
 
  def process(self, llm_sentence):
+ if llm_sentence == b"DONE":
+ self.should_listen.set()
+ yield b"DONE"
+ return
+
  language_code = None
 
  if isinstance(llm_sentence, tuple):
@@ -94,10 +99,11 @@ def process(self, llm_sentence):
  )
  except (AssertionError, RuntimeError) as e:
  logger.error(f"Error in MeloTTSHandler: {e}")
- audio_chunk = np.array([])
- if len(audio_chunk) == 0:
- self.should_listen.set()
- return
+ audio_chunk = np.zeros([self.blocksize])
+ except Exception as e:
+ logger.error(f"Unknown error in MeloTTSHandler: {e}")
+ audio_chunk = np.zeros([self.blocksize])
+
  audio_chunk = librosa.resample(audio_chunk, orig_sr=44100, target_sr=16000)
  audio_chunk = (audio_chunk * 32768).astype(np.int16)
  for i in range(0, len(audio_chunk), self.blocksize):

TTS/parler_handler.py

@@ -189,3 +189,4 @@ def process(self, llm_sentence):
  )
 
  self.should_listen.set()
+ yield b"END"

arguments_classes/module_arguments.py

@@ -11,7 +11,7 @@ class ModuleArguments:
  mode: Optional[str] = field(
  default="socket",
  metadata={
- "help": "The mode to run the pipeline in. Either 'local' or 'socket'. Default is 'socket'."
+ "help": "The mode to run the pipeline in. Either 'local', 'socket', or 'none'. Default is 'socket'."
  },
  )
  local_mac_optimal_settings: bool = field(

audio_streaming_client.py

@@ -0,0 +1,185 @@
+import threading
+from queue import Queue
+import sounddevice as sd
+import numpy as np
+import time
+from dataclasses import dataclass, field
+import websocket
+import ssl
+
+
+@dataclass
+class AudioStreamingClientArguments:
+ sample_rate: int = field(
+ default=16000, metadata={"help": "Audio sample rate in Hz. Default is 16000."}
+ )
+ chunk_size: int = field(
+ default=512,
+ metadata={"help": "The size of audio chunks in samples. Default is 512."},
+ )
+ api_url: str = field(
+ default="https://yxfmjcvuzgi123sw.us-east-1.aws.endpoints.huggingface.cloud",
+ metadata={"help": "The URL of the API endpoint."},
+ )
+ auth_token: str = field(
+ default="your_auth_token",
+ metadata={"help": "Authentication token for the API."},
+ )
+
+
+class AudioStreamingClient:
+ def __init__(self, args: AudioStreamingClientArguments):
+ self.args = args
+ self.stop_event = threading.Event()
+ self.send_queue = Queue()
+ self.recv_queue = Queue()
+ self.session_id = None
+ self.headers = {
+ "Accept": "application/json",
+ "Authorization": f"Bearer {self.args.auth_token}",
+ "Content-Type": "application/json",
+ }
+ self.session_state = (
+ "idle" # Possible states: idle, sending, processing, waiting
+ )
+ self.ws_ready = threading.Event()
+
+ def start(self):
+ print("Starting audio streaming...")
+
+ ws_url = self.args.api_url.replace("http", "ws") + "/ws"
+
+ self.ws = websocket.WebSocketApp(
+ ws_url,
+ header=[f"{key}: {value}" for key, value in self.headers.items()],
+ on_open=self.on_open,
+ on_message=self.on_message,
+ on_error=self.on_error,
+ on_close=self.on_close,
+ )
+
+ self.ws_thread = threading.Thread(
+ target=self.ws.run_forever, kwargs={"sslopt": {"cert_reqs": ssl.CERT_NONE}}
+ )
+ self.ws_thread.start()
+
+ # Wait for the WebSocket to be ready
+ self.ws_ready.wait()
+ self.start_audio_streaming()
+
+ def start_audio_streaming(self):
+ self.send_thread = threading.Thread(target=self.send_audio)
+ self.play_thread = threading.Thread(target=self.play_audio)
+
+ with sd.InputStream(
+ samplerate=self.args.sample_rate,
+ channels=1,
+ dtype="int16",
+ callback=self.audio_input_callback,
+ blocksize=self.args.chunk_size,
+ ):
+ self.send_thread.start()
+ self.play_thread.start()
+ input("Press Enter to stop streaming... \n")
+ self.on_shutdown()
+
+ def on_open(self, ws):
+ print("WebSocket connection opened.")
+ self.ws_ready.set() # Signal that the WebSocket is ready
+
+ def on_message(self, ws, message):
+ # message is bytes
+ if message == b"DONE":
+ print("listen")
+ self.session_state = "listen"
+ else:
+ if self.session_state != "processing":
+ print("processing")
+ self.session_state = "processing"
+ audio_np = np.frombuffer(message, dtype=np.int16)
+ self.recv_queue.put(audio_np)
+
+ def on_error(self, ws, error):
+ print(f"WebSocket error: {error}")
+
+ def on_close(self, ws, close_status_code, close_msg):
+ print("WebSocket connection closed.")
+
+ def on_shutdown(self):
+ self.stop_event.set()
+ self.send_thread.join()
+ self.play_thread.join()
+ self.ws.close()
+ self.ws_thread.join()
+ print("Service shutdown.")
+
+ def send_audio(self):
+ while not self.stop_event.is_set():
+ if not self.send_queue.empty():
+ chunk = self.send_queue.get()
+ if self.session_state != "processing":
+ self.ws.send(chunk.tobytes(), opcode=websocket.ABNF.OPCODE_BINARY)
+ else:
+ self.ws.send([], opcode=websocket.ABNF.OPCODE_BINARY) # handshake
+ time.sleep(0.01)
+
+ def audio_input_callback(self, indata, frames, time, status):
+ self.send_queue.put(indata.copy())
+
+ def audio_out_callback(self, outdata, frames, time, status):
+ if not self.recv_queue.empty():
+ chunk = self.recv_queue.get()
+
+ # Ensure chunk is int16 and clip to valid range
+ chunk_int16 = np.clip(chunk, -32768, 32767).astype(np.int16)
+
+ if len(chunk_int16) < len(outdata):
+ outdata[: len(chunk_int16), 0] = chunk_int16
+ outdata[len(chunk_int16) :] = 0
+ else:
+ outdata[:, 0] = chunk_int16[: len(outdata)]
+ else:
+ outdata[:] = 0
+
+ def play_audio(self):
+ with sd.OutputStream(
+ samplerate=self.args.sample_rate,
+ channels=1,
+ dtype="int16",
+ callback=self.audio_out_callback,
+ blocksize=self.args.chunk_size,
+ ):
+ while not self.stop_event.is_set():
+ time.sleep(0.1)
+
+
+if __name__ == "__main__":
+ import argparse
+
+ parser = argparse.ArgumentParser(description="Audio Streaming Client")
+ parser.add_argument(
+ "--sample_rate",
+ type=int,
+ default=16000,
+ help="Audio sample rate in Hz. Default is 16000.",
+ )
+ parser.add_argument(
+ "--chunk_size",
+ type=int,
+ default=1024,
+ help="The size of audio chunks in samples. Default is 1024.",
+ )
+ parser.add_argument(
+ "--api_url", type=str, required=True, help="The URL of the API endpoint."
+ )
+ parser.add_argument(
+ "--auth_token",
+ type=str,
+ required=True,
+ help="Authentication token for the API.",
+ )
+
+ args = parser.parse_args()
+ client_args = AudioStreamingClientArguments(**vars(args))
+ client = AudioStreamingClient(client_args)
+ client.start()

listen_and_play.py

@@ -125,4 +125,4 @@ def receive_full_chunk(conn, chunk_size):
 if __name__ == "__main__":
  parser = HfArgumentParser((ListenAndPlayArguments,))
  (listen_and_play_kwargs,) = parser.parse_args_into_dataclasses()
- listen_and_play(**vars(listen_and_play_kwargs))
+ listen_and_play(**vars(listen_and_play_kwargs))

requirements.txt

@@ -8,4 +8,5 @@ funasr>=1.1.6
 faster-whisper>=1.0.3
 modelscope>=1.17.1
 deepfilternet>=0.5.6
-openai>=1.40.1
+openai>=1.40.1
+websocket-client>=1.8.0

requirements_mac.txt

@@ -10,4 +10,5 @@ funasr>=1.1.6
 faster-whisper>=1.0.3
 modelscope>=1.17.1
 deepfilternet>=0.5.6
-openai>=1.40.1
+openai>=1.40.1
+websocket-client>=1.8.0

s2s_handler.py

@@ -0,0 +1,105 @@
+from typing import Dict, Any, List, Generator
+import torch
+import os
+import logging
+from s2s_pipeline import main, prepare_all_args, get_default_arguments, setup_logger, initialize_queues_and_events, build_pipeline
+import numpy as np
+from queue import Queue, Empty
+import threading
+import base64
+import uuid
+import torch
+
+class EndpointHandler:
+ def __init__(self, path=""):
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ lm_model_name = os.getenv('LM_MODEL_NAME', 'meta-llama/Meta-Llama-3.1-8B-Instruct')
+ chat_size = int(os.getenv('CHAT_SIZE', 10))
+
+ (
+ self.module_kwargs,
+ self.socket_receiver_kwargs,
+ self.socket_sender_kwargs,
+ self.vad_handler_kwargs,
+ self.whisper_stt_handler_kwargs,
+ self.paraformer_stt_handler_kwargs,
+ self.faster_whisper_stt_handler_kwargs,
+ self.language_model_handler_kwargs,
+ self.open_api_language_model_handler_kwargs,
+ self.mlx_language_model_handler_kwargs,
+ self.parler_tts_handler_kwargs,
+ self.melo_tts_handler_kwargs,
+ self.chat_tts_handler_kwargs,
+ self.facebook_mm_stts_handler_kwargs,
+ ) = get_default_arguments(mode='none', log_level='DEBUG', lm_model_name=lm_model_name, 
+ tts="melo", device=device, chat_size=chat_size)
+ setup_logger(self.module_kwargs.log_level)
+
+ prepare_all_args(
+ self.module_kwargs,
+ self.whisper_stt_handler_kwargs,
+ self.paraformer_stt_handler_kwargs,
+ self.faster_whisper_stt_handler_kwargs,
+ self.language_model_handler_kwargs,
+ self.open_api_language_model_handler_kwargs,
+ self.mlx_language_model_handler_kwargs,
+ self.parler_tts_handler_kwargs,
+ self.melo_tts_handler_kwargs,
+ self.chat_tts_handler_kwargs,
+ self.facebook_mm_stts_handler_kwargs,
+ )
+
+ self.queues_and_events = initialize_queues_and_events()
+
+ self.pipeline_manager = build_pipeline(
+ self.module_kwargs,
+ self.socket_receiver_kwargs,
+ self.socket_sender_kwargs,
+ self.vad_handler_kwargs,
+ self.whisper_stt_handler_kwargs,
+ self.paraformer_stt_handler_kwargs,
+ self.faster_whisper_stt_handler_kwargs,
+ self.language_model_handler_kwargs,
+ self.open_api_language_model_handler_kwargs,
+ self.mlx_language_model_handler_kwargs,
+ self.parler_tts_handler_kwargs,
+ self.melo_tts_handler_kwargs,
+ self.chat_tts_handler_kwargs,
+ self.facebook_mm_stts_handler_kwargs,
+ self.queues_and_events,
+ )
+
+ self.vad_chunk_size = 512 # Set the chunk size required by the VAD model
+ self.sample_rate = 16000 # Set the expected sample rate
+
+ def process_streaming_data(self, data: bytes) -> bytes:
+ audio_array = np.frombuffer(data, dtype=np.int16)
+
+ # Process the audio data in chunks
+ chunks = [audio_array[i:i+self.vad_chunk_size] for i in range(0, len(audio_array), self.vad_chunk_size)]
+
+ for chunk in chunks:
+ if len(chunk) == self.vad_chunk_size:
+ self.queues_and_events['recv_audio_chunks_queue'].put(chunk.tobytes())
+ elif len(chunk) < self.vad_chunk_size:
+ # Pad the last chunk if it's smaller than the required size
+ padded_chunk = np.pad(chunk, (0, self.vad_chunk_size - len(chunk)), 'constant')
+ self.queues_and_events['recv_audio_chunks_queue'].put(padded_chunk.tobytes())
+
+ # Collect the output, if any
+ try:
+ output = self.queues_and_events['send_audio_chunks_queue'].get_nowait() # improvement idea, group all available output chunks
+ if isinstance(output, np.ndarray):
+ return output.tobytes()
+ else:
+ return output
+ except Empty:
+ return None
+
+ def cleanup(self):
+ # Stop the pipeline
+ self.pipeline_manager.stop()
+
+ # Stop the output collector thread
+ self.queues_and_events['send_audio_chunks_queue'].put(b"END")
+ self.output_collector_thread.join()