ScoreSdeVeScheduler

ScoreSdeVeScheduler is a variance exploding stochastic differential equation (SDE) scheduler. It was introduced in the Score-Based Generative Modeling through Stochastic Differential Equations paper by Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, Ben Poole.

The abstract from the paper is:

Creating noise from data is easy; creating data from noise is generative modeling. We present a stochastic differential equation (SDE) that smoothly transforms a complex data distribution to a known prior distribution by slowly injecting noise, and a corresponding reverse-time SDE that transforms the prior distribution back into the data distribution by slowly removing the noise. Crucially, the reverse-time SDE depends only on the time-dependent gradient field (\aka, score) of the perturbed data distribution. By leveraging advances in score-based generative modeling, we can accurately estimate these scores with neural networks, and use numerical SDE solvers to generate samples. We show that this framework encapsulates previous approaches in score-based generative modeling and diffusion probabilistic modeling, allowing for new sampling procedures and new modeling capabilities. In particular, we introduce a predictor-corrector framework to correct errors in the evolution of the discretized reverse-time SDE. We also derive an equivalent neural ODE that samples from the same distribution as the SDE, but additionally enables exact likelihood computation, and improved sampling efficiency. In addition, we provide a new way to solve inverse problems with score-based models, as demonstrated with experiments on class-conditional generation, image inpainting, and colorization. Combined with multiple architectural improvements, we achieve record-breaking performance for unconditional image generation on CIFAR-10 with an Inception score of 9.89 and FID of 2.20, a competitive likelihood of 2.99 bits/dim, and demonstrate high fidelity generation of 1024 x 1024 images for the first time from a score-based generative model.

mindone.diffusers.ScoreSdeVeScheduler

Bases: SchedulerMixin, ConfigMixin

ScoreSdeVeScheduler is a variance exploding stochastic differential equation (SDE) scheduler.

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.

PARAMETER	DESCRIPTION
num_train_timesteps	The number of diffusion steps to train the model. TYPE: `int`, defaults to 1000 DEFAULT: 2000
snr	A coefficient weighting the step from the model_output sample (from the network) to the random noise. TYPE: `float`, defaults to 0.15 DEFAULT: 0.15
sigma_min	The initial noise scale for the sigma sequence in the sampling procedure. The minimum sigma should mirror the distribution of the data. TYPE: `float`, defaults to 0.01 DEFAULT: 0.01
sigma_max	The maximum value used for the range of continuous timesteps passed into the model. TYPE: `float`, defaults to 1348.0 DEFAULT: 1348.0
sampling_eps	The end value of sampling where timesteps decrease progressively from 1 to epsilon. TYPE: `float`, defaults to 1e-5 DEFAULT: 1e-05
correct_steps	The number of correction steps performed on a produced sample. TYPE: `int`, defaults to 1 DEFAULT: 1

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

class ScoreSdeVeScheduler(SchedulerMixin, ConfigMixin):
    """
    `ScoreSdeVeScheduler` is a variance exploding stochastic differential equation (SDE) scheduler.

    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.
        snr (`float`, defaults to 0.15):
            A coefficient weighting the step from the `model_output` sample (from the network) to the random noise.
        sigma_min (`float`, defaults to 0.01):
            The initial noise scale for the sigma sequence in the sampling procedure. The minimum sigma should mirror
            the distribution of the data.
        sigma_max (`float`, defaults to 1348.0):
            The maximum value used for the range of continuous timesteps passed into the model.
        sampling_eps (`float`, defaults to 1e-5):
            The end value of sampling where timesteps decrease progressively from 1 to epsilon.
        correct_steps (`int`, defaults to 1):
            The number of correction steps performed on a produced sample.
    """

    order = 1

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 2000,
        snr: float = 0.15,
        sigma_min: float = 0.01,
        sigma_max: float = 1348.0,
        sampling_eps: float = 1e-5,
        correct_steps: int = 1,
    ):
        # standard deviation of the initial noise distribution
        self.init_noise_sigma = sigma_max

        # setable values
        self.timesteps = None

        self.set_sigmas(num_train_timesteps, sigma_min, sigma_max, sampling_eps)

    def scale_model_input(self, sample: ms.Tensor, timestep: Optional[int] = None) -> ms.Tensor:
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
            sample (`ms.Tensor`):
                The input sample.
            timestep (`int`, *optional*):
                The current timestep in the diffusion chain.

        Returns:
            `ms.Tensor`:
                A scaled input sample.
        """
        return sample

    def set_timesteps(self, num_inference_steps: int, sampling_eps: float = None):
        """
        Sets the continuous 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.
            sampling_eps (`float`, *optional*):
                The final timestep value (overrides value given during scheduler instantiation).

        """
        sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps

        self.timesteps = ms.tensor(np.linspace(1, sampling_eps, num_inference_steps), dtype=ms.float32)

    def set_sigmas(
        self, num_inference_steps: int, sigma_min: float = None, sigma_max: float = None, sampling_eps: float = None
    ):
        """
        Sets the noise scales used for the diffusion chain (to be run before inference). The sigmas control the weight
        of the `drift` and `diffusion` components of the sample update.

        Args:
            num_inference_steps (`int`):
                The number of diffusion steps used when generating samples with a pre-trained model.
            sigma_min (`float`, optional):
                The initial noise scale value (overrides value given during scheduler instantiation).
            sigma_max (`float`, optional):
                The final noise scale value (overrides value given during scheduler instantiation).
            sampling_eps (`float`, optional):
                The final timestep value (overrides value given during scheduler instantiation).

        """
        sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min
        sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max
        sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps
        if self.timesteps is None:
            self.set_timesteps(num_inference_steps, sampling_eps)

        self.sigmas = sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
        self.discrete_sigmas = ops.exp(
            ms.tensor(np.linspace(math.log(sigma_min), math.log(sigma_max), num_inference_steps), dtype=ms.float32)
        )
        self.sigmas = [sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps]

    def get_adjacent_sigma(self, timesteps, t):
        return ops.where(
            timesteps == 0,
            ops.zeros_like(t),
            self.discrete_sigmas[timesteps - 1],
        )

    def step_pred(
        self,
        model_output: ms.Tensor,
        timestep: int,
        sample: ms.Tensor,
        generator: Optional[np.random.Generator] = None,
        return_dict: bool = False,
    ) -> Union[SdeVeOutput, 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 (`ms.Tensor`):
                The direct output from learned diffusion model.
            timestep (`int`):
                The current discrete timestep in the diffusion chain.
            sample (`ms.Tensor`):
                A current instance of a sample created by the diffusion process.
            generator (`np.random.Generator`, *optional*):
                A random number generator.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`.

        Returns:
            [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple
                is returned where the first element is the sample tensor.

        """
        if self.timesteps is None:
            raise ValueError(
                "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
            )

        broadcast_shape = sample.shape
        timestep = timestep * ops.ones(sample.shape[0])  # ops.repeat_interleave(timestep, sample.shape[0])
        timesteps = (timestep * (len(self.timesteps) - 1)).long()

        # mps requires indices to be in the same device, so we use cpu as is the default with cuda
        sigma = self.discrete_sigmas[timesteps]
        adjacent_sigma = self.get_adjacent_sigma(timesteps, timestep)
        drift = ops.zeros_like(sample)
        diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5

        # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
        # also equation 47 shows the analog from SDE models to ancestral sampling methods
        diffusion = diffusion.flatten()
        # while len(diffusion.shape) < len(sample.shape):
        #     diffusion = diffusion.unsqueeze(-1)
        diffusion = ops.reshape(diffusion, (timesteps.shape[0],) + (1,) * (len(broadcast_shape) - 1))
        drift = (drift - diffusion**2 * model_output).to(sample.dtype)

        #  equation 6: sample noise for the diffusion term of
        noise = randn_tensor(sample.shape, generator=generator, dtype=sample.dtype)
        prev_sample_mean = sample - drift  # subtract because `dt` is a small negative timestep
        # TODO is the variable diffusion the correct scaling term for the noise?
        prev_sample = (prev_sample_mean + diffusion * noise).to(sample.dtype)  # add impact of diffusion field g

        if not return_dict:
            return (prev_sample, prev_sample_mean)

        return SdeVeOutput(prev_sample=prev_sample, prev_sample_mean=prev_sample_mean)

    def step_correct(
        self,
        model_output: ms.Tensor,
        sample: ms.Tensor,
        generator: Optional[np.random.Generator] = None,
        return_dict: bool = False,
    ) -> Union[SchedulerOutput, Tuple]:
        """
        Correct the predicted sample based on the `model_output` of the network. This is often run repeatedly after
        making the prediction for the previous timestep.

        Args:
            model_output (`ms.Tensor`):
                The direct output from learned diffusion model.
            sample (`ms.Tensor`):
                A current instance of a sample created by the diffusion process.
            generator (`np.random.Generator`, *optional*):
                A random number generator.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`.

        Returns:
            [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple
                is returned where the first element is the sample tensor.

        """
        if self.timesteps is None:
            raise ValueError(
                "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
            )

        # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
        # sample noise for correction
        noise = randn_tensor(sample.shape, generator=generator, dtype=ms.float32)

        # compute step size from the model_output, the noise, and the snr
        grad_norm = ops.norm(model_output.reshape(model_output.shape[0], -1), dim=-1).mean()
        noise_norm = ops.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
        step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
        step_size = step_size * ops.ones(sample.shape[0])
        # self.repeat_scalar(step_size, sample.shape[0])

        # compute corrected sample: model_output term and noise term
        step_size = step_size.flatten()
        # while len(step_size.shape) < len(sample.shape):
        #     step_size = step_size.unsqueeze(-1)
        step_size = ops.reshape(step_size, (sample.shape[0],) + (1,) * (len(sample.shape) - 1))

        prev_sample_mean = sample + step_size * model_output
        prev_sample = (prev_sample_mean + ((step_size * 2) ** 0.5) * noise).to(sample.dtype)

        if not return_dict:
            return (prev_sample,)

        return SchedulerOutput(prev_sample=prev_sample)

    def add_noise(
        self,
        original_samples: ms.Tensor,
        noise: ms.Tensor,
        timesteps: ms.Tensor,
    ) -> ms.Tensor:
        # Make sure sigmas and timesteps have the same device and dtype as original_samples
        sigmas = self.discrete_sigmas[timesteps]
        noise = (
            noise * sigmas[:, None, None, None]
            if noise is not None
            else ops.randn_like(original_samples, dtype=original_samples.dtype) * sigmas[:, None, None, None]
        ).to(original_samples.dtype)
        noisy_samples = noise + original_samples
        return noisy_samples

    def __len__(self):
        return self.config.num_train_timesteps

mindone.diffusers.ScoreSdeVeScheduler.scale_model_input(sample, timestep=None)

Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep.

PARAMETER	DESCRIPTION
sample	The input sample. TYPE: `ms.Tensor`
timestep	The current timestep in the diffusion chain. TYPE: `int`, *optional* DEFAULT: None

RETURNS	DESCRIPTION
Tensor	ms.Tensor: A scaled input sample.

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

def scale_model_input(self, sample: ms.Tensor, timestep: Optional[int] = None) -> ms.Tensor:
    """
    Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
    current timestep.

    Args:
        sample (`ms.Tensor`):
            The input sample.
        timestep (`int`, *optional*):
            The current timestep in the diffusion chain.

    Returns:
        `ms.Tensor`:
            A scaled input sample.
    """
    return sample

mindone.diffusers.ScoreSdeVeScheduler.set_sigmas(num_inference_steps, sigma_min=None, sigma_max=None, sampling_eps=None)

Sets the noise scales used for the diffusion chain (to be run before inference). The sigmas control the weight of the drift and diffusion components of the sample update.

PARAMETER	DESCRIPTION
num_inference_steps	The number of diffusion steps used when generating samples with a pre-trained model. TYPE: `int`
sigma_min	The initial noise scale value (overrides value given during scheduler instantiation). TYPE: `float` DEFAULT: None
sigma_max	The final noise scale value (overrides value given during scheduler instantiation). TYPE: `float` DEFAULT: None
sampling_eps	The final timestep value (overrides value given during scheduler instantiation). TYPE: `float` DEFAULT: None

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

def set_sigmas(
    self, num_inference_steps: int, sigma_min: float = None, sigma_max: float = None, sampling_eps: float = None
):
    """
    Sets the noise scales used for the diffusion chain (to be run before inference). The sigmas control the weight
    of the `drift` and `diffusion` components of the sample update.

    Args:
        num_inference_steps (`int`):
            The number of diffusion steps used when generating samples with a pre-trained model.
        sigma_min (`float`, optional):
            The initial noise scale value (overrides value given during scheduler instantiation).
        sigma_max (`float`, optional):
            The final noise scale value (overrides value given during scheduler instantiation).
        sampling_eps (`float`, optional):
            The final timestep value (overrides value given during scheduler instantiation).

    """
    sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min
    sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max
    sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps
    if self.timesteps is None:
        self.set_timesteps(num_inference_steps, sampling_eps)

    self.sigmas = sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
    self.discrete_sigmas = ops.exp(
        ms.tensor(np.linspace(math.log(sigma_min), math.log(sigma_max), num_inference_steps), dtype=ms.float32)
    )
    self.sigmas = [sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps]

mindone.diffusers.ScoreSdeVeScheduler.set_timesteps(num_inference_steps, sampling_eps=None)

Sets the continuous timesteps used for the diffusion chain (to be run before inference).

PARAMETER	DESCRIPTION
num_inference_steps	The number of diffusion steps used when generating samples with a pre-trained model. TYPE: `int`
sampling_eps	The final timestep value (overrides value given during scheduler instantiation). TYPE: `float`, *optional* DEFAULT: None

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

def set_timesteps(self, num_inference_steps: int, sampling_eps: float = None):
    """
    Sets the continuous 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.
        sampling_eps (`float`, *optional*):
            The final timestep value (overrides value given during scheduler instantiation).

    """
    sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps

    self.timesteps = ms.tensor(np.linspace(1, sampling_eps, num_inference_steps), dtype=ms.float32)

mindone.diffusers.ScoreSdeVeScheduler.step_correct(model_output, sample, generator=None, return_dict=False)

Correct the predicted sample based on the model_output of the network. This is often run repeatedly after making the prediction for the previous timestep.

PARAMETER	DESCRIPTION
model_output	The direct output from learned diffusion model. TYPE: `ms.Tensor`
sample	A current instance of a sample created by the diffusion process. TYPE: `ms.Tensor`
generator	A random number generator. TYPE: `np.random.Generator`, *optional* DEFAULT: None
return_dict	Whether or not to return a [~schedulers.scheduling_sde_ve.SdeVeOutput] or tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
Union[SchedulerOutput, Tuple]	[~schedulers.scheduling_sde_ve.SdeVeOutput] or tuple: If return_dict is True, [~schedulers.scheduling_sde_ve.SdeVeOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

def step_correct(
    self,
    model_output: ms.Tensor,
    sample: ms.Tensor,
    generator: Optional[np.random.Generator] = None,
    return_dict: bool = False,
) -> Union[SchedulerOutput, Tuple]:
    """
    Correct the predicted sample based on the `model_output` of the network. This is often run repeatedly after
    making the prediction for the previous timestep.

    Args:
        model_output (`ms.Tensor`):
            The direct output from learned diffusion model.
        sample (`ms.Tensor`):
            A current instance of a sample created by the diffusion process.
        generator (`np.random.Generator`, *optional*):
            A random number generator.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`.

    Returns:
        [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`:
            If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple
            is returned where the first element is the sample tensor.

    """
    if self.timesteps is None:
        raise ValueError(
            "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
        )

    # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
    # sample noise for correction
    noise = randn_tensor(sample.shape, generator=generator, dtype=ms.float32)

    # compute step size from the model_output, the noise, and the snr
    grad_norm = ops.norm(model_output.reshape(model_output.shape[0], -1), dim=-1).mean()
    noise_norm = ops.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
    step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
    step_size = step_size * ops.ones(sample.shape[0])
    # self.repeat_scalar(step_size, sample.shape[0])

    # compute corrected sample: model_output term and noise term
    step_size = step_size.flatten()
    # while len(step_size.shape) < len(sample.shape):
    #     step_size = step_size.unsqueeze(-1)
    step_size = ops.reshape(step_size, (sample.shape[0],) + (1,) * (len(sample.shape) - 1))

    prev_sample_mean = sample + step_size * model_output
    prev_sample = (prev_sample_mean + ((step_size * 2) ** 0.5) * noise).to(sample.dtype)

    if not return_dict:
        return (prev_sample,)

    return SchedulerOutput(prev_sample=prev_sample)

mindone.diffusers.ScoreSdeVeScheduler.step_pred(model_output, timestep, sample, generator=None, return_dict=False)

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

PARAMETER	DESCRIPTION
model_output	The direct output from learned diffusion model. TYPE: `ms.Tensor`
timestep	The current discrete timestep in the diffusion chain. TYPE: `int`
sample	A current instance of a sample created by the diffusion process. TYPE: `ms.Tensor`
generator	A random number generator. TYPE: `np.random.Generator`, *optional* DEFAULT: None
return_dict	Whether or not to return a [~schedulers.scheduling_sde_ve.SdeVeOutput] or tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
Union[SdeVeOutput, Tuple]	[~schedulers.scheduling_sde_ve.SdeVeOutput] or tuple: If return_dict is True, [~schedulers.scheduling_sde_ve.SdeVeOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

def step_pred(
    self,
    model_output: ms.Tensor,
    timestep: int,
    sample: ms.Tensor,
    generator: Optional[np.random.Generator] = None,
    return_dict: bool = False,
) -> Union[SdeVeOutput, 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 (`ms.Tensor`):
            The direct output from learned diffusion model.
        timestep (`int`):
            The current discrete timestep in the diffusion chain.
        sample (`ms.Tensor`):
            A current instance of a sample created by the diffusion process.
        generator (`np.random.Generator`, *optional*):
            A random number generator.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`.

    Returns:
        [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`:
            If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple
            is returned where the first element is the sample tensor.

    """
    if self.timesteps is None:
        raise ValueError(
            "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
        )

    broadcast_shape = sample.shape
    timestep = timestep * ops.ones(sample.shape[0])  # ops.repeat_interleave(timestep, sample.shape[0])
    timesteps = (timestep * (len(self.timesteps) - 1)).long()

    # mps requires indices to be in the same device, so we use cpu as is the default with cuda
    sigma = self.discrete_sigmas[timesteps]
    adjacent_sigma = self.get_adjacent_sigma(timesteps, timestep)
    drift = ops.zeros_like(sample)
    diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5

    # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
    # also equation 47 shows the analog from SDE models to ancestral sampling methods
    diffusion = diffusion.flatten()
    # while len(diffusion.shape) < len(sample.shape):
    #     diffusion = diffusion.unsqueeze(-1)
    diffusion = ops.reshape(diffusion, (timesteps.shape[0],) + (1,) * (len(broadcast_shape) - 1))
    drift = (drift - diffusion**2 * model_output).to(sample.dtype)

    #  equation 6: sample noise for the diffusion term of
    noise = randn_tensor(sample.shape, generator=generator, dtype=sample.dtype)
    prev_sample_mean = sample - drift  # subtract because `dt` is a small negative timestep
    # TODO is the variable diffusion the correct scaling term for the noise?
    prev_sample = (prev_sample_mean + diffusion * noise).to(sample.dtype)  # add impact of diffusion field g

    if not return_dict:
        return (prev_sample, prev_sample_mean)

    return SdeVeOutput(prev_sample=prev_sample, prev_sample_mean=prev_sample_mean)

mindone.diffusers.schedulers.scheduling_sde_ve.SdeVeOutput dataclass

Bases: BaseOutput

Output class for the scheduler's step function output.

PARAMETER	DESCRIPTION
prev_sample	Computed sample (x_{t-1}) of previous timestep. prev_sample should be used as next model input in the denoising loop. TYPE: `ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images
prev_sample_mean	Mean averaged prev_sample over previous timesteps. TYPE: `ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images

Source code in mindone/diffusers/schedulers/scheduling_sde_ve.py

@dataclass
class SdeVeOutput(BaseOutput):
    """
    Output class for the scheduler's `step` function output.

    Args:
        prev_sample (`ms.Tensor` 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_mean (`ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
            Mean averaged `prev_sample` over previous timesteps.
    """

    prev_sample: ms.Tensor
    prev_sample_mean: ms.Tensor