初始提交：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,158 @@
|
||||
"""Shared, colored "Try these commands" block.
|
||||
|
||||
Used by:
|
||||
- scripts/install.py — printed at the end of install.sh
|
||||
- scripts/sa3_tflite.py — appended to `./sa3 --help`
|
||||
|
||||
Examples render with the user's `./sa3` wrapper when present (falling back to
|
||||
`.venv/bin/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
|
||||
|
||||
# Windows consoles default to a legacy code page (cp1252/cp437) that can't encode
|
||||
# the emoji/box-drawing characters below. Force UTF-8 with replacement; no-op on
|
||||
# macOS/Linux. (sa3_tflite.py/install.py do this too — reconfigure is idempotent.)
|
||||
for _stream in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
_stream.reconfigure(encoding="utf-8", errors="replace")
|
||||
except AttributeError:
|
||||
pass # exotic stream without reconfigure() — leave as-is
|
||||
|
||||
# Enable ANSI/VT escape processing in legacy Windows consoles (cmd.exe);
|
||||
# harmless elsewhere and on modern Windows Terminal.
|
||||
if os.name == "nt":
|
||||
os.system("")
|
||||
|
||||
# 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 (./sa3 on POSIX, sa3.bat on Windows), else
|
||||
the .venv python + scripts/sa3_tflite.py, else a bare `python scripts/sa3_tflite.py`.
|
||||
"""
|
||||
if os.name == "nt":
|
||||
if (SCRIPT_DIR / "sa3.bat").exists():
|
||||
return "sa3.bat", "the sa3.bat wrapper uses .venv automatically", True
|
||||
venv_py = SCRIPT_DIR / ".venv" / "Scripts" / "python.exe"
|
||||
if venv_py.exists():
|
||||
return str(venv_py.relative_to(SCRIPT_DIR)), "", False
|
||||
return "python", "", False
|
||||
wrapper = SCRIPT_DIR / "sa3"
|
||||
if wrapper.exists() and os.access(wrapper, os.X_OK):
|
||||
return "./sa3", "the ./sa3 wrapper uses .venv automatically", True
|
||||
venv_py = SCRIPT_DIR / ".venv" / "bin" / "python"
|
||||
if venv_py.exists():
|
||||
return str(venv_py.relative_to(SCRIPT_DIR)), "", False
|
||||
if Path(sys.prefix).resolve() == (SCRIPT_DIR / ".venv").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_tflite.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 ("sm-music", "sm-sfx", "medium")}
|
||||
|
||||
# ── Basic generation ─────────────────────────────────────────────
|
||||
hdr("🎵 Generate audio from a prompt")
|
||||
if have["sm-music"]:
|
||||
cmd('--prompt "lofi house loop, mellow piano" \\\n'
|
||||
f' --dit sm-music --decoder same-s --seconds 30 --out lofi.wav',
|
||||
"fast music generation")
|
||||
if have["sm-sfx"]:
|
||||
cmd('--prompt "footsteps on gravel, then a door slamming" \\\n'
|
||||
f' --dit sm-sfx --decoder same-s --seconds 20 --out sfx.wav',
|
||||
"sound-effect generation")
|
||||
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 (chunked SAME-L decode)")
|
||||
|
||||
# ── Playback ─────────────────────────────────────────────────────
|
||||
hdr("▶ Play immediately after generation")
|
||||
one_dit = "sm-music" if have["sm-music"] else ("sm-sfx" if have["sm-sfx"] else "medium")
|
||||
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 20 --out drone.wav --play',
|
||||
"writes WAV + plays it (afplay/winsound/aplay; 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 (sequential dual-pass)")
|
||||
|
||||
# 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,166 @@
|
||||
#!/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.
|
||||
|
||||
Unlike the MLX release, this TFLite/LiteRT stack is portable CPU — it runs on
|
||||
macOS/Linux, x86/ARM. No Apple-Silicon requirement.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Windows consoles default to a legacy code page (cp1252/cp437) that can't encode
|
||||
# the ✓/→ status glyphs. Force UTF-8 with replacement; no-op on macOS/Linux.
|
||||
for _stream in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
_stream.reconfigure(encoding="utf-8", errors="replace")
|
||||
except AttributeError:
|
||||
pass # exotic stream without reconfigure() — leave as-is
|
||||
|
||||
# Enable ANSI/VT escape processing in legacy Windows consoles (cmd.exe);
|
||||
# harmless elsewhere and on modern Windows Terminal.
|
||||
if os.name == "nt":
|
||||
os.system("")
|
||||
|
||||
# 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, 9)
|
||||
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 too old."""
|
||||
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 this stack 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: brew install python@3.11\n"
|
||||
f"On Debian/Ubuntu: apt install python3.11\n"
|
||||
f"On Windows: install from https://www.python.org/downloads/ "
|
||||
f"(check 'Add python.exe to PATH')\n",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
# TFLite / LiteRT (ai_edge_litert) is portable across macOS/Linux, x86/ARM —
|
||||
# no platform gate here (unlike the Metal-backed MLX release).
|
||||
|
||||
|
||||
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 TFLite 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 "
|
||||
"(sm-music, sm-sfx, medium). Without this flag, "
|
||||
"nothing is downloaded — sa3_tflite.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: {Path(sys.executable).name} install.py --download sm-music")
|
||||
print(f" or: ./install.sh --download sm-music,medium")
|
||||
|
||||
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,321 @@
|
||||
"""LoRA adapter parsing + per-layer merge math — torch/mlx-free (numpy only).
|
||||
|
||||
The TFLite route ships standalone (its own bootstrap, its own venv, runs on
|
||||
Windows), so it can't import the MLX runtime. This module is a faithful vendored
|
||||
copy of the *backend-agnostic* half of
|
||||
``optimized/mlx/models/defs/lora_merge.py`` — the safetensors parsing and the
|
||||
per-adapter-type ``W' = f(W0, adapter)`` math — with the two MLX touch-points
|
||||
replaced by pure numpy:
|
||||
|
||||
* safetensors is read with a ~25-line numpy reader (JSON header + a raw blob),
|
||||
instead of ``mx.load``;
|
||||
* everything else (``_parse_adapter``, ``_group_native``/``_group_peft``,
|
||||
``_merged_weight`` for all nine adapter types, ``_check_shapes``,
|
||||
``_svd_bases`` for the ``-xs`` variants, ``parse_lora_spec``) is copied
|
||||
verbatim, since it was already numpy-fp32.
|
||||
|
||||
KEEP IN SYNC with optimized/mlx/models/defs/lora_merge.py — the merge math must
|
||||
match the MLX path bit-for-bit (both mirror ``LoRAParametrization.*_forward`` in
|
||||
``stable_audio_3/models/lora/model.py``). What is intentionally NOT here: the
|
||||
MLX weight-dict merge (``merge_loras_into_weights``) and the per-step
|
||||
``LoraStepPlan`` — TFLite patches frozen FlatBuffer weight buffers instead (see
|
||||
``lora_patch.py``), and per-step gating is not feasible on a frozen graph.
|
||||
|
||||
Trust boundary: only ``.safetensors`` adapters are accepted; a pickle
|
||||
``.ckpt``/``.pt``/``.bin`` (which ``torch.load`` would execute code from) is
|
||||
refused — this module never imports torch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Pickle-backed extensions we refuse to load (the trust boundary).
|
||||
_PICKLE_EXTS = (".ckpt", ".pt", ".pth", ".bin")
|
||||
|
||||
# Adapter param names per type (mirrors utils._get_adapter_param_names).
|
||||
_PARAMS_FOR = {
|
||||
"lora": ("lora_A", "lora_B"),
|
||||
"dora-rows": ("lora_A", "lora_B", "magnitude"),
|
||||
"dora-cols": ("lora_A", "lora_B", "magnitude"),
|
||||
"bora": ("lora_A", "lora_B", "magnitude_r", "magnitude_c"),
|
||||
"lora-xs": ("M_xs",),
|
||||
"dora-rows-xs": ("M_xs", "magnitude"),
|
||||
"dora-cols-xs": ("M_xs", "magnitude"),
|
||||
"bora-xs": ("M_xs", "magnitude_r", "magnitude_c"),
|
||||
}
|
||||
|
||||
# Underfit saves the seconds conditioner Linear under this checkpoint name; the
|
||||
# baked TFLite DiT carries it in-graph as the unique (768, 256) FULLY_CONNECTED
|
||||
# (lora_patch maps it there).
|
||||
COND_SECONDS_LAYER = "conditioners.seconds_total.embedder.embedding.1"
|
||||
|
||||
|
||||
class LoraError(Exception):
|
||||
"""An adapter could not be loaded or applied."""
|
||||
|
||||
|
||||
# ── numpy safetensors reader (no torch, no safetensors pkg) ────────────────────
|
||||
|
||||
_ST_DTYPES = {
|
||||
"F64": np.float64, "F32": np.float32, "F16": np.float16, "BF16": None,
|
||||
"I64": np.int64, "I32": np.int32, "I16": np.int16, "I8": np.int8,
|
||||
"U8": np.uint8, "BOOL": np.bool_,
|
||||
}
|
||||
|
||||
|
||||
def _load_safetensors(path: str):
|
||||
"""Return ``(tensors: dict[str, np.ndarray fp32], metadata: dict)``.
|
||||
|
||||
Refuses pickle. bf16 is widened to fp32 via a uint16<<16 bit-cast (numpy has
|
||||
no native bf16). Every returned tensor is fp32 so the merge math is uniform.
|
||||
"""
|
||||
lower = path.lower()
|
||||
if lower.endswith(_PICKLE_EXTS):
|
||||
raise LoraError(
|
||||
f"refusing to load pickle-format adapter {os.path.basename(path)!r} — "
|
||||
f"only .safetensors adapters are accepted (a .ckpt/.pt is unpickled by "
|
||||
f"torch.load and can execute arbitrary code)"
|
||||
)
|
||||
if not lower.endswith(".safetensors"):
|
||||
raise LoraError(f"not a .safetensors adapter: {path!r}")
|
||||
with open(path, "rb") as f:
|
||||
(hlen,) = struct.unpack("<Q", f.read(8))
|
||||
hdr = json.loads(f.read(hlen))
|
||||
blob = f.read()
|
||||
meta = hdr.pop("__metadata__", {}) or {}
|
||||
out: dict[str, np.ndarray] = {}
|
||||
for k, v in hdr.items():
|
||||
if v["dtype"] not in _ST_DTYPES:
|
||||
raise LoraError(f"{os.path.basename(path)}: unsupported tensor dtype "
|
||||
f"{v['dtype']!r} for {k}")
|
||||
b0, b1 = v["data_offsets"]
|
||||
n = 1
|
||||
for d in v["shape"]:
|
||||
n *= d
|
||||
if v["dtype"] == "BF16":
|
||||
u = np.frombuffer(blob, np.uint16, count=n or 1, offset=b0)
|
||||
arr = (u.astype(np.uint32) << 16).view(np.float32)
|
||||
else:
|
||||
arr = np.frombuffer(blob[b0:b1], _ST_DTYPES[v["dtype"]])
|
||||
out[k] = arr.reshape(v["shape"]).astype(np.float32, copy=False)
|
||||
return out, meta
|
||||
|
||||
|
||||
# ── SVD bases for -xs adapters (recomputed; mirrors model.py) ──────────────────
|
||||
|
||||
def _canonicalize_svd_signs(U: np.ndarray, Vh: np.ndarray):
|
||||
"""Deterministic sign convention: largest-magnitude element of each U column
|
||||
is positive (mirrors model._canonicalize_svd_signs)."""
|
||||
max_abs_idx = np.argmax(np.abs(U), axis=0)
|
||||
signs = np.sign(U[max_abs_idx, np.arange(U.shape[1])])
|
||||
signs[signs == 0] = 1.0
|
||||
return U * signs[None, :], Vh * signs[:, None]
|
||||
|
||||
|
||||
def _svd_bases(W0: np.ndarray, rank: int):
|
||||
"""``(U[:, :rank], V[:, :rank])`` from the SVD of W0 (fan_out, fan_in), with
|
||||
V such that ``U @ diag(S) @ V.T`` reconstructs W0 (mirrors model.py)."""
|
||||
U_full, _S, Vh_full = np.linalg.svd(W0, full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
return U_full[:, :rank], Vh_full[:rank, :].T
|
||||
|
||||
|
||||
# ── per-type merge math (numpy, float32) — mirrors lora_merge._merged_weight ───
|
||||
|
||||
def _mag_2d(mag: np.ndarray, norm_dim: int) -> np.ndarray:
|
||||
mag = np.atleast_1d(np.squeeze(mag))
|
||||
return mag.reshape(-1, 1) if norm_dim == 1 else mag.reshape(1, -1)
|
||||
|
||||
|
||||
def merged_weight(W0: np.ndarray, p: dict, adapter_type: str, scaling: float) -> np.ndarray:
|
||||
"""LoRA-merged weight for one layer at full strength (lora_strength=1).
|
||||
``W0`` is (fan_out, fan_in) float32; ``p`` holds the adapter tensors."""
|
||||
if adapter_type == "lora":
|
||||
return W0 + scaling * (p["lora_B"] @ p["lora_A"])
|
||||
|
||||
if adapter_type in ("dora-rows", "dora-cols"):
|
||||
norm_dim = 1 if adapter_type == "dora-rows" else 0
|
||||
V = W0 + scaling * (p["lora_B"] @ p["lora_A"])
|
||||
V_hat = V / (np.linalg.norm(V, axis=norm_dim, keepdims=True) + 1e-12)
|
||||
return V_hat * _mag_2d(p["magnitude"], norm_dim)
|
||||
|
||||
if adapter_type == "bora":
|
||||
V = W0 + scaling * (p["lora_B"] @ p["lora_A"])
|
||||
V_r = V / (np.linalg.norm(V, axis=1, keepdims=True) + 1e-12)
|
||||
inter = p["magnitude_r"].reshape(-1, 1) * V_r
|
||||
H_c = inter / (np.linalg.norm(inter, axis=0, keepdims=True) + 1e-12)
|
||||
return H_c * p["magnitude_c"].reshape(1, -1)
|
||||
|
||||
if adapter_type.endswith("-xs"):
|
||||
rank = p["M_xs"].shape[0]
|
||||
U, V = _svd_bases(W0, rank)
|
||||
Vfull = W0 + scaling * (U @ p["M_xs"] @ V.T)
|
||||
if adapter_type == "lora-xs":
|
||||
return Vfull
|
||||
if adapter_type in ("dora-rows-xs", "dora-cols-xs"):
|
||||
norm_dim = 1 if adapter_type == "dora-rows-xs" else 0
|
||||
V_hat = Vfull / (np.linalg.norm(Vfull, axis=norm_dim, keepdims=True) + 1e-12)
|
||||
return V_hat * _mag_2d(p["magnitude"], norm_dim)
|
||||
if adapter_type == "bora-xs":
|
||||
V_r = Vfull / (np.linalg.norm(Vfull, axis=1, keepdims=True) + 1e-12)
|
||||
inter = p["magnitude_r"].reshape(-1, 1) * V_r
|
||||
H_c = inter / (np.linalg.norm(inter, axis=0, keepdims=True) + 1e-12)
|
||||
return H_c * p["magnitude_c"].reshape(1, -1)
|
||||
|
||||
raise LoraError(f"unknown adapter_type {adapter_type!r}")
|
||||
|
||||
|
||||
def check_shapes(layer: str, W0: np.ndarray, p: dict, adapter_type: str) -> None:
|
||||
"""Clear error when the adapter doesn't fit the base weight — almost always a
|
||||
wrong-base-model mismatch. Without it the failure is a raw numpy broadcast."""
|
||||
fan_out, fan_in = W0.shape
|
||||
if adapter_type.endswith("-xs"):
|
||||
rank = p["M_xs"].shape[0]
|
||||
if rank > min(fan_out, fan_in):
|
||||
raise LoraError(
|
||||
f"{layer}: LoRA-XS rank {rank} exceeds base min-dim "
|
||||
f"{min(fan_out, fan_in)} for weight {W0.shape} — wrong base model?"
|
||||
)
|
||||
return
|
||||
b_out, b_rank = p["lora_B"].shape
|
||||
a_rank, a_in = p["lora_A"].shape
|
||||
if b_out != fan_out or a_in != fan_in or a_rank != b_rank:
|
||||
raise LoraError(
|
||||
f"{layer}: adapter lora_B{p['lora_B'].shape}·lora_A{p['lora_A'].shape} "
|
||||
f"does not fit base weight {W0.shape} — wrong base model?"
|
||||
)
|
||||
|
||||
|
||||
# ── checkpoint parsing → normalized per-layer adapter ──────────────────────────
|
||||
|
||||
def _resolve_path(path: str) -> str:
|
||||
"""Accept a .safetensors file or a PEFT adapter directory (resolve to the
|
||||
adapter_model.safetensors inside it)."""
|
||||
if os.path.isdir(path):
|
||||
cand = os.path.join(path, "adapter_model.safetensors")
|
||||
if os.path.isfile(cand):
|
||||
return cand
|
||||
hits = [f for f in os.listdir(path) if f.lower().endswith(".safetensors")]
|
||||
if len(hits) == 1:
|
||||
return os.path.join(path, hits[0])
|
||||
raise LoraError(
|
||||
f"{path!r}: expected one .safetensors adapter in the directory, found {hits}"
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def parse_adapter(path: str):
|
||||
"""Load one adapter → ``(adapter_type, scaling, layers)`` where ``layers``
|
||||
maps a checkpoint layer name → its param dict (numpy float32)."""
|
||||
tensors, meta = _load_safetensors(path)
|
||||
|
||||
native_marker = ".parametrizations.weight.0."
|
||||
if any(native_marker in k for k in tensors):
|
||||
cfg = json.loads(meta.get("lora_config", "{}")) if meta else {}
|
||||
layers = _group_native(tensors)
|
||||
rank = int(cfg.get("rank") or _infer_rank(layers))
|
||||
alpha = float(cfg.get("alpha", rank))
|
||||
adapter_type = _resolve_native_type(cfg.get("adapter_type", "lora"))
|
||||
return adapter_type, alpha / rank, layers
|
||||
|
||||
peft_marker = ".lora_A.weight"
|
||||
if any(k.endswith(peft_marker) for k in tensors):
|
||||
cfg = _read_peft_config(path)
|
||||
rank = int(cfg["r"])
|
||||
alpha = float(cfg.get("lora_alpha", rank))
|
||||
use_dora = bool(cfg.get("use_dora", False))
|
||||
use_rslora = bool(cfg.get("use_rslora", False))
|
||||
adapter_type = "dora-rows" if use_dora else "lora"
|
||||
scaling = alpha / (np.sqrt(rank) if use_rslora else rank)
|
||||
return adapter_type, scaling, _group_peft(tensors)
|
||||
|
||||
raise LoraError(
|
||||
f"{os.path.basename(path)!r}: not a recognised LoRA (no SA3-native "
|
||||
f"parametrization keys and no PEFT lora_A/lora_B keys)"
|
||||
)
|
||||
|
||||
|
||||
def _group_native(tensors: dict) -> dict:
|
||||
marker = ".parametrizations.weight.0."
|
||||
layers: dict[str, dict] = {}
|
||||
for k, v in tensors.items():
|
||||
if marker not in k:
|
||||
continue
|
||||
layer, _, param = k.partition(marker)
|
||||
layers.setdefault(layer, {})[param] = v
|
||||
return layers
|
||||
|
||||
|
||||
def _group_peft(tensors: dict) -> dict:
|
||||
prefix = "base_model.model."
|
||||
layers: dict[str, dict] = {}
|
||||
for k, v in tensors.items():
|
||||
name = k[len(prefix):] if k.startswith(prefix) else k
|
||||
for suffix, param in ((".lora_A.weight", "lora_A"),
|
||||
(".lora_B.weight", "lora_B"),
|
||||
(".lora_magnitude_vector.weight", "magnitude")):
|
||||
if name.endswith(suffix):
|
||||
layers.setdefault(name[: -len(suffix)], {})[param] = v
|
||||
break
|
||||
return layers
|
||||
|
||||
|
||||
def _read_peft_config(path: str) -> dict:
|
||||
cfg_path = os.path.join(os.path.dirname(path), "adapter_config.json")
|
||||
if not os.path.isfile(cfg_path):
|
||||
raise LoraError(f"PEFT adapter at {path!r} is missing its adapter_config.json sibling")
|
||||
with open(cfg_path) as fh:
|
||||
return json.load(fh)
|
||||
|
||||
|
||||
def _infer_rank(layers: dict) -> int:
|
||||
for params in layers.values():
|
||||
if "lora_A" in params:
|
||||
return params["lora_A"].shape[0]
|
||||
if "M_xs" in params:
|
||||
return params["M_xs"].shape[0]
|
||||
raise LoraError("cannot infer LoRA rank (no lora_A / M_xs tensors)")
|
||||
|
||||
|
||||
def _resolve_native_type(adapter_type: str) -> str:
|
||||
"""Legacy 'dora' → 'dora-rows' (mirrors utils.resolve_adapter_type)."""
|
||||
return "dora-rows" if adapter_type == "dora" else adapter_type
|
||||
|
||||
|
||||
# ── CLI spec parsing ────────────────────────────────────────────────────────
|
||||
|
||||
def parse_lora_spec(tokens, default_strength: float = 1.0) -> dict:
|
||||
"""Parse one ``--lora`` group: ``PATH [strength=S]``.
|
||||
|
||||
Mirrors the MLX CLI's flag, minus ``steps=`` — per-step LoRA gating is not
|
||||
feasible on a frozen TFLite graph (weights are packed into XNNPACK's private
|
||||
buffers at allocate_tensors; there is no per-step re-merge), so ``steps=`` is
|
||||
rejected with a pointer to the MLX backend.
|
||||
"""
|
||||
if not tokens:
|
||||
raise LoraError("--lora needs an adapter path")
|
||||
spec = {"path": tokens[0], "strength": float(default_strength)}
|
||||
for tok in tokens[1:]:
|
||||
key, eq, val = tok.partition("=")
|
||||
if key == "steps":
|
||||
raise LoraError(
|
||||
"--lora steps=… (per-step gating) is not supported on the TFLite "
|
||||
"backend — the frozen graph's weights are merged once at load. "
|
||||
"Use the MLX backend (optimized/mlx) for step-gated LoRA."
|
||||
)
|
||||
if not eq or key != "strength":
|
||||
raise LoraError(
|
||||
f"--lora: unexpected token {tok!r} — one adapter per --lora flag; "
|
||||
f"the only option after the path is strength=S "
|
||||
f"(e.g. --lora a.safetensors strength=0.8)"
|
||||
)
|
||||
try:
|
||||
spec["strength"] = float(val)
|
||||
except ValueError:
|
||||
raise LoraError(f"--lora: bad strength {val!r}") from None
|
||||
return spec
|
||||
@@ -0,0 +1,392 @@
|
||||
"""Merge LoRA adapters into a baked SA3 DiT ``.tflite`` by patching its weight
|
||||
buffers — the TFLite counterpart of the MLX weight-dict merge.
|
||||
|
||||
A ``.tflite`` is a frozen FlatBuffer: there is no weight dict to edit at load, so
|
||||
LoRA inference means rewriting the FULLY_CONNECTED weight buffers in place. The
|
||||
mechanism (validated on the 5.8 GB medium DiT):
|
||||
|
||||
1. Parse the FlatBuffer header (mmap, no full read) and map every
|
||||
checkpoint layer name → the file offset/size/dtype of its FC weight buffer.
|
||||
2. Clone the base file (APFS ``cp -c`` = a few ms and no extra disk until
|
||||
written; plain copy elsewhere).
|
||||
3. For each adapter layer: read W0 from the clone's mmap, compute W' with the
|
||||
shared numpy merge math (``lora_core.merged_weight``), and overwrite the
|
||||
buffer's bytes at its exact offset. Deltas from multiple adapters accumulate
|
||||
against the ORIGINAL W0 (order-independent), matching the MLX path.
|
||||
|
||||
Patching happens BEFORE ``Interpreter(...)`` so XNNPACK packs the merged weights
|
||||
at ``allocate_tensors`` exactly as for an unpatched model — no runtime cost, no
|
||||
cache interaction. Patched files are cached by a content hash of (base, adapters,
|
||||
strengths), so repeat runs reuse them instantly.
|
||||
|
||||
Precision: fp32 (weights are direct fp32 consts) and w16a32 (weights are fp16
|
||||
consts behind a DEQUANTIZE) are supported. Quantized-int8 variants (w8a32,
|
||||
w8a8-dyn, w4a32) are refused — merging requires a dequant→merge→requant that
|
||||
sacrifices the GPTQ error-feedback grid (a phase-2 item).
|
||||
|
||||
Name mapping tiers (all hard-fail on ambiguity — never silently mis-patch):
|
||||
1. Block linears (self_attn/cross_attn/ff): name carries ``TransformerBlock_N``
|
||||
→ (block, module, leaf) + shape → unique.
|
||||
2. ``to_local_embed.seq.{0,2}``: per-block but the tflite names carry NO block
|
||||
index and their ``;NN`` counter runs BACKWARDS — block identity comes from
|
||||
FC topological order within the group, never the counter.
|
||||
3. Prefix embedders / project_in / project_out: singleton names → unique.
|
||||
4. Seconds conditioner: the unique ``(768, 256)`` FC (in-graph, patchable).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import mmap
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import lora_core as lc # noqa: E402
|
||||
|
||||
# Bump when the merge math or mapping changes so stale cache entries are missed.
|
||||
_MERGE_VERSION = "1"
|
||||
|
||||
_QUANTIZED_PRECISIONS = ("w8a32", "w8a8-dyn", "w4a32")
|
||||
|
||||
|
||||
# ── FlatBuffer FC-weight discovery ─────────────────────────────────────────────
|
||||
|
||||
def _fc_weight_buffers(buf):
|
||||
"""Yield one record per FULLY_CONNECTED weight, in graph topological order:
|
||||
``dict(name, shape, off, size, np_dtype, op_index)``.
|
||||
|
||||
Resolves the weight buffer through a DEQUANTIZE (w16a32: fp16 const) or takes
|
||||
the FC's weight input directly (fp32: fp32 const). Refuses int8 weights.
|
||||
"""
|
||||
from ai_edge_litert import schema_py_generated as schema
|
||||
BO = schema.BuiltinOperator
|
||||
TT = schema.TensorType
|
||||
_NP = {TT.FLOAT32: np.float32, TT.FLOAT16: np.float16}
|
||||
|
||||
m = schema.Model.GetRootAs(buf, 0)
|
||||
sg = m.Subgraphs(0)
|
||||
producer = {}
|
||||
for oi in range(sg.OperatorsLength()):
|
||||
op = sg.Operators(oi)
|
||||
bc = m.OperatorCodes(op.OpcodeIndex()).BuiltinCode()
|
||||
for j in range(op.OutputsLength()):
|
||||
producer[op.Outputs(j)] = (bc, op)
|
||||
|
||||
for oi in range(sg.OperatorsLength()):
|
||||
op = sg.Operators(oi)
|
||||
if m.OperatorCodes(op.OpcodeIndex()).BuiltinCode() != BO.FULLY_CONNECTED:
|
||||
continue
|
||||
w_idx = op.Inputs(1)
|
||||
wt = sg.Tensors(w_idx)
|
||||
shape = tuple(wt.Shape(i) for i in range(wt.ShapeLength()))
|
||||
name = (wt.Name() or b"?").decode()
|
||||
prod = producer.get(w_idx)
|
||||
if prod is not None and prod[0] == BO.DEQUANTIZE:
|
||||
# w16a32: FC weight ← DEQUANTIZE ← fp16 const (the patch target).
|
||||
src = sg.Tensors(prod[1].Inputs(0))
|
||||
ttype = src.Type()
|
||||
b = m.Buffers(src.Buffer())
|
||||
src_name = (src.Name() or wt.Name() or b"?").decode()
|
||||
else:
|
||||
ttype, b, src_name = wt.Type(), m.Buffers(wt.Buffer()), name
|
||||
if ttype not in _NP:
|
||||
raise lc.LoraError(
|
||||
f"FC weight {name!r} is {TT.__dict__ and _tt_name(ttype)} — LoRA merge "
|
||||
f"on the TFLite backend supports fp32 and w16a32 (fp16) weights only. "
|
||||
f"Quantized-int8 models (w8a32 / w8a8-dyn / w4a32) are not supported "
|
||||
f"(merging would destroy the GPTQ grid). Use --precision fp32 or w16a32."
|
||||
)
|
||||
off, size = b.Offset(), b.Size()
|
||||
if off == 0 or size == 0:
|
||||
raise lc.LoraError(
|
||||
f"FC weight {name!r} is not stored in an out-of-band buffer "
|
||||
f"(off={off}, size={size}) — cannot patch in place."
|
||||
)
|
||||
yield dict(name=name, dbg=src_name, shape=shape, off=off, size=size,
|
||||
np_dtype=_NP[ttype], op_index=oi)
|
||||
|
||||
|
||||
def _tt_name(code):
|
||||
from ai_edge_litert import schema_py_generated as schema
|
||||
return next((k for k, v in schema.TensorType.__dict__.items() if v == code), str(code))
|
||||
|
||||
|
||||
# ── name classification → checkpoint-layer resolution ──────────────────────────
|
||||
|
||||
_LEAF = re.compile(r"torch\.nn\.modules\.linear\.Linear_([A-Za-z0-9_]+)")
|
||||
_BLOCK = re.compile(r"TransformerBlock_(\d+)/")
|
||||
|
||||
# checkpoint (post model.-strip) block-module suffix → (module tag, tflite leaf, glu)
|
||||
_BLOCK_CKPT = {
|
||||
"self_attn.to_qkv": ("self_attn", "to_qkv", False),
|
||||
"self_attn.to_out": ("self_attn", "to_out", False),
|
||||
"cross_attn.to_q": ("cross_attn", "to_q", False),
|
||||
"cross_attn.to_kv": ("cross_attn", "to_kv", False),
|
||||
"cross_attn.to_out": ("cross_attn", "to_out", False),
|
||||
"ff.ff.0.proj": ("ff", "proj", True),
|
||||
"ff.ff.2": ("ff", "2", False),
|
||||
# to_local_embed handled separately (topological order); accept both the
|
||||
# torch (.0/.2) and MLX (.seq.0/.seq.2) spellings.
|
||||
"to_local_embed.0": ("local_embed", "0", False),
|
||||
"to_local_embed.2": ("local_embed", "2", False),
|
||||
"to_local_embed.seq.0": ("local_embed", "0", False),
|
||||
"to_local_embed.seq.2": ("local_embed", "2", False),
|
||||
}
|
||||
|
||||
# non-block singleton checkpoint names → a substring that uniquely identifies the
|
||||
# tflite FC name (verified against the exported graph).
|
||||
_SINGLETON_CKPT = {
|
||||
"transformer.project_in": "Linear_project_in",
|
||||
"transformer.project_out": "Linear_project_out",
|
||||
"to_timestep_embed.0": "Sequential_to_timestep_embed/torch.nn.modules.linear.Linear_0",
|
||||
"to_timestep_embed.2": "Sequential_to_timestep_embed/torch.nn.modules.linear.Linear_2",
|
||||
"to_cond_embed.0": "Sequential_to_cond_embed/torch.nn.modules.linear.Linear_0",
|
||||
"to_cond_embed.2": "Sequential_to_cond_embed/torch.nn.modules.linear.Linear_2",
|
||||
"to_global_embed.0": "Sequential_to_global_embed/torch.nn.modules.linear.Linear_0",
|
||||
"to_global_embed.2": "Sequential_to_global_embed/torch.nn.modules.linear.Linear_2",
|
||||
"transformer.global_cond_embedder.0": "Sequential_global_cond_embedder/torch.nn.modules.linear.Linear_0",
|
||||
"transformer.global_cond_embedder.2": "Sequential_global_cond_embedder/torch.nn.modules.linear.Linear_2",
|
||||
}
|
||||
|
||||
|
||||
def _module_tag(name: str) -> str | None:
|
||||
if "SelfAttention_self_attn" in name:
|
||||
return "self_attn"
|
||||
if "CrossAttention_cross_attn" in name:
|
||||
return "cross_attn"
|
||||
if "Seq_to_local_embed" in name or "Sequential_seq" in name:
|
||||
return "local_embed"
|
||||
if "FeedForward_ff" in name:
|
||||
return "ff"
|
||||
return None
|
||||
|
||||
|
||||
class _Mapper:
|
||||
"""Resolves checkpoint layer names → FC weight records for one base model."""
|
||||
|
||||
def __init__(self, fcs):
|
||||
self._fcs = fcs
|
||||
# tier-1 block linears: (block, module, leaf, glu) → record
|
||||
self._block = {}
|
||||
# tier-2 local_embed: leaf ('0'/'2') → [records in topological order]
|
||||
self._local = {"0": [], "2": []}
|
||||
self._seconds = [] # tier-4
|
||||
for rec in fcs: # fcs already in topological (op) order
|
||||
name = rec["name"]
|
||||
mod = _module_tag(name)
|
||||
mleaf = _LEAF.search(name)
|
||||
leaf = mleaf.group(1) if mleaf else None
|
||||
mblk = _BLOCK.search(name)
|
||||
if mod == "local_embed" and leaf in ("0", "2"):
|
||||
self._local[leaf].append(rec)
|
||||
elif mod in ("self_attn", "cross_attn", "ff") and mblk and leaf:
|
||||
glu = "_GLUWrap_0" in name
|
||||
self._block[(int(mblk.group(1)), mod, leaf, glu)] = rec
|
||||
if rec["shape"] == (768, 256):
|
||||
self._seconds.append(rec)
|
||||
|
||||
def resolve(self, layer: str, want_shape):
|
||||
"""Return the FC record for a checkpoint layer, or raise LoraError."""
|
||||
if layer == lc.COND_SECONDS_LAYER:
|
||||
cands = [r for r in self._seconds if r["shape"] == want_shape]
|
||||
return self._one(layer, want_shape, cands, "seconds conditioner")
|
||||
|
||||
stripped = layer[len("model."):] if layer.startswith("model.") else layer
|
||||
mm = re.match(r"transformer\.layers\.(\d+)\.(.+)$", stripped)
|
||||
if mm:
|
||||
block, suffix = int(mm.group(1)), mm.group(2)
|
||||
spec = _BLOCK_CKPT.get(suffix)
|
||||
if spec is None:
|
||||
raise lc.LoraError(f"{layer}: unrecognised block sub-module {suffix!r}")
|
||||
module, leaf, glu = spec
|
||||
if module == "local_embed":
|
||||
# tier 2 — topological rank within the group is the block index.
|
||||
group = self._local[leaf]
|
||||
if block >= len(group):
|
||||
raise lc.LoraError(
|
||||
f"{layer}: to_local_embed block {block} out of range "
|
||||
f"(only {len(group)} in the graph)")
|
||||
rec = group[block]
|
||||
if rec["shape"] != tuple(want_shape):
|
||||
raise lc.LoraError(
|
||||
f"{layer}: shape {rec['shape']} != adapter {tuple(want_shape)} "
|
||||
f"— wrong base model?")
|
||||
return rec
|
||||
rec = self._block.get((block, module, leaf, glu))
|
||||
cands = [rec] if rec is not None and rec["shape"] == tuple(want_shape) else []
|
||||
return self._one(layer, want_shape, cands, "block linear")
|
||||
|
||||
sub = _SINGLETON_CKPT.get(stripped)
|
||||
if sub is None:
|
||||
raise lc.LoraError(f"{layer}: no TFLite tensor mapping for this layer")
|
||||
cands = [r for r in self._fcs if sub in r["name"] and r["shape"] == tuple(want_shape)]
|
||||
return self._one(layer, want_shape, cands, "singleton")
|
||||
|
||||
@staticmethod
|
||||
def _one(layer, want_shape, cands, tier):
|
||||
if len(cands) == 1:
|
||||
return cands[0]
|
||||
raise lc.LoraError(
|
||||
f"{layer}: expected exactly 1 matching FC ({tier}, shape {tuple(want_shape)}), "
|
||||
f"found {len(cands)} — wrong base model or unsupported target module?"
|
||||
)
|
||||
|
||||
|
||||
# ── clone + patch + cache ──────────────────────────────────────────────────────
|
||||
|
||||
def _clone(src: str, dst: str) -> None:
|
||||
"""COW-clone src→dst (APFS ``cp -c``, instant) or fall back to a full copy."""
|
||||
try:
|
||||
r = subprocess.run(["cp", "-c", src, dst], capture_output=True)
|
||||
if r.returncode == 0 and os.path.exists(dst):
|
||||
return
|
||||
except (OSError, FileNotFoundError):
|
||||
pass
|
||||
shutil.copyfile(src, dst)
|
||||
|
||||
|
||||
def _cache_key(base_path: str, specs) -> str:
|
||||
h = hashlib.sha256()
|
||||
h.update(_MERGE_VERSION.encode())
|
||||
st = os.stat(base_path)
|
||||
h.update(f"{os.path.basename(base_path)}:{st.st_size}".encode())
|
||||
for path, strength in sorted((lc._resolve_path(s["path"]), s["strength"]) for s in specs):
|
||||
fh = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1 << 20), b""):
|
||||
fh.update(chunk)
|
||||
h.update(f"{fh.hexdigest()}:{strength}".encode())
|
||||
return h.hexdigest()[:12]
|
||||
|
||||
|
||||
def get_patched_dit(base_path, specs, *, family: str, precision: str,
|
||||
cache_dir=None, log=print) -> Path:
|
||||
"""Return a path to a DiT ``.tflite`` with ``specs`` (from
|
||||
:func:`lora_core.parse_lora_spec`) merged into its weights. Result is cached
|
||||
by content hash; a hit returns instantly. Raises LoraError on any mapping,
|
||||
shape, precision, or trust-boundary failure — never mis-patches silently.
|
||||
"""
|
||||
base_path = str(base_path)
|
||||
if precision in _QUANTIZED_PRECISIONS:
|
||||
raise lc.LoraError(
|
||||
f"--lora is not supported with --precision {precision} (quantized-int8) — "
|
||||
f"merging int8 weights would destroy the GPTQ grid. Use fp32 or w16a32."
|
||||
)
|
||||
|
||||
# Cache next to the family dirs (models/tflite/lora_cache/) — use the
|
||||
# base's OWN path, not resolve(): the base is usually a symlink into the HF
|
||||
# cache, and resolving it would drop the patched (multi-GB) clones there.
|
||||
cache_dir = Path(cache_dir or (Path(base_path).parent.parent / "lora_cache"))
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
key = _cache_key(base_path, specs)
|
||||
out = cache_dir / f"dit_{family}_{precision}_{key}.tflite"
|
||||
if out.exists():
|
||||
log(f"lora: reusing cached patched model {out.name}")
|
||||
return out
|
||||
|
||||
# Parse all adapters up front (fail fast before the clone).
|
||||
parsed = []
|
||||
for s in specs:
|
||||
path = lc._resolve_path(s["path"])
|
||||
adapter_type, scaling, layers = lc.parse_adapter(path)
|
||||
parsed.append((os.path.basename(path), adapter_type, scaling,
|
||||
float(s["strength"]), layers))
|
||||
log(f"lora: {os.path.basename(path)} — {adapter_type}, "
|
||||
f"scaling {scaling:g}, strength {s['strength']:g}, {len(layers)} layers")
|
||||
|
||||
# Map every adapter layer against the base model BEFORE cloning (so a
|
||||
# wrong-base adapter errors without leaving a stale clone behind).
|
||||
with open(base_path, "rb") as f:
|
||||
buf = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
|
||||
fcs = list(_fc_weight_buffers(buf))
|
||||
mapper = _Mapper(fcs)
|
||||
# accum[off] = [record, W0(fp32), summed_delta(fp32)]
|
||||
accum = {}
|
||||
matched = 0
|
||||
for name, atype, scaling, strength, layers in parsed:
|
||||
need = lc._PARAMS_FOR.get(atype, ())
|
||||
for layer, params in layers.items():
|
||||
want = (params.get("lora_B", params.get("M_xs")).shape[0],
|
||||
params["lora_A"].shape[1] if "lora_A" in params
|
||||
else params["M_xs"].shape[1])
|
||||
# For -xs, fan_in comes from the base weight; resolve then re-check.
|
||||
rec = mapper.resolve(layer, want if atype not in
|
||||
("lora-xs", "dora-rows-xs", "dora-cols-xs", "bora-xs")
|
||||
else _shape_from_leaf(mapper, layer, params))
|
||||
W0 = np.frombuffer(buf, rec["np_dtype"],
|
||||
count=rec["size"] // np.dtype(rec["np_dtype"]).itemsize,
|
||||
offset=rec["off"]).reshape(rec["shape"]).astype(np.float32)
|
||||
missing = [n for n in need if n not in params]
|
||||
if missing:
|
||||
raise lc.LoraError(f"{layer}: adapter is {atype} but missing {missing}")
|
||||
lc.check_shapes(layer, W0, params, atype)
|
||||
merged = lc.merged_weight(W0, params, atype, scaling)
|
||||
delta = strength * (merged - W0)
|
||||
if rec["off"] in accum:
|
||||
accum[rec["off"]][2] += delta
|
||||
else:
|
||||
accum[rec["off"]] = [rec, W0, delta]
|
||||
matched += 1
|
||||
buf.close()
|
||||
if not matched:
|
||||
raise lc.LoraError("no adapter layers matched this DiT (wrong base model?)")
|
||||
log(f"lora: mapped {len(accum)} weight tensor(s) across {matched} adapter layer(s)")
|
||||
|
||||
# Clone, then patch each buffer once (W0 + summed delta).
|
||||
tmp = out.with_suffix(".tflite.tmp")
|
||||
if tmp.exists():
|
||||
tmp.unlink()
|
||||
_clone(base_path, str(tmp))
|
||||
n_bytes = 0
|
||||
with open(tmp, "r+b") as fw:
|
||||
for off, (rec, W0, delta) in accum.items():
|
||||
Wp = (W0 + delta).astype(rec["np_dtype"])
|
||||
b = Wp.tobytes()
|
||||
assert len(b) == rec["size"], (rec["name"], len(b), rec["size"])
|
||||
fw.seek(off)
|
||||
fw.write(b)
|
||||
n_bytes += rec["size"]
|
||||
fw.flush()
|
||||
os.fsync(fw.fileno())
|
||||
os.replace(tmp, out)
|
||||
log(f"lora: patched {n_bytes / 1e6:.0f} MB → {out.name}")
|
||||
return out
|
||||
|
||||
|
||||
def _shape_from_leaf(mapper: _Mapper, layer: str, params: dict):
|
||||
"""For -xs adapters the fan_in isn't in the params (only M_xs r×r), so map by
|
||||
a shape probe: try each FC whose descriptor matches and return its shape.
|
||||
Resolves by re-using the mapper's block/singleton/seconds index with a
|
||||
wildcard shape, then reads the found record's shape."""
|
||||
# Cheap approach: the -xs magnitude (rows/cols) or M_xs can't give fan_in, so
|
||||
# locate the record ignoring shape via a temporary relaxation.
|
||||
stripped = layer[len("model."):] if layer.startswith("model.") else layer
|
||||
if layer == lc.COND_SECONDS_LAYER and mapper._seconds:
|
||||
return mapper._seconds[0]["shape"]
|
||||
mm = re.match(r"transformer\.layers\.(\d+)\.(.+)$", stripped)
|
||||
if mm:
|
||||
block, suffix = int(mm.group(1)), mm.group(2)
|
||||
spec = _BLOCK_CKPT.get(suffix)
|
||||
if spec:
|
||||
module, leaf, glu = spec
|
||||
if module == "local_embed":
|
||||
grp = mapper._local[leaf]
|
||||
if block < len(grp):
|
||||
return grp[block]["shape"]
|
||||
else:
|
||||
rec = mapper._block.get((block, module, leaf, glu))
|
||||
if rec:
|
||||
return rec["shape"]
|
||||
sub = _SINGLETON_CKPT.get(stripped)
|
||||
if sub:
|
||||
for r in mapper._fcs:
|
||||
if sub in r["name"]:
|
||||
return r["shape"]
|
||||
raise lc.LoraError(f"{layer}: cannot locate a base weight for this -xs layer")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,932 @@
|
||||
"""SA3 text-to-audio (+ audio-to-audio + inpainting + CFG) on CPU via the BAKED-I/O
|
||||
varlen TFLite / LiteRT models — the portable-CPU sibling of the MLX and TensorRT
|
||||
releases. CLI at feature parity with sa3_mlx.py / sa3_trt.py, sharing their flag names.
|
||||
|
||||
Baked I/O (ONNX / TensorRT convention — conditioner + patch/unpatch are IN-GRAPH):
|
||||
DiT (6-input): x[1,256,L], t[1], t5_hidden[1,256,768], t5_mask[1,256],
|
||||
seconds_total[1], local_add_cond[1,257,L] -> velocity[1,256,L]
|
||||
(feed RAW T5 outputs; prompt-padding + seconds-embed happen in-graph)
|
||||
Encoder (audio-in): audio[1,2,N] -> latents[1,256,N/4096] (patch-encode baked)
|
||||
Decoder (audio-out): latents[1,256,L] -> audio[1,2,L*4096] (unpatch baked)
|
||||
T5Gemma: t5gemma/encoder_fp16.tflite (fixed 256 text tokens; tokenizer bundled)
|
||||
|
||||
Modes (identical flags to sa3_mlx / sa3_trt):
|
||||
text-to-audio --prompt P
|
||||
audio-to-audio --prompt P --init-audio IN.wav [--init-noise-level σ]
|
||||
inpainting --prompt P --init-audio IN.wav --inpaint-range START,END
|
||||
negative CFG --prompt P --cfg N [--negative-prompt P_NEG] [--apg A]
|
||||
|
||||
CFG uses the TRT/ONNX baked-conditioner convention: the uncond branch feeds the negative
|
||||
prompt's T5 output (or an all-zero hidden+mask, which the in-graph conditioner turns into
|
||||
learned padding embeddings) — it does NOT zero cross_attn like MLX (our DiT bakes the
|
||||
conditioner, so there's no raw cross_attn). CFG runs cond+uncond as ONE batch=2 invoke on the
|
||||
variable-batch DiT (--cfg-batched, default; ~7-29%% faster on Apple-Silicon AMX) or as a
|
||||
sequential batch=1 dual-pass (--no-cfg-batched). Bit-identical at fp32/w8a32; ~80 dB at
|
||||
w16a32; w8a8-dyn diverges by design (activation scales shared across the batch).
|
||||
|
||||
All model files auto-download from HuggingFace (stabilityai/stable-audio-3-optimized)
|
||||
on first use and symlink into models/tflite/ from the HF cache. See scripts/weights.py.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse, math, os, random, re, subprocess, sys, time, wave
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
|
||||
# Windows consoles default to a legacy code page (cp1252/cp437) that can't encode
|
||||
# the banners' box-drawing / emoji characters. Force UTF-8 with replacement so no
|
||||
# print can crash the CLI. No-op on macOS/Linux, where stdout is UTF-8 already.
|
||||
for _stream in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
_stream.reconfigure(encoding="utf-8", errors="replace")
|
||||
except AttributeError:
|
||||
pass # exotic stream without reconfigure() (non-TextIOWrapper) — leave as-is
|
||||
|
||||
# Enable ANSI/VT escape processing in legacy Windows consoles (cmd.exe). Calling
|
||||
# os.system("") flips the console into VT mode; harmless elsewhere and on
|
||||
# modern Windows Terminal, which has it on already.
|
||||
if os.name == "nt":
|
||||
os.system("")
|
||||
|
||||
# The arrow-key picker below uses POSIX terminal APIs. Keep the non-interactive
|
||||
# CLI usable on Windows, where those modules are not available.
|
||||
try:
|
||||
import termios
|
||||
import tty
|
||||
except ImportError:
|
||||
termios = tty = None
|
||||
|
||||
REPO = Path(__file__).resolve().parent.parent # project root (scripts/ is one level down)
|
||||
sys.path.insert(0, str(REPO)) # so `from models.defs.* import` resolves
|
||||
sys.path.insert(0, str(REPO / "scripts")) # so `from weights import *` resolves
|
||||
|
||||
from models.defs import tflite_pipeline as P # Tokenizer, T5GemmaTFLite, build_pingpong_schedule, make_noise, sample, save_wav
|
||||
from weights import ensure_local, is_present, PRECISIONS, dit_rel, dec_rel, enc_rel
|
||||
|
||||
SAMPLE_RATE = 44100
|
||||
SAMPLES_PER_LATENT = 4096
|
||||
COND_TOKENS = 256
|
||||
COND_DIM = 768
|
||||
MIN_SIGMA = 0.01 # rf_denoiser is undefined at t≈0 → NaN below this
|
||||
|
||||
# ─── Model manifest (local rel paths; resolved via weights.ensure_local at load) ───
|
||||
# --precision picks the DiT + decoder variant (same flag as the TensorRT release's
|
||||
# --precision; values here use the wXaY convention — weights/activations bit-widths,
|
||||
# "16" = fp16, there is no int16). Encoders + T5Gemma have one precision, like TRT.
|
||||
# fp32 reference models (default — on CPU fp16 is SLOWER than fp32, so unlike
|
||||
# TRT the fastest-and-accurate choice is fp32)
|
||||
# w16a32 fp16 weights / fp32 activations — half size, ≈lossless, 1.5-3× slower
|
||||
# on an M4 Pro (legacy name: fp16; ≈ TRT's fp16 in spirit,
|
||||
# storage-only here)
|
||||
# w8a32 GPTQ weight-only int8 — ¼ size at fp32 speed (legacy: woint8)
|
||||
# w8a8-dyn GPTQ dynamic int8 — fastest (~1.2-1.3×), lowest quality (legacy: dynint8)
|
||||
# PRECISIONS + the path builders live in weights.py (single source of truth with
|
||||
# the download manifest — a precision the CLI accepts is guaranteed downloadable).
|
||||
DIT_REL = {d: dit_rel(d) for d in ("sm-music", "sm-sfx", "medium")} # fp32 (picker)
|
||||
DEC_REL = {d: dec_rel(d) for d in ("same-s", "same-l")}
|
||||
ENC_REL = {d: enc_rel(d) for d in ("same-s", "same-l")}
|
||||
T5_REL = "models/tflite/t5gemma/encoder_fp16.tflite"
|
||||
DEFAULT_DECODER = {"sm-music": "same-s", "sm-sfx": "same-s", "medium": "same-l"}
|
||||
|
||||
# SAME-L dense SWA mask is O(S^2): chunk long decodes. SAME-S has a narrow field: whole.
|
||||
SAMEL_CHUNK = 64 # latent tokens/window — throughput optimum (6.5x RT) vs the O(S^2) dense SWA mask
|
||||
SAMEL_OVERLAP = 8 # symmetric latent-token context each interior side (SAME-L needs >=8)
|
||||
|
||||
|
||||
def _free():
|
||||
import gc; gc.collect()
|
||||
|
||||
|
||||
# ─── ANSI display (match sa3_mlx.py style) ─────────────────────────────────
|
||||
_USE_COLOR = sys.stdout.isatty()
|
||||
_RULE_W = 64
|
||||
def _c(code, s): 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 red(s): return _c("31", s)
|
||||
def rule(ch="━"): print(cyan(ch * _RULE_W))
|
||||
def banner(t):
|
||||
rule(); print(f" {bold(t)}"); rule()
|
||||
def stage(idx, label, ms=None):
|
||||
head = f" {cyan(idx)} {bold(label)}"
|
||||
if ms is None:
|
||||
print(head, flush=True); return
|
||||
visible = len(f" {idx} {label}")
|
||||
dots = dim("·" * max(2, _RULE_W - visible - 9))
|
||||
print(f"{head} {dots} {yellow(f'{ms:>5.0f} ms')}", flush=True)
|
||||
def sub(t): print(f" {dim(t)}", flush=True)
|
||||
|
||||
|
||||
# ─── Interactive arrow-key picker (from sa3_mlx.py) ────────────────────────
|
||||
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 termios is None or tty is None or 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)
|
||||
for _ in options:
|
||||
print()
|
||||
try:
|
||||
tty.setcbreak(fd)
|
||||
while True:
|
||||
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_REL.keys()), default="sm-music")
|
||||
print(f" → {args.dit}")
|
||||
if args.decoder is None:
|
||||
suggested = DEFAULT_DECODER[args.dit]
|
||||
args.decoder = _arrow_pick("Choose audio decoder:", list(DEC_REL.keys()), default=suggested)
|
||||
print(f" → {args.decoder}")
|
||||
if args.seed is None:
|
||||
args.seed = random.randint(0, 2**31 - 1)
|
||||
return args
|
||||
|
||||
|
||||
def _preflight_download(args, dec) -> None:
|
||||
"""Resolve every model 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 = [T5_REL, dit_rel(args.dit, args.dit_precision), dec_rel(dec, args.decoder_precision)]
|
||||
if args.init_audio:
|
||||
needed.append(enc_rel(dec, args.encoder_precision))
|
||||
missing = [p for p in needed if not is_present(p)]
|
||||
if not missing:
|
||||
return
|
||||
print(f" Fetching {len(missing)} missing model file(s) before starting:")
|
||||
for rel in missing:
|
||||
ensure_local(rel)
|
||||
print()
|
||||
|
||||
|
||||
# ─── WAV read (16-bit PCM fast path + ffmpeg fallback) ─────────────────────
|
||||
def read_wav(path: str) -> np.ndarray:
|
||||
"""Return (2, T) float32 in [-1, 1]. 16-bit/44.1k native; else via ffmpeg. Mono→stereo."""
|
||||
try:
|
||||
with wave.open(path, "rb") as w:
|
||||
nch, sw, sr, n = w.getnchannels(), w.getsampwidth(), w.getframerate(), w.getnframes()
|
||||
if sr == SAMPLE_RATE and sw == 2:
|
||||
raw = np.frombuffer(w.readframes(n), np.int16).astype(np.float32) / 32767.0
|
||||
if nch == 1:
|
||||
return np.stack([raw, raw], 0)
|
||||
return raw.reshape(-1, nch).T[:2]
|
||||
except wave.Error:
|
||||
pass
|
||||
try:
|
||||
r = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-f", "s16le",
|
||||
"-ar", str(SAMPLE_RATE), "-ac", "2", "-"],
|
||||
capture_output=True, check=True)
|
||||
except FileNotFoundError:
|
||||
ffmpeg_hint = {"darwin": "brew install ffmpeg",
|
||||
"win32": "winget install ffmpeg (or: choco install ffmpeg)",
|
||||
}.get(sys.platform, "apt install ffmpeg")
|
||||
raise RuntimeError(f"{path}: unsupported WAV format. Install ffmpeg for 24/32-bit or "
|
||||
f"non-44.1kHz input: {ffmpeg_hint}")
|
||||
except subprocess.CalledProcessError as e:
|
||||
raise RuntimeError(f"{path}: ffmpeg failed — {e.stderr.decode().strip()}")
|
||||
raw = np.frombuffer(r.stdout, np.int16).astype(np.float32) / 32767.0
|
||||
return raw.reshape(-1, 2).T
|
||||
|
||||
|
||||
# ─── Baked DiT backend (6-input; batched OR sequential CFG + inpaint local_add_cond) ───
|
||||
def _apply_cfg(x, t, v_cond, v_uncond, cfg, apg):
|
||||
"""Combine cond/uncond velocities in denoised space (RF), optional APG. fp32.
|
||||
Matches sa3_trt_core.model_fn. Returns cfg_v so `denoised = x - t*v` is the guided one."""
|
||||
sigma = np.float32(t)
|
||||
xf = x.astype(np.float32)
|
||||
cond_d = xf - v_cond.astype(np.float32) * sigma
|
||||
uncond_d = xf - v_uncond.astype(np.float32) * sigma
|
||||
diff = cond_d - uncond_d
|
||||
if apg <= 0.0:
|
||||
cfg_diff = diff
|
||||
else:
|
||||
norm = np.sqrt((cond_d * cond_d).sum(axis=(-2, -1), keepdims=True))
|
||||
unit = cond_d / np.maximum(norm, 1e-8)
|
||||
parallel = (diff * unit).sum(axis=(-2, -1), keepdims=True) * unit
|
||||
diff_orth = diff - parallel
|
||||
cfg_diff = diff_orth if apg >= 1.0 else (apg * diff_orth + (1.0 - apg) * diff)
|
||||
cfg_d = cond_d + (cfg - 1.0) * cfg_diff
|
||||
return ((xf - cfg_d) / sigma).astype(np.float32)
|
||||
|
||||
|
||||
class BakedDiT:
|
||||
"""model_fn(x,t,cross,gcond)->v compatible with P.sample. cross/gcond are IGNORED
|
||||
(conditioning is baked in-graph, driven by the raw T5 outputs held here).
|
||||
|
||||
cfg==1.0 -> one batch=1 forward with the (cond) T5 output.
|
||||
cfg!=1.0 -> CFG combining a cond and an uncond (null_hidden/null_mask) velocity in
|
||||
denoised space with optional APG. Two backends (bit-identical output at
|
||||
fp32/w8a32; ~80 dB at w16a32; w8a8-dyn diverges — batch-shared activation
|
||||
scales; auto-falls-back to batch=1 for cfg==1.0):
|
||||
batched=True (default): ONE batch=2 invoke/step — cond=row0, uncond=row1 — on the
|
||||
canonical variable-batch DiT (its batch axis is dynamic). ~7-29%% faster
|
||||
on Apple-Silicon (the AMX matrix unit amortizes weight loads across the 2 rows).
|
||||
batched=False: SEQUENTIAL dual-pass — two batch=1 invokes/step, like TensorRT (whose
|
||||
engine is static-batch=1). Use if a backend/delegate dislikes batch=2.
|
||||
Returns cfg_v so P.sample's `denoised = x - t*v` yields the guided denoised."""
|
||||
def __init__(self, path, L, t5_hidden, t5_mask, seconds, threads=8,
|
||||
cfg=1.0, apg=1.0, null_hidden=None, null_mask=None, local_add_cond=None,
|
||||
batched=True):
|
||||
from ai_edge_litert import interpreter as tfl
|
||||
self.it = tfl.Interpreter(model_path=str(path), num_threads=threads)
|
||||
det = self.it.get_input_details()
|
||||
def pick(pred): return [d for d in det if pred(d)]
|
||||
self.i_x = pick(lambda d: len(d["shape"]) == 3 and d["shape"][1] == 256 and d["shape"][2] != COND_DIM)[0]
|
||||
self.i_lac = pick(lambda d: len(d["shape"]) == 3 and d["shape"][1] == 257)[0]
|
||||
self.i_t5h = pick(lambda d: len(d["shape"]) == 3 and d["shape"][2] == COND_DIM)[0]
|
||||
self.i_t5m = pick(lambda d: len(d["shape"]) == 2)[0]
|
||||
scalars = sorted(pick(lambda d: len(d["shape"]) == 1), key=lambda d: d["name"])
|
||||
self.i_t, self.i_sec = scalars[0], scalars[1] # args_1=t < args_4=seconds by name
|
||||
self.out = self.it.get_output_details()[0]["index"]
|
||||
# Track the currently-allocated (batch, length) so set_conditioning can skip the
|
||||
# expensive resize+allocate when only the conditioning tensors change.
|
||||
self._alloc_B = self._alloc_L = None
|
||||
self.set_conditioning(L, t5_hidden, t5_mask, seconds, cfg=cfg, apg=apg,
|
||||
null_hidden=null_hidden, null_mask=null_mask,
|
||||
local_add_cond=local_add_cond, batched=batched)
|
||||
|
||||
def set_conditioning(self, L, t5_hidden, t5_mask, seconds, *, cfg=1.0, apg=1.0,
|
||||
null_hidden=None, null_mask=None, local_add_cond=None, batched=True):
|
||||
"""(Re)bind per-generation conditioning onto this (possibly reused) interpreter.
|
||||
|
||||
The baked DiT holds the conditioning (T5 hidden/mask, seconds, local_add_cond) as
|
||||
RESIDENT input tensors — a fresh generation only changes those, not the graph. A
|
||||
cached BakedDiT can therefore be re-pointed at a new prompt/length WITHOUT rebuilding
|
||||
the interpreter or re-packing XNNPACK's weights: resize+allocate_tensors runs only
|
||||
when the batch (cfg on/off) or length L actually changes, otherwise the residents are
|
||||
overwritten in place (same reuse pattern BakedDecoder/BakedEncoder use for L). Called
|
||||
from __init__ (so the CLI's one-shot path is unchanged) and by the gradio DiT cache."""
|
||||
self.L = L
|
||||
self.cfg = float(cfg); self.apg = float(apg)
|
||||
self.n_fwd = 0
|
||||
# Batched CFG only applies when there's an uncond branch to co-batch (cfg != 1.0).
|
||||
self.batched = bool(batched) and self.cfg != 1.0
|
||||
B = 2 if self.batched else 1
|
||||
self.B = B
|
||||
# Resize batch axis to B (the canonical DiT is variable-batch) + length axis to L.
|
||||
if (self._alloc_B, self._alloc_L) != (B, L):
|
||||
self.it.resize_tensor_input(self.i_x["index"], [B, 256, L], strict=False)
|
||||
self.it.resize_tensor_input(self.i_lac["index"], [B, 257, L], strict=False)
|
||||
self.it.resize_tensor_input(self.i_t5h["index"], [B, COND_TOKENS, COND_DIM], strict=False)
|
||||
self.it.resize_tensor_input(self.i_t5m["index"], [B, COND_TOKENS], strict=False)
|
||||
self.it.resize_tensor_input(self.i_t["index"], [B], strict=False)
|
||||
self.it.resize_tensor_input(self.i_sec["index"], [B], strict=False)
|
||||
self.it.allocate_tensors()
|
||||
self._alloc_B, self._alloc_L = B, L
|
||||
|
||||
t5h = t5_hidden.astype(np.float32)
|
||||
t5m = t5_mask.astype(np.float32)
|
||||
sec = np.array([np.float32(seconds)], np.float32)
|
||||
lac = (np.zeros((1, 257, L), np.float32) if local_add_cond is None
|
||||
else local_add_cond.astype(np.float32))
|
||||
self.null_h = None if null_hidden is None else null_hidden.astype(np.float32)
|
||||
self.null_m = None if null_mask is None else null_mask.astype(np.float32)
|
||||
if self.batched:
|
||||
# row0 = cond, row1 = uncond. t5_hidden/mask differ per row; seconds + lac are
|
||||
# shared → tiled. All four are constant across steps → set resident once.
|
||||
self.it.set_tensor(self.i_t5h["index"], np.concatenate([t5h, self.null_h], axis=0))
|
||||
self.it.set_tensor(self.i_t5m["index"], np.concatenate([t5m, self.null_m], axis=0))
|
||||
self.it.set_tensor(self.i_sec["index"], np.concatenate([sec, sec], axis=0))
|
||||
self.it.set_tensor(self.i_lac["index"], np.concatenate([lac, lac], axis=0))
|
||||
else:
|
||||
self.t5h, self.t5m = t5h, t5m
|
||||
self.it.set_tensor(self.i_sec["index"], sec) # seconds + lac constant → resident
|
||||
self.it.set_tensor(self.i_lac["index"], lac)
|
||||
return self
|
||||
|
||||
def _fwd(self, x, t, t5h, t5m):
|
||||
"""One batch=1 invoke (cfg==1.0 or sequential CFG)."""
|
||||
self.it.set_tensor(self.i_x["index"], x.astype(np.float32))
|
||||
self.it.set_tensor(self.i_t["index"], np.array([np.float32(t)], np.float32))
|
||||
self.it.set_tensor(self.i_t5h["index"], t5h)
|
||||
self.it.set_tensor(self.i_t5m["index"], t5m)
|
||||
self.it.invoke()
|
||||
self.n_fwd += 1
|
||||
return self.it.get_tensor(self.out).copy()
|
||||
|
||||
def _fwd_batched(self, x, t):
|
||||
"""One batch=2 invoke -> (v_cond, v_uncond). x is the shared batch=1 state (both CFG
|
||||
branches denoise the same latent), tiled to [2,256,L]; t5h/t5m/sec/lac already resident."""
|
||||
x2 = np.concatenate([x, x], axis=0).astype(np.float32)
|
||||
self.it.set_tensor(self.i_x["index"], x2)
|
||||
self.it.set_tensor(self.i_t["index"], np.array([np.float32(t), np.float32(t)], np.float32))
|
||||
self.it.invoke()
|
||||
self.n_fwd += 2
|
||||
v2 = self.it.get_tensor(self.out)
|
||||
return v2[0:1].copy(), v2[1:2].copy()
|
||||
|
||||
def __call__(self, x, t, cross=None, gcond=None):
|
||||
if self.cfg == 1.0:
|
||||
return self._fwd(x, t, self.t5h, self.t5m)
|
||||
if self.batched:
|
||||
v_cond, v_uncond = self._fwd_batched(x, t)
|
||||
else:
|
||||
v_cond = self._fwd(x, t, self.t5h, self.t5m)
|
||||
v_uncond = self._fwd(x, t, self.null_h, self.null_m)
|
||||
return _apply_cfg(x, t, v_cond, v_uncond, self.cfg, self.apg)
|
||||
|
||||
|
||||
# ─── Baked audio-in encoder (audio -> latents; SAME-S needs even L) ────────
|
||||
class BakedEncoder:
|
||||
def __init__(self, path, threads=8, needs_even=False):
|
||||
from ai_edge_litert import interpreter as tfl
|
||||
self.it = tfl.Interpreter(model_path=str(path), num_threads=threads)
|
||||
self.i = self.it.get_input_details()[0]["index"]
|
||||
self.o = self.it.get_output_details()[0]["index"]
|
||||
self.needs_even = needs_even
|
||||
self._cur_N = None
|
||||
|
||||
def _resize(self, N):
|
||||
if self._cur_N != N:
|
||||
self.it.resize_tensor_input(self.i, [1, 2, N], strict=False)
|
||||
self.it.allocate_tensors()
|
||||
self._cur_N = N
|
||||
|
||||
def encode(self, audio, T_lat):
|
||||
"""audio: (1,2,M), M a multiple of 4096 (caller pads to even L for SAME-S).
|
||||
Returns latents (1,256,T_lat) — trimmed back to the natural (decoder-independent) T_lat."""
|
||||
self._resize(audio.shape[-1])
|
||||
self.it.set_tensor(self.i, audio.astype(np.float32))
|
||||
self.it.invoke()
|
||||
return self.it.get_tensor(self.o)[:, :, :T_lat].copy()
|
||||
|
||||
|
||||
# ─── Baked audio-out decoder (whole for SAME-S; chunked for SAME-L) ─────────
|
||||
class BakedDecoder:
|
||||
def __init__(self, path, threads=8, needs_even=False):
|
||||
from ai_edge_litert import interpreter as tfl
|
||||
self.tfl = tfl
|
||||
self.path = str(path)
|
||||
self.threads = threads
|
||||
self.needs_even = needs_even # SAME-S: T_aud=L*16 must be %32 -> pad odd L->even
|
||||
self.it = tfl.Interpreter(model_path=self.path, num_threads=threads)
|
||||
self.i = self.it.get_input_details()[0]["index"]
|
||||
self.o = self.it.get_output_details()[0]["index"]
|
||||
self.i_det = self.it.get_input_details()[0]
|
||||
self._cur_L = None
|
||||
|
||||
def _resize(self, L):
|
||||
if self._cur_L != L:
|
||||
self.it.resize_tensor_input(self.i, [1, 256, L], strict=False)
|
||||
self.it.allocate_tensors()
|
||||
self._cur_L = L
|
||||
|
||||
def decode_whole(self, latents):
|
||||
L = latents.shape[2]
|
||||
if self.needs_even and L % 2 != 0:
|
||||
# SAME-S needs even L. Pad one edge-replicated latent token, decode at L+1,
|
||||
# trim the extra token's audio. Narrow receptive field => negligible boundary
|
||||
# effect on the kept L*4096 samples. Keeps odd-L requests working without
|
||||
# changing the DiT/noise path (which stays natural-ceil, MLX/TRT-matched).
|
||||
latents = np.concatenate([latents, latents[:, :, -1:]], axis=2)
|
||||
self._resize(L + 1)
|
||||
self.it.set_tensor(self.i, latents.astype(np.float32))
|
||||
self.it.invoke()
|
||||
return self.it.get_tensor(self.o)[:, :, :L * SAMPLES_PER_LATENT].copy()
|
||||
self._resize(L)
|
||||
self.it.set_tensor(self.i, latents.astype(np.float32))
|
||||
self.it.invoke()
|
||||
return self.it.get_tensor(self.o).copy() # [1,2,L*4096]
|
||||
|
||||
def decode_chunked(self, latents, chunk, overlap, on_chunk=None):
|
||||
"""Stitch audio directly (stride = 4096 samples per latent token)."""
|
||||
B, C, L = latents.shape
|
||||
if L <= chunk:
|
||||
return self.decode_whole(latents)
|
||||
S = SAMPLES_PER_LATENT
|
||||
out = np.zeros((B, 2, L * S), np.float32)
|
||||
K = chunk - 2 * overlap
|
||||
assert K > 0, (chunk, overlap)
|
||||
core = 0
|
||||
n_windows = (L + K - 1) // K
|
||||
wi = 0
|
||||
while core < L:
|
||||
core_end = min(core + K, L)
|
||||
win_start = core - overlap
|
||||
win_end = win_start + chunk
|
||||
if win_start < 0:
|
||||
win_start, win_end = 0, chunk
|
||||
if win_end > L:
|
||||
win_end, win_start = L, L - chunk
|
||||
win = latents[:, :, win_start:win_end]
|
||||
y = self.decode_whole(win) # [1,2,chunk*4096]
|
||||
ks = (core - win_start) * S
|
||||
ke = (core_end - win_start) * S
|
||||
out[:, :, core * S: core_end * S] = y[:, :, ks:ke]
|
||||
wi += 1
|
||||
if on_chunk:
|
||||
on_chunk(wi, n_windows)
|
||||
core = core_end
|
||||
return out
|
||||
|
||||
|
||||
def valid_T_lat(seconds):
|
||||
"""seconds -> T_lat via natural ceil, DECODER-INDEPENDENT and identical to MLX/TRT
|
||||
(sa3_mlx / sa3_trt.resolve_T_lat = max(1, ceil(seconds*44100/4096))). So MLX == TRT ==
|
||||
TFLite pick the SAME length -> no length-driven divergence, and true ODD-length
|
||||
requests are honored. The DiT handles odd L natively; SAME-L takes any L; SAME-S
|
||||
(even T_aud=L*16 % 32) pads odd->even at encode/decode and trims."""
|
||||
return max(1, int(np.ceil(seconds * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||||
|
||||
|
||||
def _play_backend(platform: str | None = None) -> tuple[str, list[str] | None]:
|
||||
"""Pick the --play playback backend for `platform` (default: sys.platform)
|
||||
WITHOUT executing anything. Split out of main() so it can be import-tested.
|
||||
|
||||
Returns (kind, argv_prefix):
|
||||
("subprocess", ["afplay"]) macOS — argv_prefix + [wav_path] is run
|
||||
("subprocess", ["aplay", …]) Linux with alsa-utils installed
|
||||
("winsound", None) Windows — stdlib winsound.PlaySound (blocking)
|
||||
("none", None) no player available — caller prints the path
|
||||
"""
|
||||
plat = sys.platform if platform is None else platform
|
||||
if plat == "darwin":
|
||||
return ("subprocess", ["afplay"])
|
||||
if plat == "win32":
|
||||
return ("winsound", None)
|
||||
if plat.startswith("linux"):
|
||||
from shutil import which
|
||||
if which("aplay") is not None:
|
||||
return ("subprocess", ["aplay", "-q"])
|
||||
return ("none", None)
|
||||
|
||||
|
||||
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)
|
||||
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 + CFG) — baked-I/O varlen TFLite / CPU",
|
||||
allow_abbrev=False, # --pr/--di/--de became ambiguous once *-precision flags landed
|
||||
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] [--apg A]\n"))
|
||||
# Inputs
|
||||
ap.add_argument("--prompt", default=None,
|
||||
help="Text prompt. Empty string = unconditional. If omitted, asked via stdin.")
|
||||
ap.add_argument("--negative-prompt", default=None,
|
||||
help="Negative prompt for CFG's uncond branch. Ignored when --cfg=1.0. "
|
||||
"When unset and --cfg≠1.0, the uncond branch uses all-zero T5 hidden+mask "
|
||||
"(→ learned padding embeddings in-graph).")
|
||||
ap.add_argument("--init-audio", default=None,
|
||||
help="WAV to start from. With --init-noise-level → audio-to-audio; with "
|
||||
"--inpaint-range → inpainting. Encoder loaded automatically; audio is "
|
||||
"trimmed/zero-padded to --seconds. Any format via ffmpeg fallback.")
|
||||
ap.add_argument("--inpaint-range", default=None,
|
||||
help="Inpaint time range 'START,END' in seconds (e.g. '5,10'). Requires "
|
||||
"--init-audio. Regenerates the masked span, preserves the rest (paste-back).")
|
||||
# Models
|
||||
ap.add_argument("--dit", choices=["sm-music", "sm-sfx", "medium"], default=None,
|
||||
help="DiT model (names match sa3_mlx / sa3_trt). If omitted, prompts "
|
||||
"interactively with an arrow-key picker.")
|
||||
ap.add_argument("--decoder", choices=["same-s", "same-l"], default=None,
|
||||
help="Audio decoder. Default: same-s for sm-music/sm-sfx, same-l for medium. "
|
||||
"If omitted, prompts interactively with an arrow-key picker.")
|
||||
ap.add_argument("--precision", choices=list(PRECISIONS), default="fp32",
|
||||
help="DiT + decoder precision, wXaY = weight/activation bits (same flag "
|
||||
"as the TensorRT CLI): 'fp32' (default — reference; on CPU also the "
|
||||
"fast choice) | 'w16a32' (fp16 weights, half size, ≈lossless, 1.5-3× "
|
||||
"slower) | 'w8a32' (GPTQ int8 weights, ¼ size at fp32 speed) | "
|
||||
"'w8a8-dyn' (dynamic int8, fastest, lower quality). "
|
||||
"Applies to the DiT, codec decoder, and (for a2a/inpaint) the "
|
||||
"encoder; T5Gemma is single-precision. Per-component overrides below.")
|
||||
ap.add_argument("--dit-precision", choices=list(PRECISIONS), default=None,
|
||||
help="Override --precision for the DiT only (e.g. a quantized DiT "
|
||||
"with an fp32 codec).")
|
||||
ap.add_argument("--decoder-precision", choices=list(PRECISIONS), default=None,
|
||||
help="Override --precision for the SAME codec only. The codec runs "
|
||||
"once on a fixed latent, so its precision maps directly to "
|
||||
"audio quality (w8a32 is transparent at 40-46 dB).")
|
||||
ap.add_argument("--encoder-precision", choices=list(PRECISIONS), default=None,
|
||||
help="Override --precision for the SAME encoder (audio-to-audio / "
|
||||
"inpainting only). Latent PSNR vs fp32 (same-s/same-l): "
|
||||
"w16a32 ≈lossless (66/71 dB), w8a32 (GPTQ) 36/46 dB, "
|
||||
"w8a8-dyn 24/30 dB.")
|
||||
# LoRA (merged into the DiT weights at load; mirrors the MLX CLI's --lora)
|
||||
ap.add_argument("--lora", action="append", nargs="+", default=None,
|
||||
metavar=("ADAPTER", "strength=S"),
|
||||
help="A LoRA adapter to merge into the DiT — repeat the flag to stack "
|
||||
"several. Each --lora takes the adapter path followed by an optional "
|
||||
"strength=S (default --lora-strength). Example: "
|
||||
"--lora plini.safetensors strength=0.8 . The adapter is a .safetensors "
|
||||
"file (SA3-native train_lora.py / underfit output) or a PEFT adapter "
|
||||
"directory; ONLY .safetensors is accepted. The merged weights are "
|
||||
"written into a cached copy of the DiT (models/tflite/lora_cache/) and "
|
||||
"reused on repeat runs. Requires --dit-precision fp32 or w16a32 "
|
||||
"(quantized-int8 DiTs can't be LoRA-merged). NOTE: per-step gating "
|
||||
"(steps=) is MLX-only — the frozen TFLite graph merges once at load.")
|
||||
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 length. T_lat = ceil(seconds*44100/4096) (natural ceil, decoder-"
|
||||
"independent, matches MLX/TRT). Final WAV trimmed to exactly --seconds.")
|
||||
ap.add_argument("--steps", type=int, default=8,
|
||||
help="Pingpong sampling steps (≥1). rf_denoiser is distilled for 8 (default).")
|
||||
ap.add_argument("--seed", type=int, default=None,
|
||||
help="Random seed. If omitted, a random seed is chosen and printed at the end.")
|
||||
ap.add_argument("--init-noise-level", type=float, default=1.0,
|
||||
help="σmax — schedule's starting noise level. With --init-audio: 0.4-0.8 varies, "
|
||||
"1.0 = full regen (init ignored). Min %.2f (model NaNs at t≈0)." % MIN_SIGMA)
|
||||
ap.add_argument("--cfg", type=float, default=1.0,
|
||||
help="Classifier-Free Guidance scale. 1.0 = off (single pass). >1 toward prompt, "
|
||||
"<1 toward uncond. Any value ≠1.0 runs cond+uncond each step (see --cfg-batched).")
|
||||
ap.add_argument("--apg", type=float, default=1.0,
|
||||
help="Adaptive Projected Guidance [0..1], only when --cfg≠1.0. 1.0 = full APG "
|
||||
"(orthogonal projection), 0.0 = vanilla CFG. rf_denoiser default 1.0.")
|
||||
ap.add_argument("--cfg-batched", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="When --cfg≠1.0, run cond+uncond as ONE batch=2 invoke on the variable-batch "
|
||||
"DiT (default; ~7-29%% faster on Apple-Silicon AMX). --no-cfg-batched forces the "
|
||||
"sequential batch=1 dual-pass (like TensorRT). Bit-identical at "
|
||||
"fp32/w8a32; ~80 dB at w16a32; w8a8-dyn diverges by design (batch-"
|
||||
"shared activation scales) — use --no-cfg-batched to reproduce "
|
||||
"sequential baselines with w8a8-dyn.")
|
||||
# Runtime / output
|
||||
ap.add_argument("--threads", type=int, default=8, help="XNNPACK CPU threads.")
|
||||
ap.add_argument("--free-models", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="Free each model after its last use to lower peak RAM (default on).")
|
||||
ap.add_argument("--out", "-o", default=None,
|
||||
help="Output WAV path. Relative → output/<file>; absolute → as-is. "
|
||||
"If omitted, auto-named from the prompt + seed.")
|
||||
ap.add_argument("--play", action="store_true",
|
||||
help="Play the WAV after writing (blocking): `afplay` on macOS, "
|
||||
"winsound on Windows, `aplay` on Linux (prints the path if "
|
||||
"no player is available).")
|
||||
args = ap.parse_args()
|
||||
if args.steps < 1:
|
||||
ap.error(f"--steps must be ≥ 1 (got {args.steps})")
|
||||
# per-component overrides fall back to the shared --precision
|
||||
args.dit_precision = args.dit_precision or args.precision
|
||||
args.decoder_precision = args.decoder_precision or args.precision
|
||||
args.encoder_precision = args.encoder_precision or args.precision
|
||||
|
||||
# Parse --lora specs up front (fail fast, before any model work).
|
||||
args.lora_specs = None
|
||||
if args.lora:
|
||||
from lora_core 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))
|
||||
|
||||
# Interactive fills (match MLX/TRT).
|
||||
args = prompt_user_if_missing(args)
|
||||
if args.prompt is None:
|
||||
args.prompt = input("Prompt: ").strip()
|
||||
if args.seed is None:
|
||||
args.seed = random.randint(0, 2**31 - 1)
|
||||
|
||||
dec = args.decoder or DEFAULT_DECODER[args.dit]
|
||||
T_lat = valid_T_lat(args.seconds)
|
||||
target_dur = T_lat * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||||
|
||||
# Resolve output path — auto-name from prompt+seed when --out is not given.
|
||||
# Relative paths land in output/ (auto-created); absolute paths honored as-is.
|
||||
# A relative path that already starts with "output/" is taken as-is (relative to
|
||||
# cwd) so `--out output/foo.wav` doesn't become output/output/foo.wav.
|
||||
if args.out is None:
|
||||
slug = re.sub(r'[^a-z0-9]+', '_', args.prompt.lower()).strip('_')[:48] or "out"
|
||||
# non-fp32 runs get a precision suffix so same-prompt/seed A/B runs across
|
||||
# --precision values don't overwrite each other
|
||||
if args.dit_precision == args.decoder_precision:
|
||||
prec_tag = "" if args.precision == "fp32" else f"_{args.precision}"
|
||||
else:
|
||||
prec_tag = f"_dit-{args.dit_precision}_dec-{args.decoder_precision}"
|
||||
lora_tag = ""
|
||||
if args.lora_specs:
|
||||
names = "-".join(re.sub(r'[^a-z0-9]+', '', Path(s["path"]).stem.lower())[:12]
|
||||
for s in args.lora_specs)
|
||||
lora_tag = f"_lora-{names}"
|
||||
args.out = f"{slug}{prec_tag}{lora_tag}_{args.seed}.wav"
|
||||
out_path = Path(args.out)
|
||||
if not out_path.is_absolute() and out_path.parts[:1] != ("output",):
|
||||
out_path = REPO / "output" / out_path
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.out = str(out_path)
|
||||
|
||||
# Inpaint validation → latent range.
|
||||
inpaint_range = None
|
||||
inp_start_sec = inp_end_sec = 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)
|
||||
|
||||
sigma_max = float(args.init_noise_level)
|
||||
if sigma_max < MIN_SIGMA:
|
||||
sys.exit(f"error: --init-noise-level={sigma_max} too low (min {MIN_SIGMA}; model NaNs at t≈0)")
|
||||
mode = ("inpaint" if inpaint_range else
|
||||
"audio-to-audio" if args.init_audio else "text-to-audio")
|
||||
|
||||
# ── Preflight: download any missing model files BEFORE the banner/wall-clock. ──
|
||||
_preflight_download(args, dec)
|
||||
|
||||
# ── Banner ──
|
||||
print()
|
||||
banner(f"SA3 → TFLite/CPU {mode}")
|
||||
k = lambda s: dim(f"{s:>11}")
|
||||
print(f" {k('prompt')} {bold(repr(args.prompt))}")
|
||||
if args.negative_prompt:
|
||||
suffix = "" if args.cfg != 1.0 else dim(" (ignored: --cfg=1.0)")
|
||||
print(f" {k('neg prompt')} {bold(repr(args.negative_prompt))}{suffix}")
|
||||
line = f" {k('dit')} {magenta(args.dit)} {k('decoder')} {magenta(dec)}"
|
||||
if args.init_audio:
|
||||
enc_note = "" if args.encoder_precision == args.dit_precision else f" ({args.encoder_precision})"
|
||||
line += f" {k('encoder')} {magenta(dec + enc_note)}"
|
||||
prec_disp = (args.dit_precision if args.dit_precision == args.decoder_precision
|
||||
else f"dit {args.dit_precision} · codec {args.decoder_precision}")
|
||||
line += f" {k('precision')} {magenta(prec_disp)} {k('threads')} {args.threads}"
|
||||
print(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})')}")
|
||||
cfg_line = f" {k('σmax')} {bold(f'{sigma_max:.2f}')} {k('cfg')} {args.cfg}"
|
||||
if args.cfg != 1.0:
|
||||
cfg_line += dim(f" (apg={args.apg}, {'batched' if args.cfg_batched else 'sequential'} CFG)")
|
||||
print(cfg_line)
|
||||
if args.lora_specs:
|
||||
lora_disp = ", ".join(
|
||||
os.path.basename(s["path"].rstrip("/"))
|
||||
+ (f"×{s['strength']:g}" if s["strength"] != 1.0 else "")
|
||||
for s in args.lora_specs)
|
||||
print(f" {k('lora')} {magenta(lora_disp)}")
|
||||
print(f" {k('seconds')} {args.seconds}s {k('steps')} {args.steps} {k('seed')} {args.seed}")
|
||||
print(f" {k('T_lat')} {T_lat} {dim(f'({target_dur:.2f}s → trimmed to {args.seconds:.2f}s)')}")
|
||||
print()
|
||||
|
||||
# Stage numbering (extra stage when encoding init audio).
|
||||
N = 3 + (1 if args.init_audio else 0)
|
||||
TAG = {"t5": f"[1/{N}]"}
|
||||
if args.init_audio:
|
||||
TAG.update(enc=f"[2/{N}]", dit=f"[3/{N}]", dec=f"[4/{N}]")
|
||||
else:
|
||||
TAG.update(dit=f"[2/{N}]", dec=f"[3/{N}]")
|
||||
|
||||
t_wall = time.perf_counter()
|
||||
|
||||
# ── T5Gemma (tokenize + encode; + negative prompt for CFG) ──
|
||||
stage(TAG["t5"], "T5Gemma (tokenize + encode)")
|
||||
t0 = time.perf_counter()
|
||||
tok = P.Tokenizer()
|
||||
ids, mask = tok(args.prompt) # (1,256) int32 each
|
||||
t5 = P.T5GemmaTFLite(ensure_local(T5_REL), args.threads)
|
||||
t5_hidden = t5(ids, mask) # (1,256,768) fp32
|
||||
null_h = null_m = None
|
||||
if args.cfg != 1.0:
|
||||
if args.negative_prompt:
|
||||
n_ids, n_mask = tok(args.negative_prompt)
|
||||
null_h = t5(n_ids, n_mask)
|
||||
null_m = n_mask.astype(np.float32)
|
||||
else:
|
||||
# All-zero hidden+mask → in-graph conditioner emits learned padding embeds
|
||||
# for every position (the standard unconditional branch). No extra T5 pass.
|
||||
null_h = np.zeros((1, COND_TOKENS, COND_DIM), np.float32)
|
||||
null_m = np.zeros((1, COND_TOKENS), np.float32)
|
||||
t5_ms = (time.perf_counter() - t0) * 1000
|
||||
stage(TAG["t5"], "T5Gemma (tokenize + encode)", t5_ms)
|
||||
sub(f"t5_hidden {t5_hidden.shape} mask sum={int(mask.sum())}"
|
||||
+ (f" neg ({'prompt' if args.negative_prompt else 'zeros'})" if null_h is not None else ""))
|
||||
if args.free_models:
|
||||
del t5; _free()
|
||||
|
||||
# ── (audio-to-audio / inpaint) Encode init audio → init_latents ──
|
||||
init_latents = None
|
||||
if args.init_audio:
|
||||
stage(TAG["enc"], f"Encode init audio → latents ({dec})")
|
||||
t0 = time.perf_counter()
|
||||
# SAME-S encoder needs even L; round the encode grid up, trim latents to T_lat.
|
||||
enc_L = T_lat + 1 if (dec == "same-s" and T_lat % 2 != 0) else T_lat
|
||||
target_samples = enc_L * SAMPLES_PER_LATENT
|
||||
audio_in = read_wav(args.init_audio) # (2, T_in)
|
||||
if audio_in.shape[-1] >= target_samples:
|
||||
audio_in = audio_in[:, :target_samples]
|
||||
init_action = f"trimmed to {target_samples} samples"
|
||||
else:
|
||||
pad = target_samples - audio_in.shape[-1]
|
||||
audio_in = np.pad(audio_in, ((0, 0), (0, pad)))
|
||||
init_action = f"zero-padded by {pad} samples"
|
||||
audio_in = audio_in[None] # (1,2,target_samples)
|
||||
enc = BakedEncoder(ensure_local(enc_rel(dec, args.encoder_precision)), args.threads, needs_even=(dec == "same-s"))
|
||||
init_latents = enc.encode(audio_in, T_lat) # (1,256,T_lat)
|
||||
enc_ms = (time.perf_counter() - t0) * 1000
|
||||
stage(TAG["enc"], f"Encode init audio → latents ({dec})", enc_ms)
|
||||
sub(f"{init_action} latents {init_latents.shape}")
|
||||
if args.free_models:
|
||||
del enc; _free()
|
||||
|
||||
# ── Build inpaint local_add_cond + paste-back, and initial noise ──
|
||||
local_add_cond = None
|
||||
paste_back = None
|
||||
if inpaint_range is not None:
|
||||
s0, s1 = inpaint_range
|
||||
keep = np.ones((1, 1, T_lat), np.float32); keep[:, :, s0:s1] = 0.0 # 1=keep, 0=regen
|
||||
masked = init_latents.astype(np.float32) * keep
|
||||
local_add_cond = np.concatenate([keep, masked], axis=1) # (1,257,T_lat), TRT layout
|
||||
paste_back = (init_latents.astype(np.float32), keep)
|
||||
|
||||
x0, step_noise = P.make_noise(T_lat, args.steps, args.seed) # x0 = pure noise
|
||||
if init_latents is not None and inpaint_range is None:
|
||||
# rf_denoiser init mix (linear interp): noise = init*(1-σmax) + pure*σmax
|
||||
x0 = init_latents.astype(np.float32) * (1.0 - sigma_max) + x0 * sigma_max
|
||||
|
||||
# ── DiT load + pingpong sample ──
|
||||
stage(TAG["dit"], f"DiT — load + sample ({args.steps} steps, σmax={sigma_max:.2f})")
|
||||
t0 = time.perf_counter()
|
||||
cfg_note = (("CFG batched (1× batch=2 invoke/step)" if args.cfg_batched
|
||||
else "CFG sequential (2× batch=1 invokes/step)") if args.cfg != 1.0 else "")
|
||||
print(f" {dim('loading baked DiT ' + args.dit + ' ...')}", flush=True)
|
||||
dit_path = ensure_local(dit_rel(args.dit, args.dit_precision))
|
||||
if args.lora_specs:
|
||||
from lora_patch import get_patched_dit
|
||||
from lora_core import LoraError
|
||||
try:
|
||||
dit_path = get_patched_dit(dit_path, args.lora_specs, family=args.dit,
|
||||
precision=args.dit_precision, log=sub)
|
||||
except LoraError as e:
|
||||
sys.exit(f"error: {e}")
|
||||
backend = BakedDiT(dit_path, T_lat, t5_hidden, mask.astype(np.float32),
|
||||
args.seconds, args.threads, cfg=args.cfg, apg=args.apg,
|
||||
null_hidden=null_h, null_mask=null_m, local_add_cond=local_add_cond,
|
||||
batched=args.cfg_batched)
|
||||
load_ms = (time.perf_counter() - t0) * 1000
|
||||
sub(f"load {load_ms/1000:.1f}s" + (f" {cfg_note}" if cfg_note else ""))
|
||||
|
||||
sig = P.build_pingpong_schedule(args.steps, sigma_max)
|
||||
sched_str = " · ".join(f"{float(x):.3f}" for x in sig)
|
||||
sub(f"schedule {sched_str}")
|
||||
|
||||
step_prev = [time.perf_counter()]
|
||||
def on_step(i, total):
|
||||
now = time.perf_counter(); el = (now - step_prev[0]) * 1000; step_prev[0] = now
|
||||
if _USE_COLOR:
|
||||
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'{el:.0f} ms')}")
|
||||
sys.stdout.flush()
|
||||
else:
|
||||
print(f" sampling step {i}/{total} {el:.0f} ms", flush=True)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
latents = P.sample(backend, x0, step_noise, sig, None, None,
|
||||
on_step=on_step, paste_back=paste_back)
|
||||
samp_ms = (time.perf_counter() - t0) * 1000
|
||||
if _USE_COLOR:
|
||||
sys.stdout.write("\r\x1b[K")
|
||||
stage(TAG["dit"], f"DiT — load + sample ({args.steps} steps, σmax={sigma_max:.2f})",
|
||||
load_ms + samp_ms)
|
||||
sub(f"sample {samp_ms:.0f} ms ({samp_ms/max(args.steps,1):.0f} ms/step, "
|
||||
f"{backend.n_fwd} forwards) latents {latents.shape}")
|
||||
if args.free_models:
|
||||
del backend; _free()
|
||||
|
||||
# ── Decode (audio-out) + WAV ──
|
||||
stage(TAG["dec"], f"Decoder ({dec}, audio-out) + WAV")
|
||||
t0 = time.perf_counter()
|
||||
print(f" {dim('loading baked decoder ' + dec + ' ...')}", flush=True)
|
||||
decoder = BakedDecoder(ensure_local(dec_rel(dec, args.decoder_precision)), args.threads, needs_even=(dec == "same-s"))
|
||||
load2_ms = (time.perf_counter() - t0) * 1000
|
||||
sub(f"load {load2_ms:.0f} ms")
|
||||
|
||||
t0 = time.perf_counter()
|
||||
if dec == "same-l" and T_lat > SAMEL_CHUNK:
|
||||
print(f" {dim(f'SAME-L chunked decode (chunk={SAMEL_CHUNK}, ovl={SAMEL_OVERLAP}) ...')}", flush=True)
|
||||
def on_chunk(i, n):
|
||||
print(f" {dim(f'decode chunk {i}/{n}')}", flush=True)
|
||||
audio = decoder.decode_chunked(latents, SAMEL_CHUNK, SAMEL_OVERLAP, on_chunk=on_chunk)
|
||||
dmode = f"chunked (chunk={SAMEL_CHUNK}, ovl={SAMEL_OVERLAP})"
|
||||
else:
|
||||
print(f" {dim('whole decode ...')}", flush=True)
|
||||
audio = decoder.decode_whole(latents)
|
||||
dmode = "whole"
|
||||
dec_ms = (time.perf_counter() - t0) * 1000
|
||||
|
||||
audio_np = audio[0] # (2, L*4096)
|
||||
req = int(round(args.seconds * SAMPLE_RATE))
|
||||
if audio_np.shape[-1] > req:
|
||||
audio_np = audio_np[:, :req]
|
||||
P.save_wav(args.out, audio_np)
|
||||
stage(TAG["dec"], f"Decoder ({dec}, audio-out) + WAV", load2_ms + dec_ms)
|
||||
peak = float(np.abs(audio_np).max()); rms = float(np.sqrt((audio_np**2).mean()))
|
||||
sub(f"decode {dmode} {dec_ms:.0f} ms audio {audio_np.shape} peak {peak:.3f} rms {rms:.3f}")
|
||||
if args.free_models:
|
||||
del decoder; _free()
|
||||
|
||||
total = time.perf_counter() - t_wall
|
||||
dur = audio_np.shape[-1] / SAMPLE_RATE
|
||||
print()
|
||||
rule()
|
||||
print(f" {bold(green('done'))} {bold(f'{total:.2f}s')} wall → {dur:.1f}s audio → "
|
||||
f"{bold(yellow(f'{dur/total:.2f}× realtime'))} {dim(f'seed {args.seed}')}")
|
||||
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:
|
||||
kind, argv = _play_backend()
|
||||
try:
|
||||
if kind == "winsound":
|
||||
print(f" {bold('▶ playing')} {args.out} {dim('(Ctrl-C to stop)')}")
|
||||
import winsound
|
||||
# SND_ASYNC + an interruptible sleep instead of the synchronous
|
||||
# call: PlaySound(..., SND_FILENAME) blocks inside native code,
|
||||
# so Python can't deliver Ctrl-C until playback finishes. This
|
||||
# way Ctrl-C lands in time.sleep and SND_PURGE stops the audio.
|
||||
with wave.open(args.out) as _w:
|
||||
_dur = _w.getnframes() / _w.getframerate()
|
||||
winsound.PlaySound(args.out,
|
||||
winsound.SND_FILENAME | winsound.SND_ASYNC)
|
||||
try:
|
||||
time.sleep(_dur + 0.25)
|
||||
except KeyboardInterrupt:
|
||||
winsound.PlaySound(None, winsound.SND_PURGE)
|
||||
raise
|
||||
elif kind == "subprocess":
|
||||
print(f" {bold('▶ playing')} {args.out} {dim('(Ctrl-C to stop)')}")
|
||||
subprocess.run(argv + [args.out], check=False)
|
||||
else:
|
||||
print(f" {bold('▶ play')} no audio player found ('aplay' missing) — "
|
||||
f"file saved at {args.out}")
|
||||
except KeyboardInterrupt:
|
||||
print()
|
||||
|
||||
|
||||
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,301 @@
|
||||
"""Tests for the TFLite LoRA merge (lora_core.py + lora_patch.py).
|
||||
|
||||
Stdlib unittest, no torch / no mlx / no pytest (matches test_windows_compat.py).
|
||||
Two tiers:
|
||||
* CI-safe: safetensors reader, per-type merge math (independent re-derivation),
|
||||
spec parser, quantized-precision refusal — no model files needed.
|
||||
* Model-dependent (skipped unless a local DiT .tflite is present): the real
|
||||
plini adapter's 169/169 name-mapping bijection on medium; w16a32 fp16
|
||||
buffer discovery; and a synthetic FULL-target-set adapter on sm-music that
|
||||
exercises every mapping tier + the clone/patch/read-back path.
|
||||
|
||||
Run: python scripts/test_lora_patch.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.insert(0, _HERE)
|
||||
import lora_core as lc # noqa: E402
|
||||
|
||||
# Local model paths (present on a dev box with the tflite bundle installed;
|
||||
# absent in CI → the model-dependent tests skip). SA3_TFLITE_MODELS_DIR
|
||||
# overrides the location (e.g. pointing a worktree's tests at the main checkout).
|
||||
_CHECKOUT = os.environ.get(
|
||||
"SA3_TFLITE_MODELS_DIR",
|
||||
os.path.abspath(os.path.join(_HERE, "..", "models", "tflite")))
|
||||
_MODELS = {
|
||||
("sm-music", "fp32"): f"{_CHECKOUT}/sa3-sm-music/dit_fp32.tflite",
|
||||
("sm-music", "w16a32"): f"{_CHECKOUT}/sa3-sm-music/dit_w16a32.tflite",
|
||||
("medium", "fp32"): f"{_CHECKOUT}/sa3-m/dit_fp32.tflite",
|
||||
("sm-music", "w8a32"): f"{_CHECKOUT}/sa3-sm-music/dit_w8a32.tflite",
|
||||
}
|
||||
_PLINI = "/Users/cj/clod/speed-metal/scripts/lora_bench/plini-sa3-380.safetensors"
|
||||
|
||||
|
||||
def _have(key):
|
||||
p = _MODELS.get(key)
|
||||
return p and os.path.exists(p)
|
||||
|
||||
|
||||
def _save_native_safetensors(path, layers, cfg):
|
||||
"""Write a minimal SA3-native adapter: {layer.parametrizations.weight.0.param}."""
|
||||
tensors, blob, hdr = {}, bytearray(), {}
|
||||
for layer, params in layers.items():
|
||||
for pname, arr in params.items():
|
||||
tensors[f"{layer}.parametrizations.weight.0.{pname}"] = np.ascontiguousarray(arr, np.float32)
|
||||
for k, v in tensors.items():
|
||||
off = len(blob)
|
||||
blob += v.tobytes()
|
||||
hdr[k] = {"dtype": "F32", "shape": list(v.shape), "data_offsets": [off, len(blob)]}
|
||||
hdr["__metadata__"] = {"lora_config": json.dumps(cfg)}
|
||||
hbytes = json.dumps(hdr).encode()
|
||||
with open(path, "wb") as f:
|
||||
f.write(struct.pack("<Q", len(hbytes)))
|
||||
f.write(hbytes)
|
||||
f.write(blob)
|
||||
|
||||
|
||||
# ── CI-safe unit tests ─────────────────────────────────────────────────────────
|
||||
|
||||
class TestSafetensorsReader(unittest.TestCase):
|
||||
def test_roundtrip_fp32_and_metadata(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = os.path.join(d, "a.safetensors")
|
||||
A = np.random.default_rng(0).standard_normal((4, 8)).astype(np.float32)
|
||||
_save_native_safetensors(p, {"foo": {"lora_A": A, "lora_B": np.zeros((6, 4), np.float32)}},
|
||||
{"adapter_type": "lora", "rank": 4, "alpha": 4})
|
||||
t, meta = lc._load_safetensors(p)
|
||||
self.assertIn("foo.parametrizations.weight.0.lora_A", t)
|
||||
np.testing.assert_array_equal(t["foo.parametrizations.weight.0.lora_A"], A)
|
||||
self.assertEqual(json.loads(meta["lora_config"])["adapter_type"], "lora")
|
||||
|
||||
def test_pickle_refused(self):
|
||||
for bad in ("evil.ckpt", "evil.pt", "evil.bin", "x.npz"):
|
||||
with self.assertRaises(lc.LoraError):
|
||||
lc._load_safetensors(bad)
|
||||
|
||||
|
||||
class TestMergeMath(unittest.TestCase):
|
||||
"""Independent re-derivation of the four core types (guards the copy)."""
|
||||
|
||||
def setUp(self):
|
||||
self.rng = np.random.default_rng(1)
|
||||
self.W0 = self.rng.standard_normal((6, 8)).astype(np.float32)
|
||||
|
||||
def _p(self, t, r=3):
|
||||
p = {}
|
||||
if not t.endswith("-xs"):
|
||||
p["lora_A"] = self.rng.standard_normal((r, 8)).astype(np.float32)
|
||||
p["lora_B"] = self.rng.standard_normal((6, r)).astype(np.float32)
|
||||
else:
|
||||
p["M_xs"] = self.rng.standard_normal((r, r)).astype(np.float32)
|
||||
if "magnitude" in lc._PARAMS_FOR[t]:
|
||||
p["magnitude"] = np.linalg.norm(self.W0, axis=1 if "rows" in t else 0).astype(np.float32)
|
||||
if "magnitude_r" in lc._PARAMS_FOR[t]:
|
||||
p["magnitude_r"] = np.linalg.norm(self.W0, axis=1).astype(np.float32)
|
||||
p["magnitude_c"] = np.linalg.norm(self.W0, axis=0).astype(np.float32)
|
||||
return p
|
||||
|
||||
def test_lora(self):
|
||||
p = self._p("lora")
|
||||
ref = self.W0 + 2.0 * (p["lora_B"] @ p["lora_A"])
|
||||
np.testing.assert_allclose(lc.merged_weight(self.W0, p, "lora", 2.0), ref, atol=1e-5)
|
||||
|
||||
def test_dora_rows(self):
|
||||
p = self._p("dora-rows")
|
||||
V = self.W0 + 1.0 * (p["lora_B"] @ p["lora_A"])
|
||||
ref = V / (np.linalg.norm(V, axis=1, keepdims=True) + 1e-12) * p["magnitude"].reshape(-1, 1)
|
||||
np.testing.assert_allclose(lc.merged_weight(self.W0, p, "dora-rows", 1.0), ref, atol=1e-5)
|
||||
|
||||
def test_dora_cols(self):
|
||||
p = self._p("dora-cols")
|
||||
V = self.W0 + 1.0 * (p["lora_B"] @ p["lora_A"])
|
||||
ref = V / (np.linalg.norm(V, axis=0, keepdims=True) + 1e-12) * p["magnitude"].reshape(1, -1)
|
||||
np.testing.assert_allclose(lc.merged_weight(self.W0, p, "dora-cols", 1.0), ref, atol=1e-5)
|
||||
|
||||
def test_all_types_finite_and_shaped(self):
|
||||
for t in lc._PARAMS_FOR:
|
||||
p = self._p(t)
|
||||
lc.check_shapes("x", self.W0, p, t)
|
||||
W = lc.merged_weight(self.W0, p, t, 1.0)
|
||||
self.assertEqual(W.shape, self.W0.shape)
|
||||
self.assertTrue(np.isfinite(W).all(), t)
|
||||
|
||||
def test_shape_mismatch_raises(self):
|
||||
p = {"lora_A": np.zeros((3, 99), np.float32), "lora_B": np.zeros((6, 3), np.float32)}
|
||||
with self.assertRaises(lc.LoraError):
|
||||
lc.check_shapes("x", self.W0, p, "lora")
|
||||
|
||||
|
||||
class TestSpecParser(unittest.TestCase):
|
||||
def test_basic(self):
|
||||
self.assertEqual(lc.parse_lora_spec(["a.safetensors"])["strength"], 1.0)
|
||||
self.assertEqual(lc.parse_lora_spec(["a.safetensors", "strength=0.7"])["strength"], 0.7)
|
||||
|
||||
def test_steps_rejected_points_to_mlx(self):
|
||||
with self.assertRaises(lc.LoraError) as cm:
|
||||
lc.parse_lora_spec(["a.safetensors", "steps=2-8"])
|
||||
self.assertIn("MLX", str(cm.exception))
|
||||
|
||||
def test_bad_tokens(self):
|
||||
for bad in (["a", "b.st"], ["a", "strength=x"], []):
|
||||
with self.assertRaises(lc.LoraError):
|
||||
lc.parse_lora_spec(bad)
|
||||
|
||||
|
||||
class TestQuantizedRefusal(unittest.TestCase):
|
||||
def test_precision_refused_before_touching_files(self):
|
||||
import lora_patch as lp
|
||||
for prec in ("w8a32", "w8a8-dyn", "w4a32"):
|
||||
with self.assertRaises(lc.LoraError) as cm:
|
||||
lp.get_patched_dit("/nonexistent.tflite", [{"path": "x", "strength": 1.0}],
|
||||
family="medium", precision=prec, log=lambda _m: None)
|
||||
self.assertIn(prec, str(cm.exception))
|
||||
|
||||
|
||||
# ── model-dependent tests ────────────────────────────────────────────────────
|
||||
|
||||
@unittest.skipUnless(_have(("medium", "fp32")) and os.path.exists(_PLINI),
|
||||
"medium fp32 DiT + plini adapter not present")
|
||||
class TestRealAdapterMapping(unittest.TestCase):
|
||||
def test_plini_169_bijection_on_medium(self):
|
||||
import mmap
|
||||
import lora_patch as lp
|
||||
f = open(_MODELS[("medium", "fp32")], "rb")
|
||||
buf = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
|
||||
try:
|
||||
mapper = lp._Mapper(list(lp._fc_weight_buffers(buf)))
|
||||
_at, _sc, layers = lc.parse_adapter(_PLINI)
|
||||
offs = set()
|
||||
for layer, p in layers.items():
|
||||
want = (p["lora_B"].shape[0], p["lora_A"].shape[1])
|
||||
rec = mapper.resolve(layer, want)
|
||||
self.assertEqual(rec["shape"], want, layer)
|
||||
offs.add(rec["off"])
|
||||
self.assertEqual(len(layers), 169)
|
||||
self.assertEqual(len(offs), 169, "buffer offsets must be unique (no collisions)")
|
||||
finally:
|
||||
del mapper
|
||||
buf.close()
|
||||
f.close()
|
||||
|
||||
|
||||
@unittest.skipUnless(_have(("sm-music", "w16a32")), "sm-music w16a32 DiT not present")
|
||||
class TestW16A32Discovery(unittest.TestCase):
|
||||
def test_all_weights_fp16_via_dequant(self):
|
||||
import mmap
|
||||
import lora_patch as lp
|
||||
f = open(_MODELS[("sm-music", "w16a32")], "rb")
|
||||
buf = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
|
||||
try:
|
||||
fcs = list(lp._fc_weight_buffers(buf))
|
||||
self.assertGreater(len(fcs), 100)
|
||||
self.assertTrue(all(np.dtype(r["np_dtype"]) == np.float16 for r in fcs))
|
||||
self.assertTrue(any(r["shape"] == (768, 256) for r in fcs), "seconds cond present")
|
||||
finally:
|
||||
del fcs
|
||||
buf.close()
|
||||
f.close()
|
||||
|
||||
|
||||
@unittest.skipUnless(_have(("sm-music", "fp32")), "sm-music fp32 DiT not present")
|
||||
class TestFullSetPatchRoundTrip(unittest.TestCase):
|
||||
"""Synthetic all-target adapter (every mapping tier) → clone + patch +
|
||||
byte-exact read-back on the smaller sm-music fp32 model."""
|
||||
|
||||
def test_map_patch_readback(self):
|
||||
import mmap
|
||||
import lora_patch as lp
|
||||
base = _MODELS[("sm-music", "fp32")]
|
||||
f = open(base, "rb")
|
||||
buf = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
|
||||
fcs = list(lp._fc_weight_buffers(buf))
|
||||
mapper = lp._Mapper(fcs)
|
||||
|
||||
# Inverse-classify each FC → a checkpoint layer name, so we can build a
|
||||
# synthetic adapter that targets every tier (block linears, local_embed,
|
||||
# singleton embedders, seconds cond).
|
||||
ckpt_for = {} # off -> (ckpt_layer_name, (fan_out, fan_in))
|
||||
# tier 1
|
||||
rev_block = {("self_attn", "to_qkv", False): "self_attn.to_qkv",
|
||||
("self_attn", "to_out", False): "self_attn.to_out",
|
||||
("cross_attn", "to_q", False): "cross_attn.to_q",
|
||||
("cross_attn", "to_kv", False): "cross_attn.to_kv",
|
||||
("cross_attn", "to_out", False): "cross_attn.to_out",
|
||||
("ff", "proj", True): "ff.ff.0.proj", ("ff", "2", False): "ff.ff.2"}
|
||||
for (blk, mod, leaf, glu), rec in mapper._block.items():
|
||||
ckpt_for[rec["off"]] = (f"model.transformer.layers.{blk}.{rev_block[(mod, leaf, glu)]}", rec["shape"])
|
||||
# tier 2 (topological rank == block)
|
||||
for leaf, recs in mapper._local.items():
|
||||
for blk, rec in enumerate(recs):
|
||||
ckpt_for[rec["off"]] = (f"model.transformer.layers.{blk}.to_local_embed.seq.{leaf}", rec["shape"])
|
||||
# tier 3 singletons
|
||||
for ckpt, sub in lp._SINGLETON_CKPT.items():
|
||||
hit = [r for r in fcs if sub in r["name"]]
|
||||
if hit:
|
||||
ckpt_for[hit[0]["off"]] = (ckpt, hit[0]["shape"])
|
||||
# tier 4 seconds cond
|
||||
if mapper._seconds:
|
||||
ckpt_for[mapper._seconds[0]["off"]] = (lc.COND_SECONDS_LAYER, mapper._seconds[0]["shape"])
|
||||
|
||||
rng = np.random.default_rng(7)
|
||||
layers = {}
|
||||
for off, (name, (fo, fi)) in ckpt_for.items():
|
||||
r = 4
|
||||
layers[name] = {
|
||||
"lora_A": (rng.standard_normal((r, fi)) * 0.01).astype(np.float32),
|
||||
"lora_B": (rng.standard_normal((fo, r)) * 0.01).astype(np.float32),
|
||||
"magnitude": np.linalg.norm(
|
||||
rng.standard_normal((fo, fi)).astype(np.float32), axis=1).astype(np.float32),
|
||||
}
|
||||
buf.close()
|
||||
f.close()
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
adapter = os.path.join(d, "full.safetensors")
|
||||
_save_native_safetensors(adapter, layers,
|
||||
{"adapter_type": "dora-rows", "rank": 4, "alpha": 4})
|
||||
specs = [lc.parse_lora_spec([adapter, "strength=0.5"])]
|
||||
out = lp.get_patched_dit(base, specs, family="sm-music", precision="fp32",
|
||||
cache_dir=os.path.join(d, "cache"), log=lambda _m: None)
|
||||
self.assertTrue(out.exists())
|
||||
|
||||
# read-back: every patched buffer == fp32(W0 + 0.5*(merged-W0))
|
||||
fb = open(base, "rb"); bb = mmap.mmap(fb.fileno(), 0, access=mmap.ACCESS_READ)
|
||||
fo2 = open(out, "rb"); bo = mmap.mmap(fo2.fileno(), 0, access=mmap.ACCESS_READ)
|
||||
def owned(mm, off, size, dt):
|
||||
# copy out of the mmap so no view keeps it open at close time
|
||||
return np.array(np.frombuffer(mm, dt, count=size // np.dtype(dt).itemsize,
|
||||
offset=off))
|
||||
m2 = lp._Mapper(list(lp._fc_weight_buffers(bb)))
|
||||
checked = 0
|
||||
for name, p in layers.items():
|
||||
rec = m2.resolve(name, (p["lora_B"].shape[0], p["lora_A"].shape[1]))
|
||||
W0 = owned(bb, rec["off"], rec["size"], np.float32).reshape(rec["shape"])
|
||||
got = owned(bo, rec["off"], rec["size"], np.float32).reshape(rec["shape"])
|
||||
ref = (W0 + 0.5 * (lc.merged_weight(W0, p, "dora-rows", 1.0) - W0)).astype(np.float32)
|
||||
self.assertTrue(np.array_equal(got, ref), name)
|
||||
checked += 1
|
||||
self.assertEqual(checked, len(layers))
|
||||
# a non-target FC — any FC not in our set — must be unchanged
|
||||
untouched = [r for r in lp._fc_weight_buffers(bb) if r["off"] not in ckpt_for]
|
||||
if untouched:
|
||||
r = untouched[0]
|
||||
self.assertTrue(np.array_equal(owned(bb, r["off"], r["size"], np.uint8),
|
||||
owned(bo, r["off"], r["size"], np.uint8)),
|
||||
"untouched tensor changed")
|
||||
del m2
|
||||
bb.close(); bo.close(); fb.close(); fo2.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Cross-platform (Windows/macOS/Linux) compatibility tests for the TFLite CLI.
|
||||
|
||||
Run directly — stdlib unittest only, no pytest, NO model downloads, no network:
|
||||
|
||||
python scripts/test_windows_compat.py
|
||||
|
||||
Covers the three Windows-specific failure modes this release guards against:
|
||||
1. weights.ensure_local materializes model files without symlink privileges
|
||||
(symlink -> hardlink -> copy fallback chain);
|
||||
2. the CLI's banners/emoji print cleanly under a legacy console code page
|
||||
(the guarded sys.stdout.reconfigure(encoding="utf-8") at script top);
|
||||
3. --play dispatches to the right playback backend per platform
|
||||
(afplay / winsound / aplay) without executing anything.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
# Keep this test's own output safe under legacy Windows console code pages
|
||||
# (the subprocess tests below spawn FRESH interpreters, so this does not mask them).
|
||||
for _stream in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
_stream.reconfigure(encoding="utf-8", errors="replace")
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
SCRIPTS_DIR = Path(__file__).resolve().parent # <project>/scripts
|
||||
PROJECT_DIR = SCRIPTS_DIR.parent # <project>
|
||||
sys.path.insert(0, str(SCRIPTS_DIR)) # import weights / sa3_tflite / examples
|
||||
|
||||
import weights # noqa: E402
|
||||
|
||||
|
||||
class TestEnsureLocalLinkOrCopy(unittest.TestCase):
|
||||
"""ensure_local must materialize the target on every OS, with or without
|
||||
symlink privileges (Windows without Developer Mode raises WinError 1314)."""
|
||||
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory(prefix="sa3_citest_")
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
self.tmp = Path(self._tmp.name)
|
||||
self.fake_cache = self.tmp / "hf_cache"
|
||||
self.fake_cache.mkdir()
|
||||
# Redirect the project root so nothing is written into the repo checkout.
|
||||
self._orig_script_dir = weights.SCRIPT_DIR
|
||||
weights.SCRIPT_DIR = self.tmp / "project"
|
||||
self.addCleanup(lambda: setattr(weights, "SCRIPT_DIR", self._orig_script_dir))
|
||||
# Register a fake manifest entry + suppress the (non-ASCII) HF login tip
|
||||
# so no in-process print depends on console encoding.
|
||||
self.rel = "models/tflite/_citest/fake_model.bin"
|
||||
weights.FLAT_MANIFEST[self.rel] = "tflite/_citest/fake_model.bin"
|
||||
self.addCleanup(weights.FLAT_MANIFEST.pop, self.rel, None)
|
||||
self._orig_tip = weights._LOGIN_TIP_SHOWN
|
||||
weights._LOGIN_TIP_SHOWN = True
|
||||
self.addCleanup(lambda: setattr(weights, "_LOGIN_TIP_SHOWN", self._orig_tip))
|
||||
self.payload = b"not a real tflite model \x00\x01\x02" * 100
|
||||
|
||||
def _fake_hf_hub_download(self, repo_id, filename):
|
||||
self.assertEqual(repo_id, weights.REPO_ID)
|
||||
cached = self.fake_cache / filename.replace("/", "__")
|
||||
cached.write_bytes(self.payload)
|
||||
return str(cached)
|
||||
|
||||
def test_ensure_local_materializes_target(self):
|
||||
"""End-to-end: monkeypatched download -> target exists and is readable.
|
||||
Uses whatever link mode the OS grants (symlink where privileged,
|
||||
hardlink/copy otherwise) — all satisfy exists()/is_present()."""
|
||||
import huggingface_hub
|
||||
with mock.patch.object(huggingface_hub, "hf_hub_download",
|
||||
self._fake_hf_hub_download):
|
||||
target = weights.ensure_local(self.rel, verbose=False)
|
||||
self.assertTrue(target.exists(), f"target not materialized: {target}")
|
||||
self.assertEqual(target.read_bytes(), self.payload)
|
||||
self.assertTrue(weights.is_present(self.rel))
|
||||
|
||||
def test_fallback_chain_needs_no_symlink_privilege(self):
|
||||
"""Force symlink_to to fail (as on Windows without Developer Mode,
|
||||
WinError 1314) and verify the hardlink/copy fallback still materializes
|
||||
a readable, non-symlink target. Runs identically on every OS."""
|
||||
cached = self.fake_cache / "src.bin"
|
||||
cached.write_bytes(self.payload)
|
||||
target = self.tmp / "project" / "models" / "linked.bin"
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
boom = OSError(1314, "A required privilege is not held by the client")
|
||||
with mock.patch.object(type(target), "symlink_to", side_effect=boom):
|
||||
weights._link_or_copy(cached, target)
|
||||
self.assertTrue(target.exists())
|
||||
self.assertFalse(target.is_symlink())
|
||||
self.assertEqual(target.read_bytes(), self.payload)
|
||||
|
||||
def test_copy_fallback_when_hardlink_also_fails(self):
|
||||
"""Symlink AND hardlink refused (e.g. cross-volume) -> plain copy."""
|
||||
cached = self.fake_cache / "src2.bin"
|
||||
cached.write_bytes(self.payload)
|
||||
target = self.tmp / "project" / "models" / "copied.bin"
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
with mock.patch.object(type(target), "symlink_to",
|
||||
side_effect=OSError(1314, "no symlink privilege")), \
|
||||
mock.patch.object(weights, "_LINK_FALLBACK_SHOWN", True), \
|
||||
mock.patch("os.link", side_effect=OSError(17, "cross-device link")):
|
||||
weights._link_or_copy(cached, target)
|
||||
self.assertTrue(target.exists())
|
||||
self.assertFalse(target.is_symlink())
|
||||
self.assertEqual(target.read_bytes(), self.payload)
|
||||
|
||||
|
||||
class TestConsoleEncoding(unittest.TestCase):
|
||||
"""The banners use box-drawing + emoji. Run the scripts in FRESH interpreters
|
||||
with PYTHONUTF8/PYTHONIOENCODING stripped, so stdout is a pipe under the
|
||||
runner's native default encoding (cp1252 on Windows without the in-script
|
||||
reconfigure fix) — a UnicodeEncodeError would make these exit nonzero."""
|
||||
|
||||
@staticmethod
|
||||
def _clean_env():
|
||||
env = os.environ.copy()
|
||||
env.pop("PYTHONUTF8", None)
|
||||
env.pop("PYTHONIOENCODING", None)
|
||||
return env
|
||||
|
||||
def test_sa3_tflite_help(self):
|
||||
r = subprocess.run(
|
||||
[sys.executable, str(SCRIPTS_DIR / "sa3_tflite.py"), "--help"],
|
||||
capture_output=True, env=self._clean_env(), timeout=120)
|
||||
self.assertEqual(
|
||||
r.returncode, 0,
|
||||
f"--help failed\nstdout:\n{r.stdout.decode('utf-8', 'replace')}\n"
|
||||
f"stderr:\n{r.stderr.decode('utf-8', 'replace')}")
|
||||
out = r.stdout.decode("utf-8", "replace")
|
||||
self.assertIn("--prompt", out)
|
||||
self.assertIn("--play", out)
|
||||
|
||||
def test_examples_block(self):
|
||||
# examples.py is a module (no __main__): render the full emoji block.
|
||||
r = subprocess.run(
|
||||
[sys.executable, "-c",
|
||||
"import examples; examples.print_example_commands()"],
|
||||
capture_output=True, env=self._clean_env(), cwd=str(SCRIPTS_DIR),
|
||||
timeout=120)
|
||||
self.assertEqual(
|
||||
r.returncode, 0,
|
||||
f"examples block failed\nstdout:\n{r.stdout.decode('utf-8', 'replace')}\n"
|
||||
f"stderr:\n{r.stderr.decode('utf-8', 'replace')}")
|
||||
self.assertIn("Generate audio", r.stdout.decode("utf-8", "replace"))
|
||||
|
||||
|
||||
class TestPlayBackendDispatch(unittest.TestCase):
|
||||
"""--play platform dispatch, without playing anything."""
|
||||
|
||||
def _backend(self):
|
||||
import sa3_tflite
|
||||
return sa3_tflite._play_backend
|
||||
|
||||
def test_darwin_uses_afplay(self):
|
||||
self.assertEqual(self._backend()("darwin"), ("subprocess", ["afplay"]))
|
||||
|
||||
def test_windows_uses_winsound(self):
|
||||
kind, argv = self._backend()("win32")
|
||||
self.assertEqual(kind, "winsound")
|
||||
self.assertIsNone(argv)
|
||||
if sys.platform == "win32":
|
||||
import winsound # stdlib on Windows — the module the dispatcher relies on
|
||||
self.assertTrue(hasattr(winsound, "PlaySound"))
|
||||
|
||||
def test_linux_uses_aplay_or_prints_path(self):
|
||||
kind, argv = self._backend()("linux")
|
||||
if kind == "subprocess":
|
||||
self.assertEqual(argv[0], "aplay")
|
||||
else:
|
||||
self.assertEqual((kind, argv), ("none", None))
|
||||
|
||||
def test_current_platform_resolves(self):
|
||||
kind, argv = self._backend()()
|
||||
self.assertIn(kind, ("subprocess", "winsound", "none"))
|
||||
if kind == "subprocess":
|
||||
self.assertIsInstance(argv, list)
|
||||
else:
|
||||
self.assertIsNone(argv)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,245 @@
|
||||
"""Shared weights manifest + downloader.
|
||||
|
||||
Maps every TFLite / LiteRT model file the runtime needs to its position in the
|
||||
`stabilityai/stable-audio-3-optimized` HuggingFace repo (under `tflite/…`).
|
||||
|
||||
`install.py` calls `ensure_local` upfront for the bundles the user picks.
|
||||
`sa3_tflite.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.
|
||||
|
||||
The SentencePiece tokenizer is BUNDLED at models/tokenizer.model (the .tflite
|
||||
T5Gemma is encoder-only), so it's not in this manifest.
|
||||
"""
|
||||
|
||||
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/tflite/foo.tflite")
|
||||
# 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
|
||||
# model files (local relative path on the left, HF repo path on the right).
|
||||
# T5Gemma is in SHARED because every bundle needs it. The two small DiTs
|
||||
# share the SAME-S codec; medium uses the SAME-L codec.
|
||||
|
||||
DIT_BUNDLES: dict[str, list[tuple[str, str]]] = {
|
||||
"sm-music": [
|
||||
("models/tflite/sa3-sm-music/dit_fp32.tflite", "tflite/sa3-sm-music/dit_fp32.tflite"),
|
||||
("models/tflite/same-s/enc_fp32.tflite", "tflite/same-s/enc_fp32.tflite"),
|
||||
("models/tflite/same-s/dec_fp32.tflite", "tflite/same-s/dec_fp32.tflite"),
|
||||
],
|
||||
"sm-sfx": [
|
||||
("models/tflite/sa3-sm-sfx/dit_fp32.tflite", "tflite/sa3-sm-sfx/dit_fp32.tflite"),
|
||||
("models/tflite/same-s/enc_fp32.tflite", "tflite/same-s/enc_fp32.tflite"),
|
||||
("models/tflite/same-s/dec_fp32.tflite", "tflite/same-s/dec_fp32.tflite"),
|
||||
],
|
||||
"medium": [
|
||||
("models/tflite/sa3-m/dit_fp32.tflite", "tflite/sa3-m/dit_fp32.tflite"),
|
||||
("models/tflite/same-l/enc_fp32.tflite", "tflite/same-l/enc_fp32.tflite"),
|
||||
("models/tflite/same-l/dec_fp32.tflite", "tflite/same-l/dec_fp32.tflite"),
|
||||
],
|
||||
}
|
||||
|
||||
SHARED: list[tuple[str, str]] = [
|
||||
("models/tflite/t5gemma/encoder_fp16.tflite", "tflite/t5gemma/encoder_fp16.tflite"),
|
||||
]
|
||||
|
||||
# Human-friendly bundle sizes (for the install prompt). Exact, from HF metadata.
|
||||
BUNDLE_SIZES = {
|
||||
"sm-music": "2.3 GB (small music DiT + SAME-S codec, all fp32)",
|
||||
"sm-sfx": "2.3 GB (small sfx DiT + SAME-S codec, all fp32)",
|
||||
"medium": "9.5 GB (medium DiT + SAME-L codec, all fp32)",
|
||||
}
|
||||
# T5Gemma (shared, fp16) adds ~0.6 GB the first time any bundle is fetched.
|
||||
|
||||
# Quantized DiT + decoder variants (selected via sa3_tflite.py --precision; wXaY =
|
||||
# weight/activation bit-widths, "16" = fp16). Not part of the install bundles — they
|
||||
# lazy-download on first use. w16a32 = fp16 weights (half size, ≈lossless, slower on
|
||||
# CPU); w8a32 / w8a8-dyn = GPTQ-calibrated int8 weights.
|
||||
QUANT_PRECISIONS = ("w16a32", "w8a32", "w8a8-dyn")
|
||||
PRECISIONS = ("fp32",) + QUANT_PRECISIONS # everything --precision accepts
|
||||
DIT_SUBDIR = {"sm-music": "sa3-sm-music", "sm-sfx": "sa3-sm-sfx", "medium": "sa3-m"}
|
||||
|
||||
|
||||
def dit_rel(dit: str, precision: str = "fp32") -> str:
|
||||
"""Local rel path of a DiT model file for (family, precision)."""
|
||||
return f"models/tflite/{DIT_SUBDIR[dit]}/dit_{precision}.tflite"
|
||||
|
||||
|
||||
def dec_rel(dec: str, precision: str = "fp32") -> str:
|
||||
"""Local rel path of a SAME codec decoder file for (codec, precision)."""
|
||||
return f"models/tflite/{dec}/dec_{precision}.tflite"
|
||||
|
||||
|
||||
def enc_rel(dec: str, precision: str = "fp32") -> str:
|
||||
"""Local rel path of a SAME codec encoder file for (codec, precision).
|
||||
Encoder int8 is GPTQ-calibrated on real audio (torch-free, in-place on the
|
||||
tflite graph) — w8a32 measures 36/46 dB latent PSNR (same-s/same-l)."""
|
||||
return f"models/tflite/{dec}/enc_{precision}.tflite"
|
||||
|
||||
# Flat (local_rel_path → hf_path) lookup — used by sa3_tflite.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 _prec in QUANT_PRECISIONS:
|
||||
for _fam in DIT_SUBDIR:
|
||||
_rel = dit_rel(_fam, _prec)
|
||||
FLAT_MANIFEST[_rel] = _rel.replace("models/tflite/", "tflite/", 1)
|
||||
for _dec in ("same-s", "same-l"):
|
||||
for _rel in (dec_rel(_dec, _prec), enc_rel(_dec, _prec)):
|
||||
FLAT_MANIFEST[_rel] = _rel.replace("models/tflite/", "tflite/", 1)
|
||||
|
||||
|
||||
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
|
||||
_LINK_FALLBACK_SHOWN = False
|
||||
|
||||
|
||||
def _link_or_copy(cached: Path, target: Path) -> None:
|
||||
"""Expose the HF-cached file at `target`: symlink → hardlink → copy.
|
||||
|
||||
A symlink keeps the HF cache canonical (one copy on disk) and is used
|
||||
wherever possible. On Windows, creating symlinks requires Developer Mode
|
||||
or admin rights (WinError 1314), so fall back to a hardlink (privilege-
|
||||
free and zero-copy, but same-volume only), then to a plain copy (works
|
||||
everywhere, at the cost of a second copy on disk). Unlike a symlink, a
|
||||
hardlink/copy won't track future HF-cache updates — acceptable, because
|
||||
downloads are content-addressed: a changed file resolves to a new cache
|
||||
path, and `ensure_local` only ever links each target once.
|
||||
"""
|
||||
global _LINK_FALLBACK_SHOWN
|
||||
try:
|
||||
target.symlink_to(cached)
|
||||
return
|
||||
except OSError:
|
||||
pass
|
||||
import os
|
||||
import shutil
|
||||
try:
|
||||
os.link(cached, target) # hardlink: no privilege needed, same volume only
|
||||
mode = "hardlink"
|
||||
except OSError:
|
||||
shutil.copy2(cached, target)
|
||||
mode = "copy"
|
||||
if not _LINK_FALLBACK_SHOWN:
|
||||
_LINK_FALLBACK_SHOWN = True
|
||||
print(f" (symlinks unavailable on this system — materialized as a {mode} instead)")
|
||||
|
||||
|
||||
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 model 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)
|
||||
# Expose the file at the project-relative path the runtime expects.
|
||||
# Symlink when possible (keeps the HF cache canonical — one copy on disk);
|
||||
# falls back to hardlink → copy on systems without symlink privileges
|
||||
# (Windows without Developer Mode). The `exists() or is_symlink()`
|
||||
# early-exits above and in is_present() are satisfied by every mode.
|
||||
if target.is_symlink():
|
||||
target.unlink()
|
||||
_link_or_copy(Path(cached), target)
|
||||
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