From 73a3fbae3aaeca971315f71263798e613353e1a5 Mon Sep 17 00:00:00 2001 From: Charlelie Laurent Date: Fri, 16 Oct 2026 13:13:46 -0700 Subject: [PATCH 1/2] Add structured diffusion model inputs and outputs Signed-off-by: Charlelie Laurent --- .importlinter | 3 +- CHANGELOG.md | 11 + docs/api/diffusion/introduction.rst | 10 + docs/api/diffusion/models.rst | 62 ++++- physicsnemo/diffusion/__init__.py | 9 +- physicsnemo/diffusion/base.py | 257 ++++++++++++++++-- physicsnemo/diffusion/metrics/losses.py | 32 +++ .../preconditioners/preconditioners.py | 145 +++++++--- physicsnemo/diffusion/utils/model_wrappers.py | 60 +++- test/diffusion/helpers.py | 48 ++++ test/diffusion/test_model_io.py | 208 ++++++++++++++ 11 files changed, 769 insertions(+), 76 deletions(-) create mode 100644 test/diffusion/test_model_io.py diff --git a/.importlinter b/.importlinter index fa39832fae..891ae686a7 100644 --- a/.importlinter +++ b/.importlinter @@ -117,7 +117,8 @@ layers = samplers | metrics multi_diffusion | preconditioners | guidance noise_schedulers - utils | base + utils + base [importlinter:contract:physicsnemo-sym] name = Control Internal Dependencies in PhysicsNeMo sym diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cb26a3be9..9f560ac5bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Adds an experimental, independently implemented SMART operator with sampled geometry encoding, reusable latent contexts, and chunked field decoding. +- Adds `DiffusionExtraTimes` and `DiffusionExtraOutputs` to + `physicsnemo.diffusion` for utilities that carry named times or auxiliary + outputs. Compatible model interfaces and preconditioners can accept + `(t, extra_times)`, where `extra_times` is a `DiffusionExtraTimes` container, + and return `(prediction, extra_outputs)`, where `extra_outputs` is a + `DiffusionExtraOutputs` container. The API reserves the named field + `extra_times.r` for a destination diffusion time, as in MeanFlow. It reserves + the named field `extra_outputs.logvar` for learned loss weighting, as in sCM + and EDM2. Existing single-time, single-prediction calls remain + unchanged. Existing DSM and flow-matching losses continue to consume + single tensors. - Adds standalone FLARE++ attention and model APIs, with input-conditioned dynamic routing, plus a ``GALE_FPP`` backend for using the same mixer inside GeoTransolver. diff --git a/docs/api/diffusion/introduction.rst b/docs/api/diffusion/introduction.rst index 15863eb88f..06993b70f4 100644 --- a/docs/api/diffusion/introduction.rst +++ b/docs/api/diffusion/introduction.rst @@ -290,6 +290,16 @@ API Reference :members: :exclude-members: __init__ +:code:`DiffusionExtraTimes` +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: physicsnemo.diffusion.DiffusionExtraTimes + +:code:`DiffusionExtraOutputs` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: physicsnemo.diffusion.DiffusionExtraOutputs + :code:`Predictor` ~~~~~~~~~~~~~~~~~ diff --git a/docs/api/diffusion/models.rst b/docs/api/diffusion/models.rst index 1ae6ac54d0..eed92f216f 100644 --- a/docs/api/diffusion/models.rst +++ b/docs/api/diffusion/models.rst @@ -80,7 +80,68 @@ expects ``(x, t, condition)``. A short adapter class solves this: def forward(self, x, t, condition=None): return self.net(x, noise_labels=t, class_labels=condition) +Most diffusion and flow-matching models take one time tensor and return one +prediction tensor. +A compatible :class:`~physicsnemo.diffusion.DiffusionModel` can instead +accept ``(t, extra_times)`` with +:class:`~physicsnemo.diffusion.DiffusionExtraTimes` and return +``(prediction, extra_outputs)`` with +:class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. For example, MeanFlow +uses a destination time ``r``, while sCM and EDM2 can use ``logvar`` for loss +weighting. The primary time and prediction remain required tensors. + +The adapter below combines the destination time with the primary time for a +DiT backbone. A small +:class:`~physicsnemo.models.mlp.FullyConnected` head predicts a log-scale +from the primary time. This example does not define a canonical MeanFlow +architecture. A compatible loss decides how to use ``logvar`` during +training. Inference code that needs only a prediction can ignore it. +.. code-block:: python + + import torch + from physicsnemo.core import Module + from physicsnemo.diffusion import DiffusionExtraOutputs, DiffusionExtraTimes + from physicsnemo.models.dit import DiT + from physicsnemo.models.mlp import FullyConnected + + class ExtendedDiTAdapter(Module): + def __init__(self): + super().__init__() + self.backbone = DiT( + input_size=8, in_channels=2, out_channels=2, + patch_size=4, hidden_size=32, depth=1, num_heads=4, + ) + self.r_embed = FullyConnected( + in_features=1, layer_size=8, out_features=1, num_layers=1, + ) + self.logvar_head = FullyConnected( + in_features=1, layer_size=8, out_features=1, num_layers=1, + ) + + def forward(self, x, t, condition=None): + t, extra_times = t + # DiT takes one time tensor, so combine the two times in the adapter. + r_embed = self.r_embed(extra_times.r[:, None]).squeeze(-1) + prediction = self.backbone(x, t + r_embed, condition) + logvar = self.logvar_head(t[:, None]).squeeze(-1) + return prediction, DiffusionExtraOutputs( + logvar=logvar, batch_size=[x.shape[0]], + ) + + model = ExtendedDiTAdapter() + x = torch.randn(2, 2, 8, 8) + t = torch.ones(2) + times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2])) + prediction, extra_outputs = model(x, times) + assert prediction.shape == x.shape + assert extra_outputs.logvar.shape == t.shape + +This adapter illustrates one way to embed the reserved ``r`` time and compute +the reserved ``logvar`` output. Other implementations can use different +architectures. Its ``forward`` method implements +:class:`~physicsnemo.diffusion.DiffusionModel` and works with utilities that +support both extra times and extra outputs. .. _backbone_external_models: @@ -100,4 +161,3 @@ You can bring in models from external libraries, such as: * by using :meth:`~physicsnemo.core.module.Module.from_torch` to convert a ``torch.nn.Module`` subclass into a :class:`~physicsnemo.core.module.Module` with PhysicsNeMo serialization support - diff --git a/physicsnemo/diffusion/__init__.py b/physicsnemo/diffusion/__init__.py index d9602b252f..7b7d2bf629 100644 --- a/physicsnemo/diffusion/__init__.py +++ b/physicsnemo/diffusion/__init__.py @@ -14,4 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .base import Denoiser, DiffusionModel, Predictor, PredictorType # noqa: F401 +from .base import ( # noqa: F401 + Denoiser, + DiffusionExtraOutputs, + DiffusionExtraTimes, + DiffusionModel, + Predictor, + PredictorType, +) diff --git a/physicsnemo/diffusion/base.py b/physicsnemo/diffusion/base.py index f386ba3a32..f2acfe093b 100644 --- a/physicsnemo/diffusion/base.py +++ b/physicsnemo/diffusion/base.py @@ -19,11 +19,166 @@ from typing import Any, Literal, Protocol, runtime_checkable import torch -from jaxtyping import Float -from tensordict import TensorDict +from jaxtyping import Float, Shaped +from tensordict import TensorClass, TensorDict + + +class DiffusionExtraTimes(TensorClass): + r"""Named, batched diffusion times accompanying the primary time tensor. + + This ``tensordict.TensorClass`` container carries named extra times for + diffusion utilities. A :class:`DiffusionModel` is one possible consumer. + + Most diffusion and flow-matching models take a single diffusion time + tensor. Models that need more than one diffusion time during training + or sampling can take ``(t, extra_times)`` instead. The required primary + time ``t`` stays separate from the optional named fields in this + container. + + Every supplied time tensor has shape :math:`(B,)` and uses the same + scheduler coordinates as the primary time. Set ``batch_size=[B]`` to + match the input batch. For example, utilities reserve ``r`` for the + destination diffusion time in MeanFlow. Each field is independently + optional and defaults to ``None``. + + .. note:: + + Subclass this container to add time fields for a custom use. Do not + reuse reserved names for different meanings. The framework does not + infer the meaning of custom fields; the consuming utility must define + and handle them. + + Parameters + ---------- + r : torch.Tensor or None, optional, default=None + Destination diffusion time for MeanFlow, with shape :math:`(B,)`. + Leave unset when the model does not use this field. + batch_size : sequence of int + Set to ``[B]``, matching the primary time tensor. + + Examples + -------- + Pair a primary time with the reserved destination time for a MeanFlow + consumer. The batch sizes of both tensors and the container agree. + This example shows the tuple passed to a compatible utility: + + >>> import torch + >>> from physicsnemo.diffusion import DiffusionExtraTimes + >>> t = torch.ones(2) + >>> # The reserved r field holds the destination time. + >>> times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2])) + >>> t, extra_times = times + >>> t.shape, extra_times.r.shape + (torch.Size([2]), torch.Size([2])) + + A subclass can add independent optional times for a custom consumer. + The consumer must interpret these fields and can accept either field + alone or both together: + + >>> from jaxtyping import Float + >>> class CustomTimes(DiffusionExtraTimes): + ... # The consuming utility defines these custom times. + ... r1: Float[torch.Tensor, " B"] | None = None + ... r2: Float[torch.Tensor, " B"] | None = None + >>> first = CustomTimes(r1=torch.zeros(2), batch_size=[2]) + >>> second = CustomTimes(r2=torch.ones(2), batch_size=[2]) + >>> first.r is None, first.r2 is None, second.r1 is None + (True, True, True) + >>> first.r1.shape, second.r2.shape + (torch.Size([2]), torch.Size([2])) + >>> both = CustomTimes( + ... r1=torch.zeros(2), r2=torch.ones(2), batch_size=[2] + ... ) + >>> both.r1.shape, both.r2.shape + (torch.Size([2]), torch.Size([2])) + """ + + r: Float[torch.Tensor, " B"] | None = None + + +class DiffusionExtraOutputs(TensorClass): + r"""Named, batched auxiliary outputs accompanying a model prediction. + + This ``tensordict.TensorClass`` container carries named extra outputs for + diffusion utilities. A :class:`DiffusionModel` is one possible producer. + + Most diffusion and flow-matching models return a single tensor that + predicts clean data, noise, a score, or a velocity, as described by + :data:`PredictorType`. Models that produce auxiliary outputs for + training or sampling can return ``(prediction, extra_outputs)`` instead. + For example, sCM and EDM2 losses use a learned log-scale alongside the + prediction for adaptive loss weighting. The required prediction stays + separate from the optional named fields in this container. + + All supplied output tensors share the prediction's batch size + :math:`B`, but their remaining dimensions can differ. Set + ``batch_size=[B]`` to describe that shared batch dimension. For example, + the reserved field ``logvar`` provides a learned log-scale for adaptive + loss weighting in sCM and EDM2. It does not represent diffusion + transition variance. Each field is independently optional and defaults + to ``None``. + + .. note:: + + Subclass this container to add outputs for a custom use. Do not reuse + reserved names for different meanings. The framework does not infer + the meaning of custom fields; the consuming utility must define and + handle them. + + Parameters + ---------- + logvar : torch.Tensor or None, optional, default=None + Learned log-scale for adaptive loss weighting in sCM and EDM2, with + shape :math:`(B,)`. Leave unset when the model does not use it. + batch_size : sequence of int + Set to ``[B]``, matching the prediction's batch dimension. + + Examples + -------- + Pair a required prediction with a learned log-scale for a compatible + loss. The prediction retains its own shape while the container shares + its batch dimension: + + >>> import torch + >>> from physicsnemo.diffusion import DiffusionExtraOutputs + >>> prediction = torch.zeros(2, 3, 5) + >>> # The reserved logvar field supplies loss weighting metadata. + >>> output = ( + ... prediction, + ... DiffusionExtraOutputs(logvar=torch.zeros(2), batch_size=[2]), + ... ) + >>> prediction, extra_outputs = output + >>> prediction.shape, extra_outputs.logvar.shape + (torch.Size([2, 3, 5]), torch.Size([2])) + + A subclass can hold independent optional outputs with different feature + shapes. The consumer decides how to use each field, so these names carry + no framework-defined meaning: + + >>> from jaxtyping import Float + >>> class CustomOutputs(DiffusionExtraOutputs): + ... # The consuming utility defines these custom outputs. + ... features: Float[torch.Tensor, " B C H W"] | None = None + ... confidence: Float[torch.Tensor, " B K"] | None = None + >>> extra = CustomOutputs( + ... features=torch.zeros(2, 4, 5, 7), + ... confidence=torch.ones(2, 3), + ... batch_size=[2], + ... ) + >>> extra.features.shape, extra.confidence.shape + (torch.Size([2, 4, 5, 7]), torch.Size([2, 3])) + >>> extra.logvar is None + True + >>> features_only = CustomOutputs(features=extra.features, batch_size=[2]) + >>> features_only.confidence is None + True + """ + + logvar: Float[torch.Tensor, " B"] | None = None + PredictorType = Literal["x0", "score", "epsilon", "flow"] -"""Type of prediction produced by a diffusion model. +"""Type of prediction produced by a diffusion or flow-matching model. One of: @@ -46,8 +201,18 @@ class DiffusionModel(Protocol): into a prediction. This protocol defines the standard interface that all such models must meet. - Any model or function that implements this interface can be used with - preconditioners, losses, samplers, and other diffusion utilities. + Most models take a single diffusion time tensor and return a single + prediction tensor. Specialized models can extend either side of this + interface: + + - Pass ``(t, extra_times)`` with :class:`DiffusionExtraTimes` for models + that accept more than one diffusion time. + - Return ``(prediction, extra_outputs)`` with + :class:`DiffusionExtraOutputs` for models that produce auxiliary outputs. + + The primary time and prediction remain required tensors, separate from + the optional named fields. Use tuples with consumers that support the + supplied fields. Prediction targets follow :data:`PredictorType`. The interface is **prediction-agnostic**: whether your model predicts clean data (:math:`\mathbf{x}_0`), noise (:math:`\epsilon`), score @@ -69,6 +234,9 @@ class DiffusionModel(Protocol): Examples -------- + An ordinary model can take one time tensor and return one prediction. + The protocol checks that the model exposes the expected call interface: + >>> import torch >>> import torch.nn.functional as F >>> from physicsnemo.diffusion import DiffusionModel @@ -79,15 +247,40 @@ class DiffusionModel(Protocol): ... >>> isinstance(Model(), DiffusionModel) True + + A compatible model can also consume an extra time and return an auxiliary + output. This example keeps the prediction as the first tuple element and + uses the reserved fields with their intended meanings: + + >>> from physicsnemo.diffusion import DiffusionExtraTimes, DiffusionExtraOutputs + >>> class ExtendedModel: + ... def __call__(self, x, t, condition=None, **kwargs): + ... t, extra_times = t + ... prediction = x * (t - extra_times.r)[:, None] + ... extra_outputs = DiffusionExtraOutputs( + ... logvar=torch.zeros_like(t), batch_size=[x.shape[0]] + ... ) + ... # The required prediction stays first in the output tuple. + ... return prediction, extra_outputs + >>> x = torch.ones(2, 3) + >>> t = torch.ones(2) + >>> times = (t, DiffusionExtraTimes(r=torch.zeros(2), batch_size=[2])) + >>> prediction, extra_outputs = ExtendedModel()(x, times) + >>> prediction.shape, extra_outputs.logvar.shape + (torch.Size([2, 3]), torch.Size([2])) """ def __call__( self, x: Float[torch.Tensor, " B *dims"], - t: Float[torch.Tensor, " B"], + t: Float[torch.Tensor, " B"] + | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]], condition: Float[torch.Tensor, " B *cond_dims"] | TensorDict | None = None, **model_kwargs: Any, - ) -> Float[torch.Tensor, " B *dims"]: + ) -> ( + Float[torch.Tensor, " B *dims"] + | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]] + ): r""" Forward pass of the diffusion model. @@ -97,8 +290,10 @@ def __call__( 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 - Diffusion time or noise level tensor of shape :math:`(B,)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time for each sample, with shape :math:`(B,)`. + Pass ``(t, extra_times)`` to supply extra named times in + :class:`DiffusionExtraTimes`, each with shape :math:`(B,)`. condition : torch.Tensor, TensorDict, or None, optional, default=None Conditioning information for the model. If a Tensor or a TensorDict is passed, it should have batch size :math:`B` matching that of @@ -109,8 +304,10 @@ def __call__( Returns ------- - torch.Tensor - Model output with the same shape as ``x``. + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + Prediction with the shape of ``x``, optionally paired with a + :class:`DiffusionExtraOutputs`. Every call returns a prediction + tensor; auxiliary tensors share its batch dimension. """ ... @@ -130,7 +327,7 @@ class Predictor(Protocol): prediction. A predictor can be a trained neural network, a guidance function (e.g., classifier-free guidance, DPS-style guidance), or any combination thereof. The exact meaning of the output depends on the - predictor type and how it is used. Any callable that implements this + :data:`PredictorType` and its use. Any callable that implements this interface can be used as a predictor in sampling utilities. This protocol is typically used during inference. For training, which @@ -163,8 +360,9 @@ class Predictor(Protocol): Examples -------- - **Example 1:** Convert a trained conditional model into a predictor using - ``functools.partial``: + **Example 1:** A conditional model needs ``y`` to form its clean-data + estimate. Bind that condition with ``functools.partial`` so sampling code + can call the resulting predictor with only ``x`` and ``t``: >>> import torch >>> from functools import partial @@ -184,9 +382,10 @@ class Predictor(Protocol): >>> isinstance(x0_predictor, Predictor) True - **Example 2:** Convert the x0-predictor above into a score-predictor - (using a simple EDM-like schedule where :math:`\sigma(t) = t` and - :math:`\alpha(t) = 1`): + **Example 2:** A scheduler may need a score rather than the clean-data + estimate produced above. For the illustrative schedule + :math:`\sigma(t) = t` and :math:`\alpha(t) = 1`, convert that estimate + with the score identity while keeping the same predictor call signature: >>> def x0_to_score(x0, x_t, t): ... sigma_sq = t.view(-1, 1) ** 2 @@ -198,6 +397,24 @@ class Predictor(Protocol): >>> >>> isinstance(score_predictor, Predictor) True + + **Example 3:** A model can return :class:`DiffusionExtraOutputs` for + training, such as a ``logvar`` field. Wrap the model to select its + prediction tensor for sampling. Standard denoiser factories require a + predictor that accepts one diffusion time tensor: + + >>> from physicsnemo.diffusion import DiffusionExtraOutputs + >>> def model_with_extras(x, t, condition=None): + ... return x / 2, DiffusionExtraOutputs( + ... logvar=torch.zeros(x.shape[0]), batch_size=[x.shape[0]] + ... ) + >>> def prediction_only(x, t): + ... prediction, _ = model_with_extras(x, t) + ... return prediction + >>> prediction_only(torch.ones(2, 3), torch.ones(2)).shape + torch.Size([2, 3]) + >>> isinstance(prediction_only, Predictor) + True """ def __call__( @@ -270,8 +487,10 @@ class Denoiser(Protocol): Examples -------- - Manually creating a denoiser from an x0-predictor using a simple EDM-like - schedule (:math:`\sigma(t)=t`, :math:`\alpha(t)=1`): + This example builds a denoiser from a clean-data predictor for the + illustrative schedule :math:`\sigma(t)=t` and :math:`\alpha(t)=1`. + The returned tensor is the update term consumed by a solver. It uses one + time tensor, as required by the standard predictor and denoiser interfaces: >>> import torch >>> from physicsnemo.diffusion import Denoiser diff --git a/physicsnemo/diffusion/metrics/losses.py b/physicsnemo/diffusion/metrics/losses.py index 55aae56244..6f84d90d17 100644 --- a/physicsnemo/diffusion/metrics/losses.py +++ b/physicsnemo/diffusion/metrics/losses.py @@ -165,6 +165,14 @@ class MSEDSMLoss: :meth:`~physicsnemo.diffusion.noise_schedulers.LinearGaussianNoiseScheduler.score_to_x0` method provides a ready-made ``score_to_x0_fn``. + Notes + ----- + This loss passes one diffusion time tensor to the model and expects one + prediction tensor in return. It does not support ``(t, extra_times)`` + tuples with :class:`~physicsnemo.diffusion.DiffusionExtraTimes` or + ``(prediction, extra_outputs)`` tuples with + :class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. + Parameters ---------- model : DiffusionModel @@ -503,6 +511,14 @@ class WeightedMSEDSMLoss: For more details on prediction types, expected signatures, and additional examples, see :class:`MSEDSMLoss`. + Notes + ----- + This loss passes one diffusion time tensor to the model and expects one + prediction tensor in return. It does not support ``(t, extra_times)`` + tuples with :class:`~physicsnemo.diffusion.DiffusionExtraTimes` or + ``(prediction, extra_outputs)`` tuples with + :class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. + Parameters ---------- model : DiffusionModel @@ -710,6 +726,14 @@ class FlowMatchingLoss: :class:`~physicsnemo.diffusion.noise_schedulers.DomainParallelNoiseScheduler`; a plain scheduler with sharded data raises a ``ValueError``. + Notes + ----- + This loss passes one diffusion time tensor to the model and expects one + prediction tensor in return. It does not support ``(t, extra_times)`` + tuples with :class:`~physicsnemo.diffusion.DiffusionExtraTimes` or + ``(prediction, extra_outputs)`` tuples with + :class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. + Parameters ---------- model : DiffusionModel @@ -941,6 +965,14 @@ class WeightedFlowMatchingLoss: ``weight`` must be a ``DTensor`` on the same device mesh; a plain tensor ``weight`` raises a ``ValueError``. + Notes + ----- + This loss passes one diffusion time tensor to the model and expects one + prediction tensor in return. It does not support ``(t, extra_times)`` + tuples with :class:`~physicsnemo.diffusion.DiffusionExtraTimes` or + ``(prediction, extra_outputs)`` tuples with + :class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. + Parameters ---------- model : DiffusionModel diff --git a/physicsnemo/diffusion/preconditioners/preconditioners.py b/physicsnemo/diffusion/preconditioners/preconditioners.py index 404240330b..eeac45dd74 100644 --- a/physicsnemo/diffusion/preconditioners/preconditioners.py +++ b/physicsnemo/diffusion/preconditioners/preconditioners.py @@ -17,16 +17,21 @@ import math from abc import ABC, abstractmethod from collections.abc import Sequence +from functools import partial from typing import Any, Tuple import torch -from jaxtyping import Float +from jaxtyping import Float, Shaped from tensordict import TensorDict from torch.distributed.tensor import DTensor from torch.distributed.tensor.placement_types import Replicate from physicsnemo.core.meta import ModelMetaData from physicsnemo.core.module import Module +from physicsnemo.diffusion.base import ( + DiffusionExtraOutputs, + DiffusionExtraTimes, +) def _ensure_plain_tensor(t: torch.Tensor) -> torch.Tensor: @@ -65,8 +70,8 @@ class BaseAffinePreconditioner(Module, ABC): transformation. This class provides a standardized interface for implementing - preconditioners that use affine transformations of the model - input and output. + preconditioners that use affine transformations of the model input and + output. The preconditioner wraps a neural network model :math:`F` and applies a preconditioning formula to transform the network output to produce @@ -90,7 +95,7 @@ class BaseAffinePreconditioner(Module, ABC): The wrapped model :math:`F` must be an instance of :class:`~physicsnemo.core.Module` that satisfies the :class:`~physicsnemo.diffusion.DiffusionModel` interface, with the - following signature: + following signature for a single diffusion time and prediction tensor: .. code-block:: python @@ -101,6 +106,16 @@ class BaseAffinePreconditioner(Module, ABC): **model_kwargs: Any, ) -> torch.Tensor # Shape: (B, *) + .. note:: + + Models can also accept ``(t, extra_times)`` with + :class:`~physicsnemo.diffusion.DiffusionExtraTimes` and return + ``(prediction, extra_outputs)`` with + :class:`~physicsnemo.diffusion.DiffusionExtraOutputs`. + Every supplied time tensor uses the same time-to-noise mapping and + noise encoding as the primary time. The affine formula transforms + only the prediction tensor. Auxiliary outputs pass through unchanged. + The preconditioner is agnostic to the prediction target of the wrapped model :math:`F`. The same preconditioning formula is applied regardless of whether the model is an :math:`\mathbf{x}_0`-predictor, an @@ -128,8 +143,11 @@ class BaseAffinePreconditioner(Module, ABC): 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. - t : torch.Tensor - Diffusion time tensor of shape :math:`(B,)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time for each sample, with shape :math:`(B,)`. Pass + ``(t, extra_times)`` to supply extra named diffusion times in + :class:`~physicsnemo.diffusion.DiffusionExtraTimes`. + The preconditioner encodes every supplied time for the backbone. condition : torch.Tensor, TensorDict, or None, optional, default=None Single Tensor or a TensorDict containing conditioning tensors with batch size :math:`B` matching that of ``x``. Pass ``None`` for an @@ -140,9 +158,10 @@ class BaseAffinePreconditioner(Module, ABC): Outputs ------- - torch.Tensor - Preconditioned model output with the same shape as the original model - output. + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + Preconditioned prediction with the shape of ``x``, optionally paired + with auxiliary outputs. The affine formula transforms only the + prediction; auxiliary outputs keep their values and gradients. Note ---- @@ -182,7 +201,8 @@ class BaseAffinePreconditioner(Module, ABC): method since :math:`\sigma(t) = t` (noise level and diffusion time are the same). - We first define a simple model to wrap: + First, define a small backbone with the required call signature. Its + convolution provides a prediction with the same shape as its input: >>> import torch >>> from tensordict import TensorDict @@ -196,7 +216,8 @@ class BaseAffinePreconditioner(Module, ABC): ... def forward(self, x, t, condition=None): ... return self.net(x) - Now we define the EDM preconditioner: + The subclass below supplies the four EDM coefficients. It wraps the + backbone and keeps the original prediction shape: >>> from physicsnemo.diffusion.preconditioners import ( ... BaseAffinePreconditioner, @@ -225,9 +246,9 @@ class BaseAffinePreconditioner(Module, ABC): >>> out.shape torch.Size([2, 3, 16, 16]) - The following example shows how to override the :meth:`sigma` method to - implement a Variance Exploding (VE) preconditioner where - :math:`\sigma(t) = \sqrt{t}`. + A different schedule can override :meth:`sigma` while retaining the + preconditioner interface. This VE example uses + :math:`\sigma(t) = \sqrt{t}`: >>> class VEPreconditioner(BaseAffinePreconditioner): ... def __init__(self, model): @@ -256,17 +277,16 @@ class BaseAffinePreconditioner(Module, ABC): Some models in PhysicsNeMo have signatures that differ from the :class:`~physicsnemo.diffusion.DiffusionModel` interface. Below are - examples showing how to write thin wrappers to make them compatible - with preconditioners, including image-based conditioning via channel + examples showing thin wrappers that make them compatible with + preconditioners. They also show image conditioning through channel concatenation. **Example: Wrapping SongUNet** The :class:`~physicsnemo.models.diffusion_unets.SongUNet` model has the signature ``forward(x, noise_labels, class_labels, augment_labels)``. - We wrap it to match ``forward(x, t, condition)``, where ``condition`` - contains both class labels (1D vector) and an image to concatenate - channel-wise: + We wrap it to match ``forward(x, t, condition)``. The condition carries + class labels as a vector and an image to concatenate along channels: >>> from physicsnemo.models.diffusion_unets import SongUNet >>> from physicsnemo.diffusion import DiffusionModel @@ -309,8 +329,8 @@ class BaseAffinePreconditioner(Module, ABC): The :class:`~physicsnemo.models.dit.DiT` model has the signature ``forward(x, t, condition, ...)``. We wrap it to support - both image conditioning (via channel concatenation) and vector - conditioning: + both image conditioning through channel concatenation and vector + conditioning. The adapter passes the vector to DiT's ``condition`` input: >>> from physicsnemo.models.dit import DiT >>> class DiTWrapper(Module): @@ -355,7 +375,8 @@ class BaseAffinePreconditioner(Module, ABC): is also passed as a separate argument) is common across several diffusion use-cases. Thus for convenience, we provide the wrapper class :class:`~physicsnemo.diffusion.utils.ConcatConditionWrapper` - to save you the trouble of writing your own wrapper for this common pattern: + to use this common pattern. The example pairs it with an EDM + preconditioner and checks the output shape: >>> from physicsnemo.diffusion.preconditioners import EDMPreconditioner >>> from physicsnemo.diffusion.utils import ConcatConditionWrapper @@ -451,13 +472,34 @@ def _validate_input(self, x: torch.Tensor) -> None: tracing. The default implementation is a no-op. """ + def _encode_time( + self, + t: Float[torch.Tensor, " B"], + x: Float[torch.Tensor, " B *dims"], + c_noise: Float[torch.Tensor, " B *ones"] | None = None, + ) -> Float[torch.Tensor, " B"]: + r"""Encode a diffusion time, reusing its noise coefficient when available.""" + if c_noise is None: + shape = (-1,) + (1,) * (x.ndim - 1) + sigma = self.sigma(t).reshape(shape) # (B, 1, ...) + _, c_noise, _, _ = self.compute_coefficients(_ensure_plain_tensor(sigma)) + c_noise = _ensure_plain_tensor(c_noise) + if isinstance(x, DTensor): + c_noise = _replicate_on_mesh(c_noise, x.device_mesh) + return c_noise.flatten() # (B,) + def forward( self, - x: torch.Tensor, - t: torch.Tensor, - condition: torch.Tensor | TensorDict | None = None, + x: Float[torch.Tensor, " B *dims"], + t: Float[torch.Tensor, " B"] + | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]], + condition: Float[torch.Tensor, " B *cond_dims"] | TensorDict | None = None, **model_kwargs: Any, - ) -> torch.Tensor: + ) -> ( + Float[torch.Tensor, " B *dims"] + | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]] + ): + t, extra_times = t if isinstance(t, tuple) else (t, None) if not torch.compiler.is_compiling(): B = x.shape[0] if t.shape != (B,): @@ -465,6 +507,11 @@ def forward( f"Expected t to have shape ({B},) matching batch size of " f"x, but got {t.shape}." ) + if extra_times is not None: + if extra_times.batch_size != torch.Size([B]): + raise ValueError( + f"Expected time batch size [{B}], got {extra_times.batch_size}." + ) if isinstance(condition, TensorDict): if condition.batch_size and condition.batch_size[0] != B: raise ValueError( @@ -481,7 +528,7 @@ def forward( # Map time step to noise level via sigma method expected_shape = (-1,) + (1,) * (x.ndim - 1) - sigma_t = self.sigma(t).reshape(expected_shape) + sigma_t = self.sigma(t).reshape(expected_shape) # (B, 1, ...) # Compute preconditioning coefficients # Unwrap to plain tensors so that element-wise arithmetic between @@ -495,9 +542,7 @@ def forward( # tensors, then re-promote to replicated DTensors on x's mesh when # x is a ShardTensor so that element-wise arithmetic between # coefficients and data is type-compatible. - c_in, c_noise, c_out, c_skip = ( - _ensure_plain_tensor(c) for c in (c_in, c_noise, c_out, c_skip) - ) + c_in, c_out, c_skip = (_ensure_plain_tensor(c) for c in (c_in, c_out, c_skip)) # Only manually promote coefficients for genuine (non-Shard) DTensor # activations. A ShardTensor subclasses torch.Tensor (not DTensor), so @@ -507,28 +552,32 @@ def forward( # otherwise breaks ops like PositionalEmbedding's torch.outer). if isinstance(x, DTensor): x_mesh = x.device_mesh - c_in, c_noise, c_out, c_skip = ( - _replicate_on_mesh(c, x_mesh) for c in (c_in, c_noise, c_out, c_skip) + c_in, c_out, c_skip = ( + _replicate_on_mesh(c, x_mesh) for c in (c_in, c_out, c_skip) ) - # Forward through the underlying model + model_t = self._encode_time(t, x, c_noise=c_noise) # (B,) + if extra_times is not None: + model_t = (model_t, extra_times.apply(partial(self._encode_time, x=x))) + if condition is not None: F_x = self.model( c_in * x, - c_noise.flatten(), + model_t, condition=condition, **model_kwargs, ) else: F_x = self.model( c_in * x, - c_noise.flatten(), + model_t, **model_kwargs, ) - D_x = c_skip * x + c_out * F_x - - return D_x + if isinstance(F_x, tuple): + prediction, extra_outputs = F_x + return c_skip * x + c_out * prediction, extra_outputs + return c_skip * x + c_out * F_x class VPPreconditioner(BaseAffinePreconditioner): @@ -604,6 +653,10 @@ class VPPreconditioner(BaseAffinePreconditioner): Examples -------- + The VP coefficients map diffusion time to a noise level before calling + this convolutional backbone. The resulting prediction retains the image + shape for comparison with a clean-data target: + >>> import torch >>> from physicsnemo.core import Module >>> # Define a simple model satisfying the diffusion model interface @@ -742,6 +795,10 @@ class VEPreconditioner(BaseAffinePreconditioner): Examples -------- + This VE preconditioner wraps a convolutional backbone whose prediction + matches the image shape. Passing one time per image produces one + preconditioned prediction per image: + >>> import torch >>> from physicsnemo.core import Module >>> # Define a simple model satisfying the diffusion model interface @@ -857,6 +914,10 @@ class IDDPMPreconditioner(BaseAffinePreconditioner): Examples -------- + This iDDPM preconditioner receives a time for each image and calls a + convolutional backbone. It returns a clean-data prediction with the + same batch and spatial dimensions: + >>> import torch >>> from physicsnemo.core import Module >>> # Define a simple model satisfying the diffusion model interface @@ -1005,7 +1066,9 @@ class EDMPreconditioner(BaseAffinePreconditioner): Examples -------- - Scalar ``sigma_data`` (default): + The first example uses one ``sigma_data`` value for all channels. A + convolutional backbone supplies the raw prediction, and EDM applies its + coefficients to produce a clean-data estimate: >>> import torch >>> from physicsnemo.core import Module @@ -1024,8 +1087,8 @@ class EDMPreconditioner(BaseAffinePreconditioner): >>> out.shape torch.Size([2, 3, 16, 16]) - Per-channel ``sigma_data`` for heterogeneous channels (e.g. weather - variables with different scales): + The second example gives each channel its own ``sigma_data``. This is + useful for variables with different scales, such as weather fields: >>> precond_ch = EDMPreconditioner(model, sigma_data=[0.3, 0.5, 0.7]) >>> out_ch = precond_ch(x, t, condition=None) diff --git a/physicsnemo/diffusion/utils/model_wrappers.py b/physicsnemo/diffusion/utils/model_wrappers.py index 9837f54d3c..bed1f91b4b 100644 --- a/physicsnemo/diffusion/utils/model_wrappers.py +++ b/physicsnemo/diffusion/utils/model_wrappers.py @@ -18,10 +18,11 @@ from typing import Any import torch -from jaxtyping import Float +from jaxtyping import Float, Shaped from tensordict import TensorDict from physicsnemo.core import Module +from physicsnemo.diffusion.base import DiffusionExtraOutputs, DiffusionExtraTimes from physicsnemo.models.diffusion_unets import ( DhariwalUNet, SongUNet, @@ -88,8 +89,9 @@ class ConcatConditionWrapper(Module): x : torch.Tensor Noisy latent state of shape :math:`(B, *)` where :math:`B` is the batch size. - t : torch.Tensor - Diffusion time tensor of shape :math:`(B,)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time tensor of shape :math:`(B,)`, optionally paired with + extra times. The wrapped model must accept the supplied format. condition : torch.Tensor, TensorDict, or None, optional, default=None Conditioning data. Use a ``TensorDict`` to explicitly supply concatenated conditioning and vector conditioning as keys @@ -101,13 +103,15 @@ class ConcatConditionWrapper(Module): Outputs ------- - torch.Tensor - Model output with the same shape as ``x``. + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + The wrapped model's prediction, optionally paired with auxiliary + outputs. The wrapper returns both unchanged. Examples -------- Wrap :class:`~physicsnemo.models.diffusion_unets.SongUNet` with channel - concatenation and vector conditioning: + concatenation and vector conditioning. The wrapper routes the image to + the input channels and the vector to class labels: >>> import torch >>> from tensordict import TensorDict @@ -127,7 +131,8 @@ class ConcatConditionWrapper(Module): >>> out.shape torch.Size([2, 3, 8, 8]) - Wrap :class:`~physicsnemo.models.dit.DiT` similarly: + The same condition layout works for :class:`~physicsnemo.models.dit.DiT`. + The wrapper sends the vector to DiT's ``condition`` input: >>> from physicsnemo.models.dit import DiT >>> dit = DiT(in_channels=4, out_channels=3, condition_dim=4, input_size=8, patch_size=4) @@ -135,6 +140,27 @@ class ConcatConditionWrapper(Module): >>> out = dit_wrapper(x, t, condition) >>> out.shape torch.Size([2, 3, 8, 8]) + + A backbone that supports extra times and outputs can use tuples. The + wrapper forwards these values and returns both output components: + + >>> from physicsnemo.diffusion import DiffusionExtraTimes, DiffusionExtraOutputs + >>> class TwoTimeModel(Module): + ... def forward(self, x, t, condition=None): + ... t, extra_times = t + ... # The first three channels exclude concatenated conditioning. + ... prediction = x[:, :3] * (t - extra_times.r)[:, None, None, None] + ... # Keep the required prediction separate from auxiliary outputs. + ... return prediction, DiffusionExtraOutputs( + ... logvar=torch.zeros_like(t), batch_size=[x.shape[0]] + ... ) + >>> wrapper = ConcatConditionWrapper(TwoTimeModel()) + >>> times = (t, DiffusionExtraTimes(r=torch.zeros_like(t), batch_size=[2])) + >>> prediction, extra_outputs = wrapper(x, times, condition) + >>> prediction.shape, extra_outputs.logvar.shape + (torch.Size([2, 3, 8, 8]), torch.Size([2])) + + SongUNet and DhariwalUNet backbones still require a single time tensor. """ def __init__( @@ -151,10 +177,14 @@ def __init__( def forward( self, x: Float[torch.Tensor, " B *dims"], - t: Float[torch.Tensor, " B"], + t: Float[torch.Tensor, " B"] + | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]], condition: Float[torch.Tensor, " B *cond_dims"] | TensorDict | None = None, **model_kwargs: Any, - ) -> Float[torch.Tensor, " B *dims"]: + ) -> ( + Float[torch.Tensor, " B *dims"] + | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]] + ): r""" Forward pass for the conditioned wrapper. @@ -163,8 +193,9 @@ def forward( x : torch.Tensor Noisy latent state of shape :math:`(B, *)` where :math:`B` is the batch size. - t : torch.Tensor - Diffusion time tensor of shape :math:`(B,)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time tensor of shape :math:`(B,)`, optionally paired + with extra times. Forwarded to the backbone unchanged. condition : torch.Tensor, TensorDict, or None, optional, default=None Conditioning data. Use a ``TensorDict`` to supply ``cond_concat`` and ``cond_vec`` tensors. A plain tensor is treated as the concatenated @@ -174,9 +205,11 @@ def forward( Returns ------- - torch.Tensor - Model output with the same shape as ``x``. + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + The wrapped model's prediction and any auxiliary outputs, + returned unchanged. """ + t, extra_times = t if isinstance(t, tuple) else (t, None) cond_concat: torch.Tensor | None = None cond_vec: torch.Tensor | None = None @@ -246,6 +279,7 @@ def forward( if cond_concat is not None: x = torch.cat([x, cond_concat], dim=1) + t = t if extra_times is None else (t, extra_times) if isinstance( self.model, (SongUNet, SongUNetPosEmbd, SongUNetPosLtEmbd, DhariwalUNet) ): diff --git a/test/diffusion/helpers.py b/test/diffusion/helpers.py index 8553d34b4e..9f96aef56a 100644 --- a/test/diffusion/helpers.py +++ b/test/diffusion/helpers.py @@ -20,10 +20,15 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch +from jaxtyping import Float, Shaped from tensordict import TensorDict import physicsnemo.core from physicsnemo.core import Module +from physicsnemo.diffusion import ( + DiffusionExtraOutputs, + DiffusionExtraTimes, +) # Directory for test reference data DATA_DIR = Path(__file__).parent / "data" @@ -77,6 +82,49 @@ def forward(self, x: torch.Tensor, t: torch.Tensor, **kwargs: Any) -> torch.Tens return self.net(x) / (1 + t_bc) +class CustomExtraTimes(DiffusionExtraTimes): + r"""Additional time fields that callers can supply independently.""" + + r1: Float[torch.Tensor, " B"] | None = None + r2: Float[torch.Tensor, " B"] | None = None + + +class StructuredDiffusionModel(Module): + r"""Linear model with optional destination times and learned loss weights.""" + + def __init__(self, output: str = "tuple"): + super().__init__() + self.output = output + self.gain = torch.nn.Parameter(torch.tensor(0.75)) + + def forward( + self, + x: Float[torch.Tensor, " B *dims"], + t: Float[torch.Tensor, " B"] + | tuple[Float[torch.Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]], + condition: TensorDict | None = None, + *, + offset: float = 0.0, + ) -> ( + Float[torch.Tensor, " B *dims"] + | tuple[Float[torch.Tensor, " B *dims"], Shaped[DiffusionExtraOutputs, " B"]] + ): + r"""Return a prediction, optionally paired with its loss weight.""" + t, extra_times = t if isinstance(t, tuple) else (t, None) + t_bc = t.reshape(-1, *([1] * (x.ndim - 1))) # (B, 1, ...) + prediction = self.gain * x + t_bc + offset + if extra_times is not None: + for time in extra_times.to_tensordict().values(): + prediction = prediction + time.reshape(-1, *([1] * (x.ndim - 1))) + if condition is not None: + prediction = prediction + condition["shift"] + if self.output == "tensor": + return prediction + return prediction, DiffusionExtraOutputs( + logvar=self.gain * t, batch_size=[x.shape[0]] + ) + + def make_input( shape: Tuple[int, ...], seed: int = 42, diff --git a/test/diffusion/test_model_io.py b/test/diffusion/test_model_io.py new file mode 100644 index 0000000000..4e4681b124 --- /dev/null +++ b/test/diffusion/test_model_io.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Check primary tensors and extra fields through diffusion wrappers.""" + +import pytest +import torch +from jaxtyping import Float + +from physicsnemo.diffusion import DiffusionExtraOutputs, DiffusionExtraTimes +from physicsnemo.diffusion.preconditioners import ( + EDMPreconditioner, + IDDPMPreconditioner, + VEPreconditioner, + VPPreconditioner, +) +from physicsnemo.diffusion.utils.model_wrappers import ConcatConditionWrapper + +from .helpers import CustomExtraTimes, StructuredDiffusionModel + +CONFIGS = ( + (EDMPreconditioner, {}), + (VEPreconditioner, {}), + (VPPreconditioner, {}), + (IDDPMPreconditioner, {"M": 32}), +) +IDS = ("edm", "ve", "vp", "iddpm") + +TIME_FIELDS = (None, (), ("r",), ("r1",), ("r2",), ("r1", "r2")) +TIME_IDS = ("tensor", "empty", "r", "r1", "r2", "r1_r2") + + +def _make_time( + t: Float[torch.Tensor, " B"], fields: tuple[str, ...] | None +) -> Float[torch.Tensor, " B"] | tuple[Float[torch.Tensor, " B"], CustomExtraTimes]: + r"""Pair the primary time with the selected extra fields.""" + if fields is None: + return t + return t, CustomExtraTimes( + **{name: t / (index + 2) for index, name in enumerate(fields)}, + batch_size=t.shape, + ) + + +@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS) +class TestConsistency: + @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS) + def test_tuple_protocol(self, preconditioner, kwargs, device, fields): + model = preconditioner(StructuredDiffusionModel(), **kwargs).to(device) + x = torch.randn(2, 3, 5, 7, device=device) + t = torch.tensor([0.4, 0.8], device=device) + time = _make_time(t, fields) + calls = [] + handle = model.model.register_forward_hook( + lambda module, inputs, output: calls.append((inputs, output)) + ) + prediction, outputs = model(x, time, offset=0.2) + handle.remove() + inputs, raw = calls[0] + sigma = model.sigma(t).reshape(2, 1, 1, 1) # (B, 1, 1, 1) + c_in, c_noise, c_out, c_skip = model.compute_coefficients(sigma) + torch.testing.assert_close(inputs[0], c_in * x) + encoded_t = inputs[1][0] if fields is not None else inputs[1] + torch.testing.assert_close(encoded_t, c_noise.flatten()) + if fields is not None: + assert isinstance(inputs[1][1], CustomExtraTimes) + for index, name in enumerate(fields): + original = t / (index + 2) + sigma_extra = model.sigma(original).reshape(2, 1, 1, 1) + encoded = model.compute_coefficients(sigma_extra)[1].flatten() + torch.testing.assert_close(getattr(inputs[1][1], name), encoded) + torch.testing.assert_close(getattr(time[1], name), original) + for name in {"r", "r1", "r2"}.difference(fields): + assert getattr(inputs[1][1], name) is None + torch.testing.assert_close(prediction, c_skip * x + c_out * raw[0]) + assert outputs is raw[1] + torch.testing.assert_close(outputs.logvar, model.model.gain * encoded_t) + + def test_tensor_output(self, preconditioner, kwargs, device): + model = preconditioner(StructuredDiffusionModel(), **kwargs).to(device) + plain = preconditioner(StructuredDiffusionModel("tensor"), **kwargs).to(device) + x = torch.randn(2, 3, device=device) + t = torch.tensor([0.4, 0.8], device=device) + prediction, _ = model(x, t) + torch.testing.assert_close(prediction, plain(x, t)) + + +class TestConstructor: + def test_optional_outputs(self): + extras = DiffusionExtraOutputs(batch_size=[2]) + assert extras.logvar is None + + def test_conditioning_wrapper(self): + model = StructuredDiffusionModel() + wrapper = ConcatConditionWrapper(model) + x = torch.randn(2, 3, 5, 7) + t = torch.tensor([0.4, 0.8]) + time = (t, DiffusionExtraTimes(r=t / 2, batch_size=[2])) + condition = torch.ones(2, 1, 5, 7) + prediction, extras = wrapper(x, time, condition) + expected, expected_extras = model(torch.cat((x, condition), dim=1), time) + torch.testing.assert_close(prediction, expected) + torch.testing.assert_close(extras.logvar, expected_extras.logvar) + + +@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS) +class TestGradientFlow: + @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS) + def test_numerical_gradients(self, preconditioner, kwargs, fields): + model = preconditioner(StructuredDiffusionModel(), **kwargs).double() + x = torch.randn(2, 3, dtype=torch.float64, requires_grad=True) + t = torch.tensor([0.4, 0.8], dtype=torch.float64, requires_grad=True) + + def forward(x, t): + prediction, extras = model(x, _make_time(t, fields)) + return prediction, extras.logvar + + assert torch.autograd.gradcheck(forward, (x, t)) + + +@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS) +class TestCompile: + @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS) + def test_forward_backward(self, preconditioner, kwargs, device, fields): + model = preconditioner(StructuredDiffusionModel(), **kwargs).to(device) + x = torch.randn(2, 3, 5, 7, device=device, requires_grad=True) + t = torch.tensor([0.4, 0.8], device=device, requires_grad=True) + time = _make_time(t, fields) + expected, expected_extras = model(x, time) + expected_grads = torch.autograd.grad( + expected.sum() + expected_extras.logvar.sum(), + (x, t, model.model.gain), + retain_graph=True, + ) + compiled = torch.compile(model, backend="aot_eager", fullgraph=True) + result, extras = compiled(x, time) + with torch._dynamo.config.patch(error_on_recompile=True): + repeated, _ = compiled(x, time) + torch.testing.assert_close(repeated, result) + torch.testing.assert_close(result, expected) + torch.testing.assert_close(extras.logvar, expected_extras.logvar) + gradients = torch.autograd.grad( + result.sum() + extras.logvar.sum(), (x, t, model.model.gain) + ) + for actual, reference in zip(gradients, expected_grads): + torch.testing.assert_close(actual, reference) + + def test_conditioning_wrapper(self, preconditioner, kwargs, device): + wrapper = ConcatConditionWrapper( + preconditioner(StructuredDiffusionModel(), **kwargs) + ).to(device) + x = torch.randn(2, 3, device=device, requires_grad=True) + t = torch.tensor([0.4, 0.8], device=device) + time = (t, DiffusionExtraTimes(r=t / 2, batch_size=[2])) + expected, _ = wrapper(x, time) + compiled = torch.compile(wrapper, backend="aot_eager", fullgraph=True) + actual, _ = compiled(x, time) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + torch.autograd.grad(actual.sum(), x)[0], + torch.autograd.grad(expected.sum(), x)[0], + ) + + +class TestDiagnostics: + @pytest.mark.parametrize( + "t,r,batch_size", + ( + (torch.ones(2, 1), torch.ones(2), [2]), + (torch.ones(2), torch.ones(2), []), + ), + ) + def test_invalid_time(self, t, r, batch_size): + model = EDMPreconditioner(StructuredDiffusionModel()) + with pytest.raises(ValueError, match="shape|batch"): + model( + torch.ones(2, 3), (t, DiffusionExtraTimes(r=r, batch_size=batch_size)) + ) + + +@pytest.mark.parametrize("preconditioner,kwargs", CONFIGS, ids=IDS) +class TestNonRegression: + @pytest.mark.parametrize("fields", TIME_FIELDS, ids=TIME_IDS) + def test_checkpoint(self, preconditioner, kwargs, tmp_path, fields): + model = preconditioner(StructuredDiffusionModel(), **kwargs) + x = torch.randn(2, 3) + times = _make_time(torch.tensor([0.4, 0.8]), fields) + expected, expected_extras = model(x, times) + checkpoint = tmp_path / "model.pt" + torch.save(model.state_dict(), checkpoint) + restored = preconditioner(StructuredDiffusionModel(), **kwargs) + restored.load_state_dict(torch.load(checkpoint, weights_only=True)) + actual, extras = restored(x, times) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(extras.logvar, expected_extras.logvar) From 93133b1e96597930a82d1364d4632233b4995e92 Mon Sep 17 00:00:00 2001 From: Charlelie Laurent Date: Fri, 16 Oct 2026 13:13:54 -0700 Subject: [PATCH 2/2] Add inverse SNR and TrigFlow components Signed-off-by: Charlelie Laurent --- CHANGELOG.md | 10 + docs/api/diffusion/noise_schedulers.rst | 40 ++++ docs/api/diffusion/preconditioners.rst | 8 + .../diffusion/noise_schedulers/__init__.py | 1 + physicsnemo/diffusion/noise_schedulers/edm.py | 11 +- .../diffusion/noise_schedulers/iddpm.py | 7 +- .../noise_schedulers/linear_gaussian.py | 93 +++++++- .../noise_schedulers/rectified_flow.py | 9 +- .../diffusion/noise_schedulers/trigflow.py | 190 +++++++++++++++++ physicsnemo/diffusion/noise_schedulers/ve.py | 8 +- physicsnemo/diffusion/noise_schedulers/vp.py | 18 +- .../diffusion/preconditioners/__init__.py | 1 + .../preconditioners/preconditioners.py | 88 ++++++++ .../data/trigflow_precond_conv.mdlus | Bin 0 -> 3608 bytes .../trigflow_precond_conv_coefficients.pth | Bin 0 -> 2731 bytes .../data/trigflow_precond_conv_forward.pth | Bin 0 -> 4163 bytes .../data/trigflow_precond_conv_sigma.pth | Bin 0 -> 1845 bytes .../data/trigflow_precond_linear.mdlus | Bin 0 -> 3998 bytes .../trigflow_precond_linear_coefficients.pth | Bin 0 -> 2751 bytes .../data/trigflow_precond_linear_forward.pth | Bin 0 -> 2129 bytes .../data/trigflow_precond_linear_sigma.pth | Bin 0 -> 1859 bytes test/diffusion/test_model_io.py | 4 +- test/diffusion/test_noise_schedulers.py | 4 + test/diffusion/test_preconditioners.py | 6 +- test/diffusion/test_trigflow.py | 198 ++++++++++++++++++ 25 files changed, 683 insertions(+), 13 deletions(-) create mode 100644 physicsnemo/diffusion/noise_schedulers/trigflow.py create mode 100644 test/diffusion/data/trigflow_precond_conv.mdlus create mode 100644 test/diffusion/data/trigflow_precond_conv_coefficients.pth create mode 100644 test/diffusion/data/trigflow_precond_conv_forward.pth create mode 100644 test/diffusion/data/trigflow_precond_conv_sigma.pth create mode 100644 test/diffusion/data/trigflow_precond_linear.mdlus create mode 100644 test/diffusion/data/trigflow_precond_linear_coefficients.pth create mode 100644 test/diffusion/data/trigflow_precond_linear_forward.pth create mode 100644 test/diffusion/data/trigflow_precond_linear_sigma.pth create mode 100644 test/diffusion/test_trigflow.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f560ac5bc..7d466860e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 and EDM2. Existing single-time, single-prediction calls remain unchanged. Existing DSM and flow-matching losses continue to consume single tensors. +- Adds `TrigFlowNoiseScheduler` and `TrigFlowPreconditioner` to the diffusion + noise schedulers and preconditioners in `physicsnemo.diffusion`. Pair them + with the same `sigma_data` to train and sample models along a trigonometric + diffusion path, with the backbone predicting normalized velocity. Built-in + linear Gaussian schedulers also provide differentiable `snr_inv()` methods + to convert signal-to-noise amplitude ratios back to diffusion times. Use + these inverses to match noise levels across schedules. The amplitude ratio + is `alpha(t) / sigma(t)`. The power signal-to-noise ratio is its square. + Custom schedulers can supply their own inverse. These additions are opt-in. + Existing schedulers, preconditioners, and training workflows work as before. - Adds standalone FLARE++ attention and model APIs, with input-conditioned dynamic routing, plus a ``GALE_FPP`` backend for using the same mixer inside GeoTransolver. diff --git a/docs/api/diffusion/noise_schedulers.rst b/docs/api/diffusion/noise_schedulers.rst index dda757e7fe..90f0d58675 100644 --- a/docs/api/diffusion/noise_schedulers.rst +++ b/docs/api/diffusion/noise_schedulers.rst @@ -23,6 +23,46 @@ the coupling between noise schedule and reverse process in one place and avoids the need for schedule-specific sampler code. +Amplitude ratio and inverse +--------------------------- + +The amplitude ratio :math:`\alpha(t)/\sigma(t)` measures signal strength +relative to noise strength. The method ``scheduler.snr(t)`` returns this +ratio. The usual power signal-to-noise ratio is its square. + +Use ``scheduler.snr_inv(q)`` to recover a diffusion time from an amplitude +ratio or to match noise levels between schedules. EDM, VE, VP, rectified-flow, +continuous iDDPM, and TrigFlow provide differentiable inverses. Use finite +positive ratios on the schedule's invertible domain. Inverses do not clip +to training bounds. Add ``snr_inv`` for custom schedulers that need this +conversion, and keep it differentiable for training losses such as sCM. + +.. code-block:: python + + import torch + from physicsnemo.diffusion.noise_schedulers import TrigFlowNoiseScheduler + + scheduler = TrigFlowNoiseScheduler(sigma_data=0.5) + t = torch.tensor([0.2, 0.7]) + recovered = scheduler.snr_inv(scheduler.snr(t)) + torch.testing.assert_close(recovered, t) + +TrigFlow interpolates between data and Gaussian noise with trigonometric +coefficients. Its diffusion time is an angle in radians, whereas EDM uses +the noise standard deviation. + +Pair ``TrigFlowNoiseScheduler`` with ``TrigFlowPreconditioner`` and the same +``sigma_data``. Training samples diffusion times uniformly. Inference uses +an evenly spaced grid ending at zero. The denoising weight +:math:`1/(\sigma_d\sin(t))^2` gives unit weight to the normalized-velocity +residual with this preconditioner. Keep training diffusion times above zero +for denoising losses and below :math:`\pi/2` for reverse-process sampling. + +.. autoclass:: TrigFlowNoiseScheduler + :show-inheritance: + :members: + + Role in Training ---------------- diff --git a/docs/api/diffusion/preconditioners.rst b/docs/api/diffusion/preconditioners.rst index 8a80c59f24..5f150a6c94 100644 --- a/docs/api/diffusion/preconditioners.rst +++ b/docs/api/diffusion/preconditioners.rst @@ -140,6 +140,14 @@ API Reference :members: :exclude-members: forward +:code:`TrigFlowPreconditioner` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: physicsnemo.diffusion.preconditioners.TrigFlowPreconditioner + :show-inheritance: + :members: + :exclude-members: forward + :code:`VEPreconditioner` ~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/physicsnemo/diffusion/noise_schedulers/__init__.py b/physicsnemo/diffusion/noise_schedulers/__init__.py index 21cbe392a6..5316af9fc3 100644 --- a/physicsnemo/diffusion/noise_schedulers/__init__.py +++ b/physicsnemo/diffusion/noise_schedulers/__init__.py @@ -21,6 +21,7 @@ from .linear_gaussian import LinearGaussianNoiseScheduler # noqa: F401 from .rectified_flow import RectifiedFlowNoiseScheduler # noqa: F401 from .student_t_edm import StudentTEDMNoiseScheduler # noqa: F401 +from .trigflow import TrigFlowNoiseScheduler # noqa: F401 from .ve import VENoiseScheduler # noqa: F401 from .vp import VPNoiseScheduler # noqa: F401 diff --git a/physicsnemo/diffusion/noise_schedulers/edm.py b/physicsnemo/diffusion/noise_schedulers/edm.py index 6b0174e103..5189965ea3 100644 --- a/physicsnemo/diffusion/noise_schedulers/edm.py +++ b/physicsnemo/diffusion/noise_schedulers/edm.py @@ -71,7 +71,9 @@ class EDMNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the EDM noise scheduler: + The first sequence covers the EDM training and sampling interfaces with one + scheduler. Its final call converts a clean-data predictor into the right-hand side + used by an ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler @@ -98,7 +100,8 @@ class EDMNoiseScheduler(LinearGaussianNoiseScheduler): >>> denoiser(xN, tN).shape # ODE RHS for sampling torch.Size([4, 3, 8, 8]) - Per-channel ``sigma_data`` for heterogeneous channels: + A channel-specific ``sigma_data`` assigns different data scales to each channel. + The resulting loss weight retains a separate value for each channel. >>> sigma_per_ch = torch.tensor([0.3, 0.5, 0.7]) >>> scheduler_ch = EDMNoiseScheduler(sigma_data=sigma_per_ch) @@ -143,6 +146,10 @@ def sigma_inv( r"""Identity mapping: :math:`t = \sigma`.""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = 1/q` for amplitude ratios ``q > 0``.""" + return self.sigma_inv(q.reciprocal()) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/iddpm.py b/physicsnemo/diffusion/noise_schedulers/iddpm.py index ca33329988..126a5fc589 100644 --- a/physicsnemo/diffusion/noise_schedulers/iddpm.py +++ b/physicsnemo/diffusion/noise_schedulers/iddpm.py @@ -59,7 +59,8 @@ class IDDPMNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the iDDPM noise scheduler: + The first sequence builds noisy training samples and a reverse-time sampling state + for iDDPM. It then wraps a clean-data predictor as an ODE right-hand side. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import IDDPMNoiseScheduler @@ -131,6 +132,10 @@ def sigma_inv( r"""For iDDPM, :math:`t = \sigma` (identity mapping).""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find continuous diffusion time :math:`t = 1/q` for amplitude ratios ``q > 0``.""" + return q.reciprocal() + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py b/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py index ca5d7ae3e0..65196bbd97 100644 --- a/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py +++ b/physicsnemo/diffusion/noise_schedulers/linear_gaussian.py @@ -130,6 +130,7 @@ class LinearGaussianNoiseScheduler(ABC, NoiseScheduler): ... smin_rho = self.sigma_min**(1/self.rho) ... frac = i/(num_steps-1) ... t = (smax_rho + frac * (smin_rho - smax_rho))**self.rho + ... # Include the clean endpoint after the reverse-time grid. ... return torch.cat([t, torch.zeros(1, device=device)]) ... ... def sample_time(self, N, *, device=None, dtype=None): @@ -456,6 +457,50 @@ def snr( """ return self.alpha(t) / self.sigma(t) + def snr_inv( + self, + q: Float[Tensor, " *shape"], + ) -> Float[Tensor, " *shape"]: + r"""Find the diffusion time for a signal-to-noise amplitude ratio. + + Use this to match noise levels between diffusion schedules. The + ratio compares signal and noise amplitudes. The usual power + signal-to-noise ratio is its square. + + .. math:: + + q = \alpha(t) / \sigma(t), \qquad t = \mathrm{snr\_inv}(q). + + Parameters + ---------- + q : Tensor + Amplitude ratio, not its square. Built-in continuous inverses + support finite positive ratios without clipping to training bounds. + + Returns + ------- + Tensor + Diffusion times with the shape of ``q``. + + Notes + ----- + Override this method if your custom scheduler needs an inverse. + Keep the calculation differentiable for training losses such as sCM. + Schedulers without an inverse raise ``NotImplementedError`` here. + + Examples + -------- + For EDM, the signal coefficient is one and the noise coefficient equals time. + Inverting a signal-to-noise ratio of two gives time one half. + + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler + >>> EDMNoiseScheduler().snr_inv(torch.tensor([2.0])) + tensor([0.5000]) + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement snr_inv()." + ) + def x0_to_score( self, x0: Float[Tensor, " B *dims"], @@ -491,6 +536,10 @@ def x0_to_score( Examples -------- + A caller can convert a clean-data predictor into a score predictor when it needs + score parameterization. The conversion uses the noisy state and time to recover + the implied Gaussian score. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion @@ -562,6 +611,10 @@ def score_to_x0( Examples -------- + A caller can convert a score predictor into a clean-data predictor for a loss or + denoiser that expects that parameterization. The conversion combines the score + with the noisy state at the same time. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion @@ -620,6 +673,10 @@ def epsilon_to_score( Examples -------- + A caller can convert a noise predictor into a score predictor without changing + its backbone. The scheduler applies its time-dependent noise coefficient to the + predicted noise. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an epsilon-predictor, wrap it for manual conversion @@ -677,6 +734,9 @@ def score_to_epsilon( Examples -------- + A score predictor can also supply a noise prediction for a noise-based consumer. + The scheduler scales the score by its time-dependent noise coefficient. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion: @@ -740,6 +800,10 @@ def epsilon_to_x0( Examples -------- + A noise predictor can supply a clean-data estimate for losses that train on + clean targets. The conversion subtracts predicted noise from the noisy state + using the schedule coefficients. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an epsilon-predictor, wrap it for manual conversion @@ -801,6 +865,10 @@ def x0_to_epsilon( Examples -------- + A clean-data predictor can supply a noise estimate without retraining the + backbone. The conversion subtracts the predicted clean signal from the noisy + state. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion: @@ -867,6 +935,10 @@ def x0_to_flow( Examples -------- + A clean-data predictor can supply a velocity target for a flow-matching loss. + The conversion differentiates the schedule coefficients at the same diffusion + time. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have an x0-predictor, wrap it for manual conversion @@ -934,6 +1006,10 @@ def flow_to_x0( Examples -------- + A caller can convert a velocity predictor to a clean-data predictor when it + needs the original sample. The scheduler solves the linear-Gaussian velocity + relation using the noisy state. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a flow-predictor, wrap it for manual conversion: @@ -1003,6 +1079,10 @@ def score_to_flow( Examples -------- + A score predictor can supply the velocity expected by a flow-matching loss. The + scheduler first relates the score to the clean signal, then applies the schedule + derivatives. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a score-predictor, wrap it for manual conversion @@ -1073,6 +1153,9 @@ def flow_to_score( Examples -------- + A velocity predictor can supply the score needed by a denoiser. The scheduler + uses the noisy state and schedule derivatives to perform the conversion. + >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler >>> scheduler = EDMNoiseScheduler() >>> # If you have a flow-predictor, wrap it for manual conversion @@ -1190,7 +1273,9 @@ def get_denoiser( Examples -------- - Generate ODE RHS from a score-predictor: + A denoiser returns the reverse-process right-hand side used by a solver. This + first example starts from a score predictor and requests the probability-flow + ODE. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler @@ -1204,7 +1289,8 @@ def get_denoiser( >>> dx_dt.shape torch.Size([2, 3, 8, 8]) - Generate ODE RHS from an x0-predictor (score conversion is done internally): + A clean-data predictor can use the same ODE interface without producing a + score explicitly. The scheduler converts its output to the required score. >>> x0_pred = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2) # Toy x0-predictor >>> denoiser = scheduler.get_denoiser( @@ -1213,7 +1299,8 @@ def get_denoiser( >>> dx_dt.shape torch.Size([2, 3, 8, 8]) - Generate ODE RHS from an epsilon-predictor: + A noise predictor can also drive the ODE through the same factory. + The scheduler converts its output before evaluating the right-hand side. >>> eps_pred = lambda x, t: x * 0.1 # Toy epsilon-predictor >>> denoiser = scheduler.get_denoiser( diff --git a/physicsnemo/diffusion/noise_schedulers/rectified_flow.py b/physicsnemo/diffusion/noise_schedulers/rectified_flow.py index 9ec1935f1d..f21eba4539 100644 --- a/physicsnemo/diffusion/noise_schedulers/rectified_flow.py +++ b/physicsnemo/diffusion/noise_schedulers/rectified_flow.py @@ -86,8 +86,9 @@ class RectifiedFlowNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Construct noisy states for training, then initialize sampling with the same - flow parameterization: + Training interpolates clean data toward noise, while sampling starts near the noisy + endpoint. The flow predictor then supplies the ODE right-hand side under the same + time convention. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import ( @@ -154,6 +155,10 @@ def sigma_inv( r"""Identity mapping: :math:`t = \sigma`.""" return sigma + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = 1/(1+q)` for amplitude ratios ``q > 0``.""" + return 1 / (1 + q) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/trigflow.py b/physicsnemo/diffusion/noise_schedulers/trigflow.py new file mode 100644 index 0000000000..69e96759c3 --- /dev/null +++ b/physicsnemo/diffusion/noise_schedulers/trigflow.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Trigonometric interpolation between data and Gaussian noise.""" + +import math + +import torch +from jaxtyping import Float +from torch import Tensor + +from .linear_gaussian import LinearGaussianNoiseScheduler + + +class TrigFlowNoiseScheduler(LinearGaussianNoiseScheduler): + r"""Interpolate between clean data and Gaussian noise along a trigonometric path. + + Use this scheduler for TrigFlow denoising or consistency training, paired + with ``TrigFlowPreconditioner`` and the same ``sigma_data``. Diffusion time + is an angle in radians, whereas EDM uses the noise standard deviation. + Training draws diffusion times uniformly. Inference uses evenly spaced + diffusion times ending at zero. + + .. math:: + + x_t = \cos(t)x_0 + \sigma_d \sin(t)\epsilon, + \qquad \epsilon \sim \mathcal{N}(0,I). + + Parameters + ---------- + sigma_data : float, optional, default=0.5 + Data scale and noise standard deviation at :math:`t=\pi/2`. + Must be finite and positive. + t_min : float, optional, default=0.001 + Lower training diffusion time in radians, with + :math:`0 \leq t_{\min} < t_{\max}`. Use a positive value for + denoising loss weights, which diverge at zero. + t_max : float, optional, default=1.55 + Upper training diffusion time and first inference diffusion time, + in radians. Must be below ``pi / 2``. The default stays below that + endpoint in reduced precision. + + Examples + -------- + Use this scheduler when the model treats diffusion time as an angle and predicts + TrigFlow velocity. The example checks noisy-state shape, reverse-time grid size, and + the inverse signal-to-noise mapping. + + >>> scheduler = TrigFlowNoiseScheduler() + >>> t = scheduler.sample_time(2) + >>> scheduler.add_noise(torch.zeros(2, 3), t).shape + torch.Size([2, 3]) + >>> scheduler.timesteps(4).shape + torch.Size([5]) + >>> # Use interior angles for the inverse SNR check. + >>> t = torch.tensor([0.2, 0.7]) + >>> torch.allclose(scheduler.snr_inv(scheduler.snr(t)), t) + True + + Note + ---- + Coefficient functions accept angular times without clipping. The training + bounds do not limit coefficient evaluation. At zero, score conversions + and denoising weights are singular. At ``pi / 2``, conversions that divide + by the signal coefficient are singular. + """ + + def __init__( + self, sigma_data: float = 0.5, t_min: float = 0.001, t_max: float = 1.55 + ) -> None: + if not math.isfinite(sigma_data) or sigma_data <= 0: + raise ValueError(f"Expected finite sigma_data > 0, got {sigma_data}.") + if not 0 <= t_min < t_max < math.pi / 2: + raise ValueError( + f"Expected 0 <= t_min < t_max < pi/2, got {t_min=}, {t_max=}." + ) + self.sigma_data = sigma_data + self.t_min = t_min + self.t_max = t_max + + def alpha(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Signal coefficient: :math:`\alpha(t)=\cos(t)`.""" + return t.cos() + + def alpha_dot(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Signal derivative: :math:`\dot\alpha(t)=-\sin(t)`.""" + return -t.sin() + + def sigma(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Noise coefficient: :math:`\sigma(t)=\sigma_d\sin(t)`.""" + return self.sigma_data * t.sin() + + def sigma_dot(self, t: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Noise derivative: :math:`\dot\sigma(t)=\sigma_d\cos(t)`.""" + return self.sigma_data * t.cos() + + def sigma_inv(self, sigma: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`\arcsin(\sigma/\sigma_d)` for noise scales in ``[0, sigma_data]``.""" + return (sigma / self.sigma_data).asin() + + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`\arctan(1/(\sigma_d q))` for nonnegative amplitude ratios.""" + return torch.atan2(torch.ones_like(q), self.sigma_data * q) + + def sample_time( + self, + N: int, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> Float[Tensor, " N"]: + r"""Draw training diffusion times uniformly between ``t_min`` and ``t_max``. + + Parameters + ---------- + N : int + Number of training examples. + device : torch.device, optional + Device for the sampled times. + dtype : torch.dtype, optional + Floating-point type for the sampled times. + + Returns + ------- + Tensor + Training diffusion times in radians with shape :math:`(N,)`. + """ + return self.t_min + (self.t_max - self.t_min) * torch.rand( + N, device=device, dtype=dtype + ) + + def timesteps( + self, + num_steps: int, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> Float[Tensor, " N+1"]: + r"""Return evenly spaced inference diffusion times from ``t_max`` to zero. + + Parameters + ---------- + num_steps : int + Number of solver steps, at least one. + device : torch.device, optional + Device for the inference grid. + dtype : torch.dtype, optional + Floating-point type for the grid. + + Returns + ------- + Tensor + Decreasing diffusion times in radians with shape + :math:`(\mathrm{num\_steps}+1,)`. + """ + if num_steps < 1: + raise ValueError(f"Expected num_steps >= 1, got {num_steps}.") + return torch.linspace(self.t_max, 0, num_steps + 1, device=device, dtype=dtype) + + def loss_weight(self, t: Float[Tensor, " N"]) -> Float[Tensor, " N"]: + r"""Weight denoising residuals by :math:`1/(\sigma_d\sin(t))^2`. + + Use with ``TrigFlowPreconditioner`` to give unit weight to the + normalized-velocity residual. Use positive diffusion times to avoid + the singularity at zero. + + Parameters + ---------- + t : Tensor + Training diffusion times in radians with shape :math:`(N,)`. + + Returns + ------- + Tensor + Per-example denoising weights. + """ + return self.sigma(t).square().reciprocal() diff --git a/physicsnemo/diffusion/noise_schedulers/ve.py b/physicsnemo/diffusion/noise_schedulers/ve.py index 608bbdc1de..4351b8faea 100644 --- a/physicsnemo/diffusion/noise_schedulers/ve.py +++ b/physicsnemo/diffusion/noise_schedulers/ve.py @@ -55,7 +55,9 @@ class VENoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the VE noise scheduler: + The first sequence uses a variance-exploding schedule for both training and + sampling. The final predictor conversion supplies the right-hand side expected by an + ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import VENoiseScheduler @@ -105,6 +107,10 @@ def sigma_inv( r"""Inverse VE mapping: :math:`t = \sigma^2`.""" return sigma**2 + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time :math:`t = q^{-2}` for amplitude ratios ``q > 0``.""" + return self.sigma_inv(q.reciprocal()) + def sigma_dot( self, t: Float[Tensor, " *shape"], diff --git a/physicsnemo/diffusion/noise_schedulers/vp.py b/physicsnemo/diffusion/noise_schedulers/vp.py index ee6582822d..010f4c5cfd 100644 --- a/physicsnemo/diffusion/noise_schedulers/vp.py +++ b/physicsnemo/diffusion/noise_schedulers/vp.py @@ -66,7 +66,9 @@ class VPNoiseScheduler(LinearGaussianNoiseScheduler): Examples -------- - Basic training and sampling workflow using the VP noise scheduler: + The first sequence uses a variance-preserving schedule for both training and + sampling. The final predictor conversion supplies the right-hand side expected by an + ODE solver. >>> import torch >>> from physicsnemo.diffusion.noise_schedulers import VPNoiseScheduler @@ -158,6 +160,20 @@ def sigma_inv( discriminant = self.beta_min**2 + 2 * self.beta_d * a return 2 * a / (self.beta_min + torch.sqrt(discriminant)) + def snr_inv(self, q: Float[Tensor, " *shape"]) -> Float[Tensor, " *shape"]: + r"""Find diffusion time for finite amplitude ratios ``q > 0``. + + .. math:: + + A = \log(1+q^{-2}), \qquad + t = \frac{2A}{\beta_{\min}+\sqrt{\beta_{\min}^2+2\beta_d A}}. + + The rationalized root avoids cancellation and supports ``beta_d=0`` + when ``beta_min > 0``. See the base ``snr_inv`` contract. + """ + a = torch.logaddexp(torch.zeros_like(q), -2 * q.log()) + return 2 * a / (self.beta_min + (self.beta_min**2 + 2 * self.beta_d * a).sqrt()) + def timesteps( self, num_steps: int, diff --git a/physicsnemo/diffusion/preconditioners/__init__.py b/physicsnemo/diffusion/preconditioners/__init__.py index 848114c5f8..8dcdabe631 100644 --- a/physicsnemo/diffusion/preconditioners/__init__.py +++ b/physicsnemo/diffusion/preconditioners/__init__.py @@ -28,6 +28,7 @@ BaseAffinePreconditioner, EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) diff --git a/physicsnemo/diffusion/preconditioners/preconditioners.py b/physicsnemo/diffusion/preconditioners/preconditioners.py index eeac45dd74..3a5392ea02 100644 --- a/physicsnemo/diffusion/preconditioners/preconditioners.py +++ b/physicsnemo/diffusion/preconditioners/preconditioners.py @@ -580,6 +580,94 @@ def forward( return c_skip * x + c_out * F_x +class TrigFlowPreconditioner(BaseAffinePreconditioner): + r"""Convert normalized-velocity predictions into clean-data predictions. + + Use with ``TrigFlowNoiseScheduler`` and the same ``sigma_data`` for + TrigFlow denoising or consistency training. The backbone predicts velocity + divided by ``sigma_data``, rather than clean data or noise. Diffusion times + are angles in radians. A backbone can return a prediction paired with + ``DiffusionExtraOutputs``; its learned loss weight passes through unchanged. + + .. math:: + + D_\theta(x,t) = \cos(t)x - \sigma_d\sin(t)F_\theta(x/\sigma_d,t). + + Parameters + ---------- + model : physicsnemo.Module + Backbone predicting velocity divided by ``sigma_data``. + sigma_data : float, optional, default=0.5 + Data scale, matching the noise scheduler. Must be finite and positive. + + Forward + ------- + x : Tensor + Noisy states with shape :math:`(B, \ldots)`. + t : torch.Tensor or tuple[torch.Tensor, DiffusionExtraTimes] + Diffusion time in radians. Pass ``(t, extra_times)`` to include a + destination diffusion time ``r``. + condition : Tensor or TensorDict, optional + Conditioning passed to the backbone. + **model_kwargs : Any + Model-specific keyword arguments. + + Outputs + ------- + torch.Tensor or tuple[torch.Tensor, DiffusionExtraOutputs] + Clean-data prediction, optionally paired with its unchanged loss weight. + + Examples + -------- + The backbone in this example predicts zero normalized velocity, so the wrapper + returns the clean-data skip term. At zero angular time, that term equals the input. + + >>> class Velocity(Module): + ... def forward(self, x, t, condition=None): + ... return torch.zeros_like(x) + >>> model = TrigFlowPreconditioner(Velocity()) + >>> model(torch.ones(2, 3), torch.zeros(2)) + tensor([[1., 1., 1.], + [1., 1., 1.]]) + """ + + sigma_data: Float[torch.Tensor, ""] + + def __init__(self, model: Module, sigma_data: float = 0.5) -> None: + super().__init__(model) + if not math.isfinite(sigma_data) or sigma_data <= 0: + raise ValueError(f"Expected finite sigma_data > 0, got {sigma_data}.") + self.register_buffer("sigma_data", torch.tensor(float(sigma_data))) + + def compute_coefficients( + self, t: Float[torch.Tensor, " B *dims"], / + ) -> tuple[ + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + Float[torch.Tensor, " B *dims"], + ]: + r"""Return input, time, output, and skip coefficients for diffusion time. + + Parameters + ---------- + t : Tensor + Diffusion times in radians with shape :math:`(B, 1, \ldots)` + for broadcasting over the noisy states. + + Returns + ------- + tuple of Tensor + Coefficients ``(1 / sigma_data, t, -sigma_data * sin(t), cos(t))``. + """ + return ( + torch.ones_like(t) / self.sigma_data, + t, + -self.sigma_data * t.sin(), + t.cos(), + ) + + class VPPreconditioner(BaseAffinePreconditioner): r""" Variance Preserving (VP) preconditioner. diff --git a/test/diffusion/data/trigflow_precond_conv.mdlus b/test/diffusion/data/trigflow_precond_conv.mdlus new file mode 100644 index 0000000000000000000000000000000000000000..d5ba176859709653ac8d17e8b8350739fc5b4034 GIT binary patch literal 3608 zcmbtX4RBLc7Jg~^o3`|yLcxNiEXA#_`RfnZyqgIU&^OT9MH*xAktVOnqxtE}dnvNr zteuf+m9;twyQ2c4I}Yy1kDw#!ns;j*6c^mxqSGk?G6?SCZy^gY4?`S}ijOyKMBPIGq8nmvxF~T7v|?jqtVR<1r%idW?CiK(k5*46=8mXCpw2&6077Z!~DYNd`M^~o9T?q z1E2qk-i<K8MJ}En6-5iru%{b(s0k*r^BcUSOG;xs0fvWR7nw zx+Bb=804J-}a@&pAo$$uX ziK3UqKJ=N+%}z)5u}tm541!AdfH%8FYqc!>84ef>uL5lSlzP*4bdU)w(^)smt?R6o z#-Cx`Ztt7(1bnpjDL6v+O7|`NR4Q4157@j&2Ykn#mxh>(!3Ncf)TP&+hu5^PQg6Mx zgW4DvH$Bz46AWA&0QS1DG_+tZTv2=irsNlZ`xek(%c|Mn-rV=$`|s_OR@Zz3HsqK= z<+&1gbwDTm@$7BX;K%`>$~_9Un}fjkpH_IJ=qBiLt4z0;uSi|%e@hL|J4=n!|HJfy zrV=!sIZM6y@Ni`N={rGF(^jdpu+>CEixZUeH-qabre3=R1H7n88BY*FgQH7 zj>&SQUv#M?P=WY4(NI&w|=@X3DI`uYH< z&dHa~EME=oy%k0t`m6?wJi8Vy{D51^Jk$!-3?Gw9bNj$3bBNmW&mPI&R%F(l{0=^` zkB*#rrw=xq7U65QEK}o=T5$N~;~=$cA3VF^6)JaT8rW-IG+6SB6HxmGCoNn1IA#6h zH`G}74%6(sz0^IsFHn`2zMwAkuYf2fr&>FU=jDGg!A~TjY`2?8wvZZa}_w zkQzJR345zMBERJ3!T)-{kuF~u2f?zxO4|mnns$|L1GVj5aQ?{y;LFN(_)*V@G%ss4 z96Fj0PMYU}w(K#gV*kIvj-q~exou~pv;AAs;k;Qecgb_``;y^EGrbua3U^7m{XNv5 z$2Lm`?H5g%-e;&wM@~jYo_PRlX}u=hGHL;rSC)gXRz59#zS1n!ogR;n18-2vKG-GA zDS8xU{-K-Ndgvg`9=KrI?S3qJ2JR^Tx~3mp0a&Ix1C8=Cpto8YZukr=((Ix7f7c5h zzcY3YN7L7w%Ry%l%hW#PC0&Jy%k_rtKZ6FJ{^6bOyMOxrD<|}cm6OrSM__{-x`BV~ zz}1`KTviMDe?~&vpsh1%bq2k5iNUDX>U6sLdYw*R-_TH7r>`{_v`dWodZST~y=-cF z=}%QCE?Amr6?Qn$c_bGnT&?9p*@r06$pxrX2EB}G(ga#j&Xgp{MWjqdau;$nL2_N? zQt~7hf&v-O@g$x^0aC(~Z*tsWOpdjv&zW!|Ry&l~D6&ex%$=w8n;4tO#H5h zCL2NAA8&LK4mX$t#LX!hvCbt8*KUN#BCvi;7a!uuHkzpX=x?_pWh|K65p)aW1|i_@ zsbCnt!^bjA#nSl|Yl55`HLDw&gRC>)cX9Z5vB8S!`A7?OXd!Q0ugkMZ;GBY=^#w>5 z*U=He4IbGgLkY_F7!9@@HP%xR?fI}-hDjz|6*ol2{zt|qyKtnfBaPvxL(^cE&X=!i zS-{_oXj$*|OA%Q?Oj-keWfTj!&T*&5;rFv%0eNExhr{wvn_pS?Kp2G>OD>X$_BbBy zB(Ep>Qyc<4X&|)@EjHzIc|!uz!J%NsGY>0hNgWyed%9qe(_xXHs+3~&WKt-ajAl`>8C|zElyc=Uk-9rBuW~^slrTElnNXYWKB`Id76Svfk3VWkJs?Xd`WlSM^0t7^q$Z`?BzC&a0!eS+N59Ei48O7%ab)HE3d; zGgw0i4rk%{YA};-@jYF9!>r>~Bt6qGPH5nB5+}i;qR!!C79j@M1^BB8i<92iz&IFoVqoSW5)=|V zsieQ5?rx??B0`ctJUqFy-9ug>>5V;%fFUOyNLEgSQpth+49MSNtho~>e=}T9D|SKp z|D5zVb`nk{JL@}n^{JhG`|K*~{^DgjXJ@_<16>$bW%9Co!iG@_My`F(^gh|BWOd=kNMVS~wPsW+aP%M$4 znGj7gF*+K-0}>HEOE-H5mjgH^NzbHR#BgorGS{jx^1#v#u#efn=GsJCpaE>a0oYdf z?0j?SiCXw;4*7fPon69QS;j3AZFp~VFojx0vy;u0V#Gpz?u=~IDLd9&0IXO)J}5Xa ztYh)kVvP}VUR$vn2f)&HtT|n+SU=v~ow{Z0SaVK}TGn>M8FnZ0?ObyfB`n-Fy!AM+ zTR*=uX|H!Txcc1}YhY^UZufhQ9LI3`K=wJ@@c+nR4qV$AV4MM(^`{h|02yzL{Rd_t BrUC!} literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_conv_forward.pth b/test/diffusion/data/trigflow_precond_conv_forward.pth new file mode 100644 index 0000000000000000000000000000000000000000..b38b7888001a44ffc59d68b5680a878cffc1a2b5 GIT binary patch literal 4163 zcmbVPc|4VA8$Nc95K<9}B(lu0gp%_-$2RIvG_*%LPIV+Y9BHM^GDf4M(n(4oBc-yH z+)vH$kJ$vinPBi zW7-s&Ae^^!Nk~{!z{+qgCoD85Kzvyhup}%zDlk0gyP!ZppykTtA@0r&w0ZyM+9p(43W#e0eV{E#NEouxTc2S+)#Yku5La zE1Iykvl(o8wv>Rcl)&MHg@kZ90$x~X1bs$$5I39~Cb}x;!E+`pi9FV7s>@JBGy^z9Lt`TsmRvjV^7Ng*m8gS;6I%J!a zB7@xD(CwaSu()C@IaZWH_8k{umn&9y`rebsDK8$*?P?N*MU}$r2PwpFlN>dqXe86F z_5$YnsImT~EJsZ|)+BOD(!(VI_Bc(c4Ni8nK)d!@()H7;%5KB6K)1+dzIvzv{WBC{ zaZ5Bh9#n?4n#v#{Di$`)qEr53qG6{-KPh)UflKDMqs?RAliJrgq#@oD&)?xjJ*>^b z2+6WC@{_?^V1~zx-HF%VI}2lNONiOfG^pC~01I}#6gf}W0kxs!V5Khw<(0RfX-ghN zj?jg)R1Ub5JV(u=je+}%GWEE}5AxP5M$8Q=tfy|x_^Ng={2KlOx)rn0>hpcjzaWjc zj&mbZ7*DYM3?EkHuYKXAoy>elV?RXC6NbZNcMT-SJxVf|>k0i; z6r3+h0lMJI`D8Fe)Whw2$4g7exU}gbUs(a8XUNU(L8{rUVBHCzp9hGh#4%^>7 z0TKI0n1zRRh^vW$WmRCOu6d>dJNy{L>ID`(ih_JpVF4pfbj;>mf&F1i6rxgcE{N zaCTxlxv91TEY`gwGSf?8Vy6}I_nd=Evl4;r6H0Q0iV)wO0yj3!g&Ku;;?gq>4=eZ$ zN_vH`R_F)s%hHf}?|le3eh)7A=&*8|j-$rM74XDo9HlBZj7qq8wL++BM@5;;N8Jz7 zz%soEYSKTzxZBCF;#d^$YumwBUWv6Th0dy!(g&xkMVJ|XQB=wBhts|hXzzw+*g)kX zj<-yL<7V`;?-}Mf&p6OGnGczdgdk{pO{!xiP%qxf;e)wj zsLTz~DE&b_7zHF?Tk~7!nM^5;yLtqBT(~02*;Ylg1o`->+kNDq*#XkqHsQo-p_BRh z_vHQr8)nxf4^rr%NIV?tNP@>6AX`h}kg5WB+-pX(Uo}u}>O7{h*+J&9{kKrQjqM8U4PT>|~~5kIk;Qx^yFqcc>@R z8@?s(_XG%_!1)DxOpGT%@Y00oDe(Xs7ae zk=KG&WZq_qe|6G=^A&H%v|<_9Z&*z{V=p4nh(vgotb#*&RdKc7XEk#K8_VM~)g-ywU{R&g5b>rH}#WGpkSH5YZ*nj)F@Cgge4nAnV^!y+)DIPK=}?nylJ&4xxY zra&Kt8WiE(wOdh*nlUzLPlThpyT}Rocyy((2-GyHNJ{i%n0I#%Ih*+mZPm$x1E&f} zT~Qr$#orO#=Q8lJ$l=f|6vEZL2T{$ok#LgggXH#YWO7YDv|RT=njS0RoVFoV+By?P z=pKROOUnReD8P<)aWK*(79JW=DCEQexDa{>>jv5rj!z$SWgUft*_n9eiiae~`w+Ss zvIAn=ETQ!=4?FqJAkP}ERalKUgOcK=!L?0}L@Fi%3K{{WPuvdq%NxiP|F+8Bp-E(k zr$2ION`&01ZlsC3k#62Jq7$M22s&mAF`3#=f0sxAZ@77)Z^ zKtIL9_Aa1f0sHWDb#>_I+7H7`)vOmG6 zcCJn|bzu|vA=4OQq||}$$^fDD8TN$G1!5C>zs2P(CuBL^FFZ-mFKa>-fQGuP87Jo(MI1TA(Y zw1(?Md-e@*T-OM0eRh<-?I>#P2|Y4*fh*q3i-hYl^RcO`Dpt*zhnw^Cp>}o~q_)l^ ziYjkN%_Xvn--rY#n?jQ^-KZ?HyZa7b@fY$~} zoIn6C$lZD5R|zp3Y;9~OO=Q?JtgY>At*mYA?ClwLlkDs!Svy!+iys*dR(3YFHWHRM zsU_Mb;x;AbJlOjOHJ+ z{w3t+BVs!MvVKs=|319lYwNLM# SyrrDPU5U>q&?Nbva{mp|C~MXL literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_conv_sigma.pth b/test/diffusion/data/trigflow_precond_conv_sigma.pth new file mode 100644 index 0000000000000000000000000000000000000000..777d486de2633f69ce87fa46adfaf44b1f76c5ef GIT binary patch literal 1845 zcmbuA&x_MQ6vwCApGE~i5LDQO?XFg<*)&blU12R231SY`ics_rv&pmu)25S5ySg9< zOAj9P5AdRR^We#U!RxxPCvRT#YwbHEmSd}CU>bS;J zUK>!$_BhSOTLa5BHHv$l<7z{ZBMXH>D+?RgO1*i{)k?560!g=$ z>}vD}B%`J%$=+9yni`Runi!E2uMNYoZER4>>3Q5e*TgP1Z(9cC2zRAPPEvBJNltrw zL(S%-*3O%cObW6ww`)aoE3#{R+i|-()rQ#hSW%_E`5(<=T`vLJm*3xa;b|e~qwy#o zjW%tlrrTDh*Tp@m)R*GuWC;Lj+vL{*D-JmyU=NP>BUpgkhhp%UApWoi8CHFJ6-F&VFBnQLx-7|~=az05O zEk%;d9C9%VrOZLay?335@%TDx&Bv+L}9vha$dKy1T4av_(BuYkgE@wJp@0NgyO~anG)E&YXKQ`M>}B z?mz#`{F5TLv-hG%%i5K+bi(zp>-$_NigYJDR#=ovYXrq?3bw){WXPuRDGy4i%TcLI z>F``7s}vL$XcA>(s0^!CF(N0YNTZ>1Sap$B@4Qf#OY7*|x#}F&nW3hnGQ6;KXY5$7 z&M#Cdh#O`^4qJrX3KS)*TBCPX=;-VcwI)}=(pr6yPEi`+%-Ds6r8^Q7wp8ZRjJ*ls zoTw>MvTvCNof(I89<~uOBh&5hGpD$Oag_5|#%Yn9nlfdRim;Dj8Rt}HbUGg!n7&v` zvw|{Oov&gU7Zv7F*{e)`M2xG-F^R$s#$X3>zMQ|Fal;SZdWLUu=FT!6sf=g(7%OMl zYNej>vT_x{c&i-ce3lt&^i;|C$a$5FugS|emhnqv{7cH2@#&+mCq*I6Dvg1eV3ZsX zmX1d$Ca|2DXbeIrt<&S!GLz&yeY(RSHxpFO1gnf=#Ebdkn2>TN)Hq(E&laS}9i4t^ zY+JJjM~d{}XWNIGV4iJ3UXiX)$tnhz6J_IUCfyAxRIEj^x-v(p@3Nfne}jCqAcz;1Vlz2g@nD^Zsjh#kSWq%#V)mG7bVJO zzXp3oMy3B2R8>)ny5tqyve4_Ccfvbx-5VlUczz$(s&H(H3L3$8hs%QHusUfReAl%dwSRGlOAfyem%B(|$W3q5 zUnAmv*E0*WbRB^_*K=@}R0pGePD5urC!pMXo@A=Dk1Ni4A9M!xfUcO^lIPALaQ4?d z;KYVb!>(&@!i5XBaXubt62Vsv@a;5DX|w%Fw8!@x+PiQ%y3QzIulEMnHeL+)*9thN zZD+Vl@6m8<>RGO?^cEMg0>F~SR@jh`A-&Rn8U;V!jat9I1N(%oNb;FGlB_L&GsA1( zcE1C#Gs=LflR~7evmbNG7d{22HeW*2Nf*?0u$N2gdcj@Tx)G)Qr3OxS;d5UnB*S-S zAjA3_;jnAx5)|5@j#I2pgNr-Qb3U$>@Uo&A9Q<|_r_J<~impCFTMo$$-KQ#1>@^l0 zOCKefeI^>VeflM|_dA4omLCSLV;x|F^o{aeI~>kONsM;kuQFNK0Hn!v+b1*l?ro?)6g z5dEUL&pqqwg*v~BTwP1Qq~50vMrUi_t*uAkcOlv6Z>zevz){Jl^&B5wm5zs*qaT8x z#vkEs&yDC=<{m>q_5(?qn=f*mxflK5-DyaXuS1a@^_-}275Mz&I<6_>u|#}nCHTa> z4#u?9f{;gT;ONaH?%3KN4FCM=T;O>y0j&p5Q1T@Quy;We@(?zdJTe_;~>7Hp7ZytLQWr+ zf^E%B=tAiR!=m?^xj8-U=wj`kx$Rk1F#MB|;E=ilo*F$Lmbq7<4>rf52eOmUY2s!u zwP`9$-SR&2lYE43czulS#k-+hzjcDTw=MzbcouA}4+kseX2HPzt(;%muVH59M~1!& z)u^%g47e8I3(r+efdQ%=Xxgv*;R+8IB)we+gATS!n)a4LO<)>qEvrZJeXC(%d_Q+8 zZa*s6#KOb{`{A9Esc6SRKKwSR7QBDV$?&gx8&E>}cZTMUMlL$A10E621uYv|LDDQQ z?tWY_oa3;C%bNKoG^fCwYd1KdX?+f;x;hiu^K;On6HB1S)>?R=pa(TfHlTl;_y&#a z-zhodR)qqV9*5CuZWtysc*37%&EQrBt~E@$w+?O(Rl&8TPdJ~u%h2i=FK}VY6WH6+ z4#?iKz_o3mapS?8&#U$xd+WU?B2h+Bp$U@=vKdLJueqO7_@3Z??(l|A&ubC5=)rj| zpT~92YBcO$(*$)Fec|nxXWTsIbM8v_AJO-}T!x_{dnw_c_u=Ueqk zQY~3G+d9w|ZCL>LSV14)nqUxZQ&-qJu{^~6teh;uQq4|YsvKr3Pn)d z#Gu7Gbw0i!XiPgIxnw;RN8`K zE0o^$x5*7*Hj=FhG8W&2R-v$Tm^45nF#V6ohi*aGS|o@nHwt2@V@18sm4|V#~y&dOA2%eNY#W&n6Dv8)+t^S=W3E%UzIDCiD_@?Chmjz<>CBJ>!a zP#IY*pjqQj3U=gOcuWL#<~@X6twAr>fne8l4sacu>nZecF7oJJ5d>A1MFB~wqEpwj zqNs_*%aVX7NS)8Y0|*|>!9xv2fmc*X7vbRw_E7Ou42>u}QfaVzn(SVK?O&-7L?41j zbFe?(*_fnSLvX+Xh`~VvaI6Uk7=R%Ia2&xCIXIm6HURp(RE0qcfrcjy!bp=aY7oW@ z!YKq#=ir%IILmeU-Yu`cs^ZllJ==2s;~?y8|HCC!TPPrYQPgz2hMDw*Ho_oE41L`E z`Izh$>8%YV!H`2kTam#RWT~ty;2_PU2U>Vy;=};E1b@4+IO(kogu$3I19SI~sFFa+ zq=PMVxAR1n5RygW;mM_q0NG&DTYDG>qs}~#telEu(!+ZjP<}A1xf3UUvV5QiyC8ip zCxdM`iDuG$yE*ynzQ}I&1lZ8>ZT967V2S?JhX48h){nK&_Xd>u*)Lb>?;k#^=SmZG zq945aHU7W#=@$CGZn7-9I?H}oImfQ6cN_Ef2}9lc>ra3IXY1e0OaCq#KvQJK-`>vL#px>EdxDkB1s7W{j?-xjA<)vaK-MiC3G)||OB%Mk|m;}wxvFRi|9gQ;aL_88> zm{bb?W)g{5g1CHL`5q17m?S;RxrpKB$z`r-ljMb^wRMi!(dOzDw_r)VaD=v%NV|f$ z{3I<3FB~iEsCjk~bA1V0MEv-KbTo}-U9+>z6=U4OUUkYg>XjXAE(KQb2);5nGTi{< zW5yaU<{Y|k^9TEQdu{CDn{wu8+%I$>GkHK*XMN?<3O^J>OI_T$6Ik>%Yp yKFfK%+g5{~e`!f=iEKaLk}Dc+pUByV8~%=5@W7p|0mi+cv+b`v3~(N60-mR literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_linear_forward.pth b/test/diffusion/data/trigflow_precond_linear_forward.pth new file mode 100644 index 0000000000000000000000000000000000000000..d1a46ac1ab65fee1f8d8e4c8afce0f14abb4a0a8 GIT binary patch literal 2129 zcmWIWW@cev;NW1u02&Nd3?)UG>1jFn_FwkK$VHEZzcbg&f`j-i&R9oE^!@`8heM$t9WjdBt4*MJcI8sVOd*$t7HcT$(|J z+$DuPL4~}pdk;ik`EO)%W_fB~4BSQ1~BT2u_Q*v-kD%;1YM z1UYQ``#a4b|AH`XhZ!-*6LuIhSmP6OGSl;NQ}asPoTTx&&CCSk0Q>U7mrOv*K^V6K zG#S(hIv}aCBsCut3vN!z`0Rxxww%n8lAKhK8(M82OR@p&2VvZ9P-9Rd=mt=#);Dl- zQX;|uwm0e=gc}4M*K^7`m}u_Xr=;@NJ|b$xfsG4A>>o$2winkqxc{#5e0!CpPxfb- z7TTAmraOrLdb!$D-*TDzLEZTl_KZS1)|Z`qe2`gi})UzK(YFFp4x%St|=^}y=D^I0Dcl--YY zC`dYR;Qb`k1NY~dI@Cp~A4qeaaiD``>VYZxoAxz`so2JRcX0?|JYc`6VZ{M?!FdOG zR9y}@ED^MozIt>26r1DrYdQ8GP*XKB81nG^piV^4gOh+U7sv${=cR`d zPJd<$QUqO~2af#W)S^sa4yZ~56$|m1DQ-?;I2>ePVPb4zX=-3@U}R)wVrXP+W^Qg^ zYHng=WME-wW@=_+X=ZF{Y+?X%rPBeCb`40+L4s#FoS8rXm*% z`Y1+P0ON{~sc5AXy5Y!~-U!9;i^Ld?QPQBBkDQ@3QOpkl1|1>C!wV90W06y<8j7*k zh%gpXHlUk{oc`2N%nSh5u7sQlP4(!8A}1F!6hki&VJP;59N^8yrUNxgj#(FODyT*U l0T^8lPMi!36|4|(pgyqFd#D0PO~VQ*Yk_*$89?eGY5{BgQU(A3 literal 0 HcmV?d00001 diff --git a/test/diffusion/data/trigflow_precond_linear_sigma.pth b/test/diffusion/data/trigflow_precond_linear_sigma.pth new file mode 100644 index 0000000000000000000000000000000000000000..76f4b47c540905b17c49eb98282cae091f29f0e0 GIT binary patch literal 1859 zcmbW2O^ee&7{{mYH7YC%f(2I)+TAX=ZQ7=3*B1goa>>D35sDsSHkqx1Y0^oiTU`+Z z=|Mc{2k_`u5d0SF!IKx+lc)U(&ZOO%WKFZ|0}mn1%>Vi2>HnGZ)JztDViBB)O>hqo z7Y+9Y z7A@6u3k~fNw22~btfmtM1|UuGpENDbG9F+HY$pKhy(Z9M|AHb{R@XhXn}rCmAY+)_ zKQ85IeU|YIWl&9EV6Lw7id;Sy!1weW>13@WRo^_5K2L$fNb1MyAJ-@a%XpFs*hnB* zI;0ej(&AXfqc6iJ$^93=jf9Fd|K&aCBI~#2cr`;9#cC*ph?Y6~6qbvG0tzTR3RcTLH_S6&uUYDpuuLUF!xD5@D8i=B|3 zNGD~rW@JjE(m})XA?c%J=`&LnRS#NuE2Q41O6=Rf+* T&J!B*vo+`673eOue!liEHA&0R literal 0 HcmV?d00001 diff --git a/test/diffusion/test_model_io.py b/test/diffusion/test_model_io.py index 4e4681b124..b1afdf679a 100644 --- a/test/diffusion/test_model_io.py +++ b/test/diffusion/test_model_io.py @@ -24,6 +24,7 @@ from physicsnemo.diffusion.preconditioners import ( EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) @@ -36,8 +37,9 @@ (VEPreconditioner, {}), (VPPreconditioner, {}), (IDDPMPreconditioner, {"M": 32}), + (TrigFlowPreconditioner, {"sigma_data": 1.4}), ) -IDS = ("edm", "ve", "vp", "iddpm") +IDS = ("edm", "ve", "vp", "iddpm", "trigflow") TIME_FIELDS = (None, (), ("r",), ("r1",), ("r2",), ("r1", "r2")) TIME_IDS = ("tensor", "empty", "r", "r1", "r2", "r1_r2") diff --git a/test/diffusion/test_noise_schedulers.py b/test/diffusion/test_noise_schedulers.py index 13fb98c917..20dbf741f7 100644 --- a/test/diffusion/test_noise_schedulers.py +++ b/test/diffusion/test_noise_schedulers.py @@ -289,6 +289,10 @@ def test_concrete_snr(self, device): t = torch.tensor([0.5, 2.0], device=device) assert torch.allclose(s.snr(t), 1 / t, atol=1e-7) + def test_optional_snr_inverse(self): + with pytest.raises(NotImplementedError, match="snr_inv"): + _MinimalScheduler().snr_inv(torch.ones(2)) + def test_concrete_add_noise(self, device): s = _MinimalScheduler() x0 = make_input((2, 3, 8, 8), seed=3, device=device) diff --git a/test/diffusion/test_preconditioners.py b/test/diffusion/test_preconditioners.py index beaccfda5d..8bee845b29 100644 --- a/test/diffusion/test_preconditioners.py +++ b/test/diffusion/test_preconditioners.py @@ -27,6 +27,7 @@ BaseAffinePreconditioner, EDMPreconditioner, IDDPMPreconditioner, + TrigFlowPreconditioner, VEPreconditioner, VPPreconditioner, ) @@ -113,6 +114,7 @@ def forward( # Preconditioner configurations for parameterized tests PRECOND_CONFIGS = [ + (TrigFlowPreconditioner, {"sigma_data": 1.4}, "trigflow_precond"), ( VPPreconditioner, {"beta_d": 19.9, "beta_min": 0.1, "M": 2000}, @@ -403,7 +405,7 @@ def test_single_element_sequence_stored_as_scalar(self, simple_model): @pytest.mark.parametrize( "precond_cls,precond_kwargs,precond_name", PRECOND_CONFIGS, - ids=["VP", "VE", "iDDPM", "EDM"], + ids=["TrigFlow", "VP", "VE", "iDDPM", "EDM"], ) class TestNonRegression: """Non-regression tests parameterized across all preconditioner types.""" @@ -611,7 +613,7 @@ def create_fn(): @pytest.mark.parametrize( "precond_cls,precond_kwargs,precond_name", PRECOND_CONFIGS, - ids=["VP", "VE", "iDDPM", "EDM"], + ids=["TrigFlow", "VP", "VE", "iDDPM", "EDM"], ) class TestAllPreconditioners: """Tests that apply to all preconditioner types.""" diff --git a/test/diffusion/test_trigflow.py b/test/diffusion/test_trigflow.py new file mode 100644 index 0000000000..8a32f9acee --- /dev/null +++ b/test/diffusion/test_trigflow.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Check continuous inverse SNR and TrigFlow coordinates.""" + +import pytest +import torch + +from physicsnemo.diffusion import ( + DiffusionExtraTimes, + preconditioners, +) +from physicsnemo.diffusion import noise_schedulers as schedules + +from .helpers import StructuredDiffusionModel + +CONFIGS = ( + ("EDMNoiseScheduler", {}), + ("EDMLogUniformNoiseScheduler", {}), + ("VENoiseScheduler", {}), + ("VPNoiseScheduler", {}), + ("VPNoiseScheduler", {"beta_d": 0.0}), + ("IDDPMNoiseScheduler", {"M": 32}), + ("RectifiedFlowNoiseScheduler", {}), + ("TrigFlowNoiseScheduler", {"sigma_data": 0.5}), + ("TrigFlowNoiseScheduler", {"sigma_data": 1.4}), +) +IDS = ( + "edm", + "edm-uniform", + "ve", + "vp", + "vp-constant", + "iddpm", + "rf", + "trig", + "trig-scaled", +) + + +@pytest.mark.parametrize("name,kwargs", CONFIGS, ids=IDS) +class TestConsistency: + def test_snr_inverse(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.1, 0.4, 0.8], dtype=torch.float64, device=device) + recovered, tangent = torch.func.jvp( + lambda t: scheduler.snr_inv(scheduler.snr(t)), + (t,), + (torch.ones_like(t),), + ) + torch.testing.assert_close(recovered, t) + torch.testing.assert_close(tangent, torch.ones_like(t)) + + def test_coefficient_derivatives(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.1, 0.4, 0.8], dtype=torch.float64, device=device) + for coefficient, derivative in ( + (scheduler.alpha, scheduler.alpha_dot), + (scheduler.sigma, scheduler.sigma_dot), + ): + _, actual = torch.func.jvp(coefficient, (t,), (torch.ones_like(t),)) + torch.testing.assert_close(actual, derivative(t)) + + +@pytest.mark.parametrize("sigma_data", (0.5, 1.4)) +class TestGradientFlow: + def test_trigflow(self, sigma_data): + scheduler = schedules.TrigFlowNoiseScheduler(sigma_data=sigma_data) + model = preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data=sigma_data + ).double() + x = torch.randn(2, 3, dtype=torch.float64, requires_grad=True) + t = torch.tensor([0.4, 0.8], dtype=torch.float64, requires_grad=True) + + def forward(x, t): + prediction, extras = model( + x, (t, DiffusionExtraTimes(r=t / 2, batch_size=[2])) + ) + return ( + prediction, + extras.logvar, + scheduler.snr_inv(scheduler.snr(t)), + ) + + assert torch.autograd.gradcheck(forward, (x, t)) + + +class TestNumerics: + @pytest.mark.parametrize( + "dtype,ratio", + ( + (torch.float16, 0.003), + (torch.float32, 1e-20), + (torch.float64, 1e-200), + ), + ) + def test_vp_small_ratio(self, dtype, ratio, device): + scheduler = schedules.VPNoiseScheduler() + q = torch.tensor([ratio], dtype=dtype, device=device, requires_grad=True) + result = scheduler.snr_inv(q) + (gradient,) = torch.autograd.grad(result.sum(), (q,)) + q64 = q.detach().double() + a = torch.logaddexp(torch.zeros_like(q64), -2 * q64.log()) + root = (scheduler.beta_min**2 + 2 * scheduler.beta_d * a).sqrt() + expected = 2 * a / (scheduler.beta_min + root) + expected_gradient = -2 / (q64 * (1 + q64.square()) * root) + torch.testing.assert_close(result, expected.to(dtype)) + torch.testing.assert_close(gradient, expected_gradient.to(dtype)) + + +@pytest.mark.parametrize("sigma_data", (0.5, 1.4)) +class TestConstructor: + def test_trigflow(self, sigma_data, device): + scheduler = schedules.TrigFlowNoiseScheduler( + sigma_data=sigma_data, t_min=0.1, t_max=1.3 + ) + t = scheduler.sample_time(128, device=device, dtype=torch.float64) + assert torch.all((t >= 0.1) & (t <= 1.3)) + grid = scheduler.timesteps(8, device=device) + assert grid.shape == (9,) + torch.testing.assert_close(grid, torch.linspace(1.3, 0, 9, device=device)) + model = preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data=sigma_data + ).to(device) + x = torch.randn(2, 3, device=device) + t = torch.tensor([0.4, 0.8], device=device) + result, extras = model(x, t) + raw, raw_extras = model.model(x / sigma_data, t) + expected = t.cos()[:, None] * x - sigma_data * t.sin()[:, None] * raw + torch.testing.assert_close(result, expected) + torch.testing.assert_close(extras.logvar, raw_extras.logvar) + torch.testing.assert_close( + scheduler.loss_weight(t), 1 / (sigma_data * t.sin()) ** 2 + ) + + def test_zero_time(self, sigma_data): + scheduler = schedules.TrigFlowNoiseScheduler(sigma_data=sigma_data) + t = torch.zeros(2) + torch.testing.assert_close(scheduler.alpha(t), torch.ones_like(t)) + torch.testing.assert_close(scheduler.sigma(t), t) + torch.testing.assert_close(scheduler.alpha_dot(t), t) + torch.testing.assert_close( + scheduler.sigma_dot(t), torch.full_like(t, sigma_data) + ) + + +@pytest.mark.parametrize("name,kwargs", CONFIGS, ids=IDS) +class TestCompile: + def test_inverse(self, name, kwargs, device): + scheduler = getattr(schedules, name)(**kwargs) + t = torch.tensor([0.2, 0.7], device=device, requires_grad=True) + fn = torch.compile( + lambda t: scheduler.snr_inv(scheduler.snr(t)), + backend="aot_eager", + fullgraph=True, + ) + result = fn(t) + with torch._dynamo.config.patch(error_on_recompile=True): + torch.testing.assert_close(fn(t), result) + torch.testing.assert_close(result, t) + (gradient,) = torch.autograd.grad(result.sum(), (t,)) + torch.testing.assert_close(gradient, torch.ones_like(t)) + + +class TestDiagnostics: + @pytest.mark.parametrize( + "kwargs", + ( + {"sigma_data": 0}, + {"sigma_data": float("nan")}, + {"t_min": -0.1}, + {"t_max": 1.6}, + {"t_min": 0.7, "t_max": 0.5}, + ), + ) + def test_invalid_scheduler(self, kwargs): + with pytest.raises(ValueError): + schedules.TrigFlowNoiseScheduler(**kwargs) + + @pytest.mark.parametrize("sigma_data", (0, -1, float("nan"))) + def test_invalid_preconditioner(self, sigma_data): + with pytest.raises(ValueError): + preconditioners.TrigFlowPreconditioner( + StructuredDiffusionModel(), sigma_data + )