Spaces:
Running
on
Zero
Running
on
Zero
| # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT | |
| # except for the third-party components listed below. | |
| # Hunyuan 3D does not impose any additional limitations beyond what is outlined | |
| # in the repsective licenses of these third-party components. | |
| # Users must comply with all terms and conditions of original licenses of these third-party | |
| # components and must ensure that the usage of the third party components adheres to | |
| # all relevant laws and regulations. | |
| # For avoidance of doubts, Hunyuan 3D means the large language models and | |
| # their software and algorithms, including trained model weights, parameters (including | |
| # optimizer states), machine-learning model code, inference-enabling code, training-enabling code, | |
| # fine-tuning enabling code and other elements of the foregoing made publicly available | |
| # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT. | |
| from typing import Union, Tuple, List | |
| import numpy as np | |
| import torch | |
| from skimage import measure | |
| class Latent2MeshOutput: | |
| def __init__(self, mesh_v=None, mesh_f=None): | |
| self.mesh_v = mesh_v | |
| self.mesh_f = mesh_f | |
| def center_vertices(vertices): | |
| """Translate the vertices so that bounding box is centered at zero.""" | |
| vert_min = vertices.min(dim=0)[0] | |
| vert_max = vertices.max(dim=0)[0] | |
| vert_center = 0.5 * (vert_min + vert_max) | |
| return vertices - vert_center | |
| class SurfaceExtractor: | |
| def _compute_box_stat(self, bounds: Union[Tuple[float], List[float], float], octree_resolution: int): | |
| if isinstance(bounds, float): | |
| bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] | |
| bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6]) | |
| bbox_size = bbox_max - bbox_min | |
| grid_size = [int(octree_resolution) + 1, int(octree_resolution) + 1, int(octree_resolution) + 1] | |
| return grid_size, bbox_min, bbox_size | |
| def run(self, *args, **kwargs): | |
| return NotImplementedError | |
| def __call__(self, grid_logits, **kwargs): | |
| outputs = [] | |
| for i in range(grid_logits.shape[0]): | |
| try: | |
| vertices, faces = self.run(grid_logits[i], **kwargs) | |
| vertices = vertices.astype(np.float32) | |
| faces = np.ascontiguousarray(faces) | |
| outputs.append(Latent2MeshOutput(mesh_v=vertices, mesh_f=faces)) | |
| except Exception: | |
| import traceback | |
| traceback.print_exc() | |
| outputs.append(None) | |
| return outputs | |
| class MCSurfaceExtractor(SurfaceExtractor): | |
| def run(self, grid_logit, *, mc_level, bounds, octree_resolution, **kwargs): | |
| vertices, faces, normals, _ = measure.marching_cubes( | |
| grid_logit.cpu().numpy(), | |
| mc_level, | |
| method="lewiner" | |
| ) | |
| grid_size, bbox_min, bbox_size = self._compute_box_stat(bounds, octree_resolution) | |
| vertices = vertices / grid_size * bbox_size + bbox_min | |
| return vertices, faces | |
| class DMCSurfaceExtractor(SurfaceExtractor): | |
| def run(self, grid_logit, *, octree_resolution, **kwargs): | |
| device = grid_logit.device | |
| if not hasattr(self, 'dmc'): | |
| try: | |
| from diso import DiffDMC | |
| except: | |
| raise ImportError("Please install diso via `pip install diso`, or set mc_algo to 'mc'") | |
| self.dmc = DiffDMC(dtype=torch.float32).to(device) | |
| sdf = -grid_logit / octree_resolution | |
| sdf = sdf.to(torch.float32).contiguous() | |
| verts, faces = self.dmc(sdf, deform=None, return_quads=False, normalize=True) | |
| verts = center_vertices(verts) | |
| vertices = verts.detach().cpu().numpy() | |
| faces = faces.detach().cpu().numpy()[:, ::-1] | |
| return vertices, faces | |
| SurfaceExtractors = { | |
| 'mc': MCSurfaceExtractor, | |
| 'dmc': DMCSurfaceExtractor, | |
| } | |