| import os |
| import torch |
| import torch.nn as nn |
| from transformers import CLIPTokenizer, CLIPTextModel |
|
|
|
|
| class FrozenCLIPTextEncoder(nn.Module): |
| def __init__(self, model_path_or_id: str, device: str = "cuda"): |
| super().__init__() |
| self.device = device |
| |
| if os.path.isdir(model_path_or_id): |
| assert os.path.exists(os.path.join(model_path_or_id, "config.json")), \ |
| f"Error: Local path '{model_path_or_id}' does not contain 'config.json'. Check your weight path." |
| |
| self.tokenizer = CLIPTokenizer.from_pretrained(model_path_or_id, local_files_only=True) |
| self.text_encoder = CLIPTextModel.from_pretrained(model_path_or_id, local_files_only=True).to(self.device) |
| else: |
| self.tokenizer = CLIPTokenizer.from_pretrained(model_path_or_id) |
| self.text_encoder = CLIPTextModel.from_pretrained(model_path_or_id).to(self.device) |
| |
| self.text_encoder.eval() |
| for param in self.text_encoder.parameters(): |
| param.requires_grad = False |
|
|
| @torch.no_grad() |
| def forward(self, text: str or list[str]) -> torch.Tensor: |
| if isinstance(text, str): |
| text = [text] |
|
|
| inputs = self.tokenizer( |
| text, |
| padding="max_length", |
| truncation=True, |
| max_length=77, |
| return_tensors="pt", |
| ).to(self.device) |
|
|
| outputs = self.text_encoder(**inputs) |
|
|
| text_tokens = outputs.last_hidden_state |
| text_tokens = text_tokens / text_tokens.norm(dim=-1, keepdim=True).clamp_min(1e-6) |
|
|
| return text_tokens |
