Upload 27 files

tok/__init__.py

@@ -0,0 +1 @@
+from .ar_dtok import *

tok/ar_dtok/__init__.py

@@ -0,0 +1,2 @@
+from .bottleneck import Bottleneck, SimVectorQuantizer
+from .vqvae import VQVAE

tok/ar_dtok/ar_model.py

@@ -0,0 +1,533 @@
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .. import models
+from .generate import generate as ar_generate
+
+
+def find_multiple(n: int, k: int):
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, scale_factor=10000):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    scale_factor: the base for the scaling factor, default is 10000
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / scale_factor**omega  # Parameterized scaling factor (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+@dataclass
+class ModelArgs:
+    dim: int = 4096
+    n_layer: int = 32
+    n_head: int = 32
+
+    n_kv_head: Optional[int] = None
+    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2
+    ffn_dim_multiplier: Optional[float] = None
+    rope_base: float = 10000
+    norm_eps: float = 1e-5
+    initializer_range: float = 0.02
+    
+    token_dropout_p: float = 0.1
+    attn_dropout_p: float = 0.0
+    resid_dropout_p: float = 0.1
+    ffn_dropout_p: float = 0.1
+    drop_path_rate: float = 0.0
+
+    num_classes: int = 1000
+    class_dropout_prob: float = 0.1
+    model_type: str = 'class_cond' # clip_cond, indice_cond
+    cond_dim: int = 1152
+    cond_vocab_size: int = 8192
+
+    vocab_size: int = 8192
+    cls_token_num: int = 1
+
+    max_batch_size: int = 32
+    max_seq_len: int = 2048
+
+    use_fixed_pe: bool = False
+
+    frame_prediction: bool = False
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    @torch.autocast(device_type='cuda', enabled=False)
+    def _norm(self, x):
+        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features, out_features):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=False)
+        self.act = nn.GELU(approximate='tanh')
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=False)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+
+#################################################################################
+#                            Drop Path Implementation                           #
+#################################################################################
+
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(torch.nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f'drop_prob={round(self.drop_prob,3):0.3f}'
+
+
+#################################################################################
+#                                   AR Model                                    #
+#################################################################################
+
+class FeedForward(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        hidden_dim = 4 * config.dim
+        hidden_dim = int(2 * hidden_dim / 3)
+        # custom dim factor multiplier
+        if config.ffn_dim_multiplier is not None:
+            hidden_dim = int(config.ffn_dim_multiplier * hidden_dim)
+        hidden_dim = find_multiple(hidden_dim, config.multiple_of)
+
+        self.w1 = nn.Linear(config.dim, hidden_dim, bias=False)
+        self.w3 = nn.Linear(config.dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, config.dim, bias=False)
+        self.ffn_dropout = nn.Dropout(config.ffn_dropout_p)
+
+    def forward(self, x):
+        return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+    
+
+class KVCache(nn.Module):
+    def __init__(self, max_batch_size, max_seq_length, n_head, head_dim, dtype):
+        super().__init__()
+        cache_shape = (max_batch_size, n_head, max_seq_length, head_dim)
+        self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
+        self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))
+
+    def update(self, input_pos, k_val, v_val):
+        # input_pos: [S], k_val: [B, H, S, D]
+        assert input_pos.shape[0] == k_val.shape[2], f"{input_pos.shape[0]} != {k_val.shape[2]}"
+        k_out = self.k_cache
+        v_out = self.v_cache
+        k_out[:, :, input_pos] = k_val.to(k_out.dtype)
+        v_out[:, :, input_pos] = v_val.to(v_out.dtype)
+
+        return k_out, v_out
+
+
+class Attention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        assert config.dim % config.n_head == 0
+        self.dim = config.dim
+        self.head_dim = config.dim // config.n_head
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head if config.n_kv_head is not None else config.n_head
+        total_kv_dim = (self.n_head + 2 * self.n_kv_head) * self.head_dim
+
+        # key, query, value projections for all heads, but in a batch
+        self.wqkv = nn.Linear(config.dim, total_kv_dim, bias=False)
+        self.wo = nn.Linear(config.dim, config.dim, bias=False)
+        self.kv_cache = None
+
+        # regularization
+        self.attn_dropout_p = config.attn_dropout_p
+        self.resid_dropout = nn.Dropout(config.resid_dropout_p)
+
+    def forward(
+        self, x: torch.Tensor,
+        input_pos: Optional[torch.Tensor] = None, 
+        mask: Optional[torch.Tensor] = None
+    ):
+        bsz, seqlen, _ = x.shape
+        kv_size = self.n_kv_head * self.head_dim
+        xq, xk, xv = self.wqkv(x).split([self.dim, kv_size, kv_size], dim=-1)
+
+        xq = xq.view(bsz, seqlen, self.n_head, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_kv_head, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_kv_head, self.head_dim)
+        
+        xq, xk, xv = map(lambda x: x.transpose(1, 2), (xq, xk, xv))
+
+        if self.kv_cache is not None:
+            keys, values = self.kv_cache.update(input_pos, xk, xv)
+        else:
+            keys, values = xk, xv
+        keys = keys.repeat_interleave(self.n_head // self.n_kv_head, dim=1)
+        values = values.repeat_interleave(self.n_head // self.n_kv_head, dim=1)
+
+        output = F.scaled_dot_product_attention(
+            xq, keys, values, 
+            attn_mask=mask, 
+            is_causal=True if mask is None else False, # is_causal=False is for KV cache
+            dropout_p=self.attn_dropout_p if self.training else 0)            
+        
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)
+
+        output = self.resid_dropout(self.wo(output))
+        return output
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelArgs, drop_path: float):
+        super().__init__()
+        self.attention = Attention(config)
+        self.feed_forward = FeedForward(config)
+        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(
+        self, x: torch.Tensor, start_pos: int, mask: Optional[torch.Tensor] = None):
+        h = x + self.drop_path(self.attention(self.attention_norm(x), start_pos, mask))
+        out = h + self.drop_path(self.feed_forward(self.ffn_norm(h)))
+        return out
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+
+        # replace all negative labels with the last class (unconditional class)
+        labels = torch.where(labels < 0, self.num_classes, labels)
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+class ARModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.n_layer = config.n_layer
+        self.max_seq_length = config.max_seq_len
+        self.num_classes = config.num_classes
+        self.model_type = config.model_type
+        self.cls_token_num = config.cls_token_num
+        self.is_sampling = False
+        self.frame_prediction = config.frame_prediction
+
+        if self.model_type == 'class_cond':
+            self.cls_embedding = LabelEmbedder(config.num_classes, config.dim, config.class_dropout_prob)
+        elif self.model_type == 'clip_cond':
+            self.clip_proj = nn.Linear(config.cond_dim, config.dim)
+        elif self.model_type == 'indice_cond':
+            self.clip_proj = LabelEmbedder(config.cond_vocab_size + 1, config.dim, 0.0)
+        else:
+            raise Exception("please check model type")
+        
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
+        self.tok_dropout = nn.Dropout(config.token_dropout_p)
+
+        # transformer blocks
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.n_layer)]
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(config.n_layer):
+            self.layers.append(TransformerBlock(config, dpr[layer_id]))
+
+        # output layer
+        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+        if config.use_fixed_pe:
+            self.register_buffer('abs_pe', torch.zeros(1, config.max_seq_len + config.cls_token_num - 1, config.dim))
+            abs_pe = get_1d_sincos_pos_embed_from_grid(embed_dim=config.dim, pos=np.arange(config.max_seq_len + config.cls_token_num - 1))
+            self.abs_pe.copy_(torch.from_numpy(abs_pe).float().reshape_as(self.abs_pe))
+            print(f"Using fixed absolute PE")
+        else:
+            self.abs_pe = nn.Parameter(torch.randn(1, config.max_seq_len + config.cls_token_num - 1, config.dim) * 0.02)
+            print(f"Using learned absolute PE")
+
+        self.initialize_weights()
+
+    def initialize_weights(self):        
+        # Initialize nn.Linear and nn.Embedding
+        self.apply(self._init_weights)
+
+        # Zero-out output layers:
+        if hasattr(self.output, 'weight') and isinstance(self.output.weight, nn.Parameter):
+            nn.init.constant_(self.output.weight, 0)
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+    
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+    
+
+    @contextmanager
+    def sampling(self):
+        self.is_sampling = True
+        try:
+            yield
+        finally:
+            self.is_sampling = False
+
+
+    def setup_caches(self, max_batch_size, max_seq_length, dtype):
+        assert max_seq_length == self.max_seq_length + self.cls_token_num, f'{max_seq_length} != {self.max_seq_length} + {self.cls_token_num=}'
+
+        head_dim = self.config.dim // self.config.n_head
+        max_seq_length = find_multiple(max_seq_length, 8)
+
+        for b in self.layers:
+            b.attention.kv_cache = KVCache(max_batch_size, max_seq_length, self.config.n_head, head_dim, dtype)
+
+        causal_mask = torch.tril(torch.ones(max_seq_length, max_seq_length, dtype=torch.bool))
+        self.causal_mask = causal_mask.unsqueeze(0).repeat(max_batch_size, 1, 1)
+
+
+    def reset_caches(self):
+        for b in self.layers:
+            b.attention.kv_cache = None
+
+    def clip_embedding(self, x):
+        if self.model_type == 'clip_cond':
+            if self.training and self.config.class_dropout_prob > 0:
+                drop_ids = torch.rand(x.shape[0], device=x.device) < self.config.class_dropout_prob
+                x[drop_ids] = 0.
+            x = self.clip_proj(x.to(self.dtype)) # Linear
+        elif self.model_type == 'indice_cond':
+            if self.training and self.config.class_dropout_prob > 0:
+                drop_ids = torch.rand(x.shape[0], device=x.device) < self.config.class_dropout_prob
+                x[drop_ids] = self.config.cond_vocab_size
+            x = self.clip_proj(x, train=self.training) # Embedding
+        return x
+
+    def forward(
+        self, 
+        idx: Optional[torch.Tensor], # (b, n)
+        cond_idx: Optional[torch.Tensor],  # cond_idx_or_embed
+        input_pos:  Optional[torch.Tensor] = None, 
+        targets: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        valid: Optional[torch.Tensor] = None,
+    ):
+        if idx is not None and cond_idx is not None: # training or naive inference
+            if self.model_type == 'class_cond':
+                cond_embeddings = self.cls_embedding(cond_idx, train=self.training).unsqueeze(1)[:,:self.cls_token_num]
+            elif self.model_type in ['clip_cond', 'indice_cond']:
+                cond_embeddings = self.clip_embedding(cond_idx)
+            token_embeddings = self.tok_embeddings(idx) # (b, n, d)
+            token_embeddings = torch.cat((cond_embeddings, token_embeddings), dim=1)  # (b, cls_token_num + n, d)
+            h = self.tok_dropout(token_embeddings)
+        else:
+            if cond_idx is not None: # prefill in inference
+                if self.model_type == 'class_cond':
+                    token_embeddings = self.cls_embedding(cond_idx, train=self.training).unsqueeze(1)[:,:self.cls_token_num]
+                elif self.model_type in ['clip_cond', 'indice_cond']:
+                    token_embeddings = self.clip_embedding(cond_idx)
+            else: # decode_n_tokens(kv cache) in inference
+                token_embeddings = self.tok_embeddings(idx)
+            
+            bs = token_embeddings.shape[0]
+            mask = self.causal_mask[:bs, None, input_pos]
+            h = self.tok_dropout(token_embeddings)
+        
+        if self.is_sampling:
+            h = h + self.abs_pe[:, input_pos]
+        else:
+            h = h + self.abs_pe[:, :h.shape[1]]
+        
+        # transformer blocks
+        for layer in self.layers:
+            h = layer(h, input_pos, mask)
+        
+        # output layers
+        h = self.norm(h)
+        logits = self.output(h)
+        # if self.training or self.is_sampling:
+        if cond_idx is not None:
+        # if self.training:
+            # logits = logits[:, self.cls_token_num - 1:].contiguous()
+            logits = logits[:, cond_idx.size(1) - 1:].contiguous()
+
+        # if we are given some desired targets also calculate the loss
+        loss = None
+        if valid is not None:
+            loss_all = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), reduction='none')
+            valid_all = valid[:,None].repeat(1, targets.shape[1]).view(-1)
+            loss = (loss_all * valid_all).sum() / max(valid_all.sum(), 1)
+        elif targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    
+    @torch.inference_mode()
+    def sample(
+        self, 
+        c,
+        cfg_scale=2.0,
+        cfg_interval=-1,
+        temperature=1.0,
+        top_k=0,
+        top_p=1.0,
+        seq_length=None,
+    ):
+        seq_length = self.max_seq_length if seq_length is None else seq_length     
+        with self.sampling():
+            sampled_seqs = ar_generate(
+                self, c, seq_length,
+                cfg_scale=cfg_scale, cfg_interval=cfg_interval,
+                temperature=temperature, top_k=top_k,
+                top_p=top_p, sample_logits=True, 
+            )   
+        return sampled_seqs
+    
+
+    @classmethod
+    def from_checkpoint(cls, ckpt, load_state_dict=True):
+        if isinstance(ckpt, str):
+            assert os.path.exists(ckpt), f"checkpoint {ckpt} does not exist"
+            ckpt = torch.load(ckpt, map_location=lambda storage, loc: storage)
+        else:
+            assert isinstance(
+                ckpt, dict
+            ), f"checkpoint must be a dict or a path to a checkpoint"
+        model = models.make(ckpt["model"], load_sd=load_state_dict)
+        return model
+
+
+#################################################################################
+#                             LLAMA-ABS Configs                                 #
+#################################################################################
+
+def LLAMA_ABS_XXXL(**kwargs):
+    return ARModel(ModelArgs(n_layer=48, n_head=40, dim=2560, **kwargs)) # 3.9B
+
+def LLAMA_ABS_XXL(**kwargs):
+    return ARModel(ModelArgs(n_layer=48, n_head=24, dim=1536, **kwargs)) # 1.4B
+
+def LLAMA_ABS_XL(**kwargs):
+    return ARModel(ModelArgs(n_layer=36, n_head=20, dim=1280, **kwargs)) # 775M
+
+def LLAMA_ABS_LP(**kwargs):
+    return ARModel(ModelArgs(n_layer=30, n_head=20, dim=1280, **kwargs)) # 632M
+
+def LLAMA_ABS_L(**kwargs):
+    return ARModel(ModelArgs(n_layer=24, n_head=16, dim=1024, **kwargs)) # 343M
+
+def LLAMA_ABS_B(**kwargs):
+    return ARModel(ModelArgs(n_layer=12, n_head=12, dim=768, **kwargs)) # 111M
+
+def LLAMA_ABS_S(**kwargs):
+    return ARModel(ModelArgs(n_layer=12, n_head=6, dim=384, **kwargs)) # 21.7M
+
+ar_models = {
+    'llama-abs-S': LLAMA_ABS_S,
+    'llama-abs-B': LLAMA_ABS_B,
+    'llama-abs-L': LLAMA_ABS_L,
+    'llama-abs-LP': LLAMA_ABS_LP,
+    'llama-abs-XL': LLAMA_ABS_XL,
+    'llama-abs-XXL': LLAMA_ABS_XXL,
+    'llama-abs-XXXL': LLAMA_ABS_XXXL,
+}
+
+models.models.update(ar_models)

tok/ar_dtok/bottleneck.py

@@ -0,0 +1,184 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .. import models
+from ..models import register
+
+
+@register("bottleneck")
+class Bottleneck(nn.Module):
+    def __init__(
+        self,
+        bottleneck_dim: int,
+        input_dim: int,
+        output_dim: int,
+        token_nums: int,
+        regularizer=None,
+        **kwargs
+    ):  
+        super().__init__()
+        self.token_nums = token_nums
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        if bottleneck_dim > 0:
+            self.bottleneck_dim = bottleneck_dim
+        else:
+            assert self.input_dim == self.output_dim, "input_dim and output_dim must be the same when bottleneck_dim is not specified"
+            self.bottleneck_dim = self.input_dim
+        
+        self.project_dim = self.bottleneck_dim
+
+        if self.bottleneck_dim > 0:
+            self.in_linear = nn.Linear(self.input_dim, self.project_dim)
+            self.out_linear = nn.Linear(self.bottleneck_dim, self.output_dim)
+        else:
+            self.in_linear = self.out_linear = lambda x: x
+        
+        regularizer['args']['dim'] = self.bottleneck_dim
+        regularizer['args']['token_nums'] = self.token_nums
+        self.regularizer = models.make(regularizer)
+
+    def project_in(self, x):
+        assert len(x.shape) == 3, "Input shape must be (batch, n_tokens, e_dim)"
+        z = self.in_linear(x)
+        return z
+
+    def project_out(self, z_cat):
+        z = self.out_linear(z_cat)
+        return z
+
+    def decode(self, bottleneck_rep):
+        regularized_z = self.regularizer.decode(bottleneck_rep)
+        return self.project_out(regularized_z)
+
+    def forward(self, x):  
+        z = self.project_in(x)
+        projected_z = z
+        regularized_output = self.regularizer(z)
+        x_hat = self.project_out(regularized_output['regularized_z'])
+        bottleneck_rep = regularized_output.pop('bottleneck_rep')
+        return {
+            'output': x_hat,
+            'bottleneck_rep': bottleneck_rep,
+            'projected_z': projected_z,
+            **regularized_output,
+        }
+
+
+@register("simvq")
+class SimVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        codebook_size,
+        l2_normalized=False,
+        same_index_shape=True,
+        stochastic=False,
+        stochastic_temperature=1.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.dim = dim
+        assert isinstance(l2_normalized, bool)
+        self.l2_normalized = l2_normalized
+        self.stochastic = stochastic
+        self.eval_deterministic = False
+        self.default_stochastic_temperature = stochastic_temperature
+        
+        if self.stochastic:
+            if stochastic_temperature > 0: # fixed temperature
+                self.stochastic_temperature_inv = 1 / stochastic_temperature
+            else: # set stochastic_temperature < 0 to use learnable temperature
+                self.stochastic_temperature_inv = nn.Parameter(torch.tensor(10.0))
+
+        # for clear inference code, we remove the codebook init from LLM's embedding
+        self.embedding = nn.Embedding(self.codebook_size, self.dim)
+        self.embedding_proj = nn.Linear(self.dim, self.dim)
+
+        self.same_index_shape = same_index_shape
+
+    def set_eval_deterministic(self, deterministic=True):
+        self.eval_deterministic = deterministic
+
+    def set_stochastic_temperature(self, temperature):
+        self.stochastic_temperature_inv = 1 / temperature
+
+    @torch.autocast(device_type='cuda', enabled=False)
+    def get_emb(self):
+        emb = self.embedding_proj(self.embedding.weight)
+        if self.l2_normalized:
+            emb = F.normalize(emb, p=2, dim=-1)
+        # assert emb.dtype == torch.float32, f"Embedding weight dtype is {emb.dtype}, expected float32"
+        return emb
+
+    @torch.autocast(device_type='cuda', enabled=False)
+    def forward(self, z):
+        emb = self.get_emb()
+        z = z.to(emb)
+        # z = z.float()
+        assert len(z.shape) == 3, "Input shape must be (batch, n_tokens, e_dim)"
+        if self.l2_normalized:
+            z = F.normalize(z, p=2, dim=-1)
+
+        z_flattened = rearrange(z, 'b n d -> (b n) d')
+
+        if self.stochastic:
+            # sample the softmaxed cosine similarity
+            assert self.l2_normalized, "Stochastic sampling requires l2 normalization"
+            cos_sim = torch.einsum("bd,nd->bn", z_flattened, emb)
+            probs = F.softmax(cos_sim * self.stochastic_temperature_inv, dim=-1)
+            if self.eval_deterministic and not self.training:
+                q_indices = torch.argmax(probs, dim=-1)
+            else:
+                q_indices = torch.multinomial(probs, 1).squeeze(-1)
+        else:
+            d = (
+                torch.sum(z_flattened**2, dim=1, keepdim=True)
+                + torch.sum(emb**2, dim=1)
+                - 2
+                * torch.einsum(
+                    "bd,dn->bn", z_flattened, rearrange(emb, "n d -> d n")
+                )
+            )
+            q_indices = torch.argmin(d, dim=1)
+
+        quantized = F.embedding(q_indices, emb, self.embedding.padding_idx, self.embedding.max_norm,
+            self.embedding.norm_type, self.embedding.scale_grad_by_freq, self.embedding.sparse).view(z.shape)  # (b, n, d)
+        
+        # preserve gradients
+        quantized = z + (quantized - z).detach()
+
+        if self.same_index_shape:
+            q_indices = q_indices.reshape(quantized.shape[0], quantized.shape[1])
+
+        return_dict = {
+            'unregularized_z': z, # but l2 normalized if l2_normalized=True
+            'emb': emb, # but l2 normalized if l2_normalized=True
+            'regularized_z': quantized,
+            'bottleneck_rep': q_indices
+        }
+        return return_dict
+    
+    def get_codebook_entry(self, indices, shape=None):
+        # shape specifying (batch, height, width, channel)
+        indices_shape = indices.shape
+        indices_flatten = rearrange(indices, '... -> (...)')
+
+        # get quantized latent vectors
+        emb = self.get_emb()
+        z_q = F.embedding(indices_flatten, emb)
+        # z_q = self.embedding(indices_flatten)
+        if self.l2_normalized:
+            z_q = F.normalize(z_q, p=2, dim=-1)
+
+        if shape is not None:
+            z_q = z_q.reshape(shape)
+        else:
+            z_q = z_q.reshape([*indices_shape, self.dim])
+        return z_q
+
+    def decode(self, indices):
+        return self.get_codebook_entry(indices)

tok/ar_dtok/generate.py

@@ -0,0 +1,188 @@
+# Modified from:
+#   llamagen: https://github.com/FoundationVision/LlamaGen/blob/main/autoregressive/models/generate.py
+#   gpt-fast: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
+#   DiT:      https://github.com/facebookresearch/DiT/blob/main/models.py
+
+
+import torch
+import torch._dynamo.config
+import torch._inductor.config
+from torch.nn import functional as F
+
+
+### from https://huggingface.co/transformers/v3.2.0/_modules/transformers/generation_utils.html
+def top_k_top_p_filtering(
+    logits,
+    top_k: int = 0,
+    top_p: float = 1.0,
+    filter_value: float = -float("Inf"),
+    min_tokens_to_keep: int = 1,
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+    return logits
+
+
+def sample(logits, temperature: float=1.0, top_k: int=0, top_p: float=1.0, sample_logits=True):        
+    logits = logits[:, -1, :] / max(temperature, 1e-5)
+    if top_k > 0 or top_p < 1.0:
+        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+    # improve numerical stability of softmax
+    probs = F.softmax(logits.float(), dim=-1)
+    if sample_logits:
+        idx = torch.multinomial(probs, num_samples=1)
+    else:
+        _, idx = torch.topk(probs, k=1, dim=-1)
+    return idx, probs
+
+
+def logits_to_probs(logits, temperature: float = 1.0, top_p: float=1.0, top_k: int = None, **kwargs):
+    logits = logits / max(temperature, 1e-5)
+    if top_k > 0 or top_p < 1.0:
+        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    return probs
+
+
+def prefill(model, cond_idx: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, **sampling_kwargs):
+    if cfg_scale > 1.0:
+        logits, _ = model(None, cond_idx, input_pos)
+        logits_combined = logits
+        cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0)
+        logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
+    else:
+        logits, _ = model(None, cond_idx, input_pos)
+
+    return sample(logits, **sampling_kwargs)[0]
+
+
+def decode_one_token(model, x: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, cfg_flag: bool, **sampling_kwargs):
+    assert input_pos.shape[-1] == 1
+    if cfg_scale > 1.0:
+        x_combined = torch.cat([x, x])
+        logits, _ = model(x_combined, cond_idx=None, input_pos=input_pos)
+        logits_combined = logits
+        cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0) 
+        if cfg_flag:
+            logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
+        else:
+            logits = cond_logits
+    else:
+        logits, _ = model(x, cond_idx=None, input_pos=input_pos)
+    return sample(logits, **sampling_kwargs)
+
+
+def decode_n_tokens(
+    model, cur_token: torch.Tensor, input_pos: torch.Tensor, num_new_tokens: int, 
+    cfg_scale: float, cfg_interval: int,
+    **sampling_kwargs):
+    new_tokens, new_probs = [], []
+    cfg_flag = True
+    for i in range(num_new_tokens):
+        with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): # Actually better for Inductor to codegen attention here
+            if cfg_interval > -1 and i > cfg_interval:
+                cfg_flag = False
+            next_token, next_prob = decode_one_token(
+                model, cur_token, input_pos, cfg_scale, cfg_flag, **sampling_kwargs
+            )
+            input_pos += 1
+            new_tokens.append(next_token.clone())
+            new_probs.append(next_prob.clone())
+            cur_token = next_token.view(-1, 1)
+    
+    return new_tokens, new_probs
+
+
+@torch.no_grad()
+def generate(model, cond, max_new_tokens, emb_masks=None, cfg_scale=1.0, cfg_interval=-1, **sampling_kwargs):
+    if model.frame_prediction:
+        assert cfg_scale == 1.0, "frame prediction requires cfg_scale=1.0 (no classifier-free guidance)"
+        cond_combined = cond
+        T = cond.shape[1]
+    elif model.model_type == 'class_cond':
+        if cfg_scale > 1.0:
+            cond_null = torch.ones_like(cond) * model.num_classes
+            cond_combined = torch.cat([cond, cond_null])
+        else:
+            cond_combined = cond
+        T = 1 
+    elif model.model_type == 'clip_cond':
+        if cfg_scale > 1.0:
+            cond_null = torch.zeros_like(cond)
+            cond_combined = torch.cat([cond, cond_null])
+        else:
+            cond_combined = cond
+        T = model.cls_token_num
+    elif model.model_type == 'indice_cond':
+        if cfg_scale > 1.0:
+            cond_null = torch.ones_like(cond) * model.cond_vocab_size
+            cond_combined = torch.cat([cond, cond_null])
+        else:
+            cond_combined = cond
+        T = model.cls_token_num
+    else:
+        raise Exception("please check model type")
+
+    T_new = T + max_new_tokens
+    max_seq_length = T_new
+    max_batch_size = cond.shape[0]
+
+    device = cond.device
+    with torch.device(device):
+        max_batch_size_cfg = max_batch_size * 2 if cfg_scale > 1.0 else max_batch_size
+        model.setup_caches(max_batch_size=max_batch_size_cfg, max_seq_length=max_seq_length, dtype=model.tok_embeddings.weight.dtype)
+    
+    if emb_masks is not None:
+        assert emb_masks.shape[0] == max_batch_size
+        assert emb_masks.shape[-1] == T
+        if cfg_scale > 1.0:
+            model.causal_mask[:, :, :T] = model.causal_mask[:, :, :T] * torch.cat([emb_masks, emb_masks]).unsqueeze(1)
+        else:
+            model.causal_mask[:, :, :T] = model.causal_mask[:, :, :T] * emb_masks.unsqueeze(1)
+
+        eye_matrix = torch.eye(model.causal_mask.size(1), model.causal_mask.size(2), device=device)
+        model.causal_mask[:] = model.causal_mask * (1 - eye_matrix) + eye_matrix
+    
+    # create an empty tensor of the expected final shape and fill in the current tokens
+    seq = torch.empty((max_batch_size, T_new), dtype=torch.int, device=device)
+
+    input_pos = torch.arange(0, T, device=device)
+
+    next_token = prefill(model, cond_combined, input_pos, cfg_scale, **sampling_kwargs)
+    seq[:, T:T+1] = next_token
+
+    input_pos = torch.tensor([T], device=device, dtype=torch.int)
+    generated_tokens, _ = decode_n_tokens(model, next_token, input_pos, max_new_tokens-1, cfg_scale, cfg_interval, **sampling_kwargs)
+    seq[:, T+1:] = torch.cat(generated_tokens, dim=1)
+
+    return seq[:, T:]

tok/ar_dtok/vqvae.py

@@ -0,0 +1,457 @@
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..models import register
+from ..utils import ScalingLayer
+
+
+@register('vqvae')
+class VQVAE(nn.Module):
+    def __init__(
+        self,
+        model='VQ-16',
+        ckpt='',
+        codebook_size=16384,
+        codebook_embed_dim=8,
+        bottleneck_token_num=256,
+        input_size=256,
+        *args,
+        **kwargs,
+    ):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_embed_dim = codebook_embed_dim
+        self.bottleneck_token_num = bottleneck_token_num
+        self.input_size = input_size
+        self.model = VQ_models[model](
+            codebook_size=codebook_size,
+            codebook_embed_dim=codebook_embed_dim)
+        ckpt = torch.load(ckpt, map_location='cpu')
+        self.model.load_state_dict(ckpt['model'])
+        self.model.eval()
+
+        self.scale_layer = ScalingLayer(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    @classmethod
+    def from_checkpoint(cls, ckpt, **kwargs):
+        model = cls(ckpt=ckpt, **kwargs)
+        return model
+
+    def decode_from_bottleneck(self, z):
+        if z.ndim == 2:
+            b = z.size(0)
+            h = w = int(z.size(-1) ** 0.5)
+        z = self.model.decode_code(z, (b, self.codebook_embed_dim, h, w))
+        return self.scale_layer.inv(z)
+
+
+# Adapt from https://github.com/FoundationVision/LlamaGen/blob/main/tokenizer/tokenizer_image/vq_model.py
+@dataclass
+class ModelArgs:
+    codebook_size: int = 16384
+    codebook_embed_dim: int = 8
+    codebook_l2_norm: bool = True
+    codebook_show_usage: bool = True
+    commit_loss_beta: float = 0.25
+    entropy_loss_ratio: float = 0.0
+    
+    encoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4])
+    decoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4])
+    z_channels: int = 256
+    dropout_p: float = 0.0
+
+
+class VQModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.encoder = Encoder(ch_mult=config.encoder_ch_mult, z_channels=config.z_channels, dropout=config.dropout_p)
+        self.decoder = Decoder(ch_mult=config.decoder_ch_mult, z_channels=config.z_channels, dropout=config.dropout_p)
+
+        self.quantize = VectorQuantizer(config.codebook_size, config.codebook_embed_dim, 
+                                        config.commit_loss_beta, config.entropy_loss_ratio,
+                                        config.codebook_l2_norm, config.codebook_show_usage)
+        self.quant_conv = nn.Conv2d(config.z_channels, config.codebook_embed_dim, 1)
+        self.post_quant_conv = nn.Conv2d(config.codebook_embed_dim, config.z_channels, 1)
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, info = self.quantize(h)
+        return quant, emb_loss, info
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b, shape=None, channel_first=True):
+        quant_b = self.quantize.get_codebook_entry(code_b, shape, channel_first)
+        dec = self.decode(quant_b)
+        return dec
+
+    def forward(self, input):
+        quant, diff, _ = self.encode(input)
+        dec = self.decode(quant)
+        return dec, diff
+
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels=3, ch=128, ch_mult=(1,1,2,2,4), num_res_blocks=2, 
+                 norm_type='group', dropout=0.0, resamp_with_conv=True, z_channels=256):
+        super().__init__()
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+        # downsampling
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.conv_blocks = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            conv_block = nn.Module()
+            # res & attn
+            res_block = nn.ModuleList()
+            attn_block = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                res_block.append(ResnetBlock(block_in, block_out, dropout=dropout, norm_type=norm_type))
+                block_in = block_out
+                if i_level == self.num_resolutions - 1:
+                    attn_block.append(AttnBlock(block_in, norm_type))
+            conv_block.res = res_block
+            conv_block.attn = attn_block
+            # downsample
+            if i_level != self.num_resolutions-1:
+                conv_block.downsample = Downsample(block_in, resamp_with_conv)
+            self.conv_blocks.append(conv_block)
+
+        # middle
+        self.mid = nn.ModuleList()
+        self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+        self.mid.append(AttnBlock(block_in, norm_type=norm_type))
+        self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+
+        # end
+        self.norm_out = Normalize(block_in, norm_type)
+        self.conv_out = nn.Conv2d(block_in, z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        h = self.conv_in(x)
+        # downsampling
+        for i_level, block in enumerate(self.conv_blocks):
+            for i_block in range(self.num_res_blocks):
+                h = block.res[i_block](h)
+                if len(block.attn) > 0:
+                    h = block.attn[i_block](h)
+            if i_level != self.num_resolutions - 1:
+                h = block.downsample(h)
+        
+        # middle
+        for mid_block in self.mid:
+            h = mid_block(h)
+        
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, z_channels=256, ch=128, ch_mult=(1,1,2,2,4), num_res_blocks=2, norm_type="group",
+                 dropout=0.0, resamp_with_conv=True, out_channels=3):
+        super().__init__()
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+
+        block_in = ch*ch_mult[self.num_resolutions-1]
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+       # middle
+        self.mid = nn.ModuleList()
+        self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+        self.mid.append(AttnBlock(block_in, norm_type=norm_type))
+        self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+
+        # upsampling
+        self.conv_blocks = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            conv_block = nn.Module()
+            # res & attn
+            res_block = nn.ModuleList()
+            attn_block = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                res_block.append(ResnetBlock(block_in, block_out, dropout=dropout, norm_type=norm_type))
+                block_in = block_out
+                if i_level == self.num_resolutions - 1:
+                    attn_block.append(AttnBlock(block_in, norm_type))
+            conv_block.res = res_block
+            conv_block.attn = attn_block
+            # downsample
+            if i_level != 0:
+                conv_block.upsample = Upsample(block_in, resamp_with_conv)
+            self.conv_blocks.append(conv_block)
+
+        # end
+        self.norm_out = Normalize(block_in, norm_type)
+        self.conv_out = nn.Conv2d(block_in, out_channels, kernel_size=3, stride=1, padding=1)
+
+    @property
+    def last_layer(self):
+        return self.conv_out.weight
+    
+    def forward(self, z):
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        for mid_block in self.mid:
+            h = mid_block(h)
+        
+        # upsampling
+        for i_level, block in enumerate(self.conv_blocks):
+            for i_block in range(self.num_res_blocks + 1):
+                h = block.res[i_block](h)
+                if len(block.attn) > 0:
+                    h = block.attn[i_block](h)
+            if i_level != self.num_resolutions - 1:
+                h = block.upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_e, e_dim, beta, entropy_loss_ratio, l2_norm, show_usage):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.entropy_loss_ratio = entropy_loss_ratio
+        self.l2_norm = l2_norm
+        self.show_usage = show_usage
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+        if self.l2_norm:
+            self.embedding.weight.data = F.normalize(self.embedding.weight.data, p=2, dim=-1)
+        if self.show_usage:
+            self.register_buffer("codebook_used", nn.Parameter(torch.zeros(65536)))
+
+    def forward(self, z):
+        # reshape z -> (batch, height, width, channel) and flatten
+        z = torch.einsum('b c h w -> b h w c', z).contiguous()
+        z_flattened = z.view(-1, self.e_dim)
+        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+
+        if self.l2_norm:
+            z = F.normalize(z, p=2, dim=-1)
+            z_flattened = F.normalize(z_flattened, p=2, dim=-1)
+            embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+        else:
+            embedding = self.embedding.weight
+
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(embedding**2, dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', z_flattened, torch.einsum('n d -> d n', embedding))
+
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = embedding[min_encoding_indices].view(z.shape)
+        perplexity = None
+        min_encodings = None
+        vq_loss = None
+        commit_loss = None
+        entropy_loss = None
+        codebook_usage = 0
+
+        if self.show_usage and self.training:
+            cur_len = min_encoding_indices.shape[0]
+            self.codebook_used[:-cur_len] = self.codebook_used[cur_len:].clone()
+            self.codebook_used[-cur_len:] = min_encoding_indices
+            codebook_usage = len(torch.unique(self.codebook_used)) / self.n_e
+
+        # compute loss for embedding
+        if self.training:
+            vq_loss = torch.mean((z_q - z.detach()) ** 2) 
+            commit_loss = self.beta * torch.mean((z_q.detach() - z) ** 2) 
+            entropy_loss = self.entropy_loss_ratio * compute_entropy_loss(-d)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        z_q = torch.einsum('b h w c -> b c h w', z_q)
+
+        return z_q, (vq_loss, commit_loss, entropy_loss, codebook_usage), (perplexity, min_encodings, min_encoding_indices)
+
+    def get_codebook_entry(self, indices, shape=None, channel_first=True):
+        # shape = (batch, channel, height, width) if channel_first else (batch, height, width, channel)
+        if self.l2_norm:
+            embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+        else:
+            embedding = self.embedding.weight
+        z_q = embedding[indices]  # (b*h*w, c)
+
+        if shape is not None:
+            if channel_first:
+                z_q = z_q.reshape(shape[0], shape[2], shape[3], shape[1])
+                # reshape back to match original input shape
+                z_q = z_q.permute(0, 3, 1, 2).contiguous()
+            else:
+                z_q = z_q.view(shape)
+        return z_q
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels=None, conv_shortcut=False, dropout=0.0, norm_type='group'):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels, norm_type)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = Normalize(out_channels, norm_type)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+            else:
+                self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x+h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels, norm_type='group'):
+        super().__init__()
+        self.norm = Normalize(in_channels, norm_type)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = q.reshape(b,c,h*w)
+        q = q.permute(0,2,1)   # b,hw,c
+        k = k.reshape(b,c,h*w) # b,c,hw
+        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = F.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b,c,h*w)
+        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b,c,h,w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+def nonlinearity(x):
+    # swish
+    return x*torch.sigmoid(x)
+
+def Normalize(in_channels, norm_type='group'):
+    assert norm_type in ['group', 'batch']
+    if norm_type == 'group':
+        return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+    elif norm_type == 'batch':
+        return nn.SyncBatchNorm(in_channels)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0,1,0,1)
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+def compute_entropy_loss(affinity, loss_type="softmax", temperature=0.01):
+    flat_affinity = affinity.reshape(-1, affinity.shape[-1])
+    flat_affinity /= temperature
+    probs = F.softmax(flat_affinity, dim=-1)
+    log_probs = F.log_softmax(flat_affinity + 1e-5, dim=-1)
+    if loss_type == "softmax":
+        target_probs = probs
+    else:
+        raise ValueError("Entropy loss {} not supported".format(loss_type))
+    avg_probs = torch.mean(target_probs, dim=0)
+    avg_entropy = - torch.sum(avg_probs * torch.log(avg_probs + 1e-5))
+    sample_entropy = - torch.mean(torch.sum(target_probs * log_probs, dim=-1))
+    loss = sample_entropy - avg_entropy
+    return loss
+
+
+#################################################################################
+#                              VQ Model Configs                                 #
+#################################################################################
+def VQ_8(**kwargs):
+    return VQModel(ModelArgs(encoder_ch_mult=[1, 2, 2, 4], decoder_ch_mult=[1, 2, 2, 4], **kwargs))
+
+def VQ_16(**kwargs):
+    return VQModel(ModelArgs(encoder_ch_mult=[1, 1, 2, 2, 4], decoder_ch_mult=[1, 1, 2, 2, 4], **kwargs))
+
+VQ_models = {'VQ-16': VQ_16, 'VQ-8': VQ_8}

tok/mm_autoencoder.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+
+from tok.ar_dtok.ar_model import ARModel
+from tok.ar_dtok.vqvae import VQVAE
+from tok.ta_tok import TextAlignedTokenizer
+
+
+class MMAutoEncoder(nn.Module):
+    def __init__(self,
+        ar_path,
+        encoder_path, decoder_path, 
+        encoder_args={}, decoder_args={}):
+        super().__init__()
+        self.ar_model = ARModel.from_checkpoint(ar_path)
+
+        self.encoder = TextAlignedTokenizer.from_checkpoint(encoder_path, load_teacher=False, **encoder_args)
+        self.decoder = VQVAE.from_checkpoint(decoder_path, **decoder_args)
+
+    def ar_sample(self, x, args):
+        x = self.ar_model.sample(
+            x,
+            cfg_scale=args.get('cfg_scale', 1.0),
+            cfg_interval=args.get('cfg_interval', -1),
+            temperature=args.get('temperature', 1.0),
+            top_k=args.get('top_k', 0),
+            top_p=args.get('top_p', 1.0)
+        )
+        return x
+
+    def post_process(self, x):
+        x = x.cpu().float().clamp(0., 1.) * 255.
+        x = x.permute(0, 2, 3, 1) # [b, h, w, c]
+        x = x.to(torch.uint8)
+        return x
+    
+    def encode(self, x):
+        return self.encoder(x.to(self.encoder.dtype))['encoded']
+    
+    def get_encoder_indices(self, x):
+        # img -> encoder -> indices
+        return self.encoder(x.to(self.encoder.dtype))['bottleneck_rep']
+    
+    @torch.inference_mode()
+    def decode_from_encoder_indices(self, indices, args={}):
+        # indices -> encoder feats -> ar -> decoder
+        encoder_x = self.encoder.decode_from_bottleneck(indices)
+        ar_indices = self.ar_sample(encoder_x, args)
+        decoder_x = self.decoder.decode_from_bottleneck(ar_indices)
+        x = self.post_process(decoder_x)
+        return x
+    
+    def decode_from_vqvae_indices(self, indices):
+        decoder_x = self.decoder.decode_from_bottleneck(indices)
+        x = self.post_process(decoder_x)
+        return x
+    
+    @torch.inference_mode()
+    def forward(self, x, args={}):
+        encoder_x = self.encoder(x.to(self.encoder.dtype))['encoded']
+        ar_indices = self.ar_sample(encoder_x, args)
+        decoder_x = self.decoder.decode_from_bottleneck(ar_indices)
+        x = self.post_process(decoder_x)
+        return x

tok/models.py

@@ -0,0 +1,31 @@
+import copy
+import inspect
+
+import torch
+
+models = {}
+
+
+def register(name):
+    def decorator(cls):
+        models[name] = cls
+        return cls
+    return decorator
+
+
+def make(model_spec, args=None, load_sd=False) -> torch.nn.Module:
+    if args is not None:
+        model_args = copy.deepcopy(model_spec['args'])
+        model_args.update(args)
+    else:
+        model_args = model_spec['args']
+    model_params = inspect.signature(models[model_spec['name']]).parameters
+    if 'kwargs' not in model_params:
+        model_args = {k: v for k, v in model_args.items() if k in model_params}
+    model = models[model_spec['name']](**model_args)
+    if load_sd:
+        if ('abs_pe' in  model_spec['sd']) and hasattr(model, 'abs_pe') and model_spec['sd']['abs_pe'].shape != model.abs_pe.shape:
+            del model_spec['sd']['abs_pe']
+        msg = model.load_state_dict(model_spec['sd'], strict=False)
+        print(msg)
+    return model

tok/ta_tok.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torchvision.transforms import Resize
+from transformers import AutoConfig, AutoModel, Siglip2VisionConfig, Siglip2VisionModel
+
+from . import models
+from .utils import ScalingLayer
+
+from huggingface_hub import hf_hub_download
+
+ckpt_path = hf_hub_download(
+    repo_id="csuhan/TA-Tok",
+    filename="ta_tok.pth",
+    repo_type="model"     
+)
+
+
+class TextAlignedTokenizer(nn.Module):
+    def __init__(
+        self, 
+        bottleneck,
+        bottleneck_token_num=256,
+        input_size=384,
+        teacher='google/siglip2-so400m-patch14-384',
+        input_type='quant', # choose from ['quant', 'rec', 'indices']
+        pool_scale=1, # choose from [1, 2, 3]
+        decoder_depth=3,
+        select_layer_id=-2,
+        *args,
+        **kwargs
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.bottleneck_token_num = bottleneck_token_num
+        self.teacher = teacher
+        self.input_type = input_type
+        self.pool_scale = pool_scale
+        self.decoder_depth = decoder_depth
+        self.select_layer_id = select_layer_id
+       
+        self.bottleneck_dim = bottleneck['args']['bottleneck_dim']
+
+        self.encoder_config = AutoConfig.from_pretrained(teacher)
+        self.encoder = AutoModel.from_config(self.encoder_config).vision_model         
+        
+        self.encoder_hidden_dim = self.encoder.config.hidden_size
+
+        self.decoder_config = Siglip2VisionConfig()
+        self.decoder_config.update({
+            'patch_size': 1,
+            'num_hidden_layers': self.decoder_depth,
+            'num_channels': self.bottleneck_dim,
+            'hidden_size': self.encoder_hidden_dim,
+        })
+        self.decoder = Siglip2VisionModel(self.decoder_config)
+
+        self.encode_task_layer = nn.Sequential(
+            nn.Linear(self.encoder_hidden_dim, self.encoder_hidden_dim),
+            nn.Tanh())
+        self.decode_task_layer = nn.Sequential(
+            nn.Linear(self.encoder_hidden_dim, self.encoder_hidden_dim),
+            nn.Tanh(),
+            nn.Linear(self.encoder_hidden_dim, self.encoder_hidden_dim))
+
+        bottleneck_args = {
+            'token_nums': self.bottleneck_token_num, 
+            'input_dim': self.encoder_hidden_dim, 
+            'output_dim': self.bottleneck_dim}
+        self.bottleneck = models.make(bottleneck, args=bottleneck_args)
+
+        self.scale_layer = ScalingLayer(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])   
+        self.image_resize = Resize((self.input_size, self.input_size))
+       
+    def set_vq_eval_deterministic(self, deterministic=True):
+        self.bottleneck.regularizer.set_eval_deterministic(deterministic)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+    
+    @classmethod
+    def from_checkpoint(cls, ckpt, load_teacher=True, **kwargs):
+        ckpt = torch.load(ckpt_path, map_location='cpu')
+        ckpt_kwargs = ckpt["model"]["args"]
+        model = cls(**kwargs, **ckpt_kwargs)
+        sd = ckpt["model"]["sd"]
+        if not load_teacher:
+            sd = {k: v for k, v in sd.items() if not k.startswith('teacher')}
+        model.load_state_dict(sd, strict=True)
+        return model
+
+    def encode(self, x, **kwargs):
+        if x.ndim == 5:
+            x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.scale_layer(x)
+        if tuple(x.shape[-2:]) != (self.input_size, self.input_size):
+            x = self.image_resize(x)
+        vq_feats = self.encoder(x, output_hidden_states=True).hidden_states[self.select_layer_id]
+
+        pool_scale = self.pool_scale
+        pool_scale = kwargs.get("pool_scale", pool_scale)
+        if pool_scale != 1:
+            vq_feats = self.avg_pool(vq_feats, pool_scale)
+        vq_feats = self.encode_task_layer(vq_feats.to(x))
+        
+        bottleneck_out = self.bottleneck(vq_feats)
+        z = bottleneck_out.pop('output')
+
+        return {'encoded': z, 'pool_scale': pool_scale, 'vq_feats': vq_feats, **bottleneck_out}
+
+    def avg_pool(self, z, pool_scale=1):
+        if z.ndim == 3:
+            b, n, c = z.shape
+            p = int(n ** 0.5)
+            z = rearrange(z, 'b (p1 p2) c -> b c p1 p2', p1=p, p2=p)
+        else:
+            b, c, p, _ = z.shape
+        p_s = int(p // pool_scale)
+        z = F.avg_pool2d(
+            z,
+            kernel_size=(pool_scale, pool_scale),
+            stride=(pool_scale, pool_scale)
+        ).contiguous()
+        z = rearrange(z, 'b c p1 p2 -> b (p1 p2) c')
+        return z
+
+    def decode(self, z):
+        if z.ndim == 4:
+            z = rearrange(z, 'b c p1 p2 -> b (p1 p2) c')
+        attention_mask = torch.ones(z.shape[:2], dtype=torch.int, device=z.device)
+        p = int(z.shape[1]**0.5)
+        spatial_shape = torch.tensor([[p, p]]*z.shape[0], device=self.device)
+        z = self.decoder(z, attention_mask, spatial_shape, output_hidden_states=True).last_hidden_state
+        z = self.decode_task_layer(z)
+        return z
+
+    def decode_from_bottleneck(self, bottleneck_rep):
+        z = self.bottleneck.decode(bottleneck_rep) # (b, n, c)
+        p = int(z.shape[1]**0.5)
+        z = rearrange(z, 'b (p1 p2) c -> b c p1 p2', p1=p, p2=p)
+        return self.decode(z)
+
+    def forward(self, data, **kwargs):
+        # data: video in shape (b, c, t, h, w)
+        encode_output = self.encode(data, **kwargs)
+        vq_feats = encode_output['encoded']
+        p = int(vq_feats.shape[1] ** 0.5)
+        vq_feats = rearrange(vq_feats, 'b (h w) c -> b c h w', h=p, w=p)
+        pred_feats = self.decode(vq_feats)
+
+        if self.input_type == 'quant':
+            z = encode_output["regularized_z"] # [b, n, c]
+        elif self.input_type == 'indices':
+            z = encode_output["bottleneck_rep"] # [b, n]
+        elif self.input_type == 'rec':
+            z = pred_feats # [b, n, c]
+        encode_output['encoded'] = z
+        return encode_output

tok/utils.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
+        super().__init__()
+        self.register_buffer('shift', torch.Tensor(mean)[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor(std)[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+    
+    def inv(self, inp):
+        return inp * self.scale + self.shift