| """CLIP text encoder for AeroGen. Uses transformers only (no ldm).""" |
|
|
| import torch.nn as nn |
| from transformers import CLIPTokenizer, CLIPTextModel |
|
|
|
|
| class AeroGenCLIPTextEncoder(nn.Module): |
| """CLIP text encoder compatible with FrozenCLIPEmbedder interface. |
| Uses transformers CLIPTextModel + CLIPTokenizer. No ldm dependency. |
| """ |
|
|
| def __init__(self, version: str = "openai/clip-vit-large-patch14", device: str = "cuda", max_length: int = 77): |
| super().__init__() |
| self.tokenizer = CLIPTokenizer.from_pretrained(version) |
| self.transformer = CLIPTextModel.from_pretrained(version) |
| self.device = device |
| self.max_length = max_length |
| self.freeze() |
|
|
| def freeze(self): |
| self.transformer = self.transformer.eval() |
| for param in self.parameters(): |
| param.requires_grad = False |
|
|
| def forward(self, text): |
| if isinstance(text, str): |
| text = [text] |
| batch_encoding = self.tokenizer( |
| text, |
| truncation=True, |
| max_length=self.max_length, |
| return_length=True, |
| return_overflowing_tokens=False, |
| padding="max_length", |
| return_tensors="pt", |
| ) |
| device = next(self.parameters()).device |
| tokens = batch_encoding["input_ids"].to(device) |
| outputs = self.transformer(input_ids=tokens) |
| return outputs.last_hidden_state |
|
|
| def encode(self, text): |
| return self(text) |
|
|
