don't download sdxl anymore

Usage example - Conditional Rendering Based on App Mode:
```typescript
import { appMode } from "~/signals";

function ConditionalComponent() {
  return (
    <div className="app-container">
      {/* Show/hide components based on current mode */}
      {appMode.value === 'image' && (
        <div className="image-tools">
          <h3>Image Mode Tools</h3>
        </div>
      )}
      
      {appMode.value === 'edit' && (
        <div className="edit-tools">
          <h3>Edit Mode Tools</h3>
        </div>
      )}
  );
}
```

.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

@@ -96,6 +96,17 @@ members = [
   "crates/service/plugins/user_traits_component",
 
   # 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",
+  "crates/vendor/candle/candle-nn",
+  "crates/vendor/candle/candle-onnx",
+  "crates/vendor/candle/candle-pyo3",
+  "crates/vendor/candle/candle-transformers",
+  "crates/vendor/candle/tensor-tools",
   "crates/vendor/limitation",
   # NB(bt,2023-12-12): Newrelic keeps breaking with upgrades. Kill for now.
   #"crates/vendor/newrelic-telemetry-sdk-rust", 

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/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

@@ -27,24 +27,56 @@ memory_caching = { path = "../../lib/caching/memory_caching" }
 cudarc = "0.13.7"
 
 # External
+#openssl = { version = "0.10", features = ["vendored"] }
 anyhow = "1.0"
-bytes = "1.10"
 base64 = "0.22.1"
-candle-core = { version = "0.8.2", features = ["cuda"] }
-candle-nn = { version = "0.8.2", features = ["cuda"] }
-candle-transformers = { version = "0.8.2", features = ["cuda"] }
+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"] }
+dirs = "6.0.0"
+directories = "6.0.0"
+futures = "0.3"
 hf-hub = "0.3.2"
 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"
+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"] }
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 
 [features]
 default = ["cuda"]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/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"
+
+[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,16 @@
+use crate::ml::downloads::download_all_models::download_all_models;
+use crate::state::app_dir::AppDataRoot;
+use tauri::State;
+
+#[tauri::command]
+pub async fn download_models(
+  app_data_root: State<'_, AppDataRoot>,
+) -> Result<String, String> {
+  
+  download_all_models(&app_data_root).await
+    .map_err(|err| {
+      "error downloading models".to_string()
+    })?;
+  
+  Ok(("downloaded".to_string()))
+}

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

@@ -1,6 +1,7 @@
+use crate::ml::downloads::download_all_models::download_all_models;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
-use crate::ml::stable_diffusion::stable_diffusion_pipeline::{stable_diffusion_pipeline, Args};
+use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
 use crate::state::app_config::AppConfig;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use bytes::BytesMut;
@@ -9,33 +10,58 @@ use image::{DynamicImage, ImageFormat, ImageReader};
 use log::{error, info};
 use std::io::Cursor;
 use std::path::PathBuf;
-use tauri::State;
+use tauri::{AppHandle, State};
+use crate::state::app_dir::AppDataRoot;
 
 const PROMPT_FILENAME : &str = "prompt.txt";
-const PROMPT: &str = "A beautiful landscape with mountains and a lake";
+const PROMPT: &str = "green goblin, blood dripping from lips, detailed, photorealistic, 8k";
 
 const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
 
+const RANDOM_SEED: u32 = 42;
+
+// TODO: Mark async https://github.com/tauri-apps/tauri/discussions/7737
+
+
 /// This handler takes an image (as a base64 encoded string) and a prompt and returns
 /// an image (as a base64-encoded string).
 #[tauri::command]
-pub fn infer_image(
+pub async fn infer_image(
   image: &str,
   prompt: Option<String>,
   strength: Option<u8>,
-  model_config: State<AppConfig>,
-  model_cache: State<ModelCache>,
-  prompt_cache: State<PromptCache>,
+  model_config: State<'_, AppConfig>,
+  model_cache: State<'_, ModelCache>,
+  prompt_cache: State<'_, PromptCache>,
+  app_data_root: State<'_, AppDataRoot>,
+  app: AppHandle,
 ) -> Result<String, String> {
 
   let prompt = get_prompt_or_fallback(prompt);
 
   info!("Prompt: {}", prompt);
+  info!("Strength: {:?}", strength);
 
   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);
+  // TODO(bt): Move this to another endpoint
+  download_all_models(&app_data_root).await
+    .map_err(|err| {
+      error!("couldn't download models: {:?}", err);
+      "Couldn't download models".to_string()
+    })?;
+
+  let result = do_infer_image(
+    &prompt, 
+    image, 
+    strength, 
+    &model_config, 
+    &model_cache, 
+    prompt_cache, 
+    app, 
+    &app_data_root
+  ).await;
   
   if let Err(err) = result.as_deref() {
     error!("There was an error: {:?}", err);
@@ -44,29 +70,32 @@ pub fn infer_image(
   result
 }
 
-fn do_infer_image(
+async fn do_infer_image(
   prompt: &str,
   image: DynamicImage,
   strength: Option<u8>,
   config: &AppConfig,
   model_cache: &ModelCache,
-  prompt_cache: State<PromptCache>,
+  prompt_cache: State<'_, PromptCache>,
+  app: AppHandle,
+  app_data_root: &AppDataRoot,
 ) -> Result<String, String> {
-  println!("infer_image called; generating image with SDXL Turbo...");
-  
+
   let args = Args {
     image: &image,
     prompt: prompt.to_string(),
     uncond_prompt: NEGATIVE_PROMPT.to_string(),
-    //guidance_scale: Some(0.0),
     model_cache,
     configs: config,
     prompt_cache: &prompt_cache,
     i2i_strength: strength,
     cfg_scale: config.cfg_scale,
+    app: &app,
+    app_data_root,
+    use_flash_attn: true,
   };
 
-  match stable_diffusion_pipeline(args) {
+  match lcm_pipeline(args) {
     Ok(image) => {
       let mut bytes = Vec::with_capacity(1024*1024);
       

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

@@ -1,3 +1,4 @@
+pub mod download_models;
 pub mod flip_image;
 pub mod image_endpoint;
 pub mod test_counter;

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 {
+  Unet,
+  Vae,
+  Json,
+  ClipEncoder,
+}
+
+#[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,17 @@
 pub mod endpoints;
+pub mod events;
 pub mod ml;
 pub mod state;
+pub mod threads;
+pub mod transfer;
 
+use crate::endpoints::download_models::download_models;
+use crate::endpoints::image_endpoint::infer_image;
+use crate::endpoints::test_counter::test_counter;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
 use crate::state::app_config::AppConfig;
-use endpoints::image_endpoint::infer_image;
-use endpoints::test_counter::test_counter;
+use crate::threads::downloader_thread::downloader_thread;
 use tauri_plugin_log::Target;
 use tauri_plugin_log::TargetKind;
 
@@ -24,12 +29,10 @@ pub fn run() {
 
   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...");
 
@@ -46,14 +49,20 @@ pub fn run() {
       //      .build(),
       //  )?;
       //}
+      let app = app.handle().clone();
+
+      tauri::async_runtime::spawn(downloader_thread(app_data_root2, app));
+
       Ok(())
     })
     .manage(config)
     .manage(prompt_cache)
     .manage(model_cache)
+    .manage(app_data_root)
     .invoke_handler(tauri::generate_handler![
       test_counter, 
       infer_image,
+      download_models,
     ])
     .run(tauri::generate_context!())
     .expect("error while running tauri application");

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/downloads/download_all_models.rs

@@ -0,0 +1,15 @@
+use crate::ml::downloads::download_model_registry_file::download_model_registry_file;
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppDataRoot;
+
+pub async fn download_all_models(app_data_root: &AppDataRoot) -> anyhow::Result<()> {
+  let weights_dir = app_data_root.weights_dir();
+  
+  //download_model_registry_file(ModelType::ClipJson, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?;
+  //download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?;
+  
+  Ok(())
+}

crates/desktop/src-tauri/src/ml/downloads/download_model_registry_file.rs

@@ -0,0 +1,30 @@
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppWeightsDir;
+use log::info;
+use std::fs::File;
+use std::path::{Path, PathBuf};
+
+pub async fn download_model_registry_file(
+  model_type: ModelType,
+  weights_dir: &AppWeightsDir,
+) -> anyhow::Result<PathBuf> {
+  
+  let path = weights_dir.model_path(&model_type);
+
+  if path.exists() {
+    return Ok(path);
+  }
+
+  info!("downloading model: {:?} to {:?}", model_type, path);
+
+  let body = reqwest::get(model_type.get_download_url())
+    .await?
+    .bytes()
+    .await?;
+  
+  let _file = File::create(&path)?;
+  
+  std::fs::write(&path, body.as_ref())?;
+  
+  Ok(path)
+}

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

@@ -0,0 +1,2 @@
+pub mod download_model_registry_file;
+pub mod download_all_models;

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

@@ -3,4 +3,6 @@ pub mod model_cache;
 pub mod model_file;
 pub mod models;
 pub mod prompt_cache;
-pub mod stable_diffusion;
+pub mod stable_diffusion;
+pub mod model_type;
+pub mod downloads;

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

@@ -1,117 +1,60 @@
 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 candle_transformers::models::stable_diffusion::vae::AutoEncoderKL;
 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,
+  vae: RwLock<Option<Arc<AutoEncoderKL>>>,
+  unet2: RwLock<Option<Arc<UNetModel>>>,
 }
 
 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"))
-      }
+  pub fn new() -> Self {
+    Self { 
+      vae: RwLock::new(None),
+      unet2: RwLock::new(None),
     }
   }
   
-  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);
-          }
-        }
+  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(())
       }
     }
-    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)?))
-          }
-        }
+  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<UNetModel>) -> 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<UNetModel>>> {
+    match self.unet2.read() {
+      Err(err) => Err(anyhow!("{}", err)),
+      Ok(ref lock) => {
+        Ok(lock.as_ref().map(|v|v.clone()))
       }
-      Err(err) => Err(anyhow!("lock error: {:?}", err)),
     }
   }
   

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

@@ -4,80 +4,177 @@ extern crate accelerate_src;
 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,
-  Vae,
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum StableDiffusionVersion {
-  V1_5,
-  V1_5Inpaint,
-  V2_1,
-  V2Inpaint,
-  Xl,
-  XlInpaint,
-  Turbo,
+    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 => "runwayml/stable-diffusion-v1-5",
-      Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
-      Self::Turbo => "stabilityai/sdxl-turbo",
+    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,
-  ) -> 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)
-      }
+    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/src-tauri/src/ml/model_type.rs

@@ -0,0 +1,86 @@
+use crate::events::notification_event::NotificationModelType;
+use crate::state::app_dir::AppWeightsDir;
+use std::path::PathBuf;
+
+#[derive(Copy, Clone, Debug)]
+pub enum ModelType {
+  ClipJson,
+  SdxlTurboUnet,
+  SdxlTurboVae,
+  SdxlTurboClipEncoder,
+  SdxlTurboClipEncoder2,
+  SimianLuoLcmDreamshaperV7Unet,
+  LykonDreamshaper7Vae,
+  LykonDreamshaper7TextEncoderFp16,
+}
+
+impl ModelType {
+  pub fn get_name(&self) -> &'static str {
+    match self {
+      Self::ClipJson => "CLIP (JSON)",
+      Self::SdxlTurboUnet => "SDXL Turbo UNET",
+      Self::SdxlTurboVae => "SDXL Turbo VAE",
+      Self::SdxlTurboClipEncoder => "Clip Encoder #1",
+      Self::SdxlTurboClipEncoder2 => "Clip Encoder #2",
+      Self::SimianLuoLcmDreamshaperV7Unet => "SimianLuo LCM Dreamshaper v7 UNET",
+      Self::LykonDreamshaper7Vae => "Lykon Dreamshaper 7 VAE",
+      Self::LykonDreamshaper7TextEncoderFp16 => "Lykon Dreamshaper Text Encoder (fp16)",
+    }
+  }
+
+  pub fn get_notification_type(&self) -> NotificationModelType {
+    match self {
+      Self::ClipJson => NotificationModelType::Json,
+      Self::SdxlTurboUnet => NotificationModelType::Unet,
+      Self::SdxlTurboVae => NotificationModelType::Vae,
+      Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder,
+      Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder,
+      Self::SimianLuoLcmDreamshaperV7Unet => NotificationModelType::Unet,
+      Self::LykonDreamshaper7Vae => NotificationModelType::Vae,
+      Self::LykonDreamshaper7TextEncoderFp16 => NotificationModelType::ClipEncoder,
+    }
+  }
+
+  pub fn get_hf_id(&self) -> Option<&'static str> {
+    match self {
+      Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"),
+      Self::SdxlTurboUnet 
+      | Self::SdxlTurboVae
+      | Self::SdxlTurboClipEncoder
+      | Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"),
+      Self::SimianLuoLcmDreamshaperV7Unet => Some("SimianLuo/LCM_Dreamshaper_v7"),
+      Self::LykonDreamshaper7Vae
+      | Self::LykonDreamshaper7TextEncoderFp16 => Some("Lykon/dreamshaper-7"),
+    }
+  }
+
+  pub fn get_download_url(&self) -> &'static str {
+    match self {
+      Self::ClipJson => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/tokenizer.json",
+      Self::SdxlTurboUnet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.unet.safetensors",
+      Self::SdxlTurboVae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.vae.safetensors",
+      Self::SdxlTurboClipEncoder => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder.safetensors",
+      Self::SdxlTurboClipEncoder2 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder_2.safetensors",
+      Self::SimianLuoLcmDreamshaperV7Unet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/simianluo_lcm_dreamshaper_v7_unet.safetensors",
+      Self::LykonDreamshaper7Vae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_vae.safetensors",
+      Self::LykonDreamshaper7TextEncoderFp16 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors",
+    }
+  }
+
+  pub fn get_filename(&self) -> &'static str {
+    match self {
+      Self::ClipJson => "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
+      Self::SdxlTurboUnet => "diffusion_pytorch_model.unet.safetensors",
+      Self::SdxlTurboVae => "diffusion_pytorch_model.vae.safetensors",
+      Self::SdxlTurboClipEncoder => "clip_text_encoder.safetensors",
+      Self::SdxlTurboClipEncoder2 => "clip_text_encoder_2.safetensors",
+      Self::SimianLuoLcmDreamshaperV7Unet => "simianluo_lcm_dreamshaper_v7_unet.safetensors",
+      Self::LykonDreamshaper7Vae => "lykon_dreamshaper_7_vae.safetensors",
+      Self::LykonDreamshaper7TextEncoderFp16 => "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
+    }
+  }
+  
+  pub fn get_path(&self, weights_dir: &AppWeightsDir) -> PathBuf {
+    weights_dir.model_path(self)
+  }
+}

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

@@ -8,28 +8,23 @@ pub struct UNetModel {
 }
 
 impl UNetModel {
-  pub fn new<P: AsRef<Path>>(
-    sd_config: &StableDiffusionConfig,
-    unet_weights_file: P,
-    device: &Device,
-    dtype: DType,
-  ) -> anyhow::Result<Self> {
+  pub fn new<P: AsRef<Path>>(sd_config: &StableDiffusionConfig, unet_weights_file: P, device: &Device, dtype: DType) -> anyhow::Result<Self> {
     println!("building unet model... (1)");
-    
-    let model = sd_config.build_unet(
-      unet_weights_file, 
-      &device, 
-      4, 
-      false, 
-      dtype
-    )?;
-    
-    Ok(Self {
-      model,
-    })
+
+    let model = sd_config.build_unet(unet_weights_file, &device, 4, true, dtype)?;
+
+    Ok(Self { model })
   }
-  
+
+  pub fn new_with_inner(model: UNet2DConditionModel) -> Self {
+    Self { model }
+  }
+
   pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
     Ok(self.model.forward(&latent_model_input, timestep, &text_embeddings)?)
   }
+
+  pub fn forward_with_guidance(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor, guidance_scale_embedding: &Tensor) -> anyhow::Result<Tensor> {
+    Ok(self.model.forward_with_guidance(&latent_model_input, timestep, &text_embeddings, Some(guidance_scale_embedding))?)
+  }
 }

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

@@ -3,6 +3,7 @@ extern crate accelerate_src;
 #[cfg(feature = "mkl")]
 extern crate intel_mkl_src;
 
+use std::path::PathBuf;
 use candle_transformers::models::stable_diffusion;
 
 use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
@@ -10,27 +11,15 @@ use anyhow::Error as E;
 use candle_core::{DType, Device, Module, Tensor, D};
 use log::info;
 use tokenizers::Tokenizer;
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppWeightsDir;
 
 #[allow(clippy::too_many_arguments)]
-pub fn infer_clip_text_embeddings(
-  prompt: &str,
-  uncond_prompt: &str,
-  tokenizer: Option<String>,
-  clip_weights: Option<String>,
-  clip2_weights: Option<String>,
-  sd_version: StableDiffusionVersion,
-  sd_config: &stable_diffusion::StableDiffusionConfig,
-  use_f16: bool,
-  device: &Device,
-  dtype: DType,
-  use_guide_scale: bool,
-) -> anyhow::Result<Tensor> {
-  info!("Preparing text embeddings...");
+pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
+  info!("Preparing text embeddings... for {:?}", sd_version);
 
   let which = match sd_version {
-    StableDiffusionVersion::Xl
-    | StableDiffusionVersion::XlInpaint
-    | StableDiffusionVersion::Turbo => vec![true, false],
+    StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false],
     _ => vec![true],
   };
 
@@ -41,7 +30,7 @@ pub fn infer_clip_text_embeddings(
         &prompt,
         &uncond_prompt,
         None, // tokenizer
-        None, // clip_weights
+        clip_weights.clone(),
         None, // clip2_weights
         sd_version,
         &sd_config,
@@ -50,55 +39,45 @@ pub fn infer_clip_text_embeddings(
         dtype,
         use_guide_scale,
         *first,
+        weights_dir,
       )
     })
     .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)?;
+fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, first: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
+  info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first));
+
+  //let tokenizer_file = if first {
+  //  info!("ModelFile::Tokenizer");
+  //  ModelFile::Tokenizer
+  //} else {
+  //  info!("ModelFile::Tokenizer2");
+  //  ModelFile::Tokenizer2
+  //};
+
+  //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
+
+  let tokenizer = weights_dir.model_path(&ModelType::ClipJson);
+
+  info!("Loading Clip Tokenizer path: {:?}", tokenizer);
+
   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();
+  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
-        )
+    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)
@@ -106,37 +85,23 @@ fn do_infer_clip_text_embeddings(
   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_file = if first { ModelFile::Clip } else { ModelFile::Clip2 };
   let clip_weights = if first {
-    clip_weights_file.get(clip_weights, sd_version, use_f16)?
+    //clip_weights_file.get(clip_weights, sd_version, use_f16)?
+    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder)
   } 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()
+    //clip_weights_file.get(clip2_weights, sd_version, use_f16)?
+    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder2)
   };
+  let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
   let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
 
   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();
+    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
-            )
+      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)

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

@@ -0,0 +1,291 @@
+use std::path::PathBuf;
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use crate::events::notification_event::{NotificationEvent, NotificationModelType};
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
+use crate::ml::model_cache::ModelCache;
+use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use crate::ml::model_type::ModelType;
+use crate::ml::models::unet_model::UNetModel;
+use crate::ml::prompt_cache::PromptCache;
+use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
+use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
+use crate::state::app_config::AppConfig;
+use crate::state::app_dir::AppDataRoot;
+use anyhow::{anyhow, Error as E, Result};
+use candle_core::{DType, IndexOp, Tensor, D};
+use candle_nn::VarBuilder;
+use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
+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 image::DynamicImage;
+use log::info;
+use rand::Rng;
+use tauri::{AppHandle, Emitter};
+
+pub struct Args<'a> {
+  pub image: &'a DynamicImage,
+  pub prompt: String,
+  pub uncond_prompt: String,
+  pub cfg_scale: Option<f64>,
+  pub i2i_strength: Option<u8>,
+  pub configs: &'a AppConfig,
+  pub model_cache: &'a ModelCache,
+  pub prompt_cache: &'a PromptCache,
+  pub app_data_root: &'a AppDataRoot,
+  pub app: &'a AppHandle,
+  pub use_flash_attn: bool,
+}
+
+pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
+  let Args { prompt, uncond_prompt, cfg_scale, i2i_strength, configs, model_cache, prompt_cache, app, image, app_data_root, use_flash_attn } = args;
+
+  let weights_dir = app_data_root.weights_dir();
+
+  let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0;
+  let use_f16 = true;
+
+  //println!("Starting image generation with the following configuration:");
+  //println!("  Model: {:?}", configs.sd_version);
+  //println!("  Prompt: {}", prompt);
+  //println!("  Steps: {}", configs.scheduler_steps);
+  //println!("  Device: {:?}", configs.device);
+
+  //println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
+
+  // Use LCM Scheduler instead of Euler Ancestral for better speed and quality
+  let mut scheduler = LCMScheduler::new(configs.scheduler_steps, img2img_strength, candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default())?;
+
+  let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
+  configs.device.set_seed(seed)?;
+
+  //info!("Using seed: {}", seed);
+
+  let guidance_scale = match cfg_scale {
+    Some(guidance_scale) => guidance_scale,
+    None => match configs.sd_version {
+      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;
+
+  //info!("Using guide scale: {}", use_guide_scale);
+
+  //info!("Checking if prompt is cached");
+  let maybe_cached = prompt_cache.get_copy(&prompt)?;
+
+  let mut text_embeddings = if let Some(tensor) = maybe_cached {
+    tensor
+  } else {
+    info!("Loading clip weights...");
+    //println!("Loading clip weights for {:?} with model {:?}", configs.sd_version, configs.sd_version.repo());
+    //let api = configs.hf_api.clone();
+
+    //let clip_path = api.model(configs.sd_version.repo().to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
+
+    //println!(">>>> CLIP REPO = {:?}", configs.sd_version.repo());
+    //println!(">>>> CLIP PATH = {:?}", clip_path);
+    //println!(">>>> USE_FP16 = {:?}", use_f16);
+    //let clip_weights = Some(clip_path.to_string_lossy().to_string());
+
+    let clip_weights = weights_dir.model_path(&ModelType::LykonDreamshaper7TextEncoderFp16);
+    let clip_weights = Some(clip_weights.to_string_lossy().to_string());
+
+    info!("Prompt is NOT cached! Calculating embedding...");
+    let tensor = infer_clip_text_embeddings(
+      &prompt,
+      &uncond_prompt,
+      None, // tokenizer
+      clip_weights,
+      None, // clip2_weights
+      configs.sd_version,
+      &configs.sd_config,
+      use_f16, // use_f16
+      &configs.device,
+      configs.dtype,
+      use_guide_scale,
+      weights_dir,
+    )?;
+    prompt_cache.store_copy(&prompt, &tensor)?;
+    tensor
+  };
+
+  println!("Text embeddings shape: {:?}", text_embeddings.shape());
+
+  println!("Loading input image into tensor...");
+
+  let input_image = dynamic_image_to_tensor(image, &configs.device, configs.dtype)?;
+
+  println!("Reference image shape: {:?}", input_image.shape());
+
+  let vae_scale = get_vae_scale(configs.sd_version);
+
+  let maybe_vae = model_cache.get_vae()?;
+
+  let vae = match maybe_vae {
+    Some(vae) => vae,
+    None => {
+      //let repo = configs.sd_version.repo();
+      //println!(">>>>> VAE REPO = {:?}", repo);
+      //println!("Building VAE model from : {:?} ... (3)", repo);
+      //let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
+
+      let vae_file = weights_dir.model_path(&ModelType::LykonDreamshaper7Vae);
+
+      println!("Building VAE model from file {:?}...", &vae_file);
+
+      let vae = configs.sd_config.build_vae(vae_file, &configs.device, configs.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 = weights_dir.model_path(&ModelType::SimianLuoLcmDreamshaperV7Unet);
+
+      let in_channels = match configs.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_weights], configs.dtype, &configs.device)? };
+
+      let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
+
+      let unet = UNetModel::new_with_inner(unet);
+
+      let unet = Arc::new(unet);
+
+      model_cache.set_unet(unet.clone())?;
+
+      unet
+    },
+  };
+
+  let init_latent_dist: DiagonalGaussianDistribution = vae.encode(&input_image)?;
+
+  // Parse and calculate the img2img strength
+  let img2img_strength = match i2i_strength {
+    None => 0.75f64,
+    Some(strength) => (strength as f64) / 100.0f64,
+  };
+
+  // Create LCM guidance scale embeddings
+  let embedding_dim = 1280;
+  let guidance_scale_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &configs.device, configs.dtype)?;
+
+  // Generate latents from input image using the LCM approach
+  println!("Generating latents from input image...");
+  let init_latents = (init_latent_dist.sample()? * vae_scale)?;
+  println!("Initial latents shape: {:?}", init_latents.shape());
+
+  // Calculate timesteps for LCM with proper img2img handling
+  let timesteps = LCMScheduler::get_timesteps_for_steps(configs.scheduler_steps, img2img_strength);
+  let t_start = timesteps[0]; // First denoising step
+
+  // Add noise at the appropriate timestep for img2img
+  println!("Adding noise to latents for t_start={}", t_start);
+  let noise = init_latents.randn_like(0f64, 1f64)?;
+  let latents = scheduler.add_noise(&init_latents, noise, t_start)?;
+
+  println!("Latents initialized successfully");
+  println!("Initial noise shape: {:?}", latents.shape());
+
+  println!("Starting diffusion process...");
+
+  let mut latents = latents;
+  let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
+
+  for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+    println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
+
+    let latent_model_input = if use_guide_scale { Tensor::cat(&[&latents, &latents], 0)? } else { latents.clone() };
+
+    let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);
+
+    // Use the guidance scale embedding in the UNet inference with the proper method name
+    let mut noise_pred = match unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, &guidance_scale_embedding) {
+      Ok(pred) => pred,
+      Err(e) => {
+        println!("UNet inference failed with error: {}", e);
+        return Err(anyhow::anyhow!("UNet inference failed: {}", e));
+      },
+    };
+
+    if use_guide_scale {
+      let chunks = noise_pred.chunk(2, 0)?;
+      let (noise_pred_uncond, noise_pred_text) = (&chunks[0], &chunks[1]);
+      noise_pred = (noise_pred_uncond + (guidance_scale * (noise_pred_text - noise_pred_uncond)?)?)?;
+    }
+
+    // Apply the LCM scheduler step
+    latents = scheduler.step(&noise_pred, timestep, &latents, timestep_index, configs.scheduler_steps)?;
+  }
+
+  println!("Diffusion process completed, decoding image...");
+  let image = vae.decode(&(latents / vae_scale)?)?;
+
+  println!("VAE decode completed");
+
+  println!("Scaling image back");
+  let image = ((image / 2.)? + 0.5)?;
+
+  println!("Normalized image values");
+  let image = (image.clamp(0f32, 1.)? * 255.)?;
+
+  println!("Convert to int8");
+  let image = image.to_dtype(DType::U8)?;
+
+  println!("Converted to 8-bit format");
+
+  let image = tensor_to_image_buffer(&image.i(0)?)?;
+
+  Ok(image)
+}

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

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

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

@@ -1,19 +1,24 @@
-use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
+use crate::events::notification_event::{NotificationModelType, NotificationEvent};
 use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
 use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
 use crate::ml::model_cache::ModelCache;
 use crate::ml::model_file::StableDiffusionVersion;
+use crate::ml::model_type::ModelType;
+use crate::ml::models::unet_model::UNetModel;
 use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
 use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
 use crate::state::app_config::AppConfig;
 use anyhow::{anyhow, Error as E, Result};
 use candle_core::{DType, IndexOp, Tensor, D};
-use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
 use image::DynamicImage;
 use log::info;
 use rand::Rng;
+use tauri::{AppHandle, Emitter};
+use crate::state::app_dir::AppDataRoot;
 
 pub struct Args<'a> {
     pub image: &'a DynamicImage,
@@ -24,31 +29,48 @@ pub struct Args<'a> {
     pub configs: &'a AppConfig,
     pub model_cache: &'a ModelCache,
     pub prompt_cache: &'a PromptCache,
+    pub app_data_root: &'a AppDataRoot,
+    pub app: &'a AppHandle,
 }
 
 pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
+    let Args { 
+        prompt, 
+        uncond_prompt, 
+        cfg_scale, 
+        i2i_strength, 
+        configs, 
+        model_cache, 
+        prompt_cache, 
+        app, 
+        image,
+        app_data_root,
+    } = args;
+    
+    let weights_dir = app_data_root.weights_dir();
+    
     println!("Starting image generation with the following configuration:");
-    println!("  Model: {:?}", args.configs.sd_version);
-    println!("  Prompt: {}", args.prompt);
-    println!("  Steps: {}", args.configs.scheduler_steps);
-    println!("  Device: {:?}", args.configs.device);
+    println!("  Model: {:?}", configs.sd_version);
+    println!("  Prompt: {}", prompt);
+    println!("  Steps: {}", configs.scheduler_steps);
+    println!("  Device: {:?}", configs.device);
 
-    println!("Model dimensions: {}x{}", args.configs.sd_config.width, args.configs.sd_config.height);
+    println!("Model dimensions: {}x{}", configs.sd_config.width, configs.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 = configs.sd_config.build_scheduler(
+        configs.scheduler_steps)?;
 
-    let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
-    args.configs.device.set_seed(seed)?;
+    let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
+    configs.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 configs.sd_version {
             StableDiffusionVersion::V1_5
             | StableDiffusionVersion::V2_1
             | StableDiffusionVersion::Xl => 7.5,
@@ -59,29 +81,31 @@ 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,
+            &prompt,
+            &uncond_prompt,
             None, // tokenizer
             None, // clip_weights
             None, // clip2_weights
-            args.configs.sd_version,
-            &args.configs.sd_config,
+            configs.sd_version,
+            &configs.sd_config,
             false, // use_f16
-            &args.configs.device,
-            args.configs.dtype,
+            &configs.device,
+            configs.dtype,
             use_guide_scale,
+            weights_dir,
         )?;
-        args.prompt_cache.store_copy(&args.prompt, &tensor)?;
+        prompt_cache.store_copy(&prompt, &tensor)?;
         tensor
     };
 
@@ -90,36 +114,111 @@ 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)?;
+        image,
+        &configs.device,
+        configs.dtype)?;
 
     println!("Reference image shape: {:?}", input_image.shape());
 
-    let vae_scale = get_vae_scale(args.configs.sd_version);
+    let vae_scale = get_vae_scale(configs.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 mut notify_download_complete = false;
+            //if !vae_file.exists() {
+            if true {
+                notify_download_complete = true;
+                app.emit("notification", NotificationEvent::ModelDownloadStarted {
+                    model_name: "sdxl-turbo-vae",
+                    model_type: NotificationModelType::Vae,
+                })?;
+            }
+            
+            let vae_file = weights_dir.model_path(&ModelType::SdxlTurboVae);
+            
+            let vae = configs
+              .sd_config
+              .build_vae(vae_file, &configs.device, configs.dtype)?;
+
+            let vae = Arc::new(vae);
+            
+            model_cache.set_vae(vae.clone())?;
+            
+            if notify_download_complete {
+                app.emit("notification", NotificationEvent::ModelDownloadComplete {
+                    model_name: "sdxl-turbo-vae",
+                    model_type: NotificationModelType::Vae,
+                })?;
+            }
+            
+            vae
+        }
+    };
+
+    let maybe_unet = model_cache.get_unet()?;
+
+    let unet = match maybe_unet {
+        Some(unet) => unet,
+        None => {
+            info!("No unet found in cache; loading...");
+
+            let mut notify_download_complete = false;
+            //if !unet_file.exists() {
+            if true {
+                notify_download_complete = true;
+                app.emit("notification", NotificationEvent::ModelDownloadStarted {
+                    model_name: "sdxl-turbo-unet",
+                    model_type: NotificationModelType::Unet,
+                })?;
+            }
+            
+            let unet_file = weights_dir.model_path(&ModelType::SdxlTurboUnet);
+
+            let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
+              .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
+            
+            let unet = Arc::new(unet);
+
+            model_cache.set_unet(unet.clone())?;
+
+            if notify_download_complete {
+                app.emit("notification", NotificationEvent::ModelDownloadComplete {
+                    model_name: "sdxl-turbo-unet",
+                    model_type: NotificationModelType::Unet,
+                })?;
+            }
+
+            unet
+        }
+    };
+
+    let init_latent_dist : DiagonalGaussianDistribution = vae.encode(&input_image)?;
 
     // 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 = configs.scheduler_steps - (configs.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, configs.scheduler_steps, img2img_strength);
         start
     };
 
@@ -158,7 +257,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
         let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
 
-        let mut noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) {
+        let mut noise_pred = match unet.inference(&latent_model_input, timestep as f64, &text_embeddings) {
             Ok(pred) => pred,
             Err(e) => {
                 println!("UNet inference failed with error: {}", e);
@@ -176,7 +275,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/src-tauri/src/state/app_config.rs

@@ -1,18 +1,19 @@
 use crate::ml::model_file::StableDiffusionVersion;
+use crate::state::app_dir::AppDataRoot;
+use crate::state::os_platform::OsPlatform;
 use crate::state::yaml_config::YamlConfig;
 use candle_core::{DType, Device};
 use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
 use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
 use hf_hub::api::sync::Api;
 use log::info;
+use std::path::PathBuf;
 
-const CONFIG_FILENAME : &str = "app_config.yaml";
-
-const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
-const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
+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::Turbo;
+const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::V1_5;
 
 /// A central place to configure app-wide details.
 pub struct AppConfig {
@@ -37,13 +38,15 @@ pub struct AppConfig {
   
   /// Probably shouldn't be here.
   pub hf_api: Api,
+  
+  /// Root location for application data.
+  pub app_data_root: AppDataRoot,
 }
 
 impl AppConfig {
   
   pub fn init() -> anyhow::Result<Self> {
-    println!("Configuration filename: {}", CONFIG_FILENAME);
-    let yaml_configs = load_yaml_configs();
+    let yaml_configs = YamlConfig::load_from_config_file_recursive();
     
     println!("Configs: {:?}", yaml_configs);
     
@@ -52,18 +55,42 @@ impl AppConfig {
           info!("CUDA not available ({}), falling back to CPU", e);
           Device::Cpu
         });
+
     
     let sd_config = match SD_VERSION {
       StableDiffusionVersion::Turbo => {
         StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
       }
-      _ => {
+      StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
+        StableDiffusionConfig::v1_5(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => {
         StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
       }
+      StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => {
+        StableDiffusionConfig::sdxl(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
     };
     
     let hf_api = Api::new()?;
     
+    let os_platform = OsPlatform::get();
+
+    let default_data_root = match os_platform {
+      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(),
+    };
+    
+    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,
@@ -71,26 +98,12 @@ impl AppConfig {
       sd_config,
       image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
       image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
-      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(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,
     })
   }
 }
-
-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,186 @@
+use crate::ml::model_type::ModelType;
+use crate::state::expanduser::expanduser;
+use crate::state::os_platform::OsPlatform;
+use anyhow::anyhow;
+use directories::UserDirs;
+use std::path::{Path, PathBuf};
+use tempdir::TempDir;
+use tempfile::{Builder, NamedTempFile};
+const DEFAULT_DATA_DIR : &str = "artcraft";
+const ASSETS_SUBDIRECTORY : &str = "assets";
+const WEIGHTS_SUBDIRECTORY : &str = "weights";
+
+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 model_path(&self, model_registry: &ModelType) -> PathBuf {
+    self.path.join(model_registry.get_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/threads/downloader_thread.rs

@@ -0,0 +1,172 @@
+use crate::events::notification_event::{NotificationEvent, NotificationModelType};
+use crate::ml::model_type::ModelType;
+use crate::state::app_dir::AppDataRoot;
+use crate::transfer::download::{download_async, ProgressUpdate};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use log::{debug, error, info};
+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();
+
+  // TODO: Automatic enqueue of known important models + enqueue new models on-demand
+  download_queue.push_back(ModelType::ClipJson);
+  //download_queue.push_back(ModelType::SdxlTurboUnet);
+  //download_queue.push_back(ModelType::SdxlTurboVae);
+  download_queue.push_back(ModelType::SdxlTurboClipEncoder); // TODO(bt): Why is this still needed?
+  //download_queue.push_back(ModelType::SdxlTurboClipEncoder2);
+  download_queue.push_back(ModelType::SimianLuoLcmDreamshaperV7Unet);
+  download_queue.push_back(ModelType::LykonDreamshaper7Vae);
+  download_queue.push_back(ModelType::LykonDreamshaper7TextEncoderFp16);
+
+  let mut handles = FuturesUnordered::new();
+
+  while let Some(model) = download_queue.pop_front() {
+    let url = model.get_download_url().to_string();
+
+    let filename = app_data_root.weights_dir().model_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 result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
+      model_name: model.get_name(),
+      model_type: model.get_notification_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.get_name().to_string(),
+          model_type: model.get_notification_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/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/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/.cargo/config.toml

@@ -0,0 +1,8 @@
+[build]
+rustflags = ["-C", "target-cpu=native"]
+
+[target.wasm32-unknown-unknown]
+rustflags = ["-C", "target-feature=+simd128"]
+
+[target.x86_64-apple-darwin]
+rustflags = ["-C", "target-feature=-avx,-avx2"]

crates/vendor/candle/.github/dependabot.yml

@@ -0,0 +1,7 @@
+version: 2
+updates:
+  - package-ecosystem: "cargo"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 5

crates/vendor/candle/.github/workflows/book-cd.yml

@@ -0,0 +1,40 @@
+name: Deploy Rust book
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write  # To push a branch 
+      pull-requests: write  # To create a PR from that branch
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Install latest mdbook
+      run: |
+        tag=$(curl 'https://api.github.com/repos/rust-lang/mdbook/releases/latest' | jq -r '.tag_name')
+        url="https://github.com/rust-lang/mdbook/releases/download/${tag}/mdbook-${tag}-x86_64-unknown-linux-gnu.tar.gz"
+        mkdir mdbook
+        curl -sSL $url | tar -xz --directory=./mdbook
+        echo `pwd`/mdbook >> $GITHUB_PATH
+    - name: Deploy GitHub Pages
+      run: |
+        # This assumes your book is in the root of your repository.
+        # Just add a `cd` here if you need to change to another directory.
+        cd candle-book
+        mdbook build
+        git worktree add gh-pages
+        git config user.name "Deploy from CI"
+        git config user.email ""
+        cd gh-pages
+        # Delete the ref to avoid keeping history.
+        git update-ref -d refs/heads/gh-pages
+        rm -rf *
+        mv ../book/* .
+        git add .
+        git commit -m "Deploy $GITHUB_SHA to gh-pages"
+        git push --force --set-upstream origin gh-pages

crates/vendor/candle/.github/workflows/book.yml

@@ -0,0 +1,29 @@
+name: CI
+on: 
+  pull_request:
+
+jobs:
+  test:
+    name: Test candle-book
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write  # To push a branch 
+      pull-requests: write  # To create a PR from that branch
+    steps:
+    - uses: actions/checkout@master
+    - name: Install Rust
+      run: |
+        rustup set profile minimal
+        rustup toolchain install stable
+        rustup default stable
+    - name: Install latest mdbook
+      run: |
+        tag=$(curl 'https://api.github.com/repos/rust-lang/mdbook/releases/latest' | jq -r '.tag_name')
+        url="https://github.com/rust-lang/mdbook/releases/download/${tag}/mdbook-${tag}-x86_64-unknown-linux-gnu.tar.gz"
+        mkdir bin
+        curl -sSL $url | tar -xz --directory=bin
+        echo "$(pwd)/bin" >> $GITHUB_PATH
+    - name: Run tests
+      run: cd candle-book && cargo build && mdbook test -L ../target/debug/deps/
+
+

crates/vendor/candle/.github/workflows/ci_cuda.yaml

@@ -0,0 +1,34 @@
+name: CI / cuda
+
+on:
+  workflow_dispatch:
+  pull_request:
+
+jobs:
+  test-cuda:
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.job }}-${{ github.head_ref || github.run_id }}
+      cancel-in-progress: true
+    runs-on:
+      group: aws-g4dn-2xlarge
+    container:
+      image: nvidia/cuda:12.3.1-devel-ubuntu22.04
+      options: --gpus 0 
+    if: ${{ github.event.pull_request.head.repo.full_name == github.event.pull_request.base.repo.full_name }}
+    permissions:
+      contents: write
+      packages: write
+      # This is used to complete the identity challenge
+      # with sigstore/fulcio when running outside of PRs.
+      id-token: write
+      security-events: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+      - name: Install dependencies
+        run: apt-get update && apt install curl build-essential libssl-dev protobuf-compiler pkg-config -y
+      - name: Install Rust Stable
+        uses: actions-rust-lang/setup-rust-toolchain@v1
+      - uses: Swatinem/rust-cache@v2
+      - name: Test (cuda)
+        run: cargo test --features cuda

crates/vendor/candle/.github/workflows/python.yml

@@ -0,0 +1,68 @@
+name: PyO3-CI
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+    paths:
+      - candle-pyo3/**
+  pull_request:
+    paths:
+      - candle-pyo3/**
+
+jobs:
+  build_and_test:
+    name: Check everything builds & tests
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest] # For now, only test on Linux
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Install Rust
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+
+      - name: Install Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.11
+          architecture: "x64"
+
+      - name: Cache Cargo Registry
+        uses: actions/cache@v1
+        with:
+          path: ~/.cargo/registry
+          key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
+
+      - name: Install Protoc
+        uses: arduino/setup-protoc@v2
+        with:
+            version: "25.0"
+            repo-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Install
+        working-directory: ./candle-pyo3
+        run: |
+          python -m venv .env
+          source .env/bin/activate
+          pip install -U pip
+          pip install pytest maturin black
+          python -m maturin develop -r --features onnx
+
+      - name: Check style
+        working-directory: ./candle-pyo3
+        run: |
+          source .env/bin/activate
+          python stub.py --check
+          black --check .
+
+      - name: Run tests
+        working-directory: ./candle-pyo3
+        run: |
+          source .env/bin/activate
+          python -m pytest -s -v tests

crates/vendor/candle/.github/workflows/rust-ci.yml

@@ -0,0 +1,87 @@
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+name: Continuous integration
+
+jobs:
+  check:
+    name: Check
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macOS-latest]
+        rust: [stable]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: ${{ matrix.rust }}
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: check
+          args: --workspace
+
+  test:
+    name: Test Suite
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macOS-latest]
+        rust: [stable]
+    steps:
+      - name: Delete huge unnecessary tools folder
+        if: runner.os == 'Linux'
+        run: rm -rf /opt/hostedtoolcache
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: ${{ matrix.rust }}
+          override: true
+      - uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --workspace
+
+  fmt:
+    name: Rustfmt
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          override: true
+      - run: rustup component add rustfmt
+      - uses: actions-rs/cargo@v1
+        with:
+          command: fmt
+          args: --all -- --check
+
+  clippy:
+    name: Clippy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          override: true
+      - run: rustup component add clippy
+      - uses: actions-rs/cargo@v1
+        with:
+          command: clippy
+          args: --workspace --tests --examples -- -D warnings

crates/vendor/candle/.github/workflows/trufflehog.yml

@@ -0,0 +1,15 @@
+on:
+  push:
+
+name: Secret Leaks
+
+jobs:
+  trufflehog:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: Secret Scanning
+      uses: trufflesecurity/trufflehog@main

crates/vendor/candle/.gitignore

@@ -0,0 +1,48 @@
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+data/
+dist/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# editor config
+.helix
+.vscode
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb
+
+*tokenizer*.json
+*.npz
+
+perf.data
+flamegraph.svg
+*.dylib
+*.so
+*.swp
+*.swo
+trace-*.json
+
+candle-wasm-examples/*/build
+candle-wasm-examples/*/*.bin
+candle-wasm-examples/*/*.jpeg
+candle-wasm-examples/*/audios/*.wav
+candle-wasm-examples/**/*.safetensors
+candle-wasm-examples/**/*.gguf
+candle-wasm-examples/*/package-lock.json
+candle-wasm-examples/**/config*.json
+.DS_Store
+.idea/*
+__pycache__
+out.safetensors
+out.wav
+bria.mp3
+bria.safetensors
+bria.wav

crates/vendor/candle/.gitmodules

@@ -0,0 +1,3 @@
+[submodule "candle-examples/examples/flash-attn/cutlass"]
+	path = candle-flash-attn/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

crates/vendor/candle/.pre-commit-config.yaml

@@ -0,0 +1,15 @@
+repos:
+  - repo: https://github.com/Narsil/pre-commit-rust
+    rev: 2eed6366172ef2a5186e8785ec0e67243d7d73d0
+    hooks:
+      - id: fmt
+        name: "Rust (fmt)"
+      - id: clippy
+        name: "Rust (clippy)"
+        args:
+          [
+            "--tests",
+            "--examples",
+            "--",
+            "-Dwarnings",
+          ]

crates/vendor/candle/CHANGELOG.md

@@ -0,0 +1,113 @@
+# Changelog
+This documents the main changes to the `candle` crate.
+
+## v0.3.1 - Unreleased
+
+### Added
+
+### Modified
+
+## v0.3.0 - 2023-10-01
+
+### Added
+
+- Added the Mistral 7b v0.1 model
+  [983](https://github.com/huggingface/candle/pull/983).
+- Quantized version of the Mistral model
+  [1009](https://github.com/huggingface/candle/pull/1009).
+- Add the gelu-erf op and activation function
+  [969](https://github.com/huggingface/candle/pull/969).
+- Add the mixformer/phi-v1.5 model
+  [930](https://github.com/huggingface/candle/pull/930).
+- Add the sclice-scatter op
+  [927](https://github.com/huggingface/candle/pull/927).
+- Add the Wuerstchen diffusion model
+  [911](https://github.com/huggingface/candle/pull/911).
+
+### Modified
+
+- Support for simd128 intrinsics in some quantized vecdots
+  [982](https://github.com/huggingface/candle/pull/982).
+- Optimize the index-select cuda kernel
+  [976](https://github.com/huggingface/candle/pull/976).
+- Self-contained safetensor wrappers
+  [946](https://github.com/huggingface/candle/pull/946).
+
+## v0.2.2 - 2023-09-18
+
+### Added
+- Support for `top_p` sampling
+  [819](https://github.com/huggingface/candle/pull/819).
+- T5 model including decoding
+  [864](https://github.com/huggingface/candle/pull/864).
+- 1-d upsampling
+  [839](https://github.com/huggingface/candle/pull/839).
+
+### Modified
+- Bugfix for conv2d
+  [820](https://github.com/huggingface/candle/pull/820).
+- Support tensor based indexing using `.i`
+  [842](https://github.com/huggingface/candle/pull/842).
+
+## v0.2.1 - 2023-09-11
+
+### Added
+- Add some RNNs (GRU and LSTM) in `candle-nn`
+  [674](https://github.com/huggingface/candle/pull/674),
+  [688](https://github.com/huggingface/candle/pull/688).
+- gguf v2 support
+  [725](https://github.com/huggingface/candle/pull/725).
+- Quantized llama example in Python using the pyo3 api
+  [716](https://github.com/huggingface/candle/pull/716).
+- `candle-nn` layer for conv2d-transposed
+  [760](https://github.com/huggingface/candle/pull/760).
+- Add the Segment-Anything Model (SAM) as an example
+  [773](https://github.com/huggingface/candle/pull/773).
+- TinyViT backbone for the segment anything example
+  [787](https://github.com/huggingface/candle/pull/787).
+- Shape with holes support
+  [770](https://github.com/huggingface/candle/pull/770).
+
+### Modified
+- Dilations are now supported in conv-transpose2d.
+  [671](https://github.com/huggingface/candle/pull/671).
+- Interactive mode for the quantized model
+  [690](https://github.com/huggingface/candle/pull/690).
+- Faster softmax operation
+  [747](https://github.com/huggingface/candle/pull/747).
+- Faster convolution operations on CPU and CUDA via im2col
+  [802](https://github.com/huggingface/candle/pull/802).
+- Moving some models to a more central location
+  [796](https://github.com/huggingface/candle/pull/796).
+
+## v0.2.0 - 2023-08-30
+
+### Added
+- Add the powf op
+  [664](https://github.com/huggingface/candle/pull/664).
+- Stable Diffusion XL support
+  [647](https://github.com/huggingface/candle/pull/647).
+- Add the conv-transpose2d op
+  [635](https://github.com/huggingface/candle/pull/635).
+- Refactor the VarBuilder api
+  [627](https://github.com/huggingface/candle/pull/627).
+- Add some quantization command
+  [625](https://github.com/huggingface/candle/pull/625).
+- Support more quantized types, e.g. Q2K, Q4K, Q5K...
+  [586](https://github.com/huggingface/candle/pull/586).
+- Add pose estimation to the yolo example
+  [589](https://github.com/huggingface/candle/pull/589).
+- Api to write GGUF files
+  [585](https://github.com/huggingface/candle/pull/585).
+- Support more quantization types
+  [580](https://github.com/huggingface/candle/pull/580).
+- Add EfficientNet as an example Computer Vision model
+  [572](https://github.com/huggingface/candle/pull/572).
+- Add a group parameter to convolutions
+  [566](https://github.com/huggingface/candle/pull/566).
+- New dtype: int64
+  [563](https://github.com/huggingface/candle/pull/563).
+- Handling of the GGUF file format.
+  [559](https://github.com/huggingface/candle/pull/559).
+
+## v0.1.2 - 2023-08-21

crates/vendor/candle/Cargo.toml

@@ -0,0 +1,82 @@
+[workspace]
+members = [
+    "candle-core",
+    "candle-datasets",
+    "candle-examples",
+    "candle-book",
+    "candle-nn",
+    "candle-pyo3",
+    "candle-transformers",
+    "candle-wasm-examples/*",
+    "candle-wasm-tests",
+    "tensor-tools",
+]
+exclude = [
+   "candle-flash-attn",
+   "candle-kernels",
+   "candle-metal-kernels",
+   "candle-onnx",
+]
+resolver = "2"
+
+[workspace.package]
+version = "0.8.3"
+edition = "2021"
+description = "Minimalist ML framework."
+repository = "https://github.com/huggingface/candle"
+keywords = ["blas", "tensor", "machine-learning"]
+categories = ["science"]
+license = "MIT OR Apache-2.0"
+
+[workspace.dependencies]
+ab_glyph = "0.2.23"
+accelerate-src = { version = "0.3.2" }
+anyhow = { version = "1", features = ["backtrace"] }
+byteorder = "1.4.3"
+candle = { path = "./candle-core", package = "candle-core", version = "0.8.3" }
+candle-datasets = { path = "./candle-datasets", version = "0.8.3" }
+candle-flash-attn = { path = "./candle-flash-attn", version = "0.8.3" }
+candle-kernels = { path = "./candle-kernels", version = "0.8.3" }
+candle-metal-kernels = { path = "./candle-metal-kernels", version = "0.8.3" }
+candle-nn = { path = "./candle-nn", version = "0.8.3" }
+candle-onnx = { path = "./candle-onnx", version = "0.8.3" }
+candle-transformers = { path = "./candle-transformers", version = "0.8.3" }
+clap = { version = "4.2.4", features = ["derive"] }
+criterion = { version = "0.5.1", default-features=false }
+cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false }
+fancy-regex = "0.13.0"
+gemm = { version = "0.17.0", features = ["wasm-simd128-enable"] }
+hf-hub = "0.4.1"
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
+hound = "3.5.1"
+image = { version = "0.25.2", default-features = false, features = ["jpeg", "png"] }
+imageproc = { version = "0.24.0", default-features = false }
+intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"] }
+libc = { version = "0.2.147" }
+log = "0.4"
+memmap2 = { version = "0.9.3", features = ["stable_deref_trait"] }
+num_cpus = "1.15.0"
+num-traits = "0.2.15"
+parquet = { version = "51.0.0" }
+rand = "0.8.5"
+rand_distr = "0.4.3"
+rayon = "1.7.0"
+safetensors = "0.4.1"
+serde = { version = "1.0.171", features = ["derive"] }
+serde_plain = "1.0.2"
+serde_json = "1.0.99"
+thiserror = "1"
+tokenizers = { version = "0.19.1", default-features = false }
+tracing = "0.1.37"
+tracing-chrome = "0.7.1"
+tracing-subscriber = "0.3.7"
+ug = "0.1.0"
+ug-cuda = "0.1.0"
+ug-metal = "0.1.0"
+yoke = { version = "0.7.2", features = ["derive"] }
+zip = { version = "1.1.1", default-features = false }
+metal = { version = "0.27.0", features = ["mps"]}
+
+[profile.release-with-debug]
+inherits = "release"
+debug = true

crates/vendor/candle/LICENSE-APACHE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

crates/vendor/candle/LICENSE-MIT

@@ -0,0 +1,23 @@
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.

crates/vendor/candle/Makefile

@@ -0,0 +1,16 @@
+.PHONY: clean-ptx clean test
+
+clean-ptx:
+	find target -name "*.ptx" -type f -delete
+	echo "" > candle-kernels/src/lib.rs
+	touch candle-kernels/build.rs
+	touch candle-examples/build.rs
+	touch candle-flash-attn/build.rs
+
+clean:
+	cargo clean
+
+test:
+	cargo test
+
+all: test

crates/vendor/candle/README.md

@@ -0,0 +1,425 @@
+# candle
+[![discord server](https://dcbadge.vercel.app/api/server/hugging-face-879548962464493619)](https://discord.gg/hugging-face-879548962464493619)
+[![Latest version](https://img.shields.io/crates/v/candle-core.svg)](https://crates.io/crates/candle-core)
+[![Documentation](https://docs.rs/candle-core/badge.svg)](https://docs.rs/candle-core)
+[![License](https://img.shields.io/github/license/base-org/node?color=blue)](https://github.com/huggingface/candle/blob/main/LICENSE-MIT)
+[![License](https://img.shields.io/badge/license-Apache%202.0-blue?style=flat-square)](https://github.com/huggingface/candle/blob/main/LICENSE-APACHE)
+
+Candle is a minimalist ML framework for Rust with a focus on performance (including GPU support) 
+and ease of use. Try our online demos: 
+[whisper](https://huggingface.co/spaces/lmz/candle-whisper),
+[LLaMA2](https://huggingface.co/spaces/lmz/candle-llama2),
+[T5](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm),
+[yolo](https://huggingface.co/spaces/lmz/candle-yolo),
+[Segment
+Anything](https://huggingface.co/spaces/radames/candle-segment-anything-wasm).
+
+## Get started
+
+Make sure that you have [`candle-core`](https://github.com/huggingface/candle/tree/main/candle-core) correctly installed as described in [**Installation**](https://huggingface.github.io/candle/guide/installation.html).
+
+Let's see how to run a simple matrix multiplication.
+Write the following to your `myapp/src/main.rs` file:
+```rust
+use candle_core::{Device, Tensor};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let device = Device::Cpu;
+
+    let a = Tensor::randn(0f32, 1., (2, 3), &device)?;
+    let b = Tensor::randn(0f32, 1., (3, 4), &device)?;
+
+    let c = a.matmul(&b)?;
+    println!("{c}");
+    Ok(())
+}
+```
+
+`cargo run` should display a tensor of shape `Tensor[[2, 4], f32]`.
+
+
+Having installed `candle` with Cuda support, simply define the `device` to be on GPU:
+
+```diff
+- let device = Device::Cpu;
++ let device = Device::new_cuda(0)?;
+```
+
+For more advanced examples, please have a look at the following section.
+
+## Check out our examples
+
+These online demos run entirely in your browser:
+- [yolo](https://huggingface.co/spaces/lmz/candle-yolo): pose estimation and
+  object recognition.
+- [whisper](https://huggingface.co/spaces/lmz/candle-whisper): speech recognition.
+- [LLaMA2](https://huggingface.co/spaces/lmz/candle-llama2): text generation.
+- [T5](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm): text generation.
+- [Phi-1.5, and Phi-2](https://huggingface.co/spaces/radames/Candle-Phi-1.5-Wasm): text generation.
+- [Segment Anything Model](https://huggingface.co/spaces/radames/candle-segment-anything-wasm): Image segmentation.
+- [BLIP](https://huggingface.co/spaces/radames/Candle-BLIP-Image-Captioning): image captioning.
+
+We also provide a some command line based examples using state of the art models:
+
+- [LLaMA v1, v2, and v3](./candle-examples/examples/llama/): general LLM, includes
+  the SOLAR-10.7B variant.
+- [Falcon](./candle-examples/examples/falcon/): general LLM.
+- [Codegeex4](./candle-examples/examples/codegeex4-9b/): Code completion,code interpreter,web search,fuction calling,repository-level
+- [GLM4](./candle-examples/examples/glm4/): Open Multilingual Multimodal Chat LMs by THUDM
+- [Gemma v1 and v2](./candle-examples/examples/gemma/): 2b and 7b+/9b general LLMs from Google Deepmind.
+- [RecurrentGemma](./candle-examples/examples/recurrent-gemma/): 2b and 7b
+  Griffin based models from Google that mix attention with a RNN like state.
+- [Phi-1, Phi-1.5, Phi-2, and Phi-3](./candle-examples/examples/phi/): 1.3b,
+  2.7b, and 3.8b general LLMs with performance on par with 7b models.
+- [StableLM-3B-4E1T](./candle-examples/examples/stable-lm/): a 3b general LLM
+  pre-trained on 1T tokens of English and code datasets. Also supports
+  StableLM-2, a 1.6b LLM trained on 2T tokens, as well as the code variants.
+- [Mamba](./candle-examples/examples/mamba/): an inference only
+  implementation of the Mamba state space model.
+- [Mistral7b-v0.1](./candle-examples/examples/mistral/): a 7b general LLM with
+  better performance than all publicly available 13b models as of 2023-09-28.
+- [Mixtral8x7b-v0.1](./candle-examples/examples/mixtral/): a sparse mixture of
+  experts 8x7b general LLM with better performance than a Llama 2 70B model with
+  much faster inference.
+- [StarCoder](./candle-examples/examples/bigcode/) and
+  [StarCoder2](./candle-examples/examples/starcoder2/): LLM specialized to code generation.
+- [Qwen1.5](./candle-examples/examples/qwen/): Bilingual (English/Chinese) LLMs.
+- [RWKV v5 and v6](./candle-examples/examples/rwkv/): An RNN with transformer level LLM
+  performance.
+- [Replit-code-v1.5](./candle-examples/examples/replit-code/): a 3.3b LLM specialized for code completion.
+- [Yi-6B / Yi-34B](./candle-examples/examples/yi/): two bilingual
+  (English/Chinese) general LLMs with 6b and 34b parameters.
+- [Quantized LLaMA](./candle-examples/examples/quantized/): quantized version of
+  the LLaMA model using the same quantization techniques as
+  [llama.cpp](https://github.com/ggerganov/llama.cpp).
+
+<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/quantized/assets/aoc.gif" width="600">
+  
+- [Stable Diffusion](./candle-examples/examples/stable-diffusion/): text to
+  image generative model, support for the 1.5, 2.1, SDXL 1.0 and Turbo versions.
+
+<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg" width="200">
+
+- [Wuerstchen](./candle-examples/examples/wuerstchen/): another text to
+  image generative model.
+
+<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/wuerstchen/assets/cat.jpg" width="200">
+
+- [yolo-v3](./candle-examples/examples/yolo-v3/) and
+  [yolo-v8](./candle-examples/examples/yolo-v8/): object detection and pose
+  estimation models.
+
+<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/yolo-v8/assets/bike.od.jpg" width="200"><img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/yolo-v8/assets/bike.pose.jpg" width="200">
+- [segment-anything](./candle-examples/examples/segment-anything/): image
+  segmentation model with prompt.
+
+<img src="https://github.com/huggingface/candle/raw/main/candle-examples/examples/segment-anything/assets/sam_merged.jpg" width="200">
+
+- [SegFormer](./candle-examples/examples/segformer/): transformer based semantic segmentation model.
+- [Whisper](./candle-examples/examples/whisper/): speech recognition model.
+- [EnCodec](./candle-examples/examples/encodec/): high-quality audio compression
+  model using residual vector quantization.
+- [MetaVoice](./candle-examples/examples/metavoice/): foundational model for
+  text-to-speech.
+- [Parler-TTS](./candle-examples/examples/parler-tts/): large text-to-speech
+  model.
+- [T5](./candle-examples/examples/t5), [Bert](./candle-examples/examples/bert/),
+  [JinaBert](./candle-examples/examples/jina-bert/) : useful for sentence embeddings.
+- [DINOv2](./candle-examples/examples/dinov2/): computer vision model trained
+  using self-supervision (can be used for imagenet classification, depth
+  evaluation, segmentation).
+- [VGG](./candle-examples/examples/vgg/),
+  [RepVGG](./candle-examples/examples/repvgg): computer vision models.
+- [BLIP](./candle-examples/examples/blip/): image to text model, can be used to
+  generate captions for an image.
+- [CLIP](./candle-examples/examples/clip/): multi-model vision and language
+  model.
+- [TrOCR](./candle-examples/examples/trocr/): a transformer OCR model, with
+  dedicated submodels for hand-writing and printed recognition.
+- [Marian-MT](./candle-examples/examples/marian-mt/): neural machine translation
+  model, generates the translated text from the input text.
+- [Moondream](./candle-examples/examples/moondream/): tiny computer-vision model 
+  that can answer real-world questions about images.
+
+Run them using commands like:
+```
+cargo run --example quantized --release
+```
+
+In order to use **CUDA** add `--features cuda` to the example command line. If
+you have cuDNN installed, use `--features cudnn` for even more speedups.
+
+There are also some wasm examples for whisper and
+[llama2.c](https://github.com/karpathy/llama2.c). You can either build them with
+`trunk` or try them online:
+[whisper](https://huggingface.co/spaces/lmz/candle-whisper),
+[llama2](https://huggingface.co/spaces/lmz/candle-llama2),
+[T5](https://huggingface.co/spaces/radames/Candle-T5-Generation-Wasm),
+[Phi-1.5, and Phi-2](https://huggingface.co/spaces/radames/Candle-Phi-1.5-Wasm),
+[Segment Anything Model](https://huggingface.co/spaces/radames/candle-segment-anything-wasm).
+
+For LLaMA2, run the following command to retrieve the weight files and start a
+test server:
+```bash
+cd candle-wasm-examples/llama2-c
+wget https://huggingface.co/spaces/lmz/candle-llama2/resolve/main/model.bin
+wget https://huggingface.co/spaces/lmz/candle-llama2/resolve/main/tokenizer.json
+trunk serve --release --port 8081
+```
+And then head over to
+[http://localhost:8081/](http://localhost:8081/).
+
+<!--- ANCHOR: useful_libraries --->
+
+## Useful External Resources
+- [`candle-tutorial`](https://github.com/ToluClassics/candle-tutorial): A
+  very detailed tutorial showing how to convert a PyTorch model to Candle.
+- [`candle-lora`](https://github.com/EricLBuehler/candle-lora): Efficient and
+  ergonomic LoRA implementation for Candle. `candle-lora` has      
+  out-of-the-box LoRA support for many models from Candle, which can be found
+  [here](https://github.com/EricLBuehler/candle-lora/tree/master/candle-lora-transformers/examples).
+- [`optimisers`](https://github.com/KGrewal1/optimisers): A collection of optimisers
+  including SGD with momentum, AdaGrad, AdaDelta, AdaMax, NAdam, RAdam, and RMSprop.
+- [`candle-vllm`](https://github.com/EricLBuehler/candle-vllm): Efficient platform for inference and
+  serving local LLMs including an OpenAI compatible API server.
+- [`candle-ext`](https://github.com/mokeyish/candle-ext): An extension library to Candle that provides PyTorch functions not currently available in Candle.
+- [`candle-coursera-ml`](https://github.com/vishpat/candle-coursera-ml): Implementation of ML algorithms from Coursera's [Machine Learning Specialization](https://www.coursera.org/specializations/machine-learning-introduction) course.
+- [`kalosm`](https://github.com/floneum/floneum/tree/master/interfaces/kalosm): A multi-modal meta-framework in Rust for interfacing with local pre-trained models with support for controlled generation, custom samplers, in-memory vector databases, audio transcription, and more.
+- [`candle-sampling`](https://github.com/EricLBuehler/candle-sampling): Sampling techniques for Candle.
+- [`gpt-from-scratch-rs`](https://github.com/jeroenvlek/gpt-from-scratch-rs): A port of Andrej Karpathy's _Let's build GPT_ tutorial on YouTube showcasing the Candle API on a toy problem.
+- [`candle-einops`](https://github.com/tomsanbear/candle-einops): A pure rust implementation of the python [einops](https://github.com/arogozhnikov/einops) library.
+- [`atoma-infer`](https://github.com/atoma-network/atoma-infer): A Rust library for fast inference at scale, leveraging FlashAttention2 for efficient attention computation, PagedAttention for efficient KV-cache memory management, and multi-GPU support. It is OpenAI api compatible.
+- [`llms-from-scratch-rs`](https://github.com/nerdai/llms-from-scratch-rs): A comprehensive Rust translation of the code from Sebastian Raschka's Build an LLM from Scratch book.
+
+If you have an addition to this list, please submit a pull request.
+
+<!--- ANCHOR_END: useful_libraries --->
+
+<!--- ANCHOR: features --->
+
+## Features
+
+- Simple syntax, looks and feels like PyTorch.
+    - Model training.
+    - Embed user-defined ops/kernels, such as [flash-attention v2](https://github.com/huggingface/candle/blob/89ba005962495f2bfbda286e185e9c3c7f5300a3/candle-flash-attn/src/lib.rs#L152).
+- Backends.
+    - Optimized CPU backend with optional MKL support for x86 and Accelerate for macs.
+    - CUDA backend for efficiently running on GPUs, multiple GPU distribution via NCCL.
+    - WASM support, run your models in a browser.
+- Included models.
+    - Language Models.
+        - LLaMA v1, v2, and v3 with variants such as SOLAR-10.7B.
+        - Falcon.
+        - StarCoder, StarCoder2.
+        - Phi 1, 1.5, 2, and 3.
+        - Mamba, Minimal Mamba
+        - Gemma v1 2b and 7b+, v2 2b and 9b.
+        - Mistral 7b v0.1.
+        - Mixtral 8x7b v0.1.
+        - StableLM-3B-4E1T, StableLM-2-1.6B, Stable-Code-3B.
+        - Replit-code-v1.5-3B.
+        - Bert.
+        - Yi-6B and Yi-34B.
+        - Qwen1.5, Qwen1.5 MoE.
+        - RWKV v5 and v6.
+    - Quantized LLMs.
+        - Llama 7b, 13b, 70b, as well as the chat and code variants.
+        - Mistral 7b, and 7b instruct.
+        - Mixtral 8x7b.
+        - Zephyr 7b a and b (Mistral-7b based).
+        - OpenChat 3.5 (Mistral-7b based).
+    - Text to text.
+        - T5 and its variants: FlanT5, UL2, MADLAD400 (translation), CoEdit (Grammar correction).
+        - Marian MT (Machine Translation).
+    - Text to image.
+        - Stable Diffusion v1.5, v2.1, XL v1.0.
+        - Wurstchen v2.
+    - Image to text.
+        - BLIP.
+        - TrOCR.
+    - Audio.
+        - Whisper, multi-lingual speech-to-text.
+        - EnCodec, audio compression model.
+        - MetaVoice-1B, text-to-speech model.
+        - Parler-TTS, text-to-speech model.
+    - Computer Vision Models.
+        - DINOv2, ConvMixer, EfficientNet, ResNet, ViT, VGG, RepVGG, ConvNeXT,
+          ConvNeXTv2, MobileOne, EfficientVit (MSRA), MobileNetv4, Hiera, FastViT.
+        - yolo-v3, yolo-v8.
+        - Segment-Anything Model (SAM).
+        - SegFormer.
+- File formats: load models from safetensors, npz, ggml, or PyTorch files.
+- Serverless (on CPU), small and fast deployments.
+- Quantization support using the llama.cpp quantized types.
+
+<!--- ANCHOR_END: features --->
+
+## How to use
+
+<!--- ANCHOR: cheatsheet --->
+Cheatsheet:
+
+|            | Using PyTorch                            | Using Candle                                                     |
+|------------|------------------------------------------|------------------------------------------------------------------|
+| Creation   | `torch.Tensor([[1, 2], [3, 4]])`         | `Tensor::new(&[[1f32, 2.], [3., 4.]], &Device::Cpu)?`           |
+| Creation   | `torch.zeros((2, 2))`                    | `Tensor::zeros((2, 2), DType::F32, &Device::Cpu)?`               |
+| Indexing   | `tensor[:, :4]`                          | `tensor.i((.., ..4))?`                                           |
+| Operations | `tensor.view((2, 2))`                    | `tensor.reshape((2, 2))?`                                        |
+| Operations | `a.matmul(b)`                            | `a.matmul(&b)?`                                                  |
+| Arithmetic | `a + b`                                  | `&a + &b`                                                        |
+| Device     | `tensor.to(device="cuda")`               | `tensor.to_device(&Device::new_cuda(0)?)?`                            |
+| Dtype      | `tensor.to(dtype=torch.float16)`         | `tensor.to_dtype(&DType::F16)?`                                  |
+| Saving     | `torch.save({"A": A}, "model.bin")`      | `candle::safetensors::save(&HashMap::from([("A", A)]), "model.safetensors")?` |
+| Loading    | `weights = torch.load("model.bin")`      | `candle::safetensors::load("model.safetensors", &device)`        |
+
+<!--- ANCHOR_END: cheatsheet --->
+
+
+## Structure
+
+- [candle-core](./candle-core): Core ops, devices, and `Tensor` struct definition
+- [candle-nn](./candle-nn/): Tools to build real models
+- [candle-examples](./candle-examples/): Examples of using the library in realistic settings
+- [candle-kernels](./candle-kernels/): CUDA custom kernels
+- [candle-datasets](./candle-datasets/): Datasets and data loaders.
+- [candle-transformers](./candle-transformers): transformers-related utilities.
+- [candle-flash-attn](./candle-flash-attn): Flash attention v2 layer.
+- [candle-onnx](./candle-onnx/): ONNX model evaluation.
+
+## FAQ
+
+### Why should I use Candle?
+
+Candle's core goal is to *make serverless inference possible*. Full machine learning frameworks like PyTorch
+are very large, which makes creating instances on a cluster slow. Candle allows deployment of lightweight
+binaries.
+
+Secondly, Candle lets you *remove Python* from production workloads. Python overhead can seriously hurt performance,
+and the [GIL](https://www.backblaze.com/blog/the-python-gil-past-present-and-future/) is a notorious source of headaches.
+
+Finally, Rust is cool! A lot of the HF ecosystem already has Rust crates, like [safetensors](https://github.com/huggingface/safetensors) and [tokenizers](https://github.com/huggingface/tokenizers).
+
+
+### Other ML frameworks
+
+- [dfdx](https://github.com/coreylowman/dfdx) is a formidable crate, with shapes being included
+  in types. This prevents a lot of headaches by getting the compiler to complain about shape mismatches right off the bat.
+  However, we found that some features still require nightly, and writing code can be a bit daunting for non rust experts.
+
+  We're leveraging and contributing to other core crates for the runtime so hopefully both crates can benefit from each
+  other.
+
+- [burn](https://github.com/burn-rs/burn) is a general crate that can leverage multiple backends so you can choose the best
+  engine for your workload.
+
+- [tch-rs](https://github.com/LaurentMazare/tch-rs.git) Bindings to the torch library in Rust. Extremely versatile, but they 
+  bring in the entire torch library into the runtime. The main contributor of `tch-rs` is also involved in the development
+  of `candle`.
+
+### Common Errors
+
+#### Missing symbols when compiling with the mkl feature.
+
+If you get some missing symbols when compiling binaries/tests using the mkl
+or accelerate features, e.g. for mkl you get:
+```
+  = note: /usr/bin/ld: (....o): in function `blas::sgemm':
+          .../blas-0.22.0/src/lib.rs:1944: undefined reference to `sgemm_' collect2: error: ld returned 1 exit status
+
+  = note: some `extern` functions couldn't be found; some native libraries may need to be installed or have their path specified
+  = note: use the `-l` flag to specify native libraries to link
+  = note: use the `cargo:rustc-link-lib` directive to specify the native libraries to link with Cargo
+```
+or for accelerate:
+```
+Undefined symbols for architecture arm64:
+            "_dgemm_", referenced from:
+                candle_core::accelerate::dgemm::h1b71a038552bcabe in libcandle_core...
+            "_sgemm_", referenced from:
+                candle_core::accelerate::sgemm::h2cf21c592cba3c47 in libcandle_core...
+          ld: symbol(s) not found for architecture arm64
+```
+
+This is likely due to a missing linker flag that was needed to enable the mkl library. You
+can try adding the following for mkl at the top of your binary:
+```rust
+extern crate intel_mkl_src;
+```
+or for accelerate:
+```rust
+extern crate accelerate_src;
+```
+
+#### Cannot run the LLaMA examples: access to source requires login credentials
+
+```
+Error: request error: https://huggingface.co/meta-llama/Llama-2-7b-hf/resolve/main/tokenizer.json: status code 401
+```
+
+This is likely because you're not permissioned for the LLaMA-v2 model. To fix
+this, you have to register on the huggingface-hub, accept the [LLaMA-v2 model
+conditions](https://huggingface.co/meta-llama/Llama-2-7b-hf), and set up your
+authentication token. See issue
+[#350](https://github.com/huggingface/candle/issues/350) for more details.
+
+#### Missing cute/cutlass headers when compiling flash-attn
+
+```
+  In file included from kernels/flash_fwd_launch_template.h:11:0,
+                   from kernels/flash_fwd_hdim224_fp16_sm80.cu:5:
+  kernels/flash_fwd_kernel.h:8:10: fatal error: cute/algorithm/copy.hpp: No such file or directory
+   #include <cute/algorithm/copy.hpp>
+            ^~~~~~~~~~~~~~~~~~~~~~~~~
+  compilation terminated.
+  Error: nvcc error while compiling:
+```
+[cutlass](https://github.com/NVIDIA/cutlass) is provided as a git submodule so you may want to run the following command to check it in properly.
+```bash
+git submodule update --init
+```
+
+#### Compiling with flash-attention fails
+
+```
+/usr/include/c++/11/bits/std_function.h:530:146: error: parameter packs not expanded with ‘...’:
+```
+
+This is a bug in gcc-11 triggered by the Cuda compiler. To fix this, install a different, supported gcc version - for example gcc-10, and specify the path to the compiler in the NVCC_CCBIN environment variable.
+```
+env NVCC_CCBIN=/usr/lib/gcc/x86_64-linux-gnu/10 cargo ...
+```
+
+#### Linking error on windows when running rustdoc or mdbook tests
+
+```
+Couldn't compile the test.
+---- .\candle-book\src\inference\hub.md - Using_the_hub::Using_in_a_real_model_ (line 50) stdout ----
+error: linking with `link.exe` failed: exit code: 1181
+//very long chain of linking
+ = note: LINK : fatal error LNK1181: cannot open input file 'windows.0.48.5.lib'
+```
+
+Make sure you link all native libraries that might be located outside a project target, e.g., to run mdbook tests, you should run:
+
+```
+mdbook test candle-book -L .\target\debug\deps\ `
+-L native=$env:USERPROFILE\.cargo\registry\src\index.crates.io-6f17d22bba15001f\windows_x86_64_msvc-0.42.2\lib `
+-L native=$env:USERPROFILE\.cargo\registry\src\index.crates.io-6f17d22bba15001f\windows_x86_64_msvc-0.48.5\lib
+```
+
+#### Extremely slow model load time with WSL
+
+This may be caused by the models being loaded from `/mnt/c`, more details on
+[stackoverflow](https://stackoverflow.com/questions/68972448/why-is-wsl-extremely-slow-when-compared-with-native-windows-npm-yarn-processing).
+
+#### Tracking down errors
+
+You can set `RUST_BACKTRACE=1` to be provided with backtraces when a candle
+error is generated.
+
+#### CudaRC error
+
+If you encounter an error like this one `called `Result::unwrap()` on an `Err` value: LoadLibraryExW { source: Os { code: 126, kind: Uncategorized, message: "The specified module could not be found." } }` on windows. To fix copy and rename these 3 files (make sure they are in path). The paths depend on your cuda version.
+`c:\Windows\System32\nvcuda.dll` -> `cuda.dll`
+`c:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin\cublas64_12.dll` -> `cublas.dll`
+`c:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin\curand64_10.dll` -> `curand.dll`

crates/vendor/candle/candle-book/.gitignore

@@ -0,0 +1 @@
+book

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

@@ -0,0 +1,48 @@
+[package]
+name = "candle-book"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+repository.workspace = true
+keywords.workspace = true
+categories.workspace = true
+license.workspace = true
+readme = "README.md"
+
+[dependencies]
+accelerate-src = { workspace = true, optional = true }
+candle = { workspace = true }
+candle-datasets = { workspace = true }
+candle-nn = { workspace = true }
+candle-transformers = { workspace = true }
+candle-flash-attn = { workspace = true, optional = true }
+safetensors = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
+num-traits = { workspace = true }
+intel-mkl-src = { workspace = true, optional = true }
+cudarc = { workspace = true, optional = true }
+half = { workspace = true, optional = true }
+image = { workspace = true, optional = true }
+anyhow = { workspace = true }
+tokio = "1.43.0"
+
+[dev-dependencies]
+byteorder = { workspace = true }
+hf-hub = { workspace = true, features=["tokio"]}
+clap = { workspace = true }
+memmap2 = { workspace = true }
+rand = { workspace = true }
+tokenizers = { workspace = true, features = ["onig"] }
+tracing = { workspace = true }
+tracing-chrome = { workspace = true }
+tracing-subscriber = { workspace = true }
+# Necessary to disambiguate with tokio in wasm examples which are 1.28.1
+parquet = { workspace = true }
+image = { workspace = true }
+
+[build-dependencies]
+anyhow = { workspace = true }
+
+[features]
+default = []

crates/vendor/candle/candle-book/book.toml

@@ -0,0 +1,6 @@
+[book]
+authors = ["Nicolas Patry"]
+language = "en"
+multilingual = false
+src = "src"
+title = "Candle Documentation"

crates/vendor/candle/candle-book/src/README.md

@@ -0,0 +1,6 @@
+# Introduction
+
+{{#include ../../README.md:features}}
+
+
+This book will introduce step by step how to use `candle`.

crates/vendor/candle/candle-book/src/SUMMARY.md

@@ -0,0 +1,28 @@
+# Summary
+
+[Introduction](README.md)
+
+# User Guide
+
+- [Installation](guide/installation.md)
+- [Hello World - MNIST](guide/hello_world.md)
+- [PyTorch cheatsheet](guide/cheatsheet.md)
+
+# Reference Guide
+
+- [Running a model](inference/inference.md)
+    - [Using the hub](inference/hub.md)
+- [Error management](error_manage.md)
+- [Training](training/training.md)
+    - [Simplified](training/simplified.md)
+    - [MNIST](training/mnist.md)
+    - [Fine-tuning]()
+    - [Serialization]()
+- [Advanced Cuda usage]()
+    - [Writing a custom kernel]()
+    - [Porting a custom kernel]()
+- [Using MKL]()
+- [Creating apps]()
+    - [Creating a WASM app]()
+    - [Creating a REST api webserver]()
+    - [Creating a desktop Tauri app]()

crates/vendor/candle/candle-book/src/advanced/mkl.md

@@ -0,0 +1 @@
+# Using MKL

crates/vendor/candle/candle-book/src/apps/README.md

@@ -0,0 +1 @@
+# Creating apps

crates/vendor/candle/candle-book/src/apps/desktop.md

@@ -0,0 +1 @@
+# Creating a desktop Tauri app

crates/vendor/candle/candle-book/src/apps/rest.md

@@ -0,0 +1 @@
+# Creating a REST api webserver

crates/vendor/candle/candle-book/src/apps/wasm.md

@@ -0,0 +1 @@
+# Creating a WASM app

crates/vendor/candle/candle-book/src/chapter_1.md

@@ -0,0 +1 @@
+# Chapter 1

crates/vendor/candle/candle-book/src/cuda/README.md

@@ -0,0 +1 @@
+# Advanced Cuda usage

crates/vendor/candle/candle-book/src/cuda/porting.md

@@ -0,0 +1 @@
+# Porting a custom kernel

crates/vendor/candle/candle-book/src/cuda/writing.md

@@ -0,0 +1 @@
+# Writing a custom kernel

crates/vendor/candle/candle-book/src/error_manage.md

@@ -0,0 +1,51 @@
+# Error management
+
+You might have seen in the code base a lot of `.unwrap()` or `?`.
+If you're unfamiliar with Rust check out the [Rust book](https://doc.rust-lang.org/book/ch09-02-recoverable-errors-with-result.html)
+for more information.
+
+What's important to know though, is that if you want to know *where* a particular operation failed
+You can simply use `RUST_BACKTRACE=1` to get the location of where the model actually failed.
+
+Let's see on failing code:
+
+```rust,ignore
+let x = Tensor::zeros((1, 784), DType::F32, &device)?;
+let y = Tensor::zeros((1, 784), DType::F32, &device)?;
+let z = x.matmul(&y)?;
+```
+
+Will print at runtime:
+
+```bash
+Error: ShapeMismatchBinaryOp { lhs: [1, 784], rhs: [1, 784], op: "matmul" }
+``` 
+
+
+After adding `RUST_BACKTRACE=1`:
+
+
+```bash
+Error: WithBacktrace { inner: ShapeMismatchBinaryOp { lhs: [1, 784], rhs: [1, 784], op: "matmul" }, backtrace: Backtrace [{ fn: "candle::error::Error::bt", file: "/home/nicolas/.cargo/git/checkouts/candle-5bb8ef7e0626d693/f291065/candle-core/src/error.rs", line: 200 }, { fn: "candle::tensor::Tensor::matmul", file: "/home/nicolas/.cargo/git/checkouts/candle-5bb8ef7e0626d693/f291065/candle-core/src/tensor.rs", line: 816 }, { fn: "myapp::main", file: "./src/main.rs", line: 29 }, { fn: "core::ops::function::FnOnce::call_once", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/core/src/ops/function.rs", line: 250 }, { fn: "std::sys_common::backtrace::__rust_begin_short_backtrace", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/sys_common/backtrace.rs", line: 135 }, { fn: "std::rt::lang_start::{{closure}}", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 166 }, { fn: "core::ops::function::impls::<impl core::ops::function::FnOnce<A> for &F>::call_once", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/core/src/ops/function.rs", line: 284 }, { fn: "std::panicking::try::do_call", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 500 }, { fn: "std::panicking::try", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 464 }, { fn: "std::panic::catch_unwind", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panic.rs", line: 142 }, { fn: "std::rt::lang_start_internal::{{closure}}", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 148 }, { fn: "std::panicking::try::do_call", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 500 }, { fn: "std::panicking::try", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panicking.rs", line: 464 }, { fn: "std::panic::catch_unwind", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/panic.rs", line: 142 }, { fn: "std::rt::lang_start_internal", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 148 }, { fn: "std::rt::lang_start", file: "/rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/std/src/rt.rs", line: 165 }, { fn: "main" }, { fn: "__libc_start_main" }, { fn: "_start" }] }
+```
+
+Not super pretty at the moment, but we can see error occurred on `{ fn: "myapp::main", file: "./src/main.rs", line: 29 }`
+
+
+Another thing to note, is that since Rust is compiled it is not necessarily as easy to recover proper stacktraces
+especially in release builds. We're using [`anyhow`](https://docs.rs/anyhow/latest/anyhow/) for that.
+The library is still young, please [report](https://github.com/LaurentMazare/candle/issues) any issues detecting where an error is coming from.
+
+## Cuda error management
+
+When running a model on Cuda, you might get a stacktrace not really representing the error.
+The reason is that CUDA is async by nature, and therefore the error might be caught while you were sending totally different kernels.
+
+One way to avoid this is to use `CUDA_LAUNCH_BLOCKING=1` as an environment variable. This will force every kernel to be launched sequentially.
+You might still however see the error happening on other kernels as the faulty kernel might exit without an error but spoiling some pointer for which the error will happen when dropping the `CudaSlice` only.
+
+
+If this occurs, you can use [`compute-sanitizer`](https://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/index.html)
+This tool is like `valgrind` but for cuda. It will help locate the errors in the kernels.
+
+

crates/vendor/candle/candle-book/src/guide/cheatsheet.md

@@ -0,0 +1,3 @@
+# Pytorch cheatsheet
+
+{{#include ../../../README.md:cheatsheet}}

crates/vendor/candle/candle-book/src/guide/hello_world.md

@@ -0,0 +1,195 @@
+# Hello world!
+
+We will now create the hello world of the ML world, building a model capable of solving MNIST dataset.
+
+Open `src/main.rs` and fill in this content:
+
+```rust
+# extern crate candle_core;
+use candle_core::{Device, Result, Tensor};
+
+struct Model {
+    first: Tensor,
+    second: Tensor,
+}
+
+impl Model {
+    fn forward(&self, image: &Tensor) -> Result<Tensor> {
+        let x = image.matmul(&self.first)?;
+        let x = x.relu()?;
+        x.matmul(&self.second)
+    }
+}
+
+fn main() -> Result<()> {
+    // Use Device::new_cuda(0)?; to use the GPU.
+    let device = Device::Cpu;
+
+    let first = Tensor::randn(0f32, 1.0, (784, 100), &device)?;
+    let second = Tensor::randn(0f32, 1.0, (100, 10), &device)?;
+    let model = Model { first, second };
+
+    let dummy_image = Tensor::randn(0f32, 1.0, (1, 784), &device)?;
+
+    let digit = model.forward(&dummy_image)?;
+    println!("Digit {digit:?} digit");
+    Ok(())
+}
+```
+
+Everything should now run with:
+
+```bash
+cargo run --release
+```
+
+## Using a `Linear` layer.
+
+Now that we have this, we might want to complexify things a bit, for instance by adding `bias` and creating
+the classical `Linear` layer. We can do as such
+
+```rust
+# extern crate candle_core;
+# use candle_core::{Device, Result, Tensor};
+struct Linear{
+    weight: Tensor,
+    bias: Tensor,
+}
+impl Linear{
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = x.matmul(&self.weight)?;
+        x.broadcast_add(&self.bias)
+    }
+}
+
+struct Model {
+    first: Linear,
+    second: Linear,
+}
+
+impl Model {
+    fn forward(&self, image: &Tensor) -> Result<Tensor> {
+        let x = self.first.forward(image)?;
+        let x = x.relu()?;
+        self.second.forward(&x)
+    }
+}
+```
+
+This will change the model running code into a new function
+
+```rust
+# extern crate candle_core;
+# use candle_core::{Device, Result, Tensor};
+# struct Linear{
+#     weight: Tensor,
+#     bias: Tensor,
+# }
+# impl Linear{
+#     fn forward(&self, x: &Tensor) -> Result<Tensor> {
+#         let x = x.matmul(&self.weight)?;
+#         x.broadcast_add(&self.bias)
+#     }
+# }
+# 
+# struct Model {
+#     first: Linear,
+#     second: Linear,
+# }
+# 
+# impl Model {
+#     fn forward(&self, image: &Tensor) -> Result<Tensor> {
+#         let x = self.first.forward(image)?;
+#         let x = x.relu()?;
+#         self.second.forward(&x)
+#     }
+# }
+fn main() -> Result<()> {
+    // Use Device::new_cuda(0)?; to use the GPU.
+    // Use Device::Cpu; to use the CPU.
+    let device = Device::cuda_if_available(0)?;
+
+    // Creating a dummy model
+    let weight = Tensor::randn(0f32, 1.0, (784, 100), &device)?;
+    let bias = Tensor::randn(0f32, 1.0, (100, ), &device)?;
+    let first = Linear{weight, bias};
+    let weight = Tensor::randn(0f32, 1.0, (100, 10), &device)?;
+    let bias = Tensor::randn(0f32, 1.0, (10, ), &device)?;
+    let second = Linear{weight, bias};
+    let model = Model { first, second };
+
+    let dummy_image = Tensor::randn(0f32, 1.0, (1, 784), &device)?;
+
+    // Inference on the model
+    let digit = model.forward(&dummy_image)?;
+    println!("Digit {digit:?} digit");
+    Ok(())
+}
+```
+
+Now it works, it is a great way to create your own layers.
+But most of the classical layers are already implemented in [candle-nn](https://github.com/huggingface/candle/tree/main/candle-nn).
+
+## Using `candle_nn`.
+
+For instance [Linear](https://github.com/huggingface/candle/blob/main/candle-nn/src/linear.rs) is already there.
+This Linear is coded with PyTorch layout in mind, to reuse better existing models out there, so it uses the transpose of the weights and not the weights directly.
+
+So instead we can simplify our example:
+
+```bash
+cargo add --git https://github.com/huggingface/candle.git candle-nn
+```
+
+And rewrite our examples using it
+
+```rust
+# extern crate candle_core;
+# extern crate candle_nn;
+use candle_core::{Device, Result, Tensor};
+use candle_nn::{Linear, Module};
+
+struct Model {
+    first: Linear,
+    second: Linear,
+}
+
+impl Model {
+    fn forward(&self, image: &Tensor) -> Result<Tensor> {
+        let x = self.first.forward(image)?;
+        let x = x.relu()?;
+        self.second.forward(&x)
+    }
+}
+
+fn main() -> Result<()> {
+    // Use Device::new_cuda(0)?; to use the GPU.
+    let device = Device::Cpu;
+
+    // This has changed (784, 100) -> (100, 784) !
+    let weight = Tensor::randn(0f32, 1.0, (100, 784), &device)?;
+    let bias = Tensor::randn(0f32, 1.0, (100, ), &device)?;
+    let first = Linear::new(weight, Some(bias));
+    let weight = Tensor::randn(0f32, 1.0, (10, 100), &device)?;
+    let bias = Tensor::randn(0f32, 1.0, (10, ), &device)?;
+    let second = Linear::new(weight, Some(bias));
+    let model = Model { first, second };
+
+    let dummy_image = Tensor::randn(0f32, 1.0, (1, 784), &device)?;
+
+    let digit = model.forward(&dummy_image)?;
+    println!("Digit {digit:?} digit");
+    Ok(())
+}
+```
+
+Feel free to modify this example to use `Conv2d` to create a classical convnet instead.
+
+
+Now that we have the running dummy code we can get to more advanced topics:
+
+- [For PyTorch users](../guide/cheatsheet.md)
+- [Running existing models](../inference/inference.md)
+- [Training models](../training/training.md)
+
+

crates/vendor/candle/candle-book/src/guide/installation.md

@@ -0,0 +1,59 @@
+# Installation
+
+**With Cuda support**:
+
+1. First, make sure that Cuda is correctly installed.
+- `nvcc --version` should print information about your Cuda compiler driver.
+- `nvidia-smi --query-gpu=compute_cap --format=csv` should print your GPUs compute capability, e.g. something
+like:
+
+```bash
+compute_cap
+8.9
+```
+
+You can also compile the Cuda kernels for a specific compute cap using the 
+`CUDA_COMPUTE_CAP=<compute cap>` environment variable.
+
+If any of the above commands errors out, please make sure to update your Cuda version.
+
+2. Create a new app and add [`candle-core`](https://github.com/huggingface/candle/tree/main/candle-core) with Cuda support.
+
+Start by creating a new cargo:
+
+```bash
+cargo new myapp
+cd myapp
+```
+
+Make sure to add the `candle-core` crate with the cuda feature:
+
+```bash
+cargo add --git https://github.com/huggingface/candle.git candle-core --features "cuda"
+```
+
+Run `cargo build` to make sure everything can be correctly built.
+
+```bash
+cargo build
+```
+
+**Without Cuda support**:
+
+Create a new app and add [`candle-core`](https://github.com/huggingface/candle/tree/main/candle-core) as follows:
+
+```bash
+cargo new myapp
+cd myapp
+cargo add --git https://github.com/huggingface/candle.git candle-core
+```
+
+Finally, run `cargo build` to make sure everything can be correctly built.
+
+```bash
+cargo build
+```
+
+**With mkl support**
+
+You can also see the `mkl` feature which could be interesting to get faster inference on CPU. [Using mkl](./advanced/mkl.md)

crates/vendor/candle/candle-book/src/inference/cuda/README.md

@@ -0,0 +1 @@
+# Advanced Cuda usage

crates/vendor/candle/candle-book/src/inference/cuda/porting.md

@@ -0,0 +1 @@
+# Porting a custom kernel

crates/vendor/candle/candle-book/src/inference/cuda/writing.md

@@ -0,0 +1 @@
+# Writing a custom kernel

crates/vendor/candle/candle-book/src/inference/hub.md

@@ -0,0 +1,104 @@
+# Using the hub
+
+Install the [`hf-hub`](https://github.com/huggingface/hf-hub) crate:
+
+```bash
+cargo add hf-hub
+```
+
+Then let's start by downloading the [model file](https://huggingface.co/bert-base-uncased/tree/main).
+
+
+```rust
+# extern crate candle_core;
+# extern crate hf_hub;
+use hf_hub::api::sync::Api;
+use candle_core::Device;
+
+let api = Api::new().unwrap();
+let repo = api.model("bert-base-uncased".to_string());
+
+let weights = repo.get("model.safetensors").unwrap();
+
+let weights = candle_core::safetensors::load(weights, &Device::Cpu);
+```
+
+We now have access to all the [tensors](https://huggingface.co/bert-base-uncased?show_tensors=true) within the file.
+
+You can check all the names of the tensors [here](https://huggingface.co/bert-base-uncased?show_tensors=true)
+
+
+## Using async 
+
+`hf-hub` comes with an async API.
+
+```bash
+cargo add hf-hub --features tokio
+```
+
+```rust,ignore
+# This is tested directly in examples crate because it needs external dependencies unfortunately:
+# See [this](https://github.com/rust-lang/mdBook/issues/706)
+{{#include ../lib.rs:book_hub_1}}
+```
+
+
+## Using in a real model.
+
+Now that we have our weights, we can use them in our bert architecture:
+
+```rust
+# extern crate candle_core;
+# extern crate candle_nn;
+# extern crate hf_hub;
+# use hf_hub::api::sync::Api;
+# 
+# let api = Api::new().unwrap();
+# let repo = api.model("bert-base-uncased".to_string());
+# 
+# let weights = repo.get("model.safetensors").unwrap();
+use candle_core::{Device, Tensor, DType};
+use candle_nn::{Linear, Module};
+
+let weights = candle_core::safetensors::load(weights, &Device::Cpu).unwrap();
+
+let weight = weights.get("bert.encoder.layer.0.attention.self.query.weight").unwrap();
+let bias = weights.get("bert.encoder.layer.0.attention.self.query.bias").unwrap();
+
+let linear = Linear::new(weight.clone(), Some(bias.clone()));
+
+let input_ids = Tensor::zeros((3, 768), DType::F32, &Device::Cpu).unwrap();
+let output = linear.forward(&input_ids).unwrap();
+```
+
+For a full reference, you can check out the full [bert](https://github.com/LaurentMazare/candle/tree/main/candle-examples/examples/bert) example.
+
+## Memory mapping
+
+For more efficient loading, instead of reading the file, you could use [`memmap2`](https://docs.rs/memmap2/latest/memmap2/)
+
+**Note**: Be careful about memory mapping it seems to cause issues on [Windows, WSL](https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/5893)
+and will definitely be slower on network mounted disk, because it will issue more read calls.
+
+```rust,ignore
+{{#include ../lib.rs:book_hub_2}}
+```
+
+**Note**: This operation is **unsafe**. [See the safety notice](https://docs.rs/memmap2/latest/memmap2/struct.Mmap.html#safety).
+In practice model files should never be modified, and the mmaps should be mostly READONLY anyway, so the caveat most likely does not apply, but always keep it in mind.
+
+
+## Tensor Parallel Sharding
+
+When using multiple GPUs to use in Tensor Parallel in order to get good latency, you can load only the part of the Tensor you need.
+
+For that you need to use [`safetensors`](https://crates.io/crates/safetensors) directly.
+
+```bash
+cargo add safetensors
+```
+
+
+```rust,ignore
+{{#include ../lib.rs:book_hub_3}}
+```

crates/vendor/candle/candle-book/src/inference/inference.md

@@ -0,0 +1,7 @@
+# Running a model
+
+
+In order to run an existing model, you will need to download and use existing weights.
+Most models are already available on https://huggingface.co/ in [`safetensors`](https://github.com/huggingface/safetensors) format.
+
+Let's get started by running an old model : `bert-base-uncased`.

crates/vendor/candle/candle-book/src/lib.rs

@@ -0,0 +1,199 @@
+#[cfg(test)]
+pub mod simplified;
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use candle::{DType, Device, Tensor};
+    use parquet::file::reader::SerializedFileReader;
+
+    // NOTE: Waiting on https://github.com/rust-lang/mdBook/pull/1856
+    #[rustfmt::skip]
+    #[tokio::test]
+    async fn book_hub_1() {
+// ANCHOR: book_hub_1
+use candle::Device;
+use hf_hub::api::tokio::Api;
+
+let api = Api::new().unwrap();
+let repo = api.model("bert-base-uncased".to_string());
+
+let weights_filename = repo.get("model.safetensors").await.unwrap();
+
+let weights = candle::safetensors::load(weights_filename, &Device::Cpu).unwrap();
+// ANCHOR_END: book_hub_1
+        assert_eq!(weights.len(), 206);
+    }
+
+    #[rustfmt::skip]
+    #[test]
+    fn book_hub_2() {
+        {
+// ANCHOR: book_hub_2
+use candle::Device;
+use hf_hub::api::sync::Api;
+use memmap2::Mmap;
+use std::fs;
+
+let api = Api::new().unwrap();
+let repo = api.model("bert-base-uncased".to_string());
+let weights_filename = repo.get("model.safetensors").unwrap();
+
+let file = fs::File::open(weights_filename).unwrap();
+let mmap = unsafe { Mmap::map(&file).unwrap() };
+let weights = candle::safetensors::load_buffer(&mmap[..], &Device::Cpu).unwrap();
+// ANCHOR_END: book_hub_2
+        assert_eq!(weights.len(), 206);
+    }
+
+    // #[rustfmt::skip]
+    // #[test]
+    // fn book_hub_3() {
+    {
+// ANCHOR: book_hub_3
+use candle::{DType, Device, Tensor};
+use hf_hub::api::sync::Api;
+use memmap2::Mmap;
+use safetensors::slice::IndexOp;
+use safetensors::SafeTensors;
+use std::fs;
+
+let api = Api::new().unwrap();
+let repo = api.model("bert-base-uncased".to_string());
+let weights_filename = repo.get("model.safetensors").unwrap();
+
+let file = fs::File::open(weights_filename).unwrap();
+let mmap = unsafe { Mmap::map(&file).unwrap() };
+
+// Use safetensors directly
+let tensors = SafeTensors::deserialize(&mmap[..]).unwrap();
+let view = tensors
+    .tensor("bert.encoder.layer.0.attention.self.query.weight")
+    .unwrap();
+
+// We're going to load shard with rank 1, within a world_size of 4
+// We're going to split along dimension 0 doing VIEW[start..stop, :]
+let rank = 1;
+let world_size = 4;
+let dim = 0;
+let dtype = view.dtype();
+let mut tp_shape = view.shape().to_vec();
+let size = tp_shape[0];
+
+if size % world_size != 0 {
+    panic!("The dimension is not divisible by `world_size`");
+}
+let block_size = size / world_size;
+let start = rank * block_size;
+let stop = (rank + 1) * block_size;
+
+// Everything is expressed in tensor dimension
+// bytes offsets is handled automatically for safetensors.
+
+let iterator = view.slice(start..stop).unwrap();
+
+tp_shape[dim] = block_size;
+
+// Convert safetensors Dtype to candle DType
+let dtype: DType = dtype.try_into().unwrap();
+
+// TODO: Implement from_buffer_iterator so we can skip the extra CPU alloc.
+let raw: Vec<u8> = iterator.into_iter().flatten().cloned().collect();
+let tp_tensor = Tensor::from_raw_buffer(&raw, dtype, &tp_shape, &Device::Cpu).unwrap();
+// ANCHOR_END: book_hub_3
+        assert_eq!(view.shape(), &[768, 768]);
+        assert_eq!(tp_tensor.dims(), &[192, 768]);
+    }
+}
+
+    #[allow(unused)]
+    #[rustfmt::skip]
+    fn book_training_1() -> Result<()>{
+// ANCHOR: book_training_1
+use hf_hub::{api::sync::Api, Repo, RepoType};
+
+let dataset_id = "mnist".to_string();
+
+let api = Api::new()?;
+let repo = Repo::with_revision(
+    dataset_id,
+    RepoType::Dataset,
+    "refs/convert/parquet".to_string(),
+);
+let repo = api.repo(repo);
+let test_parquet_filename = repo.get("mnist/test/0000.parquet")?;
+let train_parquet_filename = repo.get("mnist/train/0000.parquet")?;
+let test_parquet = SerializedFileReader::new(std::fs::File::open(test_parquet_filename)?)?;
+let train_parquet = SerializedFileReader::new(std::fs::File::open(train_parquet_filename)?)?;
+// ANCHOR_END: book_training_1
+// Ignore unused
+let _train = train_parquet;
+// ANCHOR: book_training_2
+for row in test_parquet {
+    for (idx, (name, field)) in row?.get_column_iter().enumerate() {
+        println!("Column id {idx}, name {name}, value {field}");
+    }
+}
+// ANCHOR_END: book_training_2
+let test_parquet_filename = repo.get("mnist/test/0000.parquet")?;
+let train_parquet_filename = repo.get("mnist/train/0000.parquet")?;
+let test_parquet = SerializedFileReader::new(std::fs::File::open(test_parquet_filename)?)?;
+let train_parquet = SerializedFileReader::new(std::fs::File::open(train_parquet_filename)?)?;
+// ANCHOR: book_training_3
+
+let test_samples = 10_000;
+let mut test_buffer_images: Vec<u8> = Vec::with_capacity(test_samples * 784);
+let mut test_buffer_labels: Vec<u8> = Vec::with_capacity(test_samples);
+for row in test_parquet{
+    for (_name, field) in row?.get_column_iter() {
+        if let parquet::record::Field::Group(subrow) = field {
+            for (_name, field) in subrow.get_column_iter() {
+                if let parquet::record::Field::Bytes(value) = field {
+                    let image = image::load_from_memory(value.data()).unwrap();
+                    test_buffer_images.extend(image.to_luma8().as_raw());
+                }
+            }
+        }else if let parquet::record::Field::Long(label) = field {
+            test_buffer_labels.push(*label as u8);
+        }
+    }
+}
+let test_images = (Tensor::from_vec(test_buffer_images, (test_samples, 784), &Device::Cpu)?.to_dtype(DType::F32)? / 255.)?;
+let test_labels = Tensor::from_vec(test_buffer_labels, (test_samples, ), &Device::Cpu)?;
+
+let train_samples = 60_000;
+let mut train_buffer_images: Vec<u8> = Vec::with_capacity(train_samples * 784);
+let mut train_buffer_labels: Vec<u8> = Vec::with_capacity(train_samples);
+for row in train_parquet{
+    for (_name, field) in row?.get_column_iter() {
+        if let parquet::record::Field::Group(subrow) = field {
+            for (_name, field) in subrow.get_column_iter() {
+                if let parquet::record::Field::Bytes(value) = field {
+                    let image = image::load_from_memory(value.data()).unwrap();
+                    train_buffer_images.extend(image.to_luma8().as_raw());
+                }
+            }
+        }else if let parquet::record::Field::Long(label) = field {
+            train_buffer_labels.push(*label as u8);
+        }
+    }
+}
+let train_images = (Tensor::from_vec(train_buffer_images, (train_samples, 784), &Device::Cpu)?.to_dtype(DType::F32)? / 255.)?;
+let train_labels = Tensor::from_vec(train_buffer_labels, (train_samples, ), &Device::Cpu)?;
+
+let mnist = candle_datasets::vision::Dataset {
+    train_images,
+    train_labels,
+    test_images,
+    test_labels,
+    labels: 10,
+};
+
+// ANCHOR_END: book_training_3
+assert_eq!(mnist.test_images.dims(), &[10_000, 784]);
+assert_eq!(mnist.test_labels.dims(), &[10_000]);
+assert_eq!(mnist.train_images.dims(), &[60_000, 784]);
+assert_eq!(mnist.train_labels.dims(), &[60_000]);
+Ok(())
+    }
+}

crates/vendor/candle/candle-book/src/simplified.rs

@@ -0,0 +1,196 @@
+//! #A simplified example in Rust of training a neural network and then using it based on the Candle Framework by Hugging Face.
+//! Author: Evgeny Igumnov 2023 igumnovnsk@gmail.com
+//! This program implements a neural network to predict the winner of the second round of elections based on the results of the first round.
+//!
+//! ##Basic moments:
+//!
+//! A multilayer perceptron with two hidden layers is used. The first hidden layer has 4 neurons, the second has 2 neurons.
+//! The input is a vector of 2 numbers - the percentage of votes for the first and second candidates in the first stage.
+//! The output is the number 0 or 1, where 1 means that the first candidate will win in the second stage, 0 means that he will lose.
+//! For training, samples with real data on the results of the first and second stages of different elections are used.
+//! The model is trained by backpropagation using gradient descent and the cross-entropy loss function.
+//! Model parameters (weights of neurons) are initialized randomly, then optimized during training.
+//! After training, the model is tested on a deferred sample to evaluate the accuracy.
+//! If the accuracy on the test set is below 100%, the model is considered underfit and the learning process is repeated.
+//! Thus, this neural network learns to find hidden relationships between the results of the first and second rounds of voting in order to make predictions for new data.
+
+#[rustfmt::skip]
+mod tests {
+
+use candle::{DType, Result, Tensor, D, Device};
+use candle_nn::{loss, ops, Linear, Module, VarBuilder, VarMap, Optimizer};
+
+// ANCHOR: book_training_simplified1
+const VOTE_DIM: usize = 2;
+const RESULTS: usize = 1;
+const EPOCHS: usize = 10;
+const LAYER1_OUT_SIZE: usize = 4;
+const LAYER2_OUT_SIZE: usize = 2;
+const LEARNING_RATE: f64 = 0.05;
+
+#[derive(Clone)]
+pub struct Dataset {
+    pub train_votes: Tensor,
+    pub train_results: Tensor,
+    pub test_votes: Tensor,
+    pub test_results: Tensor,
+}
+
+struct MultiLevelPerceptron {
+    ln1: Linear,
+    ln2: Linear,
+    ln3: Linear,
+}
+
+impl MultiLevelPerceptron {
+    fn new(vs: VarBuilder) -> Result<Self> {
+        let ln1 = candle_nn::linear(VOTE_DIM, LAYER1_OUT_SIZE, vs.pp("ln1"))?;
+        let ln2 = candle_nn::linear(LAYER1_OUT_SIZE, LAYER2_OUT_SIZE, vs.pp("ln2"))?;
+        let ln3 = candle_nn::linear(LAYER2_OUT_SIZE, RESULTS + 1, vs.pp("ln3"))?;
+        Ok(Self { ln1, ln2, ln3 })
+    }
+
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.ln1.forward(xs)?;
+        let xs = xs.relu()?;
+        let xs = self.ln2.forward(&xs)?;
+        let xs = xs.relu()?;
+        self.ln3.forward(&xs)
+    }
+}
+
+// ANCHOR_END: book_training_simplified1
+
+
+
+// ANCHOR: book_training_simplified3
+#[tokio::test]
+async fn simplified() -> anyhow::Result<()> {
+
+    let dev = Device::cuda_if_available(0)?;
+
+    let train_votes_vec: Vec<u32> = vec![
+        15, 10,
+        10, 15,
+        5, 12,
+        30, 20,
+        16, 12,
+        13, 25,
+        6, 14,
+        31, 21,
+    ];
+    let train_votes_tensor = Tensor::from_vec(train_votes_vec.clone(), (train_votes_vec.len() / VOTE_DIM, VOTE_DIM), &dev)?.to_dtype(DType::F32)?;
+
+    let train_results_vec: Vec<u32> = vec![
+        1,
+        0,
+        0,
+        1,
+        1,
+        0,
+        0,
+        1,
+    ];
+    let train_results_tensor = Tensor::from_vec(train_results_vec, train_votes_vec.len() / VOTE_DIM, &dev)?;
+
+    let test_votes_vec: Vec<u32> = vec![
+        13, 9,
+        8, 14,
+        3, 10,
+    ];
+    let test_votes_tensor = Tensor::from_vec(test_votes_vec.clone(), (test_votes_vec.len() / VOTE_DIM, VOTE_DIM), &dev)?.to_dtype(DType::F32)?;
+
+    let test_results_vec: Vec<u32> = vec![
+        1,
+        0,
+        0,
+    ];
+    let test_results_tensor = Tensor::from_vec(test_results_vec.clone(), test_results_vec.len(), &dev)?;
+
+    let m = Dataset {
+        train_votes: train_votes_tensor,
+        train_results: train_results_tensor,
+        test_votes: test_votes_tensor,
+        test_results: test_results_tensor,
+    };
+
+    let trained_model: MultiLevelPerceptron;
+    loop {
+        println!("Trying to train neural network.");
+        match train(m.clone(), &dev) {
+            Ok(model) => {
+                trained_model = model;
+                break;
+            },
+            Err(e) => {
+                println!("Error: {}", e);
+                continue;
+            }
+        }
+
+    }
+
+    let real_world_votes: Vec<u32> = vec![
+        13, 22,
+    ];
+
+    let tensor_test_votes = Tensor::from_vec(real_world_votes.clone(), (1, VOTE_DIM), &dev)?.to_dtype(DType::F32)?;
+
+    let final_result = trained_model.forward(&tensor_test_votes)?;
+
+    let result = final_result
+        .argmax(D::Minus1)?
+        .to_dtype(DType::F32)?
+        .get(0).map(|x| x.to_scalar::<f32>())??;
+    println!("real_life_votes: {:?}", real_world_votes);
+    println!("neural_network_prediction_result: {:?}", result);
+
+    Ok(())
+
+}
+// ANCHOR_END: book_training_simplified3
+
+// ANCHOR: book_training_simplified2
+fn train(m: Dataset, dev: &Device) -> anyhow::Result<MultiLevelPerceptron> {
+    let train_results = m.train_results.to_device(dev)?;
+    let train_votes = m.train_votes.to_device(dev)?;
+    let varmap = VarMap::new();
+    let vs = VarBuilder::from_varmap(&varmap, DType::F32, dev);
+    let model = MultiLevelPerceptron::new(vs.clone())?;
+    let mut sgd = candle_nn::SGD::new(varmap.all_vars(), LEARNING_RATE)?;
+    let test_votes = m.test_votes.to_device(dev)?;
+    let test_results = m.test_results.to_device(dev)?;
+    let mut final_accuracy: f32 = 0.0;
+    for epoch in 1..EPOCHS + 1 {
+        let logits = model.forward(&train_votes)?;
+        let log_sm = ops::log_softmax(&logits, D::Minus1)?;
+        let loss = loss::nll(&log_sm, &train_results)?;
+        sgd.backward_step(&loss)?;
+
+        let test_logits = model.forward(&test_votes)?;
+        let sum_ok = test_logits
+            .argmax(D::Minus1)?
+            .eq(&test_results)?
+            .to_dtype(DType::F32)?
+            .sum_all()?
+            .to_scalar::<f32>()?;
+        let test_accuracy = sum_ok / test_results.dims1()? as f32;
+        final_accuracy = 100. * test_accuracy;
+        println!("Epoch: {epoch:3} Train loss: {:8.5} Test accuracy: {:5.2}%",
+                 loss.to_scalar::<f32>()?,
+                 final_accuracy
+        );
+        if final_accuracy == 100.0 {
+            break;
+        }
+    }
+    if final_accuracy < 100.0 {
+        Err(anyhow::Error::msg("The model is not trained well enough."))
+    } else {
+        Ok(model)
+    }
+}
+// ANCHOR_END: book_training_simplified2
+
+
+}

crates/vendor/candle/candle-book/src/training/finetuning.md

@@ -0,0 +1 @@
+# Fine-tuning

crates/vendor/candle/candle-book/src/training/mnist.md

@@ -0,0 +1,10 @@
+# MNIST
+
+So we now have downloaded the MNIST parquet files, let's put them in a simple struct.
+
+```rust,ignore
+{{#include ../lib.rs:book_training_3}}
+```
+
+The parsing of the file and putting it into single tensors requires the dataset to fit the entire memory.
+It is quite rudimentary, but simple enough for a small dataset like MNIST.

crates/vendor/candle/candle-book/src/training/serialization.md

@@ -0,0 +1 @@
+# Serialization

crates/vendor/candle/candle-book/src/training/simplified.md

@@ -0,0 +1,45 @@
+# Simplified
+
+## How its works
+
+This program implements a neural network to predict the winner of the second round of elections based on the results of the first round.
+
+Basic moments:
+
+1. A multilayer perceptron with two hidden layers is used. The first hidden layer has 4 neurons, the second has 2 neurons.
+2. The input is a vector of 2 numbers - the percentage of votes for the first and second candidates in the first stage.
+3. The output is the number 0 or 1, where 1 means that the first candidate will win in the second stage, 0 means that he will lose.
+4. For training, samples with real data on the results of the first and second stages of different elections are used.
+5. The model is trained by backpropagation using gradient descent and the cross-entropy loss function.
+6. Model parameters (weights of neurons) are initialized randomly, then optimized during training.
+7. After training, the model is tested on a deferred sample to evaluate the accuracy.
+8. If the accuracy on the test set is below 100%, the model is considered underfit and the learning process is repeated.
+
+Thus, this neural network learns to find hidden relationships between the results of the first and second rounds of voting in order to make predictions for new data.
+
+
+```rust,ignore
+{{#include ../simplified.rs:book_training_simplified1}}
+```
+
+```rust,ignore
+{{#include ../simplified.rs:book_training_simplified2}}
+```
+
+```rust,ignore
+{{#include ../simplified.rs:book_training_simplified3}}
+```
+
+
+## Example output
+
+```bash
+Trying to train neural network.
+Epoch:   1 Train loss:  4.42555 Test accuracy:  0.00%
+Epoch:   2 Train loss:  0.84677 Test accuracy: 33.33%
+Epoch:   3 Train loss:  2.54335 Test accuracy: 33.33%
+Epoch:   4 Train loss:  0.37806 Test accuracy: 33.33%
+Epoch:   5 Train loss:  0.36647 Test accuracy: 100.00%
+real_life_votes: [13, 22]
+neural_network_prediction_result: 0.0
+```

crates/vendor/candle/candle-book/src/training/training.md

@@ -0,0 +1,39 @@
+# Training
+
+
+Training starts with data. We're going to use the huggingface hub and 
+start with the Hello world dataset of machine learning, MNIST.
+
+Let's start with downloading `MNIST` from [huggingface](https://huggingface.co/datasets/mnist).
+
+This requires [`hf-hub`](https://github.com/huggingface/hf-hub).
+```bash
+cargo add hf-hub
+```
+
+This is going to be very hands-on for now.
+
+```rust,ignore
+{{#include ../../../candle-examples/src/lib.rs:book_training_1}}
+```
+
+This uses the standardized `parquet` files from the `refs/convert/parquet` branch on every dataset.
+Our handles are now [`parquet::file::serialized_reader::SerializedFileReader`].
+
+We can inspect the content of the files with:
+
+```rust,ignore
+{{#include ../../../candle-examples/src/lib.rs:book_training_2}}
+```
+
+You should see something like:
+
+```bash
+Column id 1, name label, value 6
+Column id 0, name image, value {bytes: [137, ....]
+Column id 1, name label, value 8
+Column id 0, name image, value {bytes: [137, ....]
+```
+
+So each row contains 2 columns (image, label) with image being saved as bytes.
+Let's put them into a useful struct.

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

@@ -0,0 +1,51 @@
+[package]
+name = "candle-core"
+edition = "2021"
+
+[dependencies]
+accelerate-src = { version = "0.3.2" , optional = true }
+byteorder = "1.4.3"
+candle-kernels = { path = "../candle-kernels", version = "0.8.3", optional = true }
+candle-metal-kernels = { path = "../candle-metal-kernels", version = "0.8.3" , optional = true }
+metal = { version = "0.27.0", features = ["mps"], optional = 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 }
+gemm = { version = "0.17.0", features = ["wasm-simd128-enable"] }
+half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
+intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"], optional = true }
+libc = { version = "0.2.147", optional = true }
+memmap2 = { version = "0.9.3", features = ["stable_deref_trait"] }
+num-traits = "0.2.15"
+num_cpus = "1.15.0"
+rand = "0.8.5"
+rand_distr = "0.4.3"
+rayon = "1.7.0"
+safetensors = "0.4.1"
+thiserror = "1"
+ug-cuda = { version = "0.1.0", optional = true }
+ug-metal = { version = "0.1.0", optional = true }
+yoke = { version = "0.7.2", features = ["derive"] }
+zip = { version = "1.1.1", default-features = false }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+ug = { version = "0.1.0" }
+
+[dev-dependencies]
+anyhow = { version = "1", features = ["backtrace"] }
+clap = { version = "4.2.4", features = ["derive"] }
+criterion = { version = "0.5.1", default-features=false }
+
+[features]
+default = []
+cuda = ["cudarc", "dep:candle-kernels", "dep:ug-cuda"]
+cudnn = ["cuda", "cudarc/cudnn"]
+mkl = ["dep:libc", "dep:intel-mkl-src"]
+accelerate = ["dep:libc", "dep:accelerate-src"]
+metal = ["dep:metal", "dep:candle-metal-kernels", "dep:ug-metal"]
+
+[[bench]]
+name = "bench_main"
+harness = false
+
+[[example]]
+name = "metal_basics"
+required-features = ["metal"]

crates/vendor/candle/candle-core/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

crates/vendor/candle/candle-core/README.md

@@ -0,0 +1,2 @@
+# candle
+Minimalist ML framework for Rust

crates/vendor/candle/candle-core/benches/bench_main.rs

@@ -0,0 +1,14 @@
+mod benchmarks;
+
+use criterion::criterion_main;
+
+criterion_main!(
+    benchmarks::affine::benches,
+    benchmarks::matmul::benches,
+    benchmarks::random::benches,
+    benchmarks::reduce::benches,
+    benchmarks::where_cond::benches,
+    benchmarks::conv_transpose2d::benches,
+    benchmarks::qmatmul::benches,
+    benchmarks::unary::benches
+);

crates/vendor/candle/candle-core/benches/benchmarks/affine.rs

@@ -0,0 +1,43 @@
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
+use std::time::Instant;
+
+fn run(a: &Tensor) {
+    a.affine(12.34, 56.78).unwrap();
+}
+
+fn run_affine_benchmark(c: &mut Criterion, device: &Device, dtype: DType, name: &str) {
+    let b = 1;
+    let m = 1024;
+    let k = 1024;
+
+    let tensor = Tensor::zeros((b, m, k), dtype, device).unwrap();
+
+    let flops = b * m * k * dtype.size_in_bytes();
+
+    let mut group = c.benchmark_group(device.bench_name(name));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box(&tensor));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        run_affine_benchmark(c, &device, DType::F32, "affine_f32");
+        run_affine_benchmark(c, &device, DType::F16, "affine_f16");
+        run_affine_benchmark(c, &device, DType::BF16, "affine_bf16");
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);

crates/vendor/candle/candle-core/benches/benchmarks/conv_transpose2d.rs

@@ -0,0 +1,59 @@
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
+use std::time::Instant;
+
+fn run(
+    x: &Tensor,
+    k: &Tensor,
+    padding: usize,
+    output_padding: usize,
+    stride: usize,
+    dilation: usize,
+) {
+    x.conv_transpose2d(k, padding, output_padding, stride, dilation)
+        .unwrap();
+}
+
+fn run_benchmark(c: &mut Criterion, device: &Device, dtype: DType, name: &str) {
+    let t = Tensor::arange(0.0f32, 10000.0, device)
+        .unwrap()
+        .reshape((1, 4, 50, 50))
+        .unwrap()
+        .to_dtype(dtype)
+        .unwrap();
+
+    let kernel = Tensor::arange(0.0f32, 100.0, device)
+        .unwrap()
+        .reshape((4, 1, 5, 5))
+        .unwrap()
+        .to_dtype(dtype)
+        .unwrap();
+
+    let flops = t.dims().iter().product::<usize>() * dtype.size_in_bytes();
+
+    let mut group = c.benchmark_group(device.bench_name(name));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box(&t), black_box(&kernel), 1, 0, 1, 2);
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        run_benchmark(c, &device, DType::F32, "conv_transpose2d_f32");
+        run_benchmark(c, &device, DType::F16, "conv_transpose2d_f16");
+        run_benchmark(c, &device, DType::BF16, "conv_transpose2d_bf16");
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);

crates/vendor/candle/candle-core/benches/benchmarks/matmul.rs

@@ -0,0 +1,44 @@
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
+use std::time::Instant;
+
+fn run(a: &Tensor, b: &Tensor) {
+    a.matmul(&b.t().unwrap()).unwrap();
+}
+
+fn run_bench(c: &mut Criterion, device: &Device) {
+    let b = 1;
+    let m = 1;
+    let n = 2048;
+    let k = 2048;
+
+    let dtype = DType::F32;
+    let lhs = Tensor::zeros((b, m, k), dtype, device).unwrap();
+    let rhs = Tensor::zeros((b, n, k), dtype, device).unwrap();
+
+    let flops = b * m * n * k;
+
+    let mut group = c.benchmark_group(device.bench_name("matmul"));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box(&lhs), black_box(&rhs));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        run_bench(c, &device);
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);

crates/vendor/candle/candle-core/benches/benchmarks/mod.rs

@@ -0,0 +1,70 @@
+pub(crate) mod affine;
+pub(crate) mod conv_transpose2d;
+pub(crate) mod matmul;
+pub(crate) mod qmatmul;
+pub(crate) mod random;
+pub(crate) mod reduce;
+pub(crate) mod unary;
+pub(crate) mod where_cond;
+
+use candle_core::{Device, Result};
+
+pub(crate) trait BenchDevice {
+    fn sync(&self) -> Result<()>;
+
+    fn bench_name<S: Into<String>>(&self, name: S) -> String;
+}
+
+impl BenchDevice for Device {
+    fn sync(&self) -> Result<()> {
+        match self {
+            Device::Cpu => Ok(()),
+            Device::Cuda(device) => {
+                #[cfg(feature = "cuda")]
+                return Ok(device.synchronize()?);
+                #[cfg(not(feature = "cuda"))]
+                panic!("Cuda device without cuda feature enabled: {:?}", device)
+            }
+            Device::Metal(device) => {
+                #[cfg(feature = "metal")]
+                return Ok(device.wait_until_completed()?);
+                #[cfg(not(feature = "metal"))]
+                panic!("Metal device without metal feature enabled: {:?}", device)
+            }
+        }
+    }
+
+    fn bench_name<S: Into<String>>(&self, name: S) -> String {
+        match self {
+            Device::Cpu => {
+                let cpu_type = if cfg!(feature = "accelerate") {
+                    "accelerate"
+                } else if cfg!(feature = "mkl") {
+                    "mkl"
+                } else {
+                    "cpu"
+                };
+                format!("{}_{}", cpu_type, name.into())
+            }
+            Device::Cuda(_) => format!("cuda_{}", name.into()),
+            Device::Metal(_) => format!("metal_{}", name.into()),
+        }
+    }
+}
+
+struct BenchDeviceHandler {
+    devices: Vec<Device>,
+}
+
+impl BenchDeviceHandler {
+    pub fn new() -> Result<Self> {
+        let mut devices = Vec::new();
+        if cfg!(feature = "metal") {
+            devices.push(Device::new_metal(0)?);
+        } else if cfg!(feature = "cuda") {
+            devices.push(Device::new_cuda(0)?);
+        }
+        devices.push(Device::Cpu);
+        Ok(Self { devices })
+    }
+}

crates/vendor/candle/candle-core/benches/benchmarks/qmatmul.rs

@@ -0,0 +1,72 @@
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{
+    quantized::{self, GgmlDType, QMatMul},
+    Device, Module, Tensor,
+};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
+use std::time::Instant;
+
+fn run(matmul: &QMatMul, x: &Tensor) {
+    matmul.forward(x).unwrap();
+}
+
+fn run_bench(c: &mut Criterion, device: &Device, dtype: GgmlDType) {
+    let b = 1;
+    let m = 1;
+    let n = 1024;
+    let k = 1024;
+
+    let lhs = (0..(m * k))
+        .map(|v| v as f32 / (m * k) as f32)
+        .collect::<Vec<_>>();
+    let rhs = (0..(k * n))
+        .map(|v| v as f32 / (n * k) as f32)
+        .collect::<Vec<_>>();
+
+    let lhs = Tensor::from_slice(&lhs, (m, k), device).unwrap();
+    let rhs = Tensor::from_slice(&rhs, (k, n), device).unwrap();
+
+    let qtensor = quantized::QTensor::quantize(&rhs.t().unwrap(), dtype).unwrap();
+    let matmul = quantized::QMatMul::from_qtensor(qtensor).unwrap();
+
+    let flops = b * m * n * k;
+
+    let mut group = c.benchmark_group(device.bench_name(format!("qmatmul_{:?}", dtype)));
+    group.sample_size(200);
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box(&matmul), black_box(&lhs));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        for dtype in [
+            GgmlDType::F32,
+            GgmlDType::F16,
+            GgmlDType::Q4_0,
+            GgmlDType::Q4_1,
+            GgmlDType::Q5_0,
+            GgmlDType::Q5_1,
+            GgmlDType::Q8_0,
+            GgmlDType::Q2K,
+            GgmlDType::Q3K,
+            GgmlDType::Q4K,
+            GgmlDType::Q5K,
+            GgmlDType::Q6K,
+        ] {
+            run_bench(c, &device, dtype);
+        }
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);

crates/vendor/candle/candle-core/benches/benchmarks/random.rs

@@ -0,0 +1,63 @@
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
+use std::time::Instant;
+
+fn rand_uniform(a: &Tensor) {
+    a.rand_like(-1.0, 123.0).unwrap();
+}
+
+fn rand_normal(a: &Tensor) {
+    a.randn_like(100.0, 15.0).unwrap();
+}
+
+fn run_random_bench(c: &mut Criterion, device: &Device) {
+    let b = 1;
+
+    let rows = 2048;
+    let cols = 2048;
+
+    let dtype = DType::F32;
+    let tensor = Tensor::zeros((b, rows, cols), dtype, device).unwrap();
+
+    let flops = b * rows * cols * dtype.size_in_bytes();
+
+    let mut group = c.benchmark_group(device.bench_name("random_uniform"));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |benches| {
+        benches.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                rand_uniform(black_box(&tensor));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+
+    let tensor = Tensor::zeros((b, rows, cols), dtype, device).unwrap();
+
+    let mut group = c.benchmark_group(device.bench_name("random_normal"));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |benches| {
+        benches.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                rand_normal(black_box(&tensor));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        run_random_bench(c, &device);
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);

crates/vendor/candle/candle-core/benches/benchmarks/reduce.rs

@@ -0,0 +1,158 @@
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{DType, Device, Tensor};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
+use half::{bf16, f16};
+use std::time::Instant;
+
+fn run_sum(a: &Tensor) {
+    a.sum_keepdim(2).unwrap();
+}
+fn run_arg_min(a: &Tensor) {
+    a.argmin_keepdim(2).unwrap();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    let (lo, up) = (-1000.0f32, 1000.0f32);
+    for device in handler.devices {
+        run_reduce(c, &device, (lo, up), false);
+        run_reduce(c, &device, (f16::from_f32(lo), f16::from_f32(up)), false);
+        run_reduce(c, &device, (bf16::from_f32(lo), bf16::from_f32(up)), false);
+
+        run_arg_reduce(c, &device, (lo, up), false);
+        run_arg_reduce(c, &device, (f16::from_f32(lo), f16::from_f32(up)), false);
+        run_arg_reduce(c, &device, (bf16::from_f32(lo), bf16::from_f32(up)), false);
+
+        run_reduce(c, &device, (lo, up), true);
+        run_reduce(c, &device, (f16::from_f32(lo), f16::from_f32(up)), true);
+        run_reduce(c, &device, (bf16::from_f32(lo), bf16::from_f32(up)), true);
+
+        run_arg_reduce(c, &device, (lo, up), true);
+        run_arg_reduce(c, &device, (f16::from_f32(lo), f16::from_f32(up)), true);
+        run_arg_reduce(c, &device, (bf16::from_f32(lo), bf16::from_f32(up)), true);
+    }
+}
+
+fn run_reduce<T: candle_core::FloatDType>(
+    c: &mut Criterion,
+    device: &Device,
+    (lo, up): (T, T),
+    strided: bool,
+) {
+    let b = 1;
+    let m = 1024;
+    let k = 1024;
+
+    let a = if strided {
+        Tensor::rand(lo, up, (b, m, k), &device)
+            .unwrap()
+            .transpose(0, 2)
+            .unwrap()
+    } else {
+        Tensor::rand(lo, up, (b, m, k), &device).unwrap()
+    };
+
+    let flops = b * m * k * T::DTYPE.size_in_bytes();
+
+    let name = match T::DTYPE {
+        DType::F32 => {
+            if strided {
+                "reduce_f32_strided"
+            } else {
+                "reduce_f32"
+            }
+        }
+        DType::F16 => {
+            if strided {
+                "reduce_f16_strided"
+            } else {
+                "reduce_f16"
+            }
+        }
+        DType::BF16 => {
+            if strided {
+                "reduce_bf16_strided"
+            } else {
+                "reduce_bf16"
+            }
+        }
+        _ => "unknown",
+    };
+
+    let mut group = c.benchmark_group(device.bench_name(name));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run_sum(black_box(&a));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+fn run_arg_reduce<T: candle_core::FloatDType>(
+    c: &mut Criterion,
+    device: &Device,
+    (lo, up): (T, T),
+    strided: bool,
+) {
+    let b = 1;
+    let m = 1024;
+    let k = 1024;
+
+    let a = if strided {
+        Tensor::rand(lo, up, (b, m, k), &device)
+            .unwrap()
+            .transpose(0, 2)
+            .unwrap()
+    } else {
+        Tensor::rand(lo, up, (b, m, k), &device).unwrap()
+    };
+
+    let flops = b * m * k * T::DTYPE.size_in_bytes();
+
+    let name = match T::DTYPE {
+        DType::F32 => {
+            if strided {
+                "arg_reduce_f32_strided"
+            } else {
+                "arg_reduce_f32"
+            }
+        }
+        DType::F16 => {
+            if strided {
+                "arg_reduce_f16_strided"
+            } else {
+                "arg_reduce_f16"
+            }
+        }
+        DType::BF16 => {
+            if strided {
+                "arg_reduce_bf16_strided"
+            } else {
+                "arg_reduce_bf16"
+            }
+        }
+        _ => "unknown",
+    };
+
+    let mut group = c.benchmark_group(device.bench_name(name));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run_arg_min(black_box(&a));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);