convert hunyuan3d.py to V3 schema (#10664)

comfy_api/latest/__init__.py

@@ -7,7 +7,7 @@ from comfy_api.internal.singleton import ProxiedSingleton
 from comfy_api.internal.async_to_sync import create_sync_class
 from comfy_api.latest._input import ImageInput, AudioInput, MaskInput, LatentInput, VideoInput
 from comfy_api.latest._input_impl import VideoFromFile, VideoFromComponents
-from comfy_api.latest._util import VideoCodec, VideoContainer, VideoComponents
+from comfy_api.latest._util import VideoCodec, VideoContainer, VideoComponents, MESH, VOXEL
 from . import _io as io
 from . import _ui as ui
 # from comfy_api.latest._resources import _RESOURCES as resources  #noqa: F401
@@ -104,6 +104,8 @@ class Types:
     VideoCodec = VideoCodec
     VideoContainer = VideoContainer
     VideoComponents = VideoComponents
+    MESH = MESH
+    VOXEL = VOXEL
 
 ComfyAPI = ComfyAPI_latest
 

comfy_api/latest/_io.py

@@ -27,6 +27,7 @@ from comfy_api.internal import (_ComfyNodeInternal, _NodeOutputInternal, classpr
     prune_dict, shallow_clone_class)
 from comfy_api.latest._resources import Resources, ResourcesLocal
 from comfy_execution.graph_utils import ExecutionBlocker
+from ._util import MESH, VOXEL
 
 # from comfy_extras.nodes_images import SVG as SVG_ # NOTE: needs to be moved before can be imported due to circular reference
 
@@ -656,11 +657,11 @@ class LossMap(ComfyTypeIO):
 
 @comfytype(io_type="VOXEL")
 class Voxel(ComfyTypeIO):
-    Type = Any # TODO: VOXEL class is defined in comfy_extras/nodes_hunyuan3d.py; should be moved to somewhere else before referenced directly in v3
+    Type = VOXEL
 
 @comfytype(io_type="MESH")
 class Mesh(ComfyTypeIO):
-    Type = Any # TODO: MESH class is defined in comfy_extras/nodes_hunyuan3d.py; should be moved to somewhere else before referenced directly in v3
+    Type = MESH
 
 @comfytype(io_type="HOOKS")
 class Hooks(ComfyTypeIO):

comfy_api/latest/_util/__init__.py

@@ -1,8 +1,11 @@
 from .video_types import VideoContainer, VideoCodec, VideoComponents
+from .geometry_types import VOXEL, MESH
 
 __all__ = [
     # Utility Types
     "VideoContainer",
     "VideoCodec",
     "VideoComponents",
+    "VOXEL",
+    "MESH",
 ]

comfy_api/latest/_util/geometry_types.py

@@ -0,0 +1,12 @@
+import torch
+
+
+class VOXEL:
+    def __init__(self, data: torch.Tensor):
+        self.data = data
+
+
+class MESH:
+    def __init__(self, vertices: torch.Tensor, faces: torch.Tensor):
+        self.vertices = vertices
+        self.faces = faces

comfy_extras/nodes_hunyuan3d.py

@@ -7,63 +7,79 @@ from comfy.ldm.modules.diffusionmodules.mmdit import get_1d_sincos_pos_embed_fro
 import folder_paths
 import comfy.model_management
 from comfy.cli_args import args
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, IO, Types
+from comfy_api.latest._util import MESH, VOXEL  # only for backward compatibility if someone import it from this file (will be removed later) # noqa
 
-class EmptyLatentHunyuan3Dv2:
+
+class EmptyLatentHunyuan3Dv2(IO.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {
+        return IO.Schema(
-            "required": {
+            node_id="EmptyLatentHunyuan3Dv2",
-                "resolution": ("INT", {"default": 3072, "min": 1, "max": 8192}),
+            category="latent/3d",
-                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."}),
+            inputs=[
-            }
+                IO.Int.Input("resolution", default=3072, min=1, max=8192),
-        }
+                IO.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch."),
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ]
+        )
 
-    RETURN_TYPES = ("LATENT",)
+    @classmethod
-    FUNCTION = "generate"
+    def execute(cls, resolution, batch_size) -> IO.NodeOutput:
-
-    CATEGORY = "latent/3d"
-
-    def generate(self, resolution, batch_size):
         latent = torch.zeros([batch_size, 64, resolution], device=comfy.model_management.intermediate_device())
-        return ({"samples": latent, "type": "hunyuan3dv2"}, )
+        return IO.NodeOutput({"samples": latent, "type": "hunyuan3dv2"})
 
-class Hunyuan3Dv2Conditioning:
+    generate = execute  # TODO: remove
+
+
+class Hunyuan3Dv2Conditioning(IO.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {"required": {"clip_vision_output": ("CLIP_VISION_OUTPUT",),
+        return IO.Schema(
-                             }}
+            node_id="Hunyuan3Dv2Conditioning",
+            category="conditioning/video_models",
+            inputs=[
+                IO.ClipVisionOutput.Input("clip_vision_output"),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ]
+        )
 
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
+    @classmethod
-    RETURN_NAMES = ("positive", "negative")
+    def execute(cls, clip_vision_output) -> IO.NodeOutput:
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/video_models"
-
-    def encode(self, clip_vision_output):
         embeds = clip_vision_output.last_hidden_state
         positive = [[embeds, {}]]
         negative = [[torch.zeros_like(embeds), {}]]
-        return (positive, negative)
+        return IO.NodeOutput(positive, negative)
+
+    encode = execute  # TODO: remove
 
 
-class Hunyuan3Dv2ConditioningMultiView:
+class Hunyuan3Dv2ConditioningMultiView(IO.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {"required": {},
+        return IO.Schema(
-                "optional": {"front": ("CLIP_VISION_OUTPUT",),
+            node_id="Hunyuan3Dv2ConditioningMultiView",
-                             "left": ("CLIP_VISION_OUTPUT",),
+            category="conditioning/video_models",
-                             "back": ("CLIP_VISION_OUTPUT",),
+            inputs=[
-                             "right": ("CLIP_VISION_OUTPUT",), }}
+                IO.ClipVisionOutput.Input("front", optional=True),
+                IO.ClipVisionOutput.Input("left", optional=True),
+                IO.ClipVisionOutput.Input("back", optional=True),
+                IO.ClipVisionOutput.Input("right", optional=True),
+            ],
+            outputs=[
+                IO.Conditioning.Output(display_name="positive"),
+                IO.Conditioning.Output(display_name="negative"),
+            ]
+        )
 
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
+    @classmethod
-    RETURN_NAMES = ("positive", "negative")
+    def execute(cls, front=None, left=None, back=None, right=None) -> IO.NodeOutput:
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/video_models"
-
-    def encode(self, front=None, left=None, back=None, right=None):
         all_embeds = [front, left, back, right]
         out = []
         pos_embeds = None
@@ -76,29 +92,35 @@ class Hunyuan3Dv2ConditioningMultiView:
         embeds = torch.cat(out, dim=1)
         positive = [[embeds, {}]]
         negative = [[torch.zeros_like(embeds), {}]]
-        return (positive, negative)
+        return IO.NodeOutput(positive, negative)
+
+    encode = execute  # TODO: remove
 
 
-class VOXEL:
+class VAEDecodeHunyuan3D(IO.ComfyNode):
-    def __init__(self, data):
-        self.data = data
-
-class VAEDecodeHunyuan3D:
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {"required": {"samples": ("LATENT", ),
+        return IO.Schema(
-                             "vae": ("VAE", ),
+            node_id="VAEDecodeHunyuan3D",
-                             "num_chunks": ("INT", {"default": 8000, "min": 1000, "max": 500000}),
+            category="latent/3d",
-                             "octree_resolution": ("INT", {"default": 256, "min": 16, "max": 512}),
+            inputs=[
-                             }}
+                IO.Latent.Input("samples"),
-    RETURN_TYPES = ("VOXEL",)
+                IO.Vae.Input("vae"),
-    FUNCTION = "decode"
+                IO.Int.Input("num_chunks", default=8000, min=1000, max=500000),
+                IO.Int.Input("octree_resolution", default=256, min=16, max=512),
+            ],
+            outputs=[
+                IO.Voxel.Output(),
+            ]
+        )
 
-    CATEGORY = "latent/3d"
+    @classmethod
+    def execute(cls, vae, samples, num_chunks, octree_resolution) -> IO.NodeOutput:
+        voxels = Types.VOXEL(vae.decode(samples["samples"], vae_options={"num_chunks": num_chunks, "octree_resolution": octree_resolution}))
+        return IO.NodeOutput(voxels)
+
+    decode = execute  # TODO: remove
 
-    def decode(self, vae, samples, num_chunks, octree_resolution):
-        voxels = VOXEL(vae.decode(samples["samples"], vae_options={"num_chunks": num_chunks, "octree_resolution": octree_resolution}))
-        return (voxels, )
 
 def voxel_to_mesh(voxels, threshold=0.5, device=None):
     if device is None:
@@ -396,24 +418,24 @@ def voxel_to_mesh_surfnet(voxels, threshold=0.5, device=None):
 
     return final_vertices, faces
 
-class MESH:
-    def __init__(self, vertices, faces):
-        self.vertices = vertices
-        self.faces = faces
 
-
+class VoxelToMeshBasic(IO.ComfyNode):
-class VoxelToMeshBasic:
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {"required": {"voxel": ("VOXEL", ),
+        return IO.Schema(
-                             "threshold": ("FLOAT", {"default": 0.6, "min": -1.0, "max": 1.0, "step": 0.01}),
+            node_id="VoxelToMeshBasic",
-                             }}
+            category="3d",
-    RETURN_TYPES = ("MESH",)
+            inputs=[
-    FUNCTION = "decode"
+                IO.Voxel.Input("voxel"),
+                IO.Float.Input("threshold", default=0.6, min=-1.0, max=1.0, step=0.01),
+            ],
+            outputs=[
+                IO.Mesh.Output(),
+            ]
+        )
 
-    CATEGORY = "3d"
+    @classmethod
-
+    def execute(cls, voxel, threshold) -> IO.NodeOutput:
-    def decode(self, voxel, threshold):
         vertices = []
         faces = []
         for x in voxel.data:
@@ -421,21 +443,29 @@ class VoxelToMeshBasic:
             vertices.append(v)
             faces.append(f)
 
-        return (MESH(torch.stack(vertices), torch.stack(faces)), )
+        return IO.NodeOutput(Types.MESH(torch.stack(vertices), torch.stack(faces)))
 
-class VoxelToMesh:
+    decode = execute  # TODO: remove
+
+
+class VoxelToMesh(IO.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {"required": {"voxel": ("VOXEL", ),
+        return IO.Schema(
-                             "algorithm": (["surface net", "basic"], ),
+            node_id="VoxelToMesh",
-                             "threshold": ("FLOAT", {"default": 0.6, "min": -1.0, "max": 1.0, "step": 0.01}),
+            category="3d",
-                             }}
+            inputs=[
-    RETURN_TYPES = ("MESH",)
+                IO.Voxel.Input("voxel"),
-    FUNCTION = "decode"
+                IO.Combo.Input("algorithm", options=["surface net", "basic"]),
+                IO.Float.Input("threshold", default=0.6, min=-1.0, max=1.0, step=0.01),
+            ],
+            outputs=[
+                IO.Mesh.Output(),
+            ]
+        )
 
-    CATEGORY = "3d"
+    @classmethod
-
+    def execute(cls, voxel, algorithm, threshold) -> IO.NodeOutput:
-    def decode(self, voxel, algorithm, threshold):
         vertices = []
         faces = []
 
@@ -449,7 +479,9 @@ class VoxelToMesh:
             vertices.append(v)
             faces.append(f)
 
-        return (MESH(torch.stack(vertices), torch.stack(faces)), )
+        return IO.NodeOutput(Types.MESH(torch.stack(vertices), torch.stack(faces)))
+
+    decode = execute  # TODO: remove
 
 
 def save_glb(vertices, faces, filepath, metadata=None):
@@ -581,31 +613,32 @@ def save_glb(vertices, faces, filepath, metadata=None):
     return filepath
 
 
-class SaveGLB:
+class SaveGLB(IO.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {"required": {"mesh": ("MESH", ),
+        return IO.Schema(
-                             "filename_prefix": ("STRING", {"default": "mesh/ComfyUI"}), },
+            node_id="SaveGLB",
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, }
+            category="3d",
+            is_output_node=True,
+            inputs=[
+                IO.Mesh.Input("mesh"),
+                IO.String.Input("filename_prefix", default="mesh/ComfyUI"),
+            ],
+            hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo]
+        )
 
-    RETURN_TYPES = ()
+    @classmethod
-    FUNCTION = "save"
+    def execute(cls, mesh, filename_prefix) -> IO.NodeOutput:
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "3d"
-
-    def save(self, mesh, filename_prefix, prompt=None, extra_pnginfo=None):
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
         results = []
 
         metadata = {}
         if not args.disable_metadata:
-            if prompt is not None:
+            if cls.hidden.prompt is not None:
-                metadata["prompt"] = json.dumps(prompt)
+                metadata["prompt"] = json.dumps(cls.hidden.prompt)
-            if extra_pnginfo is not None:
+            if cls.hidden.extra_pnginfo is not None:
-                for x in extra_pnginfo:
+                for x in cls.hidden.extra_pnginfo:
-                    metadata[x] = json.dumps(extra_pnginfo[x])
+                    metadata[x] = json.dumps(cls.hidden.extra_pnginfo[x])
 
         for i in range(mesh.vertices.shape[0]):
             f = f"{filename}_{counter:05}_.glb"
@@ -616,15 +649,22 @@ class SaveGLB:
                 "type": "output"
             })
             counter += 1
-        return {"ui": {"3d": results}}
+        return IO.NodeOutput(ui={"3d": results})
 
 
-NODE_CLASS_MAPPINGS = {
+class Hunyuan3dExtension(ComfyExtension):
-    "EmptyLatentHunyuan3Dv2": EmptyLatentHunyuan3Dv2,
+    @override
-    "Hunyuan3Dv2Conditioning": Hunyuan3Dv2Conditioning,
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
-    "Hunyuan3Dv2ConditioningMultiView": Hunyuan3Dv2ConditioningMultiView,
+        return [
-    "VAEDecodeHunyuan3D": VAEDecodeHunyuan3D,
+            EmptyLatentHunyuan3Dv2,
-    "VoxelToMeshBasic": VoxelToMeshBasic,
+            Hunyuan3Dv2Conditioning,
-    "VoxelToMesh": VoxelToMesh,
+            Hunyuan3Dv2ConditioningMultiView,
-    "SaveGLB": SaveGLB,
+            VAEDecodeHunyuan3D,
-}
+            VoxelToMeshBasic,
+            VoxelToMesh,
+            SaveGLB,
+        ]
+
+
+async def comfy_entrypoint() -> Hunyuan3dExtension:
+    return Hunyuan3dExtension()