TorchDiff 2.1.0__tar.gz → 2.2.0__tar.gz

{torchdiff-2.1.0 → torchdiff-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TorchDiff
-Version: 2.1.0
+Version: 2.2.0
 Summary: A PyTorch-based library for diffusion models
 Home-page: https://github.com/LoqmanSamani/TorchDiff
 Author: Loghman Samani
@@ -10,9 +10,6 @@ Project-URL: Homepage, https://loqmansamani.github.io/torchdiff
 Project-URL: Documentation, https://torchdiff.readthedio
 Project-URL: Source, https://github.com/LoqmanSamani/TorchDiff
 Keywords: diffusion models,pytorch,machine learning,deep learning
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -22,7 +19,7 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: lpips>=0.1.4
@@ -64,7 +61,7 @@ Dynamic: summary
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-red?style=plastic)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-white?style=plastic&logo=pytorch&logoColor=red)](https://pytorch.org/)
-[![Version](https://img.shields.io/badge/version-2.0.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
+[![Version](https://img.shields.io/badge/version-2.1.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
 [![Python](https://img.shields.io/badge/python-3.8%2B-blue?style=plastic&logo=python&logoColor=white)](https://www.python.org/)
 [![Downloads](https://pepy.tech/badge/torchdiff)](https://pepy.tech/project/torchdiff)
 
@@ -76,7 +73,7 @@ Dynamic: summary
 
 **TorchDiff** is a PyTorch-based library for building and experimenting with diffusion models, inspired by leading research papers.  
 
-The **TorchDiff 2.0.0** release includes implementations of five major diffusion model families:  
+The **TorchDiff 2.1.0** release includes implementations of five major diffusion model families:  
 - **DDPM** (Denoising Diffusion Probabilistic Models)  
 - **DDIM** (Denoising Diffusion Implicit Models)  
 - **SDE-based Diffusion**  
@@ -107,34 +104,71 @@ import torch.nn as nn
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
 
-from torchdiff.ddpm import VarianceSchedulerDDPM, ForwardDDPM, ReverseDDPM, TrainDDPM, SampleDDPM
-from torchdiff.utils import NoisePredictor
+from torchdiff.ddpm import (SchedulerDDPM, ForwardDDPM, 
+                            ReverseDDPM, TrainDDPM, SampleDDPM)
+from torchdiff.utils import DiffusionNetwork, mse_loss
 
-# Dataset setup
+# dataset: CIFAR10
 transform = transforms.Compose([
     transforms.Resize(32),
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))
 ])
-train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
-train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+train_dataset = datasets.CIFAR10(
+    root="./data", train=True, download=True, transform=transform
+)
+train_loader = DataLoader(
+    train_dataset, batch_size=64, shuffle=True
+)
+device = 'cuda' # gpu is used for training and sampling
+
+# model components
+diff_net = DiffusionNetwork(
+    in_channels = 3,
+    down_channels = [32, 64, 128],
+    mid_channels = [128, 128],
+    up_channels = [128, 64, 32],
+    down_sampling = [True, True],
+    time_embed_dim = 128,
+    y_embed_dim = 128,
+    num_down_blocks = 2,
+    num_mid_blocks = 2,
+    num_up_blocks = 2,
+    dropout_rate = 0.1,
+    cont_time = False # time is not continuous, if SDE models it should be true
+)
+print(sum(p.numel() for p in diff_net.parameters()))
 
-# Model components
-noise_pred = NoisePredictor(in_channels=3)
-vs = VarianceSchedulerDDPM(num_steps=1000)
-fwd, rev = ForwardDDPM(vs), ReverseDDPM(vs)
+vs = SchedulerDDPM(time_steps = 400)
+fwd = ForwardDDPM(vs, 'noise') # network is trained to predict noise
+rwd = ReverseDDPM(vs, 'noise')
 
-# Training
+# optimizer
+optim = torch.optim.Adam(diff_net.parameters(), lr=1e-5)
+
+# training algorithm
 trainer = TrainDDPM(
-    noise_predictor=noise_pred, forward_diffusion=fwd, reverse_diffusion=rev,
-    conditional_model=None, optimizer=torch.optim.Adam(noise_pred.parameters(), lr=1e-4),
-    objective=nn.MSELoss(), data_loader=train_loader, max_epochs=1, device="cpu"
+    diff_net = diff_net,
+    fwd_ddpm = fwd,
+    rwd_ddpm = rwd,
+    train_loader = train_loader,
+    optim = optim,
+    loss_fn = mse_loss,
+    max_epochs = 10,
+    device = device,
+    grad_acc = 2
 )
-trainer()
+#trainer()
 
 # Sampling
-sampler = SampleDDPM(reverse_diffusion=rev, noise_predictor=noise_pred,
-                     image_shape=(32, 32), batch_size=4, in_channels=3, device="cpu")
+sampler = SampleDDPM(
+    rwd_ddpm = rwd,
+    diff_net = diff_net,
+    img_size = (32, 32),
+    batch_size = 10,
+    in_channels = 3,
+    device = device
+)
 images = sampler()
 ```
 
@@ -164,12 +198,12 @@ DALL·E 2 architecture leveraging CLIP latents for text-to-image generation.
 TorchDiff breaks each model into reusable components:
 - **Forward Diffusion**: Adds noise to data
 - **Reverse Diffusion**: Removes noise to recover data  
-- **Variance Scheduler**: Controls noise schedules
+- **Scheduler**: Controls noise schedules
 - **Training**: Complete training pipelines
 - **Sampling**: Efficient inference and generation
 
 Additional utilities:
-- **Noise Predictor**: U-Net-like model with attention and time embeddings
+- **Diffusion Network**: U-Net-like model with attention and time embeddings
 - **Text Encoder**: Transformer-based for conditional generation
 - **Metrics**: Evaluation suite (MSE, PSNR, SSIM, FID, LPIPS)
 

{torchdiff-2.1.0 → torchdiff-2.2.0}/README.md

@@ -9,7 +9,7 @@
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-red?style=plastic)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-white?style=plastic&logo=pytorch&logoColor=red)](https://pytorch.org/)
-[![Version](https://img.shields.io/badge/version-2.0.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
+[![Version](https://img.shields.io/badge/version-2.1.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
 [![Python](https://img.shields.io/badge/python-3.8%2B-blue?style=plastic&logo=python&logoColor=white)](https://www.python.org/)
 [![Downloads](https://pepy.tech/badge/torchdiff)](https://pepy.tech/project/torchdiff)
 [![Stars](https://img.shields.io/github/stars/LoqmanSamani/TorchDiff?style=plastic&color=yellow)](https://github.com/LoqmanSamani/TorchDiff)
@@ -43,12 +43,12 @@ These models support both **conditional** (e.g., text-to-image) and **unconditio
 TorchDiff is designed with **modularity** in mind. Each model is broken down into reusable components:  
 - **Forward Diffusion**: Adds noise (e.g., `ForwardDDPM`).  
 - **Reverse Diffusion**: Removes noise to recover data (e.g., `ReverseDDPM`).  
-- **Variance Scheduler**: Controls noise schedules (e.g., `VarianceSchedulerDDPM`).  
+- **Scheduler**: Controls noise schedules (e.g., `SchedulerDDPM`).  
 - **Training**: Full training pipelines (e.g., `TrainDDPM`).  
 - **Sampling**: Efficient inference and generation (e.g., `SampleDDPM`).  
 
 Additional utilities:  
-- **Noise Predictor**: A U-Net-like model with attention and time embeddings.  
+- **Diffusion Network**: A U-Net-like model with attention and time embeddings used as main model.  
 - **Text Encoder**: Transformer-based (e.g., BERT) for conditional generation.  
 - **Metrics**: Evaluation suite including MSE, PSNR, SSIM, FID, and LPIPS.  
 
@@ -56,7 +56,7 @@ Additional utilities:
 
 ## ⚡ Quick Start  
 
-Here’s a minimal working example to train and sample with **DDPM** on dummy data:  
+Here’s a minimal working example to train and sample with **DDPM** on dummy data:
 
 ```python
 import torch
@@ -64,40 +64,72 @@ import torch.nn as nn
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
 
-from torchdiff.ddpm import VarianceSchedulerDDPM, ForwardDDPM, ReverseDDPM, TrainDDPM, SampleDDPM
-from torchdiff.utils import NoisePredictor
+from torchdiff.ddpm import (SchedulerDDPM, ForwardDDPM, 
+                            ReverseDDPM, TrainDDPM, SampleDDPM)
+from torchdiff.utils import DiffusionNetwork, mse_loss
 
-# Dataset (CIFAR10 for demo)
+# dataset: CIFAR10
 transform = transforms.Compose([
     transforms.Resize(32),
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))
 ])
-train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
-train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+train_dataset = datasets.CIFAR10(
+    root="./data", train=True, download=True, transform=transform
+)
+train_loader = DataLoader(
+    train_dataset, batch_size=64, shuffle=True
+)
+device = 'cuda' # gpu is used for training and sampling
+
+# model components
+diff_net = DiffusionNetwork(
+    in_channels = 3,
+    down_channels = [32, 64, 128],
+    mid_channels = [128, 128],
+    up_channels = [128, 64, 32],
+    down_sampling = [True, True],
+    time_embed_dim = 128,
+    y_embed_dim = 128,
+    num_down_blocks = 2,
+    num_mid_blocks = 2,
+    num_up_blocks = 2,
+    dropout_rate = 0.1,
+    cont_time = False # time is not continuous, if SDE models it should be true
+)
+print(sum(p.numel() for p in diff_net.parameters()))
 
-# Model components
-noise_pred = NoisePredictor(in_channels=3)
-vs = VarianceSchedulerDDPM(num_steps=1000)
-fwd, rev = ForwardDDPM(vs), ReverseDDPM(vs)
+vs = SchedulerDDPM(time_steps = 400)
+fwd = ForwardDDPM(vs, 'noise') # network is trained to predict noise
+rwd = ReverseDDPM(vs, 'noise')
 
-# Optimizer & loss
-optim = torch.optim.Adam(noise_pred.parameters(), lr=1e-4)
-loss_fn = nn.MSELoss()
+# optimizer
+optim = torch.optim.Adam(diff_net.parameters(), lr=1e-5)
 
-# Training
+# training algorithm
 trainer = TrainDDPM(
-    noise_predictor=noise_pred, forward_diffusion=fwd, reverse_diffusion=rev,
-    conditional_model=None, optimizer=optim, objective=loss_fn,
-    data_loader=train_loader, max_epochs=1, device="cpu"
+    diff_net = diff_net,
+    fwd_ddpm = fwd,
+    rwd_ddpm = rwd,
+    train_loader = train_loader,
+    optim = optim,
+    loss_fn = mse_loss,
+    max_epochs = 10,
+    device = device,
+    grad_acc = 2
 )
-trainer()
+#trainer()
 
 # Sampling
-sampler = SampleDDPM(reverse_diffusion=rev, noise_predictor=noise_pred,
-                     image_shape=(32, 32), batch_size=4, in_channels=3, device="cpu")
+sampler = SampleDDPM(
+    rwd_ddpm = rwd,
+    diff_net = diff_net,
+    img_size = (32, 32),
+    batch_size = 10,
+    in_channels = 3,
+    device = device
+)
 images = sampler()
-print("Generated images shape:", images.shape)
 ```
 
 For detailed examples, check the [examples/](https://github.com/LoqmanSamani/TorchDiff/tree/systembiology/examples) directory.  
@@ -130,7 +162,7 @@ pip install -r requirements.txt
 pip install .
 ```
 
-> Requires **Python 3.8+**. For GPU acceleration, ensure PyTorch is installed with the correct CUDA version.  
+> Requires **Python 3.10+**. For GPU acceleration, ensure PyTorch is installed with the correct CUDA version.  
 
 ---
 

{torchdiff-2.1.0 → torchdiff-2.2.0}/TorchDiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TorchDiff
-Version: 2.1.0
+Version: 2.2.0
 Summary: A PyTorch-based library for diffusion models
 Home-page: https://github.com/LoqmanSamani/TorchDiff
 Author: Loghman Samani
@@ -10,9 +10,6 @@ Project-URL: Homepage, https://loqmansamani.github.io/torchdiff
 Project-URL: Documentation, https://torchdiff.readthedio
 Project-URL: Source, https://github.com/LoqmanSamani/TorchDiff
 Keywords: diffusion models,pytorch,machine learning,deep learning
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -22,7 +19,7 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: lpips>=0.1.4
@@ -64,7 +61,7 @@ Dynamic: summary
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-red?style=plastic)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-white?style=plastic&logo=pytorch&logoColor=red)](https://pytorch.org/)
-[![Version](https://img.shields.io/badge/version-2.0.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
+[![Version](https://img.shields.io/badge/version-2.1.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
 [![Python](https://img.shields.io/badge/python-3.8%2B-blue?style=plastic&logo=python&logoColor=white)](https://www.python.org/)
 [![Downloads](https://pepy.tech/badge/torchdiff)](https://pepy.tech/project/torchdiff)
 
@@ -76,7 +73,7 @@ Dynamic: summary
 
 **TorchDiff** is a PyTorch-based library for building and experimenting with diffusion models, inspired by leading research papers.  
 
-The **TorchDiff 2.0.0** release includes implementations of five major diffusion model families:  
+The **TorchDiff 2.1.0** release includes implementations of five major diffusion model families:  
 - **DDPM** (Denoising Diffusion Probabilistic Models)  
 - **DDIM** (Denoising Diffusion Implicit Models)  
 - **SDE-based Diffusion**  
@@ -107,34 +104,71 @@ import torch.nn as nn
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
 
-from torchdiff.ddpm import VarianceSchedulerDDPM, ForwardDDPM, ReverseDDPM, TrainDDPM, SampleDDPM
-from torchdiff.utils import NoisePredictor
+from torchdiff.ddpm import (SchedulerDDPM, ForwardDDPM, 
+                            ReverseDDPM, TrainDDPM, SampleDDPM)
+from torchdiff.utils import DiffusionNetwork, mse_loss
 
-# Dataset setup
+# dataset: CIFAR10
 transform = transforms.Compose([
     transforms.Resize(32),
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))
 ])
-train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
-train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+train_dataset = datasets.CIFAR10(
+    root="./data", train=True, download=True, transform=transform
+)
+train_loader = DataLoader(
+    train_dataset, batch_size=64, shuffle=True
+)
+device = 'cuda' # gpu is used for training and sampling
+
+# model components
+diff_net = DiffusionNetwork(
+    in_channels = 3,
+    down_channels = [32, 64, 128],
+    mid_channels = [128, 128],
+    up_channels = [128, 64, 32],
+    down_sampling = [True, True],
+    time_embed_dim = 128,
+    y_embed_dim = 128,
+    num_down_blocks = 2,
+    num_mid_blocks = 2,
+    num_up_blocks = 2,
+    dropout_rate = 0.1,
+    cont_time = False # time is not continuous, if SDE models it should be true
+)
+print(sum(p.numel() for p in diff_net.parameters()))
 
-# Model components
-noise_pred = NoisePredictor(in_channels=3)
-vs = VarianceSchedulerDDPM(num_steps=1000)
-fwd, rev = ForwardDDPM(vs), ReverseDDPM(vs)
+vs = SchedulerDDPM(time_steps = 400)
+fwd = ForwardDDPM(vs, 'noise') # network is trained to predict noise
+rwd = ReverseDDPM(vs, 'noise')
 
-# Training
+# optimizer
+optim = torch.optim.Adam(diff_net.parameters(), lr=1e-5)
+
+# training algorithm
 trainer = TrainDDPM(
-    noise_predictor=noise_pred, forward_diffusion=fwd, reverse_diffusion=rev,
-    conditional_model=None, optimizer=torch.optim.Adam(noise_pred.parameters(), lr=1e-4),
-    objective=nn.MSELoss(), data_loader=train_loader, max_epochs=1, device="cpu"
+    diff_net = diff_net,
+    fwd_ddpm = fwd,
+    rwd_ddpm = rwd,
+    train_loader = train_loader,
+    optim = optim,
+    loss_fn = mse_loss,
+    max_epochs = 10,
+    device = device,
+    grad_acc = 2
 )
-trainer()
+#trainer()
 
 # Sampling
-sampler = SampleDDPM(reverse_diffusion=rev, noise_predictor=noise_pred,
-                     image_shape=(32, 32), batch_size=4, in_channels=3, device="cpu")
+sampler = SampleDDPM(
+    rwd_ddpm = rwd,
+    diff_net = diff_net,
+    img_size = (32, 32),
+    batch_size = 10,
+    in_channels = 3,
+    device = device
+)
 images = sampler()
 ```
 
@@ -164,12 +198,12 @@ DALL·E 2 architecture leveraging CLIP latents for text-to-image generation.
 TorchDiff breaks each model into reusable components:
 - **Forward Diffusion**: Adds noise to data
 - **Reverse Diffusion**: Removes noise to recover data  
-- **Variance Scheduler**: Controls noise schedules
+- **Scheduler**: Controls noise schedules
 - **Training**: Complete training pipelines
 - **Sampling**: Efficient inference and generation
 
 Additional utilities:
-- **Noise Predictor**: U-Net-like model with attention and time embeddings
+- **Diffusion Network**: U-Net-like model with attention and time embeddings
 - **Text Encoder**: Transformer-based for conditional generation
 - **Metrics**: Evaluation suite (MSE, PSNR, SSIM, FID, LPIPS)
 

{torchdiff-2.1.0 → torchdiff-2.2.0}/TorchDiff.egg-info/SOURCES.txt

@@ -8,38 +8,29 @@ TorchDiff.egg-info/requires.txt
 TorchDiff.egg-info/top_level.txt
 ddim/__init__.py
 ddim/forward_ddim.py
-ddim/hyper_param.py
-ddim/noise_predictor.py
 ddim/reverse_ddim.py
 ddim/sample_ddim.py
-ddim/text_encoder.py
+ddim/scheduler.py
+ddim/test_ddim.py
 ddim/train_ddim.py
 ddpm/__init__.py
 ddpm/forward_ddpm.py
-ddpm/hyper_param.py
-ddpm/noise_predictor.py
 ddpm/reverse_ddpm.py
 ddpm/sample_ddpm.py
-ddpm/text_encoder.py
+ddpm/scheduler.py
+ddpm/test_ddpm.py
 ddpm/train_ddpm.py
 ldm/__init__.py
 ldm/autoencoder.py
-ldm/forward_idm.py
-ldm/hyper_param.py
-ldm/metrics.py
-ldm/noise_predictor.py
-ldm/reverse_ldm.py
 ldm/sample_ldm.py
-ldm/text_encoder.py
 ldm/train_autoencoder.py
 ldm/train_ldm.py
 sde/__init__.py
 sde/forward_sde.py
-sde/hyper_param.py
-sde/noise_predictor.py
 sde/reverse_sde.py
 sde/sample_sde.py
-sde/text_encoder.py
+sde/scheduler.py
+sde/test_sde.py
 sde/train_sde.py
 torchdiff/__init__.py
 torchdiff/ddim.py
@@ -68,4 +59,9 @@ unclip/unclip_sampler.py
 unclip/upsampler.py
 unclip/upsampler_trainer.py
 unclip/utils.py
-unclip/val_metrics.py
+unclip/val_metrics.py
+utils/__init__.py
+utils/diff_net.py
+utils/losses.py
+utils/metrics.py
+utils/text_encoder.py

{torchdiff-2.1.0 → torchdiff-2.2.0}/TorchDiff.egg-info/top_level.txt

@@ -4,3 +4,4 @@ ldm
 sde
 torchdiff
 unclip
+utils

torchdiff-2.2.0/ddim/forward_ddim.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+
+
+
+class ForwardDDIM(nn.Module):
+    """
+    Implements the forward (noising) process of DDIM.
+
+    This module samples x_t from the forward diffusion distribution:
+
+        q(x_t | x_0) = N(x_t; sqrt(alphā_t) * x_0, (1 - alphā_t) * I)
+
+    It also computes the appropriate training target depending on the
+    prediction parameterization (noise, x0, or v-prediction).
+
+    Args:
+        scheduler: Noise scheduler containing precomputed diffusion coefficients.
+        pred_type: Type of model prediction. One of ["noise", "x0", "v"].
+    """
+    def __init__(
+            self,
+            scheduler: nn.Module,
+            pred_type: str = "noise"
+    ):
+        super().__init__()
+        valid_types = ["noise", "x0", "v"]
+        if pred_type not in valid_types:
+            raise ValueError(f"prediction_type must be one of {valid_types}, got {pred_type}")
+        self.vs = scheduler
+        self.pred_type = pred_type
+
+    def forward(self, x0: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Perform the forward diffusion step and compute the training target.
+
+        Samples x_t by adding noise to the clean input x_0 at timestep t,
+        and returns the corresponding supervision target for training.
+
+        Args:
+            x0: Clean input data of shape (batch, ...).
+            t: Discrete diffusion timesteps of shape (batch,).
+            noise: Gaussian noise of same shape as x0.
+
+        Returns:
+            xt: Noised data x_t of shape (batch, ...).
+            target: Training target corresponding to pred_type:
+                - "noise": the added noise ε
+                - "x0": the original clean input x0
+                - "v": the velocity parameterization
+        """
+        sqrt_alpha_cumprod_t = self.vs.sqrt_alphas_cumprod[t]
+        sqrt_one_minus_alpha_cumprod_t = self.vs.sqrt_one_minus_alphas_cumprod[t]
+        sqrt_alpha_cumprod_t = self.vs.get_index(sqrt_alpha_cumprod_t, x0.shape)
+        sqrt_one_minus_alpha_cumprod_t = self.vs.get_index(sqrt_one_minus_alpha_cumprod_t, x0.shape)
+        xt = sqrt_alpha_cumprod_t * x0 + sqrt_one_minus_alpha_cumprod_t * noise
+        if self.pred_type == "noise":
+            target = noise
+        elif self.pred_type == "x0":
+            target = x0
+        elif self.pred_type == "v":
+            target = sqrt_alpha_cumprod_t * noise - sqrt_one_minus_alpha_cumprod_t * x0
+        return xt, target