ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-26 09:19:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reworked timestep distribution on flowmatch sampler when training.

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-08 06:01:45 -06:00
parent acafe9984f
commit e69a520616
2 changed files with 27 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
jobs/process/BaseSDTrainProcess.py
View File

@@ -906,6 +906,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.sd.noise_scheduler.set_timesteps(
num_train_timesteps, device=self.device_torch, original_inference_steps=num_train_timesteps
)
elif self.train_config.noise_scheduler == 'flowmatch':
self.sd.noise_scheduler.set_train_timesteps(
num_train_timesteps, device=self.device_torch
)
else:
self.sd.noise_scheduler.set_timesteps(
num_train_timesteps, device=self.device_torch