trying to improve whole project

9 files changed, 1268 insertions, 0 deletions

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)