5dd3f11d92
Completes the sparse op surface Pixal3D's decoders need. spatial2channel/channel2spatial are sparse space-to-depth and its inverse. The slot index is sum_i (coord[i] % f) * f**i - axis 0 FASTEST varying, the reverse of C order. decodes it the same way; using C order in one and this in the other would misplace every child while keeping all shapes valid. SparseResBlockC2S3d widens to out_channels*8 so channel2spatial can redistribute those channels across the 8 children, with a predicted subdiv mask deciding which children exist - so the occupied set grows selectively (32 voxels -> 122, not 256). Tests cover exact round-trip, zero-fill of unoccupied slots, and selective growth.
68 lines
2.2 KiB
Python
68 lines
2.2 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 .dit import (
|
|
ProjectAttention,
|
|
DiTAttention,
|
|
ModulatedTransformerCrossBlock,
|
|
MultiHeadRMSNorm,
|
|
TimestepEmbedder,
|
|
apply_rope,
|
|
rope_phases_from_coords,
|
|
SparseDiTAttention,
|
|
SparseProjectAttention,
|
|
ModulatedSparseTransformerCrossBlock,
|
|
)
|
|
from .tensor import (
|
|
SparseTensor,
|
|
VarLenTensor,
|
|
channel2spatial,
|
|
downsample,
|
|
spatial2channel,
|
|
subdivide,
|
|
upsample,
|
|
)
|
|
from .ops import (
|
|
LayerNorm32,
|
|
SparseFeedForwardNet,
|
|
SparseGELU,
|
|
SparseGroupNorm32,
|
|
SparseLinear,
|
|
SparseMultiHeadAttention,
|
|
SparseResBlock,
|
|
SparseSiLU,
|
|
SparseTransformerBlock,
|
|
SparseTransformerCrossBlock,
|
|
SparseConvNeXtBlock3d,
|
|
SparseResBlockC2S3d,
|
|
)
|
|
|
|
__all__ = [
|
|
"SparseTensor", "VarLenTensor", "subdivide", "downsample", "upsample",
|
|
"spatial2channel", "channel2spatial",
|
|
"SubMConv3d", "build_indice_map",
|
|
"SparseLinear", "LayerNorm32", "SparseGroupNorm32", "SparseSiLU", "SparseGELU",
|
|
"SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention",
|
|
"SparseTransformerBlock", "SparseTransformerCrossBlock",
|
|
"SparseConvNeXtBlock3d", "SparseResBlockC2S3d",
|
|
# DiT / flow-model pieces
|
|
"MultiHeadRMSNorm", "apply_rope", "rope_phases_from_coords", "TimestepEmbedder", "DiTAttention",
|
|
"ModulatedTransformerCrossBlock", "ProjectAttention",
|
|
"SparseDiTAttention", "SparseProjectAttention",
|
|
"ModulatedSparseTransformerCrossBlock",
|
|
]
|
|
__version__ = "0.1.0"
|