nan
e7aeeb21d4
- Implement sync_models.py for preparing and uploading ONNX artifacts. - Create build_node.py for building and validating Blender scripts. - Add docs.py for building and previewing documentation with Mike or MkDocs. - Introduce model_exports.py for exporting PyTorch models to ONNX format. - Develop pack.py for packaging the project with dependencies for different platforms. - Add test_sync_models.py to validate sync_models functionality. - Add test_tool_commands.py to ensure repository tools have correct command scripts. - Initialize tools package with __init__.py.
100 lines
3.4 KiB
Python
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" / "sync_models.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" / "sync_models.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()
|