src/model/model_wrapper.py

from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Optional
import torch
from einops import rearrange
from lightning.pytorch import LightningModule
from lightning.pytorch.utilities import rank_zero_only
from tabulate import tabulate
from torch import Tensor, nn
import torch.nn.functional as F
from ..dataset.data_module import get_data_shim
from ..dataset.types import BatchedExample
from ..evaluation.metrics import (
    compute_lpips,
    compute_psnr,
    compute_ssim,
)
from ..global_cfg import get_cfg
from ..loss import Loss
from ..misc.benchmarker import Benchmarker
from ..misc.image_io import prep_image, save_image
from ..misc.step_tracker import StepTracker
from ..misc.utils import (
    get_rank,
    inverse_normalize,
    vis_depth_map,
)
from ..visualization.annotation import add_label
from ..visualization.layout import add_border, hcat, vcat
from .decoder.decoder import DepthRenderingMode
from lightning.pytorch.loggers.wandb import WandbLogger


@dataclass
class OptimizerCfg:
    lr: float
    warm_up_steps: int
    backbone_lr_multiplier: float


@dataclass
class TestCfg:
    output_path: Path
    align_pose: bool
    pose_align_steps: int
    rot_opt_lr: float
    trans_opt_lr: float
    compute_scores: bool
    save_image: bool
    save_video: bool
    save_compare: bool
    generate_video: bool
    mode: Literal["inference", "evaluation"]
    image_folder: str


@dataclass
class TrainCfg:
    output_path: Path
    depth_mode: DepthRenderingMode | None
    extended_visualization: bool
    print_log_every_n_steps: int
    distiller: str
    distill_max_steps: int
    pose_loss_alpha: float = 1.0
    pose_loss_delta: float = 1.0
    cxt_depth_weight: float = 0.01
    weight_pose: float = 1.0
    weight_depth: float = 1.0
    weight_normal: float = 1.0
    render_ba: bool = False
    render_ba_after_step: int = 0


class ModelWrapper(LightningModule):
    logger: Optional[WandbLogger]
    model: nn.Module
    losses: nn.ModuleList
    optimizer_cfg: OptimizerCfg
    test_cfg: TestCfg
    train_cfg: TrainCfg
    step_tracker: StepTracker | None

    def __init__(
        self,
        optimizer_cfg: OptimizerCfg,
        test_cfg: TestCfg,
        train_cfg: TrainCfg,
        model: nn.Module,
        losses: list[Loss],
        step_tracker: StepTracker | None,
    ) -> None:
        super().__init__()
        self.optimizer_cfg = optimizer_cfg
        self.test_cfg = test_cfg
        self.train_cfg = train_cfg
        self.step_tracker = step_tracker

        # Set up the model.
        self.encoder_visualizer = None
        self.model = model
        self.data_shim = get_data_shim(self.model.encoder)
        self.losses = nn.ModuleList(losses)

        # This is used for testing.
        self.benchmarker = Benchmarker()

    @staticmethod
    def flatten_indices(indices) -> list[int]:
        if isinstance(indices, Tensor):
            return [int(idx) for idx in indices.detach().cpu().reshape(-1).tolist()]
        if isinstance(indices, (list, tuple)):
            flattened = []
            for item in indices:
                flattened.extend(ModelWrapper.flatten_indices(item))
            return flattened
        return [int(indices)]

    def on_train_epoch_start(self) -> None:
        # our custom dataset and sampler has to have epoch set by calling set_epoch
        print(f"Train epoch start on rank {self.trainer.global_rank}")
        if hasattr(self.trainer.datamodule.train_loader.dataset, "set_epoch"):
            self.trainer.datamodule.train_loader.dataset.set_epoch(self.current_epoch)
        if hasattr(self.trainer.datamodule.train_loader.sampler, "set_epoch"):
            self.trainer.datamodule.train_loader.sampler.set_epoch(self.current_epoch)

    def on_validation_epoch_start(self) -> None:
        print(f"Validation epoch start on rank {self.trainer.global_rank}")
        # our custom dataset and sampler has to have epoch set by calling set_epoch
        if hasattr(self.trainer.datamodule.val_loader.dataset, "set_epoch"):
            self.trainer.datamodule.val_loader.dataset.set_epoch(self.current_epoch)
        if hasattr(self.trainer.datamodule.val_loader.sampler, "set_epoch"):
            self.trainer.datamodule.val_loader.sampler.set_epoch(self.current_epoch)

    def training_step(self, batch, batch_idx):
        if isinstance(batch, list):
            batch_combined = None
            for batch_per_dl in batch:
                if batch_combined is None:
                    batch_combined = batch_per_dl
                else:
                    for k in batch_combined.keys():
                        if isinstance(batch_combined[k], list):
                            batch_combined[k] += batch_per_dl[k]
                        elif isinstance(batch_combined[k], dict):
                            for kk in batch_combined[k].keys():
                                batch_combined[k][kk] = torch.cat(
                                    [batch_combined[k][kk], batch_per_dl[k][kk]], dim=0
                                )
                        else:
                            raise NotImplementedError
            batch = batch_combined

        batch: BatchedExample = self.data_shim(batch)
        context_image = (batch["context"]["image"] + 1) / 2

        # Run the model.
        encoder_output, output = self.model(context_image, self.global_step)

        gaussians, pred_pose_enc_list, depth_dict = (
            encoder_output.gaussians,
            encoder_output.pred_pose_enc_list,
            encoder_output.depth_dict,
        )
        distill_infos = encoder_output.distill_infos

        target_gt = (batch["context"]["image"] + 1) / 2
        num_context_views = target_gt.shape[1]

        using_index = torch.arange(num_context_views, device=gaussians.means.device)
        batch["using_index"] = using_index

        psnr_probabilistic = compute_psnr(
            rearrange(target_gt, "b v c h w -> (b v) c h w"),
            rearrange(output.color, "b v c h w -> (b v) c h w"),
        )
        self.log("train/psnr_probabilistic", psnr_probabilistic.mean().item())

        total_loss = 0

        with torch.amp.autocast("cuda", enabled=False):
            depth_loss_idx = list(get_cfg()["loss"].keys()).index("depth")
            depth_loss_fn = self.losses[depth_loss_idx].ctx_depth_loss
            loss_depth_ctx = depth_loss_fn(
                depth_dict["depth"],
                batch,
                cxt_depth_weight=self.train_cfg.cxt_depth_weight,
            )
            self.log("loss/loss_depth_ctx", loss_depth_ctx.item())
            total_loss = total_loss + loss_depth_ctx

            for loss_fn in self.losses:
                if loss_fn.name == "depth":
                    break
                loss = loss_fn.forward(
                    output, batch, gaussians, depth_dict, self.global_step
                )
                self.log(f"loss/{loss_fn.name}", loss.item())
                total_loss = total_loss + loss

            loss_ca1 = F.mse_loss(
                pred_pose_enc_list, distill_infos["pred_pose_enc_list"][:, 0:1, -1]
            )
            loss_ca = 10 * loss_ca1
            self.log("loss/loss_ca", loss_ca.item())
            total_loss = total_loss + loss_ca

        self.log("loss/total", total_loss.item())
        self.log("info/global_step", self.global_step)

        if self.step_tracker is not None:
            self.step_tracker.set_step(self.global_step)

        del batch

        return total_loss

    def on_after_backward(self):
        for name, p in self.named_parameters():
            if p.grad is None:
                continue
            grad = p.grad.detach()
            if torch.isnan(grad).any() or torch.isinf(grad).any():
                print(f"[NaN-Guard]")
                p.grad = torch.zeros_like(p.grad)
                continue

    @rank_zero_only
    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        batch: BatchedExample = self.data_shim(batch)
        total_batches = len(self.trainer.datamodule.val_loader)
        print(f"Rank {self.global_rank}, batch {batch_idx+1}/{total_batches}")
        print(
            f"validation step {self.global_step}; "
            f"scene = {batch['scene']}; "
            f"context = {batch['context']['index'].tolist()}"
        )

        # Render Gaussians.
        b, v, _, h, w = batch["context"]["image"].shape
        assert b == 1

        encoder_output, output = self.model(
            (batch["context"]["image"] + 1) / 2,
            self.global_step,
        )

        # Compute validation metrics.
        rgb_pred = output.color[0].float()
        rgb_gt = (batch["context"]["image"][0].float() + 1) / 2
        psnr = compute_psnr(rgb_gt, rgb_pred).mean()
        self.log(f"val/psnr", psnr)
        lpips = compute_lpips(rgb_gt, rgb_pred).mean()
        self.log(f"val/lpips", lpips)
        ssim = compute_ssim(rgb_gt, rgb_pred).mean()
        self.log(f"val/ssim", ssim)

        # Construct comparison image.
        context_img = inverse_normalize(batch["context"]["image"][0])
        context = []
        for i in range(context_img.shape[0]):
            context.append(context_img[i])

        depth_dict = encoder_output.depth_dict
        model_depth_pred = depth_dict["depth"].squeeze(-1)[0]
        model_depth_pred = vis_depth_map(model_depth_pred)

        depth_pred = vis_depth_map(output.depth[0])

        comparison = hcat(
            add_label(vcat(*context), "Context"),
            add_label(vcat(*rgb_gt), "Target (Ground Truth)"),
            add_label(vcat(*rgb_pred), "Rendered Target"),
            add_label(vcat(*depth_pred), "Rendered Depth"),
            add_label(vcat(*model_depth_pred), "GS Depth"),
        )

        comparison = torch.nn.functional.interpolate(
            comparison.unsqueeze(0),
            scale_factor=0.5,
            mode="bicubic",
            align_corners=False,
        ).squeeze(0)

        self.logger.log_image(
            "comparison",
            [prep_image(add_border(comparison))],
            step=self.global_step,
            caption=batch["scene"],
        )

        if self.encoder_visualizer is not None:
            for k, image in self.encoder_visualizer.visualize(
                batch["context"], self.global_step
            ).items():
                self.logger.log_image(k, [prep_image(image)], step=self.global_step)

    def test_step(self, batch, batch_idx):
        batch: BatchedExample = self.data_shim(batch)
        b, v, _, h, w = batch["target"]["image"].shape

        assert b == 1
        if batch_idx % 100 == 0:
            print(
                f"Rank {get_rank()} test step {batch_idx:0>6}; "
                f"scene = {batch['scene']}; "
                f"context = {self.flatten_indices(batch['context']['index'])}; "
                f"target = {self.flatten_indices(batch['target']['index'])}"
            )

        # Render Gaussians.
        with torch.no_grad():
            with self.benchmarker.time("encoder"):
                (
                    gaussians,
                    pred_all_extrinsic,
                    pred_context_pose,
                ) = self.model.encoder.inference(
                    (batch["context"]["image"] + 1) / 2,
                    (batch["target"]["image"] + 1) / 2,
                    global_step=self.global_step,
                )
        num_context_view = batch["context"]["image"].shape[1]

        pred_all_context_extrinsic, pred_all_target_extrinsic = (
            pred_all_extrinsic[:, :num_context_view],
            pred_all_extrinsic[:, num_context_view:],
        )
        scale_factor = (
            pred_context_pose["extrinsic"][:, :, :3, 3].mean()
            / pred_all_context_extrinsic[:, :, :3, 3].mean()
        )
        pred_all_target_extrinsic[..., :3, 3] = (
            pred_all_target_extrinsic[..., :3, 3] * scale_factor
        )
        pred_all_context_extrinsic[..., :3, 3] = (
            pred_all_context_extrinsic[..., :3, 3] * scale_factor
        )

        with self.benchmarker.time("decoder", num_calls=v):
            output = self.model.decoder.forward(
                gaussians,
                pred_all_target_extrinsic,
                pred_context_pose["intrinsic"][:, 0:1, :, :]
                .repeat(1, pred_all_target_extrinsic.shape[1], 1, 1)
                .float(),
                torch.ones(1, v, device="cuda") * 0.01,
                torch.ones(1, v, device="cuda") * 100,
                (h, w),
            )

        psnr = None
        with torch.no_grad():
            if self.test_cfg.compute_scores:
                rgb_pred = output.color[0]
                rgb_gt = batch["target"]["image"][0]
                psnr = compute_psnr(rgb_gt, rgb_pred).mean().item()
                all_metrics = {
                    f"lpips_ours": compute_lpips(rgb_gt, rgb_pred).mean().item(),
                    f"ssim_ours": compute_ssim(rgb_gt, rgb_pred).mean().item(),
                    f"psnr_ours": psnr,
                }
                methods = ["ours"]
                self.log_dict(all_metrics, prog_bar=True, sync_dist=True, on_epoch=True)
                self.print_preview_metrics(all_metrics, methods)

        # Save images.
        (scene,) = batch["scene"]
        name = get_cfg()["wandb"]["name"]
        path = self.test_cfg.output_path / name
        scene_dir = f"{psnr:.4f}_{scene}" if psnr is not None else scene
        target_indices = self.flatten_indices(batch["target"]["index"])

        for i, idx in enumerate(target_indices):
            single_color = output.color[0][i]
            res_path = path / scene_dir / "color" / f"{int(idx):0>6}.png"
            save_image(single_color, res_path)

    def on_test_end(self) -> None:
        self.benchmarker.summarize()

    def print_preview_metrics(
        self,
        metrics: dict[str, float | Tensor],
        methods: list[str] | None = None,
        overlap_tag: str | None = None,
    ) -> None:
        if getattr(self, "running_metrics", None) is None:
            self.running_metrics = metrics
            self.running_metric_steps = 1
        else:
            s = self.running_metric_steps
            self.running_metrics = {
                k: ((s * v) + metrics[k]) / (s + 1)
                for k, v in self.running_metrics.items()
            }
            self.running_metric_steps += 1

        if overlap_tag is not None:
            if getattr(self, "running_metrics_sub", None) is None:
                self.running_metrics_sub = {overlap_tag: metrics}
                self.running_metric_steps_sub = {overlap_tag: 1}
            elif overlap_tag not in self.running_metrics_sub:
                self.running_metrics_sub[overlap_tag] = metrics
                self.running_metric_steps_sub[overlap_tag] = 1
            else:
                s = self.running_metric_steps_sub[overlap_tag]
                self.running_metrics_sub[overlap_tag] = {
                    k: ((s * v) + metrics[k]) / (s + 1)
                    for k, v in self.running_metrics_sub[overlap_tag].items()
                }
                self.running_metric_steps_sub[overlap_tag] += 1

        metric_list = ["psnr", "lpips", "ssim"]

        def print_metrics(runing_metric, methods=None):
            table = []
            if methods is None:
                methods = ["ours"]

            for method in methods:
                row = [
                    f"{runing_metric[f'{metric}_{method}']:.3f}"
                    for metric in metric_list
                ]
                table.append((method, *row))

            headers = ["Method"] + metric_list
            table = tabulate(table, headers)
            print(table)

        print("All Pairs:")
        print_metrics(self.running_metrics, methods)

    def configure_optimizers(self):
        new_params, new_param_names = [], []
        for name, param in self.named_parameters():
            if not param.requires_grad:
                continue
            new_params.append(param)
            new_param_names.append(name)

        param_dicts = [
            {
                "params": new_params,
                "lr": self.optimizer_cfg.lr,
            }
        ]
        optimizer = torch.optim.AdamW(
            param_dicts, lr=self.optimizer_cfg.lr, weight_decay=0.1, betas=(0.9, 0.95)
        )
        max_steps = get_cfg()["trainer"]["max_steps"]
        warm_up_steps = self.optimizer_cfg.warm_up_steps
        if warm_up_steps > 0:
            warm_up = torch.optim.lr_scheduler.LinearLR(
                optimizer,
                start_factor=1.0 / warm_up_steps,
                end_factor=1.0,
                total_iters=warm_up_steps,
            )
            lr_scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizer,
                T_max=max_steps - warm_up_steps,
                eta_min=self.optimizer_cfg.lr * 0.1,
            )
            lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
                optimizer,
                schedulers=[warm_up, lr_scheduler_cosine],
                milestones=[warm_up_steps],
            )

        else:
            lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizer, T_max=max_steps, eta_min=self.optimizer_cfg.lr * 0.1
            )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": lr_scheduler,
                "interval": "step",
                "frequency": 1,
            },
        }