From 7764c6db12ff06f11625c89edb19db682e2078d2 Mon Sep 17 00:00:00 2001
From: LuqP2 <kelipot.l@gmail.com>
Date: Thu, 20 Aug 2026 19:10:41 -0300
Subject: [PATCH] fix: preserve 3D VAE and material properties

---
 metahub_save_3d_node.py            | 14 ++++++++++++--
 tests/test_metahub_save_3d_node.py | 23 +++++++++++++++++++++++
 tests/test_workflow_extractor.py   | 15 +++++++++++++++
 workflow_extractor.py              | 20 ++++++++++++--------
 4 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/metahub_save_3d_node.py b/metahub_save_3d_node.py
index d8b2adb..720bfaa 100644
--- a/metahub_save_3d_node.py
+++ b/metahub_save_3d_node.py
@@ -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({
diff --git a/tests/test_metahub_save_3d_node.py b/tests/test_metahub_save_3d_node.py
index a16feeb..88c4bf6 100644
--- a/tests/test_metahub_save_3d_node.py
+++ b/tests/test_metahub_save_3d_node.py
@@ -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"
diff --git a/tests/test_workflow_extractor.py b/tests/test_workflow_extractor.py
index 3b7ac51..d79b57b 100644
--- a/tests/test_workflow_extractor.py
+++ b/tests/test_workflow_extractor.py
@@ -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": {
diff --git a/workflow_extractor.py b/workflow_extractor.py
index d81a4fb..524bf02 100644
--- a/workflow_extractor.py
+++ b/workflow_extractor.py
@@ -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 = []