prompt embeds converted to bf16

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -22,7 +22,7 @@ text_encoder_model_name_or_path: "ariG23498/t5-v1-1-xxl-flax"
 prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
 prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
 frame_rate: 30
-max_sequence_length: 512
+max_sequence_length: 1024
 sampler: "from_checkpoint"
 
 # Generation parameters

src/maxdiffusion/generate_ltx2.py

@@ -97,6 +97,8 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
       frame_rate=getattr(config, "fps", 24.0),
       decode_timestep=getattr(config, "decode_timestep", 0.0),
       decode_noise_scale=getattr(config, "decode_noise_scale", None),
+      max_sequence_length=getattr(config, "max_sequence_length", 1024),
+      dtype=jnp.bfloat16 if getattr(config, "activations_dtype", "bfloat16") == "bfloat16" else jnp.float32,
   )
   return out
 

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1099,7 +1099,7 @@ def __call__(
       decode_timestep: Union[float, List[float]] = 0.0,
       decode_noise_scale: Optional[Union[float, List[float]]] = None,
       max_sequence_length: int = 1024,
-      dtype: Optional[jnp.dtype] = jnp.float32,
+      dtype: Optional[jnp.dtype] = None,
       output_type: str = "pil",
       return_dict: bool = True,
   ):