Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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] | |