@ensomi/runtime 0.0.0 → 0.1.0

package/pkg/ensomi_generate_wasm_bg.wasm.d.ts

@@ -0,0 +1,25 @@
+/* tslint:disable */
+/* eslint-disable */
+export const memory: WebAssembly.Memory;
+export const __wbg_replaydraws_free: (a: number, b: number) => void;
+export const __wbg_wasmgenerator_free: (a: number, b: number) => void;
+export const __wbg_wasmmodels_free: (a: number, b: number) => void;
+export const replaydraws_add: (a: number, b: number, c: number, d: number, e: number, f: any, g: number) => void;
+export const replaydraws_new: () => number;
+export const replaydraws_recorded: (a: number) => number;
+export const start: () => void;
+export const wasmgenerator_new: (a: number, b: any, c: number, d: number, e: number, f: number, g: number, h: number, i: number, j: bigint, k: number, l: number, m: number, n: number, o: number, p: any, q: number) => [number, number, number];
+export const wasmgenerator_step: (a: number, b: any) => [number, number, number];
+export const wasmgenerator_work: (a: number) => number;
+export const wasmmodels_capabilities: (a: number) => [number, number];
+export const wasmmodels_new: (a: number, b: number, c: any) => [number, number, number];
+export const wasmmodels_requiredFiles: (a: number, b: number) => [number, number, number];
+export const wasmmodels_resamplerFiles: (a: number) => any;
+export const __wbindgen_malloc: (a: number, b: number) => number;
+export const __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
+export const __wbindgen_exn_store: (a: number) => void;
+export const __externref_table_alloc: () => number;
+export const __wbindgen_externrefs: WebAssembly.Table;
+export const __externref_table_dealloc: (a: number) => void;
+export const __wbindgen_free: (a: number, b: number, c: number) => void;
+export const __wbindgen_start: () => void;

package/resampler/resample-11025.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 11025,
+  "target_rate": 22050,
+  "numerator": 1,
+  "phases": 2,
+  "radius": 128,
+  "dtype": "float64",
+  "shape": [
+    2,
+    257
+  ],
+  "file": "resample-11025.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 192,
+  "boundary_file": "resample-11025-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-16000.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 16000,
+  "target_rate": 22050,
+  "numerator": 320,
+  "phases": 441,
+  "radius": 128,
+  "dtype": "float64",
+  "shape": [
+    441,
+    257
+  ],
+  "file": "resample-16000.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 221,
+  "boundary_file": "resample-16000-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-192000.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 192000,
+  "target_rate": 22050,
+  "numerator": 1280,
+  "phases": 147,
+  "radius": 1115,
+  "dtype": "float64",
+  "shape": [
+    147,
+    2231
+  ],
+  "file": "resample-192000.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 2230,
+  "boundary_file": "resample-192000-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-24000.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 24000,
+  "target_rate": 22050,
+  "numerator": 160,
+  "phases": 147,
+  "radius": 140,
+  "dtype": "float64",
+  "shape": [
+    147,
+    281
+  ],
+  "file": "resample-24000.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 280,
+  "boundary_file": "resample-24000-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-32000.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 32000,
+  "target_rate": 22050,
+  "numerator": 640,
+  "phases": 441,
+  "radius": 186,
+  "dtype": "float64",
+  "shape": [
+    441,
+    373
+  ],
+  "file": "resample-32000.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 372,
+  "boundary_file": "resample-32000-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-44100.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 44100,
+  "target_rate": 22050,
+  "numerator": 2,
+  "phases": 1,
+  "radius": 256,
+  "dtype": "float64",
+  "shape": [
+    1,
+    513
+  ],
+  "file": "resample-44100.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 512,
+  "boundary_file": "resample-44100-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-48000.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 48000,
+  "target_rate": 22050,
+  "numerator": 320,
+  "phases": 147,
+  "radius": 279,
+  "dtype": "float64",
+  "shape": [
+    147,
+    559
+  ],
+  "file": "resample-48000.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 558,
+  "boundary_file": "resample-48000-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-88200.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 88200,
+  "target_rate": 22050,
+  "numerator": 4,
+  "phases": 1,
+  "radius": 512,
+  "dtype": "float64",
+  "shape": [
+    1,
+    1025
+  ],
+  "file": "resample-88200.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 1024,
+  "boundary_file": "resample-88200-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/resampler/resample-96000.json

@@ -0,0 +1,18 @@
+{
+  "sample_rate": 96000,
+  "target_rate": 22050,
+  "numerator": 640,
+  "phases": 147,
+  "radius": 558,
+  "dtype": "float64",
+  "shape": [
+    147,
+    1117
+  ],
+  "file": "resample-96000.bin",
+  "boundary_frames": 128,
+  "boundary_samples": 1116,
+  "boundary_file": "resample-96000-boundary.bin",
+  "algorithm": "soxr HQ polyphase FIR",
+  "soxr_version": "1.1.0"
+}

package/src/core.js

@@ -0,0 +1,40 @@
+import init, * as bindings from '../pkg/ensomi_generate_wasm.js';
+
+let loaded, instance = 0;
+
+/** Reads a `file:` URL in Node, whose `fetch` cannot read files. */
+export async function readFile(url) {
+    const bytes = await process.getBuiltinModule('node:fs/promises').readFile(url);
+    return new Uint8Array(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+}
+
+/**
+ * Instantiates the Wasm core once. `wasm` is the module's bytes, a `WebAssembly.Module` or a URL; the
+ * default is the `.wasm` file beside the bindings, fetched in a browser and read from the file system in Node.
+ */
+export function loadCore(wasm) {
+    const url = new URL('../pkg/ensomi_generate_wasm_bg.wasm', import.meta.url);
+    loaded ??= init({ module_or_path: wasm ?? (url.protocol === 'file:' ? readFile(url) : url) }).then(() => bindings);
+    return loaded;
+}
+
+/** Which instance of the core is current; a trap in `callCore` starts the next. */
+export const coreInstance = () => instance;
+
+/**
+ * Calls into the core. A trap (a panic, whose message the error carries, or memory or stack exhausted) can leave
+ * the instance's memory inconsistent, so the instance is replaced before the error is rethrown; objects of the
+ * old instance cannot be used after that.
+ */
+export function callCore(call) {
+    try {
+        return call();
+    }
+    catch (error) {
+        if (error instanceof WebAssembly.RuntimeError) {
+            bindings.__wbg_reset_state();
+            instance++;
+        }
+        throw error;
+    }
+}

package/src/decode_audio.js

@@ -0,0 +1,22 @@
+/** Browser decode on two frontend clocks. WebAudio's codec/resampler can differ from Python's. */
+export async function decodeAudio(arrayBuffer) {
+    if (!(arrayBuffer instanceof ArrayBuffer) || arrayBuffer.byteLength === 0)
+        throw new Error('Audio must be a nonempty ArrayBuffer');
+    const decode = async (rate) => {
+        const context = new OfflineAudioContext(1, 1, rate), buffer = await context.decodeAudioData(arrayBuffer.slice(0)), mono = new Float32Array(buffer.length);
+        for (let channel = 0; channel < buffer.numberOfChannels; channel++) {
+            const samples = buffer.getChannelData(channel);
+            for (let i = 0; i < mono.length; i++)
+                mono[i] += samples[i] / buffer.numberOfChannels;
+        }
+        return mono;
+    };
+    const [mel, beatthis] = await Promise.all([decode(24000), decode(22050)]);
+    let peak = 0;
+    for (const sample of mel)
+        peak = Math.max(peak, Math.abs(sample));
+    if (peak)
+        for (let i = 0; i < mel.length; i++)
+            mel[i] /= peak;
+    return { pcm: { samples: mel, channels: 1, mel, beatthis }, sampleRate: 24000 };
+}

package/src/export.js

@@ -0,0 +1,31 @@
+import { zipSync, strToU8 } from 'fflate';
+
+/** Throws unless the chart's objects are sorted, in lanes 0 to 3, whole milliseconds and without overlap. */
+export function validateChart(chart) {
+    if (!Number.isInteger(chart.duration_ms) || chart.duration_ms < 0)
+        throw new Error('Chart duration must be nonnegative whole milliseconds');
+    const end = [-1, -1, -1, -1];
+    let previous;
+    for (const object of chart.objects) {
+        const { lane, start_ms: start, end_ms: release } = object;
+        if (!Number.isInteger(lane) || lane < 0 || lane > 3 || !Number.isInteger(start) || start < 0 || start > chart.duration_ms)
+            throw new Error(`Invalid chart object: ${JSON.stringify(object)}`);
+        if (previous && (start < previous.start_ms || start === previous.start_ms && lane <= previous.lane))
+            throw new Error('Chart objects must be sorted by start then lane, with no duplicate heads');
+        if (start < end[lane])
+            throw new Error(`Head inside an open hold in lane ${lane} at ${start}`);
+        if (release !== undefined && (!Number.isInteger(release) || release <= start || release > chart.duration_ms))
+            throw new Error(`Invalid hold release in lane ${lane} at ${start}`);
+        end[lane] = release ?? start;
+        previous = object;
+    }
+}
+
+export function audioFilename(name) {
+    const ext = /\.([a-zA-Z0-9]+)$/.exec(name)?.[1].toLowerCase() ?? 'bin';
+    return `audio.${ext}`;
+}
+
+export function exportOsz(osu, audio, receiptJson) {
+    return zipSync({ 'generated.osu': strToU8(osu), [audioFilename(audio.name)]: [audio.bytes, { level: 0 }], 'receipt.json': strToU8(receiptJson + '\n') }, { level: 6 });
+}

package/src/generate.js

@@ -0,0 +1,73 @@
+import { callCore, coreInstance } from './core.js';
+import { exportOsz } from './export.js';
+
+/** Lets queued tasks run, such as the WebGPU result that a graph run is waiting for. */
+const pause = () => new Promise(resolve => {
+    const channel = new MessageChannel();
+    channel.port1.onmessage = () => {
+        channel.port1.close();
+        resolve();
+    };
+    channel.port2.postMessage(null);
+});
+
+/**
+ * Generates a chart. `pcm` is mono Float32Array/Float64Array or `{samples, channels, mel?, beatthis?}` with
+ * interleaved samples at `sampleRate`; `request` is a GenerateChartRequest.
+ * Options: `onProgress({phase})`, `audio` ({name, bytes}, adds an .osz), `setting` (difficulty name), `clock`
+ * (milliseconds, default `performance.now`), `replay` (a core `ReplayDraws`, moved into the generation, for
+ * conformance checks; the result then has `replayConsumed`, the number of draws used).
+ */
+export async function generateChart(models, pcm, sampleRate, request, options = {}) {
+    if (options.random || options.trace)
+        throw new Error('The Wasm core takes recorded draws as options.replay and has no trace');
+    const audio = pcm instanceof Float32Array || pcm instanceof Float64Array ? { samples: pcm, channels: 1 } : pcm;
+    const seed = request.seed ?? 0;
+    if (typeof seed !== 'number' && typeof seed !== 'bigint' || typeof seed === 'number' && (!Number.isSafeInteger(seed) || seed < 0) || typeof seed === 'bigint' && (seed < 0n || seed > 0xffffffffffffffffn))
+        throw new Error('Seed must be an unsigned 64-bit integer; use bigint above the safe-number range');
+    await models.reloadCore();
+    const { core, wasm } = models;
+    const resamplerFiles = audio.beatthis ? null : core.WasmModels.resamplerFiles(sampleRate);
+    const resampler = resamplerFiles ? await Promise.all(Array.from(resamplerFiles, name => models.runtimeSource(name))) : undefined;
+    const { seed: _seed, ...rest } = request;
+    const clock = options.clock ?? (() => performance.now());
+    // decodeAudio passes its Mel stream as the samples too; the core then takes the stream once.
+    const samples = audio.mel && audio.beatthis && audio.samples === audio.mel ? new Float32Array(0) : audio.samples;
+    const generator = callCore(() => new core.WasmGenerator(wasm, samples, audio.channels, sampleRate, audio.mel, audio.beatthis, resampler, JSON.stringify(rest), BigInt(seed), typeof seed === 'bigint', options.setting, options.audio?.name, clock, options.replay));
+    const instance = coreInstance();
+    try {
+        let outputs;
+        for (;;) {
+            const step = callCore(() => generator.step(outputs));
+            outputs = undefined;
+            if (step.kind === 'progress') {
+                options.onProgress?.({ phase: step.phase });
+            }
+            else if (step.kind === 'run') {
+                const feeds = {};
+                for (const [name, tensor] of Object.entries(step.inputs))
+                    feeds[name] = new models.ort.Tensor(tensor.type, tensor.data, tensor.dims);
+                // While the graph runs (BeatThis on the GPU), the core does what needs no outputs, yielding between
+                // slices so that the outputs are taken as soon as they arrive.
+                let running = true;
+                const run = models.sessions[step.graph].run(feeds);
+                run.then(() => { running = false; }, () => { running = false; });
+                while (running && callCore(() => generator.work()))
+                    await pause();
+                const result = await run;
+                outputs = {};
+                for (const [name, tensor] of Object.entries(result))
+                    outputs[name] = { type: tensor.type, dims: Array.from(tensor.dims), data: tensor.data };
+            }
+            else {
+                const { chart, osu, receipt, replay_consumed: replayConsumed } = step;
+                const osz = options.audio ? exportOsz(osu, options.audio, receipt.receipt_json) : undefined;
+                return { chart, receipt: { model_name: receipt.model_name, model_version: receipt.model_version, seed, star: receipt.star, receipt_json: receipt.receipt_json }, osu, ...(osz ? { osz } : {}), ...(options.replay ? { replayConsumed } : {}) };
+            }
+        }
+    }
+    finally {
+        if (coreInstance() === instance)
+            generator.free();
+    }
+}

package/src/index.d.ts

@@ -0,0 +1,157 @@
+import type * as ORT from 'onnxruntime-web';
+
+/** EnsomiProtocol 0.2 chart.proto field names and enum values. */
+export interface ChartObject {
+    lane: number;
+    start_ms: number;
+    end_ms?: number;
+}
+export interface Chart {
+    duration_ms: number;
+    objects: ChartObject[];
+}
+export interface AudioAssetRef {
+    local_path?: string;
+    uri?: string;
+    asset_id?: string;
+    audio_length_ms?: number;
+    sha256_hex?: string;
+}
+export declare enum LongNoteAmount {
+    UNSPECIFIED = 0,
+    TAP_ONLY = 1,
+    HALF = 2,
+    NATURAL = 3
+}
+export declare enum PatternKind {
+    UNSPECIFIED = 0,
+    STREAM = 1,
+    JUMPSTREAM = 2,
+    CHORDSTREAM = 3,
+    JACK = 4,
+    CHORDJACK = 5
+}
+export declare enum ControlStatus {
+    UNSPECIFIED = 0,
+    WORKS = 1,
+    EXPERIMENTAL = 2,
+    NOT_WORKING = 3
+}
+export declare enum GenerationPhase {
+    UNSPECIFIED = 0,
+    LOADING = 1,
+    TIMING = 2,
+    ROWS = 3,
+    ARRANGEMENT = 4,
+    EXPORT = 5
+}
+export interface SectionRequest {
+    start_ms: number;
+    end_ms: number;
+    star?: number;
+    pattern?: {
+        kind: PatternKind;
+        strength: number;
+    };
+}
+export interface ChartControls {
+    long_notes?: LongNoteAmount;
+    sections?: SectionRequest[];
+}
+export interface GenerateChartRequest {
+    audio?: AudioAssetRef;
+    star: number;
+    seed?: number | bigint;
+    controls?: ChartControls;
+    title?: string;
+    artist?: string;
+}
+export interface ModelCapabilities {
+    model_name: string;
+    model_version: string;
+    files: {
+        path: string;
+        sha256_hex: string;
+    }[];
+    min_star: number;
+    max_star_exclusive: number;
+    controls: {
+        name: string;
+        status: ControlStatus;
+        note: string;
+    }[];
+}
+export interface GenerationProgress {
+    phase: GenerationPhase;
+}
+export interface GenerationReceipt {
+    model_name: string;
+    model_version: string;
+    seed: number | bigint;
+    star: number;
+    receipt_json: string;
+    osz_path?: string;
+}
+export interface GenerationResult {
+    chart: Chart;
+    receipt: GenerationReceipt;
+    osu: string;
+    osz?: Uint8Array;
+}
+export interface AudioPCM {
+    /** Interleaved samples at the sampleRate passed to generateChart. With both frontend streams it may be the
+     * `mel` array itself, as `decodeAudio` gives it; the core then takes that array once. */
+    samples: Float32Array | Float64Array;
+    channels: number;
+    /** Optional already-decoded frontend streams: mono, 24000 and 22050 Hz. */
+    mel?: Float32Array;
+    beatthis?: Float32Array;
+}
+export type PCM = Float32Array | Float64Array | AudioPCM;
+export interface DecodedAudio {
+    pcm: AudioPCM;
+    sampleRate: number;
+}
+export interface EncodedAudio {
+    name: string;
+    bytes: Uint8Array;
+}
+
+export type Ort = typeof ORT;
+export type FileSource = (relativePath: string) => Promise<Uint8Array>;
+/** Reads files relative to a base URL, with `fetch` or, for a `file:` URL in Node, from the file system. */
+export declare function urlSource(baseUrl: string): FileSource;
+export declare function localFilesSource(files: Iterable<File>): FileSource;
+
+export declare class Models {
+    readonly capabilities: ModelCapabilities;
+    readonly ort: Ort;
+    readonly sessions: Record<string, ORT.InferenceSession>;
+    dispose(): Promise<void>;
+}
+export interface LoadModelsOptions {
+    ort?: Ort;
+    session?: ORT.InferenceSession.SessionOptions;
+    backend?: 'webgpu' | 'wasm';
+    /** ORT WebAssembly settings applied before sessions are created, e.g. wasmPaths and numThreads. */
+    wasm?: Partial<ORT.Env.WebAssemblyFlags>;
+    /** Reads the `resample-*` tables, needed for raw PCM at a rate other than 22050 Hz without a `beatthis` stream.
+     * Default: the package's `resampler/` directory. */
+    runtimeSource?: FileSource;
+    /** The Wasm core: bytes, a `WebAssembly.Module` or a URL. Default: the package's `pkg/ensomi_generate_wasm_bg.wasm`. */
+    core?: BufferSource | WebAssembly.Module | URL | string;
+}
+export declare function loadModels(source: FileSource | string, options?: LoadModelsOptions): Promise<Models>;
+
+export interface GenerateOptions {
+    onProgress?: (progress: GenerationProgress) => void;
+    audio?: EncodedAudio;
+    setting?: string;
+    /** Milliseconds from any fixed origin; the receipt's timings come from it. Default `performance.now`. */
+    clock?: () => number;
+    /** Recorded draws of a release recording, for conformance checks. */
+    replay?: unknown;
+}
+export declare function generateChart(models: Models, pcm: PCM, sampleRate: number, request: GenerateChartRequest, options?: GenerateOptions): Promise<GenerationResult>;
+export declare function decodeAudio(arrayBuffer: ArrayBuffer): Promise<DecodedAudio>;
+export declare function validateChart(chart: Chart): void;

package/src/index.js

@@ -0,0 +1,5 @@
+export { loadModels, urlSource, localFilesSource, Models } from './models.js';
+export { generateChart } from './generate.js';
+export { decodeAudio } from './decode_audio.js';
+export { validateChart } from './export.js';
+export * from './types.js';