1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
# Reusable VLM screenshot analysis derivation.
#
# Similar to `wait_for_text()` in NixOS VM tests.
#
# Runs a VLM (Vision Language Model) on a screenshot and asserts that the
# model's answer to a yes/no question ends with "YES".
#
# This is useful to automatically test software that is otherwise
# hard to test, e.g. "does this 3D program render the bunny correctly?".
# It is especially useful to judge screenshots made in NixOS VM tests.
{
lib,
writers,
fetchurl,
llama-cpp,
runCommand,
# VLM defaults, chosen to pick a model smart enough to be useful
# for screenshot analysis, but small enough to not consume too much RAM
# or be too slow for CI.
model ? (
fetchurl {
url = "https://huggingface.co/unsloth/gemma-4-E2B-it-GGUF/resolve/90f9618340396838ee7ff5b0ba2da27da62953d3/gemma-4-E2B-it-Q4_0.gguf";
hash = "sha256-nEwdSKRi9/iDsomErE9C02bJxXNDTqtoVT4POL9+tQw=";
}
),
mmproj ? (
fetchurl {
url = "https://huggingface.co/unsloth/gemma-4-E2B-it-GGUF/resolve/90f9618340396838ee7ff5b0ba2da27da62953d3/mmproj-F16.gguf";
hash = "sha256-FAvo14SXQfiMUHV9UpuENz7o4nBSzCI2hVtTf0qCFfo=";
}
),
# User-provided arguments:
name,
screenshot,
question,
}:
let
analysisScript =
writers.writePython3 "${name}-script" { flakeIgnore = [ "E501" ]; } # allow long lines
''
import os
import re
import subprocess
import time
out = os.environ["out"]
screenshot = "${screenshot}"
# Using JSON here even permits preserving multi-line ASCII art questions and so on.
question = ${builtins.toJSON question}
# Build the full prompt with output markers for reliable extraction.
prompt = (
"Start your output with [output-start]."
f" {question}"
" Explain what you see, and your judgment."
" On the final line, output exactly YES or NO."
" Then immediately output [output-end] with nothing else after it."
)
vlm_start = time.time()
result = subprocess.run(
[
"${lib.getExe' llama-cpp "llama-cli"}",
"--single-turn", "--no-display-prompt", "--log-verbosity", "0", "--jinja",
"--simple-io", # disables the spinner whose backspace chars would corrupt captured output
"--reasoning", "off", "--temp", "0",
"--threads", "1", # for determinism
"--n-gpu-layers", "0", # force CPU-only (results on GPUs might be different and nondeterministic, see https://github.com/ggml-org/llama.cpp/pull/16016#issuecomment-3293505238)
"--model", "${model}",
"--mmproj", "${mmproj}",
"--image", screenshot,
"-p", prompt,
],
capture_output=True,
text=True,
# `OMP_NUM_THREADS=1` prevents OpenMP from spawning extra threads in the BLAS backend
# (OpenBLAS), which causes nondeterminism with `--image`; without `--image`, `--threads 1`
# alone is already deterministic (BLAS is not used for short text prompts).
# Relevant code: https://github.com/ggml-org/llama.cpp/blob/80afa33aadcc4f71212b17e5e52904491c76b63e/ggml/src/ggml-blas/ggml-blas.cpp#L30-L148
# PR to fix it in OpenBLAS: https://github.com/OpenMathLib/OpenBLAS/pull/5808
env={**os.environ, "OMP_NUM_THREADS": "1"},
)
vlm_elapsed = time.time() - vlm_start
output = result.stdout
print(f"VLM inference took {vlm_elapsed:.1f}s")
print(f"VLM raw output: {repr(output)}")
if result.returncode != 0:
print(f"VLM stderr: {result.stderr}")
assert result.returncode == 0, f"llama-cli failed with exit code {result.returncode}"
print()
# Post-process: extract the answer between `[output-start]` and `[output-end]` markers.
# This is needed because llama-cli prints UI noise (banner,
# spinner, stats) to stdout alongside the model's response.
# TODO: Replace with `--quiet` once https://github.com/ggml-org/llama.cpp/pull/22848 is merged;
# then also remove the markers from the prompt and the extraction below.
matches = re.findall(r"\[output-start\](.*?)\[output-end\]", output, re.DOTALL)
assert matches, (
f"VLM output did not contain [output-start]...[output-end] markers."
f" Raw output: {repr(output)}"
)
answer = matches[-1].strip() # use last match (first may be prompt echo)
print("VLM answer:")
print(answer)
assert answer.upper().endswith("YES"), (
f"VLM did not confirm expected answer. Answer: {answer}"
)
os.makedirs(out, exist_ok=True)
with open(os.path.join(out, "vlm-answer.txt"), "w") as f:
f.write(answer + "\n")
os.symlink(screenshot, os.path.join(out, "screen.png"))
'';
in
runCommand name { } ''
${analysisScript}
''
|
