Add MOSS music audio captioner

extensions_built_in/captioner/AceStepCaptioner.py

@@ -425,6 +425,14 @@ class AceStepCaptioner(BaseCaptioner):
             )
         return item
 
+    @staticmethod
+    def _place_inputs(model, inputs):
+        """Move a processor batch onto the model; the thinker also reads its
+        2D pad mask from the model. Called again before every decode pass."""
+        inputs = inputs.to(model.device).to(model.dtype)
+        model._pad_mask_2d = inputs.get("attention_mask", None)
+        return inputs
+
     def _generate(
         self,
         model,
@@ -434,17 +442,17 @@ class AceStepCaptioner(BaseCaptioner):
         retries: tuple = (),
         break_loops: bool = False,
         file_path: str = "",
+        max_new_tokens: int = MAX_NEW_TOKENS,
     ) -> str:
-        inputs = inputs.to(model.device).to(model.dtype)
+        inputs = self._place_inputs(model, inputs)
         # a generate that dies between static-cache creation and its first
         # forward leaves a half-built cache that breaks every later call
         stale_cache = getattr(model, "_cache", None)
         if stale_cache is not None and not stale_cache.is_initialized:
             del model._cache
-        model._pad_mask_2d = inputs.get("attention_mask", None)
         input_len = inputs["input_ids"].shape[1]
         # greedy decode loops on repeated phrases without a penalty
-        gen_kwargs = {"max_new_tokens": MAX_NEW_TOKENS, "repetition_penalty": repetition_penalty}
+        gen_kwargs = {"max_new_tokens": max_new_tokens, "repetition_penalty": repetition_penalty}
         if model.generation_config.cache_implementation == "static":
             if input_len + 16 < STATIC_MAX_LENGTH:
                 from transformers.generation import MaxLengthCriteria, StoppingCriteriaList
@@ -453,7 +461,7 @@ class AceStepCaptioner(BaseCaptioner):
                     "repetition_penalty": repetition_penalty,
                     "max_length": STATIC_MAX_LENGTH,
                     "stopping_criteria": StoppingCriteriaList(
-                        [MaxLengthCriteria(max_length=min(input_len + MAX_NEW_TOKENS, STATIC_MAX_LENGTH))]
+                        [MaxLengthCriteria(max_length=min(input_len + max_new_tokens, STATIC_MAX_LENGTH))]
                     ),
                 }
             else:
@@ -473,7 +481,7 @@ class AceStepCaptioner(BaseCaptioner):
             kwargs = dict(gen_kwargs)
             if breaker is not None:
                 kwargs["logits_processor"] = LogitsProcessorList([breaker])
-            model._pad_mask_2d = inputs.get("attention_mask", None)
+            self._place_inputs(model, inputs)
             generated = model.generate(**inputs, **kwargs)
             return strip_eos(generated[0, input_len:]), breaker
 
@@ -533,25 +541,29 @@ class AceStepCaptioner(BaseCaptioner):
             file_path=file_path,
         )
 
-    def get_audio_caption(self, inputs) -> str:
+    def get_audio_caption(self, inputs, file_path: str = "") -> str:
         if self.caption_config.low_vram and self.model.device != torch.device("cpu"):
             self.model.to("cpu")
         if self.model2.device == torch.device("cpu"):
             self.model2.to(self.device_torch)
         return self._generate(self.model2, self.processor2, inputs, CAPTIONER_REPETITION_PENALTY)
 
+    def _transcribe_item(self, item: dict) -> str:
+        """Raw transcriber text for a prepped item: '# Languages ... # Lyrics ...'."""
+        return self.get_audio_lyrics(item["lyrics_inputs"], item["file"])
+
     def _caption_item(self, item: dict) -> str:
         if self.caption_config.caption_format == "yue2":
-            lyrics = clean_lyrics(self.get_audio_lyrics(item["lyrics_inputs"], item["file"]))
+            lyrics = clean_lyrics(self._transcribe_item(item))
             if self.caption_config.fixed_caption is not None:
                 caption = self.caption_config.fixed_caption
             else:
-                caption = self.get_audio_caption(item["caption_inputs"])
+                caption = self.get_audio_caption(item["caption_inputs"], item["file"])
             caption = " ".join(caption.split())
             return f"{caption}\n[Lyrics]\n{lyrics}"
 
         analysis = item["analysis"]
-        lyrics = self.get_audio_lyrics(item["lyrics_inputs"], item["file"])
+        lyrics = self._transcribe_item(item)
 
         language = "en"
         if "# Languages" in lyrics and "# Lyrics" in lyrics:
@@ -563,7 +575,7 @@ class AceStepCaptioner(BaseCaptioner):
         if self.caption_config.fixed_caption is not None:
             caption = self.caption_config.fixed_caption
         else:
-            caption = self.get_audio_caption(item["caption_inputs"])
+            caption = self.get_audio_caption(item["caption_inputs"], item["file"])
 
         output = f"<CAPTION>\n{caption}\n</CAPTION>\n"
         output += f"<LYRICS>\n{lyrics}\n</LYRICS>\n"

extensions_built_in/captioner/MossMusicCaptioner.py

@@ -0,0 +1,169 @@
+"""MOSS-Music captioner: one 8B audio-language model (OpenMOSS-Team/MOSS-Music-8B-Instruct)
+does both the lyric transcription and the description that AceStepCaptioner
+runs two Qwen2.5-Omni checkpoints for. Output layouts, loop breaking and the
+worker pipeline are inherited from AceStepCaptioner."""
+
+import re
+
+from optimum.quanto import freeze
+
+from toolkit.audio.moss_music import HF_REPO, load_moss_music
+from toolkit.basic import flush
+from toolkit.util.quantize import get_qtype, quantize
+
+from .AceStepCaptioner import (
+    CAPTIONER_REPETITION_PENALTY,
+    TRANSCRIBER_REPETITION_PENALTY,
+    TRANSCRIBER_RETRIES,
+    AceStepCaptionConfig,
+    AceStepCaptioner,
+    analyze_audio,
+)
+
+# timestamps are requested on purpose: each line is anchored to the input time
+# markers, so decode has to advance through the song instead of looping on a
+# repeated phrase. They are stripped from the saved caption.
+TRANSCRIBE_PROMPT = "Transcribe the lyrics of this song with timestamps."
+# the model writes multi-section essays with timings, chords and lyric quotes
+# unless framed as a generator prompt; this keeps it to one plain paragraph
+CAPTION_PROMPT = (
+    "Write a prompt that a text-to-music generator could use to recreate this track. One paragraph, "
+    "under 80 words: genre, mood, instrumentation, tempo feel, production style, vocal style. "
+    "No timestamps, section timings, chord names, or lyric quotes."
+)
+LANGUAGE_PROMPT = (
+    "What language are the lyrics sung in? Answer with only the two-letter "
+    "ISO 639-1 language code, or none if there are no lyrics."
+)
+LANGUAGE_MAX_NEW_TOKENS = 8
+
+# the model writes "[00:18.00 - 00:22.00] line"; a looping decode can run the
+# seconds field past 59 ("[00:100.00 - ...]"), so any bracketed run of time
+# characters is treated as a timestamp, plus bare mm:ss spans at a line start
+_TIME = r"(?:\d{1,2}:)?\d{1,3}:\d{1,3}(?:[.:]\d{1,3})?|\d+(?:\.\d+)?\s*s"
+_SPAN = rf"(?:{_TIME})(?:\s*(?:-{{1,3}}>?|–|—|~|to)\s*(?:{_TIME}))?"
+_TIME_CHARS = r"[\d:.\s\-–—>~s]*\d[\d:.\s\-–—>~s]*"
+TIMESTAMP_RE = re.compile(
+    rf"\[{_TIME_CHARS}\]|\({_TIME_CHARS}\)|<{_TIME_CHARS}>|^\s*{_SPAN}\s*[:\-–—|]?",
+    re.MULTILINE,
+)
+
+
+def strip_timestamps(text: str) -> str:
+    """Remove time annotations and tidy the line layout they leave behind."""
+    text = TIMESTAMP_RE.sub("", text)
+    lines = [" ".join(line.split()) for line in text.splitlines()]
+    out = []
+    for line in lines:
+        if line == "" and (not out or out[-1] == ""):
+            continue
+        out.append(line)
+    return "\n".join(out).strip()
+
+
+class MossMusicCaptionConfig(AceStepCaptionConfig):
+    def __init__(self, **kwargs):
+        kwargs.setdefault("model_name_or_path", HF_REPO)
+        super().__init__(**kwargs)
+        self.caption_prompt: str = kwargs.get("caption_prompt") or CAPTION_PROMPT
+        self.transcribe_prompt: str = kwargs.get("transcribe_prompt") or TRANSCRIBE_PROMPT
+        # MOSS transcribes the full mix better than a separated vocal stem
+        self.extract_vocals_before_transcribe = False
+        # leave the model's "[00:18.00 - 00:22.00]" line times in the saved lyrics
+        self.keep_timestamps: bool = bool(kwargs.get("keep_timestamps", False))
+
+
+class MossMusicCaptioner(AceStepCaptioner):
+    caption_config_class = MossMusicCaptionConfig
+    caption_config: MossMusicCaptionConfig
+
+    def load_model(self):
+        self.print_and_status_update("Loading MOSS-Music model")
+        self.model, self.processor = load_moss_music(
+            self.caption_config.model_name_or_path,
+            dtype=self.torch_dtype,
+            device=str(self.device_torch),
+        )
+        if self.caption_config.quantize:
+            self.print_and_status_update("Quantizing MOSS-Music model")
+            quantize(self.model, weights=get_qtype(self.caption_config.qtype))
+            freeze(self.model)
+            flush()
+        # one model does both jobs; AceStepCaptioner's second slot stays empty
+        self.model2 = None
+        self.processor2 = None
+        flush()
+
+    @staticmethod
+    def _place_inputs(model, inputs):
+        inputs = inputs.to(model.device)
+        inputs["audio_data"] = inputs["audio_data"].to(model.dtype)
+        return inputs
+
+    def _make_inputs(self, audio_data, prompt_text: str):
+        inputs = self.processor(text=prompt_text, audios=[audio_data])
+        inputs["audio_input_mask"] = inputs["input_ids"] == self.processor.audio_token_id
+        return inputs
+
+    def _prep_file(self, file_path: str) -> dict:
+        audio_data = self._load_audio(file_path)
+        is_ace_step = self.caption_config.caption_format == "ace_step"
+        item = {
+            "file": file_path,
+            "analysis": analyze_audio(file_path) if is_ace_step else None,
+            "lyrics_inputs": self._make_inputs(audio_data, self.caption_config.transcribe_prompt),
+        }
+        if is_ace_step:
+            item["language_inputs"] = self._make_inputs(audio_data, LANGUAGE_PROMPT)
+        if self.caption_config.fixed_caption is None:
+            item["caption_inputs"] = self._make_inputs(
+                audio_data, self.caption_config.caption_prompt
+            )
+        return item
+
+    def get_audio_lyrics(self, inputs, file_path: str = "") -> str:
+        text = self._generate(
+            self.model,
+            self.processor,
+            inputs,
+            TRANSCRIBER_REPETITION_PENALTY,
+            retries=TRANSCRIBER_RETRIES,
+            break_loops=True,
+            file_path=file_path,
+        )
+        return text.strip() if self.caption_config.keep_timestamps else strip_timestamps(text)
+
+    def get_audio_language(self, inputs) -> str:
+        answer = self._generate(
+            self.model,
+            self.processor,
+            inputs,
+            repetition_penalty=1.0,
+            max_new_tokens=LANGUAGE_MAX_NEW_TOKENS,
+        )
+        match = re.search(r"[a-z]{2,3}", answer.lower())
+        code = match.group(0) if match else ""
+        # instrumentals answer none; ACE-Step's default language is en
+        return code if len(code) == 2 and code != "no" else "en"
+
+    def _transcribe_item(self, item: dict) -> str:
+        lyrics = self.get_audio_lyrics(item["lyrics_inputs"], item["file"])
+        if "language_inputs" not in item:
+            return lyrics
+        # the layout AceStepCaptioner._caption_item parses from the ACE transcriber
+        language = self.get_audio_language(item["language_inputs"])
+        return f"# Languages\n{language}\n# Lyrics\n{lyrics}"
+
+    def get_audio_caption(self, inputs, file_path: str = "") -> str:
+        # a tag-list prompt can degenerate into a repeated tag; same loop
+        # breaking as transcription
+        caption = self._generate(
+            self.model,
+            self.processor,
+            inputs,
+            CAPTIONER_REPETITION_PENALTY,
+            retries=TRANSCRIBER_RETRIES,
+            break_loops=True,
+            file_path=file_path,
+        )
+        return strip_timestamps(caption)

extensions_built_in/captioner/__init__.py

@@ -13,6 +13,18 @@ class AceStepCaptionerExtension(Extension):
         return AceStepCaptioner
 
 
+class MossMusicCaptionerExtension(Extension):
+    uid = "MossMusicCaptioner"
+    name = "MOSS-Music Captioner"
+
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .MossMusicCaptioner import MossMusicCaptioner
+
+        return MossMusicCaptioner
+
+
 class Qwen3VLCaptionerExtension(Extension):
     uid = "Qwen3VLCaptioner"
     name = "Qwen 3VL Captioner"
@@ -63,6 +75,7 @@ class Ideogram4CaptionerExtension(Extension):
 
 AI_TOOLKIT_EXTENSIONS = [
     AceStepCaptionerExtension,
+    MossMusicCaptionerExtension,
     Qwen3VLCaptionerExtension,
     Qwen3OmniCaptionerExtension,
     Qwen25OmniCaptionerExtension,