VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-17 09:18:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix applying image size to post_process_loss

Browse Source
This commit is contained in:
rockerBOO
2025-03-20 15:40:26 -04:00
parent 9b3d3332a2
commit 8d5a183cc5
6 changed files with 383 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
train_network.py
View File

@@ -316,7 +316,7 @@ class NetworkTrainer:
return noise_pred, target, timesteps, None
def post_process_loss(self, loss, args, timesteps: torch.IntTensor, noise_scheduler) -> torch.FloatTensor:
def post_process_loss(self, loss, args, timesteps: torch.IntTensor, noise_scheduler, latents: Optional[torch.Tensor]) -> torch.FloatTensor:
if args.min_snr_gamma:
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
if args.scale_v_pred_loss_like_noise_pred:
@@ -442,7 +442,7 @@ class NetworkTrainer:
is_train=is_train,
)
huber_c = train_util.get_huber_threshold_if_needed(args, timesteps, noise_scheduler)
huber_c = train_util.get_huber_threshold_if_needed(args, timesteps, latents, noise_scheduler)
loss = train_util.conditional_loss(noise_pred.float(), target.float(), args.loss_type, "none", huber_c)
if weighting is not None:
loss = loss * weighting
@@ -453,7 +453,7 @@ class NetworkTrainer:
loss_weights = batch["loss_weights"] # 各sampleごとのweight
loss = loss * loss_weights
loss = self.post_process_loss(loss, args, timesteps, noise_scheduler)
loss = self.post_process_loss(loss, args, timesteps, noise_scheduler, latents)
return loss.mean()