diff --git a/management/commands/transcribe.py b/management/commands/transcribe.py index 8bf1e95..4350251 100644 --- a/management/commands/transcribe.py +++ b/management/commands/transcribe.py @@ -21,6 +21,8 @@ class Command(BaseCommand): default=None, help='import into layer (default subtitle layer)') parser.add_argument('--translate', type=str, dest='translate', default='', help='list of languages to translate: (i.e. hi:en,de:en') + parser.add_argument('--model', type=str, dest='model', + default='small', help='whisper model (small, large-v3-turbo, default: small)') parser.add_argument('--gpu', action='store_true', dest='gpu', default=False, help='user GPU (default: disabled)') parser.add_argument('--join-sentences', action='store_true', dest='join_sentences', diff --git a/transcribe.py b/transcribe.py index b750f8c..fcce57f 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: @@ -240,6 +241,7 @@ def main(**kwargs): translate, kwargs.get("gpu"), join_sentences=kwargs.get("join_sentences"), + model=kwargs.get("model", "small") ): done.items.add(item) queue.items.remove(item)