Decodig sampled SLATs vs Encoded SLATs

Hello and thanks for your great twork!

I'm trying to test the reconstruction quality of the shape-VAEs (I'm reconstructing a textureless mesh).
When running example.py and saving the intermediate mesh that is decoded by the shape decoder in
https://github.com/microsoft/TRELLIS.2/blob/903bfcf51af09b3fc6c408d2c7f2335febab7f81/trellis2/pipelines/trellis2_image_to_3d.py#L470
it looks resonably well (see below - image2mesh-decoder output ).
Unfortunately, when I'm using the shape-encoder (even on the same mesh created by image2mesh - image2mesh-full below) and decoding the SLATs (with the code below - no postprocess), i'm getting very voxelized mesh (encode-decode below).
what is the difference between the two? How come sampled SLATS are so different from encoded ones wasn't the diffusion model trained to sample encoded SLATs?

Images and code are below
Thanks for helping

input-image	image2mesh-decoder output	image2mesh-full	encode-decode

import torch
import trimesh
import o_voxel
from trellis2 import models
from trellis2.modules.sparse import SparseTensor
from trellis2.representations import MeshWithVoxel


def load_shape_models():
    shape_enc = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16")
    shape_dec = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16")
    for model in [shape_enc, shape_dec]:
        model.cuda().eval()
    
    return shape_enc, shape_dec


def normalize_mesh(mesh):
    """Normalize mesh vertices to unit cube."""
    vertices = torch.from_numpy(mesh.vertices).float()
    faces = torch.from_numpy(mesh.faces).long()
    
    center = (vertices.min(0)[0] + vertices.max(0)[0]) / 2
    scale = 0.99999 / (vertices.max(0)[0] - vertices.min(0)[0]).max()
    vertices = (vertices - center) * scale
    
    return vertices, faces


def encode_decode_shape(encoder, decoder, vertices, faces, resolution=512):
    voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
        vertices, faces, grid_size=resolution,
        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
    )
    
    coords = torch.cat([torch.zeros(len(voxel_indices), 1, dtype=torch.int32), voxel_indices], dim=1)
    vertices_sparse = SparseTensor(dual_vertices.cuda(), coords.cuda())
    intersected_sparse = SparseTensor(intersected.float().cuda(), coords.cuda())
    
    with torch.no_grad():
        latent = encoder(vertices_sparse, intersected_sparse)
        decoder.set_resolution(resolution)
        meshes, subs = decoder(latent, return_subs=True)
    
    return meshes[0], subs


def postprocess_with_default_texture(mesh, subs, resolution):
    pbr_attr_layout = {
        'base_color': slice(0, 3),
        'metallic': slice(3, 4),
        'roughness': slice(4, 5),
        'alpha': slice(5, 6),
    }
    
    mesh.fill_holes()
    
    # Create default white material
    sub = subs[-1]
    num_voxels = sub.coords.shape[0]
    default_attrs = torch.ones(num_voxels, 6, device=sub.coords.device, dtype=torch.float32)
    default_attrs[:, 3] = 0.0  # metallic
    default_attrs[:, 4] = 0.5  # roughness
    default_attrs[:, 5] = 1.0  # alpha
    
    mesh_with_voxel = MeshWithVoxel(
        mesh.vertices, mesh.faces,
        origin=[-0.5, -0.5, -0.5],
        voxel_size=1 / resolution,
        coords=sub.coords[:, 1:],
        attrs=default_attrs,
        voxel_shape=torch.Size([*sub.shape, *sub.spatial_shape]),
        layout=pbr_attr_layout
    )
    
    mesh_with_voxel.simplify(16777216)
    
    glb = o_voxel.postprocess.to_glb(
        vertices=mesh_with_voxel.vertices,
        faces=mesh_with_voxel.faces,
        attr_volume=mesh_with_voxel.attrs,
        coords=mesh_with_voxel.coords,
        attr_layout=mesh_with_voxel.layout,
        voxel_size=mesh_with_voxel.voxel_size,
        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
        decimation_target=1000000,
        texture_size=4096,
        remesh=True,
        remesh_band=1,
        remesh_project=0,
        verbose=True
    )
    return glb


def save_mesh(mesh, output_path):
    trimesh.Trimesh(
        vertices=mesh.vertices.cpu().numpy(),
        faces=mesh.faces.cpu().numpy(),
    ).export(output_path)


def main():
    shape_enc, shape_dec = load_shape_models()
    
    mesh_path = "/common_data/workspace/ariel/repos/Hunyuan3D-2.1/assets/cpt_america_out.glb"
    mesh = trimesh.load(mesh_path)
    if isinstance(mesh, trimesh.Scene):
        mesh = list(mesh.geometry.values())[0]
    vertices, faces = normalize_mesh(mesh)
    
    resolution = 1024
    shape_mesh, subs = encode_decode_shape(shape_enc, shape_dec, vertices, faces, resolution=resolution)
    save_mesh(shape_mesh, "reconstructed_raw_shape.glb")

    glb_mesh = postprocess_with_default_texture(shape_mesh, subs, resolution)
    glb_mesh.export("reconstructed_mesh_postprocessed.glb")
    


if __name__ == "__main__":
    main()

Start with example.py and trellis2/pipelines/trellis2_image_to_3d.py around line 470, then run the posted encode_decode_shape script with the supplied mesh and resolution. Compare the sampled-SLAT and encoder-decoder paths, including normalization and mesh conversion. Done means identifying and documenting the cause of the voxelized reconstruction, or confirming the needed correction.