|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from typing import Any, Dict, Optional |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from torch import nn |
|
|
|
from .attention import Attention |
|
|
|
|
|
class BasicTransformerBlock(nn.Module): |
|
r""" |
|
A basic Transformer block. |
|
|
|
Parameters: |
|
dim (`int`): The number of channels in the input and output. |
|
num_attention_heads (`int`): The number of heads to use for multi-head attention. |
|
attention_head_dim (`int`): The number of channels in each head. |
|
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. |
|
cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. |
|
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. |
|
num_embeds_ada_norm (: |
|
obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. |
|
attention_bias (: |
|
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. |
|
only_cross_attention (`bool`, *optional*): |
|
Whether to use only cross-attention layers. In this case two cross attention layers are used. |
|
double_self_attention (`bool`, *optional*): |
|
Whether to use two self-attention layers. In this case no cross attention layers are used. |
|
upcast_attention (`bool`, *optional*): |
|
Whether to upcast the attention computation to float32. This is useful for mixed precision training. |
|
norm_elementwise_affine (`bool`, *optional*, defaults to `True`): |
|
Whether to use learnable elementwise affine parameters for normalization. |
|
norm_type (`str`, *optional*, defaults to `"layer_norm"`): |
|
The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. |
|
final_dropout (`bool` *optional*, defaults to False): |
|
Whether to apply a final dropout after the last feed-forward layer. |
|
attention_type (`str`, *optional*, defaults to `"default"`): |
|
The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dim: int, |
|
num_attention_heads: int, |
|
attention_head_dim: int, |
|
dropout=0.0, |
|
cross_attention_dim: Optional[int] = None, |
|
activation_fn: str = "geglu", |
|
attention_bias: bool = False, |
|
only_cross_attention: bool = False, |
|
double_self_attention: bool = False, |
|
upcast_attention: bool = False, |
|
norm_elementwise_affine: bool = True, |
|
norm_type: str = "layer_norm", |
|
final_dropout: bool = False, |
|
): |
|
super().__init__() |
|
self.only_cross_attention = only_cross_attention |
|
|
|
assert norm_type == "layer_norm" |
|
|
|
|
|
|
|
self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) |
|
self.attn1 = Attention( |
|
query_dim=dim, |
|
heads=num_attention_heads, |
|
dim_head=attention_head_dim, |
|
dropout=dropout, |
|
bias=attention_bias, |
|
cross_attention_dim=cross_attention_dim if only_cross_attention else None, |
|
upcast_attention=upcast_attention, |
|
) |
|
|
|
|
|
if cross_attention_dim is not None or double_self_attention: |
|
|
|
|
|
|
|
self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) |
|
|
|
self.attn2 = Attention( |
|
query_dim=dim, |
|
cross_attention_dim=cross_attention_dim |
|
if not double_self_attention |
|
else None, |
|
heads=num_attention_heads, |
|
dim_head=attention_head_dim, |
|
dropout=dropout, |
|
bias=attention_bias, |
|
upcast_attention=upcast_attention, |
|
) |
|
else: |
|
self.norm2 = None |
|
self.attn2 = None |
|
|
|
|
|
self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) |
|
self.ff = FeedForward( |
|
dim, |
|
dropout=dropout, |
|
activation_fn=activation_fn, |
|
final_dropout=final_dropout, |
|
) |
|
|
|
|
|
self._chunk_size = None |
|
self._chunk_dim = 0 |
|
|
|
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): |
|
|
|
self._chunk_size = chunk_size |
|
self._chunk_dim = dim |
|
|
|
def forward( |
|
self, |
|
hidden_states: torch.FloatTensor, |
|
attention_mask: Optional[torch.FloatTensor] = None, |
|
encoder_hidden_states: Optional[torch.FloatTensor] = None, |
|
encoder_attention_mask: Optional[torch.FloatTensor] = None, |
|
) -> torch.FloatTensor: |
|
|
|
|
|
norm_hidden_states = self.norm1(hidden_states) |
|
|
|
attn_output = self.attn1( |
|
norm_hidden_states, |
|
encoder_hidden_states=encoder_hidden_states |
|
if self.only_cross_attention |
|
else None, |
|
attention_mask=attention_mask, |
|
) |
|
|
|
hidden_states = attn_output + hidden_states |
|
|
|
|
|
if self.attn2 is not None: |
|
norm_hidden_states = self.norm2(hidden_states) |
|
|
|
attn_output = self.attn2( |
|
norm_hidden_states, |
|
encoder_hidden_states=encoder_hidden_states, |
|
attention_mask=encoder_attention_mask, |
|
) |
|
hidden_states = attn_output + hidden_states |
|
|
|
|
|
norm_hidden_states = self.norm3(hidden_states) |
|
|
|
if self._chunk_size is not None: |
|
|
|
if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: |
|
raise ValueError( |
|
f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." |
|
) |
|
|
|
num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size |
|
ff_output = torch.cat( |
|
[ |
|
self.ff(hid_slice) |
|
for hid_slice in norm_hidden_states.chunk( |
|
num_chunks, dim=self._chunk_dim |
|
) |
|
], |
|
dim=self._chunk_dim, |
|
) |
|
else: |
|
ff_output = self.ff(norm_hidden_states) |
|
|
|
hidden_states = ff_output + hidden_states |
|
|
|
return hidden_states |
|
|
|
|
|
class FeedForward(nn.Module): |
|
r""" |
|
A feed-forward layer. |
|
|
|
Parameters: |
|
dim (`int`): The number of channels in the input. |
|
dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. |
|
mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. |
|
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. |
|
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. |
|
final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dim: int, |
|
dim_out: Optional[int] = None, |
|
mult: int = 4, |
|
dropout: float = 0.0, |
|
activation_fn: str = "geglu", |
|
final_dropout: bool = False, |
|
): |
|
super().__init__() |
|
inner_dim = int(dim * mult) |
|
dim_out = dim_out if dim_out is not None else dim |
|
linear_cls = nn.Linear |
|
|
|
if activation_fn == "gelu": |
|
act_fn = GELU(dim, inner_dim) |
|
if activation_fn == "gelu-approximate": |
|
act_fn = GELU(dim, inner_dim, approximate="tanh") |
|
elif activation_fn == "geglu": |
|
act_fn = GEGLU(dim, inner_dim) |
|
elif activation_fn == "geglu-approximate": |
|
act_fn = ApproximateGELU(dim, inner_dim) |
|
|
|
self.net = nn.ModuleList([]) |
|
|
|
self.net.append(act_fn) |
|
|
|
self.net.append(nn.Dropout(dropout)) |
|
|
|
self.net.append(linear_cls(inner_dim, dim_out)) |
|
|
|
if final_dropout: |
|
self.net.append(nn.Dropout(dropout)) |
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
|
for module in self.net: |
|
hidden_states = module(hidden_states) |
|
return hidden_states |
|
|
|
|
|
class GELU(nn.Module): |
|
r""" |
|
GELU activation function with tanh approximation support with `approximate="tanh"`. |
|
|
|
Parameters: |
|
dim_in (`int`): The number of channels in the input. |
|
dim_out (`int`): The number of channels in the output. |
|
approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. |
|
""" |
|
|
|
def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"): |
|
super().__init__() |
|
self.proj = nn.Linear(dim_in, dim_out) |
|
self.approximate = approximate |
|
|
|
def gelu(self, gate: torch.Tensor) -> torch.Tensor: |
|
if gate.device.type != "mps": |
|
return F.gelu(gate, approximate=self.approximate) |
|
|
|
return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to( |
|
dtype=gate.dtype |
|
) |
|
|
|
def forward(self, hidden_states): |
|
hidden_states = self.proj(hidden_states) |
|
hidden_states = self.gelu(hidden_states) |
|
return hidden_states |
|
|
|
|
|
class GEGLU(nn.Module): |
|
r""" |
|
A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202. |
|
|
|
Parameters: |
|
dim_in (`int`): The number of channels in the input. |
|
dim_out (`int`): The number of channels in the output. |
|
""" |
|
|
|
def __init__(self, dim_in: int, dim_out: int): |
|
super().__init__() |
|
linear_cls = nn.Linear |
|
|
|
self.proj = linear_cls(dim_in, dim_out * 2) |
|
|
|
def gelu(self, gate: torch.Tensor) -> torch.Tensor: |
|
if gate.device.type != "mps": |
|
return F.gelu(gate) |
|
|
|
return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype) |
|
|
|
def forward(self, hidden_states, scale: float = 1.0): |
|
args = () |
|
hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1) |
|
return hidden_states * self.gelu(gate) |
|
|
|
|
|
class ApproximateGELU(nn.Module): |
|
r""" |
|
The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2: |
|
https://arxiv.org/abs/1606.08415. |
|
|
|
Parameters: |
|
dim_in (`int`): The number of channels in the input. |
|
dim_out (`int`): The number of channels in the output. |
|
""" |
|
|
|
def __init__(self, dim_in: int, dim_out: int): |
|
super().__init__() |
|
self.proj = nn.Linear(dim_in, dim_out) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
x = self.proj(x) |
|
return x * torch.sigmoid(1.702 * x) |
|
|