diff options
| author | InvincibleDude <81354513+InvincibleDude@users.noreply.github.com> | 2023-01-29 11:36:10 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-01-29 11:36:10 +0000 |
| commit | ee3d63b6beb88e63542976a1095d4c8aa97388bd (patch) | |
| tree | cf891130682c343107ae0a0f8cec309aea16807a /modules/sub_quadratic_attention.py | |
| parent | 44c0e6b993d00bb2f441f0fde409bcb79136f034 (diff) | |
| parent | 00dab8f10defbbda579a1bc89c8d4e972c58a20d (diff) | |
| download | stable-diffusion-webui-gfx803-ee3d63b6beb88e63542976a1095d4c8aa97388bd.tar.gz stable-diffusion-webui-gfx803-ee3d63b6beb88e63542976a1095d4c8aa97388bd.tar.bz2 stable-diffusion-webui-gfx803-ee3d63b6beb88e63542976a1095d4c8aa97388bd.zip | |
Merge branch 'master' into master
Diffstat (limited to 'modules/sub_quadratic_attention.py')
| -rw-r--r-- | modules/sub_quadratic_attention.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/modules/sub_quadratic_attention.py b/modules/sub_quadratic_attention.py index 55052815..05595323 100644 --- a/modules/sub_quadratic_attention.py +++ b/modules/sub_quadratic_attention.py @@ -67,7 +67,7 @@ def _summarize_chunk( max_score, _ = torch.max(attn_weights, -1, keepdim=True) max_score = max_score.detach() exp_weights = torch.exp(attn_weights - max_score) - exp_values = torch.bmm(exp_weights, value) + exp_values = torch.bmm(exp_weights, value) if query.device.type == 'mps' else torch.bmm(exp_weights, value.to(exp_weights.dtype)).to(value.dtype) max_score = max_score.squeeze(-1) return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score) @@ -129,7 +129,7 @@ def _get_attention_scores_no_kv_chunking( ) attn_probs = attn_scores.softmax(dim=-1) del attn_scores - hidden_states_slice = torch.bmm(attn_probs, value) + hidden_states_slice = torch.bmm(attn_probs, value) if query.device.type == 'mps' else torch.bmm(attn_probs, value.to(attn_probs.dtype)).to(value.dtype) return hidden_states_slice |