219 lines
7.1 KiB
Python
219 lines
7.1 KiB
Python
from collections import OrderedDict
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
|
from transformers import T5EncoderModel, T5Tokenizer
|
|
|
|
|
|
class AdaLayerNorm(nn.Module):
|
|
def __init__(self, embedding_dim: int, time_embedding_dim: Optional[int] = None):
|
|
super().__init__()
|
|
|
|
if time_embedding_dim is None:
|
|
time_embedding_dim = embedding_dim
|
|
|
|
self.silu = nn.SiLU()
|
|
self.linear = nn.Linear(time_embedding_dim, 2 * embedding_dim, bias=True)
|
|
nn.init.zeros_(self.linear.weight)
|
|
nn.init.zeros_(self.linear.bias)
|
|
|
|
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
|
|
|
|
def forward(
|
|
self, x: torch.Tensor, timestep_embedding: torch.Tensor
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
emb = self.linear(self.silu(timestep_embedding))
|
|
shift, scale = emb.view(len(x), 1, -1).chunk(2, dim=-1)
|
|
x = self.norm(x) * (1 + scale) + shift
|
|
return x
|
|
|
|
|
|
class SquaredReLU(nn.Module):
|
|
def forward(self, x: torch.Tensor):
|
|
return torch.square(torch.relu(x))
|
|
|
|
|
|
class PerceiverAttentionBlock(nn.Module):
|
|
def __init__(
|
|
self, d_model: int, n_heads: int, time_embedding_dim: Optional[int] = None
|
|
):
|
|
super().__init__()
|
|
self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
|
|
|
|
self.mlp = nn.Sequential(
|
|
OrderedDict(
|
|
[
|
|
("c_fc", nn.Linear(d_model, d_model * 4)),
|
|
("sq_relu", SquaredReLU()),
|
|
("c_proj", nn.Linear(d_model * 4, d_model)),
|
|
]
|
|
)
|
|
)
|
|
|
|
self.ln_1 = AdaLayerNorm(d_model, time_embedding_dim)
|
|
self.ln_2 = AdaLayerNorm(d_model, time_embedding_dim)
|
|
self.ln_ff = AdaLayerNorm(d_model, time_embedding_dim)
|
|
|
|
def attention(self, q: torch.Tensor, kv: torch.Tensor):
|
|
attn_output, attn_output_weights = self.attn(q, kv, kv, need_weights=False)
|
|
return attn_output
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
latents: torch.Tensor,
|
|
timestep_embedding: torch.Tensor = None,
|
|
):
|
|
normed_latents = self.ln_1(latents, timestep_embedding)
|
|
latents = latents + self.attention(
|
|
q=normed_latents,
|
|
kv=torch.cat([normed_latents, self.ln_2(x, timestep_embedding)], dim=1),
|
|
)
|
|
latents = latents + self.mlp(self.ln_ff(latents, timestep_embedding))
|
|
return latents
|
|
|
|
|
|
class PerceiverResampler(nn.Module):
|
|
def __init__(
|
|
self,
|
|
width: int = 768,
|
|
layers: int = 6,
|
|
heads: int = 8,
|
|
num_latents: int = 64,
|
|
output_dim=None,
|
|
input_dim=None,
|
|
time_embedding_dim: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
self.output_dim = output_dim
|
|
self.input_dim = input_dim
|
|
self.latents = nn.Parameter(width**-0.5 * torch.randn(num_latents, width))
|
|
self.time_aware_linear = nn.Linear(
|
|
time_embedding_dim or width, width, bias=True
|
|
)
|
|
|
|
if self.input_dim is not None:
|
|
self.proj_in = nn.Linear(input_dim, width)
|
|
|
|
self.perceiver_blocks = nn.Sequential(
|
|
*[
|
|
PerceiverAttentionBlock(
|
|
width, heads, time_embedding_dim=time_embedding_dim
|
|
)
|
|
for _ in range(layers)
|
|
]
|
|
)
|
|
|
|
if self.output_dim is not None:
|
|
self.proj_out = nn.Sequential(
|
|
nn.Linear(width, output_dim), nn.LayerNorm(output_dim)
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, timestep_embedding: torch.Tensor = None):
|
|
learnable_latents = self.latents.unsqueeze(dim=0).repeat(len(x), 1, 1)
|
|
latents = learnable_latents + self.time_aware_linear(
|
|
torch.nn.functional.silu(timestep_embedding)
|
|
)
|
|
if self.input_dim is not None:
|
|
x = self.proj_in(x)
|
|
for p_block in self.perceiver_blocks:
|
|
latents = p_block(x, latents, timestep_embedding=timestep_embedding)
|
|
|
|
if self.output_dim is not None:
|
|
latents = self.proj_out(latents)
|
|
|
|
return latents
|
|
|
|
|
|
class T5TextEmbedder(nn.Module):
|
|
def __init__(self, pretrained_path="ybelkada/flan-t5-xl-sharded-bf16", max_length=None):
|
|
super().__init__()
|
|
self.model = T5EncoderModel.from_pretrained(pretrained_path)
|
|
self.tokenizer = T5Tokenizer.from_pretrained(pretrained_path)
|
|
self.max_length = max_length
|
|
|
|
def forward(
|
|
self, caption, text_input_ids=None, attention_mask=None, max_length=None
|
|
):
|
|
if max_length is None:
|
|
max_length = self.max_length
|
|
|
|
if text_input_ids is None or attention_mask is None:
|
|
if max_length is not None:
|
|
text_inputs = self.tokenizer(
|
|
caption,
|
|
return_tensors="pt",
|
|
add_special_tokens=True,
|
|
max_length=max_length,
|
|
padding="max_length",
|
|
truncation=True,
|
|
)
|
|
else:
|
|
text_inputs = self.tokenizer(
|
|
caption, return_tensors="pt", add_special_tokens=True
|
|
)
|
|
text_input_ids = text_inputs.input_ids
|
|
attention_mask = text_inputs.attention_mask
|
|
text_input_ids = text_input_ids.to(self.model.device)
|
|
attention_mask = attention_mask.to(self.model.device)
|
|
outputs = self.model(text_input_ids, attention_mask=attention_mask)
|
|
|
|
embeddings = outputs.last_hidden_state
|
|
return embeddings
|
|
|
|
|
|
class ELLA(nn.Module):
|
|
def __init__(
|
|
self,
|
|
time_channel=320,
|
|
time_embed_dim=768,
|
|
act_fn: str = "silu",
|
|
out_dim: Optional[int] = None,
|
|
width=768,
|
|
layers=6,
|
|
heads=8,
|
|
num_latents=64,
|
|
input_dim=2048,
|
|
):
|
|
super().__init__()
|
|
|
|
self.position = Timesteps(
|
|
time_channel, flip_sin_to_cos=True, downscale_freq_shift=0
|
|
)
|
|
self.time_embedding = TimestepEmbedding(
|
|
in_channels=time_channel,
|
|
time_embed_dim=time_embed_dim,
|
|
act_fn=act_fn,
|
|
out_dim=out_dim,
|
|
)
|
|
|
|
self.connector = PerceiverResampler(
|
|
width=width,
|
|
layers=layers,
|
|
heads=heads,
|
|
num_latents=num_latents,
|
|
input_dim=input_dim,
|
|
time_embedding_dim=time_embed_dim,
|
|
)
|
|
|
|
def forward(self, text_encode_features, timesteps):
|
|
device = text_encode_features.device
|
|
dtype = text_encode_features.dtype
|
|
|
|
ori_time_feature = self.position(timesteps.view(-1)).to(device, dtype=dtype)
|
|
ori_time_feature = (
|
|
ori_time_feature.unsqueeze(dim=1)
|
|
if ori_time_feature.ndim == 2
|
|
else ori_time_feature
|
|
)
|
|
ori_time_feature = ori_time_feature.expand(len(text_encode_features), -1, -1)
|
|
time_embedding = self.time_embedding(ori_time_feature)
|
|
|
|
encoder_hidden_states = self.connector(
|
|
text_encode_features, timestep_embedding=time_embedding
|
|
)
|
|
|
|
return encoder_hidden_states
|