diff --git a/.env.example b/.env.example index ec650c1..03d69a5 100644 --- a/.env.example +++ b/.env.example @@ -77,10 +77,7 @@ RUN_FS_MIGRATIONS= IMPORT_LIBRARY_DIR= IMPORT_LIBRARY_DIRS= -# (Required without Docker) Path to your local whisper.cpp CLI binary for STT timestamp generation -WHISPER_CPP_BIN=/whisper.cpp/build/bin/whisper-cli - -# Heavy compute backend mode for whisper alignment + PDF layout parsing. +# Heavy compute backend mode for ONNX whisper alignment + PDF layout parsing. # local = run compute in-process (default) # none = disable both capabilities (good for preview/serverless) # worker = reserved for future external worker mode (not implemented in v1) @@ -88,6 +85,22 @@ OPENREADER_COMPUTE_MODE=local # OPENREADER_COMPUTE_WORKER_URL= # OPENREADER_COMPUTE_WORKER_TOKEN= +# Optional overrides for Whisper ONNX artifacts +# Defaults target: onnx-community/whisper-base_timestamped int8 +# OPENREADER_WHISPER_MODEL_CONFIG_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/config.json +# OPENREADER_WHISPER_MODEL_GENERATION_CONFIG_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/generation_config.json +# OPENREADER_WHISPER_MODEL_TOKENIZER_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/tokenizer.json +# OPENREADER_WHISPER_MODEL_TOKENIZER_CONFIG_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/tokenizer_config.json +# OPENREADER_WHISPER_MODEL_MERGES_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/merges.txt +# OPENREADER_WHISPER_MODEL_VOCAB_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/vocab.json +# OPENREADER_WHISPER_MODEL_NORMALIZER_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/normalizer.json +# OPENREADER_WHISPER_MODEL_ADDED_TOKENS_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/added_tokens.json +# OPENREADER_WHISPER_MODEL_PREPROCESSOR_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/preprocessor_config.json +# OPENREADER_WHISPER_MODEL_SPECIAL_TOKENS_MAP_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/special_tokens_map.json +# OPENREADER_WHISPER_MODEL_ENCODER_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/onnx/encoder_model_int8.onnx +# OPENREADER_WHISPER_MODEL_DECODER_MERGED_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/onnx/decoder_model_merged_int8.onnx +# OPENREADER_WHISPER_MODEL_DECODER_WITH_PAST_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main/onnx/decoder_with_past_model_int8.onnx + # Optional overrides for PDF layout model artifacts # OPENREADER_PDF_LAYOUT_MODEL_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx # OPENREADER_PDF_LAYOUT_MODEL_DATA_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx.data diff --git a/Dockerfile b/Dockerfile index 764b387..0dcfe74 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,19 +1,4 @@ -# Stage 1: build whisper.cpp (no model download โ€“ the app handles that) -FROM alpine:3.23 AS whisper-builder - -RUN apk add --no-cache git cmake build-base - -WORKDIR /opt - -ARG TARGETARCH - -RUN git clone --depth 1 https://github.com/ggml-org/whisper.cpp.git && \ - cd whisper.cpp && \ - cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DGGML_NATIVE=OFF $( [ "$TARGETARCH" = "arm64" ] && echo "-DGGML_CPU_ARM_ARCH=armv8-a" || true ) && \ - cmake --build build -j -RUN wget -qO /tmp/whisper.cpp-LICENSE.txt "https://raw.githubusercontent.com/ggml-org/whisper.cpp/master/LICENSE" - -# Stage 1b: extract seaweedfs weed binary (for optional embedded weed mini) +# Stage 1: extract seaweedfs weed binary (for optional embedded weed mini) # Pin to 4.18 because CI observed upload regressions on 4.19. FROM chrislusf/seaweedfs:4.18 AS seaweedfs-builder RUN cp "$(command -v weed)" /tmp/weed && \ @@ -75,17 +60,12 @@ COPY --from=seaweedfs-builder /tmp/SeaweedFS-LICENSE.txt /licenses/SeaweedFS-LIC # Include static model notices for runtime-downloaded assets. COPY src/lib/server/pdf-layout/model/LICENSE.txt /licenses/pp-doclayoutv3-LICENSE.txt -# Copy the compiled whisper.cpp build output into the runtime image -# (includes whisper-cli and its shared libraries, e.g. libwhisper.so, libggml.so) -COPY --from=whisper-builder /opt/whisper.cpp/build /opt/whisper.cpp/build -COPY --from=whisper-builder /tmp/whisper.cpp-LICENSE.txt /licenses/whisper.cpp-LICENSE.txt # Copy seaweedfs weed binary for optional embedded local S3. COPY --from=seaweedfs-builder /tmp/weed /usr/local/bin/weed RUN chmod +x /usr/local/bin/weed -# Point the app at the compiled whisper-cli binary and ensure its libs are discoverable -ENV WHISPER_CPP_BIN=/opt/whisper.cpp/build/bin/whisper-cli -ENV LD_LIBRARY_PATH=/opt/whisper.cpp/build +# Include OpenAI Whisper license text for runtime-downloaded ONNX artifacts. +COPY src/lib/server/whisper/model/LICENSE.txt /licenses/openai-whisper-LICENSE.txt # Expose the port the app runs on EXPOSE 3003 diff --git a/README.md b/README.md index 210e215..63a99ec 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ OpenReader is an open source, self-host-friendly text-to-speech document reader - ๐ŸŽฏ **Multi-provider TTS** with OpenAI-compatible endpoints and cloud providers (Kokoro-FastAPI, KittenTTS-FastAPI, Orpheus-FastAPI or OpenAI, Replicate, DeepInfra). - ๐Ÿ“– **Read-along playback** for PDF/EPUB with sentence-aware narration. -- โฑ๏ธ **Word-by-word highlighting** via optional `whisper.cpp` timestamps (`OPENREADER_COMPUTE_MODE=local` + `WHISPER_CPP_BIN`). +- โฑ๏ธ **Word-by-word highlighting** via built-in ONNX Whisper alignment in local compute mode (`OPENREADER_COMPUTE_MODE=local`). - ๐Ÿงฑ **Layout-aware PDF parsing** (PP-DocLayoutV3 ONNX) with structured blocks for cleaner TTS/chaptering. - ๐Ÿ›œ **Sync + library import** to bring docs across devices and from server-mounted folders. - ๐Ÿ—‚๏ธ **Flexible storage** with embedded SeaweedFS or external S3-compatible backends. diff --git a/docs-site/docs/about/acknowledgements.md b/docs-site/docs/about/acknowledgements.md index 0a6c5ec..b47cccb 100644 --- a/docs-site/docs/about/acknowledgements.md +++ b/docs-site/docs/about/acknowledgements.md @@ -10,7 +10,7 @@ This project is built with support from the following open-source projects and t - [SQLite](https://www.sqlite.org/) - [PostgreSQL](https://www.postgresql.org/) - [SeaweedFS](https://github.com/seaweedfs/seaweedfs) -- [whisper.cpp](https://github.com/ggerganov/whisper.cpp) +- [OpenAI Whisper](https://github.com/openai/whisper) - [ffmpeg](https://ffmpeg.org) - [react-pdf](https://github.com/wojtekmaj/react-pdf) - [react-reader](https://github.com/happyr/react-reader) diff --git a/docs-site/docs/deploy/local-development.md b/docs-site/docs/deploy/local-development.md index 875c83f..1bc2183 100644 --- a/docs-site/docs/deploy/local-development.md +++ b/docs-site/docs/deploy/local-development.md @@ -114,50 +114,13 @@ sudo apt install -y libreoffice
-whisper.cpp (optional, for word-by-word highlighting) +Word-by-word highlighting (optional) -Install build dependencies: +No extra native Whisper CLI build step is required. - - +Set `OPENREADER_COMPUTE_MODE=local` to enable built-in ONNX word alignment in-process. -```bash -brew install cmake -``` - - - - -```bash -# Debian/Ubuntu example -sudo apt update -sudo apt install -y git build-essential cmake -``` - - - - -Build whisper.cpp: - -```bash -# clone and build whisper.cpp (no model download needed โ€“ OpenReader handles that) -git clone https://github.com/ggml-org/whisper.cpp.git -cd whisper.cpp -cmake -B build -cmake --build build -j --config Release - -# point OpenReader to the compiled whisper-cli binary -echo WHISPER_CPP_BIN="$(pwd)/build/bin/whisper-cli" -``` - -If you are not on Debian/Ubuntu, install equivalent packages with your distro package manager: - -- Fedora/RHEL: use `dnf` (`gcc gcc-c++ make cmake curl git tar xz`) -- Arch: use `pacman` (`base-devel cmake curl git tar xz`) - -:::tip -Set `OPENREADER_COMPUTE_MODE=local` and `WHISPER_CPP_BIN` in your `.env` to enable word-by-word highlighting. -::: +If you need mirrors or pinned artifact locations, set `OPENREADER_WHISPER_MODEL_*_URL` overrides in `.env`.
diff --git a/docs-site/docs/deploy/vercel-deployment.md b/docs-site/docs/deploy/vercel-deployment.md index 1f6f2e0..10fe73a 100644 --- a/docs-site/docs/deploy/vercel-deployment.md +++ b/docs-site/docs/deploy/vercel-deployment.md @@ -37,7 +37,7 @@ ADMIN_EMAILS=you@example.com # comma-separated; admins manage TTS + features in # Heavy compute (recommended on Vercel in v1) # local = requires native binaries/models in-process -# none = disable whisper alignment + PDF layout parsing +# none = disable ONNX whisper alignment + PDF layout parsing OPENREADER_COMPUTE_MODE=none # First-boot seed for the TTS shared provider (optional; manage in-app afterwards) @@ -95,8 +95,7 @@ Vercel deployments do not run `scripts/openreader-entrypoint.mjs`, so automatic - `/api/audiobook` - `/api/audiobook/chapter` -- `/api/audiobook/status` -- `/api/whisper` +- `/api/tts/segments/ensure` :::info `serverExternalPackages` should include `ffmpeg-static` so package paths resolve at runtime instead of being bundled into route output. @@ -113,7 +112,7 @@ FFmpeg workloads benefit from more memory/CPU. This repo includes: "$schema": "https://openapi.vercel.sh/vercel.json", "functions": { "app/api/audiobook/route.ts": { "memory": 3009 }, - "app/api/whisper/route.ts": { "memory": 3009 } + "app/api/tts/segments/ensure/route.ts": { "memory": 3009 } } } ``` @@ -130,4 +129,4 @@ Adjust memory per route if your files are larger or your plan differs. 1. Upload and read a PDF/EPUB document. 2. Confirm sync/blob fetch works across refreshes/devices. 3. Generate at least one audiobook chapter and play/download it. -4. If you later enable compute locally (`OPENREADER_COMPUTE_MODE=local`), verify word highlighting timestamps on a TTS run. +4. If you run with local compute (`OPENREADER_COMPUTE_MODE=local`) outside Vercel, verify word highlighting timestamps on a TTS run. diff --git a/docs-site/docs/introduction.md b/docs-site/docs/introduction.md index c4d09e0..53264c6 100644 --- a/docs-site/docs/introduction.md +++ b/docs-site/docs/introduction.md @@ -22,7 +22,7 @@ It supports multiple TTS providers including OpenAI, Replicate, DeepInfra, and c - [**DeepInfra**](https://deepinfra.com/models/text-to-speech): Kokoro-82M and other hosted models - [**OpenAI API**](https://platform.openai.com/docs/pricing#transcription-and-speech): `tts-1`, `tts-1-hd`, and `gpt-4o-mini-tts` - ๐Ÿ“– **Read Along Experience** - - Real-time highlighting for PDF/EPUB, with optional word-level [whisper.cpp](https://github.com/ggml-org/whisper.cpp) timestamps + - Real-time highlighting for PDF/EPUB, with built-in ONNX Whisper word-level timestamps in local compute mode - ๐Ÿ›œ **Document Storage** - Documents are persisted in server blob/object storage for consistent access - โšก **Segment-based TTS Playback** for reusable generation + preloading diff --git a/docs-site/docs/reference/environment-variables.md b/docs-site/docs/reference/environment-variables.md index f9d43cc..9c58793 100644 --- a/docs-site/docs/reference/environment-variables.md +++ b/docs-site/docs/reference/environment-variables.md @@ -53,14 +53,14 @@ For auth-enabled deployments, use **Settings โ†’ Admin** as the primary source o | `RUN_FS_MIGRATIONS` | Storage migrations | `true` | Set `false` to skip startup filesystem -> S3/DB migration pass | | `IMPORT_LIBRARY_DIR` | Library import | `docstore/library` fallback | Set a single server library root | | `IMPORT_LIBRARY_DIRS` | Library import | unset | Set multiple roots (comma/colon/semicolon separated) | -| `OPENREADER_COMPUTE_MODE` | Heavy compute backend | `local` | Set to `none` to disable whisper alignment + PDF layout parsing | +| `OPENREADER_COMPUTE_MODE` | Heavy compute backend | `local` | Set to `none` to disable ONNX word alignment + PDF layout parsing | | `OPENREADER_COMPUTE_WORKER_URL` | Heavy compute backend | unset | Reserved for future worker backend mode (`worker`) | | `OPENREADER_COMPUTE_WORKER_TOKEN` | Heavy compute backend | unset | Reserved for future worker backend mode (`worker`) | | `OPENREADER_PDF_LAYOUT_MODEL_URL` | PDF layout model | PP-DocLayoutV3 ONNX URL | Override ONNX model URL for `ensureModel()` | | `OPENREADER_PDF_LAYOUT_MODEL_DATA_URL` | PDF layout model | PP-DocLayoutV3 ONNX data URL | Override ONNX external data URL for `ensureModel()` | | `OPENREADER_PDF_LAYOUT_CONFIG_URL` | PDF layout model | PP-DocLayoutV3 config URL | Override model config URL for `ensureModel()` | | `OPENREADER_PDF_LAYOUT_PREPROCESSOR_URL` | PDF layout model | PP-DocLayoutV3 preprocessor URL | Override model preprocessor URL for `ensureModel()` | -| `WHISPER_CPP_BIN` | Word timing (local mode) | unset | Set to enable `whisper.cpp` timestamps in `OPENREADER_COMPUTE_MODE=local` | +| `OPENREADER_WHISPER_MODEL_*_URL` | Whisper ONNX model | onnx-community defaults | Optional per-artifact URL overrides for ONNX whisper-base_timestamped int8 downloads | | `FFMPEG_BIN` | Audio runtime | auto-detected (`ffmpeg-static`) | Override ffmpeg binary path | @@ -355,7 +355,7 @@ Multiple library roots for server library import. ### OPENREADER_COMPUTE_MODE -Selects the backend for heavy compute features (word alignment + PDF layout parsing). +Selects the backend for heavy compute features (ONNX word alignment + PDF layout parsing). - Default: `local` - Supported in v1: @@ -403,12 +403,29 @@ Override URL for the PP-DocLayoutV3 `preprocessor_config.json` downloaded by `en - Default: `https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/preprocessor_config.json` - You can pre-populate the model cache via `pnpm fetch-models` -### WHISPER_CPP_BIN +### OPENREADER_WHISPER_MODEL_*_URL -Absolute path to compiled `whisper.cpp` binary for word-level timestamps. +Optional per-artifact override URLs for the built-in ONNX Whisper alignment model downloader. -- Example: `/whisper.cpp/build/bin/whisper-cli` -- Required only for optional word-by-word highlighting +- Default base: `https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main` +- Default model variant: int8 (`encoder_model_int8.onnx`, `decoder_model_merged_int8.onnx`, `decoder_with_past_model_int8.onnx`) +- Use these when you need mirrors, pinned snapshots, or air-gapped fetch routing. + +Supported override vars: + +- `OPENREADER_WHISPER_MODEL_CONFIG_URL` +- `OPENREADER_WHISPER_MODEL_GENERATION_CONFIG_URL` +- `OPENREADER_WHISPER_MODEL_TOKENIZER_URL` +- `OPENREADER_WHISPER_MODEL_TOKENIZER_CONFIG_URL` +- `OPENREADER_WHISPER_MODEL_MERGES_URL` +- `OPENREADER_WHISPER_MODEL_VOCAB_URL` +- `OPENREADER_WHISPER_MODEL_NORMALIZER_URL` +- `OPENREADER_WHISPER_MODEL_ADDED_TOKENS_URL` +- `OPENREADER_WHISPER_MODEL_PREPROCESSOR_URL` +- `OPENREADER_WHISPER_MODEL_SPECIAL_TOKENS_MAP_URL` +- `OPENREADER_WHISPER_MODEL_ENCODER_URL` +- `OPENREADER_WHISPER_MODEL_DECODER_MERGED_URL` +- `OPENREADER_WHISPER_MODEL_DECODER_WITH_PAST_URL` ### FFMPEG_BIN diff --git a/docs-site/docs/reference/stack.md b/docs-site/docs/reference/stack.md index cc45fa2..b6dc596 100644 --- a/docs-site/docs/reference/stack.md +++ b/docs-site/docs/reference/stack.md @@ -34,7 +34,7 @@ title: Stack - App tables are manually maintained in Drizzle schema files - Auth tables are auto-generated by the [Better Auth](https://www.better-auth.com/) CLI and migrated alongside app tables via Drizzle - Blob storage: embedded [SeaweedFS](https://github.com/seaweedfs/seaweedfs) (`weed mini`) by default, or external S3-compatible storage via AWS SDK v3 -- Audio/processing pipeline: OpenAI-compatible TTS providers, [ffmpeg](https://ffmpeg.org/) for audiobook assembly, optional [whisper.cpp](https://github.com/ggerganov/whisper.cpp) for word timestamps +- Audio/processing pipeline: OpenAI-compatible TTS providers, [ffmpeg](https://ffmpeg.org/) for audiobook assembly, built-in ONNX Whisper (`onnx-community/whisper-base_timestamped` int8) for word timestamps ## Tooling and testing diff --git a/next.config.ts b/next.config.ts index 5179aeb..b35c2d0 100644 --- a/next.config.ts +++ b/next.config.ts @@ -14,7 +14,6 @@ const securityHeaders = [ value: 'max-age=63072000; includeSubDomains; preload', }, ]; - const nextConfig: NextConfig = { async headers() { return [ @@ -30,7 +29,13 @@ const nextConfig: NextConfig = { canvas: '@napi-rs/canvas', }, }, - serverExternalPackages: ["@napi-rs/canvas", "ffmpeg-static", "better-sqlite3"], + serverExternalPackages: [ + "@napi-rs/canvas", + "ffmpeg-static", + "better-sqlite3", + "onnxruntime-node", + "@huggingface/tokenizers", + ], outputFileTracingIncludes: { '/api/audiobook': [ './node_modules/ffmpeg-static/ffmpeg', @@ -38,7 +43,7 @@ const nextConfig: NextConfig = { '/api/audiobook/chapter': [ './node_modules/ffmpeg-static/ffmpeg', ], - '/api/whisper': [ + '/api/tts/segments/ensure': [ './node_modules/ffmpeg-static/ffmpeg', ], '/api/documents/blob/preview/ensure': [ diff --git a/package.json b/package.json index d7fec12..63197cf 100644 --- a/package.json +++ b/package.json @@ -27,6 +27,7 @@ "@aws-sdk/client-s3": "^3.1045.0", "@aws-sdk/s3-request-presigner": "^3.1045.0", "@headlessui/react": "^2.2.10", + "@huggingface/tokenizers": "^0.1.3", "@napi-rs/canvas": "^0.1.100", "@tanstack/react-query": "^5.100.10", "@types/archiver": "^7.0.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index c96e8fc..a2c265a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -22,6 +22,9 @@ importers: '@headlessui/react': specifier: ^2.2.10 version: 2.2.10(react-dom@19.2.6(react@19.2.6))(react@19.2.6) + '@huggingface/tokenizers': + specifier: ^0.1.3 + version: 0.1.3 '@napi-rs/canvas': specifier: ^0.1.100 version: 0.1.100 @@ -973,6 +976,9 @@ packages: react: ^18 || ^19 || ^19.0.0-rc react-dom: ^18 || ^19 || ^19.0.0-rc + '@huggingface/tokenizers@0.1.3': + resolution: {integrity: sha512-8rF/RRT10u+kn7YuUbUg0OF30K8rjTc78aHpxT+qJ1uWSqxT1MHi8+9ltwYfkFYJzT/oS+qw3JVfHtNMGAdqyA==} + '@humanfs/core@0.19.2': resolution: {integrity: sha512-UhXNm+CFMWcbChXywFwkmhqjs3PRCmcSa/hfBgLIb7oQ5HNb1wS0icWsGtSAUNgefHeI+eBrA8I1fxmbHsGdvA==} engines: {node: '>=18.18.0'} @@ -5210,6 +5216,8 @@ snapshots: react-dom: 19.2.6(react@19.2.6) use-sync-external-store: 1.6.0(react@19.2.6) + '@huggingface/tokenizers@0.1.3': {} + '@humanfs/core@0.19.2': dependencies: '@humanfs/types': 0.15.0 diff --git a/scripts/openreader-entrypoint.mjs b/scripts/openreader-entrypoint.mjs index 4909709..214201b 100644 --- a/scripts/openreader-entrypoint.mjs +++ b/scripts/openreader-entrypoint.mjs @@ -181,19 +181,20 @@ function spawnMainCommand(command, env) { const exitPromise = new Promise((resolve) => { child.on('error', (error) => { console.error('Failed to launch command:', error); - resolve(1); + resolve({ code: 1, signal: null, launchError: true }); }); child.on('exit', (code, signal) => { + console.error(`Main command exit event: code=${code ?? 'null'} signal=${signal ?? 'null'}.`); if (typeof code === 'number') { - resolve(code); + resolve({ code, signal: null, launchError: false }); return; } if (signal) { - resolve(1); + resolve({ code: 1, signal, launchError: false }); return; } - resolve(0); + resolve({ code: 0, signal: null, launchError: false }); }); }); @@ -421,7 +422,11 @@ async function main() { const { child, exitPromise } = spawnMainCommand(command, runtimeEnv); appProc = child; - const exitCode = await exitPromise; + const exitInfo = await exitPromise; + const exitCode = typeof exitInfo?.code === 'number' ? exitInfo.code : 1; + console.error( + `Main command finished with code=${exitInfo?.code ?? 'null'} signal=${exitInfo?.signal ?? 'null'} launchError=${Boolean(exitInfo?.launchError)}.`, + ); await shutdown('SIGTERM'); exitOnce(exitCode); diff --git a/src/app/api/whisper/route.ts b/src/app/api/whisper/route.ts deleted file mode 100644 index 5bfe152..0000000 --- a/src/app/api/whisper/route.ts +++ /dev/null @@ -1,47 +0,0 @@ -import { NextRequest, NextResponse } from 'next/server'; -import type { TTSSentenceAlignment } from '@/types/tts'; -import { auth } from '@/lib/server/auth/auth'; -import { makeWhisperCacheKey, type WhisperRequestBody } from '@/lib/server/whisper/alignment'; -import { getCompute } from '@/lib/server/compute'; - -export const runtime = 'nodejs'; - -export async function POST(req: NextRequest) { - try { - const session = await auth?.api.getSession({ headers: req.headers }); - if (auth && !session?.user) { - return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }); - } - - const body = (await req.json()) as WhisperRequestBody; - const { text, audio, lang } = body; - - if (!text || !audio || !Array.isArray(audio)) { - return NextResponse.json( - { error: 'Missing text or audio in request body' }, - { status: 400 } - ); - } - - const cacheKey = makeWhisperCacheKey(body); - const audioBuffer = new Uint8Array(audio).buffer; - - const alignments: TTSSentenceAlignment[] = (await getCompute().alignWords({ - audioBuffer, - text, - cacheKey, - lang, - })).alignments; - - return NextResponse.json({ alignments }, { status: 200 }); - } catch (error) { - console.error('Error in whisper route:', error); - return NextResponse.json( - { - error: 'WHISPER_ALIGNMENT_FAILED', - message: 'Failed to compute word-level alignment', - }, - { status: 500 } - ); - } -} diff --git a/src/lib/client/api/audiobooks.ts b/src/lib/client/api/audiobooks.ts index d737815..eb32a86 100644 --- a/src/lib/client/api/audiobooks.ts +++ b/src/lib/client/api/audiobooks.ts @@ -5,8 +5,6 @@ import type { AudiobookStatusResponse, CreateChapterPayload, VoicesResponse, - AlignmentPayload, - AlignmentResponse, TTSSegmentsEnsureRequest, TTSSegmentsEnsureResponse, } from '@/types/client'; @@ -208,23 +206,6 @@ export const getVoices = async (headers: HeadersInit): Promise = return await response.json(); }; -// --- Whisper API --- - - - -export const alignAudio = async (payload: AlignmentPayload): Promise => { - const response = await fetch('/api/whisper', { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify(payload), - }); - - if (!response.ok) return null; - return await response.json(); -}; - export const ensureTtsSegments = async ( payload: TTSSegmentsEnsureRequest, headers: TTSRequestHeaders, diff --git a/src/lib/server/compute/index.ts b/src/lib/server/compute/index.ts index f0700bb..b469a55 100644 --- a/src/lib/server/compute/index.ts +++ b/src/lib/server/compute/index.ts @@ -1,17 +1,12 @@ import type { ComputeBackend, ComputeMode } from '@/lib/server/compute/types'; import { LocalComputeBackend } from '@/lib/server/compute/local'; import { NoneComputeBackend } from '@/lib/server/compute/none'; +import { isComputeModeAvailable, readComputeMode } from '@/lib/server/compute/mode'; let backend: ComputeBackend | null = null; -function readMode(): ComputeMode { - const raw = (process.env.OPENREADER_COMPUTE_MODE || 'local').trim().toLowerCase(); - if (raw === 'local' || raw === 'none' || raw === 'worker') return raw; - return 'local'; -} - function createBackend(): ComputeBackend { - const mode = readMode(); + const mode: ComputeMode = readComputeMode(); if (mode === 'none') return new NoneComputeBackend(); if (mode === 'worker') { throw new Error( @@ -27,11 +22,5 @@ export function getCompute(): ComputeBackend { } export function isComputeAvailable(): boolean { - const mode = readMode(); - if (mode === 'worker') { - throw new Error( - 'OPENREADER_COMPUTE_MODE=worker is not implemented yet in v1. Switch to local/none or implement WorkerComputeBackend (v2).', - ); - } - return mode !== 'none'; + return isComputeModeAvailable(readComputeMode()); } diff --git a/src/lib/server/compute/local.ts b/src/lib/server/compute/local.ts index 2f95a98..a38bf6a 100644 --- a/src/lib/server/compute/local.ts +++ b/src/lib/server/compute/local.ts @@ -1,21 +1,21 @@ import type { ComputeBackend, PdfLayoutInput, WhisperAlignInput, WhisperAlignResult } from '@/lib/server/compute/types'; -import { alignAudioWithText } from '@/lib/server/whisper/alignment'; -import { parsePdf } from '@/lib/server/pdf-layout/parsePdf'; export class LocalComputeBackend implements ComputeBackend { readonly mode = 'local' as const; async alignWords(input: WhisperAlignInput): Promise { + const { alignAudioWithText } = await import('@/lib/server/whisper/alignment'); const alignments = await alignAudioWithText( input.audioBuffer, input.text, input.cacheKey, - { engine: 'whisper.cpp', lang: input.lang }, + { lang: input.lang }, ); return { alignments }; } async parsePdfLayout(input: PdfLayoutInput) { + const { parsePdf } = await import('@/lib/server/pdf-layout/parsePdf'); return parsePdf({ documentId: input.documentId, pdfBytes: input.pdfBytes }); } } diff --git a/src/lib/server/compute/mode.ts b/src/lib/server/compute/mode.ts new file mode 100644 index 0000000..170d6c9 --- /dev/null +++ b/src/lib/server/compute/mode.ts @@ -0,0 +1,16 @@ +import type { ComputeMode } from '@/lib/server/compute/types'; + +export function readComputeMode(): ComputeMode { + const raw = (process.env.OPENREADER_COMPUTE_MODE || 'local').trim().toLowerCase(); + if (raw === 'local' || raw === 'none' || raw === 'worker') return raw; + return 'local'; +} + +export function isComputeModeAvailable(mode: ComputeMode): boolean { + if (mode === 'worker') { + throw new Error( + 'OPENREADER_COMPUTE_MODE=worker is not implemented yet in v1. Switch to local/none or implement WorkerComputeBackend (v2).', + ); + } + return mode !== 'none'; +} diff --git a/src/lib/server/runtime-config.ts b/src/lib/server/runtime-config.ts index 87889d8..e1585c3 100644 --- a/src/lib/server/runtime-config.ts +++ b/src/lib/server/runtime-config.ts @@ -6,7 +6,7 @@ import { type RuntimeConfigKey, type RuntimeConfigSource, } from '@/lib/server/admin/settings'; -import { isComputeAvailable } from '@/lib/server/compute'; +import { isComputeModeAvailable, readComputeMode } from '@/lib/server/compute/mode'; export type ResolvedRuntimeConfig = RuntimeConfig & { computeAvailable: boolean; @@ -29,7 +29,7 @@ export async function getResolvedRuntimeConfig(): Promise const values = await getRuntimeConfig(); return { ...values, - computeAvailable: isComputeAvailable(), + computeAvailable: isComputeModeAvailable(readComputeMode()), }; } diff --git a/src/lib/server/whisper/alignment-mapping.ts b/src/lib/server/whisper/alignment-mapping.ts new file mode 100644 index 0000000..52c5707 --- /dev/null +++ b/src/lib/server/whisper/alignment-mapping.ts @@ -0,0 +1,46 @@ +import type { TTSSentenceAlignment, TTSSentenceWord } from '@/types/tts'; +import { preprocessSentenceForAudio } from '@/lib/shared/nlp'; + +export interface WhisperWord { + start: number; + end: number; + word: string; +} + +export function mapWordsToSentenceOffsets(sentence: string, words: WhisperWord[]): TTSSentenceAlignment { + const normalizedSentence = preprocessSentenceForAudio(sentence); + const lowerSentence = normalizedSentence.toLowerCase(); + let cursor = 0; + + const alignedWords: TTSSentenceWord[] = words.map((w) => { + const token = w.word.trim(); + if (!token) { + return { + text: '', + startSec: w.start, + endSec: w.end, + charStart: cursor, + charEnd: cursor, + }; + } + + const idx = lowerSentence.indexOf(token.toLowerCase(), cursor); + const start = idx >= 0 ? idx : cursor; + const end = Math.min(normalizedSentence.length, start + token.length); + cursor = Math.max(cursor, end); + + return { + text: token, + startSec: w.start, + endSec: w.end, + charStart: start, + charEnd: end, + }; + }).filter((word) => word.text.length > 0); + + return { + sentence, + sentenceIndex: 0, + words: alignedWords, + }; +} diff --git a/src/lib/server/whisper/alignment.ts b/src/lib/server/whisper/alignment.ts index 488b3c5..d0d8f14 100644 --- a/src/lib/server/whisper/alignment.ts +++ b/src/lib/server/whisper/alignment.ts @@ -1,21 +1,36 @@ import { createHash, randomUUID } from 'crypto'; -import { mkdtemp, writeFile, rm, access, mkdir, readFile } from 'fs/promises'; +import { mkdtemp, readFile, rm, writeFile } from 'fs/promises'; import { tmpdir } from 'os'; import { join } from 'path'; import { spawn } from 'child_process'; -import type { TTSSentenceAlignment, TTSAudioBytes, TTSAudioBuffer } from '@/types/tts'; -import { preprocessSentenceForAudio } from '@/lib/shared/nlp'; +import * as ort from 'onnxruntime-node'; +import { Tokenizer } from '@huggingface/tokenizers'; +import JSZip from 'jszip'; +import type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment } from '@/types/tts'; import { getFFmpegPath } from '@/lib/server/audiobooks/ffmpeg-bin'; +import { + mapWordsToSentenceOffsets, + type WhisperWord, +} from '@/lib/server/whisper/alignment-mapping'; +import { buildGoertzelCoefficients, goertzelPower } from '@/lib/server/whisper/spectral'; +import { + buildWordsFromTimestampedTokens, + extractTokenStartTimestamps, +} from '@/lib/server/whisper/token-timestamps'; +import { + ensureWhisperModel, + WHISPER_CONFIG_PATH, + WHISPER_GENERATION_CONFIG_PATH, + WHISPER_TOKENIZER_CONFIG_PATH, + WHISPER_TOKENIZER_PATH, + WHISPER_ENCODER_MODEL_PATH, + WHISPER_DECODER_MERGED_MODEL_PATH, + WHISPER_DECODER_WITH_PAST_MODEL_PATH, +} from '@/lib/server/whisper/ensureModel'; interface WhisperAlignmentOptions { - engine?: 'whisper.cpp'; lang?: string; -} - -interface WhisperWord { - start: number; - end: number; - word: string; + textHint?: string; } export interface WhisperRequestBody { @@ -24,362 +39,962 @@ export interface WhisperRequestBody { lang?: string; } -const alignmentCache = new Map(); - -const MODEL_NAME = 'ggml-tiny.en.bin'; -const MODEL_URL = - 'https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en.bin'; -const DOCSTORE_DIR = join(process.cwd(), 'docstore'); -const MODEL_DIR = join(DOCSTORE_DIR, 'model'); -const MODEL_PATH = join(MODEL_DIR, MODEL_NAME); -const modelReadyPromises = new Map>(); - -async function ensureModelAvailable(): Promise { - try { - await access(MODEL_PATH); - return; - } catch { - // continue - } - - const existing = modelReadyPromises.get(MODEL_PATH); - if (existing) return existing; - - const promise = (async () => { - try { - await access(MODEL_PATH); - return; - } catch { - // still missing - } - - await mkdir(MODEL_DIR, { recursive: true }); - - const res = await fetch(MODEL_URL); - if (!res.ok) { - throw new Error( - `Failed to download Whisper model from ${MODEL_URL}: ${res.status} ${res.statusText}` - ); - } - - const arrayBuffer = await res.arrayBuffer(); - await writeFile(MODEL_PATH, Buffer.from(arrayBuffer)); - })(); - - modelReadyPromises.set(MODEL_PATH, promise); - return promise; +interface WhisperRuntime { + encoder: ort.InferenceSession; + decoderMerged: ort.InferenceSession; + decoderWithPast: ort.InferenceSession; + tokenizer: Tokenizer; + promptStartToken: number; + defaultLanguageToken: number; + transcribeToken: number; + eosTokenId: number; + noTimestampsTokenId: number; + timestampBeginTokenId: number; + maxInitialTimestampIndex: number; + maxDecodeSteps: number; + suppressTokens: Set; + beginSuppressTokens: Set; + alignmentHeads: Array<[number, number]>; + prefillFetches: string[]; + stepFetches: string[]; } -async function runWhisperCpp( - wavPath: string, - opts: WhisperAlignmentOptions -): Promise { - const binary = process.env.WHISPER_CPP_BIN; - if (!binary) { - throw new Error( - 'Whisper.cpp binary path not configured. Set WHISPER_CPP_BIN to the compiled binary.' - ); +type WhisperAlignmentState = { + alignmentCache: Map; + alignmentInFlight: Map>; + runtimePromise: Promise | null; + alignMutex: Promise; + pendingAlignments: number; + officialMelFilters: Float32Array[] | null; + emptyPastFeedsTemplate: Record | null; +}; + +const WHISPER_ALIGNMENT_STATE_KEY = '__openreaderWhisperAlignmentStateV1'; +const g = globalThis as typeof globalThis & Record; +const state = (() => { + const existing = g[WHISPER_ALIGNMENT_STATE_KEY] as WhisperAlignmentState | undefined; + if (existing) return existing; + const created: WhisperAlignmentState = { + alignmentCache: new Map(), + alignmentInFlight: new Map>(), + runtimePromise: null, + alignMutex: Promise.resolve(), + pendingAlignments: 0, + officialMelFilters: null, + emptyPastFeedsTemplate: null, + }; + g[WHISPER_ALIGNMENT_STATE_KEY] = created; + return created; +})(); +const alignmentCache = state.alignmentCache; +const alignmentInFlight = state.alignmentInFlight; +const ALIGNMENT_CACHE_MAX_ENTRIES = 256; +const MAX_DECODE_STEPS_CAP = 128; +const ALIGNMENT_TIMEOUT_MS = 25000; +const FFMPEG_DECODE_TIMEOUT_MS = 10000; + +const SAMPLE_RATE = 16000; +const N_FFT = 400; +const HOP_LENGTH = 160; +const CHUNK_LENGTH_SECONDS = 30; +const N_SAMPLES = CHUNK_LENGTH_SECONDS * SAMPLE_RATE; +const N_FRAMES = N_SAMPLES / HOP_LENGTH; +const N_MELS = 80; +const WHISPER_NUM_HEADS = 8; +const WHISPER_HEAD_DIM = 64; +const WHISPER_NUM_LAYERS = 6; +const MEL_FILTER_BINS = (N_FFT / 2) + 1; + +const hannWindow = buildHannWindow(N_FFT); +const goertzelCoefficients = buildGoertzelCoefficients(MEL_FILTER_BINS, N_FFT); + +const MEL_FILTERS_NPZ_PATH = join(process.cwd(), 'src/lib/server/whisper/model/mel_filters.npz'); + +function buildHannWindow(length: number): Float32Array { + const window = new Float32Array(length); + for (let i = 0; i < length; i += 1) { + window[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / length); + } + return window; +} + +function parseNpyFloat32(bytes: Uint8Array): { shape: number[]; data: Float32Array } { + if (bytes.length < 12) { + throw new Error('Invalid NPY payload: too short'); + } + const magic = String.fromCharCode(...bytes.slice(0, 6)); + if (magic !== '\u0093NUMPY') { + throw new Error('Invalid NPY payload: missing magic header'); } - await ensureModelAvailable(); + const major = bytes[6]; + const headerLength = major <= 1 + ? new DataView(bytes.buffer, bytes.byteOffset + 8, 2).getUint16(0, true) + : new DataView(bytes.buffer, bytes.byteOffset + 8, 4).getUint32(0, true); + const headerOffset = major <= 1 ? 10 : 12; + const header = Buffer.from(bytes.slice(headerOffset, headerOffset + headerLength)).toString('latin1'); - return new Promise((resolve, reject) => { - const jsonBase = `${wavPath}.json_out`; - const jsonPath = `${jsonBase}.json`; - const args = [ - '-m', - MODEL_PATH, - '-f', - wavPath, - '-of', - jsonBase, - '-ojf', - '-np', - ]; + const descrMatch = header.match(/'descr':\s*'([^']+)'/); + if (!descrMatch || descrMatch[1] !== ' token.trim()) + .filter(Boolean) + .map((token) => Number(token)) + .filter((n) => Number.isFinite(n) && n > 0); + + const dataOffset = headerOffset + headerLength; + const dataBytes = bytes.slice(dataOffset); + const totalFloats = Math.floor(dataBytes.byteLength / 4); + const data = new Float32Array(totalFloats); + const view = new DataView(dataBytes.buffer, dataBytes.byteOffset, dataBytes.byteLength); + for (let i = 0; i < totalFloats; i += 1) { + data[i] = view.getFloat32(i * 4, true); + } + + return { shape, data }; +} + +async function loadOfficialMelFilters(): Promise { + if (state.officialMelFilters) return state.officialMelFilters; + + const npzBytes = await readFile(MEL_FILTERS_NPZ_PATH); + const zip = await JSZip.loadAsync(npzBytes); + const mel80 = zip.file('mel_80.npy'); + if (!mel80) { + throw new Error('OpenAI mel filter asset is missing mel_80.npy'); + } + + const raw = await mel80.async('uint8array'); + const parsed = parseNpyFloat32(raw); + const [rows, cols] = parsed.shape; + if (rows !== N_MELS || cols !== MEL_FILTER_BINS) { + throw new Error(`Unexpected mel filter shape: [${rows}, ${cols}]`); + } + + const filters: Float32Array[] = []; + for (let row = 0; row < rows; row += 1) { + const start = row * cols; + filters.push(parsed.data.slice(start, start + cols)); + } + + state.officialMelFilters = filters; + return filters; +} + +function pcm16ToFloat32(buffer: Buffer): Float32Array { + const view = new Int16Array(buffer.buffer, buffer.byteOffset, Math.floor(buffer.byteLength / 2)); + const out = new Float32Array(view.length); + for (let i = 0; i < view.length; i += 1) { + out[i] = view[i] / 32768; + } + return out; +} + +function padOrTrimAudio(samples: Float32Array): Float32Array { + if (samples.length === N_SAMPLES) return samples; + if (samples.length > N_SAMPLES) return samples.subarray(0, N_SAMPLES); + + const padded = new Float32Array(N_SAMPLES); + padded.set(samples, 0); + return padded; +} + +function reflectPad(audio: Float32Array, pad: number): Float32Array { + const out = new Float32Array(audio.length + (2 * pad)); + out.set(audio, pad); + + // Match PyTorch reflect padding (exclude edge sample). + for (let i = 0; i < pad; i += 1) { + out[pad - 1 - i] = audio[Math.min(audio.length - 1, i + 1)]; + out[pad + audio.length + i] = audio[Math.max(0, audio.length - 2 - i)]; + } + + return out; +} + +function computeLogMelSpectrogram(audioSamples: Float32Array): ort.Tensor { + if (!state.officialMelFilters) { + throw new Error('Whisper mel filters not loaded'); + } + + const paddedAudio = reflectPad(audioSamples, N_FFT / 2); + const stftFrames = N_FRAMES + 1; + const frameCount = N_FRAMES; + const freqBins = MEL_FILTER_BINS; + + const melSpec = Array.from({ length: N_MELS }, () => new Float32Array(frameCount)); + const frame = new Float32Array(N_FFT); + const power = new Float32Array(freqBins); + + for (let frameIndex = 0; frameIndex < stftFrames; frameIndex += 1) { + const offset = frameIndex * HOP_LENGTH; + + for (let i = 0; i < N_FFT; i += 1) { + frame[i] = (paddedAudio[offset + i] ?? 0) * hannWindow[i]; } - const child = spawn(binary, args); + for (let k = 0; k < freqBins; k += 1) { + power[k] = goertzelPower(frame, goertzelCoefficients[k]); + } + + if (frameIndex === stftFrames - 1) { + continue; + } + + for (let melIndex = 0; melIndex < N_MELS; melIndex += 1) { + const filter = state.officialMelFilters[melIndex]; + let total = 0; + for (let k = 0; k < freqBins; k += 1) { + total += filter[k] * power[k]; + } + melSpec[melIndex][frameIndex] = total; + } + } + + // Whisper normalization from openai/whisper/audio.py + let globalMaxLog = Number.NEGATIVE_INFINITY; + for (let i = 0; i < N_MELS; i += 1) { + for (let j = 0; j < frameCount; j += 1) { + const logVal = Math.log10(Math.max(1e-10, melSpec[i][j])); + if (logVal > globalMaxLog) globalMaxLog = logVal; + melSpec[i][j] = logVal; + } + } + + const floorVal = globalMaxLog - 8.0; + const flattened = new Float32Array(1 * N_MELS * frameCount); + for (let i = 0; i < N_MELS; i += 1) { + for (let j = 0; j < frameCount; j += 1) { + const clamped = Math.max(melSpec[i][j], floorVal); + flattened[(i * frameCount) + j] = (clamped + 4.0) / 4.0; + } + } + + return new ort.Tensor('float32', flattened, [1, N_MELS, frameCount]); +} + +async function decodeToPcm16(inputPath: string, outputPath: string): Promise { + await new Promise((resolve, reject) => { + const ffmpeg = spawn(getFFmpegPath(), [ + '-y', + '-i', + inputPath, + '-f', + 's16le', + '-ar', + String(SAMPLE_RATE), + '-ac', + '1', + outputPath, + ]); - let stdout = ''; let stderr = ''; - - child.stdout.on('data', (data) => { - stdout += data.toString(); - }); - - child.stderr.on('data', (data) => { + let timedOut = false; + const timer = setTimeout(() => { + timedOut = true; + ffmpeg.kill('SIGKILL'); + }, FFMPEG_DECODE_TIMEOUT_MS); + ffmpeg.stderr.on('data', (data) => { stderr += data.toString(); }); - child.on('error', (err) => { + ffmpeg.on('error', (err) => { + clearTimeout(timer); reject(err); }); - child.on('close', (code) => { - if (code !== 0) { - return reject( - new Error( - `whisper.cpp exited with code ${code}: ${stderr || stdout}` - ) - ); + ffmpeg.on('close', (code) => { + clearTimeout(timer); + if (timedOut) { + reject(new Error(`ffmpeg decode timed out after ${FFMPEG_DECODE_TIMEOUT_MS}ms`)); + return; + } + if (code === 0) { + resolve(); + } else { + reject(new Error(`ffmpeg decode failed with code ${code}: ${stderr}`)); } - - readFile(jsonPath, 'utf-8') - .then((content: string) => { - const words: WhisperWord[] = []; - const parsed = JSON.parse(content) as { - transcription?: Array<{ - text?: string; - timestamps?: { from?: string; to?: string }; - offsets?: { from?: number; to?: number }; - tokens?: Array<{ - text?: string; - timestamps?: { from?: string; to?: string }; - offsets?: { from?: number; to?: number }; - }>; - }>; - }; - - const transcription = parsed.transcription; - - const parseTimecode = (value?: string): number | null => { - if (!value) return null; - const m = value.match(/(\d+):(\d+):(\d+),(\d+)/); - if (!m) return null; - const h = Number(m[1]); - const min = Number(m[2]); - const s = Number(m[3]); - const ms = Number(m[4]); - if ( - Number.isNaN(h) - || Number.isNaN(min) - || Number.isNaN(s) - || Number.isNaN(ms) - ) { - return null; - } - return h * 3600 + min * 60 + s + ms / 1000; - }; - - if (Array.isArray(transcription)) { - for (const seg of transcription) { - const segText = (seg.text || '').trim(); - const segStartSecFromTs = parseTimecode( - seg.timestamps?.from - ); - const segEndSecFromTs = parseTimecode(seg.timestamps?.to); - const segStartSecFromMs = - typeof seg.offsets?.from === 'number' - ? seg.offsets.from / 1000 - : null; - const segEndSecFromMs = - typeof seg.offsets?.to === 'number' - ? seg.offsets.to / 1000 - : null; - - const segStartSec = - segStartSecFromTs - ?? segStartSecFromMs - ?? 0; - const segEndSec = - segEndSecFromTs - ?? segEndSecFromMs - ?? segStartSec; - - const tokens = Array.isArray(seg.tokens) - ? seg.tokens - : []; - - if (tokens.length > 0) { - for (const token of tokens) { - const rawText = token.text || ''; - const tokenText = rawText.trim(); - if (!tokenText || /^\[.*\]$/.test(tokenText)) continue; - - const tokStartSecFromTs = parseTimecode( - token.timestamps?.from - ); - const tokEndSecFromTs = parseTimecode( - token.timestamps?.to - ); - const tokStartSecFromMs = - typeof token.offsets?.from === 'number' - ? token.offsets.from / 1000 - : null; - const tokEndSecFromMs = - typeof token.offsets?.to === 'number' - ? token.offsets.to / 1000 - : null; - - const startSec = - tokStartSecFromTs - ?? tokStartSecFromMs - ?? segStartSec; - const endSec = - tokEndSecFromTs - ?? tokEndSecFromMs - ?? segEndSec; - - words.push({ - word: tokenText, - start: startSec, - end: endSec, - }); - } - } else if (segText) { - const segTokens = segText.split(/\s+/).filter(Boolean); - if (segTokens.length) { - const totalDur = Math.max(segEndSec - segStartSec, 0); - const step = - segTokens.length > 0 - ? totalDur / segTokens.length - : 0; - segTokens.forEach((token, index) => { - const wStart = - step > 0 - ? segStartSec + step * index - : segStartSec; - const wEnd = - step > 0 - ? index === segTokens.length - 1 - ? segEndSec - : segStartSec + step * (index + 1) - : segEndSec; - words.push({ - word: token, - start: wStart, - end: wEnd, - }); - }); - } - } - } - } - - resolve(words); - }) - .catch((err: unknown) => { - reject(err); - }); }); }); } -function mapWordsToSentenceOffsets( - sentence: string, - words: WhisperWord[] -): TTSSentenceAlignment { - const normalizedSentence = preprocessSentenceForAudio(sentence); - let cursor = 0; +function parseLanguageCode(lang?: string): string | null { + if (!lang) return null; + const trimmed = lang.trim().toLowerCase(); + if (!trimmed) return null; + if (trimmed.includes('-')) return trimmed.split('-')[0] || null; + if (trimmed.includes('_')) return trimmed.split('_')[0] || null; + return trimmed; +} - const alignedWords = words.map((w) => { - const token = w.word.trim(); - if (!token) { - return { - text: '', - startSec: w.start, - endSec: w.end, - charStart: cursor, - charEnd: cursor, - }; +function tensorFromInt64(values: number[]): ort.Tensor { + return new ort.Tensor('int64', BigInt64Array.from(values.map((v) => BigInt(v))), [1, values.length]); +} + +function disposeTensor(tensor: ort.Tensor | undefined | null): void { + if (!tensor) return; + try { + tensor.dispose(); + } catch { + // Best-effort cleanup: ignore disposal errors during fallback path. + } +} + +function disposeTensorMap(tensors: Record): void { + for (const tensor of Object.values(tensors)) { + disposeTensor(tensor); + } +} + +function computeAdaptiveDecodeStepLimit(maxDecodeSteps: number, textHint?: string): number { + const normalized = (textHint ?? '').trim(); + if (!normalized) return Math.min(maxDecodeSteps, 96); + + const chars = normalized.length; + const words = normalized.split(/\s+/).filter(Boolean).length; + const estTokens = Math.max(words * 3, Math.ceil(chars / 2)); + const adaptive = Math.max(64, Math.min(maxDecodeSteps, estTokens + 24)); + return adaptive; +} + +function assertWithinDeadline(deadlineMs: number): void { + if (Date.now() > deadlineMs) { + throw new Error(`Whisper alignment timed out after ${ALIGNMENT_TIMEOUT_MS}ms`); + } +} + +function makeInFlightCoalesceKey(audioBuffer: TTSAudioBuffer, text: string, lang?: string): string { + const bytes = new Uint8Array(audioBuffer); + const span = 4096; + const head = bytes.subarray(0, Math.min(span, bytes.length)); + const tailStart = Math.max(0, bytes.length - span); + const tail = bytes.subarray(tailStart); + return createHash('sha256') + .update(text) + .update('\0') + .update(lang ?? '') + .update('\0') + .update(String(bytes.length)) + .update('\0') + .update(head) + .update('\0') + .update(tail) + .digest('hex'); +} + +function buildEmptyPastFeeds() { + if (state.emptyPastFeedsTemplate) return state.emptyPastFeedsTemplate; + + const feeds: Record = {}; + const emptyDecoderPast = new Float32Array(0); + const emptyEncoderPast = new Float32Array(1 * WHISPER_NUM_HEADS * 1500 * WHISPER_HEAD_DIM); + + for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { + feeds[`past_key_values.${i}.decoder.key`] = new ort.Tensor('float32', emptyDecoderPast, [1, WHISPER_NUM_HEADS, 0, WHISPER_HEAD_DIM]); + feeds[`past_key_values.${i}.decoder.value`] = new ort.Tensor('float32', emptyDecoderPast, [1, WHISPER_NUM_HEADS, 0, WHISPER_HEAD_DIM]); + + // First pass still expects encoder KV inputs in the merged decoder graph. + feeds[`past_key_values.${i}.encoder.key`] = new ort.Tensor('float32', emptyEncoderPast, [1, WHISPER_NUM_HEADS, 1500, WHISPER_HEAD_DIM]); + feeds[`past_key_values.${i}.encoder.value`] = new ort.Tensor('float32', emptyEncoderPast, [1, WHISPER_NUM_HEADS, 1500, WHISPER_HEAD_DIM]); + } + + state.emptyPastFeedsTemplate = feeds; + return state.emptyPastFeedsTemplate; +} + +function argmax(values: Float32Array): number | null { + let bestIdx = 0; + let bestScore = Number.NEGATIVE_INFINITY; + + for (let i = 0; i < values.length; i += 1) { + const score = values[i]; + if (score > bestScore) { + bestScore = score; + bestIdx = i; + } + } + + return Number.isFinite(bestScore) ? bestIdx : null; +} + +function applyTokenSuppression(logits: Float32Array, tokens: Set) { + for (const tokenId of tokens) { + if (tokenId >= 0 && tokenId < logits.length) { + logits[tokenId] = Number.NEGATIVE_INFINITY; + } + } +} + +function logSoftmax(input: Float32Array): Float32Array { + let max = Number.NEGATIVE_INFINITY; + for (let i = 0; i < input.length; i += 1) { + if (input[i] > max) max = input[i]; + } + if (!Number.isFinite(max)) { + return new Float32Array(input.length).fill(Number.NEGATIVE_INFINITY); + } + + let sum = 0; + for (let i = 0; i < input.length; i += 1) { + sum += Math.exp(input[i] - max); + } + const logSum = Math.log(sum); + + const out = new Float32Array(input.length); + for (let i = 0; i < input.length; i += 1) { + out[i] = input[i] - max - logSum; + } + return out; +} + +function applyWhisperTimestampLogitsRules(input: { + logits: Float32Array; + generated: number[]; + beginIndex: number; + eosTokenId: number; + noTimestampsTokenId: number; + timestampBeginTokenId: number; + maxInitialTimestampIndex: number; +}) { + const { + logits, + generated, + beginIndex, + eosTokenId, + noTimestampsTokenId, + timestampBeginTokenId, + maxInitialTimestampIndex, + } = input; + + if (noTimestampsTokenId >= 0 && noTimestampsTokenId < logits.length) { + logits[noTimestampsTokenId] = Number.NEGATIVE_INFINITY; + } + + if (generated.length === beginIndex) { + const upper = Math.min(timestampBeginTokenId, logits.length); + for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + + const seq = generated.slice(beginIndex); + const lastWasTimestamp = seq.length >= 1 && seq[seq.length - 1] >= timestampBeginTokenId; + const penultimateWasTimestamp = seq.length < 2 || seq[seq.length - 2] >= timestampBeginTokenId; + + if (lastWasTimestamp) { + if (penultimateWasTimestamp) { + for (let i = timestampBeginTokenId; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } else { + const upper = Math.min(eosTokenId, logits.length); + for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + } + + if (generated.length === beginIndex && Number.isFinite(maxInitialTimestampIndex)) { + const lastAllowed = timestampBeginTokenId + maxInitialTimestampIndex; + for (let i = lastAllowed + 1; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + + const textUpper = Math.min(timestampBeginTokenId, logits.length); + if (textUpper <= 0 || textUpper >= logits.length) return; + + const logprobs = logSoftmax(logits); + + let maxTextTokenLogprob = Number.NEGATIVE_INFINITY; + for (let i = 0; i < textUpper; i += 1) { + if (logprobs[i] > maxTextTokenLogprob) maxTextTokenLogprob = logprobs[i]; + } + + let timestampProbMass = 0; + for (let i = textUpper; i < logprobs.length; i += 1) { + timestampProbMass += Math.exp(logprobs[i]); + } + const timestampLogprob = timestampProbMass > 0 ? Math.log(timestampProbMass) : Number.NEGATIVE_INFINITY; + + if (timestampLogprob > maxTextTokenLogprob) { + for (let i = 0; i < textUpper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } +} + +async function getRuntime(): Promise { + if (state.runtimePromise) return state.runtimePromise; + + state.runtimePromise = (async () => { + await ensureWhisperModel(); + await loadOfficialMelFilters(); + + const [configRaw, generationRaw, tokenizerJsonRaw, tokenizerConfigRaw] = await Promise.all([ + readFile(WHISPER_CONFIG_PATH, 'utf8'), + readFile(WHISPER_GENERATION_CONFIG_PATH, 'utf8'), + readFile(WHISPER_TOKENIZER_PATH, 'utf8'), + readFile(WHISPER_TOKENIZER_CONFIG_PATH, 'utf8'), + ]); + + const config = JSON.parse(configRaw) as { + decoder_start_token_id?: number; + eos_token_id?: number; + forced_decoder_ids?: Array<[number, number | null]>; + }; + + const generationConfig = JSON.parse(generationRaw) as { + no_timestamps_token_id?: number; + max_initial_timestamp_index?: number; + suppress_tokens?: number[]; + begin_suppress_tokens?: number[]; + max_length?: number; + alignment_heads?: Array<[number, number]>; + }; + + const tokenizer = new Tokenizer(JSON.parse(tokenizerJsonRaw), JSON.parse(tokenizerConfigRaw)); + + const promptStartToken = Number(config.decoder_start_token_id ?? 50258); + const eosTokenId = Number(config.eos_token_id ?? 50257); + const noTimestampsTokenId = Number(generationConfig.no_timestamps_token_id ?? 50363); + const timestampBeginTokenId = noTimestampsTokenId + 1; + const maxInitialTimestampIndex = Number(generationConfig.max_initial_timestamp_index ?? 50); + const configuredMaxDecodeSteps = Number(generationConfig.max_length ?? 448); + const maxDecodeSteps = Math.min(configuredMaxDecodeSteps, MAX_DECODE_STEPS_CAP); + const alignmentHeads = Array.isArray(generationConfig.alignment_heads) + ? generationConfig.alignment_heads + .filter((head): head is [number, number] => Array.isArray(head) && head.length === 2) + .map(([layer, head]) => [Number(layer), Number(head)] as [number, number]) + : []; + + const forcedDecoder = Array.isArray(config.forced_decoder_ids) ? config.forced_decoder_ids : []; + const defaultLanguageFromForced = forcedDecoder.find(([index, id]) => index === 1 && typeof id === 'number')?.[1] ?? null; + const transcribeFromForced = forcedDecoder.find(([index, id]) => index === 2 && typeof id === 'number')?.[1] ?? null; + + const defaultLanguageToken = Number(defaultLanguageFromForced ?? tokenizer.token_to_id('<|en|>') ?? 50259); + const transcribeToken = Number(transcribeFromForced ?? tokenizer.token_to_id('<|transcribe|>') ?? 50359); + + const stableSessionOptions: ort.InferenceSession.SessionOptions = { + executionProviders: ['cpu'], + graphOptimizationLevel: 'disabled', + intraOpNumThreads: 1, + interOpNumThreads: 1, + executionMode: 'sequential', + enableCpuMemArena: false, + enableMemPattern: false, + }; + + const encoder = await ort.InferenceSession.create(WHISPER_ENCODER_MODEL_PATH, stableSessionOptions); + const decoderMerged = await ort.InferenceSession.create(WHISPER_DECODER_MERGED_MODEL_PATH, stableSessionOptions); + const decoderWithPast = await ort.InferenceSession.create(WHISPER_DECODER_WITH_PAST_MODEL_PATH, stableSessionOptions); + + const alignmentLayers = [...new Set(alignmentHeads.map(([layer]) => layer))]; + const prefillFetches: string[] = ['logits']; + const stepFetches: string[] = ['logits']; + const mergedOutputNames = new Set(decoderMerged.outputNames); + const withPastOutputNames = new Set(decoderWithPast.outputNames); + + for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { + const decoderKey = `present.${i}.decoder.key`; + const decoderValue = `present.${i}.decoder.value`; + if (mergedOutputNames.has(decoderKey)) prefillFetches.push(decoderKey); + if (mergedOutputNames.has(decoderValue)) prefillFetches.push(decoderValue); + if (withPastOutputNames.has(decoderKey)) stepFetches.push(decoderKey); + if (withPastOutputNames.has(decoderValue)) stepFetches.push(decoderValue); + + const encoderKey = `present.${i}.encoder.key`; + const encoderValue = `present.${i}.encoder.value`; + if (mergedOutputNames.has(encoderKey)) prefillFetches.push(encoderKey); + if (mergedOutputNames.has(encoderValue)) prefillFetches.push(encoderValue); } - const idx = normalizedSentence - .toLowerCase() - .indexOf(token.toLowerCase(), cursor); - - const start = - idx !== -1 - ? idx - : cursor; - const end = start + token.length; - - cursor = end; + for (const layer of alignmentLayers) { + const key = `cross_attentions.${layer}`; + if (mergedOutputNames.has(key)) prefillFetches.push(key); + if (withPastOutputNames.has(key)) stepFetches.push(key); + } return { - text: token, - startSec: w.start, - endSec: w.end, - charStart: start, - charEnd: end, + encoder, + decoderMerged, + decoderWithPast, + tokenizer, + promptStartToken, + defaultLanguageToken, + transcribeToken, + eosTokenId, + noTimestampsTokenId, + timestampBeginTokenId, + maxInitialTimestampIndex, + maxDecodeSteps, + suppressTokens: new Set((generationConfig.suppress_tokens ?? []).map((v) => Number(v))), + beginSuppressTokens: new Set((generationConfig.begin_suppress_tokens ?? []).map((v) => Number(v))), + alignmentHeads, + prefillFetches, + stepFetches, }; + })().catch((error) => { + state.runtimePromise = null; + throw error; }); - return { - sentence, - sentenceIndex: 0, - words: alignedWords.filter((w) => w.text.length > 0), - }; + return state.runtimePromise; +} + +function resolveLanguageToken(runtime: WhisperRuntime, lang?: string): number { + const parsed = parseLanguageCode(lang); + if (!parsed) return runtime.defaultLanguageToken; + + const candidate = runtime.tokenizer.token_to_id(`<|${parsed}|>`); + return typeof candidate === 'number' ? candidate : runtime.defaultLanguageToken; +} + +async function runWhisperOnnx( + audioSamples: Float32Array, + opts: WhisperAlignmentOptions, + numFrames: number, + deadlineMs: number, +): Promise { + assertWithinDeadline(deadlineMs); + const runtime = await getRuntime(); + const decodeStepLimit = computeAdaptiveDecodeStepLimit(runtime.maxDecodeSteps, opts.textHint); + const mel = computeLogMelSpectrogram(audioSamples); + const encoderPast: Record = {}; + const decoderPast: Record = {}; + const crossAttentions: Record = {}; + let encoderHidden: ort.Tensor | null = null; + let outputs: Record | null = null; + + try { + const encoderOutputs = await runtime.encoder.run({ + input_features: mel, + }, ['last_hidden_state']); + encoderHidden = encoderOutputs.last_hidden_state; + + const languageToken = resolveLanguageToken(runtime, opts.lang); + const promptTokens = [ + runtime.promptStartToken, + languageToken, + runtime.transcribeToken, + ]; + + const generated: number[] = [...promptTokens]; + const emptyPastFeeds = buildEmptyPastFeeds(); + type LayerChunk = { + data: Float32Array; + heads: number; + seqLen: number; + frames: number; + }; + const selectedHeadsByLayer = new Map(); + for (const [layer, head] of runtime.alignmentHeads) { + const existing = selectedHeadsByLayer.get(layer) ?? []; + if (!existing.includes(head)) existing.push(head); + selectedHeadsByLayer.set(layer, existing); + } + for (const [layer, heads] of selectedHeadsByLayer) { + heads.sort((a, b) => a - b); + selectedHeadsByLayer.set(layer, heads); + } + const crossAttentionChunks = new Map(); + + const captureCrossAttentions = (stepOutputs: Record, prefill = false) => { + for (const [layer, selectedHeads] of selectedHeadsByLayer) { + const key = `cross_attentions.${layer}`; + const tensor = stepOutputs[key]; + if (!tensor) continue; + const [, , seqLen, frames] = tensor.dims; + const data = tensor.data as Float32Array; + const rowsToKeep = prefill ? seqLen : 1; + const seqStart = prefill ? 0 : Math.max(0, seqLen - 1); + const copied = new Float32Array(selectedHeads.length * rowsToKeep * frames); + for (let h = 0; h < selectedHeads.length; h += 1) { + const sourceHead = selectedHeads[h]!; + for (let s = 0; s < rowsToKeep; s += 1) { + const sourceSeq = seqStart + s; + for (let f = 0; f < frames; f += 1) { + const src = (((sourceHead * seqLen) + sourceSeq) * frames) + f; + const dst = (((h * rowsToKeep) + s) * frames) + f; + copied[dst] = data[src] ?? 0; + } + } + } + const list = crossAttentionChunks.get(layer) ?? []; + list.push({ data: copied, heads: selectedHeads.length, seqLen: rowsToKeep, frames }); + crossAttentionChunks.set(layer, list); + } + }; + const beginIndex = promptTokens.length; + + // Prefill: run prompt in merged decoder (non-cache branch), identical to first + // forward pass in transformers.js/transformers generation. + const prefillInputIds = tensorFromInt64(generated); + const prefillUseCacheBranch = new ort.Tensor('bool', Uint8Array.from([0]), [1]); + const prefillFeeds: Record = { + input_ids: prefillInputIds, + encoder_hidden_states: encoderHidden, + use_cache_branch: prefillUseCacheBranch, + ...emptyPastFeeds, + }; + try { + assertWithinDeadline(deadlineMs); + outputs = await runtime.decoderMerged.run(prefillFeeds, runtime.prefillFetches); + } finally { + disposeTensor(prefillInputIds); + disposeTensor(prefillUseCacheBranch); + } + captureCrossAttentions(outputs, true); + + for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { + encoderPast[`past_key_values.${i}.encoder.key`] = outputs[`present.${i}.encoder.key`]; + encoderPast[`past_key_values.${i}.encoder.value`] = outputs[`present.${i}.encoder.value`]; + decoderPast[`past_key_values.${i}.decoder.key`] = outputs[`present.${i}.decoder.key`]; + decoderPast[`past_key_values.${i}.decoder.value`] = outputs[`present.${i}.decoder.value`]; + } + + for (let step = 0; step < decodeStepLimit; step += 1) { + assertWithinDeadline(deadlineMs); + if (!outputs) break; + const logits = outputs.logits; + const logitsData = logits.data as Float32Array; + const vocabSize = logits.dims[2] ?? 0; + const offset = logitsData.length - vocabSize; + const lastLogits = logitsData.subarray(offset); + + applyTokenSuppression(lastLogits, runtime.suppressTokens); + if (generated.length === beginIndex) { + applyTokenSuppression(lastLogits, runtime.beginSuppressTokens); + } + applyWhisperTimestampLogitsRules({ + logits: lastLogits, + generated, + beginIndex, + eosTokenId: runtime.eosTokenId, + noTimestampsTokenId: runtime.noTimestampsTokenId, + timestampBeginTokenId: runtime.timestampBeginTokenId, + maxInitialTimestampIndex: runtime.maxInitialTimestampIndex, + }); + + const nextToken = argmax(lastLogits) ?? runtime.eosTokenId; + generated.push(nextToken); + if (nextToken === runtime.eosTokenId) break; + + const previousDecoderPast = { ...decoderPast }; + const stepInputIds = tensorFromInt64([nextToken]); + const stepFeeds: Record = { + input_ids: stepInputIds, + ...previousDecoderPast, + ...encoderPast, + }; + let nextOutputs: Record; + try { + assertWithinDeadline(deadlineMs); + nextOutputs = await runtime.decoderWithPast.run(stepFeeds, runtime.stepFetches); + } finally { + disposeTensor(stepInputIds); + } + captureCrossAttentions(nextOutputs, false); + + for (let i = 0; i < WHISPER_NUM_LAYERS; i += 1) { + decoderPast[`past_key_values.${i}.decoder.key`] = nextOutputs[`present.${i}.decoder.key`]; + decoderPast[`past_key_values.${i}.decoder.value`] = nextOutputs[`present.${i}.decoder.value`]; + } + + disposeTensorMap(previousDecoderPast); + disposeTensor(outputs.logits); + for (const [name, tensor] of Object.entries(outputs)) { + if (name.startsWith('cross_attentions.')) { + disposeTensor(tensor); + } + } + outputs = nextOutputs; + } + + if (crossAttentionChunks.size === 0) { + return []; + } + + const remappedAlignmentHeads: Array<[number, number]> = runtime.alignmentHeads + .map(([layer, head]) => { + const selectedHeads = selectedHeadsByLayer.get(layer) ?? []; + const remappedHead = selectedHeads.indexOf(head); + if (remappedHead < 0) return null; + return [layer, remappedHead] as [number, number]; + }) + .filter((pair): pair is [number, number] => pair !== null); + + for (let layer = 0; layer < WHISPER_NUM_LAYERS; layer += 1) { + const chunks = crossAttentionChunks.get(layer); + if (!chunks || !chunks.length) continue; + + const heads = chunks[0].heads; + const frames = chunks[0].frames; + const concatSeqLen = chunks.reduce((sum, chunk) => sum + chunk.seqLen, 0); + const merged = new Float32Array(1 * heads * concatSeqLen * frames); + let seqOffset = 0; + + for (const chunk of chunks) { + const { data, seqLen, frames: tensorFrames } = chunk; + const copyFrames = Math.min(frames, tensorFrames); + + for (let h = 0; h < heads; h += 1) { + for (let s = 0; s < seqLen; s += 1) { + for (let f = 0; f < copyFrames; f += 1) { + const src = (((h * seqLen) + s) * tensorFrames) + f; + const dst = (((h * concatSeqLen) + (seqOffset + s)) * frames) + f; + merged[dst] = data[src] ?? 0; + } + } + } + seqOffset += seqLen; + } + + crossAttentions[`cross_attentions.${layer}`] = new ort.Tensor('float32', merged, [1, heads, concatSeqLen, frames]); + } + + const tokenStartTimestamps = extractTokenStartTimestamps({ + crossAttentions, + decoderLayers: WHISPER_NUM_LAYERS, + alignmentHeads: remappedAlignmentHeads, + numFrames, + numInputIds: promptTokens.length, + timePrecision: 0.02, + sequenceLength: generated.length, + }); + + const timedWords = buildWordsFromTimestampedTokens({ + tokens: generated, + tokenStartTimestamps, + tokenizer: runtime.tokenizer, + eosTokenId: runtime.eosTokenId, + promptLength: promptTokens.length, + timestampBeginTokenId: runtime.timestampBeginTokenId, + timePrecision: 0.02, + language: parseLanguageCode(opts.lang) ?? 'english', + }); + + const maxSec = Math.max(0, numFrames * 0.02); + return timedWords.map((word) => ({ + word: word.word, + start: Math.min(maxSec, Math.max(0, word.startSec)), + end: Math.min(maxSec, Math.max(0, word.endSec)), + })); + } finally { + disposeTensor(mel); + if (outputs?.logits) disposeTensor(outputs.logits); + if (outputs) { + for (const [name, tensor] of Object.entries(outputs)) { + if (name.startsWith('cross_attentions.')) { + disposeTensor(tensor); + } + } + } + disposeTensorMap(crossAttentions); + disposeTensorMap(decoderPast); + disposeTensorMap(encoderPast); + disposeTensor(encoderHidden); + } } export async function alignAudioWithText( audioBuffer: TTSAudioBuffer, text: string, cacheKey?: string, - opts: WhisperAlignmentOptions = {} + opts: WhisperAlignmentOptions = {}, ): Promise { - if (!text.trim()) { - return []; - } + if (!text.trim()) return []; if (cacheKey && alignmentCache.has(cacheKey)) { - return alignmentCache.get(cacheKey)!; + const cached = alignmentCache.get(cacheKey)!; + alignmentCache.delete(cacheKey); + alignmentCache.set(cacheKey, cached); + return cached; } - const tmpBase = await mkdtemp(join(tmpdir(), 'openreader-whisper-')); - const inputPath = join(tmpBase, `${randomUUID()}-input.bin`); - const wavPath = join(tmpBase, `${randomUUID()}-input.wav`); + if (cacheKey) { + const inFlight = alignmentInFlight.get(cacheKey); + if (inFlight) return inFlight; + } + const inFlightKey = cacheKey ?? makeInFlightCoalesceKey(audioBuffer, text, opts.lang); + const shared = alignmentInFlight.get(inFlightKey); + if (shared) return shared; - try { - await writeFile(inputPath, Buffer.from(new Uint8Array(audioBuffer))); - - await new Promise((resolve, reject) => { - const ffmpeg = spawn(getFFmpegPath(), [ - '-y', - '-i', - inputPath, - '-ar', - '16000', - '-ac', - '1', - wavPath, - ]); - - let stderr = ''; - ffmpeg.stderr.on('data', (data) => { - stderr += data.toString(); - }); - - ffmpeg.on('error', (err) => { - reject(err); - }); - - ffmpeg.on('close', (code) => { - if (code === 0) { - resolve(); - } else { - reject( - new Error(`ffmpeg failed with code ${code}: ${stderr}`) - ); - } - }); + state.pendingAlignments += 1; + const run = (async (): Promise => { + const deadlineMs = Date.now() + ALIGNMENT_TIMEOUT_MS; + const previous = state.alignMutex; + let release!: () => void; + state.alignMutex = new Promise((resolve) => { + release = resolve; }); - const words = await runWhisperCpp(wavPath, opts); - const alignment = mapWordsToSentenceOffsets(text, words); - const result: TTSSentenceAlignment[] = [alignment]; + await previous; - if (cacheKey) { - alignmentCache.set(cacheKey, result); + // Another request with the same cache key may have completed while this one + // was waiting on the mutex. + if (cacheKey && alignmentCache.has(cacheKey)) { + const cached = alignmentCache.get(cacheKey)!; + alignmentCache.delete(cacheKey); + alignmentCache.set(cacheKey, cached); + release(); + return cached; } - return result; - } finally { - await rm(tmpBase, { recursive: true, force: true }).catch(() => {}); - } + let tmpBase = ''; + let inputPath = ''; + let pcmPath = ''; + + try { + tmpBase = await mkdtemp(join(tmpdir(), 'openreader-whisper-')); + inputPath = join(tmpBase, `${randomUUID()}-input.bin`); + pcmPath = join(tmpBase, `${randomUUID()}-input.pcm16`); + + await writeFile(inputPath, Buffer.from(new Uint8Array(audioBuffer))); + await decodeToPcm16(inputPath, pcmPath); + + const pcmBytes = await readFile(pcmPath); + const decodedSamples = pcm16ToFloat32(pcmBytes); + const effectiveSampleLength = Math.min(decodedSamples.length, N_SAMPLES); + const effectiveFrameCount = Math.max(1, Math.floor((effectiveSampleLength / HOP_LENGTH) / 2)); + const normalizedAudio = padOrTrimAudio(decodedSamples); + + const words = await runWhisperOnnx( + normalizedAudio, + { ...opts, textHint: text }, + effectiveFrameCount, + deadlineMs, + ); + const alignment = mapWordsToSentenceOffsets(text, words); + const result: TTSSentenceAlignment[] = [alignment]; + + if (cacheKey) { + if (alignmentCache.has(cacheKey)) { + alignmentCache.delete(cacheKey); + } + alignmentCache.set(cacheKey, result); + while (alignmentCache.size > ALIGNMENT_CACHE_MAX_ENTRIES) { + const oldest = alignmentCache.keys().next().value; + if (!oldest) break; + alignmentCache.delete(oldest); + } + } + + return result; + } finally { + if (tmpBase) { + await rm(tmpBase, { recursive: true, force: true }).catch(() => {}); + } + release(); + state.pendingAlignments = Math.max(0, state.pendingAlignments - 1); + } + })(); + + alignmentInFlight.set(inFlightKey, run); + run.finally(() => { + if (alignmentInFlight.get(inFlightKey) === run) { + alignmentInFlight.delete(inFlightKey); + } + }); + return run; } export function makeWhisperCacheKey(input: WhisperRequestBody): string { @@ -389,7 +1004,7 @@ export function makeWhisperCacheKey(input: WhisperRequestBody): string { text: input.text, lang: input.lang || '', audioLen: input.audio?.length || 0, - }) + }), ) .digest('hex'); } diff --git a/src/lib/server/whisper/ensureModel.ts b/src/lib/server/whisper/ensureModel.ts new file mode 100644 index 0000000..6d1ae77 --- /dev/null +++ b/src/lib/server/whisper/ensureModel.ts @@ -0,0 +1,226 @@ +import path from 'path'; +import { createHash } from 'crypto'; +import { access, copyFile, mkdir, readFile, rename, unlink, writeFile } from 'fs/promises'; +import { DOCSTORE_DIR } from '@/lib/server/storage/library-mount'; +import manifest from '@/lib/server/whisper/model/manifest.json'; + +const MODEL_DIR = path.join(DOCSTORE_DIR, 'model', 'whisper-base_timestamped'); +const STATIC_LICENSE_PATH = path.join(process.cwd(), 'src/lib/server/whisper/model/LICENSE.txt'); + +export const WHISPER_CONFIG_PATH = path.join(MODEL_DIR, 'config.json'); +export const WHISPER_GENERATION_CONFIG_PATH = path.join(MODEL_DIR, 'generation_config.json'); +export const WHISPER_TOKENIZER_PATH = path.join(MODEL_DIR, 'tokenizer.json'); +export const WHISPER_TOKENIZER_CONFIG_PATH = path.join(MODEL_DIR, 'tokenizer_config.json'); +export const WHISPER_ENCODER_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'encoder_model_int8.onnx'); +export const WHISPER_DECODER_MERGED_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'decoder_model_merged_int8.onnx'); +export const WHISPER_DECODER_WITH_PAST_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'decoder_with_past_model_int8.onnx'); + +const BASE_MODEL_URL = 'https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main'; + +const DEFAULT_URLS: Record = { + 'config.json': `${BASE_MODEL_URL}/config.json`, + 'generation_config.json': `${BASE_MODEL_URL}/generation_config.json`, + 'tokenizer.json': `${BASE_MODEL_URL}/tokenizer.json`, + 'tokenizer_config.json': `${BASE_MODEL_URL}/tokenizer_config.json`, + 'merges.txt': `${BASE_MODEL_URL}/merges.txt`, + 'vocab.json': `${BASE_MODEL_URL}/vocab.json`, + 'normalizer.json': `${BASE_MODEL_URL}/normalizer.json`, + 'added_tokens.json': `${BASE_MODEL_URL}/added_tokens.json`, + 'preprocessor_config.json': `${BASE_MODEL_URL}/preprocessor_config.json`, + 'special_tokens_map.json': `${BASE_MODEL_URL}/special_tokens_map.json`, + 'onnx/encoder_model_int8.onnx': `${BASE_MODEL_URL}/onnx/encoder_model_int8.onnx`, + 'onnx/decoder_model_merged_int8.onnx': `${BASE_MODEL_URL}/onnx/decoder_model_merged_int8.onnx`, + 'onnx/decoder_with_past_model_int8.onnx': `${BASE_MODEL_URL}/onnx/decoder_with_past_model_int8.onnx`, +}; + +const ENV_URL_OVERRIDES: Record = { + 'config.json': 'OPENREADER_WHISPER_MODEL_CONFIG_URL', + 'generation_config.json': 'OPENREADER_WHISPER_MODEL_GENERATION_CONFIG_URL', + 'tokenizer.json': 'OPENREADER_WHISPER_MODEL_TOKENIZER_URL', + 'tokenizer_config.json': 'OPENREADER_WHISPER_MODEL_TOKENIZER_CONFIG_URL', + 'merges.txt': 'OPENREADER_WHISPER_MODEL_MERGES_URL', + 'vocab.json': 'OPENREADER_WHISPER_MODEL_VOCAB_URL', + 'normalizer.json': 'OPENREADER_WHISPER_MODEL_NORMALIZER_URL', + 'added_tokens.json': 'OPENREADER_WHISPER_MODEL_ADDED_TOKENS_URL', + 'preprocessor_config.json': 'OPENREADER_WHISPER_MODEL_PREPROCESSOR_URL', + 'special_tokens_map.json': 'OPENREADER_WHISPER_MODEL_SPECIAL_TOKENS_MAP_URL', + 'onnx/encoder_model_int8.onnx': 'OPENREADER_WHISPER_MODEL_ENCODER_URL', + 'onnx/decoder_model_merged_int8.onnx': 'OPENREADER_WHISPER_MODEL_DECODER_MERGED_URL', + 'onnx/decoder_with_past_model_int8.onnx': 'OPENREADER_WHISPER_MODEL_DECODER_WITH_PAST_URL', +}; + +type ManifestEntry = { path: string; sha256?: string; size?: number }; + +export interface WhisperArtifactSpec { + path: string; + sha256?: string; + size?: number; + url: string; +} + +export interface WhisperStaticArtifactSpec { + path: string; + sha256?: string; + size?: number; + sourcePath: string; +} + +export type WhisperFetch = (input: RequestInfo | URL, init?: RequestInit) => Promise; + +const MANIFEST_FILES = manifest.files as ManifestEntry[]; +const MODEL_FILES = MANIFEST_FILES.filter((entry) => entry.path !== 'LICENSE.txt'); +const LICENSE_FILE = MANIFEST_FILES.find((entry) => entry.path === 'LICENSE.txt'); + +function normalizeExpected(entry: { sha256?: string; size?: number }): { sha256: string | null; size: number } { + return { + sha256: typeof entry.sha256 === 'string' ? entry.sha256.toLowerCase() : null, + size: Number(entry.size ?? 0), + }; +} + +function resolvePath(relativePath: string, modelDir: string): string { + return path.join(modelDir, relativePath); +} + +function resolveUrl(relativePath: string): string { + const envKey = ENV_URL_OVERRIDES[relativePath]; + const override = envKey ? process.env[envKey]?.trim() : ''; + if (override) return override; + const fallback = DEFAULT_URLS[relativePath]; + if (!fallback) { + throw new Error(`No default URL configured for Whisper model artifact: ${relativePath}`); + } + return fallback; +} + +function sha256OfBytes(bytes: Uint8Array): string { + return createHash('sha256').update(bytes).digest('hex'); +} + +function verifyBytes(bytes: Uint8Array, expected: { sha256?: string; size?: number }): boolean { + const normalized = normalizeExpected(expected); + if (Number.isFinite(normalized.size) && normalized.size > 0 && bytes.byteLength !== normalized.size) { + return false; + } + if (!normalized.sha256) return true; + return sha256OfBytes(bytes) === normalized.sha256; +} + +async function verifyFile(filePath: string, expected: { sha256?: string; size?: number }): Promise { + const bytes = await readFile(filePath); + return verifyBytes(bytes, expected); +} + +async function downloadToFile(fetchImpl: WhisperFetch, url: string, outPath: string): Promise { + const res = await fetchImpl(url); + if (!res.ok) { + throw new Error(`Download failed for ${url}: ${res.status} ${res.statusText}`); + } + const bytes = new Uint8Array(await res.arrayBuffer()); + await writeFile(outPath, bytes); +} + +export async function ensureWhisperArtifacts(options: { + modelDir: string; + artifacts: WhisperArtifactSpec[]; + staticArtifacts?: WhisperStaticArtifactSpec[]; + fetchImpl?: WhisperFetch; +}): Promise { + const { + modelDir, + artifacts, + staticArtifacts = [], + fetchImpl = fetch, + } = options; + + try { + await Promise.all(artifacts.map(async (artifact) => { + const target = resolvePath(artifact.path, modelDir); + await access(target); + const valid = await verifyFile(target, artifact); + if (!valid) { + throw new Error(`Checksum mismatch for existing Whisper artifact: ${artifact.path}`); + } + })); + + await Promise.all(staticArtifacts.map(async (artifact) => { + const target = resolvePath(artifact.path, modelDir); + await access(target); + const valid = await verifyFile(target, artifact); + if (!valid) { + throw new Error(`Checksum mismatch for existing Whisper static artifact: ${artifact.path}`); + } + })); + + return; + } catch { + // Continue to repair/download. + } + + for (const artifact of artifacts) { + const target = resolvePath(artifact.path, modelDir); + const targetDir = path.dirname(target); + const tmp = `${target}.tmp`; + + await mkdir(targetDir, { recursive: true }); + await downloadToFile(fetchImpl, artifact.url, tmp); + if (!(await verifyFile(tmp, artifact))) { + await unlink(tmp).catch(() => undefined); + throw new Error(`Whisper artifact checksum verification failed: ${artifact.path}`); + } + await rename(tmp, target); + } + + for (const artifact of staticArtifacts) { + const target = resolvePath(artifact.path, modelDir); + const targetDir = path.dirname(target); + await mkdir(targetDir, { recursive: true }); + await copyFile(artifact.sourcePath, target); + if (!(await verifyFile(target, artifact))) { + throw new Error(`Whisper static artifact checksum verification failed: ${artifact.path}`); + } + } +} + +export function createSingleflightRunner(work: () => Promise): () => Promise { + let inflight: Promise | null = null; + return async () => { + if (inflight) return inflight; + inflight = work().finally(() => { + inflight = null; + }); + return inflight; + }; +} + +async function ensureModelInternal(): Promise { + const artifacts: WhisperArtifactSpec[] = MODEL_FILES.map((entry) => ({ + path: entry.path, + sha256: entry.sha256, + size: entry.size, + url: resolveUrl(entry.path), + })); + + const staticArtifacts: WhisperStaticArtifactSpec[] = LICENSE_FILE + ? [{ + path: LICENSE_FILE.path, + sha256: LICENSE_FILE.sha256, + size: LICENSE_FILE.size, + sourcePath: STATIC_LICENSE_PATH, + }] + : []; + + await ensureWhisperArtifacts({ + modelDir: MODEL_DIR, + artifacts, + staticArtifacts, + }); + + return WHISPER_ENCODER_MODEL_PATH; +} + +const ensureWhisperModelSingleflight = createSingleflightRunner(ensureModelInternal); + +export async function ensureWhisperModel(): Promise { + return ensureWhisperModelSingleflight(); +} diff --git a/src/lib/server/whisper/model/LICENSE.txt b/src/lib/server/whisper/model/LICENSE.txt new file mode 100644 index 0000000..d255525 --- /dev/null +++ b/src/lib/server/whisper/model/LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 OpenAI + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/lib/server/whisper/model/manifest.json b/src/lib/server/whisper/model/manifest.json new file mode 100644 index 0000000..2fffd3d --- /dev/null +++ b/src/lib/server/whisper/model/manifest.json @@ -0,0 +1,76 @@ +{ + "name": "whisper-base_timestamped-int8", + "version": "onnx-community/whisper-base_timestamped@608c49e61301901684bc36cac8f74b95ff6b5a8e", + "files": [ + { + "path": "config.json", + "sha256": "f4d0608f7d918166da7edb3e188de5ef1bfe70d9802e785d271fd88111e9cf4b", + "size": 2243 + }, + { + "path": "generation_config.json", + "sha256": "61070cf8de25b1e9256e8e102ded49d8d24a8369ed36ef84fdf21549e68125a0", + "size": 3832 + }, + { + "path": "tokenizer.json", + "sha256": "27fc476bfe7f17299480be2273fc0608e4d5a99aba2ab5dec5374b4482d1a566", + "size": 2480466 + }, + { + "path": "tokenizer_config.json", + "sha256": "2e036e4dbacfdeb7242c7d4ec4149f4a16e86026048f94d1637e3a8ee9c6a573", + "size": 282682 + }, + { + "path": "merges.txt", + "sha256": "2df2990a395e35e8dfbc7511e08c12d56018d8d04691e0133e5d63b21e154dc6", + "size": 493869 + }, + { + "path": "vocab.json", + "sha256": "50d6a919f0a0601d56a04eb583c780d18553aa388254ba3158eb6a00f13e2c1a", + "size": 1036584 + }, + { + "path": "normalizer.json", + "sha256": "bf1c507dc8724ca9cf9903640dacfb69dae2f00edee4f21ceba106a7392f26dd", + "size": 52666 + }, + { + "path": "added_tokens.json", + "sha256": "9715fd2243b6f06a5858b5e32950d2853f73dd5bc201aafcf76f5082a2d8acd1", + "size": 34604 + }, + { + "path": "preprocessor_config.json", + "sha256": "a6a76d28c93edb273669eb9e0b0636a2bddbb1272c3261e47b7ca6dfdbac1b8d", + "size": 339 + }, + { + "path": "special_tokens_map.json", + "sha256": "e67ae3a0aaa99abcd9f187138e12db1f65c16a14761c50ef10eef2c174a7a691", + "size": 2194 + }, + { + "path": "onnx/encoder_model_int8.onnx", + "sha256": "152da96dd8ff3f28f3fadabc2e8960405a277846453ff94ed411fe935a72917f", + "size": 23159150 + }, + { + "path": "onnx/decoder_model_merged_int8.onnx", + "sha256": "cf9a8d5bcddc0917a0078135b484cedcaf44f28909cd91910abd29dced9171db", + "size": 53712708 + }, + { + "path": "onnx/decoder_with_past_model_int8.onnx", + "sha256": "bdd92860d0ed7dff2aca623963378cbba1b617bfae127356db1c8aa8baa930ef", + "size": 50131672 + }, + { + "path": "LICENSE.txt", + "sha256": "b5d65a59060e68c4ff940e1eddfa6f94b2d68fdf58ed7f4dd57721c997e35e9d", + "size": 1063 + } + ] +} diff --git a/src/lib/server/whisper/model/mel_filters.npz b/src/lib/server/whisper/model/mel_filters.npz new file mode 100644 index 0000000..28ea269 Binary files /dev/null and b/src/lib/server/whisper/model/mel_filters.npz differ diff --git a/src/lib/server/whisper/spectral.ts b/src/lib/server/whisper/spectral.ts new file mode 100644 index 0000000..b7223cc --- /dev/null +++ b/src/lib/server/whisper/spectral.ts @@ -0,0 +1,21 @@ +export function buildGoertzelCoefficients(freqBins: number, fftSize: number): Float64Array { + const coeffs = new Float64Array(freqBins); + for (let k = 0; k < freqBins; k += 1) { + coeffs[k] = 2 * Math.cos((2 * Math.PI * k) / fftSize); + } + return coeffs; +} + +export function goertzelPower(samples: Float32Array, coeff: number): number { + let s1 = 0; + let s2 = 0; + for (let i = 0; i < samples.length; i += 1) { + const s0 = samples[i] + (coeff * s1) - s2; + s2 = s1; + s1 = s0; + } + + const power = (s1 * s1) + (s2 * s2) - (coeff * s1 * s2); + if (!Number.isFinite(power) || power < 0) return 0; + return power; +} diff --git a/src/lib/server/whisper/token-timestamps.ts b/src/lib/server/whisper/token-timestamps.ts new file mode 100644 index 0000000..47edc1d --- /dev/null +++ b/src/lib/server/whisper/token-timestamps.ts @@ -0,0 +1,449 @@ +import type { Tokenizer } from '@huggingface/tokenizers'; +import type * as ort from 'onnxruntime-node'; + +const PUNCTUATION_REGEX = '\\p{P}\\u0021-\\u002F\\u003A-\\u0040\\u005B-\\u0060\\u007B-\\u007E'; +const PUNCTUATION_ONLY_REGEX = new RegExp(`^[${PUNCTUATION_REGEX}]+$`, 'gu'); + +type TokenTimestamp = [start: number, end: number]; + +export interface WhisperWordTiming { + word: string; + startSec: number; + endSec: number; +} + +function medianFilter(data: Float32Array, windowSize: number): Float32Array { + if (windowSize % 2 === 0 || windowSize <= 0) { + throw new Error('Window size must be a positive odd number'); + } + + const output = new Float32Array(data.length); + const buffer = new Float32Array(windowSize); + const halfWindow = Math.floor(windowSize / 2); + + for (let i = 0; i < data.length; i += 1) { + let valuesIndex = 0; + for (let j = -halfWindow; j <= halfWindow; j += 1) { + let index = i + j; + if (index < 0) { + index = Math.abs(index); + } else if (index >= data.length) { + index = (2 * (data.length - 1)) - index; + } + buffer[valuesIndex] = data[index]; + valuesIndex += 1; + } + + const sortable = Array.from(buffer); + sortable.sort((a, b) => a - b); + output[i] = sortable[halfWindow] ?? 0; + } + + return output; +} + +function dynamicTimeWarping(matrix: Float32Array[], rows: number, cols: number): [number[], number[]] { + const cost: number[][] = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(Number.POSITIVE_INFINITY)); + const trace: number[][] = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(-1)); + cost[0][0] = 0; + + for (let j = 1; j <= cols; j += 1) { + for (let i = 1; i <= rows; i += 1) { + const c0 = cost[i - 1][j - 1]; + const c1 = cost[i - 1][j]; + const c2 = cost[i][j - 1]; + let c: number; + let t: number; + if (c0 < c1 && c0 < c2) { + c = c0; + t = 0; + } else if (c1 < c0 && c1 < c2) { + c = c1; + t = 1; + } else { + c = c2; + t = 2; + } + cost[i][j] = matrix[i - 1][j - 1] + c; + trace[i][j] = t; + } + } + + for (let i = 0; i <= cols; i += 1) trace[0][i] = 2; + for (let i = 0; i <= rows; i += 1) trace[i][0] = 1; + + let i = rows; + let j = cols; + const textIndices: number[] = []; + const timeIndices: number[] = []; + while (i > 0 || j > 0) { + textIndices.push(i - 1); + timeIndices.push(j - 1); + const step = trace[i][j]; + if (step === 0) { + i -= 1; + j -= 1; + } else if (step === 1) { + i -= 1; + } else if (step === 2) { + j -= 1; + } else { + throw new Error(`Unexpected DTW trace state at [${i}, ${j}]`); + } + } + + textIndices.reverse(); + timeIndices.reverse(); + return [textIndices, timeIndices]; +} + +function round2(value: number): number { + return Math.round(value * 100) / 100; +} + +function decodeTokens(tokenizer: Pick, tokens: number[]): string { + return tokenizer.decode(tokens, { skip_special_tokens: false }); +} + +function splitTokensOnUnicode( + tokenizer: Pick, + tokens: number[], +): [string[], number[][], number[][]] { + const decodedFull = decodeTokens(tokenizer, tokens); + const replacementChar = '\uFFFD'; + const words: string[] = []; + const wordTokens: number[][] = []; + const tokenIndices: number[][] = []; + let currentTokens: number[] = []; + let currentIndices: number[] = []; + let unicodeOffset = 0; + + for (let i = 0; i < tokens.length; i += 1) { + currentTokens.push(tokens[i]); + currentIndices.push(i); + + const decoded = decodeTokens(tokenizer, currentTokens); + if ( + !decoded.includes(replacementChar) + || decodedFull[unicodeOffset + decoded.indexOf(replacementChar)] === replacementChar + ) { + words.push(decoded); + wordTokens.push(currentTokens); + tokenIndices.push(currentIndices); + currentTokens = []; + currentIndices = []; + unicodeOffset += decoded.length; + } + } + + return [words, wordTokens, tokenIndices]; +} + +function splitTokensOnSpaces( + tokenizer: Pick, + tokens: number[], + eosTokenId: number, +): [string[], number[][], number[][]] { + const [subwords, subwordTokens, subwordIndices] = splitTokensOnUnicode(tokenizer, tokens); + const words: string[] = []; + const wordTokens: number[][] = []; + const tokenIndices: number[][] = []; + + for (let i = 0; i < subwords.length; i += 1) { + const subword = subwords[i]; + const tokenList = subwordTokens[i]; + const indices = subwordIndices[i]; + const special = tokenList[0] >= eosTokenId; + const withSpace = subword.startsWith(' '); + const trimmed = subword.trim(); + const punctuation = PUNCTUATION_ONLY_REGEX.test(trimmed); + + if (special || withSpace || punctuation || words.length === 0) { + words.push(subword); + wordTokens.push([...tokenList]); + tokenIndices.push([...indices]); + } else { + const ix = words.length - 1; + words[ix] += subword; + wordTokens[ix].push(...tokenList); + tokenIndices[ix].push(...indices); + } + } + + return [words, wordTokens, tokenIndices]; +} + +function mergePunctuations( + words: string[], + tokens: number[][], + indices: number[][], + prependPunctuations = '"\'โ€œยกยฟ([{-', + appendPunctuations = '"\'.ใ€‚,๏ผŒ!๏ผ?๏ผŸ:๏ผšโ€)]}ใ€', +): [string[], number[][], number[][]] { + const newWords = words.map((w) => `${w}`); + const newTokens = tokens.map((t) => [...t]); + const newIndices = indices.map((idx) => [...idx]); + + let i = newWords.length - 2; + let j = newWords.length - 1; + while (i >= 0) { + if (newWords[i].startsWith(' ') && prependPunctuations.includes(newWords[i].trim())) { + newWords[j] = newWords[i] + newWords[j]; + newTokens[j] = [...newTokens[i], ...newTokens[j]]; + newIndices[j] = [...newIndices[i], ...newIndices[j]]; + newWords[i] = ''; + newTokens[i] = []; + newIndices[i] = []; + } else { + j = i; + } + i -= 1; + } + + i = 0; + j = 1; + while (j < newWords.length) { + if (!newWords[i].endsWith(' ') && appendPunctuations.includes(newWords[j])) { + newWords[i] += newWords[j]; + newTokens[i] = [...newTokens[i], ...newTokens[j]]; + newIndices[i] = [...newIndices[i], ...newIndices[j]]; + newWords[j] = ''; + newTokens[j] = []; + newIndices[j] = []; + } else { + i = j; + } + j += 1; + } + + return [ + newWords.filter((w) => w.length > 0), + newTokens.filter((t) => t.length > 0), + newIndices.filter((t) => t.length > 0), + ]; +} + +function combineTokensIntoWords( + tokenizer: Pick, + tokens: number[], + eosTokenId: number, + language = 'english', +): [string[], number[][], number[][]] { + let words: string[]; + let wordTokens: number[][]; + let tokenIndices: number[][]; + + if (['chinese', 'japanese', 'thai', 'lao', 'myanmar', 'zh', 'ja', 'th', 'lo', 'my'].includes(language)) { + [words, wordTokens, tokenIndices] = splitTokensOnUnicode(tokenizer, tokens); + } else { + [words, wordTokens, tokenIndices] = splitTokensOnSpaces(tokenizer, tokens, eosTokenId); + } + + return mergePunctuations(words, wordTokens, tokenIndices); +} + +export function extractTokenStartTimestamps(input: { + crossAttentions: Record; + decoderLayers: number; + alignmentHeads: Array<[number, number]>; + numFrames: number; + numInputIds: number; + timePrecision?: number; + sequenceLength: number; +}): number[] { + const { + crossAttentions, + decoderLayers, + alignmentHeads, + numFrames, + numInputIds, + timePrecision = 0.02, + sequenceLength, + } = input; + + const frameCount = Math.max(1, numFrames); + const perLayer: Float32Array[] = []; + for (let layer = 0; layer < decoderLayers; layer += 1) { + const key = `cross_attentions.${layer}`; + const tensor = crossAttentions[key]; + if (!tensor) continue; + perLayer[layer] = tensor.data as Float32Array; + } + + const selected: Float32Array[] = []; + let seqLen = 0; + let attnFrames = 0; + for (const [layer, head] of alignmentHeads) { + const flat = perLayer[layer]; + if (!flat) continue; + const layerTensor = crossAttentions[`cross_attentions.${layer}`]; + if (!layerTensor || layerTensor.dims.length < 4) continue; + const [, numHeads, currentSeqLen, currentFrames] = layerTensor.dims; + if (head >= numHeads) continue; + seqLen = currentSeqLen; + attnFrames = Math.min(currentFrames, frameCount); + const headSlice = new Float32Array(seqLen * attnFrames); + for (let s = 0; s < seqLen; s += 1) { + for (let f = 0; f < attnFrames; f += 1) { + const flatIndex = (((head * currentSeqLen) + s) * currentFrames) + f; + headSlice[(s * attnFrames) + f] = flat[flatIndex] ?? 0; + } + } + selected.push(headSlice); + } + + if (!selected.length || seqLen === 0 || attnFrames === 0) { + return new Array(sequenceLength).fill(0); + } + + const normalizedHeads = selected.map((headData) => { + const means = new Float32Array(attnFrames); + const stds = new Float32Array(attnFrames); + + for (let f = 0; f < attnFrames; f += 1) { + let sum = 0; + for (let s = 0; s < seqLen; s += 1) sum += headData[(s * attnFrames) + f]; + const mean = sum / seqLen; + means[f] = mean; + let varSum = 0; + for (let s = 0; s < seqLen; s += 1) { + const d = headData[(s * attnFrames) + f] - mean; + varSum += d * d; + } + stds[f] = Math.sqrt(varSum / seqLen) || 1; + } + + const out = new Float32Array(headData.length); + for (let s = 0; s < seqLen; s += 1) { + const row = new Float32Array(attnFrames); + for (let f = 0; f < attnFrames; f += 1) { + row[f] = (headData[(s * attnFrames) + f] - means[f]) / stds[f]; + } + const filtered = medianFilter(row, 7); + out.set(filtered, s * attnFrames); + } + return out; + }); + + const croppedRows = Math.max(0, seqLen - numInputIds); + if (croppedRows === 0) return new Array(sequenceLength).fill(0); + + const matrix: Float32Array[] = Array.from({ length: croppedRows }, () => new Float32Array(attnFrames)); + for (const headData of normalizedHeads) { + for (let r = 0; r < croppedRows; r += 1) { + const srcRow = r + numInputIds; + for (let f = 0; f < attnFrames; f += 1) { + matrix[r][f] += headData[(srcRow * attnFrames) + f]; + } + } + } + + const scale = 1 / normalizedHeads.length; + for (let r = 0; r < croppedRows; r += 1) { + for (let f = 0; f < attnFrames; f += 1) { + matrix[r][f] = -matrix[r][f] * scale; + } + } + + const [textIndices, timeIndices] = dynamicTimeWarping(matrix, croppedRows, attnFrames); + const jumps = new Array(textIndices.length).fill(false); + for (let i = 0; i < textIndices.length; i += 1) { + jumps[i] = i === 0 ? true : textIndices[i] !== textIndices[i - 1]; + } + + const jumpTimes: number[] = []; + for (let i = 0; i < jumps.length; i += 1) { + if (jumps[i]) jumpTimes.push(timeIndices[i] * timePrecision); + } + + const timestamps = new Array(sequenceLength).fill(0); + for (let i = 0; i < numInputIds && i < timestamps.length; i += 1) timestamps[i] = 0; + for (let i = 0; i < jumpTimes.length && (numInputIds + i) < timestamps.length; i += 1) { + timestamps[numInputIds + i] = jumpTimes[i]; + } + if (timestamps.length > 0 && jumpTimes.length > 0) { + timestamps[timestamps.length - 1] = jumpTimes[jumpTimes.length - 1]; + } + return timestamps; +} + +export function buildWordsFromTimestampedTokens(input: { + tokens: number[]; + tokenStartTimestamps: number[]; + tokenizer: Pick; + eosTokenId: number; + promptLength: number; + timestampBeginTokenId: number; + timePrecision?: number; + language?: string; +}): WhisperWordTiming[] { + const { + tokens, + tokenStartTimestamps, + tokenizer, + eosTokenId, + promptLength, + timestampBeginTokenId, + timePrecision = 0.02, + language = 'english', + } = input; + + const limit = Math.min(tokens.length, tokenStartTimestamps.length); + const tokenRanges: TokenTimestamp[] = []; + for (let i = 0; i < limit; i += 1) { + const start = tokenStartTimestamps[i] ?? 0; + const end = i + 1 < limit ? (tokenStartTimestamps[i + 1] ?? (start + timePrecision)) : (start + timePrecision); + tokenRanges.push([start, Math.max(start, end)]); + } + + const words: WhisperWordTiming[] = []; + let segmentStart: number | null = null; + let textTokens: number[] = []; + let textRanges: TokenTimestamp[] = []; + + const flushSegment = (segmentEnd: number | null) => { + if (!textTokens.length) return; + const [wordTexts, , tokenIndices] = combineTokensIntoWords(tokenizer, textTokens, eosTokenId, language); + for (let i = 0; i < wordTexts.length; i += 1) { + const indices = tokenIndices[i]; + if (!indices.length) continue; + const start = textRanges[indices[0]]?.[0] ?? segmentStart ?? 0; + const end = textRanges[indices[indices.length - 1]]?.[1] ?? segmentEnd ?? start; + const clampedStart = segmentStart == null ? start : Math.max(segmentStart, start); + const clampedEndBase = segmentEnd == null ? end : Math.min(segmentEnd, end); + const clampedEnd = Math.max( + clampedStart + (clampedEndBase <= clampedStart ? timePrecision : 0), + clampedEndBase, + ); + words.push({ + word: wordTexts[i].trim(), + startSec: round2(clampedStart), + endSec: round2(clampedEnd), + }); + } + textTokens = []; + textRanges = []; + }; + + for (let i = promptLength; i < limit; i += 1) { + const token = tokens[i]; + if (token === eosTokenId) break; + + if (token >= timestampBeginTokenId) { + const ts = (token - timestampBeginTokenId) * timePrecision; + if (segmentStart == null) { + segmentStart = ts; + } else { + flushSegment(ts); + segmentStart = ts; + } + continue; + } + + textTokens.push(token); + textRanges.push(tokenRanges[i]); + } + + flushSegment(null); + return words.filter((w) => w.word.length > 0); +} diff --git a/src/types/client.ts b/src/types/client.ts index aed3042..0b94363 100644 --- a/src/types/client.ts +++ b/src/types/client.ts @@ -1,8 +1,7 @@ import type { TTSAudiobookChapter, - TTSSentenceAlignment, - TTSAudioBytes, TTSAudiobookFormat, + TTSSentenceAlignment, } from '@/types/tts'; import type { TtsProviderType } from '@/lib/shared/tts-provider-catalog'; @@ -64,17 +63,6 @@ export interface VoicesResponse { voices: string[]; } -// --- Whisper API Types --- - -export interface AlignmentPayload { - text: string; - audio: TTSAudioBytes; // Array.from(new Uint8Array(arrayBuffer)) -} - -export interface AlignmentResponse { - alignments: TTSSentenceAlignment[]; -} - export interface TTSSegmentSettings { providerRef: string; providerType: TtsProviderType; diff --git a/tests/unit/whisper-alignment-mapping.spec.ts b/tests/unit/whisper-alignment-mapping.spec.ts new file mode 100644 index 0000000..aba2933 --- /dev/null +++ b/tests/unit/whisper-alignment-mapping.spec.ts @@ -0,0 +1,21 @@ +import { test, expect } from '@playwright/test'; +import { + mapWordsToSentenceOffsets, +} from '../../src/lib/server/whisper/alignment-mapping'; + +test.describe('whisper alignment mapping', () => { + test('maps words to sentence offsets with punctuation and repeated spaces', () => { + const aligned = mapWordsToSentenceOffsets('Hello, world again.', [ + { word: 'Hello', start: 0, end: 0.25 }, + { word: 'world', start: 0.25, end: 0.5 }, + { word: 'again', start: 0.5, end: 1.0 }, + ]); + + expect(aligned.words).toHaveLength(3); + expect(aligned.words[0].charStart).toBe(0); + expect(aligned.words[0].charEnd).toBe(5); + expect(aligned.words[1].charStart).toBeGreaterThan(aligned.words[0].charEnd); + expect(aligned.words[2].charStart).toBeGreaterThan(aligned.words[1].charEnd); + expect(aligned.words[2].charEnd).toBeLessThanOrEqual('Hello, world again.'.length); + }); +}); diff --git a/tests/unit/whisper-alignment-smoke.spec.ts b/tests/unit/whisper-alignment-smoke.spec.ts new file mode 100644 index 0000000..2693dc9 --- /dev/null +++ b/tests/unit/whisper-alignment-smoke.spec.ts @@ -0,0 +1,37 @@ +import { test, expect } from '@playwright/test'; +import { readFile } from 'fs/promises'; +import path from 'path'; +import { alignAudioWithText } from '../../src/lib/server/whisper/alignment'; + +test.describe('whisper alignment smoke', () => { + test('runs ONNX alignment end-to-end without decoder reshape errors', async () => { + test.setTimeout(180000); + + const audioPath = path.join(process.cwd(), 'tests/files/sample.mp3'); + const audioBytes = await readFile(audioPath); + const buffer = audioBytes.buffer.slice(audioBytes.byteOffset, audioBytes.byteOffset + audioBytes.byteLength); + + const alignments = await alignAudioWithText( + buffer, + 'This is a sample sentence used to validate whisper alignment execution.', + undefined, + { lang: 'en' }, + ); + + expect(alignments.length).toBe(1); + expect(Array.isArray(alignments[0].words)).toBe(true); + expect(alignments[0].words.length).toBeGreaterThan(0); + + let maxEnd = 0; + let positiveDurationWordCount = 0; + for (const word of alignments[0].words) { + expect(Number.isFinite(word.startSec)).toBe(true); + expect(Number.isFinite(word.endSec)).toBe(true); + expect(word.endSec).toBeGreaterThanOrEqual(word.startSec); + maxEnd = Math.max(maxEnd, word.endSec); + if (word.endSec > word.startSec) positiveDurationWordCount += 1; + } + expect(maxEnd).toBeLessThanOrEqual(10.2); + expect(positiveDurationWordCount).toBeGreaterThan(0); + }); +}); diff --git a/tests/unit/whisper-ensure-model.spec.ts b/tests/unit/whisper-ensure-model.spec.ts new file mode 100644 index 0000000..166f022 --- /dev/null +++ b/tests/unit/whisper-ensure-model.spec.ts @@ -0,0 +1,71 @@ +import { test, expect } from '@playwright/test'; +import { createHash } from 'crypto'; +import { mkdir, mkdtemp, readFile, rm, writeFile } from 'fs/promises'; +import { tmpdir } from 'os'; +import path from 'path'; +import { + createSingleflightRunner, + ensureWhisperArtifacts, +} from '../../src/lib/server/whisper/ensureModel'; + +function sha256(bytes: Uint8Array): string { + return createHash('sha256').update(bytes).digest('hex'); +} + +test.describe('whisper ensure model helpers', () => { + test('downloads and verifies artifacts, and repairs checksum mismatch', async () => { + const root = await mkdtemp(path.join(tmpdir(), 'openreader-whisper-model-test-')); + const artifactBytes = new TextEncoder().encode('artifact-content-v1'); + const artifactHash = sha256(artifactBytes); + const artifactPath = 'onnx/encoder_model_int8.onnx'; + const target = path.join(root, artifactPath); + + try { + // Seed a corrupted file to verify repair behavior. + await mkdir(path.dirname(target), { recursive: true }); + await writeFile(target, new Uint8Array([0, 1, 2, 3])); + + let fetchCount = 0; + await ensureWhisperArtifacts({ + modelDir: root, + artifacts: [ + { + path: artifactPath, + sha256: artifactHash, + size: artifactBytes.byteLength, + url: 'https://example.local/fake-artifact', + }, + ], + fetchImpl: async () => { + fetchCount += 1; + return new Response(artifactBytes, { status: 200 }); + }, + }); + + const repaired = await readFile(target); + expect(repaired.byteLength).toBe(artifactBytes.byteLength); + expect(sha256(repaired)).toBe(artifactHash); + expect(fetchCount).toBe(1); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + + test('singleflight runner deduplicates concurrent work', async () => { + let runs = 0; + const run = createSingleflightRunner(async () => { + runs += 1; + await new Promise((resolve) => setTimeout(resolve, 30)); + return 'ok'; + }); + + const [a, b, c] = await Promise.all([run(), run(), run()]); + expect(a).toBe('ok'); + expect(b).toBe('ok'); + expect(c).toBe('ok'); + expect(runs).toBe(1); + + await run(); + expect(runs).toBe(2); + }); +}); diff --git a/tests/unit/whisper-spectral.spec.ts b/tests/unit/whisper-spectral.spec.ts new file mode 100644 index 0000000..3358532 --- /dev/null +++ b/tests/unit/whisper-spectral.spec.ts @@ -0,0 +1,34 @@ +import { test, expect } from '@playwright/test'; +import { buildGoertzelCoefficients, goertzelPower } from '../../src/lib/server/whisper/spectral'; + +function dftPower(samples: Float32Array, k: number): number { + const n = samples.length; + let re = 0; + let im = 0; + for (let i = 0; i < n; i += 1) { + const angle = (-2 * Math.PI * k * i) / n; + re += samples[i] * Math.cos(angle); + im += samples[i] * Math.sin(angle); + } + return (re * re) + (im * im); +} + +test.describe('whisper spectral helpers', () => { + test('goertzel power matches direct DFT for non-power-of-two frame size', () => { + const frameSize = 400; + const bins = 201; + const coeffs = buildGoertzelCoefficients(bins, frameSize); + const samples = new Float32Array(frameSize); + for (let i = 0; i < frameSize; i += 1) { + samples[i] = Math.sin((2 * Math.PI * 37 * i) / frameSize) + (0.2 * Math.cos((2 * Math.PI * 91 * i) / frameSize)); + } + + const testBins = [0, 7, 37, 91, 150, 200]; + for (const k of testBins) { + const expected = dftPower(samples, k); + const actual = goertzelPower(samples, coeffs[k]); + const rel = Math.abs(actual - expected) / Math.max(1, Math.abs(expected)); + expect(rel).toBeLessThan(1e-5); + } + }); +}); diff --git a/tests/unit/whisper-token-timestamps.spec.ts b/tests/unit/whisper-token-timestamps.spec.ts new file mode 100644 index 0000000..22ebebd --- /dev/null +++ b/tests/unit/whisper-token-timestamps.spec.ts @@ -0,0 +1,85 @@ +import { test, expect } from '@playwright/test'; +import * as ort from 'onnxruntime-node'; +import { + buildWordsFromTimestampedTokens, + extractTokenStartTimestamps, +} from '../../src/lib/server/whisper/token-timestamps'; + +test.describe('whisper token timestamp alignment', () => { + test('extracts monotonic token timestamps from cross-attention maps', () => { + const seqLen = 6; + const frames = 10; + const heads = 8; + const data = new Float32Array(1 * heads * seqLen * frames); + for (let s = 0; s < seqLen; s += 1) { + const peak = Math.min(frames - 1, s + 1); + for (let f = 0; f < frames; f += 1) { + const val = -Math.abs(f - peak); + const idx = (((0 * seqLen) + s) * frames) + f; + data[idx] = val; + } + } + + const cross = { + 'cross_attentions.0': new ort.Tensor('float32', data, [1, heads, seqLen, frames]), + }; + + const ts = extractTokenStartTimestamps({ + crossAttentions: cross, + decoderLayers: 6, + alignmentHeads: [[0, 0]], + numFrames: frames, + numInputIds: 3, + sequenceLength: seqLen, + timePrecision: 0.02, + }); + + expect(ts).toHaveLength(seqLen); + expect(ts[0]).toBe(0); + expect(ts[1]).toBe(0); + expect(ts[2]).toBe(0); + expect(ts[3]).toBeGreaterThanOrEqual(0); + expect(ts[4]).toBeGreaterThanOrEqual(ts[3]); + expect(ts[5]).toBeGreaterThanOrEqual(ts[4]); + }); + + test('builds word timings from token timestamps with punctuation merge', () => { + const tokenText: Record = { + 100: ' hello', + 101: ' world', + 102: '!', + }; + const tokenizer = { + decode(tokens: number[]) { + return tokens.map((t) => tokenText[t] ?? '').join(''); + }, + }; + + const timestampBeginTokenId = 50364; + const tokens = [ + 1, 2, 3, + timestampBeginTokenId, + 100, 101, 102, + timestampBeginTokenId + 50, + ]; + const starts = [0, 0, 0, 0, 0.1, 0.3, 0.5, 1.0]; + + const words = buildWordsFromTimestampedTokens({ + tokens, + tokenStartTimestamps: starts, + tokenizer, + eosTokenId: 50257, + promptLength: 3, + timestampBeginTokenId, + timePrecision: 0.02, + language: 'en', + }); + + expect(words.length).toBe(2); + expect(words[0].word.toLowerCase()).toContain('hello'); + expect(words[1].word.toLowerCase()).toContain('world'); + expect(words[1].word).toContain('!'); + expect(words[0].startSec).toBeGreaterThanOrEqual(0); + expect(words[1].endSec).toBeGreaterThanOrEqual(words[1].startSec); + }); +});