support for wan2.1 in run_inference added

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -728,6 +728,11 @@ def high_noise_branch(operands):
       latents = jnp.concatenate([latents] * 2)
     timestep = jnp.broadcast_to(t, latents.shape[0])
 
+    if model_name == "wan2.1":
+      noise_pred, latents = low_noise_branch((latents, timestep, prompt_embeds))
+      latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
+      continue
+
     use_high_noise = jnp.greater_equal(t, boundary)
 
     noise_pred, latents = jax.lax.cond(