Pass min_snr_gamma_soft to apply_snr_weight and apply_snr_weight_for_flow_matching

Wire the --min_snr_gamma_soft flag (args.min_snr_gamma_soft) through to the
soft kwarg at all 21 call sites across the codebase:

apply_snr_weight (12 sites): train_network, fine_tune, train_db, train_leco,
sdxl_train, sdxl_train_leco, sdxl_train_control_net,
sdxl_train_control_net_lllite, sdxl_train_control_net_lllite_old,
train_control_net, train_textual_inversion, train_textual_inversion_XTI.

apply_snr_weight_for_flow_matching (9 sites): flux_train, sd3_train,
lumina_train, anima_train, flux_train_network, sd3_train_network,
lumina_train_network, hunyuan_image_train_network, anima_train_network.

Both functions already accept soft: bool = False. The --min_snr_gamma_soft
flag and soft-mode implementations in custom_train_functions.py were already
in place. This commit connects the training args to the function calls.

Adds test_apply_snr_weight_soft.py with 34 tests covering:
- Soft vs hard weight formula correctness for both DDPM and flow matching
- Edge cases (sigma=0, sigma=1, near-zero clamping)
- Gradient flow through soft weights
- Smoothness advantage of soft over hard at gamma threshold
- CLI argument parsing for --min_snr_gamma_soft

Author: 67372a

anima_train.py

@@ -613,7 +613,7 @@ def grad_hook(tensor: torch.Tensor):
 
                 # Min-SNR-γ for flow matching
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
 
                 if getattr(args, "contrastive_flow_matching", False) and latents.size(0) > 1:
                     # CRITICAL: .detach() prevents gradients flowing through negative samples

anima_train_network.py

@@ -1126,7 +1126,7 @@ def post_process_loss(self, loss, args, timesteps, noise_scheduler):
         if args.min_snr_gamma:
             # Anima timesteps are already in [0, 1] range (sigmas) — they are divided by 1000
             # in get_noise_pred_and_target before being returned
-            loss = apply_snr_weight_for_flow_matching(loss, timesteps, args.min_snr_gamma)
+            loss = apply_snr_weight_for_flow_matching(loss, timesteps, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
         return loss
 
     def get_sai_model_spec(self, args):

fine_tune.py

@@ -419,7 +419,7 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
                     loss = loss.mean([1, 2, 3])
 
                     if args.min_snr_gamma:
-                        loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization)
+                        loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma, args.v_parameterization, soft=args.min_snr_gamma_soft)
                     if args.scale_v_pred_loss_like_noise_pred:
                         loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
                     if args.debiased_estimation_loss:

flux_train.py

@@ -679,7 +679,7 @@ def grad_hook(parameter: torch.Tensor):
 
                 # Min-SNR-γ for flow matching
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
 
                 if args.masked_loss or ("alpha_masks" in batch and batch["alpha_masks"] is not None):
                     loss = apply_masked_loss(loss, batch)

flux_train_network.py

@@ -462,7 +462,7 @@ def post_process_loss(self, loss, args, timesteps, noise_scheduler):
         if args.min_snr_gamma:
             # Convert timesteps (in [0, 1000] range) to flow matching sigmas (in [0, 1] range)
             sigmas = timesteps / noise_scheduler.config.num_train_timesteps
-            loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+            loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
         return loss
 
     def get_sai_model_spec(self, args):

hunyuan_image_train_network.py

@@ -604,7 +604,7 @@ def post_process_loss(self, loss, args, timesteps, noise_scheduler):
         if args.min_snr_gamma:
             # Convert timesteps (in [0, 1000] range) to flow matching sigmas (in [0, 1] range)
             sigmas = timesteps / noise_scheduler.config.num_train_timesteps
-            loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+            loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
         return loss
 
     def get_sai_model_spec(self, args):

library/custom_train_functions.py

@@ -77,38 +77,65 @@ def enforce_zero_terminal_snr(betas):
     noise_scheduler.alphas = alphas
     noise_scheduler.alphas_cumprod = alphas_cumprod
 
+def apply_snr_weight(
+    loss: torch.Tensor,
+    timesteps: torch.IntTensor,
+    noise_scheduler,
+    gamma: float,
+    v_prediction: bool = False,
+    soft: bool = False,
+):
+    """Apply Min-SNR or Soft Min-SNR weighting to the loss.
 
-def apply_snr_weight(loss: torch.Tensor, timesteps: torch.IntTensor, noise_scheduler: DDPMScheduler, gamma: Number, v_prediction=False):
-    snr = torch.stack([noise_scheduler.all_snr[t] for t in timesteps])
-    min_snr_gamma = torch.minimum(snr, torch.full_like(snr, gamma))
-    if v_prediction:
-        snr_weight = torch.div(min_snr_gamma, snr + 1)
-    else:
-        snr_weight = torch.div(min_snr_gamma, snr)
+    Args:
+        loss: Per-element loss tensor.
+        timesteps: Timesteps for each sample in the batch.
+        noise_scheduler: DDPMScheduler containing the precomputed SNR values.
+        gamma: Min-SNR clipping value (typically 4.0 or 5.0).
+        v_prediction: Set to True if the model predicts velocity (v).
+        soft: Set to True to use the smooth Soft Min-SNR transition from the paper.
+    """
+    # Retrieve SNR values and move to the correct device
+    snr = torch.stack([noise_scheduler.all_snr[t] for t in timesteps]).to(device=loss.device)
 
-    snr_weight = snr_weight.to(dtype=loss.dtype, device=loss.device)
+    if soft:
+        if v_prediction:
+            snr_weight = (snr * gamma) / ((snr + gamma) * (snr + 1))
+        else:
+            snr_weight = gamma / (snr + gamma)
+    else:
+        min_snr_gamma = torch.minimum(snr, torch.full_like(snr, gamma))
+        if v_prediction:
+            snr_weight = torch.div(min_snr_gamma, snr + 1)
+        else:
+            snr_weight = torch.div(min_snr_gamma, snr)
 
-    loss = loss * snr_weight
-    return loss
+    snr_weight = snr_weight.to(dtype=loss.dtype)
 
+    # Ensure snr_weight dimensions match loss for proper broadcasting
+    while snr_weight.ndim < loss.ndim:
+        snr_weight = snr_weight.unsqueeze(-1)
 
-def apply_snr_weight_for_flow_matching(loss: torch.Tensor, sigmas: torch.Tensor, gamma: float) -> torch.Tensor:
-    """Apply Min-SNR-γ weighting for flow matching (rectified flow) models.
+    return loss * snr_weight
 
-    Computes the signal-to-noise ratio from sigma: SNR = (1 - σ)² / σ²
-    and applies the velocity-prediction weight: min(SNR, γ) / (SNR + 1).
 
-    This is the flow-matching analog of apply_snr_weight for DDPM models.
-    Flow matching velocity prediction (v = ε - x₀) is mathematically analogous
-    to v-prediction in DDPM, so the v-prediction formula is used.
+def apply_snr_weight_for_flow_matching(
+    loss: torch.Tensor, 
+    sigmas: torch.Tensor, 
+    gamma: float, 
+    soft: bool = False
+) -> torch.Tensor:
+    """Apply Min-SNR-γ or Soft Min-SNR-γ weighting for flow matching models.
 
-    Reference: https://arxiv.org/abs/2303.09556
+    Computes the signal-to-noise ratio from sigma: SNR = (1 - σ)² / σ²
+    and applies the velocity-prediction weight.
 
     Args:
         loss: Per-element loss tensor (any shape, e.g. (B,) or (B, C, H, W)).
         sigmas: Noise levels from the flow matching scheduler.
             Can be shape (B, 1, 1, 1), (B,), or broadcastable with loss.
-        gamma: Min-SNR gamma value. 5.0 is recommended by the paper.
+        gamma: Min-SNR gamma value.
+        soft: Set to True to use the smooth Soft Min-SNR transition from the paper.
 
     Returns:
         Weighted loss tensor (same shape as input loss).
@@ -119,11 +146,14 @@ def apply_snr_weight_for_flow_matching(loss: torch.Tensor, sigmas: torch.Tensor,
     # SNR in flow matching: (1 - σ)² / σ²
     snr = ((1.0 - sigma) / sigma) ** 2
 
-    # Cap SNR at gamma
-    min_snr = torch.minimum(snr, torch.full_like(snr, gamma))
-
-    # Velocity prediction weight: min(SNR, γ) / (SNR + 1)
-    snr_weight = min_snr / (snr + 1)
+    if soft:
+        # Velocity prediction weight using Soft Min-SNR
+        snr_weight = (snr * gamma) / ((snr + gamma) * (snr + 1))
+    else:
+        # Cap SNR at gamma
+        min_snr = torch.minimum(snr, torch.full_like(snr, gamma))
+        # Velocity prediction weight: min(SNR, γ) / (SNR + 1)
+        snr_weight = min_snr / (snr + 1)
 
     snr_weight = snr_weight.to(dtype=loss.dtype, device=loss.device)
 
@@ -173,6 +203,13 @@ def add_custom_train_arguments(parser: argparse.ArgumentParser, support_weighted
         default=None,
         help="gamma for reducing the weight of high loss timesteps. Lower numbers have stronger effect. 5 is recommended by paper. / 低いタイムステップでの高いlossに対して重みを減らすためのgamma値、低いほど効果が強く、論文では5が推奨",
     )
+
+    parser.add_argument(
+        "--min_snr_gamma_soft",
+        action="store_true",
+        help="Controls if min_snr_gamma uses soft implementation from https://arxiv.org/abs/2401.11605.",
+    )
+
     parser.add_argument(
         "--scale_v_pred_loss_like_noise_pred",
         action="store_true",

lumina_train.py

@@ -777,7 +777,7 @@ def grad_hook(parameter: torch.Tensor):
 
                 # Min-SNR-γ for flow matching
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
 
                 if args.masked_loss or (
                     "alpha_masks" in batch and batch["alpha_masks"] is not None

lumina_train_network.py

@@ -346,7 +346,7 @@ def post_process_loss(self, loss, args, timesteps, noise_scheduler):
         if args.min_snr_gamma:
             # Convert timesteps (in [0, 1000] range) to flow matching sigmas (in [0, 1] range)
             sigmas = timesteps / noise_scheduler.config.num_train_timesteps
-            loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+            loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
         return loss
 
     def get_sai_model_spec(self, args):

sd3_train.py

@@ -856,7 +856,7 @@ def grad_hook(parameter: torch.Tensor):
 
                 # Min-SNR-γ for flow matching (applied before spatial mean so sigmas broadcast correctly)
                 if args.min_snr_gamma:
-                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma)
+                    loss = apply_snr_weight_for_flow_matching(loss, sigmas, args.min_snr_gamma, soft=args.min_snr_gamma_soft)
 
                 loss = loss.mean([1, 2, 3])