from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import torch
from torchvision.transforms import ToPILImage
from models.vqvae import VQVAEHF
from models.clip import FrozenCLIPEmbedder
from models.var import TVARHF, sample_with_top_k_top_p_, gumbel_softmax_with_rng
class TVARPipeline:
vae_path = "michellemoorre/vae-test"
text_encoder_path = "openai/clip-vit-large-patch14"
text_encoder_2_path = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
def __init__(self, var, vae, text_encoder, text_encoder_2, device):
self.var = var
self.vae = vae
self.text_encoder = text_encoder
self.text_encoder_2 = text_encoder_2
self.var.eval()
self.vae.eval()
self.device = device
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, device="cuda"):
var = TVARHF.from_pretrained(pretrained_model_name_or_path).to(device)
vae = VQVAEHF.from_pretrained(cls.vae_path).to(device)
text_encoder = FrozenCLIPEmbedder(cls.text_encoder_path, device=device)
text_encoder_2 = FrozenCLIPEmbedder(cls.text_encoder_2_path, device=device)
return cls(var, vae, text_encoder, text_encoder_2, device)
@staticmethod
def to_image(tensor):
return [ToPILImage()(
(255 * img.cpu().detach()).to(torch.uint8))
for img in tensor]
def encode_prompt(
self,
prompt: Union[str, List[str]],
null_prompt: str = "",
encode_null: bool = True,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
encodings = [
self.text_encoder.encode(prompt),
self.text_encoder_2.encode(prompt),
]
prompt_embeds = torch.concat(
[encoding.last_hidden_state for encoding in encodings], dim=-1
)
pooled_prompt_embeds = encodings[-1].pooler_output
attn_bias = encodings[-1].attn_bias
if encode_null:
null_prompt = [null_prompt] if isinstance(null_prompt, str) else prompt
null_encodings = [
self.text_encoder.encode(null_prompt),
self.text_encoder_2.encode(null_prompt),
]
null_prompt_embeds = torch.concat(
[encoding.last_hidden_state for encoding in encodings], dim=-1
)
null_pooled_prompt_embeds = null_encodings[-1].pooler_output
null_attn_bias = null_encodings[-1].attn_bias
B, L, hidden_dim = prompt_embeds.shape
pooled_dim = pooled_prompt_embeds.shape[1]
null_prompt_embeds = null_prompt_embeds[:, :L].expand(B, L, hidden_dim).to(prompt_embeds.device)
null_pooled_prompt_embeds = null_pooled_prompt_embeds.expand(B, pooled_dim).to(pooled_prompt_embeds.device)
null_attn_bias = null_attn_bias[:, :L].expand(B, L).to(attn_bias.device)
prompt_embeds = torch.cat([prompt_embeds, null_prompt_embeds], dim=0)
pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, null_pooled_prompt_embeds], dim=0)
attn_bias = torch.cat([attn_bias, null_attn_bias], dim=0)
return prompt_embeds, pooled_prompt_embeds, attn_bias
@torch.inference_mode()
def __call__(
self,
prompt = None,
null_prompt = "",
g_seed: Optional[int] = None,
cfg=4.0,
top_k=450,
top_p=0.95,
more_smooth=False,
re=False,
re_max_depth=10,
re_start_iter=2,
return_pil=True,
encoded_prompt = None,
encoded_null_prompt = None,
) -> torch.Tensor: # returns reconstructed image (B, 3, H, W) in [0, 1]
"""
only used for inference, on autoregressive mode
:param B: batch size
:param label_B: imagenet label; if None, randomly sampled
:param g_seed: random seed
:param cfg: classifier-free guidance ratio
:param top_k: top-k sampling
:param top_p: top-p sampling
:param more_smooth: smoothing the pred using gumbel softmax; only used in visualization, not used in FID/IS benchmarking
:return: if returns_vemb: list of embedding h_BChw := vae_embed(idx_Bl), else: list of idx_Bl
"""
assert not self.var.training
var = self.var
vae = self.vae
vae_quant = self.vae.quantize
if g_seed is None:
rng = None
else:
var.rng.manual_seed(g_seed)
rng = var.rng
if encoded_prompt is not None:
assert encoded_null_prompt is not None
context, cond_vector, context_attn_bias = self.var.parse_batch(
encoded_prompt,
encoded_null_prompt,
)
else:
context, cond_vector, context_attn_bias = self.encode_prompt(prompt, null_prompt)
B = context.shape[0] // 2
cond_vector = var.text_pooler(cond_vector)
sos = cond_BD = cond_vector
lvl_pos = var.lvl_embed(var.lvl_1L)
if not var.rope:
lvl_pos += var.pos_1LC
next_token_map = (
sos.unsqueeze(1)
+ var.pos_start.expand(2 * B, var.first_l, -1)
+ lvl_pos[:, : var.first_l]
)
cur_L = 0
f_hat = sos.new_zeros(B, var.Cvae, var.patch_nums[-1], var.patch_nums[-1])
for b in var.blocks:
b.attn.kv_caching(True)
b.cross_attn.kv_caching(True)
for si, pn in enumerate(var.patch_nums): # si: i-th segment
ratio = si / var.num_stages_minus_1
cond_BD_or_gss = var.shared_ada_lin(cond_BD)
x_BLC = next_token_map
if var.rope:
freqs_cis = var.freqs_cis[:, cur_L : cur_L + pn * pn]
else:
freqs_cis = var.freqs_cis
for block in var.blocks:
x_BLC = block(
x=x_BLC,
cond_BD=cond_BD_or_gss,
attn_bias=None,
context=context,
context_attn_bias=context_attn_bias,
freqs_cis=freqs_cis,
)
cur_L += pn * pn
logits_BlV = var.get_logits(x_BLC, cond_BD)
t = cfg * ratio
logits_BlV = (1 + t) * logits_BlV[:B] - t * logits_BlV[B:]
idx_Bl = sample_with_top_k_top_p_(
logits_BlV, rng=rng, top_k=top_k, top_p=top_p, num_samples=1
)[:, :, 0]
if re and si >= re_start_iter:
selected_logits = torch.gather(logits_BlV, -1, idx_Bl.unsqueeze(-1))[:, :, 0]
mx = selected_logits.sum(dim=-1)[:, None]
for _ in range(re_max_depth):
new_idx_Bl = sample_with_top_k_top_p_(
logits_BlV, rng=rng, top_k=top_k, top_p=top_p, num_samples=1
)[:, :, 0]
selected_logits = torch.gather(logits_BlV, -1, new_idx_Bl.unsqueeze(-1))[:, :, 0]
new_mx = selected_logits.sum(dim=-1)[:, None]
idx_Bl = idx_Bl * (mx >= new_mx) + new_idx_Bl * (mx < new_mx)
mx = mx * (mx >= new_mx) + new_mx * (mx < new_mx)
if not more_smooth: # this is the default case
h_BChw = vae_quant.embedding(idx_Bl) # B, l, Cvae
else: # not used when evaluating FID/IS/Precision/Recall
gum_t = max(0.27 * (1 - ratio * 0.95), 0.005) # refer to mask-git
h_BChw = gumbel_softmax_with_rng(
logits_BlV.mul(1 + ratio), tau=gum_t, hard=False, dim=-1, rng=rng
) @ vae_quant.embedding.weight.unsqueeze(0)
h_BChw = h_BChw.transpose_(1, 2).reshape(B, var.Cvae, pn, pn)
f_hat, next_token_map = vae_quant.get_next_autoregressive_input(
si, len(var.patch_nums), f_hat, h_BChw
)
if si != var.num_stages_minus_1: # prepare for next stage
next_token_map = next_token_map.view(B, var.Cvae, -1).transpose(1, 2)
next_token_map = (
var.word_embed(next_token_map)
+ lvl_pos[:, cur_L : cur_L + var.patch_nums[si + 1] ** 2]
)
next_token_map = next_token_map.repeat(
2, 1, 1
) # double the batch sizes due to CFG
for b in var.blocks:
b.attn.kv_caching(False)
b.cross_attn.kv_caching(False)
# de-normalize, from [-1, 1] to [0, 1]
img = vae.fhat_to_img(f_hat).add(1).mul(0.5)
if return_pil:
img = self.to_image(img)
return img