e4cfa9d1e9
Fork of dgrauet/Hunyuan3D-2.1-mlx + our generate_e2e.py CLI, env-tunable remesh (HY3D_REMESH_FACES), and HARDWARE.md. Upstream training data (mini_trainset) and demo images stripped — inference needs none of it. Full upstream history: github.com/dgrauet/Hunyuan3D-2.1-mlx
408 lines
14 KiB
Python
408 lines
14 KiB
Python
"""MLX attention processors for HunyuanPaintPBR UNet.
|
|
|
|
Ports: RotaryEmbedding, AttnCore, SelfAttnProcessor2_0,
|
|
PoseRoPEAttnProcessor2_0, RefAttnProcessor2_0.
|
|
"""
|
|
|
|
import math
|
|
from typing import Optional, Tuple
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rotary Position Embeddings
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RotaryEmbedding:
|
|
"""1D and 3D rotary position embeddings for attention."""
|
|
|
|
@staticmethod
|
|
def get_1d_rotary_pos_embed(
|
|
dim: int, pos: mx.array, theta: float = 10000.0
|
|
) -> Tuple[mx.array, mx.array]:
|
|
"""Compute 1D rotary embeddings.
|
|
|
|
Args:
|
|
dim: Embedding dimension (head_dim).
|
|
pos: (L,) position indices.
|
|
theta: Base frequency.
|
|
|
|
Returns:
|
|
(cos, sin) each of shape (L, dim).
|
|
"""
|
|
freqs = 1.0 / (theta ** (mx.arange(0, dim, 2).astype(mx.float32) / dim))
|
|
# (dim//2,)
|
|
angles = pos[:, None].astype(mx.float32) * freqs[None, :] # (L, dim//2)
|
|
cos = mx.cos(angles)
|
|
sin = mx.sin(angles)
|
|
# Repeat-interleave to full dim: [c0,c0,c1,c1,...]
|
|
cos = mx.repeat(cos, 2, axis=-1) # (L, dim)
|
|
sin = mx.repeat(sin, 2, axis=-1)
|
|
return cos, sin
|
|
|
|
@staticmethod
|
|
def get_3d_rotary_pos_embed(
|
|
position: mx.array, embed_dim: int = 96, voxel_resolution: int = 128
|
|
) -> Tuple[mx.array, mx.array]:
|
|
"""Compute 3D rotary embeddings from voxel coordinates.
|
|
|
|
Args:
|
|
position: (B, L, 3) int voxel indices.
|
|
embed_dim: Total embedding dimension (split across XYZ).
|
|
voxel_resolution: Max voxel coordinate value.
|
|
|
|
Returns:
|
|
(cos, sin) each of shape (B, L, embed_dim).
|
|
"""
|
|
dim_xy = embed_dim // 8 * 3 # 36 for embed_dim=96
|
|
dim_z = embed_dim // 8 * 2 # 24
|
|
|
|
grid = mx.arange(voxel_resolution).astype(mx.float32)
|
|
xy_cos, xy_sin = RotaryEmbedding.get_1d_rotary_pos_embed(dim_xy, grid)
|
|
z_cos, z_sin = RotaryEmbedding.get_1d_rotary_pos_embed(dim_z, grid)
|
|
# Each: (voxel_resolution, dim_xy or dim_z)
|
|
|
|
# Gather per-position embeddings
|
|
px = position[..., 0] # (B, L)
|
|
py = position[..., 1]
|
|
pz = position[..., 2]
|
|
|
|
cos = mx.concatenate([xy_cos[px], xy_cos[py], z_cos[pz]], axis=-1)
|
|
sin = mx.concatenate([xy_sin[px], xy_sin[py], z_sin[pz]], axis=-1)
|
|
return cos, sin # (B, L, embed_dim)
|
|
|
|
@staticmethod
|
|
def apply_rotary_emb(
|
|
x: mx.array, freqs_cis: Tuple[mx.array, mx.array]
|
|
) -> mx.array:
|
|
"""Apply rotary embeddings to a tensor.
|
|
|
|
Args:
|
|
x: (B, num_heads, L, head_dim)
|
|
freqs_cis: (cos, sin) each broadcastable to x.
|
|
|
|
Returns:
|
|
Rotated tensor, same shape as x.
|
|
"""
|
|
cos, sin = freqs_cis
|
|
# cos/sin: (B, L, head_dim) or (L, head_dim) — broadcast to (B, 1, L, head_dim)
|
|
if cos.ndim == 2:
|
|
cos = cos[None, None, :, :] # (1, 1, L, D)
|
|
sin = sin[None, None, :, :]
|
|
elif cos.ndim == 3:
|
|
cos = cos[:, None, :, :] # (B, 1, L, D)
|
|
sin = sin[:, None, :, :]
|
|
|
|
# Rotate pairs: (x0, x1) -> (x0*cos - x1*sin, x1*cos + x0*sin)
|
|
x_even = x[..., 0::2] # (B, H, L, D//2)
|
|
x_odd = x[..., 1::2]
|
|
cos_half = cos[..., 0::2]
|
|
sin_half = sin[..., 0::2]
|
|
|
|
out_even = x_even * cos_half - x_odd * sin_half
|
|
out_odd = x_odd * cos_half + x_even * sin_half
|
|
|
|
# Interleave back: stack on last dim then flatten
|
|
return mx.stack([out_even, out_odd], axis=-1).reshape(x.shape)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ImageProjModel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class ImageProjModel(nn.Module):
|
|
"""Projects DINO features to cross-attention dimension.
|
|
|
|
Input (B, N*L, clip_dim) → Output (B, N*L*extra_tokens, cross_dim).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
cross_attention_dim: int = 768,
|
|
clip_embeddings_dim: int = 1536,
|
|
clip_extra_context_tokens: int = 4,
|
|
):
|
|
super().__init__()
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.clip_extra_context_tokens = clip_extra_context_tokens
|
|
self.proj = nn.Linear(
|
|
clip_embeddings_dim,
|
|
clip_extra_context_tokens * cross_attention_dim,
|
|
)
|
|
self.norm = nn.LayerNorm(cross_attention_dim)
|
|
|
|
def __call__(self, image_embeds: mx.array) -> mx.array:
|
|
"""
|
|
Args:
|
|
image_embeds: (B, N*L, clip_dim) or (B*N*L, clip_dim)
|
|
Returns:
|
|
(B, N*L*extra_tokens, cross_dim)
|
|
"""
|
|
has_batch = image_embeds.ndim == 3
|
|
if has_batch:
|
|
b, nl, _ = image_embeds.shape
|
|
embeds = image_embeds.reshape(-1, image_embeds.shape[-1])
|
|
else:
|
|
embeds = image_embeds
|
|
|
|
tokens = self.proj(embeds) # (B*N*L, extra * cross_dim)
|
|
tokens = tokens.reshape(-1, self.clip_extra_context_tokens,
|
|
self.cross_attention_dim)
|
|
tokens = self.norm(tokens) # (B*N*L, extra, cross_dim)
|
|
|
|
if has_batch:
|
|
tokens = tokens.reshape(b, nl * self.clip_extra_context_tokens,
|
|
self.cross_attention_dim)
|
|
return tokens
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scaled Dot-Product Attention
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def scaled_dot_product_attention(
|
|
query: mx.array,
|
|
key: mx.array,
|
|
value: mx.array,
|
|
attn_mask: Optional[mx.array] = None,
|
|
scale: Optional[float] = None,
|
|
) -> mx.array:
|
|
"""Scaled dot-product attention.
|
|
|
|
Args:
|
|
query: (B, H, L_q, D)
|
|
key: (B, H, L_k, D)
|
|
value: (B, H, L_v, D) (L_v == L_k)
|
|
attn_mask: optional (B, H, L_q, L_k) or broadcastable.
|
|
scale: scaling factor (default 1/sqrt(D)).
|
|
|
|
Returns:
|
|
(B, H, L_q, D)
|
|
"""
|
|
if scale is None:
|
|
scale = 1.0 / math.sqrt(query.shape[-1])
|
|
|
|
scores = (query @ key.transpose(0, 1, 3, 2)) * scale # (B, H, L_q, L_k)
|
|
|
|
if attn_mask is not None:
|
|
scores = scores + attn_mask
|
|
|
|
weights = mx.softmax(scores, axis=-1)
|
|
return weights @ value
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Attention Utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def reshape_for_attention(
|
|
x: mx.array, num_heads: int
|
|
) -> mx.array:
|
|
"""Reshape (B, L, C) → (B, num_heads, L, head_dim)."""
|
|
B, L, C = x.shape
|
|
head_dim = C // num_heads
|
|
return x.reshape(B, L, num_heads, head_dim).transpose(0, 2, 1, 3)
|
|
|
|
|
|
def reshape_from_attention(
|
|
x: mx.array,
|
|
) -> mx.array:
|
|
"""Reshape (B, num_heads, L, head_dim) → (B, L, C)."""
|
|
B, H, L, D = x.shape
|
|
return x.transpose(0, 2, 1, 3).reshape(B, L, H * D)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Attention Processors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class LinearProjection(nn.Module):
|
|
"""Linear layer for Q/K/V projections."""
|
|
|
|
def __init__(self, in_dim: int, out_dim: int, bias: bool = True):
|
|
super().__init__()
|
|
self.linear = nn.Linear(in_dim, out_dim, bias=bias)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return self.linear(x)
|
|
|
|
|
|
class SelfAttnProcessor(nn.Module):
|
|
"""Material-dimension self-attention.
|
|
|
|
Maintains separate Q/K/V/out projections per PBR material.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
query_dim: int = 768,
|
|
num_heads: int = 8,
|
|
pbr_settings: tuple = ("albedo", "mr"),
|
|
):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = query_dim // num_heads
|
|
self.pbr_settings = pbr_settings
|
|
|
|
# Per-material projections (albedo uses the base attn1 projections,
|
|
# additional materials get their own)
|
|
self.extra_qkv = {}
|
|
for token in pbr_settings[1:]: # skip first (uses base)
|
|
self.extra_qkv[token] = {
|
|
"to_q": nn.Linear(query_dim, query_dim),
|
|
"to_k": nn.Linear(query_dim, query_dim),
|
|
"to_v": nn.Linear(query_dim, query_dim),
|
|
"to_out": nn.Linear(query_dim, query_dim),
|
|
}
|
|
# Register as submodules
|
|
for name, mod in self.extra_qkv[token].items():
|
|
setattr(self, f"{name}_{token}", mod)
|
|
|
|
def __call__(
|
|
self,
|
|
hidden_states: mx.array,
|
|
to_q: nn.Linear,
|
|
to_k: nn.Linear,
|
|
to_v: nn.Linear,
|
|
to_out: nn.Linear,
|
|
n_views: int = 6,
|
|
) -> mx.array:
|
|
"""
|
|
Args:
|
|
hidden_states: (B, N_pbr, N_gen, L, C) — batched per material.
|
|
to_q/k/v/out: base projections (used for first material).
|
|
n_views: number of generated views.
|
|
|
|
Returns:
|
|
(B, N_pbr, N_gen, L, C)
|
|
"""
|
|
B, N_pbr, N_gen, L, C = hidden_states.shape
|
|
results = []
|
|
|
|
for i, token in enumerate(self.pbr_settings):
|
|
hs = hidden_states[:, i] # (B, N_gen, L, C)
|
|
hs = hs.reshape(B * N_gen, L, C)
|
|
|
|
if i == 0:
|
|
q, k, v, out_proj = to_q, to_k, to_v, to_out
|
|
else:
|
|
q = getattr(self, f"to_q_{token}")
|
|
k = getattr(self, f"to_k_{token}")
|
|
v = getattr(self, f"to_v_{token}")
|
|
out_proj = getattr(self, f"to_out_{token}")
|
|
|
|
query = reshape_for_attention(q(hs), self.num_heads)
|
|
key = reshape_for_attention(k(hs), self.num_heads)
|
|
value = reshape_for_attention(v(hs), self.num_heads)
|
|
|
|
attn_out = scaled_dot_product_attention(query, key, value)
|
|
attn_out = reshape_from_attention(attn_out)
|
|
attn_out = out_proj(attn_out)
|
|
|
|
results.append(attn_out.reshape(B, 1, N_gen, L, C))
|
|
|
|
return mx.concatenate(results, axis=1)
|
|
|
|
|
|
class PoseRoPEAttnProcessor(nn.Module):
|
|
"""Multiview self-attention with 3D rotary position embeddings."""
|
|
|
|
def __init__(self, query_dim: int = 768, num_heads: int = 8):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = query_dim // num_heads
|
|
self.to_q = nn.Linear(query_dim, query_dim)
|
|
self.to_k = nn.Linear(query_dim, query_dim)
|
|
self.to_v = nn.Linear(query_dim, query_dim)
|
|
self.to_out = nn.Linear(query_dim, query_dim)
|
|
|
|
def __call__(
|
|
self,
|
|
hidden_states: mx.array,
|
|
position_indices: Optional[dict] = None,
|
|
) -> mx.array:
|
|
"""
|
|
Args:
|
|
hidden_states: (B*N_pbr, N_gen*L, C) — views concatenated.
|
|
position_indices: dict with voxel_indices and voxel_resolution.
|
|
|
|
Returns:
|
|
(B*N_pbr, N_gen*L, C)
|
|
"""
|
|
B, NL, C = hidden_states.shape
|
|
|
|
query = reshape_for_attention(self.to_q(hidden_states), self.num_heads)
|
|
key = reshape_for_attention(self.to_k(hidden_states), self.num_heads)
|
|
value = reshape_for_attention(self.to_v(hidden_states), self.num_heads)
|
|
|
|
if position_indices is not None:
|
|
vi = position_indices.get("voxel_indices")
|
|
vr = position_indices.get("voxel_resolution", 128)
|
|
if vi is not None:
|
|
rope = RotaryEmbedding.get_3d_rotary_pos_embed(
|
|
vi, self.head_dim, vr
|
|
)
|
|
query = RotaryEmbedding.apply_rotary_emb(query, rope)
|
|
key = RotaryEmbedding.apply_rotary_emb(key, rope)
|
|
|
|
attn_out = scaled_dot_product_attention(query, key, value)
|
|
attn_out = reshape_from_attention(attn_out)
|
|
return self.to_out(attn_out)
|
|
|
|
|
|
class RefAttnProcessor(nn.Module):
|
|
"""Reference cross-attention with shared Q/K, material-specific V."""
|
|
|
|
def __init__(
|
|
self,
|
|
query_dim: int = 768,
|
|
num_heads: int = 8,
|
|
pbr_settings: tuple = ("albedo", "mr"),
|
|
):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = query_dim // num_heads
|
|
self.pbr_settings = pbr_settings
|
|
self.to_q = nn.Linear(query_dim, query_dim)
|
|
self.to_k = nn.Linear(query_dim, query_dim)
|
|
|
|
# Per-material value + output projections
|
|
for token in pbr_settings:
|
|
setattr(self, f"to_v_{token}", nn.Linear(query_dim, query_dim))
|
|
setattr(self, f"to_out_{token}", nn.Linear(query_dim, query_dim))
|
|
|
|
def __call__(
|
|
self,
|
|
hidden_states: mx.array,
|
|
encoder_hidden_states: mx.array,
|
|
) -> mx.array:
|
|
"""
|
|
Args:
|
|
hidden_states: (B, N_gen*L, C) — generated features (query).
|
|
encoder_hidden_states: (B, N_ref*L, C) — reference features (key/value).
|
|
|
|
Returns:
|
|
(B, N_pbr, N_gen*L, C) — per-material outputs.
|
|
"""
|
|
query = reshape_for_attention(self.to_q(hidden_states), self.num_heads)
|
|
key = reshape_for_attention(self.to_k(encoder_hidden_states), self.num_heads)
|
|
|
|
results = []
|
|
for token in self.pbr_settings:
|
|
to_v = getattr(self, f"to_v_{token}")
|
|
to_out = getattr(self, f"to_out_{token}")
|
|
value = reshape_for_attention(to_v(encoder_hidden_states), self.num_heads)
|
|
|
|
attn_out = scaled_dot_product_attention(query, key, value)
|
|
attn_out = reshape_from_attention(attn_out)
|
|
results.append(to_out(attn_out))
|
|
|
|
return mx.stack(results, axis=1) # (B, N_pbr, N_gen*L, C)
|