From b2c88e43125739dca7d17d9fa8c2ba9aa591e7fc Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 4 May 2026 20:29:37 +0800 Subject: [PATCH] refactor(speech): add ONNX type augmentation and fix type-aware lint errors (#701) --- .../pocket/onnxruntime-augmentation.d.ts | 15 ++ .../providers/local/pocket/pocket-tts-onnx.ts | 197 +++++++++--------- .../providers/local/sherpa/model-catalog.ts | 28 ++- .../local/sherpa/model-downloader.ts | 22 +- 4 files changed, 147 insertions(+), 115 deletions(-) create mode 100644 packages/server/src/server/speech/providers/local/pocket/onnxruntime-augmentation.d.ts diff --git a/packages/server/src/server/speech/providers/local/pocket/onnxruntime-augmentation.d.ts b/packages/server/src/server/speech/providers/local/pocket/onnxruntime-augmentation.d.ts new file mode 100644 index 000000000..9169cf906 --- /dev/null +++ b/packages/server/src/server/speech/providers/local/pocket/onnxruntime-augmentation.d.ts @@ -0,0 +1,15 @@ +// Type augmentation for onnxruntime-node to include runtime properties +// not exposed in the official type definitions + +import type * as ort from "onnxruntime-node"; + +declare module "onnxruntime-node" { + interface InferenceSession { + /** Input tensor names (available at runtime but not in types) */ + readonly inputNames: string[]; + /** Output tensor names (available at runtime but not in types) */ + readonly outputNames?: string[]; + /** Input metadata for shape/type info (available at runtime but not in types) */ + readonly inputMetadata?: unknown; + } +} diff --git a/packages/server/src/server/speech/providers/local/pocket/pocket-tts-onnx.ts b/packages/server/src/server/speech/providers/local/pocket/pocket-tts-onnx.ts index ae4ab9bc0..635765120 100644 --- a/packages/server/src/server/speech/providers/local/pocket/pocket-tts-onnx.ts +++ b/packages/server/src/server/speech/providers/local/pocket/pocket-tts-onnx.ts @@ -16,6 +16,10 @@ type OrtModule = typeof import("onnxruntime-node"); type OrtSession = import("onnxruntime-node").InferenceSession; type OrtTensor = import("onnxruntime-node").Tensor; +function createSessionFeed(feed: Record): Record { + return feed; +} + interface SentencePieceProcessor { encodeIds: (text: string) => number[]; load?: (modelPath: string) => unknown; @@ -52,27 +56,13 @@ function getSessionInputMeta( session: OrtSession, inputName: string, ): { type?: string; dims?: Array } | undefined { - const metaAny = (session as unknown as { inputMetadata?: unknown }).inputMetadata; - if (Array.isArray(metaAny)) { - const entry = metaAny.find( - (m) => m && typeof m === "object" && (m as { name?: string }).name === inputName, - ) as { type?: string; shape?: Array } | undefined; - if (!entry) return undefined; - return { type: entry.type, dims: entry.shape }; + const metaAny = session.inputMetadata; + if (!Array.isArray(metaAny)) { + return undefined; } - - if (metaAny && typeof metaAny === "object" && inputName in metaAny) { - const entry = (metaAny as Record)[inputName] as - | { - type?: string; - dimensions?: Array; - shape?: Array; - } - | undefined; - return { type: entry?.type, dims: entry?.dimensions ?? entry?.shape }; - } - - return undefined; + const entry = metaAny.find((m) => m.name === inputName); + if (!entry) return undefined; + return { type: String(entry.type), dims: entry.shape }; } function toBigInt64(values: number[]): BigInt64Array { @@ -108,31 +98,54 @@ function normalizeTextForPocket(text: string): string { } async function loadOrt(): Promise { - return (await import("onnxruntime-node")) as OrtModule; + return import("onnxruntime-node"); +} + +interface SentencePieceModule { + SentencePieceProcessor?: new () => SentencePieceProcessor; + default?: + | (new () => SentencePieceProcessor) + | { SentencePieceProcessor?: new () => SentencePieceProcessor }; +} + +function isSentencePieceModule(mod: unknown): mod is SentencePieceModule { + return mod !== null && typeof mod === "object"; +} + +function getSentencePieceProcessor( + mod: SentencePieceModule, +): (new () => SentencePieceProcessor) | undefined { + if (mod.SentencePieceProcessor) { + return mod.SentencePieceProcessor; + } + const defaultValue = mod.default; + if ( + defaultValue && + typeof defaultValue === "object" && + "SentencePieceProcessor" in defaultValue + ) { + return defaultValue.SentencePieceProcessor; + } + if (typeof defaultValue === "function") { + return defaultValue as new () => SentencePieceProcessor; + } + return undefined; } async function loadSentencePiece(tokenizerModelPath: string): Promise { const mod = await import("@sctg/sentencepiece-js"); - const modRecord = mod as unknown as { - SentencePieceProcessor?: new () => SentencePieceProcessor; - default?: - | (new () => SentencePieceProcessor) - | { SentencePieceProcessor?: new () => SentencePieceProcessor }; - }; - const defaultValue = modRecord.default; - const Processor = - modRecord.SentencePieceProcessor ?? - (defaultValue && typeof defaultValue === "object" && "SentencePieceProcessor" in defaultValue - ? defaultValue.SentencePieceProcessor - : undefined) ?? - (typeof defaultValue === "function" ? defaultValue : undefined); + if (!isSentencePieceModule(mod)) { + throw new Error("@sctg/sentencepiece-js module has unexpected shape"); + } + + const Processor = getSentencePieceProcessor(mod); if (!Processor) { throw new Error("Failed to load SentencePiece processor from @sctg/sentencepiece-js"); } - const sp: SentencePieceProcessor = new Processor(); + const sp = new Processor(); if (typeof sp.load === "function") { await sp.load(tokenizerModelPath); @@ -181,8 +194,7 @@ function createZeroTensorForInput( function initState(session: OrtSession, ort: OrtModule): Record { const out: Record = {}; - const inputNames = (session as unknown as { inputNames: string[] }).inputNames; - for (const name of inputNames) { + for (const name of session.inputNames) { if (name.startsWith("state_")) { out[name] = createZeroTensorForInput(ort, session, name); } @@ -203,10 +215,15 @@ function updateStateFromOutputs( } } +interface OrtTensorWithData { + data: unknown; +} + function tensorDataFloat32(t: OrtTensor): Float32Array { - const data = (t as unknown as { data: unknown }).data; + const tensorWithData = t as OrtTensorWithData; + const data = tensorWithData.data; if (data instanceof Float32Array) return data; - if (Array.isArray(data)) return Float32Array.from(data as number[]); + if (Array.isArray(data)) return Float32Array.from(data); throw new Error("Unexpected tensor data type (expected Float32Array)"); } @@ -362,11 +379,9 @@ class PocketTtsOnnxEngine { const audioTensor = new ort.Tensor("float32", floatAudio, [1, 1, floatAudio.length]); const encoded = await mimiEncoder.run({ audio: audioTensor }); - const firstOutName = (mimiEncoder as unknown as { outputNames?: string[] }).outputNames?.[0]; - const encodedRecord = encoded as unknown as Record; - const voiceEmb = firstOutName - ? encodedRecord[firstOutName] - : (Object.values(encodedRecord)[0] as OrtTensor | undefined); + const firstOutName = mimiEncoder.outputNames?.[0]; + const encodedRecord = encoded as Record; + const voiceEmb = firstOutName ? encodedRecord[firstOutName] : Object.values(encodedRecord)[0]; if (!voiceEmb) { throw new Error("PocketTTS mimi_encoder: missing output"); } @@ -399,15 +414,11 @@ class PocketTtsOnnxEngine { } private async runTextConditioner(tokenIds: OrtTensor): Promise { - const out = await this.textConditioner.run({ - token_ids: tokenIds, - } as unknown as Record); - const firstOutName = (this.textConditioner as unknown as { outputNames?: string[] }) - .outputNames?.[0]; - const outRecord = out as unknown as Record; - const t = firstOutName - ? outRecord[firstOutName] - : (Object.values(outRecord)[0] as OrtTensor | undefined); + const feed: Record = { token_ids: tokenIds }; + const out = await this.textConditioner.run(feed); + const firstOutName = this.textConditioner.outputNames?.[0]; + const outRecord = out as Record; + const t = firstOutName ? outRecord[firstOutName] : Object.values(outRecord)[0]; if (!t) throw new Error("PocketTTS text_conditioner: missing output"); return t; } @@ -420,20 +431,24 @@ class PocketTtsOnnxEngine { const emptyText = new ort.Tensor("float32", new Float32Array(0), [1, 0, 1024]); // Voice conditioning pass - const resVoice = await this.flowLmMain.run({ - sequence: emptySeq, - text_embeddings: this.voiceEmbeddings, - ...state, - } as unknown as Record); - updateStateFromOutputs(state, resVoice as unknown as Record); + const resVoice = await this.flowLmMain.run( + createSessionFeed({ + sequence: emptySeq, + text_embeddings: this.voiceEmbeddings, + ...state, + }), + ); + updateStateFromOutputs(state, resVoice as Record); // Text conditioning pass - const resText = await this.flowLmMain.run({ - sequence: emptySeq, - text_embeddings: textEmbeddings, - ...state, - } as unknown as Record); - updateStateFromOutputs(state, resText as unknown as Record); + const resText = await this.flowLmMain.run( + createSessionFeed({ + sequence: emptySeq, + text_embeddings: textEmbeddings, + ...state, + }), + ); + updateStateFromOutputs(state, resText as Record); // Autoregressive generation const curr = new Float32Array(32); @@ -444,14 +459,16 @@ class PocketTtsOnnxEngine { let eosStep: number | null = null; for (let step = 0; step < this.maxFrames; step += 1) { - const resStep = await this.flowLmMain.run({ - sequence: currTensor, - text_embeddings: emptyText, - ...state, - } as unknown as Record); + const resStep = await this.flowLmMain.run( + createSessionFeed({ + sequence: currTensor, + text_embeddings: emptyText, + ...state, + }), + ); - const outputNames = (this.flowLmMain as unknown as { outputNames?: string[] }).outputNames; - const resStepRecord = resStep as unknown as Record; + const outputNames = this.flowLmMain.outputNames; + const resStepRecord = resStep as Record; const conditioningName = outputNames?.[0] ?? Object.keys(resStepRecord)[0]; const eosName = outputNames?.[1] ?? Object.keys(resStepRecord)[1]; @@ -481,17 +498,17 @@ class PocketTtsOnnxEngine { for (const st of this.stBuffers) { const xTensor = new ort.Tensor("float32", x, [1, 32]); - const flowOut = await this.flowLmFlow.run({ - c: conditioning, - s: st.s, - t: st.t, - x: xTensor, - } as unknown as Record); - const first = (this.flowLmFlow as unknown as { outputNames?: string[] }).outputNames?.[0]; - const flowOutRecord = flowOut as unknown as Record; - const flowTensor = first - ? flowOutRecord[first] - : (Object.values(flowOutRecord)[0] as OrtTensor | undefined); + const flowOut = await this.flowLmFlow.run( + createSessionFeed({ + c: conditioning, + s: st.s, + t: st.t, + x: xTensor, + }), + ); + const first = this.flowLmFlow.outputNames?.[0]; + const flowOutRecord = flowOut as Record; + const flowTensor = first ? flowOutRecord[first] : Object.values(flowOutRecord)[0]; if (!flowTensor) throw new Error("PocketTTS flow_lm_flow: missing output"); const delta = tensorDataFloat32(flowTensor); for (let i = 0; i < x.length; i += 1) { @@ -516,18 +533,12 @@ class PocketTtsOnnxEngine { } const latent = new ort.Tensor("float32", flattened, [1, frameCount, 32]); - const out = await this.mimiDecoder.run({ - latent, - ...state, - } as unknown as Record); - const outRecord = out as unknown as Record; + const out = await this.mimiDecoder.run(createSessionFeed({ latent, ...state })); + const outRecord = out as Record; updateStateFromOutputs(state, outRecord); - const firstOutName = (this.mimiDecoder as unknown as { outputNames?: string[] }) - .outputNames?.[0]; - const audioTensor = firstOutName - ? outRecord[firstOutName] - : (Object.values(outRecord)[0] as OrtTensor | undefined); + const firstOutName = this.mimiDecoder.outputNames?.[0]; + const audioTensor = firstOutName ? outRecord[firstOutName] : Object.values(outRecord)[0]; if (!audioTensor) { throw new Error("PocketTTS mimi_decoder: missing audio output"); } diff --git a/packages/server/src/server/speech/providers/local/sherpa/model-catalog.ts b/packages/server/src/server/speech/providers/local/sherpa/model-catalog.ts index b651b70bb..14eb4e47e 100644 --- a/packages/server/src/server/speech/providers/local/sherpa/model-catalog.ts +++ b/packages/server/src/server/speech/providers/local/sherpa/model-catalog.ts @@ -134,15 +134,21 @@ type ModelIdByKind = { export type LocalSttModelId = ModelIdByKind<"stt-online"> | ModelIdByKind<"stt-offline">; export type LocalTtsModelId = ModelIdByKind<"tts">; -const ALL_MODEL_IDS = Object.keys(SHERPA_ONNX_MODEL_CATALOG) as SherpaOnnxModelId[]; - -export const LOCAL_STT_MODEL_IDS = ALL_MODEL_IDS.filter( - (id): id is LocalSttModelId => SHERPA_ONNX_MODEL_CATALOG[id].kind !== "tts", +const ALL_MODEL_IDS: SherpaOnnxModelId[] = Object.keys(SHERPA_ONNX_MODEL_CATALOG).filter( + (k): k is SherpaOnnxModelId => k in SHERPA_ONNX_MODEL_CATALOG, ); -export const LOCAL_TTS_MODEL_IDS = ALL_MODEL_IDS.filter( - (id): id is LocalTtsModelId => SHERPA_ONNX_MODEL_CATALOG[id].kind === "tts", -); +function isLocalSttModelId(id: SherpaOnnxModelId): id is LocalSttModelId { + return SHERPA_ONNX_MODEL_CATALOG[id].kind !== "tts"; +} + +function isLocalTtsModelId(id: SherpaOnnxModelId): id is LocalTtsModelId { + return SHERPA_ONNX_MODEL_CATALOG[id].kind === "tts"; +} + +export const LOCAL_STT_MODEL_IDS: LocalSttModelId[] = ALL_MODEL_IDS.filter(isLocalSttModelId); + +export const LOCAL_TTS_MODEL_IDS: LocalTtsModelId[] = ALL_MODEL_IDS.filter(isLocalTtsModelId); function resolveDefaultModelId(role: "stt"): LocalSttModelId; function resolveDefaultModelId(role: "tts"): LocalTtsModelId; @@ -160,10 +166,10 @@ function resolveDefaultModelId(role: DefaultModelRole): SherpaOnnxModelId { export const DEFAULT_LOCAL_STT_MODEL = resolveDefaultModelId("stt"); export const DEFAULT_LOCAL_TTS_MODEL = resolveDefaultModelId("tts"); -function buildAliasMap(modelIds: readonly T[]): Record { +function buildAliasMap(modelIds: readonly T[]): Record { const aliasMap: Record = {}; for (const modelId of modelIds) { - const aliases = SHERPA_ONNX_MODEL_CATALOG[modelId as SherpaOnnxModelId].aliases ?? []; + const aliases = SHERPA_ONNX_MODEL_CATALOG[modelId].aliases ?? []; for (const alias of aliases) { aliasMap[alias.trim().toLowerCase()] = modelId; } @@ -175,14 +181,14 @@ function createAliasedModelIdSchema(params: { modelIds: readonly T[]; aliases: Record; }): z.ZodType { - const validIds = new Set(params.modelIds); + const validIds = new Set(params.modelIds); return z .string() .trim() .toLowerCase() .refine( (value): value is T => - validIds.has(value as T) || Object.prototype.hasOwnProperty.call(params.aliases, value), + validIds.has(value) || Object.prototype.hasOwnProperty.call(params.aliases, value), { message: "Invalid model id", }, diff --git a/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts b/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts index 2c9b6ae98..232c55be5 100644 --- a/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts +++ b/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts @@ -55,7 +55,9 @@ async function downloadToFile(options: DownloadToFileOptions): Promise { const tmpPath = `${outputPath}.tmp-${Date.now()}`; await mkdir(path.dirname(outputPath), { recursive: true }); - const nodeStream = Readable.fromWeb(res.body as Parameters[0]); + // The fetch ReadableStream type is slightly different from what Readable.fromWeb expects + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const nodeStream = Readable.fromWeb(res.body as any); try { await pipeline(nodeStream, createWriteStream(tmpPath)); @@ -207,18 +209,16 @@ export async function ensureSherpaOnnxModels(options: { logger: pino.Logger; }): Promise> { const uniq = Array.from(new Set(options.modelIds)); - const out: Partial> = {}; - const paths = await Promise.all( - uniq.map((id) => - ensureSherpaOnnxModel({ + const entries: Array<[SherpaOnnxModelId, string]> = await Promise.all( + uniq.map(async (id) => { + const modelPath = await ensureSherpaOnnxModel({ modelsDir: options.modelsDir, modelId: id, logger: options.logger, - }), - ), + }); + return [id, modelPath] as [SherpaOnnxModelId, string]; + }), ); - for (let i = 0; i < uniq.length; i += 1) { - out[uniq[i]] = paths[i]!; - } - return out as Record; + // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion + return Object.fromEntries(entries) as Record; }