improve dispatch

Author: arrjon

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -74,7 +74,7 @@ def __init__(
         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}
+                {'name': 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/diffusion_model/dispatch.py

@@ -29,8 +29,8 @@ def _(name: str, *args, **kwargs):
 
 @find_noise_schedule.register
 def _(config: dict, *args, **kwargs):
-    name = config.get("type", "").lower()
-    params = {k: v for k, v in config.items() if k != "type"}
+    name = config.get("name", "").lower()
+    params = {k: v for k, v in config.items() if k != "name"}
     match name:
         case "cosine":
             from .noise_schedules import CosineNoiseSchedule
@@ -46,9 +46,6 @@ def _(config: dict, *args, **kwargs):
 
 @find_noise_schedule.register
 def _(cls: type, *args, **kwargs):
-    # Lazily import NoiseSchedule class and compare
-    from .noise_schedules import NoiseSchedule
-
     if issubclass(cls, NoiseSchedule):
         return cls(*args, **kwargs)
     raise TypeError(f"Expected subclass of NoiseSchedule, got {cls}")