fix remesh export resolution

export_worker.py

@@ -27,13 +27,13 @@ def export_glb(
     arrays = np.load(payload_npz)
     meta = json.loads(payload_meta.read_text(encoding="utf-8"))
     attr_layout = _deserialize_attr_layout(meta["attr_layout"])
+    resolution = int(meta["resolution"])
+    aabb = meta["aabb"]
 
     vertices = torch.from_numpy(arrays["vertices"]).cuda()
     faces = torch.from_numpy(arrays["faces"]).cuda()
     attr_volume = torch.from_numpy(arrays["attrs"]).cuda()
     coords = torch.from_numpy(arrays["coords"]).cuda()
-    grid_size = torch.from_numpy(arrays["grid_size"]).cuda()
-    aabb = torch.from_numpy(arrays["aabb"]).cuda()
 
     torch.cuda.synchronize()
     glb = o_voxel.postprocess.to_glb(
@@ -42,7 +42,7 @@ def export_glb(
         attr_volume=attr_volume,
         coords=coords,
         attr_layout=attr_layout,
-        grid_size=grid_size,
+        grid_size=resolution,
         aabb=aabb,
         decimation_target=decimation_target,
         texture_size=texture_size,

service_runtime.py

@@ -133,7 +133,7 @@ class TrellisRuntime:
             "1536": "1536_cascade",
         }[generation.resolution]
         try:
-            outputs = pipeline.run(
+            outputs, latents = pipeline.run(
                 image,
                 seed=generation.seed,
                 preprocess_image=False,
@@ -156,12 +156,14 @@ class TrellisRuntime:
                     "rescale_t": generation.tex_slat_rescale_t,
                 },
                 pipeline_type=pipeline_type,
-                return_latent=False,
+                return_latent=True,
             )
             torch.cuda.synchronize()
             mesh = outputs[0]
-            payload = self._mesh_to_payload(mesh)
+            _, _, resolution = latents
+            payload = self._mesh_to_payload(mesh, resolution)
             del outputs
+            del latents
             del mesh
             torch.cuda.empty_cache()
             return payload
@@ -171,20 +173,13 @@ class TrellisRuntime:
             raise classify_runtime_error("generate", error) from error
 
     @staticmethod
-    def _mesh_to_payload(mesh: Any) -> dict[str, Any]:
-        origin = mesh.origin.detach().cpu().numpy().astype(np.float32)
-        spatial_shape = np.array(mesh.voxel_shape[-3:], dtype=np.int32)
-        voxel_size = float(mesh.voxel_size)
-        aabb = np.stack(
-            [origin, origin + voxel_size * spatial_shape.astype(np.float32)]
-        )
+    def _mesh_to_payload(mesh: Any, resolution: int) -> dict[str, Any]:
         return {
             "vertices": mesh.vertices.detach().cpu().numpy().astype(np.float32),
             "faces": mesh.faces.detach().cpu().numpy().astype(np.int32),
             "attrs": mesh.attrs.detach().cpu().numpy().astype(np.float32),
             "coords": mesh.coords.detach().cpu().numpy().astype(np.int32),
-            "grid_size": spatial_shape,
-            "aabb": aabb.astype(np.float32),
+            "resolution": int(resolution),
             "attr_layout": {
                 key: {"start": value.start, "stop": value.stop}
                 for key, value in mesh.layout.items()
@@ -205,11 +200,17 @@ def save_export_payload(job_dir: Path, payload: dict[str, Any]) -> tuple[Path, P
         faces=payload["faces"],
         attrs=payload["attrs"],
         coords=payload["coords"],
-        grid_size=payload["grid_size"],
-        aabb=payload["aabb"],
     )
     meta_path.write_text(
-        json.dumps({"attr_layout": payload["attr_layout"]}, indent=2, sort_keys=True),
+        json.dumps(
+            {
+                "attr_layout": payload["attr_layout"],
+                "resolution": payload["resolution"],
+                "aabb": [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+            },
+            indent=2,
+            sort_keys=True,
+        ),
         encoding="utf-8",
     )
     return npz_path, meta_path