修改clip报错完成,修改upscale模块

This commit is contained in:
Gin3601
2026-05-19 15:58:28 +08:00
parent f6b9ddfdcf
commit fa36d0329f
2 changed files with 100 additions and 1 deletions
+84
View File
@@ -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
View File
@@ -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):