From 7864e5e5bcfaef9d3ce0259eb58936e1c1888631 Mon Sep 17 00:00:00 2001 From: elkoled Date: Wed, 16 Sep 2026 13:07:56 -0700 Subject: [PATCH] add more examples and platforms --- .github/test_examples.py | 43 +++++++++++++++++ .github/workflows/test.yml | 56 ++++++++++++++++++---- README.md | 98 +++++++++++++++++++++++--------------- examples/04_classify.py | 24 ++++++++++ pyproject.toml | 4 -- setup.sh | 80 +++++++++++++++++++++---------- tools/export.py | 7 +++ tools/setup.py | 33 +++++++++++++ 8 files changed, 269 insertions(+), 76 deletions(-) create mode 100644 .github/test_examples.py create mode 100644 examples/04_classify.py create mode 100644 tools/setup.py diff --git a/.github/test_examples.py b/.github/test_examples.py new file mode 100644 index 0000000..9fd4fe1 --- /dev/null +++ b/.github/test_examples.py @@ -0,0 +1,43 @@ +import os +from pathlib import Path +import subprocess +import sys +import tempfile + +import cv2 +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +os.environ['DEV'] = 'CPU' + + +def run(example, *args, cwd=ROOT): + result = subprocess.run([sys.executable, str(ROOT / 'examples' / example), *args], + cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, timeout=900) + print(result.stdout, flush=True) + result.check_returncode() + return result.stdout + + +for model in ('qwen3.5:0.8b', 'llama3.2:1b'): + assert run('01_chat.py', '--model', model, '--benchmark', '3', '--max_context', '256').count('tok/s') == 3 + +image = cv2.imread(str(ROOT / 'zidane.jpg')) +assert image is not None +for model in ('yolo', 'segment'): + with tempfile.TemporaryDirectory() as directory: + folder = Path(directory) + run('02_vision.py', str(ROOT / 'zidane.jpg'), '--model', model, '--output', str(folder / 'result.jpg')) + annotated = cv2.imread(str(folder / 'result.jpg')) + assert annotated is not None and annotated.shape == image.shape + assert np.abs(annotated.astype(float) - image).mean() > 1 + for i in range(10): + assert cv2.imwrite(str(folder / f'input-{i:05d}.jpg'), image) + run('03_camera.py', '--source', str(folder / 'input-%05d.jpg'), '--model', model, cwd=folder) + frames = sorted((folder / 'frames').glob('*.jpg')) + assert len(frames) == 10 + for frame in frames: + saved = cv2.imread(str(frame)) + assert saved is not None and saved.shape == image.shape + +assert 'Samoyed:' in run('04_classify.py') diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index b0a9ffc..431bf65 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,15 +1,55 @@ name: examples -on: [push, pull_request] +on: + push: + branches: [master] + pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: examples-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + jobs: cpu: + timeout-minutes: 60 strategy: + fail-fast: false matrix: - os: [ubuntu-22.04, ubuntu-24.04, ubuntu-24.04-arm] + os: [ubuntu-22.04, ubuntu-24.04, ubuntu-24.04-arm, macos-14, macos-15] runs-on: ${{ matrix.os }} steps: - - uses: actions/checkout@v4 - - run: ./setup.sh - - run: | - .venv/bin/python examples/01_chat.py --help - .venv/bin/python examples/02_vision.py chestnut.png - .venv/bin/python examples/02_vision.py chestnut.png --model segment + - uses: actions/checkout@v7 + - name: Setup from a clean environment + run: ./setup.sh + - name: Setup is repeatable + run: ./setup.sh + - name: Run every example + run: .venv/bin/python .github/test_examples.py + - name: Setup leaves tracked files unchanged + run: git diff --exit-code + + linux-distros: + runs-on: ubuntu-24.04 + timeout-minutes: 45 + strategy: + fail-fast: false + matrix: + image: [ubuntu:22.04, ubuntu:24.04, debian:12, debian:13, fedora:43, archlinux:base, opensuse/tumbleweed:latest] + steps: + - uses: actions/checkout@v7 + - name: Fresh container setup and inference + env: + IMAGE: ${{ matrix.image }} + run: | + docker run --rm -v "$PWD:/src:ro" "$IMAGE" bash -c ' + set -e + mkdir /work + cp -a /src/. /work/ + cd /work + ./setup.sh + ./setup.sh + .venv/bin/python .github/test_examples.py + ' diff --git a/README.md b/README.md index dd6db11..cf730b9 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # chestnut -Run llm, object detection, and segmentation with tinygrad on your PC or chestnut. +Run chat, object detection, segmentation, and image classification with tinygrad on your PC or Chestnut. ## Setup @@ -10,74 +10,96 @@ cd chestnut ./setup.sh ``` -## Run on Chestnut +## Run on PC -Plug in the 12V power and connect the USB3 cable from the USB3.2 port to your PC or comma. +Activate the environment in each new terminal: -![Chestnut connections](chestnut.png) +```sh +source .venv/bin/activate +``` + +These commands run on your PC's CPU. + +### Chat -Check the connection: +Chat with Qwen 3.5 0.8B (default) or Llama 3.2 1B. Models download on first run. ```sh -.venv/bin/python tools/usb.py +python examples/01_chat.py +python examples/01_chat.py --model llama3.2:1b ``` -Expected: `Chestnut GPU check passed.` +### YOLO -Examples default to `DEV=CPU` to run on PC. -Prefix any example command with `DEV=USB+AMD:LLVM` to use Chestnut. +Detect objects or draw segmentation masks. -## Examples +```sh +python examples/02_vision.py zidane.jpg --output boxes.jpg +python examples/02_vision.py zidane.jpg --model segment --output masks.jpg +``` -### Chat +### Camera + +Run YOLO26 detection or segmentation on a webcam. +Use `--frames 100` for more. `--source` also accepts a video path or stream URL. ```sh -.venv/bin/python examples/01_chat.py -.venv/bin/python examples/01_chat.py --model llama3.2:1b +python examples/03_camera.py --source 0 +python examples/03_camera.py --source 0 --model segment ``` -Starts an interactive chat with Qwen 3.5 0.8B by default, or Llama 3.2 1B with `--model llama3.2:1b`. Models download on first run. - -### YOLO +### Image classification -Detect objects or draw segmentation masks on a sample image. Open `boxes.jpg` or `masks.jpg` to see the result. +Print the five most likely labels with ResNet18. ```sh -.venv/bin/python examples/02_vision.py zidane.jpg --output boxes.jpg -.venv/bin/python examples/02_vision.py zidane.jpg --model segment --output masks.jpg +python examples/04_classify.py +python examples/04_classify.py photo.jpg ``` -### Camera: YOLO26 detection and segmentation +## Run on Chestnut + +Plug in the 12V power and connect the USB3 cable from Chestnut's USB3.2 port to your PC or comma. -Runs YOLO26 on webcam frames. Detection draws labeled boxes. `--model segment` adds masks around each object. -Saves 10 annotated frames to `frames/` without a live preview. Use `--frames 100` for more. +![Chestnut connections](chestnut.png) + +In the activated environment, check the connection: ```sh -.venv/bin/python examples/03_camera.py --source 0 -.venv/bin/python examples/03_camera.py --source 0 --model segment +python tools/usb.py ``` -`--source` also accepts a video path or stream URL. +Expected: `Chestnut GPU check passed.` -On a comma device with openpilot installed at `/data/openpilot`, use its camera stream: +Prefix any example command with `DEV=USB+AMD:LLVM` to run it on Chestnut's GPU: ```sh -DEV=USB+AMD:LLVM .venv/bin/python examples/03_camera.py --source comma -DEV=USB+AMD:LLVM .venv/bin/python examples/03_camera.py --source comma --model segment +DEV=USB+AMD:LLVM python examples/01_chat.py +DEV=USB+AMD:LLVM python examples/02_vision.py zidane.jpg --output boxes.jpg +DEV=USB+AMD:LLVM python examples/03_camera.py --source 0 ``` -## Performance -PC CPU: Threadripper PRO 5945WX +## Comma camera -| Model | PC CPU | Chestnut | Speedup | -| --- | ---: | ---: | ---: | -| YOLO26n | 524.57 ms | 15.12 ms | 35× | -| YOLO26n-seg | 705.44 ms | 16.44 ms | 43× | +On a comma device, run setup and activate the environment as above. Requires openpilot at `/data/openpilot`. +`--source comma` selects its camera. Add `--model segment` for segmentation. +```sh +# Comma CPU +python examples/03_camera.py --source comma -### Chat +# Chestnut GPU connected to comma +DEV=USB+AMD:LLVM python examples/03_camera.py --source comma +``` + +## Performance + +PC CPU: Threadripper PRO 5945WX. Comma CPU: Qualcomm SDM845. -| Model | PC CPU | Chestnut | Speedup | +| Model | PC CPU | Comma CPU | Chestnut GPU | | --- | ---: | ---: | ---: | -| Qwen 3.5 0.8B (Q8_0) | 2.96 tokens/s | 44.90 tokens/s | 15.2× | -| Llama 3.2 1B Instruct (Q6_K) | 0.91 tokens/s | 24.93 tokens/s | 27.2× | +| YOLO26n | 526.07 ms | 1846.42 ms | 7.02 ms | +| YOLO26n-seg | 704.56 ms | 2349.27 ms | 7.98 ms | +| ResNet18 | 246.52 ms | 610.91 ms | 11.98 ms | +| Qwen 3.5 0.8B (Q8_0) | 2.98 tokens/s | 0.76 tokens/s | 45.46 tokens/s | +| Llama 3.2 1B Instruct (Q6_K) | 0.94 tokens/s | 0.34 tokens/s | 25.21 tokens/s | diff --git a/examples/04_classify.py b/examples/04_classify.py new file mode 100644 index 0000000..c28fb22 --- /dev/null +++ b/examples/04_classify.py @@ -0,0 +1,24 @@ +import argparse +import os +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +os.environ.setdefault('DEV', 'CPU') +os.environ.setdefault('XDG_CACHE_HOME', str(ROOT / '.cache')) + +from PIL import Image +from torchvision.models import ResNet18_Weights +from tinygrad import Tensor +from tinygrad.helpers import fetch +from tinygrad.nn.onnx import OnnxRunner + +parser = argparse.ArgumentParser(description='Classify an image with ResNet18.') +parser.add_argument('image', nargs='?') +args = parser.parse_args() +image = Image.open(args.image or fetch('https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg')).convert('RGB') +weights = ResNet18_Weights.DEFAULT +inputs = weights.transforms()(image).unsqueeze(0).numpy() +model = OnnxRunner(str(ROOT / 'models/resnet18.onnx')) +scores = model({'image': Tensor(inputs)})['scores'].softmax().numpy()[0] +for i in scores.argsort()[-5:][::-1]: + print(f"{weights.meta['categories'][i]}: {scores[i]:.1%}") diff --git a/pyproject.toml b/pyproject.toml index f077760..4792447 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,3 @@ torchvision = {index = "pytorch-cpu"} name = "pytorch-cpu" url = "https://download.pytorch.org/whl/cpu" explicit = true - -[tool.ruff] -line-length = 140 -indent-width = 2 diff --git a/setup.sh b/setup.sh index 9004b5b..458c533 100755 --- a/setup.sh +++ b/setup.sh @@ -1,62 +1,90 @@ #!/usr/bin/env bash set -euo pipefail cd "$(dirname "$0")" -export PATH="$HOME/.local/bin:$PATH" +export PATH="$HOME/.local/bin:/opt/homebrew/bin:$PATH" export UV_CACHE_DIR="${UV_CACHE_DIR:-$PWD/.cache/uv}" +export UV_HTTP_TIMEOUT=200 UV_HTTP_RETRIES=5 +mkdir -p .cache as_root=() if [ "$EUID" -ne 0 ]; then as_root=(sudo); fi -run_root() { - if [ "$EUID" -ne 0 ] && ! command -v sudo >/dev/null; then - echo 'Install sudo or run ./setup.sh as root.' >&2 - return 1 - fi - "${as_root[@]}" "$@" -} +download() { curl --retry 5 --retry-delay 5 --retry-all-errors -fLsS "$1" -o "$2"; } + +case "$(uname -s):$(uname -m)" in + Linux:x86_64|Linux:aarch64) + getconf GNU_LIBC_VERSION >/dev/null || { echo 'glibc Linux is required.' >&2; exit 1; } + ;; + Darwin:arm64) + [ "$(sw_vers -productVersion | cut -d. -f1)" -ge 14 ] || { echo 'macOS 14+ is required.' >&2; exit 1; } + command -v brew >/dev/null || { echo 'Install Homebrew and rerun setup.' >&2; exit 1; } + brew list --versions llvm@21 libusb >/dev/null 2>&1 || brew install llvm@21 libusb + export PATH="/opt/homebrew/opt/llvm@21/bin:$PATH" + ;; + *) echo 'Use x86_64/aarch64 Linux or Apple Silicon macOS 14+.' >&2; exit 1 ;; +esac -if ! command -v clang >/dev/null || ! command -v curl >/dev/null || +if [ "$(uname -s)" = Linux ] && { ! command -v clang >/dev/null || ! command -v curl >/dev/null || + ! command -v awk >/dev/null || ! command -v tar >/dev/null || ! command -v gzip >/dev/null || ! ldconfig -p 2>/dev/null | grep -E 'libLLVM(-|\.so\.)(19|20|21)' >/dev/null || - ! ldconfig -p 2>/dev/null | grep -F 'libusb-1.0.so' >/dev/null; then + ! ldconfig -p 2>/dev/null | grep -F 'libusb-1.0.so' >/dev/null; }; then if command -v apt-get >/dev/null; then - run_root apt-get update - llvm_package=llvm + "${as_root[@]}" apt-get update + llvm_package= for version in 21 20 19; do if apt-cache show "libllvm$version" >/dev/null 2>&1; then llvm_package="libllvm$version"; break; fi done - run_root apt-get install -y --no-install-recommends ca-certificates curl clang "$llvm_package" libusb-1.0-0 + if [ -z "$llvm_package" ]; then + # Ubuntu 22.04 and Debian 12 need upstream LLVM for the GPU compiler. + . /etc/os-release + codename="${UBUNTU_CODENAME:-${VERSION_CODENAME:-}}" + case "$codename" in + jammy|bookworm) ;; + *) echo 'Install LLVM 19–21 and rerun setup.' >&2; exit 1 ;; + esac + "${as_root[@]}" apt-get install -y --no-install-recommends ca-certificates curl + download https://apt.llvm.org/llvm-snapshot.gpg.key .cache/llvm.asc + "${as_root[@]}" install -Dm644 .cache/llvm.asc /etc/apt/keyrings/chestnut-llvm.asc + echo "deb [signed-by=/etc/apt/keyrings/chestnut-llvm.asc] https://apt.llvm.org/$codename/ llvm-toolchain-$codename-20 main" | + "${as_root[@]}" tee /etc/apt/sources.list.d/chestnut-llvm.list >/dev/null + "${as_root[@]}" apt-get update + llvm_package=libllvm20 + fi + "${as_root[@]}" apt-get install -y --no-install-recommends ca-certificates curl clang "$llvm_package" libusb-1.0-0 gawk tar gzip elif command -v dnf >/dev/null; then - run_root dnf install -y clang llvm-libs libusb1 curl + "${as_root[@]}" dnf install -y ca-certificates clang llvm-libs libusb1 curl gawk tar gzip elif command -v pacman >/dev/null; then - run_root pacman -S --needed --noconfirm clang llvm-libs libusb curl + "${as_root[@]}" pacman -Syu --needed --noconfirm ca-certificates clang llvm20-libs libusb curl gawk tar gzip elif command -v zypper >/dev/null; then - run_root zypper --non-interactive install clang llvm libusb-1_0-0 curl + "${as_root[@]}" zypper --non-interactive refresh + "${as_root[@]}" zypper --non-interactive install ca-certificates clang libLLVM20 libusb-1_0-0 curl gawk tar gzip else - echo 'Install clang, curl, LLVM 19+, and libusb with your package manager.' >&2 + echo 'Install clang, curl, LLVM 19–21, and libusb with your package manager.' >&2 exit 1 fi fi -# Like openpilot, allow access to the device without running examples as root. -if [ -d /etc/udev/rules.d ] && command -v udevadm >/dev/null && +if [ -d /etc/udev/rules.d ] && command -v udevadm >/dev/null && command -v findmnt >/dev/null && [[ ",$(findmnt -n -o OPTIONS -T /etc/udev/rules.d)," != *,ro,* ]]; then rules='SUBSYSTEM=="usb", ATTR{idVendor}=="3801", ATTR{idProduct}=="0001", MODE="0666" SUBSYSTEM=="usb", ATTR{idVendor}=="add1", ATTR{idProduct}=="0001", MODE="0666"' if [ "$(cat /etc/udev/rules.d/11-chestnut.rules 2>/dev/null || true)" != "$rules" ]; then - run_root tee /etc/udev/rules.d/11-chestnut.rules >/dev/null <<< "$rules" - run_root udevadm control --reload-rules && run_root udevadm trigger --subsystem-match=usb || true + "${as_root[@]}" tee /etc/udev/rules.d/11-chestnut.rules >/dev/null <<< "$rules" + if "${as_root[@]}" udevadm control --reload-rules; then + "${as_root[@]}" udevadm trigger --subsystem-match=usb || true + fi fi fi if ! command -v uv >/dev/null; then - curl --retry 5 --retry-delay 5 --retry-all-errors -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh + download https://astral.sh/uv/install.sh .cache/install-uv.sh + UV_NO_MODIFY_PATH=1 sh .cache/install-uv.sh fi uv sync --locked --python 3.12 +.venv/bin/python tools/setup.py .venv/bin/python tools/export.py if [ ! -f zidane.jpg ]; then - mkdir -p .cache - curl --retry 5 --retry-delay 5 --retry-all-errors -fLsS https://ultralytics.com/images/zidane.jpg -o .cache/zidane.jpg + download https://ultralytics.com/images/zidane.jpg .cache/zidane.jpg mv .cache/zidane.jpg zidane.jpg fi -echo "Ready. PC: .venv/bin/python examples/02_vision.py zidane.jpg" -echo "Chestnut: DEV=USB+AMD:LLVM .venv/bin/python examples/02_vision.py zidane.jpg" +echo 'Ready. Activate with: source .venv/bin/activate' diff --git a/tools/export.py b/tools/export.py index df42b11..68a902a 100644 --- a/tools/export.py +++ b/tools/export.py @@ -1,8 +1,15 @@ from pathlib import Path from ultralytics import YOLO +import torch +from torchvision.models import resnet18, ResNet18_Weights folder = Path(__file__).resolve().parents[1] / 'models' folder.mkdir(exist_ok=True) for name in ('yolo26n', 'yolo26n-seg'): if not (folder / f'{name}.onnx').exists(): YOLO(str(folder / f'{name}.pt')).export(format='onnx', imgsz=320, opset=17, simplify=False, device='cpu', nms=False) + +if not (folder / 'resnet18.onnx').exists(): + model = resnet18(weights=ResNet18_Weights.DEFAULT).eval() + torch.onnx.export(model, torch.zeros(1, 3, 224, 224), folder / 'resnet18.onnx', + input_names=['image'], output_names=['scores'], opset_version=17, dynamo=False) diff --git a/tools/setup.py b/tools/setup.py new file mode 100644 index 0000000..09c19fb --- /dev/null +++ b/tools/setup.py @@ -0,0 +1,33 @@ +import os +from pathlib import Path +import re +import sys +import sysconfig + +if sys.platform == 'darwin': + candidates = [Path(f'/opt/homebrew/opt/llvm@{version}/lib/libLLVM.dylib') for version in (21, 20, 19)] +else: + directories = [Path('/usr/lib64'), Path('/usr/lib'), Path('/lib'), Path('/lib64'), + Path('/usr/lib') / sysconfig.get_config_var('MULTIARCH')] + libraries = [path for directory in directories for path in directory.glob('libLLVM*') if path.is_file()] + candidates = [path for version in (21, 20, 19) for path in libraries + if re.search(rf'libLLVM(?:-|\.so\.){version}(?:\.so|\.|$)', path.name)] + +library = next((path for path in candidates if path.is_file()), None) +if library is None: + raise SystemExit('LLVM 19–21 was not found. Install its shared library and rerun setup.') + +# Persist distro-specific LLVM paths for every .venv/bin/python command. +configuration = Path(sysconfig.get_path('purelib')) / 'chestnut_native.pth' +configuration.write_text(f'import os; os.environ.setdefault("LLVM_PATH", {str(library)!r})\n') +os.environ["LLVM_PATH"] = str(library) +os.environ["DEV"] = "CPU" + +from tinygrad import Tensor +from tinygrad.runtime.autogen import libusb +from tinygrad.runtime.support.compiler_llvm import LLVMCompiler + +assert libusb.libusb_get_version() +assert LLVMCompiler("AMDGPU", "gfx1200", "").compile('define amdgpu_kernel void @test() { ret void }') +assert Tensor([1, 2, 3, 4]).square().sum().item() == 30 +print("CPU, LLVM, and libusb checked.")