Allow compiling crossattn

I don't remember why I even disabled this... seems to work, big effect with NAG
2025-06-12 19:04:15 +03:00
parent f6df3d75ef
commit e7c39757f6
1 changed files with 1 additions and 1 deletions
@@ -643,7 +643,7 @@ class WanAttentionBlock(nn.Module):
                                    num_latent_frames=num_latent_frames, nag_params=nag_params, nag_context=nag_context)
        del e
        return x
-    @torch.compiler.disable()
+    #@torch.compiler.disable()
    def cross_attn_ffn(self, x, context, context_lens, e, clip_embed=None, grid_sizes=None, 
                       audio_proj=None, audio_context_lens=None, audio_scale=1.0, num_latent_frames=21, nag_params={}, nag_context=None):
            x = x + self.cross_attn(self.norm3(x), context, context_lens, clip_embed=clip_embed,