From f6b3ceb5374b7f85a934d88396ec75cba29acace Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 4 Mar 2025 16:23:55 -0500 Subject: [PATCH 1/4] fix bad abstraction for vae caching --- .../src-tauri/src/endpoints/image_endpoint.rs | 3 ++ .../desktop/src-tauri/src/ml/model_cache.rs | 36 +++++++++-------- .../stable_diffusion_pipeline.rs | 40 ++++++++++++++++--- 3 files changed, 57 insertions(+), 22 deletions(-) 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/ml/model_cache.rs b/crates/desktop/src-tauri/src/ml/model_cache.rs index 1cbbc76ba9..47d278f324 100644 --- a/crates/desktop/src-tauri/src/ml/model_cache.rs +++ b/crates/desktop/src-tauri/src/ml/model_cache.rs @@ -1,12 +1,11 @@ -use crate::ml::models::unet_model::UNetModel; use crate::ml::model_file::StableDiffusionVersion; +use crate::ml::models::unet_model::UNetModel; use anyhow::anyhow; use candle_core::{DType, Device, Tensor}; +use candle_transformers::models::stable_diffusion::vae::AutoEncoderKL; use candle_transformers::models::stable_diffusion::StableDiffusionConfig; use hf_hub::api::sync::Api; 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 @@ -20,7 +19,7 @@ pub struct ModelCache { sd_version: StableDiffusionVersion, unet: Arc>>, - lazy_vae: LazyLoadVaeModel, + vae: RwLock>>, } impl ModelCache { @@ -37,24 +36,29 @@ impl ModelCache { dtype, hf_api: api.clone(), unet: Arc::new(RwLock::new(None)), - lazy_vae: LazyLoadVaeModel::new( - &sd_config, - &sd_version, - &api, - &device, - dtype, - )?, + vae: RwLock::new(None), sd_config, sd_version, }) } - - pub fn vae_encode(&self, xs: &Tensor) -> anyhow::Result { - self.lazy_vae.encode(xs) + + 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(()) + } + } } - pub fn vae_decode(&self, xs: &Tensor) -> anyhow::Result { - self.lazy_vae.decode(xs) + 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 unet_inference(&self, latent_model_input: &Tensor, timestep: f64, text_embeddings: &Tensor) -> anyhow::Result { 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..ed9d57d050 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,7 +10,7 @@ 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; @@ -46,7 +46,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen()); args.configs.device.set_seed(seed)?; - println!("Using seed: {}", seed); + info!("Using seed: {}", seed); let guidance_scale = match args.cfg_scale { Some(guidance_scale) => guidance_scale, @@ -61,7 +61,7 @@ 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)?; @@ -102,7 +102,35 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { let vae_scale = get_vae_scale(args.configs.sd_version); - let init_latent_dist : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?; + let maybe_vae = args.model_cache.get_vae()?; + + let vae = match maybe_vae { + Some(vae) => vae, + None => { + info!("No vae found in cache; loading..."); + + let repo = args.configs.sd_version.repo(); + + println!("Building VAE model from : {:?} ... (3)", repo); + + let vae_file = args.configs.hf_api.model(repo.to_string()) + .get("vae/diffusion_pytorch_model.safetensors")?; + + println!("Building VAE model from file {:?}...", &vae_file); + + let vae = args.configs + .sd_config + .build_vae(vae_file, &args.configs.device, args.configs.dtype)?; + + let vae = Arc::new(vae); + + args.model_cache.set_vae(vae.clone())?; + + vae + } + }; + + 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 @@ -180,7 +208,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"); From 8124cbbbe11cffc807fc957e8e0d9b3608862333 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 4 Mar 2025 16:32:45 -0500 Subject: [PATCH 2/4] destructure --- .../stable_diffusion_pipeline.rs | 82 +++++++++++-------- 1 file changed, 47 insertions(+), 35 deletions(-) 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 ed9d57d050..0d7cf4bf75 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 @@ -29,28 +29,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)?; 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, @@ -64,28 +76,28 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { 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"); + 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,37 +106,37 @@ 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 maybe_vae = args.model_cache.get_vae()?; + 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 = args.configs.sd_version.repo(); + let repo = configs.sd_version.repo(); println!("Building VAE model from : {:?} ... (3)", repo); - let vae_file = args.configs.hf_api.model(repo.to_string()) + 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 vae = args.configs + let vae = configs .sd_config - .build_vae(vae_file, &args.configs.device, args.configs.dtype)?; + .build_vae(vae_file, &configs.device, configs.dtype)?; let vae = Arc::new(vae); - args.model_cache.set_vae(vae.clone())?; + model_cache.set_vae(vae.clone())?; vae } @@ -137,21 +149,21 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { 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 }; @@ -190,7 +202,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 model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) { Ok(pred) => pred, Err(e) => { println!("UNet inference failed with error: {}", e); From 89d2898e5f2279ee98f22372a9d3539b5ae61f55 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 4 Mar 2025 17:10:05 -0500 Subject: [PATCH 3/4] send notification events --- crates/desktop/src-tauri/src/events/mod.rs | 1 + .../src-tauri/src/events/notification_event.rs | 12 ++++++++++++ crates/desktop/src-tauri/src/lib.rs | 1 + .../stable_diffusion_pipeline.rs | 16 +++++++++++++++- 4 files changed, 29 insertions(+), 1 deletion(-) create mode 100644 crates/desktop/src-tauri/src/events/mod.rs create mode 100644 crates/desktop/src-tauri/src/events/notification_event.rs 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..0ba70b61bb 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; 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 0d7cf4bf75..1a4eb0724d 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 @@ -15,6 +15,7 @@ use image::DynamicImage; use log::info; use rand::Rng; use tauri::{AppHandle, Emitter}; +use crate::events::notification_event::NotificationEvent; pub struct Args<'a> { pub image: &'a DynamicImage, @@ -81,7 +82,6 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { let mut text_embeddings = if let Some(tensor) = maybe_cached { tensor } else { - app.emit("event", "loading model"); info!("Prompt is NOT cached! Calculating embedding..."); let tensor = infer_clip_text_embeddings( @@ -129,6 +129,14 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { .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 @@ -138,6 +146,12 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { model_cache.set_vae(vae.clone())?; + if notify_download_complete { + app.emit("notification", NotificationEvent::ModelDownloadComplete { + model_name: repo, + })?; + } + vae } }; From 26c7b29c17b7e1988b783e64a0e6d46cddd24c74 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 4 Mar 2025 17:51:14 -0500 Subject: [PATCH 4/4] cleanup unet caching, report download --- crates/desktop/src-tauri/src/lib.rs | 7 +- .../desktop/src-tauri/src/ml/model_cache.rs | 97 ++++--------------- .../stable_diffusion_pipeline.rs | 43 +++++++- 3 files changed, 61 insertions(+), 86 deletions(-) diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 0ba70b61bb..7f656b907b 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -25,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 47d278f324..195ef51f80 100644 --- a/crates/desktop/src-tauri/src/ml/model_cache.rs +++ b/crates/desktop/src-tauri/src/ml/model_cache.rs @@ -1,45 +1,23 @@ -use crate::ml::model_file::StableDiffusionVersion; use crate::ml::models::unet_model::UNetModel; use anyhow::anyhow; -use candle_core::{DType, Device, Tensor}; use candle_transformers::models::stable_diffusion::vae::AutoEncoderKL; -use candle_transformers::models::stable_diffusion::StableDiffusionConfig; -use hf_hub::api::sync::Api; use std::sync::{Arc, RwLock}; + // 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>>, 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)), + pub fn new() -> Self { + Self { vae: RwLock::new(None), - sd_config, - sd_version, - }) + unet2: RwLock::new(None), + } } pub fn set_vae(&self, vae: Arc) -> anyhow::Result<()> { @@ -60,63 +38,24 @@ impl ModelCache { } } } - - 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 set_unet(&self, vae: Arc) -> anyhow::Result<()> { + match self.unet2.write() { + Err(err) => Err(anyhow!("{}", err)), + Ok(mut lock) => { + *lock = Some(vae); + Ok(()) } } } - - 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 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())) } } - 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)?)) - } - } - } - 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 1a4eb0724d..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 @@ -16,6 +16,7 @@ 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, @@ -156,6 +157,46 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { } }; + 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. @@ -216,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 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);