cee423501c
Closes the op gap for Pixal3D. downsample now takes mode='max'|'mean'. This differs BETWEEN MODELS and is silent when wrong: LATO.2 hardcodes amax, Pixal3D parameterises it and defaults to mean. Callers must be explicit. upsample is a MASKED expansion, distinct from subdivide: subdivision.feats is an [N, factor**3] 0/1 mask, so a voxel emits 0-8 children rather than always 8. Pixal3D's decoder uses it to grow the occupied set selectively. Child-index decoding is subidx // factor**i % factor - axis 0 fastest-varying, the REVERSE of subdivide's C-order, which is exactly the kind of mixup that misplaces every child silently. Tests cover both modes differing, masked expansion landing at the right offsets, and a voxel with no flagged children contributing nothing.
39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
"""MLX implementation of the TRELLIS-lineage sparse module.
|
|
|
|
Microsoft's TRELLIS.2 sparse stack has been inherited, near-verbatim, by a growing
|
|
family of 3D generation models — LATO.2 and TencentARC's Pixal3D among them. All of
|
|
them hard-require spconv or torchsparse, neither of which has a Metal build, and that
|
|
single dependency is what keeps the whole lineage off Apple Silicon.
|
|
|
|
The blocker is one operation: submanifold 3x3x3 convolution. Every SparseConv3d in
|
|
these models is constructed `stride=1, padding=None`, which spconv dispatches to
|
|
SubMConv3d. Implement that in MLX and the rest is ordinary linear/norm/attention work.
|
|
|
|
Packaged separately from any one model so each port depends on a tested core rather
|
|
than vendoring its own copy.
|
|
"""
|
|
|
|
from .conv import SubMConv3d, build_indice_map
|
|
from .tensor import SparseTensor, VarLenTensor, downsample, subdivide, upsample
|
|
from .ops import (
|
|
LayerNorm32,
|
|
SparseFeedForwardNet,
|
|
SparseGELU,
|
|
SparseGroupNorm32,
|
|
SparseLinear,
|
|
SparseMultiHeadAttention,
|
|
SparseResBlock,
|
|
SparseSiLU,
|
|
SparseTransformerBlock,
|
|
SparseTransformerCrossBlock,
|
|
)
|
|
|
|
__all__ = [
|
|
"SparseTensor", "VarLenTensor", "subdivide", "downsample", "upsample",
|
|
"SubMConv3d", "build_indice_map",
|
|
"SparseLinear", "LayerNorm32", "SparseGroupNorm32", "SparseSiLU", "SparseGELU",
|
|
"SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention",
|
|
"SparseTransformerBlock", "SparseTransformerCrossBlock",
|
|
]
|
|
__version__ = "0.1.0"
|