""" 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