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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
//! `VisionEmbed` — the device-agnostic SigLIP vision entry point: GPU transformer layers when a
//! wgpu adapter exists, the native CPU encoder otherwise. Either way, patch embedding and the MAP
//! attention-pooling head run on the CPU encoder (shared, oracle-identical code), so ONLY the 12
//! transformer layers move to the GPU — which is exactly the stage the parity gate covers.
//!
//! Mirrors `crate::embed_engine::EmbedEngine`'s auto/cpu split for the text/embedding towers.
//! Device policy (force-CPU, force-f32) is the CALLER's decision; this module reads no
//! environment variables of its own.
use crate::GpuCtx;
use crate::encoder::EncoderGpu;
use crate::encoder_cpu::CpuEncoder;
use anyhow::Result;
use std::path::Path;
/// A loaded SigLIP vision tower on whichever device was available.
// The GPU variant additionally holds a CpuEncoder for the shared patch-embed + MAP head; the
// weights live inline in CpuEncoder either way, so the variant-size asymmetry has no runtime cost
// worth a Box (one engine per model, moved at most once).
#[allow(clippy::large_enum_variant)]
pub enum VisionEmbed {
/// GPU path: the wgpu context, the prebaked transformer plan, and the CPU encoder that owns
/// the patch embedding and MAP head (run either side of the GPU layers).
Gpu {
ctx: GpuCtx,
enc: EncoderGpu,
cpu: CpuEncoder,
image_size: usize,
patch_size: usize,
},
/// Native CPU path — also the parity oracle.
Cpu {
cpu: CpuEncoder,
image_size: usize,
patch_size: usize,
},
}
impl VisionEmbed {
/// Load on the best available device: any wgpu adapter first (GPU transformer layers), else
/// the CPU encoder. A GPU-side LOAD failure falls through to CPU with a warning — a machine
/// with a broken driver must still embed images.
pub fn auto(dir: &Path) -> Result<Self> {
Self::auto_with(dir, false)
}
/// `auto`, but forcing f32 GPU weights so the CPU oracle consumes value-identical weights —
/// the precision the parity gate needs.
pub fn auto_f32(dir: &Path) -> Result<Self> {
Self::auto_with(dir, true)
}
fn auto_with(dir: &Path, force_f32: bool) -> Result<Self> {
let (_, spec) = crate::encoder_weights::siglip_configs_from_dir(dir)?;
let (image_size, patch_size) = (spec.image_size, spec.patch_size);
match GpuCtx::new() {
Ok(ctx) => {
let load = if force_f32 {
EncoderGpu::load_siglip_vision_f32(&ctx, dir)
} else {
EncoderGpu::load_siglip_vision(&ctx, dir)
};
match load.and_then(|enc| Ok((enc, CpuEncoder::load_siglip_vision(dir)?))) {
Ok((enc, cpu)) => Ok(Self::Gpu {
ctx,
enc,
cpu,
image_size,
patch_size,
}),
Err(e) => {
eprintln!(
"inferencelayer: GPU SigLIP vision load failed ({e:#}); \
falling back to CPU"
);
Self::cpu(dir)
}
}
}
Err(_) => Self::cpu(dir),
}
}
/// Load on the CPU unconditionally (deployment policy, contention avoidance, oracle runs).
pub fn cpu(dir: &Path) -> Result<Self> {
let (_, spec) = crate::encoder_weights::siglip_configs_from_dir(dir)?;
Ok(Self::Cpu {
cpu: CpuEncoder::load_siglip_vision(dir)?,
image_size: spec.image_size,
patch_size: spec.patch_size,
})
}
/// Embed ONE image's normalized `[3·image²]` CHW pixels into the pooled, L2-normalized
/// embedding. GPU path: CPU patch-embed → GPU transformer layers → CPU MAP head. CPU path:
/// the whole tower on the CPU encoder.
pub fn encode_image(
&mut self,
pixels: &[f32],
image_size: usize,
patch_size: usize,
) -> Result<Vec<f32>> {
match self {
Self::Gpu { ctx, enc, cpu, .. } => {
let rows = cpu.patch_embed_rows(pixels, image_size, patch_size)?;
let n = rows.len() / enc.config().hidden;
let hidden = enc.run_layers_gpu(ctx, &rows, n)?;
cpu.map_head(&hidden)
}
Self::Cpu { cpu, .. } => cpu.encode_image(pixels, image_size, patch_size),
}
}
/// Embed a BATCH of images' normalized `[3·image²]` CHW pixels in one shot. GPU path: CPU
/// patch-embed each → concatenate → ONE batched transformer replay (amortizing the per-image
/// dispatch overhead) → CPU MAP head each. Batches larger than [`crate::encoder::MAX_VISION_BATCH`]
/// are processed in chunks. Per-image results are identical to [`Self::encode_image`].
pub fn encode_images(
&mut self,
images: &[&[f32]],
image_size: usize,
patch_size: usize,
) -> Result<Vec<Vec<f32>>> {
match self {
Self::Gpu { ctx, enc, cpu, .. } => {
let h = enc.config().hidden;
let n = (image_size / patch_size).pow(2);
let per = n * h;
let mut out = Vec::with_capacity(images.len());
for chunk in images.chunks(crate::encoder::MAX_VISION_BATCH) {
let mut rows = Vec::with_capacity(chunk.len() * per);
for img in chunk {
rows.extend_from_slice(&cpu.patch_embed_rows(img, image_size, patch_size)?);
}
let hidden = enc.run_layers_gpu_batched(ctx, &rows, n, chunk.len())?;
for i in 0..chunk.len() {
out.push(cpu.map_head(&hidden[i * per..(i + 1) * per])?);
}
}
Ok(out)
}
Self::Cpu { cpu, .. } => images
.iter()
.map(|img| cpu.encode_image(img, image_size, patch_size))
.collect(),
}
}
/// Input image side length in pixels (the tower's expected resize target).
pub fn image_size(&self) -> usize {
match self {
Self::Gpu { image_size, .. } | Self::Cpu { image_size, .. } => *image_size,
}
}
/// Square patch side length in pixels.
pub fn patch_size(&self) -> usize {
match self {
Self::Gpu { patch_size, .. } | Self::Cpu { patch_size, .. } => *patch_size,
}
}
/// The embedding dimensionality — the tower's hidden size (768 for SigLIP base).
pub fn hidden(&self) -> usize {
match self {
Self::Gpu { cpu, .. } | Self::Cpu { cpu, .. } => cpu.config().hidden,
}
}
/// Human-readable device tag: `"<Backend>/<adapter> (<precision>)"` on the GPU, else `"cpu"`.
pub fn device(&self) -> String {
match self {
Self::Gpu { ctx, enc, .. } => format!("{} ({})", ctx.backend, enc.precision()),
Self::Cpu { .. } => "cpu".to_string(),
}
}
}