ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Set step to the last step saved at when exiting

Browse Source
This commit is contained in:
Jaret Burkett
2025-02-23 13:21:22 -07:00
parent 60f848a877
commit 3e49337a58
4 changed files with 54 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
jobs/process/BaseSDTrainProcess.py
View File

@@ -92,6 +92,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.step_num = 0
self.start_step = 0
self.epoch_num = 0
self.last_save_step = 0
# start at 1 so we can do a sample at the start
self.grad_accumulation_step = 1
# if true, then we do not do an optimizer step. We are accumulating gradients
@@ -459,6 +460,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
step_num = ''
if step is not None:
self.last_save_step = step
# zeropad 9 digits
step_num = f"_{str(step).zfill(9)}"
@@ -1827,6 +1829,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.sd)
flush()
self.last_save_step = self.step_num
### HOOK ###
self.hook_before_train_loop()