322 lines
13 KiB
Python
322 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from collections.abc import Iterable
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.layer import Attention
|
|
|
|
ANIMA_TEXT_CONDITIONER_CONFIG = {
|
|
"source_dim": 1024,
|
|
"target_dim": 1024,
|
|
"model_dim": 1024,
|
|
"num_layers": 6,
|
|
"num_attention_heads": 16,
|
|
"mlp_ratio": 4.0,
|
|
"target_vocab_size": 32128,
|
|
"use_self_attention": True,
|
|
"use_layer_norm": False,
|
|
"min_sequence_length": 512,
|
|
}
|
|
|
|
|
|
def _rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states_1 = hidden_states[..., : hidden_states.shape[-1] // 2]
|
|
hidden_states_2 = hidden_states[..., hidden_states.shape[-1] // 2 :]
|
|
return torch.cat((-hidden_states_2, hidden_states_1), dim=-1)
|
|
|
|
|
|
def _apply_rotary_pos_emb(
|
|
hidden_states: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
unsqueeze_dim: int = 1,
|
|
) -> torch.Tensor:
|
|
cos = cos.unsqueeze(unsqueeze_dim)
|
|
sin = sin.unsqueeze(unsqueeze_dim)
|
|
return (hidden_states * cos) + (_rotate_half(hidden_states) * sin)
|
|
|
|
|
|
class AnimaRotaryEmbedding(nn.Module):
|
|
def __init__(self, head_dim: int, rope_theta: float = 10000.0) -> None:
|
|
super().__init__()
|
|
inv_freq = 1.0 / (
|
|
rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float32) / head_dim)
|
|
)
|
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
|
|
def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
|
|
inv_freq_expanded = inv_freq_expanded.to(hidden_states.device)
|
|
position_ids_expanded = position_ids[:, None, :].float()
|
|
|
|
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|
cos = emb.cos()
|
|
sin = emb.sin()
|
|
return cos.to(dtype=hidden_states.dtype), sin.to(dtype=hidden_states.dtype)
|
|
|
|
|
|
class AnimaTextConditionerAttention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
query_dim: int,
|
|
context_dim: int,
|
|
num_attention_heads: int,
|
|
attention_head_dim: int,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
inner_dim = num_attention_heads * attention_head_dim
|
|
|
|
self.num_attention_heads = num_attention_heads
|
|
self.attention_head_dim = attention_head_dim
|
|
self.q_proj = nn.Linear(query_dim, inner_dim, bias=False)
|
|
self.q_norm = nn.RMSNorm(attention_head_dim, eps=1e-6)
|
|
self.k_proj = nn.Linear(context_dim, inner_dim, bias=False)
|
|
self.k_norm = nn.RMSNorm(attention_head_dim, eps=1e-6)
|
|
self.v_proj = nn.Linear(context_dim, inner_dim, bias=False)
|
|
self.o_proj = nn.Linear(inner_dim, query_dim, bias=False)
|
|
self.attn = Attention(
|
|
num_heads=num_attention_heads,
|
|
head_size=attention_head_dim,
|
|
softmax_scale=1.0 / (attention_head_dim**0.5),
|
|
causal=False,
|
|
num_kv_heads=num_attention_heads,
|
|
prefix=prefix,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attention_mask: torch.Tensor | None = None,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
encoder_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> torch.Tensor:
|
|
encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states
|
|
input_shape = hidden_states.shape[:-1]
|
|
encoder_input_shape = encoder_hidden_states.shape[:-1]
|
|
|
|
query = self.q_proj(hidden_states)
|
|
key = self.k_proj(encoder_hidden_states)
|
|
value = self.v_proj(encoder_hidden_states)
|
|
|
|
query = query.view(*input_shape, self.num_attention_heads, self.attention_head_dim)
|
|
key = key.view(*encoder_input_shape, self.num_attention_heads, self.attention_head_dim)
|
|
value = value.view(*encoder_input_shape, self.num_attention_heads, self.attention_head_dim)
|
|
|
|
query = self.q_norm(query)
|
|
key = self.k_norm(key)
|
|
|
|
if position_embeddings is not None:
|
|
if encoder_position_embeddings is None:
|
|
raise ValueError("encoder_position_embeddings is required with rotary embeddings.")
|
|
cos, sin = position_embeddings
|
|
query = _apply_rotary_pos_emb(query, cos, sin, unsqueeze_dim=2)
|
|
cos, sin = encoder_position_embeddings
|
|
key = _apply_rotary_pos_emb(key, cos, sin, unsqueeze_dim=2)
|
|
|
|
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
|
|
hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
|
|
hidden_states = hidden_states.flatten(2, 3).contiguous()
|
|
return self.o_proj(hidden_states)
|
|
|
|
|
|
class AnimaTextConditionerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
source_dim: int,
|
|
model_dim: int,
|
|
num_attention_heads: int = 16,
|
|
mlp_ratio: float = 4.0,
|
|
use_self_attention: bool = True,
|
|
use_layer_norm: bool = False,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.use_self_attention = use_self_attention
|
|
norm_cls = nn.LayerNorm if use_layer_norm else nn.RMSNorm
|
|
norm_kwargs = {} if use_layer_norm else {"eps": 1e-6}
|
|
|
|
if use_self_attention:
|
|
self.norm_self_attn = norm_cls(model_dim, **norm_kwargs)
|
|
self.self_attn = AnimaTextConditionerAttention(
|
|
query_dim=model_dim,
|
|
context_dim=model_dim,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=model_dim // num_attention_heads,
|
|
prefix=f"{prefix}.self_attn" if prefix else "self_attn",
|
|
)
|
|
|
|
self.norm_cross_attn = norm_cls(model_dim, **norm_kwargs)
|
|
self.cross_attn = AnimaTextConditionerAttention(
|
|
query_dim=model_dim,
|
|
context_dim=source_dim,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=model_dim // num_attention_heads,
|
|
prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
|
|
)
|
|
self.norm_mlp = norm_cls(model_dim, **norm_kwargs)
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(model_dim, int(model_dim * mlp_ratio)),
|
|
nn.GELU(),
|
|
nn.Linear(int(model_dim * mlp_ratio), model_dim),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
target_attention_mask: torch.Tensor | None = None,
|
|
source_attention_mask: torch.Tensor | None = None,
|
|
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
source_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> torch.Tensor:
|
|
if self.use_self_attention:
|
|
norm_hidden_states = self.norm_self_attn(hidden_states)
|
|
attn_hidden_states = self.self_attn(
|
|
norm_hidden_states,
|
|
attention_mask=target_attention_mask,
|
|
position_embeddings=position_embeddings,
|
|
encoder_position_embeddings=position_embeddings,
|
|
)
|
|
hidden_states = hidden_states + attn_hidden_states
|
|
|
|
norm_hidden_states = self.norm_cross_attn(hidden_states)
|
|
attn_hidden_states = self.cross_attn(
|
|
norm_hidden_states,
|
|
attention_mask=source_attention_mask,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
position_embeddings=position_embeddings,
|
|
encoder_position_embeddings=source_position_embeddings,
|
|
)
|
|
hidden_states = hidden_states + attn_hidden_states
|
|
hidden_states = hidden_states + self.mlp(self.norm_mlp(hidden_states))
|
|
return hidden_states
|
|
|
|
|
|
def is_conditioner_block_module(name: str, module: nn.Module) -> bool:
|
|
return "blocks" in name and name.split(".")[-1].isdigit()
|
|
|
|
|
|
class AnimaTextConditioner(nn.Module):
|
|
_repeated_blocks = ["AnimaTextConditionerBlock"]
|
|
_hsdp_shard_conditions = [is_conditioner_block_module]
|
|
|
|
def __init__(
|
|
self,
|
|
source_dim: int = 1024,
|
|
target_dim: int = 1024,
|
|
model_dim: int = 1024,
|
|
num_layers: int = 6,
|
|
num_attention_heads: int = 16,
|
|
mlp_ratio: float = 4.0,
|
|
target_vocab_size: int = 32128,
|
|
use_self_attention: bool = True,
|
|
use_layer_norm: bool = False,
|
|
min_sequence_length: int = 512,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
super().__init__()
|
|
self.config = SimpleNamespace(
|
|
source_dim=source_dim,
|
|
target_dim=target_dim,
|
|
model_dim=model_dim,
|
|
num_layers=num_layers,
|
|
num_attention_heads=num_attention_heads,
|
|
mlp_ratio=mlp_ratio,
|
|
target_vocab_size=target_vocab_size,
|
|
use_self_attention=use_self_attention,
|
|
use_layer_norm=use_layer_norm,
|
|
min_sequence_length=min_sequence_length,
|
|
extra_config=kwargs,
|
|
)
|
|
self.embed = nn.Embedding(target_vocab_size, target_dim)
|
|
self.in_proj = nn.Linear(target_dim, model_dim) if model_dim != target_dim else nn.Identity()
|
|
self.rotary_emb = AnimaRotaryEmbedding(model_dim // num_attention_heads)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
AnimaTextConditionerBlock(
|
|
source_dim=source_dim,
|
|
model_dim=model_dim,
|
|
num_attention_heads=num_attention_heads,
|
|
mlp_ratio=mlp_ratio,
|
|
use_self_attention=use_self_attention,
|
|
use_layer_norm=use_layer_norm,
|
|
prefix=f"blocks.{i}",
|
|
)
|
|
for i in range(num_layers)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(model_dim, target_dim)
|
|
self.norm = nn.RMSNorm(target_dim, eps=1e-6)
|
|
self.gradient_checkpointing = False
|
|
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return next(self.parameters()).dtype
|
|
|
|
@staticmethod
|
|
def _prepare_attention_mask(attention_mask: torch.Tensor | None) -> torch.Tensor | None:
|
|
if attention_mask is None:
|
|
return None
|
|
return attention_mask.to(torch.bool)
|
|
|
|
def forward(
|
|
self,
|
|
source_hidden_states: torch.Tensor,
|
|
target_input_ids: torch.Tensor,
|
|
target_attention_mask: torch.Tensor | None = None,
|
|
source_attention_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
target_attention_mask = self._prepare_attention_mask(target_attention_mask)
|
|
source_attention_mask = self._prepare_attention_mask(source_attention_mask)
|
|
|
|
hidden_states = self.embed(target_input_ids).to(dtype=source_hidden_states.dtype)
|
|
hidden_states = self.in_proj(hidden_states)
|
|
|
|
position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
|
|
source_position_ids = torch.arange(source_hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
|
|
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
|
source_position_embeddings = self.rotary_emb(hidden_states, source_position_ids)
|
|
|
|
for block in self.blocks:
|
|
hidden_states = block(
|
|
hidden_states,
|
|
source_hidden_states,
|
|
target_attention_mask=target_attention_mask,
|
|
source_attention_mask=source_attention_mask,
|
|
position_embeddings=position_embeddings,
|
|
source_position_embeddings=source_position_embeddings,
|
|
)
|
|
|
|
hidden_states = self.norm(self.out_proj(hidden_states))
|
|
|
|
if target_attention_mask is not None:
|
|
hidden_states = hidden_states * target_attention_mask.to(hidden_states).unsqueeze(-1)
|
|
|
|
if hidden_states.shape[1] < self.config.min_sequence_length:
|
|
hidden_states = F.pad(hidden_states, (0, 0, 0, self.config.min_sequence_length - hidden_states.shape[1]))
|
|
|
|
return hidden_states
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
params = dict(self.named_parameters())
|
|
loaded = set()
|
|
for name, tensor in weights:
|
|
if name not in params:
|
|
continue
|
|
param = params[name]
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
weight_loader(param, tensor)
|
|
loaded.add(name)
|
|
return loaded
|