From 454722cc9730e52fc3fbfd1224f96dcd8d90fef2 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 7 Feb 2026 09:49:55 -0700 Subject: [PATCH] Add signal correction noise --- jobs/process/BaseSDTrainProcess.py | 10 ++++++++++ toolkit/config_modules.py | 1 + 2 files changed, 11 insertions(+) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 40e07514..58a3ae67 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1306,6 +1306,16 @@ class BaseSDTrainProcess(BaseTrainProcess): noise = noise * noise_multiplier + if self.train_config.do_signal_correction_noise: + batch_noise = latents.clone().to(noise.device, dtype=noise.dtype) + scn_scale = torch.randn( + batch_noise.shape[0], batch_noise.shape[1], 1, 1, + device=batch_noise.device, + dtype=batch_noise.dtype + ) + batch_noise = batch_noise * scn_scale + noise = noise + batch_noise + if self.train_config.random_noise_shift > 0.0: # get random noise -1 to 1 noise_shift = torch.randn( diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index d93f471b..bbe3887d 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -397,6 +397,7 @@ class TrainConfig: self.noise_multiplier = kwargs.get('noise_multiplier', 1.0) self.target_noise_multiplier = kwargs.get('target_noise_multiplier', 1.0) self.random_noise_multiplier = kwargs.get('random_noise_multiplier', 0.0) + self.do_signal_correction_noise = kwargs.get('do_signal_correction_noise', False) self.random_noise_shift = kwargs.get('random_noise_shift', 0.0) self.img_multiplier = kwargs.get('img_multiplier', 1.0) self.noisy_latent_multiplier = kwargs.get('noisy_latent_multiplier', 1.0)