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/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/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/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/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_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 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) + + 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 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) 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