From f9e33be9de15956972a804015977250b318de1ca Mon Sep 17 00:00:00 2001 From: pluja Date: Wed, 3 Jan 2024 14:54:04 +0100 Subject: [PATCH] add support for CPU_THREADS config (#28) --- docker-compose.yml | 3 ++- transcription-api/backends/fasterwhisper.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 26f515c..f346541 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -66,4 +66,5 @@ services: PUBLIC_API_HOST: ${WHISHPER_HOST:-} PUBLIC_WHISHPER_PROFILE: cpu WHISPER_MODELS_DIR: /app/models - UPLOAD_DIR: /app/uploads \ No newline at end of file + UPLOAD_DIR: /app/uploads + CPU_THREADS: 4 \ No newline at end of file diff --git a/transcription-api/backends/fasterwhisper.py b/transcription-api/backends/fasterwhisper.py index 0af4621..f6db2c4 100644 --- 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: