Merge remote-tracking branch 'upstream/main' into tensorrt_backend

whisper_live/__version__.py

@@ -1 +1 @@
-__version__="0.0.10"
+__version__="0.0.11"

whisper_live/client.py

@@ -50,7 +50,13 @@ class Client:
     INSTANCES = {}
 
     def __init__(
-        self, host=None, port=None, is_multilingual=False, lang=None, translate=False, model_size="small"
+        self,
+        host=None,
+        port=None,
+        is_multilingual=False,
+        lang=None,
+        translate=False,
+        model="small",
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -81,8 +87,9 @@ class Client:
         self.disconnect_if_no_response_for = 15
         self.multilingual = is_multilingual
         self.language = lang
-        self.model_size = model_size
+        self.model = model
         self.server_error = False
+
         if translate:
             self.task = "translate"
 
@@ -220,7 +227,7 @@ class Client:
                     "multilingual": self.multilingual,
                     "language": self.language,
                     "task": self.task,
-                    "model_size": self.model_size,
+                    "model": self.model,
                 }
             )
         )
@@ -505,8 +512,15 @@ class TranscriptionClient:
         transcription_client()
         ```
     """
-    def __init__(self, host, port, is_multilingual=False, lang=None, translate=False, model_size="small"):
-        self.client = Client(host, port, is_multilingual, lang, translate, model_size)
+    def __init__(self,
+        host,
+        port,
+        is_multilingual=False,
+        lang=None,
+        translate=False,
+        model="small",
+    ):
+        self.client = Client(host, port, is_multilingual, lang, translate, model)
 
     def __call__(self, audio=None, hls_url=None):
         """

whisper_live/server.py

@@ -120,6 +120,11 @@ class TranscriptionServer:
             websocket.close()
             del websocket
             return
+        
+        # validate custom model
+        if custom_model_path is not None and os.path.exists(custom_model_path):
+            logging.info(f"Using custom model {custom_model_path}")
+            options["model"] = custom_model_path
 
         if self.backend == "tensorrt":
             try:
@@ -561,7 +566,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         client_uid=None,
         model="small",
         initial_prompt=None,
-        vad_parameters=None
+        vad_parameters=None,
         ):
         """
         Initialize a ServeClient instance.
@@ -583,6 +588,7 @@ class ServeClientFasterWhisper(ServeClientBase):
             "tiny", "tiny.en", "base", "base.en", "small", "small.en",
             "medium", "medium.en", "large-v2", "large-v3",
         ]
+
         self.multilingual = multilingual
         if not os.path.exists(model):
             self.model_size_or_path = self.get_model_size(model)
@@ -633,12 +639,13 @@ class ServeClientFasterWhisper(ServeClientBase):
             )
             return None
         
-        if model_size in ["large-v2", "large-v3"]:
+        if model_size.endswith("en") and self.multilingual:
+            logging.info(f"Setting multilingual to false with {model_size} which is english only model.")
+            self.multilingual = False
+            
+        if not model_size.endswith("en") and not self.multilingual:
+            logging.info(f"Setting multilingual to true with multilingual model {model_size}.")
             self.multilingual = True
-            return model_size
-
-        if not self.multilingual:
-            model_size = model_size + ".en"
 
         return model_size