muscriptor_cli.py

"""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()