"""MuScriptor inference using llama.cpp as the transformer backend. Python computes conditioning embeddings (mel spectrogram + class embeddings), then passes them to llama-server via the /completion_with_embd endpoint which accepts continuous float vectors as a prefix before text token generation. Usage: # Terminal 1: start the server (fp32, 8k ctx, 8 parallel slots) llama-server -m model_fp32_8k.gguf --port 8081 -ngl 99 --parallel 8 # Terminal 2: run transcription python muscriptor_cli.py audio.mp3 -o out.mid """ import os os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" import argparse import io import math import sys import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from pathlib import Path import numpy as np import requests import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange from safetensors.torch import load_file _CONDITIONING_KEYS = [ "condition_provider.conditioners.self_wav.output_proj.weight", "condition_provider.conditioners.self_wav.output_proj.bias", "condition_provider.conditioners.self_wav.mel_spec_transform.mel_scale.fb", "condition_provider.conditioners.self_wav.mel_spec_transform.spectrogram.window", "condition_provider.conditioners.instrument_group.embed.weight", "condition_provider.conditioners.dataset_name.embed.weight", ] def _load_conditioning_state(weights_path: str) -> dict: state = load_file(weights_path) missing = [k for k in _CONDITIONING_KEYS if k not in state] if missing: raise KeyError(f"Missing conditioning keys in {weights_path}: {missing}") return {k: state[k] for k in _CONDITIONING_KEYS} _SAMPLE_RATE = 16000 _SEGMENT_DURATION = 5.0 _FRAME_RATE = 100 _MAX_GEN_LEN = 2000 _MODEL_CARD = 1395 @dataclass class NoteStartEvent: pitch: int start_time: float index: int instrument: str @dataclass class NoteEndEvent: end_time: float start_event: NoteStartEvent @property def start_event_index(self) -> int: return self.start_event.index @dataclass class ProgressEvent: completed: int total: int @dataclass class ChunkBoundary: seek_time: float next_seek_time: float | None @dataclass class Event: type: str value: int _DRUM_INSTRUMENT = "drums" _MINIMUM_NOTE_DURATION_SEC = 0.01 _SPECIAL_TOKENS = ("PAD", "EOS", "UNK") @dataclass class _EventRange: type: str min_value: int max_value: int def _build_event_vocab(max_shift_steps: int = 1001) -> list[Event]: ranges = ( [_EventRange(t, 0, 0) for t in _SPECIAL_TOKENS] + [_EventRange("shift", 0, max_shift_steps - 1)] + [ _EventRange("pitch", 0, 127), _EventRange("velocity", 0, 1), _EventRange("tie", 0, 0), _EventRange("program", 0, 129), _EventRange("drum", 0, 127), ] ) vocab = [] for er in ranges: for value in range(er.min_value, er.max_value + 1): vocab.append(Event(type=er.type, value=value)) return vocab def _get_instrument_program_map(): from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES, get_group_program_map gpm = get_group_program_map("FULL", misc_programs="SINGLETON_GROUPS", is_mt3=True) name_to_prog = {} for name, gid in MT3_FULL_PLUS_GROUP_NAMES.items(): if name == "drums": continue if gid in gpm and gpm[gid]: name_to_prog[name] = gpm[gid][0] return name_to_prog def _instrument_for_program(prog: int) -> str: from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES, get_group_program_map gpm = get_group_program_map("FULL", misc_programs="SINGLETON_GROUPS", is_mt3=True) for name, gid in MT3_FULL_PLUS_GROUP_NAMES.items(): if name == "drums": continue if gid in gpm and gpm[gid] and gpm[gid][0] == prog: return name.replace("_", " ") return f"program_{prog}" def _program_for_instrument(instrument: str) -> int: from muscriptor.tokenizer.notes import DRUM_PROGRAM if instrument == _DRUM_INSTRUMENT: return DRUM_PROGRAM for name, prog in _get_instrument_program_map().items(): if name.replace("_", " ") == instrument: return prog return 0 def _load_audio(path: str, target_sr: int = _SAMPLE_RATE) -> torch.Tensor: import torchaudio wav, sr = torchaudio.load(path) wav = wav.float() if wav.dim() == 1: wav = wav.unsqueeze(0) if wav.shape[0] > 1: wav = wav.mean(0, keepdim=True) if sr != target_sr: wav = torchaudio.functional.resample(wav, sr, target_sr) return wav class _MelConditioner(nn.Module): def __init__(self, dim: int): super().__init__() self.dim = dim from muscriptor.modules.mel_spectrogram import _MelSpectrogram self.mel_transform = _MelSpectrogram( sample_rate=_SAMPLE_RATE, n_fft=2048, hop_length=_SAMPLE_RATE // _FRAME_RATE, n_mels=512, power=1.0, ) self.output_proj = nn.Linear(512, dim, bias=True) def forward(self, wav: torch.Tensor, length: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: mel = self.mel_transform(wav) mel = rearrange(mel, "b 1 d t -> b t d") mel = torch.log(mel + 1e-6) embeds = self.output_proj(mel) lengths = length / (_SAMPLE_RATE // _FRAME_RATE) from muscriptor.utils.sampling import length_to_mask mask = length_to_mask(lengths, max_len=embeds.shape[1]).int() embeds = embeds * mask.float().unsqueeze(-1) return embeds, mask class _ClassConditioner(nn.Module): def __init__(self, num_classes: int, dim: int): super().__init__() self.embed = nn.Embedding(num_classes + 1, dim) def forward(self, text_val: str | None, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: if text_val is not None: ints = list(map(int, text_val.split())) else: ints = [-1] tokens = 2 + torch.LongTensor([ints]).to(device) embeds = self.embed(tokens) mask = torch.ones_like(embeds[..., 0]) return embeds, mask class MuScriptorLlama: def __init__(self, weights_path: str, server_url: str = "http://localhost:8081", device: str = "cuda", parallel: int = 8): self.device = torch.device(device) self.server_url = server_url.rstrip("/") self.parallel = parallel self.vocab = _build_event_vocab() self.card = _MODEL_CARD self.eos_id = _SPECIAL_TOKENS.index("EOS") self.initial_token_id = self.card state = _load_conditioning_state(weights_path) self.mel_cond = _MelConditioner(dim=1536) self.mel_cond.output_proj.weight.data = state[ "condition_provider.conditioners.self_wav.output_proj.weight" ] self.mel_cond.output_proj.bias.data = state[ "condition_provider.conditioners.self_wav.output_proj.bias" ] self.mel_cond.mel_transform.mel_scale.fb.data = state[ "condition_provider.conditioners.self_wav.mel_spec_transform.mel_scale.fb" ] self.mel_cond.mel_transform.spectrogram.window.data = state[ "condition_provider.conditioners.self_wav.mel_spec_transform.spectrogram.window" ] self.mel_cond = self.mel_cond.to(self.device).float() self.inst_cond = _ClassConditioner(num_classes=1000, dim=1536) self.inst_cond.embed.weight.data = state[ "condition_provider.conditioners.instrument_group.embed.weight" ] self.inst_cond = self.inst_cond.to(self.device).float() self.ds_cond = _ClassConditioner(num_classes=4, dim=1536) self.ds_cond.embed.weight.data = state[ "condition_provider.conditioners.dataset_name.embed.weight" ] self.ds_cond = self.ds_cond.to(self.device).float() self.inst_program_map = _get_instrument_program_map() def _build_condition_embeddings(self, wav: torch.Tensor, instrument_group: str | None = None) -> list[list[float]]: wav = wav.to(self.device) T = wav.shape[-1] wav_3d = wav.unsqueeze(0) length = torch.tensor([T], device=self.device) mel_emb, mel_mask = self.mel_cond(wav_3d, length) inst_emb, inst_mask = self.inst_cond(instrument_group, self.device) ds_emb, ds_mask = self.ds_cond(None, self.device) mel_emb = mel_emb[0].detach().float().cpu().numpy() inst_emb = inst_emb[0].detach().float().cpu().numpy() ds_emb = ds_emb[0].detach().float().cpu().numpy() embeddings = [] for i in range(mel_emb.shape[0]): embeddings.append(mel_emb[i].tolist()) for i in range(ds_emb.shape[0]): embeddings.append(ds_emb[i].tolist()) for i in range(inst_emb.shape[0]): embeddings.append(inst_emb[i].tolist()) return embeddings def _infer_chunk(self, chunk_index: int, embeddings: list[list[float]], temperature: float) -> tuple[int, list[int]]: body = { "prompt": [self.initial_token_id], "input_embd": embeddings, "n_predict": _MAX_GEN_LEN, "temperature": temperature, "top_k": -1, "top_p": 0.0, "stream": False, "cache_prompt": False, "return_tokens": True, "logit_bias": [ [token_id, False] for token_id in range(len(self.vocab), self.card + 1) ], } resp = requests.post( f"{self.server_url}/completion_with_embd", json=body, timeout=600, ) if resp.status_code != 200: raise RuntimeError( f"llama-server returned {resp.status_code} for chunk {chunk_index}: {resp.text[:500]}" ) data = resp.json() tokens = data.get("tokens", []) eos_pos = -1 for i, tok in enumerate(tokens): if tok == self.eos_id: eos_pos = i break if eos_pos >= 0: tokens = tokens[:eos_pos] return chunk_index, tokens def transcribe(self, audio_path: str, instrument_group: str | None = None, use_sampling: bool = False, temperature: float = 1.0, batch_size: int = 4): wav = _load_audio(audio_path) wav = wav.to(self.device) total_samples = wav.shape[-1] total_duration = total_samples / _SAMPLE_RATE segment_samples = int(_SEGMENT_DURATION * _SAMPLE_RATE) num_chunks = math.ceil(total_samples / segment_samples) print(f"[muscriptor-llama] audio: {total_duration:.1f}s -> {num_chunks} chunk(s)", file=sys.stderr) yield ProgressEvent(completed=0, total=num_chunks) # Build all conditions upfront chunk_embeddings = [] seek_times = [] for i in range(num_chunks): start = i * segment_samples chunk = wav[:, start:start + segment_samples] if chunk.shape[-1] < segment_samples: chunk = F.pad(chunk, (0, segment_samples - chunk.shape[-1])) chunk_embeddings.append(self._build_condition_embeddings(chunk, instrument_group)) seek_times.append(i * _SEGMENT_DURATION) # Parallel inference in waves of self.parallel results: dict[int, list[int]] = {} n_completed = 0 for wave_start in range(0, num_chunks, self.parallel): wave_end = min(wave_start + self.parallel, num_chunks) wave_indices = list(range(wave_start, wave_end)) with ThreadPoolExecutor(max_workers=len(wave_indices)) as executor: futures = { executor.submit( self._infer_chunk, ci, chunk_embeddings[ci], temperature ): ci for ci in wave_indices } for future in as_completed(futures): ci, tokens = future.result() results[ci] = tokens n_completed += 1 print(f"[muscriptor-llama] chunk {ci} ({seek_times[ci]:.0f}s): {len(tokens)} tokens", file=sys.stderr) yield ProgressEvent(completed=n_completed, total=num_chunks) # Yield in chunk order with boundaries for i in range(num_chunks): yield ChunkBoundary( seek_times[i], seek_times[i + 1] if i + 1 < num_chunks else None, ) yield from results[i] def transcribe_to_midi(self, audio_path: str, **kwargs) -> bytes: raw_stream = self.transcribe(audio_path, **kwargs) events = _decode_model_tokens( raw_stream, self.vocab, _instrument_for_program ) return _events_to_midi_bytes(events) def _decode_model_tokens(stream, vocab, instrument_for_program, frame_rate: int = _FRAME_RATE): open_notes = {} next_index = 0 def mint(pitch, start_time, instrument): nonlocal next_index ev = NoteStartEvent(pitch=pitch, start_time=start_time, index=next_index, instrument=instrument) next_index += 1 return ev seek_time = 0.0 next_seek_time = None start_tick = 0 tick_state = 0 program_state = None velocity_state = None in_prologue = True skip_rest = False tie_set = set() chunk_started = False for item in stream: if isinstance(item, ProgressEvent): yield item continue if isinstance(item, ChunkBoundary): if chunk_started and in_prologue: for key in list(open_notes): yield NoteEndEvent(end_time=seek_time, start_event=open_notes.pop(key)) seek_time = item.seek_time next_seek_time = item.next_seek_time start_tick = round(seek_time * frame_rate) tick_state = start_tick program_state = None velocity_state = None in_prologue = True skip_rest = False tie_set = set() chunk_started = True continue if not isinstance(item, int): continue if item >= len(vocab): continue event = vocab[item] etype = event.type if in_prologue: if etype == "tie": in_prologue = False velocity_state = None for key in list(open_notes): if key not in tie_set: yield NoteEndEvent(end_time=seek_time, start_event=open_notes.pop(key)) elif etype == "shift": in_prologue = False skip_rest = True for key in list(open_notes): yield NoteEndEvent(end_time=seek_time, start_event=open_notes.pop(key)) elif etype == "program": program_state = event.value elif etype == "pitch" and program_state is not None: tie_set.add((program_state, event.value)) continue if skip_rest: continue if etype == "shift": if event.value > 0: tick_state = start_tick + event.value elif etype == "program": program_state = event.value elif etype == "velocity": velocity_state = event.value elif etype == "drum": time_ = tick_state / frame_rate if next_seek_time is None or time_ < next_seek_time: start = mint(event.value, time_, _DRUM_INSTRUMENT) yield start yield NoteEndEvent(end_time=time_ + _MINIMUM_NOTE_DURATION_SEC, start_event=start) elif etype == "pitch": if program_state is None or velocity_state is None: continue time_ = tick_state / frame_rate if next_seek_time is not None and time_ >= next_seek_time: continue key = (program_state, event.value) if key in open_notes: yield NoteEndEvent(end_time=time_, start_event=open_notes.pop(key)) if velocity_state > 0: start = mint(event.value, time_, instrument_for_program(program_state)) open_notes[key] = start yield start if chunk_started and in_prologue: for key in list(open_notes): yield NoteEndEvent(end_time=seek_time, start_event=open_notes.pop(key)) else: for ev in list(open_notes.values()): yield NoteEndEvent(end_time=ev.start_time + _MINIMUM_NOTE_DURATION_SEC, start_event=ev) open_notes.clear() def _events_to_midi_bytes(events): from muscriptor.tokenizer.notes import Note, validate_notes, trim_overlapping_notes, DRUM_PROGRAM from muscriptor.utils.midi import notes_to_midi notes = [] open_notes = {} program_names = {} for ev in events: if isinstance(ev, ProgressEvent): continue if isinstance(ev, NoteStartEvent): is_drum = ev.instrument == _DRUM_INSTRUMENT program = _program_for_instrument(ev.instrument) program_names[program] = ev.instrument.replace("_", " ") note = Note(is_drum=is_drum, program=program, onset=ev.start_time, offset=ev.start_time, pitch=ev.pitch) open_notes[ev.index] = note elif isinstance(ev, NoteEndEvent): if ev.start_event_index not in open_notes: continue note = open_notes.pop(ev.start_event_index) note.offset = ev.end_time notes.append(note) notes = validate_notes(notes, fix=True) notes = trim_overlapping_notes(notes, sort=True) midi = notes_to_midi(notes, program_names=program_names) buf = io.BytesIO() midi.save(file=buf) return buf.getvalue() def main(): parser = argparse.ArgumentParser(description="MuScriptor inference via llama.cpp") parser.add_argument("audio", help="Audio file path") parser.add_argument("-o", "--output", default="out.mid", help="Output MIDI path") parser.add_argument("--weights", default="./conditioning_weights.safetensors", help="Conditioning weights (extracted or full safetensors)") parser.add_argument("--server", default="http://localhost:8081", help="llama-server URL") parser.add_argument("--device", default="cuda", help="Device for conditioning computation") parser.add_argument("--parallel", type=int, default=8, help="Number of parallel inference requests") parser.add_argument("--temperature", type=float, default=1.0, help="Sampling temperature") args = parser.parse_args() t0 = time.time() model = MuScriptorLlama(args.weights, server_url=args.server, device=args.device, parallel=args.parallel) print(f"[muscriptor-llama] conditioning model loaded in {time.time() - t0:.1f}s", file=sys.stderr) midi_bytes = model.transcribe_to_midi(args.audio, temperature=args.temperature) Path(args.output).write_bytes(midi_bytes) print(f"[muscriptor-llama] wrote {args.output} ({len(midi_bytes)} bytes)", file=sys.stderr) if __name__ == "__main__": main()