|
import torch |
|
import torch.nn as nn |
|
|
|
class IDEncoder(nn.Module): |
|
def __init__(self, width=1280, context_dim=2048, num_token=5): |
|
super().__init__() |
|
self.num_token = num_token |
|
self.context_dim = context_dim |
|
h1 = min((context_dim * num_token) // 4, 1024) |
|
h2 = min((context_dim * num_token) // 2, 1024) |
|
self.body = nn.Sequential( |
|
nn.Linear(width, h1), |
|
nn.LayerNorm(h1), |
|
nn.LeakyReLU(), |
|
nn.Linear(h1, h2), |
|
nn.LayerNorm(h2), |
|
nn.LeakyReLU(), |
|
nn.Linear(h2, context_dim * num_token), |
|
) |
|
|
|
for i in range(5): |
|
setattr( |
|
self, |
|
f'mapping_{i}', |
|
nn.Sequential( |
|
nn.Linear(1024, 1024), |
|
nn.LayerNorm(1024), |
|
nn.LeakyReLU(), |
|
nn.Linear(1024, 1024), |
|
nn.LayerNorm(1024), |
|
nn.LeakyReLU(), |
|
nn.Linear(1024, context_dim), |
|
), |
|
) |
|
|
|
setattr( |
|
self, |
|
f'mapping_patch_{i}', |
|
nn.Sequential( |
|
nn.Linear(1024, 1024), |
|
nn.LayerNorm(1024), |
|
nn.LeakyReLU(), |
|
nn.Linear(1024, 1024), |
|
nn.LayerNorm(1024), |
|
nn.LeakyReLU(), |
|
nn.Linear(1024, context_dim), |
|
), |
|
) |
|
|
|
def forward(self, x, y): |
|
|
|
x = self.body(x) |
|
x = x.reshape(-1, self.num_token, self.context_dim) |
|
|
|
hidden_states = () |
|
for i, emb in enumerate(y): |
|
hidden_state = getattr(self, f'mapping_{i}')(emb[:, :1]) + getattr(self, f'mapping_patch_{i}')( |
|
emb[:, 1:] |
|
).mean(dim=1, keepdim=True) |
|
hidden_states += (hidden_state,) |
|
hidden_states = torch.cat(hidden_states, dim=1) |
|
|
|
return torch.cat([x, hidden_states], dim=1) |
|
|
