pull configs and ml deps out of tauri app

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"

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"

crates/desktop/ml_models/src/lib.rs

@@ -0,0 +1 @@
+pub mod ml;

crates/desktop/ml_models/src/ml/background_removal/mod.rs

@@ -0,0 +1,2 @@
+pub mod onnx_session;
+pub mod remove_image_background;

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<ort::Session, anyhow::Error> {
+pub fn onnx_session(onnx_model_file: &Path) -> Result<ort::Session, anyhow::Error> {
   let environment = Environment::default().into_arc();
   let session = SessionBuilder::new(&environment)?
     .with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level

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<P: AsRef<Path>>(model_path: P, image: DynamicImage) -> anyhow::Result<DynamicImage> {
+  // 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<DynamicImage, anyhow::Error> {
+  let input_shape = session.inputs[0]
+    .dimensions()
+    .map(|dim| dim.unwrap())
+    .collect::<Vec<usize>>();
+  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::<f32>()?;
+  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))
+}

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;

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<Self> {
+    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,
+    })
+  }
+}

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<Path>,  Q: AsRef<Path>> {
+  pub prompt: &'a str, 
+  pub uncond_prompt: &'a str, 
+  pub tokenizer: Option<String>, 
+  pub clip_weights: Option<String>, 
+  pub clip2_weights: Option<String>, 
+  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<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
+pub fn infer_clip_text_embeddings<P: AsRef<Path>, Q: AsRef<Path>>(args: ClipArgs<'_, P, Q>) -> anyhow::Result<Tensor> {
+  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::<anyhow::Result<Vec<_>>>()?;
@@ -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<String>, clip_weights: Option<String>, clip2_weights: Option<String>, 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<Tensor> {
+fn do_infer_clip_text_embeddings(
+  prompt: &str, 
+  uncond_prompt: &str, 
+  tokenizer: Option<String>, 
+  clip_weights: Option<String>, 
+  clip2_weights: Option<String>, 
+  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<Tensor> {
   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)?;
 

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<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>> {
   pub image: &'a DynamicImage,
   pub prompt: String,
   pub uncond_prompt: String,
   pub cfg_scale: Option<f64>,
-  pub i2i_strength: Option<f64>,
-  pub configs: &'a AppConfig,
+  pub img2img_strength: Option<f64>,
   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<u64>,
+  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<RgbImage> {
-  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<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage> 
+  where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
+{
+  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<RgbImage> {
   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<RgbImage> {
   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<RgbImage> {
     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<RgbImage> {
       };
 
       // 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<RgbImage> {
 
   // 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<RgbImage> {
   let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
 
   for (timestep_index, &timestep) 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<RgbImage> {
     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<RgbImage> {
     }
 
     // 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)?)?;
 

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<Path>, Q: AsRef<Path>, R: AsRef<Path>, S: AsRef<Path>> {
     pub image: &'a DynamicImage,
     pub prompt: String,
     pub uncond_prompt: String,
     pub cfg_scale: Option<f64>,
     pub i2i_strength: Option<u8>,
-    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<u64>,
+    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<RgbImage> {
+pub fn stable_diffusion_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage> 
+  where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
+{
     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<RgbImage> {
     } 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<RgbImage> {
 
     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<RgbImage> {
         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<RgbImage> {
         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<RgbImage> {
     };
 
     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
     };
 

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"

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<RgbImage, String> {
+  
+  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)

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<DynamicImage, anyhow::Error> {
-  let input_shape = session.inputs[0]
-    .dimensions()
-    .map(|dim| dim.unwrap())
-    .collect::<Vec<usize>>();
-  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::<f32>()?;
-  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))
-}

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<RgbImage, String> {
+  
+  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)

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 {

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;
 

crates/desktop/src-tauri/src/ml/background_removal/mod.rs

@@ -1 +0,0 @@
-pub mod onnx_session;

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<f64>,
   
-  /// 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()?,
     })
   }
 }

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};

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;

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(

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));
 }

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