Skip to content

Commit

Permalink
Update subgen.py
Browse files Browse the repository at this point in the history
Cleaned up garbage collection so vram is released when the transcode queue is empty.
  • Loading branch information
McCloudS authored Oct 27, 2023
1 parent 036dd9a commit 42bf55f
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion subgen/subgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import time
import queue
import logging
import gc
from array import array

# List of packages to install
Expand Down Expand Up @@ -66,7 +67,7 @@ def convert_to_bool(in_bool):
jellyfin_userid = ""

app = Flask(__name__)
model = stable_whisper.load_faster_whisper(whisper_model, download_root=model_location, device=transcribe_device, cpu_threads=whisper_threads, num_workers=concurrent_transcriptions)
model = None
files_to_transcribe = []
subextension = '.subgen.' + whisper_model + '.' + namesublang + '.srt'
print("Transcriptions are limited to running " + str(concurrent_transcriptions) + " at a time")
Expand Down Expand Up @@ -179,6 +180,7 @@ def gen_subtitles(video_file_path: str) -> None:
try:
print(f"Transcribing file: {video_file_path}")
start_time = time.time()
model = stable_whisper.load_faster_whisper(whisper_model, download_root=model_location, device=transcribe_device, cpu_threads=whisper_threads, num_workers=concurrent_transcriptions)
result = model.transcribe_stable(video_file_path, task=transcribe_or_translate)
result.to_srt_vtt(video_file_path.rsplit('.', 1)[0] + subextension, word_level=word_level_highlight)
elapsed_time = time.time() - start_time
Expand All @@ -188,6 +190,10 @@ def gen_subtitles(video_file_path: str) -> None:
print(f"Error processing or transcribing {video_file_path}: {e}")
finally:
files_to_transcribe.remove(video_file_path)
if len(files_to_transcribe) == 0:
logging.debug("Queue is empty, clearing/releasing VRAM")
del model
gc.collect()

# Function to add a file for transcription
def add_file_for_transcription(file_path, front=True):
Expand Down

0 comments on commit 42bf55f

Please sign in to comment.