diff --git a/src/ai.ts b/src/ai.ts index 4483632..c1ef133 100644 --- a/src/ai.ts +++ b/src/ai.ts @@ -21,7 +21,7 @@ export type AiOptions = LLMOptions & { export type WhisperModel = 'tiny' | 'base' | 'small' | 'medium' | 'large'; export class Ai { - private downloads: {[key: string]: Promise} = {}; + private downloads: {[key: string]: Promise} = {}; private whisperModel!: string; /** Large Language Models */ @@ -43,11 +43,11 @@ export class Ai { */ async asr(path: string, model?: WhisperModel): Promise { if(!this.options.whisper?.binary) throw new Error('Whisper not configured'); - if(!model) model = this.options.whisper.model; - await this.downloadAsrModel(); + await this.downloadAsrModel(model); const name = Math.random().toString(36).substring(2, 10) + '-' + path.split('/').pop() + '.txt'; const output = Path.join(this.options.whisper.temp || '/tmp', name); - await $`rm -f ${output} && ${this.options.whisper.binary} -nt -np -m ${this.whisperModel} -f ${path} -otxt -of ${output}`; + console.log(`rm -f ${output} && ${this.options.whisper.binary} -nt -np -m ${model ? Path.join(this.options.whisper.path, model) : this.whisperModel} -f ${path} -otxt -of ${output}`); + await $`rm -f ${output} && ${this.options.whisper.binary} -nt -np -m ${model ? Path.join(this.options.whisper.path, model) : this.whisperModel} -f ${path} -otxt -of ${output}`; return fs.readFile(output, 'utf-8').then(text => text?.trim() || null) .finally(() => fs.rm(output, {force: true}).catch(() => {})); } @@ -56,19 +56,24 @@ export class Ai { * Downloads the specified Whisper model if it is not already present locally. * * @param {string} model Whisper model that will be downloaded - * @return {Promise} A promise that resolves once the model is downloaded and saved locally. + * @return {Promise} Absolute path to model file, resolves once downloaded */ - async downloadAsrModel(): Promise { + async downloadAsrModel(model?: string): Promise { if(!this.options.whisper?.binary) throw new Error('Whisper not configured'); - if(await fs.stat(this.whisperModel).then(() => true).catch(() => false)) return; - const model = this.whisperModel.split('/').at(-1); - if(!!this.downloads[model]) return this.downloads[model]; - this.downloads[model] = fetch(`https://huggingface.co/ggerganov/whisper.cpp/resolve/main/${model}`) - .then(resp => resp.arrayBuffer()).then(arr => Buffer.from(arr)).then(async buffer => { - await fs.writeFile(this.whisperModel, buffer); - delete this.downloads[model]; + let m; + if(model) m = model?.endsWith('.bin') ? model : model + '.bin'; + else m = this.whisperModel.split('/').at(-1); + const p = Path.join(this.options.whisper.path, m); + if(await fs.stat(p).then(() => true).catch(() => false)) return p; + if(!!this.downloads[m]) return this.downloads[m]; + this.downloads[m] = fetch(`https://huggingface.co/ggerganov/whisper.cpp/resolve/main/${m}`) + .then(resp => resp.arrayBuffer()) + .then(arr => Buffer.from(arr)).then(async buffer => { + await fs.writeFile(Path.join((this.options.whisper).path, m), buffer); + delete this.downloads[m]; + return p; }); - return this.downloads[model]; + return this.downloads[m]; } /**