notify of download start

Cargo.lock

@@ -666,12 +666,14 @@ dependencies = [
  "cudarc",
  "expanduser",
  "filesys",
+ "futures 0.3.29",
  "hf-hub 0.3.2",
  "image",
  "libloading 0.8.6",
  "log",
  "memory_caching",
  "once_cell",
+ "openssl",
  "rand 0.8.5",
  "reqwest 0.12.12",
  "serde",
@@ -680,7 +682,11 @@ dependencies = [
  "tauri",
  "tauri-build",
  "tauri-plugin-log",
+ "tempdir",
+ "tempfile",
  "tokenizers 0.15.2",
+ "tokio",
+ "tokio-util",
 ]
 
 [[package]]
@@ -3792,12 +3798,12 @@ dependencies = [
 
 [[package]]
 name = "errno"
-version = "0.3.9"
+version = "0.3.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba"
+checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d"
 dependencies = [
  "libc",
- "windows-sys 0.52.0",
+ "windows-sys 0.59.0",
 ]
 
 [[package]]
@@ -3927,9 +3933,9 @@ dependencies = [
 
 [[package]]
 name = "fastrand"
-version = "2.1.0"
+version = "2.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9fc0510504f03c51ada170672ac806f1f105a88aa97a5281117e1ddc3368e51a"
+checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
 
 [[package]]
 name = "fdeflate"
@@ -4657,6 +4663,18 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "getrandom"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8"
+dependencies = [
+ "cfg-if 1.0.0",
+ "libc",
+ "wasi 0.13.3+wasi-0.2.2",
+ "windows-targets 0.52.6",
+]
+
 [[package]]
 name = "getset"
 version = "0.1.4"
@@ -6011,7 +6029,7 @@ dependencies = [
  "chumsky",
  "email-encoding",
  "email_address",
- "fastrand 2.1.0",
+ "fastrand 2.3.0",
  "futures-util",
  "hostname 0.3.1",
  "httpdate",
@@ -6229,6 +6247,12 @@ version = "0.4.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89"
 
+[[package]]
+name = "linux-raw-sys"
+version = "0.9.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6db9c683daf087dc577b7506e9695b3d556a9f3849903fa28186283afd6809e9"
+
 [[package]]
 name = "local-channel"
 version = "0.1.5"
@@ -9664,10 +9688,23 @@ dependencies = [
  "bitflags 2.5.0",
  "errno",
  "libc",
- "linux-raw-sys",
+ "linux-raw-sys 0.4.14",
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "rustix"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f7178faa4b75a30e269c71e61c353ce2748cf3d76f0c44c393f4e60abf49b825"
+dependencies = [
+ "bitflags 2.5.0",
+ "errno",
+ "libc",
+ "linux-raw-sys 0.9.2",
+ "windows-sys 0.59.0",
+]
+
 [[package]]
 name = "rustls"
 version = "0.21.12"
@@ -11899,14 +11936,16 @@ dependencies = [
 
 [[package]]
 name = "tempfile"
-version = "3.10.1"
+version = "3.18.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "85b77fafb263dd9d05cbeac119526425676db3784113aa9295c88498cbf8bff1"
+checksum = "2c317e0a526ee6120d8dabad239c8dadca62b24b6f168914bbbc8e2fb1f0e567"
 dependencies = [
  "cfg-if 1.0.0",
- "fastrand 2.1.0",
- "rustix",
- "windows-sys 0.52.0",
+ "fastrand 2.3.0",
+ "getrandom 0.3.1",
+ "once_cell",
+ "rustix 1.0.2",
+ "windows-sys 0.59.0",
 ]
 
 [[package]]
@@ -13279,6 +13318,15 @@ version = "0.11.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
+[[package]]
+name = "wasi"
+version = "0.13.3+wasi-0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2"
+dependencies = [
+ "wit-bindgen-rt",
+]
+
 [[package]]
 name = "wasite"
 version = "0.1.0"
@@ -14048,6 +14096,15 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "wit-bindgen-rt"
+version = "0.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c"
+dependencies = [
+ "bitflags 2.5.0",
+]
+
 [[package]]
 name = "wry"
 version = "0.48.1"
@@ -14138,8 +14195,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e105d177a3871454f754b33bb0ee637ecaaac997446375fd3e5d43a2ed00c909"
 dependencies = [
  "libc",
- "linux-raw-sys",
- "rustix",
+ "linux-raw-sys 0.4.14",
+ "rustix 0.38.34",
 ]
 
 [[package]]

crates/desktop/src-tauri/Cargo.toml

@@ -28,24 +28,30 @@ cudarc = "0.13.7"
 
 # External
 anyhow = "1.0"
-bytes = "1.10"
 base64 = "0.22.1"
+bytes = "1.10"
 candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
 candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
 candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
+expanduser = "1.2.2"
+futures = "0.3"
 hf-hub = "0.3.2"
 image = "0.25.5"
-serde_yaml = "0.9"
 log = "0.4"
 once_cell = "1.20.3"
-expanduser = "1.2.2"
-rand = "0.8"
-reqwest = "0.12.12"
+openssl = { version = "0.10", features = ["vendored"] }
+rand = "0.8.5"
+reqwest = { version = "0.12.12", features = ["stream"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+serde_yaml = "0.9"
 tauri = { version = "2.2.4", features = [] }
 tauri-plugin-log = "2"
+tempdir = "0.3.7"
+tempfile = "3.18.0"
 tokenizers = "0.15.2"
+tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
+tokio-util = { version = "0.7", features = ["codec"] }
 
 [features]
 default = ["cuda"]

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)]
@@ -15,6 +17,12 @@ pub enum NotificationEvent<'a> {
     model_name: &'a str, 
     model_type: NotificationModelType,
   },
+  ModelDownloadProgress {
+    model_name: &'a str,
+    model_type: NotificationModelType,
+    currently_downloaded_bytes: usize,
+    total_file_bytes: usize,
+  },
   ModelDownloadComplete {
     model_name: &'a str,
     model_type: NotificationModelType,

crates/desktop/src-tauri/src/lib.rs

@@ -3,6 +3,7 @@ pub mod events;
 pub mod ml;
 pub mod state;
 pub mod threads;
+pub mod transfer;
 
 use crate::endpoints::download_models::download_models;
 use crate::endpoints::image_endpoint::infer_image;

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/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"),

crates/desktop/src-tauri/src/state/app_dir.rs

@@ -1,17 +1,22 @@
 use crate::state::os_platform::OsPlatform;
 use expanduser::expanduser;
 use std::path::{Path, PathBuf};
+use tempdir::TempDir;
+use tempfile::{Builder, NamedTempFile};
 use crate::ml::model_type::ModelType;
 
 const ASSETS_SUBDIRECTORY : &str = "assets";
 const WEIGHTS_SUBDIRECTORY : &str = "weights";
 
+const TEMPORARY_SUBDIRECTORY : &str = "temp";
+
 /// The path to the application data directory, which includes "asset" and "weights" data.
 #[derive(Clone)]
 pub struct AppDataRoot {
   path: PathBuf,
   assets_dir: AppAssetsDir,
   weights_dir: AppWeightsDir,
+  temp_dir: TemporaryDir,
 }
 
 #[derive(Clone)]
@@ -24,6 +29,11 @@ pub struct AppWeightsDir {
   path: PathBuf,
 }
 
+#[derive(Clone)]
+pub struct TemporaryDir {
+  path: PathBuf,
+}
+
 impl AppDataRoot {
   pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
     let mut dir = dir.as_ref().to_path_buf();
@@ -51,11 +61,13 @@ impl AppDataRoot {
     
     let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?;
     let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?;
-    
+    let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?;
+
     Ok(Self {
       path: dir,
       assets_dir,
       weights_dir,
+      temp_dir,
     })
   }
   
@@ -66,7 +78,11 @@ impl AppDataRoot {
   pub fn weights_dir(&self) -> &AppWeightsDir {
     &self.weights_dir
   }
-  
+
+  pub fn temp_dir(&self) -> &TemporaryDir {
+    &self.temp_dir
+  }
+
   pub fn path(&self) -> &Path {
     &self.path
   }
@@ -121,3 +137,35 @@ impl AppWeightsDir {
     self.path.join(model_registry.get_filename())
   }
 }
+
+impl TemporaryDir {
+  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
+    let mut dir = dir.as_ref().to_path_buf();
+    match dir.canonicalize() {
+      Ok(d) => dir = d,
+      Err(err) => {
+        println!("Error canonicalizing {:?}: {}", dir, err);
+      }
+    }
+    if !dir.exists() {
+      println!("Creating directory {:?}", dir);
+      std::fs::create_dir(&dir)?;
+    }
+    Ok(Self {
+      path: dir,
+    })
+  }
+
+  pub fn with_prefix(&self, prefix: &str) -> anyhow::Result<TempDir> {
+    let tempdir = TempDir::new_in(&self.path, prefix)?;
+    Ok(tempdir)
+  }
+
+  pub fn new_named_temp_file(&self) -> anyhow::Result<NamedTempFile> {
+    let tempfile = Builder::new()
+      .prefix("temp_")
+      .suffix(".bin")
+      .tempfile_in(&self.path)?;
+    Ok(tempfile)
+  }
+}

crates/desktop/src-tauri/src/threads/downloader_thread.rs

@@ -1,11 +1,169 @@
-use tauri::{AppHandle, Emitter};
+use crate::events::notification_event::{NotificationEvent, NotificationModelType};
+use crate::ml::model_type::ModelType;
 use crate::state::app_dir::AppDataRoot;
+use crate::transfer::download::{download_async, ProgressUpdate};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use log::{debug, error, info};
+use std::collections::VecDeque;
+use std::path::PathBuf;
+use std::sync::mpsc::Receiver;
+use std::sync::{Arc, Mutex};
+use tauri::{AppHandle, Emitter};
+use tempfile::NamedTempFile;
+use tokio::task::JoinHandle;
+
+const NOTIFICATION_CHANNEL_NAME : &str = "notification";
+const MAX_FILES : usize = 8;
+const CHUNK_SIZE : usize = 1024 * 1024;
+const PARALLEL_FAILURES : usize = 8;
+const MAX_RETRIES : usize = 30;
+
+struct ModelDownloadStatusChannel {
+  pub model_name: String,
+  pub model_type: NotificationModelType,
+  pub final_file_path: PathBuf,
+  pub receiver: Receiver<ProgressUpdate>,
+}
 
 pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! {
-  loop {
-    app.emit("testing", 1);
+
+  let mut status_check_queue : Arc<Mutex<Vec<ModelDownloadStatusChannel>>> = Arc::new(Mutex::new(Vec::new()));
+
+  {
+    let status_check_queue2 = status_check_queue.clone();
+    let app2 = app.clone();
+
+    tokio::spawn(download_notify_thread(status_check_queue2, app2));
+  }
+
+  let mut download_queue = VecDeque::new();
+
+  // TODO: Automatic enqueue of known important models + enqueue new models on-demand
+  download_queue.push_back(ModelType::ClipJson);
+  download_queue.push_back(ModelType::SdxlTurboUnet);
+  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() {
+      info!("File already exists: {:?}", &filename);
+      continue;
+    }
+
+    let headers = None;
     
-    println!(">>> Testing: {:?}", app_data_root.path());
-    std::thread::sleep(std::time::Duration::from_millis(500));
+    let (tx, rx) = std::sync::mpsc::channel();
+
+    info!("Download: {:?}", filename);
+
+    let temp_file = match app_data_root.temp_dir().new_named_temp_file() {
+      Ok(file) => file,
+      Err(err) => {
+        error!("Couldn't create a temporary file: {:?}", err);
+        continue;
+      }
+    };
+
+    let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
+      model_name: model.get_name(),
+      model_type: model.get_notification_type(),
+    });
+
+    if let Err(err) = result {
+      error!("Could not emit event: {}", err);
+    }
+
+    let task = tokio::spawn(download_async(
+      url,
+      temp_file,
+      filename.clone(),
+      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(ModelDownloadStatusChannel {
+          receiver: rx,
+          final_file_path: filename.clone(),
+          model_name: model.get_name().to_string(),
+          model_type: model.get_notification_type(),
+        })
+      }
+    }
+
+    handles.push(task);
+  }
+
+  while let Some(result) = handles.next().await {
+    match result {
+      Err(err) => error!("Error downloading: {:?}", err),
+      Ok(_) => {},
+    }
+  }
+
+  loop {
+    // TODO(bt,2025-03-10): Queue management for new files entering the queue.
+    tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
+  }
+}
+
+pub async fn download_notify_thread(queue: Arc<Mutex<Vec<ModelDownloadStatusChannel>>>, 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(update) = item.receiver.try_recv() {
+            let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadProgress {
+              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 {
+              error!("Could not emit event: {}", err);
+            }
+          }
+
+          if item.final_file_path.exists() {
+            debug!("No longer sending notifications for {:?}", &item.final_file_path);
+
+            let result = app.emit(NOTIFICATION_CHANNEL_NAME, 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);
+          }
+        }
+
+        lock.extend(reenqueue);
+      }
+    }
+
+    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
   }
 }

crates/desktop/src-tauri/src/transfer/download.rs

@@ -0,0 +1,241 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use crate::transfer::error::Error;
+use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use reqwest::header::{
+  HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, CONTENT_RANGE,
+  RANGE,
+};
+use reqwest::Url;
+use std::collections::HashMap;
+use std::fmt::Display;
+use std::fs::remove_file;
+use std::io::SeekFrom;
+use std::path::Path;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::mpsc::Sender;
+use std::time::Duration;
+use log::info;
+use tempfile::NamedTempFile;
+use tokio::fs::OpenOptions;
+use tokio::io::AsyncWriteExt;
+use tokio::io::AsyncSeekExt;
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
+
+pub struct ProgressUpdate {
+  pub complete: usize,
+  pub total_length: usize,
+}
+
+/// max_files: Number of open file handles, which determines the maximum number of parallel downloads
+/// parallel_failures:  Number of maximum failures of different chunks in parallel (cannot exceed max_files)
+/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter)
+///
+/// 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 async fn download_async<P: AsRef<Path>>(
+  url: String,
+  filename: NamedTempFile,
+  final_filename: P,
+  max_files: usize,
+  chunk_size: usize,
+  parallel_failures: usize,
+  max_retries: usize,
+  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))
+    .build()?;
+
+  let mut headers = HeaderMap::new();
+  let mut auth_token = None;
+  if let Some(input_headers) = input_headers {
+    headers.reserve(input_headers.len());
+    for (k, v) in input_headers {
+      let name: HeaderName = k
+        .try_into()
+        .map_err(|err| Error::new_err(format!("Invalid header: {err}")))?;
+      let value: HeaderValue = AsRef::<str>::as_ref(&v)
+        .try_into()
+        .map_err(|err| Error::new_err(format!("Invalid header value: {err}")))?;
+      if name == AUTHORIZATION {
+        auth_token = Some(value);
+      } else {
+        headers.insert(name, value);
+      }
+    }
+  };
+
+  let response = if let Some(token) = auth_token.as_ref() {
+    client.get(&url).header(AUTHORIZATION, token)
+  } else {
+    client.get(&url)
+  }
+    .headers(headers.clone())
+    .header(RANGE, "bytes=0-0")
+    .send()
+    .await
+    .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?
+    .error_for_status()
+    .map_err(|err| Error::new_err(err.to_string()))?;
+
+  // Only call the final redirect URL to avoid overloading the Hub with requests and also
+  // altering the download count
+  let redirected_url = response.url();
+  if Url::parse(&url)
+    .map_err(|err| Error::new_err(format!("failed to parse url: {err}")))?
+    .host()
+    == redirected_url.host()
+  {
+    if let Some(token) = auth_token {
+      headers.insert(AUTHORIZATION, token);
+    }
+  }
+
+  let content_range = response
+    .headers()
+    .get(CONTENT_RANGE)
+    .ok_or(Error::new_err("No content length".to_string()))?
+    .to_str()
+    .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
+
+  let size: Vec<&str> = content_range.split('/').collect();
+  // Content-Range: bytes 0-0/702517648
+  // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Range
+  let length: usize = size
+    .last()
+    .ok_or(Error::new_err(
+      "Error while downloading: No size was detected".to_string(),
+    ))?
+    .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));
+  let max_chunk_complete = Arc::new(AtomicUsize::new(0));
+
+  for start in (0..length).step_by(chunk_size) {
+    let url = redirected_url.to_string();
+    let filename = filename.as_ref().to_owned();
+    let client = client.clone();
+    let headers = headers.clone();
+
+    let stop = std::cmp::min(start + chunk_size - 1, length);
+    let semaphore = semaphore.clone();
+    let parallel_failures_semaphore = parallel_failures_semaphore.clone();
+    let maybe_progress = maybe_progress_sender.clone();
+    let max_chunk_complete = max_chunk_complete.clone();
+    handles.push(tokio::spawn(async move {
+      let permit = semaphore
+        .acquire_owned()
+        .await
+        .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?;
+      let mut chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await;
+      let mut i = 0;
+      if parallel_failures > 0 {
+        while let Err(dlerr) = chunk {
+          if i >= max_retries {
+            return Err(Error::new_err(format!(
+              "Failed after too many retries ({max_retries}): {dlerr}"
+            )));
+          }
+          let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
+            Error::new_err(format!(
+              "Failed too many failures in parallel ({parallel_failures}): {dlerr} ({err})"
+            ))
+          })?;
+
+          let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
+          sleep(Duration::from_millis(wait_time as u64)).await;
+
+          chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await;
+          i += 1;
+          drop(parallel_failure_permit);
+          
+          if let Some(sender) = maybe_progress.as_ref() {
+            // NB: Calling fetch_max twice. First call will return the lower value.
+            let _old = max_chunk_complete.fetch_max(stop, Ordering::SeqCst);
+            let max_chunk = max_chunk_complete.fetch_max(stop, Ordering::SeqCst);
+            let progress_update = ProgressUpdate { 
+              complete: max_chunk,
+              total_length: length,
+            };
+            if let Err(_err) = sender.send(progress_update) {
+              // Fail silently, not essential to issue progress updates
+            }
+          }
+        }
+      }
+      drop(permit);
+      chunk.map_err(|e| Error::new_err(format!("Downloading error {e}"))).and(Ok(stop - start))
+    }));
+  }
+
+  // Output the chained result
+  while let Some(result) = handles.next().await {
+    match result {
+      Ok(Ok(size)) => {
+        //if let Some(ref callback) = callback {
+        //  callback.call((size,), None)?;
+        //}
+      }
+      Ok(Err(py_err)) => {
+        return Err(py_err);
+      }
+      Err(err) => {
+        return Err(Error::new_err(format!(
+          "Error while downloading: {err}"
+        )));
+      }
+    }
+  }
+
+  let file = filename.persist(final_filename)
+    .map_err(|e| Error::new_err(format!("Could not rename to final filename: {:?}", e)))?;
+
+  file.sync_all()?;
+
+  Ok(())
+}
+
+async fn download_chunk<P: AsRef<Path>>(
+  client: &reqwest::Client,
+  url: &str,
+  filename: P,
+  start: usize,
+  stop: usize,
+  headers: HeaderMap,
+) -> Result<(), Error> {
+  // Process each socket concurrently.
+  let range = format!("bytes={start}-{stop}");
+  let mut file = OpenOptions::new()
+    .write(true)
+    .truncate(false)
+    .create(true)
+    .open(filename)
+    .await?;
+  file.seek(SeekFrom::Start(start as u64)).await?;
+  let response = client
+    .get(url)
+    .headers(headers)
+    .header(RANGE, range)
+    .send()
+    .await?
+    .error_for_status()?;
+  let content = response.bytes().await?;
+  file.write_all(&content).await?;
+  Ok(())
+}

crates/desktop/src-tauri/src/transfer/error.rs

@@ -0,0 +1,49 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use std::fmt::Display;
+use reqwest::header::ToStrError;
+
+#[derive(Debug)]
+pub enum Error {
+  Io(std::io::Error),
+  Request(reqwest::Error),
+  ToStrError(ToStrError),
+  Misc(String),
+}
+
+impl From<std::io::Error> for Error {
+  fn from(value: std::io::Error) -> Self {
+    Self::Io(value)
+  }
+}
+
+impl From<reqwest::Error> for Error {
+  fn from(value: reqwest::Error) -> Self {
+    Self::Request(value)
+  }
+}
+
+impl From<ToStrError> for Error {
+  fn from(value: ToStrError) -> Self {
+    Self::ToStrError(value)
+  }
+}
+
+impl Display for Error {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    match self {
+      Self::Io(io) => write!(f, "Io: {io}"),
+      Self::Request(req) => write!(f, "Request: {req}"),
+      Self::ToStrError(req) => write!(f, "Response non ascii: {req}"),
+      Self::Misc(err) => write!(f, "Misc error: {err}"),
+    }
+  }
+}
+
+impl std::error::Error for Error {}
+
+impl Error {
+  pub fn new_err(reason: String) -> Self {
+    Self::Misc(reason)
+  }
+}

crates/desktop/src-tauri/src/transfer/mod.rs

@@ -0,0 +1,4 @@
+pub mod download;
+pub mod upload;
+pub mod util;
+pub mod error;

crates/desktop/src-tauri/src/transfer/upload.rs

@@ -0,0 +1,189 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use crate::transfer::error::Error;
+use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME};
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use reqwest::header::CONTENT_LENGTH;
+use std::collections::HashMap;
+use std::fmt::Display;
+use std::io::SeekFrom;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::fs::OpenOptions;
+use tokio::io::{AsyncReadExt, AsyncSeekExt};
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
+use tokio_util::codec::{BytesCodec, FramedRead};
+
+/// parts_urls: Dictionary consisting of part numbers as keys and the associated url as values
+/// completion_url: The url that should be called when the upload is finished
+/// max_files: Number of open file handles, which determines the maximum number of parallel uploads
+/// parallel_failures:  Number of maximum failures of different chunks in parallel (cannot exceed max_files)
+/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter)
+///
+/// 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
+///
+/// See https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html for more information
+/// on the multipart upload
+#[allow(clippy::too_many_arguments)]
+fn multipart_upload(
+  file_path: String,
+  parts_urls: Vec<String>,
+  chunk_size: u64,
+  max_files: usize,
+  parallel_failures: usize,
+  max_retries: usize,
+) -> Result<Vec<HashMap<String, String>>, 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 {
+      upload_async(
+        file_path,
+        parts_urls,
+        chunk_size,
+        max_files,
+        parallel_failures,
+        max_retries,
+      )
+        .await
+    })
+}
+
+
+
+#[allow(clippy::too_many_arguments)]
+async fn upload_async(
+  file_path: String,
+  parts_urls: Vec<String>,
+  chunk_size: u64,
+  max_files: usize,
+  parallel_failures: usize,
+  max_retries: usize,
+) -> Result<Vec<HashMap<String, String>>, Error> {
+  let client = reqwest::Client::new();
+
+  let mut handles = FuturesUnordered::new();
+  let semaphore = Arc::new(Semaphore::new(max_files));
+  let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures));
+
+  for (part_number, part_url) in parts_urls.iter().enumerate() {
+    let url = part_url.to_string();
+    let path = file_path.to_owned();
+    let client = client.clone();
+
+    let start = (part_number as u64) * chunk_size;
+    let semaphore = semaphore.clone();
+    let parallel_failures_semaphore = parallel_failures_semaphore.clone();
+    handles.push(tokio::spawn(async move {
+      let permit = semaphore
+        .clone()
+        .acquire_owned()
+        .await
+        .map_err(|err| Error::new_err(format!("Error acquiring semaphore: {err}")))?;
+      let mut chunk = upload_chunk(&client, &url, &path, start, chunk_size).await;
+      let mut i = 0;
+      if parallel_failures > 0 {
+        while let Err(ul_err) = chunk {
+          if i >= max_retries {
+            return Err(Error::new_err(format!(
+              "Failed after too many retries ({max_retries}): {ul_err}"
+            )));
+          }
+
+          let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| {
+            Error::new_err(format!(
+              "Failed too many failures in parallel ({parallel_failures}): {ul_err} ({err})"
+            ))
+          })?;
+
+          let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME);
+          sleep(Duration::from_millis(wait_time as u64)).await;
+
+          chunk = upload_chunk(&client, &url, &path, start, chunk_size).await;
+          i += 1;
+          drop(parallel_failure_permit);
+        }
+      }
+      drop(permit);
+      chunk.map_err(|e|{
+        //match e {
+        //  Error::Io(io) => Error::new_err(format!("Io error {io}")),
+        //  Error::Request(req) => Error::new_err(format!("Error while sending chunk {req}")),
+        //  Error::ToStrError(req) => Error::new_err(format!("Response header contains non ASCII chars: {req}")),
+        //}
+        e
+      }).map(|chunk| (part_number, chunk, chunk_size))
+    }));
+  }
+
+  let mut results: Vec<HashMap<String, String>> = vec![HashMap::default(); parts_urls.len()];
+
+  while let Some(result) = handles.next().await {
+    match result {
+      Ok(Ok((part_number, headers, size))) => {
+        //if let Some(ref callback) = callback {
+        //  callback.call((size,), None)?;
+        //}
+        results[part_number] = headers;
+      }
+      Ok(Err(py_err)) => {
+        return Err(py_err);
+      }
+      Err(err) => {
+        return Err(Error::new_err(format!(
+          "Error occurred while uploading: {err}"
+        )));
+      }
+    }
+  }
+
+  Ok(results)
+}
+
+async fn upload_chunk(
+  client: &reqwest::Client,
+  url: &str,
+  path: &str,
+  start: u64,
+  chunk_size: u64,
+) -> Result<HashMap<String, String>, Error> {
+  let mut options = OpenOptions::new();
+  let mut file = options.read(true).open(path).await?;
+  let file_size = file.metadata().await?.len();
+  let bytes_transferred = std::cmp::min(file_size - start, chunk_size);
+
+  file.seek(SeekFrom::Start(start)).await?;
+  let chunk = file.take(chunk_size);
+
+  let response = client
+    .put(url)
+    .header(CONTENT_LENGTH, bytes_transferred)
+    .body(reqwest::Body::wrap_stream(FramedRead::new(
+      chunk,
+      BytesCodec::new(),
+    )))
+    .send()
+    .await?;
+  let response = response.error_for_status()?;
+  let mut headers = HashMap::new();
+  for (name, value) in response.headers().into_iter() {
+    headers.insert(name.to_string(), value.to_str()?.to_owned());
+  }
+  Ok(headers)
+}
+

crates/desktop/src-tauri/src/transfer/util.rs

@@ -0,0 +1,15 @@
+//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs
+
+use rand::{thread_rng, Rng};
+use std::fmt::Display;
+
+pub const BASE_WAIT_TIME: usize = 300;
+pub const MAX_WAIT_TIME: usize = 10_000;
+
+fn jitter() -> usize {
+  thread_rng().gen_range(0..=500)
+}
+
+pub fn exponential_backoff(base_wait_time: usize, n: usize, max: usize) -> usize {
+  (base_wait_time + n.pow(2) + jitter()).min(max)
+}

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