speedup

.idea/dictionaries/bt.xml

@@ -31,6 +31,7 @@
       <w>hifigan</w>
       <w>hparams</w>
       <w>huggingface</w>
+      <w>kasisnu</w>
       <w>lipsync</w>
       <w>mocap</w>
       <w>mochi</w>
@@ -47,6 +48,7 @@
       <w>techcrunch</w>
       <w>transcoded</w>
       <w>transcodings</w>
+      <w>unet</w>
       <w>unpermitted</w>
       <w>upserting</w>
       <w>upserts</w>

app_config.yaml

@@ -0,0 +1,7 @@
+image_width: 512
+image_height: 512
+#scheduler_steps: 15
+# NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
+scheduler_steps: 2
+scheduler_samples: 1
+

crates/desktop/src-tauri/Cargo.toml

@@ -18,11 +18,30 @@ crate-type = ["staticlib", "cdylib", "rlib"]
 tauri-build = { version = "2.0.4", features = [] }
 
 [dependencies]
-serde_json = "1.0"
-serde = { version = "1.0", features = ["derive"] }
-log = "0.4"
-tauri = { version = "2.2.4", features = [] }
-tauri-plugin-log = "2.0.0-rc"
-once_cell = "1.20.3"
-image = "0.25.5"
+
+# Internal
+filesys = { path = "../../lib/files/filesys" }
+memory_caching = { path = "../../lib/caching/memory_caching" }
+
+# External
+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"] }
+hf-hub = "0.3.2"
+image = "0.25.5"
+serde_yaml = "0.9"
+log = "0.4"
+once_cell = "1.20.3"
+rand = "0.8"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+tauri = { version = "2.2.4", features = [] }
+tauri-plugin-log = "2"
+tokenizers = "0.15.2"
+
+[features]
+default = ["cuda"]
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]

crates/desktop/src-tauri/capabilities/default.json

@@ -6,6 +6,7 @@
     "main"
   ],
   "permissions": [
-    "core:default"
+    "core:default",
+    "log:default"
   ]
-}
+}

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

@@ -1,11 +1,11 @@
-use std::io::Cursor;
+use base64::prelude::BASE64_STANDARD;
 use base64::Engine;
-use base64::prelude::{BASE64_STANDARD, BASE64_URL_SAFE};
 use image::codecs::png::{CompressionType, FilterType, PngEncoder};
 use image::ImageReader;
+use std::io::Cursor;
 
 #[tauri::command]
-pub fn infer_image(image: &str) -> Result<String, String> {
+pub fn flip_image(image: &str) -> Result<String, String> {
   println!("infer_image called; processing image...");
 
   let bytes = BASE64_STANDARD.decode(image)
@@ -18,17 +18,19 @@ pub fn infer_image(image: &str) -> Result<String, String> {
     .map_err(|err| format!("Image decode error: {}", err))?;
 
   let rotated_image = image.rotate180();
-  
+
   let mut buffer = Vec::new(); // TODO: Preallocate, better data structure
   let mut writer = Cursor::new(&mut buffer);
+
+  let encoder = PngEncoder::new_with_quality
+    (&mut writer, CompressionType::Fast, FilterType::Adaptive);
   
-  let encoder = PngEncoder::new_with_quality(&mut writer, CompressionType::Fast, FilterType::Adaptive);
   rotated_image.write_with_encoder(encoder)
     .map_err(|err| format!("Image encoding error: {}", err))?;
 
   let encoded = BASE64_STANDARD.encode(&buffer);
-  
+
   println!("Encoded: {:?}", encoded.split_at(10).0);
-  
+
   Ok(encoded)
 }

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

@@ -0,0 +1,92 @@
+use crate::endpoints::sd::{run, Args};
+use crate::ml::model_cache::ModelCache;
+use crate::ml::prompt_cache::PromptCache;
+use crate::state::app_config::AppConfig;
+use crate::state::yaml_config::YamlConfig;
+use anyhow::Result;
+use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
+use bytes::{BufMut, BytesMut};
+use image::imageops::FilterType;
+use image::{DynamicImage, EncodableLayout, ImageFormat, ImageReader};
+use log::{error, info};
+use std::io::Cursor;
+use std::path::PathBuf;
+use tauri::State;
+
+const PROMPT: &str = "A beautiful landscape with mountains and a lake";
+
+#[tauri::command]
+pub fn infer_image(
+  image: &str,
+  model_config: State<AppConfig>,
+  model_cache: State<ModelCache>,
+  prompt_cache: State<PromptCache>,
+) -> Result<String, String> {
+
+  let bytes = BASE64_STANDARD.decode(image)
+    .map_err(|err| format!("Base64 decode error: {}", err))?;
+  
+  let prompt_file = PathBuf::from("prompt.txt").canonicalize()
+    .unwrap_or_else(|_| PathBuf::from("prompt.txt"));
+  
+  let prompt = std::fs::read_to_string(&prompt_file)
+    .map_err(|err| format!("Failed to read prompt file: {}", err))
+    .unwrap_or_else(|_| {
+      error!("Failed to read prompt file: {:?}", prompt_file);
+      PROMPT.to_string()
+    })
+    .trim()
+    .to_string();
+
+  let image = ImageReader::new(Cursor::new(bytes))
+    .with_guessed_format()
+    .map_err(|err| format!("Image format error: {}", err))?
+    .decode()
+    .map_err(|err| format!("Image decode error: {}", err))?;
+  
+  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
+  let image = image.resize(512, 512, FilterType::CatmullRom);
+  
+  let result = do_infer_image(&prompt, image, &model_config, &model_cache, prompt_cache);
+  
+  if let Err(err) = result.as_deref() {
+    error!("There was an error: {:?}", err);
+  }
+  
+  result
+}
+
+fn do_infer_image(
+  prompt: &str,
+  image: DynamicImage,
+  config: &AppConfig,
+  model_cache: &ModelCache,
+  prompt_cache: State<PromptCache>,
+) -> Result<String, String> {
+  println!("infer_image called; generating image with SDXL Turbo...");
+  
+  let args = Args {
+    image: &image,
+    prompt: prompt.to_string(),
+    uncond_prompt: "".to_string(),
+    //guidance_scale: Some(0.0),
+    guidance_scale: None,
+    model_cache,
+    configs: config,
+    prompt_cache: &prompt_cache,
+  };
+
+  match run(args) {
+    Ok(image) => {
+      let mut bytes = Vec::with_capacity(1024*1024);
+      
+      image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
+        .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+      let bytes = BASE64_STANDARD.encode(bytes);
+
+      Ok(bytes)
+    },
+    Err(e) => Err(format!("Failed to generate image: {}", e)),
+  }
+}

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

@@ -0,0 +1,4 @@
+pub mod flip_image;
+pub mod image_endpoint;
+pub mod sd;
+pub mod test_counter;

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

@@ -0,0 +1,172 @@
+use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
+use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
+use crate::ml::infer_clip_text_embeddings::infer_clip_text_embeddings;
+use crate::ml::model_cache::ModelCache;
+use crate::ml::prompt_cache::PromptCache;
+use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
+use crate::state::app_config::AppConfig;
+use anyhow::{Error as E, Result};
+use candle_core::{DType, IndexOp, D};
+use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
+use image::DynamicImage;
+use log::info;
+use rand::Rng;
+
+// TODO: Clean up
+
+// Note for Kasisnu: I'm going to start using lifetimes as long as that doesn't slow your velocity
+// Basically the args o this telescopic args struct are guaranteed to live as long as the struct
+// itself with the 'a lifetime.
+pub struct Args<'a> {
+    pub image: &'a DynamicImage,
+    pub prompt: String,
+    pub uncond_prompt: String,
+    pub guidance_scale: Option<f64>,
+    pub configs: &'a AppConfig,
+    pub model_cache: &'a ModelCache,
+    pub prompt_cache: &'a PromptCache,
+}
+
+pub fn run(args: Args<'_>) -> Result<RgbImage> {
+    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 dimensions: {}x{}", args.configs.sd_config.width, args.configs.sd_config.height);
+
+    println!("Building scheduler...");
+    let mut scheduler = args.configs.sd_config.build_scheduler(
+        args.configs.scheduler_steps)?;
+
+    let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
+
+    println!("Using seed: {}", seed);
+    args.configs.device.set_seed(seed)?;
+
+    info!("Checking if prompt is cached");
+    let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?;
+    
+    let text_embeddings = if let Some(tensor) = maybe_cached {
+        info!("Prompt is cached!");
+        tensor
+    } else {
+        info!("Prompt is NOT cached! Calculating embedding...");
+        let tensor = infer_clip_text_embeddings(
+            &args.prompt,
+            &args.uncond_prompt,
+            None, // tokenizer
+            None, // clip_weights
+            None, // clip2_weights
+            args.configs.sd_version,
+            &args.configs.sd_config,
+            false, // use_f16
+            &args.configs.device,
+            args.configs.dtype,
+            false, // use_guide_scale
+        )?;
+        args.prompt_cache.store_copy(&args.prompt, &tensor)?;
+        tensor
+    };
+    
+    println!("Text embeddings shape: {:?}", text_embeddings.shape());
+
+    println!("Loading input image into tensor...");
+
+    let input_image = dynamic_image_to_tensor(
+        args.image,
+        &args.configs.device,
+        args.configs.dtype)?;
+
+    println!("Reference image shape: {:?}", input_image.shape());
+
+    let vae_scale = get_vae_scale(args.configs.sd_version);
+
+    let init_latent_dist : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?;
+
+    // TODO(bt,2025-02-18): This takes a little bit to generate.
+    println!("Generating latents from input image...");
+    let latents = (init_latent_dist.sample()? * vae_scale)?;
+    
+    //.to_device(&args.configs.device)?;
+    
+    println!("Initial latents shape: {:?}", latents.shape());
+
+    println!("Calculating start step for diffusion process...");
+    let img2img_strength = 0.75f64;
+
+    let t_start = {
+        let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize;
+
+        println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength);
+        start
+    };
+
+    let latents = if t_start < scheduler.timesteps().len() {
+        println!("Adding noise to latents...");
+        let noise = latents.randn_like(0f64, 1f64)?;
+        scheduler.add_noise(&latents, noise, scheduler.timesteps()[t_start])?
+    } else {
+        println!("Using latents directly (no noise addition needed)");
+        latents
+    };
+
+    println!("Latents initialized successfully");
+
+    let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
+
+    // Initial noise shape: [1, 4, 64, 64]
+    println!("Initial noise shape: {:?}", latents.shape());
+
+    println!("Starting diffusion process...");
+    
+    let mut latents = latents;
+
+    for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+        if timestep_index < t_start {
+            continue;
+        }
+
+        println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
+        
+        // TODO(bt,2025-08-18): Do we have to clone the tensor?
+        let latent_model_input = scheduler.scale_model_input(latents.clone(), timestep)?;
+        
+        let noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) {
+            Ok(pred) => {
+                pred
+            },
+            Err(e) => {
+                println!("UNet inference failed with error: {}", e);
+                return Err(anyhow::anyhow!("UNet inference failed: {}", e));
+            },
+        };
+
+        latents = scheduler.step(&noise_pred, timestep, &latents)?;
+    }
+
+    println!("Diffusion process completed, decoding image...");
+    let image = args.model_cache.vae_decode(&(latents / vae_scale)?)?;
+    
+    println!("VAE decode completed");
+
+    println!("Scaling image back");
+    let image = ((image / 2.)? + 0.5)?;
+    
+    // TODO(bt,2025-08-18): This normalization is slow.
+    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/endpoints/test_counter.rs

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

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

@@ -1,55 +1,60 @@
-mod image_endpoint;
+pub mod endpoints;
+pub mod ml;
+pub mod state;
 
-use std::sync::{Arc, RwLock};
-use once_cell::sync::Lazy;
-use image_endpoint::infer_image;
+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 tauri_plugin_log::Target;
+use tauri_plugin_log::TargetKind;
 
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
+  // NB: Tauri wants to install the logger itself, so we can't rely on the logger crate 
+  // until the tauri runtime begins.
+  println!("Loading model config...");
+
+  let config = AppConfig::init()
+    .expect("config should load");
+
+  let prompt_cache = PromptCache::with_capacity(8)
+    .expect("prompt cache should load");
+
+  println!("Creating model cache...");
+  
+  let model_cache = ModelCache::new(
+    config.device.clone(),
+    config.dtype,
+    config.sd_version.clone(),
+    config.sd_config.clone()
+  ).expect("Model cache should create");
+
+  println!("Initializing backend runtime...");
+
   tauri::Builder::default()
+    .plugin(tauri_plugin_log::Builder::new()
+      .level(log::LevelFilter::Info)
+      .targets(vec![Target::new(TargetKind::Stdout)])
+      .build())
     .setup(|app| {
-      if cfg!(debug_assertions) {
-        app.handle().plugin(
-          tauri_plugin_log::Builder::default()
-            .level(log::LevelFilter::Info)
-            .build(),
-        )?;
-      }
+      //if cfg!(debug_assertions) {
+      //  app.handle().plugin(
+      //    tauri_plugin_log::Builder::default()
+      //      .level(log::LevelFilter::Info)
+      //      .build(),
+      //  )?;
+      //}
       Ok(())
     })
+    .manage(config)
+    .manage(prompt_cache)
+    .manage(model_cache)
     .invoke_handler(tauri::generate_handler![
-      my_custom_command, 
-      test_round_trip, 
+      test_counter, 
       infer_image,
     ])
     .run(tauri::generate_context!())
     .expect("error while running tauri application");
 }
-
-#[tauri::command]
-fn my_custom_command() {
-  println!("I was invoked from JavaScript!");
-}
-
-static COUNTER : Lazy<Arc<RwLock<u64>>> = Lazy::new(|| Arc::new(RwLock::new(0)));
-
-#[tauri::command]
-fn test_round_trip() -> u64 {
-  let value : u64;
-  {
-    match COUNTER.write() {
-      Ok(mut counter) => {
-        *counter += 1;
-        value = *counter;
-      },
-      Err(_e) => {
-        value = 0;
-      }
-    }
-  }
-
-  println!("I was invoked from JavaScript! {:?}", value);
-
-  value
-}
-

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

@@ -0,0 +1,80 @@
+use std::path::PathBuf;
+use crate::ml::model_file::StableDiffusionVersion;
+use crate::ml::models::lazy_load_vae_model::LazyLoadVaeModel;
+use candle_core::{DType, Device, Tensor};
+use crate::ml::model_cache::ModelCache;
+
+/// Generates the mask latents, scaled mask and mask_4 for inpainting. 
+/// Returns a tuple of None if inpainting is not being used.
+#[allow(clippy::too_many_arguments)]
+pub fn create_inpainting_tensors(
+  sd_version: StableDiffusionVersion,
+  mask_path: Option<PathBuf>,
+  dtype: DType,
+  device: &Device,
+  use_guide_scale: bool,
+  model_cache: &ModelCache,
+  image: Option<Tensor>,
+  vae_scale: f64,
+) -> anyhow::Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> {
+  match sd_version {
+    StableDiffusionVersion::XlInpaint
+    | StableDiffusionVersion::V2Inpaint
+    | StableDiffusionVersion::V1_5Inpaint => {
+      let inpaint_mask = mask_path.ok_or_else(|| {
+        anyhow::anyhow!("An inpainting model was requested but mask-path is not provided.")
+      })?;
+      // Get the mask image with shape [1, 1, 128, 128]
+      let mask = mask_preprocess(inpaint_mask)?
+        .to_device(device)?
+        .to_dtype(dtype)?;
+      // Generate the masked image from the image and the mask with shape [1, 3, 1024, 1024]
+      let xmask = mask.le(0.5)?.repeat(&[1, 3, 1, 1])?.to_dtype(dtype)?;
+      let image = &image
+        .ok_or_else(|| anyhow::anyhow!(
+                    "An inpainting model was requested but img2img which is used as the input image is not provided."
+                ))?;
+      let masked_img = (image * xmask)?;
+      // Scale down the mask
+      let shape = masked_img.shape();
+      let (w, h) = (shape.dims()[3] / 8, shape.dims()[2] / 8);
+      let mask = mask.interpolate2d(w, h)?;
+      // shape: [1, 4, 128, 128]
+      let mask_latents = model_cache.vae_encode(&masked_img)?;
+      let mask_latents = (mask_latents.sample()? * vae_scale)?.to_device(device)?;
+
+      let mask_4 = mask.as_ref().repeat(&[1, 4, 1, 1])?;
+      let (mask_latents, mask) = if use_guide_scale {
+        (
+          Tensor::cat(&[&mask_latents, &mask_latents], 0)?,
+          Tensor::cat(&[&mask, &mask], 0)?,
+        )
+      } else {
+        (mask_latents, mask)
+      };
+      Ok((Some(mask_latents), Some(mask), Some(mask_4)))
+    }
+    _ => Ok((None, None, None)),
+  }
+}
+
+fn mask_preprocess<T: AsRef<std::path::Path>>(path: T) -> anyhow::Result<Tensor> {
+  let img = image::open(path)?.to_luma8();
+  let (new_width, new_height) = {
+    let (width, height) = img.dimensions();
+    (width - width % 32, height - height % 32)
+  };
+  let img = image::imageops::resize(
+    &img,
+    new_width,
+    new_height,
+    image::imageops::FilterType::CatmullRom,
+  )
+    .into_raw();
+  let mask = Tensor::from_vec(img, (new_height as usize, new_width as usize), &Device::Cpu)?
+    .unsqueeze(0)?
+    .to_dtype(DType::F32)?
+    //.div(255.0)? // TODO(bt): Need to figure out what changed in candle to no longer permit division by f64
+    .unsqueeze(0)?;
+  Ok(mask)
+}

crates/desktop/src-tauri/src/ml/image/dynamic_image_to_tensor.rs

@@ -0,0 +1,17 @@
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+use image::DynamicImage;
+
+pub fn dynamic_image_to_tensor(image: &DynamicImage, device: &Device, datatype: DType) -> Result<Tensor> {
+  let (height, width) = (image.height() as usize, image.width() as usize);
+  let height = height - height % 32;
+  let width = width - width % 32;
+  let image = image.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
+  let image = image.to_rgb8();
+  let image = image.into_raw();
+  Ok(Tensor::from_vec(image, (height, width, 3), device)?
+    .permute((2, 0, 1))?
+    .to_dtype(datatype)?
+    .affine(2. / 255., -1.)?
+    .unsqueeze(0)?)
+}

crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor.rs

@@ -0,0 +1,28 @@
+use candle_core::{DType, Device, Tensor};
+
+pub fn load_image_file_to_tensor<T: AsRef<std::path::Path>>(path: T, device: &Device) -> anyhow::Result<Tensor> {
+  let img = image::ImageReader::open(path)?.decode()?;
+  let (height, width) = (img.height() as usize, img.width() as usize);
+  println!("INPUT IMAGE WIDTH AND HEIGHT: {}x{}", width, height);
+
+  // TODO: Downscale from 1024 -- YES! THIS HACK DOES THE TRICK OF GETTING THE TENSOR SHAPE TO [1,4,64,64] !!
+  let height = height / 2;
+  let width = width  / 2;
+  
+  let height = height - height % 32;
+  let width = width - width % 32;
+  let img = img.resize_to_fill(
+    width as u32,
+    height as u32,
+    image::imageops::FilterType::CatmullRom,
+  );
+  println!("RESIZED IMAGE WIDTH AND HEIGHT: {}x{}", width, height);
+  let img = img.to_rgb8();
+  let img = img.into_raw();
+  let img = Tensor::from_vec(img, (height, width, 3), device)?
+    .permute((2, 0, 1))?
+    .to_dtype(DType::F32)?
+    .affine(2. / 255., -1.)?
+    .unsqueeze(0)?;
+  Ok(img)
+}

crates/desktop/src-tauri/src/ml/image/load_image_file_to_tensor_2.rs

@@ -0,0 +1,18 @@
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+
+pub fn load_image_file_to_tensor_2<T: AsRef<std::path::Path>>(path: T) -> Result<Tensor> {
+  let img = image::ImageReader::open(path)?.decode()?;
+  let (height, width) = (img.height() as usize, img.width() as usize);
+  let height = height - height % 32;
+  let width = width - width % 32;
+  let img = img.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
+  let img = img.to_rgb8();
+  let img = img.into_raw();
+  let img = Tensor::from_vec(img, (height, width, 3), &Device::Cpu)?
+    .permute((2, 0, 1))?
+    .to_dtype(DType::F32)?
+    .affine(2. / 255., -1.)?
+    .unsqueeze(0)?;
+  Ok(img)
+}

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

@@ -0,0 +1,5 @@
+pub mod dynamic_image_to_tensor;
+pub mod load_image_file_to_tensor;
+pub mod load_image_file_to_tensor_2;
+pub mod save_image_from_tensor;
+pub mod tensor_to_image_buffer;

crates/desktop/src-tauri/src/ml/image/save_image_from_tensor.rs

@@ -0,0 +1,28 @@
+use anyhow::{Error as E, Result};
+use candle_core::Tensor;
+use image::{ImageBuffer, Rgb};
+
+pub fn save_image_from_tensor<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
+  println!("Saving image to {:?}", p.as_ref());
+  let p = p.as_ref();
+  let (channel, height, width) = img.dims3()?;
+  println!("Image dimensions: {}x{} with {} channels", width, height, channel);
+
+  if channel != 3 {
+    anyhow::bail!("save_image expects an input of shape (3, height, width)")
+  }
+  let img = img.permute((1, 2, 0))?.flatten_all()?;
+  let pixels = img.to_vec1::<u8>()?;
+  println!("Converting tensor to image buffer...");
+
+  let image: ImageBuffer<Rgb<u8>, Vec<u8>> =
+    match ImageBuffer::from_raw(width as u32, height as u32, pixels) {
+      Some(image) => image,
+      None => anyhow::bail!("error saving image {p:?}"),
+    };
+  println!("Successfully created image buffer");
+
+  image.save(p).map_err(E::from)?;
+  println!("Successfully saved image to disk");
+  Ok(())
+}

crates/desktop/src-tauri/src/ml/image/tensor_to_image_buffer.rs

@@ -0,0 +1,18 @@
+use candle_core::Tensor;
+use image::{ImageBuffer, Rgb};
+
+pub type RgbImage = ImageBuffer<Rgb<u8>, Vec<u8>>;
+
+pub fn tensor_to_image_buffer(img: &Tensor) -> anyhow::Result<RgbImage> {
+  let (channel, height, width) = img.dims3()?;
+
+  if channel != 3 {
+    anyhow::bail!("save_image expects an input of shape (3, height, width)")
+  }
+
+  let img = img.permute((1, 2, 0))?.flatten_all()?;
+  let pixels = img.to_vec1::<u8>()?;
+
+  ImageBuffer::from_raw(width as u32, height as u32, pixels)
+    .ok_or_else(|| anyhow::anyhow!("error loading image from raw tensor"))
+}

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

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

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

@@ -0,0 +1,8 @@
+pub mod create_inpainting_tensors;
+pub mod image;
+pub mod infer_clip_text_embeddings;
+pub mod model_cache;
+pub mod model_file;
+pub mod models;
+pub mod prompt_cache;
+pub mod stable_diffusion;

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,22 @@
+use candle_core::Tensor;
+use memory_caching::arc_sieve::ArcSieve;
+
+pub struct PromptCache {
+  cache: ArcSieve<String, Tensor>,
+}
+
+impl PromptCache {
+  pub fn with_capacity(capacity: usize) -> anyhow::Result<Self> {
+    Ok(Self {
+      cache: ArcSieve::with_capacity(capacity)?,
+    })
+  }
+  
+  pub fn get_copy(&self, key: &String) -> anyhow::Result<Option<Tensor>> {
+    self.cache.get_copy(key)
+  }
+  
+  pub fn store_copy(&self, key: &String, value: &Tensor) -> anyhow::Result<()> {
+    self.cache.store_copy(key, value)
+  }
+}

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

@@ -0,0 +1,13 @@
+use crate::ml::model_file::StableDiffusionVersion;
+
+pub fn get_vae_scale(sd_version: StableDiffusionVersion) -> f64 {
+  match sd_version {
+    StableDiffusionVersion::V1_5
+    | StableDiffusionVersion::V1_5Inpaint
+    | StableDiffusionVersion::V2_1
+    | StableDiffusionVersion::V2Inpaint
+    | StableDiffusionVersion::XlInpaint
+    | StableDiffusionVersion::Xl => 0.18215,
+    StableDiffusionVersion::Turbo => 0.13025,
+  }
+}

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

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

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

@@ -0,0 +1,93 @@
+use crate::ml::model_file::StableDiffusionVersion;
+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;
+
+const CONFIG_FILENAME : &str = "app_config.yaml";
+
+const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
+const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
+
+const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
+const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;
+
+/// A central place to configure app-wide details.
+pub struct AppConfig {
+  /// We only support one device type currently, and we can't multiplex 
+  /// over several devices yet.
+  pub device: Device,
+  
+  /// The core numeric type to use for models.
+  pub dtype: DType,
+  
+  pub sd_version: StableDiffusionVersion,
+  pub sd_config: StableDiffusionConfig,
+
+  pub image_height: usize,
+  pub image_width: usize,
+  pub scheduler_steps: usize,
+  pub scheduler_samples: usize,
+  
+  pub seed: Option<u64>,
+  
+  /// Probably shouldn't be here.
+  pub hf_api: Api,
+}
+
+impl AppConfig {
+  
+  pub fn init() -> anyhow::Result<Self> {
+    println!("Configuration filename: {}", CONFIG_FILENAME);
+    let yaml_configs = load_yaml_configs();
+    
+    println!("Configs: {:?}", yaml_configs);
+    
+    let device = Device::new_cuda(0)
+      .unwrap_or_else(|e| {
+          info!("CUDA not available ({}), falling back to CPU", e);
+          Device::Cpu
+        });
+    
+    let sd_config = match SD_VERSION {
+      StableDiffusionVersion::Turbo => {
+        StableDiffusionConfig::sdxl_turbo(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      _ => {
+        StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+    };
+    
+    let hf_api = Api::new()?;
+    
+    Ok(Self {
+      device,
+      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
+      sd_version: SD_VERSION,
+      sd_config,
+      image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
+      image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
+      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(1),
+      scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
+      seed: yaml_configs.seed,
+      hf_api,
+    })
+  }
+}
+
+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/mod.rs

@@ -0,0 +1,2 @@
+pub mod yaml_config;
+pub mod app_config;

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

@@ -0,0 +1,24 @@
+use std::fs::File;
+use std::io::BufReader;
+use std::path::Path;
+
+/// 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.
+#[derive(serde::Deserialize, serde::Serialize, Clone, Debug, Default)]
+pub struct YamlConfig {
+  pub image_height: Option<usize>,
+  pub image_width: Option<usize>,
+  pub scheduler_steps: Option<usize>,
+  pub scheduler_samples: Option<usize>,
+  pub seed: Option<u64>,
+}
+
+impl YamlConfig {
+  pub fn read_from_file<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
+    let file = File::open(path)?;
+    let reader = BufReader::new(file);
+    let config = serde_yaml::from_reader(reader)?;
+    Ok(config)
+  }
+}

crates/lib/files/filesys/src/lib.rs

@@ -31,4 +31,5 @@ pub mod file_size;
 pub mod filename_concat;
 pub mod is_filesystem_full_error;
 pub mod path_to_string;
+pub mod recursively_find_file_by_name;
 pub mod rename_across_devices;

crates/lib/files/filesys/src/recursively_find_file_by_name.rs

@@ -0,0 +1,56 @@
+use std::path::{Path, PathBuf};
+use walkdir::WalkDir;
+use errors::AnyhowResult;
+
+pub fn recursively_find_file_by_name<P: AsRef<Path>, Q: AsRef<Path>>(
+  search_filename: P,
+  search_root: Q,
+  max_depth: usize,
+) -> AnyhowResult<Option<PathBuf>> {
+
+  let walker = WalkDir::new(search_root)
+    .max_depth(max_depth);
+
+  for entry in walker {
+    let entry = entry?;
+    if entry.path().is_dir() {
+      let check_path = entry.path().join(search_filename.as_ref());
+      if check_path.exists() && check_path.is_file() {
+        return Ok(Some(check_path));
+      }
+    }
+  }
+
+  Ok(None)
+}
+
+#[cfg(test)]
+mod tests {
+  use testing::test_file_path::test_file_path;
+  use crate::recursively_find_file_by_name::recursively_find_file_by_name;
+
+  #[test]
+  fn test_path_exists() {
+    let test_data_path = test_file_path("test_data/").expect("should work");
+    let result = recursively_find_file_by_name(
+      "super_mario_rpg_beware_the_forests_mushrooms.mp3",
+      test_data_path,
+      5
+    );
+    let result = result.expect("should be ok");
+    let result = result.expect("should exist");
+    assert!(result.ends_with("test_data/audio/mp3/super_mario_rpg_beware_the_forests_mushrooms.mp3"));
+  }
+
+  #[test]
+  fn test_path_does_not_exist() {
+    let test_data_path = test_file_path("test_data/").expect("should work");
+    let result = recursively_find_file_by_name(
+      "fake_name.foo",
+      test_data_path,
+      5
+    );
+    let result = result.expect("should be ok");
+    assert!(result.is_none());
+  }
+}

script/build_production_realtime_desktop.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+# Tauri doesn't let you configure the frontend project directory statically, though they do provide an
+# environment variable to pass it to the CLI. Without doing this, the tauri cli randomly walks the
+# filesystem and finds the wrong frontend code.
+
+# TODO(bt,2025-02-13): This is not the correct way to get the root dir
+root_dir=$(pwd)
+export TAURI_FRONTEND_PATH="${root_dir}/frontend/realtime"
+
+cargo tauri build
+

script/launch_realtime_desktop.sh

@@ -8,4 +8,4 @@
 root_dir=$(pwd)
 export TAURI_FRONTEND_PATH="${root_dir}/frontend/realtime"
 
-cargo tauri dev
+cargo tauri dev --no-watch