feat: add model control page (#11)

Co-authored-by: linxiaodong <[email protected]>
buxuku · Jun 14, 2024 · ede26d7 · ede26d7
1 parent 1404c80
commit ede26d7
Show file tree

Hide file tree

Showing 15 changed files with 1,001 additions and 392 deletions.
diff --git a/electron-builder.yml b/electron-builder.yml
@@ -17,16 +17,15 @@ mac:
         - arm64
         - x64
   artifactName: "${productName}_Mac_${version}_${arch}.${ext}"
-win: {
-  target: nsis,
-  requestedExecutionLevel: asInvoker,
-  artifactName: "${productName}_Windows_${version}_${arch}.${ext}",
-}
+win:
+  target: nsis
+  requestedExecutionLevel: asInvoker
+  artifactName: "${productName}_Windows_${version}_${arch}.${ext}"
 nsis:
-  oneClick: false,
-  perMachine: false,
-  allowToChangeInstallationDirectory: true,
-  deleteAppDataOnUninstall: false,
+  oneClick: false
+  perMachine: false
+  allowToChangeInstallationDirectory: true
+  deleteAppDataOnUninstall: false
 extraResources:
   - from: ./extraResources/
     to: ./extraResources/

diff --git a/main/background.ts b/main/background.ts
@@ -9,10 +9,17 @@ import {
   makeWhisper,
   checkWhisperInstalled,
   getModelsInstalled,
+  deleteModel,
+  downloadModelSync,
+  getPath,
 } from "./helpers/whisper";
 import { extractAudio } from "./helpers/ffmpeg";
 import translate from "./helpers/translate";
-import { getExtraResourcesPath, isWin32, renderTemplate } from "./helpers/utils";
+import {
+  getExtraResourcesPath,
+  isWin32,
+  renderTemplate,
+} from "./helpers/utils";
 import fs from "fs";
 
 const isProd = process.env.NODE_ENV === "production";
@@ -91,36 +98,62 @@ ipcMain.on("handleTask", async (event, { files, formData }) => {
       };
       const srtFile = path.join(
         directory,
-        `${renderTemplate(sourceSrtSaveFileName, templateData)}`,
+        `${renderTemplate(
+          saveSourceSrt ? sourceSrtSaveFileName : "${fileName}-temp",
+          templateData,
+        )}`,
       );
       const whisperModel = model?.toLowerCase();
+      event.sender.send("taskStatusChange", file, "extractAudio", "loading");
       await extractAudio(filePath, audioFile);
-      event.sender.send("extractAudio-completed", file);
+      event.sender.send("taskStatusChange", file, "extractAudio", "done");
       let mainPath = `${whisperPath}main`;
-      if(isWin32()){
-        mainPath =  path.join(getExtraResourcesPath(), 'whisper-bin-x64', 'main.exe');
+      if (isWin32()) {
+        mainPath = path.join(
+          getExtraResourcesPath(),
+          "whisper-bin-x64",
+          "main.exe",
+        );
       }
+      event.sender.send("taskStatusChange", file, "extractSubtitle", "loading");
       exec(
-        `"${mainPath}" -m "${whisperPath}models/ggml-${whisperModel}.bin" -f "${audioFile}" -osrt -of "${srtFile}"`,
+        `"${mainPath}" -m "${whisperPath}models/ggml-${whisperModel}.bin" -f "${audioFile}" -osrt -of "${srtFile}" -l ${sourceLanguage}`,
         async (error, stdout, stderr) => {
           if (error) {
             event.sender.send("message", error);
           }
-          event.sender.send("extractSubtitle-completed", file);
+          event.sender.send(
+            "taskStatusChange",
+            file,
+            "extractSubtitle",
+            "done",
+          );
           fs.unlink(audioFile, (err) => {
             if (err) {
               console.log(err);
             }
           });
           if (translateProvider !== "-1") {
+            event.sender.send(
+              "taskStatusChange",
+              file,
+              "translateSubtitle",
+              "loading",
+            );
+
             await translate(
               event,
               directory,
               fileName,
               `${srtFile}.srt`,
               formData,
             );
-            event.sender.send("translate-completed", file);
+            event.sender.send(
+              "taskStatusChange",
+              file,
+              "translateSubtitle",
+              "done",
+            );
           }
           if (!saveSourceSrt) {
             fs.unlink(`${srtFile}.srt`, (err) => {
@@ -158,3 +191,22 @@ ipcMain.on("downModel", (event, { model, source }) => {
 ipcMain.on("openUrl", (event, url) => {
   shell.openExternal(url);
 });
+
+ipcMain.handle("deleteModel", async (event, modelName) => {
+  await deleteModel(modelName);
+  return true;
+});
+
+ipcMain.handle("getSystemInfo", async (event, key) => {
+  const res = {
+    whisperInstalled: checkWhisperInstalled(),
+    modelsInstalled: getModelsInstalled(),
+    modelsPath: getPath("modelsPath"),
+  };
+  return res;
+});
+
+ipcMain.handle("downloadModel", async (event, { model, source }) => {
+  await downloadModelSync(model?.toLowerCase(), source);
+  return true;
+});
diff --git a/main/helpers/whisper.ts b/main/helpers/whisper.ts
@@ -84,12 +84,12 @@ export const downModel = async (
   try {
     let downShellPath;
     let shell: string;
-    if(isDarwin()){
-      downShellPath = path.join(modelsPath, 'download-ggml-model.sh');
-      shell = 'bash';
-    } else if(isWin32()){
-      downShellPath = path.join(modelsPath, 'download-ggml-model.cmd');
-      shell = 'cmd.exe /c'
+    if (isDarwin()) {
+      downShellPath = path.join(modelsPath, "download-ggml-model.sh");
+      shell = "bash";
+    } else if (isWin32()) {
+      downShellPath = path.join(modelsPath, "download-ggml-model.cmd");
+      shell = "cmd.exe /c";
     } else {
       throw Error("platform does not support! ");
     }
@@ -118,7 +118,7 @@ export const makeWhisper = (event) => {
   if (fs.existsSync(mainPath) || isWin32()) {
     event.sender.send("makeWhisperComplete", true);
     return;
-  };
+  }
   if (!checkWhisperInstalled()) {
     event.sender.send("message", "whisper.cpp 未下载，请先下载 whisper.cpp");
   }
@@ -135,3 +135,50 @@ export const makeWhisper = (event) => {
     event.sender.send("makeWhisperComplete", !err);
   });
 };
+
+export const deleteModel = async (model) => {
+  const modelsPath = getPath("modelsPath");
+  const modelPath = path.join(modelsPath, `ggml-${model}.bin`);
+  return new Promise((resolve, reject) => {
+    if (fs.existsSync(modelPath)) {
+      fs.unlinkSync(modelPath);
+    }
+    resolve("ok");
+  });
+};
+
+export const downloadModelSync = async (model, source) => {
+  const modelsPath = getPath("modelsPath");
+  const modelPath = path.join(modelsPath, `ggml-${model}.bin`);
+  if (fs.existsSync(modelPath)) return;
+  if (!checkWhisperInstalled()) {
+    throw Error("whisper.cpp 未下载，请先下载 whisper.cpp");
+  }
+  try {
+    let downShellPath;
+    let shell: string;
+    if (isDarwin()) {
+      downShellPath = path.join(modelsPath, "download-ggml-model.sh");
+      shell = "bash";
+    } else if (isWin32()) {
+      downShellPath = path.join(modelsPath, "download-ggml-model.cmd");
+      shell = "cmd.exe /c";
+    } else {
+      throw Error("platform does not support! ");
+    }
+    await replaceModelSource(`${downShellPath}`, 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')
+            }
+        });
+    })
+  } catch (error) {
+      console.log(error)
+  }
+};
diff --git a/main/preload.ts b/main/preload.ts
@@ -1,13 +1,12 @@
-import {
-  contextBridge,
-  ipcRenderer,
-  IpcRendererEvent,
-} from "electron";
+import { contextBridge, ipcRenderer, IpcRendererEvent } from "electron";
 
 const handler = {
   send(channel: string, value: unknown) {
     ipcRenderer.send(channel, value);
   },
+  invoke(channel: string, ...args): Promise<any> {
+      return ipcRenderer.invoke(channel, ...args)
+  },
   on(channel: string, callback: (...args: unknown[]) => void) {
     const subscription = (_event: IpcRendererEvent, ...args: unknown[]) =>
       callback(...args);

diff --git a/package.json b/package.json
@@ -2,7 +2,7 @@
   "private": true,
   "name": "video-subtitle-master",
   "description": "视频转字幕，字幕翻译软件",
-  "version": "1.0.10",
+  "version": "1.0.11",
   "author": "buxuku <[email protected]>",
   "main": "app/background.js",
   "scripts": {
@@ -14,6 +14,7 @@
   "dependencies": {
     "@ffmpeg-installer/ffmpeg": "^1.1.0",
     "@hookform/resolvers": "^3.4.0",
+    "@radix-ui/react-alert-dialog": "^1.0.5",
     "@radix-ui/react-dialog": "^1.0.5",
     "@radix-ui/react-dropdown-menu": "^2.0.6",
     "@radix-ui/react-hover-card": "^1.0.7",

diff --git a/renderer/components/DeleteModel.tsx b/renderer/components/DeleteModel.tsx
@@ -0,0 +1,44 @@
+import React from 'react';
+import {
+    AlertDialog,
+    AlertDialogAction,
+    AlertDialogCancel,
+    AlertDialogContent,
+    AlertDialogDescription,
+    AlertDialogFooter,
+    AlertDialogHeader,
+    AlertDialogTitle,
+    AlertDialogTrigger,
+} from "@/components/ui/alert-dialog"
+
+const DeleteModel = ({children, modelName, callBack}) => {
+    const [visibility, setVisibility] = React.useState(false);
+    const handleDelete = async (e) => {
+        e.preventDefault();
+        const res = await window?.ipc?.invoke('deleteModel', modelName);
+        setVisibility(false);
+        callBack && callBack();
+    }
+    return (
+        <AlertDialog open={visibility}>
+            <AlertDialogTrigger asChild onClick={() => setVisibility(true)}>
+                {children}
+            </AlertDialogTrigger>
+            <AlertDialogContent>
+                <AlertDialogHeader>
+                    <AlertDialogTitle>确认删除该模型？</AlertDialogTitle>
+                    <AlertDialogDescription>
+                        删除之后，如果你需要再次使用该模型，需要重新下载。
+                    </AlertDialogDescription>
+                </AlertDialogHeader>
+                <AlertDialogFooter>
+                    <AlertDialogCancel onClick={() => setVisibility(false)}>取消</AlertDialogCancel>
+                    <AlertDialogAction onClick={handleDelete}>删除</AlertDialogAction>
+                </AlertDialogFooter>
+            </AlertDialogContent>
+        </AlertDialog>
+
+    );
+};
+
+export default DeleteModel;
diff --git a/renderer/components/DownModel.tsx b/renderer/components/DownModel.tsx
@@ -0,0 +1,23 @@
+import React from "react";
+import { Button } from "@/components/ui/button";
+import { Loader2 } from "lucide-react";
+
+const DownModel = ({ modelName, callBack, downSource }) => {
+  const [loading, setLoading] = React.useState(false);
+  const handleDownModel = async () => {
+    setLoading(true);
+    await window?.ipc?.invoke("downloadModel", {
+        model: modelName,
+        source: downSource,
+    });
+    setLoading(false);
+    callBack && callBack();
+  };
+  return (
+    <Button onClick={handleDownModel} disabled={loading}>
+      {loading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}下载
+    </Button>
+  );
+};
+
+export default DownModel;