SOMETHING IS HAPPENING

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

crates/desktop/src-tauri/Cargo.toml

@@ -26,3 +26,14 @@ tauri-plugin-log = "2.0.0-rc"
 once_cell = "1.20.3"
 image = "0.25.5"
 base64 = "0.22.1"
+candle-core = { version = "0.8.2", features = ["cuda"] }
+candle-transformers = { version = "0.8.2", features = ["cuda"] }
+candle-nn = { version = "0.8.2", features = ["cuda"] }
+tokenizers = "0.15.2"
+hf-hub = "0.3.2"
+anyhow = "1.0"
+rand = "0.8"
+
+[features]
+default = ["cuda"]
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]

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,78 @@
+use std::io::Cursor;
+use std::path::PathBuf;
+use crate::endpoints::sd::{run, Args};
+use crate::ml::model_cache::ModelCache;
+use crate::model_config::ModelConfig;
+use anyhow::Result;
+use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
+use image::ImageReader;
+use log::error;
+use tauri::State;
+
+const PROMPT: &str = "A beautiful landscape with mountains and a lake";
+
+const HEIGHT: usize = 512;
+const WIDTH: usize = 512;
+
+#[tauri::command]
+pub fn infer_image(image: &str, model_config: State<ModelConfig>, model_cache: State<ModelCache>) -> Result<String, String> {
+
+  let bytes = BASE64_STANDARD.decode(image)
+    .map_err(|err| format!("Base64 decode error: {}", err))?;
+
+  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-14): Use byte buffers instead of hitting the filesystem.
+  let image_path = PathBuf::from("input_image.png");
+  
+  image.save(&image_path)
+    .map_err(|err| format!("Failed to save image: {}", err))?;
+  
+  let result = do_infer_image(&image_path, &model_config, &model_cache);
+  
+  if let Err(err) = result.as_deref() {
+    error!("There was an error: {:?}", err);
+  }
+  
+  result
+}
+
+fn do_infer_image(image_path: &PathBuf, config: &ModelConfig, model_cache: &ModelCache) -> Result<String, String> {
+  println!("infer_image called; generating image with SDXL Turbo...");
+  
+  let args = Args {
+    image_path,
+    api: config.hf_api.clone(),
+    prompt: PROMPT.to_string(),
+    uncond_prompt: "".to_string(),
+    cpu: config.device.is_cpu(),
+    height: Some(HEIGHT),
+    width: Some(WIDTH),
+    n_steps: Some(1),
+    num_samples: 1,
+    seed: None,
+    sd_version: config.sd_version.clone(),
+    guidance_scale: Some(0.0),
+    model_cache,
+    model_configs: config,
+  };
+
+  match run(args) {
+    Ok(_) => {
+      let img_data = std::fs::read("temp.png")
+        .map_err(|e| format!("Failed to read generated image: {}", e))?;
+        
+      let encoded = BASE64_STANDARD.encode(&img_data);
+      println!("Generated image encoded successfully");
+
+      let _ = std::fs::remove_file("temp.png");
+
+      Ok(encoded)
+    }
+    Err(e) => Err(format!("Failed to generate image: {}", e))
+  }
+}

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

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

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

@@ -0,0 +1,266 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+use std::mem::forget;
+use std::path::Path;
+use std::sync::{Arc, LockResult, RwLock, RwLockReadGuard, RwLockWriteGuard};
+use candle_transformers::models::stable_diffusion;
+
+use crate::ml::infer_clip_text_embeddings::infer_clip_text_embeddings;
+use crate::ml::model_file::StableDiffusionVersion;
+use crate::ml::save_image_from_tensor::save_image_from_tensor;
+use anyhow::{Error as E, Result};
+use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
+use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
+use hf_hub::api::sync::Api;
+use log::info;
+use once_cell::sync::Lazy;
+use rand::Rng;
+use crate::ml::create_inpainting_tensors::create_inpainting_tensors;
+use crate::ml::load_image_file_to_tensor::load_image_file_to_tensor;
+use crate::ml::model_cache::ModelCache;
+use crate::ml::models::unet_model::UNetModel;
+use crate::model_config::ModelConfig;
+// 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_path: &'a Path,
+    pub api: Api,
+    pub prompt: String,
+    pub uncond_prompt: String,
+    pub cpu: bool,
+    pub height: Option<usize>,
+    pub width: Option<usize>,
+    pub n_steps: Option<usize>,
+    pub num_samples: usize,
+    pub seed: Option<u64>,
+    pub sd_version: StableDiffusionVersion,
+    pub guidance_scale: Option<f64>,
+    pub model_cache: &'a ModelCache,
+    pub model_configs: &'a ModelConfig,
+}
+
+pub fn run(args: Args<'_>) -> Result<()> {
+    println!("Starting image generation with the following configuration:");
+    println!("  Model: {:?}", args.sd_version);
+    println!("  Prompt: {}", args.prompt);
+    println!("  Steps: {}", args.n_steps.unwrap_or(1));
+    println!("  Using CPU: {}", args.cpu);
+
+
+    println!("Model dimensions: {}x{}", args.model_configs.sd_config.width, args.model_configs.sd_config.height);
+
+    println!("Building scheduler...");
+    let mut scheduler = args.model_configs.sd_config.build_scheduler(args.n_steps.unwrap_or(1))?;
+
+    let seed = args.seed.unwrap_or_else(|| rand::thread_rng().gen());
+    println!("Using seed: {}", seed);
+
+    args.model_configs.device.set_seed(seed)?;
+
+    println!("Initializing Hugging Face API...");
+
+    let repo = args.sd_version.repo();
+    println!("Downloading model files from: {}", repo);
+    
+    //println!("Downloading VAE ...");
+    //let vae_file = args.api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
+
+    //println!("VAE Path: {:?}", &vae_file);
+
+    //println!("Downloading text encoders...");
+    //let clip_file = args.api.model(repo.to_string()).get("text_encoder/model.safetensors")?;
+    //let clip2_file = args.api.model(repo.to_string()).get("text_encoder_2/model.safetensors")?;
+    //let tokenizer = args.api.model("openai/clip-vit-large-patch14".to_string()).get("tokenizer.json")?;
+    //let tokenizer2 = args.api.model("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k".to_string()).get("tokenizer.json")?;
+
+    //println!("Building VAE model...");
+    //let vae = sd_config.build_vae(vae_file, &device, dtype)?;
+
+    //// Build text encoders
+    //println!("Building text encoders...");
+    //let text_model = stable_diffusion::build_clip_transformer(&sd_config.clip, clip_file, &device, dtype)?;
+    //let text_model2 = stable_diffusion::build_clip_transformer(sd_config.clip2.as_ref().unwrap(), clip2_file, &device, dtype)?;
+    //
+    //// Tokenize and encode text
+    //println!("Processing text prompt...");
+    //let tokenizer = tokenizers::Tokenizer::from_file(tokenizer).map_err(E::msg)?;
+    //let tokenizer2 = tokenizers::Tokenizer::from_file(tokenizer2).map_err(E::msg)?;
+    
+    println!("Preparing text embeddings...");
+    let which = match args.sd_version {
+        StableDiffusionVersion::Xl
+        | StableDiffusionVersion::XlInpaint
+        | StableDiffusionVersion::Turbo => vec![true, false],
+        _ => vec![true],
+    };
+    
+    let text_embeddings = which
+        .iter()
+        .map(|first| {
+            infer_clip_text_embeddings(
+                &args.prompt,
+                &args.uncond_prompt,
+                None, // tokenizer
+                None, // clip_weights
+                None, // clip2_weights
+                args.sd_version,
+                &args.model_configs.sd_config,
+                false, // use_f16
+                &args.model_configs.device,
+                args.model_configs.dtype,
+                false, // use_guide_scale
+                *first,
+            )
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
+    println!("Text embeddings shape: {:?}", text_embeddings.shape());
+
+    println!("Loading input image into tensor...");
+
+    let input_image = load_image_file_to_tensor(args.image_path, &args.model_configs.device)?;
+
+    // REFERENCE IMAGE shape: [1, 3, 1024, 1024]
+    println!("REFERENCE IMAGE shape: {:?}", input_image.shape());
+
+    let use_guide_scale = false; // TODO
+
+    let vae_scale = match args.model_configs.sd_version {
+        StableDiffusionVersion::V1_5
+        | StableDiffusionVersion::V1_5Inpaint
+        | StableDiffusionVersion::V2_1
+        | StableDiffusionVersion::V2Inpaint
+        | StableDiffusionVersion::XlInpaint
+        | StableDiffusionVersion::Xl => 0.18215,
+        StableDiffusionVersion::Turbo => 0.13025,
+    };
+
+
+
+    let encoded_image : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?;
+
+    let encoded_sample = encoded_image.sample()?;
+
+    // Encoded sample shape: [1, 4, 128, 128] ... so close. Just needs to be smaller.
+    println!("Encoded sample shape: {:?}", encoded_sample.shape());
+
+    let mask_latents = (encoded_sample * vae_scale)?.to_device(&args.model_configs.device)?;
+
+    println!("Mask latents shape: {:?}", mask_latents.shape());
+
+
+    // TODO: Just to trip the conditionals that force mask/inpaint generation
+    // TODO: Set this to turbo to turn off inpaint, anything else will turn it on.
+    // TODO: THIS MIGHT NOT BE POSSIBLE WITH TURBO: https://huggingface.co/stabilityai/sdxl-turbo/discussions/7
+    // TODO/SEE ALSO: https://www.reddit.com/r/StableDiffusion/comments/198cl08/major_inpainting_issues_with_sdxl_xl_turbo/
+    // TODO/SEE ALSO: Alternate weights/issue: https://github.com/huggingface/diffusers/issues/6529
+    // TODO/SEE ALSO: https://huggingface.co/spaces/OzzyGT/diffusers-fast-inpaint (THIS LOOKS AWFUL THOUGH -- let's not use this)
+    //
+    // TODO: Referencing this code: https://github.com/placrosse/candle/blob/43017539ab4f9ccb43015b456136b704ebf693e0/candle-examples/examples/stable-diffusion/main.rs#L491
+    const HACK_INPAINT_SD_VERSION : StableDiffusionVersion = StableDiffusionVersion::Turbo;
+
+    // TODO: This is inpainting. It doesn't work yet.
+    //let (mask_latents, mask, mask_4) = create_inpainting_tensors(
+    //    HACK_INPAINT_SD_VERSION, // TODO: This is a hack
+    //    Some(args.image_path.to_path_buf()), // TODO: Mask needs hardcoding
+    //    args.model_configs.dtype,
+    //    &args.model_configs.device,
+    //    use_guide_scale,
+    //    &args.model_cache,
+    //    Some(input_image),
+    //    vae_scale,
+    //)?;
+
+    println!("Generating initial noise...");
+    let timesteps = scheduler.timesteps().to_vec();
+
+    // TODO: This shape may need to change for inpainting.
+    let latents = Tensor::randn(
+        0f32,
+        1f32,
+        (1, 4, args.model_configs.sd_config.height / 8, args.model_configs.sd_config.width / 8),
+        &args.model_configs.device,
+    )?;
+
+    //let mut latents = (latents * scheduler.init_noise_sigma())?;
+    let mut latents = (mask_latents * scheduler.init_noise_sigma())?;
+
+    // Initial noise shape: [1, 4, 64, 64]
+    println!("Initial noise shape: {:?}", latents.shape());
+
+    println!("Starting diffusion process...");
+
+    for (timestep_index, &timestep) in timesteps.iter().enumerate() {
+        println!("Processing step {}/{}", timestep_index + 1, timesteps.len());
+
+        let latent_model_input = latents.clone();
+
+        // Latent input shape: [1, 4, 64, 64]
+        println!("Latent input shape: {:?}", latent_model_input.shape());
+        
+        println!("Scaling model input...");
+        let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
+
+        // Scaled input shape: [1, 4, 64, 64]
+        println!("Scaled input shape: {:?}", latent_model_input.shape());
+
+
+        // TODO: THIS IS A HACK
+        //let latent_model_input = match HACK_INPAINT_SD_VERSION {
+        //    StableDiffusionVersion::XlInpaint
+        //    | StableDiffusionVersion::V2Inpaint
+        //    | StableDiffusionVersion::V1_5Inpaint => {
+        //        info!("Concatenating input shape: {:?}", latent_model_input.shape());
+        //        info!("Mask shape: {:?}", mask.as_ref().unwrap().shape());
+        //        info!("Mask latents shape: {:?}", mask_latents.as_ref().unwrap().shape());
+        //        info!("IF THIS FAILS, REVERT THE `HACK_INPAINT_SD_VERSION` 'flag'");
+        //        Tensor::cat(
+        //            &[
+        //                &latent_model_input,
+        //                mask.as_ref().unwrap(),
+        //                mask_latents.as_ref().unwrap(),
+        //            ],
+        //            1,
+        //        )?
+        //    },
+        //    _ => latent_model_input,
+        //}
+        //  .to_device(&args.model_configs.device)?;
+
+
+        
+        println!("Running UNet inference with timestep {}...", timestep);
+
+        let noise_pred = args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings)?;
+
+        println!("Noise prediction shape: {:?}", noise_pred.shape());
+        
+        println!("Applying scheduler step...");
+        latents = scheduler.step(&noise_pred, timestep, &latents)?;
+
+        println!("Step {}/{} completed", timestep_index + 1, timesteps.len());
+    }
+
+    println!("Diffusion process completed, decoding image...");
+    let image = args.model_cache.vae_decode(&(latents / 0.13025)?)?;
+
+    println!("VAE decode completed");
+    let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
+
+    println!("Normalized image values");
+    let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
+
+    println!("Converted to 8-bit format");
+    
+    save_image_from_tensor(&image.i(0)?, "temp.png")?;
+    println!("Image generation completed successfully");
+
+    Ok(())
+}

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

@@ -1,11 +1,32 @@
-mod image_endpoint;
+pub mod ml;
+pub mod endpoints;
+pub mod model_config;
 
-use std::sync::{Arc, RwLock};
+use crate::ml::model_cache::ModelCache;
+use crate::model_config::ModelConfig;
+use endpoints::image_endpoint::infer_image;
+use log::info;
 use once_cell::sync::Lazy;
-use image_endpoint::infer_image;
+use std::sync::{Arc, RwLock};
 
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
+  info!("Loading model config...");
+
+  let model_config = ModelConfig::init()
+    .expect("config should load");
+
+  info!("Creating model cache...");
+  
+  let model_cache = ModelCache::new(
+    model_config.device.clone(),
+    model_config.dtype,
+    model_config.sd_version.clone(),
+    model_config.sd_config.clone()
+  ).expect("Model cache should create");
+
+  info!("Initializing backend runtime...");
+
   tauri::Builder::default()
     .setup(|app| {
       if cfg!(debug_assertions) {
@@ -17,6 +38,8 @@ pub fn run() {
       }
       Ok(())
     })
+    .manage(model_config)
+    .manage(model_cache)
     .invoke_handler(tauri::generate_handler![
       my_custom_command, 
       test_round_trip, 

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/infer_clip_text_embeddings.rs

@@ -0,0 +1,103 @@
+#[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};
+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,
+  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/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/mod.rs

@@ -0,0 +1,7 @@
+pub mod create_inpainting_tensors;
+pub mod infer_clip_text_embeddings;
+pub mod load_image_file_to_tensor;
+pub mod model_cache;
+pub mod model_file;
+pub mod models;
+pub mod save_image_from_tensor;

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/save_image_from_tensor.rs

@@ -0,0 +1,33 @@
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+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/model_config.rs

@@ -0,0 +1,59 @@
+use crate::ml::model_file::StableDiffusionVersion;
+use candle_core::{DType, Device};
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use hf_hub::api::sync::Api;
+use log::info;
+
+const SD_IMAGE_WIDTH: usize = 512;
+const SD_IMAGE_HEIGHT: usize = 512;
+
+const SD_NUMERIC_DATATYPE: DType = DType::F32;
+const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;
+
+/// A central place to configure app-wide model details at startup.
+/// We'll probably need much more fine-grained control later.
+pub struct ModelConfig {
+  /// We only support one device type currently, and we can't multiplex 
+  /// over several devices yet.
+  pub device: Device,
+  
+  /// The core numeric type to use for models.
+  pub dtype: DType,
+  
+  pub sd_version: StableDiffusionVersion,
+  pub sd_config: StableDiffusionConfig,
+  
+  /// Probably shouldn't be here.
+  pub hf_api: Api,
+}
+
+impl ModelConfig {
+  
+  pub fn init() -> anyhow::Result<Self> {
+    
+    let device = Device::new_cuda(0)
+      .unwrap_or_else(|e| {
+          info!("CUDA not available ({}), falling back to CPU", e);
+          Device::Cpu
+        });
+    
+    let sd_config = match SD_VERSION {
+      StableDiffusionVersion::Turbo => {
+        StableDiffusionConfig::sdxl_turbo(None, Some(SD_IMAGE_HEIGHT), Some(SD_IMAGE_WIDTH))
+      }
+      _ => {
+        StableDiffusionConfig::v2_1(None, Some(SD_IMAGE_HEIGHT), Some(SD_IMAGE_WIDTH))
+      }
+    };
+    
+    let hf_api = Api::new()?;
+    
+    Ok(Self {
+      device,
+      dtype: SD_NUMERIC_DATATYPE,
+      sd_version: SD_VERSION,
+      sd_config,
+      hf_api,
+    })
+  }
+}

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