fix(cortex): placement waits for a neuron's own startup loads

Right after a deploy restarted benjy's neuron, a HeartMuLa job
cold-loaded while the neuron was still loading its default Qwen3-8B
(36 s). That load was in no model list and its memory had barely
started to fill, so the card looked empty; both landed and the HeartMuLa
load ran out of memory.

neuron's /health already reports its activation (pending, in_progress).
Placement now counts those models' catalogue footprint as taken, and a
cold-load that does not fit while the node is still activating answers
503 + Retry-After instead of loading over it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017ohK2pULMVJctL5B1Gixdp

.gitea/workflows/build-prerelease.yml

@@ -52,6 +52,11 @@ concurrency:
 env:
   CARGO_INCREMENTAL: "0"
   CARGO_TERM_COLOR: "always"
+  # Every RPM here is built for Fedora 43, but the runners build on their
+  # host, which may run a newer release. Each Stage step fails a binary that
+  # needs a newer glibc than Fedora 43 ships, rather than letting dnf skip
+  # the upgrade at deploy.
+  GLIBC_FLOOR: "2.42"
 
 jobs:
   prepare:
@@ -315,6 +320,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/cortex artifacts/cortex
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/cortex
           ./artifacts/cortex --version || true
 
       - uses: actions/upload-artifact@v3
@@ -354,6 +360,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/helexa-bench artifacts/helexa-bench
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-bench
           ./artifacts/helexa-bench --version || true
 
       - uses: actions/upload-artifact@v3
@@ -393,6 +400,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/helexa-upstream artifacts/helexa-upstream
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-upstream
           ./artifacts/helexa-upstream --version || true
 
       - uses: actions/upload-artifact@v3
@@ -432,6 +440,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/helexa-angels artifacts/helexa-angels
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-angels
           ./artifacts/helexa-angels --version || true
 
       - uses: actions/upload-artifact@v3
@@ -467,6 +476,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/helexa-router artifacts/helexa-router
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-router
           ./artifacts/helexa-router --version || true
 
       - uses: actions/upload-artifact@v3
@@ -502,6 +512,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/helexa-tools artifacts/helexa-tools
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-tools
           ./artifacts/helexa-tools --version || true
 
       - uses: actions/upload-artifact@v3
@@ -619,6 +630,7 @@ jobs:
         run: |
           mkdir --parents artifacts
           cp target/release/neuron artifacts/neuron-${{ matrix.flavour }}
+          script/check-glibc-floor.sh "$GLIBC_FLOOR" "artifacts/neuron-${{ matrix.flavour }}"
           file "artifacts/neuron-${{ matrix.flavour }}"
 
       - uses: actions/upload-artifact@v3

.gitea/workflows/ci.yml

@@ -242,10 +242,20 @@ jobs:
 
   test:
     name: Test
-    timeout-minutes: 25
+    # The compile alone took 19 of the old 25 minutes on main (2026-10-17,
+    # 75% Rust cache hits), so any growth in a workspace crate tipped a
+    # healthy run into a timeout. 45 leaves room without hiding a hang.
+    timeout-minutes: 45
     runs-on: rust
     needs: changes
     if: needs.changes.outputs.rust == 'true'
+    # Line tables only: full debuginfo for the whole graph made the neuron
+    # test binary's compile and link the memory peak of the job, and on the
+    # 15 GB rust runners a Test sharing the host with a Clippy was killed
+    # mid-compile with no log (2026-10-17). Backtraces keep file:line.
+    env:
+      CARGO_PROFILE_DEV_DEBUG: line-tables-only
+      CARGO_PROFILE_TEST_DEBUG: line-tables-only
     steps:
       - uses: actions/checkout@v4
       # See script/ci-cargo-escalate.sh for the escalation rationale.

.gitea/workflows/deploy.yml

@@ -183,6 +183,13 @@ jobs:
           # enable --now: start the service AND enable it for boot so the
           # fleet self-heals after a host reboot.
           sudo /usr/bin/systemctl enable --now cortex.service
+          # dnf can exit 0 without installing anything (e.g. an unmet glibc
+          # requirement), so compare what is now installed with what was published.
+          now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
+          if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
+            echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
+            exit 1
+          fi
           DEPLOY
 
       # Wait for the service to either come up or wedge, then capture
@@ -244,14 +251,12 @@ jobs:
             flavour: blackwell
             load_timeout: 900
             max_prompt_tokens: 131072
-          # benjy is out of the matrix while its boot NVMe is replaced
-          # (dead 2026-09-29): every deploy failed at SSH init against a
-          # host that is not there. Restore these four lines when it is
-          # back on the mesh.
-          # - host: benjy.hanzalova.internal
-          #   flavour: ada
-          #   load_timeout: 300
-          #   max_prompt_tokens: 16384
+          # benjy: the everyday class (Qwen3-8B) plus the cold-swap
+          # targets that take turns with it -- Z-Image and ACE-Step music.
+          - host: benjy.hanzalova.internal
+            flavour: ada
+            load_timeout: 300
+            max_prompt_tokens: 16384
           - host: quadbrat.hanzalova.internal
             flavour: ampere
             load_timeout: 300
@@ -464,6 +469,13 @@ jobs:
           # enable --now: start the service AND enable it for boot so the
           # fleet self-heals after a host reboot.
           sudo /usr/bin/systemctl enable --now neuron.service
+          # dnf can exit 0 without installing anything (e.g. an unmet glibc
+          # requirement), so compare what is now installed with what was published.
+          now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
+          if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
+            echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
+            exit 1
+          fi
 
           # ── Post-deploy validation ────────────────────────────────
           # A deploy only goes green if the neuron (a) finishes loading
@@ -508,10 +520,28 @@ jobs:
                   sys.stderr.write(msg + chr(10))
               sys.exit(1)
           completed = a.get("completed", [])
-          print(completed[0] if completed else "")
+          print(" ".join(completed))
           ')
+          # Probe the first default that answers chat. A decision model
+          # (Laya on quadbrat) or an image or audio model 422s a chat
+          # request, which made every quadbrat deploy report failure.
+          chat_models=$(curl -s --max-time 10 http://localhost:13131/models | python3 -c '
+          import json, sys
+          for m in json.load(sys.stdin):
+              caps = m.get("capabilities", [])
+              if "text" in caps and not {"decision", "image", "audio"} & set(caps):
+                  print(m["id"])
+          ')
+          probe_model=""
+          for candidate in ${model}; do
+            if printf '%s\n' "${chat_models}" | grep -qxF "${candidate}"; then
+              probe_model="${candidate}"
+              break
+            fi
+          done
+          model="${probe_model}"
           if [ -z "${model}" ]; then
-            echo "no default models configured — skipping LLM probe"
+            echo "no chat-serving default model — skipping LLM probe"
             exit 0
           fi
 
@@ -842,7 +872,15 @@ jobs:
             fi
             sudo /usr/bin/systemctl daemon-reload
             sudo /usr/bin/systemctl enable --now "${pkg}.service"
+            # dnf can exit 0 without installing anything (e.g. an unmet glibc
+            # requirement), so compare what is now installed with what was published.
+            now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
+            if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
+              echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
+              failed=1
+            fi
           done
+          [ "${failed:-0}" -eq 0 ]
           DEPLOY
 
       # Router config (#287). Tracked because it holds no secrets, and
@@ -1071,6 +1109,13 @@ jobs:
           # enable --now: start the service AND enable it for boot so the
           # bench resumes collecting after a host reboot.
           sudo /usr/bin/systemctl enable --now helexa-bench.service
+          # dnf can exit 0 without installing anything (e.g. an unmet glibc
+          # requirement), so compare what is now installed with what was published.
+          now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
+          if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
+            echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
+            exit 1
+          fi
 
           # ── Post-deploy validation ────────────────────────────────
           # The bench serves a read-only API on :13132 alongside the

CLAUDE.md

@@ -1226,9 +1226,8 @@ mean-vs-reference-envelope, not single cases.
   near-ties, probabilities within a bf16 tolerance.
 - The public edge is `https://helexa.ai/v1/systemone`. It is rate-limited
   (10 r/min per IP), so replay against the router (gallumbits:8088).
-  Neither edge is reachable from inside its own LAN (no NAT reflection;
-  hanzalova's presents the OPNsense certificate), so public probes must
-  go cross-site.
+  Probing the public names has network constraints; see
+  `doc/infra/edge-ingress.md` in the private helexa/security repo.
 
 **Sharp edges:**
 - **rustingface#6:** a neuron cold load of anything rf.internal hasn't
@@ -1243,3 +1242,86 @@ retries, a bench `decision:` scenario, a dashboard panel,
 micro-batching); #342 (serving fine-tuned checkpoints); #352 (Z-Image on
 the ampere flash build, unrelated to Laya but enabled by the same
 flavour change).
+
+## 2026-10-17 addendum: music generation (epics #378, #379)
+
+Music is the fourth modality. Two engines serve candle-natively on
+benjy's RTX 4090 through the device-worker discipline (`AudioHandle`
+slab; `LoadAudio` / `DropAudio` / `GenerateAudio` jobs; CPU-side sample
+replies, FLAC encoded async-side):
+
+- **ACE-Step 1.5** (MIT): turbo DiT + Qwen3-Embedding text encoder +
+  48 kHz stereo VAE, with an optional 1.7B planner LM. Vendored from
+  candle PR #3451 into `harness/arch/acestep/`. 10–600 s.
+- **HeartMuLa oss-3B** (Apache-2.0; Llama 3.2 tokenizer): an
+  autoregressive LM (bf16) plus HeartCodec (f32), from two sibling repos
+  (`HeartCodec-oss-20260123`, `HeartMuLaGen`). Detected by
+  `model_type: heartmula`. 10–240 s.
+
+**Surface: async jobs only.** `/v1/audio/generations`: POST submits and
+returns 202 with a job; GET lists; `/{id}` polls; `/{id}/content`
+fetches the audio; DELETE cancels. **cortex owns the jobs**, persisting
+them under `[audio_jobs] dir` and re-queueing unfinished ones on restart.
+neuron renders synchronously and returns the FLAC plus an
+`x-helexa-audio-result` header. The router dispatches POST by capacity
+and fans GET/DELETE out across cortexes until one answers non-404. Wire
+types are in `cortex-core/src/audio.rs`; limits (`[models.audio]`) and
+model-scoped licence notices (`[[models.notices]]`) come from the
+catalogue and are advertised in `/v1/models`.
+
+**Formats.** Audio is rendered straight to 16-bit FLAC (no WAV step) and
+stored as that lossless master. MP3 (320 kbps CBR, LAME) is a download
+format only: `/content?format=mp3` transcodes the FLAC once and caches
+`{id}.mp3` beside it until the job expires; `response_format` sets a
+job's default download. Results stored as WAV by older neurons are
+served as WAV. cortex refuses new jobs (503 `audio_storage_full`) while the job
+directory's filesystem has under `[audio_jobs] min_free_mb` free (10 GiB
+default); shared object storage for the Pascal-scale fleet is #411.
+
+**Metering: per audio-second**, at a per-model rate stored in
+helexa-upstream's `model_rates` table (`GET /authz/v1/rates`, cached by
+cortex's `rates.rs`; lookup is the model id, then `'*'`). Unpriced
+metered submissions get 503. A job settles at the rate it was reserved at.
+
+**Live numbers (benjy 4090, deployed build, 2026-10-17):**
+
+| | Resident | 10 s warm | 30 s | Cold swap from text |
+|---|---|---|---|---|
+| ACE-Step (planner loaded, off) | ~10.2 GB | 0.56 s | — | ~1 s load |
+| HeartMuLa | 14.6 GB | 4.3 s (LM 2.4 + codec 1.9) | 11.1 s | 8 s wall incl. load |
+
+Text restores in 2.5 s. Same seed gives byte-identical output on both
+engines, and the chain render is byte-identical to the direct neuron
+render.
+
+**Sharp edges:**
+
+- **ACE-Step needs DCW.** The Haar "double" correction after every turbo
+  step: without it, latents are 90% off upstream. Parity with the
+  upstream CPU f32 reference is ~1e-5 with it.
+- **Seeding is host-side** (SplitMix64 + Box-Muller): candle's CPU
+  `set_seed` is unsupported, and a host RNG makes CPU and GPU agree.
+- **The ACE planner checkpoint has no `model.` prefix**; it is renamed at
+  load. A planner that fails to load is non-fatal. It stays off unless
+  `planner: "on"` until #384 closes the parity gaps (code-phase CFG,
+  constrained decoding).
+- **Levels follow each upstream.** ACE-Step peak-normalises to −1 dBFS;
+  HeartMuLa does not.
+- **ACE timing is one number.** `text2music` is a single call, so
+  `generate_ms` covers encode + denoise + VAE.
+- **HeartMuLa's first decode after load is slow** (~33 s, one-time
+  warm-up); every decode after that is ~2 s per 10 s of audio.
+- **benjy's root drive is slow under pressure** (DRAM-less, 90% full;
+  ~8 MB/s flush). First fetches of 20 GB take ~25 min because each
+  shard's close() waits on writeback. Before #409 a client timeout
+  cancelled the load and threw the progress away.
+- **Validation:** `script/validate-audio.sh [host] [cortex] [model]`
+  (guarded window: evicts and restores co-resident models). helexa-bench
+  has an `audio:<secs>` scenario.
+
+**Retired:** the PyTorch evaluation beta (`beta.music.helexa.ai`, #388).
+Its ratings data is archived privately.
+
+**Next:** #384 (planner parity), #385 (q8 + an 8 GB profile), and the
+Pascal epic #380 (rig work waits until rigs are named; the fallback if
+ACE-Step cannot run on Pascal is Stable Audio and/or EzAudio).

Cargo.lock

@@ -40,7 +40,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cabdc9d845d08ec7ed2d0c9de1ae4a1b198301407d55855261572761be90ec9f"
 dependencies = [
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -261,7 +261,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -272,13 +272,13 @@ checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de"
 
 [[package]]
 name = "async-trait"
-version = "0.1.89"
+version = "0.1.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
+checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 3.0.9",
 ]
 
 [[package]]
@@ -311,6 +311,15 @@ version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
 
+[[package]]
+name = "autotools"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef941527c41b0fc0dd48511a8154cd5fc7e29200a0ff8b7203c5d777dbc795cf"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "aws-lc-rs"
 version = "1.16.2"
@@ -394,7 +403,7 @@ checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -485,6 +494,12 @@ dependencies = [
  "tinyvec",
 ]
 
+[[package]]
+name = "built"
+version = "0.7.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56ed6191a7e78c36abdb16ab65341eefd73d64d303fffccdbb00d51e4205967b"
+
 [[package]]
 name = "bumpalo"
 version = "3.20.2"
@@ -508,7 +523,7 @@ checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -697,7 +712,7 @@ dependencies = [
  "heck",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -706,6 +721,12 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
 
+[[package]]
+name = "claxon"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4bfbf56724aa9eca8afa4fcfadeb479e722935bb2a0900c2d37e0cc477af0688"
+
 [[package]]
 name = "cmake"
 version = "0.1.58"
@@ -844,15 +865,20 @@ dependencies = [
  "axum",
  "bytes",
  "chrono",
+ "claxon",
  "cortex-core",
  "eventsource-stream",
+ "flacenc",
  "futures",
  "helexa-stream",
  "metrics",
  "metrics-exporter-prometheus",
+ "mp3lame-encoder",
  "reqwest",
+ "rustix",
  "serde",
  "serde_json",
+ "tempfile",
  "thiserror 2.0.18",
  "tokio",
  "tokio-stream",
@@ -861,6 +887,7 @@ dependencies = [
  "tracing",
  "url",
  "urlencoding",
+ "uuid",
 ]
 
 [[package]]
@@ -896,6 +923,15 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "crossbeam-channel"
+version = "0.5.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1"
+dependencies = [
+ "crossbeam-utils",
+]
+
 [[package]]
 name = "crossbeam-deque"
 version = "0.8.6"
@@ -966,7 +1002,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "10a2a99df6e410a8ff4245aa2006499ea662245f967cc7c0a38c83ef8eb44dbf"
 dependencies = [
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1039,7 +1075,7 @@ dependencies = [
  "proc-macro2",
  "quote",
  "strsim",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1052,7 +1088,7 @@ dependencies = [
  "proc-macro2",
  "quote",
  "strsim",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1063,7 +1099,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
 dependencies = [
  "darling_core 0.20.11",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1074,7 +1110,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d"
 dependencies = [
  "darling_core 0.23.0",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1125,7 +1161,7 @@ dependencies = [
  "darling 0.20.11",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1135,7 +1171,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
 dependencies = [
  "derive_builder_core",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1157,7 +1193,7 @@ dependencies = [
  "proc-macro2",
  "quote",
  "rustc_version",
- "syn",
+ "syn 2.0.117",
  "unicode-xid",
 ]
 
@@ -1202,7 +1238,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1337,7 +1373,7 @@ dependencies = [
  "heck",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -1477,6 +1513,22 @@ version = "0.5.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
 
+[[package]]
+name = "flacenc"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74c892c2b5fa08f967e8b5ad29121570b4762202f2402033ce08479ec65eccd0"
+dependencies = [
+ "built",
+ "crc",
+ "crossbeam-channel",
+ "heapless",
+ "md-5",
+ "num-traits",
+ "rustversion",
+ "seq-macro",
+]
+
 [[package]]
 name = "flagset"
 version = "0.4.7"
@@ -1667,7 +1719,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -2043,6 +2095,15 @@ dependencies = [
  "zerocopy",
 ]
 
+[[package]]
+name = "hash32"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606"
+dependencies = [
+ "byteorder",
+]
+
 [[package]]
 name = "hashbrown"
 version = "0.12.3"
@@ -2106,6 +2167,16 @@ dependencies = [
  "hashbrown 0.15.5",
 ]
 
+[[package]]
+name = "heapless"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad"
+dependencies = [
+ "hash32",
+ "stable_deref_trait",
+]
+
 [[package]]
 name = "heck"
 version = "0.5.0"
@@ -3096,7 +3167,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -3109,6 +3180,26 @@ dependencies = [
  "pxfm",
 ]
 
+[[package]]
+name = "mp3lame-encoder"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60cb9bdd89806317373e36ff745f264b7ed7ffc5bc5aab02dc7d1b837c16a8d4"
+dependencies = [
+ "mp3lame-sys",
+]
+
+[[package]]
+name = "mp3lame-sys"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54e3b1772db47828840702e5a2e05694527f731abadf9b931355d54035f019d8"
+dependencies = [
+ "autotools",
+ "cc",
+ "libc",
+]
+
 [[package]]
 name = "native-tls"
 version = "0.2.18"
@@ -3139,10 +3230,12 @@ dependencies = [
  "candle-nn",
  "candle-transformers",
  "clap",
+ "claxon",
  "cortex-core",
  "cudaforge",
  "cudarc 0.19.8",
  "figment",
+ "flacenc",
  "futures",
  "half",
  "hf-hub",
@@ -3168,6 +3261,7 @@ dependencies = [
  "tracing-subscriber",
  "unicode-normalization",
  "unicode-properties",
+ "zip",
 ]
 
 [[package]]
@@ -3374,7 +3468,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -3473,7 +3567,7 @@ dependencies = [
  "proc-macro2",
  "proc-macro2-diagnostics",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -3542,7 +3636,7 @@ dependencies = [
  "phf_shared",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -3571,7 +3665,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -3694,7 +3788,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
 dependencies = [
  "proc-macro2",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -3714,7 +3808,7 @@ checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
  "version_check",
  "yansi",
 ]
@@ -4066,7 +4160,7 @@ checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4191,7 +4285,7 @@ dependencies = [
  "proc-macro2",
  "quote",
  "serde_json",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4394,7 +4488,7 @@ dependencies = [
  "proc-macro2",
  "quote",
  "serde_derive_internals",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4484,7 +4578,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4495,7 +4589,7 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4581,7 +4675,7 @@ dependencies = [
  "darling 0.23.0",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4814,7 +4908,7 @@ dependencies = [
  "quote",
  "sqlx-core",
  "sqlx-macros-core",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -4837,7 +4931,7 @@ dependencies = [
  "sqlx-mysql",
  "sqlx-postgres",
  "sqlx-sqlite",
- "syn",
+ "syn 2.0.117",
  "tokio",
  "url",
 ]
@@ -5022,7 +5116,7 @@ dependencies = [
  "heck",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -5042,6 +5136,17 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "syn"
+version = "3.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
 [[package]]
 name = "sync_wrapper"
 version = "1.0.2"
@@ -5059,7 +5164,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -5145,7 +5250,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -5156,7 +5261,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -5282,7 +5387,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -5440,7 +5545,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -5800,7 +5905,7 @@ dependencies = [
  "bumpalo",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
  "wasm-bindgen-shared",
 ]
 
@@ -5990,7 +6095,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -6001,7 +6106,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -6315,7 +6420,7 @@ dependencies = [
  "heck",
  "indexmap 2.14.0",
  "prettyplease",
- "syn",
+ "syn 2.0.117",
  "wasm-metadata",
  "wit-bindgen-core",
  "wit-component",
@@ -6331,7 +6436,7 @@ dependencies = [
  "prettyplease",
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
  "wit-bindgen-core",
  "wit-bindgen-rust",
 ]
@@ -6425,7 +6530,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
  "synstructure",
 ]
 
@@ -6437,7 +6542,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
  "synstructure",
 ]
 
@@ -6458,7 +6563,7 @@ checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]
@@ -6478,7 +6583,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
  "synstructure",
 ]
 
@@ -6518,7 +6623,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn",
+ "syn 2.0.117",
 ]
 
 [[package]]

asset/monitoring/README.md

@@ -52,21 +52,10 @@ curl -s 'http://localhost:26559/api/v1/targets' | jq '.data.activeTargets[]|sele
 
 `grafana-helexa-fleet.json` is a raw dashboard model with a `datasource`
 template variable, so it binds to whichever Prometheus data source you pick
-at import. Import via the API (creds in `/etc/grafana/grafana.env`):
-
-```sh
-# on golgafrinchans; GF_SECURITY_ADMIN_USER/PASSWORD live in the root-only
-# env file, Grafana is published on :28767
-set -a; source <(sudo cat /etc/grafana/grafana.env); set +a
-jq -n --slurpfile d grafana-helexa-fleet.json \
-  '{dashboard: $d[0], overwrite: true, folderUid: null}' \
-| curl -s -u "$GF_SECURITY_ADMIN_USER:$GF_SECURITY_ADMIN_PASSWORD" \
-    -H 'Content-Type: application/json' \
-    -d @- http://localhost:28767/api/dashboards/db | jq '{status,uid,version}'
-```
-
-The dashboard lands at uid `helexa-fleet`. Re-importing with `overwrite:true`
-updates it in place.
+at import. Import it through the Grafana dashboards API with admin
+credentials (`overwrite: true` updates it in place; it lands at uid
+`helexa-fleet`). The exact command and where the admin credentials live are
+in `doc/infra/monitoring-access.md` in the private helexa/security repo.
 
 ## What the dashboard shows
 

asset/nginx/angels.helexa.ai.conf

@@ -34,8 +34,8 @@ server {
 
 server {
     # TCP 443 on the edge is owned by the stream SNI router
-    # (streams-enabled/sni-router.conf, architecture reverse-proxies.md
-    # §5), which preads the SNI and hands every local name to this http
+    # (streams-enabled/sni-router.conf; see doc/infra/edge-ingress.md in
+    # the private helexa/security repo), which preads the SNI and hands every local name to this http
     # tier over PROXY protocol. Listening on :443 directly here would
     # never be reached — the router would answer first with whichever
     # vhost happened to be default, which is exactly the symptom seen

asset/nginx/bench.internal.conf

@@ -1,6 +1,7 @@
-# Internal bench UI vhost — https://bench.internal, reachable from inside
-# the WireGuard mesh (the public bench.helexa.ai dead-ends at the OPNsense
-# LAN interface, which only port-forwards :443 from the WAN). Same SPA +
+# Internal bench UI vhost — https://bench.internal, for clients inside
+# the WireGuard mesh, where the public bench.helexa.ai doesn't resolve
+# to a usable path (see doc/infra/edge-ingress.md in the private
+# helexa/security repo). Same SPA +
 # /api→bob proxy as bench.helexa.ai, but with an internal-CA cert
 # (smallstep "lair", renewed by step@bench.timer). Mirrors the
 # *.internal vhost convention on oolon.kosherinata.internal.

asset/nginx/helexa.ai.conf

@@ -4,7 +4,7 @@
 # Installed on BOTH edge proxies (oolon.kosherinata.internal and
 # hanzalova.internal); Cloudflare DNS load balancing routes the public
 # name to the two site WAN IPs. Mesh clients use helexa.internal instead
-# (public names don't hairpin — see architecture/reverse-proxies.md §2).
+# (see doc/infra/edge-ingress.md in the private helexa/security repo).
 #
 # Same-origin backends (#167): /v1 (+ /health) reverse-proxies to
 # helexa-router and /api to helexa-upstream, both on gallumbits — so the

asset/nginx/helexa.internal.conf

@@ -1,7 +1,7 @@
-# Internal helexa website vhost — https://helexa.internal, reachable from
-# inside the WireGuard mesh (the public helexa.ai dead-ends at the
-# OPNsense LAN interface, which only port-forwards :443 from the WAN —
-# see architecture/reverse-proxies.md §2). Same SPA webroot as the
+# Internal helexa website vhost — https://helexa.internal, for clients
+# inside the WireGuard mesh, where the public name doesn't resolve to a
+# usable path (see doc/infra/edge-ingress.md in the private
+# helexa/security repo). Same SPA webroot as the
 # public vhost, but with an internal-CA cert (smallstep "lair", renewed
 # by step@helexa.timer). Installed on hanzalova.internal only; the
 # split-horizon helexa.internal record points at hanzalova's mesh IP on

cortex.example.toml

@@ -113,3 +113,22 @@ enabled = false
 # timeout_secs = 5
 # How often to flush served-usage counters to upstream for reconciliation (#58).
 # served_usage_report_interval_secs = 60
+
+# -- Async audio generation jobs (#382) -----------------------------------
+# Music renders take seconds to minutes, so /v1/audio/generations is a job
+# API: submit, poll, fetch the audio. Job records and rendered audio live
+# under `dir` and survive a restart (the packaged unit's StateDirectory
+# creates /var/lib/cortex). All values below are the defaults.
+[audio_jobs]
+# dir = "/var/lib/cortex/audio-jobs"
+# Finished jobs and their audio are deleted after this many hours.
+# retention_hours = 72
+# Renders running at once across the fleet; the rest queue in cortex.
+# max_concurrent = 4
+# Dispatch attempts per job; a busy (429/503) model does not use one up.
+# max_attempts = 2
+# How long a job may wait on a busy model before failing, in seconds.
+# queue_timeout_secs = 3600
+# New jobs are refused (503) while the job directory's filesystem has less
+# than this much free space, in MiB. 0 disables the check.
+# min_free_mb = 10240

crates/cortex-cli/build.rs

@@ -0,0 +1,10 @@
+//! Keep the cortex binary installable on the oldest glibc we package for.
+//!
+//! glibc 2.43 gave `log10f` a new symbol version, so a binary linked on a
+//! newer host references `log10f@GLIBC_2.43` and will not install on a
+//! system with an older glibc. The only caller is LAME (the MP3 encoder
+//! behind `/content?format=mp3`). Routing its calls through
+//! `__wrap_log10f` (see `glibc_compat.rs`) removes the reference.
+fn main() {
+    println!("cargo:rustc-link-arg-bins=-Wl,--wrap=log10f");
+}

crates/cortex-cli/src/glibc_compat.rs

@@ -0,0 +1,29 @@
+//! Replacements for libm symbols whose current version is newer than the
+//! glibc cortex is packaged for. `build.rs` points the linker at these.
+
+/// `log10f` computed through `log10`, which keeps its original symbol
+/// version. Rounding the f64 result to f32 is at least as accurate as
+/// libm's single-precision version.
+#[unsafe(no_mangle)]
+extern "C" fn __wrap_log10f(x: f32) -> f32 {
+    (x as f64).log10() as f32
+}
+
+#[cfg(test)]
+mod tests {
+    use super::__wrap_log10f;
+
+    #[test]
+    fn matches_log10_including_edge_cases() {
+        for x in [1e-30_f32, 0.5, 1.0, 2.0, 10.0, 1234.5, f32::MAX] {
+            assert!(
+                (__wrap_log10f(x) - x.log10()).abs() <= 1e-6 * x.log10().abs().max(1.0),
+                "{x}"
+            );
+        }
+        assert_eq!(__wrap_log10f(0.0), f32::NEG_INFINITY);
+        assert!(__wrap_log10f(-1.0).is_nan());
+        assert!(__wrap_log10f(f32::NAN).is_nan());
+        assert_eq!(__wrap_log10f(f32::INFINITY), f32::INFINITY);
+    }
+}

crates/cortex-cli/src/main.rs

@@ -1,3 +1,5 @@
+mod glibc_compat;
+
 use anyhow::Result;
 use clap::{Parser, Subcommand};
 use cortex_core::config::GatewayConfig;

crates/cortex-core/src/audio.rs

@@ -0,0 +1,386 @@
+//! Music / audio generation surface (#382).
+//!
+//! OpenAI has no music endpoint, so this is a helexa extension. It follows
+//! the images surface where that makes sense (#197) but is **asynchronous**:
+//! a song takes seconds to minutes to render, so the public API is a job API.
+//!
+//! - `POST   /v1/audio/generations`               submit → `202` + [`AudioJob`]
+//! - `GET    /v1/audio/generations`               list the caller's jobs
+//! - `GET    /v1/audio/generations/{id}`          poll → [`AudioJob`]
+//! - `GET    /v1/audio/generations/{id}/content`  the rendered audio bytes
+//! - `DELETE /v1/audio/generations/{id}`          cancel
+//!
+//! cortex owns the job record. A neuron renders exactly one generation per
+//! dispatch, synchronously, on `POST /v1/audio/generations`: the response
+//! body is the audio (`audio/flac`) and the result metadata travels in the
+//! [`RESULT_HEADER`] header as JSON ([`NeuronAudioResult`]). Keeping the
+//! bytes out of a JSON envelope matters here — ten minutes of 48 kHz stereo
+//! is over 100 MB of PCM.
+//!
+//! Audio is rendered straight to FLAC (16-bit, lossless) and stored that
+//! way. MP3 is a download format only: cortex transcodes the stored FLAC
+//! when a client asks for it.
+
+use serde::{Deserialize, Serialize};
+
+/// Header carrying [`NeuronAudioResult`] JSON on a neuron's audio response.
+pub const RESULT_HEADER: &str = "x-helexa-audio-result";
+
+/// Capability string a model advertises when it serves this surface.
+pub const CAPABILITY: &str = "audio";
+
+/// `POST /v1/audio/generations` request body. The same body is accepted by
+/// cortex (which turns it into a job) and by neuron (which renders it).
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioGenerationRequest {
+    pub model: String,
+    /// Style / caption: what the music should sound like.
+    pub prompt: String,
+    /// Lyrics for models that sing. Absent, empty, or `[Instrumental]`
+    /// means an instrumental.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub lyrics: Option<String>,
+    /// Requested length in seconds. Defaults to the model's advertised
+    /// default; validated against its advertised range.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub duration: Option<f64>,
+    /// Fixed RNG seed for reproducible renders. The seed actually used is
+    /// always reported back, so a random render can be reproduced too.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub seed: Option<u64>,
+    /// Default download format for the job's content: `"flac"` (the
+    /// default, lossless) or `"mp3"`. The render itself is always FLAC; a
+    /// content fetch can still ask for either with `?format=`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub response_format: Option<String>,
+
+    // ── model-specific extensions ───────────────────────────────
+    /// Diffusion steps (diffusion models only).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub steps: Option<u32>,
+    /// Guidance / CFG scale.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub guidance_scale: Option<f64>,
+    /// Language of the lyrics, as a code such as `"en"` or `"es"`. Models
+    /// that use it default to `"unknown"`; a song planner may infer it.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub language: Option<String>,
+    /// Song planner: `"off"`, `"auto"` (model default) or `"on"`, for
+    /// models that have one (ACE-Step's LM).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub planner: Option<String>,
+    /// Sampling temperature (autoregressive models only).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub temperature: Option<f64>,
+    /// Top-k sampling (autoregressive models only).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub top_k: Option<u32>,
+}
+
+/// Response formats the surface delivers.
+pub const SUPPORTED_FORMATS: &[&str] = &["flac", "mp3"];
+
+/// The format audio is rendered and stored in.
+pub const STORED_FORMAT: &str = "flac";
+
+/// MIME type for an audio format name, including `"wav"` for results stored
+/// before audio moved to FLAC.
+pub fn content_type(format: &str) -> &'static str {
+    match format {
+        "flac" => "audio/flac",
+        "mp3" => "audio/mpeg",
+        "wav" => "audio/wav",
+        _ => "application/octet-stream",
+    }
+}
+
+impl AudioGenerationRequest {
+    /// True when the request asks for an instrumental (no singing).
+    pub fn is_instrumental(&self) -> bool {
+        match self.lyrics.as_deref().map(str::trim) {
+            None | Some("") => true,
+            Some(l) => l.eq_ignore_ascii_case("[instrumental]"),
+        }
+    }
+
+    /// Transport-level validation every hop can do without knowing the
+    /// model: required fields, sane numbers, a known format. Model-specific
+    /// limits ([`AudioLimits`]) are checked separately.
+    pub fn validate_shape(&self) -> Result<(), String> {
+        if self.model.trim().is_empty() {
+            return Err("'model' must be non-empty".into());
+        }
+        if self.prompt.trim().is_empty() {
+            return Err("'prompt' must be non-empty".into());
+        }
+        if let Some(d) = self.duration
+            && !(d.is_finite() && d > 0.0)
+        {
+            return Err(format!("'duration' must be a positive number, got {d}"));
+        }
+        if let Some(f) = self.response_format.as_deref()
+            && !SUPPORTED_FORMATS.contains(&f)
+        {
+            return Err(format!(
+                "response_format '{f}' is not supported; use one of {SUPPORTED_FORMATS:?}"
+            ));
+        }
+        if let Some(p) = self.planner.as_deref()
+            && !matches!(p, "off" | "auto" | "on")
+        {
+            return Err(format!("planner must be 'off', 'auto' or 'on', got '{p}'"));
+        }
+        Ok(())
+    }
+
+    /// Validate against a model's advertised limits.
+    pub fn validate_limits(&self, limits: &AudioLimits) -> Result<(), String> {
+        let d = self.duration.unwrap_or(limits.default_duration);
+        if d < limits.min_duration || d > limits.max_duration {
+            return Err(format!(
+                "duration {d}s is outside this model's range {}..={}s",
+                limits.min_duration, limits.max_duration
+            ));
+        }
+        if !limits.vocals && !self.is_instrumental() {
+            return Err("this model does not sing; omit 'lyrics' or use [Instrumental]".into());
+        }
+        if let (Some(s), Some(max)) = (self.steps, limits.max_steps)
+            && (s == 0 || s > max)
+        {
+            return Err(format!("steps must be in 1..={max}, got {s}"));
+        }
+        Ok(())
+    }
+}
+
+/// What a music model can do, advertised in `/v1/models` metadata so
+/// clients (and the `/music` page) can build their forms from it.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioLimits {
+    pub min_duration: f64,
+    pub max_duration: f64,
+    pub default_duration: f64,
+    /// Whether the model sings (accepts lyrics).
+    pub vocals: bool,
+    pub sample_rate: u32,
+    pub channels: u16,
+    /// Upper bound for `steps`, when the model takes one.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_steps: Option<u32>,
+}
+
+/// A notice an upstream licence asks to be passed on, shown beside the one
+/// model it applies to (never as a platform-wide licence).
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct ModelNotice {
+    pub text: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub url: Option<String>,
+}
+
+/// Lifecycle of a job.
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum JobStatus {
+    Queued,
+    Running,
+    Succeeded,
+    Failed,
+    Cancelled,
+}
+
+impl JobStatus {
+    pub fn is_terminal(self) -> bool {
+        matches!(self, Self::Succeeded | Self::Failed | Self::Cancelled)
+    }
+}
+
+/// The public job object returned by submit, poll, list and cancel.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioJob {
+    pub id: String,
+    /// Always `"audio.generation"`.
+    pub object: String,
+    pub status: JobStatus,
+    pub model: String,
+    /// The submitted style / caption, echoed so a job list is readable on
+    /// its own.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub prompt: Option<String>,
+    /// The submitted lyrics, if any.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub lyrics: Option<String>,
+    /// Unix seconds.
+    pub created_at: u64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub started_at: Option<u64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub completed_at: Option<u64>,
+    /// Requested duration, echoed so a client can show progress context.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub duration: Option<f64>,
+    /// Present when `status == succeeded`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub result: Option<AudioResult>,
+    /// Present when `status == succeeded`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub usage: Option<AudioUsage>,
+    /// Present when `status == failed`: the #63 envelope's inner `error`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub error: Option<serde_json::Value>,
+}
+
+pub const JOB_OBJECT: &str = "audio.generation";
+
+/// The rendered audio's properties.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioResult {
+    /// Actual length in seconds.
+    pub duration: f64,
+    pub sample_rate: u32,
+    pub channels: u16,
+    /// The seed the render used (random when the request gave none).
+    pub seed: u64,
+    /// Container of the stored audio (`"flac"`).
+    pub format: String,
+    /// Size of the stored audio in bytes.
+    pub bytes: u64,
+}
+
+/// Metering block. `helexa_audio_units` is audio-seconds generated (#386).
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioUsage {
+    pub helexa_audio_units: f64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub helexa_timing: Option<AudioTiming>,
+}
+
+/// Phase timing for one render, measured on the neuron.
+#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
+pub struct AudioTiming {
+    /// Text / lyric encoding.
+    pub encode_ms: u64,
+    /// Song planner (0 when it did not run).
+    pub plan_ms: u64,
+    /// Denoising, or autoregressive token generation.
+    pub generate_ms: u64,
+    /// Latent / codec decode to waveform.
+    pub decode_ms: u64,
+}
+
+/// What a neuron reports, in [`RESULT_HEADER`], beside the audio bytes.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct NeuronAudioResult {
+    pub duration: f64,
+    pub sample_rate: u32,
+    pub channels: u16,
+    pub seed: u64,
+    pub usage: AudioUsage,
+}
+
+/// `GET /v1/audio/generations` response.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AudioJobList {
+    /// Always `"list"`.
+    pub object: String,
+    pub data: Vec<AudioJob>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn req() -> AudioGenerationRequest {
+        serde_json::from_str(r#"{"model": "m", "prompt": "lofi"}"#).unwrap()
+    }
+
+    fn limits() -> AudioLimits {
+        AudioLimits {
+            min_duration: 10.0,
+            max_duration: 600.0,
+            default_duration: 30.0,
+            vocals: true,
+            sample_rate: 48_000,
+            channels: 2,
+            max_steps: Some(20),
+        }
+    }
+
+    #[test]
+    fn minimal_request_deserializes_and_validates() {
+        let r = req();
+        assert!(r.validate_shape().is_ok());
+        assert!(r.validate_limits(&limits()).is_ok());
+        assert!(r.is_instrumental());
+    }
+
+    #[test]
+    fn instrumental_detection() {
+        let mut r = req();
+        r.lyrics = Some("  [Instrumental] ".into());
+        assert!(r.is_instrumental());
+        r.lyrics = Some("[verse]\nla la".into());
+        assert!(!r.is_instrumental());
+    }
+
+    #[test]
+    fn shape_rejects_bad_values() {
+        let mut r = req();
+        r.prompt = " ".into();
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.duration = Some(f64::NAN);
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.duration = Some(-1.0);
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.response_format = Some("ogg".into());
+        assert!(r.validate_shape().is_err());
+        r.response_format = Some("wav".into());
+        assert!(r.validate_shape().is_err());
+        r.response_format = Some("mp3".into());
+        assert!(r.validate_shape().is_ok());
+        let mut r = req();
+        r.planner = Some("maybe".into());
+        assert!(r.validate_shape().is_err());
+    }
+
+    #[test]
+    fn limits_reject_out_of_range() {
+        let l = limits();
+        let mut r = req();
+        r.duration = Some(5.0);
+        assert!(r.validate_limits(&l).is_err());
+        r.duration = Some(601.0);
+        assert!(r.validate_limits(&l).is_err());
+        r.duration = Some(600.0);
+        assert!(r.validate_limits(&l).is_ok());
+        r.steps = Some(21);
+        assert!(r.validate_limits(&l).is_err());
+        r.steps = Some(0);
+        assert!(r.validate_limits(&l).is_err());
+    }
+
+    #[test]
+    fn lyrics_refused_by_non_singing_model() {
+        let mut l = limits();
+        l.vocals = false;
+        let mut r = req();
+        r.lyrics = Some("sing this".into());
+        assert!(r.validate_limits(&l).is_err());
+        r.lyrics = Some("[Instrumental]".into());
+        assert!(r.validate_limits(&l).is_ok());
+    }
+
+    #[test]
+    fn job_status_serializes_snake_case() {
+        assert_eq!(
+            serde_json::to_string(&JobStatus::Succeeded).unwrap(),
+            "\"succeeded\""
+        );
+        assert!(JobStatus::Cancelled.is_terminal());
+        assert!(!JobStatus::Running.is_terminal());
+    }
+}

crates/cortex-core/src/catalogue.rs

@@ -89,6 +89,15 @@ pub struct ModelProfile {
     /// are unioned with this set in the gateway's `/v1/models` response.
     #[serde(default)]
     pub capabilities: Vec<String>,
+    /// What an audio model can do (#382): duration range, vocals, sample
+    /// rate. Advertised in `/v1/models` and checked by cortex before a job
+    /// is accepted. `[models.audio]` in TOML.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub audio: Option<crate::audio::AudioLimits>,
+    /// Notices an upstream licence asks to be passed on for this model,
+    /// advertised in `/v1/models` beside it. `[[models.notices]]` in TOML.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub notices: Vec<crate::audio::ModelNotice>,
 }
 
 fn default_min_devices() -> u32 {
@@ -283,6 +292,8 @@ mod tests {
             limit: None,
             cost: None,
             capabilities: vec![],
+            audio: None,
+            notices: vec![],
         }
     }
 
@@ -531,3 +542,23 @@ source = "helexa"
         assert_eq!(cat.resolve_alias("helexa/large"), "Qwen/Qwen3.6-27B");
     }
 }
+
+#[cfg(test)]
+mod repo_catalogue {
+    /// The shipped catalogue parses, and the music model carries its limits.
+    #[test]
+    fn repo_models_toml_parses_with_audio_limits() {
+        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models.toml");
+        let cat = super::ModelCatalogue::load(&path);
+        let ace = cat
+            .get("ACE-Step/Ace-Step1.5")
+            .expect("ACE-Step in models.toml");
+        let audio = ace.audio.as_ref().expect("audio limits");
+        assert!(audio.vocals && audio.max_duration >= audio.default_duration);
+        let heart = cat
+            .get("HeartMuLa/HeartMuLa-oss-3B-happy-new-year")
+            .expect("HeartMuLa in models.toml");
+        assert!(heart.audio.is_some());
+        assert_eq!(heart.notices.len(), 2, "HeartMuLa's notices (#389)");
+    }
+}

crates/cortex-core/src/config.rs

@@ -28,6 +28,75 @@ pub struct GatewayConfig {
     /// — a single operator runs purely local.
     #[serde(default)]
     pub upstream: UpstreamClientConfig,
+    /// Async music/audio generation jobs (#382).
+    #[serde(default)]
+    pub audio_jobs: AudioJobsConfig,
+}
+
+/// `[audio_jobs]` — where cortex keeps async audio generation jobs (#382).
+///
+/// A job record and its rendered audio are written under `dir`, so finished
+/// results and the job list survive a cortex restart. If `dir` cannot be
+/// created the store runs in memory only and says so at startup.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AudioJobsConfig {
+    #[serde(default = "default_audio_jobs_dir")]
+    pub dir: String,
+    /// How long a finished job and its audio are kept, in hours.
+    #[serde(default = "default_audio_retention_hours")]
+    pub retention_hours: u64,
+    /// How many renders cortex runs at once across the fleet. Each one
+    /// occupies a neuron device for its whole duration; queued jobs wait
+    /// in cortex rather than piling onto neuron admission queues.
+    #[serde(default = "default_audio_max_concurrent")]
+    pub max_concurrent: usize,
+    /// Dispatch attempts per job before it fails. A busy neuron (429/503)
+    /// does not consume an attempt; an unreachable or failing one does.
+    #[serde(default = "default_audio_max_attempts")]
+    pub max_attempts: u32,
+    /// How long a job may wait for a busy model before it fails, in seconds.
+    #[serde(default = "default_audio_queue_timeout")]
+    pub queue_timeout_secs: u64,
+    /// New jobs are refused (503 + `Retry-After`) while the filesystem
+    /// holding `dir` has less than this much free space, in MiB. Rendered
+    /// audio is written there, so this keeps a busy store from filling the
+    /// disk it shares. 0 disables the check.
+    #[serde(default = "default_audio_min_free_mb")]
+    pub min_free_mb: u64,
+}
+
+impl Default for AudioJobsConfig {
+    fn default() -> Self {
+        Self {
+            dir: default_audio_jobs_dir(),
+            retention_hours: default_audio_retention_hours(),
+            max_concurrent: default_audio_max_concurrent(),
+            max_attempts: default_audio_max_attempts(),
+            queue_timeout_secs: default_audio_queue_timeout(),
+            min_free_mb: default_audio_min_free_mb(),
+        }
+    }
+}
+
+fn default_audio_jobs_dir() -> String {
+    "/var/lib/cortex/audio-jobs".into()
+}
+fn default_audio_retention_hours() -> u64 {
+    72
+}
+fn default_audio_max_concurrent() -> usize {
+    4
+}
+fn default_audio_max_attempts() -> u32 {
+    2
+}
+fn default_audio_queue_timeout() -> u64 {
+    3600
+}
+/// 10 GiB: room for well over a hundred of the longest songs while jobs
+/// already accepted finish.
+fn default_audio_min_free_mb() -> u64 {
+    10 * 1024
 }
 
 /// `[upstream]` — the helexa-upstream authority client (#57). Locally
@@ -200,6 +269,7 @@ impl Default for GatewayConfig {
             models_config: default_models_path(),
             entitlements: EntitlementsConfig::default(),
             upstream: UpstreamClientConfig::default(),
+            audio_jobs: AudioJobsConfig::default(),
         }
     }
 }

crates/cortex-core/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod anthropic;
+pub mod audio;
 pub mod build_info;
 pub mod catalogue;
 pub mod config;

crates/cortex-core/src/node.rs

@@ -186,6 +186,17 @@ pub struct CortexModelEntry {
     /// the absent (not priced) vs `0.0` (intentionally free) distinction.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub cost: Option<ModelCost>,
+    /// What an audio model can do (#382), from its catalogue profile:
+    /// duration range, vocals, sample rate. Clients build forms from it.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub audio: Option<crate::audio::AudioLimits>,
+    /// Ledger tokens charged per second of generated audio (#386), for
+    /// audio models with a configured rate.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub audio_tokens_per_second: Option<u64>,
+    /// Notices an upstream licence asks to be passed on for this model.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub notices: Vec<crate::audio::ModelNotice>,
     /// `true` when any neuron reports this model supports tool calls.
     #[serde(default)]
     pub tool_call: bool,

crates/cortex-gateway/Cargo.toml

@@ -27,6 +27,12 @@ eventsource-stream.workspace = true
 bytes = "1"
 urlencoding = "2"
 url = "2"
+uuid = { version = "1", features = ["v4"] }
+rustix = { version = "1", default-features = false, features = ["std", "fs"] }
+claxon = "0.4.3"
+mp3lame-encoder = { version = "0.2.5", features = ["std"] }
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["test-util"] }
+tempfile = "3"
+flacenc = { version = "0.5.1", default-features = false }

crates/cortex-gateway/src/audio_codec.rs

@@ -0,0 +1,119 @@
+//! Transcoding stored audio for download.
+//!
+//! Renders are stored as 16-bit FLAC, the lossless master. MP3 is offered
+//! as a download format because some players and users expect it; it is
+//! encoded from the master on request and cached beside it.
+
+use anyhow::{Context, bail};
+use mp3lame_encoder::{Bitrate, Builder, FlushNoGap, InterleavedPcm, Quality};
+
+/// Frames handed to the MP3 encoder per call, which bounds its output
+/// buffer for long tracks.
+const CHUNK_FRAMES: usize = 1152 * 64;
+
+/// Transcode 16-bit FLAC (mono or stereo) to MP3 at 320 kbps CBR, the
+/// highest rate MP3 has: it is a copy of a lossless master, so it should
+/// lose no more than the format forces it to.
+pub fn flac_to_mp3(flac: &[u8]) -> anyhow::Result<Vec<u8>> {
+    let mut reader =
+        claxon::FlacReader::new(std::io::Cursor::new(flac)).context("reading FLAC header")?;
+    let info = reader.streaminfo();
+    if info.bits_per_sample != 16 {
+        bail!("{}-bit FLAC is not supported", info.bits_per_sample);
+    }
+    let channels = info.channels as usize;
+    if !(1..=2).contains(&channels) {
+        bail!("{channels}-channel FLAC is not supported");
+    }
+    let mut pcm = Vec::with_capacity(info.samples.unwrap_or(0) as usize * channels);
+    for s in reader.samples() {
+        pcm.push(s.context("decoding FLAC")? as i16);
+    }
+    // LAME's interleaved input is stereo; duplicate a mono channel.
+    if channels == 1 {
+        pcm = pcm.iter().flat_map(|&s| [s, s]).collect();
+    }
+
+    let mut builder = Builder::new().context("creating MP3 encoder")?;
+    builder
+        .set_num_channels(2)
+        .map_err(|e| anyhow::anyhow!("MP3 channels: {e}"))?;
+    builder
+        .set_sample_rate(info.sample_rate)
+        .map_err(|e| anyhow::anyhow!("MP3 sample rate: {e}"))?;
+    builder
+        .set_brate(Bitrate::Kbps320)
+        .map_err(|e| anyhow::anyhow!("MP3 bitrate: {e}"))?;
+    builder
+        .set_quality(Quality::NearBest)
+        .map_err(|e| anyhow::anyhow!("MP3 quality: {e}"))?;
+    let mut encoder = builder
+        .build()
+        .map_err(|e| anyhow::anyhow!("MP3 encoder: {e}"))?;
+
+    let mut out = Vec::new();
+    for chunk in pcm.chunks(CHUNK_FRAMES * 2) {
+        out.reserve(mp3lame_encoder::max_required_buffer_size(chunk.len() / 2));
+        encoder
+            .encode_to_vec(InterleavedPcm(chunk), &mut out)
+            .map_err(|e| anyhow::anyhow!("MP3 encode: {e}"))?;
+    }
+    out.reserve(7200);
+    encoder
+        .flush_to_vec::<FlushNoGap>(&mut out)
+        .map_err(|e| anyhow::anyhow!("MP3 flush: {e}"))?;
+    Ok(out)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn flac(channels: usize, frames: usize) -> Vec<u8> {
+        use flacenc::component::BitRepr;
+        use flacenc::error::Verify;
+        let samples: Vec<i32> = (0..frames * channels)
+            .map(|i| ((i as f32 * 0.03).sin() * 12_000.0) as i32)
+            .collect();
+        let config = flacenc::config::Encoder::default().into_verified().unwrap();
+        let source = flacenc::source::MemSource::from_samples(&samples, channels, 16, 48_000);
+        let stream =
+            flacenc::encode_with_fixed_block_size(&config, source, config.block_size).unwrap();
+        let mut sink = flacenc::bitsink::ByteSink::new();
+        stream.write(&mut sink).unwrap();
+        sink.into_inner()
+    }
+
+    /// The first MPEG audio frame header: sync word, then MPEG-1 Layer III
+    /// at 320 kbps / 48 kHz.
+    fn first_frame(mp3: &[u8]) -> [u8; 4] {
+        let at = mp3
+            .windows(2)
+            .position(|w| w[0] == 0xFF && w[1] & 0xE0 == 0xE0)
+            .expect("no MPEG frame sync");
+        mp3[at..at + 4].try_into().unwrap()
+    }
+
+    #[test]
+    fn stereo_flac_becomes_320k_mp3() {
+        let mp3 = flac_to_mp3(&flac(2, 48_000 * 2)).unwrap();
+        let h = first_frame(&mp3);
+        assert_eq!(h[1] & 0x1E, 0x1A, "not MPEG-1 Layer III: {h:02x?}");
+        assert_eq!(h[2] >> 4, 0xE, "not 320 kbps: {h:02x?}");
+        assert_eq!((h[2] >> 2) & 0x3, 0x1, "not 48 kHz: {h:02x?}");
+        // Two seconds at 320 kbps is 80 kB, give or take a frame.
+        assert!((70_000..90_000).contains(&mp3.len()), "{} bytes", mp3.len());
+    }
+
+    #[test]
+    fn mono_flac_is_accepted() {
+        let mp3 = flac_to_mp3(&flac(1, 48_000)).unwrap();
+        assert!(mp3.len() > 30_000, "{} bytes", mp3.len());
+    }
+
+    #[test]
+    fn garbage_is_an_error_not_a_panic() {
+        assert!(flac_to_mp3(b"RIFF....WAVEfmt ").is_err());
+        assert!(flac_to_mp3(&[]).is_err());
+    }
+}

crates/cortex-gateway/src/evictor.rs

@@ -46,6 +46,7 @@ pub async fn evict_lru_on_node(
             .models
             .values()
             .filter(|m| m.status == ModelStatus::Loaded)
+            .filter(|m| !fleet.eviction_blocked(node, &m.id))
             .filter(|m| match incoming {
                 Some(inc) => fleet.catalogue.may_displace(inc, &m.id),
                 None => {

crates/cortex-gateway/src/handlers.rs

@@ -24,6 +24,7 @@ pub fn api_routes() -> Router<Arc<CortexState>> {
         .route("/v1/completions", post(completions))
         .route("/v1/responses", post(responses))
         .route("/v1/images/generations", post(images_generations))
+        .merge(crate::audio_jobs::routes())
         .route("/v1/systemone", post(systemone))
         .route("/v1/models", get(list_models))
         .route("/v1/messages", post(anthropic_messages))
@@ -977,6 +978,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                 // `cost` stays operator-set and flows from the catalogue.
                 limit: None,
                 cost: profile.cost.clone(),
+                audio: profile.audio.clone(),
+                audio_tokens_per_second: None,
+                notices: profile.notices.clone(),
                 // Runtime-detected — will be OR-ed in Pass 2 from neuron data.
                 tool_call: false,
                 reasoning: false,
@@ -1059,6 +1063,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                     capabilities: entry.capabilities.clone(),
                     limit: entry.limit.clone(),
                     cost: None,
+                    audio: None,
+                    audio_tokens_per_second: None,
+                    notices: Vec::new(),
                     tool_call: entry.tool_call,
                     reasoning: entry.reasoning,
                     max_model_len: None,
@@ -1128,6 +1135,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                     capabilities: discovered.get(model_id).cloned().unwrap_or_default(),
                     limit: None,
                     cost: None,
+                    audio: None,
+                    audio_tokens_per_second: None,
+                    notices: Vec::new(),
                     tool_call: false,
                     reasoning: false,
                     max_model_len: None,
@@ -1169,6 +1179,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                 capabilities: target_entry.capabilities,
                 limit: target_entry.limit.clone(),
                 cost: target_entry.cost.clone(),
+                audio: target_entry.audio.clone(),
+                audio_tokens_per_second: None,
+                notices: target_entry.notices.clone(),
                 tool_call: target_entry.tool_call,
                 reasoning: target_entry.reasoning,
                 max_model_len: None,
@@ -1186,10 +1199,19 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
     // from each entry's now-settled `limit`, so vLLM-convention clients
     // (Hermes Agent et al.) can read the window without knowing helexa's
     // `limit` schema.
+    //
+    // Audio models also carry their current rate (#386), so a client can
+    // estimate a job's cost before submitting it. An alias is priced as
+    // the model it names.
+    let rates = fleet.rates.snapshot().await;
     let data: Vec<Value> = entries
         .values_mut()
         .map(|e| {
             e.sync_flat_limit();
+            if e.audio.is_some() {
+                let priced_as = catalogue.resolve_alias(&e.id);
+                e.audio_tokens_per_second = rates.audio_tokens_per_second(priced_as);
+            }
             json!(e)
         })
         .collect();
@@ -1455,7 +1477,7 @@ fn context_length_exceeded_response(
 }
 
 /// Update `last_accessed` timestamp for a model on a node (drives LRU eviction).
-async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
+pub(crate) async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
     let mut nodes = fleet.nodes.write().await;
     if let Some(node) = nodes.get_mut(node_name)
         && let Some(entry) = node.models.get_mut(model_id)

crates/cortex-gateway/src/lib.rs

@@ -1,4 +1,6 @@
 pub mod anthropic_sse;
+pub mod audio_codec;
+pub mod audio_jobs;
 pub mod auth;
 pub mod entitlements_chain;
 pub mod entitlements_local;
@@ -10,6 +12,7 @@ pub mod metering;
 pub mod metrics;
 pub mod poller;
 pub mod proxy;
+pub mod rates;
 pub mod router;
 pub mod served_usage;
 pub mod state;
@@ -52,6 +55,15 @@ pub async fn run(config: GatewayConfig) -> Result<()> {
         poller::poll_loop(poller_fleet).await;
     });
 
+    // Per-model rates (#386): keep the cached table in step with upstream.
+    if config.upstream.enabled {
+        rates::spawn_refresh(Arc::clone(&fleet.rates), &config.upstream);
+    }
+
+    // Async audio jobs (#382): reload persisted jobs, re-queue unfinished
+    // ones, start the retention sweep.
+    audio_jobs::start(Arc::clone(&fleet)).await;
+
     // Spawn the evictor (reacts to VRAM pressure events from the router).
     let evictor_fleet = Arc::clone(&fleet);
     tokio::spawn(async move {

crates/cortex-gateway/src/metering.rs

@@ -228,6 +228,42 @@ pub fn settle_image_usage(
     guard.settle(tokens);
 }
 
+/// Audio-seconds to reserve for a job: the requested duration, else the
+/// model's advertised default, else a minute. Settle corrects to the actual
+/// length the neuron reports.
+pub fn audio_reservation_seconds(requested: Option<f64>, model_default: Option<f64>) -> f64 {
+    requested
+        .or(model_default)
+        .filter(|d| d.is_finite() && *d > 0.0)
+        .unwrap_or(60.0)
+}
+
+/// Convert audio-seconds to budget tokens at a per-model rate (#386),
+/// rounding up.
+pub fn audio_seconds_to_tokens(seconds: f64, tokens_per_second: u64) -> u64 {
+    (seconds * tokens_per_second as f64).ceil() as u64
+}
+
+/// Record per-principal audio spend and settle the reservation with the
+/// audio-seconds actually generated.
+pub fn settle_audio_usage(
+    principal: &Principal,
+    guard: ReservationGuard,
+    served_usage: &crate::served_usage::ServedUsage,
+    audio_seconds: f64,
+    tokens_per_second: u64,
+) {
+    let labels = [
+        ("account", principal.account_id.clone()),
+        ("key", principal.key_id.clone()),
+    ];
+    metrics::counter!("cortex_spend_audio_milliseconds_total", &labels)
+        .increment((audio_seconds * 1000.0).round() as u64);
+    let tokens = audio_seconds_to_tokens(audio_seconds, tokens_per_second);
+    served_usage.add(&principal.account_id, &principal.key_id, tokens);
+    guard.settle(tokens);
+}
+
 /// Map a [`BudgetError`] to the #63 envelope. The provider chose the window
 /// semantics; this only translates them to HTTP.
 fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {

crates/cortex-gateway/src/poller.rs

@@ -310,7 +310,7 @@ async fn discover_capabilities(fleet: &CortexState, name: &str, endpoint: &str)
 /// Fetch `/health` and stash the activation snapshot on NodeState.
 /// Decoupled from the /models poll so a /health glitch doesn't mark
 /// the neuron unhealthy or evict the model list.
-async fn poll_health(fleet: &CortexState, name: &str, endpoint: &str) {
+pub async fn poll_health(fleet: &CortexState, name: &str, endpoint: &str) {
     let url = format!("{endpoint}/health");
     let resp = match fleet
         .http_client

crates/cortex-gateway/src/rates.rs

@@ -0,0 +1,208 @@
+//! Per-model rates for work metered in units other than tokens (#386).
+//!
+//! The budget ledger counts tokens. Audio jobs are metered in seconds of
+//! generated audio, converted to ledger tokens at a per-model rate. The
+//! rates are configuration held by the clearing house (helexa-upstream's
+//! `model_rates` table), so pricing changes without a deploy: cortex
+//! fetches the table from `/authz/v1/rates`, caches it, and refreshes it in
+//! the background.
+//!
+//! Resolution: the model's own rate, else the unit's `"*"` default. A model
+//! with neither is not priced, and cortex refuses its metered work rather
+//! than serving it free. Until the first fetch succeeds nothing is priced,
+//! for the same reason.
+//!
+//! A cortex running without upstream has no clearing house to ask. It
+//! prices every audio-second at [`LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND`],
+//! which only matters for the operator's own local key caps.
+
+use cortex_core::config::UpstreamClientConfig;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::RwLock;
+
+/// Unit name for one second of generated audio, as stored upstream.
+pub const AUDIO_SECOND: &str = "audio_second";
+
+/// The default-rate model id.
+pub const ANY_MODEL: &str = "*";
+
+/// Audio rate used when this cortex runs without an upstream authority.
+pub const LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND: u64 = 1000;
+
+/// How often the rate table is refreshed from upstream.
+const REFRESH_INTERVAL: Duration = Duration::from_secs(60);
+
+/// Retry interval while no table has loaded yet.
+const INITIAL_RETRY: Duration = Duration::from_secs(5);
+
+/// A snapshot of the rate table: `(unit, model_id) → tokens per unit`.
+#[derive(Debug, Clone, Default)]
+pub struct RateTable {
+    rates: HashMap<(String, String), u64>,
+}
+
+impl RateTable {
+    /// The table a cortex without upstream uses.
+    pub fn local_default() -> Self {
+        let mut t = Self::default();
+        t.set(
+            AUDIO_SECOND,
+            ANY_MODEL,
+            LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND,
+        );
+        t
+    }
+
+    pub fn set(&mut self, unit: &str, model_id: &str, tokens_per_unit: u64) {
+        self.rates
+            .insert((unit.to_string(), model_id.to_string()), tokens_per_unit);
+    }
+
+    /// The model's rate for `unit`, else the unit's default, else `None`.
+    pub fn rate(&self, unit: &str, model_id: &str) -> Option<u64> {
+        self.rates
+            .get(&(unit.to_string(), model_id.to_string()))
+            .or_else(|| self.rates.get(&(unit.to_string(), ANY_MODEL.to_string())))
+            .copied()
+    }
+
+    /// Ledger tokens per second of audio for `model_id`.
+    pub fn audio_tokens_per_second(&self, model_id: &str) -> Option<u64> {
+        self.rate(AUDIO_SECOND, model_id)
+    }
+}
+
+/// The live rate table, held in `CortexState`.
+pub struct Rates {
+    table: RwLock<RateTable>,
+}
+
+impl Rates {
+    /// Rates for a cortex with no upstream: the local default.
+    pub fn local() -> Self {
+        Self {
+            table: RwLock::new(RateTable::local_default()),
+        }
+    }
+
+    /// Rates for a cortex with upstream: empty (nothing priced) until the
+    /// first fetch lands.
+    pub fn upstream() -> Self {
+        Self {
+            table: RwLock::new(RateTable::default()),
+        }
+    }
+
+    pub async fn snapshot(&self) -> RateTable {
+        self.table.read().await.clone()
+    }
+
+    pub async fn replace(&self, table: RateTable) {
+        *self.table.write().await = table;
+    }
+
+    pub async fn audio_tokens_per_second(&self, model_id: &str) -> Option<u64> {
+        self.table.read().await.audio_tokens_per_second(model_id)
+    }
+}
+
+#[derive(Deserialize)]
+struct RatesResp {
+    data: Vec<RateRow>,
+}
+
+#[derive(Deserialize)]
+struct RateRow {
+    model_id: String,
+    unit: String,
+    tokens_per_unit: i64,
+}
+
+/// Fetch the full rate table from upstream's `/authz/v1/rates`.
+pub async fn fetch(
+    client: &reqwest::Client,
+    base_url: &str,
+    bearer: &str,
+) -> anyhow::Result<RateTable> {
+    let url = format!("{}/authz/v1/rates", base_url.trim_end_matches('/'));
+    let resp = client
+        .get(&url)
+        .bearer_auth(bearer)
+        .send()
+        .await?
+        .error_for_status()?;
+    let body: RatesResp = resp.json().await?;
+    let mut table = RateTable::default();
+    for row in body.data {
+        // The database forbids non-positive rates; skip one rather than
+        // trust it if that ever changes.
+        if let Ok(tokens) = u64::try_from(row.tokens_per_unit)
+            && tokens > 0
+        {
+            table.set(&row.unit, &row.model_id, tokens);
+        }
+    }
+    Ok(table)
+}
+
+/// Keep `rates` in step with upstream. Retries quickly until the first
+/// table loads, then refreshes every minute. A failed refresh keeps the
+/// last good table.
+pub fn spawn_refresh(rates: Arc<Rates>, cfg: &UpstreamClientConfig) {
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(cfg.timeout_secs.max(1)))
+        .build()
+        .expect("failed to build rates HTTP client");
+    let url = cfg.url.clone();
+    let bearer = cfg.bearer.clone();
+    tokio::spawn(async move {
+        let mut loaded = false;
+        loop {
+            match fetch(&client, &url, &bearer).await {
+                Ok(table) => {
+                    if !loaded {
+                        tracing::info!("model rates loaded from upstream");
+                    }
+                    loaded = true;
+                    rates.replace(table).await;
+                }
+                Err(e) => tracing::warn!(error = %e, "model rate refresh failed"),
+            }
+            tokio::time::sleep(if loaded {
+                REFRESH_INTERVAL
+            } else {
+                INITIAL_RETRY
+            })
+            .await;
+        }
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn model_rate_beats_default_and_absent_unit_is_unpriced() {
+        let mut t = RateTable::default();
+        assert_eq!(t.audio_tokens_per_second("m"), None);
+        t.set(AUDIO_SECOND, ANY_MODEL, 1000);
+        assert_eq!(t.audio_tokens_per_second("m"), Some(1000));
+        t.set(AUDIO_SECOND, "m", 250);
+        assert_eq!(t.audio_tokens_per_second("m"), Some(250));
+        assert_eq!(t.audio_tokens_per_second("other"), Some(1000));
+        assert_eq!(t.rate("some_other_unit", "m"), None);
+    }
+
+    #[test]
+    fn local_default_prices_every_audio_model() {
+        let t = RateTable::local_default();
+        assert_eq!(
+            t.audio_tokens_per_second("anything"),
+            Some(LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND)
+        );
+    }
+}

crates/cortex-gateway/src/router.rs

@@ -64,6 +64,10 @@ pub enum RouteError {
         "model '{model_id}' is recovering on node '{node}' (device context rebuild in progress) — retry shortly"
     )]
     ModelRecovering { model_id: String, node: String },
+    #[error(
+        "model '{model_id}' does not fit on '{node}' until a model there finishes its work — retry shortly"
+    )]
+    NodeBusy { model_id: String, node: String },
 }
 
 impl RouteError {
@@ -74,7 +78,8 @@ impl RouteError {
         match self {
             RouteError::NoHealthyNodes
             | RouteError::ModelRecovering { .. }
-            | RouteError::FeasibleNodeUnhealthy { .. } => 503,
+            | RouteError::FeasibleNodeUnhealthy { .. }
+            | RouteError::NodeBusy { .. } => 503,
             _ => 404,
         }
     }
@@ -88,7 +93,8 @@ impl RouteError {
             | RouteError::NoFeasibleNeuron { .. }
             | RouteError::ColdLoadFailed { .. }
             | RouteError::ModelRecovering { .. }
-            | RouteError::FeasibleNodeUnhealthy { .. } => "api_error",
+            | RouteError::FeasibleNodeUnhealthy { .. }
+            | RouteError::NodeBusy { .. } => "api_error",
         }
     }
 
@@ -102,6 +108,7 @@ impl RouteError {
             RouteError::ColdLoadFailed { .. } => "service_unavailable",
             RouteError::ModelRecovering { .. } => "service_unavailable",
             RouteError::FeasibleNodeUnhealthy { .. } => "service_unavailable",
+            RouteError::NodeBusy { .. } => "service_unavailable",
         }
     }
 
@@ -115,6 +122,8 @@ impl RouteError {
             RouteError::ModelRecovering { .. } => Some(2),
             RouteError::FeasibleNodeUnhealthy { .. } => Some(3),
             RouteError::NoHealthyNodes => Some(5),
+            // Clears when the resident model's requests finish.
+            RouteError::NodeBusy { .. } => Some(10),
             _ => None,
         }
     }
@@ -205,11 +214,10 @@ pub async fn resolve(
                     // Loading is gateway-synthesised from neuron's
                     // activation snapshot; it never appears on the
                     // wire from neuron's `/models`. Skip — the model
-                    // isn't actually servable yet. The pre-existing
-                    // race (catalogue cold_load fires a parallel
-                    // /models/load against the in-flight load) is no
-                    // worse than before; fixing it needs neuron-side
-                    // in-flight tracking on /models/load itself.
+                    // isn't actually servable yet. Cortex's own
+                    // cold-loads are serialised per node (priority 4);
+                    // a load neuron started by itself at activation is
+                    // still invisible to that lock.
                     ModelStatus::Loading => {}
                 }
             }
@@ -247,41 +255,109 @@ pub async fn resolve(
 
     // Priority 4: catalogue × topology cold-load.
     if let Some(profile) = fleet.catalogue.get(model_id) {
-        let (node_name, neuron_endpoint, fits_free) = pick_feasible_neuron(fleet, profile).await?;
-        // Cold-swap (#203): when the chosen node's devices don't have the
-        // free VRAM *right now*, evict unpinned LRU models there until
-        // the profile fits (bounded — a node whose evictable set can't
-        // free enough was ranked below one that can).
-        if !fits_free {
-            for _ in 0..3 {
-                match crate::evictor::evict_lru_on_node(fleet, &node_name, Some(&profile.id)).await
-                {
-                    Ok(Some(evicted)) => {
-                        tracing::info!(
-                            model = %profile.id,
-                            node = %node_name,
-                            evicted = %evicted,
-                            "cold-swap: evicted LRU model to make room"
-                        );
-                        if node_fits_free(fleet, &node_name, profile).await {
-                            break;
-                        }
-                    }
-                    Ok(None) => break,
-                    Err(e) => {
-                        tracing::warn!(node = %node_name, error = %e, "cold-swap eviction failed");
-                        break;
-                    }
+        let (node_name, neuron_endpoint, _) = pick_feasible_neuron(fleet, profile).await?;
+        // One cold-load per node at a time. The choice above rests on the
+        // last poll's VRAM, which cannot see a load in progress, so every
+        // decision that follows is re-made inside the lock from a fresh
+        // `/health`.
+        let lock = fleet.cold_load_lock(&node_name);
+        let _turn = lock.lock().await;
+
+        // Whoever held the lock may have been loading this very model.
+        if let Some((node, endpoint)) = loaded_location(fleet, model_id).await {
+            return finish(fleet, &node, &endpoint, model_id, false).await;
+        }
+
+        crate::poller::poll_health(fleet, &node_name, &neuron_endpoint).await;
+        // Cold-swap (#203): when the node's devices don't have the free
+        // VRAM right now, evict unpinned LRU models there until the
+        // profile fits (bounded — a node whose evictable set can't free
+        // enough was ranked below one that can).
+        for _ in 0..3 {
+            if node_fits_free(fleet, &node_name, profile).await {
+                break;
+            }
+            match crate::evictor::evict_lru_on_node(fleet, &node_name, Some(&profile.id)).await {
+                Ok(Some(evicted)) => {
+                    tracing::info!(
+                        model = %profile.id,
+                        node = %node_name,
+                        evicted = %evicted,
+                        "cold-swap: evicted LRU model to make room"
+                    );
+                    crate::poller::poll_health(fleet, &node_name, &neuron_endpoint).await;
+                }
+                Ok(None) => break,
+                Err(e) => {
+                    tracing::warn!(node = %node_name, error = %e, "cold-swap eviction failed");
+                    break;
                 }
             }
         }
+        // Still short, and a model that would otherwise be displaced is
+        // busy: loading now would land beside it and fail out of memory.
+        // Answer "retry shortly" instead; the room appears when its work
+        // ends.
+        if !node_fits_free(fleet, &node_name, profile).await
+            && blocked_displaceable(fleet, &node_name, profile).await
+        {
+            tracing::info!(
+                model = %profile.id,
+                node = %node_name,
+                "cold-load deferred: the model it would displace is busy"
+            );
+            return Err(RouteError::NodeBusy {
+                model_id: profile.id.clone(),
+                node: node_name,
+            });
+        }
         cold_load(fleet, &node_name, &neuron_endpoint, profile).await?;
+        fleet.note_cold_load(&node_name, &profile.id);
+        // The next cold-load in line decides from these readings.
+        crate::poller::poll_health(fleet, &node_name, &neuron_endpoint).await;
+        drop(_turn);
         return finish(fleet, &node_name, &neuron_endpoint, model_id, true).await;
     }
 
     Err(RouteError::ModelNotFound(model_id.to_string()))
 }
 
+/// A healthy node where `model_id` is loaded and servable, if any.
+async fn loaded_location(fleet: &Arc<CortexState>, model_id: &str) -> Option<(String, String)> {
+    let nodes = fleet.nodes.read().await;
+    nodes
+        .values()
+        .filter(|n| n.healthy)
+        .find(|n| {
+            n.models.get(model_id).is_some_and(|m| {
+                matches!(m.status, ModelStatus::Loaded | ModelStatus::Reloading) && m.is_servable()
+            })
+        })
+        .map(|n| (n.name.clone(), n.endpoint.clone()))
+}
+
+/// Whether `node_name` holds a model the profile may displace but that
+/// cannot be evicted right now because it is busy (see
+/// [`CortexState::eviction_blocked`]), or is still loading its startup
+/// models.
+async fn blocked_displaceable(
+    fleet: &Arc<CortexState>,
+    node_name: &str,
+    profile: &ModelProfile,
+) -> bool {
+    let nodes = fleet.nodes.read().await;
+    nodes.get(node_name).is_some_and(|node| {
+        // A startup load cannot be evicted until it finishes, and its
+        // memory is still arriving: wait for it.
+        activating(node)
+            || node.models.values().any(|m| {
+                m.status == ModelStatus::Loaded
+                    && fleet.catalogue.may_displace(&profile.id, &m.id)
+                    && fleet.eviction_blocked(node, &m.id)
+            })
+    })
+}
+
 /// True when `node_name`'s live device readings show enough free VRAM
 /// on some device for the profile. Conservative: unknown health (no
 /// /health poll yet) counts as not-fitting so the ranking prefers
@@ -289,12 +365,40 @@ pub async fn resolve(
 async fn node_fits_free(fleet: &Arc<CortexState>, node_name: &str, profile: &ModelProfile) -> bool {
     let nodes = fleet.nodes.read().await;
     nodes.get(node_name).is_some_and(|node| {
+        let reserved = activation_reserved_mb(fleet, node);
         node.device_health
             .iter()
-            .any(|d| d.vram_free_mb >= profile.vram_mb.unwrap_or(0))
+            .any(|d| d.vram_free_mb.saturating_sub(reserved) >= profile.vram_mb.unwrap_or(0))
     })
 }
 
+/// VRAM promised to models the neuron is loading on its own at startup
+/// (its `default_models`), by their catalogue footprint. Those loads are
+/// not in the neuron's model list until they finish and their memory
+/// fills in gradually, so free VRAM alone overstates what is available.
+fn activation_reserved_mb(fleet: &CortexState, node: &cortex_core::node::NodeState) -> u64 {
+    node.activation.as_ref().map_or(0, |a| {
+        a.in_progress
+            .iter()
+            .chain(a.pending.iter())
+            .filter(|id| {
+                !node
+                    .models
+                    .get(id.as_str())
+                    .is_some_and(|m| m.status == ModelStatus::Loaded)
+            })
+            .filter_map(|id| fleet.catalogue.get(id).and_then(|p| p.vram_mb))
+            .sum()
+    })
+}
+
+/// Whether the neuron is still loading its startup models.
+fn activating(node: &cortex_core::node::NodeState) -> bool {
+    node.activation
+        .as_ref()
+        .is_some_and(|a| a.in_progress.is_some() || !a.pending.is_empty())
+}
+
 /// Pick a healthy neuron whose discovered topology satisfies the
 /// profile. Preference order (#203):
 ///   1. A neuron from `profile.pinned_on` that is healthy + feasible.
@@ -332,7 +436,8 @@ async fn pick_feasible_neuron(
             .iter()
             .map(|d| d.vram_free_mb)
             .max()
-            .unwrap_or(0);
+            .unwrap_or(0)
+            .saturating_sub(activation_reserved_mb(fleet, node));
         let need = profile.vram_mb.unwrap_or(0);
         let fits_free = max_free >= need;
         // Evictable estimate: free + the VRAM of loaded models this
@@ -344,6 +449,7 @@ async fn pick_feasible_neuron(
             .values()
             .filter(|m| {
                 matches!(m.status, cortex_core::node::ModelStatus::Loaded)
+                    && !fleet.eviction_blocked(node, &m.id)
                     && fleet.catalogue.may_displace(&profile.id, &m.id)
             })
             // neuron reports vram_used_mb: null today; fall back to the
@@ -663,6 +769,8 @@ mod tests {
             limit: None,
             cost: None,
             capabilities: vec![],
+            audio: None,
+            notices: vec![],
         }
     }
 

crates/cortex-gateway/src/state.rs

@@ -52,8 +52,27 @@ pub struct CortexState {
     /// rather than an edge — it hid a GGUF that two of three nodes could
     /// describe perfectly well.
     pub capability_probe_attempts: RwLock<HashMap<(String, String), std::time::Instant>>,
+    /// Async audio generation jobs (#382).
+    pub audio_jobs: crate::audio_jobs::AudioJobs,
+    /// Per-model rates for non-token units (#386), from upstream when it
+    /// is enabled, else the local default.
+    pub rates: Arc<crate::rates::Rates>,
+    /// One lock per node, held across a cold-load's evict → load → fresh
+    /// `/health` sequence. Placement reads VRAM from the last poll, which
+    /// cannot see a load still in progress; two cold-loads deciding at
+    /// once both saw the same free card and loaded side by side into an
+    /// out-of-memory.
+    cold_load_locks: std::sync::Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
+    /// When cortex last cold-loaded each `(node, model)`. A model loaded
+    /// for a request is protected from eviction for [`COLD_LOAD_GRACE`],
+    /// long enough for that request to reach it; otherwise the next
+    /// cold-load in line evicts it before it has served anything.
+    recent_cold_loads: std::sync::Mutex<HashMap<(String, String), std::time::Instant>>,
 }
 
+/// How long a freshly cold-loaded model is protected from eviction.
+pub const COLD_LOAD_GRACE: std::time::Duration = std::time::Duration::from_secs(30);
+
 impl CortexState {
     pub fn from_config(config: &GatewayConfig) -> Self {
         let mut nodes = HashMap::new();
@@ -131,6 +150,59 @@ impl CortexState {
             served_usage: Arc::new(crate::served_usage::ServedUsage::new()),
             discovered_capabilities: RwLock::new(HashMap::new()),
             capability_probe_attempts: RwLock::new(HashMap::new()),
+            audio_jobs: crate::audio_jobs::AudioJobs::new(&config.audio_jobs),
+            rates: Arc::new(if config.upstream.enabled {
+                crate::rates::Rates::upstream()
+            } else {
+                crate::rates::Rates::local()
+            }),
+            cold_load_locks: std::sync::Mutex::new(HashMap::new()),
+            recent_cold_loads: std::sync::Mutex::new(HashMap::new()),
         }
     }
+
+    /// The lock that serialises cold-loads on `node`.
+    pub fn cold_load_lock(&self, node: &str) -> Arc<tokio::sync::Mutex<()>> {
+        let mut locks = self
+            .cold_load_locks
+            .lock()
+            .unwrap_or_else(|e| e.into_inner());
+        Arc::clone(locks.entry(node.to_string()).or_default())
+    }
+
+    /// Record that cortex has just cold-loaded `model` on `node`.
+    pub fn note_cold_load(&self, node: &str, model: &str) {
+        let mut recent = self
+            .recent_cold_loads
+            .lock()
+            .unwrap_or_else(|e| e.into_inner());
+        recent.retain(|_, at| at.elapsed() < COLD_LOAD_GRACE);
+        recent.insert(
+            (node.to_string(), model.to_string()),
+            std::time::Instant::now(),
+        );
+    }
+
+    /// Whether `model` on `node` is inside its post-cold-load grace.
+    pub fn recently_cold_loaded(&self, node: &str, model: &str) -> bool {
+        let recent = self
+            .recent_cold_loads
+            .lock()
+            .unwrap_or_else(|e| e.into_inner());
+        recent
+            .get(&(node.to_string(), model.to_string()))
+            .is_some_and(|at| at.elapsed() < COLD_LOAD_GRACE)
+    }
+
+    /// Whether a resident model must not be evicted right now: it has
+    /// requests running or queued (as of the last `/health`), or cortex
+    /// cold-loaded it moments ago for a request that has yet to arrive.
+    /// Evicting either kills work in progress.
+    pub fn eviction_blocked(&self, node: &cortex_core::node::NodeState, model: &str) -> bool {
+        let busy = node
+            .model_load
+            .get(model)
+            .is_some_and(|l| l.in_flight + l.queue_depth > 0);
+        busy || self.recently_cold_loaded(&node.name, model)
+    }
 }

crates/cortex-gateway/tests/aliases.rs

@@ -87,6 +87,7 @@ async fn test_alias_resolves_in_chat_completions() {
         models_config: models_path.to_string_lossy().to_string(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -176,6 +177,7 @@ async fn test_aliases_surface_in_v1_models() {
         models_config: models_path.to_string_lossy().to_string(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -268,6 +270,7 @@ async fn test_alias_falls_through_for_unmapped_model() {
         models_config: models_path.to_string_lossy().to_string(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));

crates/cortex-gateway/tests/auth.rs

@@ -106,6 +106,7 @@ async fn spawn_gateway(neuron_url: &str, entitlements: EntitlementsConfig) -> St
         models_config: "/dev/null".into(),
         entitlements,
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));

crates/cortex-gateway/tests/budget_enforcement.rs

@@ -82,6 +82,7 @@ async fn spawn_gateway(neuron_url: &str, key: ApiKeyConfig) -> (Arc<CortexState>
             keys: vec![key],
         },
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = Arc::new(CortexState::from_config(&config));
     {

crates/cortex-gateway/tests/cold_load.rs

@@ -0,0 +1,308 @@
+//! Cold-load placement under concurrency.
+//!
+//! The mock neuron keeps real VRAM books for one 24 GB card: loads take
+//! time, `/health` reports what is actually resident, and the card notes
+//! if it is ever asked to hold more than it has. That last flag is the
+//! failure being guarded against — two cold-loads deciding from the same
+//! stale poll, both landing, and every request on the card then failing
+//! out of memory.
+
+use axum::extract::{Path, State};
+use axum::routing::{get, post};
+use axum::{Json, Router};
+use cortex_core::config::{
+    EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
+};
+use cortex_core::discovery::{DeviceInfo, DiscoveryResponse};
+use cortex_core::node::{ModelEntry, ModelStatus};
+use cortex_gateway::router::RouteError;
+use cortex_gateway::state::CortexState;
+use serde_json::{Value, json};
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+const CARD_MB: u64 = 24_000;
+
+const CATALOGUE: &str = r#"
+[[models]]
+id = "music-a"
+harness = "candle"
+vram_mb = 14000
+
+[[models]]
+id = "music-b"
+harness = "candle"
+vram_mb = 12000
+"#;
+
+#[derive(Default)]
+struct Card {
+    resident: HashMap<String, u64>,
+    in_flight: HashMap<String, usize>,
+    overcommitted: bool,
+    unloads: Vec<String>,
+    /// A model the neuron is loading by itself at startup: not resident
+    /// yet, its memory not yet showing.
+    activating: Option<String>,
+    loads: Vec<String>,
+}
+
+fn footprint(model: &str) -> u64 {
+    match model {
+        "music-a" => 14_000,
+        "music-b" => 12_000,
+        _ => 1_000,
+    }
+}
+
+async fn spawn_neuron(card: Arc<Mutex<Card>>) -> String {
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let url = format!("http://{}", listener.local_addr().unwrap());
+    let inference = url.clone();
+    let app = Router::new()
+        .route(
+            "/models/load",
+            post(
+                |State(card): State<Arc<Mutex<Card>>>, Json(spec): Json<Value>| async move {
+                    let id = spec["model_id"].as_str().unwrap().to_string();
+                    // A load takes a while: long enough for a second
+                    // request to decide while this one is in progress.
+                    tokio::time::sleep(Duration::from_millis(300)).await;
+                    let mut card = card.lock().unwrap();
+                    card.loads.push(id.clone());
+                    card.resident.insert(id.clone(), footprint(&id));
+                    if card.resident.values().sum::<u64>() > CARD_MB {
+                        card.overcommitted = true;
+                    }
+                    Json(json!({"status": "loaded"}))
+                },
+            ),
+        )
+        .route(
+            "/models/unload",
+            post(
+                |State(card): State<Arc<Mutex<Card>>>, Json(body): Json<Value>| async move {
+                    let id = body["model_id"].as_str().unwrap().to_string();
+                    let mut card = card.lock().unwrap();
+                    card.resident.remove(&id);
+                    card.unloads.push(id);
+                    Json(json!({"status": "unloaded"}))
+                },
+            ),
+        )
+        .route(
+            "/health",
+            get(|State(card): State<Arc<Mutex<Card>>>| async move {
+                let card = card.lock().unwrap();
+                let used: u64 = card.resident.values().sum();
+                let models: Vec<Value> = card
+                    .resident
+                    .keys()
+                    .map(|id| {
+                        json!({
+                            "id": id,
+                            "in_flight": card.in_flight.get(id).copied().unwrap_or(0),
+                            "queue_depth": 0
+                        })
+                    })
+                    .collect();
+                let activation = match &card.activating {
+                    Some(id) => json!({"state": "pre_warming", "in_progress": id}),
+                    None => json!({"state": "ready"}),
+                };
+                Json(json!({
+                    "uptime_secs": 1,
+                    "activation": activation,
+                    "devices": [{
+                        "index": 0,
+                        "vram_used_mb": used,
+                        "vram_free_mb": CARD_MB.saturating_sub(used),
+                        "utilization_pct": 0,
+                        "temp_c": 40
+                    }],
+                    "models": models
+                }))
+            }),
+        )
+        .route(
+            "/models/{id}/endpoint",
+            get(move |Path(_id): Path<String>| {
+                let url = inference.clone();
+                async move { Json(json!({"url": url})) }
+            }),
+        )
+        .with_state(card);
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    url
+}
+
+/// A gateway with one healthy 24 GB node, its first `/health` already in.
+async fn gateway(neuron: &str, tag: &str) -> Arc<CortexState> {
+    let path = std::env::temp_dir().join(format!("cortex-cold-load-{tag}.toml"));
+    std::fs::write(&path, CATALOGUE).unwrap();
+    let config = GatewayConfig {
+        gateway: GatewaySettings {
+            listen: "127.0.0.1:0".into(),
+            metrics_listen: "127.0.0.1:0".into(),
+        },
+        eviction: EvictionSettings {
+            strategy: EvictionStrategy::Lru,
+            defrag_after_cycles: 0,
+        },
+        neurons: vec![NeuronEndpoint {
+            name: "gpu".into(),
+            endpoint: neuron.to_string(),
+        }],
+        models_config: path.to_string_lossy().into_owned(),
+        entitlements: Default::default(),
+        upstream: Default::default(),
+        audio_jobs: Default::default(),
+    };
+    let fleet = Arc::new(CortexState::from_config(&config));
+    {
+        let mut nodes = fleet.nodes.write().await;
+        let node = nodes.get_mut("gpu").unwrap();
+        node.healthy = true;
+        node.discovery = Some(DiscoveryResponse {
+            hostname: "gpu".into(),
+            os: "linux".into(),
+            kernel: "6".into(),
+            cuda_version: None,
+            driver_version: None,
+            devices: vec![DeviceInfo {
+                index: 0,
+                name: "RTX 4090".into(),
+                vram_total_mb: CARD_MB,
+                compute_capability: "8.9".into(),
+            }],
+            harnesses: vec!["candle".into()],
+            cuda_unavailable_reason: None,
+            max_prompt_tokens: 0,
+        });
+    }
+    cortex_gateway::poller::poll_health(&fleet, "gpu", neuron).await;
+    fleet
+}
+
+/// Mark `id` resident in the mock and in cortex's view, last used long
+/// ago — as the poller would have seen it, not as a fresh cold-load.
+async fn seed_resident(fleet: &CortexState, card: &Mutex<Card>, neuron: &str, id: &str) {
+    card.lock()
+        .unwrap()
+        .resident
+        .insert(id.into(), footprint(id));
+    {
+        let mut nodes = fleet.nodes.write().await;
+        nodes.get_mut("gpu").unwrap().models.insert(
+            id.into(),
+            ModelEntry {
+                id: id.into(),
+                status: ModelStatus::Loaded,
+                last_accessed: Some(chrono::Utc::now() - chrono::Duration::minutes(10)),
+                vram_estimate_mb: Some(footprint(id)),
+                capabilities: Vec::new(),
+                tool_call: false,
+                reasoning: false,
+                limit: None,
+                servable: None,
+                reasoning_budget: Vec::new(),
+            },
+        );
+    }
+    cortex_gateway::poller::poll_health(fleet, "gpu", neuron).await;
+}
+
+#[tokio::test]
+async fn concurrent_cold_loads_never_overcommit_the_card() {
+    let card = Arc::new(Mutex::new(Card::default()));
+    let neuron = spawn_neuron(Arc::clone(&card)).await;
+    let fleet = gateway(&neuron, "concurrent").await;
+
+    // Both decide while the card is empty; together they do not fit.
+    let (a, b) = tokio::join!(cortex_gateway::router::resolve(&fleet, "music-a"), async {
+        tokio::time::sleep(Duration::from_millis(50)).await;
+        cortex_gateway::router::resolve(&fleet, "music-b").await
+    });
+
+    let card = card.lock().unwrap();
+    assert!(!card.overcommitted, "both models were resident at once");
+    assert!(a.is_ok(), "{a:?}");
+    // music-a was cold-loaded for a request that has yet to reach it, so
+    // music-b may not take its place yet: retry shortly, not a failure.
+    match b {
+        Err(e @ RouteError::NodeBusy { .. }) => {
+            assert_eq!(e.http_status(), 503);
+            assert!(e.retry_after_secs().is_some());
+        }
+        other => panic!("expected NodeBusy, got {other:?}"),
+    }
+    assert_eq!(card.resident.keys().collect::<Vec<_>>(), ["music-a"]);
+    assert!(card.unloads.is_empty());
+}
+
+#[tokio::test]
+async fn a_model_with_requests_in_flight_is_not_evicted() {
+    let card = Arc::new(Mutex::new(Card::default()));
+    let neuron = spawn_neuron(Arc::clone(&card)).await;
+    let fleet = gateway(&neuron, "in-flight").await;
+    card.lock().unwrap().in_flight.insert("music-a".into(), 1);
+    seed_resident(&fleet, &card, &neuron, "music-a").await;
+
+    let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
+
+    assert!(matches!(b, Err(RouteError::NodeBusy { .. })), "{b:?}");
+    let card = card.lock().unwrap();
+    assert!(card.unloads.is_empty(), "evicted a busy model");
+    assert!(!card.overcommitted);
+}
+
+#[tokio::test]
+async fn an_idle_model_is_swapped_out() {
+    let card = Arc::new(Mutex::new(Card::default()));
+    let neuron = spawn_neuron(Arc::clone(&card)).await;
+    let fleet = gateway(&neuron, "idle").await;
+    seed_resident(&fleet, &card, &neuron, "music-a").await;
+
+    let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
+
+    assert!(b.is_ok(), "{b:?}");
+    let card = card.lock().unwrap();
+    assert_eq!(card.unloads, ["music-a"]);
+    assert_eq!(card.resident.keys().collect::<Vec<_>>(), ["music-b"]);
+    assert!(!card.overcommitted);
+}
+
+#[tokio::test]
+async fn a_cold_load_waits_for_the_neuron_s_own_startup_load() {
+    // The neuron has just restarted and is loading its default model.
+    // That load is in no model list yet and its memory has barely
+    // started to fill, so the card still looks empty.
+    let card = Arc::new(Mutex::new(Card {
+        activating: Some("music-a".into()),
+        ..Default::default()
+    }));
+    let neuron = spawn_neuron(Arc::clone(&card)).await;
+    let fleet = gateway(&neuron, "activating").await;
+
+    let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
+
+    assert!(matches!(b, Err(RouteError::NodeBusy { .. })), "{b:?}");
+    assert!(
+        card.lock().unwrap().loads.is_empty(),
+        "loaded over a startup load"
+    );
+
+    // Once the startup load finishes and its model is idle, the swap
+    // goes ahead as usual.
+    {
+        let mut c = card.lock().unwrap();
+        c.activating = None;
+    }
+    seed_resident(&fleet, &card, &neuron, "music-a").await;
+    let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
+    assert!(b.is_ok(), "{b:?}");
+    let card = card.lock().unwrap();
+    assert_eq!(card.unloads, ["music-a"]);
+    assert!(!card.overcommitted);
+}

crates/cortex-gateway/tests/common/mod.rs

@@ -498,6 +498,7 @@ pub async fn spawn_gateway_with_state(mock_url: &str) -> (Arc<CortexState>, Stri
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -559,6 +560,7 @@ pub async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexS
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     Arc::new(CortexState::from_config(&config))
 }

crates/cortex-gateway/tests/error_envelope.rs

@@ -90,6 +90,7 @@ async fn error_response_no_healthy_nodes() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(cortex_gateway::state::CortexState::from_config(&config));

crates/cortex-gateway/tests/eviction.rs

@@ -82,6 +82,7 @@ fn make_fleet_with_catalogue(
         models_config: path.to_string_lossy().into_owned(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     (Arc::new(CortexState::from_config(&config)), path)
 }
@@ -144,6 +145,7 @@ fn make_fleet(endpoint: &str, defrag_after: u32) -> Arc<CortexState> {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     Arc::new(CortexState::from_config(&config))
 }

crates/cortex-gateway/tests/feasibility_routing.rs

@@ -88,6 +88,7 @@ async fn fleet_with(big_healthy: bool, big_devices: usize) -> Arc<CortexState> {
         models_config: cat.to_string_lossy().into_owned(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = Arc::new(CortexState::from_config(&config));
     {

crates/cortex-gateway/tests/images.rs

@@ -165,6 +165,7 @@ async fn spawn_keyed_gateway(neuron_url: &str, hard_cap: u64) -> String {
             }],
         },
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = Arc::new(CortexState::from_config(&config));
     {

crates/cortex-gateway/tests/load_routing.rs

@@ -89,6 +89,7 @@ async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexState
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     Arc::new(CortexState::from_config(&config))
 }

crates/cortex-gateway/tests/metering.rs

@@ -54,6 +54,7 @@ async fn spawn_metered_gateway(neuron_url: &str) -> (Arc<CortexState>, String) {
             }],
         },
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -162,6 +163,7 @@ async fn anonymous_request_records_no_spend() {
         models_config: "/dev/null".into(),
         entitlements: EntitlementsConfig::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = Arc::new(CortexState::from_config(&config));
     {

crates/cortex-gateway/tests/metrics.rs

@@ -204,6 +204,7 @@ async fn test_capacity_gauges_exported_from_health_poll() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
     cortex_gateway::poller::poll_once(&fleet).await;

crates/cortex-gateway/tests/model_cost.rs

@@ -67,6 +67,7 @@ harness = "candle"
         models_config: cat_path.to_string_lossy().into_owned(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));

crates/cortex-gateway/tests/model_limits.rs

@@ -59,6 +59,7 @@ capabilities = ["text"]
         models_config: cat_path.to_string_lossy().into_owned(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));

crates/cortex-gateway/tests/poller.rs

@@ -33,6 +33,7 @@ async fn test_poller_discovers_models() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -86,6 +87,7 @@ async fn test_poller_updates_gateway_models_endpoint() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -159,6 +161,7 @@ async fn test_models_endpoint_unions_capabilities_across_nodes() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -223,6 +226,7 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -278,6 +282,7 @@ async fn test_poller_removes_stale_models() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -310,6 +315,7 @@ async fn test_poller_removes_stale_models() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet2 = Arc::new(CortexState::from_config(&config2));
@@ -397,6 +403,7 @@ async fn test_poller_captures_activation_from_health() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));
@@ -443,6 +450,7 @@ async fn test_poller_parses_recovering_status() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
 
     let fleet = Arc::new(CortexState::from_config(&config));

crates/cortex-gateway/tests/prompt_prevalidation.rs

@@ -76,6 +76,7 @@ async fn spawn_gateway(neuron: &str, context: usize) -> String {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = Arc::new(CortexState::from_config(&config));
     {

crates/cortex-gateway/tests/proxy_basic.rs

@@ -119,6 +119,7 @@ async fn test_no_healthy_nodes() {
         models_config: "/dev/null".into(),
         entitlements: Default::default(),
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
 

crates/cortex-gateway/tests/systemone.rs

@@ -119,6 +119,7 @@ async fn spawn_gateway(neuron_url: &str, default_alias: bool, hard_cap: Option<u
         models_config: "/dev/null".into(),
         entitlements,
         upstream: Default::default(),
+        audio_jobs: Default::default(),
     };
     let mut state = CortexState::from_config(&config);
     if default_alias {

crates/helexa-bench/src/config.rs

@@ -166,6 +166,11 @@ pub struct ScenarioConfig {
     /// capability. Defaults to `[1024]`.
     #[serde(default = "default_image_sizes")]
     pub image_sizes: Vec<u32>,
+    /// Song lengths (seconds) — one `audio:<secs>` scenario per entry
+    /// (#387), run only against models advertising the `audio`
+    /// capability. Defaults to `[30]`.
+    #[serde(default = "default_audio_durations")]
+    pub audio_durations: Vec<u32>,
     /// Approximate prompt size (tokens) used by the concurrency scenarios.
     #[serde(default = "default_concurrency_prompt_tokens")]
     pub concurrency_prompt_tokens: u32,
@@ -202,6 +207,7 @@ impl Default for ScenarioConfig {
             concurrency_nonstreaming_levels: Vec::new(),
             concurrency_prompt_tokens: default_concurrency_prompt_tokens(),
             image_sizes: default_image_sizes(),
+            audio_durations: default_audio_durations(),
             capability_probes: Vec::new(),
         }
     }
@@ -280,6 +286,9 @@ fn default_api_listen() -> String {
 fn default_image_sizes() -> Vec<u32> {
     vec![1024]
 }
+fn default_audio_durations() -> Vec<u32> {
+    vec![30]
+}
 
 fn default_prompt_sizes() -> Vec<u32> {
     vec![128, 4096]

crates/helexa-bench/src/scenario.rs

@@ -166,7 +166,7 @@ pub trait Scenario: Send + Sync {
 
 /// Capabilities that mark a model as serving an endpoint other than
 /// chat completions.
-const NON_CHAT_MODALITIES: [&str; 2] = ["image", "decision"];
+const NON_CHAT_MODALITIES: [&str; 3] = ["image", "decision", "audio"];
 
 /// Whether a model answers `/v1/chat/completions`.
 pub fn serves_chat(model: &ModelInfo) -> bool {
@@ -268,6 +268,93 @@ impl Scenario for ImageLatencyScenario {
     }
 }
 
+/// Fixed-seed music render (#387): one song of a configured length against
+/// the neuron's synchronous `/v1/audio/generations`, recording the
+/// server-side phase timing. The cell dimension is the song length, so the
+/// real-time factor is `seconds / total_s`. Applies only to models
+/// advertising the `audio` capability.
+pub struct AudioLatencyScenario {
+    pub id: String,
+    pub seconds: u32,
+}
+
+#[async_trait]
+impl Scenario for AudioLatencyScenario {
+    fn id(&self) -> &str {
+        &self.id
+    }
+
+    fn prompt_size(&self) -> u32 {
+        self.seconds
+    }
+
+    fn applies_to(&self, model: &ModelInfo) -> bool {
+        model.capabilities.iter().any(|c| c == "audio")
+    }
+
+    async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
+        let url = ctx
+            .chat_url
+            .replace("/v1/chat/completions", "/v1/audio/generations");
+        let body = serde_json::json!({
+            "model": ctx.model_id,
+            "prompt": "Warm lo-fi hip hop beat with mellow Rhodes chords and dusty vinyl drums, 82 BPM",
+            "duration": self.seconds,
+            "seed": 42,
+        });
+        let start = std::time::Instant::now();
+        let resp = ctx
+            .client
+            .post(&url)
+            .timeout(ctx.timeout)
+            .json(&body)
+            .send()
+            .await
+            .context("audio request failed")?;
+        let status = resp.status();
+        let result = resp
+            .headers()
+            .get("x-helexa-audio-result")
+            .and_then(|v| v.to_str().ok())
+            .map(str::to_string);
+        let bytes = resp.bytes().await.context("audio response body")?;
+        anyhow::ensure!(
+            status.is_success(),
+            "audio request {status}: {}",
+            String::from_utf8_lossy(&bytes[..bytes.len().min(512)])
+        );
+        let total_s = start.elapsed().as_secs_f64();
+        // A bench that records latency for broken output is worse than none.
+        // Neurons render FLAC; builds before that rendered WAV.
+        let flac = bytes.len() > 42 && &bytes[..4] == b"fLaC";
+        let wav = bytes.len() > 44 && &bytes[..4] == b"RIFF" && &bytes[8..12] == b"WAVE";
+        anyhow::ensure!(flac || wav, "payload is not a FLAC (or WAV) file");
+        let result: serde_json::Value =
+            serde_json::from_str(&result.context("missing x-helexa-audio-result")?)
+                .context("x-helexa-audio-result is not JSON")?;
+        let timing = &result["usage"]["helexa_timing"];
+        Ok(ScenarioMetrics {
+            ttft_s: timing["plan_ms"].as_u64().unwrap_or(0) as f64 / 1000.0,
+            decode_tps: None,
+            total_s,
+            prompt_tokens: None,
+            completion_tokens: 0,
+            prefill_ms: timing["plan_ms"].as_u64(),
+            decode_ms: timing["generate_ms"].as_u64(),
+            prefill_tokens: None,
+            reasoning_tokens: None,
+            cached_tokens: None,
+            tpot_p95_ms: None,
+            concurrency: None,
+            ttft_p95_s: None,
+            queue_wait_ms_median: None,
+            rejected: None,
+            artifact: None,
+            image_units: None,
+        })
+    }
+}
+
 /// Build the active scenario set from config: one chat-latency scenario per
 /// prompt size, plus one concurrency scenario per configured level (#89).
 /// Concurrency levels default to empty (opt-in), since a burst puts real
@@ -308,6 +395,12 @@ pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
             side_px: side,
         }) as Box<dyn Scenario>);
     }
+    for &secs in &cfg.audio_durations {
+        scenarios.push(Box::new(AudioLatencyScenario {
+            id: format!("audio:{secs}"),
+            seconds: secs,
+        }) as Box<dyn Scenario>);
+    }
     for probe in &cfg.capability_probes {
         scenarios.push(Box::new(CapabilityScenario {
             id: format!("capability:{}", probe.name),
@@ -915,6 +1008,14 @@ mod tests {
         assert!(serves_chat(&model(&["text", "vision"])));
         assert!(!serves_chat(&model(&["image"])));
         assert!(!serves_chat(&model(&["decision"])));
+        assert!(!serves_chat(&model(&["audio"])));
+        let audio = AudioLatencyScenario {
+            id: "audio:30".into(),
+            seconds: 30,
+        };
+        assert!(audio.applies_to(&model(&["audio"])));
+        assert!(!audio.applies_to(&model(&["text"])));
+        assert!(!audio.applies_to(&model(&["image"])));
         let image = ImageLatencyScenario {
             id: "image:512".into(),
             side_px: 512,
@@ -1069,6 +1170,7 @@ mod tests {
         use crate::config::{CapabilityProbe, ScenarioConfig};
         let cfg = ScenarioConfig {
             image_sizes: vec![],
+            audio_durations: vec![],
             prompt_sizes: vec![128],
             max_tokens: 64,
             concurrency_levels: vec![2, 8],

crates/helexa-bench/tests/sweep_integration.rs

@@ -89,6 +89,7 @@ fn config_for(endpoint: String, db_path: String) -> BenchConfig {
         },
         scenarios: ScenarioConfig {
             image_sizes: vec![],
+            audio_durations: vec![],
             prompt_sizes: vec![128], // single scenario keeps assertions simple
             max_tokens: 16,
             concurrency_levels: Vec::new(),

crates/helexa-router/src/audio.rs

@@ -0,0 +1,172 @@
+//! Async audio generation jobs through the router (#382).
+//!
+//! Submitting a job is ordinary capacity dispatch: the router picks a cortex
+//! that serves the model and forwards the body, and that cortex owns the job
+//! from then on. Every later call names only a job id, which carries no hint
+//! of its cortex, so those calls are asked of each cortex in turn and the
+//! first answer that is not `404` wins. Listing merges every cortex's list.
+//!
+//! Ownership is enforced by the cortex holding the job; the router forwards
+//! the caller's bearer unchanged, exactly as for inference.
+
+use crate::dispatch;
+use crate::error::envelope_response;
+use crate::state::RouterState;
+use axum::Router;
+use axum::body::{Body, Bytes};
+use axum::extract::{Path, RawQuery, State};
+use axum::http::{HeaderMap, Method, StatusCode};
+use axum::response::{IntoResponse, Json, Response};
+use axum::routing::get;
+use cortex_core::audio::{AudioJob, AudioJobList};
+use cortex_core::error_envelope::OpenAiError;
+use std::sync::Arc;
+
+const PATH: &str = "/v1/audio/generations";
+
+pub fn routes() -> Router<Arc<RouterState>> {
+    Router::new()
+        .route(PATH, get(list).post(submit))
+        .route("/v1/audio/generations/{id}", get(get_job).delete(cancel))
+        .route("/v1/audio/generations/{id}/content", get(content))
+}
+
+async fn submit(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    dispatch::dispatch(&state, PATH, headers, body).await
+}
+
+async fn get_job(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    find(&state, Method::GET, &format!("{PATH}/{id}"), &headers, &id).await
+}
+
+async fn cancel(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    find(
+        &state,
+        Method::DELETE,
+        &format!("{PATH}/{id}"),
+        &headers,
+        &id,
+    )
+    .await
+}
+
+async fn content(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+    RawQuery(query): RawQuery,
+) -> Response {
+    // The query carries the download format (`?format=mp3`).
+    let path = match query {
+        Some(q) => format!("{PATH}/{id}/content?{q}"),
+        None => format!("{PATH}/{id}/content"),
+    };
+    find(&state, Method::GET, &path, &headers, &id).await
+}
+
+/// Ask each cortex for the job until one answers with something other than
+/// `404`, and stream that answer back verbatim.
+async fn find(
+    state: &RouterState,
+    method: Method,
+    path: &str,
+    headers: &HeaderMap,
+    id: &str,
+) -> Response {
+    for ep in &state.cortexes {
+        let Some(client) = state.client_for(&ep.name) else {
+            continue;
+        };
+        let url = format!("{}{path}", ep.endpoint);
+        let resp = match client
+            .request(method.clone(), &url)
+            .headers(forward_headers(headers))
+            .send()
+            .await
+        {
+            Ok(r) => r,
+            Err(e) => {
+                tracing::warn!(cortex = %ep.name, url = %url, error = %e, "cortex unreachable");
+                continue;
+            }
+        };
+        if resp.status() == reqwest::StatusCode::NOT_FOUND {
+            continue;
+        }
+        let status =
+            StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
+        let mut out = Response::builder().status(status);
+        for (k, v) in resp.headers() {
+            if k == "transfer-encoding" || k == "connection" {
+                continue;
+            }
+            out = out.header(k, v);
+        }
+        return out
+            .body(Body::from_stream(resp.bytes_stream()))
+            .unwrap_or_else(|_| StatusCode::BAD_GATEWAY.into_response());
+    }
+    envelope_response(OpenAiError::new(
+        404,
+        "invalid_request_error",
+        "job_not_found",
+        format!("no audio job '{id}'"),
+    ))
+}
+
+/// Merge every reachable cortex's job list, newest first.
+async fn list(State(state): State<Arc<RouterState>>, headers: HeaderMap) -> Response {
+    let mut data: Vec<AudioJob> = Vec::new();
+    for ep in &state.cortexes {
+        let Some(client) = state.client_for(&ep.name) else {
+            continue;
+        };
+        let url = format!("{}{PATH}", ep.endpoint);
+        let page = client
+            .get(&url)
+            .headers(forward_headers(&headers))
+            .send()
+            .await
+            .ok()
+            .filter(|r| r.status().is_success());
+        let Some(page) = page else {
+            tracing::warn!(cortex = %ep.name, "audio job list unavailable from cortex");
+            continue;
+        };
+        if let Ok(list) = page.json::<AudioJobList>().await {
+            data.extend(list.data);
+        }
+    }
+    data.sort_by(|a, b| b.created_at.cmp(&a.created_at).then(b.id.cmp(&a.id)));
+    Json(AudioJobList {
+        object: "list".into(),
+        data,
+    })
+    .into_response()
+}
+
+fn forward_headers(headers: &HeaderMap) -> HeaderMap {
+    let mut out = HeaderMap::new();
+    for (k, v) in headers {
+        if matches!(
+            k.as_str(),
+            "host" | "content-length" | "connection" | "transfer-encoding"
+        ) {
+            continue;
+        }
+        out.append(k.clone(), v.clone());
+    }
+    out
+}

crates/helexa-router/src/catalogue.rs

@@ -168,6 +168,9 @@ mod tests {
             capabilities: vec![],
             limit: None,
             cost: None,
+            audio: None,
+            audio_tokens_per_second: None,
+            notices: Vec::new(),
             tool_call: false,
             reasoning: false,
             max_model_len: None,

crates/helexa-router/src/handlers.rs

@@ -14,6 +14,7 @@ pub fn api_routes() -> Router<Arc<RouterState>> {
         .route("/v1/chat/completions", post(chat_completions))
         .route("/v1/completions", post(completions))
         .route("/v1/images/generations", post(images_generations))
+        .merge(crate::audio::routes())
         .route("/v1/systemone", post(systemone))
         .route("/v1/responses", post(responses))
         .route("/v1/messages", post(messages))

crates/helexa-router/src/lib.rs

@@ -11,6 +11,7 @@
 //! on capacity (epic #69). A background [`poller`] keeps a live
 //! per-cortex topology (#72) that the dispatcher (#73) will route on.
 
+pub mod audio;
 pub mod catalogue;
 pub mod config;
 pub mod dispatch;

crates/helexa-router/tests/audio.rs

@@ -0,0 +1,205 @@
+//! Audio job lookups through the router (#382): a job id carries no hint of
+//! its cortex, so job calls are asked of each cortex until one answers with
+//! something other than 404, and listings merge every cortex's jobs.
+
+use axum::extract::{Path, RawQuery};
+use axum::http::{HeaderMap, StatusCode};
+use axum::response::{IntoResponse, Response};
+use axum::routing::get;
+use axum::{Json, Router};
+use helexa_router::config::{CortexEndpoint, RouterConfig};
+use helexa_router::state::RouterState;
+use serde_json::{Value, json};
+use std::sync::Arc;
+use tokio::net::TcpListener;
+
+/// A cortex that holds exactly one job, `job_id`, created at `created`.
+async fn spawn_cortex(job_id: &'static str, created: u64) -> String {
+    let job = move || {
+        json!({"id": job_id, "object": "audio.generation", "status": "succeeded",
+               "model": "m", "created_at": created})
+    };
+    let app = Router::new()
+        .route(
+            "/v1/audio/generations",
+            get(move |headers: HeaderMap| async move {
+                // Only the bearer's own jobs are listed.
+                if headers.get("authorization").is_none() {
+                    return Json(json!({"object": "list", "data": []})).into_response();
+                }
+                Json(json!({"object": "list", "data": [job()]})).into_response()
+            }),
+        )
+        .route(
+            "/v1/audio/generations/{id}",
+            get(move |Path(id): Path<String>| async move {
+                if id == job_id {
+                    Json(job()).into_response()
+                } else {
+                    not_found()
+                }
+            })
+            .delete(move |Path(id): Path<String>| async move {
+                if id == job_id {
+                    let mut j = job();
+                    j["status"] = json!("cancelled");
+                    Json(j).into_response()
+                } else {
+                    not_found()
+                }
+            }),
+        )
+        .route(
+            "/v1/audio/generations/{id}/content",
+            get(
+                move |Path(id): Path<String>, RawQuery(q): RawQuery| async move {
+                    if id == job_id {
+                        // Echo the query so the test can see it was forwarded.
+                        let (ct, suffix) = match q.as_deref() {
+                            Some("format=mp3") => ("audio/mpeg", ".mp3"),
+                            _ => ("audio/flac", ""),
+                        };
+                        (
+                            StatusCode::OK,
+                            [("content-type", ct)],
+                            format!("audio-of-{job_id}{suffix}"),
+                        )
+                            .into_response()
+                    } else {
+                        not_found()
+                    }
+                },
+            ),
+        );
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    format!("http://{addr}")
+}
+
+fn not_found() -> Response {
+    (
+        StatusCode::NOT_FOUND,
+        Json(json!({"error": {"type": "invalid_request_error", "code": "job_not_found", "message": "no"}})),
+    )
+        .into_response()
+}
+
+async fn router_over(cortexes: Vec<(&str, String)>) -> String {
+    let cfg = RouterConfig {
+        cortexes: cortexes
+            .into_iter()
+            .map(|(name, endpoint)| CortexEndpoint {
+                name: name.into(),
+                endpoint,
+                region: None,
+                tls_ca: None,
+            })
+            .collect(),
+        ..Default::default()
+    };
+    let state = Arc::new(RouterState::from_config(&cfg));
+    let app = helexa_router::handlers::api_routes().with_state(state);
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    format!("http://{addr}")
+}
+
+#[tokio::test]
+async fn job_calls_find_the_cortex_that_holds_the_job() {
+    let a = spawn_cortex("audiogen_a", 100).await;
+    let b = spawn_cortex("audiogen_b", 200).await;
+    // An unreachable cortex in the list must be skipped, not fatal.
+    let router = router_over(vec![
+        ("dead", "http://127.0.0.1:1".into()),
+        ("a", a),
+        ("b", b),
+    ])
+    .await;
+    let http = reqwest::Client::new();
+
+    let job: Value = http
+        .get(format!("{router}/v1/audio/generations/audiogen_b"))
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(job["id"], "audiogen_b");
+
+    let content = http
+        .get(format!("{router}/v1/audio/generations/audiogen_b/content"))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(content.status(), 200);
+    assert_eq!(content.headers()["content-type"], "audio/flac");
+    assert_eq!(content.text().await.unwrap(), "audio-of-audiogen_b");
+
+    // The download format in the query reaches the cortex.
+    let mp3 = http
+        .get(format!(
+            "{router}/v1/audio/generations/audiogen_b/content?format=mp3"
+        ))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(mp3.headers()["content-type"], "audio/mpeg");
+    assert_eq!(mp3.text().await.unwrap(), "audio-of-audiogen_b.mp3");
+
+    let cancelled: Value = http
+        .delete(format!("{router}/v1/audio/generations/audiogen_a"))
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(cancelled["status"], "cancelled");
+
+    let missing = http
+        .get(format!("{router}/v1/audio/generations/audiogen_nope"))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(missing.status(), 404);
+    let body: Value = missing.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "job_not_found");
+}
+
+#[tokio::test]
+async fn listing_merges_cortexes_newest_first_and_forwards_the_bearer() {
+    let a = spawn_cortex("audiogen_a", 100).await;
+    let b = spawn_cortex("audiogen_b", 200).await;
+    let router = router_over(vec![("a", a), ("b", b)]).await;
+    let http = reqwest::Client::new();
+
+    let list: Value = http
+        .get(format!("{router}/v1/audio/generations"))
+        .bearer_auth("k")
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    let ids: Vec<&str> = list["data"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|j| j["id"].as_str().unwrap())
+        .collect();
+    assert_eq!(ids, ["audiogen_b", "audiogen_a"]);
+
+    let anon: Value = http
+        .get(format!("{router}/v1/audio/generations"))
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(anon["data"].as_array().unwrap().len(), 0);
+}

crates/helexa-router/tests/dispatch.rs

@@ -37,6 +37,9 @@ fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
         capabilities: vec![],
         limit: None,
         cost: None,
+        audio: None,
+        audio_tokens_per_second: None,
+        notices: Vec::new(),
         tool_call: false,
         reasoning: false,
         max_model_len: None,

crates/helexa-upstream/migrations/0003_model_rates.sql

@@ -0,0 +1,32 @@
+-- ── Per-model rates for non-token units (#386) ───────────────────────
+--
+-- The ledger counts tokens. Work metered in another unit — audio-seconds
+-- for music generation today — is converted into ledger tokens at a
+-- per-model rate kept here, so an operator can price a model, or change
+-- its price, without a deploy. cortex fetches the table from
+-- `/authz/v1/rates` and reserves and settles at the rate it read.
+--
+-- `model_id = '*'` is the default for a unit: it applies to every model
+-- without a row of its own. A unit with no row at all, for the model or
+-- for '*', is not priced, and cortex refuses that work rather than serving
+-- it free.
+CREATE TABLE model_rates (
+    model_id         TEXT NOT NULL,
+    -- 'audio_second' — one second of generated audio.
+    unit             TEXT NOT NULL,
+    -- Ledger tokens charged per unit.
+    tokens_per_unit  BIGINT NOT NULL,
+    description      TEXT,
+    -- Audit: who changed it and when (NULL updated_by = shipped default).
+    updated_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
+    updated_by       UUID REFERENCES users(id) ON DELETE SET NULL,
+
+    PRIMARY KEY (model_id, unit),
+    CONSTRAINT model_rates_unit_known CHECK (unit IN ('audio_second')),
+    CONSTRAINT model_rates_positive CHECK (tokens_per_unit > 0)
+);
+
+-- Placeholder default while music generation is being built; set real
+-- per-model prices with an UPDATE or per-model INSERT.
+INSERT INTO model_rates (model_id, unit, tokens_per_unit, description)
+VALUES ('*', 'audio_second', 1000, 'default audio rate (placeholder)');

crates/helexa-upstream/src/authz.rs

@@ -21,7 +21,7 @@ use axum::extract::{Request, State};
 use axum::http::{StatusCode, header};
 use axum::middleware::Next;
 use axum::response::{IntoResponse, Response};
-use axum::routing::post;
+use axum::routing::{get, post};
 use axum::{Extension, Json, Router};
 use cortex_core::error_envelope::OpenAiError;
 use serde::{Deserialize, Serialize};
@@ -42,6 +42,7 @@ pub fn router(state: &AppState) -> Router<AppState> {
         .route("/authz/v1/release", post(release))
         .route("/authz/v1/snapshot", post(snapshot))
         .route("/authz/v1/served-usage", post(served_usage))
+        .route("/authz/v1/rates", get(rates))
         .layer(axum::middleware::from_fn_with_state(
             state.clone(),
             client_auth,
@@ -338,3 +339,23 @@ fn bad_request(msg: &str) -> Response {
         msg,
     ))
 }
+
+// ── rates (#386) ────────────────────────────────────────────────────
+
+#[derive(Serialize)]
+struct RatesResp {
+    data: Vec<crate::rates::ModelRate>,
+}
+
+/// `GET /authz/v1/rates` — every per-model rate, for cortex to cache and
+/// resolve locally. A database error is a 500: cortex keeps its last good
+/// table and refuses unpriced work rather than guessing.
+async fn rates(State(state): State<AppState>) -> Response {
+    match crate::rates::list(&state.pool).await {
+        Ok(data) => Json(RatesResp { data }).into_response(),
+        Err(e) => {
+            tracing::error!(error = %e, "rates query failed");
+            (StatusCode::INTERNAL_SERVER_ERROR, "rates unavailable").into_response()
+        }
+    }
+}

crates/helexa-upstream/src/lib.rs

@@ -21,6 +21,7 @@ pub mod email;
 pub mod error;
 pub mod handlers;
 pub mod ledger;
+pub mod rates;
 pub mod reconcile;
 pub mod state;
 pub mod topup;

crates/helexa-upstream/src/rates.rs

@@ -0,0 +1,35 @@
+//! Per-model rates for work metered in units other than tokens (#386).
+//!
+//! The ledger only counts tokens; a rate converts another unit (one second
+//! of generated audio) into ledger tokens. Rates live in `model_rates` so
+//! they can change without a deploy. cortex reads the whole table from
+//! `/authz/v1/rates` and resolves a model's rate itself: the model's own
+//! row, else the unit's `'*'` default, else the work is not priced.
+
+use serde::Serialize;
+use sqlx::postgres::PgPool;
+
+/// One row of `model_rates`, as served to cortex.
+#[derive(Debug, Clone, Serialize)]
+pub struct ModelRate {
+    pub model_id: String,
+    pub unit: String,
+    pub tokens_per_unit: i64,
+}
+
+/// Every configured rate.
+pub async fn list(pool: &PgPool) -> Result<Vec<ModelRate>, sqlx::Error> {
+    let rows = sqlx::query(
+        "SELECT model_id, unit, tokens_per_unit FROM model_rates ORDER BY unit, model_id",
+    )
+    .fetch_all(pool)
+    .await?;
+    Ok(rows
+        .into_iter()
+        .map(|r| ModelRate {
+            model_id: sqlx::Row::get(&r, "model_id"),
+            unit: sqlx::Row::get(&r, "unit"),
+            tokens_per_unit: sqlx::Row::get(&r, "tokens_per_unit"),
+        })
+        .collect())
+}

crates/helexa-upstream/tests/rates_pg.rs

@@ -0,0 +1,73 @@
+//! `model_rates` (#386) against a real PostgreSQL. Gated on
+//! `UPSTREAM_TEST_DATABASE_URL`, like the ledger tests: unset, every test
+//! logs a skip and passes.
+
+use helexa_upstream::db::connect_and_migrate;
+use helexa_upstream::rates;
+use sqlx::postgres::PgPool;
+
+async fn pool_or_skip(test: &str) -> Option<PgPool> {
+    let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
+        eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
+        return None;
+    };
+    Some(
+        connect_and_migrate(&url, 4)
+            .await
+            .expect("connect + migrate"),
+    )
+}
+
+#[tokio::test]
+async fn default_audio_rate_is_seeded_and_model_rows_are_listed() {
+    let Some(pool) = pool_or_skip("default_audio_rate_is_seeded").await else {
+        return;
+    };
+    let model = format!("test/model-{}", uuid::Uuid::new_v4());
+    sqlx::query(
+        "INSERT INTO model_rates (model_id, unit, tokens_per_unit) VALUES ($1, 'audio_second', 2500)",
+    )
+    .bind(&model)
+    .execute(&pool)
+    .await
+    .unwrap();
+
+    let all = rates::list(&pool).await.unwrap();
+    assert!(
+        all.iter()
+            .any(|r| r.model_id == "*" && r.unit == "audio_second" && r.tokens_per_unit > 0),
+        "seeded default missing: {all:?}"
+    );
+    assert!(
+        all.iter()
+            .any(|r| r.model_id == model && r.tokens_per_unit == 2500)
+    );
+
+    sqlx::query("DELETE FROM model_rates WHERE model_id = $1")
+        .bind(&model)
+        .execute(&pool)
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+async fn rates_must_be_positive_and_of_a_known_unit() {
+    let Some(pool) = pool_or_skip("rates_must_be_positive").await else {
+        return;
+    };
+    let insert = |unit: &'static str, tokens: i64| {
+        let pool = pool.clone();
+        async move {
+            sqlx::query(
+                "INSERT INTO model_rates (model_id, unit, tokens_per_unit) VALUES ('test/bad', $1, $2)",
+            )
+            .bind(unit)
+            .bind(tokens)
+            .execute(&pool)
+            .await
+        }
+    };
+    assert!(insert("audio_second", 0).await.is_err());
+    assert!(insert("audio_second", -5).await.is_err());
+    assert!(insert("furlongs", 10).await.is_err());
+}

crates/neuron/Cargo.toml

@@ -53,6 +53,7 @@ cuda-integration = ["cuda"]
 
 [dependencies]
 cortex-core.workspace = true
+flacenc = { version = "0.5.1", default-features = false, features = ["par"] }
 tokio.workspace = true
 axum.workspace = true
 metrics.workspace = true
@@ -101,6 +102,9 @@ cudarc = { version = "0.19", optional = true, default-features = false, features
 # compiling the `half` crate at conflicting versions.
 half = { version = "2.5", optional = true }
 tokenizers = { version = "0.22", default-features = false, features = ["onig"] }
+# Reads ACE-Step's `silence_latent.pt` (a torch zip archive of one raw
+# tensor). Already in the graph via candle; stored entries need no codecs.
+zip = { version = "8", default-features = false }
 hf-hub = { version = "0.4", features = ["tokio"] }
 # Jinja-compatible template renderer for the model's chat template
 # (standalone `chat_template.jinja` or `tokenizer_config.json::chat_template`).
@@ -144,6 +148,7 @@ base64 = "0.22"
 tokio = { workspace = true, features = ["test-util"] }
 reqwest.workspace = true
 tempfile = "3"
+claxon = "0.4.3"
 
 [build-dependencies]
 # Used by `build.rs` to compile `src/cuda/*.cu` into `libneuroncuda.a`

crates/neuron/src/api.rs

@@ -55,6 +55,7 @@ pub fn neuron_routes() -> Router<Arc<NeuronState>> {
         .route("/models/{model_id}/capabilities", get(model_capabilities))
         .route("/v1/chat/completions", post(chat_completions))
         .route("/v1/images/generations", post(images_generations))
+        .route("/v1/audio/generations", post(audio_generations))
         .route("/v1/systemone", post(systemone))
         .route("/v1/responses", post(responses))
 }
@@ -141,8 +142,22 @@ async fn load_model(
         )
             .into_response();
     }
-    let registry = state.registry.read().await;
-    match registry.load_model(&spec).await {
+    // Run the load on its own task. A client that gives up (a curl or
+    // cortex timeout during a first fetch of tens of GB) drops this
+    // handler's future; the load must not be dropped with it, or the
+    // download restarts from the last finished file and a device-worker
+    // reply arrives with nobody to receive it.
+    let load_spec = spec.clone();
+    let load_state = Arc::clone(&state);
+    let load = tokio::spawn(async move {
+        let registry = load_state.registry.read().await;
+        registry.load_model(&load_spec).await
+    });
+    let result = match load.await {
+        Ok(r) => r,
+        Err(e) => Err(anyhow::anyhow!("load task failed: {e}")),
+    };
+    match result {
         Ok(()) => Json(json!({"status": "loaded"})).into_response(),
         Err(e) => {
             // If the underlying failure is a structured preflight
@@ -670,6 +685,63 @@ async fn images_generations(
     }
 }
 
+/// `POST /v1/audio/generations` (#382): render one audio generation and
+/// return it. This is the neuron half of the async job API — cortex owns
+/// the job and calls this once per job, waiting for the whole render. The
+/// body is the audio (`audio/flac`); the result metadata travels in the
+/// [`cortex_core::audio::RESULT_HEADER`] header so the bytes need no
+/// base64 envelope.
+async fn audio_generations(
+    State(state): State<Arc<NeuronState>>,
+    headers: axum::http::HeaderMap,
+    ApiJson(req): ApiJson<cortex_core::audio::AudioGenerationRequest>,
+) -> axum::response::Response {
+    use cortex_core::audio::RESULT_HEADER;
+    use cortex_core::error_envelope::OpenAiError;
+
+    let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
+        return (
+            StatusCode::SERVICE_UNAVAILABLE,
+            Json(json!({"error": "candle harness not enabled on this neuron"})),
+        )
+            .into_response();
+    };
+    if let Err(msg) = req.validate_shape() {
+        return envelope_response(OpenAiError::new(
+            400,
+            "invalid_request_error",
+            "invalid_audio_params",
+            format!("invalid audio generation parameters: {msg}"),
+        ));
+    }
+    let principal = principal_key(&headers);
+    match candle.audio_generation(&req, principal.as_deref()).await {
+        Ok((flac, result)) => {
+            let meta = match serde_json::to_string(&result) {
+                Ok(m) => m,
+                Err(e) => {
+                    return inference_error_response(InferenceError::Other(anyhow::anyhow!(
+                        "serialise audio result: {e}"
+                    )));
+                }
+            };
+            let mut resp = axum::response::Response::new(axum::body::Body::from(flac));
+            let h = resp.headers_mut();
+            h.insert(
+                axum::http::header::CONTENT_TYPE,
+                axum::http::HeaderValue::from_static(cortex_core::audio::content_type(
+                    cortex_core::audio::STORED_FORMAT,
+                )),
+            );
+            if let Ok(v) = axum::http::HeaderValue::from_str(&meta) {
+                h.insert(RESULT_HEADER, v);
+            }
+            resp
+        }
+        Err(e) => inference_error_response(e),
+    }
+}
+
 /// `POST /v1/systemone` (#338): typed decisions over the TypeSafe Jev
 /// wire protocol, from a loaded decision model (Laya). Parsing and the
 /// transport-level limits happen here; routing, validation, admission,
@@ -847,6 +919,12 @@ fn inference_error_response(err: InferenceError) -> axum::response::Response {
         // model. Not retryable against this model — the client should
         // pick a model whose `capabilities` include the endpoint's
         // modality.
+        InferenceError::InvalidAudioParams { detail } => OpenAiError::new(
+            400,
+            "invalid_request_error",
+            "invalid_audio_params",
+            format!("invalid audio generation parameters: {detail}"),
+        ),
         InferenceError::InvalidImageParams { detail } => OpenAiError::new(
             400,
             "invalid_request_error",

crates/neuron/src/harness/arch/acestep/condition.rs

@@ -0,0 +1,831 @@
+//! Condition encoders for ACE-Step 1.5.
+//!
+//! This module implements the conditioning pipeline that processes text descriptions,
+//! lyrics, and timbre (reference audio) into condition tensors consumed by the DiT
+//! diffusion model.
+//!
+//! The main entry point is [`ConditionEncoder`], which orchestrates:
+//! - A simple linear projection for text encoder hidden states.
+//! - A [`LyricEncoder`] (8-layer bidirectional transformer) for lyric embeddings.
+//! - A [`TimbreEncoder`] (4-layer bidirectional transformer with CLS pooling) for
+//!   reference audio embeddings.
+//!
+//! All encoder layers share the same architecture: pre-norm (RmsNorm) bidirectional
+//! GQA attention followed by a gated SiLU MLP, matching the Qwen3 block design.
+
+use candle_core::{DType, Device, Module, Result, Tensor};
+use candle_nn::{Activation, VarBuilder};
+use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias};
+use candle_transformers::utils::repeat_kv;
+
+use super::AceStepConfig;
+
+// ---------------------------------------------------------------------------
+// RotaryEmbedding (standalone for encoder, same algorithm as the DiT variant)
+// ---------------------------------------------------------------------------
+
+/// Precomputed sin/cos tables for rotary position embeddings.
+///
+/// The encoder uses its own instance because it may have a different head dimension
+/// or maximum sequence length than the DiT backbone.
+#[derive(Debug, Clone)]
+pub struct RotaryEmbedding {
+    sin: Tensor,
+    cos: Tensor,
+}
+
+impl RotaryEmbedding {
+    /// Build the sin/cos lookup tables up to `max_seq_len` positions.
+    pub fn new(
+        dtype: DType,
+        head_dim: usize,
+        max_seq_len: usize,
+        rope_theta: f64,
+        dev: &Device,
+    ) -> Result<Self> {
+        let inv_freq: Vec<_> = (0..head_dim)
+            .step_by(2)
+            .map(|i| 1f32 / rope_theta.powf(i as f64 / head_dim as f64) as f32)
+            .collect();
+        let inv_freq_len = inv_freq.len();
+        let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?;
+        let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
+            .to_dtype(DType::F32)?
+            .reshape((max_seq_len, 1))?;
+        let freqs = t.matmul(&inv_freq)?;
+        Ok(Self {
+            sin: freqs.sin()?.to_dtype(dtype)?,
+            cos: freqs.cos()?.to_dtype(dtype)?,
+        })
+    }
+
+    /// Apply RoPE to Q and K tensors of shape `(B, H, L, D)`.
+    pub fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
+        let (_, _, seq_len, _) = q.dims4()?;
+        let cos = self.cos.narrow(0, offset, seq_len)?;
+        let sin = self.sin.narrow(0, offset, seq_len)?;
+        let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?;
+        let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?;
+        Ok((q_embed, k_embed))
+    }
+}
+
+// ---------------------------------------------------------------------------
+// EncoderMLP
+// ---------------------------------------------------------------------------
+
+/// Gated SiLU MLP identical to the Qwen3 feed-forward block.
+///
+/// ```text
+/// output = down_proj(silu(gate_proj(x)) * up_proj(x))
+/// ```
+#[derive(Debug, Clone)]
+pub struct EncoderMLP {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl EncoderMLP {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let hidden = cfg.encoder_hidden_size();
+        let intermediate = cfg.encoder_intermediate_size();
+        Ok(Self {
+            gate_proj: linear_no_bias(hidden, intermediate, vb.pp("gate_proj"))?,
+            up_proj: linear_no_bias(hidden, intermediate, vb.pp("up_proj"))?,
+            down_proj: linear_no_bias(intermediate, hidden, vb.pp("down_proj"))?,
+            act_fn: cfg.hidden_act,
+        })
+    }
+}
+
+impl Module for EncoderMLP {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let gate = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let up = x.apply(&self.up_proj)?;
+        (gate * up)?.apply(&self.down_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// EncoderAttention
+// ---------------------------------------------------------------------------
+
+/// Bidirectional grouped-query attention for encoder layers.
+///
+/// Uses per-head RmsNorm on Q and K (Qwen3 style) and rotary position
+/// embeddings. No KV cache is needed because the encoder runs a single
+/// forward pass over the full sequence.
+#[derive(Debug, Clone)]
+pub struct EncoderAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+}
+
+impl EncoderAttention {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let head_dim = cfg.head_dim;
+        let num_heads = cfg.encoder_num_attention_heads();
+        let num_kv_heads = cfg.encoder_num_key_value_heads();
+        let num_kv_groups = num_heads / num_kv_heads;
+        let hidden = cfg.encoder_hidden_size();
+
+        let q_proj = linear_b(
+            hidden,
+            num_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("q_proj"),
+        )?;
+        let k_proj = linear_b(
+            hidden,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("k_proj"),
+        )?;
+        let v_proj = linear_b(
+            hidden,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("v_proj"),
+        )?;
+        let o_proj = linear_b(
+            num_heads * head_dim,
+            hidden,
+            cfg.attention_bias,
+            vb.pp("o_proj"),
+        )?;
+
+        let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?;
+        let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?;
+
+        let hidden_size = head_dim * num_heads;
+
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            q_norm,
+            k_norm,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+        })
+    }
+
+    /// Run bidirectional self-attention.
+    ///
+    /// * `x` - input tensor of shape `(B, L, D)`.
+    /// * `mask` - optional additive attention mask `(B, 1, L, L)` where `-inf`
+    ///   blocks attention. Pass `None` for fully bidirectional attention.
+    /// * `rotary` - precomputed rotary embedding tables.
+    pub fn forward(
+        &self,
+        x: &Tensor,
+        mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        // Project
+        let q = self.q_proj.forward(x)?;
+        let k = self.k_proj.forward(x)?;
+        let v = self.v_proj.forward(x)?;
+
+        // Reshape to (B, H, L, D)
+        let q = q
+            .reshape((b, l, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // Per-head RmsNorm
+        let q = self.q_norm.forward(&q.flatten(0, 2)?)?.reshape((
+            b,
+            self.num_heads,
+            l,
+            self.head_dim,
+        ))?;
+        let k = self.k_norm.forward(&k.flatten(0, 2)?)?.reshape((
+            b,
+            self.num_kv_heads,
+            l,
+            self.head_dim,
+        ))?;
+
+        // Rotary position embeddings
+        let (q, k) = rotary.apply(&q, &k, 0)?;
+
+        // GQA key/value expansion
+        let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
+
+        // Scaled dot-product attention
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
+        if let Some(m) = mask {
+            scores = scores.broadcast_add(m)?;
+        }
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let ctx = probs.matmul(&v)?;
+
+        // Reshape back and output projection
+        ctx.transpose(1, 2)?
+            .reshape((b, l, self.hidden_size))?
+            .apply(&self.o_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// EncoderLayer
+// ---------------------------------------------------------------------------
+
+/// A single bidirectional pre-norm transformer layer.
+///
+/// ```text
+/// x = x + self_attn(input_layernorm(x))
+/// x = x + mlp(post_attention_layernorm(x))
+/// ```
+///
+/// Layers alternate between sliding-window and full attention based on
+/// `attention_type`, which controls the mask passed to the attention block.
+#[derive(Debug, Clone)]
+pub struct EncoderLayer {
+    self_attn: EncoderAttention,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+    mlp: EncoderMLP,
+    attention_type: String,
+}
+
+impl EncoderLayer {
+    pub fn new(cfg: &AceStepConfig, attention_type: &str, vb: VarBuilder) -> Result<Self> {
+        let hidden = cfg.encoder_hidden_size();
+        Ok(Self {
+            self_attn: EncoderAttention::new(cfg, vb.pp("self_attn"))?,
+            input_layernorm: RmsNorm::new(hidden, cfg.rms_norm_eps, vb.pp("input_layernorm"))?,
+            post_attention_layernorm: RmsNorm::new(
+                hidden,
+                cfg.rms_norm_eps,
+                vb.pp("post_attention_layernorm"),
+            )?,
+            mlp: EncoderMLP::new(cfg, vb.pp("mlp"))?,
+            attention_type: attention_type.to_string(),
+        })
+    }
+
+    /// Returns the attention type for this layer (`"sliding_attention"` or
+    /// `"full_attention"`), used to select the appropriate mask.
+    pub fn attention_type(&self) -> &str {
+        &self.attention_type
+    }
+
+    /// Forward pass with pre-norm residual connections.
+    pub fn forward(
+        &self,
+        x: &Tensor,
+        mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        // Self-attention with residual
+        let residual = x.clone();
+        let h = self.input_layernorm.forward(x)?;
+        let h = self.self_attn.forward(&h, mask, rotary)?;
+        let x = (residual + h)?;
+
+        // MLP with residual
+        let residual = x.clone();
+        let h = self.post_attention_layernorm.forward(&x)?;
+        let h = h.apply(&self.mlp)?;
+        residual + h
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Bidirectional attention mask helpers
+// ---------------------------------------------------------------------------
+
+/// Build a bidirectional sliding-window mask.
+///
+/// Positions outside the window receive `-inf`; positions inside receive `0`.
+/// The returned tensor has shape `(1, 1, seq_len, seq_len)`.
+fn sliding_window_mask(
+    seq_len: usize,
+    window: usize,
+    dtype: DType,
+    dev: &Device,
+) -> Result<Tensor> {
+    let minf = f32::NEG_INFINITY;
+    let mask: Vec<f32> = (0..seq_len)
+        .flat_map(|i| {
+            (0..seq_len).map(move |j| {
+                if (i as i64 - j as i64).unsigned_abs() as usize <= window {
+                    0.0
+                } else {
+                    minf
+                }
+            })
+        })
+        .collect();
+    Tensor::from_slice(&mask, (1, 1, seq_len, seq_len), dev)?.to_dtype(dtype)
+}
+
+/// Build an attention mask from a padding mask, optionally combined with a
+/// sliding window constraint.
+///
+/// * `attention_mask` - `(B, L)` with `1` for valid tokens, `0` for padding.
+/// * `sliding_window` - if `Some(w)`, restricts attention to a local window.
+///
+/// Returns `(B, 1, L, L)` additive mask.
+fn build_bidirectional_mask(
+    attention_mask: &Tensor,
+    seq_len: usize,
+    sliding_window: Option<usize>,
+    dtype: DType,
+    dev: &Device,
+) -> Result<Tensor> {
+    let (b, _) = attention_mask.dims2()?;
+
+    // Padding mask: convert (B, L) 1/0 mask → (B, 1, 1, L) with 0.0 for valid, -inf for padding.
+    let pad_mask = attention_mask
+        .reshape((b, 1, 1, seq_len))?
+        .to_dtype(DType::F32)?;
+    // Use where_cond to avoid 0 * -inf = NaN
+    let zeros = Tensor::zeros_like(&pad_mask)?;
+    let neginf = Tensor::full(f32::NEG_INFINITY, pad_mask.shape(), dev)?.to_dtype(DType::F32)?;
+    let pad_mask_bool = pad_mask.gt(&zeros)?; // true for valid (mask==1)
+    let pad_mask = pad_mask_bool.where_cond(&zeros, &neginf)?; // 0.0 for valid, -inf for padding
+
+    let mask = match sliding_window {
+        Some(w) => {
+            let sw = sliding_window_mask(seq_len, w, DType::F32, dev)?;
+            // Combine: both must allow the position
+            pad_mask.broadcast_add(&sw)?
+        }
+        None => pad_mask
+            .broadcast_as((b, 1, seq_len, seq_len))?
+            .contiguous()?,
+    };
+    mask.to_dtype(dtype)
+}
+
+// ---------------------------------------------------------------------------
+// LyricEncoder
+// ---------------------------------------------------------------------------
+
+/// Bidirectional transformer encoder for lyric embeddings.
+///
+/// Takes pre-embedded lyric token features (from a text encoder) and refines
+/// them through 8 layers of bidirectional self-attention.
+///
+/// Weight prefix: `lyric_encoder.*`
+#[derive(Debug, Clone)]
+pub struct LyricEncoder {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: RotaryEmbedding,
+    sliding_window: Option<usize>,
+    dtype: DType,
+    device: Device,
+}
+
+impl LyricEncoder {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hidden = cfg.encoder_hidden_size();
+
+        // Linear projection from text embedding dim to encoder hidden size (with bias).
+        let embed_tokens = linear_b(cfg.text_hidden_dim, enc_hidden, true, vb.pp("embed_tokens"))?;
+
+        let num_layers = cfg.num_lyric_encoder_hidden_layers;
+        let mut layers = Vec::with_capacity(num_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..num_layers {
+            let attn_type = if i % 2 == 0 {
+                "sliding_attention"
+            } else {
+                "full_attention"
+            };
+            layers.push(EncoderLayer::new(cfg, attn_type, vb_l.pp(i))?);
+        }
+
+        let norm = RmsNorm::new(enc_hidden, cfg.rms_norm_eps, vb.pp("norm"))?;
+
+        let rotary_emb = RotaryEmbedding::new(
+            vb.dtype(),
+            cfg.head_dim,
+            cfg.max_position_embeddings,
+            cfg.rope_theta,
+            vb.device(),
+        )?;
+
+        let sliding_window = if cfg.use_sliding_window {
+            Some(cfg.sliding_window)
+        } else {
+            None
+        };
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            rotary_emb,
+            sliding_window,
+            dtype: vb.dtype(),
+            device: vb.device().clone(),
+        })
+    }
+
+    /// Encode lyric embeddings.
+    ///
+    /// * `inputs_embeds` - `(B, T, text_hidden_dim)` pre-embedded lyric tokens.
+    /// * `attention_mask` - `(B, T)` with `1` for valid tokens, `0` for padding.
+    ///
+    /// Returns `(B, T, encoder_hidden_size)`.
+    pub fn forward(&self, inputs_embeds: &Tensor, attention_mask: &Tensor) -> Result<Tensor> {
+        let (_, seq_len, _) = inputs_embeds.dims3()?;
+
+        let mut hidden = self.embed_tokens.forward(inputs_embeds)?;
+
+        for layer in &self.layers {
+            let sw = match layer.attention_type() {
+                "sliding_attention" => self.sliding_window,
+                _ => None,
+            };
+            let mask =
+                build_bidirectional_mask(attention_mask, seq_len, sw, self.dtype, &self.device)?;
+            hidden = layer.forward(&hidden, Some(&mask), &self.rotary_emb)?;
+        }
+
+        self.norm.forward(&hidden)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// TimbreEncoder
+// ---------------------------------------------------------------------------
+
+/// Bidirectional transformer encoder for timbre (reference audio) embeddings.
+///
+/// Uses a learned CLS token prepended to the input sequence. After processing,
+/// the CLS token representation (position 0) is extracted as the timbre summary.
+///
+/// Weight prefix: `timbre_encoder.*`
+#[derive(Debug, Clone)]
+pub struct TimbreEncoder {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: RotaryEmbedding,
+    /// CLS token prepended to timbre sequence (XL models only).
+    special_token: Tensor,
+    /// Whether to prepend the CLS special_token (true for XL, false for 2B).
+    use_cls: bool,
+    sliding_window: Option<usize>,
+    dtype: DType,
+    device: Device,
+}
+
+impl TimbreEncoder {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hidden = cfg.encoder_hidden_size();
+
+        // Linear projection from timbre embedding dim to encoder hidden size (with bias).
+        let embed_tokens = linear_b(
+            cfg.timbre_hidden_dim,
+            enc_hidden,
+            true,
+            vb.pp("embed_tokens"),
+        )?;
+
+        let num_layers = cfg.num_timbre_encoder_hidden_layers;
+        let mut layers = Vec::with_capacity(num_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..num_layers {
+            let attn_type = if i % 2 == 0 {
+                "sliding_attention"
+            } else {
+                "full_attention"
+            };
+            layers.push(EncoderLayer::new(cfg, attn_type, vb_l.pp(i))?);
+        }
+
+        let norm = RmsNorm::new(enc_hidden, cfg.rms_norm_eps, vb.pp("norm"))?;
+
+        // XL models prepend a CLS special token to the timbre sequence.
+        // Base/2B models have this commented out in Python.
+        let special_token = vb.get((1, 1, enc_hidden), "special_token")?;
+
+        let rotary_emb = RotaryEmbedding::new(
+            vb.dtype(),
+            cfg.head_dim,
+            cfg.max_position_embeddings,
+            cfg.rope_theta,
+            vb.device(),
+        )?;
+
+        let sliding_window = if cfg.use_sliding_window {
+            Some(cfg.sliding_window)
+        } else {
+            None
+        };
+
+        // XL models have encoder_hidden_size != hidden_size and use the CLS token.
+        let use_cls = cfg.encoder_hidden_size.is_some();
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            rotary_emb,
+            special_token,
+            use_cls,
+            sliding_window,
+            dtype: vb.dtype(),
+            device: vb.device().clone(),
+        })
+    }
+
+    /// Unpack packed timbre embeddings into per-batch format.
+    ///
+    /// Matches the Python `unpack_timbre_embeddings` exactly:
+    /// * `timbre_embs_packed` - `(N, D)` embeddings from each reference audio segment.
+    /// * `refer_audio_order_mask` - `(N,)` 1D integer tensor mapping each segment
+    ///   to its batch index (0..B-1).
+    ///
+    /// Returns `(timbre_embs, timbre_mask)`:
+    /// - `timbre_embs`: `(B, max_count, D)` gathered embeddings.
+    /// - `timbre_mask`: `(B, max_count)` integer mask (1=valid, 0=padding).
+    fn unpack_timbre_embeddings(
+        &self,
+        timbre_embs_packed: &Tensor,
+        refer_audio_order_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        let (n, d) = timbre_embs_packed.dims2()?;
+        let mask_1d = refer_audio_order_mask.flatten_all()?;
+
+        // B = max(order_mask) + 1
+        let b = (mask_1d.max(0)?.to_scalar::<i64>()? + 1) as usize;
+
+        // Count elements per batch and find max_count
+        let mask_vec = mask_1d.to_vec1::<i64>()?;
+        let mut counts = vec![0usize; b];
+        for &idx in &mask_vec {
+            counts[idx as usize] += 1;
+        }
+        let max_count = *counts.iter().max().unwrap_or(&1);
+
+        // Build the output by placing each embedding at the right (batch, position)
+        let mut position_in_batch = vec![0usize; b];
+        let mut result = vec![0f32; b * max_count * d];
+        let mut mask_out = vec![0f32; b * max_count];
+
+        let packed_data = timbre_embs_packed.to_dtype(DType::F32)?.to_vec2::<f32>()?;
+        for i in 0..n {
+            let batch_idx = mask_vec[i] as usize;
+            let pos = position_in_batch[batch_idx];
+            position_in_batch[batch_idx] += 1;
+            let dst_offset = (batch_idx * max_count + pos) * d;
+            result[dst_offset..dst_offset + d].copy_from_slice(&packed_data[i]);
+            mask_out[batch_idx * max_count + pos] = 1.0;
+        }
+
+        let timbre_embs =
+            Tensor::from_vec(result, (b, max_count, d), &self.device)?.to_dtype(self.dtype)?;
+        let timbre_mask =
+            Tensor::from_vec(mask_out, (b, max_count), &self.device)?.to_dtype(self.dtype)?;
+
+        Ok((timbre_embs, timbre_mask))
+    }
+
+    /// Encode timbre reference audio segments.
+    ///
+    /// * `refer_audio_packed` - `(N, T, timbre_hidden_dim)` packed reference
+    ///   audio features for all segments across the batch.
+    /// * `refer_audio_order_mask` - `(N,)` 1D integer tensor mapping each
+    ///   segment to its batch index (0..B-1).
+    ///
+    /// Returns `(timbre_embs, timbre_mask)`:
+    /// - `timbre_embs`: `(B, max_count, encoder_hidden_size)`.
+    /// - `timbre_mask`: `(B, max_count)` integer mask.
+    pub fn forward(
+        &self,
+        refer_audio_packed: &Tensor,
+        refer_audio_order_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        let (n, _t, _) = refer_audio_packed.dims3()?;
+
+        // Project input features: (N, T, timbre_dim) -> (N, T, hidden_size)
+        let mut x = self.embed_tokens.forward(refer_audio_packed)?;
+
+        // XL models prepend a CLS special_token for timbre aggregation.
+        if self.use_cls {
+            let cls = self
+                .special_token
+                .broadcast_as((n, 1, x.dim(2)?))?
+                .contiguous()?;
+            x = Tensor::cat(&[&cls, &x], 1)?;
+        }
+
+        let seq_len = x.dim(1)?;
+
+        // Build attention mask
+        let ones = Tensor::ones((n, seq_len), self.dtype, &self.device)?;
+
+        for layer in &self.layers {
+            let sw = match layer.attention_type() {
+                "sliding_attention" => self.sliding_window,
+                _ => None,
+            };
+            let mask = build_bidirectional_mask(&ones, seq_len, sw, self.dtype, &self.device)?;
+            x = layer.forward(&x, Some(&mask), &self.rotary_emb)?;
+        }
+
+        x = self.norm.forward(&x)?;
+
+        // Extract position 0 as the timbre summary: (N, T, D) -> (N, D)
+        let embeds = x.narrow(1, 0, 1)?.squeeze(1)?;
+
+        self.unpack_timbre_embeddings(&embeds, refer_audio_order_mask)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// pack_sequences
+// ---------------------------------------------------------------------------
+
+/// Pack two (hidden, mask) pairs: concatenate, then reorder so valid tokens
+/// (mask=1) come before padding tokens (mask=0) in each batch element.
+///
+/// Uses a CPU-side stable gather (valid tokens first in original order, then
+/// padding tokens) to match Python's `argsort(descending=True, stable=True)`.
+/// Length computation uses I64 to avoid BF16 precision loss for long sequences.
+///
+/// * `hidden1` - `(B, L1, D)`
+/// * `hidden2` - `(B, L2, D)`
+/// * `mask1` - `(B, L1)` float with 1.0=valid, 0.0=padding
+/// * `mask2` - `(B, L2)` float
+///
+/// Returns `(hidden, mask)` with shapes `(B, L1+L2, D)` and `(B, L1+L2)`.
+pub fn pack_sequences(
+    hidden1: &Tensor,
+    hidden2: &Tensor,
+    mask1: &Tensor,
+    mask2: &Tensor,
+) -> Result<(Tensor, Tensor)> {
+    let hidden_cat = Tensor::cat(&[hidden1, hidden2], 1)?;
+    let mask_cat = Tensor::cat(&[mask1, mask2], 1)?;
+
+    let (b, l, d) = hidden_cat.dims3()?;
+
+    // Build stable sort indices on CPU: valid positions first (in order),
+    // then padding positions (in order). This is equivalent to Python's
+    // stable descending argsort on a binary mask.
+    let mask_f32 = mask_cat
+        .to_dtype(candle_core::DType::F32)?
+        .to_device(&candle_core::Device::Cpu)?;
+    let mask_data = mask_f32.to_vec2::<f32>()?;
+
+    let mut sort_indices = vec![0u32; b * l];
+    let mut valid_counts = vec![0usize; b];
+    for bi in 0..b {
+        let mut valid = Vec::new();
+        let mut invalid = Vec::new();
+        for (j, &val) in mask_data[bi].iter().enumerate().take(l) {
+            if val > 0.5 {
+                valid.push(j as u32);
+            } else {
+                invalid.push(j as u32);
+            }
+        }
+        valid_counts[bi] = valid.len();
+        let offset = bi * l;
+        for (k, &idx) in valid.iter().chain(invalid.iter()).enumerate() {
+            sort_indices[offset + k] = idx;
+        }
+    }
+
+    let sort_idx = Tensor::from_vec(sort_indices, (b, l), hidden_cat.device())?;
+    let idx_expanded = sort_idx
+        .unsqueeze(2)?
+        .broadcast_as((b, l, d))?
+        .contiguous()?;
+    let hidden_sorted = hidden_cat.gather(&idx_expanded, 1)?;
+
+    // Build new mask using integer counts (avoids BF16 precision loss)
+    let mut new_mask_data = vec![0f32; b * l];
+    for bi in 0..b {
+        for j in 0..valid_counts[bi] {
+            new_mask_data[bi * l + j] = 1.0;
+        }
+    }
+    let new_mask =
+        Tensor::from_vec(new_mask_data, (b, l), hidden_cat.device())?.to_dtype(mask1.dtype())?;
+
+    Ok((hidden_sorted, new_mask))
+}
+
+// ---------------------------------------------------------------------------
+// ConditionEncoder
+// ---------------------------------------------------------------------------
+
+/// Top-level condition encoder that produces the cross-attention context for
+/// the DiT diffusion backbone.
+///
+/// Combines three conditioning signals:
+/// 1. **Text** - style/prompt hidden states projected to encoder dimension.
+/// 2. **Lyrics** - lyric token embeddings processed by a bidirectional
+///    transformer.
+/// 3. **Timbre** - reference audio segments encoded with CLS-token pooling.
+///
+/// The outputs are packed into a single sequence with an accompanying mask.
+///
+/// Weight prefix: `encoder.*`
+#[derive(Debug, Clone)]
+pub struct ConditionEncoder {
+    text_projector: Linear,
+    lyric_encoder: LyricEncoder,
+    timbre_encoder: TimbreEncoder,
+}
+
+impl ConditionEncoder {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hidden = cfg.encoder_hidden_size();
+
+        // Text projector: no bias
+        let text_projector =
+            linear_no_bias(cfg.text_hidden_dim, enc_hidden, vb.pp("text_projector"))?;
+
+        let lyric_encoder = LyricEncoder::new(cfg, vb.pp("lyric_encoder"))?;
+        let timbre_encoder = TimbreEncoder::new(cfg, vb.pp("timbre_encoder"))?;
+
+        Ok(Self {
+            text_projector,
+            lyric_encoder,
+            timbre_encoder,
+        })
+    }
+
+    /// Encode all conditioning inputs into a single packed sequence.
+    ///
+    /// * `text_hidden_states` - `(B, T_text, text_hidden_dim)` from the text encoder.
+    /// * `text_attention_mask` - `(B, T_text)` mask for text tokens.
+    /// * `lyric_hidden_states` - `(B, T_lyric, text_hidden_dim)` pre-embedded lyrics.
+    /// * `lyric_attention_mask` - `(B, T_lyric)` mask for lyric tokens.
+    /// * `refer_audio_packed` - `(N, T_audio, timbre_hidden_dim)` packed reference
+    ///   audio features.
+    /// * `refer_audio_order_mask` - `(N,)` 1D integer tensor mapping each
+    ///   segment to its batch index.
+    ///
+    /// Returns `(enc_hidden, enc_mask)`:
+    /// - `enc_hidden`: `(B, T_lyric + max_ref + T_text, encoder_hidden_size)`.
+    /// - `enc_mask`: `(B, T_lyric + max_ref + T_text)` float mask.
+    pub fn forward(
+        &self,
+        text_hidden_states: &Tensor,
+        text_attention_mask: &Tensor,
+        lyric_hidden_states: &Tensor,
+        lyric_attention_mask: &Tensor,
+        refer_audio_packed: &Tensor,
+        refer_audio_order_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        // Project text hidden states
+        let text = self.text_projector.forward(text_hidden_states)?;
+
+        // Encode lyrics
+        let lyric = self
+            .lyric_encoder
+            .forward(lyric_hidden_states, lyric_attention_mask)?;
+
+        // Encode timbre
+        let (timbre, timbre_mask) = self
+            .timbre_encoder
+            .forward(refer_audio_packed, refer_audio_order_mask)?;
+
+        // Pack lyric + timbre
+        let (enc_hidden, enc_mask) =
+            pack_sequences(&lyric, &timbre, lyric_attention_mask, &timbre_mask)?;
+
+        // Pack (lyric+timbre) + text
+        let (enc_hidden, enc_mask) =
+            pack_sequences(&enc_hidden, &text, &enc_mask, text_attention_mask)?;
+
+        Ok((enc_hidden, enc_mask))
+    }
+}

crates/neuron/src/harness/arch/acestep/mod.rs

@@ -0,0 +1,358 @@
+//! ACE-Step 1.5 music generation model.
+//!
+//! Vendored from the candle port in huggingface/candle#3451
+//! (n0madic/candle@294086f4, MIT OR Apache-2.0), which matches the Python
+//! reference to about 2e-5. Changes here: crate paths, a local copy of
+//! candle's `qwen3` with the two accessors the port needs (`embed_tokens`,
+//! a public `clear_kv_cache`), and lint fixes. Drop the copy once the PR
+//! lands upstream and the candle pin moves past it.
+//!
+//! ACE-Step is a music generation foundation model based on a diffusion transformer (DiT)
+//! architecture with text conditioning for lyrics and style prompts. Supports two
+//! inference modes:
+//!
+//! - **DiT-only**: text encoder → DiT denoising → VAE decode (base/sft models)
+//! - **LM+DiT**: Qwen3 LM generates metadata + audio codes → DiT conditioned on
+//!   codes → VAE decode (turbo models, higher quality)
+//!
+//! ## Modules
+//!
+//! - [`pipeline`] — `AceStepPipeline` high-level API (text2music, cover)
+//! - [`model`] — `AceStepConditionGenerationModel` (condition encoder + DiT + tokenizer)
+//! - [`lm`] — `LmPipeline` for Qwen3-based audio code generation
+//! - [`dit`] — `AceStepDiTModel` diffusion transformer
+//! - [`condition`] — `ConditionEncoder` (text/lyric/timbre encoding)
+//! - [`tokenizer`] — `AudioTokenizer` / `AudioTokenDetokenizer` / `ResidualFSQ`
+//! - [`sampling`] — timestep schedules, Euler ODE/SDE, APG guidance
+//! - [`vae`] — `AutoencoderOobleck` / `OobleckDecoder` for audio ↔ latent
+//!
+//! - [HuggingFace Models](https://huggingface.co/ACE-Step)
+//! - [GitHub Repository](https://github.com/ACE-Step/ACE-Step-1.5)
+
+pub mod condition;
+pub mod dit;
+pub mod lm;
+pub mod model;
+pub mod pipeline;
+pub mod qwen3;
+pub mod sampling;
+pub mod tokenizer;
+pub mod vae;
+
+use candle_nn::Activation;
+
+// -- AceStepConfig default value functions -----------------------------------
+
+fn default_hidden_size() -> usize {
+    2048
+}
+fn default_intermediate_size() -> usize {
+    6144
+}
+fn default_num_hidden_layers() -> usize {
+    24
+}
+fn default_num_attention_heads() -> usize {
+    16
+}
+fn default_num_key_value_heads() -> usize {
+    8
+}
+fn default_head_dim() -> usize {
+    128
+}
+fn default_hidden_act() -> Activation {
+    Activation::Silu
+}
+fn default_rms_norm_eps() -> f64 {
+    1e-6
+}
+fn default_rope_theta() -> f64 {
+    1_000_000.0
+}
+fn default_use_sliding_window() -> bool {
+    true
+}
+fn default_sliding_window() -> usize {
+    128
+}
+fn default_in_channels() -> usize {
+    192
+}
+fn default_patch_size() -> usize {
+    2
+}
+fn default_text_hidden_dim() -> usize {
+    1024
+}
+fn default_audio_acoustic_hidden_dim() -> usize {
+    64
+}
+fn default_timbre_hidden_dim() -> usize {
+    64
+}
+fn default_pool_window_size() -> usize {
+    5
+}
+fn default_vocab_size() -> usize {
+    64003
+}
+fn default_max_position_embeddings() -> usize {
+    32768
+}
+fn default_fsq_dim() -> usize {
+    2048
+}
+fn default_fsq_input_levels() -> Vec<usize> {
+    vec![8, 8, 8, 5, 5, 5]
+}
+fn default_fsq_input_num_quantizers() -> usize {
+    1
+}
+fn default_attention_bias() -> bool {
+    false
+}
+fn default_attention_dropout() -> f64 {
+    0.0
+}
+fn default_num_lyric_encoder_hidden_layers() -> usize {
+    8
+}
+fn default_num_timbre_encoder_hidden_layers() -> usize {
+    4
+}
+fn default_num_attention_pooler_hidden_layers() -> usize {
+    2
+}
+fn default_num_audio_decoder_hidden_layers() -> usize {
+    24
+}
+fn default_layer_types() -> Vec<String> {
+    (0..24)
+        .map(|i| {
+            if i % 2 == 0 {
+                "sliding_attention".to_string()
+            } else {
+                "full_attention".to_string()
+            }
+        })
+        .collect()
+}
+
+/// Main configuration for the ACE-Step 1.5 DiT-based music generation model.
+///
+/// Deserializable from HuggingFace `config.json`. Supports both the base (3.5B)
+/// and XL variants via optional encoder dimension overrides.
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct AceStepConfig {
+    // -- DiT dimensions ------------------------------------------------------
+    #[serde(default = "default_hidden_size")]
+    pub hidden_size: usize,
+    #[serde(default = "default_intermediate_size")]
+    pub intermediate_size: usize,
+    #[serde(default = "default_num_hidden_layers")]
+    pub num_hidden_layers: usize,
+    #[serde(default = "default_num_attention_heads")]
+    pub num_attention_heads: usize,
+    #[serde(default = "default_num_key_value_heads")]
+    pub num_key_value_heads: usize,
+    #[serde(default = "default_head_dim")]
+    pub head_dim: usize,
+
+    // -- Encoder dimensions (XL only, optional) ------------------------------
+    #[serde(default)]
+    pub encoder_hidden_size: Option<usize>,
+    #[serde(default)]
+    pub encoder_intermediate_size: Option<usize>,
+    #[serde(default)]
+    pub encoder_num_attention_heads: Option<usize>,
+    #[serde(default)]
+    pub encoder_num_key_value_heads: Option<usize>,
+
+    // -- Common parameters ---------------------------------------------------
+    #[serde(default = "default_hidden_act")]
+    pub hidden_act: Activation,
+    #[serde(default = "default_rms_norm_eps")]
+    pub rms_norm_eps: f64,
+    #[serde(default = "default_rope_theta")]
+    pub rope_theta: f64,
+    #[serde(default = "default_use_sliding_window")]
+    pub use_sliding_window: bool,
+    #[serde(default = "default_sliding_window")]
+    pub sliding_window: usize,
+    #[serde(default = "default_in_channels")]
+    pub in_channels: usize,
+    #[serde(default = "default_patch_size")]
+    pub patch_size: usize,
+    #[serde(default = "default_text_hidden_dim")]
+    pub text_hidden_dim: usize,
+    #[serde(default = "default_audio_acoustic_hidden_dim")]
+    pub audio_acoustic_hidden_dim: usize,
+    #[serde(default = "default_timbre_hidden_dim")]
+    pub timbre_hidden_dim: usize,
+    #[serde(default = "default_pool_window_size")]
+    pub pool_window_size: usize,
+    #[serde(default = "default_vocab_size")]
+    pub vocab_size: usize,
+    #[serde(default = "default_max_position_embeddings")]
+    pub max_position_embeddings: usize,
+    #[serde(default = "default_fsq_dim")]
+    pub fsq_dim: usize,
+    #[serde(default = "default_fsq_input_levels")]
+    pub fsq_input_levels: Vec<usize>,
+    #[serde(default = "default_fsq_input_num_quantizers")]
+    pub fsq_input_num_quantizers: usize,
+    #[serde(default = "default_attention_bias")]
+    pub attention_bias: bool,
+    #[serde(default = "default_attention_dropout")]
+    pub attention_dropout: f64,
+
+    // -- Sub-model layer counts ----------------------------------------------
+    #[serde(default = "default_num_lyric_encoder_hidden_layers")]
+    pub num_lyric_encoder_hidden_layers: usize,
+    #[serde(default = "default_num_timbre_encoder_hidden_layers")]
+    pub num_timbre_encoder_hidden_layers: usize,
+    #[serde(default = "default_num_attention_pooler_hidden_layers")]
+    pub num_attention_pooler_hidden_layers: usize,
+    #[serde(default = "default_num_audio_decoder_hidden_layers")]
+    pub num_audio_decoder_hidden_layers: usize,
+
+    // -- Layer type schedule -------------------------------------------------
+    #[serde(default = "default_layer_types")]
+    pub layer_types: Vec<String>,
+
+    // -- Turbo flag ----------------------------------------------------------
+    #[serde(default)]
+    pub is_turbo: bool,
+}
+
+impl AceStepConfig {
+    /// Returns the encoder hidden size, falling back to the main `hidden_size`.
+    pub fn encoder_hidden_size(&self) -> usize {
+        self.encoder_hidden_size.unwrap_or(self.hidden_size)
+    }
+
+    /// Returns the encoder intermediate size, falling back to the main `intermediate_size`.
+    pub fn encoder_intermediate_size(&self) -> usize {
+        self.encoder_intermediate_size
+            .unwrap_or(self.intermediate_size)
+    }
+
+    /// Returns the encoder attention head count, falling back to `num_attention_heads`.
+    pub fn encoder_num_attention_heads(&self) -> usize {
+        self.encoder_num_attention_heads
+            .unwrap_or(self.num_attention_heads)
+    }
+
+    /// Returns the encoder key-value head count, falling back to `num_key_value_heads`.
+    pub fn encoder_num_key_value_heads(&self) -> usize {
+        self.encoder_num_key_value_heads
+            .unwrap_or(self.num_key_value_heads)
+    }
+}
+
+impl Default for AceStepConfig {
+    fn default() -> Self {
+        Self {
+            hidden_size: 2048,
+            intermediate_size: 6144,
+            num_hidden_layers: 24,
+            num_attention_heads: 16,
+            num_key_value_heads: 8,
+            head_dim: 128,
+            encoder_hidden_size: None,
+            encoder_intermediate_size: None,
+            encoder_num_attention_heads: None,
+            encoder_num_key_value_heads: None,
+            hidden_act: Activation::Silu,
+            rms_norm_eps: 1e-6,
+            rope_theta: 1_000_000.0,
+            use_sliding_window: true,
+            sliding_window: 128,
+            in_channels: 192,
+            patch_size: 2,
+            text_hidden_dim: 1024,
+            audio_acoustic_hidden_dim: 64,
+            timbre_hidden_dim: 64,
+            pool_window_size: 5,
+            vocab_size: 64003,
+            max_position_embeddings: 32768,
+            fsq_dim: 2048,
+            fsq_input_levels: vec![8, 8, 8, 5, 5, 5],
+            fsq_input_num_quantizers: 1,
+            attention_bias: false,
+            attention_dropout: 0.0,
+            num_lyric_encoder_hidden_layers: 8,
+            num_timbre_encoder_hidden_layers: 4,
+            num_attention_pooler_hidden_layers: 2,
+            num_audio_decoder_hidden_layers: 24,
+            layer_types: default_layer_types(),
+            is_turbo: false,
+        }
+    }
+}
+
+// -- VaeConfig default value functions ---------------------------------------
+
+fn default_vae_encoder_hidden_size() -> usize {
+    128
+}
+fn default_vae_decoder_channels() -> usize {
+    128
+}
+fn default_vae_decoder_input_channels() -> usize {
+    64
+}
+fn default_vae_audio_channels() -> usize {
+    2
+}
+fn default_vae_downsampling_ratios() -> Vec<usize> {
+    vec![2, 4, 4, 6, 10]
+}
+fn default_vae_channel_multiples() -> Vec<usize> {
+    vec![1, 2, 4, 8, 16]
+}
+fn default_vae_sampling_rate() -> usize {
+    48000
+}
+
+/// Configuration for the ACE-Step VAE (variational autoencoder) that converts
+/// between waveform audio and the latent space consumed by the DiT.
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct VaeConfig {
+    #[serde(default = "default_vae_encoder_hidden_size")]
+    pub encoder_hidden_size: usize,
+    #[serde(default = "default_vae_decoder_channels")]
+    pub decoder_channels: usize,
+    #[serde(default = "default_vae_decoder_input_channels")]
+    pub decoder_input_channels: usize,
+    #[serde(default = "default_vae_audio_channels")]
+    pub audio_channels: usize,
+    #[serde(default = "default_vae_downsampling_ratios")]
+    pub downsampling_ratios: Vec<usize>,
+    #[serde(default = "default_vae_channel_multiples")]
+    pub channel_multiples: Vec<usize>,
+    #[serde(default = "default_vae_sampling_rate")]
+    pub sampling_rate: usize,
+}
+
+impl VaeConfig {
+    /// Computes the total hop length as the product of all downsampling ratios.
+    pub fn hop_length(&self) -> usize {
+        self.downsampling_ratios.iter().product()
+    }
+}
+
+impl Default for VaeConfig {
+    fn default() -> Self {
+        Self {
+            encoder_hidden_size: 128,
+            decoder_channels: 128,
+            decoder_input_channels: 64,
+            audio_channels: 2,
+            downsampling_ratios: vec![2, 4, 4, 6, 10],
+            channel_multiples: vec![1, 2, 4, 8, 16],
+            sampling_rate: 48000,
+        }
+    }
+}

crates/neuron/src/harness/arch/acestep/model.rs

@@ -0,0 +1,520 @@
+//! Top-level ACE-Step conditional generation model.
+//!
+//! Mirrors the subset of Python's `AceStepConditionGenerationModel` needed for
+//! **single-sample DiT-only Text2Music** with base/sft models. Owns the
+//! condition encoder, DiT decoder, audio tokenizer/detokenizer, and
+//! null-condition embedding. Provides `prepare_condition` and
+//! `generate_latents` as the primary API, operating in latent space.
+//!
+//! ## Current scope vs Python
+//!
+//! Implemented:
+//! - Text2Music with base/sft (continuous schedule + APG guidance)
+//! - Text2Music with turbo (discrete schedule, no CFG)
+//! - Real tokenizer masks, encoder_attention_mask propagation
+//! - Cover mode (`is_covers`, tokenize/detokenize round-trip in prepare_condition)
+//! - `precomputed_lm_hints_25Hz` / `audio_codes` (LM-generated audio codes)
+//! - Audio cover strength blending (switch condition mid-loop with KV cache reset)
+//! - Custom timestep schedules for turbo models (mapped to VALID_TIMESTEPS)
+//!
+//! Not yet implemented (Python `generate_audio` features):
+//! - Batched inference (batch > 1)
+//! - LoRA support
+
+use candle_core::{DType, Device, Result, Tensor};
+use candle_nn::VarBuilder;
+
+use super::AceStepConfig;
+use super::condition::ConditionEncoder;
+use super::dit::AceStepDiTModel;
+use super::sampling;
+use super::tokenizer::{AudioTokenDetokenizer, AudioTokenizer};
+
+/// Inputs for the condition encoder.
+pub struct ConditionInputs<'a> {
+    /// Text encoder hidden states `(B, T_text, text_hidden_dim)`.
+    pub text_hidden_states: &'a Tensor,
+    /// Text attention mask `(B, T_text)`.
+    pub text_attention_mask: &'a Tensor,
+    /// Lyric token embeddings `(B, T_lyric, text_hidden_dim)`.
+    pub lyric_hidden_states: &'a Tensor,
+    /// Lyric attention mask `(B, T_lyric)`.
+    pub lyric_attention_mask: &'a Tensor,
+    /// Packed reference audio features `(N, T_audio, timbre_hidden_dim)`.
+    pub refer_audio_packed: &'a Tensor,
+    /// 1D order mask mapping segments to batch indices `(N,)`.
+    pub refer_audio_order_mask: &'a Tensor,
+    /// Source latents `(B, T, acoustic_dim)` — silence for text2music.
+    pub src_latents: &'a Tensor,
+    /// Chunk masks `(B, T, acoustic_dim)` — ones for full generation.
+    pub chunk_masks: &'a Tensor,
+    /// Silence latent for tokenizer padding `(1, T_max, acoustic_dim)`.
+    pub silence_latent: &'a Tensor,
+    /// Boolean tensor `(B,)` — true for cover mode items.
+    pub is_covers: &'a Tensor,
+    /// Precomputed LM hints at 25Hz `(B, T, acoustic_dim)`. Skips tokenize/detokenize.
+    pub precomputed_lm_hints_25hz: Option<&'a Tensor>,
+    /// Audio codes from LM `(B, N, 1)`. Skips tokenize, uses quantizer → detokenize.
+    pub audio_codes: Option<&'a Tensor>,
+    /// Text hidden states for the non-cover condition (used when `audio_cover_strength < 1.0`).
+    pub non_cover_text_hidden_states: Option<&'a Tensor>,
+    /// Attention mask for non-cover text.
+    pub non_cover_text_attention_mask: Option<&'a Tensor>,
+}
+
+/// Prepared conditioning for the DiT decoder.
+pub struct PreparedCondition {
+    pub encoder_hidden_states: Tensor,
+    pub encoder_attention_mask: Tensor,
+    pub context_latents: Tensor,
+}
+
+/// Options for latent generation.
+pub struct GenerateOptions {
+    pub seed: Option<u64>,
+    pub num_steps: usize,
+    pub guidance_scale: f64,
+    pub shift: f64,
+    pub is_turbo: bool,
+    /// ODE (deterministic) or SDE (stochastic) sampling.
+    pub infer_method: sampling::InferMethod,
+    /// Timestep range in which CFG/APG guidance is applied.
+    /// Outside this range, only the conditional prediction is used.
+    /// Default: `(0.0, 1.0)` — apply guidance at all timesteps.
+    pub cfg_interval: (f64, f64),
+    /// Custom timestep schedule for turbo models (1–20 values).
+    /// Each value is mapped to the nearest valid timestep. When `None`,
+    /// the shift-based schedule is used.
+    pub timesteps: Option<Vec<f64>>,
+    /// Fraction of denoising steps that use the cover condition (0.0–1.0).
+    /// At `cover_steps = round(num_steps * audio_cover_strength)` the loop
+    /// switches to the non-cover condition and clears the KV cache.
+    /// Default: `1.0` (no blending — cover condition used for all steps).
+    pub audio_cover_strength: f64,
+    /// Starting noise `(1, frames, acoustic_dim)`, replacing the seeded
+    /// draw. Lets a parity test start from the reference's saved noise.
+    pub initial_noise: Option<Tensor>,
+    /// Wavelet correction after each turbo step; on by default, as upstream.
+    pub dcw: Option<sampling::Dcw>,
+}
+
+impl Default for GenerateOptions {
+    fn default() -> Self {
+        Self {
+            seed: None,
+            num_steps: 50,
+            guidance_scale: 5.0,
+            shift: 1.0,
+            is_turbo: false,
+            infer_method: sampling::InferMethod::Ode,
+            cfg_interval: (0.0, 1.0),
+            timesteps: None,
+            audio_cover_strength: 1.0,
+            initial_noise: None,
+            dcw: Some(sampling::Dcw::default()),
+        }
+    }
+}
+
+/// Top-level ACE-Step model that owns all sub-modules and provides
+/// `prepare_condition` + `generate_latents` as the main API.
+pub struct AceStepConditionGenerationModel {
+    config: AceStepConfig,
+    encoder: ConditionEncoder,
+    decoder: AceStepDiTModel,
+    pub tokenizer: AudioTokenizer,
+    pub detokenizer: AudioTokenDetokenizer,
+    null_condition_emb: Tensor,
+}
+
+impl AceStepConditionGenerationModel {
+    /// Load all sub-modules from a single safetensors checkpoint.
+    pub fn new(config: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let encoder = ConditionEncoder::new(config, vb.pp("encoder"))?;
+        let decoder = AceStepDiTModel::new(config, vb.pp("decoder"))?;
+        let tokenizer = AudioTokenizer::new(config, vb.pp("tokenizer"))?;
+        let detokenizer = AudioTokenDetokenizer::new(config, vb.pp("detokenizer"))?;
+        // XL models use encoder_hidden_size for null_condition_emb (it operates
+        // in encoder space, not decoder space).
+        let null_condition_emb =
+            vb.get((1, 1, config.encoder_hidden_size()), "null_condition_emb")?;
+
+        Ok(Self {
+            config: config.clone(),
+            encoder,
+            decoder,
+            tokenizer,
+            detokenizer,
+            null_condition_emb,
+        })
+    }
+
+    /// Tokenize audio latents: pad to `pool_window_size` multiple, reshape, quantize.
+    ///
+    /// Returns `(quantized, indices, downsampled_mask)`.
+    fn tokenize(
+        &self,
+        hidden_states: &Tensor,
+        silence_latent: &Tensor,
+        attention_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor, Tensor)> {
+        let (b, t, d) = hidden_states.dims3()?;
+        let pws = self.tokenizer.pool_window_size();
+
+        // Pad to multiple of pool_window_size using silence_latent
+        let (hidden_states, attention_mask) = if t % pws != 0 {
+            let pad_len = pws - (t % pws);
+            let pad = silence_latent
+                .narrow(1, 0, pad_len)?
+                .broadcast_as((b, pad_len, d))?
+                .contiguous()?;
+            let hs = Tensor::cat(&[hidden_states, &pad], 1)?;
+            let mask_pad = Tensor::zeros(
+                (b, pad_len),
+                attention_mask.dtype(),
+                attention_mask.device(),
+            )?;
+            let am = Tensor::cat(&[attention_mask, &mask_pad], 1)?;
+            (hs, am)
+        } else {
+            (hidden_states.clone(), attention_mask.clone())
+        };
+
+        let t_padded = hidden_states.dim(1)?;
+        let num_patches = t_padded / pws;
+
+        // Reshape: (B, T, D) -> (B, T/pws, pws, D)
+        let reshaped = hidden_states.reshape((b, num_patches, pws, d))?;
+
+        // Downsample attention_mask: reshape (B, T) -> (B, T/pws, pws) then max
+        let mask_reshaped = attention_mask.reshape((b, num_patches, pws))?;
+        let downsampled_mask = mask_reshaped.max(2)?;
+
+        let (quantized, indices) = self.tokenizer.forward(&reshaped)?;
+        Ok((quantized, indices, downsampled_mask))
+    }
+
+    /// Detokenize quantized representations back to acoustic space.
+    fn detokenize(&self, quantized: &Tensor) -> Result<Tensor> {
+        self.detokenizer.forward(quantized)
+    }
+
+    /// Prepare conditioning: encode text/lyrics/timbre into packed encoder
+    /// hidden states, and build context latents from src_latents + chunk_masks.
+    ///
+    /// When cover mode is active (`is_covers > 0`), source latents are replaced
+    /// with tokenized-then-detokenized LM hints.
+    pub fn prepare_condition(&self, inputs: &ConditionInputs) -> Result<PreparedCondition> {
+        let (encoder_hidden_states, encoder_attention_mask) = self.encoder.forward(
+            inputs.text_hidden_states,
+            inputs.text_attention_mask,
+            inputs.lyric_hidden_states,
+            inputs.lyric_attention_mask,
+            inputs.refer_audio_packed,
+            inputs.refer_audio_order_mask,
+        )?;
+
+        // Cover mode: compute LM hints and replace src_latents where is_covers > 0.
+        // Skip entirely when is_covers is all-false (text2music) — no tokenize/detokenize needed.
+        let latent_len = inputs.src_latents.dim(1)?;
+
+        let any_covers = inputs
+            .is_covers
+            .gt(0u8)?
+            .to_dtype(candle_core::DType::U32)?
+            .sum_all()?
+            .to_scalar::<u32>()?
+            > 0;
+
+        let src_latents = if any_covers {
+            let lm_hints_25hz = if let Some(precomputed) = inputs.precomputed_lm_hints_25hz {
+                precomputed.narrow(1, 0, latent_len)?
+            } else if let Some(audio_codes) = inputs.audio_codes {
+                let quantized = self.tokenizer.get_output_from_indices(audio_codes)?;
+                let detokenized = self.detokenize(&quantized)?;
+                let det_len = detokenized.dim(1)?;
+                if det_len >= latent_len {
+                    detokenized.narrow(1, 0, latent_len)?
+                } else {
+                    let pad = inputs.silence_latent.narrow(1, 0, latent_len - det_len)?;
+                    Tensor::cat(&[&detokenized, &pad], 1)?
+                }
+            } else {
+                let b = inputs.src_latents.dim(0)?;
+                let attention_mask = Tensor::ones(
+                    (b, latent_len),
+                    inputs.src_latents.dtype(),
+                    inputs.src_latents.device(),
+                )?;
+                let (quantized, _indices, _mask) =
+                    self.tokenize(inputs.src_latents, inputs.silence_latent, &attention_mask)?;
+                let detokenized = self.detokenize(&quantized)?;
+                detokenized.narrow(1, 0, latent_len)?
+            };
+
+            let is_covers_3d = inputs
+                .is_covers
+                .gt(0u8)?
+                .unsqueeze(1)?
+                .unsqueeze(2)?
+                .broadcast_as(inputs.src_latents.shape())?;
+            is_covers_3d.where_cond(&lm_hints_25hz, inputs.src_latents)?
+        } else {
+            inputs.src_latents.clone()
+        };
+
+        // Context = [src_latents, chunk_masks] along feature dim
+        let context_latents =
+            Tensor::cat(&[&src_latents, inputs.chunk_masks], candle_core::D::Minus1)?;
+
+        Ok(PreparedCondition {
+            encoder_hidden_states,
+            encoder_attention_mask,
+            context_latents,
+        })
+    }
+
+    /// Clear cross-attention KV caches. Call between generations with different
+    /// encoder states (different prompts).
+    pub fn clear_kv_cache(&mut self) {
+        self.decoder.clear_kv_cache();
+    }
+
+    /// Generate audio latents from prepared conditions.
+    ///
+    /// Returns the denoised latent `(B, T, acoustic_dim)` ready for VAE decode.
+    /// Automatically caches cross-attention K/V on the first denoising step and
+    /// reuses them for subsequent steps (encoder states are constant).
+    ///
+    /// When `non_cover_condition` is provided and `opts.audio_cover_strength < 1.0`,
+    /// the loop switches from cover to non-cover condition partway through and
+    /// clears the KV cache at the transition.
+    pub fn generate_latents(
+        &mut self,
+        condition: &PreparedCondition,
+        non_cover_condition: Option<&PreparedCondition>,
+        opts: &GenerateOptions,
+        device: &Device,
+        dtype: DType,
+    ) -> Result<Tensor> {
+        // Clear KV cache so it gets rebuilt for this generation's encoder states.
+        self.decoder.clear_kv_cache();
+
+        let latent_frames = condition.context_latents.dim(1)?;
+        let acoustic_dim = self.config.audio_acoustic_hidden_dim;
+
+        // Drawn on the host from a seeded generator rather than with the
+        // device RNG: candle cannot seed its CPU RNG at all, and a CUDA
+        // draw differs from a CPU one. A host draw makes the same seed give
+        // the same song on every device.
+        let noise = match &opts.initial_noise {
+            Some(n) => n.to_device(device)?.to_dtype(dtype)?,
+            None => seeded_normal(opts.seed, (1, latent_frames, acoustic_dim))?
+                .to_device(device)?
+                .to_dtype(dtype)?,
+        };
+
+        let mut xt = noise;
+
+        if opts.is_turbo {
+            // Turbo: discrete schedule, no CFG, final step = x0 prediction
+            let t_schedule = match &opts.timesteps {
+                Some(ts) => sampling::get_custom_turbo_schedule(ts)?,
+                None => sampling::get_turbo_schedule(opts.shift),
+            };
+            let num_steps = t_schedule.len();
+            let cover_steps = (num_steps as f64 * opts.audio_cover_strength).round() as usize;
+            let mut active = if cover_steps == 0 {
+                if let Some(nc) = non_cover_condition {
+                    nc
+                } else {
+                    condition
+                }
+            } else {
+                condition
+            };
+
+            for step in 0..num_steps {
+                // Switch to non-cover condition at the transition point
+                if let Some(nc) = non_cover_condition
+                    && step == cover_steps
+                    && step > 0
+                {
+                    active = nc;
+                    self.decoder.clear_kv_cache();
+                }
+
+                let t_curr = t_schedule[step];
+                let t_tensor =
+                    Tensor::from_vec(vec![t_curr as f32], (1,), device)?.to_dtype(dtype)?;
+
+                let vt = self.decoder.forward(
+                    &xt,
+                    &t_tensor,
+                    &t_tensor,
+                    None,
+                    &active.encoder_hidden_states,
+                    Some(&active.encoder_attention_mask),
+                    &active.context_latents,
+                )?;
+
+                let xt_before = xt.clone();
+                if step == num_steps - 1 {
+                    xt = sampling::get_x0_from_noise(&xt, &vt, t_curr)?;
+                } else {
+                    let t_next = t_schedule[step + 1];
+                    match opts.infer_method {
+                        sampling::InferMethod::Ode => {
+                            xt = sampling::euler_step(&xt, &vt, t_curr - t_next)?;
+                        }
+                        sampling::InferMethod::Sde => {
+                            xt = sampling::sde_step(&xt, &vt, t_curr, t_next)?;
+                        }
+                    }
+                }
+                // DCW against the clean sample this step predicted,
+                // after every step including the final x0 (as upstream).
+                if let Some(dcw) = &opts.dcw {
+                    let denoised = (&xt_before - (&vt * t_curr)?)?;
+                    xt = dcw.apply(&xt, &denoised, t_curr)?;
+                }
+            }
+        } else {
+            // Base/SFT: continuous schedule with APG guidance
+            let schedule = sampling::get_schedule(opts.num_steps, opts.shift);
+            let mut momentum_buffer = sampling::MomentumBuffer::new(-0.75);
+            let do_cfg = opts.guidance_scale > 1.0;
+            let cover_steps = (opts.num_steps as f64 * opts.audio_cover_strength).round() as usize;
+
+            // Build doubled tensors for CFG from cover condition
+            let build_cfg_tensors =
+                |cond: &PreparedCondition, null_emb: &Tensor| -> Result<(Tensor, Tensor, Tensor)> {
+                    let null_enc = null_emb
+                        .broadcast_as(cond.encoder_hidden_states.shape())?
+                        .contiguous()?;
+                    if do_cfg {
+                        let enc = Tensor::cat(&[&cond.encoder_hidden_states, &null_enc], 0)?;
+                        let mask = Tensor::cat(
+                            &[&cond.encoder_attention_mask, &cond.encoder_attention_mask],
+                            0,
+                        )?;
+                        let ctx = Tensor::cat(&[&cond.context_latents, &cond.context_latents], 0)?;
+                        Ok((enc, mask, ctx))
+                    } else {
+                        Ok((
+                            cond.encoder_hidden_states.clone(),
+                            cond.encoder_attention_mask.clone(),
+                            cond.context_latents.clone(),
+                        ))
+                    }
+                };
+
+            let initial_condition = if cover_steps == 0 {
+                non_cover_condition.unwrap_or(condition)
+            } else {
+                condition
+            };
+            let (mut enc_cond, mut enc_mask_cond, mut ctx_doubled) =
+                build_cfg_tensors(initial_condition, &self.null_condition_emb)?;
+
+            let (cfg_start, cfg_end) = opts.cfg_interval;
+
+            for step in 0..opts.num_steps {
+                // Switch to non-cover condition at the transition point
+                if let Some(nc) = non_cover_condition
+                    && step == cover_steps
+                    && step > 0
+                {
+                    self.decoder.clear_kv_cache();
+                    let rebuilt = build_cfg_tensors(nc, &self.null_condition_emb)?;
+                    enc_cond = rebuilt.0;
+                    enc_mask_cond = rebuilt.1;
+                    ctx_doubled = rebuilt.2;
+                }
+
+                let t_curr = schedule[step];
+                let t_prev = schedule[step + 1];
+
+                let x_input = if do_cfg {
+                    Tensor::cat(&[&xt, &xt], 0)?
+                } else {
+                    xt.clone()
+                };
+                let batch_size = x_input.dim(0)?;
+                let t_tensor =
+                    Tensor::from_vec(vec![t_curr as f32; batch_size], (batch_size,), device)?
+                        .to_dtype(dtype)?;
+
+                let vt = self.decoder.forward(
+                    &x_input,
+                    &t_tensor,
+                    &t_tensor,
+                    None,
+                    &enc_cond,
+                    Some(&enc_mask_cond),
+                    &ctx_doubled,
+                )?;
+
+                let vt = if do_cfg {
+                    let v_cond = vt.narrow(0, 0, 1)?;
+                    let v_uncond = vt.narrow(0, 1, 1)?;
+                    let apply_guidance = t_curr >= cfg_start && t_curr <= cfg_end;
+                    if apply_guidance {
+                        sampling::apg_forward(
+                            &v_cond,
+                            &v_uncond,
+                            opts.guidance_scale,
+                            &mut momentum_buffer,
+                            2.5,
+                        )?
+                    } else {
+                        v_cond
+                    }
+                } else {
+                    vt
+                };
+
+                match opts.infer_method {
+                    sampling::InferMethod::Ode => {
+                        xt = sampling::euler_step(&xt, &vt, t_curr - t_prev)?;
+                    }
+                    sampling::InferMethod::Sde => {
+                        xt = sampling::sde_step(&xt, &vt, t_curr, t_prev)?;
+                    }
+                }
+            }
+        }
+
+        Ok(xt)
+    }
+}
+
+/// Standard-normal noise from a seeded host generator (SplitMix64 +
+/// Box-Muller). An unseeded draw takes its seed from the process's
+/// random hasher keys.
+pub fn seeded_normal(seed: Option<u64>, shape: (usize, usize, usize)) -> Result<Tensor> {
+    let mut state = seed.unwrap_or_else(|| {
+        use std::hash::BuildHasher;
+        std::collections::hash_map::RandomState::new().hash_one(0u8)
+    });
+    let mut next = move || {
+        state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
+        let mut z = state;
+        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+        z ^ (z >> 31)
+    };
+    // 53 random bits mapped into (0, 1]: never zero, so ln() is finite.
+    let mut unit = move || ((next() >> 11) as f64 + 1.0) / (1u64 << 53) as f64;
+    let n = shape.0 * shape.1 * shape.2;
+    let mut out = Vec::with_capacity(n + 1);
+    while out.len() < n {
+        let (u1, u2) = (unit(), unit());
+        let r = (-2.0 * u1.ln()).sqrt();
+        let theta = 2.0 * std::f64::consts::PI * u2;
+        out.push((r * theta.cos()) as f32);
+        out.push((r * theta.sin()) as f32);
+    }
+    out.truncate(n);
+    Tensor::from_vec(out, shape, &Device::Cpu)
+}

crates/neuron/src/harness/arch/acestep/pipeline.rs

@@ -0,0 +1,696 @@
+//! High-level ACE-Step pipeline with text encoder inside.
+//!
+//! Vendored with the rest of this module from the candle example in
+//! huggingface/candle#3451 (`candle-examples/examples/ace-step/pipeline.rs`).
+//!
+//! Accepts raw text prompts — handles tokenization, text encoding,
+//! conditioning, DiT denoising, and VAE decoding internally.
+
+use super::qwen3;
+use super::{
+    AceStepConfig, VaeConfig,
+    lm::{LmConfig, LmPipeline, TokenIds},
+    model::{AceStepConditionGenerationModel, ConditionInputs, GenerateOptions},
+    sampling,
+    vae::{AutoencoderOobleck, OobleckDecoder},
+};
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+
+/// Metadata from LM CoT output, used to enrich the DiT text prompt.
+///
+/// When the LM generates chain-of-thought metadata (BPM, keyscale, etc.),
+/// these values should be fed back into the DiT caption — matching the Python
+/// `_dict_to_meta_string` / `_extract_caption_and_language` flow.
+#[derive(Default)]
+pub struct LmMetadata {
+    pub bpm: Option<String>,
+    pub timesignature: Option<String>,
+    pub keyscale: Option<String>,
+    /// LM-generated caption — replaces the user prompt in the DiT text input.
+    pub caption: Option<String>,
+    /// Language code for lyrics (default: "en").
+    pub language: Option<String>,
+}
+
+/// Generation parameters.
+pub struct GenerationParams {
+    pub duration_secs: f64,
+    pub seed: Option<u64>,
+    pub num_steps: Option<usize>,
+    pub guidance_scale: f64,
+    pub shift: f64,
+    pub is_turbo: bool,
+    /// Precomputed LM hints at 25Hz `(1, T, acoustic_dim)`.
+    pub lm_hints_25hz: Option<Tensor>,
+    /// LM-generated metadata to enrich the DiT text prompt.
+    pub lm_metadata: Option<LmMetadata>,
+    /// Fraction of denoising steps using cover condition (0.0–1.0).
+    pub audio_cover_strength: f64,
+    /// Timestep range for CFG/APG guidance.
+    pub cfg_interval: (f64, f64),
+    /// ODE (deterministic) or SDE (stochastic) sampling.
+    pub infer_method: sampling::InferMethod,
+    /// Latent frames per VAE-decoder chunk. `None` uses the built-in default
+    /// (128 frames ≈ 5.1 s of audio). Lower values cut peak decode memory at
+    /// the cost of more compute overhead; raise when memory is ample.
+    pub vae_chunk_frames: Option<usize>,
+    /// Overlap in latent frames on each side of a VAE-decoder chunk. `None`
+    /// uses the default (16). Must satisfy `chunk > 2 * overlap`.
+    pub vae_chunk_overlap: Option<usize>,
+    /// Starting noise for the DiT, replacing the seeded draw (parity tests).
+    pub initial_noise: Option<Tensor>,
+    /// Vocal language code (e.g. `"en"`); the planner's, when it ran.
+    pub language: Option<String>,
+}
+
+impl Default for GenerationParams {
+    fn default() -> Self {
+        Self {
+            duration_secs: 10.0,
+            seed: None,
+            num_steps: None,
+            guidance_scale: 5.0,
+            shift: 1.0,
+            is_turbo: false,
+            lm_hints_25hz: None,
+            lm_metadata: None,
+            audio_cover_strength: 1.0,
+            cfg_interval: (0.0, 1.0),
+            infer_method: sampling::InferMethod::Ode,
+            vae_chunk_frames: None,
+            vae_chunk_overlap: None,
+            initial_noise: None,
+            language: None,
+        }
+    }
+}
+
+/// Stereo audio output.
+pub struct AudioOutput {
+    /// Stereo waveform `(2, T)` in `[-1, 1]`, 48 kHz.
+    pub audio: Tensor,
+    pub sample_rate: u32,
+    /// The denoised latents `(1, T, acoustic_dim)` the VAE decoded.
+    pub latents: Tensor,
+}
+
+/// High-level pipeline: text prompt → stereo audio.
+pub struct AceStepPipeline {
+    model: AceStepConditionGenerationModel,
+    text_encoder: qwen3::Model,
+    tokenizer: tokenizers::Tokenizer,
+    vae_decoder: OobleckDecoder,
+    vae_encoder: Option<AutoencoderOobleck>,
+    lm: Option<LmState>,
+    silence_latent: Tensor,
+    dit_config: AceStepConfig,
+    vae_config: VaeConfig,
+    device: Device,
+    dtype: DType,
+}
+
+struct LmState {
+    pipeline: LmPipeline,
+    tokenizer: tokenizers::Tokenizer,
+}
+
+impl AceStepPipeline {
+    /// Load all components.
+    ///
+    /// - `dit_vb` / `vae_vb` / `text_vb`: VarBuilders for each model
+    /// - `text_config`: Qwen3 config (for the embedding model, NOT the LM)
+    /// - `tokenizer_path`: `tokenizer.json` for the text encoder
+    /// - `silence_latent`: `(1, T, acoustic_dim)` pre-loaded
+    /// - `load_encoder`: load VAE encoder for cover mode
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        dit_config: &AceStepConfig,
+        vae_config: &VaeConfig,
+        text_config: &qwen3::Config,
+        dit_vb: VarBuilder,
+        vae_vb: VarBuilder,
+        text_vb: VarBuilder,
+        tokenizer_path: &std::path::Path,
+        silence_latent: Tensor,
+        load_encoder: bool,
+    ) -> Result<Self> {
+        let device = dit_vb.device().clone();
+        let dtype = dit_vb.dtype();
+
+        let text_encoder = qwen3::Model::new(text_config, text_vb)?;
+        let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
+            .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
+        let model = AceStepConditionGenerationModel::new(dit_config, dit_vb)?;
+        let vae_decoder = OobleckDecoder::new(vae_config, vae_vb.pp("decoder"))?;
+        let vae_encoder = if load_encoder {
+            Some(AutoencoderOobleck::new(vae_config, vae_vb)?)
+        } else {
+            None
+        };
+
+        Ok(Self {
+            model,
+            text_encoder,
+            tokenizer,
+            vae_decoder,
+            vae_encoder,
+            lm: None,
+            silence_latent,
+            dit_config: dit_config.clone(),
+            vae_config: vae_config.clone(),
+            device,
+            dtype,
+        })
+    }
+
+    /// Number of latent frames for a given duration.
+    pub fn latent_frames(&self, duration_secs: f64) -> usize {
+        let hop = self.vae_config.hop_length();
+        (duration_secs * self.vae_config.sampling_rate as f64 / hop as f64).ceil() as usize
+    }
+
+    /// Generate music from a text prompt and optional lyrics.
+    pub fn text2music(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        params: &GenerationParams,
+    ) -> Result<AudioOutput> {
+        let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
+            prompt,
+            lyrics,
+            params.duration_secs,
+            params.lm_metadata.as_ref(),
+            params.language.as_deref(),
+        )?;
+
+        let latent_frames = self.latent_frames(params.duration_secs);
+        let timbre_frames = 750.min(self.silence_latent.dim(1)?);
+
+        self.generate_inner(
+            &text_hs,
+            &text_mask,
+            &lyric_hs,
+            &lyric_mask,
+            &self.silence_latent.narrow(1, 0, latent_frames)?,
+            &self.silence_latent.narrow(1, 0, timbre_frames)?,
+            params,
+            latent_frames,
+        )
+    }
+
+    /// Generate a cover conditioned on reference audio `(1, 2, T)` at 48kHz.
+    ///
+    /// - `reference_audio`: structural source (melody/rhythm)
+    /// - `timbre_audio`: optional separate timbre source; if `None`, uses
+    ///   `reference_audio` for both structure and timbre (Python separates these
+    ///   as `src_audio` vs `reference_audio`)
+    ///
+    /// When `params.lm_hints_25hz` is set, uses the LM-derived hints with
+    /// `is_covers=true` (tokenize/detokenize path in prepare_condition).
+    /// Otherwise, uses raw reference latents directly as context (no lossy
+    /// FSQ round-trip).
+    pub fn cover(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        reference_audio: &Tensor,
+        timbre_audio: Option<&Tensor>,
+        params: &GenerationParams,
+    ) -> Result<AudioOutput> {
+        let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
+            prompt,
+            lyrics,
+            params.duration_secs,
+            params.lm_metadata.as_ref(),
+            params.language.as_deref(),
+        )?;
+
+        let vae = self
+            .vae_encoder
+            .as_ref()
+            .ok_or_else(|| anyhow::anyhow!("VAE encoder not loaded (pass load_encoder=true)"))?;
+        let ref_latent = vae.encode_mean(reference_audio)?.transpose(1, 2)?;
+
+        let latent_frames = self.latent_frames(params.duration_secs);
+        let ref_len = ref_latent.dim(1)?;
+
+        let src_latents = if ref_len >= latent_frames {
+            ref_latent.narrow(1, 0, latent_frames)?
+        } else {
+            let pad = self.silence_latent.narrow(1, 0, latent_frames - ref_len)?;
+            Tensor::cat(&[&ref_latent, &pad], 1)?
+        };
+
+        // Timbre: use separate timbre audio if provided, otherwise reference audio
+        let refer_audio_packed = if let Some(timbre) = timbre_audio {
+            let timbre_latent = vae.encode_mean(timbre)?.transpose(1, 2)?;
+            let timbre_len = timbre_latent.dim(1)?;
+            let timbre_frames = 750.min(timbre_len);
+            timbre_latent.narrow(1, 0, timbre_frames)?
+        } else {
+            let timbre_frames = 750.min(ref_len);
+            ref_latent.narrow(1, 0, timbre_frames)?
+        };
+
+        // Only enable is_covers (tokenize/detokenize path) when LM hints
+        // are provided. Without LM hints, raw reference latents go directly
+        // into context — the lossy FSQ round-trip on raw VAE latents produces
+        // poor results.
+        self.generate_inner(
+            &text_hs,
+            &text_mask,
+            &lyric_hs,
+            &lyric_mask,
+            &src_latents,
+            &refer_audio_packed,
+            params,
+            latent_frames,
+        )
+    }
+
+    /// Selectively re-generate a region of existing audio while preserving the rest.
+    ///
+    /// - `source_audio`: original stereo waveform `(1, 2, T)` at 48kHz
+    /// - `start_secs` / `end_secs`: time range to regenerate (seconds)
+    ///
+    /// The preserved regions keep the original audio latents; the repaint region
+    /// is replaced with silence and the chunk_mask signals the DiT to generate new
+    /// content there.
+    pub fn repaint(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        source_audio: &Tensor,
+        start_secs: f64,
+        end_secs: f64,
+        params: &GenerationParams,
+    ) -> Result<AudioOutput> {
+        let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
+            prompt,
+            lyrics,
+            params.duration_secs,
+            params.lm_metadata.as_ref(),
+            params.language.as_deref(),
+        )?;
+
+        let vae = self
+            .vae_encoder
+            .as_ref()
+            .ok_or_else(|| anyhow::anyhow!("VAE encoder not loaded (pass load_encoder=true)"))?;
+        let src_latent = vae.encode_mean(source_audio)?.transpose(1, 2)?; // (1, T_lat, D)
+
+        let latent_frames = self.latent_frames(params.duration_secs);
+        let hop = self.vae_config.hop_length() as f64;
+        let sr = self.vae_config.sampling_rate as f64;
+
+        // Convert time range to latent frame indices
+        let start_frame = ((start_secs * sr / hop).floor() as usize).min(latent_frames);
+        let end_frame = ((end_secs * sr / hop).ceil() as usize).min(latent_frames);
+
+        anyhow::ensure!(
+            end_frame > start_frame,
+            "repaint region is empty: {start_secs}s..{end_secs}s → frames {start_frame}..{end_frame}"
+        );
+
+        // Crop/pad source latents to target length
+        let src_len = src_latent.dim(1)?;
+        let mut src_latents = if src_len >= latent_frames {
+            src_latent.narrow(1, 0, latent_frames)?
+        } else {
+            let pad = self.silence_latent.narrow(1, 0, latent_frames - src_len)?;
+            Tensor::cat(&[&src_latent, &pad], 1)?
+        };
+
+        // Replace repaint region with silence in src_latents
+        let silence_region = self
+            .silence_latent
+            .narrow(1, start_frame, end_frame - start_frame)?;
+        // Build: [original[:start], silence[start:end], original[end:]]
+        let parts: Vec<Tensor> = if start_frame > 0 && end_frame < latent_frames {
+            vec![
+                src_latents.narrow(1, 0, start_frame)?,
+                silence_region,
+                src_latents.narrow(1, end_frame, latent_frames - end_frame)?,
+            ]
+        } else if start_frame == 0 {
+            vec![
+                silence_region,
+                src_latents.narrow(1, end_frame, latent_frames - end_frame)?,
+            ]
+        } else {
+            vec![src_latents.narrow(1, 0, start_frame)?, silence_region]
+        };
+        src_latents = Tensor::cat(&parts.iter().collect::<Vec<_>>(), 1)?;
+
+        // Timbre from original audio (first 750 frames)
+        let timbre_frames = 750.min(src_len);
+        let refer_audio_packed = src_latent.narrow(1, 0, timbre_frames)?;
+
+        // Build chunk_mask: 0 = preserve, 1 = regenerate
+        let acoustic_dim = self.dit_config.audio_acoustic_hidden_dim;
+        let chunk_mask_dim = self.dit_config.in_channels - 2 * acoustic_dim;
+        let mut mask_data = vec![0.0f32; latent_frames * chunk_mask_dim];
+        for t in start_frame..end_frame {
+            for c in 0..chunk_mask_dim {
+                mask_data[t * chunk_mask_dim + c] = 1.0;
+            }
+        }
+        let chunk_masks =
+            Tensor::from_vec(mask_data, (1, latent_frames, chunk_mask_dim), &self.device)?
+                .to_dtype(self.dtype)?;
+
+        self.generate_inner_with_mask(
+            &text_hs,
+            &text_mask,
+            &lyric_hs,
+            &lyric_mask,
+            &src_latents,
+            &refer_audio_packed,
+            &chunk_masks,
+            params,
+            latent_frames,
+        )
+    }
+
+    /// Unload the LM model to free GPU/RAM memory before DiT denoising.
+    ///
+    /// The LM (~1.2 GB for 0.6B, ~3.5 GB for 1.7B) is only needed for
+    /// `generate_lm_hints()`. Dropping it before the DiT pass frees memory
+    /// for the denoising loop.
+    pub fn unload_lm(&mut self) {
+        self.lm = None;
+    }
+
+    /// Load the 5Hz LM model for audio code generation.
+    ///
+    /// Call once before `generate_lm_hints()`. The LM weights and tokenizer
+    /// are loaded from a HuggingFace repo (e.g. `ACE-Step/acestep-5Hz-lm-0.6B`).
+    pub fn load_lm(
+        &mut self,
+        lm_config: &qwen3::Config,
+        lm_vb: VarBuilder,
+        lm_tokenizer_path: &std::path::Path,
+    ) -> Result<()> {
+        let lm_model = qwen3::ModelForCausalLM::new(lm_config, lm_vb)?;
+        let lm_tokenizer = tokenizers::Tokenizer::from_file(lm_tokenizer_path)
+            .map_err(|e| anyhow::anyhow!("LM tokenizer: {e}"))?;
+        let token_ids = TokenIds::discover(|s| {
+            let enc = lm_tokenizer
+                .encode(s, false)
+                .map_err(|e| candle_core::Error::Msg(format!("encode {s:?}: {e}")))?;
+            Ok(enc.get_ids().to_vec())
+        })?;
+        let pipeline = LmPipeline::new(lm_model, token_ids, self.device.clone());
+        self.lm = Some(LmState {
+            pipeline,
+            tokenizer: lm_tokenizer,
+        });
+        Ok(())
+    }
+
+    /// Generate LM audio codes and convert to 25Hz latent hints.
+    ///
+    /// Returns `(lm_hints_tensor, metadata)`. The hints can be passed as
+    /// `GenerationParams.lm_hints_25hz` to `text2music()` or `cover()`.
+    pub fn generate_lm_hints(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        duration_secs: f64,
+        lm_config: &LmConfig,
+    ) -> Result<(Tensor, std::collections::BTreeMap<String, String>)> {
+        let lm = self
+            .lm
+            .as_mut()
+            .ok_or_else(|| anyhow::anyhow!("LM not loaded — call load_lm() first"))?;
+
+        let prompt_text = LmPipeline::format_prompt(prompt, lyrics);
+        let prompt_enc = lm
+            .tokenizer
+            .encode(prompt_text.as_str(), false)
+            .map_err(|e| anyhow::anyhow!("LM encode: {e}"))?;
+
+        let decode_fn = |t: u32| -> candle_core::Result<String> {
+            lm.tokenizer
+                .decode(&[t], false)
+                .map_err(|e| candle_core::Error::Msg(format!("{e}")))
+        };
+        let encode_fn = |s: &str| -> candle_core::Result<Vec<u32>> {
+            let e = lm
+                .tokenizer
+                .encode(s, false)
+                .map_err(|e| candle_core::Error::Msg(format!("{e}")))?;
+            Ok(e.get_ids().to_vec())
+        };
+
+        let output = lm.pipeline.generate(
+            prompt_enc.get_ids(),
+            duration_secs,
+            lm_config,
+            &decode_fn,
+            &encode_fn,
+        )?;
+
+        anyhow::ensure!(
+            !output.audio_codes.is_empty(),
+            "LM generated no audio codes"
+        );
+        let hints = self.audio_codes_to_latents(&output.audio_codes)?;
+        Ok((hints, output.metadata))
+    }
+
+    /// Convert LM-generated audio codes to 25Hz latent hints.
+    pub fn audio_codes_to_latents(&self, codes: &[i64]) -> Result<Tensor> {
+        let code_tensor = Tensor::new(codes, &self.device)?
+            .unsqueeze(0)?
+            .unsqueeze(2)?;
+        let quantized = self.model.tokenizer.get_output_from_indices(&code_tensor)?;
+        let quantized = quantized.to_dtype(self.dtype)?;
+        Ok(self.model.detokenizer.forward(&quantized)?)
+    }
+
+    // ---- Internals ----
+
+    /// The text-encoder and lyric prompts, exactly as upstream builds them:
+    /// the caption (the planner's, when it ran) with its metas, and the
+    /// lyrics under their language. `language` defaults to `"unknown"`,
+    /// upstream's own default when the caller names none.
+    pub fn format_prompts(
+        prompt: &str,
+        lyrics: &str,
+        duration_secs: f64,
+        metadata: Option<&LmMetadata>,
+        language: Option<&str>,
+    ) -> (String, String) {
+        let caption = metadata
+            .and_then(|m| m.caption.as_deref())
+            .unwrap_or(prompt);
+        let bpm = metadata.and_then(|m| m.bpm.as_deref()).unwrap_or("N/A");
+        let timesig = metadata
+            .and_then(|m| m.timesignature.as_deref())
+            .unwrap_or("N/A");
+        let keyscale = metadata
+            .and_then(|m| m.keyscale.as_deref())
+            .unwrap_or("N/A");
+        let caption_text = format!(
+            "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n\
+             # Caption\n{caption}\n\n# Metas\n\
+             - bpm: {bpm}\n- timesignature: {timesig}\n- keyscale: {keyscale}\n- duration: {} seconds\n\
+             <|endoftext|>\n",
+            duration_secs as u32
+        );
+        let language = metadata
+            .and_then(|m| m.language.as_deref())
+            .or(language)
+            .unwrap_or("unknown");
+        let lyrics = if lyrics.is_empty() {
+            "[Instrumental]"
+        } else {
+            lyrics
+        };
+        let lyric_text = format!("# Languages\n{language}\n\n# Lyric\n{lyrics}<|endoftext|>");
+        (caption_text, lyric_text)
+    }
+
+    /// Token ids of both prompts, as the encoders will see them.
+    pub fn prompt_token_ids(
+        &self,
+        caption_text: &str,
+        lyric_text: &str,
+    ) -> Result<(Vec<u32>, Vec<u32>)> {
+        let text = self
+            .tokenizer
+            .encode(caption_text, true)
+            .map_err(|e| anyhow::anyhow!("tokenization: {e}"))?;
+        let lyric = self
+            .tokenizer
+            .encode(lyric_text, true)
+            .map_err(|e| anyhow::anyhow!("lyric tokenization: {e}"))?;
+        Ok((text.get_ids().to_vec(), lyric.get_ids().to_vec()))
+    }
+
+    fn encode_text(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        duration_secs: f64,
+        metadata: Option<&LmMetadata>,
+        language: Option<&str>,
+    ) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
+        let (caption_text, lyric_text) =
+            Self::format_prompts(prompt, lyrics, duration_secs, metadata, language);
+        let encoding = self
+            .tokenizer
+            .encode(caption_text.as_str(), true)
+            .map_err(|e| anyhow::anyhow!("tokenization: {e}"))?;
+        let prompt_tokens = Tensor::new(encoding.get_ids(), &self.device)?.unsqueeze(0)?;
+        let text_hidden_states = self.text_encoder.forward(&prompt_tokens, 0)?;
+        self.text_encoder.clear_kv_cache();
+
+        let lyric_encoding = self
+            .tokenizer
+            .encode(lyric_text.as_str(), true)
+            .map_err(|e| anyhow::anyhow!("lyric tokenization: {e}"))?;
+        let lyric_tokens = Tensor::new(lyric_encoding.get_ids(), &self.device)?.unsqueeze(0)?;
+        let lyric_hidden_states = self.text_encoder.embed_tokens(&lyric_tokens)?;
+
+        let text_mask = Tensor::new(encoding.get_attention_mask(), &self.device)?
+            .unsqueeze(0)?
+            .to_dtype(self.dtype)?;
+        let lyric_mask = Tensor::new(lyric_encoding.get_attention_mask(), &self.device)?
+            .unsqueeze(0)?
+            .to_dtype(self.dtype)?;
+
+        Ok((
+            text_hidden_states,
+            text_mask,
+            lyric_hidden_states,
+            lyric_mask,
+        ))
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn generate_inner(
+        &mut self,
+        text_hidden_states: &Tensor,
+        text_attention_mask: &Tensor,
+        lyric_hidden_states: &Tensor,
+        lyric_attention_mask: &Tensor,
+        src_latents: &Tensor,
+        refer_audio_packed: &Tensor,
+        params: &GenerationParams,
+        latent_frames: usize,
+    ) -> Result<AudioOutput> {
+        let acoustic_dim = self.dit_config.audio_acoustic_hidden_dim;
+        let chunk_mask_dim = self.dit_config.in_channels - 2 * acoustic_dim;
+        let chunk_masks =
+            Tensor::ones((1, latent_frames, chunk_mask_dim), self.dtype, &self.device)?;
+
+        self.generate_inner_with_mask(
+            text_hidden_states,
+            text_attention_mask,
+            lyric_hidden_states,
+            lyric_attention_mask,
+            src_latents,
+            refer_audio_packed,
+            &chunk_masks,
+            params,
+            latent_frames,
+        )
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn generate_inner_with_mask(
+        &mut self,
+        text_hidden_states: &Tensor,
+        text_attention_mask: &Tensor,
+        lyric_hidden_states: &Tensor,
+        lyric_attention_mask: &Tensor,
+        src_latents: &Tensor,
+        refer_audio_packed: &Tensor,
+        chunk_masks: &Tensor,
+        params: &GenerationParams,
+        latent_frames: usize,
+    ) -> Result<AudioOutput> {
+        let is_cover = params.lm_hints_25hz.is_some();
+
+        let condition = self.model.prepare_condition(&ConditionInputs {
+            text_hidden_states,
+            text_attention_mask,
+            lyric_hidden_states,
+            lyric_attention_mask,
+            refer_audio_packed,
+            refer_audio_order_mask: &Tensor::zeros((1,), DType::I64, &self.device)?,
+            src_latents,
+            chunk_masks,
+            silence_latent: &self.silence_latent.narrow(1, 0, latent_frames)?,
+            is_covers: &Tensor::from_vec(vec![if is_cover { 1u8 } else { 0u8 }], 1, &self.device)?,
+            precomputed_lm_hints_25hz: params.lm_hints_25hz.as_ref(),
+            audio_codes: None,
+            non_cover_text_hidden_states: None,
+            non_cover_text_attention_mask: None,
+        })?;
+
+        let num_steps = params
+            .num_steps
+            .unwrap_or(if params.is_turbo { 8 } else { 50 });
+
+        let xt = self.model.generate_latents(
+            &condition,
+            None,
+            &GenerateOptions {
+                seed: params.seed,
+                num_steps,
+                guidance_scale: params.guidance_scale,
+                shift: params.shift,
+                is_turbo: params.is_turbo,
+                audio_cover_strength: params.audio_cover_strength,
+                cfg_interval: params.cfg_interval,
+                infer_method: params.infer_method,
+                initial_noise: params.initial_noise.clone(),
+                ..Default::default()
+            },
+            &self.device,
+            self.dtype,
+        )?;
+
+        let audio = self.vae_decoder.tiled_decode(
+            &xt.transpose(1, 2)?,
+            params.vae_chunk_frames,
+            params.vae_chunk_overlap,
+        )?;
+        let audio = audio.squeeze(0)?;
+
+        Ok(AudioOutput {
+            audio,
+            sample_rate: self.vae_config.sampling_rate as u32,
+            latents: xt,
+        })
+    }
+}
+
+/// Load silence_latent from PyTorch `.pt` file → `(1, T, acoustic_dim)`.
+pub fn load_silence_latent(
+    path: &std::path::Path,
+    acoustic_dim: usize,
+    dtype: DType,
+    device: &Device,
+) -> Result<Tensor> {
+    let file = std::fs::File::open(path)?;
+    let mut archive = zip::ZipArchive::new(std::io::BufReader::new(file))?;
+    let mut raw = Vec::new();
+    std::io::Read::read_to_end(&mut archive.by_name("silence_latent/data/0")?, &mut raw)?;
+    let n_elements = raw.len() / 4;
+    let t_len = n_elements / acoustic_dim;
+    let cpu = Device::Cpu;
+    Ok(
+        Tensor::from_raw_buffer(&raw, DType::F32, &[1, acoustic_dim, t_len], &cpu)?
+            .transpose(1, 2)?
+            .contiguous()?
+            .to_dtype(dtype)?
+            .to_device(device)?,
+    )
+}

crates/neuron/src/harness/arch/acestep/qwen3.rs

@@ -0,0 +1,391 @@
+use candle_core::{DType, Device, Module, Result, Tensor};
+use candle_nn::{Activation, VarBuilder, kv_cache::ConcatKvCache};
+use candle_transformers::{
+    models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias},
+    utils::repeat_kv,
+};
+use std::sync::Arc;
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Config {
+    pub vocab_size: usize,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub head_dim: usize,
+    pub attention_bias: bool,
+    pub num_key_value_heads: usize,
+    pub max_position_embeddings: usize,
+    pub sliding_window: Option<usize>,
+    pub max_window_layers: usize,
+    pub tie_word_embeddings: bool,
+    pub rope_theta: f64,
+    pub rms_norm_eps: f64,
+    pub use_sliding_window: bool,
+    pub hidden_act: Activation,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Qwen3RotaryEmbedding {
+    sin: Tensor,
+    cos: Tensor,
+}
+
+impl Qwen3RotaryEmbedding {
+    pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result<Self> {
+        let dim = cfg.head_dim;
+        let max_seq_len = cfg.max_position_embeddings;
+        let inv_freq: Vec<_> = (0..dim)
+            .step_by(2)
+            .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32)
+            .collect();
+        let inv_freq_len = inv_freq.len();
+        let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?;
+        let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
+            .to_dtype(DType::F32)?
+            .reshape((max_seq_len, 1))?;
+        let freqs = t.matmul(&inv_freq)?;
+        Ok(Self {
+            sin: freqs.sin()?.to_dtype(dtype)?,
+            cos: freqs.cos()?.to_dtype(dtype)?,
+        })
+    }
+
+    /// Apply RoPE (q, k shape: B x H x L x D)
+    pub(crate) fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
+        let (_, _, seq_len, _) = q.dims4()?;
+        let cos = self.cos.narrow(0, offset, seq_len)?;
+        let sin = self.sin.narrow(0, offset, seq_len)?;
+        let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?;
+        let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?;
+        Ok((q_embed, k_embed))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Qwen3MLP {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl Qwen3MLP {
+    pub(crate) fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+        Ok(Self {
+            gate_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("gate_proj"))?,
+            up_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("up_proj"))?,
+            down_proj: linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj"))?,
+            act_fn: cfg.hidden_act,
+        })
+    }
+}
+
+impl Module for Qwen3MLP {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let rhs = x.apply(&self.up_proj)?;
+        (lhs * rhs)?.apply(&self.down_proj)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Qwen3Attention {
+    // projections
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    // norms
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    // hyper params
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+    // utils
+    rotary_emb: Arc<Qwen3RotaryEmbedding>,
+    kv_cache: ConcatKvCache,
+}
+
+impl Qwen3Attention {
+    pub(crate) fn new(
+        cfg: &Config,
+        rotary_emb: Arc<Qwen3RotaryEmbedding>,
+        vb: VarBuilder,
+    ) -> Result<Self> {
+        if cfg.use_sliding_window {
+            candle_core::bail!("sliding window is not supported")
+        }
+
+        let head_dim = cfg.head_dim;
+        let num_heads = cfg.num_attention_heads;
+        let num_kv_heads = cfg.num_key_value_heads;
+        let num_kv_groups = num_heads / num_kv_heads;
+
+        let q_proj = linear_b(
+            cfg.hidden_size,
+            num_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("q_proj"),
+        )?;
+        let k_proj = linear_b(
+            cfg.hidden_size,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("k_proj"),
+        )?;
+        let v_proj = linear_b(
+            cfg.hidden_size,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("v_proj"),
+        )?;
+        let o_proj = linear_b(
+            num_heads * head_dim,
+            cfg.hidden_size,
+            cfg.attention_bias,
+            vb.pp("o_proj"),
+        )?;
+
+        let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?;
+        let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?;
+
+        // Necessary because the hidden_size in the config isn't always accurate
+        let hidden_size = head_dim * cfg.num_attention_heads;
+
+        // dim=2 because we concatenate along the sequence dimension
+        // For tensors of shape [batch, heads, seq, head_dim]
+        let kv_cache = ConcatKvCache::new(2);
+
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            q_norm,
+            k_norm,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+            rotary_emb,
+            kv_cache,
+        })
+    }
+
+    pub(crate) fn forward(
+        &mut self,
+        x: &Tensor,
+        attn_mask: Option<&Tensor>,
+        offset: usize,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        // 1. Proj
+        let q = self.q_proj.forward(x)?;
+        let k = self.k_proj.forward(x)?;
+        let v = self.v_proj.forward(x)?;
+
+        // 2. Reshape: (B, L, H, D) -> (B, H, L, D)
+        let q = q
+            .reshape((b, l, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // 3. Per‑head RMSNorm
+        let q_flat = q.flatten(0, 2)?; // (B*H, L, D) -> (BHL, D) after transpose later
+        let k_flat = k.flatten(0, 2)?;
+        let q_flat = self.q_norm.forward(&q_flat)?;
+        let k_flat = self.k_norm.forward(&k_flat)?;
+        let q = q_flat.reshape((b, self.num_heads, l, self.head_dim))?;
+        let k = k_flat.reshape((b, self.num_kv_heads, l, self.head_dim))?;
+
+        // 4. RoPE
+        let (q, k) = self.rotary_emb.apply(&q, &k, offset)?;
+
+        // 5. Accumulate KV cache
+        let (k, v) = self.kv_cache.append(&k, &v)?;
+
+        // 6. GQA repeat_kv
+        let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
+
+        // 7. Attention score
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
+        if let Some(m) = attn_mask {
+            scores = scores.broadcast_add(m)?;
+        }
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let ctx = probs.matmul(&v)?; // (B, H, L, D)
+
+        // 8. Output proj
+        ctx.transpose(1, 2)?
+            .reshape((b, l, self.hidden_size))?
+            .apply(&self.o_proj)
+    }
+
+    pub(crate) fn clear_kv_cache(&mut self) {
+        self.kv_cache.reset();
+    }
+}
+
+#[derive(Debug, Clone)]
+struct DecoderLayer {
+    self_attn: Qwen3Attention,
+    mlp: Qwen3MLP,
+    ln1: RmsNorm,
+    ln2: RmsNorm,
+}
+
+impl DecoderLayer {
+    fn new(cfg: &Config, rotary: Arc<Qwen3RotaryEmbedding>, vb: VarBuilder) -> Result<Self> {
+        let self_attn = Qwen3Attention::new(cfg, rotary, vb.pp("self_attn"))?;
+        let mlp = Qwen3MLP::new(cfg, vb.pp("mlp"))?;
+        let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
+        let ln2 = RmsNorm::new(
+            cfg.hidden_size,
+            cfg.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            ln1,
+            ln2,
+        })
+    }
+
+    fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result<Tensor> {
+        let h = self.ln1.forward(x)?;
+        let h = self.self_attn.forward(&h, mask, offset)?;
+        let x = (x + h)?;
+        let h2 = self.ln2.forward(&x)?;
+        let h2 = h2.apply(&self.mlp)?;
+        x + h2
+    }
+
+    fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache();
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Model {
+    embed_tokens: candle_nn::Embedding,
+    layers: Vec<DecoderLayer>,
+    norm: RmsNorm,
+    device: Device,
+    dtype: DType,
+}
+
+impl Model {
+    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+        let embed_tokens =
+            candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?;
+        let rotary = Arc::new(Qwen3RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?);
+        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+        let vb_l = vb.pp("model.layers");
+        for i in 0..cfg.num_hidden_layers {
+            layers.push(DecoderLayer::new(cfg, rotary.clone(), vb_l.pp(i))?);
+        }
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm: RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?,
+            device: vb.device().clone(),
+            dtype: vb.dtype(),
+        })
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        for l in &mut self.layers {
+            l.clear_kv_cache();
+        }
+    }
+
+    fn causal_mask(
+        &self,
+        b: usize,
+        tgt: usize,
+        offset: usize,
+        sw: Option<usize>,
+    ) -> Result<Tensor> {
+        let minf = f32::NEG_INFINITY;
+        let mask: Vec<_> = (0..tgt)
+            .flat_map(|i| {
+                (0..(tgt + offset)).map(move |j| {
+                    let past_ok = j <= i + offset;
+                    let sw_ok = match sw {
+                        Some(w) => (i + offset) as i64 - j as i64 <= w as i64,
+                        None => true,
+                    };
+                    if past_ok && sw_ok { 0. } else { minf }
+                })
+            })
+            .collect();
+        Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype)
+    }
+
+    pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result<Tensor> {
+        let (b, l) = input.dims2()?;
+        let mut h = self.embed_tokens.forward(input)?;
+
+        let causal = if l == 1 {
+            None
+        } else {
+            Some(self.causal_mask(b, l, offset, None)?)
+        };
+
+        for layer in &mut self.layers {
+            h = layer.forward(&h, causal.as_ref(), offset)?;
+        }
+        self.norm.forward(&h)
+    }
+
+    /// Embed token IDs without running through transformer layers.
+    /// Returns raw token embeddings of shape `(B, L, hidden_size)`.
+    pub fn embed_tokens(&self, input: &Tensor) -> Result<Tensor> {
+        self.embed_tokens.forward(input)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ModelForCausalLM {
+    base: Model,
+    lm_head: Linear,
+}
+
+impl ModelForCausalLM {
+    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+        let base = Model::new(cfg, vb.clone())?;
+        let lm_head = if cfg.tie_word_embeddings {
+            Linear::from_weights(base.embed_tokens.embeddings().clone(), None)
+        } else {
+            linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?
+        };
+        Ok(Self { base, lm_head })
+    }
+
+    pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result<Tensor> {
+        let (_, l) = input.dims2()?;
+        self.base
+            .forward(input, offset)?
+            .narrow(1, l - 1, 1)?
+            .apply(&self.lm_head)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.base.clear_kv_cache();
+    }
+}

crates/neuron/src/harness/arch/acestep/sampling.rs

@@ -0,0 +1,466 @@
+//! Flow matching sampler for ACE-Step diffusion model.
+//!
+//! Implements the Euler ODE solver for rectified flow / flow matching,
+//! with Adaptive Prompt Guidance (APG) support.
+
+use candle_core::{Result, Tensor};
+
+/// All unique timesteps from shift=1,2,3 with fix_nfe=8.
+/// Turbo models were distilled on these exact values — custom timesteps are
+/// mapped to the nearest entry in this table.
+const VALID_TIMESTEPS: [f64; 20] = [
+    1.0,
+    0.9545454545454546,
+    0.9333333333333333,
+    0.9,
+    0.875,
+    0.8571428571428571,
+    0.8333333333333334,
+    0.7692307692307693,
+    0.75,
+    0.6666666666666666,
+    0.6428571428571429,
+    0.625,
+    0.5454545454545454,
+    0.5,
+    0.4,
+    0.375,
+    0.3,
+    0.25,
+    0.2222222222222222,
+    0.125,
+];
+
+/// Generate a linear timestep schedule from 1.0 to 0.0.
+///
+/// Optionally applies a time shift transformation:
+/// `t_shifted = shift * t / (1 + (shift - 1) * t)`
+pub fn get_schedule(num_steps: usize, shift: f64) -> Vec<f64> {
+    if num_steps == 0 {
+        return vec![0.0];
+    }
+    let mut schedule = Vec::with_capacity(num_steps + 1);
+    for i in 0..=num_steps {
+        let t = 1.0 - (i as f64 / num_steps as f64);
+        let t = if (shift - 1.0).abs() > 1e-6 {
+            shift * t / (1.0 + (shift - 1.0) * t)
+        } else {
+            t
+        };
+        schedule.push(t);
+    }
+    schedule
+}
+
+/// Pre-defined discrete timestep schedules for turbo models.
+/// These are the only valid schedules — turbo models were distilled on these
+/// exact timesteps and will produce poor results with continuous schedules.
+///
+/// Each schedule has `fix_nfe` steps (default 8) plus an implicit final t=0.
+pub fn get_turbo_schedule(shift: f64) -> Vec<f64> {
+    if (shift - 1.0).abs() < 0.5 {
+        // shift ≈ 1
+        vec![1.0, 0.875, 0.75, 0.625, 0.5, 0.375, 0.25, 0.125]
+    } else if (shift - 2.0).abs() < 0.5 {
+        // shift ≈ 2
+        vec![
+            1.0,
+            0.9333333333333333,
+            0.8571428571428571,
+            0.7692307692307693,
+            0.6666666666666666,
+            0.5454545454545454,
+            0.4,
+            0.2222222222222222,
+        ]
+    } else {
+        // shift ≈ 3 (default for turbo)
+        vec![
+            1.0,
+            0.9545454545454546,
+            0.9,
+            0.8333333333333334,
+            0.75,
+            0.6428571428571429,
+            0.5,
+            0.3,
+        ]
+    }
+}
+
+/// Build a custom turbo schedule from user-provided timestep values.
+///
+/// Each value is mapped to the nearest entry in `VALID_TIMESTEPS`. Trailing
+/// zeros are stripped, and the resulting schedule must have 1–20 entries.
+pub fn get_custom_turbo_schedule(timesteps: &[f64]) -> Result<Vec<f64>> {
+    let end = timesteps
+        .iter()
+        .rposition(|&t| t != 0.0)
+        .map(|i| i + 1)
+        .unwrap_or(0);
+    let timesteps = &timesteps[..end];
+    if timesteps.is_empty() || timesteps.len() > 20 {
+        candle_core::bail!(
+            "custom timesteps must have 1-20 non-zero values, got {}",
+            timesteps.len()
+        );
+    }
+    for &t in timesteps {
+        if !t.is_finite() {
+            candle_core::bail!("custom timesteps contain non-finite value: {t}");
+        }
+    }
+    Ok(timesteps
+        .iter()
+        .map(|&t| {
+            *VALID_TIMESTEPS
+                .iter()
+                .min_by(|&&a, &&b| {
+                    (a - t)
+                        .abs()
+                        .partial_cmp(&(b - t).abs())
+                        .unwrap_or(std::cmp::Ordering::Equal)
+                })
+                .unwrap()
+        })
+        .collect())
+}
+
+/// Get x0 (clean sample) from noisy sample and velocity prediction.
+/// Flow matching: x_t = t * noise + (1 - t) * x_0, v = noise - x_0
+/// Therefore: x_0 = x_t - v * t
+pub fn get_x0_from_noise(xt: &Tensor, vt: &Tensor, t: f64) -> Result<Tensor> {
+    xt - (vt * t)?
+}
+
+/// Euler ODE step for flow matching.
+///
+/// Given current sample `x_t` and velocity prediction `v_t`:
+/// `x_{t-dt} = x_t - v_t * dt`
+pub fn euler_step(x: &Tensor, v: &Tensor, dt: f64) -> Result<Tensor> {
+    x - (v * dt)?
+}
+
+/// SDE step: predict clean sample, then re-add noise at the next timestep.
+///
+/// `x_0 = get_x0_from_noise(x_t, v_t, t_curr)`
+/// `x_{t_next} = (1 - t_next) * x_0 + t_next * noise`
+pub fn sde_step(xt: &Tensor, vt: &Tensor, t_curr: f64, t_next: f64) -> Result<Tensor> {
+    let x0 = get_x0_from_noise(xt, vt, t_curr)?;
+    let noise = Tensor::randn_like(&x0, 0., 1.)?;
+    let clean_part = (&x0 * (1.0 - t_next))?;
+    let noise_part = (noise * t_next)?;
+    clean_part + noise_part
+}
+
+/// Inference method for the denoising loop.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum InferMethod {
+    /// Ordinary Differential Equation — deterministic Euler integration.
+    Ode,
+    /// Stochastic Differential Equation — predicts x0 then re-noises.
+    Sde,
+}
+
+/// Momentum buffer for APG (Adaptive Prompt Guidance).
+pub struct MomentumBuffer {
+    momentum: f64,
+    running_average: Option<Tensor>,
+}
+
+impl MomentumBuffer {
+    pub fn new(momentum: f64) -> Self {
+        Self {
+            momentum,
+            running_average: None,
+        }
+    }
+
+    /// Update running average: `running_avg = update_value + momentum * running_avg`
+    pub fn update(&mut self, update_value: &Tensor) -> Result<()> {
+        let new_average = match &self.running_average {
+            Some(ra) => (ra * self.momentum)?,
+            None => Tensor::zeros_like(update_value)?,
+        };
+        self.running_average = Some((update_value + new_average)?);
+        Ok(())
+    }
+
+    pub fn get(&self) -> Option<&Tensor> {
+        self.running_average.as_ref()
+    }
+}
+
+/// Project v0 onto v1 and decompose into parallel and orthogonal components.
+///
+/// Operates along `dim` (typically dim=1 for sequence dimension).
+/// Uses F64 precision internally for numerical stability (matching Python).
+/// On Metal, computation is moved to CPU to avoid MPS numeric issues.
+fn project(v0: &Tensor, v1: &Tensor, dim: usize) -> Result<(Tensor, Tensor)> {
+    let orig_device = v0.device().clone();
+    let orig_dtype = v0.dtype();
+
+    // Use CPU for Metal (MPS numeric issues); stay on device for CUDA.
+    let compute_device = if orig_device.is_metal() {
+        candle_core::Device::Cpu
+    } else {
+        orig_device.clone()
+    };
+    let v0 = v0
+        .to_device(&compute_device)?
+        .to_dtype(candle_core::DType::F64)?;
+    let v1 = v1
+        .to_device(&compute_device)?
+        .to_dtype(candle_core::DType::F64)?;
+
+    // L2 normalize v1 along dim
+    let v1_norm = v1.sqr()?.sum_keepdim(dim)?.sqrt()?;
+    let v1_normalized = v1.broadcast_div(&(v1_norm + 1e-8)?)?;
+
+    // Project: v0_parallel = (v0 · v1_hat) * v1_hat
+    let dot = (&v0 * &v1_normalized)?.sum_keepdim(dim)?;
+    let v0_parallel = dot.broadcast_mul(&v1_normalized)?;
+    let v0_orthogonal = (&v0 - &v0_parallel)?;
+
+    Ok((
+        v0_parallel.to_dtype(orig_dtype)?.to_device(&orig_device)?,
+        v0_orthogonal
+            .to_dtype(orig_dtype)?
+            .to_device(&orig_device)?,
+    ))
+}
+
+/// Adaptive Prompt Guidance (APG).
+///
+/// Decomposes the cond-uncond difference into parallel/orthogonal components
+/// relative to the conditional prediction, applies norm thresholding, and
+/// uses momentum for temporal smoothing.
+///
+/// Reference: ACE-Step's `apg_forward` in `apg_guidance.py`.
+pub fn apg_forward(
+    pred_cond: &Tensor,
+    pred_uncond: &Tensor,
+    guidance_scale: f64,
+    momentum_buffer: &mut MomentumBuffer,
+    norm_threshold: f64,
+) -> Result<Tensor> {
+    // dims=[1] in Python — operates along the sequence dimension
+    let dim = 1usize;
+    let orig_device = pred_cond.device().clone();
+    let orig_dtype = pred_cond.dtype();
+
+    // On Metal, move to CPU for numerical stability (Python does this for MPS).
+    // On CUDA/CPU, stay on the original device to avoid transfer overhead.
+    let compute_device = if orig_device.is_metal() {
+        candle_core::Device::Cpu
+    } else {
+        orig_device.clone()
+    };
+    let pred_cond_cpu = pred_cond.to_device(&compute_device)?;
+    let pred_uncond_cpu = pred_uncond.to_device(&compute_device)?;
+
+    let diff = (&pred_cond_cpu - &pred_uncond_cpu)?;
+
+    // Apply momentum smoothing
+    momentum_buffer.update(&diff)?;
+    let diff = momentum_buffer.get().unwrap().clone();
+
+    // Norm thresholding: clamp L2 norm along dim to prevent oversaturation
+    let diff = if norm_threshold > 0.0 {
+        let diff_norm = diff.sqr()?.sum_keepdim(dim)?.sqrt()?;
+        let ones = Tensor::ones_like(&diff_norm)?;
+        let scale = (ones * norm_threshold)?.broadcast_div(&(&diff_norm + 1e-8)?)?;
+        let scale = scale.minimum(&Tensor::ones_like(&scale)?)?;
+        diff.broadcast_mul(&scale)?
+    } else {
+        diff
+    };
+
+    // Decompose into parallel and orthogonal components
+    let (_diff_parallel, diff_orthogonal) = project(&diff, &pred_cond_cpu, dim)?;
+
+    // eta=0.0 by default: use only orthogonal component
+    let result = (pred_cond_cpu + diff_orthogonal * (guidance_scale - 1.0))?;
+    result.to_dtype(orig_dtype)?.to_device(&orig_device)
+}
+
+/// Differential Correction in Wavelet domain ("double" mode, Haar), which
+/// upstream applies by default after every turbo step: split the latent and
+/// the predicted clean sample into low and high bands along time, and push
+/// each band of the latent away from the clean sample's. The low band's
+/// strength decays with `t`, the high band's grows as `t → 0`.
+#[derive(Debug, Clone, Copy)]
+pub struct Dcw {
+    pub low_scaler: f64,
+    pub high_scaler: f64,
+}
+
+impl Default for Dcw {
+    /// Upstream's defaults for turbo (`dcw_scaler`, `dcw_high_scaler`).
+    fn default() -> Self {
+        Self {
+            low_scaler: 0.05,
+            high_scaler: 0.02,
+        }
+    }
+}
+
+impl Dcw {
+    /// Correct `x_next` (after the sampler step at `t_curr`) against
+    /// `denoised = x_t - v·t_curr`. Shapes are `(B, T, C)`.
+    pub fn apply(&self, x_next: &Tensor, denoised: &Tensor, t_curr: f64) -> Result<Tensor> {
+        let low = t_curr * self.low_scaler;
+        let high = (1.0 - t_curr) * self.high_scaler;
+        if low == 0.0 && high == 0.0 {
+            return Ok(x_next.clone());
+        }
+        let t_out = x_next.dim(1)?;
+        let (xl, xh) = haar_dwt_t(x_next)?;
+        let (yl, yh) = haar_dwt_t(denoised)?;
+        let xl = (&xl + ((&xl - &yl)? * low)?)?;
+        let xh = (&xh + ((&xh - &yh)? * high)?)?;
+        haar_idwt_t(&xl, &xh, t_out)
+    }
+}
+
+/// One-level Haar DWT along dim 1 of `(B, T, C)`, zero-padding an odd `T`
+/// as `pytorch_wavelets` does in `mode="zero"`. Returns `(low, high)`, each
+/// `(B, ceil(T/2), C)`.
+fn haar_dwt_t(x: &Tensor) -> Result<(Tensor, Tensor)> {
+    let (b, t, c) = x.dims3()?;
+    let x = if t % 2 == 1 {
+        let pad = Tensor::zeros((b, 1, c), x.dtype(), x.device())?;
+        Tensor::cat(&[x, &pad], 1)?
+    } else {
+        x.clone()
+    };
+    let half = x.dim(1)? / 2;
+    let pairs = x.reshape((b, half, 2, c))?;
+    let even = pairs.narrow(2, 0, 1)?.squeeze(2)?;
+    let odd = pairs.narrow(2, 1, 1)?.squeeze(2)?;
+    let k = std::f64::consts::FRAC_1_SQRT_2;
+    Ok((((&even + &odd)? * k)?, ((&even - &odd)? * k)?))
+}
+
+/// Inverse of [`haar_dwt_t`], trimmed back to `t_out` frames.
+fn haar_idwt_t(low: &Tensor, high: &Tensor, t_out: usize) -> Result<Tensor> {
+    let (b, half, c) = low.dims3()?;
+    let k = std::f64::consts::FRAC_1_SQRT_2;
+    let even = ((low + high)? * k)?;
+    let odd = ((low - high)? * k)?;
+    let x = Tensor::stack(&[&even, &odd], 2)?.reshape((b, half * 2, c))?;
+    x.narrow(1, 0, t_out)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_get_custom_turbo_schedule_basic() -> Result<()> {
+        // Exact valid timesteps should pass through unchanged
+        let ts = vec![1.0, 0.875, 0.75, 0.5, 0.25];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result, ts);
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_nearest_mapping() -> Result<()> {
+        // Values should be mapped to nearest valid timestep
+        let ts = vec![0.99, 0.51, 0.13];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result, vec![1.0, 0.5, 0.125]);
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_strips_trailing_zeros() -> Result<()> {
+        let ts = vec![1.0, 0.5, 0.0, 0.0];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_empty_after_strip() {
+        let ts = vec![0.0, 0.0];
+        assert!(get_custom_turbo_schedule(&ts).is_err());
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_too_long() {
+        let ts = vec![1.0; 21];
+        assert!(get_custom_turbo_schedule(&ts).is_err());
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_single() -> Result<()> {
+        let ts = vec![0.87];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result, vec![0.875]);
+        Ok(())
+    }
+
+    #[test]
+    fn test_turbo_schedule_shift1_matches_custom() -> Result<()> {
+        // The shift=1 hardcoded schedule should match custom input of same values
+        let hardcoded = get_turbo_schedule(1.0);
+        let custom = get_custom_turbo_schedule(&hardcoded)?;
+        assert_eq!(hardcoded, custom);
+        Ok(())
+    }
+
+    #[test]
+    fn test_turbo_schedule_shift3_matches_custom() -> Result<()> {
+        let hardcoded = get_turbo_schedule(3.0);
+        let custom = get_custom_turbo_schedule(&hardcoded)?;
+        assert_eq!(hardcoded, custom);
+        Ok(())
+    }
+
+    #[test]
+    fn haar_round_trips_and_dcw_matches_the_reference_formula() {
+        use candle_core::Device;
+        let d = Device::Cpu;
+        // Odd T exercises the zero pad and the trim.
+        let x = Tensor::arange(0f32, 21.0, &d)
+            .unwrap()
+            .reshape((1, 7, 3))
+            .unwrap();
+        let (l, h) = super::haar_dwt_t(&x).unwrap();
+        assert_eq!(l.dims3().unwrap(), (1, 4, 3));
+        let back = super::haar_idwt_t(&l, &h, 7).unwrap();
+        let diff = (back - &x)
+            .unwrap()
+            .abs()
+            .unwrap()
+            .max_all()
+            .unwrap()
+            .to_scalar::<f32>()
+            .unwrap();
+        assert!(diff < 1e-5, "{diff}");
+
+        // Against the band formula directly, for a 2-frame signal:
+        // L = (a+b)/√2, H = (a-b)/√2; corrected bands, then inverted.
+        let xn = Tensor::new(&[[[3f32], [1.0]]], &d).unwrap();
+        let dn = Tensor::new(&[[[1f32], [1.0]]], &d).unwrap();
+        let t = 0.5;
+        let out = Dcw::default()
+            .apply(&xn, &dn, t)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1::<f32>()
+            .unwrap();
+        let r = std::f64::consts::FRAC_1_SQRT_2;
+        let (xl, xh, yl, yh) = (4.0 * r, 2.0 * r, 2.0 * r, 0.0);
+        let (s_lo, s_hi) = (t * 0.05, (1.0 - t) * 0.02);
+        let (cl, ch) = (xl + s_lo * (xl - yl), xh + s_hi * (xh - yh));
+        let want = [((cl + ch) * r) as f32, ((cl - ch) * r) as f32];
+        assert!(
+            (out[0] - want[0]).abs() < 1e-5 && (out[1] - want[1]).abs() < 1e-5,
+            "{out:?} vs {want:?}"
+        );
+    }
+}

crates/neuron/src/harness/arch/acestep/tokenizer.rs

@@ -0,0 +1,656 @@
+//! Audio tokenizer and detokenizer for ACE-Step.
+//!
+//! Implements Finite Scalar Quantization (FSQ) based audio tokenization.
+//! These components are used in the cover/audio-code pipeline (LM + DiT hybrid
+//! mode) but are **not exercised** by the Text2Music path.
+//!
+//! **Validation status** (against Python `vector_quantize_pytorch.ResidualFSQ`):
+//! - `AudioTokenDetokenizer`: bit-exact (diff_max < 0.00001)
+//! - `ResidualFSQ` (quantizer): FSQ bound/round matches; small rounding
+//!   boundary differences (diff_max ~0.018) are expected for discrete
+//!   quantization due to float precision at level boundaries.
+//! - `AttentionPooler` / `AudioTokenizer`: structurally correct, minor
+//!   float precision differences propagate through FSQ rounding.
+
+use candle_core::{D, DType, Module, Result, Tensor};
+use candle_nn::{Activation, VarBuilder};
+use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear, linear_no_bias};
+use candle_transformers::utils::repeat_kv;
+use std::sync::Arc;
+
+use super::AceStepConfig;
+use super::dit::RotaryEmbedding;
+
+// ---------------------------------------------------------------------------
+// Encoder MLP (same SiLU-gated pattern as Qwen3MLP)
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+#[allow(clippy::upper_case_acronyms)]
+struct MLP {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl MLP {
+    fn new(
+        hidden_size: usize,
+        intermediate_size: usize,
+        act: Activation,
+        vb: VarBuilder,
+    ) -> Result<Self> {
+        Ok(Self {
+            gate_proj: linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
+            up_proj: linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?,
+            down_proj: linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?,
+            act_fn: act,
+        })
+    }
+}
+
+impl Module for MLP {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let rhs = x.apply(&self.up_proj)?;
+        (lhs * rhs)?.apply(&self.down_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Encoder Attention (bidirectional, GQA, with RoPE)
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+struct TokenizerAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+}
+
+impl TokenizerAttention {
+    fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        let num_heads = cfg.encoder_num_attention_heads();
+        let num_kv_heads = cfg.encoder_num_key_value_heads();
+        let head_dim = cfg.head_dim;
+        let hidden_size = num_heads * head_dim;
+
+        Ok(Self {
+            q_proj: linear_no_bias(enc_hs, num_heads * head_dim, vb.pp("q_proj"))?,
+            k_proj: linear_no_bias(enc_hs, num_kv_heads * head_dim, vb.pp("k_proj"))?,
+            v_proj: linear_no_bias(enc_hs, num_kv_heads * head_dim, vb.pp("v_proj"))?,
+            o_proj: linear_no_bias(num_heads * head_dim, enc_hs, vb.pp("o_proj"))?,
+            q_norm: RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?,
+            k_norm: RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups: num_heads / num_kv_heads,
+            head_dim,
+            hidden_size,
+        })
+    }
+
+    fn forward(
+        &self,
+        x: &Tensor,
+        attn_mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        let q = self.q_proj.forward(x)?;
+        let k = self.k_proj.forward(x)?;
+        let v = self.v_proj.forward(x)?;
+
+        let q = q
+            .reshape((b, l, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // Per-head RMSNorm
+        let q = self.q_norm.forward(&q.flatten(0, 1)?)?.reshape(q.shape())?;
+        let k = self.k_norm.forward(&k.flatten(0, 1)?)?.reshape(k.shape())?;
+
+        // RoPE
+        let (q, k) = rotary.apply(&q, &k, 0)?;
+
+        // GQA
+        let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
+
+        // Attention
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
+        if let Some(m) = attn_mask {
+            scores = scores.broadcast_add(m)?;
+        }
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let ctx = probs.matmul(&v)?;
+
+        ctx.transpose(1, 2)?
+            .reshape((b, l, self.hidden_size))?
+            .apply(&self.o_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Encoder Layer
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+struct EncoderLayer {
+    self_attn: TokenizerAttention,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+    mlp: MLP,
+}
+
+impl EncoderLayer {
+    fn new(cfg: &AceStepConfig, _layer_idx: usize, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        Ok(Self {
+            self_attn: TokenizerAttention::new(cfg, vb.pp("self_attn"))?,
+            input_layernorm: RmsNorm::new(enc_hs, cfg.rms_norm_eps, vb.pp("input_layernorm"))?,
+            post_attention_layernorm: RmsNorm::new(
+                enc_hs,
+                cfg.rms_norm_eps,
+                vb.pp("post_attention_layernorm"),
+            )?,
+            mlp: MLP::new(
+                enc_hs,
+                cfg.encoder_intermediate_size(),
+                cfg.hidden_act,
+                vb.pp("mlp"),
+            )?,
+        })
+    }
+
+    fn forward(
+        &self,
+        x: &Tensor,
+        attn_mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        let residual = x.clone();
+        let h = self.input_layernorm.forward(x)?;
+        let h = self.self_attn.forward(&h, attn_mask, rotary)?;
+        let x = (residual + h)?;
+
+        let residual = x.clone();
+        let h = self.post_attention_layernorm.forward(&x)?;
+        let h = h.apply(&self.mlp)?;
+        residual + h
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AttentionPooler - CLS-token based attention pooling
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+pub struct AttentionPooler {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    special_token: Tensor,
+    rotary_emb: Arc<RotaryEmbedding>,
+}
+
+impl AttentionPooler {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        let mut layers = Vec::with_capacity(cfg.num_attention_pooler_hidden_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..cfg.num_attention_pooler_hidden_layers {
+            layers.push(EncoderLayer::new(cfg, i, vb_l.pp(i))?);
+        }
+        let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?);
+        Ok(Self {
+            embed_tokens: linear(enc_hs, enc_hs, vb.pp("embed_tokens"))?,
+            layers,
+            norm: RmsNorm::new(enc_hs, cfg.rms_norm_eps, vb.pp("norm"))?,
+            special_token: vb.get((1, 1, enc_hs), "special_token")?,
+            rotary_emb,
+        })
+    }
+
+    /// Pool patches into sequence-level representations.
+    /// Input: (B, T, pool_window_size, D) → Output: (B, T, D)
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let (b, t, p, _d) = x.dims4()?;
+        let x = x.apply(&self.embed_tokens)?;
+
+        // Prepend CLS token to each patch group
+        let cls = self
+            .special_token
+            .broadcast_as((b * t, 1, x.dim(D::Minus1)?))?;
+        let x = x.reshape((b * t, p, ()))?;
+        let x = Tensor::cat(&[&cls, &x], 1)?;
+
+        let mut hidden = x;
+        for layer in &self.layers {
+            hidden = layer.forward(&hidden, None, &self.rotary_emb)?;
+        }
+        let hidden = self.norm.forward(&hidden)?;
+
+        // Extract CLS output (position 0)
+        hidden.narrow(1, 0, 1)?.squeeze(1)?.reshape((b, t, ()))
+    }
+}
+
+// ---------------------------------------------------------------------------
+// ResidualFSQ - Finite Scalar Quantization
+// ---------------------------------------------------------------------------
+
+/// Residual Finite Scalar Quantization.
+///
+/// Matches `vector_quantize_pytorch.ResidualFSQ` with a single top-level
+/// `project_in` (dim → codebook_dim) and `project_out` (codebook_dim → dim).
+/// Each FSQ layer independently quantizes each codebook dimension to its
+/// corresponding number of levels using `tanh` bounding and rounding.
+#[derive(Debug, Clone)]
+pub struct ResidualFSQ {
+    project_in: Linear,
+    project_out: Linear,
+    /// Precomputed: (levels - 1) * (1 + eps) / 2 for each codebook dim
+    half_l: Vec<f32>,
+    /// Precomputed: levels // 2
+    half_width: Vec<f32>,
+    /// Precomputed: 0.5 for even levels, 0.0 for odd
+    offset: Vec<f32>,
+    /// Precomputed: atanh(offset / half_l) shift
+    shift: Vec<f32>,
+    /// Precomputed basis for index calculation: [1, L0, L0*L1, ...]
+    basis: Vec<i64>,
+    /// Quantization levels per codebook dimension (e.g. [8, 8, 8, 5, 5, 5]).
+    levels: Vec<usize>,
+    /// Model working dtype (inferred from VarBuilder at construction time).
+    model_dtype: DType,
+}
+
+impl ResidualFSQ {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let model_dtype = vb.dtype();
+        let dim = cfg.fsq_dim;
+        let levels = &cfg.fsq_input_levels;
+        let codebook_dim = levels.len();
+
+        // Single project_in/project_out at the ResidualFSQ level
+        let project_in = linear(dim, codebook_dim, vb.pp("project_in"))?;
+        let project_out = linear(codebook_dim, dim, vb.pp("project_out"))?;
+
+        let eps = 1e-3f32;
+        let half_l: Vec<f32> = levels
+            .iter()
+            .map(|&l| (l as f32 - 1.0) * (1.0 + eps) / 2.0)
+            .collect();
+        let half_width: Vec<f32> = levels.iter().map(|&l| (l / 2) as f32).collect();
+        let offset: Vec<f32> = levels
+            .iter()
+            .map(|&l| if l % 2 == 0 { 0.5 } else { 0.0 })
+            .collect();
+        let shift: Vec<f32> = half_l
+            .iter()
+            .zip(&offset)
+            .map(|(&hl, &o)| (o / hl).atanh())
+            .collect();
+
+        let mut basis = vec![1i64];
+        for i in 0..levels.len() - 1 {
+            basis.push(basis[i] * levels[i] as i64);
+        }
+
+        Ok(Self {
+            project_in,
+            project_out,
+            half_l,
+            half_width,
+            offset,
+            shift,
+            basis,
+            levels: levels.to_vec(),
+            model_dtype,
+        })
+    }
+
+    /// FSQ bound + round: `round(tanh(z + shift) * half_l - offset) / half_width`
+    fn fsq_quantize(&self, z: &Tensor) -> Result<Tensor> {
+        let dev = z.device();
+        let half_l = Tensor::new(&self.half_l[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let offset = Tensor::new(&self.offset[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let shift = Tensor::new(&self.shift[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let half_width = Tensor::new(&self.half_width[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+
+        // bound: tanh(z + shift) * half_l - offset, then round, then / half_width
+        let bounded = z
+            .broadcast_add(&shift)?
+            .tanh()?
+            .broadcast_mul(&half_l)?
+            .broadcast_sub(&offset)?;
+        let rounded = bounded.round()?;
+        rounded.broadcast_div(&half_width)
+    }
+
+    /// Quantize continuous representations into discrete tokens.
+    /// Input: `(B, N, dim)`. Returns `(quantized_out, indices)` where
+    /// `quantized_out` has the same shape as input and `indices` is `(B, N, num_quantizers)`.
+    pub fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
+        // project_in: (B, N, dim) -> (B, N, codebook_dim)
+        let projected = x.apply(&self.project_in)?;
+
+        // FSQ quantize
+        let codes = self.fsq_quantize(&projected)?; // (B, N, codebook_dim)
+
+        // Compute indices from codes: scale_and_shift then dot with basis
+        let dev = x.device();
+        let half_width = Tensor::new(&self.half_width[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let offset = Tensor::new(&self.offset[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let basis = Tensor::new(&self.basis[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        // scale_and_shift: codes * half_width + offset
+        let shifted = codes.broadcast_mul(&half_width)?.broadcast_add(&offset)?;
+        let indices = shifted
+            .to_dtype(DType::F32)?
+            .broadcast_mul(&basis.to_dtype(DType::F32)?)?
+            .sum(2)?
+            .round()?
+            .to_dtype(DType::I64)?
+            .unsqueeze(2)?; // (B, N, 1)
+
+        // project_out: (B, N, codebook_dim) -> (B, N, dim)
+        let quantized_out = codes.apply(&self.project_out)?;
+
+        Ok((quantized_out, indices))
+    }
+
+    /// Reconstruct from quantized codes (not indices).
+    /// Input: `(B, N, codebook_dim)` → Output: `(B, N, dim)`.
+    pub fn get_output_from_codes(&self, codes: &Tensor) -> Result<Tensor> {
+        codes.apply(&self.project_out)
+    }
+
+    /// Reconstruct from packed indices.
+    /// Input: `(B, N, 1)` → Output: `(B, N, dim)`.
+    ///
+    /// Inverse of the forward path: decompose flat index into per-level values
+    /// via basis, then invert scale_and_shift, then `project_out`.
+    pub fn get_output_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
+        let dev = indices.device();
+        let indices = indices.squeeze(D::Minus1)?.to_dtype(DType::F32)?;
+        let codebook_dim = self.levels.len();
+
+        // Decompose: level_val[i] = floor(index / basis[i]) % levels[i]
+        // F32 has enough precision for indices up to 64000 (max codebook size).
+        let mut code_dims = Vec::with_capacity(codebook_dim);
+        for i in 0..codebook_dim {
+            let b = self.basis[i] as f32;
+            let l = self.levels[i] as f32;
+            let dim_val = indices.affine(1.0 / b as f64, 0.0)?.floor()?;
+            // modulo: dim_val - floor(dim_val / l) * l
+            let dim_val = (&dim_val
+                - dim_val
+                    .affine(1.0 / l as f64, 0.0)?
+                    .floor()?
+                    .affine(l as f64, 0.0)?)?;
+            code_dims.push(dim_val);
+        }
+        let shifted = Tensor::stack(&code_dims, D::Minus1)?;
+
+        // Inverse scale_and_shift (Python convention):
+        // forward: shifted = codes * half_width + half_width  →  shifted ∈ [0, levels-1]
+        // inverse: codes = (shifted - half_width) / half_width
+        let hw = Tensor::new(&self.half_width[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let codes = shifted.broadcast_sub(&hw)?.broadcast_div(&hw)?;
+
+        codes.to_dtype(self.model_dtype)?.apply(&self.project_out)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::Device;
+
+    /// Build a minimal ResidualFSQ with identity projections for testing.
+    fn test_fsq() -> Result<ResidualFSQ> {
+        let dev = &Device::Cpu;
+        let levels = vec![8, 8, 8, 5, 5, 5];
+        let codebook_dim = levels.len();
+
+        // Identity-like project_in/project_out (codebook_dim x codebook_dim)
+        let eye = Tensor::eye(codebook_dim, DType::F32, dev)?;
+        let zeros = Tensor::zeros(codebook_dim, DType::F32, dev)?;
+
+        let eps = 1e-3f32;
+        let half_l: Vec<f32> = levels
+            .iter()
+            .map(|&l| (l as f32 - 1.0) * (1.0 + eps) / 2.0)
+            .collect();
+        let half_width: Vec<f32> = levels.iter().map(|&l| (l / 2) as f32).collect();
+        let offset: Vec<f32> = levels
+            .iter()
+            .map(|&l| if l % 2 == 0 { 0.5 } else { 0.0 })
+            .collect();
+        let shift: Vec<f32> = half_l
+            .iter()
+            .zip(&offset)
+            .map(|(&hl, &o)| (o / hl).atanh())
+            .collect();
+        let mut basis = vec![1i64];
+        for i in 0..levels.len() - 1 {
+            basis.push(basis[i] * levels[i] as i64);
+        }
+
+        Ok(ResidualFSQ {
+            project_in: Linear::from_weights(eye.clone(), Some(zeros.clone())),
+            project_out: Linear::from_weights(eye, Some(zeros)),
+            half_l,
+            half_width,
+            offset,
+            shift,
+            basis,
+            levels,
+            model_dtype: DType::F32,
+        })
+    }
+
+    #[test]
+    fn test_get_output_from_indices_roundtrip() -> Result<()> {
+        // Test that manually-packed Python-convention indices correctly
+        // reconstruct codes via get_output_from_indices.
+        let fsq = test_fsq()?;
+        let dev = &Device::Cpu;
+        // levels = [8,8,8,5,5,5], half_width = [4,4,4,2,2,2]
+        // basis = [1, 8, 64, 512, 2560, 12800]
+        //
+        // Choose codes: [0.25, -0.5, 0.75, 0.5, -1.0, 0.0]
+        // Python scale_and_shift: shifted = codes * half_width + half_width
+        //   = [0.25*4+4, -0.5*4+4, 0.75*4+4, 0.5*2+2, -1.0*2+2, 0.0*2+2]
+        //   = [5, 2, 7, 3, 0, 2]
+        // index = 5*1 + 2*8 + 7*64 + 3*512 + 0*2560 + 2*12800
+        //       = 5 + 16 + 448 + 1536 + 0 + 25600 = 27605
+        let expected_codes: Vec<f32> = vec![0.25, -0.5, 0.75, 0.5, -1.0, 0.0];
+        let index: i64 = 27605;
+
+        let indices = Tensor::new(&[[index]], dev)?.unsqueeze(2)?; // (1, 1, 1)
+        let reconstructed = fsq.get_output_from_indices(&indices)?;
+        let vals = reconstructed.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
+
+        for (i, (&got, &want)) in vals.iter().zip(expected_codes.iter()).enumerate() {
+            assert!((got - want).abs() < 1e-5, "dim {i}: got={got}, want={want}");
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_output_from_indices_known_values() -> Result<()> {
+        let fsq = test_fsq()?;
+        let dev = &Device::Cpu;
+
+        // Index 0 should decompose to level_vals = [0, 0, 0, 0, 0, 0]
+        // Python convention: codes = (level_vals - half_width) / half_width
+        // For levels [8,8,8,5,5,5]: half_width=[4,4,4,2,2,2]
+        // codes = [(0-4)/4, (0-4)/4, (0-4)/4, (0-2)/2, (0-2)/2, (0-2)/2]
+        //       = [-1, -1, -1, -1, -1, -1]
+        let indices = Tensor::new(&[[0i64]], dev)?.unsqueeze(2)?; // (1, 1, 1)
+        let result = fsq.get_output_from_indices(&indices)?;
+        let vals = result.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
+        assert!((vals[0] - (-1.0)).abs() < 1e-5, "val[0]={}", vals[0]);
+        assert!((vals[3] - (-1.0)).abs() < 1e-5, "val[3]={}", vals[3]);
+
+        // Index that corresponds to center: half_width values for each level
+        // level_vals = [4, 4, 4, 2, 2, 2]
+        // packed index = 4*1 + 4*8 + 4*64 + 2*512 + 2*2560 + 2*12800
+        //             = 4 + 32 + 256 + 1024 + 5120 + 25600 = 32036
+        // codes = [0, 0, 0, 0, 0, 0]
+        let indices = Tensor::new(&[[32036i64]], dev)?.unsqueeze(2)?;
+        let result = fsq.get_output_from_indices(&indices)?;
+        let vals = result.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
+        for (i, &v) in vals.iter().enumerate() {
+            assert!(v.abs() < 1e-5, "center val[{i}]={v}, expected 0");
+        }
+        Ok(())
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AudioTokenizer
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+pub struct AudioTokenizer {
+    audio_acoustic_proj: Linear,
+    attention_pooler: AttentionPooler,
+    quantizer: ResidualFSQ,
+    pool_window_size: usize,
+}
+
+impl AudioTokenizer {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        Ok(Self {
+            audio_acoustic_proj: linear(
+                cfg.audio_acoustic_hidden_dim,
+                enc_hs,
+                vb.pp("audio_acoustic_proj"),
+            )?,
+            attention_pooler: AttentionPooler::new(cfg, vb.pp("attention_pooler"))?,
+            quantizer: ResidualFSQ::new(cfg, vb.pp("quantizer"))?,
+            pool_window_size: cfg.pool_window_size,
+        })
+    }
+
+    pub fn pool_window_size(&self) -> usize {
+        self.pool_window_size
+    }
+
+    /// Reconstruct from packed indices via the quantizer.
+    /// Input: `(B, N, 1)` → Output: `(B, N, dim)`.
+    pub fn get_output_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
+        self.quantizer.get_output_from_indices(indices)
+    }
+
+    /// Tokenize audio features into quantized tokens.
+    /// Input: (B, T_patches, pool_window_size, acoustic_dim)
+    /// Output: (quantized, indices)
+    pub fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
+        let x = x.apply(&self.audio_acoustic_proj)?;
+        let pooled = self.attention_pooler.forward(&x)?;
+        self.quantizer.forward(&pooled)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AudioTokenDetokenizer
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+pub struct AudioTokenDetokenizer {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    proj_out: Linear,
+    special_tokens: Tensor,
+    pool_window_size: usize,
+    rotary_emb: Arc<RotaryEmbedding>,
+}
+
+impl AudioTokenDetokenizer {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        let num_layers = cfg.num_attention_pooler_hidden_layers;
+
+        let mut layers = Vec::with_capacity(num_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..num_layers {
+            layers.push(EncoderLayer::new(cfg, i, vb_l.pp(i))?);
+        }
+        let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?);
+
+        Ok(Self {
+            embed_tokens: linear(enc_hs, enc_hs, vb.pp("embed_tokens"))?,
+            layers,
+            norm: RmsNorm::new(enc_hs, cfg.rms_norm_eps, vb.pp("norm"))?,
+            proj_out: linear(enc_hs, cfg.audio_acoustic_hidden_dim, vb.pp("proj_out"))?,
+            special_tokens: vb.get((1, cfg.pool_window_size, enc_hs), "special_tokens")?,
+            pool_window_size: cfg.pool_window_size,
+            rotary_emb,
+        })
+    }
+
+    /// Detokenize quantized tokens back to continuous acoustic representations.
+    /// Input: (B, T, D) → Output: (B, T * pool_window_size, acoustic_dim)
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let (b, t, _d) = x.dims3()?;
+        let x = x.apply(&self.embed_tokens)?;
+
+        // Expand each token into pool_window_size patches
+        let x = x
+            .unsqueeze(2)?
+            .broadcast_as((b, t, self.pool_window_size, x.dim(D::Minus1)?))?;
+        let special = self.special_tokens.broadcast_as((
+            b,
+            t,
+            self.pool_window_size,
+            self.special_tokens.dim(D::Minus1)?,
+        ))?;
+        let x = (x.contiguous()? + special.contiguous()?)?;
+
+        // Reshape for processing: (B*T, pool_window_size, hidden)
+        let x = x.reshape((b * t, self.pool_window_size, ()))?;
+
+        let mut hidden = x;
+        for layer in &self.layers {
+            hidden = layer.forward(&hidden, None, &self.rotary_emb)?;
+        }
+        let hidden = self.norm.forward(&hidden)?;
+        let hidden = hidden.apply(&self.proj_out)?;
+
+        // Reshape back: (B*T, P, acoustic_dim) → (B, T*P, acoustic_dim)
+        hidden.reshape((b, t * self.pool_window_size, ()))
+    }
+}

crates/neuron/src/harness/arch/acestep/vae.rs

@@ -0,0 +1,721 @@
+//! AutoencoderOobleck VAE for the ACE-Step music generation model.
+//!
+//! This module implements the Oobleck audio VAE that encodes 48kHz stereo audio
+//! into a compact latent representation and decodes it back. The architecture
+//! uses Snake1d activations, weight-normalized convolutions, and a progressive
+//! channel structure with strided up/downsampling.
+//!
+//! The encoder compresses (B, 2, T_audio) to (B, 128, T_latent) which is then
+//! split into mean and log-variance (each 64 channels) for the diagonal Gaussian.
+//! The decoder reconstructs (B, 64, T_latent) back to (B, 2, T_audio) where
+//! T_audio is approximately T_latent * 1920 (the product of downsampling ratios).
+//!
+//! Reference: diffusers `AutoencoderOobleck`
+
+use candle_core::{D, Module, Result, Tensor};
+use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, VarBuilder};
+
+use super::VaeConfig;
+use candle_transformers::models::encodec;
+
+// ---------------------------------------------------------------------------
+// Snake1d activation
+// ---------------------------------------------------------------------------
+
+/// Learnable Snake activation with alpha and beta parameters stored in log-scale.
+///
+/// Computes: `x + (exp(beta) + 1e-9).recip() * sin(exp(alpha) * x)^2`
+///
+/// Both `alpha` and `beta` have shape `(1, channels, 1)` and are exponentiated
+/// before use to ensure positivity.
+#[derive(Debug, Clone)]
+pub struct Snake1d {
+    alpha: Tensor,
+    beta: Tensor,
+}
+
+impl Snake1d {
+    pub fn new(channels: usize, vb: VarBuilder) -> Result<Self> {
+        let alpha = vb.get((1, channels, 1), "alpha")?;
+        let beta = vb.get((1, channels, 1), "beta")?;
+        Ok(Self { alpha, beta })
+    }
+}
+
+impl Module for Snake1d {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let alpha = self.alpha.exp()?;
+        let beta = self.beta.exp()?;
+        let sin = alpha.broadcast_mul(xs)?.sin()?;
+        let sin_sq = (&sin * &sin)?;
+        xs + ((&beta + 1e-9)?.recip()?.broadcast_mul(&sin_sq)?)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckResidualUnit
+// ---------------------------------------------------------------------------
+
+/// Residual unit with dilated convolutions and Snake1d activations.
+///
+/// Structure: `snake1 -> conv1(k=7, dilation=d) -> snake2 -> conv2(k=1) + skip`
+///
+/// All convolutions use weight normalization. The skip connection trims the
+/// input if the output is shorter due to dilation-induced size differences.
+#[derive(Debug, Clone)]
+pub struct OobleckResidualUnit {
+    snake1: Snake1d,
+    conv1: Conv1d,
+    snake2: Snake1d,
+    conv2: Conv1d,
+}
+
+impl OobleckResidualUnit {
+    pub fn new(dim: usize, dilation: usize, vb: VarBuilder) -> Result<Self> {
+        let pad = ((7 - 1) * dilation) / 2;
+        let snake1 = Snake1d::new(dim, vb.pp("snake1"))?;
+        let cfg1 = Conv1dConfig {
+            dilation,
+            padding: pad,
+            ..Default::default()
+        };
+        let conv1 = encodec::conv1d_weight_norm(dim, dim, 7, cfg1, vb.pp("conv1"))?;
+        let snake2 = Snake1d::new(dim, vb.pp("snake2"))?;
+        let conv2 = encodec::conv1d_weight_norm(dim, dim, 1, Default::default(), vb.pp("conv2"))?;
+        Ok(Self {
+            snake1,
+            conv1,
+            snake2,
+            conv2,
+        })
+    }
+}
+
+impl Module for OobleckResidualUnit {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let ys = xs
+            .apply(&self.snake1)?
+            .apply(&self.conv1)?
+            .apply(&self.snake2)?
+            .apply(&self.conv2)?;
+        let xs_len = xs.dim(D::Minus1)?;
+        let ys_len = ys.dim(D::Minus1)?;
+        if xs_len >= ys_len {
+            let pad = (xs_len - ys_len) / 2;
+            if pad > 0 {
+                &ys + xs.narrow(D::Minus1, pad, ys_len)?
+            } else {
+                ys + xs
+            }
+        } else {
+            let pad = (ys_len - xs_len) / 2;
+            &ys.narrow(D::Minus1, pad, xs_len)? + xs
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckEncoderBlock
+// ---------------------------------------------------------------------------
+
+/// Encoder block: 3 residual units (dilations 1, 3, 9) followed by Snake1d
+/// activation and a strided downsampling convolution.
+#[derive(Debug, Clone)]
+pub struct OobleckEncoderBlock {
+    res_unit1: OobleckResidualUnit,
+    res_unit2: OobleckResidualUnit,
+    res_unit3: OobleckResidualUnit,
+    snake1: Snake1d,
+    conv1: Conv1d,
+}
+
+impl OobleckEncoderBlock {
+    pub fn new(input_dim: usize, output_dim: usize, stride: usize, vb: VarBuilder) -> Result<Self> {
+        let res_unit1 = OobleckResidualUnit::new(input_dim, 1, vb.pp("res_unit1"))?;
+        let res_unit2 = OobleckResidualUnit::new(input_dim, 3, vb.pp("res_unit2"))?;
+        let res_unit3 = OobleckResidualUnit::new(input_dim, 9, vb.pp("res_unit3"))?;
+        let snake1 = Snake1d::new(input_dim, vb.pp("snake1"))?;
+        let cfg = Conv1dConfig {
+            stride,
+            padding: stride.div_ceil(2),
+            ..Default::default()
+        };
+        let conv1 =
+            encodec::conv1d_weight_norm(input_dim, output_dim, 2 * stride, cfg, vb.pp("conv1"))?;
+        Ok(Self {
+            res_unit1,
+            res_unit2,
+            res_unit3,
+            snake1,
+            conv1,
+        })
+    }
+}
+
+impl Module for OobleckEncoderBlock {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        xs.apply(&self.res_unit1)?
+            .apply(&self.res_unit2)?
+            .apply(&self.res_unit3)?
+            .apply(&self.snake1)?
+            .apply(&self.conv1)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckDecoderBlock
+// ---------------------------------------------------------------------------
+
+/// Decoder block: Snake1d activation and transposed convolution for upsampling,
+/// followed by 3 residual units (dilations 1, 3, 9).
+#[derive(Debug, Clone)]
+pub struct OobleckDecoderBlock {
+    snake1: Snake1d,
+    conv_t1: ConvTranspose1d,
+    res_unit1: OobleckResidualUnit,
+    res_unit2: OobleckResidualUnit,
+    res_unit3: OobleckResidualUnit,
+}
+
+impl OobleckDecoderBlock {
+    pub fn new(input_dim: usize, output_dim: usize, stride: usize, vb: VarBuilder) -> Result<Self> {
+        let snake1 = Snake1d::new(input_dim, vb.pp("snake1"))?;
+        let cfg = ConvTranspose1dConfig {
+            stride,
+            padding: stride.div_ceil(2),
+            ..Default::default()
+        };
+        let conv_t1 = encodec::conv_transpose1d_weight_norm(
+            input_dim,
+            output_dim,
+            2 * stride,
+            true,
+            cfg,
+            vb.pp("conv_t1"),
+        )?;
+        let res_unit1 = OobleckResidualUnit::new(output_dim, 1, vb.pp("res_unit1"))?;
+        let res_unit2 = OobleckResidualUnit::new(output_dim, 3, vb.pp("res_unit2"))?;
+        let res_unit3 = OobleckResidualUnit::new(output_dim, 9, vb.pp("res_unit3"))?;
+        Ok(Self {
+            snake1,
+            conv_t1,
+            res_unit1,
+            res_unit2,
+            res_unit3,
+        })
+    }
+}
+
+impl Module for OobleckDecoderBlock {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        xs.apply(&self.snake1)?
+            .apply(&self.conv_t1)?
+            .apply(&self.res_unit1)?
+            .apply(&self.res_unit2)?
+            .apply(&self.res_unit3)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckEncoder
+// ---------------------------------------------------------------------------
+
+/// Oobleck encoder that compresses stereo audio to a latent representation.
+///
+/// Channel progression (default config):
+/// `2 -> 128 -> 128 -> 256 -> 512 -> 1024 -> 2048 -> 128`
+#[derive(Debug, Clone)]
+pub struct OobleckEncoder {
+    conv1: Conv1d,
+    blocks: Vec<OobleckEncoderBlock>,
+    snake1: Snake1d,
+    conv2: Conv1d,
+}
+
+impl OobleckEncoder {
+    pub fn new(config: &VaeConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = config.encoder_hidden_size;
+
+        // Prepend 1 to channel_multiples: [1, 1, 2, 4, 8, 16]
+        let mut multiples = vec![1usize];
+        multiples.extend_from_slice(&config.channel_multiples);
+
+        let cfg1 = Conv1dConfig {
+            padding: 3,
+            ..Default::default()
+        };
+        let conv1 =
+            encodec::conv1d_weight_norm(config.audio_channels, enc_hs, 7, cfg1, vb.pp("conv1"))?;
+
+        let mut blocks = Vec::with_capacity(config.downsampling_ratios.len());
+        for (i, &stride) in config.downsampling_ratios.iter().enumerate() {
+            let in_channels = enc_hs * multiples[i];
+            let out_channels = enc_hs * multiples[i + 1];
+            let block =
+                OobleckEncoderBlock::new(in_channels, out_channels, stride, vb.pp("block").pp(i))?;
+            blocks.push(block);
+        }
+
+        let final_channels = enc_hs * multiples[multiples.len() - 1];
+        let snake1 = Snake1d::new(final_channels, vb.pp("snake1"))?;
+
+        let cfg2 = Conv1dConfig {
+            padding: 1,
+            ..Default::default()
+        };
+        let conv2 = encodec::conv1d_weight_norm(final_channels, enc_hs, 3, cfg2, vb.pp("conv2"))?;
+
+        Ok(Self {
+            conv1,
+            blocks,
+            snake1,
+            conv2,
+        })
+    }
+}
+
+impl Module for OobleckEncoder {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let mut xs = xs.apply(&self.conv1)?;
+        for block in &self.blocks {
+            xs = xs.apply(block)?;
+        }
+        xs.apply(&self.snake1)?.apply(&self.conv2)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckDecoder
+// ---------------------------------------------------------------------------
+
+/// Oobleck decoder that reconstructs stereo audio from a latent representation.
+///
+/// Channel progression (default config):
+/// `64 -> 2048 -> 1024 -> 512 -> 256 -> 128 -> 2`
+#[derive(Debug, Clone)]
+pub struct OobleckDecoder {
+    conv1: Conv1d,
+    blocks: Vec<OobleckDecoderBlock>,
+    snake1: Snake1d,
+    conv2: Conv1d,
+}
+
+impl OobleckDecoder {
+    pub fn new(config: &VaeConfig, vb: VarBuilder) -> Result<Self> {
+        let dec_ch = config.decoder_channels;
+
+        // Prepend 1 to channel_multiples: [1, 1, 2, 4, 8, 16]
+        let mut multiples = vec![1usize];
+        multiples.extend_from_slice(&config.channel_multiples);
+        let n = multiples.len(); // 6
+
+        // Upsampling ratios are reversed downsampling: [10, 6, 4, 4, 2]
+        let upsampling_ratios: Vec<usize> =
+            config.downsampling_ratios.iter().copied().rev().collect();
+
+        let cfg1 = Conv1dConfig {
+            padding: 3,
+            ..Default::default()
+        };
+        // conv1: decoder_input_channels -> dec_ch * mult[-1] (e.g. 64 -> 2048)
+        let first_channels = dec_ch * multiples[n - 1];
+        let conv1 = encodec::conv1d_weight_norm(
+            config.decoder_input_channels,
+            first_channels,
+            7,
+            cfg1,
+            vb.pp("conv1"),
+        )?;
+
+        let mut blocks = Vec::with_capacity(upsampling_ratios.len());
+        for (i, &stride) in upsampling_ratios.iter().enumerate() {
+            // Channels go from mult[N-1-i] to mult[N-2-i]
+            let in_channels = dec_ch * multiples[n - 1 - i];
+            let out_channels = dec_ch * multiples[n - 2 - i];
+            let block =
+                OobleckDecoderBlock::new(in_channels, out_channels, stride, vb.pp("block").pp(i))?;
+            blocks.push(block);
+        }
+
+        let snake1 = Snake1d::new(dec_ch, vb.pp("snake1"))?;
+
+        let cfg2 = Conv1dConfig {
+            padding: 3,
+            ..Default::default()
+        };
+        // Final conv: dec_ch -> audio_channels, NO bias
+        let conv2 = encodec::conv1d_weight_norm_no_bias(
+            dec_ch,
+            config.audio_channels,
+            7,
+            cfg2,
+            vb.pp("conv2"),
+        )?;
+
+        Ok(Self {
+            conv1,
+            blocks,
+            snake1,
+            conv2,
+        })
+    }
+
+    /// Tiled decoding for long latents that may not fit in memory.
+    ///
+    /// Splits the latent along the time axis into overlapping chunks, decodes
+    /// each chunk independently via `forward`, trims the overlap regions in
+    /// audio space, and concatenates the cores. The decoder is fully
+    /// convolutional with local padding, so cores away from chunk boundaries
+    /// are numerically identical to a single full-length decode (up to the
+    /// decoder's finite receptive field, absorbed by the overlap).
+    ///
+    /// For short latents (`T_latent <= chunk_latent_frames`) this falls back
+    /// to a plain `forward` call.
+    ///
+    /// - `chunk_latent_frames`: latent frames per chunk (default 128 ≈ 5.1s of
+    ///   audio at 48 kHz with a 1920× upsample ratio).
+    /// - `overlap_latent_frames`: overlap on each side in latent frames
+    ///   (default 16). Must satisfy `chunk > 2 * overlap`.
+    ///
+    /// Input: `(B, decoder_input_channels, T_latent)`.
+    /// Output: `(B, audio_channels, T_latent * upsample_ratio)`.
+    pub fn tiled_decode(
+        &self,
+        z: &Tensor,
+        chunk_latent_frames: Option<usize>,
+        overlap_latent_frames: Option<usize>,
+    ) -> Result<Tensor> {
+        let chunk_size = chunk_latent_frames.unwrap_or(128);
+        let overlap = overlap_latent_frames.unwrap_or(16);
+        let total_frames = z.dim(2)?;
+
+        if total_frames <= chunk_size {
+            return self.forward(z);
+        }
+
+        if chunk_size <= 2 * overlap {
+            candle_core::bail!(
+                "tiled_decode: chunk_latent_frames ({chunk_size}) must be > 2 * overlap_latent_frames ({overlap})"
+            );
+        }
+
+        let stride = chunk_size - 2 * overlap;
+        let num_steps = total_frames.div_ceil(stride);
+        let mut upsample_factor: Option<usize> = None;
+        let mut audio_chunks = Vec::with_capacity(num_steps);
+
+        for i in 0..num_steps {
+            let core_start = i * stride;
+            let core_end = (core_start + stride).min(total_frames);
+            if core_start >= core_end {
+                break;
+            }
+
+            // Expand the window with overlap on both sides, clipped at the edges.
+            let win_start = core_start.saturating_sub(overlap);
+            let win_end = (core_end + overlap).min(total_frames);
+
+            let latent_chunk = z.narrow(2, win_start, win_end - win_start)?;
+            let audio_chunk = self.forward(&latent_chunk)?;
+
+            // Determine the upsample factor from the first chunk. The decoder
+            // is strided-conv-transpose only, so audio_len is an integer
+            // multiple of latent_len.
+            let up = match upsample_factor {
+                Some(up) => up,
+                None => {
+                    let audio_len = audio_chunk.dim(2)?;
+                    let latent_len = latent_chunk.dim(2)?;
+                    if latent_len == 0 || audio_len % latent_len != 0 {
+                        candle_core::bail!(
+                            "tiled_decode: audio length {audio_len} is not a multiple of latent length {latent_len}"
+                        );
+                    }
+                    let up = audio_len / latent_len;
+                    upsample_factor = Some(up);
+                    up
+                }
+            };
+
+            // Trim the overlap in audio space to recover only the core.
+            let trim_start = (core_start - win_start) * up;
+            let trim_end = (win_end - core_end) * up;
+            let audio_len = audio_chunk.dim(2)?;
+            let core_len = audio_len - trim_start - trim_end;
+            let audio_core = audio_chunk.narrow(2, trim_start, core_len)?;
+            audio_chunks.push(audio_core);
+        }
+
+        Tensor::cat(&audio_chunks.iter().collect::<Vec<_>>(), 2)
+    }
+}
+
+impl Module for OobleckDecoder {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let mut xs = xs.apply(&self.conv1)?;
+        for block in &self.blocks {
+            xs = xs.apply(block)?;
+        }
+        xs.apply(&self.snake1)?.apply(&self.conv2)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AutoencoderOobleck
+// ---------------------------------------------------------------------------
+
+/// AutoencoderOobleck: full VAE combining encoder and decoder for audio.
+///
+/// Encodes stereo 48kHz audio `(B, 2, T)` to latent `(B, 64, T/1920)` and
+/// decodes back. Uses a diagonal Gaussian latent distribution with
+/// reparameterization sampling.
+#[derive(Debug, Clone)]
+pub struct AutoencoderOobleck {
+    encoder: OobleckEncoder,
+    decoder: OobleckDecoder,
+    config: VaeConfig,
+}
+
+impl AutoencoderOobleck {
+    pub fn new(config: &VaeConfig, vb: VarBuilder) -> Result<Self> {
+        let encoder = OobleckEncoder::new(config, vb.pp("encoder"))?;
+        let decoder = OobleckDecoder::new(config, vb.pp("decoder"))?;
+        Ok(Self {
+            encoder,
+            decoder,
+            config: config.clone(),
+        })
+    }
+
+    /// Encode audio to a latent sample using the diagonal Gaussian.
+    ///
+    /// The encoder output `(B, encoder_hidden_size, T_latent)` is split along
+    /// the channel dimension into mean and log-variance (each with
+    /// `decoder_input_channels` channels), then sampled via reparameterization:
+    /// `z = mean + exp(0.5 * logvar) * noise`.
+    pub fn encode(&self, xs: &Tensor) -> Result<Tensor> {
+        let h = self.encoder.forward(xs)?;
+        let latent_dim = self.config.decoder_input_channels;
+        let mean = h.narrow(1, 0, latent_dim)?;
+        let logvar = h.narrow(1, latent_dim, latent_dim)?;
+        let std = (logvar * 0.5)?.exp()?;
+        let noise = std.randn_like(0., 1.)?;
+        &mean + std.mul(&noise)?
+    }
+
+    /// Encode audio to the latent mean (no sampling noise).
+    ///
+    /// For deterministic encoding (e.g. when exact reconstruction matters),
+    /// this returns only the mean of the diagonal Gaussian, without
+    /// reparameterization noise.
+    pub fn encode_mean(&self, xs: &Tensor) -> Result<Tensor> {
+        let h = self.encoder.forward(xs)?;
+        let latent_dim = self.config.decoder_input_channels;
+        h.narrow(1, 0, latent_dim)
+    }
+
+    /// Decode a latent representation back to audio waveform.
+    ///
+    /// Input shape: `(B, decoder_input_channels, T_latent)`
+    /// Output shape: `(B, audio_channels, T_audio)`
+    pub fn decode(&self, z: &Tensor) -> Result<Tensor> {
+        self.decoder.forward(z)
+    }
+
+    /// Tiled decoding for long latents; see [`OobleckDecoder::tiled_decode`].
+    pub fn tiled_decode(
+        &self,
+        z: &Tensor,
+        chunk_latent_frames: Option<usize>,
+        overlap_latent_frames: Option<usize>,
+    ) -> Result<Tensor> {
+        self.decoder
+            .tiled_decode(z, chunk_latent_frames, overlap_latent_frames)
+    }
+
+    /// Tiled encoding for long audio that may not fit in GPU memory.
+    ///
+    /// Splits the audio into overlapping chunks, encodes each independently,
+    /// trims the overlap regions, and concatenates. Uses the overlap-discard
+    /// strategy to avoid boundary artifacts.
+    ///
+    /// - `chunk_samples`: audio samples per chunk (default: 48000 * 30 = 30s)
+    /// - `overlap_samples`: overlap in audio samples (default: 48000 * 2 = 2s)
+    ///
+    /// Input: `(B, channels, T_audio)`, output: `(B, latent_dim, T_latent)`
+    pub fn tiled_encode(
+        &self,
+        xs: &Tensor,
+        chunk_samples: Option<usize>,
+        overlap_samples: Option<usize>,
+    ) -> Result<Tensor> {
+        let chunk_size = chunk_samples.unwrap_or(48000 * 30);
+        let overlap = overlap_samples.unwrap_or(48000 * 2);
+        let total_samples = xs.dim(2)?;
+
+        // Short audio: encode directly
+        if total_samples <= chunk_size {
+            return self.encode(xs);
+        }
+
+        if chunk_size <= 2 * overlap {
+            candle_core::bail!(
+                "tiled_encode: chunk_size ({chunk_size}) must be > 2 * overlap ({overlap})"
+            );
+        }
+        let stride = chunk_size - 2 * overlap;
+        let num_steps = total_samples.div_ceil(stride);
+        let mut downsample_factor: Option<f64> = None;
+        let mut latent_chunks = Vec::with_capacity(num_steps);
+
+        for i in 0..num_steps {
+            let core_start = i * stride;
+            let core_end = (core_start + stride).min(total_samples);
+
+            // Window with overlap
+            let win_start = core_start.saturating_sub(overlap);
+            let win_end = (core_end + overlap).min(total_samples);
+
+            let audio_chunk = xs.narrow(2, win_start, win_end - win_start)?;
+            let latent_chunk = self.encode(&audio_chunk)?;
+
+            // Determine downsample factor from first chunk
+            let df = match downsample_factor {
+                Some(df) => df,
+                None => {
+                    let df = audio_chunk.dim(2)? as f64 / latent_chunk.dim(2)? as f64;
+                    downsample_factor = Some(df);
+                    df
+                }
+            };
+
+            // Trim overlap in latent space
+            let added_start = core_start - win_start;
+            let trim_start = (added_start as f64 / df).round() as usize;
+
+            let added_end = win_end - core_end;
+            let trim_end = (added_end as f64 / df).round() as usize;
+
+            let latent_len = latent_chunk.dim(2)?;
+            let end_idx = if trim_end > 0 {
+                latent_len - trim_end
+            } else {
+                latent_len
+            };
+            let core_len = end_idx - trim_start;
+            let latent_core = latent_chunk.narrow(2, trim_start, core_len)?;
+            latent_chunks.push(latent_core);
+        }
+
+        Tensor::cat(&latent_chunks.iter().collect::<Vec<_>>(), 2)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+    use candle_nn::{VarBuilder, VarMap};
+
+    /// Build a small OobleckDecoder with random weights on CPU. The
+    /// architecture is identical to the production decoder; we just shrink the
+    /// channel count so the test runs quickly.
+    fn tiny_decoder() -> Result<(OobleckDecoder, Device)> {
+        let device = Device::Cpu;
+        let varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let config = VaeConfig {
+            decoder_channels: 8,
+            decoder_input_channels: 4,
+            audio_channels: 2,
+            // Five blocks, as in production, but 2x each (hop 32 instead of
+            // 1920): the tiling logic is the same, and the full-size hop made
+            // this test peak at 4.6 GB, which got the CI test binary killed.
+            downsampling_ratios: vec![2, 2, 2, 2, 2],
+            channel_multiples: vec![1, 1, 1, 1, 1],
+            ..VaeConfig::default()
+        };
+        let decoder = OobleckDecoder::new(&config, vb)?;
+        Ok((decoder, device))
+    }
+
+    /// Elementwise compare two tensors tolerant to matching NaN/Inf values.
+    /// The decoder has random weights with no normalization, so activations
+    /// can explode past fp32 range — but `forward` and `tiled_decode` do the
+    /// same computation on the same slice, so whatever NaN/Inf pattern
+    /// `forward` produces, `tiled_decode` must produce in the same places.
+    fn tensors_agree(a: &Tensor, b: &Tensor, atol: f32) -> Result<bool> {
+        let a = a.flatten_all()?.to_vec1::<f32>()?;
+        let b = b.flatten_all()?.to_vec1::<f32>()?;
+        assert_eq!(a.len(), b.len(), "shape mismatch");
+        for (x, y) in a.iter().zip(b.iter()) {
+            let ok = match (x.is_nan(), y.is_nan()) {
+                (true, true) => true,
+                (false, false) => {
+                    if x.is_infinite() || y.is_infinite() {
+                        x == y
+                    } else {
+                        let scale = x.abs().max(y.abs()).max(1.0);
+                        (x - y).abs() <= atol * scale
+                    }
+                }
+                _ => false,
+            };
+            if !ok {
+                eprintln!("mismatch: forward={x} tiled={y}");
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    }
+
+    #[test]
+    fn tiled_decode_short_input_delegates_to_forward() -> Result<()> {
+        let (decoder, device) = tiny_decoder()?;
+        // total < chunk → tiled_decode must short-circuit to forward.
+        let z = Tensor::randn(0f32, 1.0, (1, 4, 64), &device)?;
+
+        let y_direct = decoder.forward(&z)?;
+        let y_tiled = decoder.tiled_decode(&z, Some(128), Some(16))?;
+
+        assert_eq!(y_tiled.dims(), y_direct.dims());
+        assert!(tensors_agree(&y_direct, &y_tiled, 1e-5)?);
+        Ok(())
+    }
+
+    #[test]
+    fn tiled_decode_matches_forward_across_chunks() -> Result<()> {
+        let (decoder, device) = tiny_decoder()?;
+        // Force multiple chunk boundaries.
+        let t_latent = 300;
+        let z = Tensor::randn(0f32, 1.0, (1, 4, t_latent), &device)?;
+
+        let y_direct = decoder.forward(&z)?;
+        let y_tiled = decoder.tiled_decode(&z, Some(64), Some(16))?;
+
+        assert_eq!(
+            y_tiled.dims(),
+            y_direct.dims(),
+            "tiled output shape must equal full-forward shape"
+        );
+        // Receptive field fits inside the 16-frame overlap, so cores should
+        // match within fp32 rounding slack. Values can be extreme due to
+        // random untrained weights; compare relative.
+        assert!(tensors_agree(&y_direct, &y_tiled, 1e-3)?);
+        Ok(())
+    }
+
+    #[test]
+    fn tiled_decode_rejects_bad_overlap() -> Result<()> {
+        let (decoder, device) = tiny_decoder()?;
+        let z = Tensor::randn(0f32, 1.0, (1, 4, 200), &device)?;
+
+        // chunk <= 2 * overlap must fail.
+        let err = decoder.tiled_decode(&z, Some(32), Some(16));
+        assert!(
+            err.is_err(),
+            "expected tiled_decode to reject chunk<=2*overlap"
+        );
+        Ok(())
+    }
+}

crates/neuron/src/harness/arch/heartmula/lm.rs

@@ -0,0 +1,750 @@
+//! The HeartMuLa language model (#391): a 3B "global" backbone over 12.5 Hz
+//! frames and a 300M "local" decoder over the eight codebooks inside a
+//! frame, both torchtune `llama3_2` transformers.
+//!
+//! One frame is generated in two stages:
+//!
+//! 1. The backbone reads the frame's input — the **sum** of nine embedding
+//!    streams (eight audio codebooks from one shared table, plus text) — and
+//!    its last hidden state predicts codebook 0 through `codebook0_head`.
+//! 2. The local decoder, reset every frame, reads `[last_h, embed(c0)]` and
+//!    then one codebook embedding per step, predicting codebooks 1–7 through
+//!    the per-codebook `audio_head` matrices.
+//!
+//! Classifier-free guidance runs a second batch row whose *text* embedding
+//! is replaced by `unconditional_text_embedding`; every head's logits are
+//! combined as `uncond + scale · (cond − uncond)` and the sampled code is
+//! fed to both rows.
+//!
+//! Reference: heartlib `heartmula/modeling_heartmula.py` and
+//! `pipelines/music_generation.py` @ 5858ca8f. Two torchtune details matter
+//! for parity: RoPE rotates *interleaved* pairs (candle's `rope_i`, not the
+//! half-split HF layout), and its frequencies carry Llama 3.1 scaling.
+
+use anyhow::Context;
+use candle_core::{D, DType, Device, IndexOp, Module, Result, Tensor};
+use candle_nn::kv_cache::KvCache;
+use candle_nn::{Linear, VarBuilder, linear_no_bias};
+use serde::Deserialize;
+
+/// One torchtune `llama3_2` shape.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct Flavor {
+    pub layers: usize,
+    pub heads: usize,
+    pub kv_heads: usize,
+    pub dim: usize,
+    pub mlp_dim: usize,
+}
+
+impl Flavor {
+    pub fn named(name: &str) -> anyhow::Result<Self> {
+        let f = match name {
+            "llama-3B" => Self {
+                layers: 28,
+                heads: 24,
+                kv_heads: 8,
+                dim: 3072,
+                mlp_dim: 8192,
+            },
+            "llama-300M" => Self {
+                layers: 3,
+                heads: 8,
+                kv_heads: 4,
+                dim: 3072,
+                mlp_dim: 8192,
+            },
+            "llama-400M" => Self {
+                layers: 4,
+                heads: 8,
+                kv_heads: 4,
+                dim: 3072,
+                mlp_dim: 8192,
+            },
+            "llama-7B" => Self {
+                layers: 32,
+                heads: 32,
+                kv_heads: 8,
+                dim: 4096,
+                mlp_dim: 14336,
+            },
+            other => anyhow::bail!("unknown HeartMuLa flavour '{other}'"),
+        };
+        Ok(f)
+    }
+
+    fn head_dim(&self) -> usize {
+        self.dim / self.heads
+    }
+}
+
+/// `config.json` of a HeartMuLa checkpoint.
+#[derive(Debug, Clone, Deserialize)]
+pub struct Config {
+    pub audio_num_codebooks: usize,
+    pub audio_vocab_size: usize,
+    pub text_vocab_size: usize,
+    pub muq_dim: usize,
+    pub backbone_flavor: String,
+    pub decoder_flavor: String,
+}
+
+/// `gen_config.json`: the special ids generation uses.
+#[derive(Debug, Clone, Deserialize)]
+pub struct GenConfig {
+    pub text_bos_id: u32,
+    pub text_eos_id: u32,
+    pub audio_eos_id: u32,
+    pub empty_id: u32,
+}
+
+/// torchtune's RMSNorm: normalise in f32, then scale (weight `scale`).
+struct RmsNorm {
+    scale: Tensor,
+    eps: f64,
+}
+
+impl RmsNorm {
+    fn new(dim: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
+        Ok(Self {
+            scale: vb.get(dim, "scale")?,
+            eps,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let dtype = x.dtype();
+        let xf = x.to_dtype(DType::F32)?;
+        let var = xf.sqr()?.mean_keepdim(D::Minus1)?;
+        let normed = xf.broadcast_div(&(var + self.eps)?.sqrt()?)?;
+        normed.to_dtype(dtype)?.broadcast_mul(&self.scale)
+    }
+}
+
+/// RoPE tables for torchtune's `Llama3ScaledRoPE`: base 500 000, scale
+/// factor 32, low/high frequency factors 1/4, original context 8192. The
+/// frequencies are computed in f32 and the per-frequency scaling in f64,
+/// as torchtune does (a Python loop over the f32 values).
+struct Rope {
+    cos: Tensor,
+    sin: Tensor,
+}
+
+const ROPE_BASE: f32 = 500_000.0;
+const ROPE_SCALE: f64 = 32.0;
+const ROPE_LOW_FREQ: f64 = 1.0;
+const ROPE_HIGH_FREQ: f64 = 4.0;
+const ROPE_OLD_CONTEXT: f64 = 8192.0;
+
+/// torchtune's Llama 3.1 frequency scaling, for one head dimension.
+pub fn llama3_scaled_theta(head_dim: usize) -> Vec<f32> {
+    let low_wavelen = ROPE_OLD_CONTEXT / ROPE_LOW_FREQ;
+    let high_wavelen = ROPE_OLD_CONTEXT / ROPE_HIGH_FREQ;
+    (0..head_dim / 2)
+        .map(|i| {
+            let freq = 1.0f32 / ROPE_BASE.powf((2 * i) as f32 / head_dim as f32);
+            let f = freq as f64;
+            let wavelen = 2.0 * std::f64::consts::PI / f;
+            let scaled = if wavelen < high_wavelen {
+                f
+            } else if wavelen > low_wavelen {
+                f / ROPE_SCALE
+            } else {
+                let smooth =
+                    (ROPE_OLD_CONTEXT / wavelen - ROPE_LOW_FREQ) / (ROPE_HIGH_FREQ - ROPE_LOW_FREQ);
+                (1.0 - smooth) * f / ROPE_SCALE + smooth * f
+            };
+            scaled as f32
+        })
+        .collect()
+}
+
+impl Rope {
+    fn new(head_dim: usize, max_seq: usize, dtype: DType, device: &Device) -> Result<Self> {
+        let theta = llama3_scaled_theta(head_dim);
+        let half = theta.len();
+        let mut cos = Vec::with_capacity(max_seq * half);
+        let mut sin = Vec::with_capacity(max_seq * half);
+        for p in 0..max_seq {
+            for &t in &theta {
+                let a = p as f32 * t;
+                cos.push(a.cos());
+                sin.push(a.sin());
+            }
+        }
+        Ok(Self {
+            cos: Tensor::from_vec(cos, (max_seq, half), device)?.to_dtype(dtype)?,
+            sin: Tensor::from_vec(sin, (max_seq, half), device)?.to_dtype(dtype)?,
+        })
+    }
+
+    /// Rotate `x` `(b, heads, seq, head_dim)` at positions `pos..pos+seq`.
+    fn apply(&self, x: &Tensor, pos: usize) -> Result<Tensor> {
+        let seq = x.dim(2)?;
+        let cos = self.cos.narrow(0, pos, seq)?;
+        let sin = self.sin.narrow(0, pos, seq)?;
+        candle_nn::rotary_emb::rope_i(&x.contiguous()?, &cos, &sin)
+    }
+}
+
+struct Attention {
+    q: Linear,
+    k: Linear,
+    v: Linear,
+    o: Linear,
+    heads: usize,
+    kv_heads: usize,
+    head_dim: usize,
+    cache: Option<KvCache>,
+}
+
+impl Attention {
+    fn new(f: &Flavor, vb: VarBuilder) -> Result<Self> {
+        let hd = f.head_dim();
+        Ok(Self {
+            q: linear_no_bias(f.dim, f.heads * hd, vb.pp("q_proj"))?,
+            k: linear_no_bias(f.dim, f.kv_heads * hd, vb.pp("k_proj"))?,
+            v: linear_no_bias(f.dim, f.kv_heads * hd, vb.pp("v_proj"))?,
+            o: linear_no_bias(f.heads * hd, f.dim, vb.pp("output_proj"))?,
+            heads: f.heads,
+            kv_heads: f.kv_heads,
+            head_dim: hd,
+            cache: None,
+        })
+    }
+
+    fn forward(&mut self, x: &Tensor, rope: &Rope, pos: usize) -> Result<Tensor> {
+        let (b, s, _) = x.dims3()?;
+        let shape = |t: Tensor, n: usize| -> Result<Tensor> {
+            t.reshape((b, s, n, self.head_dim))?.transpose(1, 2)
+        };
+        let q = rope.apply(&shape(self.q.forward(x)?, self.heads)?, pos)?;
+        let k = rope.apply(&shape(self.k.forward(x)?, self.kv_heads)?, pos)?;
+        let v = shape(self.v.forward(x)?, self.kv_heads)?.contiguous()?;
+        let (k, v) = match self.cache.as_mut() {
+            Some(cache) => cache.append(&k.contiguous()?, &v)?,
+            None => (k, v),
+        };
+        let rep = self.heads / self.kv_heads;
+        let k = candle_transformers::utils::repeat_kv(k, rep)?.contiguous()?;
+        let v = candle_transformers::utils::repeat_kv(v, rep)?.contiguous()?;
+        let kv_len = k.dim(2)?;
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut att = (q.contiguous()?.matmul(&k.t()?)? * scale)?;
+        if s > 1 {
+            // Causal over the whole cache: query i (absolute pos + i) sees
+            // keys 0..=pos+i.
+            let offset = kv_len - s;
+            let mask: Vec<f32> = (0..s)
+                .flat_map(|i| {
+                    (0..kv_len).map(move |j| {
+                        if j > offset + i {
+                            f32::NEG_INFINITY
+                        } else {
+                            0.0
+                        }
+                    })
+                })
+                .collect();
+            let mask = Tensor::from_vec(mask, (s, kv_len), x.device())?.to_dtype(att.dtype())?;
+            att = att.broadcast_add(&mask)?;
+        }
+        let att = candle_nn::ops::softmax_last_dim(&att)?;
+        let out = att.matmul(&v)?;
+        let out = out
+            .transpose(1, 2)?
+            .reshape((b, s, self.heads * self.head_dim))?;
+        self.o.forward(&out)
+    }
+}
+
+struct Mlp {
+    w1: Linear,
+    w2: Linear,
+    w3: Linear,
+}
+
+impl Mlp {
+    fn new(f: &Flavor, vb: VarBuilder) -> Result<Self> {
+        Ok(Self {
+            w1: linear_no_bias(f.dim, f.mlp_dim, vb.pp("w1"))?,
+            w2: linear_no_bias(f.mlp_dim, f.dim, vb.pp("w2"))?,
+            w3: linear_no_bias(f.dim, f.mlp_dim, vb.pp("w3"))?,
+        })
+    }
+
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let gate = candle_nn::ops::silu(&self.w1.forward(x)?)?;
+        self.w2.forward(&(gate * self.w3.forward(x)?)?)
+    }
+}
+
+struct Layer {
+    sa_norm: RmsNorm,
+    attn: Attention,
+    mlp_norm: RmsNorm,
+    mlp: Mlp,
+}
+
+/// A torchtune `TransformerDecoder` with identity token embedding and
+/// output head: it maps hidden states to normalised hidden states.
+struct Llama {
+    layers: Vec<Layer>,
+    norm: RmsNorm,
+    rope: Rope,
+}
+
+const NORM_EPS: f64 = 1e-5;
+
+impl Llama {
+    fn new(f: &Flavor, max_seq: usize, vb: VarBuilder) -> Result<Self> {
+        let mut layers = Vec::with_capacity(f.layers);
+        for i in 0..f.layers {
+            let lvb = vb.pp(format!("layers.{i}"));
+            layers.push(Layer {
+                sa_norm: RmsNorm::new(f.dim, NORM_EPS, lvb.pp("sa_norm"))?,
+                attn: Attention::new(f, lvb.pp("attn"))?,
+                mlp_norm: RmsNorm::new(f.dim, NORM_EPS, lvb.pp("mlp_norm"))?,
+                mlp: Mlp::new(f, lvb.pp("mlp"))?,
+            });
+        }
+        Ok(Self {
+            layers,
+            norm: RmsNorm::new(f.dim, NORM_EPS, vb.pp("norm"))?,
+            rope: Rope::new(f.head_dim(), max_seq, vb.dtype(), vb.device())?,
+        })
+    }
+
+    /// Fresh caches with room for `capacity` positions.
+    fn reset(&mut self, capacity: usize) {
+        for l in &mut self.layers {
+            l.attn.cache = Some(KvCache::new(2, capacity));
+        }
+    }
+
+    fn forward(&mut self, x: &Tensor, pos: usize) -> Result<Tensor> {
+        let mut h = x.clone();
+        for l in &mut self.layers {
+            let a = l.attn.forward(&l.sa_norm.forward(&h)?, &self.rope, pos)?;
+            h = (h + a)?;
+            let m = l.mlp.forward(&l.mlp_norm.forward(&h)?)?;
+            h = (h + m)?;
+        }
+        self.norm.forward(&h)
+    }
+}
+
+/// A prompt as the LM consumes it: one row per position, nine columns —
+/// eight audio codebooks then text — and the position of the reference-
+/// audio slot. In a prompt only the text column is live.
+#[derive(Debug, Clone, PartialEq)]
+pub struct Prompt {
+    pub tokens: Vec<[u32; 9]>,
+    pub muq_idx: usize,
+}
+
+/// Build the prompt exactly as heartlib does: tags lowercased and wrapped
+/// in `<tag>…</tag>`, lyrics lowercased, each tokenized with BOS/EOS
+/// ensured, then `tags ‖ [reference slot] ‖ lyrics`.
+pub fn build_prompt(
+    tokenizer: &tokenizers::Tokenizer,
+    tags: &str,
+    lyrics: &str,
+    gen_cfg: &GenConfig,
+) -> anyhow::Result<Prompt> {
+    let mut tags = tags.to_lowercase();
+    if !tags.starts_with("<tag>") {
+        tags = format!("<tag>{tags}");
+    }
+    if !tags.ends_with("</tag>") {
+        tags = format!("{tags}</tag>");
+    }
+    let encode = |text: &str| -> anyhow::Result<Vec<u32>> {
+        let mut ids = tokenizer
+            .encode(text, true)
+            .map_err(|e| anyhow::anyhow!("tokenize: {e}"))?
+            .get_ids()
+            .to_vec();
+        if ids.first() != Some(&gen_cfg.text_bos_id) {
+            ids.insert(0, gen_cfg.text_bos_id);
+        }
+        if ids.last() != Some(&gen_cfg.text_eos_id) {
+            ids.push(gen_cfg.text_eos_id);
+        }
+        Ok(ids)
+    };
+    let tags_ids = encode(&tags)?;
+    let lyric_ids = encode(&lyrics.to_lowercase())?;
+    let row = |text: u32| {
+        let mut r = [0u32; 9];
+        r[8] = text;
+        r
+    };
+    let mut tokens: Vec<[u32; 9]> = tags_ids.iter().map(|&t| row(t)).collect();
+    let muq_idx = tokens.len();
+    tokens.push([0; 9]);
+    tokens.extend(lyric_ids.iter().map(|&t| row(t)));
+    Ok(Prompt { tokens, muq_idx })
+}
+
+/// Top-k sampling with heartlib's draw: divide by temperature, mask
+/// everything below the k-th logit, softmax, then `argmax(p / Exp(1))`.
+/// With k = 1 that is exact greedy decoding.
+///
+/// The exponential draws come from a seeded host generator, so a seed
+/// reproduces a song on any device. They do not reproduce heartlib's own
+/// streams, which depend on the torch generator of the device it ran on.
+pub struct Sampler {
+    pub temperature: f64,
+    pub topk: usize,
+    state: u64,
+}
+
+impl Sampler {
+    pub fn new(temperature: f64, topk: usize, seed: u64) -> Self {
+        Self {
+            temperature,
+            topk: topk.max(1),
+            state: seed,
+        }
+    }
+
+    fn next_unit(&mut self) -> f64 {
+        self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
+        let mut z = self.state;
+        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+        z ^= z >> 31;
+        // 53 random bits in (0, 1]: never zero, so ln() is finite.
+        ((z >> 11) as f64 + 1.0) / (1u64 << 53) as f64
+    }
+
+    /// Pick a code from one row of guided logits.
+    pub fn sample(&mut self, logits: &[f32]) -> u32 {
+        let t = self.temperature as f32;
+        let scaled: Vec<f32> = logits.iter().map(|l| l / t).collect();
+        let mut sorted = scaled.clone();
+        sorted.sort_by(|a, b| b.total_cmp(a));
+        let k = self.topk.min(sorted.len());
+        let threshold = sorted[k - 1];
+        let max = sorted[0];
+        let mut best = (0usize, f64::NEG_INFINITY);
+        for (i, &l) in scaled.iter().enumerate() {
+            // Every logit consumes a draw, kept or not, as heartlib's
+            // `empty_like(probs).exponential_()` does.
+            let q = -self.next_unit().ln();
+            if l < threshold {
+                continue;
+            }
+            // Unnormalised probability is enough for an argmax.
+            let p = ((l - max) as f64).exp();
+            let score = p / q;
+            if score > best.1 {
+                best = (i, score);
+            }
+        }
+        best.0 as u32
+    }
+}
+
+/// The loaded model.
+pub struct HeartMuLa {
+    pub config: Config,
+    backbone: Llama,
+    decoder: Llama,
+    backbone_flavor: Flavor,
+    text_embeddings: Tensor,
+    audio_embeddings: Tensor,
+    unconditional_text: Tensor,
+    projection: Linear,
+    codebook0_head: Linear,
+    /// `(codebooks − 1, decoder_dim, vocab)`.
+    audio_head: Tensor,
+    /// `muq_linear(zeros)`: the reference-audio slot's conditional value.
+    muq_bias: Tensor,
+    device: Device,
+    dtype: DType,
+}
+
+/// Decoder KV length: `[last_h, c0]` then one step per further codebook.
+const DECODER_MAX_SEQ: usize = 16;
+/// The backbone's RoPE table length (torchtune `max_seq_len`).
+pub const BACKBONE_MAX_SEQ: usize = 8192;
+
+impl HeartMuLa {
+    /// Load from a checkpoint directory's `VarBuilder`.
+    pub fn new(config: Config, vb: VarBuilder) -> anyhow::Result<Self> {
+        let bf = Flavor::named(&config.backbone_flavor)?;
+        let df = Flavor::named(&config.decoder_flavor)?;
+        let backbone =
+            Llama::new(&bf, BACKBONE_MAX_SEQ, vb.pp("backbone")).context("load backbone")?;
+        let decoder = Llama::new(&df, DECODER_MAX_SEQ, vb.pp("decoder")).context("load decoder")?;
+        let v = config.audio_vocab_size;
+        let n = config.audio_num_codebooks;
+        let muq_bias = vb.get(bf.dim, "muq_linear.bias")?;
+        // The weight is only ever applied to zeros (#389); load it anyway
+        // so a checkpoint missing it is caught as malformed.
+        let _ = vb.get((bf.dim, config.muq_dim), "muq_linear.weight")?;
+        Ok(Self {
+            text_embeddings: vb.get((config.text_vocab_size, bf.dim), "text_embeddings.weight")?,
+            audio_embeddings: vb.get((v * n, bf.dim), "audio_embeddings.weight")?,
+            unconditional_text: vb.get((1, bf.dim), "unconditional_text_embedding.weight")?,
+            projection: linear_no_bias(bf.dim, df.dim, vb.pp("projection"))?,
+            codebook0_head: linear_no_bias(bf.dim, v, vb.pp("codebook0_head"))?,
+            audio_head: vb.get((n - 1, df.dim, v), "audio_head")?,
+            muq_bias,
+            backbone,
+            decoder,
+            backbone_flavor: bf,
+            device: vb.device().clone(),
+            dtype: vb.dtype(),
+            config,
+        })
+    }
+
+    pub fn device(&self) -> &Device {
+        &self.device
+    }
+
+    /// Start a new song: fresh backbone caches sized for `positions`
+    /// (prompt plus frames).
+    pub fn reset(&mut self, positions: usize) {
+        self.backbone.reset(positions.min(BACKBONE_MAX_SEQ));
+    }
+
+    /// The embedded input `(rows, seq, dim)` for `tokens`: the masked sum of
+    /// the eight audio streams and the text stream. With CFG the second row
+    /// replaces its text embedding with the unconditional one.
+    fn embed(&self, tokens: &[[u32; 9]], mask: &[[bool; 9]], cfg: bool) -> Result<Tensor> {
+        let s = tokens.len();
+        let d = self.backbone_flavor.dim;
+        let v = self.config.audio_vocab_size as u32;
+        let text_ids: Vec<u32> = tokens.iter().map(|r| r[8]).collect();
+        let text = self
+            .text_embeddings
+            .index_select(&Tensor::new(text_ids.as_slice(), &self.device)?, 0)?;
+        let audio_ids: Vec<u32> = tokens
+            .iter()
+            .flat_map(|r| (0..8).map(move |k| r[k] + k as u32 * v))
+            .collect();
+        let audio = self
+            .audio_embeddings
+            .index_select(&Tensor::new(audio_ids.as_slice(), &self.device)?, 0)?
+            .reshape((s, 8, d))?;
+        let m: Vec<f32> = mask
+            .iter()
+            .flat_map(|r| r.iter().map(|&b| if b { 1.0 } else { 0.0 }))
+            .collect();
+        let m = Tensor::from_vec(m, (s, 9), &self.device)?.to_dtype(self.dtype)?;
+        let audio_sum = audio
+            .broadcast_mul(&m.narrow(1, 0, 8)?.unsqueeze(2)?)?
+            .sum(1)?;
+        let text_mask = m.narrow(1, 8, 1)?;
+        let cond = (&audio_sum + text.broadcast_mul(&text_mask)?)?;
+        if !cfg {
+            return cond.unsqueeze(0);
+        }
+        let uncond_text = self.unconditional_text.broadcast_as((s, d))?;
+        let uncond = (&audio_sum + uncond_text.broadcast_mul(&text_mask)?)?;
+        Tensor::stack(&[cond, uncond], 0)
+    }
+
+    fn embed_audio(&self, codebook: usize, code: u32, rows: usize) -> Result<Tensor> {
+        let id = code + (codebook * self.config.audio_vocab_size) as u32;
+        let e = self
+            .audio_embeddings
+            .index_select(&Tensor::new(&[id], &self.device)?, 0)?;
+        e.unsqueeze(0)?
+            .broadcast_as((rows, 1, e.dim(1)?))?
+            .contiguous()
+    }
+
+    /// CFG-combined logits as an f32 row.
+    fn guide(logits: &Tensor, cfg: Option<f64>) -> Result<Vec<f32>> {
+        let logits = logits.to_dtype(DType::F32)?;
+        let row = match cfg {
+            Some(scale) => {
+                let cond = logits.i(0)?;
+                let uncond = logits.i(1)?;
+                (&uncond + ((cond - &uncond)? * scale)?)?
+            }
+            None => logits.i(0)?,
+        };
+        row.to_vec1::<f32>()
+    }
+
+    /// Generate one frame at position `pos`.
+    ///
+    /// `tokens`/`mask` are the frame's input rows — the whole prompt for
+    /// the first frame, then the previous frame's codes with the text column
+    /// masked out. `muq_idx` marks the reference-audio slot (prompt only).
+    /// `choose(codebook, guided_logits)` picks each code; a sampler for
+    /// generation, the recorded code for a teacher-forced parity check.
+    /// Returns the backbone's last hidden state `(rows, dim)` too.
+    #[allow(clippy::too_many_arguments)]
+    pub fn step(
+        &mut self,
+        tokens: &[[u32; 9]],
+        mask: &[[bool; 9]],
+        pos: usize,
+        muq_idx: Option<usize>,
+        cfg_scale: f64,
+        choose: &mut dyn FnMut(usize, &[f32]) -> u32,
+    ) -> anyhow::Result<([u32; 8], Tensor)> {
+        let cfg = cfg_scale > 1.0;
+        let rows = if cfg { 2 } else { 1 };
+        let guidance = cfg.then_some(cfg_scale);
+        let mut h = self.embed(tokens, mask, cfg)?;
+        if let Some(idx) = muq_idx {
+            let d = self.backbone_flavor.dim;
+            let cond = self.muq_bias.reshape((1, 1, d))?;
+            h = h.slice_assign(&[0..1, idx..idx + 1, 0..d], &cond)?;
+            if cfg {
+                let uncond = self.unconditional_text.reshape((1, 1, d))?;
+                h = h.slice_assign(&[1..2, idx..idx + 1, 0..d], &uncond)?;
+            }
+        }
+        let h = self.backbone.forward(&h, pos)?;
+        let last_h = h.i((.., h.dim(1)? - 1, ..))?.contiguous()?;
+
+        let mut codes = [0u32; 8];
+        let c0 = self.codebook0_head.forward(&last_h)?;
+        codes[0] = choose(0, &Self::guide(&c0, guidance)?);
+
+        self.decoder.reset(DECODER_MAX_SEQ);
+        let mut curr = Tensor::cat(
+            &[last_h.unsqueeze(1)?, self.embed_audio(0, codes[0], rows)?],
+            1,
+        )?;
+        let mut dpos = 0;
+        for (i, slot) in codes.iter_mut().enumerate().skip(1) {
+            let dh = self
+                .decoder
+                .forward(&self.projection.forward(&curr)?, dpos)?;
+            dpos += curr.dim(1)?;
+            let last = dh.i((.., dh.dim(1)? - 1, ..))?.contiguous()?;
+            let logits = last.matmul(&self.audio_head.i(i - 1)?)?;
+            *slot = choose(i, &Self::guide(&logits, guidance)?);
+            curr = self.embed_audio(i, *slot, rows)?;
+        }
+        Ok((codes, last_h))
+    }
+
+    /// Generate up to `max_frames` frames after the first, as heartlib
+    /// does: stop when any code of a frame reaches `audio_eos_id`.
+    pub fn generate(
+        &mut self,
+        prompt: &Prompt,
+        max_frames: usize,
+        cfg_scale: f64,
+        gen_cfg: &GenConfig,
+        sampler: &mut Sampler,
+    ) -> anyhow::Result<Vec<[u32; 8]>> {
+        let plen = prompt.tokens.len();
+        self.reset(plen + max_frames + 1);
+        let prompt_mask: Vec<[bool; 9]> = prompt
+            .tokens
+            .iter()
+            .map(|_| {
+                let mut m = [false; 9];
+                m[8] = true;
+                m
+            })
+            .collect();
+        let mut choose = |_: usize, l: &[f32]| sampler.sample(l);
+        let (mut codes, _) = self.step(
+            &prompt.tokens,
+            &prompt_mask,
+            0,
+            Some(prompt.muq_idx),
+            cfg_scale,
+            &mut choose,
+        )?;
+        let mut frames = vec![codes];
+        let mut frame_mask = [true; 9];
+        frame_mask[8] = false;
+        for i in 0..max_frames {
+            let mut row = [gen_cfg.empty_id; 9];
+            row[..8].copy_from_slice(&codes);
+            codes = self
+                .step(
+                    &[row],
+                    &[frame_mask],
+                    plen + i,
+                    None,
+                    cfg_scale,
+                    &mut choose,
+                )?
+                .0;
+            if codes.iter().any(|&c| c >= gen_cfg.audio_eos_id) {
+                break;
+            }
+            frames.push(codes);
+        }
+        Ok(frames)
+    }
+
+    pub fn dtype(&self) -> DType {
+        self.dtype
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn scaled_theta_matches_the_llama31_rule() {
+        let t = llama3_scaled_theta(128);
+        assert_eq!(t.len(), 64);
+        // Highest frequency (wavelength 2π < 2048): unscaled.
+        assert_eq!(t[0], 1.0);
+        // Lowest frequencies (wavelength > 8192): divided by 32.
+        let raw = 1.0f32 / ROPE_BASE.powf(126.0 / 128.0);
+        assert!(
+            (t[63] - raw / 32.0).abs() < 1e-12,
+            "{} vs {}",
+            t[63],
+            raw / 32.0
+        );
+        // Monotone non-increasing throughout the smooth band.
+        assert!(t.windows(2).all(|w| w[0] >= w[1]));
+    }
+
+    #[test]
+    fn greedy_sampler_is_argmax_and_topk_restricts() {
+        let logits = [0.1f32, 3.0, -1.0, 2.9];
+        let mut greedy = Sampler::new(1.0, 1, 7);
+        for _ in 0..20 {
+            assert_eq!(greedy.sample(&logits), 1);
+        }
+        let mut top2 = Sampler::new(1.0, 2, 7);
+        for _ in 0..200 {
+            let c = top2.sample(&logits);
+            assert!(c == 1 || c == 3, "{c}");
+        }
+        // Same seed, same draws.
+        let mut a = Sampler::new(1.0, 4, 99);
+        let mut b = Sampler::new(1.0, 4, 99);
+        let sa: Vec<u32> = (0..50).map(|_| a.sample(&logits)).collect();
+        let sb: Vec<u32> = (0..50).map(|_| b.sample(&logits)).collect();
+        assert_eq!(sa, sb);
+        // Both of the top two show up over many draws.
+        assert!(sa.contains(&1) || sa.contains(&3));
+    }
+
+    #[test]
+    fn flavours_are_the_torchtune_shapes() {
+        let b = Flavor::named("llama-3B").unwrap();
+        assert_eq!(
+            (b.layers, b.heads, b.kv_heads, b.head_dim()),
+            (28, 24, 8, 128)
+        );
+        let d = Flavor::named("llama-300M").unwrap();
+        assert_eq!(
+            (d.layers, d.heads, d.kv_heads, d.head_dim()),
+            (3, 8, 4, 384)
+        );
+        assert!(Flavor::named("nope").is_err());
+    }
+}

crates/neuron/src/harness/arch/heartmula/mod.rs

@@ -0,0 +1,17 @@
+//! HeartMuLa text-to-song (#379): the MuLa language model ([`lm`], #391),
+//! which turns tags and lyrics into 12.5 Hz frames of eight RVQ codes, and
+//! the HeartCodec decoder ([`codec`], #392), which turns those codes into
+//! 48 kHz stereo audio.
+//!
+//! Ported from heartlib @ 5858ca8f (Apache-2.0). The MuQ-MuLan reference-
+//! audio encoder (CC-BY-NC) is never used: text-to-song conditions on a
+//! zero reference embedding, as heartlib does (#389).
+//!
+//! Parity with heartlib on the CPU in f32 is checked by the ignored tests
+//! in `tests.rs` against the #390 traces in `testdata/heartmula`.
+
+pub mod codec;
+pub mod lm;
+
+#[cfg(test)]
+mod tests;

crates/neuron/src/harness/arch/heartmula/tests.rs

@@ -0,0 +1,360 @@
+//! Parity with heartlib on the CPU in f32, against the #390 traces.
+//!
+//! Ignored by default: they need the checkpoints and minutes of CPU. Lay
+//! the snapshot out as heartlib does —
+//! `<ckpt>/{HeartMuLa-oss-3B,HeartCodec-oss,tokenizer.json,gen_config.json}`
+//! — and run
+//!
+//!   HEARTMULA_CKPT=<ckpt> cargo test --release -p neuron --lib \
+//!     -- --ignored heartmula --nocapture --test-threads=1
+//!
+//! `HEARTMULA_CASES` narrows the cases (comma-separated).
+
+use super::codec::{self, HeartCodec};
+use super::lm::{self, HeartMuLa};
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+
+fn ckpt() -> Option<PathBuf> {
+    match std::env::var("HEARTMULA_CKPT") {
+        Ok(p) => Some(p.into()),
+        Err(_) => {
+            eprintln!("HEARTMULA_CKPT not set; skipping");
+            None
+        }
+    }
+}
+
+fn cases(default: &[&str]) -> Vec<String> {
+    match std::env::var("HEARTMULA_CASES") {
+        Ok(c) => c.split(',').map(str::to_string).collect(),
+        Err(_) => default.iter().map(|s| s.to_string()).collect(),
+    }
+}
+
+fn fixture_dir() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/heartmula")
+}
+
+/// Named tensors from a trace. The files also hold U8 masks, which are
+/// skipped (only the requested names are read).
+fn load(case: &str, names: &[&str]) -> HashMap<String, Tensor> {
+    let path = fixture_dir().join(format!("{case}__cpu_f32.safetensors"));
+    let file = unsafe { candle_core::safetensors::MmapedSafetensors::new(path) }.unwrap();
+    names
+        .iter()
+        .map(|n| (n.to_string(), file.load(n, &Device::Cpu).unwrap()))
+        .collect()
+}
+
+fn rel_l2(a: &Tensor, b: &Tensor) -> f32 {
+    let a = a.to_dtype(DType::F32).unwrap().flatten_all().unwrap();
+    let b = b.to_dtype(DType::F32).unwrap().flatten_all().unwrap();
+    let d = (&a - &b).unwrap().sqr().unwrap().sum_all().unwrap();
+    let n = b.sqr().unwrap().sum_all().unwrap();
+    (d.to_scalar::<f32>().unwrap() / n.to_scalar::<f32>().unwrap().max(1e-30)).sqrt()
+}
+
+fn shards(dir: &Path) -> Vec<PathBuf> {
+    let mut v: Vec<PathBuf> = std::fs::read_dir(dir)
+        .unwrap()
+        .filter_map(|e| e.ok().map(|e| e.path()))
+        .filter(|p| p.extension().is_some_and(|x| x == "safetensors"))
+        .collect();
+    v.sort();
+    v
+}
+
+fn read_json<T: serde::de::DeserializeOwned>(p: &Path) -> T {
+    serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
+}
+
+/// Teacher-forced LM parity: the prompt matrix exactly, then frames 0 and
+/// 1 with the recorded codes fed back, comparing every codebook's guided
+/// logits and the backbone's last hidden state.
+#[test]
+#[ignore = "needs HEARTMULA_CKPT and minutes of CPU"]
+fn heartmula_lm_matches_reference_cpu_f32() {
+    let Some(root) = ckpt() else { return };
+    let mula_dir = root.join("HeartMuLa-oss-3B");
+    let config: lm::Config = read_json(&mula_dir.join("config.json"));
+    let gen_cfg: lm::GenConfig = read_json(&root.join("gen_config.json"));
+    let tokenizer = tokenizers::Tokenizer::from_file(root.join("tokenizer.json")).unwrap();
+    let t = std::time::Instant::now();
+    let vb = unsafe {
+        VarBuilder::from_mmaped_safetensors(&shards(&mula_dir), DType::F32, &Device::Cpu).unwrap()
+    };
+    let mut model = HeartMuLa::new(config, vb).unwrap();
+    eprintln!("LM loaded in {:.1}s", t.elapsed().as_secs_f32());
+
+    for case in cases(&["instr_12s", "vocal_en_12s", "vocal_es_12s"]) {
+        let meta: serde_json::Value =
+            read_json(&fixture_dir().join(format!("{case}__cpu_f32.json")));
+        let c = &meta["case"];
+        let prompt = lm::build_prompt(
+            &tokenizer,
+            c["tags"].as_str().unwrap(),
+            c["lyrics"].as_str().unwrap(),
+            &gen_cfg,
+        )
+        .unwrap();
+        let want = load(
+            &case,
+            &[
+                "prompt_tokens",
+                "lm_greedy_frames",
+                "lm_greedy_logits",
+                "lm_greedy_backbone_last_h",
+                "lm_sampled_frames",
+                "lm_sampled_logits",
+                "lm_sampled_backbone_last_h",
+            ],
+        );
+        let rows: Vec<Vec<i64>> = want["prompt_tokens"].to_vec2().unwrap();
+        let got: Vec<Vec<i64>> = prompt
+            .tokens
+            .iter()
+            .map(|r| r.iter().map(|&x| x as i64).collect())
+            .collect();
+        assert_eq!(got, rows, "{case}: prompt tokens");
+        assert_eq!(prompt.muq_idx, meta["muq_idx"].as_u64().unwrap() as usize);
+
+        for sampler in ["greedy", "sampled"] {
+            let frames: Vec<Vec<i64>> = want[&format!("lm_{sampler}_frames")].to_vec2().unwrap();
+            let logits_ref = &want[&format!("lm_{sampler}_logits")];
+            let h_ref = &want[&format!("lm_{sampler}_backbone_last_h")];
+            let cfg = meta["lm"][sampler]["cfg_scale"].as_f64().unwrap();
+            let plen = prompt.tokens.len();
+            model.reset(plen + 4);
+            let t = std::time::Instant::now();
+            let mut worst = 0f32;
+            for f in 0..2usize {
+                let mut got_logits: Vec<Vec<f32>> = Vec::new();
+                let mut choose = |cb: usize, l: &[f32]| {
+                    got_logits.push(l.to_vec());
+                    frames[cb][f] as u32
+                };
+                let (codes, last_h) = if f == 0 {
+                    let mask: Vec<[bool; 9]> = prompt
+                        .tokens
+                        .iter()
+                        .map(|_| {
+                            let mut m = [false; 9];
+                            m[8] = true;
+                            m
+                        })
+                        .collect();
+                    model
+                        .step(
+                            &prompt.tokens,
+                            &mask,
+                            0,
+                            Some(prompt.muq_idx),
+                            cfg,
+                            &mut choose,
+                        )
+                        .unwrap()
+                } else {
+                    let mut row = [gen_cfg.empty_id; 9];
+                    for (cb, slot) in row.iter_mut().take(8).enumerate() {
+                        *slot = frames[cb][0] as u32;
+                    }
+                    let mut m = [true; 9];
+                    m[8] = false;
+                    model
+                        .step(&[row], &[m], plen, None, cfg, &mut choose)
+                        .unwrap()
+                };
+                for (cb, &code) in codes.iter().enumerate() {
+                    assert_eq!(code as i64, frames[cb][f]);
+                }
+                for (cb, l) in got_logits.iter().enumerate() {
+                    let ours = Tensor::new(l.as_slice(), &Device::Cpu).unwrap();
+                    let theirs = logits_ref.get(f).unwrap().get(cb).unwrap();
+                    let e = rel_l2(&ours, &theirs);
+                    worst = worst.max(e);
+                    assert!(
+                        e < 1e-4,
+                        "{case}/{sampler} frame {f} cb {cb}: logits rel L2 {e}"
+                    );
+                }
+                let e = rel_l2(&last_h, &h_ref.get(f).unwrap());
+                worst = worst.max(e);
+                assert!(e < 1e-4, "{case}/{sampler} frame {f}: last_h rel L2 {e}");
+            }
+            eprintln!(
+                "{case}/{sampler}: worst rel L2 {worst:.2e} over 2 frames × 8 codebooks ({:.1}s)",
+                t.elapsed().as_secs_f32()
+            );
+        }
+    }
+}
+
+/// Codec parity from the recorded greedy codes and noise: the RVQ sum, the
+/// flow condition, the velocity at step 1, the state after step 5, the
+/// final latents, and the first 3 s of waveform.
+#[test]
+#[ignore = "needs HEARTMULA_CKPT and minutes of CPU"]
+fn heartmula_codec_matches_reference_cpu_f32() {
+    let Some(root) = ckpt() else { return };
+    let dir = root.join("HeartCodec-oss");
+    let config: codec::Config = read_json(&dir.join("config.json"));
+    let vb = unsafe {
+        VarBuilder::from_mmaped_safetensors(&shards(&dir), DType::F32, &Device::Cpu).unwrap()
+    };
+    let model = HeartCodec::new(config, vb).unwrap();
+
+    for case in cases(&["instr_12s", "vocal_en_12s"]) {
+        let want = load(
+            &case,
+            &[
+                "codec_codes",
+                "codec_rvq_quantized",
+                "codec_flow_mu",
+                "codec_flow_noise",
+                "codec_flow_t_span",
+                "codec_flow_step1_v",
+                "codec_flow_step5_x",
+                "codec_flow_final",
+                "codec_waveform",
+            ],
+        );
+        let codes = want["codec_codes"].to_dtype(DType::U32).unwrap();
+        let frames = codes.dim(1).unwrap();
+        let tiled = codec::tile_codes(&codes).unwrap();
+        let t = std::time::Instant::now();
+        let (lat, trace) = model
+            .segment_latents(
+                &tiled.narrow(1, 0, 372).unwrap(),
+                &want["codec_flow_noise"],
+                None,
+                codec::DEFAULT_STEPS,
+                codec::DEFAULT_GUIDANCE,
+            )
+            .unwrap();
+        let flow_s = t.elapsed().as_secs_f32();
+        let t_span: Vec<f32> = want["codec_flow_t_span"].to_vec1().unwrap();
+        assert_eq!(trace.t_span, t_span, "{case}: t_span");
+        let checks = [
+            ("rvq", rel_l2(&trace.rvq, &want["codec_rvq_quantized"])),
+            ("mu", rel_l2(&trace.mu, &want["codec_flow_mu"])),
+            (
+                "step1_v",
+                rel_l2(&trace.step1_v, &want["codec_flow_step1_v"]),
+            ),
+            (
+                "step5_x",
+                rel_l2(&trace.step5_x, &want["codec_flow_step5_x"]),
+            ),
+            ("final", rel_l2(&lat, &want["codec_flow_final"])),
+        ];
+
+        let t = std::time::Instant::now();
+        let wav = model.latents_to_audio(&lat, Some(128)).unwrap();
+        let decode_s = t.elapsed().as_secs_f32();
+        assert_eq!(wav.dim(1).unwrap(), 744 * 1920);
+        let target = frames * 48_000 * 2 / 25;
+        assert!(target <= wav.dim(1).unwrap());
+        let wref = &want["codec_waveform"];
+        let n = wref.dim(1).unwrap();
+        let wave = rel_l2(&wav.narrow(1, 0, n).unwrap(), wref);
+
+        eprintln!(
+            "{case}: {} | waveform {wave:.2e} | flow {flow_s:.1}s, decode {decode_s:.1}s",
+            checks
+                .iter()
+                .map(|(n, e)| format!("{n} {e:.2e}"))
+                .collect::<Vec<_>>()
+                .join(", ")
+        );
+        for (name, e) in checks {
+            assert!(e < 1e-4, "{case}: {name} rel L2 {e}");
+        }
+        assert!(wave < 1e-3, "{case}: waveform rel L2 {wave}");
+    }
+}
+
+/// The scalar decoder alone, from the reference's final latents, whole and
+/// chunked: isolates the waveform stage from the flow.
+#[test]
+#[ignore = "needs HEARTMULA_CKPT and minutes of CPU"]
+fn heartmula_scalar_decoder_matches_reference_cpu_f32() {
+    let Some(root) = ckpt() else { return };
+    let dir = root.join("HeartCodec-oss");
+    let config: codec::Config = read_json(&dir.join("config.json"));
+    let vb = unsafe {
+        VarBuilder::from_mmaped_safetensors(&shards(&dir), DType::F32, &Device::Cpu).unwrap()
+    };
+    let model = HeartCodec::new(config, vb).unwrap();
+    for case in cases(&["instr_12s"]) {
+        let want = load(&case, &["codec_flow_final", "codec_waveform"]);
+        let wref = &want["codec_waveform"];
+        let n = wref.dim(1).unwrap();
+        for chunk in [None, Some(128)] {
+            let t = std::time::Instant::now();
+            let wav = model
+                .latents_to_audio(&want["codec_flow_final"], chunk)
+                .unwrap();
+            let e = rel_l2(&wav.narrow(1, 0, n).unwrap(), wref);
+            eprintln!(
+                "{case} chunk {chunk:?}: waveform rel L2 {e:.2e} ({:.1}s)",
+                t.elapsed().as_secs_f32()
+            );
+            assert!(e < 1e-3, "{case} chunk {chunk:?}: waveform rel L2 {e}");
+        }
+    }
+}
+
+/// Free-running greedy generation: frame 0 must equal the reference's
+/// (nothing has been sampled yet), later frames may drift on near-ties
+/// (#390). Reports the per-frame time.
+#[test]
+#[ignore = "needs HEARTMULA_CKPT and minutes of CPU"]
+fn heartmula_generates_greedily() {
+    let Some(root) = ckpt() else { return };
+    let mula_dir = root.join("HeartMuLa-oss-3B");
+    let config: lm::Config = read_json(&mula_dir.join("config.json"));
+    let gen_cfg: lm::GenConfig = read_json(&root.join("gen_config.json"));
+    let tokenizer = tokenizers::Tokenizer::from_file(root.join("tokenizer.json")).unwrap();
+    let vb = unsafe {
+        VarBuilder::from_mmaped_safetensors(&shards(&mula_dir), DType::F32, &Device::Cpu).unwrap()
+    };
+    let mut model = HeartMuLa::new(config, vb).unwrap();
+    for case in cases(&["vocal_en_12s"]) {
+        let meta: serde_json::Value =
+            read_json(&fixture_dir().join(format!("{case}__cpu_f32.json")));
+        let c = &meta["case"];
+        let prompt = lm::build_prompt(
+            &tokenizer,
+            c["tags"].as_str().unwrap(),
+            c["lyrics"].as_str().unwrap(),
+            &gen_cfg,
+        )
+        .unwrap();
+        let want: Vec<Vec<i64>> = load(&case, &["lm_greedy_frames"])["lm_greedy_frames"]
+            .to_vec2()
+            .unwrap();
+        let mut sampler = lm::Sampler::new(1.0, 1, 0);
+        let extra = 24;
+        let t = std::time::Instant::now();
+        let frames = model
+            .generate(&prompt, extra, 1.5, &gen_cfg, &mut sampler)
+            .unwrap();
+        let secs = t.elapsed().as_secs_f32();
+        assert_eq!(frames.len(), extra + 1);
+        let same = frames
+            .iter()
+            .enumerate()
+            .take_while(|(f, codes)| (0..8).all(|cb| codes[cb] as i64 == want[cb][*f]))
+            .count();
+        eprintln!(
+            "{case}: {} frames ({:.2}s of audio) in {secs:.1}s, {:.2}s/frame; first {same} frames equal the reference",
+            frames.len(),
+            frames.len() as f64 / 12.5,
+            secs / frames.len() as f32
+        );
+        assert!(same >= 1, "{case}: frame 0 differs from the reference");
+    }
+}

crates/neuron/src/harness/arch/mod.rs

@@ -20,6 +20,8 @@
 //! TP-aware analogues live in `harness/tp/tp_<family>.rs` and follow
 //! the pattern set by `tp_qwen3.rs`.
 
+pub mod acestep;
+pub mod heartmula;
 pub mod laya;
 pub mod phase_probe;
 pub mod qwen3_5;

crates/neuron/src/harness/audio.rs

@@ -0,0 +1,902 @@
+//! Music generation engines behind the device worker: ACE-Step 1.5 (#383)
+//! and HeartMuLa (#393).
+//!
+//! The ACE-Step repo carries every component as a subfolder of one
+//! snapshot — the turbo DiT (with its `silence_latent.pt`), the song
+//! planner LM, the Qwen3 text encoder and the VAE — so the loader takes
+//! the snapshot root and finds them by name.
+//!
+//! Same discipline as the image engine (`image.rs`): the pipeline lives in
+//! the worker's audio slab, one `Job::GenerateAudio` runs a whole song, and
+//! the reply is CPU-side samples. WAV encoding happens on the async side.
+
+use super::arch::acestep::pipeline::{AceStepPipeline, GenerationParams, LmMetadata};
+use super::arch::acestep::{AceStepConfig, VaeConfig, lm::LmConfig, qwen3, sampling};
+use anyhow::{Context, Result};
+use candle_core::{DType, Device};
+use candle_nn::VarBuilder;
+use std::path::{Path, PathBuf};
+use std::time::Instant;
+
+/// The turbo DiT's subfolder.
+pub const DIT_DIR: &str = "acestep-v15-turbo";
+/// The text encoder's subfolder.
+pub const TEXT_ENCODER_DIR: &str = "Qwen3-Embedding-0.6B";
+/// The VAE's subfolder.
+pub const VAE_DIR: &str = "vae";
+/// Planner LMs in order of preference; the first the snapshot has is used.
+pub const PLANNER_DIRS: [&str; 2] = ["acestep-5Hz-lm-1.7B", "acestep-5Hz-lm-0.6B"];
+
+/// What ACE-Step 1.5 turbo serves (#383). Upstream allows 10 minutes; the
+/// planner's code budget (3072 tokens at 5 Hz) covers that. No `steps`
+/// knob: the turbo schedule is a fixed eight steps.
+pub const ACESTEP_LIMITS: cortex_core::audio::AudioLimits = cortex_core::audio::AudioLimits {
+    min_duration: 10.0,
+    max_duration: 600.0,
+    default_duration: 30.0,
+    vocals: true,
+    sample_rate: 48_000,
+    channels: 2,
+    max_steps: None,
+};
+
+/// The turbo schedule's default step count.
+pub const DEFAULT_STEPS: usize = 8;
+/// Turbo models are trained with timestep shift 3.
+const TURBO_SHIFT: f64 = 3.0;
+/// Upstream's default output peak, in dBFS.
+const PEAK_DB: f64 = -1.0;
+
+/// Files every ACE-Step snapshot needs, relative to its root.
+pub fn required_files() -> Vec<String> {
+    vec![
+        format!("{DIT_DIR}/config.json"),
+        format!("{DIT_DIR}/model.safetensors"),
+        format!("{DIT_DIR}/silence_latent.pt"),
+        format!("{TEXT_ENCODER_DIR}/config.json"),
+        format!("{TEXT_ENCODER_DIR}/model.safetensors"),
+        format!("{TEXT_ENCODER_DIR}/tokenizer.json"),
+        format!("{VAE_DIR}/config.json"),
+        format!("{VAE_DIR}/diffusion_pytorch_model.safetensors"),
+    ]
+}
+
+/// Files a planner LM needs, relative to the snapshot root.
+pub fn planner_files(dir: &str) -> Vec<String> {
+    ["config.json", "model.safetensors", "tokenizer.json"]
+        .iter()
+        .map(|f| format!("{dir}/{f}"))
+        .collect()
+}
+
+/// True for a repo laid out like ACE-Step 1.5: the turbo DiT and the VAE as
+/// subfolders. Checked before the generic safetensors sniff, which would
+/// otherwise take the root `config.json` for a text model.
+pub fn is_acestep_layout(filenames: &[&str]) -> bool {
+    filenames.contains(&format!("{DIT_DIR}/config.json").as_str())
+        && filenames.contains(&format!("{VAE_DIR}/diffusion_pytorch_model.safetensors").as_str())
+}
+
+/// A resolved snapshot: the root directory every component sits under, and
+/// which planner (if any) it has.
+#[derive(Debug, Clone)]
+pub struct AceStepFiles {
+    pub root: PathBuf,
+    pub planner: Option<String>,
+}
+
+/// One generation's parameters, already validated against the model's limits.
+#[derive(Debug, Clone)]
+pub struct AudioGenParams {
+    pub prompt: String,
+    /// Empty for an instrumental.
+    pub lyrics: String,
+    pub duration_secs: f64,
+    pub seed: u64,
+    pub steps: Option<usize>,
+    pub guidance_scale: Option<f64>,
+    /// Run the song planner first (when the model has one).
+    pub planner: bool,
+    pub temperature: Option<f64>,
+    /// Lyric language code; `None` lets the planner (or `"unknown"`) decide.
+    pub language: Option<String>,
+    /// Top-k sampling (autoregressive models).
+    pub top_k: Option<usize>,
+}
+
+/// One generation's output: CPU-side stereo samples in `[-1, 1]`.
+#[derive(Debug, Clone)]
+pub struct AudioGenResult {
+    pub left: Vec<f32>,
+    pub right: Vec<f32>,
+    pub sample_rate: u32,
+    pub timing: cortex_core::audio::AudioTiming,
+}
+
+impl AudioGenResult {
+    pub fn duration_secs(&self) -> f64 {
+        self.left.len() as f64 / self.sample_rate.max(1) as f64
+    }
+}
+
+/// The loaded pipeline, owned by the device worker.
+pub struct AceStepEngine {
+    pipeline: AceStepPipeline,
+    has_planner: bool,
+}
+
+fn read_json<T: serde::de::DeserializeOwned>(path: &Path) -> Result<T> {
+    let f = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
+    serde_json::from_reader(f).with_context(|| format!("parse {}", path.display()))
+}
+
+impl AceStepEngine {
+    /// Build every component on `device`: bf16 on a GPU, f32 on the CPU.
+    pub fn load(files: &AceStepFiles, device: &Device) -> Result<Self> {
+        let dtype = if device.is_cpu() {
+            DType::F32
+        } else {
+            DType::BF16
+        };
+        let root = &files.root;
+        let dit = root.join(DIT_DIR);
+        let te = root.join(TEXT_ENCODER_DIR);
+        let vae = root.join(VAE_DIR);
+
+        let dit_config: AceStepConfig = read_json(&dit.join("config.json"))?;
+        let vae_config: VaeConfig = read_json(&vae.join("config.json"))?;
+        let text_config: qwen3::Config = read_json(&te.join("config.json"))?;
+
+        // SAFETY: the snapshot is a read-only cache the daemon owns; the
+        // files are not modified while mapped.
+        let dit_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[dit.join("model.safetensors")], dtype, device)?
+        };
+        let vae_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(
+                &[vae.join("diffusion_pytorch_model.safetensors")],
+                dtype,
+                device,
+            )?
+        };
+        // The embedding model's tensors carry a `model.` prefix the bare
+        // qwen3 encoder does not expect.
+        let text_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[te.join("model.safetensors")], dtype, device)?
+        }
+        .rename_f(|n| n.strip_prefix("model.").unwrap_or(n).to_string());
+
+        let silence_latent = super::arch::acestep::pipeline::load_silence_latent(
+            &dit.join("silence_latent.pt"),
+            dit_config.audio_acoustic_hidden_dim,
+            dtype,
+            device,
+        )?;
+        let mut pipeline = AceStepPipeline::new(
+            &dit_config,
+            &vae_config,
+            &text_config,
+            dit_vb,
+            vae_vb,
+            text_vb,
+            &te.join("tokenizer.json"),
+            silence_latent,
+            false,
+        )?;
+
+        // The planner is optional and runs only when asked for, so a planner
+        // that will not load leaves a working engine without one rather
+        // than failing the model.
+        let has_planner = match &files.planner {
+            Some(dir) => match Self::load_planner(&mut pipeline, &root.join(dir), dtype, device) {
+                Ok(()) => true,
+                Err(e) => {
+                    tracing::warn!(planner = %dir, error = %format!("{e:#}"), "song planner failed to load; serving without it");
+                    false
+                }
+            },
+            None => false,
+        };
+        Ok(Self {
+            pipeline,
+            has_planner,
+        })
+    }
+
+    pub fn has_planner(&self) -> bool {
+        self.has_planner
+    }
+
+    /// Load the planner LM into the pipeline. Its checkpoint stores the
+    /// decoder's tensors without the `model.` prefix the causal-LM wrapper
+    /// asks for, and ties the output head to the embeddings.
+    fn load_planner(
+        pipeline: &mut AceStepPipeline,
+        dir: &Path,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<()> {
+        let lm_config: qwen3::Config = read_json(&dir.join("config.json"))?;
+        let lm_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[dir.join("model.safetensors")], dtype, device)?
+        }
+        .rename_f(|n| n.strip_prefix("model.").unwrap_or(n).to_string());
+        pipeline.load_lm(&lm_config, lm_vb, &dir.join("tokenizer.json"))?;
+        Ok(())
+    }
+
+    /// Render one song.
+    pub fn generate(&mut self, p: &AudioGenParams) -> Result<AudioGenResult> {
+        let mut timing = cortex_core::audio::AudioTiming::default();
+
+        let (lm_hints, lm_metadata) = if p.planner && self.has_planner {
+            let t = Instant::now();
+            let cfg = LmConfig {
+                temperature: p.temperature.unwrap_or(0.85),
+                seed: p.seed,
+                ..Default::default()
+            };
+            let (hints, meta) =
+                self.pipeline
+                    .generate_lm_hints(&p.prompt, &p.lyrics, p.duration_secs, &cfg)?;
+            timing.plan_ms = t.elapsed().as_millis() as u64;
+            let metadata = LmMetadata {
+                bpm: meta.get("bpm").cloned(),
+                timesignature: meta.get("timesignature").cloned(),
+                keyscale: meta.get("keyscale").cloned(),
+                caption: meta.get("caption").cloned(),
+                language: meta
+                    .get("language")
+                    .or_else(|| meta.get("vocal_language"))
+                    .cloned(),
+            };
+            (Some(hints), Some(metadata))
+        } else {
+            (None, None)
+        };
+
+        let params = GenerationParams {
+            duration_secs: p.duration_secs,
+            seed: Some(p.seed),
+            num_steps: Some(p.steps.unwrap_or(DEFAULT_STEPS)),
+            guidance_scale: p
+                .guidance_scale
+                .unwrap_or(GenerationParams::default().guidance_scale),
+            shift: TURBO_SHIFT,
+            is_turbo: true,
+            lm_hints_25hz: lm_hints,
+            lm_metadata,
+            infer_method: sampling::InferMethod::Ode,
+            language: p.language.clone(),
+            ..Default::default()
+        };
+        let t = Instant::now();
+        let out = self.pipeline.text2music(&p.prompt, &p.lyrics, &params)?;
+        let audio = out.audio.to_dtype(DType::F32)?.to_device(&Device::Cpu)?;
+        timing.generate_ms = t.elapsed().as_millis() as u64;
+
+        let mut left = audio.get(0)?.to_vec1::<f32>()?;
+        let mut right = audio.get(1)?.to_vec1::<f32>()?;
+        if left.iter().chain(&right).any(|s| !s.is_finite()) {
+            anyhow::bail!("generation produced non-finite samples");
+        }
+        peak_normalize(&mut left, &mut right, PEAK_DB);
+        Ok(AudioGenResult {
+            left,
+            right,
+            sample_rate: out.sample_rate,
+            timing,
+        })
+    }
+}
+
+// ── HeartMuLa (#393) ─────────────────────────────────────────────────
+
+/// What HeartMuLa oss-3B serves: songs of 10 s to 4 min, as the evaluation
+/// beta did. It is autoregressive (12.5 frames per second), so there is no
+/// step knob.
+pub const HEARTMULA_LIMITS: cortex_core::audio::AudioLimits = cortex_core::audio::AudioLimits {
+    min_duration: 10.0,
+    max_duration: 240.0,
+    default_duration: 30.0,
+    vocals: true,
+    sample_rate: 48_000,
+    channels: 2,
+    max_steps: None,
+};
+
+/// The codec and the generation config live in sibling repos of the LM.
+pub const HEARTCODEC_REPO: &str = "HeartMuLa/HeartCodec-oss-20260123";
+pub const HEARTMULA_GEN_REPO: &str = "HeartMuLa/HeartMuLaGen";
+
+/// heartlib's defaults, as the beta used them.
+const HEART_CFG: f64 = 1.5;
+const HEART_TEMPERATURE: f64 = 1.0;
+const HEART_TOP_K: usize = 50;
+/// Section tags with no words: how an instrumental is asked of HeartMuLa.
+const HEART_INSTRUMENTAL: &str =
+    "[Intro]\n\n[Verse]\n\n[Chorus]\n\n[Bridge]\n\n[Chorus]\n\n[Outro]";
+/// Latent frames per scalar-decoder chunk: bounds decode memory at any length.
+const HEART_DECODE_CHUNK: usize = 256;
+
+/// True when a `config.json` describes a HeartMuLa LM. Its weight files
+/// look like any dense LLM's, so the config is what tells them apart.
+pub fn is_heartmula_config(config_json: &str) -> bool {
+    serde_json::from_str::<serde_json::Value>(config_json)
+        .ok()
+        .and_then(|v| {
+            v.get("model_type")
+                .and_then(|m| m.as_str())
+                .map(str::to_string)
+        })
+        .is_some_and(|t| t == "heartmula")
+}
+
+/// A resolved HeartMuLa model: the LM, the codec, and the generation files.
+#[derive(Debug, Clone)]
+pub struct HeartMuLaFiles {
+    pub lm_dir: PathBuf,
+    pub codec_dir: PathBuf,
+    pub tokenizer: PathBuf,
+    pub gen_config: PathBuf,
+}
+
+/// HeartMuLa's LM and codec, owned by the device worker. Both stay
+/// resident: the LM in bf16 on a GPU (f32 on the CPU), the codec in f32 as
+/// heartlib recommends.
+pub struct HeartMuLaEngine {
+    lm: super::arch::heartmula::lm::HeartMuLa,
+    codec: super::arch::heartmula::codec::HeartCodec,
+    tokenizer: tokenizers::Tokenizer,
+    gen_cfg: super::arch::heartmula::lm::GenConfig,
+}
+
+fn safetensors_in(dir: &Path) -> Result<Vec<PathBuf>> {
+    let mut v: Vec<PathBuf> = std::fs::read_dir(dir)
+        .with_context(|| format!("read {}", dir.display()))?
+        .filter_map(|e| e.ok().map(|e| e.path()))
+        .filter(|p| p.extension().is_some_and(|x| x == "safetensors"))
+        .collect();
+    v.sort();
+    anyhow::ensure!(!v.is_empty(), "no safetensors in {}", dir.display());
+    Ok(v)
+}
+
+impl HeartMuLaEngine {
+    pub fn load(files: &HeartMuLaFiles, device: &Device) -> Result<Self> {
+        use super::arch::heartmula::{codec, lm};
+        let lm_dtype = if device.is_cpu() {
+            DType::F32
+        } else {
+            DType::BF16
+        };
+        let lm_config: lm::Config = read_json(&files.lm_dir.join("config.json"))?;
+        // SAFETY: read-only snapshot files the daemon owns.
+        let lm_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&safetensors_in(&files.lm_dir)?, lm_dtype, device)?
+        };
+        let lm = lm::HeartMuLa::new(lm_config, lm_vb)?;
+        let codec_config: codec::Config = read_json(&files.codec_dir.join("config.json"))?;
+        let codec_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(
+                &safetensors_in(&files.codec_dir)?,
+                DType::F32,
+                device,
+            )?
+        };
+        let codec = codec::HeartCodec::new(codec_config, codec_vb)?;
+        let tokenizer = tokenizers::Tokenizer::from_file(&files.tokenizer)
+            .map_err(|e| anyhow::anyhow!("load tokenizer: {e}"))?;
+        let gen_cfg: lm::GenConfig = read_json(&files.gen_config)?;
+        Ok(Self {
+            lm,
+            codec,
+            tokenizer,
+            gen_cfg,
+        })
+    }
+
+    /// Render one song: tags and lyrics through the LM, then the codec.
+    pub fn generate(&mut self, p: &AudioGenParams) -> Result<AudioGenResult> {
+        use super::arch::heartmula::{codec, lm};
+        let mut timing = cortex_core::audio::AudioTiming::default();
+        let lyrics = if p.lyrics.trim().is_empty() {
+            HEART_INSTRUMENTAL
+        } else {
+            p.lyrics.as_str()
+        };
+        let t = Instant::now();
+        let prompt = lm::build_prompt(&self.tokenizer, &p.prompt, lyrics, &self.gen_cfg)?;
+        timing.encode_ms = t.elapsed().as_millis() as u64;
+
+        let t = Instant::now();
+        let max_frames = (p.duration_secs * codec::CODE_RATE).round() as usize;
+        let mut sampler = lm::Sampler::new(
+            p.temperature.unwrap_or(HEART_TEMPERATURE),
+            p.top_k.unwrap_or(HEART_TOP_K),
+            p.seed,
+        );
+        let frames = self.lm.generate(
+            &prompt,
+            max_frames,
+            p.guidance_scale.unwrap_or(HEART_CFG),
+            &self.gen_cfg,
+            &mut sampler,
+        )?;
+        timing.generate_ms = t.elapsed().as_millis() as u64;
+        anyhow::ensure!(!frames.is_empty(), "the model produced no audio frames");
+
+        let t = Instant::now();
+        let n = frames.len();
+        let flat: Vec<u32> = (0..8)
+            .flat_map(|q| frames.iter().map(move |f| f[q]))
+            .collect();
+        let codes = candle_core::Tensor::from_vec(flat, (8, n), self.codec.device())?;
+        let device = self.codec.device().clone();
+        let out_channels = self.codec.config.out_channels;
+        let mut segment = 0u64;
+        let seed = p.seed;
+        let mut noise = |len: usize| {
+            segment += 1;
+            super::arch::acestep::model::seeded_normal(
+                Some(seed.wrapping_add(segment)),
+                (1, len, out_channels),
+            )?
+            .to_device(&device)
+        };
+        let audio = self.codec.detokenize(
+            &codes,
+            codec::DEFAULT_STEPS,
+            codec::DEFAULT_GUIDANCE,
+            &mut noise,
+            Some(HEART_DECODE_CHUNK),
+        )?;
+        let audio = audio.to_dtype(DType::F32)?.to_device(&Device::Cpu)?;
+        timing.decode_ms = t.elapsed().as_millis() as u64;
+
+        let mut left = audio.get(0)?.to_vec1::<f32>()?;
+        let mut right = audio.get(1)?.to_vec1::<f32>()?;
+        if left.iter().chain(&right).any(|s| !s.is_finite()) {
+            anyhow::bail!("generation produced non-finite samples");
+        }
+        for s in left.iter_mut().chain(right.iter_mut()) {
+            *s = s.clamp(-1.0, 1.0);
+        }
+        Ok(AudioGenResult {
+            left,
+            right,
+            sample_rate: self.codec.config.sample_rate as u32,
+            timing,
+        })
+    }
+}
+
+// ── Engine selection ─────────────────────────────────────────────────
+
+/// What the worker is asked to load.
+#[derive(Debug, Clone)]
+pub enum AudioFiles {
+    AceStep(AceStepFiles),
+    HeartMuLa(HeartMuLaFiles),
+}
+
+/// A loaded music engine in the worker's audio slab.
+pub enum AudioEngine {
+    AceStep(Box<AceStepEngine>),
+    HeartMuLa(Box<HeartMuLaEngine>),
+}
+
+impl AudioEngine {
+    pub fn load(files: &AudioFiles, device: &Device) -> Result<Self> {
+        Ok(match files {
+            AudioFiles::AceStep(f) => Self::AceStep(Box::new(AceStepEngine::load(f, device)?)),
+            AudioFiles::HeartMuLa(f) => {
+                Self::HeartMuLa(Box::new(HeartMuLaEngine::load(f, device)?))
+            }
+        })
+    }
+
+    pub fn generate(&mut self, p: &AudioGenParams) -> Result<AudioGenResult> {
+        match self {
+            Self::AceStep(e) => e.generate(p),
+            Self::HeartMuLa(e) => e.generate(p),
+        }
+    }
+
+    pub fn has_planner(&self) -> bool {
+        match self {
+            Self::AceStep(e) => e.has_planner(),
+            Self::HeartMuLa(_) => false,
+        }
+    }
+
+    pub fn name(&self) -> &'static str {
+        match self {
+            Self::AceStep(_) => "ace-step",
+            Self::HeartMuLa(_) => "heartmula",
+        }
+    }
+}
+
+/// Scale both channels so the loudest sample sits at `target_db` dBFS, as
+/// upstream does before saving. Near-silence is left alone rather than
+/// amplified into noise.
+pub fn peak_normalize(left: &mut [f32], right: &mut [f32], target_db: f64) {
+    let peak = left
+        .iter()
+        .chain(right.iter())
+        .fold(0f32, |m, s| m.max(s.abs()));
+    if peak < 1e-6 {
+        return;
+    }
+    let gain = (10f64.powf(target_db / 20.0) / peak as f64) as f32;
+    for s in left.iter_mut().chain(right.iter_mut()) {
+        *s = (*s * gain).clamp(-1.0, 1.0);
+    }
+}
+
+/// Encode stereo samples as 16-bit FLAC. Lossless at the 16-bit depth the
+/// samples are quantised to, and roughly half the size of the same PCM.
+pub fn encode_flac_pcm16(left: &[f32], right: &[f32], sample_rate: u32) -> anyhow::Result<Vec<u8>> {
+    use flacenc::component::BitRepr;
+    use flacenc::error::Verify;
+
+    let frames = left.len().min(right.len());
+    let pcm = |s: f32| (s.clamp(-1.0, 1.0) * i16::MAX as f32).round() as i32;
+    let mut interleaved = Vec::with_capacity(frames * 2);
+    for i in 0..frames {
+        interleaved.push(pcm(left[i]));
+        interleaved.push(pcm(right[i]));
+    }
+    let config = flacenc::config::Encoder::default()
+        .into_verified()
+        .map_err(|(_, e)| anyhow::anyhow!("flac encoder config: {e:?}"))?;
+    let source =
+        flacenc::source::MemSource::from_samples(&interleaved, 2, 16, sample_rate as usize);
+    let stream = flacenc::encode_with_fixed_block_size(&config, source, config.block_size)
+        .map_err(|e| anyhow::anyhow!("flac encode: {e:?}"))?;
+    let mut sink = flacenc::bitsink::ByteSink::new();
+    stream
+        .write(&mut sink)
+        .map_err(|e| anyhow::anyhow!("flac write: {e:?}"))?;
+    Ok(sink.into_inner())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn heartmula_is_recognised_by_its_config() {
+        assert!(is_heartmula_config(
+            r#"{"architectures": ["HeartMuLa"], "model_type": "heartmula"}"#
+        ));
+        assert!(!is_heartmula_config(r#"{"model_type": "qwen3"}"#));
+        assert!(!is_heartmula_config("not json"));
+    }
+
+    #[test]
+    fn layout_detection() {
+        let mut files: Vec<String> = required_files();
+        files.push("config.json".into());
+        let refs: Vec<&str> = files.iter().map(String::as_str).collect();
+        assert!(is_acestep_layout(&refs));
+        assert!(!is_acestep_layout(&["config.json", "model.safetensors"]));
+        assert!(!is_acestep_layout(&["acestep-v15-turbo/config.json"]));
+    }
+
+    #[test]
+    fn peak_normalize_hits_target_and_spares_silence() {
+        let (mut l, mut r) = (vec![0.5, -0.25], vec![0.1, 0.0]);
+        peak_normalize(&mut l, &mut r, -1.0);
+        let target = 10f32.powf(-1.0 / 20.0);
+        assert!((l[0] - target).abs() < 1e-6);
+        assert!((l[1] + target / 2.0).abs() < 1e-6);
+        let (mut l, mut r) = (vec![1e-8, 0.0], vec![0.0, -1e-8]);
+        peak_normalize(&mut l, &mut r, -1.0);
+        assert_eq!(l, vec![1e-8, 0.0]);
+    }
+
+    #[test]
+    fn flac_round_trips_the_16_bit_samples() {
+        // A second of a decaying tone, plus the clamp at both rails.
+        let n = 48_000;
+        let left: Vec<f32> = (0..n)
+            .map(|i| (i as f32 * 0.05).sin() * (1.0 - i as f32 / n as f32))
+            .collect();
+        let mut right: Vec<f32> = left.iter().map(|s| s * -0.5).collect();
+        right[0] = 2.0;
+        right[1] = -2.0;
+        let flac = encode_flac_pcm16(&left, &right, 48_000).unwrap();
+        assert_eq!(&flac[..4], b"fLaC");
+        assert!(
+            flac.len() < n * 4,
+            "{} bytes is no smaller than PCM",
+            flac.len()
+        );
+
+        let mut reader = claxon::FlacReader::new(std::io::Cursor::new(&flac)).unwrap();
+        let info = reader.streaminfo();
+        assert_eq!(
+            (info.sample_rate, info.channels, info.bits_per_sample),
+            (48_000, 2, 16)
+        );
+        assert_eq!(info.samples, Some(n as u64));
+        let decoded: Vec<i32> = reader.samples().map(Result::unwrap).collect();
+        let pcm = |s: f32| (s.clamp(-1.0, 1.0) * i16::MAX as f32).round() as i32;
+        let expected: Vec<i32> = (0..n).flat_map(|i| [pcm(left[i]), pcm(right[i])]).collect();
+        assert_eq!(decoded, expected);
+        assert_eq!(
+            (decoded[1], decoded[3]),
+            (i16::MAX as i32, -(i16::MAX as i32))
+        );
+    }
+}
+
+/// End-to-end render against real weights, on the CPU. Ignored by default:
+/// it needs an ACE-Step 1.5 snapshot on disk and minutes of CPU time.
+///
+///   ACESTEP_SNAPSHOT=/path/to/Ace-Step1.5 \
+///     cargo test -p neuron --lib -- --ignored acestep_renders_on_cpu --nocapture
+#[cfg(test)]
+mod hardware {
+    use super::*;
+    use crate::harness::arch::acestep::pipeline::AceStepPipeline;
+
+    #[test]
+    #[ignore = "needs ACESTEP_SNAPSHOT and minutes of CPU"]
+    fn acestep_renders_on_cpu() {
+        let Ok(root) = std::env::var("ACESTEP_SNAPSHOT") else {
+            eprintln!("ACESTEP_SNAPSHOT not set; skipping");
+            return;
+        };
+        let files = AceStepFiles {
+            root: root.into(),
+            planner: None,
+        };
+        let t = Instant::now();
+        let mut engine = AceStepEngine::load(&files, &Device::Cpu).expect("load");
+        eprintln!("loaded in {:.1}s", t.elapsed().as_secs_f32());
+        let params = AudioGenParams {
+            prompt: "Warm lo-fi hip hop beat with mellow Rhodes chords, 82 BPM".into(),
+            lyrics: String::new(),
+            duration_secs: 10.0,
+            seed: 42,
+            steps: None,
+            guidance_scale: None,
+            planner: false,
+            temperature: None,
+            language: None,
+            top_k: None,
+        };
+        let t = Instant::now();
+        let out = engine.generate(&params).expect("generate");
+        eprintln!(
+            "rendered {:.2}s of audio in {:.1}s ({:?})",
+            out.duration_secs(),
+            t.elapsed().as_secs_f32(),
+            out.timing
+        );
+        assert_eq!(out.sample_rate, 48_000);
+        assert!(
+            (out.duration_secs() - 10.0).abs() < 0.5,
+            "{}",
+            out.duration_secs()
+        );
+        let rms = (out.left.iter().map(|s| s * s).sum::<f32>() / out.left.len() as f32).sqrt();
+        assert!(rms > 1e-3, "silent render, rms {rms}");
+        if let Ok(path) = std::env::var("ACESTEP_OUT") {
+            std::fs::write(
+                &path,
+                encode_flac_pcm16(&out.left, &out.right, out.sample_rate).expect("encode flac"),
+            )
+            .expect("write flac");
+            eprintln!("wrote {path}");
+        }
+        // Same seed, same song.
+        let again = engine.generate(&params).expect("generate again");
+        assert_eq!(out.left, again.left, "same seed must render identically");
+    }
+
+    /// Relative L2 distance between two tensors of the same shape.
+    fn rel_l2(a: &candle_core::Tensor, b: &candle_core::Tensor) -> f32 {
+        let a = a.to_dtype(DType::F32).unwrap().flatten_all().unwrap();
+        let b = b.to_dtype(DType::F32).unwrap().flatten_all().unwrap();
+        let diff = (&a - &b)
+            .unwrap()
+            .sqr()
+            .unwrap()
+            .sum_all()
+            .unwrap()
+            .to_scalar::<f32>()
+            .unwrap();
+        let norm = b
+            .sqr()
+            .unwrap()
+            .sum_all()
+            .unwrap()
+            .to_scalar::<f32>()
+            .unwrap();
+        (diff / norm.max(1e-20)).sqrt()
+    }
+
+    /// Parity with the upstream PyTorch reference (#381) on the CPU in f32,
+    /// from the reference's own noise: prompt token ids exactly, then the
+    /// denoised latents and the decoded waveform within float tolerance.
+    #[test]
+    #[ignore = "needs ACESTEP_SNAPSHOT and minutes of CPU"]
+    fn acestep_matches_reference_cpu_f32() {
+        let Ok(root) = std::env::var("ACESTEP_SNAPSHOT") else {
+            eprintln!("ACESTEP_SNAPSHOT not set; skipping");
+            return;
+        };
+        let cases = std::env::var("ACESTEP_CASES")
+            .unwrap_or_else(|_| "instr_10s_noplan,vocal_es_20s_noplan".into());
+        let mut engine = AceStepEngine::load(
+            &AceStepFiles {
+                root: root.into(),
+                planner: None,
+            },
+            &Device::Cpu,
+        )
+        .expect("load");
+        let dir =
+            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/acestep");
+        for case in cases.split(',') {
+            let meta: serde_json::Value = serde_json::from_slice(
+                &std::fs::read(dir.join(format!("{case}__cpu_f32.json"))).unwrap(),
+            )
+            .unwrap();
+            // Read tensors by name: the file also holds BOOL masks, which
+            // candle cannot load.
+            let file = unsafe {
+                candle_core::safetensors::MmapedSafetensors::new(
+                    dir.join(format!("{case}__cpu_f32.safetensors")),
+                )
+            }
+            .unwrap();
+            let want: std::collections::HashMap<&str, candle_core::Tensor> = [
+                "text_token_ids",
+                "lyric_token_ids",
+                "noise",
+                "target_latents",
+                "vae_waveform",
+            ]
+            .into_iter()
+            .map(|k| (k, file.load(k, &Device::Cpu).unwrap()))
+            .collect();
+            let c = &meta["case"];
+            let caption = c["caption"].as_str().unwrap();
+            let lyrics = c["lyrics"].as_str().unwrap();
+            let language = c["language"].as_str().unwrap();
+            let duration = c["duration"].as_f64().unwrap();
+
+            let (text, lyric) =
+                AceStepPipeline::format_prompts(caption, lyrics, duration, None, Some(language));
+            let (text_ids, lyric_ids) = engine.pipeline.prompt_token_ids(&text, &lyric).unwrap();
+            let ids = |k: &str| -> Vec<u32> {
+                want[k]
+                    .flatten_all()
+                    .unwrap()
+                    .to_vec1::<i64>()
+                    .unwrap()
+                    .into_iter()
+                    .map(|v| v as u32)
+                    .collect()
+            };
+            assert_eq!(text_ids, ids("text_token_ids"), "{case}: text token ids");
+            assert_eq!(lyric_ids, ids("lyric_token_ids"), "{case}: lyric token ids");
+
+            let params = GenerationParams {
+                duration_secs: duration,
+                seed: None,
+                num_steps: Some(DEFAULT_STEPS),
+                shift: TURBO_SHIFT,
+                is_turbo: true,
+                infer_method: sampling::InferMethod::Ode,
+                initial_noise: Some(want["noise"].clone()),
+                language: Some(language.to_string()),
+                ..Default::default()
+            };
+            let t = Instant::now();
+            let out = engine
+                .pipeline
+                .text2music(caption, lyrics, &params)
+                .unwrap();
+            let latents = rel_l2(&out.latents, &want["target_latents"]);
+            let wave_ref = &want["vae_waveform"];
+            let n = wave_ref.dim(2).unwrap();
+            let wave = out.audio.narrow(1, 0, n).unwrap().unsqueeze(0).unwrap();
+            let waveform = rel_l2(&wave, wave_ref);
+            eprintln!(
+                "{case}: latents rel L2 {latents:.2e}, waveform rel L2 {waveform:.2e} ({:.1}s)",
+                t.elapsed().as_secs_f32()
+            );
+            assert!(latents < 1e-3, "{case}: latents rel L2 {latents}");
+            assert!(waveform < 1e-2, "{case}: waveform rel L2 {waveform}");
+        }
+    }
+
+    /// End-to-end HeartMuLa render on the CPU from the real checkpoints,
+    /// laid out as heartlib does (`<ckpt>/{HeartMuLa-oss-3B,HeartCodec-oss,
+    /// tokenizer.json,gen_config.json}`). About 15.7 GB of f32 LM plus the
+    /// codec, so a host with 32 GB or more.
+    #[test]
+    #[ignore = "needs HEARTMULA_CKPT, ~24 GB of RAM and minutes of CPU"]
+    fn heartmula_renders_on_cpu() {
+        let Ok(root) = std::env::var("HEARTMULA_CKPT") else {
+            eprintln!("HEARTMULA_CKPT not set; skipping");
+            return;
+        };
+        let root = PathBuf::from(root);
+        let files = HeartMuLaFiles {
+            lm_dir: root.join("HeartMuLa-oss-3B"),
+            codec_dir: root.join("HeartCodec-oss"),
+            tokenizer: root.join("tokenizer.json"),
+            gen_config: root.join("gen_config.json"),
+        };
+        let t = Instant::now();
+        let mut engine =
+            AudioEngine::load(&AudioFiles::HeartMuLa(files), &Device::Cpu).expect("load");
+        eprintln!("loaded in {:.1}s", t.elapsed().as_secs_f32());
+        let params = AudioGenParams {
+            prompt: "indie pop, bright guitars, female vocals".into(),
+            lyrics: "[Verse]\nStreetlights hum a quiet tune".into(),
+            duration_secs: 10.0,
+            seed: 7,
+            steps: None,
+            guidance_scale: None,
+            planner: false,
+            temperature: None,
+            language: None,
+            top_k: None,
+        };
+        let t = Instant::now();
+        let out = engine.generate(&params).expect("generate");
+        eprintln!(
+            "rendered {:.2}s in {:.1}s ({:?})",
+            out.duration_secs(),
+            t.elapsed().as_secs_f32(),
+            out.timing
+        );
+        assert_eq!(out.sample_rate, 48_000);
+        assert!(
+            out.duration_secs() > 1.0 && out.duration_secs() <= 10.5,
+            "{}",
+            out.duration_secs()
+        );
+        let rms = (out.left.iter().map(|s| s * s).sum::<f32>() / out.left.len() as f32).sqrt();
+        assert!(rms > 1e-3, "silent render, rms {rms}");
+        if let Ok(path) = std::env::var("HEARTMULA_OUT") {
+            std::fs::write(
+                &path,
+                encode_flac_pcm16(&out.left, &out.right, out.sample_rate).unwrap(),
+            )
+            .unwrap();
+        }
+        let again = engine.generate(&params).expect("generate again");
+        assert_eq!(out.left, again.left, "same seed must render identically");
+    }
+
+    /// The planner LM loads from the real snapshot (its tensors carry no
+    /// `model.` prefix). Ignored: needs ACESTEP_SNAPSHOT.
+    #[test]
+    #[ignore = "needs ACESTEP_SNAPSHOT"]
+    fn acestep_planner_loads() {
+        let Ok(root) = std::env::var("ACESTEP_SNAPSHOT") else {
+            return;
+        };
+        let dir = PathBuf::from(root).join(PLANNER_DIRS[0]);
+        let cfg: qwen3::Config = read_json(&dir.join("config.json")).unwrap();
+        let vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(
+                &[dir.join("model.safetensors")],
+                DType::BF16,
+                &Device::Cpu,
+            )
+            .unwrap()
+        }
+        .rename_f(|n| n.strip_prefix("model.").unwrap_or(n).to_string());
+        qwen3::ModelForCausalLM::new(&cfg, vb).expect("planner loads");
+    }
+}

crates/neuron/src/harness/candle.rs

@@ -146,6 +146,8 @@ pub enum LoadedHandle {
     /// Laya decision model (#336): typed questions in, calibrated
     /// option scores out, no text generation.
     Decision(Arc<super::decision_model::LoadedDecisionModel>),
+    /// ACE-Step music generation pipeline (#383).
+    Audio(Arc<LoadedAudioModel>),
 }
 
 impl LoadedHandle {
@@ -155,6 +157,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => &m.model_id,
             LoadedHandle::Image(m) => &m.model_id,
+            LoadedHandle::Audio(m) => &m.model_id,
             LoadedHandle::Decision(m) => &m.model_id,
         }
     }
@@ -166,6 +169,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => &m.spec,
             LoadedHandle::Image(m) => &m.spec,
+            LoadedHandle::Audio(m) => &m.spec,
             LoadedHandle::Decision(m) => &m.spec,
         }
     }
@@ -176,6 +180,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.devices.clone(),
             LoadedHandle::Image(m) => m.devices.clone(),
+            LoadedHandle::Audio(m) => m.devices.clone(),
             LoadedHandle::Decision(m) => m.devices.clone(),
         }
     }
@@ -190,6 +195,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.poisoned.load(Ordering::Acquire),
             LoadedHandle::Image(m) => m.poisoned.load(Ordering::Acquire),
+            LoadedHandle::Audio(m) => m.poisoned.load(Ordering::Acquire),
             LoadedHandle::Decision(m) => m.poisoned.load(Ordering::Acquire),
         }
     }
@@ -202,6 +208,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => (m.admission.in_flight(), m.admission.queue_depth()),
             LoadedHandle::Image(m) => (m.admission.in_flight(), m.admission.queue_depth()),
+            LoadedHandle::Audio(m) => (m.admission.in_flight(), m.admission.queue_depth()),
             LoadedHandle::Decision(m) => (m.admission.in_flight(), m.admission.queue_depth()),
         }
     }
@@ -215,6 +222,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
             LoadedHandle::Image(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
+            LoadedHandle::Audio(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
             LoadedHandle::Decision(m) => {
                 (m.admission.max_in_flight(), m.admission.max_queue_depth())
             }
@@ -241,6 +249,10 @@ impl LoadedHandle {
                 m.admission.anon_in_flight(),
                 m.admission.anon_max_in_flight(),
             ),
+            LoadedHandle::Audio(m) => (
+                m.admission.anon_in_flight(),
+                m.admission.anon_max_in_flight(),
+            ),
             LoadedHandle::Decision(m) => (
                 m.admission.anon_in_flight(),
                 m.admission.anon_max_in_flight(),
@@ -255,6 +267,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.admission.close(),
             LoadedHandle::Image(m) => m.admission.close(),
+            LoadedHandle::Audio(m) => m.admission.close(),
             LoadedHandle::Decision(m) => m.admission.close(),
         }
     }
@@ -267,6 +280,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.admission.rejections(),
             LoadedHandle::Image(m) => m.admission.rejections(),
+            LoadedHandle::Audio(m) => m.admission.rejections(),
             LoadedHandle::Decision(m) => m.admission.rejections(),
         }
     }
@@ -279,6 +293,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => &m.admission,
             LoadedHandle::Image(m) => &m.admission,
+            LoadedHandle::Audio(m) => &m.admission,
             LoadedHandle::Decision(m) => &m.admission,
         };
         (admission.kv_budget_mb(), admission.kv_available_mb())
@@ -293,7 +308,9 @@ impl LoadedHandle {
             LoadedHandle::Tp(m) => &m.prefill_rate,
             // Image models have no token throughput; steps/sec lives in
             // the per-request timing instead.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return (0.0, 0.0),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                return (0.0, 0.0);
+            }
         };
         (ema.get().unwrap_or(0.0), ema.decode().unwrap_or(0.0))
     }
@@ -308,12 +325,17 @@ impl LoadedHandle {
         if matches!(self, LoadedHandle::Decision(_)) {
             return vec!["decision".to_string()];
         }
+        if matches!(self, LoadedHandle::Audio(_)) {
+            return vec![cortex_core::audio::CAPABILITY.to_string()];
+        }
         let mut caps = vec!["text".to_string()];
         let has_vision = match self {
             LoadedHandle::Single(m) => m.has_vision,
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.has_vision,
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => unreachable!("handled above"),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                unreachable!("handled above")
+            }
         };
         if has_vision {
             caps.push("vision".to_string());
@@ -359,7 +381,9 @@ impl LoadedHandle {
             LoadedHandle::Tp(m) => (m.last_free_mb.load(Ordering::Acquire), m.context_profile),
             // Image generation has its own VRAM path and no prefill
             // floor, so this check would say nothing useful about it.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                return None;
+            }
         };
         // `0` is the CPU-load and missing-worker sentinel — the same one
         // `check_vram` skips on. Nothing to assert either way.
@@ -440,7 +464,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.tool_call_tokens.is_some(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.tool_call_tokens.is_some(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => false,
         }
     }
 
@@ -451,7 +475,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.reasoning_tokens.is_some(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.reasoning_tokens.is_some(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => false,
         }
     }
 
@@ -462,7 +486,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.preserve_thinking,
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.preserve_thinking,
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => None,
         }
     }
 
@@ -472,7 +496,9 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.reasoning_efforts.clone(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.reasoning_efforts.clone(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => Default::default(),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                Default::default()
+            }
         }
     }
 
@@ -498,7 +524,7 @@ impl LoadedHandle {
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.query_vram_tightest_free_mb().await,
             // No context-limit derivation for image models; nothing to cache.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => return,
         };
         // Don't clobber a good cached value with a transient `0`
         // (worker gone/poisoned sentinel).
@@ -507,7 +533,7 @@ impl LoadedHandle {
                 LoadedHandle::Single(m) => m.last_free_mb.store(free, Ordering::Release),
                 #[cfg(feature = "cuda")]
                 LoadedHandle::Tp(m) => m.last_free_mb.store(free, Ordering::Release),
-                LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {}
+                LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {}
             }
         }
     }
@@ -536,7 +562,9 @@ impl LoadedHandle {
                 m.prefill_rate.get(),
             ),
             // Token limits are meaningless for image generation.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                return None;
+            }
         };
         let rate = rate.unwrap_or(cfg.bootstrap_prefill_tok_per_sec);
         // The KV budget admission will actually grant. Advertising past
@@ -546,7 +574,7 @@ impl LoadedHandle {
             LoadedHandle::Single(m) => m.admission.kv_budget_mb(),
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(m) => m.admission.kv_budget_mb(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
                 unreachable!("derived_limit returned above")
             }
         };
@@ -567,7 +595,7 @@ impl LoadedHandle {
                 LoadedHandle::Single(m) => &m.derived_input_cap,
                 #[cfg(feature = "cuda")]
                 LoadedHandle::Tp(m) => &m.derived_input_cap,
-                LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+                LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
                     unreachable!("derived_limit returned above")
                 }
             };
@@ -633,6 +661,29 @@ impl LoadedImageModel {
     }
 }
 
+/// A loaded ACE-Step music pipeline (#383). Every component lives in the
+/// device worker's audio slab, addressed via [`Self::handle`]; this
+/// async-side struct holds admission bookkeeping and the model's limits.
+/// A render is one `Job::GenerateAudio` round-trip.
+pub struct LoadedAudioModel {
+    pub model_id: String,
+    pub spec: ModelSpec,
+    pub devices: Vec<u32>,
+    pub worker: Arc<super::device_worker::DeviceWorkerHandle>,
+    pub handle: super::device_worker::AudioHandle,
+    /// Whether the snapshot carried a song planner.
+    pub has_planner: bool,
+    /// What a request may ask of this model; validated before admission.
+    pub limits: cortex_core::audio::AudioLimits,
+    /// Set when a render fails at the driver level: subsequent requests
+    /// fast-reject and `/models` reports `poisoned`.
+    pub poisoned: Arc<AtomicBool>,
+    /// Admission control (#53) with `max_in_flight` forced to 1 — a render
+    /// saturates the device. The cortex job queue is the real backlog.
+    pub admission: super::admission::AdmissionController,
+    pub inference_lock: Arc<tokio::sync::Mutex<()>>,
+}
+
 /// A loaded model with its tokenizer, device placement, and architecture-
 /// specific weights. The `arch` is `Arc<Mutex<>>` so the lock guard can be
 /// moved into `spawn_blocking` for synchronous candle forward passes.
@@ -3307,7 +3358,7 @@ impl CandleHarness {
             LoadedHandle::Tp(m) => {
                 return self.inference_tp_stream(m, request, principal, wire).await;
             }
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
                 return Err(InferenceError::WrongModality {
                     model_id: request.model.clone(),
                 });
@@ -4031,6 +4082,25 @@ impl Harness for CandleHarness {
         if matches!(plan.format, super::preflight::SourceFormat::DecisionModel) {
             return self.load_decision_model(spec, &source_id).await;
         }
+        // ACE-Step music models (#383) take the audio path: one device,
+        // unquantized (preflight rejected tp>1 and quant).
+        if matches!(plan.format, super::preflight::SourceFormat::AceStep) {
+            let files =
+                super::audio::AudioFiles::AceStep(self.resolve_audio_files(&source_id).await?);
+            return self
+                .load_audio_model(spec, files, super::audio::ACESTEP_LIMITS)
+                .await;
+        }
+        // HeartMuLa (#393): the LM repo, plus its codec and generation
+        // config from sibling repos.
+        if matches!(plan.format, super::preflight::SourceFormat::HeartMuLa) {
+            let files = super::audio::AudioFiles::HeartMuLa(
+                self.resolve_heartmula_files(&source_id).await?,
+            );
+            return self
+                .load_audio_model(spec, files, super::audio::HEARTMULA_LIMITS)
+                .await;
+        }
 
         let tp_size = spec.tensor_parallel.unwrap_or(1);
         if tp_size > 1 {
@@ -4402,6 +4472,17 @@ impl Harness for CandleHarness {
                 // them (the device worker, for a CUDA load).
                 dm.release().await;
             }
+            LoadedHandle::Audio(audio) => {
+                // Drop the pipeline on the worker thread so its CUDA
+                // tensors release on the context that allocated them.
+                if let Err(e) = audio.worker.drop_audio(audio.handle).await {
+                    tracing::warn!(
+                        model = %model_id,
+                        error = %e,
+                        "audio unload: DropAudio RPC failed (pipeline may leak in worker slab)"
+                    );
+                }
+            }
             #[cfg(feature = "cuda")]
             LoadedHandle::Tp(tp) => {
                 // Try to recover the inner TpLoadedModel so we can move
@@ -4576,6 +4657,161 @@ impl CandleHarness {
         Ok(())
     }
 
+    /// Load an ACE-Step music pipeline (#383): resolve the snapshot's
+    /// components, build them on the device worker, and register the
+    /// model with `capabilities: ["audio"]`.
+    #[cfg_attr(not(feature = "cuda"), allow(unused_variables, unreachable_code))]
+    async fn load_audio_model(
+        &self,
+        spec: &ModelSpec,
+        files: super::audio::AudioFiles,
+        limits: cortex_core::audio::AudioLimits,
+    ) -> Result<()> {
+        let devices = spec.devices.clone().unwrap_or_else(|| vec![0]);
+        let device = Self::pick_device(&devices)?;
+        let worker: Arc<super::device_worker::DeviceWorkerHandle> = match &device {
+            #[cfg(feature = "cuda")]
+            Device::Cuda(_) => self.ensure_device_worker(devices[0]).await?,
+            // On CPU builds this arm is the whole match: a song on the CPU
+            // takes many minutes and is not a serving target.
+            _ => anyhow::bail!(
+                "music generation requires a CUDA device; '{}' resolved to a CPU device",
+                spec.model_id
+            ),
+        };
+
+        let has_planner =
+            matches!(&files, super::audio::AudioFiles::AceStep(f) if f.planner.is_some());
+        let handle = worker
+            .load_audio(files, spec.model_id.clone())
+            .await
+            .map_err(|e| anyhow::anyhow!("worker load_audio: {e:#}"))?;
+
+        // Admission (#53) with the in-flight ceiling forced to 1: a render
+        // saturates the device. cortex queues jobs ahead of this.
+        let admission_cfg = crate::config::AdmissionConfig {
+            max_in_flight: 1,
+            ..self.admission_cfg.clone()
+        };
+        let loaded = LoadedAudioModel {
+            model_id: spec.model_id.clone(),
+            spec: spec.clone(),
+            devices,
+            worker,
+            handle,
+            has_planner,
+            limits,
+            poisoned: Arc::new(AtomicBool::new(false)),
+            admission: crate::harness::admission::AdmissionController::new(&admission_cfg),
+            inference_lock: Arc::new(tokio::sync::Mutex::new(())),
+        };
+        let mut models = self.models.write().await;
+        models.insert(spec.model_id.clone(), LoadedHandle::Audio(Arc::new(loaded)));
+        tracing::info!(model = %spec.model_id, planner = has_planner, "audio pipeline loaded");
+        Ok(())
+    }
+
+    /// Resolve an ACE-Step snapshot into a local root, downloading through
+    /// hf-hub as needed: every required component, plus the first planner
+    /// LM the repo has.
+    async fn resolve_audio_files(
+        &self,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<super::audio::AceStepFiles> {
+        let api = self.hf_api_for(&source_id.scheme)?;
+        let repo = api.model(source_id.repo_path());
+        let display_id = source_id.to_string();
+        let mut first = None;
+        for file in super::audio::required_files() {
+            let path = repo
+                .get(&file)
+                .await
+                .with_context(|| format!("fetch {file} from {display_id}"))?;
+            if first.is_none() {
+                first = Some((path, file));
+            }
+        }
+        // hf-hub materialises a snapshot as a tree mirroring the repo, so
+        // the root is the first file's path less its repo-relative part.
+        let (path, rel) = first.context("ACE-Step needs at least one file")?;
+        let mut root = path;
+        for _ in rel.split('/') {
+            root = root
+                .parent()
+                .context("snapshot path shorter than its repo path")?
+                .to_path_buf();
+        }
+        let mut planner = None;
+        for dir in super::audio::PLANNER_DIRS {
+            let mut ok = true;
+            for file in super::audio::planner_files(dir) {
+                if let Err(e) = repo.get(&file).await {
+                    tracing::info!(repo = %display_id, file, error = %format!("{e:#}"), "planner file unavailable");
+                    ok = false;
+                    break;
+                }
+            }
+            if ok {
+                planner = Some(dir.to_string());
+                break;
+            }
+        }
+        Ok(super::audio::AceStepFiles { root, planner })
+    }
+
+    /// Resolve a HeartMuLa LM repo (#393) and the two sibling repos it
+    /// needs — the codec and the generation config — through the same
+    /// registry, downloading as needed.
+    async fn resolve_heartmula_files(
+        &self,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<super::audio::HeartMuLaFiles> {
+        let api = self.hf_api_for(&source_id.scheme)?;
+        // Every file of a repo, fetched; returns the snapshot directory.
+        async fn fetch_all(
+            api: &hf_hub::api::tokio::Api,
+            repo_path: String,
+        ) -> Result<std::path::PathBuf> {
+            let repo = api.model(repo_path.clone());
+            let info = repo
+                .info()
+                .await
+                .with_context(|| format!("list files of {repo_path}"))?;
+            let mut dir = None;
+            for s in info.siblings {
+                let path = repo
+                    .get(&s.rfilename)
+                    .await
+                    .with_context(|| format!("fetch {} from {repo_path}", s.rfilename))?;
+                if s.rfilename == "config.json" || dir.is_none() && !s.rfilename.contains('/') {
+                    dir = path.parent().map(|p| p.to_path_buf());
+                }
+            }
+            dir.with_context(|| format!("{repo_path} has no top-level files"))
+        }
+        let lm_dir = fetch_all(&api, source_id.repo_path()).await?;
+        let codec_dir = fetch_all(&api, super::audio::HEARTCODEC_REPO.to_string()).await?;
+        let gen_repo = api.model(super::audio::HEARTMULA_GEN_REPO.to_string());
+        let tokenizer = gen_repo.get("tokenizer.json").await.with_context(|| {
+            format!(
+                "fetch tokenizer.json from {}",
+                super::audio::HEARTMULA_GEN_REPO
+            )
+        })?;
+        let gen_config = gen_repo.get("gen_config.json").await.with_context(|| {
+            format!(
+                "fetch gen_config.json from {}",
+                super::audio::HEARTMULA_GEN_REPO
+            )
+        })?;
+        Ok(super::audio::HeartMuLaFiles {
+            lm_dir,
+            codec_dir,
+            tokenizer,
+            gen_config,
+        })
+    }
+
     /// Load a Laya decision model (#336): resolve the checkpoints' files,
     /// build them on the device worker (or the CPU), and register the
     /// model with `capabilities: ["decision"]`.
@@ -4998,6 +5234,110 @@ impl CandleHarness {
             .ok_or_else(|| InferenceError::ModelNotLoaded("no decision model is loaded".into()))
     }
 
+    /// Render one audio generation (#382). The cortex job manager sends one
+    /// of these per job and waits for the whole render; the reply is the
+    /// FLAC bytes plus the result metadata cortex records on the job.
+    ///
+    /// Validated against the model's limits before admission; the render
+    /// itself is one device-worker job. A request without a seed gets a
+    /// random one, reported back so the song can be reproduced.
+    pub async fn audio_generation(
+        &self,
+        request: &cortex_core::audio::AudioGenerationRequest,
+        principal: Option<&str>,
+    ) -> Result<(Vec<u8>, cortex_core::audio::NeuronAudioResult), InferenceError> {
+        let handle = {
+            let models = self.models.read().await;
+            models.get(&request.model).cloned()
+        };
+        let audio = match handle {
+            Some(LoadedHandle::Audio(m)) => m,
+            Some(_) => {
+                return Err(InferenceError::WrongModality {
+                    model_id: request.model.clone(),
+                });
+            }
+            None => return Err(InferenceError::ModelNotLoaded(request.model.clone())),
+        };
+        if audio.poisoned.load(Ordering::Acquire) {
+            return Err(InferenceError::Other(anyhow::anyhow!(
+                "audio model '{}' is poisoned; unload and reload it",
+                request.model
+            )));
+        }
+
+        // Validate before admission so a malformed request never takes a
+        // queue slot.
+        request
+            .validate_limits(&audio.limits)
+            .map_err(|detail| InferenceError::InvalidAudioParams { detail })?;
+        // The planner runs only when asked for until its port is proven
+        // against the reference (#384); `auto` means off until then.
+        let planner = match request.planner.as_deref() {
+            Some("on") if !audio.has_planner => {
+                return Err(InferenceError::InvalidAudioParams {
+                    detail: "this model has no song planner; use planner=off or auto".into(),
+                });
+            }
+            Some("on") => true,
+            _ => false,
+        };
+        // A random seed small enough to type back in, so any render can be
+        // reproduced from what the job reports.
+        let seed = request.seed.unwrap_or_else(|| {
+            use std::hash::BuildHasher;
+            std::collections::hash_map::RandomState::new().hash_one(unix_subsec_nanos())
+                & 0x7fff_ffff
+        });
+        let params = super::audio::AudioGenParams {
+            prompt: request.prompt.clone(),
+            lyrics: if request.is_instrumental() {
+                String::new()
+            } else {
+                request.lyrics.clone().unwrap_or_default()
+            },
+            duration_secs: request.duration.unwrap_or(audio.limits.default_duration),
+            seed,
+            steps: request.steps.map(|s| s as usize),
+            guidance_scale: request.guidance_scale,
+            planner,
+            temperature: request.temperature,
+            language: request.language.clone(),
+            top_k: request.top_k.map(|k| k as usize),
+        };
+
+        let _permit = audio.admission.enter(principal).await?;
+        let _guard = audio.inference_lock.lock().await;
+        let result = audio
+            .worker
+            .generate_audio(audio.handle, params)
+            .await
+            .map_err(|e| InferenceError::Other(anyhow::anyhow!("audio generation: {e:#}")))?;
+
+        let duration = result.duration_secs();
+        let sample_rate = result.sample_rate;
+        let timing = result.timing;
+        let flac = tokio::task::spawn_blocking(move || {
+            super::audio::encode_flac_pcm16(&result.left, &result.right, result.sample_rate)
+        })
+        .await
+        .map_err(|e| InferenceError::Other(anyhow::anyhow!("flac encode task: {e:#}")))?
+        .map_err(InferenceError::Other)?;
+        Ok((
+            flac,
+            cortex_core::audio::NeuronAudioResult {
+                duration,
+                sample_rate,
+                channels: 2,
+                seed,
+                usage: cortex_core::audio::AudioUsage {
+                    helexa_audio_units: duration,
+                    helexa_timing: Some(timing),
+                },
+            },
+        ))
+    }
+
     /// Run one text-to-image generation (#198). This is the internal
     /// path the `/v1/images/generations` handler (#200) fronts;
     /// admission, validation, tokenization, and PNG encoding happen
@@ -6601,6 +6941,10 @@ pub enum InferenceError {
     /// Maps to 400 — client error, not server fault.
     #[error("invalid image generation parameters: {detail}")]
     InvalidImageParams { detail: String },
+    /// Audio generation parameters the model cannot serve (#383): a
+    /// duration out of range, lyrics for a model that does not sing.
+    #[error("invalid audio generation parameters: {detail}")]
+    InvalidAudioParams { detail: String },
     /// The loaded model's chat template could not render the request
     /// (e.g. a message / tool-call structure it rejects). Returned only
     /// when the request carried tools — silently degrading to a

crates/neuron/src/harness/capability.rs

@@ -60,13 +60,24 @@ pub fn classify_files(
         // A decision model answers typed questions (`/v1/systemone`);
         // it generates no text, so it must not claim "text".
         SourceFormat::DecisionModel => vec!["decision".to_string()],
+        // A music model (#383) renders audio on `/v1/audio/generations`.
+        SourceFormat::AceStep => vec![cortex_core::audio::CAPABILITY.to_string()],
         // A repo with no recognised weights (tokenizer-only, or an empty
         // entry) is not servable, so claiming "text" would be a lie of
         // the same kind this module exists to remove.
         SourceFormat::Empty => Vec::new(),
+        SourceFormat::HeartMuLa => vec![cortex_core::audio::CAPABILITY.to_string()],
         _ => {
+            let config = read_config();
+            // HeartMuLa (#393) is only recognisable by its config.
+            if config
+                .as_deref()
+                .is_some_and(crate::harness::audio::is_heartmula_config)
+            {
+                return vec![cortex_core::audio::CAPABILITY.to_string()];
+            }
             let mut caps = vec!["text".to_string()];
-            if read_config().is_some_and(|text| declares_vision(&text)) {
+            if config.is_some_and(|text| declares_vision(&text)) {
                 caps.push("vision".to_string());
             }
             caps
@@ -171,4 +182,28 @@ mod tests {
         let dir = std::env::temp_dir().join("helexa-capability-probe-missing");
         assert!(from_cache(&dir, "does-not/exist").is_none());
     }
+
+    #[test]
+    fn acestep_snapshot_is_audio_not_text() {
+        let mut files = crate::harness::audio::required_files();
+        files.push("config.json".into());
+        files.push("acestep-5Hz-lm-1.7B/model.safetensors".into());
+        let refs: Vec<&str> = files.iter().map(String::as_str).collect();
+        assert_eq!(classify_files(&refs, || None), vec!["audio".to_string()]);
+    }
+
+    #[test]
+    fn heartmula_snapshot_is_audio_by_its_config() {
+        let files = [
+            "config.json",
+            "model-00001-of-00004.safetensors",
+            "model.safetensors.index.json",
+        ];
+        let caps = classify_files(&files, || {
+            Some(r#"{"architectures": ["HeartMuLa"], "model_type": "heartmula"}"#.to_string())
+        });
+        assert_eq!(caps, vec!["audio".to_string()]);
+        let text = classify_files(&files, || Some(r#"{"model_type": "qwen3"}"#.to_string()));
+        assert_eq!(text, vec!["text".to_string()]);
+    }
 }

crates/neuron/src/harness/device_worker/dispatch.rs

@@ -19,7 +19,7 @@ use crate::harness::candle::ModelArch;
 #[cfg(feature = "cuda")]
 use crate::harness::device_worker::jobs::TpHandle;
 use crate::harness::device_worker::jobs::{
-    ArchHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
+    ArchHandle, AudioHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
 };
 use crate::harness::image::ZImagePipeline;
 #[cfg(feature = "cuda")]
@@ -71,6 +71,11 @@ struct DeviceWorkerState {
     decision_models: HashMap<DecisionHandle, Box<LayaModel>>,
     /// Counter for minting fresh `DecisionHandle`s.
     next_decision_handle: u64,
+    /// ACE-Step music pipeline slab (#383). Same lifecycle discipline as
+    /// `image_models`, in its own namespace.
+    audio_models: HashMap<AudioHandle, Box<crate::harness::audio::AudioEngine>>,
+    /// Counter for minting fresh `AudioHandle`s.
+    next_audio_handle: u64,
     /// Leader's NCCL state. Populated by `Job::NcclInit`; the
     /// underlying `Comm`'s libnccl handle lives bound to this thread
     /// for its entire lifetime. Subprocess workers maintain their own
@@ -813,6 +818,58 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 };
                 let _ = reply.send(result);
             }
+            Job::LoadAudio {
+                files,
+                model_id,
+                reply,
+            } => {
+                let result = crate::harness::audio::AudioEngine::load(&files, &state.device)
+                    .with_context(|| format!("load audio engine for {model_id}"))
+                    .map(|engine| {
+                        let handle = AudioHandle(state.next_audio_handle);
+                        state.next_audio_handle = state.next_audio_handle.wrapping_add(1);
+                        tracing::info!(
+                            device_index,
+                            model_id,
+                            engine = engine.name(),
+                            planner = engine.has_planner(),
+                            handle = handle.0,
+                            "device worker: audio pipeline loaded"
+                        );
+                        state.audio_models.insert(handle, Box::new(engine));
+                        handle
+                    });
+                trim_device_pool(&state);
+                let _ = reply.send(result);
+            }
+            Job::DropAudio { handle, reply } => {
+                let removed = state.audio_models.remove(&handle);
+                let was_present = removed.is_some();
+                // Explicit drop on this thread, as for DropImage.
+                drop(removed);
+                trim_device_pool(&state);
+                tracing::debug!(
+                    device_index,
+                    handle = handle.0,
+                    was_present,
+                    "device worker: audio pipeline dropped"
+                );
+                let _ = reply.send(());
+            }
+            Job::GenerateAudio {
+                handle,
+                params,
+                reply,
+            } => {
+                let result = match state.audio_models.get_mut(&handle) {
+                    Some(engine) => engine.generate(&params),
+                    None => Err(anyhow::anyhow!("no audio pipeline for handle {}", handle.0)),
+                };
+                // Song lengths vary per request, so release this render's
+                // pooled transients before the next one, as images do.
+                trim_device_pool(&state);
+                let _ = reply.send(result);
+            }
             Job::Shutdown => unreachable!("Shutdown should break above"),
         }
     }
@@ -887,6 +944,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
             next_image_handle: 1,
             decision_models: HashMap::new(),
             next_decision_handle: 1,
+            audio_models: HashMap::new(),
+            next_audio_handle: 1,
             nccl: NcclState::new(),
             tp_models: HashMap::new(),
             next_tp_handle: 1,
@@ -907,6 +966,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
             next_image_handle: 1,
             decision_models: HashMap::new(),
             next_decision_handle: 1,
+            audio_models: HashMap::new(),
+            next_audio_handle: 1,
             nccl: NcclState::new(),
         }
     }
@@ -2044,6 +2105,15 @@ fn drain_poisoned(job: Job, device_index: u32) {
         Job::Decide { reply, .. } => {
             let _ = reply.send(Err(err()));
         }
+        Job::LoadAudio { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        Job::DropAudio { reply, .. } => {
+            let _ = reply.send(());
+        }
+        Job::GenerateAudio { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
         Job::ForwardLogitsWithImages { reply, .. } => {
             let _ = reply.send(Err(err()));
         }

crates/neuron/src/harness/device_worker/jobs.rs

@@ -68,6 +68,15 @@ pub struct ImageHandle(pub u64);
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct DecisionHandle(pub u64);
 
+/// Opaque handle to a music engine ([`AudioEngine`]) stored
+/// in the worker thread's audio slab (#383). Same discipline as
+/// [`ImageHandle`]: construction, every generation and `Drop` happen on the
+/// worker thread, and `Job::DropAudio` is the only way to release it.
+///
+/// [`AudioEngine`]: crate::harness::audio::AudioEngine
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct AudioHandle(pub u64);
+
 /// Opaque handle to a prefix-cache snapshot (#11) stored worker-side
 /// next to the model slab. Scoped to the `ArchHandle` it was captured
 /// from — `Job::DropArch` drops every snapshot under its handle. The
@@ -585,6 +594,27 @@ pub enum Job {
         batch: crate::harness::arch::laya::DecisionBatch,
         reply: oneshot::Sender<Result<Vec<crate::harness::arch::laya::DecisionRow>>>,
     },
+    /// Load an ACE-Step snapshot (#383) — already on local disk — into
+    /// the audio slab. File resolution (hf-hub) happens on the async side.
+    LoadAudio {
+        files: crate::harness::audio::AudioFiles,
+        model_id: String,
+        reply: oneshot::Sender<Result<AudioHandle>>,
+    },
+    /// Remove a music pipeline from the slab and drop it on the worker
+    /// thread. Mirrors `DropImage`.
+    DropAudio {
+        handle: AudioHandle,
+        reply: oneshot::Sender<()>,
+    },
+    /// Render one song (#383): planner, text encoding, denoise and VAE
+    /// decode in one job, since a render saturates the device. The reply
+    /// is CPU-side samples — no device tensor crosses back.
+    GenerateAudio {
+        handle: AudioHandle,
+        params: crate::harness::audio::AudioGenParams,
+        reply: oneshot::Sender<Result<crate::harness::audio::AudioGenResult>>,
+    },
     /// Tell the worker to break its dispatch loop and exit. Any jobs
     /// queued after this in the channel reply `Err` to their oneshot
     /// senders (the senders are dropped on the worker's exit, which

crates/neuron/src/harness/device_worker/mod.rs

@@ -51,7 +51,9 @@ use tokio::sync::oneshot;
 
 #[cfg(feature = "cuda")]
 pub use jobs::TpHandle;
-pub use jobs::{ArchHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
+pub use jobs::{
+    ArchHandle, AudioHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId,
+};
 
 /// Errors returned by `DeviceWorkerHandle` submit methods.
 #[derive(Debug, thiserror::Error)]
@@ -337,6 +339,82 @@ impl DeviceWorkerHandle {
         }
     }
 
+    /// Load a music engine (#383, #393) on the worker thread.
+    pub async fn load_audio(
+        &self,
+        files: crate::harness::audio::AudioFiles,
+        model_id: String,
+    ) -> Result<AudioHandle, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::LoadAudio {
+                files,
+                model_id,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
+    /// Drop the music pipeline for `handle` on the worker thread. Same
+    /// poison-tolerant contract as [`Self::drop_arch`].
+    pub async fn drop_audio(&self, handle: AudioHandle) -> Result<(), WorkerError> {
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::DropAudio {
+                handle,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        reply_rx.await.map_err(|_| WorkerError::Gone {
+            device_index: self.device_index,
+        })
+    }
+
+    /// Render one song (#383) on the worker thread. The reply is CPU-side
+    /// samples; no device tensor crosses back.
+    pub async fn generate_audio(
+        &self,
+        handle: AudioHandle,
+        params: crate::harness::audio::AudioGenParams,
+    ) -> Result<crate::harness::audio::AudioGenResult, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::GenerateAudio {
+                handle,
+                params,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
     /// Load one Laya checkpoint directory (#336) on the worker thread.
     /// Returns the handle and the dtype the weights were loaded in.
     pub async fn load_decision(

crates/neuron/src/harness/mod.rs

@@ -2,6 +2,7 @@
 
 pub mod admission;
 pub mod arch;
+pub mod audio;
 pub mod candle;
 pub mod capability;
 pub mod chat_template;

crates/neuron/src/harness/preflight.rs

@@ -59,6 +59,16 @@ pub enum SourceFormat {
     /// `config.json`. Served by the decision path — single device,
     /// unquantized. Takes precedence over the safetensors sniff.
     DecisionModel,
+    /// A HeartMuLa LM repo (#393): sharded safetensors like any dense LLM,
+    /// told apart by `model_type: heartmula` in its `config.json`. Served by
+    /// the audio path with its codec from a sibling repo; single device.
+    HeartMuLa,
+    /// An ACE-Step 1.5 music repo (#383): the turbo DiT, text encoder,
+    /// VAE and planner as subfolders of one snapshot. Served by the audio
+    /// path — single device, unquantized for now. Takes precedence over
+    /// the safetensors sniff, which would read the root `config.json` as a
+    /// text model.
+    AceStep,
 }
 
 /// Output of `preflight` for a load that can proceed. Carries the
@@ -175,7 +185,17 @@ pub async fn preflight(
     };
 
     let filenames: Vec<&str> = owned_filenames.iter().map(String::as_str).collect();
-    let format = classify(&filenames);
+    let mut format = classify(&filenames);
+    // A HeartMuLa LM (#393) looks like any dense checkpoint by its file
+    // names; only its config says otherwise.
+    if matches!(format, SourceFormat::DenseSafetensors { .. })
+        && filenames.contains(&"config.json")
+        && let Ok(path) = repo.get("config.json").await
+        && std::fs::read_to_string(&path)
+            .is_ok_and(|t| crate::harness::audio::is_heartmula_config(&t))
+    {
+        format = SourceFormat::HeartMuLa;
+    }
     let tp_size = spec.tensor_parallel.unwrap_or(1);
 
     match (&format, tp_size, spec.quant.as_deref()) {
@@ -258,6 +278,26 @@ pub async fn preflight(
             tp_size,
             picked_quant_file: None,
         }),
+        (SourceFormat::AceStep | SourceFormat::HeartMuLa, tp, _) if tp > 1 => {
+            Err(PreflightError::Unsupported {
+                model_id: source_id.to_string(),
+                reason: format!(
+                    "music models load on a single device; tensor_parallel={tp} is not supported"
+                ),
+            })
+        }
+        (SourceFormat::AceStep | SourceFormat::HeartMuLa, _, Some(q)) => {
+            Err(PreflightError::Unsupported {
+                model_id: source_id.to_string(),
+                reason: format!("music models load unquantized for now; remove quant={q:?}"),
+            })
+        }
+        (SourceFormat::AceStep | SourceFormat::HeartMuLa, _, None) => Ok(PlacementPlan {
+            model_id: source_id.to_string(),
+            format: format.clone(),
+            tp_size,
+            picked_quant_file: None,
+        }),
 
         (SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
             Ok(PlacementPlan {
@@ -327,6 +367,9 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
     if is_decision_layout(filenames) {
         return SourceFormat::DecisionModel;
     }
+    if crate::harness::audio::is_acestep_layout(filenames) {
+        return SourceFormat::AceStep;
+    }
     let mut gguf_quants: Vec<String> = filenames
         .iter()
         .filter(|f| f.to_lowercase().ends_with(".gguf"))
@@ -707,6 +750,24 @@ mod tests {
                 tp_size,
                 picked_quant_file: None,
             }),
+            (SourceFormat::AceStep | SourceFormat::HeartMuLa, tp, _) if tp > 1 => {
+                Err(PreflightError::Unsupported {
+                    model_id: source_id.to_string(),
+                    reason: format!("music models load on a single device; tensor_parallel={tp}"),
+                })
+            }
+            (SourceFormat::AceStep | SourceFormat::HeartMuLa, _, Some(q)) => {
+                Err(PreflightError::Unsupported {
+                    model_id: source_id.to_string(),
+                    reason: format!("music models load unquantized for now; remove quant={q:?}"),
+                })
+            }
+            (SourceFormat::AceStep | SourceFormat::HeartMuLa, _, None) => Ok(PlacementPlan {
+                model_id: source_id.to_string(),
+                format: format.clone(),
+                tp_size,
+                picked_quant_file: None,
+            }),
             (SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
                 Ok(PlacementPlan {
                     model_id: source_id.to_string(),
@@ -751,6 +812,10 @@ mod tests {
     #[test]
     fn classify_laya_layout() {
         assert_eq!(classify(&LAYA_FILES), SourceFormat::DecisionModel);
+        let mut ace = crate::harness::audio::required_files();
+        ace.push("config.json".into());
+        let ace: Vec<&str> = ace.iter().map(String::as_str).collect();
+        assert_eq!(classify(&ace), SourceFormat::AceStep);
         // A causal LM that also ships an rl config is still a text model.
         let lm = [
             "config.json",

crates/neuron/src/harness/testdata/acestep/README.md

@@ -0,0 +1,140 @@
+# ACE-Step 1.5 reference traces
+
+Parity fixtures for neuron's ACE-Step port (#383, #384, #385), recorded by
+`script/acestep-reference.py` from the upstream PyTorch pipeline:
+
+- code: `ace-step/ACE-Step-1.5` @ `ca1e85fe9430179831e6bc6be790c332190a3866`
+- weights: `ACE-Step/Ace-Step1.5` @ `19671f406d603126926c1b7e2adc169acbcade22`
+  (`acestep-v15-turbo`, `acestep-5Hz-lm-1.7B`, `Qwen3-Embedding-0.6B`, `vae`)
+- torch 2.10.0+cu128, transformers 4.57.6, pytorch_wavelets 1.3.0
+
+Each `<case>__<trace>.safetensors` holds the tensors, and the matching
+`.json` holds the metadata: the case inputs, decoded prompts, the sampler
+settings `generate_audio` actually received, the timestep schedule, planner
+prompts and outputs, per-tensor stats (shape, dtype, NaN/Inf counts, absmax,
+mean, std), the first tensor to contain a NaN, drift against `cpu_f32`, wall
+time and peak VRAM. `index.json` summarises every trace.
+
+## Cases
+
+| case | content | duration | planner |
+|---|---|---|---|
+| `instr_10s_noplan` | instrumental | 10 s | off |
+| `vocal_en_30s_noplan` | English vocals | 30 s | off |
+| `vocal_en_30s_plan` | English vocals | 30 s | on (1.7B, greedy) |
+| `vocal_es_20s_noplan` | Spanish vocals | 20 s | off |
+| `vocal_en_120s_noplan` | English vocals | 120 s | off |
+
+"Planner off" turns off `thinking` and all three `use_cot_*` flags, so the
+LM is never called. "Planner on" runs the LM greedily (`lm_temperature=0`),
+with upstream's default CFG (2.0 for the code phase) and constrained
+decoding.
+
+## Traces
+
+| trace | device | dtype | what it is |
+|---|---|---|---|
+| `cpu_f32` | CPU | float32 | the parity reference: deterministic and complete |
+| `cuda_bf16` | RTX 4090 | bfloat16 | what upstream runs on Ampere and later |
+| `cuda_f32` | RTX 4090 | float32 | true f32 on the GPU, TF32 off: the Pascal compute recipe |
+| `cuda_f16` | RTX 4090 | float16 | upstream's own pre-Ampere path (fp16 + eager attention), forced |
+
+All traces of a case start from the **same noise**: `cpu_f32` records it,
+and the GPU traces load it instead of drawing their own. CUDA and CPU
+`torch.Generator`s differ at the same seed. A port should likewise take
+`noise` as an input when checking parity.
+
+`cuda_f32` is loaded on the CPU and then moved to the GPU. Loading on CUDA
+and casting to f32 afterwards is not equivalent. A CUDA load rounds every
+RoPE `inv_freq` buffer and the float32 `silence_latent` to bf16, and a later
+cast cannot restore them. That alone moved a single DiT forward by 11%
+(rel L2) from the CPU result.
+
+## Results
+
+| case | `cuda_f32` vs CPU | `cuda_bf16` vs CPU | `cuda_f16` |
+|---|---|---|---|
+| `instr_10s_noplan` | 6e-6 | 0.08 | clean (0.018) |
+| `vocal_es_20s_noplan` | 1e-5 | 0.17 | **NaN** |
+| `vocal_en_30s_noplan` | 8e-4 | 0.41 | **NaN** |
+| `vocal_en_120s_noplan` | 2e-4 | 0.63 | **NaN** |
+| `vocal_en_30s_plan` | — | 1.01 (different plan) | — |
+
+Values are the relative L2 error of `target_latents` against `cpu_f32`.
+
+- **fp16 vocals NaN, instrumentals don't.** This is upstream issue #1055,
+  reproduced on Ada. The first non-finite value is in the lyric encoder's
+  last layer: `layers.7.mlp.down_proj` overflows float16 (the residual
+  stream is already ~3.6k at layer 6). Keeping RoPE in f32 does not help.
+  "[Instrumental]" is a short lyric that stays in range. Upstream raises
+  rather than decoding; the trace records everything up to that point.
+- **True f32 on the GPU matches the CPU** to about 1e-5 per forward, and
+  1e-5 to 8e-4 end to end. The f32 recipe is numerically sound.
+- **Upstream's bf16 path is far from f32**: 8% to 63% end to end, growing
+  with length. Part of that is the bf16-rounded `inv_freq` and silence
+  latent described above, not only bf16 arithmetic. A candle port that
+  computes RoPE in f32 should be judged against `cpu_f32`/`cuda_f32`.
+  `cuda_bf16` is upstream's behaviour, not a tolerance target.
+- **The planner's code phase is fragile under precision change.** The CoT
+  phase (164 tokens) is identical between `cpu_f32` and `cuda_bf16`, but
+  the very first audio code differs, so the whole plan differs.
+
+Peak VRAM on the 4090 (DiT, text encoder and VAE; plus the 1.7B LM where
+loaded): bf16 with the LM loaded, 10.2 GiB; f32 without the LM, 12.6 GiB;
+fp16 without the LM, 6.2–6.7 GiB.
+
+## Tensors
+
+`cpu_f32` and `cuda_f32` store float32; `cuda_bf16` and `cuda_f16` store
+their own dtype, which is exact and half the size. Shapes are for batch 1,
+and `T` is the number of 25 Hz latent frames.
+
+| name | shape | what |
+|---|---|---|
+| `text_token_ids` | `[1, Lt]` | DiT text prompt, tokenized by the Qwen3-Embedding tokenizer (prompt text in the JSON) |
+| `text_hidden_states` | `[1, Lt, 1024]` | Qwen3-Embedding-0.6B `last_hidden_state` |
+| `lyric_token_ids` | `[1, Ll]` | lyric prompt (`# Languages\n<lang>\n\n# Lyric\n<lyrics><\|endoftext\|>`) |
+| `lyric_embeds` | `[1, Ll, 1024]` | `embed_tokens` lookup only, with no transformer pass |
+| `encoder_hidden_states` | `[1, Ll+1+Lt, 2048]` | packed condition: lyric-encoder output (8 layers), then timbre (1 token), then projected text |
+| `encoder_attention_mask` | `[1, Ll+1+Lt]` | packed mask |
+| `gen_chunk_masks`, `gen_is_covers` | | `generate_audio` inputs |
+| `gen_precomputed_lm_hints_25Hz` | `[1, T, 64]` | planner cases: LM audio codes, detokenized to 25 Hz |
+| `context_latents` | `[1, T, 128]` | src latents (silence, or LM hints) concatenated with chunk masks |
+| `decoder_attention_mask` | `[1, T]` | |
+| `noise` | `[1, T, 64]` | initial `xt` (CPU traces only; GPU traces reuse it) |
+| `step{i}_xt`, `step{i}_vt` | `[1, T, 64]` | DiT input and predicted velocity at sampler step `i` |
+| `target_latents` | `[1, T, 64]` | sampler output, after the final step's `x0` and DCW |
+| `vae_input` | `[1, 64, T]` | what the VAE decoder received |
+| `vae_waveform` | `[1, 2, 24000]` | VAE output, first 0.5 s (48 kHz stereo), before normalisation. The pipeline output is this × `normalisation_gain` (JSON) |
+| `planner{n}_input_ids`, `planner{n}_output_ids` | | LM phase `n` (0 = CoT, 1 = codes): prompt ids (with CFG, conditional and unconditional rows, left-padded) and generated ids |
+| `planner_audio_codes` | `[N]` | the 5 Hz audio codes the planner produced (N = 5 × seconds) |
+
+The lyric encoder's own output is not stored. With an all-true mask it is
+exactly the first `Ll` rows of `encoder_hidden_states`, and the JSON
+records that this was checked (`lyric_encoder_out_is_encoder_prefix`).
+
+Steps saved (`steps_saved` in the JSON):
+
+- **CPU trace:** every step for the 10 s case; steps 0, 4 and 7 for the
+  20–30 s cases; step 0 for the 120 s case.
+- **GPU traces:** `step0_vt`, `step7_vt` and `target_latents`. The 120 s
+  case keeps only `target_latents`. A step where a NaN first appears is
+  always kept.
+- `step0_xt` is never stored, because it is `noise`.
+- The 120 s case omits `context_latents` (silence plus the chunk mask) and
+  `vae_input` (`target_latents` transposed).
+
+Every tensor that was dropped still has its stats in the JSON.
+
+## Regenerating
+
+In the ACE-Step venv, with the checkpoints already on disk (fetched through
+your registry; the script refuses to download). Record `cpu_f32` first,
+because the GPU traces load its noise:
+
+```sh
+python -I script/acestep-reference.py --src <ACE-Step-1.5 checkout> --trace cpu_f32 --out <this dir>
+python -I script/acestep-reference.py --src ... --trace cuda_bf16 --out ...
+python -I script/acestep-reference.py --src ... --trace cuda_f32 --out ...
+python -I script/acestep-reference.py --src ... --trace cuda_f16 --out ...
+```

crates/neuron/src/harness/testdata/acestep/index.json

@@ -0,0 +1,238 @@
+{
+ "traces": {
+  "instr_10s_noplan__cpu_f32": {
+   "case": "instr_10s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 17.4,
+   "vram_peak_gib": null,
+   "bytes": 2569385
+  },
+  "instr_10s_noplan__cuda_bf16": {
+   "case": "instr_10s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.0816783532500267,
+   "error": null,
+   "wall_s": 5.07,
+   "vram_peak_gib": 10.15,
+   "bytes": 96321
+  },
+  "instr_10s_noplan__cuda_f16": {
+   "case": "instr_10s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.018279578536748886,
+   "error": null,
+   "wall_s": 7.37,
+   "vram_peak_gib": 6.69,
+   "bytes": 96313
+  },
+  "instr_10s_noplan__cuda_f32": {
+   "case": "instr_10s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 6.1245264078024775e-06,
+   "error": null,
+   "wall_s": 10.95,
+   "vram_peak_gib": 12.63,
+   "bytes": 384401
+  },
+  "vocal_en_120s_noplan__cpu_f32": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 3000,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 144.06,
+   "vram_peak_gib": null,
+   "bytes": 4603028
+  },
+  "vocal_en_120s_noplan__cuda_bf16": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 3000,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.6251527070999146,
+   "error": null,
+   "wall_s": 40.29,
+   "vram_peak_gib": 10.18,
+   "bytes": 384169
+  },
+  "vocal_en_120s_noplan__cuda_f16": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 3000,
+   "nan_anywhere": true,
+   "first_nan": "lyric_encoder_out",
+   "target_latents_rel_l2_vs_cpu_f32": Infinity,
+   "error": "Generation produced NaN or Inf latents (shape=[1, 3000, 64], dtype=torch.float16, device=cuda:0, nan=192000, inf=0).\nCommon causes and fixes:\n  1. LoRA/adapter trained on an older model version \u2014 retrain or update the adapter.\n  2. Checkpoint/config mismatch \u2014 verify model checkpoints match this release.\n  3. Unsupported quantization/backend \u2014 try running with --backend pt.\n  4. CPU offload left parameters on wrong device \u2014 restart and regenerate.\n  5. Float16 overflow on pre-Ampere GPU \u2014 set ACESTEP_DTYPE=float32.",
+   "wall_s": 18.43,
+   "vram_peak_gib": 6.42,
+   "bytes": 384169
+  },
+  "vocal_en_120s_noplan__cuda_f32": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 3000,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.00023412087466567755,
+   "error": null,
+   "wall_s": 80.42,
+   "vram_peak_gib": 12.64,
+   "bytes": 960249
+  },
+  "vocal_en_30s_noplan__cpu_f32": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 39.68,
+   "vram_peak_gib": null,
+   "bytes": 3636105
+  },
+  "vocal_en_30s_noplan__cuda_bf16": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.4098935127258301,
+   "error": null,
+   "wall_s": 11.14,
+   "vram_peak_gib": 10.16,
+   "bytes": 288321
+  },
+  "vocal_en_30s_noplan__cuda_f16": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 750,
+   "nan_anywhere": true,
+   "first_nan": "lyric_encoder_out",
+   "target_latents_rel_l2_vs_cpu_f32": Infinity,
+   "error": "Generation produced NaN or Inf latents (shape=[1, 750, 64], dtype=torch.float16, device=cuda:0, nan=48000, inf=0).\nCommon causes and fixes:\n  1. LoRA/adapter trained on an older model version \u2014 retrain or update the adapter.\n  2. Checkpoint/config mismatch \u2014 verify model checkpoints match this release.\n  3. Unsupported quantization/backend \u2014 try running with --backend pt.\n  4. CPU offload left parameters on wrong device \u2014 restart and regenerate.\n  5. Float16 overflow on pre-Ampere GPU \u2014 set ACESTEP_DTYPE=float32.",
+   "wall_s": 3.77,
+   "vram_peak_gib": 6.16,
+   "bytes": 288321
+  },
+  "vocal_en_30s_noplan__cuda_f32": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.0008115487289614975,
+   "error": null,
+   "wall_s": 23.05,
+   "vram_peak_gib": 12.63,
+   "bytes": 768401
+  },
+  "vocal_en_30s_plan__cpu_f32": {
+   "case": "vocal_en_30s_plan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 97.87,
+   "vram_peak_gib": null,
+   "bytes": 5055700
+  },
+  "vocal_en_30s_plan__cuda_bf16": {
+   "case": "vocal_en_30s_plan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 1.0133154392242432,
+   "error": null,
+   "wall_s": 72.41,
+   "vram_peak_gib": 10.18,
+   "bytes": 394505
+  },
+  "vocal_es_20s_noplan__cpu_f32": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 500,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 25.68,
+   "vram_peak_gib": null,
+   "bytes": 2843838
+  },
+  "vocal_es_20s_noplan__cuda_bf16": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 500,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.16632689535617828,
+   "error": null,
+   "wall_s": 7.86,
+   "vram_peak_gib": 10.17,
+   "bytes": 192321
+  },
+  "vocal_es_20s_noplan__cuda_f16": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 500,
+   "nan_anywhere": true,
+   "first_nan": "lyric_encoder_out",
+   "target_latents_rel_l2_vs_cpu_f32": Infinity,
+   "error": "Generation produced NaN or Inf latents (shape=[1, 500, 64], dtype=torch.float16, device=cuda:0, nan=32000, inf=0).\nCommon causes and fixes:\n  1. LoRA/adapter trained on an older model version \u2014 retrain or update the adapter.\n  2. Checkpoint/config mismatch \u2014 verify model checkpoints match this release.\n  3. Unsupported quantization/backend \u2014 try running with --backend pt.\n  4. CPU offload left parameters on wrong device \u2014 restart and regenerate.\n  5. Float16 overflow on pre-Ampere GPU \u2014 set ACESTEP_DTYPE=float32.",
+   "wall_s": 6.31,
+   "vram_peak_gib": 6.16,
+   "bytes": 192321
+  },
+  "vocal_es_20s_noplan__cuda_f32": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 500,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 1.0157832548429724e-05,
+   "error": null,
+   "wall_s": 15.1,
+   "vram_peak_gib": 12.63,
+   "bytes": 576401
+  }
+ }
+}

crates/neuron/src/harness/testdata/acestep/instr_10s_noplan__cpu_f32.json

@@ -0,0 +1,594 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 14568.125,
+   "std": 30829.681640625
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    72,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 32.86247253417969,
+   "mean": 0.021975461393594742,
+   "std": 3.2490692138671875
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 29993.19921875,
+   "std": 52921.25
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    15,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1767578125,
+   "mean": 2.041088191617746e-05,
+   "std": 0.03075374662876129
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    15,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.524937629699707,
+   "mean": -9.119734022533521e-05,
+   "std": 0.11492618173360825
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    88,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 58.46120834350586,
+   "mean": 0.006924757268279791,
+   "std": 3.4267163276672363
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    88
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.2051306664943695,
+   "std": 0.9637550711631775
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.20505407452583313,
+   "std": 0.963009774684906
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.068782329559326,
+   "mean": -0.001519594225101173,
+   "std": 1.0108991861343384
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    250
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    250,
+    128
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": 0.39743465185165405,
+   "std": 0.9096668362617493
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.789435386657715,
+   "mean": -0.037685856223106384,
+   "std": 1.038394808769226
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    250
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.789435386657715,
+   "mean": -0.037685852497816086,
+   "std": 1.038394808769226
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0992342233657837,
+   "mean": 0.0012270656879991293,
+   "std": 0.16220736503601074
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.068782329559326,
+   "mean": -0.001519594225101173,
+   "std": 1.0108991861343384
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.779356002807617,
+   "mean": 0.029778892174363136,
+   "std": 1.0888593196868896
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.017130374908447,
+   "mean": -0.0014519158285111189,
+   "std": 0.9958045482635498
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.739595413208008,
+   "mean": 0.0698830708861351,
+   "std": 1.1297178268432617
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.936816453933716,
+   "mean": -0.00226192525587976,
+   "std": 0.9712164402008057
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.189512729644775,
+   "mean": 0.048355162143707275,
+   "std": 1.1502469778060913
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.813173294067383,
+   "mean": -0.0036722831428050995,
+   "std": 0.9389426112174988
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.870788097381592,
+   "mean": 0.04673989862203598,
+   "std": 1.2650023698806763
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.5883963108062744,
+   "mean": -0.006106651853770018,
+   "std": 0.8947911858558655
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.094154357910156,
+   "mean": 0.03310832753777504,
+   "std": 1.3123292922973633
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.4403023719787598,
+   "mean": -0.008855825290083885,
+   "std": 0.8410162329673767
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.550355911254883,
+   "mean": 0.04433209449052811,
+   "std": 1.3689327239990234
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.3948163986206055,
+   "mean": -0.014476503245532513,
+   "std": 0.7926256656646729
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.73010540008545,
+   "mean": 0.05107048898935318,
+   "std": 1.4008506536483765
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.267222881317139,
+   "mean": -0.024307560175657272,
+   "std": 0.8039591312408447
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.258280754089355,
+   "mean": 0.04459429532289505,
+   "std": 1.5695395469665527
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.0009948957012966275,
+   "std": 0.13151653110980988
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.9028035402297974,
+   "mean": 0.001543997903354466,
+   "std": 0.23970195651054382
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nwarm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 10 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nunknown\n\n# Lyric\n[Instrumental]<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   1234
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_seed": [
+  1234
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "timesteps": [
+  1.0,
+  0.9545454382896423,
+  0.8999999761581421,
+  0.8333333134651184,
+  0.75,
+  0.6428571343421936,
+  0.5,
+  0.30000001192092896
+ ],
+ "steps_saved": {
+  "xt": [
+   1,
+   2,
+   3,
+   4,
+   5,
+   6,
+   7
+  ],
+  "vt": [
+   0,
+   1,
+   2,
+   3,
+   4,
+   5,
+   6,
+   7
+  ]
+ },
+ "lyric_encoder_out_is_encoder_prefix": true,
+ "normalisation_gain": 0.8107925653457642,
+ "case": {
+  "name": "instr_10s_noplan",
+  "caption": "warm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed",
+  "lyrics": "[Instrumental]",
+  "instrumental": true,
+  "language": "unknown",
+  "duration": 10.0,
+  "planner": false,
+  "seed": 1234
+ },
+ "trace": "cpu_f32",
+ "sample_rate": 48000,
+ "latent_frames": 250,
+ "wall_s": 17.4,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "float32",
+ "time_costs": {
+  "dit_encoder_time_cost": 0.5645701885223389,
+  "dit_diffusion_time_cost": 5.696354866027832,
+  "dit_diffusion_per_step_time_cost": 0.712044358253479,
+  "dit_total_time_cost": 10.253296375274658,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 3.9923713207244873,
+  "pipeline_total_time": 10.253296375274658
+ }
+}

crates/neuron/src/harness/testdata/acestep/instr_10s_noplan__cuda_bf16.json

@@ -0,0 +1,628 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 14568.125,
+   "std": 30829.6796875
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    72,
+    1024
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 33.0,
+   "mean": 0.021900618448853493,
+   "std": 3.2504899501800537
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 29993.201171875,
+   "std": 52921.24609375
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    15,
+    1024
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1767578125,
+   "mean": 2.0410874640219845e-05,
+   "std": 0.03075374662876129
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    15,
+    2048
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.53125,
+   "mean": -8.808567508822307e-05,
+   "std": 0.11486581712961197
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    88,
+    2048
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 58.75,
+   "mean": 0.007172065321356058,
+   "std": 3.4257657527923584
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    88
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.71875,
+   "mean": -0.20533928275108337,
+   "std": 0.9645420908927917
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.71875,
+   "mean": -0.20527024567127228,
+   "std": 0.9638134837150574
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0625,
+   "mean": -0.0015215945895761251,
+   "std": 1.0108590126037598
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    250
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    250,
+    128
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.71875,
+   "mean": 0.3973304033279419,
+   "std": 0.9101529121398926
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.8125,
+   "mean": -0.03683716058731079,
+   "std": 1.0391204357147217
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    250
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.8125,
+   "mean": -0.03683716803789139,
+   "std": 1.0391203165054321
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    480000
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.09375,
+   "mean": 0.0019740236457437277,
+   "std": 0.16050520539283752
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0625,
+   "mean": -0.0015215944731608033,
+   "std": 1.0108590126037598
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.75,
+   "mean": 0.029024723917245865,
+   "std": 1.087060570716858
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0,
+   "mean": -0.0014959202380850911,
+   "std": 0.9948191046714783
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.78125,
+   "mean": 0.0699673667550087,
+   "std": 1.131856918334961
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.90625,
+   "mean": -0.002318344544619322,
+   "std": 0.9700552821159363
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.15625,
+   "mean": 0.04690621793270111,
+   "std": 1.1479475498199463
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.796875,
+   "mean": -0.0036648816894739866,
+   "std": 0.9381669163703918
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.875,
+   "mean": 0.0431983508169651,
+   "std": 1.2662577629089355
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.5625,
+   "mean": -0.0058504026383161545,
+   "std": 0.8947983980178833
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.125,
+   "mean": 0.0328899547457695,
+   "std": 1.3014971017837524
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.421875,
+   "mean": -0.008506354875862598,
+   "std": 0.8428772687911987
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.5625,
+   "mean": 0.044462598860263824,
+   "std": 1.3666576147079468
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.359375,
+   "mean": -0.014202592894434929,
+   "std": 0.7932853102684021
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.75,
+   "mean": 0.048620857298374176,
+   "std": 1.3998051881790161
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.25,
+   "mean": -0.023524314165115356,
+   "std": 0.8046457171440125
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.3125,
+   "mean": 0.044245440512895584,
+   "std": 1.5749262571334839
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.0016085486859083176,
+   "std": 0.13078895211219788
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.90625,
+   "mean": 0.0024766994174569845,
+   "std": 0.2409236878156662
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nwarm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 10 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nunknown\n\n# Lyric\n[Instrumental]<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   1234
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_source": "cpu_f32",
+ "noise_seed": [
+  1234
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "drift_vs_cpu_f32": {
+  "target_latents": {
+   "max_abs": 1.1789467334747314,
+   "rel_l2": 0.0816783532500267,
+   "cos": 0.9966683387756348
+  },
+  "step0_vt": {
+   "max_abs": 0.9140055179595947,
+   "rel_l2": 0.0775710716843605,
+   "cos": 0.9969895482063293
+  },
+  "step1_vt": {
+   "max_abs": 1.0068531036376953,
+   "rel_l2": 0.04007680341601372,
+   "cos": 0.9992018342018127
+  },
+  "step2_vt": {
+   "max_abs": 2.0684454441070557,
+   "rel_l2": 0.09504576772451401,
+   "cos": 0.9954776763916016
+  },
+  "step3_vt": {
+   "max_abs": 1.8307780027389526,
+   "rel_l2": 0.1291956901550293,
+   "cos": 0.9916636943817139
+  },
+  "step4_vt": {
+   "max_abs": 4.880147457122803,
+   "rel_l2": 0.14142011106014252,
+   "cos": 0.989952802658081
+  },
+  "step5_vt": {
+   "max_abs": 1.726083755493164,
+   "rel_l2": 0.10134051740169525,
+   "cos": 0.9948592185974121
+  },
+  "step6_vt": {
+   "max_abs": 2.1955137252807617,
+   "rel_l2": 0.10752364993095398,
+   "cos": 0.9942163825035095
+  },
+  "step7_vt": {
+   "max_abs": 1.8701152801513672,
+   "rel_l2": 0.09535709768533707,
+   "cos": 0.9954763650894165
+  }
+ },
+ "timesteps": [
+  1.0,
+  0.953125,
+  0.8984375,
+  0.83203125,
+  0.75,
+  0.64453125,
+  0.5,
+  0.30078125
+ ],
+ "steps_saved": {
+  "xt": [],
+  "vt": [
+   0,
+   7
+  ]
+ },
+ "normalisation_gain": 0.8148579597473145,
+ "case": {
+  "name": "instr_10s_noplan",
+  "caption": "warm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed",
+  "lyrics": "[Instrumental]",
+  "instrumental": true,
+  "language": "unknown",
+  "duration": 10.0,
+  "planner": false,
+  "seed": 1234
+ },
+ "trace": "cuda_bf16",
+ "sample_rate": 48000,
+ "latent_frames": 250,
+ "wall_s": 5.07,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "bfloat16",
+ "time_costs": {
+  "dit_encoder_time_cost": 0.12479805946350098,
+  "dit_diffusion_time_cost": 1.8494272232055664,
+  "dit_diffusion_per_step_time_cost": 0.2311784029006958,
+  "dit_total_time_cost": 3.6470632553100586,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 1.6728379726409912,
+  "pipeline_total_time": 3.6470632553100586
+ },
+ "vram_peak_gib": 10.15
+}

crates/neuron/src/harness/testdata/acestep/instr_10s_noplan__cuda_f16.json

@@ -0,0 +1,628 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 14568.125,
+   "std": 30829.6796875
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    72,
+    1024
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 32.75,
+   "mean": 0.022007420659065247,
+   "std": 3.2488560676574707
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 29993.201171875,
+   "std": 52921.24609375
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    15,
+    1024
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1767578125,
+   "mean": 2.0410874640219845e-05,
+   "std": 0.03075374662876129
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    15,
+    2048
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.52734375,
+   "mean": -9.115711145568639e-05,
+   "std": 0.11492296308279037
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    88,
+    2048
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 58.46875,
+   "mean": 0.006939474958926439,
+   "std": 3.4267029762268066
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    88
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.7265625,
+   "mean": -0.20514169335365295,
+   "std": 0.9638112187385559
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.7265625,
+   "mean": -0.20506583154201508,
+   "std": 0.9630692601203918
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0703125,
+   "mean": -0.0015202447539195418,
+   "std": 1.0108911991119385
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    250
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    250,
+    128
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.7265625,
+   "mean": 0.3974291980266571,
+   "std": 0.9097001552581787
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.7890625,
+   "mean": -0.03734346851706505,
+   "std": 1.037626028060913
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    250
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.78125,
+   "mean": -0.03735503554344177,
+   "std": 1.0376352071762085
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    480000
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.1015625,
+   "mean": 0.0019519772613421082,
+   "std": 0.16201943159103394
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0703125,
+   "mean": -0.00152024463750422,
+   "std": 1.0108911991119385
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.78125,
+   "mean": 0.029700657352805138,
+   "std": 1.0890905857086182
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.01953125,
+   "mean": -0.0014485636493191123,
+   "std": 0.9958088397979736
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.75,
+   "mean": 0.07011272013187408,
+   "std": 1.1304819583892822
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.939453125,
+   "mean": -0.00227189133875072,
+   "std": 0.9710971117019653
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.1875,
+   "mean": 0.04861733317375183,
+   "std": 1.1522797346115112
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.81640625,
+   "mean": -0.0036781812086701393,
+   "std": 0.9388203024864197
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.8671875,
+   "mean": 0.046269163489341736,
+   "std": 1.26438570022583
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.591796875,
+   "mean": -0.006093707866966724,
+   "std": 0.8945609927177429
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.09375,
+   "mean": 0.03247145190834999,
+   "std": 1.3111038208007812
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.44140625,
+   "mean": -0.00878163706511259,
+   "std": 0.8409976363182068
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.5546875,
+   "mean": 0.043878648430109024,
+   "std": 1.3689857721328735
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.404296875,
+   "mean": -0.014353488571941853,
+   "std": 0.7923156023025513
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.734375,
+   "mean": 0.05066253989934921,
+   "std": 1.4011808633804321
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.265625,
+   "mean": -0.02410614863038063,
+   "std": 0.8034966588020325
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.265625,
+   "mean": 0.04412555694580078,
+   "std": 1.5696253776550293
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.001579303527250886,
+   "std": 0.13108648359775543
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.90625,
+   "mean": 0.002429787302389741,
+   "std": 0.24185852706432343
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nwarm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 10 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nunknown\n\n# Lyric\n[Instrumental]<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   1234
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_source": "cpu_f32",
+ "noise_seed": [
+  1234
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "drift_vs_cpu_f32": {
+  "target_latents": {
+   "max_abs": 0.19800710678100586,
+   "rel_l2": 0.018279578536748886,
+   "cos": 0.9998327493667603
+  },
+  "step0_vt": {
+   "max_abs": 0.2126154899597168,
+   "rel_l2": 0.01686670631170273,
+   "cos": 0.9998575448989868
+  },
+  "step1_vt": {
+   "max_abs": 0.151755690574646,
+   "rel_l2": 0.005986565724015236,
+   "cos": 0.9999818801879883
+  },
+  "step2_vt": {
+   "max_abs": 0.4102330207824707,
+   "rel_l2": 0.029562195762991905,
+   "cos": 0.999565064907074
+  },
+  "step3_vt": {
+   "max_abs": 0.4682796001434326,
+   "rel_l2": 0.023266855627298355,
+   "cos": 0.9997289180755615
+  },
+  "step4_vt": {
+   "max_abs": 1.1111046075820923,
+   "rel_l2": 0.02790045738220215,
+   "cos": 0.999610424041748
+  },
+  "step5_vt": {
+   "max_abs": 0.7564409375190735,
+   "rel_l2": 0.0248439721763134,
+   "cos": 0.9996907711029053
+  },
+  "step6_vt": {
+   "max_abs": 0.7065806984901428,
+   "rel_l2": 0.024734484031796455,
+   "cos": 0.9996935725212097
+  },
+  "step7_vt": {
+   "max_abs": 0.5175025463104248,
+   "rel_l2": 0.022454872727394104,
+   "cos": 0.9997473359107971
+  }
+ },
+ "timesteps": [
+  1.0,
+  0.95458984375,
+  0.89990234375,
+  0.83349609375,
+  0.75,
+  0.64306640625,
+  0.5,
+  0.300048828125
+ ],
+ "steps_saved": {
+  "xt": [],
+  "vt": [
+   0,
+   7
+  ]
+ },
+ "normalisation_gain": 0.8090788125991821,
+ "case": {
+  "name": "instr_10s_noplan",
+  "caption": "warm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed",
+  "lyrics": "[Instrumental]",
+  "instrumental": true,
+  "language": "unknown",
+  "duration": 10.0,
+  "planner": false,
+  "seed": 1234
+ },
+ "trace": "cuda_f16",
+ "sample_rate": 48000,
+ "latent_frames": 250,
+ "wall_s": 7.37,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "float16",
+ "time_costs": {
+  "dit_encoder_time_cost": 0.14662766456604004,
+  "dit_diffusion_time_cost": 1.8529925346374512,
+  "dit_diffusion_per_step_time_cost": 0.2316240668296814,
+  "dit_total_time_cost": 3.715280294418335,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 1.7156600952148438,
+  "pipeline_total_time": 3.715280294418335
+ },
+ "vram_peak_gib": 6.69
+}

crates/neuron/src/harness/testdata/acestep/instr_10s_noplan__cuda_f32.json

@@ -0,0 +1,628 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 14568.125,
+   "std": 30829.6796875
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    72,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 32.86249542236328,
+   "mean": 0.021975407376885414,
+   "std": 3.2490689754486084
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 29993.201171875,
+   "std": 52921.24609375
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    15,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1767578125,
+   "mean": 2.0410874640219845e-05,
+   "std": 0.03075374662876129
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    15,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.524935722351074,
+   "mean": -9.119740570895374e-05,
+   "std": 0.11492616683244705
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    88,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 58.46121597290039,
+   "mean": 0.006924731191247702,
+   "std": 3.4267163276672363
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    88
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.2051306515932083,
+   "std": 0.9637550711631775
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.20505407452583313,
+   "std": 0.9630097150802612
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.068782329559326,
+   "mean": -0.0015195925952866673,
+   "std": 1.0108991861343384
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    250
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    250,
+    128
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": 0.39743471145629883,
+   "std": 0.9096668362617493
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.789431571960449,
+   "mean": -0.0376858115196228,
+   "std": 1.0383943319320679
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    250
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.789431571960449,
+   "mean": -0.0376858115196228,
+   "std": 1.0383943319320679
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.099230408668518,
+   "mean": 0.0012270613806322217,
+   "std": 0.16220729053020477
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.068782329559326,
+   "mean": -0.001519594225101173,
+   "std": 1.0108991861343384
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.779355049133301,
+   "mean": 0.029778843745589256,
+   "std": 1.0888599157333374
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.017129898071289,
+   "mean": -0.0014519135002046824,
+   "std": 0.9958045482635498
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.739598751068115,
+   "mean": 0.06988321244716644,
+   "std": 1.129718542098999
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.9368159770965576,
+   "mean": -0.00226192525587976,
+   "std": 0.9712163805961609
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.189511775970459,
+   "mean": 0.04835524037480354,
+   "std": 1.1502478122711182
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.8131728172302246,
+   "mean": -0.003672286169603467,
+   "std": 0.9389424920082092
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.870791912078857,
+   "mean": 0.046739786863327026,
+   "std": 1.2650028467178345
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.5883960723876953,
+   "mean": -0.006106650922447443,
+   "std": 0.8947910070419312
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.094152450561523,
+   "mean": 0.03310833498835564,
+   "std": 1.3123295307159424
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.4403018951416016,
+   "mean": -0.008855823427438736,
+   "std": 0.8410159945487976
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.550355911254883,
+   "mean": 0.04433206096291542,
+   "std": 1.3689327239990234
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.3948168754577637,
+   "mean": -0.01447649672627449,
+   "std": 0.7926254272460938
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.730103492736816,
+   "mean": 0.05107041820883751,
+   "std": 1.4008506536483765
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.267223834991455,
+   "mean": -0.024307552725076675,
+   "std": 0.8039588928222656
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.258277893066406,
+   "mean": 0.044594213366508484,
+   "std": 1.5695394277572632
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.0009948954684659839,
+   "std": 0.13151691854000092
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.9028029441833496,
+   "mean": 0.001543995807878673,
+   "std": 0.2397020012140274
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nwarm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 10 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nunknown\n\n# Lyric\n[Instrumental]<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   1234
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_source": "cpu_f32",
+ "noise_seed": [
+  1234
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "drift_vs_cpu_f32": {
+  "target_latents": {
+   "max_abs": 0.0001239776611328125,
+   "rel_l2": 6.1245264078024775e-06,
+   "cos": 1.0
+  },
+  "step0_vt": {
+   "max_abs": 9.021162986755371e-05,
+   "rel_l2": 6.811034836573526e-06,
+   "cos": 1.0000004768371582
+  },
+  "step1_vt": {
+   "max_abs": 0.0001545548439025879,
+   "rel_l2": 4.860378339799354e-06,
+   "cos": 1.000000238418579
+  },
+  "step2_vt": {
+   "max_abs": 0.00017547607421875,
+   "rel_l2": 9.911213055602275e-06,
+   "cos": 1.0000001192092896
+  },
+  "step3_vt": {
+   "max_abs": 0.00027549266815185547,
+   "rel_l2": 1.0616528015816584e-05,
+   "cos": 1.000000238418579
+  },
+  "step4_vt": {
+   "max_abs": 0.00023436546325683594,
+   "rel_l2": 9.188515832647681e-06,
+   "cos": 1.000000238418579
+  },
+  "step5_vt": {
+   "max_abs": 0.0001315474510192871,
+   "rel_l2": 6.429604582081083e-06,
+   "cos": 0.9999998211860657
+  },
+  "step6_vt": {
+   "max_abs": 0.00012257695198059082,
+   "rel_l2": 6.204536930454196e-06,
+   "cos": 1.0000003576278687
+  },
+  "step7_vt": {
+   "max_abs": 0.0003447532653808594,
+   "rel_l2": 8.590050128987059e-06,
+   "cos": 1.0
+  }
+ },
+ "timesteps": [
+  1.0,
+  0.9545454382896423,
+  0.8999999761581421,
+  0.8333333134651184,
+  0.75,
+  0.6428571343421936,
+  0.5,
+  0.30000001192092896
+ ],
+ "steps_saved": {
+  "xt": [],
+  "vt": [
+   0,
+   7
+  ]
+ },
+ "normalisation_gain": 0.8107953667640686,
+ "case": {
+  "name": "instr_10s_noplan",
+  "caption": "warm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed",
+  "lyrics": "[Instrumental]",
+  "instrumental": true,
+  "language": "unknown",
+  "duration": 10.0,
+  "planner": false,
+  "seed": 1234
+ },
+ "trace": "cuda_f32",
+ "sample_rate": 48000,
+ "latent_frames": 250,
+ "wall_s": 10.95,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "float32",
+ "time_costs": {
+  "dit_encoder_time_cost": 0.26670241355895996,
+  "dit_diffusion_time_cost": 3.246095895767212,
+  "dit_diffusion_per_step_time_cost": 0.4057619869709015,
+  "dit_total_time_cost": 8.351051092147827,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 4.838252782821655,
+  "pipeline_total_time": 8.351051092147827
+ },
+ "vram_peak_gib": 12.63
+}

crates/neuron/src/harness/testdata/acestep/vocal_en_120s_noplan__cpu_f32.json

@@ -0,0 +1,579 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    69
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 13378.8984375,
+   "std": 30062.69140625
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    69,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 25.977462768554688,
+   "mean": 0.020457586273550987,
+   "std": 3.241281747817993
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    85
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 12799.2705078125,
+   "std": 28789.384765625
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    85,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1455078125,
+   "mean": -6.758134986739606e-05,
+   "std": 0.028640329837799072
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    85,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.863743782043457,
+   "mean": 0.0002219161542598158,
+   "std": 0.21818150579929352
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    155,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 59.5096435546875,
+   "mean": -0.00023797634639777243,
+   "std": 2.5121827125549316
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    155
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.20503060519695282,
+   "std": 0.9627315402030945
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    69
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    85
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.20505407452583313,
+   "std": 0.963009774684906
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.481950759887695,
+   "mean": -0.0012088901130482554,
+   "std": 1.0021259784698486
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    3000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    3000,
+    128
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": 0.397484689950943,
+   "std": 0.9090932011604309
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.9633355140686035,
+   "mean": -0.05977988615632057,
+   "std": 0.9689841270446777
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    3000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.9633355140686035,
+   "mean": -0.05977989733219147,
+   "std": 0.9689841270446777
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    5760000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.124276876449585,
+   "mean": 0.0022899911273270845,
+   "std": 0.24259647727012634
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.481950759887695,
+   "mean": -0.0012088901130482554,
+   "std": 1.0021259784698486
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.200160026550293,
+   "mean": 0.031169500201940536,
+   "std": 0.9595843553543091
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.489745140075684,
+   "mean": -0.0011380501091480255,
+   "std": 0.9892756342887878
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.930626392364502,
+   "mean": 0.0512068010866642,
+   "std": 0.9604851603507996
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.402994155883789,
+   "mean": -0.001731583965010941,
+   "std": 0.9694339036941528
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.688071250915527,
+   "mean": 0.056339722126722336,
+   "std": 1.0303702354431152
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.267680644989014,
+   "mean": -0.0033748256973922253,
+   "std": 0.9417667984962463
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.325617790222168,
+   "mean": 0.05985618755221367,
+   "std": 1.1135389804840088
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.1802473068237305,
+   "mean": -0.006492334417998791,
+   "std": 0.9044287800788879
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.624982833862305,
+   "mean": 0.0637756735086441,
+   "std": 1.1794636249542236
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.143441200256348,
+   "mean": -0.011787991970777512,
+   "std": 0.8583288788795471
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 10.063026428222656,
+   "mean": 0.0740080326795578,
+   "std": 1.2499028444290161
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.062942028045654,
+   "mean": -0.021171152591705322,
+   "std": 0.8128213286399841
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 10.117952346801758,
+   "mean": 0.07707519084215164,
+   "std": 1.2784833908081055
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.788512229919434,
+   "mean": -0.03600812703371048,
+   "std": 0.8092694282531738
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    3000,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 11.121763229370117,
+   "mean": 0.07923923432826996,
+   "std": 1.3733216524124146
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    5760000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.0018153507262468338,
+   "std": 0.19231413304805756
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0015655570896342397,
+   "mean": 0.0008253876585513353,
+   "std": 0.00031154719181358814
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nanthemic rock, driving drums, distorted guitars, male vocals, stadium\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 120 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nen\n\n# Lyric\n[Verse]\nEngines roar across the plain\nThunder rolling in the rain\n[Chorus]\nWe are the fire, we are the flame\nNothing will ever be the same\n[Verse]\nIron hearts and open roads\nCarry all our heavy loads\n[Chorus]\nWe are the fire, we are the flame\nNothing will ever be the same\n<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   99
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_seed": [
+  99
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "timesteps": [
+  1.0,
+  0.9545454382896423,
+  0.8999999761581421,
+  0.8333333134651184,
+  0.75,
+  0.6428571343421936,
+  0.5,
+  0.30000001192092896
+ ],
+ "steps_saved": {
+  "xt": [],
+  "vt": [
+   0
+  ]
+ },
+ "lyric_encoder_out_is_encoder_prefix": true,
+ "normalisation_gain": 0.7927325963973999,
+ "case": {
+  "name": "vocal_en_120s_noplan",
+  "caption": "anthemic rock, driving drums, distorted guitars, male vocals, stadium",
+  "lyrics": "[Verse]\nEngines roar across the plain\nThunder rolling in the rain\n[Chorus]\nWe are the fire, we are the flame\nNothing will ever be the same\n[Verse]\nIron hearts and open roads\nCarry all our heavy loads\n[Chorus]\nWe are the fire, we are the flame\nNothing will ever be the same\n",
+  "instrumental": false,
+  "language": "en",
+  "duration": 120.0,
+  "planner": false,
+  "seed": 99
+ },
+ "trace": "cpu_f32",
+ "sample_rate": 48000,
+ "latent_frames": 3000,
+ "wall_s": 144.06,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "float32",
+ "time_costs": {
+  "dit_encoder_time_cost": 1.936945915222168,
+  "dit_diffusion_time_cost": 43.167064905166626,
+  "dit_diffusion_per_step_time_cost": 5.395883113145828,
+  "dit_total_time_cost": 141.40626668930054,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 96.30225586891174,
+  "pipeline_total_time": 141.40626668930054
+ }
+}