(latent, T, first_stage_model)
| 138 | |
| 139 | # * Efficient Decoding |
| 140 | def get_reconstruction_from_latents(latent, T, first_stage_model): |
| 141 | recons = [] |
| 142 | loop_num = (T - 1) // 2 |
| 143 | for i in range(loop_num): |
| 144 | if i == 0: |
| 145 | start_frame, end_frame = 0, 3 |
| 146 | else: |
| 147 | start_frame, end_frame = i * 2 + 1, i * 2 + 3 |
| 148 | if i == loop_num - 1: |
| 149 | clear_fake_cp_cache = True |
| 150 | else: |
| 151 | clear_fake_cp_cache = False |
| 152 | with torch.no_grad(): |
| 153 | recon = first_stage_model.decode( |
| 154 | latent[:, :, start_frame:end_frame].contiguous(), clear_fake_cp_cache=clear_fake_cp_cache |
| 155 | ) # torch.Size([1, 16, 3 (or 2), 64, 112]) |
| 156 | |
| 157 | recons.append(recon) |
| 158 | |
| 159 | recons = torch.cat(recons, dim=2).to(torch.float32) # [b, c, t, h, w] |
| 160 | return recons |
| 161 | |
| 162 | |
| 163 | # For n_cond: 1 |
no test coverage detected