dasemu/Transcriptarr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update subgen.py

Browse Source
This commit is contained in:
McCloudS
2024-02-08 08:33:42 -07:00
committed by GitHub
parent 9395a65a49
commit 269a54cbb6
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
subgen/subgen.py
View File

@@ -49,6 +49,7 @@ path_mapping_to = os.getenv('PATH_MAPPING_TO', '/Volumes/TV')
model_location = os.getenv('MODEL_PATH', '.') model_location = os.getenv('MODEL_PATH', '.')
transcribe_folders = os.getenv('TRANSCRIBE_FOLDERS', '') transcribe_folders = os.getenv('TRANSCRIBE_FOLDERS', '')
transcribe_or_translate = os.getenv('TRANSCRIBE_OR_TRANSLATE', 'translate') transcribe_or_translate = os.getenv('TRANSCRIBE_OR_TRANSLATE', 'translate')
force_detected_language_to = os.getenv('FORCE_DETECTED_LANGUAGE_TO', '')
compute_type = os.getenv('COMPUTE_TYPE', 'auto') compute_type = os.getenv('COMPUTE_TYPE', 'auto')
if transcribe_device == "gpu": if transcribe_device == "gpu":
transcribe_device = "cuda" transcribe_device = "cuda"
@@ -285,8 +286,11 @@ def gen_subtitles(file_path: str, transcribe_or_translate_str: str, front=True)
print(f"Transcribing file: {os.path.basename(file_path)}") print(f"Transcribing file: {os.path.basename(file_path)}")
start_time = time.time() start_time = time.time()
start_model() start_model()
result = model.transcribe_stable(file_path, task=transcribe_or_translate_str) if(force_detected_language_to):
result = model.transcribe_stable(file_path, language=force_detected_language_to, task=transcribe_or_translate_str)
else:
result = model.transcribe_stable(file_path, task=transcribe_or_translate_str)
result.to_srt_vtt(get_file_name_without_extension(file_path) + subextension, word_level=word_level_highlight) result.to_srt_vtt(get_file_name_without_extension(file_path) + subextension, word_level=word_level_highlight)
elapsed_time = time.time() - start_time elapsed_time = time.time() - start_time
minutes, seconds = divmod(int(elapsed_time), 60) minutes, seconds = divmod(int(elapsed_time), 60)