1297 lines
48 KiB
Python
1297 lines
48 KiB
Python
import asyncio
|
|
import base64
|
|
import importlib.util
|
|
import io
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def _record_matting_responses(monkeypatch, node_module):
|
|
responses = []
|
|
|
|
def json_response(payload, **kwargs):
|
|
response = SimpleNamespace(payload=payload, status=kwargs.get("status", 200))
|
|
responses.append(response)
|
|
return response
|
|
|
|
monkeypatch.setattr(node_module.web, "json_response", json_response, raising=False)
|
|
return responses
|
|
|
|
|
|
def _run_matting_model_check(node_module, model_path=None):
|
|
query = {} if model_path is None else {"model_path": model_path}
|
|
request = SimpleNamespace(query=query)
|
|
return asyncio.run(node_module.check_matting_model(request))
|
|
|
|
|
|
class RecordingRoutes:
|
|
"""Minimal PromptServer route registry for integration tests."""
|
|
|
|
def __init__(self):
|
|
self.registered = {}
|
|
|
|
def _decorator(self, method, path):
|
|
def register(handler):
|
|
self.registered[(method, path)] = handler
|
|
return handler
|
|
|
|
return register
|
|
|
|
def get(self, path):
|
|
return self._decorator("GET", path)
|
|
|
|
def post(self, path):
|
|
return self._decorator("POST", path)
|
|
|
|
|
|
def _install_runtime_stubs(monkeypatch, tmp_path):
|
|
routes = RecordingRoutes()
|
|
|
|
folder_paths = ModuleType("folder_paths")
|
|
folder_paths.models_dir = str(tmp_path / "models")
|
|
monkeypatch.setitem(sys.modules, "folder_paths", folder_paths)
|
|
|
|
server = ModuleType("server")
|
|
server.PromptServer = SimpleNamespace(
|
|
instance=SimpleNamespace(routes=routes, send_sync=lambda *args, **kwargs: None)
|
|
)
|
|
monkeypatch.setitem(sys.modules, "server", server)
|
|
|
|
aiohttp = ModuleType("aiohttp")
|
|
aiohttp.web = SimpleNamespace()
|
|
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp)
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
class ToTensor:
|
|
def __call__(self, image):
|
|
array = np.array(image, copy=True)
|
|
tensor = torch.from_numpy(array).float() / 255.0
|
|
if tensor.dim() == 2:
|
|
return tensor
|
|
return tensor.permute(2, 0, 1)
|
|
|
|
torchvision = ModuleType("torchvision")
|
|
torchvision.transforms = SimpleNamespace(ToTensor=ToTensor)
|
|
monkeypatch.setitem(sys.modules, "torchvision", torchvision)
|
|
|
|
tqdm = ModuleType("tqdm")
|
|
tqdm.tqdm = lambda iterable, *args, **kwargs: iterable
|
|
monkeypatch.setitem(sys.modules, "tqdm", tqdm)
|
|
|
|
# Avoid creating project log files while the entry point is imported.
|
|
from python.log_system.logger import logger
|
|
|
|
monkeypatch.setattr(logger, "configure", lambda config: logger)
|
|
|
|
return routes
|
|
|
|
|
|
def _import_layerforge(monkeypatch, tmp_path):
|
|
routes = _install_runtime_stubs(monkeypatch, tmp_path)
|
|
package_name = "_layerforge_test_package"
|
|
package_path = PROJECT_ROOT / "__init__.py"
|
|
spec = importlib.util.spec_from_file_location(
|
|
package_name,
|
|
package_path,
|
|
submodule_search_locations=[str(PROJECT_ROOT)],
|
|
)
|
|
package = importlib.util.module_from_spec(spec)
|
|
sys.modules[package_name] = package
|
|
assert spec.loader is not None
|
|
spec.loader.exec_module(package)
|
|
|
|
return SimpleNamespace(
|
|
entrypoint=package,
|
|
node=sys.modules[f"{package_name}.python.node"],
|
|
image_utils=sys.modules[f"{package_name}.python.image_utils"],
|
|
matting_api=sys.modules[f"{package_name}.python.matting.api"],
|
|
matting_birefnet=sys.modules[f"{package_name}.python.matting.backends.birefnet"],
|
|
matting_rmbg=sys.modules[f"{package_name}.python.matting.backends.rmbg"],
|
|
matting_catalog=sys.modules[f"{package_name}.python.matting.catalog"],
|
|
matting_options=sys.modules[f"{package_name}.python.matting.options"],
|
|
matting_paths=sys.modules[f"{package_name}.python.matting.paths"],
|
|
matting_settings=sys.modules[f"{package_name}.python.matting.settings"],
|
|
matting_service=sys.modules[f"{package_name}.python.matting.service"],
|
|
matting_progress=sys.modules[f"{package_name}.python.matting.progress"],
|
|
routes=routes,
|
|
package_name=package_name,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def layerforge_runtime(monkeypatch, tmp_path):
|
|
runtime = _import_layerforge(monkeypatch, tmp_path)
|
|
yield runtime
|
|
|
|
for module_name in list(sys.modules):
|
|
if module_name == runtime.package_name or module_name.startswith(f"{runtime.package_name}."):
|
|
sys.modules.pop(module_name, None)
|
|
|
|
|
|
def test_entrypoint_exports_node_contract_and_frontend_directory(layerforge_runtime):
|
|
entrypoint = layerforge_runtime.entrypoint
|
|
node_class = entrypoint.NODE_CLASS_MAPPINGS["LayerForgeNode"]
|
|
|
|
assert entrypoint.WEB_DIRECTORY == "./js"
|
|
assert entrypoint.__all__ == [
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"WEB_DIRECTORY",
|
|
]
|
|
assert node_class.RETURN_TYPES == ("IMAGE", "MASK")
|
|
assert node_class.RETURN_NAMES == ("image", "mask")
|
|
assert node_class.FUNCTION == "process_canvas_image"
|
|
assert node_class.CATEGORY == "azNodes > LayerForge"
|
|
|
|
inputs = node_class.INPUT_TYPES()
|
|
assert set(inputs) == {"required", "optional", "hidden"}
|
|
assert {
|
|
"fit_on_add",
|
|
"show_preview",
|
|
"auto_refresh_after_generation",
|
|
"trigger",
|
|
"node_id",
|
|
} <= set(inputs["required"])
|
|
assert inputs["optional"]["input_image"] == ("IMAGE",)
|
|
assert inputs["optional"]["input_mask"] == ("MASK",)
|
|
transport_inputs = {
|
|
name: value
|
|
for name, value in inputs["optional"].items()
|
|
if name.startswith("input_image_")
|
|
}
|
|
assert len(transport_inputs) == 32
|
|
assert all(value == ("IMAGE", {"hidden": True}) for value in transport_inputs.values())
|
|
|
|
|
|
def test_entrypoint_registers_backend_route_contract(layerforge_runtime):
|
|
registered = set(layerforge_runtime.routes.registered)
|
|
expected = {
|
|
("GET", "/layerforge/canvas_ws"),
|
|
("GET", "/layerforge/get_input_data/{node_id}"),
|
|
("POST", "/layerforge/clear_input_data/{node_id}"),
|
|
("GET", "/ycnode/get_canvas_data/{node_id}"),
|
|
("GET", "/layerforge/get-latest-images/{since}"),
|
|
("GET", "/ycnode/get_latest_image"),
|
|
("POST", "/ycnode/load_image_from_path"),
|
|
("GET", "/matting/settings"),
|
|
("POST", "/matting/settings"),
|
|
("GET", "/matting/check-model"),
|
|
("GET", "/matting/progress"),
|
|
("POST", "/matting"),
|
|
}
|
|
|
|
assert expected <= registered
|
|
|
|
|
|
def test_matting_settings_are_persisted_without_exposing_the_token(layerforge_runtime, monkeypatch, tmp_path):
|
|
settings_module = layerforge_runtime.matting_settings
|
|
settings_file = tmp_path / "layerforge_settings.json"
|
|
monkeypatch.setattr(settings_module, "SETTINGS_FILE", settings_file)
|
|
|
|
saved = settings_module.save_settings(
|
|
{
|
|
"model_path": "remote:rmbg_2_0",
|
|
"mode": "mask_only_inverted",
|
|
"threshold": 0.75,
|
|
"hf_token": "hf-test-token",
|
|
}
|
|
)
|
|
|
|
assert saved["model_path"] == "remote:rmbg_2_0"
|
|
assert saved["mode"] == "mask_only_inverted"
|
|
assert saved["threshold"] == 0.75
|
|
assert settings_module.get_huggingface_token() == "hf-test-token"
|
|
public_settings = settings_module.get_public_settings()
|
|
assert public_settings["hf_token_configured"] is True
|
|
assert "hf_token" not in public_settings
|
|
assert settings_file.exists()
|
|
|
|
settings_module.save_settings({"clear_hf_token": True})
|
|
assert settings_module.get_huggingface_token() == ""
|
|
|
|
|
|
def test_tensor_input_normalization_preserves_comfyui_shapes(layerforge_runtime):
|
|
import torch
|
|
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_cache["persistent_cache"] = {}
|
|
node = node_class()
|
|
|
|
chw_image = torch.ones((1, 3, 2, 4), dtype=torch.float32)
|
|
normalized_image = node.add_image_to_canvas(chw_image)
|
|
assert tuple(normalized_image.shape) == (2, 4, 3)
|
|
|
|
small_mask = torch.zeros((1, 2, 2), dtype=torch.float32)
|
|
resized_mask = node.add_mask_to_canvas(small_mask, torch.zeros((4, 4, 3)))
|
|
assert tuple(resized_mask.shape) == (4, 4)
|
|
|
|
|
|
def test_base64_image_conversion_preserves_rgb_and_alpha(layerforge_runtime):
|
|
import torch
|
|
from PIL import Image
|
|
|
|
node_module = layerforge_runtime.image_utils
|
|
image = Image.new("RGBA", (2, 1), (255, 0, 0, 128))
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
encoded = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()
|
|
|
|
tensor, alpha = node_module.convert_base64_to_tensor(encoded)
|
|
|
|
assert tuple(tensor.shape) == (1, 3, 1, 2)
|
|
assert tuple(alpha.shape) == (1, 1, 2)
|
|
assert float(alpha[0, 0, 0]) == pytest.approx(128 / 255, abs=0.01)
|
|
|
|
rgb_tensor = torch.tensor(
|
|
[
|
|
[
|
|
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
|
[[0.0, 0.0, 1.0], [1.0, 1.0, 1.0]],
|
|
]
|
|
],
|
|
dtype=torch.float32,
|
|
)
|
|
roundtrip = node_module.convert_tensor_to_base64(rgb_tensor)
|
|
assert roundtrip.startswith("data:image/png;base64,")
|
|
|
|
roundtrip_image = Image.open(io.BytesIO(base64.b64decode(roundtrip.split(",", 1)[1])))
|
|
assert roundtrip_image.mode == "RGB"
|
|
assert roundtrip_image.size == (2, 2)
|
|
assert roundtrip_image.getpixel((0, 0)) == (255, 0, 0)
|
|
|
|
|
|
def test_tensor_input_payloads_preserve_single_and_batch_image_contract(layerforge_runtime):
|
|
import torch
|
|
from PIL import Image
|
|
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_data_storage.clear()
|
|
node_class._canvas_cache["persistent_cache"] = {}
|
|
node = node_class()
|
|
single_image = torch.tensor(
|
|
[[
|
|
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
|
[[0.0, 0.0, 1.0], [1.0, 1.0, 1.0]],
|
|
]],
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
node.process_canvas_image(False, False, False, 0, "single-input", input_image=single_image)
|
|
single_payload = node_class._canvas_data_storage["single-input_input"]
|
|
single_image_data = base64.b64decode(single_payload["input_image"].split(",", 1)[1])
|
|
single_decoded = Image.open(io.BytesIO(single_image_data))
|
|
|
|
assert single_payload["input_image_width"] == 2
|
|
assert single_payload["input_image_height"] == 2
|
|
assert single_decoded.size == (2, 2)
|
|
assert single_decoded.getpixel((0, 0)) == (255, 0, 0)
|
|
|
|
batch_image = torch.stack((single_image[0], single_image[0] * 0.5))
|
|
node.process_canvas_image(False, False, False, 0, "batch-input", input_image=batch_image)
|
|
batch_payload = node_class._canvas_data_storage["batch-input_input"]
|
|
|
|
assert "input_image" not in batch_payload
|
|
assert [item["width"] for item in batch_payload["input_images_batch"]] == [2, 2]
|
|
assert [item["height"] for item in batch_payload["input_images_batch"]] == [2, 2]
|
|
assert all(item["data"].startswith("data:image/png;base64,") for item in batch_payload["input_images_batch"])
|
|
|
|
|
|
def test_tensor_input_payloads_preserve_ordered_multiple_image_inputs(layerforge_runtime):
|
|
import torch
|
|
from PIL import Image
|
|
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_data_storage.clear()
|
|
node_class._canvas_cache["persistent_cache"] = {}
|
|
node = node_class()
|
|
|
|
red = torch.tensor(
|
|
[[[
|
|
[1.0, 0.0, 0.0],
|
|
[1.0, 0.0, 0.0],
|
|
]]],
|
|
dtype=torch.float32,
|
|
)
|
|
green = red.clone()
|
|
green[..., 0] = 0
|
|
green[..., 1] = 1
|
|
|
|
# Deliberately pass the transport inputs out of numeric order. The node
|
|
# must preserve the port order used by the virtual multi-input frontend.
|
|
node.process_canvas_image(
|
|
False,
|
|
False,
|
|
False,
|
|
0,
|
|
"multi-input",
|
|
input_image_2=green,
|
|
input_image_1=red,
|
|
)
|
|
payload = node_class._canvas_data_storage["multi-input_input"]
|
|
|
|
assert "input_image" not in payload
|
|
assert len(payload["input_images"]) == 2
|
|
decoded = [
|
|
Image.open(io.BytesIO(base64.b64decode(item["data"].split(",", 1)[1])))
|
|
for item in payload["input_images"]
|
|
]
|
|
assert decoded[0].getpixel((0, 0)) == (255, 0, 0)
|
|
assert decoded[1].getpixel((0, 0)) == (0, 255, 0)
|
|
|
|
|
|
def test_latest_image_helpers_preserve_filtering_and_time_semantics(layerforge_runtime, monkeypatch, tmp_path):
|
|
node_module = layerforge_runtime.node
|
|
node_class = node_module.LayerForgeNode
|
|
output_dir = tmp_path / "output"
|
|
output_dir.mkdir()
|
|
first_image = output_dir / "first.png"
|
|
second_image = output_dir / "second.JPG"
|
|
ignored_file = output_dir / "ignored.txt"
|
|
ignored_directory = output_dir / "folder.png"
|
|
first_image.write_bytes(b"first")
|
|
second_image.write_bytes(b"second")
|
|
ignored_file.write_bytes(b"ignored")
|
|
ignored_directory.mkdir()
|
|
|
|
monkeypatch.setattr(
|
|
node_module.folder_paths,
|
|
"get_output_directory",
|
|
lambda: str(output_dir),
|
|
raising=False,
|
|
)
|
|
creation_times = {
|
|
str(first_image): 10,
|
|
str(second_image): 20,
|
|
}
|
|
monkeypatch.setattr(node_module.os.path, "getctime", lambda path: creation_times[path])
|
|
assert node_class.get_latest_image() == str(second_image)
|
|
|
|
modification_times = {
|
|
str(first_image): 30,
|
|
str(second_image): 20,
|
|
}
|
|
monkeypatch.setattr(node_module.os.path, "getmtime", lambda path: modification_times[path])
|
|
assert node_class.get_latest_images(0) == [str(second_image), str(first_image)]
|
|
assert node_class.get_latest_images(25) == [str(first_image)]
|
|
|
|
|
|
def test_websocket_and_cached_image_decoding_preserve_image_and_mask_modes(layerforge_runtime):
|
|
from PIL import Image
|
|
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_data_storage.clear()
|
|
node_class._canvas_cache["persistent_cache"] = {}
|
|
|
|
image = Image.new("RGBA", (2, 1), (255, 0, 0, 128))
|
|
mask = Image.new("L", (2, 1), 128)
|
|
|
|
def encode(pil_image):
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format="PNG")
|
|
return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()
|
|
|
|
node_class._canvas_data_storage["decode-node"] = {
|
|
"image": encode(image),
|
|
"mask": encode(mask),
|
|
}
|
|
node = node_class()
|
|
processed_image, processed_mask = node.process_canvas_image(
|
|
False,
|
|
False,
|
|
False,
|
|
0,
|
|
"decode-node",
|
|
)
|
|
|
|
assert tuple(processed_image.shape) == (1, 1, 2, 3)
|
|
assert tuple(processed_mask.shape) == (1, 1, 2)
|
|
|
|
node.store_image(encode(image))
|
|
assert node.cached_image.mode == "RGBA"
|
|
assert node.cached_image.size == (2, 1)
|
|
|
|
|
|
def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, monkeypatch):
|
|
import torch
|
|
|
|
service_module = layerforge_runtime.matting_service
|
|
calls = {}
|
|
|
|
class NativeBiRefNet:
|
|
def encode_image(self, image):
|
|
calls["shape"] = tuple(image.shape)
|
|
calls["dtype"] = image.dtype
|
|
return torch.full((image.shape[0], image.shape[1], image.shape[2]), 0.75)
|
|
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"_get_comfy_birefnet_loader",
|
|
lambda: lambda path: NativeBiRefNet(),
|
|
)
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"_ensure_birefnet_checkpoint",
|
|
lambda model_path=None: "native-birefnet.safetensors",
|
|
)
|
|
service_module.BiRefNetMatting._model_cache.clear()
|
|
|
|
image = torch.rand((1, 3, 2, 4), dtype=torch.float32)
|
|
matted_image, alpha_mask = service_module.BiRefNetMatting().execute(
|
|
image,
|
|
model_path=None,
|
|
threshold=0,
|
|
refinement=1,
|
|
)
|
|
|
|
assert calls == {"shape": (1, 2, 4, 3), "dtype": torch.float32}
|
|
assert tuple(matted_image.shape) == (1, 3, 2, 4)
|
|
assert tuple(alpha_mask.shape) == (1, 1, 2, 4)
|
|
|
|
|
|
def test_matting_adapter_translates_comfyui_interrupt(layerforge_runtime, monkeypatch):
|
|
import torch
|
|
|
|
service_module = layerforge_runtime.matting_service
|
|
|
|
class FakeInterruptProcessingException(BaseException):
|
|
pass
|
|
|
|
comfy_module = ModuleType("comfy")
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_management.InterruptProcessingException = FakeInterruptProcessingException
|
|
comfy_module.model_management = model_management
|
|
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
|
|
class InterruptingBiRefNet:
|
|
def encode_image(self, image):
|
|
del image
|
|
raise FakeInterruptProcessingException()
|
|
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"_get_comfy_birefnet_loader",
|
|
lambda: lambda path: InterruptingBiRefNet(),
|
|
)
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"_ensure_birefnet_checkpoint",
|
|
lambda model_path=None: "interrupting-birefnet.safetensors",
|
|
)
|
|
service_module.BiRefNetMatting._model_cache.clear()
|
|
|
|
with pytest.raises(service_module.MattingInterruptedError, match="interrupted"):
|
|
service_module.BiRefNetMatting().execute(
|
|
torch.ones((1, 3, 2, 4), dtype=torch.float32),
|
|
model_path=None,
|
|
threshold=0,
|
|
refinement=1,
|
|
)
|
|
|
|
|
|
def test_matting_model_status_preserves_unsupported_and_error_responses(layerforge_runtime, monkeypatch):
|
|
api_module = layerforge_runtime.matting_api
|
|
responses = _record_matting_responses(monkeypatch, api_module)
|
|
|
|
monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: None)
|
|
unsupported = _run_matting_model_check(api_module)
|
|
|
|
assert unsupported.payload["available"] is False
|
|
assert unsupported.payload["reason"] == "unsupported_comfyui"
|
|
assert unsupported.payload["message"] == (
|
|
"This ComfyUI version does not provide the native BiRefNet background-removal loader."
|
|
)
|
|
assert unsupported.payload["models"]
|
|
assert unsupported.status == 200
|
|
|
|
monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: object())
|
|
|
|
def fail_to_read_options():
|
|
raise RuntimeError("catalog unavailable")
|
|
|
|
monkeypatch.setattr(api_module, "_get_birefnet_model_options", fail_to_read_options)
|
|
failed = _run_matting_model_check(api_module)
|
|
|
|
assert failed.payload == {
|
|
"available": False,
|
|
"reason": "error",
|
|
"message": "Error checking model status: catalog unavailable",
|
|
}
|
|
assert failed.status == 500
|
|
assert len(responses) == 2
|
|
|
|
|
|
def test_matting_model_status_preserves_remote_model_responses(layerforge_runtime, monkeypatch):
|
|
api_module = layerforge_runtime.matting_api
|
|
responses = _record_matting_responses(monkeypatch, api_module)
|
|
model_options = [{"path": "remote:portrait", "label": "Portrait"}]
|
|
remote_model = {"label": "Portrait"}
|
|
|
|
monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: object())
|
|
monkeypatch.setattr(api_module, "_get_birefnet_model_options", lambda: model_options)
|
|
monkeypatch.setattr(api_module, "_get_birefnet_remote_model", lambda path: remote_model)
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"_find_existing_birefnet_remote_checkpoint",
|
|
lambda model: "/models/portrait.safetensors",
|
|
)
|
|
|
|
ready = _run_matting_model_check(api_module, "remote:portrait")
|
|
|
|
assert ready.payload == {
|
|
"available": True,
|
|
"reason": "ready",
|
|
"message": "Selected model is ready to use",
|
|
"model_path": "/models/portrait.safetensors",
|
|
"selected_model": "Portrait",
|
|
"models": model_options,
|
|
}
|
|
assert ready.status == 200
|
|
|
|
monkeypatch.setattr(api_module, "_find_existing_birefnet_remote_checkpoint", lambda model: None)
|
|
missing = _run_matting_model_check(api_module, "remote:portrait")
|
|
|
|
assert missing.payload == {
|
|
"available": False,
|
|
"reason": "not_downloaded",
|
|
"message": "Portrait will be downloaded automatically on first use.",
|
|
"model_path": "remote:portrait",
|
|
"selected_model": "Portrait",
|
|
"models": model_options,
|
|
}
|
|
assert missing.status == 200
|
|
assert len(responses) == 2
|
|
|
|
|
|
def test_matting_model_status_preserves_local_and_automatic_responses(layerforge_runtime, monkeypatch):
|
|
api_module = layerforge_runtime.matting_api
|
|
responses = _record_matting_responses(monkeypatch, api_module)
|
|
model_options = [{"path": "/models/custom.safetensors", "label": "custom.safetensors"}]
|
|
|
|
monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: object())
|
|
monkeypatch.setattr(api_module, "_get_birefnet_model_options", lambda: model_options)
|
|
monkeypatch.setattr(api_module, "_get_birefnet_remote_model", lambda path: None)
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"_find_local_birefnet_model",
|
|
lambda model_path=None: "/models/custom.safetensors",
|
|
)
|
|
|
|
selected = _run_matting_model_check(api_module, "/models/custom.safetensors")
|
|
|
|
assert selected.payload == {
|
|
"available": True,
|
|
"reason": "ready",
|
|
"message": "Selected model is ready to use",
|
|
"model_path": "/models/custom.safetensors",
|
|
"selected_model": "/models/custom.safetensors",
|
|
"models": model_options,
|
|
}
|
|
assert selected.status == 200
|
|
|
|
monkeypatch.setattr(api_module, "_find_local_birefnet_model", lambda model_path=None: None)
|
|
unavailable = _run_matting_model_check(api_module, "/models/missing.safetensors")
|
|
|
|
assert unavailable.payload == {
|
|
"available": False,
|
|
"reason": "selected_model_unavailable",
|
|
"message": "The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.",
|
|
"model_path": "/models/missing.safetensors",
|
|
"models": model_options,
|
|
}
|
|
assert unavailable.status == 200
|
|
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"_find_local_birefnet_model",
|
|
lambda model_path=None: "/models/automatic.safetensors",
|
|
)
|
|
automatic = _run_matting_model_check(api_module)
|
|
|
|
assert automatic.payload == {
|
|
"available": True,
|
|
"reason": "ready",
|
|
"message": "Model is ready to use",
|
|
"model_path": "/models/automatic.safetensors",
|
|
"models": model_options,
|
|
}
|
|
assert automatic.status == 200
|
|
|
|
monkeypatch.setattr(api_module, "_find_local_birefnet_model", lambda model_path=None: None)
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"_get_birefnet_base_paths",
|
|
lambda: ["/models/background_removal"],
|
|
)
|
|
not_downloaded = _run_matting_model_check(api_module)
|
|
|
|
assert not_downloaded.payload == {
|
|
"available": False,
|
|
"reason": "not_downloaded",
|
|
"message": "The BiRefNet checkpoint will be downloaded automatically on first use (requires internet connection).",
|
|
"model_path": "/models/background_removal",
|
|
"models": model_options,
|
|
}
|
|
assert not_downloaded.status == 200
|
|
assert len(responses) == 4
|
|
|
|
|
|
def test_matting_adapter_supports_inverted_and_mask_only_modes(layerforge_runtime, monkeypatch):
|
|
import torch
|
|
|
|
service_module = layerforge_runtime.matting_service
|
|
|
|
class NativeBiRefNet:
|
|
def encode_image(self, image):
|
|
return torch.full((image.shape[0], image.shape[1], image.shape[2]), 0.75)
|
|
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"_get_comfy_birefnet_loader",
|
|
lambda: lambda path: NativeBiRefNet(),
|
|
)
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"_ensure_birefnet_checkpoint",
|
|
lambda model_path=None: "native-birefnet.safetensors",
|
|
)
|
|
service_module.BiRefNetMatting._model_cache.clear()
|
|
|
|
image = torch.ones((1, 3, 2, 4), dtype=torch.float32)
|
|
matting = service_module.BiRefNetMatting()
|
|
|
|
removed_foreground, inverted_mask = matting.execute(
|
|
image,
|
|
model_path=None,
|
|
threshold=0.5,
|
|
refinement=1,
|
|
mode="remove_foreground",
|
|
)
|
|
mask_preview, preview_mask = matting.execute(
|
|
image,
|
|
model_path=None,
|
|
threshold=0.5,
|
|
refinement=1,
|
|
mode="mask_only",
|
|
)
|
|
inverted_mask_preview, inverted_preview_mask = matting.execute(
|
|
image,
|
|
model_path=None,
|
|
threshold=0.5,
|
|
refinement=1,
|
|
mode="mask_only_inverted",
|
|
)
|
|
|
|
assert torch.allclose(removed_foreground, torch.zeros_like(removed_foreground))
|
|
assert torch.allclose(inverted_mask, torch.zeros_like(inverted_mask))
|
|
assert torch.allclose(mask_preview, torch.ones_like(mask_preview))
|
|
assert torch.allclose(preview_mask, torch.ones_like(preview_mask))
|
|
assert torch.allclose(inverted_mask_preview, torch.zeros_like(inverted_mask_preview))
|
|
assert torch.allclose(inverted_preview_mask, torch.zeros_like(inverted_preview_mask))
|
|
|
|
|
|
def test_matting_model_options_include_downloadable_official_variants(layerforge_runtime, monkeypatch):
|
|
birefnet_module = layerforge_runtime.matting_birefnet
|
|
options_module = layerforge_runtime.matting_options
|
|
catalog_module = layerforge_runtime.matting_catalog
|
|
|
|
monkeypatch.setattr(birefnet_module, "_iter_birefnet_checkpoint_paths", lambda: iter(()))
|
|
|
|
options = options_module._get_birefnet_model_options()
|
|
remote_options = [option for option in options if option["source"] == "remote"]
|
|
|
|
assert remote_options
|
|
assert all(option["path"].startswith("remote:") for option in remote_options)
|
|
assert all(option["downloaded"] is False for option in remote_options)
|
|
assert all(option["description"] for option in remote_options)
|
|
assert all(option["url"].startswith("https://huggingface.co/") for option in remote_options)
|
|
birefnet_options = [option for option in remote_options if option["backend"] == "birefnet"]
|
|
assert all(option["project_url"] == "https://github.com/ZhengPeng7/BiRefNet" for option in birefnet_options)
|
|
assert any(option["path"] == "remote:portrait" for option in remote_options)
|
|
portrait = next(option for option in catalog_module._BIREFNET_MODEL_CATALOG if option["id"] == "portrait")
|
|
assert portrait["local_filename"] == "BiRefNet-portrait.safetensors"
|
|
|
|
rmbg_option = next(option for option in remote_options if option["path"] == "remote:rmbg_2_0")
|
|
assert rmbg_option["backend"] == "rmbg"
|
|
assert rmbg_option["label"] == "BRIA RMBG 2.0"
|
|
assert rmbg_option["project_url"] == "https://github.com/Bria-AI/RMBG-2.0"
|
|
|
|
|
|
def test_selected_remote_matting_model_is_sent_to_downloader(layerforge_runtime, monkeypatch):
|
|
birefnet_module = layerforge_runtime.matting_birefnet
|
|
catalog_module = layerforge_runtime.matting_catalog
|
|
selected_model = next(
|
|
model for model in catalog_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait"
|
|
)
|
|
downloaded = {}
|
|
|
|
monkeypatch.setattr(birefnet_module, "_is_native_birefnet_checkpoint", lambda path: False)
|
|
|
|
def fake_download(model=None):
|
|
downloaded["model"] = model
|
|
return "downloaded-portrait.safetensors"
|
|
|
|
monkeypatch.setattr(birefnet_module, "_download_birefnet_checkpoint", fake_download)
|
|
|
|
result = birefnet_module._ensure_birefnet_checkpoint("remote:portrait")
|
|
|
|
assert result == "downloaded-portrait.safetensors"
|
|
assert downloaded["model"] is selected_model
|
|
|
|
|
|
def test_remote_matting_checkpoint_path_uses_background_removal_root(layerforge_runtime):
|
|
import folder_paths
|
|
|
|
catalog_module = layerforge_runtime.matting_catalog
|
|
paths_module = layerforge_runtime.matting_paths
|
|
model = next(model for model in catalog_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait")
|
|
|
|
expected = Path(folder_paths.models_dir) / "background_removal" / model[
|
|
"local_filename"
|
|
]
|
|
|
|
assert Path(paths_module._get_birefnet_remote_checkpoint_path(model)) == expected
|
|
|
|
|
|
def test_remote_matting_download_uses_background_removal_root(layerforge_runtime, monkeypatch, tmp_path):
|
|
birefnet_module = layerforge_runtime.matting_birefnet
|
|
catalog_module = layerforge_runtime.matting_catalog
|
|
model = next(model for model in catalog_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait")
|
|
download = {}
|
|
|
|
def fake_hf_hub_download(**kwargs):
|
|
download.update(kwargs)
|
|
downloaded_path = Path(kwargs["local_dir"]) / kwargs["filename"]
|
|
downloaded_path.write_bytes(b"checkpoint")
|
|
return str(downloaded_path)
|
|
|
|
huggingface_hub = ModuleType("huggingface_hub")
|
|
huggingface_hub.hf_hub_download = fake_hf_hub_download
|
|
monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub)
|
|
monkeypatch.setattr(birefnet_module, "_is_native_birefnet_checkpoint", lambda path: True)
|
|
|
|
result = birefnet_module._download_birefnet_checkpoint(model)
|
|
expected_dir = Path(tmp_path) / "models" / "background_removal"
|
|
|
|
assert Path(download["local_dir"]) == expected_dir
|
|
assert "layerforge_birefnet" not in Path(download["local_dir"]).parts
|
|
assert Path(result) == expected_dir / model["local_filename"]
|
|
assert Path(result).exists()
|
|
|
|
|
|
def test_model_download_progress_reports_bytes_to_the_frontend(layerforge_runtime, monkeypatch):
|
|
progress_module = layerforge_runtime.matting_progress
|
|
events = []
|
|
monkeypatch.setattr(
|
|
sys.modules["server"].PromptServer.instance,
|
|
"send_sync",
|
|
lambda *args: events.append(args),
|
|
)
|
|
|
|
progress_class = progress_module.create_huggingface_tqdm_class("Test model", node_id="42")
|
|
progress = progress_class(total=100, initial=0, unit="B", disable=True)
|
|
progress.update(25)
|
|
progress.update(75)
|
|
progress.close()
|
|
|
|
download_events = [event for event in events if event[1]["status"] == "downloading"]
|
|
assert download_events
|
|
assert download_events[-1][1]["node_id"] == "42"
|
|
assert download_events[-1][1]["progress"] == 100.0
|
|
assert download_events[-1][1]["downloaded_bytes"] == 100
|
|
assert download_events[-1][1]["total_bytes"] == 100
|
|
progress_module.send_matting_status("completed", node_id="42")
|
|
assert progress_module.get_matting_status("42")["status"] == "completed"
|
|
|
|
|
|
def test_matting_progress_status_is_available_without_websocket_delivery(layerforge_runtime):
|
|
progress_module = layerforge_runtime.matting_progress
|
|
|
|
progress_module.send_matting_status(
|
|
"downloading",
|
|
node_id="poll-node",
|
|
progress=37.5,
|
|
downloaded_bytes=375,
|
|
total_bytes=1000,
|
|
)
|
|
|
|
assert progress_module.get_matting_status("poll-node") == {
|
|
"status": "downloading",
|
|
"node_id": "poll-node",
|
|
"progress": 37.5,
|
|
"downloaded_bytes": 375,
|
|
"total_bytes": 1000,
|
|
}
|
|
progress_module.send_matting_status("completed", node_id="poll-node")
|
|
|
|
|
|
def test_rmbg_model_path_uses_background_removal_subdirectory(layerforge_runtime):
|
|
import folder_paths
|
|
|
|
catalog_module = layerforge_runtime.matting_catalog
|
|
paths_module = layerforge_runtime.matting_paths
|
|
model = catalog_module._RMBG_MODEL_CATALOG[0]
|
|
|
|
expected = Path(folder_paths.models_dir) / "background_removal" / "RMBG-2.0"
|
|
|
|
assert Path(paths_module._get_rmbg_model_directory(model)) == expected
|
|
|
|
|
|
def test_rmbg_checkpoint_is_not_selected_as_native_birefnet(layerforge_runtime, monkeypatch, tmp_path):
|
|
birefnet_module = layerforge_runtime.matting_birefnet
|
|
model_dir = Path(tmp_path) / "models" / "background_removal" / "RMBG-2.0"
|
|
model_dir.mkdir(parents=True)
|
|
checkpoint_path = model_dir / "model.safetensors"
|
|
checkpoint_path.write_bytes(b"RMBG checkpoint")
|
|
|
|
monkeypatch.setattr(birefnet_module, "_get_birefnet_base_paths", lambda: [str(model_dir.parent)])
|
|
monkeypatch.setattr(birefnet_module, "_iter_birefnet_checkpoint_paths", lambda: iter((str(checkpoint_path),)))
|
|
monkeypatch.setattr(birefnet_module, "_is_native_birefnet_checkpoint", lambda path: True)
|
|
|
|
assert birefnet_module._find_local_birefnet_model() is None
|
|
|
|
|
|
def test_rmbg_download_uses_background_removal_subdirectory(layerforge_runtime, monkeypatch, tmp_path):
|
|
rmbg_module = layerforge_runtime.matting_rmbg
|
|
catalog_module = layerforge_runtime.matting_catalog
|
|
model = catalog_module._RMBG_MODEL_CATALOG[0]
|
|
download = {}
|
|
|
|
def fake_snapshot_download(**kwargs):
|
|
download.update(kwargs)
|
|
model_dir = Path(kwargs["local_dir"])
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
for filename in (
|
|
"config.json",
|
|
"preprocessor_config.json",
|
|
"birefnet.py",
|
|
"BiRefNet_config.py",
|
|
"model.safetensors",
|
|
):
|
|
(model_dir / filename).write_bytes(b"model")
|
|
return str(model_dir)
|
|
|
|
huggingface_hub = ModuleType("huggingface_hub")
|
|
huggingface_hub.snapshot_download = fake_snapshot_download
|
|
monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub)
|
|
monkeypatch.setattr(rmbg_module, "get_huggingface_token", lambda: "hf-test-token")
|
|
|
|
result = rmbg_module._download_rmbg_model(model)
|
|
expected_dir = Path(tmp_path) / "models" / "background_removal" / "RMBG-2.0"
|
|
|
|
assert Path(download["local_dir"]) == expected_dir
|
|
assert download["repo_id"] == "briaai/RMBG-2.0"
|
|
assert "*.safetensors" in download["allow_patterns"]
|
|
assert download["token"] == "hf-test-token"
|
|
assert Path(result) == expected_dir
|
|
assert rmbg_module._is_rmbg_model_directory(result)
|
|
|
|
|
|
def test_rmbg_adapter_returns_batched_hw_mask(layerforge_runtime):
|
|
import torch
|
|
|
|
class FakeRMBGModel:
|
|
def __call__(self, image):
|
|
return (None, torch.zeros((image.shape[0], 1, 1024, 1024), device=image.device))
|
|
|
|
adapter = layerforge_runtime.matting_rmbg.RMBG2Model(
|
|
FakeRMBGModel(),
|
|
torch.device("cpu"),
|
|
)
|
|
image = torch.rand((1, 2, 3, 4), dtype=torch.float32)
|
|
|
|
mask = adapter.encode_image(image)
|
|
|
|
assert tuple(mask.shape) == (1, 2, 3)
|
|
assert torch.allclose(mask, torch.full((1, 2, 3), 0.5))
|
|
|
|
|
|
def test_rmbg_model_status_does_not_require_native_birefnet_loader(layerforge_runtime, monkeypatch):
|
|
api_module = layerforge_runtime.matting_api
|
|
_record_matting_responses(monkeypatch, api_module)
|
|
|
|
monkeypatch.setattr(api_module, "_get_rmbg_model_loader", lambda: object())
|
|
monkeypatch.setattr(api_module, "_find_existing_rmbg_model", lambda model: None)
|
|
|
|
status = _run_matting_model_check(api_module, "remote:rmbg_2_0")
|
|
|
|
assert status.payload["available"] is False
|
|
assert status.payload["reason"] == "not_downloaded"
|
|
assert status.payload["selected_model"] == "BRIA RMBG 2.0"
|
|
assert status.payload["model_path"] == "remote:rmbg_2_0"
|
|
|
|
|
|
def test_rmbg_reports_an_unsupported_transformers_api(layerforge_runtime, monkeypatch):
|
|
rmbg_module = layerforge_runtime.matting_rmbg
|
|
transformers = ModuleType("transformers")
|
|
transformers.__version__ = "0.0-test"
|
|
monkeypatch.setitem(sys.modules, "transformers", transformers)
|
|
|
|
status = rmbg_module._get_rmbg_transformers_status()
|
|
|
|
assert status["loader"] is None
|
|
assert "0.0-test" in status["message"]
|
|
assert "not supported" in status["message"]
|
|
assert "AutoModelForImageSegmentation" in status["message"]
|
|
assert rmbg_module._get_rmbg_model_loader() is None
|
|
|
|
|
|
def test_legacy_automatic_checkpoint_is_not_migrated(layerforge_runtime, monkeypatch, tmp_path):
|
|
birefnet_module = layerforge_runtime.matting_birefnet
|
|
model_dir = Path(tmp_path) / "models" / "background_removal"
|
|
model_dir.mkdir(parents=True)
|
|
legacy_path = model_dir / "model.safetensors"
|
|
friendly_path = model_dir / "BiRefNet-general.safetensors"
|
|
legacy_path.write_bytes(b"legacy checkpoint")
|
|
|
|
monkeypatch.setattr(
|
|
birefnet_module,
|
|
"_is_native_birefnet_checkpoint",
|
|
lambda path: Path(path).resolve() == friendly_path.resolve() and Path(path).is_file(),
|
|
)
|
|
|
|
result = birefnet_module._find_existing_birefnet_default_checkpoint()
|
|
|
|
assert result is None
|
|
assert legacy_path.exists()
|
|
assert not friendly_path.exists()
|
|
|
|
|
|
def test_empty_execution_returns_comfyui_compatible_fallback_tensors(layerforge_runtime):
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_data_storage.clear()
|
|
node_class._canvas_cache["persistent_cache"] = {}
|
|
node = node_class()
|
|
|
|
image, mask = node.process_canvas_image(
|
|
False,
|
|
False,
|
|
False,
|
|
0,
|
|
"test-node",
|
|
)
|
|
|
|
assert tuple(image.shape) == (1, 512, 512, 3)
|
|
assert tuple(mask.shape) == (1, 512, 512)
|
|
|
|
|
|
def test_legacy_route_handlers_cover_cache_and_file_responses(layerforge_runtime, monkeypatch, tmp_path):
|
|
from PIL import Image
|
|
|
|
routes_module = sys.modules[f"{layerforge_runtime.package_name}.python.routes"]
|
|
responses = _record_matting_responses(monkeypatch, routes_module)
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_data_storage.clear()
|
|
node_class._canvas_cache["image"] = Image.new("RGB", (2, 1), (255, 0, 0))
|
|
node_class._canvas_cache["mask"] = Image.new("L", (2, 1), 128)
|
|
|
|
class JsonRequest:
|
|
def __init__(self, payload):
|
|
self.payload = payload
|
|
self.match_info = {}
|
|
|
|
async def json(self):
|
|
return self.payload
|
|
|
|
get_input = layerforge_runtime.routes.registered[("GET", "/layerforge/get_input_data/{node_id}")]
|
|
clear_input = layerforge_runtime.routes.registered[("POST", "/layerforge/clear_input_data/{node_id}")]
|
|
get_canvas = layerforge_runtime.routes.registered[("GET", "/ycnode/get_canvas_data/{node_id}")]
|
|
get_latest_images = layerforge_runtime.routes.registered[("GET", "/layerforge/get-latest-images/{since}")]
|
|
get_latest_image = layerforge_runtime.routes.registered[("GET", "/ycnode/get_latest_image")]
|
|
load_image = layerforge_runtime.routes.registered[("POST", "/ycnode/load_image_from_path")]
|
|
|
|
node_class._canvas_data_storage["route-node_input"] = {"input_image": "image-data"}
|
|
input_request = JsonRequest(None)
|
|
input_request.match_info = {"node_id": "route-node"}
|
|
found = asyncio.run(get_input(input_request))
|
|
assert found.payload == {"success": True, "has_input": True, "data": {"input_image": "image-data"}}
|
|
|
|
cleared = asyncio.run(clear_input(input_request))
|
|
assert cleared.payload["success"] is True
|
|
assert "route-node_input" not in node_class._canvas_data_storage
|
|
missing = asyncio.run(get_input(input_request))
|
|
assert missing.payload == {"success": True, "has_input": False}
|
|
|
|
canvas_response = asyncio.run(get_canvas(JsonRequest(None)))
|
|
assert canvas_response.payload["success"] is True
|
|
assert canvas_response.payload["data"]["image"].startswith("data:image/png;base64,")
|
|
assert canvas_response.payload["data"]["mask"].startswith("data:image/png;base64,")
|
|
|
|
latest_paths = [str(tmp_path / "first.png"), str(tmp_path / "second.png")]
|
|
monkeypatch.setattr(
|
|
node_class,
|
|
"get_latest_images",
|
|
classmethod(lambda cls, since: latest_paths if since == 1.5 else []),
|
|
)
|
|
monkeypatch.setattr(routes_module, "file_to_data_url", lambda path: f"data:{path}")
|
|
latest_request = JsonRequest(None)
|
|
latest_request.match_info = {"since": "1500"}
|
|
latest_response = asyncio.run(get_latest_images(latest_request))
|
|
assert latest_response.payload == {
|
|
"success": True,
|
|
"images": [f"data:{path}" for path in latest_paths],
|
|
}
|
|
|
|
monkeypatch.setattr(node_class, "get_latest_image", classmethod(lambda cls: latest_paths[1]))
|
|
latest_image_response = asyncio.run(get_latest_image(JsonRequest(None)))
|
|
assert latest_image_response.payload == {"success": True, "image_data": f"data:{latest_paths[1]}"}
|
|
|
|
image_path = tmp_path / "loaded.png"
|
|
Image.new("RGBA", (3, 2), (255, 0, 0, 128)).save(image_path)
|
|
loaded = asyncio.run(load_image(JsonRequest({"file_path": str(image_path)})))
|
|
assert loaded.payload["success"] is True
|
|
assert loaded.payload["width"] == 3
|
|
assert loaded.payload["height"] == 2
|
|
|
|
not_found = asyncio.run(load_image(JsonRequest({"file_path": str(tmp_path / "missing.png")})))
|
|
assert not_found.status == 404
|
|
invalid_path = tmp_path / "file.txt"
|
|
invalid_path.write_text("not an image", encoding="utf-8")
|
|
invalid = asyncio.run(load_image(JsonRequest({"file_path": str(invalid_path)})))
|
|
assert invalid.status == 400
|
|
required = asyncio.run(load_image(JsonRequest({})))
|
|
assert required.status == 400
|
|
|
|
corrupt_path = tmp_path / "corrupt.png"
|
|
corrupt_path.write_bytes(b"not-an-image")
|
|
corrupt = asyncio.run(load_image(JsonRequest({"file_path": str(corrupt_path)})))
|
|
assert corrupt.status == 500
|
|
assert len(responses) >= 10
|
|
|
|
|
|
def test_matting_settings_and_progress_endpoints_validate_requests(layerforge_runtime, monkeypatch):
|
|
api_module = layerforge_runtime.matting_api
|
|
responses = _record_matting_responses(monkeypatch, api_module)
|
|
saved_payloads = []
|
|
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"get_public_settings",
|
|
lambda: {"model_path": "auto", "hf_token_configured": False},
|
|
)
|
|
monkeypatch.setattr(api_module, "get_matting_status", lambda node_id: {"node_id": node_id, "status": "idle"})
|
|
monkeypatch.setattr(api_module, "save_settings", lambda payload: saved_payloads.append(payload))
|
|
|
|
class Request:
|
|
def __init__(self, payload, query=None, error=None):
|
|
self.payload = payload
|
|
self.query = query or {}
|
|
self.error = error
|
|
|
|
async def json(self):
|
|
if self.error:
|
|
raise self.error
|
|
return self.payload
|
|
|
|
settings = asyncio.run(api_module.get_matting_settings(Request(None)))
|
|
assert settings.payload == {"settings": {"model_path": "auto", "hf_token_configured": False}}
|
|
|
|
progress = asyncio.run(api_module.get_matting_progress(Request(None, {"node_id": "node-7"})))
|
|
assert progress.payload == {"node_id": "node-7", "status": "idle"}
|
|
|
|
saved = asyncio.run(api_module.save_matting_settings(Request({"mode": "mask_only"})))
|
|
assert saved.payload == {"success": True, "settings": {"model_path": "auto", "hf_token_configured": False}}
|
|
assert saved_payloads == [{"mode": "mask_only"}]
|
|
|
|
invalid_json = asyncio.run(api_module.save_matting_settings(Request(None, error=ValueError("bad json"))))
|
|
assert invalid_json.status == 400
|
|
invalid_shape = asyncio.run(api_module.save_matting_settings(Request(["not", "an", "object"])))
|
|
assert invalid_shape.status == 400
|
|
|
|
def fail_to_save(payload):
|
|
raise RuntimeError(f"cannot save {payload}")
|
|
|
|
monkeypatch.setattr(api_module, "save_settings", fail_to_save)
|
|
failed = asyncio.run(api_module.save_matting_settings(Request({"mode": "remove_background"})))
|
|
assert failed.status == 500
|
|
assert len(responses) == 6
|
|
|
|
|
|
def test_matting_endpoint_returns_masks_and_releases_lock(layerforge_runtime, monkeypatch):
|
|
import torch
|
|
|
|
api_module = layerforge_runtime.matting_api
|
|
responses = _record_matting_responses(monkeypatch, api_module)
|
|
calls = []
|
|
|
|
class FakeMatting:
|
|
model_path = "loaded-model"
|
|
|
|
def execute(self, image, model_path, **kwargs):
|
|
calls.append((tuple(image.shape), model_path, kwargs))
|
|
return image * 0.5, torch.full((1, 1, 2, 2), 0.75)
|
|
|
|
monkeypatch.setattr(api_module, "BiRefNetMatting", FakeMatting)
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"convert_base64_to_tensor",
|
|
lambda encoded: (torch.ones((1, 3, 2, 2)), None),
|
|
)
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"convert_tensor_to_base64",
|
|
lambda tensor, alpha_mask=None, original_alpha=None: f"encoded-{tuple(tensor.shape)}",
|
|
)
|
|
|
|
class Request:
|
|
async def json(self):
|
|
return {
|
|
"image": "data:image/png;base64,stub",
|
|
"model_path": "remote:portrait",
|
|
"mode": "remove_background",
|
|
"threshold": 0.7,
|
|
"refinement": 2,
|
|
"node_id": "node-1",
|
|
}
|
|
|
|
response = asyncio.run(api_module.matting(Request()))
|
|
assert response.payload["mode"] == "remove_background"
|
|
assert response.payload["model_path"] == "loaded-model"
|
|
assert response.payload["matted_image"] == "encoded-(1, 3, 2, 2)"
|
|
assert response.payload["alpha_mask"] == "encoded-(1, 1, 2, 2)"
|
|
assert response.payload["draw_mask"] == "encoded-(1, 1, 2, 2)"
|
|
assert calls[0][1] == "remote:portrait"
|
|
assert calls[0][2]["node_id"] == "node-1"
|
|
assert api_module._matting_lock is None
|
|
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"convert_base64_to_tensor",
|
|
lambda encoded: (torch.ones((1, 3, 2, 2)), torch.ones((1, 1, 2, 2))),
|
|
)
|
|
async def mask_json():
|
|
return {
|
|
"image": "data:image/png;base64,stub",
|
|
"mode": "mask_only",
|
|
}
|
|
|
|
response = asyncio.run(
|
|
api_module.matting(
|
|
SimpleNamespace(json=mask_json)
|
|
)
|
|
)
|
|
assert response.payload["mode"] == "mask_only"
|
|
assert len(responses) == 2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("error", "status"),
|
|
[
|
|
(RuntimeError("offline"), 400),
|
|
(RuntimeError("model failed"), 500),
|
|
(ValueError("connection reset"), 400),
|
|
(ValueError("invalid input"), 500),
|
|
],
|
|
)
|
|
def test_matting_endpoint_translates_model_errors(layerforge_runtime, monkeypatch, error, status):
|
|
import torch
|
|
|
|
api_module = layerforge_runtime.matting_api
|
|
_record_matting_responses(monkeypatch, api_module)
|
|
|
|
class FailingMatting:
|
|
def execute(self, *args, **kwargs):
|
|
raise error
|
|
|
|
monkeypatch.setattr(api_module, "BiRefNetMatting", FailingMatting)
|
|
monkeypatch.setattr(
|
|
api_module,
|
|
"convert_base64_to_tensor",
|
|
lambda encoded: (torch.ones((1, 3, 2, 2)), None),
|
|
)
|
|
|
|
class Request:
|
|
async def json(self):
|
|
return {"image": "data:image/png;base64,stub"}
|
|
|
|
response = asyncio.run(api_module.matting(Request()))
|
|
assert response.status == status
|
|
assert api_module._matting_lock is None
|
|
|
|
|
|
def test_image_tensor_conversion_supports_grayscale_and_alpha(layerforge_runtime):
|
|
import torch
|
|
from PIL import Image
|
|
|
|
image_utils = layerforge_runtime.image_utils
|
|
grayscale = Image.new("L", (2, 1), 128)
|
|
buffer = io.BytesIO()
|
|
grayscale.save(buffer, format="PNG")
|
|
encoded = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()
|
|
|
|
grayscale_tensor, alpha = image_utils.convert_base64_to_tensor(encoded)
|
|
assert tuple(grayscale_tensor.shape) == (1, 3, 1, 2)
|
|
assert alpha is None
|
|
|
|
rgb_tensor = torch.ones((1, 3, 2, 2), dtype=torch.float32)
|
|
alpha_mask = torch.full((1, 1, 2, 2), 0.25, dtype=torch.float32)
|
|
original_alpha = torch.full((1, 1, 2, 2), 0.75, dtype=torch.float32)
|
|
encoded_rgba = image_utils.convert_tensor_to_base64(rgb_tensor, alpha_mask, original_alpha)
|
|
rgba = Image.open(io.BytesIO(base64.b64decode(encoded_rgba.split(",", 1)[1])))
|
|
assert rgba.mode == "RGBA"
|
|
assert rgba.getpixel((0, 0))[3] == pytest.approx(63, abs=1)
|
|
|
|
encoded_l = image_utils.convert_tensor_to_base64(torch.ones((1, 1, 2, 2), dtype=torch.float32))
|
|
grayscale_result = Image.open(io.BytesIO(base64.b64decode(encoded_l.split(",", 1)[1])))
|
|
assert grayscale_result.mode == "L"
|
|
|
|
|
|
def test_node_cache_flow_and_cleanup_helpers_preserve_state(layerforge_runtime, monkeypatch):
|
|
from PIL import Image
|
|
|
|
node_class = layerforge_runtime.node.LayerForgeNode
|
|
node_class._canvas_cache.update(
|
|
{
|
|
"image": None,
|
|
"mask": None,
|
|
"data_flow_status": {},
|
|
"persistent_cache": {},
|
|
"last_execution_id": None,
|
|
}
|
|
)
|
|
node_class._websocket_data = {}
|
|
node = node_class()
|
|
|
|
image = Image.new("RGB", (2, 1), (0, 255, 0))
|
|
mask = Image.new("L", (2, 1), 255)
|
|
node_class._canvas_cache["image"] = image
|
|
node_class._canvas_cache["mask"] = mask
|
|
node.update_persistent_cache()
|
|
assert node_class._canvas_cache["persistent_cache"] == {"image": image, "mask": mask}
|
|
|
|
node.track_data_flow("render", "complete", {"layers": 2})
|
|
assert node.get_flow_status(node.flow_id)["status"] == "complete"
|
|
assert node.get_flow_status()[node.flow_id]["data_info"] == {"layers": 2}
|
|
assert node.get_cached_data() == {"image": image, "mask": mask}
|
|
assert node.api_get_data("node-1")["success"] is True
|
|
node.store_image(image)
|
|
assert node.get_cached_image().startswith("data:image/png;base64,")
|
|
|
|
node_class._canvas_cache["persistent_cache"] = {"image": image, "mask": mask}
|
|
node_class._canvas_cache["last_execution_id"] = "execution-1"
|
|
monkeypatch.setattr(node, "get_execution_id", lambda: "execution-1")
|
|
node.restore_cache()
|
|
assert node_class._canvas_cache["image"] is image
|
|
assert node_class._canvas_cache["mask"] is mask
|
|
|
|
monkeypatch.setattr(layerforge_runtime.node.time, "time", lambda: 1000)
|
|
node_class._websocket_data = {
|
|
-1: {"timestamp": 999},
|
|
2: {"timestamp": 600},
|
|
3: {"timestamp": 950},
|
|
}
|
|
node_class._cleanup_old_websocket_data()
|
|
assert set(node_class._websocket_data) == {3}
|
|
|
|
assert node.add_image_to_canvas("invalid") is None
|
|
assert node.add_mask_to_canvas("invalid", None) is None
|
|
node_class.setup_routes()
|