duplication corrected in wan2.1 i2v

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -269,30 +269,23 @@ 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_input], axis=-1)
+    latent_model_input = jnp.concatenate([latents_input, condition], axis=-1)
     timestep = jnp.broadcast_to(t, latents_input.shape[0])
     latent_model_input = jnp.transpose(latent_model_input, (0, 4, 1, 2, 3))
 
     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_input,
         do_classifier_free_guidance=do_classifier_free_guidance,
         guidance_scale=guidance_scale,
-        encoder_hidden_states_image=image_embeds_input,
+        encoder_hidden_states_image=image_embeds,
     )
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))