52 lines
2.3 KiB
Python
52 lines
2.3 KiB
Python
import os
|
|
import sys
|
|
import torch
|
|
import types
|
|
|
|
###############################################################################
|
|
# 1) Patch torch.load for data/<filename>.pth so base code's relative paths
|
|
# are mapped to absolute paths in your local "data" folder.
|
|
###############################################################################
|
|
orig_torch_load = torch.load
|
|
REMAP_PATHS = {
|
|
"data/vgg19_gray.pth": os.path.join(os.path.dirname(__file__), "data", "vgg19_gray.pth"),
|
|
}
|
|
def custom_torch_load(path, *args, **kwargs):
|
|
if path in REMAP_PATHS:
|
|
path = REMAP_PATHS[path]
|
|
return orig_torch_load(path, *args, **kwargs)
|
|
torch.load = custom_torch_load
|
|
|
|
###############################################################################
|
|
# 2) Insert this custom node root directory into sys.path if not already present
|
|
###############################################################################
|
|
custom_node_root = os.path.dirname(os.path.abspath(__file__))
|
|
if custom_node_root not in sys.path:
|
|
sys.path.insert(0, custom_node_root)
|
|
|
|
###############################################################################
|
|
# 3) Remove any pre-existing "utils" or "models" entries from sys.modules
|
|
###############################################################################
|
|
for modname in ["utils", "models"]:
|
|
if modname in sys.modules:
|
|
del sys.modules[modname]
|
|
|
|
###############################################################################
|
|
# 4) Dynamically create "utils" and "models" packages, pointing them to our local directories.
|
|
###############################################################################
|
|
utils_dir = os.path.join(custom_node_root, "utils")
|
|
models_dir = os.path.join(custom_node_root, "models")
|
|
if os.path.isdir(utils_dir):
|
|
utils_module = types.ModuleType("utils")
|
|
utils_module.__path__ = [utils_dir]
|
|
sys.modules["utils"] = utils_module
|
|
if os.path.isdir(models_dir):
|
|
models_module = types.ModuleType("models")
|
|
models_module.__path__ = [models_dir]
|
|
sys.modules["models"] = models_module
|
|
|
|
###############################################################################
|
|
# 5) Import and register node mappings.
|
|
###############################################################################
|
|
from .DeepExemplarColorizationNodes import NODE_CLASS_MAPPINGS
|