Uses the OpenCLIP vision transformer encoder for images
| 239 | |
| 240 | |
| 241 | class FrozenOpenCLIPImageEmbedder(AbstractEncoder): |
| 242 | """ |
| 243 | Uses the OpenCLIP vision transformer encoder for images |
| 244 | """ |
| 245 | |
| 246 | def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, |
| 247 | freeze=True, layer="pooled", antialias=True, ucg_rate=0.): |
| 248 | super().__init__() |
| 249 | model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), |
| 250 | pretrained=version, ) |
| 251 | del model.transformer |
| 252 | self.model = model |
| 253 | |
| 254 | self.device = device |
| 255 | self.max_length = max_length |
| 256 | if freeze: |
| 257 | self.freeze() |
| 258 | self.layer = layer |
| 259 | if self.layer == "penultimate": |
| 260 | raise NotImplementedError() |
| 261 | self.layer_idx = 1 |
| 262 | |
| 263 | self.antialias = antialias |
| 264 | |
| 265 | self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) |
| 266 | self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) |
| 267 | self.ucg_rate = ucg_rate |
| 268 | |
| 269 | def preprocess(self, x): |
| 270 | # normalize to [0,1] |
| 271 | x = kornia.geometry.resize(x, (224, 224), |
| 272 | interpolation='bicubic', align_corners=True, |
| 273 | antialias=self.antialias) |
| 274 | x = (x + 1.) / 2. |
| 275 | # renormalize according to clip |
| 276 | x = kornia.enhance.normalize(x, self.mean, self.std) |
| 277 | return x |
| 278 | |
| 279 | def freeze(self): |
| 280 | self.model = self.model.eval() |
| 281 | for param in self.parameters(): |
| 282 | param.requires_grad = False |
| 283 | |
| 284 | @autocast |
| 285 | def forward(self, image, no_dropout=False): |
| 286 | z = self.encode_with_vision_transformer(image) |
| 287 | if self.ucg_rate > 0. and not no_dropout: |
| 288 | z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z |
| 289 | return z |
| 290 | |
| 291 | def encode_with_vision_transformer(self, img): |
| 292 | img = self.preprocess(img) |
| 293 | x = self.model.visual(img) |
| 294 | return x |
| 295 | |
| 296 | def encode(self, text): |
| 297 | return self(text) |
| 298 |
nothing calls this directly
no outgoing calls
no test coverage detected