From 174655006ce43e07d9e4ff158e521479229e90b7 Mon Sep 17 00:00:00 2001 From: Yousef Rafat <81116377+yousef-rafat@users.noreply.github.com> Date: Tue, 8 Jul 2025 00:51:21 +0300 Subject: [PATCH] fixed some bugs and rewrote OpenCV resize funcs --- comfy/ldm/hunyuan3d/model_/conditioner.py | 69 +- comfy/ldm/hunyuan3d/model_/dinov2.py | 6 +- comfy/ldm/hunyuan3d/model_/hunyuandit.py | 4 +- comfy/ldm/hunyuan3d/model_/image_processor.py | 237 ++-- comfy/ldm/hunyuan3d/model_/pipeline.py | 53 +- comfy/ldm/hunyuan3d/model_/scheduler.py | 2 +- comfy/ldm/hunyuan3d/model_/vae.py | 1010 +++++++++++++++++ comfy/ldm/hunyuan3d/vae/vae.py | 2 +- 8 files changed, 1280 insertions(+), 103 deletions(-) create mode 100644 comfy/ldm/hunyuan3d/model_/vae.py diff --git a/comfy/ldm/hunyuan3d/model_/conditioner.py b/comfy/ldm/hunyuan3d/model_/conditioner.py index d56ffd34b..93908ecb6 100644 --- a/comfy/ldm/hunyuan3d/model_/conditioner.py +++ b/comfy/ldm/hunyuan3d/model_/conditioner.py @@ -1,13 +1,18 @@ import torch import torch.nn as nn import torch.nn.functional as F -from dinov2 import Dinov2Model, DinoConfig +from dinov2 import DinoConfig, Dinov2Model # avoid using torchvision by recreating image processing functions def resize(img: torch.Tensor, size: int) -> torch.Tensor: + + batched = img.ndim == 4 - _, h, w = img.shape + if not batched: + img = img.unsqueeze(0) + + _, _, h, w = img.shape # mantain aspect ratio if h < w: @@ -17,18 +22,31 @@ def resize(img: torch.Tensor, size: int) -> torch.Tensor: new_w = size new_h = int(h * size / w) - img = img.unsqueeze(0) img = F.interpolate(img, size = (new_h, new_w), mode = 'bilinear', align_corners = False, antialias = True ) - return img.squeeze(0) + + if not batched: + img = img.squeeze(0) + + return img def center_crop(img: torch.Tensor, size: int) -> torch.Tensor: - _, h, w = img.shape + batched = img.ndim == 4 + if not batched: + img = img.unsqueeze(0) + + _, _, h, w = img.shape top = (h - size) // 2 left = (w - size) // 2 - return img[:, top:top + size, left:left + size] + + cropped = img[..., top:top + size, left:left + size] + + if not batched: + cropped = cropped.squeeze(0) + + return cropped def normalize(img: torch.Tensor, mean: list, std: list) -> torch.Tensor: @@ -47,8 +65,8 @@ class ImageEncoder(nn.Module): def __init__( self, config: DinoConfig, - use_cls_token=True, - image_size=224, + use_cls_token = True, + image_size = 518, **kwargs, ): super().__init__() @@ -77,15 +95,16 @@ class ImageEncoder(nn.Module): def forward(self, image, value_range=(-1, 1), **kwargs): + if image.ndim == 3: + image = image.unsqueeze(0) + if value_range is not None: low, high = value_range image = (image - low) / (high - low) - image = image.to(self.model.device, dtype=self.model.dtype) inputs = self.transform(image) - outputs = self.model(inputs) - - last_hidden_state = outputs.last_hidden_state + inputs = inputs.to(self.model.device, dtype=self.model.dtype) + last_hidden_state = self.model(inputs) if not self.use_cls_token: last_hidden_state = last_hidden_state[:, 1:, :] @@ -101,16 +120,16 @@ class ImageEncoder(nn.Module): batch_size, self.num_patches, self.model.config.hidden_size, - device=device, - dtype=dtype, + device = device, + dtype = dtype, ) return zero class SingleImageEncoder(nn.Module): - def __init__(self): + def __init__(self, config): super().__init__() - self.main_image_encoder = ImageEncoder() + self.main_image_encoder = ImageEncoder(config) def forward(self, image, **kwargs): outputs = { @@ -124,22 +143,22 @@ class SingleImageEncoder(nn.Module): } return outputs -def load_dino2(dino2: Dinov2Model): - - checkpoint = "" - dino2.load_state_dict(torch.load(checkpoint)) - return dino2 - def test_image_encoder(): torch.manual_seed(2025) - image_encoder = SingleImageEncoder() + config = DinoConfig() + image_encoder = SingleImageEncoder(config) - image = torch.rand(1, 3, 224, 224) + image = torch.rand(3, 224, 224) outputs = image_encoder(image) print(outputs) if __name__ == "__main__": - test_image_encoder() \ No newline at end of file + #test_image_encoder() + conditioner = SingleImageEncoder(DinoConfig()) + torch.manual_seed(2025) + image = torch.rand(1, 3, 224, 224) + outputs = conditioner(image) + print(outputs["main"].size()) \ No newline at end of file diff --git a/comfy/ldm/hunyuan3d/model_/dinov2.py b/comfy/ldm/hunyuan3d/model_/dinov2.py index 7ff44d466..79f813a1a 100644 --- a/comfy/ldm/hunyuan3d/model_/dinov2.py +++ b/comfy/ldm/hunyuan3d/model_/dinov2.py @@ -21,6 +21,8 @@ class DinoConfig(): qkv_bias: bool = True layerscale_value: float = 1.0 drop_path_rate: float = 0.0 + device: str = "cuda" + dtype = torch.float16 class Dinov2Embeddings(nn.Module): """ @@ -372,6 +374,8 @@ class Dinov2Model(nn.Module): self.embeddings = Dinov2Embeddings(config) self.encoder = Dinov2Encoder(config) + self.device = config.device + self.dtype = config.dtype self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) @@ -391,7 +395,7 @@ class Dinov2Model(nn.Module): embedding_output, head_mask = head_mask, ) - sequence_output = encoder_outputs[0] + sequence_output = encoder_outputs sequence_output = self.layernorm(sequence_output) return sequence_output \ No newline at end of file diff --git a/comfy/ldm/hunyuan3d/model_/hunyuandit.py b/comfy/ldm/hunyuan3d/model_/hunyuandit.py index d0894e46a..d45d5c275 100644 --- a/comfy/ldm/hunyuan3d/model_/hunyuandit.py +++ b/comfy/ldm/hunyuan3d/model_/hunyuandit.py @@ -31,7 +31,7 @@ class Timesteps(nn.Module): def forward(self, timesteps: torch.Tensor): - x = timesteps.float().unsqueeze(1) * self.inv_freq.unsqueeze(0) + x = timesteps.float().unsqueeze(1) * self.inv_freq.to(timesteps.device).unsqueeze(0) # scale factor if self.scale != 1.0: @@ -73,7 +73,7 @@ class TimestepEmbedder(nn.Module): cond_embed = self.cond_proj(condition) timestep_embed = timestep_embed + cond_embed - time_conditioned = self.mlp(timestep_embed) + time_conditioned = self.mlp(timestep_embed.to(self.mlp[0].weight.device)) # for broadcasting with image tokens return time_conditioned.unsqueeze(1) diff --git a/comfy/ldm/hunyuan3d/model_/image_processor.py b/comfy/ldm/hunyuan3d/model_/image_processor.py index f98d895b8..dfc0934ac 100644 --- a/comfy/ldm/hunyuan3d/model_/image_processor.py +++ b/comfy/ldm/hunyuan3d/model_/image_processor.py @@ -5,81 +5,178 @@ from PIL import Image import torch.nn.functional as F def to_tensor(image_pt): - image_pt = image_pt / 255 * 2 - 1 - if image_pt.dim() == 4: - image_pt = image_pt.permute(0, 3, 1, 2) - else: - image_pt = image_pt.permute(2, 1, 0) - return image_pt -def resize_bilinear(img: torch.Tensor, size: int) -> torch.Tensor: - # pytorch implementation of cv2.INTER_LINEAR +def resize_nearest(img: torch.Tensor, size: int) -> torch.Tensor: batched = (img.ndim == 4) + if img.ndim == 3: + img = img.unsqueeze(0) + + img = img.permute(0, 3, 1, 2) + + out = F.interpolate(img, size=size, mode='nearest') + + if not batched: + out = out.squeeze(0) + + return out + +def cubic_kernel(x, a: float = -0.75): + absx = x.abs() + absx2 = absx ** 2 + absx3 = absx ** 3 + + w = (a + 2) * absx3 - (a + 3) * absx2 + 1 + w2 = a * absx3 - 5*a * absx2 + 8*a * absx - 4*a + + return torch.where(absx <= 1, w, torch.where(absx < 2, w2, torch.zeros_like(x))) + + +def get_indices_weights(in_size, out_size, scale): + # OpenCV-style half-pixel mapping + x = torch.arange(out_size, dtype=torch.float32) + x = (x + 0.5) / scale - 0.5 + + x0 = x.floor().long() + dx = x.unsqueeze(1) - (x0.unsqueeze(1) + torch.arange(-1, 3)) + + weights = cubic_kernel(dx) + weights = weights / weights.sum(dim=1, keepdim=True) + + indices = x0.unsqueeze(1) + torch.arange(-1, 3) + indices = indices.clamp(0, in_size - 1) + + return indices, weights + + +def resize_cubic_1d(x, out_size, dim): + b, c, h, w = x.shape + in_size = h if dim == 2 else w + scale = out_size / in_size + + indices, weights = get_indices_weights(in_size, out_size, scale) + + if dim == 2: + x = x.permute(0, 1, 3, 2) + x = x.reshape(-1, h) + else: + x = x.reshape(-1, w) + + gathered = x[:, indices] + out = (gathered * weights.unsqueeze(0)).sum(dim=2) + + if dim == 2: + out = out.reshape(b, c, w, out_size).permute(0, 1, 3, 2) + else: + out = out.reshape(b, c, h, out_size) + + return out + + +def resize_cubic(img: torch.Tensor, size: tuple) -> torch.Tensor: + """ + Resize image using OpenCV-equivalent INTER_CUBIC interpolation. + Implemented in pure PyTorch + """ + if img.ndim == 3: img = img.unsqueeze(0) - B, _, H, W = img.shape - H_out = W_out = size - - xs = torch.linspace(0, H_out - 1, H_out, device = img.device) - ys = torch.linspace(0, W_out - 1, W_out, device = img.device) - - xs = (xs + 0.5) * (H / H_out) - 0.5 - ys = (ys + 0.5) * (W / W_out) - 0.5 - - # normalize - xs = 2 * xs / (H - 1) - 1 - ys = 2 * ys / (W - 1) - 1 - - # meshgrid in “ij” order: first rows (xs), then cols (ys) - grid_i, grid_j = torch.meshgrid(xs, ys, indexing='ij') - - # stack into (x,y) where x=columns, y=rows - grid = torch.stack((grid_j, grid_i), dim=-1) - grid = grid.unsqueeze(0).expand(B, -1, -1, -1) - - out = F.grid_sample(img, grid, mode = 'bilinear', - padding_mode = 'zeros', align_corners = True) - - return out if batched else out.squeeze(0) - -def resize_bicubic(img: torch.Tensor, size: int) -> torch.Tensor: - # pytorch implementation of INTER_CUBIC - - was_batched = img.ndim == 4 - if img.ndim == 3: - img = img.unsqueeze(0) - - out = F.interpolate( - img.permute(0, 3, 2, 1), - size = (size, size), - mode = "bicubic", - align_corners = True - ) - - return out if was_batched else out.squeeze(0) + img = img.permute(0, 3, 1, 2) + + out_h, out_w = size + img = resize_cubic_1d(img, out_h, dim=2) + img = resize_cubic_1d(img, out_w, dim=3) + return img def resize_area(img: torch.Tensor, size: tuple) -> torch.Tensor: - # pytorch implementation of INTER_AREA + # vectorized implementation for OpenCV's INTER_AREA using pure PyTorch + original_shape = img.shape + is_hwc = False - was_batched = img.ndim == 4 if img.ndim == 3: - img = img.unsqueeze(0) - - image = F.interpolate(img.permute(0,3,1,2).float(), (size[1], size[0]), mode = "area") - - if was_batched: - image = image.permute(0, 2, 3, 1) # return to channel last + if img.shape[0] <= 4: + img = img.unsqueeze(0) + else: + is_hwc = True + img = img.permute(2, 0, 1).unsqueeze(0) + elif img.ndim == 4: + pass else: - image = image.squeeze(0).permute(1, 2, 0) + raise ValueError("Expected image with 3 or 4 dims.") + + B, C, H, W = img.shape + out_h, out_w = size + scale_y = H / out_h + scale_x = W / out_w + + device = img.device + + # compute the grid boundries + y_start = torch.arange(out_h, device=device).float() * scale_y + y_end = y_start + scale_y + x_start = torch.arange(out_w, device=device).float() * scale_x + x_end = x_start + scale_x + + # for each output pixel, we will compute the range for it + y_start_int = torch.floor(y_start).long() + y_end_int = torch.ceil(y_end).long() + x_start_int = torch.floor(x_start).long() + x_end_int = torch.ceil(x_end).long() + + # We will build the weighted sums by iterating over contributing input pixels once + output = torch.zeros((B, C, out_h, out_w), dtype=torch.float32, device=device) + area = torch.zeros((out_h, out_w), dtype=torch.float32, device=device) + + max_kernel_h = int(torch.max(y_end_int - y_start_int).item()) + max_kernel_w = int(torch.max(x_end_int - x_start_int).item()) + + for dy in range(max_kernel_h): + for dx in range(max_kernel_w): + # compute the weights for this offset for all output pixels + + y_idx = y_start_int.unsqueeze(1) + dy + x_idx = x_start_int.unsqueeze(0) + dx + + # clamp indices to image boundaries + y_idx_clamped = torch.clamp(y_idx, 0, H - 1) + x_idx_clamped = torch.clamp(x_idx, 0, W - 1) + + # compute weights by broadcasting + y_weight = (torch.min(y_end.unsqueeze(1), y_idx_clamped.float() + 1.0) - torch.max(y_start.unsqueeze(1), y_idx_clamped.float())).clamp(min=0) + x_weight = (torch.min(x_end.unsqueeze(0), x_idx_clamped.float() + 1.0) - torch.max(x_start.unsqueeze(0), x_idx_clamped.float())).clamp(min=0) + + weight = (y_weight * x_weight) + + y_expand = y_idx_clamped.expand(out_h, out_w) + x_expand = x_idx_clamped.expand(out_h, out_w) + + + pixels = img[:, :, y_expand, x_expand] + + # unsqueeze to broadcast + w = weight.unsqueeze(0).unsqueeze(0) + + output += pixels * w + area += weight + + # Normalize by area + output /= area.unsqueeze(0).unsqueeze(0) + + if is_hwc: + return output[0].permute(1, 2, 0) + elif img.shape[0] == 1 and original_shape[0] <= 4: + return output[0] + else: + return output - return image if was_batched else image.squeeze(0) class ImageProcessorV2(nn.Module): def __init__(self, size: int = 512, border_ratio: float = None): + super().__init__() + self.size = size self.border_ratio = border_ratio @@ -98,12 +195,14 @@ class ImageProcessorV2(nn.Module): img = torch.from_numpy(img) img, mask = self.recenter(img, border_ratio = border_ratio) - img = resize_bicubic(img, size = self.size) - mask = resize_bilinear(mask.float(), size = self.size) + img = resize_cubic(img, size = (self.size, self.size)) + mask = resize_nearest(mask.float(), size = self.size) mask = mask[..., torch.newaxis] img = to_tensor(img) + mask = to_tensor(mask) + mask = mask.permute(0, 3, 1, 2) return img, mask @@ -147,7 +246,7 @@ class ImageProcessorV2(nn.Module): y2_max = y2_min + w2 # note: opencv takes columns first (opposite to pytorch and numpy that take the row first) - result[x2_min:x2_max, y2_min:y2_max] = resize_area(image[x_min:x_max, y_min:y_max], (w2, h2)) + result[x2_min:x2_max, y2_min:y2_max] = resize_area(image[x_min:x_max, y_min:y_max], (h2, w2)) bg = torch.ones((result.shape[0], result.shape[1], 3), dtype = torch.uint8) * 255 @@ -175,18 +274,22 @@ class ImageProcessorV2(nn.Module): return outputs def test_image_processor(): + + """ + implementation speed: 0.24465346336364746 + reference speed: 2.046062469482422 + atol = 4e-2: True + """ + + import time import matplotlib.pyplot as plt image_processor = ImageProcessorV2(size = 224) - import time start = time.time() outputs = image_processor(image = r"C:\Users\yrafa\Work\Hunyuan 3D\cat.jpg") print(time.time() - start) image = outputs["image"] print(image.shape) - plt.imshow(image) + plt.imshow(image.squeeze().permute(1, 2, 0).numpy()) plt.axis("off") - plt.show() - -if __name__ == "__main__": - test_image_processor() \ No newline at end of file + plt.show() \ No newline at end of file diff --git a/comfy/ldm/hunyuan3d/model_/pipeline.py b/comfy/ldm/hunyuan3d/model_/pipeline.py index 93b9aa45b..65c50847c 100644 --- a/comfy/ldm/hunyuan3d/model_/pipeline.py +++ b/comfy/ldm/hunyuan3d/model_/pipeline.py @@ -5,6 +5,7 @@ from PIL import Image from typing import List, Union from torch.utils._pytree import tree_map from torch.utils.data._utils.collate import default_collate +from vae import VAE def export_to_trimesh(mesh_output): if isinstance(mesh_output, list): @@ -23,19 +24,30 @@ def export_to_trimesh(mesh_output): return mesh_output class Hunyuan3DDiTFlowMatchingPipeline(nn.Module): - def __init__(self, model, vae, conditioner, image_processor, scheduler): - + def __init__(self, model, vae, conditioner, image_processor, scheduler, device, dtype): + super().__init__() + self.vae = vae self.model = model self.conditioner = conditioner self.image_processor = image_processor self.scheduler = scheduler + self.device = device + self.dtype = dtype def compile(self): self.vae = torch.compile(self.vae) self.model = torch.compile(self.model) self.conditioner = torch.compile(self.conditioner) + def load_ckpt(self, checkpoint_path: str): + + checkpoint = torch.load(checkpoint_path, weights_only = True) + self.model.load_state_dict(checkpoint["model"]) + self.vae.load_state_dict(checkpoint["vae"]) + self.conditioner.load_state_dict(checkpoint["conditioner"]) + + def encode_cond(self, image, additional_cond_inputs, do_classifier_free_guidance): bsz = image.shape[0] @@ -51,8 +63,23 @@ class Hunyuan3DDiTFlowMatchingPipeline(nn.Module): return cond + def to(self, device=None, dtype=None): + if dtype is not None: + self.dtype = dtype + self.vae.to(dtype=dtype) + self.model.to(dtype=dtype) + self.conditioner.to(dtype=dtype) + if device is not None: + self.device = torch.device(device) + self.vae.to(device) + self.model.to(device) + self.conditioner.to(device) + def prepare_images(self, images): + if isinstance(images, (str, Image.Image)): + return self.image_processor(images) + outputs = [] for image in images: output = self.image_processor(image) @@ -107,7 +134,7 @@ class Hunyuan3DDiTFlowMatchingPipeline(nn.Module): self.model.guidance_embed is True ) - cond_inputs = self.prepare_image(image) + cond_inputs = self.prepare_images(image) image = cond_inputs.pop('image') cond = self.encode_cond( @@ -135,7 +162,7 @@ class Hunyuan3DDiTFlowMatchingPipeline(nn.Module): latent_model_input = latents timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) - timestep = timestep / self.scheduler.num_train_timesteps + timestep = timestep / self.scheduler.num_training_timesteps noise_pred = self.model(latent_model_input, timestep, cond, guidance=guidance) if do_classifier_free_guidance: @@ -143,7 +170,7 @@ class Hunyuan3DDiTFlowMatchingPipeline(nn.Module): noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents) + latents = self.scheduler.reverse_flow(noise_pred, latents) if callback is not None and i % callback_steps == 0: step_idx = i // getattr(self.scheduler, "order", 1) @@ -153,4 +180,18 @@ class Hunyuan3DDiTFlowMatchingPipeline(nn.Module): mesh = self.vae.decode(latents, bounds = bounds, octree_res = octree_res, num_chunks = num_chunks) return export_to_trimesh(mesh) - \ No newline at end of file + +if __name__ == '__main__': + from scheduler import EulerScheduler + from conditioner import SingleImageEncoder + from image_processor import ImageProcessorV2 + from dinov2 import DinoConfig + from hunyuandit import HunYuanDiTPlain + + model = HunYuanDiTPlain(depth = 2) + + pipeline = Hunyuan3DDiTFlowMatchingPipeline(vae = VAE(), scheduler = EulerScheduler(), model = model, + conditioner = SingleImageEncoder(DinoConfig()), image_processor = ImageProcessorV2(), + device = "cpu", dtype = torch.bfloat16) + img = r"C:\Users\yrafa\Work\Hunyuan 3D\cat.jpg" + print(pipeline(img)) \ No newline at end of file diff --git a/comfy/ldm/hunyuan3d/model_/scheduler.py b/comfy/ldm/hunyuan3d/model_/scheduler.py index 0d1f777c6..a060f9623 100644 --- a/comfy/ldm/hunyuan3d/model_/scheduler.py +++ b/comfy/ldm/hunyuan3d/model_/scheduler.py @@ -2,7 +2,7 @@ import torch class EulerScheduler(torch.nn.Module): def __init__(self, num_training_timesteps: int = 1_000, shift: float = 1, - num_inference_timesteps: int = 100, inference: bool = False): + num_inference_timesteps: int = 50, inference: bool = True): super(EulerScheduler, self).__init__() # compute timestep values so we can index into them later diff --git a/comfy/ldm/hunyuan3d/model_/vae.py b/comfy/ldm/hunyuan3d/model_/vae.py new file mode 100644 index 000000000..7687b74c9 --- /dev/null +++ b/comfy/ldm/hunyuan3d/model_/vae.py @@ -0,0 +1,1010 @@ +# replaced torch.ops.torch_cluster.fps with a manual implementation +# to avoid having torch_cluster downloaded as dependency +# also the dependency takes a long time to install + +import torch +from torch import Tensor +import math + +def fps(src: Tensor, batch: Tensor, sampling_ratio: float, start_random: bool = True): + + # manually create the pointer vector + assert src.size(0) == batch.numel() + + batch_size = int(batch.max()) + 1 + deg = src.new_zeros(batch_size, dtype = torch.long) + + deg.scatter_add_(0, batch, torch.ones_like(batch)) + + ptr_vec = deg.new_zeros(batch_size + 1) + torch.cumsum(deg, 0, out=ptr_vec[1:]) + + #return fps_sampling(src, ptr_vec, ratio) + sampled_indicies = [] + + for b in range(batch_size): + # start and the end of each batch + start, end = ptr_vec[b].item(), ptr_vec[b + 1].item() + # points from the point cloud + points = src[start:end] + + num_points = points.size(0) + num_samples = max(1, math.ceil(num_points * sampling_ratio)) + + selected = torch.zeros(num_samples, device = src.device, dtype = torch.long) + distances = torch.full((num_points,), float("inf"), device = src.device) + + # select a random start point + if start_random: + farthest = torch.randint(0, num_points, (1,), device = src.device) + else: farthest = torch.tensor([0], device = src.device, dtype = torch.long) + + for i in range(num_samples): + selected[i] = farthest + centroid = points[farthest].squeeze(0) + dist = torch.norm(points - centroid, dim = 1) # compute euclidean distance + distances = torch.minimum(distances, dist) + farthest = torch.argmax(distances) + + sampled_indicies.append(torch.arange(start, end)[selected]) + + return torch.cat(sampled_indicies, dim = 0) + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + + def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + + if keep_prob > 0.0 and self.scale_by_keep: + random_tensor.div_(keep_prob) + + return x * random_tensor + +class MLP(nn.Module): + def __init__(self, width: int, ratio: int = 4, drop_path_rate: float = 0): + super().__init__() + self.gelu = nn.GELU() + self.c_fc = nn.Linear(width, width * ratio) + self.c_proj = nn.Linear(width * ratio, width) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity() + + def forward(self, x): + return self.drop_path(self.c_proj(self.gelu(self.c_fc(x)))) + + +class QKVMultiheadAttention(nn.Module): + def __init__( + self, + heads: int, + n_ctx: int, + width=None, + qk_norm=False, + norm_layer=nn.LayerNorm + ): + super().__init__() + self.heads = heads + self.n_ctx = n_ctx + self.q_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() + self.k_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() + + def forward(self, qkv): + bs, n_ctx, width = qkv.shape + attn_ch = width // self.heads // 3 + qkv = qkv.view(bs, n_ctx, self.heads, -1) + q, k, v = torch.split(qkv, attn_ch, dim=-1) + + q = self.q_norm(q) + k = self.k_norm(k) + + q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)] + out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(bs, n_ctx, -1) + return out + + +class MultiheadAttention(nn.Module): + def __init__( + self, + n_ctx: int, + width: int, + heads: int, + qkv_bias: bool, + norm_layer = nn.LayerNorm, + qk_norm: bool = False, + drop_path_rate: float = 0.0 + ): + super().__init__() + + self.c_qkv = nn.Linear(width, width * 3, bias=qkv_bias) + self.c_proj = nn.Linear(width, width) + + self.attention = QKVMultiheadAttention( + heads = heads, + n_ctx = n_ctx, + width = width, + norm_layer = norm_layer, + qk_norm = qk_norm + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity() + + def forward(self, x): + x = self.c_qkv(x) + x = self.attention(x) + x = self.drop_path(self.c_proj(x)) + return x + + +class ResAttnBlock(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + heads: int, + qkv_bias: bool = True, + norm_layer=nn.LayerNorm, + qk_norm: bool = False, + drop_path_rate: float = 0.0, + ): + super().__init__() + self.attn = MultiheadAttention( + n_ctx=n_ctx, + width=width, + heads=heads, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + qk_norm=qk_norm, + drop_path_rate=drop_path_rate + ) + self.ln_1 = norm_layer(width, elementwise_affine=True, eps=1e-6) + self.mlp = MLP(width=width, drop_path_rate=drop_path_rate) + self.ln_2 = norm_layer(width, elementwise_affine=True, eps=1e-6) + + def forward(self, x: torch.Tensor): + x = x + self.attn(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + +class Transformer(nn.Module): + def __init__(self, n_ctx: int, heads: int, width: int, depth: int, + qkv_bias: bool = True, qk_norm: bool = False, drop_path_rate: float = 0.0): + super().__init__() + + self.resblocks = nn.ModuleList([ + ResAttnBlock(n_ctx = n_ctx, + heads = heads, + width = width, + qkv_bias = qkv_bias, + qk_norm = qk_norm, + drop_path_rate = drop_path_rate) + for _ in range(depth) + ]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + for resnet in self.resblocks: + x = resnet(x) + + return x + +class QKVMultiheadCrossAttention(nn.Module): + def __init__( + self, + heads: int, + n_data = None, + width=None, + qk_norm=False, + norm_layer=nn.LayerNorm + ): + super().__init__() + self.heads = heads + self.n_data = n_data + self.q_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() + self.k_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() + + def forward(self, q, kv): + + _, n_ctx, _ = q.shape + bs, n_data, width = kv.shape + + attn_ch = width // self.heads // 2 + q = q.view(bs, n_ctx, self.heads, -1) + + kv = kv.view(bs, n_data, self.heads, -1) + k, v = torch.split(kv, attn_ch, dim=-1) + + q = self.q_norm(q) + k = self.k_norm(k) + + q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)] + out = F.scaled_dot_product_attention(q, k, v) + + out = out.transpose(1, 2).reshape(bs, n_ctx, -1) + + return out + + +class MultiheadCrossAttention(nn.Module): + def __init__( + self, + width: int, + heads: int, + qkv_bias: bool = False, + n_data = None, + norm_layer = nn.LayerNorm, + qk_norm: bool = False, + kv_cache: bool = False, + ): + super().__init__() + + self.c_q = nn.Linear(width, width, bias=qkv_bias) + self.c_kv = nn.Linear(width, width * 2, bias=qkv_bias) + self.c_proj = nn.Linear(width, width) + + self.attention = QKVMultiheadCrossAttention( + heads = heads, + n_data = n_data, + width = width, + norm_layer = norm_layer, + qk_norm = qk_norm + ) + + self.kv_cache = kv_cache + self.data = None + + def forward(self, x, data): + x = self.c_q(x) + + if self.kv_cache: + if self.data is None: + self.data = self.c_kv(data) + + data = self.data + else: + data = self.c_kv(data) + + x = self.attention(x, data) + x = self.c_proj(x) + + return x + + +class ResidualCrossAttentionBlock(nn.Module): + def __init__( + self, + width: int, + heads: int, + n_data: int = None, + mlp_expand_ratio: int = 4, + qkv_bias: bool = False, + norm_layer=nn.LayerNorm, + qk_norm: bool = False + ): + super().__init__() + + self.attn = MultiheadCrossAttention( + n_data=n_data, + width = width, + heads=heads, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + qk_norm=qk_norm + ) + + self.ln_1 = norm_layer(width, elementwise_affine = True, eps = 1e-6) + self.ln_2 = norm_layer(width, elementwise_affine = True, eps = 1e-6) + self.ln_3 = norm_layer(width, elementwise_affine = True, eps = 1e-6) + + self.mlp = MLP(width=width, ratio = mlp_expand_ratio) + + def forward(self, x: torch.Tensor, data: torch.Tensor): + x = x + self.attn(self.ln_1(x), self.ln_2(data)) + x = x + self.mlp(self.ln_3(x)) + return x + +class CrossAttentionDecoder(nn.Module): + def __init__( + self, + num_latents: int, + out_channels: int, + fourier_embedder, + width: int, + heads: int, + mlp_expand_ratio: int = 4, + downsample_ratio: int = 1, + enable_ln_post: bool = True, + qkv_bias: bool = False, + qk_norm: bool = False): + + super().__init__() + + self.enable_ln_post = enable_ln_post + self.fourier_embedder = fourier_embedder + self.downsample_ratio = downsample_ratio + + self.query_proj = nn.Linear(self.fourier_embedder.out_dim, width) + + if self.downsample_ratio != 1: + self.latents_proj = nn.Linear(width * downsample_ratio, width) + + if self.enable_ln_post == False: + qk_norm = False + + self.cross_attn_decoder = ResidualCrossAttentionBlock( + n_data=num_latents, + width=width, + mlp_expand_ratio=mlp_expand_ratio, + heads=heads, + qkv_bias=qkv_bias, + qk_norm=qk_norm + ) + + if self.enable_ln_post: + self.ln_post = nn.LayerNorm(width) + + self.output_proj = nn.Linear(width, out_channels) + self.count = 0 + + def forward(self, queries = None, query_embeddings = None, latents = None): + + if query_embeddings is None: + query_embeddings = self.query_proj(self.fourier_embedder(queries).to(latents.dtype)) + + self.count += query_embeddings.shape[1] + + if self.downsample_ratio != 1: + latents = self.latents_proj(latents) + + x = self.cross_attn_decoder(query_embeddings, latents) + + if self.enable_ln_post: + x = self.ln_post(x) + + out = self.output_proj(x) + + return out + + +class PointCrossAttention(nn.Module): + def __init__(self, + num_latents: int, + downsample_ratio: float, + pc_size: int, + pc_sharpedge_size: int, + point_feats: int, + width: int, + heads: int, + layers: int, + fourier_embedder, + normal_pe: bool = False, + qkv_bias: bool = False, + use_ln_post: bool = True, + qk_norm: bool = True): + + super().__init__() + + self.fourier_embedder = fourier_embedder + + self.pc_size = pc_size + self.normal_pe = normal_pe + self.downsample_ratio = downsample_ratio + self.pc_sharpedge_size = pc_sharpedge_size + self.num_latents = num_latents + self.point_feats = point_feats + + self.input_proj = nn.Linear(self.fourier_embedder.out_dim + point_feats, width) + + self.cross_attn = ResidualCrossAttentionBlock( + width = width, + heads = heads, + qkv_bias = qkv_bias, + qk_norm = qk_norm + ) + + self.self_attn = None + if layers > 0: + self.self_attn = Transformer( + n_ctx = num_latents, + width = width, + heads = heads, + qkv_bias = qkv_bias, + qk_norm = qk_norm, + depth = layers + ) + + if use_ln_post: + self.ln_post = nn.LayerNorm(width) + else: + self.ln_post = None + + def sample_points_and_latents(self, point_cloud: torch.Tensor, features: torch.Tensor): + + """ + Subsample points randomly from the point cloud (input_pc) + Further sample the subsampled points to get query_pc + take the fourier embeddings for both input and query pc + + Mental Note: FPS-sampled points (query_pc) act as latent tokens that attend to and learn from the broader context in input_pc. + Goal: get a smaller represenation (query_pc) to represent the entire scence structure by learning from a broader subset (input_pc). + More computationally efficient. + + Features are additional information for each point in the cloud + """ + + B, _, D = point_cloud.shape + + num_latents = int(self.num_latents) + + num_random_query = self.pc_size / (self.pc_size + self.pc_sharpedge_size) * num_latents + num_sharpedge_query = num_latents - num_random_query + + # Split random and sharpedge surface points + random_pc, sharpedge_pc = torch.split(point_cloud, [self.pc_size, self.pc_sharpedge_size], dim=1) + + # assert statements + assert random_pc.shape[1] <= self.pc_size, "Random surface points size must be less than or equal to pc_size" + assert sharpedge_pc.shape[1] <= self.pc_sharpedge_size, "Sharpedge surface points size must be less than or equal to pc_sharpedge_size" + + input_random_pc_size = int(num_random_query * self.downsample_ratio) + random_query_pc, random_input_pc, random_idx_pc, random_idx_query = \ + self.subsample(pc = random_pc, num_query = num_random_query, input_pc_size = input_random_pc_size) + + input_sharpedge_pc_size = int(num_sharpedge_query * self.downsample_ratio) + + if input_sharpedge_pc_size == 0: + sharpedge_input_pc = torch.zeros(B, 0, D, dtype = random_input_pc.dtype).to(point_cloud.device) + sharpedge_query_pc = torch.zeros(B, 0, D, dtype= random_query_pc.dtype).to(point_cloud.device) + + else: sharpedge_query_pc, sharpedge_input_pc, sharpedge_idx_pc, sharpedge_idx_query = \ + self.subsample(pc = sharpedge_pc, num_query = num_sharpedge_query, input_pc_size = input_sharpedge_pc_size) + + # concat the random and sharpedges + query_pc = torch.cat([random_query_pc, sharpedge_query_pc], dim = 1) + input_pc = torch.cat([random_input_pc, sharpedge_input_pc], dim = 1) + + query = self.fourier_embedder(query_pc) + data = self.fourier_embedder(input_pc) + + if self.point_feats > 0: + random_surface_features, sharpedge_surface_features = torch.split(features, [self.pc_size, self.pc_sharpedge_size], dim = 1) + + input_random_surface_features, query_random_features = \ + self.handle_features(features = random_surface_features, idx_pc = random_idx_pc, batch_size = B, + input_pc_size = input_random_pc_size, idx_query = random_idx_query) + + if input_sharpedge_pc_size == 0: + input_sharpedge_surface_features = torch.zeros(B, 0, self.point_feats, + dtype = input_random_surface_features.dtype, device = point_cloud.device) + + query_sharpedge_features = torch.zeros(B, 0, self.point_feats, + dtype = query_random_features.dtype, device = point_cloud.device) + else: + + input_sharpedge_surface_features, query_sharpedge_features = \ + self.handle_features(idx_pc = sharpedge_idx_pc, features = sharpedge_surface_features, + batch_size = B, idx_query = sharpedge_idx_query, input_pc_size = input_sharpedge_pc_size) + + query_features = torch.cat([query_random_features, query_sharpedge_features], dim = 1) + input_features = torch.cat([input_random_surface_features, input_sharpedge_surface_features], dim = 1) + + if self.normal_pe: + # apply the fourier embeddings on the first 3 dims (xyz) + input_features_pe = self.fourier_embedder(input_features[..., :3]) + query_features_pe = self.fourier_embedder(query_features[..., :3]) + # replace the first 3 dims with the new PE ones + input_features = torch.cat([input_features_pe, input_features[..., :3]], dim = -1) + query_features = torch.cat([query_features_pe, query_features[..., :3]], dim = -1) + + # concat at the channels dim + query = torch.cat([query, query_features], dim = -1) + data = torch.cat([data, input_features], dim = -1) + + # don't return pc_info to avoid unnecessary memory usuage + return query.view(B, -1, query.shape[-1]), data.view(B, -1, data.shape[-1]) + + def forward(self, point_cloud: torch.Tensor, features: torch.Tensor): + + query, data = self.sample_points_and_latents(point_cloud = point_cloud, features = features) + + # apply projections + query = self.input_proj(query) + data = self.input_proj(data) + + # apply cross attention between query and data + latents = self.cross_attn(query, data) + + if self.self_attn is not None: + latents = self.self_attn(latents) + + if self.ln_post is not None: + latents = self.ln_post(latents) + + return latents + + + def subsample(self, pc, num_query, input_pc_size: int): + + """ + num_query: number of points to keep after FPS + input_pc_size: number of points to select before FPS + """ + + B, _, D = pc.shape + query_ratio = num_query / input_pc_size + + # random subsampling of points inside the point cloud + idx_pc = torch.randperm(pc.shape[1], device = pc.device)[:input_pc_size] + input_pc = pc[:, idx_pc, :] + + # flatten to allow applying fps across the whole batch + flattent_input_pc = input_pc.view(B * input_pc_size, D) + + # construct a batch_down tensor to tell fps + # which points belong to which batch + N_down = int(flattent_input_pc.shape[0] / B) + batch_down = torch.arange(B).to(pc.device) + batch_down = torch.repeat_interleave(batch_down, N_down) + + idx_query = fps(flattent_input_pc, batch_down, sampling_ratio = query_ratio) + query_pc = flattent_input_pc[idx_query].view(B, -1, D) + + return query_pc, input_pc, idx_pc, idx_query + + def handle_features(self, features, idx_pc, input_pc_size, batch_size: int, idx_query): + + B = batch_size + + input_surface_features = features[:, idx_pc, :] + flattent_input_features = input_surface_features.view(B * input_pc_size, -1) + query_features = flattent_input_features[idx_query].view(B, -1, + flattent_input_features.shape[-1]) + + return input_surface_features, query_features + + def forward(self, pc, feats): + """ + + Args: + pc (torch.FloatTensor): [B, N, 3] + feats (torch.FloatTensor or None): [B, N, C] + + Returns: + + """ + + query, data = self.sample_points_and_latents(pc, feats) + + query = self.input_proj(query) + query = query + data = self.input_proj(data) + data = data + + latents = self.cross_attn(query, data) + if self.self_attn is not None: + latents = self.self_attn(latents) + + if self.ln_post is not None: + latents = self.ln_post(latents) + + return latents + +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 + +import trimesh +import torch +import numpy as np + +def normalize_mesh(mesh, scale = 0.9999): + """Normalize mesh to fit in [-scale, scale]. Translate mesh so its center is [0,0,0]""" + + bbox = mesh.bounds + center = (bbox[1] + bbox[0]) / 2 + + max_extent = (bbox[1] - bbox[0]).max() + mesh.apply_translation(-center) + mesh.apply_scale((2 * scale) / max_extent) + + return mesh + +def sample_pointcloud(mesh, num = 200000): + """ Uniformly sample points from the surface of the mesh """ + + points, face_idx = mesh.sample(num, return_index = True) + normals = mesh.face_normals[face_idx] + return torch.from_numpy(points.astype(np.float32)), torch.from_numpy(normals.astype(np.float32)) + +def detect_sharp_edges(mesh, threshold=0.985): + """Return edge indices (a, b) that lie on sharp boundaries of the mesh.""" + + V, F = mesh.vertices, mesh.faces + VN, FN = mesh.vertex_normals, mesh.face_normals + + sharp_mask = np.ones(V.shape[0]) + for i in range(3): + indices = F[:, i] + alignment = np.einsum('ij,ij->i', VN[indices], FN) + dot_stack = np.stack((sharp_mask[indices], alignment), axis=-1) + sharp_mask[indices] = np.min(dot_stack, axis=-1) + + edge_a = np.concatenate([F[:, 0], F[:, 1], F[:, 2]]) + edge_b = np.concatenate([F[:, 1], F[:, 2], F[:, 0]]) + sharp_edges = (sharp_mask[edge_a] < threshold) & (sharp_mask[edge_b] < threshold) + + return edge_a[sharp_edges], edge_b[sharp_edges] + + +def sharp_sample_pointcloud(mesh, num = 16384): + """ Sample points preferentially from sharp edges in the mesh. """ + + edge_a, edge_b = detect_sharp_edges(mesh) + V, VN = mesh.vertices, mesh.vertex_normals + + va, vb = V[edge_a], V[edge_b] + na, nb = VN[edge_a], VN[edge_b] + + edge_lengths = np.linalg.norm(vb - va, axis=-1) + weights = edge_lengths / edge_lengths.sum() + + indices = np.searchsorted(np.cumsum(weights), np.random.rand(num)) + t = np.random.rand(num, 1) + + samples = t * va[indices] + (1 - t) * vb[indices] + normals = t * na[indices] + (1 - t) * nb[indices] + + return samples.astype(np.float32), normals.astype(np.float32) + +def load_surface_sharpedge(mesh, num_points=4096, num_sharp_points=4096, sharpedge_flag = True, device = "cuda"): + """Load a surface with optional sharp-edge annotations from a trimesh mesh.""" + + try: + mesh_full = trimesh.util.concatenate(mesh.dump()) + except Exception: + mesh_full = trimesh.util.concatenate(mesh) + + mesh_full = normalize_mesh(mesh_full) + + faces = mesh_full.faces + vertices = mesh_full.vertices + origin_face_count = faces.shape[0] + + mesh_surface = trimesh.Trimesh(vertices=vertices, faces=faces[:origin_face_count]) + mesh_fill = trimesh.Trimesh(vertices=vertices, faces=faces[origin_face_count:]) + + area_surface = mesh_surface.area + area_fill = mesh_fill.area + total_area = area_surface + area_fill + + sample_num = 499712 // 2 + fill_ratio = area_fill / total_area if total_area > 0 else 0 + + num_fill = int(sample_num * fill_ratio) + num_surface = sample_num - num_fill + + surf_pts, surf_normals = sample_pointcloud(mesh_surface, num_surface) + fill_pts, fill_normals = (torch.zeros(0, 3), torch.zeros(0, 3)) if num_fill == 0 else sample_pointcloud(mesh_fill, num_fill) + + sharp_pts, sharp_normals = sharp_sample_pointcloud(mesh_surface, sample_num) + + def assemble_tensor(points, normals, label=None): + + data = torch.cat([points, normals], dim=1).half().to(device) + + if label is not None: + label_tensor = torch.full((data.shape[0], 1), float(label), dtype=torch.float16).to(device) + data = torch.cat([data, label_tensor], dim=1) + + return data + + surface = assemble_tensor(torch.cat([surf_pts.to(device), fill_pts.to(device)], dim=0), + torch.cat([surf_normals.to(device), fill_normals.to(device)], dim=0), + label = 0 if sharpedge_flag else None) + + sharp_surface = assemble_tensor(torch.from_numpy(sharp_pts), torch.from_numpy(sharp_normals), + label = 1 if sharpedge_flag else None) + + rng = np.random.default_rng() + + surface = surface[rng.choice(surface.shape[0], num_points, replace = False)] + sharp_surface = sharp_surface[rng.choice(sharp_surface.shape[0], num_sharp_points, replace = False)] + + full = torch.cat([surface, sharp_surface], dim = 0).unsqueeze(0) + + return full + +class SharpEdgeSurfaceLoader: + """ Load mesh surface and sharp edge samples. """ + + def __init__(self, num_uniform_points = 8192, num_sharp_points = 8192): + + self.num_uniform_points = num_uniform_points + self.num_sharp_points = num_sharp_points + self.total_points = num_uniform_points + num_sharp_points + + def __call__(self, mesh_input, device = "cuda"): + mesh = self._load_mesh(mesh_input) + return load_surface_sharpedge(mesh, self.num_uniform_points, self.num_sharp_points, device = device) + + @staticmethod + def _load_mesh(mesh_input): + + if isinstance(mesh_input, str): + mesh = trimesh.load(mesh_input, force="mesh", merge_primitives = True) + else: + mesh = mesh_input + + if isinstance(mesh, trimesh.Scene): + combined = None + for obj in mesh.geometry.values(): + combined = obj if combined is None else combined + obj + return combined + + return mesh + +class FourierEmbedder(nn.Module): + def __init__(self, num_freq: int = 8, input_dim: int = 3, include_pi: bool = False): + super().__init__() + + frequencies = 2.0 ** torch.arange( + num_freq, + dtype = torch.float32 + ) + + if include_pi: + frequencies *= torch.pi + + self.register_buffer("frequencies", frequencies, persistent = False) + + self.out_dim = input_dim * (num_freq * 2 + 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + embed = (x[..., None].contiguous() * self.frequencies).view(*x.shape[:-1], -1) + return torch.cat((x, embed.sin(), embed.cos()), dim = -1) + +class DiagonalGaussianDistribution: + def __init__(self, params: torch.Tensor, feature_dim: int = -1): + + # divide quant channels (8) into mean and log variance + self.mean, self.logvar = torch.chunk(params, 2, dim = feature_dim) + + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.std = torch.exp(0.5 * self.logvar) + + def sample(self): + + eps = torch.randn_like(self.std) + z = self.mean + eps * self.std + + return z + +class VAE(nn.Module): + def __init__(self, + *, + num_latents: int = 4096, + embed_dim: int = 64, + width: int = 1024, + heads: int = 16, + num_decoder_layers: int = 16, + num_encoder_layers: int = 8, + pc_size: int = 81920, + pc_sharpedge_size: int = 0, + point_feats: int = 4, + downsample_ratio: int = 20, + geo_decoder_downsample_ratio: int = 1, + geo_decoder_mlp_expand_ratio: int = 4, + geo_decoder_ln_post: bool = True, + num_frequencies: int = 8, + qkv_bias: bool = False, + qk_norm: bool = True, + drop_path_rate: float = 0.0, + include_pi: bool = False, + scale_factor: float = 1.0039506158752403 + ): + + super().__init__() + + self.latent_shape = (num_latents, embed_dim) + self.scale_factor = scale_factor + + self.fourier_embedder = FourierEmbedder(num_freq = num_frequencies, include_pi = include_pi) + + self.encoder = PointCrossAttention(layers = num_encoder_layers, + num_latents = num_latents, + downsample_ratio = downsample_ratio, + heads = heads, + pc_size = pc_size, + width = width, + point_feats = point_feats, + fourier_embedder = self.fourier_embedder, + pc_sharpedge_size = pc_sharpedge_size) + + self.transformer = Transformer( + n_ctx=num_latents, + width=width, + depth=num_decoder_layers, + heads=heads, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + drop_path_rate=drop_path_rate + ) + + self.geo_decoder = CrossAttentionDecoder( + fourier_embedder = self.fourier_embedder, + out_channels = 1, + num_latents = num_latents, + mlp_expand_ratio = geo_decoder_mlp_expand_ratio, + downsample_ratio = geo_decoder_downsample_ratio, + enable_ln_post = geo_decoder_ln_post, + width=width // geo_decoder_downsample_ratio, + heads=heads // geo_decoder_downsample_ratio, + qkv_bias = qkv_bias, + qk_norm= qk_norm + ) + + self.pre_kl = nn.Linear(width, embed_dim * 2) + self.post_kl = nn.Linear(embed_dim, width) + + self.volume_decoder = VanillaVolumeDecoder() + self.surface_extractor = SufraceExtractor() + + + def forward(self): + pass + + def encode(self, surface): + + pc, feats = surface[:, :, :3], surface[:, :, 3:] + latents = self.encoder(pc, feats) + + moments = self.pre_kl(latents) + posterior = DiagonalGaussianDistribution(moments, feature_dim = -1) + + latents = posterior.sample() + + return latents + + def decode(self, latents, to_mesh: bool = True, **kwargs): + + latents = self.post_kl(latents) + latents = self.transformer(latents) + + if not to_mesh: + return latents + + grid_logits = self.volume_decoder(latents = latents, geo_decoder = self.geo_decoder, **kwargs) + mesh = self.surface_extractor(grid_logits, **kwargs) + + return mesh + +def load_vae(vae): + + DEBUG = False + + checkpoint = "model.fp16.ckpt" + missing, unexpected = vae.load_state_dict(torch.load(checkpoint), strict = not DEBUG) + + if DEBUG: + print(f"Missing {len(missing)}: ", missing) + print(f"\nUnexpected {len(unexpected)}: ", unexpected) + + return vae + + \ No newline at end of file diff --git a/comfy/ldm/hunyuan3d/vae/vae.py b/comfy/ldm/hunyuan3d/vae/vae.py index 2c4f47cfe..9358a0ab2 100644 --- a/comfy/ldm/hunyuan3d/vae/vae.py +++ b/comfy/ldm/hunyuan3d/vae/vae.py @@ -68,7 +68,7 @@ class VAE(nn.Module): super().__init__() self.latent_shape = (num_latents, embed_dim) - self.scale = scale_factor + self.scale_factor = scale_factor self.fourier_embedder = FourierEmbedder(num_freq = num_frequencies, include_pi = include_pi)