Trying text_mask 11

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -273,6 +273,7 @@ def run_inference_2_2_i2v(
             encoder_attention_mask = jnp.concatenate([text_attention_mask, negative_text_attention_mask], axis=0)
         else:
             encoder_attention_mask = None
+        # WAN 2.2 I2V doesn't use CLIP image embeddings, so image_embeds may be None
         if image_embeds is not None:
             image_embeds = jnp.concatenate([image_embeds, image_embeds], axis=0)
         condition = jnp.concatenate([condition] * 2)
@@ -301,13 +302,6 @@ def low_noise_branch(operands):
         )
         return noise_pred, latents_out
 
-    if do_classifier_free_guidance:
-        prompt_embeds = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
-        # WAN 2.2 I2V: image_embeds may be None since it doesn't use CLIP image encoder
-        if image_embeds is not None:
-            image_embeds = jnp.concatenate([image_embeds, image_embeds], axis=0)
-        condition = jnp.concatenate([condition] * 2)
-
     for step in range(num_inference_steps):
         t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
         latents_input = latents