TorchDiff 2.0.0__py3-none-any.whl

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

sde/reverse_sde.py

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+
+
+
+
+class ReverseSDE(nn.Module):
+    """Reverse diffusion process for SDE-based generative models.
+
+    Implements the reverse diffusion process for score-based generative models using
+    Stochastic Differential Equations (SDEs), supporting Variance Exploding (VE),
+    Variance Preserving (VP), sub-Variance Preserving (sub-VP), and ODE methods, as
+    described in Song et al. (2021). The reverse process denoises a noisy input using
+    predicted noise estimates.
+
+    Parameters
+    ----------
+    hyper_params : object
+        Hyperparameter object containing SDE-specific parameters. Expected to have
+        attributes:
+        - `dt`: Time step size for SDE integration (float).
+        - `sigmas`: Sigma values for VE method (torch.Tensor, optional).
+        - `betas`: Beta values for VP, sub-VP, or ODE methods (torch.Tensor).
+        - `cum_betas`: Cumulative beta values for sub-VP method (torch.Tensor, optional).
+    method : str
+        SDE method to use. Supported methods: "ve", "vp", "sub-vp", "ode".
+
+    Attributes
+    ----------
+    hyper_params : object
+        Stores the provided hyperparameter object.
+    method : str
+        Selected SDE method.
+
+    Raises
+    ------
+    ValueError
+        If `method` is not one of the supported methods ("ve", "vp", "sub-vp", "ode").
+    """
+    def __init__(self, hyper_params, method):
+        super().__init__()
+        self.hyper_params = hyper_params
+        self.method = method
+
+    def forward(self, xt, noise, predicted_noise, time_steps):
+        """Applies the reverse SDE diffusion process to the noisy input.
+
+        Denoises the input `xt` by applying the reverse SDE process, using predicted
+        noise estimates and optional stochastic noise, according to the specified SDE
+        method at given time steps. Incorporates drift and diffusion terms as applicable.
+
+        Parameters
+        ----------
+        xt : torch.Tensor
+            Noisy input tensor at time step `t`, shape (batch_size, channels, height, width).
+        noise : torch.Tensor or None
+            Gaussian noise tensor, same shape as `xt`, used for stochasticity. If None,
+            no stochastic noise is added (e.g., for deterministic ODE).
+        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.num_steps - 1].
+
+        Returns
+        -------
+        torch.Tensor
+            Denoised tensor at the previous time step, same shape as `xt`.
+
+        Raises
+        ------
+        ValueError
+            If `method` is not one of the supported methods ("ve", "vp", "sub-vp", "ode").
+
+        Notes
+        -----
+        - For the "ve" and "ode" methods, the output is clamped to [-1e5, 1e5] to prevent
+          numerical instability.
+        - Stochastic noise (`noise`) is only added if provided and the method supports it
+          (not applicable for "ode" in non-VE cases).
+        """
+        dt = self.hyper_params.dt
+        betas = self.hyper_params.betas[time_steps].view(-1, 1, 1, 1)
+        cum_betas = self.hyper_params.cum_betas[time_steps].view(-1, 1, 1, 1)
+        if self.method == "ve":
+            sigma_t = self.hyper_params.sigmas[time_steps]
+            sigma_t_prev = self.hyper_params.sigmas[time_steps - 1] if time_steps.min() > 0 else torch.zeros_like(sigma_t)
+            sigma_diff = torch.sqrt(torch.clamp(sigma_t ** 2 - sigma_t_prev ** 2, min=0))
+            drift = -(sigma_t ** 2 - sigma_t_prev ** 2).view(-1, 1, 1, 1) * predicted_noise * dt
+            diffusion = sigma_diff.view(-1, 1, 1, 1) * noise if noise is not None else 0
+            xt = xt + drift + diffusion
+            xt = torch.clamp(xt, -1e5, 1e5)
+
+        elif self.method == "vp":
+            drift = -0.5 * betas * xt * dt - betas * predicted_noise * dt
+            diffusion = torch.sqrt(betas * dt) * noise if noise is not None else 0
+            xt = xt + drift + diffusion
+
+        elif self.method == "sub-vp":
+            drift = -0.5 * betas * xt * dt - betas * (1 - torch.exp(-2 * cum_betas)) * predicted_noise * dt
+            diffusion = torch.sqrt(betas * (1 - torch.exp(-2 * cum_betas)) * dt) * noise if noise is not None else 0
+            xt = xt + drift + diffusion
+
+        elif self.method == "ode":
+            if self.method == "ve":
+                sigma_t = self.hyper_params.sigmas[time_steps]
+                sigma_t_prev = self.hyper_params.sigmas[time_steps - 1] if time_steps.min() > 0 else torch.zeros_like(sigma_t)
+                drift = -0.5 * (sigma_t ** 2 - sigma_t_prev ** 2).view(-1, 1, 1, 1) * predicted_noise * dt
+            else:
+                drift = -0.5 * betas * xt * dt - 0.5 * betas * predicted_noise * dt
+            xt = xt + drift
+            xt = torch.clamp(xt, -1e5, 1e5)
+        else:
+            raise ValueError(f"Unknown method: {self.method}")
+        return xt