Merge pull request #378 from makaveli10/auto_convert_faster_whisper

Auto convert hf custom whisper to ct2(faster-whisper)

README.md

@@ -47,10 +47,11 @@ The server supports 3 backends `faster_whisper`, `tensorrt` and `openvino`. If r
 python3 run_server.py --port 9090 \
                       --backend faster_whisper
   
-# running with custom model
+# running with custom model and cache_dir to save auto-converted ctranslate2 models
 python3 run_server.py --port 9090 \
                       --backend faster_whisper \
                       -fw "/path/to/custom/faster/whisper/model"
+                      -c ~/.cache/whisper-live/
 ```
 
 - TensorRT backend. Currently, we recommend to only use the docker setup for TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) which works as expected. Make sure to build your TensorRT Engines before running the server with TensorRT backend.

requirements/server.txt

@@ -11,6 +11,7 @@ evaluate
 numpy<2
 openai-whisper==20240930
 tokenizers==0.20.3
+transformers[torch]
 
 # openvino
 librosa

run_server.py

@@ -31,6 +31,10 @@ if __name__ == "__main__":
     parser.add_argument('--no_single_model', '-nsm',
                         action='store_true',
                         help='Set this if every connection should instantiate its own model. Only relevant for custom model, passed using -trt or -fw.')
+    parser.add_argument('--cache_path', '-c',
+                        type=str,
+                        default="~/.cache/whisper-live/",
+                        help='Path to cache the converted ctranslate2 models.')
     args = parser.parse_args()
 
     if args.backend == "tensorrt":
@@ -51,4 +55,5 @@ if __name__ == "__main__":
         trt_multilingual=args.trt_multilingual,
         trt_py_session=args.trt_py_session,
         single_model=not args.no_single_model,
+        cache_path=args.cache_path
     )

whisper_live/backend/faster_whisper_backend.py

@@ -1,8 +1,11 @@
+import os
 import json
 import logging
 import threading
 import time
 import torch
+import ctranslate2
+from huggingface_hub import snapshot_download
 
 from whisper_live.transcriber.transcriber_faster_whisper import WhisperModel
 from whisper_live.backend.base import ServeClientBase
@@ -28,6 +31,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         no_speech_thresh=0.45,
         clip_audio=False,
         same_output_threshold=10,
+        cache_path="~/.cache/whisper-live/"
     ):
         """
         Initialize a ServeClient instance.
@@ -58,6 +62,7 @@ class ServeClientFasterWhisper(ServeClientBase):
             clip_audio,
             same_output_threshold,
         )
+        self.cache_path = cache_path
         self.model_sizes = [
             "tiny", "tiny.en", "base", "base.en", "small", "small.en",
             "medium", "medium.en", "large-v2", "large-v3", "distil-small.en",
@@ -118,38 +123,51 @@ class ServeClientFasterWhisper(ServeClientBase):
 
     def create_model(self, device):
         """
-        Instantiates a new model, sets it as the transcriber.
+        Instantiates a new model, sets it as the transcriber. If model is a huggingface model_id
+        then it is automatically converted to ctranslate2(faster_whisper) format.
         """
+        model_ref = self.model_size_or_path
+
+        if model_ref in self.model_sizes:
+            model_to_load = model_ref
+        else:
+            logging.info(f"Model not in model_sizes")
+            if os.path.isdir(model_ref) and ctranslate2.contains_model(model_ref):
+                model_to_load = model_ref
+            else:
+                local_snapshot = snapshot_download(
+                    repo_id = model_ref,
+                    repo_type = "model",
+                )
+                if ctranslate2.contains_model(local_snapshot):
+                    model_to_load = local_snapshot
+                else:
+                    cache_root = os.path.expanduser(os.path.join(self.cache_path, "whisper-ct2-models/"))
+                    os.makedirs(cache_root, exist_ok=True)
+                    safe_name = model_ref.replace("/", "--")
+                    ct2_dir = os.path.join(cache_root, safe_name)
+
+                    if not ctranslate2.contains_model(ct2_dir):
+                        logging.info(f"Converting '{model_ref}' to CTranslate2 @ {ct2_dir}")
+                        ct2_converter = ctranslate2.converters.TransformersConverter(
+                            local_snapshot, 
+                            copy_files=["tokenizer.json", "preprocessor_config.json"]
+                        )
+                        ct2_converter.convert(
+                            output_dir=ct2_dir,
+                            quantization=self.compute_type,
+                            force=False,  # skip if already up-to-date
+                        )
+                    model_to_load = ct2_dir
+
+        logging.info(f"Loading model: {model_to_load}")
         self.transcriber = WhisperModel(
-            self.model_size_or_path,
+            model_to_load,
             device=device,
             compute_type=self.compute_type,
             local_files_only=False,
         )
 
-    def check_valid_model(self, model_size):
-        """
-        Check if it's a valid whisper model size.
-
-        Args:
-            model_size (str): The name of the model size to check.
-
-        Returns:
-            str: The model size if valid, None otherwise.
-        """
-        if model_size not in self.model_sizes:
-            self.websocket.send(
-                json.dumps(
-                    {
-                        "uid": self.client_uid,
-                        "status": "ERROR",
-                        "message": f"Invalid model size {model_size}. Available choices: {self.model_sizes}"
-                    }
-                )
-            )
-            return None
-        return model_size
-
     def set_language(self, info):
         """
         Updates the language attribute based on the detected language information.

whisper_live/server.py

@@ -233,6 +233,7 @@ class TranscriptionServer:
                     no_speech_thresh=options.get("no_speech_thresh", 0.45),
                     clip_audio=options.get("clip_audio", False),
                     same_output_threshold=options.get("same_output_threshold", 10),
+                    cache_path=self.cache_path,
                 )
 
                 logging.info("Running faster_whisper backend.")
@@ -369,7 +370,8 @@ class TranscriptionServer:
             whisper_tensorrt_path=None,
             trt_multilingual=False,
             trt_py_session=False,
-            single_model=False):
+            single_model=False,
+            cache_path="~/.cache/whisper-live/"):
         """
         Run the transcription server.
 
@@ -377,6 +379,7 @@ class TranscriptionServer:
             host (str): The host address to bind the server.
             port (int): The port number to bind the server.
         """
+        self.cache_path = cache_path
         if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path):
             raise ValueError(f"Custom faster_whisper model '{faster_whisper_custom_model_path}' is not a valid path.")
         if whisper_tensorrt_path is not None and not os.path.exists(whisper_tensorrt_path):