From c1516adaac408561fb50c8d2bd5c3ba43d3b559f Mon Sep 17 00:00:00 2001
From: Steven Rafferty <mail@stevenrafferty.net>
Date: Sat, 10 Feb 2024 14:40:00 +0700
Subject: [PATCH] Added UI logic for displaying if a model is downloaded in the
 setting model dropdown

---
 src/components/SettingsModal.tsx | 32 ++++++++++++++++++++++++++++----
 src/utils/model-data.ts          | 31 +++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx
index 362a26a..adfa070 100644
--- a/src/components/SettingsModal.tsx
+++ b/src/components/SettingsModal.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useState } from "react";
+import { Dispatch, SetStateAction, useCallback, useEffect, useState } from "react";
 import { useAtom } from "jotai";
 import {
   Modal,
@@ -14,12 +14,13 @@ import {
   Tooltip,
 } from "@mantine/core";
 import {
+  ModelOption,
   whisperModelSizes,
   whisperModelsBase,
   whisperModelsQuantized,
 } from "../whisper-utils";
 import { currentModel, currentSettings, fetchTimestamp } from "../state/main";
-import { loadOrGetModel } from "../utils/model-data";
+import { getDownloadedModels, loadOrGetModel } from "../utils/model-data";
 import { WhisperModelName } from "../types";
 import { IconQuestionMark } from "@tabler/icons-react";
 
@@ -45,6 +46,8 @@ export function SettingsModal({
   const whisperModelUpdated = useCallback((_newModel: Uint8Array) => {
     setNewModel(_newModel);
   }, []);
+  const [whisperModelsBaseState, setWhisperModelsBaseState] = useState(whisperModelsBase);
+  const [whisperModelsQuantizedState, setWhisperModelsQuantizedState] = useState(whisperModelsQuantized);
 
   useEffect(() => {
     setNewModel(model);
@@ -66,6 +69,27 @@ export function SettingsModal({
     loadingProgress,
   } = useModelData(selectedModel, whisperModelUpdated);
 
+  useEffect(() => {
+    const updateModels = async (models: ModelOption[], setStateCallback: Dispatch<SetStateAction<ModelOption[]>>) => {
+      const downloadedModels = await getDownloadedModels();
+      const updatedModels = models.map((item) => {
+        const isDownloaded = downloadedModels.includes(item.value);
+        return {
+          ...item,
+          label: isDownloaded ? `${item.label} Downloaded` : item.label
+        };
+      });
+      setStateCallback(updatedModels);
+    };
+
+    (async () => {
+      await Promise.all([
+        updateModels(whisperModelsBase, setWhisperModelsBaseState),
+        updateModels(whisperModelsQuantized, setWhisperModelsQuantizedState),
+      ])
+    })();
+  }, [whisperModelLoaded]);
+
   return (
     <Modal
       opened={opened}
@@ -130,11 +154,11 @@ export function SettingsModal({
             { label: "None", value: "" },
             {
               group: "Raw",
-              items: whisperModelsBase,
+              items: whisperModelsBaseState,
             },
             {
               group: "Quantized",
-              items: whisperModelsQuantized,
+              items: whisperModelsQuantizedState,
             },
           ]}
           onChange={(e) => {
diff --git a/src/utils/model-data.ts b/src/utils/model-data.ts
index cdcc90b..ce01fdf 100644
--- a/src/utils/model-data.ts
+++ b/src/utils/model-data.ts
@@ -10,6 +10,37 @@ const dbName = "whisperModels";
 const modelBaseUrl =
   "https://link.storjshare.io/s/jueavj4qtolpgszkbp5awref22da/models";
 
+export function getDownloadedModels(): Promise<string[]> {
+  return new Promise((resolve, reject) => {
+    if (!navigator.storage || !navigator.storage.estimate) {
+      console.log("loadRemote: navigator.storage.estimate() is not supported");
+    }
+
+    const openRequest = indexedDB.open(dbName, dbVersion);
+
+    openRequest.onsuccess = function () {
+      const db = openRequest.result;
+      const tx = db.transaction(["models"], "readonly");
+      const objectStore = tx.objectStore("models");
+      const localFilesRequest = objectStore.getAllKeys();
+
+      localFilesRequest.onsuccess = function () {
+        resolve((localFilesRequest.result as string[]) || []);
+      };
+
+      localFilesRequest.onerror = function () {
+        console.error("Failed to fetch models");
+        reject(new Error("Failed to fetch models"));
+      };
+    };
+
+    openRequest.onerror = function () {
+      console.error("Failed to open request to indexedDB");
+      reject(new Error("Failed to open request to indexedDB"));
+    };
+  });
+}
+
 // TODO: this method seems to leak memory when changing models
 export function loadOrGetModel(
   selectedModel: keyof typeof whisperModelSizes | "" | undefined,