| """ | |
| Frozen CLIP text encoder for prompt conditioning. We reuse a pretrained encoder | |
| rather than training our own — the "terrible" in T.I.T.S. lives entirely in the | |
| image model, not the text understanding, otherwise prompts would just be noise | |
| and you'd never be able to tell if a bad image was a bad *rendering* of your | |
| prompt or the model just not understanding the prompt at all. | |
| """ | |
| from typing import List | |
| import torch | |
| import torch.nn as nn | |
| from transformers import CLIPTokenizer, CLIPTextModelWithProjection | |
| class FrozenCLIPTextEncoder(nn.Module): | |
| def __init__(self, model_name: str = "openai/clip-vit-base-patch32", device: str = "cuda"): | |
| super().__init__() | |
| self.tokenizer = CLIPTokenizer.from_pretrained(model_name) | |
| self.model = CLIPTextModelWithProjection.from_pretrained(model_name) | |
| self.model.to(device) | |
| self.model.eval() | |
| for p in self.model.parameters(): | |
| p.requires_grad = False | |
| self.device = device | |
| self.embed_dim = self.model.config.projection_dim # 512 for base-patch32 | |
| def encode(self, prompts: List[str]) -> torch.Tensor: | |
| tokens = self.tokenizer( | |
| prompts, padding=True, truncation=True, max_length=77, return_tensors="pt" | |
| ).to(self.device) | |
| out = self.model(**tokens) | |
| return out.text_embeds # (batch, embed_dim), already pooled/projected | |