TorchDiff 2.0.0__py3-none-any.whl

ddim/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)

ddim/reverse_ddim.py

@@ -0,0 +1,91 @@
+"""Reverse diffusion process for Denoising Diffusion Implicit Models (DDIM).
+
+This module implements the reverse diffusion process for DDIM, as described in Song et al.
+(2021, "Denoising Diffusion Implicit Models"). The reverse process iteratively denoises a
+noisy input to reconstruct the original data distribution using a subset of time steps.
+"""
+
+import torch
+import torch.nn as nn
+
+
+
+class ReverseDDIM(nn.Module):
+    """Reverse diffusion process of DDIM.
+
+    Implements the reverse diffusion process for Denoising Diffusion Implicit Models
+    (DDIM), which denoises a noisy input `xt` using a predicted noise component and a
+    subsampled time step schedule, as defined in Song et al. (2021).
+
+    Parameters
+    ----------
+    hyper_params : object
+        Hyperparameter object containing the noise schedule parameters. Expected to have
+        attributes:
+        - `tau_num_steps`: Number of subsampled time steps (int).
+        - `eta`: Noise scaling factor for the reverse process (float).
+        - `get_tau_schedule`: Method to compute the subsampled noise schedule (callable),
+          returning a tuple of (betas, alphas, alpha_bars, sqrt_alpha_cumprod,
+          sqrt_one_minus_alpha_cumprod).
+
+    Attributes
+    ----------
+    hyper_params : object
+        Stores the provided hyperparameter object.
+    """
+    def __init__(self, hyper_params):
+        super().__init__()
+        self.hyper_params = hyper_params
+
+    def forward(self, xt, predicted_noise, time_steps, prev_time_steps):
+        """Applies the reverse diffusion process to the noisy input.
+
+        Denoises the input `xt` at time step `t` to produce the previous step `xt_prev`
+        at `prev_time_steps` using the predicted noise and the DDIM reverse process.
+        Optionally includes stochastic noise scaled by `eta`.
+
+        Parameters
+        ----------
+        xt : torch.Tensor
+            Noisy input tensor at time step `t`, shape (batch_size, channels, height, width).
+        predicted_noise : torch.Tensor
+            Predicted noise tensor, 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.tau_num_steps - 1].
+        prev_time_steps : torch.Tensor
+            Tensor of previous time step indices (long), shape (batch_size,), where each
+            value is in the range [0, hyper_params.tau_num_steps - 1].
+
+        Returns
+        -------
+        tuple
+            A tuple containing:
+            - xt_prev: Denoised tensor at `prev_time_steps`, same shape as `xt`.
+            - x0: Estimated original data (t=0), same shape as `xt`.
+
+        Raises
+        ------
+        ValueError
+            If any value in `time_steps` or `prev_time_steps` is outside the valid range
+            [0, hyper_params.tau_num_steps - 1].
+        """
+        if not torch.all((time_steps >= 0) & (time_steps < self.hyper_params.tau_num_steps)):
+            raise ValueError(f"time_steps must be between 0 and {self.hyper_params.tau_num_steps - 1}")
+        if not torch.all((prev_time_steps >= 0) & (prev_time_steps < self.hyper_params.tau_num_steps)):
+            raise ValueError(f"prev_time_steps must be between 0 and {self.hyper_params.tau_num_steps - 1}")
+
+        _, _, _, tau_sqrt_alpha_cumprod, tau_sqrt_one_minus_alpha_cumprod = self.hyper_params.get_tau_schedule()
+        tau_sqrt_alpha_cumprod_t = tau_sqrt_alpha_cumprod[time_steps].to(xt.device).view(-1, 1, 1, 1)
+        tau_sqrt_one_minus_alpha_cumprod_t = tau_sqrt_one_minus_alpha_cumprod[time_steps].to(xt.device).view(-1, 1, 1, 1)
+        prev_tau_sqrt_alpha_cumprod_t = tau_sqrt_alpha_cumprod[prev_time_steps].to(xt.device).view(-1, 1, 1, 1)
+        prev_tau_sqrt_one_minus_alpha_cumprod_t = tau_sqrt_one_minus_alpha_cumprod[prev_time_steps].to(xt.device).view(-1, 1, 1, 1)
+
+        eta = self.hyper_params.eta
+        x0 = (xt - tau_sqrt_one_minus_alpha_cumprod_t * predicted_noise) / tau_sqrt_alpha_cumprod_t
+        noise_coeff = eta * ((tau_sqrt_one_minus_alpha_cumprod_t / prev_tau_sqrt_alpha_cumprod_t) *
+                             prev_tau_sqrt_one_minus_alpha_cumprod_t / torch.clamp(tau_sqrt_one_minus_alpha_cumprod_t, min=1e-8))
+        direction_coeff = torch.clamp(prev_tau_sqrt_one_minus_alpha_cumprod_t ** 2 - noise_coeff ** 2, min=1e-8).sqrt()
+        xt_prev = prev_tau_sqrt_alpha_cumprod_t * x0 + noise_coeff * torch.randn_like(xt) + direction_coeff * predicted_noise
+
+        return xt_prev, x0