| 5 | |
| 6 | |
| 7 | class MLPProjModel(torch.nn.Module): |
| 8 | def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4): |
| 9 | super().__init__() |
| 10 | |
| 11 | self.cross_attention_dim = cross_attention_dim |
| 12 | self.num_tokens = num_tokens |
| 13 | |
| 14 | self.proj = torch.nn.Sequential( |
| 15 | torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2), |
| 16 | torch.nn.GELU(), |
| 17 | torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens), |
| 18 | ) |
| 19 | self.norm = torch.nn.LayerNorm(cross_attention_dim) |
| 20 | |
| 21 | def forward(self, id_embeds): |
| 22 | x = self.proj(id_embeds) |
| 23 | x = x.reshape(-1, self.num_tokens, self.cross_attention_dim) |
| 24 | x = self.norm(x) |
| 25 | return x |
| 26 | |
| 27 | class IpAdapterModule(torch.nn.Module): |
| 28 | def __init__(self, num_attention_heads, attention_head_dim, input_dim): |