content: import random
  import tqdm
  
  import torch
  import torch.nn.functional as F
  from pytorch_msssim import SSIM, MS_SSIM
  
  import comfy.utils
  
  from .main_3DGS_renderer import GaussianSplattingRenderer
  
  from shared_utils.camera_utils import BaseCameraController, MiniCam, calculate_fovX, get_projection_matrix
  from shared_utils.image_utils import prepare_torch_img
  
  class GSParams:
      def __init__(
          self,
          training_iterations=30_000, 
          batch_size=1,
          lambda_ssim=0.2,
          lambda_alpha=3,
          lambda_offset=0,
          lambda_offset_opacity=0,
          invert_bg_prob = 0.5,
          feature_lr=0.0025, 
          opacity_lr=0.05, 
          scaling_lr=0.005, 
          rotation_lr=0.001, 
          position_lr_init=0.00016, 
          position_lr_final=0.0000016,
          position_lr_delay_mult=0.01,
          position_lr_max_steps=30_000,
          num_pts=10_000,
          K=3,
          percent_dense=0.01,
          density_start_iter=500,
          density_end_iter=15_000,
          densification_interval=100,
          opacity_reset_interval=3000,
          densify_grad_threshold=0.0002,
          sh_degree=3
      ):
          
          # training params
          self.training_iterations = training_iterations
          self.batch_size = batch_size
          self.lambda_ssim = lambda_ssim
          self.lambda_alpha = lambda_alpha
          self.lambda_offset = lambda_offset
          self.lambda_offset_opacity = lambda_offset_opacity
          self.invert_bg_prob = invert_bg_prob
          
          # learning params
          self.feature_lr = feature_lr
          self.opacity_lr = opacity_lr
          self.scaling_lr = scaling_lr
          self.rotation_lr = rotation_lr
          self.position_lr_init = position_lr_init
          self.position_lr_final = position_lr_final
          self.position_lr_delay_mult = position_lr_delay_mult
          self.position_lr_max_steps = position_lr_max_steps
          
          # densify and prune params
          self.num_pts = num_pts
          self.K = K
          self.percent_dense = percent_dense
          self.density_start_iter=density_start_iter
          self.density_end_iter=density_end_iter
          self.densification_interval=densification_interval
          self.opacity_reset_interval=opacity_reset_interval
          self.densify_grad_threshold=densify_grad_threshold
          
          # other gaussian params
          self.sh_degree = sh_degree
          
  class GaussianSplattingCameraController(BaseCameraController):
      def get_render_result(self, render_pose, bg_color, **kwargs):
          render_cam = MiniCam(render_pose, self.cam.W, self.cam.H, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, self.projection_matrix)
          return self.renderer.render(render_cam, bg_color=bg_color, **kwargs)
      
      def post_init(self):
          self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
  
  class GaussianSplatting3D:
              
      def __init__(self, gs_params=None, init_input=None, device='cuda'):
          self.device = torch.device(device)
          
          # prepare renderer for optimization
          if gs_params is None:
              gs_params = GSParams()
              
          self.renderer = GaussianSplattingRenderer(sh_degree=gs_params.sh_degree)
          self.renderer.initialize(init_input, num_pts=gs_params.num_pts)
  
          # setup training
          self.renderer.gaussians.training_setup(gs_params)
          # do not do progressive sh-level
          self.renderer.gaussians.active_sh_degree = self.renderer.gaussians.max_sh_degree
          self.optimizer = self.renderer.gaussians.optimizer
          #self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
          self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
          
          self.gs_params = gs_params
      
      def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
          self.ref_imgs_num = len(reference_images)
          
          self.ref_size_H = reference_images[0].shape[0]
          self.ref_size_W = reference_images[0].shape[1]
          
          # default camera settings
          self.cam_controller = GaussianSplattingCameraController(
              self.renderer, self.ref_size_W, self.ref_size_H, reference_orbit_camera_fovy, self.gs_params.invert_bg_prob, None, self.device
          )
  
          self.all_ref_cam_poses = reference_orbit_camera_poses
          
          # prepare reference images and masks
          ref_imgs_torch_list = []
          ref_masks_torch_list = []
          for i in range(self.ref_imgs_num):
              ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
              ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
              
          self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
          self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
      
      def training(self):
          starter = torch.cuda.Event(enable_timing=True)
          ender = torch.cuda.Event(enable_timing=True)
          starter.record()
          
          ref_imgs_masked = []
          for i in range(self.ref_imgs_num):
              ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
              
          ref_imgs_num_minus_1 = self.ref_imgs_num-1
          
          comfy_pbar = comfy.utils.ProgressBar(self.gs_params.training_iterations)
  
          for step in tqdm.trange(self.gs_params.training_iterations):
  
              ### calculate loss between reference and rendered image from known view
              loss = 0
              masked_rendered_img_batch = []
              masked_ref_img_batch = []
              masks_batch = []
              ref_masks_batch = []
  
              step_ratio = min(1, step / self.gs_params.training_iterations)
  
              # update lr
              self.renderer.gaussians.update_learning_rate(step)
  
              loss = 0
  
              for _ in range(self.gs_params.batch_size):
                  ### calculate loss between reference and rendered image from known view
                      
                  i = random.randint(0, ref_imgs_num_minus_1)
                  
                  out = self.cam_controller.render_at_pose(self.all_ref_cam_poses[i])
                  
                  image = out["image"] # [3, H, W] in [0, 1]
                  mask = out["alpha"] # [1, H, W] in [0, 1]
                  
                  ref_mask = self.ref_masks_torch[i]
                  image_masked = (image * ref_mask).unsqueeze(0)
                  
                  masked_rendered_img_batch.append(image_masked)
                  masked_ref_img_batch.append(ref_imgs_masked[i])
                  masks_batch.append(mask.unsqueeze(0))
                  ref_masks_batch.append(ref_mask.unsqueeze(0))
              
              masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
              masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
              masks_batch_torch = torch.cat(masks_batch, dim=0)
              ref_masks_batch_torch = torch.cat(ref_masks_batch, dim=0)
                  
              #loss_scaler = self.gs_params.loss_scale * step_ratio
                  
              # rgb loss            
              loss += (1 - self.gs_params.lambda_ssim) * F.l1_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
  
              # alpha loss
              loss += self.gs_params.lambda_alpha * F.mse_loss(masks_batch_torch, ref_masks_batch_torch)
              
              # D-SSIM loss
              # [1, 3, H, W] in [0, 1]
              #loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
              loss += self.gs_params.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
  
              if self.gs_params.lambda_offset > 0:
                  # Reference offset loss
                  offset_norm = self.renderer.gaussians.get_xyz_offset.norm(dim=-1, keepdim=True)
                  loss += self.gs_params.lambda_offset * torch.mean(offset_norm)
              
              if self.gs_params.lambda_offset_opacity > 0:
                  # Alpha penalty loss
                  loss += self.gs_params.lambda_offset_opacity * torch.mean(offset_norm.detach() * self.renderer.gaussians.get_opacity)
  
  
              # optimize step
              loss.backward()
              self.optimizer.step()
              self.optimizer.zero_grad()
  
              # densify and prune
              if step >= self.gs_params.density_start_iter and step <= self.gs_params.density_end_iter:
                  viewspace_point_tensor, visibility_filter, radii = out["viewspace_points"], out["visibility_filter"], out["radii"]
                  self.renderer.gaussians.max_radii2D[visibility_filter] = torch.max(self.renderer.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
                  self.renderer.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
  
                  if step % self.gs_params.densification_interval == 0:
                      self.renderer.gaussians.densify_and_prune(self.gs_params.densify_grad_threshold, min_opacity=0.005, extent=4, max_screen_size=1)
                      #self.renderer.gaussians.densify_and_prune_by_compatness(self.gs_params.K, min_opacity=0.01, extent=4, max_screen_size=1)
                      
                      #self.renderer.gaussians.densify_by_clone_and_split(self.gs_params.densify_grad_threshold, extent=4)
                      #self.renderer.gaussians.densify_by_compatness(self.gs_params.K)
                      #self.renderer.gaussians.prune(min_opacity=0.01, extent=4, max_screen_size=1, max_offset=0.1)
                  
                  if step % self.gs_params.opacity_reset_interval == 0:
                      self.renderer.gaussians.reset_opacity()
                      
              comfy_pbar.update_absolute(step + 1)
  
          ender.record()
          torch.cuda.synchronize()
          #t = starter.elapsed_time(ender)
  
          self.need_update = True