From 0f075fc45e51fa5b686d510f5aece038ceaf4e33 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 22 Mar 2026 08:30:13 -0600 Subject: [PATCH] Adjust signal amplification target. Allow signal amplification strength in config. --- extensions_built_in/sd_trainer/SDTrainer.py | 2 ++ toolkit/config_modules.py | 1 + 2 files changed, 3 insertions(+) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 2ae346b6..13442d6d 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -578,6 +578,8 @@ class SDTrainer(BaseSDTrainProcess): raise ValueError("Signal amplification is only supported for flow matching models") with torch.no_grad(): nas = 1.0 - (timesteps / 1000).to(noise.device, dtype=noise.dtype) + nas = nas * self.train_config.signal_amplification_strength + nas = nas.clamp(min=0.1) while len(nas.shape) < len(noise.shape): nas = nas.unsqueeze(-1) aug = batch.latents * nas diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 428574f4..228ad3f3 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -402,6 +402,7 @@ class TrainConfig: self.do_batch_noise_correction = kwargs.get('do_batch_noise_correction', False) self.batch_noise_correction_scale = kwargs.get('batch_noise_correction_scale', 0.1) self.do_signal_amplification = kwargs.get('do_signal_amplification', False) + self.signal_amplification_strength = kwargs.get('signal_amplification_strength', 0.5) self.signal_correction_noise_scale = kwargs.get('signal_correction_noise_scale', 1.0) self.random_noise_shift = kwargs.get('random_noise_shift', 0.0)