SNR gamma fixes for v_prediction training (#5106)

Co-authored-by: bghira <[email protected]>

Author: bghira

examples/controlnet/train_controlnet_flax.py

@@ -908,6 +908,9 @@ def compute_loss(params, minibatch, sample_rng):
             if args.snr_gamma is not None:
                 snr = jnp.array(compute_snr(timesteps))
                 snr_loss_weights = jnp.where(snr < args.snr_gamma, snr, jnp.ones_like(snr) * args.snr_gamma) / snr
+                if noise_scheduler.config.prediction_type == "v_prediction":
+                    # velocity objective prediction requires SNR weights to be floored to a min value of 1.
+                    snr_loss_weights = snr_loss_weights + 1
                 loss = loss * snr_loss_weights
 
             loss = loss.mean()

examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py

@@ -875,6 +875,9 @@ def collate_fn(examples):
                     mse_loss_weights = (
                         torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
                     )
+                    if noise_scheduler.config.prediction_type == "v_prediction":
+                        # velocity objective prediction requires SNR weights to be floored to a min value of 1.
+                        mse_loss_weights = mse_loss_weights + 1
                     # We first calculate the original loss. Then we mean over the non-batch dimensions and
                     # rebalance the sample-wise losses with their respective loss weights.
                     # Finally, we take the mean of the rebalanced loss.

examples/text_to_image/train_text_to_image.py

@@ -955,6 +955,9 @@ def collate_fn(examples):
                     mse_loss_weights = (
                         torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
                     )
+                    if noise_scheduler.config.prediction_type == "v_prediction":
+                        # velocity objective prediction requires SNR weights to be floored to a min value of 1.
+                        mse_loss_weights = mse_loss_weights + 1
                     # We first calculate the original loss. Then we mean over the non-batch dimensions and
                     # rebalance the sample-wise losses with their respective loss weights.
                     # Finally, we take the mean of the rebalanced loss.

examples/text_to_image/train_text_to_image_lora.py

@@ -786,6 +786,9 @@ def collate_fn(examples):
                     mse_loss_weights = (
                         torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
                     )
+                    if noise_scheduler.config.prediction_type == "v_prediction":
+                        # velocity objective prediction requires SNR weights to be floored to a min value of 1.
+                        mse_loss_weights = mse_loss_weights + 1
                     # We first calculate the original loss. Then we mean over the non-batch dimensions and
                     # rebalance the sample-wise losses with their respective loss weights.
                     # Finally, we take the mean of the rebalanced loss.

examples/text_to_image/train_text_to_image_lora_sdxl.py

@@ -1075,6 +1075,9 @@ def compute_time_ids(original_size, crops_coords_top_left):
                     mse_loss_weights = (
                         torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
                     )
+                    if noise_scheduler.config.prediction_type == "v_prediction":
+                        # velocity objective prediction requires SNR weights to be floored to a min value of 1.
+                        mse_loss_weights = mse_loss_weights + 1
                     # We first calculate the original loss. Then we mean over the non-batch dimensions and
                     # rebalance the sample-wise losses with their respective loss weights.
                     # Finally, we take the mean of the rebalanced loss.