File size: 2,406 Bytes
feae2c5 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | # Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich
# copy from https://github.com/CompVis/latent-diffusion/blob/main/ldm/lr_scheduler.py under the MIT license
# LICENSE is in incl_licenses directory.
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()]
|