github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Save before last sample

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-24 09:13:01 -06:00
parent 089e41dd1c
commit f743ccf7ef
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
jobs/process/BaseSDTrainProcess.py
View File

@@ -2607,12 +2607,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.progress_bar.close() self.progress_bar.close()
if self.train_config.free_u: if self.train_config.free_u:
self.sd.pipeline.disable_freeu() self.sd.pipeline.disable_freeu()
if self.accelerator.is_main_process:
self.save()
if not self.train_config.disable_sampling: if not self.train_config.disable_sampling:
self.sample(self.step_num) self.sample(self.step_num)
self.logger.commit(step=self.step_num) self.logger.commit(step=self.step_num)
print_acc("") print_acc("")
if self.accelerator.is_main_process: if self.accelerator.is_main_process:
self.save()
self.logger.finish() self.logger.finish()
self.accelerator.end_training() self.accelerator.end_training()