Add log for base64 inpainting response

Painted nodes are persistent across edit modes.

.gitignore

@@ -66,3 +66,7 @@ scraps/
 # These should be built by the platform
 windows/sqlite/sqlite3.exp
 windows/sqlite/sqlite3.lib
+
+# ArtCraft Windows DLLs
+crates/desktop/src-tauri/*.dll
+windows_dlls/

.idea/dictionaries/bt.xml

@@ -2,6 +2,7 @@
   <dictionary name="bt">
     <words>
       <w>arpabet</w>
+      <w>artcraft</w>
       <w>backfill</w>
       <w>backfilling</w>
       <w>backfills</w>

Cargo.toml

@@ -23,6 +23,8 @@ members = [
 
   # Desktop apps
   "crates/desktop/src-tauri",
+  "crates/desktop/ml_models",
+  "crates/desktop/ml_weights_registry",
 
   # Shared libraries
   "crates/lib/actix_cors_configs",
@@ -98,6 +100,7 @@ members = [
   # Vendored code
   "crates/vendor/candle/candle-core",
   "crates/vendor/candle/candle-datasets",
+  "crates/vendor/candle/candle-examples",
   "crates/vendor/candle/candle-flash-attn",
   "crates/vendor/candle/candle-kernels",
   "crates/vendor/candle/candle-metal-kernels",

app_config.yaml

@@ -1,7 +1,8 @@
-image_width: 512
-image_height: 512
+image_width: 1024
+image_height: 1024
 #scheduler_steps: 15
 # NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
-scheduler_steps: 2
+scheduler_steps: 8
 scheduler_samples: 1
-cfg_scale: 1.1
+cfg_scale: 15.0
+seed: 42

crates/desktop/ml_models/Cargo.toml

@@ -0,0 +1,79 @@
+[package]
+name = "ml_models"
+edition = "2021"
+version = "0.0.1"
+authors = ["storyteller"]
+publish = false
+description = ""
+license = ""
+repository = ""
+
+[lib]
+name = "ml_models"
+path = "src/lib.rs"
+crate-type = ["staticlib", "cdylib", "rlib"]
+
+[package.metadata.cargo-udeps.ignore]
+normal = ["my-workspace-hack"]
+
+[dependencies]
+
+# Internal
+filesys = { path = "../../lib/files/filesys" }
+memory_caching = { path = "../../lib/caching/memory_caching" }
+ml_weights_registry = { path = "../ml_weights_registry" }
+primitives = { path = "../../lib/data/primitives" }
+
+# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
+cudarc = "0.13.7"
+
+# External
+anyhow = "1.0"
+base64 = "0.22.1"
+bytes = "1.10"
+candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
+candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
+candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
+const_format = "0.2.34"
+directories = "6.0.0"
+dirs = "6.0.0"
+futures = "0.3"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
+hf-hub = "0.3.2"
+image = "0.25.5"
+log = "0.4"
+ndarray = "0.15.6"
+nvidia-checker = "0.1.2"
+once_cell = "1.20.3"
+ort = { version = "1.16.2", features = ["cuda"] }
+rand = "0.8.5"
+# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
+reqwest = { version = "0.12.12", features = ["stream"], default-features = false }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+serde_yaml = "0.9"
+tauri = { version = "2.2.4", features = [] }
+tauri-plugin-log = "2"
+tempdir = "0.3.7"
+tempfile = "3.18.0"
+tokenizers = "0.15.2"
+tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
+tokio-util = { version = "0.7", features = ["codec"] }
+
+# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
+
+[features]
+default = ["cuda"]
+cuda = [
+    "candle-core/cuda",
+    "candle-nn/cuda",
+    "candle-transformers/cuda",
+    # NB(bt,2025-03-10): Building without Flash attention on Windows for now to bypass linker issues
+    "candle-transformers/flash-attn",
+    "candle-core/cudnn"
+]
+
+[dev-dependencies]
+libloading = "0.8.6"
+serial_test = "3.2.0"
+assertor = "0.0.3"

crates/desktop/ml_models/src/lib.rs

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

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

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

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

@@ -0,0 +1,19 @@
+//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs
+
+use std::path::Path;
+use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
+
+pub fn onnx_session(onnx_model_file: &Path) -> Result<ort::Session, anyhow::Error> {
+  let environment = Environment::default().into_arc();
+  let session = SessionBuilder::new(&environment)?
+    .with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level
+    .with_intra_threads(1)? // Configure the number of threads used for inference
+    .with_execution_providers([
+      // Configure execution providers (e.g., CUDA, CoreML, CPU)
+      ExecutionProvider::CUDA(Default::default()),
+      ExecutionProvider::CoreML(Default::default()),
+      ExecutionProvider::CPU(Default::default()),
+    ])?
+    .with_model_from_file(onnx_model_file)?;
+  Ok(session)
+}

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

@@ -0,0 +1,61 @@
+use crate::ml::background_removal::onnx_session::onnx_session;
+use image::{imageops, DynamicImage};
+use ndarray::{Array, CowArray};
+use ort::Value;
+use std::path::Path;
+
+pub async fn remove_image_background<P: AsRef<Path>>(model_path: P, image: DynamicImage) -> anyhow::Result<DynamicImage> {
+  // NB: Real time needed cuDNN 9, but this uses cuDNN 8:
+  // sudo apt-get install libcudnn8-dev libcudnn8
+  // update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken
+  // update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode
+  let session = onnx_session(model_path.as_ref())?;
+  let image = process_dynamic_image(&session, image)?;
+  Ok(image)
+}
+
+fn process_dynamic_image(
+  session: &ort::Session,
+  dynamic_img: DynamicImage,
+) -> Result<DynamicImage, anyhow::Error> {
+  let input_shape = session.inputs[0]
+    .dimensions()
+    .map(|dim| dim.unwrap())
+    .collect::<Vec<usize>>();
+  let input_img = dynamic_img.into_rgba8();
+  let scaling_factor = f32::min(
+    1., // Avoid upscaling
+    f32::min(
+      input_shape[3] as f32 / input_img.width() as f32, // Width ratio
+      input_shape[2] as f32 / input_img.height() as f32, // Height ratio
+    ),
+  );
+  let mut resized_img = imageops::resize(
+    &input_img,
+    input_shape[3] as u32,
+    input_shape[2] as u32,
+    imageops::FilterType::Triangle,
+  );
+  let input_tensor = CowArray::from(
+    Array::from_shape_fn(input_shape, |indices| {
+      let mean = 128.;
+      let std = 256.;
+
+      (resized_img[(indices[3] as u32, indices[2] as u32)][indices[1]] as f32 - mean) / std
+    })
+      .into_dyn(),
+  );
+  let inputs = vec![Value::from_array(session.allocator(), &input_tensor)?];
+  let outputs = session.run(inputs)?;
+  let output_tensor = outputs[0].try_extract::<f32>()?;
+  for (indices, alpha) in output_tensor.view().indexed_iter() {
+    resized_img[(indices[3] as u32, indices[2] as u32)][3] = (alpha * 255.) as u8;
+  }
+  let output_img = imageops::resize(
+    &resized_img,
+    (input_img.width() as f32 * scaling_factor) as u32,
+    (input_img.height() as f32 * scaling_factor) as u32,
+    imageops::FilterType::Triangle,
+  );
+  Ok(DynamicImage::ImageRgba8(output_img))
+}

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

@@ -1,6 +1,7 @@
+pub mod background_removal;
 pub mod image;
 pub mod model_cache;
+pub mod model_config;
 pub mod model_file;
-pub mod models;
 pub mod prompt_cache;
 pub mod stable_diffusion;

crates/desktop/ml_models/src/ml/model_cache.rs

@@ -0,0 +1,60 @@
+use anyhow::anyhow;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
+use candle_transformers::models::stable_diffusion::vae::AutoEncoderKL;
+use std::sync::{Arc, RwLock};
+// TODO: This data structure is gross. Generalize the locking and loading semantics for reuse.
+
+// Simple registry for now. We can build complex machinery that aids in 
+// VRAM utilization, disk space saving, intelligent scheduling, etc. in the future.
+pub struct ModelCache {
+  vae: RwLock<Option<Arc<AutoEncoderKL>>>,
+  unet2: RwLock<Option<Arc<UNet2DConditionModel>>>,
+}
+
+impl ModelCache {
+  pub fn new() -> Self {
+    Self { 
+      vae: RwLock::new(None),
+      unet2: RwLock::new(None),
+    }
+  }
+  
+  pub fn set_vae(&self, vae: Arc<AutoEncoderKL>) -> anyhow::Result<()> {
+    match self.vae.write() {
+      Err(err) => Err(anyhow!("{}", err)),
+      Ok(mut lock) => {
+        *lock = Some(vae);
+        Ok(())
+      }
+    }
+  }
+
+  pub fn get_vae(&self) -> anyhow::Result<Option<Arc<AutoEncoderKL>>> {
+    match self.vae.read() {
+      Err(err) => Err(anyhow!("{}", err)),
+      Ok(ref lock) => {
+        Ok(lock.as_ref().map(|v|v.clone()))
+      }
+    }
+  }
+
+  pub fn set_unet(&self, vae: Arc<UNet2DConditionModel>) -> anyhow::Result<()> {
+    match self.unet2.write() {
+      Err(err) => Err(anyhow!("{}", err)),
+      Ok(mut lock) => {
+        *lock = Some(vae);
+        Ok(())
+      }
+    }
+  }
+
+  pub fn get_unet(&self) -> anyhow::Result<Option<Arc<UNet2DConditionModel>>> {
+    match self.unet2.read() {
+      Err(err) => Err(anyhow!("{}", err)),
+      Ok(ref lock) => {
+        Ok(lock.as_ref().map(|v|v.clone()))
+      }
+    }
+  }
+  
+}

crates/desktop/ml_models/src/ml/model_config.rs

@@ -0,0 +1,61 @@
+use crate::ml::model_file::StableDiffusionVersion;
+use candle_core::{DType, Device};
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use hf_hub::api::sync::Api;
+use log::info;
+
+const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
+const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
+const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
+const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;
+
+pub struct ModelConfig {
+
+  /// We only support one device type currently, and we can't multiplex
+  /// over several devices yet.
+  pub device: Device,
+
+  /// The core numeric type to use for models.
+  pub dtype: DType,
+
+  pub sd_version: StableDiffusionVersion,
+  pub sd_config: StableDiffusionConfig,
+
+  /// Probably shouldn't be here.
+  pub hf_api: Api,
+}
+
+impl ModelConfig {
+  pub fn init() -> anyhow::Result<Self> {
+    let device = Device::new_cuda(0)
+      .unwrap_or_else(|e| {
+          info!("CUDA not available ({}), falling back to CPU", e);
+          Device::Cpu
+        });
+
+    let sd_config = match SD_VERSION {
+      StableDiffusionVersion::Turbo => {
+        StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
+        StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
+        StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
+        StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+    };
+
+    let hf_api = Api::new()?;
+
+    Ok(Self {
+      device,
+      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
+      sd_version: SD_VERSION,
+      sd_config,
+      hf_api,
+    })
+  }
+}

crates/desktop/ml_models/src/ml/model_file.rs

@@ -0,0 +1,180 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+use candle_transformers::models::stable_diffusion;
+use anyhow::{Error as E, Result};
+use hf_hub::api::sync::Api;
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StableDiffusionVersion {
+    V1_5,
+    V1_5Inpaint,
+    V2_1,
+    V2Inpaint,
+    Xl,
+    XlInpaint,
+    Turbo,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ModelFile {
+    Tokenizer,
+    Tokenizer2,
+    Clip,
+    Clip2,
+    Unet,
+    UnetLcm,
+    Vae,
+}
+
+impl StableDiffusionVersion {
+    pub fn repo(&self) -> &'static str {
+        match self {
+            Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+            Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0",
+            Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting",
+            Self::V2_1 => "stabilityai/stable-diffusion-2-1",
+            Self::V1_5 => "Lykon/dreamshaper-7",
+            Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
+            Self::Turbo => "Lykon/dreamshaper-7",
+        }
+    }
+
+    fn unet_file(&self, use_f16: bool) -> &'static str {
+        match self {
+            Self::V1_5
+            | Self::V1_5Inpaint
+            | Self::V2_1
+            | Self::V2Inpaint
+            | Self::Xl
+            | Self::XlInpaint
+            | Self::Turbo => {
+                if use_f16 {
+                    "unet/diffusion_pytorch_model.fp16.safetensors"
+                } else {
+                    "unet/diffusion_pytorch_model.safetensors"
+                }
+            }
+        }
+    }
+
+    fn vae_file(&self, use_f16: bool) -> &'static str {
+        match self {
+            Self::V1_5
+            | Self::V1_5Inpaint
+            | Self::V2_1
+            | Self::V2Inpaint
+            | Self::Xl
+            | Self::XlInpaint
+            | Self::Turbo => {
+                if use_f16 {
+                    "vae/diffusion_pytorch_model.fp16.safetensors"
+                } else {
+                    "vae/diffusion_pytorch_model.safetensors"
+                }
+            }
+        }
+    }
+
+    fn clip_file(&self, use_f16: bool) -> &'static str {
+        match self {
+            Self::V1_5
+            | Self::V1_5Inpaint
+            | Self::V2_1
+            | Self::V2Inpaint
+            | Self::Xl
+            | Self::XlInpaint
+            | Self::Turbo => {
+                if use_f16 {
+                    "text_encoder/model.fp16.safetensors"
+                } else {
+                    "text_encoder/model.safetensors"
+                }
+            }
+        }
+    }
+
+    fn clip2_file(&self, use_f16: bool) -> &'static str {
+        match self {
+            Self::V1_5
+            | Self::V1_5Inpaint
+            | Self::V2_1
+            | Self::V2Inpaint
+            | Self::Xl
+            | Self::XlInpaint
+            | Self::Turbo => {
+                if use_f16 {
+                    "text_encoder_2/model.fp16.safetensors"
+                } else {
+                    "text_encoder_2/model.safetensors"
+                }
+            }
+        }
+    }
+}
+
+impl ModelFile {
+    pub fn get(
+        &self,
+        filename: Option<String>,
+        version: StableDiffusionVersion,
+        use_f16: bool,
+    ) -> Result<std::path::PathBuf> {
+        match filename {
+            Some(filename) => Ok(std::path::PathBuf::from(filename)),
+            None => {
+                let (repo, path) = match self {
+                    Self::Tokenizer => {
+                        let tokenizer_repo = match version {
+                            StableDiffusionVersion::V1_5
+                            | StableDiffusionVersion::V2_1
+                            | StableDiffusionVersion::V1_5Inpaint
+                            | StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32",
+                            StableDiffusionVersion::Xl
+                            | StableDiffusionVersion::XlInpaint
+                            | StableDiffusionVersion::Turbo => {
+                                // This seems similar to the patch32 version except some very small
+                                // difference in the split regex.
+                                "openai/clip-vit-large-patch14"
+                            }
+                        };
+                        (tokenizer_repo, "tokenizer.json")
+                    }
+                    Self::Tokenizer2 => {
+                        ("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json")
+                    }
+                    Self::Clip => (version.repo(), version.clip_file(use_f16)),
+                    Self::Clip2 => (version.repo(), version.clip2_file(use_f16)),
+                    Self::UnetLcm => (
+                        "SimianLuo/LCM_Dreamshaper_v7",
+                        if use_f16 {
+                            "unet/diffusion_pytorch_model.fp16.safetensors"
+                        } else {
+                            "unet/diffusion_pytorch_model.safetensors"
+                        },
+                    ),
+                    Self::Unet => (version.repo(), version.unet_file(use_f16)),
+                    Self::Vae => {
+                        // Override for SDXL when using f16 weights.
+                        // See https://github.com/huggingface/candle/issues/1060
+                        if matches!(
+                            version,
+                            StableDiffusionVersion::Xl | StableDiffusionVersion::Turbo,
+                        ) && use_f16
+                        {
+                            (
+                                "madebyollin/sdxl-vae-fp16-fix",
+                                "diffusion_pytorch_model.safetensors",
+                            )
+                        } else {
+                            (version.repo(), version.vae_file(use_f16))
+                        }
+                    }
+                };
+                let filename = Api::new()?.model(repo.to_string()).get(path)?;
+                Ok(filename)
+            }
+        }
+    }
+}

crates/desktop/ml_models/src/ml/stable_diffusion/infer_clip_text_embeddings.rs

@@ -0,0 +1,158 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+use candle_transformers::models::stable_diffusion;
+use std::path::{Path, PathBuf};
+
+use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use ml_weights_registry::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2};
+use anyhow::Error as E;
+use candle_core::{DType, Device, Module, Tensor, D};
+use log::info;
+use tokenizers::Tokenizer;
+
+
+pub struct ClipArgs<'a, P: AsRef<Path>,  Q: AsRef<Path>> {
+  pub prompt: &'a str, 
+  pub uncond_prompt: &'a str, 
+  pub tokenizer: Option<String>, 
+  pub clip_weights: Option<String>, 
+  pub clip2_weights: Option<String>, 
+  pub sd_version: StableDiffusionVersion, 
+  pub sd_config: &'a stable_diffusion::StableDiffusionConfig, 
+  pub use_f16: bool, 
+  pub device: &'a Device, 
+  pub dtype: DType, 
+  pub use_guide_scale: bool, 
+  pub clip_json_path: P,
+  pub clip_weights_path: Q,
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn infer_clip_text_embeddings<P: AsRef<Path>, Q: AsRef<Path>>(args: ClipArgs<'_, P, Q>) -> anyhow::Result<Tensor> {
+  let ClipArgs {
+    prompt, 
+    uncond_prompt, 
+    tokenizer, 
+    clip_weights, 
+    clip2_weights, 
+    sd_version, 
+    sd_config, 
+    use_f16, 
+    device, 
+    dtype, 
+    use_guide_scale, 
+    clip_json_path,
+    clip_weights_path,
+  } = args;
+
+  info!("Preparing text embeddings... for {:?}", sd_version);
+
+  let which = match sd_version {
+    StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false],
+    _ => vec![true],
+  };
+
+  let text_embeddings = which
+    .iter()
+    .map(|first| {
+      do_infer_clip_text_embeddings(
+        &prompt,
+        &uncond_prompt,
+        None, // tokenizer
+        clip_weights.clone(),
+        None, // clip2_weights
+        sd_version,
+        &sd_config,
+        false, // use_f16
+        &device,
+        dtype,
+        use_guide_scale,
+        *first,
+        clip_json_path.as_ref(),
+        clip_weights_path.as_ref(),
+      )
+    })
+    .collect::<anyhow::Result<Vec<_>>>()?;
+
+  let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
+
+  Ok(text_embeddings)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn do_infer_clip_text_embeddings(
+  prompt: &str, 
+  uncond_prompt: &str, 
+  tokenizer: Option<String>, 
+  clip_weights: Option<String>, 
+  clip2_weights: Option<String>, 
+  sd_version: StableDiffusionVersion, 
+  sd_config: &stable_diffusion::StableDiffusionConfig, 
+  use_f16: bool, 
+  device: &Device, 
+  dtype: DType, 
+  use_guide_scale: bool, 
+  first: bool, 
+  clip_json_path: &Path,
+  clip_weights_path: &Path,
+) -> anyhow::Result<Tensor> {
+  info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first));
+
+  let tokenizer = clip_json_path;
+
+  info!("Loading Clip Tokenizer path: {:?}", tokenizer);
+
+  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(),
+  };
+
+  println!("Running with prompt \"{prompt}\".");
+  let mut tokens = tokenizer.encode(prompt, true).map_err(E::msg)?.get_ids().to_vec();
+  if tokens.len() > sd_config.clip.max_position_embeddings {
+    anyhow::bail!("the prompt is too long, {} > max-tokens ({})", tokens.len(), sd_config.clip.max_position_embeddings)
+  }
+  while tokens.len() < sd_config.clip.max_position_embeddings {
+    tokens.push(pad_id)
+  }
+  let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
+
+  println!("Building the Clip transformer.");
+  
+  // TODO(bt,2025-03-23): These used to be two different clip models. We should only need one.
+  let clip_weights = if first {
+    println!(">>>>> CLIP 1");
+    clip_weights_path
+  } else {
+    println!(">>>>> CLIP 2");
+    clip_weights_path
+  };
+  
+  let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
+  let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
+
+  let text_embeddings = text_model.forward(&tokens)?;
+
+  let text_embeddings = if use_guide_scale {
+    let mut uncond_tokens = tokenizer.encode(uncond_prompt, true).map_err(E::msg)?.get_ids().to_vec();
+    if uncond_tokens.len() > sd_config.clip.max_position_embeddings {
+      anyhow::bail!("the negative prompt is too long, {} > max-tokens ({})", uncond_tokens.len(), sd_config.clip.max_position_embeddings)
+    }
+    while uncond_tokens.len() < sd_config.clip.max_position_embeddings {
+      uncond_tokens.push(pad_id)
+    }
+
+    let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?;
+    let uncond_embeddings = text_model.forward(&uncond_tokens)?;
+
+    Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)?
+  } else {
+    text_embeddings.to_dtype(dtype)?
+  };
+  Ok(text_embeddings)
+}

crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs

@@ -0,0 +1,281 @@
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use anyhow::{anyhow, Error as E, Result};
+use candle_core::{DType, Device, IndexOp, Tensor, D};
+use candle_nn::VarBuilder;
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use candle_transformers::models::stable_diffusion::lcm::{LCMScheduler, LCMSchedulerConfig};
+use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
+use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
+use crate::ml::model_cache::ModelCache;
+use crate::ml::model_config::ModelConfig;
+use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use crate::ml::prompt_cache::PromptCache;
+use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
+use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
+use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
+use image::DynamicImage;
+use log::info;
+use ml_weights_registry::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
+use rand::Rng;
+
+const DEFAULT_STRENGTH : f64 = 75.0;
+
+pub struct Args<'a, P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>> {
+  pub image: &'a DynamicImage,
+  pub prompt: String,
+  pub uncond_prompt: String,
+  pub cfg_scale: Option<f64>,
+  pub img2img_strength: Option<f64>,
+  pub model_cache: &'a ModelCache,
+  pub prompt_cache: &'a PromptCache,
+  pub use_flash_attn: bool,
+  pub model_config: &'a ModelConfig,
+  pub maybe_seed: Option<u64>,
+  pub scheduler_steps: usize,
+  pub vae_path: P1,
+  pub unet_path: P2,
+  pub clip_json_path: P3,
+  pub clip_weights_path: P4,
+}
+
+pub fn lcm_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage> 
+  where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
+{
+  let Args { 
+    prompt, 
+    uncond_prompt, 
+    cfg_scale,
+    img2img_strength, 
+    model_cache, 
+    prompt_cache, 
+    image, 
+    use_flash_attn ,
+    model_config: ModelConfig {
+      device,
+      dtype,
+      sd_version,
+      sd_config,
+      hf_api,
+    },
+    maybe_seed,
+    scheduler_steps,
+    vae_path,
+    unet_path,
+    clip_json_path,
+    clip_weights_path,
+  } = args;
+  
+  let dtype = *dtype;
+  let sd_version = *sd_version;
+
+  let img2img_strength = img2img_strength
+    .map(|s| remap_lcm_strength_range(s))
+    .unwrap_or(DEFAULT_STRENGTH);
+  
+  let img2img_strength = img2img_strength / 100.0;
+
+  let use_f16 = true;
+
+  // Use LCM Scheduler instead of Euler Ancestral for better speed and quality
+  let mut scheduler = LCMScheduler::new(scheduler_steps, img2img_strength, LCMSchedulerConfig::default())?;
+
+  let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
+
+  let seed = 42;
+  
+  device.set_seed(seed)?;
+
+  let guidance_scale = match cfg_scale {
+    Some(guidance_scale) => guidance_scale,
+    None => match 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;
+  let use_guide_scale = false;
+
+  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(ClipArgs {
+      prompt: & prompt,
+      uncond_prompt: &uncond_prompt,
+      tokenizer: None, // tokenizer
+      clip_weights: None,
+      clip2_weights: None, // clip2_weights
+      sd_version,
+      sd_config: &sd_config,
+      use_f16, // use_f16
+      device,
+      dtype,
+      use_guide_scale,
+      clip_json_path,
+      clip_weights_path,
+    })?;
+
+    prompt_cache.store_copy(&prompt, &tensor)?;
+    tensor
+  };
+
+  info!("Text embeddings shape: {:?}", text_embeddings.shape());
+
+  info!("Loading input image into tensor...");
+
+  let input_image = dynamic_image_to_tensor(image, &device, dtype)?;
+
+  info!("Reference image shape: {:?}", input_image.shape());
+
+  let vae_scale = get_vae_scale(sd_version);
+
+  let maybe_vae = model_cache.get_vae()?;
+
+  let vae = match maybe_vae {
+    Some(vae) => vae,
+    None => {
+      info!("Building VAE model from file {:?}...", vae_path.as_ref());
+      
+      let vae = sd_config.build_vae(vae_path, &device, dtype)?;
+
+      let vae = Arc::new(vae);
+
+      model_cache.set_vae(vae.clone())?;
+
+      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 in_channels = match sd_version {
+        StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
+        _ => 4,
+      };
+
+      // For LCM models, ensure that:
+      // 1. We're using the correct prediction type (epsilon)
+      // 2. The cross-attention dimension is correct (768 for Dreamshaper)
+      // 3. The in_channels matches the expected input (4 for standard, 9 for inpainting)
+      let unet_config = UNet2DConditionModelConfig {
+        // LCM specific config values
+        cross_attention_dim: 768,      // Dreamshaper v7 uses 768
+        time_cond_proj_dim: Some(256), // LCM models use 256-dimensional guidance embeddings
+        // Rest of configuration...
+        center_input_sample: false,
+        flip_sin_to_cos: true,
+        freq_shift: 0.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.0,
+        norm_num_groups: 32,
+        norm_eps: 1e-5,
+        sliced_attention_size: None,
+        use_linear_projection: false,
+        time_embed_dim: Some(1280),
+        // time_embed_dim: Some(256),
+      };
+
+      // Load the model directly from the safetensors file
+      let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_path], dtype, &device)? };
+
+      let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
+
+      let unet = Arc::new(unet);
+
+      model_cache.set_unet(unet.clone())?;
+
+      unet
+    },
+  };
+
+  let init_latent_dist: DiagonalGaussianDistribution = vae.encode(&input_image)?;
+
+  // Create LCM guidance scale embeddings
+  let embedding_dim = 1280;
+  let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?;
+
+  // Generate latents from input image using the LCM approach
+  info!("Generating latents from input image...");
+  let init_latents = (init_latent_dist.sample()? * vae_scale)?;
+  info!("Initial latents shape: {:?}", init_latents.shape());
+
+  // Calculate timesteps for LCM with proper img2img handling
+  let timesteps = LCMScheduler::get_timesteps_for_steps(scheduler_steps, img2img_strength);
+  let t_start = timesteps[0]; // First denoising step
+
+  // Add noise at the appropriate timestep for img2img
+  info!("Adding noise to latents for t_start={}", t_start);
+  let noise = init_latents.randn_like(0f64, 1f64)?;
+  let latents = scheduler.add_noise(&init_latents, noise, t_start)?;
+
+  let mut latents = latents;
+  let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
+
+  for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+    info!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
+
+    let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() };
+
+    let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);
+
+    // Use the guidance scale embedding in the UNet inference with the proper method name
+    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, Some(&guidance_scale_embedding)) {
+      Ok(pred) => pred,
+      Err(e) => {
+        info!("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, scheduler_steps)?;
+  }
+
+  info!("Diffusion process completed, decoding image...");
+  let image = vae.decode(&(latents / vae_scale)?)?;
+
+  info!("VAE decode completed");
+
+  info!("Scaling image back");
+  let image = ((image / 2.)? + 0.5)?;
+
+  info!("Normalized image values");
+  let image = (image.clamp(0f32, 1.)? * 255.)?;
+
+  info!("Convert to int8");
+  let image = image.to_dtype(DType::U8)?;
+
+  info!("Converted to 8-bit format");
+
+  let image = tensor_to_image_buffer(&image.i(0)?)?;
+
+  Ok(image)
+}

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

@@ -0,0 +1,5 @@
+pub mod get_vae_scale;
+pub mod infer_clip_text_embeddings;
+pub mod stable_diffusion_pipeline;
+pub mod lcm_pipeline;
+mod remap_lcm_strength_range;

crates/desktop/ml_models/src/ml/stable_diffusion/remap_lcm_strength_range.rs

@@ -0,0 +1,54 @@
+use primitives::numerics::remap_range_f64::{remap_range_f64, Args};
+
+/// Our LCM implementation uses strength [13, 100]. Strength below 13 will take 8 steps, 
+/// so we need to limit ourselves to this range from an input range of [0, 100].
+pub fn remap_lcm_strength_range(value: f64) -> f64 {
+  remap_range_f64(Args {
+    old_min: 0.0,
+    old_max: 100.0,
+    new_min: 13.0,
+    new_max: 100.0,
+    value,
+  })
+}
+
+#[cfg(test)]
+mod test {
+  use crate::ml::stable_diffusion::remap_lcm_strength_range::remap_lcm_strength_range;
+  use assertor::{assert_that, FloatAssertion};
+
+  #[test]
+  fn old_max() {
+    assert_eq!(remap_lcm_strength_range(100.0), 100.0);
+  }
+  
+  #[test]
+  fn old_min() {
+    assert_eq!(remap_lcm_strength_range(0.0), 13.0);
+  }
+
+  #[test]
+  fn under_bounds() {
+    assert_eq!(remap_lcm_strength_range(-100.0), 13.0);
+  }
+  
+  #[test]
+  fn over_bounds() {
+    assert_eq!(remap_lcm_strength_range(200.0), 100.0);
+  }
+
+  #[test]
+  fn other_values() {
+    assert_that!(remap_lcm_strength_range(0.0)).with_abs_tol(0.001).is_approx_equal_to(13.000);
+    assert_that!(remap_lcm_strength_range(5.0)).with_abs_tol(0.001).is_approx_equal_to(17.350);
+    assert_that!(remap_lcm_strength_range(11.0)).with_abs_tol(0.001).is_approx_equal_to(22.570);
+    assert_that!(remap_lcm_strength_range(13.0)).with_abs_tol(0.001).is_approx_equal_to(24.310);
+    assert_that!(remap_lcm_strength_range(14.0)).with_abs_tol(0.001).is_approx_equal_to(25.180);
+    assert_that!(remap_lcm_strength_range(20.0)).with_abs_tol(0.001).is_approx_equal_to(30.400);
+    assert_that!(remap_lcm_strength_range(25.0)).with_abs_tol(0.001).is_approx_equal_to(34.750);
+    assert_that!(remap_lcm_strength_range(50.0)).with_abs_tol(0.001).is_approx_equal_to(56.500);
+    assert_that!(remap_lcm_strength_range(75.0)).with_abs_tol(0.001).is_approx_equal_to(78.250);
+    assert_that!(remap_lcm_strength_range(99.0)).with_abs_tol(0.001).is_approx_equal_to(99.130);
+    assert_that!(remap_lcm_strength_range(100.0)).with_abs_tol(0.001).is_approx_equal_to(100.000);
+  }
+}

crates/desktop/ml_models/src/ml/stable_diffusion/stable_diffusion_pipeline.rs

@@ -1,4 +1,5 @@
-use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+use std::path::Path;
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 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};
@@ -6,49 +7,80 @@ use crate::ml::model_cache::ModelCache;
 use crate::ml::model_file::StableDiffusionVersion;
 use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
-use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
-use crate::state::app_config::AppConfig;
+use crate::ml::stable_diffusion::infer_clip_text_embeddings::{infer_clip_text_embeddings, ClipArgs};
 use anyhow::{anyhow, Error as E, Result};
-use candle_core::{DType, IndexOp, Tensor, D};
-use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
+use candle_core::{DType, Device, IndexOp, Tensor, D};
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
 use image::DynamicImage;
 use log::info;
 use rand::Rng;
 
-pub struct Args<'a> {
+pub struct Args<'a, P: AsRef<Path>, Q: AsRef<Path>, R: AsRef<Path>, S: AsRef<Path>> {
     pub image: &'a DynamicImage,
     pub prompt: String,
     pub uncond_prompt: String,
     pub cfg_scale: Option<f64>,
     pub i2i_strength: Option<u8>,
-    pub configs: &'a AppConfig,
     pub model_cache: &'a ModelCache,
     pub prompt_cache: &'a PromptCache,
+    pub sd_version: StableDiffusionVersion,
+    pub sd_config: StableDiffusionConfig,
+    pub device: &'a Device,
+    pub dtype: DType,
+    pub maybe_seed: Option<u64>,
+    pub scheduler_steps: usize,
+    pub sdxl_turbo_vae_path: P,
+    pub sdxl_turbo_unet_path: Q,
+    pub clip_json_path: R,
+    pub clip_weights_path: S,
 }
 
-pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
+pub fn stable_diffusion_pipeline<P1, P2, P3, P4>(args: Args<'_, P1, P2, P3, P4>) -> Result<RgbImage> 
+  where P1: AsRef<Path>, P2: AsRef<Path>, P3: AsRef<Path>, P4: AsRef<Path>
+{
+    let Args { 
+        prompt, 
+        uncond_prompt, 
+        cfg_scale, 
+        i2i_strength, 
+        model_cache, 
+        prompt_cache, 
+        image,
+        sd_version,
+        sd_config,
+        device,
+        dtype,
+        maybe_seed,
+        scheduler_steps,
+        sdxl_turbo_vae_path,
+        sdxl_turbo_unet_path,
+        clip_json_path,
+        clip_weights_path,
+    } = args;
+    
     println!("Starting image generation with the following configuration:");
-    println!("  Model: {:?}", args.configs.sd_version);
-    println!("  Prompt: {}", args.prompt);
-    println!("  Steps: {}", args.configs.scheduler_steps);
-    println!("  Device: {:?}", args.configs.device);
+    println!("  Model: {:?}", sd_version);
+    println!("  Prompt: {}", prompt);
+    println!("  Steps: {}", scheduler_steps);
+    println!("  Device: {:?}", device);
 
-    println!("Model dimensions: {}x{}", args.configs.sd_config.width, args.configs.sd_config.height);
+    println!("Model dimensions: {}x{}", sd_config.width, sd_config.height);
 
     // TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler.
     //  This is a target for performance improvement
     //  See: https://github.com/huggingface/candle/issues/1331
-    let mut scheduler = args.configs.sd_config.build_scheduler(
-        args.configs.scheduler_steps)?;
+    let mut scheduler = sd_config.build_scheduler(scheduler_steps)?;
 
-    let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
-    args.configs.device.set_seed(seed)?;
+    let seed = maybe_seed.unwrap_or_else(|| rand::thread_rng().gen());
+    
+    device.set_seed(seed)?;
 
-    println!("Using seed: {}", seed);
+    info!("Using seed: {}", seed);
 
-    let guidance_scale = match args.cfg_scale {
+    let guidance_scale = match cfg_scale {
         Some(guidance_scale) => guidance_scale,
-        None => match args.configs.sd_version {
+        None => match sd_version {
             StableDiffusionVersion::V1_5
             | StableDiffusionVersion::V2_1
             | StableDiffusionVersion::Xl => 7.5,
@@ -59,29 +91,32 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
     let use_guide_scale = guidance_scale > 1.0;
 
-    println!("Using guide scale: {}", use_guide_scale);
+    info!("Using guide scale: {}", use_guide_scale);
 
     info!("Checking if prompt is cached");
-    let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?;
+    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(
-            &args.prompt,
-            &args.uncond_prompt,
-            None, // tokenizer
-            None, // clip_weights
-            None, // clip2_weights
-            args.configs.sd_version,
-            &args.configs.sd_config,
-            false, // use_f16
-            &args.configs.device,
-            args.configs.dtype,
+        let tensor = infer_clip_text_embeddings(ClipArgs {
+            prompt: &prompt,
+            uncond_prompt: &uncond_prompt,
+            tokenizer: None, // tokenizer
+            clip_weights: None, // clip_weights
+            clip2_weights: None, // clip2_weights
+            sd_version: sd_version,
+            sd_config: &sd_config,
+            use_f16: false, // use_f16
+            device: device,
+            dtype: dtype,
             use_guide_scale,
-        )?;
-        args.prompt_cache.store_copy(&args.prompt, &tensor)?;
+            clip_json_path: clip_json_path.as_ref(),
+            clip_weights_path: clip_weights_path.as_ref(),
+        })?;
+        prompt_cache.store_copy(&prompt, &tensor)?;
         tensor
     };
 
@@ -89,37 +124,73 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
     println!("Loading input image into tensor...");
 
-    let input_image = dynamic_image_to_tensor(
-        args.image,
-        &args.configs.device,
-        args.configs.dtype)?;
+    let input_image = dynamic_image_to_tensor(image, device, dtype)?;
 
     println!("Reference image shape: {:?}", input_image.shape());
 
-    let vae_scale = get_vae_scale(args.configs.sd_version);
+    let vae_scale = get_vae_scale(sd_version);
 
-    let init_latent_dist : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?;
+    let maybe_vae = model_cache.get_vae()?;
+    
+    let vae = match maybe_vae {
+        Some(vae) => vae,
+        None => {
+            info!("No vae found in cache; loading...");
+            
+            let vae_file = sdxl_turbo_vae_path.as_ref();
+            
+            let vae = sd_config.build_vae(vae_file, &device, dtype)?;
+
+            let vae = Arc::new(vae);
+            
+            model_cache.set_vae(vae.clone())?;
+
+            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 unet_weights_file = sdxl_turbo_unet_path.as_ref();
+
+            let unet = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)
+              .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
+
+            let unet = Arc::new(unet);
+
+            model_cache.set_unet(unet.clone())?;
+
+            unet
+        }
+    };
+
+    let init_latent_dist : DiagonalGaussianDistribution = vae.encode(&input_image)?;
 
     // TODO(bt,2025-02-18): This takes a little bit to generate the sample.
     //  This is a target for performance improvement
     println!("Generating latents from input image...");
     let latents = (init_latent_dist.sample()? * vae_scale)?;
     
-    //.to_device(&args.configs.device)?;
+    //.to_device(&configs.device)?;
     
     println!("Initial latents shape: {:?}", latents.shape());
 
     println!("Calculating start step for diffusion process...");
 
-    let img2img_strength = match args.i2i_strength {
+    let img2img_strength = match i2i_strength {
         None => 0.75f64,
         Some(strength) => (strength as f64) / 100.0f64,
     };
 
     let t_start = {
-        let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize;
+        let start = scheduler_steps - (scheduler_steps as f64 * img2img_strength) as usize;
 
-        println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength);
+        println!("Starting from step {} of {} (strength: {})", start, scheduler_steps, img2img_strength);
         start
     };
 
@@ -158,7 +229,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
         let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
 
-        let mut noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) {
+        let mut noise_pred = match unet.forward(&latent_model_input, timestep as f64, &text_embeddings) {
             Ok(pred) => pred,
             Err(e) => {
                 println!("UNet inference failed with error: {}", e);
@@ -176,7 +247,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
     }
 
     println!("Diffusion process completed, decoding image...");
-    let image = args.model_cache.vae_decode(&(latents / vae_scale)?)?;
+    let image = vae.decode(&(latents / vae_scale)?)?;
     
     println!("VAE decode completed");
 

crates/desktop/ml_weights_registry/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "ml_weights_registry"
+edition = "2021"
+version = "0.0.1"
+authors = ["storyteller"]
+publish = false
+description = ""
+license = ""
+repository = ""
+
+[lib]
+name = "ml_weights_registry"
+path = "src/lib.rs"
+
+# External
+[dependencies]
+anyhow = "1.0"
+const_format = "0.2.34"

crates/desktop/ml_weights_registry/src/lib.rs

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

crates/desktop/ml_weights_registry/src/weights_registry/mod.rs

@@ -0,0 +1,3 @@
+mod weight_descriptor_builder;
+pub mod weight_descriptor;
+pub mod weights;

crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs

@@ -0,0 +1,45 @@
+
+// TODO(bt,2025-03-13): When it's time for dynamic, non-const allocation, refer to
+//  http::header::name::HeaderName: https://github.com/hyperium/http/blob/master/src/header/name.rs
+//  as a pattern that will support both compile-time constants and dynamically allocated String-based variants.
+
+#[derive(Clone, Copy)]
+pub enum WeightFunction {
+  ImageSegmentation,
+  TextEncoder,
+  TextTokenizer,
+  Unet,
+  Vae,
+}
+
+pub struct WeightDescriptor {
+  /// Name of the weight, eg. "Lykon Dreamshaper 7"
+  pub name: &'static str,
+
+  /// Name of the file on the filesystem, eg. lykon_dreamshaper_7_text_encoder.fp16.safetensors
+  pub filename: &'static str,
+
+  /// URL of the file in the R2 bucket,
+  /// eg. https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors
+  pub r2_download_url: &'static str,
+
+  pub function: WeightFunction,
+}
+
+impl WeightFunction {
+  pub fn as_descriptive_str(&self) -> &'static str {
+    match self {
+      WeightFunction::ImageSegmentation => "image segmentation",
+      WeightFunction::TextEncoder => "text encoder",
+      WeightFunction::TextTokenizer => "text tokenizer",
+      WeightFunction::Unet => "UNET",
+      WeightFunction::Vae => "VAE",
+    }
+  }
+}
+
+impl WeightDescriptor {
+  pub fn to_descriptive_name(&self) -> String {
+    format!("{} ({})", self.name, self.function.as_descriptive_str())
+  }
+}

crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor_builder.rs

@@ -0,0 +1,17 @@
+pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/";
+
+// TODO(bt,2025-03-13): Try to use keyword args from the call site.
+// NB: We're using a macro because we can't use a builder pattern and then concat string literals
+//  and variables (even if they're known to be &'static) and keep allocations static.
+macro_rules! weight {
+    ($name: literal, $file_name:literal, $weight_function:expr) => {
+      $crate::weights_registry::weight_descriptor::WeightDescriptor {
+        name: $name,
+        filename: $file_name,
+        r2_download_url: const_format::concatcp!(crate::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
+        function: $weight_function,
+      }
+    };
+}
+
+pub(crate) use weight;

crates/desktop/ml_weights_registry/src/weights_registry/weights.rs

@@ -0,0 +1,65 @@
+use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
+use crate::weights_registry::weight_descriptor_builder::weight;
+
+pub const CLIP_JSON : WeightDescriptor = weight!(
+  "CLIP",
+  "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
+  WeightFunction::TextTokenizer
+);
+
+/// For LCM-based image generation
+pub const LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16 : WeightDescriptor = weight!(
+  "Lykon Dreamshaper 7",
+  "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
+  WeightFunction::TextEncoder
+);
+
+/// For LCM-based image generation
+pub const LYKON_DEAMSHAPER_7_VAE : WeightDescriptor = weight!(
+  "Lykon Dreamshaper 7",
+  "lykon_dreamshaper_7_vae.safetensors",
+  WeightFunction::Vae
+);
+
+/// For LCM-based image generation
+pub const SIMIANLUO_LCM_DREAMSHAPER_V7_UNET : WeightDescriptor = weight!(
+  "SimianLuo LCM Dreamshaper v7",
+  "simianluo_lcm_dreamshaper_v7_unet.safetensors",
+  WeightFunction::Unet
+);
+
+/// Image segmentation and background removal
+/// dis_medium.onnx
+pub const DIS_MEDIUM_ONNX : WeightDescriptor = weight!(
+  "Dichotomous Image Segmentation Medium",
+  "dis_medium.onnx",
+  WeightFunction::ImageSegmentation
+);
+
+#[deprecated]
+pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = weight!(
+  "SDXL Turbo CLIP",
+  "clip_text_encoder.safetensors",
+  WeightFunction::TextEncoder
+);
+
+#[deprecated]
+pub const SDXL_TURBO_CLIP_TEXT_ENCODER_2 : WeightDescriptor = weight!(
+  "SDXL Turbo CLIP #2",
+  "clip_text_encoder_2.safetensors",
+  WeightFunction::TextEncoder
+);
+
+#[deprecated]
+pub const SDXL_TURBO_VAE : WeightDescriptor = weight!(
+  "SDXL Turbo",
+  "diffusion_pytorch_model.vae.safetensors",
+  WeightFunction::Vae
+);
+
+#[deprecated]
+pub const SDXL_TURBO_UNET : WeightDescriptor = weight!(
+  "SDXL Turbo",
+  "diffusion_pytorch_model.unet.safetensors",
+  WeightFunction::Unet
+);

crates/desktop/src-tauri/.taurignore

@@ -0,0 +1,4 @@
+# Don't restart the dev server on Rust changes
+crates/
+Cargo.toml
+Cargo.lock

crates/desktop/src-tauri/Cargo.toml

@@ -7,6 +7,7 @@ license = ""
 repository = ""
 edition = "2021"
 rust-version = "1.77.2"
+publish = false
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
@@ -22,29 +23,59 @@ tauri-build = { version = "2.0.4", features = [] }
 # Internal
 filesys = { path = "../../lib/files/filesys" }
 memory_caching = { path = "../../lib/caching/memory_caching" }
+ml_weights_registry = { path = "../ml_weights_registry" }
+primitives = { path = "../../lib/data/primitives" }
 
-# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
-cudarc = "0.13.7"
+# Internal ML
+# TODO: Only on Windows and Linux
+#ml_models = { path = "../ml_models", optional = true }
 
 # External
+#openssl = { version = "0.10", features = ["vendored"] }
 anyhow = "1.0"
-bytes = "1.10"
 base64 = "0.22.1"
-candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
-candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
-candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
-hf-hub = "0.3.2"
+bytes = "1.10"
+const_format = "0.2.34"
+directories = "6.0.0"
+dirs = "6.0.0"
+futures = "0.3"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 image = "0.25.5"
-serde_yaml = "0.9"
 log = "0.4"
 once_cell = "1.20.3"
-rand = "0.8"
+rand = "0.8.5"
+# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
+reqwest = { version = "0.12.12", features = ["stream"], default-features = false }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+serde_yaml = "0.9"
 tauri = { version = "2.2.4", features = [] }
 tauri-plugin-log = "2"
-tokenizers = "0.15.2"
+tempdir = "0.3.7"
+tempfile = "3.18.0"
+tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
+tokio-util = { version = "0.7", features = ["codec"] }
 
-[features]
-default = ["cuda"]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
+
+[target.'cfg(not(target_os = "macos"))'.dependencies]
+ml_models = { path = "../ml_models" }
+
+[dev-dependencies]
+libloading = "0.8.6"
+serial_test = "3.2.0"
+assertor = "0.0.3"
+
+#[target.'cfg(target_os = "macos")'.features]
+#default  = []
+#ml_models = ["dep:ml_models"]
+
+[profile.dev]
+incremental = true # Compile your binary in smaller steps.
+
+[profile.release]
+codegen-units = 1 # Allows LLVM to perform better optimization.
+lto = true # Enables link-time-optimizations.
+opt-level = 3 # `s` prioritizes small binary size. Use `3` if you prefer speed.
+panic = "abort" # Higher performance by disabling panic handlers.
+strip = true # Ensures debug symbols are removed.

crates/desktop/src-tauri/src/endpoints/download_models.rs

@@ -0,0 +1,12 @@
+use crate::state::app_dir::AppDataRoot;
+use tauri::State;
+
+#[tauri::command]
+pub async fn download_models(
+  _app_data_root: State<'_, AppDataRoot>,
+) -> Result<String, String> {
+
+  // NB: This endpoint has been removed for now as we now download models in the app lifecycle.
+
+  Ok(("downloaded".to_string()))
+}

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -1,117 +0,0 @@
-use crate::ml::model_cache::ModelCache;
-use crate::ml::prompt_cache::PromptCache;
-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;
-use image::imageops::FilterType;
-use image::{DynamicImage, ImageFormat, ImageReader};
-use log::{error, info};
-use std::io::Cursor;
-use std::path::PathBuf;
-use tauri::State;
-
-const PROMPT_FILENAME : &str = "prompt.txt";
-const PROMPT: &str = "A beautiful landscape with mountains and a lake";
-
-const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
-
-/// This handler takes an image (as a base64 encoded string) and a prompt and returns
-/// an image (as a base64-encoded string).
-#[tauri::command]
-pub fn infer_image(
-  image: &str,
-  prompt: Option<String>,
-  strength: Option<u8>,
-  model_config: State<AppConfig>,
-  model_cache: State<ModelCache>,
-  prompt_cache: State<PromptCache>,
-) -> Result<String, String> {
-
-  let prompt = get_prompt_or_fallback(prompt);
-
-  info!("Prompt: {}", prompt);
-
-  let image = hydrate_base64_image(image)
-    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
-
-  let result = do_infer_image(&prompt, image, strength, &model_config, &model_cache, prompt_cache);
-  
-  if let Err(err) = result.as_deref() {
-    error!("There was an error: {:?}", err);
-  }
-  
-  result
-}
-
-fn do_infer_image(
-  prompt: &str,
-  image: DynamicImage,
-  strength: Option<u8>,
-  config: &AppConfig,
-  model_cache: &ModelCache,
-  prompt_cache: State<PromptCache>,
-) -> 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,
-  };
-
-  match stable_diffusion_pipeline(args) {
-    Ok(image) => {
-      let mut bytes = Vec::with_capacity(1024*1024);
-      
-      image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
-        .map_err(|err| format!("failure to encode image: {:?}", err))?;
-
-      let bytes = BASE64_STANDARD.encode(bytes);
-
-      Ok(bytes)
-    },
-    Err(e) => Err(format!("Failed to generate image: {}", e)),
-  }
-}
-
-fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
-  let user_prompt = user_prompt.map(|prompt| prompt.trim().to_string())
-    .filter(|prompt| !prompt.is_empty());
-
-  if let Some(prompt) = user_prompt {
-    return prompt;
-  }
-
-  let prompt_file = PathBuf::from(PROMPT_FILENAME)
-    .canonicalize()
-    .unwrap_or_else(|_| PathBuf::from(PROMPT_FILENAME));
-
-  std::fs::read_to_string(&prompt_file)
-    .map_err(|err| format!("Failed to read prompt file: {}", err))
-    .unwrap_or_else(|_| {
-      error!("Failed to read prompt file: {:?}", prompt_file);
-      PROMPT.to_string()
-    })
-    .trim()
-    .to_string()
-}
-
-fn hydrate_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
-  let bytes = BASE64_STANDARD.decode(base64_image)?;
-
-  let image = ImageReader::new(Cursor::new(bytes))
-    .with_guessed_format()?
-    .decode()?;
-
-  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
-  let image = image.resize(512, 512, FilterType::CatmullRom);
-
-  Ok(image)
-}

crates/desktop/src-tauri/src/endpoints/inpaint_image_endpoint.rs

@@ -0,0 +1,13 @@
+use log::info;
+
+
+#[tauri::command]
+pub async fn inpaint_image(
+  image: &str,
+  mask: &str,
+  prompt: String,
+) -> Result<String, String> {
+  info!("inpaint_endpoint endpoint called.");
+
+  Ok(mask.to_string())
+}

crates/desktop/src-tauri/src/endpoints/mod.rs

@@ -1,3 +1,6 @@
+pub mod download_models;
 pub mod flip_image;
-pub mod image_endpoint;
-pub mod test_counter;
+pub mod inpaint_image_endpoint;
+pub mod realtime_image_endpoint;
+pub mod remove_background_endpoint;
+pub mod text_to_image_endpoint;

crates/desktop/src-tauri/src/endpoints/realtime_image_endpoint.rs

@@ -0,0 +1,143 @@
+#[cfg(not(target_os = "macos"))]
+use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
+
+use crate::state::app_config::AppConfig;
+use crate::state::app_dir::AppDataRoot;
+use crate::stubs::model_cache::ModelCache;
+use crate::stubs::prompt_cache::PromptCache;
+use crate::utils::image::decode_base64_image::decode_base64_image;
+use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
+use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png;
+use image::imageops::FilterType;
+use image::{DynamicImage, RgbImage};
+use log::{error, info};
+use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
+use std::path::PathBuf;
+use tauri::{AppHandle, State};
+
+
+const PROMPT_FILENAME : &str = "prompt.txt";
+const PROMPT: &str = "shiba inu, cute dog, detailed, walking in a wooded forest, photorealistic, 8k";
+
+const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
+
+const RANDOM_SEED: u32 = 42;
+
+
+/// This handler takes an image (as a base64 encoded string) and a prompt and returns
+/// an image (as a base64-encoded string).
+#[tauri::command]
+pub async fn infer_image(
+  image: &str,
+  prompt: Option<String>,
+  strength: Option<f64>,
+  model_config: State<'_, AppConfig>,
+  model_cache: State<'_, ModelCache>,
+  prompt_cache: State<'_, PromptCache>,
+  app_data_root: State<'_, AppDataRoot>,
+  app: AppHandle,
+) -> Result<String, String> {
+  info!("infer_image endpoint called.");
+
+  let prompt = get_prompt_or_fallback(prompt);
+
+  info!("Strength: {:?}; Prompt: {}", strength, prompt);
+
+  let image = decode_base64_image(image)
+    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
+  
+  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
+  let mut image = image.resize(512, 512, FilterType::CatmullRom);
+
+  #[cfg(not(target_os = "macos"))]
+  {
+    let result = infer_image_impl(
+      &prompt,
+      image,
+      strength,
+      &model_config,
+      &model_cache,
+      prompt_cache,
+      &app_data_root
+    ).await;
+
+    image = match result {
+      Ok(img) => DynamicImage::ImageRgb8(img),
+      Err(err) => {
+        error!("There was an error generating the image: {:?}", err);
+        return Err(format!("There was an error generating the image: {}", err));
+      }
+    };
+  }
+
+  let bytes = encode_dynamic_image_base64_png(image)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  info!("Inference successful; image converted to base64, serving back to browser...");
+
+  Ok(bytes)
+}
+
+#[cfg(not(target_os = "macos"))]
+async fn infer_image_impl(
+  prompt: &str,
+  image: DynamicImage,
+  strength: Option<f64>,
+  config: &AppConfig,
+  model_cache: &ModelCache,
+  prompt_cache: State<'_, PromptCache>,
+  app_data_root: &AppDataRoot,
+) -> Result<RgbImage, String> {
+  
+  let weights_dir = app_data_root.weights_dir();
+  let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
+  let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
+  let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
+  //let clip_weights_path= weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
+  let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
+
+  let args = Args {
+    image: &image,
+    prompt: prompt.to_string(),
+    uncond_prompt: NEGATIVE_PROMPT.to_string(),
+    model_cache,
+    prompt_cache: &prompt_cache,
+    img2img_strength: strength,
+    cfg_scale: config.cfg_scale,
+    use_flash_attn: true,
+    model_config: &config.model_config,
+    maybe_seed: None,
+    scheduler_steps: config.scheduler_steps,
+    vae_path: &vae_path,
+    unet_path: &unet_path,
+    clip_json_path: &clip_json_path,
+    clip_weights_path: &clip_weights_path,
+  };
+
+  let image = lcm_pipeline(args)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  Ok(image)
+}
+
+fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
+  let user_prompt = user_prompt.map(|prompt| prompt.trim().to_string())
+    .filter(|prompt| !prompt.is_empty());
+
+  if let Some(prompt) = user_prompt {
+    return prompt;
+  }
+
+  let prompt_file = PathBuf::from(PROMPT_FILENAME)
+    .canonicalize()
+    .unwrap_or_else(|_| PathBuf::from(PROMPT_FILENAME));
+
+  std::fs::read_to_string(&prompt_file)
+    .map_err(|err| format!("Failed to read prompt file: {}", err))
+    .unwrap_or_else(|_| {
+      error!("Failed to read prompt file: {:?}", prompt_file);
+      PROMPT.to_string()
+    })
+    .trim()
+    .to_string()
+}

crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs

@@ -0,0 +1,45 @@
+//! Adapted from https://github.com/dnanhkhoa/rust-background-removal
+
+#[cfg(not(target_os = "macos"))]
+use ml_models::ml::background_removal::remove_image_background::remove_image_background;
+
+use crate::state::app_dir::AppDataRoot;
+use crate::utils::image::decode_base64_image::decode_base64_image;
+use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
+use image::{imageops, DynamicImage};
+use log::info;
+use ml_weights_registry::weights_registry::weights::DIS_MEDIUM_ONNX;
+use tauri::State;
+
+/// This handler removes the background from an image.
+#[tauri::command]
+pub async fn remove_background(
+  image: &str,
+  app_data_root: State<'_, AppDataRoot>,
+) -> Result<String, String> {
+  info!("remove_background endpoint called.");
+
+  let mut image = decode_base64_image(image)
+      .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
+
+  #[cfg(not(target_os = "macos"))]
+  {
+    image = remove_background_impl(&app_data_root, image).await?;
+  }
+
+  let image = encode_dynamic_image_base64_png(image)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  Ok(image)
+}
+
+#[cfg(not(target_os = "macos"))]
+async fn remove_background_impl(app_data_root: &AppDataRoot, image: DynamicImage) -> Result<DynamicImage, String> {
+  let model_path = app_data_root.weights_dir().weight_path(&DIS_MEDIUM_ONNX);
+
+  let image = remove_image_background(model_path, image)
+      .await
+      .map_err(|err| format!("Couldn't remove image from image: {}", err))?;
+
+  Ok(image)
+}

crates/desktop/src-tauri/src/endpoints/test_counter.rs

@@ -1,24 +0,0 @@
-use std::sync::{Arc, RwLock};
-use once_cell::sync::Lazy;
-
-static COUNTER : Lazy<Arc<RwLock<u64>>> = Lazy::new(|| Arc::new(RwLock::new(0)));
-
-#[tauri::command]
-pub fn test_counter() -> u64 {
-  let value : u64;
-  {
-    match COUNTER.write() {
-      Ok(mut counter) => {
-        *counter += 1;
-        value = *counter;
-      },
-      Err(_e) => {
-        value = 0;
-      }
-    }
-  }
-
-  println!("I was invoked from JavaScript! {:?}", value);
-
-  value
-}

crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs

@@ -0,0 +1,115 @@
+#[cfg(not(target_os = "macos"))]
+use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
+
+use crate::state::app_config::AppConfig;
+use crate::state::app_dir::AppDataRoot;
+use crate::stubs::model_cache::ModelCache;
+use crate::stubs::prompt_cache::PromptCache;
+use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
+use image::imageops::FilterType;
+use image::{DynamicImage, ImageReader, RgbImage};
+use log::{error, info};
+use ml_weights_registry::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
+
+use once_cell::sync::Lazy;
+use std::io::Cursor;
+use tauri::{AppHandle, State};
+
+const RANDOM_SEED: u32 = 42;
+
+const STRENGTH: f64 = 100.0;
+const PNG_BYTES : &[u8] = include_bytes!("../../binary_includes/1024.png");
+
+// NB: We're just hacking a temporary "text to image" test endpoint leveraging the existing image to image modality.
+static IMAGE : Lazy<DynamicImage> = Lazy::new(||{
+  ImageReader::new(Cursor::new(PNG_BYTES))
+    .with_guessed_format()
+    .expect("failed to detect png format")
+    .decode()
+    .expect("failed to decode image")
+});
+
+#[tauri::command]
+pub async fn text_to_image(
+  prompt: String,
+  model_config: State<'_, AppConfig>,
+  model_cache: State<'_, ModelCache>,
+  prompt_cache: State<'_, PromptCache>,
+  app_data_root: State<'_, AppDataRoot>,
+  app: AppHandle,
+) -> Result<String, String> {
+  info!("text_to_image endpoint called.");
+
+  let mut image = IMAGE.resize(512, 512, FilterType::CatmullRom);
+
+  #[cfg(not(target_os = "macos"))]
+  {
+    let result = text_to_image_impl(
+      &prompt,
+      image,
+      Some(STRENGTH),
+      &model_config,
+      &model_cache,
+      prompt_cache,
+      app,
+      &app_data_root
+    ).await;
+
+    image = match result {
+      Ok(img) => DynamicImage::ImageRgb8(img),
+      Err(err) => {
+        error!("There was an error generating the image: {:?}", err);
+        return Err(format!("There was an error generating the image: {}", err));
+      }
+    };
+  }
+
+  let bytes = encode_dynamic_image_base64_png(image)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  info!("Inference successful; image converted to base64, serving back to browser...");
+
+  Ok(bytes)
+}
+
+#[cfg(not(target_os = "macos"))]
+async fn text_to_image_impl(
+  prompt: &str,
+  image: DynamicImage,
+  strength: Option<f64>,
+  config: &AppConfig,
+  model_cache: &ModelCache,
+  prompt_cache: State<'_, PromptCache>,
+  app: AppHandle,
+  app_data_root: &AppDataRoot,
+) -> Result<RgbImage, String> {
+  
+  let weights_dir = app_data_root.weights_dir();
+  let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
+  let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
+  let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
+  let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
+
+  let args = Args {
+    image: &image,
+    prompt: prompt.to_string(),
+    uncond_prompt: "".to_string(),
+    model_cache,
+    prompt_cache: &prompt_cache,
+    img2img_strength: strength,
+    cfg_scale: config.cfg_scale,
+    use_flash_attn: true,
+    model_config: &config.model_config,
+    maybe_seed: None,
+    scheduler_steps: config.scheduler_steps,
+    vae_path: &vae_path,
+    unet_path: &unet_path,
+    clip_json_path: &clip_json_path,
+    clip_weights_path: &clip_weights_path,
+  };
+
+  let image = lcm_pipeline(args)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  Ok(image)
+}

crates/desktop/src-tauri/src/events/mod.rs

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

crates/desktop/src-tauri/src/events/notification_event.rs

@@ -0,0 +1,30 @@
+use serde::Serialize;
+
+#[derive(Clone, Copy, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum NotificationModelType {
+  ImageSegmentation,
+  TextEncoder,
+  TextTokenizer,
+  Unet,
+  Vae,
+}
+
+#[derive(Clone, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum NotificationEvent<'a> {
+  ModelDownloadStarted {
+    model_name: &'a str, 
+    model_type: NotificationModelType,
+  },
+  ModelDownloadProgress {
+    model_name: &'a str,
+    model_type: NotificationModelType,
+    currently_downloaded_bytes: usize,
+    total_file_bytes: usize,
+  },
+  ModelDownloadComplete {
+    model_name: &'a str,
+    model_type: NotificationModelType,
+  }
+}

crates/desktop/src-tauri/src/lib.rs

@@ -1,12 +1,22 @@
 pub mod endpoints;
-pub mod ml;
+pub mod events;
 pub mod state;
+pub mod stubs;
+pub mod threads;
+pub mod transfer;
+pub mod utils;
 
-use crate::ml::model_cache::ModelCache;
-use crate::ml::prompt_cache::PromptCache;
+use crate::endpoints::download_models::download_models;
+use crate::endpoints::inpaint_image_endpoint::inpaint_image;
+use crate::endpoints::flip_image::flip_image;
+use crate::endpoints::realtime_image_endpoint::infer_image;
+use crate::endpoints::remove_background_endpoint::remove_background;
+use crate::endpoints::text_to_image_endpoint::text_to_image;
 use crate::state::app_config::AppConfig;
-use endpoints::image_endpoint::infer_image;
-use endpoints::test_counter::test_counter;
+use crate::stubs::model_cache::ModelCache;
+use crate::stubs::prompt_cache::PromptCache;
+use crate::threads::downloader_thread::downloader_thread;
+
 use tauri_plugin_log::Target;
 use tauri_plugin_log::TargetKind;
 
@@ -19,17 +29,16 @@ pub fn run() {
   let config = AppConfig::init()
     .expect("config should load");
 
+
   let prompt_cache = PromptCache::with_capacity(8)
     .expect("prompt cache should load");
 
   println!("Creating model cache...");
   
-  let model_cache = ModelCache::new(
-    config.device.clone(),
-    config.dtype,
-    config.sd_version.clone(),
-    config.sd_config.clone()
-  ).expect("Model cache should create");
+  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...");
 
@@ -39,6 +48,9 @@ pub fn run() {
       .targets(vec![Target::new(TargetKind::Stdout)])
       .build())
     .setup(|app| {
+      // TODO(bt): This is broken on windows
+      // log_environment_details();
+
       //if cfg!(debug_assertions) {
       //  app.handle().plugin(
       //    tauri_plugin_log::Builder::default()
@@ -46,14 +58,23 @@ 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, 
+      download_models,
+      flip_image,
       infer_image,
+      inpaint_image,
+      remove_background,
+      text_to_image,
     ])
     .run(tauri::generate_context!())
     .expect("error while running tauri application");

crates/desktop/src-tauri/src/main.rs

@@ -1,6 +1,15 @@
 // Prevents additional console window on Windows in release, DO NOT REMOVE!!
 #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
 
+use std::fs::File;
+
 fn main() {
+  //// NB: This is a test to see the behavior of the application with dynamic loading modules.
+  //// We should be able to execute code before loading CUDA and other dependencies.
+  //let r = File::create("program_started.txt");
+  //match r {
+  //  Err(err) => println!("Failed to create program_started.txt: {}", err),
+  //  Ok(_) => {}
+  //}
   app_lib::run();
 }

crates/desktop/src-tauri/src/ml/model_cache.rs

@@ -1,118 +0,0 @@
-use crate::ml::models::unet_model::UNetModel;
-use crate::ml::model_file::StableDiffusionVersion;
-use anyhow::anyhow;
-use candle_core::{DType, Device, Tensor};
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use hf_hub::api::sync::Api;
-use std::sync::{Arc, RwLock};
-use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
-use crate::ml::models::lazy_load_vae_model::LazyLoadVaeModel;
-// TODO: This data structure is gross. Generalize the locking and loading semantics for reuse.
-
-// Simple registry for now. We can build complex machinery that aids in 
-// VRAM utilization, disk space saving, intelligent scheduling, etc. in the future.
-pub struct ModelCache {
-  // TODO(bt,2025-02-14): Support for more than one device in the future.
-  device: Device,
-  dtype: DType,
-  hf_api: Api,
-  sd_config: StableDiffusionConfig,
-  sd_version: StableDiffusionVersion,
-  
-  unet: Arc<RwLock<Option<UNetModel>>>,
-  lazy_vae: LazyLoadVaeModel,
-}
-
-impl ModelCache {
-  pub fn new(
-    device: Device, 
-    dtype: DType,
-    sd_version: StableDiffusionVersion, 
-    sd_config: StableDiffusionConfig,
-  ) -> anyhow::Result<Self> {
-    let api = Api::new()?;
-    
-    Ok(Self { 
-      device: device.clone(),
-      dtype,
-      hf_api: api.clone(),
-      unet: Arc::new(RwLock::new(None)),
-      lazy_vae: LazyLoadVaeModel::new(
-        &sd_config,
-        &sd_version,
-        &api,
-        &device,
-        dtype,
-      )?,
-      sd_config,
-      sd_version,
-    })
-  }
-
-  pub fn vae_encode(&self, xs: &Tensor) -> anyhow::Result<DiagonalGaussianDistribution> {
-    self.lazy_vae.encode(xs)
-  }
-
-  pub fn vae_decode(&self, xs: &Tensor) -> anyhow::Result<Tensor> {
-    self.lazy_vae.decode(xs)
-  }
-  
-  pub fn unet_inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
-    println!("unet_inference");
-    match self.try_unet_inference(latent_model_input, timestep, text_embeddings) {
-      Err(err) => Err(err),
-      Ok(Some(result)) => Ok(result),
-      Ok(None) => {
-        // Model wasn't previously loaded
-        self.load_unet()?;
-        self.try_unet_inference(latent_model_input, timestep, text_embeddings)?
-          .ok_or(anyhow!("model was not previously loaded"))
-      }
-    }
-  }
-  
-  fn load_unet(&self) -> anyhow::Result<()> {
-    println!("load_unet");
-    match self.unet.write() {
-      Err(err) => return Err(anyhow!("lock error: {:?}", err)),
-      Ok(mut maybe_model) => {
-        match &*maybe_model {
-          Some(_model) => {} // Fall through
-          None => {
-            let repo = self.sd_version.repo();
-            
-            println!("Downloading UNET model files from: {} ...", repo);
-
-            let unet_file = self.hf_api.model(repo.to_string())
-              .get("unet/diffusion_pytorch_model.safetensors")
-              .map_err(|err| anyhow!("error fetching model: {:?}", err))?;
-
-            println!("Building UNET model... (3)");
-            let unet = UNetModel::new(&self.sd_config, unet_file, &self.device, self.dtype)
-              .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
-
-            println!("Built UNET...");
-            *maybe_model = Some(unet);
-          }
-        }
-      }
-    }
-    Ok(())
-  }
-
-  pub fn try_unet_inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Option<Tensor>> {
-    println!("try_unet_inference");
-    match self.unet.read() {
-      Ok(mut model) => {
-        match &*model {
-          None => Ok(None), // Model is not yet loaded
-          Some(model) => {
-            Ok(Some(model.inference(latent_model_input, timestep, text_embeddings)?))
-          }
-        }
-      }
-      Err(err) => Err(anyhow!("lock error: {:?}", err)),
-    }
-  }
-  
-}

crates/desktop/src-tauri/src/ml/model_file.rs

@@ -1,83 +0,0 @@
-#[cfg(feature = "accelerate")]
-extern crate accelerate_src;
-#[cfg(feature = "mkl")]
-extern crate intel_mkl_src;
-
-use candle_transformers::models::stable_diffusion;
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum ModelFile {
-  Tokenizer,
-  Tokenizer2,
-  Clip,
-  Clip2,
-  Unet,
-  Vae,
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum StableDiffusionVersion {
-  V1_5,
-  V1_5Inpaint,
-  V2_1,
-  V2Inpaint,
-  Xl,
-  XlInpaint,
-  Turbo,
-}
-
-impl StableDiffusionVersion {
-  pub fn repo(&self) -> &'static str {
-    match self {
-      Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
-      Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0",
-      Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting",
-      Self::V2_1 => "stabilityai/stable-diffusion-2-1",
-      Self::V1_5 => "runwayml/stable-diffusion-v1-5",
-      Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
-      Self::Turbo => "stabilityai/sdxl-turbo",
-    }
-  }
-}
-
-impl ModelFile {
-  pub fn get(
-    &self,
-    filename: Option<String>,
-    version: StableDiffusionVersion,
-    use_f16: bool,
-  ) -> anyhow::Result<std::path::PathBuf> {
-    use hf_hub::api::sync::Api;
-    match filename {
-      Some(filename) => Ok(std::path::PathBuf::from(filename)),
-      None => {
-        let (repo, path) = match self {
-          Self::Tokenizer => {
-            let tokenizer_repo = match version {
-              StableDiffusionVersion::V1_5
-              | StableDiffusionVersion::V2_1
-              | StableDiffusionVersion::V1_5Inpaint
-              | StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32",
-              StableDiffusionVersion::Xl
-              | StableDiffusionVersion::XlInpaint
-              | StableDiffusionVersion::Turbo => {
-                "openai/clip-vit-large-patch14"
-              }
-            };
-            (tokenizer_repo, "tokenizer.json")
-          }
-          Self::Tokenizer2 => {
-            ("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json")
-          }
-          Self::Clip => (version.repo(), "text_encoder/model.safetensors"),
-          Self::Clip2 => (version.repo(), "text_encoder_2/model.safetensors"),
-          Self::Unet => (version.repo(), "unet/diffusion_pytorch_model.safetensors"),
-          Self::Vae => (version.repo(), "vae/diffusion_pytorch_model.safetensors"),
-        };
-        let filename = Api::new()?.model(repo.to_string()).get(path)?;
-        Ok(filename)
-      }
-    }
-  }
-}
-

crates/desktop/src-tauri/src/ml/models/lazy_load_vae_model.rs

@@ -1,130 +0,0 @@
-use crate::ml::model_file::StableDiffusionVersion;
-use anyhow::anyhow;
-use candle_core::{DType, Device, Tensor};
-use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use hf_hub::api::sync::Api;
-use std::sync::{Arc, RwLock};
-
-pub struct LazyLoadVaeModel {
-  lazy_model: Arc<RwLock<Option<AutoEncoderKL>>>,
-  model_configs: VaeConfigs,
-}
-
-// TODO: Maybe we pass these via a central config read at startup, passed as an Arc<ConfigT>
-pub struct VaeConfigs {
-  sd_config: StableDiffusionConfig,
-  sd_version: StableDiffusionVersion,
-  device: Device,
-  dtype: DType,
-  hf_api: Api,
-}
-
-impl LazyLoadVaeModel {
-  pub fn new(
-    sd_config: &StableDiffusionConfig,
-    sd_version: &StableDiffusionVersion,
-    hf_api: &Api,
-    device: &Device,
-    dtype: DType,
-  ) -> anyhow::Result<Self> {
-    println!("building vae model... (1)");
-    
-    Ok(Self {
-      lazy_model: Arc::new(RwLock::new(None)),
-      model_configs: VaeConfigs {
-        sd_config: sd_config.clone(),
-        sd_version: sd_version.clone(),
-        device: device.clone(),
-        hf_api: hf_api.clone(),
-        dtype
-      }
-    })
-  }
-
-  pub fn encode(&self, xs: &Tensor) -> anyhow::Result<DiagonalGaussianDistribution> {
-    println!("vae encode...");
-    match self.try_encode(xs) {
-      Err(err) => Err(err),
-      Ok(Some(result)) => Ok(result),
-      Ok(None) => {
-        // Model wasn't previously loaded
-        self.load_model()?;
-        self.try_encode(xs)?
-          .ok_or(anyhow!("model was not previously loaded"))
-      }
-    }
-  }
-  
-  pub fn decode(&self, xs: &Tensor) -> anyhow::Result<Tensor> {
-    println!("vae decode...");
-    match self.try_decode(xs) {
-      Err(err) => Err(err),
-      Ok(Some(result)) => Ok(result),
-      Ok(None) => {
-        // Model wasn't previously loaded
-        self.load_model()?;
-        self.try_decode(xs)?
-          .ok_or(anyhow!("model was not previously loaded"))
-      }
-    }
-  }
-
-  fn try_encode(&self, xs: &Tensor) -> anyhow::Result<Option<DiagonalGaussianDistribution>> {
-    match self.lazy_model.read() {
-      Err(err) => Err(anyhow!("lock error: {:?}", err)),
-      Ok(model) => {
-        match &*model {
-          None => Ok(None), // Model is not yet loaded
-          Some(model) => {
-            Ok(Some(model.encode(xs)?))
-          }
-        }
-      }
-    }
-  }
-
-  fn try_decode(&self, xs: &Tensor) -> anyhow::Result<Option<Tensor>> {
-    match self.lazy_model.read() {
-      Err(err) => Err(anyhow!("lock error: {:?}", err)),
-      Ok(model) => {
-        match &*model {
-          None => Ok(None), // Model is not yet loaded
-          Some(model) => {
-            Ok(Some(model.decode(xs)?))
-          }
-        }
-      }
-    }
-  }
-
-  fn load_model(&self) -> anyhow::Result<()> {
-    println!("load vae");
-    match self.lazy_model.write() {
-      Err(err) => return Err(anyhow!("lock error: {:?}", err)),
-      Ok(mut maybe_model) => {
-        match &*maybe_model {
-          Some(_model) => {} // Already loaded; fall through
-          None => {
-            let repo = self.model_configs.sd_version.repo();
-
-            println!("Building VAE model from : {:?} ... (3)", repo);
-            
-            let vae_file = self.model_configs.hf_api.model(repo.to_string())
-              .get("vae/diffusion_pytorch_model.safetensors")?;
-
-            println!("Building VAE model from file {:?}...", &vae_file);
-            let vae = self.model_configs
-              .sd_config
-              .build_vae(vae_file, &self.model_configs.device, self.model_configs.dtype)?;
-
-            println!("Built VAE...");
-            *maybe_model = Some(vae);
-          }
-        }
-      }
-    }
-    Ok(())
-  }
- 
-}

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

@@ -1,2 +0,0 @@
-pub mod unet_model;
-pub mod lazy_load_vae_model;

crates/desktop/src-tauri/src/ml/models/unet_model.rs

@@ -1,35 +0,0 @@
-use candle_core::{DType, Device, Tensor};
-use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use std::path::Path;
-
-pub struct UNetModel {
-  model: UNet2DConditionModel,
-}
-
-impl UNetModel {
-  pub fn new<P: AsRef<Path>>(
-    sd_config: &StableDiffusionConfig,
-    unet_weights_file: P,
-    device: &Device,
-    dtype: DType,
-  ) -> anyhow::Result<Self> {
-    println!("building unet model... (1)");
-    
-    let model = sd_config.build_unet(
-      unet_weights_file, 
-      &device, 
-      4, 
-      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)?)
-  }
-}

crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs

@@ -1,153 +0,0 @@
-#[cfg(feature = "accelerate")]
-extern crate accelerate_src;
-#[cfg(feature = "mkl")]
-extern crate intel_mkl_src;
-
-use candle_transformers::models::stable_diffusion;
-
-use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
-use anyhow::Error as E;
-use candle_core::{DType, Device, Module, Tensor, D};
-use log::info;
-use tokenizers::Tokenizer;
-
-#[allow(clippy::too_many_arguments)]
-pub fn infer_clip_text_embeddings(
-  prompt: &str,
-  uncond_prompt: &str,
-  tokenizer: Option<String>,
-  clip_weights: Option<String>,
-  clip2_weights: Option<String>,
-  sd_version: StableDiffusionVersion,
-  sd_config: &stable_diffusion::StableDiffusionConfig,
-  use_f16: bool,
-  device: &Device,
-  dtype: DType,
-  use_guide_scale: bool,
-) -> anyhow::Result<Tensor> {
-  info!("Preparing text embeddings...");
-
-  let which = match sd_version {
-    StableDiffusionVersion::Xl
-    | StableDiffusionVersion::XlInpaint
-    | StableDiffusionVersion::Turbo => vec![true, false],
-    _ => vec![true],
-  };
-
-  let text_embeddings = which
-    .iter()
-    .map(|first| {
-      do_infer_clip_text_embeddings(
-        &prompt,
-        &uncond_prompt,
-        None, // tokenizer
-        None, // clip_weights
-        None, // clip2_weights
-        sd_version,
-        &sd_config,
-        false, // use_f16
-        &device,
-        dtype,
-        use_guide_scale,
-        *first,
-      )
-    })
-    .collect::<anyhow::Result<Vec<_>>>()?;
-
-  let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
-  
-  Ok(text_embeddings)
-}
-
-#[allow(clippy::too_many_arguments)]
-fn do_infer_clip_text_embeddings(
-    prompt: &str,
-    uncond_prompt: &str,
-    tokenizer: Option<String>,
-    clip_weights: Option<String>,
-    clip2_weights: Option<String>,
-    sd_version: StableDiffusionVersion,
-    sd_config: &stable_diffusion::StableDiffusionConfig,
-    use_f16: bool,
-    device: &Device,
-    dtype: DType,
-    use_guide_scale: bool,
-    first: bool,
-  ) -> anyhow::Result<Tensor> {
-  
-  let tokenizer_file = if first {
-    ModelFile::Tokenizer
-  } else {
-    ModelFile::Tokenizer2
-  };
-  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(),
-  };
-
-  println!("Running with prompt \"{prompt}\".");
-  let mut tokens = tokenizer
-    .encode(prompt, true)
-    .map_err(E::msg)?
-    .get_ids()
-    .to_vec();
-  if tokens.len() > sd_config.clip.max_position_embeddings {
-    anyhow::bail!(
-            "the prompt is too long, {} > max-tokens ({})",
-            tokens.len(),
-            sd_config.clip.max_position_embeddings
-        )
-  }
-  while tokens.len() < sd_config.clip.max_position_embeddings {
-    tokens.push(pad_id)
-  }
-  let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
-
-  println!("Building the Clip transformer.");
-  let clip_weights_file = if first {
-    ModelFile::Clip
-  } else {
-    ModelFile::Clip2
-  };
-  let clip_weights = if first {
-    clip_weights_file.get(clip_weights, sd_version, use_f16)?
-  } else {
-    clip_weights_file.get(clip2_weights, sd_version, use_f16)?
-  };
-  let clip_config = if first {
-    &sd_config.clip
-  } else {
-    sd_config.clip2.as_ref().unwrap()
-  };
-  let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
-
-  let text_embeddings = text_model.forward(&tokens)?;
-
-  let text_embeddings = if use_guide_scale {
-    let mut uncond_tokens = tokenizer
-      .encode(uncond_prompt, true)
-      .map_err(E::msg)?
-      .get_ids()
-      .to_vec();
-    if uncond_tokens.len() > sd_config.clip.max_position_embeddings {
-      anyhow::bail!(
-                "the negative prompt is too long, {} > max-tokens ({})",
-                uncond_tokens.len(),
-                sd_config.clip.max_position_embeddings
-            )
-    }
-    while uncond_tokens.len() < sd_config.clip.max_position_embeddings {
-      uncond_tokens.push(pad_id)
-    }
-
-    let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?;
-    let uncond_embeddings = text_model.forward(&uncond_tokens)?;
-
-    Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)?
-  } else {
-    text_embeddings.to_dtype(dtype)?
-  };
-  Ok(text_embeddings)
-}

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

@@ -1,3 +0,0 @@
-pub mod get_vae_scale;
-pub mod infer_clip_text_embeddings;
-pub mod stable_diffusion_pipeline;

crates/desktop/src-tauri/src/state/app_config.rs

@@ -1,31 +1,14 @@
-use crate::ml::model_file::StableDiffusionVersion;
+use crate::state::app_dir::AppDataRoot;
+use crate::state::os_platform::OsPlatform;
 use crate::state::yaml_config::YamlConfig;
-use candle_core::{DType, Device};
-use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
-use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
-use hf_hub::api::sync::Api;
-use log::info;
+use crate::stubs::model_config::ModelConfig;
 
-const CONFIG_FILENAME : &str = "app_config.yaml";
+const DEFAULT_SD_IMAGE_WIDTH: usize = 1024;
+const DEFAULT_SD_IMAGE_HEIGHT: usize = 1024;
 
-const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
-const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
-
-const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
-const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;
 
 /// A central place to configure app-wide details.
 pub struct AppConfig {
-  /// We only support one device type currently, and we can't multiplex 
-  /// over several devices yet.
-  pub device: Device,
-  
-  /// The core numeric type to use for models.
-  pub dtype: DType,
-  
-  pub sd_version: StableDiffusionVersion,
-  pub sd_config: StableDiffusionConfig,
-
   pub image_height: usize,
   pub image_width: usize,
   pub scheduler_steps: usize,
@@ -35,62 +18,45 @@ pub struct AppConfig {
 
   pub cfg_scale: Option<f64>,
   
-  /// Probably shouldn't be here.
-  pub hf_api: Api,
+  /// Root location for application data.
+  pub app_data_root: AppDataRoot,
+  
+  pub model_config: ModelConfig,
 }
 
 impl AppConfig {
   
   pub fn init() -> anyhow::Result<Self> {
-    println!("Configuration filename: {}", CONFIG_FILENAME);
-    let yaml_configs = load_yaml_configs();
+    let yaml_configs = YamlConfig::load_from_config_file_recursive();
     
     println!("Configs: {:?}", yaml_configs);
     
-    let device = Device::new_cuda(0)
-      .unwrap_or_else(|e| {
-          info!("CUDA not available ({}), falling back to CPU", e);
-          Device::Cpu
-        });
-    
-    let sd_config = match SD_VERSION {
-      StableDiffusionVersion::Turbo => {
-        StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
-      }
-      _ => {
-        StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
-      }
+    let os_platform = OsPlatform::get();
+
+    let default_data_root = 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 hf_api = Api::new()?;
-    
+    println!("Possible user-defined data root: {:?}", default_data_root);
+
+    let app_data_root = match default_data_root {
+      None => AppDataRoot::create_default()?,
+      Some(path) => AppDataRoot::create_existing(path)?,
+    };
+
+    println!("App data root: {:?}", app_data_root.path());
+
     Ok(Self {
-      device,
-      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
-      sd_version: SD_VERSION,
-      sd_config,
       image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
       image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
-      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(1),
+      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(4),
       scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
       seed: yaml_configs.seed,
       cfg_scale: yaml_configs.cfg_scale,
-      hf_api,
+      app_data_root,
+      model_config: ModelConfig::init()?,
     })
   }
 }
-
-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()
-}

crates/desktop/src-tauri/src/state/app_dir.rs

@@ -0,0 +1,187 @@
+use crate::state::expanduser::expanduser;
+use crate::state::os_platform::OsPlatform;
+use anyhow::anyhow;
+use directories::UserDirs;
+use ml_weights_registry::weights_registry::weight_descriptor::WeightDescriptor;
+use std::path::{Path, PathBuf};
+use tempdir::TempDir;
+use tempfile::{Builder, NamedTempFile};
+
+const DEFAULT_DATA_DIR : &str = "artcraft";
+const ASSETS_SUBDIRECTORY : &str = "assets";
+const WEIGHTS_SUBDIRECTORY : &str = "weights";
+
+const TEMPORARY_SUBDIRECTORY : &str = "temp";
+
+/// 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,
+  temp_dir: TemporaryDir,
+}
+
+#[derive(Clone)]
+pub struct AppAssetsDir {
+  path: PathBuf,
+}
+
+#[derive(Clone)]
+pub struct AppWeightsDir {
+  path: PathBuf,
+}
+
+#[derive(Clone)]
+pub struct TemporaryDir {
+  path: PathBuf,
+}
+
+impl AppDataRoot {
+  pub fn create_default() -> anyhow::Result<Self> {
+    let directory = get_default_data_dir()?;
+    Self::create_existing(directory)
+  }
+
+  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))?;
+    let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?;
+
+    Ok(Self {
+      path: dir,
+      assets_dir,
+      weights_dir,
+      temp_dir,
+    })
+  }
+  
+  pub fn assets_dir(&self) -> &AppAssetsDir {
+    &self.assets_dir
+  }
+  
+  pub fn weights_dir(&self) -> &AppWeightsDir {
+    &self.weights_dir
+  }
+
+  pub fn temp_dir(&self) -> &TemporaryDir {
+    &self.temp_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 weight_path(&self, descriptor: &WeightDescriptor) -> PathBuf {
+    self.path.join(descriptor.filename)
+  }
+}
+
+impl TemporaryDir {
+  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 with_prefix(&self, prefix: &str) -> anyhow::Result<TempDir> {
+    let tempdir = TempDir::new_in(&self.path, prefix)?;
+    Ok(tempdir)
+  }
+
+  pub fn new_named_temp_file(&self) -> anyhow::Result<NamedTempFile> {
+    let tempfile = Builder::new()
+      .prefix("temp_")
+      .suffix(".bin")
+      .tempfile_in(&self.path)?;
+    Ok(tempfile)
+  }
+}
+
+// eg. /home/bob/artcraft, /Users/bob/artcraft, or C:\Users\Alice\artcraft
+fn get_default_data_dir() -> anyhow::Result<PathBuf> {
+  Ok(UserDirs::new()
+      .ok_or_else(|| anyhow!("could not determine user home directory"))?
+      .home_dir()
+      .join(DEFAULT_DATA_DIR))
+}

crates/desktop/src-tauri/src/state/expanduser.rs

@@ -0,0 +1,73 @@
+use once_cell::sync::Lazy;
+use std::io;
+use std::path::{PathBuf, MAIN_SEPARATOR};
+
+static PREFIX: Lazy<String> = Lazy::new(|| format!("~{}", MAIN_SEPARATOR));
+
+/// The "expanduser" crate doesn't compile on Windows, so we replace its functionality slightly
+pub fn expanduser<P: AsRef<str>>(path: P) -> io::Result<PathBuf> {
+
+   Ok(match path.as_ref() {
+       // matches an exact "~"
+       s if s == "~" => {
+           home_dir()?
+       },
+       // matches paths that start with `~/`
+       s if s.starts_with(&*PREFIX) => {
+           let home = home_dir()?;
+           home.join(&s[2..])
+       },
+       // // matches paths that start with `~` but not `~/`, might be a `~username/` path
+       // s if s.starts_with("~") => {
+       //     let mut parts = s[1..].splitn(2, MAIN_SEPARATOR);
+       //     let user = parts.next()
+       //         .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "malformed path"))?;
+       //     let user = Passwd::from_name(&user)
+       //         .map_err(|_| io::Error::new(io::ErrorKind::Other, "error searching for user"))?
+       //         .ok_or_else(|| io::Error::new(io::ErrorKind::Other, format!("user '{}', does not exist", &user)))?;
+       //     if let Some(ref path) = parts.next() {
+       //         PathBuf::from(user.dir).join(&path)
+       //     } else {
+       //         PathBuf::from(user.dir)
+       //     }
+       // },
+       // nothing to expand, just make a PathBuf
+       s => PathBuf::from(s)
+   })
+}
+
+pub fn home_dir() -> io::Result<PathBuf> {
+    dirs::home_dir()
+        .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no home directory is set"))
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::state::expanduser::expanduser;
+    use std::env;
+    use std::path::PathBuf;
+
+    #[test]
+    #[serial_test::serial]
+    fn test_success() {
+        let old_home = env::var("HOME").expect("no home dir set");
+        let new_home = "/home/foo";
+        env::set_var("HOME", new_home);
+        let path = expanduser("~/path/to/directory");
+        env::set_var("HOME", old_home);
+        assert_eq!(path.expect("io error"), PathBuf::from("/home/foo/path/to/directory"));
+    }
+
+    #[test]
+    #[serial_test::serial]
+    fn test_only_tilde() {
+        let old_home = env::var("HOME").expect("no home dir set");
+        let new_home = "/home/foo";
+        env::set_var("HOME", new_home);
+        let pathstr = "~";
+        let path = expanduser(pathstr);
+        env::set_var("HOME", old_home);
+        assert_eq!(path.expect("io error"), PathBuf::from("/home/foo"));
+    }
+
+}

crates/desktop/src-tauri/src/state/mod.rs

@@ -1,2 +1,5 @@
+pub mod app_config;
+pub mod app_dir;
+pub mod os_platform;
 pub mod yaml_config;
-pub mod app_config;
+pub mod expanduser;

crates/desktop/src-tauri/src/state/os_platform.rs

@@ -0,0 +1,20 @@
+
+#[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,
+    }
+  }
+}

crates/desktop/src-tauri/src/state/yaml_config.rs

@@ -1,7 +1,10 @@
+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.
@@ -13,6 +16,15 @@ 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 {
@@ -22,4 +34,19 @@ 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()
+  }
 }

crates/desktop/src-tauri/src/stubs/mod.rs

@@ -0,0 +1,3 @@
+pub mod model_cache;
+pub mod model_config;
+pub mod prompt_cache;

crates/desktop/src-tauri/src/stubs/model_cache.rs

@@ -0,0 +1,14 @@
+#[cfg(not(target_os = "macos"))]
+pub use ml_models::ml::model_cache::ModelCache;
+
+#[cfg(not(not(target_os = "macos")))]
+pub struct ModelCache {
+  // Intentionally left blank
+}
+
+#[cfg(not(not(target_os = "macos")))]
+impl ModelCache {
+  pub fn new() -> Self {
+    Self {}
+  }
+}

crates/desktop/src-tauri/src/stubs/model_config.rs

@@ -0,0 +1,14 @@
+#[cfg(not(target_os = "macos"))]
+pub use ml_models::ml::model_config::ModelConfig;
+
+#[cfg(not(not(target_os = "macos")))]
+pub struct ModelConfig {
+  // Intentionally left blank
+}
+
+#[cfg(not(not(target_os = "macos")))]
+impl ModelConfig {
+  pub fn init() -> anyhow::Result<Self> {
+    Ok(Self {})
+  }
+}

crates/desktop/src-tauri/src/stubs/prompt_cache.rs

@@ -0,0 +1,15 @@
+
+#[cfg(not(target_os = "macos"))]
+pub use ml_models::ml::prompt_cache::PromptCache;
+
+#[cfg(not(not(target_os = "macos")))]
+pub struct PromptCache {
+  // Intentionally left blank
+}
+
+#[cfg(not(not(target_os = "macos")))]
+impl PromptCache {
+  pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
+    Ok(Self {})
+  }
+}

crates/desktop/src-tauri/src/threads/downloader_thread.rs

@@ -0,0 +1,181 @@
+use crate::events::notification_event::{NotificationEvent, NotificationModelType};
+use crate::state::app_dir::AppDataRoot;
+use crate::transfer::download::{download_async, ProgressUpdate};
+use futures::StreamExt;
+use futures::stream::FuturesUnordered;
+use log::{debug, error, info};
+use ml_weights_registry::weights_registry::weight_descriptor::WeightFunction;
+use ml_weights_registry::weights_registry::weights::{CLIP_JSON, DIS_MEDIUM_ONNX, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
+use std::collections::VecDeque;
+use std::path::PathBuf;
+use std::sync::mpsc::Receiver;
+use std::sync::{Arc, Mutex};
+use tauri::{AppHandle, Emitter};
+use tempfile::NamedTempFile;
+use tokio::task::JoinHandle;
+
+const NOTIFICATION_CHANNEL_NAME : &str = "notification";
+const MAX_FILES : usize = 8;
+const CHUNK_SIZE : usize = 1024 * 1024;
+const PARALLEL_FAILURES : usize = 8;
+const MAX_RETRIES : usize = 30;
+
+struct ModelDownloadStatusChannel {
+  pub model_name: String,
+  pub model_type: NotificationModelType,
+  pub final_file_path: PathBuf,
+  pub receiver: Receiver<ProgressUpdate>,
+}
+
+pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! {
+
+  let mut status_check_queue : Arc<Mutex<Vec<ModelDownloadStatusChannel>>> = Arc::new(Mutex::new(Vec::new()));
+
+  {
+    let status_check_queue2 = status_check_queue.clone();
+    let app2 = app.clone();
+
+    tokio::spawn(download_notify_thread(status_check_queue2, app2));
+  }
+
+  let mut download_queue = VecDeque::new();
+
+  download_queue.push_back(CLIP_JSON);
+  download_queue.push_back(LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
+  download_queue.push_back(LYKON_DEAMSHAPER_7_VAE);
+  download_queue.push_back(SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
+  download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER); // TODO(bt): Why is this still needed?
+  download_queue.push_back(DIS_MEDIUM_ONNX);
+  //download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER_2); // TODO(bt): REMOVE
+  //download_queue.push_back(SDXL_TURBO_VAE); // TODO(bt): REMOVE
+  //download_queue.push_back(SDXL_TURBO_UNET); // TODO(bt): REMOVE
+
+  let mut handles = FuturesUnordered::new();
+
+  while let Some(model) = download_queue.pop_front() {
+    let url = model.r2_download_url.to_string();
+
+    let filename = app_data_root.weights_dir().weight_path(&model);
+
+    if filename.exists() {
+      info!("File already exists: {:?}", &filename);
+      continue;
+    }
+
+    let headers = None;
+    
+    let (tx, rx) = std::sync::mpsc::channel();
+
+    info!("Download: {:?}", filename);
+
+    let temp_file = match app_data_root.temp_dir().new_named_temp_file() {
+      Ok(file) => file,
+      Err(err) => {
+        error!("Couldn't create a temporary file: {:?}", err);
+        continue;
+      }
+    };
+
+    let model_type = match model.function {
+      WeightFunction::TextEncoder => NotificationModelType::TextEncoder,
+      WeightFunction::TextTokenizer => NotificationModelType::TextTokenizer,
+      WeightFunction::Unet => NotificationModelType::Unet,
+      WeightFunction::Vae => NotificationModelType::Vae,
+      WeightFunction::ImageSegmentation => NotificationModelType::ImageSegmentation,
+    };
+
+    let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
+      model_name: model.name,
+      model_type,
+    });
+
+    if let Err(err) = result {
+      error!("Could not emit event: {}", err);
+    }
+
+    let task = tokio::spawn(download_async(
+      url,
+      temp_file,
+      filename.clone(),
+      MAX_FILES,
+      CHUNK_SIZE,
+      PARALLEL_FAILURES,
+      MAX_RETRIES,
+      headers,
+      Some(tx),
+    ));
+
+    match status_check_queue.lock() {
+      Err(err) => error!("Could not unlock: {:?}", err),
+      Ok(mut lock) => {
+        lock.push(ModelDownloadStatusChannel {
+          receiver: rx,
+          final_file_path: filename.clone(),
+          model_name: model.to_descriptive_name(),
+          model_type,
+        })
+      }
+    }
+
+    handles.push(task);
+  }
+
+  while let Some(result) = handles.next().await {
+    match result {
+      Err(err) => error!("Error downloading: {:?}", err),
+      Ok(_) => {},
+    }
+  }
+
+  loop {
+    // TODO(bt,2025-03-10): Queue management for new files entering the queue.
+    tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
+  }
+}
+
+pub async fn download_notify_thread(queue: Arc<Mutex<Vec<ModelDownloadStatusChannel>>>, app: AppHandle) -> ! {
+  loop {
+    match queue.lock() {
+      Err(err) => error!("Could not unlock: {:?}", err),
+      Ok(mut lock) => {
+        debug!(">>>>> Notifier has {} items", lock.len());
+
+        let mut reenqueue = Vec::with_capacity(lock.len());
+
+        for item in (*lock).drain(..) {
+          if let Ok(update) = item.receiver.try_recv() {
+            let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadProgress {
+              model_name: &item.model_name,
+              model_type: item.model_type,
+              currently_downloaded_bytes: update.complete,
+              total_file_bytes: update.total_length,
+            });
+
+            if let Err(err) = result {
+              error!("Could not emit event: {}", err);
+            }
+          }
+
+          if item.final_file_path.exists() {
+            debug!("No longer sending notifications for {:?}", &item.final_file_path);
+
+            let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadComplete {
+              model_name: &item.model_name,
+              model_type: item.model_type,
+            });
+
+            if let Err(err) = result {
+              error!("Could not emit event: {}", err);
+            }
+          } else {
+            reenqueue.push(item);
+          }
+        }
+
+        lock.extend(reenqueue);
+      }
+    }
+
+    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
+  }
+}

crates/desktop/src-tauri/src/threads/mod.rs

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

crates/desktop/src-tauri/src/transfer/download.rs

@@ -0,0 +1,241 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use crate::transfer::error::Error;
+use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use reqwest::header::{
+  HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, CONTENT_RANGE,
+  RANGE,
+};
+use reqwest::Url;
+use std::collections::HashMap;
+use std::fmt::Display;
+use std::fs::remove_file;
+use std::io::SeekFrom;
+use std::path::Path;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::mpsc::Sender;
+use std::time::Duration;
+use log::info;
+use tempfile::NamedTempFile;
+use tokio::fs::OpenOptions;
+use tokio::io::AsyncWriteExt;
+use tokio::io::AsyncSeekExt;
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
+
+pub struct ProgressUpdate {
+  pub complete: usize,
+  pub total_length: usize,
+}
+
+/// max_files: Number of open file handles, which determines the maximum number of parallel downloads
+/// parallel_failures:  Number of maximum failures of different chunks in parallel (cannot exceed max_files)
+/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter)
+///
+/// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in
+/// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads
+#[allow(clippy::too_many_arguments)]
+pub async fn download_async<P: AsRef<Path>>(
+  url: String,
+  filename: NamedTempFile,
+  final_filename: P,
+  max_files: usize,
+  chunk_size: usize,
+  parallel_failures: usize,
+  max_retries: usize,
+  input_headers: Option<HashMap<String, String>>,
+  maybe_progress_sender: Option<Sender<ProgressUpdate>>,
+) -> Result<(), Error> {
+  info!(">>> Downloading: {}", filename.as_ref().display());
+
+  let client = reqwest::Client::builder()
+    // https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526
+    .http2_keep_alive_timeout(Duration::from_secs(15))
+    .build()?;
+
+  let mut headers = HeaderMap::new();
+  let mut auth_token = None;
+  if let Some(input_headers) = input_headers {
+    headers.reserve(input_headers.len());
+    for (k, v) in input_headers {
+      let name: HeaderName = k
+        .try_into()
+        .map_err(|err| Error::new_err(format!("Invalid header: {err}")))?;
+      let value: HeaderValue = AsRef::<str>::as_ref(&v)
+        .try_into()
+        .map_err(|err| Error::new_err(format!("Invalid header value: {err}")))?;
+      if name == AUTHORIZATION {
+        auth_token = Some(value);
+      } else {
+        headers.insert(name, value);
+      }
+    }
+  };
+
+  let response = if let Some(token) = auth_token.as_ref() {
+    client.get(&url).header(AUTHORIZATION, token)
+  } else {
+    client.get(&url)
+  }
+    .headers(headers.clone())
+    .header(RANGE, "bytes=0-0")
+    .send()
+    .await
+    .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?
+    .error_for_status()
+    .map_err(|err| Error::new_err(err.to_string()))?;
+
+  // Only call the final redirect URL to avoid overloading the Hub with requests and also
+  // altering the download count
+  let redirected_url = response.url();
+  if Url::parse(&url)
+    .map_err(|err| Error::new_err(format!("failed to parse url: {err}")))?
+    .host()
+    == redirected_url.host()
+  {
+    if let Some(token) = auth_token {
+      headers.insert(AUTHORIZATION, token);
+    }
+  }
+
+  let content_range = response
+    .headers()
+    .get(CONTENT_RANGE)
+    .ok_or(Error::new_err("No content length".to_string()))?
+    .to_str()
+    .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
+
+  let size: Vec<&str> = content_range.split('/').collect();
+  // Content-Range: bytes 0-0/702517648
+  // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Range
+  let length: usize = size
+    .last()
+    .ok_or(Error::new_err(
+      "Error while downloading: No size was detected".to_string(),
+    ))?
+    .parse()
+    .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
+
+
+  info!("File {:?} is {:?} length", filename.as_ref(), length);
+
+  let mut handles = FuturesUnordered::new();
+  let semaphore = Arc::new(Semaphore::new(max_files));
+  let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures));
+  let max_chunk_complete = Arc::new(AtomicUsize::new(0));
+
+  for start in (0..length).step_by(chunk_size) {
+    let url = redirected_url.to_string();
+    let filename = filename.as_ref().to_owned();
+    let client = client.clone();
+    let headers = headers.clone();
+
+    let stop = std::cmp::min(start + chunk_size - 1, length);
+    let semaphore = semaphore.clone();
+    let parallel_failures_semaphore = parallel_failures_semaphore.clone();
+    let maybe_progress = maybe_progress_sender.clone();
+    let max_chunk_complete = max_chunk_complete.clone();
+    handles.push(tokio::spawn(async move {
+      let permit = semaphore
+        .acquire_owned()
+        .await
+        .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
+      let mut chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await;
+      let mut i = 0;
+      if parallel_failures > 0 {
+        while let Err(dlerr) = chunk {
+          if i >= max_retries {
+            return Err(Error::new_err(format!(
+              "Failed after too many retries ({max_retries}): {dlerr}"
+            )));
+          }
+          let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
+            Error::new_err(format!(
+              "Failed too many failures in parallel ({parallel_failures}): {dlerr} ({err})"
+            ))
+          })?;
+
+          let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
+          sleep(Duration::from_millis(wait_time as u64)).await;
+
+          chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await;
+          i += 1;
+          drop(parallel_failure_permit);
+          
+          if let Some(sender) = maybe_progress.as_ref() {
+            // NB: Calling fetch_max twice. First call will return the lower value.
+            let _old = max_chunk_complete.fetch_max(stop, Ordering::SeqCst);
+            let max_chunk = max_chunk_complete.fetch_max(stop, Ordering::SeqCst);
+            let progress_update = ProgressUpdate { 
+              complete: max_chunk,
+              total_length: length,
+            };
+            if let Err(_err) = sender.send(progress_update) {
+              // Fail silently, not essential to issue progress updates
+            }
+          }
+        }
+      }
+      drop(permit);
+      chunk.map_err(|e| Error::new_err(format!("Downloading error {e}"))).and(Ok(stop - start))
+    }));
+  }
+
+  // Output the chained result
+  while let Some(result) = handles.next().await {
+    match result {
+      Ok(Ok(size)) => {
+        //if let Some(ref callback) = callback {
+        //  callback.call((size,), None)?;
+        //}
+      }
+      Ok(Err(py_err)) => {
+        return Err(py_err);
+      }
+      Err(err) => {
+        return Err(Error::new_err(format!(
+          "Error while downloading: {err}"
+        )));
+      }
+    }
+  }
+
+  let file = filename.persist(final_filename)
+    .map_err(|e| Error::new_err(format!("Could not rename to final filename: {:?}", e)))?;
+
+  file.sync_all()?;
+
+  Ok(())
+}
+
+async fn download_chunk<P: AsRef<Path>>(
+  client: &reqwest::Client,
+  url: &str,
+  filename: P,
+  start: usize,
+  stop: usize,
+  headers: HeaderMap,
+) -> Result<(), Error> {
+  // Process each socket concurrently.
+  let range = format!("bytes={start}-{stop}");
+  let mut file = OpenOptions::new()
+    .write(true)
+    .truncate(false)
+    .create(true)
+    .open(filename)
+    .await?;
+  file.seek(SeekFrom::Start(start as u64)).await?;
+  let response = client
+    .get(url)
+    .headers(headers)
+    .header(RANGE, range)
+    .send()
+    .await?
+    .error_for_status()?;
+  let content = response.bytes().await?;
+  file.write_all(&content).await?;
+  Ok(())
+}

crates/desktop/src-tauri/src/transfer/error.rs

@@ -0,0 +1,49 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use std::fmt::Display;
+use reqwest::header::ToStrError;
+
+#[derive(Debug)]
+pub enum Error {
+  Io(std::io::Error),
+  Request(reqwest::Error),
+  ToStrError(ToStrError),
+  Misc(String),
+}
+
+impl From<std::io::Error> for Error {
+  fn from(value: std::io::Error) -> Self {
+    Self::Io(value)
+  }
+}
+
+impl From<reqwest::Error> for Error {
+  fn from(value: reqwest::Error) -> Self {
+    Self::Request(value)
+  }
+}
+
+impl From<ToStrError> for Error {
+  fn from(value: ToStrError) -> Self {
+    Self::ToStrError(value)
+  }
+}
+
+impl Display for Error {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    match self {
+      Self::Io(io) => write!(f, "Io: {io}"),
+      Self::Request(req) => write!(f, "Request: {req}"),
+      Self::ToStrError(req) => write!(f, "Response non ascii: {req}"),
+      Self::Misc(err) => write!(f, "Misc error: {err}"),
+    }
+  }
+}
+
+impl std::error::Error for Error {}
+
+impl Error {
+  pub fn new_err(reason: String) -> Self {
+    Self::Misc(reason)
+  }
+}

crates/desktop/src-tauri/src/transfer/mod.rs

@@ -0,0 +1,4 @@
+pub mod download;
+pub mod upload;
+pub mod util;
+pub mod error;

crates/desktop/src-tauri/src/transfer/upload.rs

@@ -0,0 +1,189 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use crate::transfer::error::Error;
+use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use reqwest::header::CONTENT_LENGTH;
+use std::collections::HashMap;
+use std::fmt::Display;
+use std::io::SeekFrom;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::fs::OpenOptions;
+use tokio::io::{AsyncReadExt, AsyncSeekExt};
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
+use tokio_util::codec::{BytesCodec, FramedRead};
+
+/// parts_urls: Dictionary consisting of part numbers as keys and the associated url as values
+/// completion_url: The url that should be called when the upload is finished
+/// max_files: Number of open file handles, which determines the maximum number of parallel uploads
+/// parallel_failures:  Number of maximum failures of different chunks in parallel (cannot exceed max_files)
+/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter)
+///
+/// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in
+/// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads
+///
+/// See https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html for more information
+/// on the multipart upload
+#[allow(clippy::too_many_arguments)]
+fn multipart_upload(
+  file_path: String,
+  parts_urls: Vec<String>,
+  chunk_size: u64,
+  max_files: usize,
+  parallel_failures: usize,
+  max_retries: usize,
+) -> Result<Vec<HashMap<String, String>>, Error> {
+  if parallel_failures > max_files {
+    return Err(Error::new_err(
+      "Error parallel_failures cannot be > max_files".to_string(),
+    ));
+  }
+  if (parallel_failures == 0) != (max_retries == 0) {
+    return Err(Error::new_err(
+      "For retry mechanism you need to set both `parallel_failures` and `max_retries`"
+        .to_string(),
+    ));
+  }
+
+  tokio::runtime::Builder::new_multi_thread()
+    .enable_all()
+    .build()?
+    .block_on(async {
+      upload_async(
+        file_path,
+        parts_urls,
+        chunk_size,
+        max_files,
+        parallel_failures,
+        max_retries,
+      )
+        .await
+    })
+}
+
+
+
+#[allow(clippy::too_many_arguments)]
+async fn upload_async(
+  file_path: String,
+  parts_urls: Vec<String>,
+  chunk_size: u64,
+  max_files: usize,
+  parallel_failures: usize,
+  max_retries: usize,
+) -> Result<Vec<HashMap<String, String>>, Error> {
+  let client = reqwest::Client::new();
+
+  let mut handles = FuturesUnordered::new();
+  let semaphore = Arc::new(Semaphore::new(max_files));
+  let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures));
+
+  for (part_number, part_url) in parts_urls.iter().enumerate() {
+    let url = part_url.to_string();
+    let path = file_path.to_owned();
+    let client = client.clone();
+
+    let start = (part_number as u64) * chunk_size;
+    let semaphore = semaphore.clone();
+    let parallel_failures_semaphore = parallel_failures_semaphore.clone();
+    handles.push(tokio::spawn(async move {
+      let permit = semaphore
+        .clone()
+        .acquire_owned()
+        .await
+        .map_err(|err| Error::new_err(format!("Error acquiring semaphore: {err}")))?;
+      let mut chunk = upload_chunk(&client, &url, &path, start, chunk_size).await;
+      let mut i = 0;
+      if parallel_failures > 0 {
+        while let Err(ul_err) = chunk {
+          if i >= max_retries {
+            return Err(Error::new_err(format!(
+              "Failed after too many retries ({max_retries}): {ul_err}"
+            )));
+          }
+
+          let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
+            Error::new_err(format!(
+              "Failed too many failures in parallel ({parallel_failures}): {ul_err} ({err})"
+            ))
+          })?;
+
+          let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
+          sleep(Duration::from_millis(wait_time as u64)).await;
+
+          chunk = upload_chunk(&client, &url, &path, start, chunk_size).await;
+          i += 1;
+          drop(parallel_failure_permit);
+        }
+      }
+      drop(permit);
+      chunk.map_err(|e|{
+        //match e {
+        //  Error::Io(io) => Error::new_err(format!("Io error {io}")),
+        //  Error::Request(req) => Error::new_err(format!("Error while sending chunk {req}")),
+        //  Error::ToStrError(req) => Error::new_err(format!("Response header contains non ASCII chars: {req}")),
+        //}
+        e
+      }).map(|chunk| (part_number, chunk, chunk_size))
+    }));
+  }
+
+  let mut results: Vec<HashMap<String, String>> = vec![HashMap::default(); parts_urls.len()];
+
+  while let Some(result) = handles.next().await {
+    match result {
+      Ok(Ok((part_number, headers, size))) => {
+        //if let Some(ref callback) = callback {
+        //  callback.call((size,), None)?;
+        //}
+        results[part_number] = headers;
+      }
+      Ok(Err(py_err)) => {
+        return Err(py_err);
+      }
+      Err(err) => {
+        return Err(Error::new_err(format!(
+          "Error occurred while uploading: {err}"
+        )));
+      }
+    }
+  }
+
+  Ok(results)
+}
+
+async fn upload_chunk(
+  client: &reqwest::Client,
+  url: &str,
+  path: &str,
+  start: u64,
+  chunk_size: u64,
+) -> Result<HashMap<String, String>, Error> {
+  let mut options = OpenOptions::new();
+  let mut file = options.read(true).open(path).await?;
+  let file_size = file.metadata().await?.len();
+  let bytes_transferred = std::cmp::min(file_size - start, chunk_size);
+
+  file.seek(SeekFrom::Start(start)).await?;
+  let chunk = file.take(chunk_size);
+
+  let response = client
+    .put(url)
+    .header(CONTENT_LENGTH, bytes_transferred)
+    .body(reqwest::Body::wrap_stream(FramedRead::new(
+      chunk,
+      BytesCodec::new(),
+    )))
+    .send()
+    .await?;
+  let response = response.error_for_status()?;
+  let mut headers = HashMap::new();
+  for (name, value) in response.headers().into_iter() {
+    headers.insert(name.to_string(), value.to_str()?.to_owned());
+  }
+  Ok(headers)
+}
+

crates/desktop/src-tauri/src/transfer/util.rs

@@ -0,0 +1,15 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use rand::{thread_rng, Rng};
+use std::fmt::Display;
+
+pub const BASE_WAIT_TIME: usize = 300;
+pub const MAX_WAIT_TIME: usize = 10_000;
+
+fn jitter() -> usize {
+  thread_rng().gen_range(0..=500)
+}
+
+pub fn exponential_backoff(base_wait_time: usize, n: usize, max: usize) -> usize {
+  (base_wait_time + n.pow(2) + jitter()).min(max)
+}

crates/desktop/src-tauri/src/utils/image/decode_base64_image.rs

@@ -0,0 +1,14 @@
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use image::{DynamicImage, ImageReader};
+use std::io::Cursor;
+
+pub fn decode_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
+  let bytes = BASE64_STANDARD.decode(base64_image)?;
+
+  let image = ImageReader::new(Cursor::new(bytes))
+    .with_guessed_format()?
+    .decode()?;
+
+  Ok(image)
+}

crates/desktop/src-tauri/src/utils/image/encode_dynamic_image_base64_png.rs

@@ -0,0 +1,14 @@
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use image::{DynamicImage, ImageFormat};
+use std::io::Cursor;
+
+pub fn encode_dynamic_image_base64_png(image: DynamicImage) -> anyhow::Result<String> {
+  let mut bytes = Vec::with_capacity(1024*1024);
+
+  image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
+
+  let bytes = BASE64_STANDARD.encode(bytes);
+
+  Ok(bytes)
+}

crates/desktop/src-tauri/src/utils/image/encode_rgb_image_base64_png.rs

@@ -0,0 +1,16 @@
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use image::{ImageFormat, RgbImage};
+use std::io::Cursor;
+
+pub fn encode_rgb_image_base64_png(
+  image: RgbImage
+) -> anyhow::Result<String> {
+  let mut bytes = Vec::with_capacity(1024*1024);
+
+  image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
+
+  let bytes = BASE64_STANDARD.encode(bytes);
+
+  Ok(bytes)
+}

crates/desktop/src-tauri/src/utils/image/mod.rs

@@ -0,0 +1,3 @@
+pub mod decode_base64_image;
+pub mod encode_dynamic_image_base64_png;
+pub mod encode_rgb_image_base64_png;

crates/desktop/src-tauri/src/utils/log_environment_details.rs

@@ -0,0 +1,11 @@
+use log::info;
+//use nvidia_checker::{get_cuda_version, get_cudnn_version, get_kernel_version, get_nvidia_driver_version, get_os_version, get_tensorrt_version};
+
+pub fn log_environment_details() {
+  //info!("NVIDIA driver version: {}", get_nvidia_driver_version(false));
+  //info!("CUDA version: {}", get_cuda_version(false));
+  //info!("cuDNN version: {}", get_cudnn_version(false));
+  //info!("TensorRT version: {}", get_tensorrt_version(false));
+  //info!("OS version: {}", get_os_version(false));
+  //info!("Kernel version: {}", get_kernel_version(false));
+}

crates/desktop/src-tauri/src/utils/mod.rs

@@ -0,0 +1,2 @@
+pub mod image;
+pub mod log_environment_details;

crates/desktop/src-tauri/tauri.conf.json

@@ -1,6 +1,6 @@
 {
   "$schema": "https://schema.tauri.app/config/2.0.0-rc",
-  "productName": "Realtime Desktop App",
+  "productName": "ArtCraft",
   "version": "0.1.0",
   "identifier": "com.unique.value.thing.why",
   "build": {
@@ -12,9 +12,9 @@
   "app": {
     "windows": [
       {
-        "title": "Realtime Desktop App Title Goes Here",
-        "width": 1200,
-        "height": 800,
+        "title": "ArtCraft",
+        "width": 2400,
+        "height": 1300,
         "resizable": true,
         "fullscreen": false
       }

crates/desktop/src-tauri/tests/test_dynamic_library_loading.rs

@@ -0,0 +1,7 @@
+
+#[test]
+fn test_dynamic_library_loading() {
+  let name = libloading::library_filename("cublas");
+
+  assert_eq!(name, "libcublas.so");
+}

crates/lib/data/primitives/Cargo.toml

@@ -32,3 +32,4 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" }
 
 [dev-dependencies]
 testing = { path = "../../testing" }
+assertor = "0.0.3"

crates/lib/data/primitives/src/numerics/mod.rs

@@ -1,5 +1,6 @@
 pub mod i32_to_u32_zero_clamped;
 pub mod i64_to_u64_zero_clamped;
+pub mod remap_range_f64;
 pub mod u64_to_i32_saturating;
 pub mod u64_to_i64_saturating;
 pub mod u64_to_u32_saturating;

crates/lib/data/primitives/src/numerics/remap_range_f64.rs

@@ -0,0 +1,131 @@
+
+pub struct Args {
+  pub old_min: f64,
+  pub old_max: f64,
+  pub new_min: f64,
+  pub new_max: f64,
+  
+  /// The value to remap to the new range.
+  pub value: f64,
+}
+
+pub fn remap_range_f64(args: Args) -> f64 {
+  let old_range = args.old_max - args.old_min;
+  let new_range = args.new_max -  args.new_min;
+
+  let new_value = (args.value -  args.old_min) * new_range / old_range + args.new_min;
+  
+  if new_value > args.new_max {
+    args.new_max
+  } else if new_value < args.new_min { 
+    args.new_min
+  } else {
+    new_value
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use crate::numerics::remap_range_f64::{remap_range_f64, Args};
+  use assertor::{assert_that, FloatAssertion};
+
+  #[test]
+  fn test_base_cases() {
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: 0.0,
+      new_max: 1.0,
+      value: 0.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(0.0);
+    
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: 0.0,
+      new_max: 1.0,
+      value: 1.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(1.0);
+  }
+
+  #[test]
+  fn test_maxima() {
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 100.0,
+      new_min: 0.0,
+      new_max: 1.0,
+      value: 100.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(1.0);
+
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: 0.0,
+      new_max: 50.0,
+      value: 1.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(50.0);
+
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: -100.0,
+      new_max: 100.0,
+      value: 1.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(100.0);
+  }
+
+  #[test]
+  fn test_minima() {
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 100.0,
+      new_min: 0.0,
+      new_max: 1.0,
+      value: 0.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(0.0);
+
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: 0.0,
+      new_max: 50.0,
+      value: 0.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(0.0);
+
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: -100.0,
+      new_max: 100.0,
+      value: 0.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(-100.0);
+  }
+
+  #[test]
+  fn test_other_values() {
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 100.0,
+      new_min: 0.0,
+      new_max: 1.0,
+      value: 50.0,
+    })).with_abs_tol(0.01).is_approx_equal_to(0.50);
+
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: 0.0,
+      new_max: 50.0,
+      value: 0.10,
+    })).with_abs_tol(0.01).is_approx_equal_to(5.0);
+
+    assert_that!(remap_range_f64(Args {
+      old_min: 0.0,
+      old_max: 1.0,
+      new_min: -100.0,
+      new_max: 100.0,
+      value: 0.25,
+    })).with_abs_tol(0.01).is_approx_equal_to(-50.0);
+  }
+}

crates/service/web/storyteller_web/src/configs/plans/plan_list.rs

@@ -118,24 +118,40 @@ pub static DEVELOPMENT_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
 
 /// All paid-for production premium plans in the system.
 pub static PRODUCTION_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
-    let mut plans = HashSet::new();
+  let mut plans = HashSet::new();
 
-    // ========== All-locale plans ==========
+  // ========== All-locale plans ==========
 
-    plans.insert(PlanBuilder::new("fakeyou_plus")
-        .plan_category(PlanCategory::Paid)
-        .stripe_product_id("prod_MoLv23HrxPiY7D")
-        .stripe_price_id("price_1M4jDCEU5se17MekOaJ92HYX")
-        .cost_per_month_dollars(7)
-        .tts_base_priority_level(20)
-        .tts_max_character_length(2048)
-        .tts_max_duration_seconds(30)
-        .web_vc_base_priority_level(20)
-        .web_vc_requires_frontend_keepalive(false)
-        .lipsync_requires_frontend_keepalive(false)
-        .build());
+  // Old "plus" price
+  plans.insert(PlanBuilder::new("fakeyou_plus")
+      .plan_category(PlanCategory::Paid)
+      .stripe_product_id("prod_MoLv23HrxPiY7D")
+      .stripe_price_id("price_1M4jDCEU5se17MekOaJ92HYX")
+      .cost_per_month_dollars(7)
+      .tts_base_priority_level(20)
+      .tts_max_character_length(2048)
+      .tts_max_duration_seconds(30)
+      .web_vc_base_priority_level(20)
+      .web_vc_requires_frontend_keepalive(false)
+      .lipsync_requires_frontend_keepalive(false)
+      .build());
 
-    plans.insert(PlanBuilder::new("fakeyou_pro")
+  // New "plus" price
+  plans.insert(PlanBuilder::new("fakeyou_plus_plan")
+      .plan_category(PlanCategory::Paid)
+      .stripe_product_id("prod_RvC5qUgLTSLGSD")
+      .stripe_price_id("price_1R1Lh8EU5se17MekaUNtWc59")
+      .cost_per_month_dollars(12)
+      .tts_base_priority_level(20)
+      .tts_max_character_length(2048)
+      .tts_max_duration_seconds(30)
+      .web_vc_base_priority_level(20)
+      .web_vc_requires_frontend_keepalive(false)
+      .lipsync_requires_frontend_keepalive(false)
+      .build());
+
+  // Old "pro" price
+  plans.insert(PlanBuilder::new("fakeyou_pro")
         .plan_category(PlanCategory::Paid)
         .stripe_product_id("prod_MoLw8nA6eFxHzc")
         .stripe_price_id("price_1M4jEQEU5se17MeksNfA0EKm")
@@ -149,7 +165,23 @@ pub static PRODUCTION_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
         .can_remove_visual_watermarks(true)
         .build());
 
-    plans.insert(PlanBuilder::new("fakeyou_elite")
+  // New "pro" price
+  plans.insert(PlanBuilder::new("fakeyou_pro_plan")
+      .plan_category(PlanCategory::Paid)
+      .stripe_product_id("prod_RvC6nB9rw6yAsI")
+      .stripe_price_id("price_1R1LhzEU5se17Mek3JLwN3Gg")
+      .cost_per_month_dollars(25)
+      .tts_base_priority_level(30)
+      .tts_max_character_length(3072)
+      .tts_max_duration_seconds(60 * 2)
+      .web_vc_base_priority_level(30)
+      .web_vc_requires_frontend_keepalive(false)
+      .lipsync_requires_frontend_keepalive(false)
+      .can_remove_visual_watermarks(true)
+      .build());
+
+  // Old "elite" price
+  plans.insert(PlanBuilder::new("fakeyou_elite")
         .plan_category(PlanCategory::Paid)
         .stripe_product_id("prod_MoLxQmLA4R24fv")
         .stripe_price_id("price_1M4jFREU5se17Mekc8pQaSKB")
@@ -163,6 +195,21 @@ pub static PRODUCTION_PREMIUM_PLANS: Lazy<HashSet<Plan>> = Lazy::new(|| {
         .can_remove_visual_watermarks(true)
         .build());
 
+  // New "elite" price
+  plans.insert(PlanBuilder::new("fakeyou_elite_plan")
+      .plan_category(PlanCategory::Paid)
+      .stripe_product_id("prod_RvC6i84YiAYStd")
+      .stripe_price_id("price_1R1LiQEU5se17MekiJdBcsHQ")
+      .cost_per_month_dollars(40)
+      .tts_base_priority_level(40)
+      .tts_max_character_length(4096)
+      .tts_max_duration_seconds(60 * 5)
+      .web_vc_base_priority_level(40)
+      .web_vc_requires_frontend_keepalive(false)
+      .lipsync_requires_frontend_keepalive(false)
+      .can_remove_visual_watermarks(true)
+      .build());
+
     // ========== Spanish plans ==========
 
     plans.insert(PlanBuilder::new("fakeyou_basic")
@@ -293,7 +340,7 @@ mod test {
     // making ourselves check twice when we make additions, removals, or other changes.
     #[test]
     fn test_number_of_plans_is_expected() {
-        assert_eq!(11, ALL_PLANS_BY_SLUG.len());
+        assert_eq!(14, ALL_PLANS_BY_SLUG.len());
     }
 
     #[test]
@@ -303,7 +350,7 @@ mod test {
             .filter(|plan| !plan.is_synthetic_plan()) // NB: Synthetic plans are production!
             .collect::<Vec<&Plan>>();
 
-        let expected = 4;
+        let expected = 7;
 
         assert_eq!(expected, production_plans.len());
         assert_eq!(expected, PRODUCTION_PREMIUM_PLANS.len());
@@ -346,6 +393,9 @@ mod test {
             "prod_MoLw8nA6eFxHzc".to_string(), // NB: Production
             "prod_MoLxQmLA4R24fv".to_string(), // NB: Production
             "prod_MoLyt8qMDmscjr".to_string(), // NB: Production
+            "prod_RvC5qUgLTSLGSD".to_string(), // NB: Production - new "plus" product
+            "prod_RvC6nB9rw6yAsI".to_string(), // NB: Production - new "pro" product
+            "prod_RvC6i84YiAYStd".to_string(), // NB: Production - new "elite" product
         ];
 
         assert_that(&product_ids).contains_all_of(&expected_product_ids.iter());
@@ -366,6 +416,9 @@ mod test {
             "price_1M4jEQEU5se17MeksNfA0EKm".to_string(), // NB: Production
             "price_1M4jFREU5se17Mekc8pQaSKB".to_string(), // NB: Production
             "price_1M4jGPEU5se17Mek2FiztNE5".to_string(), // NB: Production
+            "price_1R1Lh8EU5se17MekaUNtWc59".to_string(), // NB: Production - new "plus" price
+            "price_1R1LhzEU5se17Mek3JLwN3Gg".to_string(), // NB: Production - new "pro" price
+            "price_1R1LiQEU5se17MekiJdBcsHQ".to_string(), // NB: Production - new "elite" price
         ];
 
         assert_that(&price_ids).contains_all_of(&expected_price_ids.iter());

crates/vendor/candle/candle-examples/Cargo.toml

@@ -1,59 +1,52 @@
 [package]
 name = "candle-examples"
-version.workspace = true
-edition.workspace = true
-description.workspace = true
-repository.workspace = true
-keywords.workspace = true
-categories.workspace = true
-license.workspace = true
-readme = "README.md"
+edition = "2021"
 
 [dependencies]
-accelerate-src = { workspace = true, optional = true }
-candle = { workspace = true }
-candle-datasets = { workspace = true, optional = true }
-candle-nn = { workspace = true }
-candle-transformers = { workspace = true }
-candle-flash-attn = { workspace = true, optional = true }
-candle-onnx = { workspace = true, optional = true }
+accelerate-src = { version = "0.3.2", optional = true }
+candle = { path = "../candle-core", package = "candle-core" }
+candle-datasets = { path = "../candle-datasets", optional = true }
+candle-nn = { path = "../candle-nn", package = "candle-nn" }
+candle-transformers = { path = "../candle-transformers" }
+candle-flash-attn = { path = "../candle-flash-attn", optional = true }
+candle-onnx = { path = "../candle-onnx", optional = true }
 
 csv = "1.3.0"
-cudarc = { workspace = true, optional = true }
-half = { workspace = true, optional = true }
-hf-hub = { workspace = true, features = ["tokio"] }
-image = { workspace = true }
-intel-mkl-src = { workspace = true, optional = true }
-num-traits = { workspace = true }
+cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false, optional = true }
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"], optional = true }
+hf-hub = { version = "0.4.1", features = ["tokio"] }
+image = { version = "0.25.2", default-features = false, features = ["jpeg", "png"] }
+intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"], optional = true }
+num-traits = "0.2.15"
 palette = { version = "0.7.6", optional = true }
 enterpolation = { version = "0.2.1", optional = true}
 pyo3 = { version = "0.22.0", features = ["auto-initialize", "abi3-py311"], optional = true }
-rayon = { workspace = true }
+rayon = "1.7.0"
 rubato = { version = "0.15.0", optional = true }
-safetensors = { workspace = true }
-serde = { workspace = true }
-serde_json = { workspace = true }
+safetensors = "0.4.1"
+serde = { version = "1.0.171", features = ["derive"] }
+serde_json = "1.0.99"
 symphonia = { version = "0.5.3", features = ["all"], optional = true }
-tokenizers = { workspace = true, features = ["onig"] }
+tokenizers = { version = "0.19.1", default-features = false, features = ["onig"] }
 cpal = { version = "0.15.2", optional = true }
 pdf2image = { version = "0.1.2" , optional = true}
 
 [dev-dependencies]
-anyhow = { workspace = true }
-byteorder = { workspace = true }
-clap = { workspace = true }
-imageproc = { workspace = true }
-memmap2 = { workspace = true }
-rand = { workspace = true }
-ab_glyph = { workspace = true }
-tracing = { workspace = true }
-tracing-chrome = { workspace = true }
-tracing-subscriber = { workspace = true }
+anyhow = { version = "1", features = ["backtrace"] }
+byteorder = "1.4.3"
+clap = { version = "4.2.4", features = ["derive"] }
+imageproc = { version = "0.24.0", default-features = false }
+memmap2 = { version = "0.9.3", features = ["stable_deref_trait"] }
+rand = "0.8.5"
+ab_glyph = "0.2.23"
+tracing = "0.1.37"
+tracing-chrome = "0.7.1"
+tracing-subscriber = "0.3.7"
 # Necessary to disambiguate with tokio in wasm examples which are 1.28.1
 tokio = "1.43.0"
 
 [build-dependencies]
-anyhow = { workspace = true }
+anyhow = { version = "1", features = ["backtrace"] }
 bindgen_cuda = { version = "0.1.1", optional = true }
 
 [features]

crates/vendor/candle/candle-flash-attn/build.rs

@@ -86,6 +86,8 @@ 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") {
@@ -100,7 +102,12 @@ fn main() -> Result<()> {
     println!("cargo:rustc-link-search={}", build_dir.display());
     println!("cargo:rustc-link-lib=flashattention");
     println!("cargo:rustc-link-lib=dylib=cudart");
-    println!("cargo:rustc-link-lib=dylib=stdc++");
+
+
+    // TODO(bt,2025-03-12): Fix builds on windows
+    //println!("cargo:rustc-link-lib=dylib=stdc++");
+
+
 
     Ok(())
 }

crates/vendor/candle/candle-flash-attn/cutlass/.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,23 @@
+---
+name: Bug report
+about: Create a bug report to help us improve CUTLASS
+title: "[BUG]"
+labels: "? - Needs Triage, bug"
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**Steps/Code to reproduce bug**
+Follow this guide http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports to craft a minimal bug report. This helps us reproduce the issue you're having and resolve the issue more quickly.
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Environment details (please complete the following information):**
+ - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)] 
+
+**Additional context**
+Add any other context about the problem here.

crates/vendor/candle/candle-flash-attn/cutlass/.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,5 @@
+blank_issues_enabled: true
+contact_links:
+  - name: CUTLASS Discord
+    url: https://discord.gg/nvidiadeveloper
+    about: Come chat about using and contributing to CUTLASS!

crates/vendor/candle/candle-flash-attn/cutlass/.github/ISSUE_TEMPLATE/documentation_request.md

@@ -0,0 +1,35 @@
+---
+name: Documentation request
+about: Report incorrect or needed documentation to improve CUTLASS
+title: "[DOC]"
+labels: "? - Needs Triage, documentation"
+assignees: ''
+
+---
+
+## Report incorrect documentation
+
+**Location of incorrect documentation**
+Provide links and line numbers if applicable.
+
+**Describe the problems or issues found in the documentation**
+A clear and concise description of what you found to be incorrect.
+
+**Steps taken to verify documentation is incorrect**
+List any steps you have taken:
+
+**Suggested fix for documentation**
+Detail proposed changes to fix the documentation if you have any.
+
+---
+
+## Report needed documentation
+
+**Report needed documentation**
+A clear and concise description of what documentation you believe it is needed and why.
+
+**Describe the documentation you'd like**
+A clear and concise description of what you want to happen.
+
+**Steps taken to search for needed documentation**
+List any steps you have taken:

crates/vendor/candle/candle-flash-attn/cutlass/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for CUTLASS
+title: "[FEA]"
+labels: "? - Needs Triage, feature request"
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I wish I could use CUTLASS to do [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context, code examples, or references to existing implementations about the feature request here.

crates/vendor/candle/candle-flash-attn/cutlass/.github/ISSUE_TEMPLATE/submit_question.md

@@ -0,0 +1,10 @@
+---
+name: Submit question
+about: Ask a general question about CUTLASS
+title: "[QST]"
+labels: "? - Needs Triage, question"
+assignees: ''
+
+---
+
+**What is your question?**

crates/vendor/candle/candle-flash-attn/cutlass/.github/workflows/labeler.yml

@@ -0,0 +1,11 @@
+name: "Pull Request Labeler"
+on:
+- pull_request_target
+
+jobs:
+  triage:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/labeler@main
+      with:
+        repo-token: "${{ secrets.GITHUB_TOKEN }}"

crates/vendor/candle/candle-flash-attn/cutlass/.github/workflows/new-issues-to-triage-projects.yml

@@ -0,0 +1,35 @@
+name: Auto Assign New Issues to Triage Project
+
+on:
+  issues:
+    types: [opened]
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+jobs:
+  assign_one_project:
+    runs-on: ubuntu-latest
+    name: Assign to New Issues to Triage Project
+    steps:
+    - name: Process bug issues
+      uses: docker://takanabe/github-actions-automate-projects:v0.0.1
+      if: contains(github.event.issue.labels.*.name, 'bug') && contains(github.event.issue.labels.*.name, '? - Needs Triage')
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        GITHUB_PROJECT_URL: https://github.com/NVIDIA/cutlass
+        GITHUB_PROJECT_COLUMN_NAME: 'Needs prioritizing'
+    - name: Process feature issues
+      uses: docker://takanabe/github-actions-automate-projects:v0.0.1
+      if: contains(github.event.issue.labels.*.name, 'feature request') && contains(github.event.issue.labels.*.name, '? - Needs Triage')
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        GITHUB_PROJECT_URL: https://github.com/NVIDIA/cutlass
+        GITHUB_PROJECT_COLUMN_NAME: 'Needs prioritizing'
+    - name: Process other issues
+      uses: docker://takanabe/github-actions-automate-projects:v0.0.1
+      if: contains(github.event.issue.labels.*.name, '? - Needs Triage') && (!contains(github.event.issue.labels.*.name, 'bug') && !contains(github.event.issue.labels.*.name, 'feature request'))
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        GITHUB_PROJECT_URL: https://github.com/NVIDIA/cutlass
+        GITHUB_PROJECT_COLUMN_NAME: 'Needs prioritizing'

crates/vendor/candle/candle-flash-attn/cutlass/.github/workflows/stale.yml

@@ -0,0 +1,57 @@
+name: Mark inactive issues and pull requests
+
+on:
+  schedule:
+    - cron: "0 * * * *"
+
+jobs:
+  mark-inactive-30d:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Mark 30 day inactive issues and pull requests
+        uses: actions/stale@v3
+        with:
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          stale-issue-message: >
+            This issue has been labeled `inactive-30d` due to no recent activity in the past 30 days.
+            Please close this issue if no further response or action is needed.
+            Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.
+            This issue will be labeled `inactive-90d` if there is no activity in the next 60 days.
+          stale-issue-label: "inactive-30d"
+          exempt-issue-labels: "0 - Blocked,0 - Backlog,good first issue"
+          days-before-issue-stale: 30
+          days-before-issue-close: -1
+          stale-pr-message: >
+            This PR has been labeled `inactive-30d` due to no recent activity in the past 30 days.
+            Please close this PR if it is no longer required.
+            Otherwise, please respond with a comment indicating any updates.
+            This PR will be labeled `inactive-90d` if there is no activity in the next 60 days.
+          stale-pr-label: "inactive-30d"
+          exempt-pr-labels: "0 - Blocked,0 - Backlog,good first issue"
+          days-before-pr-stale: 30
+          days-before-pr-close: -1
+          operations-per-run: 50
+  mark-inactive-90d:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Mark 90 day inactive issues and pull requests
+        uses: actions/stale@v3
+        with:
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          stale-issue-message: >
+            This issue has been labeled `inactive-90d` due to no recent activity in the past 90 days.
+            Please close this issue if no further response or action is needed.
+            Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.
+          stale-issue-label: "inactive-90d"
+          exempt-issue-labels: "0 - Blocked,0 - Backlog,good first issue"
+          days-before-issue-stale: 90
+          days-before-issue-close: -1
+          stale-pr-message: >
+            This PR has been labeled `inactive-90d` due to no recent activity in the past 90 days.
+            Please close this PR if it is no longer required.
+            Otherwise, please respond with a comment indicating any updates.
+          stale-pr-label: "inactive-90d"
+          exempt-pr-labels: "0 - Blocked,0 - Backlog,good first issue"
+          days-before-pr-stale: 90
+          days-before-pr-close: -1
+          operations-per-run: 50

crates/vendor/candle/candle-flash-attn/cutlass/.gitignore

@@ -0,0 +1,3 @@
+# PyCache files
+__pycache__/
+cutlass_library.egg-info/

crates/vendor/candle/candle-flash-attn/cutlass/CHANGELOG.md

@@ -0,0 +1,467 @@
+# NVIDIA CUTLASS Changelog
+
+## [3.6.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.6.0) (2024-10-03)
+
+- [Hopper structured sparse GEMM](./examples/62_hopper_sparse_gemm/62_hopper_sparse_gemm.cu).
+  + [FP16](./test/unit/gemm/device/sm90_sparse_gemm_f16_f16_f32_tensor_op_f32.cu)
+  + [FP8](./test/unit/gemm/device/sm90_sparse_gemm_f8_f8_f32_tensor_op_f32.cu)
+  + [INT8](./test/unit/gemm/device/sm90_sparse_gemm_s8_s8_s32_tensor_op_s32.cu)
+  + [TF32](./test/unit/gemm/device/sm90_sparse_gemm_tf32_tf32_f32_tensor_op_f32.cu)
+- A refactor to the CUTLASS 3.x convolution `kernel::ConvUniversal` [API](./include/cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp) to bring it in line with `gemm::GemmUniversal`. Now the 3.x convolution API is no longer considered as a beta API.
+- [An improved mixed input GEMM](./examples/55_hopper_mixed_dtype_gemm/README.md) and a [lookup table implementation](./examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu) for `INT4`x`FP8` scale-only mode.
+- [EVT nodes for Top-K selection and softmax](./include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp) and [GEMM example using those](./examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu).
+- [Programmatic Dependent Launch](./include/cutlass/arch/grid_dependency_control.h) (PDL) that leverages a new Hopper feature to speedup two back-to-back kernels, and its corresponding [documentations](./media/docs/dependent_kernel_launch.md).
+- [A new debugging tool, synclog](./include/cutlass/arch/synclog.hpp), for dumping out all synchronization events from within a kernel to a file. Please see [synclog documentation](./media/docs/utilities.md#debugging-asynchronous-kernels-with-cutlasss-built-in-synclog-tool) for details.
+- A new TMA-enabled [epilogue](./include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp) for grouped GEMM that brings significant performance improvement, as well as its EVT support.
+- A SIMT-enabled pointer-array [epilogue](./include/cutlass/epilogue/collective/sm70_epilogue_vectorized_array.hpp).
+- A new [Ping-Pong kernel schedule for Grouped GEMM](./include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp) and some other optimizations.
+- [A new instantiation strategy for CUTLASS profiler kernels](./python/cutlass_library/sm90_shapes.py) along with [improved documentation for instantiation level in CUTLASS profiler](./media/docs/profiler.md#instantiating-more-kernels-with-hopper).
+- A new hardware support for comparisons and computations of [`cutlass::bfloat16_t`](./include/cutlass/bfloat16.h)
+- Fixed use of isnan on Windows for [`half_t`](./test/unit/core/functional.cu).
+  Various improvements and fixed from the community and CUTLASS team. Thanks to everyone who submitted PRs!
+
+## [3.5.1](https://github.com/NVIDIA/cutlass/releases/tag/v3.5.1) (2024-07-25)
+
+- [Minimal SM90 WGMMA + TMA GEMM example in 100 lines of code](./examples/cute/tutorial/wgmma_sm90.cu)
+- [Exposure of L2 `cache_hint`s in TMA copy atoms](./include/cute/arch/copy_sm90_tma.hpp#L48)
+- Exposure of raster order and tile swizzle extent in [CUTLASS library profiler](./media/docs/profiler.md#GEMM), and
+[example 48](./examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm.cu).
+- [TMA store based and EVT supported epilogues](./include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp) for [Hopper pointer array batched kernels](./test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_ptr_array.cu).
+- A new [`GemmSparseUniversal` API for CUTLASS 2.x Ampere kernels](./include/cutlass/gemm/device/gemm_sparse_universal.h) to enable serial and parallel split-k for sparse tensor cores and new tiny tile sizes to better support LLM inferrence:
+  + [FP16 TN](./test/unit/gemm/device/gemm_f16t_f16n_f32t_tensor_op_f32_sparse_sm80.cu#L269-L393) and [NT](./test/unit/gemm/device/gemm_f16n_f16t_f32t_tensor_op_f32_sparse_sm80.cu#L269-L411).
+  + [int8 TN](./test/unit/gemm/device/gemm_s8t_s8n_s32t_tensor_op_s32_sparse_sm80.cu#L264-L452).
+  + [int4 TN](./test/unit/gemm/device/gemm_s4t_s4n_s32t_tensor_op_s32_sparse_sm80.cu#L264-L452).
+  + [FP32 TN](./test/unit/gemm/device/gemm_f32t_f32n_f32t_tensor_op_f32_sparse_sm80.cu#L427-L642) and [NT](./test/unit/gemm/device/gemm_f32n_f32t_f32t_tensor_op_f32_sparse_sm80.cu#L427-L456).
+- [CUDA host adapter](./include/cutlass/cuda_host_adapter.hpp) extensions to support TMA descriptor construction driver APIs.
+- Inclusion of more [Hopper fprop, dgrad, and wgrad convolution kernels in CUTLASS library and profiler](./python/cutlass_library/generator.py).
+- Support for residual add (beta != 0) in convolution kernels.
+- A new convolution [epilogue](./examples/16_ampere_tensorop_conv2dfprop/ampere_tensorop_conv2dfprop.cu#L269) for CUTLASS 2.x to support non-packed NHWC output.
+- A refactor of [include files throughout CUTLASS core directories](./include/cutlass/gemm/collective/collective_mma_decl.hpp) to reduce circular dependencies and [tests to guard against them](./test/self_contained_includes/CMakeLists.txt).
+- [A guide for setting up VSCode to work well with CUTLASS](./media/docs/ide_setup.md) and [expanded code style guide](./media/docs/programming_guidelines.md).
+- Better support for MSVC as a host compiler.
+- Many performance optimizations, improvements, and bug fixes including fixes for FlashAttention-2.
+- Optimal code generation with CUDA toolkit versions 12.4 and 12.5u1.
+
+## [3.5.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.5.0) (2024-04-09)
+
+- Implicit GEMM Convolutions targeting Hopper SM90A via WGMMA + [TMA im2col](./include/cute/atom/copy_traits_sm90_im2col.hpp)
+  + Native implementation in CUTLASS 3.x using CuTe, mirroring the [same design hierarchy as that of GEMMs](./media/docs/gemm_api_3x.md).
+  + Support for 1D, 2D, and 3D convolutions in a [rank-agnostic fashion](./include/cutlass/conv/convnd_problem_shape.hpp).
+  + Support for [Fprop](./test/unit/conv/device_3x/fprop/sm90_conv3d_fprop_implicit_gemm_s8_s8_s32_tensorop_s32.cu), [Dgrad](./test/unit/conv/device_3x/dgrad/sm90_conv2d_dgrad_implicit_gemm_f16_f16_f32_tensorop_f16.cu), and [Wgrad](./test/unit/conv/device_3x/wgrad/sm90_conv1d_wgrad_implicit_gemm_f16_f16_f32_tensorop_f16.cu) algorithms
+  + [CUTLASS profiler support](./python/cutlass_library/conv3x_emitter.py) for 2D and 3D convolutions implemented via the 3.x API.
+  + NOTE: this is a beta release. Further updates to CUTLASS will include major performance improvements, feature enablement, and possible breaking changes to the API until 3.7 release. Your feedback is welcome on the design!
+- Support for [Ada (SM89) FP8 tensor cores via the 2.x API](./examples/58_ada_fp8_gemm/ada_fp8_gemm.cu). Requires CUDA 12.4 or newer.
+- [Ampere gather/scatter convolution example](./examples/59_ampere_gather_scatter_gemm/README.md) in CuTe and CUTLASS 3.x
+  + Showcasing how custom kernels can be written and optimized using CUTLASS 3.x and CuTe and the general strategy for implementing convolutions as specializations of GETTs.
+  + Implementation of a coarse grained sparse gather/scatter kernel achieving peak performance on Ampere class tensor cores.
+- 32x and 16x tile sizes are added to CUTLASS 2.x to improve the performance of narrow-tall and wide-short matrices.
+  + [Ampere FP16 TN](./test/unit/gemm/device/gemm_f16t_f16n_f16t_tensor_op_f32_sm80.cu) and [NT](./test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu#L227-L301), [Ampere INT8 TN](./test/unit/gemm/device/gemm_s8t_s8n_s8t_tensor_op_s32_sm80.cu#L392-L1342), [Ampere INT4 TN](./test/unit/gemm/device/gemm_s4t_s4n_s4t_tensor_op_s32_sm80.cu#L372-L934).
+  + [Turing FP16 TN](./test/unit/gemm/device/gemm_f16t_f16n_f16t_tensor_op_f32_sm75.cu#L55-L394), [Turing INT8 TN](./test/unit/gemm/device/gemm_s8t_s8n_s8t_tensor_op_s32_sm75.cu#L166-L537), [Turing INT4 TN](./test/unit/gemm/device/gemm_s4t_s4n_s4t_tensor_op_s32_sm75.cu#L310-L564).
+- Updates to CuTe documentation for [`cute::Tensor<>`](./media/docs/cute/03_tensor.md), [MMA atoms](./media/docs/cute/0t_mma_atom.md), and an overhauled [CuTe GEMM tutorial series](./examples/cute/tutorial).
+- Extensions to CuTe to support [L2 prefetching](./include/cute/algorithm/prefetch.hpp) and [TMA store+reductions](./include/cute/arch/copy_sm90_tma.hpp#L1337).
+- Remove C++11 requirement on a few CUTLASS 2.x API header files. All CUTLASS files now require C++17.
+- Fixes to greatly reduce build warnings.
+- Updates and bugfixes from the community (thanks!)
+
+## [3.4.1](https://github.com/NVIDIA/cutlass/releases/tag/v3.4.1) (2024-02-14)
+
+- Statically available [CUTLASS Version macros](./include/cutlass/version.h) that allow for handling API changes between CUTLASS releases on the users' side.
+- Improvements for Hopper [Group-GEMMs](./examples/57_hopper_grouped_gemm) and [Pointer-Array Batched GEMMs](./examples/56_hopper_ptr_array_batched_gemm).
+- Updates and bugfixes from the community (thanks!).
+
+## [3.4.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.4.0) (2024-01-12)
+* Expanded [Mixed-input Hopper GEMMs](./examples/55_hopper_mixed_dtype_gemm) support covering {16-bit, 8-bit} x {8-bit, 4-bit} input types with fast numerical converters and group scaling factors.
+* Performance improvements to [Mixed-input Hopper GEMMs](./examples/55_hopper_mixed_dtype_gemm)
+* Beta release of [Pointer-Array Batched GEMMs](./examples/56_hopper_ptr_array_batched_gemm) now available on Hopper GPUs utilizing TMA and WGMMA (requires CUDA 12.3 or above).
+* Beta release of [Group-GEMM](./examples/57_hopper_grouped_gemm) utilizing TMA and WGMMA (requires CUDA 12.3 or above).
+* [Ampere Sparse GEMM](./examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm_with_visitor.cu) supports Epilogue Visitor Tree (EVT) now.
+* NamedBarriers usability improvement and list of [ReservedNamedBarriers](./include/cutlass/arch/barrier.h) has been officially released.
+* Improved [CuTe documentation](./media/docs/cute/) including improved clarity and depth of [Quickstart](./media/docs/cute/00_quickstart.md), [CuTe Layout](./media/docs/cute/01_layout.md), and [CuTe Layout Algebra](./media/docs/cute/02_layout_algebra.md). Associated code comments, post-conditions, and details in [CuTe Core Unit Tests](./test/unit/cute/core/) also improved.
+
+## [3.3](https://github.com/NVIDIA/cutlass/releases/tag/v3.3.0) (2023-10-31)
+* [Mixed-input Hopper GEMMs](./examples/55_hopper_mixed_dtype_gemm) support covering 16-bit x 8-bit input operand types.
+* [Mixed-input Ampere GEMMs](https://github.com/NVIDIA/cutlass/pull/1084) with support for canonical layouts (TN). The implementation supports upcast on operandB {fp16, bf16} x {s8, u8}, and upcast on operandA {s8, u8} x {fp16, bf16}.
+* [Copy Async based Hopper GEMMs](./test/unit/gemm/device/sm90_gemm_bf16_bf16_bf16_alignx_tensor_op_f32_warpspecialized_cooperative.cu) - which support lower than 16B aligned input tensors.
+* Kernel schedules and Builder support for mixed precision and Copy Async GEMMs with < 16B aligned input tensors.  
+* Profiler support for lower-aligned Hopper GEMMs.
+* Performance Improvements to [Scatter-Gather Hopper Example](./examples/52_hopper_gather_scatter_fusion).
+* Sub-Byte type fixes and improvements.
+* EVT Support for RELU with Aux bitmap tensor store (used in dRELU). See [SM90 EVT fusions](./include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp) for details.
+* Fusion support for backprop fusions including drelu, dgelu, and dbias.
+* Support for void-C kernels and SM80 mixed-input GEMMs in the CUTLASS Python interface
+
+## [3.2.2](https://github.com/NVIDIA/cutlass/releases/tag/v3.2.2) (2023-10-25)
+* Minor patch for issue/1138
+
+## [3.2.1](https://github.com/NVIDIA/cutlass/releases/tag/v3.2.1) (2023-09-22)
+* Python support SM90 Epilogue Visitor Tree (EVT) on top of the C++ support released in 3.2.0.
+* SM80 EVT support in C++ and Python.
+* Other SM90 epilogue improvements.
+* Splitting CUTLASS library into smaller units based on operation, arch and datatypes. See [1105](https://github.com/NVIDIA/cutlass/discussions/1105) for details.
+* Making `tools/library/scripts` packageable - `tools/library/scripts` is now moving to `python/cutlass_library`. See the Python [README](./python/README.md) for details.
+* SM90 TF32 kernel improvements for all layouts.
+* SM90 rasterization direction support in the CUTLASS profiler.
+* Improvement for CUTLASS profiler build times.
+* Remove Python-C++ bindings.
+
+## [3.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.2.0) (2023-08-03)
+
+* New warp-specialized persistent FP8 GEMM kernel [kernel schedules](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp) and [mainloops](./include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp)  targeting Hopper architecture that achieve great performance with TMA, WGMMA, and threadblock clusters. An example showcasing [Hopper warp-specialized FP8 GEMMs](./examples/54_hopper_fp8_warp_specialized_gemm). FP8 GEMMs come with a fast accumulation mode. When enabled, problem execution might be faster but at the cost of lower accuracy because intermediate results will not periodically be promoted to a higher precision.
+* New [Epilogue Visitor Tree (EVT)](./examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu) support for Hopper TMA epilogues. EVTs allows for user-defined customized epilogue fusion patterns without having to write a new epilogue.
+* [Stream-K](./include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp) feature for Hopper. Note that this is only a functional implementation of stream-K, and should not be used for performance comparison. Optimizations are expected in a future release.
+* Improved CTA rasterization and support for CTA swizzling for Hopper kernels using the [Tile Scheduler](./include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp).
+* Improved performance for [warp-specialized TensorFloat-32 (TF32) GEMM kernels](test/unit/gemm/device/sm90_gemm_tf32_tf32_f32_tensor_op_f32_gmma_rs_cluster_warpspecialized.cu) targeting Hopper TMA.
+* [Hopper GEMM+Permute](./examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu), an example of fusing tensor reordering (permutation) with GEMM mainloop or epilogue.
+* New CUTLASS 2D Convolution Python interface. New [example](./examples/python/03_basic_conv2d.ipynb) here.
+* Support for Windows (MSVC) builds. Tested with Visual Studio 2019 v16.11.27 on Windows 10.0.
+* Optimal performance using [**CUDA 12.2u1**](https://developer.nvidia.com/cuda-downloads)
+* Updates and bugfixes from the community (thanks!)
+
+## [3.1.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.1.0) (2023-04-14)
+* New CUTLASS Python interface that aims to provide an ease-of-use interface for instantiating, emitting, compiling, and running CUTLASS kernels via Python. More details [here](./python/README.md) and new [examples](./examples/python).
+* New [efficient epilogues](test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cluster_warpspecialized_cooperative.cu#L783) using TMA for Hopper.
+* Support for [fused epilogues](test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cluster_warpspecialized_cooperative_bias_elementwise.cu), such Bias, ReLU and GELU, using the new efficient epilogues.
+* New [warp-specialized TensorFloat-32 (TF32) GEMM kernels](test/unit/gemm/device/sm90_gemm_tf32_tf32_f32_tensor_op_f32_gmma_rs_cluster_warpspecialized.cu) targeting Hopper TMA.
+* New [*warp-specialized persistent cooperative*](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp) kernel design that allows for larger tile sizes and improves performance on Hopper.
+* An [example](./examples/51_hopper_gett) showcasing GEMM-Like Tensor-Tensor Contraction (GETT) capability on Hopper.
+* Epilogue builders. Similar to mainloop builders (see [example 49](./examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu)), epilogue builders aim to generate the best-possible epilogue while exposing incremental opt-ins for greater customization.
+* Profiler support for overriding kernel and epilogue builder auto schedules for 3.x API kernels, allowing specific policies to be run in the CUTLASS profiler.
+* Performance optimizations for the [*warp-specialized persistent ping-pong*](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp) kernel.
+* Changes to the [GEMM API 3.x](./media/docs/gemm_api_3x.md), involving the host-facing arguments and the underlying `Params` structs.
+* [FMHA Backward Pass](./examples/41_fused_multi_head_attention/fused_multi_head_attention_backward.cu) from Meta xFormers.
+* [Streamk GEMM with Broadcast](./examples/47_ampere_gemm_universal_streamk/ampere_gemm_universal_streamk_broadcast.cu) enables epilogue broadcast with StreamK GEMM.
+* [Batched B2B GEMM](./examples/13_two_tensor_op_fusion) now can run multiple Back-to-Back GEMM with the same problem size in parallel.
+* [Batched Strided GEMV](test/unit/gemm/device/gemv.cu) support both row major and column major input matrix.
+* [Permute + GEMM fusion](./examples/39_gemm_permute) can fuse Permute with following GEMM now.  Before, we only support fusing GEMM with Permute in the epilogue.
+* [Row Broadcast](./include/cutlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h) can be fused in the epilogue.
+* The GitHub branch is renamed from `master` to `main` in this release.
+* Optimal performance using [**CUDA 12.1**](https://developer.nvidia.com/cuda-downloads)
+* Updates and bugfixes from the community (thanks!)
+
+## [3.0.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.0.0) (2023-01-23)
+* [CuTe](./media/docs/cute/00_quickstart.md), a [new core library and backend](./include/cute) for CUTLASS 3.0 that defines a single Layout vocabulary type and an associated algebra of layouts for a much more expressive and composable abstraction for tensors, sets of parallel agents, and operations by said agents on tensors.
+* [A new conceptual operation hierarchy](./media/docs/cutlass_3x_design.md) that replaces the architecture-centric hierarchy of CUTLASS 2.x and [documentation for CUTLASS 3.0's GEMM API changes](./media/docs/gemm_api_3x.md).
+* Strict API backwards compatibility that exposes both 2.x and 3.x API kernels through the same [`device::GemmUniversalAdapter`](./include/cutlass/gemm/device/gemm_universal_adapter.h) and [`kernel::GemmUniversal`](./include/cutlass/gemm/kernel/gemm_universal.hpp) types, allowing users to include both APIs in the same translation units. More information can be found in the [3.x backwards compatibility section](./media/docs/cutlass_3x_backwards_compatibility.md).
+* Updates to [Functionality](./media/docs/functionality.md) which directs users on which kernels are supported via CUTLASS-2 and CUTLASS-3.
+* Updates to [Compatibility](./README.md#compatibility) Section regarding supported compilers, operating systems, CUDA Toolkits, Hardware Architectures and [Target Architecture](./README.md#Target-Architecture).
+* New warp-specialized GEMM [kernel schedules](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp) and [mainloops](./include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp) targeting Hopper architecture that achieve great performance with TMA, WGMMA, and threadblock clusters.
+* Extensions to CUTLASS profiler to support threadblock cluster shapes in library and profiler tile configurations.
+* [CUTLASS library integration](./tools/library/src/gemm_operation_3x.hpp) for 3.x API kernels built through the new `CollectiveBuilder` API, enabling CUTLASS profiler.
+* Support for [Hopper GEMMs](./examples/48_hopper_warp_specialized_gemm) through the new 3.0 API with CuTe-based exposure of the Hopper [Tensor Memory Accelerator](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor) and [WGMMA Tensor Core](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions) features.
+* Set of examples that demonstrate the usage of the new 3.0 API to easily build GEMM kernels targeting Hopper: examples [48](./examples/48_hopper_warp_specialized_gemm), [49](./examples/49_hopper_gemm_schedules_with_collective_builder), and [50](./examples/50_hopper_gemm_with_epilogue_swizzle).
+
+## [2.11.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.11.0) (2022-11-19)
+* [Stream-K](./examples/47_ampere_gemm_universal_streamk), which is a new general way to do split-K.  It can not only improve performance, but can also significantly reduce the number of tile sizes that need to be profiled to find the best one.
+* [Fused multi-head attention Kernel](./examples/41_fused_multi_head_attention).  It has two variants: one uses batched GEMM for the fixed sequence length, and the other one uses group GEMM for the variable sequence length.  Both versions just need one kernel.
+* [Dual GEMM](./examples/45_dual_gemm), which can fuse A x B and A x C into one kernel. Two GEMMs has no producer-consumer dependency.
+* Hopper improves [double precision matrix multiplication](./test/unit/gemm/device/gemm_f64n_f64t_f64t_tensor_op_f64_sm90.cu) by 2x compared to Ampere at iso-clocks. It is supported since CUDA 11.8.
+* [BLAS3](./test/unit/gemm/device/hemm_cf64_cf64_cf64_tensor_op_f64_sm90.cu) functions with Hoppers new double precision matrix multiplication instructions.
+* [ELL Block Sparse GEMM](./examples/43_ell_block_sparse_gemm), which uses an [ELL matrix](https://developer.nvidia.com/blog/accelerating-matrix-multiplication-with-block-sparse-format-and-nvidia-tensor-cores/) to describe the sparsity of A matrix.  B and output matrices are still dense. The block size can be arbitary.
+* Optimized [Group Conv](./examples/42_ampere_tensorop_group_conv) for SingleGroup mode, which requires that the output channel per group is a multiple of Threadblock tile N.
+* [Optimized DepthWise Conv](./examples/46_depthwise_simt_conv2dfprop/depthwise_simt_conv2dfprop.cu).  Two new modes are added
+  * [kOptimized](./test/unit/conv/device/depthwise_conv2d_fprop_direct_conv_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) - use direct conv to compute instead of implicit GEMM. 
+    *  The restrictions are: 1) input ,output channel and group number should be multiple of (128 / sizeof(input element)). 2) The input filter size should be the same as the template parameter configuration.
+  * [kFixedStrideDilation](./test/unit/conv/device/depthwise_conv2d_fprop_direct_conv_fixed_stride_dilation_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) - which puts stride and dilation into templates to further improve the performance. In this mode, kernel persistents some inputs into register to squeeze more performance, so large filter/stride/dilation is not recommanded.
+    * The restrictions are: 1) input, output channel and group number should be multiple of (128 / sizeof(input element)). 2) input filter size, stride, dilation should same as the template parameter configuration. 
+* [Scripts](./examples/44_multi_gemm_ir_and_codegen) to fuse multiple back-to-back GEMM.  Its implementation was discussed in a GTC'22 Spring [talk](https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41606/).
+* [FP8 data type definition](./include/cutlass/float8.h) and [conversion routines](./include/cutlass/numeric_conversion.h#L1274-2115).
+* Updates and bugfixes from the community (thanks!).  Big shout out to Meta's [xFormers](https://github.com/facebookresearch/xformers).
+
+* **Deprecation announcement:** CUTLASS plans to deprecate the following:
+  * Maxwell and Pascal GPU architectures
+  * Ubuntu 16.04
+  * CUDA 10.2
+
+## [2.10.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.10.0) (2022-08-23)
+* [CUTLASS Python](./examples/40_cutlass_py) now supports GEMM, CONV, Group GEMM for different data types as well as different epilogue flavours.  
+* Optimizations for CUTLASS's [Grouped GEMM](./examples/24_gemm_grouped/gemm_grouped.cu) kernel.  Threadblock scheduling part is improved.  Some computation can be moved to the host side if applicable.  [Grouped Syr2k](./examples/38_syr2k_grouped/syr2k_grouped.cu) kernels are added, too.
+* Optimizations for [GEMM+Softmax](./examples/35_gemm_softmax).  All the reduction computation is fused into the previous GEMM.  More template arguments are provided to fine tune the performance.
+* [Grouped GEMM for Multihead Attention](./examples/41_multi_head_attention).  This general group gemm based MHA does not require the sequence length of all GEMMs to be the same which makes it most useful for natural language processing.
+* [GEMM + Layer norm fusion for Ampere](./examples/37_gemm_layernorm_gemm_fusion/) splits the layernorm into two parts and both of them can be fused into the GEMMs before and after separately.  In addition to use square sum to compute variance of layernorm, [Shift-K](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Computing_shifted_data) is provided if square sum raise numerical issues.
+* [GEMM Epilogue Permutation Fusion](./examples/39_gemm_permute) can apply user provided permutation layout mapping in the GEMM epilogue.
+* [Grouped convolution targeting implicit GEMM](test/unit/conv/device/group_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) introduces the first group convolution implementation to CUTLASS.  It is an Analytical implementation, not an Optimized.  The restrictions are: 1) input and output channel number should be multiple of group number. 2) split-K is not supported.  The implementation has 2 modes: 
+  * kSingleGroup: output channel per group is multiple of Threadblock tile N.
+  * kMultipleGroup: Threadblock tile N is multiple of output channel per group.
+* [Depthwise separable convolution](test/unit/conv/device/depthwise_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) introduces the first depthwise convolution which is also Analytical for now.  The restrictions are: 1) SIMT only 2) No split-K 3) input channel equals to output channel equals to group number.
+* Standalone [Layernorm](./tools/util/include/cutlass/util/device_layernorm.h) and [Pooling](./tools/util/include/cutlass/util/device_nhwc_pooling.h) kernels.
+* [Back-to-back GEMM/CONV](./examples/13_two_tensor_op_fusion) relaxes the requirement that the first GEMM K dimension needs to be the multiple of Threadblock Tile K dimension.
+* Optimal performance using [**CUDA 11.6u2**](https://developer.nvidia.com/cuda-downloads)
+* Updates and bugfixes from the community (thanks!)
+
+## [2.9.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.9.0) (2022-04-21)
+
+* [First layer Convolution kernels](./test/unit/conv/device/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) specialized for small channel counts and reduced alignment
+  * [Few channels](./include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h) specialization for reduced alignment capabilities
+  * [Fixed channels](./include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h) further specialized when channel count perfectly matches the access vector size
+  * [Unit tests](./test/unit/conv/device/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu)
+  * [Python-based instance emitter](./python/cutlass_library/generator.py) in the CUTLASS Library and support in the Profiler
+* [BLAS3](https://docs.nvidia.com/cuda/cublas/index.html#cublas-level-3-function-reference) operators accelerated by Tensor Cores
+  * Supported types: f32, cf32, f64, cf64, tf32x3, complex tf32x3
+  * [HERK](./test/unit/gemm/device/her2k_cf32h_cf32n_tensor_op_fast_f32_sm80.cu) with [emitter](./python/cutlass_library/rank_k_operation.py)
+  * [SYRK](./test/unit/gemm/device/syrk_f32n_f32t_tensor_op_fast_f32_sm80.cu) with [emitter](./python/cutlass_library/rank_k_operation.py)
+  * [SYMM](./test/unit/gemm/device/symm_f32n_f32n_tensor_op_fast_f32_ls_sm80.cu) with [emitter](./python/cutlass_library/symm_operation.py)
+  * [TRMM](./test/unit/gemm/device/trmm_f32n_f32t_f32t_tensor_op_fast_f32_ls_sm80.cu) with [emitter](./python/cutlass_library/trmm_operation.py)
+  * [Unit tests](./test/unit/gemm/device/testbed_rank_k_universal.h)
+* [CUTLASS Python](./examples/40_cutlass_py) demonstrating JIT compilation of CUTLASS kernels and a Python-based runtime using [CUDA Python](https://developer.nvidia.com/cuda-python)
+  * [Python-based runtime](./tools/library/scripts/rt.py) interoperable with existing emitters
+* [GEMM + Softmax example](./examples/35_gemm_softmax)
+* [Gather and Scatter Fusion with GEMM](./examples/36_gather_scatter_fusion) can gather inputs and scatters outputs based on indices vectors in the same GEMM kernel.
+  * It can select random rows in a row major matrix.
+  * It can select random columns in a column major matrix.
+* [Back-to-back GEMM/CONV](./examples/13_two_tensor_op_fusion) fully supports buffering the first GEMM/CONV results in the shared memory for the latter one to use.  It can eliminate register spill when the tile size is big.  Additionally, bias vector add is supported in the first GEMM/CONV.
+  * Supported kernels: GEMM and CONV.
+  * Supported types: fp16 and int8.
+  * Supported architectures: Turing and Ampere.
+* [Transposed Convolution](./examples/34_transposed_conv2d) (a.k.a Deconvolution) support which reuses Dgrad implementation.
+* [Utility functions](./tools/util/include/cutlass/util) that can pad NHWC and convert between NCHW and NHWC.
+* [Small alignment implicit gemm](https://github.com/NVIDIA/cutlass/issues/242) support for Fprop/Dgrad/Wgrad so that padding is no longer mandated to use tensor cores in these kernels.
+* Epilogue enhancement:
+  * Eliminate bank conflicts in int8 tensor core kernels.
+  * Half2 usage if epilogue compute type is fp16.
+  * More activation functions: Silu, Hardswish, Leaky Relu.
+  * New elementwise fusion pattern for [residual block](./include/cutlass/epilogue/thread/linear_combination_residual_block.h).
+* [Group GEMM](./examples/24_gemm_grouped) thread block number calculation fix which helps to launch the intended number of threadblocks to fully occupy the GPUs.
+* [Parallel GEMM splitk](https://github.com/NVIDIA/cutlass/pull/277) support in the CUTLASS profiler.
+* Optimal performance using [**CUDA 11.6u2**](https://developer.nvidia.com/cuda-downloads)
+* Updates and bugfixes from the community (thanks!)
+
+
+## [2.8.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.8.0) (2021-11-19)
+
+* **TF32x3:** emulated single-precision using Tensor Cores
+  * 45+ TFLOPs on NVIDIA A100
+  * [GEMM SDK example](./examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu) (real)
+  * [COMPLEX GEMM SDK example](./examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_3xtf32_complex_gemm.cu) (complex)
+  * [Implicit GEMM Convolution SDK example](./examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu)
+* **Mainloop fusion for Convolution:** convolution with fused per-channel scale-bias-relu
+  * [Conv Fprop SDK example](./examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu)
+  * [Conv WGrad SDK example](./examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu) 
+  * [cutlass::conv::device::ImplicitGemmConvolutionFusion](./include/cutlass/conv/device/implicit_gemm_convolution_fusion.h)
+* **Grouped GEMM:** similar to batched GEMM with distinct problem size per group
+  * [SDK example](./examples/24_gemm_grouped) with performance comparison with Batched Strided GEMM
+  * [cutlass::gemm::device::GemmGrouped](./include/cutlass/gemm/device/gemm_grouped.h)
+* [Implicit GEMM Convolution fusion](./examples/13_two_tensor_op_fusion/) supports staging 1st convolution's output accumulator in the shared memory on Turing. This allows more flexible warp tile sizes and less regsiter pressue.
+* Optimal performance using [**CUDA 11.5**](https://developer.nvidia.com/cuda-downloads)
+* Updates from the community (thanks!)
+
+* **Deprecation announcement:** CUTLASS plans to deprecate the following:
+  * Maxwell and Pascal GPU architectures
+  * Ubuntu 16.04
+  * CUDA 10.2
+
+## [2.7.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.7.0) (2021-09-24)
+  * Mainloop fusion for GEMM: [summation over A or B](./examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu)
+  * [Strided DGRAD (optimized iterators)](./include/cutlass/conv/kernel/default_conv2d_dgrad.h)
+  * [Half-precision GELU_taylor activation functions](./include/cutlass/epilogue/thread/activation.h#L196)
+    * Use these when accumulation and epilogue compute types are all `cutlass::half_t`
+  * Tuning and bug fixes to [fused GEMM + GEMM example](./examples/13_two_tensor_op_fusion/)
+  * Support for smaller than 128b aligned Convolutions: [see examples](test/unit/conv/device/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.cu#L272)
+  * Caching of results to accelerate Convolution [unit tests](test/unit/conv/device/cache_testbed_output.h)
+    * Can be enabled or disabled by running `cmake .. -DCUTLASS_TEST_ENABLE_CACHED_RESULTS=OFF`
+  * Corrections and bug fixes reported by the CUTLASS community
+    * Thank you for filing these issues!
+
+## [2.6.1](https://github.com/NVIDIA/cutlass/releases/tag/v2.6.1) (2021-09-03)
+  * Arbitrary padding and striding for CUTLASS Strided DGRAD Convolution operator (Analytic Iterators)
+  * Tuning for GEMMs fused with partial reductions
+  * Corrections and bug fixes reported by the CUTLASS community
+    * Thank you for filing these issues!
+
+## [2.6.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.6.0) (2021-07-22)
+  * Optimal performance when compiled with the [CUDA 11.4 Toolkit](https://developer.nvidia.com/cuda-toolkit)
+    * Adopt the new L2 prefetch feature in [cp.async](./include/cutlass/arch/memory.h) and [global load](./include/cutlass/arch/memory_sm80.h)
+  * Fused operators with GEMM and Convolution
+    * [Fused broadcast in epilogue](test/unit/gemm/device/gemm_with_broadcast_f16n_f16n_f16n_tensorop_f32_sm75.cu)
+    * [Fused partial reduction in epilogue](./test/unit/gemm/device/gemm_with_reduction_f16n_f16n_f16n_tensorop_f32_sm75.cu)
+  * 64b tensor strides and leading dimensions support for GEMMs
+  * Affine rank=2 matrix layouts 
+    * Row stride and column stride for matrices using [cutlass::layout::AffineRank2](./include/cutlass/layout/matrix.h)
+    * Support [FP64 tensor core](./examples/18_ampere_fp64_tensorop_affine2_gemm/ampere_fp64_tensorop_affine2_gemm.cu) and SIMT GEMM.
+  * [Batched GEMV](./test/unit/gemm/device/gemv.cu) preview implementation
+  * [New strided Dgrad](test/unit/conv/device/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.cu) implementation
+    * Accelerates over previous implementation by cutting down redundant math by 4x
+    * Support using new `Dy` and `w` analytic iterators and existing `cutlass::conv::device::ImplicitGemmConvolution` interface
+  * Quaternion-valued GEMM and Convolution in single- and double-precision (targeting CUDA Cores)
+    * Updates to [quaternion.h](./include/cutlass/quaternion.h) and [functional.h](./include/cutlass/functional.h)
+    * SDK Example for [GEMM](./examples/21_quaternion_gemm/quaternion_gemm.cu) and [Convolution](./examples/22_quaternion_conv/quaternion_conv.cu)
+    * [Unit tests for GEMM](./test/unit/gemm/device/simt_qgemm_nn_sm50.cu) and [Convolution](./test/unit/conv/device/conv2d_fprop_implicit_gemm_qf32nhwc_qf32nhwc_qf32nhwc_simt_f32_sm50.cu)
+  * Many improvements to the epilogue.
+    * Provide an [option](./include/cutlass/epilogue/threadblock/epilogue.h) to not fully unroll the epilogue to reduce the code size and improve the performance when using complicated elementwise operations
+    * Performance improvement for FP16 tensor core kernels
+    * Bug fixes
+  * Enhanced Clang support and the combination of Clang 13 and CUDA 11.4 can build and run kernels from Pascal and Ampere.   
+  * Updated minimum CUDA Toolkit requirement to 10.2
+    * [CUDA 11.4 Toolkit](https://developer.nvidia.com/cuda-toolkit) recommended
+  * Corrections and bug fixes reported by the CUTLASS community
+    * Thank you for filing these issues!
+
+## [2.5.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.5.0) (2021-02-26)
+  * Tensor reductions
+    * _m_-to-_n_ reductions of tensors with affine layout
+    * [Specializations](./test/unit/reduction/device/tensor_reduce_contiguous.cu) for reductions including contiguous dimension
+    * [Specializations](./test/unit/reduction/device/tensor_reduce_strided.cu) for reductions excluding contiguous dimension
+    * Custom reduction functors such as `cutlass::logical_and`
+    * Large tensor support, up to 2^63 elements (however, each dimension is limited to an extent of 2^31)
+  * Optimizations for 3-D convolution
+    * [Optimized tile iterators](./include/cutlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h) using precomputed delta table for 3-D convolution
+    * Full coverage of [forward](test/unit/conv/device/conv3d_fprop_implicit_gemm_f16ndhwc_f16ndhwc_f32ndhwc_tensor_op_f32_sm80.cu) and [backwards](test/unit/conv/device/conv3d_dgrad_implicit_gemm_f16ndhwc_f16ndhwc_f32ndhwc_tensor_op_f32_sm80.cu) passes for 3D convolution
+  * [Fused Convolution+Convolution example](./examples/13_two_tensor_op_fusion/README.md)
+  * Corrections and bug fixes reported by the CUTLASS community
+    * Thank you for filing these issues!
+  
+
+## [2.4.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.4.0) (2020-11-19)
+  * Implicit GEMM convolution kernels supporting CUDA and Tensor Cores on NVIDIA GPUs
+    * Operators: forward (Fprop), backward data gradient (Dgrad), and backward weight gradient (Wgrad) convolution
+    * Data type: FP32, complex<FP32>, Tensor Float 32 (TF32), BFloat16 (BF16), Float16, Int4, Int8, Int32
+    * Spatial dimensions: 1-D, 2-D, and 3-D
+    * Layout: NHWC, NCxHWx
+  * Implicit GEMM convolution components: 
+    * Global memory iterators supporting Fprop, Dgrad, and Wgrad
+    * `MmaMultistage` for implicit GEMM convolution for NVIDIA Ampere architecture
+    * `MmaPipeline` for implicit GEMM convolution for NVIDIA Volta and Turing architectures
+    * [Documentation](./media/docs/implicit_gemm_convolution.md) describing Implicit GEMM Convolution algorithm and implementation
+
+## [2.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.3.0) (2020-09-23)
+ * [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/)
+   * [Sparse Tensor Core GEMM kernels](test/unit/gemm/device/gemm_f16n_f16n_f32t_tensor_op_f32_sparse_sm80.cu):
+     * Direct access to Sparse Tensor Cores and maximum performance via [`mma.sp.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma-and-friends)
+   * Fast SGEMM targeting GeForce RTX 30-series CUDA Cores
+ * Minor Features:
+   * [Activation functions](./include/cutlass/epilogue/thread/activation.h) such as [GeLU](./include/cutlass/epilogue/thread/linear_combination_gelu.h) and [Sigmoid](./include/cutlass/epilogue/thread/linear_combination_sigmoid.h)
+   * Small [matrix](./include/cutlass/matrix.h) and [quaternion](./include/cutlass/quaternion.h) template classes in device code
+   * [Floating-point constants](./include/cutlass/constants.h)
+ * NVIDIA Ampere GPU Architecture examples and documentation:
+   * [Tensor Float 32](./examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu) and 
+   * [Sparse Tensor Cores](./examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm.cu)
+   * Documentation added on CUTLASS [efficient row-major epilogue](./media/docs/gemm_api.md#efficient-epilogue)
+
+## [2.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.2.0) (2020-06-08)
+ * [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/)
+   * Fast Tensor Core operations: 
+    * Maximum performance via [`mma.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma-and-friends)
+    * Tensor Float 32, BFloat16, and double-precision data types
+    * Mixed integer data types (int8, int4, bin1)
+   * Asynchronous copy for deep software pipelines via [`cp.async`](https://docs.nvidia.com/cuda/parallel-thread-execution)   
+   * Described in [GTC 2020 Webinar (SR 21745)](https://developer.nvidia.com/gtc/2020/video/s21745) (free registration required)
+ * Features:
+   * SDK examples showing GEMM fused with bias+relu and fused GEMM+GEMM
+   * Complex-valued GEMMs targeting NVIDIA Ampere Tensor Cores in double-precision and Tensor Float 32
+   * Gaussian complex GEMMs using 3m complex multiply algorithm
+   * Universal GEMM kernel supporting two batch modes and two algorithms for parallel reductions
+ * Policy updates:
+   * [CUDA 11 Toolkit](https://developer.nvidia.com/cuda-toolkit) needed to enable NVIDIA Ampere Architecture features
+   * Disabled F16C by default for compatibility - enable on cmake command line with `-DCUTLASS_ENABLE_F16C=ON`
+
+## [2.1.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.1.0) (2020-04-06)
+ * BLAS-style host-side API added to [CUTLASS Library](./media/docs/quickstart.md#cutlass-library)
+    * API to launch compiled kernel instances for GEMM and planar complex GEMM
+ * Planar Complex GEMM kernels targeting Volta and Turing Tensor Cores
+    * Computes complex matrix products on matrices stored as disjoint real and imaginary parts
+    * [SDK Examples of Planar Complex GEMMs](./examples/10_planar_complex/planar_complex.cu)
+ * Minor enhancements and bug fixes
+
+## [2.0.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.0.0) (2019-11-19)
+ * Substantially refactored for
+    * Better performance, particularly for native Turing Tensor Cores
+    * Robust and durable templates spanning the design space
+    * Encapsulated functionality embodying modern C++11 programming techniques
+    * Optimized containers and data types for efficient, generic, portable device code
+  * Updates to:
+    * [Quick start guide](./media/docs/quickstart.md)
+    * [Documentation](./README.md#documentation)
+    * [Utilities](./media/docs/utilities.md)
+    * [CUTLASS Profiler](./media/docs/profiler.md)
+ * Native Turing Tensor Cores
+    * Efficient GEMM kernels targeting Turing Tensor Cores
+    * Mixed-precision floating point, 8-bit integer, 4-bit integer, and binarized operands
+ * Coverage of existing CUTLASS functionality
+    * GEMM kernels targeting CUDA and Tensor Cores in NVIDIA GPUs
+    * Volta Tensor Cores through native mma.sync and through WMMA API
+    * Optimizations such as parallel reductions, threadblock rasterization, and intra-threadblock reductions
+    * Batched GEMM operations
+    * Complex-valued GEMMs
+ * **Note: a host compiler supporting C++11 or greater is required.**
+
+# CUTLASS 1.x
+
+## [1.3.2](https://github.com/NVIDIA/cutlass/releases/tag/v1.3.2) (2019-07-09)
+ * Performance improvement for Volta Tensor Cores TN and TT layouts.
+
+## [1.3.1](https://github.com/NVIDIA/cutlass/releases/tag/v1.3.1) (2019-04-09)
+ * Corrected NVRTC unit tests.
+
+## [1.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v1.3.0) (2019-03-20)
+ * Efficient GEMM kernel targeting Volta Tensor Cores via `mma.sync` instruction added in CUDA 10.1.
+
+## [1.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v1.2.0) (2018-10-26)
+ * Parallelized reductions across threadblocks ("Split-K")
+   * Improved IGEMM performance
+ * Batched strided WMMA GEMMs
+
+## [1.1.0](https://github.com/NVIDIA/cutlass/releases/tag/v1.1.0) (2018-09-19)
+  * Turing Features
+    * WMMA GEMM targeting TensorCores - INT8, INT4, 1-bit
+  * Batched Strided GEMM
+  * Threadblock rasterization strategies
+    * Improved performance for adverse problem sizes and data layouts
+  * Extended CUTLASS Core comonents
+    * Tensor views support arbitrary matrix and tensor layouts
+    * Zip iterators for structuring multiple data streams
+  * Enhanced CUTLASS utilities
+    * Reference code for tensor operations in host and device code
+    * Added HostMatrix<> for simplified matrix creation
+  * Examples
+    * Basic GEMM, tensor views, CUTLASS utilities, batched GEMM, WMMA GEMM
+
+## [1.0.1](https://github.com/NVIDIA/cutlass/releases/tag/v1.0.1) (2018-06-11)
+
+  * Intra-threadblock reduction added for small threadblock tile sizes
+    * sgemm_64x128x16, sgemm_128x128x16, sgemm_128x64x16, sgemm_128x32x16, sgemm_64x64x16, sgemm_64x32x16
+    * igemm_32x32x128
+  * GEMM _K_ residue handled during prologue prior to mainloop
+  * Replaced Google Test copy with submodule. Use `git submodule init --recursive --update`
+
+## [1.0.0](https://github.com/NVIDIA/cutlass/commit/2028ebe120aab22bfd0b2baf8902d4c9627eb33f) (2018-05-16)
+
+  * Substantial rewrite to accommodate new architecture
+  * Kernels: SGEMM, DGEMM, IGEMM, HGEMM, WMMA GEMM
+  * Unit and performance tests
+
+## [0.0.1](https://github.com/NVIDIA/cutlass/commit/d08ba8ac46e2fa3f745e070c390182edb56b2e91) (2017-12-04)
+
+  * Initial release
+
+
+## Copyright
+
+Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: BSD-3-Clause
+
+```
+  Redistribution and use in source and binary forms, with or without
+  modification, are permitted provided that the following conditions are met:
+
+  1. Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+  2. Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+  3. Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+```
+

crates/vendor/candle/candle-flash-attn/cutlass/CITATION.cff

@@ -0,0 +1,112 @@
+cff-version: 1.2.0
+title: CUTLASS
+message: >-
+  If you use this software, please cite using the
+  following metadata.
+type: software
+authors:
+  - given-names: Vijay
+    family-names: Thakkar
+    email: vithakkar@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Pradeep
+    family-names: Ramani
+    email: prramani@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Cris
+    family-names: Cecka
+    email: ccecka@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Aniket
+    family-names: Shivam
+    email: ashivam@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Honghao
+    family-names: Lu
+    email: honghaol@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Ethan
+    family-names: Yan
+    email: etyan@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Jack
+    family-names: Kosaian
+    email: jkosaian@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Mark
+    family-names: Hoemmen
+    email: mhoemmen@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Haicheng
+    family-names: Wu
+    email: haichengw@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Andrew
+    family-names: Kerr
+    email: akerr@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Matt
+    family-names: Nicely
+    email: mnicely@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Duane
+    family-names: Merrill
+    email: dumerrill@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Dustyn
+    family-names: Blasig
+    email: dblasig@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Fengqi
+    family-names: Qiao
+    email: fqiao@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Piotr
+    family-names: Majcher
+    email: pmajcher@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Paul
+    family-names: Springer
+    email: pspringer@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Markus
+    family-names: Hohnerbach
+    affiliation: NVIDIA
+    email: mhohnerbach@nvidia.com
+  - given-names: Jin
+    family-names: Wang
+    email: jinw@nvidia.com
+    affiliation: NVIDIA
+  - given-names: Manish
+    family-names: Gupta
+    affiliation: Google
+    email: manigupta@google.com
+
+
+repository-code: 'https://github.com/NVIDIA/cutlass'
+abstract: >-
+  CUTLASS is a collection of CUDA C++ template
+  abstractions for implementing high-performance
+  matrix-multiplication (GEMM) and related
+  computations at all levels and scales within CUDA.
+  It incorporates strategies for hierarchical
+  decomposition and data movement similar to those
+  used to implement cuBLAS and cuDNN. CUTLASS
+  decomposes these "moving parts" into reusable,
+  modular software components abstracted by C++
+  template classes. These thread-wide, warp-wide,
+  block-wide, and device-wide primitives can be
+  specialized and tuned via custom tiling sizes, data
+  types, and other algorithmic policy. The resulting
+  flexibility simplifies their use as building blocks
+  within custom kernels and applications.
+keywords:
+  - 'cutlass, tensor cores, cuda, cute, nvidia, gpu, linear algebra, matrix computations'
+license: BSD-3-Clause
+license-url: https://github.com/NVIDIA/cutlass/blob/v3.0.0/LICENSE.txt
+version: '3.0.0'
+date-released: '2023-01-23'
+identifiers:
+  - type: url
+    value: "https://github.com/NVIDIA/cutlass/tree/v3.0.0"
+    description: The GitHub release URL of tag 3.0.0