初始提交：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,2 @@
|
||||
# Batch files must keep CRLF line endings on every checkout (cmd.exe quirk-proofing).
|
||||
*.bat text eol=crlf
|
||||
@@ -0,0 +1,21 @@
|
||||
# Model weights — pulled from HuggingFace, never committed here.
|
||||
# (The SentencePiece tokenizer at models/tokenizer.model IS committed — it's
|
||||
# small and the .tflite T5Gemma is encoder-only.)
|
||||
*.tflite
|
||||
|
||||
# Generated audio — anywhere on the tree, including output/.
|
||||
*.wav
|
||||
*.mp3
|
||||
output/
|
||||
|
||||
# LoRA adapter library (user .safetensors, scanned by the gradio dropdown) and
|
||||
# the LoRA-merged DiT cache (patched multi-GB copies, regenerated on demand).
|
||||
loras/
|
||||
lora_cache/
|
||||
|
||||
# Python virtualenv created by install.sh
|
||||
.venv/
|
||||
|
||||
# Python bytecode caches
|
||||
__pycache__/
|
||||
*.pyc
|
||||
@@ -0,0 +1,416 @@
|
||||
# sa3_tflite — Stable Audio 3 on CPU via LiteRT / TFLite
|
||||
|
||||
Portable CPU inference for **Stable Audio 3** — the LiteRT/TFLite sibling of the
|
||||
[MLX](../mlx) (Apple Silicon) and [TensorRT](../tensorRT) (NVIDIA GPU) releases.
|
||||
No PyTorch, transformers, or stable-audio-tools at runtime — just `ai_edge_litert`
|
||||
(LiteRT) driving fully self-contained `.tflite` graphs through the XNNPACK CPU
|
||||
delegate. Runs anywhere LiteRT runs: **macOS / Linux / Windows, x86 / ARM**
|
||||
(Windows is x64-only — see [Windows](#windows)).
|
||||
|
||||
## Quick Install
|
||||
|
||||
One line on a fresh machine — installs everything and plays back ~30 seconds of
|
||||
"Impending tribal, epic orchestral buildup":
|
||||
|
||||
```bash
|
||||
curl -LsSf https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.sh | bash
|
||||
```
|
||||
|
||||
Windows — stock PowerShell (no bash/git needed; the curl|bash line above requires Git Bash or WSL on Windows):
|
||||
|
||||
```powershell
|
||||
irm https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.ps1 | iex
|
||||
```
|
||||
|
||||
Already cloned the repo? Run from inside `optimized/tflite/`:
|
||||
|
||||
```bash
|
||||
# macOS / Linux / Windows Git Bash
|
||||
./install.sh # one-time setup
|
||||
./sa3 --prompt "Impending tribal, epic orchestral buildup" --play # generates + plays
|
||||
```
|
||||
|
||||
```bat
|
||||
:: Windows cmd / PowerShell
|
||||
install.bat
|
||||
sa3.bat --prompt "Impending tribal, epic orchestral buildup" --play
|
||||
:: PowerShell users can use .\sa3.ps1 instead of sa3.bat — same flags, and
|
||||
:: Ctrl-C exits cleanly (cmd's "Terminate batch job (Y/N)?" prompt is a .bat quirk)
|
||||
```
|
||||
|
||||
## Three models, four modes
|
||||
|
||||
| `--dit` | model | best for |
|
||||
|------------|--------------------|--------------------------------|
|
||||
| `sm-music` | sa3-sm-music (50 M block) | fast music generation |
|
||||
| `sm-sfx` | sa3-sm-sfx (50 M block) | sound effects |
|
||||
| `medium` | sa3-medium-ARC (1.4 B) | higher-quality music, slower |
|
||||
|
||||
| mode | flags | example |
|
||||
|------------------|-----------------------------------------------|----------------------------------|
|
||||
| text-to-audio | `--prompt P` | new clip from a description |
|
||||
| audio-to-audio | `--prompt P --init-audio IN.wav --init-noise-level σ` | variation of an existing clip |
|
||||
| inpainting | `--prompt P --init-audio IN.wav --inpaint-range "S,E"` | regenerate one section, keep rest |
|
||||
| CFG + negative | `--cfg 3.0 --negative-prompt P_NEG` | steer toward / away from prompts |
|
||||
|
||||
```
|
||||
prompt ─▶ T5Gemma encoder ─▶ DiT pingpong sampler ─▶ SAME-S/L decoder ─▶ WAV
|
||||
▲
|
||||
optional: encoder + init audio (audio-to-audio / inpaint)
|
||||
```
|
||||
|
||||
## Precision variants (`--precision`)
|
||||
|
||||
Same flag as the TensorRT CLI's `--precision`; the values use the wXaY naming
|
||||
(weights/activations bit-widths). One flag switches the DiT, the codec decoder,
|
||||
and (for audio-to-audio / inpainting) the encoder together; T5Gemma stays
|
||||
single-precision. All
|
||||
variants keep the full feature set: variable length (any `--seconds`, odd or even
|
||||
latent counts) and variable batch (batched or sequential CFG). Missing files
|
||||
lazy-download from HuggingFace on first use.
|
||||
|
||||
| `--precision` | legacy name | size (sm DiT / medium DiT / codecs) | quality | CPU speed |
|
||||
|---------------|-------------|-------------------------------------|---------|-----------|
|
||||
| `fp32` (default) | — | 1.8 GB / 5.8 GB / 0.2–1.8 GB | reference | 1× — on CPU this is *also* the fast choice |
|
||||
| `w16a32` | fp16 | 0.9 / 2.9 / 0.1–0.9 GB | ≈lossless (fp16 weights, fp32 activations; 62–75 dB per-forward) | 1.5–3× slower, model-dependent (XNNPACK dequantizes per-matmul) |
|
||||
| `w8a32` | woint8 | 0.45 / 1.5 / 0.05–0.5 GB | GPTQ int8 weights — codecs transparent (40–46 dB); DiT gives a *different but plausible* sample | ≈fp32 |
|
||||
| `w8a8-dyn` | dynint8 | 0.45 / 1.5 / 0.05–0.5 GB | lowest (int8 weights + activations, per-invoke dynamic scales) | fastest (~1.2–1.3×) |
|
||||
|
||||
*(Names follow the wXaY convention — weight/activation bit-widths, as in LLM releases:
|
||||
quantization here touches the FULLY_CONNECTED weights; activations stay fp32 except
|
||||
`w8a8-dyn`, whose int8×int8 matmuls are what make it the only faster-than-fp32 variant.
|
||||
All bit-widths are per-channel weight grids; "16" is fp16 — there is no int16 variant.
|
||||
Speed factors measured on an Apple M4 Pro MacBook Pro, XNNPACK, 8 threads.)*
|
||||
|
||||
Two things worth knowing, both counter-intuitive:
|
||||
|
||||
- **Quantization buys *size*, not speed, on CPU.** Weight-only int8 dequantizes
|
||||
to fp32 before each matmul, so they run at fp32 speed; fp16 is *slower* than fp32.
|
||||
Only `dynint8` (quantized activations → true int8 matmuls) is faster.
|
||||
- **DiT quantization error compounds.** The 8-step sampler is chaotically sensitive:
|
||||
per-step weight error turns into a *different* (still plausible) sample rather than
|
||||
a noisy version of the fp32 one. Judge DiT precisions by ear; decoder precisions
|
||||
by PSNR (they run once on a fixed latent).
|
||||
|
||||
One CFG note: batched and sequential CFG are bit-identical for `fp32`/`w8a32` and
|
||||
inaudibly different (~80 dB) for `w16a32`, but under `w8a8-dyn` the batch=2
|
||||
invoke shares activation-quantization scales across the cond/uncond rows, so batched
|
||||
CFG yields a *different plausible sample* than sequential. Pass `--no-cfg-batched`
|
||||
with `w8a8-dyn` when you need run-to-run reproducibility against sequential baselines.
|
||||
|
||||
```bash
|
||||
# quarter-size models, same speed and near-identical quality on the decoder side
|
||||
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s --precision w8a32
|
||||
|
||||
# fastest CPU inference (quality tradeoff)
|
||||
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s --precision w8a8-dyn
|
||||
```
|
||||
|
||||
**Encoder variants** (used by audio-to-audio / inpainting) exist at the same
|
||||
precisions; int8 weights are GPTQ-calibrated on real audio, like the DiT and codec
|
||||
(measured latent PSNR vs fp32, same-s / same-l): `w16a32` 66 / 71 dB (≈lossless),
|
||||
`w8a32` 36 / 46 dB, `w8a8-dyn` 24 / 30 dB. For the most quality-sensitive
|
||||
inpainting you can still pin `--encoder-precision fp32`.
|
||||
|
||||
`--dit-precision` / `--decoder-precision` / `--encoder-precision` override the shared
|
||||
flag per component — useful because they react to quantization very differently (the
|
||||
codec's precision maps directly to audio fidelity; the DiT's changes *which* sample
|
||||
you get; the encoder's affects a2a/inpaint latents):
|
||||
|
||||
```bash
|
||||
# fastest DiT, reference-quality codec
|
||||
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s \
|
||||
--precision fp32 --dit-precision w8a8-dyn
|
||||
```
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
./install.sh
|
||||
```
|
||||
|
||||
`install.sh` is uv-based. On a fresh machine it will:
|
||||
|
||||
1. Install [uv](https://github.com/astral-sh/uv) via the official curl
|
||||
installer if it's missing (prompts y/N; `-y` skips the prompt).
|
||||
2. Create a project-local `.venv/` with managed Python 3.11.
|
||||
3. `uv pip install` the runtime deps into the venv (much faster than pip).
|
||||
4. Ask which DiT bundles to download from HuggingFace
|
||||
(`stabilityai/stable-audio-3-optimized`). Each pick pulls its matching
|
||||
audio codec; T5Gemma (the shared text encoder) is downloaded once.
|
||||
Already-present weights are skipped.
|
||||
|
||||
End-to-end on a fresh machine: **~10 seconds** + weight downloads.
|
||||
|
||||
> Don't want to pre-pick bundles? Skip install entirely and just run
|
||||
> `./sa3 --prompt …` — any missing model file is downloaded from HF on
|
||||
> first use and symlinked into `models/tflite/` from the HuggingFace cache.
|
||||
|
||||
Portable CPU (no GPU required). Python 3.9+. `./install.sh --python 3.12` to
|
||||
pin a different Python.
|
||||
|
||||
### Windows
|
||||
|
||||
The stack runs natively on Windows x64 — no WSL needed. `ai-edge-litert`
|
||||
ships `win_amd64` wheels for Python **3.10–3.13** (x64 only; no Windows-ARM
|
||||
wheels), so install a Python in that range from
|
||||
[python.org](https://www.python.org/downloads/) (check *"Add python.exe to
|
||||
PATH"*), then from `optimized\tflite\`:
|
||||
|
||||
```bat
|
||||
install.bat :: one-time setup (python -m venv + pip)
|
||||
sa3.bat --prompt "lofi house loop" --dit sm-music --decoder same-s
|
||||
```
|
||||
|
||||
`install.bat` / `sa3.bat` are the Windows twins of `./install.sh` / `./sa3`
|
||||
(plain venv + pip; uv not required). Notes:
|
||||
|
||||
- **ffmpeg is optional** — only needed for mp3 / 24-bit / non-44.1 kHz
|
||||
`--init-audio` inputs: `winget install ffmpeg` (or `choco install ffmpeg`).
|
||||
- **Symlinks**: downloaded weights are exposed via symlinks where possible;
|
||||
without Developer Mode, Windows disallows creating them, so the downloader
|
||||
automatically falls back to a hardlink (zero-copy) or a plain copy. No
|
||||
action needed either way.
|
||||
- **Long paths**: the HuggingFace cache nests deeply — if downloads fail with
|
||||
path-length errors, enable Windows long paths once (admin PowerShell):
|
||||
`Set-ItemProperty 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name LongPathsEnabled -Value 1`
|
||||
(the `LongPathsEnabled` registry switch), then restart the terminal.
|
||||
- `--play` uses the stdlib `winsound` player (macOS uses `afplay`).
|
||||
|
||||
## Run
|
||||
|
||||
`./sa3` is a thin shell wrapper around `.venv/bin/python scripts/sa3_tflite.py
|
||||
"$@"` that prompts to run `./install.sh` if uv or `.venv/` isn't set up.
|
||||
|
||||
```bash
|
||||
# Text-to-audio
|
||||
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s --out lofi.wav
|
||||
|
||||
# Sound effects
|
||||
./sa3 --prompt "footsteps on gravel" --dit sm-sfx --decoder same-s --out steps.wav
|
||||
|
||||
# Higher-quality music (medium DiT, chunked SAME-L decode)
|
||||
./sa3 --prompt "A beautiful piano arpeggio grows into a cinematic climax" \
|
||||
--dit medium --decoder same-l --seconds 30 --out piano.wav
|
||||
|
||||
# Audio-to-audio variation (σmax 0.4-0.8 typical)
|
||||
./sa3 --prompt "jazz fusion with electric piano" --dit sm-music --decoder same-s \
|
||||
--init-audio funk.wav --init-noise-level 0.7 --out funk_jazz.wav
|
||||
|
||||
# Inpaint seconds 4-7
|
||||
./sa3 --prompt "explosive drum break" --dit sm-music --decoder same-s \
|
||||
--init-audio funk.wav --inpaint-range "4,7" --out funk_drums.wav
|
||||
|
||||
# CFG + negative prompt
|
||||
./sa3 --prompt "ambient drone" --cfg 3.0 --negative-prompt "drums, vocals" \
|
||||
--dit sm-music --decoder same-s --out drone.wav
|
||||
|
||||
# Generate + play immediately (afplay/winsound/aplay; Ctrl-C stops both)
|
||||
./sa3 --prompt "rainforest" --dit sm-sfx --decoder same-s --play
|
||||
|
||||
# All options + categorised examples
|
||||
./sa3 --help
|
||||
```
|
||||
|
||||
Omit `--dit` / `--decoder` for an interactive arrow-key picker. Omit
|
||||
`--prompt` for a stdin prompt. Relative `--out` paths land in `output/`
|
||||
(auto-created); absolute paths are honoured as-is. The output path is
|
||||
printed prominently as a `▸ saved` line at the end of each run.
|
||||
|
||||
Use `--threads` to control the XNNPACK CPU thread count (default 8).
|
||||
|
||||
## Web UI (gradio)
|
||||
|
||||
A browser UI over the same TFLite backend — every generation mode plus a
|
||||
spectrogram player, a history panel, and playback options (auto-play, infinite
|
||||
radio with pre-generation, loop, hotswap, MP3/WAV save).
|
||||
|
||||
```bash
|
||||
./sa3-gradio # sm-music, opens a public share link
|
||||
./sa3-gradio --dit medium # start on the medium DiT
|
||||
./sa3-gradio --no-share # local-only (http://localhost:7860)
|
||||
./sa3-gradio --port 7871 # pick the port
|
||||
```
|
||||
|
||||
`./sa3-gradio` is a wrapper around `.venv/bin/python scripts/sa3_gradio.py`; on
|
||||
first run it offers to install the UI-only extras (`gradio`, `pillow`,
|
||||
`soundfile`) into `.venv`. In the browser:
|
||||
|
||||
- **Precision** dropdown next to the model picker — `fp32` (default) / `w16a32`
|
||||
/ `w8a32` / `w8a8-dyn`, applied to the DiT + codec.
|
||||
- **LoRA** panel — upload a `.safetensors` adapter or pick one from
|
||||
`loras/<model>/` (e.g. `loras/sa3-medium/`), set its strength, stack several;
|
||||
settings are remembered per model. The panel is disabled under a quantized
|
||||
precision (see the LoRA note above) and has no per-step range (frozen graph).
|
||||
- **Audio-to-audio / inpainting** — upload a reference clip; inpainting adds a
|
||||
start/end range. Both are combinable.
|
||||
|
||||
Generated WAV/MP3s land in `output/gradio/`. First use of a model/precision
|
||||
loads its weights (cached after).
|
||||
|
||||
### Without the wrapper
|
||||
|
||||
```bash
|
||||
.venv/bin/python scripts/sa3_tflite.py --prompt "..." --dit medium --decoder same-l
|
||||
# or, after `source .venv/bin/activate`:
|
||||
python scripts/sa3_tflite.py --prompt "..." --dit medium --decoder same-l
|
||||
```
|
||||
|
||||
## Speed & memory
|
||||
|
||||
This is a **CPU** path — it trades the GPU releases' speed for portability. The
|
||||
small models comfortably beat realtime on a modern laptop CPU; `medium` is slower
|
||||
(its DiT is ~5.8 GB fp32 and it chunk-decodes SAME-L). Use ≥ ~20 s clips: very
|
||||
short clips have too few latent tokens for the sampler to settle into a coherent
|
||||
loop. Throughput scales with `--threads` up to your physical core count (4–8 is
|
||||
the usual sweet spot; more threads on a short model adds overhead).
|
||||
|
||||
For sub-realtime latency on a supported device, prefer the GPU siblings:
|
||||
[MLX](../mlx) on Apple Silicon, [TensorRT](../tensorRT) on NVIDIA.
|
||||
|
||||
## Flag reference
|
||||
|
||||
| Flag | Default | Notes |
|
||||
|-----------------------|----------|-----------------------------------------------------------------------|
|
||||
| `--prompt` | (asks) | Text prompt; empty string = unconditional |
|
||||
| `--negative-prompt` | — | CFG uncond branch; only used when `--cfg ≠ 1.0` |
|
||||
| `--dit` | (asks) | `sm-music`, `sm-sfx`, or `medium` |
|
||||
| `--decoder` | (asks) | `same-s` (pairs with sm-*) or `same-l` (pairs with medium) |
|
||||
| `--seconds` | 30 | Output length (use ≥ ~20 s) |
|
||||
| `--steps` | 8 | Pingpong sampler steps; 1 = single forward (fastest), 8 = sweet spot |
|
||||
| `--seed` | random | Set for reproducibility; the chosen seed is printed at the end |
|
||||
| `--cfg` | 1.0 | Guidance scale; 1.0 = off, >1 toward prompt, <1 toward uncond. ≠1 runs cond+uncond each step |
|
||||
| `--apg` | 1.0 | Adaptive Projected Guidance; only matters when `--cfg ≠ 1` |
|
||||
| `--cfg-batched` | on | When `--cfg ≠ 1`, run cond+uncond as one batch=2 invoke on the variable-batch DiT (~7–29% faster on Apple-Silicon AMX). `--no-cfg-batched` → sequential batch=1 dual-pass. Bit-identical at `fp32`/`w8a32` — see [Precision variants](#precision-variants---precision) for `w16a32`/`w8a8-dyn` |
|
||||
| `--lora` | — | A `.safetensors` LoRA adapter (SA3-native/underfit or PEFT) merged into the DiT, optional `strength=S`; repeat to stack. Requires `--dit-precision fp32` or `w16a32`. See [LoRA](#lora) |
|
||||
| `--lora-strength` | 1.0 | Default strength for `--lora` adapters without their own `strength=` |
|
||||
| `--init-audio` | — | WAV (any format via ffmpeg) input for audio-to-audio / inpaint |
|
||||
| `--init-noise-level` | 1.0 | σmax; 0.4–0.8 typical for variation, 1.0 = full regen, >1 = overshoot |
|
||||
| `--inpaint-range` | — | `START,END` seconds; regenerate that span, keep the rest |
|
||||
| `--threads` | 8 | XNNPACK CPU threads (all TFLite models run on CPU) |
|
||||
| `--free-models` | on | Free each model after its last use; `--no-free-models` keeps them resident |
|
||||
| `--out` / `-o` | (auto) | Relative → `output/<file>`; absolute → as-is. 16-bit PCM stereo @ 44.1 kHz, trimmed to exactly `--seconds` |
|
||||
| `--play` | off | After writing, play the WAV: `afplay` (macOS) / `winsound` (Windows) / `aplay` (Linux); Ctrl-C stops both |
|
||||
|
||||
All `.tflite` models are **fp32** except T5Gemma, which is **fp16** (numerically
|
||||
lossless there). There is no dtype knob: on CPU, int8/fp16 weights buy size, not
|
||||
speed (XNNPACK dequantizes to fp32 to matmul), and int8 costs quality on the DiT
|
||||
— so this release ships the fp32 graphs directly. (See "Notes on the design".)
|
||||
|
||||
## LoRA
|
||||
|
||||
Apply a LoRA finetune by merging it into the DiT's weight buffers at load —
|
||||
same adapters and semantics as the MLX backend's `--lora`:
|
||||
|
||||
```bash
|
||||
# one adapter
|
||||
./sa3 --dit medium --decoder same-l --prompt "progressive metal" \
|
||||
--lora plini-sa3-380.safetensors
|
||||
|
||||
# per-adapter strength; stack several
|
||||
./sa3 --dit medium --decoder same-l --prompt "..." \
|
||||
--lora a.safetensors strength=0.8 --lora b.safetensors strength=0.5
|
||||
```
|
||||
|
||||
The adapter is a `.safetensors` file (SA3-native `train_lora.py` / underfit
|
||||
output, or a PEFT adapter directory) — pickle `.ckpt/.pt` is refused. All nine
|
||||
adapter types (lora, dora-rows/cols, bora, the four -xs) are supported; the base
|
||||
must match `--dit`. The merge is written into a cached copy of the DiT under
|
||||
`models/tflite/lora_cache/` (keyed by adapter + strength content hash) and reused
|
||||
on repeat runs, so the ~5–15 s patch cost is paid once. A medium fp32 cache entry
|
||||
is ~5.4 GB — delete `lora_cache/` to reclaim.
|
||||
|
||||
Requires `--dit-precision fp32` (default) or `w16a32`. **LoRA on the quantized
|
||||
DiTs (`w8a32` / `w8a8-dyn` / `w4a32`) isn't figured out yet** — those store
|
||||
weights as GPTQ-calibrated int8, so merging would mean dequantize → add the LoRA
|
||||
delta → re-quantize, and a naive re-quant throws away the GPTQ error-feedback
|
||||
grid (you'd get a model that's both LoRA-adapted *and* degraded). Doing it well
|
||||
needs a GPTQ pass over the merged weights; until then `--lora` refuses quantized
|
||||
precisions (the CLI errors, the web UI disables the LoRA panel). Use fp32/w16a32
|
||||
for LoRA — on CPU fp32 is the fast-and-accurate choice anyway, so this is rarely
|
||||
a real constraint. **Per-step gating (`steps=`) is also MLX-only** — a frozen
|
||||
TFLite graph merges weights once at load; use `optimized/mlx` for step-gated LoRA.
|
||||
|
||||
## Files
|
||||
|
||||
```
|
||||
sa3_tflite/
|
||||
├── sa3 ← shell wrapper (use this)
|
||||
├── install.sh ← uv bootstrap (run once)
|
||||
├── sa3.bat / install.bat ← Windows twins of sa3 / install.sh
|
||||
├── bootstrap.sh ← one-line curl installer
|
||||
├── README.md
|
||||
├── requirements.txt ← ai_edge_litert, numpy, sentencepiece, soundfile, huggingface_hub
|
||||
├── output/ ← default landing zone for generated WAVs
|
||||
├── scripts/
|
||||
│ ├── sa3_tflite.py ← orchestrator CLI (invoked by ./sa3)
|
||||
│ ├── weights.py ← weights manifest + HF auto-download
|
||||
│ ├── examples.py ← shared examples block (--help + post-install)
|
||||
│ └── install.py ← install.sh's Python half (bundle picker)
|
||||
└── models/
|
||||
├── tokenizer.model ← SentencePiece model, BUNDLED (~4 MB; T5Gemma tflite is encoder-only)
|
||||
├── defs/
|
||||
│ └── tflite_pipeline.py ← Tokenizer + T5Gemma front-end + pingpong schedule + sampler + WAV
|
||||
└── tflite/ ← .tflite models (auto-downloaded; ~2.3 GB small, ~9.5 GB medium)
|
||||
├── t5gemma/encoder_fp16.tflite 564 MB text encoder (fp16)
|
||||
├── sa3-sm-music/dit_fp32.tflite 1.8 GB small music DiT (conditioner baked in)
|
||||
├── sa3-sm-sfx/dit_fp32.tflite 1.8 GB small sfx DiT (conditioner baked in)
|
||||
├── sa3-m/dit_fp32.tflite 5.8 GB medium DiT (conditioner baked in)
|
||||
├── same-s/{enc,dec}_fp32.tflite ~220 MB each shared sm-* codec
|
||||
└── same-l/{enc,dec}_fp32.tflite ~1.8 GB each medium codec
|
||||
```
|
||||
|
||||
The DiT graphs are **baked-I/O**: the conditioner (prompt-padding + seconds
|
||||
embedder) and the patch/unpatch are compiled into the graph, so the DiT takes the
|
||||
raw T5Gemma output directly and the decoder emits audio directly. The two small
|
||||
DiTs share the SAME-S codec (bit-exact between checkpoints), so only one set of
|
||||
small-codec files is shipped.
|
||||
|
||||
## Auto-download from HuggingFace
|
||||
|
||||
Model files aren't bundled — they're pulled from
|
||||
`stabilityai/stable-audio-3-optimized` (under `tflite/…`) on first use and
|
||||
symlinked into `models/tflite/` from the HF cache. No duplication. Anonymous
|
||||
downloads work but are rate-limited; `huggingface-cli login` with a free read-only
|
||||
token lifts the cap. The SentencePiece tokenizer (`models/tokenizer.model`, ~4 MB)
|
||||
is the one weight that IS committed, since the `.tflite` T5Gemma is encoder-only.
|
||||
|
||||
## Notes on the design
|
||||
|
||||
- **Baked-I/O varlen graphs.** Each `.tflite` is a single self-contained graph
|
||||
with the conditioner and patch/unpatch in-graph, accepting a variable sequence
|
||||
length — so one file serves any `--seconds`. The DiT is a 6-input graph
|
||||
(`x, t, t5_hidden, t5_mask, seconds, local_add_cond`); feed raw T5 outputs and
|
||||
the in-graph conditioner handles prompt-padding + seconds-embedding.
|
||||
- **fp32 everywhere (except fp16 T5Gemma).** On CPU, quantizing buys size, not
|
||||
speed — XNNPACK dequantizes int8/fp16 weights to fp32 to matmul, so fp16 is
|
||||
actually *slower* and int8 gives no speedup. And the DiT will not go int8 at
|
||||
quality: per-step error compounds over the 8 chaotic sampling steps into a
|
||||
*different* (still plausible) sample, not a degraded one. So this release ships
|
||||
the fp32 graphs directly. T5Gemma fp16 is the sole exception — it's numerically
|
||||
lossless there and halves that file.
|
||||
- **Monotonic audio-to-audio schedule.** The pingpong schedule applies the LogSNR
|
||||
shift to the normalized `[1→0]` grid, then scales by σmax, so audio-to-audio
|
||||
(σmax < 1) stays monotonically decreasing while keeping all N distilled steps.
|
||||
σmax = 1.0 (text-to-audio) is bit-identical to the classic schedule.
|
||||
- **SAME-L chunked decode.** The SAME-L decoder's dense sliding-window-attention
|
||||
mask is O(T²), so long clips are decoded in overlap-8 windows of 64 latent
|
||||
tokens (the throughput optimum) and stitched. SAME-S has a narrow receptive
|
||||
field and decodes whole.
|
||||
- **CFG (`--cfg ≠ 1`)** combines a cond and an uncond velocity in denoised space
|
||||
(optional APG). The canonical DiT is **variable-batch**, so by default cond+uncond
|
||||
run as **one batch=2 invoke per step** (`--cfg-batched`) — ~7–29% faster on
|
||||
Apple-Silicon (the AMX matrix unit amortizes the weight loads across both rows;
|
||||
measured on an M4 Pro MacBook Pro).
|
||||
`--no-cfg-batched` falls back to a sequential batch=1 dual-pass (like the TensorRT
|
||||
release, whose engine is static-batch=1); the two are bit-identical at `fp32`/`w8a32`
|
||||
(~80 dB at `w16a32`; `w8a8-dyn` diverges by design — batch-shared activation scales).
|
||||
|
||||
## License & attribution
|
||||
|
||||
Model weights derived from Stability AI's Stable Audio 3 checkpoints.
|
||||
T5Gemma text encoder from Google.
|
||||
|
||||
Use of the Stable Audio 3 weights is governed by the **Stability AI
|
||||
Community License**. Please refer to the full terms at
|
||||
<https://stability.ai/license>.
|
||||
@@ -0,0 +1,88 @@
|
||||
# sa3_tflite bootstrap for Windows - Stable Audio 3 inference on CPU (LiteRT/TFLite).
|
||||
#
|
||||
# The PowerShell twin of bootstrap.sh. Hosted at:
|
||||
# https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.ps1
|
||||
#
|
||||
# Usage (default demo):
|
||||
# irm https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.ps1 | iex
|
||||
#
|
||||
# Custom args (iex cannot forward arguments - download first):
|
||||
# irm https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.ps1 -OutFile bootstrap.ps1
|
||||
# .\bootstrap.ps1 --prompt "Death Metal" --dit medium --decoder same-l
|
||||
#
|
||||
# What it does:
|
||||
# 1. Fetches the project (git clone if git is installed, else a zip pull).
|
||||
# 2. Runs install.bat (python -m venv + pip install; needs Python 3.10-3.13
|
||||
# on PATH - ai-edge-litert ships win_amd64 wheels for those versions).
|
||||
# 3. Runs sa3.bat with your args (default: 30 s demo prompt + --play).
|
||||
# Weights auto-download from HuggingFace on first use.
|
||||
|
||||
$ErrorActionPreference = "Stop"
|
||||
|
||||
# Older Windows PowerShell 5.1 hosts may default to TLS < 1.2, which GitHub
|
||||
# rejects. Harmless no-op on PowerShell 7+ / current Win10+.
|
||||
try { [Net.ServicePointManager]::SecurityProtocol = [Net.ServicePointManager]::SecurityProtocol -bor [Net.SecurityProtocolType]::Tls12 } catch {}
|
||||
|
||||
$RepoOwner = "Stability-AI"
|
||||
$RepoName = "stable-audio-3"
|
||||
$Branch = "main"
|
||||
$SubDir = "optimized\tflite"
|
||||
$LocalDir = "sa3_tflite"
|
||||
|
||||
function Step($msg) { Write-Host "`n-> $msg" -ForegroundColor Cyan }
|
||||
function Ok($msg) { Write-Host " OK $msg" -ForegroundColor Green }
|
||||
|
||||
# --- 1. fetch the project -------------------------------------------------
|
||||
if (Get-Command git -ErrorAction SilentlyContinue) {
|
||||
$WorkDir = Join-Path $RepoName $SubDir
|
||||
if (Test-Path (Join-Path $RepoName ".git")) {
|
||||
Step "Reusing existing .\$RepoName (git pull --ff-only)"
|
||||
git -C $RepoName pull --ff-only
|
||||
} elseif (Test-Path $RepoName) {
|
||||
throw ".\$RepoName exists but isn't a git repo - remove or rename it."
|
||||
} else {
|
||||
Step "git clone https://github.com/$RepoOwner/$RepoName -> .\$RepoName"
|
||||
git clone --depth=1 "https://github.com/$RepoOwner/$RepoName" $RepoName
|
||||
}
|
||||
if (-not (Test-Path $WorkDir)) { throw "Expected '$SubDir' inside the repo but didn't find it." }
|
||||
} else {
|
||||
$WorkDir = $LocalDir
|
||||
if (Test-Path (Join-Path $LocalDir "install.bat")) {
|
||||
Step "Reusing existing $LocalDir\ (delete it to re-download)"
|
||||
} elseif (Test-Path $LocalDir) {
|
||||
throw ".\$LocalDir exists but doesn't look like a sa3_tflite checkout - remove or rename it."
|
||||
} else {
|
||||
Step "git not installed - downloading $RepoOwner/$RepoName ($Branch) zip -> .\$LocalDir"
|
||||
$zip = Join-Path $env:TEMP "sa3_repo.zip"
|
||||
$extract = Join-Path $env:TEMP "sa3_extract"
|
||||
Invoke-WebRequest "https://github.com/$RepoOwner/$RepoName/archive/refs/heads/$Branch.zip" -OutFile $zip
|
||||
if (Test-Path $extract) { Remove-Item $extract -Recurse -Force }
|
||||
Expand-Archive $zip -DestinationPath $extract -Force
|
||||
Move-Item (Join-Path $extract "$RepoName-$Branch\$SubDir") $LocalDir
|
||||
Remove-Item $zip -Force
|
||||
Remove-Item $extract -Recurse -Force
|
||||
Ok "extracted to .\$LocalDir"
|
||||
}
|
||||
}
|
||||
|
||||
Set-Location $WorkDir
|
||||
Ok "ready at $(Get-Location)"
|
||||
|
||||
# --- 2. install -----------------------------------------------------------
|
||||
Step "Running install.bat"
|
||||
cmd /c install.bat
|
||||
if ($LASTEXITCODE -ne 0) { throw "install.bat failed (exit $LASTEXITCODE)." }
|
||||
|
||||
# --- 3. inference ---------------------------------------------------------
|
||||
if ($args.Count -gt 0) {
|
||||
Step "Running sa3.bat $($args -join ' ')"
|
||||
cmd /c sa3.bat @args
|
||||
} else {
|
||||
Step 'Running demo: sa3.bat --prompt "Impending tribal, epic orchestral buildup" --dit sm-music --decoder same-s --seconds 30 --play'
|
||||
Write-Host ' (for custom args, download bootstrap.ps1 first and run: .\bootstrap.ps1 --prompt "..." ...)' -ForegroundColor DarkGray
|
||||
cmd /c sa3.bat --prompt "Impending tribal, epic orchestral buildup" --dit sm-music --decoder same-s --seconds 30 --play
|
||||
}
|
||||
|
||||
Write-Host ""
|
||||
Write-Host "You are set up in $(Get-Location)" -ForegroundColor Green
|
||||
Write-Host " run .\sa3.bat --help for options" -ForegroundColor DarkGray
|
||||
+156
@@ -0,0 +1,156 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# sa3_tflite bootstrap — Stable Audio 3 inference on CPU (LiteRT/TFLite) in one command.
|
||||
#
|
||||
# Hosted at:
|
||||
# https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.sh
|
||||
#
|
||||
# Usage:
|
||||
# curl -LsSf https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.sh | bash
|
||||
# curl -LsSf https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.sh | bash -s -- --prompt "Death Metal" --dit medium --decoder same-l
|
||||
#
|
||||
# Default demo prompt is "Impending tribal, epic orchestral buildup".
|
||||
#
|
||||
# What it does:
|
||||
# 1. Verifies curl + tar are present (portable CPU stack — runs on macOS/Linux, x86/ARM).
|
||||
# 2. Fetches the project:
|
||||
# - If git is installed → `git clone --depth=1` into ./stable-audio-3/,
|
||||
# then cd into optimized/tflite/ (real repo; pullable, modifiable).
|
||||
# - If not → tarball pull via curl + tar, extracting only optimized/tflite/
|
||||
# into ./sa3_tflite/ (no git needed).
|
||||
# 3. Runs ./install.sh -y inside it (uv + Python 3.11 + venv + weight downloads).
|
||||
# 4. Runs ./sa3 with whatever args you passed (default: "Impending tribal, epic orchestral buildup" demo + --play).
|
||||
#
|
||||
set -euo pipefail
|
||||
|
||||
# uv's curl installer drops the binary at $XDG_BIN_HOME (~/.local/bin by default)
|
||||
# and updates the user's shell profile — but that profile only takes effect in
|
||||
# *new* shells. We pre-emptively put both locations on PATH so the just-installed
|
||||
# uv (and anything else from this run) is findable in the current process tree.
|
||||
export PATH="${XDG_BIN_HOME:-$HOME/.local/bin}:$HOME/.local/bin:$PATH"
|
||||
|
||||
REPO_OWNER="Stability-AI"
|
||||
REPO_NAME="stable-audio-3"
|
||||
BRANCH="main"
|
||||
SUBDIR_IN_REPO="optimized/tflite"
|
||||
LOCAL_DIR="sa3_tflite"
|
||||
DEFAULT_ARGS=(--prompt "Impending tribal, epic orchestral buildup" --dit sm-music --decoder same-s --seconds 30 --play)
|
||||
|
||||
TAR_URL="https://github.com/$REPO_OWNER/$REPO_NAME/archive/refs/heads/$BRANCH.tar.gz"
|
||||
TAR_INNER="$REPO_NAME-$BRANCH/$SUBDIR_IN_REPO"
|
||||
|
||||
# ── colours ─────────────────────────────────────────────────────────────────
|
||||
if [[ -t 1 ]]; then
|
||||
BOLD=$'\033[1m'; CYAN=$'\033[1;36m'; RED=$'\033[1;31m'
|
||||
YELLOW=$'\033[1;33m'; GREEN=$'\033[1;32m'; DIM=$'\033[2m'; RESET=$'\033[0m'
|
||||
else
|
||||
BOLD=""; CYAN=""; RED=""; YELLOW=""; GREEN=""; DIM=""; RESET=""
|
||||
fi
|
||||
step() { printf '\n%s→ %s%s\n' "$CYAN" "$1" "$RESET"; }
|
||||
fail() { printf '\n%serror%s: %s\n' "$RED" "$RESET" "$1" >&2; exit 1; }
|
||||
ok() { printf ' %s✓%s %s\n' "$GREEN" "$RESET" "$1"; }
|
||||
warn() { printf '%swarning%s: %s\n' "$YELLOW" "$RESET" "$1" >&2; }
|
||||
|
||||
# ── 1. platform note (portable — no gate) ───────────────────────────────────
|
||||
OS="$(uname -s)"; ARCH="$(uname -m)"
|
||||
ok "platform: $OS/$ARCH (LiteRT/TFLite is CPU-portable)"
|
||||
|
||||
# ── 2. preflight: curl + tar (preinstalled on macOS/most Linux) ─────────────
|
||||
for tool in curl tar; do
|
||||
command -v "$tool" >/dev/null 2>&1 || \
|
||||
fail "$tool not found on PATH. Install it and re-run."
|
||||
done
|
||||
ok "curl + tar present"
|
||||
|
||||
# ── 3. fetch the project ────────────────────────────────────────────────────
|
||||
# Prefer `git clone` if git is on the machine — the user gets a real repo
|
||||
# they can pull updates from / navigate sibling subdirs in (mlx/, tensorRT/).
|
||||
# Falls back to a tarball pull (curl + tar) if git is missing.
|
||||
|
||||
if command -v git >/dev/null 2>&1; then
|
||||
GIT_DIR="$REPO_NAME"
|
||||
WORK_DIR="$GIT_DIR/$SUBDIR_IN_REPO"
|
||||
|
||||
if [[ -d "$GIT_DIR/.git" ]]; then
|
||||
step "Reusing existing ./$GIT_DIR (git pull --ff-only)"
|
||||
git -C "$GIT_DIR" pull --ff-only
|
||||
elif [[ -e "$GIT_DIR" ]]; then
|
||||
fail "./$GIT_DIR exists but isn't a git repo — remove or rename it."
|
||||
else
|
||||
step "git clone https://github.com/$REPO_OWNER/$REPO_NAME → ./$GIT_DIR"
|
||||
git clone --depth=1 "https://github.com/$REPO_OWNER/$REPO_NAME" "$GIT_DIR"
|
||||
fi
|
||||
|
||||
[[ -d "$WORK_DIR" ]] || \
|
||||
fail "Expected '$SUBDIR_IN_REPO' inside the repo but didn't find it."
|
||||
ok "ready at ./$WORK_DIR"
|
||||
else
|
||||
# No git — pull a tarball and extract only optimized/tflite/.
|
||||
WORK_DIR="$LOCAL_DIR"
|
||||
|
||||
if [[ -d "$LOCAL_DIR" && -x "$LOCAL_DIR/install.sh" ]]; then
|
||||
step "Reusing existing $LOCAL_DIR/ (delete it to re-download)"
|
||||
else
|
||||
if [[ -e "$LOCAL_DIR" ]]; then
|
||||
fail "./$LOCAL_DIR exists but doesn't look like a sa3_tflite checkout — remove or rename it."
|
||||
fi
|
||||
step "git not installed — downloading $REPO_OWNER/$REPO_NAME ($BRANCH) tarball → ./$LOCAL_DIR"
|
||||
|
||||
TMP_TAR="$(mktemp -t sa3_repo.XXXXXX).tar.gz"
|
||||
TMP_EXTRACT="$(mktemp -d -t sa3_extract.XXXXXX)"
|
||||
trap 'rm -rf "$TMP_TAR" "$TMP_EXTRACT"' EXIT
|
||||
|
||||
# --progress-bar writes to stderr; -f makes 404/5xx a real curl error
|
||||
curl -fL --progress-bar "$TAR_URL" -o "$TMP_TAR"
|
||||
|
||||
# BSD tar (macOS) / GNU tar both extract only paths matching the pattern.
|
||||
tar -xz -f "$TMP_TAR" -C "$TMP_EXTRACT" "$TAR_INNER"
|
||||
|
||||
SRC="$TMP_EXTRACT/$TAR_INNER"
|
||||
[[ -d "$SRC" ]] || fail "Expected '$TAR_INNER' inside the tarball but didn't find it."
|
||||
mv "$SRC" "$LOCAL_DIR"
|
||||
ok "extracted $(find "$LOCAL_DIR" -type f | wc -l | tr -d ' ') files to ./$LOCAL_DIR"
|
||||
fi
|
||||
fi
|
||||
|
||||
# ── 4. install ──────────────────────────────────────────────────────────────
|
||||
cd "$WORK_DIR"
|
||||
[[ -x ./install.sh ]] || fail "install.sh missing or not executable in ./$WORK_DIR."
|
||||
step "Running ./install.sh -y"
|
||||
./install.sh -y
|
||||
|
||||
# ── 5. inference ────────────────────────────────────────────────────────────
|
||||
# Run as a subprocess (not `exec`) so we can drop the user into an
|
||||
# interactive shell here when it finishes.
|
||||
if [[ $# -gt 0 ]]; then
|
||||
step "Running ./sa3 $*"
|
||||
./sa3 "$@" || true
|
||||
else
|
||||
step "Running demo: ./sa3 ${DEFAULT_ARGS[*]}"
|
||||
printf ' %s(pass your own args via: curl -LsSf https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.sh | bash -s -- --prompt "..." ...)%s\n' "$DIM" "$RESET"
|
||||
./sa3 "${DEFAULT_ARGS[@]}" || true
|
||||
fi
|
||||
|
||||
# ── 6. drop user into an interactive shell sitting in the project dir ──────
|
||||
# A subprocess can't change its parent shell's cwd — but we CAN replace
|
||||
# ourselves with a fresh interactive shell, leaving the user at a prompt
|
||||
# inside ./$WORK_DIR. `exit` (or Ctrl-D) returns them to their original
|
||||
# shell, at their original cwd, just like a normal subshell.
|
||||
#
|
||||
# `< /dev/tty` is essential when bootstrap.sh was invoked via curl|bash:
|
||||
# stdin at this point is the (closed) pipe; an interactive shell needs a
|
||||
# real terminal. /dev/tty always refers to the user's controlling TTY.
|
||||
|
||||
if [[ ! -e /dev/tty ]]; then
|
||||
# Headless / scripted invocation — skip the shell drop.
|
||||
exit 0
|
||||
fi
|
||||
|
||||
USER_SHELL="${SHELL:-/bin/bash}"
|
||||
printf '\n%s━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━%s\n' "$BOLD" "$RESET"
|
||||
printf ' %s✓ you are now in%s %s%s%s\n' "$GREEN" "$RESET" "$BOLD" "$(pwd)" "$RESET"
|
||||
printf ' %stype %s./sa3 --help%s for options, or %sexit%s to return to your previous shell%s\n' \
|
||||
"$DIM" "$RESET$BOLD" "$RESET$DIM" "$RESET$BOLD" "$RESET$DIM" "$RESET"
|
||||
printf '%s━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━%s\n\n' "$BOLD" "$RESET"
|
||||
|
||||
exec "$USER_SHELL" -i < /dev/tty
|
||||
@@ -0,0 +1,43 @@
|
||||
@echo off
|
||||
rem SA3 TFLite installer for Windows — the .bat twin of install.sh.
|
||||
rem
|
||||
rem Mirrors install.sh's steps with plain stdlib tooling (no uv required):
|
||||
rem 1. create a project-local .venv\ (python -m venv)
|
||||
rem 2. pip install -r requirements.txt into it
|
||||
rem 3. hand off to scripts\install.py for the weight-download prompt
|
||||
rem
|
||||
rem Extra args are forwarded to install.py, e.g.:
|
||||
rem install.bat --download sm-music
|
||||
rem
|
||||
rem Note: ai-edge-litert ships win_amd64 wheels for Python 3.10-3.13 —
|
||||
rem use a Python in that range (3.11 recommended).
|
||||
setlocal
|
||||
set "SCRIPT_DIR=%~dp0"
|
||||
|
||||
where python >nul 2>nul
|
||||
if errorlevel 1 (
|
||||
echo error: python not found on PATH. 1>&2
|
||||
echo Install Python 3.10-3.13 from https://www.python.org/downloads/ 1>&2
|
||||
echo ^(check "Add python.exe to PATH" in the installer^) and re-run install.bat. 1>&2
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
if exist "%SCRIPT_DIR%.venv\Scripts\python.exe" (
|
||||
echo Reusing existing .venv\
|
||||
) else (
|
||||
echo Creating virtual environment at .venv\ ...
|
||||
python -m venv "%SCRIPT_DIR%.venv"
|
||||
if errorlevel 1 exit /b 1
|
||||
)
|
||||
|
||||
echo Installing dependencies ^(pip install -r requirements.txt^) ...
|
||||
"%SCRIPT_DIR%.venv\Scripts\python.exe" -m pip install --upgrade pip
|
||||
if errorlevel 1 exit /b 1
|
||||
"%SCRIPT_DIR%.venv\Scripts\python.exe" -m pip install -r "%SCRIPT_DIR%requirements.txt"
|
||||
if errorlevel 1 exit /b 1
|
||||
|
||||
rem install.py's pip step is skipped — deps were just installed above
|
||||
rem (same contract install.sh uses).
|
||||
set "INSTALL_SKIP_PIP=1"
|
||||
"%SCRIPT_DIR%.venv\Scripts\python.exe" "%SCRIPT_DIR%scripts\install.py" %*
|
||||
exit /b %ERRORLEVEL%
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# SA3 TFLite installer — uv-based.
|
||||
#
|
||||
# Creates a project-local .venv/ with the right Python and runtime deps,
|
||||
# then hands off to install.py for the interactive weight-download prompt.
|
||||
#
|
||||
# Portable CPU stack (LiteRT / TFLite via ai_edge_litert) — runs on
|
||||
# macOS/Linux, x86/ARM. No Apple-Silicon requirement.
|
||||
#
|
||||
# Usage:
|
||||
# ./install.sh # auto-detect uv, prompt to install if missing
|
||||
# ./install.sh -y # assume yes to "install uv?" prompt
|
||||
# ./install.sh --python VER # pin a specific Python (default: 3.11)
|
||||
#
|
||||
# After install:
|
||||
# source .venv/bin/activate # (.venv/Scripts/activate on Windows Git Bash)
|
||||
# python scripts/sa3_tflite.py --prompt "lofi house" --dit sm-music --decoder same-s
|
||||
# # or, without activating:
|
||||
# .venv/bin/python scripts/sa3_tflite.py --prompt "lofi house" --dit sm-music
|
||||
#
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
VENV_DIR="$SCRIPT_DIR/.venv"
|
||||
|
||||
# venv layout differs by OS: POSIX puts python at .venv/bin/python, Windows
|
||||
# (incl. Git Bash / MSYS) at .venv/Scripts/python.exe. Resolve whichever exists.
|
||||
venv_py() {
|
||||
if [[ -x "$VENV_DIR/bin/python" ]]; then
|
||||
echo "$VENV_DIR/bin/python"
|
||||
elif [[ -x "$VENV_DIR/Scripts/python.exe" ]]; then
|
||||
echo "$VENV_DIR/Scripts/python.exe"
|
||||
else
|
||||
echo "$VENV_DIR/bin/python" # default (error messages)
|
||||
fi
|
||||
}
|
||||
PY_VERSION_DEFAULT="3.11"
|
||||
|
||||
# ── colours ─────────────────────────────────────────────────────────────────
|
||||
if [[ -t 1 ]]; then
|
||||
BOLD=$'\033[1m'; CYAN=$'\033[1;36m'; RED=$'\033[1;31m'
|
||||
YELLOW=$'\033[1;33m'; GREEN=$'\033[1;32m'; DIM=$'\033[2m'; RESET=$'\033[0m'
|
||||
else
|
||||
BOLD=""; CYAN=""; RED=""; YELLOW=""; GREEN=""; DIM=""; RESET=""
|
||||
fi
|
||||
step() { printf '\n%s→ %s%s\n' "$CYAN" "$1" "$RESET"; }
|
||||
fail() { printf '%serror%s: %s\n' "$RED" "$RESET" "$1" >&2; }
|
||||
warn() { printf '%swarning%s: %s\n' "$YELLOW" "$RESET" "$1" >&2; }
|
||||
ok() { printf ' %s✓%s %s\n' "$GREEN" "$RESET" "$1"; }
|
||||
|
||||
# ── arg parsing ─────────────────────────────────────────────────────────────
|
||||
ASSUME_YES=0
|
||||
PY_VERSION="$PY_VERSION_DEFAULT"
|
||||
EXTRA_ARGS=()
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
-y|--yes) ASSUME_YES=1; shift ;;
|
||||
--python) PY_VERSION="$2"; shift 2 ;;
|
||||
--python=*) PY_VERSION="${1#--python=}"; shift ;;
|
||||
-h|--help)
|
||||
sed -n '2,/^set -euo/p' "$0" | sed -e '$d' -e 's/^# \{0,1\}//'
|
||||
exit 0 ;;
|
||||
*) EXTRA_ARGS+=("$1"); shift ;;
|
||||
esac
|
||||
done
|
||||
|
||||
# ── platform note (portable — CPU only, no gate) ────────────────────────────
|
||||
OS="$(uname -s)"; ARCH="$(uname -m)"
|
||||
ok "platform: $OS/$ARCH (LiteRT/TFLite is CPU-portable; XNNPACK delegate)"
|
||||
|
||||
# ── ensure uv is installed ──────────────────────────────────────────────────
|
||||
ensure_uv() {
|
||||
if command -v uv >/dev/null 2>&1; then
|
||||
ok "uv $(uv --version 2>/dev/null | awk '{print $2}') already installed at $(command -v uv)"
|
||||
return 0
|
||||
fi
|
||||
|
||||
step "uv not found — uv is required (much faster than pip, also manages Python versions)"
|
||||
if [[ "$ASSUME_YES" -ne 1 ]]; then
|
||||
printf ' Install uv now via the official installer? (curl + sh) %s[Y/n]%s ' "$DIM" "$RESET"
|
||||
read -r REPLY < /dev/tty
|
||||
case "$REPLY" in
|
||||
""|y|Y|yes|YES) ;;
|
||||
*)
|
||||
fail "install aborted — install uv manually then re-run:"
|
||||
printf ' curl -LsSf https://astral.sh/uv/install.sh | sh\n' >&2
|
||||
printf ' or: brew install uv\n' >&2
|
||||
exit 1 ;;
|
||||
esac
|
||||
fi
|
||||
|
||||
step "Installing uv (curl -LsSf https://astral.sh/uv/install.sh | sh)"
|
||||
if ! curl -LsSf https://astral.sh/uv/install.sh | sh; then
|
||||
fail "uv installer failed. Try a manual install:"
|
||||
printf ' brew install uv\n' >&2
|
||||
exit 1
|
||||
fi
|
||||
# The installer drops uv at ~/.local/bin/uv (or $XDG_BIN_HOME)
|
||||
export PATH="$HOME/.local/bin:${XDG_BIN_HOME:-}:$PATH"
|
||||
if ! command -v uv >/dev/null 2>&1; then
|
||||
fail "uv was installed but isn't on PATH. Add ~/.local/bin to PATH, restart your shell, and re-run install.sh."
|
||||
exit 1
|
||||
fi
|
||||
ok "uv $(uv --version | awk '{print $2}') installed"
|
||||
}
|
||||
|
||||
ensure_uv
|
||||
|
||||
# ── create venv (uv auto-installs the requested Python if missing) ──────────
|
||||
step "Creating virtual environment at .venv/ with Python $PY_VERSION"
|
||||
if [[ -d "$VENV_DIR" ]]; then
|
||||
EXISTING_PY=$("$(venv_py)" -c 'import sys; print(".".join(map(str, sys.version_info[:2])))' 2>/dev/null || echo "unknown")
|
||||
if [[ "$EXISTING_PY" == "$PY_VERSION"* ]]; then
|
||||
ok "reusing existing .venv (Python $EXISTING_PY)"
|
||||
else
|
||||
warn "existing .venv uses Python $EXISTING_PY (wanted $PY_VERSION) — recreating"
|
||||
rm -rf "$VENV_DIR"
|
||||
uv venv --seed --python "$PY_VERSION" "$VENV_DIR"
|
||||
fi
|
||||
else
|
||||
uv venv --seed --python "$PY_VERSION" "$VENV_DIR"
|
||||
fi
|
||||
|
||||
# ── install runtime deps ────────────────────────────────────────────────────
|
||||
step "Installing dependencies (uv pip install -r requirements.txt)"
|
||||
VIRTUAL_ENV="$VENV_DIR" uv pip install -r "$SCRIPT_DIR/requirements.txt"
|
||||
|
||||
# ── hand off to install.py ──────────────────────────────────────────────────
|
||||
# Any unrecognized args we collected (e.g. --download medium,sm-music) get
|
||||
# forwarded to install.py via EXTRA_ARGS.
|
||||
step "Bundle picker"
|
||||
INSTALL_SKIP_PIP=1 exec "$(venv_py)" "$SCRIPT_DIR/scripts/install.py" \
|
||||
"${EXTRA_ARGS[@]+"${EXTRA_ARGS[@]}"}"
|
||||
@@ -0,0 +1,4 @@
|
||||
# Extras for the gradio web UI (./sa3-gradio) — not needed by the CLI.
|
||||
gradio>=4.0
|
||||
pillow>=10.0
|
||||
soundfile>=0.12 # decode mp3/flac/24-bit/48kHz uploads without ffmpeg
|
||||
@@ -0,0 +1,5 @@
|
||||
ai_edge_litert>=1.0
|
||||
numpy>=1.24
|
||||
sentencepiece>=0.2
|
||||
soundfile>=0.12
|
||||
huggingface_hub>=0.20
|
||||
Executable
+82
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Wrapper around scripts/sa3_tflite.py — runs it via the project's .venv so no
|
||||
# activation is needed.
|
||||
#
|
||||
# If setup is incomplete (no uv, no .venv) you'll be prompted to run
|
||||
# install.sh; saying yes runs the installer and then resumes the command.
|
||||
#
|
||||
# Usage: ./sa3 --prompt "lofi house" --dit sm-music --decoder same-s --out a.wav
|
||||
#
|
||||
set -euo pipefail
|
||||
|
||||
# Pre-emptively add the uv installer's default bin dir to PATH. The user's
|
||||
# shell profile only updates new shells; this lets us find uv in the current
|
||||
# process tree (matters right after a fresh `curl|sh` install).
|
||||
export PATH="${XDG_BIN_HOME:-$HOME/.local/bin}:$HOME/.local/bin:$PATH"
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
if [[ -t 1 ]]; then
|
||||
RED=$'\033[1;31m'; YELLOW=$'\033[1;33m'; CYAN=$'\033[1;36m'
|
||||
BOLD=$'\033[1m'; DIM=$'\033[2m'; RESET=$'\033[0m'
|
||||
else
|
||||
RED=""; YELLOW=""; CYAN=""; BOLD=""; DIM=""; RESET=""
|
||||
fi
|
||||
|
||||
fail() { printf '\n%serror%s: %s\n' "$RED" "$RESET" "$1" >&2; exit 1; }
|
||||
|
||||
# Prompts the user to run install.sh. Honors `-y` / non-interactive stdin.
|
||||
# Resumes (returns 0) on success so the caller can re-attempt the run;
|
||||
# exits 1 if the user declines or install fails.
|
||||
prompt_install() {
|
||||
local reason="$1"
|
||||
printf '\n%s%s%s\n' "$YELLOW" "$reason" "$RESET" >&2
|
||||
|
||||
if [[ ! -t 0 ]]; then
|
||||
printf '%serror%s: stdin is not a TTY — re-run %s./install.sh%s manually.\n' \
|
||||
"$RED" "$RESET" "$BOLD" "$RESET" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
printf ' %sRun ./install.sh now to set up?%s %s[Y/n]%s ' \
|
||||
"$BOLD" "$RESET" "$DIM" "$RESET"
|
||||
read -r REPLY < /dev/tty
|
||||
case "$REPLY" in
|
||||
""|y|Y|yes|YES) ;;
|
||||
*) fail "install declined — re-run ./install.sh manually when ready." ;;
|
||||
esac
|
||||
|
||||
printf '\n%s→ running ./install.sh%s\n' "$CYAN" "$RESET" >&2
|
||||
if ! "$SCRIPT_DIR/install.sh"; then
|
||||
fail "installer failed — see output above."
|
||||
fi
|
||||
}
|
||||
|
||||
# ── infrastructure checks ───────────────────────────────────────────────────
|
||||
if [[ ! -f "$SCRIPT_DIR/scripts/sa3_tflite.py" ]]; then
|
||||
fail "scripts/sa3_tflite.py not found — repo files are missing or moved."
|
||||
fi
|
||||
|
||||
needs_install=0
|
||||
missing=()
|
||||
if ! command -v uv >/dev/null 2>&1; then
|
||||
needs_install=1; missing+=("uv")
|
||||
fi
|
||||
if [[ ! -d "$SCRIPT_DIR/.venv" ]]; then
|
||||
needs_install=1; missing+=(".venv/")
|
||||
fi
|
||||
|
||||
if [[ $needs_install -eq 1 ]]; then
|
||||
prompt_install "Setup incomplete — missing: ${missing[*]}"
|
||||
fi
|
||||
|
||||
# ── hand off to scripts/sa3_tflite.py ───────────────────────────────────────
|
||||
# Invoke .venv/bin/python directly rather than via `uv run`. `uv run` walks up
|
||||
# looking for a pyproject.toml and (in newer uv versions) with --no-project
|
||||
# creates an ephemeral env with no deps, ignoring our .venv. Direct invocation
|
||||
# is shorter, faster, and side-steps both bugs.
|
||||
cd "$SCRIPT_DIR"
|
||||
PY="$SCRIPT_DIR/.venv/bin/python"
|
||||
[[ -x "$PY" ]] || PY="$SCRIPT_DIR/.venv/Scripts/python.exe"
|
||||
exec "$PY" scripts/sa3_tflite.py "$@"
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env bash
|
||||
# Shortcut for running scripts/sa3_gradio.py with the project .venv.
|
||||
# All args pass through unchanged.
|
||||
#
|
||||
# Examples:
|
||||
# ./sa3-gradio # sm-music + same-s, public share link
|
||||
# ./sa3-gradio --dit medium
|
||||
# ./sa3-gradio --no-share # local-only
|
||||
# ./sa3-gradio --help
|
||||
set -e
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PY="${SCRIPT_DIR}/.venv/bin/python"
|
||||
|
||||
if [ ! -x "$PY" ]; then
|
||||
echo "No .venv found — run ./install.sh (or ./bootstrap.sh) first."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# gradio + pillow + soundfile are UI-only extras, not part of the base requirements.
|
||||
if ! "$PY" -c "import gradio, PIL, soundfile" >/dev/null 2>&1; then
|
||||
echo "The gradio UI needs a few extra packages (gradio, pillow, soundfile)."
|
||||
read -r -p "Install them into .venv now? [Y/n] " ans
|
||||
case "${ans:-Y}" in
|
||||
[Yy]*)
|
||||
if command -v uv >/dev/null 2>&1; then
|
||||
uv pip install -p "$PY" -r "${SCRIPT_DIR}/requirements-gradio.txt"
|
||||
else
|
||||
"$PY" -m pip install -r "${SCRIPT_DIR}/requirements-gradio.txt"
|
||||
fi
|
||||
;;
|
||||
*) echo "Aborted — install manually: uv pip install -p .venv/bin/python -r requirements-gradio.txt"; exit 1 ;;
|
||||
esac
|
||||
fi
|
||||
|
||||
exec "$PY" "${SCRIPT_DIR}/scripts/sa3_gradio.py" "$@"
|
||||
@@ -0,0 +1,22 @@
|
||||
@echo off
|
||||
rem Windows twin of the `sa3` bash wrapper — runs scripts\sa3_tflite.py via the
|
||||
rem project-local .venv (created by install.bat) so no activation is needed.
|
||||
rem Falls back to whatever `python` is on PATH if .venv\ doesn't exist yet.
|
||||
rem
|
||||
rem Usage: sa3.bat --prompt "lofi house" --dit sm-music --decoder same-s --out a.wav
|
||||
setlocal
|
||||
set "SCRIPT_DIR=%~dp0"
|
||||
|
||||
if not exist "%SCRIPT_DIR%scripts\sa3_tflite.py" (
|
||||
echo error: scripts\sa3_tflite.py not found - repo files are missing or moved. 1>&2
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
set "PY=%SCRIPT_DIR%.venv\Scripts\python.exe"
|
||||
if not exist "%PY%" (
|
||||
echo warning: .venv\ not found - run install.bat for one-time setup. Trying `python` from PATH. 1>&2
|
||||
set "PY=python"
|
||||
)
|
||||
|
||||
"%PY%" "%SCRIPT_DIR%scripts\sa3_tflite.py" %*
|
||||
exit /b %ERRORLEVEL%
|
||||
@@ -0,0 +1,11 @@
|
||||
# sa3.ps1 - PowerShell twin of the sa3 wrapper (and of sa3.bat).
|
||||
# PowerShell users get clean Ctrl-C handling (no cmd.exe
|
||||
# "Terminate batch job (Y/N)?" prompt, which is a .bat quirk).
|
||||
$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
|
||||
$Py = Join-Path $ScriptDir ".venv\Scripts\python.exe"
|
||||
if (-not (Test-Path $Py)) {
|
||||
$Py = Join-Path $ScriptDir ".venv/bin/python" # POSIX layout (pwsh on mac/linux)
|
||||
}
|
||||
if (-not (Test-Path $Py)) { $Py = "python" }
|
||||
& $Py (Join-Path $ScriptDir "scripts/sa3_tflite.py") @args
|
||||
exit $LASTEXITCODE
|
||||
@@ -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