| 88 | } |
| 89 | |
| 90 | class CLIPLayer(torch.nn.Module): |
| 91 | def __init__(self, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device): |
| 92 | super().__init__() |
| 93 | self.layer_norm1 = nn.LayerNorm(embed_dim, dtype=dtype, device=device) |
| 94 | self.self_attn = CLIPAttention(embed_dim, heads, dtype, device) |
| 95 | self.layer_norm2 = nn.LayerNorm(embed_dim, dtype=dtype, device=device) |
| 96 | #self.mlp = CLIPMLP(embed_dim, intermediate_size, intermediate_activation, dtype, device) |
| 97 | self.mlp = Mlp(embed_dim, intermediate_size, embed_dim, act_layer=ACTIVATIONS[intermediate_activation], dtype=dtype, device=device) |
| 98 | |
| 99 | def forward(self, x, mask=None): |
| 100 | x += self.self_attn(self.layer_norm1(x), mask) |
| 101 | x += self.mlp(self.layer_norm2(x)) |
| 102 | return x |
| 103 | |
| 104 | |
| 105 | class CLIPEncoder(torch.nn.Module): |