Phase 1: correctness — every audited defect fixed, 16 xfails flipped

README.md

@@ -210,23 +210,21 @@ final_state = trainer.fit(data, batches, epochs=2000, sampler_class=EulerAncestr
 Here is a simplified example for generating images using a trained model:
 
 ```python
-from flaxdiff.samplers import DiffusionSampler
-
-class EulerSampler(DiffusionSampler):
-    def take_next_step(self, current_samples, reconstructed_samples, pred_noise, current_step, state, next_step=None):
-        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
-        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
-        dt = next_sigma - current_sigma
-        x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / dt
-        dx = (current_samples - x_0_coeff * reconstructed_samples) / current_sigma
-        next_samples = current_samples + dx * dt
-        return next_samples, state
-
-# Create sampler
-sampler = EulerSampler(trainer.model, trainer.state.ema_params, edm_schedule, model_output_transform=trainer.model_output_transform)
-
-# Generate images
-samples = sampler.generate_images(num_images=64, diffusion_steps=100, start_step=1000, end_step=0)
+from flaxdiff.samplers.euler import EulerSampler
+
+sampler = EulerSampler(
+    trainer.model,
+    noise_schedule=karas_ve_schedule,
+    model_output_transform=trainer.model_output_transform,
+    input_config=input_config,
+)
+
+samples = sampler.generate_samples(
+    params=trainer.state.ema_params,
+    num_samples=64,
+    resolution=128,
+    diffusion_steps=100,
+)
 plotImages(samples, dpi=300)
 ```
 

flaxdiff/inference/pipeline.py

@@ -224,7 +224,6 @@ def generate_samples(
         diffusion_steps: int = 50,
         guidance_scale: float = 1.0,
         sampler_class=EulerAncestralSampler,
-        timestep_spacing: str = 'linear',
         seed: Optional[int] = None,
         start_step: Optional[int] = None,
         end_step: int = 0,
@@ -241,8 +240,6 @@ def generate_samples(
             guidance_scale=guidance_scale,
             sampler_class=sampler_class,
         )
-        if hasattr(sampler, 'timestep_spacing'):
-            sampler.timestep_spacing = timestep_spacing
         print(f"Generating samples: steps={diffusion_steps}, num_samples={num_samples}, guidance={guidance_scale}")
 
         if use_best_params:

flaxdiff/inference/utils.py

@@ -1,14 +1,7 @@
 import jax
 import jax.numpy as jnp
 import json
-from flaxdiff.schedulers import (
-    CosineNoiseScheduler,
-    KarrasVENoiseScheduler,
-)
-from flaxdiff.predictors import (
-    VPredictionTransform,
-    KarrasPredictionTransform,
-)
+from flaxdiff.predictors import get_diffusion_preset
 from flaxdiff.models.common import kernel_init
 from flaxdiff.models.simple_unet import Unet
 from flaxdiff.models.simple_vit import UViT
@@ -147,9 +140,15 @@ def map_nested_config(config):
                     new_config[key] = ACTIVATION_MAP[value]
                 elif value == 'None':
                     new_config[key] = None
-                elif '.'in value:
-                    # Ignore any other string that contains a dot
-                    print(f"Ignoring key {key} with value {value} as it contains a dot.")
+                elif value.startswith('jax.') or value.startswith('jax._src.'):
+                    # Resolve function paths like 'jax.nn.mish' by attribute walk;
+                    # jax._src paths leak from old configs that stored the live
+                    # function object instead of its name
+                    attr_path = value.replace('jax._src.nn.functions', 'jax.nn')
+                    obj = jax
+                    for part in attr_path.split('.')[1:]:
+                        obj = getattr(obj, part)
+                    new_config[key] = obj
                 else:
                     new_config[key] = value
             else:
@@ -229,24 +228,25 @@ def map_nested_config(config):
     if not model_class:
         raise ValueError(f"Unknown architecture: {architecture}. Supported architectures: {', '.join(MODEL_CLASSES.keys())}")
 
+    # Drop config keys the model no longer has fields for (older runs logged
+    # since-removed flags like use_flash_attention)
+    import dataclasses
+    valid_fields = {f.name for f in dataclasses.fields(model_class)}
+    dropped = sorted(set(model_kwargs) - valid_fields)
+    if dropped:
+        print(f"Dropping config keys not accepted by {model_class.__name__}: {dropped}")
+    model_kwargs = {k: v for k, v in model_kwargs.items() if k in valid_fields}
+
     # Instantiate the model
     model = model_class(**model_kwargs)
 
     # If using diffusers UNet, wrap it for consistent interface
     if 'diffusers' in architecture:
         model = BCHWModelWrapper(model)
 
-    # Create noise scheduler based on configuration
+    # Same preset as training, so the sampling convention always matches
     noise_schedule_type = conf.get('noise_schedule', conf.get('arguments', {}).get('noise_schedule', 'edm'))
-    if noise_schedule_type in ['edm', 'karras']:
-        # For both EDM and karras, we use the karras scheduler for inference
-        noise_schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
-        prediction_transform = KarrasPredictionTransform(sigma_data=noise_schedule.sigma_data)
-    elif noise_schedule_type == 'cosine':
-        noise_schedule = CosineNoiseScheduler(1000, beta_end=1)
-        prediction_transform = VPredictionTransform()
-    else:
-        raise ValueError(f"Unknown noise schedule: {noise_schedule_type}")
+    _, noise_schedule, prediction_transform = get_diffusion_preset(noise_schedule_type)
 
     # Prepare return dictionary with all components
     result = {

flaxdiff/models/attention.py

@@ -11,108 +11,6 @@
 import functools
 import math
 from .common import kernel_init
-try:
-    import jax.experimental.pallas.ops.tpu.flash_attention
-except (ImportError, ModuleNotFoundError):
-    pass
-
-class EfficientAttention(nn.Module):
-    """
-    Based on the pallas attention implementation.
-    """
-    query_dim: int
-    heads: int = 4
-    dim_head: int = 64
-    dtype: Optional[Dtype] = None
-    precision: PrecisionLike = None
-    use_bias: bool = True
-    # kernel_init: Callable = kernel_init(1.0)
-    force_fp32_for_softmax: bool = True
-
-    def setup(self):
-        inner_dim = self.dim_head * self.heads
-        # Weights were exported with old names {to_q, to_k, to_v, to_out}
-        dense = functools.partial(
-            nn.Dense,
-            self.heads * self.dim_head,
-            precision=self.precision, 
-            use_bias=self.use_bias, 
-            # kernel_init=self.kernel_init, 
-            dtype=self.dtype
-        )
-        self.query = dense(name="to_q")
-        self.key = dense(name="to_k")
-        self.value = dense(name="to_v")
-
-        self.proj_attn = nn.DenseGeneral(
-            self.query_dim, 
-            use_bias=False, 
-            precision=self.precision, 
-            # kernel_init=self.kernel_init,
-            dtype=self.dtype,
-            name="to_out_0"
-        )
-        # self.attnfn = make_fast_generalized_attention(qkv_dim=inner_dim, lax_scan_unroll=16)
-
-    def _reshape_tensor_to_head_dim(self, tensor):
-        batch_size, _, seq_len, dim = tensor.shape
-        head_size = self.heads
-        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
-        tensor = jnp.transpose(tensor, (0, 2, 1, 3))
-        return tensor
-
-    def _reshape_tensor_from_head_dim(self, tensor):
-        batch_size, _, seq_len, dim = tensor.shape
-        head_size = self.heads
-        tensor = jnp.transpose(tensor, (0, 2, 1, 3))
-        tensor = tensor.reshape(batch_size, 1, seq_len, dim * head_size)
-        return tensor
-
-    @nn.compact
-    def __call__(self, x:jax.Array, context=None):
-        # print(x.shape)
-        # x has shape [B, H * W, C]
-        context = x if context is None else context
-
-        orig_x_shape = x.shape
-        if len(x.shape) == 4:
-            B, H, W, C = x.shape
-            x = x.reshape((B, 1, H * W, C))
-        else:
-            B, SEQ, C = x.shape
-            x = x.reshape((B, 1, SEQ, C))
-
-        if len(context.shape) == 4:
-            B, _H, _W, _C = context.shape
-            context = context.reshape((B, 1, _H * _W, _C))
-        else:
-            B, SEQ, _C = context.shape
-            context = context.reshape((B, 1, SEQ, _C))
-
-        query = self.query(x)
-        key = self.key(context)
-        value = self.value(context)
-
-        query = self._reshape_tensor_to_head_dim(query)
-        key = self._reshape_tensor_to_head_dim(key)
-        value = self._reshape_tensor_to_head_dim(value)
-
-        hidden_states = jax.experimental.pallas.ops.tpu.flash_attention.flash_attention(
-            query, key, value, None
-        )
-
-        hidden_states = self._reshape_tensor_from_head_dim(hidden_states)
-
-
-        # hidden_states = nn.dot_product_attention(
-        #     query, key, value, dtype=self.dtype, broadcast_dropout=False, dropout_rng=None, precision=self.precision
-        # )
-
-        proj = self.proj_attn(hidden_states)
-
-        proj = proj.reshape(orig_x_shape)
-
-        return proj
 
 class NormalAttention(nn.Module):
     """
@@ -246,17 +144,13 @@ class BasicTransformerBlock(nn.Module):
     precision: PrecisionLike = None
     use_bias: bool = True
     # kernel_init: Callable = kernel_init(1.0)
-    use_flash_attention:bool = False
     use_cross_only:bool = False
     only_pure_attention:bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-4
 
     def setup(self):
-        if self.use_flash_attention:
-            attenBlock = EfficientAttention
-        else:
-            attenBlock = NormalAttention
+        attenBlock = NormalAttention
 
         self.attention1 = attenBlock(
          query_dim=self.query_dim,
@@ -281,7 +175,7 @@ def setup(self):
             force_fp32_for_softmax=self.force_fp32_for_softmax
         )
 
-        self.ff = FlaxFeedForward(dim=self.query_dim)
+        self.ff = FlaxFeedForward(dim=self.query_dim, dtype=self.dtype, precision=self.precision)
         self.norm1 = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)
         self.norm2 = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)
         self.norm3 = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)
@@ -309,7 +203,6 @@ class TransformerBlock(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     use_projection: bool = False
-    use_flash_attention:bool = False
     use_self_and_cross:bool = True
     only_pure_attention:bool = False
     force_fp32_for_softmax: bool = True
@@ -351,7 +244,6 @@ def __call__(self, x, context=None):
             precision=self.precision,
             use_bias=False,
             dtype=self.dtype,
-            use_flash_attention=self.use_flash_attention,
             use_cross_only=(not self.use_self_and_cross),
             only_pure_attention=self.only_pure_attention,
             force_fp32_for_softmax=self.force_fp32_for_softmax,

flaxdiff/models/common.py

@@ -1,5 +1,6 @@
 import jax.numpy as jnp
 import jax
+import numpy as np
 from flax import linen as nn
 from typing import Optional, Any, Callable, Sequence, Union
 from flax.typing import Dtype, PrecisionLike
@@ -99,7 +100,10 @@ class FourierEmbedding(nn.Module):
     scale:int = 16
 
     def setup(self):
-        self.freqs = jax.random.normal(jax.random.PRNGKey(42), (self.features // 2, ), dtype=jnp.float32) * self.scale
+        # Fixed frequencies via numpy so they are identical across jax versions
+        # (jax 0.5.0 changed the default PRNG and silently altered these)
+        freqs = np.random.RandomState(42).normal(size=(self.features // 2,))
+        self.freqs = jnp.asarray(freqs, dtype=jnp.float32) * self.scale
 
     def __call__(self, x):
         x = jax.lax.convert_element_type(x, jnp.float32)
@@ -176,7 +180,6 @@ def setup(self):
                 kernel_size=self.kernel_size,
                 strides=self.strides,
                 padding="SAME",
-                param_dtype=self.dtype,
                 dtype=self.dtype,
                 precision=self.precision
             )

flaxdiff/models/general.py

@@ -6,7 +6,7 @@ class BCHWModelWrapper(nn.Module):
     model: nn.Module
 
     @nn.compact
-    def __call__(self, x, temb, textcontext):
+    def __call__(self, x, temb, textcontext, train: bool = False):
         # Reshape the input to BCHW format from BHWC
         x = jnp.transpose(x, (0, 3, 1, 2))
         # Pass the input through the UNet model