Class ReMoDiffuseTransformer

diffplanner/models/transformers/remodiffuse.py:170–253

@SUBMODULES.register_module()
class ReMoDiffuseTransformer(DiffusionTransformer):
    def __init__(self, guide_scale=None, retr_guide_scale=None, retrieval_cfg=None, waypoint=False, **kwargs):
        super().__init__(**kwargs)
        self.guide_scale = guide_scale
        self.retr_guide_scale = retr_guide_scale
        self.waypoint = waypoint
        if retrieval_cfg is None:
            self.use_retrieval = False
        else:
            self.database = RetrievalDatabase(**retrieval_cfg)
            self.use_retrieval = True
        
    def get_precompute_condition(self, text=None, motion_length=None, xf_out=None, re_feat=None, device=None):
        if xf_out is None:
            xf_out = self.encode_text(text, device)
        output = {'xf_out': xf_out}
        if self.use_retrieval:
            if re_feat is None:
                re_feat = self.database(text, motion_length, self.clip, device)
            output['re_feat'] = re_feat
        return output

    def forward_train(self, h=None, src_mask=None, emb=None, xf_out=None, re_feat=None, **kwargs):
        B, T = h.shape[0], h.shape[1]
        cond_type = random.randint(0, 99)
        for module in self.temporal_decoder_blocks:
            h = module(x=h, xf=xf_out, emb=emb, src_mask=src_mask, cond_type=cond_type, re_feat=re_feat)

        output = self.out(h).view(B, T, -1).contiguous()
        return output
    
    def forward_test(self, h=None, src_mask=None, emb=None, xf_out=None, re_feat=None, **kwargs):
        B, T = h.shape[0], h.shape[1]
        scale = self.guide_scale
        
        # for waypoint heading
        if self.waypoint:
            src_mask = src_mask * 0 + 1

        # # for unconditional
        # cond_type = 0
        # for module in self.temporal_decoder_blocks: 
        #     h = module(x=h, xf=xf_out, emb=emb, src_mask=src_mask, cond_type=cond_type, re_feat=re_feat) 
        #     output = self.out(h).view(B, T, -1).contiguous() 
        # return output
        
        if abs(scale - 1.0) < 0.000001:
            cond_type = 99
            for module in self.temporal_decoder_blocks:
                h = module(x=h, xf=xf_out, emb=emb, src_mask=src_mask, cond_type=cond_type, re_feat=re_feat)
            output = self.out(h).view(B, T, -1).contiguous()
        elif not self.use_retrieval:
            h0 = h
            cond_type = 0
            for module in self.temporal_decoder_blocks:
                h0 = module(x=h0, xf=xf_out, emb=emb, src_mask=src_mask, cond_type=cond_type, re_feat=re_feat)
            out0 = self.out(h0).view(B, T, -1).contiguous()
            h1 = h