From a6b5e6545d7ffd14a00b754dfd3faafcd955b181 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 29 Mar 2026 19:04:17 -0500 Subject: [PATCH] =?UTF-8?q?feat:=20sealed=20worker=20data=20types=20?= =?UTF-8?q?=E2=80=94=20comfy=5Fapi=5Fsealed=5Fworker=20package?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- comfy_api/latest/_util/trimesh_types.py | 259 +++++++++++++++++++++++ comfy_api_sealed_worker/__init__.py | 18 ++ comfy_api_sealed_worker/npz_types.py | 27 +++ comfy_api_sealed_worker/ply_types.py | 97 +++++++++ comfy_api_sealed_worker/trimesh_types.py | 259 +++++++++++++++++++++++ 5 files changed, 660 insertions(+) create mode 100644 comfy_api/latest/_util/trimesh_types.py create mode 100644 comfy_api_sealed_worker/__init__.py create mode 100644 comfy_api_sealed_worker/npz_types.py create mode 100644 comfy_api_sealed_worker/ply_types.py create mode 100644 comfy_api_sealed_worker/trimesh_types.py diff --git a/comfy_api/latest/_util/trimesh_types.py b/comfy_api/latest/_util/trimesh_types.py new file mode 100644 index 000000000..1c649c3e0 --- /dev/null +++ b/comfy_api/latest/_util/trimesh_types.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import numpy as np + + +class TrimeshData: + """Triangular mesh payload for cross-process transfer. + + Lightweight carrier for mesh geometry that does not depend on the + ``trimesh`` library. Serializers create this on the host side; + isolated child processes convert to/from ``trimesh.Trimesh`` as needed. + + Supports both ColorVisuals (vertex_colors) and TextureVisuals + (uv + material with textures). + """ + + def __init__( + self, + vertices: np.ndarray, + faces: np.ndarray, + vertex_normals: np.ndarray | None = None, + face_normals: np.ndarray | None = None, + vertex_colors: np.ndarray | None = None, + uv: np.ndarray | None = None, + material: dict | None = None, + vertex_attributes: dict | None = None, + face_attributes: dict | None = None, + metadata: dict | None = None, + ) -> None: + self.vertices = np.ascontiguousarray(vertices, dtype=np.float64) + self.faces = np.ascontiguousarray(faces, dtype=np.int64) + self.vertex_normals = ( + np.ascontiguousarray(vertex_normals, dtype=np.float64) + if vertex_normals is not None + else None + ) + self.face_normals = ( + np.ascontiguousarray(face_normals, dtype=np.float64) + if face_normals is not None + else None + ) + self.vertex_colors = ( + np.ascontiguousarray(vertex_colors, dtype=np.uint8) + if vertex_colors is not None + else None + ) + self.uv = ( + np.ascontiguousarray(uv, dtype=np.float64) + if uv is not None + else None + ) + self.material = material + self.vertex_attributes = vertex_attributes or {} + self.face_attributes = face_attributes or {} + self.metadata = self._detensorize_dict(metadata) if metadata else {} + + @staticmethod + def _detensorize_dict(d): + """Recursively convert any tensors in a dict back to numpy arrays.""" + if not isinstance(d, dict): + return d + result = {} + for k, v in d.items(): + if hasattr(v, "numpy"): + result[k] = v.cpu().numpy() if hasattr(v, "cpu") else v.numpy() + elif isinstance(v, dict): + result[k] = TrimeshData._detensorize_dict(v) + elif isinstance(v, list): + result[k] = [ + item.cpu().numpy() if hasattr(item, "numpy") and hasattr(item, "cpu") + else item.numpy() if hasattr(item, "numpy") + else item + for item in v + ] + else: + result[k] = v + return result + + @staticmethod + def _to_numpy(arr, dtype): + if arr is None: + return None + if hasattr(arr, "numpy"): + arr = arr.cpu().numpy() if hasattr(arr, "cpu") else arr.numpy() + return np.ascontiguousarray(arr, dtype=dtype) + + @property + def num_vertices(self) -> int: + return self.vertices.shape[0] + + @property + def num_faces(self) -> int: + return self.faces.shape[0] + + @property + def has_texture(self) -> bool: + return self.uv is not None and self.material is not None + + def to_trimesh(self): + """Convert to trimesh.Trimesh (requires trimesh in the environment).""" + import trimesh + from trimesh.visual import TextureVisuals, ColorVisuals + + kwargs = {} + if self.vertex_normals is not None: + kwargs["vertex_normals"] = self.vertex_normals + if self.face_normals is not None: + kwargs["face_normals"] = self.face_normals + if self.metadata: + kwargs["metadata"] = self.metadata + + mesh = trimesh.Trimesh( + vertices=self.vertices, faces=self.faces, process=False, **kwargs + ) + + # Reconstruct visual + if self.has_texture: + material = self._dict_to_material(self.material) + mesh.visual = TextureVisuals(uv=self.uv, material=material) + elif self.vertex_colors is not None: + mesh.visual.vertex_colors = self.vertex_colors + + for k, v in self.vertex_attributes.items(): + mesh.vertex_attributes[k] = v + + for k, v in self.face_attributes.items(): + mesh.face_attributes[k] = v + + return mesh + + @staticmethod + def _material_to_dict(material) -> dict: + """Serialize a trimesh material to a plain dict.""" + import base64 + from io import BytesIO + from trimesh.visual.material import PBRMaterial, SimpleMaterial + + result = {"type": type(material).__name__, "name": getattr(material, "name", None)} + + if isinstance(material, PBRMaterial): + result["baseColorFactor"] = material.baseColorFactor + result["metallicFactor"] = material.metallicFactor + result["roughnessFactor"] = material.roughnessFactor + result["emissiveFactor"] = material.emissiveFactor + result["alphaMode"] = material.alphaMode + result["alphaCutoff"] = material.alphaCutoff + result["doubleSided"] = material.doubleSided + + for tex_name in ("baseColorTexture", "normalTexture", "emissiveTexture", + "metallicRoughnessTexture", "occlusionTexture"): + tex = getattr(material, tex_name, None) + if tex is not None: + buf = BytesIO() + tex.save(buf, format="PNG") + result[tex_name] = base64.b64encode(buf.getvalue()).decode("ascii") + + elif isinstance(material, SimpleMaterial): + result["main_color"] = list(material.main_color) if material.main_color is not None else None + result["glossiness"] = material.glossiness + if hasattr(material, "image") and material.image is not None: + buf = BytesIO() + material.image.save(buf, format="PNG") + result["image"] = base64.b64encode(buf.getvalue()).decode("ascii") + + return result + + @staticmethod + def _dict_to_material(d: dict): + """Reconstruct a trimesh material from a plain dict.""" + import base64 + from io import BytesIO + from PIL import Image + from trimesh.visual.material import PBRMaterial, SimpleMaterial + + mat_type = d.get("type", "PBRMaterial") + + if mat_type == "PBRMaterial": + kwargs = { + "name": d.get("name"), + "baseColorFactor": d.get("baseColorFactor"), + "metallicFactor": d.get("metallicFactor"), + "roughnessFactor": d.get("roughnessFactor"), + "emissiveFactor": d.get("emissiveFactor"), + "alphaMode": d.get("alphaMode"), + "alphaCutoff": d.get("alphaCutoff"), + "doubleSided": d.get("doubleSided"), + } + for tex_name in ("baseColorTexture", "normalTexture", "emissiveTexture", + "metallicRoughnessTexture", "occlusionTexture"): + if tex_name in d and d[tex_name] is not None: + img = Image.open(BytesIO(base64.b64decode(d[tex_name]))) + kwargs[tex_name] = img + return PBRMaterial(**{k: v for k, v in kwargs.items() if v is not None}) + + elif mat_type == "SimpleMaterial": + kwargs = { + "name": d.get("name"), + "glossiness": d.get("glossiness"), + } + if d.get("main_color") is not None: + kwargs["diffuse"] = d["main_color"] + if d.get("image") is not None: + kwargs["image"] = Image.open(BytesIO(base64.b64decode(d["image"]))) + return SimpleMaterial(**kwargs) + + raise ValueError(f"Unknown material type: {mat_type}") + + @classmethod + def from_trimesh(cls, mesh) -> TrimeshData: + """Create from a trimesh.Trimesh object.""" + from trimesh.visual.texture import TextureVisuals + + vertex_normals = None + if mesh._cache.cache.get("vertex_normals") is not None: + vertex_normals = np.asarray(mesh.vertex_normals) + + face_normals = None + if mesh._cache.cache.get("face_normals") is not None: + face_normals = np.asarray(mesh.face_normals) + + vertex_colors = None + uv = None + material = None + + if isinstance(mesh.visual, TextureVisuals): + if mesh.visual.uv is not None: + uv = np.asarray(mesh.visual.uv, dtype=np.float64) + if mesh.visual.material is not None: + material = cls._material_to_dict(mesh.visual.material) + else: + try: + vc = mesh.visual.vertex_colors + if vc is not None and len(vc) > 0: + vertex_colors = np.asarray(vc, dtype=np.uint8) + except Exception: + pass + + va = {} + if hasattr(mesh, "vertex_attributes") and mesh.vertex_attributes: + for k, v in mesh.vertex_attributes.items(): + va[k] = np.asarray(v) if hasattr(v, "__array__") else v + + fa = {} + if hasattr(mesh, "face_attributes") and mesh.face_attributes: + for k, v in mesh.face_attributes.items(): + fa[k] = np.asarray(v) if hasattr(v, "__array__") else v + + return cls( + vertices=np.asarray(mesh.vertices), + faces=np.asarray(mesh.faces), + vertex_normals=vertex_normals, + face_normals=face_normals, + vertex_colors=vertex_colors, + uv=uv, + material=material, + vertex_attributes=va if va else None, + face_attributes=fa if fa else None, + metadata=mesh.metadata if mesh.metadata else None, + ) diff --git a/comfy_api_sealed_worker/__init__.py b/comfy_api_sealed_worker/__init__.py new file mode 100644 index 000000000..269aa2644 --- /dev/null +++ b/comfy_api_sealed_worker/__init__.py @@ -0,0 +1,18 @@ +"""comfy_api_sealed_worker — torch-free type definitions for sealed worker children. + +Drop-in replacement for comfy_api.latest._util type imports in sealed workers +that do not have torch installed. Contains only data type definitions (TrimeshData, +PLY, NPZ, etc.) with numpy-only dependencies. + +Usage in serializers: + if _IMPORT_TORCH: + from comfy_api.latest._util.trimesh_types import TrimeshData + else: + from comfy_api_sealed_worker.trimesh_types import TrimeshData +""" + +from .trimesh_types import TrimeshData +from .ply_types import PLY +from .npz_types import NPZ + +__all__ = ["TrimeshData", "PLY", "NPZ"] diff --git a/comfy_api_sealed_worker/npz_types.py b/comfy_api_sealed_worker/npz_types.py new file mode 100644 index 000000000..a93eed68c --- /dev/null +++ b/comfy_api_sealed_worker/npz_types.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import os + + +class NPZ: + """Ordered collection of NPZ file payloads. + + Each entry in ``frames`` is a complete compressed ``.npz`` file stored + as raw bytes (produced by ``numpy.savez_compressed`` into a BytesIO). + ``save_to`` writes numbered files into a directory. + """ + + def __init__(self, frames: list[bytes]) -> None: + self.frames = frames + + @property + def num_frames(self) -> int: + return len(self.frames) + + def save_to(self, directory: str, prefix: str = "frame") -> str: + os.makedirs(directory, exist_ok=True) + for i, frame_bytes in enumerate(self.frames): + path = os.path.join(directory, f"{prefix}_{i:06d}.npz") + with open(path, "wb") as f: + f.write(frame_bytes) + return directory diff --git a/comfy_api_sealed_worker/ply_types.py b/comfy_api_sealed_worker/ply_types.py new file mode 100644 index 000000000..8beb566bc --- /dev/null +++ b/comfy_api_sealed_worker/ply_types.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import numpy as np + + +class PLY: + """Point cloud payload for PLY file output. + + Supports two schemas: + - Pointcloud: xyz positions with optional colors, confidence, view_id (ASCII format) + - Gaussian: raw binary PLY data built by producer nodes using plyfile (binary format) + + When ``raw_data`` is provided, the object acts as an opaque binary PLY + carrier and ``save_to`` writes the bytes directly. + """ + + def __init__( + self, + points: np.ndarray | None = None, + colors: np.ndarray | None = None, + confidence: np.ndarray | None = None, + view_id: np.ndarray | None = None, + raw_data: bytes | None = None, + ) -> None: + self.raw_data = raw_data + if raw_data is not None: + self.points = None + self.colors = None + self.confidence = None + self.view_id = None + return + if points is None: + raise ValueError("Either points or raw_data must be provided") + if points.ndim != 2 or points.shape[1] != 3: + raise ValueError(f"points must be (N, 3), got {points.shape}") + self.points = np.ascontiguousarray(points, dtype=np.float32) + self.colors = np.ascontiguousarray(colors, dtype=np.float32) if colors is not None else None + self.confidence = np.ascontiguousarray(confidence, dtype=np.float32) if confidence is not None else None + self.view_id = np.ascontiguousarray(view_id, dtype=np.int32) if view_id is not None else None + + @property + def is_gaussian(self) -> bool: + return self.raw_data is not None + + @property + def num_points(self) -> int: + if self.points is not None: + return self.points.shape[0] + return 0 + + @staticmethod + def _to_numpy(arr, dtype): + if arr is None: + return None + if hasattr(arr, "numpy"): + arr = arr.cpu().numpy() if hasattr(arr, "cpu") else arr.numpy() + return np.ascontiguousarray(arr, dtype=dtype) + + def save_to(self, path: str) -> str: + if self.raw_data is not None: + with open(path, "wb") as f: + f.write(self.raw_data) + return path + self.points = self._to_numpy(self.points, np.float32) + self.colors = self._to_numpy(self.colors, np.float32) + self.confidence = self._to_numpy(self.confidence, np.float32) + self.view_id = self._to_numpy(self.view_id, np.int32) + N = self.num_points + header_lines = [ + "ply", + "format ascii 1.0", + f"element vertex {N}", + "property float x", + "property float y", + "property float z", + ] + if self.colors is not None: + header_lines += ["property uchar red", "property uchar green", "property uchar blue"] + if self.confidence is not None: + header_lines.append("property float confidence") + if self.view_id is not None: + header_lines.append("property int view_id") + header_lines.append("end_header") + + with open(path, "w") as f: + f.write("\n".join(header_lines) + "\n") + for i in range(N): + parts = [f"{self.points[i, 0]} {self.points[i, 1]} {self.points[i, 2]}"] + if self.colors is not None: + r, g, b = (self.colors[i] * 255).clip(0, 255).astype(np.uint8) + parts.append(f"{r} {g} {b}") + if self.confidence is not None: + parts.append(f"{self.confidence[i]}") + if self.view_id is not None: + parts.append(f"{int(self.view_id[i])}") + f.write(" ".join(parts) + "\n") + return path diff --git a/comfy_api_sealed_worker/trimesh_types.py b/comfy_api_sealed_worker/trimesh_types.py new file mode 100644 index 000000000..1c649c3e0 --- /dev/null +++ b/comfy_api_sealed_worker/trimesh_types.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import numpy as np + + +class TrimeshData: + """Triangular mesh payload for cross-process transfer. + + Lightweight carrier for mesh geometry that does not depend on the + ``trimesh`` library. Serializers create this on the host side; + isolated child processes convert to/from ``trimesh.Trimesh`` as needed. + + Supports both ColorVisuals (vertex_colors) and TextureVisuals + (uv + material with textures). + """ + + def __init__( + self, + vertices: np.ndarray, + faces: np.ndarray, + vertex_normals: np.ndarray | None = None, + face_normals: np.ndarray | None = None, + vertex_colors: np.ndarray | None = None, + uv: np.ndarray | None = None, + material: dict | None = None, + vertex_attributes: dict | None = None, + face_attributes: dict | None = None, + metadata: dict | None = None, + ) -> None: + self.vertices = np.ascontiguousarray(vertices, dtype=np.float64) + self.faces = np.ascontiguousarray(faces, dtype=np.int64) + self.vertex_normals = ( + np.ascontiguousarray(vertex_normals, dtype=np.float64) + if vertex_normals is not None + else None + ) + self.face_normals = ( + np.ascontiguousarray(face_normals, dtype=np.float64) + if face_normals is not None + else None + ) + self.vertex_colors = ( + np.ascontiguousarray(vertex_colors, dtype=np.uint8) + if vertex_colors is not None + else None + ) + self.uv = ( + np.ascontiguousarray(uv, dtype=np.float64) + if uv is not None + else None + ) + self.material = material + self.vertex_attributes = vertex_attributes or {} + self.face_attributes = face_attributes or {} + self.metadata = self._detensorize_dict(metadata) if metadata else {} + + @staticmethod + def _detensorize_dict(d): + """Recursively convert any tensors in a dict back to numpy arrays.""" + if not isinstance(d, dict): + return d + result = {} + for k, v in d.items(): + if hasattr(v, "numpy"): + result[k] = v.cpu().numpy() if hasattr(v, "cpu") else v.numpy() + elif isinstance(v, dict): + result[k] = TrimeshData._detensorize_dict(v) + elif isinstance(v, list): + result[k] = [ + item.cpu().numpy() if hasattr(item, "numpy") and hasattr(item, "cpu") + else item.numpy() if hasattr(item, "numpy") + else item + for item in v + ] + else: + result[k] = v + return result + + @staticmethod + def _to_numpy(arr, dtype): + if arr is None: + return None + if hasattr(arr, "numpy"): + arr = arr.cpu().numpy() if hasattr(arr, "cpu") else arr.numpy() + return np.ascontiguousarray(arr, dtype=dtype) + + @property + def num_vertices(self) -> int: + return self.vertices.shape[0] + + @property + def num_faces(self) -> int: + return self.faces.shape[0] + + @property + def has_texture(self) -> bool: + return self.uv is not None and self.material is not None + + def to_trimesh(self): + """Convert to trimesh.Trimesh (requires trimesh in the environment).""" + import trimesh + from trimesh.visual import TextureVisuals, ColorVisuals + + kwargs = {} + if self.vertex_normals is not None: + kwargs["vertex_normals"] = self.vertex_normals + if self.face_normals is not None: + kwargs["face_normals"] = self.face_normals + if self.metadata: + kwargs["metadata"] = self.metadata + + mesh = trimesh.Trimesh( + vertices=self.vertices, faces=self.faces, process=False, **kwargs + ) + + # Reconstruct visual + if self.has_texture: + material = self._dict_to_material(self.material) + mesh.visual = TextureVisuals(uv=self.uv, material=material) + elif self.vertex_colors is not None: + mesh.visual.vertex_colors = self.vertex_colors + + for k, v in self.vertex_attributes.items(): + mesh.vertex_attributes[k] = v + + for k, v in self.face_attributes.items(): + mesh.face_attributes[k] = v + + return mesh + + @staticmethod + def _material_to_dict(material) -> dict: + """Serialize a trimesh material to a plain dict.""" + import base64 + from io import BytesIO + from trimesh.visual.material import PBRMaterial, SimpleMaterial + + result = {"type": type(material).__name__, "name": getattr(material, "name", None)} + + if isinstance(material, PBRMaterial): + result["baseColorFactor"] = material.baseColorFactor + result["metallicFactor"] = material.metallicFactor + result["roughnessFactor"] = material.roughnessFactor + result["emissiveFactor"] = material.emissiveFactor + result["alphaMode"] = material.alphaMode + result["alphaCutoff"] = material.alphaCutoff + result["doubleSided"] = material.doubleSided + + for tex_name in ("baseColorTexture", "normalTexture", "emissiveTexture", + "metallicRoughnessTexture", "occlusionTexture"): + tex = getattr(material, tex_name, None) + if tex is not None: + buf = BytesIO() + tex.save(buf, format="PNG") + result[tex_name] = base64.b64encode(buf.getvalue()).decode("ascii") + + elif isinstance(material, SimpleMaterial): + result["main_color"] = list(material.main_color) if material.main_color is not None else None + result["glossiness"] = material.glossiness + if hasattr(material, "image") and material.image is not None: + buf = BytesIO() + material.image.save(buf, format="PNG") + result["image"] = base64.b64encode(buf.getvalue()).decode("ascii") + + return result + + @staticmethod + def _dict_to_material(d: dict): + """Reconstruct a trimesh material from a plain dict.""" + import base64 + from io import BytesIO + from PIL import Image + from trimesh.visual.material import PBRMaterial, SimpleMaterial + + mat_type = d.get("type", "PBRMaterial") + + if mat_type == "PBRMaterial": + kwargs = { + "name": d.get("name"), + "baseColorFactor": d.get("baseColorFactor"), + "metallicFactor": d.get("metallicFactor"), + "roughnessFactor": d.get("roughnessFactor"), + "emissiveFactor": d.get("emissiveFactor"), + "alphaMode": d.get("alphaMode"), + "alphaCutoff": d.get("alphaCutoff"), + "doubleSided": d.get("doubleSided"), + } + for tex_name in ("baseColorTexture", "normalTexture", "emissiveTexture", + "metallicRoughnessTexture", "occlusionTexture"): + if tex_name in d and d[tex_name] is not None: + img = Image.open(BytesIO(base64.b64decode(d[tex_name]))) + kwargs[tex_name] = img + return PBRMaterial(**{k: v for k, v in kwargs.items() if v is not None}) + + elif mat_type == "SimpleMaterial": + kwargs = { + "name": d.get("name"), + "glossiness": d.get("glossiness"), + } + if d.get("main_color") is not None: + kwargs["diffuse"] = d["main_color"] + if d.get("image") is not None: + kwargs["image"] = Image.open(BytesIO(base64.b64decode(d["image"]))) + return SimpleMaterial(**kwargs) + + raise ValueError(f"Unknown material type: {mat_type}") + + @classmethod + def from_trimesh(cls, mesh) -> TrimeshData: + """Create from a trimesh.Trimesh object.""" + from trimesh.visual.texture import TextureVisuals + + vertex_normals = None + if mesh._cache.cache.get("vertex_normals") is not None: + vertex_normals = np.asarray(mesh.vertex_normals) + + face_normals = None + if mesh._cache.cache.get("face_normals") is not None: + face_normals = np.asarray(mesh.face_normals) + + vertex_colors = None + uv = None + material = None + + if isinstance(mesh.visual, TextureVisuals): + if mesh.visual.uv is not None: + uv = np.asarray(mesh.visual.uv, dtype=np.float64) + if mesh.visual.material is not None: + material = cls._material_to_dict(mesh.visual.material) + else: + try: + vc = mesh.visual.vertex_colors + if vc is not None and len(vc) > 0: + vertex_colors = np.asarray(vc, dtype=np.uint8) + except Exception: + pass + + va = {} + if hasattr(mesh, "vertex_attributes") and mesh.vertex_attributes: + for k, v in mesh.vertex_attributes.items(): + va[k] = np.asarray(v) if hasattr(v, "__array__") else v + + fa = {} + if hasattr(mesh, "face_attributes") and mesh.face_attributes: + for k, v in mesh.face_attributes.items(): + fa[k] = np.asarray(v) if hasattr(v, "__array__") else v + + return cls( + vertices=np.asarray(mesh.vertices), + faces=np.asarray(mesh.faces), + vertex_normals=vertex_normals, + face_normals=face_normals, + vertex_colors=vertex_colors, + uv=uv, + material=material, + vertex_attributes=va if va else None, + face_attributes=fa if fa else None, + metadata=mesh.metadata if mesh.metadata else None, + )