Force float32 for latents in run_diffusion_loop to test precision effect

prishajain1 · prishajain1 · commit e7aceb39302d · 2026-04-19T14:32:58.000+05:30
diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py
@@ -1577,6 +1577,8 @@ def run_diffusion_loop(
     scheduler_step,
     logical_axis_rules,
 ):
+  latents_jax = latents_jax.astype(jnp.float32)
+  audio_latents_jax = audio_latents_jax.astype(jnp.float32)
   transformer = nnx.merge(graphdef, state)
 
   def scan_body(carry, t, model):
