121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
"""
|
|
MLX backend for Trellis2 on Apple Silicon.
|
|
|
|
All models run in pure MLX — no PyTorch except at the mesh-extraction boundary.
|
|
Weight loading uses mx.load() for zero-copy safetensors reads.
|
|
"""
|
|
import logging
|
|
import mlx.core as mx
|
|
|
|
|
|
def setup_logging(level=logging.INFO):
|
|
"""Enable MLX backend logging. Call before pipeline.run()."""
|
|
logging.basicConfig(
|
|
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
)
|
|
logging.getLogger("mlx_backend").setLevel(level)
|
|
|
|
__all__ = [
|
|
'load_safetensors',
|
|
'remap_flow_model_weights',
|
|
'remap_vae_decoder_weights',
|
|
]
|
|
|
|
|
|
def load_safetensors(path: str, dtype=None) -> dict:
|
|
"""Load safetensors weights via MLX.
|
|
|
|
Keeps native dtype (bf16/fp16) by default. The RoPE interleaving fix
|
|
ensures numerical correctness regardless of weight precision.
|
|
Pass dtype=mx.float32 to force float32 if needed for debugging.
|
|
"""
|
|
weights = mx.load(path)
|
|
if dtype is not None:
|
|
weights = {k: v.astype(dtype) if v.dtype in (mx.bfloat16, mx.float16) else v
|
|
for k, v in weights.items()}
|
|
return weights
|
|
|
|
|
|
def remap_flow_model_weights(weights: dict) -> dict:
|
|
"""
|
|
Remap safetensors keys to MLX module paths for flow models.
|
|
|
|
MLX uses plain list indices for model.blocks (no `.layers.`):
|
|
blocks.0.xxx stays as blocks.0.xxx
|
|
|
|
But nn.Sequential-like containers use `.layers.`:
|
|
blocks.0.mlp.mlp.0.weight → blocks.0.mlp.mlp.layers.0.weight
|
|
adaLN_modulation.1.weight → adaLN_modulation.layers.1.weight
|
|
t_embedder.mlp.0.weight → t_embedder.mlp.layers.0.weight
|
|
"""
|
|
remapped = {}
|
|
for k, v in weights.items():
|
|
new_k = k
|
|
# blocks.N.xxx stays as-is (MLX uses integer indices for lists)
|
|
# Only remap Sequential containers within blocks
|
|
|
|
# Sequential .N. → .layers.N. for specific containers
|
|
new_k = _remap_sequential(new_k)
|
|
|
|
remapped[new_k] = v
|
|
return remapped
|
|
|
|
|
|
def remap_vae_decoder_weights(weights: dict) -> dict:
|
|
"""
|
|
Remap safetensors keys to MLX module paths for VAE decoders.
|
|
|
|
VAE decoder has nested lists: blocks[i][j].xxx
|
|
MLX uses plain integer indices: blocks.i.j.xxx (no .layers.)
|
|
"""
|
|
remapped = {}
|
|
for k, v in weights.items():
|
|
new_k = k
|
|
|
|
# MlxSparseLinear wraps nn.Linear as self.linear
|
|
# from_latent.weight → from_latent.linear.weight
|
|
# output_layer.weight → output_layer.linear.weight
|
|
# Also handle to_subdiv in upsample blocks
|
|
for linear_name in ['from_latent', 'output_layer', 'to_subdiv', 'skip_connection']:
|
|
if new_k.startswith(f'{linear_name}.'):
|
|
new_k = new_k.replace(f'{linear_name}.', f'{linear_name}.linear.', 1)
|
|
break
|
|
# Handle nested: blocks.X.Y.to_subdiv.weight
|
|
import re
|
|
pattern = rf'(blocks\.\d+\.\d+\.{linear_name})\.'
|
|
new_k_try = re.sub(pattern, rf'\1.linear.', new_k, count=1)
|
|
if new_k_try != new_k:
|
|
new_k = new_k_try
|
|
break
|
|
|
|
# Sequential .N. → .layers.N. for MLP-like containers
|
|
new_k = _remap_sequential(new_k)
|
|
|
|
remapped[new_k] = v
|
|
return remapped
|
|
|
|
|
|
def _remap_sequential(key: str) -> str:
|
|
"""
|
|
Convert PyTorch nn.Sequential index notation to MLX.
|
|
e.g. 'foo.0.weight' → 'foo.layers.0.weight' when '0' is a digit
|
|
within known sequential containers.
|
|
"""
|
|
import re
|
|
# Match patterns like: prefix.DIGIT.suffix where prefix ends with
|
|
# a known sequential container name
|
|
sequential_containers = [
|
|
'mlp.mlp', 'adaLN_modulation', 't_embedder.mlp', 'mlp'
|
|
]
|
|
for container in sequential_containers:
|
|
# Pattern: container.DIGIT.rest
|
|
pattern = re.escape(container) + r'\.(\d+)\.'
|
|
replacement = container + r'.layers.\1.'
|
|
key = re.sub(pattern, replacement, key)
|
|
# Also handle terminal case: container.DIGIT (no trailing dot)
|
|
pattern_end = re.escape(container) + r'\.(\d+)$'
|
|
replacement_end = container + r'.layers.\1'
|
|
key = re.sub(pattern_end, replacement_end, key)
|
|
return key
|