| |
| |
| |
| import torch |
| import importlib |
| import torch.nn as nn |
|
|
| from inspect import isfunction |
|
|
| from transformers import ( |
| AutoTokenizer, |
| T5EncoderModel, |
| ) |
|
|
|
|
| def exists(x): |
| return x is not None |
|
|
|
|
| def default(val, d): |
| if exists(val): |
| return val |
| return d() if isfunction(d) else d |
|
|
|
|
| def instantiate_from_config(config, reload=False): |
| if not "target" in config: |
| if config == "__is_first_stage__": |
| return None |
| elif config == "__is_unconditional__": |
| return None |
| raise KeyError("Expected key `target` to instantiate.") |
| return get_obj_from_str(config["target"], reload=reload)( |
| **config.get("params", dict()) |
| ) |
|
|
|
|
| def get_obj_from_str(string, reload=False): |
| module, cls = string.rsplit(".", 1) |
| if reload: |
| module_imp = importlib.import_module(module) |
| importlib.reload(module_imp) |
| return getattr(importlib.import_module(module, package=None), cls) |
|
|
|
|
| class TextEncoder(nn.Module): |
|
|
| def __init__(self, ckpt_path=None): |
| super().__init__() |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained(ckpt_path) |
| self.model = T5EncoderModel.from_pretrained(ckpt_path) |
| self.model.eval() |
| self.device = "cuda" |
|
|
| def get_uncond(self, batchsize): |
|
|
| hidden_state, _ = self([""]) |
| hidden_state = torch.cat([hidden_state] * batchsize).float() |
| attention_mask = ( |
| torch.ones((batchsize, hidden_state.size(1))).to(self.device).float() |
| ) |
| return [hidden_state, attention_mask] |
|
|
| def forward(self, batch): |
| return self.get_cond(batch) |
|
|
| def get_cond(self, prompt): |
| device = self.model.device |
| token = self.tokenizer( |
| prompt, |
| max_length=128, |
| padding=True, |
| truncation=True, |
| return_tensors="pt", |
| ) |
| input_ids, attention_mask = token.input_ids.to(device), token.attention_mask.to( |
| device |
| ) |
|
|
| with torch.no_grad(): |
| hidden_states = self.model( |
| input_ids=input_ids, attention_mask=attention_mask |
| )[0] |
|
|
| return [hidden_states, attention_mask.float()] |
|
|
