(model, latents)
| 113 | |
| 114 | @torch.no_grad() |
| 115 | def decode_latents(model, latents): |
| 116 | if not isinstance(latents, list): |
| 117 | latents = [latents] |
| 118 | for latent_file in latents: |
| 119 | if isinstance(latent_file, str): |
| 120 | # latent_path |
| 121 | latent = torch.load(latent_file).to(model.device) # torch.Size([1, 16, 13, 90, 160]) |
| 122 | |
| 123 | # * truncate |
| 124 | TRUNCATE = 7 |
| 125 | latent = latent[:, :, :TRUNCATE].contiguous() |
| 126 | |
| 127 | recon = model.first_stage_model.decode(latent).to(torch.float32) |
| 128 | |
| 129 | samples_x = recon.permute(0, 2, 1, 3, 4).contiguous() |
| 130 | samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0).cpu() |
| 131 | |
| 132 | save_path = latent_file.replace(".pt", ".mp4") |
| 133 | if mpu.get_model_parallel_rank() == 0: |
| 134 | save_video_as_grid_and_mp4(samples, save_path, fps=args.sampling_fps) |
| 135 | |
| 136 | del latent, recon, samples_x, samples |
| 137 | |
| 138 | |
| 139 | # * Efficient Decoding |
no test coverage detected