nsr/train_util_diffusion_accelerate.py

import copy
import functools
import json
import os
from pathlib import Path
from pdb import set_trace as st

import blobfile as bf
import imageio
import numpy as np
import torch as th
import torch.distributed as dist
import torchvision
from PIL import Image
from torch.nn.parallel.distributed import DistributedDataParallel as DDP
from torch.optim import AdamW
from torch.utils.tensorboard.writer import SummaryWriter
from tqdm import tqdm

from guided_diffusion import dist_util, logger
from guided_diffusion.fp16_util import MixedPrecisionTrainer
from guided_diffusion.nn import update_ema
from guided_diffusion.resample import LossAwareSampler, UniformSampler
# from .train_util import TrainLoop3DRec
from guided_diffusion.train_util import (TrainLoop, calc_average_loss,
                                         find_ema_checkpoint,
                                         find_resume_checkpoint,
                                         get_blob_logdir, log_loss_dict,
                                         log_rec3d_loss_dict,
                                         parse_resume_step_from_filename)

import dnnlib

# AMP
from accelerate import Accelerator

# from ..guided_diffusion.train_util import TrainLoop

# use_amp = False
# use_amp = True


class TrainLoop3DDiffusion(TrainLoop):

    def __init__(
            self,
            *,
            # model,
            rec_model,
            denoise_model,
            diffusion,
            loss_class,
            data,
            eval_data,
            batch_size,
            microbatch,
            lr,
            ema_rate,
            log_interval,
            eval_interval,
            save_interval,
            resume_checkpoint,
            use_fp16=False,
            fp16_scale_growth=0.001,
            schedule_sampler=None,
            weight_decay=0,
            lr_anneal_steps=0,
            iterations=10001,
            ignore_resume_opt=False,
            freeze_ae=False,
            denoised_ae=True,
            triplane_scaling_divider=10,
            use_amp=False,
            **kwargs):

        super().__init__(model=denoise_model,
                         diffusion=diffusion,
                         data=data,
                         batch_size=batch_size,
                         microbatch=microbatch,
                         lr=lr,
                         ema_rate=ema_rate,
                         log_interval=log_interval,
                         save_interval=save_interval,
                         resume_checkpoint=resume_checkpoint,
                         use_fp16=use_fp16,
                         fp16_scale_growth=fp16_scale_growth,
                         schedule_sampler=schedule_sampler,
                         lr_anneal_steps=lr_anneal_steps,
                         weight_decay=weight_decay,
                         use_amp=use_amp)
        
        self.accelerator = Accelerator()

        self.pool_512 = th.nn.AdaptiveAvgPool2d((512, 512))
        self.pool_128 = th.nn.AdaptiveAvgPool2d((128, 128))
        self.loss_class = loss_class
        self.rec_model = rec_model
        self.eval_interval = eval_interval
        self.eval_data = eval_data
        self.iterations = iterations
        # self.triplane_std = 10
        self.triplane_scaling_divider = triplane_scaling_divider

        self._load_and_sync_parameters(model=self.rec_model, model_name='rec')

        # * for loading EMA
        self.mp_trainer_rec = MixedPrecisionTrainer(
            model=self.rec_model,
            use_fp16=self.use_fp16,
            use_amp=use_amp,
            fp16_scale_growth=fp16_scale_growth,
            model_name='rec',
        )
        self.denoised_ae = denoised_ae
        if not freeze_ae:
            self.opt_rec = AdamW(self._init_optim_groups(kwargs))
        else:
            print('!! freezing AE !!')

        if dist_util.get_rank() == 0:
            self.writer = SummaryWriter(log_dir=f'{logger.get_dir()}/runs')
            print(self.opt)
            if not freeze_ae:
                print(self.opt_rec)

        # if not freeze_ae:
        if self.resume_step:
            if not ignore_resume_opt:
                self._load_optimizer_state()
            else:
                logger.warn("Ignoring optimizer state from checkpoint.")
            # Model was resumed, either due to a restart or a checkpoint
            # being specified at the command line.
            # if not freeze_ae:
            #     self.ema_params_rec = [
            #         self._load_ema_parameters(
            #             rate,
            #             self.rec_model,
            #             self.mp_trainer_rec,
            #             model_name=self.mp_trainer_rec.model_name)
            #         for rate in self.ema_rate
            #     ]
            # else:
            self.ema_params_rec = [
                self._load_ema_parameters(
                    rate,
                    self.rec_model,
                    self.mp_trainer_rec,
                    model_name=self.mp_trainer_rec.model_name)
                for rate in self.ema_rate
            ]
        else:
            if not freeze_ae:
                self.ema_params_rec = [
                    copy.deepcopy(self.mp_trainer_rec.master_params)
                    for _ in range(len(self.ema_rate))
                ]

        if self.use_ddp is True:
            self.rec_model = th.nn.SyncBatchNorm.convert_sync_batchnorm(
                self.rec_model)
            self.ddp_rec_model = DDP(
                self.rec_model,
                device_ids=[dist_util.dev()],
                output_device=dist_util.dev(),
                broadcast_buffers=False,
                bucket_cap_mb=128,
                find_unused_parameters=False,
                # find_unused_parameters=True,
            )
        else:
            self.ddp_rec_model = self.rec_model

        if freeze_ae:
            self.ddp_rec_model.eval()
            self.ddp_rec_model.requires_grad_(False)
        self.freeze_ae = freeze_ae

        # if use_amp:

    def _init_optim_groups(self, kwargs):
        optim_groups = [
            # vit encoder
            {
                'name': 'vit_encoder',
                'params': self.mp_trainer_rec.model.encoder.parameters(),
                'lr': kwargs['encoder_lr'],
                'weight_decay': kwargs['encoder_weight_decay']
            },
            # vit decoder
            {
                'name':
                'vit_decoder',
                'params':
                self.mp_trainer_rec.model.decoder.vit_decoder.parameters(),
                'lr':
                kwargs['vit_decoder_lr'],
                'weight_decay':
                kwargs['vit_decoder_wd']
            },
            {
                'name':
                'vit_decoder_pred',
                'params':
                self.mp_trainer_rec.model.decoder.decoder_pred.parameters(),
                'lr':
                kwargs['vit_decoder_lr'],
                # 'weight_decay': 0
                'weight_decay':
                kwargs['vit_decoder_wd']
            },

            # triplane decoder
            {
                'name':
                'triplane_decoder',
                'params':
                self.mp_trainer_rec.model.decoder.triplane_decoder.parameters(
                ),
                'lr':
                kwargs['triplane_decoder_lr'],
                # 'weight_decay': self.weight_decay
            },
        ]

        if self.mp_trainer_rec.model.decoder.superresolution is not None:
            optim_groups.append({
                'name':
                'triplane_decoder_superresolution',
                'params':
                self.mp_trainer_rec.model.decoder.superresolution.parameters(),
                'lr':
                kwargs['super_resolution_lr'],
            })

        return optim_groups

    def run_loop(self, batch=None):
        th.cuda.empty_cache()
        while (not self.lr_anneal_steps
               or self.step + self.resume_step < self.lr_anneal_steps):

            # let all processes sync up before starting with a new epoch of training
            dist_util.synchronize()

            # batch, cond = next(self.data)
            # if batch is None:
            batch = next(self.data)
            self.run_step(batch)
            if self.step % self.log_interval == 0 and dist_util.get_rank(
            ) == 0:
                out = logger.dumpkvs()
                # * log to tensorboard
                for k, v in out.items():
                    self.writer.add_scalar(f'Loss/{k}', v,
                                           self.step + self.resume_step)

            # if self.step % self.eval_interval == 0 and self.step != 0:
            if self.step % self.eval_interval == 0:
                if dist_util.get_rank() == 0:
                    self.eval_ddpm_sample()
                #     continue # TODO, diffusion inference
                # self.eval_loop()
                # self.eval_novelview_loop()
                # let all processes sync up before starting with a new epoch of training
                dist_util.synchronize()
                th.cuda.empty_cache()

            if self.step % self.save_interval == 0 and self.step != 0:
                self.save()
                if not self.freeze_ae:
                    self.save(self.mp_trainer_rec, 'rec')
                dist_util.synchronize()

                th.cuda.empty_cache()
                # Run for a finite amount of time in integration tests.
                if os.environ.get("DIFFUSION_TRAINING_TEST",
                                  "") and self.step > 0:
                    return

            self.step += 1

            if self.step > self.iterations:
                print('reached maximum iterations, exiting')

                # Save the last checkpoint if it wasn't already saved.
                if (self.step - 1) % self.save_interval != 0:
                    self.save()
                    if not self.freeze_ae:
                        self.save(self.mp_trainer_rec, 'rec')

                exit()

        # Save the last checkpoint if it wasn't already saved.
        if (self.step - 1) % self.save_interval != 0:
            self.save()
            if not self.freeze_ae:
                self.save(self.mp_trainer_rec, 'rec')

    def run_step(self, batch, cond=None):
        self.forward_backward(batch,
                              cond)  # type: ignore # * 3D Reconstruction step
        took_step_ddpm = self.mp_trainer.optimize(self.opt)
        if took_step_ddpm:
            self._update_ema()

        if not self.freeze_ae:
            took_step_rec = self.mp_trainer_rec.optimize(self.opt_rec)
            if took_step_rec:
                self._update_ema_rec()

        self._anneal_lr()
        self.log_step()

    def forward_backward(self, batch, *args, **kwargs):
        # return super().forward_backward(batch, *args, **kwargs)
        self.mp_trainer.zero_grad()
        # all_denoised_out = dict()
        batch_size = batch['img'].shape[0]

        for i in range(0, batch_size, self.microbatch):

            micro = {
                k: v[i:i + self.microbatch].to(dist_util.dev())
                for k, v in batch.items()
            }

            last_batch = (i + self.microbatch) >= batch_size

            # if not freeze_ae:

            # =================================== ae part ===================================
            # with th.cuda.amp.autocast(dtype=th.float16,
            #                           enabled=self.mp_trainer_rec.use_amp
            #                           and not self.freeze_ae):
            with th.cuda.amp.autocast(dtype=th.float16,
                                      enabled=False,): # ! debugging, no AMP on all the input 

                pred = self.ddp_rec_model(img=micro['img_to_encoder'],
                                          c=micro['c'])  # pred: (B, 3, 64, 64)
                if not self.freeze_ae:
                    target = micro

                    if last_batch or not self.use_ddp:
                        ae_loss, loss_dict = self.loss_class(pred,
                                                             target,
                                                             test_mode=False)
                    else:
                        with self.ddp_model.no_sync():  # type: ignore
                            ae_loss, loss_dict = self.loss_class(
                                pred, target, test_mode=False)

                    log_rec3d_loss_dict(loss_dict)
                else:
                    ae_loss = th.tensor(0.0).to(dist_util.dev())

                micro_to_denoise = pred[
                    'latent'] / self.triplane_scaling_divider  # normalize std to 1

                t, weights = self.schedule_sampler.sample(
                    micro_to_denoise.shape[0], dist_util.dev())

                # print('!!!', micro_to_denoise.dtype)
                # =================================== denoised part ===================================

                model_kwargs = {}

                # print(micro_to_denoise.min(), micro_to_denoise.max())
                compute_losses = functools.partial(
                    self.diffusion.training_losses,
                    self.ddp_model,
                    micro_to_denoise,  # x_start
                    t,
                    model_kwargs=model_kwargs,
                )

                denoised_fn = functools.partial(
                    self.diffusion.p_mean_variance,
                    self.ddp_model,
                    micro_to_denoise,  # x_start
                    t,
                    model_kwargs=model_kwargs)

            with th.cuda.amp.autocast(dtype=th.float16,
                                      enabled=self.mp_trainer.use_amp):

                if last_batch or not self.use_ddp:
                    losses = compute_losses()
                    denoised_out = denoised_fn()
                else:
                    with self.ddp_model.no_sync():  # type: ignore
                        losses = compute_losses()
                        denoised_out = denoised_fn()

                if isinstance(self.schedule_sampler, LossAwareSampler):
                    self.schedule_sampler.update_with_local_losses(
                        t, losses["loss"].detach())

                denoise_loss = (losses["loss"] * weights).mean()
                log_loss_dict(self.diffusion, t,
                              {k: v * weights
                               for k, v in losses.items()})

                # self.mp_trainer.backward(denoise_loss)
                # =================================== denosied ae part ===================================
                # if self.denoised_ae or self.step % 500 == 0:
                if self.denoised_ae:
                    with th.cuda.amp.autocast(
                            dtype=th.float16,
                            enabled=self.mp_trainer_rec.use_amp
                            and not self.freeze_ae):
                        # continue

                        denoised_ae_pred = self.ddp_rec_model(
                            img=None,
                            c=micro['c'],
                            latent=denoised_out['pred_xstart'] * self.
                            triplane_scaling_divider,  # TODO, how to define the scale automatically?
                            behaviour='triplane_dec')

                        # if self.denoised_ae:

                        if last_batch or not self.use_ddp:
                            denoised_ae_loss, loss_dict = self.loss_class(
                                denoised_ae_pred, micro, test_mode=False)
                        else:
                            with self.ddp_model.no_sync():  # type: ignore
                                denoised_ae_loss, loss_dict = self.loss_class(
                                    denoised_ae_pred, micro, test_mode=False)

                        # * rename
                        loss_dict_denoise_ae = {}
                        for k, v in loss_dict.items():
                            loss_dict_denoise_ae[f'{k}_denoised'] = v.mean()
                        log_rec3d_loss_dict(loss_dict_denoise_ae)

                else:
                    denoised_ae_loss = th.tensor(0.0).to(dist_util.dev())

                # loss = ae_loss + denoise_loss + denoised_ae_loss
                loss = denoise_loss # ! leave only denoise_loss for debugging
                # loss = ae_loss + denoise_loss
                # self.mp_trainer.backward(denosied_ae_loss)
                # if use_amp:
                # self.mp_trainer.backward(loss)
                # self.mp_trainer.scaler.scale(loss).backward()
                # else:

            # exit AMP before backward
            self.mp_trainer.backward(loss)

            # TODO, merge visualization with original AE
            # =================================== denoised AE log part ===================================

            if dist_util.get_rank() == 0 and self.step % 500 == 0:
                with th.no_grad():
                    # gt_vis = th.cat([batch['img'], batch['depth']], dim=-1)

                    gt_depth = micro['depth']
                    if gt_depth.ndim == 3:
                        gt_depth = gt_depth.unsqueeze(1)
                    gt_depth = (gt_depth - gt_depth.min()) / (gt_depth.max() -
                                                              gt_depth.min())
                    # if True:
                    pred_depth = pred['image_depth']
                    pred_depth = (pred_depth - pred_depth.min()) / (
                        pred_depth.max() - pred_depth.min())
                    pred_img = pred['image_raw']
                    gt_img = micro['img']

                    if 'image_sr' in pred:  # TODO
                        pred_img = th.cat(
                            [self.pool_512(pred_img), pred['image_sr']],
                            dim=-1)
                        gt_img = th.cat(
                            [self.pool_512(micro['img']), micro['img_sr']],
                            dim=-1)
                        pred_depth = self.pool_512(pred_depth)
                        gt_depth = self.pool_512(gt_depth)

                    gt_vis = th.cat(
                        [
                            gt_img, micro['img'],
                            gt_depth.repeat_interleave(3, dim=1)
                        ],
                        dim=-1)[0:1]  # TODO, fail to load depth. range [0, 1]

                    if not self.denoised_ae:
                        # continue

                        denoised_ae_pred = self.ddp_rec_model(
                            img=None,
                            c=micro['c'][0:1],
                            latent=denoised_out['pred_xstart'][0:1] * self.
                            triplane_scaling_divider,  # TODO, how to define the scale automatically
                            behaviour='triplane_dec')

                    # assert denoised_ae_pred is not None

                    # print(pred_img.shape)
                    # print('denoised_ae:', self.denoised_ae)

                    pred_vis = th.cat([
                        pred_img[0:1], denoised_ae_pred['image_raw'],
                        pred_depth[0:1].repeat_interleave(3, dim=1)
                    ],
                                      dim=-1)  # B, 3, H, W

                    vis = th.cat([gt_vis, pred_vis], dim=-2)[0].permute(
                        1, 2, 0).cpu()  # ! pred in range[-1, 1]
                    # vis_grid = torchvision.utils.make_grid(vis) # HWC
                    vis = vis.numpy() * 127.5 + 127.5
                    vis = vis.clip(0, 255).astype(np.uint8)
                    Image.fromarray(vis).save(
                        f'{logger.get_dir()}/{self.step+self.resume_step}denoised_{t[0].item()}.jpg'
                    )
                    print(
                        'log denoised vis to: ',
                        f'{logger.get_dir()}/{self.step+self.resume_step}denoised_{t[0].item()}.jpg'
                    )

                    th.cuda.empty_cache()

    @th.no_grad()
    # def eval_loop(self, c_list:list):
    def eval_novelview_loop(self):
        # novel view synthesis given evaluation camera trajectory
        video_out = imageio.get_writer(
            f'{logger.get_dir()}/video_novelview_{self.step+self.resume_step}.mp4',
            mode='I',
            fps=60,
            codec='libx264')

        all_loss_dict = []
        novel_view_micro = {}

        # for i in range(0, len(c_list), 1): # TODO, larger batch size for eval
        for i, batch in enumerate(tqdm(self.eval_data)):
            # for i in range(0, 8, self.microbatch):
            # c = c_list[i].to(dist_util.dev()).reshape(1, -1)
            micro = {k: v.to(dist_util.dev()) for k, v in batch.items()}

            if i == 0:
                novel_view_micro = {
                    k: v[0:1].to(dist_util.dev()).repeat_interleave(
                        micro['img'].shape[0], 0)
                    for k, v in batch.items()
                }
            else:
                # if novel_view_micro['c'].shape[0] < micro['img'].shape[0]:
                novel_view_micro = {
                    k: v[0:1].to(dist_util.dev()).repeat_interleave(
                        micro['img'].shape[0], 0)
                    for k, v in novel_view_micro.items()
                }

            pred = self.model(img=novel_view_micro['img_to_encoder'],
                              c=micro['c'])  # pred: (B, 3, 64, 64)
            # target = {
            #     'img': micro['img'],
            #     'depth': micro['depth'],
            #     'depth_mask': micro['depth_mask']
            # }
            # targe

            _, loss_dict = self.loss_class(pred, micro, test_mode=True)
            all_loss_dict.append(loss_dict)

            # ! move to other places, add tensorboard

            # pred_vis = th.cat([
            #     pred['image_raw'],
            #     -pred['image_depth'].repeat_interleave(3, dim=1)
            # ],
            #                   dim=-1)

            # normalize depth
            # if True:
            pred_depth = pred['image_depth']
            pred_depth = (pred_depth - pred_depth.min()) / (pred_depth.max() -
                                                            pred_depth.min())
            if 'image_sr' in pred:
                pred_vis = th.cat([
                    micro['img_sr'],
                    self.pool_512(pred['image_raw']), pred['image_sr'],
                    self.pool_512(pred_depth).repeat_interleave(3, dim=1)
                ],
                                  dim=-1)
            else:
                pred_vis = th.cat([
                    self.pool_128(micro['img']), pred['image_raw'],
                    pred_depth.repeat_interleave(3, dim=1)
                ],
                                  dim=-1)  # B, 3, H, W

            vis = pred_vis.permute(0, 2, 3, 1).cpu().numpy()
            vis = vis * 127.5 + 127.5
            vis = vis.clip(0, 255).astype(np.uint8)

            for j in range(vis.shape[0]):
                video_out.append_data(vis[j])

        video_out.close()

        val_scores_for_logging = calc_average_loss(all_loss_dict)
        with open(os.path.join(logger.get_dir(), 'scores_novelview.json'),
                  'a') as f:
            json.dump({'step': self.step, **val_scores_for_logging}, f)

        # * log to tensorboard
        for k, v in val_scores_for_logging.items():
            self.writer.add_scalar(f'Eval/NovelView/{k}', v,
                                   self.step + self.resume_step)

    @th.no_grad()
    # def eval_loop(self, c_list:list):
    def eval_loop(self):
        # novel view synthesis given evaluation camera trajectory
        video_out = imageio.get_writer(
            f'{logger.get_dir()}/video_{self.step+self.resume_step}.mp4',
            mode='I',
            fps=60,
            codec='libx264')
        all_loss_dict = []

        # for i in range(0, len(c_list), 1): # TODO, larger batch size for eval
        for i, batch in enumerate(tqdm(self.eval_data)):
            # for i in range(0, 8, self.microbatch):
            # c = c_list[i].to(dist_util.dev()).reshape(1, -1)
            micro = {k: v.to(dist_util.dev()) for k, v in batch.items()}

            # pred = self.model(img=micro['img_to_encoder'],
            #                   c=micro['c'])  # pred: (B, 3, 64, 64)

            # pred of rec model
            pred = self.ddp_rec_model(img=micro['img_to_encoder'],
                                      c=micro['c'])  # pred: (B, 3, 64, 64)

            pred_depth = pred['image_depth']
            pred_depth = (pred_depth - pred_depth.min()) / (pred_depth.max() -
                                                            pred_depth.min())

            if 'image_sr' in pred:
                pred_vis = th.cat([
                    micro['img_sr'],
                    self.pool_512(pred['image_raw']), pred['image_sr'],
                    self.pool_512(pred_depth).repeat_interleave(3, dim=1)
                ],
                                  dim=-1)
            else:
                pred_vis = th.cat([
                    self.pool_128(micro['img']), pred['image_raw'],
                    pred_depth.repeat_interleave(3, dim=1)
                ],
                                  dim=-1)  # B, 3, H, W

            vis = pred_vis.permute(0, 2, 3, 1).cpu().numpy()
            vis = vis * 127.5 + 127.5
            vis = vis.clip(0, 255).astype(np.uint8)

            for j in range(vis.shape[0]):
                video_out.append_data(vis[j])

        video_out.close()

        val_scores_for_logging = calc_average_loss(all_loss_dict)
        with open(os.path.join(logger.get_dir(), 'scores.json'), 'a') as f:
            json.dump({'step': self.step, **val_scores_for_logging}, f)

        # * log to tensorboard
        for k, v in val_scores_for_logging.items():
            self.writer.add_scalar(f'Eval/Rec/{k}', v,
                                   self.step + self.resume_step)

        self.eval_novelview_loop()

    def save(self, mp_trainer=None, model_name='ddpm'):
        if mp_trainer is None:
            mp_trainer = self.mp_trainer

        def save_checkpoint(rate, params):
            state_dict = mp_trainer.master_params_to_state_dict(params)
            if dist_util.get_rank() == 0:
                logger.log(f"saving model {model_name} {rate}...")
                if not rate:
                    filename = f"model_{model_name}{(self.step+self.resume_step):07d}.pt"
                else:
                    filename = f"ema_{model_name}_{rate}_{(self.step+self.resume_step):07d}.pt"
                with bf.BlobFile(bf.join(get_blob_logdir(), filename),
                                 "wb") as f:
                    th.save(state_dict, f)

        save_checkpoint(0, self.mp_trainer.master_params)
        for rate, params in zip(self.ema_rate, self.ema_params):
            save_checkpoint(rate, params)

        dist.barrier()

    def _load_and_sync_parameters(self, model=None, model_name='ddpm'):
        resume_checkpoint, self.resume_step = find_resume_checkpoint(
            self.resume_checkpoint, model_name) or self.resume_checkpoint

        if model is None:
            model = self.model
        print(resume_checkpoint)

        if resume_checkpoint and Path(resume_checkpoint).exists():
            if dist_util.get_rank() == 0:

                # ! rank 0 return will cause all other ranks to hang
                # if not Path(resume_checkpoint).exists():
                #     logger.log(
                #         f"failed to load model from checkpoint: {resume_checkpoint}, not exist"
                #     )
                #     return

                logger.log(
                    f"loading model from checkpoint: {resume_checkpoint}...")
                map_location = {
                    'cuda:%d' % 0: 'cuda:%d' % dist_util.get_rank()
                }  # configure map_location properly

                print(f'mark {model_name} loading ', flush=True)
                resume_state_dict = dist_util.load_state_dict(
                    resume_checkpoint, map_location=map_location)
                print(f'mark {model_name} loading finished', flush=True)

                model_state_dict = model.state_dict()

                for k, v in resume_state_dict.items():
                    if k in model_state_dict.keys() and v.size(
                    ) == model_state_dict[k].size():
                        model_state_dict[k] = v
                    else:
                        print('!!!! ignore key: ', k, ": ", v.size(),
                              'shape in model: ', model_state_dict[k].size())

                model.load_state_dict(model_state_dict, strict=True)
                del model_state_dict

        if dist_util.get_world_size() > 1:
            dist_util.sync_params(model.parameters())
            print(f'synced {model_name} params')

    def _update_ema_rec(self):
        for rate, params in zip(self.ema_rate, self.ema_params_rec):
            update_ema(params, self.mp_trainer_rec.master_params, rate=rate)

    def eval_ddpm_sample(self):

        args = dnnlib.EasyDict(
            dict(batch_size=1,
                 image_size=224,
                 denoise_in_channels=24,
                 clip_denoised=True,
                 class_cond=False,
                 use_ddim=False))

        model_kwargs = {}

        if args.class_cond:
            classes = th.randint(low=0,
                                 high=NUM_CLASSES,
                                 size=(args.batch_size, ),
                                 device=dist_util.dev())
            model_kwargs["y"] = classes

        diffusion = self.diffusion
        sample_fn = (diffusion.p_sample_loop
                     if not args.use_ddim else diffusion.ddim_sample_loop)

        for i in range(2):
            triplane_sample = sample_fn(
                self.ddp_model,
                (args.batch_size, args.denoise_in_channels, args.image_size,
                 args.image_size),
                clip_denoised=args.clip_denoised,
                model_kwargs=model_kwargs,
            )

            self.render_video_given_triplane(
                triplane_sample,
                name_prefix=f'{self.step + self.resume_step}_{i}')

    @th.inference_mode()
    def render_video_given_triplane(self, planes, name_prefix='0'):

        planes *= self.triplane_scaling_divider  # if setting clip_denoised=True, the sampled planes will lie in [-1,1]. Thus, values beyond [+- std] will be abandoned in this version. Move to IN for later experiments.

        # print(planes.min(), planes.max())

        # used during diffusion sampling inference
        video_out = imageio.get_writer(
            f'{logger.get_dir()}/triplane_{name_prefix}.mp4',
            mode='I',
            fps=60,
            codec='libx264')

        # for i in range(0, len(c_list), 1): # TODO, larger batch size for eval
        for i, batch in enumerate(tqdm(self.eval_data)):
            micro = {k: v.to(dist_util.dev()) for k, v in batch.items()}

            pred = self.ddp_rec_model(img=None,
                                      c=micro['c'],
                                      latent=planes,
                                      behaviour='triplane_dec')

            # if True:
            pred_depth = pred['image_depth']
            pred_depth = (pred_depth - pred_depth.min()) / (pred_depth.max() -
                                                            pred_depth.min())

            if 'image_sr' in pred:
                pred_vis = th.cat([
                    micro['img_sr'],
                    self.pool_512(pred['image_raw']), pred['image_sr'],
                    self.pool_512(pred_depth).repeat_interleave(3, dim=1)
                ],
                                  dim=-1)
            else:
                pred_vis = th.cat([
                    self.pool_128(micro['img']), pred['image_raw'],
                    pred_depth.repeat_interleave(3, dim=1)
                ],
                                  dim=-1)  # B, 3, H, W

            vis = pred_vis.permute(0, 2, 3, 1).cpu().numpy()
            vis = vis * 127.5 + 127.5
            vis = vis.clip(0, 255).astype(np.uint8)

            for j in range(vis.shape[0]):
                video_out.append_data(vis[j])

        video_out.close()
        print('logged video to: ',
              f'{logger.get_dir()}/triplane_{name_prefix}.mp4')

    @th.inference_mode()
    def render_video_noise_schedule(self, name_prefix='0'):

        # planes *= self.triplane_std # denormalize for rendering

        video_out = imageio.get_writer(
            f'{logger.get_dir()}/triplane_visnoise_{name_prefix}.mp4',
            mode='I',
            fps=30,
            codec='libx264')

        for i, batch in enumerate(tqdm(self.eval_data)):
            micro = {k: v.to(dist_util.dev()) for k, v in batch.items()}

            if i % 10 != 0:
                continue

            # ========= novel view plane settings ====
            if i == 0:
                novel_view_micro = {
                    k: v[0:1].to(dist_util.dev()).repeat_interleave(
                        micro['img'].shape[0], 0)
                    for k, v in batch.items()
                }
            else:
                # if novel_view_micro['c'].shape[0] < micro['img'].shape[0]:
                novel_view_micro = {
                    k: v[0:1].to(dist_util.dev()).repeat_interleave(
                        micro['img'].shape[0], 0)
                    for k, v in novel_view_micro.items()
                }

            latent = self.ddp_rec_model(
                img=novel_view_micro['img_to_encoder'],
                c=micro['c'])['latent']  # pred: (B, 3, 64, 64)

            x_start = latent / self.triplane_scaling_divider  # normalize std to 1
            # x_start = latent

            all_pred_vis = []
            # for t in th.range(0,
            #                   4001,
            #                   500,
            #                   dtype=th.long,
            #                   device=dist_util.dev()):  # cosine 4k steps
            for t in th.range(0,
                              1001,
                              125,
                              dtype=th.long,
                              device=dist_util.dev()):  # cosine 4k steps

                # ========= add noise according to t
                noise = th.randn_like(x_start)  # x_start is the x0 image
                x_t = self.diffusion.q_sample(
                    x_start, t, noise=noise
                )  # * add noise according to predefined schedule
                planes_x_t = (x_t * self.triplane_scaling_divider).clamp(
                    -50, 50)  # de-scaling noised x_t

                # planes_x_t = (x_t * 1).clamp(
                #     -50, 50)  # de-scaling noised x_t

                # ===== visualize
                pred = self.ddp_rec_model(
                    img=None,
                    c=micro['c'],
                    latent=planes_x_t,
                    behaviour='triplane_dec')  # pred: (B, 3, 64, 64)

                # pred_depth = pred['image_depth']
                # pred_depth = (pred_depth - pred_depth.min()) / (
                #     pred_depth.max() - pred_depth.min())
                # pred_vis = th.cat([
                #     # self.pool_128(micro['img']),
                #     pred['image_raw'],
                # ],
                #                   dim=-1)  # B, 3, H, W
                pred_vis = pred['image_raw']

                all_pred_vis.append(pred_vis)
                # TODO, make grid

            all_pred_vis = torchvision.utils.make_grid(
                th.cat(all_pred_vis, 0),
                nrow=len(all_pred_vis),
                normalize=True,
                value_range=(-1, 1),
                scale_each=True)  # normalized to [-1,1]

            vis = all_pred_vis.permute(1, 2, 0).cpu().numpy()  # H W 3

            vis = (vis * 255).clip(0, 255).astype(np.uint8)

            video_out.append_data(vis)

        video_out.close()
        print('logged video to: ',
              f'{logger.get_dir()}/triplane_visnoise_{name_prefix}.mp4')

        th.cuda.empty_cache()