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_WHISHPER_PROFILE: cpu
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}")
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: