"""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 preprocess_image(img: Image.Image, bg_color=(0, 0, 0)) -> Image.Image: """Upstream's `preprocess_image`, alpha path only. Crops to the subject's bounding box with 1.1x headroom and composites onto a flat background. This is not cosmetic: the camera solve assumes the object fills the frame (`distance_from_fov` places the camera so a unit mesh exactly spans it), so an uncropped image silently mis-scales the whole reconstruction. Background REMOVAL is not implemented — upstream calls a rembg model for images with no usable alpha. Such an image is passed through unchanged here, which will reconstruct the background along with the subject. Feed RGBA with a real matte. """ has_alpha = False if img.mode == "RGBA": alpha = np.array(img)[:, :, 3] has_alpha = not np.all(alpha == 255) scale = min(1, 1024 / max(img.size)) if scale < 1: img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS) if not has_alpha: return img.convert("RGB") arr = np.array(img) ys, xs = np.nonzero(arr[:, :, 3] > 0.8 * 255) cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2 size = int(max(xs.max() - xs.min(), ys.max() - ys.min()) * 1.1) img = img.crop((cx - size // 2, cy - size // 2, cx + size // 2, cy + size // 2)) out = np.asarray(img, dtype=np.float32) / 255.0 rgb, a = out[:, :, :3], out[:, :, 3:4] bg = np.array(bg_color, dtype=np.float32) / 255.0 return Image.fromarray((np.clip(rgb * a + bg * (1 - a), 0, 1) * 255).astype(np.uint8)) def load_image(path: str, image_size: int, preprocess: bool = True) -> np.ndarray: """Path -> [1,3,S,S] float32 in [0,1], preprocessed and LANCZOS-resized.""" img = Image.open(path) if preprocess: img = preprocess_image(img) img = img.convert("RGB").resize((image_size, image_size), Image.LANCZOS) return (np.asarray(img, dtype=np.float32) / 255.0).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] self.stage = stage self.image_size = cfg["image_size"] self.grid_resolution = cfg["grid_resolution"] self.use_naf = bool(cfg.get("use_naf_upsample")) self.naf_target = cfg.get("naf_target_size") self.device = device self.encoder = DinoV3Encoder(image_size=self.image_size, device=device) # ONE grid for both branches, at image_size — upstream reuses `self.proj_grid` # for the high-res sample too. That is not interchangeable with a grid built at # naf_target_size: the normalised coordinate carries a 1/resolution term, so # projecting at 1024 and normalising by 1024 lands ~0.001 off in [-1,1] — a # sub-pixel shift on every voxel, in a model whose whole point is pixel # alignment. Sampling is resolution-agnostic anyway; grid_sample maps [-1,1] # onto whatever feature map it is handed. self.proj_grid = ProjGrid(grid_resolution=self.grid_resolution, image_resolution=self.image_size) self.proj_channels = self.encoder.embed_dim * (2 if self.use_naf else 1) def _grid(self, grid_resolution: int | None): """ProjGrid at `grid_resolution`, cached. Always at image_size (see __init__).""" r = grid_resolution or self.grid_resolution if r == self.grid_resolution: return self.proj_grid cache = getattr(self, "_grid_cache", None) or {} if r not in cache: cache[r] = ProjGrid(grid_resolution=r, image_resolution=self.image_size) self._grid_cache = cache return cache[r] def __call__(self, image_chw01: np.ndarray, camera_angle_x: float, distance: float | None = None, mesh_scale: float = 1.0, grid_resolution: int | None = None, ) -> Tuple[Dict[str, mx.array], Dict[str, mx.array]]: """`grid_resolution` overrides this stage's lattice — the cascade backs the high-res grid off in steps when an object is too dense for max_num_tokens.""" if distance is None: distance = distance_from_fov(camera_angle_x, mesh_scale, self.image_size) grid = self._grid(grid_resolution) z_global, z_patch = self.encoder(image_chw01) z_proj = grid(z_patch, camera_angle_x, distance, mesh_scale) if self.use_naf: from .naf import upsample import numpy as _np from PIL import Image as _Image # NAF's guide is the UNNORMALISED image at the target size guide = image_chw01 if guide.shape[-1] != self.naf_target: arr = (guide[0].transpose(1, 2, 0) * 255).astype(_np.uint8) arr = _Image.fromarray(arr).resize((self.naf_target, self.naf_target), _Image.LANCZOS) guide = (_np.asarray(arr, dtype=_np.float32) / 255.0).transpose(2, 0, 1)[None] hr = upsample(z_patch, guide, self.naf_target, device=self.device or "mps") z_hr = grid(hr, camera_angle_x, distance, mesh_scale) z_proj = mx.concatenate([z_proj, z_hr], axis=-1) cond = {"global": z_global, "proj": z_proj} uncond = {"global": mx.zeros_like(z_global), "proj": mx.zeros_like(z_proj)} return cond, uncond