minor cleanup of refactory

Author: vpratz

bayesflow/experimental/diffusion_model/__init__.py

@@ -1,5 +1,7 @@
 from .diffusion_model import DiffusionModel
-from bayesflow.experimental.diffusion_model.schedules.cosine_noise_schedule import CosineNoiseSchedule
+from bayesflow.experimental.diffusion_model.schedules import CosineNoiseSchedule
+from bayesflow.experimental.diffusion_model.schedules import EDMNoiseSchedule
+from bayesflow.experimental.diffusion_model.schedules import NoiseSchedule
 from .dispatch import find_noise_schedule
 
 from ...utils._docs import _add_imports_to_all

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -105,7 +105,7 @@ def __init__(
         self._prediction_type = prediction_type
         self._loss_type = loss_type
 
-        self.schedule_kwargs = schedule_kwargs or {}
+        schedule_kwargs = schedule_kwargs or {}
         self.noise_schedule = find_noise_schedule(noise_schedule, **self.schedule_kwargs)
         self.noise_schedule.validate()
 
@@ -148,7 +148,6 @@ def get_config(self):
             "noise_schedule": self.noise_schedule,
             "prediction_type": self._prediction_type,
             "loss_type": self._loss_type,
-            "schedule_kwargs": self.schedule_kwargs,
             "integrate_kwargs": self.integrate_kwargs,
         }
         return base_config | serialize(config)
@@ -194,8 +193,9 @@ def convert_prediction_to_x(
             return x1 * z + x2 * pred
         elif self._prediction_type == "x":
             return pred
-        else:
+        elif self._prediction_type == "score":
             return (z + sigma_t**2 * pred) / alpha_t
+        raise ValueError(f"Unknown prediction type {self._prediction_type}.")
 
     def velocity(
         self,
@@ -320,12 +320,9 @@ def _forward(
         training: bool = False,
         **kwargs,
     ) -> Tensor | tuple[Tensor, Tensor]:
-        integrate_kwargs = {
-            **self.integrate_kwargs,
-            "start_time": kwargs.pop("start_time", 0.0),
-            "stop_time": kwargs.pop("stop_time", 1.0),
-            **kwargs,
-        }
+        integrate_kwargs = {"start_time": 0.0, "stop_time": 1.0}
+        integrate_kwargs = integrate_kwargs | self.integrate_kwargs
+        integrate_kwargs = integrate_kwargs | kwargs
 
         if integrate_kwargs["method"] == "euler_maruyama":
             raise ValueError("Stochastic methods are not supported for forward integration.")
@@ -373,12 +370,9 @@ def _inverse(
         training: bool = False,
         **kwargs,
     ) -> Tensor | tuple[Tensor, Tensor]:
-        integrate_kwargs = {
-            **self.integrate_kwargs,
-            "start_time": kwargs.pop("start_time", 1.0),
-            "stop_time": kwargs.pop("stop_time", 0.0),
-            **kwargs,
-        }
+        integrate_kwargs = {"start_time": 1.0, "stop_time": 0.0}
+        integrate_kwargs = integrate_kwargs | self.integrate_kwargs
+        integrate_kwargs = integrate_kwargs | kwargs
         if density:
             if integrate_kwargs["method"] == "euler_maruyama":
                 raise ValueError("Stochastic methods are not supported for density computation.")

bayesflow/experimental/diffusion_model/schedules/edm_noise_schedule.py

@@ -101,7 +101,7 @@ def derivative_log_snr(self, log_snr_t: Tensor, training: bool = False) -> Tenso
 
     def get_weights_for_snr(self, log_snr_t: Tensor) -> Tensor:
         """Get weights for the signal-to-noise ratio (snr) for a given log signal-to-noise ratio (lambda)."""
-        # for F-prediction: w = (ops.exp(-log_snr_t) + sigma_data^2) / (ops.exp(-log_snr_t)*sigma_data^2)
+        # for F-loss: w = (ops.exp(-log_snr_t) + sigma_data^2) / (ops.exp(-log_snr_t)*sigma_data^2)
         return 1 + ops.exp(-log_snr_t) / ops.square(self.sigma_data)
 
     def get_config(self):