From 4e476c00513e234989237b9536dc5efd9cac38dc Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 14 Feb 2025 18:20:21 -0500 Subject: [PATCH] lazy load vae --- .../desktop/src-tauri/src/image_endpoint.rs | 3 +- .../desktop/src-tauri/src/ml/model_cache.rs | 33 +++++- .../src/ml/models/lazy_load_vae_model.rs | 104 ++++++++++++++++++ crates/desktop/src-tauri/src/ml/models/mod.rs | 2 +- .../src/ml/models/unet_model_cached.rs | 72 ------------ crates/desktop/src-tauri/src/sd.rs | 76 ++++--------- 6 files changed, 156 insertions(+), 134 deletions(-) create mode 100644 crates/desktop/src-tauri/src/ml/models/lazy_load_vae_model.rs delete mode 100644 crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs diff --git a/crates/desktop/src-tauri/src/image_endpoint.rs b/crates/desktop/src-tauri/src/image_endpoint.rs index bd249aee18..7f72d3eb04 100644 --- a/crates/desktop/src-tauri/src/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/image_endpoint.rs @@ -49,6 +49,7 @@ pub fn infer_image(image: &str) -> Result { 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 { } }; - 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); diff --git a/crates/desktop/src-tauri/src/ml/model_cache.rs b/crates/desktop/src-tauri/src/ml/model_cache.rs index 1732e3714f..17d280024e 100644 --- a/crates/desktop/src-tauri/src/ml/model_cache.rs +++ b/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>>, + lazy_vae: LazyLoadVaeModel, } impl ModelCache { - pub fn new(device: Device, sd_version: StableDiffusionVersion, sd_config: StableDiffusionConfig) -> anyhow::Result { + pub fn new( + device: Device, + dtype: DType, + sd_version: StableDiffusionVersion, + sd_config: StableDiffusionConfig, + ) -> anyhow::Result { + 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 { + self.lazy_vae.decode(xs) + } + pub fn unet_inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result { 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..."); diff --git a/crates/desktop/src-tauri/src/ml/models/lazy_load_vae_model.rs b/crates/desktop/src-tauri/src/ml/models/lazy_load_vae_model.rs new file mode 100644 index 0000000000..74ef285cae --- /dev/null +++ b/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>>, + model_configs: VaeConfigs, +} + +// TODO: Maybe we pass these via a central config read at startup, passed as an Arc +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 { + 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 { + 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> { + 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(()) + } + +} diff --git a/crates/desktop/src-tauri/src/ml/models/mod.rs b/crates/desktop/src-tauri/src/ml/models/mod.rs index 8631dad75c..be97348798 100644 --- a/crates/desktop/src-tauri/src/ml/models/mod.rs +++ b/crates/desktop/src-tauri/src/ml/models/mod.rs @@ -1,2 +1,2 @@ pub mod unet_model; -pub mod unet_model_cached; \ No newline at end of file +pub mod lazy_load_vae_model; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs b/crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs deleted file mode 100644 index eeb136da61..0000000000 --- a/crates/desktop/src-tauri/src/ml/models/unet_model_cached.rs +++ /dev/null @@ -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 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>, -} - -impl UNetModel { - pub fn new>( - sd_config: &StableDiffusionConfig, - unet_weights_file: P, - device: &Device, - dtype: DType, - ) -> anyhow::Result { - 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 { - 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) - } - } - } -} diff --git a/crates/desktop/src-tauri/src/sd.rs b/crates/desktop/src-tauri/src/sd.rs index d8d174b23b..9aedb46764 100644 --- a/crates/desktop/src-tauri/src/sd.rs +++ b/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>>> = 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");