mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-12-15 19:53:22 +00:00
67f00b6f47
Implements a production-grade testing infrastructure with 100+ tests covering all core modules and pipelines. Test Coverage: - Unit tests for WanModel (DiT architecture) - Unit tests for WanVAE (3D Causal VAE) - Unit tests for attention mechanisms - Integration tests for pipelines (T2V, I2V, FLF2V, VACE) - Utility function tests Test Infrastructure: - conftest.py with reusable fixtures for configs, devices, and dtypes - pytest.ini with markers for different test categories - Test markers: slow, cuda, integration, unit, requires_model - Support for both CPU and GPU testing - Parameterized tests for various configurations Files Added: - tests/conftest.py - Pytest fixtures and configuration - tests/test_attention.py - Attention mechanism tests - tests/test_model.py - WanModel tests - tests/test_vae.py - VAE tests - tests/test_utils.py - Utility function tests - tests/test_pipelines.py - Pipeline integration tests - pytest.ini - Pytest configuration Test Execution: - pytest tests/ -v # Run all tests - pytest tests/ -m "not cuda" # CPU only - pytest tests/ -m "integration" # Integration tests only
133 lines
3.4 KiB
Python
133 lines
3.4 KiB
Python
"""
|
|
Pytest configuration and shared fixtures for Wan2.1 tests.
|
|
|
|
Copyright (c) 2025 Kuaishou. All rights reserved.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Dict, Any
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def device():
|
|
"""Return the device to use for testing (CPU or CUDA if available)."""
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def dtype():
|
|
"""Return the default dtype for testing."""
|
|
return torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir():
|
|
"""Create a temporary directory for test files."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
yield Path(tmpdir)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_config_14b() -> Dict[str, Any]:
|
|
"""Return a minimal 14B model configuration for testing."""
|
|
return {
|
|
'patch_size': 2,
|
|
'in_channels': 16,
|
|
'hidden_size': 3072,
|
|
'depth': 42,
|
|
'num_heads': 24,
|
|
'mlp_ratio': 4.0,
|
|
'learn_sigma': True,
|
|
'qk_norm': True,
|
|
'qk_norm_type': 'rms',
|
|
'norm_type': 'rms',
|
|
'posemb_type': 'rope2d_video',
|
|
'num_experts': 1,
|
|
'route_method': 'soft',
|
|
'router_top_k': 1,
|
|
'pooled_projection_type': 'linear',
|
|
'cap_feat_dim': 4096,
|
|
'caption_channels': 4096,
|
|
't5_feat_dim': 2048,
|
|
'text_len': 512,
|
|
'use_attention_mask': True,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_config_1_3b() -> Dict[str, Any]:
|
|
"""Return a minimal 1.3B model configuration for testing."""
|
|
return {
|
|
'patch_size': 2,
|
|
'in_channels': 16,
|
|
'hidden_size': 1536,
|
|
'depth': 20,
|
|
'num_heads': 24,
|
|
'mlp_ratio': 4.0,
|
|
'learn_sigma': True,
|
|
'qk_norm': True,
|
|
'qk_norm_type': 'rms',
|
|
'norm_type': 'rms',
|
|
'posemb_type': 'rope2d_video',
|
|
'num_experts': 1,
|
|
'route_method': 'soft',
|
|
'router_top_k': 1,
|
|
'pooled_projection_type': 'linear',
|
|
'cap_feat_dim': 4096,
|
|
'caption_channels': 4096,
|
|
't5_feat_dim': 2048,
|
|
'text_len': 512,
|
|
'use_attention_mask': True,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_vae_config() -> Dict[str, Any]:
|
|
"""Return a minimal VAE configuration for testing."""
|
|
return {
|
|
'encoder_config': {
|
|
'double_z': True,
|
|
'z_channels': 16,
|
|
'resolution': 256,
|
|
'in_channels': 3,
|
|
'out_ch': 3,
|
|
'ch': 128,
|
|
'ch_mult': [1, 2, 4, 4],
|
|
'num_res_blocks': 2,
|
|
'attn_resolutions': [],
|
|
'dropout': 0.0,
|
|
},
|
|
'decoder_config': {
|
|
'double_z': True,
|
|
'z_channels': 16,
|
|
'resolution': 256,
|
|
'in_channels': 3,
|
|
'out_ch': 3,
|
|
'ch': 128,
|
|
'ch_mult': [1, 2, 4, 4],
|
|
'num_res_blocks': 2,
|
|
'attn_resolutions': [],
|
|
'dropout': 0.0,
|
|
},
|
|
'temporal_compress_level': 4,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def skip_if_no_cuda():
|
|
"""Skip test if CUDA is not available."""
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
|
|
|
|
@pytest.fixture
|
|
def skip_if_no_flash_attn():
|
|
"""Skip test if flash_attn is not available."""
|
|
try:
|
|
import flash_attn
|
|
except ImportError:
|
|
pytest.skip("flash_attn not available")
|