ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 02:31:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added ability to split up flux across gpus (experimental). Changed the way timestep scheduling works to prep for more specific schedules.

Browse Source
This commit is contained in:
Jaret Burkett
2024-12-31 07:06:55 -07:00
parent 8ef07a9c36
commit 4723f23c0d
5 changed files with 182 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
jobs/process/BaseSDTrainProcess.py
View File

@@ -970,10 +970,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.train_config.linear_timesteps2,
self.train_config.timestep_type == 'linear',
])
timestep_type = 'linear' if linear_timesteps else None
if timestep_type is None:
timestep_type = self.train_config.timestep_type
self.sd.noise_scheduler.set_train_timesteps(
num_train_timesteps,
device=self.device_torch,
linear=linear_timesteps
timestep_type=timestep_type
)
else:
self.sd.noise_scheduler.set_timesteps(