Prevent hard clipping across audio output backends

.github/workflows/mlx.yml

@@ -30,3 +30,10 @@ jobs:
           python scripts/lora_train_mlx.py --help > /dev/null
           python scripts/pre_encode_mlx.py --help > /dev/null
           echo "CLI smoke OK"
+      - name: Standalone subtree import smoke
+        run: |
+          bundle_dir="$(mktemp -d)/sa3_mlx"
+          cp -R optimized/mlx "$bundle_dir"
+          cp stable_audio_3/audio_output.py "$bundle_dir/scripts/audio_output.py"
+          cd "$bundle_dir"
+          python scripts/sa3_mlx.py --help > /dev/null

.github/workflows/tensorrt-decoder-rewrite.yml

@@ -0,0 +1,21 @@
+name: TensorRT decoder rewrite
+
+on:
+  pull_request:
+    paths:
+      - "optimized/tensorRT/build/**"
+      - "tests/test_tensorrt_decoder_output.py"
+      - ".github/workflows/tensorrt-decoder-rewrite.yml"
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install weight-free rewrite test dependencies
+        run: python -m pip install pytest numpy "onnx>=1.18" onnxruntime
+      - name: Test decoder output rewrite
+        run: pytest tests/test_tensorrt_decoder_output.py -q

.github/workflows/tflite-cross-platform.yml

@@ -48,11 +48,20 @@ jobs:
             optimized/tflite/scripts/install.py \
             optimized/tflite/scripts/examples.py \
             optimized/tflite/scripts/test_windows_compat.py \
-            optimized/tflite/models/defs/tflite_pipeline.py
+            optimized/tflite/models/defs/tflite_pipeline.py \
+            optimized/tflite/models/defs/wav_io.py
 
       - name: LiteRT interpreter imports
         run: python -c "from ai_edge_litert.interpreter import Interpreter"
 
+      - name: Standalone subtree import smoke
+        run: |
+          bundle_dir="$(mktemp -d)/sa3_tflite"
+          cp -R optimized/tflite "$bundle_dir"
+          cp stable_audio_3/audio_output.py "$bundle_dir/models/defs/audio_output.py"
+          cd "$bundle_dir"
+          python scripts/sa3_tflite.py --help > /dev/null
+
       - name: Windows-compat test suite
         run: python optimized/tflite/scripts/test_windows_compat.py
 

README.md

@@ -142,6 +142,11 @@ audio = model.generate(
 )
 ```
 
+Generated audio uses no-boost peak attenuation by default. Set
+`output_peak_ceiling_dbfs=-1.0` for additional headroom, or use
+`output_peak_policy="raw"` when a downstream mastering pipeline needs the
+unbounded decoded waveform.
+
 **Audio-to-Audio** — Edit an existing recording using a prompt to steer style and mood:
 
 ```python

optimized/mlx/README.md

@@ -283,6 +283,7 @@ are interchangeable in either direction.
 | `--lora`              | —        | A `.safetensors` LoRA adapter (SA3-native/underfit or PEFT) with optional `strength=S` and `steps=MIN-MAX` tokens; repeat the flag to stack adapters. Full-range adapters merge at load; step-gated ones re-merge in place at step boundaries. Pickle `.ckpt/.pt` is refused. Base must match `--dit` |
 | `--lora-strength`     | 1.0      | Default strength for adapters without their own `strength=`; 0 = bit-exact bypass, >1 amplifies |
 | `--free-models`       | on       | Progressive model freeing; `--no-free-models` keeps them resident     |
+| `--peak-ceiling-dbfs` | 0        | Sample-peak ceiling in dBFS; use `-1` for additional encoding headroom |
 | `--out`               | out.wav  | Relative → `output/<file>`; absolute → as-is. 16-bit PCM stereo @ 44.1 kHz, trimmed to exactly `--seconds` |
 | `--play`              | off      | After writing, play via `afplay`; Ctrl-C stops both processes         |
 

optimized/mlx/bootstrap.sh

@@ -38,6 +38,7 @@ DEFAULT_ARGS=(--prompt "Impending tribal, epic orchestral buildup" --dit sm-musi
 
 TAR_URL="https://github.com/$REPO_OWNER/$REPO_NAME/archive/refs/heads/$BRANCH.tar.gz"
 TAR_INNER="$REPO_NAME-$BRANCH/$SUBDIR_IN_REPO"
+SHARED_AUDIO_INNER="$REPO_NAME-$BRANCH/stable_audio_3/audio_output.py"
 
 # ── colours ─────────────────────────────────────────────────────────────────
 if [[ -t 1 ]]; then
@@ -114,11 +115,13 @@ else
         curl -fL --progress-bar "$TAR_URL" -o "$TMP_TAR"
 
         # BSD tar (macOS) extracts only paths matching the pattern.
-        tar -xz -f "$TMP_TAR" -C "$TMP_EXTRACT" "$TAR_INNER"
+        tar -xz -f "$TMP_TAR" -C "$TMP_EXTRACT" \
+            "$TAR_INNER" "$SHARED_AUDIO_INNER"
 
         SRC="$TMP_EXTRACT/$TAR_INNER"
         [[ -d "$SRC" ]] || fail "Expected '$TAR_INNER' inside the tarball but didn't find it."
         mv "$SRC" "$LOCAL_DIR"
+        mv "$TMP_EXTRACT/$SHARED_AUDIO_INNER" "$LOCAL_DIR/scripts/audio_output.py"
         ok "extracted $(find "$LOCAL_DIR" -type f | wc -l | tr -d ' ') files to ./$LOCAL_DIR"
     fi
 fi

optimized/mlx/scripts/sa3_gradio.py

@@ -65,6 +65,7 @@
 from models.defs.t5gemma_mlx import T5Gemma  # noqa: E402
 from weights import ensure_local  # noqa: E402
 from spec import render_spectrogram_png  # noqa: E402
+from wav_io import audio_to_pcm16  # noqa: E402
 
 OUTPUT_DIR = REPO / "output" / "gradio"
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
@@ -903,7 +904,12 @@ def _generate_entry(dit_name, decoder_name, prompt, negative_prompt,
         if not np.isfinite(audio_np).all():
             return None, "error: model produced non-finite audio (try a higher σmax or different seed)"
 
-        pcm = (np.clip(audio_np, -1, 1) * 32767.0).astype(np.int16).T   # (T, 2)
+        raw_peak = float(np.abs(audio_np).max()) if audio_np.size else 0.0
+        pcm = audio_to_pcm16(audio_np)
+        if raw_peak > 1.0:
+            notes.append(
+                f"output peak {raw_peak:.3f} exceeded 0 dBFS — attenuated without boosting"
+            )
         basename = verbose_basename(prompt, negative_prompt, cfg, sigma_max, seed)
         out_path = OUTPUT_DIR / f"{basename}.wav"
         _save_wav(pcm, out_path)

optimized/mlx/scripts/sa3_mlx.py

@@ -30,6 +30,7 @@
     load_conditioner_from_npz,
 )
 from models.defs.t5gemma_mlx import T5Gemma
+from wav_io import dbfs_to_amplitude, save_wav
 from weights import ensure_local, is_present
 
 SAMPLE_RATE = 44100
@@ -228,20 +229,6 @@ def _stage_peak_b(label: str | None = None) -> int:
     return b
 
 
-def save_wav(path: str, audio: np.ndarray, sample_rate: int = SAMPLE_RATE):
-    """audio: (channels, T) float32 in [-1, 1]. Writes 16-bit PCM stereo WAV."""
-    if not np.isfinite(audio).all():
-        n_bad = int((~np.isfinite(audio)).sum())
-        raise RuntimeError(f"refusing to write WAV — audio contains {n_bad} non-finite samples (NaN/Inf)")
-    audio = np.clip(audio, -1.0, 1.0)
-    pcm = (audio * 32767.0).astype(np.int16).T  # (T, channels) interleaved
-    with wave.open(path, "wb") as w:
-        w.setnchannels(audio.shape[0])
-        w.setsampwidth(2)
-        w.setframerate(sample_rate)
-        w.writeframes(pcm.tobytes())
-
-
 def read_wav(path: str) -> np.ndarray:
     """Read a WAV file. Returns (2, T) float32 in [-1, 1].
 
@@ -463,6 +450,8 @@ def main():
                          "directory (auto-created); absolute paths are used as-is. "
                          "Always written as 16-bit PCM stereo at 44.1 kHz, trimmed to "
                          "exactly --seconds. If omitted, auto-named from the prompt and seed.")
+    ap.add_argument("--peak-ceiling-dbfs", type=float, default=0.0,
+                    help="Output sample-peak ceiling in dBFS, at or below 0 (default: 0).")
     ap.add_argument("--play", action="store_true",
                     help="After writing the WAV, play it through the default output device "
                          "via the macOS `afplay` binary. Blocking — the script exits when "
@@ -471,6 +460,10 @@ def main():
     args = ap.parse_args()
     if args.steps < 1:
         ap.error(f"--steps must be ≥ 1 (got {args.steps})")
+    try:
+        dbfs_to_amplitude(args.peak_ceiling_dbfs)
+    except ValueError as exc:
+        ap.error(str(exc))
 
     # Parse --lora groups into specs (fail fast, before any model loads).
     args.lora_specs = None
@@ -839,7 +832,11 @@ def _on_step(i: int, total: int):
     requested_samples = int(round(args.seconds * SAMPLE_RATE))
     if audio_np.shape[-1] > requested_samples:
         audio_np = audio_np[..., :requested_samples]
-    save_wav(args.out, audio_np)
+    save_wav(
+        args.out,
+        audio_np,
+        peak_ceiling_dbfs=args.peak_ceiling_dbfs,
+    )
     stage("[5/5]", "Unpatch + write WAV", (time.time()-t0)*1000, peak_b=_stage_peak_b("Unpatch + WAV"))
     peak = float(np.abs(audio_np).max()); rms = float(np.sqrt((audio_np**2).mean()))
     sub(f"audio {audio_np.shape}   peak {peak:.3f}   rms {rms:.3f}")

optimized/mlx/scripts/wav_io.py

@@ -0,0 +1,26 @@
+"""Compatibility import for the repository's shared WAV helpers."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+THIS_DIR = Path(__file__).resolve().parent
+FULL_REPO_HELPER_DIR = Path(__file__).resolve().parents[3] / "stable_audio_3"
+HELPER_DIR = (
+    FULL_REPO_HELPER_DIR
+    if (FULL_REPO_HELPER_DIR / "audio_output.py").is_file()
+    else THIS_DIR
+)
+if not (HELPER_DIR / "audio_output.py").is_file():
+    raise ModuleNotFoundError(
+        "shared audio_output.py is missing; reinstall the MLX bundle or use a full checkout"
+    )
+sys.path.insert(0, str(HELPER_DIR))
+
+from audio_output import (  # noqa: E402, F401
+    audio_to_pcm16,
+    dbfs_to_amplitude,
+    protect_audio_peak,
+    save_wav,
+)

optimized/tensorRT/README.md

@@ -221,6 +221,13 @@ The full-pipeline CUDA graph eliminates per-stage Python/dispatch overhead
 — each replay completes in **literally identical wall-clock time** (zero
 variance once the graph is built).
 
+Decoder engines built before the peak-protection update expose a `pcm`
+binding with hard clipping baked into the engine. The runtime detects those
+legacy engines and warns, but clipped sample ratios cannot be recovered.
+Rebuild the decoder through `build/build.py`; updated engines expose
+`audio_unbounded` FP32 and apply no-boost attenuation at runtime before PCM
+scaling and INT16 narrowing.
+
 ### Benchmark DiT step time across L values
 
 ```bash
@@ -248,6 +255,7 @@ variance once the graph is built).
 | `--quiet`            | off         | Suppress per-stage prints + NVML probes — saves ~4 ms                          |
 | `--pinned-copy`      | on          | Pinned host buffer + non_blocking DtoH for Stage 5                             |
 | `--free-models`      | off         | Free TRT engine memory after each stage's last use                             |
+| `--peak-ceiling-dbfs` | 0          | Sample-peak ceiling in dBFS; use `-1` for additional encoding headroom         |
 | `--out`              | out.wav     | Relative → `output/<file>`; absolute → as-is. 16-bit PCM stereo @ 44.1 kHz     |
 
 ## Files
@@ -272,6 +280,7 @@ optimized/tensorRT/
 │   ├── README.md                ← how to build for a new GPU arch
 │   ├── build.py                 ← interactive menu (default entry)
 │   ├── build_from_onnx.py       ← one target → ONNX → TRT engine
+│   ├── decoder_output.py        ← removes baked decoder clipping before engine build
 │   └── build_dit_profile.py     ← DiT with custom (min, opt, max) profile shapes
 └── models/                      ← .trt engines (auto-downloaded per arch; ~8 GB)
     └── sm_<cc>/                 ← arch dir matches `nvidia-smi --query-gpu=compute_cap`
@@ -301,10 +310,11 @@ invocation per sampling step handles everything.
 - **STRONGLY_TYPED T5Gemma**: built with an FP16-mixed graph (FP32
   attention island around softmax) — fixes a BF16 numerical bug where one
   specific cross-attention output token collapsed in magnitude.
-- **PCM-baked SAME-S decoder**: the int16 narrow + transpose are folded
-  into the decoder engine itself; saves ~3 ms of post-decode CPU work.
+- **Sample-major SAME-S decoder output**: transpose is folded into the decoder
+  engine; peak protection, PCM scaling, and INT16 narrowing stay in the
+  captured runtime graph so non-finites and out-of-range ratios remain visible.
 - **Mixed precision**: DiT runs FP16-mixed (FP16 trunk + FP32 RMSNorm/RoPE
-  islands + FMHA-fused FP16 attention core), decoder int32→int16, T5Gemma
+  islands + FMHA-fused FP16 attention core), decoder FP32→int16, T5Gemma
   FP16-mixed. `--quiet` skips per-stage NVML probes for an extra ~4 ms.
 - **Auto-download**: missing engines are pulled from
   `stabilityai/stable-audio-3-optimized/tensorRT/sm_<cc>/` on first use.

optimized/tensorRT/bootstrap.sh

@@ -73,6 +73,14 @@ ok "GPU: $GPU_INFO"
 if [[ -f ./install.sh && -x ./sa3 && -d ./build && -d ./scripts ]]; then
     step "Already inside an optimized/tensorRT/ checkout — using ./ in place"
     cd "$(pwd)"   # no-op; just makes the path absolute for later messages
+    if [[ ! -f ../../stable_audio_3/audio_output.py && ! -f ./scripts/audio_output.py ]]; then
+        command -v curl >/dev/null 2>&1 || \
+            fail "curl is required to complete this standalone TensorRT checkout."
+        step "Fetching shared audio output helper for the standalone checkout"
+        curl -fL --progress-bar \
+            "https://raw.githubusercontent.com/$REPO_OWNER/$REPO_NAME/$BRANCH/stable_audio_3/audio_output.py" \
+            -o ./scripts/audio_output.py
+    fi
     [[ -x ./install.sh ]] || fail "install.sh not executable in $(pwd)."
     ok "running install.sh in $(pwd)"
     ./install.sh -y

optimized/tensorRT/build/README.md

@@ -362,7 +362,8 @@ python make_dit_fp8_smalldit.py \
 | File | Role | Flow |
 |---|---|---|
 | `build.py` | Interactive menu (default entry point) | consumer |
-| `build_from_onnx.py` | One target → download ONNX from HF + compile to TRT. **For the SA3 DiTs, pulls `dit_fp16mixed.onnx` (the pre-processed island-wrapped graph)** so the consumer just needs to invoke `STRONGLY_TYPED` compilation — no `onnx-graphsurgeon` required | consumer |
+| `build_from_onnx.py` | One target → download ONNX from HF + compile to TRT. Decoder builds remove the baked output Clip first. **For the SA3 DiTs, pulls `dit_fp16mixed.onnx` (the pre-processed island-wrapped graph)** so the consumer just needs to invoke `STRONGLY_TYPED` compilation — no `onnx-graphsurgeon` required | consumer |
+| `decoder_output.py` | Rewrites decoder ONNX outputs to remove the baked `[-1, 1]` Clip, PCM scale, and integer cast and expose sample-major `audio_unbounded` FP32; runtime applies no-boost attenuation before PCM scaling and INT16 narrowing | consumer + producer |
 | `build_dit_profile.py` | Build a DiT with custom `(min, opt, max)` profile shapes (experimental — short-form / fixed-shape variants). Operates on either ONNX flavor. | consumer |
 | `build_dit_fp16mixed.py` | **Producer-side** ONNX surgery: takes the canonical FP32 `dit.onnx`, finds RMSNorm chains + attention `Softmax` + RoPE region, wraps each in `Cast(FP32) ↔ Cast(FP16)` islands, converts non-island weights to FP16, then bounds the RoPE island before QK^T (`bound_attention_core()`, `--no-bound-attn` to skip) so the attention core runs FP16 and TRT's FMHA fuser fires — 96/96 attentions on the medium DiT, 4.3× at L=4096. Writes both the modified `dit_fp16mixed.onnx` AND the TRT engine, which **must** be `STRONGLY_TYPED` (weakly-typed + `BuilderFlag.FP16` re-casts the FP32 islands and silently degrades to naive FP16). Only re-run when the model retrains or the island recipe changes. Requires `onnx` + `onnx-graphsurgeon`. | producer |
 | `build_dit_bf16.py` | **Producer-side** shared RoPE-baker for the medium `bf16` AND `fp8` engines: precomputes RoPE's cos/sin in fp64 on the host, freezes them as fp32 constant tables (`--max-t`), rewires the 96 trig sites and lets DCE delete the runtime angle chain — so the trunk runs bf16/fp8 without the long-angle drift. Weights are never loaded (keeps the input's `.data` sidecar). Handles both external `inv_freq` (fp32 `dit.onnx`) and inline (fp8-linear ONNX). Consumer compile: `build_from_onnx.py sa3-m-bf16` / `sa3-m-fp8`. Requires `onnx`. | producer |

optimized/tensorRT/build/build_from_onnx.py

@@ -31,6 +31,7 @@
 """
 import os
 import sys
+import tempfile
 import time
 from pathlib import Path
 
@@ -40,6 +41,7 @@
 sys.path.insert(0, str(SCRIPTS_DIR.parent / "scripts"))
 
 from _arch import detect_arch, arch_dir  # noqa: E402
+from decoder_output import rewrite_decoder_onnx  # noqa: E402
 
 
 HF_REPO = "stabilityai/stable-audio-3-optimized"
@@ -99,6 +101,7 @@
         "workspace_gb": 16,
         "profile":      {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]},
         "plugin":       False,
+        "unbounded_audio": True,
     },
     "same-l-encoder": {
         "onnx_hf":      ["same-l/enc_dynamic_triton_swa.onnx"],
@@ -119,6 +122,7 @@
         "workspace_gb": 16,
         "profile":      {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]},
         "plugin":       True,
+        "unbounded_audio": True,
     },
     # SA3 DiT engines: build from the pre-processed FP16-mixed ONNX hosted on
     # HF. The producer (build_dit_fp16mixed.py) does the FP32-island surgery
@@ -333,6 +337,7 @@
         "profile":      {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]},
         "plugin":       True,
         "upcast_to_fp32": True,
+        "unbounded_audio": True,
     },
     # SAME-S FP32 decoder: the canonical ONNX is already FP32 throughout
     # (no FP16 ops to upcast). Just build STRONGLY_TYPED so the engine
@@ -345,6 +350,7 @@
         "workspace_gb": 16,
         "profile":      {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]},
         "plugin":       False,
+        "unbounded_audio": True,
     },
 }
 
@@ -484,14 +490,25 @@ def build_one(name: str) -> str:
     # 1. Pull ONNX (cached by huggingface_hub)
     onnx_path = _ensure_onnx(recipe["onnx_hf"])
     print(f"  onnx: {onnx_path}", flush=True)
+    transform_tmp = None
+    if recipe.get("upcast_to_fp32") or recipe.get("unbounded_audio"):
+        transform_tmp = tempfile.TemporaryDirectory(prefix=f"sa3-{name}-")
+        transform_dir = Path(transform_tmp.name)
 
     # 1b. Optional in-process FP16→FP32 upcast for FP32 variants of FP16-mixed
     # source ONNXes (currently only SAME-L decoder needs this — DiT FP32 reads
     # the pre-existing FP32 dit.onnx directly, SAME-S canonical ONNX is already
     # FP32 throughout).
     if recipe.get("upcast_to_fp32"):
-        upcast_path = "/tmp/_build_from_onnx_fp32_upcast.onnx"
-        onnx_path = _upcast_onnx_to_fp32(onnx_path, upcast_path)
+        upcast_path = transform_dir / "decoder_fp32.onnx"
+        onnx_path = _upcast_onnx_to_fp32(onnx_path, str(upcast_path))
+
+    # Decoder ONNXes historically baked `audio.clamp(-1, 1)`, PCM scaling,
+    # and INT32 conversion into the output tail. Expose sample-major FP32 so
+    # runtime peak protection sees non-finites and preserves sample ratios.
+    if recipe.get("unbounded_audio"):
+        unbounded_path = transform_dir / "decoder_unbounded_audio.onnx"
+        onnx_path = rewrite_decoder_onnx(onnx_path, str(unbounded_path))
 
     # 2. Optional plugin import (SAME-L only — registers samel::diff_attn_swa)
     if recipe["plugin"]:
@@ -531,7 +548,10 @@ def build_one(name: str) -> str:
                   "sm_120; verify before shipping (see build/README.md)", flush=True)
     network = builder.create_network(net_flags)
     parser = trt.OnnxParser(network, logger)
-    if not parser.parse_from_file(onnx_path):
+    parsed = parser.parse_from_file(onnx_path)
+    if transform_tmp is not None:
+        transform_tmp.cleanup()
+    if not parsed:
         for i in range(parser.num_errors):
             print(f"  parse error: {parser.get_error(i)}", flush=True)
         sys.exit(2)