From e8573dad34aa45102ca322c19966ea96babe6fde Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 22 Jul 2026 10:11:08 -0600 Subject: [PATCH] Fix trailing progress bar print when stopping a job in the ui --- extensions_built_in/sd_trainer/DiffusionTrainer.py | 12 +++++++++++- run.py | 7 +++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/extensions_built_in/sd_trainer/DiffusionTrainer.py b/extensions_built_in/sd_trainer/DiffusionTrainer.py index 524db3cd..f7d2ad6d 100644 --- a/extensions_built_in/sd_trainer/DiffusionTrainer.py +++ b/extensions_built_in/sd_trainer/DiffusionTrainer.py @@ -317,7 +317,17 @@ class DiffusionTrainer(SDTrainer): super(DiffusionTrainer, self).on_error(e) if self.is_ui_trainer: try: - if self.accelerator.is_main_process and not self.is_stopping: + if isinstance(e, KeyboardInterrupt): + # SIGINT (UI stop button or ctrl+c) is a stop, not an error + self.is_stopping = True + progress_bar = getattr(self, "progress_bar", None) + if progress_bar is not None: + # silence the bar so tqdm doesn't repaint it at interpreter exit + progress_bar.disable = True + progress_bar.close() + if self.accelerator.is_main_process: + self.update_status("stopped", "Job stopped") + elif self.accelerator.is_main_process and not self.is_stopping: self.update_status("error", str(e)) self.update_db_key("step", self.last_save_step) asyncio.run(self.wait_for_all_async()) diff --git a/run.py b/run.py index af6b0d8f..a77a792a 100644 --- a/run.py +++ b/run.py @@ -128,8 +128,11 @@ def main(): except Exception as e2: print_acc(f"Error running on_error: {e2}") if not args.recover: - print_end_message(jobs_completed, jobs_failed) - raise e + print_acc("") + print_acc("========================================") + print_acc("Job stopped") + print_acc("========================================") + sys.exit(0) if __name__ == '__main__':