4090 lines
180 KiB
Python
4090 lines
180 KiB
Python
import os
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchvision.transforms as transforms
|
|
from PIL import Image, ImageSequence, ImageOps
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import json
|
|
import trimesh as Trimesh
|
|
from tqdm import tqdm
|
|
import time
|
|
import shutil
|
|
import uuid
|
|
import triton
|
|
import triton.compiler
|
|
|
|
import folder_paths
|
|
import node_helpers
|
|
import hashlib
|
|
import cv2
|
|
import gc
|
|
import copy
|
|
|
|
import pymeshlab
|
|
|
|
import cumesh as CuMesh
|
|
import o_voxel
|
|
|
|
import meshlib.mrmeshnumpy as mrmeshnumpy
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
|
|
import nvdiffrast.torch as dr
|
|
from flex_gemm.ops.grid_sample import grid_sample_3d
|
|
|
|
import comfy.model_management as mm
|
|
from comfy.utils import load_torch_file, ProgressBar, common_upscale
|
|
import comfy.utils
|
|
|
|
from .trellis2.pipelines import Trellis2ImageTo3DPipeline
|
|
from .trellis2.representations import Mesh, MeshWithVoxel
|
|
from .trellis2.modules.attention import config
|
|
from .trellis2.pipelines import samplers
|
|
from .trellis2.modules.sparse import SparseTensor
|
|
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
comfy_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
|
|
|
BASE_CACHE_DIR = Path(os.path.dirname(os.path.realpath(__file__))) / "triton_caches"
|
|
#os.environ["TRITON_ALWAYS_COMPILE"] = "1"
|
|
#os.environ["TORCHINDUCTOR_FORCE_DISABLE_CACHES"]="1"
|
|
|
|
to_pil = transforms.ToPILImage()
|
|
|
|
class AnyType(str):
|
|
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
|
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
|
|
def rotate_triton_cache():
|
|
"""
|
|
Creates a new cache directory and attempts to clean up old ones.
|
|
"""
|
|
# 1. Create the base directory if it doesn't exist
|
|
BASE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 2. Generate a unique ID for this specific run
|
|
run_id = f"cache_{int(time.time())}_{uuid.uuid4().hex[:6]}"
|
|
new_cache_path = BASE_CACHE_DIR / run_id
|
|
new_cache_path.mkdir()
|
|
|
|
# 3. Point Triton to this NEW empty folder
|
|
# This forces a recompile without needing to delete the locked file immediately
|
|
os.environ["TRITON_CACHE_DIR"] = str(new_cache_path)
|
|
print(f"[TrellisNode] 🔄 Switched to fresh Triton cache: {new_cache_path.name}")
|
|
|
|
# 4. Garbage Collection: Try to delete OLD cache folders
|
|
# We wrap this in a try/except so if Windows locks a file, we just skip it
|
|
# and leave it for the next cleanup cycle.
|
|
cleanup_old_caches(current_active=new_cache_path)
|
|
|
|
def cleanup_old_caches(current_active):
|
|
"""
|
|
Iterates through the cache folder and deletes anything that isn't the current active one.
|
|
If a file is locked by Windows, it silently fails and leaves it for later.
|
|
"""
|
|
for item in BASE_CACHE_DIR.iterdir():
|
|
if item.is_dir() and item != current_active:
|
|
try:
|
|
shutil.rmtree(item)
|
|
print(f"[TrellisNode] 🧹 Cleaned up old cache: {item.name}")
|
|
except OSError:
|
|
# This is expected on Windows! The file is locked.
|
|
# We just ignore it and try again next time the node runs.
|
|
pass
|
|
|
|
def parse_string_to_int_list(number_string):
|
|
"""
|
|
Parses a string containing comma-separated numbers into a list of integers.
|
|
|
|
Args:
|
|
number_string: A string containing comma-separated numbers (e.g., "20000,10000,5000").
|
|
|
|
Returns:
|
|
A list of integers parsed from the input string.
|
|
Returns an empty list if the input string is empty or None.
|
|
"""
|
|
if not number_string:
|
|
return []
|
|
|
|
try:
|
|
# Split the string by comma and convert each part to an integer
|
|
int_list = [int(num.strip()) for num in number_string.split(',')]
|
|
return int_list
|
|
except ValueError as e:
|
|
print(f"Error converting string to integer: {e}. Please ensure all values are valid numbers.")
|
|
return []
|
|
|
|
def reset_cuda():
|
|
# Synchronize to ensure all GPU operations complete
|
|
torch.cuda.synchronize()
|
|
|
|
# Force garbage collection of Python objects
|
|
gc.collect()
|
|
|
|
# Clear PyTorch CUDA cache
|
|
torch.cuda.empty_cache()
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
|
|
|
|
def tensor2pil(image: torch.Tensor) -> Image.Image:
|
|
"""
|
|
Accepts either:
|
|
- (H,W,C)
|
|
- (1,H,W,C)
|
|
Returns a PIL RGB/RGBA image depending on channels.
|
|
"""
|
|
if isinstance(image, torch.Tensor):
|
|
t = image.detach().cpu()
|
|
if t.ndim == 4:
|
|
# Expect (B,H,W,C); allow only B==1 here
|
|
if t.shape[0] != 1:
|
|
raise ValueError(f"tensor2pil expects batch of 1, got batch={t.shape[0]}")
|
|
t = t[0]
|
|
elif t.ndim != 3:
|
|
raise ValueError(f"tensor2pil expects (H,W,C) or (1,H,W,C), got shape={tuple(t.shape)}")
|
|
|
|
arr = (t.numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
return Image.fromarray(arr)
|
|
|
|
raise TypeError(f"tensor2pil expected torch.Tensor, got {type(image)}")
|
|
|
|
def tensor_batch_to_pil_list(images: torch.Tensor, max_views: int = 4) -> list[Image.Image]:
|
|
"""
|
|
Converts a ComfyUI IMAGE tensor (B,H,W,C) into a list of PIL images.
|
|
Caps to max_views for safety.
|
|
"""
|
|
if not isinstance(images, torch.Tensor):
|
|
raise TypeError(f"Expected torch.Tensor for IMAGE, got {type(images)}")
|
|
|
|
if images.ndim == 4:
|
|
b = int(images.shape[0])
|
|
n = min(b, int(max_views))
|
|
return [tensor2pil(images[i:i+1]) for i in range(n)]
|
|
|
|
if images.ndim == 3:
|
|
return [tensor2pil(images)]
|
|
|
|
raise ValueError(f"Unsupported IMAGE tensor shape: {tuple(images.shape)}")
|
|
|
|
def convert_tensor_images_to_pil(images):
|
|
pil_array = []
|
|
|
|
for image in images:
|
|
pil_array.append(tensor2pil(image))
|
|
|
|
return pil_array
|
|
|
|
def simplify_with_meshlib(vertices, faces, target=1000000):
|
|
current_faces_num = len(faces)
|
|
print(f'Current Faces Number: {current_faces_num}')
|
|
|
|
if current_faces_num<target:
|
|
return
|
|
|
|
settings = mrmeshpy.DecimateSettings()
|
|
faces_to_delete = current_faces_num - target
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
settings.packMesh = True
|
|
|
|
print('Generating Meshlib Mesh ...')
|
|
mesh = mrmeshnumpy.meshFromFacesVerts(faces, vertices)
|
|
print('Packing Optimally ...')
|
|
mesh.packOptimally()
|
|
print('Decimating ...')
|
|
mrmeshpy.decimateMesh(mesh, settings)
|
|
|
|
new_vertices = mrmeshnumpy.getNumpyVerts(mesh)
|
|
new_faces = mrmeshnumpy.getNumpyFaces(mesh.topology)
|
|
|
|
print(f"Reduced faces, resulting in {len(new_vertices)} vertices and {len(new_faces)} faces")
|
|
|
|
return new_vertices, new_faces
|
|
|
|
def remove_floater(mesh):
|
|
print('Removing floater ...')
|
|
faces = mesh.faces.cpu().numpy()
|
|
print(f"Current faces: {len(faces)}")
|
|
mesh_set = pymeshlab.MeshSet()
|
|
mesh_pymeshlab = pymeshlab.Mesh(vertex_matrix=mesh.vertices.cpu().numpy(), face_matrix=faces)
|
|
mesh_set.add_mesh(mesh_pymeshlab, "converted_mesh")
|
|
mesh_set = pymeshlab_remove_floater(mesh_set)
|
|
|
|
mesh_pymeshlab = mesh_set.current_mesh()
|
|
|
|
new_faces = mesh_pymeshlab.face_matrix()
|
|
print(f"After removing floater: {len(new_faces)}")
|
|
|
|
new_vertices = torch.from_numpy(mesh_pymeshlab.vertex_matrix()).contiguous().float()
|
|
new_faces = torch.from_numpy(new_faces).contiguous().int()
|
|
|
|
mesh.vertices = new_vertices
|
|
mesh.faces = new_faces
|
|
|
|
return mesh
|
|
|
|
def remove_floater2(vertices, faces):
|
|
print('Removing floater ...')
|
|
#faces = faces.cpu().numpy()
|
|
print(f"Current faces: {len(faces)}")
|
|
mesh_set = pymeshlab.MeshSet()
|
|
mesh_pymeshlab = pymeshlab.Mesh(vertex_matrix=vertices, face_matrix=faces)
|
|
mesh_set.add_mesh(mesh_pymeshlab, "converted_mesh")
|
|
mesh_set = pymeshlab_remove_floater(mesh_set)
|
|
|
|
mesh_pymeshlab = mesh_set.current_mesh()
|
|
|
|
new_faces = mesh_pymeshlab.face_matrix()
|
|
print(f"After removing floater: {len(new_faces)}")
|
|
|
|
new_vertices = mesh_pymeshlab.vertex_matrix()
|
|
|
|
return new_vertices, new_faces
|
|
|
|
def remove_mesh_infinite_vertices(mesh):
|
|
print('Removing infinite vertices ...')
|
|
vertices = mesh.vertices.cpu().numpy()
|
|
faces = mesh.faces.cpu().numpy()
|
|
|
|
trimesh = Trimesh.Trimesh(vertices=vertices,faces=faces)
|
|
print(f"Original vertex count: {len(trimesh.vertices)}")
|
|
|
|
# Remove anything outside a reasonable bounding box
|
|
limit = 1e10
|
|
valid_mask = (np.abs(trimesh.vertices) < limit).all(axis=1)
|
|
|
|
trimesh.update_vertices(valid_mask)
|
|
|
|
# Removing vertices can leave "degenerate" faces or orphan nodes
|
|
trimesh.update_faces(trimesh.nondegenerate_faces())
|
|
trimesh.remove_unreferenced_vertices()
|
|
|
|
print(f"Cleaned vertex count: {len(trimesh.vertices)}")
|
|
|
|
new_vertices = torch.from_numpy(trimesh.vertices).float()
|
|
new_faces = torch.from_numpy(trimesh.faces).int()
|
|
|
|
mesh.vertices = new_vertices
|
|
mesh.faces = new_faces
|
|
|
|
return mesh
|
|
|
|
def pymeshlab_remove_floater(mesh: pymeshlab.MeshSet):
|
|
mesh.apply_filter("compute_selection_by_small_disconnected_components_per_face",
|
|
nbfaceratio=0.005)
|
|
mesh.apply_filter("compute_selection_transfer_face_to_vertex", inclusive=False)
|
|
mesh.apply_filter("meshing_remove_selected_vertices_and_faces")
|
|
return mesh
|
|
|
|
def _batched_unsigned_distance(bvh, positions, batch_size=100000, return_uvw=False):
|
|
"""
|
|
Batch unsigned_distance queries to avoid GPU kernel timeout on large meshes.
|
|
When processing high-resolution textures (e.g., 2048x2048 = ~4M pixels) on complex
|
|
meshes, a single BVH query can cause GPU watchdog timeout. This function splits
|
|
the query into smaller batches.
|
|
Args:
|
|
bvh: The BVH structure from cumesh
|
|
positions: (N, 3) tensor of query positions
|
|
batch_size: Maximum number of queries per batch (default 100K, matching
|
|
the rasterization chunk size used elsewhere in this file)
|
|
return_uvw: Whether to return barycentric coordinates
|
|
Returns:
|
|
Same as bvh.unsigned_distance()
|
|
"""
|
|
import torch
|
|
N = positions.shape[0]
|
|
if N <= batch_size:
|
|
return bvh.unsigned_distance(positions, return_uvw=return_uvw)
|
|
|
|
distances_list = []
|
|
face_id_list = []
|
|
uvw_list = [] if return_uvw else None
|
|
|
|
for i in range(0, N, batch_size):
|
|
end = min(i + batch_size, N)
|
|
d, f, u = bvh.unsigned_distance(positions[i:end], return_uvw=return_uvw)
|
|
distances_list.append(d)
|
|
face_id_list.append(f)
|
|
if return_uvw:
|
|
uvw_list.append(u)
|
|
|
|
return (
|
|
torch.cat(distances_list),
|
|
torch.cat(face_id_list),
|
|
torch.cat(uvw_list) if return_uvw else None
|
|
)
|
|
|
|
class Trellis2LoadModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"modelname": (["microsoft/TRELLIS.2-4B","visualbruno/TRELLIS.2-4B-FP8"],{"default":"microsoft/TRELLIS.2-4B"}),
|
|
"backend": (["flash_attn","xformers","sdpa","flash_attn_3"],{"default":"flash_attn"}),
|
|
"device": (["cpu","cuda"],{"default":"cuda"}),
|
|
"low_vram": ("BOOLEAN",{"default":True}),
|
|
"keep_models_loaded": ("BOOLEAN", {"default":True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRELLIS2PIPELINE", )
|
|
RETURN_NAMES = ("pipeline", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, modelname, backend, device, low_vram, keep_models_loaded):
|
|
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # Can save GPU memory
|
|
#os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
|
|
#os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
|
|
os.environ['ATTN_BACKEND'] = backend
|
|
|
|
config.set_backend(backend)
|
|
|
|
reset_cuda()
|
|
|
|
torch.backends.cudnn.benchmark = False
|
|
|
|
model_path = os.path.join(folder_paths.models_dir, modelname)
|
|
|
|
if not os.path.exists(model_path):
|
|
print(f"Downloading model to: {model_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(
|
|
repo_id=modelname,
|
|
local_dir=model_path,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
dinov3_model_path = os.path.join(folder_paths.models_dir,"facebook","dinov3-vitl16-pretrain-lvd1689m","model.safetensors")
|
|
if not os.path.exists(dinov3_model_path):
|
|
raise Exception("Facebook Dinov3 model not found in models/facebook/dinov3-vitl16-pretrain-lvd1689m folder")
|
|
|
|
trellis_image_large_path = os.path.join(folder_paths.models_dir,"microsoft","TRELLIS-image-large","ckpts","ss_dec_conv3d_16l8_fp16.safetensors")
|
|
if not os.path.exists(trellis_image_large_path):
|
|
print('Trellis-Image-Large ss_dec_conv3d_16l8_fp16 files not found. Trying to download the files from huggingface ...')
|
|
import requests
|
|
url = "https://huggingface.co/microsoft/TRELLIS-image-large/resolve/main/ckpts/ss_dec_conv3d_16l8_fp16.json?download=true"
|
|
filename = os.path.join(folder_paths.models_dir,"microsoft","TRELLIS-image-large","ckpts","ss_dec_conv3d_16l8_fp16.json")
|
|
path = Path(filename)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
with open(filename, "wb") as f:
|
|
f.write(response.content)
|
|
print("Download ss_dec_conv3d_16l8_fp16.json complete!")
|
|
else:
|
|
raise Exception("Cannot download Trellis-Image-Large file ss_dec_conv3d_16l8_fp16.json")
|
|
|
|
url = "https://huggingface.co/microsoft/TRELLIS-image-large/resolve/main/ckpts/ss_dec_conv3d_16l8_fp16.safetensors?download=true"
|
|
filename = os.path.join(folder_paths.models_dir,"microsoft","TRELLIS-image-large","ckpts","ss_dec_conv3d_16l8_fp16.safetensors")
|
|
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
with open(filename, "wb") as f:
|
|
f.write(response.content)
|
|
print("Download ss_dec_conv3d_16l8_fp16.safetensors complete!")
|
|
else:
|
|
raise Exception("Cannot download Trellis-Image-Large file ss_dec_conv3d_16l8_fp16.safetensors")
|
|
|
|
if modelname == "visualbruno/TRELLIS.2-4B-FP8":
|
|
use_fp8 = True
|
|
else:
|
|
use_fp8 = False
|
|
|
|
pipeline = Trellis2ImageTo3DPipeline.from_pretrained(model_path, keep_models_loaded = keep_models_loaded, use_fp8=use_fp8)
|
|
pipeline.low_vram = low_vram
|
|
|
|
if device=="cuda":
|
|
if low_vram:
|
|
pipeline.cuda()
|
|
else:
|
|
pipeline.to(device)
|
|
else:
|
|
pipeline.to(device)
|
|
|
|
return (pipeline,)
|
|
|
|
class Trellis2MeshWithVoxelGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
|
|
"pipeline_type": (["512","1024","1024_cascade","1536_cascade"],{"default":"1024_cascade"}),
|
|
"sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"max_num_tokens": ("INT",{"default":49152,"min":0,"max":999999}),
|
|
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
|
|
"generate_texture_slat": ("BOOLEAN", {"default":True}),
|
|
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
|
|
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", "BVH", )
|
|
RETURN_NAMES = ("mesh", "bvh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image, seed, pipeline_type, sparse_structure_steps, shape_steps, texture_steps, max_num_tokens, max_views, sparse_structure_resolution, generate_texture_slat, use_tiled_decoder, sampler):
|
|
reset_cuda()
|
|
|
|
images = tensor_batch_to_pil_list(image, max_views=max_views)
|
|
image_in = images[0] if len(images) == 1 else images
|
|
|
|
sparse_structure_sampler_params = {"steps":sparse_structure_steps}
|
|
shape_slat_sampler_params = {"steps":shape_steps}
|
|
tex_slat_sampler_params = {"steps":texture_steps}
|
|
|
|
if generate_texture_slat:
|
|
num_steps = 5
|
|
else:
|
|
num_steps = 4
|
|
|
|
pbar = ProgressBar(num_steps)
|
|
|
|
mesh = pipeline.run(image=image_in, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens, sparse_structure_resolution = sparse_structure_resolution, max_views = max_views, generate_texture_slat = generate_texture_slat, use_tiled=use_tiled_decoder, pbar=pbar, sampler=sampler)[0]
|
|
|
|
vertices = mesh.vertices.cuda()
|
|
faces = mesh.faces.cuda()
|
|
|
|
# Build BVH for the current mesh to guide remeshing
|
|
if generate_texture_slat:
|
|
print("Building BVH for current mesh...")
|
|
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
bvh.vertices = vertices.detach().clone()
|
|
bvh.faces = faces.detach().clone()
|
|
else:
|
|
print("Not building BVH : only used for texturing")
|
|
bvh = None
|
|
|
|
return (mesh, bvh,)
|
|
|
|
class Trellis2LoadImageWithTransparency:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
|
files = folder_paths.filter_files_content_types(files, ["image"])
|
|
return {"required":
|
|
{"image": (sorted(files), {"image_upload": True})},
|
|
}
|
|
|
|
CATEGORY = "Trellis2Wrapper"
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", )
|
|
RETURN_NAMES = ("image", "mask", "image_with_alpha")
|
|
FUNCTION = "load_image"
|
|
def load_image(self, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
|
|
img = node_helpers.pillow(Image.open, image_path)
|
|
|
|
output_images = []
|
|
output_masks = []
|
|
output_images_ori = []
|
|
w, h = None, None
|
|
|
|
excluded_formats = ['MPO']
|
|
|
|
for i in ImageSequence.Iterator(img):
|
|
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
|
|
|
output_images_ori.append(pil2tensor(i))
|
|
|
|
if i.mode == 'I':
|
|
i = i.point(lambda i: i * (1 / 255))
|
|
image = i.convert("RGB")
|
|
|
|
if len(output_images) == 0:
|
|
w = image.size[0]
|
|
h = image.size[1]
|
|
|
|
if image.size[0] != w or image.size[1] != h:
|
|
continue
|
|
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
if 'A' in i.getbands():
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
elif i.mode == 'P' and 'transparency' in i.info:
|
|
mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
|
output_images.append(image)
|
|
output_masks.append(mask.unsqueeze(0))
|
|
|
|
if len(output_images) > 1 and img.format not in excluded_formats:
|
|
output_image = torch.cat(output_images, dim=0)
|
|
output_mask = torch.cat(output_masks, dim=0)
|
|
output_image_ori = torch.cat(output_images_ori, dim=0)
|
|
else:
|
|
output_image = output_images[0]
|
|
output_mask = output_masks[0]
|
|
output_image_ori = output_images_ori[0]
|
|
|
|
return (output_image, output_mask, output_image_ori)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
m = hashlib.sha256()
|
|
with open(image_path, 'rb') as f:
|
|
m.update(f.read())
|
|
return m.digest().hex()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, image):
|
|
if not folder_paths.exists_annotated_filepath(image):
|
|
return "Invalid image file: {}".format(image)
|
|
|
|
return True
|
|
|
|
class Trellis2SimplifyMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"target_face_num": ("INT",{"default":1000000,"min":1,"max":30000000}),
|
|
"method": (["Cumesh","Meshlib"],{"default":"Cumesh"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", )
|
|
RETURN_NAMES = ("mesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, target_face_num, method):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
if method=="Cumesh":
|
|
# internal testing future release
|
|
# options = {
|
|
# 'method': 'legacy'
|
|
# }
|
|
mesh_copy.simplify_with_cumesh(target = target_face_num)
|
|
elif method=="Meshlib":
|
|
mesh_copy.simplify_with_meshlib(target = target_face_num)
|
|
else:
|
|
raise Exception("Unknown simplification method")
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2SimplifyTrimesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"target_face_num": ("INT",{"default":1000000,"min":1,"max":30000000}),
|
|
"method": (["Cumesh","Meshlib"],{"default":"Cumesh"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", )
|
|
RETURN_NAMES = ("trimesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, trimesh, target_face_num, method):
|
|
mesh_copy = copy.deepcopy(trimesh)
|
|
if method=="Cumesh":
|
|
# internal testing future release
|
|
# options = {
|
|
# 'options': 'legacy'
|
|
# }
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(torch.from_numpy(mesh_copy.vertices).float().cuda(), torch.from_numpy(mesh_copy.faces).int().cuda())
|
|
cumesh.simplify(target_face_num, verbose=True)
|
|
new_vertices, new_faces = cumesh.read()
|
|
mesh_copy.vertices = new_vertices.cpu().numpy()
|
|
mesh_copy.faces = new_faces.cpu().numpy()
|
|
|
|
del cumesh
|
|
elif method=="Meshlib":
|
|
new_vertices, new_faces = simplify_with_meshlib(mesh_copy.vertices, mesh_copy.faces, target = target_face_num)
|
|
mesh_copy.vertices = new_vertices
|
|
mesh_copy.faces = new_faces
|
|
else:
|
|
raise Exception("Unknown simplification method")
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2ProgressiveSimplify:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"max_edge_length": ("FLOAT",{"default":0.00,"min":0.00,"max":99999.99,"step":0.01}),
|
|
"max_triangle_aspect_ratio": ("FLOAT",{"default":20.00,"min":0.01,"max":99999.99,"step":0.01}),
|
|
"strategy": (["Minimal Error First","Shortest Edge First"],{"default":"Minimal Error First"}),
|
|
"stabilizer": ("FLOAT",{"default":0.000001,"min":0.0,"max":0.999999,"step":0.000001}),
|
|
"touch_near_boundary_edges": ("BOOLEAN",{"default":True}),
|
|
"optimize_vertex_positions": ("BOOLEAN",{"default":True}),
|
|
"angle_based_weights": ("BOOLEAN",{"default":False}),
|
|
},
|
|
"optional": {
|
|
"trimesh": ("TRIMESH",),
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", "MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("trimesh", "mesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, max_edge_length, max_triangle_aspect_ratio, strategy, stabilizer, touch_near_boundary_edges, optimize_vertex_positions, angle_based_weights, trimesh = None, mesh = None):
|
|
if trimesh is not None:
|
|
trimesh = copy.deepcopy(trimesh)
|
|
|
|
vertices = trimesh.vertices
|
|
faces = trimesh.faces
|
|
|
|
vertices, faces = self.simplify(vertices, faces, max_edge_length, max_triangle_aspect_ratio, strategy, stabilizer, touch_near_boundary_edges, optimize_vertex_positions, angle_based_weights)
|
|
trimesh.vertices = vertices
|
|
trimesh.faces = faces
|
|
|
|
if mesh is not None:
|
|
mesh = copy.deepcopy(mesh)
|
|
|
|
vertices = mesh.vertices.cpu().numpy()
|
|
faces = mesh.faces.cpu().numpy()
|
|
|
|
vertices, faces = self.simplify(vertices, faces, max_edge_length, max_triangle_aspect_ratio, strategy, stabilizer, touch_near_boundary_edges, optimize_vertex_positions, angle_based_weights)
|
|
mesh.vertices = torch.from_numpy(vertices).float()
|
|
mesh.faces = torch.from_numpy(faces).int()
|
|
|
|
return (trimesh, mesh)
|
|
|
|
def simplify(self, vertices, faces, max_edge_length, max_triangle_aspect_ratio, strategy, stabilizer, touch_near_boundary_edges, optimize_vertex_positions, angle_based_weights):
|
|
current_faces_num = len(faces)
|
|
print(f'Current Faces Number: {current_faces_num}')
|
|
|
|
settings = mrmeshpy.DecimateSettings()
|
|
|
|
if strategy == "Minimal Error First":
|
|
settings.strategy = mrmeshpy.DecimateStrategy.MinimizeError
|
|
else:
|
|
settings.strategy = mrmeshpy.DecimateStrategy.ShortestEdgeFirst
|
|
|
|
settings.maxTriangleAspectRatio = max_triangle_aspect_ratio
|
|
settings.stabilizer = stabilizer
|
|
settings.touchNearBdEdges = touch_near_boundary_edges
|
|
settings.optimizeVertexPos = optimize_vertex_positions
|
|
settings.angleWeightedDistToPlane = angle_based_weights
|
|
settings.packMesh = True
|
|
|
|
print('Generating Meshlib Mesh ...')
|
|
mesh = mrmeshnumpy.meshFromFacesVerts(faces, vertices)
|
|
|
|
if max_edge_length == 0.0:
|
|
max_edge_length = 2.0
|
|
# for edge_id in mesh.topology.allValidEdges():
|
|
# edge_len = mesh.computeEdgeLen(edge_id)
|
|
# if edge_len > max_edge_length:
|
|
# max_edge_length = edge_len
|
|
# print(f"Calculated Max Edge Length: {max_edge_length}")
|
|
|
|
settings.maxEdgeLen = max_edge_length
|
|
settings.maxError = max_edge_length / 1000
|
|
|
|
print('Packing Optimally ...')
|
|
mesh.packOptimally()
|
|
print('Decimating ...')
|
|
mrmeshpy.decimateMesh(mesh, settings)
|
|
|
|
new_vertices = mrmeshnumpy.getNumpyVerts(mesh)
|
|
new_faces = mrmeshnumpy.getNumpyFaces(mesh.topology)
|
|
|
|
print(f"Reduced faces, resulting in {len(new_vertices)} vertices and {len(new_faces)} faces")
|
|
|
|
del mesh
|
|
gc.collect()
|
|
|
|
return new_vertices, new_faces
|
|
|
|
class Trellis2MeshWithVoxelToTrimesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"reorient_vertices":(["None","90 degrees","-90 degrees"],{"default":"90 degrees"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", )
|
|
RETURN_NAMES = ("trimesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, reorient_vertices):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
vertices_np = mesh_copy.vertices.cpu().numpy()
|
|
|
|
if reorient_vertices == '90 degrees':
|
|
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2], -vertices_np[:, 1]
|
|
elif reorient_vertices == '-90 degrees':
|
|
vertices_np[:, 1], vertices_np[:, 2] = -vertices_np[:, 2], vertices_np[:, 1]
|
|
|
|
trimesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=mesh_copy.faces.cpu().numpy(),
|
|
process=False
|
|
)
|
|
|
|
return (trimesh,)
|
|
|
|
class Trellis2ExportMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"filename_prefix": ("STRING", {"default": "3D/Trellis2"}),
|
|
"file_format": (["glb", "obj", "ply", "stl", "3mf", "dae"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","STRING",)
|
|
RETURN_NAMES = ("glb_path","relative_path",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, trimesh, filename_prefix, file_format):
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
|
|
output_glb_path = Path(full_output_folder, f'{filename}_{counter:05}_.{file_format}')
|
|
output_glb_path.parent.mkdir(exist_ok=True)
|
|
|
|
if file_format=='obj':
|
|
materialName = f"{filename}_{counter:05}_.mtl"
|
|
if hasattr(trimesh, 'visual') and hasattr(trimesh.visual, 'material') and trimesh.visual.material is not None:
|
|
trimesh.visual.material.name = f"{filename}_{counter:05}"
|
|
|
|
trimesh.export(output_glb_path, file_type=file_format, mtl_name=materialName)
|
|
else:
|
|
trimesh.export(output_glb_path, file_type=file_format)
|
|
|
|
relative_path = Path(subfolder) / f'{filename}_{counter:05}_.{file_format}'
|
|
|
|
return (str(output_glb_path), str(relative_path), )
|
|
|
|
class Trellis2PostProcessMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"remove_duplicate_faces": ("BOOLEAN",{"default":False}),
|
|
"repair_non_manifold_edges": ("BOOLEAN", {"default":False}),
|
|
"remove_non_manifold_faces": ("BOOLEAN", {"default":False}),
|
|
"remove_small_connected_components": ("BOOLEAN", {"default":False}),
|
|
"remove_small_connected_components_size": ("FLOAT", {"default":0.00001,"min":0.00001,"max":9.99999,"step":0.00001}),
|
|
"unify_faces_orientation": ("BOOLEAN", {"default":False}),
|
|
"remove_floaters": ("BOOLEAN",{"default":False}),
|
|
"remove_infinite_vertices": ("BOOLEAN",{"default":False}),
|
|
"merge_vertices": ("BOOLEAN",{"default":False}),
|
|
"merge_distance": ("FLOAT",{"default":0.0010,"min":0.0001,"max":999.9999,"step":0.0001}),
|
|
"remove_nan_vertices": ("BOOLEAN",{"default":False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self,
|
|
mesh,
|
|
remove_duplicate_faces,
|
|
repair_non_manifold_edges,
|
|
remove_non_manifold_faces,
|
|
remove_small_connected_components,
|
|
remove_small_connected_components_size,
|
|
unify_faces_orientation,
|
|
remove_floaters,
|
|
remove_infinite_vertices,
|
|
merge_vertices,
|
|
merge_distance,
|
|
remove_nan_vertices):
|
|
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
if remove_floaters:
|
|
mesh_copy = remove_floater(mesh_copy)
|
|
if remove_infinite_vertices:
|
|
mesh_copy = remove_mesh_infinite_vertices(mesh_copy)
|
|
|
|
vertices = mesh_copy.vertices
|
|
faces = mesh_copy.faces
|
|
|
|
# Move data to GPU
|
|
vertices = vertices.cuda()
|
|
faces = faces.cuda()
|
|
|
|
# Initialize CUDA mesh handler
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(vertices, faces)
|
|
print(f"Current vertices: {cumesh.num_vertices}, faces: {cumesh.num_faces}")
|
|
|
|
if remove_duplicate_faces:
|
|
print('Removing duplicate faces ...')
|
|
cumesh.remove_duplicate_faces()
|
|
|
|
if repair_non_manifold_edges:
|
|
print('Repairing non manifold edges ...')
|
|
cumesh.repair_non_manifold_edges()
|
|
|
|
if remove_non_manifold_faces:
|
|
print('Removing non manifold faces ...')
|
|
cumesh.remove_non_manifold_faces()
|
|
|
|
if remove_small_connected_components:
|
|
print('Removing small connected components ...')
|
|
cumesh.remove_small_connected_components(remove_small_connected_components_size)
|
|
|
|
if unify_faces_orientation:
|
|
print('Unifying faces orientation ...')
|
|
cumesh.unify_face_orientations()
|
|
|
|
if merge_vertices or remove_nan_vertices:
|
|
import open3d
|
|
open3d_mesh = open3d.geometry.TriangleMesh()
|
|
open3d_mesh.vertices = open3d.utility.Vector3dVector(vertices.cpu().numpy())
|
|
open3d_mesh.triangles = open3d.utility.Vector3iVector(faces.cpu().numpy().astype(np.int32))
|
|
|
|
# NaN check
|
|
print('Removing NaN vertices ...')
|
|
verts = np.asarray(open3d_mesh.vertices)
|
|
if np.any(np.isnan(verts)) or np.any(np.isinf(verts)):
|
|
print('NaN found. Cleaning them ...')
|
|
verts = np.nan_to_num(verts, nan=0.0, posinf=0.0, neginf=0.0)
|
|
open3d_mesh.vertices = open3d.utility.Vector3dVector(verts)
|
|
open3d_mesh = open3d_mesh.remove_duplicated_vertices()
|
|
open3d_mesh = open3d_mesh.remove_duplicated_triangles()
|
|
open3d_mesh = open3d_mesh.remove_degenerate_triangles()
|
|
open3d_mesh = open3d_mesh.remove_unreferenced_vertices()
|
|
|
|
#bbox = open3d_mesh.get_axis_aligned_bounding_box()
|
|
#max_extent = np.max(bbox.get_extent())
|
|
#safe_merge_distance = max_extent * 0.0005 # More conservative
|
|
#print(f"Auto-calculated merge distance: {safe_merge_distance:.6f}")
|
|
|
|
if merge_vertices:
|
|
# Merge and cleanup
|
|
open3d_mesh = open3d_mesh.merge_close_vertices(merge_distance)
|
|
open3d_mesh = open3d_mesh.remove_duplicated_vertices()
|
|
open3d_mesh = open3d_mesh.remove_duplicated_triangles()
|
|
open3d_mesh = open3d_mesh.remove_degenerate_triangles()
|
|
open3d_mesh = open3d_mesh.remove_unreferenced_vertices()
|
|
|
|
# Proper normal computation sequence
|
|
open3d_mesh.compute_triangle_normals()
|
|
open3d_mesh.compute_vertex_normals()
|
|
open3d_mesh.normalize_normals()
|
|
open3d_mesh.orient_triangles() # Orient based on computed normals
|
|
open3d_mesh.compute_vertex_normals() # Recompute after orientation
|
|
|
|
# Gentler smoothing
|
|
open3d_mesh = open3d_mesh.filter_smooth_taubin(number_of_iterations=3)
|
|
open3d_mesh.compute_vertex_normals() # Final recompute
|
|
|
|
cumesh.init(torch.from_numpy(np.asarray(open3d_mesh.vertices)).cuda().float(), torch.from_numpy(np.asarray(open3d_mesh.triangles)).cuda().int())
|
|
del open3d_mesh
|
|
|
|
print(f"After initial cleanup: {cumesh.num_vertices} vertices, {cumesh.num_faces} faces")
|
|
|
|
new_vertices, new_faces = cumesh.read()
|
|
|
|
mesh_copy.vertices = new_vertices.to(mesh_copy.device)
|
|
mesh_copy.faces = new_faces.to(mesh_copy.device)
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2UnWrapAndRasterizer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}),
|
|
"mesh_cluster_refine_iterations": ("INT",{"default":0}),
|
|
"mesh_cluster_global_iterations": ("INT",{"default":1}),
|
|
"mesh_cluster_smooth_strength": ("INT",{"default":1}),
|
|
"texture_size": ("INT",{"default":4096, "min":512, "max":16384}),
|
|
"texture_alpha_mode": (["OPAQUE","MASK","BLEND"],{"default":"OPAQUE"}),
|
|
"double_side_material": ("BOOLEAN",{"default":False}),
|
|
"bake_on_vertices": ("BOOLEAN",{"default":False}),
|
|
"use_custom_normals": ("BOOLEAN",{"default":False}),
|
|
"bvh": ("BVH",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH","IMAGE","IMAGE",)
|
|
RETURN_NAMES = ("trimesh","base_color_texture", "metallic_roughness_texture",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, mesh_cluster_threshold_cone_half_angle_rad, mesh_cluster_refine_iterations, mesh_cluster_global_iterations, mesh_cluster_smooth_strength, texture_size, texture_alpha_mode, double_side_material, bake_on_vertices,use_custom_normals,bvh):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
aabb = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]
|
|
|
|
vertices = mesh_copy.vertices
|
|
faces = mesh_copy.faces
|
|
attr_volume = mesh_copy.attrs
|
|
coords = mesh_copy.coords
|
|
attr_layout = mesh_copy.layout
|
|
voxel_size = mesh_copy.voxel_size
|
|
|
|
mesh_cluster_threshold_cone_half_angle_rad = np.radians(mesh_cluster_threshold_cone_half_angle_rad)
|
|
|
|
# --- Input Normalization (AABB, Voxel Size, Grid Size) ---
|
|
if isinstance(aabb, (list, tuple)):
|
|
aabb = np.array(aabb)
|
|
if isinstance(aabb, np.ndarray):
|
|
aabb = torch.tensor(aabb, dtype=torch.float32, device=coords.device)
|
|
|
|
# Calculate grid dimensions based on AABB and voxel size
|
|
if voxel_size is not None:
|
|
if isinstance(voxel_size, float):
|
|
voxel_size = [voxel_size, voxel_size, voxel_size]
|
|
if isinstance(voxel_size, (list, tuple)):
|
|
voxel_size = np.array(voxel_size)
|
|
if isinstance(voxel_size, np.ndarray):
|
|
voxel_size = torch.tensor(voxel_size, dtype=torch.float32, device=coords.device)
|
|
grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int()
|
|
else:
|
|
if isinstance(grid_size, int):
|
|
grid_size = [grid_size, grid_size, grid_size]
|
|
if isinstance(grid_size, (list, tuple)):
|
|
grid_size = np.array(grid_size)
|
|
if isinstance(grid_size, np.ndarray):
|
|
grid_size = torch.tensor(grid_size, dtype=torch.int32, device=coords.device)
|
|
voxel_size = (aabb[1] - aabb[0]) / grid_size
|
|
|
|
print(f"Original mesh: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
|
|
|
|
vertices = vertices.cuda()
|
|
faces = faces.cuda()
|
|
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(vertices, faces)
|
|
|
|
# Build BVH for the current mesh to guide remeshing
|
|
# if bvh == None:
|
|
# print(f"Building BVH for current mesh...")
|
|
# bvh = CuMesh.cuBVH(vertices, faces)
|
|
# bvh.vertices = vertices
|
|
# bvh.faces = faces
|
|
|
|
# --- Branch: Bake On Vertices (skip UV unwrapping and texture creation) ---
|
|
if bake_on_vertices:
|
|
print('Baking colors on vertices...')
|
|
out_vertices, out_faces = cumesh.read()
|
|
out_vertices = out_vertices.cuda()
|
|
out_faces = out_faces.cuda()
|
|
cumesh.compute_vertex_normals()
|
|
out_normals = cumesh.read_vertex_normals()
|
|
|
|
# Sample attributes directly at vertex positions from the voxel grid
|
|
# No BVH mapping needed - the voxel grid contains all the color information
|
|
vertex_attrs = grid_sample_3d(
|
|
attr_volume,
|
|
torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1),
|
|
shape=torch.Size([1, attr_volume.shape[1], *grid_size.tolist()]),
|
|
grid=((out_vertices - aabb[0]) / voxel_size).reshape(1, -1, 3),
|
|
mode='trilinear',
|
|
)
|
|
|
|
# Extract base color and alpha per vertex (vertex_attrs shape: N_vertices x C)
|
|
base_color_idx = attr_layout['base_color']
|
|
alpha_idx = attr_layout['alpha']
|
|
|
|
# Get RGB values and squeeze any extra dimensions to get (N, 3)
|
|
vertex_colors_rgb = vertex_attrs[..., base_color_idx].cpu().numpy()
|
|
vertex_colors_rgb = np.squeeze(vertex_colors_rgb) # Remove batch dims if any
|
|
if vertex_colors_rgb.ndim == 1:
|
|
vertex_colors_rgb = vertex_colors_rgb[None, :] # Ensure at least 2D
|
|
vertex_colors_rgb = np.clip(vertex_colors_rgb * 255, 0, 255).astype(np.uint8)
|
|
|
|
# Handle alpha based on texture_alpha_mode
|
|
if texture_alpha_mode == "OPAQUE":
|
|
# For OPAQUE mode, use full alpha (255)
|
|
vertex_alpha = np.full((vertex_colors_rgb.shape[0], 1), 255, dtype=np.uint8)
|
|
else:
|
|
vertex_alpha = vertex_attrs[..., alpha_idx].cpu().numpy()
|
|
vertex_alpha = np.squeeze(vertex_alpha) # Remove batch dims if any
|
|
vertex_alpha = np.clip(vertex_alpha * 255, 0, 255).astype(np.uint8)
|
|
# Ensure alpha is 2D with shape (N, 1)
|
|
if vertex_alpha.ndim == 1:
|
|
vertex_alpha = vertex_alpha[:, None]
|
|
|
|
# Combine into RGBA
|
|
vertex_colors_rgba = np.concatenate([vertex_colors_rgb, vertex_alpha], axis=-1)
|
|
|
|
print("Finalizing mesh with vertex colors...")
|
|
|
|
vertices_np = out_vertices.cpu().numpy()
|
|
faces_np = out_faces.cpu().numpy()
|
|
normals_np = out_normals.cpu().numpy()
|
|
|
|
# Swap Y and Z axes, invert Y (common conversion for GLB compatibility)
|
|
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2].copy(), -vertices_np[:, 1].copy()
|
|
normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2].copy(), -normals_np[:, 1].copy()
|
|
|
|
# Create mesh with vertex colors using ColorVisuals
|
|
if use_custom_normals:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
vertex_normals=normals_np,
|
|
vertex_colors=vertex_colors_rgba,
|
|
process=False,
|
|
)
|
|
else:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
vertex_colors=vertex_colors_rgba,
|
|
process=False,
|
|
)
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
# Return empty placeholder textures for vertex color mode
|
|
placeholder_texture = pil2tensor(Image.new('RGBA', (1, 1), (0, 0, 0, 0)))
|
|
return (textured_mesh, placeholder_texture, placeholder_texture,)
|
|
|
|
print('Unwrapping ...')
|
|
out_vertices, out_faces, out_uvs, out_vmaps = cumesh.uv_unwrap(
|
|
compute_charts_kwargs={
|
|
"threshold_cone_half_angle_rad": mesh_cluster_threshold_cone_half_angle_rad,
|
|
"refine_iterations": mesh_cluster_refine_iterations,
|
|
"global_iterations": mesh_cluster_global_iterations,
|
|
"smooth_strength": mesh_cluster_smooth_strength,
|
|
},
|
|
return_vmaps=True,
|
|
verbose=True,
|
|
)
|
|
|
|
out_vertices = out_vertices.cuda()
|
|
out_faces = out_faces.cuda()
|
|
out_uvs = out_uvs.cuda()
|
|
out_vmaps = out_vmaps.cuda()
|
|
cumesh.compute_vertex_normals()
|
|
out_normals = cumesh.read_vertex_normals()[out_vmaps]
|
|
|
|
print("Sampling attributes...")
|
|
# Setup differentiable rasterizer context
|
|
ctx = dr.RasterizeCudaContext()
|
|
# Prepare UV coordinates for rasterization (rendering in UV space)
|
|
uvs_rast = torch.cat([out_uvs * 2 - 1, torch.zeros_like(out_uvs[:, :1]), torch.ones_like(out_uvs[:, :1])], dim=-1).unsqueeze(0)
|
|
rast = torch.zeros((1, texture_size, texture_size, 4), device='cuda', dtype=torch.float32)
|
|
|
|
# Rasterize in chunks to save memory
|
|
for i in range(0, out_faces.shape[0], 100000):
|
|
rast_chunk, _ = dr.rasterize(
|
|
ctx, uvs_rast, out_faces[i:i+100000],
|
|
resolution=[texture_size, texture_size],
|
|
)
|
|
mask_chunk = rast_chunk[..., 3:4] > 0
|
|
rast_chunk[..., 3:4] += i # Store face ID in alpha channel
|
|
rast = torch.where(mask_chunk, rast_chunk, rast)
|
|
|
|
# Mask of valid pixels in texture
|
|
mask = rast[0, ..., 3] > 0
|
|
|
|
# Interpolate 3D positions in UV space (finding 3D coord for every texel)
|
|
pos = dr.interpolate(out_vertices.unsqueeze(0), rast, out_faces)[0][0]
|
|
valid_pos = pos[mask]
|
|
|
|
# Map these positions back to the *original* high-res mesh to get accurate attributes
|
|
# This corrects geometric errors introduced by simplification/remeshing
|
|
_, face_id, uvw = bvh.unsigned_distance(valid_pos, return_uvw=True)
|
|
orig_tri_verts = bvh.vertices[bvh.faces[face_id.long()]] # (N_new, 3, 3)
|
|
valid_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1)
|
|
|
|
torch.cuda.synchronize()
|
|
|
|
# Trilinear sampling from the attribute volume (Color, Material props)
|
|
attrs = torch.zeros(texture_size, texture_size, attr_volume.shape[1], device='cuda')
|
|
attrs[mask] = grid_sample_3d(
|
|
attr_volume,
|
|
torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1),
|
|
shape=torch.Size([1, attr_volume.shape[1], *grid_size.tolist()]),
|
|
grid=((valid_pos - aabb[0]) / voxel_size).reshape(1, -1, 3),
|
|
mode='trilinear',
|
|
)
|
|
|
|
# --- Texture Post-Processing & Material Construction ---
|
|
print("Finalizing mesh...")
|
|
|
|
mask = mask.cpu().numpy()
|
|
|
|
# Extract channels based on layout (BaseColor, Metallic, Roughness, Alpha)
|
|
base_color = np.clip(attrs[..., attr_layout['base_color']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
metallic = np.clip(attrs[..., attr_layout['metallic']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
roughness = np.clip(attrs[..., attr_layout['roughness']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
alpha = np.clip(attrs[..., attr_layout['alpha']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
alpha_mode = texture_alpha_mode
|
|
|
|
# Inpainting: fill gaps (dilation) to prevent black seams at UV boundaries
|
|
mask_inv = (~mask).astype(np.uint8)
|
|
base_color = cv2.inpaint(base_color, mask_inv, 3, cv2.INPAINT_NS)
|
|
metallic = cv2.inpaint(metallic, mask_inv, 1, cv2.INPAINT_NS)[..., None]
|
|
roughness = cv2.inpaint(roughness, mask_inv, 1, cv2.INPAINT_NS)[..., None]
|
|
alpha = cv2.inpaint(alpha, mask_inv, 1, cv2.INPAINT_NS)[..., None]
|
|
|
|
# Create PBR material
|
|
# Standard PBR packs Metallic and Roughness into Blue and Green channels
|
|
baseColorTexture_np = Image.fromarray(np.concatenate([base_color, alpha], axis=-1))
|
|
metallicRoughnessTexture_np = Image.fromarray(np.concatenate([np.zeros_like(metallic), roughness, metallic], axis=-1))
|
|
|
|
material = Trimesh.visual.material.PBRMaterial(
|
|
baseColorTexture=baseColorTexture_np,
|
|
baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8),
|
|
metallicRoughnessTexture=metallicRoughnessTexture_np,
|
|
metallicFactor=1.0,
|
|
roughnessFactor=1.0,
|
|
alphaMode=alpha_mode,
|
|
doubleSided=double_side_material,
|
|
)
|
|
|
|
vertices_np = out_vertices.cpu().numpy()
|
|
faces_np = out_faces.cpu().numpy()
|
|
uvs_np = out_uvs.cpu().numpy()
|
|
normals_np = out_normals.cpu().numpy()
|
|
|
|
# Swap Y and Z axes, invert Y (common conversion for GLB compatibility)
|
|
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2], -vertices_np[:, 1]
|
|
normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2], -normals_np[:, 1]
|
|
uvs_np[:, 1] = 1 - uvs_np[:, 1] # Flip UV V-coordinate
|
|
|
|
if use_custom_normals:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
vertex_normals=normals_np,
|
|
process=False,
|
|
visual=Trimesh.visual.TextureVisuals(uv=uvs_np,material=material)
|
|
)
|
|
else:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
process=False,
|
|
visual=Trimesh.visual.TextureVisuals(uv=uvs_np,material=material)
|
|
)
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
baseColorTexture = pil2tensor(baseColorTexture_np)
|
|
metallicRoughnessTexture = pil2tensor(metallicRoughnessTexture_np)
|
|
|
|
return (textured_mesh, baseColorTexture, metallicRoughnessTexture, )
|
|
|
|
class Trellis2MeshWithVoxelAdvancedGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7fffffff}),
|
|
"pipeline_type": (["512","1024","1024_cascade","1536_cascade"],{"default":"1024_cascade"}),
|
|
"sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"sparse_structure_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"sparse_structure_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":3.00,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.20,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":3.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"max_num_tokens": ("INT",{"default":999999,"min":0,"max":999999}),
|
|
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
|
|
"generate_texture_slat": ("BOOLEAN", {"default":True}),
|
|
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
|
|
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL","BVH", )
|
|
RETURN_NAMES = ("mesh", "bvh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image, seed, pipeline_type, sparse_structure_steps,
|
|
sparse_structure_guidance_strength,
|
|
sparse_structure_guidance_rescale,
|
|
sparse_structure_rescale_t,
|
|
shape_steps,
|
|
shape_guidance_strength,
|
|
shape_guidance_rescale,
|
|
shape_rescale_t,
|
|
texture_steps,
|
|
texture_guidance_strength,
|
|
texture_guidance_rescale,
|
|
texture_rescale_t,
|
|
max_num_tokens,
|
|
max_views,
|
|
sparse_structure_resolution,
|
|
generate_texture_slat,
|
|
sparse_structure_guidance_interval_start,
|
|
sparse_structure_guidance_interval_end,
|
|
shape_guidance_interval_start,
|
|
shape_guidance_interval_end,
|
|
texture_guidance_interval_start,
|
|
texture_guidance_interval_end,
|
|
use_tiled_decoder,
|
|
sampler
|
|
):
|
|
reset_cuda()
|
|
|
|
images = tensor_batch_to_pil_list(image, max_views=max_views)
|
|
image_in = images[0] if len(images) == 1 else images
|
|
|
|
sparse_structure_guidance_interval = [sparse_structure_guidance_interval_start,sparse_structure_guidance_interval_end]
|
|
shape_guidance_interval = [shape_guidance_interval_start,shape_guidance_interval_end]
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
|
|
sparse_structure_sampler_params = {"steps":sparse_structure_steps,"guidance_strength":sparse_structure_guidance_strength,"guidance_rescale":sparse_structure_guidance_rescale,"guidance_interval":sparse_structure_guidance_interval,"rescale_t":sparse_structure_rescale_t}
|
|
shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t}
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
if generate_texture_slat:
|
|
num_steps = 5
|
|
else:
|
|
num_steps = 4
|
|
|
|
pbar = ProgressBar(num_steps)
|
|
|
|
mesh = pipeline.run(image=image_in, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens, sparse_structure_resolution = sparse_structure_resolution, max_views = max_views, generate_texture_slat=generate_texture_slat, use_tiled=use_tiled_decoder, pbar=pbar, sampler=sampler)[0]
|
|
|
|
vertices = mesh.vertices.cuda()
|
|
faces = mesh.faces.cuda()
|
|
|
|
if generate_texture_slat:
|
|
# Build BVH for the current mesh to guide remeshing
|
|
print("Building BVH for current mesh...")
|
|
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
bvh.vertices = vertices.detach().clone()
|
|
bvh.faces = faces.detach().clone()
|
|
else:
|
|
print("Not building BVH : only used for texturing")
|
|
bvh = None
|
|
|
|
return (mesh,bvh,)
|
|
|
|
class Trellis2MeshWithVoxelMultiViewGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"front_image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7fffffff}),
|
|
"pipeline_type": (["512","1024","1024_cascade","1536_cascade"],{"default":"1024_cascade"}),
|
|
"sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"sparse_structure_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"sparse_structure_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":3.00,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.20,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":3.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"max_num_tokens": ("INT",{"default":999999,"min":0,"max":999999}),
|
|
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
|
|
"generate_texture_slat": ("BOOLEAN", {"default":True}),
|
|
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
|
|
"front_axis": (["z", "x"], {"default": "z"}),
|
|
"blend_temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
},
|
|
"optional": {
|
|
"back_image": ("IMAGE",),
|
|
"left_image": ("IMAGE",),
|
|
"right_image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL","BVH", )
|
|
RETURN_NAMES = ("mesh", "bvh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, front_image, seed, pipeline_type, sparse_structure_steps,
|
|
sparse_structure_guidance_strength,
|
|
sparse_structure_guidance_rescale,
|
|
sparse_structure_rescale_t,
|
|
shape_steps,
|
|
shape_guidance_strength,
|
|
shape_guidance_rescale,
|
|
shape_rescale_t,
|
|
texture_steps,
|
|
texture_guidance_strength,
|
|
texture_guidance_rescale,
|
|
texture_rescale_t,
|
|
max_num_tokens,
|
|
sparse_structure_resolution,
|
|
generate_texture_slat,
|
|
sparse_structure_guidance_interval_start,
|
|
sparse_structure_guidance_interval_end,
|
|
shape_guidance_interval_start,
|
|
shape_guidance_interval_end,
|
|
texture_guidance_interval_start,
|
|
texture_guidance_interval_end,
|
|
use_tiled_decoder,
|
|
front_axis,
|
|
blend_temperature,
|
|
sampler,
|
|
back_image=None,
|
|
left_image=None,
|
|
right_image=None):
|
|
|
|
reset_cuda()
|
|
|
|
# Convert front image tensor to PIL
|
|
front_pil = tensor2pil(front_image)
|
|
|
|
# Convert optional view image tensors to PIL
|
|
back_pil = tensor2pil(back_image) if back_image is not None else None
|
|
left_pil = tensor2pil(left_image) if left_image is not None else None
|
|
right_pil = tensor2pil(right_image) if right_image is not None else None
|
|
|
|
sparse_structure_guidance_interval = [sparse_structure_guidance_interval_start,sparse_structure_guidance_interval_end]
|
|
shape_guidance_interval = [shape_guidance_interval_start,shape_guidance_interval_end]
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
|
|
sparse_structure_sampler_params = {"steps":sparse_structure_steps,"guidance_strength":sparse_structure_guidance_strength,"guidance_rescale":sparse_structure_guidance_rescale,"guidance_interval":sparse_structure_guidance_interval,"rescale_t":sparse_structure_rescale_t}
|
|
shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t}
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
if generate_texture_slat:
|
|
num_steps = 5
|
|
else:
|
|
num_steps = 4
|
|
|
|
pbar = ProgressBar(num_steps)
|
|
|
|
mesh = pipeline.run_multiview(
|
|
front=front_pil,
|
|
back=back_pil,
|
|
left=left_pil,
|
|
right=right_pil,
|
|
seed=seed,
|
|
pipeline_type=pipeline_type,
|
|
sparse_structure_sampler_params=sparse_structure_sampler_params,
|
|
shape_slat_sampler_params=shape_slat_sampler_params,
|
|
tex_slat_sampler_params=tex_slat_sampler_params,
|
|
max_num_tokens=max_num_tokens,
|
|
sparse_structure_resolution=sparse_structure_resolution,
|
|
generate_texture_slat=generate_texture_slat,
|
|
use_tiled=use_tiled_decoder,
|
|
pbar=pbar,
|
|
front_axis=front_axis,
|
|
blend_temperature=blend_temperature,
|
|
sampler=sampler
|
|
)[0]
|
|
|
|
vertices = mesh.vertices.cuda()
|
|
faces = mesh.faces.cuda()
|
|
|
|
if generate_texture_slat:
|
|
# Build BVH for the current mesh to guide remeshing
|
|
print("Building BVH for current mesh...")
|
|
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
bvh.vertices = vertices.detach().clone()
|
|
bvh.faces = faces.detach().clone()
|
|
else:
|
|
print("Not building BVH : only used for texturing")
|
|
bvh = None
|
|
|
|
return (mesh,bvh,)
|
|
|
|
class Trellis2PostProcessAndUnWrapAndRasterizer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}),
|
|
"mesh_cluster_refine_iterations": ("INT",{"default":0}),
|
|
"mesh_cluster_global_iterations": ("INT",{"default":1}),
|
|
"mesh_cluster_smooth_strength": ("INT",{"default":1}),
|
|
"texture_size": ("INT",{"default":4096, "min":512, "max":16384}),
|
|
"remesh": ("BOOLEAN",{"default":True}),
|
|
"remesh_band": ("FLOAT",{"default":1.0}),
|
|
"remesh_project": ("FLOAT",{"default":0.0}),
|
|
"target_face_num": ("INT",{"default":2000000,"min":1,"max":16000000}),
|
|
"simplify_method": (["Cumesh","Meshlib"],{"default":"Cumesh"}),
|
|
"fill_holes": ("BOOLEAN", {"default":True}),
|
|
"texture_alpha_mode": (["OPAQUE","MASK","BLEND"],{"default":"OPAQUE"}),
|
|
"dual_contouring_resolution": (["Auto","128","256","512","1024","2048"],{"default":"1024"}),
|
|
"double_side_material": ("BOOLEAN",{"default":False}),
|
|
"remove_floaters": ("BOOLEAN",{"default":True}),
|
|
"bake_on_vertices": ("BOOLEAN",{"default":False}),
|
|
"use_custom_normals":("BOOLEAN",{"default":False}),
|
|
"bvh": ("BVH",),
|
|
"remove_inner_faces": ("BOOLEAN",{"default":True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH","IMAGE","IMAGE",)
|
|
RETURN_NAMES = ("trimesh","base_color_texture","metallic_roughness_texture",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, mesh_cluster_threshold_cone_half_angle_rad, mesh_cluster_refine_iterations, mesh_cluster_global_iterations, mesh_cluster_smooth_strength, texture_size, remesh, remesh_band, remesh_project, target_face_num, simplify_method, fill_holes, texture_alpha_mode, dual_contouring_resolution, double_side_material, remove_floaters, bake_on_vertices,use_custom_normals,bvh,remove_inner_faces):
|
|
pbar = ProgressBar(5 if not bake_on_vertices else 4)
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
aabb = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]
|
|
|
|
attr_volume = mesh_copy.attrs
|
|
coords = mesh_copy.coords
|
|
attr_layout = mesh_copy.layout
|
|
voxel_size = mesh_copy.voxel_size
|
|
|
|
mesh_cluster_threshold_cone_half_angle_rad = np.radians(mesh_cluster_threshold_cone_half_angle_rad)
|
|
|
|
# --- Input Normalization (AABB, Voxel Size, Grid Size) ---
|
|
if isinstance(aabb, (list, tuple)):
|
|
aabb = np.array(aabb)
|
|
if isinstance(aabb, np.ndarray):
|
|
aabb = torch.tensor(aabb, dtype=torch.float32, device=coords.device)
|
|
|
|
# Calculate grid dimensions based on AABB and voxel size
|
|
if voxel_size is not None:
|
|
if isinstance(voxel_size, float):
|
|
voxel_size = [voxel_size, voxel_size, voxel_size]
|
|
if isinstance(voxel_size, (list, tuple)):
|
|
voxel_size = np.array(voxel_size)
|
|
if isinstance(voxel_size, np.ndarray):
|
|
voxel_size = torch.tensor(voxel_size, dtype=torch.float32, device=coords.device)
|
|
grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int()
|
|
else:
|
|
if isinstance(grid_size, int):
|
|
grid_size = [grid_size, grid_size, grid_size]
|
|
if isinstance(grid_size, (list, tuple)):
|
|
grid_size = np.array(grid_size)
|
|
if isinstance(grid_size, np.ndarray):
|
|
grid_size = torch.tensor(grid_size, dtype=torch.int32, device=coords.device)
|
|
voxel_size = (aabb[1] - aabb[0]) / grid_size
|
|
|
|
if remove_floaters:
|
|
mesh_copy = remove_floater(mesh_copy)
|
|
|
|
vertices = mesh_copy.vertices
|
|
faces = mesh_copy.faces
|
|
|
|
vertices = vertices.cuda()
|
|
faces = faces.cuda()
|
|
|
|
# Initialize CUDA mesh handler
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(vertices, faces)
|
|
print(f"Current vertices: {cumesh.num_vertices}, faces: {cumesh.num_faces}")
|
|
|
|
pbar.update(1)
|
|
|
|
print("Cleaning mesh...")
|
|
# --- Branch 1: Standard Pipeline (Simplification & Cleaning) ---
|
|
if not remesh:
|
|
if simplify_method == 'Cumesh':
|
|
cumesh.simplify(target_face_num * 3, verbose=True)
|
|
elif simplify_method == 'Meshlib':
|
|
# GPU -> CPU -> Meshlib -> CPU -> GPU
|
|
v, f = cumesh.read()
|
|
new_vertices, new_faces = simplify_with_meshlib(v.cpu().numpy(), f.cpu().numpy(), target_face_num)
|
|
cumesh.init(torch.from_numpy(new_vertices).float().cuda(), torch.from_numpy(new_faces).int().cuda())
|
|
|
|
cumesh.remove_duplicate_faces()
|
|
cumesh.repair_non_manifold_edges()
|
|
cumesh.remove_small_connected_components(1e-5)
|
|
|
|
if simplify_method == 'Cumesh':
|
|
cumesh.simplify(target_face_num, verbose=True)
|
|
elif simplify_method == 'Meshlib':
|
|
# GPU -> CPU -> Meshlib -> CPU -> GPU
|
|
v, f = cumesh.read()
|
|
new_vertices, new_faces = simplify_with_meshlib(v.cpu().numpy(), f.cpu().numpy(), target_face_num)
|
|
cumesh.init(torch.from_numpy(new_vertices).float().cuda(), torch.from_numpy(new_faces).int().cuda())
|
|
|
|
cumesh.remove_duplicate_faces()
|
|
cumesh.repair_non_manifold_edges()
|
|
cumesh.remove_small_connected_components(1e-5)
|
|
|
|
print(f"After initial cleanup: {cumesh.num_vertices} vertices, {cumesh.num_faces} faces")
|
|
|
|
# Step 2: Unify face orientations
|
|
print('Unifying faces orientation ...')
|
|
cumesh.unify_face_orientations()
|
|
|
|
# --- Branch 2: Remeshing Pipeline ---
|
|
else:
|
|
center = aabb.mean(dim=0)
|
|
scale = (aabb[1] - aabb[0]).max().item()
|
|
|
|
if dual_contouring_resolution == "Auto":
|
|
resolution = grid_size.max().item()
|
|
print(f"Dual Contouring resolution: {resolution}")
|
|
else:
|
|
resolution = int(dual_contouring_resolution)
|
|
|
|
print('Performing Dual Contouring ...')
|
|
# Perform Dual Contouring remeshing (rebuilds topology)
|
|
cumesh.init(*CuMesh.remeshing.remesh_narrow_band_dc_quad(
|
|
vertices, faces,
|
|
center = center,
|
|
scale = scale * 1.1, # old calculation : (resolution + 3 * remesh_band) / resolution * scale,
|
|
resolution = resolution,
|
|
band = remesh_band,
|
|
project_back = remesh_project, # Snaps vertices back to original surface
|
|
verbose = True,
|
|
remove_inner_faces = remove_inner_faces,
|
|
#bvh = bvh,
|
|
))
|
|
|
|
new_vertices, new_faces = cumesh.read()
|
|
|
|
if remove_floaters:
|
|
new_vertices, new_faces = remove_floater2(new_vertices.cpu().numpy(),new_faces.cpu().numpy())
|
|
new_vertices = torch.from_numpy(new_vertices).contiguous().float().cuda()
|
|
new_faces = torch.from_numpy(new_faces).contiguous().int().cuda()
|
|
cumesh.init(new_vertices, new_faces)
|
|
|
|
print(f"After remeshing: {cumesh.num_vertices} vertices, {cumesh.num_faces} faces")
|
|
|
|
if simplify_method == 'Cumesh':
|
|
cumesh.simplify(target_face_num, verbose=True)
|
|
elif simplify_method == 'Meshlib':
|
|
# GPU -> CPU -> Meshlib -> CPU -> GPU
|
|
v, f = cumesh.read()
|
|
new_vertices, new_faces = simplify_with_meshlib(v.cpu().numpy(), f.cpu().numpy(), target_face_num)
|
|
cumesh.init(torch.from_numpy(new_vertices).float().cuda(), torch.from_numpy(new_faces).int().cuda())
|
|
|
|
print(f"After simplifying: {cumesh.num_vertices} vertices, {cumesh.num_faces} faces")
|
|
pbar.update(1)
|
|
|
|
if fill_holes:
|
|
new_vertices, new_faces = cumesh.read()
|
|
meshlib_mesh = mrmeshnumpy.meshFromFacesVerts(new_faces.detach().clone().cpu().numpy(), new_vertices.detach().clone().cpu().numpy())
|
|
hole_edges = meshlib_mesh.topology.findHoleRepresentiveEdges()
|
|
holes_filled = 0
|
|
|
|
nb_holes = len(hole_edges)
|
|
print(f"{nb_holes} holes found")
|
|
|
|
if nb_holes > 0:
|
|
progress_bar = tqdm(total=nb_holes, desc="Filling holes")
|
|
|
|
last_reported_percent = -1 # Initialize at -1 to ensure 0% triggers an update
|
|
|
|
for i, e in enumerate(hole_edges):
|
|
params = mrmeshpy.FillHoleParams()
|
|
params.metric = mrmeshpy.getUniversalMetric(meshlib_mesh)
|
|
mrmeshpy.fillHole(meshlib_mesh, e, params)
|
|
|
|
# Calculate current progress
|
|
current_step = i + 1
|
|
current_percent = int((current_step / nb_holes) * 100)
|
|
|
|
# Only update the UI if the percentage has moved up
|
|
if current_percent > last_reported_percent:
|
|
# Calculate how many holes have been filled since the last UI update
|
|
# This handles cases where 1% might represent multiple holes
|
|
if last_reported_percent == -1:
|
|
# First update
|
|
progress_bar.update(current_step)
|
|
else:
|
|
# Update by the difference since the last check
|
|
last_step = int((last_reported_percent * nb_holes) / 100)
|
|
diff = current_step - last_step
|
|
progress_bar.update(diff)
|
|
|
|
last_reported_percent = current_percent
|
|
|
|
progress_bar.close()
|
|
|
|
new_vertices = mrmeshnumpy.getNumpyVerts(meshlib_mesh)
|
|
new_faces = mrmeshnumpy.getNumpyFaces(meshlib_mesh.topology)
|
|
|
|
del meshlib_mesh
|
|
gc.collect()
|
|
|
|
cumesh.init(torch.from_numpy(new_vertices).float().to(coords.device), torch.from_numpy(new_faces).int().to(coords.device))
|
|
|
|
# --- Branch: Bake On Vertices (skip UV unwrapping and texture creation) ---
|
|
if bake_on_vertices:
|
|
print('Baking colors on vertices...')
|
|
out_vertices, out_faces = cumesh.read()
|
|
out_vertices = out_vertices.cuda()
|
|
out_faces = out_faces.cuda()
|
|
cumesh.compute_vertex_normals()
|
|
out_normals = cumesh.read_vertex_normals()
|
|
|
|
# Map vertex positions back to original mesh for accurate attribute sampling
|
|
# Use BVH to find the closest point on original mesh surface for more accurate colors
|
|
_, face_id, uvw = bvh.unsigned_distance(out_vertices, return_uvw=True)
|
|
orig_tri_verts = bvh.vertices[bvh.faces[face_id.long()]] # (N_verts, 3, 3)
|
|
mapped_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1)
|
|
|
|
# Sample attributes at mapped positions from the voxel grid
|
|
vertex_attrs = grid_sample_3d(
|
|
attr_volume,
|
|
torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1),
|
|
shape=torch.Size([1, attr_volume.shape[1], *grid_size.tolist()]),
|
|
grid=((mapped_pos - aabb[0]) / voxel_size).reshape(1, -1, 3),
|
|
mode='trilinear',
|
|
)
|
|
|
|
# Extract base color and alpha per vertex (vertex_attrs shape: N_vertices x C)
|
|
base_color_idx = attr_layout['base_color']
|
|
alpha_idx = attr_layout['alpha']
|
|
|
|
# Get RGB values and squeeze any extra dimensions to get (N, 3)
|
|
vertex_colors_rgb = vertex_attrs[..., base_color_idx].cpu().numpy()
|
|
vertex_colors_rgb = np.squeeze(vertex_colors_rgb) # Remove batch dims if any
|
|
if vertex_colors_rgb.ndim == 1:
|
|
vertex_colors_rgb = vertex_colors_rgb[None, :] # Ensure at least 2D
|
|
vertex_colors_rgb = np.clip(vertex_colors_rgb * 255, 0, 255).astype(np.uint8)
|
|
|
|
# Handle alpha based on texture_alpha_mode
|
|
if texture_alpha_mode == "OPAQUE":
|
|
# For OPAQUE mode, use full alpha (255)
|
|
vertex_alpha = np.full((vertex_colors_rgb.shape[0], 1), 255, dtype=np.uint8)
|
|
else:
|
|
vertex_alpha = vertex_attrs[..., alpha_idx].cpu().numpy()
|
|
vertex_alpha = np.squeeze(vertex_alpha) # Remove batch dims if any
|
|
vertex_alpha = np.clip(vertex_alpha * 255, 0, 255).astype(np.uint8)
|
|
# Ensure alpha is 2D with shape (N, 1)
|
|
if vertex_alpha.ndim == 1:
|
|
vertex_alpha = vertex_alpha[:, None]
|
|
|
|
# Combine into RGBA
|
|
vertex_colors_rgba = np.concatenate([vertex_colors_rgb, vertex_alpha], axis=-1)
|
|
|
|
print("Finalizing mesh with vertex colors...")
|
|
pbar.update(1)
|
|
|
|
vertices_np = out_vertices.cpu().numpy()
|
|
faces_np = out_faces.cpu().numpy()
|
|
normals_np = out_normals.cpu().numpy()
|
|
|
|
# Swap Y and Z axes, invert Y (common conversion for GLB compatibility)
|
|
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2].copy(), -vertices_np[:, 1].copy()
|
|
normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2].copy(), -normals_np[:, 1].copy()
|
|
|
|
# Create mesh with vertex colors using ColorVisuals
|
|
if use_custom_normals:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
vertex_normals=normals_np,
|
|
vertex_colors=vertex_colors_rgba,
|
|
process=False,
|
|
)
|
|
else:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
vertex_colors=vertex_colors_rgba,
|
|
process=False,
|
|
)
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
# Return empty placeholder textures for vertex color mode
|
|
placeholder_texture = pil2tensor(Image.new('RGBA', (1, 1), (0, 0, 0, 0)))
|
|
return (textured_mesh, placeholder_texture, placeholder_texture,)
|
|
|
|
# --- Standard texture baking path ---
|
|
print('Unwrapping ...')
|
|
out_vertices, out_faces, out_uvs, out_vmaps = cumesh.uv_unwrap(
|
|
compute_charts_kwargs={
|
|
"threshold_cone_half_angle_rad": mesh_cluster_threshold_cone_half_angle_rad,
|
|
"refine_iterations": mesh_cluster_refine_iterations,
|
|
"global_iterations": mesh_cluster_global_iterations,
|
|
"smooth_strength": mesh_cluster_smooth_strength,
|
|
},
|
|
return_vmaps=True,
|
|
verbose=True,
|
|
)
|
|
pbar.update(1)
|
|
|
|
out_vertices = out_vertices.cuda()
|
|
out_faces = out_faces.cuda()
|
|
out_uvs = out_uvs.cuda()
|
|
out_vmaps = out_vmaps.cuda()
|
|
cumesh.compute_vertex_normals()
|
|
out_normals = cumesh.read_vertex_normals()[out_vmaps]
|
|
|
|
print("Sampling attributes...")
|
|
# Setup differentiable rasterizer context
|
|
ctx = dr.RasterizeCudaContext()
|
|
# Prepare UV coordinates for rasterization (rendering in UV space)
|
|
uvs_rast = torch.cat([out_uvs * 2 - 1, torch.zeros_like(out_uvs[:, :1]), torch.ones_like(out_uvs[:, :1])], dim=-1).unsqueeze(0)
|
|
rast = torch.zeros((1, texture_size, texture_size, 4), device='cuda', dtype=torch.float32)
|
|
|
|
# Rasterize in chunks to save memory
|
|
for i in range(0, out_faces.shape[0], 100000):
|
|
rast_chunk, _ = dr.rasterize(
|
|
ctx, uvs_rast, out_faces[i:i+100000],
|
|
resolution=[texture_size, texture_size],
|
|
)
|
|
mask_chunk = rast_chunk[..., 3:4] > 0
|
|
rast_chunk[..., 3:4] += i # Store face ID in alpha channel
|
|
rast = torch.where(mask_chunk, rast_chunk, rast)
|
|
|
|
# Mask of valid pixels in texture
|
|
mask = rast[0, ..., 3] > 0
|
|
|
|
# Interpolate 3D positions in UV space (finding 3D coord for every texel)
|
|
pos = dr.interpolate(out_vertices.unsqueeze(0), rast, out_faces)[0][0]
|
|
valid_pos = pos[mask]
|
|
|
|
# Map these positions back to the *original* high-res mesh to get accurate attributes
|
|
# This corrects geometric errors introduced by simplification/remeshing
|
|
_, face_id, uvw = bvh.unsigned_distance(valid_pos, return_uvw=True)
|
|
orig_tri_verts = bvh.vertices[bvh.faces[face_id.long()]] # (N_new, 3, 3)
|
|
valid_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1)
|
|
|
|
torch.cuda.synchronize()
|
|
|
|
# Trilinear sampling from the attribute volume (Color, Material props)
|
|
attrs = torch.zeros(texture_size, texture_size, attr_volume.shape[1], device='cuda')
|
|
attrs[mask] = grid_sample_3d(
|
|
attr_volume,
|
|
torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1),
|
|
shape=torch.Size([1, attr_volume.shape[1], *grid_size.tolist()]),
|
|
grid=((valid_pos - aabb[0]) / voxel_size).reshape(1, -1, 3),
|
|
mode='trilinear',
|
|
)
|
|
|
|
# --- Texture Post-Processing & Material Construction ---
|
|
print("Finalizing mesh...")
|
|
pbar.update(1)
|
|
|
|
mask = mask.cpu().numpy()
|
|
|
|
# Extract channels based on layout (BaseColor, Metallic, Roughness, Alpha)
|
|
base_color = np.clip(attrs[..., attr_layout['base_color']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
metallic = np.clip(attrs[..., attr_layout['metallic']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
roughness = np.clip(attrs[..., attr_layout['roughness']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
alpha = np.clip(attrs[..., attr_layout['alpha']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
|
|
alpha_mode = texture_alpha_mode
|
|
|
|
# Inpainting: fill gaps (dilation) to prevent black seams at UV boundaries
|
|
mask_inv = (~mask).astype(np.uint8)
|
|
base_color = cv2.inpaint(base_color, mask_inv, 3, cv2.INPAINT_NS)
|
|
metallic = cv2.inpaint(metallic, mask_inv, 1, cv2.INPAINT_NS)[..., None]
|
|
roughness = cv2.inpaint(roughness, mask_inv, 1, cv2.INPAINT_NS)[..., None]
|
|
alpha = cv2.inpaint(alpha, mask_inv, 1, cv2.INPAINT_NS)[..., None]
|
|
|
|
# Create PBR material
|
|
# Standard PBR packs Metallic and Roughness into Blue and Green channels
|
|
baseColorTexture_np = Image.fromarray(np.concatenate([base_color, alpha], axis=-1))
|
|
metallicRoughnessTexture_np = Image.fromarray(np.concatenate([np.zeros_like(metallic), roughness, metallic], axis=-1))
|
|
material = Trimesh.visual.material.PBRMaterial(
|
|
baseColorTexture=baseColorTexture_np,
|
|
baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8),
|
|
metallicRoughnessTexture=metallicRoughnessTexture_np,
|
|
metallicFactor=1.0,
|
|
roughnessFactor=1.0,
|
|
alphaMode=alpha_mode,
|
|
doubleSided=double_side_material,
|
|
)
|
|
|
|
vertices_np = out_vertices.cpu().numpy()
|
|
faces_np = out_faces.cpu().numpy()
|
|
uvs_np = out_uvs.cpu().numpy()
|
|
normals_np = out_normals.cpu().numpy()
|
|
|
|
# Swap Y and Z axes, invert Y (common conversion for GLB compatibility)
|
|
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2].copy(), -vertices_np[:, 1].copy()
|
|
normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2].copy(), -normals_np[:, 1].copy()
|
|
uvs_np[:, 1] = 1 - uvs_np[:, 1] # Flip UV V-coordinate
|
|
|
|
if use_custom_normals:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
vertex_normals=normals_np,
|
|
process=False,
|
|
visual=Trimesh.visual.TextureVisuals(uv=uvs_np,material=material)
|
|
)
|
|
else:
|
|
textured_mesh = Trimesh.Trimesh(
|
|
vertices=vertices_np,
|
|
faces=faces_np,
|
|
process=False,
|
|
visual=Trimesh.visual.TextureVisuals(uv=uvs_np,material=material)
|
|
)
|
|
|
|
pbar.update(1)
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
baseColorTexture = pil2tensor(baseColorTexture_np)
|
|
metallicRoughnessTexture = pil2tensor(metallicRoughnessTexture_np)
|
|
|
|
return (textured_mesh, baseColorTexture, metallicRoughnessTexture,)
|
|
|
|
class Trellis2Remesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"remesh_band": ("FLOAT",{"default":1.0}),
|
|
"remesh_project": ("FLOAT",{"default":0.0}),
|
|
"dual_contouring_resolution": (["Auto","128","256","512","1024","2048"],{"default":"Auto"}),
|
|
"remove_floaters": ("BOOLEAN",{"default":True}),
|
|
"remove_inner_faces": ("BOOLEAN",{"default":False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, remesh_band, remesh_project, dual_contouring_resolution, remove_floaters, remove_inner_faces):
|
|
reset_cuda()
|
|
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
if remove_floaters:
|
|
mesh_copy = remove_floater(mesh_copy)
|
|
|
|
aabb = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]
|
|
|
|
vertices = mesh_copy.vertices
|
|
faces = mesh_copy.faces
|
|
attr_volume = mesh_copy.attrs
|
|
coords = mesh_copy.coords
|
|
attr_layout = mesh_copy.layout
|
|
voxel_size = mesh_copy.voxel_size
|
|
|
|
# --- Input Normalization (AABB, Voxel Size, Grid Size) ---
|
|
if isinstance(aabb, (list, tuple)):
|
|
aabb = np.array(aabb)
|
|
if isinstance(aabb, np.ndarray):
|
|
aabb = torch.tensor(aabb, dtype=torch.float32, device='cuda')
|
|
|
|
# Calculate grid dimensions based on AABB and voxel size
|
|
if voxel_size is not None:
|
|
if isinstance(voxel_size, float):
|
|
voxel_size = [voxel_size, voxel_size, voxel_size]
|
|
if isinstance(voxel_size, (list, tuple)):
|
|
voxel_size = np.array(voxel_size)
|
|
if isinstance(voxel_size, np.ndarray):
|
|
voxel_size = torch.tensor(voxel_size, dtype=torch.float32, device='cuda')
|
|
grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int()
|
|
else:
|
|
if isinstance(grid_size, int):
|
|
grid_size = [grid_size, grid_size, grid_size]
|
|
if isinstance(grid_size, (list, tuple)):
|
|
grid_size = np.array(grid_size)
|
|
if isinstance(grid_size, np.ndarray):
|
|
grid_size = torch.tensor(grid_size, dtype=torch.int32, device='cuda')
|
|
voxel_size = (aabb[1] - aabb[0]) / grid_size
|
|
|
|
# Move data to GPU
|
|
vertices = vertices.cuda()
|
|
faces = faces.cuda()
|
|
|
|
# Initialize CUDA mesh handler
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(vertices, faces)
|
|
print(f"Current vertices: {cumesh.num_vertices}, faces: {cumesh.num_faces}")
|
|
|
|
vertices, faces = cumesh.read()
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
# Build BVH for the current mesh to guide remeshing
|
|
#print(f"Building BVH for current mesh...")
|
|
#bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
|
|
print("Cleaning mesh...")
|
|
center = aabb.mean(dim=0)
|
|
scale = (aabb[1] - aabb[0]).max().item()
|
|
|
|
if dual_contouring_resolution == "Auto":
|
|
resolution = grid_size.max().item()
|
|
print(f"Dual Contouring resolution: {resolution}")
|
|
else:
|
|
resolution = int(dual_contouring_resolution)
|
|
|
|
print('Performing Dual Contouring ...')
|
|
# Perform Dual Contouring remeshing (rebuilds topology)
|
|
vertices, faces = CuMesh.remeshing.remesh_narrow_band_dc(
|
|
vertices, faces,
|
|
center = center,
|
|
scale = scale * 1.1, # old calculation (resolution + 3 * remesh_band) / resolution * scale,
|
|
resolution = resolution,
|
|
band = remesh_band,
|
|
project_back = remesh_project, # Snaps vertices back to original surface
|
|
verbose = True,
|
|
remove_inner_faces = remove_inner_faces,
|
|
#bvh = bvh,
|
|
)
|
|
|
|
if remove_floaters:
|
|
vertices, faces = remove_floater2(vertices.cpu().numpy(),faces.cpu().numpy())
|
|
vertices = torch.from_numpy(vertices).contiguous().float()
|
|
faces = torch.from_numpy(faces).contiguous().int()
|
|
|
|
print(f"After remeshing: {len(vertices)} vertices, {len(faces)} faces")
|
|
|
|
mesh_copy.vertices = vertices.to(mesh_copy.device)
|
|
mesh_copy.faces = faces.to(mesh_copy.device)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2ReconstructMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"remesh_band": ("FLOAT",{"default":1.0}),
|
|
"resolution": ([128,256,512,1024,2048],{"default":512}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, remesh_band, resolution):
|
|
reset_cuda()
|
|
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
vertices = mesh_copy.vertices.cuda()
|
|
faces = mesh_copy.faces.cuda()
|
|
|
|
# Perform Dual Contouring remeshing (rebuilds topology)
|
|
print('Reconstructing mesh ...')
|
|
vertices, faces = CuMesh.remeshing.reconstruct_mesh_dc(vertices, faces, resolution, verbose=True)
|
|
|
|
print(f"After reconstruction: {len(vertices)} vertices, {len(faces)} faces")
|
|
|
|
mesh_copy.vertices = vertices.to(mesh_copy.device)
|
|
mesh_copy.faces = faces.to(mesh_copy.device)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2ReconstructMeshWithQuad:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"remesh_band": ("FLOAT",{"default":1.0}),
|
|
"resolution": ([128,256,512,1024,2048],{"default":512}),
|
|
"remove_floaters": ("BOOLEAN",{"default":True}),
|
|
"remove_inner_faces": ("BOOLEAN",{"default":False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, remesh_band, resolution, remove_floaters, remove_inner_faces):
|
|
reset_cuda()
|
|
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
vertices = mesh_copy.vertices.cuda()
|
|
faces = mesh_copy.faces.cuda()
|
|
|
|
# Perform Dual Contouring remeshing (rebuilds topology)
|
|
print('Reconstructing mesh ...')
|
|
vertices, faces = CuMesh.remeshing.reconstruct_mesh_dc_quad(vertices, faces, resolution, verbose=True, remove_inner_faces = remove_inner_faces)
|
|
|
|
if remove_floaters:
|
|
vertices, faces = remove_floater2(vertices.cpu().numpy(),faces.cpu().numpy())
|
|
vertices = torch.from_numpy(vertices).contiguous().float()
|
|
faces = torch.from_numpy(faces).contiguous().int()
|
|
|
|
print(f"After reconstruction: {len(vertices)} vertices, {len(faces)} faces")
|
|
|
|
mesh_copy.vertices = vertices.to(mesh_copy.device)
|
|
mesh_copy.faces = faces.to(mesh_copy.device)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2MeshTexturing:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image": ("IMAGE",),
|
|
"trimesh": ("TRIMESH",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":3.00,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.20,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":3.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"resolution": ([512,1024,1536],{"default":1024}),
|
|
"texture_size": ("INT",{"default":4096,"min":512,"max":16384}),
|
|
"texture_alpha_mode": (["OPAQUE","MASK","BLEND"],{"default":"OPAQUE"}),
|
|
"double_side_material": ("BOOLEAN",{"default":False}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
"bake_on_vertices": ("BOOLEAN",{"default":False}),
|
|
"use_custom_normals": ("BOOLEAN",{"default":False}),
|
|
"mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}),
|
|
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH","IMAGE","IMAGE",)
|
|
RETURN_NAMES = ("trimesh","base_color_texture","metallic_roughness_texture",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image, trimesh, seed, texture_steps, texture_guidance_strength, texture_guidance_rescale, texture_rescale_t, resolution, texture_size, texture_alpha_mode, double_side_material, texture_guidance_interval_start, texture_guidance_interval_end, max_views,bake_on_vertices,use_custom_normals,mesh_cluster_threshold_cone_half_angle_rad, sampler):
|
|
images = tensor_batch_to_pil_list(image, max_views=max_views)
|
|
image_in = images[0] if len(images) == 1 else images
|
|
|
|
#image = tensor2pil(image)
|
|
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
textured_mesh, baseColorTexture_np, metallicRoughnessTexture_np = pipeline.texture_mesh(mesh=trimesh,
|
|
image=image_in,
|
|
seed=seed,
|
|
tex_slat_sampler_params = tex_slat_sampler_params,
|
|
resolution = resolution,
|
|
texture_size = texture_size,
|
|
texture_alpha_mode = texture_alpha_mode,
|
|
double_side_material = double_side_material,
|
|
max_views = max_views,
|
|
bake_on_vertices = bake_on_vertices,
|
|
use_custom_normals = use_custom_normals,
|
|
mesh_cluster_threshold_cone_half_angle_rad = mesh_cluster_threshold_cone_half_angle_rad,
|
|
sampler = sampler
|
|
)
|
|
|
|
baseColorTexture = pil2tensor(baseColorTexture_np)
|
|
metallicRoughnessTexture = pil2tensor(metallicRoughnessTexture_np)
|
|
|
|
return (textured_mesh, baseColorTexture, metallicRoughnessTexture, )
|
|
|
|
class Trellis2MeshTexturingMultiView:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"front_image": ("IMAGE",),
|
|
"trimesh": ("TRIMESH",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":3.00,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.20,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":3.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"resolution": ([512,1024,1536],{"default":1024}),
|
|
"texture_size": ("INT",{"default":4096,"min":512,"max":16384}),
|
|
"texture_alpha_mode": (["OPAQUE","MASK","BLEND"],{"default":"OPAQUE"}),
|
|
"double_side_material": ("BOOLEAN",{"default":False}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
"bake_on_vertices": ("BOOLEAN",{"default":False}),
|
|
"use_custom_normals": ("BOOLEAN",{"default":False}),
|
|
"mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}),
|
|
"front_axis": (["z", "x"], {"default": "z"}),
|
|
"blend_temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
},
|
|
"optional": {
|
|
"back_image": ("IMAGE",),
|
|
"left_image": ("IMAGE",),
|
|
"right_image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH","IMAGE","IMAGE",)
|
|
RETURN_NAMES = ("trimesh","base_color_texture","metallic_roughness_texture",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self,
|
|
pipeline,
|
|
front_image,
|
|
trimesh,
|
|
seed,
|
|
texture_steps,
|
|
texture_guidance_strength,
|
|
texture_guidance_rescale,
|
|
texture_rescale_t,
|
|
resolution,
|
|
texture_size,
|
|
texture_alpha_mode,
|
|
double_side_material,
|
|
texture_guidance_interval_start,
|
|
texture_guidance_interval_end,
|
|
bake_on_vertices,
|
|
use_custom_normals,
|
|
mesh_cluster_threshold_cone_half_angle_rad,
|
|
front_axis,
|
|
blend_temperature,
|
|
sampler,
|
|
back_image = None,
|
|
left_image = None,
|
|
right_image = None):
|
|
|
|
reset_cuda()
|
|
|
|
# Convert front image tensor to PIL
|
|
front_pil = tensor2pil(front_image)
|
|
|
|
# Convert optional view image tensors to PIL
|
|
back_pil = tensor2pil(back_image) if back_image is not None else None
|
|
left_pil = tensor2pil(left_image) if left_image is not None else None
|
|
right_pil = tensor2pil(right_image) if right_image is not None else None
|
|
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
textured_mesh, baseColorTexture_np, metallicRoughnessTexture_np = pipeline.texture_mesh_multiview(mesh=trimesh,
|
|
front=front_pil,
|
|
back=back_pil,
|
|
left=left_pil,
|
|
right=right_pil,
|
|
seed=seed,
|
|
tex_slat_sampler_params = tex_slat_sampler_params,
|
|
resolution = resolution,
|
|
texture_size = texture_size,
|
|
texture_alpha_mode = texture_alpha_mode,
|
|
double_side_material = double_side_material,
|
|
bake_on_vertices = bake_on_vertices,
|
|
use_custom_normals = use_custom_normals,
|
|
mesh_cluster_threshold_cone_half_angle_rad = mesh_cluster_threshold_cone_half_angle_rad,
|
|
front_axis = front_axis,
|
|
blend_temperature = blend_temperature,
|
|
sampler = sampler
|
|
)
|
|
|
|
baseColorTexture = pil2tensor(baseColorTexture_np)
|
|
metallicRoughnessTexture = pil2tensor(metallicRoughnessTexture_np)
|
|
|
|
return (textured_mesh, baseColorTexture, metallicRoughnessTexture, )
|
|
|
|
class Trellis2LoadMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"glb_path": ("STRING", {"default": "", "tooltip": "The glb path with mesh to load."}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
OUTPUT_TOOLTIPS = ("The glb model with mesh to texturize.",)
|
|
|
|
FUNCTION = "load"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
DESCRIPTION = "Loads a glb model from the given path."
|
|
|
|
def load(self, glb_path):
|
|
if not os.path.exists(glb_path):
|
|
glb_path = os.path.join(folder_paths.get_input_directory(), glb_path)
|
|
|
|
trimesh = Trimesh.load(glb_path, force="mesh")
|
|
|
|
return (trimesh,)
|
|
|
|
class Trellis2PreProcessImage:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"padding": ("INT",{"default":0,"min":0,"max":1024}),
|
|
"remove_background": ("BOOLEAN",{"default":False}),
|
|
"max_size": ("INT",{"default":2048,"min":512,"max":8192,"step":128}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
|
|
def process(self, image, padding, remove_background, max_size):
|
|
image = tensor2pil(image)
|
|
|
|
if remove_background:
|
|
from rembg import remove
|
|
image = remove(image)
|
|
|
|
image = self.preprocess_image(image, max_size)
|
|
|
|
if padding>0:
|
|
border = (int(padding), int(padding), int(padding), int(padding))
|
|
fill_color = self.parse_fill_for_image("0,0,0,255", image)
|
|
image = ImageOps.expand(image,border=border,fill=fill_color)
|
|
|
|
image = pil2tensor(image)
|
|
|
|
return (image,)
|
|
|
|
def parse_fill_for_image(self, fill: str, img):
|
|
values = [int(x.strip()) for x in fill.split(",")]
|
|
|
|
if img.mode in ("L", "P"):
|
|
return values[0]
|
|
|
|
if img.mode == "RGB":
|
|
return tuple(values[:3])
|
|
|
|
if img.mode == "RGBA":
|
|
return tuple(values[:4])
|
|
|
|
raise ValueError(f"Unsupported image mode: {img.mode}")
|
|
|
|
|
|
def preprocess_image(self, input: Image.Image, max_res) -> Image.Image:
|
|
"""
|
|
Preprocess the input image.
|
|
"""
|
|
# if has alpha channel, use it directly; otherwise, remove background
|
|
has_alpha = False
|
|
if input.mode == 'RGBA':
|
|
alpha = np.array(input)[:, :, 3]
|
|
if not np.all(alpha == 255):
|
|
has_alpha = True
|
|
max_size = max(input.size)
|
|
scale = min(1, max_res / max_size)
|
|
if scale < 1:
|
|
input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
|
|
# if has_alpha:
|
|
# output = input
|
|
# else:
|
|
# input = input.convert('RGB')
|
|
# if self.low_vram:
|
|
# self.rembg_model.to(self.device)
|
|
# output = self.rembg_model(input)
|
|
# if self.low_vram:
|
|
# self.rembg_model.cpu()
|
|
output = input
|
|
output_np = np.array(output)
|
|
alpha = output_np[:, :, 3]
|
|
bbox = np.argwhere(alpha > 0.8 * 255)
|
|
bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
|
|
center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
|
|
size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
|
|
size = int(size * 1)
|
|
bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
|
|
output = output.crop(bbox) # type: ignore
|
|
output = np.array(output).astype(np.float32) / 255
|
|
output = output[:, :, :3] * output[:, :, 3:4]
|
|
output = Image.fromarray((output * 255).astype(np.uint8))
|
|
return output
|
|
|
|
class Trellis2MeshRefiner:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"trimesh": ("TRIMESH",),
|
|
"image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7fffffff}),
|
|
"resolution": ([512,1024,1536],{"default":1024}),
|
|
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":3.00,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.20,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":3.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"max_num_tokens": ("INT",{"default":999999,"min":0,"max":999999}),
|
|
"generate_texture_slat": ("BOOLEAN", {"default":True}),
|
|
"downsampling":([16,32,64],{"default":16}),
|
|
"shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
|
|
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", "BVH", )
|
|
RETURN_NAMES = ("mesh", "bvh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, trimesh, image, seed, resolution,
|
|
shape_steps,
|
|
shape_guidance_strength,
|
|
shape_guidance_rescale,
|
|
shape_rescale_t,
|
|
texture_steps,
|
|
texture_guidance_strength,
|
|
texture_guidance_rescale,
|
|
texture_rescale_t,
|
|
max_num_tokens,
|
|
generate_texture_slat,
|
|
downsampling,
|
|
shape_guidance_interval_start,
|
|
shape_guidance_interval_end,
|
|
texture_guidance_interval_start,
|
|
texture_guidance_interval_end,
|
|
use_tiled_decoder,
|
|
max_views,
|
|
sampler):
|
|
|
|
reset_cuda()
|
|
|
|
images = tensor_batch_to_pil_list(image, max_views=max_views)
|
|
image_in = images[0] if len(images) == 1 else images
|
|
|
|
shape_guidance_interval = [shape_guidance_interval_start,shape_guidance_interval_end]
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
|
|
shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t}
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
mesh = pipeline.refine_mesh(mesh = trimesh, image=image_in, seed=seed, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, resolution = resolution, max_num_tokens = max_num_tokens, generate_texture_slat=generate_texture_slat, downsampling=downsampling, use_tiled=use_tiled_decoder, max_views = max_views, sampler = sampler)[0]
|
|
|
|
vertices = mesh.vertices.cuda()
|
|
faces = mesh.faces.cuda()
|
|
|
|
# Build BVH for the current mesh to guide remeshing
|
|
if generate_texture_slat:
|
|
print("Building BVH for current mesh...")
|
|
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
bvh.vertices = vertices.detach().clone()
|
|
bvh.faces = faces.detach().clone()
|
|
else:
|
|
print('Not building BVH, only used for texturing')
|
|
bvh = None
|
|
|
|
return (mesh, bvh,)
|
|
|
|
class Trellis2PostProcess2:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"fill_holes": ("BOOLEAN", {"default":True}),
|
|
"fix_normals": ("BOOLEAN", {"default":False}),
|
|
"fix_face_orientation": ("BOOLEAN", {"default":True}),
|
|
"remove_duplicate_faces": ("BOOLEAN",{"default":True}),
|
|
"weld_vertices": ("BOOLEAN",{"default":True}),
|
|
"weld_vertices_digits": ("INT",{"default":4,"min":1,"max":8}),
|
|
"smooth": ("BOOLEAN",{"default":False}),
|
|
"smooth_iterations": ("INT",{"default":10,"min":1,"max":99,"step":1}),
|
|
"subdivide": ("BOOLEAN",{"default":False}),
|
|
"subdivide_iterations": ("INT",{"default":1,"min":1,"max":10}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, fill_holes, fix_normals, fix_face_orientation, remove_duplicate_faces, weld_vertices, weld_vertices_digits,smooth,smooth_iterations,subdivide,subdivide_iterations):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
vertices_np = mesh_copy.vertices.cpu().numpy()
|
|
faces_np = mesh_copy.faces.cpu().numpy()
|
|
|
|
trimesh = Trimesh.Trimesh(vertices=vertices_np,faces=faces_np)
|
|
|
|
print(f"Initial mesh: {len(trimesh.faces)} faces")
|
|
#print(f"Is winding consistent? {trimesh.is_winding_consistent}")
|
|
|
|
if fix_normals:
|
|
print('Fixing normals ...')
|
|
trimesh.fix_normals()
|
|
|
|
if fix_face_orientation:
|
|
if trimesh.is_watertight:
|
|
print('Mesh is watertight, fixing inversion ...')
|
|
Trimesh.repair.fix_inversion(trimesh)
|
|
else:
|
|
print('Mesh is not watertight, cannot fix inversion')
|
|
|
|
if remove_duplicate_faces:
|
|
print('Removing duplicate faces ...')
|
|
trimesh.update_faces(trimesh.unique_faces())
|
|
|
|
if fill_holes:
|
|
print('Filling holes ...')
|
|
trimesh.fill_holes()
|
|
|
|
if weld_vertices:
|
|
vertices_count = len(trimesh.vertices)
|
|
trimesh.merge_vertices(digits_vertex=weld_vertices_digits)
|
|
new_vertices_count = len(trimesh.vertices)
|
|
nb_vertices_removed = vertices_count - new_vertices_count
|
|
faces_count = len(trimesh.faces)
|
|
trimesh.remove_unreferenced_vertices()
|
|
trimesh.update_faces(trimesh.nondegenerate_faces())
|
|
new_faces_count = len(trimesh.faces)
|
|
nb_faces_removed = faces_count - new_faces_count
|
|
print(f"Weld Vertices: Removed {nb_vertices_removed} vertices / {nb_faces_removed} faces")
|
|
|
|
if smooth:
|
|
print('Smoothing ...')
|
|
Trimesh.smoothing.filter_taubin(trimesh, lamb=0.5, nu=-0.53, iterations=smooth_iterations)
|
|
|
|
if subdivide:
|
|
print('Subdividing ...')
|
|
trimesh = trimesh.subdivide_loop(iterations=subdivide_iterations)
|
|
|
|
new_vertices = torch.from_numpy(trimesh.vertices).float()
|
|
new_faces = torch.from_numpy(trimesh.faces).int()
|
|
|
|
print(f"After postprocessing: {len(new_faces)} faces")
|
|
|
|
mesh_copy.vertices = new_vertices.to(mesh_copy.device)
|
|
mesh_copy.faces = new_faces.to(mesh_copy.device)
|
|
|
|
del trimesh
|
|
gc.collect()
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2OvoxelExportToGLB:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"resolution": ([512,1024],{"default":1024}),
|
|
"texture_size": ([512,1024,2048,4096],{"default":2048}),
|
|
"target_face_num": ("INT",{"default":2000000,"min":500,"max":16000000}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, resolution, texture_size, target_face_num):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
glb = o_voxel.postprocess.to_glb(
|
|
vertices=mesh_copy.vertices,
|
|
faces=mesh_copy.faces,
|
|
attr_volume=mesh_copy.attrs,
|
|
coords=mesh_copy.coords,
|
|
attr_layout=mesh_copy.layout,
|
|
grid_size=resolution,
|
|
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
decimation_target=target_face_num,
|
|
texture_size=texture_size,
|
|
remesh=True,
|
|
remesh_band=1,
|
|
remesh_project=0,
|
|
use_tqdm=True,
|
|
)
|
|
|
|
return (glb,)
|
|
|
|
class Trellis2TrimeshToMeshWithVoxel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"resolution": ([512,1024],{"default":1024}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", )
|
|
RETURN_NAMES = ("mesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, trimesh, resolution):
|
|
mesh_copy = trimesh.copy()
|
|
|
|
mvoxel = self.get_voxelmesh_from_trimesh(mesh_copy, resolution)
|
|
|
|
return (mvoxel,)
|
|
|
|
def get_voxelmesh_from_trimesh(self, mesh, resolution):
|
|
vertices = torch.from_numpy(mesh.vertices).float()
|
|
faces = torch.from_numpy(mesh.faces).long()
|
|
|
|
voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
|
|
vertices.cpu(), faces.cpu(),
|
|
grid_size=resolution,
|
|
aabb=[[-0.5,-0.5,-0.5],[0.5,0.5,0.5]],
|
|
face_weight=1.0,
|
|
boundary_weight=0.2,
|
|
regularization_weight=1e-2,
|
|
timing=True,
|
|
)
|
|
|
|
coords = torch.cat([torch.zeros_like(voxel_indices[:, 0:1]), voxel_indices], dim=-1)
|
|
coords = coords.cpu()
|
|
|
|
del voxel_indices
|
|
del dual_vertices
|
|
del intersected
|
|
gc.collect()
|
|
|
|
pbr_attr_layout = {
|
|
'base_color': slice(0, 3),
|
|
'metallic': slice(3, 4),
|
|
'roughness': slice(4, 5),
|
|
'alpha': slice(5, 6),
|
|
}
|
|
|
|
mvoxel = MeshWithVoxel(
|
|
vertices, faces,
|
|
origin = [-0.5, -0.5, -0.5],
|
|
voxel_size = 1 / resolution,
|
|
coords = coords,
|
|
attrs = None,
|
|
voxel_shape = None,
|
|
layout=pbr_attr_layout
|
|
)
|
|
|
|
return mvoxel
|
|
|
|
class Trellis2Continue:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"input_1": (any,),
|
|
"input_2": (any,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (any, any, )
|
|
RETURN_NAMES = ("output_1", "output_2", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, input_1, input_2):
|
|
return (input_1, input_2,)
|
|
|
|
class Trellis2MeshWithVoxelToMeshlibMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHLIB_MESH", )
|
|
RETURN_NAMES = ("meshlib_mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh):
|
|
meshlib_mesh = mrmeshnumpy.meshFromFacesVerts(mesh.faces.cpu().numpy(), mesh.vertices.cpu().numpy())
|
|
return (meshlib_mesh,)
|
|
|
|
class Trellis2FillHolesWithMeshlib:
|
|
"""Fill all holes in a mesh"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", "INT")
|
|
RETURN_NAMES = ("mesh", "holes_filled")
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
DESCRIPTION = "Fill all holes in a mesh using optimal triangulation."
|
|
|
|
def process(self, mesh):
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
mesh = mrmeshnumpy.meshFromFacesVerts(mesh_copy.faces.detach().clone().cpu().numpy(), mesh_copy.vertices.detach().clone().cpu().numpy())
|
|
|
|
hole_edges = mesh.topology.findHoleRepresentiveEdges()
|
|
holes_filled = 0
|
|
|
|
nb_holes = len(hole_edges)
|
|
print(f"{nb_holes} holes found")
|
|
|
|
if nb_holes > 0:
|
|
progress_bar = tqdm(total=nb_holes, desc="Filling holes")
|
|
pbar = ProgressBar(nb_holes)
|
|
|
|
last_reported_percent = -1 # Initialize at -1 to ensure 0% triggers an update
|
|
|
|
for i, e in enumerate(hole_edges):
|
|
params = mrmeshpy.FillHoleParams()
|
|
params.metric = mrmeshpy.getUniversalMetric(mesh)
|
|
mrmeshpy.fillHole(mesh, e, params)
|
|
|
|
# Calculate current progress
|
|
current_step = i + 1
|
|
current_percent = int((current_step / nb_holes) * 100)
|
|
|
|
# Only update the UI if the percentage has moved up
|
|
if current_percent > last_reported_percent:
|
|
# Calculate how many holes have been filled since the last UI update
|
|
# This handles cases where 1% might represent multiple holes
|
|
if last_reported_percent == -1:
|
|
# First update
|
|
progress_bar.update(current_step)
|
|
pbar.update(current_step)
|
|
else:
|
|
# Update by the difference since the last check
|
|
last_step = int((last_reported_percent * nb_holes) / 100)
|
|
diff = current_step - last_step
|
|
progress_bar.update(diff)
|
|
pbar.update(diff)
|
|
|
|
last_reported_percent = current_percent
|
|
|
|
progress_bar.close()
|
|
|
|
new_vertices = mrmeshnumpy.getNumpyVerts(mesh)
|
|
new_faces = mrmeshnumpy.getNumpyFaces(mesh.topology)
|
|
|
|
del mesh
|
|
gc.collect()
|
|
|
|
mesh_copy.vertices = torch.from_numpy(new_vertices).float().to(mesh_copy.device)
|
|
mesh_copy.faces = torch.from_numpy(new_faces).int().to(mesh_copy.device)
|
|
|
|
return (mesh_copy, holes_filled)
|
|
|
|
class Trellis2SmoothNormals:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
|
|
def process(self, trimesh):
|
|
new_mesh = trimesh.copy()
|
|
new_mesh.vertex_normals = Trimesh.smoothing.get_vertices_normals(new_mesh)
|
|
|
|
return (new_mesh,)
|
|
|
|
class Trellis2RemeshWithQuad:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"remesh_band": ("FLOAT",{"default":1.0}),
|
|
"remesh_project": ("FLOAT",{"default":0.0}),
|
|
"dual_contouring_resolution": (["Auto","128","256","512","1024","2048"],{"default":"Auto"}),
|
|
"remove_floaters": ("BOOLEAN",{"default":True}),
|
|
"remove_inner_faces": ("BOOLEAN",{"default":True})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, remesh_band, remesh_project, dual_contouring_resolution, remove_floaters, remove_inner_faces):
|
|
reset_cuda()
|
|
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
if remove_floaters:
|
|
mesh_copy = remove_floater(mesh_copy)
|
|
|
|
aabb = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]
|
|
|
|
vertices = mesh_copy.vertices
|
|
faces = mesh_copy.faces
|
|
attr_volume = mesh_copy.attrs
|
|
coords = mesh_copy.coords
|
|
attr_layout = mesh_copy.layout
|
|
voxel_size = mesh_copy.voxel_size
|
|
|
|
# --- Input Normalization (AABB, Voxel Size, Grid Size) ---
|
|
if isinstance(aabb, (list, tuple)):
|
|
aabb = np.array(aabb)
|
|
if isinstance(aabb, np.ndarray):
|
|
aabb = torch.tensor(aabb, dtype=torch.float32, device='cuda')
|
|
|
|
# Calculate grid dimensions based on AABB and voxel size
|
|
if voxel_size is not None:
|
|
if isinstance(voxel_size, float):
|
|
voxel_size = [voxel_size, voxel_size, voxel_size]
|
|
if isinstance(voxel_size, (list, tuple)):
|
|
voxel_size = np.array(voxel_size)
|
|
if isinstance(voxel_size, np.ndarray):
|
|
voxel_size = torch.tensor(voxel_size, dtype=torch.float32, device='cuda')
|
|
grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int()
|
|
else:
|
|
if isinstance(grid_size, int):
|
|
grid_size = [grid_size, grid_size, grid_size]
|
|
if isinstance(grid_size, (list, tuple)):
|
|
grid_size = np.array(grid_size)
|
|
if isinstance(grid_size, np.ndarray):
|
|
grid_size = torch.tensor(grid_size, dtype=torch.int32, device='cuda')
|
|
voxel_size = (aabb[1] - aabb[0]) / grid_size
|
|
|
|
# Move data to GPU
|
|
vertices = vertices.cuda()
|
|
faces = faces.cuda()
|
|
|
|
# Initialize CUDA mesh handler
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(vertices, faces)
|
|
print(f"Current vertices: {cumesh.num_vertices}, faces: {cumesh.num_faces}")
|
|
|
|
vertices, faces = cumesh.read()
|
|
|
|
del cumesh
|
|
gc.collect()
|
|
|
|
# Build BVH for the current mesh to guide remeshing
|
|
#print(f"Building BVH for current mesh...")
|
|
#bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
|
|
print("Cleaning mesh...")
|
|
center = aabb.mean(dim=0)
|
|
scale = (aabb[1] - aabb[0]).max().item()
|
|
|
|
if dual_contouring_resolution == "Auto":
|
|
resolution = grid_size.max().item()
|
|
print(f"Dual Contouring resolution: {resolution}")
|
|
else:
|
|
resolution = int(dual_contouring_resolution)
|
|
|
|
print('Performing Dual Contouring ...')
|
|
# Perform Dual Contouring remeshing (rebuilds topology)
|
|
vertices, faces = CuMesh.remeshing.remesh_narrow_band_dc_quad(
|
|
vertices, faces,
|
|
center = center,
|
|
scale = scale * 1.1, # old calculation (resolution + 3 * remesh_band) / resolution * scale,
|
|
resolution = resolution,
|
|
band = remesh_band,
|
|
project_back = remesh_project, # Snaps vertices back to original surface
|
|
verbose = True,
|
|
remove_inner_faces = remove_inner_faces,
|
|
#bvh = bvh,
|
|
)
|
|
|
|
if remove_floaters:
|
|
vertices, faces = remove_floater2(vertices.cpu().numpy(),faces.cpu().numpy())
|
|
vertices = torch.from_numpy(vertices).contiguous().float()
|
|
faces = torch.from_numpy(faces).contiguous().int()
|
|
|
|
print(f"After remeshing: {len(vertices)} vertices, {len(faces)} faces")
|
|
|
|
mesh_copy.vertices = vertices.to(mesh_copy.device)
|
|
mesh_copy.faces = faces.to(mesh_copy.device)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2BatchSimplifyMeshAndExport:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"target_face_num": ("STRING",{"default":"2000000,1000000,500000,100000,50000,10000,5000,2500,1000"}),
|
|
"method": (["Cumesh","Meshlib"],{"default":"Cumesh"}),
|
|
"fill_holes":("BOOLEAN",{"default":True}),
|
|
"reorient_vertices":(["None","90 degrees","-90 degrees"],{"default":"90 degrees"}),
|
|
"filename_prefix":("STRING",),
|
|
"file_format": (["glb", "obj", "ply", "stl", "3mf", "dae"],),
|
|
"weld_vertices": ("BOOLEAN",{"default":True}),
|
|
"weld_vertices_digits":("INT",{"default":4,"min":1,"max":8}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", )
|
|
RETURN_NAMES = ("lst_glb_path", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, target_face_num, method, fill_holes, reorient_vertices, filename_prefix, file_format, weld_vertices, weld_vertices_digits):
|
|
lst_output_mesh = []
|
|
list_of_faces = parse_string_to_int_list(target_face_num)
|
|
if len(list_of_faces)>0:
|
|
cumesh = CuMesh.CuMesh()
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
for target_nbfaces in list_of_faces:
|
|
print(f"Processing at {target_nbfaces} ...")
|
|
|
|
vertices = mesh_copy.vertices.detach().clone().cpu().numpy()
|
|
faces = mesh_copy.faces.detach().clone().cpu().numpy()
|
|
|
|
if method=="Cumesh":
|
|
cumesh.init(torch.from_numpy(vertices).float().cuda(), torch.from_numpy(faces).int().cuda())
|
|
# options = {
|
|
# 'method': 'legacy'
|
|
# }
|
|
cumesh.simplify(target_nbfaces, verbose=True)
|
|
vertices, faces = cumesh.read()
|
|
vertices = vertices.cpu().numpy()
|
|
faces = faces.cpu().numpy()
|
|
elif method=="Meshlib":
|
|
vertices, faces = simplify_with_meshlib(vertices, faces, target_nbfaces)
|
|
else:
|
|
raise Exception("Unknown simplification method")
|
|
|
|
if fill_holes:
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
|
|
mmesh = mrmeshnumpy.meshFromFacesVerts(faces, vertices)
|
|
|
|
hole_edges = mmesh.topology.findHoleRepresentiveEdges()
|
|
|
|
nb_holes = len(hole_edges)
|
|
print(f"{nb_holes} holes found")
|
|
|
|
if nb_holes > 0:
|
|
progress_bar = tqdm(total=nb_holes, desc="Filling holes")
|
|
|
|
last_reported_percent = -1 # Initialize at -1 to ensure 0% triggers an update
|
|
|
|
for i, e in enumerate(hole_edges):
|
|
params = mrmeshpy.FillHoleParams()
|
|
params.metric = mrmeshpy.getUniversalMetric(mmesh)
|
|
mrmeshpy.fillHole(mmesh, e, params)
|
|
|
|
# Calculate current progress
|
|
current_step = i + 1
|
|
current_percent = int((current_step / nb_holes) * 100)
|
|
|
|
# Only update the UI if the percentage has moved up
|
|
if current_percent > last_reported_percent:
|
|
# Calculate how many holes have been filled since the last UI update
|
|
# This handles cases where 1% might represent multiple holes
|
|
if last_reported_percent == -1:
|
|
# First update
|
|
progress_bar.update(current_step)
|
|
else:
|
|
# Update by the difference since the last check
|
|
last_step = int((last_reported_percent * nb_holes) / 100)
|
|
diff = current_step - last_step
|
|
progress_bar.update(diff)
|
|
|
|
last_reported_percent = current_percent
|
|
|
|
progress_bar.close()
|
|
|
|
vertices = mrmeshnumpy.getNumpyVerts(mmesh)
|
|
faces = mrmeshnumpy.getNumpyFaces(mmesh.topology)
|
|
|
|
del mmesh
|
|
gc.collect()
|
|
|
|
if reorient_vertices == '90 degrees':
|
|
vertices[:, 1], vertices[:, 2] = vertices[:, 2], -vertices[:, 1]
|
|
elif reorient_vertices == '-90 degrees':
|
|
vertices[:, 1], vertices[:, 2] = -vertices[:, 2], vertices[:, 1]
|
|
|
|
trimesh = Trimesh.Trimesh(
|
|
vertices=vertices,
|
|
faces=faces,
|
|
process=False
|
|
)
|
|
|
|
if weld_vertices:
|
|
vertices_count = len(trimesh.vertices)
|
|
trimesh.merge_vertices(digits_vertex=digits)
|
|
new_vertices_count = len(trimesh.vertices)
|
|
nb_vertices_removed = vertices_count - new_vertices_count
|
|
print(f"Weld Vertices: Removed {nb_vertices_removed} vertices")
|
|
|
|
|
|
filename_prefix_with_nbfaces = f"{filename_prefix}_{target_nbfaces}"
|
|
|
|
full_output_folder, filename, counter, subfolder, filename_prefix_with_nbfaces = folder_paths.get_save_image_path(filename_prefix_with_nbfaces, folder_paths.get_output_directory())
|
|
output_glb_path = Path(full_output_folder, f'{filename}_{counter:05}_.{file_format}')
|
|
output_glb_path.parent.mkdir(exist_ok=True)
|
|
|
|
trimesh.export(output_glb_path, file_type=file_format)
|
|
|
|
lst_output_mesh.append(str(output_glb_path))
|
|
|
|
del trimesh
|
|
|
|
del cumesh
|
|
del mesh_copy
|
|
|
|
return (lst_output_mesh,)
|
|
|
|
class Trellis2WeldVertices:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"merge_texture": ("BOOLEAN",{"default":True}),
|
|
"merge_normals": ("BOOLEAN",{"default":True}),
|
|
"digits":("INT",{"default":4,"min":1,"max":8}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
|
|
def process(self, trimesh, merge_texture, merge_normals, digits):
|
|
new_mesh = trimesh.copy()
|
|
vertices_count = len(new_mesh.vertices)
|
|
new_mesh.merge_vertices(merge_tex=merge_texture, merge_norm=merge_normals, digits_vertex=digits, digits_norm=digits, digits_uv=digits)
|
|
new_vertices_count = len(new_mesh.vertices)
|
|
nb_vertices_removed = vertices_count - new_vertices_count
|
|
faces_count = len(new_mesh.faces)
|
|
new_mesh.remove_unreferenced_vertices()
|
|
new_mesh.update_faces(new_mesh.nondegenerate_faces())
|
|
new_faces_count = len(new_mesh.faces)
|
|
nb_faces_removed = faces_count - new_faces_count
|
|
print(f"Weld Vertices: Removed {nb_vertices_removed} vertices / {nb_faces_removed} faces")
|
|
|
|
return (new_mesh,)
|
|
|
|
class Trellis2StringSelector:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"strings": ("STRING",),
|
|
"index": ("INT",{"default":0,"min":0,"max":1000}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("string",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
|
|
def process(self, strings, index):
|
|
if isinstance(strings, list):
|
|
if len(strings) == 0:
|
|
string = ""
|
|
elif len(strings)<=index:
|
|
index = len(strings)-1
|
|
|
|
string = strings[index]
|
|
elif isinstance(strings, str):
|
|
string = strings
|
|
else:
|
|
raise Exception("string must be a list of a string")
|
|
|
|
return (string,)
|
|
|
|
class Trellis2FillHolesWithCuMesh:
|
|
"""Fill all holes in a mesh"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"max_permieters": ("FLOAT",{"default":0.030,"min":0.001,"max":99.999,"step":0.001}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
|
|
def process(self, mesh, max_permieters):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
mesh_copy.fill_holes(max_hole_perimeter = max_permieters)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2LaplacianSmoothingWithOpen3d:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"iterations": ("INT",{"default":10, "min":1, "max":100}),
|
|
"method": (["Laplacian", "Taubin"],{"default":"Laplacian"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", )
|
|
RETURN_NAMES = ("mesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, iterations, method):
|
|
import open3d
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
vertices = mesh_copy.vertices.cpu().numpy()
|
|
faces = mesh_copy.faces.cpu().numpy().astype(np.int32)
|
|
|
|
open3d_mesh = open3d.geometry.TriangleMesh()
|
|
open3d_mesh.vertices = open3d.utility.Vector3dVector(vertices)
|
|
open3d_mesh.triangles = open3d.utility.Vector3iVector(faces)
|
|
|
|
if method == "Laplacian":
|
|
open3d_mesh = open3d_mesh.filter_smooth_laplacian(number_of_iterations=iterations)
|
|
elif method == "Taubin":
|
|
open3d_mesh = open3d_mesh.filter_smooth_taubin(number_of_iterations=iterations)
|
|
|
|
open3d_mesh.compute_vertex_normals()
|
|
|
|
new_vertices = np.asarray(open3d_mesh.vertices)
|
|
new_faces = np.asarray(open3d_mesh.triangles)
|
|
|
|
mesh_copy.vertices = torch.from_numpy(new_vertices).float().to(mesh_copy.device)
|
|
mesh_copy.faces = torch.from_numpy(new_faces).int().to(mesh_copy.device)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2UnWrapTrimesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}),
|
|
"mesh_cluster_refine_iterations": ("INT",{"default":0}),
|
|
"mesh_cluster_global_iterations": ("INT",{"default":1}),
|
|
"mesh_cluster_smooth_strength": ("INT",{"default":1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", )
|
|
RETURN_NAMES = ("trimesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, trimesh, mesh_cluster_threshold_cone_half_angle_rad, mesh_cluster_refine_iterations, mesh_cluster_global_iterations, mesh_cluster_smooth_strength):
|
|
mesh_cluster_threshold_cone_half_angle_rad = np.radians(mesh_cluster_threshold_cone_half_angle_rad)
|
|
|
|
mesh_copy = trimesh.copy()
|
|
|
|
vertices = torch.from_numpy(mesh_copy.vertices).float().cuda()
|
|
faces = torch.from_numpy(mesh_copy.faces).int().cuda()
|
|
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(vertices, faces)
|
|
|
|
out_vertices, out_faces, out_uvs = cumesh.uv_unwrap(
|
|
compute_charts_kwargs={
|
|
"threshold_cone_half_angle_rad": mesh_cluster_threshold_cone_half_angle_rad,
|
|
"refine_iterations": mesh_cluster_refine_iterations,
|
|
"global_iterations": mesh_cluster_global_iterations,
|
|
"smooth_strength": mesh_cluster_smooth_strength,
|
|
},
|
|
return_vmaps=False,
|
|
verbose=True,
|
|
)
|
|
|
|
del cumesh
|
|
|
|
mesh_copy.vertices = out_vertices.cpu().numpy()
|
|
mesh_copy.faces = out_faces.cpu().numpy()
|
|
#mesh_copy.visual.uv = out_uvs.cpu().numpy()
|
|
mesh_copy.visual = Trimesh.visual.TextureVisuals(uv=out_uvs.cpu().numpy())
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2MeshWithVoxelCascadeGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7fffffff}),
|
|
"pipeline_type": (["1024_cascade","1536_cascade"],{"default":"1024_cascade"}),
|
|
"sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"sparse_structure_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"sparse_structure_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"sparse_structure_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
|
|
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"low_res_shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"low_res_shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"low_res_shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"low_res_shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"low_res_shape_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"low_res_shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"low_res_shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"high_res_shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"high_res_shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"high_res_shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"high_res_shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"high_res_shape_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"high_res_shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"high_res_shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"generate_texture_slat": ("BOOLEAN", {"default":True}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"texture_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
"max_num_tokens": ("INT",{"default":999999,"min":0,"max":999999}),
|
|
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
|
|
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL","BVH", )
|
|
RETURN_NAMES = ("mesh", "bvh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image, seed, pipeline_type,
|
|
# sparse
|
|
sparse_structure_steps,
|
|
sparse_structure_guidance_strength,
|
|
sparse_structure_guidance_rescale,
|
|
sparse_structure_rescale_t,
|
|
sparse_structure_sampler,
|
|
sparse_structure_resolution,
|
|
sparse_structure_guidance_interval_start,
|
|
sparse_structure_guidance_interval_end,
|
|
# low res shape
|
|
low_res_shape_steps,
|
|
low_res_shape_guidance_strength,
|
|
low_res_shape_guidance_rescale,
|
|
low_res_shape_rescale_t,
|
|
low_res_shape_sampler,
|
|
low_res_shape_guidance_interval_start,
|
|
low_res_shape_guidance_interval_end,
|
|
# high res shape
|
|
high_res_shape_steps,
|
|
high_res_shape_guidance_strength,
|
|
high_res_shape_guidance_rescale,
|
|
high_res_shape_rescale_t,
|
|
high_res_shape_sampler,
|
|
high_res_shape_guidance_interval_start,
|
|
high_res_shape_guidance_interval_end,
|
|
# texture,
|
|
generate_texture_slat,
|
|
texture_steps,
|
|
texture_guidance_strength,
|
|
texture_guidance_rescale,
|
|
texture_rescale_t,
|
|
texture_sampler,
|
|
texture_guidance_interval_start,
|
|
texture_guidance_interval_end,
|
|
# others
|
|
max_num_tokens,
|
|
use_tiled_decoder,
|
|
max_views
|
|
):
|
|
|
|
reset_cuda()
|
|
|
|
images = tensor_batch_to_pil_list(image, max_views=max_views)
|
|
image_in = images[0] if len(images) == 1 else images
|
|
|
|
sparse_structure_guidance_interval = [sparse_structure_guidance_interval_start,sparse_structure_guidance_interval_end]
|
|
low_res_shape_guidance_interval = [low_res_shape_guidance_interval_start, low_res_shape_guidance_interval_end]
|
|
high_res_shape_guidance_interval = [high_res_shape_guidance_interval_start, high_res_shape_guidance_interval_end]
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
|
|
sparse_structure_sampler_params = {"steps":sparse_structure_steps,"guidance_strength":sparse_structure_guidance_strength,"guidance_rescale":sparse_structure_guidance_rescale,"guidance_interval":sparse_structure_guidance_interval,"rescale_t":sparse_structure_rescale_t}
|
|
low_res_shape_slat_sampler_params = {"steps":low_res_shape_steps,"guidance_strength":low_res_shape_guidance_strength,"guidance_rescale":low_res_shape_guidance_rescale,"guidance_interval":low_res_shape_guidance_interval,"rescale_t":low_res_shape_rescale_t}
|
|
high_res_shape_slat_sampler_params = {"steps":high_res_shape_steps,"guidance_strength":high_res_shape_guidance_strength,"guidance_rescale":high_res_shape_guidance_rescale,"guidance_interval":high_res_shape_guidance_interval,"rescale_t":high_res_shape_rescale_t}
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
if generate_texture_slat:
|
|
num_steps = 5
|
|
else:
|
|
num_steps = 4
|
|
|
|
pbar = ProgressBar(num_steps)
|
|
|
|
mesh = pipeline.run_cascade(image=image_in,
|
|
seed=seed,
|
|
pipeline_type=pipeline_type,
|
|
sparse_structure_sampler_params = sparse_structure_sampler_params,
|
|
low_res_shape_slat_sampler_params = low_res_shape_slat_sampler_params,
|
|
high_res_shape_slat_sampler_params = high_res_shape_slat_sampler_params,
|
|
tex_slat_sampler_params = tex_slat_sampler_params,
|
|
max_num_tokens = max_num_tokens,
|
|
sparse_structure_resolution = sparse_structure_resolution,
|
|
max_views = max_views,
|
|
generate_texture_slat=generate_texture_slat,
|
|
use_tiled=use_tiled_decoder,
|
|
pbar=pbar,
|
|
sparse_structure_sampler = sparse_structure_sampler,
|
|
low_res_shape_sampler = low_res_shape_sampler,
|
|
high_res_shape_sampler = high_res_shape_sampler,
|
|
tex_sampler = texture_sampler
|
|
)[0]
|
|
|
|
vertices = mesh.vertices.cuda()
|
|
faces = mesh.faces.cuda()
|
|
|
|
if generate_texture_slat:
|
|
# Build BVH for the current mesh to guide remeshing
|
|
print("Building BVH for current mesh...")
|
|
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
bvh.vertices = vertices.detach().clone()
|
|
bvh.faces = faces.detach().clone()
|
|
else:
|
|
print("Not building BVH : only used for texturing")
|
|
bvh = None
|
|
|
|
return (mesh,bvh,)
|
|
|
|
class Trellis2ImageCondGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image": ("IMAGE",),
|
|
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE_COND", "IMAGE_COND", "TRELLIS2PIPELINE",)
|
|
RETURN_NAMES = ("cond_512", "cond_1024", "pipeline",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image, max_views,):
|
|
images = tensor_batch_to_pil_list(image, max_views=max_views)
|
|
image_in = images[0] if len(images) == 1 else images
|
|
|
|
if isinstance(image_in, (list, tuple)):
|
|
images = list(image_in)
|
|
else:
|
|
images = [image_in]
|
|
|
|
pipeline.load_image_cond_model()
|
|
|
|
cond_512 = pipeline.get_cond(images, 512, max_views = max_views)
|
|
cond_1024 = pipeline.get_cond(images, 1024, max_views = max_views)
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_image_cond_model()
|
|
|
|
return (cond_512, cond_1024, pipeline,)
|
|
|
|
class Trellis2SparseGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image_cond": ("IMAGE_COND",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
|
|
"sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"sparse_structure_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"sparse_structure_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"sparse_structure_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
|
|
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("COORDS", "INT", "TRELLIS2PIPELINE",)
|
|
RETURN_NAMES = ("coords", "sparse_structure_resolution", "pipeline",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image_cond, seed,
|
|
# sparse
|
|
sparse_structure_steps,
|
|
sparse_structure_guidance_strength,
|
|
sparse_structure_guidance_rescale,
|
|
sparse_structure_rescale_t,
|
|
sparse_structure_sampler,
|
|
sparse_structure_resolution,
|
|
sparse_structure_guidance_interval_start,
|
|
sparse_structure_guidance_interval_end,
|
|
):
|
|
|
|
self.seed_all(seed)
|
|
|
|
sparse_structure_guidance_interval = [sparse_structure_guidance_interval_start,sparse_structure_guidance_interval_end]
|
|
sparse_structure_sampler_params = {"steps":sparse_structure_steps,"guidance_strength":sparse_structure_guidance_strength,"guidance_rescale":sparse_structure_guidance_rescale,"guidance_interval":sparse_structure_guidance_interval,"rescale_t":sparse_structure_rescale_t}
|
|
|
|
args = pipeline._pretrained_args
|
|
sparse_sampler_prefix = pipeline.GetSamplerName(sparse_structure_sampler)
|
|
pipeline.sparse_structure_sampler = getattr(samplers, f"Flow{sparse_sampler_prefix}GuidanceIntervalSampler")(**args['sparse_structure_sampler']['args'])
|
|
pipeline.load_sparse_structure_model()
|
|
coords = pipeline.sample_sparse_structure(
|
|
image_cond, sparse_structure_resolution,
|
|
1, sparse_structure_sampler_params
|
|
)
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_sparse_structure_model()
|
|
|
|
return (coords, sparse_structure_resolution, pipeline,)
|
|
|
|
def seed_all(self, seed: int = 0):
|
|
import random
|
|
"""
|
|
Set random seeds of all components.
|
|
"""
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
class Trellis2ShapeGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image_cond": ("IMAGE_COND",),
|
|
"coords": ("COORDS",),
|
|
"resolution": ([512,1024],{"default":1024}),
|
|
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"shape_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SHAPE_SLAT", "INT", "TRELLIS2PIPELINE",)
|
|
RETURN_NAMES = ("shape_slat", "resolution", "pipeline",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image_cond, coords, resolution,
|
|
# shape
|
|
shape_steps,
|
|
shape_guidance_strength,
|
|
shape_guidance_rescale,
|
|
shape_rescale_t,
|
|
shape_sampler,
|
|
shape_guidance_interval_start,
|
|
shape_guidance_interval_end,
|
|
):
|
|
|
|
shape_guidance_interval = [shape_guidance_interval_start, shape_guidance_interval_end]
|
|
shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t}
|
|
|
|
args = pipeline._pretrained_args
|
|
shape_sampler_prefix = pipeline.GetSamplerName(shape_sampler)
|
|
pipeline.shape_slat_sampler = getattr(samplers, f"Flow{shape_sampler_prefix}GuidanceIntervalSampler")(**args['shape_slat_sampler']['args'])
|
|
|
|
if resolution == 512:
|
|
pipeline.unload_shape_slat_flow_model_1024()
|
|
pipeline.load_shape_slat_flow_model_512()
|
|
shape_slat = pipeline.sample_shape_slat(
|
|
image_cond, pipeline.models['shape_slat_flow_model_512'],
|
|
coords, shape_slat_sampler_params
|
|
)
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_shape_slat_flow_model_512()
|
|
elif resolution == 1024:
|
|
pipeline.unload_shape_slat_flow_model_512()
|
|
pipeline.load_shape_slat_flow_model_1024()
|
|
shape_slat = pipeline.sample_shape_slat(
|
|
image_cond, pipeline.models['shape_slat_flow_model_1024'],
|
|
coords, shape_slat_sampler_params
|
|
)
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_shape_slat_flow_model_1024()
|
|
|
|
return (shape_slat, resolution, pipeline,)
|
|
|
|
class Trellis2ShapeCascadeGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image_cond": ("IMAGE_COND",),
|
|
"shape_slat": ("SHAPE_SLAT",),
|
|
"from_resolution": ("INT",),
|
|
"to_resolution": ([1024,1536,2048,2560,3072,3584,4096],{"default":1024}),
|
|
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
|
|
"max_num_tokens": ("INT",{"default":999999,"min":0,"max":999999}),
|
|
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"shape_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"shape_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"shape_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
|
|
"shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SHAPE_SLAT","INT","TRELLIS2PIPELINE",)
|
|
RETURN_NAMES = ("shape_slat","resolution","pipeline",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image_cond, shape_slat, from_resolution, to_resolution, sparse_structure_resolution, max_num_tokens,
|
|
# shape
|
|
shape_steps,
|
|
shape_guidance_strength,
|
|
shape_guidance_rescale,
|
|
shape_rescale_t,
|
|
shape_sampler,
|
|
shape_guidance_interval_start,
|
|
shape_guidance_interval_end,
|
|
):
|
|
|
|
shape_guidance_interval = [shape_guidance_interval_start, shape_guidance_interval_end]
|
|
shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t}
|
|
|
|
args = pipeline._pretrained_args
|
|
shape_sampler_prefix = pipeline.GetSamplerName(shape_sampler)
|
|
pipeline.shape_slat_sampler = getattr(samplers, f"Flow{shape_sampler_prefix}GuidanceIntervalSampler")(**args['shape_slat_sampler']['args'])
|
|
pipeline.load_shape_slat_flow_model_1024()
|
|
slat, hr_resolution = self.sample(pipeline, shape_slat, from_resolution, to_resolution, sparse_structure_resolution, max_num_tokens, image_cond, shape_slat_sampler_params, pipeline.models['shape_slat_flow_model_1024'])
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_shape_slat_flow_model_1024()
|
|
|
|
return (slat, hr_resolution, pipeline,)
|
|
|
|
def sample(self, pipeline, slat, lr_resolution, resolution, sparse_structure_resolution, max_num_tokens, cond, sampler_params, flow_model):
|
|
# Upsample
|
|
pipeline.load_shape_slat_decoder()
|
|
if pipeline.low_vram:
|
|
pipeline.models['shape_slat_decoder'].to(pipeline.device)
|
|
pipeline.models['shape_slat_decoder'].low_vram = True
|
|
hr_coords = pipeline.models['shape_slat_decoder'].upsample(slat, upsample_times=4)
|
|
if pipeline.low_vram:
|
|
pipeline.models['shape_slat_decoder'].cpu()
|
|
pipeline.models['shape_slat_decoder'].low_vram = False
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_shape_slat_decoder()
|
|
|
|
hr_resolution = resolution
|
|
ratio = (sparse_structure_resolution / 32)
|
|
|
|
while True:
|
|
quant_coords = torch.cat([
|
|
hr_coords[:, :1],
|
|
((hr_coords[:, 1:] + 0.5) / (lr_resolution * ratio) * (hr_resolution // 16)).int(),
|
|
], dim=1)
|
|
coords = quant_coords.unique(dim=0)
|
|
num_tokens = coords.shape[0]
|
|
if num_tokens < max_num_tokens:
|
|
if hr_resolution != resolution:
|
|
print(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
|
|
print(f"Num Tokens: {num_tokens}")
|
|
break
|
|
hr_resolution -= 128
|
|
if hr_resolution < 1024 and resolution >= 1024:
|
|
print(f"Num Tokens: {num_tokens}")
|
|
hr_resolution = 1024
|
|
break
|
|
if hr_resolution < 512:
|
|
print(f"Num Tokens: {num_tokens}")
|
|
hr_resolution = 512
|
|
break
|
|
|
|
coords_dev = coords.to(pipeline.device)
|
|
# Sample structured latent
|
|
noise = SparseTensor(
|
|
feats=torch.randn(coords.shape[0], flow_model.in_channels, device=pipeline.device),
|
|
coords=coords_dev,
|
|
)
|
|
sampler_params = {**pipeline.shape_slat_sampler_params, **sampler_params}
|
|
if pipeline.low_vram:
|
|
flow_model.to(pipeline.device)
|
|
slat = pipeline.shape_slat_sampler.sample(
|
|
flow_model,
|
|
noise,
|
|
**cond,
|
|
**sampler_params,
|
|
verbose=True,
|
|
tqdm_desc="Sampling shape SLat (HR)",
|
|
).samples
|
|
if pipeline.low_vram:
|
|
flow_model.cpu()
|
|
pipeline._cleanup_cuda()
|
|
|
|
std = torch.tensor(pipeline.shape_slat_normalization['std'])[None].to(slat.device)
|
|
mean = torch.tensor(pipeline.shape_slat_normalization['mean'])[None].to(slat.device)
|
|
slat = slat * std + mean
|
|
|
|
del coords_dev
|
|
if pipeline.low_vram:
|
|
cond = pipeline._cond_cpu(cond)
|
|
pipeline._cleanup_cuda()
|
|
|
|
return slat, hr_resolution
|
|
|
|
class Trellis2TexSlatGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"image_cond": ("IMAGE_COND",),
|
|
"shape_slat": ("SHAPE_SLAT",),
|
|
"resolution": ([512,1024],{"default":1024}),
|
|
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
|
|
"texture_guidance_strength": ("FLOAT",{"default":6.50,"min":0.00,"max":99.99,"step":0.01}),
|
|
"texture_guidance_rescale": ("FLOAT",{"default":0.05,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_rescale_t": ("FLOAT",{"default":4.00,"min":0.00,"max":9.99,"step":0.01}),
|
|
"texture_sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
|
|
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
|
|
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TEXTURE_SLAT", "TRELLIS2PIPELINE",)
|
|
RETURN_NAMES = ("texture_slat", "pipeline",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, image_cond, shape_slat, resolution,
|
|
# shape
|
|
texture_steps,
|
|
texture_guidance_strength,
|
|
texture_guidance_rescale,
|
|
texture_rescale_t,
|
|
texture_sampler,
|
|
texture_guidance_interval_start,
|
|
texture_guidance_interval_end,
|
|
):
|
|
|
|
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
|
|
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
|
|
|
|
if resolution == 512:
|
|
pipeline.unload_tex_slat_flow_model_1024()
|
|
pipeline.load_tex_slat_flow_model_512()
|
|
tex_slat = pipeline.sample_tex_slat_advanced(
|
|
image_cond, pipeline.models['tex_slat_flow_model_512'],
|
|
shape_slat, tex_slat_sampler_params, texture_sampler
|
|
)
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_tex_slat_flow_model_512()
|
|
|
|
elif resolution == 1024:
|
|
pipeline.unload_tex_slat_flow_model_512()
|
|
pipeline.load_tex_slat_flow_model_1024()
|
|
tex_slat = pipeline.sample_tex_slat_advanced(
|
|
image_cond, pipeline.models['tex_slat_flow_model_1024'],
|
|
shape_slat, tex_slat_sampler_params, texture_sampler
|
|
)
|
|
|
|
if not pipeline.keep_models_loaded:
|
|
pipeline.unload_tex_slat_flow_model_1024()
|
|
|
|
return (tex_slat, pipeline,)
|
|
|
|
class Trellis2DecodeLatents:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("TRELLIS2PIPELINE",),
|
|
"shape_slat": ("SHAPE_SLAT",),
|
|
"resolution": ("INT",),
|
|
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
|
|
},
|
|
"optional": {
|
|
"texture_slat": ("TEXTURE_SLAT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", "BVH", "TRELLIS2PIPELINE",)
|
|
RETURN_NAMES = ("mesh", "bvh", "pipeline",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, shape_slat, resolution, use_tiled_decoder, texture_slat = None):
|
|
mesh = pipeline.decode_latent(shape_slat, texture_slat, resolution, use_tiled=use_tiled_decoder)[0]
|
|
|
|
if texture_slat == None:
|
|
print("Not building BVH : only used for texturing")
|
|
bvh = None
|
|
else:
|
|
# Build BVH for the current mesh to guide remeshing
|
|
vertices = mesh.vertices.cuda()
|
|
faces = mesh.faces.cuda()
|
|
|
|
print("Building BVH for current mesh...")
|
|
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
|
|
bvh.vertices = vertices.detach().clone()
|
|
bvh.faces = faces.detach().clone()
|
|
|
|
|
|
return (mesh, bvh, pipeline,)
|
|
|
|
class Trellis2SimplifyMeshAdvanced:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESHWITHVOXEL",),
|
|
"target_face_num": ("INT",{"default":1000000,"min":1,"max":30000000}),
|
|
"thresh":("FLOAT",{"default":1e-8,"min":1e-12,"max":1e-2,"step":0.000000000001}),
|
|
"lambda_edge_length": ("FLOAT",{"default":0.01,"min":0.00,"max":1.00,"step":0.01}),
|
|
"lambda_skinny": ("FLOAT",{"default":0.001,"min":0.000,"max":0.100,"step":0.001}),
|
|
"lambda_curvature": ("FLOAT",{"default":0.050,"min":0.000,"max":0.500,"step":0.001}),
|
|
"lambda_boundary": ("FLOAT",{"default":0.050,"min":0.000,"max":0.500,"step":0.001}),
|
|
"lambda_area": ("FLOAT",{"default":0.010,"min":0.000,"max":0.100,"step":0.001}),
|
|
"qem_regularization": ("FLOAT",{"default":1e-8,"min":1e-10,"max":1e-5,"step":0.0000000001}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MESHWITHVOXEL", )
|
|
RETURN_NAMES = ("mesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, mesh, target_face_num, thresh, lambda_edge_length, lambda_skinny, lambda_curvature, lambda_boundary, lambda_area, qem_regularization):
|
|
mesh_copy = copy.deepcopy(mesh)
|
|
|
|
options = {
|
|
'method': 'advanced',
|
|
'thresh': thresh,
|
|
'lambda_edge_length': lambda_edge_length,
|
|
'lambda_skinny': lambda_skinny,
|
|
'lambda_curvature': lambda_curvature,
|
|
'lambda_boundary': lambda_boundary,
|
|
'lambda_area': lambda_area,
|
|
'qem_regularization': qem_regularization,
|
|
}
|
|
|
|
mesh_copy.simplify_with_cumesh(target = target_face_num, options = options)
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2SimplifyTrimeshAdvanced:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"target_face_num": ("INT",{"default":1000000,"min":1,"max":30000000}),
|
|
"thresh":("FLOAT",{"default":1e-8,"min":1e-12,"max":1e-2,"step":0.000000000001}),
|
|
"lambda_edge_length": ("FLOAT",{"default":0.01,"min":0.00,"max":1.00,"step":0.01}),
|
|
"lambda_skinny": ("FLOAT",{"default":0.001,"min":0.000,"max":0.100,"step":0.001}),
|
|
"lambda_curvature": ("FLOAT",{"default":0.050,"min":0.000,"max":0.500,"step":0.001}),
|
|
"lambda_boundary": ("FLOAT",{"default":0.050,"min":0.000,"max":0.500,"step":0.001}),
|
|
"lambda_area": ("FLOAT",{"default":0.010,"min":0.000,"max":0.100,"step":0.001}),
|
|
"qem_regularization": ("FLOAT",{"default":1e-8,"min":1e-10,"max":1e-5,"step":0.0000000001}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", )
|
|
RETURN_NAMES = ("trimesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, trimesh, target_face_num, thresh, lambda_edge_length, lambda_skinny, lambda_curvature, lambda_boundary, lambda_area, qem_regularization):
|
|
mesh_copy = copy.deepcopy(trimesh)
|
|
|
|
cumesh = CuMesh.CuMesh()
|
|
cumesh.init(torch.from_numpy(mesh_copy.vertices).float().cuda(), torch.from_numpy(mesh_copy.faces).int().cuda())
|
|
|
|
options = {
|
|
'method': 'advanced',
|
|
'thresh': thresh,
|
|
'lambda_edge_length': lambda_edge_length,
|
|
'lambda_skinny': lambda_skinny,
|
|
'lambda_curvature': lambda_curvature,
|
|
'lambda_boundary': lambda_boundary,
|
|
'lambda_area': lambda_area,
|
|
'qem_regularization': qem_regularization,
|
|
}
|
|
|
|
cumesh.simplify(target_face_num, verbose=True, options = options)
|
|
|
|
new_vertices, new_faces = cumesh.read()
|
|
mesh_copy.vertices = new_vertices.cpu().numpy()
|
|
mesh_copy.faces = new_faces.cpu().numpy()
|
|
|
|
del cumesh
|
|
|
|
return (mesh_copy,)
|
|
|
|
class Trellis2MultiViewTexturing:
|
|
"""
|
|
Apply texture to mesh by projecting multiple view images.
|
|
|
|
Uses angle-weighted blending: each surface receives texture from all views
|
|
that can "see" it, weighted by how directly the surface faces each camera.
|
|
|
|
Camera angles (Y-up coordinate system):
|
|
- Azimuth: rotation around Y axis
|
|
- 0° = front (looking in -Z direction)
|
|
- 90° = left (looking in -X direction)
|
|
- 180° = back (looking in +Z direction)
|
|
- 270° = right (looking in +X direction)
|
|
- Elevation: rotation around X axis
|
|
- 0° = horizontal
|
|
- 90° = top (looking in -Y direction, from above)
|
|
- -90° = bottom (looking in +Y direction, from below)
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"texture_size": ("INT", {"default": 4096, "min": 512, "max": 8192}),
|
|
"blend_texture": ("BOOLEAN", {"default":True}),
|
|
"blend_exponent": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 8.0, "step": 0.5}),
|
|
"ortho_scale": ("FLOAT", {"default": 1.0, "min": 0.05, "max": 10.0, "step": 0.01}),
|
|
"norm_size": ("FLOAT",{"default":1.15, "min":0.0, "max":9.99, "step":0.01}),
|
|
"fill_holes": ("BOOLEAN",{"default":True}),
|
|
"max_hole_size": ("INT",{"default":10,"min":0,"max":99999,"step":1}),
|
|
},
|
|
"optional": {
|
|
# Standard views
|
|
"front_image": ("IMAGE",), # az=0, el=0
|
|
"back_image": ("IMAGE",), # az=180, el=0
|
|
"left_image": ("IMAGE",), # az=90, el=0
|
|
"right_image": ("IMAGE",), # az=270, el=0
|
|
"top_image": ("IMAGE",), # az=0, el=90
|
|
"bottom_image": ("IMAGE",), # az=0, el=-90
|
|
# Custom views
|
|
"custom_images": ("IMAGE",),
|
|
"custom_azimuths": ("STRING", {"default": ""}),
|
|
"custom_elevations": ("STRING", {"default": ""}),
|
|
"camera_config": ("HY3DCAMERA",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", "IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("trimesh", "base_color", "metallic_roughness",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(
|
|
self,
|
|
trimesh,
|
|
texture_size,
|
|
blend_texture,
|
|
blend_exponent,
|
|
ortho_scale,
|
|
norm_size,
|
|
fill_holes,
|
|
max_hole_size,
|
|
baseColorTexture = None,
|
|
front_image=None,
|
|
back_image=None,
|
|
left_image=None,
|
|
right_image=None,
|
|
top_image=None,
|
|
bottom_image=None,
|
|
custom_images=None,
|
|
custom_azimuths="",
|
|
custom_elevations="",
|
|
camera_config = None
|
|
):
|
|
from .texture_projection_multiview import texture_mesh_with_multiview
|
|
|
|
reset_cuda()
|
|
|
|
# Collect views
|
|
images = []
|
|
azimuths = []
|
|
elevations = []
|
|
|
|
# Standard views with their camera angles
|
|
standard_views = [
|
|
(front_image, 0, 0, "front"),
|
|
(back_image, 180, 0, "back"),
|
|
(left_image, 90, 0, "left"),
|
|
(right_image, 270, 0, "right"),
|
|
(top_image, 0, 90, "top"),
|
|
(bottom_image, 0, -90, "bottom"),
|
|
]
|
|
|
|
for img, az, el, name in standard_views:
|
|
if img is not None:
|
|
images.append(self._tensor_to_pil(img))
|
|
azimuths.append(az)
|
|
elevations.append(el)
|
|
print(f"[MultiView] Added {name} view (az={az}, el={el})")
|
|
|
|
# Custom views
|
|
if custom_images is not None:
|
|
custom_az_list = self._parse_angles(custom_azimuths)
|
|
custom_el_list = self._parse_angles(custom_elevations)
|
|
|
|
if custom_az_list and custom_el_list:
|
|
num_custom = min(len(custom_az_list), len(custom_el_list), int(custom_images.shape[0]))
|
|
for i in range(num_custom):
|
|
images.append(self._tensor_to_pil(custom_images[i:i+1]))
|
|
azimuths.append(custom_az_list[i])
|
|
elevations.append(custom_el_list[i])
|
|
print(f"[MultiView] Added custom view {i+1} (az={custom_az_list[i]}, el={custom_el_list[i]})")
|
|
elif camera_config:
|
|
selected_camera_azims = camera_config["selected_camera_azims"]
|
|
selected_camera_elevs = camera_config["selected_camera_elevs"]
|
|
#ortho_scale = camera_config["ortho_scale"]
|
|
|
|
num_custom = min(len(selected_camera_azims), len(selected_camera_elevs), int(custom_images.shape[0]))
|
|
for i in range(num_custom):
|
|
images.append(self._tensor_to_pil(custom_images[i:i+1]))
|
|
azimuths.append(selected_camera_azims[i])
|
|
elevations.append(selected_camera_elevs[i])
|
|
print(f"[MultiView] Added custom view {i+1} (az={selected_camera_azims[i]}, el={selected_camera_elevs[i]})")
|
|
|
|
if len(images) == 0:
|
|
raise ValueError("No input images provided! Please connect at least one image.")
|
|
|
|
print(f"[MultiView] Total views: {len(images)}")
|
|
print(f"[MultiView] Azimuths: {azimuths}")
|
|
print(f"[MultiView] Elevations: {elevations}")
|
|
|
|
trimesh_obj, base_color, mr = texture_mesh_with_multiview(
|
|
trimesh,
|
|
images,
|
|
azimuths,
|
|
elevations,
|
|
texture_size=texture_size,
|
|
blend_exponent=blend_exponent,
|
|
ortho_scale=ortho_scale,
|
|
blend_texture=blend_texture,
|
|
fill_holes=fill_holes,
|
|
norm_size=norm_size,
|
|
max_hole_size=max_hole_size
|
|
)
|
|
|
|
return (trimesh_obj, pil2tensor(base_color), pil2tensor(mr))
|
|
|
|
def _tensor_to_pil(self, tensor):
|
|
"""Convert ComfyUI IMAGE tensor to PIL."""
|
|
if len(tensor.shape) == 4:
|
|
arr = (tensor[0].cpu().numpy() * 255).astype(np.uint8)
|
|
else:
|
|
arr = (tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
return Image.fromarray(arr)
|
|
|
|
def _parse_angles(self, angle_string):
|
|
"""Parse comma-separated angles into list of floats."""
|
|
if not angle_string or angle_string.strip() == "":
|
|
return []
|
|
try:
|
|
return [float(x.strip()) for x in angle_string.split(",") if x.strip()]
|
|
except ValueError:
|
|
print(f"[MultiView] Warning: Could not parse angles: {angle_string}")
|
|
return []
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Trellis2LoadModel": Trellis2LoadModel,
|
|
"Trellis2MeshWithVoxelGenerator": Trellis2MeshWithVoxelGenerator,
|
|
"Trellis2LoadImageWithTransparency": Trellis2LoadImageWithTransparency,
|
|
"Trellis2SimplifyMesh": Trellis2SimplifyMesh,
|
|
"Trellis2MeshWithVoxelToTrimesh": Trellis2MeshWithVoxelToTrimesh,
|
|
"Trellis2ExportMesh": Trellis2ExportMesh,
|
|
"Trellis2PostProcessMesh": Trellis2PostProcessMesh,
|
|
"Trellis2UnWrapAndRasterizer": Trellis2UnWrapAndRasterizer,
|
|
"Trellis2MeshWithVoxelAdvancedGenerator": Trellis2MeshWithVoxelAdvancedGenerator,
|
|
"Trellis2PostProcessAndUnWrapAndRasterizer": Trellis2PostProcessAndUnWrapAndRasterizer,
|
|
"Trellis2Remesh": Trellis2Remesh,
|
|
"Trellis2MeshTexturing": Trellis2MeshTexturing,
|
|
"Trellis2LoadMesh": Trellis2LoadMesh,
|
|
"Trellis2PreProcessImage": Trellis2PreProcessImage,
|
|
"Trellis2MeshRefiner": Trellis2MeshRefiner,
|
|
"Trellis2PostProcess2": Trellis2PostProcess2,
|
|
"Trellis2OvoxelExportToGLB": Trellis2OvoxelExportToGLB,
|
|
"Trellis2TrimeshToMeshWithVoxel": Trellis2TrimeshToMeshWithVoxel,
|
|
"Trellis2SimplifyTrimesh": Trellis2SimplifyTrimesh,
|
|
"Trellis2Continue": Trellis2Continue,
|
|
"Trellis2ProgressiveSimplify": Trellis2ProgressiveSimplify,
|
|
"Trellis2ReconstructMesh": Trellis2ReconstructMesh,
|
|
"Trellis2MeshWithVoxelToMeshlibMesh": Trellis2MeshWithVoxelToMeshlibMesh,
|
|
"Trellis2FillHolesWithMeshlib": Trellis2FillHolesWithMeshlib,
|
|
"Trellis2SmoothNormals": Trellis2SmoothNormals,
|
|
"Trellis2RemeshWithQuad": Trellis2RemeshWithQuad,
|
|
"Trellis2BatchSimplifyMeshAndExport": Trellis2BatchSimplifyMeshAndExport,
|
|
"Trellis2MeshWithVoxelMultiViewGenerator": Trellis2MeshWithVoxelMultiViewGenerator,
|
|
"Trellis2MeshTexturingMultiView": Trellis2MeshTexturingMultiView,
|
|
"Trellis2WeldVertices": Trellis2WeldVertices,
|
|
"Trellis2ReconstructMeshWithQuad": Trellis2ReconstructMeshWithQuad,
|
|
"Trellis2StringSelector": Trellis2StringSelector,
|
|
"Trellis2FillHolesWithCuMesh": Trellis2FillHolesWithCuMesh,
|
|
"Trellis2LaplacianSmoothingWithOpen3d": Trellis2LaplacianSmoothingWithOpen3d,
|
|
"Trellis2UnWrapTrimesh": Trellis2UnWrapTrimesh,
|
|
"Trellis2MeshWithVoxelCascadeGenerator": Trellis2MeshWithVoxelCascadeGenerator,
|
|
"Trellis2ImageCondGenerator": Trellis2ImageCondGenerator,
|
|
"Trellis2SparseGenerator": Trellis2SparseGenerator,
|
|
"Trellis2ShapeGenerator": Trellis2ShapeGenerator,
|
|
"Trellis2ShapeCascadeGenerator": Trellis2ShapeCascadeGenerator,
|
|
"Trellis2TexSlatGenerator": Trellis2TexSlatGenerator,
|
|
"Trellis2DecodeLatents": Trellis2DecodeLatents,
|
|
"Trellis2SimplifyMeshAdvanced": Trellis2SimplifyMeshAdvanced,
|
|
"Trellis2SimplifyTrimeshAdvanced": Trellis2SimplifyTrimeshAdvanced,
|
|
"Trellis2MultiViewTexturing": Trellis2MultiViewTexturing,
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Trellis2LoadModel": "Trellis2 - LoadModel",
|
|
"Trellis2MeshWithVoxelGenerator": "Trellis2 - Mesh With Voxel Generator",
|
|
"Trellis2LoadImageWithTransparency": "Trellis2 - Load Image with Transparency",
|
|
"Trellis2SimplifyMesh": "Trellis2 - Simplify Mesh",
|
|
"Trellis2MeshWithVoxelToTrimesh": "Trellis2 - Mesh With Voxel To Trimesh",
|
|
"Trellis2ExportMesh": "Trellis2 - Export Mesh",
|
|
"Trellis2PostProcessMesh": "Trellis2 - PostProcess Mesh (using Cumesh)",
|
|
"Trellis2UnWrapAndRasterizer": "Trellis2 - UV Unwrap and Rasterize",
|
|
"Trellis2MeshWithVoxelAdvancedGenerator": "Trellis2 - Mesh With Voxel Advanced Generator",
|
|
"Trellis2PostProcessAndUnWrapAndRasterizer": "Trellis2 - Post Process/UnWrap and Rasterize",
|
|
"Trellis2Remesh": "Trellis2 - Remesh",
|
|
"Trellis2MeshTexturing": "Trellis2 - Mesh Texturing",
|
|
"Trellis2LoadMesh": "Trellis2 - Load Mesh",
|
|
"Trellis2PreProcessImage": "Trellis2 - PreProcess Image",
|
|
"Trellis2MeshRefiner": "Trellis2 - Mesh Refiner",
|
|
"Trellis2PostProcess2": "Trellis2 - PostProcess Mesh (using Trimesh)",
|
|
"Trellis2OvoxelExportToGLB": "Trellis2 - Ovoxel Export to GLB",
|
|
"Trellis2TrimeshToMeshWithVoxel": "Trellis2 - Trimesh to Mesh with Voxel",
|
|
"Trellis2SimplifyTrimesh": "Trellis2 - Simplify Trimesh",
|
|
"Trellis2Continue": "Trellis2 - Continue",
|
|
"Trellis2ProgressiveSimplify": "Trellis2 - Progressive Simplify",
|
|
"Trellis2ReconstructMesh": "Trellis2 - Reconstruct Mesh",
|
|
"Trellis2MeshWithVoxelToMeshlibMesh": "Trellis2 - Mesh with Voxel to Meshlib Mesh",
|
|
"Trellis2FillHolesWithMeshlib": "Trellis2 - Fill Holes with Meshlib",
|
|
"Trellis2SmoothNormals": "Trellis2 - Smooth Normals",
|
|
"Trellis2RemeshWithQuad": "Trellis2 - Remesh With Quad",
|
|
"Trellis2BatchSimplifyMeshAndExport": "Trellis2 - Batch Simplify Mesh And Export",
|
|
"Trellis2MeshWithVoxelMultiViewGenerator": "Trellis2 - Mesh With Voxel Multi-View Generator",
|
|
"Trellis2MeshTexturingMultiView": "Trellis2 - Mesh Texturing Multi-View",
|
|
"Trellis2WeldVertices": "Trellis2 - Weld Vertices",
|
|
"Trellis2ReconstructMeshWithQuad": "Trellis2 - Reconstruct Mesh With Quad",
|
|
"Trellis2StringSelector": "Trellis2 - String Selector",
|
|
"Trellis2FillHolesWithCuMesh": "Trellis2 - Fill Holes with CuMesh",
|
|
"Trellis2LaplacianSmoothingWithOpen3d": "Trellis2 - Laplacian Smoothing (using open3d)",
|
|
"Trellis2UnWrapTrimesh": "Trellis2 - UnWrap Trimesh",
|
|
"Trellis2MeshWithVoxelCascadeGenerator": "Trellis2 - Mesh With Voxel Cascade Generator",
|
|
"Trellis2ImageCondGenerator": "Trellis2 - ImageCond Generator",
|
|
"Trellis2SparseGenerator": "Trellis2 - Sparse Generator",
|
|
"Trellis2ShapeGenerator": "Trellis2 - Shape Generator",
|
|
"Trellis2ShapeCascadeGenerator": "Trellis2 - Shape Cascade Generator",
|
|
"Trellis2TexSlatGenerator": "Trellis2 - Tex Slat Generator",
|
|
"Trellis2DecodeLatents": "Trellis2 - Decode Latents",
|
|
"Trellis2SimplifyMeshAdvanced": "Trellis2 - Simplify Mesh Advanced",
|
|
"Trellis2SimplifyTrimeshAdvanced": "Trellis2 - Simplify Trimesh Advanced",
|
|
"Trellis2MultiViewTexturing": "Trellis2 - Projection MultiView Texturing",
|
|
}
|