added loading indicator text

2026-10-25 03:09:30 -04:00 · 2025-03-05 10:21:03 +07:00
44 changed files with 1797 additions and 2693 deletions
@@ -2,7 +2,6 @@
  <dictionary name="bt">
    <words>
      <w>arpabet</w>
-      <w>artcraft</w>
      <w>backfill</w>
      <w>backfilling</w>
      <w>backfills</w>
@@ -664,7 +664,6 @@ dependencies = [
 "candle-nn",
 "candle-transformers",
 "cudarc",
- "expanduser",
 "filesys",
 "hf-hub 0.3.2",
 "image",
@@ -673,7 +672,6 @@ dependencies = [
 "memory_caching",
 "once_cell",
 "rand 0.8.5",
- "reqwest 0.12.12",
 "serde",
 "serde_json",
 "serde_yaml",
@@ -733,12 +731,6 @@ version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"

-[[package]]
-name = "arrayvec"
-version = "0.5.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b"
-
 [[package]]
 name = "arrayvec"
 version = "0.7.6"
@@ -1022,7 +1014,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6678909d8c5d46a42abcf571271e15fdbc0a225e3646cf23762cd415046c78bf"
 dependencies = [
 "anyhow",
- "arrayvec 0.7.6",
+ "arrayvec",
 "log",
 "nom",
 "num-rational",
@@ -1035,7 +1027,7 @@ version = "0.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e335041290c43101ca215eed6f43ec437eb5a42125573f600fc3fa42b9bddd62"
 dependencies = [
- "arrayvec 0.7.6",
+ "arrayvec",
 ]

 [[package]]
@@ -1087,7 +1079,7 @@ dependencies = [
 "rustversion",
 "serde",
 "sync_wrapper 0.1.2",
- "tower 0.4.13",
+ "tower",
 "tower-layer",
 "tower-service",
 ]
@@ -1305,17 +1297,6 @@ dependencies = [
 "wyz",
 ]

-[[package]]
-name = "blake2b_simd"
-version = "0.5.11"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "afa748e348ad3be8263be728124b24a24f268266f6f5d58af9d75f6a40b5c587"
-dependencies = [
- "arrayref",
- "arrayvec 0.5.2",
- "constant_time_eq 0.1.5",
-]
-
 [[package]]
 name = "blake2b_simd"
 version = "1.0.2"
@@ -1323,7 +1304,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "23285ad32269793932e830392f2fe2f83e26488fd3ec778883a93c8323735780"
 dependencies = [
 "arrayref",
- "arrayvec 0.7.6",
+ "arrayvec",
 "constant_time_eq 0.3.0",
 ]

@@ -3206,17 +3187,6 @@ dependencies = [
 "dirs-sys 0.4.1",
 ]

-[[package]]
-name = "dirs"
-version = "1.0.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3fd78930633bd1c6e35c4b42b1df7b0cbc6bc191146e512bb3bedf243fcc3901"
-dependencies = [
- "libc",
- "redox_users 0.3.5",
- "winapi 0.3.9",
-]
-
 [[package]]
 name = "dirs"
 version = "5.0.1"
@@ -3346,7 +3316,7 @@ dependencies = [
 "google_drive_common",
 "hashing",
 "hostname 0.4.0",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "jobs_common",
 "log",
 "mimetypes",
@@ -3867,17 +3837,6 @@ version = "2.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0"

-[[package]]
-name = "expanduser"
-version = "1.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "14e0b79235da57db6b6c2beed9af6e5de867d63a973ae3e91910ddc33ba40bc0"
-dependencies = [
- "dirs 1.0.5",
- "lazy_static",
- "pwd",
-]
-
 [[package]]
 name = "exr"
 version = "1.73.0"
@@ -4803,7 +4762,7 @@ dependencies = [
 "errors",
 "jwt-simple",
 "my-workspace-hack",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde",
 "serde_derive",
 "serde_json",
@@ -5029,7 +4988,7 @@ dependencies = [
 "native-tls",
 "num_cpus",
 "rand 0.8.5",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde",
 "serde_json",
 "thiserror 2.0.11",
@@ -5233,9 +5192,9 @@ dependencies = [

 [[package]]
 name = "httparse"
-version = "1.10.1"
+version = "1.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
+checksum = "d897f394bad6a705d5f4104762e116a75639e470d80901eed05a860a95cb1904"

 [[package]]
 name = "httpdate"
@@ -5275,9 +5234,9 @@ dependencies = [

 [[package]]
 name = "hyper"
-version = "1.6.0"
+version = "1.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cc2b571658e38e0c01b1fdca3bbbe93c00d3d71693ff2770043f8c29bc7d6f80"
+checksum = "fe575dd17d0862a9a33781c8c4696a55c320909004a67a00fb286ba8b1bc496d"
 dependencies = [
 "bytes 1.10.0",
 "futures-channel",
@@ -5301,7 +5260,7 @@ checksum = "08afdbb5c31130e3034af566421053ab03787c640246a446327f550d11bcb333"
 dependencies = [
 "futures-util",
 "http 1.1.0",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "hyper-util",
 "rustls 0.23.13",
 "rustls-pki-types",
@@ -5344,7 +5303,7 @@ checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0"
 dependencies = [
 "bytes 1.10.0",
 "http-body-util",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "hyper-util",
 "native-tls",
 "tokio",
@@ -5354,19 +5313,20 @@ dependencies = [

 [[package]]
 name = "hyper-util"
-version = "0.1.10"
+version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "df2dcfbe0677734ab2f3ffa7fa7bfd4706bfdc1ef393f2ee30184aed67e631b4"
+checksum = "7b875924a60b96e5d7b9ae7b066540b1dd1cbd90d1828f54c92e02a283351c56"
 dependencies = [
 "bytes 1.10.0",
 "futures-channel",
 "futures-util",
 "http 1.1.0",
 "http-body 1.0.0",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "pin-project-lite",
 "socket2",
 "tokio",
+ "tower",
 "tower-service",
 "tracing",
 ]
@@ -5603,7 +5563,7 @@ dependencies = [
 "rand 0.8.5",
 "redis_common",
 "redis_schema",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde",
 "serde_derive",
 "serde_json",
@@ -5900,7 +5860,7 @@ checksum = "50ae7e0018905a795d6f2a60ac32a547490abdd8df509906a8c6171e6d861711"
 dependencies = [
 "anyhow",
 "binstring",
- "blake2b_simd 1.0.2",
+ "blake2b_simd",
 "coarsetime",
 "ct-codecs",
 "ed25519-compact",
@@ -6859,7 +6819,7 @@ dependencies = [
 "actix-rt",
 "ahash 0.8.11",
 "anstream",
- "arrayvec 0.7.6",
+ "arrayvec",
 "base64 0.21.7",
 "base64 0.22.1",
 "base64ct",
@@ -6892,7 +6852,7 @@ dependencies = [
 "heck 0.4.1",
 "hmac 0.12.1",
 "hyper 0.14.28",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "hyper-util",
 "indexmap 2.2.6",
 "itertools 0.12.1",
@@ -6918,7 +6878,7 @@ dependencies = [
 "regex",
 "regex-automata",
 "regex-syntax",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "rsa",
 "rustls-pki-types",
 "semver 1.0.23",
@@ -6944,7 +6904,7 @@ dependencies = [
 "tokio",
 "tokio-stream",
 "tokio-util",
- "tower 0.4.13",
+ "tower",
 "tracing",
 "unicode-bidi",
 "unicode-normalization",
@@ -8636,16 +8596,6 @@ dependencies = [
 "version_check",
 ]

-[[package]]
-name = "pwd"
-version = "1.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "72c71c0c79b9701efe4e1e4b563b2016dd4ee789eb99badcb09d61ac4b92e4a2"
-dependencies = [
- "libc",
- "thiserror 1.0.63",
-]
-
 [[package]]
 name = "pyo3"
 version = "0.22.6"
@@ -8971,7 +8921,7 @@ checksum = "cd87ce80a7665b1cce111f8a16c1f3929f6547ce91ade6addf4ec86a8dda5ce9"
 dependencies = [
 "arbitrary",
 "arg_enum_proc_macro",
- "arrayvec 0.7.6",
+ "arrayvec",
 "av1-grain",
 "bitstream-io",
 "built",
@@ -9200,17 +9150,6 @@ dependencies = [
 "bitflags 2.5.0",
 ]

-[[package]]
-name = "redox_users"
-version = "0.3.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "de0737333e7a9502c789a36d7c7fa6092a49895d4faa31ca5df163857ded2e9d"
-dependencies = [
- "getrandom 0.1.16",
- "redox_syscall 0.1.57",
- "rust-argon2",
-]
-
 [[package]]
 name = "redox_users"
 version = "0.4.6"
@@ -9330,9 +9269,9 @@ dependencies = [

 [[package]]
 name = "reqwest"
-version = "0.12.12"
+version = "0.12.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43e734407157c3c2034e0258f5e4473ddb361b1e85f95a66690d67264d7cd1da"
+checksum = "f8f4955649ef5c38cc7f9e8aa41761d48fb9677197daea9984dc54f56aad5e63"
 dependencies = [
 "base64 0.22.1",
 "bytes 1.10.0",
@@ -9344,7 +9283,7 @@ dependencies = [
 "http 1.1.0",
 "http-body 1.0.0",
 "http-body-util",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "hyper-rustls",
 "hyper-tls 0.6.0",
 "hyper-util",
@@ -9369,7 +9308,6 @@ dependencies = [
 "tokio-native-tls",
 "tokio-rustls",
 "tokio-util",
- "tower 0.5.2",
 "tower-service",
 "url",
 "wasm-bindgen",
@@ -9497,18 +9435,6 @@ dependencies = [
 "realfft",
 ]

-[[package]]
-name = "rust-argon2"
-version = "0.8.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4b18820d944b33caa75a71378964ac46f58517c92b6ae5f762636247c09e78fb"
-dependencies = [
- "base64 0.13.1",
- "blake2b_simd 0.5.11",
- "constant_time_eq 0.1.5",
- "crossbeam-utils 0.8.20",
-]
-
 [[package]]
 name = "rust-embed"
 version = "8.0.0"
@@ -9576,7 +9502,7 @@ dependencies = [
 "minidom",
 "percent-encoding",
 "quick-xml 0.36.2",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde",
 "serde_derive",
 "serde_json",
@@ -9594,7 +9520,7 @@ version = "1.36.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b082d80e3e3cc52b2ed634388d436fe1f4de6af5786cc2de9ba9737527bdf555"
 dependencies = [
- "arrayvec 0.7.6",
+ "arrayvec",
 "borsh",
 "bytes 1.10.0",
 "num-traits",
@@ -10907,7 +10833,7 @@ dependencies = [
 "redis_common",
 "redis_schema",
 "regex",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "reusable_types",
 "rust-s3",
 "serde",
@@ -11272,7 +11198,7 @@ name = "symphonia-core"
 version = "0.5.0"
 source = "git+https://github.com/pdeljanov/Symphonia.git?rev=e1a70095c4409478c1ded499fb8dd4afa540fb90#e1a70095c4409478c1ded499fb8dd4afa540fb90"
 dependencies = [
- "arrayvec 0.7.6",
+ "arrayvec",
 "bitflags 1.3.2",
 "bytemuck",
 "lazy_static",
@@ -11285,7 +11211,7 @@ version = "0.5.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "798306779e3dc7d5231bd5691f5a813496dc79d3f56bf82e25789f2094e022c3"
 dependencies = [
- "arrayvec 0.7.6",
+ "arrayvec",
 "bitflags 1.3.2",
 "bytemuck",
 "lazy_static",
@@ -11671,7 +11597,7 @@ dependencies = [
 "percent-encoding",
 "plist",
 "raw-window-handle",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde",
 "serde_json",
 "serde_repr",
@@ -12019,7 +11945,7 @@ dependencies = [
 "easyenv",
 "log",
 "my-workspace-hack",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde_json",
 ]

@@ -12473,7 +12399,7 @@ dependencies = [
 "prost",
 "tokio",
 "tokio-stream",
- "tower 0.4.13",
+ "tower",
 "tower-layer",
 "tower-service",
 "tracing",
@@ -12508,32 +12434,17 @@ dependencies = [
 "tracing",
 ]

-[[package]]
-name = "tower"
-version = "0.5.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9"
-dependencies = [
- "futures-core",
- "futures-util",
- "pin-project-lite",
- "sync_wrapper 1.0.1",
- "tokio",
- "tower-layer",
- "tower-service",
-]
-
 [[package]]
 name = "tower-layer"
-version = "0.3.3"
+version = "0.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e"
+checksum = "c20c8dbed6283a09604c3e69b4b7eeb54e298b8a600d4d5ecb5ad39de609f1d0"

 [[package]]
 name = "tower-service"
-version = "0.3.3"
+version = "0.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3"
+checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52"

 [[package]]
 name = "tracing"
@@ -12670,7 +12581,7 @@ dependencies = [
 "google_drive_common",
 "hashing",
 "hostname 0.4.0",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "jobs_common",
 "log",
 "my-workspace-hack",
@@ -12708,7 +12619,7 @@ dependencies = [
 "my-workspace-hack",
 "mysql_queries",
 "r2d2_redis",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "serde",
 "serde_derive",
 "serde_json",
@@ -13068,7 +12979,7 @@ dependencies = [
 "actix-web",
 "mime_guess",
 "regex",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "rust-embed",
 "serde",
 "serde_json",
@@ -13197,7 +13108,7 @@ dependencies = [
 "google_drive_common",
 "hashing",
 "hostname 0.4.0",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "jobs_common",
 "log",
 "my-workspace-hack",
@@ -13227,7 +13138,7 @@ dependencies = [
 "filesys",
 "google_drive_common",
 "hostname 0.4.0",
- "hyper 1.6.0",
+ "hyper 1.3.1",
 "jobs_common",
 "log",
 "my-workspace-hack",
@@ -13404,7 +13315,7 @@ dependencies = [
 "log",
 "my-workspace-hack",
 "once_cell",
- "reqwest 0.12.12",
+ "reqwest 0.12.7",
 "rss",
 "scraper",
 "serde",
@@ -2,6 +2,6 @@ image_width: 512
 image_height: 512
 #scheduler_steps: 15
 # NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
-scheduler_steps: 4
+scheduler_steps: 2
 scheduler_samples: 1
 cfg_scale: 1.1
@@ -38,9 +38,7 @@ image = "0.25.5"
 serde_yaml = "0.9"
 log = "0.4"
 once_cell = "1.20.3"
-expanduser = "1.2.2"
 rand = "0.8"
-reqwest = "0.12.12"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 tauri = { version = "2.2.4", features = [] }
@@ -49,7 +47,7 @@ tokenizers = "0.15.2"

 [features]
 default = ["cuda"]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda", "candle-transformers/flash-attn", "candle-core/cudnn"]
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]

 [dev-dependencies]
 libloading = "0.8.6"
@@ -1,16 +0,0 @@
-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>,
-) -> Result<String, String> {
-  
-  download_all_models(&app_data_root).await
-    .map_err(|err| {
-      "error downloading models".to_string()
-    })?;
-  
-  Ok(("downloaded".to_string()))
-}
@@ -1,7 +1,6 @@
-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::ml::stable_diffusion::stable_diffusion_pipeline::{stable_diffusion_pipeline, Args};
 use crate::state::app_config::AppConfig;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use bytes::BytesMut;
@@ -11,7 +10,6 @@ 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 = "A beautiful landscape with mountains and a lake";
@@ -31,7 +29,6 @@ pub async fn infer_image(
  model_config: State<'_, AppConfig>,
  model_cache: State<'_, ModelCache>,
  prompt_cache: State<'_, PromptCache>,
-  app_data_root: State<'_, AppDataRoot>,
  app: AppHandle,
 ) -> Result<String, String> {

@@ -43,23 +40,7 @@ pub async fn infer_image(
  let image = hydrate_base64_image(image)
    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;

-  // 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, 
-    image, 
-    strength, 
-    &model_config, 
-    &model_cache, 
-    prompt_cache, 
-    app, 
-    &app_data_root
-  ).await;
+  let result = do_infer_image(&prompt, image, strength, &model_config, &model_cache, prompt_cache, app);
  
  if let Err(err) = result.as_deref() {
    error!("There was an error: {:?}", err);
@@ -68,31 +49,31 @@ pub async fn infer_image(
  result
 }

-async fn do_infer_image(
+fn do_infer_image(
  prompt: &str,
  image: DynamicImage,
  strength: Option<u8>,
  config: &AppConfig,
  model_cache: &ModelCache,
-  prompt_cache: State<'_, PromptCache>,
+  prompt_cache: State<PromptCache>,
  app: AppHandle,
-  app_data_root: &AppDataRoot,
 ) -> Result<String, String> {
-
+  println!("infer_image called; generating image with SDXL Turbo...");
+  
  let args = Args {
    image: &image,
    prompt: prompt.to_string(),
    uncond_prompt: NEGATIVE_PROMPT.to_string(),
+    //guidance_scale: Some(0.0),
    model_cache,
    configs: config,
    prompt_cache: &prompt_cache,
    i2i_strength: strength,
    cfg_scale: config.cfg_scale,
    app: &app,
-    app_data_root,
  };

-  match lcm_pipeline(args) {
+  match stable_diffusion_pipeline(args) {
    Ok(image) => {
      let mut bytes = Vec::with_capacity(1024*1024);
      
@@ -1,4 +1,3 @@
-pub mod download_models;
 pub mod flip_image;
 pub mod image_endpoint;
 pub mod test_counter;
@@ -3,7 +3,7 @@ use serde::Serialize;

 #[derive(Clone, Copy, Serialize)]
 #[serde(rename_all = "snake_case")]
-pub enum NotificationModelType {
+pub enum ModelType {
  Unet,
  Vae,
 }
@@ -13,10 +13,10 @@ pub enum NotificationModelType {
 pub enum NotificationEvent<'a> {
  ModelDownloadStarted {
    model_name: &'a str, 
-    model_type: NotificationModelType,
+    model_type: ModelType,
  },
  ModelDownloadComplete {
    model_name: &'a str,
-    model_type: NotificationModelType,
+    model_type: ModelType,
  }
 }
@@ -2,15 +2,12 @@ pub mod endpoints;
 pub mod events;
 pub mod ml;
 pub mod state;
-pub mod threads;

-use crate::endpoints::download_models::download_models;
-use crate::endpoints::image_endpoint::infer_image;
-use crate::endpoints::test_counter::test_counter;
 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 endpoints::image_endpoint::infer_image;
+use endpoints::test_counter::test_counter;
 use tauri_plugin_log::Target;
 use tauri_plugin_log::TargetKind;

@@ -29,9 +26,6 @@ pub fn run() {
  println!("Creating model cache...");
  
  let model_cache = ModelCache::new();
-  
-  let app_data_root = config.app_data_root.clone();
-  let app_data_root2 = config.app_data_root.clone();

  println!("Initializing backend runtime...");

@@ -48,20 +42,14 @@ pub fn run() {
      //      .build(),
      //  )?;
      //}
-      let app = app.handle().clone();
-
-      tauri::async_runtime::spawn(downloader_thread(app_data_root2, app));
-
      Ok(())
    })
    .manage(config)
    .manage(prompt_cache)
    .manage(model_cache)
-    .manage(app_data_root)
    .invoke_handler(tauri::generate_handler![
      test_counter, 
      infer_image,
-      download_models,
    ])
    .run(tauri::generate_context!())
    .expect("error while running tauri application");
@@ -1,15 +0,0 @@
-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(())
-}
@@ -1,30 +0,0 @@
-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)
-}
@@ -1,2 +0,0 @@
-pub mod download_model_registry_file;
-pub mod download_all_models;
@@ -3,6 +3,4 @@ pub mod model_cache;
 pub mod model_file;
 pub mod models;
 pub mod prompt_cache;
-pub mod stable_diffusion;
-pub mod model_type;
-pub mod downloads;
+pub mod stable_diffusion;
@@ -1,47 +0,0 @@
-use std::path::PathBuf;
-use crate::state::app_dir::AppWeightsDir;
-
-#[derive(Copy, Clone, Debug)]
-pub enum ModelType {
-  ClipJson,
-  SdxlTurboUnet,
-  SdxlTurboVae,
-  SdxlTurboClipEncoder,
-  SdxlTurboClipEncoder2,
-}
-
-impl ModelType {
-  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"),
-    }
-  }
-
-  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",
-    }
-  }
-
-  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",
-    }
-  }
-  
-  pub fn get_path(&self, weights_dir: &AppWeightsDir) -> PathBuf {
-    weights_dir.model_path(self)
-  }
-}
@@ -8,19 +8,28 @@ pub struct UNetModel {
 }

 impl UNetModel {
-  pub fn new<P: AsRef<Path>>(sd_config: &StableDiffusionConfig, unet_weights_file: P, device: &Device, dtype: DType) -> anyhow::Result<Self> {
+  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 })
+    
+    let model = sd_config.build_unet(
+      unet_weights_file, 
+      &device, 
+      4, 
+      false, 
+      dtype
+    )?;
+    
+    Ok(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,7 +3,6 @@ extern crate accelerate_src;
 #[cfg(feature = "mkl")]
 extern crate intel_mkl_src;

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

 use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
@@ -11,8 +10,6 @@ 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(
@@ -27,7 +24,6 @@ pub fn infer_clip_text_embeddings(
  device: &Device,
  dtype: DType,
  use_guide_scale: bool,
-  weights_dir: &AppWeightsDir,
 ) -> anyhow::Result<Tensor> {
  info!("Preparing text embeddings...");

@@ -54,7 +50,6 @@ pub fn infer_clip_text_embeddings(
        dtype,
        use_guide_scale,
        *first,
-        weights_dir,
      )
    })
    .collect::<anyhow::Result<Vec<_>>>()?;
@@ -78,25 +73,15 @@ fn do_infer_clip_text_embeddings(
    dtype: DType,
    use_guide_scale: bool,
    first: bool,
-    weights_dir: &AppWeightsDir,
  ) -> anyhow::Result<Tensor> {
  
  let tokenizer_file = if first {
-    info!("ModelFile::Tokenizer");
    ModelFile::Tokenizer
  } else {
-    info!("ModelFile::Tokenizer2");
    ModelFile::Tokenizer2
  };
-  
-  info!("Loading clip from download");
-  //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
-  let tokenizer = weights_dir.model_path(&ModelType::ClipJson);
-  
-  info!("Tokenizer path: {:?}", tokenizer);
-  
+  let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
  let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
-  
  let pad_id = match &sd_config.clip.pad_with {
    Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(),
    None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(),
@@ -127,11 +112,9 @@ fn do_infer_clip_text_embeddings(
    ModelFile::Clip2
  };
  let clip_weights = if first {
-    //clip_weights_file.get(clip_weights, sd_version, use_f16)?
-    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder)
+    clip_weights_file.get(clip_weights, sd_version, use_f16)?
  } else {
-    //clip_weights_file.get(clip2_weights, sd_version, use_f16)?
-    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder2)
+    clip_weights_file.get(clip2_weights, sd_version, use_f16)?
  };
  let clip_config = if first {
    &sd_config.clip
@@ -1,293 +0,0 @@
-use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
-
-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::state::app_config::AppConfig;
-use anyhow::{anyhow, Error as E, Result};
-use candle_core::{DType, IndexOp, Tensor, D};
-use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
-use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
-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,
-    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 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,
-    } = args;
-
-    let weights_dir = app_data_root.weights_dir();
-    
-    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,
-        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 {
-            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"
-        },
-    };
-
-    let use_guide_scale = guidance_scale > 1.0;
-
-    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!("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,
-            use_guide_scale,
-            weights_dir,
-        )?;
-        prompt_cache.store_copy(&prompt, &tensor)?;
-        tensor
-    };
-
-    println!("Text embeddings shape: {:?}", text_embeddings.shape());
-
-    println!("Loading input image into tensor...");
-
-    let input_image = dynamic_image_to_tensor(
-        image,
-        &configs.device,
-        configs.dtype)?;
-
-    println!("Reference image shape: {:?}", input_image.shape());
-
-    let vae_scale = get_vae_scale(configs.sd_version);
-
-    let maybe_vae = model_cache.get_vae()?;
-    
-    let vae = match maybe_vae {
-        Some(vae) => vae,
-        None => {
-            info!("No vae found in cache; loading...");
-            
-            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
-              .build_vae(vae_file, &configs.device, configs.dtype)?;
-
-            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
-        }
-    };
-
-    let maybe_unet = model_cache.get_unet()?;
-
-    let unet = match maybe_unet {
-        Some(unet) => unet,
-        None => {
-            info!("No unet found in cache; loading...");
-
-            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 = 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
-        }
-    };
-
-    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 = 256;
-    let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(
-        guidance_scale, 
-        embedding_dim, 
-        &configs.device, 
-        configs.dtype
-    )?;
-
-    // Generate latents from input image using the LCM approach
-    println!("Generating latents from input image...");
-    let init_latents = (init_latent_dist.sample()? * vae_scale)?;
-    println!("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 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);
-    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();
-
-    for (timestep_index, &timestep) in timesteps.iter().enumerate() {
-        println!("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()
-        };
-
-        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,
-            &guidance_scale_embedding) {
-            Ok(pred) => pred,
-            Err(e) => {
-                println!("UNet inference failed with error: {}", e);
-                return Err(anyhow::anyhow!("UNet inference failed: {}", e));
-            },
-        };
-
-        if use_guide_scale {
-            let chunks = noise_pred.chunk(2, 0)?;
-            let (noise_pred_uncond, noise_pred_text) = (&chunks[0], &chunks[1]);
-            noise_pred = (noise_pred_uncond + (guidance_scale * (noise_pred_text - noise_pred_uncond)?)?)?;
-        }
-
-        // Apply the LCM scheduler step
-        latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?;
-    }
-
-    println!("Diffusion process completed, decoding image...");
-    let image = vae.decode(&(latents / vae_scale)?)?;
-    
-    println!("VAE decode completed");
-
-    println!("Scaling image back");
-    let image = ((image / 2.)? + 0.5)?;
-    
-    println!("Normalized image values");
-    let image = (image.clamp(0f32, 1.)? * 255.)?;
-
-    println!("Convert to int8");
-    let image = image.to_dtype(DType::U8)?;
-    
-    println!("Converted to 8-bit format");
-
-    let image = tensor_to_image_buffer(&image.i(0)?)?;
-
-    Ok(image)
-}
@@ -1,4 +1,3 @@
 pub mod get_vae_scale;
 pub mod infer_clip_text_embeddings;
-pub mod stable_diffusion_pipeline;
-pub mod lcm_pipeline;
+pub mod stable_diffusion_pipeline;
@@ -1,12 +1,9 @@
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

-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;
@@ -18,7 +15,8 @@ use image::DynamicImage;
 use log::info;
 use rand::Rng;
 use tauri::{AppHandle, Emitter};
-use crate::state::app_dir::AppDataRoot;
+use crate::events::notification_event::{ModelType, NotificationEvent};
+use crate::ml::models::unet_model::UNetModel;

 pub struct Args<'a> {
    pub image: &'a DynamicImage,
@@ -29,7 +27,6 @@ pub struct Args<'a> {
    pub configs: &'a AppConfig,
    pub model_cache: &'a ModelCache,
    pub prompt_cache: &'a PromptCache,
-    pub app_data_root: &'a AppDataRoot,
    pub app: &'a AppHandle,
 }

@@ -43,12 +40,9 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
        model_cache, 
        prompt_cache, 
        app, 
-        image,
-        app_data_root,
+        image,  
    } = args;
    
-    let weights_dir = app_data_root.weights_dir();
-    
    println!("Starting image generation with the following configuration:");
    println!("  Model: {:?}", configs.sd_version);
    println!("  Prompt: {}", prompt);
@@ -103,7 +97,6 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
            &configs.device,
            configs.dtype,
            use_guide_scale,
-            weights_dir,
        )?;
        prompt_cache.store_copy(&prompt, &tensor)?;
        tensor
@@ -129,18 +122,25 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
        None => {
            info!("No vae found in cache; loading...");
            
+            let repo = configs.sd_version.repo();
+
+            println!("Building VAE model from : {:?} ... (3)", repo);
+
+            let vae_file = configs.hf_api.model(repo.to_string())
+              .get("vae/diffusion_pytorch_model.safetensors")?;
+
+            println!("Building VAE model from file {:?}...", &vae_file);
+
            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,
+                    model_name: repo,
+                    model_type: ModelType::Vae,
                })?;
            }
            
-            let vae_file = weights_dir.model_path(&ModelType::SdxlTurboVae);
-            
            let vae = configs
              .sd_config
              .build_vae(vae_file, &configs.device, configs.dtype)?;
@@ -151,8 +151,8 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
            
            if notify_download_complete {
                app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: "sdxl-turbo-vae",
-                    model_type: NotificationModelType::Vae,
+                    model_name: repo,
+                    model_type: ModelType::Vae,
                })?;
            }
            
@@ -167,17 +167,23 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
        None => {
            info!("No unet found in cache; loading...");

+            let repo = configs.sd_version.repo();
+
+            info!("Downloading UNET model files from: {} ...", repo);
+
+            let unet_file = configs.hf_api.model(repo.to_string())
+              .get("unet/diffusion_pytorch_model.safetensors")
+              .map_err(|err| anyhow!("error fetching model: {:?}", err))?;
+
            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,
+                    model_name: repo,
+                    model_type: ModelType::Unet,
                })?;
            }
-            
-            let unet_file = weights_dir.model_path(&ModelType::SdxlTurboUnet);

            let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
              .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
@@ -188,8 +194,8 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {

            if notify_download_complete {
                app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: "sdxl-turbo-unet",
-                    model_type: NotificationModelType::Unet,
+                    model_name: repo,
+                    model_type: ModelType::Unet,
                })?;
            }

@@ -1,14 +1,12 @@
 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 expanduser::expanduser;
 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;
+
+const CONFIG_FILENAME : &str = "app_config.yaml";

 const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
 const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
@@ -16,8 +14,6 @@ const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
 const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
 const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;

-const DEFAULT_DATA_DIR_UNIX : &str = "~/artcraft";
-
 /// A central place to configure app-wide details.
 pub struct AppConfig {
  /// We only support one device type currently, and we can't multiplex 
@@ -41,15 +37,13 @@ pub struct AppConfig {
  
  /// Probably shouldn't be here.
  pub hf_api: Api,
-  
-  /// Root location for application data.
-  pub app_data_root: AppDataRoot,
 }

 impl AppConfig {
  
  pub fn init() -> anyhow::Result<Self> {
-    let yaml_configs = YamlConfig::load_from_config_file_recursive();
+    println!("Configuration filename: {}", CONFIG_FILENAME);
+    let yaml_configs = load_yaml_configs();
    
    println!("Configs: {:?}", yaml_configs);
    
@@ -70,26 +64,6 @@ impl AppConfig {
    
    let hf_api = Api::new()?;
    
-    let os_platform = OsPlatform::get();
-
-    let default_data_path = match os_platform {
-      OsPlatform::Linux => yaml_configs.linux_default_data_path.as_deref(),
-      OsPlatform::MacOs => yaml_configs.mac_default_data_path.as_deref(),
-      OsPlatform::Windows => yaml_configs.windows_default_data_path.as_deref(),
-    };
-    
-    let data_root_path = match default_data_path {
-      Some(default_data_path) => default_data_path.to_string(),
-      None => match os_platform { 
-        OsPlatform::Linux | OsPlatform::MacOs => DEFAULT_DATA_DIR_UNIX.to_string(),
-        OsPlatform::Windows => ".".to_string(),
-      }
-    };
-      
-    println!("Using data root: {}", data_root_path);
-
-    let app_data_root = AppDataRoot::create_existing(data_root_path)?;
-    
    Ok(Self {
      device,
      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
@@ -97,12 +71,26 @@ impl AppConfig {
      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_steps: yaml_configs.scheduler_steps.unwrap_or(1),
      scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
      seed: yaml_configs.seed,
      cfg_scale: yaml_configs.cfg_scale,
      hf_api,
-      app_data_root,
    })
  }
 }
+
+fn load_yaml_configs() -> YamlConfig {
+  let maybe_path = recursively_find_file_by_name(CONFIG_FILENAME, "../../../", 5)
+    .ok()
+    .flatten();
+  if let Some(path) = maybe_path {
+    let path = path.canonicalize().unwrap_or(path);
+    println!("Attempting to load configs from: {:?}", &path);
+    if let Ok(config) = YamlConfig::read_from_file(path) {
+      return config;
+    }
+  }
+  println!("Loading default configs.");
+  YamlConfig::default()
+}
@@ -1,123 +0,0 @@
-use crate::state::os_platform::OsPlatform;
-use expanduser::expanduser;
-use std::path::{Path, PathBuf};
-use crate::ml::model_type::ModelType;
-
-const ASSETS_SUBDIRECTORY : &str = "assets";
-const WEIGHTS_SUBDIRECTORY : &str = "weights";
-
-/// The path to the application data directory, which includes "asset" and "weights" data.
-#[derive(Clone)]
-pub struct AppDataRoot {
-  path: PathBuf,
-  assets_dir: AppAssetsDir,
-  weights_dir: AppWeightsDir,
-}
-
-#[derive(Clone)]
-pub struct AppAssetsDir {
-  path: PathBuf,
-}
-
-#[derive(Clone)]
-pub struct AppWeightsDir {
-  path: PathBuf,
-}
-
-impl AppDataRoot {
-  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
-    let mut dir = dir.as_ref().to_path_buf();
-    
-    match OsPlatform::get() {
-      OsPlatform::Linux | OsPlatform::MacOs => {
-        if let Some(d) = dir.as_os_str().to_str() {
-          dir = expanduser(d)?;
-        }
-      },
-      OsPlatform::Windows => {}
-    }
-    
-    if !dir.is_dir() {
-      println!("Creating directory {:?}", dir);
-      std::fs::create_dir_all(&dir)?;
-    }
-
-    match dir.canonicalize() {
-      Ok(d) => dir = d,
-      Err(err) => {
-        println!("Error canonicalizing {:?}: {}", dir, err);
-      }
-    }
-    
-    let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?;
-    let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?;
-    
-    Ok(Self {
-      path: dir,
-      assets_dir,
-      weights_dir,
-    })
-  }
-  
-  pub fn assets_dir(&self) -> &AppAssetsDir {
-    &self.assets_dir
-  }
-  
-  pub fn weights_dir(&self) -> &AppWeightsDir {
-    &self.weights_dir
-  }
-  
-  pub fn path(&self) -> &Path {
-    &self.path
-  }
-}
-
-impl AppAssetsDir {
-  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
-    let mut dir = dir.as_ref().to_path_buf();
-    match dir.canonicalize() {
-      Ok(d) => dir = d,
-      Err(err) => {
-        println!("Error canonicalizing {:?}: {}", dir, err);
-      }
-    }
-    if !dir.exists() {
-      println!("Creating directory {:?}", dir);
-      std::fs::create_dir(&dir)?;
-    }
-    Ok(Self {
-      path: dir,
-    })
-  }
-
-  pub fn path(&self) -> &Path {
-    &self.path
-  }
-}
-
-impl AppWeightsDir {
-  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
-    let mut dir = dir.as_ref().to_path_buf();
-    match dir.canonicalize() {
-      Ok(d) => dir = d,
-      Err(err) => {
-        println!("Error canonicalizing {:?}: {}", dir, err);
-      }
-    }
-    if !dir.exists() {
-      println!("Creating directory {:?}", dir);
-      std::fs::create_dir(&dir)?;
-    }
-    Ok(Self {
-      path: dir,
-    })
-  }
-
-  pub fn path(&self) -> &Path {
-    &self.path
-  }
-  
-  pub fn model_path(&self, model_registry: &ModelType) -> PathBuf {
-    self.path.join(model_registry.get_filename())
-  }
-}
@@ -1,4 +1,2 @@
-pub mod app_config;
-pub mod app_dir;
-pub mod os_platform;
-pub mod yaml_config;
+pub mod yaml_config;
+pub mod app_config;
@@ -1,20 +0,0 @@
-
-#[derive(Copy,Clone,Eq,PartialEq,Debug)]
-pub enum OsPlatform {
-  Linux,
-  MacOs,
-  Windows,
-}
-
-impl OsPlatform {
-  pub fn get() -> OsPlatform {
-    // Totally scientific.
-    // https://doc.rust-lang.org/std/env/consts/constant.OS.html
-    match std::env::consts::OS {
-      "linux" => OsPlatform::Linux,
-      "windows" => OsPlatform::Windows,
-      "macos" | "apple" => OsPlatform::MacOs,
-      _ => OsPlatform::Linux,
-    }
-  }
-}
@@ -1,10 +1,7 @@
-use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
 use std::fs::File;
 use std::io::BufReader;
 use std::path::Path;

-const CONFIG_FILENAME : &str = "app_config.yaml";
-
 /// Configuration read or derived from a YAML config file at startup
 /// This is not meant to contain any runtime config values.
 /// This may not be exhaustive.
@@ -16,15 +13,6 @@ pub struct YamlConfig {
  pub scheduler_samples: Option<usize>,
  pub seed: Option<u64>,
  pub cfg_scale: Option<f64>,
-
-  /// Override the default location to store data on Linux
-  pub linux_default_data_path: Option<String>,
-
-  /// Override the default location to store data on Windows
-  pub windows_default_data_path: Option<String>,
-  
-  /// Override the default location to store data on Mac
-  pub mac_default_data_path: Option<String>,
 }

 impl YamlConfig {
@@ -34,19 +22,4 @@ impl YamlConfig {
    let config = serde_yaml::from_reader(reader)?;
    Ok(config)
  }
-
-  pub fn load_from_config_file_recursive() -> YamlConfig {
-    let maybe_path = recursively_find_file_by_name(CONFIG_FILENAME, "../../../", 5)
-      .ok()
-      .flatten();
-    if let Some(path) = maybe_path {
-      let path = path.canonicalize().unwrap_or(path);
-      println!("Attempting to load configs from: {:?}", &path);
-      if let Ok(config) = YamlConfig::read_from_file(path) {
-        return config;
-      }
-    }
-    println!("Couldn't load config file. Using default configs.");
-    YamlConfig::default()
-  }
 }
@@ -1,11 +0,0 @@
-use tauri::{AppHandle, Emitter};
-use crate::state::app_dir::AppDataRoot;
-
-pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! {
-  loop {
-    app.emit("testing", 1);
-    
-    println!(">>> Testing: {:?}", app_data_root.path());
-    std::thread::sleep(std::time::Duration::from_millis(500));
-  }
-}
@@ -1 +0,0 @@
-pub mod downloader_thread;
@@ -1,6 +1,6 @@
 {
  "$schema": "https://schema.tauri.app/config/2.0.0-rc",
-  "productName": "ArtCraft",
+  "productName": "Realtime Desktop App",
  "version": "0.1.0",
  "identifier": "com.unique.value.thing.why",
  "build": {
@@ -12,7 +12,7 @@
  "app": {
    "windows": [
      {
-        "title": "ArtCraft",
+        "title": "Realtime Desktop App Title Goes Here",
        "width": 1200,
        "height": 800,
        "resizable": true,
@@ -86,8 +86,6 @@ fn main() -> Result<()> {
        .arg("--expt-relaxed-constexpr")
        .arg("--expt-extended-lambda")
        .arg("--use_fast_math")
-        .arg("--compiler-options")
-        .arg("-fPIC")
        .arg("--verbose");

    if let Ok(target) = std::env::var("TARGET") {
@@ -32,4 +32,3 @@ metal = ["candle/metal", "dep:candle-metal-kernels", "dep:metal"]
 [[bench]]
 name = "bench_main"
 harness = false
-
@@ -84,8 +84,6 @@ impl StableDiffusionConfig {
        };
        // https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/unet/config.json
        let unet = unet_2d::UNet2DConditionModelConfig {
-            time_embed_dim: None,
-            time_cond_proj_dim: None,
            blocks: vec![
                bc(320, Some(1), 8),
                bc(640, Some(1), 8),
@@ -155,8 +153,6 @@ impl StableDiffusionConfig {
        };
        // https://huggingface.co/stabilityai/stable-diffusion-2-1/blob/main/unet/config.json
        let unet = unet_2d::UNet2DConditionModelConfig {
-            time_cond_proj_dim: None,
-            time_embed_dim: None,
            blocks: vec![
                bc(320, Some(1), 5),
                bc(640, Some(1), 10),
@@ -241,8 +237,6 @@ impl StableDiffusionConfig {
        };
        // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json
        let unet = unet_2d::UNet2DConditionModelConfig {
-            time_cond_proj_dim: None,
-            time_embed_dim: None,
            blocks: vec![
                bc(320, None, 5),
                bc(640, Some(2), 10),
@@ -312,8 +306,6 @@ impl StableDiffusionConfig {
        };
        // https://huggingface.co/stabilityai/sdxl-turbo/blob/main/unet/config.json
        let unet = unet_2d::UNet2DConditionModelConfig {
-            time_cond_proj_dim: None,
-            time_embed_dim: None,
            blocks: vec![
                bc(320, None, 5),
                bc(640, Some(2), 10),
@@ -413,8 +405,6 @@ impl StableDiffusionConfig {
        };
        // https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/unet/config.json
        let unet = unet_2d::UNet2DConditionModelConfig {
-            time_cond_proj_dim: None,
-            time_embed_dim: None,
            blocks: vec![
                bc(320, None, 5),
                bc(640, Some(2), 10),
@@ -10,10 +10,6 @@ use candle_nn as nn;
 use candle_nn::Module;
 use std::collections::HashMap;

-// fn save_tensor_as_numpy(tensor: &Tensor, filename: &str) -> Result<()> {
-//     // Use Candle's built-in NPY serialization
-//     tensor.write_npy(filename).map_err(|e| anyhow::anyhow!("{}", e))
-// }
 #[derive(Debug, Clone, Copy)]
 pub struct BlockConfig {
  pub out_channels: usize,
@@ -37,13 +33,11 @@ pub struct UNet2DConditionModelConfig {
  pub cross_attention_dim: usize,
  pub sliced_attention_size: Option<usize>,
  pub use_linear_projection: bool,
-  pub time_cond_proj_dim: Option<usize>,
-  pub time_embed_dim: Option<usize>,
 }

 impl Default for UNet2DConditionModelConfig {
  fn default() -> Self {
-    Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false, time_cond_proj_dim: Some(256), time_embed_dim: None }
+    Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false }
  }
 }

@@ -79,16 +73,8 @@ impl UNet2DConditionModel {
    let b_channels = config.blocks[0].out_channels;
    let bl_channels = config.blocks.last().unwrap().out_channels;
    let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
-    let time_embed_dim = config.time_embed_dim.unwrap_or(b_channels * 4);
+    let time_embed_dim = b_channels * 4;
    let conv_cfg = nn::Conv2dConfig { padding: 1, ..Default::default() };
-
-    // Log important configuration details
-    println!("Creating UNet2DConditionModel with {} blocks", n_blocks);
-    if let Some(proj_dim) = config.time_cond_proj_dim {
-      println!("UNet configured with time_cond_proj_dim: {}", proj_dim);
-      println!("This means the model expects guidance embeddings of dimension {}", proj_dim);
-    }
-
    let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;

    let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
@@ -255,138 +241,18 @@ impl UNet2DConditionModel {

    // Add guidance embedding if provided
    if let Some(guidance_emb) = guidance_embedding {
-      // Get dimensions of both embeddings
-      let (batch_size, time_emb_dim) = time_emb.dims2()?;
-      let (guidance_batch, guidance_emb_dim) = guidance_emb.dims2()?;
+      // Print shapes for debugging
+      println!("Guidance embedding shape: {:?}", guidance_emb.shape());
+      println!("Time embedding shape: {:?}", time_emb.shape());

-      println!("==== Guidance Embedding Integration ====");
-      println!("Time embedding dims: [{}, {}]", batch_size, time_emb_dim);
-      println!("Guidance embedding dims: [{}, {}]", guidance_batch, guidance_emb_dim);
+      // Get the batch size and embedding dimension
+      let (batch_size, emb_dim) = time_emb.dims2()?;

-      // Use time_cond_proj_dim from config to determine expected guidance embedding dimension
-      let expected_guidance_dim = self.config.time_cond_proj_dim.unwrap_or(time_emb_dim);
-      println!("Expected guidance dimension from config: {}", expected_guidance_dim);
-
-      // Check if dimensions match the expected projection dimension
-      let processed_guidance_emb = if guidance_emb_dim != time_emb_dim {
-        // Convert guidance_emb to f32 for calculations
-        let guidance_emb_f32 = guidance_emb.to_dtype(DType::F32)?;
-        println!("Converting guidance embedding to F32 for processing");
-
-        // If the guidance embedding dimension matches the expected projection dimension,
-        // we need to project it to match the time embedding dimension
-        if guidance_emb_dim == expected_guidance_dim {
-          println!("Projecting guidance embedding from {} to {}", guidance_emb_dim, time_emb_dim);
-
-          // Use linear interpolation
-          let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
-          let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
-
-          // Get guidance embedding data (now f32)
-          let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
-          println!("Reshaping guidance data to vector of length {}", guidance_data.len());
-
-          // Perform linear interpolation
-          let mut interpolated_data = Vec::with_capacity(time_emb_dim);
-          for &t in &dst_positions {
-            // Find the two source positions that surround the target position
-            let mut i = 0;
-            while i < src_positions.len() - 1 && src_positions[i + 1] < t {
-              i += 1;
-            }
-
-            if i == src_positions.len() - 1 {
-              // At or beyond the last position
-              interpolated_data.push(guidance_data[i]);
-            } else if t <= src_positions[0] {
-              // At or before the first position
-              interpolated_data.push(guidance_data[0]);
-            } else {
-              // Linear interpolation between two points
-              let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
-              let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
-              interpolated_data.push(value);
-            }
-          }
-
-          println!("Created interpolated data of length {}", interpolated_data.len());
-
-          // Create tensor and convert to time_emb's dtype
-          let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
-          let result = tensor.to_dtype(time_emb.dtype())?;
-          println!("Final projected guidance embedding shape: {:?}", result.shape());
-          result
-        } else {
-          // Unusual case - dimensions don't match expected pattern
-          println!("Warning: Guidance embedding dimension ({}) doesn't match expected projection dimension ({})", guidance_emb_dim, expected_guidance_dim);
-
-          // Fallback to generic interpolation
-          let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
-
-          let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
-          let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
-
-          let mut interpolated_data = Vec::with_capacity(time_emb_dim);
-          for &t in &dst_positions {
-            let mut i = 0;
-            while i < src_positions.len() - 1 && src_positions[i + 1] < t {
-              i += 1;
-            }
-
-            if i == src_positions.len() - 1 {
-              interpolated_data.push(guidance_data[i]);
-            } else if t <= src_positions[0] {
-              interpolated_data.push(guidance_data[0]);
-            } else {
-              let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
-              let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
-              interpolated_data.push(value);
-            }
-          }
-
-          println!("Created interpolated data of length {}", interpolated_data.len());
-          let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
-          let result = tensor.to_dtype(time_emb.dtype())?;
-          println!("Final fallback guidance embedding shape: {:?}", result.shape());
-          result
-        }
-      } else {
-        // Dimensions match, no processing needed
-        println!("Guidance dimension already matches time embedding dimension, no interpolation needed");
-        let result = guidance_emb.reshape((1, time_emb_dim))?.to_dtype(time_emb.dtype())?;
-        println!("Reshaped guidance embedding shape: {:?}", result.shape());
-        result
-      };
-
-      // Expand to match batch size if needed
-      let expanded_guidance = processed_guidance_emb.broadcast_as(time_emb.shape())?;
-      println!("Expanded guidance embedding shape: {:?}", expanded_guidance.shape());
-
-      // Save the processed and expanded guidance embeddings to disk if requested
-      if let Ok(save_dir) = std::env::var("CANDLE_SAVE_TENSORS_DIR") {
-        // Save pre-expansion (original projection)
-        let proc_path = format!("{}/unet_processed_guidance_embedding.npy", save_dir);
-        println!("Saving processed guidance embedding to {}", proc_path);
-        if let Err(e) = processed_guidance_emb.write_npy(&proc_path) {
-          println!("Error saving processed guidance embedding: {}", e);
-        } else {
-          println!("Successfully saved processed guidance embedding to {}", proc_path);
-        }
-
-        // Save post-expansion (final tensor used for addition to time_emb)
-        let exp_path = format!("{}/unet_expanded_guidance_embedding.npy", save_dir);
-        println!("Saving expanded guidance embedding to {}", exp_path);
-        if let Err(e) = expanded_guidance.write_npy(&exp_path) {
-          println!("Error saving expanded guidance embedding: {}", e);
-        } else {
-          println!("Successfully saved expanded guidance embedding to {}", exp_path);
-        }
-      }
+      // Reshape guidance embedding to match time embedding dimensions
+      let guidance_emb = guidance_emb.reshape(&[1, emb_dim])?.expand(&[batch_size, emb_dim])?;

      // Add to time embeddings
-      time_emb = (time_emb + expanded_guidance)?;
-      println!("Final time embedding with guidance shape: {:?}", time_emb.shape());
-      println!("==== End Guidance Integration ====");
+      time_emb = (time_emb + guidance_emb)?;
    }

    // 2. pre-process
@@ -546,7 +546,7 @@ export class Engine {
          canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
          canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
          shapeType: shapeNode.shapeType,
-          size: shapeNode.kNode.size(),
+          size: shapeNode.size,
          color: nodeColor.color,
          mediaLayerRef: this.mediaLayer,
          selectionManagerRef: this.selectionManager,
@@ -581,7 +581,7 @@ export class Engine {
      canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
      canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
      shapeType: type,
-      size: { width: size, height: size },
+      size: size,
      color: color,
      mediaLayerRef: this.mediaLayer,
      selectionManagerRef: this.selectionManager,
@@ -7,11 +7,9 @@ import { RealTimeDrawEngine } from "../RenderingPrimitives/RealTimeDrawEngine";
 import { MouseEventHandler } from "react";
 import { KonvaEventObject } from "konva/lib/Node";

-export type SupportedNode = Konva.Image | Konva.Group | Konva.Shape;
-
 export abstract class BaseNode {
  protected selectionManagerRef: SelectionManager;
-  public kNode: SupportedNode;
+  public kNode: Konva.Image | Konva.Group;
  protected mediaLayerRef: Konva.Layer;

  // Internal State members
@@ -33,7 +31,7 @@ export abstract class BaseNode {
  }: {
    selectionManagerRef: SelectionManager;
    mediaLayerRef: Konva.Layer;
-    kNode: SupportedNode;
+    kNode: Konva.Image | Konva.Group | Konva.Shape;
  }) {
    // console.log("Node constructed");
    // this.uuid = uuidv4();
@@ -64,7 +62,7 @@ export abstract class BaseNode {
  }
  public highlight() {
    // console.log("Highlight", this.kNode._id);
-    if (this.kNode instanceof Konva.Image || this.kNode instanceof Konva.Shape) {
+    if (this.kNode instanceof Konva.Image) {
      this.kNode.stroke(primaryOrange);
      this.kNode.strokeWidth(highlightStrokeWidth);
      return;
@@ -87,7 +85,7 @@ export abstract class BaseNode {

  public unhighlight() {
    // console.log("unHighlight", this.kNode._id);
-    if (this.kNode instanceof Konva.Image || this.kNode instanceof Konva.Shape) {
+    if (this.kNode instanceof Konva.Image) {
      this.kNode.strokeWidth(0);
      return;
    }
@@ -1,7 +1,7 @@
 import Konva from "konva";
 import { v4 as uuidv4 } from "uuid";
 import { SelectionManager } from "../NodesManagers";
-import { BaseNode, SupportedNode } from "./BaseNode";
+import { BaseNode } from "./BaseNode";
 import { Size } from "../types";
 import { LoadingVideosProvider } from "../EngineUtitlities/LoadingVideosProvider";
 import { NodeProgressEventDetail } from "../types/events";
@@ -20,7 +20,7 @@ export enum UploadStatus {
 }

 export abstract class NetworkedNode extends BaseNode {
-  public kNode: SupportedNode;
+  public kNode: Konva.Image;

  // members to deal with loading progress
  public didFinishLoading: boolean = false;
@@ -49,7 +49,7 @@ export abstract class NetworkedNode extends BaseNode {
    localFile,
    loadingVideosProviderRef,
  }: {
-    kNode: SupportedNode;
+    kNode: Konva.Image;
    selectionManagerRef: SelectionManager;
    mediaLayerRef: Konva.Layer;
    loadingVideosProviderRef?: LoadingVideosProvider;
@@ -15,7 +15,7 @@ interface ShapeNodeConstructor {
  canvasPosition: Position;
  canvasSize: Size;
  shapeType: ShapeType;
-  size: Size;
+  size: number;
  color?: string; // Hex color string
  transform?: TransformationData;
  mediaLayerRef: Konva.Layer;
@@ -29,7 +29,7 @@ export class ShapeNode extends NetworkedNode {
  }

  public shapeType: ShapeType;
-  public size: Size;
+  public size: number;
  public color: string;

  constructor({
@@ -46,7 +46,6 @@ export class ShapeNode extends NetworkedNode {
    // Create the actual shape inside the group
    const transform = NodeUtilities.getInitialTransform({
      existingTransform,
-      mediaFileSize: size,
      canvasPosition,
      canvasSize,
    });
@@ -55,48 +54,110 @@ export class ShapeNode extends NetworkedNode {

    switch (shapeType) {
      case ShapeType.CIRCLE:
-        shape = new Konva.Ellipse({
-          x: size.width,
-          y: size.height,
-          radiusX: size.width,
-          radiusY: size.height,
+        const circleShape = new Konva.Circle({
+          radius: 50,
          ...transform,
          fill: color,
          strokeScaleEnabled: false,
          draggable: true,
        });
+
+        // Convert circle to image
+        const dataURL = circleShape.toDataURL();
+        const imageObj = new Image();
+        imageObj.src = dataURL;
+
+        shape = new Konva.Image({
+          image: imageObj,
+          x: circleShape.x(),
+          y: circleShape.y(),
+          width: circleShape.width(),
+          height: circleShape.height(),
+          ...transform,
+          fill: "transparent",
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        imageObj.onload = () => {
+          this.mediaLayerRef.draw();
+          loaded();
+        };
+
+        // Clean up the temporary circle
+        circleShape.destroy();
        break;

      case ShapeType.SQUARE:
-        shape = new Konva.Rect({
-          width: size.width,
-          height: size.height,
+        const squareShape = new Konva.Rect({
+          width: 100,
+          height: 100,
          ...transform,
          fill: color,
          strokeScaleEnabled: false,
          draggable: true,
        });
+
+        // Convert square to image
+        const squareDataURL = squareShape.toDataURL();
+        const squareImageObj = new Image();
+        squareImageObj.src = squareDataURL;
+
+        shape = new Konva.Image({
+          image: squareImageObj,
+          x: squareShape.x(),
+          y: squareShape.y(),
+          width: squareShape.width(),
+          height: squareShape.height(),
+          ...transform,
+          fill: "transparent",
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        squareImageObj.onload = () => {
+          this.mediaLayerRef.draw();
+          loaded();
+        };
+
+        // Clean up the temporary square
+        squareShape.destroy();
        break;

      case ShapeType.TRIANGLE:
-        shape = new Konva.Shape({
-          x: 0,
-          y: 0,
-          width: size.width,
-          height: size.height,
+        const triangleShape = new Konva.RegularPolygon({
+          sides: 3,
+          radius: 50,
          ...transform,
          fill: color,
          strokeScaleEnabled: false,
          draggable: true,
-          sceneFunc: (context, shape) => {
-            context.beginPath();
-            context.moveTo(shape.width() / 2, 0);
-            context.lineTo(shape.width(), shape.height());
-            context.lineTo(0, shape.height());
-            context.closePath();
-            context.fillStrokeShape(shape);
-          }
        });
+
+        // Convert triangle to image
+        const triangleDataURL = triangleShape.toDataURL();
+        const triangleImageObj = new Image();
+        triangleImageObj.src = triangleDataURL;
+
+        shape = new Konva.Image({
+          image: triangleImageObj,
+          x: triangleShape.x(),
+          y: triangleShape.y(),
+          width: triangleShape.width(),
+          height: triangleShape.height(),
+          ...transform,
+          fill: "transparent",
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        triangleImageObj.onload = () => {
+          this.mediaLayerRef.draw();
+          loaded();
+        };
+
+        // Clean up the temporary triangle
+        triangleShape.destroy();
        break;

      default:
@@ -4,7 +4,7 @@ import { Shape } from "konva/lib/Shape";
 import { Group } from "konva/lib/Group";

 import { invoke } from "@tauri-apps/api/core";
-import { listen } from "@tauri-apps/api/event";
+import { listen } from '@tauri-apps/api/event';

 import { FileUtilities } from "../FileUtilities/FileUtilities";
 import { ImageNode, VideoNode, TextNode, ShapeNode, ShapeType } from "../Nodes";
@@ -12,7 +12,7 @@ import { MediaNode } from "../types";

 import { RenderTask } from "./RenderTask";
 import { OffScreenSceneCanvas } from "./OffScreenSceneCanvas";
-
+import { Image } from "@tauri-apps/api/image";

 import { PaintNode } from "../Nodes/PaintNode";

@@ -31,6 +31,7 @@ interface ServerSettings {
  lora_path?: string;
 }

+
 interface GenerateImageParams {
  image: string;
  prompt: string;
@@ -61,13 +62,9 @@ export class RealTimeDrawEngine {
  private positionPreviewY: number;

  private port: MessagePort | undefined;
+
  public captureCanvas: Konva.Rect;
-  public backgroundRasterRect: Konva.Image;
-  
-
  public previewCanvas: Konva.Image;
-  
-

  public videoLoadingCanvas: VideoNode | undefined;

@@ -91,15 +88,13 @@ export class RealTimeDrawEngine {
      y: number;
    },
  ) => void;
-  private offscreenRenderDiv: HTMLDivElement;

  private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback

  private client: WebSocketClient | null = null;
  private isConnected: boolean = false;

-
-
+  private backgroundNode: ShapeNode | null = null;
  public backgroundColor: string = "#d2d2d2";
  constructor({
    width,
@@ -189,33 +184,25 @@ export class RealTimeDrawEngine {
      fill: "white",
    });

-    this.offscreenRenderDiv = document.createElement("div");
-
-    this.backgroundRasterRect = new Konva.Image({
-      name: "backgroundRasterRect",
-      x: this.positionX,
-      y: this.positionY,
-      width: this.width,
-      height: this.height,
-      fill: this.backgroundColor,
-      image: undefined,
-      stroke: "black",
-      strokeWidth: 1,
-      draggable: false,
-    });
- 
    this.mediaLayerRef.add(this.captureCanvas);
-    this.mediaLayerRef.add(this.backgroundRasterRect);
    this.mediaLayerRef.add(this.previewCanvas);
    // send back
    this.captureCanvas.setZIndex(0);
-    this.backgroundRasterRect.setZIndex(0);
    this.previewCanvas.setZIndex(1);
    // Add mouse events for preview canvas copying
    this.previewCopyListener();
-    
-    //this.listenToServerEvents();
-   
+
+    //this.startServer();
+    // isLoadingVisible.value = true;
+    // const fakeTimer = setInterval(() => {
+    //   loadingProgress.value += 1;
+    //   if (loadingProgress.value >= 100) {
+    //     clearInterval(fakeTimer);
+    //     isLoadingVisible.value = false;
+    //   }
+    // }, 1000);
+
+    this.listenToServerEvents();
  }

  public async listenToServerEvents() {
@@ -223,18 +210,16 @@ export class RealTimeDrawEngine {
      console.log(event);
      // Handle model download events
      const payload = event.payload as any;
-
+      
      // Model download started
      if (payload.model_download_started) {
        const modelInfo = payload.model_download_started;
-        console.log(
-          `Model download started: ${modelInfo.model_name} (${modelInfo.model_type})`,
-        );
-
+        console.log(`Model download started: ${modelInfo.model_name} (${modelInfo.model_type})`);
+        
        // Set up loading indicator
        isLoadingVisible.value = true;
        loadingProgress.value = 0;
-
+        
        // Create fake progress updates
        const downloadTimer = setInterval(() => {
          loadingProgress.value += 2;
@@ -245,23 +230,23 @@ export class RealTimeDrawEngine {
          }
        }, 500);
      }
-
+      
      // Model download completed
      if (payload.model_download_complete) {
        const modelInfo = payload.model_download_complete;
-        console.log(
-          `Model download completed: ${modelInfo.model_name} (${modelInfo.model_type})`,
-        );
-
+        console.log(`Model download completed: ${modelInfo.model_name} (${modelInfo.model_type})`);
+        
        // Complete the loading progress
        loadingProgress.value = 100;
-
+        
        // Hide loading indicator after a short delay
        setTimeout(() => {
          isLoadingVisible.value = false;
        }, 1000);
      }
+  
    });
+    
  }

  private isEnabled: boolean = false;
@@ -335,16 +320,9 @@ export class RealTimeDrawEngine {
    const ctx = cursorCanvas.getContext("2d");
    if (!ctx) return;

-    // Draw the outer circle with a light stroke
+    // Draw the circle
    ctx.beginPath();
    ctx.arc(size, size, size / 2, 0, Math.PI * 2);
-    ctx.strokeStyle = "white";
-    ctx.lineWidth = 3;
-    ctx.stroke();
-
-    // Draw the inner circle with a dark stroke
-    ctx.beginPath();
-    ctx.arc(size, size, size / 2 - 1, 0, Math.PI * 2);
    ctx.strokeStyle = "black";
    ctx.lineWidth = 1;
    ctx.stroke();
@@ -807,27 +785,6 @@ export class RealTimeDrawEngine {
    }
  }

-  private cloneStageForRender(stage: Konva.Stage, layerOfInterest: Konva.Layer): Konva.Stage {
-    const stageClone = new Konva.Stage({
-      width: stage.width(),
-      height: stage.height(),
-      container: this.offscreenRenderDiv,
-    });
-
-    const renderLayer = new Konva.Layer();
-
-    // Clone all the nodes then reset them to render right
-    layerOfInterest.getChildren().forEach((node) => {
-      const dupNode = node.clone();
-      dupNode.strokeWidth(0);
-      renderLayer.add(dupNode);
-    });
-
-    stageClone.add(renderLayer);
-
-    return stageClone;
-  }
-
  private async renderFrame(config: {
    layerOfInterest: Konva.Layer; // layer where the element that you want to clip lives.
    // XY and height of a captureCanvas ( region of interest )
@@ -848,25 +805,54 @@ export class RealTimeDrawEngine {
      const y = config.y !== undefined ? config.y : Math.floor(box.y);
      const pixelRatio = config.pixelRatio || 1;

-      // Clone the required layer from the stage
-      // Set the right details (like removing highlight stroke)
-      // Then render the cloned stage to a bitmap
-      const stageClone = this.cloneStageForRender(stage, config.layerOfInterest);
-      const stageBlob = await stageClone.toBlob({
-        x: x,
-        y: y,
-        width: config.width || Math.ceil(box.width),
-        height: config.height || Math.ceil(box.height),
+      const container = config.layerOfInterest as Container<Group | Shape>;
+
+      const offScreenSceneCanvas = new OffScreenSceneCanvas({
+        width:
+          config.width || Math.ceil(box.width) || (stage ? stage.width() : 0),
+        height:
+          config.height ||
+          Math.ceil(box.height) ||
+          (stage ? stage.height() : 0),
        pixelRatio: pixelRatio,
-      }) as Blob;
+      });
+
+      const context = offScreenSceneCanvas.getContext();
+
+      const buffer = new OffScreenSceneCanvas({
+        width:
+          offScreenSceneCanvas.width / offScreenSceneCanvas.pixelRatio +
+          Math.abs(x),
+        height:
+          offScreenSceneCanvas.height / offScreenSceneCanvas.pixelRatio +
+          Math.abs(y),
+        pixelRatio: offScreenSceneCanvas.pixelRatio,
+      });
+
+      context.save();
+
+      if (x || y) {
+        context.translate(-1 * x, -1 * y);
+      }
+      container.drawScene(offScreenSceneCanvas, undefined, buffer);
+      // Not a type mistake ... DO NOT FIX
+      // @ts-ignore
+      const offscreenCanvas = offScreenSceneCanvas._canvas as OffscreenCanvas;
+      let result = undefined;

      // if config.test is true, the result is downloaded to the local files
      // config.test = true;
      if (config.test) {
-        await FileUtilities.blobToFileJpeg(stageBlob, "1");
+        const blob = await offscreenCanvas.convertToBlob({
+          quality: config.quality ?? 1.0,
+          type: "image/jpeg",
+        });
+        await FileUtilities.blobToFileJpeg(blob, "1");
+        result = blob;
+      } else {
+        result = offscreenCanvas.transferToImageBitmap();
      }
-
-      const result = await createImageBitmap(stageBlob);
+      context.restore();
      return result;
    } catch (error) {
      throw error;
@@ -979,8 +965,8 @@ export class RealTimeDrawEngine {

    // Test code
    if (false) {
-      this.outputBitmap = bitmap;
-      this.previewCanvas.image(bitmap);
+      //this.outputBitmap = bitmap;
+      //this.previewCanvas.image(bitmap);
      const base64Bitmap = await this.imageBitmapToBase64(bitmap);
      console.log(this.currentStrength);
      await this.client?.generateImage({
@@ -993,7 +979,6 @@ export class RealTimeDrawEngine {
        width: 1024,
        lora_strength: 1.0,
      });
-      this.isProcessing = false
      return;
    }

@@ -1049,19 +1034,6 @@ export class RealTimeDrawEngine {
  // Add method to create or update background
  public updateBackground(color: string) {
    this.captureCanvas.fill(color);
-
-    const captureCanvasImage = this.captureCanvas.toDataURL();
-    const imageSource = new Image();
-    imageSource.src = captureCanvasImage;
-
-    imageSource.onload = () => {      
-      this.backgroundRasterRect.fill(color);
-      this.backgroundRasterRect.image(imageSource)
-
-      this.mediaLayerRef.batchDraw();
-      this.render();
-    };
-
-
+    this.mediaLayerRef.batchDraw();
  }
 }
@@ -6,21 +6,6 @@
  :root {
    @apply select-none bg-ui-background font-sans text-white;
  }
-
-  /* Remove all focus styles */
-  button {
-    @apply outline-none ring-0;
-    -webkit-tap-highlight-color: transparent !important;
-    outline: none !important;
-    box-shadow: none !important;
-  }
-
-  button:focus {
-    @apply outline-none ring-0;
-    outline: none !important;
-    box-shadow: none !important;
-  }
-
  .pegboard {
    background-image: radial-gradient(#ffffff25 1px, transparent 0);
    background-color: #1b1b1f;
@@ -202,30 +187,3 @@
    background-position: -200% 0;
  }
 }
-
-.react-colorful__saturation {
-  margin: 16px 0;
-  border-radius: 10px !important;
-  border-bottom: none !important;
-}
-
-.react-colorful__saturation-pointer {
-  width: 24px !important;
-  height: 24px !important;
-}
-
-.react-colorful__hue,
-.react-colorful__alpha {
-  height: 20px !important;
-  border-radius: 50px !important;
-}
-
-.react-colorful__hue {
-  margin-bottom: 12px;
-}
-
-.react-colorful__hue-pointer,
-.react-colorful__alpha-pointer {
-  width: 20px !important;
-  height: 20px !important;
-}
@@ -58,12 +58,6 @@ export const ToolbarButton = ({
      className={mergedButtonClasses}
      disabled={disabled}
      autoFocus={false}
-      tabIndex={-1}
-      style={{
-        outline: "none",
-        boxShadow: "none",
-        WebkitTapHighlightColor: "transparent",
-      }}
      {...restButtonProps}
      onClick={(e) => {
        e.preventDefault();
@@ -28,17 +28,17 @@ export const Button = ({
  function getVariantClassNames(variant: string) {
    switch (variant) {
      case "tertiary": {
-        return "bg-tertiary hover:bg-tertiary-400 text-white";
+        return "bg-tertiary hover:bg-tertiary-400 text-white focus-visible:outline-tertiary-600";
      }
      case "secondary": {
-        return "bg-ui-controls hover:bg-ui-controls/80 text-white";
+        return "bg-ui-controls hover:bg-ui-controls/80 text-white focus-visible:outline-secondary";
      }
      case "action": {
-        return "bg-action hover:bg-action-500 text-white";
+        return " bg-action hover:bg-action-500 text-white focus-visible:outline-action";
      }
      case "primary":
      default: {
-        return "bg-primary hover:bg-primary-600 text-white";
+        return "bg-primary hover:bg-primary-600 text-white focus-visible:outline-primary-600";
      }
    }
  }
@@ -51,6 +51,7 @@ export const Button = ({
    "text-[15px] font-semibold rounded-lg px-3 py-2 shadow-sm",
    "flex gap-2 items-center justify-center",
    "transition-all duration-150",
+    "focus-visible:outline focus-visible:outline-2 focus-visible:outline-offset-2",
    getVariantClassNames(propsVariant),
    propsClassName,
    disabledClass,
@@ -1,13 +1,12 @@
+import { MouseEvent, MouseEventHandler, useState } from "react";
+import { HexColorPicker } from "react-colorful";
+import { EyeDropper } from "react-eyedrop";
 import { Button } from "./Button";
 import { faEyedropper } from "@fortawesome/pro-solid-svg-icons";
 import { twMerge } from "tailwind-merge";

-type ColorPickerHTMLAttributes = Omit<
-  React.HTMLAttributes<HTMLDivElement>,
-  "color" | "onChange" | "onChangeCapture"
->;
-interface ColorPickerBaseProps<T extends string>
-  extends ColorPickerHTMLAttributes {
+type ColorPickerHTMLAttributes = Omit<React.HTMLAttributes<HTMLDivElement>, "color" | "onChange" | "onChangeCapture">;
+interface ColorPickerBaseProps<T extends string> extends ColorPickerHTMLAttributes {
  color: T;
  onChange: (newColor: T) => void;
 }
@@ -18,7 +17,7 @@ export const HexEyedropPicker = ({
  dropperClassName,
  pickerClassName,
  onPickerChange,
-  Picker,
+  Picker
 }: {
  color: string;
  onDropperChange?: (newColor: string) => void;
@@ -27,8 +26,9 @@ export const HexEyedropPicker = ({
  onPickerChange: (newColor: string) => void;
  Picker: (props: Partial<ColorPickerBaseProps<string>>) => JSX.Element;
 }) => {
+
  const handleDropperClick = () => {
-    if (!("EyeDropper" in window)) {
+    if (!('EyeDropper' in window)) {
      // NOTE: This shouldn't happen cuz we're using Chromium
      console.error("EyeDropper not found in window");
      return;
@@ -46,40 +46,42 @@ export const HexEyedropPicker = ({
        // User cancelled the picker
        return null;
      }
-    };
+    }

    startEyeDrop().then((newColor) => {
      if (newColor) {
        onDropperChange?.(newColor);
      }
    });
-  };
+  }

-  // const [isDragging, setIsDragging] = useState(false);
-  // const handleOnMouseDown = () => {
-  //   setIsDragging(true);
-  // };
-  // const handleOnMouseUp = () => {
-  //   setIsDragging(false);
-  // };
+  const [isDragging, setIsDragging] = useState(false);
+  const handleOnMouseDown = () => {
+    setIsDragging(true);
+  }
+  const handleOnMouseUp = () => {
+    setIsDragging(false);
+  }

  return (
    <div className={"relative"}>
-      <Button
-        className={twMerge([
-          "z-20 size-8 rounded-xl bg-ui-background/30 text-white shadow-none duration-200 hover:bg-white/10 hover:text-white hover:shadow-xl",
-          dropperClassName,
-        ])}
-        icon={faEyedropper}
-        onClick={handleDropperClick}
-      />
      <Picker
        className={pickerClassName}
        color={color}
        onChange={onPickerChange}
-        // onMouseDown={handleOnMouseDown}
-        // onMouseUp={handleOnMouseUp}
+        onMouseDown={handleOnMouseDown}
+        onMouseUp={handleOnMouseUp}
+      />
+      <Button
+        className={twMerge([
+          "rounded-full m-1 size-8 text-ui-controls bg-transparent hover:bg-black/70 hover:text-white z-20 shadow-none hover:shadow-xl duration-200",
+          dropperClassName,
+          "absolute left-0 top-0",
+          isDragging ? "opacity-0 -z-10" : ""
+        ])}
+        icon={faEyedropper}
+        onClick={handleDropperClick}
      />
    </div>
-  );
-};
+  )
+}
@@ -91,7 +91,7 @@ export const ColorPicker = ({
    if (streamChanges) {
      onChange(color);
    }
-  };
+  }

  return (
    <Popover className="relative">
@@ -124,7 +124,7 @@ export const ColorPicker = ({
                anchor={anchor}
                className={twMerge(
                  paperWrapperStyles,
-                  "flex flex-col items-center gap-2 overflow-hidden p-4",
+                  "flex flex-col items-center gap-2 overflow-hidden",
                )}
                style={
                  {
@@ -134,14 +134,13 @@ export const ColorPicker = ({
                static
              >
                <HexEyedropPicker
+                  pickerClassName="overflow-hidden"
                  color={currColor}
                  onPickerChange={handleColorChange}
                  onDropperChange={handleColorChange}
                  Picker={HexAlphaColorPicker}
                />
                <Input
-                  className="mt-2"
-                  inputClassName="bg-ui-background/30"
                  style={{ width: "198px" }}
                  icon={faHashtag}
                  value={textInput}
@@ -5,18 +5,13 @@ import {
 } from "@fortawesome/pro-solid-svg-icons";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
-import { ChangeEvent, useEffect, useState, useRef } from "react";
+import { ChangeEvent, useEffect, useState } from "react";
 import { HexAlphaColorPicker } from "react-colorful";
 import { twMerge } from "tailwind-merge";

 import { useSignals } from "@preact/signals-react/runtime";
 import { paperWrapperStyles } from "~/components/styles";
-import {
-  BRUSH_MAX_SIZE,
-  BRUSH_MIN_SIZE,
-  paintBrushSize,
-  setPaintBrushSize,
-} from "~/signals/uiEvents/toolbarMain/paintMode";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, paintBrushSize, setPaintBrushSize } from "~/signals/uiEvents/toolbarMain/paintMode";
 import { Button } from "../Button";
 import { Input } from "../Input";
 import { Slider } from "../Slider";
@@ -57,58 +52,16 @@ export const PaintModeMenu = ({
    currColor: prevColor,
    textInput: prevColor.substring(1),
  });
-  const [isMouseDown, setIsMouseDown] = useState(false);
-  const [isMouseOutside, setIsMouseOutside] = useState(false);
-  const closeRef = useRef<(() => void) | null>(null);
-  const closeTimeoutRef = useRef<NodeJS.Timeout | null>(null);

  useEffect(() => {
    setIsOpen(defaultOpen);
  }, [defaultOpen]);

-  // Global mouse up handler
-  useEffect(() => {
-    const handleGlobalMouseUp = () => {
-      if (isMouseDown) {
-        setIsMouseDown(false);
-
-        // If mouse is outside when released, close the panel after a short delay
-        if (isMouseOutside && closeOnMouseLeave && closeRef.current) {
-          closeTimeoutRef.current = setTimeout(() => {
-            onChange(currColor);
-            setIsOpen(false);
-            closeRef.current?.();
-          }, 500);
-        }
-      }
-    };
-
-    window.addEventListener("mouseup", handleGlobalMouseUp);
-
-    return () => {
-      window.removeEventListener("mouseup", handleGlobalMouseUp);
-    };
-  }, [isMouseDown, isMouseOutside, closeOnMouseLeave, currColor, onChange]);
-
  const handleMouseLeave = (close: () => void) => {
-    closeRef.current = close;
-    setIsMouseOutside(true);
-
-    // Only close if mouse is not down and closeOnMouseLeave is true
-    if (closeOnMouseLeave && !isMouseDown) {
-      closeTimeoutRef.current = setTimeout(() => {
-        onChange(currColor);
-        setIsOpen(false);
-        close();
-      }, 200);
-    }
-  };
-
-  const handleMouseEnter = () => {
-    setIsMouseOutside(false);
-    if (closeTimeoutRef.current) {
-      clearTimeout(closeTimeoutRef.current);
-      closeTimeoutRef.current = null;
+    if (closeOnMouseLeave) {
+      onChange(currColor);
+      setIsOpen(false);
+      close();
    }
  };

@@ -139,7 +92,7 @@ export const PaintModeMenu = ({
    if (streamChanges) {
      onChange(color);
    }
-  };
+  }

  // Rerender the component when these signals change:
  useSignals();
@@ -170,16 +123,12 @@ export const PaintModeMenu = ({
            )}
          </PopoverButton>
          {(open || isOpen) && (
-            <div
-              onMouseLeave={() => handleMouseLeave(close)}
-              onMouseEnter={handleMouseEnter}
-              onMouseDown={() => setIsMouseDown(true)}
-            >
+            <div onMouseLeave={() => handleMouseLeave(close)}>
              <PopoverPanel
                anchor={anchor}
                className={twMerge(
                  paperWrapperStyles,
-                  "flex flex-col items-center gap-2 overflow-hidden p-4",
+                  "flex flex-col items-center gap-2 overflow-hidden",
                )}
                style={
                  {
@@ -189,23 +138,20 @@ export const PaintModeMenu = ({
                static
              >
                <HexEyedropPicker
+                  pickerClassName="overflow-hidden"
                  color={currColor}
                  onPickerChange={handleColorChange}
                  onDropperChange={handleColorChange}
                  Picker={HexAlphaColorPicker}
                />
                <Input
-                  className="mt-2"
-                  inputClassName="bg-ui-background/30"
                  style={{ width: "198px" }}
                  icon={faHashtag}
                  value={textInput}
                  onChange={handleTextInput}
                />
-                <div className="mt-2 flex w-full flex-col gap-2">
-                  <p className="w-full justify-start text-sm font-medium text-white">
-                    Brush Size:
-                  </p>
+                <div className="flex flex-col w-full mt-2 gap-2">
+                  <p className="text-white font-medium text-sm w-full justify-start">Brush Size:</p>
                  <Slider
                    min={BRUSH_MIN_SIZE}
                    max={BRUSH_MAX_SIZE}