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

Fix saving issue if epoch/step not in checkpoint

Browse Source
This commit is contained in:
Kohya S
2023-03-27 21:22:32 +09:00
parent 238f01bc9c
commit 895b0b6ca7
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
library/model_util.py
View File

@@ -1046,10 +1046,14 @@ def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_p
key_count = len(state_dict.keys()) key_count = len(state_dict.keys())
new_ckpt = {'state_dict': state_dict} new_ckpt = {'state_dict': state_dict}
# epoch and global_step are sometimes not int
try:
if 'epoch' in checkpoint: if 'epoch' in checkpoint:
epochs += checkpoint['epoch'] epochs += checkpoint['epoch']
if 'global_step' in checkpoint: if 'global_step' in checkpoint:
steps += checkpoint['global_step'] steps += checkpoint['global_step']
except:
pass
new_ckpt['epoch'] = epochs new_ckpt['epoch'] = epochs
new_ckpt['global_step'] = steps new_ckpt['global_step'] = steps