import added in wan_checkpointer_test.py

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -1132,11 +1132,6 @@ def __call__(
         encoder_attention_mask_img = jnp.ones((encoder_hidden_states_img.shape[0], padded_img_len), dtype=jnp.int32)
         if image_seq_len_actual < padded_img_len:
              encoder_attention_mask_img = encoder_attention_mask_img.at[:, image_seq_len_actual:].set(0)
-
-        # Extract image portion of attention mask (includes padded tokens)
-        # encoder_attention_mask_img = None
-        # if encoder_attention_mask is not None:
-        #   encoder_attention_mask_img = encoder_attention_mask[:, :padded_img_len]
       else:
         # If no image_seq_len is specified, treat all as text
         encoder_hidden_states_img = None

src/maxdiffusion/tests/wan_checkpointer_test.py

@@ -17,6 +17,7 @@
 from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2
 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p1 import WanCheckpointerI2V_2_1
 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p2 import WanCheckpointerI2V_2_2
+from maxdiffusion.pipelines.wan.wan_pipeline_i2v_2p1 import WanPipelineI2V_2_1
 
 class WanCheckpointer2_1Test(unittest.TestCase):
   """Tests for WAN 2.1 checkpointer."""