fix: preserve 3D VAE and material properties
This commit is contained in:
+12
-2
@@ -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({
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
@@ -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 = []
|
||||
|
||||
Reference in New Issue
Block a user