nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
086567f3e96ffd879d2d9ff95683691850095282
MrForExample-ComfyUI-3D-Pack/nodes.py
T
MrForExample 086567f3e9 Add Instant_NGP
2024-04-01 15:50:58 +02:00

2009 lines
77 KiB
Python
Raw Blame History

import os
import sys
import re
import math
import copy
from enum import Enum
from collections import OrderedDict
import folder_paths as comfy_paths
from omegaconf import OmegaConf
import json
import torch
from torch.utils.data import DataLoader
import torchvision.transforms.functional as TF
import numpy as np
from safetensors.torch import load_file
from einops import rearrange
from plyfile import PlyData
from PIL import Image
from .mesh_processer.mesh import Mesh
from .mesh_processer.mesh_utils import (
ply_to_points_cloud,
get_target_axis_and_scale,
switch_ply_axis_and_scale,
switch_mesh_axis_and_scale,
calculate_max_sh_degree_from_gs_ply,
marching_cubes_density_to_mesh,
color_func_to_albedo,
)
from .algorithms.main_3DGS import GaussianSplatting, GaussianSplattingCameraController, GSParams
from .algorithms.main_3DGS_renderer import GaussianSplattingRenderer
from .algorithms.diff_mesh import DiffMesh
from .algorithms.dmtet import DMTetMesh
from .algorithms.triplane_gaussian_transformers import TGS
from .algorithms.large_multiview_gaussian_model import LGM
from .algorithms.nerf_marching_cubes_converter import GSConverterNeRFMarchingCubes
from .algorithms.NeuS_runner import NeuSParams, NeuSRunner
from .algorithms.convolutional_reconstruction_model import CRMSampler
from .algorithms.Instant_NGP import InstantNGP
from .tgs.utils.config import ExperimentConfig, load_config as load_config_tgs
from .tgs.data import CustomImageOrbitDataset
from .tgs.utils.misc import todevice, get_device
from .lgm.core.options import config_defaults
from .lgm.mvdream.pipeline_mvdream import MVDreamPipeline
from .mvdiffusion.pipelines.pipeline_mvdiffusion_image import MVDiffusionImagePipeline
from .mvdiffusion.data.single_image_dataset import SingleImageDataset
from .mvdiffusion.utils.misc import load_config as load_config_wonder3d
from .tsr.system import TSR
from .crm.model.crm.model import CRM
from .shared_utils.image_utils import prepare_torch_img, torch_img_to_pil_rgba
from .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames
DIFFUSERS_PIPE_DICT = OrderedDict([
("MVDreamPipeline", MVDreamPipeline),
("Wonder3DMVDiffusionPipeline", MVDiffusionImagePipeline),
])
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
MANIFEST = {
"name": "ComfyUI-3D-Pack",
"version": (0,0,2),
"author": "Mr. For Example",
"project": "https://github.com/MrForExample/ComfyUI-3D-Pack",
"description": "An extensive node suite that enables ComfyUI to process 3D inputs (Mesh & UV Texture, etc) using cutting edge algorithms (3DGS, NeRF, etc.)",
}
SUPPORTED_3D_EXTENSIONS = (
'.obj',
'.ply',
'.glb',
)
SUPPORTED_3DGS_EXTENSIONS = (
'.ply',
)
SUPPORTED_CHECKPOINTS_EXTENSIONS = (
'.ckpt',
'.bin',
'.safetensors',
)
ELEVATION_MIN = -90
ELEVATION_MAX = 90.0
AZIMUTH_MIN = -180.0
AZIMUTH_MAX = 180.0
WEIGHT_DTYPE = torch.float16
class Preview_3DGS:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_file_path": ("STRING", {"default": '', "multiline": False}),
},
}
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = "preview_gs"
CATEGORY = "Comfy3D/Visualize"
def preview_gs(self, gs_file_path):
gs_folder_path, filename = os.path.split(gs_file_path)
if not os.path.isabs(gs_file_path):
gs_file_path = os.path.join(comfy_paths.output_directory, gs_folder_path)
if not filename.lower().endswith(SUPPORTED_3DGS_EXTENSIONS):
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3DGS file extensions: {SUPPORTED_3DGS_EXTENSIONS}").error.print()
gs_file_path = ""
previews = [
{
"filepath": gs_file_path,
}
]
return {"ui": {"previews": previews}, "result": ()}
class Preview_3DMesh:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
},
}
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = "preview_mesh"
CATEGORY = "Comfy3D/Visualize"
def preview_mesh(self, mesh_file_path):
mesh_folder_path, filename = os.path.split(mesh_file_path)
if not os.path.isabs(mesh_file_path):
mesh_file_path = os.path.join(comfy_paths.output_directory, mesh_folder_path)
if not filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
mesh_file_path = ""
previews = [
{
"filepath": mesh_file_path,
}
]
return {"ui": {"previews": previews}, "result": ()}
class Load_3D_Mesh:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh_file_path": ("STRING", {"default": '', "multiline": False}),
"resize": ("BOOLEAN", {"default": False},),
"renormal": ("BOOLEAN", {"default": True},),
"retex": ("BOOLEAN", {"default": False},),
"optimizable": ("BOOLEAN", {"default": False},),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "load_mesh"
CATEGORY = "Comfy3D/Import|Export"
def load_mesh(self, mesh_file_path, resize, renormal, retex, optimizable):
mesh = None
if not os.path.isabs(mesh_file_path):
mesh_file_path = os.path.join(comfy_paths.input_directory, mesh_file_path)
if os.path.exists(mesh_file_path):
folder, filename = os.path.split(mesh_file_path)
if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
with torch.inference_mode(not optimizable):
mesh = Mesh.load(mesh_file_path, resize, renormal, retex)
else:
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
else:
cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print()
return (mesh, )
class Load_3DGS:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_file_path": ("STRING", {"default": '', "multiline": False}),
},
}
RETURN_TYPES = (
"GS_PLY",
)
RETURN_NAMES = (
"gs_ply",
)
FUNCTION = "load_gs"
CATEGORY = "Comfy3D/Import|Export"
def load_gs(self, gs_file_path):
gs_ply = None
if not os.path.isabs(gs_file_path):
gs_file_path = os.path.join(comfy_paths.input_directory, gs_file_path)
if os.path.exists(gs_file_path):
folder, filename = os.path.split(gs_file_path)
if filename.lower().endswith(SUPPORTED_3DGS_EXTENSIONS):
gs_ply = PlyData.read(gs_file_path)
else:
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3DGS file extensions: {SUPPORTED_3DGS_EXTENSIONS}").error.print()
else:
cstr(f"[{self.__class__.__name__}] File {gs_file_path} does not exist").error.print()
return (gs_ply, )
class Save_3D_Mesh:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh": ("MESH",),
"save_path": ("STRING", {"default": 'Mesh_%Y-%m-%d-%M-%S-%f.obj', "multiline": False}),
},
}
OUTPUT_NODE = True
RETURN_TYPES = (
"STRING",
)
RETURN_NAMES = (
"save_path",
)
FUNCTION = "save_mesh"
CATEGORY = "Comfy3D/Import|Export"
def save_mesh(self, mesh, save_path):
save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__)
if save_path is not None:
mesh.write(save_path)
return (save_path, )
class Save_3DGS:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_ply": ("GS_PLY",),
"save_path": ("STRING", {"default": '3DGS_%Y-%m-%d-%M-%S-%f.ply', "multiline": False}),
},
}
OUTPUT_NODE = True
RETURN_TYPES = (
"STRING",
)
RETURN_NAMES = (
"save_path",
)
FUNCTION = "save_gs"
CATEGORY = "Comfy3D/Import|Export"
def save_gs(self, gs_ply, save_path):
save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3DGS_EXTENSIONS, self.__class__.__name__)
if save_path is not None:
gs_ply.write(save_path)
return (save_path, )
class Switch_3DGS_Axis:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_ply": ("GS_PLY",),
"axis_x_to": (["+x", "-x", "+y", "-y", "+z", "-z"],),
"axis_y_to": (["+y", "-y", "+z", "-z", "+x", "-x"],),
"axis_z_to": (["+z", "-z", "+x", "-x", "+y", "-y"],),
},
}
RETURN_TYPES = (
"GS_PLY",
)
RETURN_NAMES = (
"switched_gs_ply",
)
FUNCTION = "switch_axis_and_scale"
CATEGORY = "Comfy3D/Preprocessor"
def switch_axis_and_scale(self, gs_ply, axis_x_to, axis_y_to, axis_z_to):
switched_gs_ply = None
if axis_x_to[1] != axis_y_to[1] and axis_x_to[1] != axis_z_to[1] and axis_y_to[1] != axis_z_to[1]:
target_axis, target_scale, coordinate_invert_count = get_target_axis_and_scale([axis_x_to, axis_y_to, axis_z_to])
switched_gs_ply = switch_ply_axis_and_scale(gs_ply, target_axis, target_scale, coordinate_invert_count)
else:
cstr(f"[{self.__class__.__name__}] axis_x_to: {axis_x_to}, axis_y_to: {axis_y_to}, axis_z_to: {axis_z_to} have to be on separated axis").error.print()
return (switched_gs_ply, )
class Switch_Mesh_Axis:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh": ("MESH",),
"axis_x_to": (["+x", "-x", "+y", "-y", "+z", "-z"],),
"axis_y_to": (["+y", "-y", "+z", "-z", "+x", "-x"],),
"axis_z_to": (["+z", "-z", "+x", "-x", "+y", "-y"],),
"flip_normal": ("BOOLEAN", {"default": False},),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"switched_mesh",
)
FUNCTION = "switch_axis_and_scale"
CATEGORY = "Comfy3D/Preprocessor"
def switch_axis_and_scale(self, mesh, axis_x_to, axis_y_to, axis_z_to, flip_normal):
switched_mesh = None
if axis_x_to[1] != axis_y_to[1] and axis_x_to[1] != axis_z_to[1] and axis_y_to[1] != axis_z_to[1]:
target_axis, target_scale, coordinate_invert_count = get_target_axis_and_scale([axis_x_to, axis_y_to, axis_z_to])
switched_mesh = switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal)
else:
cstr(f"[{self.__class__.__name__}] axis_x_to: {axis_x_to}, axis_y_to: {axis_y_to}, axis_z_to: {axis_z_to} have to be on separated axis").error.print()
return (switched_mesh, )
class Convert_3DGS_To_Pointcloud:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_ply": ("GS_PLY",),
},
}
RETURN_TYPES = (
"POINTCLOUD",
)
RETURN_NAMES = (
"points_cloud",
)
FUNCTION = "convert_gs_ply"
CATEGORY = "Comfy3D/Preprocessor"
def convert_gs_ply(self, gs_ply):
points_cloud = ply_to_points_cloud(gs_ply)
return (points_cloud, )
class Convert_Mesh_To_Pointcloud:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh": ("MESH",),
},
}
RETURN_TYPES = (
"POINTCLOUD",
)
RETURN_NAMES = (
"points_cloud",
)
FUNCTION = "convert_mesh"
CATEGORY = "Comfy3D/Preprocessor"
def convert_mesh(self, mesh):
points_cloud = mesh.convert_to_pointcloud()
return (points_cloud, )
class Stack_Orbit_Camera_Poses:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"orbit_radius_start": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_radius_stop": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_radius_step": ("FLOAT", {"default": 0.1, 'step': 0.0001}),
"elevation_start": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
"elevation_stop": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
"elevation_step": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
"azimuth_start": ("FLOAT", {"default": 0.0, "min": AZIMUTH_MIN, "max": AZIMUTH_MAX, 'step': 0.0001}),
"azimuth_stop": ("FLOAT", {"default": 0.0, "min": AZIMUTH_MIN, "max": AZIMUTH_MAX, 'step': 0.0001}),
"azimuth_step": ("FLOAT", {"default": 0.0, "min": AZIMUTH_MIN, "max": AZIMUTH_MAX, 'step': 0.0001}),
"orbit_center_X_start": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_center_X_stop": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_center_X_step": ("FLOAT", {"default": 0.1, 'step': 0.0001}),
"orbit_center_Y_start": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_center_Y_stop": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_center_Y_step": ("FLOAT", {"default": 0.1, 'step': 0.0001}),
"orbit_center_Z_start": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_center_Z_stop": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_center_Z_step": ("FLOAT", {"default": 0.1, 'step': 0.0001}),
},
}
RETURN_TYPES = (
"ORBIT_CAMPOSES", # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"FLOAT",
"FLOAT",
"FLOAT",
"FLOAT",
"FLOAT",
"FLOAT",
)
RETURN_NAMES = (
"orbit_camposes",
"orbit_radius_list",
"elevation_list",
"azimuth_list",
"orbit_center_X_list",
"orbit_center_Y_list",
"orbit_center_Z_list",
)
OUTPUT_IS_LIST = (
False,
True,
True,
True,
True,
True,
True,
)
FUNCTION = "get_camposes"
CATEGORY = "Comfy3D/Preprocessor"
class Pose_Config(Enum):
STOP_LARGER_STEP_POS = 0
START_LARGER_STEP_POS = 1
START_LARGER_STEP_NEG = 2
STOP_LARGER_STEP_NEG = 3
class Pose_Type:
def __init__(self, start, stop, step, min_value=-math.inf, max_value=math.inf, is_linear = True):
if abs(step) < 0.0001:
step = 0.0001 * (-1.0 if step < 0 else 1.0)
if is_linear and ( (step > 0 and stop < start) or (step < 0 and stop > start)):
cstr(f"[{self.__class__.__name__}] stop value: {stop} cannot be reached from start value {start} with step value {step}, will reverse the sign of step value to {-step}").warning.print()
self.step = -step
else:
self.step = step
self.start = start
self.stop = stop
self.min = min_value
self.max = max_value
self.is_linear = is_linear # linear or circular (i.e. min and max value are connected, e.g. -180 & 180 degree in azimuth angle) value
def stack_camposes(self, pose_type_index=None, last_camposes=[[]]):
if pose_type_index == None:
pose_type_index = len(self.all_pose_types) - 1
if pose_type_index == -1:
return last_camposes
else:
current_pose_type = self.all_pose_types[pose_type_index]
all_camposes = []
# There are four different kind of situation we need to deal with to make this function generalize for any combination of inputs
if current_pose_type.step > 0:
if current_pose_type.start < current_pose_type.stop or current_pose_type.is_linear:
pose_config = Stack_Orbit_Camera_Poses.Pose_Config.STOP_LARGER_STEP_POS
else:
pose_config = Stack_Orbit_Camera_Poses.Pose_Config.START_LARGER_STEP_POS
else:
if current_pose_type.start > current_pose_type.stop or current_pose_type.is_linear:
pose_config = Stack_Orbit_Camera_Poses.Pose_Config.START_LARGER_STEP_NEG
else:
pose_config = Stack_Orbit_Camera_Poses.Pose_Config.STOP_LARGER_STEP_NEG
p = current_pose_type.start
p_passed_min_max_seam = False
while ( (pose_config == Stack_Orbit_Camera_Poses.Pose_Config.STOP_LARGER_STEP_POS and p <= current_pose_type.stop) or
(pose_config == Stack_Orbit_Camera_Poses.Pose_Config.START_LARGER_STEP_POS and (not p_passed_min_max_seam or p <= current_pose_type.stop)) or
(pose_config == Stack_Orbit_Camera_Poses.Pose_Config.START_LARGER_STEP_NEG and p >= current_pose_type.stop) or
(pose_config == Stack_Orbit_Camera_Poses.Pose_Config.STOP_LARGER_STEP_NEG and (not p_passed_min_max_seam or p >= current_pose_type.stop)) ):
# If current pose value surpass the either min/max value then we map its vaule to the oppsite sign
if pose_config == Stack_Orbit_Camera_Poses.Pose_Config.START_LARGER_STEP_POS and p > current_pose_type.max:
p = current_pose_type.min + p % current_pose_type.max
p_passed_min_max_seam = True
elif pose_config == Stack_Orbit_Camera_Poses.Pose_Config.STOP_LARGER_STEP_NEG and p < current_pose_type.min:
p = current_pose_type.max + p % current_pose_type.min
p_passed_min_max_seam = True
new_camposes = copy.deepcopy(last_camposes)
for campose in new_camposes:
campose.insert(0, p)
all_camposes.extend(new_camposes)
p += current_pose_type.step
return self.stack_camposes(pose_type_index-1, all_camposes)
def get_camposes(self,
orbit_radius_start,
orbit_radius_stop,
orbit_radius_step,
elevation_start,
elevation_stop,
elevation_step,
azimuth_start,
azimuth_stop,
azimuth_step,
orbit_center_X_start,
orbit_center_X_stop,
orbit_center_X_step,
orbit_center_Y_start,
orbit_center_Y_stop,
orbit_center_Y_step,
orbit_center_Z_start,
orbit_center_Z_stop,
orbit_center_Z_step):
"""
Return the combination of all the pose types interpolation values
Return values in two ways:
orbit_camposes: CAMPOSES type list can directly input to other 3D process node (e.g. GaussianSplatting)
all the camera pose types seperated in different list, becasue some 3D model's conditioner only takes a sub set of all camera pose types (e.g. StableZero123)
"""
orbit_radius_list = []
elevation_list = []
azimuth_list = []
orbit_center_X_list = []
orbit_center_Y_list = []
orbit_center_Z_list = []
self.all_pose_types = []
self.all_pose_types.append( Stack_Orbit_Camera_Poses.Pose_Type(orbit_radius_start, orbit_radius_stop, orbit_radius_step) )
self.all_pose_types.append( Stack_Orbit_Camera_Poses.Pose_Type(elevation_start, elevation_stop, elevation_step, ELEVATION_MIN, ELEVATION_MAX) )
self.all_pose_types.append( Stack_Orbit_Camera_Poses.Pose_Type(azimuth_start, azimuth_stop, azimuth_step, AZIMUTH_MIN, AZIMUTH_MAX, False) )
self.all_pose_types.append( Stack_Orbit_Camera_Poses.Pose_Type(orbit_center_X_start, orbit_center_X_stop, orbit_center_X_step) )
self.all_pose_types.append( Stack_Orbit_Camera_Poses.Pose_Type(orbit_center_Y_start, orbit_center_Y_stop, orbit_center_Y_step) )
self.all_pose_types.append( Stack_Orbit_Camera_Poses.Pose_Type(orbit_center_Z_start, orbit_center_Z_stop, orbit_center_Z_step) )
orbit_camposes = self.stack_camposes()
for campose in orbit_camposes:
orbit_radius_list.append(campose[0])
elevation_list.append(campose[1])
azimuth_list.append(campose[2])
orbit_center_X_list.append(campose[3])
orbit_center_Y_list.append(campose[4])
orbit_center_Z_list.append(campose[5])
return (orbit_camposes, orbit_radius_list, elevation_list, azimuth_list, orbit_center_X_list, orbit_center_Y_list, orbit_center_Z_list, )
class Generate_Orbit_Camera_Poses:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_images": ("IMAGE",),
"generate_pose_command": ("STRING", {
"default": "#([start_reference_image_index : end_reference_image_index], orbit_radius, elevation_angle [-90, 90], start_azimuth_angle [0, 360], end_azimuth_angle [0, 360])\n([0:30], 1.75, 0, 0, 360)",
"multiline": True
}),
},
}
RETURN_TYPES = (
"ORBIT_CAMPOSES", # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
)
RETURN_NAMES = (
"orbit_camposes",
)
FUNCTION = "get_camposes"
CATEGORY = "Comfy3D/Preprocessor"
class Slice_Camposes:
def __init__(self, start_reference_image_index, end_reference_image_index, camposes_start_to_end):
self.start_reference_image_index = start_reference_image_index
self.end_reference_image_index = end_reference_image_index
self.camposes_start_to_end = camposes_start_to_end
def get_camposes(self, reference_images, generate_pose_command):
orbit_camposes = []
self.ref_imgs_num_minus_1 = len(reference_images) - 1
# To match pattern ( [ start_reference_image_index : end_reference_image_index ] , orbit_radius, elevation_angle , start_azimuth_angle , end_azimuth_angle )
pattern = re.compile(r"\([ \t]*\[[ \t]*(\d+)[ \t]*:[ \t]*(\d+)[ \t]*\][ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*\)")
all_matches = pattern.findall(generate_pose_command)
all_slice_camposes = []
for match in all_matches:
start_reference_image_index, end_reference_image_index, orbit_radius, elevation_angle, start_azimuth_angle, end_azimuth_angle = (int(s) if i < 2 else float(s) for i, s in enumerate(match))
end_reference_image_index = min(end_reference_image_index, self.ref_imgs_num_minus_1)
if start_reference_image_index <= end_reference_image_index:
azimuth_imgs_num = end_reference_image_index - start_reference_image_index + 1
# calculate all the reference camera azimuth angles
camposes_start_to_end = []
if start_azimuth_angle > end_azimuth_angle:
azimuth_angle_interval = -(end_azimuth_angle + 360 - start_azimuth_angle) / azimuth_imgs_num
else:
azimuth_angle_interval = (end_azimuth_angle - start_azimuth_angle) / azimuth_imgs_num
now_azimuth_angle = start_azimuth_angle
for _ in range(azimuth_imgs_num):
camposes_start_to_end.append((orbit_radius, elevation_angle, now_azimuth_angle, 0.0, 0.0, 0.0))
now_azimuth_angle = (now_azimuth_angle + azimuth_angle_interval) % 360
all_slice_camposes.append(Generate_Orbit_Camera_Poses.Slice_Camposes(start_reference_image_index, end_reference_image_index, camposes_start_to_end))
else:
cstr(f"[{self.__class__.__name__}] start_reference_image_index: {start_reference_image_index} must smaller than or equal to end_reference_image_index: {end_reference_image_index}").error.print()
all_slice_camposes = sorted(all_slice_camposes, key=lambda slice_camposes:slice_camposes.start_reference_image_index)
last_end_index_plus_1 = 0
for slice_camposes in all_slice_camposes:
if last_end_index_plus_1 == slice_camposes.start_reference_image_index:
orbit_camposes.extend(slice_camposes.camposes_start_to_end)
last_end_index_plus_1 = slice_camposes.end_reference_image_index + 1
else:
orbit_camposes = []
cstr(f"[{self.__class__.__name__}] Last end_reference_image_index: {end_reference_image_index} plus 1 must equal to current start_reference_image_index: {start_reference_image_index}").error.print()
return (orbit_camposes, )
class Gaussian_Splatting_Orbit_Renderer:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_ply": ("GS_PLY",),
"render_image_width": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"render_image_height": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"render_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"render_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"render_background_color_r": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"render_background_color_g": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"render_background_color_b": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
}
}
RETURN_TYPES = (
"IMAGE",
"MASK",
)
RETURN_NAMES = (
"rendered_gs_images", # [Number of Poses, H, W, 3]
"rendered_gs_masks", # [Number of Poses, H, W, 1]
)
FUNCTION = "render_gs"
CATEGORY = "Comfy3D/Preprocessor"
def render_gs(
self,
gs_ply,
render_image_width,
render_image_height,
render_orbit_camera_poses,
render_orbit_camera_fovy,
render_background_color_r,
render_background_color_g,
render_background_color_b,
):
sh_degree, _ = calculate_max_sh_degree_from_gs_ply(gs_ply)
renderer = GaussianSplattingRenderer(sh_degree=sh_degree)
renderer.initialize(gs_ply)
cam_controller = GaussianSplattingCameraController(
renderer,
render_image_width,
render_image_height,
render_orbit_camera_fovy,
static_bg=[render_background_color_r, render_background_color_g, render_background_color_b]
)
all_rendered_images, all_rendered_masks = cam_controller.render_all_pose(render_orbit_camera_poses)
all_rendered_images = all_rendered_images.permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
all_rendered_masks = all_rendered_masks.squeeze(1) # [N, 1, H, W] -> [N, H, W]
return (all_rendered_images, all_rendered_masks)
class Gaussian_Splatting:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_images": ("IMAGE",),
"reference_masks": ("MASK",),
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"training_iterations": ("INT", {"default": 30_000, "min": 1, "max": 0xffffffffffffffff}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}),
"ms_ssim_loss_weight": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, }),
"alpha_loss_weight": ("FLOAT", {"default": 3, "min": 0.0, }),
"offset_loss_weight": ("FLOAT", {"default": 0.0, "min": 0.0, }),
"offset_opacity_loss_weight": ("FLOAT", {"default": 0.0, "min": 0.0, }),
"invert_background_probability": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}),
"feature_learning_rate": ("FLOAT", {"default": 0.0025, "min": 0.000001, "step": 0.000001}),
"opacity_learning_rate": ("FLOAT", {"default": 0.05, "min": 0.000001, "step": 0.000001}),
"scaling_learning_rate": ("FLOAT", {"default": 0.005, "min": 0.000001, "step": 0.000001}),
"rotation_learning_rate": ("FLOAT", {"default": 0.001, "min": 0.000001, "step": 0.000001}),
"position_learning_rate_init": ("FLOAT", {"default": 0.00016, "min": 0.000001, "step": 0.000001}),
"position_learning_rate_final": ("FLOAT", {"default": 0.0000016, "min": 0.0000001, "step": 0.0000001}),
"position_learning_rate_delay_mult": ("FLOAT", {"default": 0.01, "min": 0.000001, "step": 0.000001}),
"position_learning_rate_max_steps": ("INT", {"default": 30_000, "min": 1, "max": 0xffffffffffffffff}),
"initial_gaussians_num": ("INT", {"default": 10_000, "min": 1, "max": 0xffffffffffffffff}),
"K_nearest_neighbors": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"percent_dense": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}),
"density_start_iterations": ("INT", {"default": 500, "min": 0, "max": 0xffffffffffffffff}),
"density_end_iterations": ("INT", {"default": 15_000, "min": 0, "max": 0xffffffffffffffff}),
"densification_interval": ("INT", {"default": 100, "min": 1, "max": 0xffffffffffffffff}),
"opacity_reset_interval": ("INT", {"default": 3000, "min": 1, "max": 0xffffffffffffffff}),
"densify_grad_threshold": ("FLOAT", {"default": 0.0002, "min": 0.00001, "step": 0.00001}),
"gaussian_sh_degree": ("INT", {"default": 3, "min": 0}),
},
"optional": {
"points_cloud_to_initialize_gaussian": ("POINTCLOUD",),
"ply_to_initialize_gaussian": ("GS_PLY",),
"mesh_to_initialize_gaussian": ("MESH",),
}
}
RETURN_TYPES = (
"GS_PLY",
)
RETURN_NAMES = (
"gs_ply",
)
FUNCTION = "run_gs"
CATEGORY = "Comfy3D/Algorithm"
def run_gs(
self,
reference_images,
reference_masks,
reference_orbit_camera_poses,
reference_orbit_camera_fovy,
training_iterations,
batch_size,
ms_ssim_loss_weight,
alpha_loss_weight,
offset_loss_weight,
offset_opacity_loss_weight,
invert_background_probability,
feature_learning_rate,
opacity_learning_rate,
scaling_learning_rate,
rotation_learning_rate,
position_learning_rate_init,
position_learning_rate_final,
position_learning_rate_delay_mult,
position_learning_rate_max_steps,
initial_gaussians_num,
K_nearest_neighbors,
percent_dense,
density_start_iterations,
density_end_iterations,
densification_interval,
opacity_reset_interval,
densify_grad_threshold,
gaussian_sh_degree,
points_cloud_to_initialize_gaussian=None,
ply_to_initialize_gaussian=None,
mesh_to_initialize_gaussian=None,
):
gs_ply = None
ref_imgs_num = len(reference_images)
ref_masks_num = len(reference_masks)
if ref_imgs_num == ref_masks_num:
ref_cam_poses_num = len(reference_orbit_camera_poses)
if ref_imgs_num == ref_cam_poses_num:
if batch_size > ref_imgs_num:
cstr(f"[{self.__class__.__name__}] Batch size {batch_size} is bigger than number of reference images {ref_imgs_num}! Set batch size to {ref_imgs_num} instead").warning.print()
batch_size = ref_imgs_num
with torch.inference_mode(False):
gs_params = GSParams(
training_iterations,
batch_size,
ms_ssim_loss_weight,
alpha_loss_weight,
offset_loss_weight,
offset_opacity_loss_weight,
invert_background_probability,
feature_learning_rate,
opacity_learning_rate,
scaling_learning_rate,
rotation_learning_rate,
position_learning_rate_init,
position_learning_rate_final,
position_learning_rate_delay_mult,
position_learning_rate_max_steps,
initial_gaussians_num,
K_nearest_neighbors,
percent_dense,
density_start_iterations,
density_end_iterations,
densification_interval,
opacity_reset_interval,
densify_grad_threshold,
gaussian_sh_degree
)
if points_cloud_to_initialize_gaussian is not None:
gs_init_input = points_cloud_to_initialize_gaussian
elif ply_to_initialize_gaussian is not None:
gs_init_input = ply_to_initialize_gaussian
else:
gs_init_input = mesh_to_initialize_gaussian
gs = GaussianSplatting(gs_params, gs_init_input)
gs.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy)
gs.training()
gs_ply = gs.renderer.gaussians.to_ply()
else:
cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of reference camera poses {ref_cam_poses_num}").error.print()
else:
cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of masks {ref_masks_num}").error.print()
return (gs_ply, )
class Fitting_Mesh_With_Multiview_Images:
def __init__(self):
self.need_update = False
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_images": ("IMAGE",),
"reference_masks": ("MASK",),
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"mesh": ("MESH",),
"mesh_albedo_width": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"mesh_albedo_height": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"training_iterations": ("INT", {"default": 1024, "min": 1, "max": 100000}),
"batch_size": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"texture_learning_rate": ("FLOAT", {"default": 0.001, "min": 0.00001, "step": 0.00001}),
"train_mesh_geometry": ("BOOLEAN", {"default": False},),
"geometry_learning_rate": ("FLOAT", {"default": 0.0001, "min": 0.00001, "step": 0.00001}),
"ms_ssim_loss_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"remesh_after_n_iteration": ("INT", {"default": 512, "min": 128, "max": 100000}),
"force_cuda_rasterize": ("BOOLEAN", {"default": False},),
},
}
RETURN_TYPES = (
"MESH",
"IMAGE",
)
RETURN_NAMES = (
"trained_mesh",
"baked_texture", # [1, H, W, 3]
)
FUNCTION = "fitting_mesh"
CATEGORY = "Comfy3D/Algorithm"
def fitting_mesh(
self,
reference_images,
reference_masks,
reference_orbit_camera_poses,
reference_orbit_camera_fovy,
mesh,
mesh_albedo_width,
mesh_albedo_height,
training_iterations,
batch_size,
texture_learning_rate,
train_mesh_geometry,
geometry_learning_rate,
ms_ssim_loss_weight,
remesh_after_n_iteration,
force_cuda_rasterize
):
mesh.set_new_albedo(mesh_albedo_width, mesh_albedo_height)
trained_mesh = None
baked_texture = None
ref_imgs_num = len(reference_images)
ref_masks_num = len(reference_masks)
if ref_imgs_num == ref_masks_num:
ref_cam_poses_num = len(reference_orbit_camera_poses)
if ref_imgs_num == ref_cam_poses_num:
if batch_size > ref_imgs_num:
cstr(f"[{self.__class__.__name__}] Batch size {batch_size} is bigger than number of reference images {ref_imgs_num}! Set batch size to {ref_imgs_num} instead").warning.print()
batch_size = ref_imgs_num
with torch.inference_mode(False):
mesh_fitter = DiffMesh(mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, remesh_after_n_iteration, force_cuda_rasterize)
mesh_fitter.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy)
mesh_fitter.training()
trained_mesh, baked_texture = mesh_fitter.get_mesh_and_texture()
else:
cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of reference camera poses {ref_cam_poses_num}").error.print()
else:
cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of masks {ref_masks_num}").error.print()
return (trained_mesh, baked_texture, )
class Deep_Marching_Tetrahedrons:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"training_iterations": ("INT", {"default": 5000, "min": 1, "max": 100000}),
"points_cloud_fitting_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
"mesh_smoothing_weight": ("FLOAT", {"default": 0.1, "min": 0.0, "step": 0.01}),
"chamfer_faces_sample_scale": ("FLOAT", {"default": 1.0, "min": 0.01, "step": 0.01}),
"mesh_scale": ("FLOAT", {"default": 1.0, "min": 0.01, "step": 0.01}),
"grid_resolution": ([128, 64, 32], ),
"geometry_learning_rate": ("FLOAT", {"default": 0.0001, "min": 0.00001, "step": 0.00001}),
"positional_encoding_multires": ("INT", {"default": 2, "min": 2}),
"mlp_internal_dims": ("INT", {"default": 128, "min": 8}),
"mlp_hidden_layer_num": ("INT", {"default": 5, "min": 1}),
},
"optional": {
"reference_points_cloud": ("POINTCLOUD",),
"reference_images": ("IMAGE",),
"reference_masks": ("MASK",),
}
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"trained_mesh",
)
FUNCTION = "run_dmtet"
CATEGORY = "Comfy3D/Algorithm"
def run_dmtet(self, training_iterations, points_cloud_fitting_weight, mesh_smoothing_weight, chamfer_faces_sample_scale, mesh_scale, grid_resolution, geometry_learning_rate,
positional_encoding_multires, mlp_internal_dims, mlp_hidden_layer_num,
reference_points_cloud=None, reference_images=None, reference_masks=None):
with torch.inference_mode(False):
dmtet_mesh = DMTetMesh(
training_iterations,
points_cloud_fitting_weight,
mesh_smoothing_weight,
chamfer_faces_sample_scale,
mesh_scale,
grid_resolution,
geometry_learning_rate,
positional_encoding_multires,
mlp_internal_dims,
mlp_hidden_layer_num
)
if reference_points_cloud is None and reference_images is None:
cstr(f"[{self.__class__.__name__}] reference_points_cloud and reference_images cannot both be None, you need at least provide one of them!").error.print()
raise Exception("User didn't provide necessary inputs, stop executing the code!")
dmtet_mesh.training(reference_points_cloud, reference_images, reference_masks)
mesh = dmtet_mesh.get_mesh()
return (mesh, )
class Load_Triplane_Gaussian_Transformers:
checkpoints_dir = "checkpoints/tgs"
default_ckpt_name = "model_lvis_rel.ckpt"
default_repo_id = "VAST-AI/TriplaneGaussian"
tgs_config_path = "configs/tgs_config.yaml"
@classmethod
def INPUT_TYPES(cls):
cls.checkpoints_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS)
if cls.default_ckpt_name not in all_models_names:
all_models_names += [cls.default_ckpt_name]
return {
"required": {
"model_name": (all_models_names, ),
},
}
RETURN_TYPES = (
"TGS_MODEL",
)
RETURN_NAMES = (
"tgs_model",
)
FUNCTION = "load_TGS"
CATEGORY = "Comfy3D/Import|Export"
def load_TGS(self, model_name):
device = get_device()
config_path = os.path.join(ROOT_PATH, self.tgs_config_path)
cfg: ExperimentConfig = load_config_tgs(config_path)
# Download pre-trained model if it not exist locally
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
if not os.path.isfile(ckpt_path):
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
cfg.system.weights=ckpt_path
tgs_model = TGS(cfg=cfg.system).to(device)
save_path = os.path.join(ROOT_PATH, "outputs")
tgs_model.set_save_dir(save_path)
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
return (tgs_model, )
class Triplane_Gaussian_Transformers:
tgs_config_path = "configs/tgs_config.yaml"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_image": ("IMAGE", ),
"reference_mask": ("MASK",),
"tgs_model": ("TGS_MODEL", ),
"cam_dist": ("FLOAT", {"default": 1.9, "min": 0.01, "step": 0.01}),
},
}
RETURN_TYPES = (
"GS_PLY",
)
RETURN_NAMES = (
"gs_ply",
)
FUNCTION = "run_TGS"
CATEGORY = "Comfy3D/Algorithm"
def run_TGS(self, reference_image, reference_mask, tgs_model, cam_dist):
config_path = os.path.join(ROOT_PATH, self.tgs_config_path)
cfg: ExperimentConfig = load_config_tgs(config_path)
#save_path = os.path.join(ROOT_PATH, "outputs")
#tgs_model.set_save_dir(save_path)
cfg.data.cond_camera_distance = cam_dist
cfg.data.eval_camera_distance = cam_dist
dataset = CustomImageOrbitDataset(reference_image, reference_mask, cfg.data)
dataloader = DataLoader(
dataset,
batch_size=cfg.data.eval_batch_size,
shuffle=False,
collate_fn=dataset.collate
)
gs_ply = []
for batch in dataloader:
batch = todevice(batch)
gs_ply.extend(tgs_model(batch))
"""
# Output rendered video, for testing this node only
tgs_model.save_img_sequences(
"video",
"(\d+)\.png",
save_format="mp4",
fps=30,
delete=True,
)
"""
return (gs_ply[0], )
class Load_Diffusers_Pipeline:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"diffusers_pipeline_name": (list(DIFFUSERS_PIPE_DICT.keys()),),
"model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}),
},
}
RETURN_TYPES = (
"DIFFUSERS_PIPE",
)
RETURN_NAMES = (
"pipe",
)
FUNCTION = "load_diffusers_pipe"
CATEGORY = "Comfy3D/Import|Export"
def load_diffusers_pipe(self, diffusers_pipeline_name, model_name):
# resume pretrained checkpoint
ckpt_path = os.path.join(ROOT_PATH, "checkpoints", model_name)
if not os.path.exists(ckpt_path):
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {model_name}").warning.print()
from huggingface_hub import snapshot_download
snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
diffusers_pipeline_class = DIFFUSERS_PIPE_DICT[diffusers_pipeline_name]
# load image dream
pipe = diffusers_pipeline_class.from_pretrained(
ckpt_path,
torch_dtype=WEIGHT_DTYPE,
trust_remote_code=True,
).to(device)
pipe.enable_xformers_memory_efficient_attention()
return (pipe, )
class Wonder3D_MVDiffusion_Model:
wonder3d_config_path = "configs/wonder3d_config.yaml"
fix_cam_pose_dir = "mvdiffusion/data/fixed_poses/nine_views"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mvdiffusion_pipe": ("DIFFUSERS_PIPE",),
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"mv_guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
"num_inference_steps": ("INT", {"default": 50, "min": 1}),
},
}
RETURN_TYPES = (
"IMAGE",
"IMAGE",
)
RETURN_NAMES = (
"multiview_images",
"multiview_normals",
)
FUNCTION = "run_mvdiffusion"
CATEGORY = "Comfy3D/Algorithm"
def run_mvdiffusion(
self,
mvdiffusion_pipe,
reference_image, # [1, H, W, 3]
reference_mask, # [1, H, W]
seed,
mv_guidance_scale,
num_inference_steps,
):
config_path = os.path.join(ROOT_PATH, self.wonder3d_config_path)
cfg = load_config_wonder3d(config_path)
batch = self.prepare_data(reference_image, reference_mask)
mvdiffusion_pipe.set_progress_bar_config(disable=True)
seed = int(seed)
generator = torch.Generator(device=mvdiffusion_pipe.unet.device).manual_seed(seed)
# repeat (2B, Nv, 3, H, W)
imgs_in = torch.cat([batch['imgs_in']] * 2, dim=0).to(WEIGHT_DTYPE)
# (2B, Nv, Nce)
camera_embeddings = torch.cat([batch['camera_embeddings']] * 2, dim=0).to(WEIGHT_DTYPE)
task_embeddings = torch.cat([batch['normal_task_embeddings'], batch['color_task_embeddings']], dim=0).to(WEIGHT_DTYPE)
camera_embeddings = torch.cat([camera_embeddings, task_embeddings], dim=-1).to(WEIGHT_DTYPE)
# (B*Nv, 3, H, W)
imgs_in = rearrange(imgs_in, "Nv C H W -> (Nv) C H W")
# (B*Nv, Nce)
# camera_embeddings = rearrange(camera_embeddings, "B Nv Nce -> (B Nv) Nce")
out = mvdiffusion_pipe(
imgs_in,
# camera_embeddings,
generator=generator,
guidance_scale=mv_guidance_scale,
num_inference_steps=num_inference_steps,
output_type='pt',
num_images_per_prompt=1,
**cfg.pipe_validation_kwargs,
).images
num_views = out.shape[0] // 2
mv_images = out[num_views:].permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
mv_normals = out[:num_views].permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
return (mv_images, mv_normals, )
def prepare_data(self, ref_image, ref_mask):
single_image = torch_img_to_pil_rgba(ref_image, ref_mask)
abs_fix_cam_pose_dir = os.path.join(ROOT_PATH, self.fix_cam_pose_dir)
dataset = SingleImageDataset(fix_cam_pose_dir=abs_fix_cam_pose_dir, num_views=6, img_wh=[256, 256], bg_color='white', single_image=single_image)
return dataset[0]
class MVDream_Model:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mvdream_pipe": ("DIFFUSERS_PIPE",),
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"prompt": ("STRING", {
"default": "",
"multiline": True
}),
"prompt_neg": ("STRING", {
"default": "ugly, blurry, pixelated obscure, unnatural colors, poor lighting, dull, unclear, cropped, lowres, low quality, artifacts, duplicate",
"multiline": True
}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"mv_guidance_scale": ("FLOAT", {"default": 5.0, "min": 0.0, "step": 0.01}),
"num_inference_steps": ("INT", {"default": 30, "min": 1}),
"elevation": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"multiview_images",
)
FUNCTION = "run_mvdream"
CATEGORY = "Comfy3D/Algorithm"
def run_mvdream(
self,
mvdream_pipe,
reference_image, # [1, H, W, 3]
reference_mask, # [1, H, W]
prompt,
prompt_neg,
seed,
mv_guidance_scale,
num_inference_steps,
elevation,
):
if len(reference_image.shape) == 4:
reference_image = reference_image.squeeze(0)
if len(reference_mask.shape) == 3:
reference_mask = reference_mask.squeeze(0)
generator = torch.manual_seed(seed)
reference_mask = reference_mask.unsqueeze(2)
# give the white background to reference_image
reference_image = (reference_image * reference_mask + (1 - reference_mask)).detach().cpu().numpy()
# generate multi-view images
mv_images = mvdream_pipe(prompt, reference_image, generator=generator, negative_prompt=prompt_neg, guidance_scale=mv_guidance_scale, num_inference_steps=num_inference_steps, elevation=elevation)
mv_images = torch.from_numpy(np.stack([mv_images[1], mv_images[2], mv_images[3], mv_images[0]], axis=0)).float() # [4, H, W, 3], float32
return (mv_images, )
class Load_Large_Multiview_Gaussian_Model:
checkpoints_dir = "checkpoints/lgm"
default_ckpt_name = "model_fp16.safetensors"
default_repo_id = "ashawkey/LGM"
@classmethod
def INPUT_TYPES(cls):
cls.checkpoints_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS)
if cls.default_ckpt_name not in all_models_names:
all_models_names += [cls.default_ckpt_name]
return {
"required": {
"model_name": (all_models_names, ),
"lgb_config": (['big', 'default', 'small', 'tiny'], )
},
}
RETURN_TYPES = (
"LGM_MODEL",
)
RETURN_NAMES = (
"lgm_model",
)
FUNCTION = "load_LGM"
CATEGORY = "Comfy3D/Import|Export"
def load_LGM(self, model_name, lgb_config):
lgm_model = LGM(config_defaults[lgb_config])
# resume pretrained checkpoint
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
if not os.path.isfile(ckpt_path):
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
if ckpt_path.endswith('safetensors'):
ckpt = load_file(ckpt_path, device='cpu')
else:
ckpt = torch.load(ckpt_path, map_location='cpu')
lgm_model.load_state_dict(ckpt, strict=False)
# device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
lgm_model = lgm_model.half().to(device)
lgm_model.eval()
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
return (lgm_model, )
IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
class Large_Multiview_Gaussian_Model:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"multiview_images": ("IMAGE", ),
"lgm_model": ("LGM_MODEL", ),
},
}
OUTPUT_NODE = True
RETURN_TYPES = (
"GS_PLY",
)
RETURN_NAMES = (
"gs_ply",
)
FUNCTION = "run_LGM"
CATEGORY = "Comfy3D/Algorithm"
@torch.no_grad()
def run_LGM(self, multiview_images, lgm_model):
device = "cuda"
ref_image_torch = prepare_torch_img(multiview_images, lgm_model.opt.input_size, lgm_model.opt.input_size, device) # [4, 3, 256, 256]
ref_image_torch = TF.normalize(ref_image_torch, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
rays_embeddings = lgm_model.prepare_default_rays(device)
ref_image_torch = torch.cat([ref_image_torch, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, 256, 256]
with torch.autocast(device_type=device, dtype=WEIGHT_DTYPE):
# generate gaussians
gaussians = lgm_model.forward_gaussians(ref_image_torch)
# convert gaussians to ply
gs_ply = lgm_model.gs.to_ply(gaussians)
return (gs_ply, )
class Convert_3DGS_to_Mesh_with_NeRF_and_Marching_Cubes:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"gs_ply": ("GS_PLY",),
"gs_config": (['big', 'default', 'small', 'tiny'], ),
"force_cuda_rast": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = (
"MESH",
"IMAGE",
"MASK",
)
RETURN_NAMES = (
"mesh",
"imgs",
"alphas",
)
FUNCTION = "convert_gs_ply"
CATEGORY = "Comfy3D/Algorithm"
def convert_gs_ply(self, gs_ply, gs_config, force_cuda_rast):
with torch.inference_mode(False):
chosen_config = config_defaults[gs_config]
chosen_config.force_cuda_rast = force_cuda_rast
converter = GSConverterNeRFMarchingCubes(config_defaults[gs_config], gs_ply).cuda()
imgs, alphas = converter.fit_nerf()
converter.fit_mesh()
converter.fit_mesh_uv()
return(converter.get_mesh(), imgs, alphas)
class NeuS:
NeuS_config_path = "configs/NeuS.conf"
fix_cam_pose_dir = "NeuS/models/fixed_poses"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"reference_normals": ("IMAGE",),
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}), # longer time, better result. 1w will be ok for most cases
"batch_size": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
"learning_rate": ("FLOAT", {"default": 5e-4, "min": 0.00001, "step": 0.00001}),
"learning_rate_alpha": ("FLOAT", {"default": 5e-4, "min": 0.00001, "step": 0.00001}),
"color_loss_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
"mesh_smoothing_weight": ("FLOAT", {"default": 0.1, "min": 0.0, "step": 0.01}),
"mask_loss_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
"normal_loss_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
"sparse_loss_weight": ("FLOAT", {"default": 0.1, "min": 0.0, "step": 0.01}),
"warm_up_end": ("INT", {"default": 500, "min": 0, "max": 0xffffffffffffffff}),
"anneal_end": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"use_white_background": ("BOOLEAN", {"default": True},),
"geometry_extract_resolution": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
"marching_cude_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "step": 0.01}),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_NeuS"
CATEGORY = "Comfy3D/Algorithm"
def run_NeuS(
self,
reference_image, # [num_views, H, W, 3]
reference_mask, # [num_views, H, W]
reference_normals, # [num_views, H, W, 3]
training_iterations,
batch_size,
learning_rate,
learning_rate_alpha,
color_loss_weight,
mesh_smoothing_weight,
mask_loss_weight,
normal_loss_weight,
sparse_loss_weight,
warm_up_end,
anneal_end,
use_white_background,
geometry_extract_resolution,
marching_cude_threshold
):
mesh = None
num_views = reference_image.shape[0]
if num_views in [6, 5, 4]:
config_path = os.path.join(ROOT_PATH, self.NeuS_config_path)
abs_fix_cam_pose_dir = os.path.join(ROOT_PATH, self.fix_cam_pose_dir)
with torch.inference_mode(False):
NeuS_params = NeuSParams(
training_iterations,
batch_size,
learning_rate,
learning_rate_alpha,
color_loss_weight,
mesh_smoothing_weight,
mask_loss_weight,
normal_loss_weight,
sparse_loss_weight,
warm_up_end,
anneal_end,
use_white_background
)
runner = NeuSRunner(
reference_image,
reference_mask,
reference_normals,
config_path,
abs_fix_cam_pose_dir,
num_views,
NeuS_params
)
cstr(f"[{self.__class__.__name__}] Training NeuS...").msg.print()
runner.train()
cstr(f"[{self.__class__.__name__}] Extracting Mesh...").msg.print()
mesh = runner.extract_mesh(resolution=geometry_extract_resolution, threshold=marching_cude_threshold)
else:
cstr(f"[{self.__class__.__name__}] Number of views must be one of the following: 6, 5, 4. But got {num_views}").error.print()
return(mesh, )
class Load_TripoSR_Model:
checkpoints_dir = "checkpoints/tsr"
default_ckpt_name = "model.ckpt"
default_repo_id = "stabilityai/TripoSR"
tsr_config_path = "configs/tsr_config.yaml"
@classmethod
def INPUT_TYPES(cls):
cls.checkpoints_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS)
if cls.default_ckpt_name not in all_models_names:
all_models_names += [cls.default_ckpt_name]
return {
"required": {
"model_name": (all_models_names, ),
"chunk_size": ("INT", {"default": 8192, "min": 1, "max": 10000})
},
}
RETURN_TYPES = (
"TSR_MODEL",
)
RETURN_NAMES = (
"tsr_model",
)
FUNCTION = "load_TSR"
CATEGORY = "Comfy3D/Import|Export"
def load_TSR(self, model_name, chunk_size):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# resume pretrained checkpoint
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
if not os.path.isfile(ckpt_path):
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
tsr_model = TSR.from_pretrained(
weight_path = ckpt_path,
config_path = os.path.join(ROOT_PATH, self.tsr_config_path)
)
tsr_model.renderer.set_chunk_size(chunk_size)
tsr_model.to(device)
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
return (tsr_model, )
class TripoSR:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"tsr_model": ("TSR_MODEL", ),
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"geometry_extract_resolution": ("INT", {"default": 256, "min": 1, "max": 0xffffffffffffffff}),
"marching_cude_threshold": ("FLOAT", {"default": 25.0, "min": 0.0, "step": 0.01}),
}
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_TSR"
CATEGORY = "Comfy3D/Algorithm"
@torch.no_grad()
def run_TSR(self, tsr_model, reference_image, reference_mask, geometry_extract_resolution, marching_cude_threshold):
cstr(f"[{self.__class__.__name__}] Running TripoSR...").msg.print()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
mesh = None
image = reference_image[0]
mask = reference_mask[0].unsqueeze(2)
image = torch.cat((image, mask), dim=2).detach().cpu().numpy()
image = Image.fromarray(np.clip(255. * image, 0, 255).astype(np.uint8))
image = self.fill_background(image)
image = image.convert('RGB')
scene_codes = tsr_model([image], device)
meshes = tsr_model.extract_mesh(scene_codes, resolution=geometry_extract_resolution, threshold=marching_cude_threshold)
mesh = Mesh.load_trimesh(given_mesh=meshes[0])
return (mesh,)
# Default model are trained on images with this background
def fill_background(self, image):
image = np.array(image).astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = Image.fromarray((image * 255.0).astype(np.uint8))
return image
class Load_CRM_MVDiffusion_Model:
checkpoints_dir = "checkpoints/crm"
default_ckpt_name = ["pixel-diffusion.pth", "ccm-diffusion.pth"]
default_conf_name = ["configs/crm_configs/sd_v2_base_ipmv_zero_SNR.yaml", "configs/crm_configs/sd_v2_base_ipmv_chin8_zero_snr.yaml"]
default_repo_id = "Zhengyi/CRM"
@classmethod
def INPUT_TYPES(cls):
cls.checkpoints_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS)
for ckpt_name in cls.default_ckpt_name:
if ckpt_name not in all_models_names:
all_models_names += [ckpt_name]
return {
"required": {
"model_name": (all_models_names, ),
"crm_config_path": (cls.default_conf_name, ),
},
}
RETURN_TYPES = (
"CRM_MVDIFFUSION_SAMPLER",
)
RETURN_NAMES = (
"crm_mvdiffusion_sampler",
)
FUNCTION = "load_CRM"
CATEGORY = "Comfy3D/Import|Export"
def load_CRM(self, model_name, crm_config_path):
from .crm.imagedream.ldm.util import (
instantiate_from_config,
get_obj_from_str,
)
if not os.path.isabs(crm_config_path):
crm_config_path = os.path.join(ROOT_PATH, crm_config_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# resume pretrained checkpoint
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
if not os.path.isfile(ckpt_path):
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
crm_config = OmegaConf.load(crm_config_path)
crm_mvdiffusion_model = instantiate_from_config(crm_config.model)
crm_mvdiffusion_model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
crm_mvdiffusion_model = crm_mvdiffusion_model.to(device).to(WEIGHT_DTYPE)
crm_mvdiffusion_model.device = device
crm_mvdiffusion_sampler = get_obj_from_str(crm_config.sampler.target)(
crm_mvdiffusion_model, device=device, dtype=WEIGHT_DTYPE, **crm_config.sampler.params
)
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
return (crm_mvdiffusion_sampler, )
class CRM_Images_MVDiffusion_Model:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"crm_mvdiffusion_sampler": ("CRM_MVDIFFUSION_SAMPLER",),
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"prompt": ("STRING", {
"default": "3D assets",
"multiline": True
}),
"prompt_neg": ("STRING", {
"default": "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.",
"multiline": True
}),
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
"mv_guidance_scale": ("FLOAT", {"default": 5.5, "min": 0.0, "step": 0.01}),
"num_inference_steps": ("INT", {"default": 50, "min": 1}),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"multiview_images",
)
FUNCTION = "run_model"
CATEGORY = "Comfy3D/Algorithm"
def run_model(
self,
crm_mvdiffusion_sampler,
reference_image, # [1, H, W, 3]
reference_mask, # [1, H, W]
prompt,
prompt_neg,
seed,
mv_guidance_scale,
num_inference_steps,
):
pixel_img = torch_img_to_pil_rgba(reference_image, reference_mask)
pixel_img = CRMSampler.process_pixel_img(pixel_img)
multiview_images = CRMSampler.stage1_sample(
crm_mvdiffusion_sampler,
pixel_img,
prompt,
prompt_neg,
seed,
mv_guidance_scale,
num_inference_steps
)
return (multiview_images, )
class CRM_CCMs_MVDiffusion_Model:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"crm_mvdiffusion_sampler": ("CRM_MVDIFFUSION_SAMPLER",),
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"multiview_images": ("IMAGE",),
"prompt": ("STRING", {
"default": "3D assets",
"multiline": True
}),
"prompt_neg": ("STRING", {
"default": "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.",
"multiline": True
}),
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
"mv_guidance_scale": ("FLOAT", {"default": 5.5, "min": 0.0, "step": 0.01}),
"num_inference_steps": ("INT", {"default": 50, "min": 1}),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"multiview_CCMs",
)
FUNCTION = "run_model"
CATEGORY = "Comfy3D/Algorithm"
def run_model(
self,
crm_mvdiffusion_sampler,
reference_image, # [1, H, W, 3]
reference_mask, # [1, H, W]
multiview_images, # [6, H, W, 3]
prompt,
prompt_neg,
seed,
mv_guidance_scale,
num_inference_steps,
):
pixel_img = torch_img_to_pil_rgba(reference_image, reference_mask)
pixel_img = CRMSampler.process_pixel_img(pixel_img)
multiview_CCMs = CRMSampler.stage2_sample(
crm_mvdiffusion_sampler,
pixel_img,
multiview_images,
prompt,
prompt_neg,
seed,
mv_guidance_scale,
num_inference_steps
)
return(multiview_CCMs, )
class Load_Convolutional_Reconstruction_Model:
checkpoints_dir = "checkpoints/crm"
default_ckpt_name = "CRM.pth"
default_repo_id = "Zhengyi/CRM"
crm_config_path = "configs/crm_configs/specs_objaverse_total.json"
@classmethod
def INPUT_TYPES(cls):
cls.checkpoints_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS)
if cls.default_ckpt_name not in all_models_names:
all_models_names += [cls.default_ckpt_name]
return {
"required": {
"model_name": (all_models_names, ),
},
}
RETURN_TYPES = (
"CRM_MODEL",
)
RETURN_NAMES = (
"crm_model",
)
FUNCTION = "load_CRM"
CATEGORY = "Comfy3D/Import|Export"
def load_CRM(self, model_name):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# resume pretrained checkpoint
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
if not os.path.isfile(ckpt_path):
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
crm_conf = json.load(open(os.path.join(ROOT_PATH, self.crm_config_path)))
crm_model = CRM(crm_conf).to(device)
crm_model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
return (crm_model, )
class Convolutional_Reconstruction_Model:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"crm_model": ("CRM_MODEL", ),
"multiview_images": ("IMAGE",),
"multiview_CCMs": ("IMAGE",),
}
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_CRM"
CATEGORY = "Comfy3D/Algorithm"
@torch.no_grad()
def run_CRM(self, crm_model, multiview_images, multiview_CCMs):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
np_imgs = np.concatenate(multiview_images.cpu().numpy(), 1) # (256, 256*6==1536, 3)
np_xyzs = np.concatenate(multiview_CCMs.cpu().numpy(), 1) # (256, 1536, 3)
mesh = CRMSampler.generate3d(crm_model, np_imgs, np_xyzs, device)
return (mesh,)
class Instant_NGP:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"training_iterations": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
"training_resolution": ("INT", {"default": 128, "min": 128, "max": 8192}),
"marching_cude_grids_resolution": ("INT", {"default": 256, "min": 1, "max": 0xffffffffffffffff}),
"marching_cude_grids_batch_size": ("INT", {"default": 128, "min": 1, "max": 0xffffffffffffffff}),
"marching_cude_threshold": ("FLOAT", {"default": 10.0, "min": 0.0, "step": 0.01}),
"texture_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"force_cuda_rast": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_instant_ngp"
CATEGORY = "Comfy3D/Algorithm"
def run_instant_ngp(
self,
reference_image,
reference_mask,
reference_orbit_camera_poses,
reference_orbit_camera_fovy,
training_iterations,
training_resolution,
marching_cude_grids_resolution,
marching_cude_grids_batch_size,
marching_cude_threshold,
texture_resolution,
force_cuda_rast
):
with torch.inference_mode(False):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ngp = InstantNGP(training_resolution).to(device)
ngp.prepare_training(reference_image, reference_mask, reference_orbit_camera_poses, reference_orbit_camera_fovy)
ngp.fit_nerf(training_iterations)
vertices, triangles = marching_cubes_density_to_mesh(ngp.get_density, marching_cude_grids_resolution, marching_cude_grids_batch_size, marching_cude_threshold)
v = torch.from_numpy(vertices).contiguous().float().to(device)
f = torch.from_numpy(triangles).contiguous().int().to(device)
mesh = Mesh(v=v, f=f, albedo=None, device=device)
mesh.auto_normal()
mesh.auto_uv()
mesh.albedo = color_func_to_albedo(mesh, ngp.get_color, texture_resolution, device=device, force_cuda_rast=force_cuda_rast)
return (mesh, )
Reference in New Issue View Git Blame Copy Permalink