fix: preserve 3D VAE and material properties

This commit is contained in:
LuqP2
2026-08-20 19:10:41 -03:00
parent abab421456
commit 7764c6db12
4 changed files with 62 additions and 10 deletions
+12 -2
View File
@@ -120,10 +120,14 @@ def _write_glb(
uv_array = None
texture_bytes = b""
texture_has_transparency = False
if texture is not None:
texture_array = _tensor_numpy(texture, np.float32)
texture_array = np.clip(texture_array * 255.0, 0, 255).astype(np.uint8)
if texture_array.ndim == 3 and texture_array.shape[-1] in (3, 4):
texture_has_transparency = (
texture_array.shape[-1] == 4 and bool(np.any(texture_array[..., 3] < 255))
)
buffer = BytesIO()
Image.fromarray(texture_array, mode="RGBA" if texture_array.shape[-1] == 4 else "RGB").save(buffer, "PNG")
texture_bytes = buffer.getvalue()
@@ -195,14 +199,20 @@ def _write_glb(
gltf["images"] = [{"bufferView": texture_view, "mimeType": "image/png"}]
gltf["samplers"] = [{"magFilter": 9729, "minFilter": 9729, "wrapS": 33071, "wrapT": 33071}]
gltf["textures"] = [{"source": 0, "sampler": 0}]
materials.append({
material = {
"pbrMetallicRoughness": {
"baseColorTexture": {"index": 0, "texCoord": 0},
"metallicFactor": 0.0,
"roughnessFactor": 1.0,
},
"doubleSided": True,
})
}
if texture_has_transparency:
material["alphaMode"] = "BLEND"
if unlit:
material["extensions"] = {"KHR_materials_unlit": {}}
gltf["extensionsUsed"] = ["KHR_materials_unlit"]
materials.append(material)
primitive["material"] = 0
elif unlit:
materials.append({
+23
View File
@@ -92,6 +92,29 @@ def test_texture_without_usable_uvs_is_not_reported_as_embedded(tmp_path):
assert "textures" not in gltf
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"
+15
View File
@@ -69,6 +69,21 @@ def test_workflow_extractor_traces_sampler_from_model_3d_input():
assert data["seed"] == 222
def test_workflow_extractor_traces_vae_from_model_3d_input():
prompt = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}},
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}},
"3": {"class_type": "VAEDecode", "inputs": {"samples": ["2", 0], "vae": ["1", 2]}},
"4": {"class_type": "SyntheticMeshGenerator", "inputs": {"images": ["3", 0]}},
"5": {"class_type": "MetaHubSave3DModel", "inputs": {"model_3d": ["4", 0]}},
}
extractor = WorkflowExtractor(prompt)
data, _missing = extractor.extract(save_node_id="5")
assert data["vae_name"] == "model.safetensors"
def test_workflow_extractor_traces_source_image_from_model_3d_input():
prompt = {
"1": {
+12 -8
View File
@@ -494,14 +494,18 @@ class WorkflowExtractor:
save_node = self._get_node(save_node_id)
if not save_node:
return None
images_conn = save_node.get("inputs", {}).get("images")
start_node_id = self._get_connection_node_id(images_conn)
if not start_node_id:
return None
return self._bfs_upstream(
[start_node_id],
lambda n: self._class_type(n) in self.VAE_DECODE_NODES,
)
inputs = save_node.get("inputs", {})
for input_name in ("images", "model_3d"):
start_node_id = self._get_connection_node_id(inputs.get(input_name))
if not start_node_id:
continue
vae_decode_id = self._bfs_upstream(
[start_node_id],
lambda n: self._class_type(n) in self.VAE_DECODE_NODES,
)
if vae_decode_id:
return vae_decode_id
return None
def _find_nodes_by_type(self, types: List[str]) -> List[str]:
matched = []