File size: 1,828 Bytes
14b57af
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import pytest
import shutil
import os


@pytest.fixture(scope="session")
def test_paths(request, pytestconfig):
    try:
        output_path = "output"
        ckpt_path = request.param  # This will get the current parameterized item
        text_encoder_model_name_or_path = pytestconfig.getoption(
            "text_encoder_model_name_or_path"
        )
        input_image_path = pytestconfig.getoption("input_image_path")
        input_video_path = pytestconfig.getoption("input_video_path")

        config = {
            "ckpt_path": ckpt_path,
            "input_image_path": input_image_path,
            "input_video_path": input_video_path,
            "output_path": output_path,
            "text_encoder_model_name_or_path": text_encoder_model_name_or_path,
        }

        yield config

    finally:
        if os.path.exists(output_path):
            shutil.rmtree(output_path)


def pytest_generate_tests(metafunc):
    if "test_paths" in metafunc.fixturenames:
        ckpt_paths = metafunc.config.getoption("ckpt_path")
        metafunc.parametrize("test_paths", ckpt_paths, indirect=True)


def pytest_addoption(parser):
    parser.addoption(
        "--ckpt_path",
        action="append",
        default=[],
        help="Path to checkpoint files (can specify multiple)",
    )
    parser.addoption(
        "--text_encoder_model_name_or_path",
        action="store",
        default="PixArt-alpha/PixArt-XL-2-1024-MS",
        help="Path to the checkpoint file",
    )
    parser.addoption(
        "--input_image_path",
        action="store",
        default="tests/utils/woman.jpeg",
        help="Path to input image file.",
    )
    parser.addoption(
        "--input_video_path",
        action="store",
        default="tests/utils/woman.mp4",
        help="Path to input video file.",
    )