softchart / app.py
JacobLinCool's picture
Restore selectable legacy model generations
be60cf4
Raw
History Blame Contribute Delete
25.5 kB
"""SoftChart — selectable Taiko chart model generations for HF Spaces.
V1.8 is the recommended scratch-trained hierarchical model. V1.5–V1.7 remain
available through their frozen runtime so users can compare generations.
Arbitrary uploads always use time generation; an estimated beat grid may
quantize the exported TJA but is never promoted to a trusted-meter slot path.
"""
import logging
import os
import re
import shutil
import subprocess
import tempfile
from pathlib import Path
import gradio as gr
import numpy as np
import torch
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from legacy_softchart.generate import (
generate_song as generate_song_legacy,
load_hf as load_hf_legacy,
)
from legacy_softchart.grid import (
fit_grid_fixed_bpm as fit_grid_fixed_bpm_legacy,
fit_grid_piecewise as fit_grid_piecewise_legacy,
)
from legacy_softchart.hf import SoftChartPlanner
from legacy_softchart.tja import append_measure_with_gogo, gogo_measure_mask
from softchart.generate import generate_song as generate_song_v18, load_hf as load_hf_v18
from softchart.fonts import cjk_font_path
from softchart.grid import debias_to_grid, fit_grid_fixed_bpm, fit_grid_piecewise
from softchart.preview_audio import synthesize_taiko_preview
from softchart.rhythm import snap_chart
from softchart.tja_image import render_tja_image
from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR
LOGGER = logging.getLogger("softchart.space")
STATIC_DIR = Path(__file__).with_name("static")
PLAN_REPO = "JacobLinCool/softchart-planner"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
EXPECTED_PARAMETERS = 8_985_091
MODEL_SPECS = {
"v1.8": {
"repo": "JacobLinCool/softchart-v18",
"display": "V1.8 hierarchical",
"runtime": "v18",
},
"v1.7": {
"repo": "JacobLinCool/softchart-v17",
"display": "V1.7",
"runtime": "legacy",
},
"v1.7-small": {
"repo": "JacobLinCool/softchart-v17-small",
"display": "V1.7 small",
"runtime": "legacy",
},
"v1.7-tiny": {
"repo": "JacobLinCool/softchart-v17-tiny",
"display": "V1.7 tiny",
"runtime": "legacy",
},
"v1.6": {
"repo": "JacobLinCool/softchart-v16",
"display": "V1.6",
"runtime": "legacy",
},
"v1.6-small": {
"repo": "JacobLinCool/softchart-v16-small",
"display": "V1.6 small",
"runtime": "legacy",
},
"v1.6-tiny": {
"repo": "JacobLinCool/softchart-v16-tiny",
"display": "V1.6 tiny",
"runtime": "legacy",
},
"v1.5": {
"repo": "JacobLinCool/softchart-v15",
"display": "V1.5",
"runtime": "legacy",
},
}
DEFAULT_MODEL = "v1.8"
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
_MODELS = {}
_PLANNER = {}
_MEL_FB = None
_STFT_WINDOW = None
def get_model(model_choice):
spec = MODEL_SPECS.get(model_choice)
if spec is None:
raise ValueError(
f"Unknown model {model_choice!r}. Choose one of: "
f"{', '.join(MODEL_SPECS)}."
)
if model_choice not in _MODELS:
if spec["runtime"] == "v18":
model = load_hf_v18(spec["repo"], device=DEVICE)
else:
model = load_hf_legacy(spec["repo"], device=DEVICE)
capabilities = getattr(model, "_softchart_capabilities", {})
if spec["runtime"] == "v18":
required = {
"aux": True,
"beat_head": True,
"dual": True,
"hierarchical_ctx": True,
"global_ctx": False,
"plan": False,
}
mismatched = {
name: capabilities.get(name)
for name, expected in required.items()
if capabilities.get(name) is not expected
}
parameters = sum(parameter.numel() for parameter in model.parameters())
if (not getattr(model, "_hierarchical_ctx", False)
or getattr(model, "_has_plan", False) or mismatched
or parameters != EXPECTED_PARAMETERS):
raise RuntimeError(
f"{spec['repo']} is not the expected hierarchical V1.8 artifact"
)
bundle = {
"model": model,
"generate": generate_song_v18,
"fit_grid": fit_grid_piecewise,
"fit_grid_fixed": fit_grid_fixed_bpm,
"preprocess": load_logmel_v18,
"runtime": "v18",
}
else:
if (not getattr(model, "_dual", False) or model.beat is None
or not getattr(model, "_has_plan", False)):
raise RuntimeError(
f"{spec['repo']} is not a supported V1.5–V1.7 artifact"
)
bundle = {
"model": model,
"generate": generate_song_legacy,
"fit_grid": fit_grid_piecewise_legacy,
"fit_grid_fixed": fit_grid_fixed_bpm_legacy,
"preprocess": load_logmel_legacy,
"runtime": "legacy",
}
_MODELS[model_choice] = bundle
return _MODELS[model_choice]
def get_legacy_planner():
if "planner" not in _PLANNER:
_PLANNER["planner"] = (
SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval()
)
return _PLANNER["planner"]
def load_logmel_v18(path):
"""Apply the exact V1.8 FFmpeg + torch.stft preprocessing contract."""
import librosa
global _MEL_FB, _STFT_WINDOW
command = [
"ffmpeg", "-nostdin", "-hide_banner", "-loglevel", "error",
"-i", os.fspath(path), "-vn", "-ac", "2", "-ar", str(SR),
"-f", "f32le", "-acodec", "pcm_f32le", "pipe:1",
]
proc = subprocess.run(
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False,
)
if proc.returncode != 0:
detail = proc.stderr.decode("utf-8", errors="replace").strip()
raise ValueError(f"FFmpeg could not decode this upload: {detail}")
decoded = np.frombuffer(proc.stdout, dtype="<f4")
if decoded.size % 2:
raise ValueError("Decoded stereo audio contains an incomplete frame.")
wav = decoded.reshape(-1, 2).mean(axis=1, dtype=np.float32)
if wav.size < SR:
raise ValueError("Audio must be at least one second long.")
if not np.isfinite(wav).all():
raise ValueError("Decoded audio contains non-finite samples.")
if _MEL_FB is None:
fb = librosa.filters.mel(
sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=SR / 2,
)
_MEL_FB = torch.from_numpy(fb)
_STFT_WINDOW = torch.hann_window(N_FFT)
wav_tensor = torch.from_numpy(np.ascontiguousarray(wav))[None]
with torch.no_grad():
spec = torch.stft(
wav_tensor, N_FFT, hop_length=HOP, window=_STFT_WINDOW,
center=True, return_complex=True,
)[0]
mel = torch.log(_MEL_FB @ spec.abs().pow(2) + 1e-5)
mel = mel.numpy().astype(np.float32)
return mel, wav
def load_logmel_legacy(path):
"""Preserve the preprocessing contract used by V1.5–V1.7."""
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 learned_plan(planner, mel, course, bpm, downbeats=None):
duration = mel.shape[1] / FPS
beat = 60.0 / bpm
edges = (
list(downbeats[::4]) + [duration]
if downbeats is not None and len(downbeats) >= 2
else list(np.arange(0, duration, 4 * beat)) + [duration]
)
features = []
spans = []
for start, end in zip(edges, edges[1:]):
segment = mel[:, int(start * FPS):int(end * FPS)]
if segment.shape[1] < 2:
continue
flux = np.maximum(0, np.diff(segment, axis=1)).sum(0)
features.append(np.concatenate([
segment.mean(1), segment.std(1),
[flux.mean(), flux.std(), flux.max()],
]))
spans.append((round(float(start), 3), round(float(end), 3)))
if not features:
return None
course_id = {"easy": 0, "normal": 1, "hard": 2, "oni": 3}[course]
inputs = torch.tensor(np.array(features), dtype=torch.float32)[None].to(DEVICE)
with torch.no_grad():
density_logits, flag_logits = planner(
inputs, torch.tensor([course_id], device=DEVICE)
)
densities = density_logits[0].argmax(-1).cpu().numpy()
flags = flag_logits[0].argmax(-1).cpu().numpy()
return [
[start, end, int(density), int(flag)]
for (start, end), density, flag in zip(spans, densities, flags)
]
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).replace("\\", "/")).strip()
name = re.sub(r"[\x00-\x1f\x7f]+", " ", name).strip()
return name or "song.ogg"
def output_tja_path(wave_name, course, directory):
stem = os.path.splitext(os.path.basename(wave_name))[0].strip() or "softchart"
stem = re.sub(r"[<>:\"/\\|?*\x00-\x1f]+", "_", stem).strip(" ._") or "softchart"
return os.path.join(directory, 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 measure in range(n_meas):
in_gogo = append_measure_with_gogo(
lines,
"".join(
slots.get(measure * SUB + slot, "0") for slot in range(SUB)
) + ",",
measure,
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_song_structure(mel, title, course, out_path, 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))
grid = 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(grid[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(grid[1], sharex=ax0)
for start, end, density, flag in plan:
color = "#d64545" if flag == 2 else (
"#4a90d9" if flag == 1 else "#999999"
)
ax1.bar(
(start + end) / 2,
max(density, 0.15),
width=max((end - start) * 0.92, 0.01),
color=color,
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) — legacy plan: density / gap / climax")
ax1.grid(axis="x", alpha=0.18)
else:
ax0.set_xlabel("time (s)")
fig.savefig(out_path, dpi=130, bbox_inches="tight")
plt.close(fig)
return out_path
def _progress(stage, fraction, title, detail):
return {
"kind": "progress",
"stage": stage,
"progress": fraction,
"title": title,
"detail": detail,
}
def _uploaded_file(value, original_name):
if value is None:
raise ValueError("Upload a music file to begin.")
if not isinstance(original_name, str) or not original_name.strip():
raise ValueError("The upload is missing its original filename. Please select it again.")
data = value if isinstance(value, gr.FileData) else gr.FileData.model_validate(value)
path = os.path.realpath(data.path)
if not os.path.isfile(path):
raise ValueError("The uploaded file is no longer available. Please select it again.")
return path, upload_wave_name(original_name)
def _validate_controls(course, level, bpm, temperature, top_p, drum_volume):
if course not in COURSE_DENS:
raise ValueError(f"Unsupported difficulty: {course}")
if not 1 <= int(level) <= 10:
raise ValueError("Level must be between 1 and 10 stars.")
if bpm != 0 and not 30 <= float(bpm) <= 400:
raise ValueError("Manual BPM must be between 30 and 400; use 0 for auto-detection.")
if not 0.2 <= float(temperature) <= 1.2:
raise ValueError("Temperature must be between 0.2 and 1.2.")
if not 0.5 <= float(top_p) <= 1.0:
raise ValueError("Top-p must be between 0.5 and 1.0.")
if not 0 <= float(drum_volume) <= 1.5:
raise ValueError("Taiko volume must be between 0% and 150%.")
def _file_data(path, *, name=None, mime_type=None):
return gr.FileData(
path=os.path.realpath(path),
orig_name=name or os.path.basename(path),
mime_type=mime_type,
).model_dump()
app = gr.Server(
title="SoftChart",
description="Conditional Taiko chart generation with synchronized audio preview.",
)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
@app.get("/", include_in_schema=False)
async def homepage():
return FileResponse(
STATIC_DIR / "index.html",
headers={"Cache-Control": "no-cache"},
)
@app.get("/health", include_in_schema=False)
async def health():
return {
"status": "ok",
"default_model": DEFAULT_MODEL,
"available_models": list(MODEL_SPECS),
"device": DEVICE,
"models_loaded": sorted(_MODELS),
"planner_loaded": "planner" in _PLANNER,
}
@app.api(
name="generate_chart",
description="Generate a TJA chart, visual previews, and a synchronized taiko audio mix.",
concurrency_limit=1,
concurrency_id="generation-gpu",
queue=True,
stream_every=0.1,
)
def generate_chart(
audio: gr.FileData,
audio_name: str,
course: str,
level: int,
bpm_override: float,
use_beat: bool,
sampling: bool,
temperature: float,
top_p: float,
drum_volume: float,
model_choice: str = DEFAULT_MODEL,
) -> dict[str, object]:
"""Stream the actual inference stages to the custom frontend."""
workdir = tempfile.mkdtemp(prefix="softchart-request-")
try:
audio_path, wave_name = _uploaded_file(audio, audio_name)
_validate_controls(course, level, bpm_override, temperature, top_p, drum_volume)
level = int(level)
bpm_override = float(bpm_override)
temperature = float(temperature)
top_p = float(top_p)
drum_volume = float(drum_volume)
spec = MODEL_SPECS.get(model_choice)
if spec is None:
raise ValueError(
f"Unknown model {model_choice!r}. Choose one of: "
f"{', '.join(MODEL_SPECS)}."
)
yield _progress(
"loading", 0.03, "Preparing",
f"Loading SoftChart {spec['display']}.",
)
bundle = get_model(model_choice)
model = bundle["model"]
yield _progress(
"audio", 0.11, "Listening",
"Mapping rhythm, melody, and timbre.",
)
mel, wav = bundle["preprocess"](audio_path)
yield _progress(
"beat", 0.21, "Finding the beat",
"Finding beats and tempo.",
)
grid = dbs = None
if use_beat:
grid = bundle["fit_grid"](model, mel, device=DEVICE)
if grid is not None:
dbs = grid["downbeats"] if grid["ok"] else grid["db_peaks"]
if not grid["ok"]:
grid = None
if bpm_override > 0:
bpm = bpm_override
if use_beat and (grid is None or abs(grid["bpm"] - bpm) > 0.5):
fixed_grid = bundle["fit_grid_fixed"](
model, mel, bpm, device=DEVICE
)
grid = fixed_grid if fixed_grid is not None and fixed_grid["ok"] else None
if grid is not None:
dbs = grid["downbeats"]
elif grid is not None:
bpm = float(grid["bpm"])
elif dbs is not None and len(dbs) > 4:
period = float(np.median(np.diff(dbs)))
if period <= 0:
raise RuntimeError("The beat model returned an invalid downbeat interval.")
bpm = 240.0 / period
while bpm >= 210:
bpm /= 2.0
while bpm < 70:
bpm *= 2.0
else:
import librosa
bpm = float(np.atleast_1d(librosa.beat.beat_track(y=wav, sr=SR)[0])[0])
if not np.isfinite(bpm) or bpm <= 0:
raise RuntimeError("Could not determine a reliable BPM. Enter it in Advanced settings.")
if grid is None and abs(bpm - round(bpm)) < 0.06:
bpm = float(round(bpm))
plan = None
if bundle["runtime"] == "v18":
structure_title = "Reading the structure"
structure_detail = "Building whole-song hierarchical context."
else:
structure_title = "Shaping the arc"
structure_detail = "Applying the original legacy song planner."
plan = learned_plan(
get_legacy_planner(), mel, course, bpm, downbeats=dbs
)
yield _progress(
"structure", 0.34, structure_title, structure_detail,
)
yield _progress(
"generate", 0.47, "Writing the chart",
"Writing playable Taiko patterns.",
)
title = os.path.splitext(wave_name)[0]
generated = bundle["generate"](
model, mel, course, level=level,
density_bucket=COURSE_DENS[course], greedy=not sampling,
temperature=temperature, top_p=top_p, seed=0, device=DEVICE,
plan=plan,
)
generated = snap_chart(generated, bpm)
tja = write_tja(
generated, bpm, title, course, level, wave_name, dbs,
grid_fit=grid, plan=plan,
)
yield _progress(
"export", 0.81, "Rendering",
"Building the TJA and previews.",
)
tja_path = output_tja_path(wave_name, course, workdir)
with open(tja_path, "w", encoding="utf-8") as output_file:
output_file.write(tja)
chart_image_path = os.path.join(workdir, "chart.png")
structure_image_path = os.path.join(workdir, "song-structure.png")
render_tja_image(tja, out_path=chart_image_path)
render_song_structure(
mel, title, course, out_path=structure_image_path, plan=plan,
)
yield _progress(
"mix", 0.91, "Mixing",
"Mixing Taiko with your track.",
)
stem = Path(tja_path).stem
preview_path = os.path.join(workdir, f"{stem}_taiko-preview.wav")
mix_stats = synthesize_taiko_preview(
audio_path,
tja,
preview_path,
drum_gain=drum_volume,
sample_rate=44100,
)
grid_rms = round(float(grid["rms_ms"]), 1) if grid is not None else None
metrics = {
"model": model_choice,
"bpm": round(float(bpm), 1),
"notes": len(generated["hits"]),
"spans": len(generated["spans"]),
"timing": (
"hierarchical time + grid quantization"
if bundle["runtime"] == "v18"
else "legacy time + grid quantization"
),
"grid_rms_ms": grid_rms,
"preview_hits": int(mix_stats["rendered_hit_count"]),
}
yield {
"kind": "complete",
"stage": "complete",
"progress": 1.0,
"title": "Ready",
"detail": "Play it or download it.",
"metrics": metrics,
"files": {
"tja": _file_data(tja_path, mime_type="text/plain"),
"audio": _file_data(preview_path, mime_type="audio/wav"),
"chart_image": _file_data(chart_image_path, mime_type="image/png"),
"structure_image": _file_data(
structure_image_path, mime_type="image/png"
),
},
}
except Exception as exc:
LOGGER.exception("SoftChart generation failed")
detail = (str(exc) if isinstance(exc, (ValueError, RuntimeError))
else "The server could not complete this chart. Please try again shortly.")
yield {
"kind": "error",
"stage": "error",
"progress": 0.0,
"title": "Generation did not complete",
"detail": detail,
}
finally:
shutil.rmtree(workdir, ignore_errors=True)
if __name__ == "__main__":
app.launch(max_file_size="200mb")