Add bake texture related nodes and classes
This commit is contained in:
MrForExample
+5
-5
@@ -1,11 +1,9 @@
|
||||
#__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
import importlib
|
||||
import inspect
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
|
||||
nodes_filename = "nodes"
|
||||
module = importlib.import_module(f".{nodes_filename}", package=__name__)
|
||||
for name, cls in inspect.getmembers(module, inspect.isclass):
|
||||
@@ -16,4 +14,6 @@ for name, cls in inspect.getmembers(module, inspect.isclass):
|
||||
disp = f"{name}"
|
||||
|
||||
NODE_CLASS_MAPPINGS[node] = cls
|
||||
NODE_DISPLAY_NAME_MAPPINGS[node] = disp
|
||||
NODE_DISPLAY_NAME_MAPPINGS[node] = disp
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -7,7 +7,7 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
from ..mesh_processer.mesh import Mesh, safe_normalize
|
||||
from ..mesh_processer.mesh import safe_normalize
|
||||
|
||||
def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'):
|
||||
assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
|
||||
@@ -38,41 +38,41 @@ def make_divisible(x, m=8):
|
||||
return int(math.ceil(x / m) * m)
|
||||
|
||||
class Renderer(nn.Module):
|
||||
def __init__(self, opt):
|
||||
def __init__(self, mesh, force_cuda_rast):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.opt = opt
|
||||
self.mesh = mesh
|
||||
|
||||
self.mesh = Mesh.load(self.opt.mesh, resize=False)
|
||||
|
||||
if not self.opt.force_cuda_rast and (not self.opt.gui or os.name == 'nt'):
|
||||
self.glctx = dr.RasterizeGLContext()
|
||||
else:
|
||||
if force_cuda_rast or os.name != 'nt':
|
||||
self.glctx = dr.RasterizeCudaContext()
|
||||
else:
|
||||
self.glctx = dr.RasterizeGLContext()
|
||||
|
||||
# extract trainable parameters
|
||||
self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v))
|
||||
self.raw_albedo = nn.Parameter(trunc_rev_sigmoid(self.mesh.albedo))
|
||||
self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
|
||||
self.raw_albedo = nn.Parameter(trunc_rev_sigmoid(self.mesh.albedo), requires_grad=True)
|
||||
|
||||
|
||||
def get_params(self):
|
||||
def get_params(self, texture_lr, train_geo, geom_lr):
|
||||
|
||||
params = [
|
||||
{'params': self.raw_albedo, 'lr': self.opt.texture_lr},
|
||||
{'params': self.raw_albedo, 'lr': texture_lr},
|
||||
]
|
||||
|
||||
if self.opt.train_geo:
|
||||
params.append({'params': self.v_offsets, 'lr': self.opt.geom_lr})
|
||||
self.train_geo = train_geo
|
||||
if train_geo:
|
||||
params.append({'params': self.v_offsets, 'lr': geom_lr})
|
||||
|
||||
return params
|
||||
|
||||
@torch.no_grad()
|
||||
def export_mesh(self, save_path):
|
||||
self.mesh.v = (self.mesh.v + self.v_offsets).detach()
|
||||
self.mesh.albedo = torch.sigmoid(self.raw_albedo.detach())
|
||||
self.mesh.write(save_path)
|
||||
|
||||
@torch.no_grad()
|
||||
def update_mesh(self):
|
||||
self.mesh.v = (self.mesh.v + self.v_offsets).detach()
|
||||
self.mesh.albedo = torch.sigmoid(self.raw_albedo.detach())
|
||||
|
||||
def render(self, pose, proj, h0, w0, ssaa=1, bg_color=1, texture_filter='linear-mipmap-linear'):
|
||||
|
||||
@@ -86,7 +86,7 @@ class Renderer(nn.Module):
|
||||
results = {}
|
||||
|
||||
# get v
|
||||
if self.opt.train_geo:
|
||||
if self.train_geo:
|
||||
v = self.mesh.v + self.v_offsets # [N, 3]
|
||||
else:
|
||||
v = self.mesh.v
|
||||
@@ -106,9 +106,10 @@ class Renderer(nn.Module):
|
||||
|
||||
texc, texc_db = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft, rast_db=rast_db, diff_attrs='all')
|
||||
albedo = dr.texture(self.raw_albedo.unsqueeze(0), texc, uv_da=texc_db, filter_mode=texture_filter) # [1, H, W, 3]
|
||||
print(f"albedo = dr.texture: {albedo.requires_grad}")
|
||||
albedo = torch.sigmoid(albedo)
|
||||
# get vn and render normal
|
||||
if self.opt.train_geo:
|
||||
if self.train_geo:
|
||||
i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long()
|
||||
v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :]
|
||||
|
||||
@@ -133,6 +134,7 @@ class Renderer(nn.Module):
|
||||
|
||||
# antialias
|
||||
albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).squeeze(0) # [H, W, 3]
|
||||
print(f"albedo = dr.antialias: {albedo.requires_grad}")
|
||||
albedo = alpha * albedo + (1 - alpha) * bg_color
|
||||
|
||||
# ssaa
|
||||
|
||||
+133
-2
@@ -1,4 +1,135 @@
|
||||
import sys
|
||||
import random
|
||||
|
||||
from ..shared_utils.common_utils import get_persistent_directory
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
from ..diff_rast.diff_mesh_renderer import Renderer
|
||||
from ..shared_utils.camera_utils import orbit_camera, OrbitCamera
|
||||
|
||||
class DiffTextureBaker:
|
||||
|
||||
def __init__(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, mesh,
|
||||
training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize):
|
||||
self.device = torch.device("cuda")
|
||||
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
self.ref_cam_fovy = reference_orbit_camera_fovy
|
||||
|
||||
self.ref_size_H = reference_images[0].shape[0]
|
||||
self.ref_size_W = reference_images[0].shape[1]
|
||||
|
||||
self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
|
||||
|
||||
# prepare main components for optimization
|
||||
self.renderer = Renderer(mesh, force_cuda_rasterize).to(self.device)
|
||||
|
||||
self.optimizer = torch.optim.Adam(self.renderer.get_params(texture_learning_rate, train_mesh_geometry, geometry_learning_rate))
|
||||
#self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.lambda_ssim = ms_ssim_loss_weight
|
||||
|
||||
self.training_iterations = training_iterations
|
||||
|
||||
self.batch_size = batch_size
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
ref_masks_torch_list = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_torch_list.append(self.prepare_img(reference_images[i]))
|
||||
ref_masks_torch_list.append(self.prepare_img(reference_masks[i].unsqueeze(2)))
|
||||
|
||||
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
|
||||
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
|
||||
|
||||
def prepare_img(self, img):
|
||||
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
|
||||
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
|
||||
return img_new
|
||||
|
||||
def training(self):
|
||||
starter = torch.cuda.Event(enable_timing=True)
|
||||
ender = torch.cuda.Event(enable_timing=True)
|
||||
starter.record()
|
||||
|
||||
ref_imgs_masked = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
|
||||
|
||||
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
||||
|
||||
for step in range(self.training_iterations):
|
||||
|
||||
### calculate loss between reference and rendered image from known view
|
||||
loss = 0
|
||||
masked_rendered_img_batch = []
|
||||
masked_ref_img_batch = []
|
||||
for _ in range(self.batch_size):
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
|
||||
|
||||
# render output
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
ref_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
ref_cam = (ref_pose, self.cam.perspective)
|
||||
out = self.renderer.render(*ref_cam, self.ref_size_H, self.ref_size_W, ssaa=1) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
|
||||
|
||||
image = out["image"] # [H, W, 3] in [0, 1]
|
||||
image = image.permute(2, 0, 1).contiguous() # [3, H, W] in [0, 1]
|
||||
|
||||
#print(f"image.requires_grad: {image.requires_grad}")
|
||||
|
||||
image_masked = (image * self.ref_masks_torch[i]).unsqueeze(0)
|
||||
|
||||
#print(f"image_masked.requires_grad: {image_masked.requires_grad}")
|
||||
#print(f"ref_imgs_masked[i].requires_grad: {ref_imgs_masked[i].requires_grad}")
|
||||
|
||||
masked_rendered_img_batch.append(image_masked)
|
||||
masked_ref_img_batch.append(ref_imgs_masked[i])
|
||||
|
||||
masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
|
||||
masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
|
||||
|
||||
# rgb loss
|
||||
loss += (1 - self.lambda_ssim) * F.mse_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
|
||||
|
||||
# D-SSIM loss
|
||||
# [1, 3, H, W] in [0, 1]
|
||||
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
|
||||
loss += self.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
|
||||
|
||||
print(f"masked_rendered_img_batch_torch.requires_grad: {masked_rendered_img_batch_torch.requires_grad}")
|
||||
print(f"masked_ref_img_batch_torch.requires_grad: {masked_ref_img_batch_torch.requires_grad}")
|
||||
|
||||
print(f"loss.requires_grad: {loss.requires_grad}")
|
||||
|
||||
print(f"self.renderer.raw_albedo.requires_grad: {self.renderer.raw_albedo.requires_grad}")
|
||||
|
||||
# import kiui
|
||||
# kiui.lo(hor, ver)
|
||||
# kiui.vis.plot_image(image)
|
||||
|
||||
# optimize step
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.need_update = True
|
||||
|
||||
print(f"Step: {step}")
|
||||
|
||||
self.renderer.update_mesh()
|
||||
|
||||
ender.record()
|
||||
t = starter.elapsed_time(ender)
|
||||
|
||||
def get_mesh_and_texture(self):
|
||||
return (self.renderer.mesh, self.renderer.mesh.albedo, )
|
||||
@@ -49,7 +49,7 @@ class Mesh:
|
||||
if path is None:
|
||||
mesh = cls(**kwargs)
|
||||
# obj supports face uv
|
||||
elif path.endswith(".obj"):
|
||||
elif path.lower().endswith(".obj"):
|
||||
mesh = cls.load_obj(path, **kwargs)
|
||||
# trimesh only supports vertex uv, but can load more formats
|
||||
else:
|
||||
@@ -421,11 +421,11 @@ class Mesh:
|
||||
return self
|
||||
|
||||
def write(self, path):
|
||||
if path.endswith(".ply"):
|
||||
if path.lower().endswith(".ply"):
|
||||
self.write_ply(path)
|
||||
elif path.endswith(".obj"):
|
||||
elif path.lower().endswith(".obj"):
|
||||
self.write_obj(path)
|
||||
elif path.endswith(".glb") or path.endswith(".gltf"):
|
||||
elif path.lower().endswith(".glb") or path.lower().endswith(".gltf"):
|
||||
self.write_glb(path)
|
||||
else:
|
||||
raise NotImplementedError(f"format {path} not supported!")
|
||||
|
||||
@@ -1,12 +1,9 @@
|
||||
import nodes
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import os
|
||||
import math
|
||||
import re
|
||||
import folder_paths as comfy_paths
|
||||
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
from .diff_rast.diff_texturing import DiffTextureBaker
|
||||
from .shared_utils.common_utils import cstr
|
||||
from .mesh_processer.mesh import Mesh
|
||||
|
||||
@@ -18,7 +15,13 @@ MANIFEST = {
|
||||
"description": "An extensive node suite that enables ComfyUI to process 3D inputs (Mesh & UV Texture, etc) using cutting edge algorithms (3DGS, NeRF, etc.)",
|
||||
}
|
||||
|
||||
class Load_3D_File:
|
||||
SUPPORTED_3D_EXTENSIONS = (
|
||||
'.obj',
|
||||
'.ply',
|
||||
'.glb',
|
||||
)
|
||||
|
||||
class Load_3D_Mesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -38,27 +41,136 @@ class Load_3D_File:
|
||||
"mesh",
|
||||
)
|
||||
FUNCTION = "load_mesh"
|
||||
CATEGORY = "ComfyUI3D"
|
||||
|
||||
SUPPORTED_3D_EXTENSIONS = (
|
||||
'.obj',
|
||||
'.ply',
|
||||
'.glb',
|
||||
)
|
||||
CATEGORY = "ComfyUI3D/Import|Export"
|
||||
|
||||
def load_mesh(self, mesh_file_path, resize, renormal, retex):
|
||||
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, file = os.path.split(self.filepath)
|
||||
if file.lower().endswith(self.SUPPORTED_3D_EXTENSIONS):
|
||||
mesh = Mesh.load(self.opt.mesh, resize, renormal, retex)
|
||||
folder, filename = os.path.split(mesh_file_path)
|
||||
if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
mesh = Mesh.load(mesh_file_path, resize, renormal, retex)
|
||||
else:
|
||||
cstr(f"File name {file} does not end with supported 3D file extensions: {self.SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
else:
|
||||
cstr(f"File {mesh_file_path} does not exist").error.print()
|
||||
cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print()
|
||||
return (mesh, )
|
||||
|
||||
class Save_3D_Mesh:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh": ("MESH",),
|
||||
"mesh_save_path": ("STRING", {"default": 'Mesh_[time(%Y-%m-%d)].obj', "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_mesh"
|
||||
CATEGORY = "ComfyUI3D/Import|Export"
|
||||
|
||||
def save_mesh(self, mesh, mesh_save_path):
|
||||
|
||||
mesh_folder_path, filename = os.path.split(mesh_save_path)
|
||||
|
||||
if not os.path.isabs(mesh_save_path):
|
||||
mesh_folder_path = os.path.join(comfy_paths.output_directory, mesh_folder_path)
|
||||
mesh_save_path = os.path.join(mesh_folder_path, filename)
|
||||
|
||||
os.makedirs(mesh_folder_path, exist_ok=True)
|
||||
|
||||
if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
|
||||
mesh.export_mesh(mesh_save_path)
|
||||
cstr(f"[{self.__class__.__name__}] saved model to {mesh_save_path}").msg.print()
|
||||
else:
|
||||
cstr(f"File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print()
|
||||
|
||||
return {"ui": {"mesh": mesh}}
|
||||
|
||||
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
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
OUTPUT_NODE = 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 = "ComfyUI3D/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:
|
||||
|
||||
@classmethod
|
||||
@@ -76,7 +188,7 @@ class Gaussian_Splatting:
|
||||
"raw_3DGS",
|
||||
)
|
||||
FUNCTION = "run_3DGS"
|
||||
CATEGORY = "ComfyUI3D"
|
||||
CATEGORY = "ComfyUI3D/Algorithm"
|
||||
|
||||
def run_3DGS(self, reference_images):
|
||||
raw_3DGS = None
|
||||
@@ -85,25 +197,66 @@ class Gaussian_Splatting:
|
||||
|
||||
class Bake_Texture_To_Mesh:
|
||||
|
||||
def __init__(self):
|
||||
self.need_update = False
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"reference_images": ("IMAGE",),
|
||||
"mesh": ("MESH")
|
||||
"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",),
|
||||
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}),
|
||||
"batch_size": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}),
|
||||
"texture_learning_rate": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"train_mesh_geometry": ("BOOLEAN", {"default": False},),
|
||||
"geometry_learning_rate": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 0.1, "step": 0.001}),
|
||||
"ms_ssim_loss_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"force_cuda_rasterize": ("BOOLEAN", {"default": False},),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"MESH",
|
||||
"IMAGE",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"baked_texture",
|
||||
"trained_mesh",
|
||||
"baked_texture", # [1, H, W, 3]
|
||||
)
|
||||
FUNCTION = "bake_texture"
|
||||
CATEGORY = "ComfyUI3D"
|
||||
CATEGORY = "ComfyUI3D/Algorithm"
|
||||
|
||||
def bake_texture(self, reference_images, mesh):
|
||||
def bake_texture(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, mesh,
|
||||
training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize):
|
||||
|
||||
trained_mesh = None
|
||||
baked_texture = None
|
||||
|
||||
return (baked_texture, )
|
||||
|
||||
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
|
||||
|
||||
texture_baker = DiffTextureBaker(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, mesh,
|
||||
training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize)
|
||||
|
||||
texture_baker.training()
|
||||
|
||||
trained_mesh, baked_texture = texture_baker.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, )
|
||||
@@ -76,9 +76,9 @@ class cstr(str):
|
||||
print(self, **kwargs)
|
||||
|
||||
#! MESSAGE TEMPLATES
|
||||
cstr.color.add_code("msg", f"{cstr.color.BLUE}WAS Node Suite: {cstr.color.END}")
|
||||
cstr.color.add_code("warning", f"{cstr.color.BLUE}WAS Node Suite {cstr.color.LIGHTYELLOW}[WARNING] {cstr.color.END}")
|
||||
cstr.color.add_code("error", f"{cstr.color.RED}WAS Node Suite {cstr.color.END}[ERROR] {cstr.color.END}")
|
||||
cstr.color.add_code("msg", f"{cstr.color.BLUE}[ComfyUI-3D] {cstr.color.END}")
|
||||
cstr.color.add_code("warning", f"{cstr.color.LIGHTYELLOW}[ComfyUI-3D] [WARNING] {cstr.color.END}")
|
||||
cstr.color.add_code("error", f"{cstr.color.RED}[ComfyUI-3D] [ERROR] {cstr.color.END}")
|
||||
|
||||
def get_persistent_directory(folder_name):
|
||||
if sys.platform == "win32":
|
||||
|
||||
Reference in New Issue
Block a user