| 103 | |
| 104 | |
| 105 | class CLIPEncoder(torch.nn.Module): |
| 106 | def __init__(self, num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device): |
| 107 | super().__init__() |
| 108 | self.layers = torch.nn.ModuleList([CLIPLayer(embed_dim, heads, intermediate_size, intermediate_activation, dtype, device) for i in range(num_layers)]) |
| 109 | |
| 110 | def forward(self, x, mask=None, intermediate_output=None): |
| 111 | if intermediate_output is not None: |
| 112 | if intermediate_output < 0: |
| 113 | intermediate_output = len(self.layers) + intermediate_output |
| 114 | intermediate = None |
| 115 | for i, l in enumerate(self.layers): |
| 116 | x = l(x, mask) |
| 117 | if i == intermediate_output: |
| 118 | intermediate = x.clone() |
| 119 | return x, intermediate |
| 120 | |
| 121 | |
| 122 | class CLIPEmbeddings(torch.nn.Module): |