Add MOSS music audio captioner
This commit is contained in:
Jaret Burkett committed
1 parent
080873bec9
commit
a9e7752178
11 files changed
+1255
-11
No files matched your search
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
The code in this directory is adapted from MOSS-Music by the OpenMOSS Team
|
||||
(https://github.com/OpenMOSS/MOSS-Music), licensed under the Apache License,
|
||||
Version 2.0. Modifications for ai-toolkit: transformers 5.x compatibility
|
||||
(WhisperEncoderLayer tensor return, generate input preparation), removal of the
|
||||
chat-template loader, and direct Qwen2Tokenizer loading.
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,35 @@
|
||||
import torch
|
||||
|
||||
from .configuration import MossMusicConfig, MossMusicEncoderConfig
|
||||
from .modeling import MossMusicModel
|
||||
from .processing import MelConfig, MossMusicProcessor
|
||||
|
||||
HF_REPO = "OpenMOSS-Team/MOSS-Music-8B-Instruct"
|
||||
|
||||
|
||||
def load_moss_music(
|
||||
name_or_path: str = HF_REPO,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
device="cpu",
|
||||
enable_time_marker: bool = True,
|
||||
):
|
||||
"""Model + processor from a hub id or local checkpoint dir. Time markers
|
||||
(elapsed seconds interleaved into the audio tokens) are what the model was
|
||||
trained with and keep long transcriptions from looping; leave them on."""
|
||||
model = MossMusicModel.from_pretrained(name_or_path, dtype=dtype, device_map=device)
|
||||
model.eval()
|
||||
processor = MossMusicProcessor.from_pretrained(
|
||||
name_or_path, enable_time_marker=enable_time_marker
|
||||
)
|
||||
return model, processor
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HF_REPO",
|
||||
"MelConfig",
|
||||
"MossMusicConfig",
|
||||
"MossMusicEncoderConfig",
|
||||
"MossMusicModel",
|
||||
"MossMusicProcessor",
|
||||
"load_moss_music",
|
||||
]
|
||||
@@ -0,0 +1,135 @@
|
||||
from typing import List, Optional
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from transformers import PretrainedConfig, Qwen3Config
|
||||
|
||||
|
||||
@dataclass
|
||||
class MossMusicEncoderConfig:
|
||||
d_model: int = 1280
|
||||
output_dim: int = 1280
|
||||
num_mel_bins: int = 128
|
||||
encoder_layers: int = 32
|
||||
encoder_attention_heads: int = 20
|
||||
encoder_ffn_dim: int = 5120
|
||||
downsample_rate: int = 8
|
||||
downsample_hidden_size: int = 480
|
||||
encoder_attention_window_size: int = 100
|
||||
max_source_positions: int = 1500
|
||||
dropout: float = 0.1
|
||||
attention_dropout: float = 0.1
|
||||
activation_dropout: float = 0.0
|
||||
activation_function: str = "gelu"
|
||||
layer_norm_eps: float = 1e-5
|
||||
_attn_implementation: str = "eager"
|
||||
pretrained_path: str = ""
|
||||
|
||||
n_window: int = 200
|
||||
conv_chunksize: int = 64
|
||||
|
||||
deepstack_encoder_layer_indexes: List[int] = field(
|
||||
default_factory=lambda: [8, 16, 24]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, config_dict):
|
||||
if config_dict is None:
|
||||
return cls()
|
||||
allowed_keys = set(cls.__dataclass_fields__.keys())
|
||||
filtered = {k: v for k, v in config_dict.items() if k in allowed_keys}
|
||||
return cls(**filtered)
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
"d_model": self.d_model,
|
||||
"output_dim": self.output_dim,
|
||||
"num_mel_bins": self.num_mel_bins,
|
||||
"encoder_layers": self.encoder_layers,
|
||||
"encoder_attention_heads": self.encoder_attention_heads,
|
||||
"encoder_ffn_dim": self.encoder_ffn_dim,
|
||||
"downsample_rate": self.downsample_rate,
|
||||
"downsample_hidden_size": self.downsample_hidden_size,
|
||||
"encoder_attention_window_size": self.encoder_attention_window_size,
|
||||
"max_source_positions": self.max_source_positions,
|
||||
"dropout": self.dropout,
|
||||
"attention_dropout": self.attention_dropout,
|
||||
"activation_dropout": self.activation_dropout,
|
||||
"activation_function": self.activation_function,
|
||||
"layer_norm_eps": self.layer_norm_eps,
|
||||
"_attn_implementation": self._attn_implementation,
|
||||
"pretrained_path": self.pretrained_path,
|
||||
"n_window": self.n_window,
|
||||
"conv_chunksize": self.conv_chunksize,
|
||||
"deepstack_encoder_layer_indexes": list(
|
||||
self.deepstack_encoder_layer_indexes or []
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class MossMusicConfig(PretrainedConfig):
|
||||
model_type = "moss_music"
|
||||
is_composition = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
audio_config=None,
|
||||
language_config=None,
|
||||
adapter_hidden_size=8192,
|
||||
ignore_index=-100,
|
||||
deepstack_num_inject_layers: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
if isinstance(audio_config, dict):
|
||||
audio_config = MossMusicEncoderConfig.from_dict(audio_config)
|
||||
elif audio_config is None:
|
||||
audio_config = MossMusicEncoderConfig()
|
||||
|
||||
if isinstance(language_config, dict):
|
||||
language_config = Qwen3Config(**language_config)
|
||||
elif language_config is None:
|
||||
language_config = Qwen3Config()
|
||||
|
||||
self.audio_config = audio_config
|
||||
self.language_config = language_config
|
||||
self.adapter_hidden_size = adapter_hidden_size
|
||||
self.ignore_index = ignore_index
|
||||
self.deepstack_num_inject_layers = deepstack_num_inject_layers
|
||||
|
||||
_propagate_keys = {
|
||||
"num_hidden_layers",
|
||||
"eos_token_id",
|
||||
"bos_token_id",
|
||||
"vocab_size",
|
||||
"tie_word_embeddings",
|
||||
}
|
||||
for key in ("num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size"):
|
||||
kwargs.setdefault(key, getattr(language_config, key, None))
|
||||
kwargs.setdefault("tie_word_embeddings", False)
|
||||
|
||||
if hasattr(language_config, "to_dict"):
|
||||
_lang_keys = set(language_config.to_dict().keys())
|
||||
for key in list(kwargs.keys()):
|
||||
if key in _lang_keys and key not in _propagate_keys:
|
||||
kwargs.pop(key)
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def to_dict(self):
|
||||
output = super().to_dict()
|
||||
output["audio_config"] = (
|
||||
self.audio_config.to_dict()
|
||||
if hasattr(self.audio_config, "to_dict")
|
||||
else self.audio_config
|
||||
)
|
||||
output["language_config"] = (
|
||||
self.language_config.to_dict()
|
||||
if hasattr(self.language_config, "to_dict")
|
||||
else self.language_config
|
||||
)
|
||||
output["adapter_hidden_size"] = self.adapter_hidden_size
|
||||
output["ignore_index"] = self.ignore_index
|
||||
output["deepstack_num_inject_layers"] = self.deepstack_num_inject_layers
|
||||
return output
|
||||
|
||||
|
||||
__all__ = ["MossMusicEncoderConfig", "MossMusicConfig"]
|
||||
@@ -0,0 +1,568 @@
|
||||
from typing import Any, List, Optional, Tuple
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from transformers.generation.utils import GenerationMixin
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer, Qwen3Model
|
||||
from transformers.models.whisper.modeling_whisper import WhisperEncoderLayer
|
||||
|
||||
from .configuration import MossMusicConfig, MossMusicEncoderConfig
|
||||
|
||||
|
||||
class SinusoidsPositionEmbedding(nn.Module):
|
||||
def __init__(self, num_positions: int, embedding_dim: int):
|
||||
super().__init__()
|
||||
max_timescale = 10000.0
|
||||
log_timescale_increment = math.log(max_timescale) / (embedding_dim // 2 - 1)
|
||||
inv_timescales = torch.exp(
|
||||
-log_timescale_increment * torch.arange(embedding_dim // 2).float()
|
||||
)
|
||||
self.register_buffer("inv_timescales", inv_timescales, persistent=False)
|
||||
|
||||
def forward(self, seq_len: int, device: torch.device):
|
||||
scaled_time = torch.arange(
|
||||
seq_len, device=device, dtype=self.inv_timescales.dtype
|
||||
).unsqueeze(1) * self.inv_timescales.unsqueeze(0)
|
||||
sin_emb = torch.sin(scaled_time)
|
||||
cos_emb = torch.cos(scaled_time)
|
||||
pos_emb = torch.cat([sin_emb, cos_emb], dim=1)
|
||||
return pos_emb.unsqueeze(0)
|
||||
|
||||
|
||||
class MossMusicEncoder(nn.Module):
|
||||
"""Audio encoder with conv-stem downsampling and Whisper transformer layers."""
|
||||
|
||||
def __init__(self, config: MossMusicEncoderConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.gelu = nn.GELU()
|
||||
|
||||
self.conv1 = nn.Conv2d(
|
||||
1,
|
||||
config.downsample_hidden_size,
|
||||
kernel_size=(3, 3),
|
||||
stride=(2, 2),
|
||||
padding=(1, 1),
|
||||
)
|
||||
self.conv2 = nn.Conv2d(
|
||||
config.downsample_hidden_size,
|
||||
config.downsample_hidden_size,
|
||||
kernel_size=(3, 3),
|
||||
stride=(2, 2),
|
||||
padding=(1, 1),
|
||||
)
|
||||
self.conv3 = nn.Conv2d(
|
||||
config.downsample_hidden_size,
|
||||
config.downsample_hidden_size,
|
||||
kernel_size=(3, 3),
|
||||
stride=(2, 2),
|
||||
padding=(1, 1),
|
||||
)
|
||||
|
||||
# 128 mel bins / 8 = 16 after 3 convs with stride=2
|
||||
self.stem_proj = nn.Linear(config.downsample_hidden_size * 16, config.d_model)
|
||||
self.embed_positions = SinusoidsPositionEmbedding(
|
||||
config.max_source_positions, config.d_model
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[WhisperEncoderLayer(config) for _ in range(config.encoder_layers)]
|
||||
)
|
||||
self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
|
||||
self.out_proj = (
|
||||
nn.Linear(config.d_model, config.output_dim, bias=False)
|
||||
if config.output_dim != config.d_model
|
||||
else nn.Identity()
|
||||
)
|
||||
|
||||
self.deepstack_encoder_layer_indexes = list(
|
||||
config.deepstack_encoder_layer_indexes or []
|
||||
)
|
||||
self._deepstack_capture_map = {
|
||||
layer_idx: capture_idx
|
||||
for capture_idx, layer_idx in enumerate(self.deepstack_encoder_layer_indexes)
|
||||
}
|
||||
|
||||
self.n_window = int(config.n_window)
|
||||
self.chunk_frames = int(self.n_window * 2)
|
||||
self.conv_chunksize = int(config.conv_chunksize)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return self.conv1.weight.dtype
|
||||
|
||||
@staticmethod
|
||||
def _compute_downsampled_length(lengths: torch.Tensor) -> torch.Tensor:
|
||||
def conv_out_len(L):
|
||||
return (L - 1) // 2 + 1
|
||||
|
||||
return conv_out_len(conv_out_len(conv_out_len(lengths)))
|
||||
|
||||
def _encode_chunk_batch(
|
||||
self,
|
||||
input_features: torch.Tensor,
|
||||
seq_lengths: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
"""Encode a batch of (already padded) chunks through the conv stem and
|
||||
transformer layers. Returns (last_hidden, ordered_deepstack_hidden_states).
|
||||
"""
|
||||
if input_features.dim() == 2:
|
||||
input_features = input_features.unsqueeze(0)
|
||||
|
||||
downsampled_lengths = self._compute_downsampled_length(seq_lengths)
|
||||
|
||||
# [B, n_mels, T] -> [B, 1, n_mels, T]
|
||||
x = input_features.unsqueeze(1)
|
||||
x = self.gelu(self.conv1(x))
|
||||
x = self.gelu(self.conv2(x))
|
||||
x = self.gelu(self.conv3(x))
|
||||
|
||||
# [B, C, F, T] -> [B, T, C*F]
|
||||
x = x.permute(0, 3, 1, 2).contiguous().flatten(2)
|
||||
x = self.stem_proj(x)
|
||||
|
||||
max_len = int(downsampled_lengths.max().item())
|
||||
if x.size(1) > max_len:
|
||||
x = x[:, :max_len, :]
|
||||
|
||||
positions = self.embed_positions(x.shape[1], x.device)
|
||||
x = x + positions.to(x.dtype)
|
||||
|
||||
padding_mask = (
|
||||
torch.arange(x.size(1), device=x.device)[None, :]
|
||||
>= downsampled_lengths[:, None]
|
||||
)
|
||||
attention_mask = (1.0 - (~padding_mask).to(dtype=x.dtype)) * torch.finfo(
|
||||
x.dtype
|
||||
).min
|
||||
attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)
|
||||
|
||||
deepstack_hidden_states: List[Optional[torch.Tensor]] = [None] * len(
|
||||
self.deepstack_encoder_layer_indexes
|
||||
)
|
||||
for layer_idx, layer in enumerate(self.layers):
|
||||
x = layer(x, attention_mask)
|
||||
capture_idx = self._deepstack_capture_map.get(layer_idx)
|
||||
if capture_idx is not None:
|
||||
deepstack_hidden_states[capture_idx] = x
|
||||
|
||||
x = self.layer_norm(x)
|
||||
x = self.out_proj(x)
|
||||
|
||||
ordered_deepstack_hidden_states = [
|
||||
h for h in deepstack_hidden_states if h is not None
|
||||
]
|
||||
if not isinstance(self.out_proj, nn.Identity):
|
||||
ordered_deepstack_hidden_states = [
|
||||
self.out_proj(h) for h in ordered_deepstack_hidden_states
|
||||
]
|
||||
return x, ordered_deepstack_hidden_states
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_features: torch.Tensor,
|
||||
feature_lens: Optional[torch.Tensor] = None,
|
||||
output_deepstack_hidden_states: bool = True,
|
||||
) -> BaseModelOutputWithPast:
|
||||
if input_features.dim() == 3:
|
||||
if feature_lens is None:
|
||||
feature_lens = torch.full(
|
||||
(input_features.size(0),),
|
||||
input_features.size(-1),
|
||||
dtype=torch.long,
|
||||
device=input_features.device,
|
||||
)
|
||||
else:
|
||||
feature_lens = feature_lens.to(
|
||||
device=input_features.device, dtype=torch.long
|
||||
)
|
||||
valid_chunks = [
|
||||
input_features[i, :, : int(feature_lens[i].item())]
|
||||
for i in range(int(input_features.shape[0]))
|
||||
]
|
||||
input_features = torch.cat(valid_chunks, dim=1)
|
||||
elif input_features.dim() != 2:
|
||||
raise ValueError(
|
||||
f"Expected [n_mels, T] or [B, n_mels, T], got {tuple(input_features.shape)}."
|
||||
)
|
||||
|
||||
if feature_lens is None:
|
||||
feature_lens = torch.tensor(
|
||||
[int(input_features.shape[1])],
|
||||
device=input_features.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
else:
|
||||
feature_lens = feature_lens.to(
|
||||
device=input_features.device, dtype=torch.long
|
||||
)
|
||||
|
||||
chunk_frames = int(self.chunk_frames)
|
||||
chunk_num = torch.ceil(
|
||||
feature_lens.to(torch.float32) / float(chunk_frames)
|
||||
).long()
|
||||
chunk_lengths = torch.full(
|
||||
(int(chunk_num.sum().item()),),
|
||||
chunk_frames,
|
||||
dtype=torch.long,
|
||||
device=feature_lens.device,
|
||||
)
|
||||
tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
|
||||
chunk_lengths[tail_chunk_index] = feature_lens % chunk_frames
|
||||
chunk_lengths[chunk_lengths == 0] = chunk_frames
|
||||
|
||||
chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
|
||||
padded_feature = nn.utils.rnn.pad_sequence(
|
||||
chunk_list, batch_first=True
|
||||
).transpose(1, 2)
|
||||
|
||||
feature_lens_after_cnn = self._compute_downsampled_length(chunk_lengths)
|
||||
t_down_max = (
|
||||
int(feature_lens_after_cnn.max().item())
|
||||
if feature_lens_after_cnn.numel() > 0
|
||||
else 0
|
||||
)
|
||||
padded_mask_after_cnn = nn.utils.rnn.pad_sequence(
|
||||
[
|
||||
torch.ones(int(L.item()), dtype=torch.bool, device=padded_feature.device)
|
||||
for L in feature_lens_after_cnn
|
||||
],
|
||||
batch_first=True,
|
||||
)
|
||||
if padded_mask_after_cnn.shape[1] < t_down_max:
|
||||
padded_mask_after_cnn = F.pad(
|
||||
padded_mask_after_cnn,
|
||||
(0, t_down_max - padded_mask_after_cnn.shape[1]),
|
||||
value=False,
|
||||
)
|
||||
|
||||
num_deepstack = len(self.deepstack_encoder_layer_indexes)
|
||||
padded_embeds: List[torch.Tensor] = []
|
||||
deepstack_padded_embeds: List[List[torch.Tensor]] = [
|
||||
[] for _ in range(num_deepstack)
|
||||
]
|
||||
for feat_chunk, len_chunk in zip(
|
||||
padded_feature.split(self.conv_chunksize, dim=0),
|
||||
chunk_lengths.split(self.conv_chunksize, dim=0),
|
||||
):
|
||||
out, deepstack_outs = self._encode_chunk_batch(feat_chunk, len_chunk)
|
||||
if out.shape[1] < t_down_max:
|
||||
out = F.pad(out, (0, 0, 0, t_down_max - out.shape[1]))
|
||||
padded_embeds.append(out)
|
||||
if output_deepstack_hidden_states and num_deepstack > 0:
|
||||
if len(deepstack_outs) != num_deepstack:
|
||||
raise RuntimeError(
|
||||
"Deepstack output count does not match configured layer indexes."
|
||||
)
|
||||
for capture_idx, ds in enumerate(deepstack_outs):
|
||||
if ds.shape[1] < t_down_max:
|
||||
ds = F.pad(ds, (0, 0, 0, t_down_max - ds.shape[1]))
|
||||
deepstack_padded_embeds[capture_idx].append(ds)
|
||||
|
||||
if padded_embeds:
|
||||
padded_embed = torch.cat(padded_embeds, dim=0)
|
||||
else:
|
||||
padded_embed = torch.empty(
|
||||
(0, t_down_max, self.config.output_dim),
|
||||
device=padded_feature.device,
|
||||
)
|
||||
|
||||
valid_tokens = padded_embed[padded_mask_after_cnn] # [N_valid, D]
|
||||
last_hidden_state = valid_tokens.unsqueeze(0) # [1, N_valid, D]
|
||||
|
||||
deepstack_states: Optional[Tuple[torch.Tensor, ...]] = None
|
||||
if output_deepstack_hidden_states and num_deepstack > 0:
|
||||
collected: List[torch.Tensor] = []
|
||||
for chunks_list in deepstack_padded_embeds:
|
||||
if chunks_list:
|
||||
ds = torch.cat(chunks_list, dim=0)
|
||||
collected.append(ds[padded_mask_after_cnn].unsqueeze(0))
|
||||
else:
|
||||
collected.append(
|
||||
torch.empty(
|
||||
(1, 0, self.config.output_dim),
|
||||
device=padded_feature.device,
|
||||
dtype=padded_embed.dtype,
|
||||
)
|
||||
)
|
||||
deepstack_states = tuple(collected)
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=last_hidden_state,
|
||||
hidden_states=deepstack_states,
|
||||
)
|
||||
|
||||
|
||||
class GatedMLP(nn.Module):
|
||||
def __init__(self, input_size, hidden_size, output_size):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(input_size, hidden_size, bias=False)
|
||||
self.up_proj = nn.Linear(input_size, hidden_size, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_size, output_size, bias=False)
|
||||
self.act_fn = nn.SiLU()
|
||||
|
||||
def forward(self, x):
|
||||
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class MossMusicPreTrainedModel(PreTrainedModel):
|
||||
config_class = MossMusicConfig
|
||||
config: MossMusicConfig
|
||||
base_model_prefix = ""
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["Qwen3DecoderLayer"]
|
||||
_skip_keys_device_placement = ["past_key_values"]
|
||||
_supports_flash_attn = True
|
||||
_supports_sdpa = True
|
||||
_supports_flex_attn = True
|
||||
|
||||
_can_compile_fullgraph = False
|
||||
_supports_attention_backend = True
|
||||
_can_record_outputs = {"hidden_states": Qwen3DecoderLayer}
|
||||
|
||||
|
||||
class MossMusicModel(MossMusicPreTrainedModel, GenerationMixin):
|
||||
config_class = MossMusicConfig
|
||||
_tied_weights_keys: List[str] = []
|
||||
|
||||
def __init__(self, config: MossMusicConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.audio_encoder = MossMusicEncoder(config.audio_config)
|
||||
self.language_model = Qwen3Model(config.language_config)
|
||||
|
||||
self.audio_adapter = GatedMLP(
|
||||
input_size=config.audio_config.output_dim,
|
||||
hidden_size=config.adapter_hidden_size,
|
||||
output_size=config.language_config.hidden_size,
|
||||
)
|
||||
|
||||
deepstack_k = len(
|
||||
getattr(config.audio_config, "deepstack_encoder_layer_indexes", []) or []
|
||||
)
|
||||
if config.deepstack_num_inject_layers is not None:
|
||||
deepstack_k = min(deepstack_k, int(config.deepstack_num_inject_layers))
|
||||
self.deepstack_audio_merger_list = nn.ModuleList(
|
||||
[
|
||||
GatedMLP(
|
||||
input_size=config.audio_config.output_dim,
|
||||
hidden_size=config.adapter_hidden_size,
|
||||
output_size=config.language_config.hidden_size,
|
||||
)
|
||||
for _ in range(deepstack_k)
|
||||
]
|
||||
)
|
||||
|
||||
self.vocab_size = config.language_config.vocab_size
|
||||
self.lm_head = nn.Linear(
|
||||
config.language_config.hidden_size, self.vocab_size, bias=False
|
||||
)
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.language_model.get_input_embeddings()
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.language_model.set_input_embeddings(value)
|
||||
|
||||
def get_output_embeddings(self):
|
||||
return self.lm_head
|
||||
|
||||
def set_output_embeddings(self, new_embeddings):
|
||||
self.lm_head = new_embeddings
|
||||
|
||||
def get_audio_features(self, input_features, feature_lens):
|
||||
audio_outputs = self.audio_encoder(
|
||||
input_features=input_features,
|
||||
feature_lens=feature_lens,
|
||||
output_deepstack_hidden_states=True,
|
||||
)
|
||||
deepstack = (
|
||||
list(audio_outputs.hidden_states)
|
||||
if audio_outputs.hidden_states is not None
|
||||
else None
|
||||
)
|
||||
return audio_outputs.last_hidden_state, deepstack
|
||||
|
||||
def _apply_deepstack_to_hidden_states(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
audio_input_mask: torch.Tensor,
|
||||
deepstack_embeds: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
audio_input_mask = audio_input_mask.to(hidden_states.device)
|
||||
deepstack_embeds = deepstack_embeds.to(hidden_states.device, hidden_states.dtype)
|
||||
flat = deepstack_embeds.reshape(-1, deepstack_embeds.shape[-1])
|
||||
hs = hidden_states.clone()
|
||||
hs[audio_input_mask] = hs[audio_input_mask] + flat
|
||||
return hs
|
||||
|
||||
def _register_llm_deepstack_hooks(
|
||||
self,
|
||||
audio_input_mask: torch.Tensor,
|
||||
deepstack_audio_embeds: List[torch.Tensor],
|
||||
):
|
||||
if deepstack_audio_embeds is None or len(deepstack_audio_embeds) == 0:
|
||||
return []
|
||||
|
||||
layers = getattr(self.language_model, "layers", None)
|
||||
if layers is None:
|
||||
raise RuntimeError(
|
||||
"Qwen3Model does not expose `.layers`; cannot register DeepStack hooks."
|
||||
)
|
||||
|
||||
num_inject = len(deepstack_audio_embeds)
|
||||
handles = []
|
||||
|
||||
for layer_idx, layer in enumerate(layers):
|
||||
if layer_idx >= num_inject:
|
||||
break
|
||||
|
||||
def _make_llm_hook(k: int):
|
||||
def _hook(_module, _inputs, _output):
|
||||
if isinstance(_output, (tuple, list)):
|
||||
hs = _output[0]
|
||||
new_hs = self._apply_deepstack_to_hidden_states(
|
||||
hs, audio_input_mask, deepstack_audio_embeds[k]
|
||||
)
|
||||
return (new_hs,) + tuple(_output[1:])
|
||||
else:
|
||||
return self._apply_deepstack_to_hidden_states(
|
||||
_output, audio_input_mask, deepstack_audio_embeds[k]
|
||||
)
|
||||
|
||||
return _hook
|
||||
|
||||
handles.append(layer.register_forward_hook(_make_llm_hook(layer_idx)))
|
||||
|
||||
return handles
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
audio_data: Optional[torch.FloatTensor] = None,
|
||||
audio_data_seqlens: Optional[torch.Tensor] = None,
|
||||
audio_input_mask: Optional[torch.Tensor] = None,
|
||||
**kwargs: Any,
|
||||
) -> CausalLMOutputWithPast:
|
||||
# generate() still passes return_dict; transformers 5 models ignore it
|
||||
kwargs.pop("return_dict", None)
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.get_input_embeddings()(input_ids)
|
||||
|
||||
hook_handles = []
|
||||
if audio_data is not None:
|
||||
if audio_input_mask is None:
|
||||
raise ValueError(
|
||||
"audio_input_mask is required when audio_data is provided."
|
||||
)
|
||||
|
||||
audio_embeds, deepstack = self.get_audio_features(audio_data, audio_data_seqlens)
|
||||
audio_embeds = self.audio_adapter(audio_embeds)
|
||||
|
||||
audio_token_count = int(audio_input_mask.to(torch.int32).sum().item())
|
||||
if audio_token_count != int(audio_embeds.shape[1]):
|
||||
raise ValueError(
|
||||
f"Audio token count mismatch: audio_input_mask has {audio_token_count} audio tokens, "
|
||||
f"but audio_embeds has length {int(audio_embeds.shape[1])}."
|
||||
)
|
||||
|
||||
mask_expanded = audio_input_mask.unsqueeze(-1).expand_as(inputs_embeds)
|
||||
inputs_embeds = inputs_embeds.clone()
|
||||
inputs_embeds.masked_scatter_(mask_expanded, audio_embeds)
|
||||
|
||||
if deepstack is not None and len(self.deepstack_audio_merger_list) > 0:
|
||||
deepstack_audio_embeds = []
|
||||
for i, x in enumerate(deepstack[: len(self.deepstack_audio_merger_list)]):
|
||||
ds = self.deepstack_audio_merger_list[i](x)
|
||||
if int(ds.shape[1]) != audio_token_count:
|
||||
raise ValueError(
|
||||
f"DeepStack audio seq_len mismatch at index {i}: "
|
||||
f"expected {audio_token_count}, got {int(ds.shape[1])}."
|
||||
)
|
||||
deepstack_audio_embeds.append(ds)
|
||||
|
||||
try:
|
||||
hook_handles = self._register_llm_deepstack_hooks(
|
||||
audio_input_mask, deepstack_audio_embeds
|
||||
)
|
||||
except Exception:
|
||||
for h in hook_handles:
|
||||
h.remove()
|
||||
raise
|
||||
|
||||
try:
|
||||
outputs = self.language_model(
|
||||
input_ids=None,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
**kwargs,
|
||||
)
|
||||
finally:
|
||||
for h in hook_handles:
|
||||
h.remove()
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.ignore_index)
|
||||
shift_logits = shift_logits.view(-1, self.config.language_config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
def prepare_inputs_for_generation(
|
||||
self,
|
||||
input_ids,
|
||||
next_sequence_length=None,
|
||||
past_key_values=None,
|
||||
attention_mask=None,
|
||||
inputs_embeds=None,
|
||||
is_first_iteration=False,
|
||||
**kwargs,
|
||||
):
|
||||
# audio is encoded once into the prompt embeds; re-sending it on decode
|
||||
# steps would re-run the encoder against a stale audio_input_mask
|
||||
if not is_first_iteration:
|
||||
for key in ("audio_data", "audio_data_seqlens", "audio_input_mask"):
|
||||
kwargs.pop(key, None)
|
||||
return super().prepare_inputs_for_generation(
|
||||
input_ids,
|
||||
next_sequence_length=next_sequence_length,
|
||||
past_key_values=past_key_values,
|
||||
attention_mask=attention_mask,
|
||||
inputs_embeds=inputs_embeds,
|
||||
is_first_iteration=is_first_iteration,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MossMusicEncoderConfig",
|
||||
"MossMusicConfig",
|
||||
"MossMusicModel",
|
||||
]
|
||||
@@ -0,0 +1,250 @@
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import BatchEncoding, Qwen2Tokenizer
|
||||
|
||||
|
||||
@dataclass
|
||||
class MelConfig:
|
||||
mel_sr: int = 16000
|
||||
mel_dim: int = 128
|
||||
mel_n_fft: int = 400
|
||||
mel_hop_length: int = 160
|
||||
mel_dtype: torch.dtype = torch.bfloat16
|
||||
|
||||
|
||||
class MossMusicProcessor:
|
||||
"""Builds MOSS-Music inputs: Whisper log-mel features plus a Qwen chat
|
||||
prompt whose audio span is one placeholder token per 80ms of audio, with
|
||||
the elapsed seconds written in as digit tokens every 2s (time markers)."""
|
||||
|
||||
_AUDIO_SPAN_RE = re.compile(r"<\|audio_bos\|>(?:<\|AUDIO\|>)+<\|audio_eos\|>")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer,
|
||||
*,
|
||||
mel_config: Optional[MelConfig] = None,
|
||||
enable_time_marker: bool = True,
|
||||
audio_token_id: int = 151654,
|
||||
audio_start_id: int = 151669,
|
||||
audio_end_id: int = 151670,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.audio_token_id = int(audio_token_id)
|
||||
self.audio_start_id = int(audio_start_id)
|
||||
self.audio_end_id = int(audio_end_id)
|
||||
self.enable_time_marker = bool(enable_time_marker)
|
||||
self.config = mel_config or MelConfig()
|
||||
self._whisper_feature_extractor = None
|
||||
|
||||
self._digit_token_ids = {str(d): 15 + d for d in range(10)}
|
||||
self.audio_tokens_per_second = 12.5
|
||||
self.time_marker_every_seconds = 2
|
||||
self.time_marker_every_audio_tokens = int(
|
||||
self.audio_tokens_per_second * self.time_marker_every_seconds
|
||||
)
|
||||
self.model_input_names = [
|
||||
"input_ids",
|
||||
"attention_mask",
|
||||
"audio_data",
|
||||
"audio_data_seqlens",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
|
||||
tokenizer_kwargs = {}
|
||||
for key in ["cache_dir", "revision", "token", "local_files_only"]:
|
||||
if key in kwargs:
|
||||
tokenizer_kwargs[key] = kwargs[key]
|
||||
# the checkpoint's tokenizer_config names Qwen2Tokenizer; loading it
|
||||
# directly skips AutoTokenizer's remote-code prompt for the auto_map
|
||||
tokenizer = Qwen2Tokenizer.from_pretrained(
|
||||
pretrained_model_name_or_path, **tokenizer_kwargs
|
||||
)
|
||||
return cls(
|
||||
tokenizer,
|
||||
mel_config=kwargs.pop("mel_config", None),
|
||||
enable_time_marker=kwargs.pop("enable_time_marker", True),
|
||||
audio_token_id=kwargs.pop("audio_token_id", 151654),
|
||||
audio_start_id=kwargs.pop("audio_start_id", 151669),
|
||||
audio_end_id=kwargs.pop("audio_end_id", 151670),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _conv3_downsample_len(raw_mel_len: int) -> int:
|
||||
def conv_out_len(length: int) -> int:
|
||||
return (length - 1) // 2 + 1
|
||||
|
||||
return conv_out_len(conv_out_len(conv_out_len(int(raw_mel_len))))
|
||||
|
||||
def _get_whisper_feature_extractor(self):
|
||||
if self._whisper_feature_extractor is not None:
|
||||
return self._whisper_feature_extractor
|
||||
|
||||
from transformers.models.whisper.feature_extraction_whisper import (
|
||||
WhisperFeatureExtractor,
|
||||
)
|
||||
|
||||
self._whisper_feature_extractor = WhisperFeatureExtractor(
|
||||
feature_size=int(self.config.mel_dim),
|
||||
sampling_rate=int(self.config.mel_sr),
|
||||
hop_length=int(self.config.mel_hop_length),
|
||||
n_fft=int(self.config.mel_n_fft),
|
||||
)
|
||||
return self._whisper_feature_extractor
|
||||
|
||||
def _extract_mel(self, audio: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio = audio.detach().cpu().numpy()
|
||||
wav_np = np.ascontiguousarray(audio, dtype=np.float32)
|
||||
if wav_np.ndim == 2:
|
||||
wav_np = wav_np[0]
|
||||
# the raw fbank pass has no 30s chunking, so full songs come through whole
|
||||
fe = self._get_whisper_feature_extractor()
|
||||
feats = fe._np_extract_fbank_features(wav_np[None, ...], device="cpu")
|
||||
return torch.from_numpy(feats[0]).to(dtype=self.config.mel_dtype)
|
||||
|
||||
def _get_time_marker_token_ids(self, second: int) -> List[int]:
|
||||
return [self._digit_token_ids[digit] for digit in str(second)]
|
||||
|
||||
def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> List[int]:
|
||||
total_duration_seconds = audio_seq_len / self.audio_tokens_per_second
|
||||
num_full_seconds = int(total_duration_seconds)
|
||||
|
||||
token_ids: List[int] = []
|
||||
audio_tokens_consumed = 0
|
||||
for second in range(
|
||||
self.time_marker_every_seconds,
|
||||
num_full_seconds + 1,
|
||||
self.time_marker_every_seconds,
|
||||
):
|
||||
marker_pos = (
|
||||
second // self.time_marker_every_seconds
|
||||
) * self.time_marker_every_audio_tokens
|
||||
audio_segment_len = marker_pos - audio_tokens_consumed
|
||||
if audio_segment_len > 0:
|
||||
token_ids.extend([self.audio_token_id] * audio_segment_len)
|
||||
audio_tokens_consumed += audio_segment_len
|
||||
token_ids.extend(self._get_time_marker_token_ids(second))
|
||||
|
||||
remaining = audio_seq_len - audio_tokens_consumed
|
||||
if remaining > 0:
|
||||
token_ids.extend([self.audio_token_id] * remaining)
|
||||
return token_ids
|
||||
|
||||
def _build_audio_placeholder_ids(self, num_audio_tokens: int) -> List[int]:
|
||||
if self.enable_time_marker:
|
||||
return self._build_audio_tokens_with_time_markers(num_audio_tokens)
|
||||
return [self.audio_token_id] * num_audio_tokens
|
||||
|
||||
def _build_default_prompt(self, text: str, has_audio: bool) -> str:
|
||||
if has_audio:
|
||||
return (
|
||||
"<|im_start|>system\n"
|
||||
"You are a helpful assistant.<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
"<|audio_bos|><|AUDIO|><|audio_eos|>\n"
|
||||
f"{text}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
return (
|
||||
"<|im_start|>system\n"
|
||||
"You are a helpful assistant.<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{text}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
def _build_input_from_prompt(self, prompt: str, token_lens: List[int]) -> List[int]:
|
||||
spans = list(self._AUDIO_SPAN_RE.finditer(prompt))
|
||||
if len(spans) != len(token_lens):
|
||||
raise ValueError(
|
||||
f"Audio placeholder count mismatch: found {len(spans)} spans in text, "
|
||||
f"but got {len(token_lens)} audio inputs."
|
||||
)
|
||||
|
||||
input_ids: List[int] = []
|
||||
cursor = 0
|
||||
for index, match in enumerate(spans):
|
||||
prefix = prompt[cursor : match.start()]
|
||||
if prefix:
|
||||
input_ids.extend(self.tokenizer.encode(prefix, add_special_tokens=False))
|
||||
|
||||
input_ids.append(self.audio_start_id)
|
||||
input_ids.extend(self._build_audio_placeholder_ids(int(token_lens[index])))
|
||||
input_ids.append(self.audio_end_id)
|
||||
cursor = match.end()
|
||||
|
||||
suffix = prompt[cursor:]
|
||||
if suffix:
|
||||
input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False))
|
||||
return input_ids
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
text: Union[str, Sequence[str]],
|
||||
audios: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None,
|
||||
audio: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None,
|
||||
return_tensors: str = "pt",
|
||||
**kwargs,
|
||||
):
|
||||
if isinstance(text, (list, tuple)):
|
||||
if len(text) != 1:
|
||||
raise ValueError(f"Expected text batch size 1, got {len(text)}")
|
||||
prompt_text = text[0]
|
||||
else:
|
||||
prompt_text = text
|
||||
|
||||
audio_list = audios if audios is not None else audio
|
||||
audio_list = [] if audio_list is None else list(audio_list)
|
||||
|
||||
mels: List[torch.Tensor] = []
|
||||
raw_lengths: List[int] = []
|
||||
token_lens: List[int] = []
|
||||
for one_audio in audio_list:
|
||||
mel = self._extract_mel(one_audio)
|
||||
raw_len = int(mel.shape[-1])
|
||||
mels.append(mel)
|
||||
raw_lengths.append(raw_len)
|
||||
token_lens.append(self._conv3_downsample_len(raw_len))
|
||||
|
||||
if mels:
|
||||
max_length = max(raw_lengths)
|
||||
audio_batch = torch.zeros(
|
||||
(len(mels), self.config.mel_dim, max_length),
|
||||
dtype=self.config.mel_dtype,
|
||||
)
|
||||
for index, mel in enumerate(mels):
|
||||
audio_batch[index, :, : mel.shape[-1]] = mel
|
||||
seqlens_tensor = torch.tensor(raw_lengths, dtype=torch.long)
|
||||
else:
|
||||
audio_batch = None
|
||||
seqlens_tensor = None
|
||||
|
||||
if self._AUDIO_SPAN_RE.search(prompt_text) is None:
|
||||
prompt_text = self._build_default_prompt(prompt_text, has_audio=bool(audio_list))
|
||||
input_ids_list = self._build_input_from_prompt(prompt_text, token_lens)
|
||||
|
||||
input_ids_tensor = torch.tensor([input_ids_list], dtype=torch.long)
|
||||
data = {
|
||||
"input_ids": input_ids_tensor,
|
||||
"attention_mask": torch.ones_like(input_ids_tensor),
|
||||
}
|
||||
if audio_batch is not None:
|
||||
data["audio_data"] = audio_batch
|
||||
data["audio_data_seqlens"] = seqlens_tensor
|
||||
return BatchEncoding(data=data, tensor_type=return_tensors)
|
||||
|
||||
def batch_decode(self, *args, **kwargs):
|
||||
return self.tokenizer.batch_decode(*args, **kwargs)
|
||||
|
||||
def decode(self, *args, **kwargs):
|
||||
return self.tokenizer.decode(*args, **kwargs)
|
||||
|
||||
|
||||
__all__ = ["MelConfig", "MossMusicProcessor"]
|
||||
@@ -283,6 +283,13 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
|
||||
onChange={value => setJobConfig(value, 'config.process[0].caption.extract_vocals_before_transcribe')}
|
||||
/>
|
||||
)}
|
||||
{additionalSections.includes('caption.keep_timestamps') && (
|
||||
<Checkbox
|
||||
label="Keep Lyric Timestamps"
|
||||
checked={jobConfig.config.process[0].caption.keep_timestamps || false}
|
||||
onChange={value => setJobConfig(value, 'config.process[0].caption.keep_timestamps')}
|
||||
/>
|
||||
)}
|
||||
{additionalSections.includes('caption.thinking') && (
|
||||
<Checkbox
|
||||
label="Thinking"
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { CloudLora, GroupedSelectOption, SelectOption } from "@/types";
|
||||
|
||||
type CaptionGroup = 'image' | 'music' | 'video' | 'image/video/sound';
|
||||
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.extract_vocals_before_transcribe' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
|
||||
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.extract_vocals_before_transcribe' | 'caption.keep_timestamps' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
|
||||
|
||||
export interface CaptionOption {
|
||||
name: string;
|
||||
@@ -84,6 +84,13 @@ Part 2: on the next line write [Lyrics] and then the complete lyrics transcribed
|
||||
|
||||
Transcribe only what is actually sung. Be decisive. No preamble, no explanations, no markdown - output only the tag line and the lyrics block.`;
|
||||
|
||||
// MOSS-Music description prompts. The lyric transcription prompt is fixed in the
|
||||
// captioner (it asks for timestamps as loop anchors and strips them afterwards).
|
||||
// Framing the description as a generator prompt is what keeps the model to one
|
||||
// plain paragraph instead of a sectioned essay with timings and chord names.
|
||||
const mossMusicCaptionPrompt = "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.";
|
||||
const mossMusicTagsPrompt = "Describe this music as a single line of comma-separated style tags: genre, mood, vocal type and delivery, lead instruments, production style, tempo feel, era. Lowercase tags only, no sentences.";
|
||||
|
||||
// Editable ADDITIONAL INSTRUCTIONS block injected into the Ideogram system prompt.
|
||||
// Users can tweak this for dataset-specific guidance without altering the fixed
|
||||
// output contract, element/background rules, or bbox format.
|
||||
@@ -114,6 +121,32 @@ export const captionerTypes: CaptionOption[] = [
|
||||
'caption.extract_vocals_before_transcribe',
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'MossMusicCaptioner',
|
||||
label: 'MOSS-Music',
|
||||
group: 'music',
|
||||
defaults: {
|
||||
'config.process[0].caption.model_name_or_path': ['OpenMOSS-Team/MOSS-Music-8B-Instruct', defaultNameOrPath],
|
||||
'config.process[0].caption.extensions': [extensionsAudio, defaultExtensions],
|
||||
'config.process[0].caption.caption_format': ['ace_step', undefined],
|
||||
'config.process[0].caption.caption_prompt': [mossMusicCaptionPrompt, undefined],
|
||||
'config.process[0].caption.keep_timestamps': [false, undefined],
|
||||
'config.process[0].caption.compile': [true, false],
|
||||
},
|
||||
name_or_path_options: [
|
||||
{ value: 'OpenMOSS-Team/MOSS-Music-8B-Instruct', label: 'OpenMOSS-Team/MOSS-Music-8B-Instruct' },
|
||||
],
|
||||
captionPrompts: {
|
||||
'Description (ACE-Step)': mossMusicCaptionPrompt,
|
||||
'Style tags (YuE2)': mossMusicTagsPrompt,
|
||||
},
|
||||
additionalSections: [
|
||||
'caption.fixed_caption',
|
||||
'caption.caption_format',
|
||||
'caption.keep_timestamps',
|
||||
'caption.caption_prompt',
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Qwen3VLCaptioner',
|
||||
label: 'Qwen3-VL',
|
||||
|
||||
@@ -339,6 +339,7 @@ export interface CaptionProcessConfig {
|
||||
fixed_caption?: string;
|
||||
caption_format?: string;
|
||||
extract_vocals_before_transcribe?: boolean;
|
||||
keep_timestamps?: boolean;
|
||||
caption_extension?: string;
|
||||
thinking?: boolean;
|
||||
batch_size?: number;
|
||||
|
||||
Reference in new issue
Block a user