feat: llama.cpp-based paddleocr-vl-1.5

Cargo.lock

@@ -3947,6 +3947,7 @@ dependencies = [
  "flate2",
  "futures",
  "hf-hub",
+ "image",
  "keyring",
  "koharu-http",
  "koharu-runtime",
@@ -3987,6 +3988,7 @@ dependencies = [
  "imageproc",
  "keyring",
  "koharu-http",
+ "koharu-llm",
  "koharu-types",
  "libloading 0.8.9",
  "minijinja",

koharu-llm/Cargo.toml

@@ -22,6 +22,10 @@ path = "bin/llama.rs"
 name = "llm"
 path = "bin/llm.rs"
 
+[[bin]]
+name = "paddleocr-vl-llm"
+path = "bin/paddleocr-vl.rs"
+
 [dependencies]
 anyhow = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
@@ -29,6 +33,7 @@ dirs = { workspace = true }
 flate2 = { workspace = true }
 futures = { workspace = true }
 hf-hub = { workspace = true }
+image = { workspace = true }
 koharu-http = { workspace = true }
 koharu-runtime = { workspace = true }
 koharu-types = { workspace = true }

koharu-llm/bin/paddleocr-vl.rs

@@ -0,0 +1,308 @@
+use std::path::{Path, PathBuf};
+use std::time::Instant;
+
+use anyhow::{Context, Result, bail};
+use clap::{Parser, ValueEnum};
+use koharu_llm::paddleocr_vl::{PaddleOcrVl, PaddleOcrVlOutput, PaddleOcrVlTask};
+
+#[derive(Debug, Clone, Copy, ValueEnum)]
+enum TaskArg {
+    Ocr,
+    Table,
+    Formula,
+    Chart,
+    Spotting,
+    Seal,
+}
+
+impl From<TaskArg> for PaddleOcrVlTask {
+    fn from(value: TaskArg) -> Self {
+        match value {
+            TaskArg::Ocr => Self::Ocr,
+            TaskArg::Table => Self::Table,
+            TaskArg::Formula => Self::Formula,
+            TaskArg::Chart => Self::Chart,
+            TaskArg::Spotting => Self::Spotting,
+            TaskArg::Seal => Self::Seal,
+        }
+    }
+}
+
+#[derive(Parser)]
+struct Cli {
+    #[arg(long, value_name = "FILE", num_args = 1..)]
+    input: Vec<PathBuf>,
+
+    #[arg(long, value_name = "DIR")]
+    model_dir: Option<PathBuf>,
+
+    #[arg(long, value_enum, default_value = "ocr")]
+    task: TaskArg,
+
+    #[arg(long, default_value_t = 128)]
+    max_new_tokens: usize,
+
+    #[arg(long, value_name = "FILE")]
+    json_output: Option<PathBuf>,
+
+    #[arg(long, default_value_t = false)]
+    cpu: bool,
+
+    #[arg(long, value_name = "DIR")]
+    dataset_root: Option<PathBuf>,
+
+    #[arg(long, value_name = "FILE")]
+    ground_truth: Option<PathBuf>,
+
+    #[arg(long)]
+    limit: Option<usize>,
+
+    #[arg(long, default_value_t = 0)]
+    offset: usize,
+
+    #[arg(long, default_value_t = 10)]
+    sample_errors: usize,
+}
+
+#[derive(serde::Serialize)]
+#[serde(rename_all = "camelCase")]
+struct OutputEnvelope {
+    outputs: Vec<PaddleOcrVlOutput>,
+}
+
+#[derive(Debug, Clone)]
+struct GroundTruthRecord {
+    relative_path: PathBuf,
+    text: String,
+}
+
+#[derive(Debug, serde::Serialize)]
+#[serde(rename_all = "camelCase")]
+struct EvaluationError {
+    image: String,
+    expected: String,
+    predicted: String,
+    char_distance: usize,
+}
+
+#[derive(Debug, serde::Serialize)]
+#[serde(rename_all = "camelCase")]
+struct EvaluationSummary {
+    dataset_root: String,
+    ground_truth: String,
+    evaluated_examples: usize,
+    exact_matches: usize,
+    exact_match_rate: f64,
+    total_reference_chars: usize,
+    total_char_distance: usize,
+    cer: f64,
+    elapsed_seconds: f64,
+    sample_errors: Vec<EvaluationError>,
+}
+
+fn init_tracing() {
+    tracing_subscriber::fmt()
+        .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
+        .init();
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    init_tracing();
+    koharu_runtime::initialize()
+        .await
+        .context("failed to initialize runtime libraries")?;
+
+    let cli = Cli::parse();
+    let task: PaddleOcrVlTask = cli.task.into();
+
+    if cli.input.is_empty() && cli.dataset_root.is_none() {
+        bail!("provide either --input or --dataset-root");
+    }
+
+    let mut model = if let Some(model_dir) = &cli.model_dir {
+        PaddleOcrVl::load_from_dir(model_dir, cli.cpu)?
+    } else {
+        PaddleOcrVl::load(cli.cpu).await?
+    };
+
+    if let Some(dataset_root) = &cli.dataset_root {
+        let summary = evaluate_dataset(&mut model, dataset_root, &cli, task)?;
+        println!("{}", serde_json::to_string_pretty(&summary)?);
+        if let Some(path) = &cli.json_output {
+            std::fs::write(path, serde_json::to_string_pretty(&summary)?)?;
+        }
+        return Ok(());
+    }
+
+    let images = cli
+        .input
+        .iter()
+        .map(image::open)
+        .collect::<Result<Vec<_>, _>>()?;
+    let outputs = model.inference_images(&images, task, cli.max_new_tokens)?;
+    for (input, output) in cli.input.iter().zip(&outputs) {
+        if cli.input.len() > 1 {
+            println!("== {} ==", input.display());
+        }
+        println!("{}", output.text);
+    }
+
+    if let Some(path) = &cli.json_output {
+        std::fs::write(
+            path,
+            serde_json::to_string_pretty(&OutputEnvelope { outputs })?,
+        )?;
+    }
+
+    Ok(())
+}
+
+fn evaluate_dataset(
+    model: &mut PaddleOcrVl,
+    dataset_root: &Path,
+    cli: &Cli,
+    task: PaddleOcrVlTask,
+) -> Result<EvaluationSummary> {
+    let started = Instant::now();
+    let split_name = dataset_root
+        .file_name()
+        .and_then(|name| name.to_str())
+        .context("dataset root must end with a split directory like `train`")?;
+    let ground_truth = cli
+        .ground_truth
+        .clone()
+        .unwrap_or_else(|| dataset_root.with_file_name(format!("rec_gt_{split_name}.txt")));
+    let dataset_base = dataset_root
+        .parent()
+        .and_then(Path::parent)
+        .context("dataset root must look like `<base>/rec/<split>`")?;
+
+    let records = read_ground_truth(&ground_truth)?;
+    let selected = records
+        .into_iter()
+        .skip(cli.offset)
+        .take(cli.limit.unwrap_or(usize::MAX))
+        .collect::<Vec<_>>();
+    if selected.is_empty() {
+        bail!("no evaluation records selected");
+    }
+
+    let mut exact_matches = 0usize;
+    let mut total_reference_chars = 0usize;
+    let mut total_char_distance = 0usize;
+    let mut sample_errors = Vec::new();
+
+    for (index, record) in selected.iter().enumerate() {
+        let image_path = dataset_base.join(&record.relative_path);
+        let image = image::open(&image_path)
+            .with_context(|| format!("failed to open `{}`", image_path.display()))?;
+        let output = model.inference_with_max_new_tokens(&image, task, cli.max_new_tokens)?;
+        let predicted = output.text;
+        let expected = record.text.as_str();
+        let char_distance = levenshtein_chars(expected, &predicted);
+        let reference_chars = expected.chars().count();
+
+        total_reference_chars += reference_chars;
+        total_char_distance += char_distance;
+        if predicted == expected {
+            exact_matches += 1;
+        } else if sample_errors.len() < cli.sample_errors {
+            sample_errors.push(EvaluationError {
+                image: image_path.display().to_string(),
+                expected: expected.to_string(),
+                predicted,
+                char_distance,
+            });
+        }
+
+        if (index + 1) % 50 == 0 || index + 1 == selected.len() {
+            let cer = ratio(total_char_distance, total_reference_chars);
+            let exact_match_rate = ratio(exact_matches, index + 1);
+            eprintln!(
+                "evaluated {}/{}  exact={:.4}  cer={:.4}",
+                index + 1,
+                selected.len(),
+                exact_match_rate,
+                cer
+            );
+        }
+    }
+
+    Ok(EvaluationSummary {
+        dataset_root: dataset_root.display().to_string(),
+        ground_truth: ground_truth.display().to_string(),
+        evaluated_examples: selected.len(),
+        exact_matches,
+        exact_match_rate: ratio(exact_matches, selected.len()),
+        total_reference_chars,
+        total_char_distance,
+        cer: ratio(total_char_distance, total_reference_chars),
+        elapsed_seconds: started.elapsed().as_secs_f64(),
+        sample_errors,
+    })
+}
+
+fn read_ground_truth(path: &Path) -> Result<Vec<GroundTruthRecord>> {
+    let bytes = std::fs::read(path)
+        .with_context(|| format!("failed to read ground-truth file `{}`", path.display()))?;
+    let (text, _, had_errors) = encoding_rs::UTF_8.decode(&bytes);
+    if had_errors {
+        bail!("ground-truth file `{}` is not valid UTF-8", path.display());
+    }
+
+    let mut records = Vec::new();
+    for line in text.lines().filter(|line| !line.trim().is_empty()) {
+        if let Some((relative_path, value)) = line.split_once('\t') {
+            records.push(GroundTruthRecord {
+                relative_path: PathBuf::from(relative_path),
+                text: value.to_string(),
+            });
+            continue;
+        }
+
+        let previous = records.last_mut().with_context(|| {
+            format!("invalid ground-truth line without a preceding record: `{line}`")
+        })?;
+        previous.text.push('\n');
+        previous.text.push_str(line);
+    }
+
+    Ok(records)
+}
+
+fn levenshtein_chars(expected: &str, predicted: &str) -> usize {
+    let expected = expected.chars().collect::<Vec<_>>();
+    let predicted = predicted.chars().collect::<Vec<_>>();
+
+    if expected.is_empty() {
+        return predicted.len();
+    }
+    if predicted.is_empty() {
+        return expected.len();
+    }
+
+    let mut prev = (0..=predicted.len()).collect::<Vec<_>>();
+    let mut curr = vec![0usize; predicted.len() + 1];
+
+    for (i, expected_char) in expected.iter().enumerate() {
+        curr[0] = i + 1;
+        for (j, predicted_char) in predicted.iter().enumerate() {
+            let substitution_cost = usize::from(expected_char != predicted_char);
+            curr[j + 1] = (prev[j + 1] + 1)
+                .min(curr[j] + 1)
+                .min(prev[j] + substitution_cost);
+        }
+        std::mem::swap(&mut prev, &mut curr);
+    }
+
+    prev[predicted.len()]
+}
+
+fn ratio(numerator: usize, denominator: usize) -> f64 {
+    if denominator == 0 {
+        0.0
+    } else {
+        numerator as f64 / denominator as f64
+    }
+}

koharu-llm/src/lib.rs

@@ -2,6 +2,7 @@ pub mod api;
 pub mod facade;
 pub mod language;
 mod model;
+pub mod paddleocr_vl;
 pub mod prompt;
 pub mod providers;
 pub mod safe;

koharu-llm/src/model.rs

@@ -67,8 +67,7 @@ impl Llm {
     }
 
     fn load_from_path(id: ModelId, cpu: bool, model_path: PathBuf) -> Result<Self> {
-        crate::sys::initialize()
-            .context("failed to initialize llama.cpp runtime bindings")?;
+        crate::sys::initialize().context("failed to initialize llama.cpp runtime bindings")?;
 
         LOGGING_READY.call_once(|| {
             send_logs_to_tracing(LogOptions::default().with_logs_enabled(true));

koharu-llm/src/paddleocr_vl.rs

@@ -0,0 +1,579 @@
+use std::ffi::CString;
+use std::num::NonZeroU32;
+use std::path::{Path, PathBuf};
+use std::sync::Once;
+use std::time::Instant;
+
+use anyhow::{Context, Result, bail};
+use image::DynamicImage;
+use minijinja::{Environment, context};
+use serde::{Deserialize, Serialize};
+
+use crate::safe::context::params::{LlamaAttentionType, LlamaContextParams};
+use crate::safe::llama_backend::LlamaBackend;
+use crate::safe::llama_batch::LlamaBatch;
+use crate::safe::model::LlamaModel;
+use crate::safe::model::params::LlamaModelParams;
+use crate::safe::mtmd::{
+    MtmdBitmap, MtmdContext, MtmdContextParams, MtmdInputChunkType, MtmdInputText,
+};
+use crate::safe::sampling::LlamaSampler;
+use crate::safe::token::LlamaToken;
+use crate::safe::{LogOptions, send_logs_to_tracing};
+
+const HF_REPO: &str = "PaddlePaddle/PaddleOCR-VL-1.5-GGUF";
+const MODEL_FILENAME: &str = "PaddleOCR-VL-1.5.gguf";
+const MMPROJ_FILENAME: &str = "PaddleOCR-VL-1.5-mmproj.gguf";
+const DEFAULT_MEDIA_MARKER: &str = "<__media__>";
+const DEFAULT_GPU_LAYERS: u32 = 1000;
+const DEFAULT_MAX_NEW_TOKENS: usize = 128;
+const MAX_UBATCH: u32 = 512;
+
+static LOGGING_READY: Once = Once::new();
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum PaddleOcrVlTask {
+    Ocr,
+    Table,
+    Formula,
+    Chart,
+    Spotting,
+    Seal,
+}
+
+impl PaddleOcrVlTask {
+    fn prompt(self) -> &'static str {
+        match self {
+            Self::Ocr => "OCR:",
+            Self::Table => "Table Recognition:",
+            Self::Formula => "Formula Recognition:",
+            Self::Chart => "Chart Recognition:",
+            Self::Spotting => "Spotting:",
+            Self::Seal => "Seal Recognition:",
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct PaddleOcrVlOutput {
+    pub task: PaddleOcrVlTask,
+    pub text: String,
+    pub token_ids: Vec<u32>,
+    pub original_width: u32,
+    pub original_height: u32,
+    pub processed_width: u32,
+    pub processed_height: u32,
+    pub grid_thw: [u32; 3],
+    pub num_image_tokens: usize,
+}
+
+struct ModelFiles {
+    model: PathBuf,
+    mmproj: PathBuf,
+}
+
+struct RenderedPrompt {
+    text: String,
+    add_special: bool,
+}
+
+#[derive(Debug, Clone, Serialize)]
+struct PromptMessage {
+    role: &'static str,
+    content: String,
+}
+
+pub struct PaddleOcrVl {
+    backend: LlamaBackend,
+    model: LlamaModel,
+    chat_template: String,
+    bos_token: String,
+    eos_token_text: String,
+    mtmd: MtmdContext,
+    eos_token: LlamaToken,
+}
+
+impl PaddleOcrVl {
+    pub async fn load(cpu: bool) -> Result<Self> {
+        let files = download_model_files().await?;
+        tokio::task::spawn_blocking(move || Self::load_from_files(files, cpu))
+            .await
+            .context("failed to join PaddleOCR-VL loading task")?
+    }
+
+    pub fn load_from_dir(dir: impl AsRef<Path>, cpu: bool) -> Result<Self> {
+        let files = resolve_local_model_files(dir.as_ref())?;
+        Self::load_from_files(files, cpu)
+    }
+
+    fn load_from_files(files: ModelFiles, cpu: bool) -> Result<Self> {
+        crate::sys::initialize().context("failed to initialize llama.cpp runtime bindings")?;
+
+        LOGGING_READY.call_once(|| {
+            send_logs_to_tracing(LogOptions::default().with_logs_enabled(true));
+        });
+
+        let backend = LlamaBackend::init().context("unable to initialize llama.cpp backend")?;
+        let model_params = model_params(cpu, &backend);
+        let model = LlamaModel::load_from_file(&backend, &files.model, &model_params)
+            .with_context(|| format!("unable to load model from `{}`", files.model.display()))?;
+        let eos_token = model.token_eos();
+        let chat_template = model
+            .meta_val_str("tokenizer.ggml.chat_template")
+            .or_else(|_| model.meta_val_str("tokenizer.chat_template"))
+            .context("missing embedded PaddleOCR-VL chat template")?;
+        let bos_token = token_text(&model, model.token_bos());
+        let eos_token_text = token_text(&model, eos_token);
+        let mmproj_path = files
+            .mmproj
+            .to_str()
+            .with_context(|| format!("invalid mmproj path `{}`", files.mmproj.display()))?;
+        let mtmd = MtmdContext::init_from_file(
+            mmproj_path,
+            &model,
+            &MtmdContextParams {
+                use_gpu: !cpu && backend.supports_gpu_offload(),
+                print_timings: true,
+                n_threads: num_cpus::get().try_into().unwrap_or(i32::MAX),
+                media_marker: CString::new(DEFAULT_MEDIA_MARKER)
+                    .expect("default media marker contains no null bytes"),
+            },
+        )
+        .context("unable to initialize multimodal projector")?;
+
+        if !mtmd.support_vision() {
+            bail!("loaded PaddleOCR-VL projector does not advertise vision support");
+        }
+
+        tracing::info!(
+            has_encoder = unsafe { crate::sys::llama_model_has_encoder(model.model.as_ptr()) },
+            has_decoder = unsafe { crate::sys::llama_model_has_decoder(model.model.as_ptr()) },
+            decoder_start_token = model.decode_start_token().0,
+            "loaded PaddleOCR-VL model capabilities"
+        );
+
+        Ok(Self {
+            backend,
+            model,
+            chat_template,
+            bos_token,
+            eos_token_text,
+            mtmd,
+            eos_token,
+        })
+    }
+
+    pub fn inference(
+        &mut self,
+        image: &DynamicImage,
+        task: PaddleOcrVlTask,
+    ) -> Result<PaddleOcrVlOutput> {
+        self.inference_with_max_new_tokens(image, task, DEFAULT_MAX_NEW_TOKENS)
+    }
+
+    pub fn inference_with_max_new_tokens(
+        &mut self,
+        image: &DynamicImage,
+        task: PaddleOcrVlTask,
+        max_new_tokens: usize,
+    ) -> Result<PaddleOcrVlOutput> {
+        let started = Instant::now();
+        let original_width = image.width();
+        let original_height = image.height();
+        let bitmap = bitmap_from_image(image)?;
+        let prompt = self.render_prompt(task)?;
+        let chunks = self
+            .mtmd
+            .tokenize(
+                MtmdInputText {
+                    text: prompt.text,
+                    add_special: prompt.add_special,
+                    parse_special: true,
+                },
+                &[&bitmap],
+            )
+            .context("failed to tokenize multimodal OCR input")?;
+        if chunks.is_empty() {
+            bail!("multimodal tokenization produced no chunks");
+        }
+
+        let batch_tokens = max_chunk_tokens(&chunks).max(1);
+        let prompt_positions =
+            usize::try_from(chunks.total_positions()).context("prompt positions overflow usize")?;
+        let prompt_total_tokens = chunks.total_tokens();
+        let num_image_tokens = image_chunk_tokens(&chunks);
+        let ctx_params = context_params(
+            &self.mtmd,
+            prompt_positions,
+            prompt_total_tokens,
+            batch_tokens,
+            max_new_tokens,
+        )?;
+        let mut ctx = self
+            .model
+            .new_context(&self.backend, ctx_params)
+            .context("unable to create PaddleOCR-VL llama.cpp context")?;
+        let n_batch = i32::try_from(batch_tokens).context("batch size exceeds i32")?;
+
+        let prompt_started = Instant::now();
+        let n_past = chunks
+            .eval_chunks(&self.mtmd, &ctx, 0, 0, n_batch, true)
+            .context("failed to evaluate multimodal prompt")?;
+        let prompt_elapsed = prompt_started.elapsed();
+
+        let generation_started = Instant::now();
+        let mut sampler = LlamaSampler::greedy();
+        let mut decoder = encoding_rs::UTF_8.new_decoder();
+        let mut token_ids = Vec::new();
+        let mut text = String::new();
+
+        if max_new_tokens > 0 {
+            let decoder_start = self.decoder_start_token();
+            let (mut next_token, mut position) = if let Some(decoder_start) = decoder_start {
+                let mut batch = LlamaBatch::new(1, 1);
+                batch
+                    .add(decoder_start, 0, &[0], true)
+                    .context("failed to build decoder start batch")?;
+                ctx.decode(&mut batch)
+                    .context("failed to decode PaddleOCR-VL decoder start token")?;
+                (sampler.sample(&ctx, -1), 1)
+            } else {
+                (sampler.sample(&ctx, -1), n_past)
+            };
+
+            while token_ids.len() < max_new_tokens && !self.should_stop(next_token) {
+                sampler.accept(next_token);
+                token_ids
+                    .push(u32::try_from(next_token.0).context("generated token id was negative")?);
+                text.push_str(&decode_token(&self.model, next_token, &mut decoder)?);
+
+                if token_ids.len() >= max_new_tokens {
+                    break;
+                }
+
+                let mut batch = LlamaBatch::new(1, 1);
+                batch
+                    .add(next_token, position, &[0], true)
+                    .context("failed to build generation batch")?;
+                ctx.decode(&mut batch)
+                    .context("failed to decode generated OCR token")?;
+                position += 1;
+                next_token = sampler.sample(&ctx, -1);
+            }
+        }
+
+        tracing::info!(
+            task = ?task,
+            original_width,
+            original_height,
+            prompt_total_tokens,
+            prompt_positions,
+            batch_tokens,
+            num_image_tokens,
+            prompt_ms = prompt_elapsed.as_millis(),
+            generation_ms = generation_started.elapsed().as_millis(),
+            total_ms = started.elapsed().as_millis(),
+            "paddleocr-vl inference timings"
+        );
+
+        Ok(PaddleOcrVlOutput {
+            task,
+            text: text.trim().to_string(),
+            token_ids,
+            original_width,
+            original_height,
+            processed_width: original_width,
+            processed_height: original_height,
+            grid_thw: [0, 0, 0],
+            num_image_tokens,
+        })
+    }
+
+    pub fn inference_images(
+        &mut self,
+        images: &[DynamicImage],
+        task: PaddleOcrVlTask,
+        max_new_tokens: usize,
+    ) -> Result<Vec<PaddleOcrVlOutput>> {
+        let started = Instant::now();
+        let mut outputs = Vec::with_capacity(images.len());
+        for image in images {
+            outputs.push(self.inference_with_max_new_tokens(image, task, max_new_tokens)?);
+        }
+        tracing::info!(
+            images = images.len(),
+            total_ms = started.elapsed().as_millis(),
+            "paddleocr-vl batch timings"
+        );
+        Ok(outputs)
+    }
+
+    fn should_stop(&self, token: LlamaToken) -> bool {
+        token == self.eos_token || self.model.is_eog_token(token)
+    }
+
+    fn render_prompt(&self, task: PaddleOcrVlTask) -> Result<RenderedPrompt> {
+        let text = render_chat_prompt(
+            &self.chat_template,
+            &self.bos_token,
+            &self.eos_token_text,
+            task,
+        )
+        .context("failed to render PaddleOCR-VL prompt from embedded chat template")?;
+        Ok(RenderedPrompt {
+            text,
+            add_special: false,
+        })
+    }
+
+    fn decoder_start_token(&self) -> Option<LlamaToken> {
+        let has_encoder = unsafe { crate::sys::llama_model_has_encoder(self.model.model.as_ptr()) };
+        if !has_encoder {
+            return None;
+        }
+
+        let decoder_start = self.model.decode_start_token();
+        if decoder_start.0 >= 0 {
+            Some(decoder_start)
+        } else {
+            Some(self.model.token_bos())
+        }
+    }
+}
+
+pub async fn prefetch() -> Result<()> {
+    download_model_files().await?;
+    Ok(())
+}
+
+async fn download_model_files() -> Result<ModelFiles> {
+    let (model, mmproj) = tokio::try_join!(
+        koharu_http::download::model(HF_REPO, MODEL_FILENAME),
+        koharu_http::download::model(HF_REPO, MMPROJ_FILENAME),
+    )?;
+
+    Ok(ModelFiles { model, mmproj })
+}
+
+fn resolve_local_model_files(dir: &Path) -> Result<ModelFiles> {
+    let preferred_model = dir.join(MODEL_FILENAME);
+    let preferred_mmproj = dir.join(MMPROJ_FILENAME);
+    if preferred_model.exists() && preferred_mmproj.exists() {
+        return Ok(ModelFiles {
+            model: preferred_model,
+            mmproj: preferred_mmproj,
+        });
+    }
+
+    let mut model = None;
+    let mut mmproj = None;
+    for entry in std::fs::read_dir(dir)
+        .with_context(|| format!("unable to read PaddleOCR-VL model dir `{}`", dir.display()))?
+    {
+        let entry = entry?;
+        let path = entry.path();
+        if !path
+            .extension()
+            .and_then(|ext| ext.to_str())
+            .is_some_and(|ext| ext.eq_ignore_ascii_case("gguf"))
+        {
+            continue;
+        }
+
+        let filename = path
+            .file_name()
+            .and_then(|name| name.to_str())
+            .unwrap_or("");
+        if filename.contains("mmproj") {
+            if mmproj.is_none() {
+                mmproj = Some(path);
+            }
+        } else if model.is_none() {
+            model = Some(path);
+        }
+    }
+
+    Ok(ModelFiles {
+        model: model.with_context(|| {
+            format!(
+                "missing `{MODEL_FILENAME}` (or any non-mmproj GGUF) in `{}`",
+                dir.display()
+            )
+        })?,
+        mmproj: mmproj.with_context(|| {
+            format!(
+                "missing `{MMPROJ_FILENAME}` (or any mmproj GGUF) in `{}`",
+                dir.display()
+            )
+        })?,
+    })
+}
+
+fn model_params(cpu: bool, backend: &LlamaBackend) -> LlamaModelParams {
+    if !cpu && backend.supports_gpu_offload() {
+        LlamaModelParams::default().with_n_gpu_layers(DEFAULT_GPU_LAYERS)
+    } else {
+        LlamaModelParams::default()
+    }
+}
+
+fn context_params(
+    mtmd: &MtmdContext,
+    prompt_positions: usize,
+    prompt_total_tokens: usize,
+    batch_tokens: usize,
+    max_tokens: usize,
+) -> Result<LlamaContextParams> {
+    let required_ctx = prompt_positions
+        .saturating_add(max_tokens)
+        .saturating_add(1)
+        .max(
+            prompt_total_tokens
+                .saturating_add(max_tokens)
+                .saturating_add(1),
+        )
+        .max(batch_tokens.saturating_add(1));
+    let n_ctx = NonZeroU32::new(u32::try_from(required_ctx).context("context size exceeds u32")?)
+        .expect("required context size is non-zero");
+    let n_batch = u32::try_from(batch_tokens.max(1)).context("batch size exceeds u32")?;
+    let n_ubatch = if mtmd.decode_use_non_causal() {
+        n_batch
+    } else {
+        n_batch.min(MAX_UBATCH)
+    };
+
+    let mut params = LlamaContextParams::default()
+        .with_n_ctx(Some(n_ctx))
+        .with_n_batch(n_batch)
+        .with_n_ubatch(n_ubatch);
+    if mtmd.decode_use_non_causal() {
+        params = params.with_attention_type(LlamaAttentionType::NonCausal);
+    }
+    Ok(params)
+}
+
+fn bitmap_from_image(image: &DynamicImage) -> Result<MtmdBitmap> {
+    let rgb = image.to_rgb8();
+    let (width, height) = rgb.dimensions();
+    MtmdBitmap::from_image_data(width, height, &rgb.into_raw())
+        .context("failed to create MTMD bitmap from image")
+}
+
+fn build_user_message_content(task: PaddleOcrVlTask) -> String {
+    format!("{DEFAULT_MEDIA_MARKER}{}", task.prompt())
+}
+
+fn render_chat_prompt(
+    chat_template: &str,
+    bos_token: &str,
+    eos_token: &str,
+    task: PaddleOcrVlTask,
+) -> Result<String> {
+    let mut env = Environment::new();
+    env.add_filter("trim", |s: String| s.trim().to_string());
+    let tmpl = env
+        .template_from_str(chat_template)
+        .map_err(anyhow::Error::msg)
+        .context("failed to parse embedded PaddleOCR-VL chat template")?;
+    tmpl.render(context! {
+        messages => vec![PromptMessage {
+            role: "user",
+            content: build_user_message_content(task),
+        }],
+        bos_token => bos_token,
+        eos_token => eos_token,
+        add_generation_prompt => true,
+    })
+    .map_err(anyhow::Error::msg)
+    .context("failed to evaluate embedded PaddleOCR-VL chat template")
+}
+
+fn max_chunk_tokens(chunks: &crate::safe::mtmd::MtmdInputChunks) -> usize {
+    (0..chunks.len())
+        .filter_map(|index| chunks.get(index))
+        .map(|chunk| chunk.n_tokens())
+        .max()
+        .unwrap_or(0)
+}
+
+fn image_chunk_tokens(chunks: &crate::safe::mtmd::MtmdInputChunks) -> usize {
+    (0..chunks.len())
+        .filter_map(|index| chunks.get(index))
+        .filter(|chunk| chunk.chunk_type() == MtmdInputChunkType::Image)
+        .map(|chunk| chunk.n_tokens())
+        .sum()
+}
+
+fn decode_token(
+    model: &LlamaModel,
+    token: LlamaToken,
+    decoder: &mut encoding_rs::Decoder,
+) -> Result<String> {
+    model
+        .token_to_piece(token, decoder, true, None)
+        .context("failed to decode generated OCR token")
+}
+
+fn token_text(model: &LlamaModel, token: LlamaToken) -> String {
+    let mut decoder = encoding_rs::UTF_8.new_decoder();
+    match model.token_to_piece(token, &mut decoder, true, None) {
+        Ok(piece) if !piece.is_empty() => piece,
+        _ => token.to_string(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        DEFAULT_MEDIA_MARKER, PaddleOcrVlTask, build_user_message_content, render_chat_prompt,
+    };
+
+    #[test]
+    fn user_message_places_media_marker_before_task() {
+        assert_eq!(
+            build_user_message_content(PaddleOcrVlTask::Ocr),
+            format!("{DEFAULT_MEDIA_MARKER}OCR:")
+        );
+    }
+
+    #[test]
+    fn user_message_contents_match_paddleocr_tasks() {
+        assert_eq!(
+            build_user_message_content(PaddleOcrVlTask::Table),
+            "<__media__>Table Recognition:"
+        );
+        assert_eq!(
+            build_user_message_content(PaddleOcrVlTask::Formula),
+            "<__media__>Formula Recognition:"
+        );
+        assert_eq!(
+            build_user_message_content(PaddleOcrVlTask::Chart),
+            "<__media__>Chart Recognition:"
+        );
+        assert_eq!(
+            build_user_message_content(PaddleOcrVlTask::Spotting),
+            "<__media__>Spotting:"
+        );
+        assert_eq!(
+            build_user_message_content(PaddleOcrVlTask::Seal),
+            "<__media__>Seal Recognition:"
+        );
+    }
+
+    #[test]
+    fn embedded_template_render_keeps_media_marker_in_user_turn() -> anyhow::Result<()> {
+        let rendered = render_chat_prompt(
+            "{{ bos_token }}{% for message in messages %}User: {{ message['content'] }}\n{% endfor %}{% if add_generation_prompt %}Assistant:\n{% endif %}",
+            "<|begin_of_sentence|>",
+            "</s>",
+            PaddleOcrVlTask::Formula,
+        )?;
+        assert_eq!(
+            rendered,
+            "<|begin_of_sentence|>User: <__media__>Formula Recognition:\nAssistant:\n"
+        );
+        Ok(())
+    }
+}

koharu-llm/src/safe/model.rs

@@ -401,8 +401,14 @@ impl LlamaModel {
             ),
             x => x,
         }?;
-        // here the assumption is that each byte from the output may map to at most one output charakter
-        let mut output_piece = String::with_capacity(bytes.len());
+        let mut output_piece = String::new();
+        // encoding_rs treats the String capacity as the output buffer size and will not
+        // reallocate, so reserve using the decoder's current state before decoding.
+        output_piece.reserve(
+            decoder
+                .max_utf8_buffer_length(bytes.len())
+                .unwrap_or(bytes.len().saturating_mul(3).saturating_add(3)),
+        );
         // _result only tells if there is nothing more in the input, or if the output was full
         // but further decoding will happen on the next interation anyway
         let (_result, _somesize, _truthy) =

koharu-llm/src/sys/mod.rs

@@ -63,8 +63,9 @@ const LIB_NAMES: [&str; 4] = [
 static LIBRARIES: OnceLock<LoadedLibraries> = OnceLock::new();
 
 pub(crate) fn initialize() -> Result<()> {
-    let runtime_dir = koharu_runtime::llama_runtime_dir()
-        .context("failed to resolve the llama runtime directory; call `koharu_runtime::initialize()` first")?;
+    let runtime_dir = koharu_runtime::llama_runtime_dir().context(
+        "failed to resolve the llama runtime directory; call `koharu_runtime::initialize()` first",
+    )?;
 
     if !runtime_dir.exists() {
         bail!(
@@ -73,9 +74,9 @@ pub(crate) fn initialize() -> Result<()> {
         );
     }
 
-    let dir = runtime_dir.canonicalize().with_context(|| {
-        format!("failed to canonicalize `{}`", runtime_dir.display())
-    })?;
+    let dir = runtime_dir
+        .canonicalize()
+        .with_context(|| format!("failed to canonicalize `{}`", runtime_dir.display()))?;
 
     if let Some(existing) = LIBRARIES.get() {
         if existing.path == dir {

koharu-ml/Cargo.toml

@@ -13,6 +13,7 @@ publish.workspace = true
 
 [dependencies]
 koharu-http = { workspace = true }
+koharu-llm = { workspace = true }
 koharu-types = { workspace = true }
 anyhow = { workspace = true }
 clap = { workspace = true }

koharu-ml/src/facade.rs

@@ -2,12 +2,12 @@ use std::{sync::Mutex, time::Instant};
 
 use anyhow::Result;
 use image::DynamicImage;
+use koharu_llm::paddleocr_vl::{self as paddleocr_vl_llm, PaddleOcrVl, PaddleOcrVlTask};
 use koharu_types::{Document, FontPrediction, SerializableDynamicImage, TextBlock, TextDirection};
 
 use crate::comic_text_detector::{self, ComicTextDetector, crop_text_block_bbox};
 use crate::font_detector::{self, FontDetector};
 use crate::lama::{self, Lama};
-use crate::paddleocr_vl::{self, PaddleOcrVl, PaddleOcrVlTask};
 use crate::pp_doclayout_v3::{self, LayoutRegion, PPDocLayoutV3};
 
 const NEAR_BLACK_THRESHOLD: u8 = 12;
@@ -244,7 +244,7 @@ impl Model {
 pub async fn prefetch() -> Result<()> {
     pp_doclayout_v3::prefetch().await?;
     comic_text_detector::prefetch_segmentation().await?;
-    paddleocr_vl::prefetch().await?;
+    paddleocr_vl_llm::prefetch().await?;
     lama::prefetch().await?;
     font_detector::prefetch().await?;
 

koharu-runtime/src/cuda.rs

@@ -256,8 +256,10 @@ mod tests {
             }
         }
 
-        let all_dylibs: Vec<&str> =
-            WHEELS.iter().flat_map(|w| w.dylibs().iter().copied()).collect();
+        let all_dylibs: Vec<&str> = WHEELS
+            .iter()
+            .flat_map(|w| w.dylibs().iter().copied())
+            .collect();
         for dylib in &all_dylibs {
             assert!(root.join(dylib).exists());
         }

koharu-runtime/src/lib.rs

@@ -58,8 +58,7 @@ fn is_up_to_date(install_dir: &Path, source_id: &str) -> bool {
 
 fn reset_dir(dir: &Path) -> Result<()> {
     if dir.exists() {
-        fs::remove_dir_all(dir)
-            .with_context(|| format!("failed to reset `{}`", dir.display()))?;
+        fs::remove_dir_all(dir).with_context(|| format!("failed to reset `{}`", dir.display()))?;
     }
     fs::create_dir_all(dir).with_context(|| format!("failed to create `{}`", dir.display()))?;
     Ok(())

koharu-runtime/src/llama.rs

@@ -235,11 +235,13 @@ mod tests {
     fn install_dir_includes_tag_and_id() {
         let root = Path::new("/tmp/rt");
         let dir = LlamaRuntime::WindowsVulkanX64.install_dir(root);
-        assert!(dir.ends_with(
-            Path::new("llama.cpp")
-                .join(LLAMA_CPP_TAG)
-                .join("windows-vulkan-x64")
-        ));
+        assert!(
+            dir.ends_with(
+                Path::new("llama.cpp")
+                    .join(LLAMA_CPP_TAG)
+                    .join("windows-vulkan-x64")
+            )
+        );
     }
 
     #[test]

koharu/src/app.rs

@@ -292,4 +292,3 @@ pub async fn run() -> Result<()> {
 
     Ok(())
 }
-