parallel downloads

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(())
 }

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<H> {
-  pub task: JoinHandle<H>,
+struct TaskAndStatus {
   pub receiver: Receiver<ProgressUpdate>,
 }
 
 pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! {
 
-  let mut status_check_queue : Arc<Mutex<Vec<TaskAndStatus<_>>>> = Arc::new(Mutex::new(Vec::new()));
+  let mut status_check_queue : Arc<Mutex<Vec<TaskAndStatus>>> = 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(_) => {},
+    }
   }
 
   

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<HashMap<String, String>>,
-) -> 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<P: AsRef<Path>>(
   url: String,
@@ -99,6 +47,8 @@ pub async fn download_async<P: AsRef<Path>>(
   input_headers: Option<HashMap<String, String>>,
   maybe_progress_sender: Option<Sender<ProgressUpdate>>,
 ) -> 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<P: AsRef<Path>>(
     .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));

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();
    
   }