JIT whole diffusion loop

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1238,8 +1238,8 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
 
       import time
       timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
-      for i, t_val in enumerate(timesteps):
-        t = timesteps_jax[i]
+      def step_fn(carry, t):
+        latents_jax, audio_latents_jax = carry
 
         # Isolate input sharding to scan_layers=False to avoid affecting the standard path
         latents_jax_sharded = latents_jax
@@ -1293,11 +1293,16 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
         )
 
         if guidance_scale > 1.0:
-          latents_jax = jnp.concatenate([latents_step] * 2, axis=0)
-          audio_latents_jax = jnp.concatenate([audio_latents_step] * 2, axis=0)
+          new_latents_jax = jnp.concatenate([latents_step] * 2, axis=0)
+          new_audio_latents_jax = jnp.concatenate([audio_latents_step] * 2, axis=0)
         else:
-          latents_jax = latents_step
-          audio_latents_jax = audio_latents_step
+          new_latents_jax = latents_step
+          new_audio_latents_jax = audio_latents_step
+          
+        return (new_latents_jax, new_audio_latents_jax), None
+
+      initial_carry = (latents_jax, audio_latents_jax)
+      (latents_jax, audio_latents_jax), _ = jax.lax.scan(step_fn, initial_carry, timesteps_jax)