ai-upscale-module
Version:
Upscale images using AI
640 lines (592 loc) • 27 kB
JavaScript
// TODO:
// 1. Add support for mac and linux
// - d̶o̶w̶n̶l̶o̶a̶d̶ t̶h̶e̶ c̶o̶r̶r̶e̶c̶t̶ z̶i̶p̶ f̶i̶l̶e̶ (done, I think)
// - figure out how to run the upscaler
import fs from 'fs';
import AdmZip from 'adm-zip';
import fetch from 'node-fetch';
import { exec } from 'child_process';
import LargeDownload from 'large-download';
const flags = {
DOWNLOADING: "DOWNLOADING",
DOWNLOADED: "DOWNLOADED",
NOT_DOWNLOADED: "NOT_DOWNLOADED",
READY: "READY",
UNDEFINED: "UNDEFINED"
}
const modelsFileSize = 319059498;
let upscaleJobID = 0;
let execJobsCount = 0;
let downloadProgressCallback = null;
class Upscaler {
constructor(options) {
// first we check to see if the assets are downloaded (upscaler and models). This sets the flags.
// if they are not downloaded, we download them in the background
if (options == undefined || options == null) {
options = {};
}
if (options.defaultScale == undefined || options.defaultScale == null) {
options.defaultScale = 4;
}
if (options.defaultFormat == undefined || options.defaultFormat == null) {
options.defaultFormat = "jpg";
} else if (options.defaultFormat !== "jpg" && options.defaultFormat !== "png") {
options.defaultFormat = "jpg";
}
if (options.defaultOutputPath == undefined || options.defaultOutputPath == null) {
options.defaultOutputPath = "./output/upscaled/";
}
if (options.downloadProgressCallback === undefined || options.downloadProgressCallback === null) {
options.downloadProgressCallback = null;
}
if(options.maxJobs === undefined || options.maxJobs === null) {
options.maxJobs = 4;
}
if(options.defaultModel === undefined || options.defaultModel === null) {
options.defaultModel = "ultrasharp-2.0.1";
}
this.options = options;
downloadProgressCallback = options.downloadProgressCallback;
this.upscaler = {};
this.models = {};
this.models.status = flags.UNDEFINED;
this.upscaler.status = flags.UNDEFINED;
this.maxJobs = options.maxJobs;
this.defaultModel = options.defaultModel;
// console.log('Checking for assets');
this.status = "Checking for assets";
this.checkForAssets();
if (this.upscaler.status != flags.READY || this.models.status != flags.READY) {
// console.log('Assets not found. Downloading assets');
this.status = "Assets not found. Downloading assets";
this.downloadAssets().then((success) => {
//console.log('Assets downloaded');
this.checkForAssets();
this.status = "Assets downloaded, Upscaler ready";
}).catch((error) => {
// console.log('Error downloading assets');
this.status = "Error downloading assets";
});
}
this.upscaleJobs = [];
this.upscaleJobsRunningCount = 0;
this.finishedJobs = [];
this.jobRunner = null;
this.execJobs = [];
if (this.upscaler.status == flags.READY && this.models.status == flags.READY) {
this.status = flags.READY;
}
}
setDefaultModel(modelName) {
this.defaultModel = modelName;
}
getListOfModels() {
if(!this.models.status == flags.READY) return [];
let models = [];
let modelFolder = fs.readdirSync(this.models.path);
if (modelFolder.length !== 0) {
modelFolder.forEach((file, i) => {
let folder1 = null;
try {
folder1 = fs.readdirSync(this.models.path + file);
} catch (e) { } // do nothing
// check to see if "file" is a .param or .bin file
if (file.endsWith('.bin')) {
models.push(file.substring(0, file.lastIndexOf('.')));
} else if (folder1 !== undefined && folder1 !== null) { // go one more folder deep
folder1.forEach((file2, i) => {
let folder2 = null;
try {
folder2 = fs.readdirSync(this.models.path + file + '/' + file2);
} catch (e) { } // do nothing
// check to see if "file" is a .param or .bin file
if (file2.endsWith('.bin')) {
models.push(file2.substring(0, file2.lastIndexOf('.')));
}
});
}
});
}
return models;
}
checkForAssets() {
let upscalerFound = false;
let upscalerPath = "";
// find upscaler
if (fs.existsSync('./unzipped')) {
let unzippedFolder = fs.readdirSync('./unzipped');
let latestVersion = {};
latestVersion.versionNo = 0;
latestVersion.folderName = "";
if (unzippedFolder.length !== 0) {
unzippedFolder.forEach((file, i) => {
let versionNo = file.substring(file.indexOf('vulkan-') + 6, file.lastIndexOf("-"));
let versionNoArray = versionNo.split('.');
let versionNoInt = 0;
versionNoInt = parseInt(versionNoArray[2] * Math.pow(10, 6) + versionNoArray[0] * Math.pow(10, 3) + versionNoArray[1] * Math.pow(10, 0));
if (versionNoInt > latestVersion.versionNo) {
latestVersion.versionNo = versionNoInt;
latestVersion.folderName = file;
}
});
}
if (latestVersion.folderName !== "") {
let dir = `./unzipped/${latestVersion.folderName}/`;
let files = fs.readdirSync(dir);
let exeFile = files.find(file => file.endsWith('.exe'));
if (exeFile !== undefined) {
upscalerFound = true;
upscalerPath = dir + exeFile;
}
}
if (upscalerFound) {
this.upscaler.status = flags.READY;
this.upscaler.path = upscalerPath;
} else {
this.upscaler.status = flags.NOT_DOWNLOADED;
this.upscaler.path = "";
}
} else {
this.upscaler.status = flags.NOT_DOWNLOADED;
this.upscaler.path = "";
}
if (this.upscaler.status == flags.NOT_DOWNLOADED) {
console.log('Upscaler is not installed. Will attempt to acquire in background.');
}
// find upscale model
let modelsFound = false;
let modelsFolder = "";
// chekc to make sure model folder exists
if (fs.existsSync('./models')) {
let modelFolder = fs.readdirSync('./models');
if (modelFolder.length !== 0) {
modelFolder.forEach((file, i) => {
let folder1 = null;
try {
folder1 = fs.readdirSync('./models/' + file);
} catch (e) {
// do nothing
}
// check to see if "file" is a .param or .bin file
if (file.endsWith('.param') || file.endsWith('.bin')) {
modelsFound = true;
modelsFolder = './models/';
} else if (folder1 !== undefined) {
folder1.forEach((file2, i) => {
let folder2 = null;
try {
folder2 = fs.readdirSync('./models/' + file + '/' + file2);
} catch (e) {
// do nothing
}
// check to see if "file" is a .param or .bin file
if (file2.endsWith('.param') || file2.endsWith('.bin')) {
modelsFound = true;
modelsFolder = './models/' + file + '/';
} else if (folder2 !== undefined) {
folder2.forEach((file3, i) => {
// check to see if "file" is a .param or .bin file
if (file3.endsWith('.param') || file3.endsWith('.bin')) {
modelsFound = true;
modelsFolder = './models/' + file + '/' + file2 + '/';
}
});
}
});
}
});
}
}
if (modelsFound) {
this.models.status = flags.READY;
this.models.path = modelsFolder;
} else {
this.models.status = flags.NOT_DOWNLOADED;
this.models.path = "";
console.log('Models not found. Will attempt to acquire in background.');
}
}
removeZipFolder() {
// if zipped folder exists, remove it
let success = false;
try {
if (fs.existsSync('./zipped')) {
fs.rmSync('./zipped', { recursive: true });
}
// return true if successful or no zipped folder exists
success = true;
} catch (e) {
console.error('Error removing zipped folder');
}
return success;
}
async downloadAssets() {
return new Promise(async (resolve, reject) => {
const owner = 'upscayl';
const repo = 'upscayl-ncnn';
let downloadSuccess = false;
// create some folders
try {
// checked to see if zipped folder exists
if (!fs.existsSync('./zipped')) {
fs.mkdirSync('./zipped');
}
// checked to see if unzipped folder exists
if (!fs.existsSync('./unzipped')) {
fs.mkdirSync('./unzipped');
}
} catch (e) {
console.error('Error creating folders');
resolve(false);
}
// get platform name (windows, mac, linux) and set asset name
let platform = process.platform;
let assetName = "";
if (platform === 'win32') assetName = 'windows.zip';
else if (platform === 'darwin') assetName = 'macos.zip';
else if (platform === 'linux') assetName = 'ubuntu.zip';
else {
console.error('Platform not supported');
resolve(false);
}
// download windows upscaler
this.upscaler.status = flags.DOWNLOADING;
let tagName = await getLatestReleaseVersion(owner, repo);
let dlLink = await getReleaseDownloadLink(owner, repo, tagName, assetName);
downloadAndUnzip(dlLink, './zipped/realesrgan-ncnn-vulkan-' + tagName + '-' + assetName, 'unzipped/').then((success) => {
if (success) this.upscaler.status = flags.DOWNLOADED;
downloadSuccess = success;
}).catch((error) => {
resolve(false);
}).finally(() => {
if (this.models.status == flags.DOWNLOADED && this.upscaler.status == flags.DOWNLOADED) {
resolve(this.removeZipFolder());
}
});
if (this.models.status != flags.READY && this.models.status != flags.DOWNLOADED) {
this.models.status = flags.DOWNLOADING;
downloadAndUnzip('https://github.com/upscayl/custom-models/archive/refs/heads/main.zip', './zipped/main.zip', 'unzipped/').then((success) => {
if (success) {
//move unzipped folder to models folder
try {
// make sure the root folder exists
if (!fs.existsSync('./models')) {
fs.mkdirSync('./models');
}
// copy files from unzipped folder to models folder
// find the folder name
let unzippedModelsFolder = fs.readdirSync('./unzipped/custom-models-main/models/');
// check to make sure /models/custom-models-main/models/ exists
if (!fs.existsSync('./models/custom-models-main/')) {
fs.mkdirSync('./models/custom-models-main/');
}
if (!fs.existsSync('./models/custom-models-main/models/')) {
fs.mkdirSync('./models/custom-models-main/models/');
}
// copy files
let destFolder = './models/custom-models-main/models/';
unzippedModelsFolder.forEach((file, i) => {
fs.copyFileSync('./unzipped/custom-models-main/models/' + file, destFolder + file);
});
// remove extraneous files
fs.rmSync('./unzipped/custom-models-main/', { recursive: true });
downloadSuccess = true;
this.models.status = flags.DOWNLOADED;
} catch (e) {
resolve(false);
}
} else {
resolve(false);
}
}).catch((error) => {
resolve(false);
}).finally(() => {
if (this.models.status == flags.DOWNLOADED && this.upscaler.status == flags.DOWNLOADED) {
resolve(this.removeZipFolder());
}
});
}
});
}
setMaxJobs(maxJobs) {
this.maxJobs = maxJobs;
}
setDownloadProgressCallback(callback) {
downloadProgressCallback = callback;
}
async upscale(inputFile, outputPath = null, format = "", scale = -1, modelName = null) {
return new Promise(async (resolve, reject) => {
let job = {};
job.inputFile = inputFile;
job.outputPath = outputPath;
job.format = format;
job.scale = scale;
job.status = "waiting";
if (modelName !== null) job.modelName = modelName;
else job.modelName = this.defaultModel;
job.id = upscaleJobID++;
this.upscaleJobs.push(job);
if (this.jobRunner == null || this.jobRunner == undefined) this.jobRunner = this.processJobs();
resolve(job.id);
});
}
async getJobStatus(jobID) {
return new Promise(async (resolve, reject) => {
let job = this.upscaleJobs.find(job => job.id === jobID);
if (job == undefined) {
job = this.finishedJobs.find(job => job.id === jobID);
if (job == undefined)
resolve("not found");
else
resolve(job.status);
} else {
resolve(job.status);
}
});
}
getJob(jobID) {
let job = this.upscaleJobs.find(job => job.id === jobID);
if (job == undefined) {
job = this.finishedJobs.find(job => job.id === jobID);
if (job == undefined)
return null;
else
return job;
} else {
return job;
}
}
getNumberOfRunningJobs() {
return this.upscaleJobsRunningCount;
}
getNumberOfWaitingJobs() {
return this.upscaleJobs.length;
}
processJobs() {
return new Promise(async (resolve, reject) => {
let waiter = [];
while (this.upscaleJobs.length > 0 && this.upscaleJobsRunningCount < this.maxJobs) {
this.upscaleJobsRunningCount++;
let job = this.upscaleJobs.shift();
job.status = "processing";
waiter.push(this.upscaleJob(job.inputFile, job.outputPath, job.format, job.scale, job.modelName).then((success) => {
// console.log("Upscale completed/////////////////////////////////////////////////////////////////////////////////////////////////////");
if (success) {
job.status = "complete";
} else {
job.status = "failed";
}
}).catch((error) => {
// console.log("Upscale failed xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");
job.status = "failed";
console.error(error);
}).finally(() => {
this.upscaleJobsRunningCount--;
this.finishedJobs.push(job);
}));
}
Promise.all(waiter).then((obj) => {
// console.log("All jobs completed");
// console.log({obj});
if(this.upscaleJobs.length > 0) this.jobRunner = this.processJobs();
else this.jobRunner = null;
resolve();
});
});
}
async upscaleJob(inputFile, outputPath, format, scale, modelName) {
// console.log("Upscaling: ", inputFile);
if (outputPath == null) outputPath = this.options.defaultOutputPath;
if (format == "") format = this.options.defaultFormat;
if (scale == -1) scale = this.options.defaultScale;
return new Promise(async (resolve, reject) => {
if (this.upscaler.status != flags.READY || this.models.status != flags.READY) {
console.log('Upscaler is not ready');
resolve(false);
}
// check to see if inputFile exists
if (!fs.existsSync(inputFile)) {
console.log('File does not exist');
resolve(false);
}
// check to see if inputFile is a valid image
if (!inputFile.endsWith('.png')) {
console.log('File is not a valid image');
resolve(false);
}
let outputFile = inputFile.substring(inputFile.lastIndexOf('/') + 1, inputFile.lastIndexOf('.')) + '-upscaled.' + format;
outputPath = outputPath.substring(0, outputPath.lastIndexOf('/')); // outputPath without file name
//await waitSeconds(2);
try {
if (!fs.existsSync(outputPath)) {
// create output path
fs.mkdirSync(outputPath);
}
} catch (e) {
console.log(e);
process.exit(1);
}
if (format !== "jpg" && format !== "png") {
console.log('Format is not supported');
resolve(false);
}
if (scale !== 2 && scale !== 3 && scale !== 4) {
console.log('Scale is not supported');
resolve(false);
}
// console.log("About to upscale");
// run upscaler
// resolve absolute paths
this.upscaler.path = fs.realpathSync(this.upscaler.path);
inputFile = fs.realpathSync(inputFile);
outputFile = fs.realpathSync(outputPath) + '\\\\' + outputFile;
this.models.path = fs.realpathSync(this.models.path);
let execString = this.upscaler.path;
execString += " -i " + "\"" + inputFile + "\"";
execString += " -o " + "\"" + outputFile + "\"";
execString += " -f " + format;
execString += " -s " + scale;
execString += " -m " + "\"" + this.models.path + "\"";
execString += " -n " + modelName + " ";
// console.log("calling upscaler with command: ", execString);
let scalingExec = exec(execString, (err, stdout, stderr) => {
if (err) {
// console.log({err});
}
// console.log({stdout});
// console.log({stderr});
}).on('exit', (code) => {
// console.log("close code: " + code);
if(code == 0) resolve(true);
else resolve(false);
}).on('close', (code) => {
// console.log("close code: " + code);
if(code == 0) resolve(true);
else resolve(false);
});
let scalingTimeout = setTimeout(() => {
scalingExec.kill('SIGINT');
this.execJobs.splice(this.execJobs.findIndex(job => job.id === execJobsCount), 1); // TODO: test this
}, 1000 * 60 * 10);
this.execJobs.push({scalingExec, scalingTimeout, id: execJobsCount++});
});
}
}
const downloadAndUnzip = (url, zipPath, extractPath) => {
return new Promise(async (resolve, reject) => {
let waitingForFilename = true;
let modelsFile = false;
fetch(url).then((response) => {
const contentDisposition = response.headers.get("content-disposition");
// console.log("contentDisposition: ", contentDisposition);
if (contentDisposition.indexOf("custom-models-main.zip") != -1) {
modelsFile = true;
}
if (contentDisposition) {
const match = /filename=([^;]+)/.exec(contentDisposition);
if (match) {
const filename = match[1];
// console.log("File name:", filename);
} else {
// console.log("No filename found in Content-Disposition header");
}
} else {
// console.log("Content-Disposition header not found in the response");
}
}).catch((error) => {
console.error("Error:", error);
}).finally(() => { waitingForFilename = false; });
while (waitingForFilename) { await waitSeconds(1); }
try {
let downloading = true;
let downloadTotal = modelsFile ? modelsFileSize : 0;
const download = new LargeDownload({
link: url,
destination: zipPath,
timeout: 300000,
retries: 3,
onRetry: (error) => {
console.log("Download error. Retrying: ", { error }, { url }, { zipPath }, { extractPath });
},
onData: (downloaded, total) => {
// console.log( {downloaded}, {total});
downloadTotal = modelsFile ? modelsFileSize : parseInt(total);
if (!isNaN(downloadTotal)) {
// convert to MB and truncate to 2 decimal places
downloadTotal = (downloadTotal / 1000000).toFixed(2);
downloaded = (downloaded / 1000000).toFixed(2);
console.log("Download progress: ", (downloaded / downloadTotal).toFixed(2) * 100 + "%");
}
},
minSizeToShowProgress: Infinity
});
download.load().then(() => {
downloading = false;
download.onRetry = null;
}).catch(() => {
downloading = false;
download.onRetry = null;
reject(false);
});
while (downloading) {
await waitSeconds(0.5);
if (downloadProgressCallback !== undefined)if( downloadProgressCallback !== null) downloadProgressCallback();
}
console.log("Download complete");
const zip = new AdmZip(zipPath);
zip.extractAllTo(extractPath, true);
resolve(true);
} catch (error) {
console.error("Error: ", { error });
reject(false);
}
});
};
async function waitSeconds(count) {
// this holds the promise returned by the confirm function
return await new Promise((resolve) => {
setTimeout(() => {
// if the user hasn't pressed enter to cancel, cancel the confirmation promise and resolve the wait promise with false
resolve();
}, count * 1000);
});
};
async function getLatestReleaseVersion(owner, repo) {
try {
const response = await fetch(`https://api.github.com/repos/${owner}/${repo}/releases/latest`);
const data = await response.json();
return data.tag_name;
} catch (error) {
// console.error('Error:', error);
return null;
}
}
async function getReleaseDownloadLink(owner, repo, tagName, assetName) {
try {
// Get the release by tag name
const releaseResponse = await fetch(`https://api.github.com/repos/${owner}/${repo}/releases/tags/${tagName}`);
const releaseData = await releaseResponse.json();
// console.log("////////////////////////////////////////////");
// console.log(releaseData);
// Find the asset with the given name
let asset = false;
releaseData.assets.forEach((a) => {
// console.log(a.name);
if (a.name.lastIndexOf(assetName) != -1) {
// console.log("found");
asset = a;
// console.log(a.browser_download_url);
}
});
if (asset) {
// Return the download URL for the asset
return asset.browser_download_url;
} else {
// console.log(`Asset "${assetName}" not found in the release.`);
return null;
}
} catch (error) {
return null;
}
}
export default Upscaler;