From de1d5cbc39757e4eaf764266b4adbf5a442d89e4 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 10 Mar 2025 19:30:39 -0400 Subject: [PATCH] parallel downloads --- .../src/ml/downloads/download_all_models.rs | 10 ++-- .../src/threads/downloader_thread.rs | 31 +++++++--- .../src-tauri/src/transfer/download.rs | 59 ++----------------- .../RenderingPrimitives/RealTimeDrawEngine.ts | 2 +- 4 files changed, 36 insertions(+), 66 deletions(-) diff --git a/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs b/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs index 9b97957447..b83ab08495 100644 --- a/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs +++ b/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs @@ -5,11 +5,11 @@ use crate::state::app_dir::AppDataRoot; pub async fn download_all_models(app_data_root: &AppDataRoot) -> anyhow::Result<()> { let weights_dir = app_data_root.weights_dir(); - download_model_registry_file(ModelType::ClipJson, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?; + //download_model_registry_file(ModelType::ClipJson, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?; Ok(()) } diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index be423a54f5..b2d5dc6d97 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -5,7 +5,9 @@ use tauri::{AppHandle, Emitter}; use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use std::sync::mpsc::Receiver; -use log::error; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use log::{error, info}; use tokio::task::JoinHandle; use crate::events::notification_event::{NotificationEvent, NotificationModelType}; @@ -14,14 +16,13 @@ const CHUNK_SIZE : usize = 1024 * 1024; const PARALLEL_FAILURES : usize = 8; const MAX_RETRIES : usize = 30; -struct TaskAndStatus { - pub task: JoinHandle, +struct TaskAndStatus { 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 mut status_check_queue2 = status_check_queue.clone(); tokio::spawn(async move { @@ -61,28 +62,44 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! download_queue.push_back(ModelType::SdxlTurboVae); download_queue.push_back(ModelType::SdxlTurboClipEncoder); download_queue.push_back(ModelType::SdxlTurboClipEncoder2); - + + let mut handles = FuturesUnordered::new(); + while let Some(model) = download_queue.pop_front() { let url = model.get_download_url().to_string(); let filename = app_data_root.weights_dir().model_path(&model); + if filename.exists() { + continue; + } + let headers = None; let (tx, rx) = std::sync::mpsc::channel(); + info!("Download: {:?}", filename); + let task = tokio::spawn(download_async(url, filename, MAX_FILES, CHUNK_SIZE, PARALLEL_FAILURES, MAX_RETRIES, headers, Some(tx))); - + match status_check_queue.lock() { Err(err) => error!("Could not unlock: {:?}", err), Ok(mut lock) => { lock.push(TaskAndStatus { - task, receiver: rx, }) } } + + handles.push(task); + } + + while let Some(result) = handles.next().await { + match result { + Err(err) => error!("Error downloading: {:?}", err), + Ok(_) => {}, + } } diff --git a/crates/desktop/src-tauri/src/transfer/download.rs b/crates/desktop/src-tauri/src/transfer/download.rs index cd426d051e..a9d4ecace2 100644 --- a/crates/desktop/src-tauri/src/transfer/download.rs +++ b/crates/desktop/src-tauri/src/transfer/download.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::mpsc::Sender; use std::time::Duration; +use log::info; use tokio::fs::OpenOptions; use tokio::io::AsyncWriteExt; use tokio::io::AsyncSeekExt; @@ -35,59 +36,6 @@ pub struct ProgressUpdate { /// /// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in /// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads -#[allow(clippy::too_many_arguments)] -pub fn download( - url: String, - filename: String, - max_files: usize, - chunk_size: usize, - parallel_failures: usize, - max_retries: usize, - headers: Option>, -) -> Result<(), Error> { - if parallel_failures > max_files { - return Err(Error::new_err( - "Error parallel_failures cannot be > max_files".to_string(), - )); - } - if (parallel_failures == 0) != (max_retries == 0) { - return Err(Error::new_err( - "For retry mechanism you need to set both `parallel_failures` and `max_retries`" - .to_string(), - )); - } - tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build()? - .block_on(async { - download_async( - url, - filename.clone(), - max_files, - chunk_size, - parallel_failures, - max_retries, - headers, - None, - ) - .await - }) - .map_err(|err| { - let path = Path::new(&filename); - if path.exists() { - match remove_file(filename) { - Ok(_) => err, - Err(err) => { - Error::new_err(format!("Error while removing corrupted file: {err}")) - } - } - } else { - err - } - }) -} - - #[allow(clippy::too_many_arguments)] pub async fn download_async>( url: String, @@ -99,6 +47,8 @@ pub async fn download_async>( input_headers: Option>, maybe_progress_sender: Option>, ) -> Result<(), Error> { + info!(">>> Downloading: {}", filename.as_ref().display()); + let client = reqwest::Client::builder() // https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526 .http2_keep_alive_timeout(Duration::from_secs(15)) @@ -167,6 +117,9 @@ pub async fn download_async>( .parse() .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?; + + info!("File {:?} is {:?} length", filename.as_ref(), length); + let mut handles = FuturesUnordered::new(); let semaphore = Arc::new(Semaphore::new(max_files)); let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures)); diff --git a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts index 7d06b15847..66d74eeba2 100644 --- a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts +++ b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts @@ -212,7 +212,7 @@ export class RealTimeDrawEngine { // Add mouse events for preview canvas copying this.previewCopyListener(); - //this.listenToServerEvents(); + this.listenToServerEvents(); }