78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Regression tests for preserving MLX compile flags on raw Diffusers loads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
mx = pytest.importorskip("mlx.core", reason="MLX required for compile dispatch")
|
|
|
|
|
|
def _write_tiny_diffusers_checkpoint(tmp_path: Path) -> tuple[Path, Path]:
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
config = {
|
|
"num_layers": 0,
|
|
"num_attention_heads": 1,
|
|
"attention_head_dim": 2,
|
|
"ffn_dim": 4,
|
|
"in_channels": 16,
|
|
"out_channels": 16,
|
|
"patch_size": [1, 2, 2],
|
|
"freq_dim": 2,
|
|
"eps": 1e-6,
|
|
}
|
|
config_path = tmp_path / "config.json"
|
|
checkpoint_path = tmp_path / "diffusion_pytorch_model.safetensors"
|
|
config_path.write_text(json.dumps(config))
|
|
|
|
hidden = 2
|
|
weights = {
|
|
"patch_embedding.weight": torch.zeros(hidden, 64),
|
|
"condition_embedder.time_embedder.linear_1.weight": torch.zeros(hidden, 2),
|
|
"condition_embedder.time_embedder.linear_2.weight": torch.zeros(hidden, hidden),
|
|
"condition_embedder.time_proj.weight": torch.zeros(hidden * 6, hidden),
|
|
"condition_embedder.text_embedder.linear_1.weight": torch.zeros(hidden, 4096),
|
|
"condition_embedder.text_embedder.linear_2.weight": torch.zeros(hidden, hidden),
|
|
"scale_shift_table": torch.zeros(2, hidden),
|
|
"proj_out.weight": torch.zeros(64, hidden),
|
|
}
|
|
save_file(weights, checkpoint_path)
|
|
return checkpoint_path, config_path
|
|
|
|
|
|
def test_diffusers_loader_preserves_compile_flag(tmp_path: Path) -> None:
|
|
from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
|
|
|
|
checkpoint_path, config_path = _write_tiny_diffusers_checkpoint(tmp_path)
|
|
|
|
compiled = mlx_dit_from_diffusers_safetensors(
|
|
checkpoint_path,
|
|
config_path,
|
|
compile=True,
|
|
num_blocks=0,
|
|
quantization=None,
|
|
)
|
|
eager = mlx_dit_from_diffusers_safetensors(
|
|
checkpoint_path,
|
|
config_path,
|
|
compile=False,
|
|
num_blocks=0,
|
|
quantization=None,
|
|
)
|
|
|
|
assert compiled._enable_compile is True
|
|
assert eager._enable_compile is False
|
|
|
|
|
|
def test_prompt_to_video_cli_passes_compile_to_diffusers_loader() -> None:
|
|
root = Path(__file__).resolve().parents[3]
|
|
script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py"
|
|
source = script.read_text()
|
|
|
|
assert "compile=args.mlx_compile" in source
|