don't download sdxl anymore

stop downloading clip tokenizer from hf
vae wasn't cached in memory
2026-11-05 22:04:51 -05:00 · 2025-03-13 00:31:51 -04:00 · 2025-03-13 00:24:06 -04:00 · 2025-03-12 23:58:52 -04:00 · 2025-03-12 23:54:42 -04:00 · 2025-03-12 23:38:49 -04:00
54 changed files with 561 additions and 877 deletions
@@ -66,7 +66,3 @@ scraps/
 # These should be built by the platform
 windows/sqlite/sqlite3.exp
 windows/sqlite/sqlite3.lib
-
-# ArtCraft Windows DLLs
-crates/desktop/src-tauri/*.dll
-windows_dlls/
@@ -658,13 +658,11 @@ name = "app"
 version = "0.1.0"
 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",
@@ -676,10 +674,7 @@ dependencies = [
 "libloading 0.8.6",
 "log",
 "memory_caching",
- "ndarray",
 "once_cell",
- "ort",
- "primitives",
 "rand 0.8.5",
 "reqwest 0.12.12",
 "serde",
@@ -860,15 +855,6 @@ version = "1.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "10d2119f741b79fe9907f5396d19bffcb46568cfcc315e78677d731972ac7085"

-[[package]]
-name = "assertor"
-version = "0.0.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4b02f2af4042fbbed4d5732c517666629512c574c8e2edb5cfe2cf8a86564d52"
-dependencies = [
- "num-traits",
-]
-
 [[package]]
 name = "async-channel"
 version = "1.9.0"
@@ -2314,26 +2300,6 @@ dependencies = [
 "tiny-keccak",
 ]

-[[package]]
-name = "const_format"
-version = "0.2.34"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "126f97965c8ad46d6d9163268ff28432e8f6a1196a55578867832e3049df63dd"
-dependencies = [
- "const_format_proc_macros",
-]
-
-[[package]]
-name = "const_format_proc_macros"
-version = "0.2.34"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d57c2eccfb16dbac1f4e61e206105db5820c9d26c3c472bc17c774259ef7744"
-dependencies = [
- "proc-macro2",
- "quote",
- "unicode-xid",
-]
-
 [[package]]
 name = "constant_time_eq"
 version = "0.1.5"
@@ -7045,19 +7011,6 @@ dependencies = [
 "tempfile",
 ]

-[[package]]
-name = "ndarray"
-version = "0.15.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32"
-dependencies = [
- "matrixmultiply",
- "num-complex",
- "num-integer",
- "num-traits",
- "rawpointer",
-]
-
 [[package]]
 name = "ndk"
 version = "0.8.0"
@@ -7840,26 +7793,6 @@ dependencies = [
 "hashbrown 0.14.5",
 ]

-[[package]]
-name = "ort"
-version = "1.16.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "889dca4c98efa21b1ba54ddb2bde44fd4920d910f492b618351f839d8428d79d"
-dependencies = [
- "flate2",
- "half",
- "lazy_static",
- "libc",
- "ndarray",
- "tar",
- "thiserror 1.0.63",
- "tracing",
- "ureq",
- "vswhom",
- "winapi 0.3.9",
- "zip 0.6.6",
-]
-
 [[package]]
 name = "overload"
 version = "0.1.1"
@@ -8427,7 +8360,6 @@ dependencies = [
 name = "primitives"
 version = "0.0.1"
 dependencies = [
- "assertor",
 "my-workspace-hack",
 "testing",
 ]
@@ -12641,9 +12573,9 @@ dependencies = [

 [[package]]
 name = "tracing-subscriber"
-version = "0.3.18"
+version = "0.3.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b"
+checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008"
 dependencies = [
 "nu-ansi-term",
 "sharded-slab",
@@ -12952,12 +12884,6 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd"

-[[package]]
-name = "unicode-xid"
-version = "0.2.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
-
 [[package]]
 name = "unicode_categories"
 version = "0.1.1"
@@ -22,7 +22,6 @@ tauri-build = { version = "2.0.4", features = [] }
 # 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"
@@ -35,9 +34,8 @@ 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"
+directories = "6.0.0"
 futures = "0.3"
 hf-hub = "0.3.2"
 image = "0.25.5"
@@ -58,11 +56,6 @@ tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
 tokio-util = { version = "0.7", features = ["codec"] }
 half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }

-# NB: From `rust-background-removal`
-ndarray = "0.15.6"
-ort = { version = "1.16.2", features = ["cuda"] }
-# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
-
 [features]
 default = ["cuda"]
 cuda = [
@@ -77,7 +70,6 @@ cuda = [
 [dev-dependencies]
 libloading = "0.8.6"
 serial_test = "3.2.0"
-assertor = "0.0.3"

 [profile.dev]
 incremental = true # Compile your binary in smaller steps.
@@ -1,12 +1,16 @@
+use crate::ml::downloads::download_all_models::download_all_models;
 use crate::state::app_dir::AppDataRoot;
 use tauri::State;

 #[tauri::command]
 pub async fn download_models(
-  _app_data_root: State<'_, AppDataRoot>,
+  app_data_root: State<'_, AppDataRoot>,
 ) -> Result<String, String> {
-
-  // NB: This endpoint has been removed for now as we now download models in the app lifecycle.
-
+  
+  download_all_models(&app_data_root).await
+    .map_err(|err| {
+      "error downloading models".to_string()
+    })?;
+  
  Ok(("downloaded".to_string()))
 }
@@ -1,23 +1,27 @@
+use crate::ml::downloads::download_all_models::download_all_models;
 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;
-use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png;
+use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
+use bytes::BytesMut;
 use image::imageops::FilterType;
-use image::{DynamicImage, RgbImage};
+use image::{DynamicImage, ImageFormat, ImageReader};
 use log::{error, info};
+use std::io::Cursor;
 use std::path::PathBuf;
 use tauri::{AppHandle, State};
+use crate::state::app_dir::AppDataRoot;

 const PROMPT_FILENAME : &str = "prompt.txt";
-const PROMPT: &str = "shiba inu, cute dog, detailed, walking in a wooded forest, photorealistic, 8k";
+const PROMPT: &str = "green goblin, blood dripping from lips, detailed, photorealistic, 8k";

 const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";

 const RANDOM_SEED: u32 = 42;

+// TODO: Mark async https://github.com/tauri-apps/tauri/discussions/7737
+

 /// This handler takes an image (as a base64 encoded string) and a prompt and returns
 /// an image (as a base64-encoded string).
@@ -25,24 +29,28 @@ const RANDOM_SEED: u32 = 42;
 pub async fn infer_image(
  image: &str,
  prompt: Option<String>,
-  strength: Option<f64>,
+  strength: Option<u8>,
  model_config: State<'_, AppConfig>,
  model_cache: State<'_, ModelCache>,
  prompt_cache: State<'_, PromptCache>,
  app_data_root: State<'_, AppDataRoot>,
  app: AppHandle,
 ) -> Result<String, String> {
-  info!("infer_image endpoint called.");

  let prompt = get_prompt_or_fallback(prompt);

-  info!("Strength: {:?}; Prompt: {}", strength, prompt);
+  info!("Prompt: {}", prompt);
+  info!("Strength: {:?}", strength);

-  let image = decode_base64_image(image)
+  let image = hydrate_base64_image(image)
    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
-  
-  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
-  let image = image.resize(512, 512, FilterType::CatmullRom);
+
+  // TODO(bt): Move this to another endpoint
+  download_all_models(&app_data_root).await
+    .map_err(|err| {
+      error!("couldn't download models: {:?}", err);
+      "Couldn't download models".to_string()
+    })?;

  let result = do_infer_image(
    &prompt, 
@@ -55,32 +63,23 @@ pub async fn infer_image(
    &app_data_root
  ).await;
  
-  let image = match result {
-    Ok(image) => image,
-    Err(err) => {
-      error!("There was an error generating the image: {:?}", err);
-      return Err(format!("There was an error generating the image: {}", err));
-    }
-  };
-
-  let bytes = encode_rgb_image_base64_png(image)
-    .map_err(|err| format!("failure to encode image: {:?}", err))?;
-
-  info!("Inference successful; image converted to base64, serving back to browser...");
-
-  Ok(bytes)
+  if let Err(err) = result.as_deref() {
+    error!("There was an error: {:?}", err);
+  }
+  
+  result
 }

 async fn do_infer_image(
  prompt: &str,
  image: DynamicImage,
-  strength: Option<f64>,
+  strength: Option<u8>,
  config: &AppConfig,
  model_cache: &ModelCache,
  prompt_cache: State<'_, PromptCache>,
  app: AppHandle,
  app_data_root: &AppDataRoot,
-) -> Result<RgbImage, String> {
+) -> Result<String, String> {

  let args = Args {
    image: &image,
@@ -96,10 +95,19 @@ async fn do_infer_image(
    use_flash_attn: true,
  };

-  let image = lcm_pipeline(args)
-    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+  match lcm_pipeline(args) {
+    Ok(image) => {
+      let mut bytes = Vec::with_capacity(1024*1024);
+      
+      image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
+        .map_err(|err| format!("failure to encode image: {:?}", err))?;

-  Ok(image)
+      let bytes = BASE64_STANDARD.encode(bytes);
+
+      Ok(bytes)
+    },
+    Err(e) => Err(format!("Failed to generate image: {}", e)),
+  }
 }

 fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
@@ -123,3 +131,16 @@ fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
    .trim()
    .to_string()
 }
+
+fn hydrate_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
+  let bytes = BASE64_STANDARD.decode(base64_image)?;
+
+  let image = ImageReader::new(Cursor::new(bytes))
+    .with_guessed_format()?
+    .decode()?;
+
+  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
+  let image = image.resize(512, 512, FilterType::CatmullRom);
+
+  Ok(image)
+}
@@ -1,5 +1,4 @@
 pub mod download_models;
 pub mod flip_image;
 pub mod image_endpoint;
-pub mod remove_background_endpoint;
 pub mod test_counter;
@@ -1,91 +0,0 @@
-//! 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 tauri::{AppHandle, State};
-
-/// This handler removes the background from an image.
-#[tauri::command]
-pub async fn remove_background(
-  image: &str,
-  model_cache: State<'_, ModelCache>,
-  app_data_root: State<'_, AppDataRoot>,
-  app: AppHandle,
-) -> Result<String, String> {
-  info!("remove_background endpoint called.");
-
-  let image = decode_base64_image(image)
-    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
-
-  // NB: Real time needed cuDNN 9, but this uses cuDNN 8:
-  // sudo apt-get install libcudnn8-dev libcudnn8
-  // update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken
-  // update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode
-  let model_path = app_data_root.weights_dir().path().join("medium.onnx"); // TODO
-  let session = onnx_session(model_path.to_str().unwrap()) // TODO
-    .map_err(|err| format!("failure to start onnx session: {:?}", err))?;
-
-  let image = process_dynamic_image(&session, image)
-    .map_err(|err| {
-      error!("Couldn't remove background from image: {}", err);
-      format!("Couldn't remove background from image: {}", err)
-    })?;
-
-  let image = encode_dynamic_image_base64_png(image)
-    .map_err(|err| format!("failure to encode image: {:?}", err))?;
-
-  Ok(image)
-}
-
-fn process_dynamic_image(
-  session: &ort::Session,
-  dynamic_img: DynamicImage,
-) -> Result<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))
-}
@@ -6,8 +6,8 @@ use serde::Serialize;
 pub enum NotificationModelType {
  Unet,
  Vae,
-  TextEncoder,
-  TextTokenizer,
+  Json,
+  ClipEncoder,
 }

 #[derive(Clone, Serialize)]
@@ -4,11 +4,9 @@ pub mod ml;
 pub mod state;
 pub mod threads;
 pub mod transfer;
-pub mod utils;

 use crate::endpoints::download_models::download_models;
 use crate::endpoints::image_endpoint::infer_image;
-use crate::endpoints::remove_background_endpoint::remove_background;
 use crate::endpoints::test_counter::test_counter;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
@@ -62,10 +60,9 @@ pub fn run() {
    .manage(model_cache)
    .manage(app_data_root)
    .invoke_handler(tauri::generate_handler![
-      download_models,
-      infer_image,
-      remove_background,
      test_counter, 
+      infer_image,
+      download_models,
    ])
    .run(tauri::generate_context!())
    .expect("error while running tauri application");
@@ -1 +0,0 @@
-pub mod onnx_session;
@@ -1,18 +0,0 @@
-//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs
-
-use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
-
-pub fn onnx_session(onnx_model_file: &str) -> Result<ort::Session, anyhow::Error> {
-  let environment = Environment::default().into_arc();
-  let session = SessionBuilder::new(&environment)?
-    .with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level
-    .with_intra_threads(1)? // Configure the number of threads used for inference
-    .with_execution_providers([
-      // Configure execution providers (e.g., CUDA, CoreML, CPU)
-      ExecutionProvider::CUDA(Default::default()),
-      ExecutionProvider::CoreML(Default::default()),
-      ExecutionProvider::CPU(Default::default()),
-    ])?
-    .with_model_from_file(onnx_model_file)?;
-  Ok(session)
-}
@@ -0,0 +1,15 @@
+use crate::ml::downloads::download_model_registry_file::download_model_registry_file;
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppDataRoot;
+
+pub async fn download_all_models(app_data_root: &AppDataRoot) -> anyhow::Result<()> {
+  let weights_dir = app_data_root.weights_dir();
+  
+  //download_model_registry_file(ModelType::ClipJson, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?;
+  
+  Ok(())
+}
@@ -0,0 +1,30 @@
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppWeightsDir;
+use log::info;
+use std::fs::File;
+use std::path::{Path, PathBuf};
+
+pub async fn download_model_registry_file(
+  model_type: ModelType,
+  weights_dir: &AppWeightsDir,
+) -> anyhow::Result<PathBuf> {
+  
+  let path = weights_dir.model_path(&model_type);
+
+  if path.exists() {
+    return Ok(path);
+  }
+
+  info!("downloading model: {:?} to {:?}", model_type, path);
+
+  let body = reqwest::get(model_type.get_download_url())
+    .await?
+    .bytes()
+    .await?;
+  
+  let _file = File::create(&path)?;
+  
+  std::fs::write(&path, body.as_ref())?;
+  
+  Ok(path)
+}
@@ -0,0 +1,2 @@
+pub mod download_model_registry_file;
+pub mod download_all_models;
@@ -1,7 +1,8 @@
-pub mod background_removal;
 pub mod image;
 pub mod model_cache;
 pub mod model_file;
+pub mod models;
 pub mod prompt_cache;
 pub mod stable_diffusion;
-pub mod weights_registry;
+pub mod model_type;
+pub mod downloads;
@@ -1,14 +1,15 @@
+use crate::ml::models::unet_model::UNetModel;
 use anyhow::anyhow;
-use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
 use candle_transformers::models::stable_diffusion::vae::AutoEncoderKL;
 use std::sync::{Arc, RwLock};
+
 // TODO: This data structure is gross. Generalize the locking and loading semantics for reuse.

 // Simple registry for now. We can build complex machinery that aids in 
 // VRAM utilization, disk space saving, intelligent scheduling, etc. in the future.
 pub struct ModelCache {
  vae: RwLock<Option<Arc<AutoEncoderKL>>>,
-  unet2: RwLock<Option<Arc<UNet2DConditionModel>>>,
+  unet2: RwLock<Option<Arc<UNetModel>>>,
 }

 impl ModelCache {
@@ -38,7 +39,7 @@ impl ModelCache {
    }
  }

-  pub fn set_unet(&self, vae: Arc<UNet2DConditionModel>) -> anyhow::Result<()> {
+  pub fn set_unet(&self, vae: Arc<UNetModel>) -> anyhow::Result<()> {
    match self.unet2.write() {
      Err(err) => Err(anyhow!("{}", err)),
      Ok(mut lock) => {
@@ -48,7 +49,7 @@ impl ModelCache {
    }
  }

-  pub fn get_unet(&self) -> anyhow::Result<Option<Arc<UNet2DConditionModel>>> {
+  pub fn get_unet(&self) -> anyhow::Result<Option<Arc<UNetModel>>> {
    match self.unet2.read() {
      Err(err) => Err(anyhow!("{}", err)),
      Ok(ref lock) => {
@@ -0,0 +1,86 @@
+use crate::events::notification_event::NotificationModelType;
+use crate::state::app_dir::AppWeightsDir;
+use std::path::PathBuf;
+
+#[derive(Copy, Clone, Debug)]
+pub enum ModelType {
+  ClipJson,
+  SdxlTurboUnet,
+  SdxlTurboVae,
+  SdxlTurboClipEncoder,
+  SdxlTurboClipEncoder2,
+  SimianLuoLcmDreamshaperV7Unet,
+  LykonDreamshaper7Vae,
+  LykonDreamshaper7TextEncoderFp16,
+}
+
+impl ModelType {
+  pub fn get_name(&self) -> &'static str {
+    match self {
+      Self::ClipJson => "CLIP (JSON)",
+      Self::SdxlTurboUnet => "SDXL Turbo UNET",
+      Self::SdxlTurboVae => "SDXL Turbo VAE",
+      Self::SdxlTurboClipEncoder => "Clip Encoder #1",
+      Self::SdxlTurboClipEncoder2 => "Clip Encoder #2",
+      Self::SimianLuoLcmDreamshaperV7Unet => "SimianLuo LCM Dreamshaper v7 UNET",
+      Self::LykonDreamshaper7Vae => "Lykon Dreamshaper 7 VAE",
+      Self::LykonDreamshaper7TextEncoderFp16 => "Lykon Dreamshaper Text Encoder (fp16)",
+    }
+  }
+
+  pub fn get_notification_type(&self) -> NotificationModelType {
+    match self {
+      Self::ClipJson => NotificationModelType::Json,
+      Self::SdxlTurboUnet => NotificationModelType::Unet,
+      Self::SdxlTurboVae => NotificationModelType::Vae,
+      Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder,
+      Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder,
+      Self::SimianLuoLcmDreamshaperV7Unet => NotificationModelType::Unet,
+      Self::LykonDreamshaper7Vae => NotificationModelType::Vae,
+      Self::LykonDreamshaper7TextEncoderFp16 => NotificationModelType::ClipEncoder,
+    }
+  }
+
+  pub fn get_hf_id(&self) -> Option<&'static str> {
+    match self {
+      Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"),
+      Self::SdxlTurboUnet 
+      | Self::SdxlTurboVae
+      | Self::SdxlTurboClipEncoder
+      | Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"),
+      Self::SimianLuoLcmDreamshaperV7Unet => Some("SimianLuo/LCM_Dreamshaper_v7"),
+      Self::LykonDreamshaper7Vae
+      | Self::LykonDreamshaper7TextEncoderFp16 => Some("Lykon/dreamshaper-7"),
+    }
+  }
+
+  pub fn get_download_url(&self) -> &'static str {
+    match self {
+      Self::ClipJson => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/tokenizer.json",
+      Self::SdxlTurboUnet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.unet.safetensors",
+      Self::SdxlTurboVae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.vae.safetensors",
+      Self::SdxlTurboClipEncoder => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder.safetensors",
+      Self::SdxlTurboClipEncoder2 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder_2.safetensors",
+      Self::SimianLuoLcmDreamshaperV7Unet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/simianluo_lcm_dreamshaper_v7_unet.safetensors",
+      Self::LykonDreamshaper7Vae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_vae.safetensors",
+      Self::LykonDreamshaper7TextEncoderFp16 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors",
+    }
+  }
+
+  pub fn get_filename(&self) -> &'static str {
+    match self {
+      Self::ClipJson => "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
+      Self::SdxlTurboUnet => "diffusion_pytorch_model.unet.safetensors",
+      Self::SdxlTurboVae => "diffusion_pytorch_model.vae.safetensors",
+      Self::SdxlTurboClipEncoder => "clip_text_encoder.safetensors",
+      Self::SdxlTurboClipEncoder2 => "clip_text_encoder_2.safetensors",
+      Self::SimianLuoLcmDreamshaperV7Unet => "simianluo_lcm_dreamshaper_v7_unet.safetensors",
+      Self::LykonDreamshaper7Vae => "lykon_dreamshaper_7_vae.safetensors",
+      Self::LykonDreamshaper7TextEncoderFp16 => "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
+    }
+  }
+  
+  pub fn get_path(&self, weights_dir: &AppWeightsDir) -> PathBuf {
+    weights_dir.model_path(self)
+  }
+}
@@ -0,0 +1,130 @@
+use crate::ml::model_file::StableDiffusionVersion;
+use anyhow::anyhow;
+use candle_core::{DType, Device, Tensor};
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use hf_hub::api::sync::Api;
+use std::sync::{Arc, RwLock};
+
+pub struct LazyLoadVaeModel {
+  lazy_model: Arc<RwLock<Option<AutoEncoderKL>>>,
+  model_configs: VaeConfigs,
+}
+
+// TODO: Maybe we pass these via a central config read at startup, passed as an Arc<ConfigT>
+pub struct VaeConfigs {
+  sd_config: StableDiffusionConfig,
+  sd_version: StableDiffusionVersion,
+  device: Device,
+  dtype: DType,
+  hf_api: Api,
+}
+
+impl LazyLoadVaeModel {
+  pub fn new(
+    sd_config: &StableDiffusionConfig,
+    sd_version: &StableDiffusionVersion,
+    hf_api: &Api,
+    device: &Device,
+    dtype: DType,
+  ) -> anyhow::Result<Self> {
+    println!("building vae model... (1)");
+    
+    Ok(Self {
+      lazy_model: Arc::new(RwLock::new(None)),
+      model_configs: VaeConfigs {
+        sd_config: sd_config.clone(),
+        sd_version: sd_version.clone(),
+        device: device.clone(),
+        hf_api: hf_api.clone(),
+        dtype
+      }
+    })
+  }
+
+  pub fn encode(&self, xs: &Tensor) -> anyhow::Result<DiagonalGaussianDistribution> {
+    println!("vae encode...");
+    match self.try_encode(xs) {
+      Err(err) => Err(err),
+      Ok(Some(result)) => Ok(result),
+      Ok(None) => {
+        // Model wasn't previously loaded
+        self.load_model()?;
+        self.try_encode(xs)?
+          .ok_or(anyhow!("model was not previously loaded"))
+      }
+    }
+  }
+  
+  pub fn decode(&self, xs: &Tensor) -> anyhow::Result<Tensor> {
+    println!("vae decode...");
+    match self.try_decode(xs) {
+      Err(err) => Err(err),
+      Ok(Some(result)) => Ok(result),
+      Ok(None) => {
+        // Model wasn't previously loaded
+        self.load_model()?;
+        self.try_decode(xs)?
+          .ok_or(anyhow!("model was not previously loaded"))
+      }
+    }
+  }
+
+  fn try_encode(&self, xs: &Tensor) -> anyhow::Result<Option<DiagonalGaussianDistribution>> {
+    match self.lazy_model.read() {
+      Err(err) => Err(anyhow!("lock error: {:?}", err)),
+      Ok(model) => {
+        match &*model {
+          None => Ok(None), // Model is not yet loaded
+          Some(model) => {
+            Ok(Some(model.encode(xs)?))
+          }
+        }
+      }
+    }
+  }
+
+  fn try_decode(&self, xs: &Tensor) -> anyhow::Result<Option<Tensor>> {
+    match self.lazy_model.read() {
+      Err(err) => Err(anyhow!("lock error: {:?}", err)),
+      Ok(model) => {
+        match &*model {
+          None => Ok(None), // Model is not yet loaded
+          Some(model) => {
+            Ok(Some(model.decode(xs)?))
+          }
+        }
+      }
+    }
+  }
+
+  fn load_model(&self) -> anyhow::Result<()> {
+    println!("load vae");
+    match self.lazy_model.write() {
+      Err(err) => return Err(anyhow!("lock error: {:?}", err)),
+      Ok(mut maybe_model) => {
+        match &*maybe_model {
+          Some(_model) => {} // Already loaded; fall through
+          None => {
+            let repo = self.model_configs.sd_version.repo();
+
+            println!("Building VAE model from : {:?} ... (3)", repo);
+            
+            let vae_file = self.model_configs.hf_api.model(repo.to_string())
+              .get("vae/diffusion_pytorch_model.safetensors")?;
+
+            println!("Building VAE model from file {:?}...", &vae_file);
+            let vae = self.model_configs
+              .sd_config
+              .build_vae(vae_file, &self.model_configs.device, self.model_configs.dtype)?;
+
+            println!("Built VAE...");
+            *maybe_model = Some(vae);
+          }
+        }
+      }
+    }
+    Ok(())
+  }
+ 
+}
@@ -0,0 +1,2 @@
+pub mod unet_model;
+pub mod lazy_load_vae_model;
@@ -0,0 +1,30 @@
+use candle_core::{DType, Device, Tensor};
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use std::path::Path;
+
+pub struct UNetModel {
+  model: UNet2DConditionModel,
+}
+
+impl UNetModel {
+  pub fn new<P: AsRef<Path>>(sd_config: &StableDiffusionConfig, unet_weights_file: P, device: &Device, dtype: DType) -> anyhow::Result<Self> {
+    println!("building unet model... (1)");
+
+    let model = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)?;
+
+    Ok(Self { model })
+  }
+
+  pub fn new_with_inner(model: UNet2DConditionModel) -> Self {
+    Self { model }
+  }
+
+  pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
+    Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?)
+  }
+
+  pub fn forward_with_guidance(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor, guidance_scale_embedding: &Tensor) -> anyhow::Result<Tensor> {
+    Ok(self.model.forward_with_guidance(&latent_model_input, timestep, &text_embeddings, Some(guidance_scale_embedding))?)
+  }
+}
@@ -3,16 +3,16 @@ extern crate accelerate_src;
 #[cfg(feature = "mkl")]
 extern crate intel_mkl_src;

-use candle_transformers::models::stable_diffusion;
 use std::path::PathBuf;
+use candle_transformers::models::stable_diffusion;

 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;
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppWeightsDir;

 #[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> {
@@ -53,7 +53,17 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
 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> {
  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_file = if first {
+  //  info!("ModelFile::Tokenizer");
+  //  ModelFile::Tokenizer
+  //} else {
+  //  info!("ModelFile::Tokenizer2");
+  //  ModelFile::Tokenizer2
+  //};
+
+  //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
+
+  let tokenizer = weights_dir.model_path(&ModelType::ClipJson);

  info!("Loading Clip Tokenizer path: {:?}", tokenizer);

@@ -75,10 +85,13 @@ 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.");
+  let clip_weights_file = if first { ModelFile::Clip } else { ModelFile::Clip2 };
  let clip_weights = if first {
-    weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER)
+    //clip_weights_file.get(clip_weights, sd_version, use_f16)?
+    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder)
  } else {
-    weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER_2)
+    //clip_weights_file.get(clip2_weights, sd_version, use_f16)?
+    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder2)
  };
  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)?;
@@ -6,11 +6,11 @@ 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::{ModelFile, StableDiffusionVersion};
+use crate::ml::model_type::ModelType;
+use crate::ml::models::unet_model::UNetModel;
 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::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};
@@ -26,14 +26,12 @@ use log::info;
 use rand::Rng;
 use tauri::{AppHandle, Emitter};

-const DEFAULT_STRENGTH : f64 = 75.0;
-
 pub struct Args<'a> {
  pub image: &'a DynamicImage,
  pub prompt: String,
  pub uncond_prompt: String,
  pub cfg_scale: Option<f64>,
-  pub i2i_strength: Option<f64>,
+  pub i2i_strength: Option<u8>,
  pub configs: &'a AppConfig,
  pub model_cache: &'a ModelCache,
  pub prompt_cache: &'a PromptCache,
@@ -46,21 +44,26 @@ 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();
-  
-  let img2img_strength = args.i2i_strength
-    .map(|s| remap_lcm_strength_range(s))
-    .unwrap_or(DEFAULT_STRENGTH);
-  
-  let img2img_strength = img2img_strength / 100.0;

+  let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0;
  let use_f16 = true;

+  //println!("Starting image generation with the following configuration:");
+  //println!("  Model: {:?}", configs.sd_version);
+  //println!("  Prompt: {}", prompt);
+  //println!("  Steps: {}", configs.scheduler_steps);
+  //println!("  Device: {:?}", configs.device);
+
+  //println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
+
  // 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 seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
  configs.device.set_seed(seed)?;

+  //info!("Using seed: {}", seed);
+
  let guidance_scale = match cfg_scale {
    Some(guidance_scale) => guidance_scale,
    None => match configs.sd_version {
@@ -73,14 +76,26 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
  // let use_guide_scale = guidance_scale > 1.0;
  let use_guide_scale = false;

+  //info!("Using guide scale: {}", use_guide_scale);
+
+  //info!("Checking if prompt is cached");
  let maybe_cached = prompt_cache.get_copy(&prompt)?;

  let mut text_embeddings = if let Some(tensor) = maybe_cached {
    tensor
  } else {
    info!("Loading clip weights...");
+    //println!("Loading clip weights for {:?} with model {:?}", configs.sd_version, configs.sd_version.repo());
+    //let api = configs.hf_api.clone();

-    let clip_weights = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
+    //let clip_path = api.model(configs.sd_version.repo().to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
+
+    //println!(">>>> CLIP REPO = {:?}", configs.sd_version.repo());
+    //println!(">>>> CLIP PATH = {:?}", clip_path);
+    //println!(">>>> USE_FP16 = {:?}", use_f16);
+    //let clip_weights = Some(clip_path.to_string_lossy().to_string());
+
+    let clip_weights = weights_dir.model_path(&ModelType::LykonDreamshaper7TextEncoderFp16);
    let clip_weights = Some(clip_weights.to_string_lossy().to_string());

    info!("Prompt is NOT cached! Calculating embedding...");
@@ -117,7 +132,12 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
  let vae = match maybe_vae {
    Some(vae) => vae,
    None => {
-      let vae_file = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
+      //let repo = configs.sd_version.repo();
+      //println!(">>>>> VAE REPO = {:?}", repo);
+      //println!("Building VAE model from : {:?} ... (3)", repo);
+      //let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
+
+      let vae_file = weights_dir.model_path(&ModelType::LykonDreamshaper7Vae);

      println!("Building VAE model from file {:?}...", &vae_file);

@@ -138,7 +158,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 unet_weights = weights_dir.model_path(&ModelType::SimianLuoLcmDreamshaperV7Unet);

      let in_channels = match configs.sd_version {
        StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
@@ -179,6 +199,8 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {

      let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;

+      let unet = UNetModel::new_with_inner(unet);
+
      let unet = Arc::new(unet);

      model_cache.set_unet(unet.clone())?;
@@ -189,6 +211,12 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {

  let init_latent_dist: DiagonalGaussianDistribution = vae.encode(&input_image)?;

+  // Parse and calculate the img2img strength
+  let img2img_strength = match i2i_strength {
+    None => 0.75f64,
+    Some(strength) => (strength as f64) / 100.0f64,
+  };
+
  // 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)?;
@@ -207,6 +235,11 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
  let noise = init_latents.randn_like(0f64, 1f64)?;
  let latents = scheduler.add_noise(&init_latents, noise, t_start)?;

+  println!("Latents initialized successfully");
+  println!("Initial noise shape: {:?}", latents.shape());
+
+  println!("Starting diffusion process...");
+
  let mut latents = latents;
  let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();

@@ -218,7 +251,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
    let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);

    // Use the guidance scale embedding in the UNet inference with the proper method name
-    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, Some(&guidance_scale_embedding)) {
+    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, &guidance_scale_embedding) {
      Ok(pred) => pred,
      Err(e) => {
        println!("UNet inference failed with error: {}", e);
@@ -1,5 +1,4 @@
 pub mod get_vae_scale;
 pub mod infer_clip_text_embeddings;
 pub mod stable_diffusion_pipeline;
-pub mod lcm_pipeline;
-mod remap_lcm_strength_range;
+pub mod lcm_pipeline;
@@ -1,54 +0,0 @@
-use primitives::numerics::remap_range_f64::{remap_range_f64, Args};
-
-/// Our LCM implementation uses strength [13, 100]. Strength below 13 will take 8 steps, 
-/// so we need to limit ourselves to this range from an input range of [0, 100].
-pub fn remap_lcm_strength_range(value: f64) -> f64 {
-  remap_range_f64(Args {
-    old_min: 0.0,
-    old_max: 100.0,
-    new_min: 13.0,
-    new_max: 100.0,
-    value,
-  })
-}
-
-#[cfg(test)]
-mod test {
-  use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
-  use assertor::{assert_that, FloatAssertion};
-
-  #[test]
-  fn old_max() {
-    assert_eq!(remap_lcm_strength_range(100.0), 100.0);
-  }
-  
-  #[test]
-  fn old_min() {
-    assert_eq!(remap_lcm_strength_range(0.0), 13.0);
-  }
-
-  #[test]
-  fn under_bounds() {
-    assert_eq!(remap_lcm_strength_range(-100.0), 13.0);
-  }
-  
-  #[test]
-  fn over_bounds() {
-    assert_eq!(remap_lcm_strength_range(200.0), 100.0);
-  }
-
-  #[test]
-  fn other_values() {
-    assert_that!(remap_lcm_strength_range(0.0)).with_abs_tol(0.001).is_approx_equal_to(13.000);
-    assert_that!(remap_lcm_strength_range(5.0)).with_abs_tol(0.001).is_approx_equal_to(17.350);
-    assert_that!(remap_lcm_strength_range(11.0)).with_abs_tol(0.001).is_approx_equal_to(22.570);
-    assert_that!(remap_lcm_strength_range(13.0)).with_abs_tol(0.001).is_approx_equal_to(24.310);
-    assert_that!(remap_lcm_strength_range(14.0)).with_abs_tol(0.001).is_approx_equal_to(25.180);
-    assert_that!(remap_lcm_strength_range(20.0)).with_abs_tol(0.001).is_approx_equal_to(30.400);
-    assert_that!(remap_lcm_strength_range(25.0)).with_abs_tol(0.001).is_approx_equal_to(34.750);
-    assert_that!(remap_lcm_strength_range(50.0)).with_abs_tol(0.001).is_approx_equal_to(56.500);
-    assert_that!(remap_lcm_strength_range(75.0)).with_abs_tol(0.001).is_approx_equal_to(78.250);
-    assert_that!(remap_lcm_strength_range(99.0)).with_abs_tol(0.001).is_approx_equal_to(99.130);
-    assert_that!(remap_lcm_strength_range(100.0)).with_abs_tol(0.001).is_approx_equal_to(100.000);
-  }
-}
@@ -1,16 +1,16 @@
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

-use crate::events::notification_event::{NotificationEvent, NotificationModelType};
+use crate::events::notification_event::{NotificationModelType, NotificationEvent};
 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::model_type::ModelType;
+use crate::ml::models::unet_model::UNetModel;
 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 anyhow::{anyhow, Error as E, Result};
 use candle_core::{DType, IndexOp, Tensor, D};
 use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
@@ -18,6 +18,7 @@ use image::DynamicImage;
 use log::info;
 use rand::Rng;
 use tauri::{AppHandle, Emitter};
+use crate::state::app_dir::AppDataRoot;

 pub struct Args<'a> {
    pub image: &'a DynamicImage,
@@ -128,7 +129,17 @@ 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 mut notify_download_complete = false;
+            //if !vae_file.exists() {
+            if true {
+                notify_download_complete = true;
+                app.emit("notification", NotificationEvent::ModelDownloadStarted {
+                    model_name: "sdxl-turbo-vae",
+                    model_type: NotificationModelType::Vae,
+                })?;
+            }
+            
+            let vae_file = weights_dir.model_path(&ModelType::SdxlTurboVae);
            
            let vae = configs
              .sd_config
@@ -137,7 +148,14 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
            let vae = Arc::new(vae);
            
            model_cache.set_vae(vae.clone())?;
-
+            
+            if notify_download_complete {
+                app.emit("notification", NotificationEvent::ModelDownloadComplete {
+                    model_name: "sdxl-turbo-vae",
+                    model_type: NotificationModelType::Vae,
+                })?;
+            }
+            
            vae
        }
    };
@@ -149,15 +167,32 @@ 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 mut notify_download_complete = false;
+            //if !unet_file.exists() {
+            if true {
+                notify_download_complete = true;
+                app.emit("notification", NotificationEvent::ModelDownloadStarted {
+                    model_name: "sdxl-turbo-unet",
+                    model_type: NotificationModelType::Unet,
+                })?;
+            }
+            
+            let unet_file = weights_dir.model_path(&ModelType::SdxlTurboUnet);

-            let unet = configs.sd_config.build_unet(unet_weights_file, &configs.device, 4, true, configs.dtype)
+            let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
              .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
-
+            
            let unet = Arc::new(unet);

            model_cache.set_unet(unet.clone())?;

+            if notify_download_complete {
+                app.emit("notification", NotificationEvent::ModelDownloadComplete {
+                    model_name: "sdxl-turbo-unet",
+                    model_type: NotificationModelType::Unet,
+                })?;
+            }
+
            unet
        }
    };
@@ -222,7 +257,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {

        let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;

-        let mut noise_pred = match unet.forward(&latent_model_input, timestep as f64, &text_embeddings) {
+        let mut noise_pred = match unet.inference(&latent_model_input, timestep as f64, &text_embeddings) {
            Ok(pred) => pred,
            Err(e) => {
                println!("UNet inference failed with error: {}", e);
@@ -1,3 +0,0 @@
-mod weight_descriptor_builder;
-pub mod weight_descriptor;
-pub mod weights;
@@ -1,43 +0,0 @@
-
-// TODO(bt,2025-03-13): When it's time for dynamic, non-const allocation, refer to
-//  http::header::name::HeaderName: https://github.com/hyperium/http/blob/master/src/header/name.rs
-//  as a pattern that will support both compile-time constants and dynamically allocated String-based variants.
-
-#[derive(Clone, Copy)]
-pub enum WeightFunction {
-  TextEncoder,
-  TextTokenizer,
-  Unet,
-  Vae,
-}
-
-pub struct WeightDescriptor {
-  /// Name of the weight, eg. "Lykon Dreamshaper 7"
-  pub name: &'static str,
-
-  /// Name of the file on the filesystem, eg. lykon_dreamshaper_7_text_encoder.fp16.safetensors
-  pub filename: &'static str,
-
-  /// URL of the file in the R2 bucket,
-  /// eg. https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors
-  pub r2_download_url: &'static str,
-
-  pub function: WeightFunction,
-}
-
-impl WeightFunction {
-  pub fn as_descriptive_str(&self) -> &'static str {
-    match self {
-      WeightFunction::TextEncoder => "text encoder",
-      WeightFunction::TextTokenizer => "text tokenizer",
-      WeightFunction::Unet => "UNET",
-      WeightFunction::Vae => "VAE",
-    }
-  }
-}
-
-impl WeightDescriptor {
-  pub fn to_descriptive_name(&self) -> String {
-    format!("{} ({})", self.name, self.function.as_descriptive_str())
-  }
-}
@@ -1,17 +0,0 @@
-pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/";
-
-// TODO(bt,2025-03-13): Try to use keyword args from the call site.
-// NB: We're using a macro because we can't use a builder pattern and then concat string literals
-//  and variables (even if they're known to be &'static) and keep allocations static.
-macro_rules! weight {
-    ($name: literal, $file_name:literal, $weight_function:expr) => {
-      $crate::ml::weights_registry::weight_descriptor::WeightDescriptor {
-        name: $name,
-        filename: $file_name,
-        r2_download_url: const_format::concatcp!(crate::ml::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
-        function: $weight_function,
-      }
-    };
-}
-
-pub(crate) use weight;
@@ -1,57 +0,0 @@
-use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
-use crate::ml::weights_registry::weight_descriptor_builder::weight;
-
-pub const CLIP_JSON : WeightDescriptor = weight!(
-  "CLIP",
-  "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
-  WeightFunction::TextTokenizer
-);
-
-/// For LCM-based image generation
-pub const LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16 : WeightDescriptor = weight!(
-  "Lykon Dreamshaper 7",
-  "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
-  WeightFunction::TextEncoder
-);
-
-/// For LCM-based image generation
-pub const LYKON_DEAMSHAPER_7_VAE : WeightDescriptor = weight!(
-  "Lykon Dreamshaper 7",
-  "lykon_dreamshaper_7_vae.safetensors",
-  WeightFunction::Vae
-);
-
-/// For LCM-based image generation
-pub const SIMIANLUO_LCM_DREAMSHAPER_V7_UNET : WeightDescriptor = weight!(
-  "SimianLuo LCM Dreamshaper v7",
-  "simianluo_lcm_dreamshaper_v7_unet.safetensors",
-  WeightFunction::Unet
-);
-
-#[deprecated]
-pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = weight!(
-  "SDXL Turbo CLIP",
-  "clip_text_encoder.safetensors",
-  WeightFunction::TextEncoder
-);
-
-#[deprecated]
-pub const SDXL_TURBO_CLIP_TEXT_ENCODER_2 : WeightDescriptor = weight!(
-  "SDXL Turbo CLIP #2",
-  "clip_text_encoder_2.safetensors",
-  WeightFunction::TextEncoder
-);
-
-#[deprecated]
-pub const SDXL_TURBO_VAE : WeightDescriptor = weight!(
-  "SDXL Turbo",
-  "diffusion_pytorch_model.vae.safetensors",
-  WeightFunction::Vae
-);
-
-#[deprecated]
-pub const SDXL_TURBO_UNET : WeightDescriptor = weight!(
-  "SDXL Turbo",
-  "diffusion_pytorch_model.unet.safetensors",
-  WeightFunction::Unet
-);
@@ -1,4 +1,4 @@
-use crate::ml::weights_registry::weight_descriptor::WeightDescriptor;
+use crate::ml::model_type::ModelType;
 use crate::state::expanduser::expanduser;
 use crate::state::os_platform::OsPlatform;
 use anyhow::anyhow;
@@ -6,7 +6,6 @@ use directories::UserDirs;
 use std::path::{Path, PathBuf};
 use tempdir::TempDir;
 use tempfile::{Builder, NamedTempFile};
-
 const DEFAULT_DATA_DIR : &str = "artcraft";
 const ASSETS_SUBDIRECTORY : &str = "assets";
 const WEIGHTS_SUBDIRECTORY : &str = "weights";
@@ -140,9 +139,9 @@ impl AppWeightsDir {
  pub fn path(&self) -> &Path {
    &self.path
  }
-
-  pub fn weight_path(&self, descriptor: &WeightDescriptor) -> PathBuf {
-    self.path.join(descriptor.filename)
+  
+  pub fn model_path(&self, model_registry: &ModelType) -> PathBuf {
+    self.path.join(model_registry.get_filename())
  }
 }

@@ -1,4 +1,5 @@
 use crate::events::notification_event::{NotificationEvent, NotificationModelType};
+use crate::ml::model_type::ModelType;
 use crate::state::app_dir::AppDataRoot;
 use crate::transfer::download::{download_async, ProgressUpdate};
 use futures::stream::FuturesUnordered;
@@ -11,8 +12,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, 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;
@@ -40,21 +39,22 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !

  let mut download_queue = VecDeque::new();

-  download_queue.push_back(CLIP_JSON);
-  download_queue.push_back(LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
-  download_queue.push_back(LYKON_DEAMSHAPER_7_VAE);
-  download_queue.push_back(SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
-  download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER); // TODO(bt): Why is this still needed?
-  //download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER_2); // TODO(bt): REMOVE
-  //download_queue.push_back(SDXL_TURBO_VAE); // TODO(bt): REMOVE
-  //download_queue.push_back(SDXL_TURBO_UNET); // TODO(bt): REMOVE
+  // TODO: Automatic enqueue of known important models + enqueue new models on-demand
+  download_queue.push_back(ModelType::ClipJson);
+  //download_queue.push_back(ModelType::SdxlTurboUnet);
+  //download_queue.push_back(ModelType::SdxlTurboVae);
+  download_queue.push_back(ModelType::SdxlTurboClipEncoder); // TODO(bt): Why is this still needed?
+  //download_queue.push_back(ModelType::SdxlTurboClipEncoder2);
+  download_queue.push_back(ModelType::SimianLuoLcmDreamshaperV7Unet);
+  download_queue.push_back(ModelType::LykonDreamshaper7Vae);
+  download_queue.push_back(ModelType::LykonDreamshaper7TextEncoderFp16);

  let mut handles = FuturesUnordered::new();

  while let Some(model) = download_queue.pop_front() {
-    let url = model.r2_download_url.to_string();
+    let url = model.get_download_url().to_string();

-    let filename = app_data_root.weights_dir().weight_path(&model);
+    let filename = app_data_root.weights_dir().model_path(&model);

    if filename.exists() {
      info!("File already exists: {:?}", &filename);
@@ -75,16 +75,9 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
      }
    };

-    let model_type = match model.function {
-      WeightFunction::TextEncoder => NotificationModelType::TextEncoder,
-      WeightFunction::TextTokenizer => NotificationModelType::TextTokenizer,
-      WeightFunction::Unet => NotificationModelType::Unet,
-      WeightFunction::Vae => NotificationModelType::Vae,
-    };
-
    let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
-      model_name: model.name,
-      model_type,
+      model_name: model.get_name(),
+      model_type: model.get_notification_type(),
    });

    if let Err(err) = result {
@@ -109,8 +102,8 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
        lock.push(ModelDownloadStatusChannel {
          receiver: rx,
          final_file_path: filename.clone(),
-          model_name: model.to_descriptive_name(),
-          model_type,
+          model_name: model.get_name().to_string(),
+          model_type: model.get_notification_type(),
        })
      }
    }
@@ -1,14 +0,0 @@
-use base64::prelude::BASE64_STANDARD;
-use base64::Engine;
-use image::{DynamicImage, ImageReader};
-use std::io::Cursor;
-
-pub fn decode_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
-  let bytes = BASE64_STANDARD.decode(base64_image)?;
-
-  let image = ImageReader::new(Cursor::new(bytes))
-    .with_guessed_format()?
-    .decode()?;
-
-  Ok(image)
-}
@@ -1,14 +0,0 @@
-use base64::prelude::BASE64_STANDARD;
-use base64::Engine;
-use image::{DynamicImage, ImageFormat};
-use std::io::Cursor;
-
-pub fn encode_dynamic_image_base64_png(image: DynamicImage) -> anyhow::Result<String> {
-  let mut bytes = Vec::with_capacity(1024*1024);
-
-  image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
-
-  let bytes = BASE64_STANDARD.encode(bytes);
-
-  Ok(bytes)
-}
@@ -1,17 +0,0 @@
-use crate::ml::image::tensor_to_image_buffer::RgbImage;
-use base64::prelude::BASE64_STANDARD;
-use base64::Engine;
-use image::ImageFormat;
-use std::io::Cursor;
-
-pub fn encode_rgb_image_base64_png(
-  image: RgbImage
-) -> anyhow::Result<String> {
-  let mut bytes = Vec::with_capacity(1024*1024);
-
-  image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
-
-  let bytes = BASE64_STANDARD.encode(bytes);
-
-  Ok(bytes)
-}
@@ -1,3 +0,0 @@
-pub mod decode_base64_image;
-pub mod encode_dynamic_image_base64_png;
-pub mod encode_rgb_image_base64_png;
@@ -1 +0,0 @@
-pub mod image;
@@ -32,9 +32,6 @@
      "icons/128x128@2x.png",
      "icons/icon.icns",
      "icons/icon.ico"
-    ],
-    "resources": {
-      "../../../windows_dlls/*.dll": "."
-    }
+    ]
  }
 }
@@ -32,4 +32,3 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" }

 [dev-dependencies]
 testing = { path = "../../testing" }
-assertor = "0.0.3"
@@ -1,6 +1,5 @@
 pub mod i32_to_u32_zero_clamped;
 pub mod i64_to_u64_zero_clamped;
-pub mod remap_range_f64;
 pub mod u64_to_i32_saturating;
 pub mod u64_to_i64_saturating;
 pub mod u64_to_u32_saturating;
@@ -1,131 +0,0 @@
-
-pub struct Args {
-  pub old_min: f64,
-  pub old_max: f64,
-  pub new_min: f64,
-  pub new_max: f64,
-  
-  /// The value to remap to the new range.
-  pub value: f64,
-}
-
-pub fn remap_range_f64(args: Args) -> f64 {
-  let old_range = args.old_max - args.old_min;
-  let new_range = args.new_max -  args.new_min;
-
-  let new_value = (args.value -  args.old_min) * new_range / old_range + args.new_min;
-  
-  if new_value > args.new_max {
-    args.new_max
-  } else if new_value < args.new_min { 
-    args.new_min
-  } else {
-    new_value
-  }
-}
-
-#[cfg(test)]
-mod tests {
-  use crate::numerics::remap_range_f64::{remap_range_f64, Args};
-  use assertor::{assert_that, FloatAssertion};
-
-  #[test]
-  fn test_base_cases() {
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: 0.0,
-      new_max: 1.0,
-      value: 0.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(0.0);
-    
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: 0.0,
-      new_max: 1.0,
-      value: 1.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(1.0);
-  }
-
-  #[test]
-  fn test_maxima() {
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 100.0,
-      new_min: 0.0,
-      new_max: 1.0,
-      value: 100.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(1.0);
-
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: 0.0,
-      new_max: 50.0,
-      value: 1.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(50.0);
-
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: -100.0,
-      new_max: 100.0,
-      value: 1.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(100.0);
-  }
-
-  #[test]
-  fn test_minima() {
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 100.0,
-      new_min: 0.0,
-      new_max: 1.0,
-      value: 0.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(0.0);
-
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: 0.0,
-      new_max: 50.0,
-      value: 0.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(0.0);
-
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: -100.0,
-      new_max: 100.0,
-      value: 0.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(-100.0);
-  }
-
-  #[test]
-  fn test_other_values() {
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 100.0,
-      new_min: 0.0,
-      new_max: 1.0,
-      value: 50.0,
-    })).with_abs_tol(0.01).is_approx_equal_to(0.50);
-
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: 0.0,
-      new_max: 50.0,
-      value: 0.10,
-    })).with_abs_tol(0.01).is_approx_equal_to(5.0);
-
-    assert_that!(remap_range_f64(Args {
-      old_min: 0.0,
-      old_max: 1.0,
-      new_min: -100.0,
-      new_max: 100.0,
-      value: 0.25,
-    })).with_abs_tol(0.01).is_approx_equal_to(-50.0);
-  }
-}
@@ -1,36 +1,35 @@
 import Konva from "konva";

-import {uiAccess, uiEvents} from "~/signals";
-import {ImageNode, ShapeNode, ShapeType, TextNode, VideoNode} from "./Nodes";
-import {PaintNode} from "./Nodes/PaintNode";
-import {PreviewCopyNode} from "./Nodes/PreviewCopy";
-import {UndoStackManager} from "./UndoRedo";
-import {CommandManager, MatteBox, SceneManager} from "./EngineUtitlities";
+import { uiAccess, uiEvents } from "~/signals";
+import { ShapeNode } from "./Nodes";
+import { PaintNode } from "./Nodes/PaintNode";
+import { PreviewCopyNode } from "./Nodes/PreviewCopy";
+import { UndoStackManager } from "./UndoRedo";
+import { CommandManager, MatteBox, SceneManager } from "./EngineUtitlities";
 import {
-  NodeIsolator,
  NodesManager,
+  NodeIsolator,
+  NodeTransformer,
  NodesTranslationEventDetails,
  NodeTransformationEventDetails,
-  NodeTransformer,
  SelectionManager,
  SelectionManagerEvents,
  SelectorSquare,
 } from "./NodesManagers";
-import {EngineOptions, TextNodeData, VideoNodeData} from "./types";
+import { ImageNode, VideoNode, TextNode } from "./Nodes";
+import { EngineOptions, TextNodeData, VideoNodeData } from "./types";

-import {AppModes, VideoResolutions} from "./constants";
-import {ToolbarMainButtonNames} from "~/components/features/ToolbarMain/enum";
+import { AppModes, VideoResolutions } from "./constants";
+import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum";

-import {ToolbarNodeButtonNames} from "~/components/features/ToolbarNode/enums";
-import {NavigateFunction} from "react-router-dom";
-import {LoadingVideosProvider} from "./EngineUtitlities/LoadingVideosProvider";
+import { ToolbarNodeButtonNames } from "~/components/features/ToolbarNode/enums";
+import { NavigateFunction } from "react-router-dom";
+import { LoadingVideosProvider } from "./EngineUtitlities/LoadingVideosProvider";

-import {VideoExtractionHandler} from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler";
-import {RealTimeDrawEngine} from "./RenderingPrimitives/RealTimeDrawEngine";
-import {NodeColor} from "~/signals/uiEvents/toolbarNode";
-import {invoke} from "@tauri-apps/api/core";
-import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts";
-import {EncodeImageBitmapToBase64} from "~/utilities/EncodeImageBitmapToBase64.ts";
+import { VideoExtractionHandler } from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler";
+import { RealTimeDrawEngine } from "./RenderingPrimitives/RealTimeDrawEngine";
+import { NodeColor } from "~/signals/uiEvents/toolbarNode";
+import { ShapeType } from "./Nodes";


 export interface RenderingOptions {
@@ -327,9 +326,6 @@ export class Engine {
    uiEvents.toolbarNode.MOVE_LAYER_UP.onClick(() =>
      this.commandManager.moveNodesUp(),
    );
-    uiEvents.toolbarNode.REMOVE_BACKGROUND.onClick(() =>
-      this.removeBackground(),
-    );

    // Listen to Toolbar Main
    uiEvents.toolbarMain.UNDO.onClick(() => {
@@ -449,38 +445,6 @@ export class Engine {
    });
  }

-  removeBackground() {
-    const selectedNodeSet = this.selectionManager.getSelectedNodes();
-
-    if (selectedNodeSet.size !== 1) {
-      return; // Can only remove background from a single image at a time.
-    }
-
-    const [ selectedNode ] = selectedNodeSet;
-
-    const removeBgInvocation
-        = async (selectedNode: ImageNode, callback: (existingImageNode: ImageNode, createdImage: ImageBitmap) => void) => {
-
-      const image : HTMLImageElement = selectedNode.kNode.image();
-      const base64Bitmap = await EncodeImageBitmapToBase64(image);
-      const base64BitmapResponse = await invoke("remove_background", {
-        image: base64Bitmap,
-      });
-
-      const createdImage = await DecodeBase64ToImage(base64BitmapResponse as string);
-
-      callback(selectedNode, createdImage);
-    };
-
-    const _promise = removeBgInvocation(selectedNode, (selectedNode: ImageNode, createdImage: ImageBitmap) => {
-      if (!this.nodesManager.hasNode(selectedNode)) {
-        return;
-      }
-      this.commandManager.deleteSpecificNodes([selectedNode]);
-      this.addImageFromImageBitmap(createdImage);
-    })
-  }
-
  disableAllButtons() {
    const buttonNames = Object.values(ToolbarNodeButtonNames);
    for (const name of buttonNames) {
@@ -660,23 +624,6 @@ export class Engine {

    this.commandManager.createNode(imageNode);
  }
-
-  public addImageFromImageBitmap(imageBitmap: ImageBitmap) {
-    const imageNode = new ImageNode({
-      mediaLayerRef: this.mediaLayer,
-      canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
-      canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
-      imageBitmap: imageBitmap,
-      selectionManagerRef: this.selectionManager,
-      loaded: async () => {
-        await this.realTimeDrawEngine.render();
-        this.setAppMode(AppModes.SELECT);
-      },
-    });
-
-    this.commandManager.createNode(imageNode);
-  }
-
  public addPaintNode(
    canvas: HTMLCanvasElement,
    lineBounds: {
@@ -78,20 +78,6 @@ export class CommandManager {
    this.undoStackManagerRef.executeCommand(command);
    this.renderEngineRef.render();
  }
-
-  deleteSpecificNodes(nodes: MediaNode[]) {
-    const command = new DeleteCommand({
-      nodes: nodes,
-      mediaLayerRef: this.mediaLayerRef,
-      nodesManagerRef: this.nodesManagerRef,
-      nodeTransformerRef: this.nodeTransformerRef,
-      selectionManagerRef: this.selectionManagerRef,
-      renderEngineRef: this.renderEngineRef,
-    });
-    this.undoStackManagerRef.executeCommand(command);
-    this.renderEngineRef.render();
-  }
-
  toggleLockNodes() {
    const nodes = this.selectionManagerRef.getSelectedNodes();
    const node = nodes.values().next().value;
@@ -12,7 +12,6 @@ interface ImageNodeContructor {
  canvasPosition: Position;
  canvasSize: Size;
  imageFile?: File;
-  imageBitmap?: ImageBitmap;
  mediaFileToken?: string;
  mediaFileUrl?: string;
  transform?: TransformationData;
@@ -35,7 +34,6 @@ export class ImageNode extends NetworkedNode {
    mediaLayerRef,
    selectionManagerRef,
    loaded,
-    imageBitmap,
  }: ImageNodeContructor) {
    // kNodes need to be created first to guaruntee
    // that it is not undefined in parent's context
@@ -75,14 +73,6 @@ export class ImageNode extends NetworkedNode {
      this.loadImageFromUrl(mediaFileUrl);
      return;
    }
-    if (imageBitmap) {
-      this.loadImageFromImageBitmap({
-        imageBitmap: imageBitmap,
-        maxSize: canvasSize,
-        refPosition: canvasPosition,
-      });
-      return;
-    }
    console.log("image node creation is fucked");
  }
  private loadImageFromFile({
@@ -146,40 +136,6 @@ export class ImageNode extends NetworkedNode {
    this.listenToBaseKNode();
  }

-  private loadImageFromImageBitmap({
-    imageBitmap,
-    maxSize,
-    refPosition,
-  }: {
-    imageBitmap: ImageBitmap;
-    maxSize: Size;
-    refPosition: Position;
-  }) {
-    this.imageSize = {
-      width: imageBitmap.width,
-      height: imageBitmap.height,
-    };
-    const adjustedSize = NodeUtilities.adjustNodeSizeToCanvas({
-      componentSize: this.imageSize,
-      maxSize: maxSize, // TODO
-    });
-    const centerPosition = NodeUtilities.positionNodeOnCanvasCenter({
-      canvasOffset: refPosition, // TODO
-      componentSize: adjustedSize, // TODO
-      maxSize: maxSize, // TODO
-    });
-    this.kNode.image(imageBitmap);
-    this.kNode.setSize(adjustedSize);
-    this.kNode.setPosition(centerPosition);
-    this.kNode.moveToTop();
-    this.kNode.fill(transparent);
-    this.listenToBaseKNode();
-    this.mediaLayerRef.draw();
-    //this.uploadImage(imageFile);
-    this.loaded();
-    this.listenToBaseKNode();
-  }
-
  private async uploadImage(imageFile: File) {
    //this.setProgress({ progress: 10, status: UploadStatus.UPLOADING });

@@ -44,10 +44,6 @@ export class NodesManager {
    // this.printNodesAndZIndices();
  }

-  public hasNode(node: MediaNode): boolean {
-    return this.allNodes.has(node);
-  }
-
  /*************************
   *  For Debugging use
   */
@@ -94,8 +94,7 @@ export function getTextNodeButtonStates(
        buttonName === ButtonNames.AI_STYLIZE ||
        buttonName === ButtonNames.SEGMENTATION ||
        buttonName === ButtonNames.DOWNLOAD ||
-        buttonName === ButtonNames.CHROMA ||
-        buttonName === ButtonNames.REMOVE_BACKGROUND
+        buttonName === ButtonNames.CHROMA
      ) {
        buttonStates[buttonName] = {
          disabled: true,
@@ -133,8 +132,7 @@ export function getVideoNodeButtonStates(
      }
      // soon to come feature is disabled
      if (buttonName === ButtonNames.AI_STYLIZE ||
-        buttonName === ButtonNames.COLOR ||
-        buttonName === ButtonNames.REMOVE_BACKGROUND
+        buttonName === ButtonNames.COLOR
      ) {
        buttonStates[buttonName] = {
          disabled: true,
@@ -176,8 +174,7 @@ export function getShapeNodeButtonStates(
        buttonName === ButtonNames.AI_STYLIZE ||
        buttonName === ButtonNames.SEGMENTATION ||
        buttonName === ButtonNames.DOWNLOAD ||
-        buttonName === ButtonNames.CHROMA ||
-        buttonName === ButtonNames.REMOVE_BACKGROUND
+        buttonName === ButtonNames.CHROMA
      ) {
        buttonStates[buttonName] = {
          disabled: true,
@@ -221,8 +218,7 @@ export function getMultiSelectButtonStates(
        buttonName === ButtonNames.SEGMENTATION ||
        buttonName === ButtonNames.DOWNLOAD ||
        buttonName === ButtonNames.CHROMA ||
-        buttonName === ButtonNames.COLOR ||
-        buttonName === ButtonNames.REMOVE_BACKGROUND
+        buttonName === ButtonNames.COLOR
      ) {
        buttonStates[buttonName] = {
          disabled: true,
@@ -25,7 +25,6 @@ import {
  isLoadingVisible,
  loadingProgress,
 } from "~/signals/uiEvents/loadingIndicator";
-import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts";

 interface ServerSettings {
  model_path: string;
@@ -286,7 +285,7 @@ export class RealTimeDrawEngine {
      this.client.onResult(async (response) => {
        try {
          // Convert the base64 image to ImageBitmap and update the preview
-          const bitmap = await DecodeBase64ToImage(response.image);
+          const bitmap = await this.base64ToImageBitmap(response.image);
          this.outputBitmap = bitmap;
          this.previewCanvas.image(bitmap);
          this.mediaLayerRef.batchDraw();
@@ -896,6 +895,35 @@ export class RealTimeDrawEngine {
    return base64String.split(",")[1];
  }

+  private async base64ToImageBitmap(
+    base64String: string,
+  ): Promise<ImageBitmap> {
+    // Create an image element
+    const img = document.createElement("img");
+
+    // Convert base64 to data URL if it doesn't include the prefix
+    const dataUrl = base64String.startsWith("data:")
+      ? base64String
+      : `data:image/png;base64,${base64String}`;
+
+    // Create a promise to handle the image loading
+    return new Promise((resolve, reject) => {
+      img.onload = async () => {
+        try {
+          const bitmap = await createImageBitmap(img);
+          resolve(bitmap);
+        } catch (error) {
+          reject(error);
+        }
+      };
+
+      img.onerror = () => reject(new Error("Failed to load image"));
+
+      // Set the source to trigger loading
+      img.src = dataUrl;
+    });
+  }
+
  public async saveOutput() {
    if (!this.outputBitmap) {
      console.error("No output bitmap available to save");
@@ -979,7 +1007,7 @@ export class RealTimeDrawEngine {
      });

      //console.log(base64BitmapResponse);
-      const decoded = await DecodeBase64ToImage(
+      const decoded = await this.base64ToImageBitmap(
        base64BitmapResponse as string,
      );

@@ -22,11 +22,6 @@ export const ToolbarNodeButtonData = [
    icon: faPalette,
    tooltip: "Color",
  },
-  {
-    name: ButtonNames.REMOVE_BACKGROUND,
-    icon: faScalpelLineDashed,
-    tooltip: "Remove Background",
-  },
  {
    name: ButtonNames.CHROMA,
    icon: faTransporter,
@@ -8,5 +8,4 @@ export enum ToolbarNodeButtonNames {
  MOVE_LAYER_DOWN = "MOVE_LAYER_DOWN",
  DOWNLOAD = "DOWNLOAD",
  COLOR = "COLOR",
-  REMOVE_BACKGROUND = "REMOVE_BACKGROUND",
 }
@@ -3,7 +3,7 @@ import { SignaledModelButtons } from "~/KonvaRootComponent/SignaledModelButtons"

 export const ToolbarTopCenter = () => {
  return (
-    <div className="flex flex-col items-center gap-3">
+    <div className="flex flex-col gap-3">
      {/* Model selection buttons (only shown in dev mode for now) */}
      {import.meta.env.DEV && <SignaledModelButtons />}

@@ -25,9 +25,7 @@ export const ToolbarTopLeft = () => {
  console.log(appMode.value);

  return (
-    <div
-      className={twMerge("z-20 flex h-fit w-fit items-center gap-8 pl-1 pr-4")}
-    >
+    <div className={twMerge("z-20 flex h-fit w-fit items-center gap-8 px-4")}>
      <img
        src="/brand/mira-logo.png"
        alt="logo"
@@ -55,7 +55,7 @@ export const TabSelector: React.FC<TabSelectorProps> = ({
      <TabList className="glass relative inline-flex min-w-fit overflow-x-auto rounded-lg p-1 shadow-lg">
        {/* Animated indicator */}
        <div
-          className="absolute top-1 z-10 h-[calc(100%-8px)] rounded-md border-2 border-primary bg-primary/30 transition-all duration-200 ease-in-out"
+          className="absolute top-1 z-10 h-[calc(100%-8px)] rounded-md border-2 border-primary bg-primary/30 transition-all duration-300 ease-in-out"
          style={{
            left: indicatorStyle.left,
            width: indicatorStyle.width,
@@ -68,7 +68,7 @@ export const TabSelector: React.FC<TabSelectorProps> = ({
            ref={(el) => (tabsRef.current[index] = el)}
            className={({ selected }) =>
              twMerge(
-                "relative z-20 mx-0.5 min-w-max rounded-md border-2 border-transparent px-4 py-0.5 text-center font-semibold transition-all duration-200 ease-in-out",
+                "relative z-20 mx-0.5 min-w-max rounded-md border-2 border-transparent px-4 py-1 text-center font-semibold transition-all duration-300 ease-in-out",
                selected ? "text-white" : "text-gray-300 hover:text-white",
              )
            }
@@ -1,30 +0,0 @@
-/**
- * This decodes Base64 encoded PNG images into an image element.
- * @param base64String a standard (not "web safe") base64-encoded string
- */
-export const DecodeBase64ToImage = async (base64String: string) : Promise<ImageBitmap> => {
-  // Create an image element
-  const img = document.createElement("img");
-
-  // Convert base64 to data URL if it doesn't include the prefix
-  const dataUrl = base64String.startsWith("data:")
-    ? base64String
-    : `data:image/png;base64,${base64String}`;
-
-  // Create a promise to handle the image loading
-  return new Promise((resolve, reject) => {
-    img.onload = async () => {
-      try {
-        const bitmap = await createImageBitmap(img);
-        resolve(bitmap);
-      } catch (error) {
-        reject(error);
-      }
-    };
-
-    img.onerror = () => reject(new Error("Failed to load image"));
-
-    // Set the source to trigger loading
-    img.src = dataUrl;
-  });
-}
@@ -1,22 +0,0 @@
-
-export const EncodeImageBitmapToBase64 = async (imageBitmap: ImageBitmap): Promise<string> => {
-  // Create a temporary canvas
-  const canvas = document.createElement("canvas");
-  canvas.width = imageBitmap.width;
-  canvas.height = imageBitmap.height;
-
-  // Draw the ImageBitmap onto the canvas
-  const ctx = canvas.getContext("2d");
-  if (!ctx) throw new Error("Failed to get 2D context");
-
-  ctx.drawImage(imageBitmap, 0, 0);
-
-  // Convert to base64
-  const base64String = canvas.toDataURL("image/png");
-
-  // Clean up
-  canvas.remove();
-
-  // Remove the data:image/png;base64, prefix if you want just the base64 string
-  return base64String.split(",")[1];
-}
Author	SHA1	Message	Date
Brandon Thomas	b2a7a40f5f	don't download sdxl anymore	2025-03-13 00:31:51 -04:00
Brandon Thomas	9085c288bd	stop downloading clip tokenizer from hf	2025-03-13 00:24:06 -04:00
Brandon Thomas	dd71d022ae	vae wasn't cached in memory	2025-03-12 23:58:52 -04:00
Brandon Thomas	692c409e1b	load text encoder from our bucket	2025-03-12 23:54:42 -04:00
Brandon Thomas	3ed775cf76	vae and unet from cloudflare instead of hf	2025-03-12 23:38:49 -04:00
Brandon Thomas	08ea29d44c	cleanup lcm	2025-03-12 22:16:18 -04:00