annotation for denoising loop

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1292,7 +1292,7 @@ def step_fn(carry, t):
           new_latents_jax = latents_step
           new_audio_latents_jax = audio_latents_step
 
-        return (new_latents_jax.astype(latents_jax.dtype), new_audio_latents_jax.astype(audio_latents_jax.dtype)), None
+        return (new_latents_jax, new_audio_latents_jax), None
 
       if not self.transformer.scan_layers:
         activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
@@ -1301,10 +1301,13 @@ def step_fn(carry, t):
         video_embeds = jax.lax.with_sharding_constraint(video_embeds, activation_axis_names)
         audio_embeds = jax.lax.with_sharding_constraint(audio_embeds, activation_axis_names)
 
-      initial_carry = (latents_jax, audio_latents_jax)
+      initial_carry = (latents_jax.astype(jnp.float32), audio_latents_jax.astype(jnp.float32))
       with jax.named_scope("denoising_loop"):
         (latents_jax, audio_latents_jax), _ = jax.lax.scan(step_fn, initial_carry, timesteps_jax)
 
+      latents_jax = latents_jax.astype(jnp.bfloat16)
+      audio_latents_jax = audio_latents_jax.astype(jnp.bfloat16)
+
 
 
     # 8. Decode Latents