mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-12-15 11:43:21 +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
221 lines
7.6 KiB
Python
221 lines
7.6 KiB
Python
"""
|
|
Unit tests for WanVAE in Wan2.1.
|
|
|
|
Copyright (c) 2025 Kuaishou. All rights reserved.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
from wan.modules.vae import WanVAE_
|
|
|
|
|
|
class TestWanVAE:
|
|
"""Test suite for WanVAE (3D Causal VAE)."""
|
|
|
|
def test_vae_initialization(self, sample_vae_config):
|
|
"""Test VAE initialization."""
|
|
with torch.device('meta'):
|
|
vae = WanVAE_(**sample_vae_config)
|
|
|
|
assert vae is not None
|
|
assert hasattr(vae, 'encoder')
|
|
assert hasattr(vae, 'decoder')
|
|
assert vae.temporal_compress_level == sample_vae_config['temporal_compress_level']
|
|
|
|
def test_vae_encode_shape(self, sample_vae_config, device, dtype):
|
|
"""Test VAE encoding produces correct output shape."""
|
|
# Use smaller config for faster testing
|
|
config = sample_vae_config.copy()
|
|
config['encoder_config']['ch'] = 32
|
|
config['encoder_config']['ch_mult'] = [1, 2]
|
|
config['encoder_config']['num_res_blocks'] = 1
|
|
config['decoder_config']['ch'] = 32
|
|
config['decoder_config']['ch_mult'] = [1, 2]
|
|
config['decoder_config']['num_res_blocks'] = 1
|
|
|
|
vae = WanVAE_(**config).to(device).to(dtype)
|
|
vae.eval()
|
|
|
|
batch_size = 1
|
|
channels = 3
|
|
num_frames = 8
|
|
height = 64
|
|
width = 64
|
|
|
|
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
|
|
|
|
with torch.no_grad():
|
|
encoded = vae.encode(x)
|
|
|
|
# Check output shape after encoding
|
|
z_channels = config['encoder_config']['z_channels']
|
|
temporal_compress = config['temporal_compress_level']
|
|
spatial_compress = 2 ** (len(config['encoder_config']['ch_mult']) - 1)
|
|
|
|
expected_t = num_frames // temporal_compress
|
|
expected_h = height // spatial_compress
|
|
expected_w = width // spatial_compress
|
|
|
|
assert encoded.shape == (batch_size, z_channels, expected_t, expected_h, expected_w)
|
|
|
|
def test_vae_decode_shape(self, sample_vae_config, device, dtype):
|
|
"""Test VAE decoding produces correct output shape."""
|
|
config = sample_vae_config.copy()
|
|
config['encoder_config']['ch'] = 32
|
|
config['encoder_config']['ch_mult'] = [1, 2]
|
|
config['encoder_config']['num_res_blocks'] = 1
|
|
config['decoder_config']['ch'] = 32
|
|
config['decoder_config']['ch_mult'] = [1, 2]
|
|
config['decoder_config']['num_res_blocks'] = 1
|
|
|
|
vae = WanVAE_(**config).to(device).to(dtype)
|
|
vae.eval()
|
|
|
|
batch_size = 1
|
|
z_channels = config['encoder_config']['z_channels']
|
|
num_frames = 2
|
|
height = 32
|
|
width = 32
|
|
|
|
z = torch.randn(batch_size, z_channels, num_frames, height, width, device=device, dtype=dtype)
|
|
|
|
with torch.no_grad():
|
|
decoded = vae.decode(z)
|
|
|
|
# Check output shape after decoding
|
|
out_channels = config['decoder_config']['out_ch']
|
|
temporal_compress = config['temporal_compress_level']
|
|
spatial_compress = 2 ** (len(config['decoder_config']['ch_mult']) - 1)
|
|
|
|
expected_t = num_frames * temporal_compress
|
|
expected_h = height * spatial_compress
|
|
expected_w = width * spatial_compress
|
|
|
|
assert decoded.shape == (batch_size, out_channels, expected_t, expected_h, expected_w)
|
|
|
|
def test_vae_encode_decode_consistency(self, sample_vae_config, device, dtype):
|
|
"""Test that encode then decode produces similar output."""
|
|
config = sample_vae_config.copy()
|
|
config['encoder_config']['ch'] = 32
|
|
config['encoder_config']['ch_mult'] = [1, 2]
|
|
config['encoder_config']['num_res_blocks'] = 1
|
|
config['decoder_config']['ch'] = 32
|
|
config['decoder_config']['ch_mult'] = [1, 2]
|
|
config['decoder_config']['num_res_blocks'] = 1
|
|
|
|
vae = WanVAE_(**config).to(device).to(dtype)
|
|
vae.eval()
|
|
|
|
batch_size = 1
|
|
channels = 3
|
|
num_frames = 8
|
|
height = 64
|
|
width = 64
|
|
|
|
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
|
|
|
|
with torch.no_grad():
|
|
encoded = vae.encode(x)
|
|
decoded = vae.decode(encoded)
|
|
|
|
# Decoded output should have same shape as input
|
|
assert decoded.shape == x.shape
|
|
|
|
def test_vae_no_nan_encode(self, sample_vae_config, device, dtype):
|
|
"""Test that VAE encoding doesn't produce NaN values."""
|
|
config = sample_vae_config.copy()
|
|
config['encoder_config']['ch'] = 32
|
|
config['encoder_config']['ch_mult'] = [1, 2]
|
|
config['encoder_config']['num_res_blocks'] = 1
|
|
config['decoder_config']['ch'] = 32
|
|
config['decoder_config']['ch_mult'] = [1, 2]
|
|
config['decoder_config']['num_res_blocks'] = 1
|
|
|
|
vae = WanVAE_(**config).to(device).to(dtype)
|
|
vae.eval()
|
|
|
|
batch_size = 1
|
|
channels = 3
|
|
num_frames = 8
|
|
height = 64
|
|
width = 64
|
|
|
|
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
|
|
|
|
with torch.no_grad():
|
|
encoded = vae.encode(x)
|
|
|
|
assert not torch.isnan(encoded).any()
|
|
assert not torch.isinf(encoded).any()
|
|
|
|
def test_vae_no_nan_decode(self, sample_vae_config, device, dtype):
|
|
"""Test that VAE decoding doesn't produce NaN values."""
|
|
config = sample_vae_config.copy()
|
|
config['encoder_config']['ch'] = 32
|
|
config['encoder_config']['ch_mult'] = [1, 2]
|
|
config['encoder_config']['num_res_blocks'] = 1
|
|
config['decoder_config']['ch'] = 32
|
|
config['decoder_config']['ch_mult'] = [1, 2]
|
|
config['decoder_config']['num_res_blocks'] = 1
|
|
|
|
vae = WanVAE_(**config).to(device).to(dtype)
|
|
vae.eval()
|
|
|
|
batch_size = 1
|
|
z_channels = config['encoder_config']['z_channels']
|
|
num_frames = 2
|
|
height = 32
|
|
width = 32
|
|
|
|
z = torch.randn(batch_size, z_channels, num_frames, height, width, device=device, dtype=dtype)
|
|
|
|
with torch.no_grad():
|
|
decoded = vae.decode(z)
|
|
|
|
assert not torch.isnan(decoded).any()
|
|
assert not torch.isinf(decoded).any()
|
|
|
|
@pytest.mark.parametrize("num_frames", [4, 8, 16])
|
|
def test_vae_various_frame_counts(self, sample_vae_config, device, dtype, num_frames):
|
|
"""Test VAE with various frame counts."""
|
|
config = sample_vae_config.copy()
|
|
config['encoder_config']['ch'] = 32
|
|
config['encoder_config']['ch_mult'] = [1, 2]
|
|
config['encoder_config']['num_res_blocks'] = 1
|
|
config['decoder_config']['ch'] = 32
|
|
config['decoder_config']['ch_mult'] = [1, 2]
|
|
config['decoder_config']['num_res_blocks'] = 1
|
|
|
|
vae = WanVAE_(**config).to(device).to(dtype)
|
|
vae.eval()
|
|
|
|
batch_size = 1
|
|
channels = 3
|
|
height = 64
|
|
width = 64
|
|
|
|
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
|
|
|
|
with torch.no_grad():
|
|
encoded = vae.encode(x)
|
|
decoded = vae.decode(encoded)
|
|
|
|
assert decoded.shape == x.shape
|
|
assert not torch.isnan(decoded).any()
|
|
|
|
def test_vae_eval_mode(self, sample_vae_config):
|
|
"""Test that VAE can be set to eval mode."""
|
|
with torch.device('meta'):
|
|
vae = WanVAE_(**sample_vae_config)
|
|
|
|
vae.eval()
|
|
assert not vae.training
|
|
|
|
def test_vae_config_attributes(self, sample_vae_config):
|
|
"""Test that VAE has correct configuration attributes."""
|
|
with torch.device('meta'):
|
|
vae = WanVAE_(**sample_vae_config)
|
|
|
|
assert hasattr(vae, 'temporal_compress_level')
|
|
assert vae.temporal_compress_level == sample_vae_config['temporal_compress_level']
|