add support for CPU_THREADS config (#28)

This commit is contained in:
pluja
2024-01-03 14:54:04 +01:00
parent 87fa20c8cb
commit f9e33be9de
2 changed files with 5 additions and 2 deletions

View File

@@ -66,4 +66,5 @@ services:
PUBLIC_API_HOST: ${WHISHPER_HOST:-} PUBLIC_API_HOST: ${WHISHPER_HOST:-}
PUBLIC_WHISHPER_PROFILE: cpu PUBLIC_WHISHPER_PROFILE: cpu
WHISPER_MODELS_DIR: /app/models WHISPER_MODELS_DIR: /app/models
UPLOAD_DIR: /app/uploads UPLOAD_DIR: /app/uploads
CPU_THREADS: 4

View File

@@ -26,8 +26,10 @@ class FasterWhisperBackend(Backend):
raise RuntimeError(f"model not found in {local_model_path}") raise RuntimeError(f"model not found in {local_model_path}")
def load(self) -> None: 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 = 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: def get_model(self) -> None: