implement cumprod

crates/desktop/src-tauri/Cargo.toml

@@ -47,4 +47,10 @@ tokenizers = "0.15.2"
 
 [features]
 default = ["cuda"]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+cuda = [
+  "candle-core/cuda",
+  # TODO(bt,2025-02-22): cuDNN might speed up Stable Diffusion
+  # "candle-core/cudnn",
+  "candle-nn/cuda",
+  "candle-transformers/cuda",
+]

crates/desktop/src-tauri/src/ml/cumprod.rs

@@ -0,0 +1,150 @@
+use candle_core::shape::Dim;
+use candle_core::Tensor;
+
+/// Tensor Cumulative Product
+/// 
+/// Candle doesn't implement cumprod, only cumsum.
+/// This implementation is adapted from cumsum in candle-ext:
+///  https://github.com/mokeyish/candle-ext/blob/main/tests/cumsum.rs
+///
+/// See Also: https://pytorch.org/docs/stable/generated/torch.cumprod.html
+/// See Also: https://numpy.org/doc/stable/reference/generated/numpy.cumprod.html
+/// See Also: https://github.com/huggingface/candle/issues/1646
+///
+pub fn cumprod<D: Dim>(input: &Tensor, dim: D) -> candle_core::Result<Tensor> {
+  let dim = dim.to_index(input.shape(), "cumprod")?;
+  let dim_size = input.dim(dim)?;
+  let rank = input.rank();
+  if rank == 0 {
+    return Ok(input.clone());
+  }
+
+  let mut tensors = Vec::with_capacity(dim_size);
+
+  let mut a = input.clone();
+  for i in 0..dim_size {
+    if i > 0 {
+      a = a.narrow(dim, 1, dim_size - i)?;
+      let b = input.narrow(dim, 0, dim_size - i)?;
+      a = (a * b)?;
+    }
+    tensors.push(a.narrow(dim, 0, 1)?);
+  }
+
+  let cumsum = Tensor::cat(&tensors, dim)?;
+  Ok(cumsum)
+}
+
+#[cfg(test)]
+mod tests {
+  use crate::ml::cumprod::cumprod;
+  use candle_core::{Device, Tensor};
+
+  mod base_cases {
+    use super::*;
+
+    #[test]
+    fn zeroes_vector() -> anyhow::Result<()> {
+      let a = Tensor::new(&[0i64, 0, 0, 0], &Device::Cpu)?;
+      let b = cumprod(&a, 0)?;
+      assert_eq!(b.to_vec1::<i64>()?, &[0, 0, 0, 0]);
+      Ok(())
+    }
+
+    #[test]
+    fn zeroes_tensor() -> anyhow::Result<()> {
+      let a = Tensor::new(&[
+        [0i64, 0, 0],
+        [0i64, 0, 0],
+        [0i64, 0, 0],
+      ], &Device::Cpu)?;
+
+      let b = cumprod(&a, 0)?;
+
+      assert_eq!(b.to_vec2::<i64>()?, &[
+        [0, 0, 0],
+        [0, 0, 0],
+        [0, 0, 0],
+      ]);
+      Ok(())
+    }
+
+    #[test]
+    fn ones_vector() -> anyhow::Result<()> {
+      let a = Tensor::new(&[1i64, 1, 1, 1], &Device::Cpu)?;
+      let b = cumprod(&a, 0)?;
+      assert_eq!(b.to_vec1::<i64>()?, &[1, 1, 1, 1]);
+      Ok(())
+    }
+
+    #[test]
+    fn ones_tensor() -> anyhow::Result<()> {
+      let a = Tensor::new(&[
+        [1i64, 1, 1],
+        [1i64, 1, 1],
+        [1i64, 1, 1],
+      ], &Device::Cpu)?;
+
+      let b = cumprod(&a, 0)?;
+
+      assert_eq!(b.to_vec2::<i64>()?, &[
+        [1, 1, 1],
+        [1, 1, 1],
+        [1, 1, 1],
+      ]);
+      Ok(())
+    }
+  }
+
+  mod test_cases {
+    use super::*;
+
+    #[test]
+    fn vector() -> anyhow::Result<()> {
+      let a = Tensor::new(&[1i64, 2, 3, 4, 5], &Device::Cpu)?;
+      let b = cumprod(&a, 0)?;
+      assert_eq!(b.to_vec1::<i64>()?, &[1, 2, 6, 24, 120]);
+
+      let a = Tensor::new(&[0, 1i64, 2, 3, 4, 5], &Device::Cpu)?;
+      let b = cumprod(&a, 0)?;
+      assert_eq!(b.to_vec1::<i64>()?, &[0, 0, 0, 0, 0, 0]);
+      Ok(())
+    }
+
+    #[test]
+    fn tensor_dim_0() -> anyhow::Result<()> {
+      let a = Tensor::new(&[
+        [1i64, 2, 3],
+        [4i64, 5, 6],
+        [7i64, 8, 9],
+      ], &Device::Cpu)?;
+
+      let b = cumprod(&a, 0)?;
+
+      assert_eq!(b.to_vec2::<i64>()?, &[
+        [1, 2, 3],
+        [4, 10, 18],
+        [28, 80, 162],
+      ]);
+      Ok(())
+    }
+
+    #[test]
+    fn tensor_dim_1() -> anyhow::Result<()> {
+      let a = Tensor::new(&[
+        [1i64, 2, 3],
+        [4i64, 5, 6],
+        [7i64, 8, 9],
+      ], &Device::Cpu)?;
+
+      let b = cumprod(&a, 1)?;
+
+      assert_eq!(b.to_vec2::<i64>()?, &[
+        [1, 2, 6],
+        [4, 20, 120],
+        [7, 56, 504],
+      ]);
+      Ok(())
+    }
+  }
+}

crates/desktop/src-tauri/src/ml/lcm_scheduler.rs

@@ -0,0 +1,474 @@
+use anyhow::anyhow;
+use candle_core::{DType, Device, Shape, Tensor, WithDType};
+use candle_transformers::models::stable_diffusion::schedulers::Scheduler;
+use candle_transformers::models::stable_diffusion::utils::linspace;
+use log::error;
+
+/*
+Notes on LCM implementation:
+
+  - Robust technical comparison of samplers: https://stable-diffusion-art.com/samplers/
+  - Apple Swift LCM Implementation ticket: https://github.com/apple/ml-stable-diffusion/issues/319
+  - Reference Swift Implementation: https://github.com/GuernikaCore/Schedulers/blob/main/Sources/Schedulers/LCMScheduler.swift
+
+Candle is missing several "batteries" we'll have to implement, and some we might want to upstream:
+
+  - t[::-1] (syntax to reverse the tensor)
+
+  - torch/numpy.cumprod :
+      - https://numpy.org/doc/stable/reference/generated/numpy.cumprod.html
+      - https://pytorch.org/docs/stable/generated/torch.cumprod.html
+
+Remaining work:
+
+  - [ ] implement add_noise
+  - [ ] implement cumprod
+  - [ ] initialize step_index / begin_index
+  - [ ] find out how `set_timestamps` is called in diffusers and see if we want to change it to CTOR param in candle
+  - [ ] try to plug this into candle and run
+  - [ ] fix timesteps() method to cache values
+  - [ ] compare with diffusers
+  - [ ] compare with candle schedulers
+  - [ ] compare with swift schedulers
+  - [ ] test + unit test
+*/
+
+struct LcmScheduler {
+  alphas: Tensor,
+  alphas_cumprod: Tensor,
+  final_alpha_cumprod: Tensor,
+
+  betas: Tensor,
+
+  timesteps: Tensor,
+
+  // Standard deviation of the initial noise distribution
+  // NB: Only in diffusers, and seems to only hold the value "1.0". Swift implementation omits it.
+  init_noise_sigma: f64,
+
+  // TODO: This is set by the set_timesteps() method after construction in python.
+  //  We might want to handle this upfront in the constructor ourselves
+  num_inference_steps: Option<usize>,
+
+  _step_index: Option<usize>,
+  _begin_index: Option<usize>,
+}
+
+impl LcmScheduler {
+  pub fn new(
+    num_train_timestamps: i64,
+    beta_start: f64,
+    beta_end: f64,
+    beta_schedule: &str,
+    trained_betas: Option<Tensor>,
+    set_alpha_to_one: bool,
+    rescale_betas_zero_snr: bool,
+    device: &Device,
+  ) -> anyhow::Result<Self> {
+    let betas;
+
+    if let Some(trained) = trained_betas {
+      betas = trained;
+    } else {
+      betas = match beta_schedule {
+        "scaled_linear" => {
+          // this schedule is very specific to the latent diffusion model.
+          let t = linspace(
+            beta_start.powf(0.5),
+            beta_end.powf(0.5),
+            num_train_timestamps as usize
+          )?;
+          t.to_device(device)?
+        }
+        _ => {
+          return Err(anyhow!("schedule {beta_schedule} is not implemented"))
+        }
+      };
+    }
+
+    if rescale_betas_zero_snr {
+      return Err(anyhow!("rescale_betas_zero_snr is not used upstream"));
+    }
+
+    let alphas = (1.0 - &betas)?;
+
+    // TODO: Implement cumprod: https://github.com/huggingface/candle/issues/1646
+    //  https://numpy.org/doc/stable/reference/generated/numpy.cumprod.html
+    error!("THIS IS USING CUMSUM, NOT CUMPROD. WE NEED TO IMPLEMENT CUMPROD.");
+    // self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+    let alphas_cumprod = alphas.cumsum(0)?;
+
+    // TODO: Check implementation
+    // At every step in ddim, we are looking into the previous alphas_cumprod
+    // For the final step, there is no previous alphas_cumprod because we are already at 0
+    // `set_alpha_to_one` decides whether we set this parameter simply to one or
+    // whether we use the final alpha of the "non-previous" one.
+    let final_alpha_cumprod = if set_alpha_to_one {
+      initialize_scalar_tensor(1.0, device)? // NB: `set_alpha_to_one` is false in python
+    } else {
+      alphas_cumprod.get(0)?
+    };
+
+    // standard deviation of the initial noise distribution
+    let init_noise_sigma = 1.0;
+
+    let timesteps = initialize_timesteps(num_train_timestamps, device)?;
+
+    Ok(Self {
+      alphas,
+      alphas_cumprod,
+      final_alpha_cumprod,
+      betas,
+      init_noise_sigma,
+      timesteps,
+      num_inference_steps: None,
+      _step_index: None,
+      _begin_index: None,
+    })
+  }
+
+  pub fn set_begin_index(&mut self, begin_index: usize) {
+    self._begin_index = Some(begin_index);
+  }
+
+  pub fn init_step_index(&mut self, timestep: usize) -> candle_core::Result<()> {
+    match self._begin_index {
+      None => {
+        todo!("implement begin index none branch")
+      }
+      Some(begin_index) => {
+        self._step_index = Some(begin_index);
+      }
+    }
+    Ok(())
+  }
+
+  pub fn set_timesteps(&mut self, num_inference_steps: usize, device: &Device) -> candle_core::Result<()> {
+    // NB: The diffusers version of this method passes `None` for `original_inference_steps` and `timesteps`.
+    // It passes strength = 1.0, and config.original_inference_steps = 50.
+    const STRENGTH : f64 = 1.0;
+    const CONFIG_ORIGINAL_INFERENCE_STEPS : usize = 50;
+    const CONFIG_NUM_TRAIN_TIMESTEPS : usize = 1000;
+
+    let original_steps = CONFIG_ORIGINAL_INFERENCE_STEPS;
+
+    // LCM Timesteps Setting
+    // The skipping step parameter k from the paper.
+    let k = (CONFIG_NUM_TRAIN_TIMESTEPS as f64 / original_steps as f64).floor();
+
+    let end = (original_steps as f64 * STRENGTH) as i64 + 1;
+    let lcm_origin_timesteps = ((Tensor::arange::<i64>(1, end, device)? * k)? - 1.0)?;
+
+    if num_inference_steps > CONFIG_NUM_TRAIN_TIMESTEPS {
+      return Err(candle_core::Error::Msg("num_inference_steps cannot be larger than train timesteps".to_string()));
+    }
+
+    let skipping_step = (lcm_origin_timesteps.dims().len() as f64 / num_inference_steps as f64).floor() as usize;
+
+    self.num_inference_steps = Some(num_inference_steps);
+    
+    // TODO: This is not efficient
+    let lcm_origin_timesteps = lcm_origin_timesteps.to_vec1::<i64>()?
+      .into_iter()
+      .map(i64::from)
+      .rev()
+      .collect::<Vec<i64>>();
+    
+    // TODO: This might not be exact. We should unit test this
+    let inference_indices = linspace(0.0, lcm_origin_timesteps.len() as f64, num_inference_steps)?;
+    
+    let inference_indices = inference_indices.floor()?
+      .to_dtype(DType::I64)?;
+    
+    let timestamps = {
+      // TODO: This code is awful
+      let inference_indices_t: Vec<usize> = inference_indices.to_vec1::<i64>()?
+        .into_iter()
+        .map(|n| n as usize)
+        .collect();
+
+      let mut timesteps: Vec<i64> = Vec::with_capacity(inference_indices.dims().len());
+      
+      for i in inference_indices_t {
+        if let Some(val) = lcm_origin_timesteps.get(i) {
+          timesteps.push(*val);
+        }
+      }
+      
+      let shape = Shape::from_dims(&[timesteps.len()]);
+      
+      Tensor::from_vec(timesteps, shape, device)?
+        .to_dtype(DType::I64)?
+        .to_device(device)?
+    };
+
+    self.timesteps = timestamps;
+
+    Ok(())
+  }
+
+}
+
+impl Scheduler for LcmScheduler {
+  fn timesteps(&self) -> &[usize] {
+    // TODO: We can't use the implementation from DDIM. This method requires
+    //  (1) an infallible response and (2) a reference we can't dynamically allocate. We must
+    //  cache this somewhere.
+    &[]
+  }
+
+  fn add_noise(&self, original_samples: &Tensor, noise: Tensor, timestep: usize) -> candle_core::Result<Tensor> {
+    // TODO: Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+    //  Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
+    //  for the subsequent add_noise calls
+    let mut alphas_cumprod = self.alphas_cumprod.to_device(original_samples.device())?;
+    alphas_cumprod = alphas_cumprod.to_dtype(original_samples.dtype())?;
+    //let timesteps = timesteps.to_device(original_samples.device())?;
+
+    let sqrt_alpha_prod = alphas_cumprod.get(timestep)?.powf(0.5)?;
+    let mut sqrt_alpha_prod = sqrt_alpha_prod.flatten_all()?; // TODO: Verify that flatten_all works here
+
+
+    while sqrt_alpha_prod.dims().len() < original_samples.dims().len() {
+      // TODO: Verify that this works for `unsqueeze(-1)`
+      sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(sqrt_alpha_prod.dims().len())?;
+    }
+
+    let sqrt_one_minus_alpha_prod = (1.0 - alphas_cumprod.get(timestep)?)?.powf(0.5)?;
+    let mut sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten_all()?; // TODO: Verify flatten all works here.
+
+    while sqrt_one_minus_alpha_prod.dims().len() < original_samples.dims().len() {
+      // TODO: Verify that this works for `unsqueeze(-1)`
+      sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(sqrt_one_minus_alpha_prod.dims().len())?;
+    }
+
+    let noisy_samples = ((sqrt_alpha_prod * original_samples)? + (sqrt_one_minus_alpha_prod * noise)?)?;
+
+    Ok(noisy_samples)
+  }
+
+  fn init_noise_sigma(&self) -> f64 {
+    self.init_noise_sigma
+  }
+
+  /// Ensures interchangeability with schedulers that need to scale the denoising model input
+  /// depending on the current timestep. (For LCM, this is a pass through no-op!)
+  fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> candle_core::Result<Tensor> {
+    Ok(sample) // NB: The diffusers scheduling_lcm.py implementation is a pass through no-op !
+  }
+
+  /// Predict the sample from the previous timestep by reversing the SDE. This function
+  /// propagates the diffusion process from the learned model outputs (most often the
+  /// predicted noise).
+  fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> candle_core::Result<Tensor> {
+
+    if self.num_inference_steps.is_none() {
+      return Err(candle_core::Error::Msg(
+        "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+          .to_string()));
+    }
+
+    if self._step_index.is_none() {
+      self.init_step_index(timestep)?;
+    }
+
+    let step_index = self._step_index.clone().expect("unsafe unwrap; should exist per above");
+
+    // 1. Get previous step value
+    let prev_step_index = step_index + 1;
+    let prev_timestep;
+
+    if prev_step_index < self.timesteps.elem_count() {
+      prev_timestep = self.timesteps.get(timestep)?.to_scalar::<i64>()? as usize;
+    } else {
+      prev_timestep = timestep;
+    }
+
+    // 2. Compute alphas, betas
+    let alpha_prod_t = self.alphas_cumprod.get(timestep)?;
+    let alpha_prod_t_prev = if prev_timestep >= 0 {
+      self.alphas_cumprod.get(prev_timestep)?
+    } else {
+      self.final_alpha_cumprod.clone()
+    };
+
+    let beta_prod_t = (1. - &alpha_prod_t)?;
+    let beta_prod_t_prev = (1. - &alpha_prod_t_prev)?;
+
+    // 3. Get scaling for boundary conditions
+    // NB: timestep_scaling seems to be always set to 10.0 in diffusers examples, but it is passed
+    // via a config object. Hardcoding this for now until we need to change it.
+    const TIMESTEP_SCALING : f64 = 10.0;
+    let (c_skip, c_out) = get_scalings_for_boundary_condition_discrete(timestep, TIMESTEP_SCALING);
+
+    // 4. Compute the predicted original sample x_0 based on the model parameterization
+    // NB: Prediction type is "epsilon" in the diffusers examples, but that is configurable.
+    // This is the only one we'll implement for now.
+    let predicted_original_sample = ((sample - (beta_prod_t.sqrt()? * model_output)?)? / &alpha_prod_t.sqrt()?)?;
+
+    // 5. (skip clip/threshold - not used in diffusers example)
+
+    // 6. Denoise model output using boundary conditions
+    let denoised = ((c_out * predicted_original_sample)? + (c_skip * sample)?)?;
+
+    // TODO: Step 7. We have to return the previous example somehow.
+    //  Diffusers tolerates it not being sent back, but the quality degrades as if one step is absent.
+    Ok(denoised)
+  }
+}
+
+fn initialize_scalar_tensor<D: WithDType>(value: D, device: &Device) -> anyhow::Result<Tensor> {
+  let scalar = Tensor::from_vec(vec![value], Shape::from(vec![]), device)?;
+  Ok(scalar)
+}
+
+// TODO(bt,2025-02-23): Use native implementations
+fn initialize_timesteps(num_train_timesteps: i64, device: &Device) -> anyhow::Result<Tensor> {
+  // 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::<i64>(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)
+}
+
+fn get_scalings_for_boundary_condition_discrete(timestep: usize, timestep_scaling: f64) -> (f64, f64) {
+  let sigma_data = 0.5f64; // Default 0.5
+  let scaled_timestep = timestep as f64 * timestep_scaling;
+
+  let c_skip = sigma_data.powf(2.0) / (scaled_timestep.powf(2.0) + sigma_data.powf(2.0));
+  let c_out = scaled_timestep / (scaled_timestep.powf(2.0) + sigma_data.powf(2.0)).powf(0.5);
+
+  (c_skip, c_out)
+}
+
+// The Candle sources, and moreover, the Candle tests, are a great source of information for how the library works.
+//assert_eq!(t.i(..-1)?, 1.0);
+#[cfg(test)]
+mod tests {
+  use crate::ml::lcm_scheduler::{initialize_scalar_tensor, initialize_timesteps};
+  use crate::ml::lcm_scheduler::tests::test_helper_fn::*;
+  use candle_core::Device;
+
+  mod helper_function_unit_tests {
+    use crate::ml::lcm_scheduler::get_scalings_for_boundary_condition_discrete;
+    use super::*;
+
+    #[test]
+    fn test_initialize_scalar_tensor() -> anyhow::Result<()> {
+      let t = initialize_scalar_tensor(1.0, &Device::Cpu)?;
+      let empty: [usize; 0] = [];
+      assert_eq!(t.dims(), &empty);
+      assert_eq!(t.to_vec0::<f64>()?, 1.0);
+      Ok(())
+    }
+
+    #[test]
+    fn test_initialize_timesteps() -> anyhow::Result<()> {
+      let tensor = initialize_timesteps(5, &Device::Cpu)?;
+      assert_eq!(tensor.to_vec1::<i64>()?, &[4, 3, 2, 1,0]);
+      Ok(())
+    }
+
+    #[test]
+    fn test_get_scalings_for_boundary_condition_discrete() -> anyhow::Result<()> {
+      // Base cases
+      let v = get_scalings_for_boundary_condition_discrete(0, 0.0);
+      assert_eq!(v, (1.0, 0.0));
+      let v = get_scalings_for_boundary_condition_discrete(1, 0.0);
+      assert_eq!(v, (1.0, 0.0));
+      let v = get_scalings_for_boundary_condition_discrete(0, 1.0);
+      assert_eq!(v, (1.0, 0.0));
+
+      // Production code has "timestep_scaling" = 10.0
+      let v = get_scalings_for_boundary_condition_discrete(0, 10.0);
+      assert_eq!(v, (1.0, 0.0));
+
+      // Production code has "timestep_scaling" = 10.0
+      let v = get_scalings_for_boundary_condition_discrete(1, 10.0);
+      assert!((v.0 - 0.0024937655860349127).abs() < 0.01);
+      assert!((v.1 - 0.9987523388778445).abs() < 0.01);
+
+      // TODO: More test cases -
+      //let v = get_scalings_for_boundary_condition_discrete(10, 10.0);
+      //let v = get_scalings_for_boundary_condition_discrete(100, 10.0);
+
+      Ok(())
+    }
+  }
+
+  mod confirm_various_pytorch_behaviors {
+    use candle_core::IndexOp;
+    use super::*;
+
+    // NB: Non-unit tests, just to test understanding while porting code
+
+    #[test]
+    fn unsqueeze_zero() -> anyhow::Result<()> {
+      let t = tensor_1234()?;
+      let x = t.unsqueeze(0)?;
+      assert_eq!(x.to_vec2::<i64>()?, &[&[1, 2, 3, 4]]);
+      Ok(())
+    }
+
+    #[test]
+    fn unsqueeze_simulating_negative_one() -> anyhow::Result<()> {
+      let t = tensor_1234()?;
+      // NB: It looks like this matches the behavior of unsqueeze(-1)
+      // Per: https://pytorch.org/docs/stable/generated/torch.unsqueeze.html
+      let x = t.unsqueeze(t.dims().len())?;
+      assert_eq!(x.to_vec2::<i64>()?, &[&[1],&[2],&[3],&[4]]);
+      Ok(())
+    }
+
+    #[test]
+    fn one_minus_vector() -> anyhow::Result<()> {
+      let t = tensor_01234()?;
+      let x = (1. - t)?;
+      assert_eq!(x.rank(), 1); // NB: 1-dimensional output
+      assert_eq!(x.to_vec1::<i64>()?, &[1, 0, -1, -2, -3]);
+      Ok(())
+    }
+    
+    // #[test]
+    // fn reverse_vector() -> anyhow::Result<()> {
+    //   let t = tensor_1234()?;
+    //   let x = t.i(5..0)?;
+    //   assert_eq!(x.rank(), 1); // NB: 1-dimensional output
+    //   assert_eq!(x.to_vec1::<i64>()?, &[1, 0, -1, -2, -3]);
+    //   Ok(())
+    // }
+  }
+
+  pub mod test_helper_fn {
+    use candle_core::{Device, Shape, Tensor};
+
+    pub fn tensor_1234() -> anyhow::Result<Tensor> {
+      // Return an example tensor
+      let data: Vec<i64> = vec![1, 2, 3, 4];
+      let t = Tensor::from_vec(data, Shape::from(vec![4]), &Device::Cpu)?;
+      assert_eq!(t.to_vec1::<i64>()?, &[1, 2, 3, 4]);
+      Ok(t)
+    }
+
+    pub fn tensor_01234() -> anyhow::Result<Tensor> {
+      // Return an example tensor
+      let data: Vec<i64> = vec![0, 1, 2, 3, 4];
+      let t = Tensor::from_vec(data, Shape::from(vec![5]), &Device::Cpu)?;
+      assert_eq!(t.to_vec1::<i64>()?, &[0, 1, 2, 3, 4]);
+      Ok(t)
+    }
+  }
+}

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

@@ -1,4 +1,6 @@
+pub mod cumprod;
 pub mod image;
+pub mod lcm_scheduler;
 pub mod model_cache;
 pub mod model_file;
 pub mod models;

frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts

@@ -533,7 +533,7 @@ export class RealTimeDrawEngine {
     })) as ImageBitmap;
 
     // Test code
-    if (true) {
+    if (false) {
       this.outputBitmap = bitmap;
       this.previewCanvas.image(bitmap);
       this.isProcessing = false;