use std::fmt;
use std::str::FromStr;
use crate::error::Result;
use crate::types::{
AudioBuffer, DetectOutput, ImageBuffer, OcrOutput, TimedSegment, Transcript, VideoFrame,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Task {
Asr,
Tts,
Ocr,
Vlm,
Detect,
Depth,
}
impl fmt::Display for Task {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
Self::Asr => "asr",
Self::Tts => "tts",
Self::Ocr => "ocr",
Self::Vlm => "vlm",
Self::Detect => "detect",
Self::Depth => "depth",
})
}
}
impl FromStr for Task {
type Err = String;
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
match s.to_ascii_lowercase().as_str() {
"asr" => Ok(Self::Asr),
"tts" => Ok(Self::Tts),
"ocr" => Ok(Self::Ocr),
"vlm" => Ok(Self::Vlm),
"detect" => Ok(Self::Detect),
other => Err(format!(
"unknown task `{other}` (expected asr, tts, ocr, vlm, or detect)"
)),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EngineStatus {
Stub,
Experimental,
Stable,
}
impl fmt::Display for EngineStatus {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
Self::Stub => "stub",
Self::Experimental => "experimental",
Self::Stable => "stable",
})
}
}
#[derive(Debug, Clone)]
pub struct EngineInfo {
pub name: String,
pub task: Task,
pub status: EngineStatus,
pub description: String,
}
#[derive(Debug, Clone)]
pub struct AsrOptions {
pub language: Option<String>,
pub word_timestamps: bool,
pub diarize: bool,
pub persist_speakers: bool,
pub max_speakers: Option<usize>,
pub diarize_threshold: f32,
pub translate: bool,
pub vad: bool,
pub vad_threshold: f32,
pub vad_chunk_secs: f32,
pub stream_offset_secs: f64,
}
impl Default for AsrOptions {
fn default() -> Self {
Self {
language: None,
word_timestamps: false,
diarize: false,
persist_speakers: false,
max_speakers: None,
diarize_threshold: 0.80,
translate: false,
vad: true,
vad_threshold: 0.5,
vad_chunk_secs: 30.0,
stream_offset_secs: 0.0,
}
}
}
#[derive(Debug, Clone)]
pub struct TtsOptions {
pub voice: Option<String>,
pub speed: f32,
pub noise_scale: Option<f32>,
pub noise_w: Option<f32>,
pub seed: u64,
pub sentence_silence_s: f32,
}
impl Default for TtsOptions {
fn default() -> Self {
Self {
voice: None,
speed: 1.0,
noise_scale: None,
noise_w: None,
seed: 0,
sentence_silence_s: 0.2,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct OcrOptions {
pub languages: Vec<String>,
pub single_line: bool,
}
#[derive(Debug, Clone)]
pub struct DetectOptions {
pub confidence: f32,
pub max_detections: usize,
pub iou: Option<f32>,
pub classes: Vec<u32>,
}
impl Default for DetectOptions {
fn default() -> Self {
Self {
confidence: 0.25,
max_detections: 300,
iou: None,
classes: Vec::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Default)]
pub enum Decoding {
#[default]
Greedy,
Sampled {
temperature: f32,
top_p: Option<f32>,
top_k: Option<usize>,
seed: u64,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VlmPart<'a> {
Text(&'a str),
Image(&'a ImageBuffer),
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct VlmPrompt<'a> {
pub parts: Vec<VlmPart<'a>>,
}
impl<'a> VlmPrompt<'a> {
#[must_use]
pub fn single(image: &'a ImageBuffer, text: Option<&'a str>) -> Self {
let mut parts = vec![VlmPart::Image(image)];
if let Some(t) = text {
parts.push(VlmPart::Text(t));
}
Self { parts }
}
#[must_use]
pub fn image_count(&self) -> usize {
self.parts
.iter()
.filter(|p| matches!(p, VlmPart::Image(_)))
.count()
}
}
#[derive(Debug, Clone, Default)]
pub struct VlmOptions {
pub prompt: Option<String>,
pub system_prompt: Option<String>,
pub decoding: Decoding,
pub max_new_tokens: Option<usize>,
pub stop: Vec<String>,
pub frames_per_window: Option<usize>,
pub repetition_penalty: Option<f32>,
}
pub trait AsrEngine: Send + Sync {
fn info(&self) -> EngineInfo;
fn transcribe(&self, audio: &AudioBuffer, opts: &AsrOptions) -> Result<Transcript>;
}
pub trait TtsEngine: Send + Sync {
fn info(&self) -> EngineInfo;
fn synthesize(&self, text: &str, opts: &TtsOptions) -> Result<AudioBuffer>;
}
pub trait OcrEngine: Send + Sync {
fn info(&self) -> EngineInfo;
fn recognize(&self, image: &ImageBuffer, opts: &OcrOptions) -> Result<OcrOutput>;
}
#[derive(Debug, Clone)]
pub struct DepthOutput {
pub depth: Vec<f32>,
pub width: usize,
pub height: usize,
pub letterbox: Option<crate::types::Letterbox>,
}
impl DepthOutput {
#[must_use]
pub fn at(&self, x: usize, y: usize) -> Option<f32> {
if x >= self.width || y >= self.height {
return None;
}
self.depth.get(y * self.width + x).copied()
}
#[must_use]
pub fn range(&self) -> Option<(f32, f32)> {
if self.depth.is_empty() {
return None;
}
let mut lo = f32::MAX;
let mut hi = f32::MIN;
for &v in &self.depth {
if v.is_finite() {
lo = lo.min(v);
hi = hi.max(v);
}
}
(lo <= hi).then_some((lo, hi))
}
}
#[derive(Debug, Clone, Default)]
pub struct DepthOptions {
pub full_resolution: bool,
}
pub trait DepthEngine: Send + Sync {
fn info(&self) -> EngineInfo;
fn depth(&self, image: &ImageBuffer, opts: &DepthOptions) -> Result<DepthOutput>;
}
pub trait DetectEngine: Send + Sync {
fn info(&self) -> EngineInfo;
fn detect(&self, image: &ImageBuffer, opts: &DetectOptions) -> Result<DetectOutput>;
fn detect_batch(
&self,
images: &[ImageBuffer],
opts: &DetectOptions,
) -> Result<Vec<DetectOutput>> {
images.iter().map(|i| self.detect(i, opts)).collect()
}
fn class_names(&self) -> &[String];
}
pub trait VlmEngine: Send + Sync {
fn info(&self) -> EngineInfo;
fn describe(&self, prompt: &VlmPrompt<'_>, opts: &VlmOptions) -> Result<String>;
fn describe_image(&self, image: &ImageBuffer, opts: &VlmOptions) -> Result<String> {
self.describe(&VlmPrompt::single(image, opts.prompt.as_deref()), opts)
}
fn describe_video(
&self,
frames: &[VideoFrame],
opts: &VlmOptions,
) -> Result<Vec<TimedSegment<String>>>;
}
#[cfg(test)]
mod vlm_surface_tests {
use super::*;
use crate::types::PixelFormat;
fn img(byte: u8) -> ImageBuffer {
ImageBuffer {
width: 1,
height: 1,
format: PixelFormat::Rgb8,
data: vec![byte, byte, byte],
}
}
struct Recorder(std::sync::Mutex<Vec<String>>);
impl VlmEngine for Recorder {
fn info(&self) -> EngineInfo {
EngineInfo {
name: "recorder".into(),
task: Task::Vlm,
status: EngineStatus::Stub,
description: String::new(),
}
}
fn describe(&self, prompt: &VlmPrompt<'_>, _opts: &VlmOptions) -> Result<String> {
let shape: Vec<String> = prompt
.parts
.iter()
.map(|p| match p {
VlmPart::Text(t) => format!("text:{t}"),
VlmPart::Image(i) => format!("image:{}", i.data[0]),
})
.collect();
self.0.lock().unwrap().push(shape.join("|"));
Ok(shape.join("|"))
}
fn describe_video(
&self,
_frames: &[VideoFrame],
_opts: &VlmOptions,
) -> Result<Vec<TimedSegment<String>>> {
Ok(Vec::new())
}
}
#[test]
fn the_default_decoding_is_deterministic() {
assert_eq!(VlmOptions::default().decoding, Decoding::Greedy);
assert!(VlmOptions::default().stop.is_empty());
assert_eq!(VlmOptions::default().repetition_penalty, None);
}
#[test]
fn sampling_always_carries_a_seed() {
let d = Decoding::Sampled {
temperature: 0.7,
top_p: Some(0.9),
top_k: None,
seed: 42,
};
let Decoding::Sampled { seed, .. } = d else {
panic!("expected Sampled")
};
assert_eq!(seed, 42);
}
#[test]
fn describe_image_routes_through_describe() {
let e = Recorder(std::sync::Mutex::new(Vec::new()));
let opts = VlmOptions {
prompt: Some("what is this?".into()),
..VlmOptions::default()
};
let out = e.describe_image(&img(7), &opts).unwrap();
assert_eq!(out, "image:7|text:what is this?");
assert_eq!(e.0.lock().unwrap().len(), 1, "describe must have been used");
}
#[test]
fn a_prompt_with_no_instruction_is_just_the_image() {
let e = Recorder(std::sync::Mutex::new(Vec::new()));
let out = e.describe_image(&img(3), &VlmOptions::default()).unwrap();
assert_eq!(out, "image:3");
}
#[test]
fn interleaved_order_is_preserved() {
let (a, b) = (img(1), img(2));
let e = Recorder(std::sync::Mutex::new(Vec::new()));
let prompt = VlmPrompt {
parts: vec![
VlmPart::Text("before"),
VlmPart::Image(&a),
VlmPart::Text("between"),
VlmPart::Image(&b),
VlmPart::Text("after"),
],
};
assert_eq!(prompt.image_count(), 2);
let out = e.describe(&prompt, &VlmOptions::default()).unwrap();
assert_eq!(
out, "text:before|image:1|text:between|image:2|text:after",
"the sequence an engine receives must be the sequence it was given"
);
}
}