Phase 4: flow matching, min-SNR, EDM2 defaults, interval CFG, FID

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')

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

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)
 

flaxdiff/predictors/__init__.py

@@ -7,6 +7,7 @@
     CosineNoiseScheduler,
     KarrasVENoiseScheduler,
     EDMNoiseScheduler,
+    FlowMatchingScheduler,
 )
 
 ############################################################################################################
@@ -39,17 +40,28 @@ 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
         epsilon = preds
         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,25 +150,40 @@ 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.
 
     The single source of truth for which schedule pairs with which
     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

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

flaxdiff/schedulers/__init__.py

@@ -3,4 +3,5 @@
 from .cosine import *
 from .linear import *
 from .sqrt import *
-from .karras import *
+from .karras import *
+from .flow import *