1
0
mirror of https://github.com/upscayl/upscayl.git synced 2024-12-01 02:27:16 +01:00
upscayl/electron/commands/batch-upscayl.ts

175 lines
5.0 KiB
TypeScript
Raw Normal View History

2023-11-26 14:06:23 +01:00
import fs, { rmdir } from "fs";
import { getMainWindow } from "../main-window";
import {
childProcesses,
customModelsFolderPath,
2023-09-19 17:04:31 +02:00
noImageProcessing,
outputFolderPath,
saveOutputFolder,
2023-11-23 06:39:46 +01:00
setCompression,
2023-11-22 16:54:02 +01:00
setNoImageProcessing,
setStopped,
stopped,
} from "../utils/config-variables";
import logit from "../utils/logit";
import { spawnUpscayl } from "../utils/spawn-upscayl";
import { getBatchArguments } from "../utils/get-arguments";
import slash from "../utils/slash";
import { modelsPath } from "../utils/get-resource-paths";
import COMMAND from "../constants/commands";
import convertAndScale from "../utils/convert-and-scale";
import DEFAULT_MODELS from "../constants/models";
2023-11-22 16:54:02 +01:00
import { BatchUpscaylPayload } from "../../common/types/types";
import { ImageFormat } from "../utils/types";
2023-11-22 16:54:02 +01:00
const batchUpscayl = async (event, payload: BatchUpscaylPayload) => {
const mainWindow = getMainWindow();
if (!mainWindow) return;
// GET THE MODEL
const model = payload.model;
const gpuId = payload.gpuId;
const saveImageAs = payload.saveImageAs as ImageFormat;
// GET THE IMAGE DIRECTORY
let inputDir = payload.batchFolderPath;
// GET THE OUTPUT DIRECTORY
let outputDir = payload.outputPath;
if (saveOutputFolder === true && outputFolderPath) {
outputDir = outputFolderPath;
}
2023-11-22 16:54:02 +01:00
setNoImageProcessing(payload.noImageProcessing);
2023-11-23 06:39:46 +01:00
setCompression(parseInt(payload.compression));
2023-11-22 16:54:02 +01:00
const isDefaultModel = DEFAULT_MODELS.includes(model);
2023-11-26 08:22:57 +01:00
let initialScale = "4";
2023-11-10 14:15:48 +01:00
if (model.includes("x1")) {
2023-11-26 08:22:57 +01:00
initialScale = "1";
2023-11-10 14:15:48 +01:00
} else if (model.includes("x2")) {
2023-11-26 08:22:57 +01:00
initialScale = "2";
} else if (model.includes("x3")) {
2023-11-26 08:22:57 +01:00
initialScale = "3";
} else {
2023-11-26 08:22:57 +01:00
initialScale = "4";
}
2023-11-26 08:22:57 +01:00
const desiredScale = payload.scale as string;
2023-11-26 08:22:57 +01:00
outputDir +=
slash +
`upscayl_${model}_x${noImageProcessing ? initialScale : desiredScale}`;
if (!fs.existsSync(outputDir)) {
fs.mkdirSync(outputDir, { recursive: true });
}
2023-12-03 07:36:00 +01:00
// Delete .DS_Store files
fs.readdirSync(inputDir).forEach((file) => {
2023-12-03 07:36:00 +01:00
if (file === ".DS_Store") {
logit("🗑️ Deleting .DS_Store file");
fs.unlinkSync(inputDir + slash + file);
}
});
// UPSCALE
const upscayl = spawnUpscayl(
"realesrgan",
getBatchArguments(
2023-12-03 07:36:00 +01:00
inputDir,
outputDir,
isDefaultModel ? modelsPath : customModelsFolderPath ?? modelsPath,
model,
gpuId,
"png",
2023-11-26 08:22:57 +01:00
initialScale
),
logit
);
childProcesses.push(upscayl);
setStopped(false);
let failed = false;
const onData = (data: any) => {
if (!mainWindow) return;
data = data.toString();
mainWindow.webContents.send(
COMMAND.FOLDER_UPSCAYL_PROGRESS,
data.toString()
);
if (data.includes("invalid") || data.includes("failed")) {
logit("❌ INVALID GPU OR INVALID FILES IN FOLDER - FAILED");
failed = true;
upscayl.kill();
}
};
const onError = (data: any) => {
if (!mainWindow) return;
mainWindow.setProgressBar(-1);
mainWindow.webContents.send(
COMMAND.FOLDER_UPSCAYL_PROGRESS,
data.toString()
);
failed = true;
upscayl.kill();
mainWindow &&
mainWindow.webContents.send(
COMMAND.UPSCAYL_ERROR,
"Error upscaling image. Error: " + data
);
return;
};
const onClose = () => {
if (!mainWindow) return;
if (!failed && !stopped) {
logit("💯 Done upscaling");
logit("♻ Scaling and converting now...");
upscayl.kill();
mainWindow && mainWindow.webContents.send(COMMAND.SCALING_AND_CONVERTING);
2023-09-19 17:04:31 +02:00
if (noImageProcessing) {
logit("🚫 Skipping scaling and converting");
mainWindow.setProgressBar(-1);
mainWindow.webContents.send(COMMAND.FOLDER_UPSCAYL_DONE, outputDir);
return;
}
2023-11-26 14:06:23 +01:00
2023-12-03 07:36:00 +01:00
const files = fs.readdirSync(inputDir);
try {
files.forEach(async (file) => {
console.log("Filename: ", file.slice(0, -3));
await convertAndScale(
2023-12-03 07:36:00 +01:00
inputDir + slash + file,
outputDir + slash + file.slice(0, -3) + "png",
outputDir + slash + file.slice(0, -3) + saveImageAs,
2023-11-26 08:22:57 +01:00
desiredScale,
saveImageAs,
onError
);
// Remove the png file (default) if the saveImageAs is not png
if (saveImageAs !== "png") {
2023-12-03 07:36:00 +01:00
logit("Removing output PNG");
fs.unlinkSync(outputDir + slash + file.slice(0, -3) + "png");
}
});
mainWindow.webContents.send(COMMAND.FOLDER_UPSCAYL_DONE, outputDir);
} catch (error) {
logit("❌ Error processing (scaling and converting) the image.", error);
upscayl.kill();
mainWindow &&
mainWindow.webContents.send(
COMMAND.UPSCAYL_ERROR,
2023-10-14 10:41:57 +02:00
"Error processing (scaling and converting) the image. Please report this error on Upscayl GitHub Issues page.\n" +
error
);
}
} else {
upscayl.kill();
}
};
upscayl.process.stderr.on("data", onData);
upscayl.process.on("error", onError);
upscayl.process.on("close", onClose);
};
export default batchUpscayl;