Merge pull request #521 from akshaal/fix/save_state

fix: don't save state if no --save-state arg given
This commit is contained in:
Kohya S
2023-05-21 08:48:48 +09:00
committed by GitHub

View File

@@ -3167,10 +3167,11 @@ def save_sd_model_on_epoch_end_or_stepwise(
print(f"removing old model: {remove_out_dir}") print(f"removing old model: {remove_out_dir}")
shutil.rmtree(remove_out_dir) shutil.rmtree(remove_out_dir)
if on_epoch_end: if args.save_state:
save_and_remove_state_on_epoch_end(args, accelerator, epoch_no) if on_epoch_end:
else: save_and_remove_state_on_epoch_end(args, accelerator, epoch_no)
save_and_remove_state_stepwise(args, accelerator, global_step) else:
save_and_remove_state_stepwise(args, accelerator, global_step)
def save_and_remove_state_on_epoch_end(args: argparse.Namespace, accelerator, epoch_no): def save_and_remove_state_on_epoch_end(args: argparse.Namespace, accelerator, epoch_no):