app.py

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