Files
AnyImage/tests/test_sync_models.py
nan 89087fc2e2 feat: Add new Geometry Nodes groups for image relief and shadeless shader
- Implemented the "O Image Relief Plane" Geometry Nodes group to create a fixed-base image relief using depth images and various parameters.
- Created the "O Shadeless" shader node group to display color without scene lighting while preserving shadows and reflections.
- Introduced a preview feature for node groups, allowing for interactive visualization of node arrangements in an HTML format.
- Added a template HTML file for rendering node previews with customizable views and controls.
2026-09-09 21:24:57 +08:00

100 lines
3.4 KiB
Python

import importlib.util
import sys
import unittest
from pathlib import Path
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_sync_models():
path = PROJECT_ROOT / "tools" / "models" / "cli.py"
spec = importlib.util.spec_from_file_location("sync_models", path)
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
class SyncModelsTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.sync = load_sync_models()
def test_r2_is_only_an_upload_destination(self):
for model in self.sync.MODEL_FILES:
self.assertNotIn("models.omoolab.xyz", model.source_url)
if model.build:
self.assertFalse(model.source_url)
def test_local_exports_use_official_pytorch_weights(self):
builds = self.sync.MODEL_BUILDS
self.assertEqual(
set(builds),
{"REALESRGAN_GENERAL_WDN_X4V3", "SAFMN_REAL_X4"},
)
for build in builds.values():
self.assertTrue(build.source_url.startswith("https://github.com/"))
self.assertTrue(build.source_filename.endswith(".pth"))
self.assertEqual(len(build.source_sha256), 64)
def test_prepare_model_builds_missing_local_export(self):
model = next(model for model in self.sync.MODEL_FILES if model.build)
with (
patch.object(self.sync, "verify_file", return_value=False),
patch.object(self.sync, "build_model") as build,
patch.object(self.sync, "download_file") as download,
):
self.sync.prepare_model(model)
build.assert_called_once_with(model)
download.assert_not_called()
def test_prepare_model_downloads_upstream_onnx(self):
model = next(
model
for model in self.sync.MODEL_FILES
if model.source_url and not model.build
)
with (
patch.object(self.sync, "verify_file", return_value=False),
patch.object(self.sync, "download_file") as download,
):
self.sync.prepare_model(model)
download.assert_called_once_with(
model.source_url,
model.destination,
model.size,
model.sha256,
)
def test_upload_copies_declared_files_without_remote_deletion(self):
model = self.sync.ModelFile(
key="TEST",
folder="folder",
filename="model.onnx",
size=1,
sha256="hash",
)
with (
patch.object(self.sync.shutil, "which", return_value="rclone"),
patch.object(self.sync.subprocess, "run") as run,
):
self.sync.upload_models((model,))
command = run.call_args.args[0]
self.assertEqual(command[1], "copyto")
self.assertNotIn("sync", command)
self.assertEqual(command[-2], model.remote)
def test_verify_file_checks_size_and_sha256(self):
path = PROJECT_ROOT / "tools" / "models" / "cli.py"
size = path.stat().st_size
digest = self.sync.file_sha256(path)
self.assertTrue(self.sync.verify_file(path, size, digest))
with self.assertRaisesRegex(RuntimeError, "unexpected size"):
self.sync.verify_file(path, size - 1, digest)
if __name__ == "__main__":
unittest.main()