| 93 | "hidden" |
| 94 | ] |
| 95 | def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77, |
| 96 | freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32 |
| 97 | super().__init__() |
| 98 | assert layer in self.LAYERS |
| 99 | self.tokenizer = CLIPTokenizer.from_pretrained(version) |
| 100 | self.transformer = CLIPTextModel.from_pretrained(version) |
| 101 | self.device = device |
| 102 | self.max_length = max_length |
| 103 | if freeze: |
| 104 | self.freeze() |
| 105 | self.layer = layer |
| 106 | self.layer_idx = layer_idx |
| 107 | if layer == "hidden": |
| 108 | assert layer_idx is not None |
| 109 | assert 0 <= abs(layer_idx) <= 12 |
| 110 | |
| 111 | def freeze(self): |
| 112 | self.transformer = self.transformer.eval() |