Generative Imaging

Generative Imaging

When to Use

• Formulating image restoration (denoising, inpainting, super-resolution) as an inverse problem y = Hx + noise solved with a diffusion prior (DDRM)
• Sampling from a diffusion model with DDIM to reconstruct an image in fewer steps than training
• Restoring degraded scientific images — cryo-EM micrographs, MRI k-space undersampling, fluorescence microscopy — where the forward degradation operator H is known
• Building/debugging noise schedules (linear vs cosine) or checking whether ᾱₜ indexing is inverted
• Visualizing a score field ∇ₓ log p(x) to build intuition for reverse diffusion

Version Compatibility

• Python ≥3.10, NumPy ≥1.24 (all patterns below are pure NumPy — no GPU/framework required)
• Concepts follow DDPM (Ho et al. 2020), DDIM (Song et al. 2021), and DDRM (Kawar et al. 2022)
• For a trainable neural denoiser instead of the oracle used here, pair with PyTorch (see ai-science-diffusion-generative-models)

Prerequisites

• pip install numpy (add matplotlib for the plots in the examples)
• Familiarity with Gaussian diffusion basics: forward process xₜ = √ᾱₜ x₀ + √(1-ᾱₜ) ε
• Linear algebra: SVD, pseudoinverse, weighted least squares

Goal: turn a noisy/masked/downsampled measurement y = Hx₀ + σ·ε into a restored image x̂₀ using a diffusion prior instead of a hand-tuned regularizer.

Approach: run DDIM reverse steps with a denoiser that predicts x₀, then after each step project the prediction onto the measurement-consistent subspace via the SVD of H (DDRM), and feed the corrected x₀ back into the next reverse step.

import numpy as np


def linear_schedule(T=1000, beta_start=1e-4, beta_end=0.02):
    """Standard DDPM linear noise schedule."""
    betas = np.linspace(beta_start, beta_end, T)
    alphas = 1 - betas
    alpha_bars = np.cumprod(alphas)
    return betas, alphas, alpha_bars


def cosine_schedule(T=1000, s=0.008):
    """Improved cosine schedule (Nichol & Dhariwal 2021); smoother than linear."""
    t = np.arange(T + 1)
    f = np.cos((t / T + s) / (1 + s) * np.pi / 2) ** 2
    alpha_bars = f / f[0]
    betas = np.clip(1 - alpha_bars[1:] / alpha_bars[:-1], 0, 0.999)
    return betas, 1 - betas, alpha_bars[1:]


def forward_diffuse(x0, t, alpha_bars, rng=None):
    """Add noise to x0 at timestep t: xt = sqrt(abar_t) x0 + sqrt(1-abar_t) eps."""
    rng = rng or np.random.default_rng()
    ab = alpha_bars[t]
    eps = rng.standard_normal(x0.shape)
    xt = np.sqrt(ab) * x0 + np.sqrt(1 - ab) * eps
    return xt, eps


def ddim_step(xt, eps_pred, t, t_prev, alpha_bars, eta=0.0):
    """One DDIM reverse step. eta=0 -> deterministic; eta=1 -> DDPM-like stochastic."""
    ab_t = alpha_bars[t]
    ab_prev = alpha_bars[t_prev] if t_prev >= 0 else 1.0
    x0_pred = (xt - np.sqrt(1 - ab_t) * eps_pred) / (np.sqrt(ab_t) + 1e-9)
    x0_pred = np.clip(x0_pred, -1.5, 1.5)  # prevents artifact accumulation
    sigma = eta * np.sqrt((1 - ab_prev) / (1 - ab_t + 1e-9) * (1 - ab_t / ab_prev))
    noise = sigma * np.random.randn(*xt.shape) if eta > 0 else 0
    xt_prev = (np.sqrt(ab_prev) * x0_pred
               + np.sqrt(max(1 - ab_prev - sigma**2, 0)) * eps_pred
               + noise)
    return xt_prev, x0_pred

Goal: enforce that the reconstructed image stays consistent with the observed measurement y at every reverse step (DDRM data consistency).

Approach: decompose H with SVD once, then at each step blend the diffusion prior's x0_pred toward y in proportion to the singular-value SNR — well-observed directions (large s) trust y more, ill-posed directions (small/zero s) trust the prior more.

import numpy as np


def ddrm_projection(x0_pred, y_obs, H, sigma_obs=0.05):
    """
    General DDRM data-consistency step for a linear degradation y = H @ x0 + sigma_obs * noise.
    H: (m, n) degradation matrix (identity=denoising, row-subset=inpainting,
       block-average=downsampling for super-resolution). Uses the FULL V basis so that
       directions H can't observe (its null space) keep the prior's x0_pred untouched --
       a thin SVD would silently drop them.
    """
    U, s, Vt = np.linalg.svd(H, full_matrices=True)   # U:(m,m) Vt:(n,n) s:(m,) assuming m<=n
    x0_flat = x0_pred.flatten()
    v_coeff = Vt @ x0_flat                       # (n,) x0 in H's right-singular basis
    y_coeff = U.T @ y_obs                        # (m,) y in H's left-singular basis
    m = len(s)
    safe_s = np.where(s > 1e-8, s, 1.0)
    trust = np.where(s > 1e-8, s**2 / (s**2 + sigma_obs**2), 0.0)  # 0 for unobserved directions
    v_coeff_updated = v_coeff.copy()
    v_coeff_updated[:m] = (1 - trust) * v_coeff[:m] + trust * (y_coeff / safe_s)
    return (Vt.T @ v_coeff_updated).reshape(x0_pred.shape)


def ddrm_inpainting_step(x0_pred, y_masked, mask, sigma_obs=0.05):
    """Special case of DDRM projection for a binary mask (H = diag(mask)); avoids full SVD."""
    scale = 1 / (1 + sigma_obs**2)
    x0_proj = x0_pred.copy()
    observed = mask == 1
    x0_proj[observed] = x0_pred[observed] + scale * (y_masked[observed] - x0_pred[observed])
    return x0_proj


def ddim_ddrm_sample(denoiser, xT, y_obs, mask_or_H, alpha_bars, timesteps,
                      is_mask=True, sigma_obs=0.05):
    """
    DDIM reverse loop with DDRM data-consistency injected after every step.
    denoiser(xt, t) -> eps_pred must be supplied by the caller (oracle or trained network).
    """
    xt = xT.copy()
    for i, t in enumerate(timesteps[:-1]):
        t_prev = timesteps[i + 1]
        eps_pred = denoiser(xt, t)
        _, x0_pred = ddim_step(xt, eps_pred, t, t_prev, alpha_bars)
        if is_mask:
            x0_proj = ddrm_inpainting_step(x0_pred, y_obs, mask_or_H, sigma_obs)
        else:
            x0_proj = ddrm_projection(x0_pred, y_obs, mask_or_H, sigma_obs)
        # substitute the data-consistent x0 into the same DDIM combining formula
        # (do NOT round-trip through the old xt -- that algebraically cancels the correction)
        ab_prev = alpha_bars[t_prev] if t_prev >= 0 else 1.0
        xt = np.sqrt(ab_prev) * x0_proj + np.sqrt(1 - ab_prev) * eps_pred
    return xt

Goal: sanity-check a restoration setup end to end on a toy 1D "image" before wiring in a real network.

Approach: build ground truth, corrupt it with a known mask, run the loop above with an oracle denoiser (one that knows x0, standing in for a trained network), and confirm the reconstruction error drops versus the naive noisy observation.

def demo():
    """Self-check: DDRM inpainting should beat the naive masked observation."""
    rng = np.random.default_rng(0)
    n = 64
    x0_true = np.sin(np.linspace(0, 4 * np.pi, n)) * 0.8
    _, _, alpha_bars = cosine_schedule(T=1000)

    xT, _ = forward_diffuse(x0_true, 999, alpha_bars, rng)
    timesteps = np.linspace(999, 1, 20).astype(int)

    mask = np.ones(n)
    mask[n // 4: 3 * n // 4] = 0  # mask the middle half
    y_masked = x0_true * mask

    def oracle_denoiser(xt, t):
        ab = alpha_bars[t]
        return (xt - np.sqrt(ab) * x0_true) / np.sqrt(1 - ab + 1e-9)

    x_hat = ddim_ddrm_sample(oracle_denoiser, xT, y_masked, mask, alpha_bars,
                              timesteps, is_mask=True)

    mse_masked_region = ((x_hat[mask == 0] - x0_true[mask == 0]) ** 2).mean()
    assert mse_masked_region < 0.5, f"DDRM inpainting failed to converge: MSE={mse_masked_region:.3f}"
    print(f"DDRM inpainting MSE in masked region: {mse_masked_region:.4f}")

    # Cross-check the general SVD path (ddrm_projection) against a known-good identity case:
    # for H=I, projecting toward y should reduce MSE relative to the raw noisy observation.
    H_identity = np.eye(n)
    y_noisy = x0_true + 0.3 * rng.standard_normal(n)
    x0_proj = ddrm_projection(x0_true, y_noisy, H_identity, sigma_obs=0.3)
    mse_noisy = ((y_noisy - x0_true) ** 2).mean()
    mse_proj = ((x0_proj - x0_true) ** 2).mean()
    assert mse_proj < mse_noisy, "ddrm_projection should pull the estimate toward x0_true"
    print(f"ddrm_projection MSE: noisy={mse_noisy:.4f} -> projected={mse_proj:.4f}")


if __name__ == "__main__":
    demo()

Pitfalls

• Tensor shape: betas/alpha_bars must be shape (T,), not (T, 1) — broadcasting against images fails silently
• ᾱₜ indexing convention: alpha_bars[T-1] is most noisy, alpha_bars[0] is least noisy; mixing this up flips the whole schedule
• Clipping x0_pred: without clipping to a sane range (e.g. [-1, 1] or [-1.5, 1.5]), errors compound over DDIM steps and the image diverges
• SVD conditioning: near-zero singular values in H (e.g. aggressive downsampling) make the pseudoinverse unstable — regularize with sigma_obs or truncate small singular values
• Re-deriving xt after projection: skipping the re-projection step after ddrm_projection/ddrm_inpainting_step leaves the trajectory inconsistent with the corrected x0, and later steps silently undo the correction
• Cosine vs linear schedule: cosine keeps more signal for longer near t=0, which matters for high-frequency image detail; linear schedules over-noise images too quickly for restoration tasks

See Also

• ai-science-diffusion-generative-models — full DDPM/DDIM training and sampling reference
• bio-applied-deep-learning-for-biology — training the neural denoiser this skill treats as an oracle
• bio-structural-biology-modern-structure-prediction — cryo-EM/structure reconstruction context for imaging restoration
• ai-science-llm-finetuning — contrasting autoregressive generative paradigm