wisp/src/lib/transcription/engineImpl.web.ts

// WEB-ONLY transcription engine, backed by transformers.js (@huggingface/
// transformers) running Whisper in the browser (WebGPU when available, else
// multi-threaded WASM).
//
// WHY WE LOAD IT FROM A CDN AT RUNTIME (not a static import):
// transformers.js depends on onnxruntime-web, which uses a *computed* dynamic
// import (`import(/*webpackIgnore*/ a)`) and ships WASM — Metro (Expo's web
// bundler) cannot statically bundle either and fails the build. So we never let
// Metro see the package: we load the ESM build from a CDN at runtime via a
// dynamic import hidden behind `new Function` (so Metro's static analyzer can't
// trip over it). The browser resolves it natively. This keeps the JS bundle
// small and is the standard way to run transformers.js under Metro/Expo web.
//
// NOTE: the page must be cross-origin isolated (COOP + COEP) for multi-threaded
// WASM; we use COEP: credentialless so the CDN script and the Hugging Face model
// files (CORS-enabled) load without requiring CORP headers. See docker/nginx.conf.
//
// This module is web-only and is NEVER imported by any vitest test.

import { MODELS, recommendModel } from '../models/catalog';
import type { Backend, ModelId, PcmAudio, Segment, TranscribeOptions } from '../types';
import type { EngineCapabilities, TranscriptionEngine } from './engine';

// Pin the transformers.js version we load at runtime.
const TRANSFORMERS_CDN = 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0';

// `new Function` hides the dynamic import() specifier from Metro's bundler so it
// never tries to resolve/transform transformers.js or onnxruntime-web.
const runtimeImport = new Function('u', 'return import(u)') as (u: string) => Promise<TransformersModule>;

// Minimal structural types for the bits of transformers.js we use.
interface AsrChunk {
  timestamp: [number, number | null];
  text: string;
}
interface AsrOutput {
  text: string;
  chunks?: AsrChunk[];
}
type AsrPipeline = (audio: Float32Array, opts: Record<string, unknown>) => Promise<AsrOutput>;
interface PipelineOptions {
  device?: string;
  // A single dtype for all sub-models, or a per-file map (encoder_model,
  // decoder_model_merged, …) — Whisper's decoder needs special handling on WASM.
  dtype?: string | Record<string, string>;
  progress_callback?: (e: { status?: string; progress?: number }) => void;
}
interface TransformersModule {
  pipeline: (task: string, model: string, opts?: PipelineOptions) => Promise<AsrPipeline>;
  env: { allowLocalModels: boolean };
}

let libPromise: Promise<TransformersModule> | null = null;
async function lib(): Promise<TransformersModule> {
  if (!libPromise) {
    libPromise = runtimeImport(TRANSFORMERS_CDN).then((m) => {
      // Never read models off the local filesystem in the browser.
      m.env.allowLocalModels = false;
      return m;
    });
  }
  return libPromise;
}

/** Loaded ASR pipelines, keyed by model id. */
const loaded = new Map<ModelId, AsrPipeline>();
let cachedBackend: Backend | undefined;

/**
 * Mobile browsers expose WebGPU but their drivers are flaky for sustained ML
 * inference — fp16 Whisper tends to crash the GPU process after a chunk or two.
 * On mobile we deliberately use (cross-origin-isolated, multi-threaded) WASM,
 * which is slower but reliable. Desktop keeps WebGPU.
 */
function isMobileWeb(): boolean {
  if (typeof navigator === 'undefined') return false;
  const ua = navigator.userAgent || '';
  return /Android|iPhone|iPad|iPod|Mobile|Silk|Kindle/i.test(ua);
}

async function detectWebGpu(): Promise<boolean> {
  try {
    if (isMobileWeb()) return false;
    if (typeof navigator === 'undefined' || !('gpu' in navigator)) return false;
    const gpu = (navigator as { gpu?: { requestAdapter(): Promise<unknown> } }).gpu;
    const adapter = await gpu?.requestAdapter();
    return adapter != null;
  } catch {
    return false;
  }
}

async function resolveBackend(): Promise<Backend> {
  if (cachedBackend) return cachedBackend;
  cachedBackend = (await detectWebGpu()) ? 'webgpu' : 'wasm';
  return cachedBackend;
}

export const engine: TranscriptionEngine = {
  platform: 'web',

  async capabilities(): Promise<EngineCapabilities> {
    const backend = await resolveBackend();
    return {
      backend,
      supportsLiveMic: false,
      maxRecommendedModel: recommendModel({ backend }),
    };
  },

  async loadModel(modelId: ModelId, onProgress?: (p: number) => void): Promise<void> {
    if (loaded.has(modelId)) return;
    const { pipeline } = await lib();
    const webgpu = (await resolveBackend()) === 'webgpu';
    const asr = await pipeline('automatic-speech-recognition', MODELS[modelId].webRepo, {
      device: webgpu ? 'webgpu' : 'wasm',
      // WebGPU: fp16 (fast, small). WASM: an explicit fp32 decoder. The default
      // quantized decoders (q8/q4) use `MatMulNBits` ops that onnxruntime-web
      // (transformers.js 4.2.0) FAILS to load on the WASM backend with
      // "Missing required scale" — verified across Xenova + onnx-community
      // repos. fp32 avoids those ops and actually runs on any CPU (the WASM
      // path is what mobile and no-WebGPU devices use). Larger download, but it
      // works; q8 simply does not load on WASM here.
      dtype: webgpu
        ? 'fp16'
        : { encoder_model: 'fp32', decoder_model_merged: 'fp32' },
      progress_callback: (e) => {
        if (e.status === 'progress' && e.progress != null) onProgress?.(e.progress / 100);
      },
    });
    loaded.set(modelId, asr);
  },

  isModelLoaded(modelId: ModelId): boolean {
    return loaded.has(modelId);
  },

  async transcribeChunk(audio: PcmAudio, opts: TranscribeOptions): Promise<Segment[]> {
    const asr = loaded.get(opts.modelId);
    if (!asr) throw new Error(`Model "${opts.modelId}" is not loaded; call loadModel() first.`);

    // English-only Whisper models (".en") reject `language`/`task` — transformers.js
    // throws "Cannot specify `task` or `language` for an English-only model" if
    // either is passed. Only multilingual models accept them (translation also
    // requires a multilingual model). So we add those keys ONLY when multilingual,
    // and never pass an explicit `undefined` language (which trips the same check).
    const genOpts: Record<string, unknown> = {
      return_timestamps: true,
      // One window at a time; 30s matches Whisper's frame so it won't re-chunk.
      chunk_length_s: 30,
    };
    if (MODELS[opts.modelId].multilingual) {
      genOpts.task = opts.translate ? 'translate' : 'transcribe';
      if (opts.language) genOpts.language = opts.language;
    }

    const out = await asr(audio.samples, genOpts);

    // Whisper sometimes omits the END timestamp of the last utterance in a
    // window (transformers.js yields timestamp[1] === null). Don't collapse that
    // to a zero-length [t, t] segment (it breaks stitch seams + citation spans):
    // estimate the end from the next chunk's start, else the window's duration.
    const chunks = out.chunks ?? [];
    const chunkDurationSec = audio.samples.length / audio.sampleRate;
    const segments: Segment[] = [];
    for (let i = 0; i < chunks.length; i++) {
      const c = chunks[i]!;
      const [start, rawEnd] = c.timestamp;
      if (start == null) continue;
      const text = c.text.trim();
      if (text.length === 0) continue;
      let end = rawEnd;
      if (end == null) {
        const nextStart = chunks[i + 1]?.timestamp[0];
        end = nextStart != null && nextStart > start ? nextStart : chunkDurationSec;
      }
      if (end < start) end = chunkDurationSec;
      segments.push({ start, end, text });
    }
    return segments;
  },
};