JIT whole diffusion loop

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1299,7 +1299,7 @@ def step_fn(carry, t):
           new_latents_jax = latents_step
           new_audio_latents_jax = audio_latents_step
 
-        return (new_latents_jax, new_audio_latents_jax), None
+        return (new_latents_jax.astype(latents_jax.dtype), new_audio_latents_jax.astype(audio_latents_jax.dtype)), None
 
       initial_carry = (latents_jax, audio_latents_jax)
       (latents_jax, audio_latents_jax), _ = jax.lax.scan(step_fn, initial_carry, timesteps_jax)