From 1c7d7e138b768c443a9dc37de305c761a9b2026b Mon Sep 17 00:00:00 2001
From: harry <harry@wangxutech.com>
Date: Thu, 11 Apr 2024 18:27:14 +0800
Subject: [PATCH] compatible with local whisper model

---
 app/services/subtitle.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/app/services/subtitle.py b/app/services/subtitle.py
index 5ea8028..d19ac59 100644
--- a/app/services/subtitle.py
+++ b/app/services/subtitle.py
@@ -1,4 +1,5 @@
 import json
+import os.path
 import re
 
 from faster_whisper import WhisperModel
@@ -17,8 +18,13 @@ model = None
 def create(audio_file, subtitle_file: str = ""):
     global model
     if not model:
-        logger.info(f"loading model: {model_size}, device: {device}, compute_type: {compute_type}")
-        model = WhisperModel(model_size_or_path=model_size,
+        model_path = f"{utils.root_dir()}/models/whisper-{model_size}"
+        model_bin_file = f"{model_path}/model.bin"
+        if not os.path.isdir(model_path) or not os.path.isfile(model_bin_file):
+            model_path = model_size
+
+        logger.info(f"loading model: {model_path}, device: {device}, compute_type: {compute_type}")
+        model = WhisperModel(model_size_or_path=model_path,
                              device=device,
                              compute_type=compute_type)