content: """
  Adapted from https://github.com/huggingface/diffusers/blob/v0.30.3/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py.
  """
  
  import math
  from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  from diffusers.configuration_utils import ConfigMixin, register_to_config
  from diffusers.schedulers.scheduling_utils import SchedulerMixin
  from diffusers.utils import BaseOutput, logging
  from torch.distributions import LogisticNormal
  
  logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  # TODO: may move to training_utils.py
  def compute_density_for_timestep_sampling(
      weighting_scheme: str,
      batch_size: int,
      logit_mean: float = 0.0,
      logit_std: float = 1.0,
      mode_scale: float = None,
  ):
      if weighting_scheme == "logit_normal":
          # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
          u = torch.normal(
              mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu"
          )
          u = torch.nn.functional.sigmoid(u)
      elif weighting_scheme == "logit_normal_dist":
          u = (
              LogisticNormal(loc=logit_mean, scale=logit_std)
              .sample((batch_size,))[:, 0]
              .to("cpu")
          )
      elif weighting_scheme == "mode":
          u = torch.rand(size=(batch_size,), device="cpu")
          u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
      else:
          u = torch.rand(size=(batch_size,), device="cpu")
      return u
  
  
  def compute_loss_weighting(weighting_scheme: str, sigmas=None):
      """
      Computes loss weighting scheme for SD3 training.
  
      Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
  
      SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
      """
      if weighting_scheme == "sigma_sqrt":
          weighting = (sigmas**-2.0).float()
      elif weighting_scheme == "cosmap":
          bot = 1 - 2 * sigmas + 2 * sigmas**2
          weighting = 2 / (math.pi * bot)
      else:
          weighting = torch.ones_like(sigmas)
      return weighting
  
  
  @dataclass
  class RectifiedFlowSchedulerOutput(BaseOutput):
      """
      Output class for the scheduler's `step` function output.
  
      Args:
          prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
              Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
              denoising loop.
      """
  
      prev_sample: torch.FloatTensor
  
  
  class RectifiedFlowScheduler(SchedulerMixin, ConfigMixin):
      """
      The rectified flow scheduler is a scheduler that is used to propagate the diffusion process in the rectified flow.
  
      This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
      methods the library implements for all schedulers such as loading and saving.
  
      Args:
          num_train_timesteps (`int`, defaults to 1000):
              The number of diffusion steps to train the model.
          timestep_spacing (`str`, defaults to `"linspace"`):
              The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
              Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
          shift (`float`, defaults to 1.0):
              The shift value for the timestep schedule.
      """
  
      _compatibles = []
      order = 1
  
      @register_to_config
      def __init__(
          self,
          num_train_timesteps: int = 1000,
          shift: float = 1.0,
          use_dynamic_shifting: bool = False,
      ):
          # pre-compute timesteps and sigmas; no use in fact
          # NOTE that shape diffusion sample timesteps randomly or in a distribution,
          # instead of sampling from the pre-defined linspace
          timesteps = np.array(
              [
                  (1.0 - i / num_train_timesteps) * num_train_timesteps
                  for i in range(num_train_timesteps)
              ]
          )
          timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
  
          sigmas = timesteps / num_train_timesteps
          if not use_dynamic_shifting:
              # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
              sigmas = self.time_shift(sigmas)
  
          self.timesteps = sigmas * num_train_timesteps
  
          self._step_index = None
          self._begin_index = None
  
          self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
  
      @property
      def step_index(self):
          """
          The index counter for current timestep. It will increase 1 after each scheduler step.
          """
          return self._step_index
  
      @property
      def begin_index(self):
          """
          The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
          """
          return self._begin_index
  
      # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
      def set_begin_index(self, begin_index: int = 0):
          """
          Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
  
          Args:
              begin_index (`int`):
                  The begin index for the scheduler.
          """
          self._begin_index = begin_index
  
      def _sigma_to_t(self, sigma):
          return sigma * self.config.num_train_timesteps
  
      def _t_to_sigma(self, timestep):
          return timestep / self.config.num_train_timesteps
  
      def time_shift_dynamic(self, mu: float, sigma: float, t: torch.Tensor):
          return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
  
      def time_shift(self, t: torch.Tensor):
          return self.config.shift * t / (1 + (self.config.shift - 1) * t)
  
      def set_timesteps(
          self,
          num_inference_steps: int = None,
          device: Union[str, torch.device] = None,
          sigmas: Optional[List[float]] = None,
          mu: Optional[float] = None,
      ):
          """
          Sets the discrete timesteps used for the diffusion chain (to be run before inference).
  
          Args:
              num_inference_steps (`int`):
                  The number of diffusion steps used when generating samples with a pre-trained model.
              device (`str` or `torch.device`, *optional*):
                  The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
          """
  
          if self.config.use_dynamic_shifting and mu is None:
              raise ValueError(
                  " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
              )
  
          if sigmas is None:
              self.num_inference_steps = num_inference_steps
              timesteps = np.array(
                  [
                      (1.0 - i / num_inference_steps) * self.config.num_train_timesteps
                      for i in range(num_inference_steps)
                  ]
              )  # different from the original code in SD3
              sigmas = timesteps / self.config.num_train_timesteps
  
          if self.config.use_dynamic_shifting:
              sigmas = self.time_shift_dynamic(mu, 1.0, sigmas)
          else:
              sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
  
          sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
          timesteps = sigmas * self.config.num_train_timesteps
  
          self.timesteps = timesteps.to(device=device)
          self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
  
          self._step_index = None
          self._begin_index = None
  
      def index_for_timestep(self, timestep, schedule_timesteps=None):
          if schedule_timesteps is None:
              schedule_timesteps = self.timesteps
  
          indices = (schedule_timesteps == timestep).nonzero()
  
          # The sigma index that is taken for the **very** first `step`
          # is always the second index (or the last index if there is only 1)
          # This way we can ensure we don't accidentally skip a sigma in
          # case we start in the middle of the denoising schedule (e.g. for image-to-image)
          pos = 1 if len(indices) > 1 else 0
  
          return indices[pos].item()
  
      def _init_step_index(self, timestep):
          if self.begin_index is None:
              if isinstance(timestep, torch.Tensor):
                  timestep = timestep.to(self.timesteps.device)
              self._step_index = self.index_for_timestep(timestep)
          else:
              self._step_index = self._begin_index
  
      def step(
          self,
          model_output: torch.FloatTensor,
          timestep: Union[float, torch.FloatTensor],
          sample: torch.FloatTensor,
          s_churn: float = 0.0,
          s_tmin: float = 0.0,
          s_tmax: float = float("inf"),
          s_noise: float = 1.0,
          generator: Optional[torch.Generator] = None,
          return_dict: bool = True,
      ) -> Union[RectifiedFlowSchedulerOutput, Tuple]:
          """
          Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
          process from the learned model outputs (most often the predicted noise).
  
          Args:
              model_output (`torch.FloatTensor`):
                  The direct output from learned diffusion model.
              timestep (`float`):
                  The current discrete timestep in the diffusion chain.
              sample (`torch.FloatTensor`):
                  A current instance of a sample created by the diffusion process.
              s_churn (`float`):
              s_tmin  (`float`):
              s_tmax  (`float`):
              s_noise (`float`, defaults to 1.0):
                  Scaling factor for noise added to the sample.
              generator (`torch.Generator`, *optional*):
                  A random number generator.
              return_dict (`bool`):
                  Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
                  tuple.
  
          Returns:
              [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
                  If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
                  returned, otherwise a tuple is returned where the first element is the sample tensor.
          """
  
          if (
              isinstance(timestep, int)
              or isinstance(timestep, torch.IntTensor)
              or isinstance(timestep, torch.LongTensor)
          ):
              raise ValueError(
                  (
                      "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
                      " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
                      " one of the `scheduler.timesteps` as a timestep."
                  ),
              )
  
          if self.step_index is None:
              self._init_step_index(timestep)
  
          # Upcast to avoid precision issues when computing prev_sample
          sample = sample.to(torch.float32)
  
          sigma = self.sigmas[self.step_index]
          sigma_next = self.sigmas[self.step_index + 1]
  
          # Here different directions are used for the flow matching
          prev_sample = sample + (sigma - sigma_next) * model_output
  
          # Cast sample back to model compatible dtype
          prev_sample = prev_sample.to(model_output.dtype)
  
          # upon completion increase step index by one
          self._step_index += 1
  
          if not return_dict:
              return (prev_sample,)
  
          return RectifiedFlowSchedulerOutput(prev_sample=prev_sample)
  
      def scale_noise(
          self,
          original_samples: torch.Tensor,
          noise: torch.Tensor,
          timesteps: torch.IntTensor,
      ) -> torch.Tensor:
          """
          Forward function for the noise scaling in the flow matching.
          """
          sigmas = self._t_to_sigma(timesteps.to(dtype=torch.float32))
  
          while len(sigmas.shape) < len(original_samples.shape):
              sigmas = sigmas.unsqueeze(-1)
  
          return (1.0 - sigmas) * original_samples + sigmas * noise
  
      def __len__(self):
          return self.config.num_train_timesteps
  