feat: sealed worker data types — comfy_api_sealed_worker package

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,
+        )

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"]

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

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

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,
+        )