Add more examples and platforms

.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')

.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
+          '

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 |

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%}")

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

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'