add standalone apollo_m.py + working usage example

README.md

@@ -74,24 +74,23 @@ Not intended as a replacement for composers, and not evaluated for any use beyon
 
 ## How to use
 
-The model is a custom PyTorch architecture (not a `transformers` model). The architecture, tokenizer and constrained decoder live in the training notebook and are not yet packaged as a library. Loading the weights:
+The model is a custom PyTorch architecture (not a `transformers` model). Everything needed (architecture, vocab, prompt parser, constrained decoder, MIDI writer) is in `apollo_m.py` in this repo.
 
 ```python
-import json, torch
+# pip install torch huggingface_hub safetensors pretty_midi
+import sys, os
 from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-
 repo = "viknesh100/Apollo-M"
-cfg = json.load(open(hf_hub_download(repo, "config.json")))["cfg"]
-state = load_file(hf_hub_download(repo, "model.safetensors"))
-model = GPT(cfg)    # GPT = the architecture class from the training code
-model.load_state_dict({k: v.float() for k, v in state.items()}, strict=False)  # head.weight is tied to emb.weight
-model = model.half().cuda().eval()
+sys.path.insert(0, os.path.dirname(hf_hub_download(repo, "apollo_m.py")))   # add token="hf_..." if the repo is private
+from apollo_m import load
+
+m = load(repo)
+m.make("8 bars in D minor at 90 bpm, sad piano melody", n=3, seed=1)       # writes MIDI files to ./apollo_out/
 ```
 
-Generation uses grammar-constrained sampling: at each step only tokens that keep the sequence valid are allowed (`BAR` then `POS` then `PITCH` then `DUR`, increasing positions inside a bar, exactly the requested number of bars), so every output parses to valid MIDI. Defaults: temperature 0.9, nucleus p 0.92.
+Generation uses grammar-constrained sampling: at each step only tokens that keep the sequence valid are allowed (`BAR` then `POS` then `PITCH` then `DUR`, increasing positions inside a bar, exactly the requested number of bars), so every output parses to valid MIDI. Defaults: temperature 0.9, nucleus p 0.92. Override fields directly if you prefer, e.g. `m.make("8 bars", key=("A","minor"), bpm=100, program=40)`.
 
-A small rule-based parser turns a prompt such as *"8 bars in D minor at 90 bpm, sad piano melody"* into header tokens (key, bars, bpm or tempo words, time signature, genre and mood words from the label lists, instrument words). The model itself does not read free text.
+A small rule-based parser turns the prompt into header tokens (key, bars, bpm or tempo words, time signature, genre and mood words from `labels.json`, instrument words). The model itself does not read free text.
 
 ## Files
 

apollo_m.py

@@ -0,0 +1,257 @@
+"""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)