File size: 1,346 Bytes
2ca14fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """
Frozen CLIP ViT-L/14 text encoder for T.I.T.S.2.
Unlike T.I.T.S.1, which squashed the whole prompt into one 512-dim vector, this returns
the per-token hidden states (77 x 768) so the model can cross-attend word by word, plus
a pooled vector for the global conditioning path, plus the attention mask so padding
tokens are ignored.
"""
from typing import List, Tuple
import torch
import torch.nn as nn
from transformers import CLIPTextModel, CLIPTokenizer
class FrozenCLIPTextEncoder(nn.Module):
def __init__(self, model_name: str = "openai/clip-vit-large-patch14", device: str = "cuda"):
super().__init__()
self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
self.model = CLIPTextModel.from_pretrained(model_name).to(device).eval()
self.model.requires_grad_(False)
self.device = device
self.embed_dim = self.model.config.hidden_size # 768
@torch.no_grad()
def encode(self, prompts: List[str]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""-> (tokens (b,77,768), pooled (b,768), mask (b,77))"""
tok = self.tokenizer(
prompts, padding="max_length", truncation=True, max_length=77, return_tensors="pt"
).to(self.device)
out = self.model(**tok)
return out.last_hidden_state, out.pooler_output, tok.attention_mask
|