Fix shifted audio-to-audio schedules across backends

optimized/mlx/models/defs/sa3_pipeline.py

@@ -89,15 +89,18 @@ def logsnr_shift(t: mx.array, anchor_logsnr: float = -6.2, logsnr_end: float = 2
 
 def build_pingpong_schedule(steps: int, sigma_max: float = 1.0,
                              use_logsnr_shift: bool = True) -> mx.array:
-    """Linear t from sigma_max → 0 in (steps+1) points, optionally warped by LogSNRShift.
+    """LogSNR schedule of (steps+1) points from sigma_max down to zero.
 
+    Warp the normalized grid before scaling it to sigma_max so audio-to-audio
+    schedules remain monotonic and agree with the init-latent mix.
     Returns mx.array shape (steps+1,) of float32.
     """
-    t = mx.linspace(sigma_max, 0.0, steps + 1, dtype=mx.float32)
+    t = mx.linspace(1.0, 0.0, steps + 1, dtype=mx.float32)
     if use_logsnr_shift:
-        t = logsnr_shift(t)
-        # Re-anchor start to sigma_max
+        t = logsnr_shift(t) * sigma_max
         t = mx.concatenate([mx.array([sigma_max], dtype=mx.float32), t[1:]], axis=0)
+    else:
+        t = t * sigma_max
     return t
 
 

optimized/tensorRT/scripts/sa3_trt_core.py

@@ -841,13 +841,17 @@ def step_captured(self, stream):
 
 # ─── Pingpong sampler ────────────────────────────────────────────────────
 def build_pingpong_schedule(steps, sigma_max=1.0, dist_shift=None, latent_len=None):
-    """Match SAT build_schedule for RF/RF-denoiser: linspace(sigma_max, 0, steps+1),
-    optionally dist-shifted, with t[0] forced back to sigma_max so the schedule's
-    starting point aligns with the init-mix's t."""
-    t = torch.linspace(sigma_max, 0.0, steps + 1, device="cuda")
+    """LogSNR schedule of (steps+1) points from sigma_max down to zero.
+
+    Warp the normalized grid before scaling it to sigma_max so audio-to-audio
+    schedules remain monotonic and agree with the init-latent mix.
+    """
+    t = torch.linspace(1.0, 0.0, steps + 1, device="cuda")
     if dist_shift is not None and latent_len is not None:
-        t = dist_shift.shift(t, latent_len)
+        t = dist_shift.shift(t, latent_len) * sigma_max
         t[0] = sigma_max
+    else:
+        t = t * sigma_max
     return t
 
 

stable_audio_3/inference/sampling.py

@@ -35,10 +35,13 @@ def build_schedule(
     """
     n_points = steps + 1 if include_endpoint else steps
 
+    # Warp a normalized schedule, then scale it to the requested audio-to-audio
+    # starting noise level.  Warping an already-scaled schedule can move interior
+    # points above sigma_max even when the shift itself is monotonic.
     if include_endpoint:
-        t = torch.linspace(sigma_max, 0, n_points, device=device)
+        t = torch.linspace(1.0, 0, n_points, device=device)
     else:
-        t = torch.linspace(sigma_max, 0, n_points + 1, device=device)[:-1]
+        t = torch.linspace(1.0, 0, n_points + 1, device=device)[:-1]
 
     if dist_shift is not None:
         seq_len = effective_seq_len if effective_seq_len is not None else fallback_seq_len
@@ -50,16 +53,17 @@ def build_schedule(
             seq_len = max(int(seq_len), 1)
         t = dist_shift.shift(t, seq_len)
 
-        # Ensure the first timestep remains aligned with sigma_max after shifting.
-        # This keeps the schedule consistent with the initialization in sample_diffusion(),
-        # which mixes init_data using sigma_max.
-        if isinstance(t, torch.Tensor):
-            sigma_max_tensor = t.new_tensor(sigma_max)
-            if t.ndim == 1:
-                t[0] = sigma_max_tensor
-            else:
-                # For batched/per-element schedules, enforce sigma_max at the first time index.
-                t[..., 0] = sigma_max_tensor
+    t = t * sigma_max
+
+    # Keep the schedule's first point exactly aligned with sample_diffusion()'s
+    # init_data mix. Distribution shifts preserve this endpoint; assigning it
+    # explicitly also avoids small floating-point drift.
+    if isinstance(t, torch.Tensor):
+        sigma_max_tensor = t.new_tensor(sigma_max)
+        if t.ndim == 1:
+            t[0] = sigma_max_tensor
+        else:
+            t[..., 0] = sigma_max_tensor
 
     return t
 
@@ -520,4 +524,4 @@ def sample_diffusion(
                 audio_mask = torch.nn.functional.pad(audio_mask, (0, sampled.shape[-1] - audio_mask.shape[-1]), value=False)
             sampled = sampled * audio_mask.to(sampled.dtype)
 
-    return sampled
+    return sampled

tests/test_schedule.py

@@ -0,0 +1,83 @@
+import torch
+
+from stable_audio_3.inference.distribution_shift import LogSNRShift
+from stable_audio_3.inference.sampling import build_schedule
+
+
+class SquareShift:
+    """Simple endpoint-preserving warp that exposes scaled-before-warp bugs."""
+
+    def shift(self, t, _seq_len):
+        return t.square()
+
+
+def test_audio_to_audio_shift_is_normalized_then_scaled():
+    schedule = build_schedule(
+        steps=4,
+        sigma_max=0.5,
+        dist_shift=SquareShift(),
+        fallback_seq_len=1,
+    )
+
+    expected = torch.linspace(1.0, 0.0, 5).square() * 0.5
+    assert torch.equal(schedule, expected)
+    assert torch.all(schedule[:-1] >= schedule[1:])
+    assert torch.all(schedule <= 0.5)
+
+
+def test_shifted_schedule_matches_init_mix_at_both_endpoints():
+    schedule = build_schedule(
+        steps=8,
+        sigma_max=0.37,
+        dist_shift=LogSNRShift(),
+        fallback_seq_len=646,
+    )
+
+    assert schedule[0] == schedule.new_tensor(0.37)
+    assert schedule[-1].item() == 0.0
+    assert torch.all(schedule[:-1] >= schedule[1:])
+    assert torch.all(schedule >= 0.0)
+    assert torch.all(schedule <= schedule[0])
+
+
+def test_full_noise_shifted_schedule_is_unchanged():
+    grid = torch.linspace(1.0, 0.0, 9)
+    shift = LogSNRShift()
+    previous_behavior = shift.shift(grid, 646)
+    previous_behavior[0] = 1.0
+
+    schedule = build_schedule(
+        steps=8,
+        sigma_max=1.0,
+        dist_shift=shift,
+        fallback_seq_len=646,
+    )
+
+    assert torch.equal(schedule, previous_behavior)
+
+
+def test_per_element_schedules_are_bounded_by_sigma_max():
+    schedule = build_schedule(
+        steps=4,
+        sigma_max=0.6,
+        dist_shift=LogSNRShift(),
+        effective_seq_len=torch.tensor([324, 646]),
+    )
+
+    assert schedule.shape == (2, 5)
+    assert torch.equal(schedule[:, 0], torch.tensor([0.6, 0.6]))
+    assert torch.equal(schedule[:, -1], torch.zeros(2))
+    assert torch.all(schedule[:, :-1] >= schedule[:, 1:])
+    assert torch.all(schedule <= 0.6)
+
+
+def test_unshifted_schedule_preserves_endpoint_option():
+    with_endpoint = build_schedule(steps=4, sigma_max=0.5)
+    without_endpoint = build_schedule(
+        steps=4,
+        sigma_max=0.5,
+        include_endpoint=False,
+    )
+
+    assert torch.equal(with_endpoint, torch.tensor([0.5, 0.375, 0.25, 0.125, 0.0]))
+    assert torch.equal(without_endpoint, torch.tensor([0.5, 0.375, 0.25, 0.125]))