| """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 |
|
|
| |
| 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) |
|
|
| |
| 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)) |
|
|
| |
| 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] |
|
|
| |
| _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 |
|
|
| |
| 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") |
|
|
| 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) |
|
|
