1545 lines
55 KiB
Python
1545 lines
55 KiB
Python
import base64
|
|
import io
|
|
import os
|
|
import struct
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from server import BinaryEventTypes, PromptServer
|
|
|
|
class Hy3DtoTD:
|
|
#https://github.com/kijai/ComfyUI-Hunyuan3DWrapper
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"trimesh": ("TRIMESH",),
|
|
"mesh": ("MESH",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_mesh"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_mesh(self, broadcast, trimesh=None, mesh=None):
|
|
try:
|
|
if trimesh is not None and mesh is not None:
|
|
raise ValueError("Please connect only one input: trimesh or mesh, not both.")
|
|
|
|
if trimesh is not None:
|
|
vertices, colors = self.process_trimesh(trimesh)
|
|
elif mesh is not None:
|
|
vertices, colors = self.process_comfy_mesh(mesh)
|
|
else:
|
|
raise ValueError("Please connect either a trimesh or mesh input.")
|
|
|
|
return self._send_to_td(vertices, colors, broadcast)
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing mesh: {str(e)}")
|
|
|
|
|
|
def process_trimesh(self, mesh):
|
|
try:
|
|
import trimesh
|
|
|
|
if isinstance(mesh, trimesh.Scene):
|
|
meshes = []
|
|
for m in mesh.geometry.values():
|
|
if isinstance(m, trimesh.Trimesh):
|
|
meshes.append(m)
|
|
|
|
if meshes:
|
|
mesh = trimesh.util.concatenate(meshes)
|
|
else:
|
|
raise ValueError("No valid mesh data found")
|
|
|
|
if not isinstance(mesh, trimesh.Trimesh):
|
|
raise ValueError("Input data is not a valid Trimesh object")
|
|
|
|
mesh.fix_normals()
|
|
|
|
vertices = np.asarray(mesh.vertices, dtype=np.float32)
|
|
vertex_count = len(vertices)
|
|
|
|
if hasattr(mesh.visual, "vertex_colors") and mesh.visual.vertex_colors is not None:
|
|
colors = np.asarray(mesh.visual.vertex_colors)
|
|
elif hasattr(mesh.visual, "uv") and mesh.visual.uv is not None:
|
|
try:
|
|
mesh.visual = mesh.visual.to_color()
|
|
colors = np.asarray(mesh.visual.vertex_colors)
|
|
except Exception:
|
|
colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
|
|
else:
|
|
colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
|
|
|
|
colors = self._normalize_colors(colors, vertex_count)
|
|
|
|
return vertices, colors
|
|
|
|
except ImportError:
|
|
raise ImportError("Please install the trimesh library: pip install trimesh")
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing trimesh: {str(e)}")
|
|
|
|
def process_comfy_mesh(self, mesh):
|
|
try:
|
|
if isinstance(mesh, list):
|
|
if len(mesh) == 0:
|
|
raise ValueError("Empty mesh list received")
|
|
mesh = mesh[0]
|
|
|
|
vertices = self._extract_comfy_mesh_vertices(mesh)
|
|
vertex_count = len(vertices)
|
|
|
|
colors = self._extract_comfy_mesh_colors(mesh, vertex_count)
|
|
colors = self._normalize_colors(colors, vertex_count)
|
|
|
|
return vertices, colors
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing ComfyUI mesh: {str(e)}")
|
|
|
|
def _extract_comfy_mesh_vertices(self, mesh):
|
|
vertices = None
|
|
|
|
if hasattr(mesh, "vertices") and mesh.vertices is not None:
|
|
vertices = mesh.vertices
|
|
elif isinstance(mesh, dict) and "vertices" in mesh:
|
|
vertices = mesh["vertices"]
|
|
elif isinstance(mesh, (tuple, list)) and len(mesh) > 0:
|
|
vertices = mesh[0]
|
|
|
|
if vertices is None:
|
|
raise ValueError("Cannot find vertices in ComfyUI mesh")
|
|
|
|
vertices = self._to_numpy(vertices)
|
|
vertices = self._squeeze_batch(vertices)
|
|
vertices = np.asarray(vertices, dtype=np.float32)
|
|
|
|
if vertices.ndim != 2 or vertices.shape[1] < 3:
|
|
raise ValueError(f"Invalid vertices shape: {vertices.shape}")
|
|
|
|
return vertices[:, :3]
|
|
|
|
def _extract_comfy_mesh_colors(self, mesh, vertex_count):
|
|
colors = None
|
|
|
|
if hasattr(mesh, "vertex_colors") and mesh.vertex_colors is not None:
|
|
colors = mesh.vertex_colors
|
|
elif hasattr(mesh, "colors") and mesh.colors is not None:
|
|
colors = mesh.colors
|
|
elif hasattr(mesh, "visual") and mesh.visual is not None:
|
|
if hasattr(mesh.visual, "vertex_colors") and mesh.visual.vertex_colors is not None:
|
|
colors = mesh.visual.vertex_colors
|
|
elif isinstance(mesh, dict) and "vertex_colors" in mesh:
|
|
colors = mesh["vertex_colors"]
|
|
elif isinstance(mesh, dict) and "colors" in mesh:
|
|
colors = mesh["colors"]
|
|
|
|
if colors is None:
|
|
colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
|
|
|
|
return colors
|
|
|
|
def _to_numpy(self, value):
|
|
if value is None:
|
|
return None
|
|
|
|
if hasattr(value, "detach"):
|
|
value = value.detach()
|
|
if hasattr(value, "cpu"):
|
|
value = value.cpu()
|
|
if hasattr(value, "numpy"):
|
|
value = value.numpy()
|
|
|
|
return np.asarray(value)
|
|
|
|
def _squeeze_batch(self, value):
|
|
value = np.asarray(value)
|
|
|
|
if value.ndim == 3 and value.shape[0] == 1:
|
|
value = value[0]
|
|
|
|
return value
|
|
|
|
def _normalize_colors(self, colors, vertex_count):
|
|
colors = self._to_numpy(colors)
|
|
|
|
if colors is None:
|
|
colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
|
|
|
|
colors = self._squeeze_batch(colors)
|
|
colors = np.asarray(colors)
|
|
|
|
if colors.ndim == 1:
|
|
if colors.shape[0] in (3, 4):
|
|
colors = np.tile(colors, (vertex_count, 1))
|
|
else:
|
|
colors = np.column_stack([colors, colors, colors])
|
|
|
|
if colors.ndim != 2:
|
|
colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
|
|
|
|
if len(colors) != vertex_count:
|
|
if len(colors) > vertex_count:
|
|
colors = colors[:vertex_count]
|
|
elif len(colors) > 0:
|
|
last_color = colors[-1]
|
|
padding = np.tile(last_color, (vertex_count - len(colors), 1))
|
|
colors = np.vstack([colors, padding])
|
|
else:
|
|
colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
|
|
|
|
if colors.dtype != np.uint8:
|
|
if colors.size > 0 and colors.max() <= 1.0:
|
|
colors = (np.clip(colors, 0, 1) * 255).astype(np.uint8)
|
|
else:
|
|
colors = np.clip(colors, 0, 255).astype(np.uint8)
|
|
|
|
if colors.shape[1] == 1:
|
|
colors = np.repeat(colors, 3, axis=1)
|
|
|
|
if colors.shape[1] == 2:
|
|
colors = np.repeat(colors[:, :1], 3, axis=1)
|
|
|
|
if colors.shape[1] == 3:
|
|
alpha = np.full((vertex_count, 1), 255, dtype=np.uint8)
|
|
colors = np.concatenate([colors, alpha], axis=1)
|
|
|
|
if colors.shape[1] > 4:
|
|
colors = colors[:, :4]
|
|
|
|
if colors.shape[1] < 4:
|
|
default_colors = np.tile([200, 200, 200, 255], (vertex_count, 1)).astype(np.uint8)
|
|
default_colors[:, :colors.shape[1]] = colors
|
|
colors = default_colors
|
|
|
|
return colors[:, :4].astype(np.uint8)
|
|
|
|
def _send_to_td(self, vertices, colors, broadcast):
|
|
try:
|
|
vertices = np.asarray(vertices, dtype=np.float32)
|
|
vertex_count = len(vertices)
|
|
colors = self._normalize_colors(colors, vertex_count)
|
|
|
|
dtype_vertex = np.dtype([
|
|
("pos", "<f4", 3),
|
|
("color", "u1", 4)
|
|
])
|
|
|
|
vertex_data = np.empty(vertex_count, dtype=dtype_vertex)
|
|
vertex_data["pos"] = vertices[:, :3]
|
|
vertex_data["color"] = colors[:, :4]
|
|
|
|
header = (
|
|
"ply\n"
|
|
"format binary_little_endian 1.0\n"
|
|
f"element vertex {vertex_count}\n"
|
|
"property float x\n"
|
|
"property float y\n"
|
|
"property float z\n"
|
|
"property uchar red\n"
|
|
"property uchar green\n"
|
|
"property uchar blue\n"
|
|
"property uchar alpha\n"
|
|
"end_header\n"
|
|
).encode("ascii")
|
|
|
|
binary_data = header + vertex_data.tobytes()
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(1000, binary_data, sid=sid)
|
|
|
|
return {
|
|
"ui": {
|
|
"mesh": [{
|
|
"source": "websocket",
|
|
"content-type": "model/ply",
|
|
"type": "output",
|
|
}]
|
|
}
|
|
}
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Error sending point PLY to TouchDesigner: {str(e)}")
|
|
|
|
class Tripo3DtoTD:
|
|
#https://github.com/VAST-AI-Research/ComfyUI-Tripo
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_file": ("STRING",),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_mesh"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_mesh(self, model_file, broadcast):
|
|
try:
|
|
processed_mesh = self.process_mesh(model_file)
|
|
return self._send_to_td(processed_mesh, broadcast)
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing mesh: {str(e)}")
|
|
|
|
def process_mesh(self, model_file):
|
|
try:
|
|
import trimesh
|
|
if not model_file.startswith(('http://', 'https://')) and not os.path.isabs(model_file):
|
|
try:
|
|
import folder_paths
|
|
output_dir = folder_paths.get_output_directory()
|
|
local_path = os.path.join(output_dir, model_file)
|
|
if os.path.exists(local_path):
|
|
model_file = local_path
|
|
except:
|
|
pass
|
|
|
|
mesh = trimesh.load(model_file)
|
|
|
|
if isinstance(mesh, trimesh.Scene):
|
|
meshes = []
|
|
for m in mesh.geometry.values():
|
|
if isinstance(m, trimesh.Trimesh):
|
|
self._process_vertex_colors(m)
|
|
meshes.append(m)
|
|
|
|
if meshes:
|
|
mesh = trimesh.util.concatenate(meshes)
|
|
else:
|
|
raise ValueError("No valid mesh data found")
|
|
|
|
if not isinstance(mesh, trimesh.Trimesh):
|
|
raise ValueError("Input data is not a valid Trimesh object")
|
|
|
|
mesh.fix_normals()
|
|
|
|
if not hasattr(mesh.visual, 'vertex_colors') or mesh.visual.vertex_colors is None:
|
|
mesh.visual = trimesh.visual.ColorVisuals(mesh)
|
|
mesh.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(mesh.vertices), 1))
|
|
else:
|
|
if mesh.visual.vertex_colors.max() <= 1.0:
|
|
mesh.visual.vertex_colors = (mesh.visual.vertex_colors * 255).astype(np.uint8)
|
|
else:
|
|
mesh.visual.vertex_colors = mesh.visual.vertex_colors.astype(np.uint8)
|
|
|
|
return mesh
|
|
|
|
except ImportError:
|
|
raise ImportError("Please install the trimesh library: pip install trimesh")
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing mesh: {str(e)}")
|
|
|
|
def _process_vertex_colors(self, m):
|
|
try:
|
|
import trimesh
|
|
|
|
if not hasattr(m.visual, 'vertex_colors') or m.visual.vertex_colors is None:
|
|
if hasattr(m.visual, 'material'):
|
|
material = m.visual.material
|
|
if hasattr(material, 'baseColorTexture') and material.baseColorTexture is not None:
|
|
if hasattr(m.visual, 'uv') and m.visual.uv is not None:
|
|
self._extract_colors_from_texture(m)
|
|
elif hasattr(material, 'baseColorFactor'):
|
|
color = np.array(material.baseColorFactor)
|
|
if len(color) == 3:
|
|
color = np.append(color, 1.0)
|
|
color = np.clip(color, 0, 1)
|
|
color = (color * 255).astype(np.uint8)
|
|
m.visual = trimesh.visual.ColorVisuals(m)
|
|
m.visual.vertex_colors = np.tile(color, (len(m.vertices), 1))
|
|
else:
|
|
m.visual = trimesh.visual.ColorVisuals(m)
|
|
m.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(m.vertices), 1))
|
|
else:
|
|
m.visual = trimesh.visual.ColorVisuals(m)
|
|
m.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(m.vertices), 1))
|
|
else:
|
|
if m.visual.vertex_colors.max() <= 1.0:
|
|
m.visual.vertex_colors = (m.visual.vertex_colors * 255).astype(np.uint8)
|
|
else:
|
|
m.visual.vertex_colors = m.visual.vertex_colors.astype(np.uint8)
|
|
except ImportError:
|
|
raise ImportError("Please install the trimesh library: pip install trimesh")
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing vertex colors: {str(e)}")
|
|
|
|
def _extract_colors_from_texture(self, m):
|
|
try:
|
|
import trimesh
|
|
|
|
material = m.visual.material
|
|
image = material.baseColorTexture
|
|
if image is not None:
|
|
image_data = np.array(image)
|
|
if image_data.dtype != np.uint8:
|
|
image_data = (np.clip(image_data, 0, 1) * 255).astype(np.uint8)
|
|
|
|
uvs = m.visual.uv
|
|
uvs = np.clip(uvs, 0, 1)
|
|
h, w = image_data.shape[0], image_data.shape[1]
|
|
i = ((1 - uvs[:, 1]) * (h - 1)).astype(int)
|
|
j = (uvs[:, 0] * (w - 1)).astype(int)
|
|
i = np.clip(i, 0, h - 1)
|
|
j = np.clip(j, 0, w - 1)
|
|
colors = image_data[i, j, :4]
|
|
m.visual = trimesh.visual.ColorVisuals(m)
|
|
m.visual.vertex_colors = colors
|
|
else:
|
|
m.visual = trimesh.visual.ColorVisuals(m)
|
|
m.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(m.vertices), 1))
|
|
except ImportError:
|
|
raise ImportError("Please install the trimesh library: pip install trimesh")
|
|
except Exception as e:
|
|
raise ValueError(f"Error extracting colors from texture: {str(e)}")
|
|
|
|
def _send_to_td(self, mesh, broadcast):
|
|
try:
|
|
if not hasattr(mesh.visual, 'vertex_colors') or mesh.visual.vertex_colors is None:
|
|
colors = np.tile([200, 200, 200, 255], (len(mesh.vertices), 1)).astype(np.uint8)
|
|
else:
|
|
colors = mesh.visual.vertex_colors.astype(np.uint8)
|
|
|
|
header = [
|
|
"ply",
|
|
"format binary_little_endian 1.0",
|
|
f"element vertex {len(mesh.vertices)}",
|
|
"property float x",
|
|
"property float y",
|
|
"property float z",
|
|
"property uchar red",
|
|
"property uchar green",
|
|
"property uchar blue",
|
|
"property uchar alpha",
|
|
"end_header\n"
|
|
]
|
|
header = '\n'.join(header).encode('ascii')
|
|
|
|
vertices = mesh.vertices.astype('<f4')
|
|
colors = colors[:, :4].astype(np.uint8)
|
|
|
|
vertex_count = len(vertices)
|
|
binary_data = bytearray()
|
|
for i in range(vertex_count):
|
|
binary_data.extend(struct.pack('<fff', vertices[i, 0], vertices[i, 1], vertices[i, 2]))
|
|
binary_data.extend(colors[i, :4].tobytes())
|
|
|
|
buffer = io.BytesIO()
|
|
buffer.write(header)
|
|
buffer.write(binary_data)
|
|
binary_output = buffer.getvalue()
|
|
buffer.close()
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(1000, binary_output, sid=sid)
|
|
|
|
return {"ui": {
|
|
"mesh": [{
|
|
"source": "websocket",
|
|
"content-type": "model/ply",
|
|
"type": "output",
|
|
}]
|
|
}}
|
|
except ImportError:
|
|
raise ImportError("Please install the trimesh library: pip install trimesh")
|
|
except Exception as e:
|
|
raise ValueError(f"Error sending mesh: {str(e)}")
|
|
|
|
|
|
class ImagetoTD:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_images(self, images, broadcast):
|
|
results = []
|
|
for tensor in images:
|
|
array = 255.0 * tensor.cpu().numpy()
|
|
image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(
|
|
BinaryEventTypes.UNENCODED_PREVIEW_IMAGE,
|
|
["PNG", image, None],
|
|
sid,
|
|
)
|
|
results.append({
|
|
"source": "websocket",
|
|
"content-type": "image/png",
|
|
"type": "output",
|
|
})
|
|
|
|
return {"ui": {"images": results}}
|
|
|
|
|
|
class ImagetoTD_JPEG:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
"quality": ("INT", {"default": 95, "min": 1, "max": 100}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_images(self, images, broadcast, quality):
|
|
results = []
|
|
for tensor in images:
|
|
array = 255.0 * tensor.cpu().numpy()
|
|
image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
|
|
|
|
bytesIO = io.BytesIO()
|
|
type_num = 1
|
|
header = struct.pack(">I", type_num)
|
|
bytesIO.write(header)
|
|
|
|
image.save(bytesIO, format="JPEG", quality=quality)
|
|
jpeg_bytes = bytesIO.getvalue()
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
|
|
server.send_sync(
|
|
BinaryEventTypes.PREVIEW_IMAGE,
|
|
jpeg_bytes,
|
|
sid,
|
|
)
|
|
results.append({
|
|
"source": "websocket",
|
|
"content-type": "image/jpeg",
|
|
"type": "output",
|
|
})
|
|
|
|
return {"ui": {"images": results}}
|
|
|
|
class Comfy3DPacktoTD:
|
|
#https://github.com/MrForExample/ComfyUI-3D-Pack
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESH",),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_to_td"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_to_td(self, mesh, broadcast):
|
|
try:
|
|
vertices = mesh.v.cpu().numpy().astype(np.float32)
|
|
colors = None
|
|
color_source = "none"
|
|
|
|
if hasattr(mesh, 'albedo') and mesh.albedo is not None:
|
|
albedo = mesh.albedo
|
|
if hasattr(albedo, 'cpu'):
|
|
albedo = albedo.cpu()
|
|
if hasattr(albedo, 'numpy'):
|
|
albedo = albedo.numpy()
|
|
|
|
if len(albedo.shape) > 1:
|
|
if hasattr(mesh, 'vt') and mesh.vt is not None:
|
|
uvs = mesh.vt.cpu().numpy()
|
|
uvs = np.clip(uvs, 0, 1)
|
|
tex_h, tex_w = albedo.shape[:2]
|
|
|
|
mapped_u, mapped_v = uvs[:, 0], uvs[:, 1]
|
|
|
|
tex_x = mapped_u * (tex_w - 1)
|
|
tex_y = mapped_v * (tex_h - 1)
|
|
|
|
tex_x0 = np.floor(tex_x).astype(np.int32)
|
|
tex_y0 = np.floor(tex_y).astype(np.int32)
|
|
tex_x1 = np.minimum(tex_x0 + 1, tex_w - 1)
|
|
tex_y1 = np.minimum(tex_y0 + 1, tex_h - 1)
|
|
|
|
wx = tex_x - tex_x0
|
|
wy = tex_y - tex_y0
|
|
|
|
c00 = albedo[tex_y0, tex_x0]
|
|
c01 = albedo[tex_y0, tex_x1]
|
|
c10 = albedo[tex_y1, tex_x0]
|
|
c11 = albedo[tex_y1, tex_x1]
|
|
|
|
c0 = c00 * (1 - wx[:, np.newaxis]) + c01 * wx[:, np.newaxis]
|
|
c1 = c10 * (1 - wx[:, np.newaxis]) + c11 * wx[:, np.newaxis]
|
|
colors = c0 * (1 - wy[:, np.newaxis]) + c1 * wy[:, np.newaxis]
|
|
|
|
color_source = "albedo_with_uv_no_v_flip"
|
|
else:
|
|
if albedo.shape[0] == len(vertices):
|
|
colors = albedo
|
|
color_source = "albedo_direct"
|
|
else:
|
|
colors = albedo
|
|
color_source = "albedo_scalar"
|
|
|
|
if colors is None and hasattr(mesh, 'vc') and mesh.vc is not None:
|
|
colors = mesh.vc.cpu().numpy()
|
|
color_source = "vertex_colors"
|
|
|
|
if colors is None:
|
|
positions = vertices - vertices.min(axis=0)
|
|
positions = positions / positions.max(axis=0) if positions.max() > 0 else positions
|
|
colors = positions
|
|
color_source = "generated"
|
|
|
|
if colors is not None:
|
|
if colors.max() <= 1.0:
|
|
colors = (colors * 255).astype(np.uint8)
|
|
else:
|
|
colors = np.clip(colors, 0, 255).astype(np.uint8)
|
|
|
|
if colors.shape[1] == 3:
|
|
alpha = np.full((len(vertices), 1), 255, dtype=np.uint8)
|
|
colors = np.concatenate([colors, alpha], axis=1)
|
|
|
|
header = [
|
|
"ply",
|
|
"format binary_little_endian 1.0",
|
|
f"element vertex {len(vertices)}",
|
|
"property float x",
|
|
"property float y",
|
|
"property float z",
|
|
"property uchar red",
|
|
"property uchar green",
|
|
"property uchar blue",
|
|
"property uchar alpha",
|
|
"end_header\n"
|
|
]
|
|
header = '\n'.join(header).encode('ascii')
|
|
|
|
binary_data = bytearray()
|
|
for i in range(len(vertices)):
|
|
binary_data.extend(struct.pack('<fff',
|
|
float(vertices[i, 0]), float(vertices[i, 1]), float(vertices[i, 2])))
|
|
binary_data.extend(bytes([
|
|
int(colors[i, 0]),
|
|
int(colors[i, 1]),
|
|
int(colors[i, 2]),
|
|
int(colors[i, 3])
|
|
]))
|
|
|
|
buffer = io.BytesIO()
|
|
buffer.write(header)
|
|
buffer.write(binary_data)
|
|
binary_output = buffer.getvalue()
|
|
buffer.close()
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(1000, binary_output, sid=sid)
|
|
|
|
return {"ui": {
|
|
"mesh": [{
|
|
"source": "websocket",
|
|
"content-type": "model/ply",
|
|
"type": "output",
|
|
}]
|
|
}}
|
|
|
|
except Exception as e:
|
|
print(f"[Error] send_to_td: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise ValueError(f"Error sending mesh: {str(e)}")
|
|
|
|
|
|
class TripoSRtoTD:
|
|
#https://github.com/flowtyone/ComfyUI-Flowty-TripoSR
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESH",),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_to_td"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_to_td(self, mesh, broadcast):
|
|
try:
|
|
if isinstance(mesh, list):
|
|
if len(mesh) > 0:
|
|
mesh_obj = mesh[0]
|
|
else:
|
|
raise ValueError("Empty mesh list received")
|
|
else:
|
|
mesh_obj = mesh
|
|
|
|
if hasattr(mesh_obj, 'vertices'):
|
|
vertices = mesh_obj.vertices
|
|
if hasattr(vertices, 'cpu'):
|
|
vertices = vertices.cpu()
|
|
if hasattr(vertices, 'numpy'):
|
|
vertices = vertices.numpy()
|
|
vertices = vertices.astype(np.float32)
|
|
else:
|
|
raise ValueError("Cannot find vertex data in mesh object")
|
|
|
|
rotated_vertices = vertices.copy()
|
|
rotated_vertices[:, 0] = -vertices[:, 0]
|
|
rotated_vertices[:, 1], rotated_vertices[:, 2] = vertices[:, 2], vertices[:, 1]
|
|
|
|
final_vertices = rotated_vertices.copy()
|
|
final_vertices[:, 0] = rotated_vertices[:, 2]
|
|
final_vertices[:, 2] = -rotated_vertices[:, 0]
|
|
|
|
vertices = final_vertices
|
|
|
|
colors = None
|
|
|
|
if hasattr(mesh_obj, 'visual') and mesh_obj.visual is not None:
|
|
visual = mesh_obj.visual
|
|
|
|
if hasattr(visual, 'vertex_colors') and visual.vertex_colors is not None:
|
|
colors = visual.vertex_colors
|
|
if hasattr(colors, 'cpu'):
|
|
colors = colors.cpu()
|
|
if hasattr(colors, 'numpy'):
|
|
colors = colors.numpy()
|
|
|
|
elif hasattr(visual, 'face_colors') and visual.face_colors is not None:
|
|
face_colors = visual.face_colors
|
|
if hasattr(face_colors, 'cpu'):
|
|
face_colors = face_colors.cpu()
|
|
if hasattr(face_colors, 'numpy'):
|
|
face_colors = face_colors.numpy()
|
|
|
|
if hasattr(mesh_obj, 'faces') and mesh_obj.faces is not None:
|
|
faces = mesh_obj.faces
|
|
if hasattr(faces, 'cpu'):
|
|
faces = faces.cpu()
|
|
if hasattr(faces, 'numpy'):
|
|
faces = faces.numpy()
|
|
|
|
colors = np.zeros((len(vertices), 4), dtype=np.uint8)
|
|
vertex_color_count = np.zeros(len(vertices), dtype=np.int32)
|
|
|
|
for i, face in enumerate(faces):
|
|
for vertex_idx in face:
|
|
colors[vertex_idx] += face_colors[i]
|
|
vertex_color_count[vertex_idx] += 1
|
|
|
|
for i in range(len(vertices)):
|
|
if vertex_color_count[i] > 0:
|
|
colors[i] = colors[i] // vertex_color_count[i]
|
|
|
|
elif hasattr(visual, 'material') and visual.material is not None:
|
|
material = visual.material
|
|
|
|
if hasattr(material, 'diffuse') and material.diffuse is not None:
|
|
diffuse = material.diffuse
|
|
colors = np.tile(diffuse, (len(vertices), 1))
|
|
|
|
if colors is None and hasattr(mesh_obj, '_visual') and mesh_obj._visual is not None:
|
|
visual = mesh_obj._visual
|
|
|
|
if hasattr(visual, 'vertex_colors') and visual.vertex_colors is not None:
|
|
colors = visual.vertex_colors
|
|
if hasattr(colors, 'cpu'):
|
|
colors = colors.cpu()
|
|
if hasattr(colors, 'numpy'):
|
|
colors = colors.numpy()
|
|
|
|
if colors is None and hasattr(mesh_obj, 'visual') and hasattr(mesh_obj.visual,
|
|
'uv') and mesh_obj.visual.uv is not None:
|
|
if hasattr(mesh_obj.visual, 'material') and hasattr(mesh_obj.visual.material, 'image'):
|
|
uvs = mesh_obj.visual.uv
|
|
image = mesh_obj.visual.material.image
|
|
|
|
h, w = image.shape[:2]
|
|
u = np.clip(uvs[:, 0], 0, 1) * (w - 1)
|
|
v = np.clip(uvs[:, 1], 0, 1) * (h - 1)
|
|
|
|
u_int = np.round(u).astype(np.int32)
|
|
v_int = np.round(v).astype(np.int32)
|
|
|
|
u_int = np.clip(u_int, 0, w - 1)
|
|
v_int = np.clip(v_int, 0, h - 1)
|
|
|
|
colors = image[v_int, u_int]
|
|
|
|
if colors is None:
|
|
np.random.seed(42)
|
|
colors = np.random.randint(100, 200, size=(len(vertices), 3), dtype=np.uint8)
|
|
alpha = np.full((len(vertices), 1), 255, dtype=np.uint8)
|
|
colors = np.concatenate([colors, alpha], axis=1)
|
|
|
|
if colors is not None:
|
|
if len(colors.shape) == 1:
|
|
colors = np.column_stack([colors, colors, colors])
|
|
|
|
if colors.max() <= 1.0 and colors.dtype != np.uint8:
|
|
colors = (colors * 255).astype(np.uint8)
|
|
else:
|
|
colors = np.clip(colors, 0, 255).astype(np.uint8)
|
|
|
|
if colors.shape[1] == 3:
|
|
alpha = np.full((len(vertices), 1), 255, dtype=np.uint8)
|
|
colors = np.concatenate([colors, alpha], axis=1)
|
|
elif colors.shape[1] > 4:
|
|
colors = colors[:, :4]
|
|
|
|
if len(colors) != len(vertices):
|
|
if len(colors) > len(vertices):
|
|
colors = colors[:len(vertices)]
|
|
else:
|
|
last_color = colors[-1]
|
|
padding = np.tile(last_color, (len(vertices) - len(colors), 1))
|
|
colors = np.vstack([colors, padding])
|
|
|
|
header = [
|
|
"ply",
|
|
"format binary_little_endian 1.0",
|
|
f"element vertex {len(vertices)}",
|
|
"property float x",
|
|
"property float y",
|
|
"property float z",
|
|
"property uchar red",
|
|
"property uchar green",
|
|
"property uchar blue",
|
|
"property uchar alpha",
|
|
"end_header\n"
|
|
]
|
|
header = '\n'.join(header).encode('ascii')
|
|
|
|
binary_data = bytearray()
|
|
for i in range(len(vertices)):
|
|
binary_data.extend(struct.pack('<fff',
|
|
float(vertices[i, 0]), float(vertices[i, 1]), float(vertices[i, 2])))
|
|
binary_data.extend(bytes([
|
|
int(colors[i, 0]),
|
|
int(colors[i, 1]),
|
|
int(colors[i, 2]),
|
|
int(colors[i, 3])
|
|
]))
|
|
|
|
buffer = io.BytesIO()
|
|
buffer.write(header)
|
|
buffer.write(binary_data)
|
|
binary_output = buffer.getvalue()
|
|
buffer.close()
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(1000, binary_output, sid=sid)
|
|
|
|
return {"ui": {
|
|
"mesh": [{
|
|
"source": "websocket",
|
|
"content-type": "model/ply",
|
|
"type": "output",
|
|
}]
|
|
}}
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise ValueError(f"Error sending mesh: {str(e)}")
|
|
|
|
|
|
class LoadTDImage:
|
|
# This node originates from the comfyui-tooling-nodes and utilizes the "Load Image (Base64)" functionality.
|
|
# https://github.com/Acly/comfyui-tooling-nodes
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image": ("STRING", {"multiline": False})}}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
CATEGORY = "TouchDesigner"
|
|
FUNCTION = "load_image"
|
|
|
|
def load_image(self, image):
|
|
try:
|
|
imgdata = base64.b64decode(image)
|
|
img = Image.open(io.BytesIO(imgdata))
|
|
|
|
width, height = img.size
|
|
|
|
if "A" in img.getbands():
|
|
mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
|
|
mask = 1.0 - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((height, width), dtype=torch.float32, device="cpu")
|
|
|
|
if len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0)
|
|
|
|
img = img.convert("RGB")
|
|
img = np.array(img).astype(np.float32) / 255.0
|
|
img = torch.from_numpy(img)[None,]
|
|
|
|
print(f"Image shape: {img.shape}, Mask shape: {mask.shape}")
|
|
|
|
return (img, mask)
|
|
|
|
except Exception as e:
|
|
print(f"Error in LoadTDImage: {str(e)}")
|
|
default_img = torch.zeros((1, 3, 64, 64), dtype=torch.float32, device="cpu")
|
|
default_mask = torch.zeros((1, 64, 64), dtype=torch.float32, device="cpu")
|
|
return (default_img, default_mask)
|
|
|
|
|
|
class VideotoTD:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"frame_rate": ("FLOAT", {"default": 8, "min": 1, "step": 1}),
|
|
"quality": ("INT", {"default": 75, "min": 15, "max": 100, "step": 1}),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "process_video"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def process_video(self, images, frame_rate, broadcast, quality):
|
|
if images is None or (isinstance(images, torch.Tensor) and images.size(0) == 0):
|
|
raise ValueError("No valid images provided")
|
|
|
|
video_binary = self._encode_video_ffmpeg(images, frame_rate, quality)
|
|
|
|
return self._send_to_td(video_binary, broadcast)
|
|
|
|
def _encode_video_ffmpeg(self, images, frame_rate, quality):
|
|
try:
|
|
import imageio_ffmpeg
|
|
|
|
with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
|
|
temp_path = temp_file.name
|
|
|
|
if isinstance(images, torch.Tensor):
|
|
if len(images.shape) != 4:
|
|
raise ValueError(f"Expected 4D tensor, got shape {images.shape}")
|
|
|
|
images_np = images.cpu().numpy()
|
|
if images_np.max() <= 1.0:
|
|
images_np = (images_np * 255).astype(np.uint8)
|
|
else:
|
|
images_np = images_np.astype(np.uint8)
|
|
else:
|
|
raise TypeError("Images must be a torch.Tensor")
|
|
|
|
height, width = images_np[0].shape[:2]
|
|
|
|
output_params = [
|
|
'-c:v', 'libx264',
|
|
'-preset', 'medium',
|
|
'-crf', str(51 - quality // 2),
|
|
'-pix_fmt', 'yuv420p'
|
|
]
|
|
|
|
writer = imageio_ffmpeg.write_frames(
|
|
temp_path,
|
|
(width, height),
|
|
fps=frame_rate,
|
|
output_params=output_params
|
|
)
|
|
|
|
writer.send(None)
|
|
|
|
for img in images_np:
|
|
if img.shape[2] == 4: # RGBA
|
|
img = img[:, :, :3]
|
|
|
|
if not img.flags['C_CONTIGUOUS']:
|
|
img = np.ascontiguousarray(img)
|
|
|
|
writer.send(img)
|
|
|
|
writer.close()
|
|
|
|
with open(temp_path, 'rb') as f:
|
|
video_binary = f.read()
|
|
|
|
os.unlink(temp_path)
|
|
|
|
return video_binary
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error encoding video with ffmpeg: {str(e)}")
|
|
|
|
def _send_to_td(self, video_binary, broadcast):
|
|
try:
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(1001, video_binary, sid=sid)
|
|
|
|
return {"ui": {
|
|
"video": [{
|
|
"source": "websocket",
|
|
"content-type": "video/mp4",
|
|
"type": "output",
|
|
}]
|
|
}}
|
|
except Exception as e:
|
|
raise ValueError(f"Error sending video: {str(e)}")
|
|
|
|
class AudiotoTD:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "process_audio"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def process_audio(self, audio, broadcast):
|
|
if audio is None or not isinstance(audio, dict):
|
|
raise ValueError("No valid audio dict provided")
|
|
|
|
if "waveform" not in audio or "sample_rate" not in audio:
|
|
raise ValueError("Audio dict must contain 'waveform' and 'sample_rate'")
|
|
|
|
waveform = audio["waveform"]
|
|
sample_rate = int(audio["sample_rate"])
|
|
|
|
if not isinstance(waveform, torch.Tensor):
|
|
raise TypeError("audio['waveform'] must be a torch.Tensor")
|
|
|
|
if waveform.ndim != 3:
|
|
raise ValueError(
|
|
f"Expected waveform as 3D tensor [B, C, T], got shape {tuple(waveform.shape)}"
|
|
)
|
|
|
|
if waveform.size(0) > 1:
|
|
waveform = waveform[0:1, ...]
|
|
|
|
waveform = waveform[0]
|
|
|
|
with torch.no_grad():
|
|
max_abs = waveform.abs().amax()
|
|
if torch.isfinite(max_abs) and max_abs > 1.0:
|
|
waveform = waveform / max_abs
|
|
|
|
wav_np = waveform.detach().cpu().float().numpy()
|
|
|
|
if wav_np.ndim != 2:
|
|
raise ValueError(
|
|
f"Waveform must be [C, T] after squeeze, got {wav_np.shape}"
|
|
)
|
|
|
|
wav_np = np.transpose(wav_np, (1, 0))
|
|
|
|
audio_binary = self._encode_wav(wav_np, sample_rate)
|
|
|
|
return self._send_to_td(audio_binary, broadcast)
|
|
|
|
def _encode_wav(self, wav_np, sample_rate):
|
|
import wave
|
|
|
|
if wav_np.ndim == 1:
|
|
wav_np = wav_np[:, None]
|
|
|
|
channels = wav_np.shape[1]
|
|
|
|
pcm16 = np.clip(wav_np, -1.0, 1.0)
|
|
pcm16 = (pcm16 * 32767.0).astype("<i2")
|
|
|
|
if not pcm16.flags["C_CONTIGUOUS"]:
|
|
pcm16 = np.ascontiguousarray(pcm16)
|
|
|
|
buf = io.BytesIO()
|
|
|
|
with wave.open(buf, "wb") as wf:
|
|
wf.setnchannels(channels)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(sample_rate)
|
|
wf.writeframes(pcm16.tobytes())
|
|
|
|
audio_binary = buf.getvalue()
|
|
buf.close()
|
|
|
|
return audio_binary
|
|
|
|
def _send_to_td(self, audio_binary, broadcast):
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
|
|
server.send_sync(1002, audio_binary, sid=sid)
|
|
|
|
return {
|
|
"ui": {
|
|
"audio": [{
|
|
"source": "websocket",
|
|
"content-type": "audio/wav",
|
|
"type": "output",
|
|
}]
|
|
}
|
|
}
|
|
|
|
class GaussianSplattingtoTD:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {},
|
|
"optional": {
|
|
"splat": ("SPLAT", {
|
|
"tooltip": "Optional input source. Choose either this SPLAT input or ply_path. If both are provided."
|
|
}),
|
|
"ply_path": ("STRING", {
|
|
"forceInput": True,
|
|
"default": "",
|
|
"tooltip": "Optional input source. Choose either a PLY file path or the SPLAT input. If both are provided,."
|
|
}),
|
|
"broadcast": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Send to all connected ComfyUI frontend clients. If disabled, the data is sent only to the current client."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_gaussian_splatting"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_gaussian_splatting(self, splat=None, ply_path="", broadcast=False):
|
|
try:
|
|
has_splat = splat is not None
|
|
has_ply = bool(ply_path and ply_path.strip())
|
|
|
|
if has_splat and has_ply:
|
|
print(
|
|
"[GaussianSplattingtoTD] Both SPLAT and ply_path were provided. "
|
|
"Using SPLAT input and ignoring ply_path."
|
|
)
|
|
|
|
if has_splat:
|
|
gs_data = self.process_splat(splat)
|
|
elif has_ply:
|
|
gs_data = self.process_ply(ply_path)
|
|
else:
|
|
raise ValueError(
|
|
"No input source provided. Please connect a SPLAT input or provide a ply_path. "
|
|
"Choose one of them."
|
|
)
|
|
|
|
return self._send_to_td(gs_data, broadcast)
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing Gaussian Splatting data: {str(e)}")
|
|
|
|
def process_splat(self, splat):
|
|
def t2n(x):
|
|
if x is None:
|
|
return None
|
|
if isinstance(x, torch.Tensor):
|
|
return x.detach().cpu().float().numpy()
|
|
return np.asarray(x, dtype=np.float32)
|
|
|
|
positions = t2n(getattr(splat, 'positions'))
|
|
scales = t2n(getattr(splat, 'scales'))
|
|
rotations = t2n(getattr(splat, 'rotations'))
|
|
opacities = t2n(getattr(splat, 'opacities'))
|
|
sh = t2n(getattr(splat, 'sh'))
|
|
|
|
counts = getattr(splat, 'counts', None)
|
|
if counts is not None:
|
|
counts = t2n(counts)
|
|
|
|
if positions is None:
|
|
raise ValueError("SPLAT is missing required field: positions")
|
|
if scales is None:
|
|
raise ValueError("SPLAT is missing required field: scales")
|
|
if rotations is None:
|
|
raise ValueError("SPLAT is missing required field: rotations")
|
|
if opacities is None:
|
|
raise ValueError("SPLAT is missing required field: opacities")
|
|
if sh is None:
|
|
raise ValueError("SPLAT is missing required field: sh")
|
|
|
|
# Validate batched dimensions
|
|
if positions.ndim != 3:
|
|
raise ValueError(f"Expected positions shape (B, N, 3), got {positions.shape}")
|
|
if scales.ndim != 3:
|
|
raise ValueError(f"Expected scales shape (B, N, 3), got {scales.shape}")
|
|
if rotations.ndim != 3:
|
|
raise ValueError(f"Expected rotations shape (B, N, 4), got {rotations.shape}")
|
|
if opacities.ndim != 3:
|
|
raise ValueError(f"Expected opacities shape (B, N, 1), got {opacities.shape}")
|
|
if sh.ndim != 4:
|
|
raise ValueError(f"Expected sh shape (B, N, K, 3), got {sh.shape}")
|
|
|
|
positions = positions[0] # (N, 3)
|
|
scales = scales[0] # (N, 3)
|
|
rotations = rotations[0] # (N, 4)
|
|
opacities = opacities[0] # (N, 1)
|
|
sh = sh[0] # (N, K, 3)
|
|
|
|
# Use counts[0] to remove padded points if available
|
|
if counts is not None:
|
|
n_real = int(counts[0])
|
|
positions = positions[:n_real]
|
|
scales = scales[:n_real]
|
|
rotations = rotations[:n_real]
|
|
opacities = opacities[:n_real]
|
|
sh = sh[:n_real]
|
|
|
|
f_dc = sh[:, 0, :].astype(np.float32)
|
|
|
|
# Convert activated SPLAT values back to standard 3DGS PLY parameter space:
|
|
eps = 1e-6
|
|
|
|
op = np.clip(opacities.reshape(-1), eps, 1.0 - eps).astype(np.float32)
|
|
opacity_logit = np.log(op / (1.0 - op)) # inverse sigmoid
|
|
|
|
sc = np.clip(scales, eps, None).astype(np.float32)
|
|
scale_log = np.log(sc) # inverse exp
|
|
|
|
gs_data = {
|
|
'positions': positions.astype(np.float32),
|
|
'colors': f_dc,
|
|
'opacity': opacity_logit,
|
|
'scales': scale_log,
|
|
'rotations': rotations.astype(np.float32),
|
|
}
|
|
|
|
return gs_data
|
|
|
|
def process_ply(self, ply_path):
|
|
try:
|
|
import os
|
|
import numpy as np
|
|
from plyfile import PlyData
|
|
|
|
if not ply_path.startswith(('http://', 'https://')) and not os.path.isabs(ply_path):
|
|
try:
|
|
import folder_paths
|
|
output_dir = folder_paths.get_output_directory()
|
|
local_path = os.path.join(output_dir, ply_path)
|
|
if os.path.exists(local_path):
|
|
ply_path = local_path
|
|
except Exception:
|
|
pass
|
|
|
|
if not os.path.exists(ply_path):
|
|
raise ValueError(f"PLY file not found: {ply_path}")
|
|
|
|
plydata = PlyData.read(ply_path)
|
|
vertex = plydata['vertex']
|
|
|
|
required_props = [
|
|
'x', 'y', 'z',
|
|
'f_dc_0', 'f_dc_1', 'f_dc_2',
|
|
'opacity',
|
|
'scale_0', 'scale_1', 'scale_2',
|
|
'rot_0', 'rot_1', 'rot_2', 'rot_3'
|
|
]
|
|
|
|
missing_props = [
|
|
prop for prop in required_props
|
|
if prop not in vertex.data.dtype.names
|
|
]
|
|
|
|
if missing_props:
|
|
raise ValueError(f"Missing properties in PLY file: {missing_props}")
|
|
|
|
return {
|
|
'positions': np.stack([
|
|
vertex['x'],
|
|
vertex['y'],
|
|
vertex['z']
|
|
], axis=1).astype(np.float32),
|
|
|
|
'colors': np.stack([
|
|
vertex['f_dc_0'],
|
|
vertex['f_dc_1'],
|
|
vertex['f_dc_2']
|
|
], axis=1).astype(np.float32),
|
|
|
|
'opacity': vertex['opacity'].astype(np.float32),
|
|
|
|
'scales': np.stack([
|
|
vertex['scale_0'],
|
|
vertex['scale_1'],
|
|
vertex['scale_2']
|
|
], axis=1).astype(np.float32),
|
|
|
|
'rotations': np.stack([
|
|
vertex['rot_0'],
|
|
vertex['rot_1'],
|
|
vertex['rot_2'],
|
|
vertex['rot_3']
|
|
], axis=1).astype(np.float32),
|
|
}
|
|
|
|
except ImportError:
|
|
raise ImportError("Please install the plyfile library: pip install plyfile")
|
|
except Exception as e:
|
|
raise ValueError(f"Error reading PLY file: {str(e)}")
|
|
|
|
def _send_to_td(self, gs_data, broadcast):
|
|
try:
|
|
import numpy as np
|
|
from server import PromptServer
|
|
|
|
num_points = len(gs_data['positions'])
|
|
|
|
header = (
|
|
"ply\n"
|
|
"format binary_little_endian 1.0\n"
|
|
f"element vertex {num_points}\n"
|
|
"property float x\n"
|
|
"property float y\n"
|
|
"property float z\n"
|
|
"property float f_dc_0\n"
|
|
"property float f_dc_1\n"
|
|
"property float f_dc_2\n"
|
|
"property float opacity\n"
|
|
"property float scale_0\n"
|
|
"property float scale_1\n"
|
|
"property float scale_2\n"
|
|
"property float rot_0\n"
|
|
"property float rot_1\n"
|
|
"property float rot_2\n"
|
|
"property float rot_3\n"
|
|
"end_header\n"
|
|
).encode('ascii')
|
|
|
|
# Vectorized packing:
|
|
# 14 float32 values per point:
|
|
# xyz + f_dc_rgb + opacity + scale_xyz + rotation_wxyz
|
|
interleaved = np.concatenate([
|
|
gs_data['positions'], # 3
|
|
gs_data['colors'], # 3
|
|
gs_data['opacity'].reshape(-1, 1), # 1
|
|
gs_data['scales'], # 3
|
|
gs_data['rotations'], # 4
|
|
], axis=1).astype('<f4')
|
|
|
|
binary_output = header + interleaved.tobytes()
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
|
|
server.send_sync(1003, binary_output, sid=sid)
|
|
|
|
return {
|
|
"ui": {
|
|
"gaussian_splatting": [{
|
|
"source": "websocket",
|
|
"content-type": "model/ply",
|
|
"type": "output",
|
|
"point_count": num_points,
|
|
}]
|
|
}
|
|
}
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Error sending Gaussian Splatting data: {str(e)}")
|
|
|
|
class StringtoTD:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"broadcast": ("BOOLEAN", {"default": False}),
|
|
"prefix": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"text": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_string"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_string(self, broadcast, prefix, text=None):
|
|
try:
|
|
if text is None:
|
|
text = ""
|
|
|
|
final_text = str(prefix) + str(text)
|
|
text_bytes = final_text.encode("utf-8")
|
|
|
|
server = PromptServer.instance
|
|
sid = None if broadcast else server.client_id
|
|
server.send_sync(1004, text_bytes, sid=sid)
|
|
|
|
return {
|
|
"ui": {
|
|
"text": [{
|
|
"source": "websocket",
|
|
"content-type": "text/plain; charset=utf-8",
|
|
"type": "output",
|
|
}]
|
|
}
|
|
}
|
|
except Exception as e:
|
|
raise ValueError(f"Error sending string: {str(e)}")
|
|
|
|
class ImagetoTD_Tagged:
|
|
EVENT_TYPE = 1005
|
|
MAGIC = b"ITD3"
|
|
VERSION = 1
|
|
|
|
FORMAT_MAP = {
|
|
"png": 1,
|
|
"jpeg": 2,
|
|
}
|
|
|
|
CONTENT_TYPE_MAP = {
|
|
"png": "image/png",
|
|
"jpeg": "image/jpeg",
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"format": (["png", "jpeg"], {"default": "png"}),
|
|
"layer_tag": ("STRING", {"default": "depth", "multiline": False}),
|
|
"quality": ("INT", {"default": 95, "min": 1, "max": 100}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "send_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "TouchDesigner"
|
|
|
|
def send_images(self, images, format, layer_tag, quality):
|
|
results = []
|
|
|
|
if layer_tag is None:
|
|
layer_tag = ""
|
|
layer_tag = str(layer_tag)
|
|
|
|
for tensor in images:
|
|
image = self._tensor_to_pil_image(tensor)
|
|
image_bytes, content_type = self._encode_image(image, format, quality)
|
|
packet = self._build_packet(format, layer_tag, image_bytes)
|
|
|
|
server = PromptServer.instance
|
|
|
|
server.send_sync(
|
|
self.EVENT_TYPE,
|
|
packet,
|
|
sid=None,
|
|
)
|
|
|
|
results.append({
|
|
"source": "websocket",
|
|
"content-type": content_type,
|
|
"type": "output",
|
|
"layer-tag": layer_tag,
|
|
})
|
|
|
|
return {"ui": {"images": results}}
|
|
|
|
def _tensor_to_pil_image(self, tensor):
|
|
array = 255.0 * tensor.cpu().numpy()
|
|
array = np.clip(array, 0, 255).astype(np.uint8)
|
|
|
|
if array.ndim == 2:
|
|
return Image.fromarray(array, mode="L")
|
|
|
|
if array.ndim == 3:
|
|
if array.shape[2] == 1:
|
|
return Image.fromarray(array[:, :, 0], mode="L")
|
|
elif array.shape[2] == 3:
|
|
return Image.fromarray(array, mode="RGB")
|
|
elif array.shape[2] == 4:
|
|
return Image.fromarray(array, mode="RGBA")
|
|
|
|
raise ValueError(f"Unsupported image tensor shape: {array.shape}")
|
|
|
|
def _encode_image(self, image, format, quality):
|
|
bytes_io = io.BytesIO()
|
|
|
|
if format == "png":
|
|
image.save(bytes_io, format="PNG")
|
|
return bytes_io.getvalue(), self.CONTENT_TYPE_MAP["png"]
|
|
|
|
elif format == "jpeg":
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
image.save(bytes_io, format="JPEG", quality=quality)
|
|
return bytes_io.getvalue(), self.CONTENT_TYPE_MAP["jpeg"]
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported format: {format}")
|
|
|
|
def _build_packet(self, format, layer_tag, image_bytes):
|
|
"""
|
|
4s magic = b"ITD3"
|
|
B version = 1
|
|
B format = 1(png) / 2(jpeg)
|
|
H tag_len
|
|
I image_len
|
|
... tag utf-8 bytes
|
|
... image bytes
|
|
"""
|
|
if format not in self.FORMAT_MAP:
|
|
raise ValueError(f"Unsupported format: {format}")
|
|
|
|
tag_bytes = layer_tag.encode("utf-8")
|
|
fmt_code = self.FORMAT_MAP[format]
|
|
|
|
if len(tag_bytes) > 65535:
|
|
raise ValueError("layer_tag is too long, max 65535 bytes")
|
|
|
|
header = struct.pack(
|
|
">4sBBHI",
|
|
self.MAGIC,
|
|
self.VERSION,
|
|
fmt_code,
|
|
len(tag_bytes),
|
|
len(image_bytes),
|
|
)
|
|
|
|
return header + tag_bytes + image_bytes |