Added directory read and models list

electron/commands.ts

@@ -18,6 +18,7 @@ const commands = {
   UPSCAYL_VIDEO_PROGRESS: "Send Video Upscale Progress from Main to Renderer",
   FFMPEG_VIDEO_DONE: "Ran FFMpeg successfully",
   FFMPEG_VIDEO_PROGRESS: "Running FFMpeg for frame extraction",
+  CUSTOM_MODEL_FILES_LIST: "Get custom model files list",
 };
 
 export default commands;

electron/index.ts

@@ -35,7 +35,7 @@ import commands from "./commands";
 log.initialize({ preload: true });
 
 // Prepare the renderer once the app is ready
-let mainWindow;
+let mainWindow: BrowserWindow;
 app.on("ready", async () => {
   await prepareNext("./renderer");
 
@@ -89,6 +89,7 @@ app.on("window-all-closed", app.quit);
 
 log.log(app.getAppPath());
 
+// Path variables for file and folder selection
 let imagePath: string | undefined = undefined;
 let folderPath: string | undefined = undefined;
 let customModelsFolderPath: string | undefined = undefined;
@@ -104,10 +105,40 @@ ipcMain.handle(commands.SELECT_FILE, async () => {
 
   if (canceled) {
     log.log("File Operation Cancelled");
-    return "cancelled";
+    return null;
   } else {
     log.log("Selected File Path: ", filePaths[0]);
+    let isValid = false;
     imagePath = filePaths[0];
+
+    // READ SELECTED FILES
+    filePaths.forEach((file) => {
+      // log.log("Files in Folder: ", file);
+      if (
+        file.endsWith(".png") ||
+        file.endsWith(".jpg") ||
+        file.endsWith(".jpeg") ||
+        file.endsWith(".webp") ||
+        file.endsWith(".JPG") ||
+        file.endsWith(".PNG") ||
+        file.endsWith(".JPEG") ||
+        file.endsWith(".WEBP")
+      ) {
+        isValid = true;
+      }
+    });
+
+    if (!isValid) {
+      const options: MessageBoxOptions = {
+        type: "error",
+        title: "Invalid File",
+        message:
+          "The selected file is not a valid image. Make sure you select a '.png', '.jpg', or '.webp' file.",
+      };
+      dialog.showMessageBoxSync(mainWindow, options);
+      return null;
+    }
+
     // CREATE input AND upscaled FOLDER
     return filePaths[0];
   }
@@ -120,8 +151,7 @@ ipcMain.handle(commands.SELECT_FOLDER, async (event, message) => {
     defaultPath: folderPath,
   });
   if (canceled) {
-    log.log("operation cancelled");
-    return "cancelled";
+    return null;
   } else {
     log.log("Selected Folder Path: ", folderPaths[0]);
     folderPath = folderPaths[0];
@@ -129,6 +159,43 @@ ipcMain.handle(commands.SELECT_FOLDER, async (event, message) => {
   }
 });
 
+//------------------------Get Model Names-----------------------------//
+const getModels = (folderPath: string) => {
+  let models: string[] = [];
+  let isValid = false;
+
+  // READ CUSTOM MODELS FOLDER
+  fs.readdirSync(folderPath).forEach((file) => {
+    // log.log("Files in Folder: ", file);
+    if (
+      file.endsWith(".param") ||
+      file.endsWith(".PARAM") ||
+      file.endsWith(".bin") ||
+      file.endsWith(".BIN")
+    ) {
+      isValid = true;
+      const modelName = file.substring(0, file.lastIndexOf(".")) || file;
+      if (!models.includes(modelName)) {
+        models.push(modelName);
+      }
+    }
+  });
+
+  if (!isValid) {
+    const options: MessageBoxOptions = {
+      type: "error",
+      title: "Invalid Folder",
+      message:
+        "The selected folder does not contain valid model files. Make sure you select the folder that ONLY contains '.param' and '.bin' files.",
+      buttons: ["OK"],
+    };
+    dialog.showMessageBoxSync(options);
+    return null;
+  }
+
+  return models;
+};
+
 //------------------------Select Custom Models Folder---------------------//
 ipcMain.handle(commands.SELECT_CUSTOM_MODEL_FOLDER, async (event, message) => {
   const { canceled, filePaths: folderPaths } = await dialog.showOpenDialog({
@@ -137,12 +204,17 @@ ipcMain.handle(commands.SELECT_CUSTOM_MODEL_FOLDER, async (event, message) => {
     defaultPath: customModelsFolderPath,
   });
   if (canceled) {
-    log.log("operation cancelled");
-    return "cancelled";
+    return null;
   } else {
     log.log("Custom Folder Path: ", folderPaths[0]);
     customModelsFolderPath = folderPaths[0];
-    return folderPaths[0];
+
+    mainWindow.webContents.send(
+      commands.CUSTOM_MODEL_FILES_LIST,
+      getModels(customModelsFolderPath)
+    );
+
+    return customModelsFolderPath;
   }
 });
 

main/commands.js

@@ -19,5 +19,6 @@ const commands = {
     UPSCAYL_VIDEO_PROGRESS: "Send Video Upscale Progress from Main to Renderer",
     FFMPEG_VIDEO_DONE: "Ran FFMpeg successfully",
     FFMPEG_VIDEO_PROGRESS: "Running FFMpeg for frame extraction",
+    CUSTOM_MODEL_FILES_LIST: "Get custom model files list",
 };
 exports.default = commands;

main/index.js

@@ -75,48 +75,113 @@ electron_1.app.on("ready", () => __awaiter(void 0, void 0, void 0, function* ()
 // Quit the app once all windows are closed
 electron_1.app.on("window-all-closed", electron_1.app.quit);
 electron_log_1.default.log(electron_1.app.getAppPath());
+// Path variables for file and folder selection
+let imagePath = undefined;
+let folderPath = undefined;
+let customModelsFolderPath = undefined;
 //------------------------Select File-----------------------------//
 // ! DONT FORGET TO RESTART THE APP WHEN YOU CHANGE CODE HERE
 electron_1.ipcMain.handle(commands_1.default.SELECT_FILE, () => __awaiter(void 0, void 0, void 0, function* () {
     const { canceled, filePaths } = yield electron_1.dialog.showOpenDialog({
         properties: ["openFile", "multiSelections"],
+        title: "Select Image",
+        defaultPath: imagePath,
     });
     if (canceled) {
         electron_log_1.default.log("File Operation Cancelled");
-        return "cancelled";
+        return null;
     }
     else {
         electron_log_1.default.log("Selected File Path: ", filePaths[0]);
+        let isValid = false;
+        imagePath = filePaths[0];
+        // READ SELECTED FILES
+        filePaths.forEach((file) => {
+            // log.log("Files in Folder: ", file);
+            if (file.endsWith(".png") ||
+                file.endsWith(".jpg") ||
+                file.endsWith(".jpeg") ||
+                file.endsWith(".webp") ||
+                file.endsWith(".JPG") ||
+                file.endsWith(".PNG") ||
+                file.endsWith(".JPEG") ||
+                file.endsWith(".WEBP")) {
+                isValid = true;
+            }
+        });
+        if (!isValid) {
+            const options = {
+                type: "error",
+                title: "Invalid File",
+                message: "The selected file is not a valid image. Make sure you select a '.png', '.jpg', or '.webp' file.",
+            };
+            electron_1.dialog.showMessageBoxSync(mainWindow, options);
+            return null;
+        }
         // CREATE input AND upscaled FOLDER
         return filePaths[0];
     }
 }));
 //------------------------Select Folder-----------------------------//
 electron_1.ipcMain.handle(commands_1.default.SELECT_FOLDER, (event, message) => __awaiter(void 0, void 0, void 0, function* () {
-    const { canceled, filePaths } = yield electron_1.dialog.showOpenDialog({
+    const { canceled, filePaths: folderPaths } = yield electron_1.dialog.showOpenDialog({
         properties: ["openDirectory"],
+        defaultPath: folderPath,
     });
     if (canceled) {
-        electron_log_1.default.log("operation cancelled");
-        return "cancelled";
+        return null;
     }
     else {
-        electron_log_1.default.log(filePaths[0]);
-        return filePaths[0];
+        electron_log_1.default.log("Selected Folder Path: ", folderPaths[0]);
+        folderPath = folderPaths[0];
+        return folderPaths[0];
     }
 }));
+//------------------------Get Model Names-----------------------------//
+const getModels = (folderPath) => {
+    let models = [];
+    let isValid = false;
+    // READ CUSTOM MODELS FOLDER
+    fs_1.default.readdirSync(folderPath).forEach((file) => {
+        // log.log("Files in Folder: ", file);
+        if (file.endsWith(".param") ||
+            file.endsWith(".PARAM") ||
+            file.endsWith(".bin") ||
+            file.endsWith(".BIN")) {
+            isValid = true;
+            const modelName = file.substring(0, file.lastIndexOf(".")) || file;
+            if (!models.includes(modelName)) {
+                models.push(modelName);
+            }
+        }
+    });
+    if (!isValid) {
+        const options = {
+            type: "error",
+            title: "Invalid Folder",
+            message: "The selected folder does not contain valid model files. Make sure you select the folder that ONLY contains '.param' and '.bin' files.",
+            buttons: ["OK"],
+        };
+        electron_1.dialog.showMessageBoxSync(options);
+        return null;
+    }
+    return models;
+};
 //------------------------Select Custom Models Folder---------------------//
 electron_1.ipcMain.handle(commands_1.default.SELECT_CUSTOM_MODEL_FOLDER, (event, message) => __awaiter(void 0, void 0, void 0, function* () {
-    const { canceled, filePaths } = yield electron_1.dialog.showOpenDialog({
+    const { canceled, filePaths: folderPaths } = yield electron_1.dialog.showOpenDialog({
         properties: ["openDirectory"],
+        title: "Select Custom Models Folder",
+        defaultPath: customModelsFolderPath,
     });
     if (canceled) {
-        electron_log_1.default.log("operation cancelled");
-        return "cancelled";
+        return null;
     }
     else {
-        electron_log_1.default.log(filePaths[0]);
-        return filePaths[0];
+        electron_log_1.default.log("Custom Folder Path: ", folderPaths[0]);
+        customModelsFolderPath = folderPaths[0];
+        mainWindow.webContents.send(commands_1.default.CUSTOM_MODEL_FILES_LIST, getModels(customModelsFolderPath));
+        return customModelsFolderPath;
     }
 }));
 //------------------------Open Folder-----------------------------//

renderer/atoms/modelsListAtom.ts

@@ -0,0 +1,14 @@
+import { atom } from "jotai";
+
+export type TModelsList = {
+  label: string;
+  value: string;
+}[];
+
+export const modelsListAtom = atom<TModelsList>([
+  { label: "General Photo (Real-ESRGAN)", value: "realesrgan-x4plus" },
+  { label: "General Photo (Remacri)", value: "remacri" },
+  { label: "General Photo (Ultramix Balanced)", value: "ultramix_balanced" },
+  { label: "General Photo (Ultrasharp)", value: "ultrasharp" },
+  { label: "Digital Art", value: "realesrgan-x4plus-anime" },
+]);

renderer/atoms/userSettingsAtom.ts

@@ -0,0 +1,6 @@
+import { atomWithStorage } from "jotai/utils";
+
+export const customModelsPathAtom = atomWithStorage<string | null>(
+  "customModelsPath",
+  null
+);

renderer/components/LeftPaneImageSteps.tsx

@@ -1,7 +1,9 @@
+import { useAtom } from "jotai";
 import React, { useEffect, useState } from "react";
 import Select from "react-select";
 import ReactTooltip from "react-tooltip";
 import { themeChange } from "theme-change";
+import { modelsListAtom } from "../atoms/modelsListAtom";
 
 interface IProps {
   progress: string;
@@ -63,6 +65,8 @@ function LeftPaneImageSteps({
     value: null,
   });
 
+  const [modelOptions, setModelOptions] = useAtom(modelsListAtom);
+
   useEffect(() => {
     themeChange(false);
 
@@ -130,14 +134,6 @@ function LeftPaneImageSteps({
     },
   };
 
-  const modelOptions = [
-    { label: "General Photo (Real-ESRGAN)", value: "realesrgan-x4plus" },
-    { label: "General Photo (Remacri)", value: "remacri" },
-    { label: "General Photo (Ultramix Balanced)", value: "ultramix_balanced" },
-    { label: "General Photo (Ultrasharp)", value: "ultrasharp" },
-    { label: "Digital Art", value: "realesrgan-x4plus-anime" },
-  ];
-
   const availableThemes = [
     { label: "light", value: "light" },
     { label: "dark", value: "dark" },

renderer/components/SettingsTab.tsx

@@ -4,6 +4,9 @@ import ReactTooltip from "react-tooltip";
 import { themeChange } from "theme-change";
 import log from "electron-log/renderer";
 import commands from "../../electron/commands";
+import { useAtom } from "jotai";
+import { customModelsPathAtom } from "../atoms/userSettingsAtom";
+import { TModelsList, modelsListAtom } from "../atoms/modelsListAtom";
 
 interface IProps {
   batchMode: boolean;
@@ -28,6 +31,7 @@ function SettingsTab({
   setSaveImageAs,
   logData,
 }: IProps) {
+  // STATES
   const [currentModel, setCurrentModel] = useState<{
     label: string;
     value: string;
@@ -37,6 +41,22 @@ function SettingsTab({
   });
 
   const [isCopied, setIsCopied] = useState(false);
+  const [customModelsPath, setCustomModelsPath] = useAtom(customModelsPathAtom);
+  const [customModelOptions, setCustomModelOptions] = useState<TModelsList>([]);
+  const [modelOptions, setModelOptions] = useAtom(modelsListAtom);
+
+  // EFFECTS
+  useEffect(() => {
+    // CUSTOM FOLDER LISTENER
+    window.electron.on(
+      commands.CUSTOM_MODEL_FILES_LIST,
+      (_, modelsList: string[]) => {
+        modelsList.forEach((model: string) => {
+          setModelOptions((prev) => [...prev, { label: model, value: model }]);
+        });
+      }
+    );
+  }, []);
 
   useEffect(() => {
     themeChange(false);
@@ -72,15 +92,12 @@ function SettingsTab({
     console.log("Current Model: ", currentModel);
   }, [currentModel]);
 
+  // HANDLERS
   const setExportType = (format: string) => {
     setSaveImageAs(format);
     localStorage.setItem("saveImageAs", format);
   };
 
-  const handleBatchMode = () => {
-    setBatchMode((oldValue) => !oldValue);
-  };
-
   const handleGpuIdChange = (e) => {
     setGpuId(e.target.value);
     localStorage.setItem("gpuId", e.target.value);
@@ -94,33 +111,6 @@ function SettingsTab({
     }, 2000);
   };
 
-  const customStyles = {
-    option: (provided, state) => ({
-      ...provided,
-      borderBottom: "1px dotted pink",
-      color: state.isSelected ? "red" : "blue",
-      padding: 20,
-    }),
-    control: () => ({
-      // none of react-select's styles are passed to <Control />
-      width: 200,
-    }),
-    singleValue: (provided, state) => {
-      const opacity = state.isDisabled ? 0.5 : 1;
-      const transition = "opacity 300ms";
-
-      return { ...provided, opacity, transition };
-    },
-  };
-
-  const modelOptions = [
-    { label: "General Photo (Real-ESRGAN)", value: "realesrgan-x4plus" },
-    { label: "General Photo (Remacri)", value: "remacri" },
-    { label: "General Photo (Ultramix Balanced)", value: "ultramix_balanced" },
-    { label: "General Photo (Ultrasharp)", value: "ultrasharp" },
-    { label: "Digital Art", value: "realesrgan-x4plus-anime" },
-  ];
-
   const availableThemes = [
     { label: "light", value: "light" },
     { label: "dark", value: "dark" },
@@ -153,8 +143,6 @@ function SettingsTab({
     { label: "winter", value: "winter" },
   ];
 
-  useEffect(() => {}, [imagePath]);
-
   return (
     <div className="animate-step-in animate flex h-screen flex-col gap-7 overflow-y-auto p-5 overflow-x-hidden">
       {/* IMAGE FORMAT BUTTONS */}
@@ -225,13 +213,17 @@ function SettingsTab({
       {/* GPU ID INPUT */}
       <div className="flex flex-col items-start gap-2">
         <p className="text-sm font-medium">Custom Models Path:</p>
-        <p className="text-sm text-base-content/60">/fasfas/asfasf/asf/saf</p>
+        <p className="text-sm text-base-content/60">{customModelsPath}</p>
         <button
           className="btn-primary btn"
           onClick={async () => {
-            const customModelPath = window.electron.invoke(
+            const customModelPath = await window.electron.invoke(
               commands.SELECT_CUSTOM_MODEL_FOLDER
             );
+
+            if (customModelPath !== null) {
+              setCustomModelsPath(customModelPath);
+            }
           }}>
           Select Folder
         </button>

renderer/pages/index.tsx

@@ -235,7 +235,7 @@ const Home = () => {
 
     var path = await window.electron.invoke(commands.SELECT_FILE);
 
-    if (path !== "cancelled") {
+    if (path !== null) {
       SetImagePath(path);
       var dirname = path.match(/(.*)[\/\\]/)[1] || "";
       setOutputPath(dirname);
@@ -247,7 +247,7 @@ const Home = () => {
 
     var path = await window.electron.invoke(commands.SELECT_FOLDER);
 
-    if (path !== "cancelled") {
+    if (path !== null) {
       setBatchFolderPath(path);
       setOutputPath(path + "_upscayled");
     }
@@ -339,7 +339,7 @@ const Home = () => {
 
   const outputHandler = async () => {
     var path = await window.electron.invoke(commands.SELECT_FOLDER);
-    if (path !== "cancelled") {
+    if (path !== null) {
       setOutputPath(path);
     } else {
       console.log("Getting output path from input file");