trying to improve whole project

19 files changed, 1737 insertions, 1176 deletions

diff --git a/.gitignore b/.gitignore
index 0c22bca..c1eb923 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,172 +1,17 @@
-# Tensorboard
-runs/
+# Python virtual environment
+.venv/
 
-# Torch datas.
-data/
+# Plots
+plots/
 
-# Model weights.
-*.pth
+# Checkpoints
+checkpoints/
 
-# Tmp files.
-*.tmp.*
+# MLflow artifacts
+mlflow.db
 
-# Byte-compiled / optimized / DLL files
+# Python bytecode
 __pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
 
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+# Data
+data/
diff --git a/README.md b/README.md
index 93cdc90..805baae 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,79 @@
-# diffusion-mnist
-Diffusion inspired mnist like image generation.
+# MNIST Diffusion Model
+
+A diffusion-based generative model for MNIST digits implemented in PyTorch.
+
+## Project Structure
+
+```
+diffusion-mnist/
+├── main.py                 # Main entry point with CLI
+├── models/                 # Neural network architectures
+│   ├── __init__.py
+│   ├── unet.py            # UNet for diffusion model
+│   └── autoencoder.py     # Autoencoder for latent diffusion
+├── src/                   # Source code modules
+│   ├── __init__.py
+│   ├── config.py          # Configuration and hyperparameters
+│   ├── diffusion.py       # Diffusion process utilities
+│   ├── dataloader.py      # Dataset and dataloader classes
+│   ├── utils.py           # Helper functions and metrics
+│   ├── train_diffusion.py # Diffusion training script
+│   ├── train_autoencoder.py # Autoencoder training script
+│   └── sample.py          # Sampling and visualization
+├── checkpoints/           # Model checkpoints
+├── plots/                 # Generated visualizations
+├── data/                  # Dataset directory
+└── runs/                  # TensorBoard logs
+```
+
+## Installation
+
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+
+### Train Diffusion Model
+```bash
+python main.py train --epochs 10 --lr 2e-4 --batch-size 64
+```
+
+### Train with Self-Attention
+```bash
+python main.py train --epochs 10 --attention
+```
+
+### Train Autoencoder (for latent diffusion)
+```bash
+python main.py train-ae --epochs 10
+```
+
+### Generate Samples
+```bash
+python main.py sample --checkpoint checkpoints/diffusion_latest.pt
+```
+
+### Visualize Diffusion Process
+```bash
+python main.py visualize --all
+```
+
+## Configuration
+
+All hyperparameters can be found in `src/config.py`:
+- `DIFFU_STEPS`: Number of diffusion steps (default: 1000)
+- `EPOCHS`: Training epochs (default: 10)
+- `BATCH_SIZE`: Batch size (default: 64)
+- `LEARNING_RATE`: Learning rate (default: 2e-4)
+
+## Model Architecture
+
+The diffusion model uses a UNet architecture with:
+- Timestep embedding
+- Label conditioning (for class-conditional generation)
+- Optional self-attention layers
+
+## License
+
+See [LICENSE](LICENSE) for details.
diff --git a/autoencoder.py b/autoencoder.py
deleted file mode 100644
index 23c47f9..0000000
--- a/autoencoder.py
+++ /dev/null
@@ -1,154 +0,0 @@
-import torch
-from torch.utils.data import DataLoader, Dataset
-
-from torchvision import datasets, transforms
-
-import matplotlib.pyplot as plt
-
-from trainer import train
-
-
-class PrintLayer(torch.nn.Module):
-    def forward(self, x):
-        print(x.shape)
-        print(x.min())
-        print(x.max())
-        return x
-
-
-class Autoencoder(torch.nn.Module):
-    def __init__(self, input_dim, latent_dim):
-        super(Autoencoder, self).__init__()
-        self.input_dim = input_dim
-        self.latent_dim = latent_dim
-        self.encoder = torch.nn.Sequential(
-            torch.nn.Conv2d(input_dim, 16, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(16, 16, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.MaxPool2d(kernel_size=2),
-            torch.nn.Conv2d(16, 32, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 32, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.MaxPool2d(kernel_size=2),
-            torch.nn.Conv2d(32, 64, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, latent_dim, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            # 1x7x7 to 1x8x8
-            torch.nn.Conv2d(latent_dim, latent_dim, kernel_size=2, padding=1),
-            torch.nn.Sigmoid(),
-        )
-        self.decoder = torch.nn.Sequential(
-            # 1x8x8 to 1x7x7
-            torch.nn.Conv2d(latent_dim, latent_dim, kernel_size=2, padding=0),
-            torch.nn.ReLU(),
-            # main decoder
-            torch.nn.Conv2d(latent_dim, 64, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 32, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(16, 16, kernel_size=3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(16, input_dim, kernel_size=3, padding=1),
-            torch.nn.Sigmoid(),
-        )
-
-    def forward(self, x):
-        x = self.encoder(x)
-        x = self.decoder(x)
-        return x
-
-    def encode(self, x):
-        return self.encoder(x)
-
-    def decode(self, x):
-        return self.decoder(x)
-
-
-class AutoencoderDataset(Dataset):
-    def __init__(self, dataset, device='cpu'):
-        self.dataset = dataset
-        self.device = device
-        self.dummy_param = torch.nn.Parameter(torch.empty(0))
-
-    def __len__(self):
-        return len(self.dataset)
-
-    def __getitem__(self, idx):
-        data = self.dataset[idx][0].to(self.device)
-        return data, data
-
-def main():
-    torch.multiprocessing.set_start_method("spawn")
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-    # Load dataset
-    mnist = datasets.MNIST(
-        root='data',
-        train=True,
-        download=True,
-        transform=transforms.ToTensor(),
-    )
-    dataset = AutoencoderDataset(mnist, device=device)
-    dataloader = DataLoader(dataset, batch_size=64, shuffle=True,
-                            num_workers=4, persistent_workers=True)
-
-    # Initialize model
-    model = Autoencoder(input_dim=1, latent_dim=1)
-    model.load_state_dict(torch.load('mnist_autoencoder.pth'))
-    model.to(device)
-
-    # Train model
-    lr = 1e-3
-    epochs = 1
-    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
-    criterion = torch.nn.functional.binary_cross_entropy
-    train(model, dataloader, epochs, optimizer, criterion)
-
-    # Save model
-    torch.save(model.state_dict(), 'autoencoder.pth')
-
-    # Visualize results
-    n = 10
-    with torch.no_grad():
-        plt.figure(figsize=(2*n, 6))
-        for i, j in enumerate(torch.randint(0, len(dataset), (n,))):
-            x, _ = dataset[j]
-            x = x.unsqueeze(0)
-            x_latent = model.encode(x)
-            x_hat = model.decode(x_latent)
-            plt.subplot(3, n+1, i + 2)
-            plt.imshow(x.cpu().squeeze().numpy())
-            plt.axis('off')
-            plt.subplot(3, n+1, i + n + 3)
-            plt.imshow(x_latent.cpu().squeeze().numpy())
-            plt.axis('off')
-            plt.subplot(3, n+1, i + 2*n + 4)
-            plt.imshow(x_hat.cpu().squeeze().numpy())
-            plt.axis('off')
-        plt.subplot(3, n+1, 1)
-        plt.text(0.5, 0.5, 'Original', horizontalalignment='center', fontsize=12)
-        plt.axis('off')
-        plt.subplot(3, n+1, n + 2)
-        plt.text(0.5, 0.5, 'Latent', horizontalalignment='center', fontsize=12)
-        plt.axis('off')
-        plt.subplot(3, n+1, 2*n + 3)
-        plt.text(0.5, 0.5, 'Reconstructed', horizontalalignment='center', fontsize=12)
-        plt.axis('off')
-        plt.suptitle('Autoencoder')
-        plt.tight_layout()
-        plt.savefig('autoencoder.tmp.png')
-
-
-if __name__ == '__main__':
-    main()
diff --git a/main.py b/main.py
index 7f8444a..46c8a50 100644
--- a/main.py
+++ b/main.py
@@ -1,580 +1,137 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.attention as attention
-from torch.utils.data import DataLoader, Dataset
-from rich.progress import track
-import io
-from PIL import Image
-
-from torchvision import datasets, transforms
-
-from trainer import train
-from trainer.trainer import Trainer
-
-import matplotlib.pyplot as plt
-
-import numpy as np
-import os
-import cv2
-
-from autoencoder import Autoencoder
-from utils import *
-
-
-class SelfAttention(nn.Module):
-    def __init__(self, nb_channels, nb_heads):
-        super().__init__()
-        self.attention = nn.MultiheadAttention(nb_channels, nb_heads)
-
-    def forward(self, x):
-        _, c, w, h = x.shape
-        x = x.view(-1, c, w*h)
-        x = x.permute(2, 0, 1)
-        x, _ = self.attention(x, x, x)
-        x = x.permute(1, 2, 0)
-        x = x.view(-1, c, w, h)
-        return x
-
-
-class UNetAddAttUEDF(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        ## Inputs:
-        # xt: image at step t (NB_CHANNEL*IMG_SIZE*IMG_SIZE)
-        # t: step number (1)
-        # vec: one-hot vector of the label (NB_LABEL)
-
-        ## Encoder for t
-        self.encodet = nn.Linear(1, IMG_SIZE*IMG_SIZE)
-
-        ## Encoder for vec
-        self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE*IMG_SIZE)
-
-        ## UNet (2 more channels input because we concatenate xt with t and vec)
-        self.conv1 = nn.Conv2d(NB_CHANNEL+2, 64, 3, padding=1)
-        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
-        self.maxpool1 = nn.MaxPool2d(2, 2)
-        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
-        self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
-        self.maxpool2 = nn.MaxPool2d(2, 2)
-        self.att1 = SelfAttention(128, 8)
-        self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
-        self.att2 = SelfAttention(256, 8)
-        self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
-        self.att3 = SelfAttention(256, 8)
-        self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
-        self.conv7 = nn.Conv2d(256, 128, 3, padding=1)
-        self.conv8 = nn.Conv2d(128, 128, 3, padding=1)
-        self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
-        self.conv9 = nn.Conv2d(128, 64, 3, padding=1)
-        self.conv10 = nn.Conv2d(64, 64, 3, padding=1)
-        self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1)
-
-    def forward(self, xt, t, vec):
-        # Encode t and vec
-        t = F.relu(self.encodet(t / DIFFU_STEPS))
-        t = t.view(-1, 1, IMG_SIZE, IMG_SIZE)
-        vec = F.relu(self.encodevec(vec))
-        vec = vec.view(-1, 1, IMG_SIZE, IMG_SIZE)
-
-        # Concat all 3
-        x = torch.cat((xt, t, vec), dim=1)
-
-        # UNet
-        x1 = F.relu(self.conv1(x))
-        x1 = F.relu(self.conv2(x1))
-        x2 = self.maxpool1(x1)
-        x2 = F.relu(self.conv3(x2))
-        x2 = F.relu(self.conv4(x2))
-        x3 = self.maxpool2(x2)
-        x3 = self.att1(x3)
-        x3 = F.relu(self.conv5(x3))
-        x3 = self.att2(x3)
-        x3 = F.relu(self.conv6(x3))
-        x4 = self.upconv1(x3)
-        x4 = torch.cat((x4, x2), dim=1)
-        x4 = F.relu(self.conv7(x4))
-        x4 = F.relu(self.conv8(x4))
-        x5 = self.upconv2(x4)
-        x5 = torch.cat((x5, x1), dim=1)
-        x5 = F.relu(self.conv9(x5))
-        x5 = F.relu(self.conv10(x5))
-        x5 = self.conv11(x5)
-
-        return x5
-
-
-class UNetMNISTSig(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        ## Inputs:
-        # xt: image at step t (NB_CHANNEL*IMG_SIZE*IMG_SIZE)
-        # t: step number (1)
-        # vec: one-hot vector of the label (NB_LABEL)
-
-        ## Encoder for t
-        self.encodet = nn.Linear(1, IMG_SIZE*IMG_SIZE)
-
-        ## Encoder for vec
-        self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE*IMG_SIZE)
-
-        ## UNet (2 more channels input because we concatenate xt with t and vec)
-        self.conv1 = nn.Conv2d(NB_CHANNEL+2, 64, 3, padding=1)
-        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
-        self.maxpool1 = nn.MaxPool2d(2, 2)
-        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
-        self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
-        self.maxpool2 = nn.MaxPool2d(2, 2)
-        self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
-        self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
-        self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
-        self.conv7 = nn.Conv2d(256, 128, 3, padding=1)
-        self.conv8 = nn.Conv2d(128, 128, 3, padding=1)
-        self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
-        self.conv9 = nn.Conv2d(128, 64, 3, padding=1)
-        self.conv10 = nn.Conv2d(64, 64, 3, padding=1)
-        self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1)
-
-    def forward(self, xt, t, vec):
-        # Encode t and vec
-        t = F.relu(self.encodet(t / DIFFU_STEPS))
-        t = t.view(-1, 1, IMG_SIZE, IMG_SIZE)
-        vec = F.relu(self.encodevec(vec))
-        vec = vec.view(-1, 1, IMG_SIZE, IMG_SIZE)
-
-        # Concat all 3
-        x = torch.cat((xt, t, vec), dim=1)
-
-        # UNet
-        x1 = F.relu(self.conv1(x))
-        x1 = F.relu(self.conv2(x1))
-        x2 = self.maxpool1(x1)
-        x2 = F.relu(self.conv3(x2))
-        x2 = F.relu(self.conv4(x2))
-        x3 = self.maxpool2(x2)
-        x3 = F.relu(self.conv5(x3))
-        x3 = F.relu(self.conv6(x3))
-        x4 = self.upconv1(x3)
-        x4 = torch.cat((x4, x2), dim=1)
-        x4 = F.relu(self.conv7(x4))
-        x4 = F.relu(self.conv8(x4))
-        x5 = self.upconv2(x4)
-        x5 = torch.cat((x5, x1), dim=1)
-        x5 = F.relu(self.conv9(x5))
-        x5 = F.relu(self.conv10(x5))
-        x5 = self.conv11(x5)
-
-        return x5
-
-class FolderDataset(Dataset):
-    def __init__(self, path, size=(32, 32)):
-        super().__init__()
-        self.path = path
-        self.size = size
-        self.files = os.listdir(self.path)
+"""
+MNIST Diffusion Model
 
-    def __getitem__(self, index):
-        img = cv2.imread(os.path.join(self.path, self.files[index]))
-        img = cv2.resize(img, self.size)
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        img = np.transpose(img, (2, 0, 1)) / 255
-        return torch.tensor(img, dtype=torch.float32), 0
+A diffusion-based generative model for MNIST digits.
 
-    def __len__(self):
-        return len(self.files)
+Usage:
+    Train diffusion model:
+        python main.py train --epochs 10
 
+    Train autoencoder:
+        python main.py train-ae --epochs 10
 
-def q_xt_xt_1(xt_1, t):
-    t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
+    Generate samples:
+        python main.py sample --checkpoint checkpoints/diffusion_latest.pt
 
-    alpha = ALPHA[t_ind]
-    mean = torch.sqrt(alpha) * xt_1
-    std = torch.sqrt(1 - alpha)
+    Visualize diffusion process:
+        python main.py visualize --all
+"""
 
-    eps = torch.randn(xt_1.shape, device=DEVICE)
-    xt = mean + std * eps
-
-    return xt, eps
-
-
-def q_xt_x0(x0, t):
-    t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
-
-    reshaped = len(x0.shape) == 3
-    if reshaped:
-        c, w, h = x0.shape
-        b = 1
-        x0 = x0.view(b, c, w, h)
-    else:
-        b, c, w, h = x0.shape
-    t_ind = t_ind.view(b, 1, 1, 1)
-    t_ind = t_ind.expand(b, c, w, h)
-
-    alpha_bar = ALPHA_BAR[t_ind]
-    mean = torch.sqrt(alpha_bar) * x0
-    std = torch.sqrt(1 - alpha_bar)
-
-    eps = torch.randn(x0.shape, device=DEVICE)
-    xt = mean + std * eps
-
-    if reshaped:
-        xt = xt.view(c, w, h)
-
-    return xt, eps
-
-
-def p_xt_1_xt(model, xt, t, vec):
-    t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
-
-    alpha_bar_t = ALPHA_BAR[t_ind].view(-1, 1, 1, 1)
-    alpha_bar_t_1 = ALPHA_BAR[t_ind-1].view(-1, 1, 1, 1)
-    alpha_t = ALPHA[t_ind].view(-1, 1, 1, 1)
-    beta_t = BETA[t_ind].view(-1, 1, 1, 1)
-
-    beta_tilde = (1 - alpha_bar_t_1) / (1 - alpha_bar_t) * beta_t
-
-    epsilon_theta = model(xt, t, vec)
-
-    sigma_theta = torch.sqrt(beta_tilde)
-    mu_theta = (xt - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * epsilon_theta) / torch.sqrt(alpha_t)
-
-    mask_t0 = (t > 1).to(dtype=torch.float32).view(-1, 1, 1, 1)
-    noise = torch.randn(xt.shape, device=DEVICE) * mask_t0
-
-    return mu_theta + sigma_theta * noise
-
-
-def p_xt_1_xt_sig(model, xt_1, t, vec):
-    x0 = model(xt_1, t, vec)
-    xt, _ = q_xt_x0(x0, t)
-    return xt
-
-
-class DiffusionDataset(Dataset):
-    def __init__(self, dataset, autoencoder=None):
-        super().__init__()
-        self.dataset = dataset
-        self.autoencoder = autoencoder
-
-    def __getitem__(self, index):
-        # Get the image and the label.
-        img, label = self.dataset[index]
-        img = img.to(DEVICE)
-
-        # Encode the image.
-        if self.autoencoder is not None:
-            img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0)
-
-        # Normalize the image.
-        img = img * 2 - 1
-
-        # Add noise to the image.
-        t = torch.randint(1, DIFFU_STEPS, (1,), device=DEVICE)
-        xt, eps = q_xt_x0(img, t)
+import argparse
+import torch
 
-        # Convert the label to a one-hot vector.
-        vec = torch.nn.functional.one_hot(
-            torch.tensor(min(label, NB_LABEL-1)),
-            num_classes=NB_LABEL,
-        )
 
-        return (
-            # x_true
-            xt.clone().detach().to(dtype=torch.float32, device=DEVICE),
-            t.clone().detach().to(dtype=torch.float32, device=DEVICE),
-            vec.clone().detach().to(dtype=torch.float32, device=DEVICE),
-            # y_true
-            eps,
+def main():
+    parser = argparse.ArgumentParser(
+        description="MNIST Diffusion Model",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__
+    )
+    
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+    
+    # Train diffusion model
+    train_parser = subparsers.add_parser("train", help="Train diffusion model")
+    train_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
+    train_parser.add_argument("--lr", type=float, default=2e-4, help="Learning rate")
+    train_parser.add_argument("--batch-size", type=int, default=64, help="Batch size")
+    train_parser.add_argument("--attention", action="store_true", help="Use self-attention")
+    train_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint")
+    train_parser.add_argument("--name", type=str, default=None, help="Run name")
+    
+    # Train autoencoder
+    ae_parser = subparsers.add_parser("train-ae", help="Train autoencoder")
+    ae_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
+    ae_parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
+    ae_parser.add_argument("--batch-size", type=int, default=64, help="Batch size")
+    ae_parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels")
+    ae_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint")
+    
+    # Sample from model
+    sample_parser = subparsers.add_parser("sample", help="Generate samples")
+    sample_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt",
+                               help="Path to model checkpoint")
+    sample_parser.add_argument("--n-samples", type=int, default=10, help="Samples per class")
+    sample_parser.add_argument("--attention", action="store_true", help="Use attention in model")
+    
+    # Visualize diffusion
+    viz_parser = subparsers.add_parser("visualize", help="Visualize diffusion process")
+    viz_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt",
+                            help="Path to model checkpoint")
+    viz_parser.add_argument("--attention", action="store_true", help="Use attention in model")
+    viz_parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion")
+    viz_parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion")
+    viz_parser.add_argument("--all", action="store_true", help="Run all visualizations")
+    
+    args = parser.parse_args()
+    
+    if args.command is None:
+        parser.print_help()
+        return
+    
+    # Set multiprocessing start method
+    torch.multiprocessing.set_start_method("spawn", force=True)
+    
+    if args.command == "train":
+        from src.train_diffusion import train_diffusion
+        train_diffusion(
+            epochs=args.epochs,
+            learning_rate=args.lr,
+            batch_size=args.batch_size,
+            use_attention=args.attention,
+            checkpoint_path=args.checkpoint,
+            run_name=args.name,
         )
-
-    def __len__(self):
-        return len(self.dataset)
-
-
-class DiffusionDatasetSig(Dataset):
-    def __init__(self, dataset, autoencoder=None):
-        super().__init__()
-        self.dataset = dataset
-        self.autoencoder = autoencoder
-
-    def __getitem__(self, index):
-        # Get the image and the label.
-        img, label = self.dataset[index]
-        img = img.to(DEVICE)
-
-        # Encode the image.
-        if self.autoencoder is not None:
-            img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0)
-
-        # Normalize the image.
-        img = img * 2 - 1
-
-        # Add noise to the image.
-        t = torch.randint(1, DIFFU_STEPS, (1,), device=DEVICE)
-        xt, eps = q_xt_x0(img, t)
-
-        # Convert the label to a one-hot vector.
-        vec = torch.nn.functional.one_hot(
-            torch.tensor(min(label, NB_LABEL-1)),
-            num_classes=NB_LABEL,
+    
+    elif args.command == "train-ae":
+        from src.train_autoencoder import train_autoencoder
+        train_autoencoder(
+            epochs=args.epochs,
+            learning_rate=args.lr,
+            batch_size=args.batch_size,
+            latent_channels=args.latent_channels,
+            checkpoint_path=args.checkpoint,
         )
-
-        return (
-            # x_true
-            xt.clone().detach().to(dtype=torch.float32, device=DEVICE),
-            t.clone().detach().to(dtype=torch.float32, device=DEVICE),
-            vec.clone().detach().to(dtype=torch.float32, device=DEVICE),
-            # y_true
-            img.clone().detach().to(dtype=torch.float32, device=DEVICE),
+    
+    elif args.command == "sample":
+        import os
+        from src.config import DEVICE
+        from src.sample import generate_grid
+        from src.utils import load_checkpoint
+        from models import UNetMNIST
+        
+        model = UNetMNIST(use_attention=args.attention).to(DEVICE)
+        if os.path.exists(args.checkpoint):
+            model = load_checkpoint(model, os.path.basename(args.checkpoint))
+        else:
+            print(f"Warning: Checkpoint {args.checkpoint} not found.")
+        
+        generate_grid(model, n_per_class=args.n_samples)
+    
+    elif args.command == "visualize":
+        import os
+        from src.config import DEVICE
+        from src.sample import (
+            visualize_forward_diffusion,
+            visualize_backward_diffusion,
+            generate_grid,
         )
+        from src.utils import load_checkpoint
+        from models import UNetMNIST
+        
+        model = UNetMNIST(use_attention=args.attention).to(DEVICE)
+        if os.path.exists(args.checkpoint):
+            model = load_checkpoint(model, os.path.basename(args.checkpoint))
+        
+        if args.forward or args.all:
+            visualize_forward_diffusion()
+        
+        if args.backward or args.all:
+            visualize_backward_diffusion(model)
+        
+        if args.all:
+            generate_grid(model)
 
-    def __len__(self):
-        return len(self.dataset)
-
-
-def loss(y_pred, y_true):
-    return nn.MSELoss()(y_pred, y_true)
-
-
-def forward_diffusion(x0):
-    x = x0.clone()
-    xs = [x]
-    for t in range(1, DIFFU_STEPS+1):
-        x, _ = q_xt_xt_1(x, t)
-        xs.append(x)
-    return xs
-
-
-def tensor_to_image(tensor):
-    img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0)
-    img -= img.min()
-    img /= img.max()
-    return img
-
-
-def tensor_to_images(tensor):
-    img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1)
-    img -= np.min(img, axis=(1, 2, 3), keepdims=True)
-    img /= np.max(img, axis=(1, 2, 3), keepdims=True)
-    return img
-
-
-def figure_to_image(figure):
-    buf = io.BytesIO()
-    figure.savefig(buf, format='png')
-    buf.seek(0)
-    image = np.array(Image.open(buf))
-    return image
-
-
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-DIFFU_STEPS = 1000
-BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE)
-BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA))
-ALPHA = 1 - BETA
-ALPHA_BAR = torch.cumprod(ALPHA, dim=0)
-
-dataset = datasets.MNIST(
-    root="./data",
-    train=True,
-    download=True,
-    transform=transforms.ToTensor(),
-)
-# dataset = datasets.LFWPeople(
-#     root="./data",
-#     download=True,
-#     transform=transforms.Compose([
-#         transforms.Resize((64, 64)),
-#         transforms.ToTensor(),
-#     ]),
-# )
-# dataset = FolderDataset('data/lfwcrop_color/faces')
-# dataset = FolderDataset('data/edface')
-
-autoencoder = None
-# autoencoder = Autoencoder(1, 1).to(DEVICE)
-# autoencoder.load_state_dict(torch.load('autoencoder.pth'))
-# autoencoder.eval()
-
-img = dataset[0][0].to(DEVICE)
-if autoencoder is not None:
-    img = autoencoder.encode(img.unsqueeze(0)).squeeze(0)
-NB_CHANNEL, IMG_SIZE, _ = img.shape
-NB_LABEL = 10
-
-EPOCHS = 10
-LEARNING_RATE = 2e-4
-
-
-def epoch_callback(trainer: Trainer):
-    epoch_i = trainer.epoch_i
-
-    if epoch_i % 1 == 0 or epoch_i == trainer.epoch_end:
-        save_path = f'runs/{trainer.run_name}/checkpoints/{epoch_i:04}e.pt'
-        torch.save(trainer.model, save_path)
-        save_path = f'runs/{trainer.run_name}/checkpoints/last.pt'
-        torch.save(trainer.model, save_path)
-
-        print("Calculating metrics...")
-        with torch.no_grad():
-            batch_size = 64
-            n_batches = 16
-            n_samples = batch_size * n_batches
-
-            fakes = np.zeros((0, NB_CHANNEL, IMG_SIZE, IMG_SIZE))
-            for _ in range(n_batches):
-                x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE)
-                vec = torch.randint(0, NB_LABEL, (batch_size,)).to(DEVICE)
-                vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
-
-                for t in range(DIFFU_STEPS, 0, -1):
-                    t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32)
-                    x = p_xt_1_xt_sig(model, x, t_tensor, vec)
-
-                x = x.cpu().numpy()
-                x -= x.min(axis=(1, 2, 3), keepdims=True)
-                x /= x.max(axis=(1, 2, 3), keepdims=True)
-                fakes = np.concatenate((fakes, x))
-
-            reals = torch.stack([dataset[i][0] for i in range(n_samples)]).cpu().numpy()
-            reals = reals * 2 - 1
-
-            trainer.writer.add_scalar('FID/Validation', fid(reals, fakes), epoch_i)
-            trainer.writer.add_scalar('KL/Validation', kl(reals, fakes), epoch_i)
-            trainer.writer.add_scalar('RKL/Validation', kl(fakes, reals), epoch_i)
-            trainer.writer.add_scalar('JSD/Validation', jsd(reals, fakes), epoch_i)
-
-
-            fig = plt.figure(figsize=(32, 16))
-            for i in range(32):
-                plt.subplot(4, 8, i + 1)
-                plt.imshow(fakes[i].transpose(1, 2, 0))
-                plt.axis("off")
-            plt.tight_layout()
-            trainer.writer.add_image('Fakes/Validation', figure_to_image(fig), epoch_i, dataformats='HWC')
-            plt.close(fig)
-
-
-if __name__ == '__main__':
-
-    ############
-    # Training #
-    ############
-
-    torch.multiprocessing.set_start_method("spawn")
-
-    # Load the model.
-    # model = UNetMNISTSig().to(DEVICE)
-    model = torch.load('runs/20240719-154613_UNetMNISTSig/checkpoints/last.pt').to(DEVICE)
-
-    # Define the optimizer.
-    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
-
-    # Define the training dataset.
-    train_dataset = DiffusionDatasetSig(dataset, autoencoder)
-    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True,
-                              num_workers=4, persistent_workers=True)
-    criterion = loss
-    epochs = EPOCHS
-
-    # Train the model.
-    train(model, train_loader, epochs, optimizer, criterion, epoch_callbacks=[epoch_callback], save_chekpoint=False)
-
-    ############## 
-    # Evaluation #
-    ##############
-
-    with torch.no_grad():
-        # Forward diffusion
-        img, label = dataset[np.random.randint(0, len(dataset))]
-        img = img.to(DEVICE)
-        if autoencoder is not None:
-            img = autoencoder.encode(img.unsqueeze(0)).squeeze(0)
-        img = img * 2 - 1
-
-        nb_plots = 10
-        plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)]
-
-        xs = forward_diffusion(img)
-
-        plt.figure(figsize=(nb_plots, 2.5))
-        for plot_i, t in enumerate(plots_id):
-            x = xs[t]
-            plt.subplot(2, nb_plots + 1, plot_i + 2)
-            plt.title(f"t={t}")
-            plt.imshow(tensor_to_image(x), interpolation='none')
-            plt.axis("off")
-
-        for plot_i, t in enumerate(plots_id):
-            x, _ = q_xt_x0(img, t)
-            plt.subplot(2, nb_plots + 1, nb_plots + plot_i + 3)
-            plt.imshow(tensor_to_image(x), interpolation='none')
-            plt.axis("off")
-
-        plt.subplot(2, nb_plots + 1, 1)
-        plt.text(0, 0.5, "Implicit", fontsize=12)
-        plt.axis("off")
-        plt.subplot(2, nb_plots + 1, nb_plots + 2)
-        plt.text(0, 0.5, "Explicit", fontsize=12)
-        plt.axis("off")
-
-        plt.suptitle("Forward diffusion")
-        plt.tight_layout()
-        plt.savefig("plots/forward_diffusion.tmp.png")
-
-
-        # Backward diffusion
-        t_plots = np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)
-
-        n_classes = 10
-
-        x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE)
-        vec = torch.tensor([[min(i, NB_LABEL-1)] for i in range(n_classes)], dtype=torch.int64)
-        vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
-
-        plt.figure(figsize=(nb_plots, n_classes))
-        plt.suptitle("Backward diffusion")
-        for t in track(range(DIFFU_STEPS, 0, -1), description='Diffusing...'):
-            t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32)
-            x = p_xt_1_xt_sig(model, x, t_tensor, vec)
-            if t in t_plots:
-                t_plot_i = nb_plots - t_plots.tolist().index(t) - 1
-                for class_i in range(n_classes):
-                    plt.subplot(n_classes, nb_plots, t_plot_i + nb_plots * class_i + 1)
-                    if class_i == 0:
-                        plt.title(f"t={t}")
-                    plt.imshow(tensor_to_image(x[class_i]))
-                    plt.axis("off")
-        plt.tight_layout()
-        plt.savefig("plots/backward_diffusion.tmp.png")
-
-
-        # Benchmark
-        x = torch.randn(nb_plots * n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE)
-        vec = sum([[[min(i, NB_LABEL-1)]] * nb_plots for i in range(n_classes)], [])
-        vec = torch.tensor(vec, device=DEVICE, dtype=torch.int64)
-        vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
-
-        for ti in track(range(DIFFU_STEPS, 0, -1), description='Benchmarking...'):
-            t = torch.tensor([[ti]] * n_classes * nb_plots, device=DEVICE, dtype=torch.float32)
-            x = p_xt_1_xt_sig(model, x, t, vec)
-
-        x = x * 0.5 + 0.5
-        x = x.clamp(0, 1)
 
-        plt.figure(figsize=(nb_plots, n_classes))
-        for i in range(nb_plots):
-            for j in range(n_classes):
-                id = i * n_classes + j
-                img = x[id]
-                if autoencoder is not None:
-                    img = autoencoder.decode(img.unsqueeze(0)).squeeze(0)
-                plt.subplot(n_classes, nb_plots, id + 1)
-                plt.imshow(tensor_to_image(img))
-                plt.axis("off")
-        plt.tight_layout()
-        plt.savefig("plots/benchmark.tmp.png")
+if __name__ == "__main__":
+    main()
diff --git a/models/__init__.py b/models/__init__.py
new file mode 100644
index 0000000..db05aa9
--- /dev/null
+++ b/models/__init__.py
@@ -0,0 +1,6 @@
+"""
+Model definitions for MNIST Diffusion.
+"""
+
+from .unet import UNetMNIST
+from .autoencoder import AEModule, Autoencoder
diff --git a/models/autoencoder.py b/models/autoencoder.py
new file mode 100644
index 0000000..353ec89
--- /dev/null
+++ b/models/autoencoder.py
@@ -0,0 +1,103 @@
+"""
+Autoencoder model for MNIST.
+Can be used for latent diffusion.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class AEModule(nn.Module):
+    """Base class for autoencoder modules (encoder and decoder)."""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Subclasses must implement forward method.")
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode input to latent space."""
+        raise NotImplementedError("Subclasses must implement encode method.")
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode from latent space to image space."""
+        raise NotImplementedError("Subclasses must implement decode method.")
+
+
+class Autoencoder(AEModule):
+    """
+    Convolutional Autoencoder for MNIST images.
+
+    Args:
+        input_channels: Number of input image channels (1 for MNIST)
+        latent_channels: Number of channels in latent space
+    """
+
+    def __init__(self, input_channels: int = 1, latent_channels: int = 1):
+        super().__init__()
+        self.input_channels = input_channels
+        self.latent_channels = latent_channels
+
+        # Encoder: 28x28 -> 8x8
+        self.encoder = nn.Sequential(
+            nn.Conv2d(input_channels, 16, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(16, 16, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2),  # 14x14
+
+            nn.Conv2d(16, 32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2),  # 7x7
+
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, latent_channels, kernel_size=3, padding=1),
+            nn.ReLU(),
+
+            # 7x7 -> 8x8 (no activation - latent space should be unconstrained)
+            nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=1),
+        )
+
+        # Decoder: 8x8 -> 28x28
+        self.decoder = nn.Sequential(
+            # 8x8 -> 7x7
+            nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=0),
+            nn.ReLU(),
+
+            nn.Conv2d(latent_channels, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),  # 14x14
+            nn.ReLU(),
+
+            nn.Conv2d(32, 32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2),  # 28x28
+            nn.ReLU(),
+
+            nn.Conv2d(16, 16, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(16, input_channels, kernel_size=3, padding=1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Full autoencoder forward pass."""
+        z = self.encoder(x)
+        x_recon = self.decoder(z)
+        return x_recon
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode input to latent space."""
+        return self.encoder(x)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode from latent space to image space."""
+        return self.decoder(z)
diff --git a/models/unet.py b/models/unet.py
new file mode 100644
index 0000000..f21859a
--- /dev/null
+++ b/models/unet.py
@@ -0,0 +1,140 @@
+"""
+UNet model for MNIST diffusion.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sys
+sys.path.append("..")
+from src.config import IMG_SIZE, NB_CHANNEL, NB_LABEL, DIFFU_STEPS
+
+
+class SelfAttention(nn.Module):
+    """Self-attention module for UNet."""
+
+    def __init__(self, nb_channels: int, nb_heads: int):
+        super().__init__()
+        self.attention = nn.MultiheadAttention(nb_channels, nb_heads)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, c, w, h = x.shape
+        x = x.view(-1, c, w * h)
+        x = x.permute(2, 0, 1)
+        x, _ = self.attention(x, x, x)
+        x = x.permute(1, 2, 0)
+        x = x.view(-1, c, w, h)
+        return x
+
+
+class UNetMNIST(nn.Module):
+    """
+    UNet architecture for MNIST diffusion model.
+
+    Inputs:
+        xt: image at step t (NB_CHANNEL x IMG_SIZE x IMG_SIZE)
+        t: step number (1)
+        vec: one-hot vector of the label (NB_LABEL)
+
+    Output:
+        Predicted noise or denoised image (NB_CHANNEL x IMG_SIZE x IMG_SIZE)
+    """
+
+    def __init__(self, use_attention: bool = False):
+        super().__init__()
+        self.use_attention = use_attention
+
+        # Encoder for timestep t
+        self.encodet = nn.Linear(1, IMG_SIZE * IMG_SIZE)
+
+        # Encoder for label vector
+        self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE * IMG_SIZE)
+
+        # UNet encoder (2 extra channels for t and vec)
+        self.conv1 = nn.Conv2d(NB_CHANNEL + 2, 64, 3, padding=1)
+        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
+        self.maxpool1 = nn.MaxPool2d(2, 2)
+
+        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
+        self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
+        self.maxpool2 = nn.MaxPool2d(2, 2)
+
+        # Bottleneck
+        self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
+        self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
+
+        # Optional attention layers
+        if use_attention:
+            self.att1 = SelfAttention(128, 8)
+            self.att2 = SelfAttention(256, 8)
+            self.att3 = SelfAttention(256, 8)
+
+        # UNet decoder
+        self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.conv7 = nn.Conv2d(256, 128, 3, padding=1)
+        self.conv8 = nn.Conv2d(128, 128, 3, padding=1)
+
+        self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.conv9 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv10 = nn.Conv2d(64, 64, 3, padding=1)
+
+        # Output layer
+        self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1)
+
+    def forward(self, xt: torch.Tensor, t: torch.Tensor, vec: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of UNet.
+
+        Args:
+            xt: Noisy image at timestep t [B, C, H, W]
+            t: Timestep [B, 1]
+            vec: Label one-hot vector [B, NB_LABEL]
+
+        Returns:
+            Predicted noise or denoised image [B, C, H, W]
+        """
+        # Encode timestep and label
+        t_enc = F.relu(self.encodet(t / DIFFU_STEPS))
+        t_enc = t_enc.view(-1, 1, IMG_SIZE, IMG_SIZE)
+
+        vec_enc = F.relu(self.encodevec(vec))
+        vec_enc = vec_enc.view(-1, 1, IMG_SIZE, IMG_SIZE)
+
+        # Concatenate input with embeddings
+        x = torch.cat((xt, t_enc, vec_enc), dim=1)
+
+        # Encoder path
+        x1 = F.relu(self.conv1(x))
+        x1 = F.relu(self.conv2(x1))
+
+        x2 = self.maxpool1(x1)
+        x2 = F.relu(self.conv3(x2))
+        x2 = F.relu(self.conv4(x2))
+
+        # Bottleneck
+        x3 = self.maxpool2(x2)
+        if self.use_attention:
+            x3 = self.att1(x3)
+        x3 = F.relu(self.conv5(x3))
+        if self.use_attention:
+            x3 = self.att2(x3)
+        x3 = F.relu(self.conv6(x3))
+        if self.use_attention:
+            x3 = self.att3(x3)
+
+        # Decoder path with skip connections
+        x4 = self.upconv1(x3)
+        x4 = torch.cat((x4, x2), dim=1)
+        x4 = F.relu(self.conv7(x4))
+        x4 = F.relu(self.conv8(x4))
+
+        x5 = self.upconv2(x4)
+        x5 = torch.cat((x5, x1), dim=1)
+        x5 = F.relu(self.conv9(x5))
+        x5 = F.relu(self.conv10(x5))
+
+        # Output
+        out = self.conv11(x5)
+
+        return out
diff --git a/plots.py b/plots.py
deleted file mode 100644
index 0513495..0000000
--- a/plots.py
+++ /dev/null
@@ -1,174 +0,0 @@
-import matplotlib.pyplot as plt
-from matplotlib.gridspec import GridSpec
-import numpy as np
-from torchvision import datasets
-import torch
-import os
-from rich.progress import track
-
-from main import UNet, q_xt_xt_1, p_xt_1_xt, tensor_to_image
-from autoencoder import Autoencoder
-
-
-os.makedirs('plots', exist_ok=True)
-os.makedirs('plots/diffusion', exist_ok=True)
-os.makedirs('plots/diffusion_inverse', exist_ok=True)
-
-
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-DIFFU_STEPS = 1000
-BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE)
-BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA))
-ALPHA = 1 - BETA
-ALPHA_BAR = torch.cumprod(ALPHA, dim=0)
-
-NB_BINS = 50
-BIN_MIN = -4
-BIN_MAX = 4
-
-
-
-model = UNet().to(DEVICE)
-model.load_state_dict(torch.load('model.pth'))
-
-autoencoder = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8)).to(DEVICE)
-autoencoder.load_state_dict(torch.load('autoencoder.pth'))
-
-
-
-
-# plt.figure()
-# plt.plot(BETA, label='beta')
-# plt.plot(ALPHA, label='alpha')
-# plt.plot(ALPHA_BAR, label='alpha_bar')
-# plt.legend()
-# plt.title('Alpha, Beta and Alpha_bar schedules')
-# plt.savefig('plots/alpha_beta.tmp.png')
-
-
-mnist = datasets.MNIST('data', train=True, download=True)
-img, label = mnist[np.random.randint(0, len(mnist))]
-img = np.array(img) / 255 * 2 - 1
-img = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
-
-encoded = autoencoder.encode(img)
-
-img = encoded.squeeze().cpu().detach().numpy()
-
-plt.figure()
-plt.imshow(img, cmap='gray')
-plt.title('Image')
-plt.axis('off')
-plt.savefig('plots/img.tmp.png')
-
-
-plt.figure()
-plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX))
-plt.yscale('log')
-plt.title('Image histogram')
-plt.savefig('plots/img_hist.tmp.png')
-
-def norm_dist(x, mean, std):
-    return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi))
-
-x_norm = np.linspace(BIN_MIN, BIN_MAX, 100)
-y_norm = norm_dist(x_norm, 0, 1) * 8**2 / NB_BINS * (BIN_MAX - BIN_MIN)
-
-fig = plt.figure(figsize=(10, 5))
-fig.suptitle('Diffusion naturelle')
-gs = GridSpec(1, 3, figure=fig)
-ax1 = fig.add_subplot(gs[0, 0])
-ax2 = fig.add_subplot(gs[0, 1:])
-
-plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int)
-
-xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
-for t in track(range(1, DIFFU_STEPS+1)):
-    xt, eps = q_xt_xt_1(xt, t)
-
-    if t not in plots_to_save:
-        continue
-
-    xt_numpy = xt.cpu().detach().numpy()[0, 0]
-
-    ax1.clear()
-    ax1.imshow(xt_numpy, cmap='gray')
-    ax1.set_title(f'xt at t={t:04d}')
-    ax1.axis('off')
-
-    ax2.clear()
-    ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX))
-    ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)')
-    ax2.legend()
-    ax2.set_yscale('log')
-    ax2.set_ylim(y_norm.min(), 1e3)
-    ax2.set_title(f'xt histogram')
-
-    fig.savefig(f'plots/diffusion/{t:04d}.tmp.png')
-os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif')
-
-fig = plt.figure(figsize=(10, 5))
-fig.suptitle('Diffusion inverse')
-gs = GridSpec(1, 3, figure=fig)
-ax1 = fig.add_subplot(gs[0, 0])
-ax2 = fig.add_subplot(gs[0, 1:])
-
-xt = torch.randn(1, 1, 28, 28, device=DEVICE)
-vec = torch.zeros(1, 10).to(DEVICE)
-vec[0, label] = 1
-for t in track(range(DIFFU_STEPS, 0, -1)):
-    t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32)
-    xt = p_xt_1_xt(model, xt, t_tensor, vec)
-
-    if t not in plots_to_save:
-        continue
-
-    xt_numpy = xt.cpu().detach().numpy()[0, 0]
-
-    ax1.clear()
-    ax1.imshow(xt_numpy, cmap='gray')
-    ax1.set_title(f'xt at t={t:04d}')
-    ax1.axis('off')
-
-    ax2.clear()
-    ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX))
-    ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)')
-    ax2.legend()
-    ax2.set_yscale('log')
-    ax2.set_ylim(y_norm.min(), 1e3)
-    ax2.set_title(f'xt histogram')
-
-    fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png')
-os.system('convert -delay 20 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif')
-
-exit(0)
-
-tpause = {150: 'xt', 20: 'mu', 50: 'xt_1'}
-
-xT = xt = torch.randn(1, 1, 28, 28, device=DEVICE)
-vec = torch.zeros(1, 10).to(DEVICE)
-vec[0, label] = 1
-
-for t in range(DIFFU_STEPS, 0, -1):
-    t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32)
-    xt = p_xt_1_xt(model, xt, t_tensor, vec)
-
-    if t in tpause:
-        plt.figure(figsize=(5, 5))
-        plt.imshow(tensor_to_image(xt[0]), cmap='gray')
-        plt.axis('off')
-        plt.tight_layout()
-        plt.savefig(f'plots/{tpause[t]}.tmp.png')
-
-plt.figure(figsize=(5, 5))
-plt.imshow(tensor_to_image(xT[0]), cmap='gray')
-plt.axis('off')
-plt.tight_layout()
-plt.savefig('plots/xT.tmp.png')
-
-plt.figure(figsize=(5, 5))
-plt.imshow(tensor_to_image(xt[0]), cmap='gray')
-plt.axis('off')
-plt.tight_layout()
-plt.savefig('plots/x0.tmp.png')
diff --git a/requirements.txt b/requirements.txt
index d590e96..34490ec 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
-torchvision==0.15.2
-matplotlib==3.8.0
-rich==13.3.5
-PyQt5==5.15.10
-torch-trainer @ git+https://github.com/gdamms/torch-trainer.git@da70ca78adf6195f2e3add92938b9c282d5c94cb
\ No newline at end of file
+torch
+torchvision
+mlflow
+scipy
+matplotlib
+rich
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..220bb0e
--- /dev/null
+++ b/src/__init__.py
@@ -0,0 +1,50 @@
+"""
+MNIST Diffusion Model package.
+"""
+
+from .config import (
+    DEVICE,
+    IMG_SIZE,
+    NB_CHANNEL,
+    NB_LABEL,
+    DIFFU_STEPS,
+    ALPHA,
+    ALPHA_BAR,
+    BETA,
+    EPOCHS,
+    BATCH_SIZE,
+    LEARNING_RATE,
+    CHECKPOINT_DIR,
+    PLOTS_DIR,
+)
+
+from .diffusion import (
+    q_xt_xt_1,
+    q_xt_x0,
+    p_xt_1_xt,
+    p_xt_1_xt_x0_pred,
+    forward_diffusion,
+)
+
+from .dataloader import (
+    get_mnist_dataset,
+    get_diffusion_dataloader,
+    get_autoencoder_dataloader,
+    DiffusionDataset,
+    DiffusionDatasetX0,
+    AutoencoderDataset,
+)
+
+from .utils import (
+    tensor_to_image,
+    tensor_to_images,
+    save_checkpoint,
+    load_checkpoint,
+    save_plot,
+)
+
+from .metrics import (
+    fid,
+    kl_divergence,
+    jsd,
+)
diff --git a/src/config.py b/src/config.py
new file mode 100644
index 0000000..4ac56c4
--- /dev/null
+++ b/src/config.py
@@ -0,0 +1,34 @@
+"""
+Configuration file for MNIST Diffusion model.
+Contains all hyperparameters and constants.
+"""
+
+import torch
+
+# Device configuration
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Image parameters (MNIST)
+IMG_SIZE = 28
+NB_CHANNEL = 1
+NB_LABEL = 10
+
+# Diffusion parameters
+DIFFU_STEPS = 1000
+
+# Noise schedule (linear beta schedule)
+BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE)
+BETA = torch.cat((torch.tensor([0.0], device=DEVICE), BETA))
+ALPHA = 1 - BETA
+ALPHA_BAR = torch.cumprod(ALPHA, dim=0)
+
+# Training parameters
+EPOCHS = 10
+BATCH_SIZE = 64
+LEARNING_RATE = 2e-4
+NUM_WORKERS = 4
+
+# Paths
+DATA_DIR = "data"
+CHECKPOINT_DIR = "checkpoints"
+PLOTS_DIR = "plots"
diff --git a/src/dataloader.py b/src/dataloader.py
new file mode 100644
index 0000000..a08ae27
--- /dev/null
+++ b/src/dataloader.py
@@ -0,0 +1,210 @@
+"""
+Data loading utilities for MNIST diffusion training.
+"""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import datasets, transforms
+
+from .config import DEVICE, DIFFU_STEPS, NB_LABEL, DATA_DIR
+from .diffusion import q_xt_x0
+
+
+def get_mnist_dataset(train: bool = True) -> datasets.MNIST:
+    """
+    Load MNIST dataset.
+
+    Args:
+        train: If True, load training set. Otherwise load test set.
+
+    Returns:
+        MNIST dataset
+    """
+    return datasets.MNIST(
+        root=DATA_DIR,
+        train=train,
+        download=True,
+        transform=transforms.ToTensor(),
+    )
+
+
+class DiffusionDataset(Dataset):
+    """
+    Dataset wrapper for diffusion training.
+    Returns noisy image, timestep, label, and target noise.
+
+    Args:
+        dataset: Base image dataset (e.g., MNIST)
+        autoencoder: Optional autoencoder for latent diffusion
+    """
+
+    def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None):
+        super().__init__()
+        self.dataset = dataset
+        self.autoencoder = autoencoder
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
+        # Get image and label
+        img, label = self.dataset[index]
+        img = img.to(DEVICE)
+
+        # Optionally encode to latent space
+        if self.autoencoder is not None:
+            with torch.no_grad():
+                img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0)
+
+        # Normalize to [-1, 1]
+        img = img * 2 - 1
+
+        # Sample random timestep and add noise
+        t = torch.randint(1, DIFFU_STEPS + 1, (1,), device=DEVICE)
+        xt, eps = q_xt_x0(img, t)
+
+        # Convert label to one-hot vector
+        vec = torch.nn.functional.one_hot(
+            torch.tensor(min(label, NB_LABEL - 1)),
+            num_classes=NB_LABEL,
+        )
+
+        return (
+            xt.clone().detach().to(dtype=torch.float32, device=DEVICE),
+            t.clone().detach().to(dtype=torch.float32, device=DEVICE),
+            vec.clone().detach().to(dtype=torch.float32, device=DEVICE),
+            eps,  # Target: the noise that was added
+        )
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+
+class DiffusionDatasetX0(Dataset):
+    """
+    Dataset wrapper for diffusion training where model predicts x0 instead of noise.
+    Returns noisy image, timestep, label, and target clean image.
+
+    Args:
+        dataset: Base image dataset (e.g., MNIST)
+        autoencoder: Optional autoencoder for latent diffusion
+    """
+
+    def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None):
+        super().__init__()
+        self.dataset = dataset
+        self.autoencoder = autoencoder
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
+        # Get image and label
+        img, label = self.dataset[index]
+        img = img.to(DEVICE)
+
+        # Optionally encode to latent space
+        if self.autoencoder is not None:
+            with torch.no_grad():
+                img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0)
+
+        # Normalize to [-1, 1]
+        img = img * 2 - 1
+
+        # Sample random timestep and add noise
+        t = torch.randint(1, DIFFU_STEPS + 1, (1,), device=DEVICE)
+        xt, _ = q_xt_x0(img, t)
+
+        # Convert label to one-hot vector
+        vec = torch.nn.functional.one_hot(
+            torch.tensor(min(label, NB_LABEL - 1)),
+            num_classes=NB_LABEL,
+        )
+
+        return (
+            xt.clone().detach().to(dtype=torch.float32, device=DEVICE),
+            t.clone().detach().to(dtype=torch.float32, device=DEVICE),
+            vec.clone().detach().to(dtype=torch.float32, device=DEVICE),
+            img.clone().detach().to(dtype=torch.float32, device=DEVICE),  # Target: clean image
+        )
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+
+class AutoencoderDataset(Dataset):
+    """
+    Dataset wrapper for autoencoder training.
+    Returns image as both input and target.
+
+    Args:
+        dataset: Base image dataset
+    """
+
+    def __init__(self, dataset: Dataset):
+        self.dataset = dataset
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        data = self.dataset[idx][0].to(DEVICE)
+        return data, data
+
+
+def get_diffusion_dataloader(
+    predict_x0: bool = True,
+    batch_size: int = 64,
+    shuffle: bool = True,
+    num_workers: int = 4,
+    autoencoder: torch.nn.Module = None,
+) -> DataLoader:
+    """
+    Create a DataLoader for diffusion training.
+
+    Args:
+        predict_x0: If True, model predicts x0. Otherwise predicts noise.
+        batch_size: Batch size
+        shuffle: Whether to shuffle data
+        num_workers: Number of data loading workers
+        autoencoder: Optional autoencoder for latent diffusion
+
+    Returns:
+        DataLoader for training
+    """
+    mnist = get_mnist_dataset(train=True)
+
+    if predict_x0:
+        dataset = DiffusionDatasetX0(mnist, autoencoder)
+    else:
+        dataset = DiffusionDataset(mnist, autoencoder)
+
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        persistent_workers=True if num_workers > 0 else False,
+    )
+
+
+def get_autoencoder_dataloader(
+    batch_size: int = 64,
+    shuffle: bool = True,
+    num_workers: int = 4,
+) -> DataLoader:
+    """
+    Create a DataLoader for autoencoder training.
+
+    Args:
+        batch_size: Batch size
+        shuffle: Whether to shuffle data
+        num_workers: Number of data loading workers
+
+    Returns:
+        DataLoader for training
+    """
+    mnist = get_mnist_dataset(train=True)
+    dataset = AutoencoderDataset(mnist)
+
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        persistent_workers=True if num_workers > 0 else False,
+    )
diff --git a/src/diffusion.py b/src/diffusion.py
new file mode 100644
index 0000000..0d63ee1
--- /dev/null
+++ b/src/diffusion.py
@@ -0,0 +1,154 @@
+"""
+Diffusion process utilities.
+Contains forward and reverse diffusion functions.
+"""
+
+import torch
+from .config import DEVICE, ALPHA, ALPHA_BAR, BETA, DIFFU_STEPS
+
+
+def q_xt_xt_1(xt_1: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Forward diffusion step: q(x_t | x_{t-1}).
+    Adds noise to image at step t-1 to get image at step t.
+
+    Args:
+        xt_1: Image at timestep t-1 [B, C, H, W]
+        t: Timestep (int or tensor)
+
+    Returns:
+        xt: Noisy image at timestep t
+        eps: The noise that was added
+    """
+    if isinstance(t, int):
+        t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE)
+    else:
+        t_ind = t.to(dtype=torch.long, device=DEVICE)
+
+    alpha = ALPHA[t_ind]
+    mean = torch.sqrt(alpha) * xt_1
+    std = torch.sqrt(1 - alpha)
+
+    eps = torch.randn(xt_1.shape, device=DEVICE)
+    xt = mean + std * eps
+
+    return xt, eps
+
+
+def q_xt_x0(x0: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Forward diffusion: q(x_t | x_0).
+    Directly compute noisy image at any timestep t from clean image x0.
+
+    Args:
+        x0: Clean image [B, C, H, W] or [C, H, W]
+        t: Timestep tensor [B, 1] or [B]
+
+    Returns:
+        xt: Noisy image at timestep t
+        eps: The noise that was added
+    """
+    if isinstance(t, int):
+        t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE)
+    else:
+        t_ind = t.to(dtype=torch.long, device=DEVICE)
+
+    # Handle both batched and single images
+    reshaped = len(x0.shape) == 3
+    if reshaped:
+        c, w, h = x0.shape
+        b = 1
+        x0 = x0.view(b, c, w, h)
+    else:
+        b, c, w, h = x0.shape
+
+    # Reshape t for broadcasting
+    t_ind = t_ind.view(b, 1, 1, 1)
+    t_ind = t_ind.expand(b, c, w, h)
+
+    alpha_bar = ALPHA_BAR[t_ind]
+    mean = torch.sqrt(alpha_bar) * x0
+    std = torch.sqrt(1 - alpha_bar)
+
+    eps = torch.randn(x0.shape, device=DEVICE)
+    xt = mean + std * eps
+
+    if reshaped:
+        xt = xt.view(c, w, h)
+
+    return xt, eps
+
+
+def p_xt_1_xt(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor,
+              vec: torch.Tensor) -> torch.Tensor:
+    """
+    Reverse diffusion step: p(x_{t-1} | x_t).
+    Denoise image at step t to get image at step t-1.
+    Model predicts the noise.
+
+    Args:
+        model: UNet model that predicts noise
+        xt: Noisy image at timestep t [B, C, H, W]
+        t: Timestep tensor [B, 1]
+        vec: Label one-hot vector [B, NB_LABEL]
+
+    Returns:
+        xt_1: Denoised image at timestep t-1
+    """
+    t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
+
+    alpha_bar_t = ALPHA_BAR[t_ind].view(-1, 1, 1, 1)
+    alpha_bar_t_1 = ALPHA_BAR[t_ind - 1].view(-1, 1, 1, 1)
+    alpha_t = ALPHA[t_ind].view(-1, 1, 1, 1)
+    beta_t = BETA[t_ind].view(-1, 1, 1, 1)
+
+    beta_tilde = (1 - alpha_bar_t_1) / (1 - alpha_bar_t) * beta_t
+
+    # Model predicts the noise
+    epsilon_theta = model(xt, t, vec)
+
+    sigma_theta = torch.sqrt(beta_tilde)
+    mu_theta = (xt - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * epsilon_theta) / torch.sqrt(alpha_t)
+
+    # Don't add noise at t=1
+    mask_t0 = (t > 1).to(dtype=torch.float32).view(-1, 1, 1, 1)
+    noise = torch.randn(xt.shape, device=DEVICE) * mask_t0
+
+    return mu_theta + sigma_theta * noise
+
+
+def p_xt_1_xt_x0_pred(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor,
+                      vec: torch.Tensor) -> torch.Tensor:
+    """
+    Reverse diffusion step where model predicts x0 directly.
+
+    Args:
+        model: UNet model that predicts clean image x0
+        xt: Noisy image at timestep t [B, C, H, W]
+        t: Timestep tensor [B, 1]
+        vec: Label one-hot vector [B, NB_LABEL]
+
+    Returns:
+        xt_1: Denoised image at timestep t-1
+    """
+    x0_pred = model(xt, t, vec)
+    xt_1, _ = q_xt_x0(x0_pred, t - 1)
+    return xt_1
+
+
+def forward_diffusion(x0: torch.Tensor) -> list[torch.Tensor]:
+    """
+    Run full forward diffusion process.
+
+    Args:
+        x0: Clean image [C, H, W] or [B, C, H, W]
+
+    Returns:
+        List of images at each timestep [x0, x1, ..., xT]
+    """
+    x = x0.clone()
+    xs = [x]
+    for t in range(1, DIFFU_STEPS + 1):
+        x, _ = q_xt_xt_1(x, t)
+        xs.append(x)
+    return xs
diff --git a/src/metrics.py b/src/metrics.py
new file mode 100644
index 0000000..08eb3cd
--- /dev/null
+++ b/src/metrics.py
@@ -0,0 +1,98 @@
+import numpy as np
+import scipy
+
+
+def fid(reals: np.ndarray, fakes: np.ndarray) -> float:
+    """
+    Calculate Frechet Inception Distance (FID) score.
+
+    Args:
+        reals: Real images [N, C, H, W]
+        fakes: Generated images [N, C, H, W]
+
+    Returns:
+        FID score (lower is better)
+    """
+    reals = reals.reshape(reals.shape[0], -1)
+    fakes = fakes.reshape(fakes.shape[0], -1)
+
+    mu_real = np.mean(reals, axis=0)
+    mu_fake = np.mean(fakes, axis=0)
+    sigma_real = np.cov(reals, rowvar=False)
+    sigma_fake = np.cov(fakes, rowvar=False)
+
+    diff = mu_real - mu_fake
+    covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False)
+
+    if not np.isfinite(covmean).all():
+        eps = 1e-6
+        offset = np.eye(sigma_real.shape[0]) * eps
+        covmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset))
+
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
+
+    return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean)
+
+
+def kl_divergence(reals: np.ndarray, fakes: np.ndarray) -> float:
+    """
+    Calculate KL divergence between real and fake image distributions.
+
+    Args:
+        reals: Real images [N, C, H, W]
+        fakes: Generated images [N, C, H, W]
+
+    Returns:
+        KL divergence value
+    """
+    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
+    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)
+
+    hist_real = np.apply_along_axis(
+        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals
+    )
+    hist_fake = np.apply_along_axis(
+        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes
+    )
+
+    # Add smoothing
+    hist_real = hist_real + 1
+    hist_fake = hist_fake + 1
+
+    hist_real = hist_real / np.sum(hist_real)
+    hist_fake = hist_fake / np.sum(hist_fake)
+
+    return np.mean(np.log(hist_real / hist_fake))
+
+
+def jsd(reals: np.ndarray, fakes: np.ndarray) -> float:
+    """
+    Calculate Jensen-Shannon divergence between real and fake image distributions.
+
+    Args:
+        reals: Real images [N, C, H, W]
+        fakes: Generated images [N, C, H, W]
+
+    Returns:
+        JSD value
+    """
+    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
+    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)
+
+    hist_real = np.apply_along_axis(
+        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals
+    )
+    hist_fake = np.apply_along_axis(
+        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes
+    )
+
+    hist_real = hist_real + 1
+    hist_fake = hist_fake + 1
+
+    hist_real = hist_real / np.sum(hist_real)
+    hist_fake = hist_fake / np.sum(hist_fake)
+
+    hist_avg = (hist_real + hist_fake) / 2
+
+    return 0.5 * (np.mean(np.log(hist_real / hist_avg)) + np.mean(np.log(hist_fake / hist_avg)))
diff --git a/src/sample.py b/src/sample.py
new file mode 100644
index 0000000..5b438ed
--- /dev/null
+++ b/src/sample.py
@@ -0,0 +1,248 @@
+"""
+Sampling and evaluation script for trained diffusion model.
+Generate samples from a trained model and visualize results.
+"""
+
+from models import UNetMNIST
+from src.utils import ensure_dirs, tensor_to_image, load_checkpoint
+from src.dataloader import get_mnist_dataset
+from src.diffusion import p_xt_1_xt_x0_pred, forward_diffusion, q_xt_x0
+from src.config import (
+    DEVICE, DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL,
+    CHECKPOINT_DIR, PLOTS_DIR
+)
+import os
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from rich.progress import track
+
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def generate_samples(
+    model: torch.nn.Module,
+    n_samples: int = 10,
+    labels: list[int] | None = None,
+) -> torch.Tensor:
+    """
+    Generate samples from the diffusion model.
+
+    Args:
+        model: Trained UNet model
+        n_samples: Number of samples per class (or total if labels provided)
+        labels: Optional list of specific labels to generate
+
+    Returns:
+        Generated samples tensor [N, C, H, W]
+    """
+    model.eval()
+
+    if labels is None:
+        # Generate samples for all classes
+        labels = list(range(NB_LABEL)) * n_samples
+
+    n_total = len(labels)
+
+    with torch.no_grad():
+        # Start from pure noise
+        x = torch.randn(n_total, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE)
+
+        # Create one-hot label vectors
+        vec = torch.tensor(labels, device=DEVICE)
+        vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32)
+
+        # Reverse diffusion process
+        for t in track(range(DIFFU_STEPS, 0, -1), description="Generating samples"):
+            t_tensor = torch.tensor([[t]] * n_total, device=DEVICE, dtype=torch.float32)
+            x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec)
+
+        # Normalize to [0, 1]
+        x = x * 0.5 + 0.5
+        x = x.clamp(0, 1)
+
+    return x
+
+
+def visualize_forward_diffusion(save_path: str | None = None):
+    """
+    Visualize the forward diffusion process on a real image.
+
+    Args:
+        save_path: Path to save the visualization
+    """
+    ensure_dirs()
+
+    # Get a random image from MNIST
+    dataset = get_mnist_dataset(train=True)
+    idx = np.random.randint(0, len(dataset))
+    img, label = dataset[idx]
+    img = img.to(DEVICE)
+
+    # Normalize to [-1, 1]
+    img = img * 2 - 1
+
+    # Run forward diffusion
+    xs = forward_diffusion(img)
+
+    # Select timesteps to visualize
+    n_plots = 10
+    timesteps = np.linspace(1, DIFFU_STEPS, n_plots, dtype=int)
+
+    fig, axes = plt.subplots(2, n_plots + 1, figsize=(2 * n_plots, 5))
+
+    # Row labels
+    axes[0, 0].text(0.5, 0.5, 'Step-by-step', ha='center', va='center', fontsize=10)
+    axes[0, 0].axis('off')
+    axes[1, 0].text(0.5, 0.5, 'Direct', ha='center', va='center', fontsize=10)
+    axes[1, 0].axis('off')
+
+    # Plot step-by-step diffusion
+    for i, t in enumerate(timesteps):
+        axes[0, i + 1].imshow(tensor_to_image(xs[t]), cmap='gray')
+        axes[0, i + 1].set_title(f't={t}')
+        axes[0, i + 1].axis('off')
+
+        # Plot direct diffusion for comparison
+        xt, _ = q_xt_x0(img, t)
+        axes[1, i + 1].imshow(tensor_to_image(xt), cmap='gray')
+        axes[1, i + 1].axis('off')
+
+    fig.suptitle(f'Forward Diffusion Process (Label: {label})')
+    plt.tight_layout()
+
+    if save_path is None:
+        save_path = os.path.join(PLOTS_DIR, 'forward_diffusion.png')
+    fig.savefig(save_path)
+    print(f"Saved forward diffusion visualization to {save_path}")
+    plt.close(fig)
+
+
+def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = None):
+    """
+    Visualize the backward (reverse) diffusion process.
+
+    Args:
+        model: Trained UNet model
+        save_path: Path to save the visualization
+    """
+    ensure_dirs()
+    model.eval()
+
+    n_classes = NB_LABEL
+    n_timesteps = 10
+    timesteps = np.linspace(1, DIFFU_STEPS, n_timesteps, dtype=int)[::-1]
+
+    fig, axes = plt.subplots(n_classes, n_timesteps, figsize=(2 * n_timesteps, 2 * n_classes))
+
+    with torch.no_grad():
+        # Start from noise
+        x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE)
+
+        # One sample per class
+        vec = torch.arange(n_classes, device=DEVICE)
+        vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32)
+
+        for t in track(range(DIFFU_STEPS, 0, -1), description="Visualizing backward diffusion"):
+            t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32)
+            x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec)
+
+            if t in timesteps:
+                t_idx = timesteps.tolist().index(t)
+                for class_idx in range(n_classes):
+                    axes[class_idx, t_idx].imshow(tensor_to_image(x[class_idx]), cmap='gray')
+                    if class_idx == 0:
+                        axes[class_idx, t_idx].set_title(f't={t}')
+                    if t_idx == 0:
+                        axes[class_idx, t_idx].set_ylabel(f'Class {class_idx}')
+                    axes[class_idx, t_idx].set_xticks([])
+                    axes[class_idx, t_idx].set_yticks([])
+
+    fig.suptitle('Backward Diffusion Process')
+    plt.tight_layout()
+
+    if save_path is None:
+        save_path = os.path.join(PLOTS_DIR, 'backward_diffusion.png')
+    fig.savefig(save_path)
+    print(f"Saved backward diffusion visualization to {save_path}")
+    plt.close(fig)
+
+
+def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str | None = None):
+    """
+    Generate a grid of samples, organized by class.
+
+    Args:
+        model: Trained UNet model
+        n_per_class: Number of samples per class
+        save_path: Path to save the grid
+    """
+    ensure_dirs()
+
+    # Generate samples
+    labels = []
+    for class_idx in range(NB_LABEL):
+        labels.extend([class_idx] * n_per_class)
+
+    samples = generate_samples(model, labels=labels)
+    samples = samples.cpu().numpy()
+
+    # Create grid
+    fig, axes = plt.subplots(NB_LABEL, n_per_class, figsize=(n_per_class, NB_LABEL))
+
+    for class_idx in range(NB_LABEL):
+        for sample_idx in range(n_per_class):
+            idx = class_idx * n_per_class + sample_idx
+            axes[class_idx, sample_idx].imshow(samples[idx].transpose(1, 2, 0).squeeze(), cmap='gray')
+            axes[class_idx, sample_idx].axis('off')
+
+            if sample_idx == 0:
+                axes[class_idx, sample_idx].set_ylabel(f'{class_idx}')
+
+    fig.suptitle('Generated MNIST Digits')
+    plt.tight_layout()
+
+    if save_path is None:
+        save_path = os.path.join(PLOTS_DIR, 'generated_grid.png')
+    fig.savefig(save_path)
+    print(f"Saved generated grid to {save_path}")
+    plt.close(fig)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Sample from trained diffusion model")
+    parser.add_argument("--checkpoint", type=str, default=os.path.join(CHECKPOINT_DIR, "diffusion_latest.pt"),
+                        help="Path to model checkpoint")
+    parser.add_argument("--n-samples", type=int, default=10, help="Samples per class")
+    parser.add_argument("--attention", action="store_true", help="Use attention in model")
+    parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion")
+    parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion")
+    parser.add_argument("--grid", action="store_true", help="Generate sample grid")
+    parser.add_argument("--all", action="store_true", help="Run all visualizations")
+
+    args = parser.parse_args()
+
+    # Load model
+    model = UNetMNIST(use_attention=args.attention).to(DEVICE)
+
+    if os.path.exists(args.checkpoint):
+        model = load_checkpoint(model, os.path.basename(args.checkpoint))
+    else:
+        print(f"Warning: Checkpoint {args.checkpoint} not found. Using untrained model.")
+
+    # Run visualizations
+    if args.forward or args.all:
+        visualize_forward_diffusion()
+
+    if args.backward or args.all:
+        visualize_backward_diffusion(model)
+
+    if args.grid or args.all:
+        generate_grid(model, n_per_class=args.n_samples)
+
+    # Default: generate grid if no specific option selected
+    if not (args.forward or args.backward or args.grid or args.all):
+        generate_grid(model, n_per_class=args.n_samples)
diff --git a/src/train_autoencoder.py b/src/train_autoencoder.py
new file mode 100644
index 0000000..669ce47
--- /dev/null
+++ b/src/train_autoencoder.py
@@ -0,0 +1,167 @@
+"""
+Training script for MNIST autoencoder.
+"""
+
+from models import Autoencoder, AEModule
+from src.utils import ensure_dirs, save_checkpoint
+from src.dataloader import get_autoencoder_dataloader
+from src.config import DEVICE, BATCH_SIZE, NUM_WORKERS, CHECKPOINT_DIR, PLOTS_DIR
+import os
+import torch
+import torch.nn as nn
+import matplotlib.pyplot as plt
+from rich.progress import track
+import mlflow
+
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def train_autoencoder(
+    epochs: int = 10,
+    learning_rate: float = 1e-3,
+    batch_size: int = BATCH_SIZE,
+    latent_channels: int = 1,
+    checkpoint_path: str | None = None,
+    run_name: str | None = None,
+):
+    """
+    Train the autoencoder model.
+
+    Args:
+        epochs: Number of training epochs
+        learning_rate: Learning rate for optimizer
+        batch_size: Training batch size
+        latent_channels: Number of channels in latent space
+        checkpoint_path: Path to checkpoint to resume training from
+        run_name: Name for this training run (for logging)
+    """
+    ensure_dirs()
+
+    # Initialize model
+    model = Autoencoder(input_channels=1, latent_channels=latent_channels).to(DEVICE)
+
+    if checkpoint_path:
+        if not os.path.exists(checkpoint_path):
+            print(f"Checkpoint not found: {checkpoint_path}")
+            return
+        model.load_state_dict(torch.load(checkpoint_path, weights_only=True))
+
+    if run_name is None:
+        from datetime import datetime
+        run_name = f"autoencoder_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+
+    mlflow.set_experiment("MNIST Autoencoder")
+    mlflow.start_run(run_name=run_name)
+
+    # Setup training
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    # Use MSE loss instead of BCE for better reconstruction of continuous values
+    criterion = nn.MSELoss()
+
+    # Get dataloader
+    dataloader = get_autoencoder_dataloader(
+        batch_size=batch_size,
+        num_workers=NUM_WORKERS,
+    )
+
+    # Training loop
+    for epoch in range(1, epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+
+        for batch_idx, (x, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")):
+            optimizer.zero_grad()
+
+            # Forward pass
+            x_recon = model(x)
+            loss = criterion(x_recon, target)
+
+            # Backward pass
+            loss.backward()
+
+            # Gradient clipping to prevent exploding gradients
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+            optimizer.step()
+
+            epoch_loss += loss.item()
+
+        avg_loss = epoch_loss / len(dataloader)
+        mlflow.log_metric("epoch_loss", avg_loss, step=epoch)
+
+        # Save checkpoint
+        save_checkpoint(model, f"autoencoder_epoch_{epoch:03d}.pt")
+        save_checkpoint(model, "autoencoder_latest.pt")
+
+    # Visualize results
+    visualize_reconstructions(model, dataloader)
+
+    mlflow.end_run()
+    return model
+
+
+def visualize_reconstructions(model: AEModule, dataloader, n_samples: int = 10):
+    """Visualize original, latent, and reconstructed images."""
+    model.eval()
+    ensure_dirs()
+
+    # Get a batch of samples
+    x_batch, _ = next(iter(dataloader))
+    x_batch = x_batch[:n_samples]
+
+    with torch.no_grad():
+        latent = model.encode(x_batch)
+        x_recon = model.decode(latent)
+
+    # Create visualization
+    fig, axes = plt.subplots(3, n_samples + 1, figsize=(2 * n_samples, 6))
+
+    # Labels
+    axes[0, 0].text(0.5, 0.5, 'Original', ha='center', va='center', fontsize=12)
+    axes[0, 0].axis('off')
+    axes[1, 0].text(0.5, 0.5, 'Latent', ha='center', va='center', fontsize=12)
+    axes[1, 0].axis('off')
+    axes[2, 0].text(0.5, 0.5, 'Reconstructed', ha='center', va='center', fontsize=12)
+    axes[2, 0].axis('off')
+
+    # Plot images
+    for i in range(n_samples):
+        axes[0, i + 1].imshow(x_batch[i].cpu().squeeze().numpy(), cmap='gray')
+        axes[0, i + 1].axis('off')
+
+        axes[1, i + 1].imshow(latent[i].cpu().squeeze().numpy(), cmap='gray')
+        axes[1, i + 1].axis('off')
+
+        axes[2, i + 1].imshow(x_recon[i].cpu().squeeze().numpy(), cmap='gray')
+        axes[2, i + 1].axis('off')
+
+    fig.suptitle('Autoencoder Results')
+    plt.tight_layout()
+
+    save_path = os.path.join(PLOTS_DIR, 'autoencoder_results.png')
+    fig.savefig(save_path)
+    plt.close(fig)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Train MNIST autoencoder")
+    parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
+    parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Batch size")
+    parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint")
+
+    args = parser.parse_args()
+
+    torch.multiprocessing.set_start_method("spawn", force=True)
+
+    train_autoencoder(
+        epochs=args.epochs,
+        learning_rate=args.lr,
+        batch_size=args.batch_size,
+        latent_channels=args.latent_channels,
+        checkpoint_path=args.checkpoint,
+    )
diff --git a/src/train_diffusion.py b/src/train_diffusion.py
new file mode 100644
index 0000000..cb70650
--- /dev/null
+++ b/src/train_diffusion.py
@@ -0,0 +1,197 @@
+"""
+Training script for MNIST diffusion model.
+"""
+
+from models import UNetMNIST
+from src.utils import (
+    ensure_dirs, save_checkpoint, tensor_to_image,
+    figure_to_image,
+)
+from src.metrics import fid, kl_divergence, jsd
+from src.diffusion import p_xt_1_xt_x0_pred
+from src.dataloader import get_diffusion_dataloader, get_mnist_dataset
+from src.config import (
+    DEVICE, EPOCHS, LEARNING_RATE, BATCH_SIZE, NUM_WORKERS,
+    DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL, CHECKPOINT_DIR, PLOTS_DIR
+)
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import matplotlib.pyplot as plt
+from rich.progress import track
+import mlflow
+
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def train_diffusion(
+    epochs: int = EPOCHS,
+    learning_rate: float = LEARNING_RATE,
+    batch_size: int = BATCH_SIZE,
+    predict_x0: bool = True,
+    use_attention: bool = False,
+    checkpoint_path: str | None = None,
+    run_name: str | None = None,
+):
+    """
+    Train the diffusion model.
+
+    Args:
+        epochs: Number of training epochs
+        learning_rate: Learning rate for optimizer
+        batch_size: Training batch size
+        predict_x0: If True, model predicts x0. Otherwise predicts noise.
+        use_attention: If True, use self-attention in UNet
+        checkpoint_path: Path to checkpoint to resume training from
+        run_name: Name for this training run (for logging)
+    """
+    ensure_dirs()
+
+    # Setup run name and logging
+    if run_name is None:
+        from datetime import datetime
+        run_name = f"diffusion_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+
+    mlflow.set_experiment("MNIST Diffusion")
+    mlflow.start_run(run_name=run_name)
+
+    # Initialize model
+    model = UNetMNIST(use_attention=use_attention).to(DEVICE)
+
+    if checkpoint_path and os.path.exists(checkpoint_path):
+        model.load_state_dict(torch.load(checkpoint_path, weights_only=True))
+        print(f"Loaded checkpoint from {checkpoint_path}")
+
+    # Setup training
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    criterion = nn.MSELoss()
+
+    # Get dataloader
+    dataloader = get_diffusion_dataloader(
+        predict_x0=predict_x0,
+        batch_size=batch_size,
+        num_workers=NUM_WORKERS,
+    )
+
+    # Training loop
+    global_step = 0
+    for epoch in range(1, epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+
+        for batch_idx, (xt, t, vec, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")):
+            optimizer.zero_grad()
+
+            # Forward pass
+            pred = model(xt, t, vec)
+            loss = criterion(pred, target)
+
+            # Backward pass
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            global_step += 1
+
+            # Log training loss every 100 steps
+            if global_step % 100 == 0:
+                avg_loss = epoch_loss / (batch_idx + 1)
+                mlflow.log_metric("train_loss", avg_loss, step=global_step)
+
+        avg_loss = epoch_loss / len(dataloader)
+        mlflow.log_metric("epoch_loss", avg_loss, step=epoch)
+
+        # Save checkpoint every epoch
+        save_checkpoint(model, f"diffusion_epoch_{epoch:03d}.pt")
+        save_checkpoint(model, "diffusion_latest.pt")
+
+        # Generate and log sample images
+        if epoch % 1 == 0:
+            evaluate_and_log(model, epoch, predict_x0)
+
+    mlflow.end_run()
+    return model
+
+
+def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True):
+    """Generate samples and log metrics."""
+    model.eval()
+
+    with torch.no_grad():
+        # Generate samples
+        batch_size = 64
+        n_batches = 4
+
+        fakes = []
+        for _ in range(n_batches):
+            x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE)
+            vec = torch.randint(0, NB_LABEL, (batch_size,), device=DEVICE)
+            vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32)
+
+            # Reverse diffusion
+            for t in range(DIFFU_STEPS, 0, -1):
+                t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32)
+                x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec)
+
+            x = x.cpu().numpy()
+            x = (x - x.min()) / (x.max() - x.min() + 1e-8)
+            fakes.append(x)
+
+        fakes = np.concatenate(fakes)
+
+        # Get real samples for comparison
+        dataset = get_mnist_dataset(train=True)
+        n_samples = len(fakes)
+        reals = torch.stack([dataset[i][0] for i in range(n_samples)]).numpy()
+        reals = reals * 2 - 1
+
+        # Log metrics
+        fid_score = fid(reals, fakes)
+        kl_score = kl_divergence(reals, fakes)
+        jsd_score = jsd(reals, fakes)
+
+        mlflow.log_metric("FID", fid_score, step=epoch)
+        mlflow.log_metric("KL Divergence", kl_score, step=epoch)
+        mlflow.log_metric("JSD", jsd_score, step=epoch)
+
+        # Log sample images
+        fig, axes = plt.subplots(4, 8, figsize=(16, 8))
+        for i, ax in enumerate(axes.flat):
+            if i < len(fakes):
+                ax.imshow(fakes[i].transpose(1, 2, 0).squeeze(), cmap='gray')
+            ax.axis('off')
+        fig.suptitle(f"Generated Samples - Epoch {epoch}")
+        plt.tight_layout()
+
+        mlflow.log_figure(fig, f"samples_epoch_{epoch:03d}.png")
+
+        # Save to plots folder
+        fig.savefig(os.path.join(PLOTS_DIR, f"samples_epoch_{epoch:03d}.png"))
+        plt.close(fig)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Train MNIST diffusion model")
+    parser.add_argument("--epochs", type=int, default=EPOCHS, help="Number of epochs")
+    parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="Learning rate")
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Batch size")
+    parser.add_argument("--attention", action="store_true", help="Use self-attention")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint")
+    parser.add_argument("--name", type=str, default=None, help="Run name")
+
+    args = parser.parse_args()
+
+    torch.multiprocessing.set_start_method("spawn", force=True)
+
+    train_diffusion(
+        epochs=args.epochs,
+        learning_rate=args.lr,
+        batch_size=args.batch_size,
+        use_attention=args.attention,
+        checkpoint_path=args.checkpoint,
+        run_name=args.name,
+    )
diff --git a/src/utils.py b/src/utils.py
new file mode 100644
index 0000000..f727a68
--- /dev/null
+++ b/src/utils.py
@@ -0,0 +1,110 @@
+"""
+Utility functions for MNIST diffusion.
+"""
+
+import os
+import io
+import numpy as np
+import torch
+import scipy.linalg
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from .config import CHECKPOINT_DIR, PLOTS_DIR
+
+
+def ensure_dirs():
+    """Create necessary directories if they don't exist."""
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    os.makedirs(PLOTS_DIR, exist_ok=True)
+
+
+def tensor_to_image(tensor: torch.Tensor) -> np.ndarray:
+    """
+    Convert a tensor to a numpy image array.
+
+    Args:
+        tensor: Image tensor [C, H, W]
+
+    Returns:
+        Numpy array [H, W, C] normalized to [0, 1]
+    """
+    img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0)
+    img -= img.min()
+    img /= img.max() + 1e-8
+    return img
+
+
+def tensor_to_images(tensor: torch.Tensor) -> np.ndarray:
+    """
+    Convert a batch of tensors to numpy image arrays.
+
+    Args:
+        tensor: Batch of image tensors [B, C, H, W]
+
+    Returns:
+        Numpy array [B, H, W, C] normalized to [0, 1]
+    """
+    img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1)
+    img -= np.min(img, axis=(1, 2, 3), keepdims=True)
+    img /= np.max(img, axis=(1, 2, 3), keepdims=True) + 1e-8
+    return img
+
+
+def figure_to_image(figure: plt.Figure) -> np.ndarray:
+    """
+    Convert a matplotlib figure to a numpy image array.
+
+    Args:
+        figure: Matplotlib figure
+
+    Returns:
+        Numpy array of the figure image
+    """
+    buf = io.BytesIO()
+    figure.savefig(buf, format='png')
+    buf.seek(0)
+    image = np.array(Image.open(buf))
+    return image
+
+
+def save_checkpoint(model: torch.nn.Module, filename: str):
+    """
+    Save model checkpoint.
+
+    Args:
+        model: Model to save
+        filename: Filename (will be saved in CHECKPOINT_DIR)
+    """
+    ensure_dirs()
+    path = os.path.join(CHECKPOINT_DIR, filename)
+    torch.save(model.state_dict(), path)
+
+
+def load_checkpoint(model: torch.nn.Module, filename: str) -> torch.nn.Module:
+    """
+    Load model checkpoint.
+
+    Args:
+        model: Model architecture to load weights into
+        filename: Filename (loaded from CHECKPOINT_DIR)
+
+    Returns:
+        Model with loaded weights
+    """
+    path = os.path.join(CHECKPOINT_DIR, filename)
+    model.load_state_dict(torch.load(path, weights_only=True))
+    return model
+
+
+def save_plot(figure: plt.Figure, filename: str):
+    """
+    Save a matplotlib figure to the plots directory.
+
+    Args:
+        figure: Matplotlib figure to save
+        filename: Filename (will be saved in PLOTS_DIR)
+    """
+    ensure_dirs()
+    path = os.path.join(PLOTS_DIR, filename)
+    figure.savefig(path)
diff --git a/utils.py b/utils.py
deleted file mode 100644
index 9935d16..0000000
--- a/utils.py
+++ /dev/null
@@ -1,108 +0,0 @@
-import numpy as np
-import scipy.linalg
-import cv2
-
-
-def fid(reals, fakes):
-    """FID score calculation.
-
-    Args:
-        reals (numpy.array): Real images.
-        fakes (numpy.array): Fake images.
-    """
-    reals = reals.reshape(reals.shape[0], -1)
-    fakes = fakes.reshape(fakes.shape[0], -1)
-
-    mu_real = np.mean(reals, axis=0)
-    mu_fake = np.mean(fakes, axis=0)
-    sigma_real = np.cov(reals, rowvar=False)
-    sigma_fake = np.cov(fakes, rowvar=False)
-
-    diff = mu_real - mu_fake
-    covmean = np.dot(sigma_real, sigma_fake.T)
-    covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False)
-
-    if not np.isfinite(covmean).all():
-        eps=1e-6
-        offset = np.eye(sigma_real.shape[0]) * eps
-        ncovmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset))
-        covmean = ncovmean
-
-    if np.iscomplexobj(covmean):
-        covmean = covmean.real
-
-    return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean)
-
-
-def kl(reals, fakes):
-    """KL divergence calculation.
-
-    Args:
-        reals (numpy.array): Real images.
-        fakes (numpy.array): Fake images.
-    """
-    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
-    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)
-
-    hist_real = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals)
-    hist_fake = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes)
-
-    hist_real = hist_real + 1
-    hist_fake = hist_fake + 1
-
-    hist_real = hist_real / np.sum(hist_real)
-    hist_fake = hist_fake / np.sum(hist_fake)
-
-    return np.mean(np.log(hist_real / hist_fake))
-
-
-def jsd(reals, fakes):
-    """Jensen-Shannon divergence calculation.
-
-    Args:
-        reals (numpy.array): Real images.
-        fakes (numpy.array): Fake images.
-    """
-    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
-    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)
-
-    hist_real = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals)
-    hist_fake = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes)
-
-    hist_real = hist_real + 1
-    hist_fake = hist_fake + 1
-
-    hist_real = hist_real / np.sum(hist_real)
-    hist_fake = hist_fake / np.sum(hist_fake)
-
-    hist_avg = (hist_real + hist_fake) / 2
-
-    return 0.5 * (np.mean(np.log(hist_real / hist_avg)) + np.mean(np.log(hist_fake / hist_avg)))
-
-
-def haar(image):
-    # Load the Haar cascade for face detection
-    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-
-    # Convert the image to grayscale
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-
-    # Perform face detection
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=4)
-
-    # Draw rectangles around the detected faces
-    for (x, y, w, h) in faces:
-        cv2.rectangle(image, (x, y), (x+w, y+h), (255, 0, 0), 2)
-
-    # Display the result
-    cv2.imshow('Face Detection', image)
-    cv2.waitKey(0)
-    cv2.destroyAllWindows()
-
-
-if __name__ == '__main__':
-    # Load the image
-    image = cv2.imread('data/edface/500/03120500_000.png')
-
-    # Perform face detection
-    haar(image)
\ No newline at end of file