Clamp negative SDE sampler tmp

nemo/collections/audio/parts/submodules/schroedinger_bridge.py

@@ -557,6 +557,7 @@ def forward(
                 # Calculate scaling for the first-order discretization from the paper
                 weight_prev = alpha_t * sigma_t**2 / (alpha_prev * sigma_prev**2 + self.eps)
                 tmp = 1 - sigma_t**2 / (sigma_prev**2 + self.eps)
+                tmp = tmp.clamp_min(0)
                 weight_estimate = alpha_t * tmp
                 weight_z = alpha_t * sigma_t * torch.sqrt(tmp)
 

tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py

@@ -16,11 +16,90 @@
 import pytest
 import torch
 
-from nemo.collections.audio.parts.submodules.schroedinger_bridge import SBNoiseScheduleVE, SBNoiseScheduleVP, SBSampler
+from nemo.collections.audio.parts.submodules.schroedinger_bridge import (
+    SBNoiseSchedule,
+    SBNoiseScheduleVE,
+    SBNoiseScheduleVP,
+    SBSampler,
+)
 
 NUM_STEPS = [1, 5, 10, 20, 100]
 
 
+class SlightlyIncreasingSigmaNoiseSchedule(SBNoiseSchedule):
+    def __init__(
+        self,
+        sigma_base: float = 1.0,
+        delta: float = 1e-6,
+        time_min: float = 1.0 - 1e-6,
+        time_max: float = 1.0,
+        num_steps: int = 1,
+        eps: float = 1e-8,
+    ):
+        super().__init__(time_min=time_min, time_max=time_max, num_steps=num_steps, eps=eps)
+        self.sigma_base = sigma_base
+        self.delta = delta
+
+    def f(self, time: torch.Tensor) -> torch.Tensor:
+        return torch.zeros_like(time)
+
+    def g(self, time: torch.Tensor) -> torch.Tensor:
+        return torch.ones_like(time)
+
+    def alpha(self, time: torch.Tensor) -> torch.Tensor:
+        return torch.ones_like(time)
+
+    def sigma(self, time: torch.Tensor) -> torch.Tensor:
+        sigma_base = torch.full_like(time, self.sigma_base)
+        delta = torch.full_like(time, self.delta)
+        return torch.where(time < self.time_max, sigma_base + delta, sigma_base)
+
+    def copy(self):
+        return SlightlyIncreasingSigmaNoiseSchedule(
+            sigma_base=self.sigma_base,
+            delta=self.delta,
+            time_min=self.time_min,
+            time_max=self.time_max,
+            num_steps=self.num_steps,
+            eps=self.eps,
+        )
+
+
+@pytest.mark.unit
+def test_sb_sampler_sde_clamps_negative_tmp_before_sqrt():
+    class IdentityEstimator(torch.nn.Module):
+        def forward(self, input, input_length, condition):
+            return input, input_length
+
+    noise_schedule = SlightlyIncreasingSigmaNoiseSchedule()
+    sampler = SBSampler(
+        noise_schedule=noise_schedule,
+        estimator=IdentityEstimator(),
+        estimator_output='data_prediction',
+        process='sde',
+        num_steps=1,
+    )
+
+    init_state = torch.ones(1, 1, 1, 2)
+
+    time_prev = torch.tensor([sampler.time_max], device=init_state.device)
+    time = torch.tensor([sampler.time_min], device=init_state.device)
+    sigma_prev, _, _ = sampler.noise_schedule.get_sigmas(time_prev)
+    sigma_t, _, _ = sampler.noise_schedule.get_sigmas(time)
+    raw_tmp = 1 - sigma_t**2 / (sigma_prev**2 + sampler.eps)
+
+    assert raw_tmp.item() < 0
+
+    sample, sample_length = sampler.forward(prior_mean=init_state, estimator_condition=None, state_length=None)
+
+    expected_scale = sigma_t**2 / (sigma_prev**2 + sampler.eps)
+    expected = expected_scale.view(-1, 1, 1, 1) * init_state
+
+    assert sample_length is None
+    assert torch.isfinite(sample).all()
+    torch.testing.assert_close(sample, expected)
+
+
 @pytest.mark.parametrize("num_steps", NUM_STEPS)
 @pytest.mark.parametrize("process", ["sde", "ode"])
 @pytest.mark.parametrize("noise_schedule_type", ["ve", "vp"])