188 lines
6.9 KiB
Python
188 lines
6.9 KiB
Python
import json
|
|
import struct
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
|
|
import metahub_save_3d_node as node3d
|
|
|
|
|
|
class FakeFolderPaths:
|
|
def __init__(self, output_dir):
|
|
self.output_dir = output_dir
|
|
|
|
def get_output_directory(self):
|
|
return str(self.output_dir)
|
|
|
|
def get_save_image_path(self, _prefix, _output_root):
|
|
folder = self.output_dir / "3d"
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
return str(folder), "ComfyUI", 1, "3d", "3d/ComfyUI"
|
|
|
|
|
|
def _mesh(with_optional_data=True):
|
|
vertices = np.array([[[0, 0, 0], [1, 0, 0], [0, 1, 0]]], dtype=np.float32)
|
|
faces = np.array([[[0, 1, 2]]], dtype=np.int64)
|
|
return SimpleNamespace(
|
|
vertices=vertices,
|
|
faces=faces,
|
|
vertex_colors=np.array([[[1, 0, 0], [0, 1, 0], [0, 0, 1]]], dtype=np.float32) if with_optional_data else None,
|
|
uvs=np.array([[[0, 0], [1, 0], [0, 1]]], dtype=np.float32) if with_optional_data else None,
|
|
texture=np.ones((1, 2, 2, 3), dtype=np.float32) if with_optional_data else None,
|
|
unlit=False,
|
|
)
|
|
|
|
|
|
def _read_glb_json(path):
|
|
data = path.read_bytes()
|
|
assert data[:4] == b"glTF"
|
|
json_length, json_type = struct.unpack_from("<II", data, 12)
|
|
assert json_type == 0x4E4F534A
|
|
return json.loads(data[20:20 + json_length].rstrip(b" \x00"))
|
|
|
|
|
|
def test_mesh_writes_glb_sidecar_and_embedded_metadata(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(node3d, "folder_paths", FakeFolderPaths(tmp_path))
|
|
monkeypatch.setattr(node3d.args, "disable_metadata", False)
|
|
monkeypatch.setattr(node3d, "_build_metadata", lambda *args, **kwargs: {
|
|
"schema_version": 1,
|
|
"media_type": "model3d",
|
|
"prompt": "synthetic prompt",
|
|
"workflow": {"nodes": []},
|
|
"prompt_api": {},
|
|
})
|
|
|
|
result = node3d.MetaHubSave3DModel().save_model(_mesh())
|
|
|
|
model_path = tmp_path / "3d" / result["ui"]["3d"][0]["filename"]
|
|
sidecar_path = model_path.with_name(model_path.name + ".imagemetahub.json")
|
|
sidecar = json.loads(sidecar_path.read_text(encoding="utf-8"))
|
|
embedded = _read_glb_json(model_path)["asset"]["extras"]["imagemetahub_data"]
|
|
assert sidecar == embedded
|
|
assert sidecar["model_3d"]["vertexCount"] == 3
|
|
assert sidecar["model_3d"]["faceCount"] == 1
|
|
assert sidecar["model_3d"]["hasTextures"] is True
|
|
assert result["ui"]["imagemetahub_files"] == [str(model_path)]
|
|
|
|
|
|
def test_disable_metadata_writes_model_without_sidecar_or_extras(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(node3d, "folder_paths", FakeFolderPaths(tmp_path))
|
|
monkeypatch.setattr(node3d.args, "disable_metadata", True)
|
|
|
|
result = node3d.MetaHubSave3DModel().save_model(_mesh(with_optional_data=False))
|
|
|
|
model_path = tmp_path / "3d" / result["ui"]["3d"][0]["filename"]
|
|
assert not model_path.with_name(model_path.name + ".imagemetahub.json").exists()
|
|
assert "extras" not in _read_glb_json(model_path)["asset"]
|
|
|
|
|
|
def test_texture_without_usable_uvs_is_not_reported_as_embedded(tmp_path):
|
|
mesh = _mesh()
|
|
|
|
geometry = node3d._write_glb(
|
|
tmp_path / "no-uvs.glb",
|
|
mesh.vertices[0],
|
|
mesh.faces[0],
|
|
{},
|
|
texture=mesh.texture[0],
|
|
)
|
|
gltf = _read_glb_json(tmp_path / "no-uvs.glb")
|
|
|
|
assert geometry["hasTextures"] is False
|
|
assert "textures" not in gltf
|
|
|
|
|
|
def test_build_metadata_preserves_zero_denoise(monkeypatch):
|
|
class FakeExtractor:
|
|
def __init__(self, _prompt):
|
|
pass
|
|
|
|
def extract(self, save_node_id):
|
|
assert save_node_id == "3"
|
|
return {"denoise": 0.0, "lora_list": []}, set()
|
|
|
|
monkeypatch.setattr(node3d, "WorkflowExtractor", FakeExtractor)
|
|
monkeypatch.setattr(node3d.utils, "get_workflow_json", lambda _info: {})
|
|
monkeypatch.setattr(node3d.utils, "ensure_prompt_in_workflow", lambda workflow, _prompt: workflow)
|
|
monkeypatch.setattr(node3d.utils, "ensure_metahub_save_node", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(node3d.utils, "extract_loras_from_workflow", lambda _workflow: [])
|
|
monkeypatch.setattr(node3d.utils, "collect_gpu_metrics", lambda: {})
|
|
monkeypatch.setattr(node3d.utils, "collect_version_info", lambda: {})
|
|
monkeypatch.setattr(node3d.utils, "build_metadata_sources", lambda *args: {})
|
|
monkeypatch.setattr(node3d.utils, "extract_workflow_attribution", lambda *args: {})
|
|
monkeypatch.setattr(node3d.utils, "build_imh_metadata", lambda params, _workflow: params)
|
|
|
|
metadata = node3d._build_metadata({}, {}, "3", "", "", "", None)
|
|
|
|
assert metadata["denoise"] == 0.0
|
|
|
|
|
|
def test_textured_unlit_mesh_preserves_unlit_and_transparency(tmp_path):
|
|
mesh = _mesh()
|
|
mesh.unlit = True
|
|
mesh.texture = np.ones((1, 2, 2, 4), dtype=np.float32)
|
|
mesh.texture[0, 0, 0, 3] = 0.25
|
|
|
|
node3d._write_glb(
|
|
tmp_path / "textured-unlit.glb",
|
|
mesh.vertices[0],
|
|
mesh.faces[0],
|
|
{},
|
|
uvs=mesh.uvs[0],
|
|
texture=mesh.texture[0],
|
|
unlit=mesh.unlit,
|
|
)
|
|
gltf = _read_glb_json(tmp_path / "textured-unlit.glb")
|
|
material = gltf["materials"][0]
|
|
|
|
assert material["extensions"] == {"KHR_materials_unlit": {}}
|
|
assert gltf["extensionsUsed"] == ["KHR_materials_unlit"]
|
|
assert material["alphaMode"] == "BLEND"
|
|
|
|
|
|
def test_file3d_preserves_received_format_and_writes_sidecar(tmp_path, monkeypatch):
|
|
class FakeFile3D:
|
|
format = "obj"
|
|
|
|
def save_to(self, path):
|
|
with open(path, "wb") as handle:
|
|
handle.write(b"o synthetic\n")
|
|
|
|
monkeypatch.setattr(node3d, "folder_paths", FakeFolderPaths(tmp_path))
|
|
monkeypatch.setattr(node3d.args, "disable_metadata", False)
|
|
monkeypatch.setattr(node3d, "_build_metadata", lambda *args, **kwargs: {
|
|
"schema_version": 1,
|
|
"media_type": "model3d",
|
|
"workflow": {},
|
|
"prompt_api": {},
|
|
})
|
|
|
|
result = node3d.MetaHubSave3DModel().save_model(FakeFile3D())
|
|
|
|
model_path = tmp_path / "3d" / result["ui"]["3d"][0]["filename"]
|
|
assert model_path.suffix == ".obj"
|
|
assert model_path.read_bytes() == b"o synthetic\n"
|
|
assert model_path.with_name(model_path.name + ".imagemetahub.json").exists()
|
|
|
|
|
|
def test_source_node_class_is_read_from_prompt_link():
|
|
prompt = {
|
|
"7": {"class_type": "SyntheticMeshGenerator", "inputs": {}},
|
|
"8": {"class_type": "MetaHubSave3DModel", "inputs": {"model_3d": ["7", 0]}},
|
|
}
|
|
|
|
assert node3d._source_node_class(prompt, "8") == "SyntheticMeshGenerator"
|
|
|
|
|
|
def test_mesh_batch_increments_filenames(tmp_path, monkeypatch):
|
|
mesh = _mesh(with_optional_data=False)
|
|
mesh.vertices = np.repeat(mesh.vertices, 2, axis=0)
|
|
mesh.faces = np.repeat(mesh.faces, 2, axis=0)
|
|
monkeypatch.setattr(node3d, "folder_paths", FakeFolderPaths(tmp_path))
|
|
monkeypatch.setattr(node3d.args, "disable_metadata", True)
|
|
|
|
result = node3d.MetaHubSave3DModel().save_model(mesh)
|
|
|
|
filenames = [item["filename"] for item in result["ui"]["3d"]]
|
|
assert filenames == ["ComfyUI_00001_.glb", "ComfyUI_00002_.glb"]
|