206 lines
7.0 KiB
Python
206 lines
7.0 KiB
Python
"""
|
|
VAE decoder building blocks in MLX.
|
|
Matches PyTorch SparseResBlock3d, SparseConvNeXtBlock3d, etc.
|
|
"""
|
|
import logging
|
|
import time
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
from .norm import LayerNorm32
|
|
from .sparse_tensor import MlxSparseTensor
|
|
from .sparse_conv import MlxSparseConv3d
|
|
from .sparse_ops import (
|
|
MlxSparseDownsample, MlxSparseUpsample,
|
|
MlxSparseSpatial2Channel, MlxSparseChannel2Spatial,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MlxSparseLinear(nn.Module):
|
|
"""Linear layer operating on sparse tensor features."""
|
|
|
|
def __init__(self, in_features: int, out_features: int, bias: bool = True):
|
|
super().__init__()
|
|
self.linear = nn.Linear(in_features, out_features, bias=bias)
|
|
|
|
def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
|
|
return x.replace(self.linear(x.feats))
|
|
|
|
|
|
class MlxSparseConvNeXtBlock3d(nn.Module):
|
|
"""ConvNeXt-style block: conv → norm → MLP residual."""
|
|
|
|
def __init__(self, channels: int, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
|
self.conv = MlxSparseConv3d(channels, channels, 3)
|
|
hidden = int(channels * mlp_ratio)
|
|
self.mlp = MlxConvNeXtMLP(channels, hidden)
|
|
|
|
def __call__(self, x: MlxSparseTensor) -> MlxSparseTensor:
|
|
t0 = time.time()
|
|
h = self.conv(x)
|
|
h = h.replace(self.norm(h.feats))
|
|
h = h.replace(self.mlp(h.feats))
|
|
result = MlxSparseTensor(
|
|
feats=h.feats + x.feats,
|
|
coords=x.coords,
|
|
shape=x.shape,
|
|
_scale=x._scale,
|
|
_spatial_cache=x._spatial_cache,
|
|
)
|
|
dt = time.time() - t0
|
|
if dt > 0.5:
|
|
logger.debug("[MLX] ConvNeXtBlock: N=%d, ch=%d, %.2fs",
|
|
x.feats.shape[0], self.channels, dt)
|
|
return result
|
|
|
|
|
|
class MlxConvNeXtMLP(nn.Module):
|
|
"""MLP for ConvNeXt block: Linear → SiLU → Linear (zero-init)."""
|
|
|
|
def __init__(self, channels: int, hidden: int):
|
|
super().__init__()
|
|
# layers list used by weight remapping (mlp.layers.0/2)
|
|
self.layers = [
|
|
nn.Linear(channels, hidden),
|
|
None, # SiLU (not a module)
|
|
nn.Linear(hidden, channels),
|
|
]
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return self.layers[2](nn.silu(self.layers[0](x)))
|
|
|
|
|
|
class MlxSparseResBlockC2S3d(nn.Module):
|
|
"""
|
|
Residual block with Channel2Spatial upsampling.
|
|
Used in shape decoder to go from coarse to fine resolution.
|
|
"""
|
|
|
|
def __init__(self, channels: int, out_channels: int = None,
|
|
pred_subdiv: bool = True):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.out_channels = out_channels or channels
|
|
self.pred_subdiv = pred_subdiv
|
|
|
|
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
|
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
|
# conv1 outputs out_channels * 8 (spatial expansion)
|
|
self.conv1 = MlxSparseConv3d(channels, self.out_channels * 8, 3)
|
|
self.conv2 = MlxSparseConv3d(self.out_channels, self.out_channels, 3)
|
|
if pred_subdiv:
|
|
self.to_subdiv = MlxSparseLinear(channels, 8)
|
|
self.updown = MlxSparseChannel2Spatial(2)
|
|
|
|
def _skip_connection(self, x: MlxSparseTensor) -> MlxSparseTensor:
|
|
"""Repeat features to match out_channels after C2S."""
|
|
# x after C2S has channels // 8 features
|
|
feats = x.feats
|
|
c_in = feats.shape[1]
|
|
repeats = self.out_channels // c_in
|
|
if repeats > 1:
|
|
feats = mx.repeat(feats, repeats, axis=1)
|
|
return x.replace(feats[:, :self.out_channels])
|
|
|
|
def __call__(self, x: MlxSparseTensor, subdiv: MlxSparseTensor = None):
|
|
t0 = time.time()
|
|
N_in = x.feats.shape[0]
|
|
|
|
if self.pred_subdiv:
|
|
subdiv = self.to_subdiv(x)
|
|
|
|
h = x.replace(self.norm1(x.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv1(h)
|
|
|
|
subdiv_bin = subdiv.replace((subdiv.feats > 0).astype(subdiv.feats.dtype)) if subdiv is not None else None
|
|
h = self.updown(h, subdiv_bin)
|
|
x = self.updown(x, subdiv_bin)
|
|
|
|
h = h.replace(self.norm2(h.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv2(h)
|
|
|
|
skip = self._skip_connection(x)
|
|
result = MlxSparseTensor(
|
|
feats=h.feats + skip.feats,
|
|
coords=h.coords,
|
|
shape=h.shape,
|
|
_scale=h._scale,
|
|
_spatial_cache=h._spatial_cache,
|
|
)
|
|
|
|
dt = time.time() - t0
|
|
logger.debug("[MLX] ResBlockC2S: N %d → %d, ch %d → %d, %.2fs",
|
|
N_in, result.feats.shape[0], self.channels, self.out_channels, dt)
|
|
|
|
if self.pred_subdiv:
|
|
return result, subdiv
|
|
return result
|
|
|
|
|
|
class MlxSparseResBlockUpsample3d(nn.Module):
|
|
"""Residual block with nearest-neighbor upsampling."""
|
|
|
|
def __init__(self, channels: int, out_channels: int = None,
|
|
pred_subdiv: bool = True):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.out_channels = out_channels or channels
|
|
self.pred_subdiv = pred_subdiv
|
|
|
|
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
|
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
|
self.conv1 = MlxSparseConv3d(channels, self.out_channels, 3)
|
|
self.conv2 = MlxSparseConv3d(self.out_channels, self.out_channels, 3)
|
|
if channels != self.out_channels:
|
|
self.skip_connection = MlxSparseLinear(channels, self.out_channels)
|
|
else:
|
|
self.skip_connection = None
|
|
if pred_subdiv:
|
|
self.to_subdiv = MlxSparseLinear(channels, 8)
|
|
self.updown = MlxSparseUpsample(2)
|
|
|
|
def __call__(self, x: MlxSparseTensor, subdiv: MlxSparseTensor = None):
|
|
t0 = time.time()
|
|
N_in = x.feats.shape[0]
|
|
|
|
if self.pred_subdiv:
|
|
subdiv = self.to_subdiv(x)
|
|
|
|
h = x.replace(self.norm1(x.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
subdiv_bin = subdiv.replace((subdiv.feats > 0).astype(subdiv.feats.dtype)) if subdiv is not None else None
|
|
h = self.updown(h, subdiv_bin)
|
|
x_up = self.updown(x, subdiv_bin)
|
|
|
|
h = self.conv1(h)
|
|
h = h.replace(self.norm2(h.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv2(h)
|
|
|
|
if self.skip_connection is not None:
|
|
skip = self.skip_connection(x_up)
|
|
else:
|
|
skip = x_up
|
|
|
|
result = MlxSparseTensor(
|
|
feats=h.feats + skip.feats,
|
|
coords=h.coords,
|
|
shape=h.shape,
|
|
_scale=h._scale,
|
|
_spatial_cache=h._spatial_cache,
|
|
)
|
|
|
|
dt = time.time() - t0
|
|
logger.debug("[MLX] ResBlockUpsample: N %d → %d, ch %d → %d, %.2fs",
|
|
N_in, result.feats.shape[0], self.channels, self.out_channels, dt)
|
|
|
|
if self.pred_subdiv:
|
|
return result, subdiv
|
|
return result
|