Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

honor upcast_attention model setting #2365

Closed
wants to merge 1 commit into from
Closed

honor upcast_attention model setting #2365

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions ldm/models/diffusion/cross_attention_control.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,15 @@
import math
from typing import Optional, Callable

import diffusers
import psutil
import torch
import diffusers
from torch import nn
from diffusers.models.unet_2d_condition import UNet2DConditionModel
from torch import nn

from ldm.invoke.devices import torch_dtype


# adapted from bloc97's CrossAttentionControl colab
# https://github.com/bloc97/CrossAttentionControl

Expand Down Expand Up @@ -182,6 +184,9 @@ class InvokeAICrossAttentionMixin:
through both to an attention_slice_wrangler and a slicing_strategy_getter for custom attention map wrangling
and dymamic slicing strategy selection.
"""
upcast_attention = False
upcast_softmax = False

def __init__(self):
self.mem_total_gb = psutil.virtual_memory().total // (1 << 30)
self.attention_slice_wrangler = None
Expand Down Expand Up @@ -210,6 +215,10 @@ def set_attention_slice_calculated_callback(self, callback: Optional[Callable[[t
self.attention_slice_calculated_callback = callback

def einsum_lowest_level(self, query, key, value, dim, offset, slice_size):
dtype = query.dtype
if self.upcast_attention:
query = query.float()
key = key.float()
# calculate attention scores
#attention_scores = torch.einsum('b i d, b j d -> b i j', q, k)
attention_scores = torch.baddbmm(
Expand All @@ -221,7 +230,10 @@ def einsum_lowest_level(self, query, key, value, dim, offset, slice_size):
)

# calculate attention slice by taking the best scores for each latent pixel
default_attention_slice = attention_scores.softmax(dim=-1, dtype=attention_scores.dtype)
if self.upcast_softmax:
attention_scores = attention_scores.float()
default_attention_slice = attention_scores.softmax(dim=-1).to(dtype=dtype)

attention_slice_wrangler = self.attention_slice_wrangler
if attention_slice_wrangler is not None:
attention_slice = attention_slice_wrangler(self, default_attention_slice, dim, offset, slice_size)
Expand Down