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

Added flat snr gamma vs min. Fixes timestep timing

Browse Source
This commit is contained in:
Jaret Burkett committed 2023-10-29 15:41:55 -06:00
1 parent 3097865203
commit 436a09430e
4 files changed
+18 -9
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

jobs/process/BaseSDTrainProcess.py
+5 -5
View File
@@ -559,12 +559,12 @@ class BaseSDTrainProcess(BaseTrainProcess):
# for content / structure, it is best to favor earlier timesteps
# for style, it is best to favor later timesteps
timesteps = torch.rand((batch_size,), device=latents.device)
orig_timesteps = torch.rand((batch_size,), device=latents.device)
if self.train_config.content_or_style == 'style':
timesteps = timesteps ** 3 * self.sd.noise_scheduler.config['num_train_timesteps']
elif self.train_config.content_or_style == 'content':
timesteps = (1 - timesteps ** 3) * self.sd.noise_scheduler.config['num_train_timesteps']
if self.train_config.content_or_style == 'content':
timesteps = orig_timesteps ** 3 * self.sd.noise_scheduler.config['num_train_timesteps']
elif self.train_config.content_or_style == 'style':
timesteps = (1 - orig_timesteps ** 3) * self.sd.noise_scheduler.config['num_train_timesteps']
timesteps = value_map(
timesteps,