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

Added flux_shift as timestep type

Browse Source
This commit is contained in:
Jaret Burkett
2025-01-27 07:35:00 -07:00
parent 2141c6e06c
commit 34a1c6947a
3 changed files with 64 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
jobs/process/BaseSDTrainProcess.py
View File

@@ -79,7 +79,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.accelerator: Accelerator = get_accelerator()
if self.accelerator.is_local_main_process:
transformers.utils.logging.set_verbosity_warning()
diffusers.utils.logging.set_verbosity_info()
diffusers.utils.logging.set_verbosity_error()
else:
transformers.utils.logging.set_verbosity_error()
diffusers.utils.logging.set_verbosity_error()
@@ -1066,7 +1066,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.sd.noise_scheduler.set_train_timesteps(
num_train_timesteps,
device=self.device_torch,
timestep_type=timestep_type
timestep_type=timestep_type,
latents=latents
)
else:
self.sd.noise_scheduler.set_timesteps(