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