lazy load vae

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

@@ -49,6 +49,7 @@ pub fn infer_image(image: &str) -> Result<String, String> {
     let api = hf_hub::api::sync::Api::new()
       .map_err(|err| err.to_string())?;
 
+    let dtype = DType::F32;
 
     match MODEL_CACHE.write() {
         Err(err) => return Err(err.to_string()),
@@ -72,7 +73,7 @@ pub fn infer_image(image: &str) -> Result<String, String> {
                         }
                     };
                     
-                    let cache = ModelCache::new(device, sd_version.clone(), sd_config.clone())
+                    let cache = ModelCache::new(device, dtype, sd_version.clone(), sd_config.clone())
                       .map_err(|err| err.to_string())?;
 
                     *maybe_cache = Some(cache);

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

@@ -5,30 +5,53 @@ 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 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.
-#[derive(Clone)]
 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, sd_version: StableDiffusionVersion, sd_config: StableDiffusionConfig) -> anyhow::Result<Self> {
+  pub fn new(
+    device: Device, 
+    dtype: DType,
+    sd_version: StableDiffusionVersion, 
+    sd_config: StableDiffusionConfig,
+  ) -> anyhow::Result<Self> {
+    let api = Api::new()?;
+    
     Ok(Self { 
-      device,
-      hf_api: Api::new()?,
+      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_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) {
@@ -60,7 +83,7 @@ impl ModelCache {
               .map_err(|err| anyhow!("error fetching model: {:?}", err))?;
 
             println!("Building UNET model... (3)");
-            let unet = UNetModel::new(&self.sd_config, unet_file, &self.device, DType::F32)
+            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...");

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

@@ -0,0 +1,104 @@
+use candle_core::{DType, Device, Tensor};
+use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, RwLock};
+use anyhow::anyhow;
+use candle_transformers::models::stable_diffusion::vae::AutoEncoderKL;
+use hf_hub::api::sync::Api;
+use crate::ml::models::unet_model::UNetModel;
+use crate::model::model_file::StableDiffusionVersion;
+
+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 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_decode(&self, xs: &Tensor) -> anyhow::Result<Option<Tensor>> {
+    match self.lazy_model.read() {
+      Err(err) => Err(anyhow!("lock error: {:?}", err)),
+      Ok(model) => {
+        match &*model {
+          None => Ok(None), // Model is not yet loaded
+          Some(model) => {
+            Ok(Some(model.decode(xs)?))
+          }
+        }
+      }
+    }
+  }
+
+  fn load_model(&self) -> anyhow::Result<()> {
+    println!("load vae");
+    match self.lazy_model.write() {
+      Err(err) => return Err(anyhow!("lock error: {:?}", err)),
+      Ok(mut maybe_model) => {
+        match &*maybe_model {
+          Some(_model) => {} // Already loaded; fall through
+          None => {
+            let repo = self.model_configs.sd_version.repo();
+
+            println!("Building VAE model from : {:?} ... (3)", repo);
+            
+            let vae_file = self.model_configs.hf_api.model(repo.to_string())
+              .get("vae/diffusion_pytorch_model.safetensors")?;
+
+            println!("Building VAE model from file {:?}...", &vae_file);
+            let vae = self.model_configs
+              .sd_config
+              .build_vae(vae_file, &self.model_configs.device, self.model_configs.dtype)?;
+
+            println!("Built VAE...");
+            *maybe_model = Some(vae);
+          }
+        }
+      }
+    }
+    Ok(())
+  }
+  
+}

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

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

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

@@ -1,72 +0,0 @@
-use anyhow::anyhow;
-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;
-use std::sync::{Arc, RwLock};
-use log::debug;
-
-#[derive(Clone)]
-pub struct UNetModel {
-  // Kasisnu - a few notes: 
-  //
-  // (1) This struct has what Rust terms as "interior mutability". We can now pass around 
-  // "immutable" references to the struct, yet mutate what's inside with non-mut (!) functions. 
-  // The callers don't know, and because the struct has safety guarantees of "Send + Sync", it's 
-  // provably sound. This allows multiple threads / scopes to mutate the structure. They just wait 
-  // on the lock/mutex to synchronize.
-  //
-  // (2) Another cool thing we can do is mark this struct as "Clone"  since the Arc<T> pointer is
-  // super cheap to duplicate (you'd never do that with a wide type or anything beyond int64 
-  // size). It makes passing this structure along smooth as butter. The only penalty we pay is the 
-  // unlocking of the mutex (actually rwlock in this case) and the (minor) indirection behind Arc.
-  // We don't mark it Copy, though, since it does ref counting and manages its own allocation, which 
-  // shouldn't be implicit.
-  model: Arc<RwLock<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... (2)");
-    
-    let unet = sd_config.build_unet(
-      unet_weights_file, 
-      &device, 
-      4, 
-      false, 
-      dtype
-    )?;
-    
-    Ok(Self {
-      model: Arc::new(RwLock::new(unet)),
-    })
-  }
-  
-  pub fn inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result<Tensor> {
-    match self.model.read() {
-      Err(err) => return Err(anyhow!("lock error: {:?}", err)),
-      Ok(model) => {
-        debug!("Running model prediction...");
-        let inference_result = model.forward(&latent_model_input, timestep, &text_embeddings);
-
-        let noise_pred = match inference_result {
-          Ok(pred) => {
-            debug!("Inference successful");
-            pred
-          },
-          Err(e) => {
-            debug!("UNet inference failed with error: {}", e);
-            return Err(anyhow!("UNet inference failed: {}", e));
-          }
-        };
-        
-        Ok(noise_pred)
-      }
-    }
-  }
-}

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

@@ -19,8 +19,6 @@ use crate::ml::models::unet_model::UNetModel;
 
 // TODO: Clean up
 
-static UNET_MODEL: Lazy<Arc<RwLock<Option<UNetModel>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
-
 // 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.
@@ -89,29 +87,29 @@ pub fn run(args: Args<'_>) -> Result<()> {
     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!("Downloading VAE ...");
+    //let vae_file = args.api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
 
-    println!("VAE Path: {:?}", &vae_file);
+    //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!("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)?;
+    //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)?;
+    //// 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 {
@@ -155,24 +153,6 @@ pub fn run(args: Args<'_>) -> Result<()> {
     let latents = (latents * scheduler.init_noise_sigma())?;
     println!("Initial noise shape: {:?}", latents.shape());
 
-    //match UNET_MODEL.write() {
-    //    Err(err) => return Err(anyhow::Error::msg(err.to_string())),
-    //    Ok(mut maybe_model) => {
-    //        match &*maybe_model {
-    //            None => {
-    //                println!("Building UNET model...");
-    //                //let unet = sd_config.build_unet(unet_file, &device, 4, false, dtype)?;
-    //                let unet = UNetModel::new(&sd_config, unet_file, &device, dtype)?;
-    //                *maybe_model = Some(unet);
-    //            }
-    //            Some(_model) => {
-    //                println!("Model already exists");
-    //            }
-    //        }
-    //
-    //    }
-    //}
-
 
     println!("Starting diffusion process...");
     let mut latents = latents;
@@ -190,21 +170,7 @@ pub fn run(args: Args<'_>) -> Result<()> {
         
         println!("Running UNet inference with timestep {}...", timestep);
 
-        let noise_pred;
-
-        noise_pred = args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings)?;
-
-
-        //let noise_pred = match unet.forward(&latent_model_input, timestep as f64, &text_embeddings) {
-        //    Ok(pred) => {
-        //        println!("UNet inference successful");
-        //        pred
-        //    }
-        //    Err(e) => {
-        //        println!("UNet inference failed with error: {}", e);
-        //        return Err(anyhow::anyhow!("UNet inference failed: {}", e));
-        //    }
-        //};
+        let noise_pred = args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings)?;
 
         println!("Noise prediction shape: {:?}", noise_pred.shape());
         
@@ -214,7 +180,7 @@ pub fn run(args: Args<'_>) -> Result<()> {
     }
 
     println!("Diffusion process completed, decoding image...");
-    let image = vae.decode(&(latents / 0.13025)?)?;
+    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");