265 lines
8.2 KiB
Python
265 lines
8.2 KiB
Python
"""
|
|
Thin adapters wrapping MLX models to match upstream PyTorch model interfaces.
|
|
Used by create_mlx_pipeline() so the upstream Trellis2ImageTo3DPipeline can
|
|
call MLX models transparently via torch→numpy→mx→numpy→torch conversion.
|
|
"""
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import mlx.core as mx
|
|
|
|
from .sparse_tensor import MlxSparseTensor
|
|
from trellis2.modules.sparse import SparseTensor
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Conversion helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _sparse_to_mlx(st: SparseTensor) -> MlxSparseTensor:
|
|
"""Convert upstream SparseTensor to MlxSparseTensor."""
|
|
feats = mx.array(st.feats.cpu().numpy())
|
|
coords = mx.array(st.coords.cpu().numpy().astype(np.int32))
|
|
return MlxSparseTensor(feats=feats, coords=coords)
|
|
|
|
|
|
def _mlx_to_sparse(mx_st: MlxSparseTensor, ref: SparseTensor = None) -> SparseTensor:
|
|
"""Convert MlxSparseTensor back to upstream SparseTensor.
|
|
|
|
If ref is provided and coords match, reuses backend data via .replace().
|
|
"""
|
|
mx.eval(mx_st.feats)
|
|
feats = torch.from_numpy(np.array(mx_st.feats))
|
|
if ref is not None and mx_st.coords.shape == ref.coords.shape:
|
|
return ref.replace(feats=feats)
|
|
mx.eval(mx_st.coords)
|
|
coords = torch.from_numpy(np.array(mx_st.coords)).int()
|
|
return SparseTensor(feats=feats, coords=coords)
|
|
|
|
|
|
def _torch_to_mx(t: torch.Tensor) -> mx.array:
|
|
"""Convert a PyTorch CPU tensor to MLX array."""
|
|
if t.dtype == torch.bfloat16:
|
|
return mx.array(t.float().cpu().numpy()).astype(mx.bfloat16)
|
|
return mx.array(t.cpu().numpy())
|
|
|
|
|
|
def _mx_to_torch(a: mx.array) -> torch.Tensor:
|
|
"""Convert MLX array to PyTorch CPU tensor."""
|
|
mx.eval(a)
|
|
return torch.from_numpy(np.array(a))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MlxFlowModelAdapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MlxFlowModelAdapter(nn.Module):
|
|
"""Wraps dense MlxSparseStructureFlowModel or sparse MlxSLatFlowModel.
|
|
|
|
The upstream PT sampler calls model(x_t, t, cond, **kwargs).
|
|
This adapter converts tensors, runs the MLX model, and converts back.
|
|
"""
|
|
|
|
def __init__(self, mlx_model, is_sparse: bool = False):
|
|
super().__init__()
|
|
self._mlx = mlx_model
|
|
self._is_sparse = is_sparse
|
|
self.resolution = mlx_model.resolution
|
|
self.in_channels = mlx_model.in_channels
|
|
self.out_channels = mlx_model.out_channels
|
|
|
|
def forward(self, x, t, cond, **kwargs):
|
|
if self._is_sparse:
|
|
return self._forward_sparse(x, t, cond, **kwargs)
|
|
return self._forward_dense(x, t, cond)
|
|
|
|
def _forward_dense(self, x, t, cond):
|
|
mx_out = self._mlx(_torch_to_mx(x), _torch_to_mx(t), _torch_to_mx(cond))
|
|
return _mx_to_torch(mx_out)
|
|
|
|
def _forward_sparse(self, x, t, cond, **kwargs):
|
|
mx_kwargs = {}
|
|
if kwargs.get('concat_cond') is not None:
|
|
mx_kwargs['concat_cond'] = _sparse_to_mlx(kwargs['concat_cond'])
|
|
mx_out = self._mlx(
|
|
_sparse_to_mlx(x), _torch_to_mx(t), _torch_to_mx(cond), **mx_kwargs
|
|
)
|
|
return _mlx_to_sparse(mx_out, x)
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MlxStructureDecoderAdapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MlxStructureDecoderAdapter(nn.Module):
|
|
"""Wraps MlxSparseStructureDecoder for the upstream pipeline."""
|
|
|
|
def __init__(self, mlx_model):
|
|
super().__init__()
|
|
self._mlx = mlx_model
|
|
|
|
def forward(self, z_s):
|
|
return _mx_to_torch(self._mlx(_torch_to_mx(z_s)))
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MlxFlexiDualGridAdapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MlxFlexiDualGridAdapter(nn.Module):
|
|
"""Wraps MlxFlexiDualGridVaeDecoder to match upstream FlexiDualGridVaeDecoder.
|
|
|
|
Upstream returns (meshes, subs) in eval mode with return_subs=True,
|
|
where meshes is a list of Mesh objects. This adapter runs the MLX decoder,
|
|
converts outputs to torch, then does mesh extraction via o_voxel.
|
|
"""
|
|
|
|
def __init__(self, mlx_model):
|
|
super().__init__()
|
|
self._mlx = mlx_model
|
|
self.low_vram = False # upstream toggles this
|
|
|
|
@property
|
|
def resolution(self):
|
|
return self._mlx.resolution
|
|
|
|
def set_resolution(self, resolution):
|
|
self._mlx.set_resolution(resolution)
|
|
|
|
def forward(self, x, return_subs=False, **kwargs):
|
|
from trellis2.representations import Mesh
|
|
from o_voxel.convert import flexible_dual_grid_to_mesh
|
|
|
|
mx_x = _sparse_to_mlx(x)
|
|
result = self._mlx(mx_x, return_subs=return_subs)
|
|
|
|
if return_subs:
|
|
(h_mx, verts_mx, inter_mx, quad_mx), subs_mx = result
|
|
else:
|
|
(h_mx, verts_mx, inter_mx, quad_mx) = result
|
|
subs_mx = None
|
|
|
|
mx.eval(h_mx.feats, h_mx.coords, verts_mx, inter_mx, quad_mx)
|
|
|
|
# Convert to torch for mesh extraction
|
|
coords_t = torch.from_numpy(np.array(h_mx.coords[:, 1:])).int()
|
|
verts_t = torch.from_numpy(np.array(verts_mx)).float()
|
|
inter_t = torch.from_numpy(np.array(inter_mx)).bool()
|
|
quad_t = torch.from_numpy(np.array(quad_mx)).float()
|
|
|
|
mesh_verts, mesh_faces = flexible_dual_grid_to_mesh(
|
|
coords_t, verts_t, inter_t, quad_t,
|
|
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
grid_size=self._mlx.resolution,
|
|
train=False,
|
|
)
|
|
meshes = [Mesh(mesh_verts, mesh_faces)]
|
|
|
|
if return_subs:
|
|
subs_t = []
|
|
for s in subs_mx:
|
|
mx.eval(s.feats)
|
|
subs_t.append(_mlx_to_sparse(s))
|
|
return meshes, subs_t
|
|
|
|
return meshes
|
|
|
|
def upsample(self, x, upsample_times):
|
|
mx_x = _sparse_to_mlx(x)
|
|
mx_coords = self._mlx.upsample(mx_x, upsample_times)
|
|
mx.eval(mx_coords)
|
|
return torch.from_numpy(np.array(mx_coords)).int()
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MlxTexVaeDecoderAdapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MlxTexVaeDecoderAdapter(nn.Module):
|
|
"""Wraps MlxSparseUnetVaeDecoder (texture) for the upstream pipeline.
|
|
|
|
Returns SparseTensor so upstream arithmetic (* 0.5 + 0.5) works.
|
|
"""
|
|
|
|
def __init__(self, mlx_model):
|
|
super().__init__()
|
|
self._mlx = mlx_model
|
|
|
|
def forward(self, x, guide_subs=None, **kwargs):
|
|
mx_x = _sparse_to_mlx(x)
|
|
mx_guide = None
|
|
if guide_subs is not None:
|
|
mx_guide = [_sparse_to_mlx(s) for s in guide_subs]
|
|
|
|
result = self._mlx(mx_x, guide_subs=mx_guide)
|
|
mx.eval(result.feats)
|
|
return _mlx_to_sparse(result, x)
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MlxImageCondAdapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MlxImageCondAdapter:
|
|
"""Wraps MlxDINOv3FeatureExtractor for the upstream pipeline.
|
|
|
|
The upstream pipeline sets .image_size then calls model(images).
|
|
"""
|
|
|
|
def __init__(self, mlx_dino):
|
|
self._mlx = mlx_dino
|
|
self.image_size = 512
|
|
|
|
def __call__(self, images):
|
|
from PIL import Image as PILImage
|
|
resized = [
|
|
img.resize((self.image_size, self.image_size), PILImage.LANCZOS)
|
|
for img in images
|
|
]
|
|
mx_out = self._mlx(resized)
|
|
mx.eval(mx_out)
|
|
return torch.from_numpy(np.array(mx_out))
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
def cpu(self):
|
|
return self
|