From 6c35b3b976e75661f867de1f2cabc80c6cd028e1 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Tue, 15 Sep 2026 13:05:55 -0400 Subject: [PATCH 1/3] Harden release workflows and engine lifecycle --- .github/workflows/ci.yml | 12 + .github/workflows/deploy.yml | 4 +- README.md | 344 ++++++------------ package-lock.json | 46 ++- package.json | 12 +- packages/acestep/README.md | 209 ++++------- packages/dicose/README.md | 186 +++++----- scripts/build-sdk.mjs | 7 +- scripts/test-sdk.mjs | 65 ++++ src/core/resource-session.ts | 70 ++++ src/engines/musicgen-acestep/cache-lease.ts | 18 + src/engines/musicgen-acestep/index.ts | 120 +++--- .../musicgen-acestep/model-cache-migration.ts | 59 ++- .../pending-output-registry.ts | 96 +++++ .../musicgen-acestep/progress-breadcrumb.ts | 10 + src/engines/registry.ts | 92 +---- src/engines/sdk-registry.ts | 87 +++++ src/engines/stem-dicose/index.ts | 41 ++- src/engines/tts-kokoro/synth-backend.js | 153 ++++---- src/index.ts | 4 +- src/music.ts | 244 +++++++------ src/pages/playground.ts | 312 ++++++++++------ test/cache-lifecycle.test.ts | 298 +++++++++++++++ test/deploy-workflow.test.ts | 19 + test/dicose-lifecycle.test.ts | 72 ++++ test/kokoro-cleanup.test.ts | 20 + test/music-disposal.test.ts | 101 +++++ test/music-page-disposal.test.ts | 32 ++ test/music-preparation.test.ts | 140 +++++++ test/pending-output-registry.test.ts | 55 +++ test/progress-breadcrumb.test.ts | 27 ++ test/resource-session.test.ts | 118 ++++++ tsconfig.sdk.json | 2 +- vitest.config.ts | 5 + 34 files changed, 2135 insertions(+), 945 deletions(-) create mode 100644 scripts/test-sdk.mjs create mode 100644 src/core/resource-session.ts create mode 100644 src/engines/musicgen-acestep/cache-lease.ts create mode 100644 src/engines/musicgen-acestep/pending-output-registry.ts create mode 100644 src/engines/musicgen-acestep/progress-breadcrumb.ts create mode 100644 src/engines/sdk-registry.ts create mode 100644 test/cache-lifecycle.test.ts create mode 100644 test/deploy-workflow.test.ts create mode 100644 test/dicose-lifecycle.test.ts create mode 100644 test/kokoro-cleanup.test.ts create mode 100644 test/music-disposal.test.ts create mode 100644 test/music-page-disposal.test.ts create mode 100644 test/music-preparation.test.ts create mode 100644 test/pending-output-registry.test.ts create mode 100644 test/progress-breadcrumb.test.ts create mode 100644 test/resource-session.test.ts create mode 100644 vitest.config.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d22ca03..de78459 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,6 +3,7 @@ on: push: branches: [main] pull_request: + workflow_dispatch: # Rapid successive pushes to one ref cancel in-flight runs (the engines job is # the expensive one). @@ -24,6 +25,8 @@ jobs: - run: npm ci --ignore-scripts - run: npm run test:unit - run: npm run build + - run: npm run test:integration + - run: npm run sdk:test # Vendored ACE-Step runtime (packages/acestep): typecheck + weight-free # vitest suite (artifact-dependent contract tests self-exclude on machines @@ -129,3 +132,12 @@ jobs: key: hf-weights-${{ hashFiles('scripts/ci-weights.lock') }} - run: node scripts/textnorm-check.mjs - run: node scripts/vocab-rescorer-check.mjs + + deploy: + if: github.ref == 'refs/heads/main' && (github.event_name == 'push' || github.event_name == 'workflow_dispatch') + needs: [build, acestep, dicose, format, kernels, engines] + permissions: + contents: read + pages: write + id-token: write + uses: ./.github/workflows/deploy.yml diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 8efdba6..2cac08b 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -1,8 +1,6 @@ name: Deploy to GitHub Pages on: - push: - branches: [main] - workflow_dispatch: + workflow_call: # Allow one concurrent deploy; let a new push cancel an in-flight one. permissions: diff --git a/README.md b/README.md index 2604818..bc79613 100644 --- a/README.md +++ b/README.md @@ -1,33 +1,80 @@ # FluidAudio Web -Local speech AI in the browser — ASR, TTS, VAD, speaker diarization, and -music generation on **hand-written WebGPU (WGSL) + WASM-SIMD kernels**. No onnxruntime-web, no -transformers.js, no server: model weights stream from Hugging Face on first -use, cache client-side, and everything runs on the visitor's machine. This is -the browser sibling of the Swift/CoreML +Speech recognition, text to speech, music generation, and audio analysis in +your browser. Inference runs locally using custom WebGPU and WASM kernels; +audio is not uploaded. Model weights download on first use. + +This is the browser sibling of the Swift/CoreML [FluidAudio](https://github.com/FluidInference/FluidAudio) framework. -**1 hour of audio transcribed in ~12 seconds — 293× real-time — in a Chrome -tab** (Parakeet TDT 0.6B v3, multilingual; verified across three runs on the -1-hour benchmark, Chrome/macOS/WebGPU; ~199× under the node harness). - -**Live:** https://fluidinference.github.io/fluidaudio-web/ — one page per -function: [speech to text](https://fluidinference.github.io/fluidaudio-web/) -at the root, [text to speech](https://fluidinference.github.io/fluidaudio-web/tts.html), -[music generation + stem splitting](https://fluidinference.github.io/fluidaudio-web/music.html), -[other audio models](https://fluidinference.github.io/fluidaudio-web/analyze.html) -(VAD, diarization, stems), and [live captions](https://fluidinference.github.io/fluidaudio-web/live.html). -Mirror: https://fluidaudio-web.hanweng9.workers.dev. Both deploy automatically -from every commit to `main`. - -> **Why hand-written kernels?** The first iteration of this repo ran the same -> models through onnxruntime-web. Rewriting the hot paths as raw WGSL + Rust -> WASM-SIMD (see [`docs/ORT_REMOVAL.md`](docs/ORT_REMOVAL.md) and -> [`docs/RAW_WEBGPU.md`](docs/RAW_WEBGPU.md)) took Parakeet from 33× to **100×+ -> real-time in-browser** — batched-window encoding, f16 weight storage _and_ -> f16 compute (2× ALU on Apple GPUs), a 3-stage GPU/CPU pipeline, and parallel -> RNNT decode on a Web Worker pool. Every optimization is gated on -> token-identical output. +## Try it + +- [Speech to text](https://fluidinference.github.io/fluidaudio-web/) +- [Text to speech](https://fluidinference.github.io/fluidaudio-web/tts.html) +- [Music generation and stem splitting](https://fluidinference.github.io/fluidaudio-web/music.html) +- [Audio analysis](https://fluidinference.github.io/fluidaudio-web/analyze.html) +- [Live captions](https://fluidinference.github.io/fluidaudio-web/live.html) + +[Cloudflare mirror](https://fluidaudio-web.hanweng9.workers.dev). + +Use a desktop browser with WebGPU for best performance. Speech engines also +support WASM-SIMD. Music generation requires WebGPU with `shader-f16`; stem +splitting additionally requires fixed 32-wide subgroups. These two engines +have no WASM fallback. + +Downloads, cached model loading, and shader compilation all affect startup +time. Caching depends on browser storage availability and limits; some large +speech-model files may download again. Performance varies by device, browser, +model, and input. See [benchmarks](docs/BENCHMARKS.md) for measurements. + +## Music generation + +Open the music page, enter a prompt, and select **Generate song**. Leave lyrics +empty for an instrumental. When generation finishes, use the audio player or +download the stereo 48 kHz WAV. Songs can be 10 seconds to 4 minutes long. + +The first generation downloads **5.75 GB** from the FluidInference Hugging Face +mirror and caches it in browser storage (OPFS). Later generations still need +to read and prepare model data. Settings shows cache usage and lets you delete +the downloaded model. + +The site uses ACE-Step 1.5 Turbo in direct mode. The optional planner is +available in the underlying runtime but disabled on the public music page. +See the [ACE-Step README](packages/acestep/README.md) for implementation and +validation details. + +Choose **Split stems** on a finished song to get drums, bass, other, vocals, +and a derived instrumental, each with playback and download controls. DiCoSe +downloads another **623 MB** on first use. The site defaults to its faster +deterministic mode, which skips refinement. See the +[DiCoSe README](packages/dicose/README.md) for mode differences and requirements. + +## Run locally + +Requires Node.js `^20.19.0` or `>=22.12.0` and npm. From the repository root: + +```bash +npm ci +npm run acestep:build +npm run dicose:build +npm run dev +``` + +Open `http://localhost:5173/`. Build the workspace libraries before starting +the site: its imports resolve to their generated `dist/` files. + +```bash +npm run build # Build libraries, type-check, and bundle the site +npm run test:unit # Shared UI progress tests +npm run acestep:test # ACE-Step tests +npm run dicose:test # DiCoSe tests +npm run format:check # Formatting check +``` + +Model weights are downloaded at runtime and are excluded from the site build. +Set `VITE_ACE_MODEL_ORIGIN` to use another ACE model host or a local package +directory. The default host and package identities are in +[config.ts](src/engines/musicgen-acestep/config.ts). ## SDK @@ -40,219 +87,54 @@ import { ParakeetV3Engine } from "@fluidinference/fluidaudio-web/asr-parakeet"; import { decodeToMono16k } from "@fluidinference/fluidaudio-web"; const asr = new ParakeetV3Engine(); -await asr.load((p) => console.log(p.file, p.fraction)); -asr.setVocabulary(["NVIDIA", "Newrez"]); // optional: fuzzy-correct domain terms -asr.setItn(true); // optional: "twenty one" → "21" -const { text } = await asr.transcribe(await decodeToMono16k(fileArrayBuffer), { - // optional: transcription progress on long files, emitted at window boundaries - onProgress: (p) => console.log(`${(p.fraction * 100).toFixed(0)}% — ${p.processedSeconds.toFixed(0)}s / ${p.totalSeconds.toFixed(0)}s`), -}); -await asr.dispose(); +try { + await asr.load((p) => console.log(p.file, p.fraction)); + const audio = await decodeToMono16k(fileArrayBuffer); + const { text } = await asr.transcribe(audio); + console.log(text); +} finally { + await asr.dispose(); +} ``` -True streaming (EOU / Nemotron) and captions (v0.2.0): +Engine subpaths: `/asr-parakeet`, `/asr-whisper`, `/asr-nemotron`, +`/tts-kokoro` (English or Chinese), `/vad-silero`, `/diarization-sortformer`, +and `/eou-parakeet`. The published package can lag behind this repository; +check its version before using newer APIs. -```ts -import { MicCapture, segmentsToSrt } from "@fluidinference/fluidaudio-web"; -import { ParakeetEouEngine } from "@fluidinference/fluidaudio-web/eou-parakeet"; - -const engine = new ParakeetEouEngine(); -await engine.load(); - -// live: feed mic chunks, get cumulative text; events + word segments -const mic = new MicCapture(); -await mic.start(); -let pos = 0; -setInterval(async () => { - const { samples, total } = mic.since(pos); - const text = await engine.push(samples); // conformer caches carried — no re-decode - pos = total; - console.log(text, engine.streamEvents, engine.streamSegments); -}, 300); -// on stop: const final = await engine.finish(); engine.reset(); - -// batch: word timestamps → SRT captions -const r = await engine.transcribe({ samples, sampleRate: 16000 }); -const srt = segmentsToSrt(r.segments); // also: segmentsToVtt, groupCues -``` - -One tree-shakeable subpath per engine — `/asr-parakeet`, `/asr-whisper`, -`/asr-nemotron`, `/tts-kokoro` (`{ lang: "en" | "zh" }`), `/vad-silero`, -`/diarization-sortformer`, `/eou-parakeet` — plus `/registry` (enumerate -engines, instantiate via `entry.make()`), `/textnorm`, and `/vocab-rescorer`. -Requires a bundler with `new URL(..., import.meta.url)` asset + module-worker -support (Vite, webpack 5 work out of the box). The demo site consumes the -identical source, so every site gate doubles as SDK regression coverage. - -Release flow: bump `version` in the root `package.json` → `npm run sdk:pack` → -`cd dist-sdk && npm publish --access public`. - -## Engines - -Measured in-browser (Chrome/macOS, WebGPU, warm) on a real 284.5s recording via -the since-removed verify page — not a lab clip. RTFx = audio-seconds per wall-second (for TTS: -audio _generated_ per wall-second; not comparable to ASR). Each registry entry -carries a `category` (`stt` / `tts` / `analysis`) that routes it to the matching -demo page. - -| Engine | Model | RTFx | Notes | -| ------------------------ | --------------------------------- | --------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `asr-whisper` | Whisper (99 langs) | re-measuring | KV-cached decode, f16 weights; long-form chunking just landed (the prior 114× was measured on the first-30s-only bug and is retracted) | -| `asr-parakeet` | Parakeet TDT 0.6B v3 | **282×** (1hr file) | 2.15% WER LibriSpeech test-clean (core parity-gated vs the reference); worker-pool RNNT decode; opt-in ITN + custom vocabulary | -| `vad-silero` | Silero VAD v5 | 79× | WASM-SIMD (tiny sequential model by design) | -| `diarization-sortformer` | NVIDIA Sortformer 4-spk | 79× | windowed with 24-permutation overlap stitching | -| `tts-kokoro` | Kokoro 82M (en + zh) | 4.7× en / 5.6× zh | waveform corr ~0.97 vs reference; en input auto-normalized ("$4.50" is spoken, not dropped) | -| `asr-nemotron` | Nemotron 3.5 streaming (40 langs) | realtime+ | cache-aware streaming RNNT | -| `eou-parakeet` | Parakeet EOU 120M | **297×** browser-verified (1hr in 12.1s; worker-overlapped wasm decode + linear-cost stream-batch encode) | transcript + end-of-utterance events; TRUE streaming push()/finish() (bit-exact cache-carrying encode) + wasm-SIMD RNNT decode; whole-clip batch runs through the same linear-cost encoder | -| `asr-voicechat` | VoiceChat-11B STT (609M encoder) | 34.6× (1hr file) | the speech-recognition slice of NVIDIA's full-duplex VoiceChat-11B; fully-causal per-frame streaming, parity byte-identical to the torch reference; weights hosted at [`FluidInference/fluidaudio-web`](https://huggingface.co/FluidInference/fluidaudio-web) like the other engines | -| `musicgen-acestep` | ACE-Step 1.5 Turbo (3.5B + VAE) | ~1.9× (180s song in ~95s, M3, warm) | full text-to-music on [`/music`](music.html): 8-step DiT + Oobleck VAE in pure WGSL (`packages/acestep`); ~5.7 GB one-time download; requires `shader-f16`; direct mode (optional planner LLM path exists upstream, still being optimized) | -| `stem-dicose` | DiCoSe stem separation (BS-RoFormer + 1-step CD) | ~3× fast mode (30s / 48 kHz song in ~9.8s, M5 Pro); refined ~0.45× | 5 stems — drums/bass/other/vocals + derived instrumental — on [`/analyze.html`](analyze.html) and as **Split stems** on [`/music`](music.html); vendored `packages/dicose` (DiCoSe.wgsl by Hamza Qayyum); 623 MB f16 weight package; requires `shader-f16` + fixed 32-wide subgroups; mix-reconstruction NRMSE 1.5% (4-stem sum) / 9e-5 (vocals + instrumental) | -| `tts-voicechat` | VoiceChat-11B TTS “Aria” (595M backbone + 159M MoG + 763M codec) | 23.7 GPU-ms per 80 ms frame + ~16 ms host (timestamp-query, M5 Pro dawn; est. ~1.6× in-browser) — node wall is poll-bound at 0.12× (dawn ~100 ms/sync × 5 syncs/frame); WASM 0.19× | the speech-decoder slice of NVIDIA's full-duplex VoiceChat-11B as a standalone TTS voice; GPU-resident decode loop (backbone/MoG-MLP batched submits, GPU KV caches, 5 readbacks/frame for the host-side f64 PRVQ decisions); audio codes bit-exact vs the torch reference ON BOTH BACKENDS (1550/1550), waveform NRMSE 1.1e-6; codec GPU decode ~24 GPU-ms/s of audio; local-only weights (`scripts/extract-voicechat-tts.py`, ~3.5 GB) — hidden from the picker unless exported | - -First (cold) run is several× slower — WebGPU compiles pipelines and weights -download once. WebGPU is optional: every engine falls back to the same math on -WASM-SIMD (slower on the big encoders, identical outputs — cross-backend -parity is CI-gated). History and methodology: [`docs/BENCHMARKS.md`](docs/BENCHMARKS.md). - -## Music generation (ACE-Step) - -[`/music.html`](music.html) generates full songs — prompt, optional lyrics, up -to 4 minutes, stereo 48 kHz WAV — entirely client-side. The runtime is -[`packages/acestep`](packages/acestep/), a vendored npm-workspace import of -ace-step-1.5.wgsl by Hamza Qayyum (upstream repo not yet public — his live -demo is at [acestep.narcotic.sh](https://acestep.narcotic.sh); the vendored -source lives in this repo): -~100k lines of TypeScript + WGSL implementing the Qwen3 text encoder, ACE -condition encoder, 24-layer DiT, and Oobleck VAE decoder, with authenticated -streamed model packaging, bounded GPU memory, and cooperative scheduling. It -keeps its own rigorous experiment ledger (`packages/acestep/optimization/`) — -read `packages/acestep/AGENTS.md` before touching kernels there. - -The ~5.7 GB of content-addressed model packages currently stream from the -upstream author's public R2 bucket and cache in OPFS; set -`VITE_ACE_MODEL_ORIGIN` to point at a mirror or locally staged packages -(`packages/acestep/model/convert.py --profile production` reproduces the -exact tuple). The optional 0.6B planner ("thinking") path is excluded from -the served manifest until its pending optimization experiments -(OPT-0084/0085/0087) are integrated. - -A finished song offers **Split stems**: DiCoSe (also by Hamza Qayyum, -vendored at [`packages/dicose`](packages/dicose/)) separates the generated -WAV into drums, bass, other, vocals, and a derived instrumental, right in -the result panel — playable and downloadable per stem. Fast deterministic -mode by default (~3× realtime); the 623 MB weight package downloads on -first use and is released with the result panel. - -## Text processing (WASM) - -[`text-processing-rs`](https://github.com/FluidInference/text-processing-rs) -vendored as a 1 MB wasm module (pure Rust, no network): - -- **TN** (written → spoken) runs automatically on English Kokoro input: - numbers/currency/times aren't in the G2P lexicon and used to be silently - dropped from the audio. -- **ITN** (spoken → written, `"i paid four dollars and fifty cents"` → - `"i paid $4.50"`) is **opt-in** (`setItn(true)` / the STT page checkbox) — - on everyday speech it also rewrites phrases like "no one" → "no 1". - -## Quick start (repo) +Use a bundler that supports module workers and `new URL(..., import.meta.url)` +assets, such as Vite or webpack 5. The source SDK supports streaming, caption +exports, and optional Parakeet vocabulary correction and inverse text +normalization. See [streaming](docs/STREAMING.md) and [end-of-utterance detection](docs/EOU.md). -```bash -npm install -npm run dev # http://localhost:5173 — speech to text; /tts.html, /analyze.html, /live.html, /music.html -npm run build # static site → dist/ -npm run sdk:pack # publishable SDK tarball (dist-sdk/ + .tgz in repo root) +To prepare a release, bump the root `package.json` version and run +`npm run build`, `npm run sdk:test`, and `npm run sdk:pack`. The SDK test validates +the tarball in a clean consumer. The SDK registry includes only packaged engines; +VoiceChat, music generation, and DiCoSe remain site-only. -npm run acestep:check && npm run acestep:test # ACE-Step runtime (packages/acestep) gates -npm run dicose:check && npm run dicose:test # DiCoSe runtime (packages/dicose) gates -``` +## Development references -## Deploy +- [Engine catalog](src/engines/registry.ts) +- [Architecture](docs/ARCHITECTURE.md) and [adding a model](docs/PORTING.md) +- [WebGPU implementation](docs/RAW_WEBGPU.md) and [ONNX Runtime removal](docs/ORT_REMOVAL.md) +- [ACE-Step development rules](packages/acestep/AGENTS.md) and [optimization ledger](packages/acestep/optimization/LEDGER.md) +- [DiCoSe correctness audit](packages/dicose/optimization/CORRECTNESS_AUDIT.md) -`main` auto-deploys to **Cloudflare Workers** (static assets, see -`wrangler.jsonc`) via the connected Workers Builds integration — merge and it's -live. Manual: `npm run build && npx wrangler deploy`. +GitHub Pages deploys `main` only after all [CI jobs](.github/workflows/ci.yml) +pass. Cloudflare Workers uses a separate deployment integration; see the +[Worker configuration](wrangler.jsonc). -Deliberately **no COOP/COEP**: cross-origin isolation would break the -cross-origin Hugging Face weight fetches, and nothing here needs -`SharedArrayBuffer` — parallelism comes from WebGPU and the decode worker pool -(each worker gets its own weight copy). +## Credits and licenses -## Layout +Code is MIT licensed; model weights retain their upstream licenses. See +[third-party licenses](THIRD-PARTY-LICENSES.md). -``` -src/ - gpu/ the kernel library: WGSL GEMM/conv/attention/LSTM (compute.js), - WASM-SIMD twin (wasm-context.js) — one interface, two backends - engines/ one folder per model on those kernels; registry.ts is the catalog - core/ audio I/O, model cache, text normalization, shared types - index.ts SDK root (engines are subpath exports) - pages/ playground.ts — shared pick→load→run core for the demo pages - stt.ts / tts.ts / analyze.ts / live.ts / music.ts demo pages (thin consumers of the registry / music client) -packages/ - acestep/ vendored ace-step-1.5.wgsl music-gen runtime (own kernels, - scheduler, tests, and optimization ledger — see its AGENTS.md) - dicose/ vendored DiCoSe.wgsl stem-separation runtime (own kernels, - tests, and optimization ledger) -scripts/ node gates: token-identity, kernel parity, per-engine smokes -rust/ parakeet RNNT decoder + kernel lib sources (wasm32+simd128) -docs/ architecture, benchmarks, PORTING.md (add-a-model checklist), the ORT removal story -``` +Hamza Qayyum ([Narcotic Software](https://narcotic.sh)) built the ACE-Step +browser port and DiCoSe WebGPU runtime, vendored under `packages/`. FluidInference +integrated them and continues their development. His original ACE-Step demo is +at [acestep.narcotic.sh](https://acestep.narcotic.sh). -## Hard-won lessons (things that cost real debugging) - -- **Wall-clock lies under dawn/node; only `timestamp-query` tells the truth.** - Every kernel "benchmark" read ~2 ms/op until per-dispatch GPU timestamps - showed the real distribution — several optimization verdicts flipped. -- **WebGPU errors are async and silent.** A missing `shader-f16` feature - request turned every f16 GEMM into a no-op: empty transcripts at a - fake-fast RTFx. Feature-gate every `enable` directive and log - `uncapturederror`. -- **Synchronous WASM starves microtasks.** A `.then()` holding a GPU readback - couldn't fire while a 190 ms decode blocked the thread — the GPU idled after - every batch. Staging copies must ride the producing submit. -- **Measure on the target machine.** The dev box was CPU-bound where user - machines were GPU-bound and vice versa; the per-stage split in the metrics - (`mel / encode / decode`) exists because RTFx alone misdiagnosed both. -- **f16 storage ≠ f16 compute.** Halving weight bytes did nothing on a - compute-bound GPU; switching the inner loop to f16 fma (f32 accumulate per - 8-deep K-tile) bought 1.46× with token-identical output. -- **ITN is not a free win.** English inverse normalization rewrites "no one" → - "no 1" and deletes words in other languages — it shipped opt-in only because - a review pass ran the wasm on realistic sentences. -- **Gate on tokens, not maxΔ.** Every perf change here ships with a - token-identity / parity gate; two of them caught real kernel breakage that - numeric thresholds would have argued about. - -## License - -MIT (code). Model weights follow their upstream licenses (see the registry and -Hugging Face model cards). - -## Acknowledgements - -The encoder GEMM kernel geometry and the GPU TDT decoder design are adapted -from [parakeet.wgsl](https://github.com/narcotic-sh/parakeet.wgsl) by -Narcotic Software (MIT) — a fast, focused browser Parakeet implementation -that served as both inspiration and reference throughout our optimization -work. - -Music generation is built on -ace-step-1.5.wgsl ([live demo](https://acestep.narcotic.sh); upstream repo -not yet public — the full source is vendored at `packages/acestep`) by -**Hamza Qayyum** (Narcotic Software, MIT): he built the complete ACE-Step -1.5 Turbo browser port — correctness-gated WGSL kernels, model packaging, -scheduling, and the Stage-2 optimization program — and handed the project -over for integration here; we took it over, integrated, and are continuing -the optimization work. The `packages/acestep` runtime and the `/music` page's -backend seam are his code. - -Stem separation is likewise his: DiCoSe.wgsl (vendored at -`packages/dicose`, MIT) ports DiCoSe — BS-RoFormer plus one-step -consistency-distilled refinement, [karchkha/DiCoSe](https://huggingface.co/karchkha/DiCoSe) -checkpoints — to raw WebGPU WGSL with its own correctness-audited kernel -ledger, and powers both the `stem-dicose` engine and the `/music` page's -Split stems feature. - -See [THIRD-PARTY-LICENSES.md](./THIRD-PARTY-LICENSES.md) for the full -list of adapted code and licenses. +The Parakeet encoder GEMM layout and GPU decoder design are adapted from +[parakeet.wgsl](https://github.com/narcotic-sh/parakeet.wgsl). Text normalization +uses the vendored [text-processing-rs](https://github.com/FluidInference/text-processing-rs) +WASM module. diff --git a/package-lock.json b/package-lock.json index c483086..4efe7d5 100644 --- a/package-lock.json +++ b/package-lock.json @@ -18,10 +18,12 @@ }, "devDependencies": { "@kmamal/gpu": "^0.2.1", - "onnxruntime-node": "^1.27.0", + "onnxruntime-node": "^1.30.0", "prettier": "^3.9.6", "typescript": "^5.6.3", - "vite": "^6.0.0" + "vite": "^6.0.0", + "vitest": "^4.1.11", + "yaml": "^2.9.1" }, "engines": { "node": "^20.19.0 || >=22.12.0" @@ -1366,13 +1368,13 @@ "link": true }, "node_modules/adm-zip": { - "version": "0.5.18", - "resolved": "https://registry.npmjs.org/adm-zip/-/adm-zip-0.5.18.tgz", - "integrity": "sha512-ufJnssQGbxzLNS1Ho9bCtX4rQKCCvoVuDLHoJyc3F9dOGDB4BkWs2Ci0kv53lqocAEQ/Cbi+I2XCsNYGqVYqng==", + "version": "0.6.1", + "resolved": "https://registry.npmjs.org/adm-zip/-/adm-zip-0.6.1.tgz", + "integrity": "sha512-Xwrja8nx9e5o2N1my4DsKCeKpdrnACyr1wtbPxBDgGzKzKyE9kRtBFA8mWldI+RVlD7CBZNWY/wQ2+ydwOR6kQ==", "dev": true, "license": "MIT", "engines": { - "node": ">=12.0" + "node": ">=14.0" } }, "node_modules/assertion-error": { @@ -2022,16 +2024,16 @@ } }, "node_modules/onnxruntime-common": { - "version": "1.27.0", - "resolved": "https://registry.npmjs.org/onnxruntime-common/-/onnxruntime-common-1.27.0.tgz", - "integrity": "sha512-3KxL5wIVqa8Ex08jxSzncm9CMgw8CjOFyOQ7SxvG9o0cVLlhTNKXyIQuTbtX4tGPJEf73OER2xrjt4HJSBL4ow==", + "version": "1.30.0", + "resolved": "https://registry.npmjs.org/onnxruntime-common/-/onnxruntime-common-1.30.0.tgz", + "integrity": "sha512-7fdVWjAID1dVhH/G8qK3APARunV4VkBFoCQAP7qp4Wkab0mrorvmc+sqiT+mKXOzDqdjN5j+/Z9nb4gzNPWcyA==", "dev": true, "license": "MIT" }, "node_modules/onnxruntime-node": { - "version": "1.27.0", - "resolved": "https://registry.npmjs.org/onnxruntime-node/-/onnxruntime-node-1.27.0.tgz", - "integrity": "sha512-QEzGwrvNBgv4uPVdnbHsOGG4G6T96mdlcFI8aAKPjMU8wOPpVocPXb6k3QGkaZagVTv2G9Bnnbo6Z3JdXr1fQw==", + "version": "1.30.0", + "resolved": "https://registry.npmjs.org/onnxruntime-node/-/onnxruntime-node-1.30.0.tgz", + "integrity": "sha512-twhs1C2C/BFkz1yc5OY0KIU2GUq6DURO7hD4bx5Q2Qy3nAMJwRXW8xU3NVczE29VA9lolLOYepoD8fjTGOfIqw==", "dev": true, "hasInstallScript": true, "license": "MIT", @@ -2041,9 +2043,9 @@ "linux" ], "dependencies": { - "adm-zip": "^0.5.16", + "adm-zip": "^0.6.0", "global-agent": "^4.1.3", - "onnxruntime-common": "1.27.0" + "onnxruntime-common": "1.30.0" } }, "node_modules/pathe": { @@ -2551,6 +2553,22 @@ "node": ">=18" } }, + "node_modules/yaml": { + "version": "2.9.1", + "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.9.1.tgz", + "integrity": "sha512-3NxN8+78OdzbT7C/WjGsyfPAtJaN3FNDsWxv7Y7mcDsT/oOmgW8BpyQQFFBnvZE3j9Y2Sdz1ULFLezL7Eb2yFw==", + "dev": true, + "license": "ISC", + "bin": { + "yaml": "bin.mjs" + }, + "engines": { + "node": ">= 14.6" + }, + "funding": { + "url": "https://github.com/sponsors/eemeli" + } + }, "packages/acestep": { "name": "ace-step-1.5.wgsl", "version": "0.0.0", diff --git a/package.json b/package.json index e287e47..5130815 100644 --- a/package.json +++ b/package.json @@ -45,6 +45,7 @@ "build": "npm run acestep:build && npm run dicose:build && tsc --noEmit && vite build && node scripts/postbuild-strip-wasm.mjs", "preview": "vite preview", "test:unit": "node --test test/*.test.mjs", + "test:integration": "vitest run --config vitest.config.ts", "typecheck": "tsc --noEmit", "gpu:verify": "node scripts/gpu-verify.mjs", "gpu:bench": "node scripts/gpu-bench.mjs", @@ -57,11 +58,12 @@ "wasm:build": "cd rust/wasm-kernels && RUSTFLAGS='-C target-feature=+simd128' cargo build --release --target wasm32-unknown-unknown && cp target/wasm32-unknown-unknown/release/wasm_kernels.wasm ../../src/gpu/wasm-kernels.wasm", "wasm:verify": "node scripts/wasm-verify.mjs", "backend:parity": "node scripts/backend-parity.mjs", - "format": "prettier --write \"src/**/*.{ts,js}\" \"scripts/**/*.mjs\" \"test/**/*.mjs\" \"*.ts\"", - "format:check": "prettier --check \"src/**/*.{ts,js}\" \"scripts/**/*.mjs\" \"test/**/*.mjs\" \"*.ts\"", + "format": "prettier --write \"src/**/*.{ts,js}\" \"scripts/**/*.mjs\" \"test/**/*.{ts,mjs}\" \"*.ts\"", + "format:check": "prettier --check \"src/**/*.{ts,js}\" \"scripts/**/*.mjs\" \"test/**/*.{ts,mjs}\" \"*.ts\"", "ci:smoke": "node scripts/interface-conformance.mjs && node scripts/ci-smoke-parakeet.mjs && node scripts/ci-smoke-eou.mjs && node scripts/ci-smoke-kokoro.mjs", "sdk:build": "node scripts/build-sdk.mjs", "sdk:pack": "node scripts/build-sdk.mjs && cd dist-sdk && npm pack --pack-destination ..", + "sdk:test": "npm run sdk:build && node scripts/test-sdk.mjs", "conformance": "node scripts/interface-conformance.mjs", "conformance:numeric": "node scripts/backend-conformance.mjs" }, @@ -72,9 +74,11 @@ }, "devDependencies": { "@kmamal/gpu": "^0.2.1", - "onnxruntime-node": "^1.27.0", + "onnxruntime-node": "^1.30.0", "prettier": "^3.9.6", "typescript": "^5.6.3", - "vite": "^6.0.0" + "vite": "^6.0.0", + "vitest": "^4.1.11", + "yaml": "^2.9.1" } } diff --git a/packages/acestep/README.md b/packages/acestep/README.md index 15c0206..accba55 100644 --- a/packages/acestep/README.md +++ b/packages/acestep/README.md @@ -1,164 +1,91 @@ # ace-step-1.5.wgsl -An in-progress, fully custom WebGPU and WebAssembly browser port of ACE-Step -1.5 Turbo. - -The project aims to generate high-quality music locally on consumer Apple -Silicon without a server or generic browser ML runtime. It is inspired by -[`parakeet.wgsl`](https://github.com/narcotic-sh/parakeet.wgsl), including its -streamed model packaging, bounded GPU memory, specialized WGSL kernels, worker -runtime, and queue-drained cooperative GPU scheduling. - -## Project status - -Stage 1's browser correctness and human-listening gate passed on 2026-08-13. -The complete, untuned browser pipeline produced owner-approved instrumental and -vocal songs, including the default-CoT planner path. Stage 2 measured -optimization is now authorized; this remains a development baseline rather than -a release. - -The first foundation includes the pinned request/sampler and worker contracts, -browser capability reporting, browser-defined Gaussian and categorical random -transforms, authenticated golden-fixture contracts, and a deterministic -model-package converter. The runtime authenticates manifest bytes before strict -parsing, validates the canonical package schema, streams pinned assets through -bounded transactions, provides an integrity-aware resumable OPFS primitive, -and owns GPU graphs with Parakeet-style FIFO cooperative submission. GPU weight -uploads additionally enforce a 64 MiB queued-write high-water mark, drain the -queue before crossing it, and insert a real queue-empty interval between -non-final upload batches. - -The untuned correctness kernels now cover GEMM, RMSNorm, authenticated Qwen -RoPE, score-buffer-free full/sliding/causal GQA attention, KV-cache writes, -shard-aware embeddings, transformer and DiT plumbing, fixed FSQ decoding, the -pinned Haar DCW update, and the complete FP32 Oobleck VAE primitive family. -Correctness-first graph composers cover the 28-layer text encoder, Qwen cached -blocks, semantic detokenizer, direct lyric/silence-timbre conditioner, all 24 -DiT layers and sampler operations, and the 88-operation VAE decoder. The VAE -decoder is split into bounded FIFO output-domain quanta; even a single large -convolution cannot become one GPU-monopolizing dispatch. - -The planner control path reproduces left-padded one/two-row prefill, bounded KV -capacity, deterministic browser-owned filtering/sampling, the normal two-phase -prompt workflow, a PyYAML-compatible output subset, and the pinned metadata -constraint state machine. Independent Python vectors authenticate its field -tries, language choice, caption mask, and upstream's unusual early reasoning -termination. Small CPU-oracle cases for every current graph family pass in -Chrome/WebGPU on the development M3 in both package profiles. Concrete -whole-model owners now connect authenticated model acquisition, -planner-optional conditioning, all eight DiT evaluations, chunked VAE decode, -transactional OPFS audio, and cooperative/cancellable scheduling. Native -production-weight taps remain unclaimed external parity evidence; the browser -implementation is not allowed to serve as its own numerical oracle. The -owner-approved direct and planner-enabled reference-profile WAVs and their -exact settings are recorded in -[`LISTENING_CANDIDATE.md`](./LISTENING_CANDIDATE.md). - -The M3 actual-browser safety harness also exercises the largest current tiled -GEMM shape with exact sentinels while continuously checking a UI heartbeat. -Those runs establish bounded cooperative dispatch behavior only; they are not -Stage 2 performance benchmarks. - -The implementation program is in [`PLAN.md`](./PLAN.md). Its central rule is: - -1. implement and numerically validate the complete scoped browser pipeline; -2. stop for a human listening test; and -3. begin performance optimization only after explicit approval. - -The listening packet records measured M3 correctness-run latency. Stage 2 now -replaces those untuned measurements with thermally controlled baselines and a -durable experiment ledger. The headline M3 target is a warm-cache three-minute -song in under one minute without leaving the approved quality, memory, or -responsiveness envelope. M5 performance will be tested later on actual hardware -rather than claimed from scaling alone. - -## Initial scope - -- ACE-Step 1.5 Turbo, eight denoising evaluations -- text prompts and supplied lyrics -- optional full ACE 5 Hz 0.6B thinking/semantic planner workflow -- direct generation without the planner -- semantic-code FSQ projection and learned detokenizer for the planner path -- upstream-default Haar DCW correction in the eight-step Euler sampler -- default -1 dBFS peak normalization -- stereo 48 kHz WAV output -- Chrome on desktop Apple Silicon first -- uncompressed, phase-staged weights first - -Weight compression, larger planners, source-audio workflows, Safari tuning, -and iPhone support are later projects. - -The pinned direct, no-cover path evaluates a silence acoustic-tokenizer branch -upstream and then discards it through an all-false `is_covers` selection. The -browser package removes that dead encode-only branch exactly; cover/source -audio remains outside v1 scope. - -## Source identities +A custom ACE-Step 1.5 Turbo browser runtime by Hamza Qayyum, using WebGPU and +WebAssembly. It powers [FluidAudio Web's music page](../../music.html). -The correctness baseline is pinned to: +## Features and requirements + +- Text prompts, optional lyrics, and deterministic seeds. +- Direct generation or an optional 0.6B planner for metadata and semantic codes. +- Eight denoising evaluations with the pinned Haar DCW sampler. +- Stereo 48 kHz WAV output, normalized to −1 dBFS. +- Inference in a dedicated worker, with progress and cancellation. +- Authenticated model downloads and persistent OPFS caching. + +Desktop Chrome on Apple Silicon is the primary validation target. The +production profile requires WebGPU with `shader-f16`; subgroup optimizations +have a portable fallback. Source-audio editing, cover generation, and the VAE +encoder are outside the current scope. Support for phones and other browsers +must be validated on the target device. + +The FluidAudio music page uses direct generation and a **5.75 GB** model cache. +Its planner is disabled. The package's development demo also exposes the +planner and uses a different reference manifest. The site's selected model +packages are defined in [config.ts](../../src/engines/musicgen-acestep/config.ts). -- ACE-Step source commit - `6d467e4b5081ccb0abf1ec1bf4fdf9051a2d34b0`; -- `ACE-Step/Ace-Step1.5` revision - `19671f406d603126926c1b7e2adc169acbcade22`; and -- `ACE-Step/acestep-5Hz-lm-0.6B` revision - `148d8ea0225bdab342ee1ae3a354275ccd60ca80`. +## Status and validation -The browser model packages are generated reproducibly by `model/convert.py`. -`--profile production` downloads and creates the exact three-package tuple used -by the demo, failing closed unless all committed manifest identities match. -Model weights, conversion caches, and generated audio do not belong in Git. -The initial uncompressed profiles and current production tuple have canonical -manifest hashes recorded in -[`model/canonical-packages.json`](./model/canonical-packages.json). +The owner approved the initial browser listening results on August 13, 2026, +and authorized measured optimization. The project remains in development. +Approved audio identities are recorded in [LISTENING_CANDIDATE.md](LISTENING_CANDIDATE.md). -## Development +Unit tests and browser kernel checks validate their declared cases; they do +not establish whole-model equivalence to the native upstream implementation. +Native CUDA/XPU captures remain required before making that claim. -The browser-side project uses the pnpm version pinned in `package.json`: +[PLAN.md](PLAN.md) defines release criteria, including reliable operation on +the 16 GB M3, stock-Chrome validation, cancellation and cleanup, and a warm-cache +three-minute song in under one minute. Initial download time is measured +separately. See the [optimization ledger](optimization/LEDGER.md) for experiment +results and integration status. + +## Build and test + +From the FluidAudio Web repository root: ```bash -corepack enable -pnpm install --frozen-lockfile -pnpm check -pnpm test -pnpm build +npm ci +npm run acestep:check +npm run acestep:test +npm run acestep:build ``` -Generate the current production weights and run the browser demo at -`http://127.0.0.1:5174/demo/`: +These checks do not require downloading the model. For actual GPU validation, +follow the [browser test guide](test/browser/README.md). + +To run the public music UI, follow the [root quick start](../../README.md#run-locally). +To prepare weights and run the package's development demo, use `uv` with the +locked Python 3.13 environment. From `packages/acestep/`: ```bash uv run --frozen --project model --python 3.13 \ python3 model/convert.py --profile production -pnpm exec vite --host 127.0.0.1 --port 5174 +npm exec -- vite --host 127.0.0.1 --port 5174 ``` -The demo defaults to the short direct/instrumental golden request, authenticates -the canonical packed-BF16 reference manifest, uses the cooperative scheduling -profile, and keeps generated audio in transactional OPFS storage until the UI -releases it. The raw-FP16 profile is an explicitly lower-precision diagnostic -until it passes the numerical and listening gates required by `AGENTS.md`. +Open `http://127.0.0.1:5174/demo/`. See [model preparation](model/README.md) for +disk requirements and package layout. Generated weights, audio, caches, and +large profiles stay out of Git. -The model workspace is a frozen uv project on CPython 3.13. Its complete -source-to-package commands and disk requirements are documented in -[`model/README.md`](./model/README.md). The weight-free checks are: +## Implementation and source identities -```bash -uv run --python 3.13 python3 scripts/validate-golden.py -uv run --python 3.13 python3 -m unittest discover \ - -s scripts -p 'test_*.py' -v -uv run --frozen --project model --python 3.13 \ - python3 -m unittest discover -s model/tests -v -uv run --frozen --project reference --python 3.13 \ - python3 -m unittest discover -s reference/tests -v -``` +The pipeline combines a Qwen3 text encoder, ACE condition encoder, optional +planner and semantic detokenizer, 24-layer DiT, and Oobleck VAE decoder. Weights +are loaded by phase and audio is decoded in chunks to bound memory use. GPU +submissions are bounded to support cancellation and keep the UI responsive. + +The correctness baseline is pinned to: + +- ACE-Step source: `6d467e4b5081ccb0abf1ec1bf4fdf9051a2d34b0` +- Main model: `19671f406d603126926c1b7e2adc169acbcade22` +- 0.6B planner: `148d8ea0225bdab342ee1ae3a354275ccd60ca80` -The actual-browser kernel harness is documented in -[`test/browser/README.md`](./test/browser/README.md). Passing it proves only the -small declared correctness cases; it is not a performance benchmark. +[Canonical package hashes](model/canonical-packages.json) identify reproducible +converter outputs. Read [AGENTS.md](AGENTS.md) before changing kernels, +precision, scheduling, or model packaging; it defines the required numerical, +listening, and performance checks. ## License -The original source in this repository is MIT licensed. ACE-Step and Qwen -artifacts retain their respective upstream terms and attribution; see -[`THIRD_PARTY_LICENSES`](./THIRD_PARTY_LICENSES). +Original code is MIT licensed. ACE-Step and Qwen artifacts retain their upstream +terms. See [THIRD_PARTY_LICENSES](THIRD_PARTY_LICENSES). diff --git a/packages/dicose/README.md b/packages/dicose/README.md index 7f61fa1..f600f5d 100644 --- a/packages/dicose/README.md +++ b/packages/dicose/README.md @@ -1,120 +1,102 @@ # DiCoSe WebGPU -An interactive and automation-ready browser runtime for the released **[DiCoSe BS-RoFormer + one-step -consistency-distilled (CD) refinement](https://arxiv.org/abs/2412.06965)**, using the -[official model weights](https://huggingface.co/karchkha/DiCoSe). The neural graph is raw WGSL/WebGPU: -f16 storage, f32 reductions, converter-native tile-major subgroup GEMM, fused -online attention with producer-rotated K, RMSNorm, FiLM, Conv2d STFT -conditioning, complex masks, and CD affine sampling. -The CPU boundary is intentionally limited to WAV decoding, deterministic -resampling, centered Hann STFT/ISTFT, seeded noise generation, and the final -instrumental complement subtraction. - -The public API lives in `src/index.ts`. Browser inference runs in a dedicated -worker and transfers PCM/result buffers instead of blocking the page thread. - -## Inference modes - -`new DiCoSeWorkerClient()` and `separateAudio(source)` keep the released -full-resolution, one-step refined graph as the default. Fast is an explicit -quality/performance tradeoff that returns the deterministic separator before -CD refinement: +A browser stem separator by Hamza Qayyum, using the +[DiCoSe BS-RoFormer and one-step refinement model](https://arxiv.org/abs/2412.06965) +with [official weights](https://huggingface.co/karchkha/DiCoSe). -```ts -const fast = new DiCoSeWorkerClient(); -const fastResult = await fast.separateAudio(source, { - outputMode: "deterministic", -}); -``` +Inference runs in a dedicated WebGPU worker. Input audio and generated stems +stay in the browser. The runtime returns drums, bass, other, vocals, and a +derived instrumental, restored to the input's sample rate and exact frame count. -Both modes return the four neural estimates under `result.stems` and a derived -`result.instrumental`. Instrumental is computed as the decoded input mixture -minus the vocal estimate after both have been restored to the uploaded file's -sample rate and exact frame count. It adds no model pass and is deliberately -not computed by summing drums, bass, and other. +## Requirements and modes -On the supplied WAV, deterministic-only had a 5.92-s sustained median. That -number remains useful as a performance measurement, not quality evidence. -Fast uses the released deterministic checkpoint but omits learned refinement. -See `optimization/CORRECTNESS_AUDIT.md` and `optimization/LEDGER.md` for the -current evidence and dispositions. +Requires HTTPS or localhost, WebGPU with `shader-f16`, fixed 32-wide subgroups, +1 GiB GPU buffers/storage bindings, and 25,344 bytes of workgroup storage. +Use `checkSupport()` to check the device before loading the **623 MB** model. -## Model package +| Mode | API value | Behavior | +| ---- | --------------- | -------------------------------------------------------------------------- | +| Full | `refined` | Deterministic separation followed by one-step refinement; package default. | +| Fast | `deterministic` | Skips refinement; default in FluidAudio Web. | -Large checkpoints, download caches, and the generated weight blob are ignored -by Git. Run this command manually whenever the local browser package needs to -be prepared: - -```sh -pnpm model:prepare -``` +Fast trades refinement quality for speed. For long files, Full uses 50% chunk +overlap and Fast uses 10%; both process bounded 11-second model items with +normalized overlap-add. The instrumental is the restored input mixture minus +vocals, not the sum of the other stems. -It downloads the two pinned official checkpoints, verifies them, converts the -exact Full/Fast production package into `public/model/`, and verifies the -canonical generated hashes. The source download is about 4.66 GB and is cached -under `model/cache/`. The command requires `uv`; Python 3.13 and all converter -dependencies come from the locked `model/` environment. +Measurements and quality evidence are in the +[correctness audit](optimization/CORRECTNESS_AUDIT.md) and +[optimization ledger](optimization/LEDGER.md). Performance depends on the mode, +input length, and hardware. ## Run locally -```sh -pnpm dev +From the FluidAudio Web repository root: + +```bash +npm ci +npm run dicose:check +npm run dicose:test +npm run dicose:build +``` + +For the integrated site, follow the [root quick start](../../README.md#run-locally). +For the standalone package demo, run these commands from `packages/dicose/`: + +```bash +npm run model:prepare +npm run dev ``` -Open `http://127.0.0.1:5173/`, choose or drop a local WAV, select Full or Fast, -and run the separation. The page shows stage timings and -creates the four model stems plus a derived instrumental as five in-memory -stereo WAVs with playback and download controls. -The source file and generated outputs stay in the browser tab; they are not -uploaded. Inputs above 12 seconds are processed as fixed 11-second model items -with reflected context and normalized overlap-add. Full retains the upstream -50% overlap policy. Fast overlaps only the existing 10% fade region. For the -5,608,109-sample model-rate `trust_nobody.wav` input, that changes the plan from -25 chunks in Full to 13 in Fast. A fresh isolated-Chrome sustained panel measured -a 79.61-s median (71.57–113.05 s) for that Fast path; this does not meet the -30-second target, and listening remains the quality gate. Long tracks still -require serial model calls. File-based runs restore each output to the uploaded -WAV's sample rate and exact frame count before playback/download. - -For an unattended page invocation, add `?autorun=1`; the result is published -to `window.__DICOSE_BROWSER__.report` and `#result`. `?mode=benchmark` uses a -single persistent worker/model package across its warmup and measured runs. -Neither path opens a save dialog, download, or UI control. - -## Checks and isolated browser testing - -```sh -pnpm check -pnpm test -pnpm test:reference-quality -pnpm test:refined-reference-quality -pnpm test:output-mode-quality -pnpm test:webgpu -pnpm test:browser -pnpm benchmark:browser +Model preparation requires `uv`. It uses the locked Python 3.13 environment, +downloads and verifies about **4.66 GB** of source checkpoints, and writes the +browser package to `public/model/`. Source downloads are cached in +`model/cache/`; these files stay out of Git. See [model preparation](model/README.md). + +Open `http://127.0.0.1:5173/`, select a local WAV and Full or Fast, then run +separation. Each stem has playback and WAV download controls. + +## API + +The public API is exported from [src/index.ts](src/index.ts). Within this workspace: + +```ts +import { DiCoSeWorkerClient } from "dicose-wgsl"; + +const client = new DiCoSeWorkerClient(); +try { + const result = await client.separateAudio(source, { + outputMode: "deterministic", + }); + // Four model outputs: result.stems; derived output: result.instrumental. +} finally { + await client.dispose(); +} ``` -The release benchmark accepts an explicit output selector: +`source` is a `Blob` or `ArrayBuffer`. Set `manifestUrl` in the constructor to +use another model host. Consumers bundling the prebuilt library may need +`createWorker` to provide a worker that imports `dicose-wgsl/worker`; see the +[FluidAudio integration](../../src/engines/stem-dicose/index.ts). + +## Browser checks -```sh -DICOSE_BENCHMARK_OUTPUT_MODE=deterministic pnpm benchmark:browser +From `packages/dicose/`, with model files prepared: + +```bash +npm run test:webgpu +npm run test:browser +npm run test:reference-quality +npm run test:refined-reference-quality +npm run test:output-mode-quality +DICOSE_BENCHMARK_OUTPUT_MODE=deterministic npm run benchmark:browser ``` -The browser scripts automatically start Vite and a new headless Chrome process -with a freshly-created temporary `--user-data-dir`, then delete that profile, -stop Chrome, and stop Vite in `finally`. They use CDP to await the automatic -result; no user profile, click, permission prompt, or file save is involved. -`test:browser` additionally enforces the fixture's f16 deterministic-output -envelope against the upstream f32 reference. `test:reference-quality` checks -the released deterministic graph and 30 internal tensor seams; -`test:refined-reference-quality` checks the released one-step CD graph, 17 -internal CD seams, its raw model output, and the final refined stems against a -fixed-noise execution of the official PyTorch implementation. - -`Mixture_audio_1.wav` is the supplied 22.05 kHz mono fixture. Production decode -duplicates it to stereo and uses the Hann-windowed sinc geometry and defaults -from torchaudio 2.0.2 before processing 1,189 centered-STFT frames. The -deterministic model-arithmetic oracle deliberately replays its older frozen -linear input tensor so resampler and neural-graph regressions remain separate -gates. The CD sampler uses a fixed default noise seed, so an otherwise -identical run is reproducible. +The scripts start Vite and headless Chrome with a temporary profile, then clean +them up. Reference checks compare deterministic and refined outputs and +intermediate tensors against the official implementation. The fixed default +noise seed makes otherwise identical runs reproducible. + +For page automation, `?autorun=1` publishes the report to +`window.__DICOSE_BROWSER__.report` and `#result`. `?mode=benchmark` reuses one +worker across warmup and measured runs. diff --git a/scripts/build-sdk.mjs b/scripts/build-sdk.mjs index 4aebf52..a78c4f5 100644 --- a/scripts/build-sdk.mjs +++ b/scripts/build-sdk.mjs @@ -21,9 +21,8 @@ execSync("npx tsc -p tsconfig.sdk.json", { stdio: "inherit" }); console.log("── copy runtime assets (.js/.d.ts/.wasm/.bin/.json)"); const KEEP = /\.(js|d\.ts|wasm|bin|json)$/; -// musicgen-acestep imports the private ace-step-1.5.wgsl workspace package; -// asr-voicechat weights are local-only — neither ships in the SDK yet. -const SKIP_DIRS = new Set(["__pycache__", "musicgen-acestep", "asr-voicechat"]); +// These engines are site-only; the SDK registry contains only packaged engines. +const SKIP_DIRS = new Set(["__pycache__", "musicgen-acestep", "asr-voicechat", "stem-dicose", "tts-voicechat"]); let copied = 0; function walk(dir) { for (const name of readdirSync(dir)) { @@ -49,7 +48,7 @@ console.log("── manifest + docs"); const ENGINE_SUBPATHS = ["asr-parakeet", "asr-whisper", "asr-nemotron", "tts-kokoro", "vad-silero", "diarization-sortformer", "eou-parakeet"]; const exports_ = { ".": { types: "./index.d.ts", default: "./index.js" }, - "./registry": { types: "./engines/registry.d.ts", default: "./engines/registry.js" }, + "./registry": { types: "./engines/sdk-registry.d.ts", default: "./engines/sdk-registry.js" }, "./textnorm": { types: "./core/textnorm.d.ts", default: "./core/textnorm.js" }, "./vocab-rescorer": { types: "./engines/asr-parakeet/vocab-rescorer.d.ts", default: "./engines/asr-parakeet/vocab-rescorer.js" }, "./captions": { types: "./core/captions.d.ts", default: "./core/captions.js" }, diff --git a/scripts/test-sdk.mjs b/scripts/test-sdk.mjs new file mode 100644 index 0000000..ac08593 --- /dev/null +++ b/scripts/test-sdk.mjs @@ -0,0 +1,65 @@ +// Verify the tarball outside the monorepo so workspace dependencies cannot hide missing files. +import assert from "node:assert/strict"; +import { execFileSync } from "node:child_process"; +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { pathToFileURL } from "node:url"; +import { build } from "vite"; + +const consumer = mkdtempSync(join(tmpdir(), "fluidaudio-sdk-test-")); +try { + const packed = JSON.parse( + execFileSync("npm", ["pack", "--json", "--pack-destination", consumer], { + cwd: resolve("dist-sdk"), + encoding: "utf8", + }), + ); + writeFileSync(join(consumer, "package.json"), JSON.stringify({ private: true, type: "module" })); + execFileSync("npm", ["install", "--ignore-scripts", "--no-audit", "--no-fund", join(consumer, packed[0].filename)], { + cwd: consumer, + stdio: "inherit", + }); + const manifest = JSON.parse(readFileSync(join(consumer, "node_modules/@fluidinference/fluidaudio-web/package.json"), "utf8")); + // Both entry points must resolve only complete engines. Keep every exported subpath reachable. + const specifiers = Object.keys(manifest.exports).filter((key) => key !== "./package.json"); + const source = specifiers + .map((key, i) => `import * as m${i} from ${JSON.stringify("@fluidinference/fluidaudio-web" + (key === "." ? "" : key.slice(1)))};\nexport { m${i} };`) + .join("\n"); + const entry = join(consumer, "consumer.js"); + writeFileSync(entry, source); + await build({ + root: consumer, + configFile: false, + logLevel: "warn", + build: { outDir: "browser", target: "es2022", lib: { entry, formats: ["es"] } }, + worker: { format: "es" }, + }); + // Bundle JSON imports for Node, then instantiate the real engines without downloading weights. + await build({ + root: consumer, + configFile: false, + logLevel: "warn", + ssr: { noExternal: ["@fluidinference/fluidaudio-web"] }, + build: { ssr: entry, outDir: "server" }, + }); + const modules = await import(pathToFileURL(join(consumer, "server/consumer.js")).href); + const expected = ["vad-silero", "asr-parakeet", "asr-whisper", "diarization-sortformer", "tts-kokoro-en", "tts-kokoro-zh", "asr-nemotron", "eou-parakeet"]; + const rootRegistry = modules.m0.ENGINES; + assert.deepEqual( + rootRegistry.map((entry) => entry.id), + expected, + ); + const registryModule = modules[`m${specifiers.indexOf("./registry")}`]; + assert.deepEqual( + registryModule.ENGINES.map((entry) => entry.id), + expected, + ); + for (const entry of rootRegistry) { + const engine = await entry.make(); + await engine.dispose(); + } + console.log(`SDK tarball: browser build and ${expected.length} real engine factories passed.`); +} finally { + rmSync(consumer, { recursive: true, force: true }); +} diff --git a/src/core/resource-session.ts b/src/core/resource-session.ts new file mode 100644 index 0000000..7603cc9 --- /dev/null +++ b/src/core/resource-session.ts @@ -0,0 +1,70 @@ +/** Owns one disposable resource and prevents overlapping operations on it. */ +export class ResourceSession }> { + private resource: T | undefined; + private pending: Promise | undefined; + private closed = false; + private closing: Promise | undefined; + + get current(): T | undefined { + return this.resource; + } + + get busy(): boolean { + return this.pending !== undefined || this.closed; + } + + async load(factory: () => Promise, initialize: (resource: T) => Promise): Promise { + return this.exclusive(async () => { + await this.release(); + const resource = await factory(); + try { + if (this.closed) throw new DOMException("Page closed", "AbortError"); + await initialize(resource); + if (this.closed) throw new DOMException("Page closed", "AbortError"); + this.resource = resource; + return resource; + } catch (error) { + await resource.dispose(); + throw error; + } + }); + } + + async run(operation: (resource: T) => Promise): Promise { + return this.exclusive(async () => { + if (this.resource === undefined) throw new Error("Load a model first"); + return operation(this.resource); + }); + } + + async clear(): Promise { + return this.exclusive(() => this.release()); + } + + /** Waits for in-flight work before disposing; no new work may start. */ + close(): Promise { + if (this.closing !== undefined) return this.closing; + this.closed = true; + this.closing = (async () => { + await this.pending?.catch(() => undefined); + await this.release(); + })(); + return this.closing; + } + + private async release(): Promise { + await this.resource?.dispose(); + this.resource = undefined; + } + + private async exclusive(operation: () => Promise): Promise { + if (this.busy) throw new Error("An operation is already in progress or the page is closed"); + const pending = Promise.resolve().then(operation); + this.pending = pending; + try { + return await pending; + } finally { + this.pending = undefined; + } + } +} diff --git a/src/engines/musicgen-acestep/cache-lease.ts b/src/engines/musicgen-acestep/cache-lease.ts new file mode 100644 index 0000000..5274c32 --- /dev/null +++ b/src/engines/musicgen-acestep/cache-lease.ts @@ -0,0 +1,18 @@ +/** Acquire a Web Lock before returning an idempotent, awaitable release function. */ +export function acquireCacheLease(locks: LockManager, name: string, signal?: AbortSignal): Promise<() => Promise> { + return new Promise((resolve, reject) => { + const released = locks.request(name, { mode: "shared", ...(signal === undefined ? {} : { signal }) }, async () => { + await new Promise((releaseLock) => { + let releasePromise: Promise | undefined; + resolve(async () => { + if (releasePromise === undefined) { + releaseLock(); + releasePromise = released.then(() => undefined); + } + await releasePromise; + }); + }); + }); + void released.catch(reject); + }); +} diff --git a/src/engines/musicgen-acestep/index.ts b/src/engines/musicgen-acestep/index.ts index 64f0e11..e4b4162 100644 --- a/src/engines/musicgen-acestep/index.ts +++ b/src/engines/musicgen-acestep/index.ts @@ -16,7 +16,7 @@ import { import { aceProductionWorkerConfiguration } from "./config.js"; import { aceInferenceWorkerName } from "./worker-name.js"; -import { ACE_MODEL_CACHE_LIFECYCLE_LOCK, ensureCurrentAceDemoModelCache } from "./model-cache-migration.js"; +import { acquireAceDemoModelCache, deleteAceDemoModelCache } from "./model-cache-migration.js"; import { INITIAL_MODEL_DOWNLOAD_PROGRESS, updateModelDownloadProgress, type ModelDownloadProgress } from "./model-download-progress.js"; export { @@ -24,11 +24,11 @@ export { ACE_MIN_DURATION_SECONDS, aceSeed, checkSupport, - deleteAceModelCache, inspectAceModelCache, releaseAceAudioOutput, requestAceModelStoragePersistence, } from "ace-step-1.5.wgsl"; +export { deleteAceDemoModelCache as deleteAceModelCache }; export type { AceGenerationRequest, AceGenerationResult, AceModelCacheInfo, AceSupportReport } from "ace-step-1.5.wgsl"; export { isModelDownloadComplete, @@ -55,18 +55,21 @@ interface ActiveOperation { initializationRequestId: number | undefined; jobId: number | undefined; request: AceGenerationRequest | undefined; + cancelRequested: boolean; } export class AceStepMusicClient { private worker: Worker | undefined; + private cacheAcquisition: AbortController | undefined; private workerReady = false; private nextRequestId = 1; private nextJobId = 1; private active: ActiveOperation | undefined; private disposal: { requestId: number; resolve: () => void; reject: (reason: unknown) => void } | undefined; + private disposePromise: Promise | undefined; private fatalGpuDiagnostic = false; /** Releases the shared model-cache lifecycle lock held while the worker is alive. */ - private releaseRuntimeLock: (() => void) | undefined; + private releaseRuntimeLock: (() => Promise) | undefined; private downloadProgress: ModelDownloadProgress = INITIAL_MODEL_DOWNLOAD_PROGRESS; /** Diagnostics reported by the worker's ready message, when initialized. */ runtimeDiagnostics: AceRuntimeDiagnostics | undefined; @@ -84,7 +87,7 @@ export class AceStepMusicClient { * packages) on first use. Only one generation may be in flight. */ async generate(request: AceGenerationRequest, handlers: AceMusicGenerateHandlers = {}): Promise { - if (this.active !== undefined) { + if (this.active !== undefined || this.disposePromise !== undefined) { throw new Error("A generation is already in progress"); } this.downloadProgress = INITIAL_MODEL_DOWNLOAD_PROGRESS; @@ -96,11 +99,12 @@ export class AceStepMusicClient { initializationRequestId: undefined, jobId: undefined, request, + cancelRequested: false, }; if (this.workerReady && this.worker !== undefined) { this.startGeneration(); } else { - this.startInitialization(); + void this.startInitialization(); } }); } @@ -108,8 +112,14 @@ export class AceStepMusicClient { /** Cancel the in-flight initialization or generation, if any. */ cancel(): void { const active = this.active; - if (this.worker === undefined || active === undefined) return; + if (active === undefined) return; + if (this.cacheAcquisition !== undefined) { + this.resetWorker(new DOMException("Initialization cancelled", "AbortError")); + return; + } + if (this.worker === undefined) return; if (active.initializationRequestId !== undefined) { + active.cancelRequested = true; this.worker.postMessage({ type: "cancel-initialization", requestId: active.initializationRequestId, @@ -120,7 +130,17 @@ export class AceStepMusicClient { } /** Release the worker's GPU/runtime resources and terminate it. */ - async dispose(): Promise { + dispose(): Promise { + if (this.disposePromise !== undefined) return this.disposePromise; + this.disposePromise = this.disposeInner().finally(() => { + this.disposePromise = undefined; + }); + return this.disposePromise; + } + + private async disposeInner(): Promise { + this.cacheAcquisition?.abort(); + this.cacheAcquisition = undefined; const active = this.active; if (active !== undefined) { this.active = undefined; @@ -133,16 +153,22 @@ export class AceStepMusicClient { return; } const requestId = this.nextRequestId++; - await new Promise((resolve, reject) => { - this.disposal = { requestId, resolve, reject }; - current.postMessage({ type: "dispose", requestId }); - }); - current.terminate(); - if (this.worker === current) this.worker = undefined; - this.workerReady = false; - this.runtimeDiagnostics = undefined; - this.releaseRuntimeLock?.(); - this.releaseRuntimeLock = undefined; + try { + await new Promise((resolve, reject) => { + this.disposal = { requestId, resolve, reject }; + current.postMessage({ type: "dispose", requestId }); + }); + } finally { + if (this.worker === current) { + current.terminate(); + this.worker = undefined; + this.workerReady = false; + this.runtimeDiagnostics = undefined; + const release = this.releaseRuntimeLock; + this.releaseRuntimeLock = undefined; + await release?.(); + } + } } /** Terminate immediately without an orderly runtime dispose. */ @@ -150,34 +176,27 @@ export class AceStepMusicClient { this.resetWorker(new DOMException("Client terminated", "AbortError")); } - private startInitialization(): void { - // Stamp/upgrade the OPFS cache generation exactly like the /music page does - // — a client that downloads 5.75 GB without the marker would see that cache - // wiped the next time the page's migration runs on this origin. - ensureCurrentAceDemoModelCache().then( - () => this.startInitializationInner(), - (error) => this.fail(new Error(`Could not prepare model storage: ${String(error)}`), false), - ); - } - - private acquireRuntimeLock(): void { - if (this.releaseRuntimeLock !== undefined || typeof navigator.locks?.request !== "function") return; - // Shared with every runtime tab; the migration's exclusive request (a - // generation bump) waits until all runtimes shut down before wiping. - void navigator.locks.request( - ACE_MODEL_CACHE_LIFECYCLE_LOCK, - { mode: "shared" }, - () => - new Promise((resolve) => { - this.releaseRuntimeLock = resolve; - }), - ); - } - - private startInitializationInner(): void { + private async startInitialization(): Promise { const active = this.active; if (active === undefined) return; - this.acquireRuntimeLock(); + const acquisition = new AbortController(); + this.cacheAcquisition = acquisition; + try { + const release = await acquireAceDemoModelCache(acquisition.signal); + if (this.active !== active || acquisition.signal.aborted) { + await release(); + return; + } + this.cacheAcquisition = undefined; + this.releaseRuntimeLock = release; + this.startInitializationInner(active); + } catch (error) { + if (this.active !== active || acquisition.signal.aborted) return; + this.fail(new Error(`Could not prepare model storage: ${String(error)}`), true); + } + } + + private startInitializationInner(active: ActiveOperation): void { this.worker?.terminate(); this.workerReady = false; this.fatalGpuDiagnostic = false; @@ -241,6 +260,11 @@ export class AceStepMusicClient { active.initializationRequestId = undefined; this.workerReady = true; this.runtimeDiagnostics = message.diagnostics; + if (active.cancelRequested) { + this.active = undefined; + active.reject(new DOMException("Initialization cancelled", "AbortError")); + return; + } this.startGeneration(); return; case "initialization-cancelled": @@ -287,7 +311,7 @@ export class AceStepMusicClient { return; } const fatal = this.fatalGpuDiagnostic || isAceFatalGpuErrorCode(message.error.code); - this.fail(new Error(`${message.error.code}: ${message.error.message}`), fatal); + this.fail(new Error(`${message.error.code}: ${message.error.message}`), fatal || !this.workerReady); return; } } @@ -301,17 +325,21 @@ export class AceStepMusicClient { const active = this.active; this.active = undefined; if (resetWorker) { + this.cacheAcquisition?.abort(); + this.cacheAcquisition = undefined; this.worker?.terminate(); this.worker = undefined; this.workerReady = false; this.runtimeDiagnostics = undefined; - this.releaseRuntimeLock?.(); + void this.releaseRuntimeLock?.(); this.releaseRuntimeLock = undefined; } active?.reject(reason); } private resetWorker(reason: unknown): void { + this.cacheAcquisition?.abort(); + this.cacheAcquisition = undefined; if (this.disposal !== undefined) { this.disposal.reject(reason instanceof Error ? reason : new Error(String(reason))); this.disposal = undefined; @@ -322,7 +350,7 @@ export class AceStepMusicClient { this.worker = undefined; this.workerReady = false; this.runtimeDiagnostics = undefined; - this.releaseRuntimeLock?.(); + void this.releaseRuntimeLock?.(); this.releaseRuntimeLock = undefined; if (active !== undefined) active.reject(reason); } diff --git a/src/engines/musicgen-acestep/model-cache-migration.ts b/src/engines/musicgen-acestep/model-cache-migration.ts index 43f71a8..4309dd3 100644 --- a/src/engines/musicgen-acestep/model-cache-migration.ts +++ b/src/engines/musicgen-acestep/model-cache-migration.ts @@ -1,4 +1,5 @@ import { deleteAceModelCache } from "ace-step-1.5.wgsl"; +import { acquireCacheLease } from "./cache-lease.js"; /** * Bump this identity whenever the demo's selected model payload inventory @@ -17,14 +18,16 @@ const MIGRATION_LOCK_NAME = ACE_MODEL_CACHE_LIFECYCLE_LOCK; const METADATA_DIRECTORY = "ace-step-1.5.wgsl-demo-metadata-v1"; const GENERATION_MARKER_FILE = "model-cache-generation.txt"; const MAX_GENERATION_MARKER_BYTES = 256; +const localReleasePromises = new Set>(); type ModelCacheStorage = Pick; export interface AceDemoExclusiveLockManager { - request(name: string, options: Readonly<{ mode: "exclusive" }>, operation: () => Promise): Promise; + request(name: string, options: Readonly<{ mode: "exclusive"; signal?: AbortSignal }>, operation: () => Promise): Promise; } export interface AceDemoModelCacheMigrationOptions { + readonly signal?: AbortSignal; /** @internal Test seam. Production defaults to `navigator.storage`. */ readonly storage?: ModelCacheStorage | undefined; /** @internal Test seam. Production defaults to `navigator.locks`. */ @@ -65,7 +68,7 @@ export async function ensureCurrentAceDemoModelCache(options: AceDemoModelCacheM } const deleteModelCache = options.deleteModelCache ?? deleteAceModelCache; - return locks.request(MIGRATION_LOCK_NAME, { mode: "exclusive" }, async () => { + return locks.request(MIGRATION_LOCK_NAME, { mode: "exclusive", ...(options.signal === undefined ? {} : { signal: options.signal }) }, async () => { // Read inside the exclusive lock. Another current-version tab may have // completed the migration while this tab was waiting for ownership. const root = await storage.getDirectory(); @@ -93,6 +96,58 @@ export async function ensureCurrentAceDemoModelCache(options: AceDemoModelCacheM }); } +/** Hold the current generation's cache for the entire runtime lifetime. */ +export async function acquireAceDemoModelCache(signal?: AbortSignal): Promise<() => Promise> { + const locks = globalThis.navigator?.locks; + const storage = browserStorage(); + if (locks === undefined || storage === undefined) throw new Error("Browser storage and Web Locks are required"); + for (;;) { + signal?.throwIfAborted(); + const release = await acquireCacheLease(locks, MIGRATION_LOCK_NAME, signal); + try { + signal?.throwIfAborted(); + const root = await storage.getDirectory(); + const metadata = await root.getDirectoryHandle(METADATA_DIRECTORY, { create: true }); + const generation = await readGenerationMarker(metadata); + signal?.throwIfAborted(); + if (generation === ACE_DEMO_MODEL_CACHE_GENERATION) return trackLocalRelease(release); + } catch (error) { + await release(); + throw error; + } + await release(); + await ensureCurrentAceDemoModelCache({ signal }); + // Recheck under the shared lock: another version could migrate in between. + } +} + +/** Fail promptly if another tab still owns the cache; never delete underneath it. */ +export async function deleteAceDemoModelCache(): Promise { + // A release requested in this page settles slightly before Web Locks removes + // the held lock. Do not mistake that brief local transition for another tab. + await Promise.allSettled([...localReleasePromises]); + const locks = globalThis.navigator?.locks; + if (locks === undefined) throw new Error("Web Locks are required to delete the model cache safely"); + return locks.request(MIGRATION_LOCK_NAME, { mode: "exclusive", ifAvailable: true }, async (lock) => { + if (lock === null) throw new Error("Model is in use in another tab. Close that tab and retry."); + return deleteAceModelCache(); + }); +} + +function trackLocalRelease(release: () => Promise): () => Promise { + let pending: Promise | undefined; + return () => { + if (pending !== undefined) return pending; + pending = release(); + localReleasePromises.add(pending); + void pending.then( + () => localReleasePromises.delete(pending!), + () => localReleasePromises.delete(pending!), + ); + return pending; + }; +} + async function readGenerationMarker(metadata: FileSystemDirectoryHandle): Promise { try { const handle = await metadata.getFileHandle(GENERATION_MARKER_FILE); diff --git a/src/engines/musicgen-acestep/pending-output-registry.ts b/src/engines/musicgen-acestep/pending-output-registry.ts new file mode 100644 index 0000000..6509609 --- /dev/null +++ b/src/engines/musicgen-acestep/pending-output-registry.ts @@ -0,0 +1,96 @@ +const PENDING_OUTPUTS_KEY = "ace-step-pending-output-ids"; +const PENDING_OUTPUTS_LOCK = "ace-step-pending-output-registry"; +const STALE_AFTER_MS = 60 * 60 * 1000; + +interface PendingOutput { + readonly id: string; + readonly at: number; +} + +type RegistryStorage = Pick; + +export interface PendingOutputRegistryOptions { + readonly storage?: RegistryStorage; + readonly locks?: LockManager; + readonly now?: () => number; +} + +export async function recordPendingOutput(id: string, options: PendingOutputRegistryOptions = {}): Promise { + await updateRecords((records) => [...records.filter((record) => record.id !== id), { id, at: now(options) }], options); +} + +export async function forgetPendingOutput(id: string, options: PendingOutputRegistryOptions = {}): Promise { + await updateRecords((records) => records.filter((record) => record.id !== id), options); +} + +export async function reclaimOrphanedOutputs( + currentId: string | undefined, + release: (id: string) => Promise, + options: PendingOutputRegistryOptions = {}, +): Promise { + const cutoff = now(options) - STALE_AFTER_MS; + const deleted = new Set(); + for (const record of readRecords(storage(options))) { + if (record.id === currentId || record.at > cutoff) continue; + try { + await release(record.id); + deleted.add(record.id); + } catch { + // Keep the record so a later visit retries the deletion. + } + } + if (deleted.size > 0) { + await updateRecords((records) => records.filter((record) => !deleted.has(record.id)), options); + } +} + +async function updateRecords(transform: (records: PendingOutput[]) => PendingOutput[], options: PendingOutputRegistryOptions): Promise { + const target = storage(options); + const update = () => writeRecords(target, transform(readRecords(target))); + const locks = options.locks ?? globalThis.navigator?.locks; + if (locks === undefined) { + update(); + return; + } + try { + await locks.request(PENDING_OUTPUTS_LOCK, { mode: "exclusive" }, update); + } catch { + // Registry bookkeeping is best-effort when Web Locks are unavailable. + update(); + } +} + +function storage(options: PendingOutputRegistryOptions): RegistryStorage | undefined { + try { + return options.storage ?? globalThis.localStorage; + } catch { + return undefined; + } +} + +function now(options: PendingOutputRegistryOptions): number { + return (options.now ?? Date.now)(); +} + +function readRecords(target: RegistryStorage | undefined): PendingOutput[] { + if (target === undefined) return []; + try { + const parsed = JSON.parse(target.getItem(PENDING_OUTPUTS_KEY) ?? "[]") as unknown; + if (!Array.isArray(parsed)) return []; + return parsed.filter( + (record): record is PendingOutput => + typeof record === "object" && record !== null && typeof record.id === "string" && typeof record.at === "number" && Number.isFinite(record.at), + ); + } catch { + return []; + } +} + +function writeRecords(target: RegistryStorage | undefined, records: PendingOutput[]): void { + if (target === undefined) return; + try { + target.setItem(PENDING_OUTPUTS_KEY, JSON.stringify(records.slice(-20))); + } catch { + // Storage unavailable; persistent outputs wait for another cleanup path. + } +} diff --git a/src/engines/musicgen-acestep/progress-breadcrumb.ts b/src/engines/musicgen-acestep/progress-breadcrumb.ts new file mode 100644 index 0000000..fb31e41 --- /dev/null +++ b/src/engines/musicgen-acestep/progress-breadcrumb.ts @@ -0,0 +1,10 @@ +export const CRASH_BREADCRUMB_KEY = "ace-step-progress-breadcrumb"; + +/** Terminal progress must not leave a crash marker for the next visit. */ +export function writeProgressBreadcrumb(title: string, detail: string, active: boolean, storage?: Pick): void { + try { + (storage ?? globalThis.localStorage).setItem(CRASH_BREADCRUMB_KEY, JSON.stringify({ title, detail, at: Date.now(), open: active })); + } catch { + // Diagnostics must not prevent generation when storage is unavailable. + } +} diff --git a/src/engines/registry.ts b/src/engines/registry.ts index 554be4a..0e3485a 100644 --- a/src/engines/registry.ts +++ b/src/engines/registry.ts @@ -1,32 +1,6 @@ -// Single source of truth for "what engines exist", consumed by the demo -// pages (src/pages/playground.ts, filtered per page by `category`). -// -// `make` is lazy: each engine's module (and its deps) loads only when selected, -// so a broken engine can't take down the whole app at page load. - -import type { Engine } from "../core/types.js"; - -export type EngineKind = "audio" | "text"; - -/** Which demo page an engine belongs to (pages filter the registry by this). */ -export type EngineCategory = "stt" | "tts" | "analysis"; - -export interface EngineEntry { - id: string; - /** Display label. */ - label: string; - kind: EngineKind; - category: EngineCategory; - /** Large weight downloads (hundreds of MB). */ - heavy?: boolean; - /** - * Optional availability probe. Engines whose weights may not be deployed - * (e.g. local-only exports) resolve false to be hidden from pickers instead - * of failing at load time. Absent = always available. - */ - available?: () => Promise; - make: () => Promise; -} +// Website catalog: shared SDK engines plus site-only runtimes. +import { ENGINES as SDK_ENGINES, type EngineEntry } from "./sdk-registry.js"; +export type { EngineEntry, EngineKind, EngineCategory } from "./sdk-registry.js"; /** * Dev convenience: when the vite middleware serves a local weight export at @@ -47,49 +21,7 @@ export async function localWeightDir(dir: string, probeFile: string): Promise new (await import("./vad-silero/index.js")).SileroVadEngine(), - }, - { - id: "asr-parakeet", - label: "Parakeet TDT v3", - kind: "audio", - category: "stt", - make: async () => new (await import("./asr-parakeet/index.js")).ParakeetV3Engine(), - }, - { - id: "asr-whisper", - label: "Whisper (99 langs)", - kind: "audio", - category: "stt", - make: async () => new (await import("./asr-whisper/index.js")).WhisperEngine(), - }, - { - id: "diarization-sortformer", - label: "Diarization (Sortformer)", - kind: "audio", - category: "analysis", - make: async () => new (await import("./diarization-sortformer/index.js")).SortformerDiarizationEngine(), - }, - { - id: "tts-kokoro-en", - label: "Kokoro TTS — English", - kind: "text", - category: "tts", - make: async () => new (await import("./tts-kokoro/index.js")).KokoroTtsEngine({ lang: "en" }), - }, - { - id: "tts-kokoro-zh", - label: "Kokoro TTS — Chinese", - kind: "text", - category: "tts", - heavy: true, - make: async () => new (await import("./tts-kokoro/index.js")).KokoroTtsEngine({ lang: "zh" }), - }, + ...SDK_ENGINES, { id: "stem-dicose", label: "DiCoSe Stem Splitter", @@ -114,22 +46,6 @@ export const ENGINES: EngineEntry[] = [ return new (await import("./stem-dicose/index.js")).DicoseStemEngine(baseUrl ? { baseUrl } : {}); }, }, - { - id: "asr-nemotron", - label: "Nemotron 3.5 (40 langs)", - kind: "audio", - category: "stt", - heavy: true, - make: async () => new (await import("./asr-nemotron/index.js")).NemotronEngine(), - }, - { - id: "eou-parakeet", - label: "Parakeet EOU 120M", - kind: "audio", - category: "stt", - heavy: true, - make: async () => new (await import("./eou-parakeet/index.js")).ParakeetEouEngine(), - }, { id: "asr-voicechat", label: "VoiceChat 11B STT", diff --git a/src/engines/sdk-registry.ts b/src/engines/sdk-registry.ts new file mode 100644 index 0000000..f579d8d --- /dev/null +++ b/src/engines/sdk-registry.ts @@ -0,0 +1,87 @@ +// Engines shipped in the npm SDK. The website extends this catalog. + +import type { Engine } from "../core/types.js"; + +export type EngineKind = "audio" | "text"; + +/** Which demo page an engine belongs to (pages filter the registry by this). */ +export type EngineCategory = "stt" | "tts" | "analysis"; + +export interface EngineEntry { + id: string; + /** Display label. */ + label: string; + kind: EngineKind; + category: EngineCategory; + /** Large weight downloads (hundreds of MB). */ + heavy?: boolean; + /** + * Optional availability probe. Engines whose weights may not be deployed + * (e.g. local-only exports) resolve false to be hidden from pickers instead + * of failing at load time. Absent = always available. + */ + available?: () => Promise; + make: () => Promise; +} + +export const ENGINES: EngineEntry[] = [ + { + id: "vad-silero", + label: "Silero VAD", + kind: "audio", + category: "analysis", + make: async () => new (await import("./vad-silero/index.js")).SileroVadEngine(), + }, + { + id: "asr-parakeet", + label: "Parakeet TDT v3", + kind: "audio", + category: "stt", + make: async () => new (await import("./asr-parakeet/index.js")).ParakeetV3Engine(), + }, + { + id: "asr-whisper", + label: "Whisper (99 langs)", + kind: "audio", + category: "stt", + make: async () => new (await import("./asr-whisper/index.js")).WhisperEngine(), + }, + { + id: "diarization-sortformer", + label: "Diarization (Sortformer)", + kind: "audio", + category: "analysis", + make: async () => new (await import("./diarization-sortformer/index.js")).SortformerDiarizationEngine(), + }, + { + id: "tts-kokoro-en", + label: "Kokoro TTS — English", + kind: "text", + category: "tts", + make: async () => new (await import("./tts-kokoro/index.js")).KokoroTtsEngine({ lang: "en" }), + }, + { + id: "tts-kokoro-zh", + label: "Kokoro TTS — Chinese", + kind: "text", + category: "tts", + heavy: true, + make: async () => new (await import("./tts-kokoro/index.js")).KokoroTtsEngine({ lang: "zh" }), + }, + { + id: "asr-nemotron", + label: "Nemotron 3.5 (40 langs)", + kind: "audio", + category: "stt", + heavy: true, + make: async () => new (await import("./asr-nemotron/index.js")).NemotronEngine(), + }, + { + id: "eou-parakeet", + label: "Parakeet EOU 120M", + kind: "audio", + category: "stt", + heavy: true, + make: async () => new (await import("./eou-parakeet/index.js")).ParakeetEouEngine(), + }, +]; diff --git a/src/engines/stem-dicose/index.ts b/src/engines/stem-dicose/index.ts index 872e92a..998fd47 100644 --- a/src/engines/stem-dicose/index.ts +++ b/src/engines/stem-dicose/index.ts @@ -26,6 +26,9 @@ export class DicoseStemEngine implements SeparationEngine { private readonly baseUrl: string; private readonly outputMode: DiCoSeOutputMode; private client: DiCoSeWorkerClient | null = null; + private loadPromise: Promise | null = null; + private disposePromise: Promise | null = null; + private disposed = false; private onClientProgress: ((p: DiCoSeProgress) => void) | null = null; constructor(opts: DicoseStemEngineOptions = {}) { @@ -34,8 +37,21 @@ export class DicoseStemEngine implements SeparationEngine { } async load(onProgress?: ProgressCb): Promise { - if (this.client) return; + if (this.disposed) throw new DOMException("DiCoSe engine has been disposed", "AbortError"); + if (this.loadPromise !== null) return await this.loadPromise; + if (this.client !== null) return; + const loadPromise = this.loadInner(onProgress); + this.loadPromise = loadPromise; + try { + await loadPromise; + } finally { + if (this.loadPromise === loadPromise) this.loadPromise = null; + } + } + + private async loadInner(onProgress?: ProgressCb): Promise { const support = await checkSupport(); + if (this.disposed) throw new DOMException("DiCoSe engine was disposed while loading", "AbortError"); if (!support.supported) { throw new Error(`DiCoSe needs WebGPU features this browser lacks: ${support.errors.join("; ")}`); } @@ -46,16 +62,25 @@ export class DicoseStemEngine implements SeparationEngine { createWorker: () => new Worker(new URL("./worker.ts", import.meta.url), { type: "module", name: "dicose-webgpu" }), onProgress: (p) => this.onClientProgress?.(p), }); + this.client = client; this.onClientProgress = (p) => onProgress?.(loadProgress(p)); try { await client.initialize(); + if (this.client !== client) { + await this.disposePromise?.catch(() => undefined); + throw new DOMException("DiCoSe engine was disposed while loading", "AbortError"); + } } catch (err) { - void client.dispose(); + if (this.client === client) { + this.client = null; + await client.dispose().catch(() => undefined); + } else { + await this.disposePromise?.catch(() => undefined); + } throw err; } finally { this.onClientProgress = null; } - this.client = client; } /** Full-band decode via the vendored decoder: stereo preserved, WAV kept at @@ -91,10 +116,16 @@ export class DicoseStemEngine implements SeparationEngine { return stems; } - async dispose(): Promise { + dispose(): Promise { + if (this.disposePromise !== null) return this.disposePromise; + this.disposed = true; const client = this.client; this.client = null; - if (client) await client.dispose(); + if (client === null) return Promise.resolve(); + this.disposePromise = client.dispose().finally(() => { + this.disposePromise = null; + }); + return this.disposePromise; } } diff --git a/src/engines/tts-kokoro/synth-backend.js b/src/engines/tts-kokoro/synth-backend.js index 008f0c7..985f308 100644 --- a/src/engines/tts-kokoro/synth-backend.js +++ b/src/engines/tts-kokoro/synth-backend.js @@ -25,82 +25,87 @@ export async function loadKokoroBackend( { modelDir = "kokoro", voiceRepo = "onnx-community/Kokoro-82M-v1.0-ONNX", onProgress } = {}, ) { const ctx = await createContext(); - const f32 = (u8) => new Float32Array(u8.buffer, u8.byteOffset, u8.byteLength / 4); - const json = async (path) => JSON.parse(new TextDecoder().decode(await fetchCached(hfUrl(WEIGHTS_REPO, path), onProgress, path))); - const bytes = (path) => fetchCached(hfUrl(WEIGHTS_REPO, path), onProgress, path); + try { + const f32 = (u8) => new Float32Array(u8.buffer, u8.byteOffset, u8.byteLength / 4); + const json = async (path) => JSON.parse(new TextDecoder().decode(await fetchCached(hfUrl(WEIGHTS_REPO, path), onProgress, path))); + const bytes = (path) => fetchCached(hfUrl(WEIGHTS_REPO, path), onProgress, path); - // predictor/decoder/generator weights - const K = makeKokoro(ctx, f32(await bytes(`${modelDir}/weights.bin`)), await json(`${modelDir}/manifest.json`), await json(`${modelDir}/roles.json`)); + // predictor/decoder/generator weights + const K = makeKokoro(ctx, f32(await bytes(`${modelDir}/weights.bin`)), await json(`${modelDir}/manifest.json`), await json(`${modelDir}/roles.json`)); - // ALBERT frontend (same structure as scripts/gpu-kokoro.mjs) - const aman = await json(`${modelDir}/albert/manifest.json`); - const cpu = {}; - for (const k of Object.keys(aman)) cpu[k] = f32(await bytes(`${modelDir}/albert/${k}.bin`)); - const scale = 1 / Math.sqrt(ALBERT_DIMS.HEAD_DIM); - for (let i = 0; i < cpu.q_w.length; i++) cpu.q_w[i] *= scale; - for (let i = 0; i < cpu.q_b.length; i++) cpu.q_b[i] *= scale; - const up2 = (n) => ctx.upload(cpu[n], aman[n][0], aman[n][1]); - const up1 = (n) => ctx.upload(cpu[n], 1, cpu[n].length); - const albertW = { - EMBED: ALBERT_DIMS.EMBED, - map_w: up2("map_w"), - map_b: up1("map_b"), - q_w: up2("q_w"), - q_b: up1("q_b"), - k_w: up2("k_w"), - k_b: up1("k_b"), - v_w: up2("v_w"), - v_b: up1("v_b"), - dense_w: up2("dense_w"), - dense_b: up1("dense_b"), - ffn_w: up2("ffn_w"), - ffn_b: up1("ffn_b"), - ffn_out_w: up2("ffn_out_w"), - ffn_out_b: up1("ffn_out_b"), - attn_ln_w: up1("attn_ln_w"), - attn_ln_b: up1("attn_ln_b"), - full_ln_w: up1("full_ln_w"), - full_ln_b: up1("full_ln_b"), - word_emb: cpu.word_emb, - pos_emb: cpu.pos_emb, - tok_emb: cpu.tok_emb, - emb_ln_w: cpu.emb_ln_w, - emb_ln_b: cpu.emb_ln_b, - }; - const fref = await json(`${modelDir}/ref.json`); - const beW = ctx.upload(f32(await bytes(`${modelDir}/be_w.bin`)), fref.be_in, fref.be_out); - const beB = ctx.upload(f32(await bytes(`${modelDir}/be_b.bin`)), 1, fref.be_out); + // ALBERT frontend (same structure as scripts/gpu-kokoro.mjs) + const aman = await json(`${modelDir}/albert/manifest.json`); + const cpu = {}; + for (const k of Object.keys(aman)) cpu[k] = f32(await bytes(`${modelDir}/albert/${k}.bin`)); + const scale = 1 / Math.sqrt(ALBERT_DIMS.HEAD_DIM); + for (let i = 0; i < cpu.q_w.length; i++) cpu.q_w[i] *= scale; + for (let i = 0; i < cpu.q_b.length; i++) cpu.q_b[i] *= scale; + const up2 = (n) => ctx.upload(cpu[n], aman[n][0], aman[n][1]); + const up1 = (n) => ctx.upload(cpu[n], 1, cpu[n].length); + const albertW = { + EMBED: ALBERT_DIMS.EMBED, + map_w: up2("map_w"), + map_b: up1("map_b"), + q_w: up2("q_w"), + q_b: up1("q_b"), + k_w: up2("k_w"), + k_b: up1("k_b"), + v_w: up2("v_w"), + v_b: up1("v_b"), + dense_w: up2("dense_w"), + dense_b: up1("dense_b"), + ffn_w: up2("ffn_w"), + ffn_b: up1("ffn_b"), + ffn_out_w: up2("ffn_out_w"), + ffn_out_b: up1("ffn_out_b"), + attn_ln_w: up1("attn_ln_w"), + attn_ln_b: up1("attn_ln_b"), + full_ln_w: up1("full_ln_w"), + full_ln_b: up1("full_ln_b"), + word_emb: cpu.word_emb, + pos_emb: cpu.pos_emb, + tok_emb: cpu.tok_emb, + emb_ln_w: cpu.emb_ln_w, + emb_ln_b: cpu.emb_ln_b, + }; + const fref = await json(`${modelDir}/ref.json`); + const beW = ctx.upload(f32(await bytes(`${modelDir}/be_w.bin`)), fref.be_in, fref.be_out); + const beB = ctx.upload(f32(await bytes(`${modelDir}/be_b.bin`)), 1, fref.be_out); - const voiceCache = new Map(); - const getVoice = async (name) => { - if (!voiceCache.has(name)) voiceCache.set(name, f32(await fetchCached(hfUrl(voiceRepo, `voices/${name}.bin`), onProgress, `${name}.bin`))); - return voiceCache.get(name); // [510*256] - }; + const voiceCache = new Map(); + const getVoice = async (name) => { + if (!voiceCache.has(name)) voiceCache.set(name, f32(await fetchCached(hfUrl(voiceRepo, `voices/${name}.bin`), onProgress, `${name}.bin`))); + return voiceCache.get(name); // [510*256] + }; - return { - ctx, // exposed so the engine can destroy the GPUDevice on dispose - backend: ctx.backend, - /** phonemes (IPA string) → 24 kHz Float32Array. */ - async synthFromPhonemes(phonemes, voice = "af_heart", speed = 1) { - const ids = [0]; // $ BOS - for (const ch of phonemes) { - const id = vocab[ch]; - if (id !== undefined) ids.push(id); - } - // ALBERT's positional table is 512 rows — cap at 510 phonemes + BOS/EOS - // (kokoro-js tokenized with truncation:true; without this, longer input - // reads past pos_emb → NaN audio). - if (ids.length > 511) { - console.warn(`[kokoro] input truncated: ${ids.length - 1} phonemes > 510`); - ids.length = 511; - } - ids.push(0); // $ EOS - const idArr = Int32Array.from(ids); - const pack = await getVoice(voice); - const si = 256 * Math.min(Math.max(idArr.length - 2, 0), 509); - const style = pack.slice(si, si + 256); - // textEncoding's output stays GPU-resident straight into the predictor. - return await synth(K, textEncoding(ctx, idArr, albertW, beW, beB), idArr, style, { speed }); - }, - }; + return { + ctx, // exposed so the engine can destroy the GPUDevice on dispose + backend: ctx.backend, + /** phonemes (IPA string) → 24 kHz Float32Array. */ + async synthFromPhonemes(phonemes, voice = "af_heart", speed = 1) { + const ids = [0]; // $ BOS + for (const ch of phonemes) { + const id = vocab[ch]; + if (id !== undefined) ids.push(id); + } + // ALBERT's positional table is 512 rows — cap at 510 phonemes + BOS/EOS + // (kokoro-js tokenized with truncation:true; without this, longer input + // reads past pos_emb → NaN audio). + if (ids.length > 511) { + console.warn(`[kokoro] input truncated: ${ids.length - 1} phonemes > 510`); + ids.length = 511; + } + ids.push(0); // $ EOS + const idArr = Int32Array.from(ids); + const pack = await getVoice(voice); + const si = 256 * Math.min(Math.max(idArr.length - 2, 0), 509); + const style = pack.slice(si, si + 256); + // textEncoding's output stays GPU-resident straight into the predictor. + return await synth(K, textEncoding(ctx, idArr, albertW, beW, beB), idArr, style, { speed }); + }, + }; + } catch (error) { + ctx.destroy(); + throw error; + } } diff --git a/src/index.ts b/src/index.ts index 45cedb6..2e5b64c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -5,8 +5,8 @@ // The root exports the shared types, the engine registry (id → lazy factory), // and the audio-decoding helper the demo pages use. -export { ENGINES } from "./engines/registry.js"; -export type { EngineEntry, EngineKind, EngineCategory } from "./engines/registry.js"; +export { ENGINES } from "./engines/sdk-registry.js"; +export type { EngineEntry, EngineKind, EngineCategory } from "./engines/sdk-registry.js"; export type { Engine, AsrEngine, diff --git a/src/music.ts b/src/music.ts index f09bcb8..5f45321 100644 --- a/src/music.ts +++ b/src/music.ts @@ -8,7 +8,6 @@ import { aceSeed, checkSupport, - deleteAceModelCache, inspectAceModelCache, isAceFatalGpuErrorCode, isAceWorkerMessage, @@ -26,6 +25,8 @@ import moonIcon from "./engines/musicgen-acestep/assets/moon.png"; import { aceProductionWorkerConfiguration } from "./engines/musicgen-acestep/config.js"; import { aceInferenceWorkerName } from "./engines/musicgen-acestep/worker-name.js"; +import { CRASH_BREADCRUMB_KEY, writeProgressBreadcrumb } from "./engines/musicgen-acestep/progress-breadcrumb.js"; +import { forgetPendingOutput, reclaimOrphanedOutputs, recordPendingOutput } from "./engines/musicgen-acestep/pending-output-registry.js"; import { formatDecimalBytes, formatModelDownloadAmount, @@ -36,7 +37,7 @@ import { updateModelDownloadProgress, type ModelDownloadProgress, } from "./engines/musicgen-acestep/model-download-progress.js"; -import { ACE_MODEL_CACHE_LIFECYCLE_LOCK, ensureCurrentAceDemoModelCache } from "./engines/musicgen-acestep/model-cache-migration.js"; +import { acquireAceDemoModelCache, deleteAceDemoModelCache } from "./engines/musicgen-acestep/model-cache-migration.js"; import { pcmToWav } from "./core/audio.js"; import { localWeightDir } from "./engines/registry.js"; import type { DicoseStemEngine } from "./engines/stem-dicose/index.js"; @@ -99,6 +100,11 @@ let pendingRequest: AceGenerationRequest | undefined; let nextRequestId = 1; let nextJobId = 1; let busy = false; +let deletingModel = false; +let cacheAcquisition: AbortController | undefined; +let generationPreparation: AbortController | undefined; +let initializationCancelRequested = false; +const pageLifecycle = new AbortController(); let supportDetails: AceSupportReport | undefined; let cacheDetails: AceModelCacheInfo | undefined; let workerDetails: unknown; @@ -119,7 +125,7 @@ let splittingStems = false; let tooltipRenderFrame: number | undefined; let pendingTooltipPoint: { readonly clientX: number } | undefined; /** Releases the shared model-cache lifecycle lock held while a worker is alive. */ -let releaseRuntimeLock: (() => void) | undefined; +let releaseRuntimeLock: (() => Promise) | undefined; let disposal: | { readonly requestId: number; @@ -131,18 +137,13 @@ let disposal: // Crash breadcrumb: iOS jetsam kills the tab with no error event, so persist // the last progress stage; after an unclean end the next visit reports where // the previous attempt died (the only telemetry a killed tab can leave). -const CRASH_BREADCRUMB_KEY = "ace-step-progress-breadcrumb"; function recordBreadcrumb(title: string, detail: string): void { // Mirror to the console so a tethered Web Inspector (iPhone debugging) // streams the stages — the last line before "Webpage Crashed" is the // memory-kill diagnosis. console.info(`[ace] ${title}${detail ? ` — ${detail}` : ""}`); - try { - localStorage.setItem(CRASH_BREADCRUMB_KEY, JSON.stringify({ title, detail, at: Date.now(), open: true })); - } catch { - // Storage unavailable — breadcrumbs are best-effort. - } + writeProgressBreadcrumb(title, detail, busy || splittingStems); } function closeBreadcrumb(): void { @@ -249,16 +250,25 @@ function wireEvents(): void { }); window.addEventListener("pagehide", () => { + pageLifecycle.abort(); + generationPreparation?.abort(); void stemEngine?.dispose(); stemEngine = undefined; if (output !== undefined) { - URL.revokeObjectURL(output.url); + const currentOutput = output; + URL.revokeObjectURL(currentOutput.url); // Best-effort only — the reliable path is the pending-output record // reclaimed on the next visit (releaseOrphanedOutputs). - void releaseAceAudioOutput(output.storageId).then(() => forgetPendingOutput(output?.storageId)); + void releaseAceAudioOutput(currentOutput.storageId) + .then(() => forgetPendingOutput(currentOutput.storageId)) + .catch(() => undefined); } - worker?.terminate(); - releaseRuntimeLock?.(); + resetWorker(); + setBusy(false); + closeBreadcrumb(); + }); + window.addEventListener("pageshow", (event) => { + if (event.persisted) location.reload(); }); } @@ -303,35 +313,6 @@ function hideProjectTooltip(): void { githubProjectTooltip.hidden = true; } -const PENDING_OUTPUTS_KEY = "ace-step-pending-output-ids"; - -function readPendingOutputs(): { id: string; at: number }[] { - try { - const raw = localStorage.getItem(PENDING_OUTPUTS_KEY); - const list = raw === null ? [] : (JSON.parse(raw) as { id: string; at: number }[]); - return Array.isArray(list) ? list : []; - } catch { - return []; - } -} - -function writePendingOutputs(list: { id: string; at: number }[]): void { - try { - localStorage.setItem(PENDING_OUTPUTS_KEY, JSON.stringify(list.slice(-20))); - } catch { - // Storage unavailable — the OPFS entries just wait for a later visit. - } -} - -function recordPendingOutput(id: string): void { - writePendingOutputs([...readPendingOutputs().filter((p) => p.id !== id), { id, at: Date.now() }]); -} - -function forgetPendingOutput(id: string | undefined): void { - if (id === undefined) return; - writePendingOutputs(readPendingOutputs().filter((p) => p.id !== id)); -} - /** * Committed WAVs are deliberately excluded from the runtime's own cleanup, so * a navigation that skipped releaseCurrentOutput() leaves up to ~92 MB per @@ -340,24 +321,16 @@ function forgetPendingOutput(id: string | undefined): void { * them, and are retried on a later visit once stale. */ async function releaseOrphanedOutputs(): Promise { - const cutoff = Date.now() - 60 * 60 * 1000; - const keep: { id: string; at: number }[] = []; - for (const entry of readPendingOutputs()) { - if (entry.at > cutoff && entry.id !== output?.storageId) { - keep.push(entry); - continue; - } - if (entry.id === output?.storageId) continue; - await releaseAceAudioOutput(entry.id).catch(() => undefined); - } - writePendingOutputs(keep); + await reclaimOrphanedOutputs(output?.storageId, releaseAceAudioOutput); } async function initializePage(): Promise { void releaseOrphanedOutputs(); try { - await ensureCurrentAceDemoModelCache(); + const release = await acquireAceDemoModelCache(pageLifecycle.signal); + await release(); } catch (error) { + if (pageLifecycle.signal.aborted) return; supportWarning.textContent = `Could not prepare model storage: ${errorMessage(error)} Reload to retry.`; supportWarning.className = "support-warning is-error"; supportWarning.hidden = false; @@ -394,7 +367,7 @@ async function initializePage(): Promise { } async function beginGeneration(): Promise { - if (busy || supportDetails?.supported !== true) return; + if (pageLifecycle.signal.aborted || busy || deletingModel || splittingStems || supportDetails?.supported !== true) return; let request: AceGenerationRequest; try { request = readGenerationRequest(); @@ -405,12 +378,24 @@ async function beginGeneration(): Promise { } formError.hidden = true; formError.textContent = ""; + const preparation = new AbortController(); + generationPreparation = preparation; + setBusy(true); try { await releaseCurrentOutput(); + if (pageLifecycle.signal.aborted) return; } catch (error) { + setBusy(false); formError.textContent = `Could not release the previous song: ${errorMessage(error)}`; formError.hidden = false; return; + } finally { + generationPreparation = undefined; + } + if (preparation.signal.aborted) { + setBusy(false); + setDeterminateProgress(0, "Cancelled", "Song generation cancelled", ""); + return; } pendingRequest = request; generationDetails = undefined; @@ -435,7 +420,7 @@ async function beginGeneration(): Promise { startPendingGeneration(); return; } - startWorkerInitialization(); + void startWorkerInitialization(); } function readGenerationRequest(): AceGenerationRequest { @@ -477,41 +462,39 @@ function readGenerationRequest(): AceGenerationRequest { }; } -function acquireRuntimeLock(): void { - if (releaseRuntimeLock !== undefined || typeof navigator.locks?.request !== "function") return; - // Shared mode: many tabs may run concurrently; the migration's exclusive - // request (a future generation bump in a new tab) waits until every tab's - // runtime has shut down instead of deleting the cache out from under one. - void navigator.locks.request( - ACE_MODEL_CACHE_LIFECYCLE_LOCK, - { mode: "shared" }, - () => - new Promise((resolve) => { - releaseRuntimeLock = resolve; - }), - ); -} - -function startWorkerInitialization(): void { - acquireRuntimeLock(); - worker?.terminate(); - workerReady = false; - worker = new Worker(new URL("./engines/musicgen-acestep/worker.ts", import.meta.url), { - type: "module", - name: aceInferenceWorkerName(), - }); - worker.addEventListener("message", onWorkerMessage); - worker.addEventListener("error", onWorkerError); - initializationRequestId = nextRequestId++; - setIndeterminateProgress("Preparing model", "Checking WebGPU and browser storage"); - worker.postMessage({ - type: "initialize", - requestId: initializationRequestId, - configuration: aceProductionWorkerConfiguration(), - modelSource: "cache-or-network", - reportProgress: true, - reportDiagnostics: true, - }); +async function startWorkerInitialization(): Promise { + resetWorker(); + const acquisition = new AbortController(); + cacheAcquisition = acquisition; + setIndeterminateProgress("Preparing model", "Waiting for model storage"); + try { + const release = await acquireAceDemoModelCache(acquisition.signal); + if (cacheAcquisition !== acquisition || acquisition.signal.aborted) { + await release(); + return; + } + cacheAcquisition = undefined; + releaseRuntimeLock = release; + worker = new Worker(new URL("./engines/musicgen-acestep/worker.ts", import.meta.url), { + type: "module", + name: aceInferenceWorkerName(), + }); + worker.addEventListener("message", onWorkerMessage); + worker.addEventListener("error", onWorkerError); + initializationRequestId = nextRequestId++; + setIndeterminateProgress("Preparing model", "Checking WebGPU and browser storage"); + worker.postMessage({ + type: "initialize", + requestId: initializationRequestId, + configuration: aceProductionWorkerConfiguration(), + modelSource: "cache-or-network", + reportProgress: true, + reportDiagnostics: true, + }); + } catch (error) { + if (acquisition.signal.aborted) return; + failOperation(`Could not initialize the model: ${errorMessage(error)}`, true); + } } function startPendingGeneration(): void { @@ -531,9 +514,25 @@ function startPendingGeneration(): void { } function cancelActiveOperation(): void { - if (worker === undefined || !busy) return; + if (!busy) return; + if (generationPreparation !== undefined) { + generationPreparation.abort(); + cancelButton.disabled = true; + setIndeterminateProgress("Cancelling", "Releasing the previous song"); + return; + } + if (cacheAcquisition !== undefined) { + resetWorker(); + pendingRequest = undefined; + setBusy(false); + setDeterminateProgress(0, "Cancelled", "Model preparation cancelled", ""); + return; + } + if (worker === undefined) return; cancelButton.disabled = true; if (initializationRequestId !== undefined) { + initializationCancelRequested = true; + pendingRequest = undefined; worker.postMessage({ type: "cancel-initialization", requestId: initializationRequestId, @@ -568,11 +567,18 @@ function onWorkerMessage(event: MessageEvent): void { workerDetails = message.diagnostics; updateRuntimeDetails(); void refreshCacheInfo(); + if (initializationCancelRequested) { + initializationCancelRequested = false; + setBusy(false); + setDeterminateProgress(modelProgress.fraction, "Cancelled", "Song generation cancelled", ""); + return; + } startPendingGeneration(); return; case "initialization-cancelled": if (message.requestId !== initializationRequestId) return; initializationRequestId = undefined; + initializationCancelRequested = false; pendingRequest = undefined; resetWorker(); setBusy(false); @@ -633,7 +639,7 @@ function onWorkerMessage(event: MessageEvent): void { return; } const fatal = fatalGpuDiagnostic || isAceFatalGpuErrorCode(message.error.code); - failOperation(`${message.error.code}: ${message.error.message}`, fatal); + failOperation(`${message.error.code}: ${message.error.message}`, fatal || !workerReady); return; } } @@ -650,7 +656,7 @@ async function publishResult(result: AceGenerationResult): Promise { output = { url, storageId: result.audioStorageId }; resultBlob = result.audio; resultSeed = result.seed; - recordPendingOutput(result.audioStorageId); + await recordPendingOutput(result.audioStorageId); audioPlayer.src = url; audioPlayer.load(); download.href = url; @@ -696,7 +702,7 @@ async function releaseCurrentOutput(): Promise { download.removeAttribute("href"); URL.revokeObjectURL(current.url); await releaseAceAudioOutput(current.storageId); - forgetPendingOutput(current.storageId); + await forgetPendingOutput(current.storageId); } // ── Split stems (DiCoSe, engines/stem-dicose) ──────────────────────────────── @@ -706,7 +712,7 @@ async function releaseCurrentOutput(): Promise { async function splitStems(): Promise { const blob = resultBlob; - if (splittingStems || busy || blob === undefined) return; + if (splittingStems || busy || deletingModel || blob === undefined) return; splittingStems = true; updateActionAvailability(); resetStemsUi(); @@ -746,6 +752,7 @@ async function splitStems(): Promise { setDeterminateProgress(progressElement.value, "Stem split failed", errorMessage(error), ""); } finally { splittingStems = false; + closeBreadcrumb(); updateActionAvailability(); } } @@ -802,6 +809,7 @@ function releaseStemUrls(): void { function failOperation(message: string, reset: boolean): void { closeBreadcrumb(); initializationRequestId = undefined; + initializationCancelRequested = false; activeJobId = undefined; pendingRequest = undefined; if (disposal !== undefined) { @@ -816,8 +824,8 @@ function failOperation(message: string, reset: boolean): void { } function resetWorker(): void { - releaseRuntimeLock?.(); - releaseRuntimeLock = undefined; + cacheAcquisition?.abort(); + cacheAcquisition = undefined; if (disposal !== undefined) { disposal.reject(new Error("worker reset while a dispose was pending")); disposal = undefined; @@ -825,7 +833,10 @@ function resetWorker(): void { worker?.terminate(); worker = undefined; workerReady = false; + initializationCancelRequested = false; workerDetails = undefined; + void releaseRuntimeLock?.(); + releaseRuntimeLock = undefined; } async function refreshCacheInfo(): Promise { @@ -850,18 +861,21 @@ async function refreshCacheInfo(): Promise { } async function deleteDownloadedModel(): Promise { - if (busy || !cacheCanBeDeleted()) return; + if (busy || splittingStems || deletingModel || !cacheCanBeDeleted()) return; + deletingModel = true; + updateActionAvailability(); deleteModelButton.disabled = true; cacheStatus.textContent = "Releasing the runtime…"; try { await disposeWorker(); cacheStatus.textContent = "Deleting downloaded model…"; - await deleteAceModelCache(); + await deleteAceDemoModelCache(); modelProgress = INITIAL_MODEL_DOWNLOAD_PROGRESS; await refreshCacheInfo(); } catch (error) { cacheStatus.textContent = `Could not delete the model: ${errorMessage(error)}`; } finally { + deletingModel = false; updateActionAvailability(); } } @@ -874,15 +888,21 @@ async function disposeWorker(): Promise { return; } const requestId = nextRequestId++; - await new Promise((resolve, reject) => { - disposal = { requestId, resolve, reject }; - current.postMessage({ type: "dispose", requestId }); - }); - current.terminate(); - if (worker === current) worker = undefined; - workerReady = false; - releaseRuntimeLock?.(); - releaseRuntimeLock = undefined; + try { + await new Promise((resolve, reject) => { + disposal = { requestId, resolve, reject }; + current.postMessage({ type: "dispose", requestId }); + }); + } finally { + if (worker === current) { + current.terminate(); + worker = undefined; + workerReady = false; + const release = releaseRuntimeLock; + releaseRuntimeLock = undefined; + await release?.(); + } + } } function renderModelProgress(): void { @@ -899,10 +919,10 @@ function setBusy(value: boolean): void { } function updateActionAvailability(): void { - generateButton.disabled = busy || splittingStems || supportDetails?.supported !== true; + generateButton.disabled = busy || deletingModel || splittingStems || supportDetails?.supported !== true; cancelButton.disabled = !busy; - deleteModelButton.disabled = busy || !cacheCanBeDeleted(); - splitStemsButton.disabled = busy || splittingStems || resultBlob === undefined; + deleteModelButton.disabled = busy || deletingModel || splittingStems || !cacheCanBeDeleted(); + splitStemsButton.disabled = busy || deletingModel || splittingStems || resultBlob === undefined; } function cacheCanBeDeleted(): boolean { diff --git a/src/pages/playground.ts b/src/pages/playground.ts index c748f48..15584ae 100644 --- a/src/pages/playground.ts +++ b/src/pages/playground.ts @@ -11,6 +11,7 @@ import { segmentsToSrt, segmentsToVtt } from "../core/captions.js"; import { formatLoadProgress } from "../core/loadProgress.js"; import { webgpuAvailable } from "../core/webgpu.js"; import { ENGINES, type EngineCategory, type EngineEntry } from "../engines/registry.js"; +import { ResourceSession } from "../core/resource-session.js"; import { MicCapture } from "../core/mic.js"; import type { Engine, LoadProgress, SeparationEngine, TranscribeProgress } from "../core/types.js"; @@ -86,6 +87,7 @@ export function initPlayground(opts: PlaygroundOptions) { } } const progress = $("progress"); + const loadBtn = $("load"); const runBtn = $("run"); const micBtn = opts.mic ? $("mic") : null; const player = $("player") as HTMLAudioElement | null; @@ -112,7 +114,20 @@ export function initPlayground(opts: PlaygroundOptions) { }), ); + const session = new ResourceSession(); let engine: Engine | null = null; + let loading = false; + let runningFile = false; + let startingLive = false; + let pageHidden = false; + let playbackUrl: string | undefined; + function updateActions() { + const occupied = loading || runningFile || startingLive || stopping || mic.running || pageHidden; + loadBtn.disabled = occupied; + engineSel.disabled = occupied; + runBtn.disabled = occupied || engine === null; + if (micBtn) micBtn.disabled = loading || runningFile || startingLive || stopping || pageHidden || engine === null; + } function currentEntry(): EngineEntry { return ENTRIES[engineSel.value]; } @@ -129,89 +144,109 @@ export function initPlayground(opts: PlaygroundOptions) { // TTS takes text → text box; ASR/VAD/diarization take audio → file picker. $("inputLabel").textContent = kind === "text" ? "Text to synthesize" : "Audio file"; } - engineSel.addEventListener("change", () => { + engineSel.addEventListener("change", async () => { syncInputs(); - runBtn.disabled = true; - if (micBtn) micBtn.disabled = true; - if (mic.running) void stopLive(); + engine = null; + loading = true; + updateActions(); + try { + await session.clear(); + } catch (error) { + output.textContent = String(error); + } finally { + loading = false; + updateActions(); + } }); syncInputs(); - $("load").addEventListener("click", async () => { + loadBtn.addEventListener("click", async () => { + if (loading || runningFile || startingLive || stopping || mic.running || pageHidden) return; const entry = currentEntry(); + loading = true; + engine = null; + updateActions(); output.textContent = ""; progress.hidden = false; - runBtn.disabled = true; try { - const eng = await entry.make(); - engine = eng; status.textContent = `Loading ${entry.label}…`; - await eng.load((p: LoadProgress) => { - progress.value = p.fraction || 0; - status.textContent = formatLoadProgress(p); - }); + engine = await session.load(entry.make, (eng) => + eng.load((p: LoadProgress) => { + progress.value = p.fraction || 0; + status.textContent = formatLoadProgress(p); + }), + ); status.textContent = `Ready: ${entry.label}`; - runBtn.disabled = false; - if (micBtn) micBtn.disabled = !(currentEntry().kind === "audio" && typeof (engine as any)?.transcribe === "function"); } catch (err) { - status.textContent = `Load failed`; + status.textContent = "Load failed"; output.textContent = String(err); } finally { + loading = false; progress.hidden = true; + updateActions(); } }); runBtn.addEventListener("click", async () => { - if (!engine) return; + if (!engine || loading || runningFile || startingLive || stopping || mic.running || pageHidden) return; + runningFile = true; + updateActions(); const entry = currentEntry(); output.textContent = "Running…"; if (player) player.hidden = true; try { - if (entry.kind === "text") { - const text = $("text").value; - const t0 = performance.now(); - const audio = await (engine as any).synthesize(text); - const ms = performance.now() - t0; - const dur = audio.samples.length / audio.sampleRate; - const blob = pcmToWav(audio.samples, audio.sampleRate); - const url = URL.createObjectURL(blob); - if (player) { - player.src = url; - player.hidden = false; - } - const wavLink = $("wavLink") as HTMLAnchorElement | null; - if (wavLink) { - wavLink.href = url; - wavLink.download = `${engineSel.value}.wav`; - wavLink.hidden = false; - } - output.textContent = - `Synthesized ${dur.toFixed(2)}s @ ${audio.sampleRate}Hz\n` + - `⏱ ${ms.toFixed(0)}ms · RTFx ${(dur / (ms / 1000)).toFixed(1)}× · ${(text.length / (ms / 1000)).toFixed(0)} chars/s`; - } else { - const file = $("file").files?.[0]; - if (!file) { - output.textContent = "Choose an audio file first."; - return; - } - // Stem splitters take full-band stereo at the clip's native rate and - // return one audio per stem — they bypass the 16 kHz mono decode and - // the text-result path entirely. - if (isSeparation(engine)) { - await runSeparationEngine(engine, await file.arrayBuffer(), file.name); - return; + await session.run(async (engine) => { + if (entry.kind === "text") { + const text = $("text").value; + const t0 = performance.now(); + const audio = await (engine as any).synthesize(text); + const ms = performance.now() - t0; + const dur = audio.samples.length / audio.sampleRate; + const blob = pcmToWav(audio.samples, audio.sampleRate); + if (playbackUrl !== undefined) URL.revokeObjectURL(playbackUrl); + const url = URL.createObjectURL(blob); + playbackUrl = url; + if (player) { + player.src = url; + player.hidden = false; + } + const wavLink = $("wavLink") as HTMLAnchorElement | null; + if (wavLink) { + wavLink.href = url; + wavLink.download = `${engineSel.value}.wav`; + wavLink.hidden = false; + } + output.textContent = + `Synthesized ${dur.toFixed(2)}s @ ${audio.sampleRate}Hz\n` + + `⏱ ${ms.toFixed(0)}ms · RTFx ${(dur / (ms / 1000)).toFixed(1)}× · ${(text.length / (ms / 1000)).toFixed(0)} chars/s`; + } else { + const file = $("file").files?.[0]; + if (!file) { + output.textContent = "Choose an audio file first."; + return; + } + // Stem splitters take full-band stereo at the clip's native rate and + // return one audio per stem — they bypass the 16 kHz mono decode and + // the text-result path entirely. + if (isSeparation(engine)) { + await runSeparationEngine(engine, await file.arrayBuffer(), file.name); + return; + } + const audio = await decodeToMono16k(await file.arrayBuffer()); + const dur = audio.samples.length / audio.sampleRate; + lastFileName = file.name; + const t0 = performance.now(); + const result = await runAudioEngine(engine, audio); + const ms = performance.now() - t0; + output.textContent = `⏱ ${ms.toFixed(0)}ms · audio ${dur.toFixed(1)}s · RTFx ${(dur / (ms / 1000)).toFixed(1)}×\n\n` + result; + renderCaptionLinks(); } - const audio = await decodeToMono16k(await file.arrayBuffer()); - const dur = audio.samples.length / audio.sampleRate; - lastFileName = file.name; - const t0 = performance.now(); - const result = await runAudioEngine(engine, audio); - const ms = performance.now() - t0; - output.textContent = `⏱ ${ms.toFixed(0)}ms · audio ${dur.toFixed(1)}s · RTFx ${(dur / (ms / 1000)).toFixed(1)}×\n\n` + result; - renderCaptionLinks(); - } + }); } catch (err) { output.textContent = String(err); + } finally { + runningFile = false; + updateActions(); } }); @@ -228,22 +263,24 @@ export function initPlayground(opts: PlaygroundOptions) { let livePos = 0; // absolute sample index consumed by the streaming path async function liveTick() { - if (!engine || liveBusy || mic.seconds < 1) return; + if (!engine || liveBusy || stopping || pageHidden || mic.seconds < 1) return; liveBusy = true; try { - const vu = "▁▂▃▄▅▆▇█"[Math.min(7, Math.floor(mic.level * 8))]; - if (isStreaming(engine)) { - const { samples, total } = mic.since(livePos); - const text = await engine.push(samples); - livePos = total; // only after push resolves — a failed push must not skip audio - mic.dropBefore(livePos); // streaming never re-reads history; keep hours-long sessions bounded - const ev = engine.streamEvents ?? []; - output.textContent = `● LIVE ${vu} ${mic.seconds.toFixed(0)}s (true streaming)\n\n${text}${ev.length ? `\n\nevents: ${ev.map((e) => `${e.type}@${e.time}s`).join(" ")}` : ""}`; - } else { - const samples = mic.tail(LIVE_WINDOW_SEC); - const r = await (engine as any).transcribe({ samples, sampleRate: 16000 }); - output.textContent = `● LIVE ${vu} ${mic.seconds.toFixed(0)}s (showing last ${Math.min(mic.seconds, LIVE_WINDOW_SEC).toFixed(0)}s)\n\n${r.text}${r.events?.length ? `\n\nevents: ${r.events.map((e: any) => `${e.type}@${e.time}s`).join(" ")}` : ""}`; - } + await session.run(async (engine) => { + const vu = "▁▂▃▄▅▆▇█"[Math.min(7, Math.floor(mic.level * 8))]; + if (isStreaming(engine)) { + const { samples, total } = mic.since(livePos); + const text = await engine.push(samples); + livePos = total; // only after push resolves — a failed push must not skip audio + mic.dropBefore(livePos); // streaming never re-reads history; keep hours-long sessions bounded + const ev = engine.streamEvents ?? []; + output.textContent = `● LIVE ${vu} ${mic.seconds.toFixed(0)}s (true streaming)\n\n${text}${ev.length ? `\n\nevents: ${ev.map((e) => `${e.type}@${e.time}s`).join(" ")}` : ""}`; + } else { + const samples = mic.tail(LIVE_WINDOW_SEC); + const r = await (engine as any).transcribe({ samples, sampleRate: 16000 }); + output.textContent = `● LIVE ${vu} ${mic.seconds.toFixed(0)}s (showing last ${Math.min(mic.seconds, LIVE_WINDOW_SEC).toFixed(0)}s)\n\n${r.text}${r.events?.length ? `\n\nevents: ${r.events.map((e: any) => `${e.type}@${e.time}s`).join(" ")}` : ""}`; + } + }); } catch (err) { output.textContent = `live error: ${String(err)}`; } finally { @@ -252,16 +289,26 @@ export function initPlayground(opts: PlaygroundOptions) { } async function startLive() { - if (stopping) return; // a previous session's flush is still settling + if (!engine || loading || runningFile || startingLive || stopping || mic.running || pageHidden) return; + startingLive = true; + updateActions(); output.textContent = "requesting microphone…"; try { mic.clear(); livePos = 0; if (engine && isStreaming(engine)) engine.reset(); await mic.start(); + if (pageHidden) { + await mic.stop(); + return; + } } catch (err) { + await mic.stop(); output.textContent = `microphone unavailable: ${String(err)}`; return; + } finally { + startingLive = false; + updateActions(); } if (micBtn) micBtn.textContent = "⏹ Stop"; runBtn.disabled = true; @@ -274,55 +321,88 @@ export function initPlayground(opts: PlaygroundOptions) { async function stopLive() { if (stopping) return; stopping = true; - if (liveTimer) clearInterval(liveTimer); - liveTimer = null; - // An in-flight tick may be mid-push: pushing/finishing/resetting concurrently - // would interleave on the same encoder caches. Let it settle first. - while (liveBusy) await new Promise((r) => setTimeout(r, 25)); - await mic.stop(); - if (micBtn) micBtn.textContent = "🎤 Live"; - // (Run stays disabled until the flush below completes.) - // Capture: the dropdown can reassign the global `engine` while we await — - // the flush must finish/reset the engine that owned this stream. - const eng = engine; - if (eng && isStreaming(eng)) { - // Streamed all along — just flush the right-padded tail. No re-decode. - try { - const { samples, total } = mic.since(livePos); - livePos = total; - if (samples.length) await eng.push(samples); - const text = await eng.finish(); - const ev = eng.streamEvents ?? []; - output.textContent = `■ final transcript (${mic.seconds.toFixed(0)}s, true streaming)\n\n${text}${ev.length ? `\n\nevents: ${ev.map((e) => `${e.type}@${e.time}s`).join(" ")}` : ""}`; - status.textContent = "Done."; - } catch (err) { - output.textContent = String(err); - } finally { - // reset even when the flush failed — a stranded stream blocks every - // subsequent Run with the stream-active guard. + updateActions(); + try { + if (liveTimer) clearInterval(liveTimer); + liveTimer = null; + // An in-flight tick may be mid-push: pushing/finishing/resetting concurrently + // would interleave on the same encoder caches. Let it settle first. + while (liveBusy) await new Promise((r) => setTimeout(r, 25)); + await mic.stop(); + if (micBtn) micBtn.textContent = "🎤 Live"; + // (Run stays disabled until the flush below completes.) + // Capture: the dropdown can reassign the global `engine` while we await — + // the flush must finish/reset the engine that owned this stream. + const eng = engine; + if (eng && isStreaming(eng)) { + // Streamed all along — just flush the right-padded tail. No re-decode. try { - eng.reset(); - } catch { - /* disposed mid-flight */ + const { samples, total } = mic.since(livePos); + livePos = total; + if (samples.length) await eng.push(samples); + const text = await eng.finish(); + const ev = eng.streamEvents ?? []; + output.textContent = `■ final transcript (${mic.seconds.toFixed(0)}s, true streaming)\n\n${text}${ev.length ? `\n\nevents: ${ev.map((e) => `${e.type}@${e.time}s`).join(" ")}` : ""}`; + status.textContent = "Done."; + } catch (err) { + output.textContent = String(err); + } finally { + // reset even when the flush failed — a stranded stream blocks every + // subsequent Run with the stream-active guard. + try { + eng.reset(); + } catch { + /* disposed mid-flight */ + } + } + } else if (engine && mic.seconds >= 1) { + // Final pass over the WHOLE capture (the rolling view only showed the tail). + status.textContent = `transcribing full ${mic.seconds.toFixed(0)}s capture…`; + try { + const t0 = performance.now(); + const r = await (engine as any).transcribe({ samples: mic.all(), sampleRate: 16000 }); + const ms = performance.now() - t0; + output.textContent = `■ final transcript (${mic.seconds.toFixed(0)}s captured, ${ms.toFixed(0)}ms, RTFx ${(mic.seconds / (ms / 1000)).toFixed(1)}×)\n\n${r.text}`; + status.textContent = "Done."; + } catch (err) { + output.textContent = String(err); } } - } else if (engine && mic.seconds >= 1) { - // Final pass over the WHOLE capture (the rolling view only showed the tail). - status.textContent = `transcribing full ${mic.seconds.toFixed(0)}s capture…`; - try { - const t0 = performance.now(); - const r = await (engine as any).transcribe({ samples: mic.all(), sampleRate: 16000 }); - const ms = performance.now() - t0; - output.textContent = `■ final transcript (${mic.seconds.toFixed(0)}s captured, ${ms.toFixed(0)}ms, RTFx ${(mic.seconds / (ms / 1000)).toFixed(1)}×)\n\n${r.text}`; - status.textContent = "Done."; - } catch (err) { - output.textContent = String(err); - } + } catch (err) { + output.textContent = String(err); + } finally { + stopping = false; + if (micBtn) micBtn.textContent = "🎤 Live"; + updateActions(); } - runBtn.disabled = false; // only after the flush — Run mid-flush hits the stream-active guard - stopping = false; } + window.addEventListener("pagehide", () => { + pageHidden = true; + if (liveTimer) clearInterval(liveTimer); + liveTimer = null; + void mic.stop().catch(console.error); + // Disposal waits for loading/inference to finish, including late load failures. + void (async () => { + while (stopping) await new Promise((resolve) => setTimeout(resolve, 25)); + try { + await session.close(); + } finally { + // In-flight work may have published new URLs after pagehide. + if (playbackUrl !== undefined) URL.revokeObjectURL(playbackUrl); + for (const url of stemUrls) URL.revokeObjectURL(url); + lastSegments = null; + renderCaptionLinks(); + } + })().catch(console.error); + if (playbackUrl !== undefined) URL.revokeObjectURL(playbackUrl); + for (const url of stemUrls) URL.revokeObjectURL(url); + updateActions(); + }); + window.addEventListener("pageshow", (event) => { + if (event.persisted) location.reload(); + }); + micBtn?.addEventListener("click", () => { if (mic.running) void stopLive(); else void startLive(); diff --git a/test/cache-lifecycle.test.ts b/test/cache-lifecycle.test.ts new file mode 100644 index 0000000..2f148c1 --- /dev/null +++ b/test/cache-lifecycle.test.ts @@ -0,0 +1,298 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { acquireCacheLease } from "../src/engines/musicgen-acestep/cache-lease.js"; +import { AceStepMusicClient, aceSeed, type AceGenerationRequest } from "../src/engines/musicgen-acestep/index.js"; +import { + acquireAceDemoModelCache, + ACE_DEMO_MODEL_CACHE_GENERATION, + ACE_MODEL_CACHE_LIFECYCLE_LOCK, + deleteAceDemoModelCache, +} from "../src/engines/musicgen-acestep/model-cache-migration.js"; + +// FIFO Web Locks test adapter; no model or audio data are needed. +class TestLocks { + active: { mode: string }[] = []; + queue: { mode: string; grant: () => void; reject: (error: unknown) => void; signal?: AbortSignal }[] = []; + request(name: string, options: LockOptions, callback: (lock: Lock | null) => unknown): Promise { + return new Promise((resolve, reject) => { + const mode = options.mode ?? "exclusive"; + if (options.signal?.aborted) { + reject(options.signal.reason); + return; + } + const request = { + mode, + reject, + signal: options.signal, + grant: () => { + this.active.push(request); + Promise.resolve() + .then(() => callback({ name, mode } as Lock)) + .then( + (value) => { + this.active.splice(this.active.indexOf(request), 1); + this.drain(); + resolve(value); + }, + (error) => { + this.active.splice(this.active.indexOf(request), 1); + this.drain(); + reject(error); + }, + ); + }, + }; + const blocked = () => this.active.some((lock) => lock.mode === "exclusive" || mode === "exclusive"); + if (options.ifAvailable && (blocked() || this.queue.length > 0)) { + Promise.resolve() + .then(() => callback(null)) + .then(resolve, reject); + return; + } + options.signal?.addEventListener( + "abort", + () => { + const index = this.queue.indexOf(request); + if (index < 0) return; + this.queue.splice(index, 1); + reject(options.signal?.reason); + this.drain(); + }, + { once: true }, + ); + this.queue.push(request); + this.drain(); + }); + } + drain() { + for (;;) { + const next = this.queue[0]; + if (!next || this.active.some((held) => held.mode === "exclusive" || next.mode === "exclusive")) return; + this.queue.shift(); + next.grant(); + } + } + asManager() { + return this as unknown as LockManager; + } +} + +function browserStorage(generation: string | null = ACE_DEMO_MODEL_CACHE_GENERATION) { + let marker = generation; + const removeEntry = vi.fn(async () => {}); + const metadata = { + async getFileHandle(_name: string, options?: { create?: boolean }) { + if (marker === null && !options?.create) throw new DOMException("Missing", "NotFoundError"); + return { + async getFile() { + return new Blob([marker ?? ""]); + }, + async createWritable() { + return { + async write(value: string) { + marker = value; + }, + async close() {}, + async abort() {}, + }; + }, + }; + }, + }; + return { + async getDirectory() { + return { + async getDirectoryHandle() { + return metadata; + }, + removeEntry, + }; + }, + removeEntry, + }; +} + +afterEach(() => vi.unstubAllGlobals()); + +describe("music cache ownership", () => { + const request: AceGenerationRequest = { + generationProfile: "ace-turbo-v1-correctness", + prompt: "Piano instrumental", + instrumental: true, + durationSeconds: 30, + seed: aceSeed("42"), + planner: { mode: "disabled" }, + }; + + it.each(["cancel", "dispose"] as const)("%s while waiting never starts a late worker", async (action) => { + const locks = new TestLocks(); + vi.stubGlobal("navigator", { locks: locks.asManager(), storage: browserStorage() }); + const worker = vi.fn(); + vi.stubGlobal("Worker", worker); + let finish!: () => void; + const exclusive = locks.request( + ACE_MODEL_CACHE_LIFECYCLE_LOCK, + { mode: "exclusive" }, + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + const client = new AceStepMusicClient(); + const rejected = expect(client.generate(request)).rejects.toMatchObject({ name: "AbortError" }); + await vi.waitFor(() => expect(locks.queue).toHaveLength(1)); + expect(worker).not.toHaveBeenCalled(); + await client[action](); + await rejected; + finish(); + await exclusive; + await vi.waitFor(() => expect(locks.active).toHaveLength(0)); + expect(worker).not.toHaveBeenCalled(); + expect(client.busy).toBe(false); + }); + + it("releases the worker and lease after initialization errors, including retries", async () => { + const locks = new TestLocks(); + vi.stubGlobal("navigator", { locks: locks.asManager(), storage: browserStorage() }); + vi.stubGlobal("location", { search: "" }); + // Worker transport only: no model implementation or generated audio. + const workers: WorkerTransport[] = []; + class WorkerTransport extends EventTarget { + terminate = vi.fn(); + constructor() { + super(); + workers.push(this); + } + postMessage(message: { type: string; requestId: number }) { + expect(message.type).toBe("initialize"); + expect(locks.active.map((lock) => lock.mode)).toEqual(["shared"]); + queueMicrotask(() => + this.dispatchEvent( + new MessageEvent("message", { + data: { + type: "error", + requestId: message.requestId, + error: { name: "Error", code: "NETWORK_ERROR", message: "Download interrupted" }, + }, + }), + ), + ); + } + } + vi.stubGlobal("Worker", WorkerTransport); + const client = new AceStepMusicClient(); + for (let attempt = 0; attempt < 2; attempt++) { + await expect(client.generate(request)).rejects.toThrow("Download interrupted"); + await vi.waitFor(() => expect(locks.active).toHaveLength(0)); + expect(workers[attempt].terminate).toHaveBeenCalledOnce(); + } + expect(client.initialized).toBe(false); + await client.dispose(); + }); + + it("does not grant a runtime lease until an exclusive operation finishes", async () => { + const locks = new TestLocks(); + let finish!: () => void; + const exclusive = locks.request( + "cache", + { mode: "exclusive" }, + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + const ready = vi.fn(); + const lease = acquireCacheLease(locks.asManager(), "cache").then((release) => { + ready(); + return release; + }); + await vi.waitFor(() => expect(locks.queue).toHaveLength(1)); + expect(ready).not.toHaveBeenCalled(); + finish(); + await exclusive; + const release = await lease; + expect(ready).toHaveBeenCalledOnce(); + await release(); + }); + + it("cancels a queued lease without acquiring an orphaned lock later", async () => { + const locks = new TestLocks(); + let finish!: () => void; + const exclusive = locks.request( + "cache", + { mode: "exclusive" }, + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + const controller = new AbortController(); + const lease = acquireCacheLease(locks.asManager(), "cache", controller.signal); + const rejected = expect(lease).rejects.toMatchObject({ name: "AbortError" }); + await vi.waitFor(() => expect(locks.queue).toHaveLength(1)); + controller.abort(); + await rejected; + finish(); + await exclusive; + await vi.waitFor(() => expect(locks.active).toHaveLength(0)); + expect(locks.queue).toHaveLength(0); + }); + + it("allows same-version tabs to share the cache and refuses deletion until both release", async () => { + const locks = new TestLocks(); + const storage = browserStorage(); + vi.stubGlobal("navigator", { locks: locks.asManager(), storage }); + const release1 = await acquireAceDemoModelCache(); + const release2 = await acquireAceDemoModelCache(); + expect(locks.active).toHaveLength(2); + await expect(deleteAceDemoModelCache()).rejects.toThrow("in use in another tab"); + expect(storage.removeEntry).not.toHaveBeenCalled(); + await release1(); + await expect(deleteAceDemoModelCache()).rejects.toThrow("in use in another tab"); + await release2(); + await expect(deleteAceDemoModelCache()).resolves.toBe(true); + expect(storage.removeEntry).toHaveBeenCalledOnce(); + }); + + it("waits for this page's requested release before trying an immediate deletion", async () => { + const locks = new TestLocks(); + const storage = browserStorage(); + vi.stubGlobal("navigator", { locks: locks.asManager(), storage }); + const release = await acquireAceDemoModelCache(); + void release(); + await expect(deleteAceDemoModelCache()).resolves.toBe(true); + expect(storage.removeEntry).toHaveBeenCalledOnce(); + }); + + it("migrates a stale cache before publishing a runtime lease", async () => { + const locks = new TestLocks(); + const storage = browserStorage("old-generation"); + vi.stubGlobal("navigator", { locks: locks.asManager(), storage }); + const release = await acquireAceDemoModelCache(); + expect(storage.removeEntry).toHaveBeenCalledOnce(); + expect(locks.active.map((held) => held.mode)).toEqual(["shared"]); + await release(); + }); + + it("cancels preparation while a migration is blocking it", async () => { + const locks = new TestLocks(); + const storage = browserStorage(); + vi.stubGlobal("navigator", { locks: locks.asManager(), storage }); + let finish!: () => void; + const exclusive = locks.request( + ACE_MODEL_CACHE_LIFECYCLE_LOCK, + { mode: "exclusive" }, + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + const controller = new AbortController(); + const rejected = expect(acquireAceDemoModelCache(controller.signal)).rejects.toMatchObject({ name: "AbortError" }); + await vi.waitFor(() => expect(locks.queue).toHaveLength(1)); + controller.abort(); + await rejected; + finish(); + await exclusive; + expect(storage.removeEntry).not.toHaveBeenCalled(); + }); +}); diff --git a/test/deploy-workflow.test.ts b/test/deploy-workflow.test.ts new file mode 100644 index 0000000..c56eef0 --- /dev/null +++ b/test/deploy-workflow.test.ts @@ -0,0 +1,19 @@ +import { readFileSync } from "node:fs"; +import { expect, it } from "vitest"; +import { parse } from "yaml"; + +it("gates Pages on every check for the same main revision", () => { + const ci = parse(readFileSync(new URL("../.github/workflows/ci.yml", import.meta.url), "utf8")); + const deploy = parse(readFileSync(new URL("../.github/workflows/deploy.yml", import.meta.url), "utf8")); + expect(ci.jobs.deploy.needs.toSorted()).toEqual( + Object.keys(ci.jobs) + .filter((name) => name !== "deploy") + .toSorted(), + ); + expect(ci.jobs.deploy.uses).toBe("./.github/workflows/deploy.yml"); + expect(ci.jobs.deploy.if).toBe("github.ref == 'refs/heads/main' && (github.event_name == 'push' || github.event_name == 'workflow_dispatch')"); + expect(Object.keys(deploy.on)).toEqual(["workflow_call"]); + expect(ci.jobs.build.steps.map((step: { run?: string }) => step.run)).toEqual(expect.arrayContaining(["npm run test:integration", "npm run sdk:test"])); + // Checkout must use the caller's revision, not fetch a newer main after tests. + expect(deploy.jobs.build.steps.find((step: { uses?: string }) => step.uses?.startsWith("actions/checkout@")).with?.ref).toBeUndefined(); +}); diff --git a/test/dicose-lifecycle.test.ts b/test/dicose-lifecycle.test.ts new file mode 100644 index 0000000..311d7af --- /dev/null +++ b/test/dicose-lifecycle.test.ts @@ -0,0 +1,72 @@ +import { afterEach, expect, it, vi } from "vitest"; +import { DicoseStemEngine } from "../src/engines/stem-dicose/index.js"; + +interface AdapterGate { + readonly promise: Promise; + readonly resolve: (adapter: GPUAdapter) => void; +} + +function adapter(): GPUAdapter { + return { + features: new Set(["shader-f16", "subgroups"]), + info: { vendor: "test", architecture: "test", device: "test", description: "test", subgroupMinSize: 32, subgroupMaxSize: 32 }, + limits: { maxBufferSize: 1_073_741_824, maxStorageBufferBindingSize: 1_073_741_824, maxComputeWorkgroupStorageSize: 25_344 }, + } as unknown as GPUAdapter; +} + +function deferredAdapter(): AdapterGate { + let resolve!: (adapter: GPUAdapter) => void; + return { + promise: new Promise((done) => { + resolve = done; + }), + resolve, + }; +} + +afterEach(() => vi.unstubAllGlobals()); + +it("does not create a stem worker when disposed during the capability check", async () => { + const gate = deferredAdapter(); + vi.stubGlobal("isSecureContext", true); + vi.stubGlobal("navigator", { gpu: { requestAdapter: () => gate.promise } }); + const worker = vi.fn(); + vi.stubGlobal("Worker", worker); + const engine = new DicoseStemEngine(); + const load = engine.load(); + await engine.dispose(); + gate.resolve(adapter()); + await expect(load).rejects.toMatchObject({ name: "AbortError" }); + expect(worker).not.toHaveBeenCalled(); +}); + +it("disposes a stem worker whose initialization finishes after teardown", async () => { + vi.stubGlobal("isSecureContext", true); + vi.stubGlobal("navigator", { gpu: { requestAdapter: async () => adapter() } }); + vi.stubGlobal("location", { href: "http://localhost/" }); + let transport!: WorkerTransport; + class WorkerTransport extends EventTarget { + readonly messages: { type: string; id: number }[] = []; + readonly terminate = vi.fn(); + constructor() { + super(); + transport = this; + } + postMessage(message: { type: string; id: number }) { + this.messages.push(message); + } + } + vi.stubGlobal("Worker", WorkerTransport); + const engine = new DicoseStemEngine(); + const load = engine.load(); + await vi.waitFor(() => expect(transport.messages.map((message) => message.type)).toEqual(["initialize"])); + const disposal = engine.dispose(); + expect(transport.messages.map((message) => message.type)).toEqual(["initialize", "dispose"]); + const initializeId = transport.messages[0].id; + transport.dispatchEvent(new MessageEvent("message", { data: { type: "initialized", id: initializeId } })); + const disposeId = transport.messages[1].id; + transport.dispatchEvent(new MessageEvent("message", { data: { type: "disposed", id: disposeId } })); + await expect(load).rejects.toMatchObject({ name: "AbortError" }); + await disposal; + expect(transport.terminate).toHaveBeenCalledOnce(); +}); diff --git a/test/kokoro-cleanup.test.ts b/test/kokoro-cleanup.test.ts new file mode 100644 index 0000000..fae09f6 --- /dev/null +++ b/test/kokoro-cleanup.test.ts @@ -0,0 +1,20 @@ +import { afterEach, expect, it, vi } from "vitest"; +import { loadKokoroBackend } from "../src/engines/tts-kokoro/synth-backend.js"; +import { WasmContext } from "../src/gpu/wasm-context.js"; + +afterEach(() => vi.restoreAllMocks()); + +it("destroys each real WASM context after a download failure, including retries", async () => { + // Real compute context; the download fails before any model data is returned. + const destroy = vi.spyOn(WasmContext.prototype, "destroy"); + const failure = new Error("Download interrupted"); + const download = vi.fn(async () => { + throw failure; + }); + for (let attempt = 1; attempt <= 2; attempt++) { + await expect(loadKokoroBackend(download, (_repo, path) => path, {})).rejects.toBe(failure); + expect(destroy).toHaveBeenCalledTimes(attempt); + } + expect(destroy.mock.instances[0]).toBeInstanceOf(WasmContext); + expect(destroy.mock.instances[0]).not.toBe(destroy.mock.instances[1]); +}); diff --git a/test/music-disposal.test.ts b/test/music-disposal.test.ts new file mode 100644 index 0000000..089b724 --- /dev/null +++ b/test/music-disposal.test.ts @@ -0,0 +1,101 @@ +import { readFileSync } from "node:fs"; +import { expect, it, vi } from "vitest"; +import { AceStepMusicClient, aceSeed } from "../src/engines/musicgen-acestep/index.js"; + +const request = { + generationProfile: "ace-turbo-v1-correctness" as const, + prompt: "Piano instrumental", + instrumental: true, + durationSeconds: 30, + seed: aceSeed("42"), + planner: { mode: "disabled" as const }, +}; + +function readyClient() { + const client = new AceStepMusicClient(); + // Set only lifecycle state; this transport never loads or simulates a model. + const worker = { postMessage: vi.fn(), terminate: vi.fn() }; + const release = vi.fn(async () => {}); + Object.assign(client, { worker, workerReady: true, releaseRuntimeLock: release }); + const receive = (data: unknown) => Reflect.get(client, "onMessage").call(client, { data }); + return { client, worker, release, receive }; +} + +it("shares disposal until both callers settle and prevents new generation meanwhile", async () => { + const { client, worker, release, receive } = readyClient(); + const first = client.dispose(); + const second = client.dispose(); + expect(second).toBe(first); + expect(worker.postMessage).toHaveBeenCalledExactlyOnceWith({ type: "dispose", requestId: 1 }); + await expect(client.generate(request)).rejects.toThrow("already in progress"); + receive({ type: "disposed", requestId: 1 }); + await Promise.all([first, second]); + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(client.initialized).toBe(false); + await client.dispose(); + expect(worker.terminate).toHaveBeenCalledOnce(); +}); + +it("does not generate when ready was already queued before initialization cancellation", async () => { + const { client, worker, receive } = readyClient(); + let reject!: (error: unknown) => void; + const generation = new Promise((_, rejectPromise) => { + reject = rejectPromise; + }); + Object.assign(client, { + workerReady: false, + active: { + resolve: vi.fn(), + reject, + handlers: {}, + initializationRequestId: 7, + jobId: undefined, + request, + cancelRequested: false, + }, + }); + const rejected = expect(generation).rejects.toMatchObject({ name: "AbortError" }); + client.cancel(); + const receipt = JSON.parse(readFileSync(new URL("../packages/acestep/optimization/results/OPT-0073/final-180s-receipt.json", import.meta.url), "utf8")) as { + result: { diagnostics: unknown }; + }; + receive({ type: "ready", requestId: 7, diagnostics: receipt.result.diagnostics }); + await rejected; + expect(worker.postMessage).toHaveBeenCalledExactlyOnceWith({ type: "cancel-initialization", requestId: 7 }); + expect(client.busy).toBe(false); + expect(client.initialized).toBe(true); + client.terminate(); +}); + +it("releases the worker and lease even when orderly disposal fails", async () => { + const { client, worker, release, receive } = readyClient(); + const first = client.dispose(); + const second = client.dispose(); + const settled = Promise.allSettled([first, second]); + receive({ type: "error", requestId: 1, error: { name: "Error", code: "DISPOSE_FAILED", message: "Disposal failed" } }); + expect((await settled).map((result) => result.status)).toEqual(["rejected", "rejected"]); + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(client.initialized).toBe(false); + await client.dispose(); +}); + +it("cleans up when posting the disposal request throws", async () => { + const { client, worker, release } = readyClient(); + worker.postMessage.mockImplementation(() => { + throw new Error("Worker unavailable"); + }); + await expect(client.dispose()).rejects.toThrow("Worker unavailable"); + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); +}); + +it("settles all disposal callers when terminated mid-disposal", async () => { + const { client, worker, release } = readyClient(); + const settled = Promise.allSettled([client.dispose(), client.dispose()]); + client.terminate(); + expect((await settled).map((result) => result.status)).toEqual(["rejected", "rejected"]); + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); +}); diff --git a/test/music-page-disposal.test.ts b/test/music-page-disposal.test.ts new file mode 100644 index 0000000..d72ac7f --- /dev/null +++ b/test/music-page-disposal.test.ts @@ -0,0 +1,32 @@ +import { readFileSync } from "node:fs"; +import vm from "node:vm"; +import ts from "typescript"; +import { isAceWorkerMessage } from "ace-step-1.5.wgsl"; +import { expect, it, vi } from "vitest"; + +it("terminates the page worker and awaits its cache lease after disposal errors", async () => { + const source = readFileSync(new URL("../src/music.ts", import.meta.url), "utf8"); + const ast = ts.createSourceFile("music.ts", source, ts.ScriptTarget.Latest, true); + const handlers = ast.statements.filter((node) => ts.isFunctionDeclaration(node) && ["disposeWorker", "onWorkerMessage"].includes(node.name?.text ?? "")); + const release = vi.fn(async () => {}); + const worker = { postMessage: vi.fn(), terminate: vi.fn() }; + const state: Record = { + worker, + workerReady: true, + nextRequestId: 1, + disposal: undefined, + releaseRuntimeLock: release, + isAceWorkerMessage, + }; + vm.createContext(state); + vm.runInContext(ts.transpile(handlers.map((node) => node.getText(ast)).join("\n"), { target: ts.ScriptTarget.ES2022 }), state); + const disposed = state.disposeWorker(); + state.onWorkerMessage({ + data: { type: "error", requestId: 1, error: { name: "Error", code: "DISPOSE_FAILED", message: "Device cleanup failed" } }, + }); + await expect(disposed).rejects.toThrow("Device cleanup failed"); + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(state.worker).toBeUndefined(); + expect(state.workerReady).toBe(false); +}); diff --git a/test/music-preparation.test.ts b/test/music-preparation.test.ts new file mode 100644 index 0000000..f711d4c --- /dev/null +++ b/test/music-preparation.test.ts @@ -0,0 +1,140 @@ +import { readFileSync } from "node:fs"; +import vm from "node:vm"; +import ts from "typescript"; +import { isAceWorkerMessage } from "ace-step-1.5.wgsl"; +import { expect, it, vi } from "vitest"; + +// Run the actual page handlers with UI/worker transports supplied by the test. +// No models or audio are substituted; execution stops before inference starts. +function page() { + const source = readFileSync(new URL("../src/music.ts", import.meta.url), "utf8"); + const ast = ts.createSourceFile("music.ts", source, ts.ScriptTarget.Latest, true); + const names = ["beginGeneration", "cancelActiveOperation"]; + const handlers = ast.statements.filter((node) => ts.isFunctionDeclaration(node) && names.includes(node.name?.text ?? "")); + expect(handlers).toHaveLength(names.length); + let finish!: () => void; + let fail!: (error: Error) => void; + const cleanup = new Promise((resolve, reject) => { + finish = resolve; + fail = reject; + }); + const state: Record = { + AbortController, + pageLifecycle: new AbortController(), + generationPreparation: undefined, + busy: false, + deletingModel: false, + splittingStems: false, + supportDetails: { supported: true }, + readGenerationRequest: () => ({ prompt: "Piano instrumental" }), + formError: {}, + resultPanel: {}, + cancelButton: {}, + releaseCurrentOutput: vi.fn(() => cleanup), + errorMessage: String, + cacheAcquisition: undefined, + initializationRequestId: undefined, + activeJobId: undefined, + INITIAL_MODEL_DOWNLOAD_PROGRESS: {}, + isModelDownloadComplete: () => true, + cacheDetails: undefined, + requestAceModelStoragePersistence: async () => false, + recordValue: () => ({}), + updateRuntimeDetails: () => {}, + workerDetails: undefined, + workerReady: true, + worker: { postMessage: vi.fn() }, + startPendingGeneration: vi.fn(), + setDeterminateProgress: vi.fn(), + setIndeterminateProgress: vi.fn(), + }; + state.setBusy = (value: boolean) => { + state.busy = value; + }; + vm.createContext(state); + vm.runInContext(ts.transpile(handlers.map((node) => node.getText(ast)).join("\n"), { target: ts.ScriptTarget.ES2022 }), state); + return { state, finish, fail }; +} + +it("honors Cancel during output cleanup and blocks another generation until cleanup ends", async () => { + const { state, finish } = page(); + const running = state.beginGeneration(); + state.cancelActiveOperation(); + expect(state.busy).toBe(true); + await state.beginGeneration(); + expect(state.releaseCurrentOutput).toHaveBeenCalledOnce(); + finish(); + await running; + expect(state.startPendingGeneration).not.toHaveBeenCalled(); + expect(state.worker.postMessage).not.toHaveBeenCalled(); + expect(state.busy).toBe(false); + expect(state.generationPreparation).toBeUndefined(); + expect(state.setDeterminateProgress).toHaveBeenCalledWith(0, "Cancelled", "Song generation cancelled", ""); + await state.beginGeneration(); + expect(state.startPendingGeneration).toHaveBeenCalledOnce(); +}); + +it("does not start generation if cleanup fails after cancellation", async () => { + const { state, fail } = page(); + const running = state.beginGeneration(); + state.cancelActiveOperation(); + fail(new Error("Storage unavailable")); + await running; + expect(state.startPendingGeneration).not.toHaveBeenCalled(); + expect(state.busy).toBe(false); + expect(state.generationPreparation).toBeUndefined(); + expect(state.formError.textContent).toContain("Storage unavailable"); +}); + +it("does not start generation when the page closes during cleanup", async () => { + const { state, finish } = page(); + const running = state.beginGeneration(); + state.pageLifecycle.abort(); + finish(); + await running; + expect(state.startPendingGeneration).not.toHaveBeenCalled(); + expect(state.generationPreparation).toBeUndefined(); +}); + +it("does not generate when ready was already queued before page cancellation", async () => { + const source = readFileSync(new URL("../src/music.ts", import.meta.url), "utf8"); + const ast = ts.createSourceFile("music.ts", source, ts.ScriptTarget.Latest, true); + const handlers = ast.statements.filter( + (node) => ts.isFunctionDeclaration(node) && ["cancelActiveOperation", "onWorkerMessage"].includes(node.name?.text ?? ""), + ); + const worker = { postMessage: vi.fn() }; + const state: Record = { + busy: true, + worker, + cacheAcquisition: undefined, + generationPreparation: undefined, + initializationRequestId: 7, + initializationCancelRequested: false, + activeJobId: undefined, + pendingRequest: { prompt: "Piano instrumental" }, + cancelButton: {}, + workerReady: false, + workerDetails: undefined, + modelProgress: { fraction: 0.5 }, + isAceWorkerMessage, + updateRuntimeDetails: vi.fn(), + refreshCacheInfo: vi.fn(), + setBusy: vi.fn((value: boolean) => { + state.busy = value; + }), + setDeterminateProgress: vi.fn(), + startPendingGeneration: vi.fn(), + }; + vm.createContext(state); + vm.runInContext(ts.transpile(handlers.map((node) => node.getText(ast)).join("\n"), { target: ts.ScriptTarget.ES2022 }), state); + state.cancelActiveOperation(); + const receipt = JSON.parse(readFileSync(new URL("../packages/acestep/optimization/results/OPT-0073/final-180s-receipt.json", import.meta.url), "utf8")) as { + result: { diagnostics: unknown }; + }; + state.onWorkerMessage({ data: { type: "ready", requestId: 7, diagnostics: receipt.result.diagnostics } }); + expect(worker.postMessage).toHaveBeenCalledExactlyOnceWith({ type: "cancel-initialization", requestId: 7 }); + expect(state.startPendingGeneration).not.toHaveBeenCalled(); + expect(state.pendingRequest).toBeUndefined(); + expect(state.busy).toBe(false); + expect(state.workerReady).toBe(true); +}); diff --git a/test/pending-output-registry.test.ts b/test/pending-output-registry.test.ts new file mode 100644 index 0000000..ed679d4 --- /dev/null +++ b/test/pending-output-registry.test.ts @@ -0,0 +1,55 @@ +import { expect, it } from "vitest"; +import { reclaimOrphanedOutputs, recordPendingOutput } from "../src/engines/musicgen-acestep/pending-output-registry.js"; + +function memoryStorage(initial: { id: string; at: number }[] = []) { + let value = JSON.stringify(initial); + return { + getItem: () => value, + setItem: (_key: string, next: string) => { + value = next; + }, + records: () => JSON.parse(value) as { id: string; at: number }[], + }; +} + +it("preserves a record added while stale output deletion is in flight", async () => { + const storage = memoryStorage([{ id: "old-output", at: 0 }]); + let finish!: () => void; + const deletion = reclaimOrphanedOutputs( + undefined, + () => + new Promise((resolve) => { + finish = resolve; + }), + { storage, now: () => 10_000_000 }, + ); + await Promise.resolve(); + await recordPendingOutput("new-output", { storage, now: () => 10_000_000 }); + finish(); + await deletion; + expect(storage.records()).toEqual([{ id: "new-output", at: 10_000_000 }]); +}); + +it("retains stale records when output deletion fails", async () => { + const storage = memoryStorage([{ id: "retry-output", at: 0 }]); + await reclaimOrphanedOutputs( + undefined, + async () => { + throw new Error("Storage unavailable"); + }, + { storage, now: () => 10_000_000 }, + ); + expect(storage.records()).toEqual([{ id: "retry-output", at: 0 }]); +}); + +it("keeps output publication best-effort when storage rejects access", async () => { + const storage = { + getItem() { + throw new DOMException("Storage blocked", "SecurityError"); + }, + setItem() { + throw new DOMException("Storage blocked", "SecurityError"); + }, + }; + await expect(recordPendingOutput("output", { storage })).resolves.toBeUndefined(); +}); diff --git a/test/progress-breadcrumb.test.ts b/test/progress-breadcrumb.test.ts new file mode 100644 index 0000000..967ea07 --- /dev/null +++ b/test/progress-breadcrumb.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it } from "vitest"; +import { writeProgressBreadcrumb } from "../src/engines/musicgen-acestep/progress-breadcrumb.js"; + +describe("progress crash markers", () => { + it.each(["Cancelled", "Generation failed", "Stem split failed"])("records %s as a handled end", (title) => { + let saved = ""; + const storage = { + setItem: (_key: string, value: string) => { + saved = value; + }, + }; + writeProgressBreadcrumb("Downloading model", "", true, storage); + expect(JSON.parse(saved).open).toBe(true); + writeProgressBreadcrumb(title, "", false, storage); + expect(JSON.parse(saved)).toMatchObject({ title, open: false }); + }); + + it("does not interrupt work when storage rejects writes", () => { + expect(() => + writeProgressBreadcrumb("Generating song", "", true, { + setItem() { + throw new Error("Storage unavailable"); + }, + }), + ).not.toThrow(); + }); +}); diff --git a/test/resource-session.test.ts b/test/resource-session.test.ts new file mode 100644 index 0000000..f51c22b --- /dev/null +++ b/test/resource-session.test.ts @@ -0,0 +1,118 @@ +import { describe, expect, it, vi } from "vitest"; +import { ResourceSession } from "../src/core/resource-session.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +// Disposable handles exercise ownership without creating models or audio. +function handle() { + return { dispose: vi.fn(async () => {}) }; +} + +describe("resource ownership", () => { + it("releases the previous resource before constructing its replacement", async () => { + const session = new ResourceSession>(); + const first = handle(); + const next = handle(); + await session.load( + async () => first, + async () => {}, + ); + await session.load( + async () => { + expect(first.dispose).toHaveBeenCalledOnce(); + return next; + }, + async () => {}, + ); + expect(session.current).toBe(next); + await session.close(); + expect(next.dispose).toHaveBeenCalledOnce(); + }); + + it("rejects overlapping loads and runs without constructing another resource", async () => { + const session = new ResourceSession>(); + const gate = deferred(); + const item = handle(); + const load = session.load( + async () => item, + () => gate.promise, + ); + const factory = vi.fn(async () => handle()); + await expect(session.load(factory, async () => {})).rejects.toThrow("already in progress"); + await expect(session.run(async () => {})).rejects.toThrow("already in progress"); + expect(factory).not.toHaveBeenCalled(); + gate.resolve(); + await load; + await session.close(); + }); + + it("disposes partial initialization and allows a retry", async () => { + const session = new ResourceSession>(); + const item = handle(); + await expect( + session.load( + async () => item, + async () => { + throw new Error("load failed"); + }, + ), + ).rejects.toThrow("load failed"); + expect(item.dispose).toHaveBeenCalledOnce(); + expect(session.current).toBeUndefined(); + await session.load( + async () => handle(), + async () => {}, + ); + await session.close(); + }); + + it("closes a resource whose load finishes after navigation", async () => { + const session = new ResourceSession>(); + const item = handle(); + const gate = deferred(); + const entered = deferred(); + const load = session.load( + async () => item, + async () => { + entered.resolve(); + await gate.promise; + }, + ); + const rejected = expect(load).rejects.toMatchObject({ name: "AbortError" }); + await entered.promise; + const close = session.close(); + gate.resolve(); + await Promise.all([rejected, close]); + expect(item.dispose).toHaveBeenCalledOnce(); + expect(session.current).toBeUndefined(); + await expect( + session.load( + async () => handle(), + async () => {}, + ), + ).rejects.toThrow("closed"); + }); + + it("waits for in-flight use before teardown and disposes once", async () => { + const session = new ResourceSession>(); + const item = handle(); + await session.load( + async () => item, + async () => {}, + ); + const gate = deferred(); + const run = session.run(() => gate.promise); + const close = session.close(); + expect(session.close()).toBe(close); + expect(item.dispose).not.toHaveBeenCalled(); + gate.resolve(); + await Promise.all([run, close]); + expect(item.dispose).toHaveBeenCalledOnce(); + }); +}); diff --git a/tsconfig.sdk.json b/tsconfig.sdk.json index d1c7a07..6455001 100644 --- a/tsconfig.sdk.json +++ b/tsconfig.sdk.json @@ -8,5 +8,5 @@ "allowImportingTsExtensions": false }, "include": ["src/index.ts", "src/core/**/*.ts", "src/engines/**/*.ts"], - "exclude": ["src/main.ts", "src/music.ts", "src/engines/musicgen-acestep", "src/engines/asr-voicechat"] + "exclude": ["src/main.ts", "src/music.ts", "src/engines/registry.ts", "src/engines/musicgen-acestep", "src/engines/asr-voicechat", "src/engines/stem-dicose", "src/engines/tts-voicechat"] } diff --git a/vitest.config.ts b/vitest.config.ts new file mode 100644 index 0000000..890eda7 --- /dev/null +++ b/vitest.config.ts @@ -0,0 +1,5 @@ +import { defineConfig } from "vitest/config"; + +export default defineConfig({ + test: { include: ["test/*.test.ts"], testTimeout: 10_000 }, +}); From a819cb45fc9fd76a4a7b99a16f6a3649257a9531 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Tue, 15 Sep 2026 13:12:19 -0400 Subject: [PATCH 2/3] Update GitHub Actions runtimes --- .github/workflows/ci.yml | 24 ++++++++++++------------ .github/workflows/deploy.yml | 4 ++-- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index de78459..9475a86 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,8 +16,8 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm @@ -34,8 +34,8 @@ jobs: acestep: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm @@ -49,8 +49,8 @@ jobs: dicose: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm @@ -62,8 +62,8 @@ jobs: format: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm @@ -83,8 +83,8 @@ jobs: kernels: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm @@ -107,8 +107,8 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 2cac08b..000d3f1 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -15,8 +15,8 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-node@v4 + - uses: actions/checkout@v7 + - uses: actions/setup-node@v7 with: node-version: 22 cache: npm From 2e1be1ff795f8b57f854ae54c30287f3ad2b857b Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Tue, 15 Sep 2026 13:33:35 -0400 Subject: [PATCH 3/3] Fix remaining music lifecycle races --- src/engines/musicgen-acestep/index.ts | 8 +- .../pending-output-registry.ts | 37 ++++++++- .../musicgen-acestep/worker-disposal.ts | 36 +++++++++ src/music.ts | 77 ++++++++++++------- test/music-disposal.test.ts | 16 ++++ test/music-page-disposal.test.ts | 63 +++++++++++++++ test/pending-output-registry.test.ts | 43 ++++++++++- 7 files changed, 247 insertions(+), 33 deletions(-) create mode 100644 src/engines/musicgen-acestep/worker-disposal.ts diff --git a/src/engines/musicgen-acestep/index.ts b/src/engines/musicgen-acestep/index.ts index e4b4162..ca00918 100644 --- a/src/engines/musicgen-acestep/index.ts +++ b/src/engines/musicgen-acestep/index.ts @@ -18,6 +18,7 @@ import { aceProductionWorkerConfiguration } from "./config.js"; import { aceInferenceWorkerName } from "./worker-name.js"; import { acquireAceDemoModelCache, deleteAceDemoModelCache } from "./model-cache-migration.js"; import { INITIAL_MODEL_DOWNLOAD_PROGRESS, updateModelDownloadProgress, type ModelDownloadProgress } from "./model-download-progress.js"; +import { waitForWorkerDisposal, type PendingWorkerDisposal } from "./worker-disposal.js"; export { ACE_MAX_DURATION_SECONDS, @@ -65,7 +66,7 @@ export class AceStepMusicClient { private nextRequestId = 1; private nextJobId = 1; private active: ActiveOperation | undefined; - private disposal: { requestId: number; resolve: () => void; reject: (reason: unknown) => void } | undefined; + private disposal: PendingWorkerDisposal | undefined; private disposePromise: Promise | undefined; private fatalGpuDiagnostic = false; /** Releases the shared model-cache lifecycle lock held while the worker is alive. */ @@ -154,11 +155,12 @@ export class AceStepMusicClient { } const requestId = this.nextRequestId++; try { - await new Promise((resolve, reject) => { - this.disposal = { requestId, resolve, reject }; + await waitForWorkerDisposal(requestId, (pending) => { + this.disposal = pending; current.postMessage({ type: "dispose", requestId }); }); } finally { + if (this.disposal?.requestId === requestId) this.disposal = undefined; if (this.worker === current) { current.terminate(); this.worker = undefined; diff --git a/src/engines/musicgen-acestep/pending-output-registry.ts b/src/engines/musicgen-acestep/pending-output-registry.ts index 6509609..72b6808 100644 --- a/src/engines/musicgen-acestep/pending-output-registry.ts +++ b/src/engines/musicgen-acestep/pending-output-registry.ts @@ -1,5 +1,8 @@ +import { acquireCacheLease } from "./cache-lease.js"; + const PENDING_OUTPUTS_KEY = "ace-step-pending-output-ids"; const PENDING_OUTPUTS_LOCK = "ace-step-pending-output-registry"; +const OUTPUT_OWNERSHIP_LOCK_PREFIX = "ace-step-output-ownership:"; const STALE_AFTER_MS = 60 * 60 * 1000; interface PendingOutput { @@ -15,6 +18,19 @@ export interface PendingOutputRegistryOptions { readonly now?: () => number; } +/** Record a committed output and hold shared ownership until the caller releases it. */ +export async function claimPendingOutput(id: string, options: PendingOutputRegistryOptions = {}): Promise<() => Promise> { + const locks = options.locks ?? globalThis.navigator?.locks; + const release = locks === undefined ? async () => {} : await acquireCacheLease(locks, outputOwnershipLock(id)); + try { + await recordPendingOutput(id, options); + return release; + } catch (error) { + await release(); + throw error; + } +} + export async function recordPendingOutput(id: string, options: PendingOutputRegistryOptions = {}): Promise { await updateRecords((records) => [...records.filter((record) => record.id !== id), { id, at: now(options) }], options); } @@ -30,11 +46,12 @@ export async function reclaimOrphanedOutputs( ): Promise { const cutoff = now(options) - STALE_AFTER_MS; const deleted = new Set(); + const locks = options.locks ?? globalThis.navigator?.locks; for (const record of readRecords(storage(options))) { if (record.id === currentId || record.at > cutoff) continue; try { - await release(record.id); - deleted.add(record.id); + const released = await releaseIfUnowned(record.id, release, locks); + if (released) deleted.add(record.id); } catch { // Keep the record so a later visit retries the deletion. } @@ -44,6 +61,22 @@ export async function reclaimOrphanedOutputs( } } +async function releaseIfUnowned(id: string, release: (id: string) => Promise, locks: LockManager | undefined): Promise { + if (locks === undefined) { + await release(id); + return true; + } + return locks.request(outputOwnershipLock(id), { mode: "exclusive", ifAvailable: true }, async (lock) => { + if (lock === null) return false; + await release(id); + return true; + }); +} + +function outputOwnershipLock(id: string): string { + return `${OUTPUT_OWNERSHIP_LOCK_PREFIX}${id}`; +} + async function updateRecords(transform: (records: PendingOutput[]) => PendingOutput[], options: PendingOutputRegistryOptions): Promise { const target = storage(options); const update = () => writeRecords(target, transform(readRecords(target))); diff --git a/src/engines/musicgen-acestep/worker-disposal.ts b/src/engines/musicgen-acestep/worker-disposal.ts new file mode 100644 index 0000000..23a95c5 --- /dev/null +++ b/src/engines/musicgen-acestep/worker-disposal.ts @@ -0,0 +1,36 @@ +export const ACE_WORKER_DISPOSAL_TIMEOUT_MS = 30_000; + +export interface PendingWorkerDisposal { + readonly requestId: number; + readonly resolve: () => void; + readonly reject: (reason: unknown) => void; +} + +/** Start a disposal request and reject if the worker never acknowledges it. */ +export function waitForWorkerDisposal( + requestId: number, + start: (pending: PendingWorkerDisposal) => void, + timeoutMs = ACE_WORKER_DISPOSAL_TIMEOUT_MS, +): Promise { + return new Promise((resolve, reject) => { + let settled = false; + let timer: ReturnType | undefined; + const finish = (outcome: () => void) => { + if (settled) return; + settled = true; + if (timer !== undefined) clearTimeout(timer); + outcome(); + }; + const pending: PendingWorkerDisposal = { + requestId, + resolve: () => finish(resolve), + reject: (reason) => finish(() => reject(reason)), + }; + timer = setTimeout(() => pending.reject(new Error(`ACE worker disposal timed out after ${timeoutMs} ms`)), timeoutMs); + try { + start(pending); + } catch (error) { + pending.reject(error); + } + }); +} diff --git a/src/music.ts b/src/music.ts index 5f45321..69e0495 100644 --- a/src/music.ts +++ b/src/music.ts @@ -26,7 +26,7 @@ import moonIcon from "./engines/musicgen-acestep/assets/moon.png"; import { aceProductionWorkerConfiguration } from "./engines/musicgen-acestep/config.js"; import { aceInferenceWorkerName } from "./engines/musicgen-acestep/worker-name.js"; import { CRASH_BREADCRUMB_KEY, writeProgressBreadcrumb } from "./engines/musicgen-acestep/progress-breadcrumb.js"; -import { forgetPendingOutput, reclaimOrphanedOutputs, recordPendingOutput } from "./engines/musicgen-acestep/pending-output-registry.js"; +import { claimPendingOutput, forgetPendingOutput, reclaimOrphanedOutputs } from "./engines/musicgen-acestep/pending-output-registry.js"; import { formatDecimalBytes, formatModelDownloadAmount, @@ -38,6 +38,7 @@ import { type ModelDownloadProgress, } from "./engines/musicgen-acestep/model-download-progress.js"; import { acquireAceDemoModelCache, deleteAceDemoModelCache } from "./engines/musicgen-acestep/model-cache-migration.js"; +import { waitForWorkerDisposal, type PendingWorkerDisposal } from "./engines/musicgen-acestep/worker-disposal.js"; import { pcmToWav } from "./core/audio.js"; import { localWeightDir } from "./engines/registry.js"; import type { DicoseStemEngine } from "./engines/stem-dicose/index.js"; @@ -113,26 +114,21 @@ let diagnosticDetails: readonly unknown[] = []; let modelProgress: ModelDownloadProgress = INITIAL_MODEL_DOWNLOAD_PROGRESS; let coldDownload = true; let fatalGpuDiagnostic = false; -let output: { readonly url: string; readonly storageId: string } | undefined; +let output: { readonly url: string; readonly storageId: string; readonly releaseOwnership: () => Promise } | undefined; /** The generated song's WAV blob, kept for stem separation (DiCoSe decodes it directly). */ let resultBlob: Blob | undefined; let resultSeed: string | number | bigint = "song"; /** Lazy DiCoSe engine — 623 MB of weights + a ~625 MB GPU buffer, so it only * exists between a "Split stems" click and the next panel reset / pagehide. */ let stemEngine: DicoseStemEngine | undefined; +let stemDisposal: Promise | undefined; let stemUrls: string[] = []; let splittingStems = false; let tooltipRenderFrame: number | undefined; let pendingTooltipPoint: { readonly clientX: number } | undefined; /** Releases the shared model-cache lifecycle lock held while a worker is alive. */ let releaseRuntimeLock: (() => Promise) | undefined; -let disposal: - | { - readonly requestId: number; - readonly resolve: () => void; - readonly reject: (reason: unknown) => void; - } - | undefined; +let disposal: PendingWorkerDisposal | undefined; // Crash breadcrumb: iOS jetsam kills the tab with no error event, so persist // the last progress stage; after an unclean end the next visit reports where @@ -252,16 +248,15 @@ function wireEvents(): void { window.addEventListener("pagehide", () => { pageLifecycle.abort(); generationPreparation?.abort(); - void stemEngine?.dispose(); - stemEngine = undefined; + void resetStemSplitter().catch(() => undefined); if (output !== undefined) { const currentOutput = output; + output = undefined; + resultBlob = undefined; URL.revokeObjectURL(currentOutput.url); // Best-effort only — the reliable path is the pending-output record // reclaimed on the next visit (releaseOrphanedOutputs). - void releaseAceAudioOutput(currentOutput.storageId) - .then(() => forgetPendingOutput(currentOutput.storageId)) - .catch(() => undefined); + void releaseOwnedOutput(currentOutput).catch(() => undefined); } resetWorker(); setBusy(false); @@ -650,13 +645,16 @@ function onWorkerError(event: ErrorEvent): void { } async function publishResult(result: AceGenerationResult): Promise { + let releaseOwnership: (() => Promise) | undefined; try { await releaseCurrentOutput(); + releaseOwnership = await claimPendingOutput(result.audioStorageId); + if (pageLifecycle.signal.aborted) throw new DOMException("Page closed", "AbortError"); const url = URL.createObjectURL(result.audio); - output = { url, storageId: result.audioStorageId }; + output = { url, storageId: result.audioStorageId, releaseOwnership }; + releaseOwnership = undefined; resultBlob = result.audio; resultSeed = result.seed; - await recordPendingOutput(result.audioStorageId); audioPlayer.src = url; audioPlayer.load(); download.href = url; @@ -684,14 +682,22 @@ async function publishResult(result: AceGenerationResult): Promise { await refreshCacheInfo(); } catch (error) { if (output?.storageId !== result.audioStorageId) { - await releaseAceAudioOutput(result.audioStorageId).catch(() => undefined); + try { + await releaseAceAudioOutput(result.audioStorageId); + await forgetPendingOutput(result.audioStorageId); + } catch { + // The pending record remains so a later visit retries cleanup. + } finally { + await releaseOwnership?.(); + } } + if (pageLifecycle.signal.aborted) return; failOperation(`Could not publish the WAV: ${errorMessage(error)}`, false); } } async function releaseCurrentOutput(): Promise { - resetStemSplitter(); + await resetStemSplitter(); const current = output; if (current === undefined) return; output = undefined; @@ -701,8 +707,16 @@ async function releaseCurrentOutput(): Promise { audioPlayer.load(); download.removeAttribute("href"); URL.revokeObjectURL(current.url); - await releaseAceAudioOutput(current.storageId); - await forgetPendingOutput(current.storageId); + await releaseOwnedOutput(current); +} + +async function releaseOwnedOutput(current: { readonly storageId: string; readonly releaseOwnership: () => Promise }): Promise { + try { + await releaseAceAudioOutput(current.storageId); + await forgetPendingOutput(current.storageId); + } finally { + await current.releaseOwnership(); + } } // ── Split stems (DiCoSe, engines/stem-dicose) ──────────────────────────────── @@ -793,12 +807,22 @@ function resetStemsUi(): void { stemsPanel.hidden = true; } -/** Full reset for a new generation / page teardown: the separator holds a - * ~625 MB GPU weight buffer, so it never survives the result panel. */ -function resetStemSplitter(): void { +/** Full reset for a new generation / page teardown. */ +async function resetStemSplitter(): Promise { resetStemsUi(); - void stemEngine?.dispose(); + const current = stemEngine; stemEngine = undefined; + if (current === undefined) { + await stemDisposal; + return; + } + const pending = current.dispose(); + stemDisposal = pending; + try { + await pending; + } finally { + if (stemDisposal === pending) stemDisposal = undefined; + } } function releaseStemUrls(): void { @@ -889,11 +913,12 @@ async function disposeWorker(): Promise { } const requestId = nextRequestId++; try { - await new Promise((resolve, reject) => { - disposal = { requestId, resolve, reject }; + await waitForWorkerDisposal(requestId, (pending) => { + disposal = pending; current.postMessage({ type: "dispose", requestId }); }); } finally { + if (disposal?.requestId === requestId) disposal = undefined; if (worker === current) { current.terminate(); worker = undefined; diff --git a/test/music-disposal.test.ts b/test/music-disposal.test.ts index 089b724..4e944a5 100644 --- a/test/music-disposal.test.ts +++ b/test/music-disposal.test.ts @@ -1,6 +1,7 @@ import { readFileSync } from "node:fs"; import { expect, it, vi } from "vitest"; import { AceStepMusicClient, aceSeed } from "../src/engines/musicgen-acestep/index.js"; +import { ACE_WORKER_DISPOSAL_TIMEOUT_MS } from "../src/engines/musicgen-acestep/worker-disposal.js"; const request = { generationProfile: "ace-turbo-v1-correctness" as const, @@ -99,3 +100,18 @@ it("settles all disposal callers when terminated mid-disposal", async () => { expect(worker.terminate).toHaveBeenCalledOnce(); expect(release).toHaveBeenCalledOnce(); }); + +it("terminates the worker and releases its lease when disposal gets no response", async () => { + vi.useFakeTimers(); + try { + const { client, worker, release } = readyClient(); + const rejected = expect(client.dispose()).rejects.toThrow("disposal timed out"); + await vi.advanceTimersByTimeAsync(ACE_WORKER_DISPOSAL_TIMEOUT_MS); + await rejected; + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(client.initialized).toBe(false); + } finally { + vi.useRealTimers(); + } +}); diff --git a/test/music-page-disposal.test.ts b/test/music-page-disposal.test.ts index d72ac7f..f85fee3 100644 --- a/test/music-page-disposal.test.ts +++ b/test/music-page-disposal.test.ts @@ -3,6 +3,7 @@ import vm from "node:vm"; import ts from "typescript"; import { isAceWorkerMessage } from "ace-step-1.5.wgsl"; import { expect, it, vi } from "vitest"; +import { ACE_WORKER_DISPOSAL_TIMEOUT_MS, waitForWorkerDisposal } from "../src/engines/musicgen-acestep/worker-disposal.js"; it("terminates the page worker and awaits its cache lease after disposal errors", async () => { const source = readFileSync(new URL("../src/music.ts", import.meta.url), "utf8"); @@ -17,6 +18,7 @@ it("terminates the page worker and awaits its cache lease after disposal errors" disposal: undefined, releaseRuntimeLock: release, isAceWorkerMessage, + waitForWorkerDisposal, }; vm.createContext(state); vm.runInContext(ts.transpile(handlers.map((node) => node.getText(ast)).join("\n"), { target: ts.ScriptTarget.ES2022 }), state); @@ -30,3 +32,64 @@ it("terminates the page worker and awaits its cache lease after disposal errors" expect(state.worker).toBeUndefined(); expect(state.workerReady).toBe(false); }); + +it("forces page worker cleanup when disposal gets no response", async () => { + vi.useFakeTimers(); + try { + const source = readFileSync(new URL("../src/music.ts", import.meta.url), "utf8"); + const ast = ts.createSourceFile("music.ts", source, ts.ScriptTarget.Latest, true); + const handler = ast.statements.find((node) => ts.isFunctionDeclaration(node) && node.name?.text === "disposeWorker"); + expect(handler).toBeDefined(); + const release = vi.fn(async () => {}); + const worker = { postMessage: vi.fn(), terminate: vi.fn() }; + const state: Record = { + worker, + workerReady: true, + nextRequestId: 1, + disposal: undefined, + releaseRuntimeLock: release, + waitForWorkerDisposal, + }; + vm.createContext(state); + vm.runInContext(ts.transpile(handler!.getText(ast), { target: ts.ScriptTarget.ES2022 }), state); + const rejected = expect(state.disposeWorker()).rejects.toThrow("disposal timed out"); + await vi.advanceTimersByTimeAsync(ACE_WORKER_DISPOSAL_TIMEOUT_MS); + await rejected; + expect(worker.terminate).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(state.worker).toBeUndefined(); + } finally { + vi.useRealTimers(); + } +}); + +it("waits for stem GPU disposal before releasing the previous result", async () => { + const source = readFileSync(new URL("../src/music.ts", import.meta.url), "utf8"); + const ast = ts.createSourceFile("music.ts", source, ts.ScriptTarget.Latest, true); + const handlers = ast.statements.filter( + (node) => ts.isFunctionDeclaration(node) && ["releaseCurrentOutput", "resetStemSplitter"].includes(node.name?.text ?? ""), + ); + let finish!: () => void; + const disposal = new Promise((resolve) => { + finish = resolve; + }); + const dispose = vi.fn(() => disposal); + const state: Record = { + stemEngine: { dispose }, + stemDisposal: undefined, + resetStemsUi: vi.fn(), + output: undefined, + }; + vm.createContext(state); + vm.runInContext(ts.transpile(handlers.map((node) => node.getText(ast)).join("\n"), { target: ts.ScriptTarget.ES2022 }), state); + let settled = false; + const release = state.releaseCurrentOutput().then(() => { + settled = true; + }); + await Promise.resolve(); + expect(dispose).toHaveBeenCalledOnce(); + expect(settled).toBe(false); + finish(); + await release; + expect(settled).toBe(true); +}); diff --git a/test/pending-output-registry.test.ts b/test/pending-output-registry.test.ts index ed679d4..917d6aa 100644 --- a/test/pending-output-registry.test.ts +++ b/test/pending-output-registry.test.ts @@ -1,5 +1,29 @@ -import { expect, it } from "vitest"; -import { reclaimOrphanedOutputs, recordPendingOutput } from "../src/engines/musicgen-acestep/pending-output-registry.js"; +import { expect, it, vi } from "vitest"; +import { claimPendingOutput, reclaimOrphanedOutputs, recordPendingOutput } from "../src/engines/musicgen-acestep/pending-output-registry.js"; + +class OwnershipLocks { + private readonly shared = new Map(); + + async request(name: string, options: LockOptions, callback: (lock: Lock | null) => unknown): Promise { + if (options.mode === "exclusive" && options.ifAvailable && (this.shared.get(name) ?? 0) > 0) { + return await callback(null); + } + if (options.mode === "shared") this.shared.set(name, (this.shared.get(name) ?? 0) + 1); + try { + return await callback({ name, mode: options.mode ?? "exclusive" } as Lock); + } finally { + if (options.mode === "shared") { + const remaining = (this.shared.get(name) ?? 1) - 1; + if (remaining === 0) this.shared.delete(name); + else this.shared.set(name, remaining); + } + } + } + + asManager(): LockManager { + return this as unknown as LockManager; + } +} function memoryStorage(initial: { id: string; at: number }[] = []) { let value = JSON.stringify(initial); @@ -42,6 +66,21 @@ it("retains stale records when output deletion fails", async () => { expect(storage.records()).toEqual([{ id: "retry-output", at: 0 }]); }); +it("does not reclaim a stale output while another tab still owns it", async () => { + const storage = memoryStorage(); + const locks = new OwnershipLocks().asManager(); + const releaseOwnership = await claimPendingOutput("active-output", { storage, locks, now: () => 0 }); + const release = vi.fn(async () => {}); + await reclaimOrphanedOutputs(undefined, release, { storage, locks, now: () => 10_000_000 }); + expect(release).not.toHaveBeenCalled(); + expect(storage.records()).toEqual([{ id: "active-output", at: 0 }]); + + await releaseOwnership(); + await reclaimOrphanedOutputs(undefined, release, { storage, locks, now: () => 10_000_000 }); + expect(release).toHaveBeenCalledExactlyOnceWith("active-output"); + expect(storage.records()).toEqual([]); +}); + it("keeps output publication best-effort when storage rejects access", async () => { const storage = { getItem() {