src/dima/ddpmx.py

# how to use 
# ddpm = DDPM(Z_train, T=100, n_iter=20_000, key=random.PRNGKey(0))
# 
# Fast unconditional samples (try 15–30 first)
# z = ddpm.sample_dpmpp(N=4096, num_steps=20)
#
# Fast “refine latents” starting from an intermediate noise level
#z_ref = ddpm.refine_latents_dpmpp(z0, t_start=30, num_steps=20, add_noise=True)



# src/dima/ddpmx.py
from __future__ import annotations

from typing import Any, Optional

import jax
import jax.numpy as jnp
from jax import random

from flax import linen as nn
from flax.training import train_state
from flax import struct
import optax


# --- Add near the top of ddpm.py ---
import numpy as np

def _sigma_to_alpha_sigma_t(sigma: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
    """
    EDM-style sigma parameterization used by many DPM-Solver++ implementations:
      alpha_t = 1 / sqrt(1 + sigma^2)
      sigma_t = sigma * alpha_t
    so that x = alpha_t * x0 + sigma_t * eps
    """
    alpha_t = 1.0 / jnp.sqrt(1.0 + sigma**2)
    sigma_t = sigma * alpha_t
    return alpha_t, sigma_t


def _make_lu_sigma_schedule(
    sigma_start: float,
    sigma_end: float,
    num_steps: int,
) -> np.ndarray:
    """
    "Lu" schedule in lambda-space (uniform in lambda = -log(sigma)),
    which is a common default for DPM-Solver samplers.
    """
    sigma_start = float(max(sigma_start, 1e-12))
    sigma_end = float(max(sigma_end, 1e-12))
    lam_start = -np.log(sigma_start)
    lam_end = -np.log(sigma_end)
    lambdas = np.linspace(lam_start, lam_end, num_steps, dtype=np.float32)
    sigmas = np.exp(-lambdas).astype(np.float32)
    return sigmas


# ------------------------------
# Schedules & embeddings
# ------------------------------
def cosine_schedule(T: int, s: float = 0.008):
    """
    Nichol & Dhariwal cosine schedule.

    Returns:
        alpha:     (T,)
        beta:      (T,)
        alpha_bar: (T,)
    """
    steps = jnp.arange(T + 1, dtype=jnp.float32)
    f = jnp.cos(((steps / T + s) / (1.0 + s)) * jnp.pi / 2.0) ** 2
    alpha_bar_all = f / f[0]
    alpha_bar = alpha_bar_all[1:]  # (T,)
    alpha = alpha_bar / jnp.concatenate([jnp.array([1.0], dtype=jnp.float32), alpha_bar[:-1]])
    beta = 1.0 - alpha
    return alpha, beta, alpha_bar


def sinusoidal_embedding(t_idx: jnp.ndarray, dim: int) -> jnp.ndarray:
    """
    t_idx: (B,1) int32
    returns: (B,dim)
    """
    if t_idx.ndim != 2 or t_idx.shape[1] != 1:
        raise ValueError("t_idx must have shape (B,1)")
    t = t_idx.astype(jnp.float32)
    half = dim // 2
    denom = float(max(half - 1, 1))
    freqs = jnp.exp(-jnp.log(10_000.0) * jnp.arange(half, dtype=jnp.float32) / denom)
    args = t * freqs
    emb = jnp.concatenate([jnp.sin(args), jnp.cos(args)], axis=-1)
    if dim % 2 == 1:
        emb = jnp.pad(emb, ((0, 0), (0, 1)))
    return emb


# ------------------------------
# Model: epsilon predictor
# ------------------------------
class EpsMLP(nn.Module):
    """Simple MLP epsilon-predictor for DDPM in R^D."""
    hidden: int
    t_dim: int
    data_dim: int

    @nn.compact
    def __call__(self, x: jnp.ndarray, t_idx: jnp.ndarray) -> jnp.ndarray:
        # x: (B,D), t_idx: (B,1)
        t_emb = sinusoidal_embedding(t_idx, self.t_dim)
        t_h = nn.Dense(self.hidden)(t_emb)
        t_h = nn.gelu(t_h)

        h = nn.Dense(self.hidden)(x)
        h = nn.gelu(h + t_h)

        t_h2 = nn.Dense(self.hidden)(t_h)
        h = nn.Dense(self.hidden)(h)
        h = nn.gelu(h + t_h2)

        out = nn.Dense(self.data_dim)(h)
        return out


# ------------------------------
# TrainState with EMA
# ------------------------------
@struct.dataclass
class TrainStateEMA(train_state.TrainState):
    """Flax TrainState extended with EMA params."""
    ema_params: Any = struct.field(pytree_node=True)

    def apply_gradients(self, *, grads, ema_decay: float):
        updates, new_opt_state = self.tx.update(grads, self.opt_state, self.params)
        new_params = optax.apply_updates(self.params, updates)
        new_ema = optax.incremental_update(new_params, self.ema_params, step_size=1.0 - ema_decay)
        return self.replace(
            step=self.step + 1,
            params=new_params,
            opt_state=new_opt_state,
            ema_params=new_ema,
        )


# ------------------------------
# DDPM
# ------------------------------
class DDPM:
    """
    DDPM for D-dimensional latents.

    API expected by your DIMA wrapper:
      - DDPM(Z_train, ...): trains in __init__ (n_iter can be 0 to skip)
      - refine_latents(z0, t_start, key, add_noise) -> z_refined
      - __call__(...) delegates to refine_latents
      - sample(N) -> latent samples
      - attributes: state.params, state.ema_params, T, D, model.hidden, model.t_dim, beta_max, eps
    """

    def __init__(
        self,
        Z_iX: jnp.ndarray,
        *,
        T: int = 100,
        hidden_dim: int = 128,
        t_embed_dim: int = 64,
        learning_rate: float = 1e-3,
        n_iter: int = 20_000,
        ema_decay: float = 0.999,
        beta_max: float = 0.02,
        batch_size: Optional[int] = None,
        key: jax.Array = random.PRNGKey(0),
        verbose_every: int = 0,
        eps: float = 1e-5,
    ):
        Z_iX = jnp.asarray(Z_iX, dtype=jnp.float32)
        if Z_iX.ndim != 2:
            raise ValueError("Z_iX must be 2D (N,D).")

        self.D = int(Z_iX.shape[1])
        self.T = int(T)
        self.key = key
        self.ema_decay = float(ema_decay)
        self.batch_size = batch_size
        self.verbose_every = int(verbose_every)
        self.eps = float(eps)
        self.beta_max = float(beta_max)

        # schedule (cosine, clipped by beta_max)
        alpha, beta, alpha_bar = cosine_schedule(self.T)
        beta = jnp.minimum(beta, self.beta_max)
        alpha = 1.0 - beta
        alpha_bar = jnp.cumprod(alpha)

        self.alpha_s = alpha.astype(jnp.float32)
        self.beta_s = beta.astype(jnp.float32)
        self.alpha_bar_s = alpha_bar.astype(jnp.float32)

        # model + optimizer + EMA
        self.model = EpsMLP(hidden=int(hidden_dim), t_dim=int(t_embed_dim), data_dim=self.D)

        params = self.model.init(
            self.key,
            jnp.zeros((1, self.D), dtype=jnp.float32),
            jnp.zeros((1, 1), dtype=jnp.int32),
        )["params"]

        tx = optax.adam(float(learning_rate))
        self.state = TrainStateEMA.create(apply_fn=self.model.apply, params=params, tx=tx, ema_params=params)

        # train
        if int(n_iter) > 0:
            self._train(Z_iX, int(n_iter))

    # ---------- training ----------
    @staticmethod
    def _loss(params, apply_fn, x_t, t_idx, eps_true):
        eps_pred = apply_fn({"params": params}, x_t, t_idx)
        return jnp.mean((eps_pred - eps_true) ** 2)

    @staticmethod
    @jax.jit
    def _train_step(
        state: TrainStateEMA,
        x0_batch: jnp.ndarray,
        key: jax.Array,
        alpha_bar_s: jnp.ndarray,
        ema_decay: float,
        eps: float,
    ):
        B = x0_batch.shape[0]
        key, k_eps, k_t = random.split(key, 3)

        eps_noise = random.normal(k_eps, shape=x0_batch.shape)  # (B,D)
        t_idx = random.randint(k_t, shape=(B, 1), minval=0, maxval=alpha_bar_s.shape[0])

        a_bar_t = jnp.take(alpha_bar_s, t_idx.squeeze(-1))[:, None]
        a_bar_t = jnp.clip(a_bar_t, eps, 1.0)

        x_t = jnp.sqrt(a_bar_t) * x0_batch + jnp.sqrt(1.0 - a_bar_t) * eps_noise

        def loss_fn(p):
            return DDPM._loss(p, state.apply_fn, x_t, t_idx, eps_noise)

        loss, grads = jax.value_and_grad(loss_fn)(state.params)
        new_state = state.apply_gradients(grads=grads, ema_decay=ema_decay)
        return new_state, loss, key

    def _train(self, Z_iX: jnp.ndarray, n_iter: int):
        N = int(Z_iX.shape[0])
        bs = N if (self.batch_size is None) else min(int(self.batch_size), N)

        for it in range(n_iter):
            if bs >= N:
                batch = Z_iX
            else:
                self.key, k_perm = random.split(self.key)
                idx = random.permutation(k_perm, N)[:bs]
                batch = Z_iX[idx]

            self.state, loss, self.key = self._train_step(
                self.state,
                batch,
                self.key,
                self.alpha_bar_s,
                self.ema_decay,
                self.eps,
            )

            if self.verbose_every and (it % self.verbose_every == 0 or it == n_iter - 1):
                print(f"iter {it:6d}  loss {float(loss):.6f}", end="\r")

        if self.verbose_every:
            print("\ntraining complete.")

    # ---------- diffusion utilities ----------
    @staticmethod
    def _posterior_variance(alpha_s, beta_s, alpha_bar_s, t):
        a_bar_t = alpha_bar_s[t]
        a_bar_prev = jnp.where(t > 0, alpha_bar_s[t - 1], jnp.array(1.0, dtype=alpha_bar_s.dtype))
        return ((1.0 - a_bar_prev) / (1.0 - a_bar_t)) * beta_s[t]

    @staticmethod
    def _make_sampler_step(params_ema, apply_fn, alpha_s, beta_s, alpha_bar_s, eps: float):
        @jax.jit
        def step(carry, _):
            key, t, x = carry  # x: (B,D)
            key, k = random.split(key)

            alpha_t = jnp.clip(alpha_s[t], eps, 1.0)
            a_bar_t = jnp.clip(alpha_bar_s[t], eps, 1.0)

            sqrt_alpha = jnp.sqrt(alpha_t)
            sqrt_one_minus_a_bar = jnp.sqrt(jnp.clip(1.0 - a_bar_t, eps, 1.0))

            B = x.shape[0]
            t_batch = jnp.full((B, 1), t, dtype=jnp.int32)
            eps_pred = apply_fn({"params": params_ema}, x, t_batch)  # (B,D)

            # predict x0
            x0_hat = (x - sqrt_one_minus_a_bar * eps_pred) / jnp.sqrt(a_bar_t)

            a_bar_prev = jnp.where(t > 0, alpha_bar_s[t - 1], jnp.array(1.0, dtype=alpha_bar_s.dtype))
            denom = jnp.clip(1.0 - a_bar_t, eps, 1.0)

            coef1 = jnp.sqrt(jnp.clip(a_bar_prev, eps, 1.0)) * beta_s[t] / denom
            coef2 = sqrt_alpha * (1.0 - a_bar_prev) / denom

            mean = coef1 * x0_hat + coef2 * x

            beta_tilde = DDPM._posterior_variance(alpha_s, beta_s, alpha_bar_s, t)
            sigma = jnp.sqrt(jnp.clip(beta_tilde, 0.0, 1.0))

            z = random.normal(k, x.shape)
            z = jnp.where(t == 0, 0.0, z)
            x_prev = mean + sigma * z

            return (key, t - 1, x_prev), x_prev

        return step

    # ---------- API ----------
    def refine_latents(
        self,
        z0: jnp.ndarray,
        t_start: int = 10,
        key: Optional[jax.Array] = None,
        add_noise: bool = True,
    ) -> jnp.ndarray:
        """
        Refine latents by:
          (optional) forward-noise z0 to step t_start
          reverse-diffuse from t_start -> 0 using EMA params.
        """
        z0 = jnp.asarray(z0, dtype=jnp.float32)
        if z0.ndim != 2 or z0.shape[1] != self.D:
            raise ValueError(f"z0 must have shape (B,{self.D}).")
        if not (0 <= int(t_start) < self.T):
            raise ValueError(f"t_start must be in [0, {self.T-1}]")

        t_start = int(t_start)

        if key is None:
            self.key, key = random.split(self.key)
        else:
            # advance internal RNG too
            self.key, _ = random.split(key)

        # forward diffuse to t_start
        key, k_eps = random.split(key)
        eps_noise = random.normal(k_eps, z0.shape)

        a_bar_t = jnp.clip(self.alpha_bar_s[t_start], self.eps, 1.0)
        if add_noise:
            z_t = jnp.sqrt(a_bar_t) * z0 + jnp.sqrt(1.0 - a_bar_t) * eps_noise
        else:
            z_t = z0

        step = self._make_sampler_step(
            self.state.ema_params,
            self.state.apply_fn,
            self.alpha_s,
            self.beta_s,
            self.alpha_bar_s,
            self.eps,
        )

        (final_key, _, _), trace = jax.lax.scan(
            step,
            (key, t_start, z_t),
            xs=None,
            length=t_start + 1,
        )
        self.key = final_key
        return trace[-1]

    def __call__(
        self,
        z0: jnp.ndarray,
        t_start: int = 10,
        key: Optional[jax.Array] = None,
        add_noise: bool = True,
    ) -> jnp.ndarray:
        return self.refine_latents(z0, t_start=t_start, key=key, add_noise=add_noise)

    def reverse_from_T(self, x_T: jnp.ndarray) -> jnp.ndarray:
        x_T = jnp.asarray(x_T, dtype=jnp.float32)
        if x_T.ndim != 2 or x_T.shape[1] != self.D:
            raise ValueError(f"x_T must have shape (B,{self.D}).")

        step = self._make_sampler_step(
            self.state.ema_params,
            self.state.apply_fn,
            self.alpha_s,
            self.beta_s,
            self.alpha_bar_s,
            self.eps,
        )
        self.key, k0 = random.split(self.key)
        (_, _, _), trace = jax.lax.scan(
            step,
            (k0, self.T - 1, x_T),
            xs=None,
            length=self.T,
        )
        return trace[-1]

    def sample(self, N: int = 10_000) -> jnp.ndarray:
        self.key, k = random.split(self.key)
        noise = random.normal(k, (int(N), self.D))
        return self.reverse_from_T(noise)

    # --- In __init__ AFTER you compute self.alpha_bar_s ---
    # sigma_in is the "EDM-style" sigma = sigma_t / alpha_t = sqrt((1-a_bar)/a_bar)
    self.sigma_in_train = jnp.sqrt(
        jnp.clip((1.0 - self.alpha_bar_s) / jnp.clip(self.alpha_bar_s, self.eps, 1.0), self.eps, 1e12)
    ).astype(jnp.float32)

    def _make_dpmpp_schedule(self, *, num_steps: int, t_start: int) -> tuple[jnp.ndarray, jnp.ndarray]:
        """
        Returns:
          sigmas_in: (K+1,) float32, decreasing, last one is 0
          t_cont:    (K,)   float32, "continuous" time indices to feed the model
        """
        t_start = int(t_start)
        if not (0 <= t_start < self.T):
            raise ValueError(f"t_start must be in [0, {self.T-1}]")
        if int(num_steps) < 1:
            raise ValueError("num_steps must be >= 1")

        # Start/end sigmas from your trained VP schedule
        sigma_start = float(self.sigma_in_train[t_start])
        sigma_end = float(self.sigma_in_train[0])

        # Build sigma schedule (length K) and append final sigma=0
        sigmas_k = _make_lu_sigma_schedule(sigma_start, sigma_end, int(num_steps))
        sigmas = np.concatenate([sigmas_k, np.array([0.0], np.float32)], axis=0)

        # Map sigma -> continuous t by interpolating in log(sigma) over the training sigmas
        sigma_train = np.array(self.sigma_in_train).astype(np.float32)  # (T,)
        log_sig_train = np.log(np.maximum(sigma_train, 1e-12))          # increasing with t
        t_train = np.arange(self.T, dtype=np.float32)

        # For model calls, we only need times for the K "current" sigmas (exclude final 0)
        log_sig = np.log(np.maximum(sigmas[:-1], 1e-12))
        t_cont = np.interp(log_sig, log_sig_train, t_train).astype(np.float32)  # (K,)

        return jnp.array(sigmas, dtype=jnp.float32), jnp.array(t_cont, dtype=jnp.float32)

    @staticmethod
    @jax.jit
    def _dpmpp_2m_midpoint_sample(
        params_ema: Any,
        apply_fn: Any,
        x_start: jnp.ndarray,     # (B,D)
        sigmas_in: jnp.ndarray,   # (K+1,)
        t_cont: jnp.ndarray,      # (K,)
        eps: float,
    ) -> jnp.ndarray:
        """
        DPM-Solver++ (2M, midpoint) sampler.

        Uses:
          - first-order update on step 0
          - second-order midpoint update on steps 1..K-2
          - first-order update on final step K-1 (important when next sigma is 0)
        """
        # Current/next sigma for each step
        sigma_s = sigmas_in[:-1]  # (K,)
        sigma_t = sigmas_in[1:]   # (K,)

        alpha_s, sigma_s_t = _sigma_to_alpha_sigma_t(sigma_s)  # both (K,)
        alpha_t, sigma_t_t = _sigma_to_alpha_sigma_t(sigma_t)  # both (K,)

        # lambda = log(alpha) - log(sigma_t); for this parameterization it's ~ -log(sigma_in)
        # Avoid log(0) warnings: allow inf; we will not use 2nd-order formula on final step.
        lambda_s = jnp.log(alpha_s) - jnp.log(sigma_s_t)
        lambda_t = jnp.log(alpha_t) - jnp.log(sigma_t_t)

        K = t_cont.shape[0]
        is_first = jnp.arange(K) == 0
        is_last = jnp.arange(K) == (K - 1)

        def step(carry, inp):
            x, m_prev, lam_prev = carry
            (a_s, s_s, a_t, s_t, lam_s, lam_t, t_i, first_i, last_i) = inp

            B = x.shape[0]
            t_batch = jnp.full((B, 1), t_i, dtype=jnp.float32)

            # model predicts epsilon
            eps_pred = apply_fn({"params": params_ema}, x, t_batch)  # (B,D)

            # convert epsilon -> x0 (data prediction)
            a_s_b = jnp.clip(a_s, eps, 1.0)
            x0 = (x - s_s * eps_pred) / a_s_b  # (B,D)

            h = lam_t - lam_s
            exp_neg_h = jnp.exp(-h)

            # 1st-order (DDIM-like in dpmsolver++ form)
            x_first = (s_t / s_s) * x - (a_t * (exp_neg_h - 1.0)) * x0

            def do_second(_):
                # 2nd-order multistep midpoint
                h0 = lam_s - lam_prev
                # r0 = h0/h; clip to avoid division by 0 if something degenerate happens
                r0 = h0 / jnp.clip(h, 1e-12)
                D1 = (x0 - m_prev) / jnp.clip(r0, 1e-12)
                x_second = (s_t / s_s) * x - (a_t * (exp_neg_h - 1.0)) * (x0 + 0.5 * D1)
                return x_second

            x_next = jax.lax.cond(first_i | last_i, lambda _: x_first, do_second, operand=None)
            return (x_next, x0, lam_s), x_next

        # Pack per-step scalars
        xs = (
            alpha_s, sigma_s_t,
            alpha_t, sigma_t_t,
            lambda_s, lambda_t,
            t_cont, is_first, is_last
        )

        # init prev buffers (unused on first step)
        x0_init = jnp.zeros_like(x_start)
        lam_init = jnp.array(0.0, dtype=jnp.float32)

        (x_final, _, _), _trace = jax.lax.scan(step, (x_start, x0_init, lam_init), xs)
        return x_final

    # --- Public API: fast refine/sample methods ---
    def refine_latents_dpmpp(
        self,
        z0: jnp.ndarray,
        *,
        t_start: int = 10,
        num_steps: int = 20,
        key: Optional[jax.Array] = None,
        add_noise: bool = True,
    ) -> jnp.ndarray:
        z0 = jnp.asarray(z0, dtype=jnp.float32)
        if z0.ndim != 2 or z0.shape[1] != self.D:
            raise ValueError(f"z0 must have shape (B,{self.D}).")
        if not (0 <= int(t_start) < self.T):
            raise ValueError(f"t_start must be in [0, {self.T-1}]")

        if key is None:
            self.key, key = random.split(self.key)
        else:
            self.key, _ = random.split(key)

        # Forward-noise to t_start (same as your original code)
        key, k_eps = random.split(key)
        eps_noise = random.normal(k_eps, z0.shape)

        a_bar = jnp.clip(self.alpha_bar_s[int(t_start)], self.eps, 1.0)
        if add_noise:
            x_start = jnp.sqrt(a_bar) * z0 + jnp.sqrt(1.0 - a_bar) * eps_noise
        else:
            x_start = z0

        sigmas_in, t_cont = self._make_dpmpp_schedule(num_steps=int(num_steps), t_start=int(t_start))

        x_final = self._dpmpp_2m_midpoint_sample(
            self.state.ema_params,
            self.state.apply_fn,
            x_start,
            sigmas_in,
            t_cont,
            self.eps,
        )
        return x_final

    def reverse_from_T_dpmpp(self, x_T: jnp.ndarray, *, num_steps: int = 20) -> jnp.ndarray:
        x_T = jnp.asarray(x_T, dtype=jnp.float32)
        if x_T.ndim != 2 or x_T.shape[1] != self.D:
            raise ValueError(f"x_T must have shape (B,{self.D}).")

        sigmas_in, t_cont = self._make_dpmpp_schedule(num_steps=int(num_steps), t_start=self.T - 1)

        return self._dpmpp_2m_midpoint_sample(
            self.state.ema_params,
            self.state.apply_fn,
            x_T,
            sigmas_in,
            t_cont,
            self.eps,
        )

    def sample_dpmpp(self, N: int = 10_000, *, num_steps: int = 20) -> jnp.ndarray:
        self.key, k = random.split(self.key)
        x_T = random.normal(k, (int(N), self.D)).astype(jnp.float32)
        return self.reverse_from_T_dpmpp(x_T, num_steps=int(num_steps))


__all__ = ["DDPM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]