Spaces:

cfahlgren1
/

talking-pixel

Running

convert talking pixel to a sveltekit app that runs pivo-1.4m tts in the browser

3ec6b9e 5 days ago

9.33 kB

// Inference backends for PiVoEngine.
//  - createLiteBackend: zero-dependency. onnx-lite interpreter for encoder/decoder and the
//    hand-written vocoder (WASM SIMD or plain JS kernel), or the interpreter for everything.
//  - createOrtBackend: onnxruntime-web (wasm or webgpu execution provider) for all four graphs.

import { Session, parseModel, tensor, setAccelerator } from './onnx-lite.ts';
import { WasmGemm, wasmSimdSupported } from './accel.ts';
import { HandVocoder } from './vocoder.ts';
import type { VocoderKind } from './vocoder.ts';
import type { Backend } from './engine.ts';
import type { AssetBuffers } from './assets.ts';
import type { Tensor } from './onnx-lite.ts';

export type LiteCache = Record<string, Tensor>;
export type ModelBuffers = Pick<AssetBuffers, 'encoder' | 'decoder' | 'conditioning' | 'vocoder_step'>;

// The parts of onnxruntime-web the benchmark uses (it is loaded at runtime, never bundled).
export interface OrtTensor {
  data: unknown;
  dims: readonly number[];
  location?: string;
  dispose(): void;
}
export type OrtFeeds = Record<string, OrtTensor>;
export interface OrtSession {
  run(feeds: OrtFeeds): Promise<OrtFeeds>;
  release?(): void;
}
export interface OrtModule {
  Tensor: new (type: string, data: Float32Array | BigInt64Array, dims: number[]) => OrtTensor;
  InferenceSession: { create(model: Uint8Array, options: object): Promise<OrtSession> };
  env: { wasm: { wasmPaths?: string; numThreads?: number; simd?: boolean } };
}

const CACHE_NAMES = ['key0', 'value0', 'key1', 'value1'];
const STATE: [string, number][] = [['memory', 256], ['gru1', 160], ['gru2', 128], ['gru3', 128], ['frame_state', 164]];

function trimRows(data: Float32Array, heads: number, rows: number, dim: number, keep: number): Float32Array {
  const out = new Float32Array(heads * keep * dim);
  for (let h = 0; h < heads; h++) {
    out.set(data.subarray((h * rows + rows - keep) * dim, (h * rows + rows) * dim), h * keep * dim);
  }
  return out;
}

// vocoder: 'wasm' | 'js' | 'interpreted'
export function createLiteBackend(
  buffers: ModelBuffers,
  { vocoder = (HandVocoder.wasmAvailable() ? 'wasm' : 'js') as VocoderKind | 'interpreted' } = {},
): Backend<LiteCache> {
  // The plain-JS and interpreted variants stay pure JS so they can be benchmarked as such.
  setAccelerator(vocoder === 'wasm' && wasmSimdSupported() ? new WasmGemm() : null);
  const encoder = Session.fromBuffer(buffers.encoder);
  const decoder = Session.fromBuffer(buffers.decoder);
  const vocGraph = parseModel(buffers.vocoder_step);
  const condGraph = parseModel(buffers.conditioning);
  const i64 = (dims: number[], values: ArrayLike<number>) => tensor('i64', dims, Float64Array.from(values));

  let frameFn: Backend<LiteCache>['vocoderFrame'];
  let condFn: Backend<LiteCache>['conditioning'];
  let resetFn: () => void;
  if (vocoder === 'interpreted') {
    const condSession = new Session(condGraph);
    const stepSession = new Session(vocGraph);
    let state: LiteCache;
    resetFn = () => {
      state = Object.fromEntries(STATE.map(([n, w]) => [n, tensor('f32', [1, w])]));
    };
    condFn = (context, periods) =>
      condSession.run({ features: tensor('f32', [1, 3, 20], Float32Array.from(context)), periods: i64([1, 3], periods) }).conditioning.data as Float32Array;
    frameFn = (cond, period, audio, base) => {
      for (let sub = 0; sub < 4; sub++) {
        const r = stepSession.run({ conditioning: tensor('f32', [1, 80], cond.slice(sub * 80, sub * 80 + 80)), period: i64([1], [period]), ...state });
        for (const [n] of STATE) state[n] = r[`${n}_out`];
        audio.set(r.excitation.data as Float32Array, base + sub * 40);
      }
    };
  } else {
    const hv = new HandVocoder(vocGraph, condGraph, vocoder);
    resetFn = () => hv.reset();
    condFn = (context, periods) => hv.conditioning(context, periods);
    frameFn = (cond, period, audio, base) => {
      for (let sub = 0; sub < 4; sub++) hv.step(cond, sub * 80, period, audio, base + sub * 40);
    };
  }

  return {
    name: vocoder === 'interpreted' ? 'JS interpreter (all graphs)' : `onnx-lite + hand-written vocoder (${vocoder === 'wasm' ? 'WASM SIMD' : 'JS'})`,
    encoder(ids, voice) {
      const T = ids.length;
      const r = encoder.run({ tokens: i64([1, T], ids), voice: i64([1], [voice]), lengths: i64([1], [T]) });
      return { encoded: r.encoded.data as Float32Array, logDurations: r.log_durations.data as Float32Array };
    },
    emptyCache() {
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, tensor('f32', [1, 2, 0, 24])]));
    },
    decoder(chunk, n, position, cache) {
      const r = decoder.run({ expanded: tensor('f32', [1, n, 100], chunk), position: i64([], [position]), ...cache });
      return { features: r.features.data as Float32Array, cache: Object.fromEntries(CACHE_NAMES.map((name) => [name, r[`${name}_out`]])) };
    },
    trimCache(cache, keep) {
      const rows = cache.key0.dims[2];
      if (rows <= keep) return cache;
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, tensor('f32', [1, 2, keep, 24], trimRows(cache[n].data as Float32Array, 2, rows, 24, keep))]));
    },
    resetVocoder: resetFn,
    conditioning: condFn,
    vocoderFrame: frameFn,
  };
}

export async function createOrtBackend(ort: OrtModule, buffers: ModelBuffers, { ep = 'wasm' } = {}): Promise<Backend<OrtFeeds>> {
  const gpu = ep === 'webgpu';
  const options = { executionProviders: [ep], graphOptimizationLevel: 'all' };
  const create = (buf: Uint8Array, extra = {}) => ort.InferenceSession.create(new Uint8Array(buf), { ...options, ...extra });
  const encoder = await create(buffers.encoder);
  const decoder = await create(buffers.decoder);
  const conditioning = await create(buffers.conditioning);
  // On WebGPU keep the recurrent state on the GPU between steps; only excitation comes back.
  const stepOutputs = gpu ? Object.fromEntries(STATE.map(([n]) => [`${n}_out`, 'gpu-buffer'])) : undefined;
  const step = await create(buffers.vocoder_step, gpu ? { preferredOutputLocation: { ...stepOutputs, excitation: 'cpu' } } : {});

  const big = (values: number[]) => BigInt64Array.from(values, (v) => BigInt(v));
  // Reused input tensors: their backing arrays are rewritten in place before each run.
  const condBuf = new Float32Array(80);
  const condTensor = new ort.Tensor('float32', condBuf, [1, 80]);
  const periodBuf = new BigInt64Array(1);
  const periodTensor = new ort.Tensor('int64', periodBuf, [1]);
  const ctxBuf = new Float32Array(60);
  const ctxTensor = new ort.Tensor('float32', ctxBuf, [1, 3, 20]);
  const perBuf = new BigInt64Array(3);
  const perTensor = new ort.Tensor('int64', perBuf, [1, 3]);
  let state: OrtFeeds | undefined;

  const disposeState = () => {
    if (!state) return;
    for (const t of Object.values(state)) if (t.location === 'gpu-buffer') t.dispose();
  };

  return {
    name: `onnxruntime-web ${ep}`,
    async encoder(ids, voice) {
      const T = ids.length;
      const r = await encoder.run({
        tokens: new ort.Tensor('int64', big(ids), [1, T]),
        voice: new ort.Tensor('int64', big([voice]), [1]),
        lengths: new ort.Tensor('int64', big([T]), [1]),
      });
      return { encoded: r.encoded.data as Float32Array, logDurations: r.log_durations.data as Float32Array };
    },
    emptyCache() {
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, new ort.Tensor('float32', new Float32Array(0), [1, 2, 0, 24])]));
    },
    async decoder(chunk, n, position, cache) {
      const r = await decoder.run({
        expanded: new ort.Tensor('float32', chunk.slice(), [1, n, 100]),
        position: new ort.Tensor('int64', big([position]), []),
        ...cache,
      });
      return { features: r.features.data as Float32Array, cache: Object.fromEntries(CACHE_NAMES.map((name) => [name, r[`${name}_out`]])) };
    },
    trimCache(cache, keep) {
      const rows = cache.key0.dims[2];
      if (rows <= keep) return cache;
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, new ort.Tensor('float32', trimRows(cache[n].data as Float32Array, 2, rows, 24, keep), [1, 2, keep, 24])]));
    },
    resetVocoder() {
      disposeState();
      state = Object.fromEntries(STATE.map(([n, w]) => [n, new ort.Tensor('float32', new Float32Array(w), [1, w])]));
    },
    async conditioning(context, periods) {
      ctxBuf.set(context);
      for (let k = 0; k < 3; k++) perBuf[k] = BigInt(periods[k]);
      const r = await conditioning.run({ features: ctxTensor, periods: perTensor });
      return r.conditioning.data as Float32Array;
    },
    async vocoderFrame(cond, period, audio, base) {
      periodBuf[0] = BigInt(period);
      for (let sub = 0; sub < 4; sub++) {
        condBuf.set(cond.subarray(sub * 80, sub * 80 + 80));
        const r = await step.run({ conditioning: condTensor, period: periodTensor, ...state });
        const next: OrtFeeds = {};
        for (const [n] of STATE) next[n] = r[`${n}_out`];
        if (gpu) for (const t of Object.values(state as OrtFeeds)) if (t.location === 'gpu-buffer') t.dispose();
        state = next;
        audio.set(r.excitation.data as Float32Array, base + sub * 40);
      }
    },
    release() {
      disposeState();
      for (const s of [encoder, decoder, conditioning, step]) s.release?.();
    },
  };
}

	// Inference backends for PiVoEngine.
	// - createLiteBackend: zero-dependency. onnx-lite interpreter for encoder/decoder and the
	// hand-written vocoder (WASM SIMD or plain JS kernel), or the interpreter for everything.
	// - createOrtBackend: onnxruntime-web (wasm or webgpu execution provider) for all four graphs.

	import { Session, parseModel, tensor, setAccelerator } from './onnx-lite.ts';
	import { WasmGemm, wasmSimdSupported } from './accel.ts';
	import { HandVocoder } from './vocoder.ts';
	import type { VocoderKind } from './vocoder.ts';
	import type { Backend } from './engine.ts';
	import type { AssetBuffers } from './assets.ts';
	import type { Tensor } from './onnx-lite.ts';

	export type LiteCache = Record<string, Tensor>;
	export type ModelBuffers = Pick<AssetBuffers, 'encoder' \| 'decoder' \| 'conditioning' \| 'vocoder_step'>;

	// The parts of onnxruntime-web the benchmark uses (it is loaded at runtime, never bundled).
	export interface OrtTensor {
	data: unknown;
	dims: readonly number[];
	location?: string;
	dispose(): void;
	}
	export type OrtFeeds = Record<string, OrtTensor>;
	export interface OrtSession {
	run(feeds: OrtFeeds): Promise<OrtFeeds>;
	release?(): void;
	}
	export interface OrtModule {
	Tensor: new (type: string, data: Float32Array \| BigInt64Array, dims: number[]) => OrtTensor;
	InferenceSession: { create(model: Uint8Array, options: object): Promise<OrtSession> };
	env: { wasm: { wasmPaths?: string; numThreads?: number; simd?: boolean } };
	}

	const CACHE_NAMES = ['key0', 'value0', 'key1', 'value1'];
	const STATE: [string, number][] = [['memory', 256], ['gru1', 160], ['gru2', 128], ['gru3', 128], ['frame_state', 164]];

	function trimRows(data: Float32Array, heads: number, rows: number, dim: number, keep: number): Float32Array {
	const out = new Float32Array(heads * keep * dim);
	for (let h = 0; h < heads; h++) {
	out.set(data.subarray((h * rows + rows - keep) * dim, (h * rows + rows) * dim), h * keep * dim);
	}
	return out;
	}

	// vocoder: 'wasm' \| 'js' \| 'interpreted'
	export function createLiteBackend(
	buffers: ModelBuffers,
	{ vocoder = (HandVocoder.wasmAvailable() ? 'wasm' : 'js') as VocoderKind \| 'interpreted' } = {},
	): Backend<LiteCache> {
	// The plain-JS and interpreted variants stay pure JS so they can be benchmarked as such.
	setAccelerator(vocoder === 'wasm' && wasmSimdSupported() ? new WasmGemm() : null);
	const encoder = Session.fromBuffer(buffers.encoder);
	const decoder = Session.fromBuffer(buffers.decoder);
	const vocGraph = parseModel(buffers.vocoder_step);
	const condGraph = parseModel(buffers.conditioning);
	const i64 = (dims: number[], values: ArrayLike<number>) => tensor('i64', dims, Float64Array.from(values));

	let frameFn: Backend<LiteCache>['vocoderFrame'];
	let condFn: Backend<LiteCache>['conditioning'];
	let resetFn: () => void;
	if (vocoder === 'interpreted') {
	const condSession = new Session(condGraph);
	const stepSession = new Session(vocGraph);
	let state: LiteCache;
	resetFn = () => {
	state = Object.fromEntries(STATE.map(([n, w]) => [n, tensor('f32', [1, w])]));
	};
	condFn = (context, periods) =>
	condSession.run({ features: tensor('f32', [1, 3, 20], Float32Array.from(context)), periods: i64([1, 3], periods) }).conditioning.data as Float32Array;
	frameFn = (cond, period, audio, base) => {
	for (let sub = 0; sub < 4; sub++) {
	const r = stepSession.run({ conditioning: tensor('f32', [1, 80], cond.slice(sub * 80, sub * 80 + 80)), period: i64([1], [period]), ...state });
	for (const [n] of STATE) state[n] = r[`${n}_out`];
	audio.set(r.excitation.data as Float32Array, base + sub * 40);
	}
	};
	} else {
	const hv = new HandVocoder(vocGraph, condGraph, vocoder);
	resetFn = () => hv.reset();
	condFn = (context, periods) => hv.conditioning(context, periods);
	frameFn = (cond, period, audio, base) => {
	for (let sub = 0; sub < 4; sub++) hv.step(cond, sub * 80, period, audio, base + sub * 40);
	};
	}

	return {
	name: vocoder === 'interpreted' ? 'JS interpreter (all graphs)' : `onnx-lite + hand-written vocoder (${vocoder === 'wasm' ? 'WASM SIMD' : 'JS'})`,
	encoder(ids, voice) {
	const T = ids.length;
	const r = encoder.run({ tokens: i64([1, T], ids), voice: i64([1], [voice]), lengths: i64([1], [T]) });
	return { encoded: r.encoded.data as Float32Array, logDurations: r.log_durations.data as Float32Array };
	},
	emptyCache() {
	return Object.fromEntries(CACHE_NAMES.map((n) => [n, tensor('f32', [1, 2, 0, 24])]));
	},
	decoder(chunk, n, position, cache) {
	const r = decoder.run({ expanded: tensor('f32', [1, n, 100], chunk), position: i64([], [position]), ...cache });
	return { features: r.features.data as Float32Array, cache: Object.fromEntries(CACHE_NAMES.map((name) => [name, r[`${name}_out`]])) };
	},
	trimCache(cache, keep) {
	const rows = cache.key0.dims[2];
	if (rows <= keep) return cache;
	return Object.fromEntries(CACHE_NAMES.map((n) => [n, tensor('f32', [1, 2, keep, 24], trimRows(cache[n].data as Float32Array, 2, rows, 24, keep))]));
	},
	resetVocoder: resetFn,
	conditioning: condFn,
	vocoderFrame: frameFn,
	};
	}

	export async function createOrtBackend(ort: OrtModule, buffers: ModelBuffers, { ep = 'wasm' } = {}): Promise<Backend<OrtFeeds>> {
	const gpu = ep === 'webgpu';
	const options = { executionProviders: [ep], graphOptimizationLevel: 'all' };
	const create = (buf: Uint8Array, extra = {}) => ort.InferenceSession.create(new Uint8Array(buf), { ...options, ...extra });
	const encoder = await create(buffers.encoder);
	const decoder = await create(buffers.decoder);
	const conditioning = await create(buffers.conditioning);
	// On WebGPU keep the recurrent state on the GPU between steps; only excitation comes back.
	const stepOutputs = gpu ? Object.fromEntries(STATE.map(([n]) => [`${n}_out`, 'gpu-buffer'])) : undefined;
	const step = await create(buffers.vocoder_step, gpu ? { preferredOutputLocation: { ...stepOutputs, excitation: 'cpu' } } : {});

	const big = (values: number[]) => BigInt64Array.from(values, (v) => BigInt(v));
	// Reused input tensors: their backing arrays are rewritten in place before each run.
	const condBuf = new Float32Array(80);
	const condTensor = new ort.Tensor('float32', condBuf, [1, 80]);
	const periodBuf = new BigInt64Array(1);
	const periodTensor = new ort.Tensor('int64', periodBuf, [1]);
	const ctxBuf = new Float32Array(60);
	const ctxTensor = new ort.Tensor('float32', ctxBuf, [1, 3, 20]);
	const perBuf = new BigInt64Array(3);
	const perTensor = new ort.Tensor('int64', perBuf, [1, 3]);
	let state: OrtFeeds \| undefined;

	const disposeState = () => {
	if (!state) return;
	for (const t of Object.values(state)) if (t.location === 'gpu-buffer') t.dispose();
	};

	return {
	name: `onnxruntime-web ${ep}`,
	async encoder(ids, voice) {
	const T = ids.length;
	const r = await encoder.run({
	tokens: new ort.Tensor('int64', big(ids), [1, T]),
	voice: new ort.Tensor('int64', big([voice]), [1]),
	lengths: new ort.Tensor('int64', big([T]), [1]),
	});
	return { encoded: r.encoded.data as Float32Array, logDurations: r.log_durations.data as Float32Array };
	},
	emptyCache() {
	return Object.fromEntries(CACHE_NAMES.map((n) => [n, new ort.Tensor('float32', new Float32Array(0), [1, 2, 0, 24])]));
	},
	async decoder(chunk, n, position, cache) {
	const r = await decoder.run({
	expanded: new ort.Tensor('float32', chunk.slice(), [1, n, 100]),
	position: new ort.Tensor('int64', big([position]), []),
	...cache,
	});
	return { features: r.features.data as Float32Array, cache: Object.fromEntries(CACHE_NAMES.map((name) => [name, r[`${name}_out`]])) };
	},
	trimCache(cache, keep) {
	const rows = cache.key0.dims[2];
	if (rows <= keep) return cache;
	return Object.fromEntries(CACHE_NAMES.map((n) => [n, new ort.Tensor('float32', trimRows(cache[n].data as Float32Array, 2, rows, 24, keep), [1, 2, keep, 24])]));
	},
	resetVocoder() {
	disposeState();
	state = Object.fromEntries(STATE.map(([n, w]) => [n, new ort.Tensor('float32', new Float32Array(w), [1, w])]));
	},
	async conditioning(context, periods) {
	ctxBuf.set(context);
	for (let k = 0; k < 3; k++) perBuf[k] = BigInt(periods[k]);
	const r = await conditioning.run({ features: ctxTensor, periods: perTensor });
	return r.conditioning.data as Float32Array;
	},
	async vocoderFrame(cond, period, audio, base) {
	periodBuf[0] = BigInt(period);
	for (let sub = 0; sub < 4; sub++) {
	condBuf.set(cond.subarray(sub * 80, sub * 80 + 80));
	const r = await step.run({ conditioning: condTensor, period: periodTensor, ...state });
	const next: OrtFeeds = {};
	for (const [n] of STATE) next[n] = r[`${n}_out`];
	if (gpu) for (const t of Object.values(state as OrtFeeds)) if (t.location === 'gpu-buffer') t.dispose();
	state = next;
	audio.set(r.excitation.data as Float32Array, base + sub * 40);
	}
	},
	release() {
	disposeState();
	for (const s of [encoder, decoder, conditioning, step]) s.release?.();
	},
	};
	}