diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 460c994a52..bb8e11cdd7 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/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] diff --git a/crates/desktop/src-tauri/src/events/mod.rs b/crates/desktop/src-tauri/src/events/mod.rs new file mode 100644 index 0000000000..2db7fd03ce --- /dev/null +++ b/crates/desktop/src-tauri/src/events/mod.rs @@ -0,0 +1 @@ +pub mod notification_event; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/events/notification_event.rs b/crates/desktop/src-tauri/src/events/notification_event.rs new file mode 100644 index 0000000000..4ccd8aaff0 --- /dev/null +++ b/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, + } +} diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 625e8f1c96..7f656b907b 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/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..."); diff --git a/crates/desktop/src-tauri/src/ml/model_cache.rs b/crates/desktop/src-tauri/src/ml/model_cache.rs index 1cbbc76ba9..195ef51f80 100644 --- a/crates/desktop/src-tauri/src/ml/model_cache.rs +++ b/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>>, - lazy_vae: LazyLoadVaeModel, + vae: RwLock>>, + unet2: RwLock>>, } impl ModelCache { - pub fn new( - device: Device, - dtype: DType, - sd_version: StableDiffusionVersion, - sd_config: StableDiffusionConfig, - ) -> anyhow::Result { - 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 { - self.lazy_vae.encode(xs) - } - - 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) { - 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) -> 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> { - 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>> { + 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) -> 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>> { + 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)), } } diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs index 91d51a4bc8..3a656f7d93 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs +++ b/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 { + 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 { 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 { 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 { 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 { } 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");