706 lines
30 KiB
Python
706 lines
30 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 numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from diffusers.models.embeddings import Timesteps, apply_rotary_emb
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
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
|
|
from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
|
|
|
|
ANIMA_TRANSFORMER_CONFIG = {
|
|
"in_channels": 16,
|
|
"out_channels": 16,
|
|
"num_attention_heads": 16,
|
|
"attention_head_dim": 128,
|
|
"num_layers": 28,
|
|
"mlp_ratio": 4.0,
|
|
"text_embed_dim": 1024,
|
|
"adaln_lora_dim": 256,
|
|
"max_size": (128, 240, 240),
|
|
"patch_size": (1, 2, 2),
|
|
"rope_scale": (1.0, 4.0, 4.0),
|
|
"concat_padding_mask": True,
|
|
"extra_pos_embed_type": None,
|
|
}
|
|
|
|
|
|
# NOTE: We import and use diffusers' `apply_rotary_emb` instead of a custom native implementation
|
|
# to prevent numerical drift in bfloat16. Diffusers upcasts queries, keys, and rotary frequency
|
|
# tensors to float32 before computing the rotation, and casts back to bfloat16 at the end.
|
|
# Performing the entire computation in bfloat16 accumulates precision errors across the 28
|
|
# transformer blocks, which is heavily amplified by Classifier-Free Guidance (CFG).
|
|
|
|
|
|
class CosmosPatchEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
patch_size: tuple[int, int, int],
|
|
bias: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] * patch_size[2], out_channels, bias=bias)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
|
p_t, p_h, p_w = self.patch_size
|
|
hidden_states = hidden_states.reshape(
|
|
batch_size,
|
|
num_channels,
|
|
num_frames // p_t,
|
|
p_t,
|
|
height // p_h,
|
|
p_h,
|
|
width // p_w,
|
|
p_w,
|
|
)
|
|
hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7)
|
|
return self.proj(hidden_states)
|
|
|
|
|
|
class CosmosTimestepEmbedding(nn.Module):
|
|
def __init__(self, in_features: int, out_features: int) -> None:
|
|
super().__init__()
|
|
self.linear_1 = nn.Linear(in_features, out_features, bias=False)
|
|
self.activation = nn.SiLU()
|
|
self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False)
|
|
|
|
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
|
|
emb = self.linear_1(timesteps)
|
|
emb = self.activation(emb)
|
|
return self.linear_2(emb)
|
|
|
|
|
|
class CosmosEmbedding(nn.Module):
|
|
def __init__(self, embedding_dim: int, condition_dim: int) -> None:
|
|
super().__init__()
|
|
self.time_proj = Timesteps(embedding_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0)
|
|
self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
|
|
self.norm = nn.RMSNorm(embedding_dim, eps=1e-6)
|
|
|
|
def forward(self, hidden_states: torch.Tensor, timestep: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
|
|
temb = self.t_embedder(timesteps_proj)
|
|
embedded_timestep = self.norm(timesteps_proj)
|
|
return temb, embedded_timestep
|
|
|
|
|
|
class CosmosAdaLayerNorm(nn.Module):
|
|
def __init__(self, in_features: int, hidden_features: int) -> None:
|
|
super().__init__()
|
|
self.embedding_dim = in_features
|
|
self.activation = nn.SiLU()
|
|
self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)
|
|
self.linear_1 = nn.Linear(in_features, hidden_features, bias=False)
|
|
self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
embedded_timestep: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
embedded_timestep = self.activation(embedded_timestep)
|
|
embedded_timestep = self.linear_1(embedded_timestep)
|
|
embedded_timestep = self.linear_2(embedded_timestep)
|
|
|
|
if temb is not None:
|
|
embedded_timestep = embedded_timestep + temb[..., : 2 * self.embedding_dim]
|
|
|
|
shift, scale = embedded_timestep.chunk(2, dim=-1)
|
|
hidden_states = self.norm(hidden_states)
|
|
|
|
if embedded_timestep.ndim == 2:
|
|
shift, scale = (x.unsqueeze(1) for x in (shift, scale))
|
|
|
|
return hidden_states * (1 + scale) + shift
|
|
|
|
|
|
class CosmosAdaLayerNormZero(nn.Module):
|
|
def __init__(self, in_features: int, hidden_features: int | None = None) -> None:
|
|
super().__init__()
|
|
self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6)
|
|
self.activation = nn.SiLU()
|
|
self.linear_1 = (
|
|
nn.Identity() if hidden_features is None else nn.Linear(in_features, hidden_features, bias=False)
|
|
)
|
|
self.linear_2 = nn.Linear(hidden_features or in_features, 3 * in_features, bias=False)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
embedded_timestep: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
embedded_timestep = self.activation(embedded_timestep)
|
|
embedded_timestep = self.linear_1(embedded_timestep)
|
|
embedded_timestep = self.linear_2(embedded_timestep)
|
|
|
|
if temb is not None:
|
|
embedded_timestep = embedded_timestep + temb
|
|
|
|
shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
|
|
hidden_states = self.norm(hidden_states)
|
|
|
|
if embedded_timestep.ndim == 2:
|
|
shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))
|
|
|
|
hidden_states = hidden_states * (1 + scale) + shift
|
|
return hidden_states, gate
|
|
|
|
|
|
class CosmosGELU(nn.Module):
|
|
def __init__(self, dim: int, inner_dim: int, bias: bool = False) -> None:
|
|
super().__init__()
|
|
self.proj = nn.Linear(dim, inner_dim, bias=bias)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
return F.gelu(self.proj(hidden_states))
|
|
|
|
|
|
class CosmosFeedForward(nn.Module):
|
|
def __init__(self, dim: int, mult: float = 4.0, bias: bool = False) -> None:
|
|
super().__init__()
|
|
inner_dim = int(dim * mult)
|
|
self.net = nn.ModuleList(
|
|
[
|
|
CosmosGELU(dim, inner_dim, bias=bias),
|
|
nn.Identity(),
|
|
nn.Linear(inner_dim, dim, bias=bias),
|
|
]
|
|
)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
for module in self.net:
|
|
hidden_states = module(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class CosmosAttention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
query_dim: int,
|
|
cross_attention_dim: int | None = None,
|
|
heads: int = 8,
|
|
dim_head: int = 64,
|
|
out_bias: bool = False,
|
|
img_context: bool = False,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
inner_dim = heads * dim_head
|
|
cross_attention_dim = query_dim if cross_attention_dim is None else cross_attention_dim
|
|
|
|
self.heads = heads
|
|
self.dim_head = dim_head
|
|
self.out_dim = inner_dim
|
|
self.query_dim = query_dim
|
|
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
|
|
self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=False)
|
|
self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=False)
|
|
self.norm_q = nn.RMSNorm(dim_head, eps=1e-6)
|
|
self.norm_k = nn.RMSNorm(dim_head, eps=1e-6)
|
|
self.to_out = nn.ModuleList([nn.Linear(inner_dim, query_dim, bias=out_bias), nn.Identity()])
|
|
self.attn = Attention(
|
|
num_heads=heads,
|
|
head_size=dim_head,
|
|
softmax_scale=1.0 / (dim_head**0.5),
|
|
causal=False,
|
|
num_kv_heads=heads,
|
|
prefix=prefix,
|
|
)
|
|
|
|
self.img_context = img_context
|
|
if img_context:
|
|
self.q_img = nn.Linear(query_dim, inner_dim, bias=False)
|
|
self.k_img = nn.Linear(query_dim, inner_dim, bias=False)
|
|
self.v_img = nn.Linear(query_dim, inner_dim, bias=False)
|
|
self.q_img_norm = nn.RMSNorm(dim_head, eps=1e-6)
|
|
self.k_img_norm = nn.RMSNorm(dim_head, eps=1e-6)
|
|
self.img_attn = Attention(
|
|
num_heads=heads,
|
|
head_size=dim_head,
|
|
softmax_scale=1.0 / (dim_head**0.5),
|
|
causal=False,
|
|
num_kv_heads=heads,
|
|
prefix=f"{prefix}.img_attn" if prefix else "img_attn",
|
|
)
|
|
|
|
def _attention(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> torch.Tensor:
|
|
encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states
|
|
batch_size, seq_len, _ = hidden_states.shape
|
|
encoder_seq_len = encoder_hidden_states.shape[1]
|
|
|
|
query = self.to_q(hidden_states).view(batch_size, seq_len, self.heads, self.dim_head)
|
|
key = self.to_k(encoder_hidden_states).view(batch_size, encoder_seq_len, self.heads, self.dim_head)
|
|
value = self.to_v(encoder_hidden_states).view(batch_size, encoder_seq_len, self.heads, self.dim_head)
|
|
|
|
query = self.norm_q(query)
|
|
key = self.norm_k(key)
|
|
|
|
if image_rotary_emb is not None:
|
|
# We use diffusers' apply_rotary_emb to leverage its internal float32 rotation upcasting
|
|
# logic, resolving the bfloat16 cumulative precision drift vs. the reference pipeline.
|
|
query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
|
|
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1, use_real_unbind_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)
|
|
return hidden_states.flatten(2, 3).type_as(query)
|
|
|
|
def _image_attention(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
img_context: torch.Tensor,
|
|
attention_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
batch_size, seq_len, _ = hidden_states.shape
|
|
img_seq_len = img_context.shape[1]
|
|
|
|
query = self.q_img(hidden_states).view(batch_size, seq_len, self.heads, self.dim_head)
|
|
key = self.k_img(img_context).view(batch_size, img_seq_len, self.heads, self.dim_head)
|
|
value = self.v_img(img_context).view(batch_size, img_seq_len, self.heads, self.dim_head)
|
|
|
|
query = self.q_img_norm(query)
|
|
key = self.k_img_norm(key)
|
|
|
|
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
|
|
hidden_states = self.img_attn(query, key, value, attn_metadata=attn_metadata)
|
|
return hidden_states.flatten(2, 3).type_as(query)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None] | None = None,
|
|
attention_mask: torch.Tensor | tuple[torch.Tensor | None, torch.Tensor | None] | None = None,
|
|
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> torch.Tensor:
|
|
if self.img_context:
|
|
if not isinstance(encoder_hidden_states, tuple):
|
|
raise ValueError("Expected encoder_hidden_states as (text_context, img_context).")
|
|
text_context, img_context = encoder_hidden_states
|
|
text_mask, img_mask = attention_mask if attention_mask is not None else (None, None)
|
|
query_states = hidden_states
|
|
hidden_states = self._attention(query_states, text_context, text_mask, image_rotary_emb)
|
|
if img_context is not None:
|
|
hidden_states = hidden_states + self._image_attention(query_states, img_context, img_mask)
|
|
else:
|
|
hidden_states = self._attention(hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb)
|
|
|
|
hidden_states = self.to_out[0](hidden_states)
|
|
return self.to_out[1](hidden_states)
|
|
|
|
|
|
class CosmosTransformerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_attention_heads: int,
|
|
attention_head_dim: int,
|
|
cross_attention_dim: int,
|
|
mlp_ratio: float = 4.0,
|
|
adaln_lora_dim: int = 256,
|
|
out_bias: bool = False,
|
|
img_context: bool = False,
|
|
before_proj: bool = False,
|
|
after_proj: bool = False,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
hidden_size = num_attention_heads * attention_head_dim
|
|
|
|
self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
|
|
self.img_context = img_context
|
|
self.attn1 = CosmosAttention(
|
|
query_dim=hidden_size,
|
|
cross_attention_dim=None,
|
|
heads=num_attention_heads,
|
|
dim_head=attention_head_dim,
|
|
out_bias=out_bias,
|
|
prefix=f"{prefix}.attn1" if prefix else "attn1",
|
|
)
|
|
self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
|
|
self.attn2 = CosmosAttention(
|
|
query_dim=hidden_size,
|
|
cross_attention_dim=cross_attention_dim,
|
|
heads=num_attention_heads,
|
|
dim_head=attention_head_dim,
|
|
out_bias=out_bias,
|
|
img_context=img_context,
|
|
prefix=f"{prefix}.attn2" if prefix else "attn2",
|
|
)
|
|
self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim)
|
|
self.ff = CosmosFeedForward(hidden_size, mult=mlp_ratio, bias=out_bias)
|
|
|
|
self.before_proj = nn.Linear(hidden_size, hidden_size) if before_proj else None
|
|
self.after_proj = nn.Linear(hidden_size, hidden_size) if after_proj else None
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None],
|
|
embedded_timestep: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
extra_pos_emb: torch.Tensor | None = None,
|
|
attention_mask: torch.Tensor | tuple[torch.Tensor | None, torch.Tensor | None] | None = None,
|
|
controlnet_residual: torch.Tensor | None = None,
|
|
latents: torch.Tensor | None = None,
|
|
block_idx: int | None = None,
|
|
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
|
if self.before_proj is not None:
|
|
hidden_states = self.before_proj(hidden_states) + latents
|
|
|
|
if extra_pos_emb is not None:
|
|
hidden_states = hidden_states + extra_pos_emb
|
|
|
|
norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb)
|
|
attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb)
|
|
hidden_states = hidden_states + gate * attn_output
|
|
|
|
norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb)
|
|
attn_output = self.attn2(
|
|
norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
|
|
)
|
|
hidden_states = hidden_states + gate * attn_output
|
|
|
|
norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb)
|
|
ff_output = self.ff(norm_hidden_states)
|
|
hidden_states = hidden_states + gate * ff_output
|
|
|
|
if controlnet_residual is not None:
|
|
if self.after_proj is not None:
|
|
raise ValueError("controlnet_residual is incompatible with after_proj.")
|
|
hidden_states = hidden_states + controlnet_residual
|
|
|
|
if self.after_proj is not None:
|
|
return hidden_states, self.after_proj(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class CosmosRotaryPosEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
max_size: tuple[int, int, int] = (128, 240, 240),
|
|
patch_size: tuple[int, int, int] = (1, 2, 2),
|
|
base_fps: int = 24,
|
|
rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
|
|
) -> None:
|
|
super().__init__()
|
|
self.max_size = [size // patch for size, patch in zip(max_size, patch_size)]
|
|
self.patch_size = patch_size
|
|
self.base_fps = base_fps
|
|
|
|
self.dim_h = hidden_size // 6 * 2
|
|
self.dim_w = hidden_size // 6 * 2
|
|
self.dim_t = hidden_size - self.dim_h - self.dim_w
|
|
|
|
self.h_ntk_factor = rope_scale[1] ** (self.dim_h / (self.dim_h - 2))
|
|
self.w_ntk_factor = rope_scale[2] ** (self.dim_w / (self.dim_w - 2))
|
|
self.t_ntk_factor = rope_scale[0] ** (self.dim_t / (self.dim_t - 2))
|
|
|
|
def forward(self, hidden_states: torch.Tensor, fps: float | None = None) -> tuple[torch.Tensor, torch.Tensor]:
|
|
_, _, num_frames, height, width = hidden_states.shape
|
|
pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
|
|
device = hidden_states.device
|
|
|
|
h_theta = 10000.0 * self.h_ntk_factor
|
|
w_theta = 10000.0 * self.w_ntk_factor
|
|
t_theta = 10000.0 * self.t_ntk_factor
|
|
|
|
seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32)
|
|
dim_h_range = torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[: self.dim_h // 2] / self.dim_h
|
|
dim_w_range = torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[: self.dim_w // 2] / self.dim_w
|
|
dim_t_range = torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[: self.dim_t // 2] / self.dim_t
|
|
h_spatial_freqs = 1.0 / (h_theta**dim_h_range)
|
|
w_spatial_freqs = 1.0 / (w_theta**dim_w_range)
|
|
temporal_freqs = 1.0 / (t_theta**dim_t_range)
|
|
|
|
emb_h = torch.outer(seq[: pe_size[1]], h_spatial_freqs)[None, :, None, :].repeat(pe_size[0], 1, pe_size[2], 1)
|
|
emb_w = torch.outer(seq[: pe_size[2]], w_spatial_freqs)[None, None, :, :].repeat(pe_size[0], pe_size[1], 1, 1)
|
|
if fps is None:
|
|
emb_t = torch.outer(seq[: pe_size[0]], temporal_freqs)
|
|
else:
|
|
emb_t = torch.outer(seq[: pe_size[0]] / fps * self.base_fps, temporal_freqs)
|
|
emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1)
|
|
|
|
freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, 2).float()
|
|
return torch.cos(freqs), torch.sin(freqs)
|
|
|
|
|
|
class CosmosLearnablePositionalEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
max_size: tuple[int, int, int],
|
|
patch_size: tuple[int, int, int],
|
|
eps: float = 1e-6,
|
|
) -> None:
|
|
super().__init__()
|
|
self.max_size = [size // patch for size, patch in zip(max_size, patch_size)]
|
|
self.patch_size = patch_size
|
|
self.eps = eps
|
|
self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0], hidden_size))
|
|
self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1], hidden_size))
|
|
self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2], hidden_size))
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
batch_size, _, num_frames, height, width = hidden_states.shape
|
|
pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]]
|
|
|
|
emb_t = self.pos_emb_t[: pe_size[0]][None, :, None, None, :].repeat(batch_size, 1, pe_size[1], pe_size[2], 1)
|
|
emb_h = self.pos_emb_h[: pe_size[1]][None, None, :, None, :].repeat(batch_size, pe_size[0], 1, pe_size[2], 1)
|
|
emb_w = self.pos_emb_w[: pe_size[2]][None, None, None, :, :].repeat(batch_size, pe_size[0], pe_size[1], 1, 1)
|
|
emb = (emb_t + emb_h + emb_w).flatten(1, 3)
|
|
|
|
norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32)
|
|
norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel()))
|
|
return (emb / norm).type_as(hidden_states)
|
|
|
|
|
|
class AnimaTransformer3DModel(nn.Module):
|
|
_repeated_blocks = ["CosmosTransformerBlock"]
|
|
_hsdp_shard_conditions = [is_transformer_block_module]
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 16,
|
|
out_channels: int = 16,
|
|
num_attention_heads: int = 32,
|
|
attention_head_dim: int = 128,
|
|
num_layers: int = 28,
|
|
mlp_ratio: float = 4.0,
|
|
text_embed_dim: int = 1024,
|
|
adaln_lora_dim: int = 256,
|
|
max_size: tuple[int, int, int] = (128, 240, 240),
|
|
patch_size: tuple[int, int, int] = (1, 2, 2),
|
|
rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0),
|
|
concat_padding_mask: bool = True,
|
|
extra_pos_embed_type: str | None = "learnable",
|
|
use_crossattn_projection: bool = False,
|
|
crossattn_proj_in_channels: int = 1024,
|
|
encoder_hidden_states_channels: int = 1024,
|
|
controlnet_block_every_n: int | None = None,
|
|
img_context_dim_in: int | None = None,
|
|
img_context_num_tokens: int = 256,
|
|
img_context_dim_out: int = 2048,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
super().__init__()
|
|
hidden_size = num_attention_heads * attention_head_dim
|
|
self.config = SimpleNamespace(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=attention_head_dim,
|
|
num_layers=num_layers,
|
|
mlp_ratio=mlp_ratio,
|
|
text_embed_dim=text_embed_dim,
|
|
adaln_lora_dim=adaln_lora_dim,
|
|
max_size=max_size,
|
|
patch_size=patch_size,
|
|
rope_scale=rope_scale,
|
|
concat_padding_mask=concat_padding_mask,
|
|
extra_pos_embed_type=extra_pos_embed_type,
|
|
use_crossattn_projection=use_crossattn_projection,
|
|
crossattn_proj_in_channels=crossattn_proj_in_channels,
|
|
encoder_hidden_states_channels=encoder_hidden_states_channels,
|
|
controlnet_block_every_n=controlnet_block_every_n,
|
|
img_context_dim_in=img_context_dim_in,
|
|
img_context_num_tokens=img_context_num_tokens,
|
|
img_context_dim_out=img_context_dim_out,
|
|
extra_config=kwargs,
|
|
)
|
|
|
|
patch_embed_in_channels = in_channels + 1 if concat_padding_mask else in_channels
|
|
self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, hidden_size, patch_size, bias=False)
|
|
self.rope = CosmosRotaryPosEmbed(
|
|
hidden_size=attention_head_dim,
|
|
max_size=max_size,
|
|
patch_size=patch_size,
|
|
rope_scale=rope_scale,
|
|
)
|
|
self.learnable_pos_embed: CosmosLearnablePositionalEmbed | None = None
|
|
if extra_pos_embed_type == "learnable":
|
|
self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
|
|
hidden_size=hidden_size,
|
|
max_size=max_size,
|
|
patch_size=patch_size,
|
|
)
|
|
self.time_embed = CosmosEmbedding(hidden_size, hidden_size)
|
|
self.transformer_blocks = nn.ModuleList(
|
|
[
|
|
CosmosTransformerBlock(
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=attention_head_dim,
|
|
cross_attention_dim=text_embed_dim,
|
|
mlp_ratio=mlp_ratio,
|
|
adaln_lora_dim=adaln_lora_dim,
|
|
out_bias=False,
|
|
img_context=img_context_dim_in is not None and img_context_dim_in > 0,
|
|
prefix=f"transformer_blocks.{i}",
|
|
)
|
|
for i in range(num_layers)
|
|
]
|
|
)
|
|
self.norm_out = CosmosAdaLayerNorm(hidden_size, adaln_lora_dim)
|
|
self.proj_out = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * patch_size[2] * out_channels, bias=False)
|
|
|
|
self.crossattn_proj: nn.Module | None = None
|
|
if use_crossattn_projection:
|
|
self.crossattn_proj = nn.Sequential(
|
|
nn.Linear(crossattn_proj_in_channels, encoder_hidden_states_channels, bias=True),
|
|
nn.GELU(),
|
|
)
|
|
|
|
self.img_context_proj: nn.Module | None = None
|
|
if img_context_dim_in:
|
|
self.img_context_proj = nn.Sequential(
|
|
nn.Linear(img_context_dim_in, img_context_dim_out, bias=True),
|
|
nn.GELU(),
|
|
)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return next(self.parameters()).dtype
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None],
|
|
block_controlnet_hidden_states: list[torch.Tensor] | None = None,
|
|
attention_mask: torch.Tensor | tuple[torch.Tensor | None, torch.Tensor | None] | None = None,
|
|
fps: float | None = None,
|
|
condition_mask: torch.Tensor | None = None,
|
|
padding_mask: torch.Tensor | None = None,
|
|
return_dict: bool = True,
|
|
) -> Transformer2DModelOutput | tuple[torch.Tensor]:
|
|
batch_size, _, num_frames, height, width = hidden_states.shape
|
|
|
|
if condition_mask is not None:
|
|
hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
|
|
|
|
if self.config.concat_padding_mask:
|
|
if padding_mask is None:
|
|
raise ValueError("padding_mask is required when concat_padding_mask=True.")
|
|
padding_mask_resized = F.interpolate(
|
|
padding_mask,
|
|
size=list(hidden_states.shape[-2:]),
|
|
mode="nearest",
|
|
)
|
|
hidden_states = torch.cat(
|
|
[hidden_states, padding_mask_resized.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)],
|
|
dim=1,
|
|
)
|
|
|
|
if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim != 2:
|
|
attention_mask = attention_mask.squeeze(1).squeeze(1)
|
|
|
|
image_rotary_emb = self.rope(hidden_states, fps=fps)
|
|
extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.learnable_pos_embed is not None else None
|
|
|
|
p_t, p_h, p_w = self.config.patch_size
|
|
post_patch_num_frames = num_frames // p_t
|
|
post_patch_height = height // p_h
|
|
post_patch_width = width // p_w
|
|
|
|
hidden_states = self.patch_embed(hidden_states)
|
|
hidden_states = hidden_states.flatten(1, 3)
|
|
|
|
if timestep.ndim == 1:
|
|
temb, embedded_timestep = self.time_embed(hidden_states, timestep)
|
|
elif timestep.ndim == 5:
|
|
if timestep.shape != (batch_size, 1, num_frames, 1, 1):
|
|
raise ValueError(f"Expected timestep shape [B, 1, T, 1, 1], got {timestep.shape}.")
|
|
timestep = timestep.flatten()
|
|
temb, embedded_timestep = self.time_embed(hidden_states, timestep)
|
|
temb, embedded_timestep = (
|
|
x.view(batch_size, post_patch_num_frames, 1, 1, -1)
|
|
.expand(-1, -1, post_patch_height, post_patch_width, -1)
|
|
.flatten(1, 3)
|
|
for x in (temb, embedded_timestep)
|
|
)
|
|
else:
|
|
raise ValueError(f"Expected timestep shape [B, 1, T, 1, 1] or [T], got {timestep.shape}.")
|
|
|
|
text_context, img_context = (
|
|
encoder_hidden_states if isinstance(encoder_hidden_states, tuple) else (encoder_hidden_states, None)
|
|
)
|
|
if self.config.use_crossattn_projection:
|
|
if self.crossattn_proj is None:
|
|
raise ValueError("crossattn_proj is required when use_crossattn_projection=True.")
|
|
text_context = self.crossattn_proj(text_context)
|
|
if img_context is not None and self.config.img_context_dim_in:
|
|
if self.img_context_proj is None:
|
|
raise ValueError("img_context_proj is required when img_context_dim_in is set.")
|
|
img_context = self.img_context_proj(img_context)
|
|
processed_encoder_hidden_states = (
|
|
(text_context, img_context) if isinstance(encoder_hidden_states, tuple) else text_context
|
|
)
|
|
|
|
controlnet_block_index_map = {}
|
|
if block_controlnet_hidden_states is not None:
|
|
n_blocks = len(self.transformer_blocks)
|
|
controlnet_block_index_map = {
|
|
block_idx: block_controlnet_hidden_states[idx]
|
|
for idx, block_idx in list(enumerate(range(0, n_blocks, self.config.controlnet_block_every_n)))
|
|
}
|
|
|
|
for block_idx, block in enumerate(self.transformer_blocks):
|
|
hidden_states = block(
|
|
hidden_states,
|
|
processed_encoder_hidden_states,
|
|
embedded_timestep,
|
|
temb,
|
|
image_rotary_emb,
|
|
extra_pos_emb,
|
|
attention_mask,
|
|
controlnet_block_index_map.get(block_idx),
|
|
)
|
|
|
|
hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
|
|
hidden_states = self.proj_out(hidden_states)
|
|
hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
|
|
hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
|
|
hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
|
|
hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
|
|
if not return_dict:
|
|
return (hidden_states,)
|
|
|
|
return Transformer2DModelOutput(sample=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
|