fix cache, fix logging

Cargo.lock

@@ -608,9 +608,11 @@ dependencies = [
  "candle-core",
  "candle-nn",
  "candle-transformers",
+ "env_logger",
  "hf-hub",
  "image",
  "log",
+ "memory_caching",
  "once_cell",
  "rand 0.8.5",
  "serde",
@@ -3171,9 +3173,9 @@ dependencies = [
 
 [[package]]
 name = "env_logger"
-version = "0.11.3"
+version = "0.11.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38b35839ba51819680ba087cd351788c9a3c476841207e0b8cee0b04722343b9"
+checksum = "dcaee3d8e3cfc3fd92428d477bc97fc29ec8716d180c0d74c643bb26166660e0"
 dependencies = [
  "anstream",
  "anstyle",

crates/desktop/src-tauri/Cargo.toml

@@ -18,21 +18,27 @@ crate-type = ["staticlib", "cdylib", "rlib"]
 tauri-build = { version = "2.0.4", features = [] }
 
 [dependencies]
-serde_json = "1.0"
-serde = { version = "1.0", features = ["derive"] }
-log = "0.4"
-tauri = { version = "2.2.4", features = [] }
-tauri-plugin-log = "2.0.0-rc"
-once_cell = "1.20.3"
-image = "0.25.5"
+
+# Internal
+memory_caching = { path = "../../lib/caching/memory_caching" }
+
+# External
+anyhow = "1.0"
 base64 = "0.22.1"
 candle-core = { version = "0.8.2", features = ["cuda"] }
-candle-transformers = { version = "0.8.2", features = ["cuda"] }
 candle-nn = { version = "0.8.2", features = ["cuda"] }
-tokenizers = "0.15.2"
+candle-transformers = { version = "0.8.2", features = ["cuda"] }
+env_logger = "0.11.6"
 hf-hub = "0.3.2"
-anyhow = "1.0"
+image = "0.25.5"
+log = "0.4"
+once_cell = "1.20.3"
 rand = "0.8"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+tauri = { version = "2.2.4", features = [] }
+tauri-plugin-log = "2.0.0-rc"
+tokenizers = "0.15.2"
 
 [features]
 default = ["cuda"]

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -9,6 +9,7 @@ use image::imageops::FilterType;
 use image::ImageReader;
 use log::error;
 use tauri::State;
+use crate::ml::prompt_cache::PromptCache;
 
 const PROMPT: &str = "A beautiful landscape with mountains and a lake";
 
@@ -16,7 +17,12 @@ const HEIGHT: usize = 512;
 const WIDTH: usize = 512;
 
 #[tauri::command]
-pub fn infer_image(image: &str, model_config: State<ModelConfig>, model_cache: State<ModelCache>) -> Result<String, String> {
+pub fn infer_image(
+  image: &str, 
+  model_config: State<ModelConfig>, 
+  model_cache: State<ModelCache>,
+  prompt_cache: State<PromptCache>,
+) -> Result<String, String> {
 
   let bytes = BASE64_STANDARD.decode(image)
     .map_err(|err| format!("Base64 decode error: {}", err))?;
@@ -48,7 +54,7 @@ pub fn infer_image(image: &str, model_config: State<ModelConfig>, model_cache: S
   image.save(&image_path)
     .map_err(|err| format!("Failed to save image: {}", err))?;
   
-  let result = do_infer_image(&prompt, &image_path, &model_config, &model_cache);
+  let result = do_infer_image(&prompt, &image_path, &model_config, &model_cache, prompt_cache);
   
   if let Err(err) = result.as_deref() {
     error!("There was an error: {:?}", err);
@@ -57,7 +63,13 @@ pub fn infer_image(image: &str, model_config: State<ModelConfig>, model_cache: S
   result
 }
 
-fn do_infer_image(prompt: &str, image_path: &PathBuf, config: &ModelConfig, model_cache: &ModelCache) -> Result<String, String> {
+fn do_infer_image(
+  prompt: &str, 
+  image_path: &PathBuf, 
+  config: &ModelConfig, 
+  model_cache: &ModelCache,
+  prompt_cache: State<PromptCache>,
+) -> Result<String, String> {
   println!("infer_image called; generating image with SDXL Turbo...");
   
   let args = Args {
@@ -76,6 +88,7 @@ fn do_infer_image(prompt: &str, image_path: &PathBuf, config: &ModelConfig, mode
     guidance_scale: None,
     model_cache,
     model_configs: config,
+    prompt_cache: &prompt_cache,
   };
 
   match run(args) {

crates/desktop/src-tauri/src/endpoints/sd.rs

@@ -23,6 +23,7 @@ use crate::ml::load_image_file_to_tensor::load_image_file_to_tensor;
 use crate::ml::load_image_file_to_tensor_2::load_image_file_to_tensor_2;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::models::unet_model::UNetModel;
+use crate::ml::prompt_cache::PromptCache;
 use crate::model_config::ModelConfig;
 // TODO: Clean up
 
@@ -44,6 +45,7 @@ pub struct Args<'a> {
     pub guidance_scale: Option<f64>,
     pub model_cache: &'a ModelCache,
     pub model_configs: &'a ModelConfig,
+    pub prompt_cache: &'a PromptCache,
 }
 
 pub fn run(args: Args<'_>) -> Result<()> {
@@ -93,19 +95,32 @@ pub fn run(args: Args<'_>) -> Result<()> {
     //let tokenizer = tokenizers::Tokenizer::from_file(tokenizer).map_err(E::msg)?;
     //let tokenizer2 = tokenizers::Tokenizer::from_file(tokenizer2).map_err(E::msg)?;
     
-    let text_embeddings = infer_clip_text_embeddings(
-          &args.prompt,
-          &args.uncond_prompt,
-          None, // tokenizer
-          None, // clip_weights
-          None, // clip2_weights
-          args.sd_version,
-          &args.model_configs.sd_config,
-          false, // use_f16
-          &args.model_configs.device,
-          args.model_configs.dtype,
-          false, // use_guide_scale
-      )?;
+    info!("Checking if prompt is cached");
+    let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?;
+    
+    let text_embeddings = if let Some(tensor) = maybe_cached {
+        info!("Prompt is cached!");
+        tensor
+    } else {
+        info!("Prompt is NOT cached!");
+        let tensor = infer_clip_text_embeddings(
+            &args.prompt,
+            &args.uncond_prompt,
+            None, // tokenizer
+            None, // clip_weights
+            None, // clip2_weights
+            args.sd_version,
+            &args.model_configs.sd_config,
+            false, // use_f16
+            &args.model_configs.device,
+            args.model_configs.dtype,
+            false, // use_guide_scale
+        )?;
+        
+        args.prompt_cache.store_copy(&args.prompt, &tensor)?;
+        
+        tensor
+    };
     
     println!("Text embeddings shape: {:?}", text_embeddings.shape());
 

crates/desktop/src-tauri/src/lib.rs

@@ -3,19 +3,26 @@ pub mod endpoints;
 pub mod model_config;
 
 use crate::ml::model_cache::ModelCache;
+use crate::ml::prompt_cache::PromptCache;
 use crate::model_config::ModelConfig;
 use endpoints::image_endpoint::infer_image;
+use env_logger;
 use log::info;
 use once_cell::sync::Lazy;
 use std::sync::{Arc, RwLock};
 
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
+  env_logger::init();
+
   info!("Loading model config...");
 
   let model_config = ModelConfig::init()
     .expect("config should load");
 
+  let prompt_cache = PromptCache::with_capacity(8)
+    .expect("prompt cache should load");
+
   info!("Creating model cache...");
   
   let model_cache = ModelCache::new(
@@ -39,6 +46,7 @@ pub fn run() {
       Ok(())
     })
     .manage(model_config)
+    .manage(prompt_cache)
     .manage(model_cache)
     .invoke_handler(tauri::generate_handler![
       my_custom_command, 

crates/desktop/src-tauri/src/ml/mod.rs

@@ -5,4 +5,5 @@ pub mod load_image_file_to_tensor_2;
 pub mod model_cache;
 pub mod model_file;
 pub mod models;
+pub mod prompt_cache;
 pub mod save_image_from_tensor;

crates/desktop/src-tauri/src/ml/prompt_cache.rs

@@ -0,0 +1,22 @@
+use candle_core::Tensor;
+use memory_caching::arc_sieve::ArcSieve;
+
+pub struct PromptCache {
+  cache: ArcSieve<String, Tensor>,
+}
+
+impl PromptCache {
+  pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
+    Ok(Self {
+      cache: ArcSieve::with_capacity(capacity)?,
+    })
+  }
+  
+  pub fn get_copy(&self, key: &String) -> anyhow::Result<Option<Tensor>> {
+    self.cache.get_copy(key)
+  }
+  
+  pub fn store_copy(&self, key: &String, value: &Tensor) -> anyhow::Result<()> {
+    self.cache.store_copy(key, value)
+  }
+}