From 9d8d6b58d6e0ecd6be28e793455a55e561352f8d Mon Sep 17 00:00:00 2001 From: brxs Date: Fri, 31 Jul 2026 13:39:44 -0700 Subject: [PATCH 1/3] Fix hard clipping across audio output backends --- optimized/mlx/scripts/sa3_mlx.py | 15 +- optimized/mlx/scripts/wav_io.py | 11 ++ optimized/tensorRT/README.md | 12 +- optimized/tensorRT/build/README.md | 3 +- optimized/tensorRT/build/build_from_onnx.py | 11 ++ .../build/build_same_s_dec_fp16mixed.py | 18 +- optimized/tensorRT/build/decoder_output.py | 143 +++++++++++++++ optimized/tensorRT/requirements.txt | 1 + optimized/tensorRT/scripts/pt_inference.py | 15 +- optimized/tensorRT/scripts/sa3_trt.py | 70 +++++-- optimized/tensorRT/scripts/sa3_trt_core.py | 75 ++++---- stable_audio_3/audio_output.py | 171 ++++++++++++++++++ stable_audio_3/cli.py | 4 +- stable_audio_3/inference/audio_utils.py | 6 +- stable_audio_3/interface/diffusion_cond.py | 9 +- stable_audio_3/model.py | 8 +- tests/test_audio_peak_protection.py | 138 ++++++++++++++ tests/test_cli.py | 16 ++ tests/test_tensorrt_decoder_output.py | 92 ++++++++++ 19 files changed, 729 insertions(+), 89 deletions(-) create mode 100644 optimized/mlx/scripts/wav_io.py create mode 100644 optimized/tensorRT/build/decoder_output.py create mode 100644 stable_audio_3/audio_output.py create mode 100644 tests/test_audio_peak_protection.py create mode 100644 tests/test_tensorrt_decoder_output.py diff --git a/optimized/mlx/scripts/sa3_mlx.py b/optimized/mlx/scripts/sa3_mlx.py index 94b92532..7d4a2728 100644 --- a/optimized/mlx/scripts/sa3_mlx.py +++ b/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 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]. diff --git a/optimized/mlx/scripts/wav_io.py b/optimized/mlx/scripts/wav_io.py new file mode 100644 index 00000000..b250e5b1 --- /dev/null +++ b/optimized/mlx/scripts/wav_io.py @@ -0,0 +1,11 @@ +"""Compatibility import for the repository's shared WAV helpers.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) + +from audio_output import protect_audio_peak, save_wav # noqa: E402, F401 diff --git a/optimized/tensorRT/README.md b/optimized/tensorRT/README.md index 8d552261..2a478a5b 100644 --- a/optimized/tensorRT/README.md +++ b/optimized/tensorRT/README.md @@ -221,6 +221,12 @@ 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 +`pcm_unbounded` and apply no-boost attenuation at runtime before INT16 narrowing. + ### Benchmark DiT step time across L values ```bash @@ -272,6 +278,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_/ ← arch dir matches `nvidia-smi --query-gpu=compute_cap` @@ -301,8 +308,9 @@ 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. +- **PCM-baked SAME-S decoder**: PCM scaling + transpose are folded into the + decoder engine; peak protection + INT16 narrowing stay in the captured + runtime graph so out-of-range sample ratios are preserved. - **Mixed precision**: DiT runs FP16-mixed (FP16 trunk + FP32 RMSNorm/RoPE islands + FMHA-fused FP16 attention core), decoder int32→int16, T5Gemma FP16-mixed. `--quiet` skips per-stage NVML probes for an extra ~4 ms. diff --git a/optimized/tensorRT/build/README.md b/optimized/tensorRT/build/README.md index de5b1af5..7218a630 100644 --- a/optimized/tensorRT/build/README.md +++ b/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 and expose `pcm_unbounded`; runtime applies no-boost attenuation before 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 | diff --git a/optimized/tensorRT/build/build_from_onnx.py b/optimized/tensorRT/build/build_from_onnx.py index 2c5d337e..53bd5108 100755 --- a/optimized/tensorRT/build/build_from_onnx.py +++ b/optimized/tensorRT/build/build_from_onnx.py @@ -40,6 +40,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 +100,7 @@ "workspace_gb": 16, "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": False, + "unbounded_pcm": True, }, "same-l-encoder": { "onnx_hf": ["same-l/enc_dynamic_triton_swa.onnx"], @@ -119,6 +121,7 @@ "workspace_gb": 16, "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": True, + "unbounded_pcm": 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 +336,7 @@ "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": True, "upcast_to_fp32": True, + "unbounded_pcm": 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 +349,7 @@ "workspace_gb": 16, "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": False, + "unbounded_pcm": True, }, } @@ -493,6 +498,12 @@ def build_one(name: str) -> str: upcast_path = "/tmp/_build_from_onnx_fp32_upcast.onnx" onnx_path = _upcast_onnx_to_fp32(onnx_path, upcast_path) + # Decoder ONNXes historically baked `audio.clamp(-1, 1)` into the PCM + # tail. Remove it so runtime peak protection can preserve sample ratios. + if recipe.get("unbounded_pcm"): + unbounded_path = f"/tmp/_build_from_onnx_{name}_unbounded_pcm.onnx" + onnx_path = rewrite_decoder_onnx(onnx_path, unbounded_path) + # 2. Optional plugin import (SAME-L only — registers samel::diff_attn_swa) if recipe["plugin"]: print(f" registering Triton SWA plugin...", flush=True) diff --git a/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py b/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py index e66653c0..6482da1b 100644 --- a/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py +++ b/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py @@ -28,9 +28,9 @@ Inputs/outputs: - Input: `latent` (FP32, shape [1, 256, L]) — keep FP32 -- Output: `pcm` (INT32, shape [1, T, 2]) — already produced by a Cast(to=INT32) - inside the graph. The pre-Cast chain (Clip + Mul) is in trunk-FP16; the - Cast(to=INT32) is unaffected by trunk dtype since it's an explicit dtype change. +- Output: `pcm_unbounded` (INT32, shape [1, T, 2]) — the output Clip is removed + before conversion so runtime can apply no-boost attenuation before INT16 narrowing. + The pre-Cast scale remains in the graph. Usage: python build_same_s_dec_fp16mixed.py @@ -57,6 +57,7 @@ fix_dtype_mismatches, manual_convert_to_fp16, ) +from decoder_output import remove_output_hard_clip # SAME-S decoder profile — same as the canonical BF16 engine. @@ -504,6 +505,9 @@ def _inline_tensor(t): _inline_tensor(attr.t) print(f" inlined external-data tensors") + removed_clips = remove_output_hard_clip(model) + print(f" decoder peak policy: removed {removed_clips} hard Clip") + # Find FP32 islands BEFORE stripping no-op casts (so RoPE seeds are # still in the graph). blocked_names = find_fp32_islands_same_s_dec(model, mode=mode) @@ -525,13 +529,7 @@ def _inline_tensor(t): print(f" fixing dtype mismatches with autocast insertion...") fp16_model = fix_dtype_mismatches(fp16_model, blocked_names) - # SAME-S-specific: the DecWrap postprocess tail has a Clip node - # (audio.clamp(-1, 1)) whose min/max come from Cast(to=FP32) of two - # Constants. After our conversion the Slice feeding it is FP16, but - # the Cast outputs remain FP32 — the Clip is then a heterogeneous - # op. Fix by retargeting those Casts to FP16 (the min/max values are - # -1, +1, well within FP16 range). fix_dtype_mismatches doesn't - # cover Clip in its SHARED_FLOAT_DT_OPS set. + # Retain compatibility with input graphs that have other Clip nodes. fp16_model = fix_extra_dtype_mismatches(fp16_model) print(f" saving to {output_onnx}") diff --git a/optimized/tensorRT/build/decoder_output.py b/optimized/tensorRT/build/decoder_output.py new file mode 100644 index 00000000..a052a63f --- /dev/null +++ b/optimized/tensorRT/build/decoder_output.py @@ -0,0 +1,143 @@ +"""ONNX rewrite for decoder outputs that preserves out-of-range amplitudes.""" + +from __future__ import annotations + +from pathlib import Path + + +UNBOUNDED_PCM_OUTPUT = "pcm_unbounded" + + +def remove_output_hard_clip(model) -> int: + """Remove the final audio Clip and mark the PCM output as unbounded. + + The decoder's existing scale, INT32 cast, and channel transpose stay in + the graph. Runtime code can then apply the shared no-boost attenuation + policy before narrowing to INT16. Returns the number of removed Clip nodes. + """ + import numpy as np + from onnx import helper, numpy_helper + + if any(output.name == UNBOUNDED_PCM_OUTPUT for output in model.graph.output): + return 0 + + pcm_output = next( + (output for output in model.graph.output if output.name == "pcm"), None + ) + if pcm_output is None: + raise RuntimeError("decoder ONNX has no 'pcm' graph output") + + producer_by_output = { + output: node for node in model.graph.node for output in node.output + } + queue = [(pcm_output.name, 0)] + seen = set() + clips = [] + while queue: + tensor_name, distance = queue.pop(0) + if tensor_name in seen: + continue + seen.add(tensor_name) + producer = producer_by_output.get(tensor_name) + if producer is None: + continue + if producer.op_type == "Clip": + clips.append((distance, producer)) + continue + queue.extend((input_name, distance + 1) for input_name in producer.input) + + if not clips: + raise RuntimeError( + "decoder ONNX output has no upstream Clip; cannot verify peak-policy rewrite" + ) + + _, clip = min(clips, key=lambda item: item[0]) + + initializer_by_name = { + initializer.name: initializer for initializer in model.graph.initializer + } + + def constant_scalar(tensor_name): + initializer = initializer_by_name.get(tensor_name) + if initializer is not None: + value = numpy_helper.to_array(initializer) + return float(value.reshape(-1)[0]) if value.size == 1 else None + producer = producer_by_output.get(tensor_name) + if producer is None: + return None + if producer.op_type == "Cast": + return constant_scalar(producer.input[0]) + if producer.op_type == "Constant": + value_attr = next( + ( + attribute + for attribute in producer.attribute + if attribute.name == "value" + ), + None, + ) + if value_attr is not None: + value = numpy_helper.to_array(value_attr.t) + return float(value.reshape(-1)[0]) if value.size == 1 else None + return None + + minimum = constant_scalar(clip.input[1]) if len(clip.input) > 1 else None + maximum = constant_scalar(clip.input[2]) if len(clip.input) > 2 else None + for attribute in clip.attribute: + if attribute.name == "min": + minimum = float(attribute.f) + elif attribute.name == "max": + maximum = float(attribute.f) + if minimum is None or maximum is None: + raise RuntimeError("could not resolve decoder output Clip bounds") + if not np.isclose(minimum, -1.0) or not np.isclose(maximum, 1.0): + raise RuntimeError( + f"refusing to remove unexpected decoder Clip bounds [{minimum}, {maximum}]" + ) + + unclipped_input = clip.input[0] + clipped_output = clip.output[0] + for node in model.graph.node: + for index, input_name in enumerate(node.input): + if input_name == clipped_output: + node.input[index] = unclipped_input + model.graph.node.remove(clip) + + # Give rebuilt engines an explicit binding name. Runtime can distinguish + # them from legacy `pcm` engines whose destructive Clip is already baked in. + model.graph.node.append( + helper.make_node( + "Identity", + inputs=[pcm_output.name], + outputs=[UNBOUNDED_PCM_OUTPUT], + name="ExposeUnboundedPCM", + ) + ) + pcm_output.name = UNBOUNDED_PCM_OUTPUT + return 1 + + +def rewrite_decoder_onnx(input_path: str, output_path: str) -> str: + """Write a decoder ONNX whose INT32 PCM output has no baked hard clip.""" + import onnx + + model = onnx.load(input_path, load_external_data=True) + removed = remove_output_hard_clip(model) + onnx.checker.check_model(model) + + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + onnx.save_model( + model, + str(output), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=output.name + ".data", + size_threshold=1024 * 1024, + ) + print( + f" decoder peak policy: removed {removed} hard Clip; " + f"output binding -> {UNBOUNDED_PCM_OUTPUT}", + flush=True, + ) + return str(output) diff --git a/optimized/tensorRT/requirements.txt b/optimized/tensorRT/requirements.txt index 5a4c77df..a524cf09 100644 --- a/optimized/tensorRT/requirements.txt +++ b/optimized/tensorRT/requirements.txt @@ -3,6 +3,7 @@ torch # Engines are NOT compatible across minor TRT versions (10.16 won't load 10.15 engines). tensorrt==10.15.1.29 numpy +onnx>=1.18 tokenizers huggingface-hub nvidia-ml-py # process-local VRAM tracking (replaces pynvml — same module name) diff --git a/optimized/tensorRT/scripts/pt_inference.py b/optimized/tensorRT/scripts/pt_inference.py index 3b296268..b35a9f05 100644 --- a/optimized/tensorRT/scripts/pt_inference.py +++ b/optimized/tensorRT/scripts/pt_inference.py @@ -20,6 +20,9 @@ import numpy as np import torch +PROJECT_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) +from audio_output import PCM16_CEILING, protect_audio_peak # noqa: E402 TRT_REPO = Path("/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT") SCRIPTS_DIR = TRT_REPO / "scripts" @@ -289,13 +292,15 @@ def generate(self, prompt: str, *, torch.cuda.synchronize() decode_ms = (time.time() - t0) * 1000 - # 5. FP32 audio → int16 stereo PCM (same formula as TRT pipeline). - pcm_torch = (audio_fp32.clamp(-1.0, 1.0) * 32767.0).to(torch.int16) + # 5. Trim before peak protection so discarded decoder padding + # cannot attenuate the requested clip. + actual_samples = int(round(seconds * SAMPLE_RATE)) + audio_fp32 = audio_fp32[..., :actual_samples] + audio_fp32 = protect_audio_peak(audio_fp32) + pcm_torch = (audio_fp32 * PCM16_CEILING).to(torch.int16) pcm = pcm_torch.squeeze(0).T.contiguous().cpu().numpy() - # Trim to requested seconds. - actual_samples = int(round(seconds * SAMPLE_RATE)) - pcm = pcm[:actual_samples].copy() + pcm = pcm.copy() inference_ms = (time.time() - t_total) * 1000 return pcm, { diff --git a/optimized/tensorRT/scripts/sa3_trt.py b/optimized/tensorRT/scripts/sa3_trt.py index 4a6c5699..36e7d6a5 100644 --- a/optimized/tensorRT/scripts/sa3_trt.py +++ b/optimized/tensorRT/scripts/sa3_trt.py @@ -32,7 +32,7 @@ the drift in practice. """ from __future__ import annotations -import argparse, math, os, random, sys, threading, time, wave +import argparse, math, os, random, sys, threading, time, warnings, wave from pathlib import Path import numpy as np @@ -147,6 +147,8 @@ def __init__(self, t5_runner: TRTRunner, dit: DiTRunner, dec_runner: TRTRunner, self.pcm_int32_buf = None # (1, T_lat*4096, 2) int32, device self.pcm_int16_buf = None # (T_lat*4096, 2) int16, device self.pinned_host_pcm = None # (T_lat*4096, 2) int16, pinned host + self.pinned_host_peak = None # scalar float32, pinned host + self._peak_ceiling = None self.local_add_cond_buf = None # (1, 257, L) fp32, device (kept zero) self._graph = None self._built = False @@ -223,12 +225,24 @@ def build(self, sigmas, seconds: float, sigma_max: float): dec_in_dt = self.dec_runner.in_dtype["latent"] self.decoder_in_buf = canon.torch.empty(1, IO_CHANNELS, L, dtype=dec_in_dt, device="cuda") # Auto-detect output flavor like decoder_decode does. - if "pcm" in self.dec_runner.out_dtype: + if "pcm_unbounded" in self.dec_runner.out_dtype: + self._dec_out_name = "pcm_unbounded" + pcm_dt = self.dec_runner.out_dtype[self._dec_out_name] + self.pcm_int32_buf = canon.torch.empty(1, T_full, 2, dtype=pcm_dt, device="cuda") + dec_ctx.set_tensor_address("latent", self.decoder_in_buf.data_ptr()) + dec_ctx.set_tensor_address(self._dec_out_name, self.pcm_int32_buf.data_ptr()) + elif "pcm" in self.dec_runner.out_dtype: self._dec_out_name = "pcm" - pcm_dt = self.dec_runner.out_dtype["pcm"] + warnings.warn( + "legacy TensorRT decoder engine has baked hard clipping; rebuild the " + "decoder engine to get the pcm_unbounded output and preserve peak ratios", + RuntimeWarning, + stacklevel=2, + ) + pcm_dt = self.dec_runner.out_dtype[self._dec_out_name] self.pcm_int32_buf = canon.torch.empty(1, T_full, 2, dtype=pcm_dt, device="cuda") dec_ctx.set_tensor_address("latent", self.decoder_in_buf.data_ptr()) - dec_ctx.set_tensor_address("pcm", self.pcm_int32_buf.data_ptr()) + dec_ctx.set_tensor_address(self._dec_out_name, self.pcm_int32_buf.data_ptr()) else: # Legacy fp output — graph-capture supported but the postprocess # path is different. Not the production target. @@ -248,6 +262,8 @@ def build(self, sigmas, seconds: float, sigma_max: float): # this T_lat. The captured DMA copies from device to here. self.pinned_host_pcm = canon.torch.empty(T_full, 2, dtype=canon.torch.int16, pin_memory=True) + self.pinned_host_peak = canon.torch.empty((), dtype=canon.torch.float32, + pin_memory=True) # ── Warmup: replicate canonical's warmup sequence exactly so that the # decoder context's internal state at the time of the captured call @@ -353,17 +369,38 @@ def build(self, sigmas, seconds: float, sigma_max: float): self.decoder_in_buf.copy_(self.latents_out_buf) _enqueue(self.dec_runner.context, capture_stream, "decoder (mega-graph)") # Stage 5a: narrow + cast int32 → int16 (or legacy fp32 → int16) - if self._dec_out_name == "pcm": - # Belt-and-suspenders int16 clamp. New engines (from the - # FP32 clip+scale fix in the ONNX producer) already bound - # the int32 output to ±32767, so this is a no-op for them. - # Kept for backwards-compat with any older engine still in - # use, which has BF16 trunk rounding 32767 → 32768 and - # wrapping on int16 downcast (audible clicks). - self.pcm_int16_buf.copy_(self.pcm_int32_buf[0].clamp(-32767, 32767)) + if self._dec_out_name in ("pcm_unbounded", "pcm"): + source = self.pcm_int32_buf[0, :self.requested_samples] + peak = canon.audio_peak(source).float() + protected = canon.protect_audio_peak( + source, + ceiling=canon.PCM16_CEILING, + peak=peak, + validate_nonfinite=False, + emit_warning=False, + ) + self.pcm_int16_buf[:self.requested_samples].copy_( + protected.to(canon.torch.int16) + ) + self._peak_ceiling = canon.PCM16_CEILING else: - a = self._audio_legacy_buf[0].clamp(-1.0, 1.0) * 32767.0 - self.pcm_int16_buf.copy_(a.to(canon.torch.int16).T) + source = self._audio_legacy_buf[ + 0, :, :self.requested_samples + ] + peak = canon.audio_peak(source).float() + protected = canon.protect_audio_peak( + source, + peak=peak, + validate_nonfinite=False, + emit_warning=False, + ) + self.pcm_int16_buf[:self.requested_samples].copy_( + (protected * canon.PCM16_CEILING) + .to(canon.torch.int16) + .T + ) + self._peak_ceiling = 1.0 + self.pinned_host_peak.copy_(peak, non_blocking=True) # Stage 5b: DtoH (captured non_blocking copy into pinned host) self.pinned_host_pcm[:self.requested_samples].copy_( self.pcm_int16_buf[:self.requested_samples], non_blocking=True) @@ -408,6 +445,11 @@ def run(self, input_ids_cpu, attn_mask_cpu, seed: int, seconds: float): # Replay the full pipeline. self._graph.replay() stream.synchronize() + canon.report_peak_protection( + float(self.pinned_host_peak.item()), + ceiling=self._peak_ceiling, + stacklevel=3, + ) # pinned_host_pcm has been written by the DtoH; return a view. return self.pinned_host_pcm[:self.requested_samples].numpy() diff --git a/optimized/tensorRT/scripts/sa3_trt_core.py b/optimized/tensorRT/scripts/sa3_trt_core.py index 96df0065..1c812e46 100644 --- a/optimized/tensorRT/scripts/sa3_trt_core.py +++ b/optimized/tensorRT/scripts/sa3_trt_core.py @@ -10,7 +10,7 @@ If --dit or --decoder is omitted, the script prompts the user interactively. """ from __future__ import annotations -import argparse, math, os, random, sys, termios, time, tty, wave +import argparse, math, os, random, sys, termios, time, tty, warnings, wave from pathlib import Path import numpy as np @@ -19,6 +19,16 @@ SCRIPTS = Path(__file__).resolve().parent REPO = SCRIPTS.parent sys.path.insert(0, str(SCRIPTS)) +PROJECT_ROOT = REPO.parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) + +from audio_output import ( + PCM16_CEILING, + audio_peak, + protect_audio_peak, + report_peak_protection, + save_wav, +) # torch + tensorrt are imported LAZILY in main() (after CLI parsing) so that # `sa3 --help` doesn't pay the ~5 s of import cost. The silence_fd helper is @@ -631,13 +641,14 @@ def decoder_decode(runner: TRTRunner, latents: torch.Tensor) -> torch.Tensor: Two engine flavors are supported (auto-detected by output tensor name): - - Legacy (output name "audio", fp32/bf16): shape (1, 2, L*4096) audio - in [-1, 1]. The caller is responsible for clip + scale + cast to + - Legacy (output name "audio", fp32/bf16): shape (1, 2, L*4096) audio. + The caller is responsible for peak protection + scale + cast to int16 and transposing to (T, 2) interleaved PCM. - - PCM-baked (output name "pcm", int32): shape (1, L*4096, 2) PCM - already clipped + scaled to int16 range and transposed. The caller - only needs `.to(torch.int16)` to finish the conversion. Saves ~18 ms - per inference by letting TRT fuse the postprocess tail. + - Unbounded PCM (output name "pcm_unbounded", int32): shape + (1, L*4096, 2), scaled and transposed but not hard-clipped. The caller + applies peak protection before narrowing to int16. + - Legacy PCM-baked (output name "pcm", int32): the same shape, but the + engine has a baked hard clip that cannot be undone at runtime. Both engines accept any L in [32, 4096] (odd or even); SAME-S decoder at odd L matches PT eager at cos ≥ 0.99 on in-distribution latents — no @@ -648,7 +659,18 @@ def decoder_decode(runner: TRTRunner, latents: torch.Tensor) -> torch.Tensor: """ ctx = runner.context in_dt = runner.in_dtype["latent"] - out_name = "pcm" if "pcm" in runner.out_dtype else "audio" + if "pcm_unbounded" in runner.out_dtype: + out_name = "pcm_unbounded" + elif "pcm" in runner.out_dtype: + out_name = "pcm" + warnings.warn( + "legacy TensorRT decoder engine has baked hard clipping; rebuild the " + "decoder engine to get the pcm_unbounded output and preserve peak ratios", + RuntimeWarning, + stacklevel=2, + ) + else: + out_name = "audio" out_dt = runner.out_dtype[out_name] lat = latents.to(in_dt).contiguous() ctx.set_input_shape("latent", tuple(lat.shape)) @@ -1001,20 +1023,6 @@ def sample(self, initial_noise, seed=None): # ─── WAV I/O ───────────────────────────────────────────────────────────── -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 16-bit PCM @ 44.1 kHz. Returns (2, T) float32 in [-1, 1].""" with wave.open(path, "rb") as w: @@ -1614,12 +1622,10 @@ def _on_step(i: int, total: int): t0_total = time.time() # PCM conversion. Two paths: - # - PCM-baked engines (output "pcm", int32 (1, T_full, 2)): clip + scale - # + transpose are already done inside the decoder graph. The only - # remaining work is narrow int32→int16 and trim → done as part of - # the GPU→CPU copy (one fast kernel + DtoH). - # - Legacy engines (output "audio", fp32 (1, 2, T_full)): we still do - # the old clip + scale + cast + transpose on the GPU before the copy. + # - PCM engines (int32 (1, T_full, 2)): scale + transpose are already + # done inside the decoder graph. Apply peak protection before narrowing. + # - Legacy engines (output "audio", fp32 (1, 2, T_full)): protect, + # scale, cast, and transpose on the GPU before the copy. requested_samples = int(round(args.seconds * SAMPLE_RATE)) t0 = time.time() if _pcm_baked: @@ -1627,9 +1633,7 @@ def _on_step(i: int, total: int): pcm_gpu = audio[0] # (T_full, 2) int32 if pcm_gpu.shape[0] > requested_samples: pcm_gpu = pcm_gpu[:requested_samples] - # Engine output isn't clipped — values > ±32767 wrap when cast to int16 - # (audible clicks). Clamp first. - pcm_gpu = pcm_gpu.clamp(-32767, 32767).to(torch.int16) + pcm_gpu = protect_audio_peak(pcm_gpu, ceiling=PCM16_CEILING).to(torch.int16) n = pcm_gpu.shape[0] if _pinned_pcm is not None: # Non-blocking DMA straight into the pre-allocated pinned host @@ -1641,11 +1645,12 @@ def _on_step(i: int, total: int): else: pcm = pcm_gpu.contiguous().cpu().numpy() # blocking fallback else: - # legacy fp32 (1, 2, T_full): clip + scale + cast + transpose on GPU + # legacy fp32 (1, 2, T_full): protect + scale + cast + transpose on GPU audio_gpu = audio[0] # (2, T_full) fp32 if audio_gpu.shape[-1] > requested_samples: audio_gpu = audio_gpu[..., :requested_samples] - pcm_gpu = (audio_gpu.clamp(-1.0, 1.0) * 32767.0).to(torch.int16).T.contiguous() # (T, 2) + audio_gpu = protect_audio_peak(audio_gpu) + pcm_gpu = (audio_gpu * PCM16_CEILING).to(torch.int16).T.contiguous() # (T, 2) pcm = pcm_gpu.cpu().numpy() t_gpu2cpu = (time.time() - t0) * 1000 @@ -1657,9 +1662,9 @@ def _on_step(i: int, total: int): stage("[5/5]", f"WAV → {out_display}", (time.time() - t0_total) * 1000) if _pcm_baked: - sub(f"cast int32→int16 + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") + sub(f"protect/cast int32→int16 + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") else: - sub(f"clip/cast/transpose + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") + sub(f"protect/cast/transpose + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") t_full = time.time() - t_wall_start # ── Per-stage VRAM table ── diff --git a/stable_audio_3/audio_output.py b/stable_audio_3/audio_output.py new file mode 100644 index 00000000..91f8919d --- /dev/null +++ b/stable_audio_3/audio_output.py @@ -0,0 +1,171 @@ +"""Backend-neutral peak protection and PCM16 WAV serialization.""" + +from __future__ import annotations + +import math +import warnings +import wave +from typing import Any + + +PCM_FLOAT_CEILING = 1.0 +PCM16_CEILING = 32767.0 + + +def _backend_name(audio: Any) -> str: + module = type(audio).__module__.partition(".")[0] + if module in {"numpy", "torch"}: + return module + raise TypeError( + "audio must be a numpy.ndarray or torch.Tensor, " + f"got {type(audio).__module__}.{type(audio).__qualname__}" + ) + + +def _reduce_dims(audio: Any, batch_dim: int | None) -> tuple[int, ...] | None: + if batch_dim is None: + return None + if audio.ndim == 0: + raise ValueError("batch_dim cannot be used with scalar audio") + batch_dim %= audio.ndim + return tuple(dim for dim in range(audio.ndim) if dim != batch_dim) + + +def audio_peak(audio: Any, batch_dim: int | None = None) -> Any: + """Return the absolute peak, optionally reduced independently per batch item.""" + backend = _backend_name(audio) + reduce_dims = _reduce_dims(audio, batch_dim) + if backend == "numpy": + import numpy as np + + if audio.size == 0: + return 0.0 + if reduce_dims is None: + return np.abs(audio).max() + return np.abs(audio).max(axis=reduce_dims, keepdims=True) + + if audio.numel() == 0: + return 0.0 + if reduce_dims is None: + return audio.abs().amax() + return audio.abs().amax(dim=reduce_dims, keepdim=True) + + +def report_peak_protection( + peak: float, + ceiling: float = PCM_FLOAT_CEILING, + *, + n_affected: int = 1, + stacklevel: int = 2, +) -> None: + """Raise for a non-finite peak or warn when attenuation is required.""" + if not math.isfinite(peak): + raise RuntimeError("refusing to process audio with a non-finite peak (NaN/Inf)") + if peak <= ceiling: + return + + item_label = "item" if n_affected == 1 else "items" + warnings.warn( + f"audio peak {peak:.3f} exceeds the {ceiling:.3f} PCM ceiling; " + f"applying no-boost attenuation to {n_affected} {item_label} to prevent clipping", + RuntimeWarning, + stacklevel=stacklevel, + ) + + +def protect_audio_peak( + audio: Any, + ceiling: float = PCM_FLOAT_CEILING, + batch_dim: int | None = None, + *, + peak: Any | None = None, + validate_nonfinite: bool = True, + emit_warning: bool = True, +) -> Any: + """Attenuate out-of-range audio without boosting or hard clipping it. + + NumPy arrays and Torch tensors are supported. When ``batch_dim`` is + provided, each batch item is attenuated independently. Passing both + ``validate_nonfinite=False`` and ``emit_warning=False`` avoids host-side + branching, which makes the Torch path safe to capture in a CUDA graph. + """ + if ceiling <= 0: + raise ValueError(f"ceiling must be positive, got {ceiling}") + + backend = _backend_name(audio) + is_empty = audio.size == 0 if backend == "numpy" else audio.numel() == 0 + if is_empty: + return audio + + if backend == "numpy": + import numpy as np + + if validate_nonfinite: + finite = np.isfinite(audio) + if not finite.all(): + n_bad = int((~finite).sum()) + raise RuntimeError( + f"refusing to process audio containing {n_bad} non-finite samples (NaN/Inf)" + ) + peaks = audio_peak(audio, batch_dim) if peak is None else peak + over_ceiling = peaks > ceiling + has_over_ceiling = ( + bool(np.any(over_ceiling)) if (validate_nonfinite or emit_warning) else None + ) + if emit_warning and has_over_ceiling: + report_peak_protection( + float(np.max(peaks)), + ceiling, + n_affected=int(np.count_nonzero(over_ceiling)), + stacklevel=3, + ) + if has_over_ceiling is False: + return audio + scale = np.maximum(peaks / ceiling, 1.0) + return audio / scale + + import torch + + if validate_nonfinite: + finite = torch.isfinite(audio) + if not finite.all(): + n_bad = int((~finite).sum().item()) + raise RuntimeError( + f"refusing to process audio containing {n_bad} non-finite samples (NaN/Inf)" + ) + peaks = audio_peak(audio, batch_dim) if peak is None else peak + over_ceiling = peaks > ceiling + has_over_ceiling = ( + bool(over_ceiling.any()) if (validate_nonfinite or emit_warning) else None + ) + if emit_warning and has_over_ceiling: + report_peak_protection( + float(peaks.max().item()), + ceiling, + n_affected=int(over_ceiling.sum().item()), + stacklevel=3, + ) + if has_over_ceiling is False: + return audio + scale = (peaks / ceiling).clamp(min=1.0) + return audio / scale + + +def save_wav( + path: str, + audio: Any, + sample_rate: int = 44100, +) -> None: + """Write channel-first NumPy floating-point audio as 16-bit PCM WAV.""" + if _backend_name(audio) != "numpy": + raise TypeError("save_wav expects a channel-first numpy.ndarray") + + import numpy as np + + audio = protect_audio_peak(audio) + pcm = (audio * PCM16_CEILING).astype(np.int16).T + with wave.open(path, "wb") as wav: + wav.setnchannels(audio.shape[0]) + wav.setsampwidth(2) + wav.setframerate(sample_rate) + wav.writeframes(pcm.tobytes()) diff --git a/stable_audio_3/cli.py b/stable_audio_3/cli.py index 192e4628..db6f79c0 100644 --- a/stable_audio_3/cli.py +++ b/stable_audio_3/cli.py @@ -13,6 +13,7 @@ import torchaudio from stable_audio_3 import StableAudioModel +from stable_audio_3.audio_output import protect_audio_peak def _save_output(audio: torch.Tensor, sample_rate: int, output: str, batch_size: int): @@ -22,7 +23,8 @@ def _save_output(audio: torch.Tensor, sample_rate: int, output: str, batch_size: ext = ".wav" for i in range(batch_size): path = f"{base}_{i}{ext}" if batch_size > 1 else f"{base}{ext}" - torchaudio.save(path, audio[i].cpu(), sample_rate) + output_audio = protect_audio_peak(audio[i].cpu()) + torchaudio.save(path, output_audio, sample_rate) print(f"Saved: {path}") diff --git a/stable_audio_3/inference/audio_utils.py b/stable_audio_3/inference/audio_utils.py index a375b218..f670f579 100644 --- a/stable_audio_3/inference/audio_utils.py +++ b/stable_audio_3/inference/audio_utils.py @@ -50,7 +50,7 @@ def set_audio_channels(audio, target_channels): # Add channel dim if it's missing if audio.dim() == 2: audio = audio.unsqueeze(1) - + if target_channels == 1: # Convert to mono audio = audio.mean(1, keepdim=True) @@ -63,7 +63,7 @@ def set_audio_channels(audio, target_channels): return audio def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device): - + audio = audio.to(device) if in_sr != target_sr: @@ -95,4 +95,4 @@ def __call__(self, signal): end = start + self.n_samples output = signal.new_zeros([n, self.n_samples]) output[:, :min(s, self.n_samples)] = signal[:, start:end] - return output \ No newline at end of file + return output diff --git a/stable_audio_3/interface/diffusion_cond.py b/stable_audio_3/interface/diffusion_cond.py index f365122b..bbe0c1d2 100644 --- a/stable_audio_3/interface/diffusion_cond.py +++ b/stable_audio_3/interface/diffusion_cond.py @@ -10,6 +10,7 @@ from einops import rearrange +from stable_audio_3.audio_output import PCM16_CEILING, protect_audio_peak from stable_audio_3.interface.aeiou import audio_spectrogram_image from stable_audio_3.inference.distribution_shift import LogSNRShift, FluxDistributionShift, DistributionShift, IdentityDistributionShift from stable_audio_3.models.lora import has_lora @@ -142,7 +143,10 @@ def progress_callback(callback_info): if stable_audio_3_model.model.pretransform is not None: denoised = stable_audio_3_model.model.pretransform.decode(denoised) denoised = rearrange(denoised, "b d n -> d (b n)") - denoised = denoised.clamp(-1, 1).mul(32767).to(torch.int16).cpu() + denoised = protect_audio_peak( + denoised.to(torch.float32), emit_warning=False + ) + denoised = denoised.mul(PCM16_CEILING).to(torch.int16).cpu() audio_spectrogram = audio_spectrogram_image(denoised, sample_rate=sample_rate) preview_images.append((audio_spectrogram, f"Step {current_step} sigma={sigma:.3f} logSNR={log_snr:.3f}")) @@ -227,7 +231,8 @@ def progress_callback(callback_info): # Encode the audio to WAV format audio = rearrange(audio, "b d n -> d (b n)") - audio = audio.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu() + audio = protect_audio_peak(audio.to(torch.float32)) + audio = audio.mul(PCM16_CEILING).to(torch.int16).cpu() # save as wav file torchaudio.save(output_wav, audio, sample_rate) diff --git a/stable_audio_3/model.py b/stable_audio_3/model.py index ec448112..f816f6db 100644 --- a/stable_audio_3/model.py +++ b/stable_audio_3/model.py @@ -4,7 +4,8 @@ import typing as tp from torch.nn.functional import interpolate -from stable_audio_3.inference.audio_utils import prepare_audio, numpy_audio_to_tensor +from stable_audio_3.audio_output import protect_audio_peak +from stable_audio_3.inference.audio_utils import numpy_audio_to_tensor, prepare_audio from stable_audio_3.inference.sampling import sample_diffusion from stable_audio_3.loading_utils import load_autoencoder, load_diffusion_cond from stable_audio_3.model_configs import ae_models, all_models @@ -342,7 +343,7 @@ def generate( ) if not return_latents: - result = result.to(torch.float32).clamp(-1, 1) + result = result.to(torch.float32) if not return_latents and truncate_output_to_duration: if isinstance(duration, (int, float)): @@ -358,6 +359,9 @@ def generate( "Warning: Cannot truncate output to a single duration when passing a list of different durations" ) + if not return_latents: + result = protect_audio_peak(result, batch_dim=0) + return result # --- generate() helpers --- diff --git a/tests/test_audio_peak_protection.py b/tests/test_audio_peak_protection.py new file mode 100644 index 00000000..3aa7f73b --- /dev/null +++ b/tests/test_audio_peak_protection.py @@ -0,0 +1,138 @@ +import wave +import warnings +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +from optimized.mlx.scripts.wav_io import protect_audio_peak as protect_numpy_peak +from optimized.mlx.scripts.wav_io import save_wav +from stable_audio_3.audio_output import ( + PCM16_CEILING, + audio_peak, + protect_audio_peak, +) +from stable_audio_3.model import StableAudioModel + + +def test_torch_peak_protection_attenuates_batch_items_independently(): + audio = torch.tensor( + [ + [[0.0, 0.5, 1.0], [0.0, -0.5, -1.0]], + [[0.0, 1.25, 1.75], [0.0, -1.25, -1.75]], + ], + dtype=torch.float32, + ) + + with pytest.warns(RuntimeWarning, match="no-boost attenuation"): + protected = protect_audio_peak(audio, batch_dim=0) + + assert torch.equal(protected[0], audio[0]) + assert protected[1].abs().max() == 1.0 + ratio = (protected[1, 0, 1] / protected[1, 0, 2]).item() + assert ratio == pytest.approx(1.25 / 1.75) + + +def test_torch_peak_protection_rejects_non_finite_audio(): + with pytest.raises(RuntimeError, match="1 non-finite"): + protect_audio_peak(torch.tensor([0.0, torch.nan])) + + +def test_torch_peak_protection_supports_unbounded_int32_pcm(): + pcm = torch.tensor([[0, 40959, 57342], [0, -40959, -57342]], dtype=torch.int32) + + with pytest.warns(RuntimeWarning, match="peak 57342.000"): + protected = protect_audio_peak(pcm, ceiling=PCM16_CEILING) + + narrowed = protected.to(torch.int16) + assert narrowed.abs().max() == 32767 + ratio = (narrowed[0, 1].float() / narrowed[0, 2].float()).item() + assert ratio == pytest.approx(40959 / 57342, abs=1e-4) + + +def test_torch_peak_protection_has_capture_safe_branchless_mode(): + audio = torch.tensor([0.0, 1.25, 1.75]) + peak = audio_peak(audio) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + protected = protect_audio_peak( + audio, + peak=peak, + validate_nonfinite=False, + emit_warning=False, + ) + + assert protected.max() == 1.0 + assert (protected[1] / protected[2]).item() == pytest.approx(1.25 / 1.75) + + +def test_numpy_peak_protection_does_not_boost_quiet_audio(): + audio = np.array([[0.0, 0.25, -0.75]], dtype=np.float32) + + protected = protect_numpy_peak(audio) + + assert protected is audio + + +def test_mlx_wav_serializer_attenuates_instead_of_clipping(tmp_path): + audio = np.array( + [ + [0.0, 0.5, 1.0, 1.25, 1.75], + [0.0, -0.5, -1.0, -1.25, -1.75], + ], + dtype=np.float32, + ) + output = tmp_path / "out.wav" + + with pytest.warns(RuntimeWarning, match="peak 1.750"): + save_wav(str(output), audio, 44100) + + with wave.open(str(output), "rb") as wav: + pcm = np.frombuffer(wav.readframes(wav.getnframes()), dtype=np.int16) + pcm = pcm.reshape(-1, wav.getnchannels()).T + + assert pcm[0, -1] == 32767 + assert pcm[0, -2] < pcm[0, -1] + assert pcm[0, -2] == pytest.approx(32767 * 1.25 / 1.75, abs=1) + assert pcm[1, -2] == pytest.approx(-32767 * 1.25 / 1.75, abs=1) + + +class _FakePipeline: + sample_rate = 1 + io_channels = 2 + pretransform = None + diffusion_objective = None + sampling_dist_shift = None + + def __init__(self): + self.model = torch.nn.Linear(1, 1, bias=False) + + @staticmethod + def get_conditioning_inputs(_conditioning, negative=False): + return {} + + +@pytest.mark.parametrize("discarded_tail", [10.0, float("nan")]) +def test_generate_trims_decoder_padding_before_peak_protection(discarded_tail): + model = StableAudioModel.__new__(StableAudioModel) + model.model = _FakePipeline() + model.device = "cpu" + decoded = torch.tensor( + [[[0.5, -0.5, discarded_tail], [0.25, -0.25, discarded_tail]]] + ) + + with ( + patch("stable_audio_3.model.sample_diffusion", return_value=decoded), + warnings.catch_warnings(), + ): + warnings.simplefilter("error", RuntimeWarning) + result = model.generate( + conditioning_tensors={}, + duration=2, + sample_size=3, + batch_size=1, + ) + + assert torch.equal(result, decoded[..., :2]) diff --git a/tests/test_cli.py b/tests/test_cli.py index d8650779..4dce09e1 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -174,6 +174,22 @@ def test_output_single(mock_model, mock_torchaudio_save, tmp_path): assert torch.equal(saved_tensor, mock_model.generate.return_value[0].cpu()) +def test_output_attenuates_out_of_range_audio( + mock_model, mock_torchaudio_save, tmp_path +): + mock_model.generate.return_value = torch.tensor( + [[[0.0, 1.25, 1.75], [0.0, -1.25, -1.75]]], dtype=torch.float32 + ) + + with pytest.warns(RuntimeWarning, match="peak 1.750"): + _run(["-p", "test", "-o", str(tmp_path / "out.wav")]) + + saved_audio = mock_torchaudio_save.call_args.args[1] + assert saved_audio.abs().max() == 1.0 + ratio = (saved_audio[0, 1] / saved_audio[0, 2]).item() + assert ratio == pytest.approx(1.25 / 1.75) + + def test_output_batch_naming(mock_torchaudio_save, tmp_path): model = _make_model_mock(batch=3) with patch( diff --git a/tests/test_tensorrt_decoder_output.py b/tests/test_tensorrt_decoder_output.py new file mode 100644 index 00000000..567dc266 --- /dev/null +++ b/tests/test_tensorrt_decoder_output.py @@ -0,0 +1,92 @@ +import numpy as np +import pytest + +from optimized.tensorRT.build.decoder_output import ( + UNBOUNDED_PCM_OUTPUT, + remove_output_hard_clip, + rewrite_decoder_onnx, +) + +onnx = pytest.importorskip("onnx") +TensorProto = onnx.TensorProto +helper = onnx.helper +numpy_helper = onnx.numpy_helper + + +def _decoder_tail_model(): + audio = helper.make_tensor_value_info("audio", TensorProto.FLOAT, [1, 2, 4]) + pcm = helper.make_tensor_value_info("pcm", TensorProto.INT32, [1, 4, 2]) + minimum = numpy_helper.from_array(np.array(-1.0, dtype=np.float32), "minimum") + maximum = numpy_helper.from_array(np.array(1.0, dtype=np.float32), "maximum") + scale = numpy_helper.from_array(np.array(32767.0, dtype=np.float32), "scale") + nodes = [ + helper.make_node("Clip", ["audio", "minimum", "maximum"], ["clipped"]), + helper.make_node("Mul", ["clipped", "scale"], ["scaled"]), + helper.make_node( + "Cast", ["scaled"], ["pcm_channels_first"], to=TensorProto.INT32 + ), + helper.make_node("Transpose", ["pcm_channels_first"], ["pcm"], perm=[0, 2, 1]), + ] + graph = helper.make_graph( + nodes, + "decoder_tail", + [audio], + [pcm], + initializer=[minimum, maximum, scale], + ) + return helper.make_model( + graph, + ir_version=9, + opset_imports=[helper.make_opsetid("", 17)], + ) + + +def test_decoder_rewrite_removes_clip_and_marks_unbounded_output(): + model = _decoder_tail_model() + + assert remove_output_hard_clip(model) == 1 + + assert all(node.op_type != "Clip" for node in model.graph.node) + assert model.graph.output[0].name == UNBOUNDED_PCM_OUTPUT + mul = next(node for node in model.graph.node if node.op_type == "Mul") + assert mul.input[0] == "audio" + onnx.checker.check_model(model) + + +def test_decoder_rewrite_is_idempotent(): + model = _decoder_tail_model() + remove_output_hard_clip(model) + + assert remove_output_hard_clip(model) == 0 + + +def test_decoder_rewrite_writes_loadable_onnx(tmp_path): + source = tmp_path / "decoder.onnx" + output = tmp_path / "decoder_unbounded.onnx" + onnx.save(_decoder_tail_model(), source) + + assert rewrite_decoder_onnx(str(source), str(output)) == str(output) + + rewritten = onnx.load(output) + assert rewritten.graph.output[0].name == UNBOUNDED_PCM_OUTPUT + onnx.checker.check_model(rewritten) + + +def test_decoder_rewrite_preserves_out_of_range_sample_ratios(): + onnxruntime = pytest.importorskip("onnxruntime") + model = _decoder_tail_model() + remove_output_hard_clip(model) + session = onnxruntime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + audio = np.array( + [[[0.0, 1.0, 1.25, 1.75], [0.0, -1.0, -1.25, -1.75]]], + dtype=np.float32, + ) + + pcm = session.run([UNBOUNDED_PCM_OUTPUT], {"audio": audio})[0] + + assert pcm[0, -1, 0] > 32767 + assert pcm[0, -2, 0] < pcm[0, -1, 0] + ratio = pcm[0, -2, 0] / pcm[0, -1, 0] + assert ratio == pytest.approx(1.25 / 1.75, abs=1e-4) From 550d200fc10b5448dbc0db4a10097add3de44420 Mon Sep 17 00:00:00 2001 From: brxs Date: Fri, 31 Jul 2026 16:04:53 -0700 Subject: [PATCH 2/3] Harden peak protection and expose output ceiling --- README.md | 5 + optimized/mlx/README.md | 1 + optimized/mlx/scripts/sa3_mlx.py | 14 +- optimized/mlx/scripts/wav_io.py | 6 +- optimized/tensorRT/README.md | 1 + .../build/build_same_s_dec_fp16mixed.py | 7 +- optimized/tensorRT/build/decoder_output.py | 264 +++++++++++++----- optimized/tensorRT/scripts/pt_inference.py | 12 +- optimized/tensorRT/scripts/sa3_trt.py | 52 +++- optimized/tensorRT/scripts/sa3_trt_core.py | 15 +- stable_audio_3/audio_output.py | 142 ++++++++-- stable_audio_3/cli.py | 32 ++- stable_audio_3/model.py | 16 +- tests/test_audio_peak_protection.py | 86 ++++++ tests/test_cli.py | 30 ++ tests/test_tensorrt_decoder_output.py | 62 ++++ 16 files changed, 630 insertions(+), 115 deletions(-) diff --git a/README.md b/README.md index 5e4bd024..ac556710 100644 --- a/README.md +++ b/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 diff --git a/optimized/mlx/README.md b/optimized/mlx/README.md index 4b00ff8a..51821365 100644 --- a/optimized/mlx/README.md +++ b/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/`; 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 | diff --git a/optimized/mlx/scripts/sa3_mlx.py b/optimized/mlx/scripts/sa3_mlx.py index 7d4a2728..72afeca2 100644 --- a/optimized/mlx/scripts/sa3_mlx.py +++ b/optimized/mlx/scripts/sa3_mlx.py @@ -30,7 +30,7 @@ load_conditioner_from_npz, ) from models.defs.t5gemma_mlx import T5Gemma -from wav_io import save_wav +from wav_io import dbfs_to_amplitude, save_wav from weights import ensure_local, is_present SAMPLE_RATE = 44100 @@ -450,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 " @@ -458,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 @@ -826,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}") diff --git a/optimized/mlx/scripts/wav_io.py b/optimized/mlx/scripts/wav_io.py index b250e5b1..122cf41c 100644 --- a/optimized/mlx/scripts/wav_io.py +++ b/optimized/mlx/scripts/wav_io.py @@ -8,4 +8,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[3] sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) -from audio_output import protect_audio_peak, save_wav # noqa: E402, F401 +from audio_output import ( # noqa: E402, F401 + dbfs_to_amplitude, + protect_audio_peak, + save_wav, +) diff --git a/optimized/tensorRT/README.md b/optimized/tensorRT/README.md index 2a478a5b..18a108e3 100644 --- a/optimized/tensorRT/README.md +++ b/optimized/tensorRT/README.md @@ -254,6 +254,7 @@ Rebuild the decoder through `build/build.py`; updated engines expose | `--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/`; absolute → as-is. 16-bit PCM stereo @ 44.1 kHz | ## Files diff --git a/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py b/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py index 6482da1b..be6619d7 100644 --- a/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py +++ b/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py @@ -57,7 +57,7 @@ fix_dtype_mismatches, manual_convert_to_fp16, ) -from decoder_output import remove_output_hard_clip +from decoder_output import force_unbounded_pcm_tail_fp32, remove_output_hard_clip # SAME-S decoder profile — same as the canonical BF16 engine. @@ -532,6 +532,11 @@ def _inline_tensor(t): # Retain compatibility with input graphs that have other Clip nodes. fp16_model = fix_extra_dtype_mismatches(fp16_model) + # Removing the output Clip makes the PCM scale genuinely unbounded. Keep + # that small postprocess tail in FP32 so values above ~2 cannot overflow + # FP16 before runtime peak protection sees them. + force_unbounded_pcm_tail_fp32(fp16_model) + print(f" saving to {output_onnx}") try: onnx.save(fp16_model, output_onnx) diff --git a/optimized/tensorRT/build/decoder_output.py b/optimized/tensorRT/build/decoder_output.py index a052a63f..665abf17 100644 --- a/optimized/tensorRT/build/decoder_output.py +++ b/optimized/tensorRT/build/decoder_output.py @@ -8,86 +8,148 @@ UNBOUNDED_PCM_OUTPUT = "pcm_unbounded" -def remove_output_hard_clip(model) -> int: - """Remove the final audio Clip and mark the PCM output as unbounded. +def _producer_map(model): + return {output: node for node in model.graph.node for output in node.output} - The decoder's existing scale, INT32 cast, and channel transpose stay in - the graph. Runtime code can then apply the shared no-boost attenuation - policy before narrowing to INT16. Returns the number of removed Clip nodes. - """ - import numpy as np - from onnx import helper, numpy_helper - if any(output.name == UNBOUNDED_PCM_OUTPUT for output in model.graph.output): - return 0 +def _constant_scalar(model, tensor_name: str): + from onnx import numpy_helper - pcm_output = next( - (output for output in model.graph.output if output.name == "pcm"), None + initializer_by_name = { + initializer.name: initializer for initializer in model.graph.initializer + } + initializer = initializer_by_name.get(tensor_name) + if initializer is not None: + value = numpy_helper.to_array(initializer) + return float(value.reshape(-1)[0]) if value.size == 1 else None + + producer = _producer_map(model).get(tensor_name) + if producer is None: + return None + if producer.op_type == "Cast": + return _constant_scalar(model, producer.input[0]) + if producer.op_type != "Constant": + return None + value_attr = next( + (attribute for attribute in producer.attribute if attribute.name == "value"), + None, ) - if pcm_output is None: - raise RuntimeError("decoder ONNX has no 'pcm' graph output") + if value_attr is None: + return None + value = numpy_helper.to_array(value_attr.t) + return float(value.reshape(-1)[0]) if value.size == 1 else None - producer_by_output = { - output: node for node in model.graph.node for output in node.output - } - queue = [(pcm_output.name, 0)] - seen = set() - clips = [] - while queue: - tensor_name, distance = queue.pop(0) - if tensor_name in seen: - continue - seen.add(tensor_name) - producer = producer_by_output.get(tensor_name) - if producer is None: - continue - if producer.op_type == "Clip": - clips.append((distance, producer)) - continue - queue.extend((input_name, distance + 1) for input_name in producer.input) - - if not clips: + +def _attribute_ints(node, name: str): + attribute = next((item for item in node.attribute if item.name == name), None) + return tuple(attribute.ints) if attribute is not None else None + + +def _attribute_int(node, name: str): + attribute = next((item for item in node.attribute if item.name == name), None) + return int(attribute.i) if attribute is not None else None + + +def _find_pcm_tail(model, output_name: str): + """Return the verified Transpose <- Cast(INT32) <- Mul PCM tail.""" + import numpy as np + from onnx import TensorProto + + graph_output = next( + (output for output in model.graph.output if output.name == output_name), None + ) + if graph_output is None: + raise RuntimeError(f"decoder ONNX has no {output_name!r} graph output") + + producer_by_output = _producer_map(model) + output_tensor = graph_output.name + output_producer = producer_by_output.get(output_tensor) + if output_producer is not None and output_producer.op_type == "Identity": + output_tensor = output_producer.input[0] + + transpose = producer_by_output.get(output_tensor) + if transpose is None or transpose.op_type != "Transpose": + raise RuntimeError("decoder PCM output is not produced by Transpose") + if _attribute_ints(transpose, "perm") != (0, 2, 1): raise RuntimeError( - "decoder ONNX output has no upstream Clip; cannot verify peak-policy rewrite" + "decoder PCM Transpose does not use the expected [0, 2, 1] perm" ) - _, clip = min(clips, key=lambda item: item[0]) + cast = producer_by_output.get(transpose.input[0]) + if ( + cast is None + or cast.op_type != "Cast" + or _attribute_int(cast, "to") != TensorProto.INT32 + ): + raise RuntimeError("decoder PCM Transpose is not fed by Cast(to=INT32)") - initializer_by_name = { - initializer.name: initializer for initializer in model.graph.initializer + multiply = producer_by_output.get(cast.input[0]) + if multiply is None or multiply.op_type != "Mul" or len(multiply.input) != 2: + raise RuntimeError("decoder PCM Cast is not fed by the expected scale Mul") + + scalar_inputs = [ + (index, _constant_scalar(model, name)) + for index, name in enumerate(multiply.input) + ] + scale_inputs = [ + (index, value) + for index, value in scalar_inputs + if value is not None and np.isclose(value, 32767.0, rtol=0.0, atol=1.0) + ] + if len(scale_inputs) != 1: + raise RuntimeError( + "decoder PCM Mul does not have exactly one 32767 scale input" + ) + scale_index, scale = scale_inputs[0] + signal_index = 1 - scale_index + return { + "output": graph_output, + "transpose": transpose, + "cast": cast, + "multiply": multiply, + "signal_index": signal_index, + "scale_index": scale_index, + "scale": scale, + "producer_by_output": producer_by_output, } - def constant_scalar(tensor_name): - initializer = initializer_by_name.get(tensor_name) - if initializer is not None: - value = numpy_helper.to_array(initializer) - return float(value.reshape(-1)[0]) if value.size == 1 else None - producer = producer_by_output.get(tensor_name) - if producer is None: - return None - if producer.op_type == "Cast": - return constant_scalar(producer.input[0]) - if producer.op_type == "Constant": - value_attr = next( - ( - attribute - for attribute in producer.attribute - if attribute.name == "value" - ), - None, - ) - if value_attr is not None: - value = numpy_helper.to_array(value_attr.t) - return float(value.reshape(-1)[0]) if value.size == 1 else None - return None - minimum = constant_scalar(clip.input[1]) if len(clip.input) > 1 else None - maximum = constant_scalar(clip.input[2]) if len(clip.input) > 2 else None +def _clip_bounds(model, clip) -> tuple[float | None, float | None]: + minimum = _constant_scalar(model, clip.input[1]) if len(clip.input) > 1 else None + maximum = _constant_scalar(model, clip.input[2]) if len(clip.input) > 2 else None for attribute in clip.attribute: if attribute.name == "min": minimum = float(attribute.f) elif attribute.name == "max": maximum = float(attribute.f) + return minimum, maximum + + +def remove_output_hard_clip(model) -> int: + """Remove only the verified final audio Clip and mark PCM unbounded. + + The decoder's existing scale, INT32 cast, and channel transpose stay in + the graph. Runtime code can then apply the shared no-boost attenuation + policy before narrowing to INT16. Returns the number of removed Clip nodes. + """ + import numpy as np + from onnx import helper + + if any(output.name == UNBOUNDED_PCM_OUTPUT for output in model.graph.output): + _find_pcm_tail(model, UNBOUNDED_PCM_OUTPUT) + return 0 + + tail = _find_pcm_tail(model, "pcm") + multiply = tail["multiply"] + signal_index = tail["signal_index"] + producer_by_output = tail["producer_by_output"] + clip = producer_by_output.get(multiply.input[signal_index]) + if clip is None or clip.op_type != "Clip": + raise RuntimeError( + "decoder PCM scale Mul is not fed directly by the expected output Clip" + ) + + minimum, maximum = _clip_bounds(model, clip) if minimum is None or maximum is None: raise RuntimeError("could not resolve decoder output Clip bounds") if not np.isclose(minimum, -1.0) or not np.isclose(maximum, 1.0): @@ -95,12 +157,14 @@ def constant_scalar(tensor_name): f"refusing to remove unexpected decoder Clip bounds [{minimum}, {maximum}]" ) - unclipped_input = clip.input[0] clipped_output = clip.output[0] - for node in model.graph.node: - for index, input_name in enumerate(node.input): - if input_name == clipped_output: - node.input[index] = unclipped_input + consumers = [node for node in model.graph.node if clipped_output in node.input] + if consumers != [multiply]: + raise RuntimeError( + "decoder output Clip is shared; refusing to remove a semantic graph node" + ) + + multiply.input[signal_index] = clip.input[0] model.graph.node.remove(clip) # Give rebuilt engines an explicit binding name. Runtime can distinguish @@ -108,12 +172,71 @@ def constant_scalar(tensor_name): model.graph.node.append( helper.make_node( "Identity", - inputs=[pcm_output.name], + inputs=[tail["output"].name], outputs=[UNBOUNDED_PCM_OUTPUT], name="ExposeUnboundedPCM", ) ) - pcm_output.name = UNBOUNDED_PCM_OUTPUT + tail["output"].name = UNBOUNDED_PCM_OUTPUT + return 1 + + +def force_unbounded_pcm_tail_fp32(model) -> int: + """Force unbounded audio scaling to FP32 before the INT32 cast. + + FP16 can only represent finite values through 65504, so an unbounded + ``audio * 32767`` tail would overflow for peaks just above 2. This inserts + a stable FP32 boundary and restores the exact 32767 scale after any mixed- + precision graph conversion. Returns 1 when the graph changed, else 0. + """ + import numpy as np + from onnx import TensorProto, helper, numpy_helper + + tail = _find_pcm_tail(model, UNBOUNDED_PCM_OUTPUT) + multiply = tail["multiply"] + signal_index = tail["signal_index"] + scale_index = tail["scale_index"] + producer_by_output = tail["producer_by_output"] + + scale_name = "peak_protect_pcm16_scale_fp32" + cast_name = "PeakProtectPCMInputFP32" + cast_output = "pcm_unbounded_input_fp32" + signal_input = multiply.input[signal_index] + signal_producer = producer_by_output.get(signal_input) + already_cast = ( + signal_producer is not None + and signal_producer.op_type == "Cast" + and signal_producer.name == cast_name + and _attribute_int(signal_producer, "to") == TensorProto.FLOAT + ) + if already_cast and multiply.input[scale_index] == scale_name: + return 0 + + existing_node_names = {node.name for node in model.graph.node} + existing_tensor_names = {tensor.name for tensor in model.graph.initializer} | { + output for node in model.graph.node for output in node.output + } + if cast_name in existing_node_names or cast_output in existing_tensor_names: + raise RuntimeError("decoder graph already uses reserved FP32 PCM tail names") + if scale_name in existing_tensor_names: + raise RuntimeError( + "decoder graph already uses the reserved FP32 PCM scale name" + ) + + cast = helper.make_node( + "Cast", + inputs=[signal_input], + outputs=[cast_output], + name=cast_name, + to=TensorProto.FLOAT, + ) + multiply_index = list(model.graph.node).index(multiply) + model.graph.node.insert(multiply_index, cast) + model.graph.initializer.append( + numpy_helper.from_array(np.array(32767.0, dtype=np.float32), scale_name) + ) + multiply.input[signal_index] = cast_output + multiply.input[scale_index] = scale_name return 1 @@ -123,6 +246,7 @@ def rewrite_decoder_onnx(input_path: str, output_path: str) -> str: model = onnx.load(input_path, load_external_data=True) removed = remove_output_hard_clip(model) + force_unbounded_pcm_tail_fp32(model) onnx.checker.check_model(model) output = Path(output_path) diff --git a/optimized/tensorRT/scripts/pt_inference.py b/optimized/tensorRT/scripts/pt_inference.py index b35a9f05..302ec395 100644 --- a/optimized/tensorRT/scripts/pt_inference.py +++ b/optimized/tensorRT/scripts/pt_inference.py @@ -22,7 +22,11 @@ PROJECT_ROOT = Path(__file__).resolve().parents[3] sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) -from audio_output import PCM16_CEILING, protect_audio_peak # noqa: E402 +from audio_output import ( # noqa: E402 + PCM16_CEILING, + dbfs_to_amplitude, + protect_audio_peak, +) TRT_REPO = Path("/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT") SCRIPTS_DIR = TRT_REPO / "scripts" @@ -215,6 +219,7 @@ def generate(self, prompt: str, *, cfg: float = 1.0, init_audio_path: Optional[str] = None, inpaint_range: Optional[tuple] = None, + peak_ceiling_dbfs: float = 0.0, ) -> tuple[np.ndarray, dict]: if cfg != 1.0: raise NotImplementedError("CFG not yet wired through PTInference") @@ -296,7 +301,10 @@ def generate(self, prompt: str, *, # cannot attenuate the requested clip. actual_samples = int(round(seconds * SAMPLE_RATE)) audio_fp32 = audio_fp32[..., :actual_samples] - audio_fp32 = protect_audio_peak(audio_fp32) + audio_fp32 = protect_audio_peak( + audio_fp32, + ceiling=dbfs_to_amplitude(peak_ceiling_dbfs), + ) pcm_torch = (audio_fp32 * PCM16_CEILING).to(torch.int16) pcm = pcm_torch.squeeze(0).T.contiguous().cpu().numpy() diff --git a/optimized/tensorRT/scripts/sa3_trt.py b/optimized/tensorRT/scripts/sa3_trt.py index 36e7d6a5..d6002e65 100644 --- a/optimized/tensorRT/scripts/sa3_trt.py +++ b/optimized/tensorRT/scripts/sa3_trt.py @@ -127,13 +127,15 @@ class FullPipelineGraph: """ def __init__(self, t5_runner: TRTRunner, dit: DiTRunner, dec_runner: TRTRunner, - L: int, steps: int, requested_samples: int): + L: int, steps: int, requested_samples: int, + peak_ceiling: float = 1.0): self.t5_runner = t5_runner self.dit = dit self.dec_runner = dec_runner self.L = L self.steps = steps self.requested_samples = requested_samples + self.peak_ceiling = peak_ceiling # Will be allocated in build(). self.input_ids_buf = None # (1, 256) int64, device self.attn_mask_buf = None # (1, 256) int64, device @@ -148,6 +150,7 @@ def __init__(self, t5_runner: TRTRunner, dit: DiTRunner, dec_runner: TRTRunner, self.pcm_int16_buf = None # (T_lat*4096, 2) int16, device self.pinned_host_pcm = None # (T_lat*4096, 2) int16, pinned host self.pinned_host_peak = None # scalar float32, pinned host + self.valid_sample_mask = None # (T,) bool, updated before each replay self._peak_ceiling = None self.local_add_cond_buf = None # (1, 257, L) fp32, device (kept zero) self._graph = None @@ -264,6 +267,9 @@ def build(self, sigmas, seconds: float, sigma_max: float): pin_memory=True) self.pinned_host_peak = canon.torch.empty((), dtype=canon.torch.float32, pin_memory=True) + self.valid_sample_mask = canon.torch.ones( + T_full, dtype=canon.torch.bool, device="cuda" + ) # ── Warmup: replicate canonical's warmup sequence exactly so that the # decoder context's internal state at the time of the captured call @@ -371,10 +377,15 @@ def build(self, sigmas, seconds: float, sigma_max: float): # Stage 5a: narrow + cast int32 → int16 (or legacy fp32 → int16) if self._dec_out_name in ("pcm_unbounded", "pcm"): source = self.pcm_int32_buf[0, :self.requested_samples] + canon.zero_audio_padding_( + source, + self.valid_sample_mask[:self.requested_samples], + sample_dim=0, + ) peak = canon.audio_peak(source).float() protected = canon.protect_audio_peak( source, - ceiling=canon.PCM16_CEILING, + ceiling=canon.PCM16_CEILING * self.peak_ceiling, peak=peak, validate_nonfinite=False, emit_warning=False, @@ -382,14 +393,20 @@ def build(self, sigmas, seconds: float, sigma_max: float): self.pcm_int16_buf[:self.requested_samples].copy_( protected.to(canon.torch.int16) ) - self._peak_ceiling = canon.PCM16_CEILING + self._peak_ceiling = canon.PCM16_CEILING * self.peak_ceiling else: source = self._audio_legacy_buf[ 0, :, :self.requested_samples ] + canon.zero_audio_padding_( + source, + self.valid_sample_mask[:self.requested_samples], + sample_dim=1, + ) peak = canon.audio_peak(source).float() protected = canon.protect_audio_peak( source, + ceiling=self.peak_ceiling, peak=peak, validate_nonfinite=False, emit_warning=False, @@ -399,7 +416,7 @@ def build(self, sigmas, seconds: float, sigma_max: float): .to(canon.torch.int16) .T ) - self._peak_ceiling = 1.0 + self._peak_ceiling = self.peak_ceiling self.pinned_host_peak.copy_(peak, non_blocking=True) # Stage 5b: DtoH (captured non_blocking copy into pinned host) self.pinned_host_pcm[:self.requested_samples].copy_( @@ -414,7 +431,7 @@ def run(self, input_ids_cpu, attn_mask_cpu, seed: int, seconds: float): seed: int — used to seed the in-loop noise + the initial latent randn. seconds: float — duration condition. Written into dit._sec_buf before replay. - Returns: numpy.ndarray (requested_samples, 2) int16 — a view into the + Returns: numpy.ndarray (round(seconds * SAMPLE_RATE), 2) int16 — a view into the pinned host buffer (zero-copy; valid until next run() overwrites it). """ assert self._built, "call build() first" @@ -424,6 +441,14 @@ def run(self, input_ids_cpu, attn_mask_cpu, seed: int, seconds: float): # All host writes must happen on the SAME stream the graph replays on, # otherwise the replay races the input copies. with torch.cuda.stream(stream): + actual_samples = int(round(seconds * SAMPLE_RATE)) + if not 0 <= actual_samples <= self.requested_samples: + raise ValueError( + f"requested {actual_samples} samples but graph capacity is " + f"{self.requested_samples}" + ) + self.valid_sample_mask[:actual_samples].fill_(True) + self.valid_sample_mask[actual_samples:].fill_(False) # Update T5 input buffers via HtoD copy (captured by the graph # as raw device buffers — but the HtoD copy itself runs here, # OUTSIDE the graph, before replay). @@ -451,7 +476,7 @@ def run(self, input_ids_cpu, attn_mask_cpu, seed: int, seconds: float): stacklevel=3, ) # pinned_host_pcm has been written by the DtoH; return a view. - return self.pinned_host_pcm[:self.requested_samples].numpy() + return self.pinned_host_pcm[:actual_samples].numpy() # ─── Reusable inference class (CLI + gradio share this) ───────────────── @@ -491,7 +516,8 @@ def __init__(self, dit: str, decoder: str, *, default_seconds: float = 30.0, models_dir: Path | None = None, with_encoder: bool = False, - quiet: bool = False): + quiet: bool = False, + peak_ceiling_dbfs: float = 0.0): """Load engines + build a warmup graph. Args: @@ -516,6 +542,7 @@ def __init__(self, dit: str, decoder: str, *, with_encoder: also load the audio encoder TRT engine (needed for future audio-to-audio / inpaint modes) quiet: suppress per-stage print() output from canon helpers + peak_ceiling_dbfs: sample-peak ceiling in dBFS, at or below 0 """ if dit not in DIT_CHOICES: raise ValueError(f"unknown dit={dit!r}; valid: {list(DIT_CHOICES)}") @@ -561,6 +588,7 @@ def __init__(self, dit: str, decoder: str, *, self.precision = precision self.with_encoder = with_encoder self.quiet = quiet + self.peak_ceiling = canon.dbfs_to_amplitude(peak_ceiling_dbfs) # 1. Lazy-download any missing engines (precision-aware). needed = canon.get_engine_files(dit, decoder, precision, with_encoder=with_encoder) @@ -664,7 +692,7 @@ def get_graph(self, T_lat: int, steps: int, seconds: float) -> FullPipelineGraph # serves any seconds within that T_lat's range. max_samples = T_lat * SAMPLES_PER_LATENT graph = FullPipelineGraph(self.runners["t5"], self.dit, self.runners["dec"], - T_lat, steps, max_samples) + T_lat, steps, max_samples, self.peak_ceiling) graph.build(sigmas, seconds, sigma_max) canon.torch.cuda.synchronize() @@ -785,12 +813,18 @@ def main(): ap.add_argument("--pinned-copy", action=argparse.BooleanOptionalAction, default=True) ap.add_argument("--quiet", action="store_true") ap.add_argument("--out", default="out.wav") + 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("--mega-graph", action=argparse.BooleanOptionalAction, default=True, help="Capture the entire pipeline in one CUDA graph (T5+DiT+decoder+narrow+DtoH). " "On by default. Falls back to eager path for cfg≠1.0, inpaint, or audio-to-audio.") args = ap.parse_args() if args.steps < 1: ap.error(f"--steps must be ≥ 1 (got {args.steps})") + try: + peak_ceiling = canon.dbfs_to_amplitude(args.peak_ceiling_dbfs) + except ValueError as exc: + ap.error(str(exc)) # Mute display in quiet mode — match sa3_trt's behavior. if args.quiet: @@ -945,7 +979,7 @@ def _noop_vram(label): return 0 # The graph copies the full pinned-host buffer up to requested_samples; # the WAV save reads the same buffer. mega = FullPipelineGraph(runners["t5"], dit, runners["dec"], - T_lat, args.steps, requested_samples) + T_lat, args.steps, requested_samples, peak_ceiling) mega.build(sigmas, args.seconds, sigma_max) torch.cuda.synchronize() sub(f"{dim(f'warmup + capture')} {(time.time()-t0)*1000:.0f} ms") diff --git a/optimized/tensorRT/scripts/sa3_trt_core.py b/optimized/tensorRT/scripts/sa3_trt_core.py index 1c812e46..909c5a6b 100644 --- a/optimized/tensorRT/scripts/sa3_trt_core.py +++ b/optimized/tensorRT/scripts/sa3_trt_core.py @@ -25,9 +25,11 @@ from audio_output import ( PCM16_CEILING, audio_peak, + dbfs_to_amplitude, protect_audio_peak, report_peak_protection, save_wav, + zero_audio_padding_, ) # torch + tensorrt are imported LAZILY in main() (after CLI parsing) so that @@ -1198,9 +1200,15 @@ def main(): help=f"Output WAV path. Relative paths are saved under {OUTPUT_DIR}/; " f"absolute paths are used as-is. Always 16-bit PCM stereo @ 44.1 kHz, " f"trimmed to --seconds.") + ap.add_argument("--peak-ceiling-dbfs", type=float, default=0.0, + help="Output sample-peak ceiling in dBFS, at or below 0 (default: 0).") args = ap.parse_args() if args.steps < 1: ap.error(f"--steps must be ≥ 1 (got {args.steps})") + try: + peak_ceiling = dbfs_to_amplitude(args.peak_ceiling_dbfs) + except ValueError as exc: + ap.error(str(exc)) # --quiet: stub out stage/sub prints, VRAM probes, and the sampling progress # bar so we can measure pure inference cost without instrumentation overhead. @@ -1633,7 +1641,10 @@ def _on_step(i: int, total: int): pcm_gpu = audio[0] # (T_full, 2) int32 if pcm_gpu.shape[0] > requested_samples: pcm_gpu = pcm_gpu[:requested_samples] - pcm_gpu = protect_audio_peak(pcm_gpu, ceiling=PCM16_CEILING).to(torch.int16) + pcm_gpu = protect_audio_peak( + pcm_gpu, + ceiling=PCM16_CEILING * peak_ceiling, + ).to(torch.int16) n = pcm_gpu.shape[0] if _pinned_pcm is not None: # Non-blocking DMA straight into the pre-allocated pinned host @@ -1649,7 +1660,7 @@ def _on_step(i: int, total: int): audio_gpu = audio[0] # (2, T_full) fp32 if audio_gpu.shape[-1] > requested_samples: audio_gpu = audio_gpu[..., :requested_samples] - audio_gpu = protect_audio_peak(audio_gpu) + audio_gpu = protect_audio_peak(audio_gpu, ceiling=peak_ceiling) pcm_gpu = (audio_gpu * PCM16_CEILING).to(torch.int16).T.contiguous() # (T, 2) pcm = pcm_gpu.cpu().numpy() t_gpu2cpu = (time.time() - t0) * 1000 diff --git a/stable_audio_3/audio_output.py b/stable_audio_3/audio_output.py index 91f8919d..feb34110 100644 --- a/stable_audio_3/audio_output.py +++ b/stable_audio_3/audio_output.py @@ -5,11 +5,24 @@ import math import warnings import wave -from typing import Any +from typing import Any, Literal PCM_FLOAT_CEILING = 1.0 PCM16_CEILING = 32767.0 +OutputPeakPolicy = Literal["attenuate", "raw"] + + +def dbfs_to_amplitude(ceiling_dbfs: float) -> float: + """Convert a finite, non-positive dBFS ceiling to linear amplitude.""" + if not math.isfinite(ceiling_dbfs): + raise ValueError(f"ceiling_dbfs must be finite, got {ceiling_dbfs}") + if ceiling_dbfs > 0: + raise ValueError(f"ceiling_dbfs must be <= 0, got {ceiling_dbfs}") + amplitude = 10.0 ** (ceiling_dbfs / 20.0) + if amplitude == 0.0: + raise ValueError(f"ceiling_dbfs is too small to represent, got {ceiling_dbfs}") + return amplitude def _backend_name(audio: Any) -> str: @@ -27,28 +40,68 @@ def _reduce_dims(audio: Any, batch_dim: int | None) -> tuple[int, ...] | None: return None if audio.ndim == 0: raise ValueError("batch_dim cannot be used with scalar audio") + if not -audio.ndim <= batch_dim < audio.ndim: + raise ValueError( + f"batch_dim must be in [-{audio.ndim}, {audio.ndim - 1}], got {batch_dim}" + ) batch_dim %= audio.ndim return tuple(dim for dim in range(audio.ndim) if dim != batch_dim) +def _absolute_audio(audio: Any, backend: str) -> Any: + """Take an overflow-safe absolute value, promoting integer PCM first.""" + if backend == "numpy": + import numpy as np + + if np.issubdtype(audio.dtype, np.integer): + audio = audio.astype(np.float64) + return np.abs(audio) + + import torch + + if not audio.dtype.is_floating_point and not audio.dtype.is_complex: + audio = audio.to(torch.float64 if audio.dtype == torch.int64 else torch.float32) + return audio.abs() + + +def _validate_finite(audio: Any, backend: str) -> None: + if backend == "numpy": + import numpy as np + + finite = np.isfinite(audio) + if finite.all(): + return + n_bad = int((~finite).sum()) + else: + import torch + + finite = torch.isfinite(audio) + if bool(finite.all()): + return + n_bad = int((~finite).sum().item()) + raise RuntimeError( + f"refusing to process audio containing {n_bad} non-finite samples (NaN/Inf)" + ) + + def audio_peak(audio: Any, batch_dim: int | None = None) -> Any: """Return the absolute peak, optionally reduced independently per batch item.""" backend = _backend_name(audio) reduce_dims = _reduce_dims(audio, batch_dim) if backend == "numpy": - import numpy as np - if audio.size == 0: return 0.0 + absolute = _absolute_audio(audio, backend) if reduce_dims is None: - return np.abs(audio).max() - return np.abs(audio).max(axis=reduce_dims, keepdims=True) + return absolute.max() + return absolute.max(axis=reduce_dims, keepdims=True) if audio.numel() == 0: return 0.0 + absolute = _absolute_audio(audio, backend) if reduce_dims is None: - return audio.abs().amax() - return audio.abs().amax(dim=reduce_dims, keepdim=True) + return absolute.amax() + return absolute.amax(dim=reduce_dims, keepdim=True) def report_peak_protection( @@ -101,12 +154,7 @@ def protect_audio_peak( import numpy as np if validate_nonfinite: - finite = np.isfinite(audio) - if not finite.all(): - n_bad = int((~finite).sum()) - raise RuntimeError( - f"refusing to process audio containing {n_bad} non-finite samples (NaN/Inf)" - ) + _validate_finite(audio, backend) peaks = audio_peak(audio, batch_dim) if peak is None else peak over_ceiling = peaks > ceiling has_over_ceiling = ( @@ -124,15 +172,8 @@ def protect_audio_peak( scale = np.maximum(peaks / ceiling, 1.0) return audio / scale - import torch - if validate_nonfinite: - finite = torch.isfinite(audio) - if not finite.all(): - n_bad = int((~finite).sum().item()) - raise RuntimeError( - f"refusing to process audio containing {n_bad} non-finite samples (NaN/Inf)" - ) + _validate_finite(audio, backend) peaks = audio_peak(audio, batch_dim) if peak is None else peak over_ceiling = peaks > ceiling has_over_ceiling = ( @@ -151,10 +192,67 @@ def protect_audio_peak( return audio / scale +def apply_output_peak_policy( + audio: Any, + policy: OutputPeakPolicy = "attenuate", + *, + ceiling_dbfs: float = 0.0, + batch_dim: int | None = None, + emit_warning: bool = True, +) -> Any: + """Apply the public output policy while always rejecting non-finite audio.""" + if policy not in {"attenuate", "raw"}: + raise ValueError( + f"output peak policy must be 'attenuate' or 'raw', got {policy!r}" + ) + backend = _backend_name(audio) + if policy == "raw": + _validate_finite(audio, backend) + return audio + return protect_audio_peak( + audio, + ceiling=dbfs_to_amplitude(ceiling_dbfs), + batch_dim=batch_dim, + validate_nonfinite=True, + emit_warning=emit_warning, + ) + + +def zero_audio_padding_(audio: Any, valid_sample_mask: Any, sample_dim: int) -> Any: + """Zero invalid samples in-place using a 1-D boolean validity mask.""" + backend = _backend_name(audio) + if _backend_name(valid_sample_mask) != backend: + raise TypeError("audio and valid_sample_mask must use the same backend") + if not -audio.ndim <= sample_dim < audio.ndim: + raise ValueError( + f"sample_dim must be in [-{audio.ndim}, {audio.ndim - 1}], got {sample_dim}" + ) + sample_dim %= audio.ndim + if valid_sample_mask.ndim != 1: + raise ValueError("valid_sample_mask must be one-dimensional") + if valid_sample_mask.shape[0] != audio.shape[sample_dim]: + raise ValueError( + "valid_sample_mask length must equal the audio sample dimension" + ) + + mask_shape = [1] * audio.ndim + mask_shape[sample_dim] = valid_sample_mask.shape[0] + valid_sample_mask = valid_sample_mask.reshape(mask_shape) + if backend == "numpy": + import numpy as np + + np.copyto(audio, 0, where=~valid_sample_mask) + else: + audio.masked_fill_(~valid_sample_mask, 0) + return audio + + def save_wav( path: str, audio: Any, sample_rate: int = 44100, + *, + peak_ceiling_dbfs: float = 0.0, ) -> None: """Write channel-first NumPy floating-point audio as 16-bit PCM WAV.""" if _backend_name(audio) != "numpy": @@ -162,7 +260,7 @@ def save_wav( import numpy as np - audio = protect_audio_peak(audio) + audio = protect_audio_peak(audio, ceiling=dbfs_to_amplitude(peak_ceiling_dbfs)) pcm = (audio * PCM16_CEILING).astype(np.int16).T with wave.open(path, "wb") as wav: wav.setnchannels(audio.shape[0]) diff --git a/stable_audio_3/cli.py b/stable_audio_3/cli.py index db6f79c0..1328ef86 100644 --- a/stable_audio_3/cli.py +++ b/stable_audio_3/cli.py @@ -13,17 +13,24 @@ import torchaudio from stable_audio_3 import StableAudioModel -from stable_audio_3.audio_output import protect_audio_peak +from stable_audio_3.audio_output import dbfs_to_amplitude, protect_audio_peak -def _save_output(audio: torch.Tensor, sample_rate: int, output: str, batch_size: int): +def _save_output( + audio: torch.Tensor, + sample_rate: int, + output: str, + batch_size: int, + peak_ceiling_dbfs: float = 0.0, +): """Save generated audio tensor(s) to disk.""" + peak_ceiling = dbfs_to_amplitude(peak_ceiling_dbfs) base, ext = os.path.splitext(output) if not ext: ext = ".wav" for i in range(batch_size): path = f"{base}_{i}{ext}" if batch_size > 1 else f"{base}{ext}" - output_audio = protect_audio_peak(audio[i].cpu()) + output_audio = protect_audio_peak(audio[i].cpu(), ceiling=peak_ceiling) torchaudio.save(path, output_audio, sample_rate) print(f"Saved: {path}") @@ -99,6 +106,12 @@ def main(): default="output.wav", help="Output file path (default: output.wav)", ) + parser.add_argument( + "--peak-ceiling-dbfs", + type=float, + default=0.0, + help="Output sample-peak ceiling in dBFS, at or below 0 (default: 0)", + ) # Audio-to-Audio parser.add_argument( @@ -175,6 +188,10 @@ def main(): ) args = parser.parse_args() + try: + dbfs_to_amplitude(args.peak_ceiling_dbfs) + except ValueError as exc: + parser.error(str(exc)) # --- Validate inpaint args --- if (args.inpaint_starts is None) != (args.inpaint_ends is None): @@ -287,9 +304,16 @@ def main(): inpaint_mask_start_seconds=inpaint_start, inpaint_mask_end_seconds=inpaint_end, chunked_decode=chunked_decode, + output_peak_ceiling_dbfs=args.peak_ceiling_dbfs, ) - _save_output(audio, model.model.sample_rate, args.output, batch_size) + _save_output( + audio, + model.model.sample_rate, + args.output, + batch_size, + peak_ceiling_dbfs=args.peak_ceiling_dbfs, + ) if __name__ == "__main__": diff --git a/stable_audio_3/model.py b/stable_audio_3/model.py index f816f6db..574b5565 100644 --- a/stable_audio_3/model.py +++ b/stable_audio_3/model.py @@ -4,7 +4,7 @@ import typing as tp from torch.nn.functional import interpolate -from stable_audio_3.audio_output import protect_audio_peak +from stable_audio_3.audio_output import OutputPeakPolicy, apply_output_peak_policy from stable_audio_3.inference.audio_utils import numpy_audio_to_tensor, prepare_audio from stable_audio_3.inference.sampling import sample_diffusion from stable_audio_3.loading_utils import load_autoencoder, load_diffusion_cond @@ -106,6 +106,8 @@ def generate( dist_shift=None, return_latents: bool = False, chunked_decode: tp.Optional[bool] = None, + output_peak_policy: OutputPeakPolicy = "attenuate", + output_peak_ceiling_dbfs: float = 0.0, **sampler_kwargs, ) -> torch.Tensor: """ @@ -146,6 +148,11 @@ def generate( return_latents: Whether to return the latents used for generation instead of the decoded audio. chunked_decode: Whether to decode latents in overlapping chunks to reduce peak VRAM. True forces chunked decoding on, False forces it off, None (default) uses the value set in the model config. + output_peak_policy: ``"attenuate"`` preserves sample ratios while fitting + output under the configured ceiling. ``"raw"`` returns the unbounded + decoded waveform for callers that handle mastering themselves. + output_peak_ceiling_dbfs: Sample-peak ceiling in dBFS when + ``output_peak_policy="attenuate"``. Must be finite and <= 0. **sampler_kwargs: Additional keyword arguments to pass to the sampler. """ @@ -360,7 +367,12 @@ def generate( ) if not return_latents: - result = protect_audio_peak(result, batch_dim=0) + result = apply_output_peak_policy( + result, + output_peak_policy, + ceiling_dbfs=output_peak_ceiling_dbfs, + batch_dim=0, + ) return result diff --git a/tests/test_audio_peak_protection.py b/tests/test_audio_peak_protection.py index 3aa7f73b..30af5090 100644 --- a/tests/test_audio_peak_protection.py +++ b/tests/test_audio_peak_protection.py @@ -10,8 +10,11 @@ from optimized.mlx.scripts.wav_io import save_wav from stable_audio_3.audio_output import ( PCM16_CEILING, + apply_output_peak_policy, audio_peak, + dbfs_to_amplitude, protect_audio_peak, + zero_audio_padding_, ) from stable_audio_3.model import StableAudioModel @@ -51,6 +54,23 @@ def test_torch_peak_protection_supports_unbounded_int32_pcm(): assert ratio == pytest.approx(40959 / 57342, abs=1e-4) +@pytest.mark.parametrize("backend", ["numpy", "torch"]) +def test_integer_peak_protection_handles_minimum_int32(backend): + minimum = np.iinfo(np.int32).min + if backend == "numpy": + audio = np.array([0, minimum], dtype=np.int32) + else: + audio = torch.tensor([0, minimum], dtype=torch.int32) + + protected = protect_audio_peak( + audio, + ceiling=PCM16_CEILING, + emit_warning=False, + ) + + assert float(audio_peak(protected)) == pytest.approx(PCM16_CEILING) + + def test_torch_peak_protection_has_capture_safe_branchless_mode(): audio = torch.tensor([0.0, 1.25, 1.75]) peak = audio_peak(audio) @@ -76,6 +96,54 @@ def test_numpy_peak_protection_does_not_boost_quiet_audio(): assert protected is audio +def test_output_peak_policy_can_return_validated_raw_audio(): + audio = torch.tensor([0.0, 1.75]) + + assert apply_output_peak_policy(audio, "raw") is audio + with pytest.raises(RuntimeError, match="non-finite"): + apply_output_peak_policy(torch.tensor([torch.nan]), "raw") + + +def test_dbfs_ceiling_adds_headroom_without_boosting(): + ceiling = dbfs_to_amplitude(-1.0) + loud = np.array([0.0, 2.0], dtype=np.float32) + quiet = np.array([0.0, ceiling / 2], dtype=np.float32) + + protected = apply_output_peak_policy( + loud, + ceiling_dbfs=-1.0, + emit_warning=False, + ) + + assert np.abs(protected).max() == pytest.approx(ceiling) + assert ( + apply_output_peak_policy( + quiet, + ceiling_dbfs=-1.0, + emit_warning=False, + ) + is quiet + ) + + +@pytest.mark.parametrize("ceiling_dbfs", [0.1, float("inf"), float("nan")]) +def test_dbfs_ceiling_rejects_invalid_values(ceiling_dbfs): + with pytest.raises(ValueError, match="ceiling_dbfs"): + dbfs_to_amplitude(ceiling_dbfs) + + +@pytest.mark.parametrize("sample_dim", [0, 1]) +def test_zero_audio_padding_excludes_invalid_tail(sample_dim): + audio = torch.tensor([[0.5, 0.25], [10.0, 10.0]]) + if sample_dim == 1: + audio = audio.T.contiguous() + valid = torch.tensor([True, False]) + + zero_audio_padding_(audio, valid, sample_dim=sample_dim) + + assert float(audio.abs().max()) == pytest.approx(0.5) + + def test_mlx_wav_serializer_attenuates_instead_of_clipping(tmp_path): audio = np.array( [ @@ -136,3 +204,21 @@ def test_generate_trims_decoder_padding_before_peak_protection(discarded_tail): ) assert torch.equal(result, decoded[..., :2]) + + +def test_generate_can_return_raw_unbounded_audio(): + model = StableAudioModel.__new__(StableAudioModel) + model.model = _FakePipeline() + model.device = "cpu" + decoded = torch.tensor([[[0.5, 1.75], [-0.25, -1.25]]]) + + with patch("stable_audio_3.model.sample_diffusion", return_value=decoded): + result = model.generate( + conditioning_tensors={}, + duration=2, + sample_size=2, + batch_size=1, + output_peak_policy="raw", + ) + + assert torch.equal(result, decoded) diff --git a/tests/test_cli.py b/tests/test_cli.py index 4dce09e1..80b63ed4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -129,6 +129,7 @@ def test_text_to_audio_defaults(mock_model): assert kwargs["init_audio"] is None assert kwargs["inpaint_audio"] is None assert kwargs["chunked_decode"] is None + assert kwargs["output_peak_ceiling_dbfs"] == 0.0 def test_generation_flags(mock_model): @@ -190,6 +191,35 @@ def test_output_attenuates_out_of_range_audio( assert ratio == pytest.approx(1.25 / 1.75) +def test_peak_ceiling_flag_is_shared_by_generation_and_save( + mock_model, mock_torchaudio_save, tmp_path +): + mock_model.generate.return_value = torch.tensor( + [[[0.0, 1.0], [0.0, -1.0]]], dtype=torch.float32 + ) + + with pytest.warns(RuntimeWarning, match="0.501 PCM ceiling"): + _run( + [ + "-p", + "test", + "--peak-ceiling-dbfs", + "-6", + "-o", + str(tmp_path / "out.wav"), + ] + ) + + assert mock_model.generate.call_args.kwargs["output_peak_ceiling_dbfs"] == -6.0 + saved_audio = mock_torchaudio_save.call_args.args[1] + assert saved_audio.abs().max() == pytest.approx(10 ** (-6 / 20)) + + +def test_peak_ceiling_rejects_positive_dbfs(mock_model): + with pytest.raises(SystemExit): + _run(["-p", "test", "--peak-ceiling-dbfs", "0.1"]) + + def test_output_batch_naming(mock_torchaudio_save, tmp_path): model = _make_model_mock(batch=3) with patch( diff --git a/tests/test_tensorrt_decoder_output.py b/tests/test_tensorrt_decoder_output.py index 567dc266..1cd523fb 100644 --- a/tests/test_tensorrt_decoder_output.py +++ b/tests/test_tensorrt_decoder_output.py @@ -3,6 +3,7 @@ from optimized.tensorRT.build.decoder_output import ( UNBOUNDED_PCM_OUTPUT, + force_unbounded_pcm_tail_fp32, remove_output_hard_clip, rewrite_decoder_onnx, ) @@ -69,6 +70,21 @@ def test_decoder_rewrite_writes_loadable_onnx(tmp_path): rewritten = onnx.load(output) assert rewritten.graph.output[0].name == UNBOUNDED_PCM_OUTPUT + multiply = next(node for node in rewritten.graph.node if node.op_type == "Mul") + signal_producer = next( + node for node in rewritten.graph.node if multiply.input[0] in node.output + ) + assert signal_producer.op_type == "Cast" + assert ( + next(attr.i for attr in signal_producer.attribute if attr.name == "to") + == TensorProto.FLOAT + ) + scale = next( + initializer + for initializer in rewritten.graph.initializer + if initializer.name == "peak_protect_pcm16_scale_fp32" + ) + assert numpy_helper.to_array(scale).item() == 32767.0 onnx.checker.check_model(rewritten) @@ -76,6 +92,7 @@ def test_decoder_rewrite_preserves_out_of_range_sample_ratios(): onnxruntime = pytest.importorskip("onnxruntime") model = _decoder_tail_model() remove_output_hard_clip(model) + force_unbounded_pcm_tail_fp32(model) session = onnxruntime.InferenceSession( model.SerializeToString(), providers=["CPUExecutionProvider"] ) @@ -90,3 +107,48 @@ def test_decoder_rewrite_preserves_out_of_range_sample_ratios(): assert pcm[0, -2, 0] < pcm[0, -1, 0] ratio = pcm[0, -2, 0] / pcm[0, -1, 0] assert ratio == pytest.approx(1.25 / 1.75, abs=1e-4) + + +def test_decoder_rewrite_refuses_an_unrelated_upstream_clip(): + model = _decoder_tail_model() + zero = numpy_helper.from_array(np.array(0.0, dtype=np.float32), "zero") + model.graph.initializer.append(zero) + multiply = next(node for node in model.graph.node if node.op_type == "Mul") + multiply.input[0] = "processed" + clip_index = next( + index for index, node in enumerate(model.graph.node) if node.op_type == "Clip" + ) + model.graph.node.insert( + clip_index + 1, + helper.make_node("Add", ["clipped", "zero"], ["processed"]), + ) + + with pytest.raises(RuntimeError, match="not fed directly"): + remove_output_hard_clip(model) + + +def test_fp16_mixed_tail_is_promoted_before_unbounded_scale(): + onnxruntime = pytest.importorskip("onnxruntime") + model = _decoder_tail_model() + remove_output_hard_clip(model) + model.graph.input[0].type.tensor_type.elem_type = TensorProto.FLOAT16 + scale = next(item for item in model.graph.initializer if item.name == "scale") + scale.CopyFrom( + numpy_helper.from_array(np.array(32768.0, dtype=np.float16), "scale") + ) + + assert force_unbounded_pcm_tail_fp32(model) == 1 + assert force_unbounded_pcm_tail_fp32(model) == 0 + onnx.checker.check_model(model) + session = onnxruntime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + audio = np.array( + [[[0.0, 2.0, 2.5, 3.0], [0.0, -2.0, -2.5, -3.0]]], + dtype=np.float16, + ) + + pcm = session.run([UNBOUNDED_PCM_OUTPUT], {"audio": audio})[0] + + assert pcm[0, -1, 0] == 3 * 32767 + assert pcm[0, -1, 1] == -3 * 32767 From 9d0f559a58815ce44aa4ae5a1047f79f41090257 Mon Sep 17 00:00:00 2001 From: brxs Date: Fri, 31 Jul 2026 16:37:01 -0700 Subject: [PATCH 3/3] Complete peak protection across optimized backends --- .github/workflows/mlx.yml | 7 + .../workflows/tensorrt-decoder-rewrite.yml | 21 ++ .github/workflows/tflite-cross-platform.yml | 11 +- optimized/mlx/bootstrap.sh | 5 +- optimized/mlx/scripts/sa3_gradio.py | 8 +- optimized/mlx/scripts/wav_io.py | 15 +- optimized/tensorRT/README.md | 11 +- optimized/tensorRT/bootstrap.sh | 8 + optimized/tensorRT/build/README.md | 2 +- optimized/tensorRT/build/build_from_onnx.py | 33 ++- .../build/build_same_s_dec_fp16mixed.py | 10 +- optimized/tensorRT/build/decoder_output.py | 262 ++++++++++++------ optimized/tensorRT/scripts/pt_inference.py | 6 +- optimized/tensorRT/scripts/sa3_trt.py | 51 +++- optimized/tensorRT/scripts/sa3_trt_core.py | 63 +++-- optimized/tensorRT/scripts/wav_io.py | 30 ++ optimized/tflite/README.md | 1 + optimized/tflite/bootstrap.ps1 | 3 + optimized/tflite/bootstrap.sh | 5 +- .../tflite/models/defs/tflite_pipeline.py | 17 +- optimized/tflite/models/defs/wav_io.py | 27 ++ optimized/tflite/scripts/sa3_gradio.py | 8 +- optimized/tflite/scripts/sa3_tflite.py | 13 +- .../tflite/scripts/test_windows_compat.py | 44 +++ stable_audio_3/audio_output.py | 23 +- tests/test_audio_peak_protection.py | 29 ++ tests/test_tensorrt_decoder_output.py | 166 ++++++++--- 27 files changed, 691 insertions(+), 188 deletions(-) create mode 100644 .github/workflows/tensorrt-decoder-rewrite.yml create mode 100644 optimized/tensorRT/scripts/wav_io.py create mode 100644 optimized/tflite/models/defs/wav_io.py diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 02d7135e..a981a511 100644 --- a/.github/workflows/mlx.yml +++ b/.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 diff --git a/.github/workflows/tensorrt-decoder-rewrite.yml b/.github/workflows/tensorrt-decoder-rewrite.yml new file mode 100644 index 00000000..4bfed3a5 --- /dev/null +++ b/.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 diff --git a/.github/workflows/tflite-cross-platform.yml b/.github/workflows/tflite-cross-platform.yml index 3400c02a..38d9fa6e 100644 --- a/.github/workflows/tflite-cross-platform.yml +++ b/.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 diff --git a/optimized/mlx/bootstrap.sh b/optimized/mlx/bootstrap.sh index cae93c00..0c64de89 100755 --- a/optimized/mlx/bootstrap.sh +++ b/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 diff --git a/optimized/mlx/scripts/sa3_gradio.py b/optimized/mlx/scripts/sa3_gradio.py index 6ff207cf..bc91f068 100644 --- a/optimized/mlx/scripts/sa3_gradio.py +++ b/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) diff --git a/optimized/mlx/scripts/wav_io.py b/optimized/mlx/scripts/wav_io.py index 122cf41c..a4ccedb1 100644 --- a/optimized/mlx/scripts/wav_io.py +++ b/optimized/mlx/scripts/wav_io.py @@ -5,10 +5,21 @@ import sys from pathlib import Path -PROJECT_ROOT = Path(__file__).resolve().parents[3] -sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) +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, diff --git a/optimized/tensorRT/README.md b/optimized/tensorRT/README.md index 18a108e3..93e304d8 100644 --- a/optimized/tensorRT/README.md +++ b/optimized/tensorRT/README.md @@ -225,7 +225,8 @@ 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 -`pcm_unbounded` and apply no-boost attenuation at runtime before INT16 narrowing. +`audio_unbounded` FP32 and apply no-boost attenuation at runtime before PCM +scaling and INT16 narrowing. ### Benchmark DiT step time across L values @@ -309,11 +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**: PCM scaling + transpose are folded into the - decoder engine; peak protection + INT16 narrowing stay in the captured - runtime graph so out-of-range sample ratios are preserved. +- **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_/` on first use. diff --git a/optimized/tensorRT/bootstrap.sh b/optimized/tensorRT/bootstrap.sh index 9a78f142..1f2430b9 100755 --- a/optimized/tensorRT/bootstrap.sh +++ b/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 diff --git a/optimized/tensorRT/build/README.md b/optimized/tensorRT/build/README.md index 7218a630..b811e60e 100644 --- a/optimized/tensorRT/build/README.md +++ b/optimized/tensorRT/build/README.md @@ -363,7 +363,7 @@ python make_dit_fp8_smalldit.py \ |---|---|---| | `build.py` | Interactive menu (default entry point) | 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 and expose `pcm_unbounded`; runtime applies no-boost attenuation before INT16 narrowing | consumer + producer | +| `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 | diff --git a/optimized/tensorRT/build/build_from_onnx.py b/optimized/tensorRT/build/build_from_onnx.py index 53bd5108..b707ad2e 100755 --- a/optimized/tensorRT/build/build_from_onnx.py +++ b/optimized/tensorRT/build/build_from_onnx.py @@ -31,6 +31,7 @@ """ import os import sys +import tempfile import time from pathlib import Path @@ -100,7 +101,7 @@ "workspace_gb": 16, "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": False, - "unbounded_pcm": True, + "unbounded_audio": True, }, "same-l-encoder": { "onnx_hf": ["same-l/enc_dynamic_triton_swa.onnx"], @@ -121,7 +122,7 @@ "workspace_gb": 16, "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": True, - "unbounded_pcm": 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 @@ -336,7 +337,7 @@ "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": True, "upcast_to_fp32": True, - "unbounded_pcm": 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 @@ -349,7 +350,7 @@ "workspace_gb": 16, "profile": {"latent": [(1, 256, 32), (1, 256, 1292), (1, 256, 4096)]}, "plugin": False, - "unbounded_pcm": True, + "unbounded_audio": True, }, } @@ -489,20 +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)` into the PCM - # tail. Remove it so runtime peak protection can preserve sample ratios. - if recipe.get("unbounded_pcm"): - unbounded_path = f"/tmp/_build_from_onnx_{name}_unbounded_pcm.onnx" - onnx_path = rewrite_decoder_onnx(onnx_path, unbounded_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"]: @@ -542,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) diff --git a/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py b/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py index be6619d7..e827a96e 100644 --- a/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py +++ b/optimized/tensorRT/build/build_same_s_dec_fp16mixed.py @@ -28,9 +28,9 @@ Inputs/outputs: - Input: `latent` (FP32, shape [1, 256, L]) — keep FP32 -- Output: `pcm_unbounded` (INT32, shape [1, T, 2]) — the output Clip is removed - before conversion so runtime can apply no-boost attenuation before INT16 narrowing. - The pre-Cast scale remains in the graph. +- Output: `audio_unbounded` (FP32, shape [1, T, 2]) — the output Clip, PCM + scale, and integer cast are removed so runtime can apply no-boost attenuation + before PCM scaling and INT16 narrowing. Usage: python build_same_s_dec_fp16mixed.py @@ -57,7 +57,7 @@ fix_dtype_mismatches, manual_convert_to_fp16, ) -from decoder_output import force_unbounded_pcm_tail_fp32, remove_output_hard_clip +from decoder_output import force_unbounded_audio_output_fp32, remove_output_hard_clip # SAME-S decoder profile — same as the canonical BF16 engine. @@ -535,7 +535,7 @@ def _inline_tensor(t): # Removing the output Clip makes the PCM scale genuinely unbounded. Keep # that small postprocess tail in FP32 so values above ~2 cannot overflow # FP16 before runtime peak protection sees them. - force_unbounded_pcm_tail_fp32(fp16_model) + force_unbounded_audio_output_fp32(fp16_model) print(f" saving to {output_onnx}") try: diff --git a/optimized/tensorRT/build/decoder_output.py b/optimized/tensorRT/build/decoder_output.py index 665abf17..31c4c633 100644 --- a/optimized/tensorRT/build/decoder_output.py +++ b/optimized/tensorRT/build/decoder_output.py @@ -2,10 +2,16 @@ from __future__ import annotations +import os +import tempfile +import uuid +import warnings from pathlib import Path -UNBOUNDED_PCM_OUTPUT = "pcm_unbounded" +UNBOUNDED_AUDIO_OUTPUT = "audio_unbounded" +_FP32_CAST_NAME = "PeakProtectAudioOutputFP32" +_FP32_CAST_OUTPUT = "audio_unbounded_channels_first_fp32" def _producer_map(model): @@ -64,7 +70,9 @@ def _find_pcm_tail(model, output_name: str): producer_by_output = _producer_map(model) output_tensor = graph_output.name output_producer = producer_by_output.get(output_tensor) + output_identity = None if output_producer is not None and output_producer.op_type == "Identity": + output_identity = output_producer output_tensor = output_producer.input[0] transpose = producer_by_output.get(output_tensor) @@ -104,6 +112,7 @@ def _find_pcm_tail(model, output_name: str): signal_index = 1 - scale_index return { "output": graph_output, + "output_identity": output_identity, "transpose": transpose, "cast": cast, "multiply": multiply, @@ -114,6 +123,40 @@ def _find_pcm_tail(model, output_name: str): } +def _find_unbounded_audio_tail(model): + """Return the verified sample-major floating-point output Transpose.""" + graph_output = next( + ( + output + for output in model.graph.output + if output.name == UNBOUNDED_AUDIO_OUTPUT + ), + None, + ) + if graph_output is None: + raise RuntimeError( + f"decoder ONNX has no {UNBOUNDED_AUDIO_OUTPUT!r} graph output" + ) + + producer_by_output = _producer_map(model) + output_tensor = graph_output.name + output_producer = producer_by_output.get(output_tensor) + if output_producer is not None and output_producer.op_type == "Identity": + output_tensor = output_producer.input[0] + transpose = producer_by_output.get(output_tensor) + if transpose is None or transpose.op_type != "Transpose": + raise RuntimeError("unbounded decoder audio is not produced by Transpose") + if _attribute_ints(transpose, "perm") != (0, 2, 1): + raise RuntimeError( + "unbounded decoder audio Transpose does not use the expected [0, 2, 1] perm" + ) + return { + "output": graph_output, + "transpose": transpose, + "producer_by_output": producer_by_output, + } + + def _clip_bounds(model, clip) -> tuple[float | None, float | None]: minimum = _constant_scalar(model, clip.input[1]) if len(clip.input) > 1 else None maximum = _constant_scalar(model, clip.input[2]) if len(clip.input) > 2 else None @@ -126,17 +169,17 @@ def _clip_bounds(model, clip) -> tuple[float | None, float | None]: def remove_output_hard_clip(model) -> int: - """Remove only the verified final audio Clip and mark PCM unbounded. + """Remove only the verified final Clip and expose float sample-major audio. - The decoder's existing scale, INT32 cast, and channel transpose stay in - the graph. Runtime code can then apply the shared no-boost attenuation - policy before narrowing to INT16. Returns the number of removed Clip nodes. + The destructive scale and integer cast are removed with the Clip. Runtime + code applies the shared no-boost attenuation policy to floating-point audio, + then scales and narrows to INT16. Returns the number of removed Clip nodes. """ import numpy as np - from onnx import helper + from onnx import TensorProto, helper - if any(output.name == UNBOUNDED_PCM_OUTPUT for output in model.graph.output): - _find_pcm_tail(model, UNBOUNDED_PCM_OUTPUT) + if any(output.name == UNBOUNDED_AUDIO_OUTPUT for output in model.graph.output): + _find_unbounded_audio_tail(model) return 0 tail = _find_pcm_tail(model, "pcm") @@ -164,104 +207,167 @@ def remove_output_hard_clip(model) -> int: "decoder output Clip is shared; refusing to remove a semantic graph node" ) - multiply.input[signal_index] = clip.input[0] - model.graph.node.remove(clip) + multiply_consumers = [ + node for node in model.graph.node if multiply.output[0] in node.input + ] + if multiply_consumers != [tail["cast"]]: + raise RuntimeError( + "decoder PCM scale is shared; refusing to remove a semantic graph node" + ) + cast_consumers = [ + node for node in model.graph.node if tail["cast"].output[0] in node.input + ] + if cast_consumers != [tail["transpose"]]: + raise RuntimeError( + "decoder PCM Cast is shared; refusing to remove a semantic graph node" + ) - # Give rebuilt engines an explicit binding name. Runtime can distinguish - # them from legacy `pcm` engines whose destructive Clip is already baked in. - model.graph.node.append( + existing_node_names = {node.name for node in model.graph.node} + existing_tensor_names = {tensor.name for tensor in model.graph.initializer} | { + output for node in model.graph.node for output in node.output + } + if ( + _FP32_CAST_NAME in existing_node_names + or _FP32_CAST_OUTPUT in existing_tensor_names + ): + raise RuntimeError("decoder graph already uses reserved unbounded-audio names") + + for node in (clip, multiply, tail["cast"]): + model.graph.node.remove(node) + transpose = tail["transpose"] + transpose_index = list(model.graph.node).index(transpose) + model.graph.node.insert( + transpose_index, helper.make_node( - "Identity", - inputs=[tail["output"].name], - outputs=[UNBOUNDED_PCM_OUTPUT], - name="ExposeUnboundedPCM", - ) + "Cast", + inputs=[clip.input[0]], + outputs=[_FP32_CAST_OUTPUT], + name=_FP32_CAST_NAME, + to=TensorProto.FLOAT, + ), ) - tail["output"].name = UNBOUNDED_PCM_OUTPUT - return 1 + transpose.input[0] = _FP32_CAST_OUTPUT + terminal = ( + tail["output_identity"] if tail["output_identity"] is not None else transpose + ) + terminal.output[0] = UNBOUNDED_AUDIO_OUTPUT + tail["output"].name = UNBOUNDED_AUDIO_OUTPUT + tail["output"].type.tensor_type.elem_type = TensorProto.FLOAT + return 1 -def force_unbounded_pcm_tail_fp32(model) -> int: - """Force unbounded audio scaling to FP32 before the INT32 cast. - - FP16 can only represent finite values through 65504, so an unbounded - ``audio * 32767`` tail would overflow for peaks just above 2. This inserts - a stable FP32 boundary and restores the exact 32767 scale after any mixed- - precision graph conversion. Returns 1 when the graph changed, else 0. - """ - import numpy as np - from onnx import TensorProto, helper, numpy_helper - tail = _find_pcm_tail(model, UNBOUNDED_PCM_OUTPUT) - multiply = tail["multiply"] - signal_index = tail["signal_index"] - scale_index = tail["scale_index"] - producer_by_output = tail["producer_by_output"] +def force_unbounded_audio_output_fp32(model) -> int: + """Restore the explicit FP32 output boundary after mixed-precision conversion.""" + from onnx import TensorProto, helper - scale_name = "peak_protect_pcm16_scale_fp32" - cast_name = "PeakProtectPCMInputFP32" - cast_output = "pcm_unbounded_input_fp32" - signal_input = multiply.input[signal_index] - signal_producer = producer_by_output.get(signal_input) - already_cast = ( - signal_producer is not None - and signal_producer.op_type == "Cast" - and signal_producer.name == cast_name - and _attribute_int(signal_producer, "to") == TensorProto.FLOAT - ) - if already_cast and multiply.input[scale_index] == scale_name: - return 0 + tail = _find_unbounded_audio_tail(model) + transpose = tail["transpose"] + signal_input = transpose.input[0] + producer = tail["producer_by_output"].get(signal_input) + if producer is not None and producer.name == _FP32_CAST_NAME: + if producer.op_type != "Cast": + raise RuntimeError("reserved unbounded-audio node is not a Cast") + to_attribute = next( + (item for item in producer.attribute if item.name == "to"), None + ) + if to_attribute is None: + raise RuntimeError("unbounded-audio Cast has no target dtype") + if to_attribute.i == TensorProto.FLOAT: + return 0 + to_attribute.i = TensorProto.FLOAT + tail["output"].type.tensor_type.elem_type = TensorProto.FLOAT + return 1 existing_node_names = {node.name for node in model.graph.node} existing_tensor_names = {tensor.name for tensor in model.graph.initializer} | { output for node in model.graph.node for output in node.output } - if cast_name in existing_node_names or cast_output in existing_tensor_names: - raise RuntimeError("decoder graph already uses reserved FP32 PCM tail names") - if scale_name in existing_tensor_names: - raise RuntimeError( - "decoder graph already uses the reserved FP32 PCM scale name" - ) + if ( + _FP32_CAST_NAME in existing_node_names + or _FP32_CAST_OUTPUT in existing_tensor_names + ): + raise RuntimeError("decoder graph already uses reserved unbounded-audio names") - cast = helper.make_node( - "Cast", - inputs=[signal_input], - outputs=[cast_output], - name=cast_name, - to=TensorProto.FLOAT, - ) - multiply_index = list(model.graph.node).index(multiply) - model.graph.node.insert(multiply_index, cast) - model.graph.initializer.append( - numpy_helper.from_array(np.array(32767.0, dtype=np.float32), scale_name) + transpose_index = list(model.graph.node).index(transpose) + model.graph.node.insert( + transpose_index, + helper.make_node( + "Cast", + inputs=[signal_input], + outputs=[_FP32_CAST_OUTPUT], + name=_FP32_CAST_NAME, + to=TensorProto.FLOAT, + ), ) - multiply.input[signal_index] = cast_output - multiply.input[scale_index] = scale_name + transpose.input[0] = _FP32_CAST_OUTPUT + tail["output"].type.tensor_type.elem_type = TensorProto.FLOAT return 1 def rewrite_decoder_onnx(input_path: str, output_path: str) -> str: - """Write a decoder ONNX whose INT32 PCM output has no baked hard clip.""" + """Write a decoder ONNX exposing sample-major FP32 audio without clipping.""" import onnx model = onnx.load(input_path, load_external_data=True) removed = remove_output_hard_clip(model) - force_unbounded_pcm_tail_fp32(model) + force_unbounded_audio_output_fp32(model) onnx.checker.check_model(model) output = Path(output_path) output.parent.mkdir(parents=True, exist_ok=True) - onnx.save_model( - model, - str(output), - save_as_external_data=True, - all_tensors_to_one_file=True, - location=output.name + ".data", - size_threshold=1024 * 1024, - ) + sidecar_prefix = output.name + ".data" + old_sidecars = set() + if output.exists(): + old_model = onnx.load(str(output), load_external_data=False) + for initializer in old_model.graph.initializer: + for entry in initializer.external_data: + if entry.key != "location": + continue + location = entry.value + if location == sidecar_prefix or location.startswith( + sidecar_prefix + "." + ): + old_sidecars.add(output.parent / location) + + sidecar_name = f"{sidecar_prefix}.{uuid.uuid4().hex}" + sidecar = output.parent / sidecar_name + with tempfile.TemporaryDirectory( + prefix=f".{output.name}-", dir=output.parent + ) as staging_dir: + staged_output = Path(staging_dir) / output.name + staged_sidecar = Path(staging_dir) / sidecar_name + onnx.save_model( + model, + str(staged_output), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=sidecar_name, + size_threshold=1024 * 1024, + ) + sidecar_published = False + if staged_sidecar.exists(): + os.replace(staged_sidecar, sidecar) + sidecar_published = True + try: + os.replace(staged_output, output) + except BaseException: + if sidecar_published: + sidecar.unlink(missing_ok=True) + raise + + for old_sidecar in old_sidecars - {sidecar}: + try: + old_sidecar.unlink(missing_ok=True) + except OSError as exc: + warnings.warn( + f"could not remove superseded ONNX sidecar {old_sidecar}: {exc}", + RuntimeWarning, + stacklevel=2, + ) print( f" decoder peak policy: removed {removed} hard Clip; " - f"output binding -> {UNBOUNDED_PCM_OUTPUT}", + f"output binding -> {UNBOUNDED_AUDIO_OUTPUT}", flush=True, ) return str(output) diff --git a/optimized/tensorRT/scripts/pt_inference.py b/optimized/tensorRT/scripts/pt_inference.py index 302ec395..49b89035 100644 --- a/optimized/tensorRT/scripts/pt_inference.py +++ b/optimized/tensorRT/scripts/pt_inference.py @@ -20,9 +20,9 @@ import numpy as np import torch -PROJECT_ROOT = Path(__file__).resolve().parents[3] -sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) -from audio_output import ( # noqa: E402 +LOCAL_SCRIPTS_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(LOCAL_SCRIPTS_DIR)) +from wav_io import ( # noqa: E402 PCM16_CEILING, dbfs_to_amplitude, protect_audio_peak, diff --git a/optimized/tensorRT/scripts/sa3_trt.py b/optimized/tensorRT/scripts/sa3_trt.py index d6002e65..d194a886 100644 --- a/optimized/tensorRT/scripts/sa3_trt.py +++ b/optimized/tensorRT/scripts/sa3_trt.py @@ -109,9 +109,9 @@ class FullPipelineGraph: dit_engine.execute_async_v3 ... pingpong math reads dit._vel_buf / dit._x_buf / noise_bufs[i] ... 5. decoder_in_buf <- final_latents_buf - 6. decoder engine reads decoder_in_buf, writes pcm_int32_buf - (1, T_lat*4096, 2) int32 - 7. pcm_int16_buf <- pcm_int32_buf.to(int16) (narrow cast) + 6. decoder engine reads decoder_in_buf, writes audio_unbounded_buf + (1, T_lat*4096, 2) fp32 + 7. peak-protect, scale, and narrow into pcm_int16_buf 8. pinned_host_pcm.copy_(pcm_int16_buf[:requested_samples], non_blocking) Per-inference (replay path): @@ -146,6 +146,7 @@ def __init__(self, t5_runner: TRTRunner, dit: DiTRunner, dec_runner: TRTRunner, self._sigma_next_bufs = None self.latents_out_buf = None # (1, 256, L) fp32, device — final latent self.decoder_in_buf = None # (1, 256, L) fp32, device — what dec reads + self.audio_unbounded_buf = None # (1, T_lat*4096, 2) fp32, device self.pcm_int32_buf = None # (1, T_lat*4096, 2) int32, device self.pcm_int16_buf = None # (T_lat*4096, 2) int16, device self.pinned_host_pcm = None # (T_lat*4096, 2) int16, pinned host @@ -228,17 +229,21 @@ def build(self, sigmas, seconds: float, sigma_max: float): dec_in_dt = self.dec_runner.in_dtype["latent"] self.decoder_in_buf = canon.torch.empty(1, IO_CHANNELS, L, dtype=dec_in_dt, device="cuda") # Auto-detect output flavor like decoder_decode does. - if "pcm_unbounded" in self.dec_runner.out_dtype: - self._dec_out_name = "pcm_unbounded" - pcm_dt = self.dec_runner.out_dtype[self._dec_out_name] - self.pcm_int32_buf = canon.torch.empty(1, T_full, 2, dtype=pcm_dt, device="cuda") + if "audio_unbounded" in self.dec_runner.out_dtype: + self._dec_out_name = "audio_unbounded" + audio_dt = self.dec_runner.out_dtype[self._dec_out_name] + self.audio_unbounded_buf = canon.torch.empty( + 1, T_full, 2, dtype=audio_dt, device="cuda" + ) dec_ctx.set_tensor_address("latent", self.decoder_in_buf.data_ptr()) - dec_ctx.set_tensor_address(self._dec_out_name, self.pcm_int32_buf.data_ptr()) + dec_ctx.set_tensor_address( + self._dec_out_name, self.audio_unbounded_buf.data_ptr() + ) elif "pcm" in self.dec_runner.out_dtype: self._dec_out_name = "pcm" warnings.warn( "legacy TensorRT decoder engine has baked hard clipping; rebuild the " - "decoder engine to get the pcm_unbounded output and preserve peak ratios", + "decoder engine to get the audio_unbounded output and preserve peak ratios", RuntimeWarning, stacklevel=2, ) @@ -374,8 +379,29 @@ def build(self, sigmas, seconds: float, sigma_max: float): # Stage 4: Decoder self.decoder_in_buf.copy_(self.latents_out_buf) _enqueue(self.dec_runner.context, capture_stream, "decoder (mega-graph)") - # Stage 5a: narrow + cast int32 → int16 (or legacy fp32 → int16) - if self._dec_out_name in ("pcm_unbounded", "pcm"): + # Stage 5a: peak-protect before PCM scaling / INT16 narrowing. + if self._dec_out_name == "audio_unbounded": + source = self.audio_unbounded_buf[ + 0, :self.requested_samples + ] + canon.zero_audio_padding_( + source, + self.valid_sample_mask[:self.requested_samples], + sample_dim=0, + ) + peak = canon.audio_peak(source).float() + protected = canon.protect_audio_peak( + source, + ceiling=self.peak_ceiling, + peak=peak, + validate_nonfinite=False, + emit_warning=False, + ) + self.pcm_int16_buf[:self.requested_samples].copy_( + (protected * canon.PCM16_CEILING).to(canon.torch.int16) + ) + self._peak_ceiling = self.peak_ceiling + elif self._dec_out_name == "pcm": source = self.pcm_int32_buf[0, :self.requested_samples] canon.zero_audio_padding_( source, @@ -718,7 +744,8 @@ def generate(self, prompt: str, *, """Generate one audio clip. Returns (pcm_int16, timing_dict). Returns: - pcm: (T_samples, 2) int16 numpy array, T_samples = round(seconds*44100) + pcm: (T_samples, 2) int16 numpy array, + T_samples = round(seconds * SAMPLE_RATE) timing: dict with 'inference_ms', 'graph_build_ms' (0 if cache hit), 'realtime', 'seed', 'T_lat', 'samples' """ diff --git a/optimized/tensorRT/scripts/sa3_trt_core.py b/optimized/tensorRT/scripts/sa3_trt_core.py index 909c5a6b..1d6d6a99 100644 --- a/optimized/tensorRT/scripts/sa3_trt_core.py +++ b/optimized/tensorRT/scripts/sa3_trt_core.py @@ -19,10 +19,8 @@ SCRIPTS = Path(__file__).resolve().parent REPO = SCRIPTS.parent sys.path.insert(0, str(SCRIPTS)) -PROJECT_ROOT = REPO.parents[1] -sys.path.insert(0, str(PROJECT_ROOT / "stable_audio_3")) -from audio_output import ( +from wav_io import ( PCM16_CEILING, audio_peak, dbfs_to_amplitude, @@ -641,14 +639,14 @@ def encode_chunked(runner: TRTRunner, audio: torch.Tensor, *, def decoder_decode(runner: TRTRunner, latents: torch.Tensor) -> torch.Tensor: """SAME-S/L decoder. - Two engine flavors are supported (auto-detected by output tensor name): + Three engine flavors are supported (auto-detected by output tensor name): - Legacy (output name "audio", fp32/bf16): shape (1, 2, L*4096) audio. The caller is responsible for peak protection + scale + cast to int16 and transposing to (T, 2) interleaved PCM. - - Unbounded PCM (output name "pcm_unbounded", int32): shape - (1, L*4096, 2), scaled and transposed but not hard-clipped. The caller - applies peak protection before narrowing to int16. + - Unbounded audio (output name "audio_unbounded", fp32): shape + (1, L*4096, 2), transposed but neither hard-clipped nor PCM-scaled. The + caller applies peak protection before scaling and narrowing to int16. - Legacy PCM-baked (output name "pcm", int32): the same shape, but the engine has a baked hard clip that cannot be undone at runtime. @@ -656,18 +654,18 @@ def decoder_decode(runner: TRTRunner, latents: torch.Tensor) -> torch.Tensor: odd L matches PT eager at cos ≥ 0.99 on in-distribution latents — no chunking needed. - Returns whatever the engine emits (caller branches on .dtype to decide - what postprocessing — if any — is still needed in Stage 5). + Returns whatever the engine emits; the caller uses dtype and layout to + select the remaining Stage-5 postprocessing. """ ctx = runner.context in_dt = runner.in_dtype["latent"] - if "pcm_unbounded" in runner.out_dtype: - out_name = "pcm_unbounded" + if "audio_unbounded" in runner.out_dtype: + out_name = "audio_unbounded" elif "pcm" in runner.out_dtype: out_name = "pcm" warnings.warn( "legacy TensorRT decoder engine has baked hard clipping; rebuild the " - "decoder engine to get the pcm_unbounded output and preserve peak ratios", + "decoder engine to get the audio_unbounded output and preserve peak ratios", RuntimeWarning, stacklevel=2, ) @@ -1414,6 +1412,15 @@ def _stage_vram(label): return 0 _pinned_pcm.copy_(pcm_w, non_blocking=True) else: _ = pcm_w.cpu().numpy() + elif dec_out.shape[-1] == 2: + protected_w = protect_audio_peak( + dec_out[0], validate_nonfinite=False, emit_warning=False + ) + pcm_w = (protected_w * PCM16_CEILING).to(torch.int16) + if _pinned_pcm is not None: + _pinned_pcm.copy_(pcm_w, non_blocking=True) + else: + _ = pcm_w.cpu().numpy() else: _ = dec_out.cpu().numpy() if _w_audio is not None: @@ -1607,10 +1614,11 @@ def _on_step(i: int, total: int): if args.free_models: runners["dec"].free(); del runners["dec"] decode_ms = (time.time() - t0) * 1000 - _pcm_baked = audio.dtype == torch.int32 + _legacy_pcm_baked = audio.dtype == torch.int32 + _sample_major_float = audio.dtype.is_floating_point and audio.shape[-1] == 2 stage("[4/5]", f"Decoder ({args.decoder})", decode_ms) sub(f"audio {tuple(audio.shape)} {audio.dtype}" - f"{' (pcm baked-in)' if _pcm_baked else ''}") + f"{' (legacy pcm baked-in)' if _legacy_pcm_baked else ''}") _stage_vram("Decode") # ── End of inference wall clock (WAV save excluded — that's I/O) ── @@ -1630,13 +1638,28 @@ def _on_step(i: int, total: int): t0_total = time.time() # PCM conversion. Two paths: - # - PCM engines (int32 (1, T_full, 2)): scale + transpose are already - # done inside the decoder graph. Apply peak protection before narrowing. + # - New engines (fp32 (1, T_full, 2)): transpose is already done inside + # the decoder graph. Protect before PCM scaling and narrowing. + # - Legacy PCM engines (int32 (1, T_full, 2)): clipping, scale, and + # transpose are already baked in and cannot be undone. # - Legacy engines (output "audio", fp32 (1, 2, T_full)): protect, # scale, cast, and transpose on the GPU before the copy. requested_samples = int(round(args.seconds * SAMPLE_RATE)) t0 = time.time() - if _pcm_baked: + if _sample_major_float: + audio_gpu = audio[0] + if audio_gpu.shape[0] > requested_samples: + audio_gpu = audio_gpu[:requested_samples] + audio_gpu = protect_audio_peak(audio_gpu, ceiling=peak_ceiling) + pcm_gpu = (audio_gpu * PCM16_CEILING).to(torch.int16).contiguous() + n = pcm_gpu.shape[0] + if _pinned_pcm is not None: + _pinned_pcm[:n].copy_(pcm_gpu, non_blocking=True) + torch.cuda.synchronize() + pcm = _pinned_pcm[:n].numpy() + else: + pcm = pcm_gpu.cpu().numpy() + elif _legacy_pcm_baked: # (1, T_full, 2) int32 → (T, 2) int16 on GPU pcm_gpu = audio[0] # (T_full, 2) int32 if pcm_gpu.shape[0] > requested_samples: @@ -1672,8 +1695,10 @@ def _on_step(i: int, total: int): t_disk = (time.time() - t0) * 1000 stage("[5/5]", f"WAV → {out_display}", (time.time() - t0_total) * 1000) - if _pcm_baked: - sub(f"protect/cast int32→int16 + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") + if _sample_major_float: + sub(f"protect/scale/cast + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") + elif _legacy_pcm_baked: + sub(f"legacy int32→int16 + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") else: sub(f"protect/cast/transpose + GPU→CPU {t_gpu2cpu:.0f} ms · disk write {t_disk:.0f} ms") diff --git a/optimized/tensorRT/scripts/wav_io.py b/optimized/tensorRT/scripts/wav_io.py new file mode 100644 index 00000000..b8022da2 --- /dev/null +++ b/optimized/tensorRT/scripts/wav_io.py @@ -0,0 +1,30 @@ +"""Compatibility imports 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; rerun the TensorRT bootstrap or use a full checkout" + ) +sys.path.insert(0, str(HELPER_DIR)) + +from audio_output import ( # noqa: E402, F401 + PCM16_CEILING, + audio_peak, + dbfs_to_amplitude, + protect_audio_peak, + report_peak_protection, + save_wav, + zero_audio_padding_, +) diff --git a/optimized/tflite/README.md b/optimized/tflite/README.md index 9fd20034..d1b27601 100644 --- a/optimized/tflite/README.md +++ b/optimized/tflite/README.md @@ -289,6 +289,7 @@ For sub-realtime latency on a supported device, prefer the GPU siblings: | `--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 | +| `--peak-ceiling-dbfs` | 0 | Sample-peak ceiling in dBFS; use `-1` for additional encoding headroom | | `--out` / `-o` | (auto) | Relative → `output/`; 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 | diff --git a/optimized/tflite/bootstrap.ps1 b/optimized/tflite/bootstrap.ps1 index 1b6f2652..d006b191 100644 --- a/optimized/tflite/bootstrap.ps1 +++ b/optimized/tflite/bootstrap.ps1 @@ -59,6 +59,9 @@ if (Get-Command git -ErrorAction SilentlyContinue) { if (Test-Path $extract) { Remove-Item $extract -Recurse -Force } Expand-Archive $zip -DestinationPath $extract -Force Move-Item (Join-Path $extract "$RepoName-$Branch\$SubDir") $LocalDir + Copy-Item ` + (Join-Path $extract "$RepoName-$Branch\stable_audio_3\audio_output.py") ` + (Join-Path $LocalDir "models\defs\audio_output.py") Remove-Item $zip -Force Remove-Item $extract -Recurse -Force Ok "extracted to .\$LocalDir" diff --git a/optimized/tflite/bootstrap.sh b/optimized/tflite/bootstrap.sh index a7bdeeb2..14105f5e 100755 --- a/optimized/tflite/bootstrap.sh +++ b/optimized/tflite/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 @@ -104,11 +105,13 @@ else 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" + 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/models/defs/audio_output.py" ok "extracted $(find "$LOCAL_DIR" -type f | wc -l | tr -d ' ') files to ./$LOCAL_DIR" fi fi diff --git a/optimized/tflite/models/defs/tflite_pipeline.py b/optimized/tflite/models/defs/tflite_pipeline.py index 4b117f5e..7a319de7 100644 --- a/optimized/tflite/models/defs/tflite_pipeline.py +++ b/optimized/tflite/models/defs/tflite_pipeline.py @@ -12,11 +12,12 @@ repo's tflite_pipeline.py and are intentionally dropped here. """ from __future__ import annotations -import wave from pathlib import Path from typing import Callable import numpy as np +from .wav_io import save_wav as _save_wav + # This file lives in /models/defs/. The bundled SentencePiece model sits at # /models/tokenizer.model — resolve it relative to this file so the tokenizer # works regardless of the caller's cwd. @@ -30,12 +31,14 @@ # ───────────────────────── WAV ───────────────────────── -def save_wav(path, audio): # audio: (2, T) float32 in [-1,1] - audio = np.clip(np.asarray(audio, np.float32), -1, 1) - pcm = (audio * 32767.0).astype(np.int16).T # (T, 2) interleaved - with wave.open(str(path), "wb") as w: - w.setnchannels(audio.shape[0]); w.setsampwidth(2); w.setframerate(SAMPLE_RATE) - w.writeframes(pcm.tobytes()) +def save_wav(path, audio, *, peak_ceiling_dbfs=0.0): + """Write channel-first float audio with shared no-boost peak protection.""" + _save_wav( + str(path), + np.asarray(audio, np.float32), + SAMPLE_RATE, + peak_ceiling_dbfs=peak_ceiling_dbfs, + ) # ───────────────────────── Tokenizer (SentencePiece, bundled) ───────────────────────── diff --git a/optimized/tflite/models/defs/wav_io.py b/optimized/tflite/models/defs/wav_io.py new file mode 100644 index 00000000..d37935dc --- /dev/null +++ b/optimized/tflite/models/defs/wav_io.py @@ -0,0 +1,27 @@ +"""Compatibility imports 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[4] / "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 TFLite 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, +) diff --git a/optimized/tflite/scripts/sa3_gradio.py b/optimized/tflite/scripts/sa3_gradio.py index 98f4ba4d..7585e47e 100644 --- a/optimized/tflite/scripts/sa3_gradio.py +++ b/optimized/tflite/scripts/sa3_gradio.py @@ -61,6 +61,7 @@ from lora_patch import get_patched_dit # noqa: E402 from weights import ensure_local, PRECISIONS, dit_rel, dec_rel, enc_rel # noqa: E402 from spec import render_spectrogram_png # noqa: E402 +from models.defs.wav_io import audio_to_pcm16 # noqa: E402 OUTPUT_DIR = REPO / "output" / "gradio" OUTPUT_DIR.mkdir(parents=True, exist_ok=True) @@ -824,7 +825,12 @@ def _generate_entry(dit_name, precision, 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, precision) out_path = OUTPUT_DIR / f"{basename}.wav" _save_wav(pcm, out_path) diff --git a/optimized/tflite/scripts/sa3_tflite.py b/optimized/tflite/scripts/sa3_tflite.py index 9c020ae3..8b249ca9 100644 --- a/optimized/tflite/scripts/sa3_tflite.py +++ b/optimized/tflite/scripts/sa3_tflite.py @@ -60,6 +60,7 @@ 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 models.defs.wav_io import dbfs_to_amplitude from weights import ensure_local, is_present, PRECISIONS, dit_rel, dec_rel, enc_rel SAMPLE_RATE = 44100 @@ -607,6 +608,8 @@ def main(): ap.add_argument("--out", "-o", default=None, help="Output WAV path. Relative → output/; absolute → as-is. " "If omitted, auto-named from the prompt + 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="Play the WAV after writing (blocking): `afplay` on macOS, " "winsound on Windows, `aplay` on Linux (prints the path if " @@ -614,6 +617,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)) # 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 @@ -877,7 +884,11 @@ def on_chunk(i, n): 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) + P.save_wav( + args.out, + audio_np, + peak_ceiling_dbfs=args.peak_ceiling_dbfs, + ) 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}") diff --git a/optimized/tflite/scripts/test_windows_compat.py b/optimized/tflite/scripts/test_windows_compat.py index 00bbf151..0b5fe2b2 100644 --- a/optimized/tflite/scripts/test_windows_compat.py +++ b/optimized/tflite/scripts/test_windows_compat.py @@ -19,9 +19,13 @@ import sys import tempfile import unittest +import warnings +import wave from pathlib import Path from unittest import mock +import numpy as np + # 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): @@ -33,6 +37,7 @@ SCRIPTS_DIR = Path(__file__).resolve().parent # /scripts PROJECT_DIR = SCRIPTS_DIR.parent # sys.path.insert(0, str(SCRIPTS_DIR)) # import weights / sa3_tflite / examples +sys.path.insert(0, str(PROJECT_DIR)) # import models.defs.* import weights # noqa: E402 @@ -133,8 +138,24 @@ def test_sa3_tflite_help(self): f"stderr:\n{r.stderr.decode('utf-8', 'replace')}") out = r.stdout.decode("utf-8", "replace") self.assertIn("--prompt", out) + self.assertIn("--peak-ceiling-dbfs", out) self.assertIn("--play", out) + def test_positive_peak_ceiling_is_rejected_before_model_loading(self): + r = subprocess.run( + [ + sys.executable, + str(SCRIPTS_DIR / "sa3_tflite.py"), + "--peak-ceiling-dbfs", + "0.1", + ], + capture_output=True, + env=self._clean_env(), + timeout=120, + ) + self.assertEqual(r.returncode, 2) + self.assertIn("ceiling_dbfs must be <= 0", r.stderr.decode("utf-8", "replace")) + def test_examples_block(self): # examples.py is a module (no __main__): render the full emoji block. r = subprocess.run( @@ -183,5 +204,28 @@ def test_current_platform_resolves(self): self.assertIsNone(argv) +class TestWavPeakProtection(unittest.TestCase): + def test_tflite_writer_attenuates_instead_of_hard_clipping(self): + from models.defs.tflite_pipeline import save_wav + + audio = np.array( + [[0.0, 1.25, 1.75], [0.0, -1.25, -1.75]], dtype=np.float32 + ) + with tempfile.TemporaryDirectory(prefix="sa3_pcm_test_") as tmp: + path = Path(tmp) / "protected.wav" + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + save_wav(path, audio) + with wave.open(str(path), "rb") as wav_file: + pcm = np.frombuffer( + wav_file.readframes(wav_file.getnframes()), dtype=np.int16 + ).reshape(-1, wav_file.getnchannels()) + + self.assertEqual(int(pcm[-1, 0]), 32767) + self.assertAlmostEqual( + int(pcm[-2, 0]), 32767 * 1.25 / 1.75, delta=1 + ) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/stable_audio_3/audio_output.py b/stable_audio_3/audio_output.py index feb34110..6faf918a 100644 --- a/stable_audio_3/audio_output.py +++ b/stable_audio_3/audio_output.py @@ -255,15 +255,24 @@ def save_wav( peak_ceiling_dbfs: float = 0.0, ) -> None: """Write channel-first NumPy floating-point audio as 16-bit PCM WAV.""" - if _backend_name(audio) != "numpy": - raise TypeError("save_wav expects a channel-first numpy.ndarray") - - import numpy as np - - audio = protect_audio_peak(audio, ceiling=dbfs_to_amplitude(peak_ceiling_dbfs)) - pcm = (audio * PCM16_CEILING).astype(np.int16).T + pcm = audio_to_pcm16(audio, peak_ceiling_dbfs=peak_ceiling_dbfs) with wave.open(path, "wb") as wav: wav.setnchannels(audio.shape[0]) wav.setsampwidth(2) wav.setframerate(sample_rate) wav.writeframes(pcm.tobytes()) + + +def audio_to_pcm16( + audio: Any, + *, + peak_ceiling_dbfs: float = 0.0, +) -> Any: + """Protect channel-first NumPy audio and return interleaved PCM16.""" + if _backend_name(audio) != "numpy": + raise TypeError("audio_to_pcm16 expects a channel-first numpy.ndarray") + + import numpy as np + + audio = protect_audio_peak(audio, ceiling=dbfs_to_amplitude(peak_ceiling_dbfs)) + return (audio * PCM16_CEILING).astype(np.int16).T diff --git a/tests/test_audio_peak_protection.py b/tests/test_audio_peak_protection.py index 30af5090..73b602e8 100644 --- a/tests/test_audio_peak_protection.py +++ b/tests/test_audio_peak_protection.py @@ -11,6 +11,7 @@ from stable_audio_3.audio_output import ( PCM16_CEILING, apply_output_peak_policy, + audio_to_pcm16, audio_peak, dbfs_to_amplitude, protect_audio_peak, @@ -167,6 +168,34 @@ def test_mlx_wav_serializer_attenuates_instead_of_clipping(tmp_path): assert pcm[1, -2] == pytest.approx(-32767 * 1.25 / 1.75, abs=1) +def test_shared_pcm_conversion_attenuates_without_changing_sample_ratios(): + audio = np.array([[0.0, 1.25, 1.75], [0.0, -1.25, -1.75]], dtype=np.float32) + + with pytest.warns(RuntimeWarning, match="peak 1.750"): + pcm = audio_to_pcm16(audio) + + assert pcm.shape == (3, 2) + assert pcm[-1, 0] == 32767 + assert pcm[-2, 0] == pytest.approx(32767 * 1.25 / 1.75, abs=1) + + +def test_tflite_wav_serializer_uses_shared_peak_policy(tmp_path): + from optimized.tflite.models.defs.tflite_pipeline import save_wav as save_tflite_wav + + audio = np.array([[0.0, 1.25, 1.75], [0.0, -1.25, -1.75]], dtype=np.float32) + output = tmp_path / "tflite.wav" + + with pytest.warns(RuntimeWarning, match="peak 1.750"): + save_tflite_wav(output, audio, peak_ceiling_dbfs=-1.0) + + with wave.open(str(output), "rb") as wav: + pcm = np.frombuffer(wav.readframes(wav.getnframes()), dtype=np.int16) + pcm = pcm.reshape(-1, wav.getnchannels()).T + + assert pcm[0, -1] == pytest.approx(32767 * dbfs_to_amplitude(-1.0), abs=1) + assert pcm[0, -2] == pytest.approx(pcm[0, -1] * 1.25 / 1.75, abs=1) + + class _FakePipeline: sample_rate = 1 io_channels = 2 diff --git a/tests/test_tensorrt_decoder_output.py b/tests/test_tensorrt_decoder_output.py index 1cd523fb..6a0dfe48 100644 --- a/tests/test_tensorrt_decoder_output.py +++ b/tests/test_tensorrt_decoder_output.py @@ -1,9 +1,12 @@ +import os +from unittest.mock import patch + import numpy as np import pytest from optimized.tensorRT.build.decoder_output import ( - UNBOUNDED_PCM_OUTPUT, - force_unbounded_pcm_tail_fp32, + UNBOUNDED_AUDIO_OUTPUT, + force_unbounded_audio_output_fp32, remove_output_hard_clip, rewrite_decoder_onnx, ) @@ -14,7 +17,7 @@ numpy_helper = onnx.numpy_helper -def _decoder_tail_model(): +def _decoder_tail_model(*, large_initializer_value=None): audio = helper.make_tensor_value_info("audio", TensorProto.FLOAT, [1, 2, 4]) pcm = helper.make_tensor_value_info("pcm", TensorProto.INT32, [1, 4, 2]) minimum = numpy_helper.from_array(np.array(-1.0, dtype=np.float32), "minimum") @@ -28,12 +31,20 @@ def _decoder_tail_model(): ), helper.make_node("Transpose", ["pcm_channels_first"], ["pcm"], perm=[0, 2, 1]), ] + initializers = [minimum, maximum, scale] + if large_initializer_value is not None: + initializers.append( + numpy_helper.from_array( + np.full(300_000, large_initializer_value, dtype=np.float32), + "large_external_weight", + ) + ) graph = helper.make_graph( nodes, "decoder_tail", [audio], [pcm], - initializer=[minimum, maximum, scale], + initializer=initializers, ) return helper.make_model( graph, @@ -48,9 +59,15 @@ def test_decoder_rewrite_removes_clip_and_marks_unbounded_output(): assert remove_output_hard_clip(model) == 1 assert all(node.op_type != "Clip" for node in model.graph.node) - assert model.graph.output[0].name == UNBOUNDED_PCM_OUTPUT - mul = next(node for node in model.graph.node if node.op_type == "Mul") - assert mul.input[0] == "audio" + assert model.graph.output[0].name == UNBOUNDED_AUDIO_OUTPUT + assert model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT + assert all(node.op_type != "Mul" for node in model.graph.node) + assert all( + node.op_type != "Cast" + or next(attr.i for attr in node.attribute if attr.name == "to") + != TensorProto.INT32 + for node in model.graph.node + ) onnx.checker.check_model(model) @@ -69,22 +86,19 @@ def test_decoder_rewrite_writes_loadable_onnx(tmp_path): assert rewrite_decoder_onnx(str(source), str(output)) == str(output) rewritten = onnx.load(output) - assert rewritten.graph.output[0].name == UNBOUNDED_PCM_OUTPUT - multiply = next(node for node in rewritten.graph.node if node.op_type == "Mul") + assert rewritten.graph.output[0].name == UNBOUNDED_AUDIO_OUTPUT + transpose = next( + node for node in rewritten.graph.node if node.op_type == "Transpose" + ) signal_producer = next( - node for node in rewritten.graph.node if multiply.input[0] in node.output + node for node in rewritten.graph.node if transpose.input[0] in node.output ) assert signal_producer.op_type == "Cast" assert ( next(attr.i for attr in signal_producer.attribute if attr.name == "to") == TensorProto.FLOAT ) - scale = next( - initializer - for initializer in rewritten.graph.initializer - if initializer.name == "peak_protect_pcm16_scale_fp32" - ) - assert numpy_helper.to_array(scale).item() == 32767.0 + assert all(node.op_type != "Mul" for node in rewritten.graph.node) onnx.checker.check_model(rewritten) @@ -92,7 +106,7 @@ def test_decoder_rewrite_preserves_out_of_range_sample_ratios(): onnxruntime = pytest.importorskip("onnxruntime") model = _decoder_tail_model() remove_output_hard_clip(model) - force_unbounded_pcm_tail_fp32(model) + force_unbounded_audio_output_fp32(model) session = onnxruntime.InferenceSession( model.SerializeToString(), providers=["CPUExecutionProvider"] ) @@ -101,14 +115,35 @@ def test_decoder_rewrite_preserves_out_of_range_sample_ratios(): dtype=np.float32, ) - pcm = session.run([UNBOUNDED_PCM_OUTPUT], {"audio": audio})[0] + unbounded = session.run([UNBOUNDED_AUDIO_OUTPUT], {"audio": audio})[0] - assert pcm[0, -1, 0] > 32767 - assert pcm[0, -2, 0] < pcm[0, -1, 0] - ratio = pcm[0, -2, 0] / pcm[0, -1, 0] + assert unbounded[0, -1, 0] == 1.75 + assert unbounded[0, -2, 0] < unbounded[0, -1, 0] + ratio = unbounded[0, -2, 0] / unbounded[0, -1, 0] assert ratio == pytest.approx(1.25 / 1.75, abs=1e-4) +def test_decoder_rewrite_preserves_nonfinite_and_extreme_float_values(): + onnxruntime = pytest.importorskip("onnxruntime") + model = _decoder_tail_model() + remove_output_hard_clip(model) + session = onnxruntime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + audio = np.array( + [[[np.nan, np.inf, -np.inf, 70_000.0], [0.0, 1.0, -1.0, -70_000.0]]], + dtype=np.float32, + ) + + unbounded = session.run([UNBOUNDED_AUDIO_OUTPUT], {"audio": audio})[0] + + assert np.isnan(unbounded[0, 0, 0]) + assert np.isposinf(unbounded[0, 1, 0]) + assert np.isneginf(unbounded[0, 2, 0]) + assert unbounded[0, 3, 0] == 70_000.0 + assert unbounded[0, 3, 1] == -70_000.0 + + def test_decoder_rewrite_refuses_an_unrelated_upstream_clip(): model = _decoder_tail_model() zero = numpy_helper.from_array(np.array(0.0, dtype=np.float32), "zero") @@ -127,18 +162,18 @@ def test_decoder_rewrite_refuses_an_unrelated_upstream_clip(): remove_output_hard_clip(model) -def test_fp16_mixed_tail_is_promoted_before_unbounded_scale(): +def test_fp16_mixed_output_is_promoted_to_fp32(): onnxruntime = pytest.importorskip("onnxruntime") model = _decoder_tail_model() remove_output_hard_clip(model) model.graph.input[0].type.tensor_type.elem_type = TensorProto.FLOAT16 - scale = next(item for item in model.graph.initializer if item.name == "scale") - scale.CopyFrom( - numpy_helper.from_array(np.array(32768.0, dtype=np.float16), "scale") - ) + output_cast = next(node for node in model.graph.node if node.op_type == "Cast") + next( + attr for attr in output_cast.attribute if attr.name == "to" + ).i = TensorProto.FLOAT16 - assert force_unbounded_pcm_tail_fp32(model) == 1 - assert force_unbounded_pcm_tail_fp32(model) == 0 + assert force_unbounded_audio_output_fp32(model) == 1 + assert force_unbounded_audio_output_fp32(model) == 0 onnx.checker.check_model(model) session = onnxruntime.InferenceSession( model.SerializeToString(), providers=["CPUExecutionProvider"] @@ -148,7 +183,76 @@ def test_fp16_mixed_tail_is_promoted_before_unbounded_scale(): dtype=np.float16, ) - pcm = session.run([UNBOUNDED_PCM_OUTPUT], {"audio": audio})[0] + unbounded = session.run([UNBOUNDED_AUDIO_OUTPUT], {"audio": audio})[0] - assert pcm[0, -1, 0] == 3 * 32767 - assert pcm[0, -1, 1] == -3 * 32767 + assert unbounded.dtype == np.float32 + assert unbounded[0, -1, 0] == 3.0 + assert unbounded[0, -1, 1] == -3.0 + + +def test_rewrite_replaces_external_sidecar_instead_of_appending(tmp_path): + source = tmp_path / "decoder.onnx" + output = tmp_path / "decoder_unbounded.onnx" + onnx.save(_decoder_tail_model(large_initializer_value=1.0), source) + + rewrite_decoder_onnx(str(source), str(output)) + first_sidecar = _external_sidecar(output) + first_size = first_sidecar.stat().st_size + rewrite_decoder_onnx(str(source), str(output)) + second_sidecar = _external_sidecar(output) + + assert second_sidecar != first_sidecar + assert second_sidecar.stat().st_size == first_size + assert not first_sidecar.exists() + onnx.checker.check_model(onnx.load(output, load_external_data=True)) + + +def _external_sidecar(model_path): + model = onnx.load(model_path, load_external_data=False) + locations = { + entry.value + for initializer in model.graph.initializer + for entry in initializer.external_data + if entry.key == "location" + } + assert len(locations) == 1 + return model_path.parent / locations.pop() + + +def test_failed_model_swap_keeps_previous_model_and_sidecar_consistent(tmp_path): + first_source = tmp_path / "decoder_v1.onnx" + second_source = tmp_path / "decoder_v2.onnx" + output = tmp_path / "decoder_unbounded.onnx" + onnx.save(_decoder_tail_model(large_initializer_value=1.0), first_source) + onnx.save(_decoder_tail_model(large_initializer_value=2.0), second_source) + rewrite_decoder_onnx(str(first_source), str(output)) + first_sidecar = _external_sidecar(output) + + real_replace = os.replace + replace_calls = 0 + + def fail_second_replace(source, destination): + nonlocal replace_calls + replace_calls += 1 + if replace_calls == 2: + raise OSError("injected model-swap failure") + return real_replace(source, destination) + + with ( + patch( + "optimized.tensorRT.build.decoder_output.os.replace", + side_effect=fail_second_replace, + ), + pytest.raises(OSError, match="injected model-swap failure"), + ): + rewrite_decoder_onnx(str(second_source), str(output)) + + assert _external_sidecar(output) == first_sidecar + assert first_sidecar.exists() + loaded = onnx.load(output, load_external_data=True) + weight = next( + item + for item in loaded.graph.initializer + if item.name == "large_external_weight" + ) + assert np.all(numpy_helper.to_array(weight) == 1.0)