Add pipeline

arpg.py

@@ -0,0 +1,636 @@
+# Modified from:
+#   LlamaGen:   https://github.com/FoundationVision/LlamaGen/
+#   YOCO:       https://github.com/microsoft/unilm/tree/master/YOCO
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torch.nn import functional as F
+from einops import rearrange
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+from transformers.configuration_utils import PretrainedConfig
+
+
+def find_multiple(n: int, k: int):
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+def batch_seq_shuffle(x, orders=None):
+    assert x.ndim >= 2, "The input should contain at least two dimensions, batch and length"
+    bs, seq_len = x.shape[:2]
+    
+    if orders is None:
+        orders = torch.rand(bs, seq_len, device=x.device).argsort(dim=1)
+    
+    orders_expand = orders.view(*orders.shape, *(1,) * (x.ndim - orders.ndim))
+    shuffled_data = torch.gather(x, 1, orders_expand.expand(*x.shape))
+    
+    return shuffled_data, orders
+
+
+# @dataclass
+class ModelArgs(PretrainedConfig):
+    def __init__(
+        self,
+        dim: int = 4096,
+        n_layer: int = 32,
+        n_head: int = 32,
+        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 = 'c2i',
+        vocab_size: int = 16384,
+        cls_token_num: int = 1,
+        block_size: int = 256,
+    ):
+        self.dim = dim
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.multiple_of = multiple_of
+        self.ffn_dim_multiplier = ffn_dim_multiplier
+        self.rope_base = rope_base
+        self.norm_eps = norm_eps
+        self.initializer_range = initializer_range
+            
+        self.token_dropout_p = token_dropout_p
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout_p = resid_dropout_p
+        self.ffn_dropout_p = ffn_dropout_p
+        self.drop_path_rate = drop_path_rate
+
+        self.num_classes = num_classes
+        self.class_dropout_prob = class_dropout_prob
+        self.model_type = model_type
+        self.vocab_size = vocab_size
+        self.cls_token_num = cls_token_num
+        self.block_size = block_size
+
+
+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))
+
+    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 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 Attention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        assert config.dim % config.n_head == 0
+        self.dim = config.dim
+        self.n_head = config.n_head
+        self.head_dim = config.dim // config.n_head
+
+        self.to_q = nn.Linear(config.dim, config.dim, bias=False)
+        self.to_k = nn.Linear(config.dim, config.dim, bias=False)
+        self.to_v = nn.Linear(config.dim, config.dim, bias=False)
+        
+        self.proj = nn.Linear(config.dim, config.dim, bias=False)
+
+        self.attn_drop = config.attn_dropout_p
+        self.proj_drop = nn.Dropout(config.resid_dropout_p)
+        
+        self.kv_cache = False
+        self.k_cache = None
+        self.v_cache = None
+        
+    def reset_kv_cache(self):
+        self.k_cache = None
+        self.v_cache = None
+        
+    def update_kv_cache(self, k: torch.Tensor, v: torch.Tensor):
+        if self.k_cache is None and self.v_cache is None:
+            k_cache = k
+            v_cache = v
+        else:
+            k_cache = torch.cat([self.k_cache, k], dim=-2)
+            v_cache = torch.cat([self.v_cache, v], dim=-2)
+
+        self.k_cache = k_cache
+        self.v_cache = v_cache
+
+        return k_cache, v_cache
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor = None
+    ):
+
+        q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=self.n_head), (q, k, v))
+        
+        q = apply_rotary_emb(q, freqs_cis)
+        k = apply_rotary_emb(k, freqs_cis)
+
+        q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+        if self.kv_cache:
+            k, v = self.update_kv_cache(k, v)
+            
+        output = F.scaled_dot_product_attention(
+            q, k, v, 
+            attn_mask=None, 
+            is_causal=True if self.training else False,
+            dropout_p=self.attn_drop if self.training else 0
+        )            
+        output = rearrange(output, 'b h n d -> b n (h d)').contiguous()
+        output = self.proj_drop(self.proj(output))
+        return output
+    
+    
+class CrossAttention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        assert config.dim % config.n_head == 0
+        self.dim = config.dim
+        self.n_head = config.n_head
+        self.head_dim = config.dim // config.n_head
+
+        self.to_q = nn.Linear(config.dim, config.dim, bias=False)
+        
+        self.proj = nn.Linear(config.dim, config.dim, bias=False)
+
+        self.attn_drop = config.attn_dropout_p
+        self.proj_drop = nn.Dropout(config.resid_dropout_p)
+        
+        self.kv_cache = False
+        self.k_cache = None
+        self.v_cache = None
+        
+    def reset_kv_cache(self):
+        self.k_cache = None
+        self.v_cache = None
+        
+    def update_kv_cache(self, k: torch.Tensor, v: torch.Tensor):
+        if self.k_cache is None and self.v_cache is None:
+            k_cache = k
+            v_cache = v
+        else:
+            k_cache = torch.cat([self.k_cache, k], dim=-2)
+            v_cache = torch.cat([self.v_cache, v], dim=-2)
+
+        self.k_cache = k_cache
+        self.v_cache = v_cache
+
+        return k_cache, v_cache
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        freqs_cis: torch.Tensor = None
+    ):
+        q = self.to_q(x)
+        q = rearrange(q, 'b n (h d) -> b n h d', h=self.n_head)
+        
+        # target-aware
+        q = apply_rotary_emb(q, freqs_cis[:, -q.shape[1]:, ...])
+
+        q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+        if self.kv_cache:
+            k, v = self.update_kv_cache(k, v)
+            
+        output = F.scaled_dot_product_attention(
+            q, k, v, 
+            attn_mask=None, 
+            is_causal=True if self.training else False,
+            dropout_p=self.attn_drop if self.training else 0
+        )            
+        output = rearrange(output, 'b h n d -> b n (h d)').contiguous()
+        output = self.proj_drop(self.proj(output))
+        return output
+
+
+class SelfDecoder(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.attn = Attention(config)
+        self.ffn = FeedForward(config)
+        
+        self.attn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor = None
+    ):
+        h = x + self.attn(x=self.attn_norm(x), freqs_cis=freqs_cis[:, :x.shape[1], ...])
+        out = h + self.ffn(self.ffn_norm(h))
+        
+        return out
+    
+
+class CrossDecoder(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.attn = CrossAttention(config)
+        self.ffn = FeedForward(config)
+        
+        self.attn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        
+    def forward(
+        self,
+        x: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        freqs_cis: torch.Tensor = None
+    ):
+        h = x + self.attn(x=self.attn_norm(x), k=k, v=v, freqs_cis=freqs_cis)
+        out = h + self.ffn(self.ffn_norm(h))
+        
+        return out
+    
+    
+class Decoder_Decoder(nn.Module):
+    def __init__(self, config: ModelArgs, n_layer):
+        super().__init__()
+        self.config = config
+        self.self_dec = nn.ModuleList([SelfDecoder(config) for _ in range(n_layer//2)])
+        self.cross_dec = nn.ModuleList([CrossDecoder(config) for _ in range(n_layer//2)])
+        
+        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.to_k = nn.Linear(config.dim, config.dim, bias=False)
+        self.to_v = nn.Linear(config.dim, config.dim, bias=False)
+        
+        self.kv_cache = False
+        self.k_cache = None
+        self.v_cache = None
+        
+    def reset_kv_cache(self):
+        self.k_cache = None
+        self.v_cache = None
+        
+    def update_kv_cache(self, k: torch.Tensor, v: torch.Tensor, head_first=False):
+        t_dim = 2 if head_first else 1
+        
+        if self.k_cache is None and self.v_cache is None:
+            k_cache = k
+            v_cache = v
+        else:
+            k_cache = torch.cat([self.k_cache, k], dim=t_dim)
+            v_cache = torch.cat([self.v_cache, v], dim=t_dim)
+
+        self.k_cache = k_cache
+        self.v_cache = v_cache
+
+        return k_cache, v_cache
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        q: torch.Tensor,
+        freqs_cis: torch.Tensor = None
+    ):
+        for layer in self.self_dec:
+            x = layer(x=x, freqs_cis=freqs_cis)
+            
+        x_norm = self.norm(x)
+        k = self.to_k(x_norm)
+        v = self.to_v(x_norm)
+
+        k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=self.config.n_head), (k, v))
+        k = apply_rotary_emb(k, freqs_cis[:, :k.shape[1], ...])
+
+        if self.kv_cache:
+            k, v = self.update_kv_cache(k, v)
+        
+        for layer in self.cross_dec:
+            q = layer(x=q, k=k, v=v, freqs_cis=freqs_cis)
+            
+        return q
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.image_seq_len = config.block_size
+        
+        """
+        ref: https://github.com/bytedance/1d-tokenizer/blob/main/modeling/rar.py
+        Token space:
+            [0, vocab_size - 1]                         : those are the learned quantized image tokens
+            [vocab_size]                                : the mask token id
+            [vocab_size + 1, vocab_size + num_classes]  : the imagenet class tokens
+            [vocab_size + num_classes + 1]              : the class drop label
+            [vocab_size + num_classes + 2]              : the drop token for scg
+        """
+        self.embeddings = nn.Embedding(config.vocab_size + 1 + config.num_classes + 1 + 1, config.dim)
+        self.embed_drop = nn.Dropout(config.token_dropout_p)
+        
+        self.mask_token_id = config.vocab_size
+        self.none_conds_id = config.vocab_size + config.num_classes + 1
+        self.none_token_id = config.vocab_size + config.num_classes + 2
+        
+        # 2-pass decoder
+        self.layers = Decoder_Decoder(config, config.n_layer)
+
+        # output layer
+        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.head = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+        # 2d rotary pos embedding
+        grid_size = int(self.image_seq_len ** 0.5)
+        self.freqs_cis = precompute_freqs_cis_2d(grid_size, config.dim // config.n_head, config.rope_base, config.cls_token_num)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):        
+        # Initialize nn.Linear and nn.Embedding
+        self.apply(self._init_weights)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.head.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)
+
+    def setup_kv_cache(self, enable=True):
+        for block in self.layers.self_dec:
+            block.attn.kv_cache = enable
+            block.attn.reset_kv_cache()
+
+        self.layers.kv_cache = enable
+        self.layers.reset_kv_cache()
+    
+    def preprocess_condition(self, condition, cond_drop_prob=0.0):
+        # Set class condition to None condition
+        drop_label_mask = torch.rand_like(condition, dtype=torch.float) < cond_drop_prob
+        condition = condition + self.config.vocab_size + 1  # [0, 999] -> [codebook_size + 1, codebook_size + 999]
+        condition[drop_label_mask] = self.none_conds_id
+        
+        if condition.ndim == 1:
+            condition = condition.unsqueeze(-1)
+        
+        return condition
+    
+    def forward_shared(self, input_ids, freqs_cis, num_query=None):
+        embedds = self.embeddings(input_ids)
+        
+        x = self.embed_drop(embedds)
+        num_query = input_ids.shape[-1] if num_query == None else num_query
+        queries = self.embeddings(torch.full((input_ids.shape[0], num_query), self.mask_token_id, device=input_ids.device))
+        
+        x = self.layers(x, queries, freqs_cis=freqs_cis)
+        logits = self.head(self.norm(x)).float()
+        
+        return logits
+
+    def forward(self, input_ids, condition, targets=None, debug=False):
+        # shift class id and dropout for classifier-free guidance
+        condition = self.preprocess_condition(condition, cond_drop_prob=self.config.class_dropout_prob)
+        
+        # shuffle input
+        shuffled_ids, orders = batch_seq_shuffle(input_ids)
+        
+        # shuffle RoPE
+        freqs_cis = self.freqs_cis.unsqueeze(0).repeat(input_ids.shape[0], 1, 1, 1).to(input_ids.device)
+        fixed_freqs_cis = freqs_cis[:, :1, ...]
+        shuffled_freqs_cis = batch_seq_shuffle(freqs_cis[:, 1:, ...], orders)[0]
+        freqs_cis = torch.cat([fixed_freqs_cis, shuffled_freqs_cis], dim=1)
+
+        # teacher-forcing input
+        logits = self.forward_shared(torch.cat([condition, shuffled_ids[:, :-1]], dim=-1), freqs_cis)
+        
+        loss = None
+        if targets is not None:
+            targets = batch_seq_shuffle(targets, orders)[0]
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+
+        return logits, loss
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        condition,
+        guidance_scale=4.0,
+        cfg_schedule='linear',
+        sample_schedule='arccos',
+        temperature=1.0,
+        top_k=0,
+        top_p=1,
+        seq_len=256,
+        num_iter=64,
+    ):
+        device = condition.device
+        num_samples = condition.shape[0]
+        freqs_cis_ = self.freqs_cis.unsqueeze(0).to(device)
+        
+        # shift condition id
+        condition = self.preprocess_condition(condition, cond_drop_prob=0.0)
+
+        # generate a random order
+        orders = torch.rand(256, device=device).argsort(dim=0) + 1
+
+        last_pos = 0
+        last_range = range(0, 1)  # for class token, hardcode
+        sequences = []
+        
+        self.setup_kv_cache(enable=True)
+        for step in range(num_iter):
+            if sample_schedule == 'arccos':
+                mask_ratio = np.arccos(1. * (step + 1) / num_iter) / (math.pi * 0.5)
+            elif sample_schedule == 'cosine':
+                mask_ratio = np.cos(math.pi / 2. * (step + 1) / num_iter)
+            else:
+                raise NotImplementedError
+
+            mask_len = int(seq_len * mask_ratio)
+            mask_len = max(1, min(seq_len - last_pos - 1, mask_len))
+            
+            num_pred = seq_len - last_pos - mask_len
+            if step == num_iter - 1:
+                num_pred = seq_len - last_pos
+
+            next_range = orders[range(last_pos, last_pos + num_pred)]
+            last_pos += num_pred
+            
+            if cfg_schedule == 'linear':
+                cfg_scale = 1.0 + (guidance_scale - 1.0) * last_pos / seq_len
+            elif cfg_schedule == 'constant':
+                cfg_scale = guidance_scale
+            else:
+                raise NotImplementedError
+
+            """
+            1. Since the cached key has already had rotary embedding applied,
+               we only need to input the current position's frequencies for key.
+            2. We need the next position's frequencies for query to achieve target-aware guidance.
+            """
+            freqs_cis = torch.cat([
+                freqs_cis_[:, last_range, ...],
+                freqs_cis_[:, next_range, ...]], dim=1
+            )
+            if guidance_scale != 0:
+                if step == 0:
+                    input_ids = torch.cat([condition, torch.full_like(condition, self.none_conds_id)], dim=0)
+                else:
+                    input_ids = torch.cat([sequences[-1], sequences[-1]], dim=0)
+
+                logits = self.forward_shared(input_ids, freqs_cis, num_pred)
+                cond_logits, uncond_logits = logits[:num_samples], logits[num_samples:]
+                logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
+            else:
+                raise NotImplementedError
+
+            # keep the logits of last n-tokens
+            logits = logits[:, -num_pred:] / 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 = F.softmax(logits, dim=-1)
+            sampled = torch.multinomial(probs.flatten(0, 1), num_samples=1)
+            sequences.append(sampled.reshape(num_samples, -1))
+            
+            last_range = next_range
+            
+        self.setup_kv_cache(enable=False)
+
+        sequences = torch.cat(sequences, dim=-1)
+        return sequences[:, orders.argsort(dim=0)]
+
+
+# https://github.com/pytorch-labs/gpt-fast/blob/main/model.py 
+def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000, cls_token_num=120):
+    freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+    t = torch.arange(seq_len, device=freqs.device)
+    freqs = torch.outer(t, freqs) # (seq_len, head_dim // 2)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+    cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1) # (cls_token_num+seq_len, head_dim // 2, 2)
+    cond_cache = torch.cat([torch.zeros(cls_token_num, n_elem // 2, 2), cache]) # (cls_token_num+seq_len, head_dim // 2, 2)
+    return cond_cache 
+
+
+def precompute_freqs_cis_2d(grid_size: int, n_elem: int, base: int = 10000, cls_token_num=120):
+    # split the dimension into half, one for x and one for y
+    half_dim = n_elem // 2
+    freqs = 1.0 / (base ** (torch.arange(0, half_dim, 2)[: (half_dim // 2)].float() / half_dim))
+    t = torch.arange(grid_size, device=freqs.device)
+    freqs = torch.outer(t, freqs) # (grid_size, head_dim // 2)
+    freqs_grid = torch.concat([
+        freqs[:, None, :].expand(-1, grid_size, -1),
+        freqs[None, :, :].expand(grid_size, -1, -1),
+    ], dim=-1)  # (grid_size, grid_size, head_dim // 2)
+    cache_grid = torch.stack([torch.cos(freqs_grid), torch.sin(freqs_grid)], dim=-1) # (grid_size, grid_size, head_dim // 2, 2)
+    cache = cache_grid.flatten(0, 1)
+    cond_cache = torch.cat([torch.zeros(cls_token_num, n_elem // 2, 2), cache]) # (cls_token_num+grid_size**2, head_dim // 2, 2)
+    return cond_cache 
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor):
+    # x: (bs, seq_len, n_head, head_dim)
+    # freqs_cis (seq_len, head_dim // 2, 2)
+    xshaped = x.float().reshape(*x.shape[:-1], -1, 2) # (bs, seq_len, n_head, head_dim//2, 2)
+    freqs_cis = freqs_cis.view(-1, xshaped.size(1), 1, xshaped.size(3), 2) # (1, seq_len, 1, head_dim//2, 2)
+    x_out2 = torch.stack([
+            xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+            xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+    ], dim=-1)
+    x_out2 = x_out2.flatten(3)
+    return x_out2.type_as(x)
+
+
+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 ARPG_XXL(**kwargs):
+    return Transformer(ModelArgs(n_layer=48, n_head=24, dim=1536, **kwargs))
+
+def ARPG_XL(**kwargs):
+    return Transformer(ModelArgs(n_layer=36, n_head=20, dim=1280, **kwargs))
+
+def ARPG_L(**kwargs):
+    return Transformer(ModelArgs(n_layer=24, n_head=16, dim=1024, **kwargs))
+
+
+ARPG_models = {'ARPG-L': ARPG_L, 'ARPG-XL': ARPG_XL, 'ARPG-XXL': ARPG_XXL}

pipeline.py

@@ -0,0 +1,111 @@
+from diffusers import DiffusionPipeline
+import torch
+import numpy as np
+import importlib.util
+import sys
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+import os
+from torchvision.utils import save_image
+from PIL import Image
+from safetensors.torch import load_file
+from .vq_model import VQ_models
+from .arpg import ARPG_models
+
+# inheriting from DiffusionPipeline for HF
+class ARPGModel(DiffusionPipeline):
+
+    def __init__(self):
+        super().__init__()
+
+    @torch.no_grad()
+    def __call__(self, *args, **kwargs):
+        """
+        This method downloads the model and VAE components,
+        then executes the forward pass based on the user's input.
+        """
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # init the mar model architecture
+        model_type = kwargs.get("model_type", "ARPG-XXL")
+
+        # download the pretrained model and set diffloss parameters
+        if model_type == "ARPG-L":
+            model_path = "arpg_300m.pt"
+        elif model_type == "ARPG-XL":
+            model_path = "arpg_700m.pt"
+        elif model_type == "ARPG-XXL":
+            model_path = "arpg_1b.pt"
+        else:
+            raise NotImplementedError
+        # download and load the model weights (.safetensors or .pth)
+        model_checkpoint_path = hf_hub_download(
+            repo_id=kwargs.get("repo_id", "hp-l33/ARPG"),
+            filename=kwargs.get("model_filename", model_path)
+        )
+
+        model_fn = ARPG_models[model_type]
+
+        model = model_fn(
+          num_classes=1000,
+          vocab_size=16384
+        ).cuda()
+
+        # use safetensors
+        state_dict = load_file(model_checkpoint_path)['state_dict']
+        model.load_state_dict(state_dict)
+        model.eval()
+
+        # download and load the vae
+        vae_checkpoint_path = hf_hub_download(
+            repo_id=kwargs.get("repo_id", "FoundationVision/LlamaGen"),
+            filename=kwargs.get("vae_filename", "vq_ds16_c2i.pt")
+        )
+
+        vae = VQ_models['VQ-16']()
+
+        vae_state_dict = load_file(vae_checkpoint_path)['model']
+        vae.load_state_dict(vae_state_dict)
+        vae = vae.to(device).eval()
+
+        # set up user-specified or default values for generation
+        seed = kwargs.get("seed", 6)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+
+        num_steps = kwargs.get("num_steps", 64)
+        cfg_scale = kwargs.get("cfg_scale", 4)
+        cfg_schedule = kwargs.get("cfg_schedule", "constant")
+        sample_schedule = kwargs.get("sample_schedule", "arccos")
+        temperature = kwargs.get("temperature", 1.0)
+        top_k = kwargs.get("top_k", 600)
+        class_labels = kwargs.get("class_labels", [207, 360, 388, 113, 355, 980, 323, 979])
+
+        # generate the tokens and images
+        with torch.cuda.amp.autocast():
+            sampled_tokens = model.generate(
+               condition=torch.Tensor(class_labels).long().cuda(),
+               num_iter=num_steps,
+               guidance_scale=cfg_scale,
+               cfg_schedule=cfg_schedule,
+               sample_schedule=sample_schedule,
+               temperature=temperature,
+               top_k=top_k,
+            )
+            sampled_images = vae.decode_code(sampled_tokens, shape=(len(class_labels), 8, 16, 16))
+
+        output_dir = kwargs.get("output_dir", "./")
+        os.makedirs(output_dir, exist_ok=True)
+    
+        # save the images
+        image_path = os.path.join(output_dir, "sampled_image.png")
+        samples_per_row = kwargs.get("samples_per_row", 4)
+    
+        save_image(
+            torch.clamp(127.5 * sampled_images + 128.0, 0, 255), image_path, nrow=int(samples_per_row), normalize=False
+        )
+    
+        # return as a pil image
+        image = Image.open(image_path)
+    
+        return image

vq_model.py

@@ -0,0 +1,459 @@
+# Modified from:
+#   taming-transformers: https://github.com/CompVis/taming-transformers
+#   maskgit: https://github.com/google-research/maskgit
+
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, reduce
+
+
+@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
+    num_res_blocks: int = 4
+
+
+
+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 = rearrange(q, 'b c h w -> b (h w) c')
+        k = rearrange(k, 'b c h w -> b (h w) c')
+        v = rearrange(v, 'b c h w -> b (h w) c')
+        
+        # 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_ = nn.functional.scaled_dot_product_attention(q, k, v)
+        h_ = rearrange(h_, 'b (h w) c -> b c h w', h=h, w=w)
+        # 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
+
+
+def compute_entropy_loss2(
+    logits,
+    temperature=0.01,
+    sample_minimization_weight=1.0,
+    batch_maximization_weight=1.0,
+    eps=1e-5,
+):
+    """
+    Entropy loss of unnormalized logits
+
+    logits: Affinities are over the last dimension
+
+    https://github.com/google-research/magvit/blob/05e8cfd6559c47955793d70602d62a2f9b0bdef5/videogvt/train_lib/losses.py#L279
+    LANGUAGE MODEL BEATS DIFFUSION — TOKENIZER IS KEY TO VISUAL GENERATION (2024)
+    """
+    probs = F.softmax(logits / temperature, -1)
+    log_probs = F.log_softmax(logits / temperature + eps, -1)
+
+    avg_probs = reduce(probs, "... D -> D", "mean")
+
+    avg_entropy = -torch.sum(avg_probs * torch.log(avg_probs + eps))
+
+    sample_entropy = -torch.sum(probs * log_probs, -1)
+    sample_entropy = torch.mean(sample_entropy)
+
+    loss = (sample_minimization_weight * sample_entropy) - (
+        batch_maximization_weight * avg_entropy
+    )
+
+    return sample_entropy, avg_entropy, loss
+
+
+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}