Remove unused files and fix some bugs
This commit is contained in:
Gaohan123 committed
1 parent
6248bbf646
commit
6cf73c4e80
21 files changed
+548
-1891
No files matched your search
@@ -1,86 +0,0 @@
|
||||
"""
|
||||
Shared fixtures and configuration for vLLM-omni tests.
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.config import (
|
||||
DiTCacheConfig,
|
||||
DiTConfig,
|
||||
create_ar_stage_config,
|
||||
create_dit_stage_config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def device():
|
||||
"""Get available device for testing."""
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_ar_stage_config():
|
||||
"""Sample AR stage configuration for testing."""
|
||||
return create_ar_stage_config(
|
||||
stage_id=0,
|
||||
model_path="test-ar-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_dit_stage_config():
|
||||
"""Sample DiT stage configuration for testing."""
|
||||
dit_config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=10,
|
||||
guidance_scale=7.5,
|
||||
)
|
||||
|
||||
return create_dit_stage_config(
|
||||
stage_id=1,
|
||||
model_path="test-dit-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"],
|
||||
dit_config=dit_config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_stage_configs(sample_ar_stage_config, sample_dit_stage_config):
|
||||
"""Sample stage configurations for testing."""
|
||||
return [sample_ar_stage_config, sample_dit_stage_config]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vllm_config():
|
||||
"""Mock vLLM configuration."""
|
||||
config = Mock()
|
||||
config.model = "test-model"
|
||||
config.tensor_parallel_size = 1
|
||||
config.pipeline_parallel_size = 1
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_dit_cache_config():
|
||||
"""Mock DiT cache configuration."""
|
||||
from vllm_omni.config import DiTCacheTensor
|
||||
|
||||
cache_tensors = [
|
||||
DiTCacheTensor(
|
||||
name="test_tensor", shape=[1, 512, 512], dtype="float32", persistent=True
|
||||
)
|
||||
]
|
||||
|
||||
return DiTCacheConfig(
|
||||
cache_tensors=cache_tensors,
|
||||
max_cache_size=1024 * 1024 * 1024, # 1GB
|
||||
cache_strategy="fifo",
|
||||
enable_optimization=True,
|
||||
)
|
||||
@@ -1,160 +0,0 @@
|
||||
"""
|
||||
Unit tests for configuration modules.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config import (
|
||||
DiTCacheConfig,
|
||||
DiTCacheTensor,
|
||||
DiTConfig,
|
||||
OmniStageConfig,
|
||||
create_ar_stage_config,
|
||||
create_dit_stage_config,
|
||||
)
|
||||
|
||||
|
||||
class TestOmniStageConfig:
|
||||
def test_ar_stage_config_creation(self):
|
||||
"""Test AR stage configuration creation."""
|
||||
config = OmniStageConfig(
|
||||
stage_id=0,
|
||||
engine_type="AR",
|
||||
model_path="test-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
)
|
||||
assert config.engine_type == "AR"
|
||||
assert config.stage_id == 0
|
||||
assert config.model_path == "test-model"
|
||||
assert config.input_modalities == ["text"]
|
||||
assert config.output_modalities == ["text"]
|
||||
|
||||
def test_dit_stage_config_creation(self):
|
||||
"""Test DiT stage configuration creation."""
|
||||
dit_config = DiTConfig(num_inference_steps=50)
|
||||
config = OmniStageConfig(
|
||||
stage_id=1,
|
||||
engine_type="DiT",
|
||||
model_path="test-dit-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"],
|
||||
dit_config=dit_config,
|
||||
)
|
||||
assert config.engine_type == "DiT"
|
||||
assert config.dit_config is not None
|
||||
assert config.dit_config.num_inference_steps == 50
|
||||
|
||||
def test_invalid_engine_type(self):
|
||||
"""Test validation of engine type."""
|
||||
with pytest.raises(ValueError):
|
||||
OmniStageConfig(
|
||||
stage_id=0,
|
||||
engine_type="INVALID",
|
||||
model_path="test-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
)
|
||||
|
||||
def test_empty_modalities(self):
|
||||
"""Test validation of empty modalities."""
|
||||
with pytest.raises(ValueError):
|
||||
OmniStageConfig(
|
||||
stage_id=0,
|
||||
engine_type="AR",
|
||||
model_path="test-model",
|
||||
input_modalities=[],
|
||||
output_modalities=["text"],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
OmniStageConfig(
|
||||
stage_id=0,
|
||||
engine_type="AR",
|
||||
model_path="test-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=[],
|
||||
)
|
||||
|
||||
def test_dit_requires_config(self):
|
||||
"""Test that DiT engine requires dit_config."""
|
||||
with pytest.raises(ValueError):
|
||||
OmniStageConfig(
|
||||
stage_id=1,
|
||||
engine_type="DiT",
|
||||
model_path="test-dit-model",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"],
|
||||
)
|
||||
|
||||
|
||||
class TestDiTConfig:
|
||||
def test_dit_config_defaults(self):
|
||||
"""Test DiT configuration with defaults."""
|
||||
config = DiTConfig(num_inference_steps=50)
|
||||
assert config.use_diffusers is False
|
||||
assert config.diffusers_pipeline is None
|
||||
assert config.guidance_scale == 7.5
|
||||
assert config.height == 512
|
||||
assert config.width == 512
|
||||
|
||||
def test_diffusers_config(self):
|
||||
"""Test DiT configuration with diffusers."""
|
||||
config = DiTConfig(
|
||||
num_inference_steps=50,
|
||||
use_diffusers=True,
|
||||
diffusers_pipeline="stable-diffusion",
|
||||
)
|
||||
assert config.use_diffusers is True
|
||||
assert config.diffusers_pipeline == "stable-diffusion"
|
||||
|
||||
|
||||
class TestDiTCacheConfig:
|
||||
def test_cache_config_creation(self):
|
||||
"""Test cache configuration creation."""
|
||||
cache_tensors = [
|
||||
DiTCacheTensor(
|
||||
name="test_tensor",
|
||||
shape=[1, 512, 512],
|
||||
dtype="float32",
|
||||
persistent=True,
|
||||
)
|
||||
]
|
||||
|
||||
config = DiTCacheConfig(
|
||||
cache_tensors=cache_tensors,
|
||||
max_cache_size=1024 * 1024 * 1024,
|
||||
cache_strategy="fifo",
|
||||
)
|
||||
|
||||
assert len(config.cache_tensors) == 1
|
||||
assert config.max_cache_size == 1024 * 1024 * 1024
|
||||
assert config.cache_strategy == "fifo"
|
||||
assert config.enable_optimization is True
|
||||
|
||||
|
||||
class TestHelperFunctions:
|
||||
def test_create_ar_stage_config(self):
|
||||
"""Test create_ar_stage_config helper function."""
|
||||
config = create_ar_stage_config(stage_id=0, model_path="test-model")
|
||||
|
||||
assert config.stage_id == 0
|
||||
assert config.engine_type == "AR"
|
||||
assert config.model_path == "test-model"
|
||||
assert config.input_modalities == ["text"]
|
||||
assert config.output_modalities == ["text"]
|
||||
|
||||
def test_create_dit_stage_config(self):
|
||||
"""Test create_dit_stage_config helper function."""
|
||||
dit_config = DiTConfig(num_inference_steps=50)
|
||||
|
||||
config = create_dit_stage_config(
|
||||
stage_id=1, model_path="test-dit-model", dit_config=dit_config
|
||||
)
|
||||
|
||||
assert config.stage_id == 1
|
||||
assert config.engine_type == "DiT"
|
||||
assert config.model_path == "test-dit-model"
|
||||
assert config.input_modalities == ["text"]
|
||||
assert config.output_modalities == ["image"]
|
||||
assert config.dit_config is not None
|
||||
Reference in new issue
Block a user