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
208 lines
6.9 KiB
Python
208 lines
6.9 KiB
Python
"""Tests for MLX diffusion model components (Phase 1 + Phase 2)."""
|
|
|
|
import sys
|
|
import os
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
from hunyuanpaintpbr_mlx.unet.attn_processor_mlx import (
|
|
ImageProjModel,
|
|
PoseRoPEAttnProcessor,
|
|
RefAttnProcessor,
|
|
RotaryEmbedding,
|
|
SelfAttnProcessor,
|
|
reshape_for_attention,
|
|
reshape_from_attention,
|
|
scaled_dot_product_attention,
|
|
)
|
|
from hunyuanpaintpbr_mlx.scheduler_mlx import (
|
|
SchedulerConfig,
|
|
UniPCMultistepSchedulerMLX,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RotaryEmbedding
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestRotaryEmbedding:
|
|
def test_1d_shape(self):
|
|
pos = mx.arange(32)
|
|
cos, sin = RotaryEmbedding.get_1d_rotary_pos_embed(96, pos)
|
|
assert cos.shape == (32, 96)
|
|
assert sin.shape == (32, 96)
|
|
|
|
def test_1d_values_normalized(self):
|
|
pos = mx.arange(16)
|
|
cos, sin = RotaryEmbedding.get_1d_rotary_pos_embed(64, pos)
|
|
mx.synchronize()
|
|
# cos^2 + sin^2 should be 1 for each position
|
|
magnitude = np.array(cos * cos + sin * sin)
|
|
np.testing.assert_allclose(magnitude, 1.0, atol=1e-5)
|
|
|
|
def test_3d_shape(self):
|
|
position = mx.zeros((2, 100, 3), dtype=mx.int32)
|
|
cos, sin = RotaryEmbedding.get_3d_rotary_pos_embed(position, 96, 128)
|
|
assert cos.shape == (2, 100, 96)
|
|
assert sin.shape == (2, 100, 96)
|
|
|
|
def test_apply_rotary_preserves_shape(self):
|
|
x = mx.random.normal((2, 8, 64, 96))
|
|
pos = mx.arange(64)
|
|
cos, sin = RotaryEmbedding.get_1d_rotary_pos_embed(96, pos)
|
|
result = RotaryEmbedding.apply_rotary_emb(x, (cos, sin))
|
|
assert result.shape == x.shape
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ImageProjModel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestImageProjModel:
|
|
def test_output_shape(self):
|
|
model = ImageProjModel(
|
|
cross_attention_dim=768,
|
|
clip_embeddings_dim=1536,
|
|
clip_extra_context_tokens=4,
|
|
)
|
|
x = mx.random.normal((1, 100, 1536))
|
|
out = model(x)
|
|
mx.synchronize()
|
|
assert out.shape == (1, 400, 768) # 100 * 4 = 400
|
|
|
|
def test_2d_input(self):
|
|
model = ImageProjModel(
|
|
cross_attention_dim=768,
|
|
clip_embeddings_dim=1536,
|
|
clip_extra_context_tokens=4,
|
|
)
|
|
x = mx.random.normal((50, 1536))
|
|
out = model(x)
|
|
mx.synchronize()
|
|
assert out.shape == (50, 4, 768)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Attention utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestAttentionUtils:
|
|
def test_reshape_roundtrip(self):
|
|
x = mx.random.normal((2, 64, 768))
|
|
reshaped = reshape_for_attention(x, 8)
|
|
assert reshaped.shape == (2, 8, 64, 96) # 768/8 = 96
|
|
back = reshape_from_attention(reshaped)
|
|
assert back.shape == (2, 64, 768)
|
|
mx.synchronize()
|
|
np.testing.assert_allclose(np.array(back), np.array(x), atol=1e-5)
|
|
|
|
def test_sdpa_shape(self):
|
|
q = mx.random.normal((2, 8, 64, 96))
|
|
k = mx.random.normal((2, 8, 32, 96))
|
|
v = mx.random.normal((2, 8, 32, 96))
|
|
out = scaled_dot_product_attention(q, k, v)
|
|
mx.synchronize()
|
|
assert out.shape == (2, 8, 64, 96)
|
|
|
|
def test_sdpa_self_attention(self):
|
|
q = mx.random.normal((1, 1, 4, 8))
|
|
out = scaled_dot_product_attention(q, q, q)
|
|
mx.synchronize()
|
|
assert out.shape == (1, 1, 4, 8)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Attention Processors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestSelfAttnProcessor:
|
|
def test_output_shape(self):
|
|
import mlx.nn as nn
|
|
proc = SelfAttnProcessor(query_dim=768, num_heads=8, pbr_settings=("albedo", "mr"))
|
|
to_q = nn.Linear(768, 768)
|
|
to_k = nn.Linear(768, 768)
|
|
to_v = nn.Linear(768, 768)
|
|
to_out = nn.Linear(768, 768)
|
|
hs = mx.random.normal((1, 2, 6, 64, 768))
|
|
out = proc(hs, to_q, to_k, to_v, to_out, n_views=6)
|
|
mx.synchronize()
|
|
assert out.shape == (1, 2, 6, 64, 768)
|
|
|
|
|
|
class TestPoseRoPEAttnProcessor:
|
|
def test_output_shape(self):
|
|
proc = PoseRoPEAttnProcessor(query_dim=768, num_heads=8)
|
|
hs = mx.random.normal((2, 384, 768)) # 2 materials, 6*64=384 tokens
|
|
out = proc(hs)
|
|
mx.synchronize()
|
|
assert out.shape == (2, 384, 768)
|
|
|
|
def test_with_rope(self):
|
|
proc = PoseRoPEAttnProcessor(query_dim=768, num_heads=8)
|
|
hs = mx.random.normal((2, 384, 768))
|
|
pos_idx = {
|
|
"voxel_indices": mx.zeros((2, 384, 3), dtype=mx.int32),
|
|
"voxel_resolution": 128,
|
|
}
|
|
out = proc(hs, position_indices=pos_idx)
|
|
mx.synchronize()
|
|
assert out.shape == (2, 384, 768)
|
|
|
|
|
|
class TestRefAttnProcessor:
|
|
def test_output_shape(self):
|
|
proc = RefAttnProcessor(query_dim=768, num_heads=8, pbr_settings=("albedo", "mr"))
|
|
hs = mx.random.normal((1, 384, 768))
|
|
ref = mx.random.normal((1, 64, 768))
|
|
out = proc(hs, ref)
|
|
mx.synchronize()
|
|
assert out.shape == (1, 2, 384, 768)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scheduler
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestUniPCScheduler:
|
|
def test_set_timesteps(self):
|
|
sched = UniPCMultistepSchedulerMLX()
|
|
sched.set_timesteps(15)
|
|
assert len(sched.timesteps) == 15
|
|
assert sched.timesteps[0] > sched.timesteps[-1] # descending
|
|
|
|
def test_step_shape(self):
|
|
sched = UniPCMultistepSchedulerMLX()
|
|
sched.set_timesteps(15)
|
|
sample = mx.random.normal((4, 4, 64, 64))
|
|
noise_pred = mx.random.normal((4, 4, 64, 64))
|
|
t = sched.timesteps[0]
|
|
result = sched.step(noise_pred, t, sample)
|
|
mx.synchronize()
|
|
assert result.shape == (4, 4, 64, 64)
|
|
|
|
def test_multi_step_denoising(self):
|
|
sched = UniPCMultistepSchedulerMLX()
|
|
sched.set_timesteps(5)
|
|
sample = mx.random.normal((1, 4, 8, 8))
|
|
for t in sched.timesteps:
|
|
noise = mx.random.normal(sample.shape)
|
|
sample = sched.step(noise, int(t), sample)
|
|
mx.synchronize()
|
|
assert sample.shape == (1, 4, 8, 8)
|
|
# After denoising, values should be finite
|
|
assert not mx.isnan(sample).any().item()
|
|
|
|
def test_scale_model_input_identity(self):
|
|
sched = UniPCMultistepSchedulerMLX()
|
|
sample = mx.ones((1, 4, 8, 8))
|
|
result = sched.scale_model_input(sample, 500)
|
|
np.testing.assert_array_equal(np.array(result), np.array(sample))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|