diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py index c8a487420..c03524246 100644 --- a/comfy_extras/nodes_save_3d.py +++ b/comfy_extras/nodes_save_3d.py @@ -1,12 +1,10 @@ -"""Save-side 3D nodes: mesh packing/slicing helpers + GLB writer + SaveGLB node. - -Pairs with nodes_load_3d.py (load-side counterpart). -""" +"""Save-side 3D nodes: mesh packing/slicing helpers + GLB writer + SaveGLB node.""" import json import logging import os import struct +from io import BytesIO import numpy as np from PIL import Image @@ -22,7 +20,7 @@ def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=Non # Pack lists of (Nᵢ, *) vertex/face/color/uv tensors into padded batched tensors, # stashing per-item lengths as runtime attrs so consumers can recover the real slice. # colors and uvs are 1:1 with vertices, so they're padded to max_vertices and read with vertex_counts. - # texture is (B, H, W, 3) — passed through unchanged since it doesn't need ragged packing. + # texture is (B, H, W, 3) — passed through unchanged batch_size = len(vertices) max_vertices = max(v.shape[0] for v in vertices) max_faces = max(f.shape[0] for f in faces) @@ -122,8 +120,7 @@ def save_glb(vertices, faces, filepath, metadata=None, faces_np = faces_signed.astype(np.uint32) texture_png_bytes = None if texture_image is not None: - import io as _io - buf = _io.BytesIO() + buf = BytesIO() texture_image.save(buf, format="PNG") texture_png_bytes = buf.getvalue() @@ -143,8 +140,13 @@ def save_glb(vertices, faces, filepath, metadata=None, colors_buffer_padded = pad_to_4_bytes(colors_buffer) texture_buffer_padded = pad_to_4_bytes(texture_buffer) - buffer_data = (vertices_buffer_padded + indices_buffer_padded - + uvs_buffer_padded + colors_buffer_padded + texture_buffer_padded) + buffer_data = b"".join([ + vertices_buffer_padded, + indices_buffer_padded, + uvs_buffer_padded, + colors_buffer_padded, + texture_buffer_padded, + ]) vertices_byte_length = len(vertices_buffer) vertices_byte_offset = 0 @@ -270,7 +272,7 @@ def save_glb(vertices, faces, filepath, metadata=None, if materials: gltf["materials"] = materials - if metadata is not None: + if metadata: gltf["asset"]["extras"] = metadata # Convert the JSON to bytes @@ -282,8 +284,7 @@ def save_glb(vertices, faces, filepath, metadata=None, gltf_json_padded = pad_json_to_4_bytes(gltf_json) - # Create the GLB header - # Magic glTF + # Create the GLB header (a 4-byte ASCII magic identifier glTF) glb_header = struct.pack('<4sII', b'glTF', 2, 12 + 8 + len(gltf_json_padded) + 8 + len(buffer_data)) # Create JSON chunk header (chunk type 0) @@ -355,21 +356,22 @@ class SaveGLB(IO.ComfyNode): "subfolder": subfolder, "type": "output" }) + counter += 1 else: # Handle Mesh input - save vertices and faces as GLB; carry optional UVs / colors / texture. texture_b = getattr(mesh, "texture", None) + texture_np = None + if texture_b is not None: + texture_np = (texture_b.clamp(0.0, 1.0).cpu().numpy() * 255).astype(np.uint8) + assert texture_np.ndim == 4 and texture_np.shape[-1] == 3, ( + f"texture must be (B, H, W, 3) RGB, got shape {tuple(texture_np.shape)}" + ) for i in range(mesh.vertices.shape[0]): vertices_i, faces_i, v_colors, uvs_i = get_mesh_batch_item(mesh, i) if vertices_i.shape[0] == 0 or faces_i.shape[0] == 0: logging.warning(f"SaveGLB: skipping empty mesh at batch index {i}") continue - tex_img = None - if texture_b is not None: - arr = (texture_b[i].clamp(0.0, 1.0).cpu().numpy() * 255).astype(np.uint8) - assert arr.ndim == 3 and arr.shape[-1] == 3, ( - f"texture[{i}] must be (H, W, 3) RGB, got shape {tuple(arr.shape)}" - ) - tex_img = Image.fromarray(arr, mode="RGB") + tex_img = Image.fromarray(texture_np[i], mode="RGB") if texture_np is not None else None f = f"{filename}_{counter:05}_.glb" save_glb(vertices_i, faces_i, os.path.join(full_output_folder, f), metadata, uvs=uvs_i,