From 9db1733d2e3a088d2163137b8025a02f709c105d Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Wed, 22 Jul 2026 23:59:48 -0400 Subject: [PATCH 1/6] chore: delete dead code - favor_fastattn.py (722 lines, zero importers, hard-depends on undeclared gin/absl, contains a live typo) - unet_3d.py + unet_3d_blocks.py (950 lines, unreachable since the models/__init__ import was commented out; video support will be rebuilt from the shared backbone) - vit_common: AdaLNZero and PositionalEncoding (nothing constructs them) - common: TimeEmbedding (superseded by FourierEmbedding, NaN for features=2) and PixelShuffle (only reference was a commented line) - utils: RMSNorm (nn.RMSNorm is used everywhere), the never-called map_nested_config that silently dropped every non-string value, and the string maps that fed it - trainer: the duplicated form_global_array/convert_to_global_tree copy (the flaxdiff.utils one is the survivor), dead tensorboard stubs, write-only last_save_time, a 39-line commented-out block - simple_vit: an if/else with identical branches and the dead use_remat flag Co-Authored-By: Claude Fable 5 --- flaxdiff/models/autoencoder/autoencoder.py | 2 +- flaxdiff/models/common.py | 29 - flaxdiff/models/favor_fastattn.py | 723 ------------------ flaxdiff/models/simple_vit.py | 12 +- flaxdiff/models/unet_3d.py | 446 ----------- flaxdiff/models/unet_3d_blocks.py | 505 ------------ flaxdiff/models/vit_common.py | 67 -- flaxdiff/trainer/general_diffusion_trainer.py | 39 - flaxdiff/trainer/simple_trainer.py | 49 +- flaxdiff/utils.py | 118 --- 10 files changed, 6 insertions(+), 1984 deletions(-) delete mode 100644 flaxdiff/models/favor_fastattn.py delete mode 100644 flaxdiff/models/unet_3d.py delete mode 100644 flaxdiff/models/unet_3d_blocks.py diff --git a/flaxdiff/models/autoencoder/autoencoder.py b/flaxdiff/models/autoencoder/autoencoder.py index ab2454d..a307da1 100644 --- a/flaxdiff/models/autoencoder/autoencoder.py +++ b/flaxdiff/models/autoencoder/autoencoder.py @@ -3,7 +3,7 @@ from flax import linen as nn from typing import Dict, Callable, Sequence, Any, Union, Optional import einops -from ..common import kernel_init, ConvLayer, Upsample, Downsample, PixelShuffle +from ..common import kernel_init, ConvLayer, Upsample, Downsample from dataclasses import dataclass from abc import ABC, abstractmethod diff --git a/flaxdiff/models/common.py b/flaxdiff/models/common.py index b732eae..97c2eea 100644 --- a/flaxdiff/models/common.py +++ b/flaxdiff/models/common.py @@ -66,35 +66,6 @@ def __call__(self, x): return(conv.apply({'params': {'kernel': standardized_kernel, 'bias': bias}},x)) -class PixelShuffle(nn.Module): - scale: int - - @nn.compact - def __call__(self, x): - up = einops.rearrange( - x, - pattern="b h w (h2 w2 c) -> b (h h2) (w w2) c", - h2=self.scale, - w2=self.scale, - ) - return up - -class TimeEmbedding(nn.Module): - features:int - nax_positions:int=10000 - - def setup(self): - half_dim = self.features // 2 - emb = jnp.log(self.nax_positions) / (half_dim - 1) - emb = jnp.exp(-emb * jnp.arange(half_dim, dtype=jnp.float32)) - self.embeddings = emb - - def __call__(self, x): - x = jax.lax.convert_element_type(x, jnp.float32) - emb = x[:, None] * self.embeddings[None, :] - emb = jnp.concatenate([jnp.sin(emb), jnp.cos(emb)], axis=-1) - return emb - class FourierEmbedding(nn.Module): features:int scale:int = 16 diff --git a/flaxdiff/models/favor_fastattn.py b/flaxdiff/models/favor_fastattn.py deleted file mode 100644 index fd53844..0000000 --- a/flaxdiff/models/favor_fastattn.py +++ /dev/null @@ -1,723 +0,0 @@ -# coding=utf-8 -# Copyright 2024 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Core Fast Attention Module for Flax. - -Implementation of the approximate fast softmax and generalized -attention mechanism leveraging structured random feature maps [RFM] techniques -and low rank decomposition of the attention matrix. -""" -# pylint: disable=invalid-name, missing-function-docstring, line-too-long - -import abc -from collections.abc import Iterable # pylint: disable=g-importing-member -import functools -from absl import logging -import gin -import jax -from jax import lax -from jax import random -import jax.numpy as jnp - -import numpy as onp - -# Nonlinear mappings encoding different attention kernels. -gin.external_configurable(jnp.cos, 'jcos') -gin.external_configurable(jnp.sin, 'jsin') -gin.external_configurable(jnp.tanh, 'jtanh') -gin.external_configurable(jax.nn.sigmoid, 'jsigmoid') -gin.external_configurable( - lambda x: jax.nn.gelu(x, approximate=False), 'jgelu' -) # Needs to be exact, although might be slower. See https://github.com/google/jax/issues/4428. -gin.external_configurable(lambda x: x * x * (x > 0.0), 'jrequ') -gin.external_configurable(jnp.exp, 'jexp') -gin.external_configurable(lambda x: x, 'jidentity') -gin.external_configurable( - lambda x: (jnp.exp(x)) * (x <= 0.0) + (x + 1.0) * (x > 0.0), 'jshiftedelu' -) # Nonlinearity used in "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention" (https://arxiv.org/abs/2006.16236). - - -def nonnegative_softmax_kernel_feature_creator(data, - projection_matrix, - attention_dims_t, - batch_dims_t, - precision, - is_query, - normalize_data=True, - eps=0.0001): - """Constructs nonnegative kernel features for fast softmax attention. - - - Args: - data: input for which features are computes - projection_matrix: random matrix used to compute features - attention_dims_t: tuple of attention dimensions - batch_dims_t: tuple of batch dimensions - precision: precision parameter - is_query: predicate indicating whether input data corresponds to queries or - keys - normalize_data: predicate indicating whether data should be normalized, - eps: numerical stabilizer. - - Returns: - Random features for fast softmax attention. - """ - - if normalize_data: - # We have e^{qk^T/sqrt{d}} = e^{q_norm k_norm^T}, where - # w_norm = w * data_normalizer for w in {q,k}. - data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1]))) - else: - data_normalizer = 1.0 - ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0]) - data_mod_shape = data.shape[0:len(batch_dims_t)] + projection_matrix.shape - data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix - - data_dash = lax.dot_general( - data_normalizer * data, - data_thick_random_matrix, - (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), - (batch_dims_t, batch_dims_t)), - precision=precision) - - diag_data = jnp.square(data) - diag_data = jnp.sum(diag_data, axis=data.ndim - 1) - diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer - diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1) - - last_dims_t = (len(data_dash.shape) - 1,) - if is_query: - data_dash = ratio * ( - jnp.exp(data_dash - diag_data - - jnp.max(data_dash, axis=last_dims_t, keepdims=True)) + eps) - else: - data_dash = ratio * ( - jnp.exp(data_dash - diag_data - jnp.max( - data_dash, axis=last_dims_t + attention_dims_t, keepdims=True)) + - eps) - - return data_dash - - -def sincos_softmax_kernel_feature_creator(data, - projection_matrix, - attention_dims_t, - batch_dims_t, - precision, - normalize_data=True): - """Constructs kernel sin-cos features for fast softmax attention. - - - Args: - data: input for which features are computes - projection_matrix: random matrix used to compute features - attention_dims_t: tuple of attention dimensions - batch_dims_t: tuple of batch dimensions - precision: precision parameter - normalize_data: predicate indicating whether data should be normalized. - - Returns: - Random features for fast softmax attention. - """ - if normalize_data: - # We have: exp(qk^T/sqrt{d}) = exp(|q|^2/2sqrt{d}) * exp(|k|^2/2sqrt{d}) * - # exp(-(|q*c-k*c|^2)/2), where c = 1.0 / sqrt{sqrt{d}}. - data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1]))) - else: - data_normalizer = 1.0 - ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0]) - data_mod_shape = data.shape[0:len(batch_dims_t)] + projection_matrix.shape - data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix - - data_dash = lax.dot_general( - data_normalizer * data, - data_thick_random_matrix, - (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), - (batch_dims_t, batch_dims_t)), - precision=precision) - data_dash_cos = ratio * jnp.cos(data_dash) - data_dash_sin = ratio * jnp.sin(data_dash) - data_dash = jnp.concatenate((data_dash_cos, data_dash_sin), axis=-1) - - # Constructing D_data and data^{'} - diag_data = jnp.square(data) - diag_data = jnp.sum(diag_data, axis=data.ndim - 1) - diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer - diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1) - # Additional renormalization for numerical stability - data_renormalizer = jnp.max(diag_data, attention_dims_t, keepdims=True) - diag_data -= data_renormalizer - diag_data = jnp.exp(diag_data) - data_prime = data_dash * diag_data - return data_prime - - -def generalized_kernel_feature_creator(data, projection_matrix, batch_dims_t, - precision, kernel_fn, kernel_epsilon, - normalize_data): - """Constructs kernel features for fast generalized attention. - - - Args: - data: input for which features are computes - projection_matrix: matrix used to compute features - batch_dims_t: tuple of batch dimensions - precision: precision parameter - kernel_fn: kernel function used - kernel_epsilon: additive positive term added to every feature for numerical - stability - normalize_data: predicate indicating whether data should be normalized. - - Returns: - Random features for fast generalized attention. - """ - if normalize_data: - data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1]))) - else: - data_normalizer = 1.0 - if projection_matrix is None: - return kernel_fn(data_normalizer * data) + kernel_epsilon - else: - data_mod_shape = data.shape[0:len(batch_dims_t)] + projection_matrix.shape - data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix - data_dash = lax.dot_general( - data_normalizer * data, - data_thick_random_matrix, - (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), - (batch_dims_t, batch_dims_t)), - precision=precision) - data_prime = kernel_fn(data_dash) + kernel_epsilon - return data_prime - - -@gin.configurable -def make_fast_softmax_attention(qkv_dim, - renormalize_attention=True, - numerical_stabilizer=0.000001, - nb_features=256, - ortho_features=True, - ortho_scaling=0.0, - redraw_features=True, - unidirectional=False, - nonnegative_features=True, - lax_scan_unroll=1): - """Construct a fast softmax attention method.""" - logging.info( - 'Fast softmax attention: %s features and orthogonal=%s, renormalize=%s', - nb_features, ortho_features, renormalize_attention) - if ortho_features: - matrix_creator = functools.partial( - GaussianOrthogonalRandomMatrix, - nb_features, - qkv_dim, - scaling=ortho_scaling) - else: - matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, - nb_features, qkv_dim) - if nonnegative_features: - - def kernel_feature_creator(data, - projection_matrix, - attention_dims_t, - batch_dims_t, - precision, - is_query, - normalize_data=True): - return nonnegative_softmax_kernel_feature_creator( - data, projection_matrix, attention_dims_t, batch_dims_t, precision, - is_query, normalize_data, numerical_stabilizer) - else: - - def kernel_feature_creator(data, - projection_matrix, - attention_dims_t, - batch_dims_t, - precision, - is_query, - normalize_data=True): - del is_query - return sincos_softmax_kernel_feature_creator(data, projection_matrix, - attention_dims_t, - batch_dims_t, precision, - normalize_data) - - attention_fn = FastAttentionviaLowRankDecomposition( - matrix_creator, - kernel_feature_creator, - renormalize_attention=renormalize_attention, - numerical_stabilizer=numerical_stabilizer, - redraw_features=redraw_features, - unidirectional=unidirectional, - lax_scan_unroll=lax_scan_unroll).dot_product_attention - return attention_fn - - -@gin.configurable -def make_fast_generalized_attention(qkv_dim, - renormalize_attention=True, - numerical_stabilizer=0.0, - nb_features=256, - features_type='deterministic', - kernel_fn=jax.nn.relu, - kernel_epsilon=0.001, - redraw_features=False, - unidirectional=False, - lax_scan_unroll=1): - """Construct a fast generalized attention menthod.""" - logging.info('Fast generalized attention.: %s features and renormalize=%s', - nb_features, renormalize_attention) - if features_type == 'ortho': - matrix_creator = functools.partial( - GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=False) - elif features_type == 'iid': - matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, - nb_features, qkv_dim) - elif features_type == 'deterministic': - matrix_creator = None - else: - raise ValueError('Unknown feature value type') - - def kernel_feature_creator(data, - projection_matrix, - attention_dims_t, - batch_dims_t, - precision, - is_query, - normalize_data=False): - del attention_dims_t - del is_query - return generalized_kernel_feature_creator(data, projection_matrix, - batch_dims_t, precision, - kernel_fn, kernel_epsilon, - normalize_data) - - attention_fn = FastAttentionviaLowRankDecomposition( - matrix_creator, - kernel_feature_creator, - renormalize_attention=renormalize_attention, - numerical_stabilizer=numerical_stabilizer, - redraw_features=redraw_features, - unidirectional=unidirectional, - lax_scan_unroll=lax_scan_unroll).dot_product_attention - return attention_fn - - -class RandomMatrix(object): - r"""Abstract class providing a method for constructing 2D random arrays. - - Class is responsible for constructing 2D random arrays. - """ - - __metaclass__ = abc.ABCMeta - - @abc.abstractmethod - def get_2d_array(self): - raise NotImplementedError('Abstract method') - - -class GaussianUnstructuredRandomMatrix(RandomMatrix): - - def __init__(self, nb_rows, nb_columns, key): - self.nb_rows = nb_rows - self.nb_columns = nb_columns - self.key = key - - def get_2d_array(self): - return random.normal(self.key, (self.nb_rows, self.nb_columns)) - - -class GaussianOrthogonalRandomMatrix(RandomMatrix): - r"""Class providing a method to create Gaussian orthogonal matrix. - - Class is responsible for constructing 2D Gaussian orthogonal arrays. - """ - - def __init__(self, nb_rows, nb_columns, key, scaling=0): - self.nb_rows = nb_rows - self.nb_columns = nb_columns - self.key = key - self.scaling = scaling - - def get_2d_array(self): - nb_full_blocks = int(self.nb_rows / self.nb_columns) - block_list = [] - rng = self.key - for _ in range(nb_full_blocks): - rng, rng_input = jax.random.split(rng) - unstructured_block = random.normal(rng_input, - (self.nb_columns, self.nb_columns)) - q, _ = jnp.linalg.qr(unstructured_block) - q = jnp.transpose(q) - block_list.append(q) - remaining_rows = self.nb_rows - nb_full_blocks * self.nb_columns - if remaining_rows > 0: - rng, rng_input = jax.random.split(rng) - unstructured_block = random.normal(rng_input, - (self.nb_columns, self.nb_columns)) - q, _ = jnp.linalg.qr(unstructured_block) - q = jnp.transpose(q) - block_list.append(q[0:remaining_rows]) - final_matrix = jnp.vstack(block_list) - - if self.scaling == 0: - multiplier = jnp.linalg.norm( - random.normal(self.key, (self.nb_rows, self.nb_columns)), axis=1) - elif self.scaling == 1: - multiplier = jnp.sqrt(float(self.nb_columns)) * jnp.ones((self.nb_rows)) - else: - raise ValueError('Scaling must be one of {0, 1}. Was %s' % self._scaling) - - return jnp.matmul(jnp.diag(multiplier), final_matrix) - - -class FastAttention(object): - r"""Abstract class providing a method for fast attention. - - Class is responsible for providing a method for fast - approximate attention. - """ - - __metaclass__ = abc.ABCMeta - - @abc.abstractmethod - def dot_product_attention(self, - query, - key, - value, - dtype=jnp.float32, - bias=None, - mask=None, - axis=None, - broadcast_dropout=True, - dropout_rng=None, - dropout_rate=0., - deterministic=False, - precision=None): - """Computes dot-product attention given query, key, and value. - - This is the core function for applying fast approximate dot-product - attention. It calculates the attention weights given query and key and - combines the values using the attention weights. This function supports - multi-dimensional inputs. - - - Args: - query: queries for calculating attention with shape of [batch_size, dim1, - dim2, ..., dimN, num_heads, mem_channels]. - key: keys for calculating attention with shape of [batch_size, dim1, dim2, - ..., dimN, num_heads, mem_channels]. - value: values to be used in attention with shape of [batch_size, dim1, - dim2,..., dimN, num_heads, value_channels]. - dtype: the dtype of the computation (default: float32) - bias: bias for the attention weights. This can be used for incorporating - autoregressive mask, padding mask, proximity bias. - mask: mask for the attention weights. This can be used for incorporating - autoregressive masks. - axis: axises over which the attention is applied. - broadcast_dropout: bool: use a broadcasted dropout along batch dims. - dropout_rng: JAX PRNGKey: to be used for dropout. - dropout_rate: dropout rate. - deterministic: bool, deterministic or not (to apply dropout). - precision: numerical precision of the computation see `jax.lax.Precision` - for details. - - Returns: - Output of shape [bs, dim1, dim2, ..., dimN,, num_heads, value_channels]. - """ - raise NotImplementedError('Abstract method') - - -def _numerator(z_slice_shape, precision, unroll=1): - - def fwd(qs, ks, vs): - - def body(p, qkv): - (q, k, v) = qkv - p += jnp.einsum('...m,...d->...md', k, v, precision=precision) - X_slice = jnp.einsum('...m,...md->...d', q, p, precision=precision) - return p, X_slice - - init_value = jnp.zeros(z_slice_shape) - p, W = lax.scan(body, init_value, (qs, ks, vs), unroll=unroll) - return W, (p, qs, ks, vs) - - def bwd(pqkv, W_ct): - - def body(carry, qkv_xct): - p, p_ct = carry - q, k, v, x_ct = qkv_xct - q_ct = jnp.einsum('...d,...md->...m', x_ct, p, precision=precision) - p_ct += jnp.einsum('...d,...m->...md', x_ct, q, precision=precision) - k_ct = jnp.einsum('...md,...d->...m', p_ct, v, precision=precision) - v_ct = jnp.einsum('...md,...m->...d', p_ct, k, precision=precision) - p -= jnp.einsum('...m,...d->...md', k, v, precision=precision) - return (p, p_ct), (q_ct, k_ct, v_ct) - - p, qs, ks, vs = pqkv - _, (qs_ct, ks_ct, vs_ct) = lax.scan( - body, (p, jnp.zeros_like(p)), (qs, ks, vs, W_ct), - reverse=True, - unroll=unroll) - return qs_ct, ks_ct, vs_ct - - @jax.custom_vjp - def _numerator_impl(qs, ks, vs): - W, _ = fwd(qs, ks, vs) - return W - - _numerator_impl.defvjp(fwd, bwd) - - return _numerator_impl - - -def _denominator(t_slice_shape, precision, unroll=1): - - def fwd(qs, ks): - - def body(p, qk): - q, k = qk - p += k - x = jnp.einsum('...m,...m->...', q, p, precision=precision) - return p, x - - p = jnp.zeros(t_slice_shape) - p, R = lax.scan(body, p, (qs, ks), unroll=unroll) - return R, (qs, ks, p) - - def bwd(qkp, R_ct): - - def body(carry, qkx): - p, p_ct = carry - q, k, x_ct = qkx - q_ct = jnp.einsum('...,...m->...m', x_ct, p, precision=precision) - p_ct += jnp.einsum('...,...m->...m', x_ct, q, precision=precision) - k_ct = p_ct - p -= k - return (p, p_ct), (q_ct, k_ct) - - qs, ks, p = qkp - _, (qs_ct, ks_ct) = lax.scan( - body, (p, jnp.zeros_like(p)), (qs, ks, R_ct), - reverse=True, - unroll=unroll) - return (qs_ct, ks_ct) - - @jax.custom_vjp - def _denominator_impl(qs, ks): - R, _ = fwd(qs, ks) - return R - - _denominator_impl.defvjp(fwd, bwd) - - return _denominator_impl - - -class FastAttentionviaLowRankDecomposition(FastAttention): - r"""Class providing a method for fast attention via low rank decomposition. - - Class is responsible for providing a method for fast - dot-product attention with the use of low rank decomposition (e.g. with - random feature maps). - """ - - def __init__(self, - matrix_creator, - kernel_feature_creator, - renormalize_attention, - numerical_stabilizer, - redraw_features, - unidirectional, - lax_scan_unroll=1): # For optimal GPU performance, set to 16. - rng = random.PRNGKey(0) - self.matrix_creator = matrix_creator - self.projection_matrix = self.draw_weights(rng) - self.kernel_feature_creator = kernel_feature_creator - self.renormalize_attention = renormalize_attention - self.numerical_stabilizer = numerical_stabilizer - self.redraw_features = redraw_features - self.unidirectional = unidirectional - self.lax_scan_unroll = lax_scan_unroll - - def draw_weights(self, key): - if self.matrix_creator is None: - return None - matrixrng, _ = random.split(key) - projection_matrix = self.matrix_creator(key=matrixrng).get_2d_array() - return projection_matrix - - def dot_product_attention(self, - query, - key, - value, - dtype=jnp.float32, - bias=None, - mask=None, - axis=None, - broadcast_dropout=True, - dropout_rng=None, - dropout_rate=0., - deterministic=False, - precision=None): - - assert key.shape[:-1] == value.shape[:-1] - assert (query.shape[0:1] == key.shape[0:1] and - query.shape[-1] == key.shape[-1]) - if axis is None: - axis = tuple(range(1, key.ndim - 2)) - if not isinstance(axis, Iterable): - axis = (axis,) - assert key.ndim == query.ndim - assert key.ndim == value.ndim - for ax in axis: - if not (query.ndim >= 3 and 1 <= ax < query.ndim - 2): - raise ValueError('Attention axis must be between the batch ' - 'axis and the last-two axes.') - n = key.ndim - - # Constructing projection tensor. - if self.redraw_features: - # TODO(kchoro): Get rid of the constant below. - query_seed = lax.convert_element_type( - jnp.ceil(jnp.sum(query) * 10000000.0), jnp.int32) - rng = random.PRNGKey(query_seed) - self.projection_matrix = self.draw_weights(rng) - - # batch_dims is , num_heads> - batch_dims = tuple(onp.delete(range(n), axis + (n - 1,))) - # q & k -> (bs, , num_heads, , channels) - qk_perm = batch_dims + axis + (n - 1,) - k_extra_perm = axis + batch_dims + (n - 1,) - key_extra = key.transpose(k_extra_perm) - key = key.transpose(qk_perm) - query = query.transpose(qk_perm) - # v -> (bs, , num_heads, , channels) - v_perm = batch_dims + axis + (n - 1,) - value = value.transpose(v_perm) - batch_dims_t = tuple(range(len(batch_dims))) - attention_dims_t = tuple( - range(len(batch_dims), - len(batch_dims) + len(axis))) - - # Constructing tensors Q^{'} and K^{'}. - query_prime = self.kernel_feature_creator(query, self.projection_matrix, - attention_dims_t, batch_dims_t, - precision, True) - key_prime = self.kernel_feature_creator(key, self.projection_matrix, - attention_dims_t, batch_dims_t, - precision, False) - - if self.unidirectional: - index = attention_dims_t[0] - z_slice_shape = key_prime.shape[0:len(batch_dims_t)] + ( - key_prime.shape[-1],) + (value.shape[-1],) - - numerator_fn = _numerator(z_slice_shape, precision, self.lax_scan_unroll) - W = numerator_fn( - jnp.moveaxis(query_prime, index, 0), - jnp.moveaxis(key_prime, index, 0), jnp.moveaxis(value, index, 0)) - - # Constructing W = (Q^{'}(K^{'})^{T})_{masked}V - W = jnp.moveaxis(W, 0, index) - - if not self.renormalize_attention: - # Unidirectional, not-normalized attention. - perm_inv = _invert_perm(qk_perm) - result = W.transpose(perm_inv) - return result - else: - # Unidirectional, normalized attention. - thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones( - key_extra.shape[0:len(axis)]) - - index = attention_dims_t[0] - t_slice_shape = key_prime.shape[0:len(batch_dims_t)] + ( - key_prime.shape[-1],) - denominator_fn = _denominator(t_slice_shape, precision, - self.lax_scan_unroll) - R = denominator_fn( - jnp.moveaxis(query_prime, index, 0), - jnp.moveaxis(key_prime, index, 0)) - - R = jnp.moveaxis(R, 0, index) - else: - contract_query = tuple( - range(len(batch_dims) + len(axis), - len(batch_dims) + len(axis) + 1)) - contract_z = tuple(range(len(batch_dims), len(batch_dims) + 1)) - # Constructing Z = (K^{'})^{T}V - # Z (bs, , num_heads, channels_m, channels_v) - Z = lax.dot_general( - key_prime, - value, - ((attention_dims_t, attention_dims_t), (batch_dims_t, batch_dims_t)), - precision=precision) - # Constructing W = Q^{'}Z = Q^{'}(K^{'})^{T}V - # q (bs, , num_heads, , channels_m) - # Z (bs, , num_heads, channels_m, channels_v) - # W (bs, , num_heads, , channels_v) - W = lax.dot_general( - query_prime, - Z, ((contract_query, contract_z), (batch_dims_t, batch_dims_t)), - precision=precision) - if not self.renormalize_attention: - # Bidirectional, not-normalized attention. - perm_inv = _invert_perm(qk_perm) - result = W.transpose(perm_inv) - return result - else: - # Bidirectional, normalized attention. - thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones( - key_extra.shape[0:len(axis)]) - contract_key = tuple( - range(len(batch_dims), - len(batch_dims) + len(axis))) - contract_thick_all_ones = tuple( - range(thick_all_ones.ndim - len(axis), thick_all_ones.ndim)) - # Construct T = (K^{'})^{T} 1_L - # k (bs, , num_heads, , channels) - T = lax.dot_general( - key_prime, - thick_all_ones, ((contract_key, contract_thick_all_ones), - (batch_dims_t, batch_dims_t)), - precision=precision) - - # Construct partition function: R = Q^{'} T = Q^{'}(K^{'})^{T} 1_L - # q_p (bs, , num_heads, , channs_m) - # T (bs, , num_heads, channels_m) - R = lax.dot_general( - query_prime, - T, (((query_prime.ndim - 1,), (T.ndim - 1,)), - (batch_dims_t, range(0, - len(T.shape) - 1))), - precision=precision) - - R = R + 2 * self.numerical_stabilizer * ( - jnp.abs(R) <= self.numerical_stabilizer) - R = jnp.reciprocal(R) - R = jnp.expand_dims(R, len(R.shape)) - # W (bs, , num_heads, , channels_v) - # R (bs, , num_heads, , extra_channel) - result = W * R - # back to (bs, dim1, dim2, ..., dimN, num_heads, channels) - perm_inv = _invert_perm(qk_perm) - result = result.transpose(perm_inv) - return result - - -def _invert_perm(perm): - perm_inv = [0] * len(perm) - for i, j in enumerate(perm): - perm_inv[j] = i - return tuple(perm_inv) \ No newline at end of file diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py index 0aca9fe..5daaa76 100644 --- a/flaxdiff/models/simple_vit.py +++ b/flaxdiff/models/simple_vit.py @@ -11,7 +11,7 @@ from flax.typing import Dtype, PrecisionLike from functools import partial from .hilbert import hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify -from .vit_common import _rotate_half, unpatchify, PatchEmbedding, apply_rotary_embedding, RotaryEmbedding, RoPEAttention, AdaLNZero, AdaLNParams +from .vit_common import unpatchify, PatchEmbedding, RotaryEmbedding, RoPEAttention, AdaLNParams from .simple_dit import DiTBlock @@ -35,20 +35,14 @@ class UViT(nn.Module): explicitly_add_residual: bool = True # Passed to TransformerBlock norm_epsilon: float = 1e-5 # Adjusted default use_hilbert: bool = False # Toggle Hilbert patch reorder - use_remat: bool = False # Add flag to use remat def setup(self): assert self.num_layers % 2 == 0, "num_layers must be even for U-Net structure" half_layers = self.num_layers // 2 # --- Norm Layer --- - if self.norm_groups > 0: - print(f"Warning: norm_groups > 0 not fully supported with standard LayerNorm fallback in UViT setup. Using LayerNorm.") - self.norm_factory = partial( - nn.LayerNorm, epsilon=self.norm_epsilon, dtype=self.dtype) - else: - self.norm_factory = partial( - nn.LayerNorm, epsilon=self.norm_epsilon, dtype=self.dtype) + self.norm_factory = partial( + nn.LayerNorm, epsilon=self.norm_epsilon, dtype=self.dtype) # --- Input Path --- self.patch_embed = PatchEmbedding( diff --git a/flaxdiff/models/unet_3d.py b/flaxdiff/models/unet_3d.py deleted file mode 100644 index 7635680..0000000 --- a/flaxdiff/models/unet_3d.py +++ /dev/null @@ -1,446 +0,0 @@ -from typing import Dict, Optional, Tuple, Union - -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from diffusers.configuration_utils import ConfigMixin, flax_register_to_config -from diffusers.utils import BaseOutput -from diffusers.models.embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps -from diffusers.models.modeling_flax_utils import FlaxModelMixin - -from .unet_3d_blocks import ( - FlaxCrossAttnDownBlock3D, - FlaxCrossAttnUpBlock3D, - FlaxDownBlock3D, - FlaxUNetMidBlock3DCrossAttn, - FlaxUpBlock3D, -) - - -@flax_register_to_config -class FlaxUNet3DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - A conditional 3D UNet model for video diffusion. - - Parameters: - sample_size (`int` or `Tuple[int, int, int]`, *optional*, defaults to (16, 32, 32)): - The spatial and temporal size of the input sample. Can be provided as a single integer for square spatial size and fixed temporal size. - in_channels (`int`, *optional*, defaults to 4): - The number of channels in the input sample. - out_channels (`int`, *optional*, defaults to 4): - The number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to ("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")): - The tuple of downsample blocks to use. - up_block_types (`Tuple[str]`, *optional*, defaults to ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")): - The tuple of upsample blocks to use. - block_out_channels (`Tuple[int]`, *optional*, defaults to (320, 640, 1280, 1280)): - The tuple of output channels for each block. - layers_per_block (`int`, *optional*, defaults to 2): - The number of layers per block. - attention_head_dim (`int`, *optional*, defaults to 8): - The dimension of the attention heads. - cross_attention_dim (`int`, *optional*, defaults to 1280): - The dimension of the cross attention features. - dropout (`float`, *optional*, defaults to 0): - Dropout probability for down, up and bottleneck blocks. - use_linear_projection (`bool`, *optional*, defaults to False): - Whether to use linear projection in attention blocks. - dtype (`jnp.dtype`, *optional*, defaults to jnp.float32): - The dtype of the model weights. - flip_sin_to_cos (`bool`, *optional*, defaults to True): - Whether to flip the sin to cos in the time embedding. - freq_shift (`int`, *optional*, defaults to 0): - The frequency shift to apply to the time embedding. - use_memory_efficient_attention (`bool`, *optional*, defaults to False): - Whether to use memory-efficient attention. - split_head_dim (`bool`, *optional*, defaults to False): - Whether to split the head dimension into a new axis for the self-attention computation. - """ - - sample_size: Union[int, Tuple[int, int, int]] = (16, 32, 32) - in_channels: int = 4 - out_channels: int = 4 - down_block_types: Tuple[str, ...] = ( - "CrossAttnDownBlock3D", - "CrossAttnDownBlock3D", - "CrossAttnDownBlock3D", - "DownBlock3D", - ) - up_block_types: Tuple[str, ...] = ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D") - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280) - layers_per_block: int = 2 - attention_head_dim: Union[int, Tuple[int, ...]] = 8 - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None - cross_attention_dim: int = 1280 - dropout: float = 0.0 - use_linear_projection: bool = False - dtype: jnp.dtype = jnp.float32 - flip_sin_to_cos: bool = True - freq_shift: int = 0 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1 - addition_embed_type: Optional[str] = None - addition_time_embed_dim: Optional[int] = None - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - if isinstance(self.sample_size, int): - sample_size = (self.sample_size, self.sample_size, self.sample_size) - else: - sample_size = self.sample_size - - # Shape: [batch, frames, height, width, channels] - sample_shape = (1, sample_size[0], sample_size[1], sample_size[2], self.in_channels) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - timesteps = jnp.ones((1,), dtype=jnp.int32) - encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - added_cond_kwargs = None - if self.addition_embed_type == "text_time": - # For text-time conditioning for video diffusion - text_embeds_dim = self.cross_attention_dim - time_ids_dims = 6 # Default value for video models - added_cond_kwargs = { - "text_embeds": jnp.zeros((1, text_embeds_dim), dtype=jnp.float32), - "time_ids": jnp.zeros((1, time_ids_dims), dtype=jnp.float32), - } - - return self.init(rngs, sample, timesteps, encoder_hidden_states, added_cond_kwargs)["params"] - - def setup(self) -> None: - block_out_channels = self.block_out_channels - time_embed_dim = block_out_channels[0] * 4 - - if self.num_attention_heads is not None: - raise ValueError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue. " - "Use `attention_head_dim` instead." - ) - - # Default behavior: if num_attention_heads is not set, use attention_head_dim - num_attention_heads = self.num_attention_heads or self.attention_head_dim - - # input - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3, 3), - strides=(1, 1, 1), - padding=((1, 1), (1, 1), (1, 1)), - dtype=self.dtype, - ) - - # time - self.time_proj = FlaxTimesteps( - block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.freq_shift - ) - self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - - # Handle attention head configurations - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(self.down_block_types) - - # transformer layers per block - transformer_layers_per_block = self.transformer_layers_per_block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(self.down_block_types) - - # addition embed types - if self.addition_embed_type == "text_time": - if self.addition_time_embed_dim is None: - raise ValueError( - f"addition_embed_type {self.addition_embed_type} requires `addition_time_embed_dim` to not be None" - ) - self.add_time_proj = FlaxTimesteps(self.addition_time_embed_dim, self.flip_sin_to_cos, self.freq_shift) - self.add_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - else: - self.add_embedding = None - - # down blocks - down_blocks = [] - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlock3D": - down_block = FlaxCrossAttnDownBlock3D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - num_attention_heads=num_attention_heads[i], - add_downsample=not is_final_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=False, # We don't use only cross attention in 3D UNet - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - elif down_block_type == "DownBlock3D": - down_block = FlaxDownBlock3D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - else: - raise ValueError(f"Unknown down block type: {down_block_type}") - - down_blocks.append(down_block) - self.down_blocks = down_blocks - - # mid block - self.mid_block = FlaxUNetMidBlock3DCrossAttn( - in_channels=block_out_channels[-1], - dropout=self.dropout, - num_attention_heads=num_attention_heads[-1], - transformer_layers_per_block=transformer_layers_per_block[-1], - use_linear_projection=self.use_linear_projection, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - - # up blocks - up_blocks = [] - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - output_channel = reversed_block_out_channels[0] - reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - - for i, up_block_type in enumerate(self.up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - is_final_block = i == len(block_out_channels) - 1 - - if up_block_type == "CrossAttnUpBlock3D": - up_block = FlaxCrossAttnUpBlock3D( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - num_layers=self.layers_per_block + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - num_attention_heads=reversed_num_attention_heads[i], - add_upsample=not is_final_block, - dropout=self.dropout, - use_linear_projection=self.use_linear_projection, - only_cross_attention=False, # We don't use only cross attention in 3D UNet - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - elif up_block_type == "UpBlock3D": - up_block = FlaxUpBlock3D( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - num_layers=self.layers_per_block + 1, - add_upsample=not is_final_block, - dropout=self.dropout, - dtype=self.dtype, - ) - else: - raise ValueError(f"Unknown up block type: {up_block_type}") - - up_blocks.append(up_block) - prev_output_channel = output_channel - self.up_blocks = up_blocks - - # out - self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.conv_out = nn.Conv( - self.out_channels, - kernel_size=(3, 3, 3), - strides=(1, 1, 1), - padding=((1, 1), (1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__( - self, - sample: jnp.ndarray, - timesteps: Union[jnp.ndarray, float, int], - encoder_hidden_states: jnp.ndarray, - frame_encoder_hidden_states: Optional[jnp.ndarray] = None, - added_cond_kwargs: Optional[Union[Dict, FrozenDict]] = None, - down_block_additional_residuals: Optional[Tuple[jnp.ndarray, ...]] = None, - mid_block_additional_residual: Optional[jnp.ndarray] = None, - return_dict: bool = True, - train: bool = False, - ) -> Union[jnp.ndarray]: - r""" - Args: - sample (`jnp.ndarray`): (batch, frames, height, width, channels) noisy inputs tensor - timestep (`jnp.ndarray` or `float` or `int`): timesteps - encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states - frame_encoder_hidden_states (`jnp.ndarray`, *optional*): - (batch_size, frames, sequence_length, hidden_size) per-frame encoder hidden states - added_cond_kwargs: (`dict`, *optional*): - Additional embeddings to add to the time embeddings - down_block_additional_residuals: (`tuple` of `jnp.ndarray`, *optional*): - Additional residual connections for down blocks - mid_block_additional_residual: (`jnp.ndarray`, *optional*): - Additional residual connection for mid block - return_dict (`bool`, *optional*, defaults to `True`): - Whether to return a dict or tuple - train (`bool`, *optional*, defaults to `False`): - Training mode flag for dropout - """ - # Extract the number of frames from the input - batch, num_frames, height, width, channels = sample.shape - - # 1. Time embedding - if not isinstance(timesteps, jnp.ndarray): - timesteps = jnp.array([timesteps], dtype=jnp.int32) - elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0: - timesteps = timesteps.astype(dtype=jnp.float32) - timesteps = jnp.expand_dims(timesteps, 0) - - t_emb = self.time_proj(timesteps) - t_emb = self.time_embedding(t_emb) - - # Repeat time embedding for each frame - t_emb = jnp.repeat(t_emb, repeats=num_frames, axis=0) - - - # additional embeddings - if self.add_embedding is not None and added_cond_kwargs is not None: - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - "text_embeds must be provided for text_time addition_embed_type" - ) - if "time_ids" not in added_cond_kwargs: - raise ValueError( - "time_ids must be provided for text_time addition_embed_type" - ) - - text_embeds = added_cond_kwargs["text_embeds"] - time_ids = added_cond_kwargs["time_ids"] - - # Compute time embeds - time_embeds = self.add_time_proj(jnp.ravel(time_ids)) - time_embeds = jnp.reshape(time_embeds, (text_embeds.shape[0], -1)) - - # Concatenate text and time embeds - add_embeds = jnp.concatenate([text_embeds, time_embeds], axis=-1) - - # Project to time embedding dimension - aug_emb = self.add_embedding(add_embeds) - t_emb = t_emb + aug_emb - - # 2. Pre-process input - reshape from [B, F, H, W, C] to [B*F, H, W, C] for 2D operations - sample = sample.reshape(batch * num_frames, height, width, channels) - sample = self.conv_in(sample) - - # Process encoder hidden states - repeat for each frame and combine with frame-specific conditioning if provided - if encoder_hidden_states is not None: - # Repeat video-wide conditioning for each frame: (B, S, X) -> (B*F, S, X) - encoder_hidden_states_expanded = jnp.repeat( - encoder_hidden_states, repeats=num_frames, axis=0 - ) - - # If we have frame-specific conditioning, reshape and concatenate with video conditioning - if frame_encoder_hidden_states is not None: - # Reshape from (B, F, S, X) to (B*F, S, X) - frame_encoder_hidden_states = frame_encoder_hidden_states.reshape( - batch * num_frames, *frame_encoder_hidden_states.shape[2:] - ) - - # Concatenate along the sequence dimension - encoder_hidden_states_combined = jnp.concatenate( - [encoder_hidden_states_expanded, frame_encoder_hidden_states], - axis=1 - ) - else: - encoder_hidden_states_combined = encoder_hidden_states_expanded - else: - encoder_hidden_states_combined = None - - # 3. Down blocks - down_block_res_samples = (sample,) - for down_block in self.down_blocks: - if isinstance(down_block, FlaxCrossAttnDownBlock3D): - sample, res_samples = down_block( - sample, - t_emb, - encoder_hidden_states_combined, - num_frames=num_frames, - deterministic=not train - ) - else: - sample, res_samples = down_block( - sample, - t_emb, - num_frames=num_frames, - deterministic=not train - ) - down_block_res_samples += res_samples - - # Add additional residuals if provided - if down_block_additional_residuals is not None: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample += down_block_additional_residual - new_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. Mid block - sample = self.mid_block( - sample, - t_emb, - encoder_hidden_states_combined, - num_frames=num_frames, - deterministic=not train - ) - - # Add mid block residual if provided - if mid_block_additional_residual is not None: - sample += mid_block_additional_residual - - # 5. Up blocks - for up_block in self.up_blocks: - res_samples = down_block_res_samples[-(self.layers_per_block + 1) :] - down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)] - if isinstance(up_block, FlaxCrossAttnUpBlock3D): - sample = up_block( - sample, - res_hidden_states_tuple=res_samples, - temb=t_emb, - encoder_hidden_states=encoder_hidden_states_combined, - num_frames=num_frames, - deterministic=not train, - ) - else: - sample = up_block( - sample, - res_hidden_states_tuple=res_samples, - temb=t_emb, - num_frames=num_frames, - deterministic=not train - ) - - # 6. Post-process - sample = self.conv_norm_out(sample) - sample = nn.silu(sample) - sample = self.conv_out(sample) - - # Reshape back to [B, F, H, W, C] - sample = sample.reshape(batch, num_frames, height, width, self.out_channels) - return sample \ No newline at end of file diff --git a/flaxdiff/models/unet_3d_blocks.py b/flaxdiff/models/unet_3d_blocks.py deleted file mode 100644 index bb0d127..0000000 --- a/flaxdiff/models/unet_3d_blocks.py +++ /dev/null @@ -1,505 +0,0 @@ -from typing import Tuple, Optional - -import flax.linen as nn -import jax -import jax.numpy as jnp - -from diffusers.models.attention_flax import ( - FlaxBasicTransformerBlock, - FlaxTransformer2DModel, -) - -from diffusers.models.resnet_flax import ( - FlaxResnetBlock2D, - FlaxUpsample2D, - FlaxDownsample2D, -) - -from diffusers.models.unets.unet_2d_blocks_flax import ( - FlaxCrossAttnDownBlock2D, - FlaxDownBlock2D, - FlaxUNetMidBlock2DCrossAttn, - FlaxUpBlock2D, - FlaxCrossAttnUpBlock2D, -) - -class FlaxTransformerTemporalModel(nn.Module): - """ - Transformer for temporal attention in 3D UNet. - """ - in_channels: int - n_heads: int - d_head: int - depth: int = 1 - dropout: float = 0.0 - only_cross_attention: bool = False - dtype: jnp.dtype = jnp.float32 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - - def setup(self): - self.norm = nn.GroupNorm(num_groups=32, epsilon=1e-5) - - inner_dim = self.n_heads * self.d_head - self.proj_in = nn.Dense(inner_dim, dtype=self.dtype) - # Use existing FlaxBasicTransformerBlock from diffusers - self.transformer_blocks = [ - FlaxBasicTransformerBlock( - inner_dim, - self.n_heads, - self.d_head, - dropout=self.dropout, - only_cross_attention=self.only_cross_attention, - dtype=self.dtype, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - ) - for _ in range(self.depth) - ] - - self.proj_out = nn.Dense(inner_dim, dtype=self.dtype) - - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states: jnp.ndarray, context: jnp.ndarray, num_frames: int, deterministic=True): - # Save original shape for later reshaping - batch_depth, height, width, channels = hidden_states.shape - batch = batch_depth // num_frames - - # Reshape to (batch, depth, height, width, channels) - hidden_states = hidden_states.reshape(batch, num_frames, height, width, channels) - residual = hidden_states - - # Apply normalization - hidden_states = self.norm(hidden_states) - - # Reshape for temporal attention: (batch, depth, height, width, channels) -> - # (batch*height*width, depth, channels) - hidden_states = hidden_states.transpose(0, 2, 3, 1, 4) - hidden_states = hidden_states.reshape(batch * height * width, num_frames, channels) - - # Project input - hidden_states = self.proj_in(hidden_states) - - # Apply transformer blocks - for block in self.transformer_blocks: - hidden_states = block(hidden_states, context=context, deterministic=deterministic) - - # Project output - hidden_states = self.proj_out(hidden_states) - - # Reshape back to original shape - hidden_states = hidden_states.reshape(batch, height, width, num_frames, channels) - hidden_states = hidden_states.transpose(0, 3, 1, 2, 4) - - # Add residual connection - hidden_states = hidden_states + residual - - # Reshape back to (batch*depth, height, width, channels) - hidden_states = hidden_states.reshape(batch_depth, height, width, channels) - - return hidden_states - -class TemporalConvLayer(nn.Module): - in_channels: int - out_channels: Optional[int] = None - dropout: float = 0.0 - norm_num_groups: int = 32 - dtype: jnp.dtype = jnp.float32 - - @nn.compact - def __call__(self, x: jnp.ndarray, num_frames: int, deterministic=True) -> jnp.ndarray: - """ - Args: - x: shape (B*F, H, W, C) - num_frames: number of frames F per batch element - - Returns: - A jnp.ndarray of shape (B*F, H, W, C) - """ - out_channels = self.out_channels or self.in_channels - bf, h, w, c = x.shape - b = bf // num_frames - - # Reshape to [B, F, H, W, C], interpret F as "depth" for 3D conv - x = x.reshape(b, num_frames, h, w, c) - identity = x - - # conv1: in_channels -> out_channels - x = nn.GroupNorm(num_groups=self.norm_num_groups)(x) - x = nn.silu(x) - x = nn.Conv(features=out_channels, kernel_size=(3, 1, 1), - dtype=self.dtype, - padding=((1,1), (0,0), (0,0)))(x) - - # conv2: out_channels -> in_channels - x = nn.GroupNorm(num_groups=self.norm_num_groups)(x) - x = nn.silu(x) - x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic) - x = nn.Conv(features=self.in_channels, kernel_size=(3, 1, 1), - dtype=self.dtype, - padding=((1,1), (0,0), (0,0)))(x) - - # conv3: in_channels -> in_channels - x = nn.GroupNorm(num_groups=self.norm_num_groups)(x) - x = nn.silu(x) - x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic) - x = nn.Conv(features=self.in_channels, kernel_size=(3, 1, 1), - dtype=self.dtype, - padding=((1,1), (0,0), (0,0)))(x) - - # conv4 (zero-init): in_channels -> in_channels - x = nn.GroupNorm(num_groups=self.norm_num_groups)(x) - x = nn.silu(x) - x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic) - x = nn.Conv( - features=self.in_channels, - kernel_size=(3, 1, 1), - padding=((1,1), (0,0), (0,0)), - kernel_init=nn.initializers.zeros, - bias_init=nn.initializers.zeros, - dtype=self.dtype, - )(x) - - # Residual connection and reshape back to (B*F, H, W, C) - x = identity + x - x = x.reshape(bf, h, w, c) - return x - - -class FlaxCrossAttnDownBlock3D(FlaxCrossAttnDownBlock2D): - """ - Cross attention 3D downsampling block. - """ - - def setup(self): - resnets = [] - temp_convs = [] - attentions = [] - temp_attentions = [] - - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - temp_conv = TemporalConvLayer( - in_channels=self.out_channels, - out_channels=self.out_channels, - dropout=self.dropout, - dtype=self.dtype, - ) - temp_convs.append(temp_conv) - attn_block = FlaxTransformer2DModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=self.only_cross_attention, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - temp_attn_block = FlaxTransformerTemporalModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - dropout=self.dropout, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - temp_attentions.append(temp_attn_block) - - self.temp_convs = temp_convs - self.temp_attentions = temp_attentions - self.resnets = resnets - self.attentions = attentions - - if self.add_downsample: - # self.downsamplers_0 = FlaxDownsample3D(self.out_channels, dtype=self.dtype) - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, temb, encoder_hidden_states, num_frames, deterministic=True): - output_states = () - - for resnet, temp_conv, attn, temp_attn in zip(self.resnets, self.temp_convs, self.attentions, self.temp_attentions): - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = temp_conv(hidden_states, num_frames=num_frames, deterministic=deterministic) - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - hidden_states = temp_attn(hidden_states, None, num_frames=num_frames, deterministic=deterministic) - output_states += (hidden_states,) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - output_states += (hidden_states,) - - return hidden_states, output_states - - -class FlaxDownBlock3D(FlaxDownBlock2D): - """ - Basic downsampling block without attention. - """ - def setup(self): - resnets = [] - temp_convs = [] - - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - temp_conv = TemporalConvLayer( - in_channels=self.out_channels, - out_channels=self.out_channels, - dropout=self.dropout, - dtype=self.dtype, - ) - temp_convs.append(temp_conv) - self.temp_convs = temp_convs - self.resnets = resnets - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, temb, num_frames, deterministic=True): - output_states = () - - for resnet, temp_conv in zip(self.resnets, self.temp_convs): - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = temp_conv(hidden_states, num_frames=num_frames, deterministic=deterministic) - output_states += (hidden_states,) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - output_states += (hidden_states,) - - return hidden_states, output_states - - -class FlaxCrossAttnUpBlock3D(FlaxCrossAttnUpBlock2D): - """ - Cross attention 3D upsampling block. - """ - - def setup(self): - resnets = [] - temp_convs = [] - attentions = [] - temp_attentions = [] - - for i in range(self.num_layers): - res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels - resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - temp_conv = TemporalConvLayer( - in_channels=self.out_channels, - out_channels=self.out_channels, - dropout=self.dropout, - dtype=self.dtype, - ) - temp_convs.append(temp_conv) - attn_block = FlaxTransformer2DModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=self.only_cross_attention, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - temp_attn_block = FlaxTransformerTemporalModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - dropout=self.dropout, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - temp_attentions.append(temp_attn_block) - - self.resnets = resnets - self.attentions = attentions - self.temp_convs = temp_convs - self.temp_attentions = temp_attentions - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, res_hidden_states_tuple, temb, encoder_hidden_states, num_frames, deterministic=True): - for resnet, temp_conv, attn, temp_attn in zip(self.resnets, self.temp_convs, self.attentions, self.temp_attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1) - - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = temp_conv(hidden_states, num_frames=num_frames, deterministic=deterministic) - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - hidden_states = temp_attn(hidden_states, None, num_frames=num_frames, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUpBlock3D(FlaxUpBlock2D): - """ - Basic upsampling block without attention. - """ - def setup(self): - resnets = [] - temp_convs = [] - - for i in range(self.num_layers): - res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels - resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - temp_conv = TemporalConvLayer( - in_channels=self.out_channels, - out_channels=self.out_channels, - dropout=self.dropout, - dtype=self.dtype, - ) - temp_convs.append(temp_conv) - - self.resnets = resnets - self.temp_convs = temp_convs - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, res_hidden_states_tuple, temb, num_frames, deterministic=True): - for resnet, temp_conv in zip(self.resnets, self.temp_convs): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1) - - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = temp_conv(hidden_states, num_frames=num_frames, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUNetMidBlock3DCrossAttn(FlaxUNetMidBlock2DCrossAttn): - """ - Middle block with cross-attention for 3D UNet. - """ - def setup(self): - # there is always at least one resnet - resnets = [ - FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - ] - temp_convs = [ - TemporalConvLayer( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout=self.dropout, - dtype=self.dtype, - ) - ] - - attentions = [] - temp_attentions = [] - - for _ in range(self.num_layers): - attn_block = FlaxTransformer2DModel( - in_channels=self.in_channels, - n_heads=self.num_attention_heads, - d_head=self.in_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - temp_block = FlaxTransformerTemporalModel( - in_channels=self.in_channels, - n_heads=self.num_attention_heads, - d_head=self.in_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - dropout=self.dropout, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - temp_attentions.append(temp_block) - - res_block = FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - temp_conv = TemporalConvLayer( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout=self.dropout, - dtype=self.dtype, - ) - temp_convs.append(temp_conv) - - self.temp_convs = temp_convs - self.temp_attentions = temp_attentions - self.resnets = resnets - self.attentions = attentions - - def __call__(self, hidden_states, temb, encoder_hidden_states, num_frames, deterministic=True): - hidden_states = self.resnets[0](hidden_states, temb, deterministic=deterministic) - hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames, deterministic=deterministic) - - for attn, temp_attn, resnet, temp_conv in zip( - self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:] - ): - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - hidden_states = temp_attn(hidden_states, None, num_frames=num_frames, deterministic=deterministic) - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = temp_conv(hidden_states, num_frames=num_frames, deterministic=deterministic) - - return hidden_states diff --git a/flaxdiff/models/vit_common.py b/flaxdiff/models/vit_common.py index 6fb3761..632f832 100644 --- a/flaxdiff/models/vit_common.py +++ b/flaxdiff/models/vit_common.py @@ -37,22 +37,6 @@ def __call__(self, x): return x -class PositionalEncoding(nn.Module): - max_len: int - embedding_dim: int - - @nn.compact - def __call__(self, x): - pe = self.param('pos_encoding', - jax.nn.initializers.zeros, - (1, self.max_len, self.embedding_dim)) - return x + pe[:, :x.shape[1], :] - - -# --- Rotary Positional Embedding (RoPE) --- -# Adapted from https://github.com/google-deepmind/ring_attention/blob/main/ring_attention/layers/rotary.py - - def _rotate_half(x: jax.Array) -> jax.Array: """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -186,57 +170,6 @@ def __call__(self, x, context=None, freqs_cis=None): # --- adaLN-Zero --- -class AdaLNZero(nn.Module): - features: int - dtype: Optional[Dtype] = None - precision: PrecisionLike = None - norm_epsilon: float = 1e-5 # Standard LayerNorm epsilon - - @nn.compact - def __call__(self, x, conditioning): - # Project conditioning signal to get scale and shift parameters - # Conditioning shape: [B, D_cond] -> [B, 1, ..., 1, 6 * features] for broadcasting - # Or [B, 1, 6*features] if x is [B, S, F] - - # Ensure conditioning has seq dim if x does - # x=[B,S,F], cond=[B,D_cond] - if x.ndim == 3 and conditioning.ndim == 2: - conditioning = jnp.expand_dims( - conditioning, axis=1) # cond=[B,1,D_cond] - - # Project conditioning to get 6 params per feature (scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn) - # Using nn.DenseGeneral for flexibility if needed, but nn.Dense is fine if cond is [B, D_cond] or [B, 1, D_cond] - ada_params = nn.Dense( - features=6 * self.features, - dtype=self.dtype, - precision=self.precision, - # Initialize projection to zero (Zero init) - kernel_init=nn.initializers.zeros, - name="ada_proj" - )(conditioning) - - # Split into scale, shift, gate for MLP and Attention - scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( - ada_params, 6, axis=-1) - - scale_mlp = jnp.clip(scale_mlp, -10.0, 10.0) - shift_mlp = jnp.clip(shift_mlp, -10.0, 10.0) - # Apply Layer Normalization - norm = nn.LayerNorm(epsilon=self.norm_epsilon, - use_scale=False, use_bias=False, dtype=self.dtype) - # norm = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype) # Alternative: RMSNorm - - norm_x = norm(x) - - # Modulate for Attention path - x_attn = norm_x * (1 + scale_attn) + shift_attn - - # Modulate for MLP path - x_mlp = norm_x * (1 + scale_mlp) + shift_mlp - - # Return modulated outputs and gates - return x_attn, gate_attn, x_mlp, gate_mlp - class AdaLNParams(nn.Module): # Renamed for clarity features: int dtype: Optional[Dtype] = None diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index 147bc22..554bad9 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -64,45 +64,6 @@ def generate_modelname( # model_name += f"-{'.'.join([cond.encoder.key for cond in input_config.conditions])}" - # # Create a sorted representation of model config for consistent hashing - # def sort_dict_recursively(d): - # if isinstance(d, dict): - # return {k: sort_dict_recursively(d[k]) for k in sorted(d.keys())} - # elif isinstance(d, list): - # return [sort_dict_recursively(v) for v in d] - # else: - # return d - - # # Extract model config and sort it - # model_config = serialize_model(model) - # sorted_model_config = sort_dict_recursively(model_config) - - # # Convert to JSON string with sorted keys for consistent hash - # try: - # config_json = json.dumps(sorted_model_config) - # except TypeError: - # # Handle non-serializable objects - # def make_serializable(obj): - # if isinstance(obj, dict): - # return {k: make_serializable(v) for k, v in obj.items()} - # elif isinstance(obj, list): - # return [make_serializable(v) for v in obj] - # else: - # try: - # # Test if object is JSON serializable - # json.dumps(obj) - # return obj - # except TypeError: - # return str(obj) - - # serializable_config = make_serializable(sorted_model_config) - # config_json = json.dumps(serializable_config) - - # # Generate a hash of the configuration - # config_hash = hashlib.md5(config_json.encode('utf-8')).hexdigest()[:8] - - # # Construct the model name - # model_name = f"{model_name}-{config_hash}" return model_name class GeneralDiffusionTrainer(DiffusionTrainer): diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py index 1d23acd..8b69f10 100644 --- a/flaxdiff/trainer/simple_trainer.py +++ b/flaxdiff/trainer/simple_trainer.py @@ -23,11 +23,10 @@ from termcolor import colored from typing import Dict, Callable, Sequence, Any, Union, Tuple from flax.training.dynamic_scale import DynamicScale -from flaxdiff.utils import RandomMarkovState +from flaxdiff.utils import RandomMarkovState, convert_to_global_tree from flax.training import dynamic_scale as dynamic_scale_lib from dataclasses import dataclass import shutil -import gc PROCESS_COLOR_MAP = { 0: "green", @@ -40,30 +39,6 @@ 7: "light_cyan" } -def _build_global_shape_and_sharding( - local_shape: tuple[int, ...], global_mesh: Mesh -) -> tuple[tuple[int, ...], jax.sharding.NamedSharding]: - sharding = jax.sharding.NamedSharding(global_mesh, P(global_mesh.axis_names)) - global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:] - return global_shape, sharding - - -def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array: - """Put local sharded array into local devices""" - global_shape, sharding = _build_global_shape_and_sharding(np.shape(array), global_mesh) - try: - local_device_arrays = np.split(array, len(global_mesh.local_devices), axis=0) - except ValueError as array_split_error: - raise ValueError( - f"Unable to put to devices shape {array.shape} with " - f"local device count {len(global_mesh.local_devices)} " - ) from array_split_error - local_device_buffers = jax.device_put(local_device_arrays, global_mesh.local_devices) - return jax.make_array_from_single_device_arrays(global_shape, sharding, local_device_buffers) - -def convert_to_global_tree(global_mesh, pytree): - return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree) - @struct.dataclass class Metrics(metrics.Collection): accuracy: metrics.Accuracy @@ -310,13 +285,6 @@ def checkpoint_path(self): os.makedirs(path) return path - def tensorboard_path(self): - experiment_name = self.name - path = os.path.join(os.path.abspath('./tensorboard'), experiment_name) - if not os.path.exists(path): - os.makedirs(path) - return path - def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=False): checkpointer = orbax.checkpoint.PyTreeCheckpointer() options = orbax.checkpoint.CheckpointManagerOptions( @@ -441,17 +409,6 @@ def config(self): "input_shapes": self.input_shapes } - def init_tensorboard(self, batch_size, steps_per_epoch, epochs): - from flax.metrics import tensorboard - summary_writer = tensorboard.SummaryWriter(self.tensorboard_path()) - summary_writer.hparams({ - **self.config(), - "steps_per_epoch": steps_per_epoch, - "epochs": epochs, - "batch_size": batch_size - }) - return summary_writer - def validation_loop( self, val_state: SimpleTrainState, @@ -510,7 +467,6 @@ def train_loop( 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') @@ -563,8 +519,7 @@ def train_loop( print(f"Devices: {len(jax.devices())}") # To sync the devices self.save(current_epoch, current_step, train_state, rng_state) print(f"Saving done by process index {process_index}") - last_save_time = time.time() - print(colored(f"Epoch done on index {process_index} => {current_epoch} Loss: {epoch_loss/train_steps_per_epoch}", 'green')) + print(colored(f"Epoch done on index {process_index} => {current_epoch} Loss: {epoch_loss/train_steps_per_epoch}", 'green')) if pbar is not None: pbar.close() return epoch_loss, current_step, train_state, rng_state diff --git a/flaxdiff/utils.py b/flaxdiff/utils.py index b1263b0..eda6110 100644 --- a/flaxdiff/utils.py +++ b/flaxdiff/utils.py @@ -10,53 +10,6 @@ from flaxdiff.inputs import TextEncoder, CLIPTextEncoder # Setup mappings for dtype, precision, and activation -DTYPE_MAP = { - 'bfloat16': jnp.bfloat16, - 'float32': jnp.float32, - 'jax.numpy.float32': jnp.float32, - 'jax.numpy.bfloat16': jnp.bfloat16, - 'None': None, - None: None, -} - -PRECISION_MAP = { - 'high': jax.lax.Precision.HIGH, - 'HIGH': jax.lax.Precision.HIGH, - 'default': jax.lax.Precision.DEFAULT, - 'DEFAULT': jax.lax.Precision.DEFAULT, - 'highest': jax.lax.Precision.HIGHEST, - 'HIGHEST': jax.lax.Precision.HIGHEST, - 'None': None, - None: None, -} - -ACTIVATION_MAP = { - 'swish': jax.nn.swish, - 'silu': jax.nn.silu, - 'jax._src.nn.functions.silu': jax.nn.silu, - 'mish': jax.nn.mish, -} - -def map_nested_config(config): - new_config = {} - for key, value in config.items(): - if isinstance(value, dict): - new_config[key] = map_nested_config(value) - elif isinstance(value, str): - if value in DTYPE_MAP: - new_config[key] = DTYPE_MAP[value] - elif value in PRECISION_MAP: - new_config[key] = PRECISION_MAP[value] - elif value in ACTIVATION_MAP: - 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.") - return new_config - def serialize_model(model: nn.Module): """ Serializes the model to a dictionary format. @@ -165,77 +118,6 @@ def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array: def convert_to_global_tree(global_mesh, pytree): return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree) -class RMSNorm(nn.Module): - """ - From "Root Mean Square Layer Normalization" by https://arxiv.org/abs/1910.07467 - - Adapted from flax.linen.LayerNorm - """ - - epsilon: float = 1e-6 - dtype: Any = jnp.float32 - param_dtype: Any = jnp.float32 - use_scale: bool = True - scale_init: Any = jax.nn.initializers.ones - - @nn.compact - def __call__(self, x): - reduction_axes = (-1,) - feature_axes = (-1,) - - rms_sq = self._compute_rms_sq(x, reduction_axes) - - return self._normalize( - self, - x, - rms_sq, - reduction_axes, - feature_axes, - self.dtype, - self.param_dtype, - self.epsilon, - self.use_scale, - self.scale_init, - ) - - def _compute_rms_sq(self, x, axes): - x = jnp.asarray(x, jnp.promote_types(jnp.float32, jnp.result_type(x))) - rms_sq = jnp.mean(jax.lax.square(x), axes) - return rms_sq - - def _normalize( - self, - mdl, - x, - rms_sq, - reduction_axes, - feature_axes, - dtype, - param_dtype, - epsilon, - use_scale, - scale_init, - ): - reduction_axes = nn.normalization._canonicalize_axes(x.ndim, reduction_axes) - feature_axes = nn.normalization._canonicalize_axes(x.ndim, feature_axes) - stats_shape = list(x.shape) - for axis in reduction_axes: - stats_shape[axis] = 1 - rms_sq = rms_sq.reshape(stats_shape) - feature_shape = [1] * x.ndim - reduced_feature_shape = [] - for ax in feature_axes: - feature_shape[ax] = x.shape[ax] - reduced_feature_shape.append(x.shape[ax]) - mul = jax.lax.rsqrt(rms_sq + epsilon) - if use_scale: - scale = mdl.param( - "scale", scale_init, reduced_feature_shape, param_dtype - ).reshape(feature_shape) - mul *= scale - y = mul * x - return jnp.asarray(y, dtype) - class AutoTextTokenizer: def __init__(self, tensor_type="pt", modelname="openai/clip-vit-large-patch14"): from transformers import AutoTokenizer From 4fd490763f056bcdd41293e4d26c1f6279b60e2e Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 00:04:53 -0400 Subject: [PATCH 2/6] refactor: one model registry as the single source of truth for construction training.py and the inference pipeline each carried their own architecture table, their own dtype/precision/activation string maps (three copies total, already drifted: a 'highes' typo on one side, extra keys on another), and their own +2d/+hilbert/+zigzag suffix parser. A model was therefore constructed through two different code paths at train and inference time. flaxdiff.models.registry now owns the architecture table, the string mapping, suffix canonicalization, and field-filtered construction. Model configs stay as plain strings all the way into the logged wandb config, and build_model resolves them identically on both sides - the logged config IS the construction input. The hand-maintained --architecture choices list (which had already drifted from the inference table) is gone; the registry validates. Co-Authored-By: Claude Fable 5 --- flaxdiff/inference/utils.py | 142 +++--------------------------------- flaxdiff/models/registry.py | 132 +++++++++++++++++++++++++++++++++ tests/test_config.py | 30 ++++++++ training.py | 108 ++++----------------------- 4 files changed, 188 insertions(+), 224 deletions(-) create mode 100644 flaxdiff/models/registry.py diff --git a/flaxdiff/inference/utils.py b/flaxdiff/inference/utils.py index e45036c..36e16d5 100644 --- a/flaxdiff/inference/utils.py +++ b/flaxdiff/inference/utils.py @@ -1,32 +1,18 @@ import jax import jax.numpy as jnp import json -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 -from flaxdiff.models.general import BCHWModelWrapper -from flaxdiff.models.autoencoder.diffusers import StableDiffusionVAE -from flaxdiff.inputs import DiffusionInputConfig, ConditionalInputConfig -from flaxdiff.utils import defaultTextEncodeModel -from diffusers import FlaxUNet2DConditionModel +import os +import warnings + import wandb -from flaxdiff.models.simple_unet import Unet -from flaxdiff.models.simple_vit import UViT -from flaxdiff.models.general import BCHWModelWrapper +from orbax.checkpoint import CheckpointManager, CheckpointManagerOptions, PyTreeCheckpointer + +from flaxdiff.predictors import get_diffusion_preset +from flaxdiff.models.registry import build_model, canonicalize_architecture, map_config_strings from flaxdiff.models.autoencoder.diffusers import StableDiffusionVAE from flaxdiff.inputs import DiffusionInputConfig, ConditionalInputConfig from flaxdiff.utils import defaultTextEncodeModel -from flaxdiff.models.simple_vit import UViT, SimpleUDiT -from flaxdiff.models.simple_dit import SimpleDiT -from flaxdiff.models.simple_mmdit import SimpleMMDiT, HierarchicalMMDiT -from flaxdiff.models.ssm_dit import HybridSSMAttentionDiT -from orbax.checkpoint import CheckpointManager, CheckpointManagerOptions, PyTreeCheckpointer -import os - -import warnings - def get_wandb_run(wandb_run: str, project, entity): """ Try to get the wandb run for the given experiment name and project. @@ -82,97 +68,12 @@ def parse_config(config, overrides=None): conf = merged_config # Setup mappings for dtype, precision, and activation - DTYPE_MAP = { - 'bfloat16': jnp.bfloat16, - 'float32': jnp.float32, - 'jax.numpy.float32': jnp.float32, - 'jax.numpy.bfloat16': jnp.bfloat16, - 'None': None, - None: None, - } - - PRECISION_MAP = { - 'high': jax.lax.Precision.HIGH, - 'HIGH': jax.lax.Precision.HIGH, - 'default': jax.lax.Precision.DEFAULT, - 'DEFAULT': jax.lax.Precision.DEFAULT, - 'highest': jax.lax.Precision.HIGHEST, - 'HIGHEST': jax.lax.Precision.HIGHEST, - 'None': None, - None: None, - } - - ACTIVATION_MAP = { - 'swish': jax.nn.swish, - 'silu': jax.nn.silu, - 'jax._src.nn.functions.silu': jax.nn.silu, - 'mish': jax.nn.mish, - } - - # Get model class based on architecture - MODEL_CLASSES = { - 'unet': Unet, - 'uvit': UViT, - 'diffusers_unet_simple': FlaxUNet2DConditionModel, - 'simple_dit': SimpleDiT, - 'simple_mmdit': SimpleMMDiT, - 'simple_udit': SimpleUDiT, - 'hierarchical_mmdit': HierarchicalMMDiT, - # +2d/+hilbert/+zigzag suffixes are stripped before this lookup; - # the corresponding flags come in via model_config kwargs - 'hybrid_dit': HybridSSMAttentionDiT, - } - - # Map all the leaves of the model config, converting strings to appropriate types - def map_nested_config(config): - new_config = {} - for key, value in config.items(): - if isinstance(value, dict): - new_config[key] = map_nested_config(value) - elif isinstance(value, list): - new_config[key] = [map_nested_config(item) if isinstance(item, dict) else item for item in value] - elif isinstance(value, str): - if value in DTYPE_MAP: - new_config[key] = DTYPE_MAP[value] - elif value in PRECISION_MAP: - new_config[key] = PRECISION_MAP[value] - elif value in ACTIVATION_MAP: - new_config[key] = ACTIVATION_MAP[value] - elif value == 'None': - new_config[key] = None - 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: - new_config[key] = value - return new_config - # Parse architecture and model config model_config = conf['model'] # Get architecture type architecture = conf.get('architecture', conf.get('arguments', {}).get('architecture', 'unet')) - - # Strip +2d/+hilbert/+zigzag suffixes for the MODEL_CLASSES lookup - the - # matching flags (use_2d_fusion, use_hilbert, use_zigzag) are already in - # model_config so the model still gets them as kwargs. - if isinstance(architecture, str): - canonical = architecture - for suffix in ('+2d', '+hilbert', '+zigzag'): - canonical = canonical.replace(suffix, '') - if canonical != architecture: - print(f"Canonicalized architecture: '{architecture}' -> '{canonical}' " - f"(suffix flags read from model_config)") - architecture = canonical - + # Handle autoencoder autoencoder_name = conf.get('autoencoder', conf.get('arguments', {}).get('autoencoder')) autoencoder_opts_str = conf.get('autoencoder_opts', conf.get('arguments', {}).get('autoencoder_opts', '{}')) @@ -188,7 +89,7 @@ def map_nested_config(config): if autoencoder_name == 'stable_diffusion': print("Using Stable Diffusion Autoencoder for Latent Diffusion Modeling") - autoencoder_opts = map_nested_config(autoencoder_opts) + autoencoder_opts = map_config_strings(autoencoder_opts) autoencoder = StableDiffusionVAE(**autoencoder_opts) input_config = conf.get('input_config', None) @@ -220,29 +121,8 @@ def map_nested_config(config): # Deserialize the input config if it's a string input_config = DiffusionInputConfig.deserialize(input_config) - model_kwargs = map_nested_config(model_config) - - print(f"Model kwargs after mapping: {model_kwargs}") - - model_class = MODEL_CLASSES.get(architecture) - if not model_class: - raise ValueError(f"Unknown architecture: {architecture}. Supported architectures: {', '.join(MODEL_CLASSES.keys())}") - - # Drop config keys the model no longer has fields for (older runs logged - # since-removed flags like use_flash_attention) - import dataclasses - valid_fields = {f.name for f in dataclasses.fields(model_class)} - dropped = sorted(set(model_kwargs) - valid_fields) - if dropped: - print(f"Dropping config keys not accepted by {model_class.__name__}: {dropped}") - model_kwargs = {k: v for k, v in model_kwargs.items() if k in valid_fields} - - # Instantiate the model - model = model_class(**model_kwargs) - - # If using diffusers UNet, wrap it for consistent interface - if 'diffusers' in architecture: - model = BCHWModelWrapper(model) + model = build_model(architecture, model_config) + model_kwargs = map_config_strings(model_config) # Same preset as training, so the sampling convention always matches noise_schedule_type = conf.get('noise_schedule', conf.get('arguments', {}).get('noise_schedule', 'edm')) diff --git a/flaxdiff/models/registry.py b/flaxdiff/models/registry.py new file mode 100644 index 0000000..74b320d --- /dev/null +++ b/flaxdiff/models/registry.py @@ -0,0 +1,132 @@ +"""Single source of truth for architecture names and model construction. + +Both training.py and the inference pipeline build models through here, so a +config logged at training time always reconstructs the same model at +inference time. Architecture names may carry +2d/+hilbert/+zigzag suffixes; +they canonicalize to the base architecture plus the matching config flags. +""" + +import dataclasses + +import jax +import jax.numpy as jnp + +from .simple_unet import Unet +from .simple_vit import UViT, SimpleUDiT +from .simple_dit import SimpleDiT +from .simple_mmdit import SimpleMMDiT, HierarchicalMMDiT +from .ssm_dit import HybridSSMAttentionDiT + +MODEL_REGISTRY = { + 'unet': Unet, + 'uvit': UViT, + 'simple_udit': SimpleUDiT, + 'simple_dit': SimpleDiT, + 'simple_mmdit': SimpleMMDiT, + 'hierarchical_mmdit': HierarchicalMMDiT, + 'hybrid_dit': HybridSSMAttentionDiT, +} + +ARCHITECTURE_SUFFIX_FLAGS = { + '+2d': 'use_2d_fusion', + '+hilbert': 'use_hilbert', + '+zigzag': 'use_zigzag', +} + +DTYPE_MAP = { + 'bfloat16': jnp.bfloat16, + 'float32': jnp.float32, + 'jax.numpy.float32': jnp.float32, + 'jax.numpy.bfloat16': jnp.bfloat16, + 'None': None, + None: None, +} + +PRECISION_MAP = { + 'high': jax.lax.Precision.HIGH, + 'HIGH': jax.lax.Precision.HIGH, + 'default': jax.lax.Precision.DEFAULT, + 'DEFAULT': jax.lax.Precision.DEFAULT, + 'highest': jax.lax.Precision.HIGHEST, + 'HIGHEST': jax.lax.Precision.HIGHEST, + 'None': None, + None: None, +} + +ACTIVATION_MAP = { + 'swish': jax.nn.swish, + 'silu': jax.nn.silu, + 'mish': jax.nn.mish, + 'gelu': jax.nn.gelu, + 'relu': jax.nn.relu, +} + + +def canonicalize_architecture(architecture: str) -> tuple[str, dict]: + """Strip +2d/+hilbert/+zigzag suffixes into their matching config flags.""" + flags = {} + canonical = architecture + for suffix, flag in ARCHITECTURE_SUFFIX_FLAGS.items(): + if suffix in canonical: + canonical = canonical.replace(suffix, '') + flags[flag] = True + return canonical, flags + + +def map_config_strings(config: dict): + """Convert the string leaves of a logged config back to objects. + + dtypes, precisions and activations are stored as strings in the wandb + config; function paths like 'jax.nn.mish' (or the 'jax._src.nn.functions.*' + that older configs contain) resolve by attribute walk. + """ + if isinstance(config, dict): + return {k: map_config_strings(v) for k, v in config.items()} + if isinstance(config, list): + return [map_config_strings(v) for v in config] + if isinstance(config, str): + if config in DTYPE_MAP: + return DTYPE_MAP[config] + if config in PRECISION_MAP: + return PRECISION_MAP[config] + if config in ACTIVATION_MAP: + return ACTIVATION_MAP[config] + if config == 'None': + return None + if config.startswith('jax.') or config.startswith('jax._src.'): + attr_path = config.replace('jax._src.nn.functions', 'jax.nn') + obj = jax + for part in attr_path.split('.')[1:]: + obj = getattr(obj, part) + return obj + return config + + +def build_model(architecture: str, config: dict): + """Construct a model from an architecture name and a plain config dict. + + Config keys the model has no field for are dropped with a notice (older + runs logged since-removed flags), so old configs keep reconstructing. + """ + canonical, flags = canonicalize_architecture(architecture) + + if canonical == 'diffusers_unet_simple': + # Third-party linen model, kept behind a lazy import and the BCHW seam + from diffusers import FlaxUNet2DConditionModel + from .general import BCHWModelWrapper + return BCHWModelWrapper(FlaxUNet2DConditionModel(**map_config_strings(config))) + + model_class = MODEL_REGISTRY.get(canonical) + if model_class is None: + raise ValueError( + f"Unknown architecture: {architecture}. " + f"Supported: {', '.join(MODEL_REGISTRY.keys())}") + + kwargs = {**map_config_strings(config), **flags} + valid_fields = {f.name for f in dataclasses.fields(model_class)} + dropped = sorted(set(kwargs) - valid_fields) + if dropped: + print(f"Dropping config keys not accepted by {model_class.__name__}: {dropped}") + kwargs = {k: v for k, v in kwargs.items() if k in valid_fields} + + return model_class(**kwargs) diff --git a/tests/test_config.py b/tests/test_config.py index dea975d..292d9f4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -108,3 +108,33 @@ def test_training_and_inference_share_schedule_presets(): 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) + + +def test_registry_builds_every_architecture(): + """Every registry entry must construct from a minimal string config - + the same call path training and inference share.""" + from flaxdiff.models.registry import MODEL_REGISTRY, build_model + + base = {"emb_features": 64, "dtype": "float32", "precision": "default", + "output_channels": 3, "patch_size": 4, "num_layers": 2, "num_heads": 2} + per_arch = { + "unet": {"feature_depths": [16, 32], "attention_configs": [None, None], + "num_res_blocks": 1, "num_middle_res_blocks": 1, + "activation": "swish", "norm_groups": 4}, + "uvit": {"num_layers": 4}, + "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}, + } + for name in MODEL_REGISTRY: + config = {**base, **per_arch.get(name, {})} + model = build_model(name, config) + assert type(model).__name__ == MODEL_REGISTRY[name].__name__ + + +def test_registry_suffix_canonicalization(): + from flaxdiff.models.registry import canonicalize_architecture + + name, flags = canonicalize_architecture("hybrid_dit+2d+hilbert") + assert name == "hybrid_dit" + assert flags == {"use_2d_fusion": True, "use_hilbert": True} diff --git a/training.py b/training.py index 834380a..39c8614 100644 --- a/training.py +++ b/training.py @@ -5,14 +5,8 @@ from functools import partial import flax.training.dynamic_scale import jax.experimental.multihost_utils -from flaxdiff.models.common import kernel_init -from flaxdiff.models.simple_unet import Unet -from flaxdiff.models.simple_vit import UViT, SimpleUDiT -from flaxdiff.models.simple_dit import SimpleDiT -from flaxdiff.models.simple_mmdit import SimpleMMDiT, HierarchicalMMDiT -from flaxdiff.models.ssm_dit import HybridSSMAttentionDiT +from flaxdiff.models.registry import build_model, canonicalize_architecture from flaxdiff.predictors import get_diffusion_preset -from diffusers import FlaxUNet2DConditionModel from flaxdiff.samplers.euler import EulerAncestralSampler import struct as st import flax @@ -109,29 +103,9 @@ def boolean_string(s): parser.add_argument('--noise_schedule', type=str, default='edm', choices=['cosine', 'karras', 'edm'], help='Noise schedule') -parser.add_argument('--architecture', type=str, - choices=[ - "unet", - # "uvit", - "simple_udit", - "diffusers_unet_simple", - "simple_dit", - "simple_mmdit", - "hierarchical_mmdit", - # 'uvit-hilbert', - 'simple_dit+hilbert', - 'simple_dit+zigzag', - "simple_udit+hilbert", - "simple_mmdit+hilbert", - "hierarchical_mmdit+hilbert", - "hybrid_dit", - "hybrid_dit+hilbert", - "hybrid_dit+zigzag", - "hybrid_dit+2d", - "hybrid_dit+2d+hilbert", - "hybrid_dit+2d+zigzag", - ], - default="unet", help='Architecture to use') +# Any name from flaxdiff.models.registry, optionally with +2d/+hilbert/+zigzag +# suffixes; validated by build_model against the registry itself +parser.add_argument('--architecture', type=str, default="unet", help='Architecture to use') 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') @@ -233,26 +207,6 @@ def main(args): print(f"Number of devices: {jax.device_count()}") print(f"Local devices: {jax.local_devices()}") - DTYPE_MAP = { - 'bfloat16': jnp.bfloat16, - 'float32': jnp.float32, - 'None': None, - None: None, - } - - PRECISION_MAP = { - 'high': jax.lax.Precision.HIGH, - 'default': jax.lax.Precision.DEFAULT, - 'highes': jax.lax.Precision.HIGHEST, - 'None': None, - None: None, - } - - ACTIVATION_MAP = { - 'swish': jax.nn.swish, - 'mish': jax.nn.mish, - } - OPTIMIZER_MAP = { 'adam' : optax.adam, 'adamw' : optax.adamw, @@ -263,8 +217,10 @@ def main(args): if args.checkpoint_fs == 'gcs': CHECKPOINT_DIR = f"gs://{CHECKPOINT_DIR}" - DTYPE = DTYPE_MAP[args.dtype] - PRECISION = PRECISION_MAP[args.precision] + # Model configs carry plain strings; build_model resolves them, so the + # logged wandb config is exactly the construction input + DTYPE = args.dtype + PRECISION = args.precision GRAIN_WORKER_COUNT = args.GRAIN_WORKER_COUNT GRAIN_READ_THREAD_COUNT = args.GRAIN_READ_THREAD_COUNT @@ -339,23 +295,10 @@ def main(args): INPUT_CHANNELS = 4 DIFFUSION_INPUT_SIZE = DIFFUSION_INPUT_SIZE // 8 - use_hilbert = args.use_hilbert - use_zigzag = args.use_zigzag - use_2d_fusion = args.use_2d_fusion - architecture_name = args.architecture - # parse '+2d' first so it can coexist with '+hilbert' or '+zigzag' - if '+2d' in architecture_name: - architecture_name = architecture_name.replace('+2d', '') - print("Will use 2D state fusion (Spatial-Mamba-style) in SSM blocks") - use_2d_fusion = True - if 'hilbert' in architecture_name: - architecture_name = architecture_name.split('+')[0] - print("Will use Hilbert Patch Reordering") - use_hilbert = True - elif 'zigzag' in architecture_name: - architecture_name = architecture_name.split('+')[0] - print("Will use Zigzag (serpentine) Patch Reordering") - use_zigzag = True + architecture_name, suffix_flags = canonicalize_architecture(args.architecture) + use_hilbert = args.use_hilbert or suffix_flags.get('use_hilbert', False) + use_zigzag = args.use_zigzag or suffix_flags.get('use_zigzag', False) + use_2d_fusion = args.use_2d_fusion or suffix_flags.get('use_2d_fusion', False) assert not (use_hilbert and use_zigzag), "use_hilbert and use_zigzag are mutually exclusive" if 'diffusers' in architecture_name: @@ -369,9 +312,8 @@ def main(args): } - MODEL_ARCHITECUTRES = { + ARCHITECTURE_KWARGS = { "unet": { - "class": Unet, "kwargs": { "feature_depths": args.feature_depths, "attention_configs": attention_configs, @@ -383,7 +325,6 @@ def main(args): }, }, "uvit": { - "class": UViT, "kwargs": { "patch_size": args.patch_size, "num_layers": args.num_layers, @@ -396,7 +337,6 @@ def main(args): }, }, "simple_udit": { - "class": SimpleUDiT, "kwargs": { "patch_size": args.patch_size, "num_layers": args.num_layers, @@ -407,7 +347,6 @@ def main(args): }, }, "simple_dit": { - "class": SimpleDiT, "kwargs": { "patch_size": args.patch_size, "num_layers": args.num_layers, @@ -419,7 +358,6 @@ def main(args): }, }, "simple_mmdit": { - "class": SimpleMMDiT, "kwargs": { "patch_size": args.patch_size, "num_layers": args.num_layers, @@ -430,7 +368,6 @@ def main(args): }, }, "hierarchical_mmdit": { - "class": HierarchicalMMDiT, "kwargs": { "base_patch_size": args.patch_size // 2, # Use half the patch size for base "emb_features": (args.emb_features - 256, args.emb_features, args.emb_features + 256), # Default dims per stage @@ -442,7 +379,6 @@ def main(args): }, }, "hybrid_dit": { - "class": HybridSSMAttentionDiT, "kwargs": { "patch_size": args.patch_size, "num_layers": args.num_layers, @@ -457,7 +393,6 @@ def main(args): }, }, "diffusers_unet_simple": { - "class": FlaxUNet2DConditionModel, "kwargs": { "sample_size": DIFFUSION_INPUT_SIZE, "in_channels": INPUT_CHANNELS, @@ -472,8 +407,7 @@ def main(args): } } - model_architecture = MODEL_ARCHITECUTRES[architecture_name]['class'] - model_config.update(MODEL_ARCHITECUTRES[architecture_name]['kwargs']) + model_config.update(ARCHITECTURE_KWARGS[architecture_name]['kwargs']) if architecture_name == 'uvit': model_config['emb_features'] = 768 @@ -564,19 +498,7 @@ def main(args): print("Experiment_Name:", experiment_name) - # 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: - from flaxdiff.models.general import BCHWModelWrapper - model = BCHWModelWrapper(model) + model = build_model(architecture_name, model_config) learning_rate = CONFIG['learning_rate'] optimizer = OPTIMIZER_MAP[args.optimizer] From 15e77edb39bd902ecc3b7ba1b47f63902ed67b5e Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 00:11:04 -0400 Subject: [PATCH 3/6] refactor(models): one shared DiT backbone and one modulated block The patchify/scan-order/sincos/conditioning/final-head skeleton was repeated near-verbatim in SimpleDiT, HybridSSMAttentionDiT, SimpleUDiT and SimpleMMDiT, and the adaLN-Zero block body was copy-pasted three times (attention, SSM, and dual-cond variants). dit_common.py now owns the sandwich (PatchSequenceEmbed / ConditioningEmbed / PatchSequenceOutput) and a single ModulatedBlock with a pluggable mixer (attention | S5 ssm); the S5 layers move to s5.py so the block can use them without importing the full hybrid model. SimpleDiT (291 -> 90 lines) and HybridSSMAttentionDiT (597 -> 142) are now just block arrangements; SimpleUDiT builds its U-skips from the same block. Also: unpatchify takes an explicit patch grid so non-square inputs work (the last remaining xfail from the audit - the whole suite is now xfail-free), and attention gains optional SD3-style qk-norm (RMSNorm on q/k) for bf16 logit safety, exposed on every DiT-family model. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/attention.py | 8 + flaxdiff/models/dit_common.py | 308 +++++++++++++++++++ flaxdiff/models/s5.py | 234 +++++++++++++++ flaxdiff/models/simple_dit.py | 286 +++--------------- flaxdiff/models/simple_vit.py | 8 +- flaxdiff/models/ssm_dit.py | 548 +++------------------------------- flaxdiff/models/vit_common.py | 15 +- tests/test_models.py | 1 - 8 files changed, 651 insertions(+), 757 deletions(-) create mode 100644 flaxdiff/models/dit_common.py create mode 100644 flaxdiff/models/s5.py diff --git a/flaxdiff/models/attention.py b/flaxdiff/models/attention.py index 8106bf5..1b50812 100644 --- a/flaxdiff/models/attention.py +++ b/flaxdiff/models/attention.py @@ -24,6 +24,7 @@ class NormalAttention(nn.Module): use_bias: bool = True # 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) def setup(self): inner_dim = self.dim_head * self.heads @@ -40,6 +41,10 @@ def setup(self): self.key = dense(name="to_k") self.value = dense(name="to_v") + if self.qk_norm: + self.q_norm = nn.RMSNorm(dtype=self.dtype, name="q_norm") + self.k_norm = nn.RMSNorm(dtype=self.dtype, name="k_norm") + self.proj_attn = nn.DenseGeneral( self.query_dim, axis=(-2, -1), @@ -64,6 +69,9 @@ def __call__(self, x, context=None): query = self.query(x) key = self.key(context) value = self.value(context) + if self.qk_norm: + 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, diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py new file mode 100644 index 0000000..ed97f31 --- /dev/null +++ b/flaxdiff/models/dit_common.py @@ -0,0 +1,308 @@ +""" +Shared machinery for the DiT family. + +Every DiT-style model here is the same sandwich: patchify in some scan order, +add a 2D sincos position signal, run adaLN-Zero modulated blocks over the +token sequence, and unpatchify back. The only real differences between the +models are the token mixer inside the block (attention or S5 SSM) and how the +blocks are arranged (plain stack, U-shaped skips, hybrid patterns). This +module owns the sandwich; the model files just arrange blocks. +""" + +import jax +import jax.numpy as jnp +from flax import linen as nn +from typing import Optional +from flax.typing import Dtype, PrecisionLike + +from .common import FourierEmbedding, TimeProjection +from .vit_common import PatchEmbedding, RotaryEmbedding, RoPEAttention, AdaLNParams, unpatchify +from .s5 import S5Layer, BidirectionalS5Layer, SpatialFusionConv +from .hilbert import ( + hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify, + zigzag_indices, zigzag_patchify, + build_2d_sincos_pos_embed, +) + +SCAN_ORDERS = ('raster', 'hilbert', 'zigzag') + + +def scan_indices(scan_order: str, H_P: int, W_P: int): + """Forward permutation for a scan order (None for raster).""" + if scan_order == 'hilbert': + return hilbert_indices(H_P, W_P) + if scan_order == 'zigzag': + return zigzag_indices(H_P, W_P) + return None + + +class PatchSequenceEmbed(nn.Module): + """Patchify in raster/hilbert/zigzag order and add the 2D sincos signal. + + Returns (tokens, inv_idx) - inv_idx restores row-major order on the way + out and is None for raster. + """ + patch_size: int + emb_features: int + scan_order: str = 'raster' + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + def setup(self): + assert self.scan_order in SCAN_ORDERS, f"Unknown scan order {self.scan_order}" + if self.scan_order == 'raster': + self.patch_embed = PatchEmbedding( + patch_size=self.patch_size, + embedding_dim=self.emb_features, + dtype=self.dtype, + precision=self.precision, + name="patch_embed", + ) + else: + # Raw patch extraction + Dense instead of the Conv-based embedding + self.scan_proj = nn.Dense( + features=self.emb_features, + dtype=self.dtype, + precision=self.precision, + name="hilbert_projection", + ) + + def __call__(self, x): + 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" + H_P, W_P = H // self.patch_size, W // self.patch_size + + inv_idx = None + if self.scan_order == 'hilbert': + patches_raw, inv_idx = hilbert_patchify(x, self.patch_size) + tokens = self.scan_proj(patches_raw) + elif self.scan_order == 'zigzag': + patches_raw, inv_idx = zigzag_patchify(x, self.patch_size) + tokens = self.scan_proj(patches_raw) + else: + tokens = self.patch_embed(x) + + # Fixed 2D sincos position embedding (order-invariant spatial signal). + # For hilbert/zigzag, reorder the row-major embedding to match the scan + # so each token gets the sincos for its real 2D position. + pos_embed = build_2d_sincos_pos_embed(self.emb_features, H_P, W_P) + pos_embed = jnp.asarray(pos_embed, dtype=tokens.dtype) + idx = scan_indices(self.scan_order, H_P, W_P) + if idx is not None: + pos_embed = pos_embed[idx] + tokens = tokens + pos_embed[None, :, :] + return tokens, inv_idx + + +class ConditioningEmbed(nn.Module): + """Fourier time embedding + mean-pooled text projection, summed into the + single conditioning vector the adaLN modulation consumes.""" + emb_features: int + mlp_ratio: int = 4 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + def setup(self): + self.time_embed = nn.Sequential([ + FourierEmbedding(features=self.emb_features), + TimeProjection(features=self.emb_features * self.mlp_ratio), + nn.Dense(features=self.emb_features, dtype=self.dtype, precision=self.precision), + ], name="time_embed") + self.text_proj = nn.Dense( + features=self.emb_features, dtype=self.dtype, + precision=self.precision, name="text_context_proj") + + def __call__(self, temb, textcontext=None): + cond_emb = self.time_embed(temb) + if textcontext is not None: + text_emb = self.text_proj(textcontext) + cond_emb = cond_emb + jnp.mean(text_emb, axis=1) + return cond_emb + + +class PatchSequenceOutput(nn.Module): + """Final norm + zero-init fp32 head + unpatchify for any scan order and + any (non-square included) patch grid.""" + patch_size: int + output_channels: int + learn_sigma: bool = False + norm_epsilon: float = 1e-5 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + def setup(self): + self.final_norm = nn.LayerNorm( + epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") + output_dim = self.patch_size * self.patch_size * self.output_channels + if self.learn_sigma: + output_dim *= 2 + self.final_proj = nn.Dense( + features=output_dim, + dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 + precision=self.precision, + kernel_init=nn.initializers.zeros, + name="final_proj", + ) + + def __call__(self, tokens, inv_idx, H, W): + x_out = self.final_proj(self.final_norm(tokens)) + if self.learn_sigma: + x_out, _x_logvar = jnp.split(x_out, 2, axis=-1) + if inv_idx is not None: + return hilbert_unpatchify(x_out, inv_idx, self.patch_size, H, W, self.output_channels) + return unpatchify(x_out, channels=self.output_channels, + H_P=H // self.patch_size, W_P=W // self.patch_size) + + +class ModulatedBlock(nn.Module): + """adaLN-Zero modulated residual block with a pluggable token mixer. + + mixer='attention' gives the standard DiT block (RoPE attention); + mixer='ssm' replaces attention with a bidirectional S5 scan, optionally + followed by Spatial-Mamba style 2D state fusion. freqs_cis is unused by + the SSM mixer but kept in the interface so blocks are interchangeable. + """ + features: int + num_heads: int + rope_emb: RotaryEmbedding = None + mixer: str = 'attention' + 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 + use_gating: bool = True + qk_norm: bool = False + # ssm mixer options + ssm_state_dim: int = 64 + bidirectional_ssm: bool = True + use_2d_fusion: bool = False + scan_order: str = 'raster' # needed to un-permute for the 2D fusion conv + + def setup(self): + assert self.mixer in ('attention', 'ssm'), f"Unknown mixer {self.mixer}" + hidden_features = int(self.features * self.mlp_ratio) + + self.ada_params_module = AdaLNParams( + self.features, dtype=self.dtype, precision=self.precision) + self.norm1 = nn.LayerNorm( + epsilon=self.norm_epsilon, use_scale=False, use_bias=False, + dtype=self.dtype, name="norm1") + self.norm2 = nn.LayerNorm( + epsilon=self.norm_epsilon, use_scale=False, use_bias=False, + dtype=self.dtype, name="norm2") + + if self.mixer == 'attention': + self.attention = RoPEAttention( + query_dim=self.features, + heads=self.num_heads, + dim_head=self.features // self.num_heads, + dtype=self.dtype, + precision=self.precision, + use_bias=True, + qk_norm=self.qk_norm, + force_fp32_for_softmax=self.force_fp32_for_softmax, + rope_emb=self.rope_emb, + ) + else: + ssm_cls = BidirectionalS5Layer if self.bidirectional_ssm else S5Layer + self.ssm = ssm_cls( + features=self.features, + state_dim=self.ssm_state_dim, + dtype=self.dtype, + precision=self.precision, + name="ssm", + ) + if self.use_2d_fusion: + assert self.scan_order in SCAN_ORDERS, f"Unknown scan_order {self.scan_order}" + self.spatial_fusion = SpatialFusionConv( + features=self.features, + dilations=(1, 2, 3), + kernel_size=3, + dtype=self.dtype, + precision=self.precision, + name="spatial_fusion", + ) + + self.mlp = nn.Sequential([ + nn.Dense(features=hidden_features, dtype=self.dtype, precision=self.precision), + nn.gelu, + 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 + # square patch grid; S is a static python int at trace time + import math + H_P = math.isqrt(S) + W_P = H_P + assert H_P * W_P == S, ( + f"2D fusion requires a square patch grid; got S={S} which is not a " + f"perfect square.") + + # Index arrays are constant at JIT time so computing both directions is free + scan_fwd = scan_indices(self.scan_order, H_P, W_P) + if scan_fwd is not None: + scan_inv = inverse_permutation(scan_fwd, S) + ssm_rm = ssm_output[:, scan_inv, :] + else: + ssm_rm = ssm_output + + y_2d = ssm_rm.reshape(B, H_P, W_P, F) + y_fused_2d = self.spatial_fusion(y_2d) + y_fused_rm = y_fused_2d.reshape(B, S, F) + + if scan_fwd is not None: + y_fused = y_fused_rm[:, scan_fwd, :] + else: + y_fused = y_fused_rm + return y_fused + + @nn.compact + def __call__(self, x, conditioning, freqs_cis, train: bool = False): + scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( + self.ada_params_module(conditioning), 6, axis=-1 + ) + + # --- Mixer path (attention or SSM) --- + residual = x + x_modulated = self.norm1(x) * (1 + scale_attn) + shift_attn + if self.mixer == 'attention': + mixer_output = self.attention(x_modulated, context=None, freqs_cis=freqs_cis) + else: + mixer_output = self.ssm(x_modulated) + if self.use_2d_fusion: + mixer_output = self._apply_2d_fusion(mixer_output) + mixer_output = self.dropout(mixer_output, deterministic=not train) + + if self.use_gating: + x = residual + gate_attn * mixer_output + else: + x = residual + mixer_output + + # --- MLP path --- + residual = x + x_mlp_modulated = self.norm2(x) * (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 + else: + x = residual + mlp_output + return x + + +def neutralized_rope_freqs(rope: RotaryEmbedding, seq_len: int, scan_order: str): + """RoPE frequencies for the sequence, neutralized to identity for + hilbert/zigzag scans where the 1D index is not a 2D position (the 2D + sincos embedding already carries position there).""" + freqs_cos, freqs_sin = rope(seq_len=seq_len) + if scan_order != 'raster': + freqs_cos = jnp.ones_like(freqs_cos) + freqs_sin = jnp.zeros_like(freqs_sin) + return freqs_cos, freqs_sin diff --git a/flaxdiff/models/s5.py b/flaxdiff/models/s5.py new file mode 100644 index 0000000..95abce7 --- /dev/null +++ b/flaxdiff/models/s5.py @@ -0,0 +1,234 @@ +""" +S5 state-space layers (diagonal SSM with associative_scan, HiPPO init) and the +Spatial-Mamba style 2D state fusion conv. Extracted from ssm_dit.py so the +shared DiT block can use them without importing the full hybrid model. +""" + +import jax +import jax.numpy as jnp +from flax import linen as nn +from typing import Optional, Tuple +from flax.typing import Dtype, PrecisionLike + +# --- S5 SSM Layer --- + +def hippo_log_a_real_init(key, shape, dtype=jnp.float32): + """HiPPO-diag init: A_real_n = -(n + 0.5), stored as log of the negative.""" + state_dim = shape[0] + n = jnp.arange(state_dim, dtype=dtype) + return jnp.log(n + 0.5).astype(dtype) + + +def hippo_a_imag_init(key, shape, dtype=jnp.float32): + """HiPPO-diag init: A_imag_n = pi * n.""" + state_dim = shape[0] + n = jnp.arange(state_dim, dtype=dtype) + return (jnp.pi * n).astype(dtype) + + +class S5Layer(nn.Module): + """S5 layer with diagonal complex state matrix. + x_k = A * x_{k-1} + B * u_k + y_k = Re(C * x_k) + D * u_k + """ + features: int + state_dim: int = 64 + dt_min: float = 0.001 + dt_max: float = 0.1 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + @nn.compact + 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 + log_A_real = self.param( + 'log_A_real', + hippo_log_a_real_init, + (self.state_dim,) + ) + A_imag = self.param( + 'A_imag', + hippo_a_imag_init, + (self.state_dim,) + ) + + # B: input-to-state projection [state_dim, F] + B_re = self.param( + 'B_re', + nn.initializers.lecun_normal(), + (self.state_dim, F) + ) + B_im = self.param( + 'B_im', + nn.initializers.lecun_normal(), + (self.state_dim, F) + ) + + # C: state-to-output projection [F, state_dim], lecun_normal as in S5 + C_re = self.param( + 'C_re', + nn.initializers.lecun_normal(), + (F, self.state_dim) + ) + C_im = self.param( + 'C_im', + nn.initializers.lecun_normal(), + (F, self.state_dim) + ) + + # D: skip connection, N(0,1) per channel as in S5 + D = self.param('D', nn.initializers.normal(stddev=1.0), (F,)) + + # dt: discretization timestep, learned per state dim so each state + # channel can model its own time scale + log_dt = self.param( + 'log_dt', + lambda key, shape: jax.random.uniform( + key, shape, + minval=jnp.log(self.dt_min), + maxval=jnp.log(self.dt_max) + ), + (self.state_dim,) + ) + dt = jnp.exp(log_dt) # [state_dim] + + # Construct complex A and discretize + A_real = -jnp.exp(log_A_real) # negative real part for stability + A_diag = A_real + 1j * A_imag # [state_dim] + + # ZOH discretization: A_bar = exp(A * dt), B_bar = (A_bar - I) * A^{-1} * B + A_bar = jnp.exp(A_diag * dt) # [state_dim], complex + + B_complex = B_re + 1j * B_im + B_bar = ((A_bar[:, None] - 1.0) / (A_diag[:, None] + 1e-8)) * B_complex # [state_dim, F] + + C_complex = C_re + 1j * C_im + + # --- Parallel Scan --- + # x_k = A_bar * x_{k-1} + B_bar @ u_k via associative scan with + # (a1, b1) * (a2, b2) = (a1 * a2, a2 * b1 + b2) + u_float = u.astype(jnp.float32) + Bu = jnp.einsum('bsf,nf->bsn', u_float, B_bar) # [B, S, state_dim] + + A_bar_expanded = jnp.broadcast_to(A_bar[None, None, :], (B, S, self.state_dim)) + + def binary_operator(e1, e2): + a1, b1 = e1 + a2, b2 = e2 + return a1 * a2, a2 * b1 + b2 + + _, x_states = jax.lax.associative_scan( + binary_operator, + (A_bar_expanded, Bu), + axis=1 + ) + # x_states: [B, S, state_dim] (complex) + + # y_k = Re(C @ x_k) + D * u_k + y_complex = jnp.einsum('fn,bsn->bsf', C_complex, x_states) # [B, S, F] + y = y_complex.real + + # skip connection + y = y + D[None, None, :] * u_float # [B, S, F] + + # cast back to input dtype + if self.dtype is not None: + y = y.astype(self.dtype) + else: + y = y.astype(u.dtype) + + return y + + +# --- Bidirectional S5 --- + +class BidirectionalS5Layer(nn.Module): + """Runs forward and backward S5 scans, concats and projects back to features. + Patches have no inherent direction, so scan both ways. + """ + features: int + state_dim: int = 64 + dt_min: float = 0.001 + dt_max: float = 0.1 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + @nn.compact + def __call__(self, u): + # u: [B, S, F] + y_fwd = S5Layer( + features=self.features, + state_dim=self.state_dim, + dt_min=self.dt_min, + dt_max=self.dt_max, + dtype=self.dtype, + precision=self.precision, + name="s5_forward" + )(u) + + # backward scan: reverse input, scan, reverse output + u_rev = jnp.flip(u, axis=1) + y_bwd_rev = S5Layer( + features=self.features, + state_dim=self.state_dim, + dt_min=self.dt_min, + dt_max=self.dt_max, + dtype=self.dtype, + precision=self.precision, + name="s5_backward" + )(u_rev) + y_bwd = jnp.flip(y_bwd_rev, axis=1) + + y_cat = jnp.concatenate([y_fwd, y_bwd], axis=-1) # [B, S, 2F] + + y = nn.Dense( + features=self.features, + dtype=self.dtype, + precision=self.precision, + name="out_proj" + )(y_cat) + + return y + + +# --- 2D state fusion (Spatial-Mamba style) --- + +class SpatialFusionConv(nn.Module): + """Multi-dilation depthwise 2D convs summed as a residual over the SSM output grid. + The 1D scan scrambles 2D locality; this recovers a direction-balanced local + receptive field. Kernels are zero-init so the fusion starts as a pass-through. + """ + features: int + dilations: Tuple[int, ...] = (1, 2, 3) + kernel_size: int = 3 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + @nn.compact + def __call__(self, y_2d): + # y_2d: [B, H_P, W_P, F], SSM output reshaped to a row-major grid + out = y_2d + for dil in self.dilations: + dw = nn.Conv( + features=self.features, + kernel_size=(self.kernel_size, self.kernel_size), + strides=(1, 1), + padding='SAME', + kernel_dilation=(dil, dil), + feature_group_count=self.features, # depthwise + use_bias=False, + kernel_init=nn.initializers.zeros, + dtype=self.dtype, + precision=self.precision, + name=f"dwconv_dil{dil}", + )(y_2d) + out = out + dw + return out + + +# --- SSM DiT Block --- diff --git a/flaxdiff/models/simple_dit.py b/flaxdiff/models/simple_dit.py index 70fa752..3aa7542 100644 --- a/flaxdiff/models/simple_dit.py +++ b/flaxdiff/models/simple_dit.py @@ -1,109 +1,17 @@ -import jax import jax.numpy as jnp from flax import linen as nn -from typing import Callable, Any, Optional, Tuple, Sequence, Union -import einops -from functools import partial - -# Re-use existing components if they are suitable -from .vit_common import PatchEmbedding, unpatchify, RotaryEmbedding, RoPEAttention, AdaLNParams -from .common import kernel_init, FourierEmbedding, TimeProjection -# Using NormalAttention for RoPE integration -from .attention import NormalAttention +from typing import Optional from flax.typing import Dtype, PrecisionLike -# Use our improved Hilbert implementation -from .hilbert import ( - hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify, - zigzag_indices, zigzag_patchify, zigzag_unpatchify, - build_2d_sincos_pos_embed, +from .dit_common import ( + PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, + ModulatedBlock, neutralized_rope_freqs, ) +from .vit_common import RotaryEmbedding -# --- DiT Block --- -class DiTBlock(nn.Module): - features: int - num_heads: int - rope_emb: RotaryEmbedding - 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 - use_gating: bool = True # Add flag to easily disable gating - - def setup(self): - hidden_features = int(self.features * self.mlp_ratio) - # Get modulation parameters (scale, shift, gates) - self.ada_params_module = AdaLNParams( # Use the modified module - self.features, dtype=self.dtype, precision=self.precision) - - # Layer Norms - one before Attn, one before MLP - self.norm1 = nn.LayerNorm(epsilon=self.norm_epsilon, use_scale=False, use_bias=False, dtype=self.dtype, name="norm1") - self.norm2 = nn.LayerNorm(epsilon=self.norm_epsilon, use_scale=False, use_bias=False, dtype=self.dtype, name="norm2") - - self.attention = RoPEAttention( - query_dim=self.features, - heads=self.num_heads, - dim_head=self.features // self.num_heads, - dtype=self.dtype, - precision=self.precision, - use_bias=True, - force_fp32_for_softmax=self.force_fp32_for_softmax, - rope_emb=self.rope_emb - ) - - self.mlp = nn.Sequential([ - nn.Dense(features=hidden_features, dtype=self.dtype, precision=self.precision), - nn.gelu, # Or swish as specified in SimpleDiT? Consider consistency. - 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, 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( - self.ada_params_module(conditioning), 6, axis=-1 - ) - - # --- Attention Path --- - residual = x - norm_x_attn = self.norm1(x) - # 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 - else: - x = residual + attn_output # Original DiT style without gate - - # --- MLP Path --- - residual = x - norm_x_mlp = self.norm2(x) # Apply second LayerNorm - # 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 - else: - x = residual + mlp_output # Original DiT style without gate - - return x - - -# --- Patch Embedding (reuse or define if needed) --- -# Assuming PatchEmbedding exists in simple_vit.py and is suitable - -# --- DiT --- class SimpleDiT(nn.Module): + """Standard DiT: a plain stack of adaLN-Zero attention blocks.""" output_channels: int = 3 patch_size: int = 16 emb_features: int = 768 @@ -113,178 +21,68 @@ class SimpleDiT(nn.Module): dropout_rate: float = 0.0 # Typically 0 for diffusion dtype: Optional[Dtype] = None precision: PrecisionLike = None - # Passed down, but RoPEAttention uses NormalAttention 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 + learn_sigma: bool = False + qk_norm: bool = False + 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.patch_embed = PatchEmbedding( + self.embed = PatchSequenceEmbed( patch_size=self.patch_size, - embedding_dim=self.emb_features, + emb_features=self.emb_features, + scan_order=self.scan_order, dtype=self.dtype, - precision=self.precision + precision=self.precision, + ) + self.conditioning = ConditioningEmbed( + emb_features=self.emb_features, + mlp_ratio=self.mlp_ratio, + dtype=self.dtype, + precision=self.precision, ) - - assert not (self.use_hilbert and self.use_zigzag), \ - "use_hilbert and use_zigzag are mutually exclusive" - - # Add projection layer for Hilbert/zigzag reordered patches (raw patch - # extraction + Dense instead of the Conv-based PatchEmbedding) - if self.use_hilbert or self.use_zigzag: - self.hilbert_proj = nn.Dense( - features=self.emb_features, - dtype=self.dtype, - precision=self.precision, - name="hilbert_projection" - ) - - # Time embedding projection - self.time_embed = nn.Sequential([ - FourierEmbedding(features=self.emb_features), - TimeProjection(features=self.emb_features * - self.mlp_ratio), # Project to MLP dim - nn.Dense(features=self.emb_features, dtype=self.dtype, - precision=self.precision) # Final projection - ]) - - # Text context projection (if used) - # Assuming textcontext is already projected to some dimension, project it to match emb_features - # This might need adjustment based on how text context is provided - self.text_proj = nn.Dense(features=self.emb_features, dtype=self.dtype, - precision=self.precision, name="text_context_proj") - - # Rotary Positional Embedding - # Max length needs to be estimated or set large enough. - # For images, seq len = (H/P) * (W/P). Example: 256/16 * 256/16 = 16*16 = 256 - # Add 1 if a class token is used, or more for text tokens if concatenated. - # Let's assume max seq len accommodates patches + time + text tokens if needed, or just patches. - # If only patches use RoPE, max_len = max_image_tokens - # If time/text are concatenated *before* blocks, max_len needs to include them. - # DiT typically applies PE only to patch tokens. Let's follow that. - # max_len should be max number of patches. - # Example: max image size 512x512, patch 16 -> (512/16)^2 = 32^2 = 1024 patches self.rope = RotaryEmbedding( - dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) # Dim per head - - # Transformer Blocks + dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) self.blocks = [ - DiTBlock( + ModulatedBlock( features=self.emb_features, num_heads=self.num_heads, + rope_emb=self.rope, + 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, - rope_emb=self.rope, # Pass RoPE instance + qk_norm=self.qk_norm, name=f"dit_block_{i}" ) for i in range(self.num_layers) ] - - # Final Layer (Normalization + Linear Projection) - self.final_norm = nn.LayerNorm( - epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") - # self.final_norm = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") - - # Predict patch pixels + potentially sigma - output_dim = self.patch_size * self.patch_size * self.output_channels - if self.learn_sigma: - output_dim *= 2 # Predict both mean and variance (or log_variance) - - self.final_proj = nn.Dense( - features=output_dim, - dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 + 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, - kernel_init=nn.initializers.zeros, # Initialize final layer to zero - name="final_proj" ) @nn.compact 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" - - # Compute dimensions in terms of patches - H_P = H // self.patch_size - W_P = W // self.patch_size - - # 1. Patch Embedding (raster, hilbert or zigzag order) - scan_inv_idx = None - if self.use_hilbert: - patches_raw, scan_inv_idx = hilbert_patchify(x, self.patch_size) - patches = self.hilbert_proj(patches_raw) - elif self.use_zigzag: - patches_raw, scan_inv_idx = zigzag_patchify(x, self.patch_size) - patches = self.hilbert_proj(patches_raw) - else: - patches = self.patch_embed(x) - - # keep the old variable name for downstream code - hilbert_inv_idx = scan_inv_idx + x_seq, inv_idx = self.embed(x) + cond_emb = self.conditioning(temb, textcontext) + freqs_cis = neutralized_rope_freqs(self.rope, x_seq.shape[1], self.scan_order) - num_patches = patches.shape[1] - - # Add a fixed 2D sincos position embedding (order-invariant spatial signal). - # For hilbert/zigzag, reorder the row-major embedding to match the scan so - # each token gets the sincos for its real 2D position. - pos_embed_rm = build_2d_sincos_pos_embed(self.emb_features, H_P, W_P) - pos_embed_rm = jnp.asarray(pos_embed_rm, dtype=patches.dtype) - if self.use_hilbert: - scan_idx = hilbert_indices(H_P, W_P) - pos_embed = pos_embed_rm[scan_idx] - elif self.use_zigzag: - scan_idx = zigzag_indices(H_P, W_P) - pos_embed = pos_embed_rm[scan_idx] - else: - pos_embed = pos_embed_rm - patches = patches + pos_embed[None, :, :] - - x_seq = patches - - # 2. Prepare Conditioning Signal (Time + Text Context) - t_emb = self.time_embed(temb) # Shape: [B, emb_features] - - cond_emb = t_emb - if textcontext is not None: - # Shape: [B, num_text_tokens, emb_features] - text_emb = self.text_proj(textcontext) - # Pool or select text embedding (e.g., mean pool or use CLS token) - # Assuming mean pooling for simplicity - # Shape: [B, emb_features] - text_emb_pooled = jnp.mean(text_emb, axis=1) - cond_emb = cond_emb + text_emb_pooled # Combine time and text embeddings - - # 3. Apply RoPE - # Get RoPE frequencies for the sequence length (number of patches) - # Shape [num_patches, D_head/2] - freqs_cos, freqs_sin = self.rope(seq_len=num_patches) - # With hilbert/zigzag the sequence index is not a 2D position and RoPE - # distances would be wrong, so override RoPE with identity freqs (cos=1, - # sin=0) - the 2D sincos above already carries position. - if self.use_hilbert or self.use_zigzag: - freqs_cos = jnp.ones_like(freqs_cos) - freqs_sin = jnp.zeros_like(freqs_sin) - - # 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), train=train) - - # 5. Final Layer - x_out = self.final_norm(x_seq) - # Shape: [B, num_patches, patch_pixels (*2 if learn_sigma)] - x_out = self.final_proj(x_out) + x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=freqs_cis, train=train) - # 6. Unpatchify (hilbert_unpatchify works for any inv_idx, so zigzag too) - if self.use_hilbert or self.use_zigzag: - if self.learn_sigma: - x_mean, _x_logvar = jnp.split(x_out, 2, axis=-1) - return hilbert_unpatchify(x_mean, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - return hilbert_unpatchify(x_out, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - if self.learn_sigma: - x_mean, _x_logvar = jnp.split(x_out, 2, axis=-1) - return unpatchify(x_mean, channels=self.output_channels) - return unpatchify(x_out, channels=self.output_channels) + return self.output(x_seq, inv_idx, H, W) diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py index 5daaa76..7dc4b2d 100644 --- a/flaxdiff/models/simple_vit.py +++ b/flaxdiff/models/simple_vit.py @@ -12,7 +12,7 @@ from functools import partial from .hilbert import hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify from .vit_common import unpatchify, PatchEmbedding, RotaryEmbedding, RoPEAttention, AdaLNParams -from .simple_dit import DiTBlock +from .dit_common import ModulatedBlock class UViT(nn.Module): @@ -304,7 +304,7 @@ def setup(self): ) self.down_blocks = [ - DiTBlock( + ModulatedBlock( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, @@ -318,7 +318,7 @@ def setup(self): ) for i in range(half_layers) ] - self.mid_block = DiTBlock( + self.mid_block = ModulatedBlock( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, @@ -340,7 +340,7 @@ def setup(self): ) for i in range(half_layers) ] self.up_blocks = [ - DiTBlock( + ModulatedBlock( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py index 94930d7..3ee4b03 100644 --- a/flaxdiff/models/ssm_dit.py +++ b/flaxdiff/models/ssm_dit.py @@ -1,395 +1,20 @@ """ -SSM (S5) based DiT blocks and a hybrid SSM-attention DiT. -S5 uses a diagonal state space with associative_scan, so it runs well on TPUs. +Hybrid SSM-attention DiT: interleaves linear-time S5 blocks with attention +blocks in a configurable ratio. The S5 layers themselves live in s5.py; the +block and the patchify/conditioning/output machinery live in dit_common.py. """ -import jax import jax.numpy as jnp from flax import linen as nn -from typing import Callable, Any, Optional, Tuple, Sequence, Union -import einops -from functools import partial - -from .vit_common import PatchEmbedding, unpatchify, RotaryEmbedding, RoPEAttention, AdaLNParams -from .common import kernel_init, FourierEmbedding, TimeProjection -from .attention import NormalAttention +from typing import Optional, Sequence from flax.typing import Dtype, PrecisionLike -from .hilbert import ( - hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify, - zigzag_indices, zigzag_patchify, zigzag_unpatchify, - build_2d_sincos_pos_embed, +from .dit_common import ( + PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, + ModulatedBlock, neutralized_rope_freqs, ) -from .simple_dit import DiTBlock - - -# --- S5 SSM Layer --- - -def hippo_log_a_real_init(key, shape, dtype=jnp.float32): - """HiPPO-diag init: A_real_n = -(n + 0.5), stored as log of the negative.""" - state_dim = shape[0] - n = jnp.arange(state_dim, dtype=dtype) - return jnp.log(n + 0.5).astype(dtype) - - -def hippo_a_imag_init(key, shape, dtype=jnp.float32): - """HiPPO-diag init: A_imag_n = pi * n.""" - state_dim = shape[0] - n = jnp.arange(state_dim, dtype=dtype) - return (jnp.pi * n).astype(dtype) - - -class S5Layer(nn.Module): - """S5 layer with diagonal complex state matrix. - x_k = A * x_{k-1} + B * u_k - y_k = Re(C * x_k) + D * u_k - """ - features: int - state_dim: int = 64 - dt_min: float = 0.001 - dt_max: float = 0.1 - dtype: Optional[Dtype] = None - precision: PrecisionLike = None - - @nn.compact - 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 - log_A_real = self.param( - 'log_A_real', - hippo_log_a_real_init, - (self.state_dim,) - ) - A_imag = self.param( - 'A_imag', - hippo_a_imag_init, - (self.state_dim,) - ) - - # B: input-to-state projection [state_dim, F] - B_re = self.param( - 'B_re', - nn.initializers.lecun_normal(), - (self.state_dim, F) - ) - B_im = self.param( - 'B_im', - nn.initializers.lecun_normal(), - (self.state_dim, F) - ) - - # C: state-to-output projection [F, state_dim], lecun_normal as in S5 - C_re = self.param( - 'C_re', - nn.initializers.lecun_normal(), - (F, self.state_dim) - ) - C_im = self.param( - 'C_im', - nn.initializers.lecun_normal(), - (F, self.state_dim) - ) - - # D: skip connection, N(0,1) per channel as in S5 - D = self.param('D', nn.initializers.normal(stddev=1.0), (F,)) - - # dt: discretization timestep, learned per state dim so each state - # channel can model its own time scale - log_dt = self.param( - 'log_dt', - lambda key, shape: jax.random.uniform( - key, shape, - minval=jnp.log(self.dt_min), - maxval=jnp.log(self.dt_max) - ), - (self.state_dim,) - ) - dt = jnp.exp(log_dt) # [state_dim] - - # Construct complex A and discretize - A_real = -jnp.exp(log_A_real) # negative real part for stability - A_diag = A_real + 1j * A_imag # [state_dim] - - # ZOH discretization: A_bar = exp(A * dt), B_bar = (A_bar - I) * A^{-1} * B - A_bar = jnp.exp(A_diag * dt) # [state_dim], complex - - B_complex = B_re + 1j * B_im - B_bar = ((A_bar[:, None] - 1.0) / (A_diag[:, None] + 1e-8)) * B_complex # [state_dim, F] - - C_complex = C_re + 1j * C_im - - # --- Parallel Scan --- - # x_k = A_bar * x_{k-1} + B_bar @ u_k via associative scan with - # (a1, b1) * (a2, b2) = (a1 * a2, a2 * b1 + b2) - u_float = u.astype(jnp.float32) - Bu = jnp.einsum('bsf,nf->bsn', u_float, B_bar) # [B, S, state_dim] - - A_bar_expanded = jnp.broadcast_to(A_bar[None, None, :], (B, S, self.state_dim)) - - def binary_operator(e1, e2): - a1, b1 = e1 - a2, b2 = e2 - return a1 * a2, a2 * b1 + b2 - - _, x_states = jax.lax.associative_scan( - binary_operator, - (A_bar_expanded, Bu), - axis=1 - ) - # x_states: [B, S, state_dim] (complex) +from .vit_common import RotaryEmbedding - # y_k = Re(C @ x_k) + D * u_k - y_complex = jnp.einsum('fn,bsn->bsf', C_complex, x_states) # [B, S, F] - y = y_complex.real - - # skip connection - y = y + D[None, None, :] * u_float # [B, S, F] - - # cast back to input dtype - if self.dtype is not None: - y = y.astype(self.dtype) - else: - y = y.astype(u.dtype) - - return y - - -# --- Bidirectional S5 --- - -class BidirectionalS5Layer(nn.Module): - """Runs forward and backward S5 scans, concats and projects back to features. - Patches have no inherent direction, so scan both ways. - """ - features: int - state_dim: int = 64 - dt_min: float = 0.001 - dt_max: float = 0.1 - dtype: Optional[Dtype] = None - precision: PrecisionLike = None - - @nn.compact - def __call__(self, u): - # u: [B, S, F] - y_fwd = S5Layer( - features=self.features, - state_dim=self.state_dim, - dt_min=self.dt_min, - dt_max=self.dt_max, - dtype=self.dtype, - precision=self.precision, - name="s5_forward" - )(u) - - # backward scan: reverse input, scan, reverse output - u_rev = jnp.flip(u, axis=1) - y_bwd_rev = S5Layer( - features=self.features, - state_dim=self.state_dim, - dt_min=self.dt_min, - dt_max=self.dt_max, - dtype=self.dtype, - precision=self.precision, - name="s5_backward" - )(u_rev) - y_bwd = jnp.flip(y_bwd_rev, axis=1) - - y_cat = jnp.concatenate([y_fwd, y_bwd], axis=-1) # [B, S, 2F] - - y = nn.Dense( - features=self.features, - dtype=self.dtype, - precision=self.precision, - name="out_proj" - )(y_cat) - - return y - - -# --- 2D state fusion (Spatial-Mamba style) --- - -class SpatialFusionConv(nn.Module): - """Multi-dilation depthwise 2D convs summed as a residual over the SSM output grid. - The 1D scan scrambles 2D locality; this recovers a direction-balanced local - receptive field. Kernels are zero-init so the fusion starts as a pass-through. - """ - features: int - dilations: Tuple[int, ...] = (1, 2, 3) - kernel_size: int = 3 - dtype: Optional[Dtype] = None - precision: PrecisionLike = None - - @nn.compact - def __call__(self, y_2d): - # y_2d: [B, H_P, W_P, F], SSM output reshaped to a row-major grid - out = y_2d - for dil in self.dilations: - dw = nn.Conv( - features=self.features, - kernel_size=(self.kernel_size, self.kernel_size), - strides=(1, 1), - padding='SAME', - kernel_dilation=(dil, dil), - feature_group_count=self.features, # depthwise - use_bias=False, - kernel_init=nn.initializers.zeros, - dtype=self.dtype, - precision=self.precision, - name=f"dwconv_dil{dil}", - )(y_2d) - out = out + dw - return out - - -# --- SSM DiT Block --- - -class SSMDiTBlock(nn.Module): - """Same interface as DiTBlock, but attention replaced with bidirectional S5. - freqs_cis is accepted for interface compat but unused by the SSM. - """ - features: int - num_heads: int # Not used by SSM, kept for interface compat - rope_emb: RotaryEmbedding # Not used by SSM, kept for interface compat - state_dim: int = 64 - mlp_ratio: int = 4 - dropout_rate: float = 0.0 - dtype: Optional[Dtype] = None - precision: PrecisionLike = None - force_fp32_for_softmax: bool = True # Ignored, interface compat - norm_epsilon: float = 1e-5 - use_gating: bool = True - bidirectional: bool = True - use_2d_fusion: bool = False # 2D state fusion (SpatialFusionConv) after the scan - scan_order: str = 'raster' # parent model's scan order, needed to un-permute for the conv - - def setup(self): - hidden_features = int(self.features * self.mlp_ratio) - - # AdaLN modulation, same as DiTBlock - self.ada_params_module = AdaLNParams( - self.features, dtype=self.dtype, precision=self.precision) - - self.norm1 = nn.LayerNorm( - epsilon=self.norm_epsilon, use_scale=False, use_bias=False, - dtype=self.dtype, name="norm1") - self.norm2 = nn.LayerNorm( - epsilon=self.norm_epsilon, use_scale=False, use_bias=False, - dtype=self.dtype, name="norm2") - - # S5 SSM layer (replaces attention) - ssm_cls = BidirectionalS5Layer if self.bidirectional else S5Layer - self.ssm = ssm_cls( - features=self.features, - state_dim=self.state_dim, - dtype=self.dtype, - precision=self.precision, - name="ssm" - ) - - # optional 2D state fusion after the SSM scan - if self.use_2d_fusion: - assert self.scan_order in ('raster', 'hilbert', 'zigzag'), \ - f"Unknown scan_order {self.scan_order}" - self.spatial_fusion = SpatialFusionConv( - features=self.features, - dilations=(1, 2, 3), - kernel_size=3, - dtype=self.dtype, - precision=self.precision, - name="spatial_fusion", - ) - - self.mlp = nn.Sequential([ - nn.Dense(features=hidden_features, dtype=self.dtype, precision=self.precision), - nn.gelu, - 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 - # square patch grid; S is a static python int at trace time - import math - H_P = math.isqrt(S) - W_P = H_P - assert H_P * W_P == S, ( - f"2D fusion requires a square patch grid; got S={S} which is not a " - f"perfect square.") - - # hilbert_indices/zigzag_indices give the forward perm scan_idx[h] = k - # (row-major index of the h-th scan token). Index arrays are constant - # at JIT time so computing both directions is free. - if self.scan_order == 'hilbert': - scan_fwd = hilbert_indices(H_P, W_P) # [S], scan→rowmajor - scan_inv = inverse_permutation(scan_fwd, S) # [S], rowmajor→scan - elif self.scan_order == 'zigzag': - scan_fwd = zigzag_indices(H_P, W_P) - scan_inv = inverse_permutation(scan_fwd, S) - else: # raster - scan_fwd = None - scan_inv = None - - if scan_fwd is not None: - # to row-major: rowmajor_tokens[k] = scan_tokens[scan_inv[k]] - ssm_rm = ssm_output[:, scan_inv, :] - else: - ssm_rm = ssm_output - - y_2d = ssm_rm.reshape(B, H_P, W_P, F) - y_fused_2d = self.spatial_fusion(y_2d) - y_fused_rm = y_fused_2d.reshape(B, S, F) - - if scan_fwd is not None: - # back to scan order: scan_tokens[h] = rowmajor_tokens[scan_fwd[h]] - y_fused = y_fused_rm[:, scan_fwd, :] - else: - y_fused = y_fused_rm - - return y_fused - - @nn.compact - 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 - ) - - # --- SSM Path (replaces Attention Path) --- - residual = x - norm_x = self.norm1(x) - x_modulated = norm_x * (1 + scale_attn) + shift_attn - ssm_output = self.ssm(x_modulated) - - 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 - else: - x = residual + ssm_output - - # --- MLP Path --- - residual = x - 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 - else: - x = residual + mlp_output - - return x - - -# --- Hybrid SSM-Attention DiT --- class HybridSSMAttentionDiT(nn.Module): """DiT that interleaves SSM blocks with attention blocks in a configurable ratio. @@ -408,6 +33,7 @@ class HybridSSMAttentionDiT(nn.Module): force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 learn_sigma: bool = False + qk_norm: bool = False use_hilbert: bool = False use_zigzag: bool = False # ZigMa-style serpentine scan block_pattern: Optional[Sequence[str]] = None # e.g., ['ssm','ssm','ssm','attn'] @@ -415,6 +41,12 @@ class HybridSSMAttentionDiT(nn.Module): bidirectional_ssm: bool = True use_2d_fusion: bool = False # 2D state fusion in SSM blocks (see SpatialFusionConv) + @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 _build_block_pattern(self): """Generate block pattern from ratio string.""" if self.block_pattern is not None: @@ -431,167 +63,77 @@ def _build_block_pattern(self): return pattern def setup(self): - self.patch_embed = PatchEmbedding( + self.embed = PatchSequenceEmbed( patch_size=self.patch_size, - embedding_dim=self.emb_features, + emb_features=self.emb_features, + scan_order=self.scan_order, dtype=self.dtype, - precision=self.precision + precision=self.precision, + ) + self.conditioning = ConditioningEmbed( + emb_features=self.emb_features, + mlp_ratio=self.mlp_ratio, + dtype=self.dtype, + precision=self.precision, ) - - assert not (self.use_hilbert and self.use_zigzag), \ - "use_hilbert and use_zigzag are mutually exclusive" - - if self.use_hilbert or self.use_zigzag: - self.hilbert_proj = nn.Dense( - features=self.emb_features, - dtype=self.dtype, - precision=self.precision, - name="hilbert_projection" - ) - - # Time embedding - self.time_embed = nn.Sequential([ - FourierEmbedding(features=self.emb_features), - TimeProjection(features=self.emb_features * self.mlp_ratio), - nn.Dense(features=self.emb_features, dtype=self.dtype, precision=self.precision) - ]) - - # Text context projection - self.text_proj = nn.Dense( - features=self.emb_features, dtype=self.dtype, - precision=self.precision, name="text_context_proj") - - # RoPE (used by attention blocks, passed through SSM blocks) self.rope = RotaryEmbedding( dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) - # Build hybrid block sequence pattern = self._build_block_pattern() blocks = [] for i, block_type in enumerate(pattern): if block_type == 'ssm': - if self.use_hilbert: - scan_order = 'hilbert' - elif self.use_zigzag: - scan_order = 'zigzag' - else: - scan_order = 'raster' - blocks.append(SSMDiTBlock( + blocks.append(ModulatedBlock( features=self.emb_features, num_heads=self.num_heads, rope_emb=self.rope, - state_dim=self.ssm_state_dim, + mixer='ssm', mlp_ratio=self.mlp_ratio, dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, norm_epsilon=self.norm_epsilon, - bidirectional=self.bidirectional_ssm, + ssm_state_dim=self.ssm_state_dim, + bidirectional_ssm=self.bidirectional_ssm, use_2d_fusion=self.use_2d_fusion, - scan_order=scan_order, + scan_order=self.scan_order, name=f"ssm_block_{i}" )) else: # 'attn' - blocks.append(DiTBlock( + blocks.append(ModulatedBlock( features=self.emb_features, num_heads=self.num_heads, rope_emb=self.rope, + 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, + force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, + qk_norm=self.qk_norm, name=f"dit_block_{i}" )) self.blocks = blocks - # Final layer - self.final_norm = nn.LayerNorm( - epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") - - output_dim = self.patch_size * self.patch_size * self.output_channels - if self.learn_sigma: - output_dim *= 2 - - self.final_proj = nn.Dense( - features=output_dim, - dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 + 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, - kernel_init=nn.initializers.zeros, - name="final_proj" ) @nn.compact 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 - - H_P = H // self.patch_size - W_P = W // self.patch_size - - # 1. Patch Embedding - hilbert_inv_idx = None - if self.use_hilbert: - patches_raw, hilbert_inv_idx = hilbert_patchify(x, self.patch_size) - patches = self.hilbert_proj(patches_raw) - elif self.use_zigzag: - patches_raw, hilbert_inv_idx = zigzag_patchify(x, self.patch_size) - patches = self.hilbert_proj(patches_raw) - else: - patches = self.patch_embed(x) - - num_patches = patches.shape[1] - - # 2D sincos position embedding - the SSM blocks ignore RoPE so they need - # an explicit positional signal. For hilbert/zigzag, reorder the row-major - # embedding to match the scan so each patch gets its real 2D position. - pos_embed_2d_rm = build_2d_sincos_pos_embed(self.emb_features, H_P, W_P) - pos_embed_2d_rm = jnp.asarray(pos_embed_2d_rm, dtype=patches.dtype) - if self.use_hilbert: - scan_idx = hilbert_indices(H_P, W_P) - pos_embed_2d = pos_embed_2d_rm[scan_idx] - elif self.use_zigzag: - scan_idx = zigzag_indices(H_P, W_P) - pos_embed_2d = pos_embed_2d_rm[scan_idx] - else: - pos_embed_2d = pos_embed_2d_rm - patches = patches + pos_embed_2d[None, :, :] + x_seq, inv_idx = self.embed(x) + cond_emb = self.conditioning(temb, textcontext) + freqs_cis = neutralized_rope_freqs(self.rope, x_seq.shape[1], self.scan_order) - x_seq = patches - - # 2. Conditioning - t_emb = self.time_embed(temb) - cond_emb = t_emb - if textcontext is not None: - text_emb = self.text_proj(textcontext) - text_emb_pooled = jnp.mean(text_emb, axis=1) - cond_emb = cond_emb + text_emb_pooled - - # 3. RoPE frequencies for the attention blocks. With hilbert/zigzag the - # sequence index is not a 2D position and RoPE distances would be wrong, - # so use identity freqs (cos=1, sin=0) - the 2D sincos above carries position. - freqs_cos, freqs_sin = self.rope(seq_len=num_patches) - if self.use_hilbert or self.use_zigzag: - freqs_cos = jnp.ones_like(freqs_cos) - freqs_sin = jnp.zeros_like(freqs_sin) - - # 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), train=train) - - # 5. Final output - x_out = self.final_norm(x_seq) - x_out = self.final_proj(x_out) + x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=freqs_cis, train=train) - # 6. Unpatchify - if self.use_hilbert or self.use_zigzag: - if self.learn_sigma: - x_mean, _x_logvar = jnp.split(x_out, 2, axis=-1) - return hilbert_unpatchify(x_mean, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - return hilbert_unpatchify(x_out, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - if self.learn_sigma: - x_mean, _x_logvar = jnp.split(x_out, 2, axis=-1) - return unpatchify(x_mean, channels=self.output_channels) - return unpatchify(x_out, channels=self.output_channels) + return self.output(x_seq, inv_idx, H, W) diff --git a/flaxdiff/models/vit_common.py b/flaxdiff/models/vit_common.py index 632f832..4fc9cfc 100644 --- a/flaxdiff/models/vit_common.py +++ b/flaxdiff/models/vit_common.py @@ -7,13 +7,15 @@ from .attention import NormalAttention -def unpatchify(x, channels=3): +def unpatchify(x, channels=3, H_P=None, W_P=None): patch_size = int((x.shape[2] // channels) ** 0.5) - h = w = int(x.shape[1] ** .5) - assert h * w == x.shape[1] and patch_size ** 2 * \ - channels == x.shape[2], f"Invalid shape: {x.shape}, should be {h*w}, {patch_size**2*channels}" + if H_P is None or W_P is None: + # No grid given - only valid for square grids + H_P = W_P = int(x.shape[1] ** .5) + assert H_P * W_P == x.shape[1] and patch_size ** 2 * \ + channels == x.shape[2], f"Invalid shape: {x.shape}, should be {H_P*W_P}, {patch_size**2*channels}" x = einops.rearrange( - x, 'B (h w) (p1 p2 C) -> B (h p1) (w p2) C', h=h, p1=patch_size, p2=patch_size) + x, 'B (h w) (p1 p2 C) -> B (h p1) (w p2) C', h=H_P, w=W_P, p1=patch_size, p2=patch_size) return x @@ -128,6 +130,9 @@ def __call__(self, x, context=None, freqs_cis=None): query = self.query(x) # [B, S, H, D] key = self.key(context) # [B, S_ctx, H, D] value = self.value(context) # [B, S_ctx, H, D] + if self.qk_norm: + query = self.q_norm(query) + key = self.k_norm(key) if freqs_cis is None and self.rope_emb is not None: seq_len_q = query.shape[1] # Use query's sequence length diff --git a/tests/test_models.py b/tests/test_models.py index 7747b49..bfdd96c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -61,7 +61,6 @@ def test_simple_dit_bf16_outputs_fp32(rng): assert out.dtype == jnp.float32 -@pytest.mark.xfail(strict=True, reason="bug: unpatchify assumes a square patch grid") def test_simple_dit_non_square(rng): model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) x = jax.random.normal(rng, (2, 16, 64, 3)) From c68b7dd5fa5dffc47d096fd166d60580b3065261 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 00:16:53 -0400 Subject: [PATCH 4/6] feat(models): real dual-stream MM-DiT SimpleMMDiT was MM-DiT in name only: text was mean-pooled into the adaLN vector and never entered the token sequence, attention and MLP ran in parallel off a single norm (verified by ablation - zeroing attention changed the MLP contribution by ~2e-6), so it was a text-modulated DiT at roughly half its nominal depth. The block is now a true SD3-style dual-stream MM-DiT: per-modality qkv/mlp/modulation weights, joint attention over concat([txt, img]) with RoPE on the image tokens only, sequential attn -> mlp, and an adaLN-modulated final layer. HierarchicalMMDiT is ported onto the same block (raster only - its merge/expand grid reshapes assume row-major order, so the old hilbert option was scrambling the neighborhoods being merged). AdaLNParams gains the missing SiLU before its zero-init projection, per the DiT paper. A behavior test pins the dual-stream property: changing text tokens must reach the image stream through joint attention. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/dit_common.py | 29 +- flaxdiff/models/simple_mmdit.py | 758 +++++++++++--------------------- flaxdiff/models/vit_common.py | 6 +- tests/test_models.py | 25 ++ 4 files changed, 301 insertions(+), 517 deletions(-) diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py index ed97f31..bf4ce78 100644 --- a/flaxdiff/models/dit_common.py +++ b/flaxdiff/models/dit_common.py @@ -127,26 +127,41 @@ class PatchSequenceOutput(nn.Module): patch_size: int output_channels: int learn_sigma: bool = False + modulated: bool = False # adaLN shift/scale on the final norm (DiT FinalLayer) norm_epsilon: float = 1e-5 dtype: Optional[Dtype] = None precision: PrecisionLike = None - def setup(self): - self.final_norm = nn.LayerNorm( - epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") + @nn.compact + def __call__(self, tokens, inv_idx, H, W, conditioning=None): + features = tokens.shape[-1] + x_out = nn.LayerNorm( + epsilon=self.norm_epsilon, use_scale=not self.modulated, + use_bias=not self.modulated, dtype=self.dtype, name="final_norm")(tokens) + if self.modulated: + assert conditioning is not None, "modulated output head needs the conditioning vector" + if conditioning.ndim == 2: + conditioning = jnp.expand_dims(conditioning, axis=1) + shift, scale = jnp.split(nn.Dense( + features=2 * features, + dtype=self.dtype, + precision=self.precision, + kernel_init=nn.initializers.zeros, + name="final_ada_proj", + )(nn.silu(conditioning)), 2, axis=-1) + x_out = x_out * (1 + scale) + shift + output_dim = self.patch_size * self.patch_size * self.output_channels if self.learn_sigma: output_dim *= 2 - self.final_proj = nn.Dense( + x_out = nn.Dense( features=output_dim, dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 precision=self.precision, kernel_init=nn.initializers.zeros, name="final_proj", - ) + )(x_out) - def __call__(self, tokens, inv_idx, H, W): - x_out = self.final_proj(self.final_norm(tokens)) if self.learn_sigma: x_out, _x_logvar = jnp.split(x_out, 2, axis=-1) if inv_idx is not None: diff --git a/flaxdiff/models/simple_mmdit.py b/flaxdiff/models/simple_mmdit.py index 904701e..0b25cb4 100644 --- a/flaxdiff/models/simple_mmdit.py +++ b/flaxdiff/models/simple_mmdit.py @@ -1,174 +1,139 @@ +""" +MM-DiT (SD3-style multi-modal DiT) and a hierarchical variant. + +The block is a true dual-stream MM-DiT: text and image tokens keep separate +qkv/mlp/modulation weights and mix through a single joint attention over the +concatenated sequence. The previous implementation mean-pooled the text into +the adaLN vector (text never entered the token sequence) and ran attention +and MLP in parallel off one norm, which roughly halved the effective depth - +it was a text-modulated DiT, not an MM-DiT. +""" + import jax import jax.numpy as jnp from flax import linen as nn -from typing import Callable, Any, Optional, Tuple, Sequence, Union, List +from typing import Optional, Sequence import einops -from functools import partial from flax.typing import Dtype, PrecisionLike -# Imports from local modules -from .vit_common import PatchEmbedding, unpatchify, RotaryEmbedding, RoPEAttention -from .common import kernel_init, FourierEmbedding, TimeProjection -from .attention import NormalAttention # Base for RoPEAttention -# Replace common.hilbert_indices with improved implementation from hilbert.py -from .hilbert import hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify - -# --- MM-DiT AdaLN-Zero --- -class MMAdaLNZero(nn.Module): - """ - Adaptive Layer Normalization Zero (AdaLN-Zero) tailored for MM-DiT. - Projects time and text embeddings separately, combines them, and then - generates modulation parameters (scale, shift, gate) for attention and MLP paths. - """ - features: int - dtype: Optional[Dtype] = None - precision: PrecisionLike = None - norm_epsilon: float = 1e-5 - use_mean_pooling: bool = True # Whether to use mean pooling for sequence inputs - - @nn.compact - def __call__(self, x, t_emb, text_emb): - # x shape: [B, S, F] - # t_emb shape: [B, D_t] - # text_emb shape: [B, S_text, D_text] or [B, D_text] - - # First normalize the input features - norm = nn.LayerNorm(epsilon=self.norm_epsilon, - use_scale=False, use_bias=False, dtype=self.dtype) - norm_x = norm(x) # Shape: [B, S, F] - - # Process time embedding: ensure it has a sequence dimension for later broadcasting - if t_emb.ndim == 2: # [B, D_t] - t_emb = jnp.expand_dims(t_emb, axis=1) # [B, 1, D_t] - - # Process text embedding: if it has a sequence dimension different from x - if text_emb.ndim == 2: # [B, D_text] - text_emb = jnp.expand_dims(text_emb, axis=1) # [B, 1, D_text] - elif text_emb.ndim == 3 and self.use_mean_pooling and text_emb.shape[1] != x.shape[1]: - # Mean pooling is standard in MM-DiT for handling different sequence lengths - text_emb = jnp.mean( - text_emb, axis=1, keepdims=True) # [B, 1, D_text] - - # Project time embedding - t_params = nn.Dense( - features=6 * self.features, - dtype=self.dtype, - precision=self.precision, - kernel_init=nn.initializers.zeros, # Zero init is standard in AdaLN-Zero - name="ada_t_proj" - )(t_emb) # Shape: [B, 1, 6*F] - - # Project text embedding - text_params = nn.Dense( - features=6 * self.features, - dtype=self.dtype, - precision=self.precision, - kernel_init=nn.initializers.zeros, # Zero init - name="ada_text_proj" - )(text_emb) # Shape: [B, 1, 6*F] or [B, S_text, 6*F] - - # If text_params still has a sequence dim different from t_params, mean pool it - if t_params.shape[1] != text_params.shape[1]: - text_params = jnp.mean(text_params, axis=1, keepdims=True) - - # Combine parameters (summing is standard in MM-DiT) - ada_params = t_params + text_params # Shape: [B, 1, 6*F] - - # Split into scale, shift, gate for MLP and Attention - scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split( - ada_params, 6, axis=-1) # Each shape: [B, 1, F] - - scale_mlp = jnp.clip(scale_mlp, -10.0, 10.0) - shift_mlp = jnp.clip(shift_mlp, -10.0, 10.0) - # Apply modulation for Attention path (broadcasting handled by JAX) - x_attn = norm_x * (1 + scale_attn) + shift_attn +from .dit_common import ( + PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, + neutralized_rope_freqs, +) +from .vit_common import RotaryEmbedding, AdaLNParams, apply_rotary_embedding - # Apply modulation for MLP path - x_mlp = norm_x * (1 + scale_mlp) + shift_mlp - # Return modulated outputs and gates - return x_attn, gate_attn, x_mlp, gate_mlp - - -# --- MM-DiT Block --- class MMDiTBlock(nn.Module): - """ - A Transformer block adapted for MM-DiT, using MMAdaLNZero for conditioning. + """Dual-stream MM-DiT block: per-modality weights, joint attention. + + Both streams are modulated adaLN-Zero style from the same conditioning + vector but with separate parameters, then attend jointly over + concat([txt, img]) and go through separate sequential MLPs. """ features: int num_heads: int - rope_emb: RotaryEmbedding # Pass RoPE module mlp_ratio: int = 4 dropout_rate: float = 0.0 dtype: Optional[Dtype] = None precision: PrecisionLike = None - # Keep option, though RoPEAttention doesn't use it force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 + qk_norm: bool = False def setup(self): hidden_features = int(self.features * self.mlp_ratio) - # Use the new MMAdaLNZero block - self.ada_ln_zero = MMAdaLNZero( - self.features, dtype=self.dtype, precision=self.precision, norm_epsilon=self.norm_epsilon) - - # RoPEAttention remains the same - self.attention = RoPEAttention( - query_dim=self.features, - heads=self.num_heads, - dim_head=self.features // self.num_heads, - dtype=self.dtype, - precision=self.precision, - use_bias=True, # Bias is common in DiT attention proj - force_fp32_for_softmax=self.force_fp32_for_softmax, - rope_emb=self.rope_emb # Pass RoPE module instance - ) - - # Standard MLP block remains the same - self.mlp = nn.Sequential([ - nn.Dense(features=hidden_features, dtype=self.dtype, - precision=self.precision), - nn.gelu, # Consider swish/silu if preferred - nn.Dense(features=self.features, dtype=self.dtype, - precision=self.precision) - ]) + dim_head = self.features // self.num_heads + qkv = lambda name: nn.DenseGeneral( + features=[self.num_heads, dim_head], axis=-1, + dtype=self.dtype, precision=self.precision, use_bias=True, name=name) + out_proj = lambda name: nn.DenseGeneral( + self.features, axis=(-2, -1), + dtype=self.dtype, precision=self.precision, name=name) + mlp = lambda name: nn.Sequential([ + nn.Dense(features=hidden_features, dtype=self.dtype, precision=self.precision), + nn.gelu, + nn.Dense(features=self.features, dtype=self.dtype, precision=self.precision), + ], name=name) + norm = lambda name: nn.LayerNorm( + epsilon=self.norm_epsilon, use_scale=False, use_bias=False, + dtype=self.dtype, name=name) + + # image stream + self.img_ada = AdaLNParams(self.features, dtype=self.dtype, precision=self.precision) + self.img_norm1, self.img_norm2 = norm("img_norm1"), norm("img_norm2") + self.img_q, self.img_k, self.img_v = qkv("img_to_q"), qkv("img_to_k"), qkv("img_to_v") + self.img_out = out_proj("img_out") + self.img_mlp = mlp("img_mlp") + + # text stream + self.txt_ada = AdaLNParams(self.features, dtype=self.dtype, precision=self.precision) + self.txt_norm1, self.txt_norm2 = norm("txt_norm1"), norm("txt_norm2") + self.txt_q, self.txt_k, self.txt_v = qkv("txt_to_q"), qkv("txt_to_k"), qkv("txt_to_v") + self.txt_out = out_proj("txt_out") + self.txt_mlp = mlp("txt_mlp") + + if self.qk_norm: + self.img_q_norm = nn.RMSNorm(dtype=self.dtype, name="img_q_norm") + self.img_k_norm = nn.RMSNorm(dtype=self.dtype, name="img_k_norm") + self.txt_q_norm = nn.RMSNorm(dtype=self.dtype, name="txt_q_norm") + self.txt_k_norm = nn.RMSNorm(dtype=self.dtype, name="txt_k_norm") self.dropout = nn.Dropout(rate=self.dropout_rate) @nn.compact - 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] - - residual = x - - # Apply MMAdaLNZero with separate time and text embeddings - x_attn, gate_attn, x_mlp, gate_mlp = self.ada_ln_zero( - x, t_emb, text_emb) - - # 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 + def __call__(self, img, txt, conditioning, freqs_cis, train: bool = False): + S_txt = txt.shape[1] + i_scale_mlp, i_shift_mlp, i_gate_mlp, i_scale_attn, i_shift_attn, i_gate_attn = jnp.split( + self.img_ada(conditioning), 6, axis=-1) + t_scale_mlp, t_shift_mlp, t_gate_mlp, t_scale_attn, t_shift_attn, t_gate_attn = jnp.split( + self.txt_ada(conditioning), 6, axis=-1) + + # --- Joint attention --- + img_h = self.img_norm1(img) * (1 + i_scale_attn) + i_shift_attn + txt_h = self.txt_norm1(txt) * (1 + t_scale_attn) + t_shift_attn + + q_i, k_i, v_i = self.img_q(img_h), self.img_k(img_h), self.img_v(img_h) + q_t, k_t, v_t = self.txt_q(txt_h), self.txt_k(txt_h), self.txt_v(txt_h) + if self.qk_norm: + q_i, k_i = self.img_q_norm(q_i), self.img_k_norm(k_i) + q_t, k_t = self.txt_q_norm(q_t), self.txt_k_norm(k_t) + + # RoPE carries 2D position for image tokens only + if freqs_cis is not None: + freqs_cos, freqs_sin = freqs_cis + q_i = einops.rearrange(q_i, 'b s h d -> b h s d') + k_i = einops.rearrange(k_i, 'b s h d -> b h s d') + q_i = apply_rotary_embedding(q_i, freqs_cos, freqs_sin) + k_i = apply_rotary_embedding(k_i, freqs_cos, freqs_sin) + q_i = einops.rearrange(q_i, 'b h s d -> b s h d') + k_i = einops.rearrange(k_i, 'b h s d -> b s h d') + + q = jnp.concatenate([q_t, q_i], axis=1) + 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) + txt_attn, img_attn = attn[:, :S_txt], attn[:, S_txt:] + + img_attn = self.dropout(self.img_out(img_attn), deterministic=not train) + txt_attn = self.dropout(self.txt_out(txt_attn), deterministic=not train) + img = img + i_gate_attn * img_attn + txt = txt + t_gate_attn * txt_attn + + # --- Sequential MLPs --- + img_h = self.img_norm2(img) * (1 + i_scale_mlp) + i_shift_mlp + img = img + i_gate_mlp * self.dropout(self.img_mlp(img_h), deterministic=not train) + txt_h = self.txt_norm2(txt) * (1 + t_scale_mlp) + t_shift_mlp + txt = txt + t_gate_mlp * self.dropout(self.txt_mlp(txt_h), deterministic=not train) + + return img, txt -# --- SimpleMMDiT --- class SimpleMMDiT(nn.Module): - """ - A Simple Multi-Modal Diffusion Transformer (MM-DiT) implementation. - Integrates time and text conditioning using separate projections within - each transformer block, following the MM-DiT approach. Uses RoPE for - patch positional encoding. - """ + """SD3-style MM-DiT: a plain stack of dual-stream blocks.""" output_channels: int = 3 patch_size: int = 16 emb_features: int = 768 @@ -178,48 +143,39 @@ class SimpleMMDiT(nn.Module): dropout_rate: float = 0.0 # Typically 0 for diffusion dtype: Optional[Dtype] = None precision: PrecisionLike = None - # Passed down, but RoPEAttention uses NormalAttention 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 + learn_sigma: bool = False + qk_norm: bool = False + 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.patch_embed = PatchEmbedding( + self.embed = PatchSequenceEmbed( patch_size=self.patch_size, - embedding_dim=self.emb_features, + 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 + precision=self.precision, ) - - # Time embedding projection (output dim: emb_features) - self.time_embed = nn.Sequential([ - FourierEmbedding(features=self.emb_features), - TimeProjection(features=self.emb_features * - self.mlp_ratio), # Intermediate projection - nn.Dense(features=self.emb_features, dtype=self.dtype, - precision=self.precision) # Final projection - ], name="time_embed") - - # Add projection layer for Hilbert patches - if self.use_hilbert: - self.hilbert_proj = nn.Dense( - features=self.emb_features, - dtype=self.dtype, - precision=self.precision, - name="hilbert_projection" - ) - # Text context projection (output dim: emb_features) - # Input dim depends on the text encoder output, assumed to be handled externally - self.text_proj = nn.Dense(features=self.emb_features, dtype=self.dtype, - precision=self.precision, name="text_context_proj") - - # Rotary Positional Embedding (for patches) - # Dim per head, max_len should cover max number of patches + # text tokens enter the sequence, so they need their own projection + self.txt_embed = nn.Dense( + features=self.emb_features, dtype=self.dtype, + precision=self.precision, name="txt_embed") self.rope = RotaryEmbedding( dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) - - # Transformer Blocks (use MMDiTBlock) self.blocks = [ MMDiTBlock( features=self.emb_features, @@ -230,107 +186,35 @@ def setup(self): precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax, norm_epsilon=self.norm_epsilon, - rope_emb=self.rope, # Pass RoPE instance + qk_norm=self.qk_norm, name=f"mmdit_block_{i}" ) for i in range(self.num_layers) ] - - # Final Layer (Normalization + Linear Projection) - self.final_norm = nn.LayerNorm( - epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") - # self.final_norm = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") # Alternative - - # Predict patch pixels + potentially sigma - output_dim = self.patch_size * self.patch_size * self.output_channels - if self.learn_sigma: - output_dim *= 2 # Predict both mean and variance (or log_variance) - - self.final_proj = nn.Dense( - features=output_dim, - dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 + self.output = PatchSequenceOutput( + patch_size=self.patch_size, + output_channels=self.output_channels, + learn_sigma=self.learn_sigma, + modulated=True, + norm_epsilon=self.norm_epsilon, + dtype=self.dtype, precision=self.precision, - kernel_init=nn.initializers.zeros, # Initialize final layer to zero - name="final_proj" ) @nn.compact 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" + B, H, W, C = x.shape + + img, inv_idx = self.embed(x) + txt = self.txt_embed(textcontext) + cond_emb = self.conditioning(temb, textcontext) + freqs_cis = neutralized_rope_freqs(self.rope, img.shape[1], self.scan_order) - # 1. Patch Embedding - if self.use_hilbert: - # Use hilbert_patchify which handles both patchification and reordering - patches_raw, hilbert_inv_idx = hilbert_patchify( - x, self.patch_size) # Shape [B, S, P*P*C] - # Apply projection - # Shape [B, S, emb_features] - patches = self.hilbert_proj(patches_raw) - else: - # Shape: [B, num_patches, emb_features] - patches = self.patch_embed(x) - hilbert_inv_idx = None - - num_patches = patches.shape[1] - x_seq = patches - - # 2. Prepare Conditioning Signals - t_emb = self.time_embed(temb) # Shape: [B, emb_features] - # Assuming textcontext is [B, context_seq_len, context_dim] or [B, context_dim] - # If [B, context_seq_len, context_dim], usually mean/pool or take CLS token first. - # Assuming textcontext is already pooled/CLS token: [B, context_dim] - text_emb = self.text_proj(textcontext) # Shape: [B, emb_features] - - # 3. Apply RoPE Frequencies (only to patch tokens) - seq_len = x_seq.shape[1] - freqs_cos, freqs_sin = self.rope(seq_len) # Shapes: [S, D_head/2] - - # 4. Apply Transformer Blocks 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), train=train) - - # 5. Final Layer - x_seq = self.final_norm(x_seq) - # Shape: [B, num_patches, P*P*C (*2 if learn_sigma)] - x_seq = self.final_proj(x_seq) - - # 6. Unpatchify - if self.use_hilbert: - # For Hilbert mode, we need to use the specialized unpatchify function - if self.learn_sigma: - # Split into mean and variance predictions - x_mean, x_logvar = jnp.split(x_seq, 2, axis=-1) - x_image = hilbert_unpatchify( - x_mean, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - # If needed, also unpack the logvar - # logvar_image = hilbert_unpatchify(x_logvar, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - # return x_image, logvar_image - return x_image - else: - x_image = hilbert_unpatchify( - x_seq, hilbert_inv_idx, self.patch_size, H, W, self.output_channels) - return x_image - else: - # Standard patch ordering - use the existing unpatchify function - if self.learn_sigma: - # Split into mean and variance predictions - x_mean, x_logvar = jnp.split(x_seq, 2, axis=-1) - x = unpatchify(x_mean, channels=self.output_channels) - # Return both mean and logvar if needed by the loss function - # For now, just returning the mean prediction like standard diffusion models - # logvar = unpatchify(x_logvar, channels=self.output_channels) - # return x, logvar - return x - else: - # Shape: [B, H, W, C] - x = unpatchify(x_seq, channels=self.output_channels) - return x - - -# --- Hierarchical MM-DiT components --- + img, txt = block(img, txt, conditioning=cond_emb, freqs_cis=freqs_cis, train=train) + + return self.output(img, inv_idx, H, W, conditioning=cond_emb) + class PatchMerging(nn.Module): """ @@ -428,19 +312,18 @@ def __call__(self, x, H_patches, W_patches): return expanded, new_H, new_W -# --- Hierarchical MM-DiT --- class HierarchicalMMDiT(nn.Module): - """ - A Hierarchical Multi-Modal Diffusion Transformer (MM-DiT) implementation - based on the PixArt-α architecture. Processes images at multiple resolutions - with skip connections between encoder and decoder paths. - Follows a U-Net like structure: Fine -> Coarse (Encoder) -> Coarse -> Fine (Decoder). + """U-shaped MM-DiT: dual-stream blocks per stage with patch merging on the + way down and expansion + skip fusion on the way up. + + Raster order only - the merge/expand grid reshapes assume row-major token + order, so a hilbert scan would scramble the neighborhoods being merged. """ output_channels: int = 3 base_patch_size: int = 8 # Patch size at the *finest* resolution level (stage 0) - emb_features: Sequence[int] = (512, 768, 1024) # Feature dimensions for stages 0, 1, 2 (fine to coarse) - num_layers: Sequence[int] = (4, 4, 14) # Layers per stage (can be asymmetric encoder/decoder if needed) - num_heads: Sequence[int] = (8, 12, 16) # Heads per stage (fine to coarse) + emb_features: Sequence[int] = (512, 768, 1024) # Feature dims for stages, fine to coarse + num_layers: Sequence[int] = (4, 4, 14) # Layers per stage, fine to coarse + num_heads: Sequence[int] = (8, 12, 16) # Heads per stage, fine to coarse mlp_ratio: int = 4 dropout_rate: float = 0.0 dtype: Optional[Dtype] = None @@ -448,72 +331,47 @@ class HierarchicalMMDiT(nn.Module): force_fp32_for_softmax: bool = True norm_epsilon: float = 1e-5 learn_sigma: bool = False - use_hilbert: bool = False + qk_norm: bool = False def setup(self): assert len(self.emb_features) == len(self.num_layers) == len(self.num_heads), \ "Feature dimensions, layers, and heads must have the same number of stages" - num_stages = len(self.emb_features) - # 1. Initial Patch Embedding (FINEST level - stage 0) - self.patch_embed = PatchEmbedding( + self.embed = PatchSequenceEmbed( patch_size=self.base_patch_size, - embedding_dim=self.emb_features[0], # Finest embedding dim + emb_features=self.emb_features[0], + scan_order='raster', dtype=self.dtype, - precision=self.precision + precision=self.precision, ) - - # 2. Time/Text Embeddings (Projected for each stage) - # Base projection to largest dimension first for stability/capacity - base_t_emb_dim = self.emb_features[-1] - self.time_embed_base = nn.Sequential([ - FourierEmbedding(features=base_t_emb_dim), - TimeProjection(features=base_t_emb_dim * self.mlp_ratio), - nn.Dense(features=base_t_emb_dim, dtype=self.dtype, precision=self.precision) - ], name="time_embed_base") - self.text_proj_base = nn.Dense( - features=base_t_emb_dim, + # Base conditioning at the finest dim, projected per stage + self.conditioning = ConditioningEmbed( + emb_features=self.emb_features[0], + mlp_ratio=self.mlp_ratio, dtype=self.dtype, precision=self.precision, - name="text_context_proj_base" ) - # Projections for each stage (0 to N-1) - self.t_emb_projs = [ - nn.Dense(features=self.emb_features[i], dtype=self.dtype, precision=self.precision, name=f"t_emb_proj_stage{i}") + self.cond_projs = [ + nn.Dense(features=self.emb_features[i], dtype=self.dtype, + precision=self.precision, name=f"cond_proj_stage{i}") for i in range(num_stages) ] - self.text_emb_projs = [ - nn.Dense(features=self.emb_features[i], dtype=self.dtype, precision=self.precision, name=f"text_emb_proj_stage{i}") + # Per-stage text streams (dims differ per stage) + self.txt_embeds = [ + nn.Dense(features=self.emb_features[i], dtype=self.dtype, + precision=self.precision, name=f"txt_embed_stage{i}") for i in range(num_stages) ] - - # 3. Hilbert projection (if used, applied after initial patch embedding) - if self.use_hilbert: - self.hilbert_proj = nn.Dense( - features=self.emb_features[0], # Match finest embedding dim - dtype=self.dtype, - precision=self.precision, - name="hilbert_projection" - ) - - # 4. RoPE embeddings for each stage (0 to N-1) self.ropes = [ RotaryEmbedding( dim=self.emb_features[i] // self.num_heads[i], - max_seq_len=4096, # Adjust if needed based on max patch count per stage - dtype=self.dtype, - name=f"rope_stage_{i}" - ) + max_seq_len=4096, dtype=self.dtype, name=f"rope_stage_{i}") for i in range(num_stages) ] - # 5. --- Encoder Path (Fine to Coarse) --- - encoder_blocks = [] - patch_mergers = [] - for stage in range(num_stages): - # Blocks for this stage - stage_blocks = [ + def stage_blocks(stage, prefix): + return [ MMDiTBlock( features=self.emb_features[stage], num_heads=self.num_heads[stage], @@ -521,204 +379,88 @@ def setup(self): dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision, - 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}" - ) - # Assuming symmetric layers for now, adjust if needed (e.g., self.num_encoder_layers[stage]) - for i in range(self.num_layers[stage]) + qk_norm=self.qk_norm, + name=f"{prefix}_block_stage{stage}_{i}" + ) for i in range(self.num_layers[stage]) ] - encoder_blocks.append(stage_blocks) - # Patch Merging layer (except for the last/coarsest stage) - if stage < num_stages - 1: - patch_mergers.append( - PatchMerging( - out_features=self.emb_features[stage + 1], # Target next stage dim - dtype=self.dtype, - precision=self.precision, - norm_epsilon=self.norm_epsilon, - name=f"patch_merger_{stage}" - ) - ) - self.encoder_blocks = encoder_blocks - self.patch_mergers = patch_mergers - - # 6. --- Decoder Path (Coarse to Fine) --- - decoder_blocks = [] - patch_expanders = [] - fusion_layers = [] - # Iterate from second coarsest stage (N-2) down to finest (0) - for stage in range(num_stages - 2, -1, -1): - # Patch Expanding layer (Expands from stage+1 to stage) - patch_expanders.append( - PatchExpanding( - out_features=self.emb_features[stage], # Target current stage dim - dtype=self.dtype, - precision=self.precision, - norm_epsilon=self.norm_epsilon, - name=f"patch_expander_{stage}" # Naming indicates target stage - ) - ) - # Fusion layer (Combines skip[stage] and expanded[stage+1]->[stage]) - fusion_layers.append( - nn.Sequential([ # Use Sequential for Norm + Dense - nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name=f"fusion_norm_{stage}"), - nn.Dense( - features=self.emb_features[stage], # Output current stage dim - dtype=self.dtype, - precision=self.precision, - name=f"fusion_dense_{stage}" - ) - ]) - ) - - # Blocks for this stage (stage N-2 down to 0) - # Assuming symmetric layers for now - stage_blocks = [ - MMDiTBlock( - features=self.emb_features[stage], - num_heads=self.num_heads[stage], - 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, - rope_emb=self.ropes[stage], - name=f"decoder_block_stage{stage}_{i}" - ) - for i in range(self.num_layers[stage]) - ] - # Append blocks in order: stage N-2, N-3, ..., 0 - decoder_blocks.append(stage_blocks) - - self.patch_expanders = patch_expanders - self.fusion_layers = fusion_layers - self.decoder_blocks = decoder_blocks - - # Note: The lists expanders, fusion_layers, decoder_blocks are now ordered - # corresponding to stages N-2, N-3, ..., 0. - - # 7. Final Layer - self.final_norm = nn.LayerNorm( - epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm") - - # Output projection to pixels (at finest resolution) - output_dim = self.base_patch_size * self.base_patch_size * self.output_channels - if self.learn_sigma: - output_dim *= 2 # Predict both mean and variance - - self.final_proj = nn.Dense( - features=output_dim, - dtype=jnp.float32, # fp32 output head - the loss is computed in fp32 + # --- Encoder path (fine to coarse) --- + self.encoder_blocks = [stage_blocks(s, "encoder") for s in range(num_stages)] + self.patch_mergers = [ + PatchMerging( + out_features=self.emb_features[s + 1], + dtype=self.dtype, + precision=self.precision, + norm_epsilon=self.norm_epsilon, + name=f"patch_merger_{s}" + ) for s in range(num_stages - 1) + ] + + # --- Decoder path (coarse to fine), ordered for stages N-2, ..., 0 --- + decoder_stages = list(range(num_stages - 2, -1, -1)) + self.patch_expanders = [ + PatchExpanding( + out_features=self.emb_features[s], + dtype=self.dtype, + precision=self.precision, + norm_epsilon=self.norm_epsilon, + name=f"patch_expander_{s}" + ) for s in decoder_stages + ] + self.fusion_layers = [ + nn.Sequential([ + nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype), + nn.Dense(features=self.emb_features[s], dtype=self.dtype, + precision=self.precision), + ], name=f"fusion_{s}") for s in decoder_stages + ] + self.decoder_blocks = [stage_blocks(s, "decoder") for s in decoder_stages] + + self.output = PatchSequenceOutput( + patch_size=self.base_patch_size, + output_channels=self.output_channels, + learn_sigma=self.learn_sigma, + modulated=True, + norm_epsilon=self.norm_epsilon, + dtype=self.dtype, precision=self.precision, - kernel_init=nn.initializers.zeros, # Zero init - name="final_proj" ) + @nn.compact def __call__(self, x, temb, textcontext, train: bool = False): + assert textcontext is not None, "textcontext must be provided" B, H, W, C = x.shape num_stages = len(self.emb_features) - finest_patch_size = self.base_patch_size - - # Assertions - assert H % (finest_patch_size * (2**(num_stages - 1))) == 0 and \ - W % (finest_patch_size * (2**(num_stages - 1))) == 0, \ - f"Image dimensions ({H},{W}) must be divisible by effective coarsest patch size {finest_patch_size * (2**(num_stages - 1))}" - assert textcontext is not None, "textcontext must be provided" - - # 1. Initial Patch Embedding (Finest Level - stage 0) - H_patches = H // finest_patch_size - W_patches = W // finest_patch_size - total_patches = H_patches * W_patches # Calculate total patches - hilbert_inv_idx = None - if self.use_hilbert: - # Calculate Hilbert indices and inverse permutation for the *finest* grid - fine_idx = hilbert_indices(H_patches, W_patches) - # Pass the total number of patches as total_size - hilbert_inv_idx = inverse_permutation(fine_idx, total_size=total_patches) # Store for unpatchify - - # Apply Hilbert patchify at the finest level - patches_raw, _ = hilbert_patchify(x, finest_patch_size) # We already have inv_idx - x_seq = self.hilbert_proj(patches_raw) # Shape [B, S_fine, emb[0]] - else: - x_seq = self.patch_embed(x) # Shape [B, S_fine, emb[0]] - - # 2. Prepare Conditioning Signals for each stage - t_emb_base = self.time_embed_base(temb) - text_emb_base = self.text_proj_base(textcontext) - t_embs = [proj(t_emb_base) for proj in self.t_emb_projs] # List for stages 0 to N-1 - text_embs = [proj(text_emb_base) for proj in self.text_emb_projs] # List for stages 0 to N-1 - - # --- Encoder Path (Fine to Coarse: stages 0 to N-1) --- - skip_features = {} - current_H_patches, current_W_patches = H_patches, W_patches + assert H % (self.base_patch_size * (2**(num_stages - 1))) == 0 and \ + W % (self.base_patch_size * (2**(num_stages - 1))) == 0, \ + f"Image dimensions ({H},{W}) must be divisible by effective coarsest patch size {self.base_patch_size * (2**(num_stages - 1))}" + + img, _ = self.embed(x) + cond_base = self.conditioning(temb, textcontext) + conds = [proj(cond_base) for proj in self.cond_projs] + txts = [embed(textcontext) for embed in self.txt_embeds] + + # --- Encoder path --- + H_P, W_P = H // self.base_patch_size, W // self.base_patch_size + skips = {} for stage in range(num_stages): - # Apply RoPE for current stage - seq_len = x_seq.shape[1] - freqs_cos, freqs_sin = self.ropes[stage](seq_len) - - # Apply blocks for this stage + freqs_cis = self.ropes[stage](seq_len=img.shape[1]) + txt = txts[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), train=train) - - # Store skip features (before merging) - skip_features[stage] = x_seq - - # Apply Patch Merging (if not the last/coarsest stage) + img, txt = block(img, txt, conditioning=conds[stage], freqs_cis=freqs_cis, train=train) + skips[stage] = img if stage < num_stages - 1: - x_seq, current_H_patches, current_W_patches = self.patch_mergers[stage]( - x_seq, current_H_patches, current_W_patches - ) - - # --- Bottleneck --- - # x_seq now holds the output of the coarsest stage (stage N-1) - - # --- Decoder Path (Coarse to Fine: stages N-2 down to 0) --- - # Decoder lists (expanders, fusion, blocks) are ordered for stages N-2, ..., 0 - for i, stage in enumerate(range(num_stages - 2, -1, -1)): # stage = N-2, N-3, ..., 0; i = 0, 1, ..., N-2 - # Apply Patch Expanding (Expand from stage+1 feature map to stage feature map) - x_seq, current_H_patches, current_W_patches = self.patch_expanders[i]( - x_seq, current_H_patches, current_W_patches - ) - - # Fusion with skip connection from corresponding encoder stage - skip = skip_features[stage] - x_seq = jnp.concatenate([x_seq, skip], axis=-1) # Concatenate along feature dim - x_seq = self.fusion_layers[i](x_seq) # Apply fusion (Norm + Dense) - - # Apply RoPE for current stage - seq_len = x_seq.shape[1] - freqs_cos, freqs_sin = self.ropes[stage](seq_len) - - # 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), train=train) - - # --- Final Layer --- - # x_seq should now be at the finest resolution (stage 0 features) - x_seq = self.final_norm(x_seq) - x_seq = self.final_proj(x_seq) # Project to patch pixel values - - # --- Unpatchify --- - if self.use_hilbert: - # Use the inverse Hilbert index calculated for the *finest* grid - assert hilbert_inv_idx is not None, "Hilbert inverse index should exist if use_hilbert is True" - if self.learn_sigma: - x_mean, x_logvar = jnp.split(x_seq, 2, axis=-1) - out = hilbert_unpatchify(x_mean, hilbert_inv_idx, finest_patch_size, H, W, self.output_channels) - # Optionally return logvar: logvar_image = hilbert_unpatchify(x_logvar, hilbert_inv_idx, finest_patch_size, H, W, self.output_channels) - else: - out = hilbert_unpatchify(x_seq, hilbert_inv_idx, finest_patch_size, H, W, self.output_channels) - else: - # Standard unpatchify - if self.learn_sigma: - x_mean, x_logvar = jnp.split(x_seq, 2, axis=-1) - out = unpatchify(x_mean, channels=self.output_channels) - # Optionally return logvar: logvar = unpatchify(x_logvar, channels=self.output_channels) - else: - out = unpatchify(x_seq, channels=self.output_channels) - - return out + img, H_P, W_P = self.patch_mergers[stage](img, H_P, W_P) + + # --- Decoder path --- + for i, stage in enumerate(range(num_stages - 2, -1, -1)): + img, H_P, W_P = self.patch_expanders[i](img, H_P, W_P) + img = self.fusion_layers[i](jnp.concatenate([img, skips[stage]], axis=-1)) + freqs_cis = self.ropes[stage](seq_len=img.shape[1]) + txt = txts[stage] + for block in self.decoder_blocks[i]: + img, txt = block(img, txt, conditioning=conds[stage], freqs_cis=freqs_cis, train=train) + + return self.output(img, None, H, W, conditioning=conds[0]) diff --git a/flaxdiff/models/vit_common.py b/flaxdiff/models/vit_common.py index 4fc9cfc..66eab3e 100644 --- a/flaxdiff/models/vit_common.py +++ b/flaxdiff/models/vit_common.py @@ -186,14 +186,16 @@ def __call__(self, conditioning): if conditioning.ndim == 2: conditioning = jnp.expand_dims(conditioning, axis=1) - # Project conditioning to get 6 params per feature + # SiLU then a zero-init projection, as in the DiT paper - without the + # nonlinearity every block's 6 modulation vectors are just an affine + # map of the same shared vector ada_params = nn.Dense( features=6 * self.features, dtype=self.dtype, precision=self.precision, kernel_init=nn.initializers.zeros, name="ada_proj" - )(conditioning) + )(nn.silu(conditioning)) # Return all params (or split if preferred, but maybe return tuple/dict) # Shape: [B, 1, 6*F] return ada_params # Or split and return tuple: jnp.split(ada_params, 6, axis=-1) diff --git a/tests/test_models.py b/tests/test_models.py index bfdd96c..4df3335 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -136,3 +136,28 @@ def test_dropout_is_active_in_train_mode(rng): e0 = model.apply(params, x, temb, textcontext) e1 = model.apply(params, x, temb, textcontext) assert jnp.allclose(e0, e1), "eval mode must be deterministic" + + +def test_hierarchical_mmdit_forward(rng): + from flaxdiff.models.simple_mmdit import HierarchicalMMDiT + model = HierarchicalMMDiT(base_patch_size=2, emb_features=(32, 64, 96), + num_layers=(1, 1, 1), num_heads=(2, 2, 2), mlp_ratio=2) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + + +def test_mmdit_is_dual_stream(rng): + """Text must participate in the token sequence: zeroing the text context + must change the image output through joint attention, not just through + the pooled conditioning vector.""" + model = SimpleMMDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + x, temb, textcontext = small_inputs(rng) + params = model.init(rng, x, temb, textcontext) + params = jax.tree.map(lambda p: p + 0.02, params) + + text_a = jnp.ones_like(textcontext) + text_b = jnp.concatenate([jnp.ones_like(textcontext[:, :38]) * 2.0, text_a[:, 38:]], axis=1) + 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" From bc0104ef820203706381ce5121adc8873395f2b0 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 00:22:04 -0400 Subject: [PATCH 5/6] refactor(trainer): collapse the trainer hierarchy to one concrete trainer GeneralDiffusionTrainer overrode every method the legacy DiffusionTrainer middle class added except TrainState/generate_states/fit - the middle class was a fork whose train/validation code was dead (its validation step could not even construct a sampler post the API change) yet still shipped. Those three live pieces now sit directly on GeneralDiffusionTrainer, the middle class and the broken AutoEncoderTrainer (a copy-paste of the diffusion step referencing attributes that were never set) are deleted, and SimpleTrainer keeps only generic infra with abstract step definitions instead of toy classification steps no real path used. The trainer smoke tests now run the real GeneralDiffusionTrainer end to end (tiny DiT, unconditional synthetic data): fit without wandb, and a save/restore round-trip asserting step counter, adam moments and EMA params all survive. Co-Authored-By: Claude Fable 5 --- flaxdiff/trainer/__init__.py | 3 +- flaxdiff/trainer/autoencoder_trainer.py | 181 --------- flaxdiff/trainer/diffusion_trainer.py | 355 ------------------ flaxdiff/trainer/general_diffusion_trainer.py | 78 +++- flaxdiff/trainer/simple_trainer.py | 50 +-- tests/test_trainer.py | 61 +-- 6 files changed, 107 insertions(+), 621 deletions(-) delete mode 100644 flaxdiff/trainer/autoencoder_trainer.py delete mode 100644 flaxdiff/trainer/diffusion_trainer.py diff --git a/flaxdiff/trainer/__init__.py b/flaxdiff/trainer/__init__.py index b89f89c..871700a 100644 --- a/flaxdiff/trainer/__init__.py +++ b/flaxdiff/trainer/__init__.py @@ -1,3 +1,2 @@ from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics -from .diffusion_trainer import DiffusionTrainer, TrainState -from .general_diffusion_trainer import GeneralDiffusionTrainer, ConditionalInputConfig \ No newline at end of file +from .general_diffusion_trainer import GeneralDiffusionTrainer, TrainState, ConditionalInputConfig diff --git a/flaxdiff/trainer/autoencoder_trainer.py b/flaxdiff/trainer/autoencoder_trainer.py deleted file mode 100644 index f7b2918..0000000 --- a/flaxdiff/trainer/autoencoder_trainer.py +++ /dev/null @@ -1,181 +0,0 @@ -from flax import linen as nn -import jax -from typing import Callable -from dataclasses import field -import jax.numpy as jnp -import optax -from jax.sharding import Mesh, PartitionSpec as P -from jax.experimental.shard_map import shard_map -from typing import Dict, Callable, Sequence, Any, Union, Tuple - -from ..schedulers import NoiseScheduler -from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform - -from flaxdiff.utils import RandomMarkovState - -from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics -from .diffusion_trainer import TrainState -from flaxdiff.models.autoencoder.autoencoder import AutoEncoder - -class AutoEncoderTrainer(SimpleTrainer): - def __init__(self, - model: nn.Module, - input_shape: Union[int, int, int], - latent_dim: int, - spatial_scale: int, - optimizer: optax.GradientTransformation, - rngs: jax.random.PRNGKey, - name: str = "Autoencoder", - **kwargs - ): - super().__init__( - model=model, - input_shapes={"image": input_shape}, - optimizer=optimizer, - rngs=rngs, - name=name, - **kwargs - ) - self.latent_dim = latent_dim - self.spatial_scale = spatial_scale - - - 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 - ) -> 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) - - state = TrainState.create( - apply_fn=model.apply, - params=new_state['params'], - ema_params=new_state['ema_params'], - tx=optimizer, - rngs=rngs, - metrics=Metrics.empty() - ) - - 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 - - def _define_train_step(self, batch_size, null_labels_seq, text_embedder): - noise_schedule: NoiseScheduler = self.noise_schedule - model = self.model - model_output_transform = self.model_output_transform - loss_fn = self.loss_fn - unconditional_prob = self.unconditional_prob - - # Determine the number of unconditional samples - num_unconditional = int(batch_size * unconditional_prob) - - nS, nC = null_labels_seq.shape - null_labels_seq = jnp.broadcast_to( - null_labels_seq, (batch_size, nS, nC)) - - distributed_training = self.distributed_training - - autoencoder = self.autoencoder - - # @jax.jit - def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, local_device_index): - """Train for a single step.""" - rng_state, subkey = rng_state.get_random_key() - subkey = jax.random.fold_in(subkey, local_device_index.reshape()) - local_rng_state = RandomMarkovState(subkey) - - images = batch['image'] - - if autoencoder is not None: - # Convert the images to latent space - local_rng_state, rngs = local_rng_state.get_random_key() - images = autoencoder.encode(images, rngs) - else: - # normalize image - images = (images - 127.5) / 127.5 - - output = text_embedder.encode_from_tokens(batch['text']) - label_seq = output.last_hidden_state - - # Generate random probabilities to decide how much of this batch will be unconditional - - label_seq = jnp.concat( - [null_labels_seq[:num_unconditional], label_seq[num_unconditional:]], axis=0) - - noise_level, local_rng_state = noise_schedule.generate_timesteps(images.shape[0], local_rng_state) - - local_rng_state, rngs = local_rng_state.get_random_key() - noise: jax.Array = jax.random.normal(rngs, shape=images.shape) - - rates = noise_schedule.get_rates(noise_level) - 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_output_transform.pred_transform( - noisy_images, preds, rates) - nloss = loss_fn(preds, expected_output) - # nloss = jnp.mean(nloss, axis=1) - nloss *= noise_schedule.get_weights(noise_level) - nloss = jnp.mean(nloss) - loss = nloss - return loss - - loss, grads = jax.value_and_grad(model_loss)(train_state.params) - if distributed_training: - grads = jax.lax.pmean(grads, "data") - loss = jax.lax.pmean(loss, "data") - train_state = train_state.apply_gradients(grads=grads) - train_state = train_state.apply_ema(self.ema_decay) - return train_state, loss, rng_state - - if distributed_training: - train_step = shard_map(train_step, mesh=self.mesh, in_specs=(P(), P(), P('data'), P('data')), - out_specs=(P(), P(), P())) - train_step = jax.jit(train_step) - - return train_step - - def _define_compute_metrics(self): - @jax.jit - def compute_metrics(state: TrainState, expected, pred): - loss = jnp.mean(jnp.square(pred - expected)) - metric_updates = state.metrics.single_from_model_output(loss=loss) - metrics = state.metrics.merge(metric_updates) - state = state.replace(metrics=metrics) - return state - return compute_metrics - - def fit(self, data, steps_per_epoch, epochs): - null_labels_full = data['null_labels_full'] - local_batch_size = data['local_batch_size'] - text_embedder = data['model'] - super().fit(data, steps_per_epoch, epochs, { - "batch_size": local_batch_size, "null_labels_seq": null_labels_full, "text_embedder": text_embedder}) - -def boolean_string(s): - if type(s) == bool: - return s - return s == 'True' - diff --git a/flaxdiff/trainer/diffusion_trainer.py b/flaxdiff/trainer/diffusion_trainer.py deleted file mode 100644 index adbb834..0000000 --- a/flaxdiff/trainer/diffusion_trainer.py +++ /dev/null @@ -1,355 +0,0 @@ -import flax -from flax import linen as nn -import jax -from typing import Callable -from dataclasses import field -import jax.numpy as jnp -import traceback -import optax -import functools -from jax.sharding import Mesh, PartitionSpec as P -from jax.experimental.shard_map import shard_map -from typing import Dict, Callable, Sequence, Any, Union, Tuple, Type - -from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple -from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform -from ..samplers.common import DiffusionSampler -from ..samplers.ddim import DDIMSampler - -from flaxdiff.utils import RandomMarkovState - -from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics - -from flaxdiff.models.autoencoder.autoencoder import AutoEncoder -from flax.training import dynamic_scale as dynamic_scale_lib -from flaxdiff.inputs import TextEncoder, ConditioningEncoder - -class TrainState(SimpleTrainState): - rngs: jax.random.PRNGKey - ema_params: dict - - def apply_ema(self, decay: float = 0.999): - new_ema_params = jax.tree_util.tree_map( - lambda ema, param: decay * ema + (1 - decay) * param, - self.ema_params, - self.params, - ) - return self.replace(ema_params=new_ema_params) - -from flaxdiff.models.autoencoder.autoencoder import AutoEncoder - -class DiffusionTrainer(SimpleTrainer): - noise_schedule: NoiseScheduler - model_output_transform: DiffusionPredictionTransform - ema_decay: float = 0.999 - native_resolution: int = None - - def __init__(self, - model: nn.Module, - input_shapes: Dict[str, Tuple[int]], - optimizer: optax.GradientTransformation, - noise_schedule: NoiseScheduler, - rngs: jax.random.PRNGKey, - unconditional_prob: float = 0.12, - name: str = "Diffusion", - model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(), - autoencoder: AutoEncoder = None, - encoder: ConditioningEncoder = None, - native_resolution: int = None, - ema_decay: float = 0.999, - **kwargs - ): - super().__init__( - model=model, - input_shapes=input_shapes, - optimizer=optimizer, - rngs=rngs, - name=name, - **kwargs - ) - self.noise_schedule = noise_schedule - self.model_output_transform = model_output_transform - self.unconditional_prob = unconditional_prob - self.ema_decay = ema_decay - - if native_resolution is None: - if 'image' in input_shapes: - native_resolution = input_shapes['image'][1] - elif 'x' in input_shapes: - native_resolution = input_shapes['x'][1] - elif 'sample' in input_shapes: - native_resolution = input_shapes['sample'][1] - else: - raise ValueError("No image input shape found in input shapes") - if autoencoder is not None: - native_resolution = native_resolution * 8 - - self.native_resolution = native_resolution - - self.autoencoder = autoencoder - self.encoder = encoder - - def generate_states( - self, - optimizer: optax.GradientTransformation, - rngs: jax.random.PRNGKey, - model: nn.Module = None, - use_dynamic_scale: bool = False - ) -> Tuple[TrainState, TrainState]: - print("Generating states for DiffusionTrainer") - rngs, subkey = jax.random.split(rngs) - - input_vars = self.get_input_ones() - params = model.init(subkey, **input_vars) - - state = TrainState.create( - apply_fn=model.apply, - params=params, - ema_params=params, - tx=optimizer, - rngs=rngs, - metrics=Metrics.empty(), - dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None - ) - return state, state - - def _define_train_step(self, batch_size): - noise_schedule: NoiseScheduler = self.noise_schedule - model = self.model - model_output_transform = self.model_output_transform - loss_fn = self.loss_fn - unconditional_prob = self.unconditional_prob - - null_labels_full = self.encoder([""]) - null_labels_seq = jnp.array(null_labels_full[0], dtype=jnp.float16) - - conditioning_encoder = self.encoder - - nS, nC = null_labels_seq.shape - null_labels_seq = jnp.broadcast_to( - null_labels_seq, (batch_size, nS, nC)) - - distributed_training = self.distributed_training - - autoencoder = self.autoencoder - - # @jax.jit - def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, local_device_index): - """Train for a single step.""" - rng_state, subkey = rng_state.get_random_key() - subkey = jax.random.fold_in(subkey, local_device_index.reshape()) - local_rng_state = RandomMarkovState(subkey) - - images = batch['image'] - - local_batch_size = images.shape[0] - - # First get the standard deviation of the images - # std = jnp.std(images, axis=(1, 2, 3)) - # is_non_zero = (std > 0) - - images = jnp.array(images, dtype=jnp.float32) - # normalize image - images = (images - 127.5) / 127.5 - - if autoencoder is not None: - # Convert the images to latent space - local_rng_state, rngs = local_rng_state.get_random_key() - images = autoencoder.encode(images, rngs) - - label_seq = conditioning_encoder.encode_from_tokens(batch['text']) - - # Generate random probabilities to decide how much of this batch will be unconditional - local_rng_state, uncond_key = local_rng_state.get_random_key() - # Efficient way to determine unconditional samples for JIT compatibility - uncond_mask = jax.random.bernoulli( - uncond_key, - shape=(local_batch_size,), - p=unconditional_prob - ) - num_unconditional = jnp.sum(uncond_mask).astype(jnp.int32) - - label_seq = jnp.concatenate([null_labels_seq[:num_unconditional], label_seq[num_unconditional:]], axis=0) - - noise_level, local_rng_state = noise_schedule.generate_timesteps(local_batch_size, local_rng_state) - - 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) - - rates = noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(images)) - 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, - 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 - nloss *= noise_schedule.get_weights(noise_level, get_coeff_shapes_tuple(nloss)) - nloss = jnp.mean(nloss) - loss = nloss - return loss - - - if train_state.dynamic_scale is not None: - # dynamic scale takes care of averaging gradients across replicas - grad_fn = train_state.dynamic_scale.value_and_grad( - model_loss, axis_name="data" - ) - dynamic_scale, is_fin, loss, grads = grad_fn(train_state.params) - train_state = train_state.replace(dynamic_scale=dynamic_scale) - else: - grad_fn = jax.value_and_grad(model_loss) - loss, grads = grad_fn(train_state.params) - if distributed_training: - grads = jax.lax.pmean(grads, "data") - - new_state = train_state.apply_gradients(grads=grads) - - if train_state.dynamic_scale is not None: - # if is_fin == False the gradients contain Inf/NaNs and optimizer state and - # params should be restored (= skip this step). - select_fn = functools.partial(jnp.where, is_fin) - new_state = new_state.replace( - opt_state=jax.tree_util.tree_map( - select_fn, new_state.opt_state, train_state.opt_state - ), - params=jax.tree_util.tree_map( - select_fn, new_state.params, train_state.params - ), - ) - - new_state = new_state.apply_ema(self.ema_decay) - - if distributed_training: - loss = jax.lax.pmean(loss, "data") - return new_state, loss, rng_state - - if distributed_training: - train_step = shard_map( - train_step, - mesh=self.mesh, - in_specs=(P(), P(), P('data'), P('data')), - out_specs=(P(), P(), P()), - ) - train_step = jax.jit( - train_step, - donate_argnums=(2) - ) - - return train_step - - def _define_validation_step(self, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): - model = self.model - encoder = self.encoder - autoencoder = self.autoencoder - - null_labels_full = encoder([""]) - null_labels_full = null_labels_full.astype(jnp.float16) - # null_labels_seq = jnp.array(null_labels_full[0], dtype=jnp.float16) - - if self.native_resolution is not None: - image_size = self.native_resolution - elif 'image' in self.input_shapes: - image_size = self.input_shapes['image'][1] - elif 'x' in self.input_shapes: - image_size = self.input_shapes['x'][1] - elif 'sample' in self.input_shapes: - image_size = self.input_shapes['sample'][1] - else: - raise ValueError("No image input shape found in input shapes") - - sampler = sampler_class( - model=model, - noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule, - model_output_transform=self.model_output_transform, - null_labels_seq=null_labels_full, - autoencoder=autoencoder, - guidance_scale=3.0, - ) - - def generate_samples( - val_state: TrainState, - batch, - sampler: DiffusionSampler, - diffusion_steps: int, - ): - labels_seq = encoder.encode_from_tokens(batch) - labels_seq = jnp.array(labels_seq, dtype=jnp.float16) - samples = sampler.generate_images( - params=val_state.ema_params, - resolution=image_size, - num_samples=len(labels_seq), - diffusion_steps=diffusion_steps, - end_step=0, - priors=None, - model_conditioning_inputs=(labels_seq,), - ) - return samples - - return sampler, generate_samples - - def validation_loop( - self, - val_state: SimpleTrainState, - val_step_fn: Callable, - val_ds, - val_steps_per_epoch, - current_step, - diffusion_steps=200, - ): - sampler, generate_samples = val_step_fn - - # sampler = generate_sampler(val_state) - - val_ds = iter(val_ds()) if val_ds else None - # Evaluation step - try: - samples = generate_samples( - val_state, - next(val_ds), - sampler, - diffusion_steps, - ) - - # Put each sample on wandb - if self.wandb is not None and self.wandb: - import numpy as np - from wandb import Image as wandbImage - wandb_images = [] - for i in range(samples.shape[0]): - # convert the sample to numpy - sample = np.array(samples[i]) - # denormalize the image - sample = (sample + 1) * 127.5 - sample = np.clip(sample, 0, 255).astype(np.uint8) - # add the image to the list - wandb_images.append(sample) - # log the images to wandb - self.wandb.log({ - f"sample_{i}": wandbImage(sample, caption=f"Sample {i} at step {current_step}") - }, step=current_step) - except Exception as e: - print("Error logging images to wandb", e) - traceback.print_exc() - - def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): - local_batch_size = data['local_batch_size'] - validation_step_args = { - "sampler_class": sampler_class, - "sampling_noise_schedule": sampling_noise_schedule, - } - super().fit( - data, - train_steps_per_epoch=training_steps_per_epoch, - epochs=epochs, - train_step_args={"batch_size": local_batch_size}, - val_steps_per_epoch=val_steps_per_epoch, - validation_step_args=validation_step_args, - ) diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index 554bad9..b651d9d 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -22,11 +22,22 @@ from flaxdiff.models.autoencoder.autoencoder import AutoEncoder from flax.training import dynamic_scale as dynamic_scale_lib - -# Reuse the TrainState from the DiffusionTrainer -from .diffusion_trainer import TrainState, DiffusionTrainer import shutil + +class TrainState(SimpleTrainState): + rngs: jax.random.PRNGKey + ema_params: dict + + def apply_ema(self, decay: float = 0.999): + new_ema_params = jax.tree_util.tree_map( + lambda ema, param: decay * ema + (1 - decay) * param, + self.ema_params, + self.params, + ) + return self.replace(ema_params=new_ema_params) + + from flaxdiff.metrics.common import EvaluationMetric def generate_modelname( @@ -66,7 +77,7 @@ def generate_modelname( return model_name -class GeneralDiffusionTrainer(DiffusionTrainer): +class GeneralDiffusionTrainer(SimpleTrainer): """ General trainer for diffusion models supporting both images and videos. @@ -91,6 +102,7 @@ def __init__(self, wandb_config: Dict[str, Any] = None, eval_metrics: List[EvaluationMetric] = None, best_tracker_metric: str = "train/best_loss", + ema_decay: float = 0.999, **kwargs ): """ @@ -142,18 +154,25 @@ def __init__(self, self.modelname = modelname wandb_config['config']['modelname'] = modelname + self.noise_schedule = noise_schedule + self.model_output_transform = model_output_transform + self.unconditional_prob = unconditional_prob + self.autoencoder = autoencoder + self.ema_decay = ema_decay + + if native_resolution is None: + sample_shape = input_config.sample_data_shape + native_resolution = sample_shape[-2] + if autoencoder is not None: + native_resolution = native_resolution * autoencoder.downscale_factor + self.native_resolution = native_resolution + super().__init__( model=model, input_shapes=input_shapes, optimizer=optimizer, - noise_schedule=noise_schedule, - unconditional_prob=unconditional_prob, - autoencoder=autoencoder, - model_output_transform=model_output_transform, rngs=rngs, name=name, - native_resolution=native_resolution, - encoder=None, # Don't use the default encoder from the parent class wandb_config=wandb_config, **kwargs ) @@ -175,6 +194,45 @@ def __init__(self, for m in eval_metrics: self.metric_higher_is_better[f"val/{m.name}"] = getattr(m, 'higher_is_better', False) + def generate_states( + self, + optimizer: optax.GradientTransformation, + rngs: jax.random.PRNGKey, + model: nn.Module = None, + use_dynamic_scale: bool = False + ) -> Tuple[TrainState, TrainState]: + print("Generating states for DiffusionTrainer") + rngs, subkey = jax.random.split(rngs) + + input_vars = self.get_input_ones() + params = model.init(subkey, **input_vars) + + state = TrainState.create( + apply_fn=model.apply, + params=params, + ema_params=params, + tx=optimizer, + rngs=rngs, + metrics=Metrics.empty(), + dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None + ) + return state, state + + def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): + local_batch_size = data['local_batch_size'] + validation_step_args = { + "sampler_class": sampler_class, + "sampling_noise_schedule": sampling_noise_schedule, + } + return super().fit( + data, + train_steps_per_epoch=training_steps_per_epoch, + epochs=epochs, + train_step_args={"batch_size": local_batch_size}, + val_steps_per_epoch=val_steps_per_epoch, + validation_step_args=validation_step_args, + ) + def _is_video_data(self): sample_data_shape = self.input_config.sample_data_shape return len(sample_data_shape) == 5 diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py index 8b69f10..e9ff541 100644 --- a/flaxdiff/trainer/simple_trainer.py +++ b/flaxdiff/trainer/simple_trainer.py @@ -349,52 +349,10 @@ def save(self, epoch=0, step=0, state=None, rngstate=None): print("Error saving checkpoint outer", e) def _define_train_step(self, **kwargs): - model = self.model - loss_fn = self.loss_fn - distributed_training = self.distributed_training - - def train_step(train_state: SimpleTrainState, rng_state: RandomMarkovState, batch, local_device_indexes): - """Train for a single step.""" - images = batch['image'] - labels = batch['label'] - - def model_loss(params): - preds = model.apply(params, images) - expected_output = labels - nloss = loss_fn(preds, expected_output) - loss = jnp.mean(nloss) - return loss - loss, grads = jax.value_and_grad(model_loss)(train_state.params) - if distributed_training: - grads = jax.lax.pmean(grads, "data") - train_state = train_state.apply_gradients(grads=grads) - return train_state, loss, rng_state - - if distributed_training: - train_step = shard_map(train_step, mesh=self.mesh, in_specs=(P(), P(), P('data'), P('data')), out_specs=(P(), P('data'), P())) - train_step = jax.pmap(train_step) - return train_step - - def _define_validation_step(self): - model = self.model - loss_fn = self.loss_fn - distributed_training = self.distributed_training - - def validation_step(state: SimpleTrainState, batch): - preds = model.apply(state.params, batch['image']) - expected_output = batch['label'] - loss = jnp.mean(loss_fn(preds, expected_output)) - if distributed_training: - loss = jax.lax.pmean(loss, "data") - metric_updates = state.metrics.single_from_model_output( - loss=loss, logits=preds, labels=expected_output) - metrics = state.metrics.merge(metric_updates) - state = state.replace(metrics=metrics) - return state - if distributed_training: - validation_step = shard_map(validation_step, mesh=self.mesh, in_specs=(P(), P('data')), out_specs=(P())) - validation_step = jax.pmap(validation_step) - return validation_step + raise NotImplementedError("Subclasses must define their train step") + + def _define_validation_step(self, **kwargs): + raise NotImplementedError("Subclasses must define their validation step") def summary(self): input_vars = self.get_input_ones() diff --git a/tests/test_trainer.py b/tests/test_trainer.py index a33c15f..630436a 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,7 +1,8 @@ """Trainer smoke tests: training without wandb, and checkpoint save/restore. -These run on CPU with a tiny regression model through SimpleTrainer, which -owns the state/checkpoint/loop machinery shared by all trainers. +These run the real GeneralDiffusionTrainer on CPU with a tiny DiT and an +unconditional synthetic dataset - the same code path a real run takes, minus +wandb and conditioning encoders. """ import jax @@ -9,22 +10,27 @@ import numpy as np import optax import pytest -from flax import linen as nn -from flaxdiff.trainer.simple_trainer import SimpleTrainer +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.models.simple_dit import SimpleDiT +from flaxdiff.predictors import get_diffusion_preset +from flaxdiff.trainer import GeneralDiffusionTrainer - -class TinyModel(nn.Module): - @nn.compact - def __call__(self, x): - return nn.Dense(2)(x) +RES = 8 def make_trainer(tmp_path, name="smoke", load_from_checkpoint=None): - return SimpleTrainer( - model=TinyModel(), - input_shapes={"x": (4,)}, + train_schedule, _, transform = get_diffusion_preset("edm") + return GeneralDiffusionTrainer( + model=SimpleDiT(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="image", + sample_data_shape=(RES, RES, 3), + conditions=[], + ), rngs=jax.random.PRNGKey(0), name=name, wandb_config=None, @@ -35,21 +41,18 @@ def make_trainer(tmp_path, name="smoke", load_from_checkpoint=None): def batch_iterator(): + # uint8-range images, as the data pipeline provides them + images = np.tile(np.linspace(0, 255, RES, dtype=np.float32)[None, :, None, None], (8, 1, RES, 3)) while True: - yield {"image": jnp.ones((8, 4)), "label": jnp.zeros((8, 2))} + yield {"image": jnp.asarray(images)} def test_fit_without_wandb(tmp_path): trainer = make_trainer(tmp_path) - data = {"train": batch_iterator, "train_len": 32} - 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 + data = {"train": batch_iterator, "train_len": 32, "local_batch_size": 8} + state = trainer.fit(data, training_steps_per_epoch=4, epochs=1, val_steps_per_epoch=0) + assert state is not None + assert int(state.step) == 4 def test_save_writes_checkpoint(tmp_path): @@ -62,14 +65,18 @@ def test_save_writes_checkpoint(tmp_path): def test_restore_preserves_optimizer_state(tmp_path): trainer = make_trainer(tmp_path) - # One real update so the adam moments and step counter are non-trivial + # One real update so the adam moments, step counter and EMA are non-trivial grads = jax.tree.map(jnp.ones_like, trainer.state.params) - trainer.state = trainer.state.apply_gradients(grads=grads) + trainer.state = trainer.state.apply_gradients(grads=grads).apply_ema(0.99) trainer.save(epoch=0, step=1) restored = make_trainer(tmp_path, load_from_checkpoint=trainer.checkpoint_path()) assert int(restored.state.step) == 1, "optimizer step counter was reset" - old_mu = jax.tree.leaves(trainer.state.opt_state) - new_mu = jax.tree.leaves(restored.state.opt_state) - assert all(np.allclose(a, b) for a, b in zip(old_mu, new_mu)), "adam moments were reset" + old_opt = jax.tree.leaves(trainer.state.opt_state) + new_opt = jax.tree.leaves(restored.state.opt_state) + assert all(np.allclose(a, b) for a, b in zip(old_opt, new_opt)), "adam moments were reset" + + 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" From 18796ae605bdfe6f9ed8e5a7182be03f341e7ffc Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Thu, 23 Jul 2026 00:22:57 -0400 Subject: [PATCH 6/6] feat: legacy checkpoint conversion script for the new param layout Remaps the top-level module paths the backbone consolidation moved (embed/, conditioning/, output/) so pre-consolidation DiT-family checkpoints load against the new models. UNet/UViT need no conversion; MMDiT checkpoints are refused implicitly since the architecture itself changed. Documents the FourierEmbedding frequency provenance caveat for pre-jax-0.5 checkpoints. Co-Authored-By: Claude Fable 5 --- scripts/convert_legacy_checkpoint.py | 85 ++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 scripts/convert_legacy_checkpoint.py diff --git a/scripts/convert_legacy_checkpoint.py b/scripts/convert_legacy_checkpoint.py new file mode 100644 index 0000000..434d665 --- /dev/null +++ b/scripts/convert_legacy_checkpoint.py @@ -0,0 +1,85 @@ +"""One-time conversion of pre-consolidation checkpoints to the new param layout. + +The DiT backbone consolidation moved the patchify/conditioning/output modules +into submodules (embed/, conditioning/, output/), which renames their param +paths. This remaps an old orbax checkpoint in place and re-saves it so it +loads against the new models. UNet and UViT checkpoints need no conversion. +MMDiT checkpoints are NOT convertible - the architecture itself changed +(dual-stream rewrite), not just the naming. + +Usage: + python scripts/convert_legacy_checkpoint.py + +Note on time conditioning: FourierEmbedding frequencies used to come from +jax's PRNG, whose default implementation changed in jax 0.5.0. Checkpoints +trained before that already sample with subtly different time conditioning on +modern jax; the frequencies are now fixed numpy values, so models trained +from this version onward are stable. This script cannot repair the old +provenance - expect slightly off conditioning on pre-0.5.0 checkpoints. +""" + +import sys + +import numpy as np +import orbax.checkpoint + +# Top-level module renames from the DiT backbone consolidation +RENAMES = { + 'patch_embed': ('embed', 'patch_embed'), + 'hilbert_projection': ('embed', 'hilbert_projection'), + 'time_embed': ('conditioning', 'time_embed'), + 'text_context_proj': ('conditioning', 'text_context_proj'), + 'final_norm': ('output', 'final_norm'), + 'final_proj': ('output', 'final_proj'), +} + + +def convert_params(params: dict) -> dict: + """Remap the top-level module names of a DiT-family param tree.""" + converted = {} + for key, value in params.items(): + if key in RENAMES: + parent, child = RENAMES[key] + converted.setdefault(parent, {})[child] = value + else: + converted[key] = value + return converted + + +def convert_state(state: dict) -> dict: + state = dict(state) + for key in ('params', 'ema_params'): + if key in state and isinstance(state[key], dict) and 'params' in state[key]: + state[key] = {**state[key], 'params': convert_params(state[key]['params'])} + return state + + +def main(): + if len(sys.argv) != 3: + print(__doc__) + sys.exit(1) + checkpoint_dir, output_dir = sys.argv[1], sys.argv[2] + + checkpointer = orbax.checkpoint.PyTreeCheckpointer() + manager = orbax.checkpoint.CheckpointManager( + checkpoint_dir, checkpointer, + orbax.checkpoint.CheckpointManagerOptions(create=False)) + step = manager.latest_step() + print(f"Converting checkpoint at step {step} from {checkpoint_dir}") + ckpt = manager.restore(step) + + for key in ('state', 'best_state'): + if key in ckpt and ckpt[key] is not None: + ckpt[key] = convert_state(ckpt[key]) + print(f"Converted {key}") + + out_manager = orbax.checkpoint.CheckpointManager( + output_dir, orbax.checkpoint.PyTreeCheckpointer(), + orbax.checkpoint.CheckpointManagerOptions(create=True)) + out_manager.save(step, ckpt) + out_manager.wait_until_finished() + print(f"Saved converted checkpoint to {output_dir} at step {step}") + + +if __name__ == '__main__': + main()