# Copyright 2022 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import flax.linen as nn

import jax.numpy as jnp

class FlaxAttentionBlock(nn.Module):

query_dim: int

heads: int = 8

dim_head: int = 64

dropout: float = 0.0

dtype: jnp.dtype = jnp.float32

def setup(self):

inner_dim = self.dim_head * self.heads

self.scale = self.dim_head**-0.5

# Weights were exported with old names {to_q, to_k, to_v, to_out}

self.query = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_q")

self.key = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_k")

self.value = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_v")

self.proj_attn = nn.Dense(self.query_dim, dtype=self.dtype, name="to_out_0")

def reshape_heads_to_batch_dim(self, tensor):

batch_size, seq_len, dim = tensor.shape

head_size = self.heads

tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)

tensor = jnp.transpose(tensor, (0, 2, 1, 3))

tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)

return tensor

def reshape_batch_dim_to_heads(self, tensor):

batch_size, seq_len, dim = tensor.shape

head_size = self.heads

tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)

tensor = jnp.transpose(tensor, (0, 2, 1, 3))

tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size)

return tensor

def __call__(self, hidden_states, context=None, deterministic=True):

context = hidden_states if context is None else context

query_proj = self.query(hidden_states)

key_proj = self.key(context)

value_proj = self.value(context)

query_states = self.reshape_heads_to_batch_dim(query_proj)

key_states = self.reshape_heads_to_batch_dim(key_proj)

value_states = self.reshape_heads_to_batch_dim(value_proj)

# compute attentions

attention_scores = jnp.einsum("b i d, b j d->b i j", query_states, key_states)

attention_scores = attention_scores * self.scale

attention_probs = nn.softmax(attention_scores, axis=2)

# attend to values

hidden_states = jnp.einsum("b i j, b j d -> b i d", attention_probs, value_states)

hidden_states = self.reshape_batch_dim_to_heads(hidden_states)

hidden_states = self.proj_attn(hidden_states)

return hidden_states

class FlaxBasicTransformerBlock(nn.Module):

dim: int

n_heads: int

d_head: int

dropout: float = 0.0

dtype: jnp.dtype = jnp.float32

def setup(self):

# self attention

self.attn1 = FlaxAttentionBlock(self.dim, self.n_heads, self.d_head, self.dropout, dtype=self.dtype)

# cross attention

self.attn2 = FlaxAttentionBlock(self.dim, self.n_heads, self.d_head, self.dropout, dtype=self.dtype)

self.ff = FlaxGluFeedForward(dim=self.dim, dropout=self.dropout, dtype=self.dtype)

self.norm1 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)

self.norm2 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)

self.norm3 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)

def __call__(self, hidden_states, context, deterministic=True):

# self attention

residual = hidden_states

hidden_states = self.attn1(self.norm1(hidden_states), deterministic=deterministic)

hidden_states = hidden_states + residual

# cross attention

residual = hidden_states

hidden_states = self.attn2(self.norm2(hidden_states), context, deterministic=deterministic)

hidden_states = hidden_states + residual

# feed forward

residual = hidden_states

hidden_states = self.ff(self.norm3(hidden_states), deterministic=deterministic)

hidden_states = hidden_states + residual

return hidden_states

class FlaxSpatialTransformer(nn.Module):

in_channels: int

n_heads: int

d_head: int

depth: int = 1

dropout: float = 0.0

dtype: jnp.dtype = jnp.float32

def setup(self):

self.norm = nn.GroupNorm(num_groups=32, epsilon=1e-5)

inner_dim = self.n_heads * self.d_head

self.proj_in = nn.Conv(

inner_dim,

kernel_size=(1, 1),

strides=(1, 1),

padding="VALID",

dtype=self.dtype,

)

self.transformer_blocks = [

FlaxBasicTransformerBlock(inner_dim, self.n_heads, self.d_head, dropout=self.dropout, dtype=self.dtype)

for _ in range(self.depth)

]

self.proj_out = nn.Conv(

inner_dim,

kernel_size=(1, 1),

strides=(1, 1),

padding="VALID",

dtype=self.dtype,

)

def __call__(self, hidden_states, context, deterministic=True):

batch, height, width, channels = hidden_states.shape

residual = hidden_states

hidden_states = self.norm(hidden_states)

hidden_states = self.proj_in(hidden_states)

hidden_states = hidden_states.reshape(batch, height * width, channels)

for transformer_block in self.transformer_blocks:

hidden_states = transformer_block(hidden_states, context, deterministic=deterministic)

hidden_states = hidden_states.reshape(batch, height, width, channels)

hidden_states = self.proj_out(hidden_states)

hidden_states = hidden_states + residual

return hidden_states

class FlaxGluFeedForward(nn.Module):

dim: int

dropout: float = 0.0

dtype: jnp.dtype = jnp.float32

def setup(self):

# The second linear layer needs to be called

# net_2 for now to match the index of the Sequential layer

self.net_0 = FlaxGEGLU(self.dim, self.dropout, self.dtype)

self.net_2 = nn.Dense(self.dim, dtype=self.dtype)

def __call__(self, hidden_states, deterministic=True):

hidden_states = self.net_0(hidden_states)

hidden_states = self.net_2(hidden_states)

return hidden_states

class FlaxGEGLU(nn.Module):

dim: int

dropout: float = 0.0

dtype: jnp.dtype = jnp.float32

def setup(self):

inner_dim = self.dim * 4

self.proj = nn.Dense(inner_dim * 2, dtype=self.dtype)

def __call__(self, hidden_states, deterministic=True):

hidden_states = self.proj(hidden_states)

hidden_linear, hidden_gelu = jnp.split(hidden_states, 2, axis=2)

return hidden_linear * nn.gelu(hidden_gelu)