ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-13 22:49:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

allow flipping for point of interesting autocropping. allow num repeats. Fixed some bugs with new free u

Browse Source
This commit is contained in:
Jaret Burkett
2023-10-12 21:02:47 -06:00
parent 4e3b2c2569
commit 38e441a29c
5 changed files with 27 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
jobs/process/BaseSDTrainProcess.py
View File

@@ -928,6 +928,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
flush()
# self.step_num = 0
for step in range(self.step_num, self.train_config.steps):
if self.train_config.free_u:
self.sd.pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.1, b2=1.2)
self.progress_bar.unpause()
with torch.no_grad():
# if is even step and we have a reg dataset, use that
@@ -1000,6 +1002,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
if is_sample_step:
self.progress_bar.pause()
# print above the progress bar
if self.train_config.free_u:
self.sd.pipeline.disable_freeu()
self.sample(self.step_num)
self.progress_bar.unpause()
@@ -1047,6 +1051,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
# flush()
self.progress_bar.close()
if self.train_config.free_u:
self.sd.pipeline.disable_freeu()
self.sample(self.step_num + 1)
print("")
self.save()