TorchDiff 2.0.0__py3-none-any.whl

ddpm/noise_predictor.py

@@ -0,0 +1,521 @@
+import torch
+import torch.nn as nn
+
+
+
+class NoisePredictor(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            down_channels,
+            mid_channels,
+            up_channels,
+            down_sampling,
+            time_embed_dim,
+            y_embed_dim, # output embedding dimension in text conditional net
+            num_down_blocks,
+            num_mid_blocks,
+            num_up_blocks,
+            dropout_rate=0.1,
+            down_sampling_factor=2,
+            where_y=True,
+            y_to_all=False
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.down_channels = down_channels
+        self.mid_channels = mid_channels
+        self.up_channels = up_channels
+        self.down_sampling = down_sampling
+        self.time_embed_dim = time_embed_dim
+        self.y_embed_dim = y_embed_dim
+        self.num_down_blocks = num_down_blocks
+        self.num_mid_blocks = num_mid_blocks
+        self.num_up_blocks = num_up_blocks
+        self.dropout_rate = dropout_rate
+        self.where_y = where_y
+        self.up_sampling = list(reversed(self.down_sampling))
+        self.conv1 = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.down_channels[0],
+            kernel_size=3,
+            padding=1
+        )
+        # initial time embedding projection
+        self.time_projection = nn.Sequential(
+            nn.Linear(in_features=self.time_embed_dim, out_features=self.time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(in_features=self.time_embed_dim, out_features=self.time_embed_dim)
+        )
+        # down blocks
+        self.down_blocks = nn.ModuleList([
+            DownBlock(
+                in_channels=self.down_channels[i],
+                out_channels=self.down_channels[i+1],
+                time_embed_dim=self.time_embed_dim,
+                y_embed_dim=y_embed_dim,
+                num_layers=self.num_down_blocks,
+                down_sampling_factor=down_sampling_factor,
+                down_sample=self.down_sampling[i],
+                dropout_rate=self.dropout_rate,
+                y_to_all=y_to_all
+            ) for i in range(len(self.down_channels)-1)
+        ])
+        # middle blocks
+        self.mid_blocks = nn.ModuleList([
+            MiddleBlock(
+                in_channels=self.mid_channels[i],
+                out_channels=self.mid_channels[i + 1],
+                time_embed_dim=self.time_embed_dim,
+                y_embed_dim=y_embed_dim,
+                num_layers=self.num_mid_blocks,
+                dropout_rate=self.dropout_rate,
+                y_to_all=y_to_all
+            ) for i in range(len(self.mid_channels) - 1)
+        ])
+        # up blocks
+        skip_channels = list(reversed(self.down_channels))
+        self.up_blocks = nn.ModuleList([
+            UpBlock(
+                in_channels=self.up_channels[i],
+                out_channels=self.up_channels[i+1],
+                skip_channels=skip_channels[i],
+                time_embed_dim=self.time_embed_dim,
+                y_embed_dim=y_embed_dim,
+                num_layers=self.num_up_blocks,
+                up_sampling_factor=down_sampling_factor,
+                up_sampling=self.up_sampling[i],
+                dropout_rate=self.dropout_rate,
+                y_to_all=y_to_all
+            ) for i in range(len(self.up_channels)-1)
+        ])
+        # final convolution layer
+        self.conv2 = nn.Sequential(
+            nn.GroupNorm(num_groups=8, num_channels=self.up_channels[-1]),
+            nn.Dropout(p=self.dropout_rate),
+            nn.Conv2d(in_channels=self.up_channels[-1], out_channels=self.in_channels, kernel_size=3, padding=1)
+        )
+
+    def initialize_weights(self):
+        """Initialize model weights for better training stability"""
+        for module in self.modules():
+            if isinstance(module, (nn.Conv2d, nn.Linear, nn.ConvTranspose2d)):
+                nn.init.kaiming_normal_(module.weight, a=0.2, nonlinearity='leaky_relu')
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, x, t, y=None):
+
+        if not self.where_y and y is not None:
+            x = torch.cat(tensors=[x, y], dim=1)
+        output = self.conv1(x)
+        time_embed = GetEmbeddedTime(embed_dim=self.time_embed_dim)(time_steps=t)
+        time_embed = self.time_projection(time_embed)
+        skip_connections = []
+
+        for i, down in enumerate(self.down_blocks):
+            skip_connections.append(output)
+            output = down(x=output, embed_time=time_embed, y=y)
+        for i, mid in enumerate(self.mid_blocks):
+            output = mid(x=output, embed_time=time_embed, y=y)
+        for i, up in enumerate(self.up_blocks):
+            skip_connection = skip_connections.pop()
+            output = up(x=output, skip_connection=skip_connection, embed_time=time_embed, y=y)
+
+        output = self.conv2(output)
+        return output
+#-----------------------------------------------------------------------------
+class DownBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, time_embed_dim, y_embed_dim,num_layers, down_sampling_factor,  down_sample, dropout_rate, y_to_all):
+        super().__init__()
+        self.num_layers = num_layers
+        self.y_to_all = y_to_all
+        self.conv1 = nn.ModuleList([
+            Conv3(
+                in_channels=in_channels if i==0 else out_channels,
+                out_channels=out_channels,
+                num_groups=8,
+                kernel_size=3,
+                norm=True,
+                activation=True,
+                dropout_rate=dropout_rate
+            ) for i in range(self.num_layers)
+        ])
+        self.conv2 = nn.ModuleList([
+            Conv3(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                num_groups=8,
+                kernel_size=3,
+                norm=True,
+                activation=True,
+                dropout_rate=dropout_rate
+            ) for _ in range(self.num_layers)
+        ])
+        self.time_embedding = nn.ModuleList([
+            TimeEmbedding(
+                output_dim=out_channels,
+                embed_dim=time_embed_dim
+            ) for _ in range(self.num_layers)
+        ])
+        self.attention = nn.ModuleList([
+            Attention(
+                in_channels=out_channels,
+                y_embed_dim= y_embed_dim,
+                num_groups=8,
+                num_heads=4,
+                dropout_rate=dropout_rate
+            ) for _ in range(self.num_layers)
+        ])
+        self.down_sampling = DownSampling(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            down_sampling_factor=down_sampling_factor,
+            conv_block=True,
+            max_pool=True
+        ) if down_sample else nn.Identity()
+        self.resnet = nn.ModuleList([
+            nn.Conv2d(
+                in_channels=in_channels if i == 0 else out_channels,
+                out_channels=out_channels,
+                kernel_size=1
+            ) for i in range(num_layers)
+
+        ])
+
+    def forward(self, x, embed_time, y):
+        #print("down-block input shape:", x.size())
+        output = x
+        for i in range(self.num_layers):
+            resnet_input = output
+            output = self.conv1[i](output)
+            output = output + self.time_embedding[i](embed_time)[:, :, None, None]
+            output = self.conv2[i](output)
+            output = output + self.resnet[i](resnet_input)
+            if y is not None and not self.y_to_all and i == 0:
+                out_attn = self.attention[i](output, y)
+                output = output + out_attn
+            elif y is not None and self.y_to_all:
+                out_attn = self.attention[i](output, y)
+                output = output + out_attn
+            elif y is None and self.y_to_all:
+                out_attn = self.attention[i](output)
+                output = output + out_attn
+            elif y is None and not self.y_to_all and i == 0:
+                out_attn = self.attention[i](output)
+                output = output + out_attn
+
+        output = self.down_sampling(output)
+        #print("down-block output shape:", output.size())
+        return output
+#------------------------------------------------------------------------------
+class MiddleBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, time_embed_dim,  y_embed_dim, num_layers, dropout_rate, y_to_all=False):
+        super().__init__()
+        self.num_layers = num_layers
+        self.y_to_all = y_to_all
+        self.conv1 = nn.ModuleList([
+            Conv3(
+                in_channels=in_channels if i == 0 else out_channels,
+                out_channels=out_channels,
+                num_groups=8,
+                kernel_size=3,
+                norm=True,
+                activation=True,
+                dropout_rate=dropout_rate
+            ) for i in range(self.num_layers+1)
+        ])
+        self.conv2 = nn.ModuleList([
+            Conv3(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                num_groups=8,
+                kernel_size=3,
+                norm=True,
+                activation=True,
+                dropout_rate=dropout_rate
+            ) for _ in range(self.num_layers+1)
+        ])
+        self.time_embedding = nn.ModuleList([
+            TimeEmbedding(
+                output_dim=out_channels,
+                embed_dim=time_embed_dim
+            ) for _ in range(self.num_layers+1)
+        ])
+        self.attention = nn.ModuleList([
+            Attention(
+                in_channels=out_channels,
+                y_embed_dim=y_embed_dim,
+                num_groups=8,
+                num_heads=4,
+                dropout_rate=dropout_rate
+            ) for _ in range(self.num_layers + 1)
+        ])
+        self.resnet = nn.ModuleList([
+            nn.Conv2d(
+                in_channels=in_channels if i == 0 else out_channels,
+                out_channels=out_channels,
+                kernel_size=1
+            ) for i in range(num_layers+1)
+        ])
+
+    def forward(self, x, embed_time, y=None):
+        #print("mid-input shape:", x.size())
+        output = x
+        resnet_input = output
+        output = self.conv1[0](output)
+        output = output + self.time_embedding[0](embed_time)[:, :, None, None]
+        output = self.conv2[0](output)
+        output = output + self.resnet[0](resnet_input)
+        for i in range(self.num_layers):
+            if y is not None and not self.y_to_all and i == 0:
+                out_attn = self.attention[i](output, y)
+                output = output + out_attn
+            elif y is not None and self.y_to_all:
+                out_attn = self.attention[i](output, y)
+                output = output + out_attn
+            elif y is None and self.y_to_all:
+                out_attn = self.attention[i](output)
+                output = output + out_attn
+            elif y is None and not self.y_to_all and i == 0:
+                out_attn = self.attention[i](output)
+                output = output + out_attn
+            resnet_input = output
+            output = self.conv1[i + 1](output)
+            output = output + self.time_embedding[i + 1](embed_time)[:, :, None, None]
+            output = self.conv2[i + 1](output)
+            output = output + self.resnet[i+1](resnet_input)
+        #print("mid-block output shape:", output.size())
+
+        return output
+#------------------------------------------------------------------------------
+class UpBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, skip_channels, time_embed_dim,  y_embed_dim, num_layers, up_sampling_factor, up_sampling=True, dropout_rate=0.2, y_to_all=False):
+        super().__init__()
+        self.num_layers = num_layers
+        self.y_to_all = y_to_all
+        effective_in_channels = in_channels//2 + skip_channels
+        self.conv1 = nn.ModuleList([
+            Conv3(
+                in_channels=effective_in_channels  if i == 0 else out_channels,
+                out_channels=out_channels,
+                num_groups=8,
+                kernel_size=3,
+                norm=True,
+                activation=True,
+                dropout_rate=dropout_rate
+            ) for i in range(self.num_layers)
+        ])
+        self.conv2 = nn.ModuleList([
+            Conv3(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                num_groups=8,
+                kernel_size=3,
+                norm=True,
+                activation=True,
+                dropout_rate=dropout_rate
+            ) for _ in range(self.num_layers)
+        ])
+        self.time_embedding = nn.ModuleList([
+            TimeEmbedding(
+                output_dim=out_channels,
+                embed_dim=time_embed_dim
+            ) for _ in range(self.num_layers)
+        ])
+        self.attention = nn.ModuleList([
+            Attention(
+                in_channels=out_channels,
+                y_embed_dim=y_embed_dim,
+                num_groups=8,
+                num_heads=4,
+                dropout_rate=dropout_rate
+            ) for _ in range(self.num_layers)
+        ])
+        self.up_sampling = UpSampling(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            up_sampling_factor=up_sampling_factor,
+            conv_block=True,
+            up_sampling=True
+        ) if up_sampling else nn.Identity()
+        self.resnet = nn.ModuleList([
+            nn.Conv2d(
+                in_channels=effective_in_channels  if i == 0 else out_channels,
+                out_channels=out_channels,
+                kernel_size=1
+            ) for i in range(num_layers)
+
+        ])
+
+    def forward(self, x, skip_connection, embed_time, y=None):
+        #print("up-block input shape:", x.size())
+        x = self.up_sampling(x)
+        x = torch.cat(tensors=[x, skip_connection], dim=1)
+        output = x
+        for i in range(self.num_layers):
+            resnet_input = output
+            output = self.conv1[i](output)
+            output = output + self.time_embedding[i](embed_time)[:, :, None, None]
+            output = self.conv2[i](output)
+            output = output + self.resnet[i](resnet_input)
+            if y is not None and not self.y_to_all and i == 0:
+                out_attn = self.attention[i](output, y)
+                output = output + out_attn
+            elif y is not None and self.y_to_all:
+                out_attn = self.attention[i](output, y)
+                output = output + out_attn
+            elif y is None and self.y_to_all:
+                out_attn = self.attention[i](output)
+                output = output + out_attn
+            elif y is None and not self.y_to_all and i == 0:
+                out_attn = self.attention[i](output)
+                output = output + out_attn
+        #print("up-block output shape:", output.size())
+        return output
+#------------------------------------------------------------------------
+class Conv3(nn.Module):
+    def __init__(self, in_channels, out_channels, num_groups=8, kernel_size=3, norm=True, activation=True, dropout_rate=0.2):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2)
+        self.group_norm = nn.GroupNorm(num_groups=num_groups, num_channels=out_channels) if norm else nn.Identity()
+        self.activation = nn.SiLU() if activation else nn.Identity()
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward(self, batch):
+        batch = self.conv(batch)
+        batch = self.group_norm(batch)
+        batch = self.activation(batch)
+        batch = self.dropout(batch)
+        return batch
+#----------------------------------------------------------------
+class TimeEmbedding(nn.Module):
+    def __init__(self, output_dim, embed_dim):
+        super().__init__()
+        self.embedding = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(in_features=embed_dim, out_features=output_dim)
+        )
+    def forward(self, batch):
+        return self.embedding(batch)
+#----------------------------------------------------------------
+class GetEmbeddedTime(nn.Module):
+    def __init__(self, embed_dim):
+        super().__init__()
+        assert embed_dim % 2 == 0, "The embedding dimension must be divisible by two"
+        self.embed_dim = embed_dim
+
+    def forward(self, time_steps):
+        i = torch.arange(start=0, end=self.embed_dim // 2, dtype=torch.float32, device=time_steps.device)
+        factor = 10000 ** (2 * i / self.embed_dim)
+        embed_time = time_steps[:, None] / factor
+        embed_time = torch.cat(tensors=[torch.sin(embed_time), torch.cos(embed_time)], dim=-1)
+        return embed_time
+#----------------------------------------------------------------
+class Attention(nn.Module):
+    def __init__(self, in_channels, y_embed_dim=768, num_heads=4, num_groups=8, dropout_rate=0.1):
+        super().__init__()
+        self.in_channels = in_channels
+        self.y_embed_dim = y_embed_dim
+        self.num_heads = num_heads
+        self.dropout_rate = dropout_rate
+        self.attention = nn.MultiheadAttention(embed_dim=in_channels, num_heads=num_heads, dropout=dropout_rate, batch_first=True)
+        self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=in_channels)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.y_projection = nn.Linear(y_embed_dim, in_channels)
+
+    def forward(self, x, y=None):
+        batch_size, channels, h, w = x.shape
+        assert channels == self.in_channels, f"Expected {self.in_channels} channels, got {channels}"
+        x_reshaped = x.view(batch_size, channels, h * w).permute(0, 2, 1)
+        if y is not None:
+            y = self.y_projection(y)
+            if y.dim() != 3:
+                if y.dim() == 2:
+                    y = y.unsqueeze(1)
+                else:
+                    raise ValueError(
+                        f"Expected y to be 2D or 3D after projection, got {y.dim()}D with shape {y.shape}"
+                    )
+            if y.shape[-1] != self.in_channels:
+                raise ValueError(
+                    f"Expected y's embedding dim to match in_channels ({self.in_channels}), got {y.shape[-1]}"
+                )
+            out, _ = self.attention(x_reshaped, y, y)
+        else:
+            out, _ = self.attention(x_reshaped, x_reshaped, x_reshaped)
+        out = out.permute(0, 2, 1).view(batch_size, channels, h, w)
+        out = self.norm(out)
+        out = self.dropout(out)
+        return out
+#-----------------------------------------------------------------
+class DownSampling(nn.Module):
+    def __init__(self, in_channels, out_channels, down_sampling_factor, conv_block=True, max_pool=True):
+        super().__init__()
+        self.conv_block = conv_block
+        self.max_pool = max_pool
+        self.down_sampling_factor = down_sampling_factor
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels // 2 if max_pool else out_channels,
+                      kernel_size=3, stride=down_sampling_factor, padding=1)
+        ) if conv_block else nn.Identity()
+        self.pool = nn.Sequential(
+            nn.MaxPool2d(kernel_size=down_sampling_factor, stride=down_sampling_factor),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels//2 if conv_block else out_channels,
+                      kernel_size=1, stride=1, padding=0)
+        ) if max_pool else nn.Identity()
+
+    def forward(self, batch):
+        if not self.conv_block:
+            return self.pool(batch)
+        if not self.max_pool:
+            return self.conv(batch)
+        return torch.cat(tensors=[self.conv(batch), self.pool(batch)], dim=1)
+#--------------------------------------------------------------------------
+class UpSampling(nn.Module):
+    def __init__(self, in_channels, out_channels, up_sampling_factor, conv_block=True, up_sampling=True):
+        super().__init__()
+        self.conv_block = conv_block
+        self.up_sampling = up_sampling
+        self.up_sampling_factor = up_sampling_factor
+        half_out_channels = out_channels // 2
+        self.conv = nn.Sequential(
+            nn.ConvTranspose2d(
+                in_channels=in_channels,
+                out_channels=half_out_channels if up_sampling else out_channels,
+                kernel_size=3,
+                stride=up_sampling_factor,
+                padding=1,
+                output_padding=up_sampling_factor - 1
+            ),
+            nn.Conv2d(
+                in_channels=half_out_channels if up_sampling else out_channels,
+                out_channels=half_out_channels if up_sampling else out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0
+            )
+        ) if conv_block else nn.Identity()
+
+        self.up_sample = nn.Sequential(
+            nn.Upsample(scale_factor=up_sampling_factor, mode="nearest"),
+            nn.Conv2d(in_channels=in_channels, out_channels=half_out_channels if conv_block else out_channels,
+                      kernel_size=1, stride=1, padding=0)
+        ) if up_sampling else nn.Identity()
+
+    def forward(self, batch):
+        if not self.conv_block:
+            return self.up_sample(batch)
+        if not self.up_sampling:
+            return self.conv(batch)
+        conv_output = self.conv(batch)
+        up_sample_output = self.up_sample(batch)
+        if conv_output.shape[2:] != up_sample_output.shape[2:]:
+            _, _, h, w = conv_output.shape
+            up_sample_output = torch.nn.functional.interpolate(
+                up_sample_output,
+                size=(h, w),
+                mode='nearest'
+            )
+        return torch.cat(tensors=[conv_output, up_sample_output], dim=1)

ddpm/reverse_ddpm.py

@@ -0,0 +1,102 @@
+"""Reverse diffusion process for Denoising Diffusion Probabilistic Models (DDPM).
+
+This module implements the reverse diffusion process as described in the DDPM paper
+(Ho et al., 2020, "Denoising Diffusion Probabilistic Models"). The reverse process
+gradually denoises a noisy input to reconstruct the original data distribution.
+"""
+
+
+import torch
+import torch.nn as nn
+
+
+
+class ReverseDDPM(nn.Module):
+    """Reverse diffusion process of DDPM.
+
+    Implements the reverse diffusion process for DDPM, which iteratively denoises a
+    noisy input `xt` using a predicted noise component, as defined in Ho et al. (2020).
+    The process relies on a noise schedule that can be either fixed or trainable,
+    specified through the provided hyperparameters.
+
+    Parameters
+    ----------
+    hyper_params : object
+        Hyperparameter object containing the noise schedule parameters. Expected to have
+        attributes:
+        - `num_steps`: Number of diffusion steps (int).
+        - `trainable_beta`: Whether the noise schedule is trainable (bool).
+        - `betas`: Noise schedule parameters (torch.Tensor, optional if trainable_beta is True).
+        - `alphas`: Precomputed alpha values (torch.Tensor, optional if trainable_beta is False).
+        - `alpha_bars`: Precomputed cumulative product of alphas (torch.Tensor, optional if trainable_beta is False).
+        - `compute_schedule`: Method to compute the noise schedule (callable, optional if trainable_beta is True).
+
+    Attributes
+    ----------
+    hyper_params : object
+        Stores the provided hyperparameter object for use in the reverse process.
+    """
+    def __init__(self, hyper_params):
+        super().__init__()
+        self.hyper_params = hyper_params # hyperparameters class
+
+    def forward(self, xt, predicted_noise, time_steps):
+        """Applies the reverse diffusion process to the noisy input.
+
+        Denoises the input `xt` by computing the mean of the reverse process
+        distribution using the predicted noise and optionally adding stochastic noise
+        for time steps greater than 0, as per the DDPM reverse process.
+
+        Parameters
+        ----------
+        xt : torch.Tensor
+            Noisy input tensor at time step `t`, of shape (batch_size, channels, height, width).
+        predicted_noise : torch.Tensor
+            Predicted noise tensor, of the same shape as `xt`, typically output by a neural network.
+        time_steps : torch.Tensor
+            Tensor of time step indices (long), shape (batch_size,), where each value
+            is in the range [0, hyper_params.num_steps - 1].
+
+        Returns
+        -------
+        torch.Tensor
+            Denoised tensor `xt_minus_1` at time step `t-1`, with the same shape as `xt`.
+            For time_steps == 0, returns the mean of the reverse process without added noise.
+
+        Raises
+        ------
+        ValueError
+            If any value in `time_steps` is outside the valid range
+            [0, hyper_params.num_steps - 1].
+        """
+        if not torch.all((time_steps >= 0) & (time_steps < self.hyper_params.num_steps)):
+            raise ValueError(f"time_steps must be between 0 and {self.hyper_params.num_steps - 1}")
+
+        if self.hyper_params.trainable_beta:
+            betas_t, alphas_t, alpha_bars_t, _, _ = self.hyper_params.compute_schedule(self.hyper_params.betas)
+            betas_t = betas_t[time_steps].to(xt.device)
+            alphas_t = alphas_t[time_steps].to(xt.device)
+            alpha_bars_t = alpha_bars_t[time_steps].to(xt.device)
+            alpha_bars_t_minus_1 = alpha_bars_t[time_steps - 1].to(xt.device) if time_steps.any() else None
+        else:
+            betas_t = self.hyper_params.betas[time_steps].to(xt.device)
+            alphas_t = self.hyper_params.alphas[time_steps].to(xt.device)
+            alpha_bars_t = self.hyper_params.alpha_bars[time_steps].to(xt.device)
+            alpha_bars_t_minus_1 = self.hyper_params.alpha_bars[time_steps - 1].to(xt.device) if time_steps.any() else None
+
+        sqrt_alphas_t = torch.sqrt(alphas_t).view(-1, 1, 1, 1)
+        sqrt_one_minus_alpha_bars_t = torch.sqrt(1 - alpha_bars_t).view(-1, 1, 1, 1)
+        betas_t = betas_t.view(-1, 1, 1, 1)
+
+        mu = (xt - (betas_t / sqrt_one_minus_alpha_bars_t) * predicted_noise) / sqrt_alphas_t
+
+        mask = (time_steps == 0)
+        if mask.all():
+            return mu
+
+        variance = (1 - alpha_bars_t_minus_1) / (1 - alpha_bars_t) * betas_t.squeeze()
+        std = torch.sqrt(variance).view(-1, 1, 1, 1)
+
+        z = torch.randn_like(xt).to(xt.device)
+        xt_minus_1 = mu + (~mask).float().view(-1, 1, 1, 1) * std * z
+        return xt_minus_1