From 98f2ce9c3146795e6313e4851f6ea7ad2e280b30 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 17 Feb 2025 17:48:58 -0500 Subject: [PATCH] fix cache, fix logging --- Cargo.lock | 6 ++- crates/desktop/src-tauri/Cargo.toml | 26 +++++++----- .../src-tauri/src/endpoints/image_endpoint.rs | 19 +++++++-- crates/desktop/src-tauri/src/endpoints/sd.rs | 41 +++++++++++++------ crates/desktop/src-tauri/src/lib.rs | 8 ++++ crates/desktop/src-tauri/src/ml/mod.rs | 1 + .../desktop/src-tauri/src/ml/prompt_cache.rs | 22 ++++++++++ 7 files changed, 95 insertions(+), 28 deletions(-) create mode 100644 crates/desktop/src-tauri/src/ml/prompt_cache.rs diff --git a/Cargo.lock b/Cargo.lock index bddcc471b8..98f7affe53 100644 --- a/Cargo.lock +++ b/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", diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 7dd9a5e191..52db39a4bf 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/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"] diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index ab54221021..c38d90ab3b 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/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, model_cache: State) -> Result { +pub fn infer_image( + image: &str, + model_config: State, + model_cache: State, + prompt_cache: State, +) -> Result { 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, 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, model_cache: S result } -fn do_infer_image(prompt: &str, image_path: &PathBuf, config: &ModelConfig, model_cache: &ModelCache) -> Result { +fn do_infer_image( + prompt: &str, + image_path: &PathBuf, + config: &ModelConfig, + model_cache: &ModelCache, + prompt_cache: State, +) -> Result { 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) { diff --git a/crates/desktop/src-tauri/src/endpoints/sd.rs b/crates/desktop/src-tauri/src/endpoints/sd.rs index e47142ecde..b884e8fe5c 100644 --- a/crates/desktop/src-tauri/src/endpoints/sd.rs +++ b/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, 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()); diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 0a649cf8d5..14ed9d6ff3 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/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, diff --git a/crates/desktop/src-tauri/src/ml/mod.rs b/crates/desktop/src-tauri/src/ml/mod.rs index 712b6728e4..9fd2560f35 100644 --- a/crates/desktop/src-tauri/src/ml/mod.rs +++ b/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; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/prompt_cache.rs b/crates/desktop/src-tauri/src/ml/prompt_cache.rs new file mode 100644 index 0000000000..6816d8b798 --- /dev/null +++ b/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, +} + +impl PromptCache { + pub fn with_capacity(capacity: usize) -> anyhow::Result { + Ok(Self { + cache: ArcSieve::with_capacity(capacity)?, + }) + } + + pub fn get_copy(&self, key: &String) -> anyhow::Result> { + self.cache.get_copy(key) + } + + pub fn store_copy(&self, key: &String, value: &Tensor) -> anyhow::Result<()> { + self.cache.store_copy(key, value) + } +}