18179bc479
All seven Pixal3D models now load and run. shape_dec 292/292, tex_dec 284/284, both with zero missing/unmapped/mismatched keys - the 8-param difference between them is exactly the four to_subdiv layers, since tex_dec has pred_subdiv=False. Behaviour is right: 12 latent voxels grow SELECTIVELY through four stages (12 -> 91 -> 193 -> 358 -> 1150) rather than x8 each time, which would have reached 49,152. Scale lands at 1/16. shape_dec emits the hardcoded 7 channels and its vertex head produces offsets inside the [-0.5, 1.5] band its sigmoid+voxel_margin allows. tex_dec, guided by shape_dec's masks, reproduces exactly the same voxel count. VERIFICATION CAVEAT, recorded in the README: these two are the only models using sparse conv, so upstream cannot run here and there is no numerical oracle. Unlike the five models verified at correlation 1.0, these are checked structurally and behaviourally only. Weaker evidence, and labelled as such rather than presented alongside the verified results.
245 lines
8.5 KiB
Python
245 lines
8.5 KiB
Python
"""Pixal3D sparse VAE decoders in MLX — `shape_dec` and `tex_dec`.
|
|
|
|
Both are the same U-Net: stages of `SparseConvNeXtBlock3d`, each stage ending in a
|
|
`SparseResBlockC2S3d` that upsamples via channel->spatial. The occupied set grows
|
|
selectively at every stage, driven by a predicted subdivision mask, which is how the
|
|
decoder reaches 256^3-equivalent detail without ever materialising a dense grid.
|
|
|
|
model_channels [1024, 512, 256, 128, 64]
|
|
num_blocks [4, 16, 8, 4, 0] <- last stage is empty and has no up-block
|
|
block_type SparseConvNeXtBlock3d
|
|
up_block_type SparseResBlockC2S3d
|
|
|
|
`FlexiDualGridVaeDecoder` (shape_dec) subclasses it and hardcodes **7** output channels
|
|
regardless of config — 3 for vertex offsets (squashed through a sigmoid widened by
|
|
`voxel_margin`) plus 4 more — so its `out_channels` is not a config field at all.
|
|
`SparseUnetVaeDecoder` (tex_dec) uses the config's `out_channels` (6) and sets
|
|
`pred_subdiv=False`, meaning it needs subdivision masks supplied from outside rather
|
|
than predicting its own.
|
|
|
|
NOTE ON VERIFICATION: unlike the flow models and ss_dec, these use sparse convolution,
|
|
so upstream cannot run here (no spconv on Metal) and there is no numerical oracle. They
|
|
are built from blocks that ARE individually tested, but the assembled graphs are checked
|
|
only for shape, finiteness and plausible behaviour. Treat them as less certain than
|
|
anything reporting a correlation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
from trellis_sparse_mlx import (
|
|
SparseConvNeXtBlock3d,
|
|
SparseLinear,
|
|
SparseResBlockC2S3d,
|
|
SparseTensor,
|
|
)
|
|
|
|
_BLOCKS = {
|
|
"SparseConvNeXtBlock3d": SparseConvNeXtBlock3d,
|
|
"SparseResBlockC2S3d": SparseResBlockC2S3d,
|
|
}
|
|
|
|
|
|
class SparseUnetVaeDecoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
out_channels: int,
|
|
model_channels: List[int],
|
|
latent_channels: int,
|
|
num_blocks: List[int],
|
|
block_type: List[str],
|
|
up_block_type: List[str],
|
|
block_args: Optional[List[Dict[str, Any]]] = None,
|
|
pred_subdiv: bool = True,
|
|
**_ignored,
|
|
):
|
|
super().__init__()
|
|
self.out_channels = out_channels
|
|
self.model_channels = list(model_channels)
|
|
self.num_blocks = list(num_blocks)
|
|
self.pred_subdiv = pred_subdiv
|
|
block_args = block_args or [{} for _ in num_blocks]
|
|
|
|
self.from_latent = SparseLinear(latent_channels, model_channels[0])
|
|
self.output_layer = SparseLinear(model_channels[-1], out_channels)
|
|
|
|
stages = []
|
|
for i in range(len(num_blocks)):
|
|
stage = [
|
|
_BLOCKS[block_type[i]](model_channels[i], **block_args[i])
|
|
for _ in range(num_blocks[i])
|
|
]
|
|
if i < len(num_blocks) - 1:
|
|
stage.append(
|
|
_BLOCKS[up_block_type[i]](
|
|
model_channels[i],
|
|
model_channels[i + 1],
|
|
pred_subdiv=pred_subdiv,
|
|
)
|
|
)
|
|
stages.append(stage)
|
|
self.blocks = stages
|
|
|
|
def __call__(
|
|
self,
|
|
x: SparseTensor,
|
|
guide_subs: Optional[List[SparseTensor]] = None,
|
|
return_subs: bool = False,
|
|
):
|
|
h = self.from_latent(x)
|
|
subs = []
|
|
for i, stage in enumerate(self.blocks):
|
|
for j, blk in enumerate(stage):
|
|
is_up = i < len(self.blocks) - 1 and j == len(stage) - 1
|
|
if is_up:
|
|
if self.pred_subdiv:
|
|
h, sub = blk(h)
|
|
subs.append(sub)
|
|
else:
|
|
h = blk(h, subdiv=guide_subs[i] if guide_subs else None)
|
|
else:
|
|
h = blk(h)
|
|
|
|
# Same parameterless final LayerNorm as every other model in this family.
|
|
f = h.feats
|
|
f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(
|
|
mx.var(f, -1, keepdims=True) + 1e-5
|
|
)
|
|
out = self.output_layer(h.replace(f))
|
|
return (out, subs) if return_subs else out
|
|
|
|
|
|
class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
|
|
"""shape_dec. Emits 7 channels; the first 3 are vertex offsets within the voxel."""
|
|
|
|
def __init__(
|
|
self,
|
|
resolution: int,
|
|
model_channels: List[int],
|
|
latent_channels: int,
|
|
num_blocks: List[int],
|
|
block_type: List[str],
|
|
up_block_type: List[str],
|
|
block_args: Optional[List[Dict[str, Any]]] = None,
|
|
voxel_margin: float = 0.5,
|
|
**_ignored,
|
|
):
|
|
# out_channels is hardcoded to 7 upstream, not read from config
|
|
super().__init__(
|
|
7,
|
|
model_channels,
|
|
latent_channels,
|
|
num_blocks,
|
|
block_type,
|
|
up_block_type,
|
|
block_args,
|
|
pred_subdiv=True,
|
|
)
|
|
self.resolution = resolution
|
|
self.voxel_margin = voxel_margin
|
|
|
|
def decode_vertices(self, h: SparseTensor) -> mx.array:
|
|
"""Channels 0:3 -> vertex offsets, sigmoid widened by `voxel_margin`.
|
|
|
|
`(1 + 2m)*sigmoid(v) - m` maps to [-m, 1+m], letting a vertex sit slightly
|
|
outside its own voxel — that overshoot is what lets the dual grid represent
|
|
surfaces that do not pass exactly through cell centres.
|
|
"""
|
|
m = self.voxel_margin
|
|
return (1 + 2 * m) * mx.sigmoid(h.feats[..., 0:3]) - m
|
|
|
|
|
|
# ------------------------------------------------------------------ loading
|
|
|
|
def _remap(k: str) -> str:
|
|
"""Upstream: blocks.<stage>.<idx>.<...>; SparseLinear wraps an nn.Linear."""
|
|
for name in ("from_latent", "output_layer", "to_subdiv"):
|
|
k = k.replace(f"{name}.weight", f"{name}.linear.weight")
|
|
k = k.replace(f"{name}.bias", f"{name}.linear.bias")
|
|
k = k.replace(".mlp.0.", ".mlp_0.").replace(".mlp.2.", ".mlp_2.")
|
|
k = k.replace(".conv1.conv.", ".conv1.").replace(".conv2.conv.", ".conv2.")
|
|
k = k.replace(".conv.conv.", ".conv.")
|
|
k = k.replace(".norm.weight", ".norm.weight").replace(".norm1.weight", ".norm1.weight")
|
|
return k
|
|
|
|
|
|
def load(weights_path: str | Path, config_path: str | Path):
|
|
cfg = json.loads(Path(config_path).read_text())
|
|
name, args = cfg["name"], dict(cfg.get("args", {}))
|
|
args.pop("use_fp16", None)
|
|
cls = {
|
|
"FlexiDualGridVaeDecoder": FlexiDualGridVaeDecoder,
|
|
"SparseUnetVaeDecoder": SparseUnetVaeDecoder,
|
|
}[name]
|
|
model = cls(**args)
|
|
|
|
w = mx.load(str(weights_path))
|
|
flat = dict(_flatten(model.parameters()))
|
|
mapped, unmapped, mismatched = {}, [], []
|
|
for k, v in w.items():
|
|
m = _remap(k)
|
|
if m in flat:
|
|
if flat[m].shape != v.shape:
|
|
mismatched.append(f"{k} -> {m}: want {flat[m].shape} got {v.shape}")
|
|
continue
|
|
mapped[m] = v
|
|
else:
|
|
unmapped.append(f"{k} -> {m} {tuple(v.shape)}")
|
|
missing = [k for k in flat if k not in mapped]
|
|
if mapped:
|
|
model.update(_unflatten(mapped))
|
|
return model, {
|
|
"config": name,
|
|
"loaded": len(mapped),
|
|
"params": len(flat),
|
|
"missing": missing,
|
|
"unmapped": unmapped,
|
|
"mismatched": mismatched,
|
|
}
|
|
|
|
|
|
def _flatten(tree, prefix=""):
|
|
if isinstance(tree, dict):
|
|
for k, v in tree.items():
|
|
yield from _flatten(v, f"{prefix}{k}.")
|
|
elif isinstance(tree, list):
|
|
for i, v in enumerate(tree):
|
|
yield from _flatten(v, f"{prefix}{i}.")
|
|
elif isinstance(tree, mx.array):
|
|
yield prefix[:-1], tree
|
|
|
|
|
|
def _unflatten(flat: dict):
|
|
root: dict = {}
|
|
for key, val in flat.items():
|
|
parts = key.split(".")
|
|
node = root
|
|
for i, p in enumerate(parts[:-1]):
|
|
nxt = parts[i + 1]
|
|
default = [] if nxt.isdigit() else {}
|
|
if isinstance(node, list):
|
|
idx = int(p)
|
|
while len(node) <= idx:
|
|
node.append({})
|
|
if isinstance(default, list) and not isinstance(node[idx], list):
|
|
node[idx] = default
|
|
node = node[idx]
|
|
else:
|
|
if p not in node or not isinstance(node[p], (dict, list)):
|
|
node[p] = default
|
|
node = node[p]
|
|
if isinstance(node, list):
|
|
idx = int(parts[-1])
|
|
while len(node) <= idx:
|
|
node.append(None)
|
|
node[idx] = val
|
|
else:
|
|
node[parts[-1]] = val
|
|
return root
|