content: # Copyright (c) 2023, Tencent Inc
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import numpy as np
  import torch
  import torch.nn as nn
  import nvdiffrast.torch as dr
  from einops import rearrange, repeat
  
  from .encoder.dino_wrapper import DinoWrapper
  from .decoder.transformer import TriplaneTransformer
  from .renderer.synthesizer_mesh import TriplaneSynthesizer
  from .geometry.camera.perspective_camera import PerspectiveCamera
  from .geometry.render.neural_render import NeuralRender
  from .geometry.rep_3d.flexicubes_geometry import FlexiCubesGeometry
  from ..utils.mesh_util import xatlas_uvmap
  
  
  class InstantMesh(nn.Module):
      """
      Full model of the large reconstruction model.
      """
      def __init__(
          self, 
          encoder_freeze: bool = False, 
          encoder_model_name: str = 'facebook/dino-vitb16', 
          encoder_feat_dim: int = 768,
          transformer_dim: int = 1024, 
          transformer_layers: int = 16, 
          transformer_heads: int = 16,
          triplane_low_res: int = 32, 
          triplane_high_res: int = 64, 
          triplane_dim: int = 80,
          rendering_samples_per_ray: int = 128,
          grid_res: int = 128, 
          grid_scale: float = 2.0,
      ):
          super().__init__()
          
          # attributes
          self.grid_res = grid_res
          self.grid_scale = grid_scale
          self.deformation_multiplier = 4.0
  
          # modules
          self.encoder = DinoWrapper(
              model_name=encoder_model_name,
              freeze=encoder_freeze,
          )
  
          self.transformer = TriplaneTransformer(
              inner_dim=transformer_dim, 
              num_layers=transformer_layers, 
              num_heads=transformer_heads,
              image_feat_dim=encoder_feat_dim,
              triplane_low_res=triplane_low_res, 
              triplane_high_res=triplane_high_res, 
              triplane_dim=triplane_dim,
          )
          
          self.synthesizer = TriplaneSynthesizer(
              triplane_dim=triplane_dim, 
              samples_per_ray=rendering_samples_per_ray,
          )
  
      def init_flexicubes_geometry(self, device, fovy=50.0):
          camera = PerspectiveCamera(fovy=fovy, device=device)
          renderer = NeuralRender(device, camera_model=camera)
          self.geometry = FlexiCubesGeometry(
              grid_res=self.grid_res, 
              scale=self.grid_scale, 
              renderer=renderer, 
              render_type='neural_render',
              device=device,
          )
  
      def forward_planes(self, images, cameras):
          # images: [B, V, C_img, H_img, W_img]
          # cameras: [B, V, 16]
          B = images.shape[0]
  
          # encode images
          image_feats = self.encoder(images, cameras)
          image_feats = rearrange(image_feats, '(b v) l d -> b (v l) d', b=B)
          
          # decode triplanes
          planes = self.transformer(image_feats)
  
          return planes
      
      def get_sdf_deformation_prediction(self, planes):
          '''
          Predict SDF and deformation for tetrahedron vertices
          :param planes: triplane feature map for the geometry
          '''
          init_position = self.geometry.verts.unsqueeze(0).expand(planes.shape[0], -1, -1)
          
          # Step 1: predict the SDF and deformation
          sdf, deformation, weight = torch.utils.checkpoint.checkpoint(
              self.synthesizer.get_geometry_prediction,
              planes, 
              init_position, 
              self.geometry.indices,
              use_reentrant=False,
          )
  
          # Step 2: Normalize the deformation to avoid the flipped triangles.
          deformation = 1.0 / (self.grid_res * self.deformation_multiplier) * torch.tanh(deformation)
          sdf_reg_loss = torch.zeros(sdf.shape[0], device=sdf.device, dtype=torch.float32)
  
          ####
          # Step 3: Fix some sdf if we observe empty shape (full positive or full negative)
          sdf_bxnxnxn = sdf.reshape((sdf.shape[0], self.grid_res + 1, self.grid_res + 1, self.grid_res + 1))
          sdf_less_boundary = sdf_bxnxnxn[:, 1:-1, 1:-1, 1:-1].reshape(sdf.shape[0], -1)
          pos_shape = torch.sum((sdf_less_boundary > 0).int(), dim=-1)
          neg_shape = torch.sum((sdf_less_boundary < 0).int(), dim=-1)
          zero_surface = torch.bitwise_or(pos_shape == 0, neg_shape == 0)
          if torch.sum(zero_surface).item() > 0:
              update_sdf = torch.zeros_like(sdf[0:1])
              max_sdf = sdf.max()
              min_sdf = sdf.min()
              update_sdf[:, self.geometry.center_indices] += (1.0 - min_sdf)  # greater than zero
              update_sdf[:, self.geometry.boundary_indices] += (-1 - max_sdf)  # smaller than zero
              new_sdf = torch.zeros_like(sdf)
              for i_batch in range(zero_surface.shape[0]):
                  if zero_surface[i_batch]:
                      new_sdf[i_batch:i_batch + 1] += update_sdf
              update_mask = (new_sdf == 0).float()
              # Regulraization here is used to push the sdf to be a different sign (make it not fully positive or fully negative)
              sdf_reg_loss = torch.abs(sdf).mean(dim=-1).mean(dim=-1)
              sdf_reg_loss = sdf_reg_loss * zero_surface.float()
              sdf = sdf * update_mask + new_sdf * (1 - update_mask)
  
          # Step 4: Here we remove the gradient for the bad sdf (full positive or full negative)
          final_sdf = []
          final_def = []
          for i_batch in range(zero_surface.shape[0]):
              if zero_surface[i_batch]:
                  final_sdf.append(sdf[i_batch: i_batch + 1].detach())
                  final_def.append(deformation[i_batch: i_batch + 1].detach())
              else:
                  final_sdf.append(sdf[i_batch: i_batch + 1])
                  final_def.append(deformation[i_batch: i_batch + 1])
          sdf = torch.cat(final_sdf, dim=0)
          deformation = torch.cat(final_def, dim=0)
          return sdf, deformation, sdf_reg_loss, weight
      
      def get_geometry_prediction(self, planes=None):
          '''
          Function to generate mesh with give triplanes
          :param planes: triplane features
          '''
          # Step 1: first get the sdf and deformation value for each vertices in the tetrahedon grid.
          sdf, deformation, sdf_reg_loss, weight = self.get_sdf_deformation_prediction(planes)
          v_deformed = self.geometry.verts.unsqueeze(dim=0).expand(sdf.shape[0], -1, -1) + deformation
          tets = self.geometry.indices
          n_batch = planes.shape[0]
          v_list = []
          f_list = []
          flexicubes_surface_reg_list = []
          
          # Step 2: Using marching tet to obtain the mesh
          for i_batch in range(n_batch):
              verts, faces, flexicubes_surface_reg = self.geometry.get_mesh(
                  v_deformed[i_batch], 
                  sdf[i_batch].squeeze(dim=-1),
                  with_uv=False, 
                  indices=tets, 
                  weight_n=weight[i_batch].squeeze(dim=-1),
                  is_training=self.training,
              )
              flexicubes_surface_reg_list.append(flexicubes_surface_reg)
              v_list.append(verts)
              f_list.append(faces)
          
          flexicubes_surface_reg = torch.cat(flexicubes_surface_reg_list).mean()
          flexicubes_weight_reg = (weight ** 2).mean()
          
          return v_list, f_list, sdf, deformation, v_deformed, (sdf_reg_loss, flexicubes_surface_reg, flexicubes_weight_reg)
      
      def get_texture_prediction(self, planes, tex_pos, hard_mask=None):
          '''
          Predict Texture given triplanes
          :param planes: the triplane feature map
          :param tex_pos: Position we want to query the texture field
          :param hard_mask: 2D silhoueete of the rendered image
          '''
          tex_pos = torch.cat(tex_pos, dim=0)
          if not hard_mask is None:
              tex_pos = tex_pos * hard_mask.float()
          batch_size = tex_pos.shape[0]
          tex_pos = tex_pos.reshape(batch_size, -1, 3)
          ###################
          # We use mask to get the texture location (to save the memory)
          if hard_mask is not None:
              n_point_list = torch.sum(hard_mask.long().reshape(hard_mask.shape[0], -1), dim=-1)
              sample_tex_pose_list = []
              max_point = n_point_list.max()
              expanded_hard_mask = hard_mask.reshape(batch_size, -1, 1).expand(-1, -1, 3) > 0.5
              for i in range(tex_pos.shape[0]):
                  tex_pos_one_shape = tex_pos[i][expanded_hard_mask[i]].reshape(1, -1, 3)
                  if tex_pos_one_shape.shape[1] < max_point:
                      tex_pos_one_shape = torch.cat(
                          [tex_pos_one_shape, torch.zeros(
                              1, max_point - tex_pos_one_shape.shape[1], 3,
                              device=tex_pos_one_shape.device, dtype=torch.float32)], dim=1)
                  sample_tex_pose_list.append(tex_pos_one_shape)
              tex_pos = torch.cat(sample_tex_pose_list, dim=0)
  
          tex_feat = torch.utils.checkpoint.checkpoint(
              self.synthesizer.get_texture_prediction,
              planes, 
              tex_pos,
              use_reentrant=False,
          )
  
          if hard_mask is not None:
              final_tex_feat = torch.zeros(
                  planes.shape[0], hard_mask.shape[1] * hard_mask.shape[2], tex_feat.shape[-1], device=tex_feat.device)
              expanded_hard_mask = hard_mask.reshape(hard_mask.shape[0], -1, 1).expand(-1, -1, final_tex_feat.shape[-1]) > 0.5
              for i in range(planes.shape[0]):
                  final_tex_feat[i][expanded_hard_mask[i]] = tex_feat[i][:n_point_list[i]].reshape(-1)
              tex_feat = final_tex_feat
  
          return tex_feat.reshape(planes.shape[0], hard_mask.shape[1], hard_mask.shape[2], tex_feat.shape[-1])
      
      def render_mesh(self, mesh_v, mesh_f, cam_mv, render_size=256):
          '''
          Function to render a generated mesh with nvdiffrast
          :param mesh_v: List of vertices for the mesh
          :param mesh_f: List of faces for the mesh
          :param cam_mv:  4x4 rotation matrix
          :return:
          '''
          return_value_list = []
          for i_mesh in range(len(mesh_v)):
              return_value = self.geometry.render_mesh(
                  mesh_v[i_mesh],
                  mesh_f[i_mesh].int(),
                  cam_mv[i_mesh],
                  resolution=render_size,
                  hierarchical_mask=False
              )
              return_value_list.append(return_value)
  
          return_keys = return_value_list[0].keys()
          return_value = dict()
          for k in return_keys:
              value = [v[k] for v in return_value_list]
              return_value[k] = value
  
          mask = torch.cat(return_value['mask'], dim=0)
          hard_mask = torch.cat(return_value['hard_mask'], dim=0)
          tex_pos = return_value['tex_pos']
          depth = torch.cat(return_value['depth'], dim=0)
          normal = torch.cat(return_value['normal'], dim=0)
          return mask, hard_mask, tex_pos, depth, normal
      
      def forward_geometry(self, planes, render_cameras, render_size=256):
          '''
          Main function of our Generator. It first generate 3D mesh, then render it into 2D image
          with given `render_cameras`.
          :param planes: triplane features
          :param render_cameras: cameras to render generated 3D shape
          '''
          B, NV = render_cameras.shape[:2]
  
          # Generate 3D mesh first
          mesh_v, mesh_f, sdf, deformation, v_deformed, sdf_reg_loss = self.get_geometry_prediction(planes)
  
          # Render the mesh into 2D image (get 3d position of each image plane)
          cam_mv = render_cameras
          run_n_view = cam_mv.shape[1]
          antilias_mask, hard_mask, tex_pos, depth, normal = self.render_mesh(mesh_v, mesh_f, cam_mv, render_size=render_size)
  
          tex_hard_mask = hard_mask
          tex_pos = [torch.cat([pos[i_view:i_view + 1] for i_view in range(run_n_view)], dim=2) for pos in tex_pos]
          tex_hard_mask = torch.cat(
              [torch.cat(
                  [tex_hard_mask[i * run_n_view + i_view: i * run_n_view + i_view + 1]
                   for i_view in range(run_n_view)], dim=2)
                  for i in range(planes.shape[0])], dim=0)
  
          # Querying the texture field to predict the texture feature for each pixel on the image
          tex_feat = self.get_texture_prediction(planes, tex_pos, tex_hard_mask)
          background_feature = torch.ones_like(tex_feat)      # white background
  
          # Merge them together
          img_feat = tex_feat * tex_hard_mask + background_feature * (1 - tex_hard_mask)
  
          # We should split it back to the original image shape
          img_feat = torch.cat(
              [torch.cat(
                  [img_feat[i:i + 1, :, render_size * i_view: render_size * (i_view + 1)]
                   for i_view in range(run_n_view)], dim=0) for i in range(len(tex_pos))], dim=0)
  
          img = img_feat.clamp(0, 1).permute(0, 3, 1, 2).unflatten(0, (B, NV))
          antilias_mask = antilias_mask.permute(0, 3, 1, 2).unflatten(0, (B, NV))
          depth = -depth.permute(0, 3, 1, 2).unflatten(0, (B, NV))        # transform negative depth to positive
          normal = normal.permute(0, 3, 1, 2).unflatten(0, (B, NV))
  
          out = {
              'img': img,
              'mask': antilias_mask,
              'depth': depth,
              'normal': normal,
              'sdf': sdf,
              'mesh_v': mesh_v,
              'mesh_f': mesh_f,
              'sdf_reg_loss': sdf_reg_loss,
          }
          return out
  
      def forward(self, images, cameras, render_cameras, render_size: int):
          # images: [B, V, C_img, H_img, W_img]
          # cameras: [B, V, 16]
          # render_cameras: [B, M, D_cam_render]
          # render_size: int
          B, M = render_cameras.shape[:2]
  
          planes = self.forward_planes(images, cameras)
          out = self.forward_geometry(planes, render_cameras, render_size=render_size)
  
          return {
              'planes': planes,
              **out
          }
  
      def extract_mesh(
          self, 
          planes: torch.Tensor, 
          use_texture_map: bool = False,
          texture_resolution: int = 1024,
          **kwargs,
      ):
          '''
          Extract a 3D mesh from FlexiCubes. Only support batch_size 1.
          :param planes: triplane features
          :param use_texture_map: use texture map or vertex color
          :param texture_resolution: the resolution of texure map
          '''
          assert planes.shape[0] == 1
          device = planes.device
  
          # predict geometry first
          mesh_v, mesh_f, sdf, deformation, v_deformed, sdf_reg_loss = self.get_geometry_prediction(planes)
          vertices, faces = mesh_v[0], mesh_f[0]
  
          if not use_texture_map:
              # query vertex colors
              vertices_tensor = vertices.unsqueeze(0)
              vertices_colors = self.synthesizer.get_texture_prediction(
                  planes, vertices_tensor).clamp(0, 1).squeeze(0).cpu().numpy()
              vertices_colors = (vertices_colors * 255).astype(np.uint8)
  
              return vertices.cpu().numpy(), faces.cpu().numpy(), vertices_colors
  
          # use x-atlas to get uv mapping for the mesh
          ctx = dr.RasterizeCudaContext(device=device)
          uvs, mesh_tex_idx, gb_pos, tex_hard_mask = xatlas_uvmap(
              self.geometry.renderer.ctx, vertices, faces, resolution=texture_resolution)
          tex_hard_mask = tex_hard_mask.float()
  
          # query the texture field to get the RGB color for texture map
          tex_feat = self.get_texture_prediction(
              planes, [gb_pos], tex_hard_mask)
          background_feature = torch.zeros_like(tex_feat)
          img_feat = torch.lerp(background_feature, tex_feat, tex_hard_mask)
          texture_map = img_feat.permute(0, 3, 1, 2).squeeze(0)
  
          return vertices, faces, uvs, mesh_tex_idx, texture_map