Refactored diffusion sampler

physicsnemo/diffusion/__init__.py

@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base import DiffusionModel  # noqa: F401
+from .base import DiffusionDenoiser, DiffusionModel  # noqa: F401

physicsnemo/diffusion/base.py

@@ -98,3 +98,100 @@ def __call__(
             Model output with the same shape as ``x``.
         """
         ...
+
+
+@runtime_checkable
+class DiffusionDenoiser(Protocol):
+    r"""
+    Protocol defining a denoiser interface for diffusion model inference.
+
+    This is the minimal interface required for sampling from a diffusion model,
+    and any object that implements this interface can be used as a denoiser in
+    sampling utilities.
+
+    A denoiser is a callable that takes a noisy state ``x`` and a noise level
+    (or diffusion time) ``t``, and typically returns a denoising term. A
+    denoising could be the right hand side for
+    ODE/SDE-based sampling, directly the denoised latent state for discrete
+    Markov chain-based sampling, etc. This interface is generic and it does not
+    make any assumption about the nature of the denoising term. It is expected
+    to be used in conjunction with a compatible
+    :class:`~physicsnemo.diffusion.samplers.solvers.Solver` to perform the
+    actual sampling.
+
+    This protocol is used during inference. For training,
+    which often requires additional inputs like conditioning, use the more
+    general :class:`DiffusionModel` protocol instead.
+
+    A :class:`DiffusionDenoiser` can be obtained from a
+    :class:`DiffusionModel` by partially applying the ``condition`` and
+    any other keyword arguments. For example:
+
+    .. code-block:: python
+
+        from functools import partial
+        from tensordict import TensorDict
+
+        class MyDiffusionModel:
+            def __call__(self, x, t, condition):
+                # Model forward pass using x, t, and condition
+                return x * 0.9 + condition["y"]
+
+        model = MyDiffusionModel()
+        my_condition = TensorDict({"y": torch.randn(batch_size, 10)})
+
+        # Create denoiser by partially applying the condition
+        denoiser = partial(model, condition=my_condition)
+
+        # Now denoiser(x, t) implements the DiffusionDenoiser interface.
+
+    See Also
+    --------
+    :func:`~physicsnemo.diffusion.samplers.sample` : The sampling function
+        that uses this denoiser interface.
+    :class:`DiffusionModel` : The full diffusion model interface with
+        conditioning support.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from physicsnemo.diffusion import DiffusionDenoiser
+    >>>
+    >>> class SimpleDenoiser:
+    ...     def __call__(self, x, t):
+    ...         # A trivial denoiser that returns the input unchanged
+    ...         return x
+    ...
+    >>> denoiser = SimpleDenoiser()
+    >>> isinstance(denoiser, DiffusionDenoiser)
+    True
+    """
+
+    def __call__(
+        self,
+        x: Float[torch.Tensor, "B *dims"],  # noqa: F821
+        t: Float[torch.Tensor, "B "],  # noqa: F821
+    ) -> Float[torch.Tensor, "B *dims"]:  # noqa: F821
+        r"""
+        Function to produce a denoising output at the given noise level.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Noisy latent state of shape :math:`(B, *)` where :math:`B` is the
+            batch size and :math:`*` denotes any number of additional
+            dimensions (e.g., channels and spatial dimensions).
+        t : torch.Tensor
+            Batched diffusion time tensor of shape :math:`(B,)`.
+            All batch elements in the latent state ``x`` typically share the
+            same diffusion time values, but ``t`` is still required to be a
+            batched tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Denoising output. Can be the right hand side for ODE/SDE-based
+            sampling, the denoised latent state for discrete Markov chain-based
+            sampling, etc.
+        """
+        ...

physicsnemo/diffusion/noise_schedulers/__init__.py

@@ -14,13 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import warnings
-
-from physicsnemo.core.warnings import FutureFeatureWarning
-
-warnings.warn(
-    "The 'physicsnemo.diffusion.noise_schedulers' module is a placeholder for "
-    "future functionality that will be implemented in an upcoming release.",
-    FutureFeatureWarning,
-    stacklevel=2,
+from .noise_schedulers import (  # noqa: F401
+    EDMNoiseScheduler,
+    IDDPMNoiseScheduler,
+    LinearGaussianNoiseScheduler,
+    NoiseScheduler,
+    StudentTEDMNoiseScheduler,
+    VENoiseScheduler,
+    VPNoiseScheduler,
 )