"""SoftChart — Gradio demo for Hugging Face Spaces. Generate a Taiko no Tatsujin chart from any audio file, using the full system: - SoftChartGenerator (main model, plan-conditioned) - SoftChartPlanner (auto song-level planning) - SoftChartBeat (beat/downbeat for barline anchoring) All models load from the Hub via from_pretrained. MIT licensed. """ import os import re import tempfile import gradio as gr import numpy as np import torch from softchart.generate import generate_song, generate_song_slot, load_hf from softchart.fonts import cjk_font_path from softchart.grid import debias_to_grid, fit_grid_fixed_bpm, fit_grid_piecewise from softchart.hf import SoftChartPlanner from softchart.rhythm import snap_chart from softchart.tja import append_measure_with_gogo, gogo_measure_mask, write_tja_slots from softchart.tja_image import render_tja_image from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR V15_REPO = os.environ.get("SC_V15", "JacobLinCool/softchart-v15") PLAN_REPO = os.environ.get("SC_PLAN", "JacobLinCool/softchart-planner") DEVICE = "cuda" if torch.cuda.is_available() else "cpu" COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7} CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4", "roll": "5", "roll_big": "6", "balloon": "7"} SUB = 96 DEFAULT_TEMPERATURE = 0.8 DEFAULT_TOP_P = 0.9 _MODELS = {} def get_models(): if not _MODELS: u = load_hf(V15_REPO, device=DEVICE) if not getattr(u, "_dual", False) or u.beat is None: raise RuntimeError(f"{V15_REPO} must be a dual-mode generator with a beat head") _MODELS.update(gen=u, slot=u, beat=u) try: _MODELS["plan"] = SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval() except Exception: _MODELS["plan"] = None return _MODELS def load_logmel(path): import librosa wav, _ = librosa.load(path, sr=SR, mono=True) fb = librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=SR / 2) spec = torch.stft(torch.from_numpy(wav), N_FFT, hop_length=HOP, window=torch.hann_window(N_FFT), center=True, return_complex=True) mel = np.log(fb @ spec.abs().pow(2).numpy() + 1e-5).astype(np.float32) return mel, wav def auto_plan(mel, bpm, downbeats=None): T = mel.shape[1] dur = T / FPS flux = np.concatenate([[0], np.maximum(0, np.diff(mel, axis=1)).sum(0)]) beat = 60.0 / bpm edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \ else list(np.arange(0, dur, 4 * beat)) + [dur] vals = [float(flux[int(a * FPS):int(b * FPS)].mean()) if int(b * FPS) > int(a * FPS) else 0.0 for a, b in zip(edges, edges[1:])] if not vals: return None vals = np.array(vals) lo, hi = np.percentile(vals, 15), np.percentile(vals, 92) peak = int(np.argmax(vals)) plan = [] for i, (a, b) in enumerate(zip(edges, edges[1:])): frac = (vals[i] - lo) / max(hi - lo, 1e-6) d8 = int(np.clip(round(frac * 7), 0, 7)) fl = 1 if (vals[i] <= lo and 0 < i < len(vals) - 1) else (2 if i == peak and vals[i] >= hi else 0) plan.append([round(a, 3), round(b, 3), d8, fl]) return plan def learned_plan(planner, mel, course, bpm, downbeats=None): dur = mel.shape[1] / FPS beat = 60.0 / bpm edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \ else list(np.arange(0, dur, 4 * beat)) + [dur] feats, spans = [], [] for a, b in zip(edges, edges[1:]): seg = mel[:, int(a * FPS):int(b * FPS)] if seg.shape[1] < 2: continue fx = np.maximum(0, np.diff(seg, axis=1)).sum(0) feats.append(np.concatenate([seg.mean(1), seg.std(1), [fx.mean(), fx.std(), fx.max()]])) spans.append((round(float(a), 3), round(float(b), 3))) if not feats: return None cid = {"easy": 0, "normal": 1, "hard": 2, "oni": 3}[course] x = torch.tensor(np.array(feats), dtype=torch.float32)[None].to(DEVICE) with torch.no_grad(): pd, pf = planner(x, torch.tensor([cid], device=DEVICE)) d8 = pd[0].argmax(-1).cpu().numpy() fl = pf[0].argmax(-1).cpu().numpy() return [[a, b, int(d), int(f)] for (a, b), d, f in zip(spans, d8, fl)] def group_quantize(times, phase, grid, min_run=3): n = len(times) slots = [0] * n i = 0 while i < n: j = i while j + 1 < n: ioi = times[j + 1] - times[j] ref = (times[j] - times[i]) / (j - i) if j > i else ioi if 0.02 < ioi < 1.2 and abs(ioi - ref) < 0.22 * max(ref, 1e-6): j += 1 else: break if j - i + 1 >= min_run: k = max(1, int(round((times[j] - times[i]) / (j - i) / grid))) anchor = int(round((times[i] - phase) / grid)) for m in range(j - i + 1): slots[i + m] = anchor + m * k else: for m in range(i, j + 1): slots[m] = int(round((times[m] - phase) / grid)) i = j + 1 return slots def upload_wave_name(audio_path): name = os.path.basename(str(audio_path)).strip() name = re.sub(r"[\r\n]+", " ", name) return name or "song.ogg" def output_tja_path(wave_name, course): stem = os.path.splitext(os.path.basename(wave_name))[0].strip() or "softchart" stem = re.sub(r"[^0-9A-Za-z._ -]+", "_", stem).strip(" ._") or "softchart" return os.path.join(tempfile.mkdtemp(), f"{stem}_{course}.tja") def write_tja(gen, bpm, title, course, level, wave, downbeats=None, grid_fit=None, plan=None): hits = sorted((h["t"], CHAR[h["type"]]) for h in gen["hits"]) beat = 60.0 / bpm grid = beat / (SUB / 4) bias = 0.0 if grid_fit is not None and hits: # authoritative fitted grid: barlines ARE the fitted downbeats. # De-bias the generator's systematic latency (global shift only), # then anchor slot 0 on the last fitted barline at/before the first note. times, bias = debias_to_grid([t for t, _ in hits], grid_fit["phase"], grid) phase = grid_fit["phase"] + float(np.floor((times[0] - grid_fit["phase"]) / (4 * beat))) * 4 * beat q_times = list(times) else: times = np.array([t for t, _ in hits]) if hits else np.array([0.0]) cands = np.arange(0, beat, grid / 4) phase = float(cands[int(np.argmin([np.mean(np.abs(((times - o) / grid) - np.round((times - o) / grid))) for o in cands]))]) q_times = [t for t, _ in hits] slot_idx = group_quantize(q_times, phase, grid) if slot_idx and min(slot_idx) < 0: # note quantized just before the anchor barline: pull back whole bars # so nothing is dropped (barline alignment is preserved mod SUB) nb = int(np.ceil(-min(slot_idx) / SUB)) slot_idx = [s + nb * SUB for s in slot_idx] phase -= nb * SUB * grid slots = {} for idx, (t, ch) in zip(slot_idx, hits): if idx >= 0 and idx not in slots: slots[idx] = ch for sp in gen["spans"]: i0 = int(round((sp["t0"] - bias - phase) / grid)) i1 = int(round((sp["t1"] - bias - phase) / grid)) while i0 in slots: i0 += 1 while i1 in slots or i1 <= i0: i1 += 1 if i0 >= 0: slots[i0] = CHAR[sp["type"]] slots[i1] = "8" if slots and grid_fit is None: # Gridless anchoring: shift so the first note sits on a detected # downbeat if one is nearby, otherwise on the first barline. first_t = min(slots) * grid + phase anchor_t = None if downbeats is not None and len(downbeats): near = downbeats[downbeats <= first_t + 0.12] if len(near) and first_t - near[-1] < 4 * beat: anchor_t = near[-1] shift = int(round((anchor_t - phase) / grid)) if anchor_t is not None else min(slots) if shift: slots = {k - shift: v for k, v in slots.items()} phase += shift * grid n_meas = (max(slots) // SUB + 1) if slots else 1 measure_starts = phase + np.arange(n_meas + 1, dtype=float) * (4 * beat) gogo_mask = gogo_measure_mask(plan, measure_starts, n_meas) lines = [] in_gogo = False for m in range(n_meas): in_gogo = append_measure_with_gogo( lines, "".join(slots.get(m * SUB + k, "0") for k in range(SUB)) + ",", m, gogo_mask, in_gogo) if in_gogo: lines.append("#GOGOEND") balloons = [10] * sum(1 for s in gen["spans"] if s["type"] == "balloon") return "\n".join([ f"TITLE:{title} (SoftChart)", f"BPM:{bpm:g}", f"WAVE:{wave}", f"OFFSET:{-phase:.3f}", f"COURSE:{'Oni' if course == 'oni' else course.capitalize()}", f"LEVEL:{level}", f"BALLOON:{','.join(map(str, balloons))}" if balloons else "BALLOON:", "", "#START", *lines, "#END"]) + "\n" def render_audio_plan(mel, title, course, plan=None): import matplotlib matplotlib.use("Agg") from matplotlib import font_manager import matplotlib.pyplot as plt font_path = cjk_font_path() if font_path is not None: font_manager.fontManager.addfont(font_path) matplotlib.rcParams["font.family"] = font_manager.FontProperties(fname=font_path).get_name() matplotlib.rcParams["axes.unicode_minus"] = False dur = mel.shape[1] / FPS fig = plt.figure(figsize=(13, 4.4 if plan else 3.2)) gs = fig.add_gridspec(2 if plan else 1, 1, height_ratios=[3.0, 1.0] if plan else [1], hspace=0.14 if plan else 0.0) ax0 = fig.add_subplot(gs[0]) ax0.imshow(mel, aspect="auto", origin="lower", cmap="magma", extent=[0, dur, 0, N_MELS]) ax0.set_ylabel("mel") ax0.set_title(f"{title} — {course} | full-song mel spectrogram") ax0.set_xlim(0, dur) ax0.grid(axis="x", alpha=0.18) if plan: ax0.set_xticklabels([]) ax1 = fig.add_subplot(gs[1], sharex=ax0) for a, b, d, f in plan: c = "#d64545" if f == 2 else ("#4a90d9" if f == 1 else "#999999") ax1.bar((a + b) / 2, max(d, 0.15), width=max((b - a) * 0.92, 0.01), color=c, alpha=0.85) ax1.set_xlim(0, dur) ax1.set_ylim(0, 8) ax1.set_yticks([0, 4, 8]) ax1.set_ylabel("plan", fontsize=8) ax1.set_xlabel("time (s) — plan: grey=density blue=gap red=climax") ax1.grid(axis="x", alpha=0.18) else: ax0.set_xlabel("time (s)") out = tempfile.mktemp(suffix=".png") fig.savefig(out, dpi=130, bbox_inches="tight") plt.close(fig) return out def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_planner, sampling, temperature, top_p, progress=gr.Progress()): if audio is None: raise gr.Error("Please upload an audio file.") def P(frac, desc): # progress is a no-op when called outside a Gradio event try: progress(frac, desc=desc) except Exception: pass P(0.01, "Loading models (first run downloads them from the Hub)…") M = get_models() P(0.08, "Computing log-mel…") mel, wav = load_logmel(audio) grid = dbs = None if use_beat and M["beat"] is not None: P(0.14, "Fitting beat grid…") # global robust (period, phase) fit over the whole song — much more # precise than per-peak use (each raw peak carries ~±23 ms bin noise) grid = fit_grid_piecewise(M["beat"], mel, device=DEVICE) if grid is not None: # rigid synthesized barlines when the fit is trustworthy; raw peaks # (plan-block edges only, no anchoring) when it is not dbs = grid["downbeats"] if grid["ok"] else grid["db_peaks"] if not grid["ok"]: grid = None if bpm_override and bpm_override > 0: bpm = float(bpm_override) if use_beat and M["beat"] is not None and (grid is None or abs(grid["bpm"] - bpm) > 0.5): # the user KNOWS the tempo: phase-only fit with a trusted period — # much easier than the free fit, often unlocks the slot-exact path g2 = fit_grid_fixed_bpm(M["beat"], mel, bpm, device=DEVICE) grid = g2 if (g2 is not None and g2["ok"]) else None if grid is not None: dbs = grid["downbeats"] elif grid is not None: bpm = grid["bpm"] # already integer-snapped when the residual allows elif dbs is not None and len(dbs) > 4: period = float(np.median(np.diff(dbs))) # downbeat gap = one 4/4 bar bpm = 240.0 / period if period > 0 else 0.0 while bpm >= 210: # octave guard bpm /= 2.0 while 0 < bpm < 70: bpm *= 2.0 else: import librosa bpm = float(np.atleast_1d(librosa.beat.beat_track(y=wav, sr=SR)[0])[0]) if grid is None and abs(bpm - round(bpm)) < 0.06: bpm = float(round(bpm)) temperature = float(np.clip(temperature or DEFAULT_TEMPERATURE, 0.2, 1.2)) top_p = float(np.clip(top_p or DEFAULT_TOP_P, 0.5, 1.0)) plan = None P(0.28, "Planning song structure…") if any(getattr(m, "_has_plan", False) for m in (M["gen"], M["slot"]) if m is not None): if use_planner and M["plan"] is not None: plan = learned_plan(M["plan"], mel, course, bpm, dbs) elif auto_plan_on: plan = auto_plan(mel, bpm, dbs) wave_name = upload_wave_name(audio) title = os.path.splitext(wave_name)[0] slot_used = False if grid is not None and M["slot"] is not None: # slot-exact path: the decoder emits TJA lattice indices directly — # no quantization, tuplets even by construction g = generate_song_slot( M["slot"], mel, grid, course, level=int(level), density_bucket=COURSE_DENS[course], greedy=not sampling, seed=0, temperature=temperature, top_p=top_p, device=DEVICE, plan=plan, on_progress=lambda d, t: P(0.32 + 0.58 * d / max(t, 1), f"Generating notes (slot-exact)… {d}/{t} windows")) tja = write_tja_slots(g, grid, title, course, int(level), wave_name, plan=plan) slot_used = True else: g = generate_song( M["gen"], mel, course, level=int(level), density_bucket=COURSE_DENS[course], greedy=not sampling, temperature=temperature, top_p=top_p, seed=0, device=DEVICE, plan=plan, on_progress=lambda d, t: P(0.32 + 0.58 * d / max(t, 1), f"Generating notes… {d}/{t} windows")) g = snap_chart(g, bpm) tja = write_tja(g, bpm, title, course, int(level), wave_name, dbs, grid_fit=grid, plan=plan) P(0.92, "Writing TJA…") tja_path = output_tja_path(wave_name, course) with open(tja_path, "w", encoding="utf-8") as f: f.write(tja) P(0.95, "Rendering preview…") audio_plan_img = render_audio_plan(mel, title, course, plan=plan) tja_img = render_tja_image(tja) if grid is not None: grid_info = (f" · grid: rms {grid['rms_ms']:.1f}ms ({grid['inlier_frac']:.0%} inlier)" + (f" · {grid['n_segments']} tempo segs" if grid.get("piecewise") else "") + (" · slot-exact" if slot_used else "")) elif use_beat and M["beat"] is not None: grid_info = " · grid: unreliable, barline anchoring off" else: grid_info = "" sampling_info = f" · sampling T {temperature:.2f}, top-p {top_p:.2f}" if sampling else "" info = (f"BPM {bpm:.1f} · {len(g['hits'])} notes · {len(g['spans'])} spans" + grid_info + sampling_info + (f" · plan: {sum(1 for b in plan if b[3]==1)} gaps, {sum(1 for b in plan if b[3]==2)} climax" if plan else "")) return tja_path, info, tja_img, audio_plan_img with gr.Blocks(title="SoftChart — AI Taiko chart generator") as demo: gr.Markdown( "# 🥁 SoftChart\n" "Generate a **Taiko no Tatsujin** chart from any song. 8M-param plan-conditioned " "model + learned planner + beat-anchoring — all MIT-licensed.\n\n" "Outputs a full-song mel/plan preview, a full TJA chart image, and a playable `.tja`. " "*MIT-licensed.*" ) with gr.Row(): with gr.Column(): audio = gr.Audio(type="filepath", label="Song (any format)") course = gr.Dropdown(["easy", "normal", "hard", "oni"], value="oni", label="Difficulty (course)") level = gr.Slider(1, 10, value=9, step=1, label="Level (1–10 stars)") bpm = gr.Number(label="BPM (0 = auto-detect)", value=0) with gr.Row(): auto_plan_on = gr.Checkbox(label="Auto-plan (heuristic structure)", value=True) use_planner = gr.Checkbox(label="Learned planner", value=True) with gr.Row(): use_beat = gr.Checkbox(label="Beat-anchor barlines", value=True) sampling = gr.Checkbox(label="Sampling (diverse) vs greedy (best)", value=False) with gr.Row(): temperature = gr.Slider(0.2, 1.2, value=DEFAULT_TEMPERATURE, step=0.05, label="Temperature", interactive=False) top_p = gr.Slider(0.5, 1.0, value=DEFAULT_TOP_P, step=0.01, label="Top-p", interactive=False) btn = gr.Button("Generate chart", variant="primary") with gr.Column(): out_tja = gr.File(label="Download .tja") out_info = gr.Textbox(label="Result", interactive=False) out_chart = gr.Image(label="TJA chart") out_img = gr.Image(label="Full-song mel + plan") sampling.change( lambda enabled: (gr.update(interactive=enabled), gr.update(interactive=enabled)), sampling, [temperature, top_p]) btn.click(generate, [audio, course, level, bpm, auto_plan_on, use_beat, use_planner, sampling, temperature, top_p], [out_tja, out_info, out_chart, out_img]) if __name__ == "__main__": # show_api=False avoids a gradio_client schema-introspection bug # ("bool is not iterable") that can appear on some versions demo.launch(show_api=False)