# 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, Callable
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import repeat
from tqdm import tqdm
from .attention_blocks import CrossAttentionDecoder
from .attention_processors import FlashVDMCrossAttentionProcessor, FlashVDMTopMCrossAttentionProcessor
from ...utils import logger
def extract_near_surface_volume_fn(input_tensor: torch.Tensor, alpha: float):
device = input_tensor.device
D = input_tensor.shape[0]
signed_val = 0.0
# 添加偏移并处理无效值
val = input_tensor + alpha
valid_mask = val > -9000 # 假设-9000是无效值
# 改进的邻居获取函数（保持维度一致）
def get_neighbor(t, shift, axis):
"""根据指定轴进行位移并保持维度一致"""
if shift == 0:
return t.clone()
# 确定填充轴（输入为[D, D, D]对应z,y,x轴）
pad_dims = [0, 0, 0, 0, 0, 0] # 格式：[x前，x后，y前，y后，z前，z后]
# 根据轴类型设置填充
if axis == 0: # x轴（最后一个维度）
pad_idx = 0 if shift > 0 else 1
pad_dims[pad_idx] = abs(shift)
elif axis == 1: # y轴（中间维度）
pad_idx = 2 if shift > 0 else 3
pad_dims[pad_idx] = abs(shift)
elif axis == 2: # z轴（第一个维度）
pad_idx = 4 if shift > 0 else 5
pad_dims[pad_idx] = abs(shift)
# 执行填充（添加batch和channel维度适配F.pad）
padded = F.pad(t.unsqueeze(0).unsqueeze(0), pad_dims[::-1], mode='replicate') # 反转顺序适配F.pad
# 构建动态切片索引
slice_dims = [slice(None)] * 3 # 初始化为全切片
if axis == 0: # x轴（dim=2）
if shift > 0:
slice_dims[0] = slice(shift, None)
else:
slice_dims[0] = slice(None, shift)
elif axis == 1: # y轴（dim=1）
if shift > 0:
slice_dims[1] = slice(shift, None)
else:
slice_dims[1] = slice(None, shift)
elif axis == 2: # z轴（dim=0）
if shift > 0:
slice_dims[2] = slice(shift, None)
else:
slice_dims[2] = slice(None, shift)
# 应用切片并恢复维度
padded = padded.squeeze(0).squeeze(0)
sliced = padded[slice_dims]
return sliced
# 获取各方向邻居（确保维度一致）
left = get_neighbor(val, 1, axis=0) # x方向
right = get_neighbor(val, -1, axis=0)
back = get_neighbor(val, 1, axis=1) # y方向
front = get_neighbor(val, -1, axis=1)
down = get_neighbor(val, 1, axis=2) # z方向
up = get_neighbor(val, -1, axis=2)
# 处理边界无效值（使用where保持维度一致）
def safe_where(neighbor):
return torch.where(neighbor > -9000, neighbor, val)
left = safe_where(left)
right = safe_where(right)
back = safe_where(back)
front = safe_where(front)
down = safe_where(down)
up = safe_where(up)
# 计算符号一致性（转换为float32确保精度）
sign = torch.sign(val.to(torch.float32))
neighbors_sign = torch.stack([
torch.sign(left.to(torch.float32)),
torch.sign(right.to(torch.float32)),
torch.sign(back.to(torch.float32)),
torch.sign(front.to(torch.float32)),
torch.sign(down.to(torch.float32)),
torch.sign(up.to(torch.float32))
], dim=0)
# 检查所有符号是否一致
same_sign = torch.all(neighbors_sign == sign, dim=0)
# 生成最终掩码
mask = (~same_sign).to(torch.int32)
return mask * valid_mask.to(torch.int32)
def generate_dense_grid_points(
bbox_min: np.ndarray,
bbox_max: np.ndarray,
octree_resolution: int,
indexing: str = "ij",
):
length = bbox_max - bbox_min
num_cells = octree_resolution
x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)
y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)
z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)
[xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)
xyz = np.stack((xs, ys, zs), axis=-1)
grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]
return xyz, grid_size, length
class VanillaVolumeDecoder:
@torch.no_grad()
def __call__(
self,
latents: torch.FloatTensor,
geo_decoder: Callable,
bounds: Union[Tuple[float], List[float], float] = 1.01,
num_chunks: int = 10000,
octree_resolution: int = None,
enable_pbar: bool = True,
**kwargs,
):
device = latents.device
dtype = latents.dtype
batch_size = latents.shape[0]
# 1. generate query points
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])
xyz_samples, grid_size, length = generate_dense_grid_points(
bbox_min=bbox_min,
bbox_max=bbox_max,
octree_resolution=octree_resolution,
indexing="ij"
)
xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype).contiguous().reshape(-1, 3)
# 2. latents to 3d volume
batch_logits = []
for start in tqdm(range(0, xyz_samples.shape[0], num_chunks), desc=f"Volume Decoding",
disable=not enable_pbar):
chunk_queries = xyz_samples[start: start + num_chunks, :]
chunk_queries = repeat(chunk_queries, "p c -> b p c", b=batch_size)
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((batch_size, *grid_size)).float()
return grid_logits
class HierarchicalVolumeDecoding:
@torch.no_grad()
def __call__(
self,
latents: torch.FloatTensor,
geo_decoder: Callable,
bounds: Union[Tuple[float], List[float], float] = 1.01,
num_chunks: int = 10000,
mc_level: float = 0.0,
octree_resolution: int = None,
min_resolution: int = 63,
enable_pbar: bool = True,
**kwargs,
):
device = latents.device
dtype = latents.dtype
resolutions = []
if octree_resolution < min_resolution:
resolutions.append(octree_resolution)
while octree_resolution >= min_resolution:
resolutions.append(octree_resolution)
octree_resolution = octree_resolution // 2
resolutions.reverse()
# 1. generate query points
if isinstance(bounds, float):
bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
bbox_min = np.array(bounds[0:3])
bbox_max = np.array(bounds[3:6])
bbox_size = bbox_max - bbox_min
xyz_samples, grid_size, length = generate_dense_grid_points(
bbox_min=bbox_min,
bbox_max=bbox_max,
octree_resolution=resolutions[0],
indexing="ij"
)
dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype)
dilate.weight = torch.nn.Parameter(torch.ones(dilate.weight.shape, dtype=dtype, device=device))
grid_size = np.array(grid_size)
xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype).contiguous().reshape(-1, 3)
# 2. latents to 3d volume
batch_logits = []
batch_size = latents.shape[0]
for start in tqdm(range(0, xyz_samples.shape[0], num_chunks),
desc=f"Hierarchical Volume Decoding [r{resolutions[0] + 1}]"):
queries = xyz_samples[start: start + num_chunks, :]
batch_queries = repeat(queries, "p c -> b p c", b=batch_size)
logits = geo_decoder(queries=batch_queries, latents=latents)
batch_logits.append(logits)
grid_logits = torch.cat(batch_logits, dim=1).view((batch_size, grid_size[0], grid_size[1], grid_size[2]))
for octree_depth_now in resolutions[1:]:
grid_size = np.array([octree_depth_now + 1] * 3)
resolution = bbox_size / octree_depth_now
next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device)
next_logits = torch.full(next_index.shape, -10000., dtype=dtype, device=device)
curr_points = extract_near_surface_volume_fn(grid_logits.squeeze(0), mc_level)
curr_points += grid_logits.squeeze(0).abs() < 0.95
if octree_depth_now == resolutions[-1]:
expand_num = 0
else:
expand_num = 1
for i in range(expand_num):
curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0)
(cidx_x, cidx_y, cidx_z) = torch.where(curr_points > 0)
next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1
for i in range(2 - expand_num):
next_index = dilate(next_index.unsqueeze(0)).squeeze(0)
nidx = torch.where(next_index > 0)
next_points = torch.stack(nidx, dim=1)
next_points = (next_points * torch.tensor(resolution, dtype=next_points.dtype, device=device) +
torch.tensor(bbox_min, dtype=next_points.dtype, device=device))
batch_logits = []
for start in tqdm(range(0, next_points.shape[0], num_chunks),
desc=f"Hierarchical Volume Decoding [r{octree_depth_now + 1}]"):
queries = next_points[start: start + num_chunks, :]
batch_queries = repeat(queries, "p c -> b p c", b=batch_size)
logits = geo_decoder(queries=batch_queries.to(latents.dtype), latents=latents)
batch_logits.append(logits)
grid_logits = torch.cat(batch_logits, dim=1)
next_logits[nidx] = grid_logits[0, ..., 0]
grid_logits = next_logits.unsqueeze(0)
grid_logits[grid_logits == -10000.] = float('nan')
return grid_logits
class FlashVDMVolumeDecoding:
def __init__(self, topk_mode='mean'):
if topk_mode not in ['mean', 'merge']:
raise ValueError(f'Unsupported topk_mode {topk_mode}, available: {["mean", "merge"]}')
if topk_mode == 'mean':
self.processor = FlashVDMCrossAttentionProcessor()
else:
self.processor = FlashVDMTopMCrossAttentionProcessor()
@torch.no_grad()
def __call__(
self,
latents: torch.FloatTensor,
geo_decoder: CrossAttentionDecoder,
bounds: Union[Tuple[float], List[float], float] = 1.01,
num_chunks: int = 10000,
mc_level: float = 0.0,
octree_resolution: int = None,
min_resolution: int = 63,
mini_grid_num: int = 4,
enable_pbar: bool = True,
**kwargs,
):
processor = self.processor
geo_decoder.set_cross_attention_processor(processor)
device = latents.device
dtype = latents.dtype
resolutions = []
if octree_resolution < min_resolution:
resolutions.append(octree_resolution)
while octree_resolution >= min_resolution:
resolutions.append(octree_resolution)
octree_resolution = octree_resolution // 2
resolutions.reverse()
resolutions[0] = round(resolutions[0] / mini_grid_num) * mini_grid_num - 1
for i, resolution in enumerate(resolutions[1:]):
resolutions[i + 1] = resolutions[0] * 2 ** (i + 1)
logger.info(f"FlashVDMVolumeDecoding Resolution: {resolutions}")
# 1. generate query points
if isinstance(bounds, float):
bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
bbox_min = np.array(bounds[0:3])
bbox_max = np.array(bounds[3:6])
bbox_size = bbox_max - bbox_min
xyz_samples, grid_size, length = generate_dense_grid_points(
bbox_min=bbox_min,
bbox_max=bbox_max,
octree_resolution=resolutions[0],
indexing="ij"
)
dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype)
dilate.weight = torch.nn.Parameter(torch.ones(dilate.weight.shape, dtype=dtype, device=device))
grid_size = np.array(grid_size)
# 2. latents to 3d volume
xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype)
batch_size = latents.shape[0]
mini_grid_size = xyz_samples.shape[0] // mini_grid_num
xyz_samples = xyz_samples.view(
mini_grid_num, mini_grid_size,
mini_grid_num, mini_grid_size,
mini_grid_num, mini_grid_size, 3
).permute(
0, 2, 4, 1, 3, 5, 6
).reshape(
-1, mini_grid_size * mini_grid_size * mini_grid_size, 3
)
batch_logits = []
num_batchs = max(num_chunks // xyz_samples.shape[1], 1)
for start in tqdm(range(0, xyz_samples.shape[0], num_batchs),
desc=f"FlashVDM Volume Decoding", disable=not enable_pbar):
queries = xyz_samples[start: start + num_batchs, :]
batch = queries.shape[0]
batch_latents = repeat(latents.squeeze(0), "p c -> b p c", b=batch)
processor.topk = True
logits = geo_decoder(queries=queries, latents=batch_latents)
batch_logits.append(logits)
grid_logits = torch.cat(batch_logits, dim=0).reshape(
mini_grid_num, mini_grid_num, mini_grid_num,
mini_grid_size, mini_grid_size,
mini_grid_size
).permute(0, 3, 1, 4, 2, 5).contiguous().view(
(batch_size, grid_size[0], grid_size[1], grid_size[2])
)
for octree_depth_now in resolutions[1:]:
grid_size = np.array([octree_depth_now + 1] * 3)
resolution = bbox_size / octree_depth_now
next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device)
next_logits = torch.full(next_index.shape, -10000., dtype=dtype, device=device)
curr_points = extract_near_surface_volume_fn(grid_logits.squeeze(0), mc_level)
curr_points += grid_logits.squeeze(0).abs() < 0.95
if octree_depth_now == resolutions[-1]:
expand_num = 0
else:
expand_num = 1
for i in range(expand_num):
curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0)
(cidx_x, cidx_y, cidx_z) = torch.where(curr_points > 0)
next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1
for i in range(2 - expand_num):
next_index = dilate(next_index.unsqueeze(0)).squeeze(0)
nidx = torch.where(next_index > 0)
next_points = torch.stack(nidx, dim=1)
next_points = (next_points * torch.tensor(resolution, dtype=torch.float32, device=device) +
torch.tensor(bbox_min, dtype=torch.float32, device=device))
query_grid_num = 6
min_val = next_points.min(axis=0).values
max_val = next_points.max(axis=0).values
vol_queries_index = (next_points - min_val) / (max_val - min_val) * (query_grid_num - 0.001)
index = torch.floor(vol_queries_index).long()
index = index[..., 0] * (query_grid_num ** 2) + index[..., 1] * query_grid_num + index[..., 2]
index = index.sort()
next_points = next_points[index.indices].unsqueeze(0).contiguous()
unique_values = torch.unique(index.values, return_counts=True)
grid_logits = torch.zeros((next_points.shape[1]), dtype=latents.dtype, device=latents.device)
input_grid = [[], []]
logits_grid_list = []
start_num = 0
sum_num = 0
for grid_index, count in zip(unique_values[0].cpu().tolist(), unique_values[1].cpu().tolist()):
if sum_num + count < num_chunks or sum_num == 0:
sum_num += count
input_grid[0].append(grid_index)
input_grid[1].append(count)
else:
processor.topk = input_grid
logits_grid = geo_decoder(queries=next_points[:, start_num:start_num + sum_num], latents=latents)
start_num = start_num + sum_num
logits_grid_list.append(logits_grid)
input_grid = [[grid_index], [count]]
sum_num = count
if sum_num > 0:
processor.topk = input_grid
logits_grid = geo_decoder(queries=next_points[:, start_num:start_num + sum_num], latents=latents)
logits_grid_list.append(logits_grid)
logits_grid = torch.cat(logits_grid_list, dim=1)
grid_logits[index.indices] = logits_grid.squeeze(0).squeeze(-1)
next_logits[nidx] = grid_logits
grid_logits = next_logits.unsqueeze(0)
grid_logits[grid_logits == -10000.] = float('nan')
return grid_logits