From b1bf3e4381f7d49572ee012c79b36a91da360466 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 9 Sep 2026 11:26:28 -0600 Subject: [PATCH] Add a config to force pure noise step at a set rate --- jobs/process/BaseSDTrainProcess.py | 5 +++++ toolkit/config_modules.py | 3 +++ 2 files changed, 8 insertions(+) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 9fa8e62..8220d8a 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1296,6 +1296,11 @@ class BaseSDTrainProcess(BaseTrainProcess): timestep_indices = timestep_indices.long() else: raise ValueError(f"Unknown content_or_style {content_or_style}") + + if self.train_config.first_timestep_chance > 0.0: + # index 0 is full noise; per-sample chance to force it + force_first = torch.rand((batch_size,), device=timestep_indices.device) < self.train_config.first_timestep_chance + timestep_indices = torch.where(force_first, torch.zeros_like(timestep_indices), timestep_indices) with self.timer('convert_timestep_indices_to_timesteps'): # convert the timestep_indices to a timestep timesteps = self.sd.noise_scheduler.timesteps[timestep_indices.long()] diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 2291b32..07a549e 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -552,6 +552,9 @@ class TrainConfig: self.target_norm_std = kwargs.get('target_norm_std', None) self.target_norm_std_value = kwargs.get('target_norm_std_value', 1.0) self.timestep_type = kwargs.get('timestep_type', 'sigmoid') # sigmoid, linear, lognorm_blend, next_sample, weighted, one_step + + self.first_timestep_chance = kwargs.get('first_timestep_chance', 0.0) + self.next_sample_timesteps = kwargs.get('next_sample_timesteps', 8) self.linear_timesteps = kwargs.get('linear_timesteps', False) self.linear_timesteps2 = kwargs.get('linear_timesteps2', False)