Optimize blob and wav file saving

No need to save the entire audio in memory

Author: xenova

src/utils/audio.js

@@ -9,10 +9,10 @@
 
 import { getFile } from './hub.js';
 import { FFT, max } from './maths.js';
-import { calculateReflectOffset, saveBlob } from './core.js';
+import { calculateReflectOffset } from './core.js';
+import { saveBlob } from './io.js';
 import { apis } from '../env.js';
 import { Tensor, matmul } from './tensor.js';
-import fs from 'node:fs';
 
 /**
  * Helper function to read audio from a path/URL.
@@ -733,23 +733,24 @@ export function window_function(window_length, name, { periodic = true, frame_le
 }
 
 /**
- * Encode audio data to a WAV file.
+ * Efficiently encode audio data to a WAV file.
  * WAV file specs : https://en.wikipedia.org/wiki/WAV#WAV_File_header
  *
  * Adapted from https://www.npmjs.com/package/audiobuffer-to-wav
- * @param {Float32Array} samples The audio samples.
+ * @param {Float32Array[]} chunks The audio samples.
  * @param {number} rate The sample rate.
- * @returns {ArrayBuffer} The WAV audio buffer.
+ * @returns {Blob} The WAV file as a Blob.
  */
-function encodeWAV(samples, rate) {
-    let offset = 44;
-    const buffer = new ArrayBuffer(offset + samples.length * 4);
+function encodeWAV(chunks, rate) {
+    const totalLength = chunks.reduce((acc, chunk) => acc + chunk.length, 0);
+
+    const buffer = new ArrayBuffer(44);
     const view = new DataView(buffer);
 
     /* RIFF identifier */
     writeString(view, 0, 'RIFF');
     /* RIFF chunk length */
-    view.setUint32(4, 36 + samples.length * 4, true);
+    view.setUint32(4, 36 + totalLength * 4, true);
     /* RIFF type */
     writeString(view, 8, 'WAVE');
     /* format chunk identifier */
@@ -771,13 +772,10 @@ function encodeWAV(samples, rate) {
     /* data chunk identifier */
     writeString(view, 36, 'data');
     /* data chunk length */
-    view.setUint32(40, samples.length * 4, true);
-
-    for (let i = 0; i < samples.length; ++i, offset += 4) {
-        view.setFloat32(offset, samples[i], true);
-    }
+    view.setUint32(40, totalLength * 4, true);
 
-    return buffer;
+    // @ts-expect-error TS2322
+    return new Blob([buffer, ...chunks], { type: "audio/wav" });
 }
 
 function writeString(view, offset, string) {
@@ -789,7 +787,7 @@ function writeString(view, offset, string) {
 export class RawAudio {
     /**
      * Create a new `RawAudio` object.
-     * @param {Float32Array} audio Audio data
+     * @param {Float32Array|Float32Array[]} audio Audio data, either as a single `Float32Array` chunk or multiple `Float32Array` chunks.
      * @param {number} sampling_rate Sampling rate of the audio data
      */
     constructor(audio, sampling_rate) {
@@ -798,44 +796,49 @@ export class RawAudio {
     }
 
     /**
-     * Convert the audio to a wav file buffer.
-     * @returns {ArrayBuffer} The WAV file.
+     * Get the audio data, accumulating all chunks if necessary.
+     * @returns {Float32Array} The audio data.
      */
-    toWav() {
-        return encodeWAV(this.audio, this.sampling_rate);
+    get data() {
+        if (Array.isArray(this.audio)) {
+            if (this.audio.length === 0) {
+                return new Float32Array(0);
+            }
+            if (this.audio.length === 1) {
+                return this.audio[0];
+            }
+            // Concatenate all chunks into a single Float32Array
+            const totalLength = this.audio.reduce((acc, chunk) => acc + chunk.length, 0);
+            const result = new Float32Array(totalLength);
+            let offset = 0;
+            for (const chunk of this.audio) {
+                result.set(chunk, offset);
+                offset += chunk.length;
+            }
+            return result;
+        } else {
+            return this.audio;
+        }
     }
 
     /**
      * Convert the audio to a blob.
      * @returns {Blob}
      */
     toBlob() {
-        const wav = this.toWav();
-        const blob = new Blob([wav], { type: 'audio/wav' });
-        return blob;
+        let audio = this.audio;
+        if (audio instanceof Float32Array) {
+            audio = [audio]; // Ensure audio is an array of chunks
+        }
+        return encodeWAV(audio, this.sampling_rate);
     }
 
     /**
      * Save the audio to a wav file.
      * @param {string} path
+     * @returns {Promise<void>}
      */
     async save(path) {
-        let fn;
-
-        if (apis.IS_BROWSER_ENV) {
-            if (apis.IS_WEBWORKER_ENV) {
-                throw new Error('Unable to save a file from a Web Worker.');
-            }
-            fn = saveBlob;
-        } else if (apis.IS_FS_AVAILABLE) {
-            fn = async (/** @type {string} */ path, /** @type {Blob} */ blob) => {
-                let buffer = await blob.arrayBuffer();
-                fs.writeFileSync(path, Buffer.from(buffer));
-            };
-        } else {
-            throw new Error('Unable to save because filesystem is disabled in this environment.');
-        }
-
-        await fn(path, this.toBlob());
+        return saveBlob(path, this.toBlob());
     }
 }

src/utils/core.js

@@ -187,32 +187,6 @@ export function calculateReflectOffset(i, w) {
     return Math.abs(((i + w) % (2 * w)) - w);
 }
 
-/**
- * Save blob file on the web.
- * @param {string} path The path to save the blob to
- * @param {Blob} blob The blob to save
- */
-export function saveBlob(path, blob) {
-    // Convert the canvas content to a data URL
-    const dataURL = URL.createObjectURL(blob);
-
-    // Create an anchor element with the data URL as the href attribute
-    const downloadLink = document.createElement('a');
-    downloadLink.href = dataURL;
-
-    // Set the download attribute to specify the desired filename for the downloaded image
-    downloadLink.download = path;
-
-    // Trigger the download
-    downloadLink.click();
-
-    // Clean up: remove the anchor element from the DOM
-    downloadLink.remove();
-
-    // Revoke the Object URL to free up memory
-    URL.revokeObjectURL(dataURL);
-}
-
 /**
  *
  * @param {Object} o

src/utils/image.js

@@ -7,10 +7,11 @@
  * @module utils/image
  */
 
-import { isNullishDimension, saveBlob } from './core.js';
+import { isNullishDimension } from './core.js';
 import { getFile } from './hub.js';
 import { apis } from '../env.js';
 import { Tensor } from './tensor.js';
+import { saveBlob } from './io.js';
 
 // Will be empty (or not used) if running in browser or web-worker
 import sharp from 'sharp';
@@ -769,6 +770,7 @@ export class RawImage {
     /**
      * Save the image to the given path.
      * @param {string} path The path to save the image to.
+     * @returns {Promise<void>}
      */
     async save(path) {
         if (IS_BROWSER_OR_WEBWORKER) {
@@ -782,15 +784,19 @@ export class RawImage {
             // Convert image to Blob
             const blob = await this.toBlob(mime);
 
-            saveBlob(path, blob);
-        } else if (!apis.IS_FS_AVAILABLE) {
-            throw new Error('Unable to save the image because filesystem is disabled in this environment.');
-        } else {
+            return saveBlob(path, blob);
+        } else if (apis.IS_FS_AVAILABLE) {
             const img = this.toSharp();
-            return await img.toFile(path);
+            await img.toFile(path);
+        } else {
+            throw new Error('Unable to save the image because filesystem is disabled in this environment.');
         }
     }
 
+    /**
+     * Convert the image to a Sharp instance.
+     * @returns {import('sharp').Sharp} The Sharp instance.
+     */
     toSharp() {
         if (IS_BROWSER_OR_WEBWORKER) {
             throw new Error('toSharp() is only supported in server-side environments.');

src/utils/io.js

@@ -0,0 +1,52 @@
+
+import fs from 'node:fs';
+import { Readable } from 'node:stream';
+import { pipeline as pipe } from 'node:stream/promises';
+
+import { apis } from "../env.js";
+
+/**
+ * Save blob file.
+ * @param {string} path The path to save the blob to
+ * @param {Blob} blob The blob to save
+ * @returns {Promise<void>} A promise that resolves when the blob has been saved
+ */
+export async function saveBlob(path, blob) {
+    if (apis.IS_BROWSER_ENV) {
+        if (apis.IS_WEBWORKER_ENV) {
+            throw new Error('Unable to save a file from a Web Worker.');
+        }
+        // Convert the canvas content to a data URL
+        const dataURL = URL.createObjectURL(blob);
+        
+        // Create an anchor element with the data URL as the href attribute
+        const downloadLink = document.createElement('a');
+        downloadLink.href = dataURL;
+        
+        // Set the download attribute to specify the desired filename for the downloaded image
+        downloadLink.download = path;
+        
+        // Trigger the download
+        downloadLink.click();
+        
+        // Clean up: remove the anchor element from the DOM
+        downloadLink.remove();
+        
+        // Revoke the Object URL to free up memory
+        URL.revokeObjectURL(dataURL);
+
+    } else if (apis.IS_FS_AVAILABLE) {
+        // Convert Blob to a Node.js Readable Stream
+        const webStream = blob.stream();
+        const nodeStream = Readable.fromWeb(webStream);
+
+        // Create the file write stream
+        const fileStream = fs.createWriteStream(path);
+
+        // Pipe the readable stream to the file write stream
+        await pipe(nodeStream, fileStream);
+
+    } else {
+        throw new Error('Unable to save because filesystem is disabled in this environment.');
+    }
+}

tests/utils/audio.test.js

@@ -0,0 +1,80 @@
+import { RawAudio } from "../../src/utils/audio.js";
+
+/**
+ * Helper function to generate a sine wave.
+ * @param {number} length Length of the audio in samples.
+ * @param {number} freq Frequency of the sine wave.
+ * @param {number} sampling_rate Sampling rate.
+ * @returns {Float32Array} The generated sine wave.
+ */
+function generateSineWave(length, freq, sampling_rate) {
+    const audio = new Float32Array(length);
+    for (let i = 0; i < length; ++i) {
+        audio[i] = Math.sin((2 * Math.PI * freq * i) / sampling_rate);
+    }
+    return audio;
+}
+
+describe("Audio utilities", () => {
+    describe("RawAudio", () => {
+        it("should create RawAudio from a single Float32Array", () => {
+            const sampling_rate = 16000;
+            const audioData = generateSineWave(1000, 440, sampling_rate);
+            const rawAudio = new RawAudio(audioData, sampling_rate);
+
+            expect(rawAudio.sampling_rate).toBe(sampling_rate);
+            expect(rawAudio.data).toBeInstanceOf(Float32Array);
+            expect(rawAudio.data).toEqual(audioData);
+            expect(rawAudio.data.length).toBe(1000);
+        });
+
+        it("should create RawAudio from multiple Float32Array chunks", () => {
+            const sampling_rate = 16000;
+            const chunk1 = generateSineWave(500, 440, sampling_rate);
+            const chunk2 = generateSineWave(500, 880, sampling_rate);
+            const rawAudio = new RawAudio([chunk1, chunk2], sampling_rate);
+
+            expect(rawAudio.sampling_rate).toBe(sampling_rate);
+            expect(rawAudio.data).toBeInstanceOf(Float32Array);
+            expect(rawAudio.data.length).toBe(1000);
+
+            // Check if concatenation is correct
+            const combined = new Float32Array(1000);
+            combined.set(chunk1, 0);
+            combined.set(chunk2, 500);
+            expect(rawAudio.data).toEqual(combined);
+        });
+
+        it("should handle empty array of chunks", () => {
+            const rawAudio = new RawAudio([], 16000);
+            expect(rawAudio.data).toBeInstanceOf(Float32Array);
+            expect(rawAudio.data.length).toBe(0);
+        });
+
+        it("should convert to Blob (WAV)", () => {
+            const sampling_rate = 16000;
+            const audioData = generateSineWave(1000, 440, sampling_rate);
+            const rawAudio = new RawAudio(audioData, sampling_rate);
+
+            const blob = rawAudio.toBlob();
+            expect(blob).toBeInstanceOf(Blob);
+            expect(blob.type).toBe("audio/wav");
+            
+            // WAV header is 44 bytes
+            // 1000 samples * 4 bytes/sample (float32) = 4000 bytes
+            expect(blob.size).toBe(4044); 
+        });
+
+        it("should convert to Blob (WAV) from chunks", () => {
+            const sampling_rate = 16000;
+            const chunk1 = generateSineWave(500, 440, sampling_rate);
+            const chunk2 = generateSineWave(500, 880, sampling_rate);
+            const rawAudio = new RawAudio([chunk1, chunk2], sampling_rate);
+
+            const blob = rawAudio.toBlob();
+            expect(blob).toBeInstanceOf(Blob);
+            expect(blob.type).toBe("audio/wav");
+            expect(blob.size).toBe(4044); // 44 header + 4000 data
+        });
+    });
+});