finish transformer

Author: jfacevedo-google

src/maxdiffusion/configs/base_flux.yml

@@ -24,6 +24,9 @@ save_config_to_gcs: False
 log_period: 100
 
 pretrained_model_name_or_path: 'black-forest-labs/FLUX.1-dev'
+clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax'
+t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax'
+
 unet_checkpoint: ''
 revision: 'refs/pr/95'
 # This will convert the weights to this dtype.
@@ -41,7 +44,7 @@ precision: "DEFAULT"
 # Set true to load weights from pytorch
 from_pt: False
 split_head_dim: True
-attention: 'dot_product' # Supported attention: dot_product, flash
+attention: 'flash' # Supported attention: dot_product, flash
 flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
@@ -171,7 +174,7 @@ max_train_steps: 200
 num_train_epochs: 1
 seed: 0
 output_dir: 'sdxl-model-finetuned'
-per_device_batch_size: 2
+per_device_batch_size: 1
 
 warmup_steps_fraction: 0.0
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.

src/maxdiffusion/generate_flux.py

@@ -16,22 +16,30 @@
 
 from typing import Any, Callable, Dict, List, Optional, Union, Sequence
 from absl import app
-
+import functools
 import numpy as np
 import jax
+from jax.sharding import Mesh, PositionalSharding
 import jax.numpy as jnp
 from chex import Array
 from transformers import (
   CLIPTokenizer,
   FlaxCLIPTextModel,
   T5TokenizerFast,
-  T5EncoderModel
+  T5EncoderModel,
+  FlaxT5EncoderModel
 )
 
 from maxdiffusion import FlaxAutoencoderKL
 from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
-
 from maxdiffusion import pyconfig
+from max_utils import (
+  device_put_replicated,
+  get_memory_allocations,
+  create_device_mesh,
+  get_flash_block_sizes,
+  get_precision
+)
 
 def prepare_latent_image_ids(height, width):
   latent_image_ids = jnp.zeros((height, width, 3))
@@ -133,19 +141,17 @@ def get_t5_prompt_embeds(
     truncation=True,
     return_length=False,
     return_overflowing_tokens=False,
-    return_tensors="pt"
+    return_tensors="np"
   )
   text_input_ids = text_inputs.input_ids
-
   prompt_embeds = text_encoder(text_input_ids, output_hidden_states=False)[0]
   dtype = text_encoder.dtype
-  prompt_embeds = prompt_embeds.to(dtype=dtype)
+  prompt_embeds = prompt_embeds.astype(dtype)
 
   _, seq_len, _ = prompt_embeds.shape
-
   # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
-  prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-  prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+  prompt_embeds = jnp.tile(prompt_embeds, (1, num_images_per_prompt, 1))
+  prompt_embeds = jnp.reshape(prompt_embeds, (batch_size * num_images_per_prompt, seq_len, -1))
 
   return prompt_embeds
 
@@ -178,15 +184,16 @@ def encode_prompt(
     tokenizer=t5_tokenizer,
     text_encoder=t5_text_encoder
   )
-  prompt_embeds = jnp.asarray(prompt_embeds.detach().numpy())
 
   text_ids = jnp.zeros((prompt_embeds.shape[0], prompt_embeds.shape[1], 3)).astype(jnp.bfloat16)
   return prompt_embeds, pooled_prompt_embeds, text_ids
 
 def run(config):
   from maxdiffusion.models.flux.util import load_flow_model
 
-  rng = jax.random.PRNGKey(config.seed)
+  rng = jax.random.key(config.seed)
+  devices_array = create_device_mesh(config)
+  mesh = Mesh(devices_array, config.mesh_axes)
 
   per_host_number_of_images = 1#config.per_device_batch_size * jax.local_device_count()
 
@@ -201,11 +208,18 @@ def run(config):
   )
   vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
 
-  # LOAD UNET
-
+  # LOAD TRANSFORMER
+  flash_block_sizes = get_flash_block_sizes(config)
   transformer = FluxTransformer2DModel.from_config(
     config.pretrained_model_name_or_path,
-    subfolder="transformer"
+    subfolder="transformer",
+    mesh=mesh,
+    split_head_dim=config.split_head_dim,
+    attention_kernel=config.attention,
+    flash_block_sizes=flash_block_sizes,
+    dtype=config.activations_dtype,
+    weights_dtype=config.weights_dtype,
+    precision=get_precision(config)
   )
 
   num_channels_latents = transformer.in_channels // 4
@@ -242,34 +256,40 @@ def run(config):
     dtype=config.weights_dtype
   )
 
-  t5_encoder_pt = T5EncoderModel.from_pretrained(
-    config.pretrained_model_name_or_path,
-    subfolder="text_encoder_2",
+  t5_encoder = FlaxT5EncoderModel.from_pretrained(
+    config.clip_model_name_or_path,
+    dtype=config.weights_dtype
   )
-
   t5_tokenizer = T5TokenizerFast.from_pretrained(
     config.pretrained_model_name_or_path,
     subfolder="tokenizer_2",
   )
 
+  encoders_sharding = PositionalSharding(devices_array).replicate()
+  partial_device_put_replicated = functools.partial(device_put_replicated, sharding=encoders_sharding)
+  clip_text_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), clip_text_encoder.params)
+  clip_text_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, clip_text_encoder.params)
+  t5_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), t5_encoder.params)
+  t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, t5_encoder.params)
+
   prompt_embeds, pooled_prompt_embeds, text_ids = encode_prompt(
     prompt=config.prompt,
     prompt_2=config.prompt_2,
     clip_tokenizer=clip_tokenizer,
     clip_text_encoder=clip_text_encoder,
     t5_tokenizer=t5_tokenizer,
-    t5_text_encoder=t5_encoder_pt,
+    t5_text_encoder=t5_encoder,
     num_images_per_prompt=per_host_number_of_images
   )
 
   def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timesteps, guidance, pooled_prompt_embeds):
-    print("latents.shape: ", latents.shape)
-    print("latent_image_ids.shape: ", latent_image_ids.shape)
-    print("text_ids.shape: ", text_ids.shape)
-    print("prompt_embeds: ", prompt_embeds.shape)
-    print("timesteps.shape: ", timesteps.shape)
-    print("guidance.shape: ", guidance.shape)
-    print("pooled_prompt_embeds.shape: ", pooled_prompt_embeds.shape)
+    print("latents.shape: ", latents.shape, latents.dtype)
+    print("latent_image_ids.shape: ", latent_image_ids.shape, latent_image_ids.dtype)
+    print("text_ids.shape: ", text_ids.shape, text_ids.dtype)
+    print("prompt_embeds: ", prompt_embeds.shape, prompt_embeds.dtype)
+    print("timesteps.shape: ", timesteps.shape, timesteps.dtype)
+    print("guidance.shape: ", guidance.shape, guidance.dtype)
+    print("pooled_prompt_embeds.shape: ", pooled_prompt_embeds.shape, pooled_prompt_embeds.dtype)
 
   timesteps = jnp.asarray([1.0], dtype=jnp.bfloat16)
   guidance = jnp.asarray([3.5], dtype=jnp.bfloat16)
@@ -282,17 +302,19 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
     guidance,
     pooled_prompt_embeds
   )
-  
-  transformer_params = transformer.init(
-    {"params" : rng},
-    img=latents,
-    img_ids=latent_image_ids,
-    txt=prompt_embeds,
-    txt_ids=text_ids,
-    timesteps=timesteps,
-    guidance=guidance,
-    y=pooled_prompt_embeds
-  )["params"]
+  get_memory_allocations()
+  transformer_params = transformer.init_weights(rng, True)
+  # transformer_params = transformer.init(
+  #   {"params" : rng},
+  #   img=latents,
+  #   img_ids=latent_image_ids,
+  #   txt=prompt_embeds,
+  #   txt_ids=text_ids,
+  #   timesteps=timesteps,
+  #   guidance=guidance,
+  #   y=pooled_prompt_embeds
+  # )["params"]
+  get_memory_allocations()
   breakpoint()
 
 

src/maxdiffusion/models/flux/modules/layers.py

@@ -177,10 +177,9 @@ def __call__(self, vec: Array) -> tuple[ModulationOut, ModulationOut | None]:
     )(nn.silu(vec))
 
     out = jnp.split(lin[:, None, :], multiplier, axis=-1)
-
     return (
       ModulationOut(*out[:3]),
-      ModulationOut(*out[3:] if self.double else None)
+      ModulationOut(*out[3:]) if self.double else None
     )
 
 class SingleStreamBlock(nn.Module):
@@ -209,7 +208,6 @@ def __call__(self, x: Array, vec: Array, pe: Array) -> Array:
       precision=self.precision
     )(vec)
     x_mod = (1 + mod.scale) * nn.LayerNorm(
-      self.hidden_size,
       use_scale=False,
       use_bias=False,
       epsilon=1e-6,
@@ -261,7 +259,7 @@ def __call__(self, x: Array, vec: Array, pe: Array) -> Array:
         ("embed", "heads")
       ),
       name="linear2"
-    )(jnp.concatenate((attn, nn.genu(mlp)), 2))
+    )(jnp.concatenate((attn, nn.gelu(mlp)), 2))
     return x + mod.gate * output    
 
 class DoubleStreamBlock(nn.Module):
@@ -279,7 +277,7 @@ class DoubleStreamBlock(nn.Module):
   attention_kernel: str = "dot_product"
 
   @nn.compact
-  def __call__(self, img: Array, txt: Array, vec: Array, pe: Array):
+  def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array, Array]:
 
     mlp_hidden_dim = int(self.hidden_size * self.mlp_ratio)
 
@@ -422,7 +420,7 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array):
     )
 
     # calculate the txt blocks
-    txt = txt + txt_mod1.gate * nn.Dense(
+    txt_proj = nn.Dense(
       self.hidden_size,
       use_bias=True,
       dtype=self.dtype,
@@ -433,6 +431,8 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array):
         ("heads", "embed")
       ),
     )(txt_attn)
+    txt = txt + txt_mod1.gate * txt_proj
+
     txt = txt + txt_mod2.gate * nn.Sequential(
       [
         nn.Dense(
@@ -466,4 +466,54 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array):
       )(txt) + txt_mod2.shift
     )
 
-    return img, txt
+    return img, txt
+
+class LastLayer(nn.Module):
+  hidden_size: int
+  patch_size: int
+  out_channels: int
+  dtype: jnp.dtype = jnp.float32
+  weights_dtype: jnp.dtype = jnp.float32
+  precision: jax.lax.Precision = None
+
+  @nn.compact
+  def __call__(self, x: Array, vec: Array) -> Array:
+    shift, scale = jnp.split(
+      nn.Sequential(
+        [
+          nn.silu,
+          nn.Dense(
+            2 * self.hidden_size,
+            use_bias=True,
+            param_dtype=self.weights_dtype,
+            dtype=self.dtype,
+            precision=self.precision,
+            kernel_init=nn.with_logical_partitioning(
+              nn.initializers.lecun_normal(),
+              ("embed", "heads")
+            )
+          )
+        ]
+      )(vec), 2, axis=1
+    )
+    norm_final = nn.LayerNorm(
+      epsilon=1e-6,
+      use_scale=False,
+      use_bias=False,
+      param_dtype=self.weights_dtype,
+      name="norm_final"
+    )(x)
+    x = (1 + scale[:, None, :]) * norm_final + shift[:, None, :]
+    x = nn.Dense(
+      self.patch_size * self.patch_size * self.out_channels,
+      use_bias=True,
+      param_dtype=self.weights_dtype,
+            dtype=self.dtype,
+            precision=self.precision,
+      kernel_init=nn.with_logical_partitioning(
+        nn.initializers.lecun_normal(),
+        ("heads", "embed")
+      ),
+      name="linear"
+    )
+    return x