Added button to cancel the current upscale (kill esrgan process)

Code/IO/IOUtils.cs

@@ -142,16 +142,23 @@ public static void ClearDir(string path)
             if (!Directory.Exists(path))
                 return;
             if (Logger.doLogIo) Logger.Log("[IOUtils] Clearing " + path);
-            DirectoryInfo directoryInfo = new DirectoryInfo(path);
-            FileInfo[] files = directoryInfo.GetFiles();
-            foreach (FileInfo fileInfo in files)
+            try
             {
-                fileInfo.Delete();
+                DirectoryInfo directoryInfo = new DirectoryInfo(path);
+                FileInfo[] files = directoryInfo.GetFiles();
+                foreach (FileInfo fileInfo in files)
+                {
+                    fileInfo.Delete();
+                }
+                DirectoryInfo[] directories = directoryInfo.GetDirectories();
+                foreach (DirectoryInfo directoryInfo2 in directories)
+                {
+                    directoryInfo2.Delete(recursive: true);
+                }
             }
-            DirectoryInfo[] directories = directoryInfo.GetDirectories();
-            foreach (DirectoryInfo directoryInfo2 in directories)
+            catch (Exception e)
             {
-                directoryInfo2.Delete(recursive: true);
+                Logger.Log($"Failed to clear {path}: {e.Message}");
             }
         }
 

Code/Main/MainForm.cs

@@ -79,8 +79,23 @@ private async void MainForm_Load(object sender, EventArgs e)
 			flowPanelRight.AutoScroll = true;
 
 			initialized = true;
+			BusyCheckLoop();
 		}
 
+		public async void BusyCheckLoop ()
+        {
+            while (true)
+            {
+				if (ActiveForm != this)
+				{
+					await Task.Delay(100);
+					continue;
+				}
+				cancelBtn.Visible = (Program.busy && Program.currentEsrganProcess != null && !Program.currentEsrganProcess.HasExited);
+				await Task.Delay(100);
+			}
+        }
+
 		public async Task CheckInstallation ()
         {
 			await Installer.Init();
@@ -708,5 +723,10 @@ private void MainForm_Resize(object sender, EventArgs e)
 			lastLeftScrollbarState = leftScrollbar;
 			lastRightScrollbarState = rightScrollbar;
 		}
+
+        private void cancelBtn_Click(object sender, EventArgs e)
+        {
+			Program.KillEsrgan();
+        }
     }
 }

Code/Main/Program.cs

@@ -33,6 +33,9 @@ internal static class Program
         public static List<Form> currentTemporaryForms = new List<Form>();  // Temp forms that get closed when something gets cancelled
         public static List<MsgBox> openMessageBoxes = new List<MsgBox>();  // Temp forms that get closed when something gets cancelled
 
+        public static Process currentEsrganProcess;
+        public static bool cancelled = false;
+
         public static bool busy;
 
         [STAThread]
@@ -63,14 +66,22 @@ private static void Main()
             Application.Run(new MainForm());
         }
 
-        static void Lock()
+        public static void KillEsrgan (bool cleanup = true)
         {
-
+            if (currentEsrganProcess == null || currentEsrganProcess.HasExited)
+                return;
+            cancelled = true;
+            OSUtils.KillProcessTree(currentEsrganProcess.Id);
+            if (cleanup)
+            {
+                IOUtils.ClearDir(Paths.imgInPath);
+                IOUtils.ClearDir(Paths.imgOutPath);
+                IOUtils.ClearDir(Paths.imgOutNcnnPath);
+            }
         }
 
         public static MsgBox ShowMessage(string msg, string title = "Message")
         {
-            //MessageBox.Show(msg, title, MessageBoxButtons.OK, MessageBoxIcon.None, MessageBoxDefaultButton.Button1, (MessageBoxOptions)0x40000);
             DialogQueue.Init();
             MsgBox msgBox = new MsgBox(msg.Replace("\n", Environment.NewLine), title);
             DialogQueue.ShowDialog(msgBox);

Code/OS/ESRGAN.cs

@@ -19,13 +19,12 @@ namespace Cupscale.OS
 {
     internal class ESRGAN
     {
-        private static Process currentProcess;
-
         public enum PreviewMode { None, Cutout, FullImage }
         public enum Backend { CUDA, CPU, NCNN }
 
         public static async Task DoUpscale(string inpath, string outpath, ModelData mdl, string tilesize, bool alpha, PreviewMode mode, Backend backend, bool showTileProgress = true)
         {
+            Program.cancelled = false;      // Reset cancel flag
             bool useJoey = Config.GetInt("esrganVer") == 0;
             try
             {
@@ -71,12 +70,14 @@ public static async Task DoUpscale(string inpath, string outpath, ModelData mdl,
             }
             catch (Exception e)
             {
+                Program.mainForm.SetProgress(0f, "Cancelled.");
+                if (Program.cancelled)
+                    return;
                 if (e.Message.Contains("No such file"))
                     Program.ShowMessage("An error occured during upscaling.\nThe upscale process seems to have exited before completion!", "Error");
                 else
                     Program.ShowMessage("An error occured during upscaling.", "Error");
                 Logger.Log("[ESRGAN] Upscaling Error: " + e.Message + "\n" + e.StackTrace);
-                Program.mainForm.SetProgress(0f, "Cancelled.");
             }
 
         }
@@ -166,7 +167,7 @@ public static async Task Run(string inpath, string outpath, string modelArg, str
                 esrganProcess.OutputDataReceived += OutputHandler;
                 esrganProcess.ErrorDataReceived += OutputHandler;
             }
-            currentProcess = esrganProcess;
+            Program.currentEsrganProcess = esrganProcess;
             esrganProcess.Start();
             if (!showWindow)
             {
@@ -225,7 +226,7 @@ public static async Task RunJoey(string inpath, string outpath, string modelArg,
                 esrganProcess.OutputDataReceived += OutputHandler;
                 esrganProcess.ErrorDataReceived += OutputHandler;
             }
-            currentProcess = esrganProcess;
+            Program.currentEsrganProcess = esrganProcess;
             esrganProcess.Start();
             if (!showWindow)
             {
@@ -250,16 +251,13 @@ public static async Task RunJoey(string inpath, string outpath, string modelArg,
         private static void OutputHandler(object sendingProcess, DataReceivedEventArgs output)
         {
             if (output == null || output.Data == null)
-            {
                 return;
-            }
+
             string data = output.Data;
             Logger.Log("[Python] " + data);
             if (data.ToLower().Contains("error"))
             {
-                if (currentProcess != null && !currentProcess.HasExited)
-                    currentProcess.Kill();
-
+                Program.KillEsrgan();
                 Program.ShowMessage("Error occurred: \n\n" + data + "\n\nThe ESRGAN process was killed to avoid lock-ups.", "Error");
             }
 
@@ -337,8 +335,8 @@ public static async Task RunNcnn(string inpath, string outpath, string modelPath
             {
                 ncnnProcess.OutputDataReceived += NcnnOutputHandler;
                 ncnnProcess.ErrorDataReceived += NcnnOutputHandler;
-            } 
-            currentProcess = ncnnProcess;
+            }
+            Program.currentEsrganProcess = ncnnProcess;
             ncnnProcess.Start();
             if (!showWindow)
             {
@@ -366,9 +364,7 @@ private static void NcnnOutputHandler(object sendingProcess, DataReceivedEventAr
             Logger.Log("[NCNN] " + data.Replace("\n", " ").Replace("\r", " "));
             if (data.Contains("failed"))
             {
-                if (currentProcess != null && !currentProcess.HasExited)
-                    currentProcess.Kill();
-
+                Program.KillEsrgan();
                 Program.ShowMessage("Error occurred: \n\n" + data + "\n\nThe ESRGAN-NCNN process was killed to avoid lock-ups.", "Error");
             }
 

Code/OS/OSUtils.cs

@@ -3,6 +3,7 @@
 using System.Security.Principal;
 using System;
 using System.Diagnostics;
+using System.Management;
 
 namespace Cupscale.OS
 {
@@ -52,5 +53,30 @@ public static Process NewProcess(bool hidden, string filename = "cmd.exe")
             Process proc = new Process();
             return SetStartInfo(proc, hidden, filename);
         }
+
+        public static void KillProcessTree(int pid)
+        {
+            ManagementObjectSearcher processSearcher = new ManagementObjectSearcher
+              ("Select * From Win32_Process Where ParentProcessID=" + pid);
+            ManagementObjectCollection processCollection = processSearcher.Get();
+
+            try
+            {
+                Process proc = Process.GetProcessById(pid);
+                if (!proc.HasExited) proc.Kill();
+            }
+            catch (ArgumentException)
+            {
+                // Process already exited.
+            }
+
+            if (processCollection != null)
+            {
+                foreach (ManagementObject mo in processCollection)
+                {
+                    KillProcessTree(Convert.ToInt32(mo["ProcessID"])); //kill child processes(also kills childrens of childrens etc.)
+                }
+            }
+        }
     }
 }

Code/UI/BatchUpscaleUI.cs

@@ -194,7 +194,8 @@ public static async Task Run (bool preprocess, bool postProcess = true, string o
             sw.Restart();
             await Task.WhenAll(tasks);
 
-            Program.mainForm.SetProgress(0, $"Done - Upscaling took {(sw.ElapsedMilliseconds / 1000f).ToString("0")}s");
+            if(!Program.cancelled)
+                Program.mainForm.SetProgress(0, $"Done - Upscaling took {(sw.ElapsedMilliseconds / 1000f).ToString("0")}s");
             Program.mainForm.SetBusy(false);
         }
 

Code/UI/PreviewUI.cs

@@ -5,6 +5,7 @@
 using System.IO;
 using System.Threading.Tasks;
 using System.Windows.Forms;
+using Cupscale.Forms;
 using Cupscale.ImageUtils;
 using Cupscale.IO;
 using Cupscale.Main;
@@ -79,7 +80,7 @@ public static async Task UpscaleImage()
                 if (useCpu) backend = ESRGAN.Backend.CPU;
                 if (useNcnn) backend = ESRGAN.Backend.NCNN;
                 await ESRGAN.DoUpscale(Paths.imgInPath, Paths.imgOutPath, mdl, Config.Get("tilesize"), Config.GetBool("alpha"), ESRGAN.PreviewMode.None, backend);
-                if(backend == ESRGAN.Backend.NCNN)
+                if (backend == ESRGAN.Backend.NCNN)
                     outImg = Directory.GetFiles(Paths.imgOutPath, "*.png*", SearchOption.AllDirectories)[0];
                 else
                     outImg = Directory.GetFiles(Paths.imgOutPath, "*.tmp", SearchOption.AllDirectories)[0];
@@ -91,12 +92,15 @@ public static async Task UpscaleImage()
             }
             catch (Exception e)
             {
+                Program.mainForm.SetProgress(0f, "Cancelled.");
+                if (Program.cancelled)
+                    return;
                 if (e.StackTrace.Contains("Index"))
                     Program.ShowMessage("The upscale process seems to have exited before completion!", "Error");
                 Logger.ErrorMessage("An error occured during upscaling:", e);
-                Program.mainForm.SetProgress(0f, "Cancelled.");
             }
-            Program.mainForm.SetProgress(0, $"Done - Upscaling took {(sw.ElapsedMilliseconds / 1000f).ToString("0.0")}s");
+            if (!Program.cancelled)
+                Program.mainForm.SetProgress(0, $"Done - Upscaling took {(sw.ElapsedMilliseconds / 1000f).ToString("0.0")}s");
             Program.mainForm.SetBusy(false);
         }
 
@@ -211,7 +215,8 @@ public static async void UpscalePreview(bool fullImage = false)
                 ModelData mdl = new ModelData(null, null, ModelData.ModelMode.Advanced);
                 await ESRGAN.DoUpscale(Paths.previewPath, Paths.previewOutPath, mdl, tilesize, alpha, prevMode, backend);
             }
-            Program.mainForm.SetProgress(0, $"Done - Upscaling took {(sw.ElapsedMilliseconds / 1000f).ToString("0.0")}s");
+            if (!Program.cancelled)
+                Program.mainForm.SetProgress(0, $"Done - Upscaling took {(sw.ElapsedMilliseconds / 1000f).ToString("0.0")}s");
             Program.mainForm.SetBusy(false);
         }