初始提交：3 款纯音频游戏 + 共享引擎 + 完整文档与资产管线
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
This commit is contained in:
SoundGame Dev
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark — wall time + peak RAM for SA3 across model sizes and clip lengths.
|
||||
|
||||
Runs each (DiT, decoder, seconds) cell as a fresh ./sa3 subprocess so peak
|
||||
RAM is isolated per run. Weights are pre-warmed via ensure_local() before
|
||||
timing starts, so HF download time is NOT charged against the benchmark.
|
||||
|
||||
Usage:
|
||||
uv run --no-project python scripts/benchmark.py
|
||||
# or
|
||||
./sa3 --help # see the regular CLI; this file is separate
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
PROJECT_DIR = SCRIPT_DIR.parent
|
||||
sys.path.insert(0, str(SCRIPT_DIR))
|
||||
from weights import ensure_local # noqa: E402
|
||||
|
||||
# ── what to benchmark ──────────────────────────────────────────────────────
|
||||
CONFIGS = [
|
||||
("sm-music", "same-s"),
|
||||
("medium", "same-l"),
|
||||
]
|
||||
SECONDS = [5, 30, 120, 380]
|
||||
PROMPT = "Impending tribal, epic orchestral buildup"
|
||||
SEED = 42
|
||||
|
||||
# Weights we need on disk before benchmarks start. T5Gemma is shared.
|
||||
PREFLIGHT_WEIGHTS = [
|
||||
"models/mlx/t5gemma_f16.npz",
|
||||
"models/mlx/dit_sm-music_f16.npz",
|
||||
"models/mlx/dit_medium_f16.npz",
|
||||
"models/mlx/same_s_decoder_f32.npz",
|
||||
"models/mlx/same_l_decoder_f32.npz",
|
||||
]
|
||||
|
||||
|
||||
# ── hardware info ──────────────────────────────────────────────────────────
|
||||
def hwinfo() -> str:
|
||||
try:
|
||||
chip = subprocess.check_output(
|
||||
["sysctl", "-n", "machdep.cpu.brand_string"], text=True
|
||||
).strip()
|
||||
except Exception:
|
||||
chip = platform.processor() or platform.machine() or "unknown"
|
||||
try:
|
||||
mem_bytes = int(subprocess.check_output(
|
||||
["sysctl", "-n", "hw.memsize"], text=True).strip())
|
||||
mem_gb = f"{mem_bytes / (1024**3):.0f} GB"
|
||||
except Exception:
|
||||
mem_gb = "?"
|
||||
return f"{chip} · {mem_gb} RAM"
|
||||
|
||||
|
||||
# ── output parsing ─────────────────────────────────────────────────────────
|
||||
_ANSI = re.compile(r"\x1b\[[0-9;]*m")
|
||||
_DONE = re.compile(
|
||||
r"done\s+([\d.]+)s\s+wall\s+→\s+([\d.]+)s\s+audio\s+→\s+"
|
||||
r"([\d.]+)×\s+realtime\s+peak\s+RAM\s+([\d.]+)\s+(MB|GB)"
|
||||
)
|
||||
|
||||
def parse_done(stdout: str) -> dict | None:
|
||||
plain = _ANSI.sub("", stdout)
|
||||
m = _DONE.search(plain)
|
||||
if not m:
|
||||
return None
|
||||
return {
|
||||
"wall_s": float(m.group(1)),
|
||||
"audio_s": float(m.group(2)),
|
||||
"realtime": float(m.group(3)),
|
||||
"peak_ram": f"{m.group(4)} {m.group(5)}",
|
||||
}
|
||||
|
||||
|
||||
# ── single benchmark cell ──────────────────────────────────────────────────
|
||||
def run_one(dit: str, decoder: str, seconds: int) -> dict | None:
|
||||
out = Path(tempfile.gettempdir()) / f"bench_{dit}_{decoder}_{seconds}s.wav"
|
||||
cmd = [
|
||||
str(PROJECT_DIR / "sa3"),
|
||||
"--prompt", PROMPT,
|
||||
"--dit", dit,
|
||||
"--decoder", decoder,
|
||||
"--seconds", str(seconds),
|
||||
"--seed", str(SEED),
|
||||
"--out", str(out),
|
||||
]
|
||||
label = f" {dit:9s} {decoder:7s} {seconds:>4}s"
|
||||
print(f"{label} → running …", flush=True)
|
||||
t0 = time.time()
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
cmd, capture_output=True, text=True, cwd=str(PROJECT_DIR),
|
||||
timeout=20 * 60, # 20-minute hard cap per cell
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f"{label} → TIMEOUT after {time.time()-t0:.0f}s")
|
||||
return None
|
||||
elapsed = time.time() - t0
|
||||
out.unlink(missing_ok=True)
|
||||
|
||||
if proc.returncode != 0:
|
||||
last_err = (proc.stderr.strip().split("\n") or ["(no stderr)"])[-1][:120]
|
||||
print(f"{label} → FAILED exit={proc.returncode} {last_err}")
|
||||
return None
|
||||
parsed = parse_done(proc.stdout)
|
||||
if not parsed:
|
||||
print(f"{label} → FAILED (could not parse 'done' line after {elapsed:.0f}s)")
|
||||
return None
|
||||
print(
|
||||
f"{label} → {parsed['wall_s']:7.2f}s wall "
|
||||
f"{parsed['realtime']:5.2f}× RT peak {parsed['peak_ram']}"
|
||||
)
|
||||
return {"dit": dit, "decoder": decoder, "seconds": seconds, **parsed}
|
||||
|
||||
|
||||
# ── ASCII table ────────────────────────────────────────────────────────────
|
||||
def render_table(rows: list[dict]) -> str:
|
||||
headers = ["model", "decoder", "seconds", "wall (s)", "×realtime", "peak RAM"]
|
||||
widths = [9, 7, 7, 9, 9, 10]
|
||||
|
||||
def hline(l, m, r):
|
||||
return l + m.join("─" * (w + 2) for w in widths) + r
|
||||
def row(cells):
|
||||
parts = []
|
||||
for i, c in enumerate(cells):
|
||||
# right-align numeric columns (idx 2-5)
|
||||
parts.append(c.rjust(widths[i]) if i >= 2 else c.ljust(widths[i]))
|
||||
return "│ " + " │ ".join(parts) + " │"
|
||||
|
||||
lines = [hline("┌", "┬", "┐"), row(headers), hline("├", "┼", "┤")]
|
||||
prev_dit = None
|
||||
for r in rows:
|
||||
# Visual separator when the model changes
|
||||
if prev_dit is not None and r["dit"] != prev_dit:
|
||||
lines.append(hline("├", "┼", "┤"))
|
||||
prev_dit = r["dit"]
|
||||
lines.append(row([
|
||||
r["dit"], r["decoder"], str(r["seconds"]),
|
||||
f"{r['wall_s']:.2f}",
|
||||
f"{r['realtime']:.2f}×",
|
||||
r["peak_ram"],
|
||||
]))
|
||||
lines.append(hline("└", "┴", "┘"))
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ── main ───────────────────────────────────────────────────────────────────
|
||||
def main() -> int:
|
||||
print()
|
||||
print("━" * 72)
|
||||
print(" Benchmark — Stable Audio 3 (MLX) on Apple Silicon")
|
||||
print("━" * 72)
|
||||
print(f" Hardware: {hwinfo()}")
|
||||
print(f" Prompt: {PROMPT!r}")
|
||||
print(f" Seed: {SEED}")
|
||||
print(f" Cells: {len(CONFIGS)} configs × {len(SECONDS)} durations = "
|
||||
f"{len(CONFIGS) * len(SECONDS)} runs")
|
||||
print()
|
||||
|
||||
print("→ Preflight: ensuring all weights are local (downloads NOT timed)")
|
||||
for rel in PREFLIGHT_WEIGHTS:
|
||||
ensure_local(rel)
|
||||
print(" done.\n")
|
||||
|
||||
print("→ Running benchmarks (each subprocess is isolated for clean peak-RAM)")
|
||||
rows: list[dict] = []
|
||||
for dit, decoder in CONFIGS:
|
||||
for s in SECONDS:
|
||||
r = run_one(dit, decoder, s)
|
||||
if r:
|
||||
rows.append(r)
|
||||
print()
|
||||
|
||||
if not rows:
|
||||
print("No successful runs — nothing to tabulate.")
|
||||
return 1
|
||||
|
||||
print("━" * 72)
|
||||
print(f" Results ({len(rows)}/{len(CONFIGS)*len(SECONDS)} cells succeeded)")
|
||||
print(f" Hardware: {hwinfo()}")
|
||||
print("━" * 72)
|
||||
print()
|
||||
print(render_table(rows))
|
||||
print()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Shared, colored "Try these commands" block.
|
||||
|
||||
Used by:
|
||||
- scripts/install.py — printed at the end of install.sh
|
||||
- scripts/sa3_mlx.py — appended to `./sa3 --help`
|
||||
|
||||
Examples render with the user's `./sa3` wrapper when present (falling back to
|
||||
`uv run python` or bare `python`), and only show entries for DiT bundles that
|
||||
the user actually has installed. Bundles not yet on disk get a "re-run
|
||||
install.sh, or just use them — weights auto-download" note at the bottom.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# This file lives in <project>/scripts/. SCRIPT_DIR points at the project
|
||||
# root (where ./sa3, ./install.sh, models/, .venv/ live).
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
# ── ANSI colours (safe no-ops when stdout isn't a TTY) ──────────────────────
|
||||
def _c(code: str) -> str:
|
||||
return code if sys.stdout.isatty() else ""
|
||||
|
||||
BOLD = _c("\033[1m")
|
||||
CYAN = _c("\033[1;36m")
|
||||
GREEN = _c("\033[1;32m")
|
||||
YELLOW = _c("\033[1;33m")
|
||||
DIM = _c("\033[2m")
|
||||
RESET = _c("\033[0m")
|
||||
|
||||
|
||||
def _have(binary: str) -> bool:
|
||||
from shutil import which
|
||||
return which(binary) is not None
|
||||
|
||||
|
||||
def _py_invocation() -> tuple[str, str, bool]:
|
||||
"""Return (command-to-print, tip-or-empty, is_wrapper).
|
||||
|
||||
Prefer the ./sa3 wrapper if present, else `uv run python scripts/sa3_mlx.py`,
|
||||
else fall back to a bare `python scripts/sa3_mlx.py`.
|
||||
"""
|
||||
wrapper = SCRIPT_DIR / "sa3"
|
||||
if wrapper.exists() and os.access(wrapper, os.X_OK):
|
||||
return "./sa3", "the ./sa3 wrapper handles .venv + uv automatically", True
|
||||
if _have("uv"):
|
||||
return "uv run python", "uv run finds .venv/ automatically — no activation needed", False
|
||||
venv_dir = SCRIPT_DIR / ".venv"
|
||||
if Path(sys.prefix).resolve() == venv_dir.resolve():
|
||||
return ".venv/bin/python", "source .venv/bin/activate # to run `python` directly", False
|
||||
return "python", "", False
|
||||
|
||||
|
||||
def print_example_commands(header: str | None = None) -> None:
|
||||
"""Print the categorized example-commands block.
|
||||
|
||||
`header` is the title line inside the rule block. Defaults to a neutral
|
||||
"Examples:" for --help; install.sh passes "✓ Install complete. …".
|
||||
"""
|
||||
from weights import bundle_status
|
||||
|
||||
py, tip, is_wrapper = _py_invocation()
|
||||
prefix = py if is_wrapper else f"{py} scripts/sa3_mlx.py"
|
||||
|
||||
def hdr(text: str) -> None:
|
||||
print(f"\n {CYAN}{text}{RESET}")
|
||||
def cmd(args: str, comment: str = "") -> None:
|
||||
line = f"{prefix} {args}"
|
||||
if comment:
|
||||
print(f" {GREEN}$ {line}{RESET} {DIM}# {comment}{RESET}")
|
||||
else:
|
||||
print(f" {GREEN}$ {line}{RESET}")
|
||||
|
||||
if header is None:
|
||||
header = f"{BOLD}Examples:{RESET}"
|
||||
|
||||
print(f"\n{BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{RESET}")
|
||||
print(f" {header}")
|
||||
print(f"{BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{RESET}")
|
||||
|
||||
if tip:
|
||||
print(f"\n {DIM}tip: {tip}{RESET}")
|
||||
|
||||
have = {name: bundle_status(name) == (4, 4) for name in ("medium", "sm-music", "sm-sfx")}
|
||||
|
||||
# ── Basic generation ─────────────────────────────────────────────
|
||||
hdr("🎵 Generate audio from a prompt")
|
||||
if have["medium"]:
|
||||
cmd('--prompt "A beautiful piano arpeggio grows into a cinematic climax" \\\n'
|
||||
f' --dit medium --decoder same-l --seconds 30 --out piano.wav',
|
||||
"highest quality (~5s for 10s clip on M1)")
|
||||
if have["sm-music"]:
|
||||
cmd('--prompt "lofi house loop, 120 BPM" \\\n'
|
||||
f' --dit sm-music --decoder same-s --seconds 15 --out lofi.wav',
|
||||
"fast music generation (~1s wall for 10s clip)")
|
||||
if have["sm-sfx"]:
|
||||
cmd('--prompt "footsteps on gravel, then a door slamming" \\\n'
|
||||
f' --dit sm-sfx --decoder same-s --seconds 8 --out sfx.wav',
|
||||
"sound-effect generation")
|
||||
|
||||
# ── Playback ─────────────────────────────────────────────────────
|
||||
hdr("▶ Play immediately after generation")
|
||||
one_dit = "medium" if have["medium"] else ("sm-music" if have["sm-music"] else "sm-sfx")
|
||||
one_dec = "same-l" if one_dit == "medium" else "same-s"
|
||||
cmd(f'--prompt "ambient drone" --dit {one_dit} --decoder {one_dec} \\\n'
|
||||
f' --seconds 10 --out drone.wav --play',
|
||||
"writes WAV + plays via afplay (Ctrl-C stops both)")
|
||||
|
||||
# ── Audio-to-audio + inpaint ─────────────────────────────────────
|
||||
hdr("🎚️ Audio-to-audio & inpainting (requires an input WAV)")
|
||||
cmd(f'--prompt "jazz fusion with electric piano" --dit {one_dit} --decoder {one_dec} \\\n'
|
||||
f' --init-audio funk.wav --init-noise-level 0.7 --out funk_jazz.wav',
|
||||
"variation: 0.4-0.8 typical, higher = more change")
|
||||
cmd(f'--prompt "explosive drum break" --dit {one_dit} --decoder {one_dec} \\\n'
|
||||
f' --init-audio funk.wav --inpaint-range "4,7" --out funk_drums.wav',
|
||||
"regenerate seconds 4-7, keep rest")
|
||||
|
||||
# ── CFG & negative prompts ───────────────────────────────────────
|
||||
hdr("🎯 Steer with CFG + negative prompts")
|
||||
cmd(f'--prompt "ambient drone" --cfg 3.0 \\\n'
|
||||
f' --negative-prompt "drums, vocals, distortion" \\\n'
|
||||
f' --dit {one_dit} --decoder {one_dec} --out clean_drone.wav',
|
||||
"cfg > 1.0 toward prompt, neg pushes away")
|
||||
|
||||
# Bundles not installed (offer the on-demand path)
|
||||
missing = [name for name, ok in have.items() if not ok]
|
||||
if missing:
|
||||
print(f"\n {YELLOW}note:{RESET} bundles not installed: {', '.join(missing)}.")
|
||||
print(f" Re-run {BOLD}./install.sh{RESET} to pick them up, or just use them in"
|
||||
f" {BOLD}./sa3{RESET} —")
|
||||
print(f" weights auto-download from HuggingFace on first use.")
|
||||
|
||||
print(f"\n{BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{RESET}\n")
|
||||
@@ -0,0 +1,90 @@
|
||||
"""Convert a Stable Audio 3 *-base torch checkpoint → the MLX training npz.
|
||||
|
||||
Torch-free (safetensors.numpy). Produces `dit_<model>-base_f16.npz` — the DiT
|
||||
weights + baked conditioner that `lora_train_mlx.py --dit-weights` expects.
|
||||
LoRA training uses the BASE checkpoint (not the shipped ARC weights inference
|
||||
uses); these npz are auto-downloaded by weights.py, or reproduced here from the
|
||||
HF `stabilityai/stable-audio-3-*-base` repos.
|
||||
|
||||
python scripts/export_base_npz.py \
|
||||
--base-dir <sa3-*-base> # dir with model.safetensors + model_config.json
|
||||
--output models/mlx/dit_<model>-base_f16.npz
|
||||
[--validate-against <existing npz>] # assert bit-parity with a known-good npz
|
||||
|
||||
Conversion (identical to the shipped inference-npz layout): strip the DiT
|
||||
wrapper prefix `model.model.`, transpose the two Conv1d weights
|
||||
[out,in,k]→[out,k,in], rename RMSNorm `.gamma`→`.weight`, remap
|
||||
`.to_local_embed.{0,2}`→`.seq.{0,2}`, and fold the 3 conditioner tensors under
|
||||
the `cond.` prefix. Same for small (sm-music/sm-sfx) and medium.
|
||||
"""
|
||||
import argparse
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
from safetensors.numpy import load_file
|
||||
|
||||
_DIT_PREFIX = "model.model."
|
||||
_CONV_WEIGHTS = ("preprocess_conv.weight", "postprocess_conv.weight")
|
||||
_COND_MAP = {
|
||||
"conditioner.conditioners.prompt.padding_embedding": "cond.padding_embedding",
|
||||
"conditioner.conditioners.seconds_total.embedder.embedding.1.weight": "cond.seconds_total_weight",
|
||||
"conditioner.conditioners.seconds_total.embedder.embedding.1.bias": "cond.seconds_total_bias",
|
||||
}
|
||||
|
||||
|
||||
def convert(base_dir):
|
||||
sd = load_file(f"{base_dir}/model.safetensors")
|
||||
out = {}
|
||||
for k, arr in sd.items():
|
||||
if k in _COND_MAP:
|
||||
out[_COND_MAP[k]] = mx.array(arr.astype(np.float32))
|
||||
continue
|
||||
if not k.startswith(_DIT_PREFIX):
|
||||
continue # pretransform / T5 backbone — not part of the DiT npz
|
||||
sk = k[len(_DIT_PREFIX):]
|
||||
a = arr.astype(np.float32)
|
||||
if sk in _CONV_WEIGHTS: # Conv1d [out,in,k] -> MLX [out,k,in]
|
||||
out[sk] = mx.array(a.transpose(0, 2, 1))
|
||||
elif sk.endswith(".gamma"): # RMSNorm .gamma -> .weight
|
||||
out[sk[: -len(".gamma")] + ".weight"] = mx.array(a)
|
||||
else:
|
||||
sk = sk.replace(".to_local_embed.0.", ".to_local_embed.seq.0.")
|
||||
sk = sk.replace(".to_local_embed.2.", ".to_local_embed.seq.2.")
|
||||
out[sk] = mx.array(a)
|
||||
return out
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--base-dir", required=True)
|
||||
ap.add_argument("--output", required=True)
|
||||
ap.add_argument("--validate-against", default=None,
|
||||
help="Assert key-set + value parity with a known-good npz.")
|
||||
args = ap.parse_args()
|
||||
|
||||
wd = convert(args.base_dir)
|
||||
print(f"converted {len(wd)} keys ({sum(1 for k in wd if k.startswith('cond'))} conditioner)")
|
||||
|
||||
if args.validate_against:
|
||||
ref = dict(mx.load(args.validate_against))
|
||||
extra, missing = set(wd) - set(ref), set(ref) - set(wd)
|
||||
assert not extra and not missing, f"key mismatch: extra={extra}, missing={missing}"
|
||||
worst = max(
|
||||
float(mx.max(mx.abs(wd[k].astype(mx.float32) - ref[k].astype(mx.float32))))
|
||||
for k in ref
|
||||
)
|
||||
print(f"validate: {len(ref)} keys match, max|Δ| = {worst:.2e}")
|
||||
# Tolerance is fp16-scale: the source ckpt is fp32 and both npz are fp16,
|
||||
# so a handful of the specially-transformed keys (conv transpose, the
|
||||
# to_local_embed rename) can differ by ~1 fp16 ULP (~4e-3) depending on
|
||||
# the rounding path. A real structural bug (wrong transpose/rename) gives
|
||||
# an O(0.1-1) diff or a shape/key mismatch — both caught well below this.
|
||||
assert worst < 1e-2, "values diverge from the reference npz (structural bug)"
|
||||
|
||||
mx.savez(args.output, **{k: v.astype(mx.float16) for k, v in wd.items()})
|
||||
print(f"wrote {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,157 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Internal install plumbing — do NOT invoke directly. Use ../install.sh from the project root.
|
||||
|
||||
Called by install.sh after it has set up uv + .venv. Responsibilities:
|
||||
1. (Optionally) pip-install requirements.txt — skipped when INSTALL_SKIP_PIP=1
|
||||
is set (install.sh sets this since uv already handled deps).
|
||||
2. Ask which DiT bundles to download, then fetch them from HuggingFace.
|
||||
3. Print example commands the user can copy-paste.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# This file lives in <project>/scripts/. SCRIPT_DIR points at the project
|
||||
# root (where models/, requirements.txt, .venv/, sa3 wrapper live).
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent)) # `from weights import …`
|
||||
|
||||
MIN_PY = (3, 10)
|
||||
SKIP_PIP = os.environ.get("INSTALL_SKIP_PIP") == "1"
|
||||
|
||||
|
||||
def step(msg: str) -> None:
|
||||
print(f"\n\033[1;36m→ {msg}\033[0m")
|
||||
|
||||
|
||||
def check_environment() -> None:
|
||||
"""Bail out early with a clear message if the interpreter is wrong."""
|
||||
py = sys.version_info
|
||||
if (py.major, py.minor) < MIN_PY:
|
||||
print(
|
||||
f"\n\033[1;31merror\033[0m: this script is running under Python "
|
||||
f"{py.major}.{py.minor}, but MLX requires Python "
|
||||
f"{MIN_PY[0]}.{MIN_PY[1]}+.\n\n"
|
||||
f"Re-run install.py with a Python {MIN_PY[0]}.{MIN_PY[1]}+ interpreter, e.g.:\n"
|
||||
f" /path/to/python3.11 install.py\n\n"
|
||||
f"On macOS you can get one via:\n"
|
||||
f" brew install python@3.11\n"
|
||||
f" /opt/homebrew/bin/python3.11 install.py\n"
|
||||
f"or with pyenv:\n"
|
||||
f" pyenv install 3.11.10 && pyenv local 3.11.10 && python install.py\n",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
if platform.system() != "Darwin" or platform.machine() != "arm64":
|
||||
print(
|
||||
f"\n\033[1;33mwarning\033[0m: this stack is Apple-Silicon-only "
|
||||
f"(MLX is Metal-backed). Detected {platform.system()}/{platform.machine()} "
|
||||
f"— pip install will probably fail when it tries to fetch MLX.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
def pip_install_requirements() -> None:
|
||||
if SKIP_PIP:
|
||||
step("Dependencies already installed by install.sh (uv) — skipping pip step")
|
||||
return
|
||||
step("Installing Python dependencies")
|
||||
req = SCRIPT_DIR / "requirements.txt"
|
||||
|
||||
# Prefer pip; fall back to uv if pip isn't available in this interpreter
|
||||
# (e.g. user is running install.py directly inside a uv-created venv
|
||||
# that wasn't seeded with pip).
|
||||
pip_available = subprocess.run(
|
||||
[sys.executable, "-m", "pip", "--version"],
|
||||
capture_output=True,
|
||||
).returncode == 0
|
||||
|
||||
if pip_available:
|
||||
cmd = [sys.executable, "-m", "pip", "install", "-r", str(req)]
|
||||
elif _have("uv"):
|
||||
cmd = ["uv", "pip", "install", "--python", sys.executable, "-r", str(req)]
|
||||
else:
|
||||
print(
|
||||
f"\n\033[1;31merror\033[0m: neither pip nor uv is available in this "
|
||||
f"interpreter ({sys.executable}).\n"
|
||||
f"Install pip with `python -m ensurepip --upgrade`, or use ./install.sh "
|
||||
f"which sets everything up via uv.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
print(f" $ {' '.join(cmd)}")
|
||||
subprocess.check_call(cmd)
|
||||
|
||||
|
||||
def _have(binary: str) -> bool:
|
||||
"""True if the given binary is on PATH."""
|
||||
from shutil import which
|
||||
return which(binary) is not None
|
||||
|
||||
|
||||
def main() -> None:
|
||||
import argparse
|
||||
ap = argparse.ArgumentParser(
|
||||
description="SA3 MLX install — non-interactive bundle downloader + post-install help.",
|
||||
epilog="Normally invoked via ../install.sh, not directly.",
|
||||
)
|
||||
ap.add_argument("--download", default="",
|
||||
metavar="BUNDLES",
|
||||
help="Comma-separated list of bundles to pre-download "
|
||||
"(medium, sm-music, sm-sfx). Without this flag, "
|
||||
"nothing is downloaded — sa3_mlx.py will fetch any "
|
||||
"missing weights from HuggingFace on first use.")
|
||||
cli = ap.parse_args()
|
||||
|
||||
check_environment()
|
||||
pip_install_requirements()
|
||||
|
||||
# Import after pip install so a fresh checkout works.
|
||||
from weights import DIT_BUNDLES, SHARED, BUNDLE_SIZES, bundle_status, ensure_local
|
||||
|
||||
step("Current weights status")
|
||||
for name in DIT_BUNDLES:
|
||||
present, total = bundle_status(name)
|
||||
mark = "✓" if present == total else " "
|
||||
print(f" [{mark}] {name:9s} {present}/{total} files present ({BUNDLE_SIZES[name]})")
|
||||
|
||||
chosen: list[str] = []
|
||||
if cli.download.strip():
|
||||
chosen = [b.strip() for b in cli.download.split(",") if b.strip()]
|
||||
unknown = [b for b in chosen if b not in DIT_BUNDLES]
|
||||
if unknown:
|
||||
print(f"\n\033[1;31merror\033[0m: unknown bundle(s): {', '.join(unknown)}. "
|
||||
f"Choices: {', '.join(DIT_BUNDLES)}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
if chosen:
|
||||
step(f"Downloading {len(chosen)} bundle(s): {', '.join(chosen)}")
|
||||
seen: set[str] = set()
|
||||
for name in chosen:
|
||||
print(f"\n[{name}]")
|
||||
items = DIT_BUNDLES[name] + SHARED
|
||||
for rel, _hf in items:
|
||||
if rel in seen:
|
||||
continue
|
||||
seen.add(rel)
|
||||
ensure_local(rel)
|
||||
else:
|
||||
step("No --download set — weights will auto-download on first ./sa3 use")
|
||||
print(f" To pre-download instead, pass: {sys.executable.split('/')[-1]} install.py --download medium")
|
||||
print(f" or: ./install.sh --download medium,sm-music")
|
||||
|
||||
from examples import print_example_commands, BOLD, GREEN, RESET
|
||||
print_example_commands(header=f"{BOLD}{GREEN}✓ Install complete.{RESET} Try these commands:")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
print("\nAborted.")
|
||||
sys.exit(130)
|
||||
@@ -0,0 +1,884 @@
|
||||
"""SA3 LoRA training on Apple Silicon — pure MLX, underfit conventions.
|
||||
|
||||
The MLX counterpart of underfit's raw-PyTorch training loop
|
||||
(github.com/dada-bots/underfit, underfit/training/loop.py), built on the
|
||||
trainable adapters in models/defs/lora.py. Every training convention and
|
||||
default mirrors underfit (see ../TRAINING_CONVENTIONS.md):
|
||||
|
||||
- rectified-flow velocity target (noise − x), signal-only masked MSE,
|
||||
per-sample-then-mean reduction
|
||||
- uniform timestep sampler + the SA3 models' "full" distribution shift
|
||||
(min_length 256, max_length 4096)
|
||||
- CFG dropout 0.1: per-sample, the whole cross-attention conditioning
|
||||
(prompt + seconds token) is zeroed; global_cond is kept
|
||||
- AdamW, single param group, torch defaults (betas 0.9/0.999, eps 1e-8,
|
||||
weight decay 0.0), no schedule, no clipping; LR has NO default and must
|
||||
be supplied (--lr)
|
||||
- adapters train in fp32 over a frozen fp16 base
|
||||
- step-driven loop; checkpoints every --checkpoint-every (1000) steps and
|
||||
at the end, saved BEFORE anything else can fail, named
|
||||
{run_label}-step={step}-epoch={epoch}.safetensors with lora_config
|
||||
metadata {rank, alpha, adapter_type, include, exclude, step, epoch,
|
||||
base_model}
|
||||
- resume offsets: explicit flags > checkpoint metadata > filename tokens
|
||||
- per-step loss_by_timestep.bin telemetry (struct "Iff": step, t_mean,
|
||||
loss_mean), flushed every 10 steps
|
||||
|
||||
Dataset: pre-encoded latents from scripts/pre_encode_mlx.py (npy + json
|
||||
pairs), cropped to --latent-crop-length with random crop, tiny datasets
|
||||
oversampled with replacement (batch_size*100 draws/epoch) — the underfit
|
||||
overfitting workflow. Prompts are built per-sample with underfit's tag
|
||||
augmentation. seconds_total stays the FULL source duration after cropping
|
||||
(deliberate underfit convention).
|
||||
|
||||
Performance:
|
||||
- the train step (loss + grad + optimizer update) is mx.compile'd by
|
||||
default (MLX docs "Compile > Training graphs" pattern; --no-compile for
|
||||
the eager path). Batch shapes are constant, so it compiles once.
|
||||
- T5Gemma prompt conditioning is cached per exact prompt string (per-run,
|
||||
in-memory; --no-t5-cache to disable) — the tiny-dataset workflow repeats
|
||||
a handful of prompts, so the hit rate is ≈100% after the first epoch.
|
||||
- --grad-checkpoint recomputes each transformer block's activations in the
|
||||
backward pass (mlx-lm style mx.checkpoint) — memory for big crops on
|
||||
small Macs, at the cost of roughly one extra forward.
|
||||
- wired memory limit is raised to the device's recommended working-set
|
||||
size at startup (mx.set_wired_limit).
|
||||
|
||||
Usage:
|
||||
python scripts/lora_train_mlx.py --dit sm-music \
|
||||
--latents-dir output/latents/my-set --lr 1e-4 \
|
||||
--name my-lora --save-dir output/runs --max-steps 2000
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import struct
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
SCRIPTS_DIR = Path(__file__).resolve().parent
|
||||
REPO = SCRIPTS_DIR.parent
|
||||
sys.path.insert(0, str(REPO))
|
||||
sys.path.insert(0, str(SCRIPTS_DIR))
|
||||
|
||||
import mlx.core as mx # noqa: E402
|
||||
import mlx.nn as nn # noqa: E402
|
||||
import mlx.optimizers as optim # noqa: E402
|
||||
from mlx.utils import tree_flatten # noqa: E402
|
||||
|
||||
from sa3_mlx import DIT_CHOICES, load_dit # noqa: E402
|
||||
from weights import ensure_local # noqa: E402
|
||||
from sa3_mlx import T5GEMMA_NPZ_REL # noqa: E402
|
||||
from models.defs.sa3_pipeline import ( # noqa: E402
|
||||
apply_prompt_padding, load_conditioner_from_npz,
|
||||
)
|
||||
from models.defs.t5gemma_mlx import T5Gemma # noqa: E402
|
||||
from models.defs.training import ( # noqa: E402
|
||||
rectified_flow_loss, sample_training_timesteps, shift_training_timesteps,
|
||||
)
|
||||
from models.defs.lora import ( # noqa: E402
|
||||
TrainableSecondsEmbedder, inject_from_lora_config, iter_trainable_lora_layers,
|
||||
load_lora_checkpoint, load_trainable_lora_state, save_lora_checkpoint,
|
||||
underfit_lora_config,
|
||||
)
|
||||
from models.defs import demo_mlx as demo # noqa: E402
|
||||
from models.defs.latent_dataset import ( # noqa: E402
|
||||
PreEncodedLatentDataset, iterate_batches,
|
||||
)
|
||||
|
||||
# underfit registry conventions per model
|
||||
BASE_MODEL_NAMES = {"sm-music": "sa3-sm-music", "sm-sfx": "sa3-sm-sfx",
|
||||
"medium": "sa3-medium"}
|
||||
LATENT_CROP_DEFAULTS = {"sm-music": 1300, "sm-sfx": 1300, "medium": 4096}
|
||||
DIST_SHIFT_DEFAULT = {"shift_type": "full",
|
||||
"options": {"min_length": 256, "max_length": 4096}}
|
||||
T5_CACHE_CAP = 512 # prompt-conditioning cache entries (oldest evicted first)
|
||||
|
||||
|
||||
def _rm(path):
|
||||
"""Remove a file, ignoring a missing one (temp-ckpt cleanup)."""
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def grad_checkpoint(layer):
|
||||
"""
|
||||
Update all instances of type(layer) to use gradient checkpointing.
|
||||
|
||||
Verbatim port of mlx-lm's mlx_lm/tuner/trainer.py:grad_checkpoint —
|
||||
activations of every instance of the layer's class are recomputed during
|
||||
the backward pass instead of being kept alive, trading ~one extra forward
|
||||
for a much smaller peak working set.
|
||||
"""
|
||||
fn = type(layer).__call__
|
||||
|
||||
def checkpointed_fn(model, *args, **kwargs):
|
||||
def inner_fn(params, *args, **kwargs):
|
||||
model.update(params)
|
||||
return fn(model, *args, **kwargs)
|
||||
|
||||
return mx.checkpoint(inner_fn)(model.trainable_parameters(), *args, **kwargs)
|
||||
|
||||
type(layer).__call__ = checkpointed_fn
|
||||
|
||||
|
||||
class TrainBundle(nn.Module):
|
||||
"""DiT + (optional) trainable seconds conditioner under one module root so
|
||||
a single value_and_grad covers every adapter parameter. The bundle's
|
||||
trainable_parameters() are the LoRA params only (bases frozen at inject)."""
|
||||
|
||||
def __init__(self, dit, secs):
|
||||
super().__init__()
|
||||
self.dit = dit
|
||||
self.secs = secs # TrainableSecondsEmbedder or None
|
||||
|
||||
|
||||
def parse_args():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--dit", required=True, choices=list(DIT_CHOICES.keys()))
|
||||
ap.add_argument("--dit-weights", default=None,
|
||||
help="Override the DiT weight file — REQUIRED in practice for "
|
||||
"training parity with underfit: train on the BASE "
|
||||
"(rectified-flow) checkpoint, not the shipped ARC "
|
||||
"(rf_denoiser) weights that inference uses. Convert the "
|
||||
"HF *-base model.safetensors to npz first (441-key "
|
||||
"layout; see TRAINING_CONVENTIONS.md).")
|
||||
ap.add_argument("--latents-dir", required=True,
|
||||
help="Root of pre-encoded npy+json pairs (scripts/pre_encode_mlx.py)")
|
||||
ap.add_argument("--lr", type=float, required=True,
|
||||
help="Learning rate (underfit has NO default — must be supplied)")
|
||||
ap.add_argument("--name", default="underfit-run", help="Run name (defaults.ini)")
|
||||
ap.add_argument("--save-dir", default=str(REPO / "output" / "runs"))
|
||||
ap.add_argument("--batch-size", type=int, default=1)
|
||||
ap.add_argument("--seed", type=int, default=42)
|
||||
ap.add_argument("--max-steps", type=int, default=10_000_000_000,
|
||||
help="Absolute global-step target (resume-aware)")
|
||||
ap.add_argument("--checkpoint-every", type=int, default=1000)
|
||||
# adapter config — underfit dashboard defaults
|
||||
ap.add_argument("--adapter-type", default="dora-rows")
|
||||
ap.add_argument("--rank", type=int, default=16)
|
||||
ap.add_argument("--alpha", type=float, default=None, help="Default: = rank")
|
||||
ap.add_argument("--include", default=None,
|
||||
help="Comma-separated substring filters (underfit convention)")
|
||||
ap.add_argument("--exclude", default=None,
|
||||
help="Comma-separated; default = underfit's dashboard exclude list")
|
||||
ap.add_argument("--no-conditioner-lora", action="store_true",
|
||||
help="Skip the seconds-conditioner adapter (underfit adapts it "
|
||||
"by default — e.g. the plini checkpoint carries its delta)")
|
||||
ap.add_argument("--bora-mode", default="speed", choices=["speed", "memory"],
|
||||
help="bora/bora-xs forward: 'speed' (default) caches W0² per "
|
||||
"adapted layer (+1 weight copy) for the reformulated "
|
||||
"no-materialize forward; 'memory' keeps the original "
|
||||
"full-weight forward with no cache. Ignored for other "
|
||||
"adapter types.")
|
||||
# data
|
||||
ap.add_argument("--latent-crop-length", type=int, default=None,
|
||||
help="Default per model: 1300 (sm-*) / 4096 (medium)")
|
||||
ap.add_argument("--random-crop", action=argparse.BooleanOptionalAction, default=True)
|
||||
ap.add_argument("--prompt-config", default=None,
|
||||
help="JSON file with underfit prompt_config (trigger, tag options…)")
|
||||
# schedule / conditioning
|
||||
ap.add_argument("--timestep-sampler", default="uniform",
|
||||
choices=["uniform", "logit_normal", "trunc_logit_normal",
|
||||
"log_snr", "log_snr_uniform"])
|
||||
ap.add_argument("--dist-shift", default="full",
|
||||
choices=["none", "full", "flux", "logsnr"])
|
||||
ap.add_argument("--cfg-dropout-prob", type=float, default=0.1)
|
||||
ap.add_argument("--use-effective-length", action=argparse.BooleanOptionalAction,
|
||||
default=False,
|
||||
help="Shift timesteps by the EFFECTIVE latent length "
|
||||
"ceil(int(seconds_total*44100)/4096) per sample instead "
|
||||
"of the crop length (underfit model flag "
|
||||
"use_effective_length_for_schedule; True in the SA3 "
|
||||
"templates). Only affects the 'full'/'flux'/'logsnr' shift.")
|
||||
# optimizer (underfit reads these from optimizer_configs; torch AdamW defaults)
|
||||
ap.add_argument("--beta1", type=float, default=0.9)
|
||||
ap.add_argument("--beta2", type=float, default=0.999,
|
||||
help="AdamW β2 (torch default 0.999; SA3 templates use 0.95)")
|
||||
ap.add_argument("--eps", type=float, default=1e-8)
|
||||
ap.add_argument("--weight-decay", type=float, default=0.0,
|
||||
help="AdamW decoupled weight decay (torch default 0.0; SA3 "
|
||||
"templates use 0.01)")
|
||||
ap.add_argument("--gradient-clip-val", type=float, default=0.0,
|
||||
help="Global grad-norm clip over the adapter params before "
|
||||
"the optimizer step (0 = off). underfit's dashboard "
|
||||
"passes 1.0; the grad_norm metric is then post-clip, "
|
||||
"matching underfit's torch loop.")
|
||||
# LR schedule (underfit's optimizer_configs.diffusion.scheduler; InverseLR
|
||||
# with warmup — at step 0 the SA3 template LR is ~5e-7, not the nominal lr)
|
||||
ap.add_argument("--lr-scheduler", default="none", choices=["none", "inverse"],
|
||||
help="'inverse' = SAT-dev InverseLR: "
|
||||
"lr(step) = (1-warmup^(step+1)) * "
|
||||
"max(final, lr*(1+step/inv_gamma)^-power)")
|
||||
ap.add_argument("--inv-gamma", type=float, default=1e6)
|
||||
ap.add_argument("--lr-power", type=float, default=0.5)
|
||||
ap.add_argument("--lr-warmup", type=float, default=0.0,
|
||||
help="InverseLR exponential-warmup base in [0,1) (template 0.995)")
|
||||
ap.add_argument("--lr-final", type=float, default=0.0)
|
||||
# performance
|
||||
ap.add_argument("--compile", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="mx.compile the train step (loss+grad+optimizer update; "
|
||||
"MLX 'Compile > Training graphs' pattern). "
|
||||
"--no-compile restores the eager path.")
|
||||
ap.add_argument("--grad-checkpoint", action="store_true",
|
||||
help="Gradient-checkpoint every transformer block (mlx-lm "
|
||||
"style mx.checkpoint): activations are recomputed in "
|
||||
"the backward pass — much lower peak memory for big "
|
||||
"crops, ~one extra forward of compute")
|
||||
ap.add_argument("--t5-cache", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="Cache the padded T5Gemma prompt conditioning per exact "
|
||||
f"prompt string (per-run, in-memory, max {T5_CACHE_CAP} "
|
||||
"entries)")
|
||||
# resume
|
||||
ap.add_argument("--lora-ckpt-path", default=None)
|
||||
ap.add_argument("--step-offset", type=int, default=None)
|
||||
ap.add_argument("--epoch-offset", type=int, default=None)
|
||||
# demos (underfit convention: RF Euler inference on the base model + trained
|
||||
# LoRA, decode → mp3, baseline at step 0 then every --demo-every steps)
|
||||
ap.add_argument("--demo-every", type=int, default=0,
|
||||
help="Generate demos every N steps (0 = off). A baseline "
|
||||
"set is also rendered at step 0.")
|
||||
ap.add_argument("--demo-config", default=None,
|
||||
help="JSON list of demo entries. Each: {prompt, cfg=7, "
|
||||
"seed=0, steps=50, lora_strength?, lora_interval_max?}.")
|
||||
ap.add_argument("--demo-decoder", default=None,
|
||||
choices=["same-s", "same-l"],
|
||||
help="Decoder for demos (default: the DiT's default_decoder).")
|
||||
ap.add_argument("--demo-dir", default=None,
|
||||
help="Where to write demo_<i>_<step>.mp3 (default: cwd, "
|
||||
"which the dashboard sets to <run>/demos/).")
|
||||
ap.add_argument("--arc-weights", default=None,
|
||||
help="DiT weights for ARC demos (entries with arc=true): the "
|
||||
"SHIPPED rf_denoiser npz, NOT the training base. Default "
|
||||
"auto-downloads models/mlx/dit_<dit>_f16.npz. The trained "
|
||||
"LoRA is merged into a fresh copy of these weights and "
|
||||
"sampled with the pingpong integrator.")
|
||||
return ap.parse_args()
|
||||
|
||||
|
||||
def resolve_offsets(args, ckpt_config) -> tuple[int, int]:
|
||||
"""Underfit's resume ladder: explicit flags > checkpoint metadata >
|
||||
filename step=/epoch= tokens > zero."""
|
||||
if args.step_offset is not None or args.epoch_offset is not None:
|
||||
return int(args.step_offset or 0), int(args.epoch_offset or 0)
|
||||
if ckpt_config:
|
||||
step = ckpt_config.get("step")
|
||||
epoch = ckpt_config.get("epoch")
|
||||
if step is not None:
|
||||
return int(step), int(epoch or 0)
|
||||
if args.lora_ckpt_path:
|
||||
name = Path(args.lora_ckpt_path).name
|
||||
ms = re.search(r"step=(\d+)", name)
|
||||
me = re.search(r"epoch=(\d+)", name)
|
||||
if ms:
|
||||
return int(ms.group(1)), int(me.group(1)) if me else 0
|
||||
return 0, 0
|
||||
|
||||
|
||||
def build_conditioning(t5, padding_emb, prompts, max_len=256,
|
||||
cache=None, cache_stats=None):
|
||||
"""Frozen prompt-side conditioning: T5Gemma embeddings with the learned
|
||||
padding embedding applied. Returns [B, 256, 768] fp32 (the trainable
|
||||
seconds token is concatenated later, inside the grad scope).
|
||||
|
||||
When ``cache`` (a dict) is given, the padded embedding is cached per
|
||||
EXACT prompt string. The cache is per-run only (in-memory, never
|
||||
persisted) and stores evaluated [1, max_len, 768] fp32 mx arrays; it is
|
||||
capped at T5_CACHE_CAP entries, evicting oldest-inserted first. The
|
||||
tiny-dataset underfit workflow repeats a handful of prompts, so after
|
||||
the first epoch the T5Gemma forward is skipped entirely (~100% hits)."""
|
||||
prompts = list(prompts)
|
||||
if cache is None:
|
||||
embeds, mask = t5.encode(prompts, max_len=max_len)
|
||||
mx.eval(embeds, mask)
|
||||
padded = apply_prompt_padding(embeds.astype(mx.float32), mask,
|
||||
padding_emb.astype(mx.float32))
|
||||
return mx.stop_gradient(padded)
|
||||
|
||||
missing = [p for p in dict.fromkeys(prompts) if p not in cache]
|
||||
if cache_stats is not None:
|
||||
n_miss = sum(1 for p in prompts if p in set(missing))
|
||||
cache_stats["misses"] += n_miss
|
||||
cache_stats["hits"] += len(prompts) - n_miss
|
||||
if missing:
|
||||
embeds, mask = t5.encode(missing, max_len=max_len)
|
||||
padded = apply_prompt_padding(embeds.astype(mx.float32), mask,
|
||||
padding_emb.astype(mx.float32))
|
||||
mx.eval(padded)
|
||||
for i, p in enumerate(missing):
|
||||
while len(cache) >= T5_CACHE_CAP:
|
||||
cache.pop(next(iter(cache))) # evict oldest
|
||||
row = padded[i:i + 1]
|
||||
mx.eval(row)
|
||||
cache[p] = row
|
||||
print(f" t5-cache: encoded {len(missing)} new prompt(s) "
|
||||
f"({len(cache)} cached)")
|
||||
return mx.stop_gradient(mx.concatenate([cache[p] for p in prompts],
|
||||
axis=0))
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
np_rng = np.random.default_rng(args.seed)
|
||||
mx.random.seed(args.seed)
|
||||
|
||||
# Let Metal wire (pin) up to the recommended working set — avoids paging
|
||||
# stalls when the training working set approaches physical memory.
|
||||
if hasattr(mx, "set_wired_limit") and hasattr(mx, "device_info"):
|
||||
wired = mx.device_info().get("max_recommended_working_set_size")
|
||||
if wired:
|
||||
mx.set_wired_limit(int(wired))
|
||||
|
||||
alpha = float(args.alpha) if args.alpha is not None else float(args.rank)
|
||||
lora_config = underfit_lora_config()
|
||||
lora_config.update({"adapter_type": args.adapter_type, "rank": args.rank,
|
||||
"alpha": alpha})
|
||||
if args.include is not None:
|
||||
lora_config["include"] = [s.strip() for s in args.include.split(",") if s.strip()]
|
||||
if args.exclude is not None:
|
||||
lora_config["exclude"] = [s.strip() for s in args.exclude.split(",") if s.strip()]
|
||||
|
||||
crop_len = args.latent_crop_length or LATENT_CROP_DEFAULTS[args.dit]
|
||||
base_model = BASE_MODEL_NAMES[args.dit]
|
||||
|
||||
# ── run dirs (underfit layout: <save>/<name>/<uuid8>/checkpoints) ────────
|
||||
run_label = re.sub(r"-\d{14}$", "", args.name) if args.name else None
|
||||
session = uuid.uuid4().hex[:8]
|
||||
ckpt_dir = Path(args.save_dir) / args.name / session / "checkpoints"
|
||||
ckpt_dir.mkdir(parents=True, exist_ok=True)
|
||||
print(f"Engine: MLX · device: {mx.default_device()}")
|
||||
print(f"run: {args.name} session {session}")
|
||||
print(f"checkpoints → {ckpt_dir}")
|
||||
|
||||
# ── dataset ──────────────────────────────────────────────────────────────
|
||||
prompt_config = None
|
||||
if args.prompt_config:
|
||||
prompt_config = json.loads(Path(args.prompt_config).read_text())
|
||||
dataset = PreEncodedLatentDataset(
|
||||
args.latents_dir, crop_len, random_crop=args.random_crop,
|
||||
prompt_config=prompt_config, seed=args.seed)
|
||||
print(f"dataset: {len(dataset)} latent file(s), crop {crop_len} latents"
|
||||
+ (" (oversampling with replacement — tiny dataset)"
|
||||
if len(dataset) < args.batch_size else ""))
|
||||
# Batches per epoch — mirrors iterate_batches: tiny sets (< batch) oversample
|
||||
# to batch_size*100 draws (→ 100 batches); else ceil(n / batch). Used as the
|
||||
# tqdm total so the progress bar (and the dashboard's collapse signature)
|
||||
# matches underfit's torch dataloader.
|
||||
_n_idx = (args.batch_size * 100 if len(dataset) < args.batch_size
|
||||
else len(dataset))
|
||||
steps_per_epoch = max(1, (_n_idx + args.batch_size - 1) // args.batch_size)
|
||||
|
||||
# ── models ───────────────────────────────────────────────────────────────
|
||||
# LoRA training uses the BASE checkpoint (rectified_flow), NOT the shipped
|
||||
# ARC weights inference uses. Default to the base-model npz (auto-downloaded
|
||||
# from HF via weights.TRAINING_BASE); --dit-weights overrides with a path.
|
||||
t0 = time.time()
|
||||
import importlib
|
||||
if args.dit_weights:
|
||||
base_path = args.dit_weights
|
||||
else:
|
||||
base_path = str(ensure_local(f"models/mlx/dit_{args.dit}-base_f16.npz"))
|
||||
mod = importlib.import_module(DIT_CHOICES[args.dit]["loader"])
|
||||
dit_model = mod.load_dit(base_path, T_lat=crop_len, dtype=mx.float16,
|
||||
compile_=False)
|
||||
cond_src = base_path
|
||||
print(f"DiT loaded ({time.time()-t0:.1f}s, fp16 base, T_lat={crop_len})")
|
||||
t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL))) # frozen
|
||||
padding_emb, secs_embedder = load_conditioner_from_npz(cond_src, prefix="cond.")
|
||||
|
||||
# ── inject adapters ─────────────────────────────────────────────────────
|
||||
report, saved_config = inject_from_lora_config(
|
||||
dit_model, lora_config, checkpoint_prefix="model.",
|
||||
bora_mode=args.bora_mode)
|
||||
print(f"lora: {report.layer_count} DiT layer(s), {report.adapter_type}, "
|
||||
f"rank {args.rank}, alpha {alpha:g} "
|
||||
f"({report.trainable_parameters/1e6:.2f}M trainable)")
|
||||
|
||||
if args.grad_checkpoint:
|
||||
# Patch the TransformerBlock class (both DiT defs expose the blocks at
|
||||
# dit.transformer.layers) so every block recomputes in backward.
|
||||
grad_checkpoint(dit_model.transformer.layers[0])
|
||||
print(f"grad-checkpoint: on "
|
||||
f"({len(dit_model.transformer.layers)} transformer blocks)")
|
||||
|
||||
secs_module = None
|
||||
if not args.no_conditioner_lora:
|
||||
secs_module = TrainableSecondsEmbedder(secs_embedder.W, secs_embedder.b)
|
||||
try:
|
||||
cond_report, _ = inject_from_lora_config(
|
||||
secs_module, lora_config,
|
||||
checkpoint_prefix="conditioners.seconds_total.",
|
||||
bora_mode=args.bora_mode)
|
||||
print(f"lora: conditioner seconds embedder adapted "
|
||||
f"({cond_report.layer_count} layer)")
|
||||
except ValueError:
|
||||
secs_module = None # filtered out by include/exclude
|
||||
bundle = TrainBundle(dit_model, secs_module)
|
||||
|
||||
# ── resume ───────────────────────────────────────────────────────────────
|
||||
step_offset = epoch_offset = 0
|
||||
if args.lora_ckpt_path:
|
||||
sd, ckpt_cfg = load_lora_checkpoint(args.lora_ckpt_path)
|
||||
restored = load_trainable_lora_state(bundle, sd)
|
||||
step_offset, epoch_offset = resolve_offsets(args, ckpt_cfg)
|
||||
print(f"resumed {restored} adapter layer(s) from "
|
||||
f"{Path(args.lora_ckpt_path).name} (step {step_offset}, "
|
||||
f"epoch {epoch_offset})")
|
||||
else:
|
||||
step_offset, epoch_offset = resolve_offsets(args, None)
|
||||
|
||||
max_new_steps = max(0, args.max_steps - step_offset)
|
||||
|
||||
# ── optimizer: AdamW, single group; betas/eps/wd from config (torch
|
||||
# AdamW's decoupled weight decay — matches MLX's AdamW update exactly) ─────
|
||||
def scheduled_lr(step):
|
||||
"""SAT-dev InverseLR at 0-based step (== torch last_epoch on a fresh run,
|
||||
whose scheduler restarts even on resume). --lr-scheduler none → flat."""
|
||||
if args.lr_scheduler != "inverse":
|
||||
return args.lr
|
||||
warmup_factor = 1.0 - args.lr_warmup ** (step + 1)
|
||||
lr_mult = (1.0 + step / args.inv_gamma) ** (-args.lr_power)
|
||||
return warmup_factor * max(args.lr_final, args.lr * lr_mult)
|
||||
|
||||
optimizer = optim.AdamW(learning_rate=scheduled_lr(0),
|
||||
betas=[args.beta1, args.beta2],
|
||||
eps=args.eps, weight_decay=args.weight_decay)
|
||||
|
||||
# ── training-time local conditioning (underfit/upstream convention) ──────
|
||||
# The SA3 models are the diffusion_cond_inpaint variant: the trainer feeds
|
||||
# pure generation as an all-ONES inpaint mask + zero context (verified at
|
||||
# 80.7 dB forward parity vs the torch loop; the inference path's
|
||||
# local_add_cond=None ≡ zeros is an inference-only convention and trains
|
||||
# against the wrong conditioning regime — 4x loss difference).
|
||||
lac_const = mx.array(
|
||||
np.concatenate([np.ones((1, 1, crop_len), np.float32),
|
||||
np.zeros((1, 256, crop_len), np.float32)],
|
||||
axis=1).transpose(0, 2, 1)).astype(mx.float16)
|
||||
|
||||
# ── loss (adapters in the bundle are the only trainable params) ──────────
|
||||
# `seconds_in` is the raw seconds batch [B] fp32 when the trainable
|
||||
# seconds conditioner is active; with --no-conditioner-lora it is the
|
||||
# PRE-computed frozen seconds token [B, 1, 768] (built outside the step so
|
||||
# the compiled step stays pure mx — no numpy inside the traced graph).
|
||||
def loss_fn(bundle, latents, timesteps, loss_mask, prompt_cond,
|
||||
seconds_in, drop_mask):
|
||||
if bundle.secs is not None:
|
||||
sec_tok = bundle.secs(seconds_in) # [B, 1, 768] fp32
|
||||
else:
|
||||
sec_tok = seconds_in # frozen, precomputed
|
||||
cross = mx.concatenate([prompt_cond, sec_tok.astype(mx.float32)], axis=1)
|
||||
# CFG dropout (dit.py:441): whole cross_attn → zeros per sample;
|
||||
# global_cond (the seconds token) is kept.
|
||||
cross = mx.where(drop_mask[:, None, None], mx.zeros_like(cross), cross)
|
||||
global_cond = sec_tok[:, 0, :]
|
||||
cross16 = cross.astype(mx.float16)
|
||||
global16 = global_cond.astype(mx.float16)
|
||||
|
||||
def model_fn(noised, t):
|
||||
return bundle.dit(noised, t.astype(noised.dtype), cross16, global16,
|
||||
local_add_cond=lac_const)
|
||||
|
||||
return rectified_flow_loss(model_fn, latents, timesteps,
|
||||
loss_mask=loss_mask)
|
||||
|
||||
value_and_grad = nn.value_and_grad(bundle, loss_fn)
|
||||
|
||||
# ── train step (compiled by default — MLX "Compile > Training graphs") ──
|
||||
# Captured state: model params (the LoRA adapters the optimizer mutates),
|
||||
# optimizer moments, and the global PRNG key — rectified_flow_loss draws
|
||||
# its noise INSIDE the step. Everything np/python (timestep sampling, CFG
|
||||
# dropout draw, T5 encode, telemetry) stays OUTSIDE the step; batch
|
||||
# shapes/dtypes are constant (fixed crop + batch), so this traces once.
|
||||
# optimizer.init() materializes the moment arrays up front so the captured
|
||||
# state structure never changes (no recompile on step 2).
|
||||
optimizer.init(bundle.trainable_parameters())
|
||||
state = [bundle.state, optimizer.state, mx.random.state]
|
||||
clip_val = float(args.gradient_clip_val or 0.0)
|
||||
|
||||
def _global_l2(tree):
|
||||
"""sqrt(Σ ‖leaf‖²) over a param/grad tree in fp32 — the single global L2
|
||||
across all adapter leaves that underfit's _compute_grad_and_lora_norms
|
||||
reports (loop.py:243)."""
|
||||
total = mx.zeros((), dtype=mx.float32)
|
||||
for _, v in tree_flatten(tree):
|
||||
total = total + (v.astype(mx.float32) ** 2).sum()
|
||||
return mx.sqrt(total)
|
||||
|
||||
def train_step(latents, timesteps, loss_mask, prompt_cond, seconds_in,
|
||||
drop):
|
||||
loss, grads = value_and_grad(bundle, latents, timesteps, loss_mask,
|
||||
prompt_cond, seconds_in, drop)
|
||||
if clip_val > 0.0:
|
||||
grads, _ = optim.clip_grad_norm(grads, clip_val)
|
||||
# grad_norm on the POST-clip grads, lora_magnitude on the PRE-update
|
||||
# params — the order + definition underfit's torch loop uses (clip →
|
||||
# compute norms → optimizer.step, loop.py:628-637).
|
||||
grad_norm = _global_l2(grads)
|
||||
lora_mag = _global_l2(bundle.trainable_parameters())
|
||||
optimizer.update(bundle, grads)
|
||||
return loss, grad_norm, lora_mag
|
||||
|
||||
if args.compile:
|
||||
train_step = partial(mx.compile, inputs=state, outputs=state)(train_step)
|
||||
|
||||
# ── telemetry (underfit's loss_by_timestep.bin, struct "Iff") ────────────
|
||||
tele = open("loss_by_timestep.bin", "ab")
|
||||
|
||||
def checkpoint(step, epoch):
|
||||
fname = (f"{run_label}-step={step}-epoch={epoch}.safetensors"
|
||||
if run_label else f"step={step}-epoch={epoch}.safetensors")
|
||||
path = ckpt_dir / fname
|
||||
save_lora_checkpoint(
|
||||
bundle, path,
|
||||
include=lora_config.get("include"),
|
||||
exclude=lora_config.get("exclude"),
|
||||
extra_config={"step": int(step), "epoch": int(epoch),
|
||||
"base_model": base_model})
|
||||
print(f"\n ✓ checkpoint {path.name}")
|
||||
return path
|
||||
|
||||
# ── train loop (step-driven; epoch = one dataloader pass) ────────────────
|
||||
raw_step = 0
|
||||
epoch = epoch_offset
|
||||
last_saved_step = None
|
||||
t5_cache = {} if args.t5_cache else None
|
||||
t5_stats = {"hits": 0, "misses": 0}
|
||||
|
||||
# ── demos (underfit RF + ARC paths) ──────────────────────────────────────
|
||||
# RF/base entries: Euler inference on the BASE model + trained LoRA. ARC
|
||||
# entries (arc=true): the trained LoRA merged into the SHIPPED rf_denoiser
|
||||
# weights, sampled with the pingpong integrator — underfit's "train on base,
|
||||
# demo on ARC" story (its torch loop swaps in the full ARC model; here we
|
||||
# load the ARC npz with the LoRA merged at load). Both decode to mp3.
|
||||
# Baseline at the start, then every --demo-every steps + a final render.
|
||||
# Idempotent (a resume never re-renders an existing demo).
|
||||
demo_entries = None
|
||||
if args.demo_every > 0:
|
||||
if args.demo_config:
|
||||
demo_entries = json.loads(Path(args.demo_config).read_text())
|
||||
n_arc = sum(1 for e in demo_entries if e.get("arc"))
|
||||
print(f"demos: {len(demo_entries)} prompt(s) "
|
||||
f"({n_arc} ARC) every {args.demo_every} "
|
||||
f"step(s) → {args.demo_dir or '.'}")
|
||||
else:
|
||||
print("WARNING: --demo-every set but no --demo-config — demos disabled")
|
||||
demo_dir = args.demo_dir or "."
|
||||
demo_decoder_name = args.demo_decoder or DIT_CHOICES[args.dit]["default_decoder"]
|
||||
_model_default_decoder = DIT_CHOICES[args.dit]["default_decoder"]
|
||||
if demo_entries and demo_decoder_name != _model_default_decoder:
|
||||
print(f"demos: decoding with {demo_decoder_name.upper()} instead of the "
|
||||
f"model's {_model_default_decoder.upper()} (faster demos)")
|
||||
_demo_dec = {} # lazy: (decoder, chunk_fn, chunk_cfg)
|
||||
|
||||
def _demo_T_lat(entry):
|
||||
"""Per-demo latent length. The dashboard sets `duration` (seconds) on a
|
||||
demo entry only when it isn't at the crop length (see index.html
|
||||
demo_cond serialization) — full-length demos carry it, crop-length ones
|
||||
omit it. The DiT is length-agnostic (dit_mlx*), so one loaded model
|
||||
serves both without reloading."""
|
||||
dur = entry.get("duration")
|
||||
if dur:
|
||||
return max(2, int(round(float(dur) * demo.SAMPLE_RATE
|
||||
/ demo.SAMPLES_PER_LATENT)))
|
||||
return crop_len
|
||||
|
||||
def _demo_conditioning(prompt, seconds_val):
|
||||
"""cross_attn [1,257,768] fp16 + global_cond [1,768] fp16 for one demo
|
||||
prompt. Model-independent (T5Gemma + the trained seconds conditioner),
|
||||
so RF (base) and ARC demos share it."""
|
||||
prompt_cond = build_conditioning(t5, padding_emb, [prompt],
|
||||
cache=t5_cache, cache_stats=t5_stats)
|
||||
if bundle.secs is not None:
|
||||
sec_tok = bundle.secs(mx.array([seconds_val], dtype=mx.float32))
|
||||
else:
|
||||
sec_tok = secs_embedder([seconds_val])
|
||||
sec_tok = mx.stop_gradient(sec_tok.astype(mx.float32))
|
||||
cross = mx.concatenate([prompt_cond, sec_tok], axis=1).astype(mx.float16)
|
||||
gcond = sec_tok[:, 0, :].astype(mx.float16)
|
||||
return cross, gcond
|
||||
|
||||
def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg):
|
||||
"""Merge the trained LoRA into a fresh copy of the SHIPPED rf_denoiser
|
||||
weights and sample with the pingpong integrator. Loaded once for the
|
||||
whole ARC set (length-agnostic), freed after (peak = base + ARC DiT)."""
|
||||
arc_path = args.arc_weights or str(
|
||||
ensure_local(f"models/mlx/dit_{args.dit}_f16.npz"))
|
||||
# Scratch LoRA for merging into the ARC DiT — write it to the SESSION
|
||||
# dir (ckpt_dir.parent), NOT the scanned checkpoints/ dir, so it never
|
||||
# shows up as a checkpoint (dotfile too, belt + suspenders).
|
||||
tmp_ckpt = ckpt_dir.parent / f".arc_demo_tmp_{step:08d}.safetensors"
|
||||
save_lora_checkpoint(
|
||||
bundle, tmp_ckpt, include=lora_config.get("include"),
|
||||
exclude=lora_config.get("exclude"),
|
||||
extra_config={"step": int(step), "base_model": base_model})
|
||||
t_a = time.time()
|
||||
try:
|
||||
arc_dit = mod.load_dit(arc_path, T_lat=crop_len, dtype=mx.float16,
|
||||
compile_=False, lora_paths=[str(tmp_ckpt)],
|
||||
lora_strength=1.0, lora_log=lambda *a, **k: None)
|
||||
print(f" ARC model ({args.dit}) + LoRA merged "
|
||||
f"({time.time()-t_a:.1f}s)")
|
||||
except Exception as e:
|
||||
print(f" ARC demos skipped: could not load "
|
||||
f"{Path(arc_path).name} with LoRA merged: {e}", flush=True)
|
||||
_rm(tmp_ckpt)
|
||||
return
|
||||
try:
|
||||
for i, entry in arc_pending:
|
||||
prompt = entry.get("prompt", "")
|
||||
cfg = float(entry.get("cfg", 1))
|
||||
seed = int(entry.get("seed", 0))
|
||||
steps = int(entry.get("steps", 8))
|
||||
T_lat = _demo_T_lat(entry)
|
||||
seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE
|
||||
cross, gcond = _demo_conditioning(prompt, seconds_val)
|
||||
null = mx.zeros_like(cross) if cfg != 1.0 else None
|
||||
model_fn = demo.make_rf_model_fn(arc_dit, cross, gcond, cfg=cfg,
|
||||
null_cross_attn=null)
|
||||
sigmas = demo.build_sigmas(steps)
|
||||
noise = mx.random.normal((1, 256, T_lat), dtype=mx.float16,
|
||||
key=mx.random.key(seed))
|
||||
print(f" ▸ ARC demo {i}: {T_lat} latents (~{seconds_val:.0f}s audio), "
|
||||
f"{steps} steps", flush=True)
|
||||
try:
|
||||
t_gen = time.time()
|
||||
latent = demo.pingpong_sample(model_fn, noise, sigmas, seed=seed,
|
||||
desc=f"ARC demo {i}")
|
||||
audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg,
|
||||
latent, T_lat)
|
||||
meta = {"prompt": prompt, "cfg": cfg, "seed": seed,
|
||||
"steps": steps, "step": int(step), "arc": True}
|
||||
demo.save_demo_mp3(audio, i, step, demo.SAMPLE_RATE,
|
||||
demo_dir, meta)
|
||||
print(f" ♪ ARC demo {i} @ step {step}: {prompt[:44]!r} "
|
||||
f"→ demo_{i}_{step:08d}.mp3 ({time.time()-t_gen:.0f}s)", flush=True)
|
||||
except Exception as e:
|
||||
print(f" ARC demo {i} failed: {e}", flush=True)
|
||||
finally:
|
||||
del arc_dit
|
||||
if hasattr(mx, "clear_cache"):
|
||||
mx.clear_cache()
|
||||
_rm(tmp_ckpt)
|
||||
|
||||
def run_demos(step):
|
||||
if not demo_entries:
|
||||
return
|
||||
os.makedirs(demo_dir, exist_ok=True)
|
||||
pending = [(i, e) for i, e in enumerate(demo_entries)
|
||||
if not os.path.exists(
|
||||
os.path.join(demo_dir, f"demo_{i}_{step:08d}.mp3"))]
|
||||
if not pending:
|
||||
return
|
||||
if "dec" not in _demo_dec:
|
||||
from sa3_mlx import load_decoder
|
||||
t_d = time.time()
|
||||
_demo_dec["dec"] = load_decoder(demo_decoder_name, dtype=mx.float32)
|
||||
print(f" demo decoder {demo_decoder_name} loaded ({time.time()-t_d:.1f}s)")
|
||||
decoder, chunk_fn, chunk_cfg = _demo_dec["dec"]
|
||||
adapters = list(iter_trainable_lora_layers(bundle))
|
||||
rf_pending = [(i, e) for i, e in pending if not e.get("arc")]
|
||||
arc_pending = [(i, e) for i, e in pending if e.get("arc")]
|
||||
|
||||
# ── RF/base demos (Euler on the base model + trained LoRA) ──
|
||||
for i, entry in rf_pending:
|
||||
prompt = entry.get("prompt", "")
|
||||
cfg = float(entry.get("cfg", 7))
|
||||
seed = int(entry.get("seed", 0))
|
||||
steps = int(entry.get("steps", 50))
|
||||
strength = entry.get("lora_strength")
|
||||
interval = entry.get("lora_interval_max")
|
||||
T_lat = _demo_T_lat(entry)
|
||||
seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE
|
||||
cross, gcond = _demo_conditioning(prompt, seconds_val)
|
||||
null = mx.zeros_like(cross) if cfg != 1.0 else None
|
||||
model_fn = demo.make_rf_model_fn(bundle.dit, cross, gcond, cfg=cfg,
|
||||
null_cross_attn=null)
|
||||
sigmas = demo.build_sigmas(steps)
|
||||
noise = mx.random.normal((1, 256, T_lat), dtype=mx.float16,
|
||||
key=mx.random.key(seed))
|
||||
snap = before = None
|
||||
if interval is not None:
|
||||
snap = demo.snapshot_adapters(adapters)
|
||||
s_on = float(strength) if strength is not None else 1.0
|
||||
|
||||
def before(_i, sigma, s_on=s_on, interval=interval, snap=snap):
|
||||
demo.scale_adapters(snap, s_on if sigma <= interval else 0.0)
|
||||
elif strength is not None and float(strength) != 1.0:
|
||||
snap = demo.snapshot_adapters(adapters)
|
||||
demo.scale_adapters(snap, float(strength))
|
||||
print(f" ▸ demo {i}: {T_lat} latents (~{seconds_val:.0f}s audio), "
|
||||
f"{steps} steps", flush=True)
|
||||
try:
|
||||
t_gen = time.time()
|
||||
latent = demo.rf_euler_sample(model_fn, noise, sigmas,
|
||||
before_step=before, desc=f"demo {i}")
|
||||
audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, latent, T_lat)
|
||||
meta = {"prompt": prompt, "cfg": cfg, "seed": seed,
|
||||
"steps": steps, "step": int(step)}
|
||||
if strength is not None:
|
||||
meta["lora_strength"] = strength
|
||||
if interval is not None:
|
||||
meta["lora_interval_max"] = interval
|
||||
demo.save_demo_mp3(audio, i, step, demo.SAMPLE_RATE, demo_dir, meta)
|
||||
print(f" ♪ demo {i} @ step {step}: {prompt[:48]!r} "
|
||||
f"→ demo_{i}_{step:08d}.mp3 ({time.time()-t_gen:.0f}s)", flush=True)
|
||||
except Exception as e:
|
||||
print(f" demo {i} failed: {e}", flush=True)
|
||||
finally:
|
||||
if snap is not None:
|
||||
demo.scale_adapters(snap, 1.0) # restore trained factors
|
||||
|
||||
# ── ARC demos (pingpong on the shipped ARC weights + trained LoRA) ──
|
||||
if arc_pending:
|
||||
_run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg)
|
||||
|
||||
if hasattr(mx, "clear_cache"):
|
||||
mx.clear_cache()
|
||||
|
||||
if hasattr(mx, "reset_peak_memory"):
|
||||
mx.reset_peak_memory() # measure the training loop, not weight loading
|
||||
t_start = time.time()
|
||||
print(f"training: max {max_new_steps} new step(s) "
|
||||
f"(global target {args.max_steps}), lr {args.lr:g}, "
|
||||
f"batch {args.batch_size}, cfg-dropout {args.cfg_dropout_prob}, "
|
||||
f"compile {'on' if args.compile else 'off'}")
|
||||
run_demos(step_offset) # baseline (untrained / resumed state)
|
||||
try:
|
||||
while raw_step < max_new_steps:
|
||||
# tqdm per epoch — byte-identical to underfit's torch loop
|
||||
# (loop.py:516) so the dashboard's progress collapse, history parser,
|
||||
# and grad-norm / lora-magnitude charts work unchanged. desc
|
||||
# "Step N, Epoch E"; postfix train/loss,lr,grad_norm,lora_magnitude.
|
||||
pbar = tqdm(
|
||||
iterate_batches(dataset, args.batch_size, shuffle=True,
|
||||
seed=args.seed + epoch),
|
||||
total=steps_per_epoch,
|
||||
desc=f"Step {raw_step + step_offset}, Epoch {epoch}",
|
||||
mininterval=0, miniters=1, file=sys.stdout,
|
||||
)
|
||||
for batch in pbar:
|
||||
if raw_step >= max_new_steps:
|
||||
break
|
||||
global_step = raw_step + step_offset + 1
|
||||
|
||||
latents = mx.array(batch["latents"]).astype(mx.float16)
|
||||
loss_mask = mx.array(batch["padding_mask"])
|
||||
seconds = mx.array(np.asarray(batch["seconds_total"],
|
||||
dtype=np.float32))
|
||||
|
||||
# timesteps: sampler + model-config distribution shift. The shift
|
||||
# sequence length is the EFFECTIVE latent length per sample when
|
||||
# --use-effective-length (underfit use_effective_length_for_schedule,
|
||||
# True in the SA3 templates), else the crop length.
|
||||
t_np = sample_training_timesteps(
|
||||
args.timestep_sampler, latents.shape[0], rng=np_rng)
|
||||
if args.dist_shift != "none":
|
||||
if args.use_effective_length:
|
||||
seq_len = np.asarray(
|
||||
[int(math.ceil(int(float(s) * 44100) / 4096))
|
||||
for s in batch["seconds_total"]])
|
||||
else:
|
||||
seq_len = latents.shape[-1]
|
||||
t_np = shift_training_timesteps(
|
||||
t_np, seq_len, shift_type=args.dist_shift,
|
||||
options=DIST_SHIFT_DEFAULT["options"]
|
||||
if args.dist_shift == "full" else None)
|
||||
timesteps = mx.array(t_np)
|
||||
|
||||
prompt_cond = build_conditioning(t5, padding_emb,
|
||||
batch["prompt"],
|
||||
cache=t5_cache,
|
||||
cache_stats=t5_stats)
|
||||
drop = mx.array((np_rng.random(latents.shape[0])
|
||||
< args.cfg_dropout_prob))
|
||||
if bundle.secs is not None:
|
||||
seconds_in = seconds
|
||||
else:
|
||||
# frozen conditioner: build the seconds token OUTSIDE the
|
||||
# (possibly compiled) step — see loss_fn
|
||||
seconds_in = mx.stop_gradient(secs_embedder(
|
||||
[float(s) for s in batch["seconds_total"]]))
|
||||
|
||||
# InverseLR (or flat): the scheduler steps per optimizer step,
|
||||
# last_epoch = raw_step (0-based new-step index), matching torch's
|
||||
# fresh-on-resume scheduler. Setting learning_rate updates the leaf
|
||||
# in optimizer.state, which the compiled step reads via inputs=state.
|
||||
lr_now = scheduled_lr(raw_step)
|
||||
if args.lr_scheduler != "none":
|
||||
optimizer.learning_rate = lr_now
|
||||
|
||||
loss, grad_norm, lora_mag = train_step(
|
||||
latents, timesteps, loss_mask, prompt_cond, seconds_in, drop)
|
||||
mx.eval(state, loss, grad_norm, lora_mag)
|
||||
|
||||
raw_step += 1
|
||||
loss_v = float(loss)
|
||||
tele.write(struct.pack("Iff", global_step,
|
||||
float(t_np.mean()), loss_v))
|
||||
if raw_step % 10 == 0:
|
||||
tele.flush()
|
||||
|
||||
# Metrics on the tqdm postfix — the dashboard parses
|
||||
# train/loss=, train/lr=, train/grad_norm=, train/lora_magnitude=
|
||||
# per step (server.py _HISTORY_RE/_LR_RE/_GN_RE/_LM_RE), and the
|
||||
# "Step N," prefix gives it the authoritative global step.
|
||||
pbar.set_postfix({
|
||||
"train/loss": f"{loss_v:.6f}",
|
||||
"train/lr": f"{lr_now:.3e}",
|
||||
"train/grad_norm": f"{float(grad_norm):.6f}",
|
||||
"train/lora_magnitude": f"{float(lora_mag):.6f}",
|
||||
})
|
||||
pbar.set_description(f"Step {global_step}, Epoch {epoch}")
|
||||
|
||||
save_now = (args.checkpoint_every > 0
|
||||
and global_step % args.checkpoint_every == 0)
|
||||
demo_now = (args.demo_every > 0
|
||||
and global_step % args.demo_every == 0)
|
||||
if save_now or demo_now:
|
||||
# Clear + disable the bar so checkpoint/demo prints land on
|
||||
# their own lines (matches the torch loop, loop.py:680).
|
||||
pbar.clear()
|
||||
pbar.disable = True
|
||||
try:
|
||||
if save_now:
|
||||
checkpoint(global_step, epoch) # save BEFORE demos
|
||||
last_saved_step = global_step
|
||||
if demo_now:
|
||||
run_demos(global_step)
|
||||
finally:
|
||||
pbar.disable = False
|
||||
pbar.refresh()
|
||||
pbar.close()
|
||||
epoch += 1
|
||||
except KeyboardInterrupt:
|
||||
print("\ninterrupted — saving final checkpoint")
|
||||
|
||||
global_step = raw_step + step_offset
|
||||
if raw_step > 0 and global_step != last_saved_step:
|
||||
checkpoint(global_step, epoch)
|
||||
if raw_step > 0:
|
||||
run_demos(global_step) # final render
|
||||
tele.close()
|
||||
if t5_cache is not None:
|
||||
print(f"t5 cache: {t5_stats['hits']} hit(s) / "
|
||||
f"{t5_stats['misses']} miss(es), {len(t5_cache)} cached")
|
||||
print(f"peak memory: {mx.get_peak_memory()/2**30:.2f} GB"
|
||||
+ (" (grad-checkpoint on)" if args.grad_checkpoint else ""))
|
||||
print(f"done: {raw_step} step(s) in {time.time()-t_start:.0f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,107 @@
|
||||
"""§9 forward-parity harness — MLX side + diff.
|
||||
|
||||
Reads shared_inputs.npz + torch_out.npz (from parity_forward_torch.py) and runs
|
||||
the SAME controlled forward through the MLX base DiT in fp32, then reports three
|
||||
PSNRs that LOCALIZE any mismatch:
|
||||
|
||||
(1) DiT-only — feed torch's cross/global + the shared noised into the MLX DiT.
|
||||
Isolates the DiT forward convention (weights are the same fp16 source→fp32,
|
||||
so ~79 dB is the cross-framework fp32 floor; a convention bug reads ≪40 dB).
|
||||
(2) conditioning — MLX-native cross/global vs torch's (T5 + seconds; fp16-T5 bound).
|
||||
(3) end-to-end — full MLX-native prediction + loss vs torch.
|
||||
|
||||
Also asserts the torch model's pretransform.scale == 1.0 (the softnorm no-scale
|
||||
convention the MLX trainer relies on). Run in the MLX venv:
|
||||
|
||||
python parity_forward_mlx.py --base-npz <dit_sm-music-base_f16.npz> --workdir <dir>
|
||||
"""
|
||||
import argparse
|
||||
import importlib
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def psnr(a, b):
|
||||
a = np.asarray(a, np.float64); b = np.asarray(b, np.float64)
|
||||
mse = np.mean((a - b) ** 2)
|
||||
if mse == 0:
|
||||
return float("inf")
|
||||
return 20 * np.log10(max(np.abs(a).max(), np.abs(b).max()) / np.sqrt(mse))
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--base-npz", required=True, help="MLX base DiT weights (dit_<model>-base_f16.npz)")
|
||||
ap.add_argument("--workdir", required=True)
|
||||
ap.add_argument("--dit-module", default="models.defs.dit_mlx")
|
||||
ap.add_argument("--prompt", default="techno")
|
||||
ap.add_argument("--seconds", type=float, default=30.0)
|
||||
args = ap.parse_args()
|
||||
|
||||
import mlx.core as mx
|
||||
from sa3_mlx import T5GEMMA_NPZ_REL
|
||||
from weights import ensure_local
|
||||
from models.defs.sa3_pipeline import apply_prompt_padding, load_conditioner_from_npz
|
||||
from models.defs.t5gemma_mlx import T5Gemma
|
||||
|
||||
tt = np.load(f"{args.workdir}/torch_out.npz")
|
||||
inp = np.load(f"{args.workdir}/shared_inputs.npz", allow_pickle=True)
|
||||
assert abs(float(tt["scale"]) - 1.0) < 1e-9, \
|
||||
f"pretransform.scale={float(tt['scale'])} != 1.0 — the MLX trainer's no-scale " \
|
||||
"softnorm convention is WRONG for this model (real bug)."
|
||||
|
||||
latents = mx.array(inp["latents"]).astype(mx.float32)
|
||||
noise = mx.array(inp["noise"]).astype(mx.float32)
|
||||
t = mx.array(inp["t"]).astype(mx.float32)
|
||||
B, C, T = latents.shape
|
||||
dit = importlib.import_module(args.dit_module).load_dit(
|
||||
args.base_npz, T_lat=T, dtype=mx.float32, compile_=False)
|
||||
|
||||
noised = latents * (1.0 - t)[:, None, None] + noise * t[:, None, None]
|
||||
target = noise - latents
|
||||
lac = mx.concatenate([mx.ones((B, 1, T)), mx.zeros((B, C, T))], axis=1).transpose(0, 2, 1).astype(mx.float32)
|
||||
|
||||
# (1) DiT-only: inject torch's cross/global
|
||||
pred1 = dit(noised, t, mx.array(tt["cross"]).astype(mx.float32),
|
||||
mx.array(tt["gcond"]).astype(mx.float32), local_add_cond=lac)
|
||||
mx.eval(pred1)
|
||||
|
||||
# (2) MLX-native conditioning
|
||||
t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL)))
|
||||
padding_emb, secs = load_conditioner_from_npz(args.base_npz, prefix="cond.")
|
||||
embeds, mask = t5.encode([args.prompt] * B, max_len=256)
|
||||
prompt_cond = apply_prompt_padding(embeds.astype(mx.float32), mask, padding_emb.astype(mx.float32))
|
||||
sec_tok = mx.stop_gradient(secs([args.seconds] * B).astype(mx.float32))
|
||||
cross_m = mx.concatenate([prompt_cond, sec_tok], axis=1)
|
||||
gcond_m = sec_tok[:, 0, :]
|
||||
mx.eval(cross_m, gcond_m)
|
||||
|
||||
# (3) end-to-end
|
||||
pred3 = dit(noised, t, cross_m, gcond_m, local_add_cond=lac)
|
||||
mx.eval(pred3)
|
||||
loss_m = float(mx.mean((pred3 - target) ** 2))
|
||||
|
||||
# (4) BACKWARD: ∂loss/∂noised through the DiT (with torch's cross/global so
|
||||
# the backward is compared on identical inputs — isolates the DiT backward).
|
||||
cross_t = mx.array(tt["cross"]).astype(mx.float32)
|
||||
gcond_t = mx.array(tt["gcond"]).astype(mx.float32)
|
||||
|
||||
def loss_of_noised(nz):
|
||||
pred = dit(nz, t, cross_t, gcond_t, local_add_cond=lac)
|
||||
return mx.mean((pred - target) ** 2)
|
||||
|
||||
grad_m = mx.grad(loss_of_noised)(noised)
|
||||
mx.eval(grad_m)
|
||||
p_grad = psnr(np.array(grad_m), tt["grad_noised"]) if "grad_noised" in tt else float("nan")
|
||||
|
||||
print(f"noised sanity PSNR = {psnr(np.array(noised), tt['noised']):.1f} dB (inf = bit-identical)")
|
||||
print(f"(1) DiT-only forward = {psnr(np.array(pred1), tt['prediction']):6.2f} dB [conv bug ≪40; fp32 floor ~79]")
|
||||
print(f"(2) conditioning = cross {psnr(np.array(cross_m), tt['cross']):.2f} dB / global {psnr(np.array(gcond_m), tt['gcond']):.2f} dB")
|
||||
print(f"(3) end-to-end pred = {psnr(np.array(pred3), tt['prediction']):6.2f} dB")
|
||||
print(f"(4) BACKWARD d loss/d noised PSNR = {p_grad:6.2f} dB [full DiT backward; conv bug ≪40]")
|
||||
print(f"loss: torch {float(tt['loss']):.6f} MLX {loss_m:.6f} |Δ| {abs(float(tt['loss'])-loss_m):.2e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,105 @@
|
||||
"""§9 forward-parity harness — torch (MPS) side.
|
||||
|
||||
Controlled single forward of the SA3 base model, replicating underfit's
|
||||
training-step forward (underfit/training/loop.py) on INJECTED inputs so the
|
||||
result is deterministic and diffable against the MLX trainer's forward
|
||||
(parity_forward_mlx.py). fp32, cfg_dropout=0.
|
||||
|
||||
Run in an env with `stable_audio_3` importable (the sa3 repo venv):
|
||||
|
||||
python parity_forward_torch.py \
|
||||
--base-dir <sa3-sm-music-base> # dir with model_config.json + model.safetensors
|
||||
--latents <clip.npy> # a [256,T] pre-encoded latent
|
||||
--workdir <dir> # writes shared_inputs.npz + torch_out.npz
|
||||
|
||||
Then run parity_forward_mlx.py (MLX venv) with the same --workdir.
|
||||
|
||||
The harness localizes any mismatch: it saves cross/global/noised/prediction/
|
||||
loss + the model's pretransform.scale (the MLX side asserts scale==1.0 for the
|
||||
softnorm no-scale convention; a non-1.0 value would be a real bug). Weights on
|
||||
both sides come from the same fp16 source upcast to fp32, so agreement is the
|
||||
cross-framework fp32 floor (~79 dB); a convention bug reads far lower.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
from stable_audio_3.factory import create_diffusion_cond_from_config
|
||||
from stable_audio_3.loading_utils import remap_state_dict_keys
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--base-dir", required=True,
|
||||
help="Base checkpoint dir (model_config.json + model.safetensors)")
|
||||
ap.add_argument("--latents", required=True, help="A [256,T] pre-encoded latent .npy")
|
||||
ap.add_argument("--workdir", required=True)
|
||||
ap.add_argument("--crop", type=int, default=256)
|
||||
ap.add_argument("--prompt", default="techno")
|
||||
ap.add_argument("--seconds", type=float, default=30.0)
|
||||
ap.add_argument("--noise-seed", type=int, default=12345)
|
||||
ap.add_argument("--device", default=None)
|
||||
args = ap.parse_args()
|
||||
os.makedirs(args.workdir, exist_ok=True)
|
||||
device = args.device or ("mps" if torch.backends.mps.is_available() else "cpu")
|
||||
|
||||
# Shared injected inputs (numpy — both engines consume these verbatim).
|
||||
clip = np.load(args.latents)[:, : args.crop]
|
||||
lat = np.stack([clip, clip, clip], 0).astype(np.float32) # B=3
|
||||
noise = np.random.default_rng(args.noise_seed).standard_normal(lat.shape).astype(np.float32)
|
||||
t = np.array([0.25, 0.5, 0.75], dtype=np.float32)
|
||||
np.savez(f"{args.workdir}/shared_inputs.npz", latents=lat, noise=noise, t=t,
|
||||
prompt=args.prompt, seconds_total=args.seconds)
|
||||
|
||||
cfg = json.load(open(f"{args.base_dir}/model_config.json"))
|
||||
print("building base model (fp32)...", flush=True)
|
||||
model = create_diffusion_cond_from_config(cfg)
|
||||
sd = load_file(f"{args.base_dir}/model.safetensors")
|
||||
sd = remap_state_dict_keys(sd, model.state_dict())
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False)
|
||||
print(f"loaded: {len(missing)} missing / {len(unexpected)} unexpected", flush=True)
|
||||
model = model.to(device).eval().requires_grad_(False)
|
||||
|
||||
scale = float(getattr(model.pretransform, "scale", 1.0)) if model.pretransform is not None else 1.0
|
||||
print(f"pretransform_scale = {scale} diffusion_objective = {model.diffusion_objective}", flush=True)
|
||||
|
||||
latents = torch.tensor(lat).to(device)
|
||||
noise_t = torch.tensor(noise).to(device)
|
||||
tt = torch.tensor(t).to(device)
|
||||
B, C, T = latents.shape
|
||||
diffusion_input = latents / scale if scale != 1.0 else latents
|
||||
noised = diffusion_input * (1 - tt)[:, None, None] + noise_t * tt[:, None, None]
|
||||
target = noise_t - diffusion_input
|
||||
|
||||
metadata = [{"prompt": args.prompt, "seconds_total": args.seconds, "seconds_start": 0.0}
|
||||
for _ in range(B)]
|
||||
conditioning = model.conditioner(metadata, device)
|
||||
conditioning["inpaint_mask"] = [torch.ones(B, 1, T, device=device)]
|
||||
conditioning["inpaint_masked_input"] = [torch.zeros(B, C, T, device=device)]
|
||||
with torch.no_grad():
|
||||
ci = model.get_conditioning_inputs(conditioning)
|
||||
|
||||
# BACKWARD: gradient of the loss w.r.t. the noised input — the full DiT
|
||||
# backward sweep (every interior gradient the adapter grads are built from).
|
||||
noised.requires_grad_(True)
|
||||
output = model(noised, tt, cond=conditioning, cfg_dropout_prob=0.0,
|
||||
padding_mask=torch.ones(B, T, dtype=torch.bool, device=device))
|
||||
loss = ((output.float() - target.float()) ** 2).mean() # full mask → signal-only == mean
|
||||
loss.backward()
|
||||
grad_noised = noised.grad
|
||||
print(f"loss = {float(loss):.8f} |grad_noised| = {float(grad_noised.abs().mean()):.6e}", flush=True)
|
||||
np.savez(f"{args.workdir}/torch_out.npz",
|
||||
prediction=output.detach().float().cpu().numpy(), loss=float(loss),
|
||||
cross=ci["cross_attn_cond"].float().cpu().numpy(),
|
||||
gcond=ci["global_cond"].float().cpu().numpy(),
|
||||
noised=noised.detach().float().cpu().numpy(),
|
||||
target=target.float().cpu().numpy(),
|
||||
grad_noised=grad_noised.float().cpu().numpy(), scale=scale)
|
||||
print(f"wrote {args.workdir}/torch_out.npz", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,441 @@
|
||||
"""Pre-encode audio files into SAME latents for LoRA finetuning (MLX, torch-free).
|
||||
|
||||
Offline audio → latent CLI mirroring underfit's ``dataset_processing/pre_encode.py``
|
||||
conventions, ported to the standalone MLX runtime: single-process, soundfile-based
|
||||
loading (no torch / torchaudio), and ``models.defs.audio_encoding.encode_audio``
|
||||
doing the padding / chunking / masking.
|
||||
|
||||
Every audio file under ``--audio-dir`` (recursive) becomes a ``.npy`` + ``.json``
|
||||
pair at the same relative path under ``--output-dir``:
|
||||
|
||||
<output-dir>/<relpath-with-npy-suffix> latents [D, T] float32
|
||||
<output-dir>/<relpath-with-json-suffix> metadata (underfit field parity)
|
||||
|
||||
Conventions matched to underfit:
|
||||
* Whole files, no clip splitting; cap at ``--max-duration`` seconds (default
|
||||
600) with the sample cap aligned DOWN to a multiple of 4096 (ds_ratio).
|
||||
* 44.1 kHz stereo — mono is duplicated to stereo, >2 channels keep the first
|
||||
two, other sample rates are linearly resampled. NO loudness/peak norm.
|
||||
* Encoder runs fp32; ``chunked=True`` for clips longer than 30 s.
|
||||
* Metadata: ``path, relpath, src_relpath, seconds_total`` (rounded 3dp),
|
||||
``seconds_start: 0, audio_samples, latent_shape``, latent-resolution 0/1
|
||||
``padding_mask``, merged with tags extracted from sidecars.
|
||||
* Skip files whose outputs already exist unless ``--overwrite``.
|
||||
* ``<output-dir>/details.json`` records the global encode params.
|
||||
|
||||
Latent scaling note: the MLX SAME encoders output DiT-space (softnorm) latents
|
||||
directly — the equivalent of torch's unscaled-latents-then-divide-by-scale
|
||||
convention with scale 1.0 — so latents are saved as-is and the trainer uses
|
||||
them directly (no pretransform-scale division at train time).
|
||||
|
||||
Tag extraction (ported from underfit's ``extract_tags``): JSON sidecar
|
||||
(``<stem>.json`` next to the audio, or ``<parent>/json/<stem>.json``) keeps all
|
||||
non-empty string/number values; else a ``.txt`` sidecar (same dir or
|
||||
``<parent>/txt/``) whose whole stripped content becomes the ``prompt`` tag.
|
||||
Embedded ID3/Vorbis tags are a DELIBERATE OMISSION — they need an extra
|
||||
dependency (``audio_metadata``), and sidecars are the underfit-recommended
|
||||
path anyway.
|
||||
|
||||
Usage:
|
||||
python pre_encode_mlx.py --audio-dir ~/clips --output-dir ~/latents \
|
||||
--codec same-s [--max-duration 600] [--overwrite]
|
||||
|
||||
``--codec same-s`` pairs with the sm-music / sm-sfx DiTs, ``same-l`` with medium.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
REPO = Path(__file__).resolve().parent.parent # optimized/mlx (scripts/ is one level down)
|
||||
if str(REPO) not in sys.path:
|
||||
sys.path.insert(0, str(REPO))
|
||||
if str(REPO / "scripts") not in sys.path:
|
||||
sys.path.insert(0, str(REPO / "scripts"))
|
||||
|
||||
import mlx.core as mx # noqa: E402
|
||||
|
||||
from models.defs.audio_encoding import ( # noqa: E402
|
||||
SAME_ENCODER_PAD_MODULO,
|
||||
SAMPLES_PER_LATENT,
|
||||
encode_audio,
|
||||
)
|
||||
|
||||
SAMPLE_RATE = 44100
|
||||
DS_RATIO = SAMPLES_PER_LATENT # 4096 samples per latent
|
||||
CHUNKED_THRESHOLD_SECONDS = 30.0
|
||||
|
||||
AUDIO_EXTS = {".wav", ".mp3", ".flac", ".ogg", ".aif", ".aiff", ".m4a"}
|
||||
_MIN_AUDIO_SIZE = 4096 # skip files smaller than this (macOS resource forks, corrupt)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Audio discovery
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def find_audio_files(root) -> list[Path]:
|
||||
"""Recursively find audio files under ``root``, sorted by path.
|
||||
|
||||
Skips macOS resource forks (``._*``) and files too small to be real audio.
|
||||
"""
|
||||
files = []
|
||||
for dirpath, _, filenames in os.walk(root):
|
||||
for fn in filenames:
|
||||
if fn.startswith("._"):
|
||||
continue
|
||||
fp = Path(dirpath) / fn
|
||||
if fp.suffix.lower() in AUDIO_EXTS:
|
||||
try:
|
||||
if fp.stat().st_size < _MIN_AUDIO_SIZE:
|
||||
continue
|
||||
except OSError:
|
||||
continue
|
||||
files.append(fp)
|
||||
files.sort()
|
||||
return files
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Audio loading (torch-free)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_audio(path) -> np.ndarray:
|
||||
"""Load an audio file → ``(2, T)`` float32 at 44.1 kHz.
|
||||
|
||||
soundfile decodes; mono is duplicated to stereo, >2 channels keep the
|
||||
first two; non-44.1 kHz input is linearly resampled at
|
||||
align_corners=False positions (same pattern as sa3_gradio.read_audio_any).
|
||||
No loudness/peak normalization.
|
||||
"""
|
||||
import soundfile as sf
|
||||
|
||||
data, sr = sf.read(str(path), dtype="float32", always_2d=True) # (T, ch)
|
||||
y = data.T
|
||||
y = np.stack([y[0], y[0]]) if y.shape[0] == 1 else y[:2]
|
||||
if sr != SAMPLE_RATE:
|
||||
in_len = y.shape[-1]
|
||||
new_len = int(round(in_len * SAMPLE_RATE / sr))
|
||||
scale = in_len / new_len
|
||||
pos = np.clip((np.arange(new_len) + 0.5) * scale - 0.5, 0, in_len - 1)
|
||||
y = np.stack([np.interp(pos, np.arange(in_len), ch) for ch in y])
|
||||
return np.ascontiguousarray(y, dtype=np.float32)
|
||||
|
||||
|
||||
def max_samples_for_duration(max_duration: float) -> int:
|
||||
"""Sample cap for ``max_duration`` seconds, aligned DOWN to ds_ratio (4096)."""
|
||||
return (int(max_duration * SAMPLE_RATE) // DS_RATIO) * DS_RATIO
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sidecar tag extraction (ported from underfit pre_encode.extract_tags;
|
||||
# embedded ID3/Vorbis tags deliberately omitted — see module docstring)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def extract_tags(filepath) -> dict:
|
||||
"""Extract tag fields for an audio file from sidecar files.
|
||||
|
||||
Priority:
|
||||
1. JSON sidecar (``<stem>.json`` same dir, or ``<parent>/json/<stem>.json``)
|
||||
— all non-empty string/number values kept (custom keys preserved).
|
||||
2. Plain ``.txt`` sidecar (same dir or ``<parent>/txt/``) — whole stripped
|
||||
content becomes the ``prompt`` tag.
|
||||
Returns ``{}`` when neither sidecar exists (no embedded-tag fallback).
|
||||
"""
|
||||
fp = Path(filepath)
|
||||
stem = fp.stem
|
||||
|
||||
# 1. JSON sidecar — same dir, then sibling json/ dir
|
||||
sidecar_candidates = [
|
||||
fp.with_suffix(".json"),
|
||||
fp.parent.parent / "json" / (stem + ".json"),
|
||||
]
|
||||
for sidecar in sidecar_candidates:
|
||||
if sidecar.exists():
|
||||
try:
|
||||
with open(sidecar) as f:
|
||||
sc = json.load(f)
|
||||
tags_out = {
|
||||
k: str(v)
|
||||
for k, v in sc.items()
|
||||
if v and isinstance(v, (str, int, float))
|
||||
}
|
||||
if tags_out:
|
||||
return tags_out
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 2. Plain .txt sidecar (SA3 convention) — same dir, then sibling txt/ dir.
|
||||
txt_candidates = [
|
||||
fp.with_suffix(".txt"),
|
||||
fp.parent.parent / "txt" / (stem + ".txt"),
|
||||
]
|
||||
for txt in txt_candidates:
|
||||
if txt.exists():
|
||||
try:
|
||||
content = txt.read_text(encoding="utf-8", errors="replace").strip()
|
||||
if content:
|
||||
return {"prompt": content}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core encode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def encode_file(
|
||||
encoder,
|
||||
audio_np: np.ndarray,
|
||||
*,
|
||||
pad_modulo: int,
|
||||
actual_samples: int | None = None,
|
||||
):
|
||||
"""Encode one ``(2, T)`` float32 waveform → (latents, padding_mask, seconds).
|
||||
|
||||
Returns:
|
||||
latents_np: ``[D, T_latent]`` float32 (softnorm/DiT space, saved as-is)
|
||||
padding_mask: latent-resolution 0/1 int list, length ``T_latent`` —
|
||||
zeros over the alignment-padding tail
|
||||
seconds_total: ``round(actual_samples / 44100, 3)``
|
||||
|
||||
``chunked=True`` is used automatically for clips longer than 30 s.
|
||||
"""
|
||||
audio_np = np.asarray(audio_np)
|
||||
if audio_np.ndim != 2 or audio_np.shape[0] != 2:
|
||||
raise ValueError(f"Expected audio shaped (2, T), got {audio_np.shape}.")
|
||||
if actual_samples is None:
|
||||
actual_samples = int(audio_np.shape[-1])
|
||||
duration = actual_samples / SAMPLE_RATE
|
||||
encoded = encode_audio(
|
||||
encoder,
|
||||
audio_np[None, ...],
|
||||
valid_sample_lengths=[actual_samples],
|
||||
pad_modulo=pad_modulo,
|
||||
chunked=duration > CHUNKED_THRESHOLD_SECONDS,
|
||||
)
|
||||
latents_np = np.asarray(encoded.latents.astype(mx.float32))[0] # [D, T]
|
||||
padding_mask = [int(v) for v in np.asarray(encoded.padding_mask)[0]]
|
||||
seconds_total = round(actual_samples / SAMPLE_RATE, 3)
|
||||
return latents_np, padding_mask, seconds_total
|
||||
|
||||
|
||||
def build_metadata(
|
||||
fpath,
|
||||
rel,
|
||||
out_rel,
|
||||
*,
|
||||
actual_samples: int,
|
||||
latent_shape,
|
||||
padding_mask,
|
||||
seconds_total,
|
||||
) -> dict:
|
||||
"""Underfit-parity metadata dict for one encoded file (tags merged last)."""
|
||||
meta = {
|
||||
"path": str(fpath),
|
||||
"relpath": str(out_rel),
|
||||
"src_relpath": str(rel),
|
||||
"seconds_total": seconds_total,
|
||||
"seconds_start": 0,
|
||||
"audio_samples": int(actual_samples),
|
||||
"latent_shape": [int(s) for s in latent_shape],
|
||||
"padding_mask": padding_mask,
|
||||
}
|
||||
meta.update(extract_tags(fpath))
|
||||
return meta
|
||||
|
||||
|
||||
def run(
|
||||
audio_dir,
|
||||
output_dir,
|
||||
codec: str,
|
||||
encoder,
|
||||
*,
|
||||
max_duration: float = 600.0,
|
||||
overwrite: bool = False,
|
||||
exclude: set | None = None,
|
||||
) -> dict:
|
||||
"""Encode every audio file under ``audio_dir`` into ``output_dir``.
|
||||
|
||||
Returns stats: ``{"encoded", "skipped", "errors", "count"}``.
|
||||
"""
|
||||
if codec not in SAME_ENCODER_PAD_MODULO:
|
||||
raise ValueError(f"Unknown codec {codec!r}; expected same-s or same-l.")
|
||||
pad_modulo = SAME_ENCODER_PAD_MODULO[codec]
|
||||
|
||||
audio_dir = Path(audio_dir).expanduser().resolve()
|
||||
output_dir = Path(output_dir).expanduser().resolve()
|
||||
files = find_audio_files(audio_dir)
|
||||
if exclude:
|
||||
files = [f for f in files if str(f.relative_to(audio_dir)) not in exclude]
|
||||
max_samples = max_samples_for_duration(max_duration)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
details = {
|
||||
"codec": codec,
|
||||
"sample_rate": SAMPLE_RATE,
|
||||
"max_duration": max_duration,
|
||||
"max_samples": max_samples,
|
||||
"pad_modulo": pad_modulo,
|
||||
"audio_dir": str(audio_dir),
|
||||
"count": len(files),
|
||||
}
|
||||
with open(output_dir / "details.json", "w") as f:
|
||||
json.dump(details, f, indent=2)
|
||||
|
||||
total = len(files)
|
||||
encoded = skipped = errors = 0
|
||||
t0 = time.time()
|
||||
for index, fpath in enumerate(files, 1):
|
||||
rel = fpath.relative_to(audio_dir)
|
||||
out_rel = rel.with_suffix(".npy")
|
||||
npy_path = output_dir / out_rel
|
||||
json_path = npy_path.with_suffix(".json")
|
||||
tag = f"[{index}/{total}]"
|
||||
if npy_path.exists() and json_path.exists() and not overwrite:
|
||||
skipped += 1
|
||||
print(f" {tag} SKIP {out_rel} (exists)")
|
||||
continue
|
||||
try:
|
||||
audio = load_audio(fpath)
|
||||
actual_samples = int(audio.shape[-1])
|
||||
if actual_samples > max_samples:
|
||||
audio = audio[:, :max_samples]
|
||||
actual_samples = max_samples
|
||||
|
||||
latents_np, padding_mask, seconds_total = encode_file(
|
||||
encoder, audio, pad_modulo=pad_modulo, actual_samples=actual_samples
|
||||
)
|
||||
|
||||
npy_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
np.save(str(npy_path), latents_np)
|
||||
meta = build_metadata(
|
||||
fpath,
|
||||
rel,
|
||||
out_rel,
|
||||
actual_samples=actual_samples,
|
||||
latent_shape=latents_np.shape,
|
||||
padding_mask=padding_mask,
|
||||
seconds_total=seconds_total,
|
||||
)
|
||||
with open(json_path, "w") as f:
|
||||
json.dump(meta, f)
|
||||
|
||||
encoded += 1
|
||||
shape_str = "x".join(str(s) for s in latents_np.shape)
|
||||
print(f" {tag} {out_rel} {seconds_total:.1f}s -> [{shape_str}]")
|
||||
except Exception as e:
|
||||
errors += 1
|
||||
print(f" {tag} ERROR {rel}: {e}")
|
||||
|
||||
elapsed = time.time() - t0
|
||||
mins, secs = divmod(int(elapsed), 60)
|
||||
print(
|
||||
f"\nDone in {mins}m{secs:02d}s: {encoded} encoded, {skipped} skipped, "
|
||||
f"{errors} errors (of {total})"
|
||||
)
|
||||
print(f"Output: {output_dir}")
|
||||
return {"encoded": encoded, "skipped": skipped, "errors": errors, "count": total}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_codec_encoder(codec: str):
|
||||
"""Load the real SAME encoder for ``codec`` in fp32. Returns (encoder, pad_modulo)."""
|
||||
from sa3_mlx import load_encoder # lazy: pulls weight-download machinery
|
||||
|
||||
return load_encoder(codec, mx.float32)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Pre-encode audio into SAME latents for MLX LoRA finetuning",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-dir", "-i", required=True,
|
||||
help="Directory of audio files (searched recursively)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir", "-o", required=True,
|
||||
help="Output root; each audio file becomes <output-dir>/<relpath>.npy + .json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--codec", required=True, choices=sorted(SAME_ENCODER_PAD_MODULO),
|
||||
help="Autoencoder: same-s pairs with sm-music/sm-sfx, same-l with medium",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-duration", type=float, default=600.0,
|
||||
help="Crop files longer than N seconds (default: 600; aligned down to 4096 samples)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite", action="store_true",
|
||||
help="Re-encode files whose outputs already exist",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--exclude-file", type=str, default=None,
|
||||
help="Text file of relpaths (one per line, relative to --audio-dir) to skip",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
audio_dir = Path(args.audio_dir).expanduser().resolve()
|
||||
if not audio_dir.is_dir():
|
||||
print(f"Error: not a directory: {audio_dir}")
|
||||
sys.exit(1)
|
||||
files = find_audio_files(audio_dir)
|
||||
if not files:
|
||||
print(f"No audio files found in {audio_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
exclude = set()
|
||||
if args.exclude_file and Path(args.exclude_file).is_file():
|
||||
exclude = {ln.strip() for ln in Path(args.exclude_file).read_text().splitlines()
|
||||
if ln.strip()}
|
||||
n_before = len(files)
|
||||
if exclude:
|
||||
files = [f for f in files if str(f.relative_to(audio_dir)) not in exclude]
|
||||
if not files:
|
||||
print(f"All {n_before} file(s) excluded by {args.exclude_file}")
|
||||
sys.exit(1)
|
||||
|
||||
max_samples = max_samples_for_duration(args.max_duration)
|
||||
import mlx.core as mx
|
||||
print(f"Engine: MLX · device: {mx.default_device()}")
|
||||
print(f"Codec: {args.codec} (pad_modulo={SAME_ENCODER_PAD_MODULO[args.codec]}, fp32)")
|
||||
print(f"Input: {audio_dir} ({len(files)} audio files"
|
||||
+ (f", {n_before - len(files)} excluded)" if exclude else ")"))
|
||||
print(f"Output: {Path(args.output_dir).expanduser().resolve()}")
|
||||
print(f"Max: {max_samples / SAMPLE_RATE:.1f}s ({max_samples:,} samples)\n")
|
||||
|
||||
print("Loading encoder...")
|
||||
encoder, pad_modulo = load_codec_encoder(args.codec)
|
||||
assert pad_modulo == SAME_ENCODER_PAD_MODULO[args.codec]
|
||||
|
||||
run(
|
||||
audio_dir,
|
||||
args.output_dir,
|
||||
args.codec,
|
||||
encoder,
|
||||
max_duration=args.max_duration,
|
||||
overwrite=args.overwrite,
|
||||
exclude=exclude,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,891 @@
|
||||
"""SA3 text-to-audio inference in pure MLX (no stable-audio-tools needed).
|
||||
|
||||
Pipeline:
|
||||
prompt → T5Gemma → conditioning (cross_attn + global) → DiT pingpong (8 steps)
|
||||
→ SAME-S / SAME-L decoder → patched-pretransform unpatch → WAV
|
||||
|
||||
Usage:
|
||||
python sa3_mlx.py --prompt "lofi house loop" --dit small --decoder same-s --seconds 30 --out out.wav
|
||||
|
||||
If --dit or --decoder is omitted, the script prompts the user interactively.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import argparse, math, os, random, subprocess, sys, termios, time, tty, wave
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
|
||||
REPO = Path(__file__).resolve().parent.parent # project root (scripts/ is one level down)
|
||||
sys.path.insert(0, str(REPO)) # so `from models.defs.*` resolves
|
||||
sys.path.insert(0, str(REPO / "scripts")) # so `from weights import *` resolves
|
||||
|
||||
from models.defs.sa3_pipeline import (
|
||||
SecondsTotalEmbedder,
|
||||
apply_prompt_padding,
|
||||
build_pingpong_schedule,
|
||||
sample_flow_pingpong,
|
||||
patched_decode,
|
||||
load_conditioner_from_npz,
|
||||
)
|
||||
from models.defs.t5gemma_mlx import T5Gemma
|
||||
from weights import ensure_local, is_present
|
||||
|
||||
SAMPLE_RATE = 44100
|
||||
SAMPLES_PER_LATENT = 4096 # PatchedPretransform downsample × SAME 16× expansion
|
||||
|
||||
# ─── Display helpers (ANSI color when stdout is a TTY) ───────────────────
|
||||
_USE_COLOR = sys.stdout.isatty()
|
||||
_RULE_W = 64
|
||||
|
||||
def _c(code: str, s: str) -> str:
|
||||
return f"\x1b[{code}m{s}\x1b[0m" if _USE_COLOR else s
|
||||
|
||||
def bold(s): return _c("1", s)
|
||||
def dim(s): return _c("2", s)
|
||||
def cyan(s): return _c("36", s)
|
||||
def yellow(s): return _c("33", s)
|
||||
def green(s): return _c("32", s)
|
||||
def magenta(s):return _c("35", s)
|
||||
|
||||
def rule(char="━", color=cyan):
|
||||
print(color(char * _RULE_W))
|
||||
|
||||
def banner(title: str):
|
||||
rule()
|
||||
print(f" {bold(title)}")
|
||||
rule()
|
||||
|
||||
def _fmt_mem(b: int) -> str:
|
||||
if b >= 1024**3: return f"{b/1024**3:5.2f} GB"
|
||||
return f"{b/1024**2:5.0f} MB"
|
||||
|
||||
|
||||
def stage(idx_total: str, label: str, ms: float | None = None, peak_b: int | None = None):
|
||||
"""Right-aligned timing column with dotted fill. `peak_b` is accepted for
|
||||
backward-compat / tracking-via-caller but is NOT rendered inline — peak RAM
|
||||
is shown only in the 'Peak RAM by stage' summary at the end."""
|
||||
head = f" {cyan(idx_total)} {bold(label)}"
|
||||
if ms is None:
|
||||
print(head)
|
||||
return
|
||||
visible = len(f" {idx_total} {label}")
|
||||
fill = max(2, _RULE_W - visible - 9)
|
||||
dots = dim("·" * fill)
|
||||
print(f"{head} {dots} {yellow(f'{ms:>5.0f} ms')}")
|
||||
|
||||
def sub(text: str):
|
||||
print(f" {dim(text)}")
|
||||
|
||||
# Each DiT .npz file has the conditioner weights baked in under a "cond." key prefix.
|
||||
# The codec architecture (same-s vs same-l) is fixed per DiT family.
|
||||
DIT_CHOICES = {
|
||||
"sm-music": {"loader": "models.defs.dit_mlx",
|
||||
"ckpt": "models/mlx/dit_sm-music_f16.npz",
|
||||
"default_decoder": "same-s"},
|
||||
"sm-sfx": {"loader": "models.defs.dit_mlx",
|
||||
"ckpt": "models/mlx/dit_sm-sfx_f16.npz",
|
||||
"default_decoder": "same-s"},
|
||||
"medium": {"loader": "models.defs.dit_mlx_medium",
|
||||
"ckpt": "models/mlx/dit_medium_f16.npz",
|
||||
"default_decoder": "same-l"},
|
||||
}
|
||||
DECODER_CHOICES = {
|
||||
"same-s": ("models.defs.same_s_decoder", "decode_chunked", (8, 2),
|
||||
"models/mlx/same_s_decoder_f32.npz"),
|
||||
"same-l": ("models.defs.same_l_decoder", "decode_chunked", (128, 8),
|
||||
"models/mlx/same_l_decoder_f32.npz"),
|
||||
}
|
||||
# Encoder paired with each decoder (used for --init-audio).
|
||||
ENCODER_CHOICES = {
|
||||
"same-s": ("models.defs.same_s_encoder", 32, "models/mlx/same_s_encoder_f32.npz"),
|
||||
"same-l": ("models.defs.same_l_encoder", 16, "models/mlx/same_l_encoder_f32.npz"),
|
||||
}
|
||||
T5GEMMA_NPZ_REL = "models/mlx/t5gemma_f16.npz"
|
||||
|
||||
|
||||
def _arrow_pick(prompt: str, options: list[str], default: str | None = None) -> str:
|
||||
"""Tiny arrow-key picker — no external deps, posix termios only.
|
||||
|
||||
Up/Down to move, Enter to select, Ctrl-C to abort. Falls back to a
|
||||
numeric prompt when stdin isn't a TTY (piped input, CI, etc.).
|
||||
"""
|
||||
if not sys.stdin.isatty():
|
||||
print(prompt)
|
||||
for i, o in enumerate(options):
|
||||
mark = "*" if o == default else " "
|
||||
print(f" {mark} [{i}] {o}")
|
||||
s = input(f"Choose [0-{len(options)-1}] (Enter for default): ").strip()
|
||||
if s == "":
|
||||
return default or options[0]
|
||||
if s.isdigit() and 0 <= int(s) < len(options):
|
||||
return options[int(s)]
|
||||
return s if s in options else (default or options[0])
|
||||
|
||||
idx = options.index(default) if default in options else 0
|
||||
fd = sys.stdin.fileno()
|
||||
old = termios.tcgetattr(fd)
|
||||
print(prompt)
|
||||
# leave room for the options
|
||||
for _ in options:
|
||||
print()
|
||||
try:
|
||||
tty.setcbreak(fd)
|
||||
while True:
|
||||
# move up `len(options)` lines and redraw
|
||||
sys.stdout.write(f"\x1b[{len(options)}A")
|
||||
for i, o in enumerate(options):
|
||||
if i == idx:
|
||||
sys.stdout.write(f"\x1b[2K\x1b[36m▶ {o}\x1b[0m\n")
|
||||
else:
|
||||
sys.stdout.write(f"\x1b[2K {o}\n")
|
||||
sys.stdout.flush()
|
||||
ch = sys.stdin.read(1)
|
||||
if ch == "\x1b":
|
||||
seq = sys.stdin.read(2)
|
||||
if seq == "[A":
|
||||
idx = (idx - 1) % len(options)
|
||||
elif seq == "[B":
|
||||
idx = (idx + 1) % len(options)
|
||||
elif ch in ("\n", "\r"):
|
||||
return options[idx]
|
||||
elif ch == "\x03":
|
||||
raise KeyboardInterrupt
|
||||
finally:
|
||||
termios.tcsetattr(fd, termios.TCSADRAIN, old)
|
||||
|
||||
|
||||
def prompt_user_if_missing(args):
|
||||
"""Interactive arrow-key selection when --dit / --decoder aren't supplied."""
|
||||
if args.dit is None:
|
||||
args.dit = _arrow_pick("Choose DiT model:", list(DIT_CHOICES.keys()), default="medium")
|
||||
print(f" → {args.dit}")
|
||||
if args.decoder is None:
|
||||
suggested = DIT_CHOICES[args.dit]["default_decoder"]
|
||||
args.decoder = _arrow_pick("Choose audio decoder:", list(DECODER_CHOICES.keys()), default=suggested)
|
||||
print(f" → {args.decoder}")
|
||||
if args.seed is None:
|
||||
args.seed = random.randint(0, 2**31 - 1)
|
||||
return args
|
||||
|
||||
|
||||
def load_dit(dit_name: str, T_lat: int, dtype, lora_specs=None, num_steps=None):
|
||||
cfg = DIT_CHOICES[dit_name]
|
||||
ckpt = ensure_local(cfg["ckpt"])
|
||||
import importlib, io, contextlib
|
||||
mod = importlib.import_module(cfg["loader"])
|
||||
# The loader's own chatter is swallowed, but the LoRA merge summary is routed
|
||||
# to stderr (not redirected) so it stays visible to callers that capture it.
|
||||
lora_log = lambda m: print(m, file=sys.stderr)
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
model = mod.load_dit(str(ckpt), T_lat=T_lat, dtype=dtype, compile_=False,
|
||||
lora_specs=lora_specs, num_steps=num_steps,
|
||||
lora_log=lora_log)
|
||||
return model, str(ckpt)
|
||||
|
||||
|
||||
def load_decoder(decoder_name: str, dtype):
|
||||
spec = DECODER_CHOICES[decoder_name]
|
||||
weights_path = ensure_local(spec[3])
|
||||
import importlib, io, contextlib
|
||||
mod = importlib.import_module(spec[0])
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
model = mod.load_model(weights_path=str(weights_path), dtype=dtype, compile_=False)
|
||||
chunk_fn = getattr(mod, spec[1])
|
||||
chunk_cfg = spec[2]
|
||||
return model, chunk_fn, chunk_cfg
|
||||
|
||||
|
||||
# Peak unified-memory tracker (works as both pre-0.34 mx.metal.* and new mx.* API)
|
||||
def _reset_peak_mem():
|
||||
fn = getattr(mx, "reset_peak_memory", None) or getattr(mx.metal, "reset_peak_memory", None)
|
||||
if fn: fn()
|
||||
|
||||
def _get_peak_mem_bytes() -> int:
|
||||
fn = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
|
||||
return int(fn()) if fn else 0
|
||||
|
||||
|
||||
def _free_to_pool():
|
||||
"""gc.collect() + flush MLX's memory cache so freed weights actually return to the OS."""
|
||||
import gc
|
||||
gc.collect()
|
||||
fn = getattr(mx, "clear_cache", None) or getattr(getattr(mx, "metal", None), "clear_cache", None)
|
||||
if fn: fn()
|
||||
|
||||
|
||||
_CUMULATIVE_PEAK_B = 0
|
||||
_STAGE_PEAKS: list[tuple[str, int]] = [] # [(stage_label, peak_bytes)]
|
||||
def _stage_peak_b(label: str | None = None) -> int:
|
||||
"""Read current peak, track cumulative max + per-stage log, reset for next stage."""
|
||||
global _CUMULATIVE_PEAK_B
|
||||
b = _get_peak_mem_bytes()
|
||||
_CUMULATIVE_PEAK_B = max(_CUMULATIVE_PEAK_B, b)
|
||||
if label is not None:
|
||||
_STAGE_PEAKS.append((label, b))
|
||||
_reset_peak_mem()
|
||||
return b
|
||||
|
||||
|
||||
def save_wav(path: str, audio: np.ndarray, sample_rate: int = SAMPLE_RATE):
|
||||
"""audio: (channels, T) float32 in [-1, 1]. Writes 16-bit PCM stereo WAV."""
|
||||
if not np.isfinite(audio).all():
|
||||
n_bad = int((~np.isfinite(audio)).sum())
|
||||
raise RuntimeError(f"refusing to write WAV — audio contains {n_bad} non-finite samples (NaN/Inf)")
|
||||
audio = np.clip(audio, -1.0, 1.0)
|
||||
pcm = (audio * 32767.0).astype(np.int16).T # (T, channels) interleaved
|
||||
with wave.open(path, "wb") as w:
|
||||
w.setnchannels(audio.shape[0])
|
||||
w.setsampwidth(2)
|
||||
w.setframerate(sample_rate)
|
||||
w.writeframes(pcm.tobytes())
|
||||
|
||||
|
||||
def read_wav(path: str) -> np.ndarray:
|
||||
"""Read a WAV file. Returns (2, T) float32 in [-1, 1].
|
||||
|
||||
Handles 16-bit PCM at 44.1 kHz natively. Falls back to ffmpeg for any
|
||||
other format (24-bit, 32-bit float, 48 kHz, etc.)
|
||||
Mono input is duplicated to stereo.
|
||||
"""
|
||||
try:
|
||||
with wave.open(path, "rb") as w:
|
||||
nch, sw, sr, nframes = w.getnchannels(), w.getsampwidth(), w.getframerate(), w.getnframes()
|
||||
if sr == SAMPLE_RATE and sw == 2:
|
||||
raw = np.frombuffer(w.readframes(nframes), dtype=np.int16).astype(np.float32) / 32767.0
|
||||
if nch == 1:
|
||||
return np.stack([raw, raw], axis=0) # (2, T)
|
||||
return raw.reshape(-1, nch).T[:2] # (2, T)
|
||||
except wave.Error:
|
||||
pass # unsupported format (32-bit float, 24-bit PCM, etc.)
|
||||
|
||||
# Fallback: decode via ffmpeg — handles any sample rate, bit depth, or format
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["ffmpeg", "-v", "error", "-i", path,
|
||||
"-f", "s16le", "-ar", str(SAMPLE_RATE), "-ac", "2", "-"],
|
||||
capture_output=True, check=True,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
raise RuntimeError(
|
||||
f"{path}: unsupported WAV format. Install ffmpeg to handle 24-bit/32-bit/48kHz audio:\n"
|
||||
f" brew install ffmpeg\n"
|
||||
f"Or convert manually: ffmpeg -i \"{path}\" -ar {SAMPLE_RATE} -ac 2 -sample_fmt s16 resampled.wav"
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
raise RuntimeError(f"{path}: ffmpeg failed — {e.stderr.decode().strip()}")
|
||||
raw = np.frombuffer(result.stdout, dtype=np.int16).astype(np.float32) / 32767.0
|
||||
return raw.reshape(-1, 2).T
|
||||
|
||||
|
||||
def load_encoder(decoder_name: str, dtype):
|
||||
"""Mirror load_decoder but for the matching encoder."""
|
||||
spec = ENCODER_CHOICES[decoder_name]
|
||||
weights_path = ensure_local(spec[2])
|
||||
import importlib, io, contextlib
|
||||
mod = importlib.import_module(spec[0])
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
model = mod.load_model(weights_path=str(weights_path), dtype=dtype, compile_=False)
|
||||
pad_modulo = spec[1]
|
||||
return model, pad_modulo
|
||||
|
||||
|
||||
def _preflight_download(args) -> None:
|
||||
"""Resolve every weight file this run needs and download any that are
|
||||
missing — BEFORE the banner prints and BEFORE the wall-clock starts.
|
||||
Network time then isn't charged against "×realtime" and the user sees
|
||||
download progress as a clearly separate setup step."""
|
||||
needed = [
|
||||
T5GEMMA_NPZ_REL,
|
||||
DIT_CHOICES[args.dit]["ckpt"],
|
||||
DECODER_CHOICES[args.decoder][3],
|
||||
]
|
||||
if args.init_audio:
|
||||
needed.append(ENCODER_CHOICES[args.decoder][2])
|
||||
missing = [p for p in needed if not is_present(p)]
|
||||
if not missing:
|
||||
return
|
||||
print(f" Fetching {len(missing)} missing weight file(s) before starting:")
|
||||
for rel in missing:
|
||||
ensure_local(rel)
|
||||
print()
|
||||
|
||||
|
||||
def patch_audio(audio: np.ndarray, patch_size: int = 256) -> np.ndarray:
|
||||
"""Patched-pretransform encode: (B, 2, T_audio) → (B, 512, T_audio/256).
|
||||
Mirrors rearrange("b c (l h) -> b (c h) l", h=patch_size)."""
|
||||
B, C, T = audio.shape
|
||||
assert T % patch_size == 0, f"audio length {T} not multiple of patch_size {patch_size}"
|
||||
L = T // patch_size
|
||||
x = audio.reshape(B, C, L, patch_size).transpose(0, 1, 3, 2) # (B, C, patch, L)
|
||||
return x.reshape(B, C * patch_size, L) # (B, 512, L)
|
||||
|
||||
|
||||
class _HelpfulParser(argparse.ArgumentParser):
|
||||
"""argparse that prints full help (not just usage) when a flag is unknown / invalid,
|
||||
and tacks the shared example-commands block onto the end of -h / --help."""
|
||||
def error(self, message):
|
||||
sys.stderr.write(f"\nerror: {message}\n\n")
|
||||
self.print_help(sys.stderr)
|
||||
sys.exit(2)
|
||||
def print_help(self, file=None):
|
||||
super().print_help(file)
|
||||
# Append the colored example-commands block. Same content as install.sh's
|
||||
# final summary; rendered to stdout regardless of `file` so the colors and
|
||||
# emojis don't get split between two streams when help goes to stderr.
|
||||
try:
|
||||
from examples import print_example_commands
|
||||
print_example_commands()
|
||||
except Exception:
|
||||
pass # never let an examples-block failure mask the actual --help
|
||||
|
||||
|
||||
def main():
|
||||
ap = _HelpfulParser(
|
||||
description="SA3 text-to-audio (+ audio-to-audio + inpainting) in pure MLX",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=(
|
||||
"modes\n"
|
||||
" text-to-audio --prompt P\n"
|
||||
" audio-to-audio --prompt P --init-audio IN.wav [--init-noise-level σ]\n"
|
||||
" inpainting --prompt P --init-audio IN.wav --inpaint-range START,END\n"
|
||||
" negative CFG --prompt P --cfg N --negative-prompt P_NEG\n"
|
||||
"\nrun `sa3_mlx.py --help` for per-flag details."
|
||||
),
|
||||
)
|
||||
# ── Inputs ────────────────────────────────────────────────────────────────
|
||||
ap.add_argument("--prompt", default=None,
|
||||
help="Text prompt describing the audio to generate. "
|
||||
"Empty string is valid (unconditional generation). "
|
||||
"If omitted, the script asks interactively via stdin.")
|
||||
ap.add_argument("--negative-prompt", default=None,
|
||||
help="Negative prompt for CFG's unconditional branch. "
|
||||
"When --cfg=1.0 this flag has no effect (no uncond pass is run). "
|
||||
"When unset and --cfg ≠ 1.0, the uncond branch uses zero embeddings.")
|
||||
ap.add_argument("--init-audio", default=None,
|
||||
help="Path to a WAV file (44.1 kHz, 16-bit PCM, stereo or mono) to use as the "
|
||||
"starting point. Enables audio-to-audio mode (with --init-noise-level) or "
|
||||
"inpainting mode (with --inpaint-range). Encoder is loaded automatically. "
|
||||
"Audio is trimmed or zero-padded to match --seconds.")
|
||||
ap.add_argument("--inpaint-range", default=None,
|
||||
help="Inpainting time range as 'START,END' in seconds (e.g. '5,10'). "
|
||||
"Requires --init-audio. The model regenerates the masked range while "
|
||||
"preserving the rest of the input exactly (via paste-back).")
|
||||
|
||||
# ── Models ────────────────────────────────────────────────────────────────
|
||||
ap.add_argument("--dit", choices=list(DIT_CHOICES.keys()), default=None,
|
||||
help="DiT model. 'small' = sa3-sm-music (faster, smaller). "
|
||||
"'medium' = sa3-medium-ARC (larger, higher quality). "
|
||||
"If omitted, prompts interactively with arrow-key picker.")
|
||||
ap.add_argument("--decoder", choices=list(DECODER_CHOICES.keys()), default=None,
|
||||
help="Audio decoder. 'same-s' pairs with sa3-small (50 M params). "
|
||||
"'same-l' pairs with sa3-medium (426 M params). "
|
||||
"If omitted, prompts interactively with arrow-key picker.")
|
||||
ap.add_argument("--dit-dtype", default="fp16", choices=["fp32", "fp16"],
|
||||
help="DiT compute dtype. Default fp16 — validated transparent at ~50-57 dB PSNR "
|
||||
"vs FP32. Halves DiT memory and speeds sampling ~25%%. The decoder always "
|
||||
"runs FP32 (SAME-S catastrophically cancels at fp16 due to differential "
|
||||
"attention) and T5Gemma is always fp16. Set to fp32 only if you need "
|
||||
"bit-exact reproducibility against the PyTorch reference.")
|
||||
ap.add_argument("--t5gemma-npz", default=None,
|
||||
help="Path to the bundled T5Gemma FP16 .npz (weights + tokenizer). "
|
||||
"Default points at models/mlx/t5gemma_f16.npz next to this script; "
|
||||
"auto-downloaded from HuggingFace if not present.")
|
||||
ap.add_argument("--lora", action="append", nargs="+", default=None,
|
||||
metavar=("ADAPTER", "KEY=VAL"),
|
||||
help="A LoRA adapter to apply to the DiT — repeat the flag for several. "
|
||||
"Each --lora takes the adapter path followed by optional key=value "
|
||||
"options: strength=S (default --lora-strength) and steps=RANGE, a "
|
||||
"1-based inclusive sampling-step range: 2-8, 2- (2..last), -4 (1..4) "
|
||||
"or 3 (just step 3; default: all steps). Example: "
|
||||
"--lora plini.safetensors strength=0.8 steps=2- . Adapters covering "
|
||||
"every step are merged at load; step-gated ones are re-merged in "
|
||||
"place at the (few) step boundaries — ~80 ms each on medium, no "
|
||||
"extra memory, every step runs at full speed. The adapter is a "
|
||||
".safetensors file (SA3-native train_lora.py / underfit output) or "
|
||||
"a PEFT adapter directory (with its adapter_config.json). ONLY "
|
||||
".safetensors is accepted — a pickle .ckpt/.pt is refused (it would "
|
||||
"execute code on load). The adapter's base must match --dit.")
|
||||
ap.add_argument("--lora-strength", type=float, default=1.0,
|
||||
help="Default strength for --lora adapters without their own strength= "
|
||||
"(default 1.0). 0 disables (bit-identical to no LoRA); >1 amplifies.")
|
||||
|
||||
# ── Sampling ──────────────────────────────────────────────────────────────
|
||||
ap.add_argument("--seconds", type=float, default=30.0,
|
||||
help="Output audio length in seconds. T_lat (latent positions) is derived as "
|
||||
"ceil(seconds * 44100 / 4096) — a natural ceil that depends ONLY on "
|
||||
"--seconds (and --seed via the DiT), never on --decoder. This keeps the "
|
||||
"sampled latent (and hence the music) identical across decoders for a "
|
||||
"given prompt/seed. Final WAV is trimmed to exactly --seconds.")
|
||||
ap.add_argument("--steps", type=int, default=8,
|
||||
help="Number of pingpong sampling steps. Minimum 1 (single forward pass — fastest, "
|
||||
"lowest quality). rf_denoiser is distilled for 8 (default — sweet spot). "
|
||||
">8 gives diminishing returns and may overshoot. The LogSNR schedule "
|
||||
"auto-computes for any N: steps=N produces N+1 sigma values from σmax→0. "
|
||||
"Sample wall time scales ~linearly with steps; quality/coherence improves "
|
||||
"noticeably from 1→4 and 4→8, less from 8→16.")
|
||||
ap.add_argument("--seed", type=int, default=None,
|
||||
help="Random seed (any int). Set this for reproducible outputs. "
|
||||
"If omitted, a random seed is chosen and printed in the final 'done' line.")
|
||||
ap.add_argument("--init-noise-level", type=float, default=1.0,
|
||||
help="σmax — the schedule's starting noise level (always honored regardless of mode). "
|
||||
"Valid range: [0.01, ∞). Below 0.01 the script errors because the model is "
|
||||
"undefined at t≈0 and produces NaN. With --init-audio: 0.4-0.8 is typical for "
|
||||
"variation; 1.0 = full regeneration (init ignored). Without --init-audio: 1.0 "
|
||||
"is standard text-to-audio; <1.0 is a creative effect (model sees pure noise but "
|
||||
"with t=σmax timestep — mildly OOD); >1.0 is 'guidance overshoot' (more diverse, "
|
||||
"increasingly weird).")
|
||||
ap.add_argument("--cfg", type=float, default=1.0,
|
||||
help="Classifier-Free Guidance scale. 1.0 = off (single forward pass, the rf_denoiser "
|
||||
"default). >1.0 pushes toward the prompt (classic CFG). [0, 1) pulls toward the "
|
||||
"unconditional / negative branch (less prompt-aligned, more diverse). <0 actively "
|
||||
"pushes AWAY from the prompt. Any value ≠ 1.0 costs ~2× per step (batched cond + "
|
||||
"uncond forward).")
|
||||
ap.add_argument("--apg", type=float, default=1.0,
|
||||
help="Adaptive Projected Guidance scale [0..1], only matters when --cfg ≠ 1.0. "
|
||||
"1.0 = full APG (project the cond−uncond difference orthogonal to cond_denoised; "
|
||||
"prevents over-saturation at high CFG). 0.0 = vanilla CFG (use the full "
|
||||
"difference). Intermediate values blend the two. rf_denoiser default is 1.0.")
|
||||
|
||||
# ── Memory / runtime ──────────────────────────────────────────────────────
|
||||
ap.add_argument("--free-models", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="Free each model after its last use (T5Gemma after step 2, encoder after "
|
||||
"step 3a, DiT after step 3b) to minimize peak RAM. Lowers decode-stage peak "
|
||||
"by 3-5 GB on sa3-medium since the DiT (2.7 GB FP16) doesn't sit idle while "
|
||||
"the decoder runs. Use --no-free-models to keep everything resident — only "
|
||||
"useful if you plan to call models multiple times in one run (currently the "
|
||||
"script doesn't). Default: on.")
|
||||
|
||||
# ── Output ────────────────────────────────────────────────────────────────
|
||||
ap.add_argument("--out", default=None,
|
||||
help="Output WAV path. Relative paths land in the project's output/ "
|
||||
"directory (auto-created); absolute paths are used as-is. "
|
||||
"Always written as 16-bit PCM stereo at 44.1 kHz, trimmed to "
|
||||
"exactly --seconds. If omitted, auto-named from the prompt and seed.")
|
||||
ap.add_argument("--play", action="store_true",
|
||||
help="After writing the WAV, play it through the default output device "
|
||||
"via the macOS `afplay` binary. Blocking — the script exits when "
|
||||
"playback finishes. Ctrl-C stops playback and exits the script "
|
||||
"(SIGINT is delivered to both processes).")
|
||||
args = ap.parse_args()
|
||||
if args.steps < 1:
|
||||
ap.error(f"--steps must be ≥ 1 (got {args.steps})")
|
||||
|
||||
# Parse --lora groups into specs (fail fast, before any model loads).
|
||||
args.lora_specs = None
|
||||
if args.lora:
|
||||
from models.defs.lora_merge import parse_lora_spec, LoraError
|
||||
try:
|
||||
args.lora_specs = [parse_lora_spec(toks, args.lora_strength)
|
||||
for toks in args.lora]
|
||||
except LoraError as e:
|
||||
ap.error(str(e))
|
||||
|
||||
args = prompt_user_if_missing(args)
|
||||
if args.prompt is None:
|
||||
args.prompt = input("Prompt: ").strip()
|
||||
|
||||
# Resolve output path — auto-name from prompt+seed when --out is not given.
|
||||
if args.out is None:
|
||||
import re as _re
|
||||
slug = _re.sub(r'[^a-z0-9]+', '_', args.prompt.lower()).strip('_')[:48]
|
||||
args.out = f"{slug}_{args.seed}.wav" if slug else f"out_{args.seed}.wav"
|
||||
out_path = Path(args.out)
|
||||
if not out_path.is_absolute():
|
||||
out_path = REPO / "output" / out_path
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.out = str(out_path)
|
||||
# Empty prompt is allowed — T5Gemma will produce padding-only embeddings,
|
||||
# which (with the learned padding_embedding) is the unconditional case.
|
||||
dtype = mx.float32 if args.dit_dtype == "fp32" else mx.float16
|
||||
# T_lat is a DECODER-INDEPENDENT function of --seconds only (natural ceil, no
|
||||
# even-bump). This matches the TensorRT pipeline (sa3_trt.py::resolve_T_lat) so
|
||||
# the same prompt/seed/seconds yields the SAME latent — and hence the same music —
|
||||
# regardless of --decoder. The old code bumped T_lat to even for same-s, which made
|
||||
# make_noise(T_lat, seed) draw a different noise tensor and thus sample a completely
|
||||
# different latent just by switching decoder. SAME-S's internal-length constraint
|
||||
# (T_lat*17 must align to 34, i.e. even) is satisfied by the decode dispatch below:
|
||||
# odd T_lat is routed through decode_chunked, whose every window is an even kernel.
|
||||
T_lat = max(1, math.ceil(args.seconds * SAMPLE_RATE / SAMPLES_PER_LATENT))
|
||||
target_dur = T_lat * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||||
|
||||
# Inpaint validation + parameter mapping
|
||||
inpaint_range = None
|
||||
if args.inpaint_range is not None:
|
||||
if args.init_audio is None:
|
||||
sys.exit("error: --inpaint-range requires --init-audio (the audio to inpaint into)")
|
||||
try:
|
||||
s_str, e_str = args.inpaint_range.split(",")
|
||||
inp_start_sec = float(s_str.strip()); inp_end_sec = float(e_str.strip())
|
||||
except ValueError:
|
||||
sys.exit(f"error: --inpaint-range must be 'START,END' in seconds; got {args.inpaint_range!r}")
|
||||
if not (0 <= inp_start_sec < inp_end_sec <= args.seconds):
|
||||
sys.exit(f"error: invalid inpaint range {inp_start_sec}-{inp_end_sec}s "
|
||||
f"(must satisfy 0 <= start < end <= {args.seconds}s)")
|
||||
inp_start_lat = max(0, int(round(inp_start_sec * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||||
inp_end_lat = min(T_lat, int(round(inp_end_sec * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||||
inpaint_range = (inp_start_lat, inp_end_lat)
|
||||
|
||||
# σmax is always taken from --init-noise-level, regardless of mode. The mode
|
||||
# only changes how σmax is *used* (init mix vs schedule-only vs inpaint init).
|
||||
sigma_max = float(args.init_noise_level)
|
||||
mode = ("inpaint" if inpaint_range else
|
||||
"audio-to-audio" if args.init_audio else "text-to-audio")
|
||||
|
||||
# Reject very low σmax in any mode — the diffusion model was never trained
|
||||
# at t≈0 and produces NaN/garbage there.
|
||||
MIN_SIGMA = 0.01
|
||||
if sigma_max < MIN_SIGMA:
|
||||
sys.exit(
|
||||
f"error: --init-noise-level={sigma_max} is too low (σmax min is {MIN_SIGMA}). "
|
||||
f"The rf_denoiser model is undefined at t≈0; below {MIN_SIGMA} the schedule "
|
||||
f"collapses and the model emits NaN. Pass --init-noise-level=1.0 for normal "
|
||||
f"text-to-audio, or use the input WAV directly if you want it unchanged."
|
||||
)
|
||||
|
||||
# ── Preflight: ensure every weight file this run needs is on disk BEFORE
|
||||
# we print the banner or start the wall-time clock. ensure_local() prints
|
||||
# a "↓ downloading …" line per missing file and symlinks from the HF cache.
|
||||
_preflight_download(args)
|
||||
|
||||
global _CUMULATIVE_PEAK_B
|
||||
_CUMULATIVE_PEAK_B = 0
|
||||
_STAGE_PEAKS.clear()
|
||||
_reset_peak_mem()
|
||||
t_wall_start = time.time()
|
||||
|
||||
print()
|
||||
banner(f"SA3 → MLX {mode}")
|
||||
k = lambda s: dim(f"{s:>10}")
|
||||
v = lambda s, w=10: f"{s:<{w}}"
|
||||
print(f" {k('prompt')} {bold(repr(args.prompt))}")
|
||||
if args.negative_prompt:
|
||||
suffix = "" if args.cfg != 1.0 else dim(" (ignored: --cfg=1.0, no uncond branch)")
|
||||
print(f" {k('neg prompt')} {bold(repr(args.negative_prompt))}{suffix}")
|
||||
dit_decoder_line = f" {k('dit')} {magenta(v(args.dit))} {k('decoder')} {magenta(v(args.decoder))}"
|
||||
if args.init_audio:
|
||||
dit_decoder_line += f" {k('encoder')} {magenta(v(args.decoder))}"
|
||||
print(dit_decoder_line)
|
||||
if args.init_audio:
|
||||
print(f" {k('init audio')} {bold(args.init_audio)}")
|
||||
if inpaint_range:
|
||||
s0, s1 = inpaint_range
|
||||
print(f" {k('inpaint')} {bold(f'{inp_start_sec:.2f}s..{inp_end_sec:.2f}s')} "
|
||||
f"{dim(f'(latent {s0}..{s1} of {T_lat})')}")
|
||||
print(f" {k('σmax')} {bold(f'{sigma_max:.2f}')}")
|
||||
print(f" {k('seconds')} {v(f'{args.seconds}s')} {k('steps')} {v(args.steps)} {k('seed')} {args.seed}")
|
||||
cfg_label = f"{args.cfg}"
|
||||
if args.cfg != 1.0:
|
||||
cfg_label += f" (apg={args.apg})"
|
||||
print(f" {k('DiT dtype')} {v(args.dit_dtype)} {k('cfg')} {cfg_label}")
|
||||
print(f" {k('T_lat')} {T_lat} {dim(f'({target_dur:.2f}s → trimmed to {args.seconds}s)')}")
|
||||
print()
|
||||
|
||||
# ── 1. Text encoder ──
|
||||
t0 = time.time()
|
||||
t5_path = args.t5gemma_npz if args.t5gemma_npz else str(ensure_local(T5GEMMA_NPZ_REL))
|
||||
enc = T5Gemma.from_npz(t5_path)
|
||||
embeds, mask = enc.encode([args.prompt], max_len=256)
|
||||
mx.eval(embeds, mask)
|
||||
stage("[1/5]", "T5Gemma encode", (time.time()-t0)*1000, peak_b=_stage_peak_b("T5Gemma encode"))
|
||||
sub(f"embeds {embeds.shape} {embeds.dtype}")
|
||||
|
||||
# ── 2. Conditioning ──
|
||||
# Conditioner weights are baked into the DiT .npz under "cond." prefix.
|
||||
t0 = time.time()
|
||||
padding_emb, secs_embedder = load_conditioner_from_npz(
|
||||
str(ensure_local(DIT_CHOICES[args.dit]["ckpt"])), prefix="cond.")
|
||||
if args.lora_specs:
|
||||
# The conditioner is loaded straight from the npz, so the DiT-side merge
|
||||
# never reaches it — apply any seconds-conditioner deltas (underfit
|
||||
# adapters carry one) to the live embedder here. Whole-generation:
|
||||
# conditioning runs once, before sampling, so steps= cannot gate it.
|
||||
from models.defs.lora_merge import apply_conditioner_lora
|
||||
secs_embedder.W, _n_cond = apply_conditioner_lora(secs_embedder.W,
|
||||
args.lora_specs)
|
||||
if _n_cond:
|
||||
sub(f"lora conditioner delta applied ({_n_cond} adapter(s), whole generation)")
|
||||
embeds = embeds.astype(dtype)
|
||||
embeds_padded = apply_prompt_padding(embeds, mask, padding_emb.astype(dtype))
|
||||
seconds_embed = secs_embedder(args.seconds).astype(dtype) # (1, 1, 768)
|
||||
cross_attn = mx.concatenate([embeds_padded, seconds_embed], axis=1) # (1, 257, 768)
|
||||
global_cond = seconds_embed[:, 0, :] # (1, 768)
|
||||
|
||||
# Negative / null cross_attn for CFG (used whenever cfg != 1.0)
|
||||
null_cross_attn = None
|
||||
if args.cfg != 1.0:
|
||||
if args.negative_prompt:
|
||||
neg_embeds, neg_mask = enc.encode([args.negative_prompt], max_len=256)
|
||||
mx.eval(neg_embeds, neg_mask)
|
||||
neg_embeds = neg_embeds.astype(dtype)
|
||||
neg_padded = apply_prompt_padding(neg_embeds, neg_mask, padding_emb.astype(dtype))
|
||||
null_cross_attn = mx.concatenate([neg_padded, seconds_embed], axis=1)
|
||||
else:
|
||||
# Match upstream: zeros_like(cross_attn) for the uncond branch.
|
||||
null_cross_attn = mx.zeros_like(cross_attn)
|
||||
mx.eval(null_cross_attn)
|
||||
|
||||
mx.eval(cross_attn, global_cond)
|
||||
stage("[2/5]", "Conditioning", (time.time()-t0)*1000, peak_b=_stage_peak_b("Conditioning"))
|
||||
sub(f"cross_attn {cross_attn.shape} global {global_cond.shape}"
|
||||
+ (f" neg_cross_attn ready ({'prompt' if args.negative_prompt else 'zeros'})" if null_cross_attn is not None else ""))
|
||||
|
||||
# T5Gemma weights are no longer needed (we have the embeddings). Free 537 MB.
|
||||
if args.free_models:
|
||||
del enc
|
||||
_free_to_pool()
|
||||
|
||||
# ── 3a. (audio-to-audio only) Encode init_audio → init_latents ──
|
||||
init_latents = None
|
||||
if args.init_audio:
|
||||
stage("[3a]", f"Encoding init audio → latents")
|
||||
t0 = time.time()
|
||||
enc_model, pad_mod = load_encoder(args.decoder, mx.float32)
|
||||
sub(f"encoder load {(time.time()-t0)*1000:.0f} ms")
|
||||
|
||||
t0 = time.time()
|
||||
# The ENCODER (SAME-S) needs T_audio_patches divisible by `pad_mod` (32) → an
|
||||
# even latent length. T_lat itself stays the natural-ceil (decoder-independent)
|
||||
# value; we only round the encode grid UP to the modulo, then trim the latent
|
||||
# back to T_lat. This keeps the DiT/noise path identical across decoders.
|
||||
enc_T_lat = T_lat
|
||||
if (T_lat * 16) % pad_mod != 0:
|
||||
enc_T_lat = math.ceil((T_lat * 16) / pad_mod) * pad_mod // 16
|
||||
target_samples = enc_T_lat * SAMPLES_PER_LATENT
|
||||
audio_np = read_wav(args.init_audio) # (2, T_audio_in)
|
||||
if audio_np.shape[-1] >= target_samples:
|
||||
audio_np = audio_np[:, :target_samples]
|
||||
init_action = f"trimmed to {target_samples} samples"
|
||||
else:
|
||||
pad = target_samples - audio_np.shape[-1]
|
||||
audio_np = np.pad(audio_np, ((0, 0), (0, pad)), mode="constant")
|
||||
init_action = f"zero-padded by {pad} samples"
|
||||
audio_np = audio_np[None, ...] # (1, 2, T)
|
||||
sub(f"read+prep ({init_action}) {(time.time()-t0)*1000:.0f} ms")
|
||||
|
||||
# Patch + encode (encoder always runs FP32 — softnorm-bottleneck-sensitive)
|
||||
t0 = time.time()
|
||||
patches_np = patch_audio(audio_np, patch_size=256) # (1, 512, enc_T_lat*16)
|
||||
# Sanity: T_audio_patches must be divisible by encoder's required modulo
|
||||
assert patches_np.shape[-1] % pad_mod == 0, (
|
||||
f"T_audio_patches={patches_np.shape[-1]} not divisible by {pad_mod} "
|
||||
f"(decoder={args.decoder})"
|
||||
)
|
||||
init_latents = enc_model(mx.array(patches_np))[..., :T_lat] # trim to natural T_lat
|
||||
mx.eval(init_latents)
|
||||
_stage_peak_b('Init audio encode')
|
||||
sub(f"encode {(time.time()-t0)*1000:.0f} ms latents {init_latents.shape}")
|
||||
init_latents = init_latents.astype(dtype)
|
||||
del enc_model
|
||||
|
||||
# ── 3b. DiT pingpong sampling ──
|
||||
stage("[3/5]", f"DiT — load + sample ({args.steps} steps, σmax={sigma_max:.2f})")
|
||||
if args.lora_specs:
|
||||
for s in args.lora_specs:
|
||||
rng = s["steps"]
|
||||
rng_str = ("all steps" if rng is None else
|
||||
f"steps {rng[0] or 1}-{rng[1] if rng[1] is not None else args.steps}")
|
||||
sub(f"lora {os.path.basename(s['path'].rstrip('/'))} "
|
||||
f"(strength {s['strength']:g}, {rng_str})")
|
||||
t0 = time.time(); dit_model, _ = load_dit(args.dit, T_lat=T_lat, dtype=dtype,
|
||||
lora_specs=args.lora_specs, num_steps=args.steps)
|
||||
_stage_peak_b('DiT load')
|
||||
sub(f"load {time.time()-t0:.1f}s")
|
||||
|
||||
sigmas = build_pingpong_schedule(args.steps, sigma_max=sigma_max, use_logsnr_shift=True)
|
||||
sched_str = " · ".join(f"{float(x):.3f}" for x in sigmas)
|
||||
sub(f"schedule {sched_str}")
|
||||
|
||||
key = mx.random.key(args.seed)
|
||||
pure_noise = mx.random.normal((1, 256, T_lat), dtype=dtype, key=key)
|
||||
if init_latents is not None and inpaint_range is None:
|
||||
# rf_denoiser init_data mix (linear interpolation, matches upstream):
|
||||
# noise = init_data * (1 - σmax) + pure_noise * σmax
|
||||
noise = init_latents * (1.0 - sigma_max) + pure_noise * sigma_max
|
||||
sub(f"init: latent * {1-sigma_max:.2f} + noise * {sigma_max:.2f}")
|
||||
else:
|
||||
noise = pure_noise
|
||||
mx.eval(noise)
|
||||
|
||||
# Build local_add_cond for inpainting: [B, T_lat, 257] = (mask, masked_input)
|
||||
local_add_cond = None
|
||||
paste_back = None
|
||||
if inpaint_range is not None:
|
||||
s0, s1 = inpaint_range
|
||||
# mask: 1 = keep (context), 0 = inpaint (regenerate); shape (1, 1, T_lat)
|
||||
mask_np = np.ones((1, 1, T_lat), dtype=np.float32)
|
||||
mask_np[:, :, s0:s1] = 0.0
|
||||
mask = mx.array(mask_np)
|
||||
# masked_input: init_latents with the inpaint region zeroed
|
||||
masked_input = init_latents.astype(mx.float32) * mask # (1, 256, T_lat)
|
||||
# concat over channel dim → (1, 257, T_lat); DiT expects batch-last-channel → transpose
|
||||
lac = mx.concatenate([mask, masked_input], axis=1).transpose(0, 2, 1).astype(dtype)
|
||||
local_add_cond = lac
|
||||
# Paste-back keeps unchanged regions bit-exact (model preserves them well but
|
||||
# paste-back makes it guaranteed).
|
||||
paste_back = (init_latents, mask)
|
||||
sub(f"local_add_cond {tuple(lac.shape)} inpaint mask: {s0}..{s1} of {T_lat} latents "
|
||||
f"({(s1-s0)/T_lat*100:.0f}% regenerated)")
|
||||
|
||||
def model_fn(x, t):
|
||||
if args.cfg == 1.0:
|
||||
return dit_model(x, t, cross_attn, global_cond, local_add_cond=local_add_cond)
|
||||
|
||||
# Batched CFG: one forward over cat([x, x]) on the batch dim.
|
||||
x2 = mx.concatenate([x, x], axis=0)
|
||||
t2 = mx.concatenate([t, t], axis=0)
|
||||
cross2 = mx.concatenate([cross_attn, null_cross_attn], axis=0)
|
||||
global2 = mx.concatenate([global_cond, global_cond], axis=0)
|
||||
lac2 = None if local_add_cond is None else mx.concatenate([local_add_cond, local_add_cond], axis=0)
|
||||
v_batched = dit_model(x2, t2, cross2, global2, local_add_cond=lac2)
|
||||
cond_v, uncond_v = mx.split(v_batched, 2, axis=0)
|
||||
|
||||
# Convert to denoised space (RF: denoised = x - σ * v, with σ = t)
|
||||
sigma = t.reshape(-1, 1, 1).astype(mx.float32)
|
||||
cond_d = x.astype(mx.float32) - cond_v.astype(mx.float32) * sigma
|
||||
uncond_d = x.astype(mx.float32) - uncond_v.astype(mx.float32) * sigma
|
||||
diff = cond_d - uncond_d
|
||||
|
||||
if args.apg <= 0.0:
|
||||
cfg_diff = diff
|
||||
else:
|
||||
# Project diff onto direction orthogonal to cond_d (per-sample, over (C, T)).
|
||||
# fp32 throughout — fp16 has catastrophic-cancellation risk in the projection.
|
||||
norm = mx.sqrt((cond_d * cond_d).sum(axis=(-2, -1), keepdims=True))
|
||||
unit = cond_d / mx.maximum(norm, 1e-8)
|
||||
parallel = (diff * unit).sum(axis=(-2, -1), keepdims=True) * unit
|
||||
diff_orth = diff - parallel
|
||||
cfg_diff = diff_orth if args.apg >= 1.0 else (args.apg * diff_orth + (1.0 - args.apg) * diff)
|
||||
|
||||
cfg_d = cond_d + (args.cfg - 1.0) * cfg_diff
|
||||
cfg_v = (x.astype(mx.float32) - cfg_d) / sigma
|
||||
return cfg_v.astype(x.dtype)
|
||||
|
||||
# Live in-place progress while sampling (TTY only — quiet log files).
|
||||
t_step_prev = [time.time()]
|
||||
def _on_step(i: int, total: int):
|
||||
if not _USE_COLOR: # _USE_COLOR ≡ stdout.isatty() set at top of this file
|
||||
return
|
||||
now = time.time()
|
||||
elapsed = (now - t_step_prev[0]) * 1000
|
||||
t_step_prev[0] = now
|
||||
bar_w = 20
|
||||
filled = int(round(bar_w * i / total))
|
||||
bar = cyan("█" * filled) + dim("·" * (bar_w - filled))
|
||||
sys.stdout.write(f"\r\x1b[K {dim('sampling')} {bar} "
|
||||
f"{bold(f'step {i}/{total}')} {yellow(f'{elapsed:.0f} ms')}")
|
||||
sys.stdout.flush()
|
||||
|
||||
t0 = time.time()
|
||||
_lora_plan = getattr(dit_model, "_lora_plan", None)
|
||||
latents = sample_flow_pingpong(model_fn, noise, sigmas, seed=args.seed + 1,
|
||||
paste_back=paste_back, on_step=_on_step,
|
||||
before_step=_lora_plan.sync if _lora_plan else None)
|
||||
mx.eval(latents)
|
||||
sample_ms = (time.time()-t0)*1000
|
||||
# Clear the progress line, then print the final summary in its place.
|
||||
if _USE_COLOR:
|
||||
sys.stdout.write("\r\x1b[K")
|
||||
_stage_peak_b('DiT sample')
|
||||
sub(f"sample {sample_ms:.0f} ms ({sample_ms/max(args.steps,1):.0f} ms/step)")
|
||||
|
||||
# DiT weights are no longer needed (latents are computed). Free 2.7 GB for sa3-medium.
|
||||
if args.free_models:
|
||||
del dit_model, model_fn
|
||||
_free_to_pool()
|
||||
|
||||
# ── 4. Decode latents → audio patches ──
|
||||
# Decoder always fp32 — SAME-S needs it (FP16 catastrophically cancels at differential attn);
|
||||
# SAME-L works at fp16 but we use fp32 for max quality.
|
||||
stage("[4/5]", f"Decoder ({args.decoder}, FP32)")
|
||||
t0 = time.time()
|
||||
dec_dtype = mx.float32
|
||||
decoder, chunk_fn, (chunk, ovl) = load_decoder(args.decoder, dec_dtype)
|
||||
_stage_peak_b('Decoder load')
|
||||
sub(f"load {(time.time()-t0)*1000:.0f} ms")
|
||||
latents_fp32 = latents.astype(dec_dtype)
|
||||
t0 = time.time()
|
||||
kernel = chunk + 2 * ovl
|
||||
if T_lat > kernel:
|
||||
# Natural-ceil T_lat may be odd; decode_chunked windows are all even kernels,
|
||||
# so odd T_lat decodes correctly here (this is the common path, incl. 30s→323).
|
||||
patches = chunk_fn(decoder, latents_fp32, chunk, ovl)
|
||||
decode_mode = f"chunked (chunk={chunk}, ovl={ovl})"
|
||||
elif T_lat % 2 == 0:
|
||||
patches = decoder(latents_fp32)
|
||||
decode_mode = "un-chunked"
|
||||
elif T_lat > 6:
|
||||
# Odd T_lat in (6, kernel]: a chunk=2,ovl=2 kernel (=6 < T_lat) keeps every
|
||||
# window even for SAME-S while never zero-padding real latents.
|
||||
patches = chunk_fn(decoder, latents_fp32, 2, 2)
|
||||
decode_mode = "chunked (chunk=2, ovl=2)"
|
||||
else:
|
||||
# Tiny odd T_lat ≤ 6 (sub-0.5s): no even chunk kernel fits. Reflect-pad one
|
||||
# latent to make T even, un-chunked decode, then trim the extra 16 patches.
|
||||
# SAME-L has no even constraint and takes this path only for symmetry.
|
||||
latents_even = mx.concatenate([latents_fp32, latents_fp32[..., -1:]], axis=-1)
|
||||
patches = decoder(latents_even)[..., : T_lat * 16]
|
||||
decode_mode = "un-chunked (odd-pad→trim)"
|
||||
mx.eval(patches)
|
||||
_stage_peak_b('Decode')
|
||||
sub(f"decode {decode_mode} → {(time.time()-t0)*1000:.0f} ms patches {patches.shape}")
|
||||
|
||||
# ── 5. Unpatch → audio + save WAV ──
|
||||
t0 = time.time()
|
||||
audio = patched_decode(patches, patch_size=256, channels=2) # (1, 2, T_lat*4096)
|
||||
mx.eval(audio)
|
||||
audio_np = np.array(audio.astype(mx.float32))[0] # (2, T_lat*4096)
|
||||
requested_samples = int(round(args.seconds * SAMPLE_RATE))
|
||||
if audio_np.shape[-1] > requested_samples:
|
||||
audio_np = audio_np[..., :requested_samples]
|
||||
save_wav(args.out, audio_np)
|
||||
stage("[5/5]", "Unpatch + write WAV", (time.time()-t0)*1000, peak_b=_stage_peak_b("Unpatch + WAV"))
|
||||
peak = float(np.abs(audio_np).max()); rms = float(np.sqrt((audio_np**2).mean()))
|
||||
sub(f"audio {audio_np.shape} peak {peak:.3f} rms {rms:.3f}")
|
||||
|
||||
# ── Totals ──
|
||||
total = time.time() - t_wall_start
|
||||
audio_dur = audio_np.shape[-1] / SAMPLE_RATE
|
||||
peak_str = _fmt_mem(_CUMULATIVE_PEAK_B)
|
||||
|
||||
# Per-stage peak RAM table
|
||||
print()
|
||||
print(f" {bold('Peak RAM by stage')}")
|
||||
name_w = max(len(n) for n, _ in _STAGE_PEAKS) if _STAGE_PEAKS else 0
|
||||
sorted_for_emphasis = max((b for _, b in _STAGE_PEAKS), default=0)
|
||||
for name, b in _STAGE_PEAKS:
|
||||
bar_units = int(round(b / max(sorted_for_emphasis, 1) * 24))
|
||||
bar = cyan("█" * bar_units) + dim("·" * (24 - bar_units))
|
||||
mark = bold(" ← max") if b == sorted_for_emphasis else ""
|
||||
print(f" {dim(name.ljust(name_w))} {bar} {_fmt_mem(b)}{mark}")
|
||||
|
||||
print()
|
||||
rule()
|
||||
print(f" {bold(green('done'))} "
|
||||
f"{bold(f'{total:.2f}s')} wall → "
|
||||
f"{audio_dur:.1f}s audio → "
|
||||
f"{bold(yellow(f'{audio_dur/total:.2f}× realtime'))} "
|
||||
f"{dim(f'peak RAM {peak_str}')} "
|
||||
f"{dim(f'seed {args.seed}')}")
|
||||
# Prominent path block: show relative-to-cwd if shorter (more readable when
|
||||
# the user is already in the project), absolute otherwise (more copy-pasteable).
|
||||
abs_out = os.path.abspath(args.out)
|
||||
try:
|
||||
rel_out = os.path.relpath(abs_out)
|
||||
except ValueError:
|
||||
rel_out = abs_out
|
||||
shown = rel_out if len(rel_out) <= len(abs_out) and not rel_out.startswith("..") else abs_out
|
||||
print(f" {bold(green('▸ saved'))} {bold(shown)} {dim(f'({abs_out})' if shown != abs_out else '')}".rstrip())
|
||||
rule()
|
||||
|
||||
if args.play:
|
||||
try:
|
||||
print(f" {bold('▶ playing')} {args.out} {dim('(Ctrl-C to stop)')}")
|
||||
subprocess.run(["afplay", args.out], check=False)
|
||||
except KeyboardInterrupt:
|
||||
print() # newline after the ^C echo
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Mel-spectrogram renderer — numpy-only port of the TensorRT release's spec.py.
|
||||
|
||||
Same algorithm and constants as optimized/tensorRT/scripts/spec.py (the 3-band
|
||||
tinted stereo mel spectrogram from our run dashboards), with the torch pieces
|
||||
(STFT, resample, filterbank matmul) rewritten in numpy so the MLX release stays
|
||||
torch-free. Numerics match: periodic Hann window, center=True reflect padding,
|
||||
Slaney mel scale, librosa-compatible filterbank, linear resample with
|
||||
align_corners=False sampling positions.
|
||||
|
||||
Exposes render_spectrogram_png(pcm_int16, sample_rate, width, height).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import io
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# 3-band tint
|
||||
SPEC_BANDS = [
|
||||
(0, 200, (1.0, 0.0, 0.0)), # Bass -> Red
|
||||
(200, 1500, (0.0, 1.0, 0.0)), # Mid -> Green
|
||||
(1500, 16000, (0.0, 0.0, 1.0)), # High -> Blue
|
||||
]
|
||||
_BAND_COLORS = np.array([c for _, _, c in SPEC_BANDS], dtype=np.float32)
|
||||
|
||||
# Slaney mel-scale constants (linear below 1 kHz, log above), matches librosa default.
|
||||
_F_SP = 200.0 / 3
|
||||
_MIN_LOG_HZ = 1000.0
|
||||
_MIN_LOG_MEL = _MIN_LOG_HZ / _F_SP
|
||||
_LOGSTEP = math.log(6.4) / 27.0
|
||||
|
||||
|
||||
def _hz_to_mel(hz):
|
||||
hz = np.asarray(hz, dtype=np.float64)
|
||||
log_term = np.log(np.maximum(hz, _MIN_LOG_HZ) / _MIN_LOG_HZ) / _LOGSTEP
|
||||
return np.where(hz >= _MIN_LOG_HZ, _MIN_LOG_MEL + log_term, hz / _F_SP)
|
||||
|
||||
|
||||
def _mel_to_hz(mels):
|
||||
mels = np.asarray(mels, dtype=np.float64)
|
||||
return np.where(mels >= _MIN_LOG_MEL,
|
||||
_MIN_LOG_HZ * np.exp(_LOGSTEP * (mels - _MIN_LOG_MEL)),
|
||||
_F_SP * mels)
|
||||
|
||||
|
||||
def _mel_frequencies(n_mels, fmax, fmin=0.0):
|
||||
return _mel_to_hz(np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels))
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _mel_filterbank(n_mels, n_fft, sr, fmax):
|
||||
pts = _mel_to_hz(np.linspace(_hz_to_mel(0.0), _hz_to_mel(fmax), n_mels + 2))
|
||||
fft_f = np.linspace(0, sr / 2, n_fft // 2 + 1)
|
||||
filt = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
|
||||
for i in range(n_mels):
|
||||
lo, ce, hi = pts[i], pts[i + 1], pts[i + 2]
|
||||
left = (fft_f - lo) / max(ce - lo, 1e-10)
|
||||
right = (hi - fft_f) / max(hi - ce, 1e-10)
|
||||
filt[i] = np.maximum(0, np.minimum(left, right))
|
||||
enorm = (2.0 / (pts[2:n_mels + 2] - pts[0:n_mels])).astype(np.float32)
|
||||
filt *= enorm[:, None]
|
||||
return filt
|
||||
|
||||
|
||||
def _stft_power(y, n_fft, hop):
|
||||
"""|STFT|^2 with periodic Hann + center=True reflect pad (torch.stft-compatible).
|
||||
Returns (n_fft//2+1, n_frames) float32."""
|
||||
y = np.asarray(y, dtype=np.float32)
|
||||
pad = n_fft // 2
|
||||
y = np.pad(y, pad, mode="reflect")
|
||||
n_frames = 1 + (len(y) - n_fft) // hop
|
||||
idx = np.arange(n_fft)[None, :] + hop * np.arange(n_frames)[:, None]
|
||||
# periodic Hann (torch.hann_window default), not numpy's symmetric hanning
|
||||
win = (0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(n_fft) / n_fft)).astype(np.float32)
|
||||
frames = y[idx] * win[None, :]
|
||||
spec = np.fft.rfft(frames, n=n_fft, axis=1)
|
||||
return (spec.real ** 2 + spec.imag ** 2).T.astype(np.float32)
|
||||
|
||||
|
||||
def _melspectrogram(y_ch, sr, n_mels=30, fmax=16000, hop_length=2048, n_fft=2048):
|
||||
power = _stft_power(y_ch, n_fft, hop_length)
|
||||
return _mel_filterbank(n_mels, n_fft, sr, fmax) @ power
|
||||
|
||||
|
||||
def _power_to_db(S, top_db=80.0):
|
||||
log_spec = 10.0 * np.log10(np.maximum(S, 1e-10))
|
||||
return np.maximum(log_spec - log_spec.max(), -top_db)
|
||||
|
||||
|
||||
def _mel_channel(y_ch, sr, n_mels=30):
|
||||
"""dB-scaled mel + band-tinted RGB for one mono channel."""
|
||||
S = _melspectrogram(y_ch, sr, n_mels=n_mels, fmax=16000, hop_length=2048)
|
||||
|
||||
S_db = _power_to_db(S)
|
||||
np.clip(S_db, -60, 0, out=S_db)
|
||||
S_db += 60.0
|
||||
S_db /= 60.0
|
||||
np.power(S_db, 0.6, out=S_db)
|
||||
|
||||
mel_f = _mel_frequencies(n_mels, fmax=16000)
|
||||
n_frames = S.shape[1]
|
||||
band_norms = np.empty((3, n_frames), dtype=np.float32)
|
||||
for i, (flo, fhi, _) in enumerate(SPEC_BANDS):
|
||||
mask = (mel_f >= flo) & (mel_f < fhi)
|
||||
if mask.any():
|
||||
power = np.sum(S[mask], axis=0)
|
||||
db = 10.0 * np.log10(power + 1e-10)
|
||||
np.clip(db, -20, None, out=db)
|
||||
db -= -20
|
||||
mx = db.max()
|
||||
if mx > 0:
|
||||
db /= mx
|
||||
band_norms[i] = db
|
||||
else:
|
||||
band_norms[i] = 0.0
|
||||
|
||||
rgb = band_norms.T @ _BAND_COLORS
|
||||
for c in range(3):
|
||||
mx = rgb[:, c].max()
|
||||
if mx > 0:
|
||||
rgb[:, c] /= mx
|
||||
|
||||
return S_db, rgb
|
||||
|
||||
|
||||
def _resample_to_32k(pcm_int16, sr_in):
|
||||
"""Resample (T, 2) int16 PCM to (2, T') float32 at 32 kHz.
|
||||
Linear interpolation at align_corners=False sample positions (matches the
|
||||
torch.nn.functional.interpolate call in the TRT original)."""
|
||||
target_sr = 32000
|
||||
y = pcm_int16.astype(np.float32).T / 32768.0 # (2, T)
|
||||
if sr_in == target_sr:
|
||||
return y, target_sr
|
||||
in_len = y.shape[-1]
|
||||
new_len = int(round(in_len * target_sr / sr_in))
|
||||
scale = in_len / new_len
|
||||
pos = np.clip((np.arange(new_len, dtype=np.float64) + 0.5) * scale - 0.5, 0, in_len - 1)
|
||||
out = np.stack([np.interp(pos, np.arange(in_len), ch) for ch in y]).astype(np.float32)
|
||||
return out, target_sr
|
||||
|
||||
|
||||
def render_spectrogram_png(pcm_int16, sample_rate=44100, width=1200, height=240,
|
||||
n_mels=30) -> bytes:
|
||||
"""Render the 3-band tinted stereo mel spectrogram to PNG bytes.
|
||||
|
||||
Args:
|
||||
pcm_int16: (T, 2) int16 numpy array — stereo PCM.
|
||||
sample_rate: input PCM sample rate (default 44100).
|
||||
width, height: final image size.
|
||||
n_mels: mel filterbank size (30 matches the dashboards).
|
||||
|
||||
Returns: PNG bytes ready to base64 / save / etc.
|
||||
"""
|
||||
if pcm_int16.ndim != 2 or pcm_int16.shape[1] != 2:
|
||||
raise ValueError(f"expected (T, 2) PCM, got {pcm_int16.shape}")
|
||||
y, sr = _resample_to_32k(pcm_int16, sample_rate)
|
||||
S_L, rgb_L = _mel_channel(y[0], sr, n_mels=n_mels)
|
||||
S_R, rgb_R = _mel_channel(y[1], sr, n_mels=n_mels)
|
||||
|
||||
nf = min(S_L.shape[1], S_R.shape[1])
|
||||
S_L, S_R = S_L[:, :nf], S_R[:, :nf]
|
||||
rgb_L, rgb_R = rgb_L[:nf], rgb_R[:nf]
|
||||
nm = S_L.shape[0]
|
||||
|
||||
# L channel: flip vertically so bass-band is at bottom.
|
||||
S_L = S_L[::-1]
|
||||
|
||||
img = np.empty((nm * 2, nf, 3), dtype=np.float32)
|
||||
img[:nm] = S_L[:, :, np.newaxis] * rgb_L[np.newaxis, :, :]
|
||||
img[nm:] = S_R[:, :, np.newaxis] * rgb_R[np.newaxis, :, :]
|
||||
|
||||
np.clip(img, 0, 1, out=img)
|
||||
img *= 255
|
||||
pil = Image.fromarray(img.astype(np.uint8)).resize((width, height), Image.LANCZOS)
|
||||
buf = io.BytesIO()
|
||||
pil.save(buf, format="PNG", optimize=True)
|
||||
return buf.getvalue()
|
||||
@@ -0,0 +1,221 @@
|
||||
"""Sanity-check every shipped npz and every CLI config of sa3_mlx.py.
|
||||
|
||||
Run from inside the sa3_mlx/ folder:
|
||||
python test_all_configs.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import os, subprocess, sys, time, wave, tempfile
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
|
||||
|
||||
REPO = Path(__file__).resolve().parent.parent # project root (scripts/ is one level down)
|
||||
PY = sys.executable
|
||||
|
||||
|
||||
# ─── expected weight files ────────────────────────────────────────────────
|
||||
|
||||
EXPECTED_NPZ = {
|
||||
"models/mlx/dit_sm-music_f16.npz": {"min_keys": 400, "must_have_prefix": "cond."},
|
||||
"models/mlx/dit_sm-sfx_f16.npz": {"min_keys": 400, "must_have_prefix": "cond."},
|
||||
"models/mlx/dit_medium_f16.npz": {"min_keys": 500, "must_have_prefix": "cond."},
|
||||
"models/mlx/same_s_encoder_f32.npz": {"min_keys": 100},
|
||||
"models/mlx/same_s_decoder_f32.npz": {"min_keys": 100},
|
||||
"models/mlx/same_l_encoder_f32.npz": {"min_keys": 200},
|
||||
"models/mlx/same_l_decoder_f32.npz": {"min_keys": 200},
|
||||
"models/mlx/t5gemma_f16.npz": {"min_keys": 130, "must_have_key": "TOKENIZER_MODEL"},
|
||||
}
|
||||
|
||||
|
||||
# ─── helpers ──────────────────────────────────────────────────────────────
|
||||
|
||||
def check_npz(path: Path, spec: dict) -> tuple[bool, str]:
|
||||
"""Open the npz and verify it has the expected shape/keys."""
|
||||
try:
|
||||
z = np.load(path, allow_pickle=True)
|
||||
except Exception as e:
|
||||
return False, f"load failed: {e}"
|
||||
keys = list(z.files)
|
||||
if len(keys) < spec["min_keys"]:
|
||||
return False, f"only {len(keys)} keys (expected ≥ {spec['min_keys']})"
|
||||
if "must_have_prefix" in spec:
|
||||
prefix = spec["must_have_prefix"]
|
||||
if not any(k.startswith(prefix) for k in keys):
|
||||
return False, f"no keys with prefix {prefix!r}"
|
||||
if "must_have_key" in spec:
|
||||
if spec["must_have_key"] not in keys:
|
||||
return False, f"missing key {spec['must_have_key']!r}"
|
||||
# Sanity: actually read one tensor without error
|
||||
try:
|
||||
for k in keys[:3]:
|
||||
_ = np.asarray(z[k])
|
||||
except Exception as e:
|
||||
return False, f"tensor read failed: {e}"
|
||||
return True, f"{len(keys)} keys"
|
||||
|
||||
|
||||
def read_wav(path: str) -> tuple[np.ndarray, int]:
|
||||
with wave.open(path, "rb") as w:
|
||||
nch = w.getnchannels()
|
||||
sr = w.getframerate()
|
||||
n = w.getnframes()
|
||||
pcm = np.frombuffer(w.readframes(n), dtype=np.int16).reshape(-1, nch).T.astype(np.float32) / 32767.0
|
||||
return pcm, sr
|
||||
|
||||
|
||||
def run_cli(name: str, extra_args: list[str], expected_seconds: float) -> tuple[bool, str]:
|
||||
"""Run sa3_mlx.py with the given args; sanity-check the output WAV."""
|
||||
out = Path(tempfile.gettempdir()) / f"sa3_test_{name}.wav"
|
||||
cmd = [PY, str(REPO / "scripts" / "sa3_mlx.py"),
|
||||
"--seed", "42",
|
||||
"--out", str(out)] + extra_args
|
||||
if "--steps" not in extra_args: # default to 4 steps for speed
|
||||
cmd += ["--steps", "4"]
|
||||
t0 = time.time()
|
||||
try:
|
||||
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=180,
|
||||
cwd=str(REPO))
|
||||
except subprocess.TimeoutExpired:
|
||||
return False, f"TIMEOUT after 180s"
|
||||
dt = time.time() - t0
|
||||
if proc.returncode != 0:
|
||||
last_err = proc.stderr.strip().split("\n")[-1] if proc.stderr.strip() else "(no stderr)"
|
||||
return False, f"exit {proc.returncode}: {last_err[:120]}"
|
||||
if not out.exists():
|
||||
return False, "no output WAV produced"
|
||||
|
||||
pcm, sr = read_wav(str(out))
|
||||
# Quality gates: correct duration, non-NaN, not all-silent.
|
||||
actual_secs = pcm.shape[-1] / sr
|
||||
if abs(actual_secs - expected_seconds) > 0.1:
|
||||
return False, f"duration {actual_secs:.2f}s ≠ requested {expected_seconds:.2f}s"
|
||||
if not np.isfinite(pcm).all():
|
||||
return False, "audio contains NaN/Inf"
|
||||
peak = float(np.abs(pcm).max())
|
||||
rms = float(np.sqrt((pcm ** 2).mean()))
|
||||
if peak < 0.005 or rms < 0.0005:
|
||||
return False, f"effectively silent (peak={peak:.4f} rms={rms:.4f})"
|
||||
|
||||
out.unlink(missing_ok=True)
|
||||
return True, f"{dt:.1f}s peak={peak:.3f} rms={rms:.3f}"
|
||||
|
||||
|
||||
# ─── tests ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_weights():
|
||||
print("\n[ phase 1 ] verify every shipped npz loads cleanly\n")
|
||||
fails = []
|
||||
for relpath, spec in EXPECTED_NPZ.items():
|
||||
full = REPO / relpath
|
||||
if not full.exists():
|
||||
print(f" ✗ {relpath:42s} FILE MISSING")
|
||||
fails.append(relpath)
|
||||
continue
|
||||
ok, info = check_npz(full, spec)
|
||||
mark = "✓" if ok else "✗"
|
||||
print(f" {mark} {relpath:42s} {info}")
|
||||
if not ok:
|
||||
fails.append(relpath)
|
||||
return fails
|
||||
|
||||
|
||||
def test_configs():
|
||||
print("\n[ phase 2 ] end-to-end CLI configurations (steps=4, 3s clips)\n")
|
||||
SECONDS = "3"
|
||||
EXPECTED_SEC = 3.0
|
||||
|
||||
# Will use this as init-audio for audio-to-audio + inpaint tests.
|
||||
# Generate one short clip with sm-music first.
|
||||
init_wav = Path(tempfile.gettempdir()) / "sa3_test_init.wav"
|
||||
print(" [setup] generating a small init audio for a2a/inpaint tests …")
|
||||
subprocess.run([PY, str(REPO / "scripts" / "sa3_mlx.py"),
|
||||
"--prompt", "drums",
|
||||
"--dit", "sm-music", "--decoder", "same-s",
|
||||
"--seconds", SECONDS, "--seed", "1", "--steps", "4",
|
||||
"--out", str(init_wav)],
|
||||
capture_output=True, check=True, cwd=str(REPO))
|
||||
print(f" [setup] init audio at {init_wav}\n")
|
||||
|
||||
matrix = [
|
||||
# text-to-audio: each DiT × natural decoder
|
||||
("t2a-music", ["--prompt", "lofi house", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS]),
|
||||
("t2a-sfx", ["--prompt", "rain on roof", "--dit", "sm-sfx", "--decoder", "same-s", "--seconds", SECONDS]),
|
||||
("t2a-medium", ["--prompt", "piano solo", "--dit", "medium", "--decoder", "same-l", "--seconds", SECONDS]),
|
||||
|
||||
# cross-pairings (DiT/decoder mismatch — should still produce audio, may sound odd)
|
||||
("cross-music-L", ["--prompt", "lofi", "--dit", "sm-music", "--decoder", "same-l", "--seconds", SECONDS]),
|
||||
("cross-medium-S", ["--prompt", "piano", "--dit", "medium", "--decoder", "same-s", "--seconds", SECONDS]),
|
||||
|
||||
# mode: audio-to-audio
|
||||
("a2a-music", ["--prompt", "jazz fusion", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--init-audio", str(init_wav), "--init-noise-level", "0.6"]),
|
||||
("a2a-medium", ["--prompt", "ambient", "--dit", "medium", "--decoder", "same-l", "--seconds", SECONDS,
|
||||
"--init-audio", str(init_wav), "--init-noise-level", "0.6"]),
|
||||
|
||||
# mode: inpainting
|
||||
("inpaint-music", ["--prompt", "guitar solo", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--init-audio", str(init_wav), "--inpaint-range", "1,2"]),
|
||||
|
||||
# CFG variants
|
||||
("cfg3-music", ["--prompt", "techno beat", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--cfg", "3.0"]),
|
||||
("cfg3-neg", ["--prompt", "techno beat", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--cfg", "3.0", "--negative-prompt", "vocals, drums"]),
|
||||
("cfg0.5", ["--prompt", "techno", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--cfg", "0.5"]),
|
||||
("cfg-apg0-vanilla",["--prompt", "techno", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--cfg", "3.0", "--apg", "0.0"]),
|
||||
|
||||
# dtype: fp32 DiT
|
||||
("fp32-music", ["--prompt", "lofi", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--dit-dtype", "fp32"]),
|
||||
|
||||
# different step counts
|
||||
("steps1", ["--prompt", "lofi", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--steps", "1"]),
|
||||
("steps16", ["--prompt", "lofi", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--steps", "16"]),
|
||||
|
||||
# --no-free-models
|
||||
("no-free", ["--prompt", "lofi", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS,
|
||||
"--no-free-models"]),
|
||||
|
||||
# empty prompt (unconditional)
|
||||
("empty-prompt", ["--prompt", "", "--dit", "sm-music", "--decoder", "same-s", "--seconds", SECONDS]),
|
||||
]
|
||||
|
||||
fails = []
|
||||
for name, args in matrix:
|
||||
ok, info = run_cli(name, args, EXPECTED_SEC)
|
||||
mark = "✓" if ok else "✗"
|
||||
print(f" {mark} {name:24s} {info}")
|
||||
if not ok:
|
||||
fails.append((name, info))
|
||||
|
||||
init_wav.unlink(missing_ok=True)
|
||||
return fails
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 68)
|
||||
print("sa3_mlx full-stack self-test")
|
||||
print("=" * 68)
|
||||
npz_fails = test_weights()
|
||||
config_fails = test_configs()
|
||||
|
||||
print("\n" + "=" * 68)
|
||||
print("SUMMARY")
|
||||
print("=" * 68)
|
||||
if not npz_fails and not config_fails:
|
||||
print("✓ ALL PASS")
|
||||
return 0
|
||||
print(f"npz failures : {len(npz_fails)}")
|
||||
for f in npz_fails: print(f" - {f}")
|
||||
print(f"config failures: {len(config_fails)}")
|
||||
for name, info in config_fails: print(f" - {name}: {info}")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,446 @@
|
||||
"""Unit tests for models/defs/lora_merge.py — the LoRA merge-at-load math.
|
||||
|
||||
Pure-numpy/MLX, no model weights or torch needed. Run either way:
|
||||
|
||||
python scripts/test_lora_merge.py # standalone (no pytest needed)
|
||||
pytest scripts/test_lora_merge.py
|
||||
|
||||
Correctness is established by:
|
||||
* the zero-init -> identity invariant for every adapter type (a strong guard on
|
||||
axis / reshape / transpose bugs: with B/M_xs zeroed and magnitudes at their
|
||||
init values, every forward must return W0);
|
||||
* an exact independent reconstruction for standard LoRA and PEFT;
|
||||
* the Conv1d (out,k,in) <-> PyTorch (out,in,k) layout round-trip;
|
||||
* the to_local_embed name remap;
|
||||
* --lora-strength scaling and the bit-exact strength-0 bypass;
|
||||
* the trust boundary (pickle refusal) and base-mismatch / 0-merge handling.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
|
||||
# Import lora_merge the same way the CLI does (REPO = optimized/mlx on sys.path).
|
||||
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, REPO)
|
||||
from models.defs import lora_merge as lm # noqa: E402
|
||||
|
||||
rng = np.random.default_rng(0)
|
||||
|
||||
|
||||
def _save_native(path, layer, params, cfg):
|
||||
d = {f"{layer}.parametrizations.weight.0.{k}": mx.array(v.astype(np.float32))
|
||||
for k, v in params.items()}
|
||||
mx.save_safetensors(path, d, metadata={"lora_config": json.dumps(cfg)})
|
||||
|
||||
|
||||
def _init_params(adapter_type, W0, rank, nonzero):
|
||||
"""Init so a zeroed core -> identity; if nonzero, give a real core."""
|
||||
fo, fi = W0.shape
|
||||
p = {}
|
||||
if not adapter_type.endswith("-xs"):
|
||||
p["lora_A"] = rng.standard_normal((rank, fi)).astype(np.float32)
|
||||
p["lora_B"] = (rng.standard_normal((fo, rank)) if nonzero
|
||||
else np.zeros((fo, rank))).astype(np.float32)
|
||||
else:
|
||||
p["M_xs"] = (rng.standard_normal((rank, rank)) if nonzero
|
||||
else np.zeros((rank, rank))).astype(np.float32)
|
||||
if "magnitude" in lm._PARAMS_FOR[adapter_type]:
|
||||
nd = 1 if "rows" in adapter_type else 0
|
||||
p["magnitude"] = np.linalg.norm(W0, axis=nd).astype(np.float32)
|
||||
if "magnitude_r" in lm._PARAMS_FOR[adapter_type]:
|
||||
p["magnitude_r"] = np.linalg.norm(W0, axis=1).astype(np.float32)
|
||||
p["magnitude_c"] = np.linalg.norm(W0, axis=0).astype(np.float32)
|
||||
return p
|
||||
|
||||
|
||||
def _np(arr):
|
||||
return np.array(arr.astype(mx.float32))
|
||||
|
||||
|
||||
def test_trust_boundary_refuses_pickle():
|
||||
for bad in ("evil.ckpt", "evil.pt", "evil.bin", "weights.npz"):
|
||||
try:
|
||||
lm._load_safetensors(bad)
|
||||
assert False, f"should have refused {bad}"
|
||||
except lm.LoraError:
|
||||
pass
|
||||
|
||||
|
||||
def test_all_types_zero_init_identity_and_nonzero_delta():
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
for atype in lm._PARAMS_FOR:
|
||||
path = os.path.join(tmp, f"{atype}.safetensors")
|
||||
# zero-init -> identity
|
||||
_save_native(path, "foo", _init_params(atype, W0, 4, nonzero=False),
|
||||
{"adapter_type": atype, "rank": 4, "alpha": 4})
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
lm.merge_loras_into_weights(w, [path])
|
||||
assert np.allclose(_np(w["foo.weight"]), W0, atol=1e-4), f"{atype} identity"
|
||||
# nonzero core -> finite, changed
|
||||
_save_native(path, "foo", _init_params(atype, W0, 4, nonzero=True),
|
||||
{"adapter_type": atype, "rank": 4, "alpha": 4})
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
lm.merge_loras_into_weights(w, [path])
|
||||
got = _np(w["foo.weight"])
|
||||
assert np.isfinite(got).all() and not np.allclose(got, W0), f"{atype} delta"
|
||||
|
||||
|
||||
def test_standard_lora_exact_and_strength():
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
A = rng.standard_normal((4, 8)).astype(np.float32)
|
||||
B = rng.standard_normal((6, 4)).astype(np.float32)
|
||||
expect = W0 + (8.0 / 4.0) * (B @ A)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "std.safetensors")
|
||||
_save_native(path, "foo", {"lora_A": A, "lora_B": B},
|
||||
{"adapter_type": "lora", "rank": 4, "alpha": 8})
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
lm.merge_loras_into_weights(w, [path])
|
||||
assert np.allclose(_np(w["foo.weight"]), expect, atol=1e-4)
|
||||
# strength 0.5 halves the delta; strength 0 is the original
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
lm.merge_loras_into_weights(w, [path], strength=0.5)
|
||||
assert np.allclose(_np(w["foo.weight"]), W0 + 0.5 * (8.0 / 4.0) * (B @ A), atol=1e-4)
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
lm.merge_loras_into_weights(w, [path], strength=0.0)
|
||||
assert np.allclose(_np(w["foo.weight"]), W0, atol=1e-6)
|
||||
|
||||
|
||||
def test_conv1d_layout_round_trip():
|
||||
out_c, k_c, in_c = 4, 1, 3
|
||||
Wc = rng.standard_normal((out_c, k_c, in_c)).astype(np.float32) # MLX layout
|
||||
Ac = rng.standard_normal((2, in_c * k_c)).astype(np.float32)
|
||||
Bc = rng.standard_normal((out_c, 2)).astype(np.float32)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "conv.safetensors")
|
||||
_save_native(path, "conv", {"lora_A": Ac, "lora_B": Bc},
|
||||
{"adapter_type": "lora", "rank": 2, "alpha": 2})
|
||||
w = {"conv.weight": mx.array(Wc)}
|
||||
lm.merge_loras_into_weights(w, [path])
|
||||
got = _np(w["conv.weight"])
|
||||
expect = Wc + (Bc @ Ac).reshape(out_c, in_c, k_c).transpose(0, 2, 1)
|
||||
assert got.shape == Wc.shape and np.allclose(got, expect, atol=1e-4)
|
||||
|
||||
|
||||
def test_to_local_embed_remap():
|
||||
Wt = rng.standard_normal((5, 5)).astype(np.float32)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "tle.safetensors")
|
||||
_save_native(path, "transformer.layers.0.to_local_embed.0",
|
||||
{"lora_A": rng.standard_normal((2, 5)).astype(np.float32),
|
||||
"lora_B": rng.standard_normal((5, 2)).astype(np.float32)},
|
||||
{"adapter_type": "lora", "rank": 2, "alpha": 2})
|
||||
w = {"transformer.layers.0.to_local_embed.seq.0.weight": mx.array(Wt)}
|
||||
stats = lm.merge_loras_into_weights(w, [path])
|
||||
assert stats["merged"] == 1 and not stats["skipped"]
|
||||
|
||||
|
||||
def test_peft_format_exact():
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
A = rng.standard_normal((4, 8)).astype(np.float32)
|
||||
B = rng.standard_normal((6, 4)).astype(np.float32)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
with open(os.path.join(tmp, "adapter_config.json"), "w") as fh:
|
||||
json.dump({"peft_type": "LORA", "r": 4, "lora_alpha": 8, "use_dora": False}, fh)
|
||||
mx.save_safetensors(
|
||||
os.path.join(tmp, "adapter_model.safetensors"),
|
||||
{"base_model.model.foo.lora_A.weight": mx.array(A),
|
||||
"base_model.model.foo.lora_B.weight": mx.array(B)},
|
||||
metadata={"format": "pt"})
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
stats = lm.merge_loras_into_weights(w, [tmp]) # directory input
|
||||
assert stats["merged"] == 1
|
||||
assert np.allclose(_np(w["foo.weight"]), W0 + (8.0 / 4.0) * (B @ A), atol=1e-4)
|
||||
|
||||
|
||||
def test_unknown_layers_skipped_base_untouched():
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "std.safetensors")
|
||||
_save_native(path, "foo",
|
||||
{"lora_A": rng.standard_normal((4, 8)).astype(np.float32),
|
||||
"lora_B": rng.standard_normal((6, 4)).astype(np.float32)},
|
||||
{"adapter_type": "lora", "rank": 4, "alpha": 4})
|
||||
w = {"other.weight": mx.array(W0)}
|
||||
msgs = []
|
||||
stats = lm.merge_loras_into_weights(w, [path], log=msgs.append)
|
||||
assert stats["merged"] == 0 and stats["skipped"] == ["foo"]
|
||||
assert np.allclose(_np(w["other.weight"]), W0)
|
||||
# a 0-merge run must warn, not look like a successful no-op
|
||||
assert any("WARNING" in m and "0 layers" in m for m in msgs)
|
||||
|
||||
|
||||
def test_base_mismatch_raises_clear_error():
|
||||
# base is (6, 8) but the adapter was trained for an (10, 8) layer
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "mismatch.safetensors")
|
||||
_save_native(path, "foo",
|
||||
{"lora_A": rng.standard_normal((4, 8)).astype(np.float32),
|
||||
"lora_B": rng.standard_normal((10, 4)).astype(np.float32)},
|
||||
{"adapter_type": "lora", "rank": 4, "alpha": 4})
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
try:
|
||||
lm.merge_loras_into_weights(w, [path])
|
||||
assert False, "should have raised on base mismatch"
|
||||
except lm.LoraError as e:
|
||||
assert "foo" in str(e) and "base" in str(e).lower()
|
||||
|
||||
|
||||
def test_svd_bases_orthonormal():
|
||||
Wsvd = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
U, V = lm._svd_bases(Wsvd, rank=6)
|
||||
assert np.allclose(U.T @ U, np.eye(6), atol=1e-4)
|
||||
assert np.allclose(V.T @ V, np.eye(6), atol=1e-4)
|
||||
|
||||
|
||||
def test_underfit_wrapper_prefixes():
|
||||
"""Underfit (github.com/dada-bots/underfit) saves full-wrapper layer names:
|
||||
DiT layers as ``model.transformer...`` and the seconds conditioner as
|
||||
``conditioners.seconds_total.embedder.embedding.1``. Both must land on the
|
||||
bare npz keys; bare names keep working."""
|
||||
W_dit = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
W_cond = rng.standard_normal((5, 7)).astype(np.float32)
|
||||
A_d = rng.standard_normal((4, 8)).astype(np.float32)
|
||||
B_d = rng.standard_normal((6, 4)).astype(np.float32)
|
||||
A_c = rng.standard_normal((4, 7)).astype(np.float32)
|
||||
B_c = rng.standard_normal((5, 4)).astype(np.float32)
|
||||
cfg = {"adapter_type": "lora", "rank": 4, "alpha": 4}
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "underfit.safetensors")
|
||||
d = {}
|
||||
for layer, (A, B) in {
|
||||
"model.transformer.layers.0.self_attn.to_qkv": (A_d, B_d),
|
||||
"conditioners.seconds_total.embedder.embedding.1": (A_c, B_c),
|
||||
}.items():
|
||||
d[f"{layer}.parametrizations.weight.0.lora_A"] = mx.array(A)
|
||||
d[f"{layer}.parametrizations.weight.0.lora_B"] = mx.array(B)
|
||||
mx.save_safetensors(path, d, metadata={"lora_config": json.dumps(cfg)})
|
||||
w = {"transformer.layers.0.self_attn.to_qkv.weight": mx.array(W_dit),
|
||||
"cond.seconds_total_weight": mx.array(W_cond)}
|
||||
stats = lm.merge_loras_into_weights(w, [path])
|
||||
assert stats["merged"] == 2 and not stats["skipped"], stats
|
||||
assert np.allclose(_np(w["transformer.layers.0.self_attn.to_qkv.weight"]),
|
||||
W_dit + B_d @ A_d, atol=1e-4)
|
||||
assert np.allclose(_np(w["cond.seconds_total_weight"]),
|
||||
W_cond + B_c @ A_c, atol=1e-4)
|
||||
# bare names still pass through, and model.-prefixed to_local_embed remaps
|
||||
assert lm._layer_to_npz_key("transformer.layers.3.ff.ff.2") == \
|
||||
"transformer.layers.3.ff.ff.2.weight"
|
||||
assert lm._layer_to_npz_key("model.transformer.layers.3.to_local_embed.0") == \
|
||||
"transformer.layers.3.to_local_embed.seq.0.weight"
|
||||
|
||||
|
||||
def test_lora_spec_parser():
|
||||
s = lm.parse_lora_spec(["a.safetensors"], default_strength=0.5)
|
||||
assert s == {"path": "a.safetensors", "strength": 0.5, "steps": None}
|
||||
s = lm.parse_lora_spec(["a.safetensors", "strength=0.8", "steps=2-8"])
|
||||
assert s["strength"] == 0.8 and s["steps"] == (2, 8)
|
||||
assert lm.parse_lora_spec(["a", "steps=2-"])["steps"] == (2, None)
|
||||
assert lm.parse_lora_spec(["a", "steps=-4"])["steps"] == (None, 4)
|
||||
assert lm.parse_lora_spec(["a", "steps=3"])["steps"] == (3, 3)
|
||||
for bad in (["a.st", "b.st"], # two paths in one flag
|
||||
["a", "steps=8-2"], # min > max
|
||||
["a", "steps=0-3"], # 1-based
|
||||
["a", "steps=-"], # empty range
|
||||
["a", "strength=loud"], # not a float
|
||||
[]):
|
||||
try:
|
||||
lm.parse_lora_spec(bad)
|
||||
assert False, f"should have rejected {bad}"
|
||||
except lm.LoraError:
|
||||
pass
|
||||
# resolve against an 8-step schedule (0-based inclusive out)
|
||||
assert lm.resolve_steps(None, 8) == (0, 7)
|
||||
assert lm.resolve_steps((2, 8), 8) == (1, 7)
|
||||
assert lm.resolve_steps((2, None), 8) == (1, 7)
|
||||
assert lm.resolve_steps((None, 4), 8) == (0, 3)
|
||||
assert lm.resolve_steps((2, 99), 8) == (1, 7) # clamp
|
||||
assert lm.resolve_steps((9, None), 8) is None # misses the schedule
|
||||
|
||||
|
||||
class _StubModule:
|
||||
"""Bare object with .weight — stands in for nn.Linear in plan tests."""
|
||||
def __init__(self, w):
|
||||
self.weight = w
|
||||
|
||||
|
||||
def _stub_tree(weights):
|
||||
"""Build a walkable module tree for keys like 'foo.weight' and
|
||||
'layers.0.lin.weight' (digits = list containers, as in the DiT)."""
|
||||
class _NS: # namespace node
|
||||
pass
|
||||
root = _NS()
|
||||
for key, w in weights.items():
|
||||
parts = key.split(".")[:-1] # drop trailing 'weight'
|
||||
obj = root
|
||||
for j, p in enumerate(parts[:-1]):
|
||||
nxt = parts[j + 1]
|
||||
if p.isdigit():
|
||||
obj = obj[int(p)]
|
||||
else:
|
||||
if not hasattr(obj, p):
|
||||
setattr(obj, p, [] if nxt.isdigit() else _NS())
|
||||
obj = getattr(obj, p)
|
||||
if isinstance(obj, list) and nxt.isdigit():
|
||||
while len(obj) <= int(nxt):
|
||||
obj.append(_NS())
|
||||
last = parts[-1]
|
||||
mod = _StubModule(w)
|
||||
if last.isdigit():
|
||||
obj[int(last)] = mod
|
||||
else:
|
||||
setattr(obj, last, mod)
|
||||
return root
|
||||
|
||||
|
||||
def test_step_plan_matches_reference():
|
||||
"""For every adapter type: gate two overlapping adapters across an 8-step
|
||||
schedule and check the in-place weights at EVERY step against the from-
|
||||
scratch reference W0 + Σ_active strength·(merged − W0)."""
|
||||
num_steps = 8
|
||||
K1, K2 = "foo.weight", "layers.0.lin.weight"
|
||||
W = {K1: rng.standard_normal((6, 8)).astype(np.float32),
|
||||
K2: rng.standard_normal((5, 7)).astype(np.float32)}
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
for atype in lm._PARAMS_FOR:
|
||||
p1 = os.path.join(tmp, f"{atype}-1.safetensors")
|
||||
p2 = os.path.join(tmp, f"{atype}-2.safetensors")
|
||||
prm = {k: _init_params(atype, w, 3, nonzero=True) for k, w in W.items()}
|
||||
d = {}
|
||||
for k, w in W.items():
|
||||
layer = k[: -len(".weight")]
|
||||
for pk, pv in prm[k].items():
|
||||
d[f"{layer}.parametrizations.weight.0.{pk}"] = mx.array(pv)
|
||||
mx.save_safetensors(p1, d, metadata={"lora_config": json.dumps(
|
||||
{"adapter_type": atype, "rank": 3, "alpha": 3})})
|
||||
prm2 = {K1: _init_params("lora", W[K1], 3, nonzero=True)}
|
||||
d2 = {f"foo.parametrizations.weight.0.{pk}": mx.array(pv)
|
||||
for pk, pv in prm2[K1].items()}
|
||||
mx.save_safetensors(p2, d2, metadata={"lora_config": json.dumps(
|
||||
{"adapter_type": "lora", "rank": 3, "alpha": 3})})
|
||||
|
||||
specs = [{"path": p1, "strength": 0.7, "steps": (2, 5)},
|
||||
{"path": p2, "strength": 1.3, "steps": (4, None)}]
|
||||
weights = {k: mx.array(w) for k, w in W.items()}
|
||||
plan = lm.prepare_loras(weights, specs, num_steps=num_steps)
|
||||
assert plan is not None and len(plan.layers) == 2
|
||||
|
||||
model = _stub_tree(weights)
|
||||
plan.attach(model)
|
||||
mods = {K1: model.foo, K2: model.layers[0].lin}
|
||||
|
||||
def ref(i):
|
||||
out = {k: W[k].copy() for k in W}
|
||||
if 1 <= i <= 4: # adapter 1 (steps 2-5, 0-based 1..4)
|
||||
for k in W:
|
||||
m = lm._merged_weight(W[k], prm[k], atype, 1.0)
|
||||
out[k] += 0.7 * (m - W[k])
|
||||
if i >= 3: # adapter 2 (steps 4-, 0-based 3..7)
|
||||
m = lm._merged_weight(W[K1], prm2[K1], "lora", 1.0)
|
||||
out[K1] += 1.3 * (m - W[K1])
|
||||
return out
|
||||
|
||||
for i in range(num_steps):
|
||||
plan.sync(i)
|
||||
expect = ref(i)
|
||||
for k in W:
|
||||
got = _np(mods[k].weight)
|
||||
err = np.abs(got - expect[k]).max()
|
||||
assert err < 1e-3, f"{atype} step {i} {k}: max|d|={err:.2e}"
|
||||
|
||||
|
||||
def test_step_plan_round_trip_drift():
|
||||
"""50 on/off gating cycles on fp16 weights must not accumulate drift
|
||||
(exact-inverse pairing — measured frozen-after-cycle-1 in the design
|
||||
benchmarks; this guards the implementation)."""
|
||||
K = "foo.weight"
|
||||
W0 = rng.standard_normal((32, 48)).astype(np.float16)
|
||||
prm = _init_params("dora-rows", W0.astype(np.float32), 4, nonzero=True)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "d.safetensors")
|
||||
_save_native(path, "foo", prm, {"adapter_type": "dora-rows", "rank": 4, "alpha": 4})
|
||||
weights = {K: mx.array(W0)}
|
||||
plan = lm.prepare_loras(weights, [{"path": path, "strength": 1.0,
|
||||
"steps": (2, None)}], num_steps=2)
|
||||
model = _stub_tree(weights)
|
||||
plan.attach(model)
|
||||
drift1 = None
|
||||
for cyc in range(50):
|
||||
plan.sync(1) # on
|
||||
plan.sync(0) # off
|
||||
d = np.abs(_np(model.foo.weight) - W0.astype(np.float32)).max()
|
||||
if cyc == 0:
|
||||
drift1 = d
|
||||
assert d <= drift1 + 1e-6, f"drift grew: cycle1 {drift1:.2e} → cycle50 {d:.2e}"
|
||||
rel = (np.linalg.norm(_np(model.foo.weight) - W0.astype(np.float32))
|
||||
/ np.linalg.norm(W0.astype(np.float32)))
|
||||
assert rel < 5e-3, f"round-trip rel drift {rel:.2e}"
|
||||
|
||||
|
||||
def test_full_interval_uses_merge_path():
|
||||
"""Specs without steps= (or covering every step) → no plan; weights match
|
||||
the plain merge_loras_into_weights result exactly."""
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
prm = _init_params("lora", W0, 4, nonzero=True)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "l.safetensors")
|
||||
_save_native(path, "foo", prm, {"adapter_type": "lora", "rank": 4, "alpha": 4})
|
||||
wa = {"foo.weight": mx.array(W0)}
|
||||
wb = {"foo.weight": mx.array(W0)}
|
||||
plan = lm.prepare_loras(wa, [{"path": path, "strength": 0.6,
|
||||
"steps": (1, 8)}], num_steps=8)
|
||||
assert plan is None
|
||||
lm.merge_loras_into_weights(wb, [path], strength=0.6)
|
||||
assert np.array_equal(_np(wa["foo.weight"]), _np(wb["foo.weight"]))
|
||||
|
||||
|
||||
def test_gate_all_and_clear_restore_base():
|
||||
"""The gradio path: gate_all=True routes even full-range adapters through
|
||||
the plan so clear() can restore a cached model to base in place; an adapter
|
||||
matching zero layers raises a clear wrong-base error."""
|
||||
W0 = rng.standard_normal((6, 8)).astype(np.float32)
|
||||
prm = _init_params("dora-rows", W0, 4, nonzero=True)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "g.safetensors")
|
||||
_save_native(path, "foo", prm, {"adapter_type": "dora-rows", "rank": 4, "alpha": 4})
|
||||
w = {"foo.weight": mx.array(W0)}
|
||||
plan = lm.prepare_loras(w, [{"path": path, "strength": 1.0, "steps": None}],
|
||||
num_steps=8, gate_all=True)
|
||||
assert plan is not None, "gate_all must plan-manage full-range adapters"
|
||||
model = _stub_tree(w)
|
||||
plan.attach(model)
|
||||
assert not np.allclose(_np(model.foo.weight), W0), "step-1 state not applied"
|
||||
plan.clear()
|
||||
assert np.abs(_np(model.foo.weight) - W0).max() < 1e-4, "clear() must restore base"
|
||||
# an adapter whose layers don't exist in the target → clear error
|
||||
w2 = {"bar.weight": mx.array(W0)}
|
||||
try:
|
||||
lm.prepare_loras(w2, [{"path": path, "strength": 1.0, "steps": (2, None)}],
|
||||
num_steps=8)
|
||||
assert False, "zero-match adapter should raise"
|
||||
except lm.LoraError as e:
|
||||
assert "different base model" in str(e)
|
||||
|
||||
|
||||
def _main():
|
||||
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
||||
failed = []
|
||||
for t in tests:
|
||||
try:
|
||||
t()
|
||||
print(f" ok {t.__name__}")
|
||||
except Exception as e: # noqa: BLE001
|
||||
print(f" FAIL {t.__name__}: {e}")
|
||||
failed.append(t.__name__)
|
||||
print("\n" + ("ALL PASS" if not failed else f"FAILURES: {failed}"))
|
||||
return 1 if failed else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(_main())
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Shared weights manifest + downloader.
|
||||
|
||||
Maps every weight file the runtime needs to its position in the
|
||||
`stabilityai/stable-audio-3-optimized` HuggingFace repo.
|
||||
|
||||
`install.py` calls `ensure_local` upfront for the bundles the user picks.
|
||||
`sa3_mlx.py` calls `ensure_local` lazily, just before each model loads — so
|
||||
a fresh checkout with no weights still works if the user is willing to
|
||||
wait for the first run to download them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ID = "stabilityai/stable-audio-3-optimized"
|
||||
# weights.py lives in <project>/scripts/; SCRIPT_DIR points at the project
|
||||
# root so the local rel paths in the manifest ("models/mlx/foo.npz") resolve
|
||||
# against the actual project layout.
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
# Bundles the install script offers to the user. Each maps to a list of
|
||||
# weight files (local relative path on the left, HF repo path on the right).
|
||||
# T5Gemma is in SHARED because every bundle needs it.
|
||||
|
||||
DIT_BUNDLES: dict[str, list[tuple[str, str]]] = {
|
||||
"medium": [
|
||||
("models/mlx/dit_medium_f16.npz", "MLX/dit_medium_f16.npz"),
|
||||
("models/mlx/same_l_decoder_f32.npz", "MLX/same_l_decoder_f32.npz"),
|
||||
("models/mlx/same_l_encoder_f32.npz", "MLX/same_l_encoder_f32.npz"),
|
||||
],
|
||||
"sm-music": [
|
||||
("models/mlx/dit_sm-music_f16.npz", "MLX/dit_sm-music_f16.npz"),
|
||||
("models/mlx/same_s_decoder_f32.npz", "MLX/same_s_decoder_f32.npz"),
|
||||
("models/mlx/same_s_encoder_f32.npz", "MLX/same_s_encoder_f32.npz"),
|
||||
],
|
||||
"sm-sfx": [
|
||||
("models/mlx/dit_sm-sfx_f16.npz", "MLX/dit_sm-sfx_f16.npz"),
|
||||
("models/mlx/same_s_decoder_f32.npz", "MLX/same_s_decoder_f32.npz"),
|
||||
("models/mlx/same_s_encoder_f32.npz", "MLX/same_s_encoder_f32.npz"),
|
||||
],
|
||||
}
|
||||
|
||||
SHARED: list[tuple[str, str]] = [
|
||||
("models/mlx/t5gemma_f16.npz", "MLX/t5gemma_f16.npz"),
|
||||
]
|
||||
|
||||
# Base-model DiT weights — needed only for LoRA *training* (lora_train_mlx.py
|
||||
# trains on the BASE checkpoint, not the ARC weights inference uses). Not part
|
||||
# of any install bundle; auto-downloaded lazily by the trainer via ensure_local.
|
||||
TRAINING_BASE: list[tuple[str, str]] = [
|
||||
("models/mlx/dit_sm-music-base_f16.npz", "MLX/dit_sm-music-base_f16.npz"),
|
||||
("models/mlx/dit_sm-sfx-base_f16.npz", "MLX/dit_sm-sfx-base_f16.npz"),
|
||||
("models/mlx/dit_medium-base_f16.npz", "MLX/dit_medium-base_f16.npz"),
|
||||
]
|
||||
|
||||
# Human-friendly bundle sizes (for the install prompt).
|
||||
BUNDLE_SIZES = {
|
||||
"medium": "5.9 GB (medium DiT + SAME-L codec)",
|
||||
"sm-music": "1.3 GB (small music DiT + SAME-S codec)",
|
||||
"sm-sfx": "1.3 GB (small sfx DiT + SAME-S codec)",
|
||||
}
|
||||
|
||||
# Flat (local_rel_path → hf_path) lookup — used by sa3_mlx.py for lazy
|
||||
# auto-download at load time.
|
||||
FLAT_MANIFEST: dict[str, str] = {}
|
||||
for _items in DIT_BUNDLES.values():
|
||||
for _rel, _hf in _items:
|
||||
FLAT_MANIFEST[_rel] = _hf
|
||||
for _rel, _hf in SHARED:
|
||||
FLAT_MANIFEST[_rel] = _hf
|
||||
for _rel, _hf in TRAINING_BASE:
|
||||
FLAT_MANIFEST[_rel] = _hf
|
||||
|
||||
|
||||
def _hf_token_configured() -> bool:
|
||||
"""True if any HF token is set — env var or cached login on disk."""
|
||||
import os
|
||||
if os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN"):
|
||||
return True
|
||||
try:
|
||||
from huggingface_hub import get_token # huggingface_hub ≥ 0.19
|
||||
return bool(get_token())
|
||||
except ImportError:
|
||||
try:
|
||||
from huggingface_hub import HfFolder
|
||||
return bool(HfFolder.get_token())
|
||||
except Exception:
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
_LOGIN_TIP_SHOWN = False
|
||||
|
||||
def _show_hf_login_tip_once() -> None:
|
||||
"""Print a one-time login suggestion if no HF token is configured.
|
||||
|
||||
Anonymous downloads work but have a ~50 GB/day soft cap on HF's LFS CDN
|
||||
and lower aggregate bandwidth — a free token effectively removes both.
|
||||
Stays silent if a token is already in place.
|
||||
"""
|
||||
global _LOGIN_TIP_SHOWN
|
||||
if _LOGIN_TIP_SHOWN:
|
||||
return
|
||||
_LOGIN_TIP_SHOWN = True
|
||||
if _hf_token_configured():
|
||||
return
|
||||
import sys
|
||||
YEL = "\033[1;33m" if sys.stdout.isatty() else ""
|
||||
BOLD = "\033[1m" if sys.stdout.isatty() else ""
|
||||
DIM = "\033[2m" if sys.stdout.isatty() else ""
|
||||
RST = "\033[0m" if sys.stdout.isatty() else ""
|
||||
print()
|
||||
print(f" {YEL}⚠ not logged in to HuggingFace{RST} — anonymous downloads work but are")
|
||||
print(f" rate-limited (~50 GB/day cap on the LFS CDN). For faster, higher-limit")
|
||||
print(f" downloads, log in once with a free read-only token:")
|
||||
print()
|
||||
print(f" 1. create an account at {BOLD}https://huggingface.co/join{RST}")
|
||||
print(f" 2. generate a token at {BOLD}https://huggingface.co/settings/tokens{RST}")
|
||||
print(f" {DIM}('Read' scope is enough){RST}")
|
||||
print(f" 3. save it on this machine — pick one:")
|
||||
print(f" {BOLD}hf auth login{RST} {DIM}# modern (huggingface_hub ≥ 1.0){RST}")
|
||||
print(f" {BOLD}huggingface-cli login{RST} {DIM}# classic; still works{RST}")
|
||||
print(f" {BOLD}export HF_TOKEN=hf_xxx{RST} {DIM}# one-off / scripts{RST}")
|
||||
print()
|
||||
|
||||
|
||||
def ensure_local(local_rel_path: str, verbose: bool = True) -> Path:
|
||||
"""Resolve a weight file to an absolute local path, downloading if missing.
|
||||
|
||||
Files are streamed into the HuggingFace cache (~/.cache/huggingface/hub/)
|
||||
and symlinked into the project at `local_rel_path` so the on-disk layout
|
||||
looks the same whether the file was bundled or downloaded.
|
||||
"""
|
||||
target = SCRIPT_DIR / local_rel_path
|
||||
if target.exists() or target.is_symlink():
|
||||
return target
|
||||
|
||||
if local_rel_path not in FLAT_MANIFEST:
|
||||
raise FileNotFoundError(
|
||||
f"{local_rel_path} is not in the weights manifest — can't auto-download."
|
||||
)
|
||||
|
||||
# First-download tip: nudge users toward logging in to HF for better limits.
|
||||
# No-op if a token is already configured.
|
||||
_show_hf_login_tip_once()
|
||||
|
||||
hf_filename = FLAT_MANIFEST[local_rel_path]
|
||||
if verbose:
|
||||
print(f" ↓ downloading {hf_filename} (from {REPO_ID})")
|
||||
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
"huggingface_hub is required to auto-download weights.\n"
|
||||
"Run: pip install huggingface_hub\n"
|
||||
"Or run the install.py script in this directory."
|
||||
) from e
|
||||
|
||||
cached = hf_hub_download(repo_id=REPO_ID, filename=hf_filename)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Symlink keeps the HF cache canonical (one copy on disk) while exposing
|
||||
# the file at the project-relative path the runtime expects.
|
||||
if target.is_symlink():
|
||||
target.unlink()
|
||||
target.symlink_to(cached)
|
||||
return target
|
||||
|
||||
|
||||
def is_present(local_rel_path: str) -> bool:
|
||||
"""True if the file exists locally (does not trigger a download)."""
|
||||
p = SCRIPT_DIR / local_rel_path
|
||||
return p.exists() or p.is_symlink()
|
||||
|
||||
|
||||
def bundle_status(bundle: str) -> tuple[int, int]:
|
||||
"""Returns (present_count, total_count) for the bundle (including SHARED)."""
|
||||
items = DIT_BUNDLES[bundle] + SHARED
|
||||
present = sum(1 for rel, _ in items if is_present(rel))
|
||||
return present, len(items)
|
||||
Reference in New Issue
Block a user