#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
use super::GeneratedContent;
use super::{MediaBuildOptions, MediaError, MediaKind, MediaProducer, Producer};
#[cfg(feature = "audio-transcribe")]
const ASR_PROMPT: &str = "Transcribe this audio recording. Output only the spoken words, verbatim.";
#[cfg(feature = "image-vision")]
const VLM_PROMPT: &str = "Describe this image in one or two sentences.";
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
const TEMPERATURE: f64 = 0.0;
#[cfg(feature = "audio-transcribe")]
const ASR_MAX_TOKENS: u32 = 512;
#[cfg(feature = "image-vision")]
const VLM_MAX_TOKENS: u32 = 128;
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
fn resolved(kind: MediaKind) -> Result<crate::model_choice::ModelChoice, MediaError> {
Ok(crate::model_choice::resolve(kind.task())?)
}
#[cfg(all(test, feature = "audio-transcribe"))]
pub(crate) const ASR_MODEL: &str = MediaKind::Audio.model();
#[cfg(all(test, feature = "image-vision"))]
pub(crate) const VLM_MODEL: &str = MediaKind::Vision.model();
#[must_use]
pub fn available() -> Vec<Producer> {
#[cfg_attr(
not(any(feature = "audio-transcribe", feature = "image-vision")),
allow(unused_mut)
)]
let mut out = Vec::new();
#[cfg(feature = "audio-transcribe")]
if let Some(p) = asr_producer() {
out.push(p);
}
#[cfg(feature = "image-vision")]
if let Some(p) = vlm_producer() {
out.push(p);
}
out
}
pub fn installed(opts: MediaBuildOptions) -> Result<Vec<Box<dyn MediaProducer>>, MediaError> {
let mut out: Vec<Box<dyn MediaProducer>> = Vec::new();
if opts.audio {
out.push(audio_producer()?);
}
if opts.vision {
out.push(vision_producer()?);
}
Ok(out)
}
#[cfg(feature = "audio-transcribe")]
fn audio_producer() -> Result<Box<dyn MediaProducer>, MediaError> {
let model = resolved(MediaKind::Audio)?.require_installed()?;
let producer = registry_producer(MediaKind::Audio, model, ASR_PROMPT, ASR_MAX_TOKENS)
.ok_or_else(|| MediaError::ModelMissing {
model: model.to_owned(),
})?;
Ok(Box::new(LlamaProducer {
producer,
engine: asr_engine,
}))
}
#[cfg(not(feature = "audio-transcribe"))]
fn audio_producer() -> Result<Box<dyn MediaProducer>, MediaError> {
Err(MediaError::NoProducer {
kind: MediaKind::Audio.as_str(),
feature: "audio-transcribe",
})
}
#[cfg(feature = "image-vision")]
fn vision_producer() -> Result<Box<dyn MediaProducer>, MediaError> {
let model = resolved(MediaKind::Vision)?.require_installed()?;
let producer = registry_producer(MediaKind::Vision, model, VLM_PROMPT, VLM_MAX_TOKENS)
.ok_or_else(|| MediaError::ModelMissing {
model: model.to_owned(),
})?;
Ok(Box::new(LlamaProducer {
producer,
engine: vlm_engine,
}))
}
#[cfg(not(feature = "image-vision"))]
fn vision_producer() -> Result<Box<dyn MediaProducer>, MediaError> {
Err(MediaError::NoProducer {
kind: MediaKind::Vision.as_str(),
feature: "image-vision",
})
}
#[cfg(feature = "audio-transcribe")]
fn asr_producer() -> Option<Producer> {
let model = resolved(MediaKind::Audio).ok()?.model?;
registry_producer(MediaKind::Audio, model, ASR_PROMPT, ASR_MAX_TOKENS)
}
#[cfg(feature = "image-vision")]
fn vlm_producer() -> Option<Producer> {
let model = resolved(MediaKind::Vision).ok()?.model?;
registry_producer(MediaKind::Vision, model, VLM_PROMPT, VLM_MAX_TOKENS)
}
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
fn registry_producer(
kind: MediaKind,
model: &str,
prompt: &str,
max_tokens: u32,
) -> Option<Producer> {
let spec = crate::models::find(model)?;
let variant = spec.variant_for(crate::models::Platform::host())?;
let dir = crate::models::model_dir(model);
if !variant.files.iter().all(|f| dir.join(f.name).exists()) {
return None;
}
let file = |name: &str| variant.files.iter().find(|f| f.name == name);
let weights = file("model.gguf")?;
let mmproj = file("mmproj.gguf")?;
Some(Producer {
kind,
model: model.to_owned(),
model_digest: weights.sha256.to_owned(),
quantisation: quantisation_of(weights.url),
mmproj_digest: mmproj.sha256.to_owned(),
prompt: prompt.to_owned(),
temperature: TEMPERATURE,
max_tokens,
})
}
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
fn quantisation_of(url: &str) -> String {
let name = url.rsplit('/').next().unwrap_or(url);
let stem = name.strip_suffix(".gguf").unwrap_or(name);
for segment in stem.rsplit('-') {
let upper = segment.to_ascii_uppercase();
let quantised = (upper.starts_with('Q') || upper.starts_with("IQ"))
&& upper.chars().any(|c| c.is_ascii_digit())
|| matches!(upper.as_str(), "F16" | "F32" | "BF16");
if quantised {
return upper;
}
}
"unknown".to_owned()
}
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
struct LlamaProducer {
producer: Producer,
engine: fn() -> Option<std::sync::Arc<rto_llama::llama::LlamaEngine>>,
}
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
impl MediaProducer for LlamaProducer {
fn producer(&self) -> &Producer {
&self.producer
}
fn generate(&self, _path: &str, bytes: &[u8]) -> Option<GeneratedContent> {
use rto_llama::Engine as _;
#[cfg(feature = "image-vision")]
if self.producer.kind == MediaKind::Vision && !crate::extract::image_dimensions_ok(bytes) {
return None;
}
let (images, audio) = match self.producer.kind {
MediaKind::Vision => (vec![bytes.to_vec()], Vec::new()),
MediaKind::Audio => (Vec::new(), vec![bytes.to_vec()]),
};
let engine = (self.engine)()?;
let completion = engine
.chat(&rto_llama::ChatRequest {
model: self.producer.model.clone(),
messages: vec![rto_llama::Message {
role: "user".to_owned(),
content: self.producer.prompt.clone(),
}],
images,
audio,
#[allow(
clippy::cast_possible_truncation,
reason = "engine API is f32; the stored identity keeps the f64"
)]
temperature: self.producer.temperature as f32,
max_tokens: self.producer.max_tokens,
})
.ok()?;
let text = completion.content.trim();
(!text.is_empty()).then(|| GeneratedContent {
text: text.to_owned(),
confidence: None,
})
}
}
#[cfg(feature = "audio-transcribe")]
static ASR_ENGINE: rto_llama::EngineSlot<rto_llama::llama::LlamaEngine> =
rto_llama::EngineSlot::new();
#[cfg(feature = "image-vision")]
static VLM_ENGINE: rto_llama::EngineSlot<rto_llama::llama::LlamaEngine> =
rto_llama::EngineSlot::new();
#[cfg(feature = "audio-transcribe")]
pub(crate) fn asr_engine() -> Option<std::sync::Arc<rto_llama::llama::LlamaEngine>> {
ASR_ENGINE.get_or_init(|| build_engine(resolved(MediaKind::Audio).ok()?.model?))
}
#[cfg(all(test, feature = "audio-transcribe"))]
pub(crate) fn asr_content(bytes: &[u8]) -> Option<String> {
let producer = asr_producer()?;
LlamaProducer {
producer,
engine: asr_engine,
}
.generate("fixture.wav", bytes)
.map(|c| c.text)
}
#[cfg(feature = "image-vision")]
pub(crate) fn vlm_engine() -> Option<std::sync::Arc<rto_llama::llama::LlamaEngine>> {
VLM_ENGINE.get_or_init(|| build_engine(resolved(MediaKind::Vision).ok()?.model?))
}
#[cfg(all(test, feature = "image-vision"))]
pub(crate) fn vlm_content(bytes: &[u8]) -> Option<String> {
let producer = vlm_producer()?;
LlamaProducer {
producer,
engine: vlm_engine,
}
.generate("fixture.png", bytes)
.map(|c| c.text)
}
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
fn build_engine(model: &str) -> Option<rto_llama::llama::LlamaEngine> {
let dir = crate::models::model_dir(model);
let (gguf, mmproj) = (dir.join("model.gguf"), dir.join("mmproj.gguf"));
if !gguf.exists() || !mmproj.exists() {
return None;
}
rto_llama::llama::LlamaEngine::new(
vec![rto_llama::llama::Served {
name: model.to_owned(),
path: gguf,
mmproj: Some(mmproj),
}],
0,
)
.ok()
}
#[cfg(feature = "image-vision")]
pub(crate) fn release_vlm_engine() -> bool {
VLM_ENGINE.release()
}
#[cfg(not(feature = "image-vision"))]
pub(crate) fn release_vlm_engine() -> bool {
false
}
#[cfg(feature = "audio-transcribe")]
pub(crate) fn release_asr_engine() -> bool {
ASR_ENGINE.release()
}
#[cfg(not(feature = "audio-transcribe"))]
pub(crate) fn release_asr_engine() -> bool {
false
}
#[cfg(test)]
mod tests {
#[cfg(not(all(feature = "audio-transcribe", feature = "image-vision")))]
use super::MediaError;
use super::{MediaBuildOptions, available, installed};
#[cfg(not(all(feature = "audio-transcribe", feature = "image-vision")))]
fn refusal(opts: MediaBuildOptions) -> MediaError {
match installed(opts) {
Err(e) => e,
Ok(ok) => panic!("expected a refusal, got {} producer(s)", ok.len()),
}
}
#[test]
#[cfg(not(feature = "audio-transcribe"))]
fn an_audio_build_without_the_feature_is_refused_by_name() {
let err = refusal(MediaBuildOptions {
audio: true,
vision: false,
..MediaBuildOptions::default()
});
assert_eq!(
err,
MediaError::NoProducer {
kind: "audio",
feature: "audio-transcribe",
}
);
assert!(
err.to_string().contains("--features audio-transcribe"),
"the message must name the rebuild: {err}"
);
}
#[test]
#[cfg(not(feature = "image-vision"))]
fn a_vision_build_without_the_feature_is_refused_by_name() {
let err = refusal(MediaBuildOptions {
audio: false,
vision: true,
..MediaBuildOptions::default()
});
assert!(
err.to_string().contains("--features image-vision"),
"the message must name the rebuild: {err}"
);
}
#[test]
fn requesting_no_modality_needs_no_producer() {
let none = installed(MediaBuildOptions {
audio: false,
vision: false,
..MediaBuildOptions::default()
})
.expect("asking for nothing cannot fail");
assert!(none.is_empty());
}
#[test]
#[cfg(not(any(feature = "audio-transcribe", feature = "image-vision")))]
fn a_default_build_offers_no_producer() {
assert!(available().is_empty());
}
#[test]
#[cfg(any(feature = "audio-transcribe", feature = "image-vision"))]
fn quantisation_is_read_from_the_pinned_file_name() {
use super::quantisation_of;
assert_eq!(
quantisation_of(
"https://huggingface.co/x/resolve/main/Voxtral-Mini-3B-2507-Q4_K_M.gguf"
),
"Q4_K_M"
);
assert_eq!(
quantisation_of(
"https://huggingface.co/x/resolve/main/mmproj-SmolVLM-500M-Instruct-Q8_0.gguf"
),
"Q8_0"
);
assert_eq!(quantisation_of("https://x/model-F16.gguf"), "F16");
assert_eq!(quantisation_of("https://x/SmolVLM-500M-3B.gguf"), "unknown");
}
#[test]
fn available_producers_match_the_compiled_modalities() {
use super::MediaKind;
for producer in available() {
let compiled_in = match producer.kind {
MediaKind::Audio => cfg!(feature = "audio-transcribe"),
MediaKind::Vision => cfg!(feature = "image-vision"),
};
assert!(
compiled_in,
"a {} producer was offered by a build without its feature",
producer.kind
);
producer.validate().expect("an offered producer is valid");
}
}
}