From 2d18ca4454165dece8d028e80e56d7973a9473ea Mon Sep 17 00:00:00 2001
From: Arthur Brussee <arthur.brussee@gmail.com>
Date: Tue, 19 Nov 2024 11:17:35 +0000
Subject: [PATCH] Add training time to stats panel

---
 Cargo.lock                              |  1 +
 Cargo.toml                              |  1 +
 crates/brush-viewer/Cargo.toml          |  1 +
 crates/brush-viewer/src/panels/stats.rs | 12 +++++++++++-
 4 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3e6d85d4..7d0f85e2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -864,6 +864,7 @@ dependencies = [
  "egui_tiles",
  "glam",
  "gloo-timers",
+ "humantime",
  "image",
  "log",
  "parking_lot",
diff --git a/Cargo.toml b/Cargo.toml
index 4640aeb3..c21aacc8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -39,6 +39,7 @@ tokio-util = { version = "0.7.12", features = ["io"] }
 reqwest = { version = "0.12.9", features = ["stream"] }
 
 web-time = "1.1.0"
+humantime = "2.1.0"
 
 tracing-wasm = "0.2.1"
 async-fn-stream = "0.2.2"
diff --git a/crates/brush-viewer/Cargo.toml b/crates/brush-viewer/Cargo.toml
index e2784c6d..5cc69d9d 100644
--- a/crates/brush-viewer/Cargo.toml
+++ b/crates/brush-viewer/Cargo.toml
@@ -31,6 +31,7 @@ tracing-subscriber.workspace = true
 rand.workspace = true
 
 web-time.workspace = true
+humantime.workspace = true
 
 log.workspace = true
 gloo-timers.workspace = true
diff --git a/crates/brush-viewer/src/panels/stats.rs b/crates/brush-viewer/src/panels/stats.rs
index 110398e5..b985ba5b 100644
--- a/crates/brush-viewer/src/panels/stats.rs
+++ b/crates/brush-viewer/src/panels/stats.rs
@@ -4,7 +4,7 @@ use crate::{
 };
 use burn_jit::cubecl::Runtime;
 use burn_wgpu::{WgpuDevice, WgpuRuntime};
-use std::{collections::VecDeque, sync::Arc};
+use std::{collections::VecDeque, sync::Arc, time::Duration};
 use web_time::Instant;
 use wgpu::AdapterInfo;
 
@@ -19,6 +19,7 @@ pub(crate) struct StatsPanel {
     training_started: bool,
     num_splats: usize,
 
+    start_load_time: Instant,
     adapter_info: AdapterInfo,
 }
 
@@ -33,6 +34,7 @@ impl StatsPanel {
             last_eval_psnr: None,
             training_started: false,
             num_splats: 0,
+            start_load_time: Instant::now(),
             adapter_info,
         }
     }
@@ -66,6 +68,7 @@ impl ViewerPanel for StatsPanel {
     fn on_message(&mut self, message: crate::viewer::ViewerMessage, _: &mut ViewerContext) {
         match message {
             ViewerMessage::StartLoading { training } => {
+                self.start_load_time = Instant::now();
                 self.last_train_step = (Instant::now(), 0);
                 self.train_iter_per_s = 0.0;
                 self.train_iter_history.clear();
@@ -129,6 +132,13 @@ impl ViewerPanel for StatsPanel {
                         "--".to_owned()
                     });
                     ui.end_row();
+
+                    ui.label("Training time");
+                    // Round duration to seconds.
+                    let elapsed =
+                        Duration::from_secs((Instant::now() - self.start_load_time).as_secs());
+                    ui.label(format!("{}", humantime::Duration::from(elapsed)));
+                    ui.end_row();
                 }
 
                 let client = WgpuRuntime::client(&self.device);