Blame for code/bifurcation/methods/edm/core.py

"""
core.py — EDM core math for conditional diffusion.

Reimplements the design of NVlabs/edm
(https://github.com/NVlabs/edm; specifically ``training/training_loop.py``)
following Karras et al. 2022, "Elucidating the Design Space of
Diffusion-Based Generative Models", NeurIPS 2022.

Method-only. No architecture, no data loading, no training loop.

Six components:
  - σ sampling at training time
  - preconditioning (c_skip, c_out, c_in, c_noise)
  - denoiser wrapper D_θ
  - training loss with λ(σ) weighting
  - σ schedule for sampling (ρ-parameterized)
  - Heun's 2nd-order sampler with Euler fallback at σ=0
"""

import torch
import math


# --- 1. Training-time σ sampling ---

def sample_training_sigma(batch_size, P_mean, P_std, device):
    """ see Table 1 in EDM paper. Formula for noise distribution under Training section.
    Returns: tensor of shape (batch_size,).
    """
    ln_sigma = torch.empty(batch_size, device=device).normal_(mean=P_mean, std=P_std)
    sigma = torch.exp(ln_sigma)
    return sigma


# --- 2. Preconditioning (EDM Table 1) ---

def c_skip(sigma, sigma_data):
    """Skip-connection scaling."""
    c_skip = sigma_data**2 / (sigma**2 +sigma_data**2)
    return c_skip

def c_out(sigma, sigma_data):
    """Output scaling."""
    c_out = sigma * sigma_data / (torch.sqrt(sigma_data**2+sigma**2))
    return c_out

def c_in(sigma, sigma_data):
    """Input scaling."""
    c_in=1 / (torch.sqrt(sigma**2+sigma_data**2))
    return c_in

def c_noise(sigma):
    """Time-like scalar fed into the network."""
    c_noise=0.25 * torch.log(sigma)
    return c_noise




# --- 3. Denoiser wrapper ---
def denoise(F_theta, x_noisy, sigma, x_0, sigma_data):
    """
    Wraps the raw network F_theta into the preconditioned denoiser D_theta. This corresponds to equation 7 in the EDM paper.
    D_θ(x, σ, x_0) = c_skip(σ)·x + c_out(σ) · F_θ(c_in(σ)·x, c_noise(σ), x_0)
    Used identically in both training (target reconstruction) and sampling (Heuns method)
    Returns predicted clean x_T (same shape as x_noisy).
    """
    # Step 1: normalize the noisy input
    x_normalized = c_in(sigma, sigma_data) * x_noisy

    # Step 2: embed the noise level
    sigma_embedded = c_noise(sigma)

    # Step 3: feed normalized inputs + conditioning through the network
    F_out = F_theta(x_normalized, sigma_embedded, x_0)

    # Step 4: combine via skip connection + scaled output
    D_theta = c_skip(sigma, sigma_data) * x_noisy + c_out(sigma, sigma_data) * F_out

    return D_theta


# --- 4. Loss ---

def lambda_weight(sigma, sigma_data):
    """Per-σ loss weighting (EDM Table 1, section training)."""
    lambda_sigma = (sigma**2 + sigma_data**2)/(sigma*sigma_data)**2
    return lambda_sigma

def edm_loss(F_theta, x_T_batch, x_0_batch, sigma_data, P_mean, P_std):
    """
    Compute the EDM training loss for one batch.

    Args:
        F_theta: the raw network
        x_T_batch: clean data, shape [B, D]
        x_0_batch: conditioning, shape [B, D]
        sigma_data: fixed scalar
        P_mean, P_std: log-normal sampling hyperparameters

    Returns:
        scalar loss
    """
    batch_size = x_T_batch.shape[0]
    device = x_T_batch.device

    # Step 1: sample sigma per example
    sigma = sample_training_sigma(batch_size, P_mean, P_std, device)
    # reshape for broadcasting: [B] -> [B, 1] for vector data
    sigma = sigma.view(-1, 1)

    # Step 2: sample noise
    epsilon = torch.randn_like(x_T_batch)

    # Step 3: add noise (EDM's simple forward process)
    x_noisy = x_T_batch + sigma * epsilon

    # Step 4: run through denoiser
    D_theta = denoise(F_theta, x_noisy, sigma, x_0_batch, sigma_data)

    # Step 5: compute weighted squared error
    squared_error = (D_theta - x_T_batch).pow(2).sum(dim=-1)  # [B]
    weighted_loss = lambda_weight(sigma.squeeze(-1), sigma_data) * squared_error  # [B]

    # Step 6: average over batch
    return weighted_loss.mean()


# --- 5. Sampling-time σ schedule ---

def sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device):
    """ρ-parameterized power-law schedule.

    Returns tensor of length N+1 with σ_0 = σ_max, σ_N = 0,
    and σ_i for 0 < i < N as in EDM Eq. 5.
    """
    # i goes from 0 to N-1
    i = torch.arange(N, device=device, dtype=torch.float32)

    # Apply the power-law formula
    inv_rho = 1.0 / rho
    sigma_max_inv = sigma_max ** inv_rho
    sigma_min_inv = sigma_min ** inv_rho
    sigmas = (sigma_max_inv + (i / (N - 1)) * (sigma_min_inv - sigma_max_inv)) ** rho

    # Append the final σ = 0 (we want to reach truly clean data)
    sigmas = torch.cat([sigmas, torch.zeros(1, device=device)])
    return sigmas


# --- 6. Heun sampler (EDM Algorithm 1) ---

def heun_sample(network, x_0, n_samples, N, sigma_min, sigma_max, rho, sigma_data, device):
    """Sample from p(x_T | x_0) via Heun's 2nd-order method.

    Initializes at x ~ N(0, σ_max² · I), iterates the σ schedule with
    Heun steps, falls back to Euler at the final step where σ_next = 0.

    Returns: tensor of shape (n_samples, x_0.shape[-1]).
    """
    # 1. Build the σ schedule (length N+1, ending at 0)
    sigmas = sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device)

    # 2. Broadcast x_0 to (n_samples, D) so every generated sample uses
    #    the same conditioning but different starting noise.
    if x_0.ndim == 1:
        x_0_batch = x_0.unsqueeze(0).expand(n_samples, -1)
    else:
        x_0_batch = x_0
    D_dim = x_0_batch.shape[-1]

    # 3. Initialize at pure noise scaled to σ_max.
    #    y₀ ~ N(0, σ_max²·I) means y₀ = σ_max · ε for ε ~ N(0, I).
    y = sigma_max * torch.randn(n_samples, D_dim, device=device)

    # 4. Iterate Heun steps along the schedule.
    for i in range(N):
        sigma_cur = sigmas[i]       # current noise level
        sigma_next = sigmas[i + 1]  # target noise level for this step

        # Reshape σ to match batch shape for broadcasting
        sigma_cur_b = sigma_cur.expand(n_samples, 1)

        # --- Predictor: Euler step ---
        # Evaluate the denoiser at the current point
        D_cur = denoise(network, y, sigma_cur_b, x_0_batch, sigma_data)
        # Compute the slope dy/dσ at the current point
        d_cur = (y - D_cur) / sigma_cur
        # Euler step toward σ_next (note: σ_next − σ_cur is negative since σ decreases)
        y_pred = y + (sigma_next - sigma_cur) * d_cur

        # --- Corrector: Heun's average slope ---
        if sigma_next > 0:
            # Evaluate the denoiser at the predicted next point
            sigma_next_b = sigma_next.expand(n_samples, 1)
            D_next = denoise(network, y_pred, sigma_next_b, x_0_batch, sigma_data)
            # Slope at the predicted point
            d_next = (y_pred - D_next) / sigma_next
            # Take the average slope and redo the step
            y = y + (sigma_next - sigma_cur) * 0.5 * (d_cur + d_next)
        else:
            # Final step: σ_next = 0 would divide by zero in d_next.
            # Just use the Euler estimate. This is the only Euler-style step.
            y = y_pred
    return y