Skip to content

Commit 8fc30f8

Browse files
catboxanoncatboxanon
committed
Fix training for V-pred and ztSNR
1) Updates debiased estimation loss function for V-pred. 2) Prevents now-deprecated scaling of loss if ztSNR is enabled.
1 parent 012e7e6 commit 8fc30f8

10 files changed

Lines changed: 26 additions & 18 deletions

‎fine_tune.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -383,10 +383,10 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
383383

384384
if args.min_snr_gamma:
385385
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
386-
if args.scale_v_pred_loss_like_noise_pred:
386+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
387387
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
388388
if args.debiased_estimation_loss:
389-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
389+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
390390

391391
loss = loss.mean() # mean over batch dimension
392392
else:

‎library/custom_train_functions.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,10 +96,13 @@ def add_v_prediction_like_loss(loss, timesteps, noise_scheduler, v_pred_like_los
9696
return loss
9797

9898

99-
def apply_debiased_estimation(loss, timesteps, noise_scheduler):
99+
def apply_debiased_estimation(loss, timesteps, noise_scheduler, v_prediction=False):
100100
snr_t = torch.stack([noise_scheduler.all_snr[t] for t in timesteps]) # batch_size
101101
snr_t = torch.minimum(snr_t, torch.ones_like(snr_t) * 1000) # if timestep is 0, snr_t is inf, so limit it to 1000
102-
weight = 1 / torch.sqrt(snr_t)
102+
if v_prediction:
103+
weight = 1 / (snr_t + 1)
104+
else:
105+
weight = 1 / torch.sqrt(snr_t)
103106
loss = weight * loss
104107
return loss
105108

‎library/train_util.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3731,6 +3731,11 @@ def verify_training_args(args: argparse.Namespace):
37313731
raise ValueError(
37323732
"scale_v_pred_loss_like_noise_pred can be enabled only with v_parameterization / scale_v_pred_loss_like_noise_predはv_parameterizationが有効なときのみ有効にできます"
37333733
)
3734+
3735+
if args.scale_v_pred_loss_like_noise_pred and args.zero_terminal_snr:
3736+
raise ValueError(
3737+
"zero_terminal_snr enabled. scale_v_pred_loss_like_noise_pred will not be used / zero_terminal_snrが有効です。scale_v_pred_loss_like_noise_predは使用されません"
3738+
)
37343739

37353740
if args.v_pred_like_loss and args.v_parameterization:
37363741
raise ValueError(

‎sdxl_train.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -725,12 +725,12 @@ def optimizer_hook(parameter: torch.Tensor):
725725

726726
if args.min_snr_gamma:
727727
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
728-
if args.scale_v_pred_loss_like_noise_pred:
728+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
729729
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
730730
if args.v_pred_like_loss:
731731
loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
732732
if args.debiased_estimation_loss:
733-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
733+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
734734

735735
loss = loss.mean() # mean over batch dimension
736736
else:

‎sdxl_train_control_net_lllite.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -474,12 +474,12 @@ def remove_model(old_ckpt_name):
474474

475475
if args.min_snr_gamma:
476476
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
477-
if args.scale_v_pred_loss_like_noise_pred:
477+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
478478
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
479479
if args.v_pred_like_loss:
480480
loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
481481
if args.debiased_estimation_loss:
482-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
482+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
483483

484484
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
485485

‎sdxl_train_control_net_lllite_old.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -434,12 +434,12 @@ def remove_model(old_ckpt_name):
434434

435435
if args.min_snr_gamma:
436436
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
437-
if args.scale_v_pred_loss_like_noise_pred:
437+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
438438
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
439439
if args.v_pred_like_loss:
440440
loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
441441
if args.debiased_estimation_loss:
442-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
442+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
443443

444444
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
445445

‎train_db.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -370,10 +370,10 @@ def train(args):
370370

371371
if args.min_snr_gamma:
372372
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
373-
if args.scale_v_pred_loss_like_noise_pred:
373+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
374374
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
375375
if args.debiased_estimation_loss:
376-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
376+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
377377

378378
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
379379

‎train_network.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -993,12 +993,12 @@ def remove_model(old_ckpt_name):
993993

994994
if args.min_snr_gamma:
995995
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
996-
if args.scale_v_pred_loss_like_noise_pred:
996+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
997997
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
998998
if args.v_pred_like_loss:
999999
loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
10001000
if args.debiased_estimation_loss:
1001-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
1001+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
10021002

10031003
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
10041004

‎train_textual_inversion.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -598,12 +598,12 @@ def remove_model(old_ckpt_name):
598598

599599
if args.min_snr_gamma:
600600
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
601-
if args.scale_v_pred_loss_like_noise_pred:
601+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
602602
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
603603
if args.v_pred_like_loss:
604604
loss = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, args.v_pred_like_loss)
605605
if args.debiased_estimation_loss:
606-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
606+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
607607

608608
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
609609

‎train_textual_inversion_XTI.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -483,10 +483,10 @@ def remove_model(old_ckpt_name):
483483
loss = loss * loss_weights
484484
if args.min_snr_gamma:
485485
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
486-
if args.scale_v_pred_loss_like_noise_pred:
486+
if args.scale_v_pred_loss_like_noise_pred and not args.zero_terminal_snr:
487487
loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
488488
if args.debiased_estimation_loss:
489-
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler)
489+
loss = apply_debiased_estimation(loss, timesteps, noise_scheduler, args.v_parameterization)
490490

491491
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
492492

0 commit comments

Comments
 (0)