fix: shared llama.cpp backend

koharu-llm/bin/llm.rs

@@ -1,6 +1,9 @@
+use std::sync::Arc;
+
 use clap::Parser;
 use tracing_subscriber::fmt::format::FmtSpan;
 
+use koharu_llm::safe::llama_backend::LlamaBackend;
 use koharu_llm::{GenerateOptions, Language, Llm, ModelId};
 
 #[derive(Parser, Debug)]
@@ -72,6 +75,7 @@ fn init_tracing() {
 
 async fn initialize_runtime() -> anyhow::Result<()> {
     koharu_runtime::initialize().await?;
+    koharu_llm::sys::initialize()?;
     Ok(())
 }
 
@@ -81,8 +85,9 @@ async fn main() -> anyhow::Result<()> {
 
     let args = Args::parse();
     initialize_runtime().await?;
+    let backend = Arc::new(LlamaBackend::init()?);
 
-    let mut llm = Llm::load(args.model, args.cpu).await?;
+    let mut llm = Llm::load(args.model, args.cpu, backend).await?;
     let target_language = Language::parse(&args.locale).unwrap_or(Language::English);
 
     let opts = GenerateOptions {

koharu-llm/bin/paddleocr-vl.rs

@@ -1,9 +1,11 @@
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 use std::time::Instant;
 
 use anyhow::{Context, Result, bail};
 use clap::{Parser, ValueEnum};
 use koharu_llm::paddleocr_vl::{PaddleOcrVl, PaddleOcrVlOutput, PaddleOcrVlTask};
+use koharu_llm::safe::llama_backend::LlamaBackend;
 
 #[derive(Debug, Clone, Copy, ValueEnum)]
 enum TaskArg {
@@ -112,18 +114,20 @@ async fn main() -> Result<()> {
     koharu_runtime::initialize()
         .await
         .context("failed to initialize runtime libraries")?;
+    koharu_llm::sys::initialize().context("failed to initialize llama.cpp runtime bindings")?;
 
     let cli = Cli::parse();
     let task: PaddleOcrVlTask = cli.task.into();
+    let backend = Arc::new(LlamaBackend::init().context("unable to initialize llama.cpp backend")?);
 
     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)?
+        PaddleOcrVl::load_from_dir(model_dir, cli.cpu, Arc::clone(&backend))?
     } else {
-        PaddleOcrVl::load(cli.cpu).await?
+        PaddleOcrVl::load(cli.cpu, backend).await?
     };
 
     if let Some(dataset_root) = &cli.dataset_root {

koharu-llm/src/facade.rs

@@ -6,7 +6,8 @@ use tokio::sync::{RwLock, broadcast};
 use koharu_types::{Document, LlmState, LlmStateStatus, TextBlock};
 
 use crate::{
-    GenerateOptions, Language, Llm, ModelId, language::tags as language_tags, supported_locales,
+    GenerateOptions, Language, Llm, ModelId, language::tags as language_tags,
+    safe::llama_backend::LlamaBackend, supported_locales,
 };
 
 pub use crate::prefetch;
@@ -74,12 +75,7 @@ pub struct Model {
     state: Arc<RwLock<State>>,
     state_tx: broadcast::Sender<LlmState>,
     cpu: bool,
-}
-
-impl Default for Model {
-    fn default() -> Self {
-        Self::new(false)
-    }
+    backend: Arc<LlamaBackend>,
 }
 
 pub trait Translatable {
@@ -441,11 +437,12 @@ impl Translatable for TextBlock {
 }
 
 impl Model {
-    pub fn new(cpu: bool) -> Self {
+    pub fn new(cpu: bool, backend: Arc<LlamaBackend>) -> Self {
         Self {
             state: Arc::new(RwLock::new(State::Empty)),
             state_tx: broadcast::channel(64).0,
             cpu,
+            backend,
         }
     }
 
@@ -482,8 +479,9 @@ impl Model {
         let state_cloned = self.state.clone();
         let state_tx = self.state_tx.clone();
         let cpu = self.cpu;
+        let backend = self.backend.clone();
         tokio::spawn(async move {
-            let res = Llm::load(id, cpu).await;
+            let res = Llm::load(id, cpu, backend).await;
             match res {
                 Ok(llm) => {
                     let mut guard = state_cloned.write().await;

koharu-llm/src/model.rs

@@ -1,5 +1,6 @@
 use std::num::NonZeroU32;
 use std::path::PathBuf;
+use std::sync::Arc;
 use std::sync::Once;
 use std::time::Instant;
 
@@ -24,7 +25,7 @@ const SAKURA_QWEN_CORRECT_EOS_ID: i32 = 151645;
 
 pub struct Llm {
     model_id: ModelId,
-    backend: LlamaBackend,
+    backend: Arc<LlamaBackend>,
     model: LlamaModel,
     prompt_renderer: PromptRenderer,
     eos_token: LlamaToken,
@@ -58,24 +59,28 @@ impl Default for GenerateOptions {
 }
 
 impl Llm {
-    pub async fn load(id: ModelId, cpu: bool) -> Result<Self> {
+    pub async fn load(id: ModelId, cpu: bool, backend: Arc<LlamaBackend>) -> Result<Self> {
         let model_path = id.get().await?;
 
-        tokio::task::spawn_blocking(move || Self::load_from_path(id, cpu, model_path))
+        tokio::task::spawn_blocking(move || Self::load_from_path(id, cpu, model_path, backend))
             .await
             .context("failed to join llama.cpp model loading task")?
     }
 
-    fn load_from_path(id: ModelId, cpu: bool, model_path: PathBuf) -> Result<Self> {
+    fn load_from_path(
+        id: ModelId,
+        cpu: bool,
+        model_path: PathBuf,
+        backend: Arc<LlamaBackend>,
+    ) -> 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, &model_path, &model_params)
+        let model_params = model_params(cpu, backend.as_ref());
+        let model = LlamaModel::load_from_file(backend.as_ref(), &model_path, &model_params)
             .with_context(|| format!("unable to load model from `{}`", model_path.display()))?;
 
         let chat_template = model
@@ -126,7 +131,7 @@ impl Llm {
         let mut ctx = self
             .model
             .new_context(
-                &self.backend,
+                self.backend.as_ref(),
                 context_params(prompt_tokens.len(), opts.max_tokens)?,
             )
             .context("unable to create llama.cpp context")?;

koharu-llm/src/paddleocr_vl.rs

@@ -1,6 +1,7 @@
 use std::ffi::CString;
 use std::num::NonZeroU32;
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 use std::sync::Once;
 use std::time::Instant;
 
@@ -86,7 +87,7 @@ struct PromptMessage {
 }
 
 pub struct PaddleOcrVl {
-    backend: LlamaBackend,
+    backend: Arc<LlamaBackend>,
     model: LlamaModel,
     chat_template: String,
     bos_token: String,
@@ -96,28 +97,31 @@ pub struct PaddleOcrVl {
 }
 
 impl PaddleOcrVl {
-    pub async fn load(cpu: bool) -> Result<Self> {
+    pub async fn load(cpu: bool, backend: Arc<LlamaBackend>) -> Result<Self> {
         let files = download_model_files().await?;
-        tokio::task::spawn_blocking(move || Self::load_from_files(files, cpu))
+        tokio::task::spawn_blocking(move || Self::load_from_files(files, cpu, backend))
             .await
             .context("failed to join PaddleOCR-VL loading task")?
     }
 
-    pub fn load_from_dir(dir: impl AsRef<Path>, cpu: bool) -> Result<Self> {
+    pub fn load_from_dir(
+        dir: impl AsRef<Path>,
+        cpu: bool,
+        backend: Arc<LlamaBackend>,
+    ) -> Result<Self> {
         let files = resolve_local_model_files(dir.as_ref())?;
-        Self::load_from_files(files, cpu)
+        Self::load_from_files(files, cpu, backend)
     }
 
-    fn load_from_files(files: ModelFiles, cpu: bool) -> Result<Self> {
+    fn load_from_files(files: ModelFiles, cpu: bool, backend: Arc<LlamaBackend>) -> 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)
+        let model_params = model_params(cpu, backend.as_ref());
+        let model = LlamaModel::load_from_file(backend.as_ref(), &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
@@ -134,7 +138,7 @@ impl PaddleOcrVl {
             mmproj_path,
             &model,
             &MtmdContextParams {
-                use_gpu: !cpu && backend.supports_gpu_offload(),
+                use_gpu: !cpu && backend.as_ref().supports_gpu_offload(),
                 print_timings: true,
                 n_threads: num_cpus::get().try_into().unwrap_or(i32::MAX),
                 media_marker: CString::new(DEFAULT_MEDIA_MARKER)
@@ -213,7 +217,7 @@ impl PaddleOcrVl {
         )?;
         let mut ctx = self
             .model
-            .new_context(&self.backend, ctx_params)
+            .new_context(self.backend.as_ref(), ctx_params)
             .context("unable to create PaddleOCR-VL llama.cpp context")?;
         let n_batch = i32::try_from(batch_tokens).context("batch size exceeds i32")?;
 

koharu-llm/src/sys/mod.rs

@@ -62,7 +62,7 @@ const LIB_NAMES: [&str; 4] = [
 
 static LIBRARIES: OnceLock<LoadedLibraries> = OnceLock::new();
 
-pub(crate) fn initialize() -> Result<()> {
+pub fn initialize() -> Result<()> {
     let runtime_dir = koharu_runtime::llama_runtime_dir().context(
         "failed to resolve the llama runtime directory; call `koharu_runtime::initialize()` first",
     )?;

koharu-llm/tests/llm.rs

@@ -1,7 +1,9 @@
 use std::path::PathBuf;
+use std::sync::Arc;
 
 use strum::IntoEnumIterator;
 
+use koharu_llm::safe::llama_backend::LlamaBackend;
 use koharu_llm::{GenerateOptions, Language, Llm, ModelId};
 
 async fn initialize_runtime() -> anyhow::Result<()> {
@@ -23,9 +25,11 @@ async fn llm_generates_text_for_all_models() -> anyhow::Result<()> {
 
     koharu_http::hf_hub::set_cache_dir(model_dir)?;
     initialize_runtime().await?;
+    koharu_llm::sys::initialize()?;
+    let backend = Arc::new(LlamaBackend::init()?);
 
     for model in ModelId::iter() {
-        let mut llm = Llm::load(model, false).await?;
+        let mut llm = Llm::load(model, false, Arc::clone(&backend)).await?;
         let opts = GenerateOptions {
             max_tokens: 100,
             temperature: 0.3,

koharu-ml/src/facade.rs

@@ -1,8 +1,12 @@
-use std::{sync::Mutex, time::Instant};
+use std::{
+    sync::{Arc, Mutex},
+    time::Instant,
+};
 
 use anyhow::Result;
 use image::DynamicImage;
 use koharu_llm::paddleocr_vl::{self as paddleocr_vl_llm, PaddleOcrVl, PaddleOcrVlTask};
+use koharu_llm::safe::llama_backend::LlamaBackend;
 use koharu_types::{Document, FontPrediction, SerializableDynamicImage, TextBlock, TextDirection};
 
 use crate::comic_text_detector::{self, ComicTextDetector, crop_text_block_bbox};
@@ -83,11 +87,11 @@ pub struct Model {
 }
 
 impl Model {
-    pub async fn new(cpu: bool) -> Result<Self> {
+    pub async fn new(cpu: bool, backend: Arc<LlamaBackend>) -> Result<Self> {
         Ok(Self {
             layout_detector: PPDocLayoutV3::load(cpu).await?,
             segmenter: ComicTextDetector::load_segmentation_only(cpu).await?,
-            ocr: Mutex::new(PaddleOcrVl::load(cpu).await?),
+            ocr: Mutex::new(PaddleOcrVl::load(cpu, backend).await?),
             lama: Lama::load(cpu).await?,
             font_detector: FontDetector::load(cpu).await?,
         })

koharu/src/app.rs

@@ -2,13 +2,14 @@ use std::{path::PathBuf, sync::Arc};
 
 use anyhow::{Context, Result};
 use clap::Parser;
-use once_cell::sync::Lazy;
+use once_cell::sync::{Lazy, OnceCell};
 use rfd::MessageDialog;
 use tauri::{Manager, WebviewWindowBuilder};
 use tokio::{net::TcpListener, sync::RwLock};
 use tracing_subscriber::fmt::format::FmtSpan;
 
 use koharu_llm::facade;
+use koharu_llm::safe::llama_backend::LlamaBackend;
 use koharu_ml::{cuda_is_available, device};
 use koharu_pipeline::AppResources;
 use koharu_renderer::facade::Renderer;
@@ -21,6 +22,7 @@ static APP_ROOT: Lazy<PathBuf> = Lazy::new(|| {
         .unwrap_or_default()
 });
 static MODEL_ROOT: Lazy<PathBuf> = Lazy::new(|| APP_ROOT.join("models"));
+static LLAMA_BACKEND: OnceCell<Arc<LlamaBackend>> = OnceCell::new();
 
 #[derive(Parser)]
 #[command(version = crate::version::APP_VERSION, about)]
@@ -113,6 +115,15 @@ async fn prefetch() -> Result<()> {
     Ok(())
 }
 
+fn shared_llama_backend() -> Result<Arc<LlamaBackend>> {
+    let backend = LLAMA_BACKEND.get_or_try_init(|| -> Result<Arc<LlamaBackend>> {
+        koharu_llm::sys::initialize().context("failed to initialize llama.cpp runtime bindings")?;
+        let backend = LlamaBackend::init().context("unable to initialize llama.cpp backend")?;
+        Ok(Arc::new(backend))
+    })?;
+    Ok(Arc::clone(backend))
+}
+
 fn warning(headless: bool, title: &str, description: &str) {
     tracing::warn!("{description}");
 
@@ -173,12 +184,13 @@ async fn build_resources(cpu: bool, headless: bool) -> Result<AppResources> {
         tracing::info!("CUDA is available and runtime packages were initialized");
     }
 
+    let llama_backend = shared_llama_backend()?;
     let ml = Arc::new(
-        koharu_ml::facade::Model::new(cpu)
+        koharu_ml::facade::Model::new(cpu, Arc::clone(&llama_backend))
             .await
             .context("Failed to initialize ML model")?,
     );
-    let llm = Arc::new(facade::Model::new(cpu));
+    let llm = Arc::new(facade::Model::new(cpu, llama_backend));
     let renderer = Arc::new(Renderer::new().context("Failed to initialize renderer")?);
     let state = Arc::new(RwLock::new(State::default()));