From 9c0222bdd15768675f1fab472365b4588f356960 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 11:28:48 -0400 Subject: [PATCH 1/4] feat(models): fused attention done right, one param tree across implementations Restores the flash-attention capability that Phase 1 deleted, built properly this time. One helper (scaled_dot_product_attention) backs every attention module: the flax reference path by default, jax.nn's fused entrypoint ('xla'/'cudnn', dispatching to the cudnn flash kernel on supported GPUs), and the pallas TPU flash kernel with the 1/sqrt(d) scale passed explicitly - the deleted EfficientAttention passed none, inflating logits 8x at head_dim 64, and used a different param tree so its checkpoints were unusable on the normal path. The projections never change with the implementation, so checkpoints are interchangeable across hardware; attention_impl is exposed on every model (DiT family, MM-DiT, UViT/UDiT, UNet) and as --attention_impl in training.py. A parity test pins reference-vs-fused agreement on the same params. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/attention.py | 51 +++++++++++++++++++++++++++++---- flaxdiff/models/dit_common.py | 2 ++ flaxdiff/models/simple_dit.py | 2 ++ flaxdiff/models/simple_mmdit.py | 14 ++++++--- flaxdiff/models/simple_unet.py | 7 +++-- flaxdiff/models/simple_vit.py | 8 ++++++ flaxdiff/models/ssm_dit.py | 2 ++ flaxdiff/models/vit_common.py | 10 +++---- tests/test_models.py | 14 +++++++++ training.py | 4 +++ 10 files changed, 96 insertions(+), 18 deletions(-) diff --git a/flaxdiff/models/attention.py b/flaxdiff/models/attention.py index 1b50812..bc878d8 100644 --- a/flaxdiff/models/attention.py +++ b/flaxdiff/models/attention.py @@ -12,6 +12,39 @@ import math from .common import kernel_init +def scaled_dot_product_attention(query, key, value, dtype=None, precision=None, + force_fp32_for_softmax=True, implementation=None): + """The one attention kernel path for every attention module. + + Inputs are [B, S, H, D]. The param trees of the callers never change with + the implementation, so checkpoints are interchangeable across hardware: + + - None: flax reference attention (einsum + softmax). Runs everywhere and + honors force_fp32_for_softmax; the default. + - 'xla' / 'cudnn': jax.nn.dot_product_attention, which dispatches to the + fused cudnn flash kernel on supported GPUs. + - 'tpu': the pallas TPU flash kernel, with the 1/sqrt(d) scale passed + explicitly (the deleted EfficientAttention passed none, which inflated + the logits by sqrt(d) and made its checkpoints poisonous). + """ + if implementation is None: + return nn.dot_product_attention( + query, key, value, dtype=dtype, broadcast_dropout=False, + dropout_rng=None, precision=precision, + force_fp32_for_softmax=force_fp32_for_softmax, deterministic=True) + if implementation in ('xla', 'cudnn'): + return jax.nn.dot_product_attention(query, key, value, implementation=implementation) + if implementation == 'tpu': + from jax.experimental.pallas.ops.tpu.flash_attention import flash_attention + # pallas wants [B, H, S, D] + q = jnp.moveaxis(query, -2, -3) + k = jnp.moveaxis(key, -2, -3) + v = jnp.moveaxis(value, -2, -3) + out = flash_attention(q, k, v, None, sm_scale=1.0 / math.sqrt(query.shape[-1])) + return jnp.moveaxis(out, -3, -2) + raise ValueError(f"Unknown attention implementation: {implementation}") + + class NormalAttention(nn.Module): """ Simple implementation of the normal attention. @@ -25,6 +58,7 @@ class NormalAttention(nn.Module): # kernel_init: Callable = kernel_init(1.0) force_fp32_for_softmax: bool = True qk_norm: bool = False # RMSNorm on q/k per head (SD3-style bf16 logit safety) + attention_impl: Optional[str] = None # None (reference) | 'xla' | 'cudnn' | 'tpu' def setup(self): inner_dim = self.dim_head * self.heads @@ -73,10 +107,10 @@ def __call__(self, x, context=None): query = self.q_norm(query) key = self.k_norm(key) - hidden_states = nn.dot_product_attention( - query, key, value, dtype=self.dtype, broadcast_dropout=False, - dropout_rng=None, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, - deterministic=True + hidden_states = scaled_dot_product_attention( + query, key, value, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + implementation=self.attention_impl, ) proj = self.proj_attn(hidden_states) proj = proj.reshape(orig_x_shape) @@ -156,6 +190,7 @@ class BasicTransformerBlock(nn.Module): only_pure_attention:bool = False force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-4 + attention_impl: Optional[str] = None def setup(self): attenBlock = NormalAttention @@ -169,7 +204,8 @@ def setup(self): use_bias=self.use_bias, dtype=self.dtype, # kernel_init=self.kernel_init, - force_fp32_for_softmax=self.force_fp32_for_softmax + force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl ) self.attention2 = attenBlock( query_dim=self.query_dim, @@ -180,7 +216,8 @@ def setup(self): use_bias=self.use_bias, dtype=self.dtype, # kernel_init=self.kernel_init, - force_fp32_for_softmax=self.force_fp32_for_softmax + force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl ) self.ff = FlaxFeedForward(dim=self.query_dim, dtype=self.dtype, precision=self.precision) @@ -214,6 +251,7 @@ class TransformerBlock(nn.Module): use_self_and_cross:bool = True only_pure_attention:bool = False force_fp32_for_softmax: bool = True + attention_impl: Optional[str] = None # kernel_init: Callable = kernel_init(1.0) norm_inputs: bool = True explicitly_add_residual: bool = True @@ -255,6 +293,7 @@ def __call__(self, x, context=None): 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, + attention_impl=self.attention_impl, norm_epsilon=self.norm_epsilon # kernel_init=self.kernel_init )(projected_x, context) diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py index bf4ce78..848d1b9 100644 --- a/flaxdiff/models/dit_common.py +++ b/flaxdiff/models/dit_common.py @@ -190,6 +190,7 @@ class ModulatedBlock(nn.Module): norm_epsilon: float = 1e-5 use_gating: bool = True qk_norm: bool = False + attention_impl: Optional[str] = None # None (reference) | 'xla' | 'cudnn' | 'tpu' # ssm mixer options ssm_state_dim: int = 64 bidirectional_ssm: bool = True @@ -218,6 +219,7 @@ def setup(self): precision=self.precision, use_bias=True, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, force_fp32_for_softmax=self.force_fp32_for_softmax, rope_emb=self.rope_emb, ) diff --git a/flaxdiff/models/simple_dit.py b/flaxdiff/models/simple_dit.py index 3aa7542..bc59f09 100644 --- a/flaxdiff/models/simple_dit.py +++ b/flaxdiff/models/simple_dit.py @@ -25,6 +25,7 @@ class SimpleDiT(nn.Module): norm_epsilon: float = 1e-5 learn_sigma: bool = False qk_norm: bool = False + attention_impl: Optional[str] = None use_hilbert: bool = False use_zigzag: bool = False @@ -63,6 +64,7 @@ def setup(self): force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, name=f"dit_block_{i}" ) for i in range(self.num_layers) ] diff --git a/flaxdiff/models/simple_mmdit.py b/flaxdiff/models/simple_mmdit.py index 0b25cb4..6ebc43e 100644 --- a/flaxdiff/models/simple_mmdit.py +++ b/flaxdiff/models/simple_mmdit.py @@ -20,6 +20,7 @@ PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, neutralized_rope_freqs, ) +from .attention import scaled_dot_product_attention from .vit_common import RotaryEmbedding, AdaLNParams, apply_rotary_embedding @@ -39,6 +40,7 @@ class MMDiTBlock(nn.Module): force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 qk_norm: bool = False + attention_impl: Optional[str] = None def setup(self): hidden_features = int(self.features * self.mlp_ratio) @@ -112,10 +114,10 @@ def __call__(self, img, txt, conditioning, freqs_cis, train: bool = False): k = jnp.concatenate([k_t, k_i], axis=1) v = jnp.concatenate([v_t, v_i], axis=1) - attn = nn.dot_product_attention( - q, k, v, dtype=self.dtype, broadcast_dropout=False, dropout_rng=None, - precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, - deterministic=True) + attn = scaled_dot_product_attention( + q, k, v, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + implementation=self.attention_impl) txt_attn, img_attn = attn[:, :S_txt], attn[:, S_txt:] img_attn = self.dropout(self.img_out(img_attn), deterministic=not train) @@ -147,6 +149,7 @@ class SimpleMMDiT(nn.Module): norm_epsilon: float = 1e-5 learn_sigma: bool = False qk_norm: bool = False + attention_impl: Optional[str] = None use_hilbert: bool = False use_zigzag: bool = False @@ -187,6 +190,7 @@ def setup(self): force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, name=f"mmdit_block_{i}" ) for i in range(self.num_layers) ] @@ -332,6 +336,7 @@ class HierarchicalMMDiT(nn.Module): norm_epsilon: float = 1e-5 learn_sigma: bool = False qk_norm: bool = False + attention_impl: Optional[str] = None def setup(self): assert len(self.emb_features) == len(self.num_layers) == len(self.num_heads), \ @@ -382,6 +387,7 @@ def stage_blocks(stage, prefix): force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, name=f"{prefix}_block_stage{stage}_{i}" ) for i in range(self.num_layers[stage]) ] diff --git a/flaxdiff/models/simple_unet.py b/flaxdiff/models/simple_unet.py index 90cc982..e2684ca 100644 --- a/flaxdiff/models/simple_unet.py +++ b/flaxdiff/models/simple_unet.py @@ -20,6 +20,7 @@ class Unet(nn.Module): dtype: Optional[Dtype] = None precision: PrecisionLike = None named_norms: bool = False # This is for backward compatibility reasons; older checkpoints have named norms + attention_impl: Optional[str] = None def setup(self): if self.norm_groups > 0: @@ -72,7 +73,7 @@ def __call__(self, x, temb, textcontext, train: bool = False): named_norms=self.named_norms )(x, temb) 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), + x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), attention_impl=self.attention_impl, dim_head=dim_in // attention_config['heads'], use_projection=attention_config.get("use_projection", False), use_self_and_cross=attention_config.get("use_self_and_cross", True), @@ -112,7 +113,7 @@ def __call__(self, x, temb, textcontext, train: bool = False): named_norms=self.named_norms )(x, temb) 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), + x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), attention_impl=self.attention_impl, dim_head=middle_dim_out // middle_attention['heads'], use_linear_attention=False, use_projection=middle_attention.get("use_projection", False), @@ -157,7 +158,7 @@ def __call__(self, x, temb, textcontext, train: bool = False): named_norms=self.named_norms )(x, temb) 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), + x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), attention_impl=self.attention_impl, dim_head=dim_out // attention_config['heads'], use_projection=attention_config.get("use_projection", False), use_self_and_cross=attention_config.get("use_self_and_cross", True), diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py index 7dc4b2d..1554afb 100644 --- a/flaxdiff/models/simple_vit.py +++ b/flaxdiff/models/simple_vit.py @@ -25,6 +25,7 @@ class UViT(nn.Module): # Passed to TransformerBlock (likely False for UViT) use_self_and_cross: bool = False force_fp32_for_softmax: bool = True # Passed to TransformerBlock + attention_impl: Optional[str] = None # Used in final convs if add_residualblock_output activation: Callable = jax.nn.swish norm_groups: int = 8 @@ -88,6 +89,7 @@ def setup(self): dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl, only_pure_attention=False, norm_inputs=self.norm_inputs, explicitly_add_residual=self.explicitly_add_residual, norm_epsilon=self.norm_epsilon, @@ -101,6 +103,7 @@ def setup(self): dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl, only_pure_attention=False, norm_inputs=self.norm_inputs, explicitly_add_residual=self.explicitly_add_residual, norm_epsilon=self.norm_epsilon, @@ -122,6 +125,7 @@ def setup(self): dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl, only_pure_attention=False, norm_inputs=self.norm_inputs, explicitly_add_residual=self.explicitly_add_residual, norm_epsilon=self.norm_epsilon, @@ -257,6 +261,7 @@ class SimpleUDiT(nn.Module): dtype: Optional[Dtype] = None # e.g., jnp.float32 or jnp.bfloat16 precision: PrecisionLike = None force_fp32_for_softmax: bool = True # Passed to DiTBlock -> RoPEAttention + attention_impl: Optional[str] = None norm_epsilon: float = 1e-5 learn_sigma: bool = False use_hilbert: bool = False @@ -312,6 +317,7 @@ def setup(self): dtype=self.dtype, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, name=f"down_block_{i}" @@ -326,6 +332,7 @@ def setup(self): dtype=self.dtype, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, name="mid_block" @@ -348,6 +355,7 @@ def setup(self): dtype=self.dtype, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, + attention_impl=self.attention_impl, norm_epsilon=self.norm_epsilon, rope_emb=self.rope, name=f"up_block_{i}" diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py index 3ee4b03..749aaf0 100644 --- a/flaxdiff/models/ssm_dit.py +++ b/flaxdiff/models/ssm_dit.py @@ -34,6 +34,7 @@ class HybridSSMAttentionDiT(nn.Module): norm_epsilon: float = 1e-5 learn_sigma: bool = False qk_norm: bool = False + attention_impl: Optional[str] = None use_hilbert: bool = False use_zigzag: bool = False # ZigMa-style serpentine scan block_pattern: Optional[Sequence[str]] = None # e.g., ['ssm','ssm','ssm','attn'] @@ -113,6 +114,7 @@ def setup(self): force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm, + attention_impl=self.attention_impl, name=f"dit_block_{i}" )) self.blocks = blocks diff --git a/flaxdiff/models/vit_common.py b/flaxdiff/models/vit_common.py index 66eab3e..08819d3 100644 --- a/flaxdiff/models/vit_common.py +++ b/flaxdiff/models/vit_common.py @@ -5,7 +5,7 @@ import einops from flax.typing import Dtype, PrecisionLike -from .attention import NormalAttention +from .attention import NormalAttention, scaled_dot_product_attention def unpatchify(x, channels=3, H_P=None, W_P=None): patch_size = int((x.shape[2] // channels) ** 0.5) @@ -158,10 +158,10 @@ def __call__(self, x, context=None, freqs_cis=None): query = einops.rearrange(query, 'b h s d -> b s h d') key = einops.rearrange(key, 'b h s d -> b s h d') - hidden_states = nn.dot_product_attention( - query, key, value, dtype=self.dtype, broadcast_dropout=False, - dropout_rng=None, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, - deterministic=True + hidden_states = scaled_dot_product_attention( + query, key, value, dtype=self.dtype, precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + implementation=self.attention_impl, ) proj = self.proj_attn(hidden_states) diff --git a/tests/test_models.py b/tests/test_models.py index 4df3335..5560e49 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -161,3 +161,17 @@ def test_mmdit_is_dual_stream(rng): out_a = model.apply(params, x, temb, text_a) out_b = model.apply(params, x, temb, text_b) assert not jnp.allclose(out_a, out_b), "text tokens do not reach the image stream" + + +def test_attention_impl_parity(rng): + """Every attention implementation must share one param tree and produce + the same outputs. 'xla' (the jax.nn fused entrypoint) is verifiable on + CPU; cudnn/tpu dispatch to the same wrapper.""" + ref = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + xla = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2, + attention_impl="xla") + x, temb, textcontext = small_inputs(rng) + params = ref.init(rng, x, temb, textcontext) + out_ref = ref.apply(params, x, temb, textcontext) + out_xla = xla.apply(params, x, temb, textcontext) + assert jnp.max(jnp.abs(out_ref - out_xla)) < 1e-4 diff --git a/training.py b/training.py index 39c8614..ab6f25f 100644 --- a/training.py +++ b/training.py @@ -109,6 +109,9 @@ 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('--attention_impl', type=str, default=None, + choices=[None, 'xla', 'cudnn', 'tpu'], + help='Fused attention kernel: None (reference), cudnn (GPU flash), tpu (pallas flash)') 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') @@ -309,6 +312,7 @@ def main(args): "dtype": DTYPE, "precision": PRECISION, "output_channels": INPUT_CHANNELS, + "attention_impl": args.attention_impl, } From 586164308c4583e760e04ed3240a9c72e601265a Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 11:33:02 -0400 Subject: [PATCH 2/4] feat(models): video DiT on the shared backbone Restores the video-modeling capability whose old carrier (the diffusers-derived UNet3D) was deleted: it was never wired into any registry or training path, and its upstream flax blocks no longer exist. The replacement is a factorized spatial-temporal DiT built entirely from the shared machinery - per-frame PatchSequenceEmbed (all scan orders work), a spatial ModulatedBlock over each frame's tokens and a temporal ModulatedBlock over the frame axis per layer, RoPE on the genuine 1D time axis, and the shared modulated output head. Registered as 'video_dit' and wired end to end: DiffusionInputConfig now keeps the time axis for 4-tuple video shapes (it silently dropped T before, so no video model could even initialize through the trainer). Tests: forward shapes, a temporal-mixing check (perturbing frame 0 must change the prediction for frame 2), registry construction, and a 5D end-to-end run through the real trainer. Co-Authored-By: Claude Fable 5 --- flaxdiff/inputs/__init__.py | 5 +- flaxdiff/models/registry.py | 2 + flaxdiff/models/video_dit.py | 123 +++++++++++++++++++++++++++++++++++ tests/test_models.py | 29 +++++++++ tests/test_trainer.py | 33 ++++++++++ 5 files changed, 191 insertions(+), 1 deletion(-) create mode 100644 flaxdiff/models/video_dit.py diff --git a/flaxdiff/inputs/__init__.py b/flaxdiff/inputs/__init__.py index 49c4659..c1a53de 100644 --- a/flaxdiff/inputs/__init__.py +++ b/flaxdiff/inputs/__init__.py @@ -88,6 +88,7 @@ def get_input_shapes( ) -> Dict[str, Tuple[int, ...]]: """Get the shapes of the input data.""" if len(self.sample_data_shape) == 3: + T = None H, W, C = self.sample_data_shape elif len(self.sample_data_shape) == 4: T, H, W, C = self.sample_data_shape @@ -99,8 +100,10 @@ def get_input_shapes( W = W // downscale_factor C = autoencoder.latent_channels + # Video models init against the full (T, H, W, C) sample shape + sample_shape = (H, W, C) if T is None else (T, H, W, C) input_shapes = { - sample_model_key: (H, W, C), + sample_model_key: sample_shape, time_embeddings_model_key: (), } for cond in self.conditions: diff --git a/flaxdiff/models/registry.py b/flaxdiff/models/registry.py index 74b320d..5665b8f 100644 --- a/flaxdiff/models/registry.py +++ b/flaxdiff/models/registry.py @@ -16,6 +16,7 @@ from .simple_dit import SimpleDiT from .simple_mmdit import SimpleMMDiT, HierarchicalMMDiT from .ssm_dit import HybridSSMAttentionDiT +from .video_dit import VideoDiT MODEL_REGISTRY = { 'unet': Unet, @@ -25,6 +26,7 @@ 'simple_mmdit': SimpleMMDiT, 'hierarchical_mmdit': HierarchicalMMDiT, 'hybrid_dit': HybridSSMAttentionDiT, + 'video_dit': VideoDiT, } ARCHITECTURE_SUFFIX_FLAGS = { diff --git a/flaxdiff/models/video_dit.py b/flaxdiff/models/video_dit.py new file mode 100644 index 0000000..8f2f7b6 --- /dev/null +++ b/flaxdiff/models/video_dit.py @@ -0,0 +1,123 @@ +""" +Video DiT with factorized spatial-temporal attention, built from the shared +DiT machinery. This replaces the old diffusers-derived UNet3D, which was +never wired into the registry and whose upstream flax blocks no longer exist. + +Each layer is a spatial ModulatedBlock over the patch tokens of every frame +followed by a temporal ModulatedBlock over the frame axis of every patch +position - the standard factorized design, so compute stays linear in T for +the spatial half and linear in S for the temporal half. +""" + +import jax.numpy as jnp +from flax import linen as nn +from typing import Optional +from flax.typing import Dtype, PrecisionLike + +from .dit_common import ( + PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, + ModulatedBlock, neutralized_rope_freqs, +) +from .vit_common import RotaryEmbedding + + +class VideoDiT(nn.Module): + """Factorized spatial-temporal DiT over (B, T, H, W, C) inputs.""" + output_channels: int = 3 + patch_size: int = 16 + emb_features: int = 768 + num_layers: int = 12 # Each layer is one spatial + one temporal block + num_heads: int = 12 + mlp_ratio: int = 4 + dropout_rate: float = 0.0 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + force_fp32_for_softmax: bool = True + norm_epsilon: float = 1e-5 + learn_sigma: bool = False + qk_norm: bool = False + attention_impl: Optional[str] = None + use_hilbert: bool = False + use_zigzag: bool = False + + @property + def scan_order(self): + assert not (self.use_hilbert and self.use_zigzag), \ + "use_hilbert and use_zigzag are mutually exclusive" + return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster' + + def setup(self): + self.embed = PatchSequenceEmbed( + patch_size=self.patch_size, + emb_features=self.emb_features, + scan_order=self.scan_order, + dtype=self.dtype, + precision=self.precision, + ) + self.conditioning = ConditioningEmbed( + emb_features=self.emb_features, + mlp_ratio=self.mlp_ratio, + dtype=self.dtype, + precision=self.precision, + ) + dim_head = self.emb_features // self.num_heads + self.spatial_rope = RotaryEmbedding( + dim=dim_head, max_seq_len=4096, dtype=self.dtype, name="spatial_rope") + self.temporal_rope = RotaryEmbedding( + dim=dim_head, max_seq_len=1024, dtype=self.dtype, name="temporal_rope") + + def block(name): + return ModulatedBlock( + features=self.emb_features, + num_heads=self.num_heads, + mixer='attention', + mlp_ratio=self.mlp_ratio, + dropout_rate=self.dropout_rate, + dtype=self.dtype, + precision=self.precision, + force_fp32_for_softmax=self.force_fp32_for_softmax, + norm_epsilon=self.norm_epsilon, + qk_norm=self.qk_norm, + attention_impl=self.attention_impl, + name=name, + ) + + self.spatial_blocks = [block(f"spatial_block_{i}") for i in range(self.num_layers)] + self.temporal_blocks = [block(f"temporal_block_{i}") for i in range(self.num_layers)] + + self.output = PatchSequenceOutput( + patch_size=self.patch_size, + output_channels=self.output_channels, + learn_sigma=self.learn_sigma, + norm_epsilon=self.norm_epsilon, + dtype=self.dtype, + precision=self.precision, + ) + + @nn.compact + def __call__(self, x, temb, textcontext=None, train: bool = False): + B, T, H, W, C = x.shape + + # Per-frame patchify; the permutation is identical for every frame + frames = x.reshape(B * T, H, W, C) + tokens, inv_idx = self.embed(frames) # [B*T, S, F] + S = tokens.shape[1] + + cond_emb = self.conditioning(temb, textcontext) # [B, F] + cond_spatial = jnp.repeat(cond_emb, T, axis=0) # [B*T, F] + cond_temporal = jnp.repeat(cond_emb, S, axis=0) # [B*S, F] + + freqs_spatial = neutralized_rope_freqs(self.spatial_rope, S, self.scan_order) + # Time is a genuine 1D axis, RoPE applies directly + freqs_temporal = self.temporal_rope(seq_len=T) + + for spatial, temporal in zip(self.spatial_blocks, self.temporal_blocks): + tokens = spatial(tokens, conditioning=cond_spatial, freqs_cis=freqs_spatial, train=train) + # [B*T, S, F] -> [B*S, T, F] + tokens = tokens.reshape(B, T, S, -1).transpose(0, 2, 1, 3).reshape(B * S, T, -1) + tokens = temporal(tokens, conditioning=cond_temporal, freqs_cis=freqs_temporal, train=train) + # back to [B*T, S, F] + tokens = tokens.reshape(B, S, T, -1).transpose(0, 2, 1, 3).reshape(B * T, S, -1) + + out_frames = self.output(tokens, inv_idx, H, W) # [B*T, H, W, C] + return out_frames.reshape(B, T, H, W, self.output_channels) diff --git a/tests/test_models.py b/tests/test_models.py index 5560e49..d8c18db 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -175,3 +175,32 @@ def test_attention_impl_parity(rng): out_ref = ref.apply(params, x, temb, textcontext) out_xla = xla.apply(params, x, temb, textcontext) assert jnp.max(jnp.abs(out_ref - out_xla)) < 1e-4 + + +def test_video_dit_forward(rng): + """Factorized ST video model over (B, T, H, W, C), the replacement for + the never-wired diffusers-derived UNet3D.""" + from flaxdiff.models.video_dit import VideoDiT + model = VideoDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + x = jax.random.normal(rng, (2, 3, 16, 16, 3)) + temb = jnp.ones((2,)) + textcontext = jnp.ones((2, 77, 768), dtype=jnp.float32) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + assert jnp.all(jnp.isfinite(out)) + + +def test_video_dit_temporal_mixing(rng): + """Temporal blocks must actually mix across frames: perturbing frame 0 + must change the prediction for frame 2.""" + from flaxdiff.models.video_dit import VideoDiT + model = VideoDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + x = jax.random.normal(rng, (1, 3, 16, 16, 3)) + temb = jnp.ones((1,)) + params = model.init(rng, x, temb, None) + params = jax.tree.map(lambda p: p + 0.02, params) + + out_a = model.apply(params, x, temb, None) + x_perturbed = x.at[:, 0].add(1.0) + out_b = model.apply(params, x_perturbed, temb, None) + assert not jnp.allclose(out_a[:, 2], out_b[:, 2]), "no information flow across frames" diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 630436a..0f21004 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -80,3 +80,36 @@ def test_restore_preserves_optimizer_state(tmp_path): old_ema = jax.tree.leaves(trainer.state.ema_params) new_ema = jax.tree.leaves(restored.state.ema_params) assert all(np.allclose(a, b) for a, b in zip(old_ema, new_ema)), "ema params were reset" + + +def test_fit_video_model(tmp_path): + """Video end to end: 5D batches through the real trainer and the video DiT.""" + from flaxdiff.models.video_dit import VideoDiT + + train_schedule, _, transform = get_diffusion_preset("edm") + trainer = GeneralDiffusionTrainer( + model=VideoDiT(patch_size=4, emb_features=16, num_layers=1, num_heads=2, mlp_ratio=1), + optimizer=optax.adam(1e-3), + noise_schedule=train_schedule, + model_output_transform=transform, + input_config=DiffusionInputConfig( + sample_data_key="video", + sample_data_shape=(3, RES, RES, 3), + conditions=[], + ), + rngs=jax.random.PRNGKey(0), + name="video-smoke", + wandb_config=None, + distributed_training=False, + checkpoint_base_path=str(tmp_path), + ) + + def video_batches(): + frames = np.tile(np.linspace(0, 255, RES, dtype=np.float32)[None, None, :, None, None], + (4, 3, 1, RES, 3)) + while True: + yield {"video": jnp.asarray(frames)} + + data = {"train": video_batches, "train_len": 16, "local_batch_size": 4} + state = trainer.fit(data, training_steps_per_epoch=2, epochs=1, val_steps_per_epoch=0) + assert int(state.step) == 2 From 0827b011c7ea8359705891ce98e556b62a9185b7 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 11:56:34 -0400 Subject: [PATCH 3/4] feat(models): UNet3D rebuilt from our own blocks, inflatable from trained 2D checkpoints The old diffusers-derived UNet3D could never have loaded our trained UNets: different blocks, different param tree, and its upstream flax modules no longer exist. This one mirrors our Unet exactly - same blocks, same explicit module names, same auto-name ordering - so the spatial param paths are identical, and inflate_unet_params drops a trained 2D checkpoint straight into the 3D tree. The interleaved temporal attention blocks (RoPE over the frame axis at every resolution level) are zero-initialized, so an inflated model reproduces the image model frame by frame and training only has to learn motion - the AnimateDiff recipe, applied to our own pretrained models. Verified by tests: an inflated UNet3D matches the 2D Unet per-frame to 1e-5 on the same inputs, temporal mixing activates once the gate moves off zero, and the registry builds it ('unet_3d', wired in training.py alongside 'video_dit'). Co-Authored-By: Claude Fable 5 --- flaxdiff/models/registry.py | 2 + flaxdiff/models/unet_3d.py | 315 ++++++++++++++++++++++++++++++++++++ tests/test_config.py | 3 + tests/test_models.py | 50 ++++++ training.py | 22 +++ 5 files changed, 392 insertions(+) create mode 100644 flaxdiff/models/unet_3d.py diff --git a/flaxdiff/models/registry.py b/flaxdiff/models/registry.py index 5665b8f..2923b9a 100644 --- a/flaxdiff/models/registry.py +++ b/flaxdiff/models/registry.py @@ -17,6 +17,7 @@ from .simple_mmdit import SimpleMMDiT, HierarchicalMMDiT from .ssm_dit import HybridSSMAttentionDiT from .video_dit import VideoDiT +from .unet_3d import UNet3D MODEL_REGISTRY = { 'unet': Unet, @@ -27,6 +28,7 @@ 'hierarchical_mmdit': HierarchicalMMDiT, 'hybrid_dit': HybridSSMAttentionDiT, 'video_dit': VideoDiT, + 'unet_3d': UNet3D, } ARCHITECTURE_SUFFIX_FLAGS = { diff --git a/flaxdiff/models/unet_3d.py b/flaxdiff/models/unet_3d.py new file mode 100644 index 0000000..781795f --- /dev/null +++ b/flaxdiff/models/unet_3d.py @@ -0,0 +1,315 @@ +""" +UNet3D: the 2D UNet inflated for video, AnimateDiff-style. + +The forward mirrors Unet exactly - same blocks, same explicit module names, +same auto-name ordering - processing every frame through the spatial path, +with zero-initialized temporal attention blocks interleaved at each +resolution level. Zero init means a freshly inflated model reproduces the 2D +UNet frame by frame, so a pretrained image checkpoint (inflate_unet_params) +is the starting point and training only has to learn motion. + +This replaces the old diffusers-derived UNet3D, which was never wired into +any training path and could not load our checkpoints (different blocks, +different param tree, upstream deleted). +""" + +import jax +import jax.numpy as jnp +from flax import linen as nn +from flax.typing import Dtype, PrecisionLike +from typing import Callable, Optional +from functools import partial + +from .common import ConvLayer, Downsample, Upsample, FourierEmbedding, TimeProjection, ResidualBlock +from .attention import TransformerBlock +from .vit_common import RotaryEmbedding, RoPEAttention + + +class TemporalBlock(nn.Module): + """Temporal self-attention over the frame axis at every spatial position. + + The output projection is zero-initialized, so the block is an exact + identity at init and an inflated model starts as the per-frame 2D model. + """ + features: int + heads: int = 8 + norm_epsilon: float = 1e-5 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + @nn.compact + def __call__(self, x, frames: int): + # x: [B*T, H, W, C] -> tokens over T per spatial position + BT, H, W, C = x.shape + B = BT // frames + h = x.reshape(B, frames, H * W, C) + h = h.transpose(0, 2, 1, 3).reshape(B * H * W, frames, C) + + h = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)(h) + rope = RotaryEmbedding( + dim=C // self.heads, max_seq_len=1024, dtype=self.dtype, name="temporal_rope") + h = RoPEAttention( + query_dim=C, + heads=self.heads, + dim_head=C // self.heads, + dtype=self.dtype, + precision=self.precision, + use_bias=True, + rope_emb=rope, + name="temporal_attention", + )(h, freqs_cis=rope(frames)) + # zero-init gate: identity at init, so inflation preserves the 2D model + h = nn.Dense( + features=C, dtype=self.dtype, precision=self.precision, + kernel_init=nn.initializers.zeros, name="temporal_out")(h) + + h = h.reshape(B, H * W, frames, C).transpose(0, 2, 1, 3).reshape(BT, H, W, C) + return x + h + + +class UNet3D(nn.Module): + """Video UNet over (B, T, H, W, C): the 2D Unet forward per frame, with a + TemporalBlock at every resolution level. Spatial param paths are + identical to Unet, so 2D checkpoints inflate directly.""" + output_channels:int=3 + emb_features:int=64*4 + feature_depths:list=(64, 128, 256, 512) + attention_configs:list=({"heads":8}, {"heads":8}, {"heads":8}, {"heads":8}) + num_res_blocks:int=2 + num_middle_res_blocks:int=1 + activation:Callable = jax.nn.swish + norm_groups:int=8 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + named_norms: bool = False + attention_impl: Optional[str] = None + temporal_heads: int = 8 + + def setup(self): + if self.norm_groups > 0: + norm = partial(nn.GroupNorm, self.norm_groups) + self.conv_out_norm = norm(name="GroupNorm_0") if self.named_norms else norm() + else: + norm = partial(nn.RMSNorm, 1e-5) + self.conv_out_norm = norm() + + @nn.compact + def __call__(self, x, temb, textcontext, train: bool = False): + B, T, H, W, C = x.shape + x = x.reshape(B * T, H, W, C) + temb = jnp.repeat(temb, T, axis=0) + if textcontext is not None: + textcontext = jnp.repeat(textcontext, T, axis=0) + + temb = FourierEmbedding(features=self.emb_features)(temb) + temb = TimeProjection(features=self.emb_features)(temb) + + feature_depths = self.feature_depths + attention_configs = self.attention_configs + + conv_type = up_conv_type = down_conv_type = middle_conv_type = "conv" + + x = ConvLayer( + conv_type, + features=self.feature_depths[0], + kernel_size=(3, 3), + strides=(1, 1), + dtype=self.dtype, + precision=self.precision + )(x) + downs = [x] + + # Downscaling blocks + for i, (dim_out, attention_config) in enumerate(zip(feature_depths, attention_configs)): + dim_in = x.shape[-1] + for j in range(self.num_res_blocks): + x = ResidualBlock( + down_conv_type, + name=f"down_{i}_residual_{j}", + features=dim_in, + kernel_size=(3, 3), + strides=(1, 1), + activation=self.activation, + norm_groups=self.norm_groups, + dtype=self.dtype, + precision=self.precision, + named_norms=self.named_norms + )(x, temb) + if attention_config is not None and j == self.num_res_blocks - 1: + x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), attention_impl=self.attention_impl, + dim_head=dim_in // attention_config['heads'], + 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), + only_pure_attention=attention_config.get("only_pure_attention", True), + force_fp32_for_softmax=attention_config.get("force_fp32_for_softmax", False), + norm_inputs=attention_config.get("norm_inputs", True), + explicitly_add_residual=attention_config.get("explicitly_add_residual", True), + name=f"down_{i}_attention_{j}")(x, textcontext) + downs.append(x) + x = TemporalBlock( + features=x.shape[-1], + heads=self.temporal_heads, + dtype=self.dtype, + precision=self.precision, + name=f"down_{i}_temporal" + )(x, T) + if i != len(feature_depths) - 1: + x = Downsample( + features=dim_out, + scale=2, + activation=self.activation, + name=f"down_{i}_downsample", + dtype=self.dtype, + precision=self.precision + )(x) + + # Middle Blocks + middle_dim_out = self.feature_depths[-1] + middle_attention = self.attention_configs[-1] + for j in range(self.num_middle_res_blocks): + x = ResidualBlock( + middle_conv_type, + name=f"middle_res1_{j}", + features=middle_dim_out, + kernel_size=(3, 3), + strides=(1, 1), + activation=self.activation, + norm_groups=self.norm_groups, + dtype=self.dtype, + precision=self.precision, + named_norms=self.named_norms + )(x, temb) + if middle_attention is not None and j == self.num_middle_res_blocks - 1: + x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), attention_impl=self.attention_impl, + dim_head=middle_dim_out // middle_attention['heads'], + use_linear_attention=False, + use_projection=middle_attention.get("use_projection", False), + use_self_and_cross=False, + precision=middle_attention.get("precision", self.precision), + only_pure_attention=middle_attention.get("only_pure_attention", True), + force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False), + norm_inputs=middle_attention.get("norm_inputs", True), + explicitly_add_residual=middle_attention.get("explicitly_add_residual", True), + name=f"middle_attention_{j}")(x, textcontext) + x = TemporalBlock( + features=x.shape[-1], + heads=self.temporal_heads, + dtype=self.dtype, + precision=self.precision, + name=f"middle_temporal_{j}" + )(x, T) + x = ResidualBlock( + middle_conv_type, + name=f"middle_res2_{j}", + features=middle_dim_out, + kernel_size=(3, 3), + strides=(1, 1), + activation=self.activation, + norm_groups=self.norm_groups, + dtype=self.dtype, + precision=self.precision, + named_norms=self.named_norms + )(x, temb) + + # Upscaling Blocks + for i, (dim_out, attention_config) in enumerate(zip(reversed(feature_depths), reversed(attention_configs))): + for j in range(self.num_res_blocks): + x = jnp.concatenate([x, downs.pop()], axis=-1) + kernel_size = (3, 3) + x = ResidualBlock( + up_conv_type, + name=f"up_{i}_residual_{j}", + features=dim_out, + kernel_size=kernel_size, + strides=(1, 1), + activation=self.activation, + norm_groups=self.norm_groups, + dtype=self.dtype, + precision=self.precision, + named_norms=self.named_norms + )(x, temb) + if attention_config is not None and j == self.num_res_blocks - 1: + x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), attention_impl=self.attention_impl, + dim_head=dim_out // attention_config['heads'], + 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), + only_pure_attention=attention_config.get("only_pure_attention", True), + force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False), + norm_inputs=attention_config.get("norm_inputs", True), + explicitly_add_residual=attention_config.get("explicitly_add_residual", True), + name=f"up_{i}_attention_{j}")(x, textcontext) + x = TemporalBlock( + features=x.shape[-1], + heads=self.temporal_heads, + dtype=self.dtype, + precision=self.precision, + name=f"up_{i}_temporal" + )(x, T) + if i != len(feature_depths) - 1: + x = Upsample( + features=feature_depths[-i], + scale=2, + activation=self.activation, + name=f"up_{i}_upsample", + dtype=self.dtype, + precision=self.precision + )(x) + + x = ConvLayer( + conv_type, + features=self.feature_depths[0], + kernel_size=(3, 3), + strides=(1, 1), + dtype=self.dtype, + precision=self.precision + )(x) + + x = jnp.concatenate([x, downs.pop()], axis=-1) + + x = ResidualBlock( + conv_type, + name="final_residual", + features=self.feature_depths[0], + kernel_size=(3,3), + strides=(1, 1), + activation=self.activation, + norm_groups=self.norm_groups, + dtype=self.dtype, + precision=self.precision, + named_norms=self.named_norms + )(x, temb) + + x = self.conv_out_norm(x) + x = self.activation(x) + + noise_out = ConvLayer( + conv_type, + features=self.output_channels, + kernel_size=(3, 3), + strides=(1, 1), + dtype=self.dtype, + precision=self.precision + )(x) + return noise_out.reshape(B, T, H, W, self.output_channels) + + +def inflate_unet_params(params_2d, params_3d): + """Copy a trained 2D Unet param tree into a UNet3D init. + + Spatial module names are identical between the two models, so every 2D + leaf lands on its 3D counterpart; the temporal blocks keep their + (zero-init) fresh params. The result reproduces the 2D model frame by + frame until training moves the temporal weights. + """ + def merge(dst, src): + out = dict(dst) + for key, value in src.items(): + if isinstance(value, dict): + assert key in dst, f"2D module '{key}' has no counterpart in the 3D tree" + out[key] = merge(dst[key], value) + else: + out[key] = value + return out + return merge(params_3d, params_2d) diff --git a/tests/test_config.py b/tests/test_config.py index 292d9f4..35e1ed9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -125,6 +125,9 @@ def test_registry_builds_every_architecture(): "simple_udit": {"num_layers": 4}, "hierarchical_mmdit": {"emb_features": (32, 64, 96), "num_layers": (1, 1, 1), "num_heads": (2, 2, 2), "base_patch_size": 2}, + "unet_3d": {"feature_depths": [16, 32], "attention_configs": [None, None], + "num_res_blocks": 1, "num_middle_res_blocks": 1, + "activation": "swish", "norm_groups": 4, "temporal_heads": 2}, } for name in MODEL_REGISTRY: config = {**base, **per_arch.get(name, {})} diff --git a/tests/test_models.py b/tests/test_models.py index d8c18db..d0f1e24 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -204,3 +204,53 @@ def test_video_dit_temporal_mixing(rng): x_perturbed = x.at[:, 0].add(1.0) out_b = model.apply(params, x_perturbed, temb, None) assert not jnp.allclose(out_a[:, 2], out_b[:, 2]), "no information flow across frames" + + +def test_unet3d_inflation_reproduces_2d_unet(rng): + """A UNet3D inflated from a 2D Unet checkpoint must reproduce the 2D model + frame by frame exactly - the temporal blocks are zero-initialized, so + training starts from the pretrained image model and only learns motion.""" + from flaxdiff.models.unet_3d import UNet3D, inflate_unet_params + + config = dict( + emb_features=64, + feature_depths=[16, 32], + 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, + ) + model_2d = Unet(**config) + model_3d = UNet3D(**config, temporal_heads=2) + + x = jax.random.normal(rng, (2, 3, 16, 16, 3)) + temb = jnp.ones((2,)) + textcontext = jnp.ones((2, 77, 768), dtype=jnp.float32) + + params_2d = model_2d.init(rng, x[:, 0], temb, textcontext) + params_3d = model_3d.init(jax.random.PRNGKey(7), x, temb, textcontext) + inflated = {"params": inflate_unet_params(params_2d["params"], params_3d["params"])} + + out_3d = model_3d.apply(inflated, x, temb, textcontext) + frames_2d = jnp.stack( + [model_2d.apply(params_2d, x[:, t], temb, textcontext) for t in range(3)], axis=1) + assert out_3d.shape == x.shape + assert jnp.max(jnp.abs(out_3d - frames_2d)) < 1e-5, "inflated UNet3D does not match the 2D model" + + +def test_unet3d_temporal_mixing_after_training_signal(rng): + """Once the temporal gate is nudged off zero, frames must exchange information.""" + from flaxdiff.models.unet_3d import UNet3D + + model = UNet3D(emb_features=64, feature_depths=[16, 32], + attention_configs=[None, None], num_res_blocks=1, + num_middle_res_blocks=1, temporal_heads=2) + x = jax.random.normal(rng, (1, 3, 16, 16, 3)) + temb = jnp.ones((1,)) + textcontext = jnp.ones((1, 77, 768), dtype=jnp.float32) + params = model.init(rng, x, temb, textcontext) + params = jax.tree.map(lambda p: p + 0.02, params) + + out_a = model.apply(params, x, temb, textcontext) + out_b = model.apply(params, x.at[:, 0].add(1.0), temb, textcontext) + assert not jnp.allclose(out_a[:, 2], out_b[:, 2]), "no information flow across frames" diff --git a/training.py b/training.py index ab6f25f..9ede2e7 100644 --- a/training.py +++ b/training.py @@ -328,6 +328,28 @@ def main(args): "norm_groups": args.norm_groups, }, }, + "unet_3d": { + "kwargs": { + "feature_depths": args.feature_depths, + "attention_configs": attention_configs, + "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, + }, + }, + "video_dit": { + "kwargs": { + "patch_size": args.patch_size, + "num_layers": args.num_layers, + "num_heads": args.num_heads, + "dropout_rate": args.dropout_rate, + "mlp_ratio": args.mlp_ratio, + "use_hilbert": use_hilbert, + "use_zigzag": use_zigzag, + }, + }, "uvit": { "kwargs": { "patch_size": args.patch_size, From 8ce2f44172a9a866bc661505954459e6babc821d Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:23:19 -0400 Subject: [PATCH 4/4] feat(inputs): generalized conditioning encoders, audio as a first-class modality The base ConditioningEncoder was text-shaped: tokenize() hardcoded max_length/truncation and encode_from_tokens() hardcoded input_ids/attention_mask, so no other modality could use it. Both are now abstract and the text specifics live in TextEncoder, leaving a base class that only knows 'raw data -> tokens -> embeddings'. Audio arrives as AudioEncoder + HFAudioEncoder, which resolves any HF audio model through the Auto* classes and forwards whatever keys its processor emits (input_values for wav2vec2/HuBERT, input_features for Whisper/AST) straight to the model - so swapping the audio model is a config change and nothing more. Formats are already general upstream: audio_utils.read_audio decodes via ffmpeg and resamples to the rate the processor declares. AutoAudioProcessor is the data-pipeline counterpart of AutoTextTokenizer, and the voxceleb2 audio-video transform now uses it instead of calling AutoAudioTokenizer, which never existed - its class body was commented out, so that path raised NameError and flake8's F821 had been failing CI on main since 2026-06-12. Also fixes CLIPTextEncoder.from_modelname discarding its own modelname argument. Co-Authored-By: Claude Fable 5 --- flaxdiff/data/sources/videos.py | 16 +++-- flaxdiff/inputs/encoders.py | 123 +++++++++++++++++++++++++++----- flaxdiff/utils.py | 29 +++++++- tests/test_encoders.py | 78 ++++++++++++++++++++ 4 files changed, 219 insertions(+), 27 deletions(-) create mode 100644 tests/test_encoders.py diff --git a/flaxdiff/data/sources/videos.py b/flaxdiff/data/sources/videos.py index 4533d6f..7d48edd 100644 --- a/flaxdiff/data/sources/videos.py +++ b/flaxdiff/data/sources/videos.py @@ -1,7 +1,7 @@ import cv2 import jax.numpy as jnp import grain.python as pygrain -from flaxdiff.utils import AutoTextTokenizer +from flaxdiff.utils import AutoTextTokenizer, AutoAudioProcessor from typing import Dict, Any, Callable, List, Optional, Tuple import random import os @@ -172,6 +172,7 @@ def create_transform( sequence_length: int = 16, audio_frame_padding: int = 3, method: Any = cv2.INTER_AREA, + audio_modelname: str = "facebook/wav2vec2-base-960h", ) -> Callable[[], pygrain.MapTransform]: """Create a transform for video datasets. @@ -179,6 +180,8 @@ def create_transform( frame_size: Size to scale video frames to. sequence_length: Number of frames to sample from each video. method: Interpolation method for resizing. + audio_modelname: HF audio model whose feature extractor prepares + the audio conditioning inputs. Returns: A callable that returns a pygrain.MapTransform. @@ -188,7 +191,8 @@ def create_transform( class AudioVideoTransform(pygrain.RandomMapTransform): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.tokenize = AutoAudioTokenizer(tensor_type="np") + self.tokenize = AutoAudioProcessor( + tensor_type="np", modelname=audio_modelname) def random_map(self, element, rng: np.random.Generator) -> Dict[str, jnp.array]: video_path = element["video_path"] @@ -201,14 +205,14 @@ def random_map(self, element, rng: np.random.Generator) -> Dict[str, jnp.array]: random_seed=random_seed, ) - # Process caption + # Feature-extract the audio; key names differ per model, so + # pass the processor's output through untouched results = self.tokenize(full_audio) - + return { "video": video_frames, "audio": { - "input_ids": results['input_ids'][0], - "attention_mask": results['attention_mask'][0], + **{key: value[0] for key, value in results.items()}, "full_audio": full_audio, "framewise_audio": framewise_audio, } diff --git a/flaxdiff/inputs/encoders.py b/flaxdiff/inputs/encoders.py index 1987ed6..8680724 100644 --- a/flaxdiff/inputs/encoders.py +++ b/flaxdiff/inputs/encoders.py @@ -1,38 +1,42 @@ import jax.numpy as jnp import flax.linen as nn -from typing import Callable +from typing import Any, Callable from dataclasses import dataclass from abc import ABC, abstractmethod @dataclass class ConditioningEncoder(ABC): - model: nn.Module + """A modality's path from raw data to conditioning embeddings. + + `tokenize` turns raw data into model inputs (runs on CPU in the data + pipeline workers), `encode_from_tokens` turns those into embeddings on + device. Modalities differ in both, so both are abstract here - nothing in + this base class assumes text. + """ + model: Any tokenizer: Callable - + @property + @abstractmethod def key(self): - name = self.tokenizer.__name__ - # Remove the 'Encoder' suffix from the name and lowercase it - if name.endswith("Encoder"): - name = name[:-7].lower() - return name + """Batch/config key this encoder reads and writes.""" + pass def __call__(self, data): tokens = self.tokenize(data) outputs = self.encode_from_tokens(tokens) return outputs - + + @abstractmethod def encode_from_tokens(self, tokens): - outputs = self.model(input_ids=tokens['input_ids'], - attention_mask=tokens['attention_mask']) - last_hidden_state = outputs.last_hidden_state - return last_hidden_state - + """Embed already-tokenized inputs.""" + pass + + @abstractmethod def tokenize(self, data): - tokens = self.tokenizer(data, padding="max_length", - max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="np") - return tokens - + """Turn raw data into model inputs.""" + pass + @abstractmethod def serialize(self): """Serialize the encoder configuration.""" @@ -50,6 +54,17 @@ class TextEncoder(ConditioningEncoder): @property def key(self): return "text" + + def tokenize(self, data): + tokens = self.tokenizer(data, padding="max_length", + max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="np") + return tokens + + def encode_from_tokens(self, tokens): + outputs = self.model(input_ids=tokens['input_ids'], + attention_mask=tokens['attention_mask']) + last_hidden_state = outputs.last_hidden_state + return last_hidden_state @dataclass class CLIPTextEncoder(TextEncoder): @@ -64,7 +79,6 @@ def from_modelname(modelname: str = "openai/clip-vit-large-patch14", backend: st FlaxCLIPTextModel, AutoTokenizer, ) - modelname = "openai/clip-vit-large-patch14" if backend == "jax": model = FlaxCLIPTextModel.from_pretrained( modelname, dtype=jnp.bfloat16) @@ -93,6 +107,77 @@ def deserialize(serialized_config): backend = serialized_config["backend"] return CLIPTextEncoder.from_modelname(modelname=modelname, backend=backend) +@dataclass +class AudioEncoder(ConditioningEncoder): + """Audio conditioning.""" + @property + def key(self): + return "audio" + + +@dataclass +class HFAudioEncoder(AudioEncoder): + """Any HuggingFace audio model, resolved through the Auto* classes. + + Works with wav2vec2/HuBERT (`input_values`), Whisper/AST + (`input_features`) and anything else whose processor and model are + registered with transformers - whatever keys the processor emits are + forwarded to the model unchanged, so swapping the audio model is a config + change and nothing more. Audio *file* formats are handled upstream by + `flaxdiff.data.sources.audio_utils.read_audio`, which decodes via ffmpeg + and resamples to `sampling_rate`. + """ + modelname: str + backend: str + sampling_rate: int + + @staticmethod + def from_modelname(modelname: str = "facebook/wav2vec2-base-960h", + backend: str = "jax", sampling_rate: int = None): + from transformers import AutoFeatureExtractor + processor = AutoFeatureExtractor.from_pretrained(modelname) + if backend == "jax": + from transformers import FlaxAutoModel + model = FlaxAutoModel.from_pretrained(modelname, dtype=jnp.bfloat16) + else: + from transformers import AutoModel + model = AutoModel.from_pretrained(modelname) + if sampling_rate is None: + # The processor knows the rate its model was trained at + sampling_rate = getattr(processor, "sampling_rate", 16000) + return HFAudioEncoder( + model=model, + tokenizer=processor, + modelname=modelname, + backend=backend, + sampling_rate=sampling_rate, + ) + + def tokenize(self, data): + return dict(self.tokenizer( + data, sampling_rate=self.sampling_rate, + padding=True, return_tensors="np")) + + def encode_from_tokens(self, tokens): + outputs = self.model(**tokens) + return outputs.last_hidden_state + + def serialize(self): + return { + "modelname": self.modelname, + "backend": self.backend, + "sampling_rate": self.sampling_rate, + } + + @staticmethod + def deserialize(serialized_config): + return HFAudioEncoder.from_modelname( + modelname=serialized_config["modelname"], + backend=serialized_config["backend"], + sampling_rate=serialized_config.get("sampling_rate"), + ) + CONDITIONAL_ENCODERS_REGISTRY = { "text": CLIPTextEncoder, + "audio": HFAudioEncoder, } diff --git a/flaxdiff/utils.py b/flaxdiff/utils.py index eda6110..3017607 100644 --- a/flaxdiff/utils.py +++ b/flaxdiff/utils.py @@ -137,8 +137,33 @@ def __call__(self, inputs): def __repr__(self): return self.__class__.__name__ + '()' - -# class AutoAudioTokenizer: + + +class AutoAudioProcessor: + """Turn raw audio waveforms into model inputs, for any HF audio model. + + The audio counterpart of AutoTextTokenizer: it runs on CPU in the grain + workers so the device only sees ready tensors. Whatever keys the model's + feature extractor emits (`input_values` for wav2vec2/HuBERT, + `input_features` for Whisper/AST, ...) are passed through unchanged, so + switching audio models needs no change here. + """ + def __init__(self, tensor_type="np", modelname="facebook/wav2vec2-base-960h", + sampling_rate=None): + from transformers import AutoFeatureExtractor + self.processor = AutoFeatureExtractor.from_pretrained(modelname) + self.tensor_type = tensor_type + self.modelname = modelname + # The processor knows the rate its model was trained at + self.sampling_rate = sampling_rate or getattr(self.processor, "sampling_rate", 16000) + + def __call__(self, audio): + features = self.processor(audio, sampling_rate=self.sampling_rate, + padding=True, return_tensors=self.tensor_type) + return dict(features) + + def __repr__(self): + return self.__class__.__name__ + '()' def defaultTextEncodeModel(modelname = "openai/clip-vit-large-patch14", backend="jax"): """Default text encoder model.""" diff --git a/tests/test_encoders.py b/tests/test_encoders.py new file mode 100644 index 0000000..cb2531e --- /dev/null +++ b/tests/test_encoders.py @@ -0,0 +1,78 @@ +"""Conditioning encoder abstraction tests. + +The base class must stay modality-agnostic: text and audio differ in both +tokenization and embedding, and adding a modality must not require touching +anything shared. These use stub processors/models so no weights are +downloaded; the network-marked test covers a real HF audio model. +""" + +import numpy as np +import pytest + +from flaxdiff.inputs import CONDITIONAL_ENCODERS_REGISTRY +from flaxdiff.inputs.encoders import ( + ConditioningEncoder, TextEncoder, CLIPTextEncoder, AudioEncoder, HFAudioEncoder, +) + + +class StubOutput: + def __init__(self, last_hidden_state): + self.last_hidden_state = last_hidden_state + + +def test_registry_exposes_both_modalities(): + assert CONDITIONAL_ENCODERS_REGISTRY["text"] is CLIPTextEncoder + assert CONDITIONAL_ENCODERS_REGISTRY["audio"] is HFAudioEncoder + + +def test_encoder_keys(): + """Each modality declares its own batch key; the intermediate classes stay + abstract until a concrete backend implements serialize/deserialize.""" + assert CLIPTextEncoder(model=None, tokenizer=None, modelname="m", backend="jax").key == "text" + assert HFAudioEncoder(model=None, tokenizer=None, modelname="m", + backend="jax", sampling_rate=16000).key == "audio" + for cls in (TextEncoder, AudioEncoder): + assert cls.__abstractmethods__, f"{cls.__name__} should stay abstract" + + +def test_base_class_is_modality_agnostic(): + """tokenize/encode_from_tokens must be abstract - a base class that + assumes input_ids/attention_mask cannot serve audio.""" + for name in ("tokenize", "encode_from_tokens", "key"): + assert name in ConditioningEncoder.__abstractmethods__ + + +@pytest.mark.parametrize("feature_key", ["input_values", "input_features"]) +def test_audio_encoder_passes_processor_keys_through(feature_key): + """wav2vec2/HuBERT emit input_values, Whisper/AST emit input_features. + The encoder must forward whatever the processor produced, so swapping the + audio model is a config change and nothing else.""" + seen = {} + + class StubProcessor: + sampling_rate = 16000 + + def __call__(self, audio, sampling_rate=None, padding=None, return_tensors=None): + seen["sampling_rate"] = sampling_rate + return {feature_key: np.zeros((1, 4), dtype=np.float32), + "attention_mask": np.ones((1, 4), dtype=np.int32)} + + def stub_model(**kwargs): + seen["forwarded"] = sorted(kwargs) + return StubOutput(np.zeros((1, 4, 8), dtype=np.float32)) + + encoder = HFAudioEncoder(model=stub_model, tokenizer=StubProcessor(), + modelname="stub", backend="jax", sampling_rate=16000) + embeddings = encoder(np.zeros(16000, dtype=np.float32)) + + assert seen["sampling_rate"] == 16000 + assert seen["forwarded"] == sorted([feature_key, "attention_mask"]) + assert embeddings.shape == (1, 4, 8) + + +def test_audio_encoder_roundtrips_config(): + encoder = HFAudioEncoder(model=None, tokenizer=None, modelname="facebook/wav2vec2-base-960h", + backend="jax", sampling_rate=24000) + config = encoder.serialize() + assert config == {"modelname": "facebook/wav2vec2-base-960h", + "backend": "jax", "sampling_rate": 24000}