latest commit

bazarr-ai-sub-generator/cli.py

@@ -89,6 +89,13 @@ def main():
         help="What is the origin language of the video? \
                                If unset, it is detected automatically.",
     )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Number of concurrent workers for processing episodes. \
+                               Increase for better CUDA utilization with multiple episodes.",
+    )
 
     args = parser.parse_args().__dict__
 

bazarr-ai-sub-generator/main.py

@@ -2,6 +2,8 @@ import os
 import warnings
 import tempfile
 import time
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from utils.files import filename, write_srt
 from utils.ffmpeg import get_audio, add_subtitles_to_mp4
 from utils.bazarr import get_wanted_episodes, get_episode_details, sync_series
@@ -9,11 +11,43 @@ from utils.sonarr import update_show_in_sonarr
 from utils.whisper import WhisperAI
 
 
+def process_episode(episode, model_args, args, audio_channel, sample_interval, processing_episodes, completed_episodes):
+    """Process a single episode for subtitle generation."""
+    episode_id = episode["sonarrEpisodeId"]
+
+    try:
+        # Double-check that this episode is still wanted before processing
+        current_wanted = get_wanted_episodes()
+        still_wanted = any(ep["sonarrEpisodeId"] == episode_id for ep in current_wanted["data"])
+
+        if not still_wanted:
+            processing_episodes.discard(episode_id)
+            return f"Skipped (no longer wanted): {episode['seriesTitle']} - {episode['episode_number']}"
+
+        print(f"Processing {episode['seriesTitle']} - {episode['episode_number']}")
+        episode_data = get_episode_details(episode_id)
+        audios = get_audio([episode_data["path"]], audio_channel, sample_interval)
+        subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
+
+        add_subtitles_to_mp4(subtitles)
+        update_show_in_sonarr(episode["sonarrSeriesId"])
+        time.sleep(5)
+        sync_series()
+
+        processing_episodes.discard(episode_id)
+        completed_episodes.append(episode_id)
+        return f"Completed: {episode['seriesTitle']} - {episode['episode_number']}"
+    except Exception as e:
+        processing_episodes.discard(episode_id)
+        return f"Failed {episode['seriesTitle']} - {episode['episode_number']}: {str(e)}"
+
+
 def process(args: dict):
     model_name: str = args.pop("model")
     language: str = args.pop("language")
     sample_interval: str = args.pop("sample_interval")
     audio_channel: str = args.pop("audio_channel")
+    workers: int = args.pop("workers", 1)
 
     if model_name.endswith(".en"):
         warnings.warn(
@@ -33,16 +67,38 @@ def process(args: dict):
     print(
         f"Found {list_of_episodes_needing_subtitles['total']} episodes needing subtitles."
     )
-    for episode in list_of_episodes_needing_subtitles["data"]:
-        print(f"Processing {episode['seriesTitle']} - {episode['episode_number']}")
-        episode_data = get_episode_details(episode["sonarrEpisodeId"])
-        audios = get_audio([episode_data["path"]], audio_channel, sample_interval)
-        subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
+    print(f"Processing with {workers} concurrent worker(s)...")
 
-        add_subtitles_to_mp4(subtitles)
-        update_show_in_sonarr(episode["sonarrSeriesId"])
-        time.sleep(5)
-        sync_series()
+    # Thread-safe tracking of episodes being processed and completed
+    processing_episodes = set()
+    completed_episodes_list = []
+    total_episodes = len(list_of_episodes_needing_subtitles["data"])
+
+    # Filter episodes to avoid duplicates and respect concurrent processing limits
+    episodes_to_process = []
+    for episode in list_of_episodes_needing_subtitles["data"]:
+        episode_id = episode["sonarrEpisodeId"]
+        if episode_id not in processing_episodes:
+            processing_episodes.add(episode_id)
+            episodes_to_process.append(episode)
+
+    print(f"Starting processing of {len(episodes_to_process)} unique episodes...")
+
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        # Submit episodes for processing with tracking sets
+        future_to_episode = {
+            executor.submit(process_episode, episode, model_args, args, audio_channel, sample_interval, processing_episodes, completed_episodes_list): episode
+            for episode in episodes_to_process
+        }
+
+        # Collect results as they complete
+        completed_count = 0
+        for future in as_completed(future_to_episode):
+            completed_count += 1
+            result = future.result()
+            print(f"[{completed_count}/{total_episodes}] {result}")
+
+    print(f"Processing complete. {len(completed_episodes_list)} episodes processed successfully.")
 
 
 def get_subtitles(