feat: show progress when downloading whisper and models

buxuku · Jun 18, 2024 · 61aa51f · 61aa51f
1 parent 2e587c1
commit 61aa51f
Show file tree

Hide file tree

Showing 10 changed files with 228 additions and 146 deletions.
diff --git a/main/background.ts b/main/background.ts
@@ -5,7 +5,6 @@ import serve from "electron-serve";
 import { createWindow } from "./helpers";
 import {
   install,
-  downModel,
   makeWhisper,
   checkWhisperInstalled,
   getModelsInstalled,
@@ -183,9 +182,6 @@ ipcMain.on("makeWhisper", (event) => {
   makeWhisper(event);
 });
 
-ipcMain.on("downModel", (event, { model, source }) => {
-  downModel(event, model, source);
-});
 
 ipcMain.on("openUrl", (event, url) => {
   shell.openExternal(url);

diff --git a/main/helpers/whisper.ts b/main/helpers/whisper.ts
@@ -61,6 +61,11 @@ export const install = (event, source) => {
       url: repoUrl,
       singleBranch: true,
       depth: 1,
+      onProgress: (res) => {
+        if (res.total) {
+          event.sender.send("installWhisperProgress", res.phase, res.loaded / res.total);
+        }
+      },
     })
     .then((res) => {
       if (checkWhisperInstalled()) {
@@ -80,50 +85,6 @@ export const install = (event, source) => {
     });
 };
 
-export const downModel = async (
-  event,
-  whisperModel,
-  source = "hf-mirror.com",
-) => {
-  const { modelsPath } = getPath();
-  const modelName = whisperModel?.toLowerCase();
-  const modelPath = path.join(modelsPath, `ggml-${modelName}.bin`);
-  if (fs.existsSync(modelPath)) return;
-  if (!checkWhisperInstalled()) {
-    event.sender.send("message", "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("完成模型下载地址替换", modelName);
-    console.log("正在安装 whisper.cpp 模型");
-    exec(`${shell} "${downShellPath}" ${modelName}`, (err, stdout) => {
-      if (err) {
-        event.sender.send("message", err);
-      } else {
-        event.sender.send("message", `模型 ${modelName} 下载完成`);
-      }
-      event.sender.send("downModelComplete", !err);
-      event.sender.send("getSystemInfoComplete", {
-        whisperInstalled: checkWhisperInstalled(),
-        modelsInstalled: getModelsInstalled(),
-      });
-    });
-  } catch (error) {
-    event.sender.send("message", error);
-  }
-};
-
 export const makeWhisper = (event) => {
   const { whisperPath, mainPath } = getPath();
   if (fs.existsSync(mainPath) || isWin32()) {
@@ -133,6 +94,7 @@ export const makeWhisper = (event) => {
   if (!checkWhisperInstalled()) {
     event.sender.send("message", "whisper.cpp 未下载，请先下载 whisper.cpp");
   }
+  event.sender.send("beginMakeWhisper", true);
   exec(`make -C "${whisperPath}"`, (err, stdout) => {
     if (err) {
       event.sender.send("message", err);
@@ -168,22 +130,23 @@ export const downloadModelSync = async (model, source, onProcess) => {
   try {
     let downShellPath;
     let shell: string;
+    let args = [];
     if (isDarwin()) {
       downShellPath = path.join(modelsPath, "download-ggml-model.sh");
       shell = "bash";
+      args = [`${downShellPath}`, `${model}`];
     } else if (isWin32()) {
       downShellPath = path.join(modelsPath, "download-ggml-model.cmd");
-      shell = "cmd.exe /c";
+      shell = "cmd";
+      args = [`/c`, `${downShellPath}`, `${model}`];
     } else {
       throw Error("platform does not support! ");
     }
     await replaceModelSource(`${downShellPath}`, source);
     console.log("完成模型下载地址替换", model);
     console.log("正在安装 whisper.cpp 模型");
     try {
-      await runCommand(`${shell}`, [`${downShellPath}`, `${model}`], (data) =>
-        onProcess(data),
-      );
+      await runCommand(`${shell}`, args, (data) => onProcess(data));
     } catch (error) {
       await deleteModel(model);
       throw error;

diff --git a/renderer/components/DownModel.tsx b/renderer/components/DownModel.tsx
@@ -1,40 +1,48 @@
-import React, { useEffect } from "react";
-import { Button } from "@/components/ui/button";
-import { Loader2 } from "lucide-react";
+import React, { useEffect, FC, PropsWithChildren } from "react";
 
-const DownModel = ({ modelName, callBack, downSource }) => {
+interface IProps extends PropsWithChildren {
+  modelName: string;
+  callBack?: () => void;
+  downSource?: string;
+}
+
+const DownModel: FC<IProps> = (props) => {
+  const { modelName, callBack, downSource = "hf-mirror", children } = props;
   const [loading, setLoading] = React.useState(false);
   const [progress, setProgress] = React.useState(0);
   useEffect(() => {
     window?.ipc?.on("downloadProgress", (model, progress: number) => {
       if (model === modelName) {
         setProgress(progress);
         setLoading(progress < 100);
-        if(progress >= 100) {
-            console.log(progress, 'progress')
-            callBack && callBack();
+        if (progress >= 100) {
+          callBack && callBack();
         }
       }
     });
   }, []);
-  const handleDownModel = async () => {
+  const handleDownModel = async (source = downSource) => {
     setLoading(true);
     await window?.ipc?.invoke("downloadModel", {
       model: modelName,
-      source: downSource,
+      source,
     });
     setLoading(false);
   };
   return (
-    <Button onClick={handleDownModel} disabled={loading}>
-      {loading ? (
-        <>
-          <Loader2 className="mr-2 h-4 w-4 animate-spin" /> {progress} %
-        </>
-      ) : (
-        "下载"
-      )}
-    </Button>
+    <span className="inline-block">
+      {React.isValidElement<{
+        loading?: boolean;
+        progress?: number;
+        handleDownModel?: (source?: string) => void;
+      }>(children)
+        ? React.cloneElement(children, {
+            loading,
+            progress,
+            handleDownModel,
+          })
+        : children}
+    </span>
   );
 };
 

diff --git a/renderer/components/DownModelButton.tsx b/renderer/components/DownModelButton.tsx
@@ -0,0 +1,29 @@
+import React, { FC } from "react";
+import { Button } from "@/components/ui/button";
+import { Loader2 } from "lucide-react";
+
+interface IProps {
+  loading?: boolean;
+  progress?: number;
+  handleDownModel?: () => void;
+}
+
+const DownModelButton: FC<IProps> = ({
+  loading,
+  progress,
+  handleDownModel,
+}) => {
+  return (
+    <Button disabled={loading} onClick={() => handleDownModel()}>
+      {loading ? (
+        <>
+          <Loader2 className="mr-2 h-4 w-4 animate-spin" /> {progress} %
+        </>
+      ) : (
+        "下载"
+      )}
+    </Button>
+  );
+};
+
+export default DownModelButton;
diff --git a/renderer/components/DownModelDropdown.tsx b/renderer/components/DownModelDropdown.tsx
@@ -0,0 +1,69 @@
+import React, { FC } from "react";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuLabel,
+  DropdownMenuSeparator,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { Button } from "@/components/ui/button";
+import { Loader2 } from "lucide-react";
+
+interface IProps {
+  loading?: boolean;
+  progress?: number;
+  handleDownModel?: (source: string) => void;
+  setShowGuide?: (type: boolean) => void;
+  installComplete?: boolean;
+  whisperLoading?: boolean;
+}
+
+const DownModelDropdown: FC<IProps> = ({
+  loading,
+  progress,
+  handleDownModel,
+  setShowGuide,
+  installComplete,
+  whisperLoading,
+}) => {
+  return (
+    <DropdownMenu>
+      <DropdownMenuTrigger asChild>
+        <Button
+          variant="outline"
+          disabled={loading || !installComplete || whisperLoading}
+          className="w-24"
+        >
+          {loading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
+          {loading ? `${progress}%` : "下载"}
+        </Button>
+      </DropdownMenuTrigger>
+      <DropdownMenuContent className="w-[225px]">
+        <DropdownMenuLabel>请选择下载源</DropdownMenuLabel>
+        <DropdownMenuSeparator />
+        <DropdownMenuItem
+          className="cursor-pointer hover:bg-gray-100"
+          onClick={() => handleDownModel("hf-mirror")}
+        >
+          国内镜像源(较快)
+        </DropdownMenuItem>
+        <DropdownMenuItem
+          className="cursor-pointer hover:bg-gray-100"
+          onClick={() => handleDownModel("huggingface")}
+        >
+          huggingface官方源(较慢)
+        </DropdownMenuItem>
+        <DropdownMenuSeparator />
+        <DropdownMenuItem
+          className="cursor-pointer hover:bg-gray-100"
+          onClick={() => setShowGuide(false)}
+        >
+          稍后下载
+        </DropdownMenuItem>
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+};
+
+export default DownModelDropdown;
diff --git a/renderer/components/DownModelLink.tsx b/renderer/components/DownModelLink.tsx
@@ -0,0 +1,26 @@
+import React, { FC } from "react";
+
+interface IProps {
+  loading?: boolean;
+  progress?: number;
+  handleDownModel?: () => void;
+}
+const DownModelLink: FC<IProps> = ({ loading, progress, handleDownModel }) => {
+  return (
+    <span className="inline-block">
+      该模型未下载，
+      {loading ? (
+        `正在下载中 ${progress}%...`
+      ) : (
+        <a
+          className="cursor-pointer text-blue-500"
+          onClick={() => handleDownModel()}
+        >
+          立即下载
+        </a>
+      )}
+    </span>
+  );
+};
+
+export default DownModelLink;