mirror of
https://git.datalinker.icu/kijai/ComfyUI-Hunyuan3DWrapper.git
synced 2026-01-23 09:44:25 +08:00
Windows fixes
This commit is contained in:
kijai
parent
74dd85173f
commit
4dbf46c580
@ -23,6 +23,7 @@
|
||||
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
|
||||
|
||||
import tempfile
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import pymeshlab
|
||||
@ -63,32 +64,70 @@ def remove_floater(mesh: pymeshlab.MeshSet):
|
||||
|
||||
|
||||
def pymeshlab2trimesh(mesh: pymeshlab.MeshSet):
|
||||
with tempfile.NamedTemporaryFile(suffix='.ply', delete=True) as temp_file:
|
||||
mesh.save_current_mesh(temp_file.name)
|
||||
mesh = trimesh.load(temp_file.name)
|
||||
# 检查加载的对象类型
|
||||
if isinstance(mesh, trimesh.Scene):
|
||||
combined_mesh = trimesh.Trimesh()
|
||||
# 如果是Scene，遍历所有的geometry并合并
|
||||
for geom in mesh.geometry.values():
|
||||
combined_mesh = trimesh.util.concatenate([combined_mesh, geom])
|
||||
mesh = combined_mesh
|
||||
return mesh
|
||||
# Create temp directory with explicit permissions
|
||||
temp_dir = os.path.join(os.getcwd(), 'temp')
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
temp_path = os.path.join(temp_dir, 'temp_mesh.ply')
|
||||
|
||||
# Save and load mesh
|
||||
mesh.save_current_mesh(temp_path)
|
||||
loaded_mesh = trimesh.load(temp_path)
|
||||
|
||||
# Check loaded object type
|
||||
if isinstance(loaded_mesh, trimesh.Scene):
|
||||
combined_mesh = trimesh.Trimesh()
|
||||
# If Scene, iterate through all geometries and combine
|
||||
for geom in loaded_mesh.geometry.values():
|
||||
combined_mesh = trimesh.util.concatenate([combined_mesh, geom])
|
||||
loaded_mesh = combined_mesh
|
||||
|
||||
# Cleanup
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
|
||||
return loaded_mesh
|
||||
|
||||
except Exception as e:
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
raise Exception(f"Error in pymeshlab2trimesh: {str(e)}")
|
||||
|
||||
|
||||
def trimesh2pymeshlab(mesh: trimesh.Trimesh):
|
||||
with tempfile.NamedTemporaryFile(suffix='.ply', delete=True) as temp_file:
|
||||
# Create temp directory with explicit permissions
|
||||
temp_dir = os.path.join(os.getcwd(), 'temp')
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
temp_path = os.path.join(temp_dir, 'temp_mesh.ply')
|
||||
|
||||
# Handle scene with multiple geometries
|
||||
if isinstance(mesh, trimesh.scene.Scene):
|
||||
temp_mesh = None
|
||||
for idx, obj in enumerate(mesh.geometry.values()):
|
||||
if idx == 0:
|
||||
temp_mesh = obj
|
||||
else:
|
||||
temp_mesh = temp_mesh + obj
|
||||
mesh = temp_mesh
|
||||
mesh.export(temp_file.name)
|
||||
mesh = pymeshlab.MeshSet()
|
||||
mesh.load_new_mesh(temp_file.name)
|
||||
return mesh
|
||||
|
||||
# Export and load mesh
|
||||
mesh.export(temp_path)
|
||||
mesh_set = pymeshlab.MeshSet()
|
||||
mesh_set.load_new_mesh(temp_path)
|
||||
|
||||
# Cleanup
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
|
||||
return mesh_set
|
||||
|
||||
except Exception as e:
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
raise Exception(f"Error in trimesh2pymeshlab: {str(e)}")
|
||||
|
||||
|
||||
def export_mesh(input, output):
|
||||
@ -148,10 +187,22 @@ class DegenerateFaceRemover:
|
||||
) -> Union[pymeshlab.MeshSet, trimesh.Trimesh, Latent2MeshOutput]:
|
||||
ms = import_mesh(mesh)
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix='.ply', delete=True) as temp_file:
|
||||
ms.save_current_mesh(temp_file.name)
|
||||
# Create temp file with explicit closing
|
||||
temp_file = tempfile.NamedTemporaryFile(suffix='.ply', delete=False)
|
||||
temp_file_path = temp_file.name
|
||||
temp_file.close()
|
||||
|
||||
try:
|
||||
ms.save_current_mesh(temp_file_path)
|
||||
ms = pymeshlab.MeshSet()
|
||||
ms.load_new_mesh(temp_file.name)
|
||||
ms.load_new_mesh(temp_file_path)
|
||||
finally:
|
||||
# Ensure temp file is removed
|
||||
if os.path.exists(temp_file_path):
|
||||
try:
|
||||
os.remove(temp_file_path)
|
||||
except:
|
||||
pass
|
||||
|
||||
mesh = export_mesh(mesh, ms)
|
||||
return mesh
|
||||
|
||||
@ -1,32 +0,0 @@
|
||||
# Open Source Model Licensed under the Apache License Version 2.0
|
||||
# and Other Licenses of the Third-Party Components therein:
|
||||
# The below Model in this distribution may have been modified by THL A29 Limited
|
||||
# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
|
||||
|
||||
# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
|
||||
# The below software and/or models in this distribution may have been
|
||||
# modified by THL A29 Limited ("Tencent Modifications").
|
||||
# All Tencent Modifications are Copyright (C) THL A29 Limited.
|
||||
|
||||
# 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 .hierarchy import BuildHierarchy, BuildHierarchyWithColor
|
||||
from .io_obj import LoadObj, LoadObjWithTexture
|
||||
from .render import rasterize, interpolate
|
||||
'''
|
||||
from .io_glb import *
|
||||
from .io_obj import *
|
||||
from .render import *
|
||||
@ -1,248 +0,0 @@
|
||||
# Open Source Model Licensed under the Apache License Version 2.0
|
||||
# and Other Licenses of the Third-Party Components therein:
|
||||
# The below Model in this distribution may have been modified by THL A29 Limited
|
||||
# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
|
||||
|
||||
# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
|
||||
# The below software and/or models in this distribution may have been
|
||||
# modified by THL A29 Limited ("Tencent Modifications").
|
||||
# All Tencent Modifications are Copyright (C) THL A29 Limited.
|
||||
|
||||
# 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.
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image as PILImage
|
||||
from pygltflib import GLTF2
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
|
||||
# Function to extract buffer data
|
||||
def get_buffer_data(gltf, buffer_view):
|
||||
buffer = gltf.buffers[buffer_view.buffer]
|
||||
buffer_data = gltf.get_data_from_buffer_uri(buffer.uri)
|
||||
byte_offset = buffer_view.byteOffset if buffer_view.byteOffset else 0
|
||||
byte_length = buffer_view.byteLength
|
||||
return buffer_data[byte_offset:byte_offset + byte_length]
|
||||
|
||||
|
||||
# Function to extract attribute data
|
||||
def get_attribute_data(gltf, accessor_index):
|
||||
accessor = gltf.accessors[accessor_index]
|
||||
buffer_view = gltf.bufferViews[accessor.bufferView]
|
||||
buffer_data = get_buffer_data(gltf, buffer_view)
|
||||
|
||||
comptype = {5120: np.int8, 5121: np.uint8, 5122: np.int16, 5123: np.uint16, 5125: np.uint32, 5126: np.float32}
|
||||
dtype = comptype[accessor.componentType]
|
||||
|
||||
t2n = {'SCALAR': 1, 'VEC2': 2, 'VEC3': 3, 'VEC4': 4, 'MAT2': 4, 'MAT3': 9, 'MAT4': 16}
|
||||
num_components = t2n[accessor.type]
|
||||
|
||||
# Calculate the correct slice of data
|
||||
byte_offset = accessor.byteOffset if accessor.byteOffset else 0
|
||||
byte_stride = buffer_view.byteStride if buffer_view.byteStride else num_components * np.dtype(dtype).itemsize
|
||||
count = accessor.count
|
||||
|
||||
# Extract the attribute data
|
||||
attribute_data = np.zeros((count, num_components), dtype=dtype)
|
||||
for i in range(count):
|
||||
start = byte_offset + i * byte_stride
|
||||
end = start + num_components * np.dtype(dtype).itemsize
|
||||
attribute_data[i] = np.frombuffer(buffer_data[start:end], dtype=dtype)
|
||||
|
||||
return attribute_data
|
||||
|
||||
|
||||
# Function to extract image data
|
||||
def get_image_data(gltf, image, folder):
|
||||
if image.uri:
|
||||
if image.uri.startswith('data:'):
|
||||
# Data URI
|
||||
header, encoded = image.uri.split(',', 1)
|
||||
data = base64.b64decode(encoded)
|
||||
else:
|
||||
# External file
|
||||
fn = image.uri
|
||||
if not os.path.isabs(fn):
|
||||
fn = folder + '/' + fn
|
||||
with open(fn, 'rb') as f:
|
||||
data = f.read()
|
||||
else:
|
||||
buffer_view = gltf.bufferViews[image.bufferView]
|
||||
data = get_buffer_data(gltf, buffer_view)
|
||||
return data
|
||||
|
||||
|
||||
# Function to convert triangle strip to triangles
|
||||
def convert_triangle_strip_to_triangles(indices):
|
||||
triangles = []
|
||||
for i in range(len(indices) - 2):
|
||||
if i % 2 == 0:
|
||||
triangles.append([indices[i], indices[i + 1], indices[i + 2]])
|
||||
else:
|
||||
triangles.append([indices[i], indices[i + 2], indices[i + 1]])
|
||||
return np.array(triangles).reshape(-1, 3)
|
||||
|
||||
|
||||
# Function to convert triangle fan to triangles
|
||||
def convert_triangle_fan_to_triangles(indices):
|
||||
triangles = []
|
||||
for i in range(1, len(indices) - 1):
|
||||
triangles.append([indices[0], indices[i], indices[i + 1]])
|
||||
return np.array(triangles).reshape(-1, 3)
|
||||
|
||||
|
||||
# Function to get the transformation matrix from a node
|
||||
def get_node_transform(node):
|
||||
if node.matrix:
|
||||
return np.array(node.matrix).reshape(4, 4).T
|
||||
else:
|
||||
T = np.eye(4)
|
||||
if node.translation:
|
||||
T[:3, 3] = node.translation
|
||||
if node.rotation:
|
||||
R_mat = R.from_quat(node.rotation).as_matrix()
|
||||
T[:3, :3] = R_mat
|
||||
if node.scale:
|
||||
S = np.diag(node.scale + [1])
|
||||
T = T @ S
|
||||
return T
|
||||
|
||||
|
||||
def get_world_transform(gltf, node_index, parents, world_transforms):
|
||||
if parents[node_index] == -2:
|
||||
return world_transforms[node_index]
|
||||
|
||||
node = gltf.nodes[node_index]
|
||||
if parents[node_index] == -1:
|
||||
world_transforms[node_index] = get_node_transform(node)
|
||||
parents[node_index] = -2
|
||||
return world_transforms[node_index]
|
||||
|
||||
parent_index = parents[node_index]
|
||||
parent_transform = get_world_transform(gltf, parent_index, parents, world_transforms)
|
||||
world_transforms[node_index] = parent_transform @ get_node_transform(node)
|
||||
parents[node_index] = -2
|
||||
return world_transforms[node_index]
|
||||
|
||||
|
||||
def LoadGlb(path):
|
||||
# Load the GLB file using pygltflib
|
||||
gltf = GLTF2().load(path)
|
||||
|
||||
primitives = []
|
||||
images = {}
|
||||
# Iterate through the meshes in the GLB file
|
||||
|
||||
world_transforms = [np.identity(4) for i in range(len(gltf.nodes))]
|
||||
parents = [-1 for i in range(len(gltf.nodes))]
|
||||
for node_index, node in enumerate(gltf.nodes):
|
||||
for idx in node.children:
|
||||
parents[idx] = node_index
|
||||
# for i in range(len(gltf.nodes)):
|
||||
# get_world_transform(gltf, i, parents, world_transform)
|
||||
|
||||
for node_index, node in enumerate(gltf.nodes):
|
||||
if node.mesh is not None:
|
||||
world_transform = get_world_transform(gltf, node_index, parents, world_transforms)
|
||||
# Iterate through the primitives in the mesh
|
||||
mesh = gltf.meshes[node.mesh]
|
||||
for primitive in mesh.primitives:
|
||||
# Access the attributes of the primitive
|
||||
attributes = primitive.attributes.__dict__
|
||||
mode = primitive.mode if primitive.mode is not None else 4 # Default to TRIANGLES
|
||||
result = {}
|
||||
if primitive.indices is not None:
|
||||
indices = get_attribute_data(gltf, primitive.indices)
|
||||
if mode == 4: # TRIANGLES
|
||||
face_indices = indices.reshape(-1, 3)
|
||||
elif mode == 5: # TRIANGLE_STRIP
|
||||
face_indices = convert_triangle_strip_to_triangles(indices)
|
||||
elif mode == 6: # TRIANGLE_FAN
|
||||
face_indices = convert_triangle_fan_to_triangles(indices)
|
||||
else:
|
||||
continue
|
||||
result['F'] = face_indices
|
||||
|
||||
# Extract vertex positions
|
||||
if 'POSITION' in attributes and attributes['POSITION'] is not None:
|
||||
positions = get_attribute_data(gltf, attributes['POSITION'])
|
||||
# Apply the world transformation to the positions
|
||||
positions_homogeneous = np.hstack([positions, np.ones((positions.shape[0], 1))])
|
||||
transformed_positions = (world_transform @ positions_homogeneous.T).T[:, :3]
|
||||
result['V'] = transformed_positions
|
||||
|
||||
# Extract vertex colors
|
||||
if 'COLOR_0' in attributes and attributes['COLOR_0'] is not None:
|
||||
colors = get_attribute_data(gltf, attributes['COLOR_0'])
|
||||
if colors.shape[-1] > 3:
|
||||
colors = colors[..., :3]
|
||||
result['VC'] = colors
|
||||
|
||||
# Extract UVs
|
||||
if 'TEXCOORD_0' in attributes and not attributes['TEXCOORD_0'] is None:
|
||||
uvs = get_attribute_data(gltf, attributes['TEXCOORD_0'])
|
||||
result['UV'] = uvs
|
||||
|
||||
if primitive.material is not None:
|
||||
material = gltf.materials[primitive.material]
|
||||
if material.pbrMetallicRoughness is not None and material.pbrMetallicRoughness.baseColorTexture is not None:
|
||||
texture_index = material.pbrMetallicRoughness.baseColorTexture.index
|
||||
texture = gltf.textures[texture_index]
|
||||
image_index = texture.source
|
||||
if not image_index in images:
|
||||
image = gltf.images[image_index]
|
||||
image_data = get_image_data(gltf, image, os.path.dirname(path))
|
||||
pil_image = PILImage.open(io.BytesIO(image_data))
|
||||
if pil_image.mode != 'RGB':
|
||||
pil_image = pil_image.convert('RGB')
|
||||
images[image_index] = pil_image
|
||||
result['TEX'] = image_index
|
||||
elif material.emissiveTexture is not None:
|
||||
texture_index = material.emissiveTexture.index
|
||||
texture = gltf.textures[texture_index]
|
||||
image_index = texture.source
|
||||
if not image_index in images:
|
||||
image = gltf.images[image_index]
|
||||
image_data = get_image_data(gltf, image, os.path.dirname(path))
|
||||
pil_image = PILImage.open(io.BytesIO(image_data))
|
||||
if pil_image.mode != 'RGB':
|
||||
pil_image = pil_image.convert('RGB')
|
||||
images[image_index] = pil_image
|
||||
result['TEX'] = image_index
|
||||
else:
|
||||
if material.pbrMetallicRoughness is not None:
|
||||
base_color = material.pbrMetallicRoughness.baseColorFactor
|
||||
else:
|
||||
base_color = np.array([0.8, 0.8, 0.8], dtype=np.float32)
|
||||
result['MC'] = base_color
|
||||
|
||||
primitives.append(result)
|
||||
|
||||
return primitives, images
|
||||
|
||||
|
||||
def RotatePrimitives(primitives, transform):
|
||||
for i in range(len(primitives)):
|
||||
if 'V' in primitives[i]:
|
||||
primitives[i]['V'] = primitives[i]['V'] @ transform.T
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
path = 'data/test.glb'
|
||||
LoadGlb(path)
|
||||
@ -1,76 +0,0 @@
|
||||
# Open Source Model Licensed under the Apache License Version 2.0
|
||||
# and Other Licenses of the Third-Party Components therein:
|
||||
# The below Model in this distribution may have been modified by THL A29 Limited
|
||||
# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
|
||||
|
||||
# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
|
||||
# The below software and/or models in this distribution may have been
|
||||
# modified by THL A29 Limited ("Tencent Modifications").
|
||||
# All Tencent Modifications are Copyright (C) THL A29 Limited.
|
||||
|
||||
# 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.
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def LoadObj(fn):
|
||||
lines = [l.strip() for l in open(fn)]
|
||||
vertices = []
|
||||
faces = []
|
||||
for l in lines:
|
||||
words = [w for w in l.split(' ') if w != '']
|
||||
if len(words) == 0:
|
||||
continue
|
||||
if words[0] == 'v':
|
||||
v = [float(words[i]) for i in range(1, 4)]
|
||||
vertices.append(v)
|
||||
elif words[0] == 'f':
|
||||
f = [int(words[i]) - 1 for i in range(1, 4)]
|
||||
faces.append(f)
|
||||
|
||||
return np.array(vertices).astype('float32'), np.array(faces).astype('int32')
|
||||
|
||||
|
||||
def LoadObjWithTexture(fn, tex_fn):
|
||||
lines = [l.strip() for l in open(fn)]
|
||||
vertices = []
|
||||
vertex_textures = []
|
||||
faces = []
|
||||
face_textures = []
|
||||
for l in lines:
|
||||
words = [w for w in l.split(' ') if w != '']
|
||||
if len(words) == 0:
|
||||
continue
|
||||
if words[0] == 'v':
|
||||
v = [float(words[i]) for i in range(1, len(words))]
|
||||
vertices.append(v)
|
||||
elif words[0] == 'vt':
|
||||
v = [float(words[i]) for i in range(1, len(words))]
|
||||
vertex_textures.append(v)
|
||||
elif words[0] == 'f':
|
||||
f = []
|
||||
ft = []
|
||||
for i in range(1, len(words)):
|
||||
t = words[i].split('/')
|
||||
f.append(int(t[0]) - 1)
|
||||
ft.append(int(t[1]) - 1)
|
||||
for i in range(2, len(f)):
|
||||
faces.append([f[0], f[i - 1], f[i]])
|
||||
face_textures.append([ft[0], ft[i - 1], ft[i]])
|
||||
|
||||
tex_image = cv2.cvtColor(cv2.imread(tex_fn), cv2.COLOR_BGR2RGB)
|
||||
return np.array(vertices).astype('float32'), np.array(vertex_textures).astype('float32'), np.array(faces).astype(
|
||||
'int32'), np.array(face_textures).astype('int32'), tex_image
|
||||
@ -1,41 +0,0 @@
|
||||
# Open Source Model Licensed under the Apache License Version 2.0
|
||||
# and Other Licenses of the Third-Party Components therein:
|
||||
# The below Model in this distribution may have been modified by THL A29 Limited
|
||||
# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
|
||||
|
||||
# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
|
||||
# The below software and/or models in this distribution may have been
|
||||
# modified by THL A29 Limited ("Tencent Modifications").
|
||||
# All Tencent Modifications are Copyright (C) THL A29 Limited.
|
||||
|
||||
# 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.
|
||||
|
||||
import custom_rasterizer_kernel
|
||||
import torch
|
||||
|
||||
|
||||
def rasterize(pos, tri, resolution, clamp_depth=torch.zeros(0), use_depth_prior=0):
|
||||
assert (pos.device == tri.device)
|
||||
findices, barycentric = custom_rasterizer_kernel.rasterize_image(pos[0], tri, clamp_depth, resolution[1],
|
||||
resolution[0], 1e-6, use_depth_prior)
|
||||
return findices, barycentric
|
||||
|
||||
|
||||
def interpolate(col, findices, barycentric, tri):
|
||||
f = findices - 1 + (findices == 0)
|
||||
vcol = col[0, tri.long()[f.long()]]
|
||||
result = barycentric.view(*barycentric.shape, 1) * vcol
|
||||
result = torch.sum(result, axis=-2)
|
||||
return result.view(1, *result.shape)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@ -12,7 +12,7 @@ void rasterizeTriangleCPU(int idx, float* vt0, float* vt1, float* vt2, int width
|
||||
for (int py = y_min; py < y_max + 1; ++py) {
|
||||
if (py < 0 || py >= height)
|
||||
continue;
|
||||
float vt[2] = {px + 0.5, py + 0.5};
|
||||
float vt[2] = {px + 0.5f, py + 0.5f};
|
||||
float baryCentricCoordinate[3];
|
||||
calculateBarycentricCoordinate(vt0, vt1, vt2, vt, baryCentricCoordinate);
|
||||
if (isBarycentricCoordInBounds(baryCentricCoordinate)) {
|
||||
@ -100,24 +100,24 @@ std::vector<torch::Tensor> rasterize_image_cpu(torch::Tensor V, torch::Tensor F,
|
||||
auto INT64_options = torch::TensorOptions().dtype(torch::kInt64).requires_grad(false);
|
||||
auto findices = torch::zeros({height, width}, options);
|
||||
INT64 maxint = (INT64)MAXINT * (INT64)MAXINT + (MAXINT - 1);
|
||||
auto z_min = torch::ones({height, width}, INT64_options) * (long)maxint;
|
||||
auto z_min = torch::ones({height, width}, INT64_options) * (int64_t)maxint;
|
||||
|
||||
if (!use_depth_prior) {
|
||||
for (int i = 0; i < num_faces; ++i) {
|
||||
rasterizeImagecoordsKernelCPU(V.data_ptr<float>(), F.data_ptr<int>(), 0,
|
||||
(INT64*)z_min.data_ptr<long>(), occlusion_truncation, width, height, num_vertices, num_faces, i);
|
||||
(INT64*)z_min.data_ptr<int64_t>(), occlusion_truncation, width, height, num_vertices, num_faces, i);
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < num_faces; ++i)
|
||||
rasterizeImagecoordsKernelCPU(V.data_ptr<float>(), F.data_ptr<int>(), D.data_ptr<float>(),
|
||||
(INT64*)z_min.data_ptr<long>(), occlusion_truncation, width, height, num_vertices, num_faces, i);
|
||||
(INT64*)z_min.data_ptr<int64_t>(), occlusion_truncation, width, height, num_vertices, num_faces, i);
|
||||
}
|
||||
|
||||
auto float_options = torch::TensorOptions().dtype(torch::kFloat32).requires_grad(false);
|
||||
auto barycentric = torch::zeros({height, width, 3}, float_options);
|
||||
for (int i = 0; i < width * height; ++i)
|
||||
barycentricFromImgcoordCPU(V.data_ptr<float>(), F.data_ptr<int>(),
|
||||
findices.data_ptr<int>(), (INT64*)z_min.data_ptr<long>(), width, height, num_vertices, num_faces, barycentric.data_ptr<float>(), i);
|
||||
findices.data_ptr<int>(), (INT64*)z_min.data_ptr<int64_t>(), width, height, num_vertices, num_faces, barycentric.data_ptr<float>(), i);
|
||||
|
||||
return {findices, barycentric};
|
||||
}
|
||||
|
||||
@ -108,20 +108,20 @@ std::vector<torch::Tensor> rasterize_image_gpu(torch::Tensor V, torch::Tensor F,
|
||||
auto INT64_options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA, device_id).requires_grad(false);
|
||||
auto findices = torch::zeros({height, width}, options);
|
||||
INT64 maxint = (INT64)MAXINT * (INT64)MAXINT + (MAXINT - 1);
|
||||
auto z_min = torch::ones({height, width}, INT64_options) * (long)maxint;
|
||||
auto z_min = torch::ones({height, width}, INT64_options) * (int64_t)maxint;
|
||||
|
||||
if (!use_depth_prior) {
|
||||
rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr<float>(), F.data_ptr<int>(), 0,
|
||||
(INT64*)z_min.data_ptr<long>(), occlusion_truncation, width, height, num_vertices, num_faces);
|
||||
(INT64*)z_min.data_ptr<int64_t>(), occlusion_truncation, width, height, num_vertices, num_faces);
|
||||
} else {
|
||||
rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr<float>(), F.data_ptr<int>(), D.data_ptr<float>(),
|
||||
(INT64*)z_min.data_ptr<long>(), occlusion_truncation, width, height, num_vertices, num_faces);
|
||||
(INT64*)z_min.data_ptr<int64_t>(), occlusion_truncation, width, height, num_vertices, num_faces);
|
||||
}
|
||||
|
||||
auto float_options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, device_id).requires_grad(false);
|
||||
auto barycentric = torch::zeros({height, width, 3}, float_options);
|
||||
barycentricFromImgcoordGPU<<<(width * height + 255)/256, 256>>>(V.data_ptr<float>(), F.data_ptr<int>(),
|
||||
findices.data_ptr<int>(), (INT64*)z_min.data_ptr<long>(), width, height, num_vertices, num_faces, barycentric.data_ptr<float>());
|
||||
findices.data_ptr<int>(), (INT64*)z_min.data_ptr<int64_t>(), width, height, num_vertices, num_faces, barycentric.data_ptr<float>());
|
||||
|
||||
return {findices, barycentric};
|
||||
}
|
||||
|
||||
Binary file not shown.
Loading…
x
Reference in New Issue
Block a user