""" Visualization + render test for TRELLIS.2. Run this instead of app.py to: 1. Generate a mesh with full pipeline visualizations saved at every stage. 2. Render the resulting mesh with render_utils.render_snapshot — the same call that app.py uses — and save every render frame to disk. This is the SMOKING GUN test: if the mesh looks correct in the decode-stage visualizations but the render images look wrong (15-30% coverage), the bug is inside the renderer / nvdiffrast path, not in the pipeline. """ import os os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" import cv2 import numpy as np import torch from PIL import Image from trellis2.pipelines import Trellis2ImageTo3DPipeline from trellis2.renderers import EnvMap from trellis2.utils import render_utils # --------------------------------------------------------------------------- # Config — edit these to taste # --------------------------------------------------------------------------- IMAGE_PATH = "assets/example_image/T2.png" PIPELINE = "1024_cascade" # '512' | '1024' | '1024_cascade' | '1536_cascade' SEED = 42 NVIEWS = 8 # number of render frames (matches app.py STEPS=8) RENDER_RES = 1024 # render resolution (matches app.py) VIZ_DIR = "visualizations_render_test" # --------------------------------------------------------------------------- def save_render_frames(images: dict, out_dir: str, prefix: str = "render"): """ Save every key × every frame from render_snapshot output to disk. images is a dict like: {'shaded': [np.uint8 (H,W,3), ...], 'normal': [...], 'base_color': [...], 'metallic': [...], 'roughness': [...], 'alpha': [...]} """ os.makedirs(out_dir, exist_ok=True) saved = [] for key, frames in images.items(): for i, frame in enumerate(frames): path = os.path.join(out_dir, f"{prefix}_{key}_{i:03d}.png") img = Image.fromarray(frame) img.save(path) saved.append(path) return saved def make_contact_sheet(images: dict, out_dir: str, prefix: str = "contact"): """ Build one wide contact sheet per render key and save it. Useful for seeing all views at once without opening 40 files. """ os.makedirs(out_dir, exist_ok=True) paths = [] for key, frames in images.items(): if not frames: continue row = np.concatenate(frames, axis=1) # stack horizontally path = os.path.join(out_dir, f"{prefix}_{key}_all_views.png") Image.fromarray(row).save(path) paths.append(path) print(f" Contact sheet [{key}]: {row.shape[1]}x{row.shape[0]} -> {path}") return paths def export_obj(mesh, path: str): """ Export mesh vertices and faces directly to a Wavefront .obj file. Uses NO nvdiffrast, NO flex_gemm, NO ROCm GPU ops — pure Python/numpy. Load in Blender to verify 100% geometry completeness. """ verts = mesh.vertices.detach().cpu().numpy() # [N, 3] faces = mesh.faces.detach().cpu().numpy() # [F, 3] os.makedirs(os.path.dirname(path) or ".", exist_ok=True) with open(path, "w") as f: f.write("# TRELLIS.2 raw mesh export (no renderer, no GLB pipeline)\n") f.write(f"# {verts.shape[0]} vertices, {faces.shape[0]} faces\n\n") for v in verts: f.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n") f.write("\n") for tri in faces: # .obj uses 1-based indices f.write(f"f {tri[0]+1} {tri[1]+1} {tri[2]+1}\n") print(f" OBJ export: {verts.shape[0]} vertices, {faces.shape[0]} faces -> {path}") def main(): print("=" * 70) print("TRELLIS.2 render smoke-test") print("=" * 70) # ------------------------------------------------------------------ # Load pipeline # ------------------------------------------------------------------ print("\nLoading pipeline...") pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B") pipeline.cuda() # ------------------------------------------------------------------ # Load HDR environment maps (same as app.py) # ------------------------------------------------------------------ print("Loading environment maps...") envmap = { 'forest': EnvMap(torch.tensor( cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda' )), 'sunset': EnvMap(torch.tensor( cv2.cvtColor(cv2.imread('assets/hdri/sunset.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda' )), 'courtyard': EnvMap(torch.tensor( cv2.cvtColor(cv2.imread('assets/hdri/courtyard.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda' )), } image = Image.open(IMAGE_PATH) # ------------------------------------------------------------------ # Run pipeline with ALL stage visualizations enabled # ------------------------------------------------------------------ print(f"\nRunning pipeline (type={PIPELINE}, seed={SEED}) ...") print(f"Stage visualizations will be saved to: {VIZ_DIR}/\n") mesh = pipeline.run( image, seed=SEED, pipeline_type=PIPELINE, visualize_sparse_structure=False, visualize_save_dir=VIZ_DIR, ) print("\nPipeline complete. Mesh object:", type(mesh[0]).__name__) print(f" vertices : {mesh[0].vertices.shape}") print(f" faces : {mesh[0].faces.shape}") if hasattr(mesh[0], 'coords'): print(f" vox coords: {mesh[0].coords.shape}") if hasattr(mesh[0], 'attrs'): print(f" vox attrs : {mesh[0].attrs.shape}") # ------------------------------------------------------------------ # BYPASS TEST: export raw .obj — zero GPU rendering, zero nvdiffrast # Load in Blender to verify geometry is 100% complete before blaming # the renderer or GLB pipeline. # ------------------------------------------------------------------ obj_path = os.path.join(VIZ_DIR, "raw_mesh.obj") print(f"\n{'='*70}") print("Exporting raw .obj (no renderer, no nvdiffrast) ...") print(f"{'='*70}") export_obj(mesh[0], obj_path) print(f" -> Open {obj_path} in Blender to check geometry completeness.") print(f" If this is complete but renders are 15-30%, the bug is in nvdiffrast/rasterizer.\n") # ------------------------------------------------------------------ # SMOKING GUN: render with render_snapshot — identical to app.py # ------------------------------------------------------------------ print(f"\n{'='*70}") print(f"Rendering {NVIEWS} views at {RENDER_RES}x{RENDER_RES} ...") print(f"{'='*70}\n") render_out_dir = os.path.join(VIZ_DIR, "render_frames") os.makedirs(render_out_dir, exist_ok=True) images = render_utils.render_snapshot( mesh[0], # same single-mesh call as app.py resolution=RENDER_RES, r=2, fov=36, nviews=NVIEWS, envmap=envmap, ) print(f"\nRender keys returned: {list(images.keys())}") for key, frames in images.items(): print(f" {key}: {len(frames)} frames, each {frames[0].shape}") # Save every individual frame print(f"\nSaving individual frames to {render_out_dir}/ ...") saved = save_render_frames(images, render_out_dir, prefix="render") print(f" Saved {len(saved)} frame files.") # Save contact sheets (one image per render key, all views side by side) print(f"\nBuilding contact sheets ...") make_contact_sheet(images, render_out_dir, prefix="contact") print(f"\n{'='*70}") print(f"All done. Check {render_out_dir}/ for render output.") print(f" contact_shaded_all_views.png <- the key one to look at") print(f" contact_base_color_all_views.png <- color without lighting") print(f" contact_normal_all_views.png <- surface normals") print(f"{'='*70}\n") if __name__ == "__main__": main()