bd0f3f1941
The pixel-aligned conditioning is the ONLY thing separating this port from the
trellis2_mlx operator already in MODELBEAST — upstream's main branch is the
TRELLIS.2 backbone, so everything else here is TRELLIS.2 with a different head.
This lands that head.
proj.py ProjGrid, project_points, bilinear_sample, distance_from_fov — MLX
dino.py DINOv3 ViT-L/16 left in torch on MPS (run once per image, outside the
25-step loop; transformers gives exact parity for free)
cond.py encode_image_proj equivalent -> {'global','proj'} + zero uncond
The extractor has no sparse conv, so upstream RUNS on CPU torch here and is a real
oracle. All 12 checks diff against it, not against a transcription:
bilinear_sample vs grid_sample max diff 2.4e-07 corr 1.00000000
project_points pixels/depth/mask exact
ProjGrid forward (ss, 16^3) max diff 1.9e-05 corr 1.00000000
extractor global tokens max diff 0.0e+00 corr 1.00000000
extractor proj features max diff 4.8e-06 corr 1.00000000
Three details that a plain transcription gets wrong and eyeballing cannot catch:
grid_sample's align_corners=False maps a normalised coord to ((c+1)*size-1)/2, not
(c+1)/2*(size-1) — half a texel, invisible until you compare; padding_mode='border'
clamps the SOURCE INDEX before corners are taken, not the corners after, which
changes the weights on every silhouette edge (tested with deliberately out-of-range
grid coords); and the camera looks down -Z, so a sign slip still yields a plausible
grid that samples the mirror image.
Also corrects a shape assumption from the earlier smoke test: 'global' is CLS + 4
register tokens = [B,5,1024], NOT the 1370 image tokens. The patch tokens go to the
proj branch. That asymmetry IS the architecture.
Note the parameterless final layer_norm in extract_features — not model.norm, which
has weights. Same trap as the ss_flow bug: no checkpoint trace, 200x output error.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
69 lines
2.9 KiB
Python
69 lines
2.9 KiB
Python
"""Image -> conditioning dicts, the MLX side of upstream's `encode_image_proj`.
|
|
|
|
Produces exactly the two-key contract the DiT blocks expect:
|
|
|
|
cond = {'global': [B, 1+regs, D], 'proj': [B, R^3, C]}
|
|
uncond = zeros of the same shapes (classifier-free guidance negative)
|
|
|
|
`global` is cross-attended; `proj` is projected and ADDED, one token per voxel. The
|
|
token counts differ (5 vs R^3) and that asymmetry IS the architecture — see proj.py.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict, Tuple
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from .dino import DinoV3Encoder
|
|
from .proj import ProjGrid, distance_from_fov
|
|
|
|
# The four extractor configurations upstream builds in inference.py. Only `ss` is wired
|
|
# so far; the other three additionally need the NAF upsampler, which doubles their
|
|
# proj_channels to embed_dim*2 by concatenating a high-res branch.
|
|
CONFIGS = {
|
|
"ss": dict(image_size=512, grid_resolution=16, use_naf_upsample=False),
|
|
"shape_512": dict(image_size=512, grid_resolution=32, use_naf_upsample=True, naf_target_size=512),
|
|
"shape_1024": dict(image_size=1024, grid_resolution=64, use_naf_upsample=True, naf_target_size=512),
|
|
"tex_1024": dict(image_size=1024, grid_resolution=64, use_naf_upsample=True, naf_target_size=1024),
|
|
}
|
|
|
|
|
|
def load_image(path: str, image_size: int) -> np.ndarray:
|
|
"""Path -> [1,3,S,S] float32 in [0,1], LANCZOS-resized like upstream."""
|
|
img = Image.open(path).convert("RGB").resize((image_size, image_size), Image.LANCZOS)
|
|
arr = np.asarray(img, dtype=np.float32) / 255.0
|
|
return arr.transpose(2, 0, 1)[None]
|
|
|
|
|
|
class ProjConditioner:
|
|
"""One extractor stage: DINOv3 + a ProjGrid at that stage's voxel resolution."""
|
|
|
|
def __init__(self, stage: str = "ss", device: str | None = None):
|
|
cfg = CONFIGS[stage]
|
|
if cfg.get("use_naf_upsample"):
|
|
raise NotImplementedError(
|
|
f"stage '{stage}' needs the NAF upsampler (proj_channels = embed_dim*2); "
|
|
"only 'ss' is wired so far"
|
|
)
|
|
self.stage = stage
|
|
self.image_size = cfg["image_size"]
|
|
self.encoder = DinoV3Encoder(image_size=self.image_size, device=device)
|
|
self.proj_grid = ProjGrid(grid_resolution=cfg["grid_resolution"],
|
|
image_resolution=self.image_size)
|
|
|
|
def __call__(self, image_chw01: np.ndarray, camera_angle_x: float,
|
|
distance: float | None = None, mesh_scale: float = 1.0
|
|
) -> Tuple[Dict[str, mx.array], Dict[str, mx.array]]:
|
|
if distance is None:
|
|
distance = distance_from_fov(camera_angle_x, mesh_scale, self.image_size)
|
|
|
|
z_global, z_patch = self.encoder(image_chw01)
|
|
z_proj = self.proj_grid(z_patch, camera_angle_x, distance, mesh_scale)
|
|
|
|
cond = {"global": z_global, "proj": z_proj}
|
|
uncond = {"global": mx.zeros_like(z_global), "proj": mx.zeros_like(z_proj)}
|
|
return cond, uncond
|