removed attn_mask from FlaxWanAttn call

prishajain1 · prishajain1 · commit 8ef6a0da44f4 · 2026-01-10T20:27:22.000+05:30
diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py
@@ -408,7 +408,6 @@ def __call__(
           attn_output = self.attn2(
               hidden_states=norm_hidden_states,
               encoder_hidden_states=encoder_hidden_states,
-              encoder_attention_mask=encoder_attention_mask,
               deterministic=deterministic,
               rngs=rngs,
           )

Original file line number	Diff line number	Diff line change
`@@ -408,7 +408,6 @@ def __call__(`
`408`	`408`	`attn_output = self.attn2(`
`409`	`409`	`hidden_states=norm_hidden_states,`
`410`	`410`	`encoder_hidden_states=encoder_hidden_states,`
`411`		`- encoder_attention_mask=encoder_attention_mask,`
`412`	`411`	`deterministic=deterministic,`
`413`	`412`	`rngs=rngs,`
`414`	`413`	`)`