From 990b7ff3e5199054e8a57514e22eec8c4a4071e3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Mar 2025 20:18:57 -0400 Subject: [PATCH 1/5] package split for mac (part one) (#459) --- Cargo.lock | 53 +++++- crates/desktop/ml_models/Cargo.toml | 78 +++++++++ crates/desktop/ml_models/src/lib.rs | 1 + .../src/ml/background_removal/mod.rs | 2 + .../src/ml/background_removal/onnx_session.rs | 3 +- .../remove_image_background.rs | 61 +++++++ .../src/ml/image/dynamic_image_to_tensor.rs | 0 .../src/ml/image/load_image_file_to_tensor.rs | 0 .../ml/image/load_image_file_to_tensor_2.rs | 0 .../src/ml/image/mod.rs | 0 .../src/ml/image/save_image_from_tensor.rs | 0 .../src/ml/image/tensor_to_image_buffer.rs | 0 .../{src-tauri => ml_models}/src/ml/mod.rs | 1 + .../src/ml/model_cache.rs | 0 .../desktop/ml_models/src/ml/model_config.rs | 61 +++++++ .../src/ml/model_file.rs | 0 .../src/ml/prompt_cache.rs | 0 .../src/ml/stable_diffusion/get_vae_scale.rs | 0 .../infer_clip_text_embeddings.rs | 69 +++++++- .../src/ml/stable_diffusion/lcm_pipeline.rs | 157 ++++++++++-------- .../src/ml/stable_diffusion/mod.rs | 0 .../remap_lcm_strength_range.rs | 0 .../stable_diffusion_pipeline.rs | 111 +++++++------ .../src/ml/weights_registry/mod.rs | 0 .../ml/weights_registry/weight_descriptor.rs | 0 .../weight_descriptor_builder.rs | 0 .../src/ml/weights_registry/weights.rs | 0 crates/desktop/src-tauri/Cargo.toml | 25 +-- .../src/endpoints/realtime_image_endpoint.rs | 28 +++- .../endpoints/remove_background_endpoint.rs | 71 +------- .../src/endpoints/text_to_image_endpoint.rs | 25 ++- .../src/events/notification_event.rs | 1 - crates/desktop/src-tauri/src/lib.rs | 5 +- .../src/ml/background_removal/mod.rs | 1 - .../desktop/src-tauri/src/state/app_config.rs | 55 +----- crates/desktop/src-tauri/src/state/app_dir.rs | 2 +- .../src/threads/downloader_thread.rs | 4 +- .../image/encode_rgb_image_base64_png.rs | 3 +- .../src/utils/log_environment_details.rs | 14 +- crates/desktop/src-tauri/tauri.conf.json | 5 +- 40 files changed, 527 insertions(+), 309 deletions(-) create mode 100644 crates/desktop/ml_models/Cargo.toml create mode 100644 crates/desktop/ml_models/src/lib.rs create mode 100644 crates/desktop/ml_models/src/ml/background_removal/mod.rs rename crates/desktop/{src-tauri => ml_models}/src/ml/background_removal/onnx_session.rs (88%) create mode 100644 crates/desktop/ml_models/src/ml/background_removal/remove_image_background.rs rename crates/desktop/{src-tauri => ml_models}/src/ml/image/dynamic_image_to_tensor.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/image/load_image_file_to_tensor.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/image/load_image_file_to_tensor_2.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/image/mod.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/image/save_image_from_tensor.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/image/tensor_to_image_buffer.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/mod.rs (87%) rename crates/desktop/{src-tauri => ml_models}/src/ml/model_cache.rs (100%) create mode 100644 crates/desktop/ml_models/src/ml/model_config.rs rename crates/desktop/{src-tauri => ml_models}/src/ml/model_file.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/prompt_cache.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/stable_diffusion/get_vae_scale.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/stable_diffusion/infer_clip_text_embeddings.rs (68%) rename crates/desktop/{src-tauri => ml_models}/src/ml/stable_diffusion/lcm_pipeline.rs (67%) rename crates/desktop/{src-tauri => ml_models}/src/ml/stable_diffusion/mod.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/stable_diffusion/remap_lcm_strength_range.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/stable_diffusion/stable_diffusion_pipeline.rs (71%) rename crates/desktop/{src-tauri => ml_models}/src/ml/weights_registry/mod.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/weights_registry/weight_descriptor.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/weights_registry/weight_descriptor_builder.rs (100%) rename crates/desktop/{src-tauri => ml_models}/src/ml/weights_registry/weights.rs (100%) delete mode 100644 crates/desktop/src-tauri/src/ml/background_removal/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 4c877cc071..190ad0c1fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -661,25 +661,18 @@ dependencies = [ "assertor", "base64 0.22.1", "bytes 1.10.0", - "candle-core", - "candle-nn", - "candle-transformers", "const_format", - "cudarc", "directories 6.0.0", "dirs 6.0.0", "filesys", "futures 0.3.29", "half", - "hf-hub 0.3.2", "image", "libloading 0.8.6", "log", "memory_caching", - "ndarray", - "nvidia-checker", + "ml_models", "once_cell", - "ort", "primitives", "rand 0.8.5", "reqwest 0.12.12", @@ -692,7 +685,6 @@ dependencies = [ "tauri-plugin-log", "tempdir", "tempfile", - "tokenizers 0.15.2", "tokio", "tokio-util", ] @@ -6804,6 +6796,49 @@ dependencies = [ "ws2_32-sys", ] +[[package]] +name = "ml_models" +version = "0.0.1" +dependencies = [ + "anyhow", + "assertor", + "base64 0.22.1", + "bytes 1.10.0", + "candle-core", + "candle-nn", + "candle-transformers", + "const_format", + "cudarc", + "directories 6.0.0", + "dirs 6.0.0", + "filesys", + "futures 0.3.29", + "half", + "hf-hub 0.3.2", + "image", + "libloading 0.8.6", + "log", + "memory_caching", + "ndarray", + "nvidia-checker", + "once_cell", + "ort", + "primitives", + "rand 0.8.5", + "reqwest 0.12.12", + "serde", + "serde_json", + "serde_yaml", + "serial_test", + "tauri", + "tauri-plugin-log", + "tempdir", + "tempfile", + "tokenizers 0.15.2", + "tokio", + "tokio-util", +] + [[package]] name = "mockall" version = "0.12.1" diff --git a/crates/desktop/ml_models/Cargo.toml b/crates/desktop/ml_models/Cargo.toml new file mode 100644 index 0000000000..dc8e074868 --- /dev/null +++ b/crates/desktop/ml_models/Cargo.toml @@ -0,0 +1,78 @@ +[package] +name = "ml_models" +edition = "2021" +version = "0.0.1" +authors = ["storyteller"] +publish = false +description = "" +license = "" +repository = "" + +[lib] +name = "ml_models" +path = "src/lib.rs" +crate-type = ["staticlib", "cdylib", "rlib"] + +[package.metadata.cargo-udeps.ignore] +normal = ["my-workspace-hack"] + +[dependencies] + +# Internal +filesys = { path = "../../lib/files/filesys" } +memory_caching = { path = "../../lib/caching/memory_caching" } +primitives = { path = "../../lib/data/primitives" } + +# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332 +cudarc = "0.13.7" + +# External +anyhow = "1.0" +base64 = "0.22.1" +bytes = "1.10" +candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] } +candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] } +candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] } +const_format = "0.2.34" +directories = "6.0.0" +dirs = "6.0.0" +futures = "0.3" +half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] } +hf-hub = "0.3.2" +image = "0.25.5" +log = "0.4" +ndarray = "0.15.6" +nvidia-checker = "0.1.2" +once_cell = "1.20.3" +ort = { version = "1.16.2", features = ["cuda"] } +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" +serde_yaml = "0.9" +tauri = { version = "2.2.4", features = [] } +tauri-plugin-log = "2" +tempdir = "0.3.7" +tempfile = "3.18.0" +tokenizers = "0.15.2" +tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } +tokio-util = { version = "0.7", features = ["codec"] } + +# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] } + +[features] +default = ["cuda"] +cuda = [ + "candle-core/cuda", + "candle-nn/cuda", + "candle-transformers/cuda", + # NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues + "candle-transformers/flash-attn", + "candle-core/cudnn" +] + +[dev-dependencies] +libloading = "0.8.6" +serial_test = "3.2.0" +assertor = "0.0.3" diff --git a/crates/desktop/ml_models/src/lib.rs b/crates/desktop/ml_models/src/lib.rs new file mode 100644 index 0000000000..bf0fcba93b --- /dev/null +++ b/crates/desktop/ml_models/src/lib.rs @@ -0,0 +1 @@ +pub mod ml; diff --git a/crates/desktop/ml_models/src/ml/background_removal/mod.rs b/crates/desktop/ml_models/src/ml/background_removal/mod.rs new file mode 100644 index 0000000000..c6e8e6eeb1 --- /dev/null +++ b/crates/desktop/ml_models/src/ml/background_removal/mod.rs @@ -0,0 +1,2 @@ +pub mod onnx_session; +pub mod remove_image_background; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs b/crates/desktop/ml_models/src/ml/background_removal/onnx_session.rs similarity index 88% rename from crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs rename to crates/desktop/ml_models/src/ml/background_removal/onnx_session.rs index 9b87d0e823..e365b7eb13 100644 --- a/crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs +++ b/crates/desktop/ml_models/src/ml/background_removal/onnx_session.rs @@ -1,8 +1,9 @@ //! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs +use std::path::Path; use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder}; -pub fn onnx_session(onnx_model_file: &str) -> Result { +pub fn onnx_session(onnx_model_file: &Path) -> Result { let environment = Environment::default().into_arc(); let session = SessionBuilder::new(&environment)? .with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level diff --git a/crates/desktop/ml_models/src/ml/background_removal/remove_image_background.rs b/crates/desktop/ml_models/src/ml/background_removal/remove_image_background.rs new file mode 100644 index 0000000000..808a2e90d5 --- /dev/null +++ b/crates/desktop/ml_models/src/ml/background_removal/remove_image_background.rs @@ -0,0 +1,61 @@ +use crate::ml::background_removal::onnx_session::onnx_session; +use image::{imageops, DynamicImage}; +use ndarray::{Array, CowArray}; +use ort::Value; +use std::path::Path; + +pub async fn remove_image_background>(model_path: P, image: DynamicImage) -> anyhow::Result { + // 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 session = onnx_session(model_path.as_ref())?; + let image = process_dynamic_image(&session, image)?; + 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)) +} diff --git a/crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs b/crates/desktop/ml_models/src/ml/image/dynamic_image_to_tensor.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs rename to crates/desktop/ml_models/src/ml/image/dynamic_image_to_tensor.rs diff --git a/crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor.rs b/crates/desktop/ml_models/src/ml/image/load_image_file_to_tensor.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor.rs rename to crates/desktop/ml_models/src/ml/image/load_image_file_to_tensor.rs diff --git a/crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs b/crates/desktop/ml_models/src/ml/image/load_image_file_to_tensor_2.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs rename to crates/desktop/ml_models/src/ml/image/load_image_file_to_tensor_2.rs diff --git a/crates/desktop/src-tauri/src/ml/image/mod.rs b/crates/desktop/ml_models/src/ml/image/mod.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/image/mod.rs rename to crates/desktop/ml_models/src/ml/image/mod.rs diff --git a/crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs b/crates/desktop/ml_models/src/ml/image/save_image_from_tensor.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs rename to crates/desktop/ml_models/src/ml/image/save_image_from_tensor.rs diff --git a/crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs b/crates/desktop/ml_models/src/ml/image/tensor_to_image_buffer.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs rename to crates/desktop/ml_models/src/ml/image/tensor_to_image_buffer.rs diff --git a/crates/desktop/src-tauri/src/ml/mod.rs b/crates/desktop/ml_models/src/ml/mod.rs similarity index 87% rename from crates/desktop/src-tauri/src/ml/mod.rs rename to crates/desktop/ml_models/src/ml/mod.rs index fb2f2e33e0..e9c0a3f46a 100644 --- a/crates/desktop/src-tauri/src/ml/mod.rs +++ b/crates/desktop/ml_models/src/ml/mod.rs @@ -1,6 +1,7 @@ pub mod background_removal; pub mod image; pub mod model_cache; +pub mod model_config; pub mod model_file; pub mod prompt_cache; pub mod stable_diffusion; diff --git a/crates/desktop/src-tauri/src/ml/model_cache.rs b/crates/desktop/ml_models/src/ml/model_cache.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/model_cache.rs rename to crates/desktop/ml_models/src/ml/model_cache.rs diff --git a/crates/desktop/ml_models/src/ml/model_config.rs b/crates/desktop/ml_models/src/ml/model_config.rs new file mode 100644 index 0000000000..a7467fd319 --- /dev/null +++ b/crates/desktop/ml_models/src/ml/model_config.rs @@ -0,0 +1,61 @@ +use crate::ml::model_file::StableDiffusionVersion; +use candle_core::{DType, Device}; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; +use hf_hub::api::sync::Api; +use log::info; + +const DEFAULT_SD_IMAGE_WIDTH: usize = 1024; +const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024; +const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32; +const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5; + +pub struct ModelConfig { + + /// We only support one device type currently, and we can't multiplex + /// over several devices yet. + pub device: Device, + + /// The core numeric type to use for models. + pub dtype: DType, + + pub sd_version: StableDiffusionVersion, + pub sd_config: StableDiffusionConfig, + + /// Probably shouldn't be here. + pub hf_api: Api, +} + +impl ModelConfig { + pub fn init() -> anyhow::Result { + let device = Device::new_cuda(0) + .unwrap_or_else(|e| { + info!("CUDA not available ({}), falling back to CPU", e); + Device::Cpu + }); + + let sd_config = match SD_VERSION { + StableDiffusionVersion::Turbo => { + StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) + } + StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => { + StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) + } + StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => { + StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) + } + StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => { + StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) + } + }; + + let hf_api = Api::new()?; + + Ok(Self { + device, + dtype: DEFAULT_SD_NUMERIC_DATATYPE, + sd_version: SD_VERSION, + sd_config, + hf_api, + }) + } +} \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/model_file.rs b/crates/desktop/ml_models/src/ml/model_file.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/model_file.rs rename to crates/desktop/ml_models/src/ml/model_file.rs diff --git a/crates/desktop/src-tauri/src/ml/prompt_cache.rs b/crates/desktop/ml_models/src/ml/prompt_cache.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/prompt_cache.rs rename to crates/desktop/ml_models/src/ml/prompt_cache.rs diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/get_vae_scale.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/get_vae_scale.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/stable_diffusion/get_vae_scale.rs rename to crates/desktop/ml_models/src/ml/stable_diffusion/get_vae_scale.rs diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs similarity index 68% rename from crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs rename to crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs index 8f682ad8d0..4c1591bc44 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs @@ -4,18 +4,50 @@ extern crate accelerate_src; extern crate intel_mkl_src; use candle_transformers::models::stable_diffusion; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use crate::ml::model_file::{ModelFile, StableDiffusionVersion}; use crate::ml::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2}; -use crate::state::app_dir::AppWeightsDir; use anyhow::Error as E; use candle_core::{DType, Device, Module, Tensor, D}; use log::info; use tokenizers::Tokenizer; + +pub struct ClipArgs<'a, P: AsRef, Q: AsRef> { + pub prompt: &'a str, + pub uncond_prompt: &'a str, + pub tokenizer: Option, + pub clip_weights: Option, + pub clip2_weights: Option, + pub sd_version: StableDiffusionVersion, + pub sd_config: &'a stable_diffusion::StableDiffusionConfig, + pub use_f16: bool, + pub device: &'a Device, + pub dtype: DType, + pub use_guide_scale: bool, + pub clip_json_path: P, + pub clip_weights_path: Q, +} + #[allow(clippy::too_many_arguments)] -pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option, clip_weights: Option, clip2_weights: Option, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, weights_dir: &AppWeightsDir) -> anyhow::Result { +pub fn infer_clip_text_embeddings, Q: AsRef>(args: ClipArgs<'_, P, Q>) -> anyhow::Result { + let ClipArgs { + prompt, + uncond_prompt, + tokenizer, + clip_weights, + clip2_weights, + sd_version, + sd_config, + use_f16, + device, + dtype, + use_guide_scale, + clip_json_path, + clip_weights_path, + } = args; + info!("Preparing text embeddings... for {:?}", sd_version); let which = match sd_version { @@ -39,7 +71,8 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: dtype, use_guide_scale, *first, - weights_dir, + clip_json_path.as_ref(), + clip_weights_path.as_ref(), ) }) .collect::>>()?; @@ -50,10 +83,25 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: } #[allow(clippy::too_many_arguments)] -fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option, clip_weights: Option, clip2_weights: Option, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, first: bool, weights_dir: &AppWeightsDir) -> anyhow::Result { +fn do_infer_clip_text_embeddings( + prompt: &str, + uncond_prompt: &str, + tokenizer: Option, + clip_weights: Option, + clip2_weights: Option, + sd_version: StableDiffusionVersion, + sd_config: &stable_diffusion::StableDiffusionConfig, + use_f16: bool, + device: &Device, + dtype: DType, + use_guide_scale: bool, + first: bool, + clip_json_path: &Path, + clip_weights_path: &Path, +) -> anyhow::Result { info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first)); - let tokenizer = weights_dir.weight_path(&CLIP_JSON); + let tokenizer = clip_json_path; info!("Loading Clip Tokenizer path: {:?}", tokenizer); @@ -75,11 +123,16 @@ fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: O let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?; println!("Building the Clip transformer."); + + // TODO(bt,2025-03-23): These used to be two different clip models. We should only need one. let clip_weights = if first { - weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER) + println!(">>>>> CLIP 1"); + clip_weights_path } else { - weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER_2) + println!(">>>>> CLIP 2"); + clip_weights_path }; + let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() }; let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?; diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs similarity index 67% rename from crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs rename to crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs index dfed61b3f6..515717b9b1 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/lcm_pipeline.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs @@ -1,53 +1,80 @@ -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; -use crate::events::notification_event::{NotificationEvent, NotificationModelType}; use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor; use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage}; use crate::ml::model_cache::ModelCache; +use crate::ml::model_config::ModelConfig; use crate::ml::model_file::{ModelFile, StableDiffusionVersion}; use crate::ml::prompt_cache::PromptCache; use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale; -use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings; +use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs}; use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range; use crate::ml::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; -use crate::state::app_config::AppConfig; -use crate::state::app_dir::AppDataRoot; use anyhow::{anyhow, Error as E, Result}; -use candle_core::{DType, IndexOp, Tensor, D}; +use candle_core::{DType, Device, IndexOp, Tensor, D}; use candle_nn::VarBuilder; -use candle_transformers::models::stable_diffusion::lcm::LCMScheduler; +use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig}; use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig; use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel; use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig; use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution}; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; use image::DynamicImage; use log::info; use rand::Rng; -use tauri::{AppHandle, Emitter}; const DEFAULT_STRENGTH : f64 = 75.0; -pub struct Args<'a> { +pub struct Args<'a, P1: AsRef, P2: AsRef, P3: AsRef, P4: AsRef> { pub image: &'a DynamicImage, pub prompt: String, pub uncond_prompt: String, pub cfg_scale: Option, - pub i2i_strength: Option, - pub configs: &'a AppConfig, + pub img2img_strength: Option, pub model_cache: &'a ModelCache, pub prompt_cache: &'a PromptCache, - pub app_data_root: &'a AppDataRoot, - pub app: &'a AppHandle, pub use_flash_attn: bool, + pub model_config: &'a ModelConfig, + pub maybe_seed: Option, + pub scheduler_steps: usize, + pub vae_path: P1, + pub unet_path: P2, + pub clip_json_path: P3, + pub clip_weights_path: P4, } -pub fn lcm_pipeline(args: Args<'_>) -> Result { - let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, configs, model_cache, prompt_cache, app, image, app_data_root, use_flash_attn } = args; - - let weights_dir = app_data_root.weights_dir(); +pub fn lcm_pipeline(args: Args<'_, P1, P2, P3, P4>) -> Result + where P1: AsRef, P2: AsRef, P3: AsRef, P4: AsRef +{ + let Args { + prompt, + uncond_prompt, + cfg_scale, + img2img_strength, + model_cache, + prompt_cache, + image, + use_flash_attn , + model_config: ModelConfig { + device, + dtype, + sd_version, + sd_config, + hf_api, + }, + maybe_seed, + scheduler_steps, + vae_path, + unet_path, + clip_json_path, + clip_weights_path, + } = args; - let img2img_strength = args.i2i_strength + let dtype = *dtype; + let sd_version = *sd_version; + + let img2img_strength = img2img_strength .map(|s| remap_lcm_strength_range(s)) .unwrap_or(DEFAULT_STRENGTH); @@ -56,14 +83,17 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { let use_f16 = true; // Use LCM Scheduler instead of Euler Ancestral for better speed and quality - let mut scheduler = LCMScheduler::new(configs.scheduler_steps, img2img_strength, candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default())?; + let mut scheduler = LCMScheduler::new(scheduler_steps, img2img_strength, LCMSchedulerConfig::default())?; - let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen()); - configs.device.set_seed(seed)?; + let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen()); + + let seed = 42; + + device.set_seed(seed)?; let guidance_scale = match cfg_scale { Some(guidance_scale) => guidance_scale, - None => match configs.sd_version { + None => match sd_version { StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::Xl => 7.5, StableDiffusionVersion::Turbo => 0., _ => 0., // NB: Not sure what the other model families should use, so sticking with "0" @@ -78,50 +108,45 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { let mut text_embeddings = if let Some(tensor) = maybe_cached { tensor } else { - info!("Loading clip weights..."); - - let clip_weights = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16); - let clip_weights = Some(clip_weights.to_string_lossy().to_string()); - info!("Prompt is NOT cached! Calculating embedding..."); - let tensor = infer_clip_text_embeddings( - &prompt, - &uncond_prompt, - None, // tokenizer - clip_weights, - None, // clip2_weights - configs.sd_version, - &configs.sd_config, + let tensor = infer_clip_text_embeddings(ClipArgs { + prompt: & prompt, + uncond_prompt: &uncond_prompt, + tokenizer: None, // tokenizer + clip_weights: None, + clip2_weights: None, // clip2_weights + sd_version, + sd_config: &sd_config, use_f16, // use_f16 - &configs.device, - configs.dtype, + device, + dtype, use_guide_scale, - weights_dir, - )?; + clip_json_path, + clip_weights_path, + })?; + prompt_cache.store_copy(&prompt, &tensor)?; tensor }; - println!("Text embeddings shape: {:?}", text_embeddings.shape()); + info!("Text embeddings shape: {:?}", text_embeddings.shape()); - println!("Loading input image into tensor..."); + info!("Loading input image into tensor..."); - let input_image = dynamic_image_to_tensor(image, &configs.device, configs.dtype)?; + let input_image = dynamic_image_to_tensor(image, &device, dtype)?; - println!("Reference image shape: {:?}", input_image.shape()); + info!("Reference image shape: {:?}", input_image.shape()); - let vae_scale = get_vae_scale(configs.sd_version); + let vae_scale = get_vae_scale(sd_version); let maybe_vae = model_cache.get_vae()?; let vae = match maybe_vae { Some(vae) => vae, None => { - let vae_file = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE); - - println!("Building VAE model from file {:?}...", &vae_file); - - let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.dtype)?; + info!("Building VAE model from file {:?}...", vae_path.as_ref()); + + let vae = sd_config.build_vae(vae_path, &device, dtype)?; let vae = Arc::new(vae); @@ -138,9 +163,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { None => { info!("No unet found in cache; loading..."); - let unet_weights = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET); - - let in_channels = match configs.sd_version { + let in_channels = match sd_version { StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9, _ => 4, }; @@ -175,7 +198,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { }; // Load the model directly from the safetensors file - let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], configs.dtype, &configs.device)? }; + let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_path], dtype, &device)? }; let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?; @@ -191,19 +214,19 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { // Create LCM guidance scale embeddings let embedding_dim = 1280; - let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &configs.device, configs.dtype)?; + let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?; // Generate latents from input image using the LCM approach - println!("Generating latents from input image..."); + info!("Generating latents from input image..."); let init_latents = (init_latent_dist.sample()? * vae_scale)?; - println!("Initial latents shape: {:?}", init_latents.shape()); + info!("Initial latents shape: {:?}", init_latents.shape()); // Calculate timesteps for LCM with proper img2img handling - let timesteps = LCMScheduler::get_timesteps_for_steps(configs.scheduler_steps, img2img_strength); + let timesteps = LCMScheduler::get_timesteps_for_steps(scheduler_steps, img2img_strength); let t_start = timesteps[0]; // First denoising step // Add noise at the appropriate timestep for img2img - println!("Adding noise to latents for t_start={}", t_start); + info!("Adding noise to latents for t_start={}", t_start); let noise = init_latents.randn_like(0f64, 1f64)?; let latents = scheduler.add_noise(&init_latents, noise, t_start)?; @@ -211,7 +234,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect(); for (timestep_index, ×tep) in timesteps.iter().enumerate() { - println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep); + info!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep); let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() }; @@ -221,7 +244,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, Some(&guidance_scale_embedding)) { Ok(pred) => pred, Err(e) => { - println!("UNet inference failed with error: {}", e); + info!("UNet inference failed with error: {}", e); return Err(anyhow::anyhow!("UNet inference failed: {}", e)); }, }; @@ -233,24 +256,24 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result { } // Apply the LCM scheduler step - latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?; + latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, scheduler_steps)?; } - println!("Diffusion process completed, decoding image..."); + info!("Diffusion process completed, decoding image..."); let image = vae.decode(&(latents / vae_scale)?)?; - println!("VAE decode completed"); + info!("VAE decode completed"); - println!("Scaling image back"); + info!("Scaling image back"); let image = ((image / 2.)? + 0.5)?; - println!("Normalized image values"); + info!("Normalized image values"); let image = (image.clamp(0f32, 1.)? * 255.)?; - println!("Convert to int8"); + info!("Convert to int8"); let image = image.to_dtype(DType::U8)?; - println!("Converted to 8-bit format"); + info!("Converted to 8-bit format"); let image = tensor_to_image_buffer(&image.i(0)?)?; diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/mod.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/stable_diffusion/mod.rs rename to crates/desktop/ml_models/src/ml/stable_diffusion/mod.rs diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/remap_lcm_strength_range.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/remap_lcm_strength_range.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/stable_diffusion/remap_lcm_strength_range.rs rename to crates/desktop/ml_models/src/ml/stable_diffusion/remap_lcm_strength_range.rs diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/stable_diffusion_pipeline.rs similarity index 71% rename from crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs rename to crates/desktop/ml_models/src/ml/stable_diffusion/stable_diffusion_pipeline.rs index 5a833870b7..7bc2da1d1b 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/stable_diffusion_pipeline.rs @@ -1,75 +1,86 @@ +use std::path::Path; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; -use crate::events::notification_event::{NotificationEvent, NotificationModelType}; use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor; use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage}; use crate::ml::model_cache::ModelCache; use crate::ml::model_file::StableDiffusionVersion; use crate::ml::prompt_cache::PromptCache; use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale; -use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings; -use crate::ml::weights_registry::weights::{SDXL_TURBO_UNET, SDXL_TURBO_VAE}; -use crate::state::app_config::AppConfig; -use crate::state::app_dir::AppDataRoot; +use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs}; use anyhow::{anyhow, Error as E, Result}; -use candle_core::{DType, IndexOp, Tensor, D}; +use candle_core::{DType, Device, IndexOp, Tensor, D}; use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution}; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; use image::DynamicImage; use log::info; use rand::Rng; -use tauri::{AppHandle, Emitter}; -pub struct Args<'a> { +pub struct Args<'a, P: AsRef, Q: AsRef, R: AsRef, S: AsRef> { pub image: &'a DynamicImage, pub prompt: String, pub uncond_prompt: String, pub cfg_scale: Option, pub i2i_strength: Option, - pub configs: &'a AppConfig, pub model_cache: &'a ModelCache, pub prompt_cache: &'a PromptCache, - pub app_data_root: &'a AppDataRoot, - pub app: &'a AppHandle, + pub sd_version: StableDiffusionVersion, + pub sd_config: StableDiffusionConfig, + pub device: &'a Device, + pub dtype: DType, + pub maybe_seed: Option, + pub scheduler_steps: usize, + pub sdxl_turbo_vae_path: P, + pub sdxl_turbo_unet_path: Q, + pub clip_json_path: R, + pub clip_weights_path: S, } -pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { +pub fn stable_diffusion_pipeline(args: Args<'_, P1, P2, P3, P4>) -> Result + where P1: AsRef, P2: AsRef, P3: AsRef, P4: AsRef +{ let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, - configs, model_cache, prompt_cache, - app, image, - app_data_root, + sd_version, + sd_config, + device, + dtype, + maybe_seed, + scheduler_steps, + sdxl_turbo_vae_path, + sdxl_turbo_unet_path, + clip_json_path, + clip_weights_path, } = args; - let weights_dir = app_data_root.weights_dir(); - println!("Starting image generation with the following configuration:"); - println!(" Model: {:?}", configs.sd_version); + println!(" Model: {:?}", sd_version); println!(" Prompt: {}", prompt); - println!(" Steps: {}", configs.scheduler_steps); - println!(" Device: {:?}", configs.device); + println!(" Steps: {}", scheduler_steps); + println!(" Device: {:?}", device); - println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height); + println!("Model dimensions: {}x{}", sd_config.width, sd_config.height); // TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler. // This is a target for performance improvement // See: https://github.com/huggingface/candle/issues/1331 - let mut scheduler = configs.sd_config.build_scheduler( - configs.scheduler_steps)?; + let mut scheduler = sd_config.build_scheduler(scheduler_steps)?; - let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen()); - configs.device.set_seed(seed)?; + let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen()); + + device.set_seed(seed)?; info!("Using seed: {}", seed); let guidance_scale = match cfg_scale { Some(guidance_scale) => guidance_scale, - None => match configs.sd_version { + None => match sd_version { StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::Xl => 7.5, @@ -90,20 +101,21 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { } else { info!("Prompt is NOT cached! Calculating embedding..."); - let tensor = infer_clip_text_embeddings( - &prompt, - &uncond_prompt, - None, // tokenizer - None, // clip_weights - None, // clip2_weights - configs.sd_version, - &configs.sd_config, - false, // use_f16 - &configs.device, - configs.dtype, + let tensor = infer_clip_text_embeddings(ClipArgs { + prompt: &prompt, + uncond_prompt: &uncond_prompt, + tokenizer: None, // tokenizer + clip_weights: None, // clip_weights + clip2_weights: None, // clip2_weights + sd_version: sd_version, + sd_config: &sd_config, + use_f16: false, // use_f16 + device: device, + dtype: dtype, use_guide_scale, - weights_dir, - )?; + clip_json_path: clip_json_path.as_ref(), + clip_weights_path: clip_weights_path.as_ref(), + })?; prompt_cache.store_copy(&prompt, &tensor)?; tensor }; @@ -112,14 +124,11 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { println!("Loading input image into tensor..."); - let input_image = dynamic_image_to_tensor( - image, - &configs.device, - configs.dtype)?; + let input_image = dynamic_image_to_tensor(image, device, dtype)?; println!("Reference image shape: {:?}", input_image.shape()); - let vae_scale = get_vae_scale(configs.sd_version); + let vae_scale = get_vae_scale(sd_version); let maybe_vae = model_cache.get_vae()?; @@ -128,11 +137,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { None => { info!("No vae found in cache; loading..."); - let vae_file = weights_dir.weight_path(&SDXL_TURBO_VAE); + let vae_file = sdxl_turbo_vae_path.as_ref(); - let vae = configs - .sd_config - .build_vae(vae_file, &configs.device, configs.dtype)?; + let vae = sd_config.build_vae(vae_file, &device, dtype)?; let vae = Arc::new(vae); @@ -149,9 +156,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { None => { info!("No unet found in cache; loading..."); - let unet_weights_file = weights_dir.weight_path(&SDXL_TURBO_UNET); + let unet_weights_file = sdxl_turbo_unet_path.as_ref(); - let unet = configs.sd_config.build_unet(unet_weights_file, &configs.device, 4, true, configs.dtype) + let unet = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype) .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?; let unet = Arc::new(unet); @@ -181,9 +188,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { }; let t_start = { - let start = configs.scheduler_steps - (configs.scheduler_steps as f64 * img2img_strength) as usize; + let start = scheduler_steps - (scheduler_steps as f64 * img2img_strength) as usize; - println!("Starting from step {} of {} (strength: {})", start, configs.scheduler_steps, img2img_strength); + println!("Starting from step {} of {} (strength: {})", start, scheduler_steps, img2img_strength); start }; diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/mod.rs b/crates/desktop/ml_models/src/ml/weights_registry/mod.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/weights_registry/mod.rs rename to crates/desktop/ml_models/src/ml/weights_registry/mod.rs diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs b/crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs rename to crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor.rs diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs b/crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor_builder.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs rename to crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor_builder.rs diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs b/crates/desktop/ml_models/src/ml/weights_registry/weights.rs similarity index 100% rename from crates/desktop/src-tauri/src/ml/weights_registry/weights.rs rename to crates/desktop/ml_models/src/ml/weights_registry/weights.rs diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index a3a354df2e..ef66671ade 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -7,6 +7,7 @@ license = "" repository = "" edition = "2021" rust-version = "1.77.2" +publish = false # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -24,29 +25,23 @@ filesys = { path = "../../lib/files/filesys" } memory_caching = { path = "../../lib/caching/memory_caching" } primitives = { path = "../../lib/data/primitives" } -# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332 -cudarc = "0.13.7" +# Internal ML +# TODO: Only on Windows and Linux +ml_models = { path = "../ml_models" } # External #openssl = { version = "0.10", features = ["vendored"] } anyhow = "1.0" base64 = "0.22.1" bytes = "1.10" -candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] } -candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] } -candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] } const_format = "0.2.34" directories = "6.0.0" dirs = "6.0.0" futures = "0.3" half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] } -hf-hub = "0.3.2" image = "0.25.5" log = "0.4" -ndarray = "0.15.6" -nvidia-checker = "0.1.2" once_cell = "1.20.3" -ort = { version = "1.16.2", features = ["cuda"] } 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 } @@ -57,23 +52,11 @@ tauri = { version = "2.2.4", features = [] } tauri-plugin-log = "2" tempdir = "0.3.7" tempfile = "3.18.0" -tokenizers = "0.15.2" tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } tokio-util = { version = "0.7", features = ["codec"] } # tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] } -[features] -default = ["cuda"] -cuda = [ - "candle-core/cuda", - "candle-nn/cuda", - "candle-transformers/cuda", -# NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues - "candle-transformers/flash-attn", - "candle-core/cudnn" -] - [dev-dependencies] libloading = "0.8.6" serial_test = "3.2.0" diff --git a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs index 9709736ba0..33b29afedb 100644 --- a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs @@ -1,6 +1,3 @@ -use crate::ml::model_cache::ModelCache; -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 crate::utils::image::decode_base64_image::decode_base64_image; @@ -8,6 +5,10 @@ use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_pn use image::imageops::FilterType; use image::{DynamicImage, RgbImage}; use log::{error, info}; +use ml_models::ml::model_cache::ModelCache; +use ml_models::ml::prompt_cache::PromptCache; +use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; +use ml_models::ml::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use std::path::PathBuf; use tauri::{AppHandle, State}; @@ -51,7 +52,6 @@ pub async fn infer_image( &model_config, &model_cache, prompt_cache, - app, &app_data_root ).await; @@ -78,22 +78,32 @@ async fn do_infer_image( config: &AppConfig, model_cache: &ModelCache, prompt_cache: State<'_, PromptCache>, - app: AppHandle, app_data_root: &AppDataRoot, ) -> Result { + + let weights_dir = app_data_root.weights_dir(); + let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE); + let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET); + let clip_json_path = weights_dir.weight_path(&CLIP_JSON); + //let clip_weights_path= weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16); + let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER); let args = Args { image: &image, prompt: prompt.to_string(), uncond_prompt: NEGATIVE_PROMPT.to_string(), model_cache, - configs: config, prompt_cache: &prompt_cache, - i2i_strength: strength, + img2img_strength: strength, cfg_scale: config.cfg_scale, - app: &app, - app_data_root, use_flash_attn: true, + model_config: &config.model_config, + maybe_seed: None, + scheduler_steps: config.scheduler_steps, + vae_path: &vae_path, + unet_path: &unet_path, + clip_json_path: &clip_json_path, + clip_weights_path: &clip_weights_path, }; let image = lcm_pipeline(args) diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index e487519b34..72621249cb 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -1,16 +1,15 @@ //! 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, info}; -use ndarray::{Array, CowArray}; -use ort::Value; +use ml_models::ml::background_removal::onnx_session::onnx_session; +use ml_models::ml::background_removal::remove_image_background::remove_image_background; +use ml_models::ml::model_cache::ModelCache; +use ml_models::ml::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use tauri::{AppHandle, State}; -use crate::ml::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; /// This handler removes the background from an image. #[tauri::command] @@ -25,68 +24,14 @@ pub async fn remove_background( 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().weight_path(&DIS_MEDIUM_ONNX); - 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 = remove_image_background(model_path, image) + .await + .map_err(|err| format!("Couldn't remove image 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/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs index 081625fb81..e7b9b2c5a1 100644 --- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -1,6 +1,3 @@ -use crate::ml::model_cache::ModelCache; -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 crate::utils::image::decode_base64_image::decode_base64_image; @@ -8,10 +5,14 @@ use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_pn use image::imageops::FilterType; use image::{DynamicImage, ImageReader, RgbImage}; use log::{error, info}; +use ml_models::ml::model_cache::ModelCache; +use ml_models::ml::prompt_cache::PromptCache; +use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; use once_cell::sync::Lazy; use std::io::Cursor; use std::path::PathBuf; use tauri::{AppHandle, State}; +use ml_models::ml::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; const RANDOM_SEED: u32 = 42; @@ -77,19 +78,29 @@ async fn do_infer_image( app: AppHandle, app_data_root: &AppDataRoot, ) -> Result { + + let weights_dir = app_data_root.weights_dir(); + let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE); + let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET); + let clip_json_path = weights_dir.weight_path(&CLIP_JSON); + let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER); let args = Args { image: &image, prompt: prompt.to_string(), uncond_prompt: "".to_string(), model_cache, - configs: config, prompt_cache: &prompt_cache, - i2i_strength: strength, + img2img_strength: strength, cfg_scale: config.cfg_scale, - app: &app, - app_data_root, use_flash_attn: true, + model_config: &config.model_config, + maybe_seed: None, + scheduler_steps: config.scheduler_steps, + vae_path: &vae_path, + unet_path: &unet_path, + clip_json_path: &clip_json_path, + clip_weights_path: &clip_weights_path, }; let image = lcm_pipeline(args) diff --git a/crates/desktop/src-tauri/src/events/notification_event.rs b/crates/desktop/src-tauri/src/events/notification_event.rs index 5bfaa5c491..bacbcd73c2 100644 --- a/crates/desktop/src-tauri/src/events/notification_event.rs +++ b/crates/desktop/src-tauri/src/events/notification_event.rs @@ -1,6 +1,5 @@ use serde::Serialize; - #[derive(Clone, Copy, Serialize)] #[serde(rename_all = "snake_case")] pub enum NotificationModelType { diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index de8cbed04c..a80c30796a 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -1,6 +1,5 @@ pub mod endpoints; pub mod events; -pub mod ml; pub mod state; pub mod threads; pub mod transfer; @@ -11,11 +10,11 @@ use crate::endpoints::realtime_image_endpoint::infer_image; use crate::endpoints::remove_background_endpoint::remove_background; use crate::endpoints::test_counter::test_counter; use crate::endpoints::text_to_image_endpoint::text_to_image; -use crate::ml::model_cache::ModelCache; -use crate::ml::prompt_cache::PromptCache; use crate::state::app_config::AppConfig; use crate::threads::downloader_thread::downloader_thread; use crate::utils::log_environment_details::log_environment_details; +use ml_models::ml::model_cache::ModelCache; +use ml_models::ml::prompt_cache::PromptCache; use tauri_plugin_log::Target; use tauri_plugin_log::TargetKind; diff --git a/crates/desktop/src-tauri/src/ml/background_removal/mod.rs b/crates/desktop/src-tauri/src/ml/background_removal/mod.rs deleted file mode 100644 index 2aed53ac31..0000000000 --- a/crates/desktop/src-tauri/src/ml/background_removal/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod onnx_session; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/state/app_config.rs b/crates/desktop/src-tauri/src/state/app_config.rs index 929579e28b..37bca072fa 100644 --- a/crates/desktop/src-tauri/src/state/app_config.rs +++ b/crates/desktop/src-tauri/src/state/app_config.rs @@ -1,32 +1,14 @@ -use crate::ml::model_file::StableDiffusionVersion; use crate::state::app_dir::AppDataRoot; use crate::state::os_platform::OsPlatform; use crate::state::yaml_config::YamlConfig; -use candle_core::{DType, Device}; -use candle_transformers::models::stable_diffusion::StableDiffusionConfig; -use filesys::recursively_find_file_by_name::recursively_find_file_by_name; -use hf_hub::api::sync::Api; -use log::info; -use std::path::PathBuf; +use ml_models::ml::model_config::ModelConfig; const DEFAULT_SD_IMAGE_WIDTH: usize = 1024; const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024; -const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32; -const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5; /// A central place to configure app-wide details. pub struct AppConfig { - /// We only support one device type currently, and we can't multiplex - /// over several devices yet. - pub device: Device, - - /// The core numeric type to use for models. - pub dtype: DType, - - pub sd_version: StableDiffusionVersion, - pub sd_config: StableDiffusionConfig, - pub image_height: usize, pub image_width: usize, pub scheduler_steps: usize, @@ -36,11 +18,10 @@ pub struct AppConfig { pub cfg_scale: Option, - /// Probably shouldn't be here. - pub hf_api: Api, - /// Root location for application data. pub app_data_root: AppDataRoot, + + pub model_config: ModelConfig, } impl AppConfig { @@ -50,30 +31,6 @@ impl AppConfig { println!("Configs: {:?}", yaml_configs); - let device = Device::new_cuda(0) - .unwrap_or_else(|e| { - info!("CUDA not available ({}), falling back to CPU", e); - Device::Cpu - }); - - - let sd_config = match SD_VERSION { - StableDiffusionVersion::Turbo => { - StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) - } - StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => { - StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) - } - StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => { - StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) - } - StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => { - StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH)) - } - }; - - let hf_api = Api::new()?; - let os_platform = OsPlatform::get(); let default_data_root = match os_platform { @@ -92,18 +49,14 @@ impl AppConfig { println!("App data root: {:?}", app_data_root.path()); Ok(Self { - device, - dtype: DEFAULT_SD_NUMERIC_DATATYPE, - sd_version: SD_VERSION, - sd_config, image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT), image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH), scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(4), scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15), seed: yaml_configs.seed, cfg_scale: yaml_configs.cfg_scale, - hf_api, app_data_root, + model_config: ModelConfig::init()?, }) } } diff --git a/crates/desktop/src-tauri/src/state/app_dir.rs b/crates/desktop/src-tauri/src/state/app_dir.rs index c553f0c095..61b02f9977 100644 --- a/crates/desktop/src-tauri/src/state/app_dir.rs +++ b/crates/desktop/src-tauri/src/state/app_dir.rs @@ -1,8 +1,8 @@ -use crate::ml::weights_registry::weight_descriptor::WeightDescriptor; use crate::state::expanduser::expanduser; use crate::state::os_platform::OsPlatform; use anyhow::anyhow; use directories::UserDirs; +use ml_models::ml::weights_registry::weight_descriptor::WeightDescriptor; use std::path::{Path, PathBuf}; use tempdir::TempDir; use tempfile::{Builder, NamedTempFile}; diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 8364ff02ec..0f7d8a5c07 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -4,6 +4,8 @@ use crate::transfer::download::{download_async, ProgressUpdate}; use futures::stream::FuturesUnordered; use futures::StreamExt; use log::{debug, error, info}; +use ml_models::ml::weights_registry::weight_descriptor::WeightFunction; +use ml_models::ml::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use std::collections::VecDeque; use std::path::PathBuf; use std::sync::mpsc::Receiver; @@ -11,8 +13,6 @@ use std::sync::{Arc, Mutex}; use tauri::{AppHandle, Emitter}; use tempfile::NamedTempFile; use tokio::task::JoinHandle; -use crate::ml::weights_registry::weight_descriptor::WeightFunction; -use crate::ml::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; const NOTIFICATION_CHANNEL_NAME : &str = "notification"; const MAX_FILES : usize = 8; 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 index defb66f7c2..fcd248eb2f 100644 --- 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 @@ -1,7 +1,6 @@ -use crate::ml::image::tensor_to_image_buffer::RgbImage; use base64::prelude::BASE64_STANDARD; use base64::Engine; -use image::ImageFormat; +use image::{ImageFormat, RgbImage}; use std::io::Cursor; pub fn encode_rgb_image_base64_png( diff --git a/crates/desktop/src-tauri/src/utils/log_environment_details.rs b/crates/desktop/src-tauri/src/utils/log_environment_details.rs index 3b54104246..755894daa5 100644 --- a/crates/desktop/src-tauri/src/utils/log_environment_details.rs +++ b/crates/desktop/src-tauri/src/utils/log_environment_details.rs @@ -1,11 +1,11 @@ use log::info; -use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version}; +//use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version}; pub fn log_environment_details() { - info!("NVIDIA driver version: {}", get_nvidia_driver_version(false)); - info!("CUDA version: {}", get_cuda_version(false)); - info!("cuDNN version: {}", get_cudnn_version(false)); - info!("TensorRT version: {}", get_tensorrt_version(false)); - info!("OS version: {}", get_os_version(false)); - info!("Kernel version: {}", get_kernel_version(false)); + //info!("NVIDIA driver version: {}", get_nvidia_driver_version(false)); + //info!("CUDA version: {}", get_cuda_version(false)); + //info!("cuDNN version: {}", get_cudnn_version(false)); + //info!("TensorRT version: {}", get_tensorrt_version(false)); + //info!("OS version: {}", get_os_version(false)); + //info!("Kernel version: {}", get_kernel_version(false)); } diff --git a/crates/desktop/src-tauri/tauri.conf.json b/crates/desktop/src-tauri/tauri.conf.json index efc6948567..f408df7602 100644 --- a/crates/desktop/src-tauri/tauri.conf.json +++ b/crates/desktop/src-tauri/tauri.conf.json @@ -32,9 +32,6 @@ "icons/128x128@2x.png", "icons/icon.icns", "icons/icon.ico" - ], - "resources": { - "../../../windows_dlls/*.dll": "." - } + ] } } From cd3c864eff07724a1bd5189e266099e29ca591a1 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Mar 2025 21:46:25 -0400 Subject: [PATCH 2/5] runs on mac --- Cargo.lock | 10 ++++ Cargo.toml | 2 + crates/desktop/ml_models/Cargo.toml | 1 + crates/desktop/ml_models/src/ml/mod.rs | 3 +- .../infer_clip_text_embeddings.rs | 2 +- .../src/ml/stable_diffusion/lcm_pipeline.rs | 20 +++---- crates/desktop/ml_weights_registry/Cargo.toml | 18 ++++++ crates/desktop/ml_weights_registry/src/lib.rs | 1 + .../src}/weights_registry/mod.rs | 0 .../weights_registry/weight_descriptor.rs | 0 .../weight_descriptor_builder.rs | 4 +- .../src}/weights_registry/weights.rs | 4 +- crates/desktop/src-tauri/Cargo.toml | 9 ++- .../src/endpoints/realtime_image_endpoint.rs | 58 ++++++++++-------- .../endpoints/remove_background_endpoint.rs | 39 ++++++++---- .../src/endpoints/text_to_image_endpoint.rs | 60 +++++++++++-------- crates/desktop/src-tauri/src/lib.rs | 7 ++- .../desktop/src-tauri/src/state/app_config.rs | 2 +- crates/desktop/src-tauri/src/state/app_dir.rs | 2 +- crates/desktop/src-tauri/src/stubs/mod.rs | 3 + .../src-tauri/src/stubs/model_cache.rs | 14 +++++ .../src-tauri/src/stubs/model_config.rs | 14 +++++ .../src-tauri/src/stubs/prompt_cache.rs | 15 +++++ .../src/threads/downloader_thread.rs | 6 +- frontend/realtime/yarn.lock | 22 +++---- 25 files changed, 218 insertions(+), 98 deletions(-) create mode 100644 crates/desktop/ml_weights_registry/Cargo.toml create mode 100644 crates/desktop/ml_weights_registry/src/lib.rs rename crates/desktop/{ml_models/src/ml => ml_weights_registry/src}/weights_registry/mod.rs (100%) rename crates/desktop/{ml_models/src/ml => ml_weights_registry/src}/weights_registry/weight_descriptor.rs (100%) rename crates/desktop/{ml_models/src/ml => ml_weights_registry/src}/weights_registry/weight_descriptor_builder.rs (73%) rename crates/desktop/{ml_models/src/ml => ml_weights_registry/src}/weights_registry/weights.rs (91%) create mode 100644 crates/desktop/src-tauri/src/stubs/mod.rs create mode 100644 crates/desktop/src-tauri/src/stubs/model_cache.rs create mode 100644 crates/desktop/src-tauri/src/stubs/model_config.rs create mode 100644 crates/desktop/src-tauri/src/stubs/prompt_cache.rs diff --git a/Cargo.lock b/Cargo.lock index 190ad0c1fe..fd1550a587 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -672,6 +672,7 @@ dependencies = [ "log", "memory_caching", "ml_models", + "ml_weights_registry", "once_cell", "primitives", "rand 0.8.5", @@ -6819,6 +6820,7 @@ dependencies = [ "libloading 0.8.6", "log", "memory_caching", + "ml_weights_registry", "ndarray", "nvidia-checker", "once_cell", @@ -6839,6 +6841,14 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "ml_weights_registry" +version = "0.0.1" +dependencies = [ + "anyhow", + "const_format", +] + [[package]] name = "mockall" version = "0.12.1" diff --git a/Cargo.toml b/Cargo.toml index ba5ddaee9b..8bc0a9c791 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,8 @@ members = [ # Desktop apps "crates/desktop/src-tauri", + "crates/desktop/ml_models", + "crates/desktop/ml_weights_registry", # Shared libraries "crates/lib/actix_cors_configs", diff --git a/crates/desktop/ml_models/Cargo.toml b/crates/desktop/ml_models/Cargo.toml index dc8e074868..a22cb0b0b0 100644 --- a/crates/desktop/ml_models/Cargo.toml +++ b/crates/desktop/ml_models/Cargo.toml @@ -21,6 +21,7 @@ normal = ["my-workspace-hack"] # Internal filesys = { path = "../../lib/files/filesys" } memory_caching = { path = "../../lib/caching/memory_caching" } +ml_weights_registry = { path = "../ml_weights_registry" } primitives = { path = "../../lib/data/primitives" } # Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332 diff --git a/crates/desktop/ml_models/src/ml/mod.rs b/crates/desktop/ml_models/src/ml/mod.rs index e9c0a3f46a..c889cdbae0 100644 --- a/crates/desktop/ml_models/src/ml/mod.rs +++ b/crates/desktop/ml_models/src/ml/mod.rs @@ -4,5 +4,4 @@ pub mod model_cache; pub mod model_config; pub mod model_file; pub mod prompt_cache; -pub mod stable_diffusion; -pub mod weights_registry; \ No newline at end of file +pub mod stable_diffusion; \ No newline at end of file diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs index 4c1591bc44..96fb0fddeb 100644 --- a/crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs @@ -7,7 +7,7 @@ use candle_transformers::models::stable_diffusion; use std::path::{Path, PathBuf}; use crate::ml::model_file::{ModelFile, StableDiffusionVersion}; -use crate::ml::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2}; +use ml_weights_registry::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2}; use anyhow::Error as E; use candle_core::{DType, Device, Module, Tensor, D}; use log::info; diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs index 515717b9b1..041dd3bdf7 100644 --- a/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs @@ -1,6 +1,15 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use anyhow::{anyhow, Error as E, Result}; +use candle_core::{DType, Device, IndexOp, Tensor, D}; +use candle_nn::VarBuilder; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; +use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig}; +use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig; +use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel; +use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig; +use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution}; use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor; use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage}; use crate::ml::model_cache::ModelCache; @@ -10,18 +19,9 @@ use crate::ml::prompt_cache::PromptCache; use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale; use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs}; use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range; -use crate::ml::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; -use anyhow::{anyhow, Error as E, Result}; -use candle_core::{DType, Device, IndexOp, Tensor, D}; -use candle_nn::VarBuilder; -use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig}; -use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig; -use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel; -use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig; -use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution}; -use candle_transformers::models::stable_diffusion::StableDiffusionConfig; use image::DynamicImage; use log::info; +use ml_weights_registry::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use rand::Rng; const DEFAULT_STRENGTH : f64 = 75.0; diff --git a/crates/desktop/ml_weights_registry/Cargo.toml b/crates/desktop/ml_weights_registry/Cargo.toml new file mode 100644 index 0000000000..7d3455d42a --- /dev/null +++ b/crates/desktop/ml_weights_registry/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "ml_weights_registry" +edition = "2021" +version = "0.0.1" +authors = ["storyteller"] +publish = false +description = "" +license = "" +repository = "" + +[lib] +name = "ml_weights_registry" +path = "src/lib.rs" + +# External +[dependencies] +anyhow = "1.0" +const_format = "0.2.34" diff --git a/crates/desktop/ml_weights_registry/src/lib.rs b/crates/desktop/ml_weights_registry/src/lib.rs new file mode 100644 index 0000000000..ec8cab49f0 --- /dev/null +++ b/crates/desktop/ml_weights_registry/src/lib.rs @@ -0,0 +1 @@ +pub mod weights_registry; diff --git a/crates/desktop/ml_models/src/ml/weights_registry/mod.rs b/crates/desktop/ml_weights_registry/src/weights_registry/mod.rs similarity index 100% rename from crates/desktop/ml_models/src/ml/weights_registry/mod.rs rename to crates/desktop/ml_weights_registry/src/weights_registry/mod.rs diff --git a/crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor.rs b/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs similarity index 100% rename from crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor.rs rename to crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs diff --git a/crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor_builder.rs b/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor_builder.rs similarity index 73% rename from crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor_builder.rs rename to crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor_builder.rs index 78f0a25fa4..61d7ce4e24 100644 --- a/crates/desktop/ml_models/src/ml/weights_registry/weight_descriptor_builder.rs +++ b/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor_builder.rs @@ -5,10 +5,10 @@ pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d64 // and variables (even if they're known to be &'static) and keep allocations static. macro_rules! weight { ($name: literal, $file_name:literal, $weight_function:expr) => { - $crate::ml::weights_registry::weight_descriptor::WeightDescriptor { + $crate::weights_registry::weight_descriptor::WeightDescriptor { name: $name, filename: $file_name, - r2_download_url: const_format::concatcp!(crate::ml::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name), + r2_download_url: const_format::concatcp!(crate::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name), function: $weight_function, } }; diff --git a/crates/desktop/ml_models/src/ml/weights_registry/weights.rs b/crates/desktop/ml_weights_registry/src/weights_registry/weights.rs similarity index 91% rename from crates/desktop/ml_models/src/ml/weights_registry/weights.rs rename to crates/desktop/ml_weights_registry/src/weights_registry/weights.rs index f2d35f50da..5c3e0cef93 100644 --- a/crates/desktop/ml_models/src/ml/weights_registry/weights.rs +++ b/crates/desktop/ml_weights_registry/src/weights_registry/weights.rs @@ -1,5 +1,5 @@ -use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction}; -use crate::ml::weights_registry::weight_descriptor_builder::weight; +use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction}; +use crate::weights_registry::weight_descriptor_builder::weight; pub const CLIP_JSON : WeightDescriptor = weight!( "CLIP", diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index ef66671ade..83d504da8c 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -23,11 +23,12 @@ tauri-build = { version = "2.0.4", features = [] } # Internal filesys = { path = "../../lib/files/filesys" } memory_caching = { path = "../../lib/caching/memory_caching" } +ml_weights_registry = { path = "../ml_weights_registry" } primitives = { path = "../../lib/data/primitives" } # Internal ML # TODO: Only on Windows and Linux -ml_models = { path = "../ml_models" } +ml_models = { path = "../ml_models", optional = true } # External #openssl = { version = "0.10", features = ["vendored"] } @@ -62,6 +63,12 @@ libloading = "0.8.6" serial_test = "3.2.0" assertor = "0.0.3" +[target.'cfg(target_os = "macos")'.features] +default = [] + +[target.'cfg(not(target_os = "macos"))'.features] +default = ["dep:ml_models"] + [profile.dev] incremental = true # Compile your binary in smaller steps. diff --git a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs index 33b29afedb..ad87bc7957 100644 --- a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs @@ -1,17 +1,23 @@ +#[cfg(feature = "ml_models")] +use { + ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}, +}; + 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 crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png; use image::imageops::FilterType; use image::{DynamicImage, RgbImage}; use log::{error, info}; -use ml_models::ml::model_cache::ModelCache; -use ml_models::ml::prompt_cache::PromptCache; -use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; -use ml_models::ml::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; +use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use std::path::PathBuf; use tauri::{AppHandle, State}; + const PROMPT_FILENAME : &str = "prompt.txt"; const PROMPT: &str = "shiba inu, cute dog, detailed, walking in a wooded forest, photorealistic, 8k"; @@ -43,27 +49,30 @@ pub async fn infer_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 mut image = image.resize(512, 512, FilterType::CatmullRom); - let result = do_infer_image( - &prompt, - image, - strength, - &model_config, - &model_cache, - prompt_cache, - &app_data_root - ).await; - - 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)); - } - }; + #[cfg(feature = "ml_models")] + { + let result = infer_image_impl( + &prompt, + image, + strength, + &model_config, + &model_cache, + prompt_cache, + &app_data_root + ).await; - let bytes = encode_rgb_image_base64_png(image) + image = match result { + Ok(img) => DynamicImage::ImageRgb8(img), + 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_dynamic_image_base64_png(image) .map_err(|err| format!("failure to encode image: {:?}", err))?; info!("Inference successful; image converted to base64, serving back to browser..."); @@ -71,7 +80,8 @@ pub async fn infer_image( Ok(bytes) } -async fn do_infer_image( +#[cfg(feature = "ml_models")] +async fn infer_image_impl( prompt: &str, image: DynamicImage, strength: Option, diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index 72621249cb..6291f57948 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -1,37 +1,50 @@ //! Adapted from https://github.com/dnanhkhoa/rust-background-removal +#[cfg(feature = "ml_models")] +use { + ml_models::ml::background_removal::onnx_session::onnx_session, + ml_models::ml::background_removal::remove_image_background::remove_image_background, +}; + 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 crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; use image::{imageops, DynamicImage}; use log::{error, info}; -use ml_models::ml::background_removal::onnx_session::onnx_session; -use ml_models::ml::background_removal::remove_image_background::remove_image_background; -use ml_models::ml::model_cache::ModelCache; -use ml_models::ml::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; +use ml_weights_registry::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; 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 { info!("remove_background endpoint called."); - let image = decode_base64_image(image) - .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; + let mut image = decode_base64_image(image) + .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; - let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX); - - let image = remove_image_background(model_path, image) - .await - .map_err(|err| format!("Couldn't remove image from image: {}", err))?; + #[cfg(feature = "ml_models")] + { + image = remove_background_impl(&app_data_root, image).await?; + } let image = encode_dynamic_image_base64_png(image) .map_err(|err| format!("failure to encode image: {:?}", err))?; Ok(image) } + +#[cfg(feature = "ml_models")] +async fn remove_background_impl(app_data_root: &AppDataRoot, image: DynamicImage) -> Result { + let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX); + + let image = remove_image_background(model_path, image) + .await + .map_err(|err| format!("Couldn't remove image from image: {}", err))?; + + Ok(image) +} diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs index e7b9b2c5a1..c1fff1ee4b 100644 --- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -1,18 +1,24 @@ +#[cfg(feature = "ml_models")] +use { + ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}, +}; + 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 crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png; use image::imageops::FilterType; use image::{DynamicImage, ImageReader, RgbImage}; use log::{error, info}; -use ml_models::ml::model_cache::ModelCache; -use ml_models::ml::prompt_cache::PromptCache; -use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; +use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; + use once_cell::sync::Lazy; use std::io::Cursor; use std::path::PathBuf; use tauri::{AppHandle, State}; -use ml_models::ml::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; const RANDOM_SEED: u32 = 42; @@ -39,28 +45,31 @@ pub async fn text_to_image( ) -> Result { info!("text_to_image endpoint called."); - let image = IMAGE.resize(512, 512, FilterType::CatmullRom); + let mut image = IMAGE.resize(512, 512, FilterType::CatmullRom); - let result = do_infer_image( - &prompt, - image, - Some(STRENGTH), - &model_config, - &model_cache, - prompt_cache, - app, - &app_data_root - ).await; - - 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)); - } - }; + #[cfg(feature = "ml_models")] + { + let result = text_to_image_impl( + &prompt, + image, + Some(STRENGTH), + &model_config, + &model_cache, + prompt_cache, + app, + &app_data_root + ).await; - let bytes = encode_rgb_image_base64_png(image) + image = match result { + Ok(img) => DynamicImage::ImageRgb8(img), + 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_dynamic_image_base64_png(image) .map_err(|err| format!("failure to encode image: {:?}", err))?; info!("Inference successful; image converted to base64, serving back to browser..."); @@ -68,7 +77,8 @@ pub async fn text_to_image( Ok(bytes) } -async fn do_infer_image( +#[cfg(feature = "ml_models")] +async fn text_to_image_impl( prompt: &str, image: DynamicImage, strength: Option, diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index a80c30796a..169a4c9c6f 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -1,6 +1,7 @@ pub mod endpoints; pub mod events; pub mod state; +pub mod stubs; pub mod threads; pub mod transfer; pub mod utils; @@ -11,10 +12,11 @@ use crate::endpoints::remove_background_endpoint::remove_background; use crate::endpoints::test_counter::test_counter; use crate::endpoints::text_to_image_endpoint::text_to_image; use crate::state::app_config::AppConfig; +use crate::stubs::model_cache::ModelCache; +use crate::stubs::prompt_cache::PromptCache; use crate::threads::downloader_thread::downloader_thread; use crate::utils::log_environment_details::log_environment_details; -use ml_models::ml::model_cache::ModelCache; -use ml_models::ml::prompt_cache::PromptCache; + use tauri_plugin_log::Target; use tauri_plugin_log::TargetKind; @@ -27,6 +29,7 @@ pub fn run() { let config = AppConfig::init() .expect("config should load"); + let prompt_cache = PromptCache::with_capacity(8) .expect("prompt cache should load"); diff --git a/crates/desktop/src-tauri/src/state/app_config.rs b/crates/desktop/src-tauri/src/state/app_config.rs index 37bca072fa..fe2e602ddc 100644 --- a/crates/desktop/src-tauri/src/state/app_config.rs +++ b/crates/desktop/src-tauri/src/state/app_config.rs @@ -1,7 +1,7 @@ use crate::state::app_dir::AppDataRoot; use crate::state::os_platform::OsPlatform; use crate::state::yaml_config::YamlConfig; -use ml_models::ml::model_config::ModelConfig; +use crate::stubs::model_config::ModelConfig; const DEFAULT_SD_IMAGE_WIDTH: usize = 1024; const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024; diff --git a/crates/desktop/src-tauri/src/state/app_dir.rs b/crates/desktop/src-tauri/src/state/app_dir.rs index 61b02f9977..0d515edf75 100644 --- a/crates/desktop/src-tauri/src/state/app_dir.rs +++ b/crates/desktop/src-tauri/src/state/app_dir.rs @@ -2,7 +2,7 @@ use crate::state::expanduser::expanduser; use crate::state::os_platform::OsPlatform; use anyhow::anyhow; use directories::UserDirs; -use ml_models::ml::weights_registry::weight_descriptor::WeightDescriptor; +use ml_weights_registry::weights_registry::weight_descriptor::WeightDescriptor; use std::path::{Path, PathBuf}; use tempdir::TempDir; use tempfile::{Builder, NamedTempFile}; diff --git a/crates/desktop/src-tauri/src/stubs/mod.rs b/crates/desktop/src-tauri/src/stubs/mod.rs new file mode 100644 index 0000000000..62ecaceb95 --- /dev/null +++ b/crates/desktop/src-tauri/src/stubs/mod.rs @@ -0,0 +1,3 @@ +pub mod model_cache; +pub mod model_config; +pub mod prompt_cache; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/stubs/model_cache.rs b/crates/desktop/src-tauri/src/stubs/model_cache.rs new file mode 100644 index 0000000000..9737dc1a07 --- /dev/null +++ b/crates/desktop/src-tauri/src/stubs/model_cache.rs @@ -0,0 +1,14 @@ +#[cfg(feature = "ml_models")] +use ml_models::ml::model_cache::ModelCache; + +#[cfg(not(feature = "ml_models"))] +pub struct ModelCache { + // Intentionally left blank +} + +#[cfg(not(feature = "ml_models"))] +impl ModelCache { + pub fn new() -> Self { + Self {} + } +} diff --git a/crates/desktop/src-tauri/src/stubs/model_config.rs b/crates/desktop/src-tauri/src/stubs/model_config.rs new file mode 100644 index 0000000000..0047bb375a --- /dev/null +++ b/crates/desktop/src-tauri/src/stubs/model_config.rs @@ -0,0 +1,14 @@ +#[cfg(feature = "ml_models")] +use ml_models::ml::model_config::ModelConfig; + +#[cfg(not(feature = "ml_models"))] +pub struct ModelConfig { + // Intentionally left blank +} + +#[cfg(not(feature = "ml_models"))] +impl ModelConfig { + pub fn init() -> anyhow::Result { + Ok(Self {}) + } +} diff --git a/crates/desktop/src-tauri/src/stubs/prompt_cache.rs b/crates/desktop/src-tauri/src/stubs/prompt_cache.rs new file mode 100644 index 0000000000..907a657246 --- /dev/null +++ b/crates/desktop/src-tauri/src/stubs/prompt_cache.rs @@ -0,0 +1,15 @@ + +#[cfg(feature = "ml_models")] +pub use ml_models::ml::prompt_cache::PromptCache; + +#[cfg(not(feature = "ml_models"))] +pub struct PromptCache { + // Intentionally left blank +} + +#[cfg(not(feature = "ml_models"))] +impl PromptCache { + pub fn with_capacity(capacity: usize) -> anyhow::Result { + Ok(Self {}) + } +} diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 0f7d8a5c07..b9bf7813d3 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -1,11 +1,11 @@ use crate::events::notification_event::{NotificationEvent, NotificationModelType}; use crate::state::app_dir::AppDataRoot; use crate::transfer::download::{download_async, ProgressUpdate}; -use futures::stream::FuturesUnordered; use futures::StreamExt; +use futures::stream::FuturesUnordered; use log::{debug, error, info}; -use ml_models::ml::weights_registry::weight_descriptor::WeightFunction; -use ml_models::ml::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; +use ml_weights_registry::weights_registry::weight_descriptor::WeightFunction; +use ml_weights_registry::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use std::collections::VecDeque; use std::path::PathBuf; use std::sync::mpsc::Receiver; diff --git a/frontend/realtime/yarn.lock b/frontend/realtime/yarn.lock index 53d6272415..0f745335c0 100644 --- a/frontend/realtime/yarn.lock +++ b/frontend/realtime/yarn.lock @@ -222,11 +222,6 @@ resolved "https://registry.npmjs.org/@emotion/unitless/-/unitless-0.7.5.tgz" integrity sha512-OWORNpfjMsSSUBVrRBVGECkhWcULOAJz9ZW8uK9qgxD+87M7jHRcvh/A96XXNhXTLmKcoYSQtBEX7lHMO7YRwg== -"@esbuild/win32-x64@0.21.5": - version "0.21.5" - resolved "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.21.5.tgz" - integrity sha512-tQd/1efJuzPC6rCFwEvLtci/xNFcTZknmXs98FYDfGE4wP9ClFV98nyKrzJKVPMhdDnjzLhdUyMX4PsQAPjwIw== - "@eslint-community/eslint-utils@^4.2.0", "@eslint-community/eslint-utils@^4.4.0": version "4.4.0" resolved "https://registry.npmjs.org/@eslint-community/eslint-utils/-/eslint-utils-4.4.0.tgz" @@ -741,10 +736,10 @@ resolved "https://registry.npmjs.org/@remix-run/router/-/router-1.19.1.tgz" integrity sha512-S45oynt/WH19bHbIXjtli6QmwNYvaz+vtnubvNpNDvUOoA/OWh6j1OikIP3G+v5GHdxyC6EXoChG3HgYGEUfcg== -"@rollup/rollup-win32-x64-msvc@4.24.0": +"@rollup/rollup-darwin-arm64@4.24.0": version "4.24.0" - resolved "https://registry.npmjs.org/@rollup/rollup-win32-x64-msvc/-/rollup-win32-x64-msvc-4.24.0.tgz" - integrity sha512-fbMkAF7fufku0N2dE5TBXcNlg0pt0cJue4xBRE2Qc5Vqikxr4VCgKj/ht6SMdFcOacVA9rqF70APJ8RN/4vMJw== + resolved "https://registry.npmjs.org/@rollup/rollup-darwin-arm64/-/rollup-darwin-arm64-4.24.0.tgz" + integrity sha512-bIv+X9xeSs1XCk6DVvkO+S/z8/2AMt/2lMqdQbMrmVpgFvXlmde9mLcbQpztXm1tajC3raFDqegsH18HQPMYtA== "@sigstore/bundle@^2.3.2": version "2.3.2" @@ -816,10 +811,10 @@ resolved "https://registry.npmjs.org/@tauri-apps/api/-/api-2.2.0.tgz" integrity sha512-R8epOeZl1eJEl603aUMIGb4RXlhPjpgxbGVEaqY+0G5JG9vzV/clNlzTeqc+NLYXVqXcn8mb4c5b9pJIUDEyAg== -"@tauri-apps/cli-win32-x64-msvc@2.2.7": +"@tauri-apps/cli-darwin-arm64@2.2.7": version "2.2.7" - resolved "https://registry.npmjs.org/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.2.7.tgz" - integrity sha512-ZB8Kw90j8Ld+9tCWyD2fWCYfIrzbQohJ4DJSidNwbnehlZzP7wAz6Z3xjsvUdKtQ3ibtfoeTqVInzCCEpI+pWg== + resolved "https://registry.npmjs.org/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.2.7.tgz" + integrity sha512-54kcpxZ3X1Rq+pPTzk3iIcjEVY4yv493uRx/80rLoAA95vAC0c//31Whz75UVddDjJfZvXlXZ3uSZ+bnCOnt0A== "@tauri-apps/cli@^2.2.7": version "2.2.7" @@ -2082,6 +2077,11 @@ fs-minipass@^3.0.0, fs-minipass@^3.0.3: dependencies: minipass "^7.0.3" +fsevents@~2.3.2, fsevents@~2.3.3: + version "2.3.3" + resolved "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz" + integrity sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw== + function-bind@^1.1.2: version "1.1.2" resolved "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz" From 456bfd7bb90cec20fadda1646436f92ef2c6b9f3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Mar 2025 22:31:58 -0400 Subject: [PATCH 3/5] trigger with --feature ml_models (need to make default) --- crates/desktop/src-tauri/Cargo.toml | 16 ++++++++++++---- .../desktop/src-tauri/src/stubs/model_cache.rs | 2 +- .../desktop/src-tauri/src/stubs/model_config.rs | 2 +- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 83d504da8c..bf1cddb99f 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -28,7 +28,7 @@ primitives = { path = "../../lib/data/primitives" } # Internal ML # TODO: Only on Windows and Linux -ml_models = { path = "../ml_models", optional = true } +#ml_models = { path = "../ml_models", optional = true } # External #openssl = { version = "0.10", features = ["vendored"] } @@ -58,16 +58,24 @@ tokio-util = { version = "0.7", features = ["codec"] } # tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] } +[target.'cfg(not(target_os = "macos"))'.dependencies] +ml_models = { path = "../ml_models" } + [dev-dependencies] libloading = "0.8.6" serial_test = "3.2.0" assertor = "0.0.3" -[target.'cfg(target_os = "macos")'.features] -default = [] +#[target.'cfg(target_os = "macos")'.features] +#default = [] +#ml_models = ["dep:ml_models"] [target.'cfg(not(target_os = "macos"))'.features] -default = ["dep:ml_models"] +default = ["ml_models"] +#ml_models = ["dep:ml_models"] + +[features] +ml_models = [] [profile.dev] incremental = true # Compile your binary in smaller steps. diff --git a/crates/desktop/src-tauri/src/stubs/model_cache.rs b/crates/desktop/src-tauri/src/stubs/model_cache.rs index 9737dc1a07..8b252a100e 100644 --- a/crates/desktop/src-tauri/src/stubs/model_cache.rs +++ b/crates/desktop/src-tauri/src/stubs/model_cache.rs @@ -1,5 +1,5 @@ #[cfg(feature = "ml_models")] -use ml_models::ml::model_cache::ModelCache; +pub use ml_models::ml::model_cache::ModelCache; #[cfg(not(feature = "ml_models"))] pub struct ModelCache { diff --git a/crates/desktop/src-tauri/src/stubs/model_config.rs b/crates/desktop/src-tauri/src/stubs/model_config.rs index 0047bb375a..55fa3e236a 100644 --- a/crates/desktop/src-tauri/src/stubs/model_config.rs +++ b/crates/desktop/src-tauri/src/stubs/model_config.rs @@ -1,5 +1,5 @@ #[cfg(feature = "ml_models")] -use ml_models::ml::model_config::ModelConfig; +pub use ml_models::ml::model_config::ModelConfig; #[cfg(not(feature = "ml_models"))] pub struct ModelConfig { From 8d51d5b8abd13de01aa8fa58228d46ce9e5540ac Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Mar 2025 23:00:29 -0400 Subject: [PATCH 4/5] can't trigger default features; switch to platform switches --- crates/desktop/src-tauri/Cargo.toml | 7 ------- .../src-tauri/src/endpoints/realtime_image_endpoint.rs | 6 +++--- .../src-tauri/src/endpoints/remove_background_endpoint.rs | 6 +++--- .../src-tauri/src/endpoints/text_to_image_endpoint.rs | 6 +++--- crates/desktop/src-tauri/src/stubs/model_cache.rs | 6 +++--- crates/desktop/src-tauri/src/stubs/model_config.rs | 6 +++--- crates/desktop/src-tauri/src/stubs/prompt_cache.rs | 6 +++--- 7 files changed, 18 insertions(+), 25 deletions(-) diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index bf1cddb99f..9426e2d526 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -70,13 +70,6 @@ assertor = "0.0.3" #default = [] #ml_models = ["dep:ml_models"] -[target.'cfg(not(target_os = "macos"))'.features] -default = ["ml_models"] -#ml_models = ["dep:ml_models"] - -[features] -ml_models = [] - [profile.dev] incremental = true # Compile your binary in smaller steps. diff --git a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs index ad87bc7957..5730a959be 100644 --- a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs @@ -1,4 +1,4 @@ -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] use { ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}, }; @@ -51,7 +51,7 @@ pub async fn infer_image( // TODO(bt,2025-02-17): Running out of vram with full image buffer size let mut image = image.resize(512, 512, FilterType::CatmullRom); - #[cfg(feature = "ml_models")] + #[cfg(not(target_os = "macos"))] { let result = infer_image_impl( &prompt, @@ -80,7 +80,7 @@ pub async fn infer_image( Ok(bytes) } -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] async fn infer_image_impl( prompt: &str, image: DynamicImage, diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index 6291f57948..576d7823d7 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -1,6 +1,6 @@ //! Adapted from https://github.com/dnanhkhoa/rust-background-removal -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] use { ml_models::ml::background_removal::onnx_session::onnx_session, ml_models::ml::background_removal::remove_image_background::remove_image_background, @@ -27,7 +27,7 @@ pub async fn remove_background( let mut image = decode_base64_image(image) .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; - #[cfg(feature = "ml_models")] + #[cfg(not(target_os = "macos"))] { image = remove_background_impl(&app_data_root, image).await?; } @@ -38,7 +38,7 @@ pub async fn remove_background( Ok(image) } -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] async fn remove_background_impl(app_data_root: &AppDataRoot, image: DynamicImage) -> Result { let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX); diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs index c1fff1ee4b..c08299305e 100644 --- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -1,4 +1,4 @@ -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] use { ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}, }; @@ -47,7 +47,7 @@ pub async fn text_to_image( let mut image = IMAGE.resize(512, 512, FilterType::CatmullRom); - #[cfg(feature = "ml_models")] + #[cfg(not(target_os = "macos"))] { let result = text_to_image_impl( &prompt, @@ -77,7 +77,7 @@ pub async fn text_to_image( Ok(bytes) } -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] async fn text_to_image_impl( prompt: &str, image: DynamicImage, diff --git a/crates/desktop/src-tauri/src/stubs/model_cache.rs b/crates/desktop/src-tauri/src/stubs/model_cache.rs index 8b252a100e..ebd7e163d3 100644 --- a/crates/desktop/src-tauri/src/stubs/model_cache.rs +++ b/crates/desktop/src-tauri/src/stubs/model_cache.rs @@ -1,12 +1,12 @@ -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] pub use ml_models::ml::model_cache::ModelCache; -#[cfg(not(feature = "ml_models"))] +#[cfg(not(not(target_os = "macos")))] pub struct ModelCache { // Intentionally left blank } -#[cfg(not(feature = "ml_models"))] +#[cfg(not(not(target_os = "macos")))] impl ModelCache { pub fn new() -> Self { Self {} diff --git a/crates/desktop/src-tauri/src/stubs/model_config.rs b/crates/desktop/src-tauri/src/stubs/model_config.rs index 55fa3e236a..d0a6740727 100644 --- a/crates/desktop/src-tauri/src/stubs/model_config.rs +++ b/crates/desktop/src-tauri/src/stubs/model_config.rs @@ -1,12 +1,12 @@ -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] pub use ml_models::ml::model_config::ModelConfig; -#[cfg(not(feature = "ml_models"))] +#[cfg(not(not(target_os = "macos")))] pub struct ModelConfig { // Intentionally left blank } -#[cfg(not(feature = "ml_models"))] +#[cfg(not(not(target_os = "macos")))] impl ModelConfig { pub fn init() -> anyhow::Result { Ok(Self {}) diff --git a/crates/desktop/src-tauri/src/stubs/prompt_cache.rs b/crates/desktop/src-tauri/src/stubs/prompt_cache.rs index 907a657246..daf4a4d2db 100644 --- a/crates/desktop/src-tauri/src/stubs/prompt_cache.rs +++ b/crates/desktop/src-tauri/src/stubs/prompt_cache.rs @@ -1,13 +1,13 @@ -#[cfg(feature = "ml_models")] +#[cfg(not(target_os = "macos"))] pub use ml_models::ml::prompt_cache::PromptCache; -#[cfg(not(feature = "ml_models"))] +#[cfg(not(not(target_os = "macos")))] pub struct PromptCache { // Intentionally left blank } -#[cfg(not(feature = "ml_models"))] +#[cfg(not(not(target_os = "macos")))] impl PromptCache { pub fn with_capacity(capacity: usize) -> anyhow::Result { Ok(Self {}) From 81c256e4f866983077245783c55725877b5da37f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Mar 2025 23:08:07 -0400 Subject: [PATCH 5/5] cleanup --- crates/desktop/src-tauri/src/endpoints/mod.rs | 1 - .../src/endpoints/realtime_image_endpoint.rs | 4 +--- .../endpoints/remove_background_endpoint.rs | 13 ++++------ .../src-tauri/src/endpoints/test_counter.rs | 24 ------------------- .../src/endpoints/text_to_image_endpoint.rs | 9 ++----- crates/desktop/src-tauri/src/lib.rs | 5 ++-- 6 files changed, 9 insertions(+), 47 deletions(-) delete mode 100644 crates/desktop/src-tauri/src/endpoints/test_counter.rs diff --git a/crates/desktop/src-tauri/src/endpoints/mod.rs b/crates/desktop/src-tauri/src/endpoints/mod.rs index f49bdce06e..cee0112b94 100644 --- a/crates/desktop/src-tauri/src/endpoints/mod.rs +++ b/crates/desktop/src-tauri/src/endpoints/mod.rs @@ -2,5 +2,4 @@ pub mod download_models; pub mod flip_image; pub mod realtime_image_endpoint; pub mod remove_background_endpoint; -pub mod test_counter; pub mod text_to_image_endpoint; diff --git a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs index 5730a959be..7671f2a376 100644 --- a/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs @@ -1,7 +1,5 @@ #[cfg(not(target_os = "macos"))] -use { - ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}, -}; +use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; use crate::state::app_config::AppConfig; use crate::state::app_dir::AppDataRoot; diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index 576d7823d7..5c3bfa4eb6 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs @@ -1,20 +1,15 @@ //! Adapted from https://github.com/dnanhkhoa/rust-background-removal #[cfg(not(target_os = "macos"))] -use { - ml_models::ml::background_removal::onnx_session::onnx_session, - ml_models::ml::background_removal::remove_image_background::remove_image_background, -}; +use ml_models::ml::background_removal::remove_image_background::remove_image_background; 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 crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; use image::{imageops, DynamicImage}; -use log::{error, info}; -use ml_weights_registry::weights_registry::weights::{DIS_MEDIUM_ONNX, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; -use tauri::{AppHandle, State}; +use log::info; +use ml_weights_registry::weights_registry::weights::DIS_MEDIUM_ONNX; +use tauri::State; /// This handler removes the background from an image. #[tauri::command] diff --git a/crates/desktop/src-tauri/src/endpoints/test_counter.rs b/crates/desktop/src-tauri/src/endpoints/test_counter.rs deleted file mode 100644 index f377b8c90e..0000000000 --- a/crates/desktop/src-tauri/src/endpoints/test_counter.rs +++ /dev/null @@ -1,24 +0,0 @@ -use std::sync::{Arc, RwLock}; -use once_cell::sync::Lazy; - -static COUNTER : Lazy>> = Lazy::new(|| Arc::new(RwLock::new(0))); - -#[tauri::command] -pub fn test_counter() -> u64 { - let value : u64; - { - match COUNTER.write() { - Ok(mut counter) => { - *counter += 1; - value = *counter; - }, - Err(_e) => { - value = 0; - } - } - } - - println!("I was invoked from JavaScript! {:?}", value); - - value -} diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs index c08299305e..2c41079cf0 100644 --- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs @@ -1,23 +1,18 @@ #[cfg(not(target_os = "macos"))] -use { - ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}, -}; +use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args}; 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 crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png; -use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png; use image::imageops::FilterType; use image::{DynamicImage, ImageReader, RgbImage}; use log::{error, info}; -use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; +use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET}; use once_cell::sync::Lazy; use std::io::Cursor; -use std::path::PathBuf; use tauri::{AppHandle, State}; const RANDOM_SEED: u32 = 42; diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 169a4c9c6f..90a4da4342 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -7,15 +7,14 @@ pub mod transfer; pub mod utils; use crate::endpoints::download_models::download_models; +use crate::endpoints::flip_image::flip_image; use crate::endpoints::realtime_image_endpoint::infer_image; use crate::endpoints::remove_background_endpoint::remove_background; -use crate::endpoints::test_counter::test_counter; use crate::endpoints::text_to_image_endpoint::text_to_image; use crate::state::app_config::AppConfig; use crate::stubs::model_cache::ModelCache; use crate::stubs::prompt_cache::PromptCache; use crate::threads::downloader_thread::downloader_thread; -use crate::utils::log_environment_details::log_environment_details; use tauri_plugin_log::Target; use tauri_plugin_log::TargetKind; @@ -70,9 +69,9 @@ pub fn run() { .manage(app_data_root) .invoke_handler(tauri::generate_handler![ download_models, + flip_image, infer_image, remove_background, - test_counter, text_to_image, ]) .run(tauri::generate_context!())