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/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/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/dit_common.py b/flaxdiff/models/dit_common.py new file mode 100644 index 0000000..bf4ce78 --- /dev/null +++ b/flaxdiff/models/dit_common.py @@ -0,0 +1,323 @@ +""" +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 + modulated: bool = False # adaLN shift/scale on the final norm (DiT FinalLayer) + norm_epsilon: float = 1e-5 + dtype: Optional[Dtype] = None + precision: PrecisionLike = None + + @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 + 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) + + 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/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/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/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_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/simple_vit.py b/flaxdiff/models/simple_vit.py index 0aca9fe..7dc4b2d 100644 --- a/flaxdiff/models/simple_vit.py +++ b/flaxdiff/models/simple_vit.py @@ -11,8 +11,8 @@ 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 .simple_dit import DiTBlock +from .vit_common import unpatchify, PatchEmbedding, RotaryEmbedding, RoPEAttention, AdaLNParams +from .dit_common import ModulatedBlock class UViT(nn.Module): @@ -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( @@ -310,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, @@ -324,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, @@ -346,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/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..66eab3e 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 @@ -37,22 +39,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] @@ -144,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 @@ -186,57 +175,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 @@ -248,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/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 147bc22..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( @@ -64,48 +75,9 @@ 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): +class GeneralDiffusionTrainer(SimpleTrainer): """ General trainer for diffusion models supporting both images and videos. @@ -130,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 ): """ @@ -181,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 ) @@ -214,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 1d23acd..e9ff541 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( @@ -381,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() @@ -441,17 +367,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 +425,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 +477,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 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() 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/tests/test_models.py b/tests/test_models.py index 7747b49..4df3335 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)) @@ -137,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" 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" 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]