feat: support openai whisper (#12)

buxuku · Jun 15, 2024 · 890d0fc · 890d0fc
1 parent a42f25d
commit 890d0fc
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 11 deletions.
diff --git a/main/background.ts b/main/background.ts
@@ -12,6 +12,7 @@ import {
   deleteModel,
   downloadModelSync,
   getPath,
+  checkOpanAiWhisper
 } from "./helpers/whisper";
 import { extractAudio } from "./helpers/ffmpeg";
 import translate from "./helpers/translate";
@@ -115,12 +116,17 @@ ipcMain.on("handleTask", async (event, { files, formData }) => {
           "main.exe",
         );
       }
+      let runShell = `"${mainPath}" -m "${whisperPath}models/ggml-${whisperModel}.bin" -f "${audioFile}" -osrt -of "${srtFile}" -l ${sourceLanguage}`
+      const hasOpenAiWhiaper = await checkOpanAiWhisper();
+      if (hasOpenAiWhiaper) {
+        runShell = `whisper "${audioFile}" --model ${whisperModel} --device cuda --output_format srt --output_dir ${directory} --language ${sourceLanguage}`
+      }
       event.sender.send("taskStatusChange", file, "extractSubtitle", "loading");
-      exec(
-        `"${mainPath}" -m "${whisperPath}models/ggml-${whisperModel}.bin" -f "${audioFile}" -osrt -of "${srtFile}" -l ${sourceLanguage}`,
+      exec(runShell,
         async (error, stdout, stderr) => {
           if (error) {
             event.sender.send("message", error);
+            return;
           }
           event.sender.send(
             "taskStatusChange",

diff --git a/main/helpers/whisper.ts b/main/helpers/whisper.ts
@@ -1,11 +1,12 @@
-import { exec } from "child_process";
+import { exec, spawn } from "child_process";
 import { app } from "electron";
 import path from "path";
 import fs from "fs";
 import git from "isomorphic-git";
 import http from "isomorphic-git/http/node";
 import replaceModelSource from "./model-source";
 import { isDarwin, isWin32 } from "./utils";
+import { log } from "console";
 
 export const getPath = (key?: string) => {
   const userDataPath = app.getPath("userData");
@@ -170,15 +171,40 @@ export const downloadModelSync = async (model, source) => {
     console.log("完成模型下载地址替换", model);
     console.log("正在安装 whisper.cpp 模型");
     return new Promise((resolve, reject) => {
-        exec(`${shell} "${downShellPath}" ${model}`, (err, stdout) => {
-            if (err) {
-               reject(err);
-            } else {
-                resolve('ok')
-            }
-        });
+      exec(`${shell} "${downShellPath}" ${model}`, (err, stdout) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve('ok')
+        }
+      });
     })
   } catch (error) {
-      console.log(error)
+    console.log(error)
   }
 };
+
+
+export async function checkOpanAiWhisper() {
+  return new Promise((resolve, reject) => {
+    let env = process.env;
+    env.PYTHONIOENCODING = 'UTF-8';
+    const childProcess = spawn('whisper', ['-h'], { env: env });
+    childProcess.on('error', (error: { code: string }) => {
+      if (error.code === 'ENOENT') {
+        resolve(false);
+      } else {
+        reject(error);
+      }
+    });
+    childProcess.on('exit', (code) => {
+      console.log('code: ', code);
+      if (code === 0) {
+        console.log('openai whisper ready')
+        resolve(true);
+      } else {
+        resolve(false);
+      }
+    });
+  });
+}