| 3 | import torch.nn.functional as F |
| 4 | |
| 5 | class DecoderWrapper(nn.Module): |
| 6 | def __init__(self, original_vae): |
| 7 | super().__init__() |
| 8 | self.decoder = original_vae.decoder |
| 9 | self.decoder_latent_proj = original_vae.decoder_latent_proj |
| 10 | self.skel_embedding = original_vae.skel_embedding |
| 11 | self.final_layer = original_vae.final_layer |
| 12 | self.query_pos_decoder = original_vae.query_pos_decoder |
| 13 | self.arch = original_vae.arch |
| 14 | self.latent_std = original_vae.latent_std |
| 15 | |
| 16 | self.h_dim = original_vae.h_dim |
| 17 | |
| 18 | |
| 19 | def forward(self, z, history_motion): |
| 20 | # 使用固定的 nfuture 或从输入张量获取 |
| 21 | bs = history_motion.shape[0] |
| 22 | |
| 23 | device = next(self.parameters()).device |
| 24 | z = z.to(device) |
| 25 | history_motion = history_motion.to(device) |
| 26 | |
| 27 | nfuture=8 |
| 28 | |
| 29 | # 这部分与原始 decode 方法相同 |
| 30 | z = self.decoder_latent_proj(z) |
| 31 | # device = z.device |
| 32 | queries = torch.zeros(nfuture, bs, self.h_dim).to(device) |
| 33 | history_embedding = self.skel_embedding(history_motion).permute(1, 0, 2).to(device) |
| 34 | |
| 35 | if self.arch == "all_encoder": |
| 36 | xseq = torch.cat((z, history_embedding, queries), dim=0) |
| 37 | xseq = self.query_pos_decoder(xseq) |
| 38 | output = self.decoder(xseq)[-nfuture:] |
| 39 | elif self.arch == "encoder_decoder": |
| 40 | xseq = torch.cat((history_embedding, queries), dim=0) |
| 41 | xseq = self.query_pos_decoder(xseq) |
| 42 | output = self.decoder(tgt=xseq, memory=z) |
| 43 | output = output[-nfuture:] |
| 44 | |
| 45 | output = self.final_layer(output) |
| 46 | feats = output.permute(1, 0, 2) |
| 47 | return feats |
nothing calls this directly
no outgoing calls
no test coverage detected