restructure noise schedules, test find_noise_schedule

Author: vpratz

bayesflow/experimental/diffusion_model/__init__.py

@@ -1,5 +1,8 @@
 from .diffusion_model import DiffusionModel
-from .noise_schedules import CosineNoiseSchedule, EDMNoiseSchedule, NoiseSchedule
+from .noise_schedule import NoiseSchedule
+from .cosine_noise_schedule import CosineNoiseSchedule
+from .edm_noise_schedule import EDMNoiseSchedule
+from .dispatch import find_noise_schedule
 
 from ...utils._docs import _add_imports_to_all
 

bayesflow/experimental/diffusion_model/cosine_noise_schedule.py

@@ -0,0 +1,82 @@
+import math
+from typing import Union, Literal
+
+from keras import ops
+
+from bayesflow.types import Tensor
+from bayesflow.utils.serialization import deserialize, serializable
+
+from .noise_schedule import NoiseSchedule
+
+
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
+class CosineNoiseSchedule(NoiseSchedule):
+    """Cosine noise schedule for diffusion models. This schedule is based on the cosine schedule from [1].
+
+    [1] Diffusion Models Beat GANs on Image Synthesis: Dhariwal and Nichol (2022)
+    """
+
+    def __init__(
+        self,
+        min_log_snr: float = -15,
+        max_log_snr: float = 15,
+        shift: float = 0.0,
+        weighting: Literal["sigmoid", "likelihood_weighting"] = "sigmoid",
+    ):
+        """
+        Initialize the cosine noise schedule.
+
+        Parameters
+        ----------
+        min_log_snr : float, optional
+            The minimum log signal-to-noise ratio (lambda). Default is -15.
+        max_log_snr : float, optional
+            The maximum log signal-to-noise ratio (lambda). Default is 15.
+        shift : float, optional
+            Shift the log signal-to-noise ratio (lambda) by this amount. Default is 0.0.
+            For images, use shift = log(base_resolution / d), where d is the used resolution of the image.
+        weighting : Literal["sigmoid", "likelihood_weighting"], optional
+            The type of weighting function to use for the noise schedule. Default is "sigmoid".
+        """
+        super().__init__(name="cosine_noise_schedule", variance_type="preserving", weighting=weighting)
+        self._shift = shift
+        self.log_snr_min = min_log_snr
+        self.log_snr_max = max_log_snr
+
+        self._t_min = self.get_t_from_log_snr(log_snr_t=self.log_snr_max, training=True)
+        self._t_max = self.get_t_from_log_snr(log_snr_t=self.log_snr_min, training=True)
+
+    def _truncated_t(self, t: Tensor) -> Tensor:
+        return self._t_min + (self._t_max - self._t_min) * t
+
+    def get_log_snr(self, t: Union[float, Tensor], training: bool) -> Tensor:
+        """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
+        t_trunc = self._truncated_t(t)
+        return -2 * ops.log(ops.tan(math.pi * t_trunc * 0.5)) + 2 * self._shift
+
+    def get_t_from_log_snr(self, log_snr_t: Union[Tensor, float], training: bool) -> Tensor:
+        """Get the diffusion time (t) from the log signal-to-noise ratio (lambda)."""
+        # SNR = -2 * log(tan(pi*t/2)) => t = 2/pi * arctan(exp(-snr/2))
+        return 2 / math.pi * ops.arctan(ops.exp((2 * self._shift - log_snr_t) * 0.5))
+
+    def derivative_log_snr(self, log_snr_t: Tensor, training: bool) -> Tensor:
+        """Compute d/dt log(1 + e^(-snr(t))), which is used for the reverse SDE."""
+        t = self.get_t_from_log_snr(log_snr_t=log_snr_t, training=training)
+
+        # Compute the truncated time t_trunc
+        t_trunc = self._truncated_t(t)
+        dsnr_dx = -(2 * math.pi) / ops.sin(math.pi * t_trunc)
+
+        # Using the chain rule on f(t) = log(1 + e^(-snr(t))):
+        # f'(t) = - (e^{-snr(t)} / (1 + e^{-snr(t)})) * dsnr_dt
+        dsnr_dt = dsnr_dx * (self._t_max - self._t_min)
+        factor = ops.exp(-log_snr_t) / (1 + ops.exp(-log_snr_t))
+        return -factor * dsnr_dt
+
+    def get_config(self):
+        return dict(min_log_snr=self.log_snr_min, max_log_snr=self.log_snr_max, shift=self._shift)
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))

bayesflow/experimental/diffusion_model/dispatch.py

@@ -1,10 +1,10 @@
 from functools import singledispatch
-from .noise_schedules import NoiseSchedule
+from .noise_schedule import NoiseSchedule
 
 
 @singledispatch
 def find_noise_schedule(arg, *args, **kwargs):
-    raise TypeError(f"Unknown noise schedule: {arg!r}")
+    raise TypeError(f"Not a noise schedule: {arg!r}. Please pass an object of type 'NoiseSchedule'.")
 
 
 @find_noise_schedule.register
@@ -16,11 +16,11 @@ def _(noise_schedule: NoiseSchedule):
 def _(name: str, *args, **kwargs):
     match name.lower():
         case "cosine":
-            from .noise_schedules import CosineNoiseSchedule
+            from .cosine_noise_schedule import CosineNoiseSchedule
 
             return CosineNoiseSchedule()
         case "edm":
-            from .noise_schedules import EDMNoiseSchedule
+            from .edm_noise_schedule import EDMNoiseSchedule
 
             return EDMNoiseSchedule()
         case other:
@@ -33,11 +33,11 @@ def _(config: dict, *args, **kwargs):
     params = {k: v for k, v in config.items() if k != "name"}
     match name:
         case "cosine":
-            from .noise_schedules import CosineNoiseSchedule
+            from .cosine_noise_schedule import CosineNoiseSchedule
 
             return CosineNoiseSchedule(**params)
         case "edm":
-            from .noise_schedules import EDMNoiseSchedule
+            from .edm_noise_schedule import EDMNoiseSchedule
 
             return EDMNoiseSchedule(**params)
         case other:
@@ -49,8 +49,3 @@ def _(cls: type, *args, **kwargs):
     if issubclass(cls, NoiseSchedule):
         return cls(*args, **kwargs)
     raise TypeError(f"Expected subclass of NoiseSchedule, got {cls}")
-
-
-@find_noise_schedule.register
-def _(schedule: type, *args, **kwargs):
-    return schedule

bayesflow/experimental/diffusion_model/noise_schedules.py renamed to bayesflow/experimental/diffusion_model/edm_noise_schedule.py

@@ -1,84 +1,12 @@
 import math
-from typing import Union, Literal
+from typing import Union
 
 from keras import ops
 
 from bayesflow.types import Tensor
 from bayesflow.utils.serialization import deserialize, serializable
 
-from .noise_schedule_base import NoiseSchedule
-
-
-@serializable("bayesflow.experimental")
-class CosineNoiseSchedule(NoiseSchedule):
-    """Cosine noise schedule for diffusion models. This schedule is based on the cosine schedule from [1].
-
-    [1] Diffusion Models Beat GANs on Image Synthesis: Dhariwal and Nichol (2022)
-    """
-
-    def __init__(
-        self,
-        min_log_snr: float = -15,
-        max_log_snr: float = 15,
-        shift: float = 0.0,
-        weighting: Literal["sigmoid", "likelihood_weighting"] = "sigmoid",
-    ):
-        """
-        Initialize the cosine noise schedule.
-
-        Parameters
-        ----------
-        min_log_snr : float, optional
-            The minimum log signal-to-noise ratio (lambda). Default is -15.
-        max_log_snr : float, optional
-            The maximum log signal-to-noise ratio (lambda). Default is 15.
-        shift : float, optional
-            Shift the log signal-to-noise ratio (lambda) by this amount. Default is 0.0.
-            For images, use shift = log(base_resolution / d), where d is the used resolution of the image.
-        weighting : Literal["sigmoid", "likelihood_weighting"], optional
-            The type of weighting function to use for the noise schedule. Default is "sigmoid".
-        """
-        super().__init__(name="cosine_noise_schedule", variance_type="preserving", weighting=weighting)
-        self._shift = shift
-        self.log_snr_min = min_log_snr
-        self.log_snr_max = max_log_snr
-
-        self._t_min = self.get_t_from_log_snr(log_snr_t=self.log_snr_max, training=True)
-        self._t_max = self.get_t_from_log_snr(log_snr_t=self.log_snr_min, training=True)
-
-    def _truncated_t(self, t: Tensor) -> Tensor:
-        return self._t_min + (self._t_max - self._t_min) * t
-
-    def get_log_snr(self, t: Union[float, Tensor], training: bool) -> Tensor:
-        """Get the log signal-to-noise ratio (lambda) for a given diffusion time."""
-        t_trunc = self._truncated_t(t)
-        return -2 * ops.log(ops.tan(math.pi * t_trunc * 0.5)) + 2 * self._shift
-
-    def get_t_from_log_snr(self, log_snr_t: Union[Tensor, float], training: bool) -> Tensor:
-        """Get the diffusion time (t) from the log signal-to-noise ratio (lambda)."""
-        # SNR = -2 * log(tan(pi*t/2)) => t = 2/pi * arctan(exp(-snr/2))
-        return 2 / math.pi * ops.arctan(ops.exp((2 * self._shift - log_snr_t) * 0.5))
-
-    def derivative_log_snr(self, log_snr_t: Tensor, training: bool) -> Tensor:
-        """Compute d/dt log(1 + e^(-snr(t))), which is used for the reverse SDE."""
-        t = self.get_t_from_log_snr(log_snr_t=log_snr_t, training=training)
-
-        # Compute the truncated time t_trunc
-        t_trunc = self._truncated_t(t)
-        dsnr_dx = -(2 * math.pi) / ops.sin(math.pi * t_trunc)
-
-        # Using the chain rule on f(t) = log(1 + e^(-snr(t))):
-        # f'(t) = - (e^{-snr(t)} / (1 + e^{-snr(t)})) * dsnr_dt
-        dsnr_dt = dsnr_dx * (self._t_max - self._t_min)
-        factor = ops.exp(-log_snr_t) / (1 + ops.exp(-log_snr_t))
-        return -factor * dsnr_dt
-
-    def get_config(self):
-        return dict(min_log_snr=self.log_snr_min, max_log_snr=self.log_snr_max, shift=self._shift)
-
-    @classmethod
-    def from_config(cls, config, custom_objects=None):
-        return cls(**deserialize(config, custom_objects=custom_objects))
+from .noise_schedule import NoiseSchedule
 
 
 # disable module check, use potential module after moving from experimental

bayesflow/experimental/diffusion_model/noise_schedule_base.py renamed to bayesflow/experimental/diffusion_model/noise_schedule.py

@@ -2,6 +2,7 @@
 import pytest
 
 from bayesflow.utils import find_inference_network, find_distribution, find_summary_network
+from bayesflow.experimental.diffusion_model import find_noise_schedule
 
 
 # --- Tests for find_inference_network.py ---
@@ -168,3 +169,72 @@ def test_find_summary_network_unknown_name():
 def test_find_summary_network_invalid_type():
     with pytest.raises(TypeError):
         find_summary_network(0.1234)
+
+
+def test_find_noise_schedule_by_name():
+    from bayesflow.experimental.diffusion_model import CosineNoiseSchedule, EDMNoiseSchedule
+
+    schedule = find_noise_schedule("cosine")
+    assert isinstance(schedule, CosineNoiseSchedule)
+
+    schedule = find_noise_schedule("edm")
+    assert isinstance(schedule, EDMNoiseSchedule)
+
+
+def test_find_noise_schedule_unknown_name():
+    with pytest.raises(ValueError):
+        find_noise_schedule("unknown_noise_schedule")
+
+
+def test_pass_noise_schedule():
+    from bayesflow.experimental.diffusion_model import NoiseSchedule
+
+    class CustomNoiseSchedule(NoiseSchedule):
+        def __init__(self):
+            pass
+
+        def get_log_snr(self, t, training):
+            pass
+
+        def get_t_from_log_snr(self, log_snr_t, training):
+            pass
+
+        def derivative_log_snr(self, log_snr_t, training):
+            pass
+
+    schedule = CustomNoiseSchedule()
+    assert schedule is find_noise_schedule(schedule)
+
+
+def test_pass_noise_schedule_type():
+    from bayesflow.experimental.diffusion_model import EDMNoiseSchedule
+
+    schedule = find_noise_schedule(EDMNoiseSchedule, sigma_data=10.0)
+    assert isinstance(schedule, EDMNoiseSchedule)
+    assert schedule.sigma_data == 10.0
+
+
+def test_find_noise_schedule_by_dict():
+    from bayesflow.experimental.diffusion_model import CosineNoiseSchedule, EDMNoiseSchedule
+
+    schedule = find_noise_schedule({"name": "cosine"})
+    assert isinstance(schedule, CosineNoiseSchedule)
+
+    schedule = find_noise_schedule({"name": "edm", "sigma_data": 10})
+    assert isinstance(schedule, EDMNoiseSchedule)
+    assert schedule.sigma_data == 10
+
+
+def test_find_noise_schedule_unknown_name_in_dict():
+    with pytest.raises(ValueError):
+        find_noise_schedule({"name": "unknown_noise_schedule"})
+
+
+def test_find_noise_schedule_invalid_class():
+    with pytest.raises(TypeError):
+        find_noise_schedule(int)
+
+
+def test_find_noise_schedule_invalid_object():
+    with pytest.raises(TypeError):
+        find_noise_schedule(1.0)