nsr/cvD/nvsD_canoD_multiview.py

import functools
import json
import os
from pathlib import Path
from pdb import set_trace as st
import torchvision
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 tqdm import tqdm

from guided_diffusion.fp16_util import MixedPrecisionTrainer
from guided_diffusion import dist_util, logger
from guided_diffusion.train_util import (calc_average_loss,
                                         log_rec3d_loss_dict,
                                         find_resume_checkpoint)

from torch.optim import AdamW

from .nvsD_canoD import TrainLoop3DcvD_nvsD_canoD


class TrainLoop3DcvD_nvsD_canoD_multiview(TrainLoop3DcvD_nvsD_canoD):

    def __init__(self,
                 *,
                 rec_model,
                 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,
                 weight_decay=0,
                 lr_anneal_steps=0,
                 iterations=10001,
                 load_submodule_name='',
                 ignore_resume_opt=False,
                 use_amp=False,
                 **kwargs):
        super().__init__(rec_model=rec_model,
                         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,
                         weight_decay=weight_decay,
                         lr_anneal_steps=lr_anneal_steps,
                         iterations=iterations,
                         load_submodule_name=load_submodule_name,
                         ignore_resume_opt=ignore_resume_opt,
                         use_amp=use_amp,
                         **kwargs)
        assert not self.mp_trainer_rec.use_amp, 'amp may lead to grad nan?'

    def forward_G_rec(self, batch):  # update G

        self.mp_trainer_rec.zero_grad()
        self.rec_model.requires_grad_(True)

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

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

        target_cano = {}

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

            for k, v in micro.items():
                if k[:2] == 'nv':
                    orig_key = k.replace('nv_', '')
                    # target_nvs[orig_key] = v
                    target_cano[orig_key] = micro[orig_key]

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

            with th.autocast(device_type='cuda',
                             dtype=th.float16,
                             enabled=self.mp_trainer_rec.use_amp):

                pred = self.rec_model(
                    img=micro['img_to_encoder'], c=micro['c']
                )  # render novel view for first half of the batch for D loss

                target_for_rec = micro
                cano_pred = pred

                with self.rec_model.no_sync():  # type: ignore
                    loss, loss_dict, fg_mask = self.loss_class(
                        cano_pred,
                        target_for_rec,
                        test_mode=False,
                        step=self.step + self.resume_step,
                        return_fg_mask=True)

                if 'image_sr' in cano_pred:
                    raise NotImplementedError()
                    # concat both resolution
                    vision_aided_loss = self.ddp_cano_cvD(
                        th.cat([
                            th.nn.functional.interpolate(
                                cano_pred['image_raw'],
                                size=cano_pred['image_sr'].shape[2:],
                                mode='bilinear',
                                align_corners=False,
                                antialias=True),
                            cano_pred['image_sr'],
                        ],
                               dim=1),  # 6 channel input
                        for_G=True).mean()  # [B, 1] shape

                else:
                    vision_aided_loss = self.ddp_cano_cvD(
                        cano_pred['image_raw'],
                        for_G=True).mean()  # [B, 1] shape

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

                d_weight = th.tensor(self.loss_class.opt.rec_cvD_lambda).to(
                    dist_util.dev())
                # d_weight = calculate_adaptive_weight(
                #     loss,
                #     vision_aided_loss,
                #     last_layer,
                #     disc_weight_max=0.1) * self.loss_class.opt.rec_cvD_lambda
                loss += vision_aided_loss * d_weight

                loss_dict.update({
                    'vision_aided_loss/G_rec':
                    (vision_aided_loss * d_weight).detach(),
                    'd_weight':
                    d_weight
                })

                log_rec3d_loss_dict(loss_dict)

            self.mp_trainer_rec.backward(
                loss)  # no nvs cvD loss, following VQ3D

            # DDP some parameters no grad:
            # for name, p in self.ddp_model.named_parameters():
            #     if p.grad is None:
            #         print(f"(in rec)found rec unused param: {name}")

            # ! move to other places, add tensorboard

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

            #         pred_vis = th.cat(
            #             [pred_img,
            #              pred_depth.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}_rec.jpg'
            #         )
            #         print(
            #             'log vis to: ',
            #             f'{logger.get_dir()}/{self.step+self.resume_step}_rec.jpg'
            #         )

            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)

                    def norm_depth(pred_depth):  # to [-1,1]
                        # pred_depth = pred['image_depth']
                        pred_depth = (pred_depth - pred_depth.min()) / (
                            pred_depth.max() - pred_depth.min())
                        return -(pred_depth * 2 - 1)

                    pred_img = pred['image_raw'].clip(-1, 1)
                    gt_img = micro['img']

                    # infer novel view also
                    pred_nv_img = self.rec_model(
                        img=micro['img_to_encoder'],
                        c=self.novel_view_poses)  # pred: (B, 3, 64, 64)

                    # if 'depth' in micro:
                    gt_depth = micro['depth']
                    if gt_depth.ndim == 3:
                        gt_depth = gt_depth.unsqueeze(1)
                    gt_depth = norm_depth(gt_depth)
                    # gt_depth = (gt_depth - gt_depth.min()) / (gt_depth.max() -
                    #                                           gt_depth.min())
                    # if True:
                    if 'image_depth' in pred:
                        # pred_depth = pred['image_depth']
                        # pred_depth = (pred_depth - pred_depth.min()) / (
                        #     pred_depth.max() - pred_depth.min())
                        pred_depth = norm_depth(pred['image_depth'])
                        pred_nv_depth = norm_depth(pred_nv_img['image_depth'])
                    else:
                        pred_depth = th.zeros_like(gt_depth)
                        pred_nv_depth = th.zeros_like(gt_depth)

                    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)

                    if gt_img.shape[-1] == 64:
                        gt_depth = self.pool_64(gt_depth)
                    elif gt_img.shape[-1] == 128:
                        gt_depth = self.pool_128(gt_depth)
                    # else:
                    # gt_depth = self.pool_64(gt_depth)

                    # st()
                    pred_vis = th.cat(
                        [pred_img,
                         pred_depth.repeat_interleave(3, dim=1)],
                        dim=-1)  # B, 3, H, W

                    pred_vis_nv = th.cat([
                        pred_nv_img['image_raw'].clip(-1, 1),
                        pred_nv_depth.repeat_interleave(3, dim=1)
                    ],
                                         dim=-1)  # B, 3, H, W
                    pred_vis = th.cat([pred_vis, pred_vis_nv],
                                      dim=-2)  # cat in H dim

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

                    # vis = th.cat([gt_vis, pred_vis], dim=-2)[0].permute(
                    vis = th.cat([gt_vis, pred_vis], dim=-2)
                    # .permute(
                    #     0, 2, 3, 1).cpu()
                    vis_tensor = torchvision.utils.make_grid(
                        vis, nrow=vis.shape[-1] // 64)  # HWC
                    torchvision.utils.save_image(
                        vis_tensor,
                        f'{logger.get_dir()}/{self.step+self.resume_step}.jpg',
                        normalize=True,
                        value_range=(-1, 1))
                    # 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}.jpg')

                    logger.log(
                        'log vis to: ',
                        f'{logger.get_dir()}/{self.step+self.resume_step}.jpg')

    def forward_G_nvs(self, batch):  # update G

        self.mp_trainer_rec.zero_grad()
        self.rec_model.requires_grad_(True)

        self.ddp_cano_cvD.requires_grad_(False)
        self.ddp_nvs_cvD.requires_grad_(False)  # only use novel view D

        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()).contiguous()
                for k, v in batch.items()
            }
            target_nvs = {}

            for k, v in micro.items():
                if k[:2] == 'nv':
                    orig_key = k.replace('nv_', '')
                    target_nvs[orig_key] = v
                    # target_cano[orig_key] = micro[orig_key]

            with th.autocast(device_type='cuda',
                             dtype=th.float16,
                             enabled=self.mp_trainer_rec.use_amp):

                nvs_pred = self.rec_model(
                    img=micro['img_to_encoder'],
                    c=micro['nv_c'],
                )  # predict novel view here
                # c=th.cat([
                #     micro['c'][1:],
                #     micro['c'][:1],
                # ]))  # ! render novel views only for D loss

                # add cvD supervision

                if 'image_sr' in nvs_pred:
                    raise NotImplementedError()
                    # concat sr and raw results
                    vision_aided_loss = self.ddp_nvs_cvD(
                        # pred_nv['image_sr'],
                        # 0.5 * pred_nv['image_sr'] + 0.5 * th.nn.functional.interpolate(pred_nv['image_raw'], size=pred_nv['image_sr'].shape[2:], mode='bilinear'),
                        th.cat([
                            th.nn.functional.interpolate(
                                nvs_pred['image_raw'],
                                size=nvs_pred['image_sr'].shape[2:],
                                mode='bilinear',
                                align_corners=False,
                                antialias=True),
                            nvs_pred['image_sr'],
                        ],
                               dim=1),
                        for_G=True).mean()  # ! for debugging

                    # supervise sr only
                    # vision_aided_loss = self.ddp_nvs_cvD(
                    #     # pred_nv['image_sr'],
                    #     # 0.5 * pred_nv['image_sr'] + 0.5 * th.nn.functional.interpolate(pred_nv['image_raw'], size=pred_nv['image_sr'].shape[2:], mode='bilinear'),
                    #     th.cat([nvs_pred['image_sr'],
                    #     th.nn.functional.interpolate(nvs_pred['image_raw'], size=nvs_pred['image_sr'].shape[2:], mode='bilinear',
                    #                         align_corners=False,
                    #                         antialias=True),]),
                    #     for_G=True).mean()  # ! for debugging

                    # pred_nv['image_raw'], for_G=True).mean()  # [B, 1] shape
                else:
                    vision_aided_loss = self.ddp_nvs_cvD(
                        nvs_pred['image_raw'],
                        for_G=True).mean()  # [B, 1] shape

                # ! add nv reconstruction loss
                with self.rec_model.no_sync():  # type: ignore
                    loss, loss_dict, fg_mask = self.loss_class(
                        nvs_pred,
                        target_nvs,
                        step=self.step + self.resume_step,
                        test_mode=False,
                        return_fg_mask=True,
                        conf_sigma_l1=None,
                        conf_sigma_percl=None)
                    log_rec3d_loss_dict(loss_dict)

                loss += vision_aided_loss * self.loss_class.opt.nvs_cvD_lambda

                log_rec3d_loss_dict({
                    'vision_aided_loss/G_nvs':
                    vision_aided_loss * self.loss_class.opt.nvs_cvD_lambda,
                    **{f'{k}_nv': v for k, v in loss_dict.items()}
                })

            self.mp_trainer_rec.backward(loss)

            # ! move to other places, add tensorboard

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

                    def norm_depth(pred_depth):  # to [-1,1]
                        # pred_depth = pred['image_depth']
                        pred_depth = (pred_depth - pred_depth.min()) / (
                            pred_depth.max() - pred_depth.min())
                        return -(pred_depth * 2 - 1)

                    gt_depth = micro['depth']
                    if gt_depth.ndim == 3:
                        gt_depth = gt_depth.unsqueeze(1)
                    gt_depth = norm_depth(gt_depth)

                    # if True:
                    # pred_depth = nvs_pred['image_depth']
                    # pred_depth = (pred_depth - pred_depth.min()) / (
                    #     pred_depth.max() - pred_depth.min())
                    pred_depth = norm_depth(nvs_pred['image_depth'])
                    pred_img = nvs_pred['image_raw']
                    gt_img = micro['img']

                    if 'image_sr' in nvs_pred:

                        if nvs_pred['image_sr'].shape[-1] == 512:
                            pred_img = th.cat([
                                self.pool_512(pred_img), nvs_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 nvs_pred['image_sr'].shape[-1] == 256:
                            pred_img = th.cat([
                                self.pool_256(pred_img), nvs_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), nvs_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)

                    if gt_img.shape[-1] == 64:
                        gt_depth = self.pool_64(gt_depth)
                    elif gt_img.shape[-1] == 128:
                        gt_depth = self.pool_128(gt_depth)

                    # else:
                    #     gt_img = self.pool_64(gt_img)
                    #     gt_depth = self.pool_64(gt_depth)

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

                    pred_vis = th.cat(
                        [pred_img,
                         pred_depth.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 = th.cat([gt_vis, pred_vis], dim=-2)

                    vis = torchvision.utils.make_grid(
                        vis,
                        normalize=True,
                        scale_each=True,
                        value_range=(-1, 1)).cpu().permute(1, 2, 0)  # H W 3
                    vis = vis.numpy() * 255
                    vis = vis.clip(0, 255).astype(np.uint8)

                    # print(vis.shape)

                    Image.fromarray(vis).save(
                        f'{logger.get_dir()}/{self.step+self.resume_step}_nvs.jpg'
                    )
                    print(
                        'log vis to: ',
                        f'{logger.get_dir()}/{self.step+self.resume_step}_nvs.jpg'
                    )