add support for CPU_THREADS config (#28)

2024-01-03 14:54:04 +01:00
parent 87fa20c8cb
commit f9e33be9de
2 changed files with 5 additions and 2 deletions
--- a/transcription-api/backends/fasterwhisper.py
+++ b/transcription-api/backends/fasterwhisper.py
@@ -26,8 +26,10 @@ class FasterWhisperBackend(Backend):
            raise RuntimeError(f"model not found in {local_model_path}")
        
    def load(self) -> None:
+        # Get CPU threads env variable or default to 4
+        cpu_threads = int(os.environ.get("CPU_THREADS", 4))
        self.model = WhisperModel(
-            self.model_path(), device=self.device, compute_type=self.quantization
+            self.model_path(), device=self.device, compute_type=self.quantization, cpu_threads=cpu_threads
        )

    def get_model(self) -> None: