start adding flux transformer.

Author: jfacevedo-google

src/maxdiffusion/models/embeddings_flax.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-
+from typing import List, Union
+import jax
 import flax.linen as nn
 import jax.numpy as jnp
 
@@ -96,3 +97,163 @@ def __call__(self, timesteps):
     return get_sinusoidal_embeddings(
         timesteps, embedding_dim=self.dim, flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.freq_shift
     )
+
+def get_1d_rotary_pos_embed(
+  dim: int,
+  pos: Union[jnp.array, int],
+  theta: float = 10000.0,
+  use_real=False,
+  linear_factor=1.0,
+  ntk_factor=1.0,
+  repeat_interleave_real=True,
+  freqs_dtype=jnp.float32
+):
+  """
+  Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+  """
+  assert dim % 2 == 0
+
+  if isinstance(pos, int):
+    pos = jnp.arange(pos)
+  
+  theta = theta * ntk_factor
+  freqs = (
+    1.0
+    / (theta ** (jnp.arange(0, dim, 2, dtype=freqs_dtype)[: (dim // 2)] / dim))
+    / linear_factor
+  )
+  freqs = jnp.outer(pos, freqs)
+  if use_real and repeat_interleave_real:
+    freqs_cos = jnp.cos(freqs).repeat(2, axis=1).astype(jnp.float32)
+    freqs_sin = jnp.sin(freqs).repeat(2, axis=1).astype(jnp.float32)
+    return freqs_cos, freqs_sin
+  elif use_real:
+    freqs_cos = jnp.concatenate([jnp.cos(freqs), jnp.cos(freqs)], axis=-1).astype(jnp.float32)
+    freqs_sin = jnp.concatenate([jnp.sin(freqs), jnp.sin(freqs)], axis=-1).astype(jnp.float32) 
+    return freqs_cos, freqs_sin
+  else:
+    raise ValueError(f"use_real {use_real} and repeat_interleave_real {repeat_interleave_real} is not supported")
+
+class PixArtAlphaTextProjection(nn.Module):
+  """
+  Projects caption embeddings. Also handles dropout for classifier-free guidance.
+
+  Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+  """
+
+  hidden_size: int
+  out_features: int = None
+  act_fn: str ='gelu_tanh'
+  dtype: jnp.dtype = jnp.float32
+  weights_dtype: jnp.dtype = jnp.float32
+  precision: jax.lax.Precision = None
+
+  @nn.compact
+  def __call__(self, caption):
+    hidden_states = nn.Dense(
+      self.hidden_size,
+      use_bias=True,
+      dtype=self.dtype,
+      param_dtype=self.weights_dtype,
+      precision=self.precision
+      )(caption)
+    
+    if self.act_fn == 'gelu_tanh':
+      act_1 = nn.gelu
+    elif self.act_fn == 'silu':
+      act_1 = nn.swish
+    else:
+      raise ValueError(f"Unknown activation function: {self.act_fn}")
+    hidden_states = act_1(hidden_states)
+
+    hidden_states = nn.Dense(self.out_features)(hidden_states)
+    return hidden_states
+
+
+class FluxPosEmbed(nn.Module):
+  theta: int
+  axes_dim: List[int]
+
+  @nn.compact
+  def __call__(self, ids):
+    n_axes = ids.shape[-1]
+    cos_out = []
+    sin_out = []
+    pos = ids.astype(jnp.float32)
+    freqs_dtype = jnp.float32
+    for i in range(n_axes):
+      cos, sin = get_1d_rotary_pos_embed(
+        self.axes_dim[i], pos[:i],
+        repeat_interleave_real=True,
+        use_real=True,
+        freqs_dtype=freqs_dtype
+      )
+      cos_out.append(cos)
+      sin_out.append(sin)
+    
+    freqs_cos = jnp.concatenate(cos_out, axis=-1)
+    freqs_sin = jnp.concatenate(sin_out, axis=-1)
+    return freqs_cos, freqs_sin
+
+class CombinedTimestepTextProjEmbeddings(nn.Module):
+  embedding_dim: int
+  pooled_projection_dim: int
+  dtype: jnp.dtype = jnp.float32
+  weights_dtype: jnp.dtype = jnp.float32
+  precision: jax.lax.Precision = None
+
+  @nn.compact
+  def __call__(self, timestep, pooled_projection):
+    timesteps_proj = FlaxTimesteps(dim=256, flip_sin_to_cos=True, freq_shift=0)(timestep)
+    timestep_emb = FlaxTimestepEmbedding(
+      time_embed_dim=self.embedding_dim,
+      dtype=self.dtype,
+      weights_dtype=self.weights_dtype,
+      precision=self.precision
+    )(timesteps_proj.astype(pooled_projection.dtype))
+    
+    pooled_projections = PixArtAlphaTextProjection(
+      self.embedding_dim,
+      act_fn='silu',
+      dtype=self.dtype,
+      weights_dtype=self.weights_dtype,
+    )(pooled_projection)
+
+    conditioning = timestep_emb + pooled_projection
+    return conditioning
+
+class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
+  embedding_dim: int
+  pooled_projection_dim: int
+  dtype: jnp.dtype = jnp.float32
+  weights_dtype: jnp.dtype = jnp.float32
+  precision: jax.lax.Precision = None
+
+  @nn.compact
+  def __call__(self, timestep, guidance, pooled_projection):
+    timesteps_proj = FlaxTimesteps(dim=256, flip_sin_to_cos=True, freq_shift=0)(timestep)
+    timestep_emb = FlaxTimestepEmbedding(
+      time_embed_dim=self.embedding_dim,
+      dtype=self.dtype,
+      weights_dtype=self.weights_dtype
+    )(timesteps_proj.astype(pooled_projection.dtype))
+    
+    guidance_proj = FlaxTimesteps(dim=256, flip_sin_to_cos=True, freq_shift=0)(guidance)
+    guidance_emb = FlaxTimestepEmbedding(
+      time_embed_dim=self.embedding_dim,
+      dtype=self.dtype,
+      weights_dtype=self.weights_dtype
+    )(guidance_proj.astype(pooled_projection.dtype))
+
+    time_guidance_emb = timestep_emb + guidance_emb
+
+    pooled_projections = PixArtAlphaTextProjection(
+      self.embedding_dim,
+      act_fn='silu',
+      dtype=self.dtype,
+      weights_dtype=self.weights_dtype,
+      precision=self.precision
+      )(pooled_projection)
+    conditioning = time_guidance_emb + pooled_projections
+
+    return conditioning

src/maxdiffusion/models/normalization_flax.py

@@ -14,12 +14,30 @@
  limitations under the License.
  """
 
-'''This script is used an example of how to shard the UNET on TPU.'''
-
+import jax
 import jax.numpy as jnp
 import flax.linen as nn
 
-class FlaxAdaLayerNormZeroSingle(nn.Module):
+class AdaLayerNormContinuous(nn.Module):
+    embedding_dim: int
+    elementwise_affine: bool = True
+    eps: float = 1e-5
+    bias: bool = True
+    norm_type: str = "layer_norm"
+    dtype: jnp.dtype = jnp.float32
+    weights_dtype: jnp.dtype = jnp.float32
+    precision: jax.lax.Precision = None
+
+    @nn.compact
+    def __call__(self, x, conditioning_embedding):
+        assert self.norm_type == 'layer_norm'
+        emb = nn.Dense(self.embedding_dim * 2, use_bias=self.bias)(nn.sigmoid(conditioning_embedding))
+        scale, shift = jnp.split(emb, 2, axis=1)
+        x = nn.LayerNorm(epsilon=self.eps, use_bias=self.elementwise_affine, use_scale=self.elementwise_affine)(x)
+        x *= (1 + scale[:, None, :]) + shift[:, None, :]
+        return x
+
+class AdaLayerNormZeroSingle(nn.Module):
     r"""
     Norm layer adaptive layer norm zero (adaLN-Zero).
 
@@ -30,11 +48,20 @@ class FlaxAdaLayerNormZeroSingle(nn.Module):
     embedding_dim: int
     norm_type: str = "layer_norm"
     bias: bool = True
+    dtype: jnp.dtype = jnp.float32
+    weights_dtype: jnp.dtype = jnp.float32
+    precision: jax.lax.Precision = None
 
     @nn.compact
     def __call__(self, x, emb):
         emb = nn.silu(emb)
-        emb = nn.Dense(3 * self.embedding_dim, use_bias=self.bias)(emb)
+        emb = nn.Dense(
+              3 * self.embedding_dim,
+              use_bias=self.bias,
+              dtype=self.dtype,
+              param_dtype=self.weights_dtype,
+              precision=self.precision
+            )(emb)
         shift_msa, scale_msa, gate_msa = jnp.split(emb, 3, axis=1)
         if self.norm_type == "layer_norm":
             x = nn.LayerNorm(epsilon=1e-6, use_bias=False, use_scale=False)(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]