308 lines
13 KiB
Python
308 lines
13 KiB
Python
"""Tests for Wan weight conversion utilities."""
|
|
|
|
import logging
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Transformer Weight Conversion Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestSanitizeTransformerWeights:
|
|
def test_patch_embedding_reshape(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"patch_embedding.weight": mx.random.normal((5120, 16, 1, 2, 2)),
|
|
"patch_embedding.bias": mx.random.normal((5120,)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
assert "patch_embedding_proj.weight" in out
|
|
assert "patch_embedding_proj.bias" in out
|
|
assert out["patch_embedding_proj.weight"].shape == (5120, 16 * 1 * 2 * 2)
|
|
|
|
def test_text_embedding_rename(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"text_embedding.0.weight": mx.zeros((64, 32)),
|
|
"text_embedding.0.bias": mx.zeros((64,)),
|
|
"text_embedding.2.weight": mx.zeros((64, 64)),
|
|
"text_embedding.2.bias": mx.zeros((64,)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
assert "text_embedding_0.weight" in out
|
|
assert "text_embedding_0.bias" in out
|
|
assert "text_embedding_1.weight" in out
|
|
assert "text_embedding_1.bias" in out
|
|
|
|
def test_time_embedding_rename(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"time_embedding.0.weight": mx.zeros((64, 32)),
|
|
"time_embedding.2.weight": mx.zeros((64, 64)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
assert "time_embedding_0.weight" in out
|
|
assert "time_embedding_1.weight" in out
|
|
|
|
def test_time_projection_rename(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"time_projection.1.weight": mx.zeros((384, 64)),
|
|
"time_projection.1.bias": mx.zeros((384,)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
assert "time_projection.weight" in out
|
|
assert "time_projection.bias" in out
|
|
|
|
def test_ffn_rename(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"blocks.0.ffn.0.weight": mx.zeros((128, 64)),
|
|
"blocks.0.ffn.0.bias": mx.zeros((128,)),
|
|
"blocks.0.ffn.2.weight": mx.zeros((64, 128)),
|
|
"blocks.0.ffn.2.bias": mx.zeros((64,)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
assert "blocks.0.ffn.fc1.weight" in out
|
|
assert "blocks.0.ffn.fc1.bias" in out
|
|
assert "blocks.0.ffn.fc2.weight" in out
|
|
assert "blocks.0.ffn.fc2.bias" in out
|
|
|
|
def test_freqs_skipped(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"freqs": mx.zeros((1024, 64, 2)),
|
|
"blocks.0.norm1.weight": mx.zeros((64,)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
assert "freqs" not in out
|
|
assert "blocks.0.norm1.weight" in out
|
|
|
|
def test_passthrough_keys(self):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"blocks.0.self_attn.q.weight": mx.zeros((64, 64)),
|
|
"blocks.0.self_attn.k.weight": mx.zeros((64, 64)),
|
|
"blocks.0.self_attn.v.weight": mx.zeros((64, 64)),
|
|
"blocks.0.self_attn.o.weight": mx.zeros((64, 64)),
|
|
"blocks.0.modulation": mx.zeros((1, 6, 64)),
|
|
"head.head.weight": mx.zeros((64, 64)),
|
|
"head.modulation": mx.zeros((1, 2, 64)),
|
|
}
|
|
out = sanitize_wan_transformer_weights(weights)
|
|
for key in weights:
|
|
assert key in out
|
|
|
|
def test_no_unconsumed_keys(self, caplog):
|
|
from mlx_video.convert_wan import sanitize_wan_transformer_weights
|
|
weights = {
|
|
"patch_embedding.weight": mx.random.normal((5120, 16, 1, 2, 2)),
|
|
"patch_embedding.bias": mx.random.normal((5120,)),
|
|
"text_embedding.0.weight": mx.zeros((64, 32)),
|
|
"text_embedding.2.weight": mx.zeros((64, 64)),
|
|
"time_embedding.0.weight": mx.zeros((64, 32)),
|
|
"time_embedding.2.weight": mx.zeros((64, 64)),
|
|
"time_projection.1.weight": mx.zeros((384, 64)),
|
|
"blocks.0.ffn.0.weight": mx.zeros((128, 64)),
|
|
"blocks.0.ffn.2.weight": mx.zeros((64, 128)),
|
|
"blocks.0.self_attn.q.weight": mx.zeros((64, 64)),
|
|
"blocks.0.modulation": mx.zeros((1, 6, 64)),
|
|
"head.head.weight": mx.zeros((64, 64)),
|
|
"freqs": mx.zeros((1024, 64, 2)),
|
|
}
|
|
with caplog.at_level(logging.WARNING, logger="mlx_video.convert_wan"):
|
|
sanitize_wan_transformer_weights(weights)
|
|
assert "Unconsumed" not in caplog.text
|
|
|
|
|
|
class TestSanitizeT5Weights:
|
|
def test_gate_rename(self):
|
|
from mlx_video.convert_wan import sanitize_wan_t5_weights
|
|
weights = {
|
|
"blocks.0.ffn.gate.0.weight": mx.zeros((128, 64)),
|
|
"blocks.0.ffn.fc1.weight": mx.zeros((128, 64)),
|
|
"blocks.0.ffn.fc2.weight": mx.zeros((64, 128)),
|
|
}
|
|
out = sanitize_wan_t5_weights(weights)
|
|
assert "blocks.0.ffn.gate_proj.weight" in out
|
|
assert "blocks.0.ffn.fc1.weight" in out
|
|
assert "blocks.0.ffn.fc2.weight" in out
|
|
|
|
def test_passthrough(self):
|
|
from mlx_video.convert_wan import sanitize_wan_t5_weights
|
|
weights = {
|
|
"token_embedding.weight": mx.zeros((100, 64)),
|
|
"blocks.0.attn.q.weight": mx.zeros((64, 64)),
|
|
"norm.weight": mx.zeros((64,)),
|
|
}
|
|
out = sanitize_wan_t5_weights(weights)
|
|
for key in weights:
|
|
assert key in out
|
|
|
|
def test_no_unconsumed_keys(self, caplog):
|
|
from mlx_video.convert_wan import sanitize_wan_t5_weights
|
|
weights = {
|
|
"token_embedding.weight": mx.zeros((100, 64)),
|
|
"blocks.0.ffn.gate.0.weight": mx.zeros((128, 64)),
|
|
"blocks.0.ffn.fc1.weight": mx.zeros((128, 64)),
|
|
"blocks.0.ffn.fc2.weight": mx.zeros((64, 128)),
|
|
"norm.weight": mx.zeros((64,)),
|
|
}
|
|
with caplog.at_level(logging.WARNING, logger="mlx_video.convert_wan"):
|
|
sanitize_wan_t5_weights(weights)
|
|
assert "Unconsumed" not in caplog.text
|
|
|
|
|
|
class TestSanitizeVAEWeights:
|
|
def test_conv3d_transpose(self):
|
|
from mlx_video.convert_wan import sanitize_wan_vae_weights
|
|
weights = {
|
|
"decoder.conv1.weight": mx.zeros((8, 4, 3, 3, 3)), # [O, I, D, H, W]
|
|
}
|
|
out = sanitize_wan_vae_weights(weights)
|
|
assert out["decoder.conv1.weight"].shape == (8, 3, 3, 3, 4) # [O, D, H, W, I]
|
|
|
|
def test_conv2d_transpose(self):
|
|
from mlx_video.convert_wan import sanitize_wan_vae_weights
|
|
weights = {
|
|
"decoder.proj.weight": mx.zeros((16, 8, 3, 3)), # [O, I, H, W]
|
|
}
|
|
out = sanitize_wan_vae_weights(weights)
|
|
assert out["decoder.proj.weight"].shape == (16, 3, 3, 8) # [O, H, W, I]
|
|
|
|
def test_non_conv_passthrough(self):
|
|
from mlx_video.convert_wan import sanitize_wan_vae_weights
|
|
weights = {
|
|
"decoder.norm.weight": mx.zeros((64,)), # 1D, no transpose
|
|
"decoder.bias": mx.zeros((16,)),
|
|
}
|
|
out = sanitize_wan_vae_weights(weights)
|
|
assert out["decoder.norm.weight"].shape == (64,)
|
|
assert out["decoder.bias"].shape == (16,)
|
|
|
|
def test_mixed_weights(self):
|
|
from mlx_video.convert_wan import sanitize_wan_vae_weights
|
|
weights = {
|
|
"conv3d.weight": mx.zeros((8, 4, 3, 3, 3)), # 5D
|
|
"conv2d.weight": mx.zeros((8, 4, 3, 3)), # 4D
|
|
"linear.weight": mx.zeros((8, 4)), # 2D
|
|
"norm.weight": mx.zeros((8,)), # 1D
|
|
}
|
|
out = sanitize_wan_vae_weights(weights)
|
|
assert out["conv3d.weight"].shape == (8, 3, 3, 3, 4)
|
|
assert out["conv2d.weight"].shape == (8, 3, 3, 4)
|
|
assert out["linear.weight"].shape == (8, 4)
|
|
assert out["norm.weight"].shape == (8,)
|
|
|
|
def test_no_unconsumed_keys(self, caplog):
|
|
from mlx_video.convert_wan import sanitize_wan_vae_weights
|
|
weights = {
|
|
"decoder.conv1.weight": mx.zeros((8, 4, 3, 3, 3)),
|
|
"decoder.proj.weight": mx.zeros((16, 8, 3, 3)),
|
|
"decoder.norm.weight": mx.zeros((64,)),
|
|
"decoder.bias": mx.zeros((16,)),
|
|
}
|
|
with caplog.at_level(logging.WARNING, logger="mlx_video.convert_wan"):
|
|
sanitize_wan_vae_weights(weights)
|
|
assert "Unconsumed" not in caplog.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Wan2.1 Conversion Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestWan21Convert:
|
|
"""Tests for Wan2.1 conversion support."""
|
|
|
|
def test_auto_detect_wan21(self, tmp_path):
|
|
"""Auto-detect single-model directory as Wan2.1."""
|
|
# Create a Wan2.1-style directory (no low_noise_model subdir)
|
|
(tmp_path / "dummy.safetensors").touch()
|
|
# The auto-detect logic: no low_noise_model dir → 2.1
|
|
from pathlib import Path
|
|
low = tmp_path / "low_noise_model"
|
|
assert not low.exists()
|
|
# Simulates auto detection
|
|
version = "2.2" if low.exists() else "2.1"
|
|
assert version == "2.1"
|
|
|
|
def test_auto_detect_wan22(self, tmp_path):
|
|
"""Auto-detect dual-model directory as Wan2.2."""
|
|
(tmp_path / "low_noise_model").mkdir()
|
|
(tmp_path / "high_noise_model").mkdir()
|
|
from pathlib import Path
|
|
low = tmp_path / "low_noise_model"
|
|
assert low.exists()
|
|
version = "2.2" if low.exists() else "2.1"
|
|
assert version == "2.2"
|
|
|
|
def test_wan21_config_saved_correctly(self):
|
|
"""Verify config dict has correct fields for Wan2.1."""
|
|
from mlx_video.models.wan.config import WanModelConfig
|
|
config = WanModelConfig.wan21_t2v_14b()
|
|
d = config.to_dict()
|
|
assert d["model_version"] == "2.1"
|
|
assert d["dual_model"] is False
|
|
assert d["sample_steps"] == 50
|
|
assert d["sample_shift"] == 5.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Encoder Weight Sanitization Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestSanitizeEncoderWeights:
|
|
"""Tests for sanitize_wan22_vae_weights with include_encoder."""
|
|
|
|
def test_exclude_encoder_by_default(self):
|
|
from mlx_video.models.wan.vae22 import sanitize_wan22_vae_weights
|
|
|
|
weights = {
|
|
"encoder.conv1.weight": mx.zeros((8, 1, 3, 3, 3)),
|
|
"conv1.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
"conv2.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
}
|
|
out = sanitize_wan22_vae_weights(weights, include_encoder=False)
|
|
assert "conv2.weight" in out
|
|
assert not any("encoder" in k or k.startswith("conv1") for k in out)
|
|
|
|
def test_include_encoder(self):
|
|
from mlx_video.models.wan.vae22 import sanitize_wan22_vae_weights
|
|
|
|
weights = {
|
|
"encoder.conv1.weight": mx.zeros((8, 1, 3, 3, 3)),
|
|
"conv1.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
"conv2.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
}
|
|
out = sanitize_wan22_vae_weights(weights, include_encoder=True)
|
|
assert "encoder.conv1.weight" in out
|
|
assert "conv1.weight" in out
|
|
assert "conv2.weight" in out
|
|
|
|
def test_no_unconsumed_keys(self, caplog):
|
|
from mlx_video.models.wan.vae22 import sanitize_wan22_vae_weights
|
|
|
|
weights = {
|
|
"encoder.conv1.weight": mx.zeros((8, 1, 3, 3, 3)),
|
|
"conv1.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
"conv2.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
}
|
|
with caplog.at_level(logging.WARNING, logger="mlx_video.models.wan.vae22"):
|
|
sanitize_wan22_vae_weights(weights, include_encoder=True)
|
|
assert "Unconsumed" not in caplog.text
|
|
|
|
def test_no_unconsumed_keys_exclude_encoder(self, caplog):
|
|
from mlx_video.models.wan.vae22 import sanitize_wan22_vae_weights
|
|
|
|
weights = {
|
|
"encoder.conv1.weight": mx.zeros((8, 1, 3, 3, 3)),
|
|
"conv1.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
"conv2.weight": mx.zeros((8, 1, 1, 1, 8)),
|
|
}
|
|
with caplog.at_level(logging.WARNING, logger="mlx_video.models.wan.vae22"):
|
|
sanitize_wan22_vae_weights(weights, include_encoder=False)
|
|
assert "Unconsumed" not in caplog.text
|