1 files changed, 154 insertions, 0 deletions
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