Class FineMoGenTransformer

mogen/models/transformers/finemogen.py:228–358

@SUBMODULES.register_module()
class FineMoGenTransformer(DiffusionTransformer):

    def __init__(self,
                 scale_func_cfg=None,
                 pose_encoder_cfg=None,
                 pose_decoder_cfg=None,
                 moe_route_loss_weight=1.0,
                 template_kl_loss_weight=0.0001,
                 **kwargs):
        super().__init__(**kwargs)
        self.scale_func_cfg = scale_func_cfg
        self.joint_embed = PoseEncoder(**pose_encoder_cfg)
        self.out = zero_module(PoseDecoder(**pose_decoder_cfg))
        self.moe_route_loss_weight = moe_route_loss_weight
        self.template_kl_loss_weight = template_kl_loss_weight

    def build_temporal_blocks(self, sa_block_cfg, ca_block_cfg, ffn_cfg):
        self.temporal_decoder_blocks = nn.ModuleList()
        for i in range(self.num_layers):
            if isinstance(ffn_cfg, list):
                ffn_cfg_block = ffn_cfg[i]
            else:
                ffn_cfg_block = ffn_cfg
            self.temporal_decoder_blocks.append(
                DecoderLayer(ca_block_cfg=ca_block_cfg, ffn_cfg=ffn_cfg_block))

    def scale_func(self, timestep):
        scale = self.scale_func_cfg['scale']
        w = (1 - (1000 - timestep) / 1000) * scale + 1
        output = {'text_coef': w, 'none_coef': 1 - w}
        return output

    def aux_loss(self):
        aux_loss = 0
        kl_loss = 0
        for module in self.temporal_decoder_blocks:
            if hasattr(module.ca_block, 'aux_loss'):
                aux_loss = aux_loss + module.ca_block.aux_loss
            if hasattr(module.ca_block, 'kl_loss'):
                kl_loss = kl_loss + module.ca_block.kl_loss
        losses = {}
        if aux_loss > 0:
            losses['moe_route_loss'] = aux_loss * self.moe_route_loss_weight
        if kl_loss > 0:
            losses['template_kl_loss'] = kl_loss * self.template_kl_loss_weight
        return losses

    def get_precompute_condition(self,
                                 text=None,
                                 motion_length=None,
                                 xf_out=None,
                                 re_dict=None,
                                 device=None,
                                 sample_idx=None,
                                 clip_feat=None,
                                 **kwargs):
        if xf_out is None:
            xf_out = self.encode_text(text, clip_feat, device)