From 79ed5177eba769df70151fce00b28d40a1a1f1c4 Mon Sep 17 00:00:00 2001 From: Julien 'Lta' BALLET Date: Thu, 3 Oct 2024 11:07:05 +0200 Subject: [PATCH] Support min/max timesteps in train_controlnet.py --- train_controlnet.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/train_controlnet.py b/train_controlnet.py index c9ac6c5a8..19cc79f03 100644 --- a/train_controlnet.py +++ b/train_controlnet.py @@ -444,8 +444,11 @@ def remove_model(old_ckpt_name): ) # Sample a random timestep for each image + min_timestep = args.min_timestep if args.min_timestep is not None else 0 + sched_max_timestep = noise_scheduler.config.num_train_timesteps + max_timestep = args.max_timestep if args.max_timestep is not None else sched_max_timestep timesteps, huber_c = train_util.get_timesteps_and_huber_c( - args, 0, noise_scheduler.config.num_train_timesteps, noise_scheduler, b_size, latents.device + args, min_timestep, max_timestep, noise_scheduler, b_size, latents.device ) # Add noise to the latents according to the noise magnitude at each timestep