diff --git a/README.md b/README.md index 3eee41c..f3d6beb 100644 --- a/README.md +++ b/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) ``` diff --git a/flaxdiff/inference/pipeline.py b/flaxdiff/inference/pipeline.py index 27c7e15..acd72aa 100644 --- a/flaxdiff/inference/pipeline.py +++ b/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: diff --git a/flaxdiff/inference/utils.py b/flaxdiff/inference/utils.py index cdf598d..e45036c 100644 --- a/flaxdiff/inference/utils.py +++ b/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,6 +228,15 @@ 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) @@ -236,17 +244,9 @@ def map_nested_config(config): 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 = { diff --git a/flaxdiff/models/attention.py b/flaxdiff/models/attention.py index a7f761e..8106bf5 100644 --- a/flaxdiff/models/attention.py +++ b/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, diff --git a/flaxdiff/models/common.py b/flaxdiff/models/common.py index 2b5b84f..b732eae 100644 --- a/flaxdiff/models/common.py +++ b/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 ) diff --git a/flaxdiff/models/general.py b/flaxdiff/models/general.py index 19d2f71..cf7b09c 100644 --- a/flaxdiff/models/general.py +++ b/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 diff --git a/flaxdiff/models/simple_dit.py b/flaxdiff/models/simple_dit.py index d7624f0..70fa752 100644 --- a/flaxdiff/models/simple_dit.py +++ b/flaxdiff/models/simple_dit.py @@ -28,7 +28,6 @@ class DiTBlock(nn.Module): dropout_rate: float = 0.0 dtype: Optional[Dtype] = None precision: PrecisionLike = None - use_flash_attention: bool = False # Keep placeholder force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 use_gating: bool = True # Add flag to easily disable gating @@ -60,8 +59,10 @@ def setup(self): nn.Dense(features=self.features, dtype=self.dtype, precision=self.precision) ]) + self.dropout = nn.Dropout(rate=self.dropout_rate) + @nn.compact - def __call__(self, x, conditioning, freqs_cis): + def __call__(self, x, conditioning, freqs_cis, train: bool = False): # Get scale/shift/gate parameters # Shape: [B, 1, 6*F] -> split into 6 of [B, 1, F] scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( @@ -74,6 +75,7 @@ def __call__(self, x, conditioning, freqs_cis): # Modulate after norm x_attn_modulated = norm_x_attn * (1 + scale_attn) + shift_attn attn_output = self.attention(x_attn_modulated, context=None, freqs_cis=freqs_cis) + attn_output = self.dropout(attn_output, deterministic=not train) if self.use_gating: x = residual + gate_attn * attn_output @@ -86,6 +88,7 @@ def __call__(self, x, conditioning, freqs_cis): # Modulate after norm x_mlp_modulated = norm_x_mlp * (1 + scale_mlp) + shift_mlp mlp_output = self.mlp(x_mlp_modulated) + mlp_output = self.dropout(mlp_output, deterministic=not train) if self.use_gating: x = residual + gate_mlp * mlp_output @@ -111,14 +114,11 @@ class SimpleDiT(nn.Module): dtype: Optional[Dtype] = None precision: PrecisionLike = None # Passed down, but RoPEAttention uses NormalAttention - use_flash_attention: bool = False force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 learn_sigma: bool = False # Option to predict sigma like in DiT paper use_hilbert: bool = False # Toggle Hilbert patch reorder use_zigzag: bool = False # Toggle zigzag (serpentine) patch reorder - norm_groups: int = 0 - activation: Callable = jax.nn.swish def setup(self): self.patch_embed = PatchEmbedding( @@ -178,7 +178,6 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, # Pass RoPE instance @@ -198,14 +197,14 @@ def setup(self): self.final_proj = nn.Dense( features=output_dim, - dtype=self.dtype, + dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 precision=self.precision, kernel_init=nn.initializers.zeros, # Initialize final layer to zero name="final_proj" ) @nn.compact - def __call__(self, x, temb, textcontext=None): + def __call__(self, x, temb, textcontext=None, train: bool = False): B, H, W, C = x.shape assert H % self.patch_size == 0 and W % self.patch_size == 0, "Image dimensions must be divisible by patch size" @@ -272,7 +271,7 @@ def __call__(self, x, temb, textcontext=None): # 4. Apply Transformer Blocks with adaLN-Zero conditioning for block in self.blocks: - x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin)) + x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin), train=train) # 5. Final Layer x_out = self.final_norm(x_seq) diff --git a/flaxdiff/models/simple_mmdit.py b/flaxdiff/models/simple_mmdit.py index 5f27bf8..904701e 100644 --- a/flaxdiff/models/simple_mmdit.py +++ b/flaxdiff/models/simple_mmdit.py @@ -103,7 +103,6 @@ class MMDiTBlock(nn.Module): dtype: Optional[Dtype] = None precision: PrecisionLike = None # Keep option, though RoPEAttention doesn't use it - use_flash_attention: bool = False force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 @@ -134,8 +133,10 @@ def setup(self): precision=self.precision) ]) + self.dropout = nn.Dropout(rate=self.dropout_rate) + @nn.compact - def __call__(self, x, t_emb, text_emb, freqs_cis): + def __call__(self, x, t_emb, text_emb, freqs_cis, train: bool = False): # x shape: [B, S, F] # t_emb shape: [B, D_t] or [B, 1, D_t] # text_emb shape: [B, D_text] or [B, 1, D_text] @@ -149,10 +150,12 @@ def __call__(self, x, t_emb, text_emb, freqs_cis): # Attention block (remains the same) attn_output = self.attention( x_attn, context=None, freqs_cis=freqs_cis) # Self-attention only + attn_output = self.dropout(attn_output, deterministic=not train) x = residual + gate_attn * attn_output # MLP block (remains the same) mlp_output = self.mlp(x_mlp) + mlp_output = self.dropout(mlp_output, deterministic=not train) x = x + gate_mlp * mlp_output return x @@ -176,13 +179,10 @@ class SimpleMMDiT(nn.Module): dtype: Optional[Dtype] = None precision: PrecisionLike = None # Passed down, but RoPEAttention uses NormalAttention - use_flash_attention: bool = False force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 learn_sigma: bool = False # Option to predict sigma like in DiT paper use_hilbert: bool = False # Toggle Hilbert patch reorder - norm_groups: int = 0 - activation: Callable = jax.nn.swish def setup(self): self.patch_embed = PatchEmbedding( @@ -228,7 +228,6 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, # Pass RoPE instance @@ -248,14 +247,14 @@ def setup(self): self.final_proj = nn.Dense( features=output_dim, - dtype=self.dtype, + dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 precision=self.precision, kernel_init=nn.initializers.zeros, # Initialize final layer to zero name="final_proj" ) @nn.compact - def __call__(self, x, temb, textcontext): # textcontext is required + def __call__(self, x, temb, textcontext, train: bool = False): # textcontext is required B, H, W, C = x.shape assert H % self.patch_size == 0 and W % self.patch_size == 0, "Image dimensions must be divisible by patch size" assert textcontext is not None, "textcontext must be provided for SimpleMMDiT" @@ -291,7 +290,7 @@ def __call__(self, x, temb, textcontext): # textcontext is required for block in self.blocks: # Pass t_emb and text_emb separately to the block x_seq = block(x_seq, t_emb, text_emb, - freqs_cis=(freqs_cos, freqs_sin)) + freqs_cis=(freqs_cos, freqs_sin), train=train) # 5. Final Layer x_seq = self.final_norm(x_seq) @@ -446,13 +445,10 @@ class HierarchicalMMDiT(nn.Module): dropout_rate: float = 0.0 dtype: Optional[Dtype] = None precision: PrecisionLike = None - use_flash_attention: bool = False force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 learn_sigma: bool = False use_hilbert: bool = False - norm_groups: int = 0 # Not used in this structure, maybe remove later - activation: Callable = jax.nn.swish # Not used directly here, used in MLP inside MMDiTBlock def setup(self): assert len(self.emb_features) == len(self.num_layers) == len(self.num_heads), \ @@ -525,8 +521,7 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, - force_fp32_for_softmax=self.force_fp32_for_softmax, + force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.ropes[stage], name=f"encoder_block_stage{stage}_{i}" @@ -589,8 +584,7 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, - force_fp32_for_softmax=self.force_fp32_for_softmax, + force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.ropes[stage], name=f"decoder_block_stage{stage}_{i}" @@ -618,13 +612,13 @@ def setup(self): self.final_proj = nn.Dense( features=output_dim, - dtype=self.dtype, + dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 precision=self.precision, kernel_init=nn.initializers.zeros, # Zero init name="final_proj" ) - def __call__(self, x, temb, textcontext): + def __call__(self, x, temb, textcontext, train: bool = False): B, H, W, C = x.shape num_stages = len(self.emb_features) finest_patch_size = self.base_patch_size @@ -668,7 +662,7 @@ def __call__(self, x, temb, textcontext): # Apply blocks for this stage for block in self.encoder_blocks[stage]: - x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin)) + x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin), train=train) # Store skip features (before merging) skip_features[stage] = x_seq @@ -701,7 +695,7 @@ def __call__(self, x, temb, textcontext): # Apply blocks for this stage for block in self.decoder_blocks[i]: # Use index i for the decoder block list - x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin)) + x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin), train=train) # --- Final Layer --- # x_seq should now be at the finest resolution (stage 0 features) diff --git a/flaxdiff/models/simple_unet.py b/flaxdiff/models/simple_unet.py index ae9e3c4..90cc982 100644 --- a/flaxdiff/models/simple_unet.py +++ b/flaxdiff/models/simple_unet.py @@ -30,7 +30,7 @@ def setup(self): self.conv_out_norm = norm() @nn.compact - def __call__(self, x, temb, textcontext): + def __call__(self, x, temb, textcontext, train: bool = False): # print("embedding features", self.emb_features) temb = FourierEmbedding(features=self.emb_features)(temb) temb = TimeProjection(features=self.emb_features)(temb) @@ -74,7 +74,6 @@ def __call__(self, x, temb, textcontext): if attention_config is not None and j == self.num_res_blocks - 1: # Apply attention only on the last block x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), dim_head=dim_in // attention_config['heads'], - use_flash_attention=attention_config.get("flash_attention", False), use_projection=attention_config.get("use_projection", False), use_self_and_cross=attention_config.get("use_self_and_cross", True), precision=attention_config.get("precision", self.precision), @@ -115,7 +114,6 @@ def __call__(self, x, temb, textcontext): if middle_attention is not None and j == self.num_middle_res_blocks - 1: # Apply attention only on the last block x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), dim_head=middle_dim_out // middle_attention['heads'], - use_flash_attention=middle_attention.get("flash_attention", False), use_linear_attention=False, use_projection=middle_attention.get("use_projection", False), use_self_and_cross=False, @@ -161,7 +159,6 @@ def __call__(self, x, temb, textcontext): if attention_config is not None and j == self.num_res_blocks - 1: # Apply attention only on the last block x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), dim_head=dim_out // attention_config['heads'], - use_flash_attention=attention_config.get("flash_attention", False), use_projection=attention_config.get("use_projection", False), use_self_and_cross=attention_config.get("use_self_and_cross", True), precision=attention_config.get("precision", self.precision), diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py index 92f130c..0aca9fe 100644 --- a/flaxdiff/models/simple_vit.py +++ b/flaxdiff/models/simple_vit.py @@ -21,9 +21,7 @@ class UViT(nn.Module): emb_features: int = 768 num_layers: int = 12 # Should be even for U-Net structure num_heads: int = 12 - dropout_rate: float = 0.1 # Dropout is often 0 in diffusion models use_projection: bool = False # In TransformerBlock MLP - use_flash_attention: bool = False # Passed to TransformerBlock # Passed to TransformerBlock (likely False for UViT) use_self_and_cross: bool = False force_fp32_for_softmax: bool = True # Passed to TransformerBlock @@ -94,7 +92,7 @@ def setup(self): heads=self.num_heads, dim_head=self.emb_features // self.num_heads, dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, - use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross, + use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, only_pure_attention=False, norm_inputs=self.norm_inputs, explicitly_add_residual=self.explicitly_add_residual, @@ -107,7 +105,7 @@ def setup(self): heads=self.num_heads, dim_head=self.emb_features // self.num_heads, dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, - use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross, + use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, only_pure_attention=False, norm_inputs=self.norm_inputs, explicitly_add_residual=self.explicitly_add_residual, @@ -128,7 +126,7 @@ def setup(self): heads=self.num_heads, dim_head=self.emb_features // self.num_heads, dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, - use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross, + use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, only_pure_attention=False, norm_inputs=self.norm_inputs, explicitly_add_residual=self.explicitly_add_residual, @@ -172,7 +170,7 @@ def setup(self): ) @nn.compact - def __call__(self, x, temb, textcontext=None): + def __call__(self, x, temb, textcontext=None, train: bool = False): original_img = x B, H, W, C = original_img.shape H_P = H // self.patch_size @@ -182,13 +180,10 @@ def __call__(self, x, temb, textcontext=None): hilbert_inv_idx = None if self.use_hilbert: - patches_raw, hilbert_inv_idx_calc = hilbert_patchify( - x, self.patch_size) + # hilbert_patchify already returns the patches in hilbert order + # along with the inverse permutation for the output path + patches_raw, hilbert_inv_idx = hilbert_patchify(x, self.patch_size) x_patches = self.hilbert_proj(patches_raw) - idx = hilbert_indices(H_P, W_P) - hilbert_inv_idx = inverse_permutation( - idx, total_size=num_patches) - x_patches = x_patches[:, idx, :] else: x_patches = self.patch_embed(x) @@ -267,7 +262,6 @@ class SimpleUDiT(nn.Module): dropout_rate: float = 0.0 # Typically 0 for diffusion dtype: Optional[Dtype] = None # e.g., jnp.float32 or jnp.bfloat16 precision: PrecisionLike = None - use_flash_attention: bool = False # Passed to DiTBlock -> RoPEAttention force_fp32_for_softmax: bool = True # Passed to DiTBlock -> RoPEAttention norm_epsilon: float = 1e-5 learn_sigma: bool = False @@ -323,7 +317,6 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, @@ -338,7 +331,6 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, @@ -361,7 +353,6 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, @@ -385,7 +376,7 @@ def setup(self): ) @nn.compact - def __call__(self, x, temb, textcontext=None): + def __call__(self, x, temb, textcontext=None, train: bool = False): B, H, W, C = x.shape H_P = H // self.patch_size W_P = W // self.patch_size @@ -396,10 +387,8 @@ def __call__(self, x, temb, textcontext=None): hilbert_inv_idx = None if self.use_hilbert: - patches_raw, _ = hilbert_patchify(x, self.patch_size) + patches_raw, hilbert_inv_idx = hilbert_patchify(x, self.patch_size) x_seq = self.hilbert_proj(patches_raw) - idx = hilbert_indices(H_P, W_P) - hilbert_inv_idx = inverse_permutation(idx, total_size=num_patches) else: x_seq = self.patch_embed(x) @@ -415,16 +404,16 @@ def __call__(self, x, temb, textcontext=None): skips = [] for i in range(self.num_layers // 2): - x_seq = self.down_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None) + x_seq = self.down_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None, train=train) skips.append(x_seq) - x_seq = self.mid_block(x_seq, conditioning=cond_emb, freqs_cis=None) + x_seq = self.mid_block(x_seq, conditioning=cond_emb, freqs_cis=None, train=train) for i in range(self.num_layers // 2): skip_conn = skips.pop() x_seq = jnp.concatenate([x_seq, skip_conn], axis=-1) x_seq = self.up_dense[i](x_seq) - x_seq = self.up_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None) + x_seq = self.up_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None, train=train) x_out = self.final_norm(x_seq) x_out = self.final_proj(x_out) diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py index 3d2c2c8..94930d7 100644 --- a/flaxdiff/models/ssm_dit.py +++ b/flaxdiff/models/ssm_dit.py @@ -55,6 +55,7 @@ class S5Layer(nn.Module): def __call__(self, u): # u: [B, S, F] B, S, F = u.shape + assert F == self.features, f"S5Layer built for {self.features} features, got {F}" # A: diagonal complex state matrix, HiPPO init, parameterized as # log of the negative real part for stability @@ -257,7 +258,6 @@ class SSMDiTBlock(nn.Module): dropout_rate: float = 0.0 dtype: Optional[Dtype] = None precision: PrecisionLike = None - use_flash_attention: bool = False # Ignored, interface compat force_fp32_for_softmax: bool = True # Ignored, interface compat norm_epsilon: float = 1e-5 use_gating: bool = True @@ -308,6 +308,8 @@ def setup(self): nn.Dense(features=self.features, dtype=self.dtype, precision=self.precision) ]) + self.dropout = nn.Dropout(rate=self.dropout_rate) + def _apply_2d_fusion(self, ssm_output): """Un-permute scan-ordered SSM output to a 2D grid, fuse, re-permute back.""" B, S, F = ssm_output.shape @@ -351,7 +353,7 @@ def _apply_2d_fusion(self, ssm_output): return y_fused @nn.compact - def __call__(self, x, conditioning, freqs_cis): + def __call__(self, x, conditioning, freqs_cis, train: bool = False): # Get scale/shift/gate parameters scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( self.ada_params_module(conditioning), 6, axis=-1 @@ -365,6 +367,7 @@ def __call__(self, x, conditioning, freqs_cis): if self.use_2d_fusion: ssm_output = self._apply_2d_fusion(ssm_output) + ssm_output = self.dropout(ssm_output, deterministic=not train) if self.use_gating: x = residual + gate_attn * ssm_output @@ -376,6 +379,7 @@ def __call__(self, x, conditioning, freqs_cis): norm_x_mlp = self.norm2(x) x_mlp_modulated = norm_x_mlp * (1 + scale_mlp) + shift_mlp mlp_output = self.mlp(x_mlp_modulated) + mlp_output = self.dropout(mlp_output, deterministic=not train) if self.use_gating: x = residual + gate_mlp * mlp_output @@ -401,14 +405,11 @@ class HybridSSMAttentionDiT(nn.Module): dropout_rate: float = 0.0 dtype: Optional[Dtype] = None precision: PrecisionLike = None - use_flash_attention: bool = False force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 learn_sigma: bool = False use_hilbert: bool = False use_zigzag: bool = False # ZigMa-style serpentine scan - norm_groups: int = 0 - activation: Callable = jax.nn.swish block_pattern: Optional[Sequence[str]] = None # e.g., ['ssm','ssm','ssm','attn'] ssm_attention_ratio: str = "3:1" # e.g., "3:1", "1:1", "all-ssm", "all-attn" bidirectional_ssm: bool = True @@ -500,8 +501,7 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - use_flash_attention=self.use_flash_attention, - force_fp32_for_softmax=self.force_fp32_for_softmax, + force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, name=f"dit_block_{i}" )) @@ -517,14 +517,14 @@ def setup(self): self.final_proj = nn.Dense( features=output_dim, - dtype=self.dtype, + dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 precision=self.precision, kernel_init=nn.initializers.zeros, name="final_proj" ) @nn.compact - def __call__(self, x, temb, textcontext=None): + def __call__(self, x, temb, textcontext=None, train: bool = False): B, H, W, C = x.shape assert H % self.patch_size == 0 and W % self.patch_size == 0 @@ -579,7 +579,7 @@ def __call__(self, x, temb, textcontext=None): # 4. Hybrid blocks (SSM and attention interleaved) for block in self.blocks: - x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin)) + x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin), train=train) # 5. Final output x_out = self.final_norm(x_seq) diff --git a/flaxdiff/predictors/__init__.py b/flaxdiff/predictors/__init__.py index 3aeb5d2..68390f2 100644 --- a/flaxdiff/predictors/__init__.py +++ b/flaxdiff/predictors/__init__.py @@ -1,6 +1,13 @@ from typing import Union import jax.numpy as jnp -from ..schedulers import NoiseScheduler, GeneralizedNoiseScheduler, get_coeff_shapes_tuple +from ..schedulers import ( + NoiseScheduler, + GeneralizedNoiseScheduler, + get_coeff_shapes_tuple, + CosineNoiseScheduler, + KarrasVENoiseScheduler, + EDMNoiseScheduler, +) ############################################################################################################ # Prediction Transforms @@ -93,4 +100,38 @@ def pred_transform(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray], eps def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8) -> jnp.ndarray: _, sigma = rates c_in = 1 / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon) - return c_in \ No newline at end of file + return c_in + +############################################################################################################ +# Noise schedule presets +############################################################################################################ + +def get_diffusion_preset( + name: str, + timesteps: int = 1000, + sigma_min: float = 0.002, + sigma_max: float = 80.0, + rho: float = 7.0, + sigma_data: float = 0.5, +) -> 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. + """ + 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) + 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 == 'cosine': + train = CosineNoiseScheduler(timesteps, beta_end=1) + sample = CosineNoiseScheduler(timesteps, beta_end=1) + transform = VPredictionTransform() + 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 69d1559..910202c 100644 --- a/flaxdiff/samplers/common.py +++ b/flaxdiff/samplers/common.py @@ -7,10 +7,8 @@ from typing import List, Tuple, Dict, Any, Optional from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform -from ..schedulers import NoiseScheduler +from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple from ..utils import RandomMarkovState, MarkovState, clip_images -from jax.experimental.shard_map import shard_map -from jax.sharding import Mesh, PartitionSpec as P from flaxdiff.models.autoencoder import AutoEncoder from flaxdiff.inputs import DiffusionInputConfig @@ -25,7 +23,6 @@ def __init__( input_config: DiffusionInputConfig, guidance_scale: float = 0.0, autoencoder: AutoEncoder = None, - timestep_spacing: str = 'linear', ): """Initialize the diffusion sampler. @@ -36,27 +33,16 @@ def __init__( model_output_transform: Transform for model predictions guidance_scale: Scale for classifier-free guidance (0.0 means disabled) autoencoder: Optional autoencoder for latent diffusion - timestep_spacing: Strategy for timestep spacing in sampling - 'linear' - Default equal spacing - 'quadratic' - Emphasizes early steps - 'karras' - Based on EDM paper, better with fewer steps - 'exponential' - Concentrates steps near the end """ self.model = model self.noise_schedule = noise_schedule self.model_output_transform = model_output_transform self.guidance_scale = guidance_scale self.autoencoder = autoencoder - self.timestep_spacing = timestep_spacing self.input_config = input_config unconditionals = input_config.get_unconditionals() - # For Karras spacing if needed - if hasattr(noise_schedule, 'min_inv_rho') and hasattr(noise_schedule, 'max_inv_rho'): - self.min_inv_rho = noise_schedule.min_inv_rho - self.max_inv_rho = noise_schedule.max_inv_rho - if self.guidance_scale > 0: # Classifier free guidance print("Using classifier-free guidance") @@ -65,7 +51,7 @@ def sample_model(params, x_t, t, *conditioning_inputs): # Concatenate unconditional and conditional inputs x_t_cat = jnp.concatenate([x_t] * 2, axis=0) t_cat = jnp.concatenate([t] * 2, axis=0) - rates_cat = self.noise_schedule.get_rates(t_cat) + rates_cat = self.noise_schedule.get_rates(t_cat, get_coeff_shapes_tuple(x_t_cat)) c_in_cat = self.model_output_transform.get_input_scale(rates_cat) final_conditionals = [] @@ -92,7 +78,7 @@ def sample_model(params, x_t, t, *conditioning_inputs): else: # Unconditional sampling def sample_model(params, x_t, t, *conditioning_inputs): - rates = self.noise_schedule.get_rates(t) + rates = self.noise_schedule.get_rates(t, get_coeff_shapes_tuple(x_t)) c_in = self.model_output_transform.get_input_scale(rates) model_output = self.model.apply( params, @@ -136,7 +122,7 @@ def sample_step( Returns: Tuple of (new samples, updated state) """ - step_ones = jnp.ones((len(current_samples),), dtype=jnp.int32) + step_ones = jnp.ones((len(current_samples),), dtype=jnp.float32) current_step = step_ones * current_step next_step = step_ones * next_step @@ -175,74 +161,30 @@ def take_next_step( raise NotImplementedError("Subclasses must implement take_next_step method") - def scale_steps(self, steps): - """Scale timesteps to match the noise schedule's range.""" - scale_factor = self.noise_schedule.max_timesteps / 1000 - return steps * scale_factor + def get_steps(self, start_step, end_step, diffusion_steps): + """Get the descending sequence of timesteps for the diffusion process. + Timesteps are in the noise schedule's own domain: [0, max_timesteps], + so [0, 1] for the continuous schedules and [0, T] for the discrete ones. + Note that for a Karras schedule, linear spacing in t is already the + rho-spacing in sigma from the EDM paper. - def get_steps(self, start_step, end_step, diffusion_steps): - """Get the sequence of timesteps for the diffusion process. - Args: - start_step: Starting timestep (typically the max) + start_step: Starting timestep (typically max_timesteps) end_step: Ending timestep (typically 0) diffusion_steps: Number of steps to use - + Returns: Array of timesteps for sampling """ step_range = start_step - end_step if diffusion_steps is None or diffusion_steps == 0: - diffusion_steps = step_range - diffusion_steps = min(diffusion_steps, step_range) - - # Linear spacing (default) - if getattr(self, 'timestep_spacing', 'linear') == 'linear': - steps = jnp.linspace( - end_step, start_step, - diffusion_steps, dtype=jnp.int16 - )[::-1] - - # Quadratic spacing (emphasizes early steps) - elif self.timestep_spacing == 'quadratic': - steps = jnp.linspace(0, 1, diffusion_steps) ** 2 - steps = (start_step - end_step) * steps + end_step - steps = jnp.asarray(steps, dtype=jnp.int16)[::-1] - - # Karras spacing from the EDM paper - often gives better results with fewer steps - elif self.timestep_spacing == 'karras': - # Implementation based on the EDM paper's recommendations - sigma_min = end_step / start_step - sigma_max = 1.0 - rho = 7.0 # Karras paper default, controls the distribution - - # Create log-spaced steps in sigma space - sigmas = jnp.exp(jnp.linspace( - jnp.log(sigma_max), jnp.log(sigma_min), diffusion_steps - )) - steps = jnp.clip( - (sigmas ** (1 / rho) - self.min_inv_rho) / - (self.max_inv_rho - self.min_inv_rho), - 0, 1 - ) * start_step - steps = jnp.asarray(steps, dtype=jnp.int16) - - # Exponential spacing (concentrates steps near the end) - elif self.timestep_spacing == 'exponential': - steps = jnp.linspace(0, 1, diffusion_steps) - steps = jnp.exp(steps * jnp.log((start_step + 1) / (end_step + 1))) * (end_step + 1) - 1 - steps = jnp.clip(steps, end_step, start_step) - steps = jnp.asarray(steps, dtype=jnp.int16)[::-1] - - # Fallback to linear spacing - else: - steps = jnp.linspace( - end_step, start_step, - diffusion_steps, dtype=jnp.int16 - )[::-1] - - return steps + diffusion_steps = int(step_range) + if self.noise_schedule.max_timesteps > 1: + # Discrete schedules cannot take more steps than they have entries + diffusion_steps = min(diffusion_steps, int(step_range)) + + return jnp.linspace(start_step, end_step, diffusion_steps, dtype=jnp.float32) def generate_samples( @@ -374,8 +316,8 @@ def sample_step(sample_model_fn, state: RandomMarkovState, samples, current_step steps = self.get_steps(start_step, end_step, diffusion_steps) for i in tqdm.tqdm(range(0, len(steps))): - current_step = self.scale_steps(steps[i]) - next_step = self.scale_steps(steps[i+1] if i+1 < len(steps) else 0) + current_step = steps[i] + next_step = steps[i+1] if i+1 < len(steps) else end_step if i != len(steps) - 1: samples, rngstate = sample_step( @@ -397,7 +339,6 @@ def _get_noise_parameters(self, resolution, start_step): Returns: Tuple of (variance, image_size, image_channels) """ - start_step = self.scale_steps(start_step) alpha_n, sigma_n = self.noise_schedule.get_rates(start_step) variance = jnp.sqrt(alpha_n ** 2 + sigma_n ** 2) diff --git a/flaxdiff/samplers/ddim.py b/flaxdiff/samplers/ddim.py index 4921952..6bcc235 100644 --- a/flaxdiff/samplers/ddim.py +++ b/flaxdiff/samplers/ddim.py @@ -32,18 +32,18 @@ def take_next_step( # Extract random noise if needed for stochastic sampling if self.eta > 0: - # For DDIM, we need to compute the variance coefficient - # This is based on the original DDIM paper's formula - # When eta=0, it's deterministic DDIM, when eta=1.0 it approaches DDPM - sigma_tilde = self.eta * sigma_next * (1 - alpha_t**2 / alpha_next**2).sqrt() / (1 - alpha_t**2).sqrt() + # DDIM paper eq. 16: eta=0 is deterministic DDIM, eta=1.0 approaches DDPM. + # The direction term must shrink to keep the marginal variance right. + sigma_tilde = self.eta * (sigma_next / sigma_t) * jnp.sqrt( + jnp.maximum(1 - alpha_t**2 / alpha_next**2, 0.0)) state, noise_key = state.get_random_key() noise = jax.random.normal(noise_key, current_samples.shape) - # Add the stochastic component - stochastic_term = sigma_tilde * noise + direction_coeff = jnp.sqrt(jnp.maximum(sigma_next**2 - sigma_tilde**2, 0.0)) + new_samples = (alpha_next * reconstructed_samples + + direction_coeff * pred_noise + + sigma_tilde * noise) else: - stochastic_term = 0 - - # Direct DDIM update formula - new_samples = alpha_next * reconstructed_samples + sigma_next * pred_noise + stochastic_term + # Direct DDIM update formula + new_samples = alpha_next * reconstructed_samples + sigma_next * pred_noise return new_samples, state diff --git a/flaxdiff/samplers/ddpm.py b/flaxdiff/samplers/ddpm.py index 7c07d46..1f5da65 100644 --- a/flaxdiff/samplers/ddpm.py +++ b/flaxdiff/samplers/ddpm.py @@ -1,37 +1,34 @@ import jax import jax.numpy as jnp from .common import DiffusionSampler +from ..schedulers import get_coeff_shapes_tuple from ..utils import MarkovState, RandomMarkovState + class DDPMSampler(DiffusionSampler): - def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, - pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: - mean = self.noise_schedule.get_posterior_mean(reconstructed_samples, current_samples, current_step) - variance = self.noise_schedule.get_posterior_variance(steps=current_step) - - state, rng = state.get_random_key() - # Now sample from the posterior - noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32) + """Exact ancestral sampler for the reverse diffusion SDE. - return mean + noise * variance, state - - def generate_images(self, num_images=16, diffusion_steps=1000, start_step: int = None, *args, **kwargs): - return super().generate_images(num_images=num_images, diffusion_steps=diffusion_steps, start_step=start_step, *args, **kwargs) - -class SimpleDDPMSampler(DiffusionSampler): + This is the simplified (but algebraically exact) form of the DDPM + posterior, phrased purely in terms of signal/noise rates so it works for + any schedule and any step stride, not just t -> t-1. + """ def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: state, rng = state.get_random_key() noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32) - # Compute noise rates and signal rates only once - current_signal_rate, current_noise_rate = self.noise_schedule.get_rates(current_step) - next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step) - + shape = get_coeff_shapes_tuple(current_samples) + current_signal_rate, current_noise_rate = self.noise_schedule.get_rates(current_step, shape) + next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step, shape) + pred_noise_coeff = ((next_noise_rate ** 2) * current_signal_rate) / (current_noise_rate * next_signal_rate) - + noise_ratio_squared = (next_noise_rate ** 2) / (current_noise_rate ** 2) signal_ratio_squared = (current_signal_rate ** 2) / (next_signal_rate ** 2) gamma = jnp.sqrt(noise_ratio_squared * (1 - signal_ratio_squared)) - + next_samples = next_signal_rate * reconstructed_samples + pred_noise_coeff * pred_noise + noise * gamma - return next_samples, state \ No newline at end of file + return next_samples, state + +# The two used to be separate implementations; the posterior-table variant +# crashed for batched steps and could not handle strided sampling +SimpleDDPMSampler = DDPMSampler diff --git a/flaxdiff/samplers/euler.py b/flaxdiff/samplers/euler.py index 62438b7..d7011f9 100644 --- a/flaxdiff/samplers/euler.py +++ b/flaxdiff/samplers/euler.py @@ -1,14 +1,16 @@ import jax import jax.numpy as jnp from .common import DiffusionSampler +from ..schedulers import get_coeff_shapes_tuple from ..utils import RandomMarkovState class EulerSampler(DiffusionSampler): # Basically a DDIM Sampler but parameterized as an ODE def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: - current_alpha, current_sigma = self.noise_schedule.get_rates(current_step) - next_alpha, next_sigma = self.noise_schedule.get_rates(next_step) + shape = get_coeff_shapes_tuple(current_samples) + current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape) + next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape) dt = next_sigma - current_sigma @@ -23,8 +25,9 @@ class SimplifiedEulerSampler(DiffusionSampler): """ def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: - _, current_sigma = self.noise_schedule.get_rates(current_step) - _, next_sigma = self.noise_schedule.get_rates(next_step) + shape = get_coeff_shapes_tuple(current_samples) + _, current_sigma = self.noise_schedule.get_rates(current_step, shape) + _, next_sigma = self.noise_schedule.get_rates(next_step, shape) dt = next_sigma - current_sigma @@ -38,8 +41,9 @@ class EulerAncestralSampler(DiffusionSampler): """ def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: - current_alpha, current_sigma = self.noise_schedule.get_rates(current_step) - next_alpha, next_sigma = self.noise_schedule.get_rates(next_step) + shape = get_coeff_shapes_tuple(current_samples) + current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape) + next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape) sigma_up = (next_sigma**2 * (current_sigma**2 - next_sigma**2) / current_sigma**2) ** 0.5 sigma_down = (next_sigma**2 - sigma_up**2) ** 0.5 diff --git a/flaxdiff/samplers/heun_sampler.py b/flaxdiff/samplers/heun_sampler.py index d553614..4baa1c3 100644 --- a/flaxdiff/samplers/heun_sampler.py +++ b/flaxdiff/samplers/heun_sampler.py @@ -1,14 +1,16 @@ import jax import jax.numpy as jnp from .common import DiffusionSampler +from ..schedulers import get_coeff_shapes_tuple from ..utils import RandomMarkovState class HeunSampler(DiffusionSampler): def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: # Get the noise and signal rates for the current and next steps - current_alpha, current_sigma = self.noise_schedule.get_rates(current_step) - next_alpha, next_sigma = self.noise_schedule.get_rates(next_step) + shape = get_coeff_shapes_tuple(current_samples) + current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape) + next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape) dt = next_sigma - current_sigma x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / dt @@ -20,8 +22,15 @@ def take_next_step(self, current_samples, reconstructed_samples, model_condition estimated_x_0, _, _ = sample_model_fn(next_samples_0, next_step, *model_conditioning_inputs) # Estimate the refined derivative using the midpoint (Heun's method) - dx_1 = (next_samples_0 - x_0_coeff * estimated_x_0) / next_sigma - # Compute the final next samples by averaging the initial and refined derivatives - final_next_samples = current_samples + 0.5 * (dx_0 + dx_1) * dt + safe_next_sigma = jnp.where(next_sigma > 0, next_sigma, 1.0) + dx_1 = (next_samples_0 - x_0_coeff * estimated_x_0) / safe_next_sigma + # Compute the final next samples by averaging the initial and refined derivatives. + # When sigma reaches 0 there is no derivative there - fall back to the euler step + # (Karras et al. 2022, Algorithm 2) + final_next_samples = jnp.where( + next_sigma > 0, + current_samples + 0.5 * (dx_0 + dx_1) * dt, + next_samples_0, + ) return final_next_samples, state diff --git a/flaxdiff/samplers/multistep_dpm.py b/flaxdiff/samplers/multistep_dpm.py index b3e7fe9..cb184e8 100644 --- a/flaxdiff/samplers/multistep_dpm.py +++ b/flaxdiff/samplers/multistep_dpm.py @@ -1,6 +1,7 @@ import jax import jax.numpy as jnp from .common import DiffusionSampler +from ..schedulers import get_coeff_shapes_tuple from ..utils import RandomMarkovState class MultiStepDPM(DiffusionSampler): @@ -8,11 +9,17 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.history = [] + def generate_samples(self, *args, **kwargs): + # The multistep history is only valid within a single trajectory + self.history = [] + return super().generate_samples(*args, **kwargs) + def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs, pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]: # Get the noise and signal rates for the current and next steps - current_alpha, current_sigma = self.noise_schedule.get_rates(current_step) - next_alpha, next_sigma = self.noise_schedule.get_rates(next_step) + shape = get_coeff_shapes_tuple(current_samples) + current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape) + next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape) dt = next_sigma - current_sigma diff --git a/flaxdiff/samplers/rk4_sampler.py b/flaxdiff/samplers/rk4_sampler.py index a4cdbe8..1763b88 100644 --- a/flaxdiff/samplers/rk4_sampler.py +++ b/flaxdiff/samplers/rk4_sampler.py @@ -2,13 +2,14 @@ import jax.numpy as jnp from .common import DiffusionSampler from ..utils import RandomMarkovState, MarkovState -from ..schedulers import GeneralizedNoiseScheduler +from ..schedulers import GeneralizedNoiseScheduler, get_coeff_shapes_tuple class RK4Sampler(DiffusionSampler): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) assert issubclass(type(self.noise_schedule), GeneralizedNoiseScheduler), "Noise schedule must be a GeneralizedNoiseScheduler" - @jax.jit + # Not jitted here: sample_model_fn is a python callable, the model call + # inside it is already jitted def get_derivative(sample_model_fn, x_t, sigma, state:RandomMarkovState, model_conditioning_inputs) -> tuple[jnp.ndarray, RandomMarkovState]: t = self.noise_schedule.get_timesteps(sigma) x_0, eps, _ = sample_model_fn(x_t, t, *model_conditioning_inputs) @@ -17,11 +18,12 @@ def get_derivative(sample_model_fn, x_t, sigma, state:RandomMarkovState, model_c self.get_derivative = get_derivative def sample_step(self, sample_model_fn, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]: - step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.int32) + step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.float32) current_step = step_ones * current_step next_step = step_ones * next_step - _, current_sigma = self.noise_schedule.get_rates(current_step) - _, next_sigma = self.noise_schedule.get_rates(next_step) + shape = get_coeff_shapes_tuple(current_samples) + _, current_sigma = self.noise_schedule.get_rates(current_step, shape) + _, next_sigma = self.noise_schedule.get_rates(next_step, shape) dt = next_sigma - current_sigma diff --git a/flaxdiff/schedulers/karras.py b/flaxdiff/schedulers/karras.py index b651c75..b9e09d5 100644 --- a/flaxdiff/schedulers/karras.py +++ b/flaxdiff/schedulers/karras.py @@ -19,9 +19,10 @@ def get_sigmas(self, steps) -> jnp.ndarray: def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray: sigma = self.get_sigmas(steps) - # Add epsilon for numerical stability - epsilon = 1e-6 - weights = ((sigma ** 2 + self.sigma_data ** 2) / ((sigma * self.sigma_data) ** 2 + epsilon)) + # 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 + # sigma_min where (sigma * sd)^2 == 1e-6) + weights = 1 / self.sigma_data ** 2 + 1 / sigma ** 2 return weights.reshape(shape) def transform_inputs(self, x, steps, num_discrete_chunks=1000) -> tuple[jnp.ndarray, jnp.ndarray]: diff --git a/flaxdiff/trainer/diffusion_trainer.py b/flaxdiff/trainer/diffusion_trainer.py index 2004da6..adbb834 100644 --- a/flaxdiff/trainer/diffusion_trainer.py +++ b/flaxdiff/trainer/diffusion_trainer.py @@ -93,42 +93,25 @@ def generate_states( self, optimizer: optax.GradientTransformation, rngs: jax.random.PRNGKey, - existing_state: dict = None, - existing_best_state: dict = None, model: nn.Module = None, - param_transforms: Callable = None, use_dynamic_scale: bool = False ) -> Tuple[TrainState, TrainState]: print("Generating states for DiffusionTrainer") rngs, subkey = jax.random.split(rngs) - if existing_state == None: - input_vars = self.get_input_ones() - params = model.init(subkey, **input_vars) - new_state = {"params": params, "ema_params": params} - else: - new_state = existing_state - - if param_transforms is not None: - params = param_transforms(params) + input_vars = self.get_input_ones() + params = model.init(subkey, **input_vars) state = TrainState.create( apply_fn=model.apply, - params=new_state['params'], - ema_params=new_state['ema_params'], + params=params, + ema_params=params, tx=optimizer, rngs=rngs, metrics=Metrics.empty(), dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None ) - - if existing_best_state is not None: - best_state = state.replace( - params=existing_best_state['params'], ema_params=existing_best_state['ema_params']) - else: - best_state = state - - return state, best_state + return state, state def _define_train_step(self, batch_size): noise_schedule: NoiseScheduler = self.noise_schedule @@ -192,6 +175,7 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc local_rng_state, rngs = local_rng_state.get_random_key() noise: jax.Array = jax.random.normal(rngs, shape=images.shape, dtype=jnp.float32) + local_rng_state, dropout_key = local_rng_state.get_random_key() # Make sure image is also float32 images = images.astype(jnp.float32) @@ -200,7 +184,10 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc noisy_images, c_in, expected_output = model_output_transform.forward_diffusion(images, noise, rates) def model_loss(params): - preds = model.apply(params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq) + preds = model.apply( + params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq, + train=True, rngs={'dropout': dropout_key}, + ) preds = model_output_transform.pred_transform(noisy_images, preds, rates) nloss = loss_fn(preds, expected_output) # Ignore the loss contribution of images with zero standard deviation @@ -300,7 +287,6 @@ def generate_samples( resolution=image_size, num_samples=len(labels_seq), diffusion_steps=diffusion_steps, - start_step=1000, end_step=0, priors=None, model_conditioning_inputs=(labels_seq,), @@ -333,7 +319,7 @@ def validation_loop( ) # Put each sample on wandb - if getattr(self, 'wandb', None) is not None and self.wandb: + if self.wandb is not None and self.wandb: import numpy as np from wandb import Image as wandbImage wandb_images = [] diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index 08c6e64..147bc22 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -279,6 +279,8 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc # Generate noise local_rng_state, noise_key = local_rng_state.get_random_key() noise = jax.random.normal(noise_key, shape=data.shape, dtype=jnp.float32) + + local_rng_state, dropout_key = local_rng_state.get_random_key() # Forward diffusion process rates = noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(data)) @@ -288,7 +290,10 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc def model_loss(params): # Apply model inputs = noise_schedule.transform_inputs(noisy_data * c_in, noise_level) - preds = model.apply(params, *inputs, *all_conditional_inputs) + preds = model.apply( + params, *inputs, *all_conditional_inputs, + train=True, rngs={'dropout': dropout_key}, + ) # Transform predictions and calculate loss preds = model_output_transform.pred_transform(noisy_data, preds, rates) @@ -312,8 +317,8 @@ def model_loss(params): # Handle NaN/Inf gradients select_fn = functools.partial(jnp.where, is_finite) new_state = new_state.replace( - opt_state=jax.tree_map(select_fn, new_state.opt_state, train_state.opt_state), - params=jax.tree_map(select_fn, new_state.params, train_state.params) + opt_state=jax.tree.map(select_fn, new_state.opt_state, train_state.opt_state), + params=jax.tree.map(select_fn, new_state.params, train_state.params) ) else: # Standard gradient computation @@ -392,7 +397,6 @@ def generate_samples( num_samples=batch_size, sequence_length=sequence_length, # Will be None for images diffusion_steps=diffusion_steps, - start_step=1000, end_step=0, priors=None, model_conditioning_inputs=tuple(model_conditioning_inputs), @@ -467,7 +471,7 @@ def validation_loop( if i == 0: print(f"Evaluation started for process index {process_index}") # Log samples to wandb - if getattr(self, 'wandb', None) is not None and self.wandb: + if self.wandb is not None and self.wandb: import numpy as np # Process samples differently based on dimensionality @@ -489,7 +493,7 @@ def validation_loop( prev = self.best_val_metrics[final_key] self.best_val_metrics[final_key] = max(prev, value) if higher_is_better else min(prev, value) # Log the best validation metrics - if getattr(self, 'wandb', None) is not None and self.wandb: + if self.wandb is not None and self.wandb: # Log the metrics for key, value in metrics.items(): if isinstance(value, jnp.ndarray): @@ -712,12 +716,12 @@ def save(self, epoch=0, step=0, state=None, rngstate=None): aliases.append("best") self.push_to_registry(aliases=aliases) print("Model pushed to registry successfully with aliases:", aliases) + # Only delete after a successful registry push - the local + # checkpoint is the only copy otherwise + shutil.rmtree(checkpoint, ignore_errors=True) + print(f"Checkpoint deleted at {checkpoint}") else: - print("Current run is not one of the best runs. Not saving model.") - - # Only delete after successful registry push - shutil.rmtree(checkpoint, ignore_errors=True) - print(f"Checkpoint deleted at {checkpoint}") + print("Current run is not one of the best runs. Keeping local checkpoint.") except Exception as e: print(f"Error during registry operations: {e}") print(f"Checkpoint preserved at {checkpoint}") diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py index b31a33c..1d23acd 100644 --- a/flaxdiff/trainer/simple_trainer.py +++ b/flaxdiff/trainer/simple_trainer.py @@ -161,7 +161,6 @@ def __init__(self, name: str = "Simple", load_from_checkpoint: str = None, loss_fn=optax.l2_loss, - param_transforms: Callable = None, wandb_config: Dict[str, Any] = None, distributed_training: bool = None, checkpoint_base_path: str = "./checkpoints", @@ -186,6 +185,7 @@ def __init__(self, load_directly_from_dir = False + self.wandb = None if wandb_config is not None and jax.process_index() == 0: import wandb run = wandb.init(resume='allow', **wandb_config) @@ -234,27 +234,12 @@ def __init__(self, self.checkpointer = orbax.checkpoint.CheckpointManager( self.checkpoint_path(), async_checkpointer, options) - if load_from_checkpoint is not None: - latest_step, old_state, old_best_state, rngstate = self.load(load_from_checkpoint, checkpoint_step, load_directly_from_dir) - else: - latest_step, old_state, old_best_state, rngstate = 0, None, None, None - - self.latest_step = latest_step - - if train_start_step_override is not None: - self.latest_step = train_start_step_override - print(f"Overriding start step to {self.latest_step}") - - if rngstate: - self.rngstate = RandomMarkovState(**rngstate) - else: - self.rngstate = RandomMarkovState(rngs) - + self.rngstate = RandomMarkovState(rngs) self.rngstate, subkey = self.rngstate.get_random_key() if train_state == None: state, best_state = self.generate_states( - optimizer, subkey, old_state, old_best_state, model, param_transforms, use_dynamic_scale + optimizer, subkey, model, use_dynamic_scale ) self.init_state(state, best_state) else: @@ -262,6 +247,14 @@ def __init__(self, self.best_state = train_state self.best_loss = 1e9 + self.latest_step = 0 + if load_from_checkpoint is not None: + self.latest_step = self.load(load_from_checkpoint, checkpoint_step, load_directly_from_dir) + + if train_start_step_override is not None: + self.latest_step = train_start_step_override + print(f"Overriding start step to {self.latest_step}") + def get_input_ones(self): return {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()} @@ -269,20 +262,14 @@ def generate_states( self, optimizer: optax.GradientTransformation, rngs: jax.random.PRNGKey, - existing_state: dict = None, - existing_best_state: dict = None, model: nn.Module = None, - param_transforms: Callable = None, use_dynamic_scale: bool = False ) -> Tuple[SimpleTrainState, SimpleTrainState]: print("Generating states for SimpleTrainer") rngs, subkey = jax.random.split(rngs) - if existing_state == None: - input_vars = self.get_input_ones() - params = model.init(subkey, **input_vars) - else: - params = existing_state['params'] + input_vars = self.get_input_ones() + params = model.init(subkey, **input_vars) state = SimpleTrainState.create( apply_fn=model.apply, @@ -291,13 +278,7 @@ def generate_states( metrics=Metrics.empty(), dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None ) - if existing_best_state is not None: - best_state = state.replace( - params=existing_best_state['params']) - else: - best_state = state - - return state, best_state + return state, state def init_state( self, @@ -353,18 +334,29 @@ def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=Fal checkpoint_path if checkpoint_path else self.checkpoint_path(), f"{step}") self.loaded_checkpoint_path = loaded_checkpoint_path - ckpt = checkpointer.restore(step) if not load_directly_from_dir else checkpointer.restore(checkpoint_path) + + # Restore against the freshly-initialized states as a template so orbax + # rebuilds the exact pytree types - optimizer state and step included. + # Restoring untyped used to silently discard opt_state and reset the + # step counter (and with it the lr schedule) on every resume. + template = { + 'rngs': self.get_rngstate(), + 'state': self.get_state(), + 'best_state': self.get_best_state(), + 'best_loss': np.array(self.best_loss), + 'epoch': 0, + } + ckpt = checkpointer.restore(step, items=template) if not load_directly_from_dir else checkpointer.restore(checkpoint_path, items=template) - state = ckpt['state'] - best_state = ckpt['best_state'] - rngstate = ckpt['rngs'] - # Convert the state to a TrainState - self.best_loss = ckpt['best_loss'] + self.state = ckpt['state'] + self.best_state = ckpt['best_state'] + self.rngstate = ckpt['rngs'] + self.best_loss = float(ckpt['best_loss']) if self.best_loss == 0: # It cant be zero as that must have been some problem self.best_loss = 1e9 - print(f"Loaded model from checkpoint at step {step}", ckpt['best_loss']) - return step, state, best_state, rngstate + print(f"Loaded model from checkpoint at step {step}", self.best_loss) + return step def save(self, epoch=0, step=0, state=None, rngstate=None): print(f"Saving model at epoch {epoch} step {step}") @@ -516,11 +508,14 @@ def train_loop( global_device_indexes = 0 epoch_loss = 0 + bad_loss_steps = 0 current_epoch = current_step // train_steps_per_epoch last_save_time = time.time() if process_index == 0: pbar = tqdm.tqdm(total=train_steps_per_epoch, desc=f'\t\tEpoch {current_epoch}', ncols=100, unit='step') + else: + pbar = None for i in range(train_steps_per_epoch): batch = next(train_ds) @@ -539,41 +534,18 @@ def train_loop( # loss = jax.experimental.multihost_utils.process_allgather(loss) loss = jnp.mean(loss) # Just to make sure its a scaler value - if loss <= 1e-8 or jnp.isnan(loss) or jnp.isinf(loss): - # If the loss is too low or NaN/Inf, log the issue and attempt recovery - print(colored(f"Abnormal loss at step {current_step}: {loss}", 'red')) - - # Check model parameters for NaN/Inf values - params = train_state.params - has_nan_or_inf = False - - if isinstance(params, dict): - for key, value in params.items(): - if isinstance(value, jnp.ndarray): - if jnp.isnan(value).any() or jnp.isinf(value).any(): - print(colored(f"NaN/inf values found in params[{key}] at step {current_step}", 'red')) - has_nan_or_inf = True - break - - if not has_nan_or_inf: - print(colored(f"Model parameters seem valid despite abnormal loss", 'yellow')) - - # Try to recover - clear JAX caches and collect garbage - gc.collect() - if hasattr(jax, "clear_caches"): - jax.clear_caches() - - # If we have a best state and the loss is truly invalid, consider restoring - if (loss <= 1e-8 or jnp.isnan(loss) or jnp.isinf(loss)) and self.best_state is not None: - print(colored(f"Attempting recovery by resetting model to last best state", 'yellow')) - train_state = self.best_state - loss = self.best_loss - else: - # If we can't recover, skip this step but continue training - print(colored(f"Unable to recover - continuing with current state", 'yellow')) - if loss <= 1e-8: - loss = 1.0 # Set to a reasonable default to continue training - + if not jnp.isfinite(loss): + # No silent recovery: a diverged run must fail loudly, not be + # papered over with a stale best_state and a cosmetic loss value + print(colored(f"Non-finite loss at step {current_step}: {loss}", 'red')) + bad_loss_steps += 1 + if bad_loss_steps >= 5: + raise RuntimeError( + f"Loss has been non-finite for {bad_loss_steps} consecutive steps, stopping" + ) + else: + bad_loss_steps = 0 + epoch_loss += loss current_step += 1 if i % 100 == 0: diff --git a/tests/test_config.py b/tests/test_config.py index b7ef699..dea975d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -18,9 +18,7 @@ def make_config(model_overrides=None, arguments_overrides=None): "emb_features": 64, "dtype": "bfloat16", "precision": "high", - "activation": "swish", "output_channels": 3, - "norm_groups": 4, "patch_size": 4, "num_layers": 2, "num_heads": 2, @@ -56,7 +54,14 @@ def test_parse_config_rebuilds_model(): assert model.num_layers == 2 assert model.dtype == jnp.bfloat16 assert model.precision == jax.lax.Precision.HIGH - assert model.activation is jax.nn.swish + + +def test_parse_config_drops_removed_fields(): + """Configs from older runs carry since-removed flags; reconstruction must + drop them instead of crashing.""" + result = parse_config(make_config(model_overrides={ + "use_flash_attention": False, "norm_groups": 8, "activation": "swish"})) + assert type(result["model"]).__name__ == "SimpleDiT" def test_parse_config_noise_schedule_selection(): @@ -67,11 +72,39 @@ def test_parse_config_noise_schedule_selection(): assert type(cosine["noise_schedule"]).__name__ == "CosineNoiseScheduler" -@pytest.mark.xfail(strict=True, reason="bug: parse_config silently drops any string value containing a dot") -def test_parse_config_preserves_dotted_values(): - """String config values with a dot in them (module paths, filenames, version - strings) must survive the round trip, not vanish into class defaults.""" - result = parse_config(make_config(model_overrides={"activation": "jax.nn.mish"})) - model = result["model"] - # a dropped key falls back to the class default (swish) instead of erroring - assert model.activation is jax.nn.mish, "activation was silently dropped" +def test_parse_config_resolves_dotted_values(): + """Function paths like 'jax.nn.mish' (and the 'jax._src.nn.functions.silu' + that old configs contain from the aliasing bug) must resolve to the actual + function, not silently vanish into class defaults.""" + config = make_config(arguments_overrides={"architecture": "uvit"}) + config["model"]["activation"] = "jax.nn.mish" + result = parse_config(config) + assert result["model"].activation is jax.nn.mish + + config = make_config(arguments_overrides={"architecture": "uvit"}) + config["model"]["activation"] = "jax._src.nn.functions.silu" + result = parse_config(config) + assert result["model"].activation is jax.nn.silu + + +def test_training_and_inference_share_schedule_presets(): + """--noise_schedule must mean the same thing at train and inference time. + Both sides now build from get_diffusion_preset.""" + from flaxdiff.predictors import get_diffusion_preset + + train, sample, transform = get_diffusion_preset("edm") + assert type(train).__name__ == "EDMNoiseScheduler" + assert type(sample).__name__ == "KarrasVENoiseScheduler" + assert type(transform).__name__ == "KarrasPredictionTransform" + + train, sample, transform = get_diffusion_preset("cosine") + assert type(train).__name__ == "CosineNoiseScheduler" + assert type(sample).__name__ == "CosineNoiseScheduler" + assert type(transform).__name__ == "VPredictionTransform" + + # parse_config must agree with the preset for every name + for name in ("edm", "karras", "cosine"): + _, sample, transform = get_diffusion_preset(name) + result = parse_config(make_config(arguments_overrides={"noise_schedule": name})) + assert type(result["noise_schedule"]) is type(sample) + assert type(result["prediction_transform"]) is type(transform) diff --git a/tests/test_models.py b/tests/test_models.py index c86f173..7747b49 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -34,7 +34,7 @@ def test_unet_forward(rng): model = Unet( emb_features=64, feature_depths=[16, 32], - attention_configs=[None, {"heads": 2, "dtype": jnp.float32, "flash_attention": False, + attention_configs=[None, {"heads": 2, "dtype": jnp.float32, "use_projection": False, "use_self_and_cross": False}], num_res_blocks=1, num_middle_res_blocks=1, @@ -53,7 +53,6 @@ def test_simple_dit_forward(rng): assert out.dtype == jnp.float32 -@pytest.mark.xfail(strict=True, reason="bug: DiT final_proj keeps compute dtype, bf16 output vs fp32 loss") def test_simple_dit_bf16_outputs_fp32(rng): model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2, dtype=jnp.bfloat16) @@ -101,17 +100,40 @@ def test_hybrid_dit_forward(rng, kwargs): assert out.shape == x.shape -@pytest.mark.xfail(strict=True, reason="bug: UViT applies the hilbert permutation twice") -def test_uvit_hilbert_matches_raster_information(rng): - """A zero-layer sanity check: with the permutation applied and inverted once, - patch content must land back at its own spatial position. UViT permutes twice - on the way in but inverts once on the way out, scrambling the output.""" - from flaxdiff.models.hilbert import hilbert_patchify, hilbert_unpatchify, hilbert_indices, inverse_permutation +def test_uvit_hilbert_forward(rng): + """UViT used to apply the hilbert permutation twice on the way in but + invert it only once on the way out, scrambling every output spatially.""" + model = UViT(patch_size=4, emb_features=64, num_layers=4, num_heads=2, use_hilbert=True) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape - x = jax.random.normal(rng, (1, 16, 16, 3)) + +def test_hilbert_patchify_roundtrip(rng): + """patchify returns patches in hilbert order plus the inverse permutation; + unpatchify with that permutation must be an exact identity.""" + from flaxdiff.models.hilbert import hilbert_patchify, hilbert_unpatchify + + x = jax.random.normal(rng, (2, 16, 16, 3)) patches, inv_idx = hilbert_patchify(x, 4) - # UViT's forward applies idx again on the already-permuted patches - idx = hilbert_indices(4, 4) - double_permuted = patches[:, idx, :] - rec = hilbert_unpatchify(double_permuted, inv_idx, 4, 16, 16, 3) + rec = hilbert_unpatchify(patches, inv_idx, 4, 16, 16, 3) assert jnp.allclose(rec, x) + + +def test_dropout_is_active_in_train_mode(rng): + """--dropout_rate used to be threaded into every block and applied nowhere.""" + model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, + mlp_ratio=2, dropout_rate=0.5) + x, temb, textcontext = small_inputs(rng) + params = model.init(rng, x, temb, textcontext) + # nudge every param off init: the zero-initialized adaLN gates and final + # projection make a fresh DiT output exactly zero, hiding dropout + params = jax.tree.map(lambda p: p + 0.02, params) + + d0 = model.apply(params, x, temb, textcontext, train=True, rngs={"dropout": jax.random.PRNGKey(1)}) + d1 = model.apply(params, x, temb, textcontext, train=True, rngs={"dropout": jax.random.PRNGKey(2)}) + assert not jnp.allclose(d0, d1), "different dropout rngs must give different outputs" + + e0 = model.apply(params, x, temb, textcontext) + e1 = model.apply(params, x, temb, textcontext) + assert jnp.allclose(e0, e1), "eval mode must be deterministic" diff --git a/tests/test_samplers.py b/tests/test_samplers.py index 0e9eac0..9156c91 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -105,7 +105,6 @@ def generate(model, sampler, **kwargs): num_samples=256, resolution=8, diffusion_steps=100, - start_step=1000, rngstate=RandomMarkovState(jax.random.PRNGKey(2)), ) defaults.update(kwargs) @@ -126,20 +125,14 @@ def test_karras_sampler_converges(sampler_class): assert_gaussian_stats(samples) -@pytest.mark.xfail(strict=True, reason="bug: default start_step uses max_timesteps but the step domain is hardcoded to [0, 1000]") def test_karras_sampler_default_start_step(): - """With no explicit start_step the sampler must still denoise from sigma_max. - Today max_timesteps=1 collapses the schedule to a single no-op step and - the 'samples' are just the initial noise.""" + """With no explicit start_step the sampler must denoise from sigma_max.""" model, sampler = make_karras_sampler(EulerSampler) samples = generate(model, sampler, start_step=None) assert_gaussian_stats(samples) -@pytest.mark.xfail(strict=True, reason="bug: sample_model reshapes rates with the default (-1,1,1,1), breaking 5D video shapes for every sampler") def test_ddim_video_samples(): - """DDIM's own take_next_step is 5D-safe, but the shared sample_model wrapper - still scales the input with 4D-reshaped rates, so video is broken everywhere.""" model, sampler = make_karras_sampler(DDIMSampler) params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,))) samples = sampler.generate_samples( @@ -148,14 +141,12 @@ def test_ddim_video_samples(): resolution=8, sequence_length=3, diffusion_steps=50, - start_step=1000, rngstate=RandomMarkovState(jax.random.PRNGKey(2)), ) assert samples.shape == (2, 3, 8, 8, 3) assert jnp.all(jnp.isfinite(samples)) -@pytest.mark.xfail(strict=True, reason="bug: samplers reshape rates with the default (-1,1,1,1), breaking 5D video shapes") def test_euler_video_samples(): model, sampler = make_karras_sampler(EulerSampler) params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,))) @@ -165,34 +156,29 @@ def test_euler_video_samples(): resolution=8, sequence_length=3, diffusion_steps=50, - start_step=1000, rngstate=RandomMarkovState(jax.random.PRNGKey(2)), ) assert samples.shape == (2, 3, 8, 8, 3) -@pytest.mark.xfail(strict=True, reason="bug: DDPMSampler casts the step array through int(), crashing for batch > 1") def test_ddpm_sampler_converges(): model, sampler = make_vp_sampler(DDPMSampler) samples = generate(model, sampler, diffusion_steps=1000) assert_gaussian_stats(samples) -@pytest.mark.xfail(strict=True, reason="bug: DDIM eta > 0 calls .sqrt() on a jax array and uses the wrong variance split") def test_ddim_eta_converges(): model, sampler = make_vp_sampler(DDIMSampler, eta=0.5) samples = generate(model, sampler) assert_gaussian_stats(samples) -@pytest.mark.xfail(strict=True, reason="bug: RK4Sampler jits over its python-callable argument") def test_rk4_sampler_runs(): model, sampler = make_karras_sampler(RK4Sampler) samples = generate(model, sampler, diffusion_steps=20) assert jnp.all(jnp.isfinite(samples)) -@pytest.mark.xfail(strict=True, reason="bug: MultiStepDPM keeps stale derivative history across generate calls") def test_multistep_dpm_reentrant(): model, sampler = make_karras_sampler(MultiStepDPM) first = generate(model, sampler) @@ -200,9 +186,3 @@ def test_multistep_dpm_reentrant(): assert jnp.allclose(first, second, atol=1e-5) -@pytest.mark.parametrize("spacing", ["quadratic", "karras", "exponential"]) -@pytest.mark.xfail(strict=True, reason="bug: non-linear timestep spacings produce log(0) or duplicate int16 steps (dt=0)") -def test_timestep_spacing_produces_finite_samples(spacing): - model, sampler = make_karras_sampler(EulerSampler, timestep_spacing=spacing) - samples = generate(model, sampler, diffusion_steps=50) - assert_gaussian_stats(samples, tol=0.1) diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py index b54cc54..44e05fc 100644 --- a/tests/test_schedulers.py +++ b/tests/test_schedulers.py @@ -66,7 +66,6 @@ def test_karras_weights_match_edm_lambda(): assert jnp.allclose(got, expected, rtol=1e-3) -@pytest.mark.xfail(strict=True, reason="bug: epsilon guard distorts the weight near sigma_min") def test_karras_weights_at_sigma_min(): schedule = KarrasVENoiseScheduler(1, sigma_min=0.002, sigma_max=80, rho=7, sigma_data=0.5) sigma = jnp.array([0.002]) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index f8d86f6..a33c15f 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -39,11 +39,17 @@ def batch_iterator(): yield {"image": jnp.ones((8, 4)), "label": jnp.zeros((8, 2))} -@pytest.mark.xfail(strict=True, reason="bug: self.wandb is only assigned when wandb_config is set, training without wandb crashes") def test_fit_without_wandb(tmp_path): trainer = make_trainer(tmp_path) data = {"train": batch_iterator, "train_len": 32} - trainer.fit(data, train_steps_per_epoch=4, epochs=1, val_steps_per_epoch=0) + state = trainer.fit(data, train_steps_per_epoch=4, epochs=2, val_steps_per_epoch=0) + # the tiny regression must actually learn something + final_loss = float(jnp.mean(optax.l2_loss( + state.apply_fn(state.params, jnp.ones((8, 4))), jnp.zeros((8, 2))))) + initial = make_trainer(tmp_path, name="fresh") + initial_loss = float(jnp.mean(optax.l2_loss( + initial.state.apply_fn(initial.state.params, jnp.ones((8, 4))), jnp.zeros((8, 2))))) + assert final_loss < initial_loss def test_save_writes_checkpoint(tmp_path): @@ -53,7 +59,6 @@ def test_save_writes_checkpoint(tmp_path): assert trainer.checkpointer.latest_step() == 1 -@pytest.mark.xfail(strict=True, reason="bug: restore rebuilds the TrainState with tx.init, discarding opt_state and the step counter") def test_restore_preserves_optimizer_state(tmp_path): trainer = make_trainer(tmp_path) diff --git a/training.py b/training.py index 3b6606c..834380a 100644 --- a/training.py +++ b/training.py @@ -11,12 +11,7 @@ from flaxdiff.models.simple_dit import SimpleDiT from flaxdiff.models.simple_mmdit import SimpleMMDiT, HierarchicalMMDiT from flaxdiff.models.ssm_dit import HybridSSMAttentionDiT -try: - import jax.experimental.pallas.ops.tpu.flash_attention -except (ImportError, ModuleNotFoundError): - pass -from flaxdiff.predictors import VPredictionTransform, EpsilonPredictionTransform, DiffusionPredictionTransform, DirectPredictionTransform, KarrasPredictionTransform -from flaxdiff.schedulers import CosineNoiseScheduler, NoiseScheduler, GeneralizedNoiseScheduler, KarrasVENoiseScheduler, EDMNoiseScheduler +from flaxdiff.predictors import get_diffusion_preset from diffusers import FlaxUNet2DConditionModel from flaxdiff.samplers.euler import EulerAncestralSampler import struct as st @@ -30,6 +25,7 @@ from datetime import datetime import json +import hashlib # For CLIP import argparse from dataclasses import dataclass @@ -139,7 +135,6 @@ def boolean_string(s): parser.add_argument('--emb_features', type=int, default=256, help='Embedding features') parser.add_argument('--feature_depths', type=int, nargs='+', default=[64, 128, 256, 512], help='Feature depths') parser.add_argument('--attention_heads', type=int, default=8, help='Number of attention heads') -parser.add_argument('--flash_attention', type=boolean_string, default=False, help='Use Flash Attention') parser.add_argument('--use_projection', type=boolean_string, default=False, help='Use projection') parser.add_argument('--use_self_and_cross', type=boolean_string, default=True, help='Use self and cross attention') parser.add_argument('--only_pure_attention', type=boolean_string, default=True, help='Use only pure attention or proper transformer in the attention blocks') @@ -314,14 +309,14 @@ def main(args): if args.attention_heads > 0: attention_configs += [ { - "heads": args.attention_heads, "dtype": DTYPE, "flash_attention": args.flash_attention, + "heads": args.attention_heads, "dtype": DTYPE, "use_projection": args.use_projection, "use_self_and_cross": args.use_self_and_cross, "only_pure_attention": args.only_pure_attention, }, ] * (len(args.feature_depths) - 2) attention_configs += [ { - "heads": args.attention_heads, "dtype": DTYPE, "flash_attention": False, + "heads": args.attention_heads, "dtype": DTYPE, "use_projection": False, "use_self_and_cross": args.use_self_and_cross, "only_pure_attention": args.only_pure_attention }, @@ -370,9 +365,7 @@ def main(args): "emb_features": args.emb_features, "dtype": DTYPE, "precision": PRECISION, - "activation": args.activation, "output_channels": INPUT_CHANNELS, - "norm_groups": args.norm_groups, } @@ -385,6 +378,8 @@ def main(args): "num_res_blocks": args.num_res_blocks, "num_middle_res_blocks": args.num_middle_res_blocks, "named_norms": args.named_norms, + "activation": args.activation, + "norm_groups": args.norm_groups, }, }, "uvit": { @@ -393,9 +388,9 @@ def main(args): "patch_size": args.patch_size, "num_layers": args.num_layers, "num_heads": args.num_heads, - "dropout_rate": args.dropout_rate, "add_residualblock_output": args.add_residualblock_output, - "use_flash_attention": args.flash_attention, + "activation": args.activation, + "norm_groups": args.norm_groups, "use_self_and_cross": args.use_self_and_cross, "use_hilbert": use_hilbert, }, @@ -407,7 +402,6 @@ def main(args): "num_layers": args.num_layers, "num_heads": args.num_heads, "dropout_rate": args.dropout_rate, - "use_flash_attention": args.flash_attention, "mlp_ratio": args.mlp_ratio, "use_hilbert": use_hilbert, }, @@ -419,7 +413,6 @@ def main(args): "num_layers": args.num_layers, "num_heads": args.num_heads, "dropout_rate": args.dropout_rate, - "use_flash_attention": args.flash_attention, "mlp_ratio": args.mlp_ratio, "use_hilbert": use_hilbert, "use_zigzag": use_zigzag, @@ -432,7 +425,6 @@ def main(args): "num_layers": args.num_layers, "num_heads": args.num_heads, "dropout_rate": args.dropout_rate, - "use_flash_attention": args.flash_attention, "mlp_ratio": args.mlp_ratio, "use_hilbert": use_hilbert, }, @@ -445,7 +437,6 @@ def main(args): "num_layers": (args.num_layers // 3, args.num_layers // 2, args.num_layers), # Default layers per stage "num_heads": (args.num_heads - 2, args.num_heads, args.num_heads + 2), # Default heads per stage "dropout_rate": args.dropout_rate, - "use_flash_attention": args.flash_attention, "mlp_ratio": args.mlp_ratio, "use_hilbert": use_hilbert, }, @@ -457,7 +448,6 @@ def main(args): "num_layers": args.num_layers, "num_heads": args.num_heads, "dropout_rate": args.dropout_rate, - "use_flash_attention": args.flash_attention, "mlp_ratio": args.mlp_ratio, "use_hilbert": use_hilbert, "use_zigzag": use_zigzag, @@ -476,7 +466,6 @@ def main(args): "block_out_channels":args.feature_depths, "cross_attention_dim":args.emb_features, "dtype": DTYPE, - "use_memory_efficient_attention": args.flash_attention, "attention_head_dim": args.attention_heads, "only_cross_attention": not args.use_self_and_cross, }, @@ -490,7 +479,8 @@ def main(args): model_config['emb_features'] = 768 sorted_args_json = json.dumps(vars(args), sort_keys=True) - arguments_hash = hash(sorted_args_json) + # hash() is randomized per process; identical args must map to the same experiment + arguments_hash = hashlib.sha256(sorted_args_json.encode()).hexdigest()[:16] text_encoder = defaultTextEncodeModel() @@ -543,10 +533,7 @@ def main(args): batches = batches if args.steps_per_epoch is None else args.steps_per_epoch - cosine_schedule = CosineNoiseScheduler(1000, beta_end=1) - karas_ve_schedule = KarrasVENoiseScheduler( - 1, sigma_max=80, rho=7, sigma_data=0.5) - edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset(args.noise_schedule) if args.experiment_name is not None: experiment_name = args.experiment_name @@ -577,10 +564,14 @@ def main(args): print("Experiment_Name:", experiment_name) - if 'activation' in model_config: - model_config['activation'] = ACTIVATION_MAP[model_config['activation']] - - model = model_architecture(**model_config) + # Map the activation string to the actual function only for construction. + # model_config is aliased into CONFIG (already logged to wandb), so mutating + # it here used to leak a live function object into the stored config. + construction_config = dict(model_config) + if 'activation' in construction_config: + construction_config['activation'] = ACTIVATION_MAP[construction_config['activation']] + + model = model_architecture(**construction_config) # If using the Diffusers UNet, we need to wrap it if 'diffusers' in architecture_name: @@ -618,11 +609,10 @@ def main(args): trainer = GeneralDiffusionTrainer( model, optimizer=solver, input_config=input_config, - noise_schedule=edm_schedule, + noise_schedule=train_schedule, rngs=jax.random.PRNGKey(4), name=experiment_name, - model_output_transform=KarrasPredictionTransform( - sigma_data=edm_schedule.sigma_data), + model_output_transform=prediction_transform, load_from_checkpoint=args.load_from_checkpoint, wandb_config=wandb_config, distributed_training=args.distributed_training, @@ -656,7 +646,7 @@ def get_val_dataset(): training_steps_per_epoch=batches, epochs=CONFIG['epochs'], sampler_class=EulerAncestralSampler, - sampling_noise_schedule=karas_ve_schedule, + sampling_noise_schedule=sampling_schedule, val_steps_per_epoch=args.val_steps_per_epoch, )