335 lines
13 KiB
Python
335 lines
13 KiB
Python
"""Tests for LoRA loading and application."""
|
|
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
class TestLoRATypes:
|
|
"""Test LoRA data structures."""
|
|
|
|
def test_lora_weights_scale(self):
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
w = LoRAWeights(
|
|
lora_A=mx.zeros((16, 64)),
|
|
lora_B=mx.zeros((128, 16)),
|
|
rank=16,
|
|
alpha=32.0,
|
|
module_name="test",
|
|
)
|
|
assert w.scale == 2.0
|
|
|
|
def test_lora_weights_scale_default(self):
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
w = LoRAWeights(
|
|
lora_A=mx.zeros((16, 64)),
|
|
lora_B=mx.zeros((128, 16)),
|
|
rank=16,
|
|
alpha=16.0,
|
|
module_name="test",
|
|
)
|
|
assert w.scale == 1.0
|
|
|
|
def test_applied_lora_delta(self):
|
|
from mlx_video.lora.types import AppliedLoRA, LoRAWeights
|
|
|
|
lora_a = mx.ones((2, 4))
|
|
lora_b = mx.ones((8, 2))
|
|
w = LoRAWeights(lora_A=lora_a, lora_B=lora_b, rank=2, alpha=2.0, module_name="test")
|
|
applied = AppliedLoRA(weights=w, strength=0.5)
|
|
delta = applied.compute_delta()
|
|
# scale=1.0, strength=0.5, B@A = [[2,2,2,2]]*8 (each row sum of 2 ones)
|
|
expected = 0.5 * mx.ones((8, 4)) * 2.0
|
|
assert mx.allclose(delta, expected).item()
|
|
|
|
|
|
class TestLoRALoader:
|
|
"""Test LoRA weight loading from safetensors."""
|
|
|
|
def _make_lora_file(self, tmp_dir, module_names, rank=4, in_dim=64, out_dim=128, key_format="AB"):
|
|
"""Helper to create a mock LoRA safetensors file."""
|
|
weights = {}
|
|
for name in module_names:
|
|
if key_format == "AB":
|
|
weights[f"{name}.lora_A.weight"] = mx.random.normal((rank, in_dim))
|
|
weights[f"{name}.lora_B.weight"] = mx.random.normal((out_dim, rank))
|
|
else:
|
|
weights[f"{name}.lora_down.weight"] = mx.random.normal((rank, in_dim))
|
|
weights[f"{name}.lora_up.weight"] = mx.random.normal((out_dim, rank))
|
|
path = Path(tmp_dir) / "test_lora.safetensors"
|
|
mx.save_safetensors(str(path), weights)
|
|
return path
|
|
|
|
def test_load_lora_a_b_format(self):
|
|
from mlx_video.lora.loader import load_lora_weights
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = self._make_lora_file(tmp, ["blocks.0.self_attn.q"], key_format="AB")
|
|
lora_weights = load_lora_weights(path)
|
|
assert "blocks.0.self_attn.q" in lora_weights
|
|
w = lora_weights["blocks.0.self_attn.q"]
|
|
assert w.rank == 4
|
|
assert w.alpha == 4.0 # default: alpha == rank
|
|
assert w.lora_A.shape == (4, 64)
|
|
assert w.lora_B.shape == (128, 4)
|
|
|
|
def test_load_lora_down_up_format(self):
|
|
from mlx_video.lora.loader import load_lora_weights
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = self._make_lora_file(
|
|
tmp, ["blocks.0.self_attn.q"], key_format="down_up"
|
|
)
|
|
lora_weights = load_lora_weights(path)
|
|
assert "blocks.0.self_attn.q" in lora_weights
|
|
|
|
def test_load_multiple_modules(self):
|
|
from mlx_video.lora.loader import load_lora_weights
|
|
|
|
modules = [
|
|
"blocks.0.self_attn.q",
|
|
"blocks.0.self_attn.k",
|
|
"blocks.0.ffn.fc1",
|
|
]
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = self._make_lora_file(tmp, modules)
|
|
lora_weights = load_lora_weights(path)
|
|
assert len(lora_weights) == 3
|
|
for name in modules:
|
|
assert name in lora_weights
|
|
|
|
def test_load_with_alpha(self):
|
|
from mlx_video.lora.loader import load_lora_weights
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
weights = {
|
|
"test.lora_A.weight": mx.random.normal((8, 64)),
|
|
"test.lora_B.weight": mx.random.normal((128, 8)),
|
|
"test.alpha": mx.array(16.0),
|
|
}
|
|
path = Path(tmp) / "lora.safetensors"
|
|
mx.save_safetensors(str(path), weights)
|
|
lora_weights = load_lora_weights(path)
|
|
assert lora_weights["test"].alpha == 16.0
|
|
assert lora_weights["test"].rank == 8
|
|
assert lora_weights["test"].scale == 2.0
|
|
|
|
def test_file_not_found(self):
|
|
from mlx_video.lora.loader import load_lora_weights
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
load_lora_weights(Path("/nonexistent/lora.safetensors"))
|
|
|
|
|
|
class TestWanKeyNormalization:
|
|
"""Test Wan2.2 LoRA key normalization."""
|
|
|
|
def _wan_model_keys(self):
|
|
"""Simulate typical Wan2.2 MLX model weight keys."""
|
|
keys = set()
|
|
for i in range(2):
|
|
for layer in ["self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o",
|
|
"cross_attn.q", "cross_attn.k", "cross_attn.v", "cross_attn.o"]:
|
|
keys.add(f"blocks.{i}.{layer}.weight")
|
|
keys.add(f"blocks.{i}.ffn.fc1.weight")
|
|
keys.add(f"blocks.{i}.ffn.fc2.weight")
|
|
keys.add("text_embedding_0.weight")
|
|
keys.add("text_embedding_1.weight")
|
|
keys.add("time_embedding_0.weight")
|
|
keys.add("time_embedding_1.weight")
|
|
keys.add("time_projection.weight")
|
|
keys.add("patch_embedding_proj.weight")
|
|
return keys
|
|
|
|
def test_direct_match(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
assert _normalize_wan_lora_key("blocks.0.self_attn.q", keys) == "blocks.0.self_attn.q"
|
|
|
|
def test_strip_diffusion_model_prefix(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
result = _normalize_wan_lora_key("diffusion_model.blocks.0.self_attn.q", keys)
|
|
assert result == "blocks.0.self_attn.q"
|
|
|
|
def test_strip_model_prefix(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
result = _normalize_wan_lora_key("model.diffusion_model.blocks.0.self_attn.k", keys)
|
|
assert result == "blocks.0.self_attn.k"
|
|
|
|
def test_ffn_key_mapping(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
assert _normalize_wan_lora_key("blocks.0.ffn.0", keys) == "blocks.0.ffn.fc1"
|
|
assert _normalize_wan_lora_key("blocks.0.ffn.2", keys) == "blocks.0.ffn.fc2"
|
|
|
|
def test_text_embedding_mapping(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
assert _normalize_wan_lora_key("text_embedding.0", keys) == "text_embedding_0"
|
|
assert _normalize_wan_lora_key("text_embedding.2", keys) == "text_embedding_1"
|
|
|
|
def test_time_embedding_mapping(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
assert _normalize_wan_lora_key("time_embedding.0", keys) == "time_embedding_0"
|
|
assert _normalize_wan_lora_key("time_embedding.2", keys) == "time_embedding_1"
|
|
|
|
def test_time_projection_mapping(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
assert _normalize_wan_lora_key("time_projection.1", keys) == "time_projection"
|
|
|
|
def test_patch_embedding_mapping(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
assert _normalize_wan_lora_key("patch_embedding", keys) == "patch_embedding_proj"
|
|
|
|
def test_combined_prefix_and_ffn(self):
|
|
from mlx_video.lora.apply import _normalize_wan_lora_key
|
|
|
|
keys = self._wan_model_keys()
|
|
result = _normalize_wan_lora_key("diffusion_model.blocks.1.ffn.0", keys)
|
|
assert result == "blocks.1.ffn.fc1"
|
|
|
|
|
|
class TestApplyLoRA:
|
|
"""Test LoRA delta application to weights."""
|
|
|
|
def test_preserves_bfloat16_dtype(self):
|
|
"""LoRA delta must not promote bfloat16 weights to float32."""
|
|
from mlx_video.lora.apply import apply_lora_to_linear
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
original = mx.ones((8, 4), dtype=mx.bfloat16)
|
|
# LoRA weights in float32 (typical when loaded from safetensors)
|
|
lora_a = mx.ones((2, 4), dtype=mx.float32) * 0.1
|
|
lora_b = mx.ones((8, 2), dtype=mx.float32) * 0.1
|
|
w = LoRAWeights(lora_A=lora_a, lora_B=lora_b, rank=2, alpha=2.0, module_name="test")
|
|
result = apply_lora_to_linear(original, [(w, 1.0)])
|
|
assert result.dtype == mx.bfloat16, f"Expected bfloat16, got {result.dtype}"
|
|
|
|
def test_preserves_float16_dtype(self):
|
|
from mlx_video.lora.apply import apply_lora_to_linear
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
original = mx.ones((8, 4), dtype=mx.float16)
|
|
lora_a = mx.ones((2, 4), dtype=mx.float32) * 0.1
|
|
lora_b = mx.ones((8, 2), dtype=mx.float32) * 0.1
|
|
w = LoRAWeights(lora_A=lora_a, lora_B=lora_b, rank=2, alpha=2.0, module_name="test")
|
|
result = apply_lora_to_linear(original, [(w, 1.0)])
|
|
assert result.dtype == mx.float16, f"Expected float16, got {result.dtype}"
|
|
|
|
def test_apply_single_lora(self):
|
|
from mlx_video.lora.apply import apply_lora_to_linear
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
original = mx.ones((8, 4))
|
|
lora_a = mx.ones((2, 4)) * 0.1
|
|
lora_b = mx.ones((8, 2)) * 0.1
|
|
w = LoRAWeights(lora_A=lora_a, lora_B=lora_b, rank=2, alpha=2.0, module_name="test")
|
|
result = apply_lora_to_linear(original, [(w, 1.0)])
|
|
# delta = 1.0 * (B @ A) = ones(8,2)*0.1 @ ones(2,4)*0.1 = 0.02 * ones(8,4)
|
|
expected = original + 0.02 * mx.ones((8, 4))
|
|
assert mx.allclose(result, expected, atol=1e-6).item()
|
|
|
|
def test_apply_multiple_loras(self):
|
|
from mlx_video.lora.apply import apply_lora_to_linear
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
original = mx.zeros((8, 4))
|
|
w1 = LoRAWeights(
|
|
lora_A=mx.ones((2, 4)),
|
|
lora_B=mx.ones((8, 2)),
|
|
rank=2, alpha=2.0, module_name="a",
|
|
)
|
|
w2 = LoRAWeights(
|
|
lora_A=mx.ones((2, 4)) * 2,
|
|
lora_B=mx.ones((8, 2)) * 2,
|
|
rank=2, alpha=4.0, module_name="b",
|
|
)
|
|
result = apply_lora_to_linear(original, [(w1, 1.0), (w2, 0.5)])
|
|
# w1 delta: 1.0 * 1.0 * (ones(8,2) @ ones(2,4)) = 2 * ones(8,4)
|
|
# w2 delta: 2.0 * 0.5 * (2*ones(8,2) @ 2*ones(2,4)) = 1.0 * 8*ones(8,4) = 8
|
|
delta1 = mx.ones((8, 4)) * 2.0
|
|
delta2 = mx.ones((8, 4)) * 8.0
|
|
expected = delta1 + delta2
|
|
assert mx.allclose(result, expected, atol=1e-5).item()
|
|
|
|
def test_apply_loras_to_weights_dict(self):
|
|
from mlx_video.lora.apply import apply_loras_to_weights
|
|
from mlx_video.lora.types import LoRAWeights
|
|
|
|
model_weights = {
|
|
"blocks.0.self_attn.q.weight": mx.ones((128, 64)),
|
|
"blocks.0.self_attn.k.weight": mx.ones((128, 64)),
|
|
"blocks.0.ffn.fc1.weight": mx.ones((256, 64)),
|
|
}
|
|
w = LoRAWeights(
|
|
lora_A=mx.ones((4, 64)) * 0.01,
|
|
lora_B=mx.ones((128, 4)) * 0.01,
|
|
rank=4, alpha=4.0, module_name="blocks.0.self_attn.q",
|
|
)
|
|
module_to_loras = {"blocks.0.self_attn.q": [(w, 1.0)]}
|
|
result = apply_loras_to_weights(model_weights, module_to_loras)
|
|
# Only q should be modified
|
|
assert not mx.array_equal(
|
|
result["blocks.0.self_attn.q.weight"],
|
|
model_weights["blocks.0.self_attn.q.weight"],
|
|
).item()
|
|
assert mx.array_equal(
|
|
result["blocks.0.self_attn.k.weight"],
|
|
model_weights["blocks.0.self_attn.k.weight"],
|
|
).item()
|
|
|
|
|
|
class TestEndToEnd:
|
|
"""End-to-end LoRA loading and application."""
|
|
|
|
def test_load_and_apply_loras(self):
|
|
from mlx_video.convert_wan import load_and_apply_loras
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
# Create mock LoRA safetensors
|
|
rank = 4
|
|
weights = {
|
|
"blocks.0.self_attn.q.lora_A.weight": mx.random.normal((rank, 64)),
|
|
"blocks.0.self_attn.q.lora_B.weight": mx.random.normal((128, rank)),
|
|
}
|
|
lora_path = Path(tmp) / "test.safetensors"
|
|
mx.save_safetensors(str(lora_path), weights)
|
|
|
|
# Create mock model weights
|
|
model_weights = {
|
|
"blocks.0.self_attn.q.weight": mx.ones((128, 64)),
|
|
"blocks.0.self_attn.k.weight": mx.ones((128, 64)),
|
|
}
|
|
|
|
result = load_and_apply_loras(
|
|
model_weights, [(str(lora_path), 1.0)]
|
|
)
|
|
|
|
# q weight should be modified, k unchanged
|
|
assert not mx.array_equal(
|
|
result["blocks.0.self_attn.q.weight"],
|
|
model_weights["blocks.0.self_attn.q.weight"],
|
|
).item()
|
|
assert mx.array_equal(
|
|
result["blocks.0.self_attn.k.weight"],
|
|
model_weights["blocks.0.self_attn.k.weight"],
|
|
).item()
|