(
self,
model="ViT-L/14",
jit=False,
device='cuda' if torch.cuda.is_available() else 'cpu',
antialias=True,
ucg_rate=0.
)
| 11 | # Models in ["ViT-B/32", "ViT-B/16", "ViT-L/14", "ViT-L/14@336px"] |
| 12 | class ClipImageEmbedder(nn.Module): |
| 13 | def __init__( |
| 14 | self, |
| 15 | model="ViT-L/14", |
| 16 | jit=False, |
| 17 | device='cuda' if torch.cuda.is_available() else 'cpu', |
| 18 | antialias=True, |
| 19 | ucg_rate=0. |
| 20 | ): |
| 21 | super().__init__() |
| 22 | from clip import load as load_clip |
| 23 | self.model, _ = load_clip(name=model, device=device, jit=jit) |
| 24 | |
| 25 | self.antialias = antialias |
| 26 | |
| 27 | self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) |
| 28 | self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) |
| 29 | self.ucg_rate = ucg_rate |
| 30 | |
| 31 | self.init_uncond() |
| 32 | |
| 33 | def init_uncond(self, path="nulltext.npy"): |
| 34 | try: |
no test coverage detected