Skip to content

Commit 78d9bfd

Browse files
jtydhr88jtydhr88
committed
CORE-329 feat: add Save 3D (Advanced) node family
1 parent ba3f697 commit 78d9bfd

1 file changed

Lines changed: 157 additions & 2 deletions

File tree

‎comfy_extras/nodes_save_3d.py‎

Lines changed: 157 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313

1414
import folder_paths
1515
from comfy.cli_args import args
16-
from comfy_api.latest import ComfyExtension, IO, Types
16+
from comfy_api.latest import ComfyExtension, IO, Types, UI
1717

1818

1919
def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False):
@@ -406,10 +406,165 @@ def execute(cls, mesh: Types.MESH | Types.File3D, filename_prefix: str) -> IO.No
406406
return IO.NodeOutput(ui={"3d": results})
407407

408408

409+
def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str:
410+
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
411+
filename_prefix, folder_paths.get_output_directory()
412+
)
413+
ext = model_3d.format or "glb"
414+
saved_filename = f"{filename}_{counter:05}.{ext}"
415+
model_3d.save_to(os.path.join(full_output_folder, saved_filename))
416+
return f"{subfolder}/{saved_filename}" if subfolder else saved_filename
417+
418+
419+
def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
420+
model_file = _save_file3d_to_output(model_3d, filename_prefix)
421+
camera_info_input = kwargs.get("camera_info", None)
422+
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
423+
model_3d_info_input = kwargs.get("model_3d_info", None)
424+
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
425+
return IO.NodeOutput(
426+
model_3d,
427+
model_3d_info,
428+
camera_info,
429+
width,
430+
height,
431+
ui=UI.PreviewUI3DAdvanced(model_file, camera_info, model_3d_info),
432+
)
433+
434+
435+
class Save3DAdvanced(IO.ComfyNode):
436+
@classmethod
437+
def define_schema(cls):
438+
return IO.Schema(
439+
node_id="Save3DAdvanced",
440+
display_name="Save 3D (Advanced)",
441+
search_aliases=["save 3d", "export 3d model", "save mesh advanced"],
442+
category="3d",
443+
is_experimental=True,
444+
is_output_node=True,
445+
inputs=[
446+
IO.MultiType.Input(
447+
"model_3d",
448+
types=[
449+
IO.File3DGLB,
450+
IO.File3DGLTF,
451+
IO.File3DFBX,
452+
IO.File3DOBJ,
453+
IO.File3DSTL,
454+
IO.File3DUSDZ,
455+
IO.File3DAny,
456+
],
457+
tooltip="3D model file from an upstream 3D node.",
458+
),
459+
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
460+
IO.Load3D.Input("viewport_state"),
461+
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
462+
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
463+
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
464+
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
465+
],
466+
outputs=[
467+
IO.File3DAny.Output(display_name="model_3d"),
468+
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
469+
IO.Load3DCamera.Output(display_name="camera_info"),
470+
IO.Int.Output(display_name="width"),
471+
IO.Int.Output(display_name="height"),
472+
],
473+
)
474+
475+
@classmethod
476+
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
477+
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
478+
479+
480+
class SaveGaussianSplat(IO.ComfyNode):
481+
@classmethod
482+
def define_schema(cls):
483+
return IO.Schema(
484+
node_id="SaveGaussianSplat",
485+
display_name="Save Splat",
486+
search_aliases=["save splat", "save gaussian splat", "export gaussian", "export splat"],
487+
category="3d",
488+
is_experimental=True,
489+
is_output_node=True,
490+
inputs=[
491+
IO.MultiType.Input(
492+
"model_3d",
493+
types=[
494+
IO.File3DSplatAny,
495+
IO.File3DPLY,
496+
IO.File3DSPLAT,
497+
IO.File3DSPZ,
498+
IO.File3DKSPLAT,
499+
],
500+
tooltip="A gaussian splat 3D file.",
501+
),
502+
503+
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
504+
IO.Load3D.Input("viewport_state"),
505+
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
506+
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
507+
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
508+
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
509+
],
510+
outputs=[
511+
IO.File3DSplatAny.Output(display_name="model_3d"),
512+
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
513+
IO.Load3DCamera.Output(display_name="camera_info"),
514+
IO.Int.Output(display_name="width"),
515+
IO.Int.Output(display_name="height"),
516+
],
517+
)
518+
519+
@classmethod
520+
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
521+
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
522+
523+
524+
class SavePointCloud(IO.ComfyNode):
525+
@classmethod
526+
def define_schema(cls):
527+
return IO.Schema(
528+
node_id="SavePointCloud",
529+
display_name="Save Point Cloud",
530+
search_aliases=["save point cloud", "save pointcloud", "export point cloud"],
531+
category="3d",
532+
is_experimental=True,
533+
is_output_node=True,
534+
inputs=[
535+
IO.MultiType.Input(
536+
"model_3d",
537+
types=[
538+
IO.File3DPointCloudAny,
539+
IO.File3DPLY,
540+
],
541+
tooltip="Point cloud file (.ply)",
542+
),
543+
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
544+
IO.Load3D.Input("viewport_state"),
545+
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
546+
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
547+
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
548+
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
549+
],
550+
outputs=[
551+
IO.File3DPointCloudAny.Output(display_name="model_3d"),
552+
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
553+
IO.Load3DCamera.Output(display_name="camera_info"),
554+
IO.Int.Output(display_name="width"),
555+
IO.Int.Output(display_name="height"),
556+
],
557+
)
558+
559+
@classmethod
560+
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
561+
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
562+
563+
409564
class Save3DExtension(ComfyExtension):
410565
@override
411566
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
412-
return [SaveGLB]
567+
return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud]
413568

414569

415570
async def comfy_entrypoint() -> Save3DExtension:

0 commit comments

Comments
 (0)