nan
e8129fe495
- Added validation documentation for AI model tiers. - Introduced `background.py` for background removal logic, including inference functions for alpha and foreground extraction. - Created `onnx_birefnet.py` to handle BiRefNet model sessions and alpha inference. - Developed tests for background models, ensuring proper dispatch and session reuse. - Added tests for BiRefNet ONNX model to validate alpha output and error handling for invalid values. - Implemented model export scripts for BiRefNet and HAT Sharper, ensuring proper weight compression and licensing. - Created `deform_sampling.py` for deformable convolution operations used in BiRefNet. - Added utility functions for compressing ONNX weights without altering graph arithmetic. - Introduced a reference plane node for depth field adjustments in geometry processing.
109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
from tests.anyimage.server.support import load_server_module, FakeModels
|
|
|
|
|
|
class ModelInferenceTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.moge2_model = load_server_module("models.moge")
|
|
|
|
cls.upscale_model = load_server_module("models.upscale")
|
|
|
|
cls.onnx_upscale = load_server_module("models.onnx_upscale")
|
|
|
|
|
|
def test_moge2_batch_releases_only_after_final_image(self):
|
|
prediction = {
|
|
"points": np.ones((2, 2, 3), dtype=np.float32),
|
|
"depth": np.ones((2, 2), dtype=np.float32),
|
|
"normal": np.ones((2, 2, 3), dtype=np.float32),
|
|
"mask": np.ones((2, 2), dtype=bool),
|
|
"intrinsics": np.eye(3, dtype=np.float32),
|
|
}
|
|
with patch.object(
|
|
self.moge2_model,
|
|
"infer_one",
|
|
return_value=prediction,
|
|
) as infer_one:
|
|
frames = self.moge2_model.infer(
|
|
FakeModels(),
|
|
{"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
|
|
[Path("a.png"), Path("b.png")],
|
|
include_points=True,
|
|
)
|
|
|
|
self.assertEqual(len(frames), 2)
|
|
self.assertEqual(
|
|
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
|
|
[False, True],
|
|
)
|
|
|
|
|
|
def test_upscale_resource_error_releases_session_and_is_readable(self):
|
|
from PIL import Image
|
|
|
|
class Models:
|
|
def __init__(self, directory):
|
|
self.model_directory = directory
|
|
self.released = False
|
|
|
|
def directory(self, _model):
|
|
return self.model_directory
|
|
|
|
def get_upscale(self, _directory, _device):
|
|
return object()
|
|
|
|
def release_upscale(self):
|
|
self.released = True
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "input.png"
|
|
Image.new("RGB", (8, 8)).save(path)
|
|
models = Models(Path(directory))
|
|
with patch.object(
|
|
self.onnx_upscale,
|
|
"infer",
|
|
side_effect=self.onnx_upscale.UpscaleResourceError(
|
|
"provider memory exhausted"
|
|
),
|
|
):
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"Reduce Maximum AI Input Size or select CPU",
|
|
):
|
|
self.upscale_model.infer_one(
|
|
models,
|
|
{
|
|
"model": "REALESRGAN_GENERAL_WDN_X4V3",
|
|
"device": "directml",
|
|
"max_input_size": 2048,
|
|
},
|
|
path,
|
|
)
|
|
|
|
self.assertTrue(models.released)
|
|
|
|
|
|
upscale = load_server_module("models.upscale")
|
|
|
|
@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "HAT_GAN_X4_SHARPER", "REALESRGAN_X4PLUS"])
|
|
def test_upscale_models_resize_prediction_then_restore_alpha(tmp_path, key):
|
|
rng = np.random.default_rng(1)
|
|
source = Image.fromarray(rng.integers(0,256,(5,7,4),dtype=np.uint8))
|
|
path = tmp_path / "source.png"
|
|
source.save(path)
|
|
prediction = Image.fromarray(rng.integers(0,256,(20,28,3),dtype=np.uint8))
|
|
manager = SimpleNamespace(directory=lambda _: tmp_path, get_upscale=lambda *args: object())
|
|
with patch("server.models.onnx_upscale.infer", return_value=prediction) as infer:
|
|
output = upscale.infer_one(manager, {"model":key,"device":"cpu"}, path)
|
|
assert infer.call_args.kwargs["tile_border"] == (48 if key == "HAT_GAN_X4_SHARPER" else 16)
|
|
np.testing.assert_array_equal(output.convert("RGB"), prediction.resize((14,10), Image.Resampling.LANCZOS))
|
|
np.testing.assert_array_equal(output.getchannel("A"), source.getchannel("A").resize((14,10), Image.Resampling.LANCZOS))
|