AI on the phone

Overview

This notebook runs machine-learning models on the phone's CPU, with no GPU or cloud service: object detection on a camera photo, speech to text, and a chat with a small language model. First it looks at how slowly plain Elixir tensors run on it. It builds on the earlier notebooks: Wi-Fi (01), the camera (05) and audio (10).

The stack

The AI stack has three kinds of parts:

• Capability packages do the work, each a native library wrapped for the BEAM: llama_cpp_ex runs language models (llama.cpp), whisper_cpp transcribes speech (whisper.cpp, with Silero voice detection) and evision runs the YOLO object detector (OpenCV).
• nerves_ai is the composer. It puts one small facade in front of each capability (NervesAI.LLM, NervesAI.STT, NervesAI.Vision), downloads model files with checksums (NervesAI.Models), reads and writes WAV files (NervesAI.Audio.WAV).
• nerves_data_resize grows the data partition on first boot so models fit.

Nothing is configured to start at boot. This notebook downloads each model when you need it and starts the matching capability with Kino.start_child/1.

The Snapdragon 632 has two clusters of four Kryo 250 cores: cpu0-3 are the slower "Silver" cores and cpu4-7 the faster "Gold" cores. Neither cluster has the dot-product or half-precision instructions that newer ARM chips use to speed up model maths, so the native libraries are built for plain ARMv8.0. The engines use four threads each.

Two things were measured and dropped because they made no measurable difference on this phone: a "run fast" lease holding the Gold cluster at full speed around a burst (schedutil already ramps the clock within milliseconds), and pinning the VM's threads to the Gold cluster (the kernel already schedules this CPU-bound work onto the Gold cores on its own — alternating pinned and unpinned runs of the pipelines below showed no timing difference).

Tensors on the BEAM

Nx is Elixir's tensor library. This firmware sets NxEigen.Backend as its default backend (config :nx, default_backend: NxEigen.Backend): the Eigen C++ library with NEON, built for this CPU, rather than Nx's plain Elixir backend, which is far too slow on a phone. The object detector below does its image work in OpenCV; Nx is there for your own tensor code and for libraries such as nx_image and nx_signal.

Models: downloaded when you need them

The next steps use trained models. They're too big for the firmware and none is downloaded automatically: each step fetches its model the first time with NervesAI.Models.ensure_one/3. It checks the SHA-256, writes the file atomically to /data/models and skips the download when the file is already there.

The phone needs internet access for the first run. If this doesn't say :internet, connect to Wi-Fi with the Wi-Fi notebook:

VintageNet.get(["connection"])

These are two models known to run on this phone, with their checksums:

catalog = %{
  tinyllama: [
    source:
      {:hf, "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"},
    path: "/data/models/tinyllama.gguf",
    sha256: "9fecc3b3cd76bba89d504f29b616eedf7da85b96540e490ca5824d3f7d2776a0"
  ],
  yolov5n: [
    source: {:url, "https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.onnx"},
    path: "/data/models/yolov5n.onnx",
    sha256: "04f0e55c26f58d17145b36045780fe1250d5bd2187543e11568e5141d05b3262"
  ]
}

Map.keys(catalog)

Two helpers for the steps below: fetch downloads a model (printing a heartbeat so a multi-minute download on slow Wi-Fi doesn't look hung) and returns its path, and start starts a capability under this notebook's supervisor and waits for its model to load.

fetch = fn id, spec ->
  heartbeat =
    Task.async(fn ->
      Stream.interval(5_000)
      |> Stream.each(fn n -> IO.puts("#{id}: still downloading... (#{(n + 1) * 5} s)") end)
      |> Stream.run()
    end)

  result = NervesAI.Models.ensure_one(id, spec)
  Task.shutdown(heartbeat, :brutal_kill)

  case result do
    {:ok, path} ->
      path

    {:error, reason} ->
      raise "Couldn't download #{id}: #{inspect(reason)}. " <>
              "Is the phone online? Check with the Wi-Fi notebook."
  end
end

start = fn adapter, config ->
  case Kino.start_child(adapter.child_spec(config)) do
    {:ok, _pid} -> :ok
    {:error, {:already_started, _pid}} -> :ok
    {:error, reason} -> raise "Couldn't start #{inspect(adapter)}: #{inspect(reason)}"
  end

  ready =
    Enum.any?(1..240, fn _ ->
      adapter.ready?(config) or (Process.sleep(500) && false)
    end)

  if ready, do: :ready, else: raise("#{inspect(adapter)} didn't load its model; see the log")
end

Find objects in a photo (YOLO)

YOLOv5n finds and labels objects in a picture: people, chairs, bottles, laptops and 76 other everyday categories (the COCO dataset; NervesAI.Vision.COCO.labels/0 lists them). The model is 4 MB. NervesAI.Vision.Evision runs it with OpenCV and NervesAI.Vision is the interface to call.

yolo_path = fetch.(:yolov5n, catalog.yolov5n)

start.(NervesAI.Vision.Evision, path: yolo_path)

Take a photo and detect

Point the back of the phone at something and run the cell. It takes a 2×2-binned photo (quicker than full resolution) and lists what the model found. Each box is [x1, y1, x2, y2] in pixels of the photo.

photo = "/tmp/detect.jpg"
IO.puts("Taking the photo...")
{:ok, _} = Fp3Camera.snap(:rear, photo, binned: true)
IO.puts("Looking for objects...")

detections = NervesAI.Vision.detect(NervesAI.Vision, photo)

for d <- detections do
  %{object: d.label, confidence: Float.round(d.score, 2), box: inspect(Enum.map(d.box, &round/1))}
end
|> Kino.DataTable.new()

detect/3 takes conf: (minimum score, default 0.25) and iou: options. This cell draws the boxes as an SVG overlay on the photo.

See it

{height, width, _channels} = Evision.imread(photo).shape

boxes =
  for %{box: [x1, y1, x2, y2], label: label, score: score} <- detections do
    """
    <rect x="#{x1}" y="#{y1}" width="#{x2 - x1}" height="#{y2 - y1}"
          fill="none" stroke="yellow" stroke-width="#{width / 250}"/>
    <text x="#{x1 + 6}" y="#{y1 + width / 30}" fill="yellow" font-size="#{width / 32}"
          font-family="sans-serif">#{label} #{round(score * 100)}%</text>
    """
  end

jpeg = Base.encode64(File.read!(photo))

Kino.HTML.new("""
<svg viewBox="0 0 #{width} #{height}" style="max-width: 100%">
  <image href="data:image/jpeg;base64,#{jpeg}" width="#{width}" height="#{height}"/>
  #{Enum.join(boxes)}
</svg>
""")

Live detection

The same model on live video. The camera ISP hands over 640×480 frames at 30 fps, and each one is shown with the latest boxes drawn on it, so the video stays real time. Detection takes about 300 ms, so it runs on its own beside the video, always on the newest frame, and the boxes jump a few times a second to keep up. OpenCV gets 4 of the 8 cores for it and the video keeps the others.

Press Stop to end it.

camera = :rear
video = Kino.Frame.new() |> Kino.render()
stop = Kino.Control.button("Stop") |> Kino.render()
Kino.Control.subscribe(stop, :stop)

Evision.setNumThreads(4)

# Asks for a frame whenever it's free, and sends back what it found in it
parent = self()

detector =
  spawn_link(fn ->
    detect = fn detect ->
      send(parent, {:free, self()})

      receive do
        {:frame, bgr} ->
          send(parent, {:detections, NervesAI.Vision.detect(NervesAI.Vision, bgr)})
          detect.(detect)

        :stop ->
          :ok
      end
    end

    detect.(detect)
  end)

draw = fn bgr, detections ->
  Enum.reduce(detections, bgr, fn %{box: [x1, y1, x2, y2], label: label, score: score}, img ->
    img
    |> Evision.rectangle({round(x1), round(y1)}, {round(x2), round(y2)}, {0, 255, 255},
      thickness: 2
    )
    |> Evision.putText("#{label} #{round(score * 100)}%", {round(x1) + 4, round(y1) + 16},
      Evision.Constant.cv_FONT_HERSHEY_SIMPLEX(), 0.5, {0, 255, 255}
    )
  end)
end

{:ok, feed} = Fp3Camera.subscribe(camera, width: 640, height: 480)

show = fn show, detections, free ->
  receive do
    {:stop, _event} ->
      :stopped

    {:detections, detections} ->
      show.(show, detections, free)

    {:free, detector} ->
      show.(show, detections, detector)

    {:camera_frame, %{data: nv12, width: w, height: h}} ->
      bgr =
        nv12
        |> Evision.Mat.from_binary({:u, 8}, div(h * 3, 2), w, 1)
        |> Evision.cvtColor(Evision.Constant.cv_COLOR_YUV2BGR_NV12())

      if free, do: send(free, {:frame, bgr})

      jpeg =
        Evision.imencode(".jpg", draw.(bgr, detections),
          params: [Evision.Constant.cv_IMWRITE_JPEG_QUALITY(), 70]
        )

      Kino.Frame.render(video, Kino.Image.new(jpeg, :jpeg))
      show.(show, detections, nil)
  end
end

show.(show, [], nil)
Fp3Camera.unsubscribe(feed)
send(detector, :stop)

Speech to text (Whisper)

Whisper turns speech into text. This is OpenAI's tiny.en model (English only, about 75 MB) in the format whisper.cpp reads, run by NervesAI.STT.WhisperCpp. A second, tiny model (about 0.9 MB), Silero, detects where speech is so silence is skipped before Whisper sees it.

whisper_path =
  fetch.(:whisper,
    source: {:hf, "ggerganov/whisper.cpp", "ggml-tiny.en.bin"},
    path: "/data/models/ggml-tiny.en.bin",
    sha256: "921e4cf8686fdd993dcd081a5da5b6c365bfde1162e72b08d75ac75289920b1f"
  )

vad_path =
  fetch.(:silero,
    source: {:hf, "ggml-org/whisper-vad", "ggml-silero-v5.1.2.bin"},
    path: "/data/models/ggml-silero-v5.1.2.bin",
    sha256: "29940d98d42b91fbd05ce489f3ecf7c72f0a42f027e4875919a28fb4c04ea2cf"
  )

start.(NervesAI.STT.WhisperCpp,
  path: whisper_path,
  vad_path: vad_path,
  language: "en"
)

Transcribe a recording

The firmware ships an 11-second excerpt of John F. Kennedy's 1961 inaugural address. Nothing here converts sample rates, so the audio has to be 16 kHz; this file is.

sample = Application.app_dir(:nerves_livebook_fp3, "priv/audio/jfk.wav")
IO.puts("Transcribing...")
{:ok, %{text: text, segments: segments}} = NervesAI.STT.transcribe(NervesAI.STT, {:wav, sample})

IO.puts(text)
Kino.DataTable.new(segments)

Transcribe yourself

Records 5 seconds from the phone's main microphone. Speak (in English) while it runs. The microphone needs a system image whose kernel drives the WCD9335 audio codec; without it, arecord fails here.

The recording is asked for at 16 kHz mono, and ALSA's plughw device converts from the microphone's rate. This has not been tested on the device. If arecord rejects -r 16000, record at 48 kHz as the audio notebook (10) does; NervesAI.STT.transcribe/3 then answers {:error, {:unsupported_rate, 48000}} and the recording needs a resampling step first.

microphone = [
  {"MultiMedia2 Mixer SLIMBUS_0_TX", :on},
  {"SLIM TX7 MUX", "DEC7"},
  {"ADC MUX7", "DMIC"},
  {"DMIC MUX7", "DMIC1"},
  {"AIF1_CAP Mixer SLIM TX7", :on}
]

for {name, value} <- microphone, do: {name, ExAudio.cset(name, value)}

arecord = fn seconds, path ->
  System.cmd("arecord", ~w(-D plughw:0,1 -f S16_LE -r 16000 -c 1 -d #{seconds} #{path}),
    stderr_to_stdout: true
  )
end

# The first capture after boot returns silence; a short one wakes the path up.
arecord.(1, "/tmp/warmup.wav")
IO.puts("Recording for 5 seconds: speak now!")
{output, status} = arecord.(5, "/tmp/me.wav")
if status != 0, do: IO.puts(output)

IO.puts("Transcribing...")
NervesAI.STT.transcribe(NervesAI.STT, {:wav, "/tmp/me.wav"})

Chat with a language model (TinyLlama)

TinyLlama 1.1B Chat is a small language model: about 670 MB, the biggest step. NervesAI.LLM.LlamaCpp runs it with llama.cpp. The model file carries its own chat template, so you pass the conversation as a list of messages and the library formats the prompt the way TinyLlama was trained.

llm_path = fetch.(:tinyllama, catalog.tinyllama)

start.(NervesAI.LLM.LlamaCpp, path: llm_path, n_ctx: 1024)

Ask a question

Change question and run the cell again. The answer is printed piece by piece as the model writes it. max_tokens caps its length.

question = "What three things should I see in Brussels?"

messages = [
  %{role: "system", content: "You are a helpful, concise travel guide."},
  %{role: "user", content: question}
]

started = System.monotonic_time(:millisecond)

pieces =
  NervesAI.LLM.stream_chat(NervesAI.LLM, messages, max_tokens: 128, temp: 0.7)
  |> Enum.reduce(0, fn
    text, count when is_binary(text) ->
      IO.write(text)
      count + 1

    {:error, reason}, count ->
      IO.puts("\n[generation failed: #{inspect(reason)}]")
      count
  end)

seconds = (System.monotonic_time(:millisecond) - started) / 1000
IO.puts("\n\n#{pieces} pieces in #{Float.round(seconds, 1)} s")

NervesAI.LLM.chat/3 takes the same arguments and returns the whole answer at once.

Already on this phone

The models downloaded so far (sizes in MB):

case File.ls("/data/models") do
  {:ok, files} ->
    for f <- Enum.sort(files), do: {f, div(File.stat!(Path.join("/data/models", f)).size, 1_000_000)}

  {:error, :enoent} ->
    []
end