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