diff --git a/crates/desktop/src-tauri/src/events/notification_event.rs b/crates/desktop/src-tauri/src/events/notification_event.rs index 4a7d422c81..6ce4d9c182 100644 --- a/crates/desktop/src-tauri/src/events/notification_event.rs +++ b/crates/desktop/src-tauri/src/events/notification_event.rs @@ -6,6 +6,8 @@ use serde::Serialize; pub enum NotificationModelType { Unet, Vae, + Json, + ClipEncoder, } #[derive(Clone, Serialize)] diff --git a/crates/desktop/src-tauri/src/ml/model_type.rs b/crates/desktop/src-tauri/src/ml/model_type.rs index 597acc3ffc..69b95acb09 100644 --- a/crates/desktop/src-tauri/src/ml/model_type.rs +++ b/crates/desktop/src-tauri/src/ml/model_type.rs @@ -1,5 +1,6 @@ -use std::path::PathBuf; +use crate::events::notification_event::NotificationModelType; use crate::state::app_dir::AppWeightsDir; +use std::path::PathBuf; #[derive(Copy, Clone, Debug)] pub enum ModelType { @@ -11,6 +12,26 @@ pub enum ModelType { } impl ModelType { + pub fn get_name(&self) -> &'static str { + match self { + Self::ClipJson => "CLIP (JSON)", + Self::SdxlTurboUnet => "SDXL Turbo UNET", + Self::SdxlTurboVae => "SDXL Turbo VAE", + Self::SdxlTurboClipEncoder => "Clip Encoder #1", + Self::SdxlTurboClipEncoder2 => "Clip Encoder #2", + } + } + + pub fn get_notification_type(&self) -> NotificationModelType { + match self { + Self::ClipJson => NotificationModelType::Json, + Self::SdxlTurboUnet => NotificationModelType::Unet, + Self::SdxlTurboVae => NotificationModelType::Vae, + Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder, + Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder, + } + } + pub fn get_hf_id(&self) -> Option<&'static str> { match self { Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"), diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 609165b0e3..f9d61b1b5c 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -18,14 +18,16 @@ const CHUNK_SIZE : usize = 1024 * 1024; const PARALLEL_FAILURES : usize = 8; const MAX_RETRIES : usize = 30; -struct TaskAndStatus { - pub filename: PathBuf, +struct ModelDownloadStatusChannel { + pub model_name: String, + pub model_type: NotificationModelType, + pub final_file_path: PathBuf, pub receiver: Receiver, } pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! { - let mut status_check_queue : Arc>> = Arc::new(Mutex::new(Vec::new())); + let mut status_check_queue : Arc>> = Arc::new(Mutex::new(Vec::new())); let status_check_queue2 = status_check_queue.clone(); tokio::spawn(download_notify_thread(status_check_queue2, app)); @@ -80,9 +82,11 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! match status_check_queue.lock() { Err(err) => error!("Could not unlock: {:?}", err), Ok(mut lock) => { - lock.push(TaskAndStatus { - filename: filename.clone(), + lock.push(ModelDownloadStatusChannel { receiver: rx, + final_file_path: filename.clone(), + model_name: model.get_name().to_string(), + model_type: model.get_notification_type(), }) } } @@ -103,23 +107,22 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! } } -pub async fn download_notify_thread(queue: Arc>>, app: AppHandle) -> ! { +pub async fn download_notify_thread(queue: Arc>>, app: AppHandle) -> ! { loop { match queue.lock() { Err(err) => error!("Could not unlock: {:?}", err), Ok(mut lock) => { - debug!(">>>>> Notifier has {} items", lock.len()); let mut reenqueue = Vec::with_capacity(lock.len()); for item in (*lock).drain(..) { - if let Ok(item) = item.receiver.try_recv() { + if let Ok(update) = item.receiver.try_recv() { let result = app.emit("notification", NotificationEvent::ModelDownloadProgress { - model_name: "model", - model_type: NotificationModelType::Unet, - currently_downloaded_bytes: item.complete, - total_file_bytes: item.total_length, + model_name: &item.model_name, + model_type: item.model_type, + currently_downloaded_bytes: update.complete, + total_file_bytes: update.total_length, }); if let Err(err) = result { @@ -127,8 +130,17 @@ pub async fn download_notify_thread(queue: Arc>>, app: } } - if item.filename.exists() { - debug!("No longer sending notifications for {:?}", &item.filename); + if item.final_file_path.exists() { + debug!("No longer sending notifications for {:?}", &item.final_file_path); + + let result = app.emit("notification", NotificationEvent::ModelDownloadComplete { + model_name: &item.model_name, + model_type: item.model_type, + }); + + if let Err(err) = result { + error!("Could not emit event: {}", err); + } } else { reenqueue.push(item); } @@ -141,4 +153,3 @@ pub async fn download_notify_thread(queue: Arc>>, app: tokio::time::sleep(std::time::Duration::from_millis(500)).await; } } -