From 29b33cd3270e5cf8870fb9ddcb206db8298646c5 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 25 Mar 2025 22:43:25 -0400 Subject: [PATCH] save image handler --- Cargo.lock | 14 +++++-- crates/desktop/src-tauri/Cargo.toml | 4 +- crates/desktop/src-tauri/src/endpoints/mod.rs | 1 + .../src/endpoints/save_image_endpoint.rs | 38 +++++++++++++++++++ crates/desktop/src-tauri/src/lib.rs | 6 ++- crates/desktop/src-tauri/src/state/app_dir.rs | 13 +++++++ 6 files changed, 70 insertions(+), 6 deletions(-) create mode 100644 crates/desktop/src-tauri/src/endpoints/save_image_endpoint.rs diff --git a/Cargo.lock b/Cargo.lock index 22bbfe0b9b..03de899227 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -661,7 +661,9 @@ dependencies = [ "assertor", "base64 0.22.1", "bytes 1.10.0", + "chrono", "const_format", + "crockford", "cudarc", "directories 6.0.0", "dirs 6.0.0", @@ -1974,9 +1976,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chrono" -version = "0.4.39" +version = "0.4.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e36cc9d416881d2e24f9a963be5fb1cd90966419ac844274161d10488b3e825" +checksum = "1a7964611d71df112cb1730f2ee67324fcf4d0fc6606acbbe9bfe06df124637c" dependencies = [ "android-tzdata", "iana-time-zone", @@ -1984,7 +1986,7 @@ dependencies = [ "num-traits", "serde", "wasm-bindgen", - "windows-targets 0.52.6", + "windows-link", ] [[package]] @@ -13876,6 +13878,12 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "windows-link" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76840935b766e1b0a05c0066835fb9ec80071d4c09a16f6bd5f7e655e3c14c38" + [[package]] name = "windows-registry" version = "0.2.0" diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 6ffaa00c3a..cb7d638d41 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -21,6 +21,7 @@ tauri-build = { version = "2.0.4", features = [] } [dependencies] # Internal +crockford = { path = "../../lib/data/crockford" } filesys = { path = "../../lib/files/filesys" } memory_caching = { path = "../../lib/caching/memory_caching" } ml_weights_registry = { path = "../ml_weights_registry" } @@ -34,9 +35,11 @@ cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand" # External #openssl = { version = "0.10", features = ["vendored"] } +# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls anyhow = "1.0" base64 = "0.22.1" bytes = "1.10" +chrono = "0.4.40" const_format = "0.2.34" directories = "6.0.0" dirs = "6.0.0" @@ -47,7 +50,6 @@ log = "0.4" once_cell = "1.20.3" ort = { version = "1.16.2", features = ["cuda", "load-dynamic"] } rand = "0.8.5" -# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls reqwest = { version = "0.12.12", features = ["stream"], default-features = false } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" diff --git a/crates/desktop/src-tauri/src/endpoints/mod.rs b/crates/desktop/src-tauri/src/endpoints/mod.rs index 3a293cea38..dcf8a0cc98 100644 --- a/crates/desktop/src-tauri/src/endpoints/mod.rs +++ b/crates/desktop/src-tauri/src/endpoints/mod.rs @@ -3,4 +3,5 @@ pub mod flip_image; pub mod inpaint_image_endpoint; pub mod realtime_image_endpoint; pub mod remove_background_endpoint; +pub mod save_image_endpoint; pub mod text_to_image_endpoint; diff --git a/crates/desktop/src-tauri/src/endpoints/save_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/save_image_endpoint.rs new file mode 100644 index 0000000000..39312b7d58 --- /dev/null +++ b/crates/desktop/src-tauri/src/endpoints/save_image_endpoint.rs @@ -0,0 +1,38 @@ +use crate::state::app_config::AppConfig; +use crate::state::app_dir::AppDataRoot; +use crate::stubs::model_cache::ModelCache; +use crate::stubs::prompt_cache::PromptCache; +use crate::utils::image::decode_base64_image::decode_base64_image; +use chrono::Local; +use crockford::crockford_entropy_lower; +use image::{ImageFormat, RgbImage}; +use log::info; +use tauri::{AppHandle, State}; + +const FILENAME_LENGTH : usize = 10; + +/// This handler takes an image (as a base64 encoded string) and saves it. +#[tauri::command] +pub async fn save_image( + image: &str, + app_data_root: State<'_, AppDataRoot>, +) -> Result { + info!("save_image endpoint called."); + + let image = decode_base64_image(image) + .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; + + let assets_dir = app_data_root.assets_dir(); + + let date_save_dir = assets_dir.make_or_get_current_date_dir() + .map_err(|err| format!("Couldn't create assets save directory: {}", err))?; + + let entropy = crockford_entropy_lower(FILENAME_LENGTH); // TODO: These are ugly. A monotonic counter would be better. + let filename = format!("{}.png", entropy); + let filename = date_save_dir.join(filename); + + image.save_with_format(filename, ImageFormat::Png) + .map_err(|err| format!("Couldn't save image: {}", err))?; + + Ok("".to_string()) +} diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 5252c940bd..e8f7557102 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -7,10 +7,11 @@ pub mod transfer; pub mod utils; use crate::endpoints::download_models::download_models; -use crate::endpoints::inpaint_image_endpoint::inpaint_image; use crate::endpoints::flip_image::flip_image; +use crate::endpoints::inpaint_image_endpoint::inpaint_image; use crate::endpoints::realtime_image_endpoint::infer_image; use crate::endpoints::remove_background_endpoint::remove_background; +use crate::endpoints::save_image_endpoint::save_image; use crate::endpoints::text_to_image_endpoint::text_to_image; use crate::state::app_config::AppConfig; use crate::stubs::model_cache::ModelCache; @@ -75,8 +76,9 @@ pub fn run() { flip_image, infer_image, inpaint_image, - text_to_image, remove_background, + save_image, + text_to_image, text_to_image, ]) .run(tauri::generate_context!()) diff --git a/crates/desktop/src-tauri/src/state/app_dir.rs b/crates/desktop/src-tauri/src/state/app_dir.rs index 4fd31f0961..2cbad5d4d5 100644 --- a/crates/desktop/src-tauri/src/state/app_dir.rs +++ b/crates/desktop/src-tauri/src/state/app_dir.rs @@ -4,6 +4,7 @@ use anyhow::anyhow; use directories::UserDirs; use ml_weights_registry::weights_registry::weight_descriptor::WeightDescriptor; use std::path::{Path, PathBuf}; +use chrono::{DateTime, Local}; use tempdir::TempDir; use tempfile::{Builder, NamedTempFile}; @@ -137,6 +138,18 @@ impl AppAssetsDir { pub fn path(&self) -> &Path { &self.path } + + /// Return the current date asset directory + /// If it doesn't exist, create it. + pub fn make_or_get_current_date_dir(&self) -> anyhow::Result { + let format: DateTime = Local::now(); + let date_directory_name = format.format("%Y-%d-%m").to_string(); + let full_path = self.path.join(date_directory_name); + if !full_path.exists() { + std::fs::create_dir(&full_path)?; + } + Ok(full_path) + } } impl AppWeightsDir {