mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
246 lines
8.1 KiB
Python
246 lines
8.1 KiB
Python
import math
|
|
|
|
import diffusers
|
|
import torch
|
|
|
|
if torch.backends.mps.is_available():
|
|
torch.empty = torch.zeros
|
|
|
|
|
|
_torch_layer_norm = torch.nn.functional.layer_norm
|
|
|
|
|
|
def new_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
|
|
if input.device.type == "mps" and input.dtype == torch.float16:
|
|
input = input.float()
|
|
if weight is not None:
|
|
weight = weight.float()
|
|
if bias is not None:
|
|
bias = bias.float()
|
|
return _torch_layer_norm(input, normalized_shape, weight, bias, eps).half()
|
|
else:
|
|
return _torch_layer_norm(input, normalized_shape, weight, bias, eps)
|
|
|
|
|
|
torch.nn.functional.layer_norm = new_layer_norm
|
|
|
|
|
|
_torch_tensor_permute = torch.Tensor.permute
|
|
|
|
|
|
def new_torch_tensor_permute(input, *dims):
|
|
result = _torch_tensor_permute(input, *dims)
|
|
if input.device == "mps" and input.dtype == torch.float16:
|
|
result = result.contiguous()
|
|
return result
|
|
|
|
|
|
torch.Tensor.permute = new_torch_tensor_permute
|
|
|
|
|
|
_torch_lerp = torch.lerp
|
|
|
|
|
|
def new_torch_lerp(input, end, weight, *, out=None):
|
|
if input.device.type == "mps" and input.dtype == torch.float16:
|
|
input = input.float()
|
|
end = end.float()
|
|
if isinstance(weight, torch.Tensor):
|
|
weight = weight.float()
|
|
if out is not None:
|
|
out_fp32 = torch.zeros_like(out, dtype=torch.float32)
|
|
else:
|
|
out_fp32 = None
|
|
result = _torch_lerp(input, end, weight, out=out_fp32)
|
|
if out is not None:
|
|
out.copy_(out_fp32.half())
|
|
del out_fp32
|
|
return result.half()
|
|
|
|
else:
|
|
return _torch_lerp(input, end, weight, out=out)
|
|
|
|
|
|
torch.lerp = new_torch_lerp
|
|
|
|
|
|
_torch_interpolate = torch.nn.functional.interpolate
|
|
|
|
|
|
def new_torch_interpolate(
|
|
input,
|
|
size=None,
|
|
scale_factor=None,
|
|
mode="nearest",
|
|
align_corners=None,
|
|
recompute_scale_factor=None,
|
|
antialias=False,
|
|
):
|
|
if input.device.type == "mps" and input.dtype == torch.float16:
|
|
return _torch_interpolate(
|
|
input.float(), size, scale_factor, mode, align_corners, recompute_scale_factor, antialias
|
|
).half()
|
|
else:
|
|
return _torch_interpolate(input, size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
|
|
|
|
|
|
torch.nn.functional.interpolate = new_torch_interpolate
|
|
|
|
# TODO: refactor it
|
|
_SlicedAttnProcessor = diffusers.models.attention_processor.SlicedAttnProcessor
|
|
|
|
|
|
class ChunkedSlicedAttnProcessor:
|
|
r"""
|
|
Processor for implementing sliced attention.
|
|
|
|
Args:
|
|
slice_size (`int`, *optional*):
|
|
The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and
|
|
`attention_head_dim` must be a multiple of the `slice_size`.
|
|
"""
|
|
|
|
def __init__(self, slice_size):
|
|
assert isinstance(slice_size, int)
|
|
slice_size = 1 # TODO: maybe implement chunking in batches too when enough memory
|
|
self.slice_size = slice_size
|
|
self._sliced_attn_processor = _SlicedAttnProcessor(slice_size)
|
|
|
|
def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
|
|
if self.slice_size != 1 or attn.upcast_attention:
|
|
return self._sliced_attn_processor(attn, hidden_states, encoder_hidden_states, attention_mask)
|
|
|
|
residual = hidden_states
|
|
|
|
input_ndim = hidden_states.ndim
|
|
|
|
if input_ndim == 4:
|
|
batch_size, channel, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
|
|
|
batch_size, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
|
|
if attn.group_norm is not None:
|
|
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
|
|
|
query = attn.to_q(hidden_states)
|
|
dim = query.shape[-1]
|
|
query = attn.head_to_batch_dim(query)
|
|
|
|
if encoder_hidden_states is None:
|
|
encoder_hidden_states = hidden_states
|
|
elif attn.norm_cross:
|
|
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
key = attn.head_to_batch_dim(key)
|
|
value = attn.head_to_batch_dim(value)
|
|
|
|
batch_size_attention, query_tokens, _ = query.shape
|
|
hidden_states = torch.zeros(
|
|
(batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
|
|
)
|
|
|
|
chunk_tmp_tensor = torch.empty(
|
|
self.slice_size, query.shape[1], key.shape[1], dtype=query.dtype, device=query.device
|
|
)
|
|
|
|
for i in range(batch_size_attention // self.slice_size):
|
|
start_idx = i * self.slice_size
|
|
end_idx = (i + 1) * self.slice_size
|
|
|
|
query_slice = query[start_idx:end_idx]
|
|
key_slice = key[start_idx:end_idx]
|
|
attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None
|
|
|
|
self.get_attention_scores_chunked(
|
|
attn,
|
|
query_slice,
|
|
key_slice,
|
|
attn_mask_slice,
|
|
hidden_states[start_idx:end_idx],
|
|
value[start_idx:end_idx],
|
|
chunk_tmp_tensor,
|
|
)
|
|
|
|
hidden_states = attn.batch_to_head_dim(hidden_states)
|
|
|
|
# linear proj
|
|
hidden_states = attn.to_out[0](hidden_states)
|
|
# dropout
|
|
hidden_states = attn.to_out[1](hidden_states)
|
|
|
|
if input_ndim == 4:
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
|
|
|
if attn.residual_connection:
|
|
hidden_states = hidden_states + residual
|
|
|
|
hidden_states = hidden_states / attn.rescale_output_factor
|
|
|
|
return hidden_states
|
|
|
|
def get_attention_scores_chunked(self, attn, query, key, attention_mask, hidden_states, value, chunk):
|
|
# batch size = 1
|
|
assert query.shape[0] == 1
|
|
assert key.shape[0] == 1
|
|
assert value.shape[0] == 1
|
|
assert hidden_states.shape[0] == 1
|
|
|
|
# dtype = query.dtype
|
|
if attn.upcast_attention:
|
|
query = query.float()
|
|
key = key.float()
|
|
|
|
# out_item_size = query.dtype.itemsize
|
|
# if attn.upcast_attention:
|
|
# out_item_size = torch.float32.itemsize
|
|
out_item_size = query.element_size()
|
|
if attn.upcast_attention:
|
|
out_item_size = 4
|
|
|
|
chunk_size = 2**29
|
|
|
|
out_size = query.shape[1] * key.shape[1] * out_item_size
|
|
chunks_count = min(query.shape[1], math.ceil((out_size - 1) / chunk_size))
|
|
chunk_step = max(1, int(query.shape[1] / chunks_count))
|
|
|
|
key = key.transpose(-1, -2)
|
|
|
|
def _get_chunk_view(tensor, start, length):
|
|
if start + length > tensor.shape[1]:
|
|
length = tensor.shape[1] - start
|
|
# print(f"view: [{tensor.shape[0]},{tensor.shape[1]},{tensor.shape[2]}] - start: {start}, length: {length}")
|
|
return tensor[:, start : start + length]
|
|
|
|
for chunk_pos in range(0, query.shape[1], chunk_step):
|
|
if attention_mask is not None:
|
|
torch.baddbmm(
|
|
_get_chunk_view(attention_mask, chunk_pos, chunk_step),
|
|
_get_chunk_view(query, chunk_pos, chunk_step),
|
|
key,
|
|
beta=1,
|
|
alpha=attn.scale,
|
|
out=chunk,
|
|
)
|
|
else:
|
|
torch.baddbmm(
|
|
torch.zeros((1, 1, 1), device=query.device, dtype=query.dtype),
|
|
_get_chunk_view(query, chunk_pos, chunk_step),
|
|
key,
|
|
beta=0,
|
|
alpha=attn.scale,
|
|
out=chunk,
|
|
)
|
|
chunk = chunk.softmax(dim=-1)
|
|
torch.bmm(chunk, value, out=_get_chunk_view(hidden_states, chunk_pos, chunk_step))
|
|
|
|
# del chunk
|
|
|
|
|
|
diffusers.models.attention_processor.SlicedAttnProcessor = ChunkedSlicedAttnProcessor
|