From 06887ce14321d8637d41def0674adb4feaca7fc0 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:44 -0400 Subject: [PATCH 1/5] feat(schedulers): flow matching, min-SNR-gamma weighting, EDM2 defaults Three changes to the training objective layer, sharing one contract. Flow matching slots into the existing scheduler/predictor seams rather than next to them: FlowMatchingScheduler is the linear path alpha = 1 - t, sigma = t with SD3's logit-normal timesteps and Flux's resolution shift, and FlowMatchPredictionTransform predicts the velocity u = eps - x0. Because the path is linear, the unmodified DDIM and Euler updates already are the flow Euler step; there is a test proving the identity exactly. min-SNR-gamma (Hang et al. 2023) is defined on the x_0 loss, so converting it into the space the trainer actually computes the loss in needs the parameterization. get_weights becomes a template over get_schedule_weights and asks the transform for that conversion, which also makes the schedule/transform pairing explicit instead of implicit. No second weighting system. EDMNoiseScheduler's P_mean/P_std were hardcoded to EDM1's -1.2/1.2; they are now constructor arguments defaulting to EDM2's -0.4/1.0. Co-Authored-By: Claude Fable 5 --- flaxdiff/predictors/__init__.py | 68 ++++++++- flaxdiff/schedulers/__init__.py | 3 +- flaxdiff/schedulers/common.py | 30 +++- flaxdiff/schedulers/cosine.py | 2 +- flaxdiff/schedulers/discrete.py | 2 +- flaxdiff/schedulers/flow.py | 55 +++++++ flaxdiff/schedulers/karras.py | 18 ++- tests/test_flow_matching.py | 249 ++++++++++++++++++++++++++++++++ tests/test_schedulers.py | 87 ++++++++++- 9 files changed, 490 insertions(+), 24 deletions(-) create mode 100644 flaxdiff/schedulers/flow.py create mode 100644 tests/test_flow_matching.py diff --git a/flaxdiff/predictors/__init__.py b/flaxdiff/predictors/__init__.py index 68390f2..befce3b 100644 --- a/flaxdiff/predictors/__init__.py +++ b/flaxdiff/predictors/__init__.py @@ -7,6 +7,7 @@ CosineNoiseScheduler, KarrasVENoiseScheduler, EDMNoiseScheduler, + FlowMatchingScheduler, ) ############################################################################################################ @@ -39,6 +40,14 @@ def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray]) -> jnp.ndarray: return 1 + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + """||target error||^2 / ||x_0 error||^2 at the given SNR. + + Loss weights (min-SNR-gamma and friends) are defined on the x_0 loss; + dividing by this converts them into the space the model is trained in. + """ + return 1.0 + class EpsilonPredictionTransform(DiffusionPredictionTransform): def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]: # preds is the predicted noise @@ -46,10 +55,13 @@ def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) signal_rates, noise_rates = rates x_0 = (x_t - epsilon * noise_rates) / signal_rates return x_0, epsilon - + def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: return epsilon + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + return snr + class DirectPredictionTransform(DiffusionPredictionTransform): def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]: # Here the model predicts x_0 directly @@ -76,7 +88,27 @@ def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: v = signal_rate * epsilon - noise_rate * x_0 variance = signal_rate**2 + noise_rate**2 return v / jnp.sqrt(variance) - + + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + return snr + 1 + +class FlowMatchPredictionTransform(DiffusionPredictionTransform): + """Rectified flow velocity: the model predicts u = epsilon - x_0, the + constant velocity of the linear path, so both endpoints are one step away. + """ + def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]: + signal_rate, noise_rate = rates + x_0 = x_t - noise_rate * preds + epsilon = x_t + signal_rate * preds + return x_0, epsilon + + def get_target(self, x_0, epsilon, rates) ->jnp.ndarray: + return epsilon - x_0 + + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + # x_0 error is t times the velocity error, and t = 1 / (1 + sqrt(SNR)) + return (1 + jnp.sqrt(snr)) ** 2 + class KarrasPredictionTransform(DiffusionPredictionTransform): def __init__(self, sigma_data=0.5) -> None: super().__init__() @@ -102,6 +134,11 @@ def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8) c_in = 1 / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon) return c_in + def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray: + # x_0 error is c_out times the raw error, and alpha = 1 here so + # sigma^2 = 1 / SNR + return 1 / self.sigma_data ** 2 + snr + ############################################################################################################ # Noise schedule presets ############################################################################################################ @@ -113,6 +150,10 @@ def get_diffusion_preset( sigma_max: float = 80.0, rho: float = 7.0, sigma_data: float = 0.5, + P_mean: float = -0.4, + P_std: float = 1.0, + shift: float = 1.0, + min_snr_gamma: float = None, ) -> tuple[NoiseScheduler, NoiseScheduler, DiffusionPredictionTransform]: """Named (train schedule, sampling schedule, prediction transform) presets. @@ -120,18 +161,29 @@ def get_diffusion_preset( parameterization. Both training and inference build from here, so a model is always sampled with the same convention it was trained with. """ + # Only the training schedule ever produces loss weights, so only it needs + # to know the parameterization if name == 'edm': - train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) - sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) transform = KarrasPredictionTransform(sigma_data=sigma_data) - elif name == 'karras': - train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) + train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data, + P_mean=P_mean, P_std=P_std, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) + elif name == 'karras': transform = KarrasPredictionTransform(sigma_data=sigma_data) + train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) + sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data) elif name == 'cosine': - train = CosineNoiseScheduler(timesteps, beta_end=1) - sample = CosineNoiseScheduler(timesteps, beta_end=1) transform = VPredictionTransform() + train = CosineNoiseScheduler(timesteps, beta_end=1, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) + sample = CosineNoiseScheduler(timesteps, beta_end=1) + elif name in ('flow', 'flow_matching'): + transform = FlowMatchPredictionTransform() + train = FlowMatchingScheduler(shift=shift, + prediction_transform=transform, min_snr_gamma=min_snr_gamma) + sample = FlowMatchingScheduler(shift=shift) else: raise ValueError(f"Unknown noise schedule preset: {name}") return train, sample, transform diff --git a/flaxdiff/schedulers/__init__.py b/flaxdiff/schedulers/__init__.py index 32594e4..278bdb9 100644 --- a/flaxdiff/schedulers/__init__.py +++ b/flaxdiff/schedulers/__init__.py @@ -3,4 +3,5 @@ from .cosine import * from .linear import * from .sqrt import * -from .karras import * \ No newline at end of file +from .karras import * +from .flow import * diff --git a/flaxdiff/schedulers/common.py b/flaxdiff/schedulers/common.py index 3061542..e4c2459 100644 --- a/flaxdiff/schedulers/common.py +++ b/flaxdiff/schedulers/common.py @@ -18,11 +18,17 @@ def __init__(self, timesteps, dtype=jnp.float32, clip_min=-1.0, clip_max=1.0, + min_snr_gamma: float = None, + prediction_transform=None, *args, **kwargs): self.max_timesteps = timesteps self.dtype = dtype self.clip_min = clip_min self.clip_max = clip_max + if min_snr_gamma is not None and prediction_transform is None: + raise ValueError("min_snr_gamma needs the prediction transform it will be trained against") + self.min_snr_gamma = min_snr_gamma + self.prediction_transform = prediction_transform if type(timesteps) == int and timesteps > 1: timestep_generator = lambda rng, batch_size, max_timesteps = timesteps: jax.random.randint(rng, (batch_size,), 0, max_timesteps) else: @@ -33,10 +39,28 @@ def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.n state, rng = state.get_random_key() timesteps = self.timestep_generator(rng, batch_size, self.max_timesteps) return timesteps, state - + + def get_snr(self, steps) -> jnp.ndarray: + signal_rates, noise_rates = self.get_rates(steps, shape=(-1,)) + return (signal_rates / noise_rates) ** 2 + def get_weights(self, steps, shape=(-1, 1, 1, 1)): + """Per-sample loss weight, in the target space of the paired parameterization. + + min-SNR-gamma (Hang et al. 2023) is defined as min(SNR, gamma) on the + x_0 loss; the transform converts that into whichever space the trainer + actually computes the loss in. It replaces the schedule's own weighting + rather than stacking on top of it. + """ + if self.min_snr_gamma is None: + return self.get_schedule_weights(steps, shape) + snr = self.get_snr(steps) + weights = jnp.minimum(snr, self.min_snr_gamma) / self.prediction_transform.target_error_scale(snr) + return weights.reshape(shape) + + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)): raise NotImplementedError - + def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]: raise NotImplementedError @@ -77,7 +101,7 @@ def __init__(self, timesteps, sigma_min=0.002, sigma_max=80.0, sigma_data=1, *ar self.sigma_max = sigma_max self.sigma_data = sigma_data - def get_weights(self, steps, shape=(-1, 1, 1, 1)): + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)): sigma = self.get_sigmas(steps) return (1 + (1 / (1 + ((1 - sigma ** 2)/(sigma ** 2)))) / (self.sigma_max ** 2)).reshape(shape) diff --git a/flaxdiff/schedulers/cosine.py b/flaxdiff/schedulers/cosine.py index 00de7dd..e1f0bd9 100644 --- a/flaxdiff/schedulers/cosine.py +++ b/flaxdiff/schedulers/cosine.py @@ -34,7 +34,7 @@ def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarra noise_rates = jnp.sin((jnp.pi * steps) / (2 * self.max_timesteps)) return reshape_rates((signal_rates, noise_rates), shape=shape) - def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: alpha, sigma = self.get_rates(steps, shape=shape) return 1 / (1 + (alpha ** 2 / sigma ** 2)) \ No newline at end of file diff --git a/flaxdiff/schedulers/discrete.py b/flaxdiff/schedulers/discrete.py index ee1a01c..a7af264 100644 --- a/flaxdiff/schedulers/discrete.py +++ b/flaxdiff/schedulers/discrete.py @@ -47,7 +47,7 @@ def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.n def get_p2_weights(self, k, gamma): return (k + self.alpha_cumprod / (1 - self.alpha_cumprod)) ** -gamma - def get_weights(self, steps, shape=(-1, 1, 1, 1)): + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)): steps = jnp.int16(steps) return self.p2_loss_weights[steps].reshape(shape) diff --git a/flaxdiff/schedulers/flow.py b/flaxdiff/schedulers/flow.py new file mode 100644 index 0000000..7909fd2 --- /dev/null +++ b/flaxdiff/schedulers/flow.py @@ -0,0 +1,55 @@ +import math +import jax +import jax.numpy as jnp +from ..utils import RandomMarkovState +from .continuous import ContinuousNoiseScheduler +from .common import reshape_rates + + +def compute_resolution_shift(sequence_length, base_seq_len=256, max_seq_len=4096, + base_shift=0.5, max_shift=1.15) -> float: + """Flux-style resolution dependent timestep shift. + + Longer token sequences carry more redundancy, so the trajectory has to + spend more of its budget at high noise for the global structure to settle. + mu is interpolated linearly in sequence length and the shift is exp(mu). + """ + slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) + mu = base_shift + slope * (sequence_length - base_seq_len) + return math.exp(mu) + + +class FlowMatchingScheduler(ContinuousNoiseScheduler): + """Rectified flow / conditional flow matching on the linear path. + + x_t = (1 - t) * x_0 + t * epsilon for t in [0, 1], so alpha + sigma = 1 and + the model input needs no scaling. Timesteps are drawn logit-normal as in + SD3, which concentrates training on the middle of the trajectory where the + velocity is hardest to predict. + """ + def __init__(self, shift: float = 1.0, logit_mean: float = 0.0, logit_std: float = 1.0, + *args, **kwargs): + super().__init__(*args, **kwargs) + self.shift = shift + self.logit_mean = logit_mean + self.logit_std = logit_std + + def shift_timesteps(self, steps) -> jnp.ndarray: + return self.shift * steps / (1 + (self.shift - 1) * steps) + + def generate_timesteps(self, batch_size, state: RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]: + state, rng = state.get_random_key() + normal = jax.random.normal(rng, (batch_size,), dtype=self.dtype) + return jax.nn.sigmoid(normal * self.logit_std + self.logit_mean), state + + def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]: + t = self.shift_timesteps(jnp.asarray(steps, dtype=self.dtype)) + return reshape_rates((1 - t, t), shape=shape) + + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: + return jnp.ones_like(jnp.asarray(steps, dtype=self.dtype)).reshape(shape) + + def transform_inputs(self, x, steps) -> tuple[jnp.ndarray, jnp.ndarray]: + # The Fourier time embedding is tuned for discrete-style timesteps, so + # the [0, 1] flow time is rescaled into that range + return x, self.shift_timesteps(jnp.asarray(steps, dtype=self.dtype)) * 1000 diff --git a/flaxdiff/schedulers/karras.py b/flaxdiff/schedulers/karras.py index b9e09d5..078165e 100644 --- a/flaxdiff/schedulers/karras.py +++ b/flaxdiff/schedulers/karras.py @@ -17,7 +17,7 @@ def get_sigmas(self, steps) -> jnp.ndarray: sigmas = (self.max_inv_rho + ramp * (self.min_inv_rho - self.max_inv_rho)) ** self.rho return sigmas - def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: + def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: sigma = self.get_sigmas(steps) # EDM lambda(sigma) = (sigma^2 + sd^2) / (sigma * sd)^2, written in a # form that needs no epsilon guard (the old guard halved the weight at @@ -64,13 +64,21 @@ def get_sigmas(self, steps) -> jnp.ndarray: return self.sigmas[steps] class EDMNoiseScheduler(KarrasVENoiseScheduler): - def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, *args, **kwargs): + """Training sigmas drawn from exp(N(P_mean, P_std^2)). + + Defaults are EDM2's (Karras et al. 2024); EDM1's -1.2/1.2 concentrated too + much mass on low noise levels for larger models. Pass them explicitly to + reproduce an EDM1 run. + """ + def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, + P_mean=-0.4, P_std=1.0, *args, **kwargs): super().__init__(timesteps=timesteps, sigma_min=sigma_min, sigma_max=sigma_max, sigma_data=sigma_data, *args, **kwargs) + self.P_mean = P_mean + self.P_std = P_std - def get_sigmas(self, steps, std=1.2, mean=-1.2) -> jnp.ndarray: + def get_sigmas(self, steps) -> jnp.ndarray: space = steps / self.max_timesteps - # space = jax.scipy.special.erfinv(self.erf_sigma_min + steps * (self.erf_sigma_max - self.erf_sigma_min)) - return jnp.exp(space * std + mean) + return jnp.exp(space * self.P_std + self.P_mean) def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]: state, rng = state.get_random_key() diff --git a/tests/test_flow_matching.py b/tests/test_flow_matching.py new file mode 100644 index 0000000..b4efa6f --- /dev/null +++ b/tests/test_flow_matching.py @@ -0,0 +1,249 @@ +"""Flow matching on the linear (rectified flow) path. + +Covers the schedule invariants, the exact velocity round-trip, the claim that +the existing DDIM/Euler samplers already integrate the flow ODE, and a toy +end-to-end run proving the objective actually learns a distribution. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest +from flax import linen as nn + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.predictors import FlowMatchPredictionTransform, get_diffusion_preset +from flaxdiff.samplers.ddim import DDIMSampler +from flaxdiff.samplers.euler import EulerSampler +from flaxdiff.schedulers import FlowMatchingScheduler +from flaxdiff.schedulers.flow import compute_resolution_shift +from flaxdiff.schedulers.common import get_coeff_shapes_tuple +from flaxdiff.utils import RandomMarkovState + +STEPS = jnp.array([0.05, 0.3, 0.6, 0.95]) + + +def test_linear_path_rates(): + schedule = FlowMatchingScheduler() + alpha, sigma = schedule.get_rates(STEPS, shape=(-1,)) + assert jnp.allclose(alpha + sigma, 1.0, atol=1e-6) + assert jnp.allclose(sigma, STEPS, atol=1e-6) + # No input preconditioning on the linear path + assert FlowMatchPredictionTransform().get_input_scale((alpha, sigma)) == 1 + + +def test_endpoints_are_data_and_noise(rng): + schedule = FlowMatchingScheduler() + key0, key1 = jax.random.split(rng) + x0 = jax.random.normal(key0, (4, 8, 8, 3)) + noise = jax.random.normal(key1, (4, 8, 8, 3)) + assert jnp.allclose(schedule.add_noise(x0, noise, jnp.zeros((4,))), x0, atol=1e-6) + assert jnp.allclose(schedule.add_noise(x0, noise, jnp.ones((4,))), noise, atol=1e-6) + + +def test_timesteps_are_logit_normal(rng): + schedule = FlowMatchingScheduler(logit_mean=-0.3, logit_std=1.4) + steps, _ = schedule.generate_timesteps(50000, RandomMarkovState(rng)) + assert jnp.all((steps > 0) & (steps < 1)) + logits = jnp.log(steps) - jnp.log1p(-steps) + assert abs(float(jnp.mean(logits)) - (-0.3)) < 0.05 + assert abs(float(jnp.std(logits)) - 1.4) < 0.05 + + +def test_resolution_shift_is_identity_at_one(): + schedule = FlowMatchingScheduler(shift=1.0) + assert jnp.allclose(schedule.shift_timesteps(STEPS), STEPS, atol=1e-7) + + +@pytest.mark.parametrize("shift", [0.5, 1.0, 3.0]) +def test_resolution_shift_is_monotonic_and_fixes_endpoints(shift): + schedule = FlowMatchingScheduler(shift=shift) + t = jnp.linspace(0.0, 1.0, 101) + shifted = schedule.shift_timesteps(t) + assert jnp.all(jnp.diff(shifted) > 0) + assert float(shifted[0]) == pytest.approx(0.0, abs=1e-7) + assert float(shifted[-1]) == pytest.approx(1.0, abs=1e-7) + # A shift above 1 moves every interior timestep towards higher noise + if shift >= 1: + assert jnp.all(shifted >= t - 1e-7) + else: + assert jnp.all(shifted <= t + 1e-7) + + +def test_resolution_shift_grows_with_sequence_length(): + shifts = [compute_resolution_shift(n) for n in (256, 1024, 4096)] + assert shifts == sorted(shifts) + assert shifts[0] == pytest.approx(np.exp(0.5)) + assert shifts[-1] == pytest.approx(np.exp(1.15)) + + +def test_timestep_conditioning_is_scaled_to_the_embedding_range(): + schedule = FlowMatchingScheduler(shift=2.0) + x = jnp.zeros((4, 8, 8, 3)) + _, temb = schedule.transform_inputs(x, STEPS) + assert jnp.allclose(temb, schedule.shift_timesteps(STEPS) * 1000) + + +def test_velocity_roundtrip_is_exact(rng): + schedule = FlowMatchingScheduler() + transform = FlowMatchPredictionTransform() + key0, key1 = jax.random.split(rng) + x0 = jax.random.normal(key0, (4, 8, 8, 3)) + noise = jax.random.normal(key1, (4, 8, 8, 3)) + rates = schedule.get_rates(STEPS, get_coeff_shapes_tuple(x0)) + + xt, _, target = transform.forward_diffusion(x0, noise, rates) + assert jnp.allclose(target, noise - x0, atol=1e-6) + + recovered_x0, recovered_noise = transform.backward_diffusion(xt, target, rates) + assert jnp.max(jnp.abs(recovered_x0 - x0)) < 1e-5 + assert jnp.max(jnp.abs(recovered_noise - noise)) < 1e-5 + + +def test_preset_wires_flow_matching(): + for name in ('flow', 'flow_matching'): + train, sample, transform = get_diffusion_preset(name, shift=2.0) + assert isinstance(train, FlowMatchingScheduler) + assert isinstance(sample, FlowMatchingScheduler) + assert isinstance(transform, FlowMatchPredictionTransform) + assert train.shift == 2.0 and sample.shift == 2.0 + + +############################################################################################################ +# The existing samplers already integrate the flow ODE +############################################################################################################ + +class ConstantVelocity(nn.Module): + """Stands in for a trained model; take_next_step never calls it.""" + @nn.compact + def __call__(self, x, temb): + return x + + +def _flow_sampler(sampler_class): + schedule = FlowMatchingScheduler() + return sampler_class( + model=ConstantVelocity(), + noise_schedule=schedule, + model_output_transform=FlowMatchPredictionTransform(), + input_config=DiffusionInputConfig(sample_data_key="image", sample_data_shape=(8, 8, 3), conditions=[]), + guidance_scale=0.0, + ) + + +@pytest.mark.parametrize("sampler_class", [EulerSampler, DDIMSampler]) +def test_sampler_step_is_the_flow_euler_step(sampler_class, rng): + """x_{t+dt} = x_t + u * dt exactly, for the unmodified samplers.""" + sampler = _flow_sampler(sampler_class) + transform = FlowMatchPredictionTransform() + schedule = sampler.noise_schedule + + key0, key1 = jax.random.split(rng) + x_t = jax.random.normal(key0, (4, 8, 8, 3)) + velocity = jax.random.normal(key1, (4, 8, 8, 3)) + current_step = jnp.full((4,), 0.8) + next_step = jnp.full((4,), 0.6) + + rates = schedule.get_rates(current_step, get_coeff_shapes_tuple(x_t)) + x0, eps = transform.backward_diffusion(x_t, velocity, rates) + + stepped, _ = sampler.take_next_step( + current_samples=x_t, + reconstructed_samples=x0, + model_conditioning_inputs=(), + pred_noise=eps, + current_step=current_step, + state=RandomMarkovState(rng), + sample_model_fn=None, + next_step=next_step, + ) + expected = x_t + velocity * (0.6 - 0.8) + assert jnp.max(jnp.abs(stepped - expected)) < 1e-5 + + +############################################################################################################ +# Toy end-to-end: a two-mode gaussian mixture in the plane +############################################################################################################ + +MODE_CENTERS = jnp.array([[-0.5, -0.5], [0.5, 0.5]]) +MODE_STD = 0.08 + + +def sample_mixture(key, n): + """Two well-separated modes in the leading two channels; the third channel + carries no mode information, so the sampler's fixed channel count does not + turn this into a harder problem.""" + mode_key, noise_key = jax.random.split(key) + modes = jax.random.bernoulli(mode_key, 0.5, (n,)).astype(jnp.int32) + centers = jnp.concatenate([MODE_CENTERS[modes], jnp.zeros((n, 1))], axis=-1) + return (centers + MODE_STD * jax.random.normal(noise_key, (n, 3))).reshape(n, 1, 1, 3) + + +class ToyVelocityMLP(nn.Module): + features: int = 128 + + @nn.compact + def __call__(self, x, temb): + t = jnp.reshape(temb, (-1, 1)) / 1000.0 + freqs = jnp.arange(1, 5, dtype=jnp.float32) * jnp.pi + h = jnp.concatenate([x.reshape(x.shape[0], -1), t, jnp.sin(t * freqs), jnp.cos(t * freqs)], axis=-1) + h = nn.swish(nn.Dense(self.features)(h)) + h = nn.swish(nn.Dense(self.features)(h)) + return nn.Dense(3)(h).reshape(x.shape) + + +def test_flow_matching_learns_a_two_mode_mixture(): + schedule, sampling_schedule, transform = get_diffusion_preset('flow') + model = ToyVelocityMLP() + key = jax.random.PRNGKey(0) + params = model.init(key, jnp.zeros((1, 1, 1, 3)), jnp.zeros((1,))) + optimizer = optax.adam(3e-3) + opt_state = optimizer.init(params) + + def loss_fn(params, x0, noise, steps): + rates = schedule.get_rates(steps, get_coeff_shapes_tuple(x0)) + x_t, c_in, target = transform.forward_diffusion(x0, noise, rates) + preds = model.apply(params, *schedule.transform_inputs(x_t * c_in, steps)) + weights = schedule.get_weights(steps, get_coeff_shapes_tuple(x0)) + return jnp.mean(weights * (preds - target) ** 2) + + @jax.jit + def train_step(params, opt_state, rng_state): + rng_state, data_key = rng_state.get_random_key() + rng_state, noise_key = rng_state.get_random_key() + x0 = sample_mixture(data_key, 512) + noise = jax.random.normal(noise_key, x0.shape) + steps, rng_state = schedule.generate_timesteps(512, rng_state) + loss, grads = jax.value_and_grad(loss_fn)(params, x0, noise, steps) + updates, opt_state = optimizer.update(grads, opt_state, params) + return optax.apply_updates(params, updates), opt_state, rng_state, loss + + rng_state = RandomMarkovState(jax.random.PRNGKey(1)) + for _ in range(1500): + params, opt_state, rng_state, loss = train_step(params, opt_state, rng_state) + assert float(loss) < 1.0, "flow matching loss did not come down" + + sampler = EulerSampler( + model=model, + noise_schedule=sampling_schedule, + model_output_transform=transform, + input_config=DiffusionInputConfig(sample_data_key="image", sample_data_shape=(1, 1, 3), conditions=[]), + guidance_scale=0.0, + ) + samples = sampler.generate_samples( + params, num_samples=2048, resolution=1, diffusion_steps=64, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + ).reshape(-1, 3) + + assignment = samples[:, 0] > 0 + fraction = float(jnp.mean(assignment)) + assert 0.35 < fraction < 0.65, f"modes not balanced: {fraction:.2f}" + + for mode, center in enumerate(MODE_CENTERS): + members = samples[assignment == bool(mode)] + assert jnp.max(jnp.abs(jnp.mean(members[:, :2], axis=0) - center)) < 0.04 + assert abs(float(jnp.std(members[:, :2])) - MODE_STD) < 0.04 + # The third channel is a single zero-centred gaussian, not a mixture + assert abs(float(jnp.mean(samples[:, 2]))) < 0.04 + assert abs(float(jnp.std(samples[:, 2])) - MODE_STD) < 0.04 diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py index 44e05fc..67c745e 100644 --- a/tests/test_schedulers.py +++ b/tests/test_schedulers.py @@ -9,6 +9,11 @@ import jax.numpy as jnp import pytest +from flaxdiff.predictors import ( + EpsilonPredictionTransform, + VPredictionTransform, + get_diffusion_preset, +) from flaxdiff.schedulers import ( CosineNoiseScheduler, LinearNoiseSchedule, @@ -75,12 +80,84 @@ def test_karras_weights_at_sigma_min(): assert jnp.allclose(got, expected, rtol=1e-2) -def test_edm_lognormal_sigma_distribution(rng): - """EDM training sigmas must follow exp(N(-1.2, 1.2^2)) when timesteps=1.""" +@pytest.mark.parametrize("P_mean,P_std", [(-0.4, 1.0), (-1.2, 1.2)]) +def test_edm_lognormal_sigma_distribution(rng, P_mean, P_std): + """EDM training sigmas follow exp(N(P_mean, P_std^2)), defaulting to EDM2.""" from flaxdiff.utils import RandomMarkovState - schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5, P_mean=P_mean, P_std=P_std) steps, _ = schedule.generate_timesteps(20000, RandomMarkovState(rng)) log_sigma = jnp.log(schedule.get_sigmas(steps)) - assert abs(float(jnp.mean(log_sigma)) - (-1.2)) < 0.05 - assert abs(float(jnp.std(log_sigma)) - 1.2) < 0.05 + assert abs(float(jnp.mean(log_sigma)) - P_mean) < 0.05 + assert abs(float(jnp.std(log_sigma)) - P_std) < 0.05 + + +def test_edm_defaults_to_edm2_distribution(): + schedule = EDMNoiseScheduler(1) + assert (schedule.P_mean, schedule.P_std) == (-0.4, 1.0) + + +############################################################################################################ +# min-SNR-gamma loss weighting (Hang et al. 2023) +############################################################################################################ + +MIN_SNR_STEPS = jnp.array([10, 200, 400, 600, 800, 990]) + + +def make_min_snr_schedule(transform, gamma): + return CosineNoiseScheduler(1000, min_snr_gamma=gamma, prediction_transform=transform) + + +def test_min_snr_needs_a_parameterization(): + with pytest.raises(ValueError): + CosineNoiseScheduler(1000, min_snr_gamma=5.0) + + +def test_min_snr_weights_keep_the_requested_shape(): + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0) + assert schedule.get_weights(MIN_SNR_STEPS, shape=(-1, 1, 1, 1)).shape == (len(MIN_SNR_STEPS), 1, 1, 1) + assert schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)).shape == (len(MIN_SNR_STEPS),) + + +def test_min_snr_epsilon_weights_match_the_paper(): + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0) + snr = schedule.get_snr(MIN_SNR_STEPS) + expected = jnp.minimum(snr, 5.0) / snr + assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), expected, rtol=1e-5) + + +def test_min_snr_v_weights_match_the_paper(): + schedule = make_min_snr_schedule(VPredictionTransform(), 5.0) + snr = schedule.get_snr(MIN_SNR_STEPS) + expected = jnp.minimum(snr, 5.0) / (snr + 1) + assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), expected, rtol=1e-5) + + +def test_min_snr_weights_are_capped_and_non_increasing_in_snr(): + """The whole point: high-SNR (low noise) steps stop dominating the gradient.""" + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0) + # ascending timesteps are descending SNR, so weights must be non-decreasing + weights = schedule.get_weights(jnp.arange(1, 1000, 10), shape=(-1,)) + assert jnp.all(jnp.diff(weights) >= -1e-6) + assert jnp.all(weights <= 1.0 + 1e-6) + + +def test_min_snr_gamma_infinity_is_the_unweighted_case(): + schedule = make_min_snr_schedule(EpsilonPredictionTransform(), float('inf')) + assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), 1.0, atol=1e-6) + + +def test_min_snr_is_off_by_default(): + schedule = CosineNoiseScheduler(1000) + assert jnp.allclose( + schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), + schedule.get_schedule_weights(MIN_SNR_STEPS, shape=(-1,)), + ) + + +@pytest.mark.parametrize("name", ['cosine', 'edm', 'karras', 'flow']) +def test_preset_wires_min_snr_into_the_training_schedule_only(name): + train, sample, transform = get_diffusion_preset(name, min_snr_gamma=5.0) + assert train.min_snr_gamma == 5.0 + assert train.prediction_transform is transform + assert sample.min_snr_gamma is None From 8ad7ac9e1d33ac5c8cb5a71f0236b4e3a69f8aaf Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:52 -0400 Subject: [PATCH 2/5] feat(samplers): interval-limited classifier-free guidance Guidance costs diversity at high noise and buys nothing at low noise (Kynkaanniemi et al. 2024). guidance_start/guidance_stop bound the fraction of the trajectory over which it applies; outside the interval the scale drops to 1, which is exactly the plain conditional prediction. The default interval is the full range, so existing samplers are unchanged. Co-Authored-By: Claude Fable 5 --- flaxdiff/samplers/common.py | 27 ++++++-- tests/test_samplers.py | 119 +++++++++++++++++++++++++++++++++++- 2 files changed, 141 insertions(+), 5 deletions(-) diff --git a/flaxdiff/samplers/common.py b/flaxdiff/samplers/common.py index 910202c..7caab39 100644 --- a/flaxdiff/samplers/common.py +++ b/flaxdiff/samplers/common.py @@ -22,31 +22,49 @@ def __init__( model_output_transform: DiffusionPredictionTransform, input_config: DiffusionInputConfig, guidance_scale: float = 0.0, + guidance_start: float = 0.0, + guidance_stop: float = 1.0, autoencoder: AutoEncoder = None, ): """Initialize the diffusion sampler. - + Args: model: Neural network model params: Model parameters noise_schedule: Noise scheduler model_output_transform: Transform for model predictions guidance_scale: Scale for classifier-free guidance (0.0 means disabled) + guidance_start: Fraction of the trajectory after which guidance turns on + guidance_stop: Fraction of the trajectory after which guidance turns off autoencoder: Optional autoencoder for latent diffusion """ self.model = model self.noise_schedule = noise_schedule self.model_output_transform = model_output_transform self.guidance_scale = guidance_scale + self.guidance_start = guidance_start + self.guidance_stop = guidance_stop self.autoencoder = autoencoder self.input_config = input_config - + unconditionals = input_config.get_unconditionals() - + if self.guidance_scale > 0: # Classifier free guidance print("Using classifier-free guidance") + def guidance_at(t): + """Interval-limited guidance (Kynkaanniemi et al. 2024). + + Guidance hurts at high noise and buys nothing at low noise, so + outside [guidance_start, guidance_stop] the scale drops to 1, + which is exactly the plain conditional prediction. Progress runs + from 0 at pure noise to 1 at the clean sample. + """ + progress = 1.0 - t / self.noise_schedule.max_timesteps + inside = (progress >= guidance_start) & (progress <= guidance_stop) + return jnp.where(inside, guidance_scale, 1.0) + def sample_model(params, x_t, t, *conditioning_inputs): # Concatenate unconditional and conditional inputs x_t_cat = jnp.concatenate([x_t] * 2, axis=0) @@ -71,7 +89,8 @@ def sample_model(params, x_t, t, *conditioning_inputs): # Split model output into unconditional and conditional parts model_output_cond, model_output_uncond = jnp.split(model_output, 2, axis=0) - model_output = model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond) + scale = guidance_at(t).reshape(get_coeff_shapes_tuple(model_output_cond)) + model_output = model_output_uncond + scale * (model_output_cond - model_output_uncond) x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule) return x_0, eps, model_output diff --git a/tests/test_samplers.py b/tests/test_samplers.py index 9156c91..3fa8507 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -7,12 +7,15 @@ shape bugs in one place, with no training involved. """ +from dataclasses import dataclass + import jax import jax.numpy as jnp import pytest from flax import linen as nn -from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.inputs import ConditionalInputConfig, DiffusionInputConfig +from flaxdiff.inputs.encoders import ConditioningEncoder from flaxdiff.predictors import EpsilonPredictionTransform, KarrasPredictionTransform from flaxdiff.samplers.ddim import DDIMSampler from flaxdiff.samplers.ddpm import DDPMSampler, SimpleDDPMSampler @@ -186,3 +189,117 @@ def test_multistep_dpm_reentrant(): assert jnp.allclose(first, second, atol=1e-5) +############################################################################################################ +# Interval-limited classifier-free guidance (Kynkaanniemi et al. 2024) +############################################################################################################ + +@dataclass +class LabelEncoder(ConditioningEncoder): + """Smallest conditioning seam that still gives CFG a real conditional and + unconditional branch to interpolate between.""" + @property + def key(self): + return "label" + + def tokenize(self, data): + return jnp.asarray(data, dtype=jnp.float32).reshape(-1, 1) + + def encode_from_tokens(self, tokens): + return tokens + + def serialize(self): + return {} + + @staticmethod + def deserialize(serialized_config): + raise NotImplementedError + + +conditional_input_config = DiffusionInputConfig( + sample_data_key="image", + sample_data_shape=(8, 8, 3), + conditions=[ + ConditionalInputConfig( + encoder=LabelEncoder(model=None, tokenizer=None), + conditioning_data_key="label", + unconditional_input=0.0, + ) + ], +) + + +class ConditionalVPOracle(nn.Module): + """VPOracle offset by the label, so the guided output reads back the scale + that was actually applied.""" + schedule: CosineNoiseScheduler + + @nn.compact + def __call__(self, x, temb, label): + alpha, sigma = self.schedule.get_rates(temb, get_coeff_shapes_tuple(x)) + eps = x * sigma / (alpha**2 * DATA_STD**2 + sigma**2) + return eps + label.reshape(get_coeff_shapes_tuple(x)) + + +def make_guided_sampler(**kwargs): + schedule = CosineNoiseScheduler(1000) + model = ConditionalVPOracle(schedule=schedule) + sampler = DDIMSampler( + model=model, + noise_schedule=schedule, + model_output_transform=EpsilonPredictionTransform(), + input_config=conditional_input_config, + **kwargs, + ) + params = model.init( + jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,)), jnp.ones((1, 1)) + ) + return params, sampler + + +@pytest.mark.parametrize("progress,inside", [(0.1, False), (0.5, True), (0.9, False)]) +def test_interval_cfg_applies_only_inside_the_interval(progress, inside): + params, full = make_guided_sampler(guidance_scale=3.0) + _, interval = make_guided_sampler(guidance_scale=3.0, guidance_start=0.4, guidance_stop=0.6) + _, unguided = make_guided_sampler(guidance_scale=0.0) + + x_t = jax.random.normal(jax.random.PRNGKey(3), (4, 8, 8, 3)) + labels = jnp.full((4, 1), 0.7) + t = jnp.full((4,), (1.0 - progress) * 1000) + + def output(sampler): + return sampler.sample_model(params, x_t, t, labels)[2] + + matches_full = bool(jnp.allclose(output(interval), output(full), atol=1e-5)) + matches_unguided = bool(jnp.allclose(output(interval), output(unguided), atol=1e-5)) + assert matches_full is inside + assert matches_unguided is not inside + + +def test_interval_cfg_defaults_to_the_full_range(): + params, default = make_guided_sampler(guidance_scale=3.0) + _, explicit = make_guided_sampler(guidance_scale=3.0, guidance_start=0.0, guidance_stop=1.0) + + x_t = jax.random.normal(jax.random.PRNGKey(3), (4, 8, 8, 3)) + labels = jnp.full((4, 1), 0.7) + t = jnp.full((4,), 500.0) + assert jnp.allclose( + default.sample_model(params, x_t, t, labels)[2], + explicit.sample_model(params, x_t, t, labels)[2], + atol=1e-6, + ) + + +def test_empty_guidance_interval_generates_the_unguided_samples(): + params, empty = make_guided_sampler(guidance_scale=3.0, guidance_start=0.9, guidance_stop=0.1) + _, unguided = make_guided_sampler(guidance_scale=0.0) + + def run(sampler): + return sampler.generate_samples( + params, num_samples=16, resolution=8, diffusion_steps=25, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + model_conditioning_inputs=(jnp.full((16, 1), 0.7),), + ) + + assert jnp.allclose(run(empty), run(unguided), atol=1e-5) + + From aadc65988e3cda9c3b8a5dc943a36ddb9f3deab2 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:52 -0400 Subject: [PATCH 3/5] feat(metrics): FID on the vendored InceptionV3 The jax-fid port in metrics/inception.py has been sitting unused. Wire it into an EvaluationMetric the same way metrics/images.py builds the CLIP ones, sharing a single cached copy of the weights. Per-batch FID is noisy and only meaningful as a trend across checkpoints, which the docstring says outright. Co-Authored-By: Claude Fable 5 --- flaxdiff/metrics/fid.py | 75 +++++++++++++++++++++++++++++++++++++++++ tests/test_metrics.py | 53 +++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+) create mode 100644 flaxdiff/metrics/fid.py create mode 100644 tests/test_metrics.py diff --git a/flaxdiff/metrics/fid.py b/flaxdiff/metrics/fid.py new file mode 100644 index 0000000..d11d13b --- /dev/null +++ b/flaxdiff/metrics/fid.py @@ -0,0 +1,75 @@ +from .common import EvaluationMetric +import jax +import jax.numpy as jnp +import numpy as np + + +# The FID InceptionV3 is ~90MB of weights; one copy is enough for every metric +# built from this module. +_inception_cache: dict = {} + + +def _get_inception(): + """Cached (model, params) for the pool3 feature extractor.""" + if 'inception' not in _inception_cache: + from .inception import InceptionV3 + print("[metrics] Loading InceptionV3 FID weights (cached for reuse)...") + model = InceptionV3(pretrained=True) + params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 299, 299, 3))) + _inception_cache['inception'] = (model, params) + return _inception_cache['inception'] + + +def frechet_distance(mu_a, sigma_a, mu_b, sigma_b, eps=1e-6) -> float: + """Frechet distance between two multivariate gaussians. + + Runs on the host through scipy: the matrix square root of the covariance + product has no jax equivalent, and FID is computed once per validation + batch so the transfer is irrelevant. + """ + from scipy import linalg + + mu_a, mu_b = np.atleast_1d(mu_a), np.atleast_1d(mu_b) + sigma_a, sigma_b = np.atleast_2d(sigma_a), np.atleast_2d(sigma_b) + + covmean, _ = linalg.sqrtm(sigma_a.dot(sigma_b), disp=False) + if not np.isfinite(covmean).all(): + # Singular product covariance, nudge the diagonal as in the reference + # implementations rather than returning a nan + offset = np.eye(sigma_a.shape[0]) * eps + covmean = linalg.sqrtm((sigma_a + offset).dot(sigma_b + offset)) + if np.iscomplexobj(covmean): + covmean = covmean.real + + diff = mu_a - mu_b + return float(diff.dot(diff) + np.trace(sigma_a) + np.trace(sigma_b) - 2 * np.trace(covmean)) + + +def _gaussian_stats(activations): + activations = np.asarray(activations, dtype=np.float64) + return activations.mean(axis=0), np.cov(activations, rowvar=False) + + +def get_fid_metric(): + """FID between the generated batch and the real batch, lower is better. + + Per-batch FID is noisy at typical validation batch sizes and is only + meaningful as a relative trend across checkpoints, not as a headline number + comparable to published FID-50k. + """ + model, params = _get_inception() + + @jax.jit + def activations(images): + # Inception wants [-1, 1] at 299x299; pool3 output is [B, 1, 1, 2048] + resized = jax.image.resize(images, (images.shape[0], 299, 299, 3), method='bilinear') + features = model.apply(params, resized, train=False) + return features.reshape(features.shape[0], -1) + + def fid_metric(generated: jnp.ndarray, batch): + original = (jnp.asarray(batch['image'], dtype=jnp.float32) - 127.5) / 127.5 + mu_gen, sigma_gen = _gaussian_stats(activations(generated)) + mu_real, sigma_real = _gaussian_stats(activations(original)) + return frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real) + + return EvaluationMetric(function=fid_metric, name='fid') diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..3801da9 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,53 @@ +"""Evaluation metric tests. + +The Frechet distance itself is checked against closed forms that need no +weights; the end-to-end InceptionV3 path downloads the FID checkpoint and is +network-marked. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from flaxdiff.metrics.common import EvaluationMetric +from flaxdiff.metrics.fid import frechet_distance, get_fid_metric + + +def test_frechet_distance_of_a_distribution_with_itself_is_zero(rng): + features = np.asarray(jax.random.normal(rng, (256, 16))) + mu, sigma = features.mean(axis=0), np.cov(features, rowvar=False) + assert frechet_distance(mu, sigma, mu, sigma) == pytest.approx(0.0, abs=1e-6) + + +def test_frechet_distance_of_shifted_gaussians_is_the_squared_mean_gap(): + sigma = np.eye(8) + mu_a = np.zeros(8) + mu_b = np.full(8, 0.5) + assert frechet_distance(mu_a, sigma, mu_b, sigma) == pytest.approx(8 * 0.25, abs=1e-6) + + +def test_frechet_distance_grows_with_covariance_mismatch(): + mu = np.zeros(8) + identity = np.eye(8) + near = frechet_distance(mu, identity, mu, identity * 1.5) + far = frechet_distance(mu, identity, mu, identity * 4.0) + assert 0 < near < far + + +@pytest.mark.network +def test_fid_metric_scores_real_images_better_than_noise(rng): + metric = get_fid_metric() + assert isinstance(metric, EvaluationMetric) + assert metric.name == 'fid' and metric.higher_is_better is False + + key_real, key_noise = jax.random.split(rng) + real = jax.random.randint(key_real, (8, 64, 64, 3), 0, 256, dtype=jnp.int32).astype(jnp.uint8) + batch = {'image': real} + + # Generated samples live in [-1, 1]; the same images should score far + # closer to the batch than unrelated noise does + matching = metric.function((jnp.asarray(real, jnp.float32) - 127.5) / 127.5, batch) + unrelated = metric.function(jax.random.normal(key_noise, (8, 64, 64, 3)), batch) + assert np.isfinite(matching) and np.isfinite(unrelated) + assert matching < unrelated From 29d1a2debb2e4731caa404c15aca5c9fce70dc01 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:58 -0400 Subject: [PATCH 4/5] feat(autoencoder): per-dataset latent normalization on the shared seam Latents are normalized as (z - latent_shift) * latent_scale and inverted on decode, the SD3 convention, so a caller can hand in a dataset's own mean and 1/std instead of relying on the VAE's single fixed constant. The SD VAE's scaling_factor now rides on that same seam rather than being applied inside its jitted frame functions, so there is one normalization path and the defaults reproduce the previous behavior exactly. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/autoencoder/autoencoder.py | 22 ++++++-- flaxdiff/models/autoencoder/diffusers.py | 19 ++++--- tests/test_vae.py | 58 ++++++++++++++++++++-- 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/flaxdiff/models/autoencoder/autoencoder.py b/flaxdiff/models/autoencoder/autoencoder.py index a307da1..233eaec 100644 --- a/flaxdiff/models/autoencoder/autoencoder.py +++ b/flaxdiff/models/autoencoder/autoencoder.py @@ -10,11 +10,19 @@ @dataclass class AutoEncoder(ABC): """Base class for autoencoder models with video support. - + This class defines the interface for autoencoders and provides video handling functionality, allowing child classes to focus on implementing the core encoding/decoding for individual frames. + + Latents are normalized as (z - latent_shift) * latent_scale on the way out + and inverted on the way in, the SD3 convention. The defaults are the + identity; set them to the dataset's own latent mean and 1/std so the + diffusion model sees roughly unit-variance, zero-mean inputs. """ + latent_shift: float = 0.0 + latent_scale: float = 1.0 + @abstractmethod def __encode__(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray: """Abstract method for encoding a batch of images. @@ -72,13 +80,15 @@ def encode(self, x: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwa # Encode all frames latent = self.__encode__(x_reshaped, key=key, **kwargs) - + # Reshape back to include temporal dimension [B, T, h, w, c] latent_shape = latent.shape - return latent.reshape(batch_size, seq_len, *latent_shape[1:]) + latent = latent.reshape(batch_size, seq_len, *latent_shape[1:]) else: # Standard image processing - return self.__encode__(x, key=key, **kwargs) + latent = self.__encode__(x, key=key, **kwargs) + + return (latent - self.latent_shift) * self.latent_scale def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwargs) -> jnp.ndarray: """Decode latent representations, with special handling for video data. @@ -95,9 +105,11 @@ def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwa Returns: Decoded output with the same batch and temporal dimensions as input """ + z = z / self.latent_scale + self.latent_shift + # Check for video data (5D tensor) is_video = len(z.shape) == 5 - + if is_video: # Extract dimensions for reshaping batch_size, seq_len, height, width, channels = z.shape diff --git a/flaxdiff/models/autoencoder/diffusers.py b/flaxdiff/models/autoencoder/diffusers.py index 43ee03c..ab8015b 100644 --- a/flaxdiff/models/autoencoder/diffusers.py +++ b/flaxdiff/models/autoencoder/diffusers.py @@ -12,7 +12,8 @@ """ class StableDiffusionVAE(AutoEncoder): - def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16): + def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16, + latent_shift=None, latent_scale=None): pretrained = load_pretrained_vae(modelname) config = pretrained["config"] @@ -61,10 +62,14 @@ def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=dtype, ) - # Older VAE configs predate the scaling_factor key; 0.18215 is the SD default - scaling_factor = config.get("scaling_factor", 0.18215) - print(f"Scaling factor: {scaling_factor}") - + # The VAE's own latent normalization rides on the AutoEncoder seam, so a + # caller can override it with per-dataset statistics without a second + # scaling path. Older configs predate these keys; 0.0 and 0.18215 are + # the SD defaults. + self.latent_shift = config.get("shift_factor", 0.0) if latent_shift is None else latent_shift + self.latent_scale = config.get("scaling_factor", 0.18215) if latent_scale is None else latent_scale + print(f"Latent shift: {self.latent_shift}, latent scale: {self.latent_scale}") + def encode_single_frame(images, rngkey: jax.random.PRNGKey = None): latents = enc.apply({"params": params['encoder']}, images, deterministic=True) latents = quant_conv.apply({"params": params['quant_conv']}, latents) @@ -75,11 +80,9 @@ def encode_single_frame(images, rngkey: jax.random.PRNGKey = None): latents = mean + std * jax.random.normal(rngkey, mean.shape, dtype=mean.dtype) else: latents, _ = jnp.split(latents, 2, axis=-1) - latents *= scaling_factor return latents - + def decode_single_frame(latents): - latents = (1.0 / scaling_factor) * latents latents = post_quant_conv.apply({"params": params['post_quant_conv']}, latents) return dec.apply({"params": params['decoder']}, latents) diff --git a/tests/test_vae.py b/tests/test_vae.py index 96904ec..3b648e0 100644 --- a/tests/test_vae.py +++ b/tests/test_vae.py @@ -1,7 +1,7 @@ -"""Vendored Stable Diffusion VAE tests. +"""Vendored Stable Diffusion VAE tests, plus the latent normalization seam. -Marked as network tests: they download the pretrained weights from the -HuggingFace Hub on first run. Excluded in CI (-m "not network"). +Everything touching the pretrained weights is network-marked: they download +from the HuggingFace Hub on first run. Excluded in CI (-m "not network"). """ import jax @@ -9,7 +9,47 @@ import numpy as np import pytest -pytestmark = pytest.mark.network +from flaxdiff.models.autoencoder import AutoEncoder + + +class IdentityAutoEncoder(AutoEncoder): + """Latents are the input, so only the normalization seam is under test.""" + def __encode__(self, x, **kwargs): + return x + + def __decode__(self, z, **kwargs): + return z + + def serialize(self): + return {} + + +def test_latent_normalization_defaults_to_the_identity(rng): + autoencoder = IdentityAutoEncoder() + x = jax.random.normal(rng, (2, 8, 8, 4)) + assert jnp.allclose(autoencoder.encode(x), x) + assert jnp.allclose(autoencoder.decode(x), x) + + +@pytest.mark.parametrize("shape", [(2, 8, 8, 4), (2, 3, 8, 8, 4)]) +def test_latent_normalization_shifts_and_scales_roundtrip(rng, shape): + """SD3-style shift+scale: latents come out centred and rescaled, and + decoding inverts it exactly, for images and for video.""" + autoencoder = IdentityAutoEncoder(latent_shift=0.25, latent_scale=4.0) + x = jax.random.normal(rng, shape) + latent = autoencoder.encode(x) + assert jnp.allclose(latent, (x - 0.25) * 4.0, atol=1e-6) + assert jnp.allclose(autoencoder.decode(latent), x, atol=1e-5) + + +def test_latent_normalization_whitens_a_known_distribution(rng): + """The point of per-dataset stats: the diffusion model sees zero mean and + unit variance instead of whatever the encoder happens to produce.""" + x = 3.0 + 5.0 * jax.random.normal(rng, (4096, 1, 1, 4)) + autoencoder = IdentityAutoEncoder(latent_shift=float(jnp.mean(x)), latent_scale=1.0 / float(jnp.std(x))) + latent = autoencoder.encode(x) + assert abs(float(jnp.mean(latent))) < 1e-4 + assert abs(float(jnp.std(latent)) - 1.0) < 1e-4 @pytest.fixture(scope="module") @@ -18,11 +58,21 @@ def vae(): return StableDiffusionVAE(dtype=jnp.float32) +@pytest.mark.network def test_vae_shapes(vae): assert vae.downscale_factor == 8 assert vae.latent_channels == 4 +@pytest.mark.network +def test_vae_uses_the_latent_normalization_seam(vae): + """The SD scaling factor rides on the shared seam, so there is one + normalization path a caller can override with dataset statistics.""" + assert vae.latent_scale == pytest.approx(0.18215) + assert vae.latent_shift == 0.0 + + +@pytest.mark.network def test_vae_roundtrip_reconstructs(vae, rng): # A smooth image should survive the encode/decode roundtrip well ramp = jnp.linspace(-0.8, 0.8, 64) From 40f69d7149b84817338d6bc1a5b74b86849940d2 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:45:58 -0400 Subject: [PATCH 5/5] feat(training): expose the flow matching, min-SNR and FID options --noise_schedule gains flow/flow_matching, plus --min_snr_gamma, --flow_shift, and 'fid' as a validation metric. Co-Authored-By: Claude Fable 5 --- training.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/training.py b/training.py index 9ede2e7..b42bb43 100644 --- a/training.py +++ b/training.py @@ -101,7 +101,11 @@ def boolean_string(s): default='/home/mrwhite0racle/gcs_mount', help="Dataset location path") parser.add_argument('--noise_schedule', type=str, default='edm', - choices=['cosine', 'karras', 'edm'], help='Noise schedule') + choices=['cosine', 'karras', 'edm', 'flow', 'flow_matching'], help='Noise schedule') +parser.add_argument('--min_snr_gamma', type=float, default=None, + help='min-SNR-gamma loss weighting (Hang et al. 2023). 5.0 is the paper default; unset keeps the schedule own weighting.') +parser.add_argument('--flow_shift', type=float, default=1.0, + help='Resolution shift for the flow matching schedule. See flaxdiff.schedulers.flow.compute_resolution_shift.') # Any name from flaxdiff.models.registry, optionally with +2d/+hilbert/+zigzag # suffixes; validated by build_model against the registry itself @@ -469,6 +473,10 @@ def main(args): from flaxdiff.metrics.images import get_clip_score_metric print("Using CLIPScore (val/clip_score, higher is better) for validation") eval_metrics.append(get_clip_score_metric()) + if 'fid' in args.val_metrics: + from flaxdiff.metrics.fid import get_fid_metric + print("Using per-batch FID (val/fid) for validation") + eval_metrics.append(get_fid_metric()) CONFIG = { "model": model_config, @@ -493,7 +501,9 @@ def main(args): batches = batches if args.steps_per_epoch is None else args.steps_per_epoch - train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset(args.noise_schedule) + train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset( + args.noise_schedule, shift=args.flow_shift, min_snr_gamma=args.min_snr_gamma, + ) if args.experiment_name is not None: experiment_name = args.experiment_name