| 94 | |
| 95 | |
| 96 | class FrozenOpenCLIPImageEmbedder(nn.Module): |
| 97 | def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, |
| 98 | freeze=True, layer="pooled", antialias=True, ucg_rate=0.): |
| 99 | super().__init__() |
| 100 | model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), |
| 101 | pretrained=version, ) |
| 102 | del model.transformer |
| 103 | self.model = model |
| 104 | |
| 105 | self.device = device |
| 106 | self.max_length = max_length |
| 107 | if freeze: |
| 108 | self.freeze() |
| 109 | self.layer = layer |
| 110 | if self.layer == "penultimate": |
| 111 | raise NotImplementedError() |
| 112 | self.layer_idx = 1 |
| 113 | |
| 114 | self.antialias = antialias |
| 115 | |
| 116 | self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) |
| 117 | self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) |
| 118 | self.ucg_rate = ucg_rate |
| 119 | |
| 120 | def preprocess(self, x): |
| 121 | # resize to 224, normalize to [0,1] and re-normalize according to clip |
| 122 | x = torch.nn.functional.interpolate(x, size=(224, 224), mode='bilinear', align_corners=False) |
| 123 | assert x.min() >= -1. and x.max() <= 1. |
| 124 | x = (x + 1.) / 2. |
| 125 | x = (x - self.mean.reshape(1,-1,1,1)) / self.std.reshape(1,-1,1,1) |
| 126 | return x |
| 127 | |
| 128 | def freeze(self): |
| 129 | self.model = self.model.eval() |
| 130 | for param in self.parameters(): |
| 131 | param.requires_grad = False |
| 132 | |
| 133 | def forward(self, image, no_dropout=False): |
| 134 | z = self.encode_with_vision_transformer(image) |
| 135 | if self.ucg_rate > 0. and not no_dropout: |
| 136 | z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z |
| 137 | return z.unsqueeze(1) |
| 138 | |
| 139 | def encode_with_vision_transformer(self, img): |
| 140 | img = self.preprocess(img) |
| 141 | x = self.model.visual(img) |
| 142 | return x |
| 143 | |
| 144 | def encode(self, text): |
| 145 | return self(text) |
| 146 | |
| 147 | |
| 148 | class FrozenOpenCLIPEmbedder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected