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
185 lines
6.1 KiB
Python
185 lines
6.1 KiB
Python
"""Tests for MLX RRDBNet (RealESRGAN) super-resolution model."""
|
|
|
|
import os
|
|
import sys
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import pytest
|
|
|
|
# Add hy3dpaint to path for imports
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
from utils.image_super_utils_mlx import (
|
|
RRDBNetMLX,
|
|
ResidualDenseBlock,
|
|
RRDB,
|
|
_nearest_upsample_2x,
|
|
imageSuperNetMLX,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Architecture tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestResidualDenseBlock:
|
|
"""Tests for the ResidualDenseBlock module."""
|
|
|
|
def test_creation(self):
|
|
rdb = ResidualDenseBlock(num_feat=64, num_grow_ch=32)
|
|
assert rdb.conv1 is not None
|
|
assert rdb.conv5 is not None
|
|
|
|
def test_forward_shape(self):
|
|
rdb = ResidualDenseBlock(num_feat=64, num_grow_ch=32)
|
|
x = mx.zeros((1, 16, 16, 64))
|
|
out = rdb(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 16, 16, 64), f"Expected (1,16,16,64), got {out.shape}"
|
|
|
|
def test_residual_connection(self):
|
|
"""Output should differ from input (conv weights are random)."""
|
|
rdb = ResidualDenseBlock(num_feat=64, num_grow_ch=32)
|
|
x = mx.ones((1, 4, 4, 64))
|
|
out = rdb(x)
|
|
mx.eval(out)
|
|
# The residual path adds x, so output should not be zero
|
|
assert not mx.allclose(out, mx.zeros_like(out)).item()
|
|
|
|
|
|
class TestRRDB:
|
|
"""Tests for the RRDB module."""
|
|
|
|
def test_creation(self):
|
|
rrdb = RRDB(num_feat=64, num_grow_ch=32)
|
|
assert len(rrdb.rdb1.conv1.weight.shape) == 4
|
|
|
|
def test_forward_shape(self):
|
|
rrdb = RRDB(num_feat=64, num_grow_ch=32)
|
|
x = mx.zeros((1, 8, 8, 64))
|
|
out = rrdb(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 8, 8, 64), f"Expected (1,8,8,64), got {out.shape}"
|
|
|
|
|
|
class TestRRDBNetMLX:
|
|
"""Tests for the full RRDBNet model."""
|
|
|
|
def test_creation_default(self):
|
|
model = RRDBNetMLX()
|
|
assert model.scale == 4
|
|
assert len(model.body) == 23
|
|
|
|
def test_creation_custom(self):
|
|
model = RRDBNetMLX(num_in_ch=3, num_out_ch=3, num_feat=32, num_block=4, num_grow_ch=16, scale=4)
|
|
assert len(model.body) == 4
|
|
|
|
def test_forward_shape_scale4(self):
|
|
"""Input 64x64x3 should produce 256x256x3 with scale=4."""
|
|
model = RRDBNetMLX(num_in_ch=3, num_out_ch=3, num_feat=32, num_block=2, num_grow_ch=16, scale=4)
|
|
x = mx.zeros((1, 64, 64, 3))
|
|
out = model(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 256, 256, 3), f"Expected (1,256,256,3), got {out.shape}"
|
|
|
|
def test_forward_shape_scale2(self):
|
|
"""Input 64x64x3 should produce 128x128x3 with scale=2."""
|
|
model = RRDBNetMLX(num_in_ch=3, num_out_ch=3, num_feat=32, num_block=2, num_grow_ch=16, scale=2)
|
|
x = mx.zeros((1, 64, 64, 3))
|
|
out = model(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 128, 128, 3), f"Expected (1,128,128,3), got {out.shape}"
|
|
|
|
def test_forward_nonzero_input(self):
|
|
"""Model should produce non-zero output for non-zero input."""
|
|
model = RRDBNetMLX(num_in_ch=3, num_out_ch=3, num_feat=32, num_block=2, num_grow_ch=16, scale=4)
|
|
x = mx.ones((1, 16, 16, 3)) * 0.5
|
|
out = model(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 64, 64, 3)
|
|
assert not mx.allclose(out, mx.zeros_like(out)).item()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Utility tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestNearestUpsample:
|
|
"""Tests for the nearest-neighbor upsampling function."""
|
|
|
|
def test_2x_upsample_shape(self):
|
|
x = mx.zeros((1, 8, 8, 3))
|
|
out = _nearest_upsample_2x(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 16, 16, 3)
|
|
|
|
def test_2x_upsample_values(self):
|
|
"""Each pixel should be replicated in a 2x2 block."""
|
|
x = mx.array([[[[1.0, 2.0, 3.0]]]]) # (1, 1, 1, 3)
|
|
out = _nearest_upsample_2x(x)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 2, 2, 3)
|
|
expected = mx.array([[[[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]],
|
|
[[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]]])
|
|
assert mx.allclose(out, expected).item()
|
|
|
|
|
|
class TestPixelUnshuffle:
|
|
"""Tests for the pixel_unshuffle static method."""
|
|
|
|
def test_unshuffle_scale2(self):
|
|
x = mx.zeros((1, 8, 8, 3))
|
|
out = RRDBNetMLX._pixel_unshuffle(x, 2)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 4, 4, 12)
|
|
|
|
def test_unshuffle_scale4(self):
|
|
x = mx.zeros((1, 16, 16, 3))
|
|
out = RRDBNetMLX._pixel_unshuffle(x, 4)
|
|
mx.eval(out)
|
|
assert out.shape == (1, 4, 4, 48)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Weight loading test (conditional on weight file availability)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _resolve_realesrgan_weights():
|
|
"""Local override, else fetch from HF (same repo as the paint model)."""
|
|
env = os.environ.get("REALESRGAN_MLX_WEIGHTS")
|
|
if env and os.path.exists(env):
|
|
return env
|
|
try:
|
|
from huggingface_hub import hf_hub_download
|
|
return hf_hub_download(
|
|
repo_id="dgrauet/hunyuan3d-2.1-mlx",
|
|
filename="realesrgan_x4plus.safetensors",
|
|
)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
WEIGHTS_PATH = _resolve_realesrgan_weights()
|
|
|
|
|
|
@pytest.mark.skipif(WEIGHTS_PATH is None, reason="Could not resolve RealESRGAN MLX weights (no env var, HF download failed)")
|
|
class TestWeightLoading:
|
|
"""Tests that require converted MLX weights."""
|
|
|
|
def test_load_weights(self):
|
|
wrapper = imageSuperNetMLX(WEIGHTS_PATH)
|
|
assert wrapper.model is not None
|
|
|
|
def test_inference_with_weights(self):
|
|
from PIL import Image
|
|
|
|
wrapper = imageSuperNetMLX(WEIGHTS_PATH)
|
|
# Create a small test image
|
|
img = Image.fromarray(np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8))
|
|
out = wrapper(img)
|
|
assert out.size == (128, 128), f"Expected (128,128), got {out.size}"
|