From ee5efc378e0c2f37f8dc10077190f4a95bbe47ee Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 23 Feb 2025 14:56:49 -0500 Subject: [PATCH] implement timestep initialization --- .../desktop/src-tauri/src/ml/lcm_scheduler.rs | 38 +++++++++++++------ 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs b/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs index 226e4beb4b..7661022ee1 100644 --- a/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs +++ b/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs @@ -131,7 +131,7 @@ impl LcmScheduler { // NB: This is generating a range and reversing it. // self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - let timesteps = Tensor::arange(0i64, num_train_timestamps, device)?; + let timesteps = initialize_timesteps(num_train_timestamps, device)?; /* ##### All instance variables: #### @@ -347,23 +347,39 @@ impl Scheduler for LcmScheduler { } } + +// TODO(bt,2025-02-23): Use native implementations +fn initialize_timesteps(num_train_timesteps: i64, device: &Device) -> anyhow::Result { + // Testing LCM scheduler assumptions: + // [::-1] is shorthand for reversing a numpy vector, eg. [0, 1,2,3,4] -> [4, 3, 2, 1, 0] + // self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) + //let tensor = Tensor::arange::(0, 5, &Device::Cpu)?; // This works, but we can't reverse/flip it! + + // https://github.com/tracel-ai/burn/pull/1468 flip impl? Will this work in place? + // https://github.com/huggingface/candle/issues/1875 + + let data = (0..num_train_timesteps) + .map(i64::from) + .rev() + .collect(); + + let shape = Shape::from_dims(&[num_train_timesteps as usize]); + + let tensor = Tensor::from_vec(data, shape, device)?; + + Ok(tensor) +} + // The Candle sources, and moreover, the Candle tests, are a great source of information for how the library works. #[cfg(test)] mod tests { use candle_core::{Device, Tensor}; + use crate::ml::lcm_scheduler::initialize_timesteps; #[test] fn test_implementation_of_reverse() -> anyhow::Result<()> { - // Testing LCM scheduler assumptions: - // [::-1] is shorthand for reversing a numpy vector, eg. [0, 1,2,3,4] -> [4, 3, 2, 1, 0] - // self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - - let tensor = Tensor::arange::(0, 5, &Device::Cpu)?; - - assert_eq!(tensor.to_vec1::()?, &[0,1,2,3,4]); - - // TODO: Reverse the shape - + let tensor = initialize_timesteps(5, &Device::Cpu)?; + assert_eq!(tensor.to_vec1::()?, &[4,3,2,1,0]); Ok(()) } } \ No newline at end of file