[diffusion] perf: improve black-forest-labs/FLUX.2-dev (#14040)

2025-11-27 14:49:52 +08:00
parent 077ca70ee4
commit 6edffc6391
8 changed files with 278 additions and 318 deletions
@@ -148,7 +148,6 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
            )
            self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)

-        # Scaled dot product attention
        self.attn = USPAttention(
            num_heads=num_heads,
            head_size=self.head_dim,