From 0265c8e14888ee28c538751fa6d4fc19f91e6673 Mon Sep 17 00:00:00 2001 From: Nayam Amarshe <25067102+NayamAmarshe@users.noreply.github.com> Date: Sun, 9 Apr 2023 10:46:15 +0530 Subject: [PATCH] Added directory read and models list --- electron/commands.ts | 1 + electron/index.ts | 86 +++++++++++++++++++-- main/commands.js | 1 + main/index.js | 87 +++++++++++++++++++--- renderer/atoms/modelsListAtom.ts | 14 ++++ renderer/atoms/userSettingsAtom.ts | 6 ++ renderer/components/LeftPaneImageSteps.tsx | 12 +-- renderer/components/SettingsTab.tsx | 62 +++++++-------- renderer/pages/index.tsx | 6 +- 9 files changed, 211 insertions(+), 64 deletions(-) create mode 100644 renderer/atoms/modelsListAtom.ts create mode 100644 renderer/atoms/userSettingsAtom.ts diff --git a/electron/commands.ts b/electron/commands.ts index ee334c3..b42e72a 100644 --- a/electron/commands.ts +++ b/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; diff --git a/electron/index.ts b/electron/index.ts index 85f48c2..a1e9ab8 100644 --- a/electron/index.ts +++ b/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; } }); diff --git a/main/commands.js b/main/commands.js index 810b487..9bae7a6 100644 --- a/main/commands.js +++ b/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; diff --git a/main/index.js b/main/index.js index 6cf7dc5..96a7307 100644 --- a/main/index.js +++ b/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-----------------------------// diff --git a/renderer/atoms/modelsListAtom.ts b/renderer/atoms/modelsListAtom.ts new file mode 100644 index 0000000..6d7dfab --- /dev/null +++ b/renderer/atoms/modelsListAtom.ts @@ -0,0 +1,14 @@ +import { atom } from "jotai"; + +export type TModelsList = { + label: string; + value: string; +}[]; + +export const modelsListAtom = atom([ + { 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" }, +]); diff --git a/renderer/atoms/userSettingsAtom.ts b/renderer/atoms/userSettingsAtom.ts new file mode 100644 index 0000000..ae32caf --- /dev/null +++ b/renderer/atoms/userSettingsAtom.ts @@ -0,0 +1,6 @@ +import { atomWithStorage } from "jotai/utils"; + +export const customModelsPathAtom = atomWithStorage( + "customModelsPath", + null +); diff --git a/renderer/components/LeftPaneImageSteps.tsx b/renderer/components/LeftPaneImageSteps.tsx index 6b68f24..415bf88 100644 --- a/renderer/components/LeftPaneImageSteps.tsx +++ b/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" }, diff --git a/renderer/components/SettingsTab.tsx b/renderer/components/SettingsTab.tsx index d67321e..dc94e67 100644 --- a/renderer/components/SettingsTab.tsx +++ b/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([]); + 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 - 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 (
{/* IMAGE FORMAT BUTTONS */} @@ -225,13 +213,17 @@ function SettingsTab({ {/* GPU ID INPUT */}

Custom Models Path:

-

/fasfas/asfasf/asf/saf

+

{customModelsPath}

diff --git a/renderer/pages/index.tsx b/renderer/pages/index.tsx index 1954e76..80b1ad8 100644 --- a/renderer/pages/index.tsx +++ b/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");