From e504a31efddf3ef0059d6572fab6d29fd2c97740 Mon Sep 17 00:00:00 2001 From: Lotus CI Date: Sun, 13 Sep 2026 01:22:20 -0400 Subject: [PATCH] =?UTF-8?q?fix(lotus):=20denoise=20=E2=80=94=20prefetch=20?= =?UTF-8?q?assets,=20RNNoise=20fallback=20+=20host=20notify,=20suspend=20w?= =?UTF-8?q?hile=20muted,=20mono/delay=20graph,=20typed=20modules?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Assets (context, worklets, wasm, DFN core) are prepared as soon as the flag is seen, so init() under LiveKit's trackChangeLock only wires already-loaded pieces; resume timeout 3 s -> 500 ms (#7). - init failure retries once with rnnoise; success/failure is reported to the host as io.lotus.denoise_state so the UI can reflect reality (#8). - Mic TrackMuted/TrackUnmuted suspend/resume the processor's context so no inference runs on silence (#9). - Every node is explicit mono; the dry path gets a per-model DelayNode so the floor mix no longer comb-filters (#24, #25). - DTLN/DFN dynamic imports are typed and their exports asserted at load, feeding the #8 fallback instead of failing silently (#26). Unit-tested (13 tests across the two files). Fixes #7 Fixes #8 Fixes #9 Fixes #24 Fixes #25 Fixes #26 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01PPmy3tPq869XDW4njjVaKA --- src/lotus/lotusDenoise.test.ts | 143 ++++++++- src/lotus/lotusDenoise.ts | 104 ++++++- src/lotus/lotusDenoiseProcessor.test.ts | 210 +++++++++++++- src/lotus/lotusDenoiseProcessor.ts | 370 ++++++++++++++++++++---- 4 files changed, 740 insertions(+), 87 deletions(-) diff --git a/src/lotus/lotusDenoise.test.ts b/src/lotus/lotusDenoise.test.ts index 8006c8a2..5b43599c 100644 --- a/src/lotus/lotusDenoise.test.ts +++ b/src/lotus/lotusDenoise.test.ts @@ -10,21 +10,40 @@ import { ParticipantEvent, Track } from "livekit-client"; import { type CallViewModel } from "../state/CallViewModel/CallViewModel"; import { startLotusDenoise } from "./lotusDenoise"; +import { LotusWidgetActions } from "./lotusActions"; // Track constructed instances so tests can assert exactly one processor was // built across racing apply() calls, and can assert the pending one is // destroyed on early teardown. `vi.hoisted` is required because `vi.mock` // factories are hoisted above this file's other top-level statements. const instances = vi.hoisted( - () => [] as { destroy: ReturnType }[], + () => + [] as { + config: { model: string }; + destroy: ReturnType; + setMicMuted: ReturnType; + }[], ); +// Ordered log of the heavy-lifting calls, to assert prepare runs before +// setProcessor (#7). +const calls = vi.hoisted(() => [] as string[]); vi.mock("./lotusDenoiseProcessor", () => ({ LotusDenoiseProcessor: class { public destroy = vi.fn().mockResolvedValue(undefined); - public constructor() { + public setMicMuted = vi.fn(); + public constructor(public config: { model: string }) { instances.push(this); } }, + prepareDenoiseAssets: vi.fn(async () => { + calls.push("prepare"); + await Promise.resolve(); + }), + releasePreparedDenoiseAssets: vi.fn().mockResolvedValue(undefined), +})); +const hostSend = vi.hoisted(() => vi.fn().mockResolvedValue({})); +vi.mock("../widget", () => ({ + widget: { api: { transport: { send: hostSend } } }, })); /** A promise plus externally-callable resolve, for controlling ordering. */ @@ -42,19 +61,23 @@ function deferred(): { function makeRoomAndVm(mic: { getProcessor: () => unknown; setProcessor: ReturnType; + isMuted?: boolean; }): { vm: CallViewModel; room: { localParticipant: Record }; firePublished: () => void; + fire: (event: string, arg?: unknown) => void; } { - const handlers = new Map void>(); + const handlers = new Map void>(); const localParticipant = { getTrackPublication: ( source: Track.Source, ): { track: typeof mic } | undefined => source === Track.Source.Microphone ? { track: mic } : undefined, - on: (event: string, cb: () => void): Map void> => - handlers.set(event, cb), + on: ( + event: string, + cb: (arg?: unknown) => void, + ): Map void> => handlers.set(event, cb), off: (event: string): boolean => handlers.delete(event), }; const room = { localParticipant }; @@ -74,11 +97,19 @@ function makeRoomAndVm(mic: { vm, room, firePublished: () => handlers.get(ParticipantEvent.LocalTrackPublished)?.(), + fire: (event, arg) => handlers.get(event)?.(arg), }; } +/** Let the async apply() chain (attach → retry → host notify) settle. */ +async function flush(): Promise { + for (let i = 0; i < 8; i++) await Promise.resolve(); +} + beforeEach(() => { instances.length = 0; + calls.length = 0; + hostSend.mockClear(); // `lotusParam`/`lotusFlag` cache the URL params at first read; seed the hash // before startLotusDenoise() so the dedicated `lotusDenoiseSource` flag reads on. window.location.hash = "#/room?lotusDenoiseSource=1"; @@ -95,6 +126,7 @@ describe("startLotusDenoise", () => { const mic = { getProcessor: (): unknown => attached, setProcessor: vi.fn(async (p: unknown) => { + calls.push("setProcessor"); await setProcessorDeferred.promise; attached = p; }), @@ -102,6 +134,8 @@ describe("startLotusDenoise", () => { const { vm, firePublished } = makeRoomAndVm(mic); startLotusDenoise(vm); + // [#7] The heavy assets were kicked off before setProcessor() was called. + expect(calls).toEqual(["prepare", "setProcessor"]); // Simulate a second LocalTrackPublished (e.g. camera) firing before the // first setProcessor() has resolved. firePublished(); @@ -138,4 +172,103 @@ describe("startLotusDenoise", () => { setProcessorDeferred.resolve(); await Promise.resolve(); }); + + test("reports the active model to the host once attached (#8)", async () => { + let attached: unknown; + const mic = { + getProcessor: (): unknown => attached, + setProcessor: vi.fn(async (p: unknown) => { + await Promise.resolve(); + attached = p; + }), + }; + const { vm } = makeRoomAndVm(mic); + startLotusDenoise(vm); + await flush(); + expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, { + active: true, + model: "rnnoise", + }); + }); + + test("retries once with rnnoise when the requested model fails to init (#8)", async () => { + window.location.hash = + "#/room?lotusDenoiseSource=1&lotusModel=deepfilternet"; + let attached: unknown; + const mic = { + getProcessor: (): unknown => attached, + setProcessor: vi.fn(async (p: { config: { model: string } }) => { + await Promise.resolve(); + if (p.config.model === "deepfilternet") + throw new Error("dfn assets missing"); + attached = p; + }), + }; + const { vm } = makeRoomAndVm(mic); + startLotusDenoise(vm); + await flush(); + + expect(instances.map((i) => i.config.model)).toEqual([ + "deepfilternet", + "rnnoise", + ]); + expect(mic.setProcessor).toHaveBeenCalledTimes(2); + expect(hostSend).toHaveBeenCalledTimes(1); + expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, { + active: true, + model: "rnnoise", + }); + }); + + test("notifies the host with the error when the rnnoise fallback also fails (#8)", async () => { + window.location.hash = "#/room?lotusDenoiseSource=1&lotusModel=dtln"; + const mic = { + getProcessor: (): unknown => undefined, + setProcessor: vi.fn(async () => { + await Promise.resolve(); + throw new Error("no worklet for you"); + }), + }; + const { vm } = makeRoomAndVm(mic); + startLotusDenoise(vm); + await flush(); + + expect(instances.map((i) => i.config.model)).toEqual(["dtln", "rnnoise"]); + expect(hostSend).toHaveBeenCalledTimes(1); + expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, { + active: false, + error: "no worklet for you", + model: "rnnoise", + }); + }); + + test("mirrors mic TrackMuted/TrackUnmuted onto the processor (#9)", async () => { + let attached: unknown; + const mic = { + isMuted: true, + getProcessor: (): unknown => attached, + setProcessor: vi.fn(async (p: unknown) => { + await Promise.resolve(); + attached = p; + }), + }; + const { vm, fire } = makeRoomAndVm(mic); + startLotusDenoise(vm); + // Seeded from the mic's current state before setProcessor(). + expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(true); + await flush(); + + const micPub = { source: Track.Source.Microphone }; + const camPub = { source: Track.Source.Camera }; + fire(ParticipantEvent.TrackUnmuted, micPub); + expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(false); + fire(ParticipantEvent.TrackMuted, micPub); + expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(true); + + // Camera mute must not touch the audio context. + instances[0].setMicMuted.mockClear(); + fire(ParticipantEvent.TrackMuted, camPub); + fire(ParticipantEvent.TrackUnmuted, camPub); + expect(instances[0].setMicMuted).not.toHaveBeenCalled(); + }); }); diff --git a/src/lotus/lotusDenoise.ts b/src/lotus/lotusDenoise.ts index 7c87e6d0..c9c00ff3 100644 --- a/src/lotus/lotusDenoise.ts +++ b/src/lotus/lotusDenoise.ts @@ -10,14 +10,23 @@ import { ParticipantEvent, type Room as LivekitRoom, Track, + type TrackPublication, } from "livekit-client"; import { logger } from "matrix-js-sdk/lib/logger"; import { type CallViewModel } from "../state/CallViewModel/CallViewModel"; -import { lotusFlag, lotusParam } from "./lotusWidget"; +import { + LotusWidgetActions, + lotusFlag, + lotusParam, + lotusSendToHost, +} from "./lotusWidget"; import { type LotusDenoiseConfig, + type LotusDenoiseModel, LotusDenoiseProcessor, + prepareDenoiseAssets, + releasePreparedDenoiseAssets, } from "./lotusDenoiseProcessor"; /** @@ -93,6 +102,12 @@ export function startLotusDenoise(vm: CallViewModel): () => void { : 0.15, }; + // [lotus #7] Start fetching wasm / creating the AudioContext / addModule-ing + // the worklet (and loading the DFN model) NOW, before any mic track exists. + // `setProcessor()` holds LiveKit's trackChangeLock while awaiting `init()`, + // so anything still loading there freezes mute/unmute/device-switch. + void prepareDenoiseAssets(config); + const micOf = (room: LivekitRoom): LocalAudioTrack | undefined => room.localParticipant.getTrackPublication(Track.Source.Microphone) ?.track as LocalAudioTrack | undefined; @@ -105,23 +120,73 @@ export function startLotusDenoise(vm: CallViewModel): () => void { // AudioContext + model load) before the first has attached. Track an // in-flight setProcessor per room and skip `apply()` while one is pending. const pendingProcessors = new Map(); + let stopped = false; + + const attach = async ( + room: LivekitRoom, + mic: LocalAudioTrack, + model: LotusDenoiseModel, + ): Promise => { + const processor = new LotusDenoiseProcessor({ ...config, model }); + pendingProcessors.set(room, processor); + // [lotus #9] Seed the mute state so a processor attached while already + // muted starts suspended rather than burning inference on silence. + processor.setMicMuted(mic.isMuted); + try { + await mic.setProcessor(processor); + } finally { + // Only clear if we're still the pending entry (a teardown that ran + // while this was in flight may have already replaced/removed it). + if (pendingProcessors.get(room) === processor) + pendingProcessors.delete(room); + } + }; const apply = (room: LivekitRoom): void => { const mic = micOf(room); if (!mic || mic.getProcessor() || pendingProcessors.has(room)) return; - const processor = new LotusDenoiseProcessor(config); - pendingProcessors.set(room, processor); - void mic - .setProcessor(processor) - .catch((e) => logger.warn("[lotus] denoise setProcessor failed", e)) - .finally(() => { - // Only clear if we're still the pending entry (a teardown that ran - // while this was in flight may have already replaced/removed it). - if (pendingProcessors.get(room) === processor) - pendingProcessors.delete(room); - }); + void (async () => { + let model = config.model; + try { + try { + await attach(room, mic, model); + } catch (e) { + logger.warn("[lotus] denoise setProcessor failed", e); + // [lotus #8] A failed init leaves LiveKit with no processor (it only + // assigns after init resolves), so retry once with the smallest, + // most portable tier before giving up. + if (model === "rnnoise" || stopped) throw e; + model = "rnnoise"; + await attach(room, mic, model); + } + // [lotus #8] Tell the host what is ACTUALLY running so its toggle + // reflects reality (including the fallback model). + lotusSendToHost(LotusWidgetActions.DenoiseState, { + active: true, + model, + }); + } catch (e) { + logger.warn("[lotus] denoise unavailable; publishing raw mic", e); + lotusSendToHost(LotusWidgetActions.DenoiseState, { + active: false, + error: e instanceof Error ? e.message : String(e), + model, + }); + } + })(); }; + // [lotus #9] Suspend/resume the processor's own AudioContext with the mic's + // mute state (LiveKit keeps the muted track flowing silent frames, so the + // worklet would otherwise run inference the whole time the user is muted). + const onMuteChange = + (room: LivekitRoom, muted: boolean) => + (pub: TrackPublication): void => { + if (pub.source !== Track.Source.Microphone) return; + const p = pendingProcessors.get(room) ?? micOf(room)?.getProcessor(); + if (p instanceof LotusDenoiseProcessor) p.setMicMuted(muted); + }; + const roomListeners = new Map void>(); let rooms: LivekitRoom[] = []; @@ -143,22 +208,29 @@ export function startLotusDenoise(vm: CallViewModel): () => void { // Re-attach on every (re)publish — this is what makes denoise survive // reconnects (A7), unlike the old getUserMedia patch. const onPublished = (): void => apply(room); + const onMuted = onMuteChange(room, true); + const onUnmuted = onMuteChange(room, false); room.localParticipant.on( ParticipantEvent.LocalTrackPublished, onPublished, ); - roomListeners.set(room, () => + room.localParticipant.on(ParticipantEvent.TrackMuted, onMuted); + room.localParticipant.on(ParticipantEvent.TrackUnmuted, onUnmuted); + roomListeners.set(room, () => { room.localParticipant.off( ParticipantEvent.LocalTrackPublished, onPublished, - ), - ); + ); + room.localParticipant.off(ParticipantEvent.TrackMuted, onMuted); + room.localParticipant.off(ParticipantEvent.TrackUnmuted, onUnmuted); + }); apply(room); } } }); return () => { + stopped = true; sub.unsubscribe(); for (const off of roomListeners.values()) off(); roomListeners.clear(); @@ -175,5 +247,7 @@ export function startLotusDenoise(vm: CallViewModel): () => void { } } pendingProcessors.clear(); + // [lotus #7] Close any prepared-but-never-claimed context (e.g. no mic). + void releasePreparedDenoiseAssets(); }; } diff --git a/src/lotus/lotusDenoiseProcessor.test.ts b/src/lotus/lotusDenoiseProcessor.test.ts index 35573d8f..5a8a1cad 100644 --- a/src/lotus/lotusDenoiseProcessor.test.ts +++ b/src/lotus/lotusDenoiseProcessor.test.ts @@ -5,19 +5,97 @@ SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial Please see LICENSE in the repository root for full details. */ -import { describe, expect, test } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { type AudioProcessorOptions } from "livekit-client"; -import { LotusDenoiseProcessor } from "./lotusDenoiseProcessor"; +import { + assertModuleExport, + type LotusDenoiseConfig, + LotusDenoiseProcessor, + prepareDenoiseAssets, + releasePreparedDenoiseAssets, +} from "./lotusDenoiseProcessor"; -function makeProcessor(): LotusDenoiseProcessor { - return new LotusDenoiseProcessor({ - model: "rnnoise", - assetBase: "https://example.invalid/denoise/", - gate: false, - gateThreshold: -45, - floor: 0.15, +const baseConfig: LotusDenoiseConfig = { + model: "rnnoise", + assetBase: "https://example.invalid/denoise/", + gate: false, + gateThreshold: -45, + floor: 0.15, +}; + +function makeProcessor( + overrides: Partial = {}, +): LotusDenoiseProcessor { + return new LotusDenoiseProcessor({ ...baseConfig, ...overrides }); +} + +/** + * Minimal fake AudioContext: tracks state, records suspend/resume, and lets a + * test fire `statechange` like the browser would after an OS interruption. + */ +class FakeAudioContext extends EventTarget { + public static created: FakeAudioContext[] = []; + public state: AudioContextState = "running"; + public readonly sampleRate: number; + public readonly audioWorklet = { + addModule: vi.fn().mockResolvedValue(undefined), + }; + public suspend = vi.fn(async (): Promise => { + await Promise.resolve(); + this.state = "suspended"; }); + public resume = vi.fn(async (): Promise => { + await Promise.resolve(); + this.state = "running"; + }); + public close = vi.fn(async (): Promise => { + await Promise.resolve(); + this.state = "closed"; + }); + public constructor(opts?: { sampleRate?: number }) { + super(); + this.sampleRate = opts?.sampleRate ?? 48_000; + FakeAudioContext.created.push(this); + } + /** Simulate an external (OS/browser) suspension. */ + public externallySuspend(): void { + this.state = "suspended"; + this.dispatchEvent(new Event("statechange")); + } +} + +beforeEach(() => { + FakeAudioContext.created = []; + vi.stubGlobal("AudioContext", FakeAudioContext); + // The flat wasm prefetch just warms a cache; give it a 200 so prepare() + // succeeds without a network. + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue({ + ok: true, + arrayBuffer: (): ArrayBuffer => new ArrayBuffer(8), + }), + ); +}); + +afterEach(async () => { + await releasePreparedDenoiseAssets(); + vi.unstubAllGlobals(); +}); + +/** Skip the real graph (needs worklets) — init() only needs a track back. */ +function stubGraph(processor: LotusDenoiseProcessor): void { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (processor as any).buildGraph = async (): Promise => { + await Promise.resolve(); + return { + source: { disconnect: vi.fn() }, + nodes: [], + disposes: [], + track: { stop: vi.fn() }, + }; + }; } describe("LotusDenoiseProcessor.restart", () => { @@ -53,3 +131,117 @@ describe("LotusDenoiseProcessor.restart", () => { expect(processor.processedTrack).not.toBe(rawTrack); }); }); + +describe("prepareDenoiseAssets (#7)", () => { + test("init() adopts the context/worklet prepared before setProcessor instead of loading inline", async () => { + await prepareDenoiseAssets(baseConfig); + // The heavy parts ran up front: one context, worklet module added, wasm fetched. + expect(FakeAudioContext.created).toHaveLength(1); + const prepared = FakeAudioContext.created[0]; + expect(prepared.audioWorklet.addModule).toHaveBeenCalledWith( + `${baseConfig.assetBase}rnnoiseWorklet.js`, + ); + expect(fetch).toHaveBeenCalled(); + + const processor = makeProcessor(); + stubGraph(processor); + await processor.init({ + track: {} as MediaStreamTrack, + } as unknown as AudioProcessorOptions); + + // No second AudioContext / addModule inside init (i.e. under the lock). + expect(FakeAudioContext.created).toHaveLength(1); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + expect((processor as any).ctx).toBe(prepared); + + // The claim is one-shot: a later processor prepares its own. + const second = makeProcessor(); + stubGraph(second); + await second.init({ + track: {} as MediaStreamTrack, + } as unknown as AudioProcessorOptions); + expect(FakeAudioContext.created).toHaveLength(2); + + await processor.destroy(); + await second.destroy(); + }); + + test("releasePreparedDenoiseAssets closes an unclaimed prepared context", async () => { + await prepareDenoiseAssets(baseConfig); + const prepared = FakeAudioContext.created[0]; + await releasePreparedDenoiseAssets(); + expect(prepared.close).toHaveBeenCalled(); + }); +}); + +describe("LotusDenoiseProcessor.setMicMuted (#9)", () => { + test("suspends the owned context on mute, resumes on unmute, and the statechange watcher leaves the mute suspension alone", async () => { + const processor = makeProcessor(); + stubGraph(processor); + await processor.init({ + track: {} as MediaStreamTrack, + } as unknown as AudioProcessorOptions); + const ctx = FakeAudioContext.created[0]; + expect(ctx.state).toBe("running"); + + processor.setMicMuted(true); + await Promise.resolve(); + expect(ctx.suspend).toHaveBeenCalledTimes(1); + expect(ctx.state).toBe("suspended"); + + // The browser fires statechange for our own suspend too; the watcher must + // NOT undo it while muted. + ctx.resume.mockClear(); + ctx.dispatchEvent(new Event("statechange")); + await Promise.resolve(); + expect(ctx.resume).not.toHaveBeenCalled(); + expect(ctx.state).toBe("suspended"); + + processor.setMicMuted(false); + await Promise.resolve(); + expect(ctx.resume).toHaveBeenCalledTimes(1); + expect(ctx.state).toBe("running"); + + // An EXTERNAL suspension while unmuted is still healed by the watcher. + ctx.resume.mockClear(); + ctx.externallySuspend(); + await Promise.resolve(); + expect(ctx.resume).toHaveBeenCalledTimes(1); + + await processor.destroy(); + }); + + test("a processor attached while already muted starts suspended", async () => { + const processor = makeProcessor(); + stubGraph(processor); + processor.setMicMuted(true); + await processor.init({ + track: {} as MediaStreamTrack, + } as unknown as AudioProcessorOptions); + const ctx = FakeAudioContext.created[0]; + expect(ctx.state).toBe("suspended"); + await processor.destroy(); + }); +}); + +describe("assertModuleExport (#26)", () => { + test("returns the module when the export is a function", () => { + const mod = { createNoiseSuppressionAudioWorklet: (): void => undefined }; + expect( + assertModuleExport(mod, "createNoiseSuppressionAudioWorklet", "x.js"), + ).toBe(mod); + }); + + test("throws a clear error naming the module and export when missing", () => { + expect(() => + assertModuleExport( + { other: 1 }, + "DeepFilterNet3Core", + "dfn/index.esm.js", + ), + ).toThrow(/dfn\/index\.esm\.js does not export DeepFilterNet3Core/); + expect(() => assertModuleExport(undefined, "X", "y.js")).toThrow( + /does not export X/, + ); + }); +}); diff --git a/src/lotus/lotusDenoiseProcessor.ts b/src/lotus/lotusDenoiseProcessor.ts index 3cef9f63..1f63385d 100644 --- a/src/lotus/lotusDenoiseProcessor.ts +++ b/src/lotus/lotusDenoiseProcessor.ts @@ -95,7 +95,9 @@ async function fetchWasm(url: string): Promise { * lands, and a still-suspended context degrades to (temporary) silence that the * watcher heals, not a hang. */ -async function resumeCtx(ctx: AudioContext, timeoutMs = 3_000): Promise { +// [lotus] 500 ms, not 3 s: this still runs under LiveKit's trackChangeLock (#7), +// and the statechange watcher heals a still-suspended context later anyway. +async function resumeCtx(ctx: AudioContext, timeoutMs = 500): Promise { await Promise.race([ ctx.resume().catch(() => undefined), new Promise((resolve) => setTimeout(resolve, timeoutMs)), @@ -127,6 +129,226 @@ interface Graph { track: MediaStreamTrack; } +// [lotus #26] Minimal local contracts for the two dynamically-imported ESM +// helpers (not bundled here — see the CONTRACT note above). Their exports are +// asserted at runtime so an asset bump that renames/removes one fails loudly +// (and flows into the rnnoise fallback in lotusDenoise.ts) instead of as a +// vague TypeError deep inside `init()`. +interface DtlnModule { + createNoiseSuppressionAudioWorklet: ( + ctx: AudioContext, + opts: { bypassUntilReady: boolean }, + ) => Promise; +} +interface DfnCore { + initialize: () => Promise; + createAudioWorkletNode: (ctx: AudioContext) => Promise; + destroy: () => void; +} +interface DfnModule { + DeepFilterNet3Core: new (opts: { + sampleRate: number; + noiseReductionLevel: number; + assetConfig: { cdnUrl: string }; + }) => DfnCore; +} + +/** Throw a clear error if a dynamic-import module lacks an expected export. */ +export function assertModuleExport( + mod: unknown, + name: string, + url: string, +): T { + const exp = (mod as Record | null | undefined)?.[name]; + if (typeof exp !== "function") + throw new Error( + `denoise: ${url} does not export ${name} (got ${typeof exp}) — asset/version mismatch`, + ); + return mod as T; +} + +async function loadDfnCore(config: LotusDenoiseConfig): Promise { + const base = config.assetBase; + const url = `${base}deepfilternet/index.esm.js`; + const dfnBase = new URL(`${base}deepfilternet`, window.location.href).href; + const mod = assertModuleExport( + await import(/* @vite-ignore */ url), + "DeepFilterNet3Core", + url, + ); + const core = new mod.DeepFilterNet3Core({ + sampleRate: 48_000, + // 60, not 80: full-strength suppression is the main source of the + // "over-processed" character; a lower level keeps voice natural while + // the dry/wet floor handles the noise tail. + noiseReductionLevel: 60, + assetConfig: { cdnUrl: dfnBase }, + }); + await core.initialize(); + return core; +} + +async function loadDtlnModule(config: LotusDenoiseConfig): Promise { + const url = `${config.assetBase}workadventure/audio-worklet.js`; + return assertModuleExport( + await import(/* @vite-ignore */ url), + "createNoiseSuppressionAudioWorklet", + url, + ); +} + +/** Which wasm file a flat model uses (SIMD build when supported). */ +function flatWasmFiles(model: "rnnoise" | "speex"): { + primary: string; + fallback?: string; +} { + const flat = FLAT[model]; + const useSimd = model === "rnnoise" && !!flat.simdWasm && supportsSimd(); + return useSimd + ? { primary: flat.simdWasm!, fallback: flat.wasm } + : { primary: flat.wasm }; +} + +// [lotus #24] Force every node in the graph to a single, explicitly-downmixed +// channel. Without `channelCountMode: "explicit"` the default ("max") IGNORES +// `channelCount`, so a stereo capture device would feed 2 channels into a +// worklet configured with `maxChannels: 1` and sum a stereo dry copy against a +// mono wet one at the destination. +const MONO: AudioNodeOptions = { + channelCount: 1, + channelCountMode: "explicit", + channelInterpretation: "speakers", +}; + +// [lotus #25] Algorithmic latency of each model in samples at its native rate, +// used to delay the DRY copy of the floor mix so it lines up with the wet path +// (otherwise the sum comb-filters — a hollow/phasey colouration on voice). +// - rnnoise: 480-sample (10 ms @ 48 kHz) frames; the sapphi worklet buffers +// 128-sample quanta up to one frame, so the wet path lags by one frame. +// - speex: the sapphi speex worklet uses the same 480-sample framing. +// - dtln: 512-sample block / 128 hop @ 16 kHz (~32 ms) per the DTLN paper — +// best-known, unmeasured (the floor is not mixed for dtln, see buildGraph). +// - deepfilternet: 480-sample hop + 2-frame lookahead @ 48 kHz (~30 ms) per +// DeepFilterNet3 — best-known, unmeasured (floor not mixed for dfn either). +const DRY_DELAY_SAMPLES: Record = { + rnnoise: 480, + speex: 480, + dtln: 512, + deepfilternet: 1440, +}; + +/** + * Create the model-rate context and register the flat/gate worklet modules. + * Closes the context (and rethrows) on any failure so nothing half-built leaks. + */ +async function createModelContext( + config: LotusDenoiseConfig, +): Promise { + const rate = sampleRateFor(config.model); + const ctx = new AudioContext({ sampleRate: rate }); + try { + if (ctx.sampleRate !== rate) + throw new Error(`denoise: got ${ctx.sampleRate}Hz, need ${rate}Hz`); + // Flat models register via addModule here; DTLN/DeepFilterNet bring their + // own processor via the dynamic-imported helper (see buildMlNode). + if (config.model === "rnnoise" || config.model === "speex") + await ctx.audioWorklet.addModule( + config.assetBase + FLAT[config.model].script, + ); + if (config.gate) + await ctx.audioWorklet.addModule(config.assetBase + GATE.script); + return ctx; + } catch (e) { + await ctx.close().catch(() => undefined); + throw e; + } +} + +// [lotus #7] Everything heavy that `init()` needs but that does NOT depend on +// the mic track: the AudioContext + worklet modules, the flat wasm binary, and +// (DFN) the fully-initialised model core. `LocalAudioTrack.setProcessor()` +// holds LiveKit's `trackChangeLock` while awaiting `init()`, so every +// mute/unmute/device-switch queues behind it — prepare these as soon as the +// flag is seen (before any track exists) so `init()` only wires them up. +interface PreparedAssets { + ctx: AudioContext; + dfnCore?: DfnCore; +} +const preparedAssets = new Map>(); +const preparedKey = (c: LotusDenoiseConfig): string => + `${c.model}|${c.gate ? 1 : 0}|${c.assetBase}`; + +async function prepareUncached( + config: LotusDenoiseConfig, +): Promise { + const ctx = await createModelContext(config); + try { + let dfnCore: DfnCore | undefined; + if (config.model === "rnnoise" || config.model === "speex") { + const { primary, fallback } = flatWasmFiles(config.model); + // Warm the wasm cache; a SIMD miss is fine — buildMlNode falls back. + await fetchWasm(config.assetBase + primary).catch(async () => + fallback ? fetchWasm(config.assetBase + fallback) : undefined, + ); + } else if (config.model === "dtln") { + await loadDtlnModule(config); // warms the browser's module map + } else { + dfnCore = await loadDfnCore(config); + } + return { ctx, dfnCore }; + } catch (e) { + await ctx.close().catch(() => undefined); + throw e; + } +} + +/** + * Prefetch/prepare the assets for `config` (idempotent per model). Never + * rejects: a failed prepare is evicted so `init()` simply loads inline and + * surfaces the real error there. + */ +export async function prepareDenoiseAssets( + config: LotusDenoiseConfig, +): Promise { + const key = preparedKey(config); + let p = preparedAssets.get(key); + if (!p) { + p = prepareUncached(config); + void p.catch((e) => { + if (preparedAssets.get(key) === p) preparedAssets.delete(key); + logger.warn(`[lotus] denoise prepare failed (${config.model})`, e); + }); + preparedAssets.set(key, p); + } + await p.catch(() => undefined); +} + +/** Take (one-shot) the prepared assets for `config`, if any were prepared. */ +function claimPreparedAssets( + config: LotusDenoiseConfig, +): Promise | undefined { + const key = preparedKey(config); + const p = preparedAssets.get(key); + if (p) preparedAssets.delete(key); + return p; +} + +/** Close any prepared-but-unclaimed contexts (call on feature teardown). */ +export async function releasePreparedDenoiseAssets(): Promise { + const all = [...preparedAssets.values()]; + preparedAssets.clear(); + await Promise.all( + all.map(async (p) => + p + .then(async (a) => { + safeCall(() => a.dfnCore?.destroy()); + if (a.ctx.state !== "closed") await a.ctx.close(); + }) + .catch(() => undefined), + ), + ); +} + /** * A LiveKit audio TrackProcessor that runs Lotus ML noise suppression * (RNNoise / Speex / DTLN / DeepFilterNet) on the local microphone track, as a @@ -151,9 +373,31 @@ export class LotusDenoiseProcessor implements TrackProcessor< private ctx?: AudioContext; private graph?: Graph; private ctxStateHandler?: () => void; + private preparedDfnCore?: DfnCore; + // [lotus #9] True while the mic is muted: we suspend our own context so the + // worklet stops running inference on silence, and the statechange watcher + // must not "heal" that intentional suspension. + private micMuted = false; public constructor(private readonly config: LotusDenoiseConfig) {} + /** + * [lotus #9] Mirror the mic's mute state onto the owned context. EC uses + * `stopMicTrackOnMute: false`, so a muted mic keeps producing (silent) frames + * and the ML worklet would otherwise keep running full inference for the + * whole time the user is muted. + */ + public setMicMuted(muted: boolean): void { + this.micMuted = muted; + const ctx = this.ctx; + if (!ctx || ctx.state === "closed") return; + if (muted) { + if (ctx.state === "running") void ctx.suspend().catch(() => undefined); + } else if (ctx.state === "suspended" && this.graph) { + void ctx.resume().catch(() => undefined); + } + } + public async init(_opts: AudioProcessorOptions): Promise { try { await this.ensureContext(); @@ -163,6 +407,9 @@ export class LotusDenoiseProcessor implements TrackProcessor< // Don't orphan the owned context if graph construction fails (browsers // cap live AudioContexts, so repeated failed inits could exhaust them). // The caller degrades to the raw mic; we just release our resources. + const core = this.preparedDfnCore; + this.preparedDfnCore = undefined; + if (core) safeCall(() => core.destroy()); await this.closeContext(); throw e; } @@ -194,6 +441,9 @@ export class LotusDenoiseProcessor implements TrackProcessor< this.disposeGraph(this.graph); this.graph = undefined; this.processedTrack = undefined; + const core = this.preparedDfnCore; + this.preparedDfnCore = undefined; + if (core) safeCall(() => core.destroy()); await this.closeContext(); } @@ -209,7 +459,7 @@ export class LotusDenoiseProcessor implements TrackProcessor< if (ctx.state !== "closed") await ctx.close().catch(() => undefined); } - /** Create (once) the model-rate context + register the flat worklet modules. */ + /** Adopt the prepared context (or create one) + install the state watcher. */ private async ensureContext(): Promise { const rate = sampleRateFor(this.config.model); if ( @@ -217,36 +467,43 @@ export class LotusDenoiseProcessor implements TrackProcessor< this.ctx.state !== "closed" && this.ctx.sampleRate === rate ) { - if (this.ctx.state === "suspended") await resumeCtx(this.ctx); + if (this.ctx.state === "suspended" && !this.micMuted) + await resumeCtx(this.ctx); return; } await this.closeContext(); - const ctx = new AudioContext({ sampleRate: rate }); + // [lotus #7] Prefer the context/modules/model prepared before + // setProcessor() was called; only load inline if nothing was prepared + // (e.g. the rnnoise fallback path, or a second processor after a + // republish). + const claimed = await claimPreparedAssets(this.config)?.catch( + () => undefined, + ); + let ctx: AudioContext; + if (claimed && claimed.ctx.state !== "closed") { + ctx = claimed.ctx; + this.preparedDfnCore = claimed.dfnCore; + } else { + ctx = await createModelContext(this.config); + } try { - if (ctx.sampleRate !== rate) - throw new Error(`denoise: got ${ctx.sampleRate}Hz, need ${rate}Hz`); - // Auto-resume if the OS/browser suspends the context mid-call (mobile // backgrounding, audio interruption): the dest node otherwise emits - // silence with no recovery. Only resume while a graph is live. + // silence with no recovery. Only resume while a graph is live and the + // suspension isn't our own mute suspension (#9). const onStateChange = (): void => { - if (ctx.state === "suspended" && this.graph) + if (ctx.state === "suspended" && this.graph && !this.micMuted) void ctx.resume().catch(() => undefined); }; ctx.addEventListener("statechange", onStateChange); - - // Flat models register via addModule here; DTLN/DeepFilterNet bring their - // own processor via the dynamic-imported helper (see buildMlNode). - if (this.config.model === "rnnoise" || this.config.model === "speex") - await ctx.audioWorklet.addModule( - this.config.assetBase + FLAT[this.config.model].script, - ); - if (this.config.gate) - await ctx.audioWorklet.addModule(this.config.assetBase + GATE.script); // The action can arrive via host postMessage, not a gesture in this // iframe, so the context can start suspended — resume without hanging. - if (ctx.state === "suspended") await resumeCtx(ctx); + if (ctx.state === "suspended" && !this.micMuted) await resumeCtx(ctx); + // Attached while already muted (#9): don't let a prepared, running + // context burn inference until the first unmute. + else if (ctx.state === "running" && this.micMuted) + await ctx.suspend().catch(() => undefined); this.ctx = ctx; this.ctxStateHandler = onStateChange; @@ -260,7 +517,7 @@ export class LotusDenoiseProcessor implements TrackProcessor< private async buildGraph(track: MediaStreamTrack): Promise { const ctx = this.ctx!; const source = ctx.createMediaStreamSource(new MediaStream([track])); - const dest = ctx.createMediaStreamDestination(); + const dest = new MediaStreamAudioDestinationNode(ctx, MONO); const nodes: AudioNode[] = []; const disposes: (() => void)[] = []; @@ -277,6 +534,7 @@ export class LotusDenoiseProcessor implements TrackProcessor< // made the threshold operate on pre-denoise levels. Gate the residual. if (this.config.gate) { const gate = new AudioWorkletNode(ctx, GATE.name, { + ...MONO, processorOptions: { openThreshold: this.config.gateThreshold, closeThreshold: this.config.gateThreshold - 5, @@ -289,11 +547,11 @@ export class LotusDenoiseProcessor implements TrackProcessor< nodes.push(gate); } - // Only mix a dry floor for the LOW-LATENCY flat models (RNNoise/Speex). - // DTLN/DeepFilterNet add tens of ms of algorithmic latency, so summing an - // undelayed dry copy would comb-filter the voice — for those we rely on - // the model's own level (e.g. DFN noiseReductionLevel) instead. RNNoise is - // also where the "robotic/underwater" reports come from, so this targets it. + // Only mix a dry floor for the flat models (RNNoise/Speex), whose + // framing latency is known exactly (DRY_DELAY_SAMPLES); the DTLN/DFN + // figures are best-known estimates, so for those we rely on the model's + // own level (e.g. DFN noiseReductionLevel) instead. RNNoise is also where + // the "robotic/underwater" reports come from, so this targets it. const lowLatency = this.config.model === "rnnoise" || this.config.model === "speex"; const floor = lowLatency @@ -305,17 +563,25 @@ export class LotusDenoiseProcessor implements TrackProcessor< // (kills the "underwater"/pumping artifact). During speech (denoised ≈ // original) the two sum back to ~unity; in noise-only gaps the output // floors at `floor` × original instead of digital silence. - const wetGain = ctx.createGain(); - wetGain.gain.value = 1 - floor; + const wetGain = new GainNode(ctx, { ...MONO, gain: 1 - floor }); wetHead.connect(wetGain); wetGain.connect(dest); nodes.push(wetGain); - const dryGain = ctx.createGain(); - dryGain.gain.value = floor; - source.connect(dryGain); + // [lotus #25] Delay the dry copy by the model's algorithmic latency so + // it sums in phase with the (framed, hence delayed) wet path instead + // of comb-filtering against it. + const delaySec = DRY_DELAY_SAMPLES[this.config.model] / ctx.sampleRate; + const dryDelay = new DelayNode(ctx, { + ...MONO, + maxDelayTime: Math.max(delaySec, 1 / ctx.sampleRate), + delayTime: delaySec, + }); + const dryGain = new GainNode(ctx, { ...MONO, gain: floor }); + source.connect(dryDelay); + dryDelay.connect(dryGain); dryGain.connect(dest); - nodes.push(dryGain); + nodes.push(dryDelay, dryGain); } else { wetHead.connect(dest); } @@ -350,48 +616,36 @@ export class LotusDenoiseProcessor implements TrackProcessor< if (model === "dtln") { // Self-contained ESM that resolves its own processor + LiteRT wasm + // TFLite models. bypassUntilReady passes raw audio until the model loads. - const mod = await import( - /* @vite-ignore */ `${base}workadventure/audio-worklet.js` - ); - return (await mod.createNoiseSuppressionAudioWorklet(ctx, { + const mod = await loadDtlnModule(this.config); + return await mod.createNoiseSuppressionAudioWorklet(ctx, { bypassUntilReady: true, - })) as MlNode; + }); } if (model === "deepfilternet") { - const dfnBase = new URL(`${base}deepfilternet`, window.location.href) - .href; - const mod = await import( - /* @vite-ignore */ `${base}deepfilternet/index.esm.js` - ); - const core = new mod.DeepFilterNet3Core({ - sampleRate: 48_000, - // 60, not 80: full-strength suppression is the main source of the - // "over-processed" character; a lower level keeps voice natural while - // the dry/wet floor handles the noise tail. - noiseReductionLevel: 60, - assetConfig: { cdnUrl: dfnBase }, - }); - await core.initialize(); - const node = (await core.createAudioWorkletNode(ctx)) as AudioNode; + // [lotus #7] Use the core initialised by prepareDenoiseAssets() if we + // have one (first graph); later rebuilds (restart) load a fresh core. + const prepared = this.preparedDfnCore; + this.preparedDfnCore = undefined; + const core = prepared ?? (await loadDfnCore(this.config)); + const node = await core.createAudioWorkletNode(ctx); return { node, dispose: () => safeCall(() => core.destroy()) }; } // Flat sapphi worklet (rnnoise/speex). const flat = FLAT[model]; - const useSimd = model === "rnnoise" && !!flat.simdWasm && supportsSimd(); - const wasmFile = useSimd ? flat.simdWasm! : flat.wasm; + const { primary, fallback } = flatWasmFiles(model); let wasmBinary: ArrayBuffer; try { - wasmBinary = await fetchWasm(base + wasmFile); + wasmBinary = await fetchWasm(base + primary); } catch (e) { - if (useSimd) { - wasmCache.delete(base + wasmFile); - wasmBinary = await fetchWasm(base + flat.wasm); // fall back to non-SIMD + if (fallback) { + wasmCache.delete(base + primary); + wasmBinary = await fetchWasm(base + fallback); // fall back to non-SIMD } else throw e; } const node = new AudioWorkletNode(ctx, flat.name, { - channelCount: 1, + ...MONO, numberOfInputs: 1, numberOfOutputs: 1, processorOptions: { maxChannels: 1, wasmBinary },