theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-18 22:08:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: prevent division by zero in trajectory imitation loss at last step

Browse Source
This commit is contained in:
Mr-Neutr0n
2026-02-11 19:51:25 +05:30
parent 1b47e1dc22
commit 0e6976a0ae
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
diffsynth/diffusion/loss.py
View File

@@ -91,7 +91,7 @@ class TrajectoryImitationLoss(torch.nn.Module):
progress_id_teacher = torch.argmin((timesteps_teacher - pipe.scheduler.timesteps[progress_id + 1]).abs())
latents_ = trajectory_teacher[progress_id_teacher]
target = (latents_ - inputs_shared["latents"]) / (sigma_ - sigma)
target = (latents_ - inputs_shared["latents"]) / (sigma_ - sigma).clamp(min=1e-6)
loss = loss + torch.nn.functional.mse_loss(noise_pred.float(), target.float()) * pipe.scheduler.training_weight(timestep)
return loss