(self, z)
| 987 | |
| 988 | @torch.inference_mode() |
| 989 | def decode_original(self, z): |
| 990 | # b (nc cf) c h w -> (b nc) cf c h w -> decode -> (b nc) c cf h w -> b (nc cf) c h w |
| 991 | chunks = list(z.split(self.latent_len, dim=1)) |
| 992 | |
| 993 | if self.world_size > 1: |
| 994 | chunks_total_num = len(chunks) |
| 995 | max_num_per_rank = (chunks_total_num + self.world_size - 1) // self.world_size |
| 996 | rank = torch.distributed.get_rank() |
| 997 | chunks_ = chunks[max_num_per_rank * rank : max_num_per_rank * (rank + 1)] |
| 998 | if len(chunks_) < max_num_per_rank: |
| 999 | chunks_.extend(chunks[:max_num_per_rank-len(chunks_)]) |
| 1000 | chunks = chunks_ |
| 1001 | |
| 1002 | for i in range(len(chunks)): |
| 1003 | chunks[i] = self.decode_naive(chunks[i], True).permute(0,2,1,3,4) |
| 1004 | x = torch.cat(chunks, dim=1) |
| 1005 | |
| 1006 | if self.world_size > 1: |
| 1007 | x_ = torch.empty([x.size(0), (self.world_size * max_num_per_rank) * self.frame_len, *x.shape[2:]], dtype=x.dtype, device=x.device) |
| 1008 | torch.distributed.all_gather_into_tensor(x_, x) |
| 1009 | x = x_[:, : chunks_total_num * self.frame_len] |
| 1010 | |
| 1011 | x = self.mix(x) |
| 1012 | return x |
| 1013 | |
| 1014 | def mix(self, x, smooth_scale = 0.6): |
| 1015 | remain_scale = smooth_scale |
nothing calls this directly
no test coverage detected