better memory organization (wip)

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

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

@@ -1,11 +1,22 @@
+use std::env::args;
+use std::sync::{Arc, RwLock};
 use crate::model::model_file::StableDiffusionVersion;
 use crate::sd::{self, Args};
 use anyhow::Result;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
-use candle_core::Device;
+use candle_core::{DType, Device};
+use candle_transformers::models::stable_diffusion;
+use once_cell::sync::Lazy;
+use crate::ml::model_cache::ModelCache;
+use crate::ml::models::unet_model::UNetModel;
 
 const PROMPT: &str = "A beautiful landscape with mountains and a lake";
 
+const HEIGHT: usize = 512;
+const WIDTH: usize = 512;
+
+static MODEL_CACHE: Lazy<Arc<RwLock<Option<ModelCache>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
+
 #[tauri::command]
 pub fn infer_image(image: &str) -> Result<String, String> {
     println!("infer_image called; generating image with SDXL Turbo...");
@@ -22,30 +33,101 @@ pub fn infer_image(image: &str) -> Result<String, String> {
         }
     };
     
-    let args = Args {
-        prompt: PROMPT.to_string(),
-        uncond_prompt: "".to_string(),
-        cpu: use_cpu,
-        height: Some(512),
-        width: Some(512),
-        n_steps: Some(1),
-        num_samples: 1,
-        seed: None,
-        sd_version: StableDiffusionVersion::Turbo,
-        guidance_scale: Some(0.0),
+    let sd_version = StableDiffusionVersion::Turbo;
+
+    let sd_config = match sd_version {
+        StableDiffusionVersion::Turbo => {
+            println!("Configuring SDXL Turbo");
+            stable_diffusion::StableDiffusionConfig::sdxl_turbo(None, Some(HEIGHT), Some(WIDTH))
+        }
+        _ => {
+            println!("Configuring SD v2.1");
+            stable_diffusion::StableDiffusionConfig::v2_1(None, Some(HEIGHT), Some(WIDTH))
+        }
     };
+    
+    let api = hf_hub::api::sync::Api::new()
+      .map_err(|err| err.to_string())?;
+
+
+    match MODEL_CACHE.write() {
+        Err(err) => return Err(err.to_string()),
+        Ok(mut maybe_cache) => {
+            match &*maybe_cache {
+                None => {
+                    let device = if use_cpu {
+                        println!("Using CPU for computation");
+                        Device::Cpu
+                    } else {
+                        println!("Attempting to use CUDA device");
+                        match Device::new_cuda(0) {
+                            Ok(cuda_device) => {
+                                println!("Successfully initialized CUDA device");
+                                cuda_device
+                            }
+                            Err(e) => {
+                                println!("Failed to initialize CUDA device: {}. Falling back to CPU", e);
+                                Device::Cpu
+                            }
+                        }
+                    };
+                    
+                    let cache = ModelCache::new(device, sd_version.clone(), sd_config.clone())
+                      .map_err(|err| err.to_string())?;
+
+                    *maybe_cache = Some(cache);
+                }
+                Some(_model) => {
+                    println!("Model cache already exists");
+                }
+            }
 
-    match sd::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))
     }
+    
+    
+
+    match MODEL_CACHE.read() {
+        Err(err) => return Err(err.to_string()),
+        Ok(maybe_model) => {
+            match &*maybe_model {
+                None => {
+                    Err("model does not exist".to_string())
+                }
+                Some(model_cache) => {
+
+                    let args = Args {
+                        api,
+                        prompt: PROMPT.to_string(),
+                        uncond_prompt: "".to_string(),
+                        cpu: use_cpu,
+                        height: Some(HEIGHT),
+                        width: Some(WIDTH),
+                        n_steps: Some(1),
+                        num_samples: 1,
+                        seed: None,
+                        sd_version,
+                        guidance_scale: Some(0.0),
+                        model_cache,
+                    };
+
+                    match sd::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/lib.rs

@@ -2,6 +2,7 @@ mod flip_image;
 mod image_endpoint;
 mod sd;
 pub mod model;
+pub mod ml;
 
 use std::sync::{Arc, RwLock};
 use once_cell::sync::Lazy;

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

@@ -0,0 +1,2 @@
+pub mod model_cache;
+pub mod models;

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

@@ -0,0 +1,90 @@
+use crate::ml::models::unet_model::UNetModel;
+use crate::model::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};
+
+// 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,
+  hf_api: Api,
+  sd_config: StableDiffusionConfig,
+  sd_version: StableDiffusionVersion,
+  unet: Arc<RwLock<Option<UNetModel>>>,
+}
+
+impl ModelCache {
+  pub fn new(device: Device, sd_version: StableDiffusionVersion, sd_config: StableDiffusionConfig) -> anyhow::Result<Self> {
+    Ok(Self { 
+      device,
+      hf_api: Api::new()?,
+      unet: Arc::new(RwLock::new(None)),
+      sd_config,
+      sd_version,
+    })
+  }
+  
+  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, DType::F32)
+              .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/models/mod.rs

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

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/models/unet_model_cached.rs

@@ -32,7 +32,7 @@ impl UNetModel {
     device: &Device,
     dtype: DType,
   ) -> anyhow::Result<Self> {
-    debug!("building unet model...");
+    println!("building unet model... (2)");
     
     let unet = sd_config.build_unet(
       unet_weights_file, 

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

@@ -2,5 +2,3 @@ pub mod gpu_detector;
 pub mod infer_clip_text_embeddings;
 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

@@ -1,21 +0,0 @@
-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/sd.rs

@@ -11,15 +11,21 @@ use crate::model::model_file::StableDiffusionVersion;
 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 hf_hub::api::sync::Api;
 use once_cell::sync::Lazy;
 use rand::Rng;
-use crate::model::model_cache::ModelCache;
-use crate::model::unet_model::UNetModel;
+use crate::ml::model_cache::ModelCache;
+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)));
 
-pub struct Args {
+// 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 api: Api,
     pub prompt: String,
     pub uncond_prompt: String,
     pub cpu: bool,
@@ -30,9 +36,10 @@ pub struct Args {
     pub seed: Option<u64>,
     pub sd_version: StableDiffusionVersion,
     pub guidance_scale: Option<f64>,
+    pub model_cache: &'a ModelCache,
 }
 
-pub fn run(args: Args) -> Result<()> {
+pub fn run(args: Args<'_>) -> Result<()> {
     println!("Starting image generation with the following configuration:");
     println!("  Model: {:?}", args.sd_version);
     println!("  Prompt: {}", args.prompt);
@@ -72,35 +79,30 @@ 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);
     
-    println!("Downloading VAE and UNet...");
-    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!("Downloading VAE ...");
+    let vae_file = args.api.model(repo.to_string()).get("vae/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")?;
+    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)?;
@@ -153,30 +155,30 @@ 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");
-                }
-            }
-            
-        }
-    }
-    
+    //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;
-    
-    
-    
+
+
+
     for (timestep_index, &timestep) in timesteps.iter().enumerate() {
         println!("Processing step {}/{}", timestep_index + 1, timesteps.len());
         let latent_model_input = latents.clone();
@@ -187,24 +189,12 @@ 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_MODEL.read() {
-            Err(err) => return Err(anyhow::Error::msg(err.to_string())),
-            Ok(maybe_model) => {
-                match &*maybe_model {
-                    None => {
-                        return Err(anyhow::Error::msg("Model does not exist"));
-                    }
-                    Some(unet) => {
-                        noise_pred = unet.inference(&latent_model_input, timestep as f64, &text_embeddings)?;
-                    }
-                }
-            }
-        }
-        
-        
+
+        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");
@@ -215,7 +205,7 @@ pub fn run(args: Args) -> Result<()> {
         //        return Err(anyhow::anyhow!("UNet inference failed: {}", e));
         //    }
         //};
-        
+
         println!("Noise prediction shape: {:?}", noise_pred.shape());
         
         println!("Applying scheduler step...");