Skip to content

Commit

Permalink
fix SanaMultiscaleLinearAttention apply_quadratic_attention bf16
Browse files Browse the repository at this point in the history
  • Loading branch information
chenjy2003 committed Jan 16, 2025
1 parent 477937e commit bba83a4
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/diffusers/models/attention_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -899,7 +899,7 @@ def apply_quadratic_attention(self, query: torch.Tensor, key: torch.Tensor, valu
scores = torch.matmul(key.transpose(-1, -2), query)
scores = scores.to(dtype=torch.float32)
scores = scores / (torch.sum(scores, dim=2, keepdim=True) + self.eps)
hidden_states = torch.matmul(value, scores)
hidden_states = torch.matmul(value, scores.to(value.dtype))
return hidden_states

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
Expand Down

0 comments on commit bba83a4

Please sign in to comment.