wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-15 11:37:32 +08:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

Small adjustments

Browse Source
This commit is contained in:
kijai 2026-05-13 17:06:25 +03:00
parent aac9a56d79
commit be79f67188
1 changed files with 21 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
comfy_extras/nodes_save_3d.py
View File

@ -1,12 +1,10 @@
"""Save-side 3D nodes: mesh packing/slicing helpers + GLB writer + SaveGLB node. """Save-side 3D nodes: mesh packing/slicing helpers + GLB writer + SaveGLB node."""
Pairs with nodes_load_3d.py (load-side counterpart).
"""
import json import json
import logging import logging
import os import os
import struct import struct
from io import BytesIO
import numpy as np import numpy as np
from PIL import Image 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, # 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. # 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. # 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) batch_size = len(vertices)
max_vertices = max(v.shape[0] for v in vertices) max_vertices = max(v.shape[0] for v in vertices)
max_faces = max(f.shape[0] for f in faces) 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) faces_np = faces_signed.astype(np.uint32)
texture_png_bytes = None texture_png_bytes = None
if texture_image is not None: if texture_image is not None:
import io as _io buf = BytesIO()
buf = _io.BytesIO()
texture_image.save(buf, format="PNG") texture_image.save(buf, format="PNG")
texture_png_bytes = buf.getvalue() 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) colors_buffer_padded = pad_to_4_bytes(colors_buffer)
texture_buffer_padded = pad_to_4_bytes(texture_buffer) texture_buffer_padded = pad_to_4_bytes(texture_buffer)
buffer_data = (vertices_buffer_padded + indices_buffer_padded buffer_data = b"".join([
+ uvs_buffer_padded + colors_buffer_padded + texture_buffer_padded) vertices_buffer_padded,
indices_buffer_padded,
uvs_buffer_padded,
colors_buffer_padded,
texture_buffer_padded,
])
vertices_byte_length = len(vertices_buffer) vertices_byte_length = len(vertices_buffer)
vertices_byte_offset = 0 vertices_byte_offset = 0
@ -270,7 +272,7 @@ def save_glb(vertices, faces, filepath, metadata=None,
if materials: if materials:
gltf["materials"] = materials gltf["materials"] = materials
if metadata is not None: if metadata:
gltf["asset"]["extras"] = metadata gltf["asset"]["extras"] = metadata
# Convert the JSON to bytes # 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) gltf_json_padded = pad_json_to_4_bytes(gltf_json)
# Create the GLB header # Create the GLB header (a 4-byte ASCII magic identifier glTF)
# Magic glTF
glb_header = struct.pack('<4sII', b'glTF', 2, 12 + 8 + len(gltf_json_padded) + 8 + len(buffer_data)) 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) # Create JSON chunk header (chunk type 0)
@ -355,21 +356,22 @@ class SaveGLB(IO.ComfyNode):
"subfolder": subfolder, "subfolder": subfolder,
"type": "output" "type": "output"
}) })
counter += 1
else: else:
# Handle Mesh input - save vertices and faces as GLB; carry optional UVs / colors / texture. # Handle Mesh input - save vertices and faces as GLB; carry optional UVs / colors / texture.
texture_b = getattr(mesh, "texture", None) 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]): for i in range(mesh.vertices.shape[0]):
vertices_i, faces_i, v_colors, uvs_i = get_mesh_batch_item(mesh, i) 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: if vertices_i.shape[0] == 0 or faces_i.shape[0] == 0:
logging.warning(f"SaveGLB: skipping empty mesh at batch index {i}") logging.warning(f"SaveGLB: skipping empty mesh at batch index {i}")
continue continue
tex_img = None tex_img = Image.fromarray(texture_np[i], mode="RGB") if texture_np is not None else 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")
f = f"{filename}_{counter:05}_.glb" f = f"{filename}_{counter:05}_.glb"
save_glb(vertices_i, faces_i, os.path.join(full_output_folder, f), metadata, save_glb(vertices_i, faces_i, os.path.join(full_output_folder, f), metadata,
uvs=uvs_i, uvs=uvs_i,