Add isolated pre-simplification geometry replay

scripts/replay_geometry_export.py

@@ -0,0 +1,40 @@
+"""Diagnostic replay of saved tensors; capture or substitute pre-simplification remesh."""
+from pathlib import Path
+import json
+import numpy as np
+
+
+def replay(raw_path, layout_path, output, filtered_path=None):
+    import torch
+    import cumesh
+    import o_voxel
+    output = Path(output)
+    output.mkdir(parents=True, exist_ok=True)
+    with np.load(raw_path, allow_pickle=False) as data:
+        arrays = {key: data[key].copy() for key in data.files}
+    layout = {key: slice(*value) for key, value in json.loads(Path(layout_path).read_text()).items()}
+    original = cumesh.remeshing.remesh_narrow_band_dc
+    class Captured(Exception):
+        pass
+    def interception(*args, **kwargs):
+        if filtered_path:
+            with np.load(filtered_path, allow_pickle=False) as data:
+                return (torch.from_numpy(data['vertices'].astype(np.float32)).cuda(),
+                        torch.from_numpy(data['faces'].astype(np.int32)).cuda())
+        vertices, faces = original(*args, **kwargs)
+        np.savez_compressed(output / 'pre-simplify.npz', vertices=vertices.cpu().numpy(), faces=faces.cpu().numpy())
+        raise Captured()
+    cumesh.remeshing.remesh_narrow_band_dc = interception
+    try:
+        result = o_voxel.postprocess.to_glb(
+            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(),
+            attr_layout=layout, grid_size=int(arrays['grid_size']) if arrays['grid_size'].ndim == 0
+            else arrays['grid_size'].tolist(), aabb=[[-.5, -.5, -.5], [.5, .5, .5]],
+            decimation_target=50000, texture_size=1024, remesh=True, remesh_band=1, remesh_project=0)
+        result.export(output / 'filtered-remesh.glb')
+    except Captured:
+        return str(output / 'pre-simplify.npz')
+    finally:
+        cumesh.remeshing.remesh_narrow_band_dc = original
+    return str(output / 'filtered-remesh.glb')

studio/native.py

@@ -132,3 +132,9 @@ def compare_geometry_gpu(image, seed, output):
         capture(meshes[0], latents[2], pipeline.pbr_attr_layout, output)
         compare(output)
     return output
+
+
+@spaces.GPU(duration=240)
+def replay_geometry_gpu(raw, layout, output, filtered=None):
+    from scripts.replay_geometry_export import replay
+    return replay(raw, layout, output, filtered)

studio/worker.py

@@ -88,6 +88,10 @@ class WorkerEndpoints:
         self.models.geometry(image, int(seed), output)
         return str(output)
 
+    def replay_geometry(self, raw, layout, filtered):
+        from .native import replay_geometry_gpu
+        return replay_geometry_gpu(raw, layout, str(self.directory()), filtered)
+
     def compare_geometry(self, image, seed):
         from .native import compare_geometry_gpu
         with Image.open(image) as source:
@@ -162,6 +166,12 @@ def build_worker(role, models, output_root="outputs/worker"):
         gr.Button("Generate").click(getattr(endpoints, role), inputs, outputs, api_name=role,
                                     concurrency_limit=1)
         if role == "geometry":
+            raw = gr.File(label="Raw tensor NPZ", type="filepath")
+            layout = gr.File(label="Attribute layout JSON", type="filepath")
+            filtered = gr.File(label="Filtered remesh NPZ (optional)", type="filepath")
+            replay_output = gr.File(label="Replay output")
+            gr.Button("Replay saved tensors").click(endpoints.replay_geometry, [raw, layout, filtered], replay_output,
+                api_name="replay_geometry", concurrency_limit=1, concurrency_id="geometry")
             diagnostic = gr.File(label="Same-input export comparison")
             gr.Button("Compare export modes").click(endpoints.compare_geometry, inputs, diagnostic,
                 api_name="compare_geometry", concurrency_limit=1, concurrency_id="geometry")