(self, captions, lengths, clip_model, device, idx=None)
| 152 | return src_mask |
| 153 | |
| 154 | def forward(self, captions, lengths, clip_model, device, idx=None): |
| 155 | B = len(captions) |
| 156 | all_indexes = [] |
| 157 | for b_ix in range(B): |
| 158 | length = int(lengths[b_ix]) |
| 159 | if idx is None: |
| 160 | batch_indexes = self.retrieve(captions[b_ix], length, |
| 161 | clip_model, device) |
| 162 | else: |
| 163 | batch_indexes = self.retrieve(captions[b_ix], length, |
| 164 | clip_model, device, idx[b_ix]) |
| 165 | all_indexes.extend(batch_indexes) |
| 166 | all_indexes = np.array(all_indexes) |
| 167 | all_motions = torch.Tensor(self.motions[all_indexes]).to(device) |
| 168 | all_m_lengths = torch.Tensor(self.m_lengths[all_indexes]).long() |
| 169 | |
| 170 | T = all_motions.shape[1] |
| 171 | src_mask = self.generate_src_mask(T, all_m_lengths).to(device) |
| 172 | raw_src_mask = src_mask.clone() |
| 173 | re_motion = self.motion_proj(all_motions) + \ |
| 174 | self.motion_pos_embedding.unsqueeze(0) |
| 175 | for module in self.motion_encoder_blocks: |
| 176 | re_motion = module(x=re_motion, src_mask=src_mask.unsqueeze(-1)) |
| 177 | re_motion = re_motion.view(B, self.num_retrieval, T, -1).contiguous() |
| 178 | # stride |
| 179 | re_motion = re_motion[:, :, ::self.stride, :].contiguous() |
| 180 | |
| 181 | src_mask = src_mask[:, ::self.stride].contiguous() |
| 182 | src_mask = src_mask.view(B, self.num_retrieval, -1).contiguous() |
| 183 | |
| 184 | T = 77 |
| 185 | all_text_seq_features = torch.Tensor( |
| 186 | self.clip_seq_features[all_indexes]).to(device) |
| 187 | all_text_seq_features = all_text_seq_features.permute(1, 0, 2) |
| 188 | re_text = self.text_encoder(all_text_seq_features) |
| 189 | re_text = re_text.permute(1, 0, 2) |
| 190 | re_text = re_text.view(B, self.num_retrieval, T, -1).contiguous() |
| 191 | re_text = re_text[:, :, -1:, :].contiguous() |
| 192 | |
| 193 | re_dict = dict(re_text=re_text, |
| 194 | re_motion=re_motion, |
| 195 | re_mask=src_mask, |
| 196 | raw_motion=all_motions, |
| 197 | raw_motion_length=all_m_lengths, |
| 198 | raw_motion_mask=raw_src_mask) |
| 199 | return re_dict |
| 200 | |
| 201 | |
| 202 | @SUBMODULES.register_module() |
nothing calls this directly
no test coverage detected