106 lines
3.9 KiB
Python
106 lines
3.9 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_infers_one_geometry_frame(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:
|
|
frame = self.moge2_model.infer(
|
|
FakeModels(),
|
|
{"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
|
|
Path("input.png"),
|
|
include_points=True,
|
|
)
|
|
|
|
self.assertEqual(frame.depth.shape, (2, 2))
|
|
infer_one.assert_called_once()
|
|
|
|
|
|
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))
|