timestep error fix wan 2.1

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -269,20 +269,30 @@ def loop_body(step, vals):
 
     latents_input = latents
     if do_classifier_free_guidance:
+        condition_input = jnp.concatenate([condition, condition], axis=0)
         latents_input = jnp.concatenate([latents, latents], axis=0)
+    else:
+        condition_input = condition
 
-    latent_model_input = jnp.concatenate([latents_input, condition], axis=-1)
+    latent_model_input = jnp.concatenate([latents_input, condition_input], axis=-1)
     timestep = jnp.broadcast_to(t, latents_input.shape[0])
     latent_model_input = jnp.transpose(latent_model_input, (0, 4, 1, 2, 3))
-    timestep = jnp.broadcast_to(t, latents.shape[0])  
+
+    prompt_embeds_input = prompt_embeds
+    image_embeds_input = image_embeds
+    if do_classifier_free_guidance:
+        prompt_embeds_input = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
+        if image_embeds is not None:
+             image_embeds_input = jnp.concatenate([image_embeds, image_embeds], axis=0)
+
 
 
     noise_pred, latents = transformer_forward_pass(
         graphdef, sharded_state, rest_of_state,
-        latent_model_input, timestep, prompt_embeds,
+        latent_model_input, timestep, prompt_embeds_input,
         do_classifier_free_guidance=do_classifier_free_guidance,
         guidance_scale=guidance_scale,
-        encoder_hidden_states_image=image_embeds,
+        encoder_hidden_states_image=image_embeds_input,
     )
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))