One Shot Build Script For Running 28M AI Model On ESP32-S3

build_and_flash.sh

#!/usr/bin/env bash
# esp32-ai: clone, train, export, build, and flash — end to end
# Reproduces github.com/slvDev/esp32-ai from a bare board.
# Requires: ESP32-S3 N16R8 (16MB flash, 8MB PSRAM).
#
# Usage:
#   ./build_and_flash.sh              # full run: data prep, train, export, build, flash
#   ./build_and_flash.sh --skip-train # reuse an existing firmware/model/model.bin, just build + flash
set -euo pipefail

SKIP_TRAIN=0
for arg in "$@"; do
  case "$arg" in
    --skip-train) SKIP_TRAIN=1 ;;
    *) echo "Unknown option: $arg"; exit 1 ;;
  esac
done

REPO_URL="https://github.com/slvDev/esp32-ai.git"
REPO_DIR="esp32-ai"

# ---- -1. Clone the repo if we're not already inside it ----
if [ ! -f "pyproject.toml" ] && [ ! -f "$REPO_DIR/pyproject.toml" ]; then
  echo ">> Cloning $REPO_URL ..."
  git clone "$REPO_URL"
fi
if [ -f "$REPO_DIR/pyproject.toml" ]; then
  cd "$REPO_DIR"
fi
echo ">> Working in $(pwd)"

# ---- 0. Auto-detect the board's port ----
echo ">> Looking for a connected ESP32-S3..."
CANDIDATE_PORTS=(/dev/cu.usbmodem*)

if [ ! -e "${CANDIDATE_PORTS[0]}" ]; then
  echo "No /dev/cu.usbmodem* device found. Is the board plugged in?"
  exit 1
elif [ "${#CANDIDATE_PORTS[@]}" -eq 1 ]; then
  ESP32_PORT="${CANDIDATE_PORTS[0]}"
  echo ">> Found one port: $ESP32_PORT"
else
  echo ">> Found multiple candidate ports:"
  select choice in "${CANDIDATE_PORTS[@]}"; do
    if [ -n "$choice" ]; then
      ESP32_PORT="$choice"
      break
    fi
  done
fi
export ESP32_PORT
echo ">> Using port: $ESP32_PORT"

uvx --from esptool esptool.py --port "$ESP32_PORT" flash-id | head -20
# Expect: Detected flash size: 16MB, Embedded PSRAM 8MB

# ---- 1. Toolchain ----
echo ">> Installing toolchain..."
brew install arduino-cli
arduino-cli config init || echo "Config already exists, continuing..."
arduino-cli config add board_manager.additional_urls \
  https://espressif.github.io/arduino-esp32/package_esp32_index.json
arduino-cli core update-index
arduino-cli core install esp32:esp32@3.3.10
arduino-cli lib install "Adafruit SH110X"

uv sync

# ---- 2. Data ----
if [ "$SKIP_TRAIN" -eq 1 ]; then
  echo ">> --skip-train set, skipping data prep..."
else
  echo ">> Preparing TinyStories data..."
  uv run python data/prepare.py --vocab 32768
  uv run python src/gen_assets.py
fi

# ---- 3. Train (~13 min on Apple silicon MPS) ----
if [ "$SKIP_TRAIN" -eq 1 ]; then
  echo ">> --skip-train set, skipping training..."
else
  echo ">> Training model (this takes a while)..."
  uv run python src/train.py --arm ple --vocab 32768 --d-model 96 --n-layers 6 \
    --ple-dim 128 --target-core 560000 --batch-size 16 --seq-len 256 \
    --steps 5000 --seed 0 --tag cleandeploy
fi

# ---- 4. Export and verify on host ----
if [ "$SKIP_TRAIN" -eq 1 ]; then
  if [ ! -f "firmware/model/model.bin" ]; then
    echo "No firmware/model/model.bin found. Run once without --skip-train first."
    exit 1
  fi
  echo ">> --skip-train set, reusing existing firmware/model/model.bin"
else
  echo ">> Exporting and verifying against PyTorch golden..."
  uv run python src/export.py
  cc -O3 -o /tmp/esp32-llm-verify firmware/host_verify/verify.c -lm
  /tmp/esp32-llm-verify firmware/model/model.bin firmware/model/golden.txt
  # Expect: PASS: C matches PyTorch golden
fi

# ---- 5. Build firmware ----
echo ">> Compiling firmware..."

SKETCH_FILE="firmware/esp32_llm/esp32_llm.ino"
echo ">> Disabling USE_DISPLAY (no OLED wired up in this setup)..."
sed -i.bak -E "s/(#define[[:space:]]+USE_DISPLAY[[:space:]]+)1/\10/" "$SKETCH_FILE"

FQBN='esp32:esp32:esp32s3:UploadSpeed=921600,USBMode=hwcdc,CDCOnBoot=default,UploadMode=default,CPUFreq=240,FlashMode=qio,FlashSize=16M,PartitionScheme=custom,PSRAM=opi,DebugLevel=info'

arduino-cli compile \
  --fqbn "$FQBN" \
  --build-property compiler.optimization_flags=-O3 \
  --build-path /tmp/esp32-llm-uart \
  firmware/esp32_llm

# ---- 6. Flash: app, then model (model only needs to happen once) ----
echo ">> Flashing app..."
arduino-cli upload -p "$ESP32_PORT" \
  --fqbn "$FQBN" \
  --input-dir /tmp/esp32-llm-uart firmware/esp32_llm

echo ">> Flashing model into custom partition at 0x110000 (~1m45s)..."
if [ "$SKIP_TRAIN" -eq 1 ]; then
  echo ">> --skip-train set, assuming model partition is already flashed. Skipping."
else
  uvx --from esptool esptool.py --chip esp32s3 --port "$ESP32_PORT" --baud 921600 \
    write-flash 0x110000 firmware/model/model.bin
fi

# ---- 7. Run ----
echo ">> Opening serial monitor. Press the board's RESET button now to see generation from the start."
arduino-cli monitor -p "$ESP32_PORT" --config baudrate=115200