VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #521 from akshaal/fix/save_state

Browse Source 
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
parent c971d9319c 0c942106bf
commit bc909e8359
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
library/train_util.py
View File

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