diff --git a/Code/IO/IOUtils.cs b/Code/IO/IOUtils.cs index 7b40f72..84bb541 100644 --- a/Code/IO/IOUtils.cs +++ b/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}"); } } diff --git a/Code/Main/MainForm.Designer.cs b/Code/Main/MainForm.Designer.cs index 7add827..30adeec 100644 --- a/Code/Main/MainForm.Designer.cs +++ b/Code/Main/MainForm.Designer.cs @@ -95,6 +95,7 @@ private void InitializeComponent() this.vramLabel = new System.Windows.Forms.Label(); this.flowPanelRight = new System.Windows.Forms.FlowLayoutPanel(); this.upscalePanel = new System.Windows.Forms.Panel(); + this.cancelBtn = new HTAlt.WinForms.HTButton(); this.saveMergedPreviewBtn = new HTAlt.WinForms.HTButton(); this.label20 = new System.Windows.Forms.Label(); this.openOutFolderBtn = new HTAlt.WinForms.HTButton(); @@ -965,7 +966,7 @@ private void InitializeComponent() this.label5.Name = "label5"; this.label5.Size = new System.Drawing.Size(364, 34); this.label5.TabIndex = 1; - this.label5.Text = "Cupscale [Build 29 - 11/10/20]"; + this.label5.Text = "Cupscale [Build 31 - 11/12/20]"; this.label5.TextAlign = System.Drawing.ContentAlignment.MiddleLeft; // // panel7 @@ -1157,6 +1158,7 @@ private void InitializeComponent() this.upscalePanel.Anchor = ((System.Windows.Forms.AnchorStyles)((System.Windows.Forms.AnchorStyles.Top | System.Windows.Forms.AnchorStyles.Right))); this.upscalePanel.BackColor = System.Drawing.Color.FromArgb(((int)(((byte)(40)))), ((int)(((byte)(40)))), ((int)(((byte)(40))))); this.upscalePanel.BorderStyle = System.Windows.Forms.BorderStyle.FixedSingle; + this.upscalePanel.Controls.Add(this.cancelBtn); this.upscalePanel.Controls.Add(this.saveMergedPreviewBtn); this.upscalePanel.Controls.Add(this.label20); this.upscalePanel.Controls.Add(this.openOutFolderBtn); @@ -1166,6 +1168,23 @@ private void InitializeComponent() this.upscalePanel.Size = new System.Drawing.Size(310, 135); this.upscalePanel.TabIndex = 0; // + // cancelBtn + // + this.cancelBtn.Anchor = ((System.Windows.Forms.AnchorStyles)(((System.Windows.Forms.AnchorStyles.Bottom | System.Windows.Forms.AnchorStyles.Left) + | System.Windows.Forms.AnchorStyles.Right))); + this.cancelBtn.BackColor = System.Drawing.Color.FromArgb(((int)(((byte)(64)))), ((int)(((byte)(64)))), ((int)(((byte)(64))))); + this.cancelBtn.FlatAppearance.BorderSize = 0; + this.cancelBtn.FlatStyle = System.Windows.Forms.FlatStyle.Flat; + this.cancelBtn.ForeColor = System.Drawing.Color.Tomato; + this.cancelBtn.Location = new System.Drawing.Point(6, 96); + this.cancelBtn.Name = "cancelBtn"; + this.cancelBtn.Size = new System.Drawing.Size(297, 30); + this.cancelBtn.TabIndex = 13; + this.cancelBtn.Text = "Cancel Upscale"; + this.cancelBtn.UseVisualStyleBackColor = false; + this.cancelBtn.Visible = false; + this.cancelBtn.Click += new System.EventHandler(this.cancelBtn_Click); + // // saveMergedPreviewBtn // this.saveMergedPreviewBtn.Anchor = ((System.Windows.Forms.AnchorStyles)(((System.Windows.Forms.AnchorStyles.Bottom | System.Windows.Forms.AnchorStyles.Left) @@ -1429,10 +1448,10 @@ private void InitializeComponent() "Pixels Width", "Pixels Longer Side", "Pixels Shorter Side"}); - this.postResizeMode.Location = new System.Drawing.Point(121, 44); + this.postResizeMode.Location = new System.Drawing.Point(124, 44); this.postResizeMode.Margin = new System.Windows.Forms.Padding(8); this.postResizeMode.Name = "postResizeMode"; - this.postResizeMode.Size = new System.Drawing.Size(182, 21); + this.postResizeMode.Size = new System.Drawing.Size(179, 21); this.postResizeMode.TabIndex = 6; this.postResizeMode.SelectedIndexChanged += new System.EventHandler(this.postResizeMode_SelectedIndexChanged); // @@ -1494,7 +1513,7 @@ private void InitializeComponent() this.postResizeScale.Location = new System.Drawing.Point(5, 44); this.postResizeScale.Margin = new System.Windows.Forms.Padding(8); this.postResizeScale.Name = "postResizeScale"; - this.postResizeScale.Size = new System.Drawing.Size(122, 21); + this.postResizeScale.Size = new System.Drawing.Size(103, 21); this.postResizeScale.TabIndex = 2; // // label9 @@ -1622,10 +1641,10 @@ private void InitializeComponent() "Mitchell", "Nearest Neighbor", "Bicubic"}); - this.preResizeFilter.Location = new System.Drawing.Point(8, 94); + this.preResizeFilter.Location = new System.Drawing.Point(5, 94); this.preResizeFilter.Margin = new System.Windows.Forms.Padding(8); this.preResizeFilter.Name = "preResizeFilter"; - this.preResizeFilter.Size = new System.Drawing.Size(295, 21); + this.preResizeFilter.Size = new System.Drawing.Size(298, 21); this.preResizeFilter.TabIndex = 4; // // preResizeScale @@ -1644,7 +1663,7 @@ private void InitializeComponent() this.preResizeScale.Location = new System.Drawing.Point(5, 44); this.preResizeScale.Margin = new System.Windows.Forms.Padding(8); this.preResizeScale.Name = "preResizeScale"; - this.preResizeScale.Size = new System.Drawing.Size(122, 21); + this.preResizeScale.Size = new System.Drawing.Size(103, 21); this.preResizeScale.TabIndex = 2; // // label18 @@ -1798,9 +1817,9 @@ private void InitializeComponent() this.refreshPreviewCutoutBtn.FlatAppearance.BorderSize = 0; this.refreshPreviewCutoutBtn.FlatStyle = System.Windows.Forms.FlatStyle.Flat; this.refreshPreviewCutoutBtn.ForeColor = System.Drawing.Color.White; - this.refreshPreviewCutoutBtn.Location = new System.Drawing.Point(184, 62); + this.refreshPreviewCutoutBtn.Location = new System.Drawing.Point(185, 62); this.refreshPreviewCutoutBtn.Name = "refreshPreviewCutoutBtn"; - this.refreshPreviewCutoutBtn.Size = new System.Drawing.Size(172, 30); + this.refreshPreviewCutoutBtn.Size = new System.Drawing.Size(171, 30); this.refreshPreviewCutoutBtn.TabIndex = 11; this.refreshPreviewCutoutBtn.Text = "Refresh Preview (Cutout)"; this.refreshPreviewCutoutBtn.UseVisualStyleBackColor = false; @@ -2632,5 +2651,6 @@ private void InitializeComponent() private TextBox videoLogBox; private Label label35; private ComboBox videoOutputFormat; + private HTAlt.WinForms.HTButton cancelBtn; } } diff --git a/Code/Main/MainForm.cs b/Code/Main/MainForm.cs index 1ced6a1..62a880d 100644 --- a/Code/Main/MainForm.cs +++ b/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(); + } } } diff --git a/Code/Main/Program.cs b/Code/Main/Program.cs index 12210e2..fe1ce0d 100644 --- a/Code/Main/Program.cs +++ b/Code/Main/Program.cs @@ -33,6 +33,9 @@ internal static class Program public static List
currentTemporaryForms = new List(); // Temp forms that get closed when something gets cancelled public static List openMessageBoxes = new List(); // 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); diff --git a/Code/OS/ESRGAN.cs b/Code/OS/ESRGAN.cs index 5b373a3..f57c9cf 100644 --- a/Code/OS/ESRGAN.cs +++ b/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"); } diff --git a/Code/OS/OSUtils.cs b/Code/OS/OSUtils.cs index 5d1bd65..62b11dc 100644 --- a/Code/OS/OSUtils.cs +++ b/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.) + } + } + } } } diff --git a/Code/UI/BatchUpscaleUI.cs b/Code/UI/BatchUpscaleUI.cs index 8f9c203..2ca2008 100644 --- a/Code/UI/BatchUpscaleUI.cs +++ b/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); } diff --git a/Code/UI/PreviewUI.cs b/Code/UI/PreviewUI.cs index b181ad7..253dc09 100644 --- a/Code/UI/PreviewUI.cs +++ b/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); }