LlamaGen text-to-image training and inference code

[Masaaki-75]

Thanks for open-sourcing this great work!

Just to confirm, the training code of LlamaGen (https://github.com/FoundationVision/UniTok/tree/main/eval/llamagen) in this repo is actually not the adapted version for UniTok, right? (since the UniTok model is not presented in tokenizer/tokenizer_image/vq_model.py) That means, if we are to use UniTok to train LlamaGen, we need to add the UniTok model class in the vq_model.py.

[machuofan]

Hi @Masaaki-75, for class-to-image generation, the LlamaGen training code has already been customized for UniTok. For example, UniTok is adopted in the image preprocessing stage. You do not need to make further changes.

[Masaaki-75]

Thanks for the reply! Actually, I was trying to adapt UniTok to text-to-image generation task. I somehow manage starting the training by modifying the forward methods in AutoregressiveHead and the Transformer (don't know if correctly set though). For example:

    def forward(
            self,
            idx: torch.Tensor,
            cond_idx: 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
            # [B, Lt, 768]
            cond_embeddings = self.cls_embedding(cond_idx, train=self.training)[:, :self.cls_token_num]
            token_embeddings = self.tok_embeddings(idx)  # [B, Li-1, 768]
            token_embeddings = torch.cat((cond_embeddings, token_embeddings), dim=1)  # [B, Lt+Li-1, 768]
            h = self.tok_dropout(token_embeddings)
            self.freqs_cis = self.freqs_cis.to(h.device)
        else:
            if cond_idx is not None:  # prefill in inference
                token_embeddings = self.cls_embedding(cond_idx, train=self.training)[:, :self.cls_token_num]
            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)
            self.freqs_cis = self.freqs_cis
        
        if self.training:
            freqs_cis = self.freqs_cis[:token_embeddings.shape[1]]
        else:
            freqs_cis = self.freqs_cis[input_pos]

        # transformer blocks
        for layer in self.layers:
            # [B, Lt+Li-1, 768]
            h = layer(h, freqs_cis, input_pos, mask)

        # output layers
        h = self.norm(h)

        if self.training:  
            # [NOTE] I modified here, dropping all but last text tokens to ensure
            # token length alignment between `h` and `target` in the following 
            # `forward_train` method.
            h = h[:, self.cls_token_num - 1:].contiguous()

        if targets is None:
            return h
        else:
            logits = self.output.forward_train(h, targets=targets).contiguous()

        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

        return logits, loss

This training code runs smoothly. However, when i was trying to sample from the trained GPT model, it complained about shape mismatching in setting up KVcache during prefilling. I am wondering if you could share any insights on that?

File "xxx/unitok/eval/generation/autoregressive/models/generate.py", line 214, in generate
    next_token = prefill(model, cond_combined, input_pos, cfg_scale, **sampling_kwargs)
  File "xxxx/unitok/eval/generation/autoregressive/models/generate.py", line 95, in prefill
    logits = model.output.forward_test(next_token, input_pos=start_pos, mask=mask)
  File "xxx/unitok/eval/generation/autoregressive/models/gpt.py", line 387, in forward_test
    h = layer(h, freqs_cis=None, start_pos=input_pos, mask=mask)
  File "/opt/conda/envs/llava/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llava/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "xxx/unitok/eval/generation/autoregressive/models/gpt.py", line 331, in forward
    h = x + self.drop_path(self.attention(self.attention_norm(x), freqs_cis, start_pos, mask, full_attn))
  File "/opt/conda/envs/llava/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llava/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "xxx/unitok/eval/generation/autoregressive/models/gpt.py", line 302, in forward
    keys, values = self.kv_cache.update(input_pos, xk, xv)
  File "xxx/unitok/eval/generation/autoregressive/models/gpt.py", line 250, in update
    assert input_pos.shape[0] == k_val.shape[2]

[machuofan]

There could be many potential bugs leading to this error. Could you please print out the shape of your kv cache and input_pos?

[Masaaki-75]

The error was originally raised in the forward_test method (https://github.com/FoundationVision/UniTok/blob/main/eval/llamagen/autoregressive/models/generate.py#L90-L93):

        start_pos = torch.tensor([i], dtype=torch.int)
        mask = model.output.head_causal_mask[:bs, None, start_pos]
        if i == 0:
            logits = model.output.forward_test(next_token, input_pos=start_pos, mask=mask)

and (https://github.com/FoundationVision/UniTok/blob/main/eval/llamagen/autoregressive/models/gpt.py#L328):

    def forward_test(self, base_tokens, idx=None, input_pos=None, mask=None):
        if idx is not None:
            h = self.codebooks[input_pos - 1](idx)
        else:
            h = base_tokens  # [B, Lt, D]. this should be the T5 embedding at the very start of the prefilling stage.
        for layer in self.layers:
            h = layer(h, freqs_cis=None, start_pos=input_pos, mask=mask)
        h = self.norm(h)
        logits = self.linear_head(h)
        return logits

eventually leading to the problem in KVCache. The shape of k_val is like [B, 8, 120, 768] while the input_pos (essentially the start_pos in generate.py) is [1].

---

Update 1

Actually, later I changed h = base_tokens in forward_test to h = base_tokens[:, -1:, :] (still not sure if it is correct to do so. Please correct me if i am wrong), and got this clear. But I then encountered another shape mismatch:

  File "xxxx/unitok/eval/generation/autoregressive/models/generate.py", line 221, in generate
    generated_tokens = decode_n_tokens(model, next_token, input_pos, max_new_tokens - 1, cfg_scale, cfg_interval,
  File "xxxxx/unitok/eval/generation/autoregressive/models/generate.py", line 162, in decode_n_tokens
    next_token = decode_one_token(model, cur_token, input_pos, cfg_scale, cfg_flag, **sampling_kwargs)
  File "xxxxx/unitok/eval/generation/autoregressive/models/generate.py", line 135, in decode_one_token
    logits = model.output.forward_test(next_token, idx=pred_idx, input_pos=start_pos, mask=mask)
  File "xxxxx/external/unitok/eval/generation/autoregressive/models/gpt.py", line 392, in forward_test
    h = layer(h, freqs_cis=None, start_pos=input_pos, mask=mask)
...
  File "xxxx/unitok/eval/generation/autoregressive/models/gpt.py", line 332, in forward
    h = x + self.drop_path(self.attention(self.attention_norm(x), freqs_cis, start_pos, mask, full_attn))
...
  File "xxxx/unitok/eval/generation/autoregressive/models/gpt.py", line 303, in forward
    keys, values = self.kv_cache.update(input_pos, xk, xv)
  File "xxxx/unitok/eval/generation/autoregressive/models/gpt.py", line 254, in update
    k_out[:, :, input_pos] = k_val
RuntimeError: shape mismatch: value tensor of shape [64, 12, 1, 64] cannot be broadcast to indexing result of shape [32, 12, 1, 64]

even when i set the cfg_scale to 1.

---

Update 2

Turns out it is caused by https://github.com/FoundationVision/UniTok/blob/main/eval/llamagen/autoregressive/models/generate.py#L95
https://github.com/FoundationVision/UniTok/blob/main/eval/llamagen/autoregressive/models/generate.py#L129

pred_idx = torch.cat([pred_idx, pred_idx])

What are these two lines for?

[Masaaki-75]

Update 3

I guess the pred_idx = torch.cat([pred_idx, pred_idx]) lines are only for cfg_scale > 1, so i add the if condition and the sampling triggers no error. however, the sampled images seemed broken (basically it shows only a rough outline, without much content). And I don't know whether it is caused by training or sampling. It would be great if anyone can share insights on that.