CompVis / taming-transformers

Taming Transformers for High-Resolution Image Synthesis
https://arxiv.org/abs/2012.09841
MIT License
5.83k stars 1.15k forks source link

Simply concatenating the source indices and condition indices even when they are from different codebooks? #228

Open abcqmars opened 1 year ago

abcqmars commented 1 year ago

Hi, I got a puzzle when I was reading the code of Net2NetTransformer cond_transformer.py:80

    def forward(self, x, c):
        # one step to produce the logits
        _, z_indices = self.encode_to_z(x)
        _, c_indices = self.encode_to_c(c)

        if self.training and self.pkeep < 1.0:
            mask = torch.bernoulli(self.pkeep*torch.ones(z_indices.shape,
                                                         device=z_indices.device))
            mask = mask.round().to(dtype=torch.int64)
            r_indices = torch.randint_like(z_indices, self.transformer.config.vocab_size)
            a_indices = mask*z_indices+(1-mask)*r_indices
        else:
            a_indices = z_indices

        cz_indices = torch.cat((c_indices, a_indices), dim=1)

        # target includes all sequence elements (no need to handle first one
        # differently because we are conditioning)
        target = z_indices
        # make the prediction
        logits, _ = self.transformer(cz_indices[:, :-1])

The condition like depth images has thier own specifically trained codebook, so their indices should have different meanings with that of rgb image code book. But in this code snippet above, It seems that we simply concatenate them and feed into transformer.

Where the transformer will process them in one shared token embeding layer as shown in mingpt.py:162

   def forward(self, idx, embeddings=None, targets=None):
        # forward the GPT model
        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector

Is this implementation correct? Or I have a wrong understanding of it?

zch42 commented 11 months ago

I have the same question

qinxiangyujiayou commented 9 months ago

I have the same question