Performance improvements to captioner. Add ability to set a default caption for ace captioner to avoid needing to caption the audio so we only transcribe

extensions_built_in/captioner/AceStepCaptioner.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import librosa
 import numpy as np
 import torch
@@ -9,7 +11,7 @@ from optimum.quanto import freeze
 from toolkit.basic import flush
 from toolkit.util.quantize import quantize, get_qtype
 
-from .BaseCaptioner import BaseCaptioner
+from .BaseCaptioner import BaseCaptioner, CaptionConfig
 import transformers
 import logging
 import warnings
@@ -95,7 +97,16 @@ def analyze_audio(audio_path):
     }
 
 
+class AceStepCaptionConfig(CaptionConfig):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.fixed_caption: Optional[str] = kwargs.get("fixed_caption", None)
+
+
 class AceStepCaptioner(BaseCaptioner):
+    caption_config_class = AceStepCaptionConfig
+    caption_config: AceStepCaptionConfig
+
     def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
         super(AceStepCaptioner, self).__init__(process_id, job, config, **kwargs)
 
@@ -118,27 +129,31 @@ class AceStepCaptioner(BaseCaptioner):
         )
         if self.caption_config.low_vram:
             self.model.to("cpu")
+        
+        self.model2 = None
+        self.processor2 = None
 
-        # load captioner model
-        self.print_and_status_update("Loading captioner model")
-        self.model2 = Qwen2_5OmniForConditionalGeneration.from_pretrained(
-            self.caption_config.model_name_or_path2,
-            dtype=self.torch_dtype,
-            device_map="cpu",
-        )
-        self.model2.to(self.device_torch)
-        self.model2.disable_talker()
-        if self.caption_config.quantize:
-            self.print_and_status_update("Quantizing captioner model")
-            quantize(self.model2, weights=get_qtype(self.caption_config.qtype))
-            freeze(self.model2)
-            flush()
-        self.processor2 = Qwen2_5OmniProcessor.from_pretrained(
-            self.caption_config.model_name_or_path2,
-        )
+        if self.caption_config.fixed_caption is not None:
+            # load captioner model
+            self.print_and_status_update("Loading captioner model")
+            self.model2 = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+                self.caption_config.model_name_or_path2,
+                dtype=self.torch_dtype,
+                device_map="cpu",
+            )
+            self.model2.to(self.device_torch)
+            self.model2.disable_talker()
+            if self.caption_config.quantize:
+                self.print_and_status_update("Quantizing captioner model")
+                quantize(self.model2, weights=get_qtype(self.caption_config.qtype))
+                freeze(self.model2)
+                flush()
+            self.processor2 = Qwen2_5OmniProcessor.from_pretrained(
+                self.caption_config.model_name_or_path2,
+            )
 
-        if self.caption_config.low_vram:
-            self.model2.to("cpu")
+            if self.caption_config.low_vram:
+                self.model2.to("cpu")
         flush()
 
     def run_qwen_audio(self, model, processor, audio_data, sr, prompt_text):
@@ -228,7 +243,10 @@ class AceStepCaptioner(BaseCaptioner):
                 lyrics = lyrics.split("# Lyrics")[1].strip()
 
             # get the caption from the audio
-            caption = self.get_audio_caption(audio_data)
+            if self.caption_config.fixed_caption is not None:
+                caption = self.caption_config.fixed_caption
+            else:
+                caption = self.get_audio_caption(audio_data)
 
             output = f"<CAPTION>\n{caption}\n</CAPTION>\n"
             output += f"<LYRICS>\n{lyrics}\n</LYRICS>\n"

extensions_built_in/captioner/BaseCaptioner.py

@@ -46,6 +46,8 @@ class CaptionConfig:
 
 
 class BaseCaptioner(BaseExtensionProcess):
+    caption_config_class = CaptionConfig
+
     def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
         super(BaseCaptioner, self).__init__(process_id, job, config, **kwargs)
         self.sqlite_db_path = self.config.get("sqlite_db_path", "./aitk_db.db")
@@ -74,7 +76,7 @@ class BaseCaptioner(BaseExtensionProcess):
             self._stop_watcher_started = False
             # self.start_stop_watcher(interval_sec=2.0)
 
-        self.caption_config = CaptionConfig(**self.get_conf("caption", {}))
+        self.caption_config = self.caption_config_class(**self.get_conf("caption", {}))
         self.model = None
         self.processor = None
         self.model2 = None
@@ -85,19 +87,20 @@ class BaseCaptioner(BaseExtensionProcess):
 
     def run(self):
         super(BaseCaptioner, self).run()
-        self.start_stop_watcher()
-        self.update_status("running", "Loading Model")
-        self.load_model()
-        self.update_status("running", "Looking for files")
-        self.find_files()
-        self.update_status("running", f"Captioning {len(self.file_paths)} files")
-        self.run_caption_loop()
-        self.update_status("completed", "Captioning completed")
-        print("")
+        with torch.no_grad():
+            self.start_stop_watcher()
+            self.update_status("running", "Loading Model")
+            self.load_model()
+            self.update_status("running", "Looking for files")
+            self.find_files()
+            self.update_status("running", f"Captioning {len(self.file_paths)} files")
+            self.run_caption_loop()
+            self.update_status("completed", "Captioning completed")
+            print("")
 
-        print("****************************************************")
-        print("Captioning complete")
-        print("****************************************************")
+            print("****************************************************")
+            print("Captioning complete")
+            print("****************************************************")
 
     def run_caption_loop(self):
         for file_path in tqdm.tqdm(