79ef81988c
image_to_occupancy() runs the structure stage on an actual photo: preprocess ->
DINOv3 -> proj back-projection -> ss_flow -> ss_dec -> 64^3 occupancy.
VERIFICATION THAT MATTERS: scripts/run_structure.py re-projects the occupied voxels
through the same camera and compares against the input alpha matte. On the upstream
sample that is silhouette IoU 0.842 with 12948 voxels occupied (4.94% of 64^3). This
is the model's own headline claim, so it is the right thing to assert — 'it ran
without crashing' would pass just as happily on a generic blob.
Two real bugs this phase found, neither visible without reading the shipped configs:
1. THE SAMPLER WAS MISSING guidance_rescale. The checkpoint's own pipeline.json sets
0.7 for the structure stage and 0.5 for shape_slat, so this fires at the model's
DEFAULT settings — omitting it silently overcooks every structure prediction. Now
implemented (Lin et al. CFG rescale) and diffed against upstream's
ClassifierFreeGuidanceSamplerMixin, run directly rather than reimplemented.
2. The sampler defaults were wrong: the real ss stage is steps=12 / rescale_t=5.0 /
guidance 7.5 / interval [0.6,1.0], not the steps=25 / rescale_t=3.0 the smoke test
assumed. All three stages' real params now live in pipeline.py, read from
pipeline.json rather than guessed.
TIMINGS, measured with interleaved reps after warmup (the first pass attributed the
same 11s of residual warmup to both 'rescale' and 'torch contention'; it was neither):
cold run 89.3s
warm, full settings 16.5s
warm, CFG off 9.2s -> CFG costs 1.80x, as expected for 10/12
steps falling inside the guidance interval
guidance_rescale ~0s -> free
torch/MPS contention ~0s -> DINOv3 can stay resident
peak memory 6.8GB
THE FINDING THAT SHAPES THE OPERATOR: warmup is ~71s against ~17s of actual compute,
i.e. 4x the work. A MODELBEAST operator MUST hold the models resident across jobs
rather than fork per job — the trellis2 lane shows the same shape (47.9s cold vs 2.5s
warm pipeline_load). Cost this in before optimising any kernel.
17/17 tests green (12 proj + 5 sampler).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
106 lines
4.6 KiB
Python
106 lines
4.6 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]
|
|
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
|