Add cutout shell and volume node assets for depth volume processing

- Implemented `cutout_shell.py` to build front normals and shell positions for balloon geometry.
- Created `cutout_volume.py` to establish smooth center fields and manage balloon surface attributes.
- Introduced `pisa_sr_export.py` for exporting PiSA-SR components with dependencies and asset finalization.
This commit is contained in:
2026-09-08 17:18:03 +08:00
parent 7dec3b7ca4
commit 196b964efc
71 changed files with 3440 additions and 1988 deletions

144
tests/test_model_pisa_sr.py Normal file
View File

@@ -0,0 +1,144 @@
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