88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
import importlib.util
|
|
import json
|
|
import pathlib
|
|
import sys
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_plugin_package():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"lls_node_test_upscale",
|
|
ROOT / "__init__.py",
|
|
submodule_search_locations=[str(ROOT)],
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules["lls_node_test_upscale"] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class RecordingUpscaleSwitcher:
|
|
def __init__(self, cls):
|
|
self._impl = cls()
|
|
self.calls = []
|
|
|
|
def _upscale_with_pytorch(self, image, scale, interpolation):
|
|
self.calls.append(("pytorch", image, scale, interpolation))
|
|
return ("PYTORCH_RESULT",)
|
|
|
|
def _upscale_with_model(self, image, model_name, tile, overlap):
|
|
self.calls.append(("model", image, model_name, tile, overlap))
|
|
return ("MODEL_RESULT",)
|
|
|
|
|
|
class TestUpscaleSwitcher(unittest.TestCase):
|
|
def test_schema_defaults_to_pytorch_when_no_upscale_models_exist(self):
|
|
load_plugin_package()
|
|
from lls_node_test_upscale.upscale import nodes as upscale_nodes
|
|
|
|
with mock.patch.object(
|
|
upscale_nodes,
|
|
"_get_upscale_model_names",
|
|
return_value=[upscale_nodes.NO_UPSCALE_MODEL_PLACEHOLDER],
|
|
):
|
|
schema = upscale_nodes.LLSUpscaleSwitcher.INPUT_TYPES()
|
|
|
|
self.assertEqual(schema["required"]["mode"][1]["default"], "interpolation")
|
|
|
|
def test_placeholder_model_name_falls_back_to_pytorch(self):
|
|
load_plugin_package()
|
|
from lls_node_test_upscale.upscale import nodes as upscale_nodes
|
|
|
|
recorder = RecordingUpscaleSwitcher(upscale_nodes.LLSUpscaleSwitcher)
|
|
with mock.patch.object(
|
|
upscale_nodes.LLSUpscaleSwitcher,
|
|
"_upscale_with_pytorch",
|
|
new=RecordingUpscaleSwitcher._upscale_with_pytorch,
|
|
), mock.patch.object(
|
|
upscale_nodes.LLSUpscaleSwitcher,
|
|
"_upscale_with_model",
|
|
new=RecordingUpscaleSwitcher._upscale_with_model,
|
|
):
|
|
image, upscale_info = upscale_nodes.LLSUpscaleSwitcher.upscale(
|
|
recorder,
|
|
image="image",
|
|
mode="upscale_model",
|
|
scale=2.0,
|
|
interpolation="bilinear",
|
|
model_name=upscale_nodes.NO_UPSCALE_MODEL_PLACEHOLDER,
|
|
tile=512,
|
|
overlap=32,
|
|
)
|
|
|
|
payload = json.loads(upscale_info)
|
|
self.assertEqual(image, "PYTORCH_RESULT")
|
|
self.assertEqual(payload["mode"], "interpolation")
|
|
self.assertEqual(
|
|
recorder.calls,
|
|
[("pytorch", "image", 2.0, "bilinear")],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|