修改clip报错完成,修改upscale模块
This commit is contained in:
@@ -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()
|
||||
+16
-1
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user