Spaces:
Running on Zero
Running on Zero
File size: 2,405 Bytes
0d8b898 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import torch
from fireredtts3.llm.modules import (
DiTBlock,
FinalLayer,
RotaryEmbedding,
)
class PatchEncoder(torch.nn.Module):
def __init__(
self,
# In & out
in_dim: int,
out_dim: int,
# Model config
patch_size: int = 4,
hidden_size: int = 1024,
mlp_ratio: int = 3,
depth: int = 8,
num_heads: int = 8,
):
super().__init__()
self.in_dim = in_dim
self.patch_size = patch_size
self.hidden_size = hidden_size
self.out_dim = out_dim
# [CLS] token
self.cls_tok = torch.nn.Parameter(torch.zeros(1, 1, hidden_size))
self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
self.blocks = torch.nn.ModuleList([
DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)
])
# Input & output proj
self.in_proj = (
torch.nn.Linear(in_dim, hidden_size)
if in_dim != hidden_size else
torch.nn.Identity()
)
self.out_proj = FinalLayer(hidden_size, out_dim)
def forward(self, inputs_embeds: torch.Tensor):
"""Patch encoder aggregating {patch_size} latents into one.
Args:
inputs_embeds(torch.Tensor): shape (b=1, t, c).
Returns:
hidden_states(torch.Tensor): shape (b=1, t//patch_size, c).
"""
assert inputs_embeds.shape[1] % self.patch_size == 0, \
'inputs_embeds.shape={} patch_size={}'.format(inputs_embeds.shape, self.patch_size)
inputs_embeds = self.in_proj(inputs_embeds)
# Patchify, (b=1, t, c) -> (t//patch_size, patch_size, c)
hidden_states = inputs_embeds.reshape(-1, self.patch_size, self.hidden_size)
cls_tok = self.cls_tok.expand(hidden_states.shape[0], -1, -1) # (b*t//patch_size, 1, c)
hidden_states = torch.cat([cls_tok, hidden_states], dim=1) # (b*t//patch_size, 1+patch_size, c)
# NOTE full attention
rope = self.rotary_embed.forward_from_seq_len(hidden_states.shape[1])
for block in self.blocks:
hidden_states = block(hidden_states, None, rope)
hidden_states = self.out_proj(hidden_states)
hidden_states = hidden_states[:, 0] # (t//patch_size, c)
hidden_states = hidden_states.unsqueeze(0)
return hidden_states
|