"""View-aligned ("proj") conditioning — the one thing that makes Pixal3D not TRELLIS.2. Upstream's `main` branch IS the TRELLIS.2 backbone; the entire difference is that pixel features are BACK-PROJECTED into the voxel grid through the camera instead of being cross-attended. So everything in this file is the actual novelty of the model. The chain per image, once, outside the denoising loop: voxel grid [-1,1]^3 -> Blender-aligned -> camera space -> pixel coords -> bilinear sample of the DINOv3 patch map -> [B, R^3, C] Two conventions here are load-bearing and are the easy things to get silently wrong: * The camera looks down **-Z** (Blender), so depth is `-z_cam` and the perspective divide uses `-z_cam`. Getting the sign wrong still produces a plausible-looking grid — it just samples the mirror image. * `F.grid_sample(align_corners=False, padding_mode='border')` has exact semantics that a naive bilinear lerp does NOT reproduce. See `bilinear_sample`. Everything is diffed against upstream in tests/test_proj.py — the extractor has no sparse conv, so upstream runs on CPU torch here and is a real oracle. Do not reason about correctness in this file; measure it. """ from __future__ import annotations import mlx.core as mx import numpy as np # Blender's sensor model, straight from upstream's project_points_to_image_batch. SENSOR_WIDTH_MM = 32.0 FOCAL_NUMERATOR = 16.0 # == SENSOR_WIDTH_MM / 2, i.e. a half-width of sensor # Rotates the [-1,1]^3 sampling grid into Blender's axis convention (Y up -> Z up). _BLENDER_ROT = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], dtype=np.float32) # Camera sits on -Y looking at the origin; row 1 col 3 is overwritten with -distance. _FRONT_VIEW = np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 0.0, -1.0, -2.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]], dtype=np.float32) def compute_f_pixels(camera_angle_x: float, resolution: int) -> float: """Horizontal FOV (radians) -> focal length in pixels.""" focal_length = FOCAL_NUMERATOR / np.tan(camera_angle_x / 2.0) return float(focal_length * resolution / SENSOR_WIDTH_MM) def distance_from_fov(camera_angle_x: float, mesh_scale: float = 1.0, image_resolution: int = 512, extend_pixel: int = 0) -> float: """Camera distance that makes a unit mesh exactly fill the frame at this FOV. Upstream solves this by projecting the grid corner (-1, 0, 0) and demanding it land on the left image edge. Reproduced here with upstream's fixed grid/target points rather than generalised, because those are the only values inference ever passes. """ gp = _BLENDER_ROT.T @ np.array([-1.0, 0.0, 0.0], dtype=np.float32) gp = gp / mesh_scale / 2.0 xw, yw = float(gp[0]), float(gp[1]) x_target = 0.0 - extend_pixel f_pixels = compute_f_pixels(camera_angle_x, image_resolution) x_ndc = x_target - image_resolution / 2.0 return f_pixels * xw / x_ndc - yw def project_points(points_3d: mx.array, transform_matrix: mx.array, camera_angle_x: float, resolution: int): """[B,N,3] world points -> ([B,N,2] pixel coords, [B,N] depth, [B,N] valid mask). Pixel coords are in [0, resolution); the mask marks points inside the frame AND in front of the camera. The mask is returned for completeness — upstream computes it and then samples with border padding regardless, so it is not applied. """ b, n, _ = points_3d.shape ones = mx.ones((b, n, 1), dtype=points_3d.dtype) points_h = mx.concatenate([points_3d, ones], axis=-1) # [B,N,4] # A 4x4 inverse, once per image — numpy is exact and dodges MLX linalg stream rules. w2c = mx.array(np.linalg.inv(np.asarray(transform_matrix, dtype=np.float64)).astype(np.float32)) points_cam = (points_h @ mx.swapaxes(w2c, -2, -1))[..., :3] x_cam, y_cam, z_cam = points_cam[..., 0], points_cam[..., 1], points_cam[..., 2] depth = -z_cam # Blender cameras look down -Z f_px = compute_f_pixels(camera_angle_x, resolution) denom = -z_cam + 1e-8 x_pixel = f_px * x_cam / denom + resolution / 2.0 y_pixel = -(f_px * y_cam / denom) + resolution / 2.0 # image Y grows downward valid = ((x_pixel >= 0) & (x_pixel < resolution) & (y_pixel >= 0) & (y_pixel < resolution) & (depth > 0)) return mx.stack([x_pixel, y_pixel], axis=-1), depth, valid def bilinear_sample(fmap: mx.array, grid: mx.array) -> mx.array: """`F.grid_sample(align_corners=False, padding_mode='border', mode='bilinear')`. fmap [B,H,W,C] (channels-last — MLX-natural, and it lets us skip upstream's two permutes since the DINOv3 patch map already arrives BHWC), grid [B,K,2] in [-1,1] as (x, y). Returns [B,K,C]. The two details that make this NOT a plain lerp, both from ATen's grid_sampler_compute_source_index: * align_corners=False maps normalised c to `((c + 1) * size - 1) / 2`, which puts -1 at the OUTER EDGE of the first texel rather than at its centre. Using `(c + 1) / 2 * (size - 1)` instead — the align_corners=True formula — is the classic off-by-half-a-texel bug and it survives eyeballing. * padding_mode='border' clamps the SOURCE INDEX before the corners are taken, not the corners afterwards. Clamping afterwards changes the interpolation weights for out-of-frame points instead of flattening them to the edge texel. """ b, h, w, c = fmap.shape k = grid.shape[1] gx, gy = grid[..., 0], grid[..., 1] ix = ((gx + 1.0) * w - 1.0) / 2.0 iy = ((gy + 1.0) * h - 1.0) / 2.0 ix = mx.clip(ix, 0, w - 1) # border: clamp first, then take corners iy = mx.clip(iy, 0, h - 1) x0 = mx.floor(ix) y0 = mx.floor(iy) wx = ix - x0 # weights from the UNclamped-corner fraction wy = iy - y0 x0i = mx.clip(x0.astype(mx.int32), 0, w - 1) x1i = mx.clip(x0i + 1, 0, w - 1) y0i = mx.clip(y0.astype(mx.int32), 0, h - 1) y1i = mx.clip(y0i + 1, 0, h - 1) flat = fmap.reshape(b, h * w, c) def gather(yi, xi): idx = (yi * w + xi) # [B,K] idx = mx.broadcast_to(idx[..., None], (b, k, c)) # take_along_axis wants rank match return mx.take_along_axis(flat, idx, axis=1) nw, ne = gather(y0i, x0i), gather(y0i, x1i) sw, se = gather(y1i, x0i), gather(y1i, x1i) wx, wy = wx[..., None], wy[..., None] top = nw + (ne - nw) * wx bot = sw + (se - sw) * wx return top + (bot - top) * wy class ProjGrid: """The [-1,1]^3 sampling lattice and its projection into image space. Stateless apart from the lattice itself, which depends only on `grid_resolution`. """ def __init__(self, grid_resolution: int = 16, image_resolution: int = 512): self.grid_resolution = grid_resolution self.image_resolution = image_resolution one = np.linspace(-1.0, 1.0, grid_resolution, dtype=np.float32) gx, gy, gz = np.meshgrid(one, one, one, indexing="ij") pts = np.stack((gx, gy, gz), axis=-1) @ _BLENDER_ROT.T self.grid_points = mx.array(pts.reshape(-1, 3)) # [R^3, 3] def __call__(self, fmap: mx.array, camera_angle_x: float, distance: float, mesh_scale: float = 1.0) -> mx.array: """fmap [B,H,W,C] -> proj features [B, R^3, C]. `transform_matrix` is deliberately not a parameter: upstream asserts it is None on every inference path and builds the front view from `distance`, so accepting one here would be dead surface that silently diverges from the checkpoint. """ b = fmap.shape[0] pts = mx.broadcast_to(self.grid_points, (b, *self.grid_points.shape)) pts = pts / mesh_scale / 2.0 tm = _FRONT_VIEW.copy() tm[1, 3] = -distance tm = mx.array(np.broadcast_to(tm, (b, 4, 4)).copy()) pixels, _, _ = project_points(pts, tm, camera_angle_x, self.image_resolution) grid = (pixels + 0.5) / self.image_resolution * 2.0 - 1.0 return bilinear_sample(fmap, grid)