6f40fae0cc
Completes the conditioning. shape_512 / shape_1024 / tex_1024 run a second HIGH-RES
branch — NAF upsamples the DINOv3 patch map to 512/1024 guided by the RGB image, the
proj grid samples that too, and the branches concatenate. That is why those stages
have proj_channels = embed_dim*2 (2048).
CORRECTS AN EARLIER CLAIM: I said natten is never imported and can be skipped. True of
Pixal3D's own source — but natten is a dependency of NAF (valeoai/NAF), which arrives
at RUNTIME via torch.hub and is not vendored. That is what README Step 3 is for. The
warning was real; the reason was one level down.
natten is a dead end on Apple Silicon regardless:
cutlass-fna requires libnatten, which the arm64 build does not produce
flex-fna CPU only ('not on a CUDA, ROCm, or CPU device: mps'), AND refuses
different head dims for QK vs V — which is exactly NAF's shape
(qk=64, v=256). Worked around, CPU took 243s at 256px and was
OOM-KILLED (exit 137) at the 512 the pipeline actually needs.
REPLACED BY AN EXACT REDUCTION, not an approximation. NAF resizes K/V from the 32x32
patch map with nearest-exact and then dilates by exactly the upsample factor, so the
dilated high-res neighborhood samples one position per low-res cell and collapses to a
plain clamped 9x9 neighborhood on the 32x32 grid, shared by every high-res pixel in
that cell. Verified against natten's own kernel at three dilations:
LR 16 -> HR 64 (dil 4) max diff 7.153e-07
LR 16 -> HR 128 (dil 8) max diff 7.153e-07
LR 32 -> HR 256 (dil 8) max diff 7.153e-07 (float32 epsilon)
Grouping queries by low-res cell also avoids materialising the high-res neighborhood,
which would be ~87GB of gathered V at 512. Result, on MPS:
32 -> 512 2.4s (natten: OOM-killed)
64 -> 512 0.2s
64 -> 1024 0.9s
Also fixes a caching bug that stranded NAF on whichever device loaded first, and one
in my own wiring: the high-res branch must reuse the SAME ProjGrid at image_size, not
a new one at naf_target_size. The normalised coordinate carries a 1/resolution term,
so the latter lands ~0.001 off in [-1,1] — a sub-pixel shift on every voxel, in the
one model whose entire premise is pixel alignment.
27/27 green (15 proj incl. the NAF stage vs upstream at corr 1.00000000, 5 sampler,
5 naf vs natten, 2 decoders).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
131 lines
6.0 KiB
Python
131 lines
6.0 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 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 __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)
|
|
|
|
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 = self.proj_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
|