RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Move guidance loss target selection when given a range before preservation so it is avaliable during preservation.

This commit is contained in:
Jaret Burkett 2026-08-11 06:57:33 -06:00
parent ab5fef8970
commit 8d4beedd04
1 changed files with 12 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
extensions_built_in/sd_trainer/SDTrainer.py
View File

@ -1925,6 +1925,18 @@ class SDTrainer(BaseSDTrainProcess):
else:
self.adapter.set_reference_images(None)
if self.train_config.do_guidance_loss and isinstance(self.train_config.guidance_loss_target, list):
batch_size = noisy_latents.shape[0]
# update the guidance value, random float between guidance_loss_target[0] and guidance_loss_target[1]
# sample before the prior prediction so the prior, main, uncond, and
# preservation passes all run at the same guidance values
self._guidance_loss_target_batch = [
random.uniform(
self.train_config.guidance_loss_target[0],
self.train_config.guidance_loss_target[1]
) for _ in range(batch_size)
]
prior_pred = None
do_inverted_masked_prior = False
@ -2032,16 +2044,6 @@ class SDTrainer(BaseSDTrainProcess):
pred_kwargs['down_block_additional_residuals'] = down_block_res_samples
pred_kwargs['mid_block_additional_residual'] = mid_block_res_sample
if self.train_config.do_guidance_loss and isinstance(self.train_config.guidance_loss_target, list):
batch_size = noisy_latents.shape[0]
# update the guidance value, random float between guidance_loss_target[0] and guidance_loss_target[1]
self._guidance_loss_target_batch = [
random.uniform(
self.train_config.guidance_loss_target[0],
self.train_config.guidance_loss_target[1]
) for _ in range(batch_size)
]
self.before_unet_predict()
if unconditional_embeds is not None: