github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Make DOP run in a single backward pass, should be faster and more stable. Show dop loss on ui

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-08 13:55:47 -06:00
parent 687def6f7a
commit cac3815b2c
2 changed files with 16 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -2057,10 +2057,6 @@ class SDTrainer(BaseSDTrainProcess):
)
if self.train_config.diff_output_preservation or self.train_config.blank_prompt_preservation:
# send the loss backwards otherwise checkpointing will fail
self.accelerator.backward(loss)
normal_loss = loss.detach() # dont send backward again
with torch.no_grad():
if self.train_config.diff_output_preservation:
preservation_embeds = self.diff_output_preservation_embeds.expand_to_batch(noisy_latents.shape[0])
@@ -2081,13 +2077,10 @@ class SDTrainer(BaseSDTrainProcess):
)
multiplier = self.train_config.diff_output_preservation_multiplier if self.train_config.diff_output_preservation else self.train_config.blank_prompt_preservation_multiplier
preservation_loss = torch.nn.functional.mse_loss(preservation_pred, prior_pred) * multiplier
self.accelerator.backward(preservation_loss)
self.additional_logs['loss/normal'] = loss.item()
self.additional_logs['loss/preservation'] = preservation_loss.item()
loss = loss + preservation_loss
loss = normal_loss + preservation_loss
loss = loss.clone().detach()
# require grad again so the backward wont fail
loss.requires_grad_(True)
# check if nan
if torch.isnan(loss):
print_acc("loss is nan")