nan
f48bbb3d05
- Implement tests for image layer functionality including alpha curve and depth layer behavior. - Create tests for image plane evaluation and modifier input handling. - Add tests to verify the saved asset inventory and its persistence. - Introduce tests for the documentation command to ensure correct project version deployment. - Validate the extension manifest structure and required permissions. - Develop tests for the packaging process, ensuring correct platform handling and dependency management. - Enhance image depth plane group functionality with new split threshold input and surface processing.
126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
import hashlib
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from server.models import onnx_pisa_sr as pisa
|
|
from server.model_manager import ModelManager, ModelSessionCache
|
|
from server import model_catalog
|
|
|
|
|
|
class FakePisa:
|
|
def __init__(self):
|
|
self.noise = np.zeros((1, 4, 64, 64), dtype=np.float16)
|
|
self.stages = []
|
|
self.closed = False
|
|
self.decoder_inputs = []
|
|
|
|
def activate(self, component):
|
|
self.component = component
|
|
self.stages.append(component)
|
|
|
|
def run(self, tensor, final=False):
|
|
if self.component == "encoder":
|
|
return np.zeros((1,8,64,64), dtype=np.float16)
|
|
if self.component == "unet-pixel":
|
|
return np.full_like(tensor, .1)
|
|
if self.component == "unet-full":
|
|
return np.full_like(tensor, .2)
|
|
self.decoder_inputs.append(tensor.copy())
|
|
return np.zeros((1,3,512,512), dtype=np.float16)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
@pytest.mark.parametrize("size", [(1,1), (7,5), (256,256), (513,257)])
|
|
def test_preset_output_size_and_prediction_combination(size):
|
|
model = FakePisa()
|
|
output = pisa.infer(model, Image.new("RGB", size, (32,64,96)))
|
|
assert output.size == (size[0]*2, size[1]*2)
|
|
assert model.stages == ["encoder", "unet-pixel", "unet-full", "decoder"]
|
|
for latent in model.decoder_inputs:
|
|
np.testing.assert_allclose(latent, -.17, atol=.0002)
|
|
assert np.isfinite(np.asarray(output)).all()
|
|
|
|
|
|
def test_cancel_cleans_intermediate_files_and_session(tmp_path):
|
|
model = FakePisa()
|
|
original = pisa.TemporaryDirectory
|
|
def temporary(**kwargs):
|
|
return original(dir=tmp_path, **kwargs)
|
|
checks = 0
|
|
def cancel():
|
|
nonlocal checks
|
|
checks += 1
|
|
return checks == 4
|
|
with patch.object(pisa, "TemporaryDirectory", side_effect=temporary):
|
|
with pytest.raises(RuntimeError, match="cancelled"):
|
|
pisa.infer(model, Image.new("RGB", (256,256)), cancel)
|
|
assert model.closed
|
|
assert not list(tmp_path.iterdir())
|
|
|
|
|
|
def test_stage_releases_previous_session_before_loading(tmp_path):
|
|
np.save(tmp_path / "noise.npy", np.zeros((1,4,64,64), dtype=np.float16))
|
|
model = pisa.PisaSession(tmp_path, "cpu")
|
|
loaded = []
|
|
def create(path, device):
|
|
assert model.session is None
|
|
loaded.append(path.stem)
|
|
return SimpleNamespace(get_providers=lambda: ["CPUExecutionProvider"], disable_fallback=lambda: None)
|
|
with patch.object(pisa, "create_session", side_effect=create):
|
|
model.activate("unet-pixel")
|
|
model.activate("unet-full")
|
|
assert loaded == ["unet-pixel", "unet-full"]
|
|
|
|
|
|
def test_resource_error_closes_pipeline():
|
|
model = FakePisa()
|
|
with patch.object(model, "run", side_effect=RuntimeError("out of memory")):
|
|
with pytest.raises(pisa.UpscaleResourceError):
|
|
pisa.infer(model, Image.new("RGB", (256,256)))
|
|
assert model.closed
|
|
|
|
|
|
def test_pisa_readiness_checks_every_asset_and_corruption(tmp_path):
|
|
content = b"verified"
|
|
digest = hashlib.sha256(content).hexdigest()
|
|
files = tuple((name, len(content), digest) for name in ("encoder.onnx", "noise.npy"))
|
|
spec = SimpleNamespace(directory_name="pisa-sr", r2_files=files)
|
|
directory = tmp_path / "pisa-sr"
|
|
directory.mkdir()
|
|
manager = ModelManager(tmp_path)
|
|
with patch.object(model_catalog, "get_downloadable_model", return_value=spec):
|
|
assert not manager.ready("PISA_SR")
|
|
for name, _, _ in files:
|
|
(directory / name).write_bytes(content)
|
|
assert manager.ready("PISA_SR")
|
|
(directory / "noise.npy").write_bytes(b"damaged!")
|
|
assert not manager.ready("PISA_SR")
|
|
|
|
|
|
def test_pisa_cache_reuse_and_clear(tmp_path):
|
|
cache = ModelSessionCache()
|
|
with patch.object(pisa, "PisaSession", return_value=object()) as create:
|
|
first = cache.get_upscale(tmp_path / "pisa-sr", "cpu")
|
|
assert cache.get_upscale(tmp_path / "pisa-sr", "cpu") is first
|
|
assert create.call_count == 1
|
|
cache.close()
|
|
assert cache.upscale_model is None
|
|
|
|
|
|
def test_coreml_returns_explicit_support_reason(tmp_path):
|
|
with pytest.raises(RuntimeError, match="select CPU"):
|
|
pisa.PisaSession(tmp_path, "coreml")
|
|
|
|
|
|
def test_clear_closes_pisa_components_even_with_external_reference():
|
|
cache = ModelSessionCache()
|
|
model = FakePisa()
|
|
cache.upscale_model = model
|
|
cache.close()
|
|
assert model.closed
|