diff --git a/transcribe.py b/transcribe.py index b750f8c..186f0e0 100644 --- a/transcribe.py +++ b/transcribe.py @@ -92,7 +92,7 @@ def prepare_joint_annotations(result, target_length=200): return annotations -def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=False): +def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=False, model="small"): language = None if "language" not in item.data: language = None @@ -108,6 +108,7 @@ def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=Fa run_py = os.path.join(os.path.dirname(os.path.abspath(__file__)), "run_whisper.py") cmd = ["/opt/whisper-timestamped/bin/python", run_py] + cmd += ["--model", model] if language: cmd += ["--language", language] if translate and language in translate: @@ -148,7 +149,7 @@ def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=Fa return True -def extract_subtitles_cmd(item, user, layer, translate, gpu=False): +def extract_subtitles_cmd(item, user, layer, translate, gpu=False, model="small"): if "language" not in item.data: language = None else: @@ -159,7 +160,7 @@ def extract_subtitles_cmd(item, user, layer, translate, gpu=False): src = item.streams()[0].media.path tmp = tempfile.mkdtemp() - cmd = ["/opt/whisper-timestamped/bin/whisper_timestamped", "--model", "turbo"] + cmd = ["/opt/whisper-timestamped/bin/whisper_timestamped", "--model", model] if language: cmd += ["--language", language] if translate and language in translate: