echogarden 2.8.6 → 2.9.0

package/src/recognition/WhisperSTT.ts

@@ -1,9 +1,10 @@
+import chalk from 'chalk'
 import type * as Onnx from 'onnxruntime-node'
 
 import { Logger } from '../utilities/Logger.js'
-import { computeMelSpectogramUsingFilterbanks, Filterbank } from '../dsp/MelSpectogram.js'
+import { computeMelSpectrogramUsingFilterbanks, Filterbank } from '../dsp/MelSpectrogram.js'
 import { clip, getIntegerRange, getTopKIndexes, splitFloat32Array, yieldToEventLoop } from '../utilities/Utilities.js'
-import { indexOfMax, logOfVector, logSumExp, meanOfVector, softmax, sumOfSquaresForVector, sumVector } from '../math/VectorMath.js'
+import { indexOfMax, logOfVector, logSumExp, meanOfVector, medianOfVector, softmax, sumAndSumOfSquaresOfVector, sumOfSquaresOfVector, sumVector } from '../math/VectorMath.js'
 
 import { alignDTWWindowed } from '../alignment/DTWSequenceAlignmentWindowed.js'
 import { extendDeep } from '../utilities/ObjectUtilities.js'
@@ -11,10 +12,9 @@ import { Timeline, TimelineEntry } from '../utilities/Timeline.js'
 import { AlignmentPath } from '../alignment/SpeechAlignment.js'
 import { getRawAudioDuration, RawAudio, sliceRawAudio } from '../audio/AudioUtilities.js'
 import { readFileAsUtf8 } from '../utilities/FileSystem.js'
-import type { LanguageDetectionResults } from '../api/API.js'
+import { logLevelGreaterOrEqualTo, type LanguageDetectionResults } from '../api/API.js'
 import { formatLanguageCodeWithName, getShortLanguageCode, languageCodeToName } from '../utilities/Locale.js'
 import { loadPackage } from '../utilities/PackageManager.js'
-import chalk from 'chalk'
 import { XorShift32PRNG } from '../utilities/RandomGenerator.js'
 import { detectSpeechLanguageByParts } from '../api/SpeechLanguageDetection.js'
 import { type Tiktoken } from 'tiktoken/lite'
@@ -25,6 +25,7 @@ import { dmlProviderAvailable, getOnnxSessionOptions, makeOnnxLikeFloat32Tensor,
 import { murmurHash3_int32Input } from '../utilities/Hashing.js'
 import { containsInvalidCodepoint, getTokenRepetitionScore } from '../utilities/StringUtilities.js'
 import { joinPath } from '../utilities/PathUtilities.js'
+import { Timer } from '../utilities/Timer.js'
 
 export async function recognize(
 	sourceRawAudio: RawAudio,
@@ -32,11 +33,12 @@ export async function recognize(
 	modelDir: string,
 	task: WhisperTask,
 	sourceLanguage: string,
-	options: WhisperOptions) {
+	options: WhisperOptions,
+	onPart?: WhisperPartCallback) {
 
 	options = extendDeep(defaultWhisperOptions, options)
 
-	if (sourceRawAudio.sampleRate != 16000) {
+	if (sourceRawAudio.sampleRate !== 16000) {
 		throw new Error('Source audio must have a sample rate of 16000 Hz')
 	}
 
@@ -46,7 +48,7 @@ export async function recognize(
 		throw new Error(`The language ${formatLanguageCodeWithName(sourceLanguage)} is not supported by the Whisper engine.`)
 	}
 
-	if (isEnglishOnlyModel(modelName) && sourceLanguage != 'en') {
+	if (isEnglishOnlyModel(modelName) && sourceLanguage !== 'en') {
 		throw new Error(`The model '${modelName}' can only be used with English inputs. However, the given source language was ${languageCodeToName(sourceLanguage)}.`)
 	}
 
@@ -60,7 +62,7 @@ export async function recognize(
 
 	// Workaround issue with large-v3-turbo that produces invalid results when a prompt is passed to it.
 	// Always disable autoprompting for that model.
-	if (options.autoPromptParts && modelName == 'large-v3-turbo') {
+	if (options.autoPromptParts && modelName === 'large-v3-turbo') {
 		options.autoPromptParts = false
 	}
 
@@ -81,7 +83,7 @@ export async function recognize(
 		decoderProviders,
 		seed)
 
-	const result = await whisper.recognize(sourceRawAudio, task, sourceLanguage, options)
+	const result = await whisper.recognize(sourceRawAudio, task, sourceLanguage, options, undefined, onPart)
 
 	return result
 }
@@ -96,7 +98,7 @@ export async function align(
 
 	options = extendDeep(defaultWhisperAlignmentOptions, options)
 
-	if (sourceRawAudio.sampleRate != 16000) {
+	if (sourceRawAudio.sampleRate !== 16000) {
 		throw new Error('Source audio must have a sample rate of 16000 Hz')
 	}
 
@@ -106,7 +108,7 @@ export async function align(
 		throw new Error(`The language ${formatLanguageCodeWithName(sourceLanguage)} is not supported by the Whisper engine.`)
 	}
 
-	if (isEnglishOnlyModel(modelName) && sourceLanguage != 'en') {
+	if (isEnglishOnlyModel(modelName) && sourceLanguage !== 'en') {
 		throw new Error(`The model '${modelName}' can only be used with English inputs. However, the given source language was ${languageCodeToName(sourceLanguage)}.`)
 	}
 
@@ -139,7 +141,7 @@ export async function alignEnglishTranslation(
 
 	options = extendDeep(defaultWhisperAlignmentOptions, options)
 
-	if (sourceRawAudio.sampleRate != 16000) {
+	if (sourceRawAudio.sampleRate !== 16000) {
 		throw new Error('Source audio must have a sample rate of 16000 Hz')
 	}
 
@@ -184,7 +186,7 @@ export async function detectLanguage(
 
 	options = extendDeep(defaultWhisperLanguageDetectionOptions, options)
 
-	if (sourceRawAudio.sampleRate != 16000) {
+	if (sourceRawAudio.sampleRate !== 16000) {
 		throw new Error('Source audio must have a sample rate of 16000 Hz')
 	}
 
@@ -232,7 +234,7 @@ export async function detectVoiceActivity(
 
 	options = extendDeep(defaultWhisperVADOptions, options)
 
-	if (sourceRawAudio.sampleRate != 16000) {
+	if (sourceRawAudio.sampleRate !== 16000) {
 		throw new Error('Source audio must have a sample rate of 16000 Hz')
 	}
 
@@ -365,6 +367,7 @@ export class Whisper {
 		language: string,
 		options: WhisperOptions,
 		logitFilter?: WhisperLogitFilter,
+		onPart?: WhisperPartCallback,
 	) {
 		await this.initializeIfNeeded()
 
@@ -441,7 +444,11 @@ export class Whisper {
 			let {
 				decodedTokens: partTokens,
 				decodedTokensConfidence: partTokensConfidence,
-				decodedTokensCrossAttentionQKs: partCrossAttentionQKs,
+				decodedTokensCrossAttentionQKs: partTokensCrossAttentionQKs,
+
+				decodedTokensDecodingTime: partTokensDecodingTime,
+				decodedTokensInferenceTime: partTokensInferenceTime,
+				decodedTokensOverheadTime: partTokensOverheadTime,
 			} = await this.decodeTokens(
 				audioPartFeatures,
 				initialTokens,
@@ -470,19 +477,19 @@ export class Whisper {
 
 			await logger.startAsync(`Extract timeline for part (timestamp accuracy: ${options.timestampAccuracy!})`)
 
-			if (partTokens.length != partCrossAttentionQKs.length) {
-				throw new Error('Unexpected: partTokens.length != partCrossAttentionQKs.length')
+			if (partTokens.length !== partTokensCrossAttentionQKs.length) {
+				throw new Error('Unexpected: partTokens.length !== partCrossAttentionQKs.length')
 			}
 
 			// Prepare tokens
 			partTokens = partTokens.slice(initialTokens.length)
 			partTokensConfidence = partTokensConfidence.slice(initialTokens.length)
-			partCrossAttentionQKs = partCrossAttentionQKs.slice(initialTokens.length)
+			partTokensCrossAttentionQKs = partTokensCrossAttentionQKs.slice(initialTokens.length)
 
 			// Find alignment path
 			let alignmentHeads: number[] | undefined
 
-			if (options.timestampAccuracy === 'medium' || options.model == 'large-v3-turbo') {
+			if (options.timestampAccuracy === 'medium' || options.model === 'large-v3-turbo') {
 				alignmentHeads = this.alignmentHeadIndexes
 			} else if (options.timestampAccuracy === 'high') {
 				alignmentHeads = undefined
@@ -490,11 +497,18 @@ export class Whisper {
 				throw new Error(`Unsupported timestamp accuracy '${options.timestampAccuracy}', can only be 'medium' or 'high'.`)
 			}
 
-			const alignmentPath = await this.findAlignmentPathFromQKs(partCrossAttentionQKs, partTokens, 0, segmentFrameCount, alignmentHeads)
+			const alignmentPath = await this.findAlignmentPathFromQKs(partTokensCrossAttentionQKs, partTokens, 0, segmentFrameCount, alignmentHeads)
 
 			// Generate timeline from alignment path
 			const partTimeline = await this.getTokenTimelineFromAlignmentPath(alignmentPath, partTokens, segmentStartTime, segmentEndTime, partTokensConfidence)
 
+			if (onPart) {
+				const partWordTimeline = this.tokenTimelineToWordTimeline(partTimeline, language)
+				const partTranscript = this.tokensToText(partTokens)
+
+				onPart(partTranscript, partTimeline, partWordTimeline)
+			}
+
 			// Add tokens to output
 			allDecodedTokens.push(...partTokens)
 			timeline.push(...partTimeline)
@@ -514,6 +528,15 @@ export class Whisper {
 			audioOffset = audioEndOffset
 
 			logger.end()
+
+			if (logLevelGreaterOrEqualTo('trace')) {
+				const promptDecodingTime = partTokensDecodingTime[0]
+				const medianTokenDecodingTime = medianOfVector(partTokensDecodingTime.slice(1))
+				const medianTokenInferenceTime = medianOfVector(partTokensInferenceTime.slice(1))
+				const medianOverheadTime = medianOfVector(partTokensOverheadTime.slice(1))
+
+				logger.log(`${chalk.blueBright('Context')}: ${initialTokens.length + partTokens.length} tokens (${initialTokens.length} prompt, ${partTokens.length} decoded)\n${chalk.blueBright('Prompt decode time')}: ${promptDecodingTime.toFixed(1)}ms\n${chalk.blueBright('Median token decode time')}: ${medianTokenDecodingTime.toFixed(1)}ms (${medianTokenInferenceTime.toFixed(1)}ms inference, ${medianOverheadTime.toFixed(2)}ms overhead)`, 'trace')
+			}
 		}
 
 		// Convert token timeline to word timeline
@@ -812,8 +835,17 @@ export class Whisper {
 
 		const maxTokensPerPart = Math.min(options.maxTokensPerPart!, largestMaximumTokensPerPart)
 
+		let decodedTokensInferenceTime: number[] = []
+		let decodedTokensDecodingTime: number[] = []
+
+		const tokenDecodingTimeTimer = new Timer()
+
 		// Start decoding loop
 		for (let decodedTokenCount = 0; decodedTokenCount < maxTokensPerPart; decodedTokenCount++) {
+			if (decodedTokenCount > 0) {
+				decodedTokensDecodingTime.push(tokenDecodingTimeTimer.getElapsedTimeAndRestart())
+			}
+
 			const isInitialState = decodedTokens.length === initialTokens.length
 			const atLeastOneTextTokenDecoded = decodedTokens.slice(initialTokens.length).some(token => this.isTextToken(token))
 
@@ -847,9 +879,13 @@ export class Whisper {
 				offset: offsetTensor
 			}
 
-			// Run decoder model
+			//// Infer with ONNX decoder model
+			const tokenInferenceTimeTimer = new Timer()
+
 			const decoderOutputs = await this.textDecoder!.run(decoderInputs)
 
+			decodedTokensInferenceTime.push(tokenInferenceTimeTimer.elapsedTime)
+
 			// Extract decoder model results
 			const logitsBuffer = decoderOutputs['logits'].data as Float32Array
 			kvCacheTensor = decoderOutputs['output_kv_cache'] as any
@@ -1142,6 +1178,10 @@ export class Whisper {
 			await yieldToEventLoop()
 		}
 
+		if (decodedTokensDecodingTime.length === decodedTokensInferenceTime.length - 1) {
+			decodedTokensDecodingTime.push(tokenDecodingTimeTimer.getElapsedTimeAndRestart())
+		}
+
 		// If at least two timestamp tokens were decoded and it's not the final part,
 		// truncate up to the last timestamp token
 		if (timestampTokenSeenCount >= 2 && !isFinalPart) {
@@ -1151,8 +1191,13 @@ export class Whisper {
 			decodedTokensTimestampLogits = decodedTokensTimestampLogits.slice(0, sliceEndTokenIndex)
 			decodedTokensCrossAttentionQKs = decodedTokensCrossAttentionQKs.slice(0, sliceEndTokenIndex)
 			decodedTokensConfidence = decodedTokensConfidence.slice(0, sliceEndTokenIndex)
+
+			decodedTokensDecodingTime = decodedTokensDecodingTime.slice(0, sliceEndTokenIndex)
+			decodedTokensInferenceTime = decodedTokensInferenceTime.slice(0, sliceEndTokenIndex)
 		}
 
+		const decodedTokensOverheadTime = decodedTokensDecodingTime.map((time, index) => time - decodedTokensInferenceTime[index])
+
 		logger.write('\n')
 		logger.end()
 
@@ -1162,6 +1207,10 @@ export class Whisper {
 			decodedTokensTimestampLogits,
 			decodedTokensConfidence,
 			decodedTokensCrossAttentionQKs,
+
+			decodedTokensDecodingTime,
+			decodedTokensInferenceTime,
+			decodedTokensOverheadTime,
 		}
 	}
 
@@ -1186,7 +1235,7 @@ export class Whisper {
 		const maxAudioSamples = sampleRate * 30
 		const maxAudioFrames = 3000
 
-		if (sampleRate != 16000) {
+		if (sampleRate !== 16000) {
 			throw new Error('Audio must have a sample rate of 16000 Hz')
 		}
 
@@ -1194,8 +1243,7 @@ export class Whisper {
 			throw new Error(`Audio part is longer than 30 seconds`)
 		}
 
-		// Compute a mel spectogram
-		await logger.startAsync('Extract mel spectogram from audio part')
+		await logger.startAsync('Extract Mel spectrogram from audio part')
 
 		// Pad audio samples to ensure that have a duration of 30 seconds
 		const paddedAudioSamples = new Float32Array(maxAudioSamples)
@@ -1203,40 +1251,40 @@ export class Whisper {
 
 		const rawAudioPart: RawAudio = { audioChannels: [paddedAudioSamples], sampleRate }
 
-		const { melSpectogram } = await computeMelSpectogramUsingFilterbanks(rawAudioPart, fftOrder, fftWindowSize, fftHopLength, filterbanks)
+		// Compute Mel spectrogram
+		const { melSpectrogram } = await computeMelSpectrogramUsingFilterbanks(rawAudioPart, fftOrder, fftWindowSize, fftHopLength, filterbanks)
 
-		await logger.startAsync('Normalize mel spectogram')
+		// Flatten, transpose, apply logarithm and normalize Mel spectrogram
+		await logger.startAsync('Process Mel spectrogram')
 
-		const logMelSpectogram = melSpectogram.map(spectrum => spectrum.map(mel => Math.log10(Math.max(mel, 1e-10))))
+		const flattenedLogMelSpectrogram = new Float32Array(maxAudioFrames * filterbankCount)
 
-		// Find maximum log mel value in the spectrum
 		let maxLogMel = -Infinity
 
-		for (const spectrum of logMelSpectogram) {
-			for (const mel of spectrum) {
-				if (mel > maxLogMel) {
-					maxLogMel = mel
+		for (let i = 0; i < filterbankCount; i++) {
+			for (let j = 0; j < maxAudioFrames; j++) {
+				const mel = melSpectrogram[j][i]
+				const logMel = Math.log10(Math.max(mel, 1e-10))
+
+				if (logMel > maxLogMel) {
+					maxLogMel = logMel
 				}
+
+				flattenedLogMelSpectrogram[(i * maxAudioFrames) + j] = logMel
 			}
 		}
 
-		// Normalize log mel spectogram (based on Python reference code)
-		const normalizedLogMelSpectogram = logMelSpectogram.map(spectrum => spectrum.map(
-			logMel => (Math.max(logMel, maxLogMel - 8) + 4) / 4))
-
-		// Flatten the normalized log mel spectogram
-		const flattenedNormalizedLogMelSpectogram = new Float32Array(maxAudioFrames * filterbankCount)
+		for (let i = 0; i < flattenedLogMelSpectrogram.length; i++) {
+			const logMel = flattenedLogMelSpectrogram[i]
+			const normalizedLogMel = (Math.max(logMel, maxLogMel - 8) + 4) / 4
 
-		for (let i = 0; i < filterbankCount; i++) {
-			for (let j = 0; j < maxAudioFrames; j++) {
-				flattenedNormalizedLogMelSpectogram[(i * maxAudioFrames) + j] = normalizedLogMelSpectogram[j][i]
-			}
+			flattenedLogMelSpectrogram[i] = normalizedLogMel
 		}
 
 		// Run the encoder model
-		await logger.startAsync('Encode mel spectogram with Whisper encoder model')
+		await logger.startAsync('Encode Mel spectrogram with Whisper encoder model')
 
-		const inputTensor = new Onnx.Tensor('float32', flattenedNormalizedLogMelSpectogram, [1, filterbankCount, maxAudioFrames])
+		const inputTensor = new Onnx.Tensor('float32', flattenedLogMelSpectrogram, [1, filterbankCount, maxAudioFrames])
 
 		const encoderInputs = { mel: inputTensor }
 
@@ -1267,7 +1315,7 @@ export class Whisper {
 			return isSeparatorCharacter(text[text.length - 1])
 		}
 
-		if (language != 'zh' && language != 'ja') {
+		if (language !== 'zh' && language !== 'ja') {
 			tokenTimeline = tokenTimeline.filter(entry => this.isTextToken(entry.id!))
 		}
 
@@ -1282,7 +1330,7 @@ export class Whisper {
 			const text = entry.text
 			const previousEntryText = previousEntry?.text
 
-			if (groups.length == 0 ||
+			if (groups.length === 0 ||
 				text === '' ||
 				startsWithSeparatorCharacter(text) ||
 				(previousEntryText != null && endsWithSeparatorCharacter(previousEntryText))) {
@@ -1346,14 +1394,14 @@ export class Whisper {
 	}
 
 	async getTokenTimelineFromAlignmentPath(alignmentPath: AlignmentPath, tokens: number[], startTimeOffset: number, endTimeOffset: number, tokensConfidence?: number[], correctionAmount = 0.0) {
-		if (alignmentPath.length == 0) {
+		if (alignmentPath.length === 0) {
 			return []
 		}
 
 		const tokenTimeline: Timeline = []
 
 		for (let pathIndex = 0; pathIndex < alignmentPath.length; pathIndex++) {
-			if (pathIndex != 0 && alignmentPath[pathIndex].source == alignmentPath[pathIndex - 1].source) {
+			if (pathIndex !== 0 && alignmentPath[pathIndex].source === alignmentPath[pathIndex - 1].source) {
 				continue
 			}
 
@@ -1448,8 +1496,10 @@ export class Whisper {
 				let countOfAllWeightsForHead = 0
 
 				for (const tokenFrames of head) {
-					sumOfAllWeightsForHead += sumVector(tokenFrames)
-					sumOfAllSquaredWeightsForHead += sumOfSquaresForVector(tokenFrames)
+					const { sum, sumOfSquares } = sumAndSumOfSquaresOfVector(tokenFrames)
+
+					sumOfAllWeightsForHead += sum
+					sumOfAllSquaredWeightsForHead += sumOfSquares
 					countOfAllWeightsForHead += tokenFrames.length
 				}
 
@@ -1604,7 +1654,7 @@ export class Whisper {
 
 		const onnxSessionOptions = getOnnxSessionOptions({ executionProviders: this.encoderExecutionProviders })
 
-		const onnxProvidersString = onnxSessionOptions.executionProviders!.length > 0 ? `${ onnxSessionOptions.executionProviders!.join(', ') }` : `default`
+		const onnxProvidersString = onnxSessionOptions.executionProviders!.length > 0 ? `${onnxSessionOptions.executionProviders!.join(', ')}` : `default`
 
 		await logger.startAsync(`Create encoder inference session for model '${this.modelName}' (ONNX provider: ${onnxProvidersString})`)
 
@@ -1642,15 +1692,15 @@ export class Whisper {
 	getKvDimensions(groupCount: number, length: number) {
 		const modelName = this.modelName
 
-		if (modelName == 'tiny' || modelName == 'tiny.en') {
+		if (modelName === 'tiny' || modelName === 'tiny.en') {
 			return [8, groupCount, length, 384]
-		} else if (modelName == 'base' || modelName == 'base.en') {
+		} else if (modelName === 'base' || modelName === 'base.en') {
 			return [12, groupCount, length, 512]
-		} else if (modelName == 'small' || modelName == 'small.en') {
+		} else if (modelName === 'small' || modelName === 'small.en') {
 			return [24, groupCount, length, 768]
-		} else if (modelName == 'medium' || modelName == 'medium.en') {
+		} else if (modelName === 'medium' || modelName === 'medium.en') {
 			return [48, groupCount, length, 1024]
-		} else if (modelName == 'large-v1' || modelName == 'large-v2' || modelName == 'large-v3' || modelName == 'large-v3-turbo') {
+		} else if (modelName === 'large-v1' || modelName === 'large-v2' || modelName === 'large-v3' || modelName === 'large-v3-turbo') {
 			return [64, groupCount, length, 1280]
 		} else {
 			throw new Error(`Unsupported model: ${modelName}`)
@@ -1664,7 +1714,7 @@ export class Whisper {
 
 		if (this.isMultiligualModel) {
 			const languageToken = this.tokenConfig.languageTokensStart + languageIdLookup[language]
-			const taskToken = task == 'translate' ? this.tokenConfig.translateTaskToken : this.tokenConfig.transcribeTaskToken
+			const taskToken = task === 'translate' ? this.tokenConfig.translateTaskToken : this.tokenConfig.transcribeTaskToken
 
 			tokens = [startOfTextToken, languageToken, taskToken]
 		} else {
@@ -1900,7 +1950,7 @@ export async function loadPackagesAndGetPaths(modelName: WhisperModelName | unde
 		if (languageCode) {
 			const shortLanguageCode = getShortLanguageCode(languageCode)
 
-			modelName = shortLanguageCode == 'en' ? 'tiny.en' : 'tiny'
+			modelName = shortLanguageCode === 'en' ? 'tiny.en' : 'tiny'
 		} else {
 			modelName = 'tiny'
 		}
@@ -1918,7 +1968,7 @@ export async function loadPackagesAndGetPaths(modelName: WhisperModelName | unde
 }
 
 export function normalizeWhisperModelName(modelName: WhisperModelName, languageCode: string | undefined): WhisperModelName {
-	if (languageCode != 'en' && modelName.endsWith('.en')) {
+	if (languageCode !== 'en' && modelName.endsWith('.en')) {
 		const originalModelName = modelName
 		modelName = modelName.slice(0, modelName.length - 3) as WhisperModelName
 
@@ -1964,7 +2014,20 @@ export type WhisperTokenData = {
 
 export type WhisperLogitFilter = (logits: number[], decodedTokens: number[], isFirstPart: boolean, isFinalPart: boolean) => number[]
 
-export type WhisperModelName = 'tiny' | 'tiny.en' | 'base' | 'base.en' | 'small' | 'small.en' | 'medium' | 'medium.en' | 'large-v1' | 'large-v2' | 'large-v3' | 'large-v3-turbo'
+export type WhisperModelName =
+	'tiny' |
+	'tiny.en' |
+	'base' |
+	'base.en' |
+	'small' |
+	'small.en' |
+	'medium' |
+	'medium.en' |
+	'large-v1' |
+	'large-v2' |
+	'large-v3' |
+	'large-v3-turbo'
+
 export type WhisperTask = 'transcribe' | 'translate' | 'detect-language'
 
 export const modelNameToPackageName: { [modelName in WhisperModelName]: string } = {
@@ -2084,17 +2147,18 @@ const languageIdLookup: { [s: string]: number } = {
 	'ba': 96,
 	'jw': 97,
 	'su': 98,
+	//'yue': 99
 }
 
 const alignmentHeadsIndexes: { [name in WhisperModelName]: number[] } = {
-	'tiny.en': [6, 12, 17, 18, 19, 20, 21, 22,],
 	'tiny': [14, 18, 20, 21, 22, 23,],
-	'base.en': [27, 39, 41, 45, 47,],
+	'tiny.en': [6, 12, 17, 18, 19, 20, 21, 22,],
 	'base': [25, 34, 35, 39, 41, 42, 44, 46,],
-	'small.en': [78, 84, 87, 92, 98, 101, 103, 108, 112, 116, 118, 120, 121, 122, 123, 126, 131, 134, 136,],
+	'base.en': [27, 39, 41, 45, 47,],
 	'small': [63, 69, 96, 100, 103, 104, 108, 115, 117, 125,],
-	'medium.en': [180, 225, 236, 238, 244, 256, 260, 265, 284, 286, 295, 298, 303, 320, 323, 329, 334, 348,],
+	'small.en': [78, 84, 87, 92, 98, 101, 103, 108, 112, 116, 118, 120, 121, 122, 123, 126, 131, 134, 136,],
 	'medium': [223, 244, 255, 257, 320, 372,],
+	'medium.en': [180, 225, 236, 238, 244, 256, 260, 265, 284, 286, 295, 298, 303, 320, 323, 329, 334, 348,],
 	'large-v1': [199, 222, 224, 237, 447, 451, 457, 462, 475,],
 	'large-v2': [212, 277, 331, 332, 333, 355, 356, 364, 371, 379, 391, 422, 423, 443, 449, 452, 465, 467, 473, 505, 521, 532, 555,],
 	'large-v3': [140, 217, 258, 272, 321, 354, 391, 424, 481, 506,],
@@ -2613,4 +2677,6 @@ export const defaultWhisperVADOptions: WhisperVADOptions = {
 	decoderProvider: undefined,
 }
 
-type WhisperTimestampAccuracy = 'medium' | 'high'
+export type WhisperTimestampAccuracy = 'medium' | 'high'
+
+export type WhisperPartCallback = (partTranscript: string, partTokenTimeline: Timeline, partWordTimeline: Timeline) => void

package/src/source-separation/MDXNetSourceSeparation.ts

@@ -1,6 +1,6 @@
 import type * as Onnx from 'onnxruntime-node'
 import { getEmptyRawAudio, RawAudio } from '../audio/AudioUtilities.js'
-import { getWindowWeights, stftrGenerator, stiftr, WindowType } from '../dsp/FFT.js'
+import { getWindowWeights, createStftrGenerator, stiftr, WindowType } from '../dsp/FFT.js'
 import { logToStderr } from '../utilities/Utilities.js'
 import { Logger } from '../utilities/Logger.js'
 import { OnnxExecutionProvider, dmlProviderAvailable, getOnnxSessionOptions } from '../utilities/OnnxUtilities.js'
@@ -74,8 +74,8 @@ export class MDXNet {
 		const fftSizeReciprocal = 1 / fftSize
 
 		// Initialize generators for STFT frames for each channel
-		const fftFramesLeftGenerator = stftrGenerator(rawAudio.audioChannels[0], fftSize, fftWindowSize, fftHopSize, fftWindowType)
-		const fftFramesRightGenerator = stftrGenerator(rawAudio.audioChannels[1], fftSize, fftWindowSize, fftHopSize, fftWindowType)
+		const fftFramesLeftGenerator = await createStftrGenerator(rawAudio.audioChannels[0], fftSize, fftWindowSize, fftHopSize, fftWindowType)
+		const fftFramesRightGenerator = await createStftrGenerator(rawAudio.audioChannels[1], fftSize, fftWindowSize, fftHopSize, fftWindowType)
 
 		// Initial windowed lists to store recently computed STFT frames
 		const fftFramesLeftWindowedList = new WindowedList<Float32Array>(segmentSize)
@@ -95,13 +95,13 @@ export class MDXNet {
 			}
 
 			while (fftFramesLeftWindowedList.endOffset < segmentEndFrameOffset) {
-				const nextLeftFrameResult = await fftFramesLeftGenerator.next()
+				const nextLeftFrameResult = fftFramesLeftGenerator.next()
 
 				if (nextLeftFrameResult.done) {
 					break
 				}
 
-				const nextRightFrameResult = await fftFramesRightGenerator.next()
+				const nextRightFrameResult = fftFramesRightGenerator.next()
 
 				if (nextRightFrameResult.done) {
 					break

package/src/speech-embeddings/WavToVec2BertFeatureEmbeddings.ts

@@ -2,7 +2,7 @@ import type * as Onnx from 'onnxruntime-node'
 import { OnnxExecutionProvider, getOnnxSessionOptions } from '../utilities/OnnxUtilities.js';
 
 import { RawAudio } from "../audio/AudioUtilities.js";
-import { computeMelSpectogram } from "../dsp/MelSpectogram.js";
+import { computeMelSpectrogram } from "../dsp/MelSpectrogram.js";
 import { Logger } from '../utilities/Logger.js';
 import { concatFloat32Arrays, splitFloat32Array } from '../utilities/Utilities.js';
 import { applyEmphasis } from '../dsp/MFCC.js';
@@ -31,7 +31,7 @@ export class Wav2Vec2BertFeatureEmbeddings {
 
 		rawAudio.audioChannels[0] = applyEmphasis(rawAudio.audioChannels[0], 0.97)
 
-		const { melSpectogram } = await computeMelSpectogram(
+		const { melSpectrogram } = await computeMelSpectrogram(
 			rawAudio,
 			512,
 			400,
@@ -42,8 +42,8 @@ export class Wav2Vec2BertFeatureEmbeddings {
 			'povey')
 
 		// Ensure even length
-		if (melSpectogram.length % 2 != 0) {
-			melSpectogram.push(new Float32Array(80))
+		if (melSpectrogram.length % 2 != 0) {
+			melSpectrogram.push(new Float32Array(80))
 		}
 
 		// Normalize filterbanks
@@ -51,24 +51,24 @@ export class Wav2Vec2BertFeatureEmbeddings {
 			let sum = 0
 			let sumOfSquares = 0
 
-			for (let i = 0; i < melSpectogram.length; i++) {
-				const value = melSpectogram[i][filterbankIndex]
+			for (let i = 0; i < melSpectrogram.length; i++) {
+				const value = melSpectrogram[i][filterbankIndex]
 
 				sum += value
 				sumOfSquares += value ** 2
 			}
 
-			const mean = sum / melSpectogram.length
-			const normalizationFactor = 1 / (Math.sqrt(sumOfSquares / melSpectogram.length) + 1e-40)
+			const mean = sum / melSpectrogram.length
+			const normalizationFactor = 1 / (Math.sqrt(sumOfSquares / melSpectrogram.length) + 1e-40)
 
-			for (let i = 0; i < melSpectogram.length; i++) {
-				melSpectogram[i][filterbankIndex] -= mean
-				melSpectogram[i][filterbankIndex] *= normalizationFactor
+			for (let i = 0; i < melSpectrogram.length; i++) {
+				melSpectrogram[i][filterbankIndex] -= mean
+				melSpectrogram[i][filterbankIndex] *= normalizationFactor
 			}
 		}
 
 		// Flatten
-		const flattenedMelSpectogram = concatFloat32Arrays(melSpectogram)
+		const flattenedMelSpectrogram = concatFloat32Arrays(melSpectrogram)
 
 		// Initialize session
 		await this.initializeSessionIfNeeded()
@@ -77,9 +77,9 @@ export class Wav2Vec2BertFeatureEmbeddings {
 
 		const Onnx = await import('onnxruntime-node')
 
-		const inputTensor = new Onnx.Tensor('float32', flattenedMelSpectogram, [1, melSpectogram.length / 2, 80 * 2])
+		const inputTensor = new Onnx.Tensor('float32', flattenedMelSpectrogram, [1, melSpectrogram.length / 2, 80 * 2])
 
-		const attentionMask = new Int32Array(melSpectogram.length / 2).fill(1)
+		const attentionMask = new Int32Array(melSpectrogram.length / 2).fill(1)
 		const attentionMaskTensor = new Onnx.Tensor('int32', attentionMask, [1, attentionMask.length])
 
 		// Run inference

package/src/utilities/Utilities.ts

@@ -130,13 +130,15 @@ export function sleep(timeMs: number) {
 	return new Promise<void>((resolve) => {
 		const tickCallback = () => {
 			if (timer.elapsedTime < timeMs) {
-				setImmediate(tickCallback)
+				//setImmediate(tickCallback)
+				setTimeout(tickCallback, 0)
 			} else {
 				resolve()
 			}
 		}
 
-		setImmediate(tickCallback)
+		//setImmediate(tickCallback)
+		setTimeout(tickCallback, 0)
 	})
 }
 

package/dist/dsp/MelSpectogram.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"MelSpectogram.js","sourceRoot":"","sources":["../../src/dsp/MelSpectogram.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,MAAM,EAAE,MAAM,wBAAwB,CAAA;AAC/C,OAAO,KAAK,GAAG,MAAM,UAAU,CAAA;AAE/B,MAAM,CAAC,KAAK,UAAU,oBAAoB,CAAC,QAAkB,EAAE,QAAgB,EAAE,UAAkB,EAAE,SAAiB,EAAE,eAAuB,EAAE,gBAAwB,EAAE,gBAAwB,EAAE,aAA6B,MAAM;IACvO,MAAM,MAAM,GAAG,IAAI,MAAM,EAAE,CAAA;IAE3B,MAAM,CAAC,KAAK,CAAC,wBAAwB,CAAC,CAAA;IACtC,MAAM,QAAQ,GAAG,CAAC,QAAQ,GAAG,CAAC,CAAC,GAAG,CAAC,CAAA;IACnC,MAAM,gBAAgB,GAAG,QAAQ,CAAC,UAAU,GAAG,CAAC,CAAA;IAChD,MAAM,cAAc,GAAG,GAAG,CAAC,iBAAiB,CAAC,QAAQ,EAAE,gBAAgB,CAAC,CAAA;IAExE,MAAM,iBAAiB,GAAG,UAAU,CAAC,gBAAgB,CAAC,CAAA;IACtD,MAAM,iBAAiB,GAAG,UAAU,CAAC,gBAAgB,CAAC,CAAA;IAEtD,MAAM,4BAA4B,GAAG,kCAAkC,CAAC,eAAe,EAAE,iBAAiB,EAAE,iBAAiB,CAAC,CAAA;IAC9H,MAAM,cAAc,GAAG,iBAAiB,CAAC,cAAc,EAAE,4BAA4B,EAAE,iBAAiB,EAAE,iBAAiB,CAAC,CAAA;IAE5H,MAAM,CAAC,GAAG,EAAE,CAAA;IAEZ,OAAO,oCAAoC,CAAC,QAAQ,EAAE,QAAQ,EAAE,UAAU,EAAE,SAAS,EAAE,cAAc,EAAE,UAAU,CAAC,CAAA;AACnH,CAAC;AAED,MAAM,CAAC,KAAK,UAAU,oCAAoC,CAAC,QAAkB,EAAE,QAAgB,EAAE,UAAkB,EAAE,SAAiB,EAAE,WAAyB,EAAE,aAA6B,MAAM;IACrM,MAAM,MAAM,GAAG,IAAI,MAAM,EAAE,CAAA;IAE3B,MAAM,CAAC,KAAK,CAAC,yBAAyB,CAAC,CAAA;IACvC,MAAM,YAAY,GAAG,QAAQ,CAAC,aAAa,CAAC,CAAC,CAAC,CAAA;IAC9C,MAAM,SAAS,GAAG,MAAM,GAAG,CAAC,KAAK,CAAC,YAAY,EAAE,QAAQ,EAAE,UAAU,EAAE,SAAS,EAAE,UAAU,CAAC,CAAA;IAE5F,MAAM,CAAC,KAAK,CAAC,wCAAwC,CAAC,CAAA;IACtD,MAAM,aAAa,GAAG,wBAAwB,CAAC,SAAS,EAAE,WAAW,CAAC,CAAA;IAEtE,MAAM,CAAC,GAAG,EAAE,CAAA;IAEZ,OAAO,EAAE,aAAa,EAAE,SAAS,EAAE,CAAA;AACpC,CAAC;AAED,MAAM,UAAU,wBAAwB,CAAC,SAAyB,EAAE,cAA4B;IAC/F,OAAO,SAAS,CAAC,GAAG,CAAC,QAAQ,CAAC,EAAE;QAC/B,MAAM,aAAa,GAAG,GAAG,CAAC,uBAAuB,CAAC,QAAQ,CAAC,CAAA;QAC3D,OAAO,0BAA0B,CAAC,aAAa,EAAE,cAAc,CAAC,CAAA;IACjE,CAAC,CAAC,CAAA;AACH,CAAC;AAED,MAAM,UAAU,0BAA0B,CAAC,aAA2B,EAAE,WAAyB;IAChG,MAAM,eAAe,GAAG,WAAW,CAAC,MAAM,CAAA;IAC1C,MAAM,WAAW,GAAG,IAAI,YAAY,CAAC,eAAe,CAAC,CAAA;IAErD,KAAK,IAAI,YAAY,GAAG,CAAC,EAAE,YAAY,GAAG,eAAe,EAAE,YAAY,EAAE,EAAE,CAAC;QAC3E,MAAM,UAAU,GAAG,WAAW,CAAC,YAAY,CAAC,CAAA;QAC5C,MAAM,oBAAoB,GAAG,UAAU,CAAC,UAAU,CAAA;QAClD,MAAM,iBAAiB,GAAG,UAAU,CAAC,OAAO,CAAA;QAE5C,IAAI,oBAAoB,KAAK,CAAC,CAAC,EAAE,CAAC;YACjC,SAAQ;QACT,CAAC;QAED,IAAI,YAAY,GAAG,CAAC,CAAA;QAEpB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,iBAAiB,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACnD,MAAM,kBAAkB,GAAG,oBAAoB,GAAG,CAAC,CAAA;YAEnD,IAAI,kBAAkB,IAAI,aAAa,CAAC,MAAM,EAAE,CAAC;gBAChD,MAAK;YACN,CAAC;YAED,MAAM,MAAM,GAAG,iBAAiB,CAAC,CAAC,CAAC,CAAA;YACnC,MAAM,kBAAkB,GAAG,aAAa,CAAC,kBAAkB,CAAC,CAAA;YAE5D,YAAY,IAAI,MAAM,GAAG,kBAAkB,CAAA;QAC5C,CAAC;QAED,WAAW,CAAC,YAAY,CAAC,GAAG,YAAY,CAAA;IACzC,CAAC;IAED,OAAO,WAAW,CAAA;AACnB,CAAC;AAED,MAAM,UAAU,iBAAiB,CAAC,0BAAwC,EAAE,oBAAkC,EAAE,iBAAyB,EAAE,iBAAyB;IACnK,MAAM,eAAe,GAAG,oBAAoB,CAAC,MAAM,CAAA;IACnD,MAAM,2BAA2B,GAAG,0BAA0B,CAAC,GAAG,CAAC,WAAW,CAAC,EAAE,CAAC,UAAU,CAAC,WAAW,CAAC,CAAC,CAAA;IAE1G,MAAM,WAAW,GAAiB,EAAE,CAAA;IAEpC,KAAK,IAAI,eAAe,GAAG,CAAC,EAAE,eAAe,GAAG,eAAe,EAAE,eAAe,EAAE,EAAE,CAAC;QACpF,MAAM,eAAe,GAAG,oBAAoB,CAAC,eAAe,CAAC,CAAA;QAE7D,MAAM,aAAa,GAAG,eAAe,GAAG,CAAC,CAAC,CAAC,CAAC,oBAAoB,CAAC,eAAe,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,iBAAiB,CAAA;QACzG,MAAM,cAAc,GAAG,eAAe,GAAG,eAAe,GAAG,CAAC,CAAC,CAAC,CAAC,oBAAoB,CAAC,eAAe,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,iBAAiB,CAAA;QAE5H,MAAM,KAAK,GAAG,cAAc,GAAG,aAAa,CAAA;QAC5C,MAAM,SAAS,GAAG,KAAK,GAAG,CAAC,CAAA;QAE3B,IAAI,UAAU,GAAG,CAAC,CAAC,CAAA;QACnB,IAAI,OAAO,GAAa,EAAE,CAAA;QAE1B,IAAI,SAAS,GAAG,CAAC,CAAA;QAEjB,KAAK,IAAI,sBAAsB,GAAG,CAAC,EAAE,sBAAsB,GAAG,2BAA2B,CAAC,MAAM,EAAE,sBAAsB,EAAE,EAAE,CAAC;YAC5H,MAAM,6BAA6B,GAAG,2BAA2B,CAAC,sBAAsB,CAAC,CAAA;YAEzF,IAAI,MAAM,GAAG,CAAC,CAAA;YAEd,IAAI,6BAA6B,IAAI,aAAa,IAAI,6BAA6B,IAAI,eAAe,EAAE,CAAC;gBACxG,MAAM,GAAG,CAAC,6BAA6B,GAAG,aAAa,CAAC,GAAG,SAAS,CAAA;YACrE,CAAC;iBAAM,IAAI,6BAA6B,GAAG,eAAe,IAAI,6BAA6B,IAAI,cAAc,EAAE,CAAC;gBAC/G,MAAM,GAAG,CAAC,cAAc,GAAG,6BAA6B,CAAC,GAAG,SAAS,CAAA;YACtE,CAAC;YAED,IAAI,MAAM,GAAG,CAAC,EAAE,CAAC;gBAChB,IAAI,UAAU,IAAI,CAAC,CAAC,EAAE,CAAC;oBACtB,UAAU,GAAG,sBAAsB,CAAA;gBACpC,CAAC;gBAED,OAAO,CAAC,IAAI,CAAC,MAAM,CAAC,CAAA;gBACpB,SAAS,IAAI,MAAM,CAAA;YACpB,CAAC;iBAAM,IAAI,UAAU,IAAI,CAAC,CAAC,EAAE,CAAC;gBAC7B,MAAK;YACN,CAAC;QACF,CAAC;QAED,OAAO,GAAG,OAAO,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,CAAC,MAAM,GAAG,SAAS,CAAC,CAAA;QAEnD,WAAW,CAAC,IAAI,CAAC,EAAE,UAAU,EAAE,OAAO,EAAE,CAAC,CAAA;IAC1C,CAAC;IAED,OAAO,WAAW,CAAA;AACnB,CAAC;AAED,MAAM,UAAU,kCAAkC,CAAC,YAAoB,EAAE,iBAAyB,EAAE,iBAAyB;IAC5H,MAAM,WAAW,GAAG,CAAC,iBAAiB,GAAG,iBAAiB,CAAC,GAAG,CAAC,YAAY,GAAG,CAAC,CAAC,CAAA;IAEhF,MAAM,iBAAiB,GAAG,IAAI,YAAY,CAAC,YAAY,CAAC,CAAA;IAExD,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,YAAY,EAAE,CAAC,EAAE,EAAE,CAAC;QACvC,iBAAiB,CAAC,CAAC,CAAC,GAAG,iBAAiB,GAAG,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,WAAW,CAAC,CAAA;IACnE,CAAC;IAED,OAAO,iBAAiB,CAAA;AACzB,CAAC;AAED,MAAM,UAAU,UAAU,CAAC,SAAiB;IAC3C,OAAO,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,GAAG,CAAC,SAAS,GAAG,KAAK,CAAC,CAAC,CAAA;AACtD,CAAC;AAED,MAAM,UAAU,UAAU,CAAC,GAAW;IACrC,OAAO,KAAK,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,IAAI,EAAE,GAAG,GAAG,MAAM,CAAC,GAAG,GAAG,CAAC,CAAA;AACpD,CAAC"}