81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
from pathlib import Path
|
|
|
|
from setuptools import find_packages, setup
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
OPS_DIR = Path("wall_x") / "model" / "core" / "ops"
|
|
CSRC_DIR = OPS_DIR / "csrc"
|
|
|
|
PUBLIC_SCRIPTS = [
|
|
"scripts/compute_norm_stats.py",
|
|
"scripts/convert_lerobot_v3_to_v21.py",
|
|
"scripts/merge_lerobot_v21_datasets.py",
|
|
"scripts/draw_openloop_plot.py",
|
|
"scripts/fake_inference.py",
|
|
"scripts/infer_libero.py",
|
|
"scripts/merge_sharded_weights.py",
|
|
"scripts/merge_tokenizer.py",
|
|
"scripts/run_libero.sh",
|
|
"scripts/run_serving.sh",
|
|
]
|
|
|
|
|
|
def read_readme() -> str:
|
|
readme = ROOT / "README.md"
|
|
return readme.read_text(encoding="utf-8") if readme.exists() else ""
|
|
|
|
|
|
def build_ext_modules():
|
|
binding = CSRC_DIR / "binding.cu"
|
|
if not (ROOT / binding).exists():
|
|
raise RuntimeError(
|
|
"CUDA operator sources are missing. The OSS export must include "
|
|
"wall_x/model/core/ops/csrc/binding.cu."
|
|
)
|
|
|
|
cuda_sources = [binding] + sorted(
|
|
path for path in (ROOT / CSRC_DIR).rglob("*.cu") if path.name != "binding.cu"
|
|
)
|
|
cuda_sources = [
|
|
path if path == binding else path.relative_to(ROOT) for path in cuda_sources
|
|
]
|
|
|
|
return [
|
|
CUDAExtension(
|
|
name="wall_x.model.core.ops._cuda_ext_bin",
|
|
sources=[str(path) for path in cuda_sources],
|
|
include_dirs=[str(CSRC_DIR), str(CSRC_DIR / "common")],
|
|
extra_compile_args={
|
|
"cxx": ["-O3", "-std=c++17"],
|
|
"nvcc": ["-O3", "--use_fast_math", "-std=c++17"],
|
|
},
|
|
)
|
|
]
|
|
|
|
|
|
setup(
|
|
name="wall_x",
|
|
version="1.1.0",
|
|
description="Training and inference code for WALL open-source embodied models.",
|
|
long_description=read_readme(),
|
|
long_description_content_type="text/markdown",
|
|
author="X-Square Robot",
|
|
url="https://github.com/X-Square-Robot/wall-x",
|
|
python_requires=">=3.10",
|
|
packages=find_packages(exclude=("tests", "tests.*")),
|
|
scripts=PUBLIC_SCRIPTS,
|
|
ext_modules=build_ext_modules(),
|
|
cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
|
|
classifiers=[
|
|
"Development Status :: 4 - Beta",
|
|
"Intended Audience :: Developers",
|
|
"Intended Audience :: Science/Research",
|
|
"License :: OSI Approved :: Apache Software License",
|
|
"Programming Language :: Python :: 3",
|
|
"Programming Language :: Python :: 3.10",
|
|
"Programming Language :: Python :: 3.11",
|
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
],
|
|
)
|