From c50e5abf7ba3132793ac8f4d9cd645d0b09d24cc Mon Sep 17 00:00:00 2001 From: an-lee Date: Thu, 20 Jun 2024 14:23:11 +0800 Subject: [PATCH] Fix whisper should abort when cancel (#692) * add whisper abort * abort whisper when exit meida page * add abort controller for youtubedr --- enjoy/src/main/whisper.ts | 21 +++++++++++-- enjoy/src/main/youtubedr.ts | 30 ++++++++++++++----- enjoy/src/preload.ts | 3 ++ .../context/media-player-provider.tsx | 22 +++++++++----- .../src/renderer/hooks/use-transcriptions.tsx | 6 ++++ enjoy/src/types/enjoy-app.d.ts | 1 + 6 files changed, 66 insertions(+), 17 deletions(-) diff --git a/enjoy/src/main/whisper.ts b/enjoy/src/main/whisper.ts index f66fce85..776e3bb8 100644 --- a/enjoy/src/main/whisper.ts +++ b/enjoy/src/main/whisper.ts @@ -22,6 +22,7 @@ class Whipser { private binMain: string; private bundledModelsDir: string; public config: WhisperConfigType; + private abortController: AbortController; constructor() { const customWhisperPath = path.join( @@ -87,6 +88,9 @@ class Whipser { } async check() { + this.abortController?.abort(); + this.abortController = new AbortController(); + const model = this.currentModel(); logger.debug(`Checking whisper model: ${model.savePath}`); @@ -107,6 +111,7 @@ class Whipser { commands.join(" "), { timeout: PROCESS_TIMEOUT, + signal: this.abortController.signal, }, (error, stdout, stderr) => { if (error) { @@ -149,6 +154,9 @@ class Whipser { ): Promise> { logger.debug("transcribing from local"); + this.abortController?.abort(); + this.abortController = new AbortController(); + const { blob } = params; let { file } = params; @@ -199,6 +207,7 @@ class Whipser { const command = spawn(this.binMain, commandArguments, { timeout: PROCESS_TIMEOUT, + signal: this.abortController.signal, }); return new Promise((resolve, reject) => { @@ -224,8 +233,8 @@ class Whipser { reject(err); }); - command.on("close", () => { - if (fs.pathExistsSync(outputFile)) { + command.on("close", (code) => { + if (code === 0 && fs.pathExistsSync(outputFile)) { resolve(fs.readJson(outputFile)); } else { reject(new Error("Transcription failed")); @@ -234,6 +243,10 @@ class Whipser { }); } + abort() { + this.abortController?.abort(); + } + registerIpcHandlers() { ipcMain.handle("whisper-config", async () => { return this.config; @@ -305,6 +318,10 @@ class Whipser { }); } }); + + ipcMain.handle("whisper-abort", async (_event) => { + return await this.abort(); + }); } } diff --git a/enjoy/src/main/youtubedr.ts b/enjoy/src/main/youtubedr.ts index dc79badd..65e0f559 100644 --- a/enjoy/src/main/youtubedr.ts +++ b/enjoy/src/main/youtubedr.ts @@ -48,9 +48,11 @@ const validPathDomains = /^https?:\/\/(youtu\.be\/|(www\.)?youtube\.com\/(embed|v|shorts)\/)/; const ONE_MINUTE = 1000 * 60; // 1 minute +const TEN_MINUTES = 1000 * 60 * 10; // 10 minutes class Youtubedr { private binFile: string; + private abortController: AbortController | null = null; constructor() { this.binFile = path.join( @@ -96,6 +98,9 @@ class Youtubedr { webContents?: Electron.WebContents; } = {} ): Promise { + this.abortController?.abort(); + this.abortController = new AbortController(); + const { quality, filename = this.getYtVideoId(url) + ".mp4", @@ -116,13 +121,20 @@ class Youtubedr { let currentSpeed = ""; return new Promise((resolve, reject) => { - const proc = spawn(this.binFile, [ - "download", - url, - `--quality=${quality || "medium"}`, - `--filename=${filename}`, - `--directory=${directory}`, - ]); + const proc = spawn( + this.binFile, + [ + "download", + url, + `--quality=${quality || "medium"}`, + `--filename=${filename}`, + `--directory=${directory}`, + ], + { + timeout: TEN_MINUTES, + signal: this.abortController.signal, + } + ); proc.stdout.on("data", (data) => { const output = data.toString(); @@ -239,6 +251,10 @@ class Youtubedr { return false; } }; + + abortDownload() { + this.abortController?.abort(); + } } export default new Youtubedr(); diff --git a/enjoy/src/preload.ts b/enjoy/src/preload.ts index 31d96a40..116b676a 100644 --- a/enjoy/src/preload.ts +++ b/enjoy/src/preload.ts @@ -463,6 +463,9 @@ contextBridge.exposeInMainWorld("__ENJOY_APP__", { onProgress: ( callback: (event: IpcRendererEvent, progress: number) => void ) => ipcRenderer.on("whisper-on-progress", callback), + abort: () => { + return ipcRenderer.invoke("whisper-abort"); + }, removeProgressListeners: () => { ipcRenderer.removeAllListeners("whisper-on-progress"); }, diff --git a/enjoy/src/renderer/context/media-player-provider.tsx b/enjoy/src/renderer/context/media-player-provider.tsx index e1cc313a..c9ffd070 100644 --- a/enjoy/src/renderer/context/media-player-provider.tsx +++ b/enjoy/src/renderer/context/media-player-provider.tsx @@ -178,6 +178,7 @@ export const MediaPlayerProvider = ({ generateTranscription, transcribing, transcribingProgress, + abortGenerateTranscription, } = useTranscriptions(media); const { @@ -559,6 +560,18 @@ export const MediaPlayerProvider = ({ }; }, [media?.src, ref, mediaProvider, layout?.playerHeight]); + /* cache last segment index */ + useEffect(() => { + if (!media) return; + if (!currentSegmentIndex) return; + + setCachedSegmentIndex(currentSegmentIndex); + }, [currentSegmentIndex]); + + /* + * Update layout when window is resized + * Abort transcription when component is unmounted + */ useEffect(() => { calculateHeight(); @@ -572,17 +585,10 @@ export const MediaPlayerProvider = ({ return () => { EnjoyApp.window.removeListeners(); + abortGenerateTranscription(); }; }, []); - /* cache last segment index */ - useEffect(() => { - if (!media) return; - if (!currentSegmentIndex) return; - - setCachedSegmentIndex(currentSegmentIndex); - }, [currentSegmentIndex]); - return ( <> { }; }, [transcription, media]); + const abortGenerateTranscription = () => { + EnjoyApp.whisper.abort(); + setTranscribing(false); + }; + return { transcription, transcribingProgress, transcribing, generateTranscription, + abortGenerateTranscription, }; }; diff --git a/enjoy/src/types/enjoy-app.d.ts b/enjoy/src/types/enjoy-app.d.ts index c52e1340..ac7615cb 100644 --- a/enjoy/src/types/enjoy-app.d.ts +++ b/enjoy/src/types/enjoy-app.d.ts @@ -269,6 +269,7 @@ type EnjoyAppType = { } ) => Promise>; onProgress: (callback: (event, progress: number) => void) => void; + abort: () => Promise; removeProgressListeners: () => Promise; }; ffmpeg: {