Code update

PiperOrigin-RevId: 600584870

swirl_dynamics/projects/debiasing/rectified_flow/main.py

@@ -23,7 +23,6 @@
 from ml_collections import config_flags
 import optax
 from orbax import checkpoint
-from swirl_dynamics.lib.diffusion import unets
 from swirl_dynamics.projects.debiasing.rectified_flow import data_utils
 from swirl_dynamics.projects.debiasing.rectified_flow import models
 from swirl_dynamics.projects.debiasing.rectified_flow import trainers
@@ -106,7 +105,7 @@ def main(argv):
   )
 
   # Setting up the neural network for the flow model.
-  flow_model = unets.PreconditionedDenoiser(
+  flow_model = models.RescaledUnet(
       out_channels=config.out_channels,
       num_channels=config.num_channels,
       downsample_ratio=config.downsample_ratio,
@@ -117,7 +116,6 @@ def main(argv):
       use_attention=config.use_attention,
       use_position_encoding=config.use_position_encoding,
       num_heads=config.num_heads,
-      sigma_data=1.0,  # standard deviation of the entire dataset.
   )
 
   model = models.ReFlowModel(

swirl_dynamics/projects/debiasing/rectified_flow/models.py

@@ -44,6 +44,7 @@
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
+from swirl_dynamics.lib.diffusion import unets
 from swirl_dynamics.templates import models
 from swirl_dynamics.templates import trainers
 
@@ -218,3 +219,42 @@ def _flow(
       )
 
     return _flow
+
+
+class RescaledUnet(unets.UNet):
+  """Rescaled flow model.
+
+  Attributes:
+    time_rescale: Factor for rescaling the time, which normally is in [0, 1] and
+      the input to the UNet which is a noise level, which has a much wider range
+      of values.
+  """
+
+  time_rescale: float = 1000.0
+
+  @nn.compact
+  def __call__(
+      self,
+      x: Array,
+      sigma: Array,
+      cond: dict[str, Array] | None = None,
+      *,
+      is_training: bool,
+  ) -> Array:
+    """Runs rescaled Unet with noise input."""
+    if sigma.ndim < 1:
+      sigma = jnp.broadcast_to(sigma, (x.shape[0],))
+
+    if sigma.ndim != 1 or x.shape[0] != sigma.shape[0]:
+      raise ValueError(
+          "sigma must be 1D and have the same leading (batch) dimension as x"
+          f" ({x.shape[0]})!"
+      )
+
+    time = sigma * self.time_rescale
+
+    f_x = super().__call__(
+        x, time, cond, is_training=is_training
+    )
+
+    return f_x