From eb1f6cef48f65cd0b45341b425cb87f363e9060f Mon Sep 17 00:00:00 2001 From: Valentin Pratz Date: Tue, 15 Jul 2025 10:44:15 +0000 Subject: [PATCH] move DiffusionModel from experimental to networks Stabilizes the DiffusionModel class. A deprecation warning for the DiffusionModel class in the experimental module was added. --- bayesflow/experimental/__init__.py | 2 +- .../experimental/diffusion_model/__init__.py | 18 +++++++++++------- bayesflow/networks/__init__.py | 3 ++- bayesflow/networks/diffusion_model/__init__.py | 9 +++++++++ .../diffusion_model/diffusion_model.py | 2 +- .../diffusion_model/dispatch.py | 0 .../diffusion_model/schedules/__init__.py | 0 .../schedules/cosine_noise_schedule.py | 0 .../schedules/edm_noise_schedule.py | 0 .../schedules/noise_schedule.py | 0 examples/Likelihood_Estimation.ipynb | 2 +- tests/test_networks/conftest.py | 12 ++++++------ .../test_diffusion_model/conftest.py | 4 ++-- tests/test_networks/test_inference_networks.py | 2 +- tests/test_utils/test_dispatch.py | 8 ++++---- 15 files changed, 38 insertions(+), 24 deletions(-) create mode 100644 bayesflow/networks/diffusion_model/__init__.py rename bayesflow/{experimental => networks}/diffusion_model/diffusion_model.py (99%) rename bayesflow/{experimental => networks}/diffusion_model/dispatch.py (100%) rename bayesflow/{experimental => networks}/diffusion_model/schedules/__init__.py (100%) rename bayesflow/{experimental => networks}/diffusion_model/schedules/cosine_noise_schedule.py (100%) rename bayesflow/{experimental => networks}/diffusion_model/schedules/edm_noise_schedule.py (100%) rename bayesflow/{experimental => networks}/diffusion_model/schedules/noise_schedule.py (100%) diff --git a/bayesflow/experimental/__init__.py b/bayesflow/experimental/__init__.py index 6c0b6828f..25126f22a 100644 --- a/bayesflow/experimental/__init__.py +++ b/bayesflow/experimental/__init__.py @@ -9,4 +9,4 @@ from ..utils._docs import _add_imports_to_all -_add_imports_to_all(include_modules=["diffusion_model"]) +_add_imports_to_all() diff --git a/bayesflow/experimental/diffusion_model/__init__.py b/bayesflow/experimental/diffusion_model/__init__.py index 6e27f5273..7d8e693c4 100644 --- a/bayesflow/experimental/diffusion_model/__init__.py +++ b/bayesflow/experimental/diffusion_model/__init__.py @@ -1,9 +1,13 @@ -from .diffusion_model import DiffusionModel -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 bayesflow.networks import DiffusionModel as StabilizedDiffusionModel -from ...utils._docs import _add_imports_to_all -_add_imports_to_all(include_modules=[]) +def DiffusionModel(*args, **kwargs): + from warnings import warn + + warn( + "DiffusionModel has been stabilized and moved to bayesflow.networks. " + "Please switch your imports to the new location. This reference will be " + "removed in a future version.", + FutureWarning, + ) + return StabilizedDiffusionModel(*args, **kwargs) diff --git a/bayesflow/networks/__init__.py b/bayesflow/networks/__init__.py index 0632b1b66..f71d4b536 100644 --- a/bayesflow/networks/__init__.py +++ b/bayesflow/networks/__init__.py @@ -7,6 +7,7 @@ from .consistency_models import ConsistencyModel from .coupling_flow import CouplingFlow from .deep_set import DeepSet +from .diffusion_model import DiffusionModel from .flow_matching import FlowMatching from .inference_network import InferenceNetwork from .point_inference_network import PointInferenceNetwork @@ -19,4 +20,4 @@ from ..utils._docs import _add_imports_to_all -_add_imports_to_all(include_modules=[]) +_add_imports_to_all(include_modules=["diffusion_model"]) diff --git a/bayesflow/networks/diffusion_model/__init__.py b/bayesflow/networks/diffusion_model/__init__.py new file mode 100644 index 000000000..341c84c62 --- /dev/null +++ b/bayesflow/networks/diffusion_model/__init__.py @@ -0,0 +1,9 @@ +from .diffusion_model import DiffusionModel +from .schedules import CosineNoiseSchedule +from .schedules import EDMNoiseSchedule +from .schedules import NoiseSchedule +from .dispatch import find_noise_schedule + +from ...utils._docs import _add_imports_to_all + +_add_imports_to_all(include_modules=[]) diff --git a/bayesflow/experimental/diffusion_model/diffusion_model.py b/bayesflow/networks/diffusion_model/diffusion_model.py similarity index 99% rename from bayesflow/experimental/diffusion_model/diffusion_model.py rename to bayesflow/networks/diffusion_model/diffusion_model.py index 404ad2f18..91a05fbff 100644 --- a/bayesflow/experimental/diffusion_model/diffusion_model.py +++ b/bayesflow/networks/diffusion_model/diffusion_model.py @@ -4,7 +4,7 @@ import keras from keras import ops -from bayesflow.networks import InferenceNetwork +from ..inference_network import InferenceNetwork from bayesflow.types import Tensor, Shape from bayesflow.utils import ( expand_right_as, diff --git a/bayesflow/experimental/diffusion_model/dispatch.py b/bayesflow/networks/diffusion_model/dispatch.py similarity index 100% rename from bayesflow/experimental/diffusion_model/dispatch.py rename to bayesflow/networks/diffusion_model/dispatch.py diff --git a/bayesflow/experimental/diffusion_model/schedules/__init__.py b/bayesflow/networks/diffusion_model/schedules/__init__.py similarity index 100% rename from bayesflow/experimental/diffusion_model/schedules/__init__.py rename to bayesflow/networks/diffusion_model/schedules/__init__.py diff --git a/bayesflow/experimental/diffusion_model/schedules/cosine_noise_schedule.py b/bayesflow/networks/diffusion_model/schedules/cosine_noise_schedule.py similarity index 100% rename from bayesflow/experimental/diffusion_model/schedules/cosine_noise_schedule.py rename to bayesflow/networks/diffusion_model/schedules/cosine_noise_schedule.py diff --git a/bayesflow/experimental/diffusion_model/schedules/edm_noise_schedule.py b/bayesflow/networks/diffusion_model/schedules/edm_noise_schedule.py similarity index 100% rename from bayesflow/experimental/diffusion_model/schedules/edm_noise_schedule.py rename to bayesflow/networks/diffusion_model/schedules/edm_noise_schedule.py diff --git a/bayesflow/experimental/diffusion_model/schedules/noise_schedule.py b/bayesflow/networks/diffusion_model/schedules/noise_schedule.py similarity index 100% rename from bayesflow/experimental/diffusion_model/schedules/noise_schedule.py rename to bayesflow/networks/diffusion_model/schedules/noise_schedule.py diff --git a/examples/Likelihood_Estimation.ipynb b/examples/Likelihood_Estimation.ipynb index e73184cee..d1e8c44c3 100644 --- a/examples/Likelihood_Estimation.ipynb +++ b/examples/Likelihood_Estimation.ipynb @@ -200,7 +200,7 @@ "source": [ "workflow = bf.BasicWorkflow(\n", " simulator=simulator,\n", - " inference_network=bf.experimental.DiffusionModel(),\n", + " inference_network=bf.networks.DiffusionModel(),\n", " inference_variables=\"x\",\n", " inference_conditions=\"theta\",\n", " initial_learning_rate=1e-3,\n", diff --git a/tests/test_networks/conftest.py b/tests/test_networks/conftest.py index a122756d0..cb3f33db0 100644 --- a/tests/test_networks/conftest.py +++ b/tests/test_networks/conftest.py @@ -5,7 +5,7 @@ @pytest.fixture() def diffusion_model_edm_F(): - from bayesflow.experimental import DiffusionModel + from bayesflow.networks import DiffusionModel return DiffusionModel( subnet=MLP([8, 8]), @@ -17,7 +17,7 @@ def diffusion_model_edm_F(): @pytest.fixture() def diffusion_model_edm_velocity(): - from bayesflow.experimental import DiffusionModel + from bayesflow.networks import DiffusionModel return DiffusionModel( subnet=MLP([8, 8]), @@ -29,7 +29,7 @@ def diffusion_model_edm_velocity(): @pytest.fixture() def diffusion_model_edm_noise(): - from bayesflow.experimental import DiffusionModel + from bayesflow.networks import DiffusionModel return DiffusionModel( subnet=MLP([8, 8]), @@ -41,7 +41,7 @@ def diffusion_model_edm_noise(): @pytest.fixture() def diffusion_model_cosine_F(): - from bayesflow.experimental import DiffusionModel + from bayesflow.networks import DiffusionModel return DiffusionModel( subnet=MLP([8, 8]), @@ -53,7 +53,7 @@ def diffusion_model_cosine_F(): @pytest.fixture() def diffusion_model_cosine_velocity(): - from bayesflow.experimental import DiffusionModel + from bayesflow.networks import DiffusionModel return DiffusionModel( subnet=MLP([8, 8]), @@ -65,7 +65,7 @@ def diffusion_model_cosine_velocity(): @pytest.fixture() def diffusion_model_cosine_noise(): - from bayesflow.experimental import DiffusionModel + from bayesflow.networks import DiffusionModel return DiffusionModel( subnet=MLP([8, 8]), diff --git a/tests/test_networks/test_diffusion_model/conftest.py b/tests/test_networks/test_diffusion_model/conftest.py index 916a8e3b8..b1ee915ae 100644 --- a/tests/test_networks/test_diffusion_model/conftest.py +++ b/tests/test_networks/test_diffusion_model/conftest.py @@ -3,14 +3,14 @@ @pytest.fixture() def cosine_noise_schedule(): - from bayesflow.experimental.diffusion_model.schedules import CosineNoiseSchedule + from bayesflow.networks.diffusion_model.schedules import CosineNoiseSchedule return CosineNoiseSchedule(min_log_snr=-12, max_log_snr=12, shift=0.1, weighting="likelihood_weighting") @pytest.fixture() def edm_noise_schedule(): - from bayesflow.experimental.diffusion_model.schedules import EDMNoiseSchedule + from bayesflow.networks.diffusion_model.schedules import EDMNoiseSchedule return EDMNoiseSchedule(sigma_data=10.0, sigma_min=1e-5, sigma_max=85.0) diff --git a/tests/test_networks/test_inference_networks.py b/tests/test_networks/test_inference_networks.py index cc00b95a1..a16743f69 100644 --- a/tests/test_networks/test_inference_networks.py +++ b/tests/test_networks/test_inference_networks.py @@ -88,7 +88,7 @@ def test_cycle_consistency(generative_inference_network, random_samples, random_ # cycle-consistency means the forward and inverse methods are inverses of each other import bayesflow as bf - if isinstance(generative_inference_network, bf.experimental.DiffusionModel): + if isinstance(generative_inference_network, bf.networks.DiffusionModel): pytest.skip(reason="test unstable for untrained diffusion models") try: forward_output, forward_log_density = generative_inference_network( diff --git a/tests/test_utils/test_dispatch.py b/tests/test_utils/test_dispatch.py index a0742d551..788b3c13a 100644 --- a/tests/test_utils/test_dispatch.py +++ b/tests/test_utils/test_dispatch.py @@ -2,7 +2,7 @@ import pytest from bayesflow.utils import find_inference_network, find_distribution, find_network, find_summary_network -from bayesflow.experimental.diffusion_model import find_noise_schedule +from bayesflow.networks.diffusion_model import find_noise_schedule # --- Tests for find__network.py --- @@ -247,7 +247,7 @@ def test_find_summary_network_invalid_type(): def test_find_noise_schedule_by_name(): - from bayesflow.experimental.diffusion_model.schedules import CosineNoiseSchedule, EDMNoiseSchedule + from bayesflow.networks.diffusion_model.schedules import CosineNoiseSchedule, EDMNoiseSchedule schedule = find_noise_schedule("cosine") assert isinstance(schedule, CosineNoiseSchedule) @@ -262,7 +262,7 @@ def test_find_noise_schedule_unknown_name(): def test_pass_noise_schedule(): - from bayesflow.experimental.diffusion_model.schedules.noise_schedule import NoiseSchedule + from bayesflow.networks.diffusion_model.schedules.noise_schedule import NoiseSchedule class CustomNoiseSchedule(NoiseSchedule): def __init__(self): @@ -282,7 +282,7 @@ def derivative_log_snr(self, log_snr_t, training): def test_pass_noise_schedule_type(): - from bayesflow.experimental.diffusion_model.schedules import EDMNoiseSchedule + from bayesflow.networks.diffusion_model.schedules import EDMNoiseSchedule schedule = find_noise_schedule(EDMNoiseSchedule, sigma_data=10.0) assert isinstance(schedule, EDMNoiseSchedule)