content: import io
  import math
  from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from dataclasses import dataclass
  from datetime import datetime
  from typing import List, Optional, Union
  
  import numpy as np
  import nvdiffrast.torch as dr
  import torch
  import torch.nn.functional as F
  import trimesh
  from PIL import Image
  from torch import BoolTensor, FloatTensor
  
  from .utils import tensor_to_image
  
  
  def dot(x: torch.FloatTensor, y: torch.FloatTensor) -> torch.FloatTensor:
      return torch.sum(x * y, -1, keepdim=True)
  
  
  @dataclass
  class TexturedMesh:
      v_pos: torch.FloatTensor
      t_pos_idx: torch.LongTensor
  
      # texture coordinates
      v_tex: Optional[torch.FloatTensor] = None
      t_tex_idx: Optional[torch.LongTensor] = None
  
      # texture map
      texture: Optional[torch.FloatTensor] = None
  
      # vertices, faces after vertex merging
      _stitched_v_pos: Optional[torch.FloatTensor] = None
      _stitched_t_pos_idx: Optional[torch.LongTensor] = None
  
      _v_nrm: Optional[torch.FloatTensor] = None
      _v_tang: Optional[torch.FloatTensor] = None
  
      @property
      def v_nrm(self) -> torch.FloatTensor:
          if self._v_nrm is None:
              self._v_nrm = self._compute_vertex_normal()
          return self._v_nrm
  
      @property
      def v_tang(self) -> torch.FloatTensor:
          if self._v_tang is None:
              self._v_tang = self._compute_tangent()
          return self._v_tang
  
      def set_vertex_normal(self, v_nrm: torch.FloatTensor) -> None:
          assert v_nrm.shape == self.v_pos.shape
          self._v_nrm = v_nrm.to(self.v_pos)
  
      def set_stitched_mesh(
          self, v_pos: torch.FloatTensor, t_pos_idx: torch.LongTensor
      ) -> None:
          self._stitched_v_pos = v_pos
          self._stitched_t_pos_idx = t_pos_idx
  
      @property
      def stitched_v_pos(self) -> torch.FloatTensor:
          if self._stitched_v_pos is None:
              print("Warning: Stitched vertices not available, using original vertices!")
              return self.v_pos
          return self._stitched_v_pos
  
      @property
      def stitched_t_pos_idx(self) -> torch.LongTensor:
          if self._stitched_t_pos_idx is None:
              print("Warning: Stitched faces not available, using original faces!")
              return self.t_pos_idx
          return self._stitched_t_pos_idx
  
      @property
      def uv_size(self) -> Optional[int]:
          if self.texture is None:
              return None
          return self.texture.shape[0]
  
      def _compute_vertex_normal(self) -> torch.FloatTensor:
          if self._stitched_v_pos is None or self._stitched_t_pos_idx is None:
              print(
                  "Warning: Stitched vertices and faces not available, computing vertex normals on original mesh, which can be erroneous!"
              )
              v_pos, t_pos_idx = self.v_pos, self.t_pos_idx
          else:
              v_pos, t_pos_idx = self._stitched_v_pos, self._stitched_t_pos_idx
  
          i0 = t_pos_idx[:, 0]
          i1 = t_pos_idx[:, 1]
          i2 = t_pos_idx[:, 2]
  
          v0 = v_pos[i0, :]
          v1 = v_pos[i1, :]
          v2 = v_pos[i2, :]
  
          face_normals = torch.cross(v1 - v0, v2 - v0)
  
          # Splat face normals to vertices
          v_nrm = torch.zeros_like(v_pos)
          v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
          v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
          v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
  
          # Normalize, replace zero (degenerated) normals with some default value
          v_nrm = torch.where(
              dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
          )
          v_nrm = F.normalize(v_nrm, dim=1)
  
          if torch.is_anomaly_enabled():
              assert torch.all(torch.isfinite(v_nrm))
  
          return v_nrm
  
      def _compute_tangent(self) -> torch.FloatTensor:
          vn_idx = [None] * 3
          pos = [None] * 3
          tex = [None] * 3
          for i in range(0, 3):
              pos[i] = self.v_pos[self.t_pos_idx[:, i]]
              tex[i] = self.v_tex[self.t_tex_idx[:, i]]
              vn_idx[i] = self.t_pos_idx[:, i]
  
          tangents = torch.zeros_like(self._v_nrm)
          tansum = torch.zeros_like(self._v_nrm)
  
          # Compute tangent space for each triangle
          uve1 = tex[1] - tex[0]
          uve2 = tex[2] - tex[0]
          pe1 = pos[1] - pos[0]
          pe2 = pos[2] - pos[0]
  
          nom = pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2]
          denom = uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1]
  
          # Avoid division by zero for degenerated texture coordinates
          tang = nom / torch.where(
              denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6)
          )
  
          # Update all 3 vertices
          for i in range(0, 3):
              idx = vn_idx[i][:, None].repeat(1, 3)
              tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
              tansum.scatter_add_(
                  0, idx, torch.ones_like(tang)
              )  # tansum[n_i] = tansum[n_i] + 1
          tangents = tangents / tansum
  
          # Normalize and make sure tangent is perpendicular to normal
          tangents = F.normalize(tangents, dim=1)
          tangents = F.normalize(
              tangents - dot(tangents, self._v_nrm) * self._v_nrm, dim=1
          )
  
          v_tng = tangents
  
          if torch.is_anomaly_enabled():
              assert torch.all(torch.isfinite(v_tng))
  
          return v_tng
  
      def to(self, device: Optional[str] = None):
          self.v_pos = self.v_pos.to(device)
          self.t_pos_idx = self.t_pos_idx.to(device)
          if self.v_tex is not None:
              self.v_tex = self.v_tex.to(device)
          if self.t_tex_idx is not None:
              self.t_tex_idx = self.t_tex_idx.to(device)
          if self.texture is not None:
              self.texture = self.texture.to(device)
          if self._stitched_v_pos is not None:
              self._stitched_v_pos = self._stitched_v_pos.to(device)
          if self._stitched_t_pos_idx is not None:
              self._stitched_t_pos_idx = self._stitched_t_pos_idx.to(device)
          if self._v_nrm is not None:
              self._v_nrm = self._v_nrm.to(device)
          if self._v_tang is not None:
              self._v_tang = self._v_tang.to(device)
  
  
  @contextmanager
  def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor):
      texture_ = mesh.texture
      mesh.texture = texture
      try:
          yield
      finally:
          mesh.texture = texture_
  
  
  def load_mesh(
      mesh_path: str,
      rescale: bool = False,
      move_to_center: bool = False,
      scale: float = 0.5,
      flip_uv: bool = True,
      merge_vertices: bool = True,
      default_uv_size: Optional[int] = None,
      shape_init_mesh_up: str = "+y",
      shape_init_mesh_front: str = "+x",
      front_x_to_y: bool = False,
      device: Optional[str] = None,
      return_transform: bool = False,
  ) -> TexturedMesh:
      if mesh_path.endswith(".npz"):
  
          class Mesh:
              vertices = None
              faces = None
  
          data = np.load(mesh_path)
          mesh = Mesh()
          mesh.vertices = data["vertices"]
          mesh.faces = data["faces"]
          merge_vertices = False
      else:
          scene = trimesh.load(mesh_path, force="mesh", process=False)
          if isinstance(scene, trimesh.Trimesh):
              mesh = scene
          elif isinstance(scene, trimesh.scene.Scene):
              mesh = trimesh.Trimesh()
              for obj in scene.geometry.values():
                  mesh = trimesh.util.concatenate([mesh, obj])
          else:
              raise ValueError(f"Unknown mesh type at {mesh_path}.")
  
      vertex_normals = getattr(mesh, "vertex_normals", None)
  
      # move to center
      transform_offset = None
      if move_to_center:
          centroid = mesh.vertices.mean(0)
          mesh.vertices = mesh.vertices - centroid
          transform_offset = centroid  # record the offset
  
      # rescale
      transform_scale = None
      if rescale:
          max_scale = np.abs(mesh.vertices).max()
          mesh.vertices = mesh.vertices / max_scale * scale
          transform_scale = max_scale / scale
  
      dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
      dir2vec = {
          "+x": np.array([1, 0, 0]),
          "+y": np.array([0, 1, 0]),
          "+z": np.array([0, 0, 1]),
          "-x": np.array([-1, 0, 0]),
          "-y": np.array([0, -1, 0]),
          "-z": np.array([0, 0, -1]),
      }
      if shape_init_mesh_up not in dirs or shape_init_mesh_front not in dirs:
          raise ValueError(
              f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
          )
      if shape_init_mesh_up[1] == shape_init_mesh_front[1]:
          raise ValueError(
              "shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
          )
      z_, x_ = (
          dir2vec[shape_init_mesh_up],
          dir2vec[shape_init_mesh_front],
      )
      y_ = np.cross(z_, x_)
      std2mesh = np.stack([x_, y_, z_], axis=0).T
      mesh2std = np.linalg.inv(std2mesh)
      mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
      if vertex_normals is not None:
          vertex_normals = np.dot(mesh2std, vertex_normals.T).T
      if front_x_to_y:
          x = mesh.vertices[:, 1].copy()
          y = -mesh.vertices[:, 0].copy()
          mesh.vertices[:, 0] = x
          mesh.vertices[:, 1] = y
          if vertex_normals is not None:
              vx = vertex_normals[:, 1].copy()
              vy = -vertex_normals[:, 0].copy()
              vertex_normals[:, 0] = vx
              vertex_normals[:, 1] = vy
  
      v_pos = torch.tensor(mesh.vertices, dtype=torch.float32)
      t_pos_idx = torch.tensor(mesh.faces, dtype=torch.int64)
  
      if hasattr(mesh, "visual") and hasattr(mesh.visual, "uv"):
          v_tex = torch.tensor(mesh.visual.uv, dtype=torch.float32)
          if flip_uv:
              v_tex[:, 1] = 1.0 - v_tex[:, 1]
          t_tex_idx = t_pos_idx.clone()
          if (
              default_uv_size is not None
              or getattr(mesh.visual.material, "baseColorTexture", None) is None
          ):
              assert default_uv_size is not None
              texture = torch.zeros(
                  (default_uv_size, default_uv_size, 3), dtype=torch.float32
              )
          else:
              texture = torch.tensor(
                  np.array(mesh.visual.material.baseColorTexture) / 255.0,
                  dtype=torch.float32,
              )[..., :3]
      else:
          v_tex = None
          t_tex_idx = None
          texture = None
  
      textured_mesh = TexturedMesh(
          v_pos=v_pos,
          t_pos_idx=t_pos_idx,
          v_tex=v_tex,
          t_tex_idx=t_tex_idx,
          texture=texture,
      )
  
      if vertex_normals is not None:
          v_nrm = F.normalize(torch.tensor(vertex_normals, dtype=torch.float32), dim=-1)
          textured_mesh.set_vertex_normal(v_nrm)
  
      if merge_vertices and vertex_normals is None:
          # only merge vertices when vertex normals are not available
          mesh.merge_vertices(merge_tex=True)
          textured_mesh.set_stitched_mesh(
              torch.tensor(mesh.vertices, dtype=torch.float32),
              torch.tensor(mesh.faces, dtype=torch.int64),
          )
      else:
          textured_mesh.set_stitched_mesh(textured_mesh.v_pos, textured_mesh.t_pos_idx)
  
      textured_mesh.to(device)
  
      if return_transform:
          return textured_mesh, transform_offset, transform_scale
      else:
          return textured_mesh
  
  
  def replace_mesh_texture_and_save_trimesh(
      input_path: str,
      output_path: str,
      texture: Union[Image.Image, np.ndarray, torch.Tensor],
      metallic_roughness_texture: Optional[
          Union[Image.Image, np.ndarray, torch.Tensor]
      ] = None,
      normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
      texture_format: str = "JPEG",
      task_id: str = "",
      **kwargs,
  ) -> None:
      tmesh = trimesh.load(input_path, force="mesh", process=False)
  
      def convert_image(tensor):
          texture = tensor_to_image(tensor)
          texture.format = texture_format
          return texture
  
      texture = convert_image(texture)
      if metallic_roughness_texture is not None:
          metallic_roughness_texture = convert_image(metallic_roughness_texture)
  
      if normal_texture is not None:
          normal_texture = convert_image(normal_texture)
  
      if not isinstance(tmesh.visual.material, trimesh.visual.material.PBRMaterial):
          material = trimesh.visual.texture.PBRMaterial(
              baseColorFactor=[1.0, 1.0, 1.0, 1.0],
              roughnessFactor=0.9 if metallic_roughness_texture is None else 1.0,
              metallicFactor=0.0 if metallic_roughness_texture is None else 1.0,
              baseColorTexture=texture,
              metallicRoughnessTexture=metallic_roughness_texture,
              normalTexture=normal_texture,
          )
          tmesh.visual = trimesh.visual.TextureVisuals(
              uv=tmesh.visual.uv, image=texture, material=material
          )
      else:
          tmesh.visual.material.baseColorTexture = texture
          tmesh.visual.material.baseColorFactor = [1.0, 1.0, 1.0, 1.0]
          if metallic_roughness_texture is None:
              tmesh.visual.material.roughnessFactor = 0.9
              tmesh.visual.material.metallicFactor = 0.0
          else:
              tmesh.visual.material.roughnessFactor = 1.0
              tmesh.visual.material.metallicFactor = 1.0
              tmesh.visual.material.metallicRoughnessTexture = metallic_roughness_texture
              tmesh.visual.material.normalTexture = normal_texture
  
      def p_func(tree, task_id_in=task_id):
          if "nodes" in tree and tree["nodes"][0]["name"].find("world") >= 0:
              tree["nodes"].pop(0)
          tree["asset"]["generator"] = "https://github.com/huanngzh/MV-Adapter"
          change_dict = {
              "nodes": "node",
              "meshes": "mesh",
              "materials": "mat",
              "images": "image",
          }
          for k, v in change_dict.items():
              if k in tree:
                  for index in range(len(tree[k])):
                      tree[k][index]["name"] = f"mvadapter_{v}_{task_id_in}"
  
      tmesh.export(output_path, tree_postprocessor=p_func)
  
  
  # def replace_mesh_texture_and_save_gltflib(
  #     input_path: str,
  #     output_path: str,
  #     texture: Union[Image.Image, np.ndarray, torch.Tensor],
  #     metallic_roughness_texture: Optional[
  #         Union[Image.Image, np.ndarray, torch.Tensor]
  #     ] = None,
  #     normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
  #     texture_format: str = "JPEG",
  #     normal_strength: float = 1.0,
  #     task_id: str = "",
  #     **kwargs,
  # ) -> None:
  #     from gltflib import GLTF
  #     from gltflib import Image as gltfImage
  #     from gltflib import NormalTextureInfo, Texture, TextureInfo
  
  #     def convert_image(tensor):
  #         texture = tensor_to_image(tensor)
  #         texture.format = texture_format
  #         return texture
  
  #     def add_texture(gltf: GLTF, image: Image.Image) -> int:
  #         jpeg_buffer = io.BytesIO()
  #         image.save(jpeg_buffer, format=texture_format)
  #         new_texture_data = bytearray(jpeg_buffer.getvalue())
  #         _, offset, bytelen = gltf._create_or_extend_glb_resource(new_texture_data)
  #         bufferView = gltf._create_embedded_image_buffer_view(offset, bytelen)
  #         image = gltfImage(mimeType="image/jpeg", bufferView=bufferView)
  #         if gltf.model.images:
  #             gltf.model.images.append(image)
  #         else:
  #             gltf.model.images = [image]
  #         texture = Texture(source=len(gltf.model.images) - 1)
  #         if gltf.model.textures:
  #             gltf.model.textures.append(texture)
  #         else:
  #             gltf.model.textures = [texture]
  #         return len(gltf.model.textures) - 1
  
  #     gltf = GLTF.load(input_path)
  
  #     texture = convert_image(texture)
  #     if metallic_roughness_texture is not None:
  #         metallic_roughness_texture = convert_image(metallic_roughness_texture)
  
  #     if normal_texture is not None:
  #         normal_texture = convert_image(normal_texture)
  
  #     pbr_material = gltf.model.materials[0].pbrMetallicRoughness
  
  #     pbr_material.baseColorTexture = TextureInfo(index=add_texture(gltf, texture))
  #     pbr_material.baseColorFactor = [1, 1, 1, 1]
  
  #     if metallic_roughness_texture is not None:
  #         pbr_material.metallicRoughnessTexture = TextureInfo(
  #             index=add_texture(gltf, metallic_roughness_texture)
  #         )
  #         pbr_material.metallicFactor = 1.0
  #         pbr_material.roughnessFactor = 1.0
  #     else:
  #         pbr_material.metallicFactor = 0.0
  #         pbr_material.roughnessFactor = 0.9
  
  #     if normal_texture is not None:
  #         gltf.model.materials[0].normalTexture = NormalTextureInfo(
  #             index=add_texture(gltf, normal_texture), scale=normal_strength
  #         )
  
  #     gltf.model.nodes[0].name = f"mvadapter_node_{task_id}"
  #     gltf.model.meshes[0].name = f"mvadapter_mesh_{task_id}"
  #     gltf.model.materials[0].name = f"mvadapter_material_{task_id}"
  #     gltf.model.asset.generator = "https://github.com/huanngzh/MV-Adapter"
  #     gltf.export(output_path)
  
  
  def replace_mesh_texture_and_save(
      input_path: str,
      output_path: str,
      texture: Union[Image.Image, np.ndarray, torch.Tensor],
      metallic_roughness_texture: Optional[
          Union[Image.Image, np.ndarray, torch.Tensor]
      ] = None,
      normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
      normal_strength: float = 1.0,
      texture_format: str = "JPEG",
      task_id: str = "",
      backend: str = "trimesh",
  ) -> None:
      if backend == "trimesh":
          replace_mesh_texture_and_save_trimesh(
              input_path,
              output_path,
              texture,
              metallic_roughness_texture=metallic_roughness_texture,
              normal_texture=normal_texture,
              texture_format=texture_format,
              task_id=task_id,
          )
      # elif backend == "gltflib":
      #     replace_mesh_texture_and_save_gltflib(
      #         input_path,
      #         output_path,
      #         texture,
      #         metallic_roughness_texture=metallic_roughness_texture,
      #         normal_texture=normal_texture,
      #         normal_strength=normal_strength,
      #         task_id=task_id,
      #     )
      else:
          raise NotImplementedError
  