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": "." - } + ] } }