From 9d279587b82260886480ef45844c25c50bf3371d Mon Sep 17 00:00:00 2001 From: Volpeon Date: Mon, 2 Jan 2023 17:02:41 +0100 Subject: Fix --- train_dreambooth.py | 3 +-- train_ti.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/train_dreambooth.py b/train_dreambooth.py index c0fe328..05f6cb5 100644 --- a/train_dreambooth.py +++ b/train_dreambooth.py @@ -1068,14 +1068,13 @@ def main(): # Create the pipeline using using the trained modules and save it. if accelerator.is_main_process: print("Finished! Saving final checkpoint and resume state.") - checkpointer.save_samples(0, args.sample_steps) + checkpointer.save_samples(global_step, args.sample_steps) checkpointer.save_model() accelerator.end_training() except KeyboardInterrupt: if accelerator.is_main_process: print("Interrupted, saving checkpoint and resume state...") - checkpointer.save_samples(0, args.sample_steps) checkpointer.save_model() accelerator.end_training() quit() diff --git a/train_ti.py b/train_ti.py index f1dbed1..5cbef06 100644 --- a/train_ti.py +++ b/train_ti.py @@ -1006,7 +1006,7 @@ def main(): if accelerator.is_main_process: print("Finished! Saving final checkpoint and resume state.") checkpointer.checkpoint(global_step + global_step_offset, "end") - checkpointer.save_samples(global_step_offset, args.sample_steps) + checkpointer.save_samples(global_step + global_step_offset, args.sample_steps) save_args(basepath, args, { "global_step": global_step + global_step_offset }) @@ -1016,7 +1016,6 @@ def main(): if accelerator.is_main_process: print("Interrupted, saving checkpoint and resume state...") checkpointer.checkpoint(global_step + global_step_offset, "end") - checkpointer.save_samples(global_step_offset, args.sample_steps) save_args(basepath, args, { "global_step": global_step + global_step_offset }) -- cgit v1.2.3-70-g09d2