28 lines
1.1 KiB
Python
28 lines
1.1 KiB
Python
import sys
|
|
import types
|
|
from pathlib import Path
|
|
import os
|
|
import importlib.util
|
|
|
|
# During unit tests we provide lightweight stubs for ComfyUI and the main package
|
|
if os.environ.get('UNIT_TEST'):
|
|
# Stub comfy package
|
|
import tests.mock_comfy as mock_comfy
|
|
sys.modules.setdefault('comfy', mock_comfy)
|
|
|
|
repo_root = Path(__file__).resolve().parent
|
|
pkg = types.ModuleType('EternalKernelLiteGraphNodes')
|
|
pkg.__path__ = [str(repo_root)]
|
|
sys.modules.setdefault('EternalKernelLiteGraphNodes', pkg)
|
|
custom_nodes = types.ModuleType('custom_nodes')
|
|
custom_nodes.EternalKernelLiteGraphNodes = pkg
|
|
sys.modules.setdefault('custom_nodes', custom_nodes)
|
|
sys.modules.setdefault('custom_nodes.EternalKernelLiteGraphNodes', pkg)
|
|
|
|
# Provide a very small torchvision stub if torchvision is missing
|
|
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')
|
|
sys.modules['torchvision'] = tv_stub
|