AnyImage/tests/test_model_pisa_sr.py

import hashlib
import sys
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
import pytest
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src/anyimage"))
from server.models import onnx_pisa_sr as pisa
from server.models import upscale
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")


@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4", "REALESRGAN_X4PLUS"])
def test_existing_models_resize_full_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):
        output = upscale.infer_one(manager, {"model":key,"device":"cpu"}, path, preserve_alpha=True)
    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))


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