wip model caching

.idea/dictionaries/bt.xml

@@ -47,6 +47,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/src/model/mod.rs

@@ -1,3 +1,6 @@
-pub mod save_image_from_tensor;
+pub mod gpu_detector;
 pub mod infer_clip_text_embeddings;
-pub mod model_file;
+pub mod model_file;
+pub mod save_image_from_tensor;
+pub mod unet_model;
+pub mod model_cache;

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

@@ -0,0 +1,21 @@
+use std::sync::{Arc, RwLock};
+use once_cell::sync::Lazy;
+use crate::model::unet_model::UNetModel;
+
+
+static MODEL_CACHE: Lazy<Arc<RwLock<ModelCache>>> = Lazy::new(|| Arc::new(RwLock::new(ModelCache::new())));
+
+
+// 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 {
+  unet: Option<UNetModel>,
+}
+
+impl ModelCache {
+  fn new() -> Self {
+    Self { 
+      unet: None 
+    }
+  }
+}

crates/desktop/src-tauri/src/model/unet_model.rs

@@ -0,0 +1,72 @@
+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> {
+    debug!("building unet model...");
+    
+    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

@@ -11,6 +11,7 @@ use crate::model::save_image_from_tensor::save_image_from_tensor;
 use anyhow::{Error as E, Result};
 use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
 use rand::Rng;
+use crate::model::unet_model::UNetModel;
 
 pub struct Args {
     pub prompt: String,
@@ -65,12 +66,14 @@ pub fn run(args: Args) -> Result<()> {
 
     println!("Building scheduler...");
     let mut scheduler = 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);
     device.set_seed(seed)?;
 
     println!("Initializing Hugging Face API...");
     let api = hf_hub::api::sync::Api::new()?;
+    
     let repo = args.sd_version.repo();
     println!("Downloading model files from: {}", repo);
     
@@ -78,15 +81,21 @@ pub fn run(args: Args) -> Result<()> {
     let vae_file = api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
     let unet_file = api.model(repo.to_string()).get("unet/diffusion_pytorch_model.safetensors")?;
     
+    println!("VAE Path: {:?}", &vae_file);
+    println!("UNET Path: {:?}", &unet_file);
+    
     println!("Downloading text encoders...");
     let clip_file = api.model(repo.to_string()).get("text_encoder/model.safetensors")?;
     let clip2_file = api.model(repo.to_string()).get("text_encoder_2/model.safetensors")?;
     let tokenizer = api.model("openai/clip-vit-large-patch14".to_string()).get("tokenizer.json")?;
     let tokenizer2 = api.model("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k".to_string()).get("tokenizer.json")?;
 
-    println!("Building models...");
+    println!("Building VAE model...");
     let vae = sd_config.build_vae(vae_file, &device, dtype)?;
-    let unet = sd_config.build_unet(unet_file, &device, 4, false, dtype)?;
+    
+    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)?;
     
     // Build text encoders
     println!("Building text encoders...");
@@ -152,16 +161,20 @@ pub fn run(args: Args) -> Result<()> {
         println!("Scaled input shape: {:?}", latent_model_input.shape());
         
         println!("Running UNet inference with timestep {}...", timestep);
-        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 = 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));
+        //    }
+        //};
+        
         println!("Noise prediction shape: {:?}", noise_pred.shape());
         
         println!("Applying scheduler step...");

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