Adopt proposed FastGen loss hooks conditionally

.importlinter

@@ -113,11 +113,13 @@ exhaustive = True
 containers=
    physicsnemo.diffusion
 layers = 
+   distillation
    generate
    samplers | metrics
-   multi_diffusion | preconditioners | guidance
+   multi_diffusion | preconditioners | guidance | _fastgen
    noise_schedulers
-   utils | base
+   utils
+   base
 
 [importlinter:contract:physicsnemo-sym]
 name = Control Internal Dependencies in PhysicsNeMo sym

CHANGELOG.md

@@ -15,6 +15,37 @@ 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 alpha distillation losses and an optional strategy layer in
+  `physicsnemo.diffusion.distillation` for sCM, MeanFlow, DMD2, and f-distill.
+  Use `DistillationLoss` or `WeightedDistillationLoss` in your own training
+  loop, or let `DistillationStrategy` manage the optimizer loop. Define custom
+  losses through `BaseDistillationLoss` without depending on FastGen.
+  Use `CheckpointStateBundle` to resume models, optimizers, and loss statistics.
+  These alpha APIs may change without backward compatibility. Distillation is
+  opt-in, with optional dependencies available through extras. Existing diffusion
+  training and inference remain unchanged, and distilled models continue to use
+  PhysicsNeMo inference APIs.
+- 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 `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.

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`
 ~~~~~~~~~~~~~~~~~
 

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
-

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
 ----------------
 

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`
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/api_diffusion.rst

@@ -13,6 +13,7 @@ PhysicsNeMo Diffusion
    api/diffusion/noise_schedulers.rst
    api/diffusion/preconditioners.rst
    api/diffusion/metrics.rst
+   api/diffusion/distillation.rst
    api/diffusion/samplers.rst
    api/diffusion/multi_diffusion.rst
-   api/diffusion/utils.rst
+   api/diffusion/utils.rst

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,
+)