diff --git a/Cargo.lock b/Cargo.lock index 666ff764f..7dc762487 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -676,7 +676,9 @@ dependencies = [ "libloading 0.8.6", "log", "memory_caching", + "ndarray", "once_cell", + "ort", "primitives", "rand 0.8.5", "reqwest 0.12.12", @@ -7043,6 +7045,19 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ndarray" +version = "0.15.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "rawpointer", +] + [[package]] name = "ndk" version = "0.8.0" @@ -7825,6 +7840,26 @@ dependencies = [ "hashbrown 0.14.5", ] +[[package]] +name = "ort" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "889dca4c98efa21b1ba54ddb2bde44fd4920d910f492b618351f839d8428d79d" +dependencies = [ + "flate2", + "half", + "lazy_static", + "libc", + "ndarray", + "tar", + "thiserror 1.0.63", + "tracing", + "ureq", + "vswhom", + "winapi 0.3.9", + "zip 0.6.6", +] + [[package]] name = "overload" version = "0.1.1" @@ -12606,9 +12641,9 @@ dependencies = [ [[package]] name = "tracing-subscriber" -version = "0.3.19" +version = "0.3.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" dependencies = [ "nu-ansi-term", "sharded-slab", diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 9275426da..5da606b9e 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -58,6 +58,11 @@ tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } tokio-util = { version = "0.7", features = ["codec"] } half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] } +# NB: From `rust-background-removal` +ndarray = "0.15.6" +ort = { version = "1.16.2", features = ["cuda"] } +# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] } + [features] default = ["cuda"] cuda = [ diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 506874bd4..fa0db4e54 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -3,12 +3,11 @@ use crate::ml::prompt_cache::PromptCache; use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; use crate::state::app_config::AppConfig; use crate::state::app_dir::AppDataRoot; -use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; -use bytes::BytesMut; +use crate::utils::image::decode_base64_image::decode_base64_image; +use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png; use image::imageops::FilterType; -use image::{DynamicImage, ImageFormat, ImageReader}; +use image::{DynamicImage, RgbImage}; use log::{error, info}; -use std::io::Cursor; use std::path::PathBuf; use tauri::{AppHandle, State}; @@ -38,9 +37,12 @@ pub async fn infer_image( info!("Strength: {:?}; Prompt: {}", strength, prompt); - let image = hydrate_base64_image(image) + let image = decode_base64_image(image) .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; + // TODO(bt,2025-02-17): Running out of vram with full image buffer size + let image = image.resize(512, 512, FilterType::CatmullRom); + let result = do_infer_image( &prompt, image, @@ -52,11 +54,20 @@ pub async fn infer_image( &app_data_root ).await; - if let Err(err) = result.as_deref() { - error!("There was an error: {:?}", err); - } - - result + let image = match result { + Ok(image) => image, + Err(err) => { + error!("There was an error generating the image: {:?}", err); + return Err(format!("There was an error generating the image: {}", err)); + } + }; + + let bytes = encode_rgb_image_base64_png(image) + .map_err(|err| format!("failure to encode image: {:?}", err))?; + + info!("Inference successful; image converted to base64, serving back to browser..."); + + Ok(bytes) } async fn do_infer_image( @@ -68,7 +79,7 @@ async fn do_infer_image( prompt_cache: State<'_, PromptCache>, app: AppHandle, app_data_root: &AppDataRoot, -) -> Result { +) -> Result { let args = Args { image: &image, @@ -84,21 +95,10 @@ async fn do_infer_image( use_flash_attn: true, }; - match lcm_pipeline(args) { - Ok(image) => { - let mut bytes = Vec::with_capacity(1024*1024); - - image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png) - .map_err(|err| format!("failure to encode image: {:?}", err))?; + let image = lcm_pipeline(args) + .map_err(|err| format!("failure to encode image: {:?}", err))?; - let bytes = BASE64_STANDARD.encode(bytes); - - info!("Inference successful; image converted to base64, serving back to browser..."); - - Ok(bytes) - }, - Err(e) => Err(format!("Failed to generate image: {}", e)), - } + Ok(image) } fn get_prompt_or_fallback(user_prompt: Option) -> String { @@ -122,29 +122,3 @@ fn get_prompt_or_fallback(user_prompt: Option) -> String { .trim() .to_string() } - -fn hydrate_base64_image(base64_image: &str) -> anyhow::Result { - let bytes = BASE64_STANDARD.decode(base64_image)?; - - let image = ImageReader::new(Cursor::new(bytes)) - .with_guessed_format()? - .decode()?; - - // TODO(bt,2025-02-17): Running out of vram with full image buffer size - let image = image.resize(512, 512, FilterType::CatmullRom); - - Ok(image) -} - -fn float_to_u8(value: f64) -> u8 { - const MAX : f64 = u8::MAX as f64; - const MIN : f64 = 0f64; - let value = value.round(); - if value >= MAX { - u8::MAX - } else if value <= MIN { - 0 - } else { - value as u8 - } -} diff --git a/crates/desktop/src-tauri/src/endpoints/mod.rs b/crates/desktop/src-tauri/src/endpoints/mod.rs index fc4c6dac2..d365df449 100644 --- a/crates/desktop/src-tauri/src/endpoints/mod.rs +++ b/crates/desktop/src-tauri/src/endpoints/mod.rs @@ -1,4 +1,5 @@ pub mod download_models; pub mod flip_image; pub mod image_endpoint; +pub mod remove_background_endpoint; pub mod test_counter; diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs new file mode 100644 index 000000000..8d8b7aa5c --- /dev/null +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -0,0 +1,90 @@ +//! Adapted from https://github.com/dnanhkhoa/rust-background-removal + +use crate::ml::background_removal::onnx_session::onnx_session; +use crate::ml::model_cache::ModelCache; +use crate::state::app_dir::AppDataRoot; +use crate::utils::image::decode_base64_image::decode_base64_image; +use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; +use image::{imageops, DynamicImage}; +use log::error; +use ndarray::{Array, CowArray}; +use ort::Value; +use tauri::{AppHandle, State}; + +/// This handler removes the background from an image. +#[tauri::command] +pub async fn remove_background( + image: &str, + model_cache: State<'_, ModelCache>, + app_data_root: State<'_, AppDataRoot>, + app: AppHandle, +) -> Result { + + let image = decode_base64_image(image) + .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; + + // NB: Real time needed cuDNN 9, but this uses cuDNN 8: + // sudo apt-get install libcudnn8-dev libcudnn8 + // update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken + // update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode + let model_path = app_data_root.weights_dir().path().join("medium.onnx"); // TODO + let session = onnx_session(model_path.to_str().unwrap()) // TODO + .map_err(|err| format!("failure to start onnx session: {:?}", err))?; + + let image = process_dynamic_image(&session, image) + .map_err(|err| { + error!("Couldn't remove background from image: {}", err); + format!("Couldn't remove background from image: {}", err) + })?; + + let image = encode_dynamic_image_base64_png(image) + .map_err(|err| format!("failure to encode image: {:?}", err))?; + + Ok(image) +} + +fn process_dynamic_image( + session: &ort::Session, + dynamic_img: DynamicImage, +) -> Result { + let input_shape = session.inputs[0] + .dimensions() + .map(|dim| dim.unwrap()) + .collect::>(); + let input_img = dynamic_img.into_rgba8(); + let scaling_factor = f32::min( + 1., // Avoid upscaling + f32::min( + input_shape[3] as f32 / input_img.width() as f32, // Width ratio + input_shape[2] as f32 / input_img.height() as f32, // Height ratio + ), + ); + let mut resized_img = imageops::resize( + &input_img, + input_shape[3] as u32, + input_shape[2] as u32, + imageops::FilterType::Triangle, + ); + let input_tensor = CowArray::from( + Array::from_shape_fn(input_shape, |indices| { + let mean = 128.; + let std = 256.; + + (resized_img[(indices[3] as u32, indices[2] as u32)][indices[1]] as f32 - mean) / std + }) + .into_dyn(), + ); + let inputs = vec![Value::from_array(session.allocator(), &input_tensor)?]; + let outputs = session.run(inputs)?; + let output_tensor = outputs[0].try_extract::()?; + for (indices, alpha) in output_tensor.view().indexed_iter() { + resized_img[(indices[3] as u32, indices[2] as u32)][3] = (alpha * 255.) as u8; + } + let output_img = imageops::resize( + &resized_img, + (input_img.width() as f32 * scaling_factor) as u32, + (input_img.height() as f32 * scaling_factor) as u32, + imageops::FilterType::Triangle, + ); + Ok(DynamicImage::ImageRgba8(output_img)) +} \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 8fb77af74..a5f00e13e 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -4,9 +4,11 @@ pub mod ml; pub mod state; pub mod threads; pub mod transfer; +pub mod utils; use crate::endpoints::download_models::download_models; use crate::endpoints::image_endpoint::infer_image; +use crate::endpoints::remove_background_endpoint::remove_background; use crate::endpoints::test_counter::test_counter; use crate::ml::model_cache::ModelCache; use crate::ml::prompt_cache::PromptCache; @@ -60,9 +62,10 @@ pub fn run() { .manage(model_cache) .manage(app_data_root) .invoke_handler(tauri::generate_handler![ - test_counter, - infer_image, download_models, + infer_image, + remove_background, + test_counter, ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/crates/desktop/src-tauri/src/ml/background_removal/mod.rs b/crates/desktop/src-tauri/src/ml/background_removal/mod.rs new file mode 100644 index 000000000..2aed53ac3 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/background_removal/mod.rs @@ -0,0 +1 @@ +pub mod onnx_session; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs b/crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs new file mode 100644 index 000000000..9b87d0e82 --- /dev/null +++ b/crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs @@ -0,0 +1,18 @@ +//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs + +use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder}; + +pub fn onnx_session(onnx_model_file: &str) -> Result { + let environment = Environment::default().into_arc(); + let session = SessionBuilder::new(&environment)? + .with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level + .with_intra_threads(1)? // Configure the number of threads used for inference + .with_execution_providers([ + // Configure execution providers (e.g., CUDA, CoreML, CPU) + ExecutionProvider::CUDA(Default::default()), + ExecutionProvider::CoreML(Default::default()), + ExecutionProvider::CPU(Default::default()), + ])? + .with_model_from_file(onnx_model_file)?; + Ok(session) +} diff --git a/crates/desktop/src-tauri/src/ml/mod.rs b/crates/desktop/src-tauri/src/ml/mod.rs index 47fcf99b0..fb2f2e33e 100644 --- a/crates/desktop/src-tauri/src/ml/mod.rs +++ b/crates/desktop/src-tauri/src/ml/mod.rs @@ -1,3 +1,4 @@ +pub mod background_removal; pub mod image; pub mod model_cache; pub mod model_file; diff --git a/crates/desktop/src-tauri/src/utils/image/decode_base64_image.rs b/crates/desktop/src-tauri/src/utils/image/decode_base64_image.rs new file mode 100644 index 000000000..d8ebb2b14 --- /dev/null +++ b/crates/desktop/src-tauri/src/utils/image/decode_base64_image.rs @@ -0,0 +1,14 @@ +use base64::prelude::BASE64_STANDARD; +use base64::Engine; +use image::{DynamicImage, ImageReader}; +use std::io::Cursor; + +pub fn decode_base64_image(base64_image: &str) -> anyhow::Result { + let bytes = BASE64_STANDARD.decode(base64_image)?; + + let image = ImageReader::new(Cursor::new(bytes)) + .with_guessed_format()? + .decode()?; + + Ok(image) +} diff --git a/crates/desktop/src-tauri/src/utils/image/encode_dynamic_image_base64_png.rs b/crates/desktop/src-tauri/src/utils/image/encode_dynamic_image_base64_png.rs new file mode 100644 index 000000000..cb68e228a --- /dev/null +++ b/crates/desktop/src-tauri/src/utils/image/encode_dynamic_image_base64_png.rs @@ -0,0 +1,14 @@ +use base64::prelude::BASE64_STANDARD; +use base64::Engine; +use image::{DynamicImage, ImageFormat}; +use std::io::Cursor; + +pub fn encode_dynamic_image_base64_png(image: DynamicImage) -> anyhow::Result { + let mut bytes = Vec::with_capacity(1024*1024); + + image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?; + + let bytes = BASE64_STANDARD.encode(bytes); + + Ok(bytes) +} diff --git a/crates/desktop/src-tauri/src/utils/image/encode_rgb_image_base64_png.rs b/crates/desktop/src-tauri/src/utils/image/encode_rgb_image_base64_png.rs new file mode 100644 index 000000000..defb66f7c --- /dev/null +++ b/crates/desktop/src-tauri/src/utils/image/encode_rgb_image_base64_png.rs @@ -0,0 +1,17 @@ +use crate::ml::image::tensor_to_image_buffer::RgbImage; +use base64::prelude::BASE64_STANDARD; +use base64::Engine; +use image::ImageFormat; +use std::io::Cursor; + +pub fn encode_rgb_image_base64_png( + image: RgbImage +) -> anyhow::Result { + let mut bytes = Vec::with_capacity(1024*1024); + + image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?; + + let bytes = BASE64_STANDARD.encode(bytes); + + Ok(bytes) +} diff --git a/crates/desktop/src-tauri/src/utils/image/mod.rs b/crates/desktop/src-tauri/src/utils/image/mod.rs new file mode 100644 index 000000000..3b7d7a6b7 --- /dev/null +++ b/crates/desktop/src-tauri/src/utils/image/mod.rs @@ -0,0 +1,3 @@ +pub mod decode_base64_image; +pub mod encode_dynamic_image_base64_png; +pub mod encode_rgb_image_base64_png; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/utils/mod.rs b/crates/desktop/src-tauri/src/utils/mod.rs new file mode 100644 index 000000000..b1b6b4f00 --- /dev/null +++ b/crates/desktop/src-tauri/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod image; \ No newline at end of file