use super::*; #[derive(Debug)] pub(super) enum ModelDownload { Idle, Active(ActiveDownload), Complete(ManagedArtifactId, ModelOperation, DownloadProgress), Failed(ManagedArtifactId, String, DownloadProgress), } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(super) enum ModelOperation { Download, Validate, } #[derive(Debug)] pub(super) struct ActiveDownload { pub(super) artifact: ManagedArtifactId, pub(super) operation: ModelOperation, pub(super) progress: DownloadProgress, pub(super) sampled_at: Instant, pub(super) sampled_bytes: u64, pub(super) bytes_per_second: f64, pub(super) verified_bytes: Arc, pub(super) cancel: Arc, pub(super) result: mpsc::Receiver>, pub(super) stopping: bool, } impl App { pub(super) fn open_model_manager(&mut self) -> Task { if let Some(id) = self.model_manager_window { return window::gain_focus(id); } let (id, open) = window::open(window::Settings { size: Size::new(760.0, 560.0), min_size: Some(Size::new(620.0, 420.0)), icon: Some(app_icon()), ..Default::default() }); self.model_manager_window = Some(id); open.map(Message::ModelManagerOpened) } pub(super) fn start_model_operation( &mut self, artifact: ManagedArtifactId, operation: ModelOperation, ) { if matches!(self.model_download, ModelDownload::Active(_)) { return; } let models_path = models_path(); let progress = match operation { ModelOperation::Download => model::artifact_download_progress(artifact, &models_path), ModelOperation::Validate => model::artifact_verification_progress(artifact, 0), }; let cancel = Arc::new(AtomicBool::new(false)); let worker_cancel = Arc::clone(&cancel); let verified_bytes = Arc::new(AtomicU64::new(0)); let worker_verified_bytes = Arc::clone(&verified_bytes); let (result_sender, result_receiver) = mpsc::channel(); let thread_name = match operation { ModelOperation::Download => "model-download", ModelOperation::Validate => "model-validation", }; if let Err(error) = thread::Builder::new() .name(thread_name.to_owned()) .spawn(move || { let result = match operation { ModelOperation::Download => model::download_managed_artifact( artifact, &models_path, &worker_cancel, &worker_verified_bytes, ), ModelOperation::Validate => model::validate_managed_artifact( artifact, &models_path, &worker_cancel, &worker_verified_bytes, ), }; let _ = result_sender.send(result); }) { self.error = Some(format!("Could not start {thread_name}: {error}")); return; } self.pending_model_delete = None; self.model_download = ModelDownload::Active(ActiveDownload { artifact, operation, sampled_at: Instant::now(), sampled_bytes: progress.completed(), progress, bytes_per_second: 0.0, verified_bytes, cancel, result: result_receiver, stopping: false, }); } } impl App { pub(super) fn update_download_progress(&mut self) { let ModelDownload::Active(download) = &mut self.model_download else { return; }; let now = Instant::now(); let verified = download.verified_bytes.load(Ordering::Relaxed); let mut progress = match download.operation { ModelOperation::Download => { model::artifact_download_progress(download.artifact, &models_path()) } ModelOperation::Validate => { model::artifact_verification_progress(download.artifact, verified) } }; if download.operation == ModelOperation::Download && let Some(verification) = &mut progress.verification { verification.verified = verified.min(verification.total); } let elapsed = now.duration_since(download.sampled_at).as_secs_f64(); let completed = progress.completed(); let phase_changed = progress.phase != download.progress.phase; let transferred = completed.saturating_sub(download.sampled_bytes); if phase_changed { download.bytes_per_second = 0.0; } else if transferred > 0 && elapsed > 0.0 { let current = transferred as f64 / elapsed; download.bytes_per_second = if download.bytes_per_second == 0.0 { current } else { download.bytes_per_second * 0.75 + current * 0.25 }; } download.progress = progress; download.sampled_at = now; download.sampled_bytes = completed; let result = match download.result.try_recv() { Ok(result) => Some(result), Err(TryRecvError::Empty) => None, Err(TryRecvError::Disconnected) => { Some(Err("Download worker stopped unexpectedly.".into())) } }; let artifact = download.artifact; let operation = download.operation; let progress = download.progress.clone(); if let Some(result) = result { match result { Ok(DownloadOutcome::Complete) => { let progress = model::artifact_download_progress(artifact, &models_path()); self.model_download = ModelDownload::Complete(artifact, operation, progress); self.error = None; } Ok(DownloadOutcome::Stopped) => { self.model_download = ModelDownload::Idle; self.error = None; } Err(error) => { self.error = Some(error.clone()); self.model_download = ModelDownload::Failed(artifact, error, progress); } } } } }