feat: 💄 GEOMETRY is now a dict

- For now only "mesh" and "material" keys are interopable from py <-> js
- Small fixes from merge (wip)

nodes/debug.py

@@ -1,10 +1,8 @@
-import base64
-import io
+import base64, io
 from pathlib import Path
 from typing import Optional
 
-import folder_paths
-import torch
+import folder_paths, torch
 
 from ..log import log
 from ..utils import tensor2pil
@@ -43,18 +41,29 @@ def process_list(anything):
         )
 
     elif isinstance(first_element, torch.Tensor):
-        text.append(f"List of Tensors: {first_element.shape} (x{len(anything)})")
+        text.append(
+            f"List of Tensors: {first_element.shape} (x{len(anything)})"
+        )
 
     return {"text": text}
 
 
 def process_dict(anything):
-    text = []
+    if "mesh" in anything:
+        m = {"geometry": {}}
+        m["geometry"]["mesh"] = mesh_to_json(anything["mesh"])
+        if "material" in anything:
+            m["geometry"]["material"] = anything["material"]
+        return m
+
+    res = []
     if "samples" in anything:
-        is_empty = "(empty)" if torch.count_nonzero(anything["samples"]) == 0 else ""
-        text.append(f"Latent Samples: {anything['samples'].shape} {is_empty}")
+        is_empty = (
+            "(empty)" if torch.count_nonzero(anything["samples"]) == 0 else ""
+        )
+        res.append(f"Latent Samples: {anything['samples'].shape} {is_empty}")
 
-    return {"text": text}
+    return {"text": res}
 
 
 def process_bool(anything):
@@ -65,6 +74,7 @@ def process_text(anything):
     return {"text": [str(anything)]}
 
 
+# NOT USED ANYMORE
 def process_geometry(anything):
     return {"geometry": [mesh_to_json(anything)]}
 
@@ -102,8 +112,6 @@ def do_debug(self, output_to_console, **kwargs):
             bool: process_bool,
             o3d.geometry.Geometry: process_geometry,
         }
-        if output_to_console:
-            print("bouh!")
 
         for anything in kwargs.values():
             processor = processors.get(type(anything))
@@ -118,11 +126,21 @@ def do_debug(self, output_to_console, **kwargs):
             processed_data = processor(anything)
 
             for ui_key, ui_value in processed_data.items():
-                output["ui"][ui_key].extend(ui_value)
+                if isinstance(ui_value, list):
+                    output["ui"][ui_key].extend(ui_value)
+                else:
+                    output["ui"][ui_key].append(ui_value)
             # log.debug(
             #     f"Processed input {k}, found {len(processed_data.get('b64_images', []))} images and {len(processed_data.get('text', []))} text items."
             # )
 
+        if output_to_console:
+            from rich.console import Console
+
+            cons = Console()
+            cons.print("OUTPUT:")
+            cons.print(output)
+
         return output
 
 

nodes/geo_tools.py

@@ -1,10 +1,15 @@
-import itertools
+import copy, itertools, json, os
+
 import numpy as np
 import open3d as o3d
-import json
+
 from ..utils import log
 
-import os
+
+def spread_geo(geo, *, cp=False):
+    mesh = geo["mesh"] if not cp else copy.copy(geo["mesh"])
+    material = geo.get("material", {})
+    return (mesh, material)
 
 
 def euler_to_rotation_matrix(x_deg, y_deg, z_deg):
@@ -14,13 +19,19 @@ def euler_to_rotation_matrix(x_deg, y_deg, z_deg):
     z = np.radians(z_deg)
 
     # Rotation matrix around x-axis
-    Rx = np.array([[1, 0, 0], [0, np.cos(x), -np.sin(x)], [0, np.sin(x), np.cos(x)]])
+    Rx = np.array(
+        [[1, 0, 0], [0, np.cos(x), -np.sin(x)], [0, np.sin(x), np.cos(x)]]
+    )
 
     # Rotation matrix around y-axis
-    Ry = np.array([[np.cos(y), 0, np.sin(y)], [0, 1, 0], [-np.sin(y), 0, np.cos(y)]])
+    Ry = np.array(
+        [[np.cos(y), 0, np.sin(y)], [0, 1, 0], [-np.sin(y), 0, np.cos(y)]]
+    )
 
     # Rotation matrix around z-axis
-    Rz = np.array([[np.cos(z), -np.sin(z), 0], [np.sin(z), np.cos(z), 0], [0, 0, 1]])
+    Rz = np.array(
+        [[np.cos(z), -np.sin(z), 0], [np.sin(z), np.cos(z), 0], [0, 0, 1]]
+    )
 
     return Rz @ Ry @ Rx
 
@@ -75,7 +86,9 @@ def create_grid(scale=(1, 1, 1), rows=10, columns=10):
     # Create vertices
     vertices = []
     for i in np.linspace(-dy / 2, dy / 2, rows + 1):
-        vertices.extend([j, 0, i] for j in np.linspace(-dx / 2, dx / 2, columns + 1))
+        vertices.extend(
+            [j, 0, i] for j in np.linspace(-dx / 2, dx / 2, columns + 1)
+        )
     # Generate triangles
     triangles = []
     for i, j in itertools.product(range(rows), range(columns)):
@@ -101,7 +114,9 @@ def create_box(scale=(1, 1, 1), divisions=(1, 1, 1)):
     vertices = []
     for i in np.linspace(-dx / 2, dx / 2, div_x + 1):
         for j in np.linspace(-dy / 2, dy / 2, div_y + 1):
-            vertices.extend([i, j, k] for k in np.linspace(-dz / 2, dz / 2, div_z + 1))
+            vertices.extend(
+                [i, j, k] for k in np.linspace(-dz / 2, dz / 2, div_z + 1)
+            )
     # Generate triangles for the box faces
     triangles = []
     for x, y in itertools.product(range(div_x), range(div_y)):
@@ -240,20 +255,125 @@ def create_torus(torus_radius=1, ring_radius=0.5, rows=10, columns=10):
 #     CATEGORY = "mtb/uv"
 
 
+def default_material(color=None):
+    return {
+        "color": color or "#00ff00",
+        "roughness": 1.0,
+        "metalness": 0.0,
+        "emissive": "#000000",
+        "displacementScale": 1.0,
+        "displacementMap": None,
+    }
+
+
+class MTB_Material:
+    """Make a std material."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        base = default_material()
+        return {
+            "required": {
+                "color": ("COLOR", {"default": base["color"]}),
+                "roughness": (
+                    "FLOAT",
+                    {
+                        "default": base["roughness"],
+                        "min": 0.005,
+                        "max": 4.0,
+                        "step": 0.01,
+                    },
+                ),
+                "metalness": (
+                    "FLOAT",
+                    {
+                        "default": base["metalness"],
+                        "min": 0.0,
+                        "max": 1.0,
+                        "step": 0.01,
+                    },
+                ),
+                "emissive": ("COLOR", {"default": base["emissive"]}),
+                "displacementScale": (
+                    "FLOAT",
+                    {"default": 1.0, "min": -10.0, "max": 10.0},
+                ),
+            },
+            "optional": {"displacementMap": ("IMAGE",)},
+        }
+
+    RETURN_TYPES = ("GEO_MATERIAL",)
+    RETURN_NAMES = ("material",)
+    FUNCTION = "make_material"
+    CATEGORY = "mtb/3D"
+
+    def make_material(
+        self, **kwargs
+    ):  # color, roughness, metalness, emissive, displacementScalen displacementMap=None):
+        return (kwargs,)
+
+
+class MTB_ApplyMaterial:
+    """Apply a Material to a geometry."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {"geometry": ("GEOMETRY",), "color": ("COLOR",)},
+            "optional": {"material": ("GEO_MATERIAL",)},
+        }
+
+    RETURN_TYPES = ("GEOMETRY",)
+    RETURN_NAMES = ("geometry",)
+    FUNCTION = "apply"
+    CATEGORY = "mtb/3D"
+
+    def apply(
+        self,
+        geometry,
+        color,
+        material=None,
+    ):
+        if material is None:
+            material = default_material(color)
+        #
+        geometry["material"] = material
+
+        return (geometry,)
+
+
 class TransformGeometry:
-    """Transforms the input geometry"""
+    """Transforms the input geometry."""
 
     @classmethod
     def INPUT_TYPES(cls):
         return {
             "required": {
                 "mesh": ("GEOMETRY",),
-                "position_x": ("FLOAT", {"default": 0.0, "step": 0.1}),
-                "position_y": ("FLOAT", {"default": 0.0, "step": 0.1}),
-                "position_z": ("FLOAT", {"default": 0.0, "step": 0.1}),
-                "rotation_x": ("FLOAT", {"default": 0.0, "step": 1}),
-                "rotation_y": ("FLOAT", {"default": 0.0, "step": 1}),
-                "rotation_z": ("FLOAT", {"default": 0.0, "step": 1}),
+                "position_x": (
+                    "FLOAT",
+                    {"default": 0.0, "step": 0.1, "min": -10000, "max": 10000},
+                ),
+                "position_y": (
+                    "FLOAT",
+                    {"default": 0.0, "step": 0.1, "min": -10000, "max": 10000},
+                ),
+                "position_z": (
+                    "FLOAT",
+                    {"default": 0.0, "step": 0.1, "min": -10000, "max": 10000},
+                ),
+                "rotation_x": (
+                    "FLOAT",
+                    {"default": 0.0, "step": 1, "min": -10000, "max": 10000},
+                ),
+                "rotation_y": (
+                    "FLOAT",
+                    {"default": 0.0, "step": 1, "min": -10000, "max": 10000},
+                ),
+                "rotation_z": (
+                    "FLOAT",
+                    {"default": 0.0, "step": 1, "min": -10000, "max": 10000},
+                ),
                 "scale_x": ("FLOAT", {"default": 1.0, "step": 0.1}),
                 "scale_y": ("FLOAT", {"default": 1.0, "step": 0.1}),
                 "scale_z": ("FLOAT", {"default": 1.0, "step": 0.1}),
@@ -267,7 +387,7 @@ def INPUT_TYPES(cls):
 
     def transform_geometry(
         self,
-        mesh,
+        mesh: o3d.geometry.TriangleMesh,
         position_x=0.0,
         position_y=0.0,
         position_z=0.0,
@@ -290,12 +410,21 @@ def transform_geometry(
         rotation = (rotation_x, rotation_y, rotation_z)
         scale = np.array([scale_x, scale_y, scale_z])
 
-        transformation_matrix = get_transformation_matrix(position, rotation, scale)
-        return (mesh.transform(transformation_matrix),)
+        transformation_matrix = get_transformation_matrix(
+            position, rotation, scale
+        )
+        mesh, material = spread_geo(mesh, cp=True)
+
+        return (
+            {
+                "mesh": mesh.transform(transformation_matrix),
+                "material": material,
+            },
+        )
 
 
 class GeometrySphere:
-    """Makes a Sphere 3D geometry"""
+    """Makes a Sphere 3D geometry.."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -320,11 +449,11 @@ def make_sphere(self, create_uv_map, radius, resolution):
         )
         mesh.compute_vertex_normals()
 
-        return (mesh,)
+        return ({"mesh": mesh},)
 
 
 class GeometryTest:
-    """Fetches an Open3D data geometry"""
+    """Fetches an Open3D data geometry.."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -358,11 +487,11 @@ def fetch_data(self, name):
         model = getattr(o3d.data, name)()
         mesh = o3d.io.read_triangle_mesh(model.path)
         mesh.compute_vertex_normals()
-        return (mesh,)
+        return ({"mesh": mesh},)
 
 
 class GeometryBox:
-    """Makes a Box 3D geometry"""
+    """Makes a Box 3D geometry."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -408,11 +537,11 @@ def make_box(
             (width, height, depth), (divisions_x, divisions_y, divisions_z)
         )
 
-        return (mesh,)
+        return ({"mesh": mesh},)
 
 
 class LoadGeometry:
-    """Load a 3D geometry"""
+    """Load a 3D geometry."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -428,35 +557,47 @@ def load_geo(self, path):
             raise ValueError(f"Path {path} does not exist")
 
         mesh = o3d.io.read_triangle_mesh(path)
+
+        if len(mesh.vertices) == 0:
+            mesh = o3d.io.read_triangle_model(path)
+            mesh_count = len(mesh.meshes)
+            if mesh_count == 0:
+                raise ValueError("Couldn't parse input file")
+
+            if mesh_count > 1:
+                log.warn(
+                    f"Found {mesh_count} meshes, only the first will be used..."
+                )
+
+            mesh = mesh.meshes[0].mesh
+
         mesh.compute_vertex_normals()
 
         return {
-            "result": (mesh,),
+            "result": ({"mesh": mesh},),
         }
 
 
 class GeometryInfo:
-    """Retrieve information about a 3D geometry"""
+    """Retrieve information about a 3D geometry."""
 
     @classmethod
     def INPUT_TYPES(cls):
         return {"required": {"geometry": ("GEOMETRY", {})}}
 
-    RETURN_TYPES = ("INT", "INT")
-    RETURN_NAMES = ("num_vertices", "num_triangles")
+    RETURN_TYPES = ("INT", "INT", "MATERIAL")
+    RETURN_NAMES = ("num_vertices", "num_triangles", "material")
     FUNCTION = "get_info"
     CATEGORY = "mtb/3D"
 
     def get_info(self, geometry):
-        log.debug(geometry)
-        return (
-            len(geometry.vertices),
-            len(geometry.triangles),
-        )
+        mesh, material = spread_geo(geometry)
+        log.debug(mesh)
+        return (len(mesh.vertices), len(mesh.triangles), material)
 
 
 class GeometryDecimater:
-    """Optimized the geometry to match the target number of triangles"""
+    """Optimized the geometry to match the target number of triangles."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -473,14 +614,16 @@ def INPUT_TYPES(cls):
     CATEGORY = "mtb/3D"
 
     def decimate(self, mesh, target):
-        mesh = mesh.simplify_quadric_decimation(target_number_of_triangles=target)
+        mesh = mesh.simplify_quadric_decimation(
+            target_number_of_triangles=target
+        )
         mesh.compute_vertex_normals()
 
-        return (mesh,)
+        return ({"mesh": mesh},)
 
 
 class GeometrySceneSetup:
-    """Scene setup for the renderer"""
+    """Scene setup for the renderer."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -496,11 +639,11 @@ def INPUT_TYPES(cls):
     CATEGORY = "mtb/3D"
 
     def setup(self, mesh, target):
-        return ({"geometry": mesh, "camera": cam},)
+        return ({"geometry": {"mesh": mesh}, "camera": cam},)
 
 
 class GeometryRender:
-    """Renders a Geometry to an image"""
+    """Renders a Geometry to an image."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -536,6 +679,9 @@ def render(self, geometry, width, height, background, camera):
     GeometryTest,
     GeometryDecimater,
     GeometrySphere,
+    GeometryRender,
     TransformGeometry,
     GeometryBox,
+    MTB_ApplyMaterial,
+    MTB_Material,
 ]

requirements.txt

@@ -6,4 +6,5 @@ rembg
 imageio_ffmpeg
 rich
 rich_argparse
-matplotlib
+matplotlib
+open3d==0.17.0 

web/debug.js

@@ -10,9 +10,7 @@
 import { app } from '../../scripts/app.js'
 
 import * as shared from './comfy_shared.js'
-import { log } from './comfy_shared.js'
 import { MtbWidgets } from './mtb_widgets.js'
-import { o3d_to_three } from './geometry_nodes.js'
 
 // TODO: respect inputs order...
 

web/geometry_nodes.js

@@ -27,7 +27,9 @@ export const make_wireframe = (mesh) => {
   return wireframe
 }
 
-export const o3d_to_three = (data) => {
+export const o3d_to_three = (data, material_opts) => {
+
+  material_opts = material_opts || { color: "0x00ff00" }
   // Parse the JSON data
   const meshData = JSON.parse(data)
 
@@ -58,13 +60,16 @@ export const o3d_to_three = (data) => {
     const uvs = new Float32Array(meshData.triangle_uvs.flat())
     geometry.setAttribute('uv', new THREE.BufferAttribute(uvs, 2))
   }
-  const material_opts = {
-    // wireframe: true,
+
+  if (material_opts.displacementB64 != undefined) {
+    material_opts.displacementMap = THREE.ImageUtils.loadTexture(material_opts.displacementB64)
+    delete (material_opts.displacementB64)
   }
+
   // For visualization, you might choose to use the MeshPhongMaterial to get the benefit of lighting with normals
   const material = meshData.vertex_colors
-    ? new THREE.MeshPhongMaterial({ ...material_opts, vertexColors: true })
-    : new THREE.MeshPhongMaterial({ ...material_opts, color: 0x00ff00 })
+    ? new THREE.MeshStandardMaterial({ ...material_opts, vertexColors: true })
+    : new THREE.MeshStandardMaterial({ ...material_opts })
 
   const threeMesh = new THREE.Mesh(geometry, material)
 

web/mtb_widgets.js

@@ -406,7 +406,7 @@ export const MtbWidgets = {
         const animate = () => {
           requestAnimationFrame(animate)
           if (this.mesh && this.animate) {
-            this.group.rotation.x += 0.005
+            //this.group.rotation.x += 0.005
             this.group.rotation.y += 0.005
           }
           this.renderer.render(this.scene, this.camera)
@@ -432,9 +432,16 @@ export const MtbWidgets = {
         }
         return [width, width]
       },
+      onRemoved: function () {
+        if (this.inputEl) {
+          this.inputEl.remove()
+        }
+      }
     }
     log('Creating canvas')
     w.inputEl = document.createElement('canvas')
+    w.inputEl.width = 768
+    w.inputEl.height = 768
 
     // add context menu with "animate" and "show wireframe"
     w.inputEl.addEventListener('contextmenu', (e) => {
@@ -483,9 +490,9 @@ export const MtbWidgets = {
       document.body.appendChild(w.menu)
     })
 
-    w.initThreeJS(w.inputEl, val)
+    w.initThreeJS(w.inputEl)
 
-    w.mesh = o3d_to_three(val)
+    w.mesh = o3d_to_three(val?.mesh ? val.mesh : val, val?.material)
     w.mesh_wireframe = make_wireframe(w.mesh)
     w.group = new THREE.Group()