Upload folder using huggingface_hub

handler.py

@@ -2,9 +2,9 @@ from typing import Dict, Any, List, Generator
 import torch
 import os
 import logging
-from s2s_pipeline import main, rename_args, get_default_arguments, setup_logger, initialize_queues_and_events, build_pipeline
+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
+from queue import Queue, Empty
 import threading
 
 class EndpointHandler:
@@ -21,16 +21,19 @@ class EndpointHandler:
             self.parler_tts_handler_kwargs,
             self.melo_tts_handler_kwargs,
             self.chat_tts_handler_kwargs,
-        ) = get_default_arguments()
+        ) = get_default_arguments(device='cpu', mode='none', tts='melo', stt='whisper-mlx')
         setup_logger(self.module_kwargs.log_level)
 
-        rename_args(self.whisper_stt_handler_kwargs, "stt")
-        rename_args(self.paraformer_stt_handler_kwargs, "paraformer_stt")
-        rename_args(self.language_model_handler_kwargs, "lm")
-        rename_args(self.mlx_language_model_handler_kwargs, "mlx_lm")
-        rename_args(self.parler_tts_handler_kwargs, "tts")
-        rename_args(self.melo_tts_handler_kwargs, "melo")
-        rename_args(self.chat_tts_handler_kwargs, "chat_tts")
+        prepare_all_args(
+            self.module_kwargs,
+            self.whisper_stt_handler_kwargs,
+            self.paraformer_stt_handler_kwargs,
+            self.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.queues_and_events = initialize_queues_and_events()
 
@@ -54,17 +57,21 @@ class EndpointHandler:
         # Add a new queue for collecting the final output
         self.final_output_queue = Queue()
 
-        # Start a thread to collect the final output
-        self.output_collector_thread = threading.Thread(target=self._collect_output)
-        self.output_collector_thread.start()
-
     def _collect_output(self):
         while True:
-            output = self.queues_and_events['send_audio_chunks_queue'].get()
-            if output == b"END":
-                self.final_output_queue.put(b"END")
+            try:
+                output = self.queues_and_events['send_audio_chunks_queue'].get(timeout=5)  # 2-second timeout
+                if isinstance(output, (str, bytes)) and output in (b"END", "END"):
+                    self.final_output_queue.put("END")
+                    break
+                elif isinstance(output, np.ndarray):
+                    self.final_output_queue.put(output.tobytes())
+                else:
+                    self.final_output_queue.put(output)
+            except Empty:
+                # If no output for 2 seconds, assume processing is complete
+                self.final_output_queue.put("END")
                 break
-            self.final_output_queue.put(output)
 
     def __call__(self, data: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
         """
@@ -74,6 +81,10 @@ class EndpointHandler:
         Returns:
             Generator[Dict[str, Any], None, None]: A generator yielding output chunks from the model or pipeline.
         """
+        # Start a thread to collect the final output
+        self.output_collector_thread = threading.Thread(target=self._collect_output)
+        self.output_collector_thread.start()
+
         input_type = data.get("input_type", "text")
         input_data = data.get("input", "")
 
@@ -89,12 +100,18 @@ class EndpointHandler:
         else:
             raise ValueError(f"Unsupported input type: {input_type}")
 
-        # Stream the output chunks
+        # Collect all output chunks
+        output_chunks = []
         while True:
             chunk = self.final_output_queue.get()
-            if chunk == b"END":
+            if chunk == "END":
                 break
-            yield {"output": chunk}
+            output_chunks.append(chunk)
+
+        # Combine all audio chunks into a single byte string
+        combined_audio = b''.join(output_chunks)
+
+        return {"output": combined_audio}
 
     def cleanup(self):
         # Stop the pipeline

s2s_pipeline.py

@@ -65,8 +65,8 @@ def rename_args(args, prefix):
 
     args.__dict__["gen_kwargs"] = gen_kwargs
 
-def get_default_arguments():
-    return (
+def get_default_arguments(**kwargs):
+    default_args = [
         ModuleArguments(),
         SocketReceiverArguments(),
         SocketSenderArguments(),
@@ -78,7 +78,14 @@ def get_default_arguments():
         ParlerTTSHandlerArguments(),
         MeloTTSHandlerArguments(),
         ChatTTSHandlerArguments(),
-    )
+    ]
+    # Update arguments with provided kwargs
+    for arg_obj in default_args:
+        for key, value in kwargs.items():
+            if hasattr(arg_obj, key):
+                setattr(arg_obj, key, value)
+
+    return tuple(default_args)
 
 def parse_arguments():
     parser = HfArgumentParser(
@@ -241,7 +248,7 @@ def build_pipeline(
         )
         comms_handlers = [local_audio_streamer]
         should_listen.set()
-    else:
+    elif module_kwargs.mode == "socket":
         from connections.socket_receiver import SocketReceiver
         from connections.socket_sender import SocketSender
 
@@ -261,6 +268,9 @@ def build_pipeline(
                 port=socket_sender_kwargs.send_port,
             ),
         ]
+    else:
+        comms_handlers = []
+        should_listen.set()
 
     vad = VADHandler(
         stop_event,

test.py

@@ -0,0 +1,7 @@
+from handler import EndpointHandler
+
+endpoint = EndpointHandler('')
+
+for x in endpoint({'text': 'how are you?'}):
+    print('passed')
+    print(x)