gradio_models.py

import logging
import math
import os
import threading
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 torchvision.transforms.functional import gaussian_blur
from PIL import Image

from gsplat import (
    project_gaussians_2d_scale_rot,
    rasterize_gaussians_no_tiles,
    rasterize_gaussians_sum,
)
from utils.image_utils import (
    compute_image_gradients,
    get_grid,
    get_psnr,
    load_images,
    to_output_format,
)
from utils.misc_utils import set_random_seed
from utils.quantization_utils import ste_quantize
from utils.saliency_utils import get_smap


class StreamingResults:
    """Container for streaming training results"""

    def __init__(self):
        self.step = 0
        self.total_steps = 0
        self.current_render = None
        self.current_gaussian_id = None
        self.initialization_map = None  # Single map for current initialization mode
        self.final_render = None
        self.final_checkpoint_path = None
        self.training_logs = []
        self.metrics = {
            "psnr": 0.0,
            "ssim": 0.0,
            "loss": 0.0,
            "render_time": 0.0,
            "total_time": 0.0,
        }
        self.is_complete = False
        # Store all step results for interactive browsing
        self.step_renders = {}  # {step: PIL_Image}
        self.step_gaussian_ids = {}  # {step: PIL_Image}
        # For async visualization generation
        self.vis_lock = threading.Lock()


class GradioStreamingHandler(logging.Handler):
    """Custom logging handler that captures logs for Gradio streaming"""

    def __init__(self, results_container: StreamingResults):
        super().__init__()
        self.results = results_container

    def emit(self, record):
        log_entry = self.format(record)
        self.results.training_logs.append(log_entry)
        # Keep only last 100 log entries to avoid memory issues
        if len(self.results.training_logs) > 100:
            self.results.training_logs = self.results.training_logs[-100:]


class GradioGaussianSplatting2D(nn.Module):
    """Gradio-optimized version of GaussianSplatting2D with streaming capabilities"""

    def __init__(self, args, results_container: StreamingResults):
        super(GradioGaussianSplatting2D, self).__init__()
        self.results = results_container
        self.evaluate = args.eval
        set_random_seed(seed=args.seed)

        # Device setup
        if torch.cuda.is_available():
            torch.cuda.set_device(0)
            self.device = torch.device("cuda:0")
        else:
            self.device = torch.device("cpu")
        self.dtype = torch.float32

        # Initialize components
        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 = getattr(args, "log_dir", "temp_gradio_logs")
        self.vis_gaussians = args.vis_gaussians
        self.save_image_steps = args.save_image_steps
        self.eval_steps = args.eval_steps

        # Set up streaming logger
        self.worklog = logging.getLogger("GradioImageGS")
        self.worklog.handlers.clear()  # Remove existing handlers

        # Add our streaming handler
        stream_handler = GradioStreamingHandler(self.results)
        stream_handler.setFormatter(
            logging.Formatter(fmt="[{asctime}] {message}", style="{")
        )
        self.worklog.addHandler(stream_handler)
        self.worklog.setLevel(logging.INFO)

        self.worklog.info(
            f"Start optimizing {args.num_gaussians:d} Gaussians for '{args.input_path}'"
        )

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

        self.num_pixels = self.img_h * self.img_w

    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 = (
            int(self.gt_images.shape[1]),
            int(self.gt_images.shape[2]),
        )
        self.tile_bounds = (
            int((self.img_w + self.block_w - 1) // self.block_w),
            int((self.img_h + self.block_h - 1) // self.block_h),
            int(1),
        )

    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
        self.eps = 1e-7 if args.disable_tiles else 1e-4
        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

        # Initialize parameters
        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)

        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}%"
        )

    def _init_loss(self, args):
        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.results.total_steps = (
            args.max_steps
        )  # Set total steps for streaming progress
        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 initialization
            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()
                self.xy.copy_(self._sample_pos(prob=self.saliency))
            else:  # random mode
                selected = np.random.choice(
                    self.num_pixels, self.num_gaussians, replace=False, p=None
                )
                self.xy.copy_(self.pixel_xy.detach().clone()[selected])
                # For random mode, create a simple random noise pattern
                if self.init_mode == "random":
                    random_pattern = np.random.rand(self.img_h, self.img_w)
                    self.results.initialization_map = Image.fromarray(
                        (random_pattern * 255).astype(np.uint8)
                    )

            # Scale initialization
            self.scale.fill_(
                self.init_scale if self.disable_inverse_scale else 1.0 / self.init_scale
            )

            # Feature initialization
            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()

        # Store gradient map for streaming (only if this is the selected initialization mode)
        if self.init_mode == "gradient":
            self.results.initialization_map = Image.fromarray(
                (g_norm * 255).astype(np.uint8)
            )

        g_norm = np.power(g_norm.reshape(-1), 2.0)
        self.image_gradients = g_norm / g_norm.sum()
        self.worklog.info("Image gradient map computed")

    def _compute_smap(self):
        smap = get_smap(
            torch.pow(self.gt_images.detach().clone(), 1.0 / self.gamma),
            "models",
            self.smap_filter_size,
        )

        # Store saliency map for streaming (only if this is the selected initialization mode)
        if self.init_mode == "saliency":
            self.results.initialization_map = Image.fromarray(
                (smap * 255).astype(np.uint8)
            )

        self.saliency = (smap / smap.sum()).reshape(-1)
        self.worklog.info("Saliency map computed")

    def _get_target_features(self, positions):
        with torch.no_grad():
            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)
        return target_features

    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, int(img_h), int(img_w), tile_bounds
        )
        xy, radii, conics, num_tiles_hit = tmp

        # Ensure correct tensor types for CUDA backend
        # Note: The custom gsplat CUDA code expects int32 for both radii and num_tiles_hit
        num_tiles_hit = num_tiles_hit.to(dtype=torch.int32)
        radii = radii.to(dtype=torch.int32)

        if not self.disable_tiles:
            enable_topk_norm = not self.disable_topk_norm

            out_image = rasterize_gaussians_sum(
                xy,
                radii,
                conics,
                num_tiles_hit,
                feat,
                int(img_h),
                int(img_w),
                int(self.block_h),
                int(self.block_w),
                enable_topk_norm,
            )
        else:
            tmp = xy, conics, feat, int(img_h), int(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, int(img_h), int(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 _tensor_to_pil_image(self, tensor_image):
        """Convert tensor image to PIL Image for streaming"""
        if tensor_image is None:
            return None

        # Convert to numpy and apply gamma correction
        image_np = (
            torch.pow(torch.clamp(tensor_image, 0.0, 1.0), 1.0 / self.gamma)
            .detach()
            .cpu()
            .numpy()
        )

        # Convert to uint8 format
        image_formatted = to_output_format(image_np, gamma=None)
        return Image.fromarray(image_formatted)

    def _create_gaussian_id_visualization(self):
        """Create Gaussian ID visualization as PIL Image using rasterization with vis_feat"""
        if not self.vis_gaussians:
            return None

        try:
            # Use vis_feat for ID visualization (this creates unique colors per Gaussian)
            feat = self.vis_feat * self.feat.norm(dim=-1, keepdim=True)

            # Render with ID features
            scale = self._get_scale()
            xy, rot = self.xy, self.rot

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

            tmp = project_gaussians_2d_scale_rot(
                xy, scale, rot, int(self.img_h), int(self.img_w), self.tile_bounds
            )
            xy, radii, conics, num_tiles_hit = tmp

            # Ensure correct tensor types for CUDA backend
            # Note: The custom gsplat CUDA code expects int32 for both radii and num_tiles_hit
            num_tiles_hit = num_tiles_hit.to(dtype=torch.int32)
            radii = radii.to(dtype=torch.int32)

            if not self.disable_tiles:
                enable_topk_norm = not self.disable_topk_norm
                out_image = rasterize_gaussians_sum(
                    xy,
                    radii,
                    conics,
                    num_tiles_hit,
                    feat,
                    int(self.img_h),
                    int(self.img_w),
                    int(self.block_h),
                    int(self.block_w),
                    enable_topk_norm,
                )
            else:
                tmp = xy, conics, feat, int(self.img_h), int(self.img_w)
                out_image = rasterize_gaussians_no_tiles(*tmp)

            out_image = (
                out_image.view(-1, int(self.img_h), int(self.img_w), self.feat_dim)
                .permute(0, 3, 1, 2)
                .contiguous()
            ).squeeze(dim=0)

            return self._tensor_to_pil_image(out_image)

        except Exception as e:
            self.worklog.error(f"Error creating Gaussian ID visualization: {e}")
            return None

    def optimize(self):
        """Main optimization loop with streaming updates"""
        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

        # Initialize attributes needed for evaluation
        self.l1_loss = None
        self.l2_loss = None
        self.ssim_loss = None
        self.stop_requested = False

        # Initial render and update
        with torch.no_grad():
            images, _ = self.forward(self.img_h, self.img_w, self.tile_bounds)
            self.results.current_render = self._tensor_to_pil_image(images)
            if self.vis_gaussians:
                try:
                    self.results.current_gaussian_id = (
                        self._create_gaussian_id_visualization()
                    )
                    self.worklog.info(
                        f"Initial visualizations created - Render: {'✓' if self.results.current_render else '✗'}, ID: {'✓' if self.results.current_gaussian_id else '✗'}"
                    )
                except Exception as e:
                    self.worklog.error(f"Error creating initial visualizations: {e}")
                    self.results.current_gaussian_id = None

            # Store initial results (step 0)
            self.results.step_renders[0] = self.results.current_render
            if self.vis_gaussians:
                self.results.step_gaussian_ids[0] = self.results.current_gaussian_id

        for step in range(self.start_step, self.max_steps + 1):
            self.step = step
            self.results.step = step

            self.optimizer.zero_grad()

            # Forward pass
            images, render_time = self.forward(self.img_h, self.img_w, self.tile_bounds)
            self.render_time_accum += render_time

            # Compute loss
            begin = perf_counter()
            self._get_total_loss(images)
            self.total_loss.backward()
            self.optimizer.step()
            self.total_time_accum += perf_counter() - begin + render_time

            # Update streaming results
            with torch.no_grad():
                if step % self.eval_steps == 0:
                    self._evaluate_and_update_stream(images)

                # Update render image more frequently, but visualizations less frequently
                render_update_freq = max(
                    50, self.save_image_steps // 2
                )  # Render updates every 50 steps
                vis_update_freq = max(
                    200, self.save_image_steps
                )  # Visualizations every 200 steps

                if step % render_update_freq == 0:
                    render_img = self._tensor_to_pil_image(images)
                    self.results.current_render = render_img

                # Only update Gaussian ID visualization less frequently
                if step % vis_update_freq == 0 and self.vis_gaussians:
                    # Generate Gaussian ID visualization asynchronously
                    def generate_gaussian_id_async():
                        try:
                            with self.results.vis_lock:
                                gaussian_id_vis = (
                                    self._create_gaussian_id_visualization()
                                )
                                self.results.current_gaussian_id = gaussian_id_vis

                        except Exception as e:
                            self.worklog.error(
                                f"Error creating Gaussian ID visualization at step {step}: {e}"
                            )
                            with self.results.vis_lock:
                                self.results.current_gaussian_id = None

                    # Start async visualization generation
                    vis_thread = threading.Thread(target=generate_gaussian_id_async)
                    vis_thread.daemon = True
                    vis_thread.start()

                # Store results for interactive browsing only at save_image_steps intervals
                if step % self.save_image_steps == 0:
                    # Store the current render for browsing
                    if self.results.current_render:
                        self.results.step_renders[step] = self.results.current_render

                    # Store Gaussian ID visualization for browsing
                    if self.vis_gaussians and self.results.current_gaussian_id:
                        self.results.step_gaussian_ids[step] = (
                            self.results.current_gaussian_id
                        )

                # Progressive optimization
                if (
                    not self.disable_prog_optim
                    and step % self.add_steps == 0
                    and self.num_gaussians < self.total_num_gaussians
                ):
                    self._add_gaussians(self.max_add_num)

                # Learning rate schedule
                terminate = False
                if (
                    not self.disable_lr_schedule
                    and self.num_gaussians == self.total_num_gaussians
                    and step % self.eval_steps == 0
                ):
                    terminate = self._lr_schedule()

                if terminate or self.stop_requested:
                    if self.stop_requested:
                        self.worklog.info("Training stopped by user request")
                    break

        # Final updates
        with torch.no_grad():
            images, _ = self.forward(self.img_h, self.img_w, self.tile_bounds)
            self.results.final_render = self._tensor_to_pil_image(images)

            # Save final checkpoint and store path
            self._save_final_checkpoint()

            self.results.is_complete = True
            self.worklog.info("Optimization completed")

    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_and_update_stream(self, images):
        """Evaluate current state and update streaming results"""
        gamma_corrected_images = torch.pow(
            torch.clamp(images, 0.0, 1.0), 1.0 / self.gamma
        )
        gamma_corrected_gt = torch.pow(self.gt_images, 1.0 / self.gamma)

        psnr = get_psnr(gamma_corrected_images, gamma_corrected_gt).item()
        ssim = fused_ssim(
            gamma_corrected_images.unsqueeze(0), gamma_corrected_gt.unsqueeze(0)
        ).item()

        self.psnr_curr, self.ssim_curr = psnr, ssim

        # Update metrics
        self.results.metrics.update(
            {
                "psnr": psnr,
                "ssim": ssim,
                "loss": self.total_loss.item(),
                "render_time": self.render_time_accum,
                "total_time": self.total_time_accum,
            }
        )

        # Log progress
        loss_results = f"Loss: {self.total_loss.item():.4f}"
        if self.l1_loss is not None:
            loss_results += f", L1: {self.l1_loss.item():.4f}"
        if self.l2_loss is not None:
            loss_results += f", L2: {self.l2_loss.item():.4f}"
        if self.ssim_loss is not None:
            loss_results += f", SSIM: {self.ssim_loss.item():.4f}"

        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: {psnr:.2f} | SSIM: {ssim:.4f}"
        )

    def _save_final_checkpoint(self):
        """Save final checkpoint and store the path"""
        if self.quantize:
            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))

        # Create checkpoint directory
        ckpt_dir = os.path.join(self.log_dir, "checkpoints")
        os.makedirs(ckpt_dir, exist_ok=True)

        psnr = self.results.metrics.get("psnr", 0.0)
        ssim = self.results.metrics.get("ssim", 0.0)

        ckpt_data = {
            "step": self.step,
            "psnr": psnr,
            "ssim": ssim,
            "bytes": getattr(self, "num_bytes", 0),
            "time": self.total_time_accum,
            "state_dict": self.state_dict(),
            "optim_state_dict": self.optimizer.state_dict(),
        }

        ckpt_path = os.path.join(ckpt_dir, f"ckpt_step-{self.step:d}.pt")
        torch.save(ckpt_data, ckpt_path)
        self.results.final_checkpoint_path = ckpt_path

        self.worklog.info(f"Final checkpoint saved: {ckpt_path}")

    def _lr_schedule(self):
        """Learning rate scheduling logic"""
        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}")
            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):
        """Add Gaussians during progressive optimization"""
        add_num = min(
            add_num, self.max_add_num, self.total_num_gaussians - self.num_gaussians
        )
        if add_num <= 0:
            return

        # Compute error map for new Gaussian placement
        raw_images, _ = self.forward(self.img_h, self.img_w, self.tile_bounds)
        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
        )

        # Create 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)

        # Update parameters
        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()

        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)

        # 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})"
        )