content: import tqdm
  import random
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import nerfacc
  
  import comfy.utils
  
  from pytorch_msssim import SSIM, MS_SSIM
  
  from kiui.op import safe_normalize
  from kiui.cam import orbit_camera
  from kiui.nn import MLP, trunc_exp
  
  from shared_utils.image_utils import prepare_torch_img
  
  class InstantNGP(nn.Module):
      def __init__(self, resolution=128, device="cuda"):
          super().__init__()
          from kiui.gridencoder import GridEncoder
          
          self.device = torch.device(device)
          self.ref_size_H = resolution
          self.ref_size_W = resolution
          
          self.render_step_size = 5e-3
          self.aabb = torch.tensor([-1.0, -1.0, -1.0, 1.0, 1.0, 1.0], device=self.device)
          self.estimator = nerfacc.OccGridEstimator(roi_aabb=self.aabb, resolution=64, levels=1)
  
          self.encoder_density = GridEncoder(num_levels=12) # VMEncoder(output_dim=16, mode='sum')
          self.encoder = GridEncoder(num_levels=12)
          self.mlp_density = MLP(self.encoder_density.output_dim, 1, 32, 2, bias=False)
          self.mlp = MLP(self.encoder.output_dim, 3, 32, 2, bias=False)
          
      def get_rays(self, pose, h, w, fovy, opengl=True):
  
          x, y = torch.meshgrid(
              torch.arange(w, device=pose.device),
              torch.arange(h, device=pose.device),
              indexing="xy",
          )
          x = x.flatten()
          y = y.flatten()
  
          cx = w * 0.5
          cy = h * 0.5
  
          focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
  
          camera_dirs = F.pad(
              torch.stack(
                  [
                      (x - cx + 0.5) / focal,
                      (y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
                  ],
                  dim=-1,
              ),
              (0, 1),
              value=(-1.0 if opengl else 1.0),
          )  # [hw, 3]
  
          rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1)  # [hw, 3]
          rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
  
          rays_o = rays_o.view(h, w, 3)
          rays_d = safe_normalize(rays_d).view(h, w, 3)
  
          return rays_o, rays_d
      
      def get_color(self, xs):
          return torch.sigmoid(self.mlp(self.encoder(xs.to(self.device))))
      
      def get_density(self, xs):
          # xs: [..., 3]
          xs = xs.to(self.device)
          prefix = xs.shape[:-1]
          xs = xs.view(-1, 3)
          feats = self.encoder_density(xs)
          density = trunc_exp(self.mlp_density(feats))
          density = density.view(*prefix, 1)
          return density
      
      def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
          self.ref_imgs_num = len(reference_images)
  
          self.all_ref_cam_poses = reference_orbit_camera_poses
          self.ref_cam_fovy = reference_orbit_camera_fovy
          
          # 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) # [N, 3, H, W]
          self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0).squeeze(1) # [N, H, W]
          
      def render_nerf(self, pose, bg_color=1):
          
          pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
          
          # get rays
          rays_o, rays_d = self.get_rays(pose, self.ref_size_H, self.ref_size_W, self.ref_cam_fovy)
          hw = rays_o.shape[0] * rays_o.shape[1]
          rays_o = rays_o.view(hw, 3)
          rays_d = rays_d.view(hw, 3)
          
          # update occ grid
          if self.training:
              def occ_eval_fn(xs):
                  sigmas = self.get_density(xs)
                  return self.render_step_size * sigmas
              
              self.estimator.update_every_n_steps(self.render_step, occ_eval_fn=occ_eval_fn, occ_thre=0.01, n=8)
              self.render_step += 1
  
          # render
          def sigma_fn(t_starts, t_ends, ray_indices):
              t_origins = rays_o[ray_indices]
              t_dirs = rays_d[ray_indices]
              xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
              sigmas = self.get_density(xs)
              return sigmas.squeeze(-1)
  
          with torch.no_grad():
              ray_indices, t_starts, t_ends = self.estimator.sampling(
                  rays_o,
                  rays_d,
                  sigma_fn=sigma_fn,
                  near_plane=0.01,
                  far_plane=100,
                  render_step_size=self.render_step_size,
                  stratified=self.training,
                  cone_angle=0,
              )
  
          t_origins = rays_o[ray_indices]
          t_dirs = rays_d[ray_indices]
          xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
          sigmas = self.get_density(xs).squeeze(-1)
          rgbs = torch.sigmoid(self.mlp(self.encoder(xs)))
  
          n_rays=rays_o.shape[0]
          weights, t, alphas = nerfacc.render_weight_from_density(t_starts, t_ends, sigmas, ray_indices=ray_indices, n_rays=n_rays)
          color = nerfacc.accumulate_along_rays(weights, values=rgbs, ray_indices=ray_indices, n_rays=n_rays)
          alpha = nerfacc.accumulate_along_rays(weights, values=None, ray_indices=ray_indices, n_rays=n_rays)
  
          color = color + (1.0 - alpha) * bg_color
  
          color = color.view(self.ref_size_H, self.ref_size_W, 3).clamp(0, 1).permute(2, 0, 1).contiguous()
          alpha = alpha.view(self.ref_size_H, self.ref_size_W).clamp(0, 1).contiguous()
          
          return color, alpha
  
      def fit_nerf(self, iters=512, bg_color=1):
  
          optimizer = torch.optim.Adam([
              {'params': self.encoder_density.parameters(), 'lr': 1e-2},
              {'params': self.encoder.parameters(), 'lr': 1e-2},
              {'params': self.mlp_density.parameters(), 'lr': 1e-3},
              {'params': self.mlp.parameters(), 'lr': 1e-3},
          ])
  
          print(f"[INFO] fitting nerf...")
          self.render_step = 0
          
          ref_imgs_num_minus_1 = self.ref_imgs_num-1
          
          comfy_pbar = comfy.utils.ProgressBar(iters)
          pbar = tqdm.trange(iters)
          for step in pbar:
              
              i = random.randint(0, ref_imgs_num_minus_1)
  
              radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
              
              orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
              pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
              
              image_gt = self.ref_imgs_torch[i]   # [3, H, W]
              alpha_gt = self.ref_masks_torch[i]  # [H, W]
              image_pred, alpha_pred = self.render_nerf(pose, bg_color)
  
              # if i % 200 == 0:
              #     kiui.vis.plot_image(image_gt, alpha_gt, image_pred, alpha_pred)
              
              loss_mse = F.mse_loss(image_pred, image_gt) + 0.1 * F.mse_loss(alpha_pred, alpha_gt)
              loss = loss_mse #+ 0.1 * self.encoder_density.tv_loss() #+ 0.0001 * self.encoder_density.density_loss()
              #loss += self.lambda_ssim * (1 - self.ms_ssim_loss(image_gt, image_pred))
  
              loss.backward()
              self.encoder_density.grad_total_variation(1e-8)
          
              optimizer.step()
              optimizer.zero_grad()
  
              pbar.set_description(f"NeRF Fitting Loss = {loss_mse.item():.6f}")
              comfy_pbar.update_absolute(step + 1)
              
          torch.cuda.synchronize()
          
          print(f"[INFO] finished fitting nerf!")