"""
Definition of Infinity transformer model.
"""
import math
import random
import time
from contextlib import nullcontext
from functools import partial
from typing import List, Optional, Tuple, Union, Dict, Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models import register_model
from torch.utils.checkpoint import checkpoint
from PIL import Image
import numpy as np
from torch.nn.attention.flex_attention import flex_attention
import utils.dist as dist
from utils.dist import for_visualize
from models.basic import flash_attn_func, flash_fused_op_installed, AdaLNBeforeHead, CrossAttnBlock, SelfAttnBlock, CrossAttention, FastRMSNorm, precompute_rope2d_freqs_grid
from utils import misc
from models.flex_attn import FlexAttn
from utils.dynamic_resolution import dynamic_resolution_h_w, h_div_w_templates
try:
from models.fused_op import fused_ada_layer_norm, fused_ada_rms_norm
except:
fused_ada_layer_norm, fused_ada_rms_norm = None, None
class MultiInpIdentity(nn.Module):
def forward(self, x, *args, **kwargs):
return x
class TextAttentivePool(nn.Module):
def __init__(self, Ct5: int, D: int):
super().__init__()
self.Ct5, self.D = Ct5, D
if D > 4096:
self.head_dim = 64
else:
self.head_dim = 128
self.num_heads = Ct5 // self.head_dim
self.ca = CrossAttention(for_attn_pool=True, embed_dim=self.D, kv_dim=Ct5, num_heads=self.num_heads)
def forward(self, ca_kv):
return self.ca(None, ca_kv).squeeze(1)
class SharedAdaLin(nn.Linear):
def forward(self, cond_BD):
C = self.weight.shape[0] // 6
return super().forward(cond_BD).reshape(-1, 1, 6, C) # B16C
class MultipleLayers(nn.Module):
def __init__(self, ls, num_blocks_in_a_chunk, index):
super().__init__()
self.module = nn.ModuleList()
for i in range(index, index+num_blocks_in_a_chunk):
self.module.append(ls[i])
def forward(self, x, cond_BD, ca_kv, attn_bias_or_two_vector, attn_fn=None, scale_schedule=None, checkpointing_full_block=False, rope2d_freqs_grid=None):
h = x
for m in self.module:
if checkpointing_full_block:
h = torch.utils.checkpoint.checkpoint(m, h, cond_BD, ca_kv, attn_bias_or_two_vector, attn_fn, scale_schedule, rope2d_freqs_grid, use_reentrant=False)
else:
h = m(h, cond_BD, ca_kv, attn_bias_or_two_vector, attn_fn, scale_schedule, rope2d_freqs_grid)
return h
class Infinity(nn.Module):
def __init__(
self, vae_local,
text_channels=0, text_maxlen=0, # text-cond generation
selecting_idx=None, # class-cond generation
embed_dim=1024, depth=16, num_heads=16, mlp_ratio=4., # model's architecture
drop_rate=0., drop_path_rate=0., # drop out and drop path
norm_eps=1e-6, rms_norm=False, # norm layer
shared_aln=False, head_aln=True, # adaptive norm
cond_drop_rate=0.1, # for classifier-free guidance
rand_uncond=False,
cross_attn_layer_scale=-1., nm0=False, tau=1, cos_attn=True, swiglu=False,
raw_scale_schedule=(1, 2, 3, 4, 5, 6, 8, 10, 13, 16),
head_depth=1,
top_p=0.0, top_k=0.0,
customized_flash_attn=False, fused_mlp=False, fused_norm=False,
block_chunks=1,
checkpointing=None,
pad_to_multiplier=0,
use_flex_attn=False,
batch_size=2,
add_lvl_embeding_only_first_block=1,
use_bit_label=1,
rope2d_each_sa_layer=0,
rope2d_normalized_by_hw=0,
pn=None,
train_h_div_w_list=None,
video_frames=1,
always_training_scales=20,
apply_spatial_patchify = 0,
inference_mode=False,
):
# set hyperparameters
self.C = embed_dim
self.inference_mode = inference_mode
self.apply_spatial_patchify = apply_spatial_patchify
if self.apply_spatial_patchify:
self.d_vae = vae_local.embed_dim * 4
else:
self.d_vae = vae_local.embed_dim
self.use_bit_label = use_bit_label
self.codebook_dim = self.d_vae
self.V = (self.codebook_dim * 2) if self.use_bit_label else vae_local.vocab_size
self.bit_mask = vae_local.quantizer.lfq.mask if self.use_bit_label else None
self.Ct5 = text_channels
self.depth = depth
self.num_heads = num_heads
self.batch_size = batch_size
self.mlp_ratio = mlp_ratio
self.cond_drop_rate = cond_drop_rate
self.norm_eps = norm_eps
self.prog_si = -1
self.pn = pn
self.train_h_div_w_list = train_h_div_w_list if train_h_div_w_list else h_div_w_templates
self.video_frames = video_frames
self.always_training_scales = always_training_scales
assert add_lvl_embeding_only_first_block in [0,1]
self.add_lvl_embeding_only_first_block = add_lvl_embeding_only_first_block
assert rope2d_each_sa_layer in [0,1]
self.rope2d_each_sa_layer = rope2d_each_sa_layer
self.rope2d_normalized_by_hw = rope2d_normalized_by_hw
print(f'self.codebook_dim: {self.codebook_dim}, self.add_lvl_embeding_only_first_block: {self.add_lvl_embeding_only_first_block}, \
self.use_bit_label: {self.use_bit_label}, self.rope2d_each_sa_layer: {rope2d_each_sa_layer}, self.rope2d_normalized_by_hw: {self.rope2d_normalized_by_hw}')
head_up_method = ''
word_patch_size = 1 if head_up_method in {'', 'no'} else 2
if word_patch_size > 1:
assert all(raw_pn % word_patch_size == 0 for raw_pn in raw_scale_schedule), f'raw_scale_schedule={raw_scale_schedule}, not compatible with word_patch_size={word_patch_size}'
self.checkpointing = checkpointing
self.pad_to_multiplier = max(1, pad_to_multiplier)
customized_kernel_installed = any('Infinity' in arg_name for arg_name in flash_attn_func.__code__.co_varnames)
self.customized_flash_attn = customized_flash_attn and customized_kernel_installed
if customized_flash_attn and not customized_kernel_installed:
import inspect, warnings
file_path = inspect.getsourcefile(flash_attn_func)
line_number = inspect.getsourcelines(flash_attn_func)[1]
info = (
f'>>>>>> Customized FlashAttention2 is not installed or compiled, but specified in args by --flash=1. Set customized_flash_attn = False. <<<<<<\n'
f'>>>>>> `flash_attn_func` is in [line {line_number}] [file {file_path}] <<<<<<\n'
f'>>>>>> {flash_attn_func.__code__.co_varnames=} <<<<<<\n'
)
warnings.warn(info, ImportWarning)
print(info, flush=True)
self.raw_scale_schedule = raw_scale_schedule # 'raw' means before any patchifying
self.first_l = 1
# solve top-p top-k sampling hyperparameters
self.top_p, self.top_k = max(min(top_p, 1), 0), (round(top_k * self.V) if 0 < top_k < 1 else round(top_k))
if self.top_p < 1e-5: self.top_p = 0
if self.top_k >= self.V or self.top_k <= 0: self.top_k = 0
t = torch.zeros(dist.get_world_size(), device=dist.get_device())
t[dist.get_rank()] = float(flash_fused_op_installed)
dist.barrier()
dist.allreduce(t)
assert round(t.sum().item()) in {0, dist.get_world_size()}, f'flash_fused_op_installed: {t}'
super().__init__()
self.rng = torch.Generator(device=dist.get_device())
self.maybe_record_function = nullcontext
self.text_maxlen = text_maxlen
self.t2i = text_channels != 0
# [inp & position embedding]
init_std = math.sqrt(1 / self.C / 3)
self.norm0_cond = nn.Identity()
if self.t2i:
self.selecting_idx = None
self.num_classes = 0
self.D = self.C
cfg_uncond = torch.empty(self.text_maxlen, self.Ct5)
rng = torch.Generator(device='cpu')
rng.manual_seed(0)
torch.nn.init.trunc_normal_(cfg_uncond, std=1.2, generator=rng)
cfg_uncond /= self.Ct5 ** 0.5
if rand_uncond:
self.register_buffer('cfg_uncond', cfg_uncond)
else:
self.cfg_uncond = nn.Parameter(cfg_uncond)
self.text_norm = FastRMSNorm(self.Ct5, elementwise_affine=True, eps=norm_eps)
self.text_proj_for_sos = TextAttentivePool(self.Ct5, self.D)
self.text_proj_for_ca = nn.Sequential(
nn.Linear(self.Ct5, self.D),
nn.GELU(approximate='tanh'),
nn.Linear(self.D, self.D),
)
else: # class-label cond
if selecting_idx is None:
num_classes = 1000
print(f'======= WARNING: selecting_idx not specified, set to 1/{num_classes} @ {dist.get_device()} =======')
selecting_idx = torch.full((1, num_classes), fill_value=1/num_classes, dtype=torch.float32, device=dist.get_device())
self.selecting_idx = selecting_idx
self.num_classes = selecting_idx.shape[-1]
self.D = self.C
self.class_emb = nn.Embedding(self.num_classes + 1, self.C)
nn.init.trunc_normal_(self.class_emb.weight.data, mean=0, std=init_std)
self.pos_start = nn.Parameter(torch.empty(1, self.first_l, self.C))
nn.init.trunc_normal_(self.pos_start.data, mean=0, std=init_std)
if self.rope2d_each_sa_layer:
rope2d_freqs_grid = precompute_rope2d_freqs_grid(dim=self.C//self.num_heads, dynamic_resolution_h_w=dynamic_resolution_h_w, pad_to_multiplier=self.pad_to_multiplier, rope2d_normalized_by_hw=self.rope2d_normalized_by_hw)
self.rope2d_freqs_grid = rope2d_freqs_grid
else:
raise ValueError(f'self.rope2d_each_sa_layer={self.rope2d_each_sa_layer} not implemented')
self.lvl_embed = nn.Embedding(15, self.C)
nn.init.trunc_normal_(self.lvl_embed.weight.data, mean=0, std=init_std)
# [input layers] input norm && input embedding
norm_layer = partial(FastRMSNorm if rms_norm else nn.LayerNorm, eps=norm_eps)
self.norm0_ve = norm_layer(self.d_vae) if nm0 else nn.Identity()
self.word_embed = nn.Linear(self.d_vae, self.C)
# [shared adaptive layernorm mapping network]
self.shared_ada_lin = nn.Sequential(nn.SiLU(inplace=False), SharedAdaLin(self.D, 6*self.C)) if shared_aln else nn.Identity()
# fused norm
if fused_norm:
fused_norm_func = fused_ada_rms_norm if rms_norm else fused_ada_layer_norm
if fused_norm_func is not None: # pre-compile
B = 2
x = torch.randn(B, 1, self.C).requires_grad_(True)
scale = torch.randn(B, 1, self.C).mul_(0.01).requires_grad_(True)
shift = torch.randn(B, 1, self.C).mul_(0.01).requires_grad_(True)
# fused_norm_func(C=self.C, eps=self.norm_eps, x=x, scale=scale, shift=shift).mean().backward()
del B, x, scale, shift
else:
fused_norm_func = None
# [backbone and head]
self.use_flex_attn = use_flex_attn
self.attn_fn_compile_dict = {}
self.batch_size = batch_size
if self.use_flex_attn:
self.attn_fn_compile_dict = self.compile_flex_attn()
self.drop_path_rate = drop_path_rate
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # dpr means drop path rate (linearly increasing)
self.unregistered_blocks = []
for block_idx in range(depth):
block = (CrossAttnBlock if self.t2i else SelfAttnBlock)(
embed_dim=self.C, kv_dim=self.D, cross_attn_layer_scale=cross_attn_layer_scale, cond_dim=self.D, act=True, shared_aln=shared_aln, norm_layer=norm_layer,
num_heads=num_heads, mlp_ratio=mlp_ratio, drop=drop_rate, drop_path=dpr[block_idx], tau=tau, cos_attn=cos_attn,
swiglu=swiglu, customized_flash_attn=self.customized_flash_attn, fused_mlp=fused_mlp, fused_norm_func=fused_norm_func,
checkpointing_sa_only=self.checkpointing == 'self-attn',
use_flex_attn=use_flex_attn, batch_size=batch_size, pad_to_multiplier=pad_to_multiplier, rope2d_normalized_by_hw=rope2d_normalized_by_hw,
)
self.unregistered_blocks.append(block)
# [head]
V = self.V
if head_aln:
self.head_nm = AdaLNBeforeHead(self.C, self.D, act=True, norm_layer=norm_layer, fused_norm_func=fused_norm_func)
self.head = nn.Linear(self.C, V) if head_depth == 1 else nn.Sequential(nn.Linear(self.C, self.C, bias=True), nn.GELU(approximate='tanh'), nn.Linear(self.C, V))
else:
self.head_nm = MultiInpIdentity()
self.head = nn.Sequential(norm_layer(self.C), nn.Linear(self.C, V)) if head_depth == 1 else nn.Sequential(norm_layer(self.C), nn.Linear(self.C, self.C, bias=True), nn.GELU(approximate='tanh'), nn.Linear(self.C, V))
self.num_block_chunks = block_chunks or 1
self.num_blocks_in_a_chunk = depth // block_chunks
print(f"{self.num_blocks_in_a_chunk=}, {depth=}, {block_chunks=}")
assert self.num_blocks_in_a_chunk * block_chunks == depth
if self.num_block_chunks == 1:
self.blocks = nn.ModuleList(self.unregistered_blocks)
else:
self.block_chunks = nn.ModuleList()
for i in range(self.num_block_chunks):
self.block_chunks.append(MultipleLayers(self.unregistered_blocks, self.num_blocks_in_a_chunk, i*self.num_blocks_in_a_chunk))
print(
f'\n[constructor] ==== customized_flash_attn={self.customized_flash_attn} (using_flash={sum((b.sa.using_flash if self.t2i else b.attn.using_flash) for b in self.unregistered_blocks)}/{self.depth}), fused_mlp={fused_mlp} (fused_mlp={sum(b.ffn.fused_mlp_func is not None for b in self.unregistered_blocks)}/{self.depth}) ==== \n'
f' [Infinity config ] embed_dim={embed_dim}, num_heads={num_heads}, depth={depth}, mlp_ratio={mlp_ratio}, swiglu={swiglu} num_blocks_in_a_chunk={self.num_blocks_in_a_chunk}\n'
f' [drop ratios] drop_rate={drop_rate}, drop_path_rate={drop_path_rate:g} ({torch.linspace(0, drop_path_rate, depth)})',
end='\n\n', flush=True
)
def compile_flex_attn(self):
attn_fn_compile_dict = {}
for h_div_w in self.train_h_div_w_list:
h_div_w_template = h_div_w_templates[np.argmin(np.abs(float(h_div_w) - h_div_w_templates))]
full_scale_schedule = dynamic_resolution_h_w[h_div_w_template][self.pn]['scales']
if self.inference_mode:
apply_flex_attn_scales = list(range(1, 1+len(full_scale_schedule)))
mask_type = "infinity_infer_mask_with_kv_cache"
auto_padding = True
else:
mask_type = 'var'
auto_padding = False
apply_flex_attn_scales = [min(self.always_training_scales, len(full_scale_schedule))]
for scales_num in apply_flex_attn_scales:
print(f'====== apply flex attn hdivw: {h_div_w} scales: {scales_num} ======')
scale_schedule = full_scale_schedule[:scales_num]
scale_schedule = [ (min(t, self.video_frames//4+1), h, w) for (t,h, w) in scale_schedule]
patchs_nums_tuple = tuple(scale_schedule)
SEQ_L = sum( pt * ph * pw for pt, ph, pw in patchs_nums_tuple)
aligned_L = SEQ_L+ (self.pad_to_multiplier - SEQ_L % self.pad_to_multiplier) if SEQ_L % self.pad_to_multiplier != 0 else SEQ_L
attn_fn = FlexAttn(block_scales = patchs_nums_tuple,
mask_type = mask_type,
B = self.batch_size,
H = self.num_heads,
L = aligned_L,
auto_padding=auto_padding)
attn_fn_compile_dict[patchs_nums_tuple] = attn_fn
if self.video_frames > 1: # append image attn_fn when self.video_frames > 1 (namely videos)
scale_schedule = [ (1, h, w) for (t,h, w) in scale_schedule]
patchs_nums_tuple = tuple(scale_schedule)
SEQ_L = sum( pt * ph * pw for pt, ph, pw in patchs_nums_tuple)
aligned_L = SEQ_L+ (self.pad_to_multiplier - SEQ_L % self.pad_to_multiplier) if SEQ_L % self.pad_to_multiplier != 0 else SEQ_L
attn_fn = FlexAttn(block_scales = patchs_nums_tuple,
mask_type = mask_type,
B = self.batch_size,
H = self.num_heads,
L = aligned_L)
attn_fn_compile_dict[patchs_nums_tuple] = attn_fn
return attn_fn_compile_dict
def get_logits(self, h: torch.Tensor, cond_BD: Optional[torch.Tensor]):
"""
:param h: hidden_state, shaped (B or batch_size, L or seq_len, C or hidden_dim)
:param cond_BD: shaped (B or batch_size, D or cond_dim)
:param tau: temperature
:return: logits, shaped (B or batch_size, V or vocabulary_size)
"""
with torch.amp.autocast('cuda', enabled=False):
return self.head(self.head_nm(h.float(), cond_BD.float()))
def add_lvl_embeding(self, feature, scale_ind, scale_schedule, need_to_pad=0):
bs, seq_len, c = feature.shape
patch_t, patch_h, patch_w = scale_schedule[scale_ind]
t_mul_h_mul_w = patch_t * patch_h * patch_w
assert t_mul_h_mul_w + need_to_pad == seq_len
feature[:, :t_mul_h_mul_w] += self.lvl_embed(scale_ind*torch.ones((bs, t_mul_h_mul_w),dtype=torch.int).to(feature.device))
return feature
def add_lvl_embeding_for_x_BLC(self, x_BLC, scale_schedule, need_to_pad=0):
ptr = 0
x_BLC_list = []
for scale_ind, patch_t_h_w in enumerate(scale_schedule):
scale_seq_len = np.array(patch_t_h_w).prod()
x_BLC_this_scale = x_BLC[:,ptr:ptr+scale_seq_len] # shape: [bs, patch_h*patch_w, c]
ptr += scale_seq_len
x_BLC_this_scale = self.add_lvl_embeding(x_BLC_this_scale, scale_ind, scale_schedule)
x_BLC_list.append(x_BLC_this_scale)
assert x_BLC.shape[1] == (ptr + need_to_pad), f'{x_BLC.shape[1]} != {ptr} + {need_to_pad}'
x_BLC_list.append(x_BLC[:,ptr:])
x_BLC = torch.cat(x_BLC_list, dim=1)
return x_BLC
def forward(self, label_B_or_BLT: Union[torch.LongTensor, Tuple[torch.FloatTensor, torch.IntTensor, int]], x_BLC_wo_prefix: torch.Tensor, scale_schedule: List[Tuple[int]],
cfg_infer=False,
**kwargs,
) -> Union[torch.Tensor, List[torch.Tensor]]: # returns logits_BLV
"""
label_B_or_BLT: label_B or (kv_compact, cu_seqlens_k, max_seqlen_k)
:return: logits BLV, V is vocab_size
"""
if cfg_infer:
return self.autoregressive_infer_cfg(label_B_or_BLT=label_B_or_BLT, scale_schedule=scale_schedule, **kwargs)
x_BLC_wo_prefix = x_BLC_wo_prefix.float() # input should be float32
B = x_BLC_wo_prefix.shape[0]
# [1. get input sequence x_BLC]
with torch.amp.autocast('cuda', enabled=False):
kv_compact, lens, cu_seqlens_k, max_seqlen_k = label_B_or_BLT
# drop cond
total = 0
for le in lens:
if random.random() < self.cond_drop_rate:
kv_compact[total:total+le] = self.cfg_uncond[:le]
total += le
must_on_graph = self.cfg_uncond[0, 0] * 0
kv_compact = self.text_norm(kv_compact).contiguous()
sos = cond_BD = self.text_proj_for_sos((kv_compact, cu_seqlens_k, max_seqlen_k)).float().contiguous() # cond_BD should be float32
kv_compact = self.text_proj_for_ca(kv_compact).contiguous()
kv_compact[0, 0] += must_on_graph
ca_kv = kv_compact, cu_seqlens_k, max_seqlen_k
cond_BD_or_gss = self.shared_ada_lin(cond_BD).contiguous() # gss: gamma, scale, shift; cond_BD_or_gss should be float32
sos = sos.unsqueeze(1).expand(B, 1, -1) + self.pos_start.expand(B, 1, -1)
x_BLC = torch.cat((sos, self.word_embed(self.norm0_ve(x_BLC_wo_prefix))), dim=1)
# [1.1. pad the seqlen dim]
l_end = x_BLC.shape[1]
need_to_pad = (l_end + self.pad_to_multiplier - 1) // self.pad_to_multiplier * self.pad_to_multiplier - l_end # 0
if self.customized_flash_attn:
Infinity_visible_kvlen = self.Infinity_visible_kvlen[:l_end]
Infinity_invisible_qlen = self.Infinity_invisible_qlen[:l_end]
attn_bias_or_two_vector = (Infinity_visible_kvlen, Infinity_invisible_qlen)
# todo: solve need_to_pad here
elif self.use_flex_attn:
if need_to_pad:
x_BLC = F.pad(x_BLC, (0, 0, 0, need_to_pad))
assert x_BLC.shape[-1] % 128 == 0, 'x_BLC.shape[-1] % 128 != 0'
attn_bias_or_two_vector = None
else:
d: torch.Tensor = torch.cat([torch.full((pn[0]*pn[1]*pn[2],), i) for i, pn in enumerate(scale_schedule)]).view(1, l_end, 1)
dT = d.transpose(1, 2) # dT: 11L
attn_bias_for_masking = torch.where(d >= dT, 0., -torch.inf).reshape(1, 1, l_end, l_end)
attn_bias = attn_bias_for_masking[:, :, :l_end, :l_end].contiguous() # attn_bias: 11LL
if need_to_pad:
attn_bias = F.pad(attn_bias, (0, need_to_pad, 0, need_to_pad), value=-torch.inf)
attn_bias[0, 0, l_end:, 0] = 0
x_BLC = F.pad(x_BLC, (0, 0, 0, need_to_pad))
attn_bias_or_two_vector = attn_bias.type_as(x_BLC).to(x_BLC.device)
if self.use_flex_attn:
attn_fn = self.attn_fn_compile_dict[tuple(scale_schedule)]
else:
attn_fn = None
# [2. block loop]
SelfAttnBlock.forward, CrossAttnBlock.forward
checkpointing_full_block = self.checkpointing == 'full-block' and self.training
if self.num_block_chunks == 1:
for i, b in enumerate(self.blocks):
if self.add_lvl_embeding_only_first_block and i == 0:
x_BLC = self.add_lvl_embeding_for_x_BLC(x_BLC, scale_schedule, need_to_pad)
if not self.add_lvl_embeding_only_first_block:
x_BLC = self.add_lvl_embeding_for_x_BLC(x_BLC, scale_schedule, need_to_pad)
if checkpointing_full_block:
x_BLC = torch.utils.checkpoint.checkpoint(b, x_BLC, cond_BD_or_gss, ca_kv, attn_bias_or_two_vector, attn_fn, scale_schedule, self.rope2d_freqs_grid, use_reentrant=False)
else:
x_BLC = b(x=x_BLC, cond_BD=cond_BD_or_gss, ca_kv=ca_kv, attn_bias_or_two_vector=attn_bias_or_two_vector, attn_fn=attn_fn, scale_schedule=scale_schedule, rope2d_freqs_grid=self.rope2d_freqs_grid)
else:
for i, chunk in enumerate(self.block_chunks): # this path
if self.add_lvl_embeding_only_first_block and i == 0:
x_BLC = self.add_lvl_embeding_for_x_BLC(x_BLC, scale_schedule, need_to_pad)
if not self.add_lvl_embeding_only_first_block:
x_BLC = self.add_lvl_embeding_for_x_BLC(x_BLC, scale_schedule, need_to_pad)
x_BLC = chunk(x=x_BLC, cond_BD=cond_BD_or_gss, ca_kv=ca_kv, attn_bias_or_two_vector=attn_bias_or_two_vector, attn_fn=attn_fn, scale_schedule=scale_schedule, checkpointing_full_block=checkpointing_full_block, rope2d_freqs_grid=self.rope2d_freqs_grid)
# [3. unpad the seqlen dim, and then get logits]
return self.get_logits(x_BLC[:, :l_end], cond_BD) # return logits BLV, V is vocab_size
@torch.no_grad()
def autoregressive_infer_cfg(
self,
vae=None,
scale_schedule=None,
label_B_or_BLT=None,
B=1, negative_label_B_or_BLT=None, force_gt_Bhw=None,
g_seed=None, cfg_list=[], tau_list=[], cfg_sc=3, top_k=0, top_p=0.0,
returns_vemb=0, ratio_Bl1=None, gumbel=0, norm_cfg=False,
cfg_exp_k: float=0.0, cfg_insertion_layer=[-5],
vae_type=0, softmax_merge_topk=-1, ret_img=False,
trunk_scale=1000,
gt_leak=0, gt_ls_Bl=None,
inference_mode=False,
save_img_path=None,
sampling_per_bits=1,
): # returns List[idx_Bl]
if g_seed is None: rng = None
else: self.rng.manual_seed(g_seed); rng = self.rng
assert len(cfg_list) >= len(scale_schedule)
assert len(tau_list) >= len(scale_schedule)
# scale_schedule is used by infinity, vae_scale_schedule is used by vae if there exists a spatial patchify,
# we need to convert scale_schedule to vae_scale_schedule by multiply 2 to h and w
if self.apply_spatial_patchify:
vae_scale_schedule = [(pt, 2*ph, 2*pw) for pt, ph, pw in scale_schedule]
else:
vae_scale_schedule = scale_schedule
kv_compact, lens, cu_seqlens_k, max_seqlen_k = label_B_or_BLT
if any(np.array(cfg_list) != 1):
bs = 2*B
if not negative_label_B_or_BLT:
kv_compact_un = kv_compact.clone()
total = 0
for le in lens:
kv_compact_un[total:total+le] = (self.cfg_uncond)[:le]
total += le
kv_compact = torch.cat((kv_compact, kv_compact_un), dim=0)
cu_seqlens_k = torch.cat((cu_seqlens_k, cu_seqlens_k[1:]+cu_seqlens_k[-1]), dim=0)
else:
kv_compact_un, lens_un, cu_seqlens_k_un, max_seqlen_k_un = negative_label_B_or_BLT
kv_compact = torch.cat((kv_compact, kv_compact_un), dim=0)
cu_seqlens_k = torch.cat((cu_seqlens_k, cu_seqlens_k_un[1:]+cu_seqlens_k[-1]), dim=0)
max_seqlen_k = max(max_seqlen_k, max_seqlen_k_un)
else:
bs = B
kv_compact = self.text_norm(kv_compact)
sos = cond_BD = self.text_proj_for_sos((kv_compact, cu_seqlens_k, max_seqlen_k)) # sos shape: [2, 4096]
kv_compact = self.text_proj_for_ca(kv_compact) # kv_compact shape: [304, 4096]
ca_kv = kv_compact, cu_seqlens_k, max_seqlen_k
last_stage = sos.unsqueeze(1).expand(bs, 1, -1) + self.pos_start.expand(bs, 1, -1)
with torch.amp.autocast('cuda', enabled=False):
cond_BD_or_gss = self.shared_ada_lin(cond_BD.float()).float().contiguous()
accu_BChw, cur_L, ret = None, 0, [] # current length, list of reconstructed images
idx_Bl_list, idx_Bld_list = [], []
if inference_mode:
for b in self.unregistered_blocks: (b.sa if isinstance(b, CrossAttnBlock) else b.attn).kv_caching(True)
else:
assert self.num_block_chunks > 1
for block_chunk_ in self.block_chunks:
for module in block_chunk_.module.module:
(module.sa if isinstance(module, CrossAttnBlock) else module.attn).kv_caching(True)
abs_cfg_insertion_layers = []
add_cfg_on_logits, add_cfg_on_probs = False, False
leng = len(self.unregistered_blocks)
for item in cfg_insertion_layer:
if item == 0: # add cfg on logits
add_cfg_on_logits = True
elif item == 1: # add cfg on probs
add_cfg_on_probs = True # todo in the future, we may want to add cfg on logits and probs
elif item < 0: # determine to add cfg at item-th layer's output
assert leng+item > 0, f'cfg_insertion_layer: {item} is not valid since len(unregistered_blocks)={self.num_block_chunks}'
abs_cfg_insertion_layers.append(leng+item)
else:
raise ValueError(f'cfg_insertion_layer: {item} is not valid')
num_stages_minus_1 = len(scale_schedule)-1
summed_codes = 0
for si, pn in enumerate(scale_schedule): # si: i-th segment
cfg = cfg_list[si]
if si >= trunk_scale:
break
cur_L += np.array(pn).prod()
need_to_pad = 0
attn_fn = None
if self.use_flex_attn:
# need_to_pad = (self.pad_to_multiplier - cur_L % self.pad_to_multiplier) % self.pad_to_multiplier
# if need_to_pad:
# last_stage = F.pad(last_stage, (0, 0, 0, need_to_pad))
attn_fn = self.attn_fn_compile_dict.get(tuple(scale_schedule[:(si+1)]), None)
# assert self.attn_bias_for_masking[:, :, last_L:cur_L, :cur_L].sum() == 0, f'AR with {(self.attn_bias_for_masking[:, :, last_L:cur_L, :cur_L] != 0).sum()} / {self.attn_bias_for_masking[:, :, last_L:cur_L, :cur_L].numel()} mask item'
layer_idx = 0
for block_idx, b in enumerate(self.block_chunks):
# last_stage shape: [4, 1, 2048], cond_BD_or_gss.shape: [4, 1, 6, 2048], ca_kv[0].shape: [64, 2048], ca_kv[1].shape [5], ca_kv[2]: int
if self.add_lvl_embeding_only_first_block and block_idx == 0:
last_stage = self.add_lvl_embeding(last_stage, si, scale_schedule, need_to_pad=need_to_pad)
if not self.add_lvl_embeding_only_first_block:
last_stage = self.add_lvl_embeding(last_stage, si, scale_schedule, need_to_pad=need_to_pad)
for m in b.module:
last_stage = m(x=last_stage, cond_BD=cond_BD_or_gss, ca_kv=ca_kv, attn_bias_or_two_vector=None, attn_fn=attn_fn, scale_schedule=scale_schedule, rope2d_freqs_grid=self.rope2d_freqs_grid, scale_ind=si)
if (cfg != 1) and (layer_idx in abs_cfg_insertion_layers):
# print(f'add cfg={cfg} on {layer_idx}-th layer output')
last_stage = cfg * last_stage[:B] + (1-cfg) * last_stage[B:]
last_stage = torch.cat((last_stage, last_stage), 0)
layer_idx += 1
if (cfg != 1) and add_cfg_on_logits:
# print(f'add cfg on add_cfg_on_logits')
logits_BlV = self.get_logits(last_stage, cond_BD).mul(1/tau_list[si])
logits_BlV = cfg * logits_BlV[:B] + (1-cfg) * logits_BlV[B:]
else:
logits_BlV = self.get_logits(last_stage[:B], cond_BD[:B]).mul(1/tau_list[si])
if self.use_bit_label:
tmp_bs, tmp_seq_len = logits_BlV.shape[:2]
logits_BlV = logits_BlV.reshape(tmp_bs, -1, 2)
idx_Bld = sample_with_top_k_top_p_also_inplace_modifying_logits_(logits_BlV, rng=rng, top_k=top_k or self.top_k, top_p=top_p or self.top_p, num_samples=1)[:, :, 0]
idx_Bld = idx_Bld.reshape(tmp_bs, tmp_seq_len, -1)
else:
idx_Bl = sample_with_top_k_top_p_also_inplace_modifying_logits_(logits_BlV, rng=rng, top_k=top_k or self.top_k, top_p=top_p or self.top_p, num_samples=1)[:, :, 0]
if vae_type != 0:
assert returns_vemb
if si < gt_leak:
idx_Bld = gt_ls_Bl[si]
else:
assert pn[0] == 1
idx_Bld = idx_Bld.reshape(B, pn[1], pn[2], -1) # shape: [B, h, w, d] or [B, h, w, 4d]
if self.apply_spatial_patchify: # unpatchify operation
idx_Bld = idx_Bld.permute(0,3,1,2) # [B, 4d, h, w]
idx_Bld = torch.nn.functional.pixel_shuffle(idx_Bld, 2) # [B, d, 2h, 2w]
idx_Bld = idx_Bld.permute(0,2,3,1) # [B, 2h, 2w, d]
idx_Bld = idx_Bld.unsqueeze(1) # [B, 1, h, w, d] or [B, 1, 2h, 2w, d]
idx_Bld_list.append(idx_Bld)
codes = vae.quantizer.lfq.indices_to_codes(idx_Bld, label_type='bit_label') # [B, d, 1, h, w] or [B, d, 1, 2h, 2w]
if si != num_stages_minus_1:
summed_codes += F.interpolate(codes, size=vae_scale_schedule[-1], mode=vae.quantizer.z_interplote_up)
last_stage = F.interpolate(summed_codes, size=vae_scale_schedule[si+1], mode=vae.quantizer.z_interplote_down) # [B, d, 1, h, w] or [B, d, 1, 2h, 2w]
last_stage = last_stage.squeeze(-3) # [B, d, h, w] or [B, d, 2h, 2w]
if self.apply_spatial_patchify: # patchify operation
last_stage = torch.nn.functional.pixel_unshuffle(last_stage, 2) # [B, 4d, h, w]
last_stage = last_stage.reshape(*last_stage.shape[:2], -1) # [B, d, h*w] or [B, 4d, h*w]
last_stage = torch.permute(last_stage, [0,2,1]) # [B, h*w, d] or [B, h*w, 4d]
else:
summed_codes += codes
else:
if si < gt_leak:
idx_Bl = gt_ls_Bl[si]
h_BChw = self.quant_only_used_in_inference[0].embedding(idx_Bl).float() # BlC
# h_BChw = h_BChw.float().transpose_(1, 2).reshape(B, self.d_vae, scale_schedule[si][0], scale_schedule[si][1])
h_BChw = h_BChw.transpose_(1, 2).reshape(B, self.d_vae, scale_schedule[si][0], scale_schedule[si][1], scale_schedule[si][2])
ret.append(h_BChw if returns_vemb != 0 else idx_Bl)
idx_Bl_list.append(idx_Bl)
if si != num_stages_minus_1:
accu_BChw, last_stage = self.quant_only_used_in_inference[0].one_step_fuse(si, num_stages_minus_1+1, accu_BChw, h_BChw, scale_schedule)
if si != num_stages_minus_1:
last_stage = self.word_embed(self.norm0_ve(last_stage))
last_stage = last_stage.repeat(bs//B, 1, 1)
if inference_mode:
for b in self.unregistered_blocks: (b.sa if isinstance(b, CrossAttnBlock) else b.attn).kv_caching(False)
else:
assert self.num_block_chunks > 1
for block_chunk_ in self.block_chunks:
for module in block_chunk_.module.module:
(module.sa if isinstance(module, CrossAttnBlock) else module.attn).kv_caching(False)
if not ret_img:
return ret, idx_Bl_list, []
if vae_type != 0:
img = vae.decode(summed_codes.squeeze(-3))
else:
img = vae.viz_from_ms_h_BChw(ret, scale_schedule=scale_schedule, same_shape=True, last_one=True)
img = (img + 1) / 2
img = img.permute(0, 2, 3, 1).mul_(255).to(torch.uint8).flip(dims=(3,))
return ret, idx_Bl_list, img
@for_visualize
def vis_key_params(self, ep):
return
def load_state_dict(self, state_dict: Dict[str, Any], strict=False, assign=False):
for k in state_dict:
if 'cfg_uncond' in k:
old, new = state_dict[k], self.cfg_uncond.data
min_tlen = min(old.shape[0], new.shape[0])
if min_tlen == old.shape[0]:
state_dict[k] = torch.cat((old.to(device=new.device, dtype=new.dtype), new[min_tlen:]))
else:
state_dict[k] = old[:min_tlen]
for buf_name in ('lvl_1L', 'attn_bias_for_masking', 'Infinity_visible_kvlen', 'Infinity_invisible_qlen'):
state_dict.pop(buf_name, None)
if hasattr(self, buf_name):
state_dict[buf_name] = getattr(self, buf_name)
return super().load_state_dict(state_dict=state_dict, strict=strict, assign=assign)
def special_init(
self,
aln_init: float,
aln_gamma_init: float,
scale_head: float,
scale_proj: int,
):
# init head's norm
if isinstance(self.head_nm, AdaLNBeforeHead):
self.head_nm.ada_lin[-1].weight.data.mul_(aln_init) # there's no gamma for head
if hasattr(self.head_nm.ada_lin[-1], 'bias') and self.head_nm.ada_lin[-1].bias is not None:
self.head_nm.ada_lin[-1].bias.data.zero_()
# init head's proj
if scale_head >= 0:
if isinstance(self.head, nn.Linear):
self.head.weight.data.mul_(scale_head)
self.head.bias.data.zero_()
elif isinstance(self.head, nn.Sequential):
self.head[-1].weight.data.mul_(scale_head)
self.head[-1].bias.data.zero_()
depth = len(self.unregistered_blocks)
for block_idx, sab in enumerate(self.unregistered_blocks):
sab: Union[SelfAttnBlock, CrossAttnBlock]
# init proj
scale = 1 / math.sqrt(2*depth if scale_proj == 1 else 2*(1 + block_idx))
if scale_proj == 1:
if self.t2i:
sab.sa.proj.weight.data.mul_(scale)
sab.ca.proj.weight.data.mul_(scale)
else:
sab.attn.proj.weight.data.mul_(scale)
sab.ffn.fc2.weight.data.mul_(scale)
# if sab.using_swiglu:
# nn.init.ones_(sab.ffn.fcg.bias)
# nn.init.trunc_normal_(sab.ffn.fcg.weight, std=1e-5)
# init ada_lin
if hasattr(sab, 'ada_lin'):
lin = sab.ada_lin[-1]
lin.weight.data[:2*self.C].mul_(aln_gamma_init) # init gamma
lin.weight.data[2*self.C:].mul_(aln_init) # init scale and shift
if hasattr(lin, 'bias') and lin.bias is not None:
lin.bias.data.zero_()
elif hasattr(sab, 'ada_gss'):
sab.ada_gss.data[:, :, :2, :].mul_(aln_gamma_init) # init gamma
sab.ada_gss.data[:, :, 2:, :].mul_(aln_init) # init scale and shift
def extra_repr(self):
return f'drop_path_rate={self.drop_path_rate}'
def get_layer_id_and_scale_exp(self, para_name: str):
raise NotImplementedError
def sample_with_top_k_top_p_also_inplace_modifying_logits_(logits_BlV: torch.Tensor, top_k: int = 0, top_p: float = 0.0, rng=None, num_samples=1) -> torch.Tensor: # return idx, shaped (B, l)
B, l, V = logits_BlV.shape
if top_k > 0:
top_k = min(top_k, V)
idx_to_remove = logits_BlV < logits_BlV.topk(top_k, largest=True, sorted=False, dim=-1)[0].amin(dim=-1, keepdim=True)
logits_BlV.masked_fill_(idx_to_remove, -torch.inf)
if top_p > 0:
sorted_logits, sorted_idx = logits_BlV.sort(dim=-1, descending=False)
sorted_idx_to_remove = sorted_logits.softmax(dim=-1).cumsum_(dim=-1) <= (1 - top_p)
sorted_idx_to_remove[..., -1:] = False
logits_BlV.masked_fill_(sorted_idx_to_remove.scatter(sorted_idx.ndim - 1, sorted_idx, sorted_idx_to_remove), -torch.inf)
# sample (have to squeeze cuz multinomial can only be used on 2D tensor)
replacement = num_samples >= 0
num_samples = abs(num_samples)
return torch.multinomial(logits_BlV.softmax(dim=-1).view(-1, V), num_samples=num_samples, replacement=replacement, generator=rng).view(B, l, num_samples)
def sampling_with_top_k_top_p_also_inplace_modifying_probs_(probs_BlV: torch.Tensor, top_k: int = 0, top_p: float = 0.0, rng=None, num_samples=1) -> torch.Tensor: # return idx, shaped (B, l)
B, l, V = probs_BlV.shape
if top_k > 0:
top_k = min(top_k, V)
idx_to_remove = probs_BlV < probs_BlV.topk(top_k, largest=True, sorted=False, dim=-1)[0].amin(dim=-1, keepdim=True)
probs_BlV.masked_fill_(idx_to_remove, 0)
if top_p > 0:
sorted_probs, sorted_idx = probs_BlV.sort(dim=-1, descending=False)
sorted_idx_to_remove = sorted_probs.softmax(dim=-1).cumsum_(dim=-1) <= (1 - top_p)
sorted_idx_to_remove[..., -1:] = False
probs_BlV.masked_fill_(sorted_idx_to_remove.scatter(sorted_idx.ndim - 1, sorted_idx, sorted_idx_to_remove), 0)
# sample (have to squeeze cuz multinomial can only be used on 2D tensor)
probs_BlV = probs_BlV / probs_BlV.sum(-1, keepdims=True)
replacement = num_samples >= 0
num_samples = abs(num_samples)
return torch.multinomial(probs_BlV.view(-1, V), num_samples=num_samples, replacement=replacement, generator=rng).view(B, l, num_samples)
def get_params_num(d, w, mlp):
m = round(mlp * w / 256) * 256
s = d * (w**2 * 8 + w*m * 2) # sa+ca, mlp
s += w**2 * 6 # saln
s += 4096 * w # pred
s += 32 * w # we
Ct5 = 4096
s += Ct5*w * 4 # T5 attn pool
s += Ct5*w + w*w # T5 mlp
return f'{s/1e9:.2f}B'
TIMM_KEYS = {'img_size', 'pretrained', 'pretrained_cfg', 'pretrained_cfg_overlay', 'global_pool'}
@register_model
def infinity_2b(depth=32, embed_dim=2048, num_heads=2048//128, drop_path_rate=0.1, **kwargs): return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
@register_model
def infinity_20b(depth=58, embed_dim=4608, num_heads=4608//128, drop_path_rate=0.25, **kwargs): return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
# model configuration for scaling Infinity transformer
@register_model
def infinity_layer12(depth=12, embed_dim=768, num_heads=8, drop_path_rate=0.1, **kwargs):
return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
@register_model
def infinity_layer16(depth=16, embed_dim=1152, num_heads=12, drop_path_rate=0.1, **kwargs):
return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
@register_model
def infinity_layer24(depth=24, embed_dim=1536, num_heads=16, drop_path_rate=0.1, **kwargs):
return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
@register_model
def infinity_layer32(depth=32, embed_dim=2080, num_heads=20, drop_path_rate=0.1, **kwargs):
return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
@register_model
def infinity_layer40(depth=40, embed_dim=2688, num_heads=24, drop_path_rate=0.1, **kwargs):
return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})
@register_model
def infinity_layer48(depth=48, embed_dim=3360, num_heads=28, drop_path_rate=0.1, **kwargs):
return Infinity(depth=depth, embed_dim=embed_dim, num_heads=num_heads, mlp_ratio=4, drop_path_rate=drop_path_rate, **{k: v for k, v in kwargs.items() if k not in TIMM_KEYS})