diff --git a/nemo/collections/audio/parts/submodules/schroedinger_bridge.py b/nemo/collections/audio/parts/submodules/schroedinger_bridge.py index 77a971ed8923..92d60bad3031 100644 --- a/nemo/collections/audio/parts/submodules/schroedinger_bridge.py +++ b/nemo/collections/audio/parts/submodules/schroedinger_bridge.py @@ -44,6 +44,7 @@ def __init__( eps: float = 1e-8, ): super().__init__() + self.register_buffer('_time_reference', torch.empty(0), persistent=False) # min and max time if time_min < 0: @@ -86,17 +87,30 @@ def generate_time(self, size: int, device: torch.device) -> torch.Tensor: time = torch.rand(size, device=device) * self.time_delta + self.time_min return time + def _time_max_like(self, ref: Optional[torch.Tensor] = None) -> torch.Tensor: + """Return `time_max` on the same device and dtype as `ref`.""" + if ref is None: + ref = self._time_reference + + return ref.new_tensor([self.time_max]) + @property def alpha_t_max(self): """Return alpha_t at t_max.""" - t_max = torch.tensor([self.time_max], device=alpha.device) - return self.alpha(t_max) + return self.alpha(self._time_max_like()) + + def alpha_t_max_like(self, ref: torch.Tensor) -> torch.Tensor: + """Return alpha_t at t_max on the same device and dtype as `ref`.""" + return self.alpha(self._time_max_like(ref)) @property def sigma_t_max(self): """Return sigma_t at t_max.""" - t_max = torch.tensor([self.time_max], device=alpha.device) - return self.sigma(t_max) + return self.sigma(self._time_max_like()) + + def sigma_t_max_like(self, ref: torch.Tensor) -> torch.Tensor: + """Return sigma_t at t_max on the same device and dtype as `ref`.""" + return self.sigma(self._time_max_like(ref)) @abstractmethod def f(self, time: torch.Tensor) -> torch.Tensor: @@ -147,7 +161,7 @@ def alpha_bar_from_alpha(self, alpha: torch.Tensor) -> (torch.Tensor, torch.Tens Returns: Tensors the same size as alpha, representing alpha_bar and alpha_t_max. """ - alpha_t_max = self.alpha(torch.tensor([self.time_max], device=alpha.device)) + alpha_t_max = self.alpha_t_max_like(alpha) alpha_bar = alpha / (alpha_t_max + self.eps) return alpha_bar, alpha_t_max @@ -189,7 +203,7 @@ def sigma_bar_from_sigma(self, sigma: torch.Tensor) -> (torch.Tensor, torch.Tens Returns: Tensors the same size as sigma, representing sigma_bar and sigma_t_max. """ - sigma_t_max = self.sigma(torch.tensor([self.time_max], device=sigma.device)) + sigma_t_max = self.sigma_t_max_like(sigma) sigma_bar_sq = sigma_t_max**2 - sigma**2 return torch.sqrt(sigma_bar_sq + self.eps), sigma_t_max diff --git a/tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py b/tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py index 8148115cbf76..c5826e752d0b 100644 --- a/tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py +++ b/tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py @@ -19,6 +19,40 @@ from nemo.collections.audio.parts.submodules.schroedinger_bridge import SBNoiseScheduleVE, SBNoiseScheduleVP, SBSampler NUM_STEPS = [1, 5, 10, 20, 100] +DEVICES = [ + pytest.param(torch.device("cpu"), id="cpu"), + pytest.param( + torch.device("cuda"), + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), + id="cuda", + ), +] + + +@pytest.mark.unit +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64], ids=["float32", "float64"]) +@pytest.mark.parametrize("device", DEVICES) +def test_sb_noise_schedule_t_max_accessors_match_reference_tensor(device, dtype): + noise_schedule = SBNoiseScheduleVE(k=2.0, c=0.5, num_steps=5).to(device=device, dtype=dtype) + ref = torch.ones(1, device=device, dtype=dtype) + + alpha_t_max = noise_schedule.alpha_t_max + sigma_t_max = noise_schedule.sigma_t_max + alpha_t_max_like = noise_schedule.alpha_t_max_like(ref) + sigma_t_max_like = noise_schedule.sigma_t_max_like(ref) + + expected_time_max = ref.new_tensor([noise_schedule.time_max]) + expected_alpha_t_max = noise_schedule.alpha(expected_time_max) + expected_sigma_t_max = noise_schedule.sigma(expected_time_max) + + for value in (alpha_t_max, sigma_t_max, alpha_t_max_like, sigma_t_max_like): + assert value.device == device + assert value.dtype == dtype + + torch.testing.assert_close(alpha_t_max, expected_alpha_t_max) + torch.testing.assert_close(sigma_t_max, expected_sigma_t_max) + torch.testing.assert_close(alpha_t_max_like, expected_alpha_t_max) + torch.testing.assert_close(sigma_t_max_like, expected_sigma_t_max) @pytest.mark.parametrize("num_steps", NUM_STEPS)