| """Small native API; model objects remain editable and hookable PyTorch modules.""" |
|
|
| from contextlib import nullcontext |
| from dataclasses import dataclass |
| import json |
| from pathlib import Path |
| from typing import Sequence |
|
|
| import torch |
|
|
| from .audio import Request, load_audio, normalize_prompt |
|
|
|
|
| @dataclass |
| class AudioSlot: |
| """One ordered source; token_positions [U_i] index the padded input row.""" |
|
|
| sample_index: int |
| audio_index: int |
| token_positions: torch.Tensor |
| samples_16k: int |
| windows: int | None |
|
|
|
|
| @dataclass |
| class PreparedBatch: |
| """Ready-to-forward tensors, validated user prompt strings and source slots. |
| |
| inputs contains input_ids/attention_mask [B,L]. MF adds mel [W,128,3000] |
| and frame mask [W,3000]. MOSS adds mel [A,128,F_max], lengths [A], |
| audio_input_mask [B,L], audio_group_sizes [B], audio_token_lens [A]. |
| All tensors live on the model device; floating features use model dtype. |
| slots identify valid audio-token positions, never padding or timestamp digits. |
| """ |
|
|
| inputs: dict[str, torch.Tensor] |
| slots: list[AudioSlot] |
| prompts: list[str] |
| windows_per_source: list[int] |
|
|
|
|
| @dataclass |
| class AudioEmbedding: |
| """One source's unpooled token features, with padding already removed. |
| |
| encoder [U_i,1280], projected [U_i,D], token_positions [U_i]. D=3584 |
| for MF, 4096 for MOSS. MOSS has three deepstack_encoder [U_i,1280] and |
| deepstack_projected [U_i,4096] tensors; MF has empty DeepStack tuples. |
| Encoder != projected != language-contextual representations. |
| """ |
|
|
| sample_index: int |
| audio_index: int |
| token_positions: torch.Tensor |
| encoder: torch.Tensor |
| projected: torch.Tensor |
| deepstack_encoder: tuple[torch.Tensor, ...] |
| deepstack_projected: tuple[torch.Tensor, ...] |
|
|
|
|
| class Session: |
| """Single-device inference session; not a training or concurrent-serving API. |
| |
| model is the actual torch module (not a remote text client). Users may edit |
| its forward or attach hooks. Default methods disable gradients; grad=True |
| enables ordinary autograd for feature/forward research without checkpointed |
| training. Do not invoke the same MOSS instance concurrently: DeepStack hooks |
| are scoped to one forward and removed on success or failure. |
| """ |
|
|
| def __init__(self, model, processor, family: str): |
| if family not in {"mf2601", "moss_music"}: |
| raise ValueError("Require mf2601 or moss_music") |
| self.model = model |
| self.processor = processor |
| self.family = family |
| if family == "mf2601": |
| from . import mf2601 as backend |
| self.context_length = int(model.config.text_config.max_position_embeddings) |
| else: |
| from . import moss_music as backend |
| self.context_length = int(model.config.language_config.max_position_embeddings) |
| self.backend = backend |
|
|
| def prepare(self, requests: Sequence[Request]) -> PreparedBatch: |
| """Decode/expand B single-turn requests, return tensors and ordered slots. |
| |
| Reject marker mismatches, empty/nonfinite waveforms, unsupported native |
| tokens and prompts leaving no output context. No truncation of tokens, |
| source concatenation, deduplication or sorting is performed. |
| """ |
| if not requests: |
| raise ValueError("At least one Request is required") |
| prompts = [normalize_prompt(r.prompt, len(r.audios), label_single_audio=self.family != "mf2601") for r in requests] |
| groups = [[load_audio(a, crop_seconds=1200 if self.family == "mf2601" else None) for a in r.audios] for r in requests] |
| inputs, lens, windows = self.backend.prepare(prompts, groups, self.processor) |
| if inputs["input_ids"].shape[1] >= self.context_length: |
| raise ValueError("Expanded multimodal prompt leaves no output context") |
| device, dtype = self.model.device, self.model.dtype |
| inputs = {k: v.to(device=device, dtype=dtype if v.is_floating_point() else v.dtype) for k, v in inputs.items()} |
| token_id = self.model.config.audio_token_id if self.family == "mf2601" else self.processor.audio_token_id |
| slots = [] |
| source_index = 0 |
| for sample_index, group in enumerate(groups): |
| positions = torch.where((inputs["input_ids"][sample_index] == token_id) & inputs["attention_mask"][sample_index].bool())[0] |
| offset = 0 |
| for audio_index, audio in enumerate(group): |
| count = lens[source_index] |
| source_positions = positions[offset:offset + count] |
| if len(source_positions) != count or count <= 0: |
| raise ValueError("Audio marker expansion differs from source token lengths") |
| slots.append(AudioSlot(sample_index, audio_index, source_positions, len(audio), windows[source_index] if windows else None)) |
| offset += count |
| source_index += 1 |
| if offset != len(positions): |
| raise ValueError("Unexpected extra audio tokens in expanded input") |
| if source_index != len(lens): |
| raise ValueError("Some source lengths were not consumed") |
| return PreparedBatch(inputs, slots, prompts, windows) |
|
|
| def _autocast(self): |
| if self.model.device.type == "cuda" and self.model.dtype in {torch.bfloat16, torch.float16}: |
| return torch.autocast("cuda", dtype=self.model.dtype) |
| return nullcontext() |
|
|
| def forward(self, batch: PreparedBatch, *, output_hidden_states: bool = False, grad: bool = False): |
| """Return native ModelOutput: logits [B,1,V], optionally hidden [B,L,D]. |
| |
| Runs prefill only with cache disabled. Hidden-state index 0 is embedded |
| input; 1..K correspond to decoder-layer boundaries (last includes final |
| normalization). Audio positions are language-contextual features, not raw |
| audio. Full hidden states can be expensive for long prompts. |
| """ |
| with torch.set_grad_enabled(grad), self._autocast(): |
| return self.model(**batch.inputs, use_cache=False, return_dict=True, |
| output_hidden_states=output_hidden_states, logits_to_keep=1) |
|
|
| def extract_audio_embeddings(self, batch: PreparedBatch, *, grad: bool = False) -> list[AudioEmbedding]: |
| """Encode sources without running the LLM; return one AudioEmbedding/slot. |
| |
| See AudioEmbedding for exact shapes. No mean pooling/detach is applied; |
| grad=True preserves the autograd graph through encoder and projectors. |
| """ |
| if not batch.slots: |
| return [] |
| with torch.set_grad_enabled(grad), self._autocast(): |
| raw, projected, deep_raw, deep_projected = self.backend.audio_features(self.model, batch.inputs, batch.windows_per_source) |
| total = sum(len(s.token_positions) for s in batch.slots) |
| if raw.shape[0] != total or projected.shape[0] != total: |
| raise ValueError("Feature lengths differ from expanded source token counts") |
| if any(s.shape[0] != total for s in (*deep_raw, *deep_projected)): |
| raise ValueError("DeepStack token count differs from final encoder stream") |
| results = [] |
| offset = 0 |
| for slot in batch.slots: |
| stop = offset + len(slot.token_positions) |
| results.append(AudioEmbedding(slot.sample_index, slot.audio_index, slot.token_positions, |
| raw[offset:stop], projected[offset:stop], tuple(s[offset:stop] for s in deep_raw), |
| tuple(s[offset:stop] for s in deep_projected))) |
| offset = stop |
| return results |
|
|
| def extract_contextual_embeddings(self, batch: PreparedBatch, layers: Sequence[int] = (-1,), *, grad: bool = False): |
| """Return list[dict] per source: requested layer -> tensor [U_i,D]. |
| |
| Uses forward output_hidden_states, including MOSS DeepStack injection. |
| token_positions refer to the padded input row. Negative layer indexes |
| follow Python indexing; invalid indexes fail, never silently clamp. |
| """ |
| if not layers: |
| raise ValueError("At least one hidden-state layer is required") |
| outputs = self.forward(batch, output_hidden_states=True, grad=grad) |
| if outputs.hidden_states is None: |
| raise RuntimeError("Native model did not return language hidden states") |
| return [{"sample_index": s.sample_index, "audio_index": s.audio_index, |
| "token_positions": s.token_positions, |
| "layers": {i: outputs.hidden_states[i][s.sample_index, s.token_positions] for i in layers}} |
| for s in batch.slots] |
|
|
| def generate(self, requests: Sequence[Request], *, profile: str = "greedy", max_new_tokens: int = 16384) -> list[str]: |
| """Return one generated string/request using KV cache; no retry/rewrite. |
| |
| greedy decodes deterministically; sampled uses the base model's sampling |
| defaults. In a batch, the token budget is capped by the longest prompt; |
| use batch size 1 near the context limit. |
| """ |
| if type(max_new_tokens) is not int or max_new_tokens < 1: |
| raise ValueError("max_new_tokens must be a positive integer") |
| if profile not in {"greedy", "sampled"}: |
| raise ValueError("Unknown generation profile") |
| batch = self.prepare(requests) |
| width = batch.inputs["input_ids"].shape[1] |
| budget = min(max_new_tokens, self.context_length - width) |
| from transformers import GenerationConfig |
| original = self.model.generation_config |
| sampled = ({"temperature": 0.7 if self.family == "mf2601" else 1.0, |
| "top_p": 0.9 if self.family == "mf2601" else 0.8, "top_k": 50} |
| if profile == "sampled" else {}) |
| config = GenerationConfig( |
| do_sample=profile == "sampled", max_new_tokens=budget, |
| repetition_penalty=1.0, |
| num_beams=1, num_return_sequences=1, use_cache=True, |
| bos_token_id=original.bos_token_id, eos_token_id=original.eos_token_id, |
| pad_token_id=self.processor.tokenizer.pad_token_id, |
| **sampled, |
| ) |
| |
| devices = [self.model.device] if self.model.device.type == "cuda" else [] |
| with torch.random.fork_rng(devices=devices), torch.no_grad(), self._autocast(): |
| torch.manual_seed(3407) |
| output = self.model.generate(**batch.inputs, generation_config=config, logits_to_keep=1) |
| return [s.strip() for s in self.processor.batch_decode(output[:, width:], skip_special_tokens=True)] |
|
|
|
|
| def load_model(model_path: str | Path, *, device: str = "cuda:0", dtype: torch.dtype = torch.bfloat16, |
| attention: str = "sdpa") -> Session: |
| """Load a downloaded model folder (e.g. from snapshot_download). |
| |
| The returned Session exposes .model and .processor. No remote code or global |
| patches are used. Missing, unexpected or mismatched tensors fail explicitly. |
| """ |
| path = Path(model_path).expanduser() |
| if not path.is_dir(): |
| raise FileNotFoundError("Pass a local model directory") |
| if attention not in {"sdpa", "eager"}: |
| raise ValueError("Supported attention: sdpa or eager") |
| cfg = json.loads((path / "config.json").read_text()) |
| kwargs = dict(local_files_only=True, dtype=dtype, attn_implementation=attention, output_loading_info=True) |
| if cfg.get("architectures") == ["MusicFlamingoForConditionalGeneration"]: |
| from transformers import AutoProcessor, MusicFlamingoForConditionalGeneration |
| processor = AutoProcessor.from_pretrained(path, local_files_only=True) |
| model, loading = MusicFlamingoForConditionalGeneration.from_pretrained(path, **kwargs) |
| from .mf2601 import patch_zero_token_window_timestamps |
| patch_zero_token_window_timestamps(model) |
| family = "mf2601" |
| elif cfg.get("architectures") == ["MossMusicModel"]: |
| from transformers import Qwen2Tokenizer |
| from .vendor.configuration_moss_music import MossMusicConfig |
| from .vendor.modeling_moss_music import MossMusicModel |
| from .vendor.processing_moss_music import MelConfig, MossMusicProcessor |
| from .moss_patches import (patch_moss_music_batch_forward, patch_moss_music_audio_encoder_official_chunks, |
| patch_whisper_encoder_layer_outputs, reset_moss_music_sinusoid_buffer) |
|
|
| class AttrDict(dict): |
| def __getattr__(self, key): |
| try: |
| return self[key] |
| except KeyError as exc: |
| raise AttributeError(key) from exc |
|
|
| markers = {cfg.get("audio_encoder_chunking"), cfg["audio_config"].get("audio_encoder_chunking")} |
| if markers != {"official_chunked_400"}: |
| raise ValueError("Expected a checkpoint trained with official 400-frame audio-encoder chunking") |
| config = MossMusicConfig.from_dict(cfg) |
| config.audio_config = AttrDict(config.audio_config) |
| |
| |
| config.audio_config["_attn_implementation"] = "eager" |
| config.language_config._attn_implementation = attention |
| config._attn_implementation = attention |
| tokenizer = Qwen2Tokenizer.from_pretrained(path, local_files_only=True) |
| processor = MossMusicProcessor(tokenizer, mel_config=MelConfig(), enable_time_marker=True) |
| model, loading = MossMusicModel.from_pretrained(path, config=config, **kwargs) |
| reset_moss_music_sinusoid_buffer(model) |
| patch_whisper_encoder_layer_outputs(model) |
| patch_moss_music_audio_encoder_official_chunks(model) |
| patch_moss_music_batch_forward(model) |
| family = "moss_music" |
| else: |
| raise ValueError("Unknown architecture; refuse to guess the model family") |
| failures = {k: loading.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") if loading.get(k)} |
| if failures: |
| raise RuntimeError(f"Incomplete checkpoint loading: {failures}") |
| model.to(device).eval() |
| |
| |
| if family == "moss_music": |
| reset_moss_music_sinusoid_buffer(model) |
| model.config.use_cache = True |
| return Session(model, processor, family) |
|
|
