automatic/pipelines/vibe/edit_head.py

"""Meta Connector."""

import torch
from torch import nn
from torch.nn import functional as func


class Expert(nn.Module):
    """Position-wise Feed-Forward Networks.

    This consists of two linear transformations with a ReLU activation in between.
    """

    def __init__(self, d_model: int, d_expert: int) -> None:
        """Initialize the Expert layer.

        Args:
            d_model (int): The dimension of the model.
            d_expert (int): The dimension of the expert.
        """
        super().__init__()
        self.d_model = d_model
        self.d_expert = d_expert

        # Linear transformation y = xW+b
        self.fc1 = nn.Linear(self.d_model, self.d_expert, bias=True)
        self.fc2 = nn.Linear(self.d_expert, self.d_model, bias=True)
        self.fc3 = nn.Linear(self.d_model, self.d_expert, bias=True)

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize the weights of the Expert layer."""
        # init weights
        nn.init.xavier_uniform_(self.fc1.weight)
        nn.init.zeros_(self.fc1.bias)
        nn.init.xavier_uniform_(self.fc2.weight)
        nn.init.zeros_(self.fc2.bias)
        nn.init.xavier_uniform_(self.fc3.weight)
        nn.init.zeros_(self.fc3.bias)

    def forward(self, embedding: torch.Tensor) -> torch.Tensor:
        """Forward pass of the Expert layer.

        Args:
            embedding (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The output tensor.
        """
        inner_emb = func.silu(self.fc1(embedding)) * self.fc3(embedding)
        return self.fc2(inner_emb)


class MetaAttentionBlock(nn.Module):
    """Meta Attention Block."""

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.05) -> None:
        """Initialize the Meta Attention Block.

        Args:
            d_model (int): The dimension of the embeddings.
            n_heads (int): The number of attention heads.
            dropout (float): The dropout rate.
        """
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)

        # Dropout layers for residual connections
        self.dropout_attn = nn.Dropout(dropout)
        self.dropout_ffn = nn.Dropout(dropout)

        self.mlp = Expert(d_model, d_model * 4)

        self.input_norm = nn.LayerNorm(d_model)
        self.output_norm = nn.LayerNorm(d_model)

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights for attention and MLP blocks."""
        # Initialize MultiheadAttention weights
        # The in_proj_weight contains Q, K, V concatenated
        nn.init.xavier_uniform_(self.attn.in_proj_weight)
        if self.attn.in_proj_bias is not None:
            nn.init.zeros_(self.attn.in_proj_bias)

        # Initialize the output projection of MultiheadAttention
        nn.init.xavier_uniform_(self.attn.out_proj.weight)
        if self.attn.out_proj.bias is not None:
            nn.init.zeros_(self.attn.out_proj.bias)

        nn.init.ones_(self.input_norm.weight)
        nn.init.zeros_(self.input_norm.bias)
        nn.init.ones_(self.output_norm.weight)
        nn.init.zeros_(self.output_norm.bias)

    def forward(self, embedding: torch.Tensor) -> torch.Tensor:
        """Forward pass of the MetaAttentionBlock.

        Args:
            embedding (torch.Tensor): The input embeddings. Shape: (batch_size, seq_len, hidden_dim).

        Returns:
            torch.Tensor: The output embeddings. Shape: (batch_size, seq_len, hidden_dim).
        """
        # apply layer normalization to embeddings
        normed_embedding = self.input_norm(embedding)

        # apply attention and skip connection
        attn_output, _ = self.attn(normed_embedding, normed_embedding, normed_embedding, need_weights=False)
        embedding = embedding + self.dropout_attn(attn_output)

        output = self.mlp(self.output_norm(embedding))
        return embedding + self.dropout_ffn(output)


class MetaConnector(nn.Module):
    """Meta Connector."""

    def __init__(
        self,
        input_dim: int = 4096,
        hidden_dim: int = 2048,
        output_dim: int = 768,
        n_stacks: int = 8,
        heads: int = 16,
        *,
        use_input_proj: bool = True,
    ) -> None:
        """Initialize the MetaConnector.

        Args:
            input_dim (int): The dimension of the input.
            hidden_dim (int): The dimension of the hidden state.
            output_dim (int): The dimension of the output.
            n_stacks (int): The number of attention layers.
            heads (int): The number of attention heads.
            use_input_proj (bool): Whether to use input projection.
        """
        super().__init__()
        self.output_dim = output_dim
        self.input_dim = input_dim
        self.use_input_proj = use_input_proj

        # Input Projection
        if use_input_proj:
            self.input_proj = nn.Linear(input_dim, hidden_dim)
        else:
            hidden_dim = input_dim

        # Attention blocks
        self.attn_blocks = nn.Sequential(*[MetaAttentionBlock(hidden_dim, heads) for _ in range(n_stacks)])

        # Output Projection
        self.output_proj = nn.Sequential(nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, output_dim))

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights."""
        # Initialize input_proj
        if self.use_input_proj:
            nn.init.xavier_uniform_(self.input_proj.weight)
            if self.input_proj.bias is not None:
                nn.init.zeros_(self.input_proj.bias)

        # Initialize output_proj
        for module in self.output_proj:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

    def forward(self, hidden_states: torch.Tensor, meta_queries: torch.Tensor) -> dict[str, torch.Tensor]:
        """Forward pass of the Meta Connector.

        Args:
            hidden_states (torch.Tensor): The input hidden states. Shape: (batch_size, seq_len, hidden_dim).
            meta_queries (torch.Tensor): The input meta queries. Shape: (batch_size, seq_len, hidden_dim).

        Returns:
            torch.Tensor: The output embeddings. Shape: (batch_size, seq_len, hidden_dim).
        """
        hidden_states = hidden_states + meta_queries
        if self.use_input_proj:
            hidden_states = self.input_proj(hidden_states)

        # apply attention blocks
        for attn_block in self.attn_blocks:
            hidden_states = attn_block(hidden_states)

        # apply output projection
        return self.output_proj(hidden_states)