Spaces:

OpenVoiceOS
/

wakehubert-wakewords-space

Running

Load models, results and test clips at runtime from the model repository and the plugin repository

afd62aa verified about 1 month ago

7.55 kB

// Wake-word scoring, a port of the scoring path of ovos-ww-plugin-wakeforge
// (WakeForgeHotwordPlugin.update / found_wake_word, OnnxWindowedWakeWord.push with
// isolated windows, PredictionSmoother with the ready-model defaults).
// The worker's URL carries the page's __sign parameter; every file of the Space it fetches carries it too, since
// a private Space's iframe may have no cookie to authenticate them. Absolute URLs do not carry it.
const MODEL_BASE = "https://huggingface.co/OpenVoiceOS/wakehubert-wakewords/resolve/main/";
const SIGN = new URL(import.meta.url).searchParams.get("__sign");
const FETCH_TIMEOUT_MS = 90_000;

function signed(path) {
  if (/^https?:/.test(path)) return path;
  const url = new URL("./" + path, import.meta.url);
  url.search = "";
  if (SIGN) url.searchParams.set("__sign", SIGN);
  return url.href;
}

async function fetchBytes(path) {
  const ctl = new AbortController();
  const timer = setTimeout(() => ctl.abort(), FETCH_TIMEOUT_MS);
  try {
    const r = await fetch(signed(path), {signal: ctl.signal});
    if (!r.ok) throw new Error(`${path}: HTTP ${r.status}${r.status === 401 || r.status === 403 ? " (not signed in to this Space)" : ""}`);
    return new Uint8Array(await r.arrayBuffer());
  } catch (err) {
    if (err.name === "AbortError") throw new Error(`${path}: no response after ${FETCH_TIMEOUT_MS / 1000} s`);
    if (err instanceof TypeError) throw new Error(`${path}: ${err.message}`);
    throw err;
  } finally {
    clearTimeout(timer);
  }
}

let ort = null;
let ortLoading = null;

function loadOrt() {
  ortLoading ??= (async () => {
    const mjs = "vendor/ort/ort.wasm.min.mjs";
    let mod;
    try {
      mod = await import(signed(mjs));
    } catch (err) {
      throw new Error(`${mjs} could not be loaded: ${err.message}`);
    }
    mod.env.wasm.numThreads = 1;
    mod.env.wasm.wasmBinary = await fetchBytes("vendor/ort/ort-wasm-simd-threaded.wasm");
    mod.env.wasm.wasmPaths = {mjs: signed("vendor/ort/ort-wasm-simd-threaded.mjs"),
                              wasm: signed("vendor/ort/ort-wasm-simd-threaded.wasm")};
    ort = mod;
  })();
  ortLoading.catch(() => { ortLoading = null; });
  return ortLoading;
}

const SAMPLE_RATE = 16000;
const HOP = 320;                   // featurizer hop (config.json output.hop_samples)
const WINDOW_FRAMES = 75;          // every head's window_frames metadata
const WINDOW = WINDOW_FRAMES * HOP; // 24,000 samples, 1.5 s
const BLOCK = 1280;                // 80 ms, already a multiple of the hop
const DEBOUNCE_SEC = 2.0;
const SMOOTHER_FRAME_RATE = 10.0;  // PredictionSmoother's default frame_rate_hz
const FEATURIZER_FILES = {"wakehubert-int8": "featurizer/wakehubert_int8.onnx",
                          "wakehubert": "featurizer/wakehubert.onnx"};

const sessions = new Map();
let models = null;

function session(path) {
  if (!sessions.has(path)) {
    const p = (async () => {
      await loadOrt();
      const bytes = await fetchBytes(path);
      try {
        return await ort.InferenceSession.create(bytes, {executionProviders: ["wasm"]});
      } catch (err) {
        throw new Error(`${path} could not be opened: ${err.message}`);
      }
    })();
    p.catch(() => sessions.delete(path));
    sessions.set(path, p);
  }
  return sessions.get(path);
}

async function modelEntry(name) {
  if (!models) {
    const json = JSON.parse(new TextDecoder().decode(await fetchBytes(`${MODEL_BASE}models.json`)));
    models = Object.fromEntries(json.models.map((m) => [m.name, m]));
  }
  const entry = models[name];
  if (!entry) throw new Error(`unknown model ${name}`);
  return entry;
}

class Detector {
  constructor(feat, head, threshold) {
    this.feat = feat;
    this.head = head;
    this.threshold = threshold;
    this.debounceFrames = Math.trunc(DEBOUNCE_SEC * SMOOTHER_FRAME_RATE);
    this.quietSamples = 0;
    this.reset();
  }

  static async create(name, threshold) {
    const entry = await modelEntry(name);
    if (!FEATURIZER_FILES[entry.featurizer]) throw new Error(`${name}: featurizer ${entry.featurizer} is not bundled`);
    const [feat, head] = await Promise.all([session(FEATURIZER_FILES[entry.featurizer]),
                                            session(MODEL_BASE + entry.file)]);
    return new Detector(feat, head, threshold ?? entry.default_threshold);
  }

  // found_wake_word() -> reset(): smoother reset and the engine's window zeroed.
  // The plugin's quiet period survives it.
  reset() {
    this.audio = new Float32Array(WINDOW);
    this.consecutive = 0;
    this.framesSinceTrigger = this.debounceFrames;
  }

  async score() {
    const wave = new ort.Tensor("float32", this.audio.slice(), [1, WINDOW]);
    const f = await this.feat.run({[this.feat.inputNames[0]]: wave});
    const features = f[this.feat.outputNames[0]];
    const h = await this.head.run({[this.head.inputNames[0]]: features});
    const logit = Number(h[this.head.outputNames[0]].data[0]);
    features.dispose?.();
    return 1.0 / (1.0 + Math.exp(-logit));
  }

  // One 1280-sample block of int16-valued audio scaled by 1/32768.
  async push(block) {
    this.audio.copyWithin(0, BLOCK);
    this.audio.set(block, WINDOW - BLOCK);
    const score = await this.score();
    let detected = false;
    if (this.quietSamples > 0) {
      this.quietSamples -= BLOCK;
    } else {
      // PredictionSmoother.update with method "max", window 1, patience 1.
      this.framesSinceTrigger += 1;
      this.consecutive = score >= this.threshold ? this.consecutive + 1 : 0;
      if (this.consecutive >= 1 && this.framesSinceTrigger >= this.debounceFrames) {
        this.framesSinceTrigger = 0;
        this.consecutive = 0;
        detected = true;
        this.quietSamples = Math.trunc(DEBOUNCE_SEC * SAMPLE_RATE);
      }
    }
    if (detected) this.reset();
    return {score, detected};
  }
}

let live = null;
let chain = Promise.resolve();

function serial(fn) {
  const p = chain.then(fn);
  chain = p.catch(() => {});
  return p;
}

async function handle(msg) {
  switch (msg.type) {
    case "load": {
      const t0 = performance.now();
      await Detector.create(msg.model);
      return {loadMs: performance.now() - t0};
    }
    case "live-start":
      live = await Detector.create(msg.model, msg.threshold);
      return {};
    case "live-stop":
      live = null;
      return {};
    case "threshold":
      if (live) live.threshold = msg.threshold;
      return {};
    case "block": {
      if (!live) return {skipped: true};
      const t0 = performance.now();
      const r = await live.push(msg.samples);
      return {...r, ms: performance.now() - t0, t: msg.t};
    }
    case "clip": {
      const det = await Detector.create(msg.model, msg.threshold);
      const n = Math.floor(msg.samples.length / BLOCK);
      const scores = new Float64Array(n);
      const detections = [];
      const t0 = performance.now();
      for (let i = 0; i < n; i++) {
        const r = await det.push(msg.samples.subarray(i * BLOCK, (i + 1) * BLOCK));
        scores[i] = r.score;
        if (r.detected) detections.push(i);
      }
      return {scores, detections, ms: performance.now() - t0, threshold: det.threshold};
    }
    default:
      throw new Error(`unknown message ${msg.type}`);
  }
}

self.onmessage = (e) => {
  const msg = e.data;
  serial(() => handle(msg)).then(
    (result) => self.postMessage({id: msg.id, type: msg.type, ok: true, ...result}),
    (err) => self.postMessage({id: msg.id, type: msg.type, ok: false, error: String(err?.message || err)}));
};

	// Wake-word scoring, a port of the scoring path of ovos-ww-plugin-wakeforge
	// (WakeForgeHotwordPlugin.update / found_wake_word, OnnxWindowedWakeWord.push with
	// isolated windows, PredictionSmoother with the ready-model defaults).
	// The worker's URL carries the page's __sign parameter; every file of the Space it fetches carries it too, since
	// a private Space's iframe may have no cookie to authenticate them. Absolute URLs do not carry it.
	const MODEL_BASE = "https://huggingface.co/OpenVoiceOS/wakehubert-wakewords/resolve/main/";
	const SIGN = new URL(import.meta.url).searchParams.get("__sign");
	const FETCH_TIMEOUT_MS = 90_000;

	function signed(path) {
	if (/^https?:/.test(path)) return path;
	const url = new URL("./" + path, import.meta.url);
	url.search = "";
	if (SIGN) url.searchParams.set("__sign", SIGN);
	return url.href;
	}

	async function fetchBytes(path) {
	const ctl = new AbortController();
	const timer = setTimeout(() => ctl.abort(), FETCH_TIMEOUT_MS);
	try {
	const r = await fetch(signed(path), {signal: ctl.signal});
	if (!r.ok) throw new Error(`${path}: HTTP ${r.status}${r.status === 401 \|\| r.status === 403 ? " (not signed in to this Space)" : ""}`);
	return new Uint8Array(await r.arrayBuffer());
	} catch (err) {
	if (err.name === "AbortError") throw new Error(`${path}: no response after ${FETCH_TIMEOUT_MS / 1000} s`);
	if (err instanceof TypeError) throw new Error(`${path}: ${err.message}`);
	throw err;
	} finally {
	clearTimeout(timer);
	}
	}

	let ort = null;
	let ortLoading = null;

	function loadOrt() {
	ortLoading ??= (async () => {
	const mjs = "vendor/ort/ort.wasm.min.mjs";
	let mod;
	try {
	mod = await import(signed(mjs));
	} catch (err) {
	throw new Error(`${mjs} could not be loaded: ${err.message}`);
	}
	mod.env.wasm.numThreads = 1;
	mod.env.wasm.wasmBinary = await fetchBytes("vendor/ort/ort-wasm-simd-threaded.wasm");
	mod.env.wasm.wasmPaths = {mjs: signed("vendor/ort/ort-wasm-simd-threaded.mjs"),
	wasm: signed("vendor/ort/ort-wasm-simd-threaded.wasm")};
	ort = mod;
	})();
	ortLoading.catch(() => { ortLoading = null; });
	return ortLoading;
	}

	const SAMPLE_RATE = 16000;
	const HOP = 320; // featurizer hop (config.json output.hop_samples)
	const WINDOW_FRAMES = 75; // every head's window_frames metadata
	const WINDOW = WINDOW_FRAMES * HOP; // 24,000 samples, 1.5 s
	const BLOCK = 1280; // 80 ms, already a multiple of the hop
	const DEBOUNCE_SEC = 2.0;
	const SMOOTHER_FRAME_RATE = 10.0; // PredictionSmoother's default frame_rate_hz
	const FEATURIZER_FILES = {"wakehubert-int8": "featurizer/wakehubert_int8.onnx",
	"wakehubert": "featurizer/wakehubert.onnx"};

	const sessions = new Map();
	let models = null;

	function session(path) {
	if (!sessions.has(path)) {
	const p = (async () => {
	await loadOrt();
	const bytes = await fetchBytes(path);
	try {
	return await ort.InferenceSession.create(bytes, {executionProviders: ["wasm"]});
	} catch (err) {
	throw new Error(`${path} could not be opened: ${err.message}`);
	}
	})();
	p.catch(() => sessions.delete(path));
	sessions.set(path, p);
	}
	return sessions.get(path);
	}

	async function modelEntry(name) {
	if (!models) {
	const json = JSON.parse(new TextDecoder().decode(await fetchBytes(`${MODEL_BASE}models.json`)));
	models = Object.fromEntries(json.models.map((m) => [m.name, m]));
	}
	const entry = models[name];
	if (!entry) throw new Error(`unknown model ${name}`);
	return entry;
	}

	class Detector {
	constructor(feat, head, threshold) {
	this.feat = feat;
	this.head = head;
	this.threshold = threshold;
	this.debounceFrames = Math.trunc(DEBOUNCE_SEC * SMOOTHER_FRAME_RATE);
	this.quietSamples = 0;
	this.reset();
	}

	static async create(name, threshold) {
	const entry = await modelEntry(name);
	if (!FEATURIZER_FILES[entry.featurizer]) throw new Error(`${name}: featurizer ${entry.featurizer} is not bundled`);
	const [feat, head] = await Promise.all([session(FEATURIZER_FILES[entry.featurizer]),
	session(MODEL_BASE + entry.file)]);
	return new Detector(feat, head, threshold ?? entry.default_threshold);
	}

	// found_wake_word() -> reset(): smoother reset and the engine's window zeroed.
	// The plugin's quiet period survives it.
	reset() {
	this.audio = new Float32Array(WINDOW);
	this.consecutive = 0;
	this.framesSinceTrigger = this.debounceFrames;
	}

	async score() {
	const wave = new ort.Tensor("float32", this.audio.slice(), [1, WINDOW]);
	const f = await this.feat.run({[this.feat.inputNames[0]]: wave});
	const features = f[this.feat.outputNames[0]];
	const h = await this.head.run({[this.head.inputNames[0]]: features});
	const logit = Number(h[this.head.outputNames[0]].data[0]);
	features.dispose?.();
	return 1.0 / (1.0 + Math.exp(-logit));
	}

	// One 1280-sample block of int16-valued audio scaled by 1/32768.
	async push(block) {
	this.audio.copyWithin(0, BLOCK);
	this.audio.set(block, WINDOW - BLOCK);
	const score = await this.score();
	let detected = false;
	if (this.quietSamples > 0) {
	this.quietSamples -= BLOCK;
	} else {
	// PredictionSmoother.update with method "max", window 1, patience 1.
	this.framesSinceTrigger += 1;
	this.consecutive = score >= this.threshold ? this.consecutive + 1 : 0;
	if (this.consecutive >= 1 && this.framesSinceTrigger >= this.debounceFrames) {
	this.framesSinceTrigger = 0;
	this.consecutive = 0;
	detected = true;
	this.quietSamples = Math.trunc(DEBOUNCE_SEC * SAMPLE_RATE);
	}
	}
	if (detected) this.reset();
	return {score, detected};
	}
	}

	let live = null;
	let chain = Promise.resolve();

	function serial(fn) {
	const p = chain.then(fn);
	chain = p.catch(() => {});
	return p;
	}

	async function handle(msg) {
	switch (msg.type) {
	case "load": {
	const t0 = performance.now();
	await Detector.create(msg.model);
	return {loadMs: performance.now() - t0};
	}
	case "live-start":
	live = await Detector.create(msg.model, msg.threshold);
	return {};
	case "live-stop":
	live = null;
	return {};
	case "threshold":
	if (live) live.threshold = msg.threshold;
	return {};
	case "block": {
	if (!live) return {skipped: true};
	const t0 = performance.now();
	const r = await live.push(msg.samples);
	return {...r, ms: performance.now() - t0, t: msg.t};
	}
	case "clip": {
	const det = await Detector.create(msg.model, msg.threshold);
	const n = Math.floor(msg.samples.length / BLOCK);
	const scores = new Float64Array(n);
	const detections = [];
	const t0 = performance.now();
	for (let i = 0; i < n; i++) {
	const r = await det.push(msg.samples.subarray(i * BLOCK, (i + 1) * BLOCK));
	scores[i] = r.score;
	if (r.detected) detections.push(i);
	}
	return {scores, detections, ms: performance.now() - t0, threshold: det.threshold};
	}
	default:
	throw new Error(`unknown message ${msg.type}`);
	}
	}

	self.onmessage = (e) => {
	const msg = e.data;
	serial(() => handle(msg)).then(
	(result) => self.postMessage({id: msg.id, type: msg.type, ok: true, ...result}),
	(err) => self.postMessage({id: msg.id, type: msg.type, ok: false, error: String(err?.message \|\| err)}));
	};