batch text encoding.

Author: jfacevedo-google

src/maxdiffusion/configs/base_flux_dev.yml

@@ -111,7 +111,7 @@ diffusion_scheduler_config: {
 base_output_directory: ""
 
 # Hardware
-hardware: 'gpu' # Supported hardware types are 'tpu', 'gpu'
+hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
 
 # Parallelism
 mesh_axes: ['data', 'fsdp', 'tensor']

src/maxdiffusion/pipelines/flux/flux_pipeline.py

@@ -190,6 +190,8 @@ def get_t5_prompt_embeds(
       tokenizer: AutoTokenizer,
       text_encoder: FlaxT5EncoderModel,
       max_sequence_length: int = 512,
+      encode_in_batches=False,
+      encode_batch_size=None,
   ):
 
     prompt = [prompt] if isinstance(prompt, str) else prompt
@@ -205,13 +207,23 @@ def get_t5_prompt_embeds(
         return_tensors="np",
     )
     text_input_ids = text_inputs.input_ids
-    prompt_embeds = text_encoder(text_input_ids, attention_mask=None, output_hidden_states=False)["last_hidden_state"]
+    if encode_in_batches:
+      prompt_embeds = None
+      for i in range(0, text_input_ids.shape[0], encode_batch_size):
+        batch_prompt_embeds = text_encoder(text_input_ids[i:i+encode_batch_size], attention_mask=None, output_hidden_states=False)["last_hidden_state"]
+        if prompt_embeds == None:
+          prompt_embeds = batch_prompt_embeds
+        else:
+          prompt_embeds = jnp.concatenate([prompt_embeds, batch_prompt_embeds])
+    else:
+      prompt_embeds = text_encoder(text_input_ids, attention_mask=None, output_hidden_states=False)["last_hidden_state"]
+      _, seq_len, _ = prompt_embeds.shape
+      # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+      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))
+
     dtype = text_encoder.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 = 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
 
@@ -226,7 +238,12 @@ def encode_prompt(
       t5_text_encoder: FlaxT5EncoderModel,
       num_images_per_prompt: int = 1,
       max_sequence_length: int = 512,
+      encode_in_batches: bool = False,
+      encode_batch_size: int = None
   ):
+    
+    if encode_in_batches:
+      assert encode_in_batches is not None
 
     prompt = [prompt] if isinstance(prompt, str) else prompt
     prompt_2 = prompt or prompt_2
@@ -242,6 +259,8 @@ def encode_prompt(
         tokenizer=t5_tokenizer,
         text_encoder=t5_text_encoder,
         max_sequence_length=max_sequence_length,
+        encode_in_batches=encode_in_batches,
+        encode_batch_size=encode_batch_size
     )
 
     text_ids = jnp.zeros((prompt_embeds.shape[0], prompt_embeds.shape[1], 3)).astype(jnp.bfloat16)

src/maxdiffusion/train_flux.py

@@ -24,14 +24,13 @@
     mllog_utils,
 )
 
-from maxdiffusion.trainers.flux_trainer import FluxTrainer
-
 from maxdiffusion.train_utils import (
     validate_train_config,
 )
 
 
 def train(config):
+  from maxdiffusion.trainers.flux_trainer import FluxTrainer
   trainer = FluxTrainer(config)
   trainer.start_training()
 

src/maxdiffusion/trainers/flux_trainer.py

@@ -105,6 +105,8 @@ def start_training(self):
     if self.config.dataset_type == "grain":
       data_iterator = self.restore_data_iterator_state(data_iterator)
 
+    # don't need this anymore, clear some memory.
+    del pipeline.t5_encoder
 
     flux_state, flux_state_mesh_shardings, flux_learning_rate_scheduler = self.create_flux_state(
         # ambiguous here, but if self.params.get("unet") doesn't exist
@@ -138,7 +140,7 @@ def start_training(self):
     )
     # 6. save final checkpoint
     # Hook
-    self.post_training_steps(pipeline, params, train_states, "after_training")
+    #self.post_training_steps(pipeline, params, train_states, "after_training")
 
   def get_shaped_batch(self, config, pipeline=None):
     """Return the shape of the batch - this is what eval_shape would return for the
@@ -267,7 +269,9 @@ def load_dataset(self, pipeline, params, train_states):
                           clip_tokenizer=pipeline.clip_tokenizer,
                           t5_tokenizer=pipeline.t5_tokenizer,
                           clip_text_encoder=pipeline.clip_encoder,
-                          t5_text_encoder=pipeline.t5_encoder
+                          t5_text_encoder=pipeline.t5_encoder,
+                          encode_in_batches=True,
+                          encode_batch_size=16
                           )
     pack_latents_p = partial(pipeline.pack_latents)
     prepare_latent_image_ids_p = partial(pipeline.prepare_latent_image_ids)