58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
os.environ.setdefault('UNIT_TEST', 'True')
|
|
|
|
# Ensure optional test stubs are loaded
|
|
repo_root = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(repo_root))
|
|
try:
|
|
import sitecustomize # noqa: F401
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
if sys.path[0] == str(repo_root):
|
|
sys.path.pop(0)
|
|
sys.path.append(str(repo_root))
|
|
|
|
import sys
|
|
import types
|
|
import importlib.util
|
|
|
|
# Provide a minimal torchvision stub only when torchvision is not installed
|
|
if importlib.util.find_spec('torchvision') is None:
|
|
tv_stub = types.ModuleType('torchvision')
|
|
tv_stub.datasets = types.ModuleType('datasets')
|
|
tv_stub.models = types.ModuleType('models')
|
|
tv_stub.transforms = types.ModuleType('transforms')
|
|
tv_stub.transforms.ToTensor = lambda: lambda x: x
|
|
sys.modules['torchvision'] = tv_stub
|
|
|
|
if importlib.util.find_spec('PIL') is None:
|
|
pil_stub = types.ModuleType('PIL')
|
|
pil_stub.Image = types.SimpleNamespace(open=lambda *a, **k: None)
|
|
sys.modules['PIL'] = pil_stub
|
|
|
|
# Provide folder_paths stub if module unavailable
|
|
if 'folder_paths' not in sys.modules:
|
|
fp_stub = types.SimpleNamespace(
|
|
get_filename_list=lambda *a, **k: [],
|
|
get_input_directory=lambda: '',
|
|
get_output_directory=lambda: '',
|
|
get_temp_directory=lambda: ''
|
|
)
|
|
sys.modules['folder_paths'] = fp_stub
|
|
|
|
# Minimal torchaudio stub to satisfy imports
|
|
if importlib.util.find_spec('torchaudio') is None:
|
|
ta_stub = types.ModuleType('torchaudio')
|
|
def _load(*a, **k):
|
|
raise NotImplementedError('torchaudio.load stub')
|
|
ta_stub.load = _load
|
|
sys.modules['torchaudio'] = ta_stub
|
|
|
|
# Basic OpenCV stub
|
|
if importlib.util.find_spec('cv2') is None:
|
|
sys.modules['cv2'] = types.ModuleType('cv2')
|