echogarden 1.6.2 → 1.7.0

package/src/source-separation/MDXNetSourceSeparation.ts

@@ -1,28 +1,32 @@
 import type * as Onnx from 'onnxruntime-node'
-import { RawAudio } from '../audio/AudioUtilities.js';
-import { binBufferToComplex, complexToBinBuffer, getWindowWeights, stftr, stiftr } from '../dsp/FFT.js';
-import { ComplexNumber } from '../math/VectorMath.js';
-import { logToStderr } from '../utilities/Utilities.js';
-import { Logger } from '../utilities/Logger.js';
-import { OnnxExecutionProvider, getOnnxSessionOptions } from '../utilities/OnnxUtilities.js';
+import { getEmptyRawAudio, RawAudio } from '../audio/AudioUtilities.js'
+import { getWindowWeights, stftrGenerator, 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'
+import chalk from 'chalk'
+import { WindowedList } from '../utilities/WindowedList.js'
 
 const log = logToStderr
 
 export async function isolate(
 	rawAudio: RawAudio,
 	modelFilePath: string,
+	modelProfile: MDXNetModelProfile,
 	executionProviders: OnnxExecutionProvider[]) {
 
-	const model = new MDXNet(modelFilePath, executionProviders)
+	const model = new MDXNet(modelFilePath, modelProfile, executionProviders)
 
 	return model.processAudio(rawAudio)
 }
 
 export class MDXNet {
 	session?: Onnx.InferenceSession
+	onnxSessionOptions?: Onnx.InferenceSession.SessionOptions
 
 	constructor(
 		public readonly modelFilePath: string,
+		public readonly modelProfile: MDXNetModelProfile,
 		public readonly executionProviders: OnnxExecutionProvider[]) {
 	}
 
@@ -31,77 +35,104 @@ export class MDXNet {
 			throw new Error(`Input audio must be stereo`)
 		}
 
-		if (rawAudio.sampleRate != 44100) {
-			throw new Error(`Input audio must have a sample rate of 44100 Hz`)
+		if (rawAudio.sampleRate !== this.modelProfile.sampleRate) {
+			throw new Error(`Input audio must have a sample rate of ${this.modelProfile.sampleRate} Hz`)
 		}
 
+		if (rawAudio.audioChannels[0].length === 0) {
+			return getEmptyRawAudio(rawAudio.audioChannels.length, rawAudio.sampleRate)
+		}
+
+		const logger = new Logger()
+
+		await logger.startAsync(`Initialize MDXNet model`)
+
 		await this.initializeSessionIfNeeded()
 
 		const Onnx = await import('onnxruntime-node')
 
-		const logger = new Logger()
-
-		const session = this.session!
+		const sampleRate = this.modelProfile.sampleRate
+		const fftSize = this.modelProfile.fftSize
+		const fftWindowSize = this.modelProfile.fftWindowSize
+		const fftHopSize = this.modelProfile.fftHopSize
+		const fftWindowType = this.modelProfile.fftWindowType
 
-		const sampleRate = rawAudio.sampleRate
-		const fftSize = 6144
-		const fftCount = 2048
-		const fftWindowSize = fftSize
-		const fftHopSize = 1024
+		const binCount = this.modelProfile.binCount
 
-		const segmentSize = 256
-		const segmentHopSize = 240
+		const segmentSize = this.modelProfile.segmentSize
+		const segmentHopSize = this.modelProfile.segmentHopSize
 
 		const sampleCount = rawAudio.audioChannels[0].length
 
-		logger.start('Compute STFT of full waveform')
+		const fftSizeReciprocal = 1 / fftSize
 
-		const fftFramesLeft = await stftr(rawAudio.audioChannels[0], fftSize, fftWindowSize, fftHopSize, 'hann')
-		const fftFramesRight = await stftr(rawAudio.audioChannels[1], fftSize, fftWindowSize, fftHopSize, 'hann')
+		// 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 fftFramesLeftComplex = fftFramesLeft.map(frame => binBufferToComplex(frame).slice(0, fftCount))
-		const fftFramesRightComplex = fftFramesRight.map(frame => binBufferToComplex(frame).slice(0, fftCount))
+		// Initial windowed lists to store recently computed STFT frames
+		const fftFramesLeftWindowedList = new WindowedList<Float32Array>(segmentSize)
+		const fftFramesRightWindowedList = new WindowedList<Float32Array>(segmentSize)
 
 		const audioForSegments: Float32Array[][] = []
 
-		for (let segmentOffset = 0; segmentOffset < fftFramesLeft.length; segmentOffset += segmentHopSize) {
-			const timePosition = segmentOffset * (fftHopSize / sampleRate)
+		for (let segmentStartFrameOffset = 0; ; segmentStartFrameOffset += segmentHopSize) {
+			const segmentEndFrameOffset = segmentStartFrameOffset + segmentSize
+
+			const timePosition = segmentStartFrameOffset * (fftHopSize / sampleRate)
+
+			await logger.startAsync(`Compute STFT of segment at time position ${timePosition.toFixed(2)}`, undefined, chalk.magentaBright)
+
+			while (fftFramesLeftWindowedList.endOffset < segmentEndFrameOffset) {
+				const nextLeftFrameResult = await fftFramesLeftGenerator.next()
+
+				if (nextLeftFrameResult.done) {
+					break
+				}
+
+				const nextRightFrameResult = await fftFramesRightGenerator.next()
+
+				if (nextRightFrameResult.done) {
+					break
+				}
+
+				fftFramesLeftWindowedList.add(nextLeftFrameResult.value)
+				fftFramesRightWindowedList.add(nextRightFrameResult.value)
+			}
+
+			const fftFramesForSegment = [
+				fftFramesLeftWindowedList.slice(segmentStartFrameOffset, segmentEndFrameOffset),
+				fftFramesRightWindowedList.slice(segmentStartFrameOffset, segmentEndFrameOffset)
+			]
 
-			logger.start(`Process segment at time position ${timePosition.toFixed(2)}`)
+			const segmentLength = fftFramesForSegment[0].length
 
-			const fftFramesLeftComplexForSegment = fftFramesLeftComplex.slice(segmentOffset, segmentOffset + segmentSize)
-			const fftFramesRightComplexForSegment = fftFramesRightComplex.slice(segmentOffset, segmentOffset + segmentSize)
+			const isLastSegment = segmentLength < segmentSize
 
-			const segmentLength = fftFramesLeftComplexForSegment.length
+			await logger.startAsync(`Reshape STFT frames`)
 
-			const flattenedInputTensor = new Float32Array(1 * 4 * fftCount * segmentSize)
+			const flattenedInputTensor = new Float32Array(1 * 4 * binCount * segmentSize)
 
 			{
 				let writePosition = 0
 
-				for (let tensorChannelIndex = 0; tensorChannelIndex < 4; tensorChannelIndex++) {
-					const isEvenTensorChannelIndex = tensorChannelIndex % 2 === 0
-					const inChannelIndex = tensorChannelIndex < 2 ? 0 : 1
+				// 4 tensor elements are structured as:
+				// <Channel 0 real> <Channel 0 imaginary> <Channel 1 real> <Channel 1 imaginary>
+				for (let tensorElementIndex = 0; tensorElementIndex < 4; tensorElementIndex++) {
+					const isRealComponentTensorElementIndex = tensorElementIndex % 2 === 0
+					const audioChannelIndex = tensorElementIndex < 2 ? 0 : 1
 
-					for (let binIndex = 0; binIndex < fftCount; binIndex++) {
+					for (let binIndex = 0; binIndex < binCount; binIndex++) {
 						for (let frameIndex = 0; frameIndex < segmentSize; frameIndex++) {
 							let value = 0
 
-							if (frameIndex < segmentLength && binIndex >= 0) {
-								let frame: ComplexNumber[]
-
-								if (inChannelIndex === 0) {
-									frame = fftFramesLeftComplexForSegment[frameIndex]
-								} else {
-									frame = fftFramesRightComplexForSegment[frameIndex]
-								}
-
-								const bin = frame[binIndex]
+							if (frameIndex < segmentLength) {
+								const frame = fftFramesForSegment[audioChannelIndex][frameIndex]
 
-								if (isEvenTensorChannelIndex) {
-									value = bin.real
+								if (isRealComponentTensorElementIndex) {
+									value = frame[binIndex << 1]
 								} else {
-									value = bin.imaginary
+									value = frame[(binIndex << 1) + 1]
 								}
 							}
 
@@ -111,51 +142,47 @@ export class MDXNet {
 				}
 			}
 
-			const inputTensor = new Onnx.Tensor('float32', flattenedInputTensor, [1, 4, 2048, 256])
+			await logger.startAsync(`Process with MDXNet model (ONNX provider: ${this.onnxSessionOptions!.executionProviders!.join(', ')})`)
 
-			logger.start('Run MDXNet model')
+			const inputTensor = new Onnx.Tensor('float32', flattenedInputTensor, [1, 4, binCount, segmentSize])
 
-			const { output: outputTensor } = await session.run({ input: inputTensor })
+			const { output: outputTensor } = await this.session!.run({ input: inputTensor })
 
-			logger.start('Process MDXNet model output')
+			await logger.startAsync('Reshape processed frames')
 
 			const flattenedOutputTensor = outputTensor.data as Float32Array
 
-			const outputChannelComplexFrames: ComplexNumber[][][] = []
+			const outputSegmentFramesForChannel: Float32Array[][] = []
 
 			{
 				for (let outChannelIndex = 0; outChannelIndex < 2; outChannelIndex++) {
-					const framesForChannel: ComplexNumber[][] = []
+					const framesForChannel: Float32Array[] = []
 
-					for (let frameIndex = 0; frameIndex < 256; frameIndex++) {
-						const frame: ComplexNumber[] = []
-
-						for (let binIndex = 0; binIndex < fftSize; binIndex++) {
-							frame.push({ real: 0, imaginary: 0 })
-						}
+					for (let frameIndex = 0; frameIndex < segmentSize; frameIndex++) {
+						const frame = new Float32Array(fftSize)
 
 						framesForChannel.push(frame)
 					}
 
-					outputChannelComplexFrames.push(framesForChannel)
+					outputSegmentFramesForChannel.push(framesForChannel)
 				}
 
 				let readPosition = 0
 
 				for (let tensorChannelIndex = 0; tensorChannelIndex < 4; tensorChannelIndex++) {
-					const isEvenTensorChannelIndex = tensorChannelIndex % 2 === 0
-					const outChannelIndex = tensorChannelIndex < 2 ? 0 : 1
+					const isRealTensorChannelIndex = tensorChannelIndex % 2 === 0
+					const audioChannelIndex = tensorChannelIndex < 2 ? 0 : 1
 
-					const binsForOutputChannel = outputChannelComplexFrames[outChannelIndex]
+					const framesForOutputChannel = outputSegmentFramesForChannel[audioChannelIndex]
 
-					for (let binIndex = 0; binIndex < 2048; binIndex++) {
-						for (let frameIndex = 0; frameIndex < 256; frameIndex++) {
-							const bin = binsForOutputChannel[frameIndex][binIndex]
+					for (let binIndex = 0; binIndex < binCount; binIndex++) {
+						for (let frameIndex = 0; frameIndex < segmentSize; frameIndex++) {
+							const outFrame = framesForOutputChannel[frameIndex]
 
-							if (isEvenTensorChannelIndex) {
-								bin.real = flattenedOutputTensor[readPosition++]
+							if (isRealTensorChannelIndex) {
+								outFrame[binIndex << 1] = flattenedOutputTensor[readPosition++] * fftSizeReciprocal
 							} else {
-								bin.imaginary = flattenedOutputTensor[readPosition++]
+								outFrame[(binIndex << 1) + 1] = flattenedOutputTensor[readPosition++] * fftSizeReciprocal
 							}
 						}
 					}
@@ -164,29 +191,29 @@ export class MDXNet {
 
 			const outputAudioChannels: Float32Array[] = []
 
-			logger.start(`Compute inverse STFT for segment`)
-			for (let channelIndex = 0; channelIndex < 2; channelIndex++) {
-				const fftSizeReciprocal = 1 / fftSize
-
-				let outputChannelFlattenedFrames = outputChannelComplexFrames[channelIndex]
-					.map(frame => complexToBinBuffer(frame).map(value => value * fftSizeReciprocal))
+			await logger.startAsync(`Compute inverse STFT of model output for segment`)
 
+			for (let channelIndex = 0; channelIndex < 2; channelIndex++) {
 				const samples = await stiftr(
-					outputChannelFlattenedFrames,
+					outputSegmentFramesForChannel[channelIndex],
 					fftSize,
 					fftWindowSize,
 					fftHopSize,
-					'hann')
+					fftWindowType)
 
 				outputAudioChannels.push(samples)
 			}
 
 			audioForSegments.push(outputAudioChannels)
+
+			if (isLastSegment) {
+				break
+			}
 		}
 
 		// Join segments using overlapping Hann windows
-		logger.start(`Join segments`)
-		const concatenatedAudioChannels = [new Float32Array(sampleCount), new Float32Array(sampleCount)]
+		await logger.startAsync(`Join segments`)
+		const joinedSegments = [new Float32Array(sampleCount), new Float32Array(sampleCount)]
 
 		{
 			const segmentCount = audioForSegments.length
@@ -213,7 +240,7 @@ export class MDXNet {
 					const weight = windowWeights[segmentSampleOffset]
 
 					for (let channelIndex = 0; channelIndex < 2; channelIndex++) {
-						concatenatedAudioChannels[channelIndex][sampleIndex] += segmentSamples[channelIndex][segmentSampleOffset] * weight
+						joinedSegments[channelIndex][sampleIndex] += segmentSamples[channelIndex][segmentSampleOffset] * weight
 					}
 
 					sumOfWeightsForSample[sampleIndex] += weight
@@ -222,12 +249,12 @@ export class MDXNet {
 
 			for (let sampleIndex = 0; sampleIndex < sampleCount; sampleIndex++) {
 				for (let channelIndex = 0; channelIndex < 2; channelIndex++) {
-					concatenatedAudioChannels[channelIndex][sampleIndex] /= sumOfWeightsForSample[sampleIndex] + 1e-8
+					joinedSegments[channelIndex][sampleIndex] /= sumOfWeightsForSample[sampleIndex] + 1e-8
 				}
 			}
 		}
 
-		const isolatedRawAudio: RawAudio = { audioChannels: concatenatedAudioChannels, sampleRate }
+		const isolatedRawAudio: RawAudio = { audioChannels: joinedSegments, sampleRate }
 
 		logger.end()
 
@@ -241,8 +268,79 @@ export class MDXNet {
 
 		const Onnx = await import('onnxruntime-node')
 
-		const onnxSessionOptions = getOnnxSessionOptions({ executionProviders: this.executionProviders })
+		this.onnxSessionOptions = getOnnxSessionOptions({ executionProviders: this.executionProviders })
+
+		this.session = await Onnx.InferenceSession.create(this.modelFilePath, this.onnxSessionOptions)
+	}
+}
+
+export function getDefaultMDXNetProviders() : OnnxExecutionProvider[] {
+	if (dmlProviderAvailable()) {
+		return ['dml', 'cpu']
+	} else {
+		return []
+	}
+}
+
+export function getProfileForMDXNetModelName(modelName: MDXNetModelName) {
+	if (['UVR_MDXNET_1_9703', 'UVR_MDXNET_2_9682', 'UVR_MDXNET_3_9662', 'UVR_MDXNET_KARA'].includes(modelName)) {
+		return mdxNetModelProfile1
+	}
 
-		this.session = await Onnx.InferenceSession.create(this.modelFilePath, onnxSessionOptions)
+	if (['UVR_MDXNET_Main', 'Kim_Vocal_1', 'Kim_Vocal_2'].includes(modelName)) {
+		return mdxNetModelProfile2
 	}
+
+	throw new Error(`Unsupported model name: '${modelName}'`)
+}
+
+export const mdxNetModelProfile1: MDXNetModelProfile = {
+	sampleRate: 44100,
+
+	fftSize: 6144,
+	fftWindowSize: 6144,
+	fftHopSize: 1024,
+	fftWindowType: 'hann',
+
+	binCount: 2048,
+
+	segmentSize: 256,
+	segmentHopSize: 224,
+}
+
+export const mdxNetModelProfile2: MDXNetModelProfile = {
+	sampleRate: 44100,
+
+	fftSize: 7680,
+	fftWindowSize: 7680,
+	fftHopSize: 1024,
+	fftWindowType: 'hann',
+
+	binCount: 3072,
+
+	segmentSize: 256,
+	segmentHopSize: 224,
 }
+
+export interface MDXNetModelProfile {
+	sampleRate: number
+
+	fftSize: number
+	fftWindowSize: number
+	fftHopSize: number
+	fftWindowType: WindowType
+
+	binCount: number
+
+	segmentSize: number
+	segmentHopSize: number
+}
+
+export type MDXNetModelName =
+	'UVR_MDXNET_1_9703' |
+	'UVR_MDXNET_2_9682' |
+	'UVR_MDXNET_3_9662' |
+	'UVR_MDXNET_KARA' |
+	'UVR_MDXNET_Main' |
+	'Kim_Vocal_1' |
+	'Kim_Vocal_2'

package/src/subtitles/Subtitles.ts

@@ -1,11 +1,12 @@
 import { htmlToText } from 'html-to-text'
 
-import { formatHMS, formatMS, secondsToHMS, secondsToMS, startsWithAnyOf } from '../utilities/Utilities.js'
+import { secondsToHMS, secondsToMS } from '../utilities/Utilities.js'
 import { isWord, isWordOrSymbolWord } from '../nlp/Segmentation.js'
 import { charactersToWriteAhead } from '../audio/AudioPlayer.js'
 import { Timeline, TimelineEntry } from '../utilities/Timeline.js'
 import { readFile } from '../utilities/FileSystem.js'
 import { deepClone } from '../utilities/ObjectUtilities.js'
+import { formatHMS, formatMS, startsWithAnyOf } from '../utilities/StringUtilities.js'
 
 export async function subtitlesFileToText(filename: string) {
 	return subtitlesToText(await readFile(filename, 'utf8'))

package/src/synthesis/EspeakTTS.ts

@@ -1,10 +1,10 @@
-import { concatFloat32Arrays, logToStderr, objToString, simplifyPunctuationCharacters } from '../utilities/Utilities.js'
+import { concatFloat32Arrays, logToStderr, formatObjectToString } from '../utilities/Utilities.js'
 import { int16PcmToFloat32 } from '../audio/AudioBufferConversion.js'
 import { Logger } from '../utilities/Logger.js'
 import { WasmMemoryManager } from '../utilities/WasmMemoryManager.js'
 import { RawAudio, getEmptyRawAudio } from '../audio/AudioUtilities.js'
 import { playAudioWithTimelinePhones } from '../audio/AudioPlayer.js'
-import { getNormalizedFragmentsForSpeech } from '../nlp/TextNormalizer.js'
+import { getNormalizedFragmentsForSpeech, simplifyPunctuationCharacters } from '../nlp/TextNormalizer.js'
 import { ipaPhoneToKirshenbaum } from '../nlp/PhoneConversion.js'
 import { splitToWords, wordCharacterPattern } from '../nlp/Segmentation.js'
 import { Lexicon, tryGetFirstLexiconSubstitution } from '../nlp/Lexicon.js'
@@ -276,11 +276,11 @@ export async function synthesizeFragments(fragments: string[], espeakOptions: Es
 				const markerIndex = parseInt(markerName.substring(2))
 
 				if (markerIndex != wordIndex) {
-					throw new Error(`Word start marker for index ${wordIndex} is not consistent with word index. The words were: ${objToString(fragments)}`)
+					throw new Error(`Word start marker for index ${wordIndex} is not consistent with word index. The words were: ${formatObjectToString(fragments)}`)
 				}
 
 				if (currentPhoneTimeline.length > 0) {
-					throw new Error(`Word entry ${wordIndex} already has phones before its start marker was seen. The words were: ${objToString(fragments)}`)
+					throw new Error(`Word entry ${wordIndex} already has phones before its start marker was seen. The words were: ${formatObjectToString(fragments)}`)
 				}
 
 				currentWordEntry.startTime = eventTime
@@ -289,7 +289,7 @@ export async function synthesizeFragments(fragments: string[], espeakOptions: Es
 				const markerIndex = parseInt(markerName.substring(2))
 
 				if (markerIndex != wordIndex) {
-					throw new Error(`Word end marker for index ${wordIndex} is not consistent with word index. The words were: ${objToString(fragments)}`)
+					throw new Error(`Word end marker for index ${wordIndex} is not consistent with word index. The words were: ${formatObjectToString(fragments)}`)
 				}
 
 				currentWordEntry.startTime = currentTokenTimeline[0].startTime

package/src/tests/Test.ts

@@ -1,4 +1,3 @@
-import { makeTarballsForInstalledPackages } from '../utilities/TarballMaker.js'
 import { logToStderr, setupProgramTerminationListeners } from '../utilities/Utilities.js'
 
 const log = logToStderr
@@ -8,7 +7,4 @@ setupProgramTerminationListeners()
 //process.env.NODE_TLS_REJECT_UNAUTHORIZED = '0'
 //process.env.http_proxy = 'http://localhost:8080'
 
-//await makeTarballsForInstalledPackages(true)
-
 process.exit(0)
-