From b00a1f2ed03afd688394b92712e195164f0d90e2 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:30:14 -0400 Subject: [PATCH 1/3] Pin SA3 to current upstream and version SAME-L plugins Move the managed SA3 source from the old FP8 development lineage to current upstream plus the rebased TensorRT integer-attribute fix. The current upstream SAME-L plugin bakes its AOT/JIT implementation into the serialized engine. Include the plugin revision and implementation choice in the engine name and metadata so stale decoders are ignored instead of being paired with incompatible plugin code. Keep the plugin revision stable across unrelated vendor bumps so they do not force a 1.2 GB decoder rebuild. Existing engines remain beside the new artifact for rollback. Also accept git worktrees as DEMON_SA3_SRC overrides. --- acestep/engine/sa3_context.py | 3 +- acestep/engine/sa3_helpers.py | 9 +++- acestep/engine/sa3_trt.py | 28 ++++++++-- acestep/engine/trt/sa3_build.py | 21 ++++++-- tests/unit/test_sa3_trt_engine_identity.py | 60 ++++++++++++++++++++++ tests/unit/test_sa3_vendor.py | 19 +++++++ 6 files changed, 130 insertions(+), 10 deletions(-) create mode 100644 tests/unit/test_sa3_trt_engine_identity.py diff --git a/acestep/engine/sa3_context.py b/acestep/engine/sa3_context.py index 681ec49e..e2232c24 100644 --- a/acestep/engine/sa3_context.py +++ b/acestep/engine/sa3_context.py @@ -254,7 +254,8 @@ class SA3SAMEWindowCodec: Two execution paths, identical interface: * **TRT** when ``use_trt`` and the built window engine exists - (``same_l_decode_window_t*``): ~9-10 ms per ~1 s window, latent + (``same_l_decode_window__t*``): ~9-10 ms per ~1 s + window, latent scaled by ``pretransform.scale`` before the call (spike ``scale_mode="pretransform"``, rel_rms ~8e-3 vs eager full). * **Eager** fallback: the spike's ``decode_sa3_latent_window`` diff --git a/acestep/engine/sa3_helpers.py b/acestep/engine/sa3_helpers.py index d0f55c89..c6ada2b3 100644 --- a/acestep/engine/sa3_helpers.py +++ b/acestep/engine/sa3_helpers.py @@ -25,7 +25,10 @@ # DEMON tracks this fork branch until the SA3 TensorRT/FP8 producer work # merges upstream. The hash is the reproducibility boundary for installs. SA3_VENDOR_URL = "https://github.com/ryanontheinside/stable-audio-3" -SA3_VENDOR_SHA = "03992120a3e296562271f209f4dd61dcfb55afff" +SA3_VENDOR_SHA = "c07698548567fe6f163806f692d282bbaa57aba3" +# Revision that last changed the source compiled into the SAME-L TensorRT +# engine. Keep this stable across vendor bumps that only touch other code. +SA3_SAME_L_PLUGIN_REVISION = "c07698548567fe6f163806f692d282bbaa57aba3" SA3_VENDOR_ENV = "DEMON_SA3_SRC" SA3_VENDOR_DIRNAME = "stable-audio-3" @@ -152,7 +155,9 @@ def ensure_sa3_vendor( reported as an error instead of being overwritten. """ vendor = sa3_vendor_dir() - if not (vendor / ".git").is_dir(): + # A normal clone has a .git directory; a git worktree has a .git file + # pointing at the parent repository. Both are valid developer overrides. + if not (vendor / ".git").exists(): if check_only: raise FileNotFoundError(f"SA3 vendor source is missing at {vendor}") vendor.parent.mkdir(parents=True, exist_ok=True) diff --git a/acestep/engine/sa3_trt.py b/acestep/engine/sa3_trt.py index 468f4803..c548a953 100644 --- a/acestep/engine/sa3_trt.py +++ b/acestep/engine/sa3_trt.py @@ -48,6 +48,7 @@ from __future__ import annotations import math +import os import re import sys import threading @@ -58,7 +59,7 @@ from acestep import paths from acestep.engine.obs import logger -from acestep.engine.sa3_helpers import sa3_vendor_dir +from acestep.engine.sa3_helpers import SA3_SAME_L_PLUGIN_REVISION, sa3_vendor_dir IO_CHANNELS = 256 T5_TOKENS = 256 @@ -95,7 +96,10 @@ _DIT_REFIT_DIR_RE = re.compile( r"^(?P.+_dit)_refit_l(?P\d+)_(?P\d+)_(?P\d+)$" ) -_SAME_L_DIR_RE = re.compile(r"^same_l_decode_window_t(?P\d+)_(?P\d+)_(?P\d+)$") +_SAME_L_DIR_RE = re.compile( + r"^same_l_decode_window_(?P[a-z0-9_]+)_t" + r"(?P\d+)_(?P\d+)_(?P\d+)$" +) # Deserialized-engine process cache. Engines are immutable post-load and # support multiple execution contexts, so sharing one deserialization @@ -108,6 +112,23 @@ _SAME_PLUGIN_REGISTERED = False +def same_l_plugin_build_tag() -> str: + """Identity of the plugin implementation compiled into a SAME-L engine. + + The upstream plugin is part of the serialized TensorRT engine, so changing + the vendored source revision or its selected AOT backend requires a new + engine even when the ONNX graph is unchanged. + """ + requested_plugin = os.environ.get("SA3_SWA_PLUGIN", "aot").strip().lower() + plugin = "jit" if requested_plugin == "jit" else "aot" + if plugin == "jit": + implementation = "jit" + else: + requested_backend = os.environ.get("SA3_SWA_AOT", "mma").strip().lower() + implementation = "mma" if requested_backend == "mma" else "ptx" + return f"{plugin}_{implementation}_v{SA3_SAME_L_PLUGIN_REVISION[:12]}" + + def trt_engines_dir() -> Path: return paths.models_dir() / "sa3" / "trt_engines" @@ -286,9 +307,10 @@ def find_same_l_window_engine() -> Optional[tuple]: base = trt_engines_dir() if not base.is_dir(): return None + expected_tag = same_l_plugin_build_tag() for sub in base.iterdir(): m = _SAME_L_DIR_RE.match(sub.name) - if not m: + if not m or m.group("tag") != expected_tag: continue f = sub / f"{sub.name}.trt" if f.is_file(): diff --git a/acestep/engine/trt/sa3_build.py b/acestep/engine/trt/sa3_build.py index 16a1feaf..502a5b6e 100644 --- a/acestep/engine/trt/sa3_build.py +++ b/acestep/engine/trt/sa3_build.py @@ -87,6 +87,7 @@ SAMPLES_PER_LATENT, T5_TOKENS, _register_same_plugin, + same_l_plugin_build_tag, trt_engines_dir, ) @@ -218,10 +219,11 @@ class SameLWindowBuildConfig: max_latents: int workspace_gb: float = 16.0 onnx_files: list[str] = field(default_factory=lambda: list(SAME_L_ONNX_FILES)) + plugin_build_tag: str = field(default_factory=same_l_plugin_build_tag) def engine_name(self) -> str: return ( - f"same_l_decode_window_t{self.min_latents}" + f"same_l_decode_window_{self.plugin_build_tag}_t{self.min_latents}" f"_{self.opt_latents}_{self.max_latents}" ) @@ -292,6 +294,7 @@ def _build_strongly_typed_engine( workspace_gb: float, profile_shapes: dict[str, tuple[tuple, tuple, tuple]], refit: bool = False, + python_plugin_preference: str | None = None, ) -> None: """Parse + build one STRONGLY_TYPED engine and serialize it to disk. @@ -304,9 +307,16 @@ def _build_strongly_typed_engine( trt_logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(trt_logger) - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) - ) + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + if python_plugin_preference == "aot": + network_flags |= 1 << int( + trt.NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS + ) + elif python_plugin_preference == "jit": + network_flags |= 1 << int( + trt.NetworkDefinitionCreationFlag.PREFER_JIT_PYTHON_PLUGINS + ) + network = builder.create_network(network_flags) parser = trt.OnnxParser(network, trt_logger) if not parser.parse_from_file(onnx_path): for i in range(parser.num_errors): @@ -494,6 +504,9 @@ def _build_same_l_window_engine( profile_shapes={ "latent": ((1, IO_CHANNELS, lo), (1, IO_CHANNELS, opt), (1, IO_CHANNELS, hi)), }, + python_plugin_preference=( + "jit" if config.plugin_build_tag.startswith("jit_") else "aot" + ), ) _write_metadata(engine_path=engine_path, expected=expected, env=env) elapsed = time.time() - t0 diff --git a/tests/unit/test_sa3_trt_engine_identity.py b/tests/unit/test_sa3_trt_engine_identity.py new file mode 100644 index 00000000..f78d595e --- /dev/null +++ b/tests/unit/test_sa3_trt_engine_identity.py @@ -0,0 +1,60 @@ +from pathlib import Path + +from acestep.engine import sa3_trt +from acestep.engine.trt.sa3_build import SameLWindowBuildConfig + + +def _engine(root: Path, name: str) -> Path: + path = root / "sa3" / "trt_engines" / name / f"{name}.trt" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"engine") + return path + + +def test_same_l_engine_name_carries_plugin_build_identity(monkeypatch): + monkeypatch.delenv("SA3_SWA_AOT", raising=False) + config = SameLWindowBuildConfig(32, 56, 96) + + assert config.plugin_build_tag == sa3_trt.same_l_plugin_build_tag() + assert config.engine_name() == ( + f"same_l_decode_window_{config.plugin_build_tag}_t32_56_96" + ) + + +def test_same_l_discovery_ignores_legacy_engine(monkeypatch, tmp_path): + monkeypatch.setenv("ACESTEP_MODELS_DIR", str(tmp_path)) + monkeypatch.delenv("SA3_SWA_AOT", raising=False) + _engine(tmp_path, "same_l_decode_window_t32_56_96") + + assert sa3_trt.find_same_l_window_engine() is None + + +def test_same_l_discovery_selects_current_plugin_engine(monkeypatch, tmp_path): + monkeypatch.setenv("ACESTEP_MODELS_DIR", str(tmp_path)) + monkeypatch.delenv("SA3_SWA_AOT", raising=False) + tag = sa3_trt.same_l_plugin_build_tag() + expected = _engine(tmp_path, f"same_l_decode_window_{tag}_t32_56_96") + + assert sa3_trt.find_same_l_window_engine() == (expected, 32, 96) + + +def test_same_l_aot_backend_is_part_of_identity(monkeypatch): + monkeypatch.setenv("SA3_SWA_PLUGIN", "aot") + monkeypatch.setenv("SA3_SWA_AOT", "mma") + mma = sa3_trt.same_l_plugin_build_tag() + monkeypatch.setenv("SA3_SWA_AOT", "ptx") + ptx = sa3_trt.same_l_plugin_build_tag() + + assert mma != ptx + assert mma.startswith("aot_mma_v") + assert ptx.startswith("aot_ptx_v") + + +def test_same_l_plugin_kind_is_part_of_identity(monkeypatch): + monkeypatch.setenv("SA3_SWA_PLUGIN", "aot") + aot = sa3_trt.same_l_plugin_build_tag() + monkeypatch.setenv("SA3_SWA_PLUGIN", "jit") + jit = sa3_trt.same_l_plugin_build_tag() + + assert aot != jit + assert jit.startswith("jit_jit_v") diff --git a/tests/unit/test_sa3_vendor.py b/tests/unit/test_sa3_vendor.py index 11c4bf76..154a9a20 100644 --- a/tests/unit/test_sa3_vendor.py +++ b/tests/unit/test_sa3_vendor.py @@ -148,6 +148,25 @@ def fake_git(args: list[str], cwd: Path | None = None) -> str: ) in calls +def test_ensure_vendor_accepts_git_worktree_override(monkeypatch, tmp_path): + vendor = tmp_path / "stable-audio-3-worktree" + vendor.mkdir() + (vendor / ".git").write_text("gitdir: ../repo/.git/worktrees/vendor\n") + monkeypatch.setenv(sa3_helpers.SA3_VENDOR_ENV, str(vendor)) + + def fake_git(args: list[str], cwd: Path | None = None) -> str: + assert cwd == vendor + if args == ["status", "--porcelain"]: + return "" + if args == ["rev-parse", "HEAD"]: + return sa3_helpers.SA3_VENDOR_SHA + raise AssertionError(f"unexpected git call: {args}") + + monkeypatch.setattr(sa3_helpers, "_git", fake_git) + + assert sa3_helpers.ensure_sa3_vendor() == vendor + + def test_ensure_vendor_refuses_dirty_wrong_commit(monkeypatch, tmp_path): monkeypatch.delenv(sa3_helpers.SA3_VENDOR_ENV, raising=False) monkeypatch.setenv("ACESTEP_MODELS_DIR", str(tmp_path)) From edbd8e6870cacf9924c8f564faaeb051d27e19cd Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:02:30 -0400 Subject: [PATCH 2/3] fix(sa3): pin the vendored SA3 to the monotonic audio-to-audio schedule Move SA3_VENDOR_SHA to the combined fork pin (current Stability main + the TRT integer-attribute fix + the normalized-then-scaled schedule ordering). The schedule change is host-side Python only; the SAME-L plugin source is unchanged between the previous pin and this one, so SA3_SAME_L_PLUGIN_REVISION stays put and the existing versioned decoder engine is reused. --- acestep/engine/sa3_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/acestep/engine/sa3_helpers.py b/acestep/engine/sa3_helpers.py index c6ada2b3..9687f3b0 100644 --- a/acestep/engine/sa3_helpers.py +++ b/acestep/engine/sa3_helpers.py @@ -25,7 +25,7 @@ # DEMON tracks this fork branch until the SA3 TensorRT/FP8 producer work # merges upstream. The hash is the reproducibility boundary for installs. SA3_VENDOR_URL = "https://github.com/ryanontheinside/stable-audio-3" -SA3_VENDOR_SHA = "c07698548567fe6f163806f692d282bbaa57aba3" +SA3_VENDOR_SHA = "960da1f8cbe205ab3b702edbfabd91113ab22473" # Revision that last changed the source compiled into the SAME-L TensorRT # engine. Keep this stable across vendor bumps that only touch other code. SA3_SAME_L_PLUGIN_REVISION = "c07698548567fe6f163806f692d282bbaa57aba3" From 9068ce3b7f4630c002ec4c182c07f816a7f75b2e Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:33:05 -0400 Subject: [PATCH 3/3] Map the SA3 denoise knob onto measured audio change Keep the upstream-corrected monotonic schedule, but map the product dial to entry sigma using the inverse of the measured change curve. This spreads useful movement across the control instead of concentrating it near the top. The mapping is deliberately presented as a product choice, not another sampler fix: its referee is validated for ordering rather than equal perceptibility, and one global curve moves different material at different rates. DEMON_SA3_DENOISE_MAPPING=identity bypasses it while retaining the upstream schedule correction. --- acestep/engine/sa3_context.py | 4 +- acestep/engine/sa3_denoise_mapping.py | 77 ++++++++++++++++++++++++++ acestep/streaming/sa3_backend.py | 9 +-- tests/unit/test_sa3_denoise_mapping.py | 53 ++++++++++++++++++ 4 files changed, 138 insertions(+), 5 deletions(-) create mode 100644 acestep/engine/sa3_denoise_mapping.py create mode 100644 tests/unit/test_sa3_denoise_mapping.py diff --git a/acestep/engine/sa3_context.py b/acestep/engine/sa3_context.py index e2232c24..3f573a64 100644 --- a/acestep/engine/sa3_context.py +++ b/acestep/engine/sa3_context.py @@ -187,9 +187,11 @@ def make_schedule_builder( def _builder(denoise: float) -> torch.Tensor: import stable_audio_3.inference.sampling as sampling + from .sa3_denoise_mapping import map_denoise_to_entry_sigma + schedule = sampling.build_schedule( steps=int(steps), - sigma_max=float(denoise), + sigma_max=map_denoise_to_entry_sigma(float(denoise)), dist_shift=prepared["dist_shift"], effective_seq_len=prepared["effective_seq_len"], fallback_seq_len=prepared["fallback_seq_len"], diff --git a/acestep/engine/sa3_denoise_mapping.py b/acestep/engine/sa3_denoise_mapping.py new file mode 100644 index 00000000..c0e8cbd6 --- /dev/null +++ b/acestep/engine/sa3_denoise_mapping.py @@ -0,0 +1,77 @@ +"""Optional product mapping for the Stable Audio 3 denoise control. + +The corrected upstream schedule makes ``sigma_max`` mathematically monotonic, +but measured audio change is still concentrated near the top of its range. In +a 41-point sweep over nine clips, less than one fifth of the available change +occurred below sigma 0.70. The knot table below inverts that measured curve so +equal control movements target roughly equal movements in the referee score. + +The referee combines harmonic and rhythmic change and reproduced two separate +listening-order judgments that a loudness-based measure did not. It validates +the ordering and the location of the dead region, not a claim that every step +is equally perceptible. This is also one global mapping: sparse acoustic input +can move faster than dense electronic material. + +Set ``DEMON_SA3_DENOISE_MAPPING=identity`` to bypass the product mapping while +retaining the upstream monotonic-schedule bugfix. +""" + +from __future__ import annotations + +import os +import typing as tp + +__all__ = [ + "denoise_mapping_mode", + "dial_to_entry_sigma", + "map_denoise_to_entry_sigma", +] + + +# Dial position -> entry sigma, obtained by inverting the measured change curve. +# Both endpoints remain exact: zero preserves the source and one starts from +# pure noise. Values between measured knots interpolate linearly. +_CALIBRATION: tp.Sequence[tuple[float, float]] = ( + (0.00, 0.0000), + (0.10, 0.3786), + (0.20, 0.5762), + (0.30, 0.6853), + (0.40, 0.7321), + (0.50, 0.7683), + (0.65, 0.8376), + (0.80, 0.8843), + (1.00, 1.0000), +) + + +def denoise_mapping_mode() -> str: + """Return ``calibrated`` (default) or the rollback mode ``identity``.""" + mode = os.environ.get("DEMON_SA3_DENOISE_MAPPING", "calibrated").strip().lower() + if mode not in ("calibrated", "identity"): + raise ValueError( + "DEMON_SA3_DENOISE_MAPPING must be calibrated|identity, " + f"got {mode!r}" + ) + return mode + + +def dial_to_entry_sigma(dial: float) -> float: + """Interpolate the measured monotonic mapping, clamped to ``[0, 1]``.""" + position = min(max(float(dial), 0.0), 1.0) + if position <= _CALIBRATION[0][0]: + return _CALIBRATION[0][1] + + for (p0, s0), (p1, s1) in zip(_CALIBRATION, _CALIBRATION[1:]): + if position <= p1: + fraction = (position - p0) / (p1 - p0) + return s0 + (s1 - s0) * fraction + + return _CALIBRATION[-1][1] + + +def map_denoise_to_entry_sigma(dial: float) -> float: + """Apply the selected product mapping without changing endpoint semantics.""" + clamped = min(max(float(dial), 0.0), 1.0) + if denoise_mapping_mode() == "identity": + return clamped + return dial_to_entry_sigma(clamped) diff --git a/acestep/streaming/sa3_backend.py b/acestep/streaming/sa3_backend.py index 4c6e0a09..0fd69482 100644 --- a/acestep/streaming/sa3_backend.py +++ b/acestep/streaming/sa3_backend.py @@ -143,10 +143,11 @@ def sa3_knob_specs(loras: tuple | list = ()) -> list: KnobSpec( "sa3_denoise", default=1.0, max_val=1.0, group="sa3", description=( - "SA3 init_noise_level: fresh-noise vs source-anchor mix " - "at slot init (1.0 = generate from pure noise, lower = " - "closer cover of the source). Distinct from ACE's " - "'denoise' (k1 strength), hence the prefix." + "Measured SA3 audio-change amount: 1.0 generates from pure " + "noise, while lower values stay progressively closer to the " + "source. Mapped onto init_noise_level so useful change is " + "spread across the dial. Distinct from ACE's 'denoise' " + "(k1 strength), hence the prefix." ), ), KnobSpec( diff --git a/tests/unit/test_sa3_denoise_mapping.py b/tests/unit/test_sa3_denoise_mapping.py new file mode 100644 index 00000000..7e0db2fa --- /dev/null +++ b/tests/unit/test_sa3_denoise_mapping.py @@ -0,0 +1,53 @@ +import pytest + +from acestep.engine.sa3_denoise_mapping import ( + denoise_mapping_mode, + dial_to_entry_sigma, + map_denoise_to_entry_sigma, +) + + +def test_mapping_is_monotonic_and_preserves_endpoints(): + values = [dial_to_entry_sigma(i / 100) for i in range(101)] + + assert values[0] == 0.0 + assert values[-1] == 1.0 + assert all(a <= b for a, b in zip(values, values[1:])) + + +def test_mapping_matches_measured_knots(): + assert dial_to_entry_sigma(0.1) == pytest.approx(0.3786) + assert dial_to_entry_sigma(0.5) == pytest.approx(0.7683) + assert dial_to_entry_sigma(0.8) == pytest.approx(0.8843) + + +def test_mapping_interpolates_between_knots(): + midpoint = (0.7683 + 0.8376) / 2 + assert dial_to_entry_sigma(0.575) == pytest.approx(midpoint) + + +def test_mapping_clamps_inputs(): + assert dial_to_entry_sigma(-1) == 0.0 + assert dial_to_entry_sigma(2) == 1.0 + + +def test_default_mapping_lifts_the_measured_dead_region(monkeypatch): + monkeypatch.delenv("DEMON_SA3_DENOISE_MAPPING", raising=False) + + assert denoise_mapping_mode() == "calibrated" + assert map_denoise_to_entry_sigma(0.2) > 0.5 + assert map_denoise_to_entry_sigma(0.5) > 0.7 + + +def test_identity_mode_is_a_bugfix_preserving_rollback(monkeypatch): + monkeypatch.setenv("DEMON_SA3_DENOISE_MAPPING", "identity") + + assert map_denoise_to_entry_sigma(0.2) == 0.2 + assert map_denoise_to_entry_sigma(0.5) == 0.5 + + +def test_invalid_mode_fails_loudly(monkeypatch): + monkeypatch.setenv("DEMON_SA3_DENOISE_MAPPING", "mystery") + + with pytest.raises(ValueError, match=r"calibrated\|identity"): + map_denoise_to_entry_sigma(0.5)