Better mesh

app.py

@@ -19,7 +19,7 @@ from typing import Union
 from shap_e.diffusion.sample import sample_latents
 from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
 from shap_e.models.download import load_model, load_config
-from shap_e.util.notebooks import create_pan_cameras, decode_latent_images
+from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, decode_latent_mesh
 
 from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
 from shap_e.models.transmitter.base import Transmitter, VectorDecoder
@@ -45,22 +45,6 @@ def set_state(s):
 def get_state():
     return state
 
-@torch.no_grad()
-def decode_latent_images_foo(
-    xm: Union[Transmitter, VectorDecoder],
-    latent: torch.Tensor,
-    cameras: DifferentiableCameraBatch,
-    rendering_mode: str = "stf",
-):
-    decoded = xm.renderer.render_views(
-        AttrDict(cameras=cameras),
-        params=(xm.encoder if isinstance(xm, Transmitter) else xm).bottleneck_to_params(
-            latent[None]
-        ),
-        options=AttrDict(rendering_mode=rendering_mode, render_with_direction=False),
-    )
-    return decoded
-
 def to_video(frames: list[Image.Image], fps: int = 5) -> str:
         out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
         writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps)
@@ -113,18 +97,8 @@ def generate_3D(input, grid_size=64):
 
     cameras = create_pan_cameras(size, device)
 
-    x=decode_latent_images_foo(xm, latents[0], cameras, rendering_mode=render_mode)
-    mesh=x['meshes'][0]
-    rm=x['raw_meshes'][0]
-
-    rm.vertex_channels["R"]=mesh.vertex_colors[:,0]
-    rm.vertex_channels["G"]=mesh.vertex_colors[:,1]
-    rm.vertex_channels["B"]=mesh.vertex_colors[:,2]
-
-    tm=rm.tri_mesh()
-
-    with open("/tmp/mesh.ply",'wb') as f:
-        tm.write_ply(f)
+    with open(f'/tmp/mesh.ply', 'wb') as f:
+        decode_latent_mesh(xm, latents[0]).tri_mesh().write_ply(f)
 
 
     set_state('Converting to point cloud...')