Spaces:
Running on Zero
Running on Zero
File size: 2,100 Bytes
ccf7a99 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """FLAN-T5-Large prompt encoder used by MiniT2I."""
from __future__ import annotations
from typing import Optional
import torch
_T5_NAME = "google/flan-t5-large"
_SINGLETON: dict[tuple, "T5PromptEncoder"] = {}
class T5PromptEncoder:
def __init__(
self,
model_name: str = _T5_NAME,
device: str | torch.device = "cuda",
dtype: torch.dtype = torch.float32,
max_len: int = 256,
cache_dir: Optional[str] = None,
):
from transformers import AutoTokenizer, T5EncoderModel
self.max_len = max_len
self.device = torch.device(device)
self.dtype = dtype
self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
self.model = (
T5EncoderModel.from_pretrained(model_name, cache_dir=cache_dir)
.to(self.device, dtype=dtype)
.eval()
)
for p in self.model.parameters():
p.requires_grad_(False)
@torch.no_grad()
def encode(self, prompts: list[str], max_len: Optional[int] = None):
max_len = int(max_len or self.max_len)
tok = self.tokenizer(
list(prompts),
padding="max_length",
truncation=True,
max_length=max_len,
return_tensors="pt",
)
input_ids = tok.input_ids.to(self.device)
attn_mask = tok.attention_mask.to(self.device)
out = self.model(input_ids=input_ids, attention_mask=attn_mask)
context = out.last_hidden_state.to(self.dtype)
return context, attn_mask.to(self.dtype)
def get_prompt_encoder(
model_name: str = _T5_NAME,
device: str | torch.device = "cuda",
dtype: torch.dtype = torch.float32,
max_len: int = 256,
cache_dir: Optional[str] = None,
) -> T5PromptEncoder:
"""Cached accessor so repeated calls reuse a single loaded T5 encoder."""
key = (model_name, str(device), dtype, max_len, cache_dir)
if key not in _SINGLETON:
_SINGLETON[key] = T5PromptEncoder(model_name, device, dtype, max_len, cache_dir)
return _SINGLETON[key]
|