mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-23 03:50:04 +08:00
106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
import torch
|
|
from skimage import measure
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
|
|
@dataclass
|
|
class Latent2MeshOutput():
|
|
# mesh for vertices and faces
|
|
mesh_v: None
|
|
mesh_f: None
|
|
|
|
class SufraceExtractor():
|
|
def compute_box_stat(self, bounds, octree_resolution: int):
|
|
|
|
# if float, turn it into a cube
|
|
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, grid_logit, *, bounds, octree_res, **kwargs):
|
|
# grid_logit from volume decoder
|
|
# use marching cube algo to turn an sdf to a mesh
|
|
vertices, faces, _, _ = measure.marching_cubes(grid_logit.cpu().numpy(),
|
|
0.0,
|
|
method = "lewiner")
|
|
|
|
grid_size, bbox_min, bbox_size = self.compute_box_stat(bounds = bounds, octree_resolution = octree_res)
|
|
vertices = vertices / grid_size * bbox_size + bbox_min
|
|
|
|
return vertices, faces
|
|
|
|
def __call__(self, grid_logits, **kwds):
|
|
|
|
outputs = []
|
|
# loop over the batches
|
|
for i in range(grid_logits.shape[0]):
|
|
try:
|
|
# process each batch
|
|
vertices, faces = self.run(grid_logits[i], **kwds)
|
|
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
|
|
|
|
################################################
|
|
# Volume Decoder
|
|
################################################
|
|
|
|
class VanillaVolumeDecoder():
|
|
@torch.no_grad()
|
|
def __call__(self, latents: torch.Tensor, geo_decoder: callable, octree_res: int, bounds = 1.01,
|
|
num_chunks: int = 10_000):
|
|
|
|
if isinstance(bounds, float):
|
|
bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
|
|
|
|
bbox_min, bbox_max = torch.tensor(bounds[:3]), torch.tensor(bounds[3:])
|
|
|
|
x = torch.linspace(bbox_min[0], bbox_max[0], int(octree_res) + 1, dtype = torch.float32)
|
|
y = torch.linspace(bbox_min[1], bbox_max[1], int(octree_res) + 1, dtype = torch.float32)
|
|
z = torch.linspace(bbox_min[2], bbox_max[2], int(octree_res) + 1, dtype = torch.float32)
|
|
|
|
[xs, ys, zs] = torch.meshgrid(x, y, z, indexing = "ij")
|
|
xyz = torch.stack((xs, ys, zs), axis=-1).to(latents.device, dtype = latents.dtype).contiguous().reshape(-1, 3)
|
|
grid_size = [int(octree_res) + 1, int(octree_res) + 1, int(octree_res) + 1]
|
|
|
|
batch_logits = []
|
|
for start in range(0, xyz.shape[0], num_chunks):
|
|
chunk_queries = xyz[start: start + num_chunks, :]
|
|
chunk_queries = chunk_queries.unsqueeze(0).repeat(latents.shape[0], 1, 1)
|
|
logits = geo_decoder(queries = chunk_queries, latents = latents)
|
|
batch_logits.append(logits)
|
|
|
|
grid_logits = torch.cat(batch_logits, dim = 1)
|
|
grid_logits = grid_logits.view((latents.shape[0], *grid_size)).float()
|
|
|
|
return grid_logits
|
|
|
|
def export_to_trimesh(mesh_output):
|
|
import trimesh
|
|
|
|
if isinstance(mesh_output, list):
|
|
outputs = []
|
|
for mesh in mesh_output:
|
|
if mesh is None:
|
|
outputs.append(None)
|
|
else:
|
|
mesh.mesh_f = mesh.mesh_f[:, ::-1]
|
|
mesh_output = trimesh.Trimesh(mesh.mesh_v, mesh.mesh_f)
|
|
outputs.append(mesh_output)
|
|
return outputs
|
|
else:
|
|
mesh_output.mesh_f = mesh_output.mesh_f[:, ::-1]
|
|
mesh_output = trimesh.Trimesh(mesh_output.mesh_v, mesh_output.mesh_f)
|
|
return mesh_output
|
|
|