nsr/lsgm/train_util_diffusion_lsgm.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 calculate_adaptive_weight

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


class TrainLoop3DDiffusionLSGM(TrainLoop3DDiffusion,TrainLoop3DcvD_nvsD_canoD):
    def __init__(self, *, 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, 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, denoised_ae=denoised_ae, triplane_scaling_divider=triplane_scaling_divider, use_amp=use_amp, diffusion_input_size=diffusion_input_size, **kwargs)

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

        if step == 'diffusion_step_rec':
            self.forward_diffusion(batch, behaviour='diffusion_step_rec')
            _ = self.mp_trainer_rec.optimize(self.opt_rec) # TODO, update two groups of parameters
            took_step_ddpm = self.mp_trainer.optimize(self.opt) # TODO, update two groups of parameters

            if took_step_ddpm:
                self._update_ema()  # g_ema # TODO, ema only needs to track ddpm, remove ema tracking in rec

        elif step == 'd_step_rec':
            self.forward_D(batch, behaviour='rec')
            # _ = self.mp_trainer_cvD.optimize(self.opt_cvD)
            _ = self.mp_trainer_canonical_cvD.optimize(self.opt_cano_cvD)

        elif step == 'diffusion_step_nvs':
            self.forward_diffusion(batch, behaviour='diffusion_step_nvs')
            _ = self.mp_trainer_rec.optimize(self.opt_rec) # TODO, update two groups of parameters
            took_step_ddpm = self.mp_trainer.optimize(self.opt) # TODO, update two groups of parameters

            if took_step_ddpm:
                self._update_ema()  # g_ema

        elif step == 'd_step_nvs':
            self.forward_D(batch, behaviour='nvs')
            _ = self.mp_trainer_cvD.optimize(self.opt_cvD)

        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, cond = next(self.data)
            # if batch is None:
            # batch = next(self.data)
            # self.run_step(batch, 'g_step_rec')

            batch = next(self.data)
            self.run_step(batch, step='diffusion_step_rec')

            batch = next(self.data)
            self.run_step(batch, 'd_step_rec')

            # batch = next(self.data)
            # self.run_step(batch, 'g_step_nvs')

            batch = next(self.data)
            self.run_step(batch, step='diffusion_step_nvs')

            batch = next(self.data)
            self.run_step(batch, 'd_step_nvs')

            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_loop()
                    # self.eval_novelview_loop()
                    # let all processes sync up before starting with a new epoch of training
                th.cuda.empty_cache()
                dist_util.synchronize()

            if self.step % self.save_interval == 0:
                self.save(self.mp_trainer, self.mp_trainer.model_name)
                self.save(self.mp_trainer_rec, self.mp_trainer_rec.model_name)
                self.save(self.mp_trainer_cvD, 'cvD')
                self.save(self.mp_trainer_canonical_cvD, 'cano_cvD')

                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)
                    self.save(self.mp_trainer_rec, self.mp_trainer_rec.model_name)
                    self.save(self.mp_trainer_cvD, 'cvD')
                    self.save(self.mp_trainer_canonical_cvD, 'cano_cvD')

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

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

        self.ddp_cano_cvD.requires_grad_(False)
        self.ddp_nvs_cvD.requires_grad_(False)

        self.ddp_model.requires_grad_(True)
        self.ddp_rec_model.requires_grad_(True)

        # if behaviour != 'diff' and 'rec' in behaviour:
        # if behaviour != 'diff' and 'rec' in behaviour: # pure diffusion step
        #     self.ddp_rec_model.requires_grad_(True)
        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; 
        # else:

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

        # ! no 'sds' step now, both add sds grad back to ViT

        # assert behaviour != 'sds'
        # if behaviour == 'sds':
        # else:
        #     self.ddp_ddpm_model.requires_grad_(True)

        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

            vae_nelbo_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())
            d_weight = 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 not self.freeze_ae):

                # apply vae
                vae_out = self.ddp_rec_model(
                    img=micro['img_to_encoder'],
                    c=micro['c'],
                    behaviour='enc_dec_wo_triplane')  # pred: (B, 3, 64, 64)
                

                if behaviour == 'diffusion_step_rec':
                    target = micro
                    pred = self.ddp_rec_model(latent=vae_out,
                                              c=micro['c'],
                                              behaviour='triplane_dec')

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

                    last_layer = self.ddp_rec_model.module.decoder.triplane_decoder.decoder.net[  # type: ignore
                        -1].weight  # type: ignore

                    if 'image_sr' in pred:
                        vision_aided_loss = self.ddp_cano_cvD(
                            0.5 * pred['image_sr'] +
                            0.5 * th.nn.functional.interpolate(
                                pred['image_raw'],
                                size=pred['image_sr'].shape[2:],
                                mode='bilinear'),
                            for_G=True).mean()  # [B, 1] shape
                    else:
                        vision_aided_loss = self.ddp_cano_cvD(
                            pred['image_raw'], for_G=True
                        ).mean(
                        )   # [B, 1] shape

                    d_weight = calculate_adaptive_weight(
                        vae_nelbo_loss,
                        vision_aided_loss,
                        last_layer,
                        # disc_weight_max=1) * 1
                        disc_weight_max=1) * self.loss_class.opt.rec_cvD_lambda
                    # d_weight = self.loss_class.opt.rec_cvD_lambda # since decoder is fixed here. set to 0.001
                    
                    vision_aided_loss *= d_weight

                    # d_weight = self.loss_class.opt.rec_cvD_lambda
                    loss_dict.update({
                        'vision_aided_loss/G_rec':
                        vision_aided_loss,
                        'd_weight_G_rec':
                        d_weight,
                    })

                    log_rec3d_loss_dict(loss_dict)

                elif behaviour == 'diffusion_step_nvs':

                    novel_view_c = th.cat([micro['c'][1:], micro['c'][:1]])

                    pred = self.ddp_rec_model(latent=vae_out,
                                              c=novel_view_c,
                                              behaviour='triplane_dec')

                    if 'image_sr' in pred:
                        vision_aided_loss = self.ddp_nvs_cvD(
                            # pred_for_rec['image_sr'],
                            0.5 * pred['image_sr'] +
                            0.5 * th.nn.functional.interpolate(
                                pred['image_raw'],
                                size=pred['image_sr'].shape[2:],
                                mode='bilinear'),
                            for_G=True).mean()  # [B, 1] shape
                    else:
                        vision_aided_loss = self.ddp_nvs_cvD(
                            pred['image_raw'], for_G=True
                        ).mean(
                        )  # [B, 1] shape

                    d_weight = self.loss_class.opt.nvs_cvD_lambda
                    vision_aided_loss *= d_weight

                    log_rec3d_loss_dict({
                        'vision_aided_loss/G_nvs':
                        vision_aided_loss,
                    })

                    # ae_loss = th.tensor(0.0).to(dist_util.dev())

                # elif behaviour == 'diff':
                #     self.ddp_rec_model.requires_grad_(False)
                #     # assert self.ddp_rec_model.module.requires_grad == False, 'freeze ddpm_rec for pure diff step'
                else:
                    raise NotImplementedError(behaviour)
                #     assert behaviour == 'sds'

                # pred = None

                # if behaviour != 'sds': # also train diffusion
                # assert pred is not None

                # TODO, train diff and sds together, available?
                eps = vae_out[self.latent_name] 

                # if behaviour != 'sds':
                # micro_to_denoise.detach_()
                eps.requires_grad_(True) # single stage diffusion

                t, weights = self.schedule_sampler.sample(
                    eps.shape[0], dist_util.dev())
                noise = th.randn(size=vae_out.size(), device='cuda')  # note that this noise value is currently shared!

                model_kwargs = {}

                # ? 
                # or directly use SSD NeRF version?
                # get diffusion quantities for p (sgm prior) sampling scheme and reweighting for q (vae)

                # ! handle the sampling 

                # 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 = \
                    diffusion.iw_quantities(args.batch_size, args.time_eps, args.iw_sample_p, args.iw_subvp_like_vp_sde)
                eps_t_p = diffusion.sample_q(vae_out, noise, var_t_p, m_t_p)

                # 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 = \
                        diffusion.iw_quantities(args.batch_size, args.time_eps, args.iw_sample_q, args.iw_subvp_like_vp_sde)
                    eps_t_q = diffusion.sample_q(vae_out, noise, var_t_q, m_t_q)

                    eps_t_p = eps_t_p.detach().requires_grad_(True)
                    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)
                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

                # mixing normal trick
                # TODO, create a new partial training_losses function 
                mixing_component = diffusion.mixing_component(eps_t, var_t, t, enabled=dae.mixed_prediction) # TODO, which should I use?
                params = utils.get_mixed_prediction(dae.mixed_prediction, pred_params, dae.mixing_logit, mixing_component)

                # nelbo loss with kl balancing




                # ! remainign parts of cross entropy in likelihook training

                cross_entropy_per_var += diffusion.cross_entropy_const(args.time_eps)
                cross_entropy = th.sum(cross_entropy_per_var, dim=[1, 2, 3])
                cross_entropy += remaining_neg_log_p_total  # for remaining scales if there is any
                all_neg_log_p = vae.decompose_eps(cross_entropy_per_var)
                all_neg_log_p.extend(remaining_neg_log_p_per_ver)  # add the remaining neg_log_p
                kl_all_list, kl_vals_per_group, kl_diag_list = utils.kl_per_group_vada(all_log_q, all_neg_log_p)


                kl_coeff = 1.0

                # ! calculate p/q loss; 
                # ? no spectral regularizer here
                # ? try adding grid_clip and sn later on.
                q_loss = th.mean(nelbo_loss) 
                p_loss = th.mean(p_objective) 

                # backpropagate q_loss for vae and update vae params, if trained
                if args.train_vae:
                    grad_scalar.scale(q_loss).backward(retain_graph=utils.different_p_q_objectives(args.iw_sample_p, args.iw_sample_q))
                    utils.average_gradients(vae.parameters(), args.distributed)
                    if args.grad_clip_max_norm > 0.:  # apply gradient clipping
                        grad_scalar.unscale_(vae_optimizer)
                        th.nn.utils.clip_grad_norm_(vae.parameters(), max_norm=args.grad_clip_max_norm)
                    grad_scalar.step(vae_optimizer)

                # if we use different p and q objectives or are not training the vae, discard gradients and backpropagate p_loss
                if utils.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
                        dae_optimizer.zero_grad()

                    # compute gradients with unweighted loss
                    grad_scalar.scale(p_loss).backward()

                # update dae parameters
                utils.average_gradients(dae.parameters(), args.distributed)
                if args.grad_clip_max_norm > 0.:         # apply gradient clipping
                    grad_scalar.unscale_(dae_optimizer)
                    th.nn.utils.clip_grad_norm_(dae.parameters(), max_norm=args.grad_clip_max_norm)
                grad_scalar.step(dae_optimizer)


                # 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']:
                    l2_term_p, l2_term_q = th.chunk(l2_term, chunks=2, dim=0)
                    p_objective = th.sum(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.sum(obj_weight_t_p * l2_term, dim=[1, 2, 3])
                    # cross_entropy_per_var = obj_weight_t_q * l2_term

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

                # ! DDPM step
                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()

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

                denoise_loss = (losses["loss"] * weights).mean()

                x_t = losses.pop('x_t')
                model_output = losses.pop('model_output')
                diffusion_target = losses.pop('diffusion_target')
                alpha_bar = losses.pop('alpha_bar')

                log_loss_dict(self.diffusion, t,
                              {k: v * weights
                               for k, v in losses.items()})

                # if behaviour == 'sds':
                # ! calculate sds grad, and add to the grad of

                # if 'rec' in behaviour and self.loss_class.opt.sds_lamdba > 0:  # only enable sds along with rec step
                #     w = (
                #         1 - alpha_bar**2
                #     ) / self.triplane_scaling_divider * self.loss_class.opt.sds_lamdba  # https://github.com/ashawkey/stable-dreamfusion/issues/106
                #     sds_grad = denoise_loss.clone().detach(
                #     ) * w  # * https://pytorch.org/docs/stable/generated/th.Tensor.detach.html. detach() returned Tensor share the same storage with previous one. add clone() here.

                #     # ae_loss = AddGradient.apply(latent[self.latent_name], sds_grad) # add sds_grad during backward

                #     def sds_hook(grad_to_add):

                #         def modify_grad(grad):
                #             return grad + grad_to_add  # add the sds grad to the original grad for BP

                #         return modify_grad

                #     eps[self.latent_name].register_hook(
                #         sds_hook(sds_grad))  # merge sds grad with rec/nvs ae step

                loss = vae_nelbo_loss + denoise_loss + vision_aided_loss  # caluclate loss within AMP

            # ! cvD loss

            # exit AMP before backward
            self.mp_trainer_rec.backward(loss)
            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 and behaviour != 'diff':
                with th.no_grad():
                    # gt_vis = th.cat([batch['img'], batch['depth']], dim=-1)

                    # st()

                    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'], micro['img'],
                            gt_depth.repeat_interleave(3, dim=1)
                        ],
                        dim=-1)[0:1]  # TODO, fail to load depth. range [0, 1]

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

                    # if denoised_out is None:
                    # if not self.denoised_ae:
                    # denoised_out = denoised_fn()

                    if self.diffusion.model_mean_type == ModelMeanType.START_X:
                        pred_xstart = model_output
                    else:  # * used here
                        pred_xstart = self.diffusion._predict_xstart_from_eps(
                            x_t=x_t, t=t, eps=model_output)

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

                    # denoised_out = denoised_ae_pred

                    # if not self.denoised_ae:
                    #     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=self.render_latent_behaviour)
                    # else:
                    #     assert denoised_ae_pred is not None
                    #     denoised_ae_pred['image_raw'] = denoised_ae_pred[
                    #         'image_raw'][0:1]

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

                    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
                    # s

                    vis = th.cat([gt_vis, pred_vis], dim=-2)[0].permute(
                        1, 2, 0).cpu()  # ! pred in range[-1, 1]

                    # vis = th.cat([
                    #     self.pool_128(micro['img']), x_t[:, :3, ...],
                    #     denoised_out['pred_xstart'][:, :3, ...]
                    # ],
                    #              dim=-1)[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'
                    )

                    th.cuda.empty_cache()