scheduler_flow_bridge.py

"""ReFlowSET's Design-B flow bridge and its explicit-Euler solver.

Forward (training) process, with ``eps ~ N(0, I)`` and ``z_e`` the EO latent::

    z_t = (1 - t) * eps + t * z_e          (bridge.py:311, sigma_b = 0)
    u*  = z_e - eps                        (bridge.py:328 at sigma_b = 0)

Sampling starts from ``z_0 ~ N(0, I)`` and integrates the predicted velocity
with explicit Euler on a uniform grid ``linspace(0, t_end, nfe + 1)``
(bridge.py:519, 536).  The bridge is deterministic: ``sigma_b = 0``, so no
stochastic term ever executes, and the only randomness in a sample is the
initial noise draw.

**Time direction.**  ``t = 0`` is NOISE and ``t = 1`` is DATA, and the solver
integrates ``t`` **ascending** (bridge.py:86-88).  That is the opposite of
`diffusers`' ``sigma`` convention: setting ``sigma := 1 - t`` recovers
``FlowMatchEulerDiscreteScheduler``'s interpolation, but then this bridge's
velocity is the **negative** of the diffusers flow-matching target and the
network must still be fed ``1 - sigma``.  This scheduler keeps ReFlowSET's own
sign and direction so neither flip is needed; ``timesteps`` therefore *increase*
from 0 towards 1, unlike every noise-schedule scheduler in `diffusers`.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Union

import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput


@dataclass
class FlowBridgeSchedulerOutput(BaseOutput):
    """Output of :meth:`FlowBridgeScheduler.step`.

    Args:
        prev_sample: The bridge state at the next time on the grid.
    """

    prev_sample: torch.Tensor


class FlowBridgeScheduler(SchedulerMixin, ConfigMixin):
    """Explicit-Euler solver for ReFlowSET's Design-B flow bridge.

    Args:
        t_end: End time of the integration grid (1.0 — the EO endpoint).  The
            model is evaluated at ``linspace(0, t_end, nfe + 1)[:-1]`` and the
            final Euler step lands on ``t_end``; the network is never queried at
            ``t = t_end``.
    """

    order = 1

    @register_to_config
    def __init__(self, t_end: float = 1.0) -> None:
        if not 0.0 < t_end <= 1.0:
            raise ValueError(f"t_end must lie in (0, 1], got {t_end}")
        self._grid: Optional[torch.Tensor] = None
        self._step_index: Optional[int] = None
        self.num_inference_steps: Optional[int] = None

    @property
    def timesteps(self) -> torch.Tensor:
        """The ``nfe`` bridge times at which the model is evaluated, ascending."""
        if self._grid is None:
            raise ValueError("call set_timesteps() before reading timesteps")
        return self._grid[:-1]

    @property
    def step_index(self) -> Optional[int]:
        """Index of the next grid interval; ``None`` until the first :meth:`step`."""
        return self._step_index

    def set_timesteps(
        self,
        num_inference_steps: int,
        device: Optional[Union[str, torch.device]] = None,
    ) -> None:
        """Build the uniform grid ``linspace(0, t_end, num_inference_steps + 1)``.

        Args:
            num_inference_steps: NFE — the number of velocity evaluations.
                50 reproduces the paper's main results; 4 is the efficiency
                operating point.
            device: Device the grid is built on.

        There is no shift, no dynamic shifting, no Karras or exponential
        spacing, and no timestep-spacing option: the reference solver uses a
        plain uniform grid (bridge.py:519).
        """
        if num_inference_steps < 1:
            raise ValueError(f"num_inference_steps must be >= 1, got {num_inference_steps}")
        self.num_inference_steps = num_inference_steps
        self._grid = torch.linspace(
            0.0, self.config.t_end, num_inference_steps + 1, device=device, dtype=torch.float32
        )
        self._step_index = 0

    def step(
        self,
        model_output: torch.Tensor,
        timestep: Union[float, torch.Tensor],
        sample: torch.Tensor,
        return_dict: bool = True,
    ) -> Union[FlowBridgeSchedulerOutput, tuple[torch.Tensor]]:
        """One explicit-Euler step: ``z + (t_next - t_cur) * v`` (bridge.py:536).

        Args:
            model_output: The predicted velocity ``dz/dt`` at ``timestep``,
                already classifier-free-guided by the caller.
            timestep: The current bridge time.  Present for API compatibility and
                checked against the grid; the step size comes from the grid.
            sample: The current bridge state.
            return_dict: Return a :class:`FlowBridgeSchedulerOutput` instead of a
                tuple.

        Steps must be taken in order, starting from the first entry of
        :attr:`timesteps`.
        """
        if self._grid is None or self._step_index is None:
            raise ValueError("call set_timesteps() before step()")
        if self._step_index >= self.num_inference_steps:
            raise ValueError(
                f"already took {self.num_inference_steps} steps; call set_timesteps() again"
            )
        t_cur, t_next = self._grid[self._step_index], self._grid[self._step_index + 1]
        if not torch.isclose(torch.as_tensor(timestep, dtype=torch.float32).to(t_cur.device), t_cur):
            raise ValueError(
                f"step {self._step_index} expects timestep {t_cur.item()}, got {float(timestep)}; "
                "the flow bridge must be integrated in ascending grid order"
            )

        # The state is carried in float32 even if the model ran lower (bridge.py:515-517).
        dtype = sample.dtype if sample.dtype in (torch.float32, torch.float64) else torch.float32
        prev_sample = sample.to(dtype) + (t_next - t_cur) * model_output.to(dtype)
        prev_sample = prev_sample.to(sample.dtype)

        self._step_index += 1
        if not return_dict:
            return (prev_sample,)
        return FlowBridgeSchedulerOutput(prev_sample=prev_sample)

    def add_noise(
        self,
        original_samples: torch.Tensor,
        noise: torch.Tensor,
        timesteps: torch.Tensor,
    ) -> torch.Tensor:
        """The training-side bridge state ``z_t = (1 - t) * eps + t * z_e`` (bridge.py:311).

        Args:
            original_samples: The EO latent ``z_e`` (the ``t = 1`` endpoint).
            noise: ``eps ~ N(0, I)`` (the ``t = 0`` endpoint).
            timesteps: Bridge times in ``[0, 1]``, broadcastable over the batch.
        """
        t = timesteps.to(original_samples.device, original_samples.dtype)
        t = t.view(-1, *([1] * (original_samples.ndim - 1)))
        return (1.0 - t) * noise + t * original_samples

    def get_velocity(
        self,
        sample: torch.Tensor,
        noise: torch.Tensor,
        timesteps: torch.Tensor,
    ) -> torch.Tensor:
        """The training target ``u* = z_e - eps`` (bridge.py:328 at ``sigma_b = 0``).

        Constant along the path, hence independent of ``timesteps``; the argument
        is kept for `diffusers` API compatibility.
        """
        del timesteps
        return sample - noise