Cast scheduler output back to input dtype in scan body

prishajain1 · prishajain1 · commit b59c0943d8d5 · 2026-04-17T23:21:45.000+05:30
diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py
@@ -1641,9 +1641,12 @@ def scan_body(carry, t, model):
             latents_step, _ = scheduler_step(
                 s_state, noise_pred, t, latents_step, return_dict=False
             )
+            latents_step = latents_step.astype(latents.dtype)
+
             audio_latents_step, _ = scheduler_step(
                 s_state, noise_pred_audio, t, audio_latents_step, return_dict=False
             )
+            audio_latents_step = audio_latents_step.astype(audio_latents.dtype)
 
             if guidance_scale > 1.0:
                 latents_next = jnp.concatenate([latents_step] * 2, axis=0)

Original file line number	Diff line number	Diff line change
`@@ -1641,9 +1641,12 @@ def scan_body(carry, t, model):`
`1641`	`1641`	`latents_step, _ = scheduler_step(`
`1642`	`1642`	`s_state, noise_pred, t, latents_step, return_dict=False`
`1643`	`1643`	`)`
	`1644`	`+ latents_step = latents_step.astype(latents.dtype)`
	`1645`	`+`
`1644`	`1646`	`audio_latents_step, _ = scheduler_step(`
`1645`	`1647`	`s_state, noise_pred_audio, t, audio_latents_step, return_dict=False`
`1646`	`1648`	`)`
	`1649`	`+ audio_latents_step = audio_latents_step.astype(audio_latents.dtype)`
`1647`	`1650`
`1648`	`1651`	`if guidance_scale > 1.0:`
`1649`	`1652`	`latents_next = jnp.concatenate([latents_step] * 2, axis=0)`