model.py

import logging
import math
import os
import sys
from time import perf_counter

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from fused_ssim import fused_ssim
from lpips import LPIPS
from pytorch_msssim import MS_SSIM
from torchvision.transforms.functional import gaussian_blur

from gsplat import (
    project_gaussians_2d_scale_rot,
    rasterize_gaussians_no_tiles,
    rasterize_gaussians_sum,
)
from utils.flip import LDRFLIPLoss
from utils.image_utils import (
    compute_image_gradients,
    get_grid,
    get_psnr,
    load_images,
    save_image,
    separate_image_channels,
    visualize_added_gaussians,
    visualize_gaussians,
)
from utils.misc_utils import clean_dir, get_latest_ckpt_step, save_cfg, set_random_seed
from utils.quantization_utils import ste_quantize
from utils.saliency_utils import get_smap


class GaussianSplatting2D(nn.Module):
    def __init__(self, args):
        super(GaussianSplatting2D, self).__init__()
        self.evaluate = args.eval
        set_random_seed(seed=args.seed)
        # Ensure we're using the correct CUDA device
        if torch.cuda.is_available():
            torch.cuda.set_device(0)  # Force device 0
            self.device = torch.device("cuda:0")
        else:
            self.device = torch.device("cpu")
        self.dtype = torch.float32
        self._init_logging(args)
        self._init_target(args)
        self._init_bit_precision(args)
        self._init_gaussians(args)
        self._init_loss(args)
        self._init_optimization(args)
        # Initialization
        if self.evaluate:
            self.ckpt_file = args.ckpt_file
            self._load_model()
        else:
            self._init_pos_scale_feat(args)

    def _init_logging(self, args):
        self.log_dir = args.log_dir
        self.log_level = args.log_level
        self.ckpt_dir = os.path.join(self.log_dir, "checkpoints")
        self.train_dir = os.path.join(self.log_dir, "train")
        self.eval_dir = os.path.join(self.log_dir, "eval")
        self.vis_gaussians = args.vis_gaussians
        self.save_image_steps = args.save_image_steps
        self.save_ckpt_steps = args.save_ckpt_steps
        self.eval_steps = args.eval_steps
        if not self.evaluate:
            clean_dir(path=self.log_dir)
            os.makedirs(self.log_dir, exist_ok=False)
            os.makedirs(self.ckpt_dir, exist_ok=False)
            os.makedirs(self.train_dir, exist_ok=False)
        else:
            os.makedirs(self.eval_dir, exist_ok=True)
        self._gen_logger(args)
        if not self.evaluate:
            save_cfg(path=f"{self.log_dir}/cfg_train.yaml", args=args)

    def _gen_logger(self, args):
        log_fname = "log_train"
        if self.evaluate:
            log_fname = "log_eval"
        log_level = getattr(logging, self.log_level, logging.INFO)
        logging.basicConfig(level=log_level)
        self.worklog = logging.getLogger("Image-GS Logger")
        self.worklog.propagate = False
        datefmt = "%Y/%m/%d %H:%M:%S"
        fileHandler = logging.FileHandler(
            f"{self.log_dir}/{log_fname}.txt", mode="a", encoding="utf8"
        )
        fileHandler.setFormatter(
            logging.Formatter(fmt="[{asctime}] {message}", datefmt=datefmt, style="{")
        )
        consoleHandler = logging.StreamHandler(sys.stdout)
        consoleHandler.setFormatter(
            logging.Formatter(
                fmt="\x1b[32m[{asctime}] \x1b[0m{message}", datefmt=datefmt, style="{"
            )
        )
        self.worklog.handlers = [fileHandler, consoleHandler]
        action = "rendering" if self.evaluate else "optimizing"
        self.worklog.info(
            f"Start {action} {args.num_gaussians:d} Gaussians for '{args.input_path}'"
        )
        self.worklog.info("***********************************************")

    def _init_target(self, args):
        self.gamma = args.gamma
        self.downsample = args.downsample
        if self.downsample:
            self.downsample_ratio = float(args.downsample_ratio)
        self.block_h, self.block_w = (
            16,
            16,
        )  # Warning: Must match hardcoded value in CUDA kernel, modify with caution
        self._load_target_images(path=os.path.join(args.data_root, args.input_path))
        if self.downsample:
            self.gt_images_upsampled = self.gt_images
            self.img_h_upsampled, self.img_w_upsampled = self.img_h, self.img_w
            self.tile_bounds_upsampled = self.tile_bounds
            self._load_target_images(
                path=os.path.join(args.data_root, args.input_path),
                downsample_ratio=self.downsample_ratio,
            )
            if not self.evaluate:
                path = f"{self.log_dir}/gt_upsample-{self.downsample_ratio:.1f}_res-{self.img_h_upsampled:d}x{self.img_w_upsampled:d}"
                self._separate_and_save_images(
                    images=self.gt_images_upsampled,
                    channels=self.input_channels,
                    path=path,
                )
        self.num_pixels = self.img_h * self.img_w
        if not self.evaluate:
            path = f"{self.log_dir}/gt_res-{self.img_h:d}x{self.img_w:d}"
            self._separate_and_save_images(
                images=self.gt_images, channels=self.input_channels, path=path
            )

    def _load_target_images(self, path, downsample_ratio=None):
        self.gt_images, self.input_channels, self.image_fnames = load_images(
            load_path=path, downsample_ratio=downsample_ratio, gamma=self.gamma
        )
        self.gt_images = torch.from_numpy(self.gt_images).to(
            dtype=self.dtype, device=self.device
        )
        self.img_h, self.img_w = self.gt_images.shape[1:]
        self.tile_bounds = (
            (self.img_w + self.block_w - 1) // self.block_w,
            (self.img_h + self.block_h - 1) // self.block_h,
            1,
        )

    def _separate_and_save_images(self, images, channels, path):
        images_sep = separate_image_channels(images=images, input_channels=channels)
        for idx, image in enumerate(images_sep, 1):
            suffix = "" if len(images_sep) == 1 else f"_{idx:d}"
            save_image(image, f"{path}{suffix}.png", gamma=self.gamma)

    def _init_bit_precision(self, args):
        self.quantize = args.quantize
        self.pos_bits = args.pos_bits
        self.scale_bits = args.scale_bits
        self.rot_bits = args.rot_bits
        self.feat_bits = args.feat_bits

    def _init_gaussians(self, args):
        self.num_gaussians = args.num_gaussians
        self.total_num_gaussians = args.num_gaussians
        self.disable_prog_optim = args.disable_prog_optim
        if not self.disable_prog_optim and not self.evaluate:
            self.initial_ratio = args.initial_ratio
            self.add_times = args.add_times
            self.add_steps = args.add_steps
            self.num_gaussians = math.ceil(
                self.initial_ratio * self.total_num_gaussians
            )
            self.max_add_num = math.ceil(
                float(self.total_num_gaussians - self.num_gaussians) / self.add_times
            )
            min_steps = self.add_steps * self.add_times + args.post_min_steps
            if args.max_steps < min_steps:
                self.worklog.info(
                    f"Max steps ({args.max_steps:d}) is too small for progressive optimization. Resetting to {min_steps:d}"
                )
                args.max_steps = min_steps
        self.topk = (
            args.topk
        )  # Warning: Must match hardcoded value in CUDA kernel, modify with caution
        self.eps = (
            1e-7 if args.disable_tiles else 1e-4
        )  # Warning: Must match hardcoded value in CUDA kernel, modify with caution
        self.init_scale = args.init_scale
        self.disable_topk_norm = args.disable_topk_norm
        self.disable_inverse_scale = args.disable_inverse_scale
        self.disable_color_init = args.disable_color_init
        self.xy = nn.Parameter(
            torch.rand(self.num_gaussians, 2, dtype=self.dtype, device=self.device),
            requires_grad=True,
        )
        self.scale = nn.Parameter(
            torch.ones(self.num_gaussians, 2, dtype=self.dtype, device=self.device),
            requires_grad=True,
        )
        self.rot = nn.Parameter(
            torch.zeros(self.num_gaussians, 1, dtype=self.dtype, device=self.device),
            requires_grad=True,
        )
        self.feat_dim = sum(self.input_channels)
        self.feat = nn.Parameter(
            torch.rand(
                self.num_gaussians, self.feat_dim, dtype=self.dtype, device=self.device
            ),
            requires_grad=True,
        )
        self.vis_feat = nn.Parameter(
            torch.rand_like(self.feat), requires_grad=False
        )  # Only used for Gaussian ID visualization
        self._log_compression_rate()

    def _log_compression_rate(self):
        bytes_uncompressed = float(self.gt_images.numel())
        bpp_uncompressed = float(8 * self.feat_dim)
        self.worklog.info(
            f"Uncompressed: {bytes_uncompressed / 1e3:.2f} KB | {bpp_uncompressed:.3f} bpp | 8.0 bppc"
        )
        bits_compressed = (
            2 * self.pos_bits
            + 2 * self.scale_bits
            + self.rot_bits
            + self.feat_dim * self.feat_bits
        ) * self.total_num_gaussians
        bytes_compressed = bits_compressed / 8.0
        bpp_compressed = float(bits_compressed) / self.num_pixels
        bppc_compressed = bpp_compressed / self.feat_dim
        self.num_bytes = bytes_compressed
        self.worklog.info(
            f"Compressed: {bytes_compressed / 1e3:.2f} KB | {bpp_compressed:.3f} bpp | {bppc_compressed:.3f} bppc"
        )
        self.worklog.info(
            f"Compression rate: {bpp_uncompressed / bpp_compressed:.2f}x | {100.0 * bpp_compressed / bpp_uncompressed:.2f}%"
        )
        self.worklog.info("***********************************************")

    def _init_loss(self, args):
        self.l1_loss = None
        self.l2_loss = None
        self.ssim_loss = None
        self.l1_loss_ratio = args.l1_loss_ratio
        self.l2_loss_ratio = args.l2_loss_ratio
        self.ssim_loss_ratio = args.ssim_loss_ratio

    def _init_optimization(self, args):
        self.disable_tiles = args.disable_tiles
        self.start_step = 1
        self.max_steps = args.max_steps
        self.pos_lr = args.pos_lr
        self.scale_lr = args.scale_lr
        self.rot_lr = args.rot_lr
        self.feat_lr = args.feat_lr
        self.optimizer = torch.optim.Adam(
            [
                {"params": self.xy, "lr": self.pos_lr},
                {"params": self.scale, "lr": self.scale_lr},
                {"params": self.rot, "lr": self.rot_lr},
                {"params": self.feat, "lr": self.feat_lr},
            ]
        )
        self.disable_lr_schedule = args.disable_lr_schedule
        if not self.disable_lr_schedule:
            self.decay_ratio = args.decay_ratio
            self.check_decay_steps = args.check_decay_steps
            self.max_decay_times = args.max_decay_times
            self.decay_threshold = args.decay_threshold

    def _init_pos_scale_feat(self, args):
        self.init_mode = args.init_mode
        self.init_random_ratio = args.init_random_ratio
        self.pixel_xy = (
            get_grid(h=self.img_h, w=self.img_w)
            .to(dtype=self.dtype, device=self.device)
            .reshape(-1, 2)
        )
        with torch.no_grad():
            # Position
            if self.init_mode == "gradient":
                self._compute_gmap()
                self.xy.copy_(self._sample_pos(prob=self.image_gradients))
            elif self.init_mode == "saliency":
                self.smap_filter_size = args.smap_filter_size
                self._compute_smap(path="models")
                self.xy.copy_(self._sample_pos(prob=self.saliency))
            else:
                selected = np.random.choice(
                    self.num_pixels, self.num_gaussians, replace=False, p=None
                )
                self.xy.copy_(self.pixel_xy.detach().clone()[selected])
            # Scale
            self.scale.fill_(
                self.init_scale if self.disable_inverse_scale else 1.0 / self.init_scale
            )
            # Feature
            if not self.disable_color_init:
                self.feat.copy_(
                    self._get_target_features(positions=self.xy).detach().clone()
                )

    def _sample_pos(self, prob):
        num_random = round(self.init_random_ratio * self.num_gaussians)
        selected_random = np.random.choice(
            self.num_pixels, num_random, replace=False, p=None
        )
        selected_other = np.random.choice(
            self.num_pixels, self.num_gaussians - num_random, replace=False, p=prob
        )
        return torch.cat(
            [
                self.pixel_xy.detach().clone()[selected_random],
                self.pixel_xy.detach().clone()[selected_other],
            ],
            dim=0,
        )

    def _compute_gmap(self):
        gy, gx = compute_image_gradients(
            np.power(self.gt_images.detach().cpu().clone().numpy(), 1.0 / self.gamma)
        )
        g_norm = np.hypot(gy, gx).astype(np.float32)
        g_norm = g_norm / g_norm.max()
        save_image(g_norm, f"{self.log_dir}/gmap_res-{self.img_h:d}x{self.img_w:d}.png")
        g_norm = np.power(g_norm.reshape(-1), 2.0)
        self.image_gradients = g_norm / g_norm.sum()
        self.worklog.info("Image gradient map successfully saved")
        self.worklog.info("***********************************************")

    def _compute_smap(self, path):
        smap = get_smap(
            torch.pow(self.gt_images.detach().clone(), 1.0 / self.gamma),
            path,
            self.smap_filter_size,
        )
        save_image(smap, f"{self.log_dir}/smap_res-{self.img_h:d}x{self.img_w:d}.png")
        self.saliency = (smap / smap.sum()).reshape(-1)
        self.worklog.info("Saliency map successfully saved")
        self.worklog.info("***********************************************")

    def _get_target_features(self, positions):
        with torch.no_grad():
            # gt_images [1, C, H, W]; positions [1, 1, P, 2]; top-left [-1, -1]; bottom-right [1, 1]
            target_features = F.grid_sample(
                self.gt_images.unsqueeze(0),
                positions[None, None, ...] * 2.0 - 1.0,
                align_corners=False,
            )
            target_features = target_features[0, :, 0, :].permute(1, 0)  # [P, C]
        return target_features

    def _load_model(self):
        if self.ckpt_file != "":
            ckpt_path = os.path.join(self.ckpt_dir, self.ckpt_file)
        else:
            latest_step = get_latest_ckpt_step(self.ckpt_dir)
            if latest_step == -1:
                raise FileNotFoundError(f"No checkpoint found in '{self.ckpt_dir}'")
            ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step-{latest_step:d}.pt")
        checkpoint = torch.load(ckpt_path, weights_only=False)
        self.load_state_dict(checkpoint["state_dict"])
        self.optimizer.load_state_dict(checkpoint["optim_state_dict"])
        self.start_step = checkpoint["step"] + 1
        self.worklog.info(f"Checkpoint '{ckpt_path}' successfully loaded")
        self.worklog.info("***********************************************")

    def _save_model(self):
        if self.quantize:
            self._quantize()
        psnr, ssim = self._evaluate(log=False, upsample=False)
        self._evaluate_extra()
        ckpt_data = {
            "step": self.step,
            "psnr": psnr,
            "ssim": ssim,
            "lpips": self.lpips_final,
            "flip": self.flip_final,
            "msssim": self.msssim_final,
            "bytes": self.num_bytes,
            "time": self.total_time_accum,
            "state_dict": self.state_dict(),
            "optim_state_dict": self.optimizer.state_dict(),
        }
        save_path = f"{self.ckpt_dir}/ckpt_step-{self.step:d}.pt"
        torch.save(ckpt_data, save_path)
        self.worklog.info(f"Checkpoint 'ckpt_step-{self.step:d}.pt' successfully saved")
        self.worklog.info(
            f"PSNR: {psnr:.2f} | SSIM: {ssim:.4f} | LPIPS: {self.lpips_final:.4f} | FLIP: {self.flip_final:.4f} | MS-SSIM: {self.msssim_final:.4f}"
        )
        self.worklog.info("***********************************************")

    def _quantize(self):
        with torch.no_grad():
            self.xy.copy_(ste_quantize(self.xy, self.pos_bits))
            self.scale.copy_(ste_quantize(self.scale, self.scale_bits))
            self.rot.copy_(ste_quantize(self.rot, self.rot_bits))
            self.feat.copy_(ste_quantize(self.feat, self.feat_bits))

    def render(self, render_height=None):
        img_h, img_w = self.img_h, self.img_w
        if render_height is not None:
            img_h, img_w = render_height, round((float(render_height) / img_h) * img_w)
        tile_bounds = (
            (img_w + self.block_w - 1) // self.block_w,
            (img_h + self.block_h - 1) // self.block_h,
            1,
        )
        upsample_ratio = float(img_h) / self.img_h
        with torch.no_grad():
            num_prep_runs = 2
            for _ in range(num_prep_runs):
                self.forward(img_h, img_w, tile_bounds, upsample_ratio, benchmark=True)
            images, render_time = self.forward(
                img_h, img_w, tile_bounds, upsample_ratio
            )
            path = f"{self.eval_dir}/render_upsample-{upsample_ratio:.1f}_res-{img_h:d}x{img_w:d}"
            self._separate_and_save_images(
                images=images, channels=self.input_channels, path=path
            )
        self.worklog.info(f"Step: {self.start_step - 1:d} | Time: {render_time:.6f} s")
        self.worklog.info(f"Rendering at resolution ({img_h:d}, {img_w:d}) completed")
        self.worklog.info("***********************************************")

    def benchmark_render_time(self, num_reps, render_height=None):
        img_h, img_w = self.img_h, self.img_w
        if render_height is not None:
            img_h, img_w = render_height, round((float(render_height) / img_h) * img_w)
        tile_bounds = (
            (img_w + self.block_w - 1) // self.block_w,
            (img_h + self.block_h - 1) // self.block_h,
            1,
        )
        upsample_ratio = float(img_h) / self.img_h
        with torch.no_grad():
            render_time_all = np.zeros(num_reps, dtype=np.float32)
            num_prep_runs = 2
            for _ in range(num_prep_runs):
                self.forward(img_h, img_w, tile_bounds, upsample_ratio, benchmark=True)
            for rid in range(num_reps):
                render_time = self.forward(
                    img_h, img_w, tile_bounds, upsample_ratio, benchmark=True
                )
                render_time_all[rid] = render_time
        return render_time_all

    def forward(self, img_h, img_w, tile_bounds, upsample_ratio=None, benchmark=False):
        scale = self._get_scale(upsample_ratio=upsample_ratio)
        xy, rot, feat = self.xy, self.rot, self.feat
        if self.quantize:
            xy, scale, rot, feat = (
                ste_quantize(xy, self.pos_bits),
                ste_quantize(scale, self.scale_bits),
                ste_quantize(rot, self.rot_bits),
                ste_quantize(feat, self.feat_bits),
            )
        begin = perf_counter()
        tmp = project_gaussians_2d_scale_rot(xy, scale, rot, img_h, img_w, tile_bounds)
        xy, radii, conics, num_tiles_hit = tmp
        if not self.disable_tiles:
            enable_topk_norm = not self.disable_topk_norm
            tmp = (
                xy,
                radii,
                conics,
                num_tiles_hit,
                feat,
                img_h,
                img_w,
                self.block_h,
                self.block_w,
                enable_topk_norm,
            )
            out_image = rasterize_gaussians_sum(*tmp)
        else:
            tmp = xy, conics, feat, img_h, img_w
            out_image = rasterize_gaussians_no_tiles(*tmp)
        render_time = perf_counter() - begin
        if benchmark:
            return render_time
        out_image = (
            out_image.view(-1, img_h, img_w, self.feat_dim)
            .permute(0, 3, 1, 2)
            .contiguous()
        )
        return out_image.squeeze(dim=0), render_time

    def _get_scale(self, upsample_ratio=None):
        scale = self.scale
        if not self.disable_inverse_scale:
            scale = 1.0 / scale
        if upsample_ratio is not None:
            scale = upsample_ratio * scale
        return scale

    def _visualize_gaussian_id(self, img_h, img_w, tile_bounds, upsample_ratio=None):
        scale = self._get_scale(upsample_ratio=upsample_ratio)
        xy, rot, feat = self.xy, self.rot, self.feat
        if self.quantize:
            xy, scale, rot, feat = (
                ste_quantize(xy, self.pos_bits),
                ste_quantize(scale, self.scale_bits),
                ste_quantize(rot, self.rot_bits),
                ste_quantize(feat, self.feat_bits),
            )
        feat = self.vis_feat * feat.norm(dim=-1, keepdim=True)
        tmp = project_gaussians_2d_scale_rot(xy, scale, rot, img_h, img_w, tile_bounds)
        xy, radii, conics, num_tiles_hit = tmp
        if not self.disable_tiles:
            enable_topk_norm = not self.disable_topk_norm
            tmp = (
                xy,
                radii,
                conics,
                num_tiles_hit,
                feat,
                img_h,
                img_w,
                self.block_h,
                self.block_w,
                enable_topk_norm,
            )
            out_image = rasterize_gaussians_sum(*tmp)
        else:
            tmp = xy, conics, feat, img_h, img_w
            out_image = rasterize_gaussians_no_tiles(*tmp)
        out_image = (
            out_image.view(-1, img_h, img_w, self.feat_dim)
            .permute(0, 3, 1, 2)
            .contiguous()
        )
        return out_image.squeeze(dim=0)

    def optimize(self):
        self.psnr_curr, self.ssim_curr = 0.0, 0.0
        self.best_psnr, self.best_ssim = 0.0, 0.0
        self.decay_times, self.no_improvement_steps = 0, 0
        self.render_time_accum, self.total_time_accum = 0.0, 0.0
        self.lpips_final, self.flip_final, self.msssim_final = 1.0, 1.0, 0.0

        self.step = 0
        with torch.no_grad():
            self._log_images(log_final=False, plot_gaussians=self.vis_gaussians)
        for step in range(self.start_step, self.max_steps + 1):
            self.step = step
            self.optimizer.zero_grad()
            # Rendering
            images, render_time = self.forward(self.img_h, self.img_w, self.tile_bounds)
            self.render_time_accum += render_time
            # Optimization
            begin = perf_counter()
            self._get_total_loss(images)
            self.total_loss.backward()
            self.optimizer.step()
            self.total_time_accum += perf_counter() - begin + render_time
            # Logging
            terminate = False
            with torch.no_grad():
                if self.step % self.eval_steps == 0:
                    self._evaluate(log=True, upsample=False)
                    if (
                        not self.disable_lr_schedule
                        and self.num_gaussians == self.total_num_gaussians
                    ):
                        terminate = self._lr_schedule()
                if self.step % self.save_image_steps == 0:
                    self._log_images(log_final=False, plot_gaussians=self.vis_gaussians)
                if (
                    self.step % self.save_ckpt_steps == 0
                    and self.num_gaussians == self.total_num_gaussians
                ):
                    self._save_model()
                if (
                    not self.disable_prog_optim
                    and self.step % self.add_steps == 0
                    and self.num_gaussians < self.total_num_gaussians
                ):
                    self._add_gaussians(
                        self.max_add_num, plot_gaussians=self.vis_gaussians
                    )
                if terminate:
                    break
        with torch.no_grad():
            self._log_images(log_final=True, plot_gaussians=self.vis_gaussians)
            self._save_model()
        self.worklog.info("Optimization completed")
        self.worklog.info("***********************************************")
        self.worklog.info(
            f"Mean scale: {self._get_scale().mean().item():.4f} (pixel) | {self.scale.mean().item():.4f} (raw)"
        )
        self.worklog.info("***********************************************")
        return self.psnr_curr, self.ssim_curr

    def _get_total_loss(self, images):
        self.total_loss = 0
        if self.l1_loss_ratio > 1e-7:
            self.l1_loss = self.l1_loss_ratio * F.l1_loss(images, self.gt_images)
            self.total_loss += self.l1_loss
        else:
            self.l1_loss = None
        if self.l2_loss_ratio > 1e-7:
            self.l2_loss = self.l2_loss_ratio * F.mse_loss(images, self.gt_images)
            self.total_loss += self.l2_loss
        else:
            self.l2_loss = None
        if self.ssim_loss_ratio > 1e-7:
            self.ssim_loss = self.ssim_loss_ratio * (
                1 - fused_ssim(images.unsqueeze(0), self.gt_images.unsqueeze(0))
            )
            self.total_loss += self.ssim_loss
        else:
            self.ssim_loss = None

    def _evaluate(self, log=True, upsample=False):
        if upsample:  # Do not log performance metrics for upsampled images
            log = False
        images = torch.pow(
            torch.clamp(self._render_images(upsample=upsample), 0.0, 1.0),
            1.0 / self.gamma,
        )
        gt_images = torch.pow(
            self.gt_images_upsampled if upsample else self.gt_images, 1.0 / self.gamma
        )
        psnr = get_psnr(images, gt_images).item()
        ssim = fused_ssim(images.unsqueeze(0), gt_images.unsqueeze(0)).item()
        if log:
            self.psnr_curr, self.ssim_curr = psnr, ssim
            loss_results = f"Loss: {self.total_loss.item():.4f}"
            loss_results += (
                f", L1: {self.l1_loss.item():.4f}" if self.l1_loss is not None else ""
            )
            loss_results += (
                f", L2: {self.l2_loss.item():.4f}" if self.l2_loss is not None else ""
            )
            loss_results += (
                f", SSIM: {self.ssim_loss.item():.4f}"
                if self.ssim_loss is not None
                else ""
            )
            time_results = f"Total: {self.total_time_accum:.2f} s | Render: {self.render_time_accum:.2f} s"
            self.worklog.info(
                f"Step: {self.step:d} | {time_results} | {loss_results} | PSNR: {self.psnr_curr:.2f} | SSIM: {self.ssim_curr:.4f}"
            )
        return psnr, ssim

    def _evaluate_extra(self):
        images = torch.pow(
            torch.clamp(self._render_images(upsample=False), 0.0, 1.0), 1.0 / self.gamma
        )[None, ...]
        gt_images = torch.pow(self.gt_images, 1.0 / self.gamma)[None, ...]
        msssim_metric = (
            MS_SSIM(data_range=1.0, size_average=True, channel=self.feat_dim)
            .to(device=self.device)
            .eval()
        )
        self.msssim_final = msssim_metric(images, gt_images).item()
        lpips_metric = LPIPS(net="alex").to(device=self.device).eval()
        flip_metric = LDRFLIPLoss().to(device=self.device).eval()
        num_channels = 1 if self.feat_dim < 3 else 3
        self.lpips_final = lpips_metric(
            images[:, :num_channels], gt_images[:, :num_channels]
        ).item()
        if self.feat_dim >= 3:
            self.flip_final = flip_metric(images[:, :3], gt_images[:, :3]).item()

    def _log_images(self, log_final=False, plot_gaussians=False):
        images = self._render_images(upsample=False)
        if log_final:
            path = f"{self.log_dir}/render_res-{self.img_h:d}x{self.img_w:d}"
            self._separate_and_save_images(
                images=images, channels=self.input_channels, path=path
            )
        psnr, ssim = self._evaluate(log=False, upsample=False)
        path = f"{self.train_dir}/render_step-{self.step:d}_psnr-{psnr:.2f}_ssim-{ssim:.4f}_res-{self.img_h:d}x{self.img_w:d}"
        self._separate_and_save_images(
            images=images, channels=self.input_channels, path=path
        )
        if plot_gaussians:
            path = f"{self.train_dir}/gaussian_step-{self.step:d}_psnr-{psnr:.2f}_ssim-{ssim:.4f}_res-{self.img_h:d}x{self.img_w:d}"
            visualize_gaussians(
                path,
                self.xy,
                self._get_scale(),
                self.rot,
                self.feat,
                self.img_h,
                self.img_w,
                self.input_channels,
                alpha=0.8,
                gamma=self.gamma,
            )
            images = self._visualize_gaussian_id(
                self.img_h, self.img_w, self.tile_bounds
            )
            path = f"{self.train_dir}/gaussian-id_step-{self.step:d}_psnr-{psnr:.2f}_ssim-{ssim:.4f}_res-{self.img_h:d}x{self.img_w:d}"
            self._separate_and_save_images(
                images=images, channels=self.input_channels, path=path
            )
        if self.downsample:
            images = self._render_images(upsample=True)
            psnr, ssim = self._evaluate(log=False, upsample=True)
            img_h, img_w = self.img_h_upsampled, self.img_w_upsampled
            path = f"{self.train_dir}/render_upsample-{self.downsample_ratio:.1f}_step-{self.step:d}_psnr-{psnr:.2f}_ssim-{ssim:.4f}_res-{img_h:d}x{img_w:d}"
            self._separate_and_save_images(
                images=images, channels=self.input_channels, path=path
            )

    def _render_images(self, upsample=False):
        if upsample:
            images, _ = self.forward(
                self.img_h_upsampled,
                self.img_w_upsampled,
                self.tile_bounds_upsampled,
                upsample_ratio=self.downsample_ratio,
            )
        else:
            images, _ = self.forward(self.img_h, self.img_w, self.tile_bounds)
        return images

    def _lr_schedule(self):
        if (
            self.psnr_curr <= self.best_psnr + 100 * self.decay_threshold
            or self.ssim_curr <= self.best_ssim + self.decay_threshold
        ):
            self.no_improvement_steps += self.eval_steps
            if self.no_improvement_steps >= self.check_decay_steps:
                self.no_improvement_steps = 0
                self.decay_times += 1
                if self.decay_times > self.max_decay_times:
                    return True
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] /= self.decay_ratio
                self.worklog.info(f"Learning rate decayed by {self.decay_ratio:.1f}")
                self.worklog.info("***********************************************")
            return False
        else:
            self.best_psnr = self.psnr_curr
            self.best_ssim = self.ssim_curr
            self.no_improvement_steps = 0
            return False

    def _add_gaussians(self, add_num, plot_gaussians=False):
        add_num = min(
            add_num, self.max_add_num, self.total_num_gaussians - self.num_gaussians
        )
        if add_num <= 0:
            return
        raw_images = self._render_images(upsample=False)
        images = torch.pow(torch.clamp(raw_images, 0.0, 1.0), 1.0 / self.gamma)
        gt_images = torch.pow(self.gt_images, 1.0 / self.gamma)
        kernel_size = round(np.sqrt(self.img_h * self.img_w) // 400)
        if kernel_size >= 1:
            kernel_size = max(3, kernel_size)
            kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size
            gt_images = gaussian_blur(img=gt_images, kernel_size=kernel_size)
        diff_map = (gt_images - images).detach().clone()
        error_map = torch.pow(torch.abs(diff_map).mean(dim=0).reshape(-1), 2.0)
        sample_prob = (error_map / error_map.sum()).cpu().numpy()
        selected = np.random.choice(
            self.num_pixels, add_num, replace=False, p=sample_prob
        )
        # New Gaussians
        new_xy = self.pixel_xy.detach().clone()[selected]
        new_scale = torch.ones(add_num, 2, dtype=self.dtype, device=self.device)
        init_scale = self.init_scale
        new_scale.fill_(init_scale if self.disable_inverse_scale else 1.0 / init_scale)
        new_rot = torch.zeros(add_num, 1, dtype=self.dtype, device=self.device)
        new_feat = diff_map.permute(1, 2, 0).reshape(-1, self.feat_dim)[selected]
        new_vis_feat = torch.rand_like(new_feat)
        # Old Gaussians
        old_xy = self.xy.detach().clone()
        old_scale = self.scale.detach().clone()
        old_rot = self.rot.detach().clone()
        old_feat = self.feat.detach().clone()
        old_vis_feat = self.vis_feat.detach().clone()
        # Update trainable parameters
        self.num_gaussians += add_num
        all_xy = torch.cat([old_xy, new_xy], dim=0)
        all_scale = torch.cat([old_scale, new_scale], dim=0)
        all_rot = torch.cat([old_rot, new_rot], dim=0)
        all_feat = torch.cat([old_feat, new_feat], dim=0)
        all_vis_feat = torch.cat([old_vis_feat, new_vis_feat], dim=0)
        self.xy = nn.Parameter(all_xy, requires_grad=True)
        self.scale = nn.Parameter(all_scale, requires_grad=True)
        self.rot = nn.Parameter(all_rot, requires_grad=True)
        self.feat = nn.Parameter(all_feat, requires_grad=True)
        self.vis_feat = nn.Parameter(all_vis_feat, requires_grad=False)
        # Plot Gaussians
        if plot_gaussians:
            path = f"{self.train_dir}/add-gaussian_step-{self.step:d}_num-{self.num_gaussians:d}_res-{self.img_h:d}x{self.img_w:d}"
            every_n = max(1, self.total_num_gaussians // 2000)
            size = (self.img_h * self.img_w) / 1e4
            visualize_added_gaussians(
                path,
                raw_images,
                old_xy,
                new_xy,
                self.input_channels,
                size=size,
                every_n=every_n,
                alpha=0.8,
                gamma=self.gamma,
            )
        # Update optimizer
        self.optimizer = torch.optim.Adam(
            [
                {"params": self.xy, "lr": self.pos_lr},
                {"params": self.scale, "lr": self.scale_lr},
                {"params": self.rot, "lr": self.rot_lr},
                {"params": self.feat, "lr": self.feat_lr},
            ]
        )
        self.worklog.info(
            f"Step: {self.step:d} | Adding {add_num:d} Gaussians ({self.num_gaussians - add_num:d} -> {self.num_gaussians:d})"
        )
        self.worklog.info("***********************************************")