diff --git a/tests/test_upscale_switcher.py b/tests/test_upscale_switcher.py new file mode 100644 index 0000000..6481c8c --- /dev/null +++ b/tests/test_upscale_switcher.py @@ -0,0 +1,84 @@ +import importlib.util +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"], "pytorch") + + 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, + ): + result = 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, + ) + + self.assertEqual(result, ("PYTORCH_RESULT",)) + self.assertEqual( + recorder.calls, + [("pytorch", "image", 2.0, "bilinear")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/upscale/nodes.py b/upscale/nodes.py index 0a488d5..1067e1f 100644 --- a/upscale/nodes.py +++ b/upscale/nodes.py @@ -73,6 +73,14 @@ def _get_upscale_model_names() -> list[str]: return names if names else [NO_UPSCALE_MODEL_PLACEHOLDER] +def _has_real_upscale_models(model_names: list[str]) -> bool: + return any(name and name != NO_UPSCALE_MODEL_PLACEHOLDER for name in model_names) + + +def _is_missing_upscale_model_selection(model_name: str | None) -> bool: + return model_name in (None, "", NO_UPSCALE_MODEL_PLACEHOLDER) + + def _resolve_upscale_model_path(model_name: str) -> str: _require_comfy_modules() if model_name == NO_UPSCALE_MODEL_PLACEHOLDER: @@ -212,10 +220,11 @@ class LLSUpscaleSwitcher: @classmethod def INPUT_TYPES(cls): model_names = _sorted_unique(_get_upscale_model_names()) + default_mode = "upscale_model" if _has_real_upscale_models(model_names) else "pytorch" return { "required": { "image": ("IMAGE",), - "mode": (UPSCALE_MODE_CHOICES, {"default": "upscale_model"}), + "mode": (UPSCALE_MODE_CHOICES, {"default": default_mode}), "scale": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1}), "interpolation": (INTERPOLATION_CHOICES, {"default": "bilinear"}), "model_name": (model_names,), @@ -229,6 +238,12 @@ class LLSUpscaleSwitcher: raise RuntimeError(f"Unsupported mode '{mode}'. Expected one of {UPSCALE_MODE_CHOICES}.") if mode == "pytorch": return self._upscale_with_pytorch(image=image, scale=scale, interpolation=interpolation) + if _is_missing_upscale_model_selection(model_name): + print( + "[LLS] WARNING: No valid upscale model is selected. " + "Falling back to PyTorch interpolation mode." + ) + return self._upscale_with_pytorch(image=image, scale=scale, interpolation=interpolation) return self._upscale_with_model(image=image, model_name=model_name, tile=tile, overlap=overlap) def _upscale_with_pytorch(self, image, scale: float, interpolation: str):