| |
| """Crop licensed demo clips and prepare notes with the fixed browser ONNX model.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import hashlib |
| import json |
| from pathlib import Path |
| import re |
| import shutil |
| import subprocess |
| import time |
|
|
| import numpy as np |
| import onnxruntime as ort |
| import soundfile as sf |
| import torch |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def digest(path: Path) -> str: |
| with path.open("rb") as stream: |
| return hashlib.file_digest(stream, "sha256").hexdigest() |
|
|
|
|
| def generate(sessions, features): |
| encoder, cross_cache, decoder = sessions |
| encoded = encoder.run(None, {"spectrogram": features})[0] |
| cross = cross_cache.run(None, {"encoder_hidden_states": encoded}) |
| cache = [np.zeros((13, 6, 0, 64), np.float32) for _ in range(16)] |
| token = np.zeros((1, 13, 1), np.int64) |
| unfinished = np.ones_like(token) |
| windows = [] |
| for position in range(256): |
| feeds = {"token_ids": token, "position_ids": np.asarray([position], np.int64)} |
| for layer in range(8): |
| for suffix, values in (("self_key", cache[2*layer]), ("self_value", cache[2*layer+1]), |
| ("cross_key", cross[2*layer]), ("cross_value", cross[2*layer+1])): |
| feeds[f"layer_{layer}_{suffix}"] = values |
| result = decoder.run(None, feeds) |
| token = result[0].argmax(-1).astype(np.int64) |
| visible = token * unfinished |
| windows.append(visible) |
| cache = result[1:] |
| unfinished *= visible != 1 |
| if not unfinished.any(): |
| break |
| return np.concatenate(windows, axis=-1)[0].tolist(), int(unfinished.sum()) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--dataset-root", type=Path, default=Path("/home/jack/workspace/datasets")) |
| parser.add_argument("--graph-dir", type=Path, default=ROOT / "progress/evaluations/yptf_web_fp32_d26") |
| parser.add_argument("--soundfont", type=Path, default=Path("/home/jack/workspace/models/soundfonts/GeneralUser_GS.sf2")) |
| parser.add_argument("--output-dir", type=Path, default=ROOT / "demo/web/public/examples") |
| args = parser.parse_args() |
| config = json.loads((ROOT / "configs/demo/public_examples.json").read_text()) |
| for row in config["examples"]: |
| if not re.fullmatch(r"[a-z0-9-]+", row["id"]): |
| raise SystemExit(f"Example ID must be safe in a URL: {row['id']}") |
| manifest = json.loads((args.graph_dir / "manifest.json").read_text()) |
| if manifest["precision"] != "fp32": |
| raise SystemExit("Public examples require the checked FP32 browser export") |
| options = ort.SessionOptions() |
| options.intra_op_num_threads = 4 |
| options.inter_op_num_threads = 1 |
| torch.set_num_threads(4) |
| sessions = [] |
| for name in ("encoder", "cross_cache", "decoder_step"): |
| path = args.graph_dir / f"{name}.onnx" |
| if digest(path) != manifest["graphs"][name]["sha256"]: |
| raise SystemExit(f"Graph hash mismatch: {path}") |
| sessions.append(ort.InferenceSession(str(path), sess_options=options, providers=["CPUExecutionProvider"])) |
|
|
| |
| with (args.dataset_root / "fma/fma_metadata/tracks.csv").open() as stream: |
| reader = csv.reader(stream) |
| groups, fields = next(reader), next(reader) |
| license_index = list(zip(groups, fields)).index(("track", "license")) |
| licenses = {int(row[0]): row[license_index] for row in reader if row[0].isdigit()} |
| for row in config["examples"]: |
| if row["collection"] == "Free Music Archive": |
| assert licenses[row["track_id"]] == row["license_metadata"], row["id"] |
| assert row["license_metadata"] in {"Attribution 3.0 United States", "Attribution 3.0 International"} |
|
|
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| catalog = {"schema_version": 1, "default_example": config["default_example"], "examples": []} |
| credits = ["# Demo audio credits", "", "All clips are cropped to at most 12 seconds and converted to mono 16 kHz PCM.", "Prepared notes are Hendrix predictions, not dataset annotations.", ""] |
| hann = torch.hann_window(2048) |
| started = time.monotonic() |
| for row in config["examples"]: |
| audio_path = args.output_dir / f"{row['id']}.wav" |
| source = args.dataset_root / row["source"] |
| subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-ss", str(row["offset_seconds"]), |
| "-i", str(source), "-t", str(row["duration_seconds"]), "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(audio_path)], check=True) |
| audio, rate = sf.read(audio_path, dtype="float32") |
| assert rate == 16000 and audio.ndim == 1 |
| duration = len(audio) / rate |
| tokens, truncations = [], 0 |
| for start in range(0, len(audio), 32767): |
| window = np.zeros(32767, np.float32) |
| segment = audio[start:start+32767] |
| window[:len(segment)] = segment |
| magnitudes = torch.stft(torch.from_numpy(window), n_fft=2048, hop_length=300, |
| window=hann, return_complex=True).abs()[1:].transpose(0, 1) |
| features = torch.log(magnitudes + 1e-5) |
| features = (features - features.min()) / (features.max() - features.min() + 0.008) |
| predicted, unfinished = generate(sessions, features.unsqueeze(0).numpy()) |
| tokens.append(predicted) |
| truncations += bool(unfinished) |
| decoded = subprocess.run(["node", "--experimental-strip-types", str(ROOT / "scripts/decode_demo_tokens.mjs")], |
| input=json.dumps({"windows": tokens, "segmentSeconds": 32767/16000, "duration": duration}), |
| text=True, capture_output=True, check=True) |
| notes = json.loads(decoded.stdout) |
| payload = {"schema_version": 1, "duration": duration, "notes": notes, |
| "checkpoint_sha256": manifest["checkpoint"]["sha256"], "audio_sha256": digest(audio_path), |
| "runtime": "Fixed FP32 ONNX browser graphs, CPU, shared browser note decoder", "truncated_windows": truncations} |
| note_path = args.output_dir / f"{row['id']}.json" |
| note_path.write_text(json.dumps(payload, separators=(",", ":")) + "\n") |
| peaks = [float(np.sqrt(np.mean(chunk**2))) for chunk in np.array_split(audio, 56)] |
| maximum = max(peaks) or 1 |
| peaks = [round((value/maximum)**0.65, 3) for value in peaks] |
| public = {**row, "audio_url": f"/examples/{audio_path.name}", "notes_url": f"/examples/{note_path.name}", |
| "duration": duration, "peaks": peaks, "note_count": len(notes), "audio_sha256": digest(audio_path), |
| "notes_sha256": digest(note_path), "source_sha256": digest(source)} |
| catalog["examples"].append(public) |
| credits += [f"## {row['title']}", f"Artist/creator: {row['artist']}", |
| f"Source: [{row['collection']}]({row['source_url']}), `{row['source']}`.", |
| f"License: [{row['license_name']}]({row['license_url']}).", |
| f"Change: {row['offset_seconds']}–{row['offset_seconds']+duration:g} seconds, mono 16 kHz excerpt.", ""] |
| print(f"{row['id']}: {duration:g}s, {len(notes)} notes, {truncations} truncated windows", flush=True) |
| (args.output_dir / "catalog.json").write_text(json.dumps(catalog, indent=2) + "\n") |
| (args.output_dir / "CREDITS.md").write_text("\n".join(credits)) |
| fonts = args.output_dir.parent / "soundfonts" |
| fonts.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(args.soundfont, fonts / "GeneralUser_GS.sf2") |
| shutil.copyfile(ROOT / "deploy/huggingface/GeneralUser-LICENSE.txt", fonts / "LICENSE.txt") |
| print(f"Prepared {len(catalog['examples'])} examples in {time.monotonic()-started:.1f}s: {args.output_dir}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
