fix serializable

Author: arrjon

bayesflow/experimental/diffusion_model.py

@@ -21,7 +21,7 @@
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 
-@serializable
+@serializable("bayesflow.experimental")
 class DiffusionModel(InferenceNetwork):
     """Diffusion Model as described in this overview paper [1].
 
@@ -49,7 +49,7 @@ def __init__(
         *,
         subnet: str | type = "mlp",
         integrate_kwargs: dict[str, any] = None,
-        noise_schedule: Literal["edm", "cosine"] | type = "edm",
+        noise_schedule: Literal["edm", "cosine"] | dict | type = "edm",
         prediction_type: Literal["velocity", "noise", "F"] = "F",
         **kwargs,
     ):
@@ -69,8 +69,10 @@ def __init__(
             callable network. Default is "mlp".
         integrate_kwargs : dict[str, any], optional
             Additional keyword arguments for the integration process. Default is None.
-        noise_schedule : Literal['edm', 'cosine'] or type, optional
+        noise_schedule : Literal['edm', 'cosine'], dict or type, optional
             The noise schedule used for the diffusion process. Can be "cosine" or "edm" or a custom noise schedule.
+            You can also pass a dictionary with the configuration for the noise schedule, e.g.,
+                {'type': cosine, 's_shift_cosine': 1.0}
             Default is "edm".
         prediction_type: Literal['velocity', 'noise', 'F'], optional
             The type of prediction used in the diffusion model. Can be "velocity", "noise" or "F" (EDM).

bayesflow/experimental/noise_schedules.py

@@ -8,7 +8,7 @@
 from bayesflow.utils.serialization import deserialize, serializable
 
 
-@serializable
+@serializable("bayesflow.experimental")
 class NoiseSchedule(ABC):
     r"""Noise schedule for diffusion models. We follow the notation from [1].
 
@@ -155,7 +155,7 @@ def validate(self):
             raise ValueError("dt/t log_snr(1) must be finite.")
 
 
-@serializable
+@serializable("bayesflow.experimental")
 class CosineNoiseSchedule(NoiseSchedule):
     """Cosine noise schedule for diffusion models. This schedule is based on the cosine schedule from [1].
     For images, use s_shift_cosine = log(base_resolution / d), where d is the used resolution of the image.
@@ -209,7 +209,7 @@ def from_config(cls, config, custom_objects=None):
         return cls(**deserialize(config, custom_objects=custom_objects))
 
 
-@serializable
+@serializable("bayesflow.experimental")
 class EDMNoiseSchedule(NoiseSchedule):
     """EDM noise schedule for diffusion models. This schedule is based on the EDM paper [1].
     This should be used with the F-prediction type in the diffusion model.