diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2023-01-25 16:12:29 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-01-25 16:12:29 +0000 |
| commit | 1574e967297586d013e4cfbb6628eae595c9fba2 (patch) | |
| tree | 99374009f63cf73cadc713b02db5fbba0701e516 /modules/sub_quadratic_attention.py | |
| parent | 1982ef68900fe3c5eee704dfbda5416c1bb5470b (diff) | |
| parent | e3b53fd295aca784253dfc8668ec87b537a72f43 (diff) | |
| download | stable-diffusion-webui-gfx803-1574e967297586d013e4cfbb6628eae595c9fba2.tar.gz stable-diffusion-webui-gfx803-1574e967297586d013e4cfbb6628eae595c9fba2.tar.bz2 stable-diffusion-webui-gfx803-1574e967297586d013e4cfbb6628eae595c9fba2.zip | |
Merge pull request #6510 from brkirch/unet16-upcast-precision
Add upcast options, full precision sampling from float16 UNet and upcasting attention for inference using SD 2.1 models without --no-half
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 |