nsr/lsgm/train_util_diffusion_lsgm_noD.py

"""
Modified from:
https://github.com/NVlabs/LSGM/blob/main/training_obj_joint.py
"""
import copy
import functools
import json
import os
from pathlib import Path
from pdb import set_trace as st
from typing import Any

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)
from guided_diffusion.gaussian_diffusion import ModelMeanType

import dnnlib
from dnnlib.util import requires_grad
from dnnlib.util import calculate_adaptive_weight

from ..train_util_diffusion import TrainLoop3DDiffusion
from ..cvD.nvsD_canoD import TrainLoop3DcvD_nvsD_canoD

from guided_diffusion.continuous_diffusion_utils import get_mixed_prediction, different_p_q_objectives, kl_per_group_vada, kl_balancer
# import utils as lsgm_utils


class TrainLoop3DDiffusionLSGM_noD(TrainLoop3DDiffusion):
    def __init__(self,
                 *,
                 rec_model,
                 denoise_model,
                 diffusion,
                 sde_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,
                 diffusion_input_size=224,
                 **kwargs):
        super().__init__(
            rec_model=rec_model,
            denoise_model=denoise_model,
            diffusion=diffusion,
            loss_class=loss_class,
            data=data,
            eval_data=eval_data,
            batch_size=batch_size,
            microbatch=microbatch,
            lr=lr,
            ema_rate=ema_rate,
            log_interval=log_interval,
            eval_interval=eval_interval,
            save_interval=save_interval,
            resume_checkpoint=resume_checkpoint,
            use_fp16=use_fp16,
            fp16_scale_growth=fp16_scale_growth,
            schedule_sampler=schedule_sampler,
            weight_decay=weight_decay,
            lr_anneal_steps=lr_anneal_steps,
            iterations=iterations,
            ignore_resume_opt=ignore_resume_opt,
            #  freeze_ae=freeze_ae,
            freeze_ae=not sde_diffusion.args.train_vae,
            denoised_ae=denoised_ae,
            triplane_scaling_divider=triplane_scaling_divider,
            use_amp=use_amp,
            diffusion_input_size=diffusion_input_size,
            **kwargs)

        assert sde_diffusion is not None
        sde_diffusion.args.batch_size = batch_size
        self.sde_diffusion = sde_diffusion
        self.latent_name = 'latent_normalized_2Ddiffusion'  # normalized triplane latent
        self.render_latent_behaviour = 'decode_after_vae'  # directly render using triplane operations

        self.pool_512 = th.nn.AdaptiveAvgPool2d((512, 512))
        self.pool_256 = th.nn.AdaptiveAvgPool2d((256, 256))
        self.pool_128 = th.nn.AdaptiveAvgPool2d((128, 128))
        self.pool_64 = th.nn.AdaptiveAvgPool2d((64, 64))

        self.ddp_ddpm_model = self.ddp_model

        # if sde_diffusion.args.joint_train:
        # assert sde_diffusion.args.train_vae

    def run_step(self, batch, step='diffusion_step_rec'):

        # if step == 'diffusion_step_rec':

        self.forward_diffusion(batch, behaviour='diffusion_step_rec')

        # if took_step_ddpm:
        self._update_ema()

        self._anneal_lr()
        self.log_step()

    def run_loop(self):
        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 = next(self.data)
            self.run_step(batch, step='diffusion_step_rec')

            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()
                    if self.sde_diffusion.args.train_vae:
                        self.eval_loop()

                th.cuda.empty_cache()
                dist_util.synchronize()

            if self.step % self.save_interval == 0:
                self.save(self.mp_trainer, self.mp_trainer.model_name)
                if self.sde_diffusion.args.train_vae:
                    self.save(self.mp_trainer_rec,
                              self.mp_trainer_rec.model_name)

                # dist_util.synchronize()
                # 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(self.mp_trainer, self.mp_trainer.model_name)
                    if self.sde_diffusion.args.train_vae:
                        self.save(self.mp_trainer_rec,
                                  self.mp_trainer_rec.model_name)

                exit()

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

    # ! duplicated code, needs refactor later
    def ddpm_step(self, eps, t, logsnr, model_kwargs={}):
        """helper function for ddpm predictions; returns predicted eps, x0 and logsnr
        """
        args = self.sde_diffusion.args
        pred_params = self.ddp_ddpm_model(eps, t, **model_kwargs)
        # pred_params = self.ddp_model(eps, t, **model_kwargs)
        if args.pred_type == 'eps':
            pred_eps = pred_params
            pred_x0 = self.sde_diffusion._predict_x0_from_eps(
                eps, pred_params, logsnr)  # for VAE loss, denosied latent
        elif args.pred_type == 'x0':
            # ! transform to pred_eps format for mixing_component
            pred_x0 = pred_params
            pred_eps = self.sde_diffusion._predict_eps_from_x0(
                eps, pred_params, logsnr)
        else:
            raise NotImplementedError(f'{args.pred_type} not implemented.')

        return pred_eps, pred_x0, logsnr

    # def apply_model(self, p_sample_batch, model_kwargs={}):
    #     # args = self.sde_diffusion.args
    #     noise, eps_t_p, t_p, logsnr_p, obj_weight_t_p, var_t_p = (
    #         p_sample_batch[k] for k in ('noise', 'eps_t_p', 't_p', 'logsnr_p',
    #                                     'obj_weight_t_p', 'var_t_p'))

    #     pred_eps_p, pred_x0_p, logsnr_p = self.ddpm_step(
    #         eps_t_p, t_p, logsnr_p, model_kwargs)

    #     # ! batchify for mixing_component
    #     # mixing normal trick
    #     mixing_component = self.sde_diffusion.mixing_component(
    #         eps_t_p, var_t_p, t_p, enabled=True)  # TODO, which should I use?
    #     pred_eps_p = get_mixed_prediction(
    #         True, pred_eps_p,
    #         self.ddp_ddpm_model(x=None,
    #                             timesteps=None,
    #                             get_attr='mixing_logit'), mixing_component)

    #     # ! eps loss equivalent to snr weighting of x0 loss, see "progressive distillation"
    #     with self.ddp_ddpm_model.no_sync():  # type: ignore
    #         l2_term_p = th.square(pred_eps_p - noise)  # ? weights

    #     p_eps_objective = th.mean(obj_weight_t_p * l2_term_p)

    #     log_rec3d_loss_dict(
    #         dict(mixing_logit=self.ddp_ddpm_model(
    #             x=None, timesteps=None, get_attr='mixing_logit').detach(), ))

    #     return {
    #         'pred_eps_p': pred_eps_p,
    #         'eps_t_p': eps_t_p,
    #         'p_eps_objective': p_eps_objective,
    #         'pred_x0_p': pred_x0_p,
    #         'logsnr_p': logsnr_p
    #     }

    def forward_diffusion(self, batch, behaviour='rec', *args, **kwargs):
        """
        add sds grad to all ae predicted x_0 
        """
        args = self.sde_diffusion.args

        # self.ddp_ddpm_model.requires_grad_(True)
        requires_grad(self.ddp_rec_model.module, args.train_vae)
        # self.ddp_rec_model.requires_grad_(args.train_vae)

        if args.train_vae:
            for param in self.ddp_rec_model.module.decoder.triplane_decoder.parameters(  # type: ignore
            ):  # type: ignore
                param.requires_grad_(
                    False
                )  # ! disable triplane_decoder grad in each iteration indepenently;

        self.mp_trainer_rec.zero_grad()
        self.mp_trainer.zero_grad()

        batch_size = batch['img'].shape[0]

        # # update ddpm params
        # took_step_ddpm = self.mp_trainer_ddpm.optimize(
        #     self.opt_ddpm)  # TODO, update two groups of parameters

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

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

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

            q_vae_recon_loss = th.tensor(0.0).to(dist_util.dev())
            # vision_aided_loss = th.tensor(0.0).to(dist_util.dev())
            # denoise_loss = th.tensor(0.0).to(dist_util.dev())

            # =================================== ae part ===================================
            with th.cuda.amp.autocast(dtype=th.float16,
                                      enabled=self.mp_trainer.use_amp):
                #   and args.train_vae):

                assert behaviour == 'diffusion_step_rec'

                # ! train vae with CE; ddpm fixed
                requires_grad(self.ddp_model.module, False)
                # if args.train_vae:
                #     assert args.add_rendering_loss
                with th.set_grad_enabled(args.train_vae):
                    vae_out = self.ddp_rec_model(
                        img=micro['img_to_encoder'],
                        c=micro['c'],
                        # behaviour='enc_dec_wo_triplane'
                        behaviour='encoder_vae',
                    )  # pred: (B, 3, 64, 64)
                    # TODO, no need to render if not SSD; no need to do ViT decoder if only the latent is needed. update later

                # TODO, train diff and sds together, available?
                all_log_q = [vae_out['log_q_2Ddiffusion']]
                eps = vae_out[self.latent_name]
                eps.requires_grad_(True)  # single stage diffusion

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

                noise = th.randn(
                    size=eps.size(), device=eps.device
                )  # note that this noise value is currently shared!
                model_kwargs = {}

                # get diffusion quantities for p (sgm prior) sampling scheme and reweighting for q (vae)
                t_p, var_t_p, m_t_p, obj_weight_t_p, obj_weight_t_q, g2_t_p = \
                    self.sde_diffusion.iw_quantities(args.iw_sample_p)
                eps_t_p = self.sde_diffusion.sample_q(eps, noise, var_t_p,
                                                      m_t_p)
                logsnr_p = self.sde_diffusion.log_snr(m_t_p,
                                                      var_t_p)  # for p only

                # in case we want to train q (vae) with another batch using a different sampling scheme for times t
                if args.iw_sample_q in ['ll_uniform', 'll_iw']:
                    t_q, var_t_q, m_t_q, obj_weight_t_q, _, g2_t_q = \
                        self.sde_diffusion.iw_quantities(args.iw_sample_q)
                    eps_t_q = self.sde_diffusion.sample_q(
                        eps, noise, var_t_q, m_t_q)

                    eps_t_p = eps_t_p.detach().requires_grad_(
                        True)  # ! p just not updated here
                    eps_t = th.cat([eps_t_p, eps_t_q], dim=0)
                    var_t = th.cat([var_t_p, var_t_q], dim=0)
                    t = th.cat([t_p, t_q], dim=0)
                    noise = th.cat([noise, noise], dim=0)
                    # logsnr = self.sde_diffusion.log_snr(m_t_q, var_t_p)
                else:
                    eps_t, m_t, var_t, t, g2_t = eps_t_p, m_t_p, var_t_p, t_p, g2_t_p

                # run the diffusion model
                eps_t.requires_grad_(True)  # 2*BS, 12, 16, 16
                pred_params = self.ddp_model(eps_t, t, **model_kwargs)

                if args.pred_type == 'eps':
                    pred_eps = pred_params
                elif args.pred_type == 'x0':
                    # ! transform to pred_eps format for mixing_component
                    pred_eps = self.sde_diffusion._predict_eps_from_x0(
                        eps_t, pred_params, logsnr_p)
                else:
                    raise NotImplementedError(
                        f'{args.pred_type} not implemented.')

                # mixing normal trick
                mixing_component = self.sde_diffusion.mixing_component(
                    eps_t, var_t, t, enabled=True)  # TODO, which should I use?
                pred_eps = get_mixed_prediction(
                    # True, pred_params,
                    True,
                    pred_eps,
                    self.ddp_model(x=None,
                                   timesteps=None,
                                   get_attr='mixing_logit'),
                    mixing_component)

                # ! eps loss equivalent to snr weighting of x0 loss, see "progressive distillation"
                if last_batch or not self.use_ddp:
                    l2_term = th.square(pred_eps - noise)
                else:
                    with self.ddp_model.no_sync():  # type: ignore
                        l2_term = th.square(pred_eps - noise)  # ? weights

                # nelbo loss with kl balancing
                # ! remainign parts of cross entropy in likelihook training
                # unpack separate objectives, in case we want to train q (vae) using a different sampling scheme for times t
                if args.iw_sample_q in ['ll_uniform',
                                        'll_iw']:  # ll_iw by default
                    l2_term_p, l2_term_q = th.chunk(l2_term, chunks=2, dim=0)
                    p_objective = th.mean(obj_weight_t_p * l2_term_p,
                                          dim=[1, 2, 3])
                    cross_entropy_per_var = obj_weight_t_q * l2_term_q
                else:
                    p_objective = th.mean(obj_weight_t_p * l2_term,
                                          dim=[1, 2, 3])
                    cross_entropy_per_var = obj_weight_t_q * l2_term

                cross_entropy_per_var += self.sde_diffusion.cross_entropy_const(
                    args.sde_time_eps)
                all_neg_log_p = [cross_entropy_per_var
                                 ]  # since only one vae group

                kl_all_list, kl_vals_per_group, kl_diag_list = kl_per_group_vada(
                    all_log_q, all_neg_log_p)  # return the mean of two terms

                # nelbo loss with kl balancing
                balanced_kl, kl_coeffs, kl_vals = kl_balancer(kl_all_list,
                                                              kl_coeff=1.0,
                                                              kl_balance=False,
                                                              alpha_i=None)

                # ! update vae for CE
                # ! single stage diffusion for rec side 1: bind vae prior and diffusion prior
                if args.train_vae:
                    # if args.add_rendering_loss:
                    # if args.joint_train:
                    with th.set_grad_enabled(args.train_vae):
                        target = micro
                        pred = self.ddp_rec_model(
                            latent=vae_out,
                            # latent={
                            #     **vae_out, self.latent_name: pred_x0,
                            #     'latent_name': self.latent_name
                            # },
                            c=micro['c'],
                            behaviour=self.render_latent_behaviour)

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

                        log_rec3d_loss_dict(loss_dict)

                # ! calculate p/q loss;
                nelbo_loss = balanced_kl + q_vae_recon_loss
                q_loss = th.mean(nelbo_loss)
                p_loss = th.mean(p_objective)

                log_rec3d_loss_dict(
                    dict(
                        q_vae_recon_loss=q_vae_recon_loss,
                        p_loss=p_loss,
                        balanced_kl=balanced_kl,
                        mixing_logit=self.ddp_model(
                            x=None, timesteps=None,
                            get_attr='mixing_logit').detach(),
                    ))

                # ! single stage diffusion for rec side 2: generative feature
                if args.p_rendering_loss:
                    with th.set_grad_enabled(args.train_vae):

                        # ! transform fro pred_eps format back to pred_x0, for p only.
                        pred_x0 = self.sde_diffusion._predict_x0_from_eps(
                            eps_t_p, pred_eps[:eps_t_p.shape[0]],
                            logsnr_p)  # for VAE loss, denosied latent

                        target = micro
                        pred = self.ddp_rec_model(
                            # latent=vae_out,
                            latent={
                                **vae_out, self.latent_name: pred_x0,
                                'latent_name': self.latent_name
                            },
                            c=micro['c'],
                            behaviour=self.render_latent_behaviour)

                        # vae reconstruction loss
                        if last_batch or not self.use_ddp:
                            p_vae_recon_loss, loss_dict = self.loss_class(
                                pred, target, test_mode=False)
                        else:
                            with self.ddp_model.no_sync():  # type: ignore
                                p_vae_recon_loss, loss_dict = self.loss_class(
                                    pred, target, test_mode=False)
                        log_rec3d_loss_dict(
                            dict(p_vae_recon_loss=p_vae_recon_loss, ))

            # ! backpropagate q_loss for vae and update vae params, if trained
            if args.train_vae:
                self.mp_trainer_rec.backward(
                    q_loss,
                    retain_graph=different_p_q_objectives(
                        args.iw_sample_p, args.iw_sample_q))

            # if we use different p and q objectives or are not training the vae, discard gradients and backpropagate p_loss
            if different_p_q_objectives(
                    args.iw_sample_p, args.iw_sample_q) or not args.train_vae:
                if args.train_vae:
                    # discard current gradients computed by weighted loss for VAE
                    self.mp_trainer_rec.zero_grad()

                self.mp_trainer.backward(p_loss)

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

            if dist_util.get_rank(
            ) == 0 and self.step % 500 == 0 and behaviour != 'diff':

                with th.no_grad():

                    if not args.train_vae:
                        vae_out.pop('posterior')  # for calculating kl loss
                        vae_out_for_pred = {
                            k: v[0:1].to(dist_util.dev()) if isinstance(
                                v, th.Tensor) else v
                            for k, v in vae_out.items()
                        }

                        pred = self.ddp_rec_model(
                            latent=vae_out_for_pred,
                            c=micro['c'][0:1],
                            behaviour=self.render_latent_behaviour)
                    assert isinstance(pred, dict)
                    assert pred is not None

                    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())

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

                    if 'image_depth' in pred:
                        pred_depth = pred['image_depth']
                        pred_depth = (pred_depth - pred_depth.min()) / (
                            pred_depth.max() - pred_depth.min())
                    else:
                        pred_depth = th.zeros_like(gt_depth)

                    pred_img = pred['image_raw']
                    gt_img = micro['img']

                    if 'image_sr' in pred:
                        if pred['image_sr'].shape[-1] == 512:
                            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)

                        elif pred['image_sr'].shape[-1] == 256:
                            pred_img = th.cat(
                                [self.pool_256(pred_img), pred['image_sr']],
                                dim=-1)
                            gt_img = th.cat(
                                [self.pool_256(micro['img']), micro['img_sr']],
                                dim=-1)
                            pred_depth = self.pool_256(pred_depth)
                            gt_depth = self.pool_256(gt_depth)

                        else:
                            pred_img = th.cat(
                                [self.pool_128(pred_img), pred['image_sr']],
                                dim=-1)
                            gt_img = th.cat(
                                [self.pool_128(micro['img']), micro['img_sr']],
                                dim=-1)
                            gt_depth = self.pool_128(gt_depth)
                            pred_depth = self.pool_128(pred_depth)
                    else:
                        gt_img = self.pool_64(gt_img)
                        gt_depth = self.pool_64(gt_depth)

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

                    # eps_t_p_3D = eps_t_p.reshape(batch_size, eps_t_p.shape[1]//3, 3, -1) # B C 3 L

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

                    # ! test time, use discrete diffusion model
                    params_p, _ = th.chunk(pred_eps, chunks=2,
                                           dim=0)  # get predicted noise

                    # TODO, implement for SDE difusion?
                    # ! two values isclose(rtol=1e-03, atol=1e-04)
                    # pred_xstart = self.diffusion._predict_xstart_from_eps(
                    #     x_t=eps_t_p,
                    #     t=th.tensor(t_p.detach() *
                    #                 self.diffusion.num_timesteps).long(),
                    #     eps=params_p)

                    pred_x0 = self.sde_diffusion._predict_x0_from_eps(
                        eps_t_p, params_p,
                        logsnr_p)  # for VAE loss, denosied latent

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

                    pred_vis = th.cat([
                        pred_img[0:1], noised_ae_pred['image_raw'][0:1],
                        denoised_ae_pred['image_raw'][0:1],
                        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()}_{behaviour}.jpg'
                    )
                    print(
                        'log denoised vis to: ',
                        f'{logger.get_dir()}/{self.step+self.resume_step}denoised_{t[0].item()}_{behaviour}.jpg'
                    )
                    del vis, pred_vis, pred_x0, pred_eps, micro, vae_out

                    th.cuda.empty_cache()

    # ! copied from train_util.py
    # TODO, needs to lint the class inheritance chain later.
    @th.inference_mode()
    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.rec_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:

                if pred['image_sr'].shape[-1] == 512:

                    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)

                elif pred['image_sr'].shape[-1] == 256:

                    pred_vis = th.cat([
                        micro['img_sr'],
                        self.pool_256(pred['image_raw']), pred['image_sr'],
                        self.pool_256(pred_depth).repeat_interleave(3, dim=1)
                    ],
                                      dim=-1)

                else:
                    pred_vis = th.cat([
                        micro['img_sr'],
                        self.pool_128(pred['image_raw']),
                        self.pool_128(pred['image_sr']),
                        self.pool_128(pred_depth).repeat_interleave(3, dim=1)
                    ],
                                      dim=-1)

            else:
                pred_vis = th.cat([
                    self.pool_64(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)
        del video_out, vis, pred_vis, pred, micro
        th.cuda.empty_cache()

    # @th.no_grad()
    # def eval_loop(self, c_list:list):
    @th.inference_mode()
    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 = []
        self.rec_model.eval()

        # 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.rec_model(img=micro['img_to_encoder'],
                                  c=micro['c'])  # pred: (B, 3, 64, 64)
            # target = {
            #     'img': micro['img'],
            #     'depth': micro['depth'],
            #     'depth_mask': micro['depth_mask']
            # }

            # if last_batch or not self.use_ddp:
            #     loss, loss_dict = self.loss_class(pred, target)
            # else:
            #     with self.ddp_model.no_sync():  # type: ignore
            _, loss_dict = self.loss_class(pred, micro, test_mode=True)
            all_loss_dict.append(loss_dict)

            # ! move to other places, add tensorboard
            # gt_vis = th.cat([micro['img'], micro['img']], dim=-1) # TODO, fail to load depth. range [0, 1]
            # pred_vis = th.cat([
            #     pred['image_raw'],
            #     -pred['image_depth'].repeat_interleave(3, dim=1)
            # ],
            #                   dim=-1)
            # vis = th.cat([gt_vis, pred_vis], dim=-2)[0].permute(1,2,0).cpu().numpy() # ! pred in range[-1, 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:

                if pred['image_sr'].shape[-1] == 512:

                    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)

                elif pred['image_sr'].shape[-1] == 256:
                    pred_vis = th.cat([
                        micro['img_sr'],
                        self.pool_256(pred['image_raw']), pred['image_sr'],
                        self.pool_256(pred_depth).repeat_interleave(3, dim=1)
                    ],
                                      dim=-1)

                else:
                    pred_vis = th.cat([
                        micro['img_sr'],
                        self.pool_128(pred['image_raw']),
                        self.pool_128(pred['image_sr']),
                        self.pool_128(pred_depth).repeat_interleave(3, dim=1)
                    ],
                                      dim=-1)

            else:
                pred_vis = th.cat([
                    self.pool_64(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)

        del video_out, vis, pred_vis, pred, micro
        th.cuda.empty_cache()
        self.eval_novelview_loop()
        self.rec_model.train()

    # for compatablity with p_sample, to lint
    def apply_model_inference(self, x_noisy, t, c=None, model_kwargs={}):
        # control = self.ddp_control_model(x=x_noisy,
        #                                  hint=th.cat(c['c_concat'], 1),
        #                                  timesteps=t,
        #                                  context=None)
        # control = [c * scale for c, scale in zip(control, self.control_scales)]
        pred_params = self.ddp_ddpm_model(x_noisy, t, 
                                    **model_kwargs
                                    )


        assert args.pred_type == 'eps'
        # mixing normal trick
        mixing_component = self.sde_diffusion.mixing_component(
            eps, var_t, t, enabled=True)  # TODO, which should I use?
        pred_eps = get_mixed_prediction(
            True, pred_eps,
            self.ddp_ddpm_model(x=None, timesteps=None, get_attr='mixing_logit'), mixing_component)

        return pred_params

    @th.inference_mode()
    def eval_ddpm_sample(self):

        args = dnnlib.EasyDict(
            dict(
                batch_size=1,
                image_size=self.diffusion_input_size,
                denoise_in_channels=self.ddp_rec_model.module.decoder.
                triplane_decoder.out_chans,  # type: ignore
                clip_denoised=False,
                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(1):
            triplane_sample = sample_fn(
                # self.ddp_model,
                self,
                (
                    args.batch_size,
                    self.ddp_rec_model.module.decoder.ldm_z_channels *
                    3,  # type: ignore
                    self.diffusion_input_size,
                    self.diffusion_input_size),
                clip_denoised=args.clip_denoised,
                model_kwargs=model_kwargs,
                mixing_normal=True,  # !
            )
            th.cuda.empty_cache()

            self.render_video_given_triplane(
                triplane_sample,
                name_prefix=f'{self.step + self.resume_step}_{i}')
            # st()
            del triplane_sample
            th.cuda.empty_cache()