Merge rust

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

@@ -16,6 +16,9 @@ const PROMPT: &str = "A beautiful landscape with mountains and a lake";
 
 const NEGATIVE_PROMPT: &str = "bad quality, bad faces, poor quality, blurry faces, watermark";
 
+// 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]

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,12 @@
+use serde::Serialize;
+
+#[derive(Clone, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum NotificationEvent<'a> {
+  ModelDownloadStarted {
+    model_name: &'a str, 
+  },
+  ModelDownloadComplete {
+    model_name: &'a str,
+  }
+}

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

@@ -1,4 +1,5 @@
 pub mod endpoints;
+pub mod events;
 pub mod ml;
 pub mod state;
 
@@ -24,12 +25,7 @@ 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();
 
   println!("Initializing backend runtime...");
 

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/stable_diffusion/stable_diffusion_pipeline.rs

@@ -1,4 +1,4 @@
-use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
 use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
@@ -10,11 +10,13 @@ use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_emb
 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::events::notification_event::NotificationEvent;
+use crate::ml::models::unet_model::UNetModel;
 
 pub struct Args<'a> {
     pub image: &'a DynamicImage,
@@ -29,28 +31,40 @@ pub struct Args<'a> {
 }
 
 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,  
+    } = args;
+    
     println!("Starting image generation with the following configuration:");
-    println!("  Model: {:?}", args.configs.sd_version);
-    println!("  Prompt: {}", args.prompt);
-    println!("  Steps: {}", args.configs.scheduler_steps);
-    println!("  Device: {:?}", args.configs.device);
+    println!("  Model: {:?}", 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,
@@ -61,31 +75,30 @@ 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 {
-        args.app.emit("event", "loading model");
         
         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,
         )?;
-        args.prompt_cache.store_copy(&args.prompt, &tensor)?;
+        prompt_cache.store_copy(&prompt, &tensor)?;
         tensor
     };
 
@@ -94,36 +107,118 @@ 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 repo = configs.sd_version.repo();
+
+            println!("Building VAE model from : {:?} ... (3)", repo);
+
+            let vae_file = configs.hf_api.model(repo.to_string())
+              .get("vae/diffusion_pytorch_model.safetensors")?;
+
+            println!("Building VAE model from file {:?}...", &vae_file);
+
+            let mut notify_download_complete = false;
+            if !vae_file.exists() {
+                notify_download_complete = true;
+                app.emit("notification", NotificationEvent::ModelDownloadStarted {
+                    model_name: repo,
+                })?;
+            }
+            
+            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: repo,
+                })?;
+            }
+            
+            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 repo = configs.sd_version.repo();
+
+            info!("Downloading UNET model files from: {} ...", repo);
+
+            let unet_file = configs.hf_api.model(repo.to_string())
+              .get("unet/diffusion_pytorch_model.safetensors")
+              .map_err(|err| anyhow!("error fetching model: {:?}", err))?;
+
+            let mut notify_download_complete = false;
+            if !unet_file.exists() {
+                notify_download_complete = true;
+                app.emit("notification", NotificationEvent::ModelDownloadStarted {
+                    model_name: repo,
+                })?;
+            }
+
+            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: repo,
+                })?;
+            }
+
+            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
     };
 
@@ -162,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);
@@ -180,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");