- 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.
146 lines
4.3 KiB
Python
146 lines
4.3 KiB
Python
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import tomllib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from zipfile import ZIP_DEFLATED, ZipFile
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
SOURCE_DIR = PROJECT_ROOT / "src" / "anyimage"
|
|
DIST_DIR = PROJECT_ROOT / "dist"
|
|
PYTHON_VERSION = "3.11"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Platform:
|
|
pip_platform: str
|
|
|
|
|
|
PLATFORM_CONFIGS = {
|
|
"windows-x64": Platform(pip_platform="win_amd64"),
|
|
"macos-arm64": Platform(pip_platform="macosx_12_0_arm64"),
|
|
"linux-x64": Platform(pip_platform="manylinux2014_x86_64"),
|
|
}
|
|
PLATFORMS = tuple(PLATFORM_CONFIGS)
|
|
|
|
|
|
def project_dependencies():
|
|
"""Return Extension wheel requirements declared by the project."""
|
|
project = tomllib.loads(
|
|
(PROJECT_ROOT / "pyproject.toml").read_text(encoding="utf-8")
|
|
)["project"]
|
|
return project.get("dependencies", [])
|
|
|
|
|
|
def download_wheels(platform, destination):
|
|
"""Download default dependencies as wheels for one target platform."""
|
|
dependencies = project_dependencies()
|
|
if not dependencies:
|
|
return []
|
|
config = PLATFORM_CONFIGS[platform]
|
|
command = [
|
|
sys.executable,
|
|
"-m",
|
|
"pip",
|
|
"download",
|
|
*dependencies,
|
|
"--dest",
|
|
str(destination),
|
|
"--find-links",
|
|
str(SOURCE_DIR / "wheels"),
|
|
"--only-binary=:all:",
|
|
f"--python-version={PYTHON_VERSION}",
|
|
f"--platform={config.pip_platform}",
|
|
]
|
|
subprocess.run(command, check=True)
|
|
wheels = sorted(destination.glob("*.whl"))
|
|
if not wheels:
|
|
raise RuntimeError(f"No wheels downloaded for {platform}")
|
|
return wheels
|
|
|
|
|
|
def manifest_for_platform(platform, wheels):
|
|
"""Return manifest text for one platform and its downloaded wheels."""
|
|
text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
|
|
lines = text.splitlines()
|
|
table_start = next(
|
|
(i for i, line in enumerate(lines) if line.startswith("[")),
|
|
len(lines),
|
|
)
|
|
head = lines[:table_start]
|
|
tail = lines[table_start:]
|
|
head = [
|
|
line
|
|
for line in head
|
|
if not line.strip().startswith(("platforms =", "wheels ="))
|
|
]
|
|
insert_at = next(
|
|
(
|
|
i + 1
|
|
for i, line in enumerate(head)
|
|
if line.strip().startswith("blender_version_min")
|
|
),
|
|
len(head),
|
|
)
|
|
head.insert(insert_at, f'platforms = ["{platform}"]')
|
|
wheel_names = ", ".join(
|
|
f'"./wheels/{wheel.name}"' for wheel in wheels
|
|
)
|
|
head.insert(insert_at + 1, f"wheels = [{wheel_names}]")
|
|
return "\n".join(head + tail) + "\n"
|
|
|
|
|
|
def _build_extension(platform, wheels):
|
|
manifest_text = manifest_for_platform(platform, wheels)
|
|
manifest = tomllib.loads(manifest_text)
|
|
output_path = DIST_DIR / (
|
|
f"{manifest['name']}.v{manifest['version']}.{platform}.zip"
|
|
)
|
|
DIST_DIR.mkdir(exist_ok=True)
|
|
with ZipFile(output_path, "w", ZIP_DEFLATED) as archive:
|
|
for source_path in sorted(SOURCE_DIR.rglob("*")):
|
|
if not source_path.is_file():
|
|
continue
|
|
if "__pycache__" in source_path.parts:
|
|
continue
|
|
if (
|
|
source_path.suffix.startswith(".blend")
|
|
and source_path.suffix != ".blend"
|
|
):
|
|
continue
|
|
relative_path = source_path.relative_to(SOURCE_DIR)
|
|
if relative_path.parts[0] == "wheels":
|
|
continue
|
|
if source_path.name == "blender_manifest.toml":
|
|
continue
|
|
archive.write(source_path, relative_path)
|
|
archive.writestr("blender_manifest.toml", manifest_text)
|
|
for wheel in wheels:
|
|
archive.write(wheel, Path("wheels") / wheel.name)
|
|
return output_path
|
|
|
|
|
|
def build_extension(platform):
|
|
if platform not in PLATFORMS:
|
|
raise ValueError(f"Unsupported platform: {platform}")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
wheels = download_wheels(platform, Path(directory))
|
|
return _build_extension(platform, wheels)
|
|
|
|
|
|
def build_all():
|
|
return [build_extension(platform) for platform in PLATFORMS]
|
|
|
|
|
|
def main():
|
|
extension_paths = build_all()
|
|
for extension_path in extension_paths:
|
|
size_mib = extension_path.stat().st_size / (1024 * 1024)
|
|
print(f"Built {extension_path} ({size_mib:.2f} MiB)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|