The gather-scatter sparse conv should be fairly generic. Any model using 3x3x3 or 5x5x5 sparse convolutions on voxel grids could use it directly.
The main thing that's TRELLIS-specific is the neighbor cache key format, but that's a few lines to adapt.
The SDPA attention swap is even more reusable - it's just padding variable-length sequences into batches and calling torch.nn.functional.scaled_dot_product_attention.
The main thing that's TRELLIS-specific is the neighbor cache key format, but that's a few lines to adapt.
The SDPA attention swap is even more reusable - it's just padding variable-length sequences into batches and calling torch.nn.functional.scaled_dot_product_attention.