From ec3bb3aca362b53f08498a3f48a995bcfd2be0ec Mon Sep 17 00:00:00 2001
From: Nayam Amarshe <25067102+NayamAmarshe@users.noreply.github.com>
Date: Fri, 11 Aug 2023 15:15:27 +0530
Subject: [PATCH] Fix batch upscale scale

---
 electron/index.ts                             | 54 ++++++++++++++++---
 .../settings-tab/ToggleOverwrite.tsx          |  9 +++-
 renderer/pages/index.tsx                      |  2 +-
 3 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/electron/index.ts b/electron/index.ts
index 9237deb..5f8d0a2 100644
--- a/electron/index.ts
+++ b/electron/index.ts
@@ -167,9 +167,9 @@ app.on("ready", async () => {
   // GET OVERWRITE SETTINGS FROM LOCAL STORAGE
   mainWindow.webContents
     .executeJavaScript('localStorage.getItem("overwrite");', true)
-    .then((lastSavedQuality: string | null) => {
-      if (lastSavedQuality !== null) {
-        quality = parseInt(lastSavedQuality);
+    .then((lastSavedOverwrite: boolean | null) => {
+      if (lastSavedOverwrite !== null) {
+        overwrite = lastSavedOverwrite;
       }
     });
 });
@@ -511,7 +511,7 @@ ipcMain.on(commands.FOLDER_UPSCAYL, async (event, payload) => {
   const model = payload.model;
   const gpuId = payload.gpuId;
   const saveImageAs = payload.saveImageAs;
-  const scale = payload.scale as string;
+  // const scale = payload.scale as string;
 
   // GET THE IMAGE DIRECTORY
   let inputDir = payload.batchFolderPath;
@@ -522,11 +522,29 @@ ipcMain.on(commands.FOLDER_UPSCAYL, async (event, payload) => {
     outputDir = outputFolderPath;
   }
 
+  const isDefaultModel = defaultModels.includes(model);
+
+  let scale = "4";
+  if (model.includes("x2")) {
+    scale = "2";
+  } else if (model.includes("x3")) {
+    scale = "3";
+  } else {
+    scale = "4";
+  }
+
+  outputDir += `_${model}_x${payload.scale}`;
   if (!fs.existsSync(outputDir)) {
     fs.mkdirSync(outputDir, { recursive: true });
   }
 
-  const isDefaultModel = defaultModels.includes(model);
+  // Delete .DS_Store files
+  fs.readdirSync(inputDir).forEach((file) => {
+    if (file === ".DS_Store") {
+      logit("🗑️ Deleting .DS_Store file");
+      fs.unlinkSync(inputDir + slash + file);
+    }
+  });
 
   // UPSCALE
   const upscayl = spawnUpscayl(
@@ -537,7 +555,7 @@ ipcMain.on(commands.FOLDER_UPSCAYL, async (event, payload) => {
       isDefaultModel ? modelsPath : customModelsFolderPath ?? modelsPath,
       model,
       gpuId,
-      saveImageAs,
+      "png",
       scale
     ),
     logit
@@ -555,7 +573,7 @@ ipcMain.on(commands.FOLDER_UPSCAYL, async (event, payload) => {
       commands.FOLDER_UPSCAYL_PROGRESS,
       data.toString()
     );
-    if (data.includes("invalid gpu") || data.includes("failed")) {
+    if (data.includes("invalid") || data.includes("failed")) {
       logit("❌ INVALID GPU OR INVALID FILES IN FOLDER - FAILED");
       failed = true;
       upscayl.kill();
@@ -575,6 +593,28 @@ ipcMain.on(commands.FOLDER_UPSCAYL, async (event, payload) => {
     if (!mainWindow) return;
     if (!failed && !stopped) {
       logit("💯 Done upscaling");
+      logit("♻ Scaling and converting now...");
+      // Get number of files in output folder
+      const files = fs.readdirSync(inputDir);
+      files.forEach(async (file) => {
+        console.log("Filename: ", file.slice(0, -3));
+        // Resize the image to the original size
+        const originalImage = await Jimp.read(inputDir + slash + file);
+        const newImage = await Jimp.read(
+          outputDir + slash + file.slice(0, -3) + "png"
+        );
+        newImage
+          .quality(100 - quality)
+          .scaleToFit(
+            originalImage.getWidth() * parseInt(payload.scale),
+            originalImage.getHeight() * parseInt(payload.scale)
+          )
+          .write(outputDir + slash + file);
+        if (saveImageAs !== "png") {
+          fs.unlinkSync(outputDir + slash + file.slice(0, -3) + "png");
+        }
+      });
+
       mainWindow.webContents.send(commands.FOLDER_UPSCAYL_DONE, outputDir);
     } else {
       upscayl.kill();
diff --git a/renderer/components/settings-tab/ToggleOverwrite.tsx b/renderer/components/settings-tab/ToggleOverwrite.tsx
index b95960d..4e2c810 100644
--- a/renderer/components/settings-tab/ToggleOverwrite.tsx
+++ b/renderer/components/settings-tab/ToggleOverwrite.tsx
@@ -1,4 +1,4 @@
-import React from "react";
+import React, { useEffect } from "react";
 
 type ToggleOverwriteProps = {
   overwrite: boolean;
@@ -6,6 +6,13 @@ type ToggleOverwriteProps = {
 };
 
 const ToggleOverwrite = ({ overwrite, setOverwrite }: ToggleOverwriteProps) => {
+  useEffect(() => {
+    if (!localStorage.getItem("overwrite")) {
+      localStorage.setItem("overwrite", JSON.stringify(overwrite));
+    } else {
+      setOverwrite(localStorage.getItem("overwrite"));
+    }
+  }, []);
   return (
     <div className="flex flex-col gap-2">
       <p className="text-sm font-medium">OVERWRITE PREVIOUS UPSCALE</p>
diff --git a/renderer/pages/index.tsx b/renderer/pages/index.tsx
index 4d5d255..b84cd88 100644
--- a/renderer/pages/index.tsx
+++ b/renderer/pages/index.tsx
@@ -670,7 +670,7 @@ const Home = () => {
               All done!
             </p>
             <button
-              className="bg-gradient-blue rounded-lg p-3 font-medium text-white/90 transition-colors"
+              className="btn bg-gradient-blue rounded-lg p-3 font-medium text-white/90 transition-colors"
               onClick={openFolderHandler}>
               Open Upscayled Folder
             </button>