apollo_m.py

"""Apollo-M: standalone inference. Load the final model from the Hub and turn a text prompt into MIDI.
    from apollo_m import load
    m = load(); m.make("8 bars in D minor at 90 bpm, sad piano melody", n=3, seed=1)
"""
import os, re, json, math, time
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
import pretty_midi
from pathlib import Path

REPO, SEQ_LEN = "viknesh100/Apollo-M", 1024

# ---------------- vocab (identical to training) ----------------
NOTE_NAMES = ["C","C#","D","D#","E","F","F#","G","G#","A","A#","B"]
TS_LIST = [(4,4),(3,4),(2,4),(6,8),(12,8)]
STEPS_PER_BAR = {ts: ts[0]*16//ts[1] for ts in TS_LIST}
BASE_FIELDS = ["TEMPO","TS","KEY","BARS","DENS","REG"]
NAMES = ["PAD","BOS","EOS","BAR"]
NAMES += [f"POS_{i}" for i in range(32)]
NAMES += [f"PITCH_{i}" for i in range(128)]
NAMES += [f"DUR_{i}" for i in range(1,33)]
NAMES += [f"TEMPO_{i}" for i in range(24)]
NAMES += [f"TS_{a}_{b}" for a,b in TS_LIST]
NAMES += [f"KEY_{i}" for i in range(24)]
NAMES += [f"BARS_{i}" for i in range(1,17)]
NAMES += [f"DENS_{i}" for i in range(8)]
NAMES += [f"REG_{i}" for i in range(8)]
NAMES += [f"UNK_{f}" for f in BASE_FIELDS]
NAMES += [f"EXTRA_{i}" for i in range(64)]
while len(NAMES) % 8: NAMES.append(f"PADV_{len(NAMES)}")
T = {n:i for i,n in enumerate(NAMES)}; VOCAB = len(NAMES)
BOS, EOS, BAR_ID = T["BOS"], T["EOS"], T["BAR"]
POS0, PITCH0, DUR1 = T["POS_0"], T["PITCH_0"], T["DUR_1"]
FIELDS = BASE_FIELDS + ["GENRE","MOOD","INSTR"]
FIELD_BASE = {"TEMPO":T["TEMPO_0"],"TS":T["TS_4_4"],"KEY":T["KEY_0"],"BARS":T["BARS_1"],"DENS":T["DENS_0"],"REG":T["REG_0"],
              "GENRE":T["EXTRA_0"],"MOOD":T["EXTRA_24"],"INSTR":T["EXTRA_40"]}
FIELD_N = {"TEMPO":24,"TS":len(TS_LIST),"KEY":24,"BARS":16,"DENS":8,"REG":8,"GENRE":24,"MOOD":16,"INSTR":16}
UNK_IDS = [T[f"UNK_{f}"] for f in BASE_FIELDS] + [T["EXTRA_56"], T["EXTRA_57"], T["EXTRA_58"]]
HEADER_LEN = 1 + len(FIELDS)

# ---------------- model (identical to training) ----------------
class RMSNorm(nn.Module):
    def __init__(s, d, eps=1e-5): super().__init__(); s.w = nn.Parameter(torch.ones(d)); s.eps = eps
    def forward(s, x):
        xf = x.float(); return (xf*torch.rsqrt(xf.pow(2).mean(-1,keepdim=True)+s.eps)).type_as(x)*s.w.to(x.dtype)

def rope_cache(T_, hd, base=10000.0):
    inv = 1.0/(base**(torch.arange(0,hd,2).float()/hd)); f = torch.outer(torch.arange(T_).float(), inv)
    return f.cos(), f.sin()
def apply_rope(x, cos, sin):
    x1, x2 = x[...,::2], x[...,1::2]; c, s_ = cos[:x.size(2)].to(x.dtype), sin[:x.size(2)].to(x.dtype)
    return torch.stack([x1*c-x2*s_, x1*s_+x2*c], -1).flatten(-2)

class Block(nn.Module):
    def __init__(s, c):
        super().__init__(); d = c["d_model"]; s.h = c["n_head"]; s.dp = c["dropout"]
        s.n1, s.n2 = RMSNorm(d), RMSNorm(d)
        s.qkv = nn.Linear(d, 3*d, bias=False); s.o = nn.Linear(d, d, bias=False)
        s.w1 = nn.Linear(d, c["d_ff"], bias=False); s.w3 = nn.Linear(d, c["d_ff"], bias=False); s.w2 = nn.Linear(c["d_ff"], d, bias=False)
        s.drop = nn.Dropout(c["dropout"])
    def forward(s, x, cos, sin):
        B,T_,C = x.shape
        q,k,v = s.qkv(s.n1(x)).view(B,T_,3,s.h,C//s.h).permute(2,0,3,1,4)
        q, k = apply_rope(q,cos,sin), apply_rope(k,cos,sin)
        y = F.scaled_dot_product_attention(q,k,v,is_causal=True,dropout_p=s.dp if s.training else 0.0)
        x = x + s.drop(s.o(y.transpose(1,2).reshape(B,T_,C)))
        h = s.n2(x); return x + s.drop(s.w2(F.silu(s.w1(h))*s.w3(h)))

class GPT(nn.Module):
    def __init__(s, c):
        super().__init__(); s.c = c
        s.emb = nn.Embedding(c["vocab"], c["d_model"]); s.blocks = nn.ModuleList(Block(c) for _ in range(c["n_layer"]))
        s.norm = RMSNorm(c["d_model"]); s.head = nn.Linear(c["d_model"], c["vocab"], bias=False); s.head.weight = s.emb.weight
        cos, sin = rope_cache(c["max_len"], c["d_model"]//c["n_head"])
        s.register_buffer("cos", cos, persistent=False); s.register_buffer("sin", sin, persistent=False)
    def forward(s, idx):
        x = s.emb(idx)
        for b in s.blocks: x = b(x, s.cos, s.sin)
        return s.head(s.norm(x))

# ---------------- MIDI codec ----------------
def decode(toks):
    toks = [int(t) for t in toks]; notes, bar, pos, pitch = [], -1, None, None
    for t in toks[HEADER_LEN:]:
        if t == EOS: break
        if t == BAR_ID: bar += 1; pos = pitch = None
        elif POS0 <= t < POS0+32: pos = t-POS0
        elif PITCH0 <= t < PITCH0+128: pitch = t-PITCH0
        elif DUR1 <= t < DUR1+32:
            if bar >= 0 and pos is not None and pitch is not None: notes.append((bar,pos,pitch,t-DUR1+1))
            pos = pitch = None
    return notes

def to_midi(toks, path, ts=(4,4), bpm=120, program=0):
    spb = STEPS_PER_BAR[ts]; notes = sorted(decode(toks))
    pm = pretty_midi.PrettyMIDI(initial_tempo=bpm)
    pm.time_signature_changes.append(pretty_midi.TimeSignature(ts[0], ts[1], 0))
    inst = pretty_midi.Instrument(program=int(program)); sec = 60.0/bpm/4
    for i,(b,p,pi,d) in enumerate(notes):
        s = b*spb+p
        if i+1 < len(notes): d = min(d, notes[i+1][0]*spb+notes[i+1][1]-s)
        inst.notes.append(pretty_midi.Note(velocity=90, pitch=pi, start=s*sec, end=(s+max(1,d))*sec))
    pm.instruments.append(inst); pm.write(str(path)); return len(notes)

def detect_key(pcs):
    maj = np.array([6.35,2.23,3.48,2.33,4.38,4.09,2.52,5.19,2.39,3.66,2.29,2.88])
    mnr = np.array([6.33,2.68,3.52,5.38,2.60,3.53,2.54,4.75,3.98,2.69,3.34,3.17])
    if pcs.std() == 0: return 0
    best = (-9, 0)
    for mode, prof in enumerate((maj, mnr)):
        for t in range(12):
            c = np.corrcoef(pcs, np.roll(prof, t))[0,1]
            if c > best[0]: best = (c, mode*12+t)
    return best[1]

# ---------------- prompt parser (rule-based) ----------------
_PC = {"C":0,"D":2,"E":4,"F":5,"G":7,"A":9,"B":11}
def _key_from(letter, acc, mode):
    pc = (_PC[letter.upper()] + (1 if acc in ("#","♯","sharp") else -1 if acc in ("b","♭","flat") else 0)) % 12
    return (NOTE_NAMES[pc], mode)
def parse_key(t):
    m = re.search(r"\b([A-Ga-g])\s*(#|♯|b|♭|sharp|flat)?\s*(major|minor|maj|min)\b", t, re.I)
    if m: return _key_from(m.group(1), (m.group(2) or "").lower(), "minor" if m.group(3).lower().startswith("min") else "major")
    m = re.search(r"\b([A-G])(#|♯|b|♭)?m\b", t)
    if m: return _key_from(m.group(1), m.group(2) or "", "minor")
    m = re.search(r"\bkey\s*(?:of|:)?\s*([A-G])(#|♯|b|♭)?(?![a-z#])", t)
    if m: return _key_from(m.group(1), m.group(2) or "", "major")
    return None

TEMPO_WORDS = [("very slow",55),("very fast",170),("slow",70),("ballad",70),("adagio",65),("andante",85),("moderate",105),
               ("medium",105),("mid-tempo",105),("upbeat",125),("allegro",130),("fast",140),("presto",170)]
INSTR_WORDS = {"piano":0,"electric piano":4,"harpsichord":6,"vibraphone":11,"marimba":12,"organ":19,"accordion":21,"harmonica":22,
               "acoustic guitar":24,"nylon guitar":24,"steel guitar":25,"guitar":25,"electric guitar":27,"bass":33,
               "violin":40,"viola":41,"cello":42,"strings":48,"harp":46,"choir":52,"voice":53,"trumpet":56,"trombone":57,"horn":60,"brass":61,
               "saxophone":65,"sax":65,"oboe":68,"clarinet":71,"flute":73,"recorder":74,"pan flute":75,"synth lead":80,"synth":80,
               "pad":88,"banjo":105,"sitar":104,"bell":14,"bells":14}
ALIASES = [("hip hop","hiphop"),("new age","newage"),("easy listening","easylistening"),("lo-fi","lofi"),("r&b","rnb")]

class Apollo:
    def __init__(self, model, labels, dev): self.model, self.labels, self.dev = model, labels, dev

    def parse(self, text):
        low = " " + text.lower() + " "
        for a, b in ALIASES: low = low.replace(a, b)
        spec = dict(bpm=None, ts=(4,4), key=parse_key(" " + text + " "), bars=8, dens=None, reg=None, genre=None, mood=None, program=None)
        m = re.search(r"\b(\d{1,2})\s*-?\s*bars?\b", low)
        if m: spec["bars"] = int(np.clip(int(m.group(1)), 1, 16))
        m = re.search(r"\b(\d{2,3})\s*(?:bpm|beats per minute)\b", low) or re.search(r"\btempo\s*(?:of|:|=)?\s*(\d{2,3})\b", low)
        if m: spec["bpm"] = int(np.clip(int(m.group(1)), 40, 279))
        else:
            for w, v in TEMPO_WORDS:
                if re.search(rf"\b{re.escape(w)}\b", low): spec["bpm"] = v; break
        m = re.search(r"\b(\d{1,2})\s*/\s*(\d{1,2})\b", low)
        if m and (int(m.group(1)), int(m.group(2))) in TS_LIST: spec["ts"] = (int(m.group(1)), int(m.group(2)))
        elif "waltz" in low: spec["ts"] = (3,4)
        for w, v in (("sparse",1),("simple",2),("busy",6),("dense",6)):
            if re.search(rf"\b{w}\b", low): spec["dens"] = v; break
        if "low register" in low: spec["reg"] = 1
        elif "high register" in low: spec["reg"] = 6
        for g in sorted(self.labels["genre"], key=len, reverse=True):
            if re.search(rf"\b{re.escape(g)}\b", low): spec["genre"] = g; break
        for md in sorted(self.labels["mood"], key=len, reverse=True):
            if re.search(rf"\b{re.escape(md)}\b", low): spec["mood"] = md; break
        for w in sorted(INSTR_WORDS, key=len, reverse=True):
            if re.search(rf"\b{re.escape(w)}\b", low): spec["program"] = INSTR_WORDS[w]; break
        return spec

    @torch.no_grad()
    def generate(self, bpm=None, ts=(4,4), key=None, bars=8, dens=None, reg=None, genre=None, mood=None, program=None,
                 temp=0.9, top_p=0.92, seed=None, max_new=1500):
        dev = self.dev
        g = torch.Generator(device=dev); g.manual_seed(seed if seed is not None else int(time.time()))
        spb = STEPS_PER_BAR[ts]
        hdr = [BOS,
               UNK_IDS[0] if bpm is None else FIELD_BASE["TEMPO"]+int(np.clip((bpm-40)//10,0,23)),
               FIELD_BASE["TS"]+TS_LIST.index(ts),
               UNK_IDS[2] if key is None else FIELD_BASE["KEY"]+(12 if key[1]=="minor" else 0)+NOTE_NAMES.index(key[0]),
               (FIELD_BASE["BARS"]+bars-1) if bars else UNK_IDS[3],
               UNK_IDS[4] if dens is None else FIELD_BASE["DENS"]+dens,
               UNK_IDS[5] if reg is None else FIELD_BASE["REG"]+reg,
               UNK_IDS[6] if genre is None else FIELD_BASE["GENRE"]+self.labels["genre"].index(genre),
               UNK_IDS[7] if mood is None else FIELD_BASE["MOOD"]+self.labels["mood"].index(mood),
               UNK_IDS[8] if program is None else FIELD_BASE["INSTR"]+int(program)//8]
        lo_b, hi_b = (bars, bars) if bars else (2, 16)
        out = list(hdr); state, bars_done, last_pos = "start", 0, -1
        for _ in range(max_new):
            x = torch.tensor([out[-SEQ_LEN:]], device=dev)
            logits = self.model(x)[0,-1].float()
            allow = torch.zeros(VOCAB, dtype=torch.bool, device=dev)
            if state == "start": allow[BAR_ID] = True
            elif state == "pos": allow[PITCH0:PITCH0+128] = True
            elif state == "pitch": allow[DUR1:DUR1+32] = True
            else:
                if last_pos+1 < spb: allow[POS0+last_pos+1:POS0+spb] = True
                if bars_done < hi_b: allow[BAR_ID] = True
                if bars_done >= lo_b and last_pos >= 0: allow[EOS] = True
                if not allow.any(): allow[EOS] = True
            l = logits.masked_fill(~allow, float("-inf"))/temp
            p = torch.softmax(l, -1); sp, si = p.sort(descending=True)
            sp[(sp.cumsum(0)-sp) > top_p] = 0; sp /= sp.sum()
            tok = int(si[torch.multinomial(sp, 1, generator=g)]); out.append(tok)
            if tok == EOS: break
            if tok == BAR_ID: bars_done += 1; last_pos = -1; state = "free"
            elif POS0 <= tok < POS0+32: last_pos = tok-POS0; state = "pos"
            elif PITCH0 <= tok < PITCH0+128: state = "pitch"
            else: state = "free"
        if out[-1] != EOS: out.append(EOS)
        return out

    def make(self, prompt, n=3, seed=None, outdir="apollo_out", temp=0.9, top_p=0.92, **override):
        """prompt -> n MIDI files. Extra kwargs override parsed fields, e.g. bpm=100, key=("D","minor"), program=40."""
        spec = self.parse(prompt); spec.update(override)
        print("parsed:", {k: v for k, v in spec.items() if v is not None})
        Path(outdir).mkdir(parents=True, exist_ok=True)
        slug = re.sub(r"[^a-z0-9]+", "_", prompt.lower()).strip("_")[:40] or "melody"
        base = seed if seed is not None else int(time.time()) % 100000
        paths = []
        for i in range(n):
            toks = self.generate(temp=temp, top_p=top_p, seed=base+i, **spec)
            path = f"{outdir}/{slug}_{base+i}.mid"
            nn_ = to_midi(toks, path, ts=spec["ts"], bpm=spec["bpm"] or 120, program=spec["program"] if spec["program"] is not None else 0)
            info = f"{nn_} notes"
            notes = decode(toks)
            if len(notes) >= 8:
                pcs = np.zeros(12)
                for _, _, p_, d in notes: pcs[p_ % 12] += d
                k = detect_key(pcs); info += f" | detected key {NOTE_NAMES[k % 12]} {'minor' if k >= 12 else 'major'}"
            print(f"  seed {base+i}: {path} | {info}"); paths.append(path)
        return paths

# ---------------- loader ----------------
def _token(token=None):
    if token: return token
    try:
        from google.colab import userdata
        return userdata.get("HF_TOKEN")
    except Exception:
        return os.environ.get("HF_TOKEN")           # None is fine for a public repo

def load(repo=REPO, device=None, token=None):
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file
    tok = _token(token); dev = device or ("cuda" if torch.cuda.is_available() else "cpu")
    get = lambda f: hf_hub_download(repo, f, token=tok)
    cfg = json.load(open(get("config.json")))["cfg"]
    assert json.load(open(get("vocab.json"))) == NAMES, "vocab.json in the repo differs from the vocab in this file"
    labels = json.load(open(get("labels.json")))
    model = GPT(cfg)
    missing, unexpected = model.load_state_dict({k: v.float() for k, v in load_file(get("model.safetensors")).items()}, strict=False)
    assert set(missing) <= {"head.weight"} and not unexpected, f"weights do not match the architecture: missing={missing} unexpected={unexpected}"
    model = (model.half() if dev == "cuda" else model.float()).to(dev).eval()
    print(f"loaded {repo} on {dev} | {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")
    return Apollo(model, labels, dev)

if __name__ == "__main__":
    import sys
    load().make(" ".join(sys.argv[1:]) or "8 bars in D minor at 90 bpm, sad piano melody", n=3, seed=1)