From e914257dd670b5fa17855bd63c7fc9e787f12b56 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 3 Jun 2025 17:59:00 -0400 Subject: [PATCH] Undo num_timesteps change --- library/flux_train_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index c841f816..25145744 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -471,7 +471,7 @@ def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None): def get_noisy_model_input_and_timesteps( - args, noise_scheduler, latents, noise, device, dtype, num_timesteps=1000 + args, noise_scheduler, latents, noise, device, dtype ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, _, h, w = latents.shape sigmas = None @@ -503,7 +503,7 @@ def get_noisy_model_input_and_timesteps( sigmas = sigmas.sigmoid() mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) # we are pre-packed so must adjust for packed size sigmas = time_shift(mu, 1.0, sigmas) - timesteps = sigmas * num_timesteps + timesteps = noise_scheduler._sigma_to_t(sigmas) else: # Sample a random timestep for each image # for weighting schemes where we sample timesteps non-uniformly