refiners/tests/conftest.py
Benjamin Trom 282578ddc0 add Segment Anything (SAM) to foundational models
Note: dense prompts (i.e. masks) support is still partial (see MaskEncoder)

Co-authored-by: Cédric Deltheil <cedric@deltheil.me>
2023-09-21 11:44:30 +02:00

36 lines
898 B
Python

import os
import torch
from pathlib import Path
from pytest import fixture
PARENT_PATH = Path(__file__).parent
@fixture(scope="session")
def test_device() -> torch.device:
test_device = os.getenv("REFINERS_TEST_DEVICE")
if not test_device:
return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
return torch.device(test_device)
@fixture(scope="session")
def test_weights_path() -> Path:
from_env = os.getenv("REFINERS_TEST_WEIGHTS_DIR")
return Path(from_env) if from_env else PARENT_PATH / "weights"
@fixture(scope="session")
def test_e2e_path() -> Path:
return PARENT_PATH / "e2e"
@fixture(scope="session")
def test_textual_inversion_path() -> Path:
return PARENT_PATH / "foundationals" / "clip" / "test_concepts_ref"
@fixture(scope="session")
def test_sam_path() -> Path:
return PARENT_PATH / "foundationals" / "segment_anything"