Skip to content

Commit eea3502

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

1 file changed

Lines changed: 156 additions & 2 deletions

File tree

‎comfy_extras/nodes_save_3d.py‎

Lines changed: 156 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,164 @@ 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.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
460+
IO.Load3D.Input("viewport_state"),
461+
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
462+
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
463+
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
464+
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
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+
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
503+
IO.Load3D.Input("viewport_state"),
504+
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
505+
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
506+
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
507+
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
508+
],
509+
outputs=[
510+
IO.File3DSplatAny.Output(display_name="model_3d"),
511+
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
512+
IO.Load3DCamera.Output(display_name="camera_info"),
513+
IO.Int.Output(display_name="width"),
514+
IO.Int.Output(display_name="height"),
515+
],
516+
)
517+
518+
@classmethod
519+
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
520+
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
521+
522+
523+
class SavePointCloud(IO.ComfyNode):
524+
@classmethod
525+
def define_schema(cls):
526+
return IO.Schema(
527+
node_id="SavePointCloud",
528+
display_name="Save Point Cloud",
529+
search_aliases=["save point cloud", "save pointcloud", "export point cloud"],
530+
category="3d",
531+
is_experimental=True,
532+
is_output_node=True,
533+
inputs=[
534+
IO.MultiType.Input(
535+
"model_3d",
536+
types=[
537+
IO.File3DPointCloudAny,
538+
IO.File3DPLY,
539+
],
540+
tooltip="Point cloud file (.ply)",
541+
),
542+
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
543+
IO.Load3D.Input("viewport_state"),
544+
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
545+
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
546+
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
547+
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
548+
],
549+
outputs=[
550+
IO.File3DPointCloudAny.Output(display_name="model_3d"),
551+
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
552+
IO.Load3DCamera.Output(display_name="camera_info"),
553+
IO.Int.Output(display_name="width"),
554+
IO.Int.Output(display_name="height"),
555+
],
556+
)
557+
558+
@classmethod
559+
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
560+
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
561+
562+
409563
class Save3DExtension(ComfyExtension):
410564
@override
411565
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
412-
return [SaveGLB]
566+
return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud]
413567

414568

415569
async def comfy_entrypoint() -> Save3DExtension:

0 commit comments

Comments
 (0)