Method forward

mogen/models/transformers/mdm.py:140–184

motion: B, T, D timesteps: [batch_size] (int)

(self, motion, timesteps, text_feat=None, **kwargs)

        return motion

    def forward(self, motion, timesteps, text_feat=None, **kwargs):
        """
        motion: B, T, D
        timesteps: [batch_size] (int)
        """
        B, T, D = motion.shape
        if text_feat is None:
            enc_text = self.get_precompute_condition(**kwargs)['text_feat']
        else:
            enc_text = text_feat
        if self.training:
            # T, B, D
            motion = self.poseEmbedding(motion).permute(1, 0, 2)

            emb = self.embed_timestep(timesteps)  # [1, bs, d]
            emb += self.embed_text(self.mask_cond(enc_text, force_mask=False))

            xseq = self.sequence_pos_encoder(torch.cat((emb, motion), axis=0))
            output = self.seqTransEncoder(xseq)[1:]

            # B, T, D
            output = self.poseFinal(output).permute(1, 0, 2)
            return output
        else:
            # T, B, D
            motion = self.poseEmbedding(motion).permute(1, 0, 2)

            emb = self.embed_timestep(timesteps)  # [1, bs, d]
            emb_uncond = emb + \
                self.embed_text(self.mask_cond(enc_text, force_mask=True))
            emb_text = emb + \
                self.embed_text(self.mask_cond(enc_text, force_mask=False))

            xseq = self.sequence_pos_encoder(
                torch.cat((emb_uncond, motion), axis=0))
            xseq_text = self.sequence_pos_encoder(
                torch.cat((emb_text, motion), axis=0))
            output = self.seqTransEncoder(xseq)[1:]
            output_text = self.seqTransEncoder(xseq_text)[1:]
            # B, T, D
            output = self.poseFinal(output).permute(1, 0, 2)
            output_text = self.poseFinal(output_text).permute(1, 0, 2)
            scale = self.guide_scale
            output = output + scale * (output_text - output)
            return output


class PositionalEncoding(nn.Module):