x: B, T, D
(self, x, timesteps, length=None, text=None, xf_proj=None, xf_out=None)
| 405 | return src_mask |
| 406 | |
| 407 | def forward(self, x, timesteps, length=None, text=None, xf_proj=None, xf_out=None): |
| 408 | """ |
| 409 | x: B, T, D |
| 410 | """ |
| 411 | B, T = x.shape[0], x.shape[1] |
| 412 | if text is not None and len(text) != B: |
| 413 | index = x.device.index |
| 414 | text = text[index * B: index * B + B] |
| 415 | if xf_proj is None or xf_out is None: |
| 416 | xf_proj, xf_out = self.encode_text(text, x.device) |
| 417 | |
| 418 | emb = self.time_embed(timestep_embedding(timesteps, self.latent_dim)) + xf_proj |
| 419 | |
| 420 | # B, T, latent_dim |
| 421 | h = self.joint_embed(x) |
| 422 | h = h + self.sequence_embedding.unsqueeze(0)[:, :T, :] |
| 423 | |
| 424 | src_mask = self.generate_src_mask(T, length).to(x.device).unsqueeze(-1) |
| 425 | for module in self.temporal_decoder_blocks: |
| 426 | h = module(h, xf_out, emb, src_mask) |
| 427 | |
| 428 | output = self.out(h).view(B, T, -1).contiguous() |
| 429 | return output |
nothing calls this directly
no test coverage detected