Merge pull request #7 from karl0ss/reworked

add support for single file and recursive folder

bazarr-ai-sub-generator/cli.py

@@ -15,16 +15,6 @@ def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
-    # parser.add_argument(
-    #     "--audio_channel", default="0", type=int, help="audio channel index to use"
-    # )
-    # parser.add_argument(
-    #     "--sample_interval",
-    #     type=str2timeinterval,
-    #     default=None,
-    #     help="generate subtitles for a specific \
-    #                           fragment of the video (e.g. 01:02:05-01:03:45)",
-    # )
     parser.add_argument(
         "--model",
         default="small",
@@ -55,6 +45,18 @@ def main():
     #     help="Type to use for computation. \
     #                           See https://opennmt.net/CTranslate2/quantization.html.",
     # )
+    parser.add_argument(
+        "--file",
+        type=str,
+        default=None,
+        help="Process a single file"
+    )
+    parser.add_argument(
+        "--folder",
+        type=str,
+        default=None,
+        help="Process all videos in folder"
+    )
     parser.add_argument(
         "--show",
         type=str,
@@ -70,7 +72,12 @@ def main():
         help="What is the origin language of the video? \
                               If unset, it is detected automatically.",
     )
-
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default="whisper",
+        choices=["whisper", "faster_whisper"],
+    )
     args = parser.parse_args().__dict__
 
     process(args)

bazarr-ai-sub-generator/main.py

@@ -6,28 +6,39 @@ 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
 from utils.sonarr import update_show_in_sonarr
-# from utils.faster_whisper import WhisperAI
+from utils.faster_whisper import WhisperAI as fasterWhisperAI
 from utils.whisper import WhisperAI
 from utils.decorator import measure_time
 
-def process(args: dict):
-    
-    model_name: str = args.pop("model")
-    language: str = args.pop("language")
-    show: str = args.pop("show")
-    
-    if model_name.endswith(".en"):
-        warnings.warn(
-            f"{model_name} is an English-only model, forcing English detection."
-        )
-        args["language"] = "en"
-    # if translate task used and language argument is set, then use it
-    elif language != "auto":
-        args["language"] = language
 
-    model_args = {}
+
-    model_args["device"] = args.pop("device")
+def folder_flow(folder, model_args, args, backend):
-    
+        print(f"Processing {folder}")
+        files = os.listdir(folder)
+        for file in files:
+            print(f"processing {file}")
+            path = folder+file
+            try:
+                audios = get_audio([path], 0, None)
+                subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend)
+
+                add_subtitles_to_mp4(subtitles)
+                time.sleep(5)
+            except Exception as ex:
+                print(f"skipping file due to - {ex}")
+
+def file_flow(show, model_args, args, backend):
+        print(f"Processing {show}")
+        try:
+            audios = get_audio([show], 0, None)
+            subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend)
+
+            add_subtitles_to_mp4(subtitles)
+            time.sleep(5)
+        except Exception as ex:
+            print(f"skipping file due to - {ex}")
+
+def bazzar_flow(show, model_args, args, backend):
     list_of_episodes_needing_subtitles = get_wanted_episodes(show)
     print(
         f"Found {list_of_episodes_needing_subtitles['total']} episodes needing subtitles."
@@ -37,7 +48,7 @@ def process(args: dict):
         episode_data = get_episode_details(episode["sonarrEpisodeId"])
         try:
             audios = get_audio([episode_data["path"]], 0, None)
-            subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
+            subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend)
 
             add_subtitles_to_mp4(subtitles)
             update_show_in_sonarr(episode["sonarrSeriesId"])
@@ -46,12 +57,15 @@ def process(args: dict):
         except Exception as ex:
             print(f"skipping file due to - {ex}")
 
+
 @measure_time
 def get_subtitles(
-    audio_paths: list, output_dir: str, model_args: dict, transcribe_args: dict
+    audio_paths: list, output_dir: str, model_args: dict, transcribe_args: dict, backend: str
 ):
-    model = WhisperAI(model_args, transcribe_args)
+    if backend == 'whisper':
-
+        model = WhisperAI(model_args, transcribe_args)
+    else:
+        model = fasterWhisperAI(model_args, transcribe_args)
     subtitles_path = {}
 
     for path, audio_path in audio_paths.items():
@@ -66,3 +80,32 @@ def get_subtitles(
         subtitles_path[path] = srt_path
 
     return subtitles_path
+
+
+def process(args: dict):
+
+    model_name: str = args.pop("model")
+    language: str = args.pop("language")
+    show: str = args.pop("show")
+    file: str = args.pop("file")
+    folder: str = args.pop("folder")
+    backend: str = args.pop("backend")
+
+    if model_name.endswith(".en"):
+        warnings.warn(
+            f"{model_name} is an English-only model, forcing English detection."
+        )
+        args["language"] = "en"
+    # if translate task used and language argument is set, then use it
+    elif language != "auto":
+        args["language"] = language
+
+    model_args = {}
+    model_args["device"] = args.pop("device")
+
+    if file:
+        file_flow(file, model_args, args, backend)
+    elif folder:
+        folder_flow(folder, model_args, args, backend)
+    else:
+        bazzar_flow(show, model_args, args, backend)