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/Engine.ts

@@ -1,5 +1,6 @@
 import Konva from "konva";
 
+
 import { uiAccess, uiEvents } from "~/signals";
 import { ShapeNode } from "./Nodes";
 
@@ -98,6 +99,7 @@ export class Engine {
     this.stage.add(this.nodeIsolationLayer);
     this.stage.add(this.uiLayer);
 
+
     // Konva Transformer
     this.nodeTransformer = new NodeTransformer();
     this.uiLayer.add(this.nodeTransformer.getKonvaNode());
@@ -118,7 +120,7 @@ export class Engine {
       height: VideoResolutions.SQUARE_1024.height,
       mediaLayerRef: this.mediaLayer,
       bgLayerRef: this.bgLayer,
-      offScreenCanvas: this.offScreenCanvas,
+      offScreenCanvas: this.offScreenCanvas
     });
 
     // Collection of all Nodes
@@ -178,8 +180,13 @@ export class Engine {
     // hence, lastly, setup these events
     this.setupEventSystem();
     this.setAppMode(AppModes.SELECT);
+
+
+
   }
 
+
+
   private setAppMode(newAppMode: AppModes) {
     this.appMode = newAppMode;
     switch (this.appMode) {
@@ -361,18 +368,7 @@ export class Engine {
 
     uiEvents.toolbarMain.SAVE.onClick(async (/*event*/) => {
       await this.realTimeDrawEngine.saveOutput();
-    });
 
-    uiEvents.toolbarMain.SELECT.onClick(() => {
-      console.log("Toolbar Main >> Select");
-    });
-
-    uiEvents.toolbarMain.PAINT.onClick(() => {
-      console.log("Toolbar Main >> Paint");
-    });
-
-    uiEvents.toolbarMain.ERASER.onClick(() => {
-      console.log("Toolbar Main >> Eraser");
     });
 
     uiEvents.toolbarMain.PREVIEW.onClick(async () => {
@@ -403,6 +399,8 @@ export class Engine {
       this.addImage(image);
     });
 
+
+
     uiEvents.onAddTextToEngine((textdata) => {
       this.addText(textdata);
     });
@@ -467,7 +465,7 @@ export class Engine {
         return;
       }
       this.changeNodeColor(nodeColor);
-    });
+    })
   }
 
   disableAllButtons() {
@@ -503,7 +501,7 @@ export class Engine {
   }
 
   // Sandbox is quickly a way to test your idea.
-  public async sandbox() {}
+  public async sandbox() { }
 
   public onMessage(event: MessageEvent) {
     console.log("Message From Shared Worker");
@@ -531,9 +529,9 @@ export class Engine {
       x: 10,
       y: 80,
       text: "",
-      fontSize: 18,
+      fontSize: 24,
       fontFamily: "Source Sans 3",
-      fill: "white",
+      fill: "black",
     });
     const anim = new Konva.Animation((frame) => {
       if (frame && import.meta.env.DEV) {
@@ -551,34 +549,32 @@ export class Engine {
   }
 
   private changeNodeColor(nodeColor: NodeColor) {
-    console.log("Changing node color:", nodeColor);
+    console.log('Changing node color:', nodeColor);
     // Find the node that is selected with the id
-    const idSelector = "#" + nodeColor.kNodeId;
+    const idSelector = '#' + nodeColor.kNodeId;
     const selectedNode = this.mediaLayer.findOne(idSelector);
-    console.log("Selected node:", selectedNode);
-    console.log("ID selector:", idSelector);
-    console.log("Media layer:", this.mediaLayer);
+    console.log('Selected node:', selectedNode);
+    console.log('ID selector:', idSelector);
+    console.log('Media layer:', this.mediaLayer);
     // TODO: Change the color of the node
     // Turns out shapes are image nodes... how do we change the image's fill??
     if (selectedNode) {
       // Get the shape node from the real time drawing engine's image nodes
-      const shapeNode = this.realTimeDrawEngine.findImageNodeById(
-        nodeColor.kNodeId,
-      );
+      const shapeNode = this.realTimeDrawEngine.findImageNodeById(nodeColor.kNodeId);
       if (shapeNode && shapeNode instanceof ShapeNode) {
         // Create new shape node with updated color
         const newShapeNode = new ShapeNode({
           canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
-          canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
+          canvasSize: this.realTimeDrawEngine.captureCanvas.size(), 
           shapeType: shapeNode.shapeType,
           size: shapeNode.size,
           color: nodeColor.color,
           mediaLayerRef: this.mediaLayer,
-          selectionManagerRef: this.selectionManager,
+          selectionManagerRef: this.selectionManager
         });
         newShapeNode.kNode.position(shapeNode.kNode.position());
         newShapeNode.kNode.zIndex(shapeNode.kNode.zIndex());
-
+        
         // Remove old node and add new one
         shapeNode.kNode.destroy();
         this.commandManager.createNode(newShapeNode);
@@ -605,7 +601,7 @@ export class Engine {
       size: size,
       color: color,
       mediaLayerRef: this.mediaLayer,
-      selectionManagerRef: this.selectionManager,
+      selectionManagerRef: this.selectionManager
     });
     shapeNode.kNode.zIndex(1); // Replace desiredZIndex with a number
 
@@ -669,7 +665,7 @@ export class Engine {
           if (
             longestVideoNode === undefined ||
             node.videoComponent.duration >
-              longestVideoNode.videoComponent.duration
+            longestVideoNode.videoComponent.duration
           ) {
             longestVideoNode = node;
           }

frontend/realtime/src/KonvaApp/Nodes/ShapeNode.ts

@@ -23,6 +23,7 @@ interface ShapeNodeConstructor {
 }
 
 export class ShapeNode extends NetworkedNode {
+  
   retry(): void {
     throw new Error("Method not implemented.");
   }
@@ -35,7 +36,7 @@ export class ShapeNode extends NetworkedNode {
     canvasSize,
     shapeType,
     size,
-    color = "#2D81FF", // Default red if no color provided
+    color = "#ff0000", // Default red if no color provided
     transform: existingTransform,
     mediaLayerRef,
     selectionManagerRef,
@@ -48,7 +49,7 @@ export class ShapeNode extends NetworkedNode {
     });
 
     let shape = null;
-
+    
     switch (shapeType) {
       case ShapeType.CIRCLE:
         const circleShape = new Konva.Circle({
@@ -71,7 +72,7 @@ export class ShapeNode extends NetworkedNode {
           width: circleShape.width(),
           height: circleShape.height(),
           ...transform,
-          fill: "transparent",
+          fill: 'transparent',
           strokeScaleEnabled: false,
           draggable: true,
         });
@@ -106,7 +107,7 @@ export class ShapeNode extends NetworkedNode {
           width: squareShape.width(),
           height: squareShape.height(),
           ...transform,
-          fill: "transparent",
+          fill: 'transparent',
           strokeScaleEnabled: false,
           draggable: true,
         });
@@ -141,7 +142,7 @@ export class ShapeNode extends NetworkedNode {
           width: triangleShape.width(),
           height: triangleShape.height(),
           ...transform,
-          fill: "transparent",
+          fill: 'transparent',
           strokeScaleEnabled: false,
           draggable: true,
         });

frontend/realtime/src/KonvaApp/Nodes/constants.ts

@@ -4,12 +4,12 @@ export enum NodeType {
   IMAGE = "image",
   TEXT = "text",
   VIDEO = "video",
-  SHAPE = "shape",
+  SHAPE = "shape"
 }
 export const transparent = Colors.transparent;
 export const primaryOrange = Colors.primaryOrange;
 
-export const highlightStrokeWidth = 2;
+export const highlightStrokeWidth = 10;
 export const minNodeSize = {
   width: 200,
   height: 200,

frontend/realtime/src/KonvaApp/NodesManagers/SelectorSquare.ts

@@ -40,9 +40,8 @@ export class SelectorSquare {
     stageRef: Konva.Stage;
   }) {
     this.kSquare = new Konva.Rect({
-      fill: "rgba(45, 129, 255, 0.3)",
-      stroke: "#00a1ff",
-      strokeWidth: 1,
+      fill: "	rgb(250, 128, 114, 0.3)",
+      stroke: "salmon",
       visible: false,
       // disable events to not interrupt with events
       listening: false,

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;

frontend/realtime/src/KonvaApp/constants.ts

@@ -2,8 +2,6 @@ export enum AppModes {
   INIT = "INIT",
   PREVIEW = "PREVIEW",
   SELECT = "SELECT",
-  PAINT = "PAINT",
-  ERASER = "ERASER",
   RENDERING = "RENDERING",
 }
 
@@ -16,5 +14,5 @@ export const VideoResolutions = {
 export const Colors = {
   transparent: "rgba(0,0,0,0)",
   translucentBlack: "rgba(0,0,0,0.3)",
-  primaryOrange: "rgba(45, 129, 255, 1)",
+  primaryOrange: "rgba(230,100,98,1)",
 };

frontend/realtime/src/KonvaRootComponent/SignaledPromptSlider.tsx

@@ -1,28 +1,22 @@
+
 import { useSignalEffect } from "@preact/signals-react";
-import { Slider } from "~/components/ui";
-import { dispatchUiEvents } from "~/signals";
+import { Slider } from "~/components/ui"
+import { dispatchUiEvents } from "~/signals"
 import { promptStrength } from "~/signals/uiEvents/promptSettings";
 
 export const SignaledPromptSlider = () => {
-  const sliderValue = promptStrength.value;
-  const onSliderValueChanged =
-    dispatchUiEvents.promptSettings.setPromptStrength;
 
-  // This makes component re-render whenever the signal value changes
-  useSignalEffect(() => {
-    promptStrength.value;
-  });
+    const sliderValue = promptStrength.value
+    const onSliderValueChanged = dispatchUiEvents.promptSettings.setPromptStrength
 
-  return (
-    <div className="fixed bottom-[138px] left-1/2 h-4 w-[400px] -translate-x-1/2 shadow-lg">
-      <Slider
-        min={0}
-        max={100}
-        value={sliderValue}
-        onChange={onSliderValueChanged}
-        step={1}
-        innerLabel={`Strength - ${sliderValue}%`}
-      />
-    </div>
-  );
-};
+    // This makes component re-render whenever the signal value changes
+    useSignalEffect(() => {
+        promptStrength.value;
+    })
+
+    return (
+        <div className="fixed bottom-80 left-1/2 -translate-x-1/2 h-4 w-96">
+            <Slider min={0} max={100} value={sliderValue} onChange={onSliderValueChanged} step={1} />
+        </div>
+    )
+}

frontend/realtime/src/KonvaRootComponent/SignaledPromptText.tsx

@@ -1,17 +1,31 @@
 import { useSignalEffect } from "@preact/signals-react";
-import { useRef } from "react";
+import { useCallback, useEffect, useRef } from "react";
 import { ResizeType, Textarea } from "~/components/ui";
 import { dispatchUiEvents } from "~/signals";
 import { promptText } from "~/signals/uiEvents/promptSettings";
 
-export const SignaledPromptText = ({}) => {
+export const SignaledPromptText = ({
+}) => {
   const textAreaRef = useRef<HTMLTextAreaElement>(null);
+  const focusOnTextArea = useCallback(() => {
+    if (textAreaRef.current === null) {
+      return;
+    }
+    textAreaRef.current.focus();
+  }, []);
+  
+  useEffect(() => {
+    focusOnTextArea();
+  }, []);
 
   const TextAreaStates = {
-    placeholder: "Type your prompt here...",
+    style: {
+      width: "530px",
+      height: "200px",
+    },
+    placeholder: "",
     label: "",
     resize: "none" as ResizeType,
-    rows: 3,
   };
 
   const textContent = promptText.value;
@@ -22,18 +36,11 @@ export const SignaledPromptText = ({}) => {
   // This makes component re-render whenever the signal value changes
   useSignalEffect(() => {
     promptText.value;
-  });
+  })
 
   return (
-    <div className="fixed bottom-5 left-1/2 w-[400px] -translate-x-1/2 text-[16px] font-medium shadow-lg">
-      <Textarea
-        className="glass rounded-xl"
-        {...TextAreaStates}
-        ref={textAreaRef}
-        onChange={onTextChanged}
-        value={textContent}
-        spellCheck={false}
-      />
+    <div className="fixed bottom-24 left-1/2 -translate-x-1/2">
+      <Textarea {...TextAreaStates} ref={textAreaRef} onChange={onTextChanged} value={textContent} />
     </div>
   );
 };

frontend/realtime/src/KonvaRootComponent/SignaledToolbarMain.tsx

@@ -50,13 +50,6 @@ export const SignaledToolbarMain = ({
         return () => dispatchUiEvents.addShapeToEngine({ shape: "triangle" });
       case ToolbarMainButtonNames.ADD_SQUARE:
         return () => dispatchUiEvents.addShapeToEngine({ shape: "square" });
-      case ToolbarMainButtonNames.SELECT:
-      case ToolbarMainButtonNames.ERASER:
-      case ToolbarMainButtonNames.PAINT:
-        return (e: React.MouseEvent<HTMLButtonElement, MouseEvent>) => {
-          toolbarMain.setMode(buttonName);
-          dispatchers[buttonName](e);
-        };
       default:
         return dispatchers[buttonName];
     }
@@ -83,7 +76,7 @@ export const SignaledToolbarMain = ({
     dispatchers.loadingBarRetry(e);
   };
   return (
-    <div className="fixed left-4 top-1/2 -translate-y-1/2">
+    <div className="fixed bottom-4 left-1/2 -translate-x-1/2">
       <Transition
         as="div"
         className={twMerge(

frontend/realtime/src/app/global.css

@@ -4,11 +4,10 @@
 
 @layer base {
   :root {
-    @apply bg-ui-background font-sans text-white;
+    @apply bg-ui-background font-sans text-[#222222];
   }
   .pegboard {
-    background-image: radial-gradient(#ffffff25 1px, transparent 0);
-    background-color: #1b1b1f;
+    background-image: radial-gradient(#3f3f466b 1px, transparent 0);
     background-size: 36px 36px;
     @apply bg-center;
   }
@@ -42,24 +41,7 @@
   }
 
   input {
-    @apply bg-ui-controls border-0 px-2 py-1 outline-none -outline-offset-2 focus:outline-primary-500 focus:ring-0;
-  }
-
-  textarea {
-    @apply bg-ui-controls placeholder:text-white/50;
-
-    /* Webkit scrollbar styles */
-    &::-webkit-scrollbar {
-      @apply w-1;
-    }
-
-    &::-webkit-scrollbar-track {
-      @apply bg-transparent;
-    }
-
-    &::-webkit-scrollbar-thumb {
-      @apply rounded-full bg-white/20 hover:bg-white/30;
-    }
+    @apply border-0 px-2 py-1 outline-none -outline-offset-2 focus:outline-primary-500 focus:ring-0;
   }
 }
 
@@ -105,11 +87,6 @@
     background-position: 100% 0;
     animation: highlight-sweep 3s infinite linear;
   }
-
-  .glass {
-    @apply relative border border-ui-border shadow-lg backdrop-blur-xl transition-colors duration-300;
-    @apply bg-ui-controls/50 hover:bg-ui-controls/90 !important;
-  }
 }
 
 @keyframes glow {

frontend/realtime/src/app/pages/main/index.tsx

@@ -8,7 +8,6 @@ import { ToolbarUserProfile } from "~/components/features";
 //Components of the Konva App are all in the KonvaComponent
 import { KonvaRootComponent } from "~/KonvaRootComponent";
 import { ToolbarTopLeft } from "~/components/features/ToolbarTopLeft";
-import { ToolbarTopRight } from "~/components/features/ToolbarTopRight";
 
 export const Main = withProtectionRoute(() => {
   // This is a hook that will log the number of times the component has rerendered
@@ -26,7 +25,6 @@ export const Main = withProtectionRoute(() => {
       />
       <div className="absolute top-0 flex w-full justify-between px-2">
         <ToolbarTopLeft />
-        <ToolbarTopRight />
         {/* <ToolbarUserProfile /> */}
       </div>
     </div>

frontend/realtime/src/components/features/DialogAddImage/DialogAddImage.tsx

@@ -52,7 +52,7 @@ export const DialogAddImage = ({
             "w-full max-w-2xl",
           )}
         >
-          <DialogTitle className="text-2xl font-bold">Upload Image</DialogTitle>
+          <DialogTitle className="text-3xl font-bold">Upload Image</DialogTitle>
           <div className="flex flex-col rounded-lg border-2 border-dashed border-ui-border">
             <FileUploader
               title=""

frontend/realtime/src/components/features/DialogEditText/DialogEditText.tsx

@@ -73,7 +73,7 @@ export const DialogEditText = ({
     <Dialog open={isOpen} onClose={closeCallback} className="relative z-50">
       <div className={dialogBackgroundStyles}>
         <DialogPanel className={twMerge(paperWrapperStyles, dialogPanelStyles)}>
-          <DialogTitle className="text-2xl font-bold">Edit Text</DialogTitle>
+          <DialogTitle className="text-3xl font-bold">Edit Text</DialogTitle>
           <TextEditor
             text={text}
             formatData={textFormatData}

frontend/realtime/src/components/features/FileUploader/DragAndDropZone.tsx

@@ -24,7 +24,7 @@ export const DragAndDropZone = ({ file, fileTypes }: Props) => {
 
   // const fileName = file && getFileName(file).toUpperCase();
   const wrapperClassName = twMerge(
-    "group cursor-pointer px-6 py-12 bg-ui-controls hover:bg-white/25 transition-colors duration-150 ease-in-out rounded-lg",
+    "group cursor-pointer px-6 py-12 bg-gray-100 hover:bg-gray-200/60 transition-colors duration-150 ease-in-out",
     !file && "flex flex-col items-center justify-center gap-",
     file && "flex items-center gap-2 py-6",
     // "rounded-lg border-2 border-dashed border-ui-border",

frontend/realtime/src/components/features/ToolbarButton.tsx

@@ -44,12 +44,12 @@ export const ToolbarButton = ({
   }
 
   const mergedButtonClasses = twMerge(
-    "rounded-lg hover:bg-white/15 transition-all duration-100 border-2 border-transparent text-white/80 px-2",
-    children ? "w-fit flex items-center gap-2.5 text-nowrap" : "size-10",
-    active && "bg-primary/30 border-2 border-primary hover:bg-primary/30",
+    "rounded-lg py-2 px-3 hover:bg-gray-200/50 transition-all duration-100",
+    children ? "w-fit flex items-center gap-2 text-nowrap" : "size-10",
+    active && "pointer-events-none text-primary ",
     disabled && "pointer-events-none text-secondary-300",
     buttonProps.prominent &&
-      "border border-primary-400/30 text-primary-500 hover:bg-primary-100/40 hover:border-primary-400/60",
+      "highlight-button border border-primary-400/30 text-primary-500 hover:bg-primary-100/40 hover:border-primary-400/60",
     customButtonClassNames,
   );
 
@@ -57,7 +57,6 @@ export const ToolbarButton = ({
     <button
       className={mergedButtonClasses}
       disabled={disabled}
-      autoFocus={false}
       {...restButtonProps}
       onClick={(e) => {
         e.preventDefault();
@@ -74,11 +73,7 @@ export const ToolbarButton = ({
     </button>
   );
   if (tooltip) {
-    return (
-      <Tooltip position="right" tip={tooltip}>
-        {Button}
-      </Tooltip>
-    );
+    return <Tooltip tip={tooltip}>{Button}</Tooltip>;
   }
   return Button;
 };

frontend/realtime/src/components/features/ToolbarMain/ToolbarMain.tsx

@@ -1,27 +1,29 @@
+import { Fragment, MouseEventHandler } from "react";
+import { twMerge } from "tailwind-merge";
+import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import {
   faArrowRotateLeft,
   faArrowRotateRight,
   faCircle,
-  faEraser,
   faFilePlus,
+  faFilm,
+  faFloppyDisk,
   faImage,
   faLocationArrow,
-  faPaintbrush,
+  faMagicWandSparkles,
   faShapes,
   faSquare,
   faText,
   faTriangle,
 } from "@fortawesome/pro-solid-svg-icons";
-import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
-import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
-import { Fragment, MouseEventHandler } from "react";
-import { twMerge } from "tailwind-merge";
 
 import { ToolbarButton } from "../ToolbarButton";
 
 // style and constants
 import { paperWrapperStyles, toolTipStyles } from "~/components/styles";
 import { ToolbarMainButtonNames } from "./enum";
+import { ButtonPreviewAndRender } from "./ButtonPreviewAndRender";
 
 export const ToolbarMain = ({
   disabled = false,
@@ -40,14 +42,13 @@ export const ToolbarMain = ({
     <>
       <div
         className={twMerge(
-          "flex h-fit flex-col items-center gap-y-2",
+          "flex w-fit items-center divide-x divide-ui-border",
           paperWrapperStyles,
-          "glass",
           disabled &&
-            "pointer-events-none cursor-default bg-ui-border shadow-md",
+          "pointer-events-none cursor-default bg-ui-border shadow-md",
         )}
       >
-        <div className="flex flex-col items-center gap-2">
+        <div className="flex items-center gap-2 pr-2">
           <ToolbarButton
             icon={faLocationArrow}
             iconProps={{ className: "fa-flip-horizontal" }}
@@ -61,111 +62,62 @@ export const ToolbarMain = ({
           />
 
           <Popover className="relative">
-            {({ close }) => (
-              <>
-                <PopoverButton as={Fragment}>
-                  <button
-                    data-tooltip="Add Shapes"
-                    className={twMerge(
-                      "size-10 rounded-lg border-2 border-transparent text-white/80 transition-all duration-100 hover:bg-white/15",
-                      toolTipStyles.base,
-                      toolTipStyles.right,
-                    )}
-                  >
-                    <FontAwesomeIcon icon={faShapes} />
-                  </button>
-                </PopoverButton>
-                <PopoverPanel
-                  anchor="right"
-                  className={twMerge(
-                    "flex flex-row [--anchor-gap:12px]",
-                    paperWrapperStyles,
-                  )}
-                >
-                  <ToolbarButton
-                    icon={faSquare}
-                    buttonProps={{
-                      ...buttonProps.ADD_SQUARE,
-                      onClick: (e) => {
-                        buttonProps.ADD_SQUARE.onClick?.(e);
-                        close();
-                      },
-                    }}
-                  />
-                  <ToolbarButton
-                    icon={faCircle}
-                    buttonProps={{
-                      ...buttonProps.ADD_CIRCLE,
-                      onClick: (e) => {
-                        buttonProps.ADD_CIRCLE.onClick?.(e);
-                        close();
-                      },
-                    }}
-                  />
-                  <ToolbarButton
-                    icon={faTriangle}
-                    buttonProps={{
-                      ...buttonProps.ADD_TRIANGLE,
-                      onClick: (e) => {
-                        buttonProps.ADD_TRIANGLE.onClick?.(e);
-                        close();
-                      },
-                    }}
-                  />
-                </PopoverPanel>
-              </>
-            )}
+            <PopoverButton as={Fragment}>
+              <button
+                data-tooltip="Add Shapes"
+                className={twMerge(
+                  "size-10 rounded-lg p-2 hover:bg-gray-200/50",
+                  toolTipStyles,
+                )}
+              >
+                <FontAwesomeIcon icon={faShapes} />
+              </button>
+            </PopoverButton>
+            <PopoverPanel
+              anchor="bottom"
+              className={twMerge(
+                // "absolute bottom-full left-1/2 z-10 mb-2 -translate-x-1/2",
+                "flex flex-col [--anchor-gap:16px]",
+                paperWrapperStyles,
+              )}
+            >
+              <ToolbarButton icon={faSquare} buttonProps={buttonProps.ADD_SQUARE} />
+              <ToolbarButton icon={faCircle} buttonProps={buttonProps.ADD_CIRCLE} />
+              <ToolbarButton icon={faTriangle} buttonProps={buttonProps.ADD_TRIANGLE} />
+            </PopoverPanel>
           </Popover>
 
           <Popover className="relative">
             <PopoverButton as={Fragment}>
               <button
-                data-tooltip="Add Media"
+                data-tooltip="Add..."
                 className={twMerge(
-                  "size-10 rounded-lg border-2 border-transparent text-white/80 transition-all duration-100 hover:bg-white/15",
-                  toolTipStyles.base,
-                  toolTipStyles.right,
+                  "size-10 rounded-lg p-2 hover:bg-gray-200/50",
+                  toolTipStyles,
                 )}
               >
                 <FontAwesomeIcon icon={faFilePlus} />
               </button>
             </PopoverButton>
             <PopoverPanel
-              anchor="right"
+              anchor="bottom"
               className={twMerge(
                 // "absolute bottom-full left-1/2 z-10 mb-2 -translate-x-1/2",
-                "flex [--anchor-gap:12px]",
+                "flex flex-col [--anchor-gap:16px]",
                 paperWrapperStyles,
               )}
             >
               <ToolbarButton icon={faImage} buttonProps={buttonProps.ADD_IMAGE}>
-                <span className="text-[16px]">Add Image</span>
+                Add Image
               </ToolbarButton>
               {/* <ToolbarButton icon={faFilm} buttonProps={buttonProps.ADD_VIDEO}>
                 Add Video
               </ToolbarButton> */}
             </PopoverPanel>
           </Popover>
+
         </div>
-
-        <hr className="w-full border-t border-white/15" />
-
-        <div className="flex flex-col items-center gap-2">
-          <ToolbarButton
-            icon={faPaintbrush}
-            buttonProps={buttonProps.PAINT}
-            tooltip="Paint Brush"
-          />
-          <ToolbarButton
-            icon={faEraser}
-            buttonProps={buttonProps.ERASER}
-            tooltip="Eraser"
-          />
-        </div>
-
-        <hr className="w-full border-t border-white/15" />
-
-        <div className="flex flex-col items-center gap-2">
+        <div className="flex items-center gap-2 pl-2">
           <ToolbarButton
             icon={faArrowRotateLeft}
             buttonProps={buttonProps.UNDO}
@@ -176,11 +128,15 @@ export const ToolbarMain = ({
             buttonProps={buttonProps.REDO}
             tooltip="Redo"
           />
-          {/* <ToolbarButton
+          <ToolbarButton
             icon={faFloppyDisk}
             buttonProps={buttonProps.SAVE}
             tooltip="Save"
-          /> */}
+          />
+          <ButtonPreviewAndRender
+            buttonPreviewProps={buttonProps.PREVIEW}
+            buttonRenderProps={buttonProps.DOWNLOAD}
+          />
         </div>
       </div>
     </>

frontend/realtime/src/components/features/ToolbarMain/enum.ts

@@ -7,8 +7,6 @@ export enum ToolbarMainButtonNames {
   ADD_CIRCLE = "ADD_CIRCLE",
   ADD_TRIANGLE = "ADD_TRIANGLE",
   ADD_SQUARE = "ADD_SQUARE",
-  PAINT = "PAINT",
-  ERASER = "ERASER",
   AI_STYLIZE = "AI_STYLIZE",
   UNDO = "UNDO",
   REDO = "REDO",

frontend/realtime/src/components/features/ToolbarTopLeft/index.tsx

@@ -1,27 +1,28 @@
 import { twMerge } from "tailwind-merge";
 
-// import { paperWrapperStyles } from "~/components/styles";
-// import { faPlus, faQuestion } from "@fortawesome/pro-solid-svg-icons";
-// import { ToolbarButton } from "~/components/features/ToolbarButton";
+import { paperWrapperStyles } from "~/components/styles";
+import { faPlus, faQuestion } from "@fortawesome/pro-solid-svg-icons";
+import { ToolbarButton } from "~/components/features/ToolbarButton";
 
 export const ToolbarTopLeft = () => {
   return (
     <div
       className={twMerge(
+        paperWrapperStyles,
         "z-20 mt-2 flex h-fit w-fit items-center gap-1 py-2.5 pl-4",
       )}
     >
       <img
-        src="/brand/mira-logo.png"
+        src="/brand/Storyteller-Logo-Black.png"
         alt="logo"
-        className="h-[30px] select-none pb-1"
+        className="mr-4 h-8"
       />
-      {/* <ToolbarButton icon={faPlus}>
+      <ToolbarButton icon={faPlus}>
         <span className="font-semibold">New Board</span>
       </ToolbarButton>
       <ToolbarButton icon={faQuestion}>
         <span className="font-semibold">Help</span>
-      </ToolbarButton> */}
+      </ToolbarButton>
     </div>
   );
 };

frontend/realtime/src/components/features/ToolbarTopRight/index.tsx

@@ -1,31 +0,0 @@
-import { faSave } from "@fortawesome/pro-solid-svg-icons";
-import { twMerge } from "tailwind-merge";
-import { Button } from "~/components/ui/Button";
-import { toolbarMain } from "~/signals/uiAccess/toolbarMain";
-import { dispatchers } from "~/signals/uiEvents/toolbarMain";
-import { ToolbarMainButtonNames } from "../ToolbarMain/enum";
-
-// import { paperWrapperStyles } from "~/components/styles";
-// import { faPlus, faQuestion } from "@fortawesome/pro-solid-svg-icons";
-// import { ToolbarButton } from "~/components/features/ToolbarButton";
-
-export const ToolbarTopRight = () => {
-  const saveButtonState =
-    toolbarMain.signal.value.buttonStates[ToolbarMainButtonNames.SAVE];
-
-  return (
-    <div
-      className={twMerge(
-        "z-20 mt-2 flex h-fit w-fit items-center gap-1 p-4 py-2.5",
-      )}
-    >
-      <Button
-        icon={faSave}
-        onClick={dispatchers[ToolbarMainButtonNames.SAVE]}
-        disabled={saveButtonState?.disabled}
-      >
-        Save Image
-      </Button>
-    </div>
-  );
-};

frontend/realtime/src/components/styles.ts

@@ -1,35 +1,26 @@
 import { twMerge } from "tailwind-merge";
 
 export const paperWrapperStyles =
-  "rounded-xl border border-ui-border p-1.5 shadow-lg bg-ui-panel";
+  "rounded-xl border border-ui-border bg-ui-panel p-2 shadow-lg";
 
 export const transitionTimingStyles = "transition duration-150 ease-in-out";
 
 export const dialogBackgroundStyles =
   "fixed inset-0 flex w-screen items-center justify-center bg-black/50";
 
-export const toolTipStyles = {
-  base: twMerge(
-    "relative",
-    "after:absolute after:z-40",
-    "after:content-[attr(data-tooltip)] after:text-black after:text-nowrap after:text-sm after:font-medium after:text-white/80 after:glass",
-    "after:rounded-lg after:border after:border-ui-border after:bg-ui-panel after:px-2.5 after:py-1 after:shadow-xl",
-    "after:hidden hover:after:block",
-  ),
-  top: twMerge(
-    "after:left-1/2 after:bottom-full after:-translate-x-1/2 after:mb-3.5",
-  ),
-  bottom: twMerge(
-    "after:left-1/2 after:top-full after:-translate-x-1/2 after:mt-3.5",
-  ),
-  left: twMerge(
-    "after:top-1/2 after:right-full after:-translate-y-1/2 after:mr-3.5",
-  ),
-  right: twMerge(
-    "after:top-1/2 after:left-full after:-translate-y-1/2 after:ml-3.5",
-  ),
-};
+export const toolTipStyles = twMerge(
+  "relative",
+  "before:w-0 before:h-0 before:absolute",
+  "before:left-1/2 before:bottom-full before:mb-[1px] before:-translate-x-1/2 before:z-50",
+  "before:border-l-8 before:border-l-transparent",
+  "before:border-r-8 before:border-r-transparent",
+  "before:border-t-8 before:border-t-white",
+  "after:absolute after:left-1/2 after:bottom-full after:-translate-x-1/2 after:z-40",
+  "after:content-[attr(data-tooltip)] after:text-black after:text-nowrap",
+  "after:rounded-lg after:border after:border-ui-border after:bg-ui-panel after:px-2.5 after:py-1.5 after:mb-2 after:shadow-xl",
+  "after:hidden before:hidden hover:after:block hover:before:block",
+);
 
 export const dialogPanelStyles = twMerge(
-  "flex max-w-5xl flex-col justify-between gap-4 p-6 bg-[#1C1C20]",
+  "flex max-w-5xl flex-col justify-between gap-4 p-6",
 );

frontend/realtime/src/components/ui/Button.tsx

@@ -31,14 +31,14 @@ export const Button = ({
         return "bg-tertiary hover:bg-tertiary-400 text-white focus-visible:outline-tertiary-600";
       }
       case "secondary": {
-        return "bg-ui-controls hover:bg-ui-controls/80 text-white focus-visible:outline-secondary";
+        return "bg-gray-600 hover:bg-gray-500 text-white focus-visible:outline-secondary";
       }
       case "action": {
         return " bg-action hover:bg-action-500 text-white focus-visible:outline-action";
       }
       case "primary":
       default: {
-        return "bg-primary hover:bg-primary-600 text-white focus-visible:outline-primary-600";
+        return "bg-primary hover:bg-primary-400 text-white focus-visible:outline-primary-600";
       }
     }
   }

frontend/realtime/src/components/ui/ComboBox.tsx

@@ -29,25 +29,25 @@ export const Combobox = ({ options, value, onChange }: ComboBoxInterface) => {
       onChange={onChange}
       onClose={() => setQuery("")}
     >
-      <div className="relative flex h-10 rounded-md border border-ui-border">
+      <div className="relative flex h-10 rounded-md border">
         <ComboboxInput
-          className="rounded-lg px-3"
+          className="rounded-md"
           displayValue={(font: string) => font}
           onChange={(event) => setQuery(event.target.value)}
         />
-        <ComboboxButton className="absolute right-0 top-0 flex h-10 items-center justify-center px-3 transition-colors duration-300 hover:text-primary">
-          <FontAwesomeIcon icon={faChevronDown} className="size-4" />
+        <ComboboxButton className="absolute right-0 top-0 flex h-10 items-center justify-center hover:text-primary">
+          <FontAwesomeIcon icon={faChevronDown} className="size-4 p-2" />
         </ComboboxButton>
       </div>
       <ComboboxOptions
         anchor="bottom"
-        className="mt-1 w-[var(--input-width)] rounded-lg border border-ui-border empty:invisible"
+        className="mt-1 w-[var(--input-width)] rounded-lg border empty:invisible"
       >
         {filteredOptions.map((option, idx) => (
           <ComboboxOption
             key={idx}
             value={option}
-            className="bg-ui-controls cursor-pointer px-3 py-2 data-[focus]:bg-gray-600 data-[selected]:font-medium"
+            className="cursor-pointer bg-white p-2 data-[focus]:bg-primary-200 data-[selected]:font-medium"
           >
             {option}
           </ComboboxOption>

frontend/realtime/src/components/ui/Slider.tsx

@@ -12,7 +12,6 @@ interface SliderProps {
   showTooltip?: boolean;
   tooltipContent?: ReactNode;
   suffix?: string;
-  innerLabel?: string;
 }
 
 export const Slider = ({
@@ -25,7 +24,6 @@ export const Slider = ({
   showTooltip,
   tooltipContent,
   suffix,
-  innerLabel,
 }: SliderProps) => {
   const [localValue, setLocalValue] = useState(value);
   const sliderRef = useRef<HTMLDivElement>(null);
@@ -90,40 +88,25 @@ export const Slider = ({
     <div
       ref={sliderRef}
       className={twMerge(
-        "glass group relative h-7 w-full cursor-pointer overflow-hidden rounded-lg border border-ui-border",
-        isDragging && "!bg-ui-controls/90",
+        "relative h-2.5 w-full cursor-pointer rounded-full bg-gray-500",
         className,
       )}
       onMouseDown={handleMouseDown}
     >
       <div
-        className={twMerge(
-          "absolute h-7 bg-primary/30 transition-colors duration-300 group-hover:bg-primary/50",
-          isDragging && "!bg-primary/50",
-        )}
+        className="absolute h-2.5 rounded-full bg-primary"
         style={{ width: `${percentage}%` }}
-      >
-        {innerLabel && (
-          <span
-            className={twMerge(
-              "absolute top-1/2 ml-2.5 -translate-y-1/2 text-nowrap text-sm font-medium text-white/60 transition-colors duration-300 group-hover:text-white",
-              isDragging && "!text-white",
-            )}
-          >
-            {innerLabel}
-          </span>
+      ></div>
+      <div
+        className={twMerge(
+          "absolute -top-1 h-[18px] w-[18px] -translate-x-1/2 transform rounded-full bg-gray-200 transition-colors",
+          isDragging ? "bg-white" : "bg-gray-300 hover:bg-white",
         )}
-        <div
-          className={twMerge(
-            "absolute right-0 top-1/2 mr-1.5 h-3.5 w-0.5 -translate-y-1/2",
-            isDragging ? "bg-white" : "bg-white/50",
-          )}
-          onMouseDown={handleMouseDown}
-          onMouseEnter={handleMouseEnter}
-          onMouseLeave={handleMouseLeave}
-        />
-      </div>
-
+        style={{ left: `${percentage}%` }}
+        onMouseDown={handleMouseDown}
+        onMouseEnter={handleMouseEnter}
+        onMouseLeave={handleMouseLeave}
+      />
       {showTooltip && (
         <Transition
           as={"div"}

frontend/realtime/src/components/ui/TextArea.tsx

@@ -15,7 +15,7 @@ const resizeStyles = {
   horizontal: "resize-x",
   vertical: "resize-y",
   undefined: "",
-};
+}
 
 export interface TextareaInterface
   extends TextareaHTMLAttributes<HTMLTextAreaElement> {
@@ -27,6 +27,7 @@ export const Textarea = forwardRef<HTMLTextAreaElement, TextareaInterface>(
     { className, label, resize = "vertical", id, ...rest }: TextareaInterface,
     ref,
   ) => {
+
     const resizeStyle = resizeStyles[resize];
 
     return (
@@ -37,16 +38,17 @@ export const Textarea = forwardRef<HTMLTextAreaElement, TextareaInterface>(
           ref={ref}
           id={id ? id : label ? kebabCase(label) : undefined}
           className={twMerge(
-            "rounded-lg border border-ui-border px-4 py-2.5 text-white outline-none focus:outline-none",
+            "border-ui-panel-border bg-ui-controls rounded-lg border px-3 py-2",
             className,
             resizeStyle,
           )}
           style={{
             outline: "2px solid transparent",
+            transition: "outline-color 0.15s ease-in-out",
           }}
           onFocus={(e) => {
             // disableHotkeyInput(DomLevels.INPUT);
-            e.currentTarget.style.outlineColor = "#2d81ff";
+            e.currentTarget.style.outlineColor = "#e66462";
           }}
           onBlur={(e) => {
             // enableHotkeyInput(DomLevels.INPUT);

frontend/realtime/src/components/ui/TextEditor/ButtonsAlignment.tsx

@@ -1,6 +1,6 @@
 import {
   faAlignCenter,
-  // faAlignJustify,
+  faAlignJustify,
   faAlignLeft,
   faAlignRight,
 } from "@fortawesome/pro-solid-svg-icons";

frontend/realtime/src/components/ui/TextEditor/ColorPicker.tsx

@@ -1,10 +1,6 @@
 import { ChangeEvent, useState } from "react";
 import { twMerge } from "tailwind-merge";
-import {
-  faFont,
-  faHashtag,
-  IconDefinition,
-} from "@fortawesome/pro-solid-svg-icons";
+import { faFont, faHashtag, IconDefinition } from "@fortawesome/pro-solid-svg-icons";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
 import { HexColorPicker } from "react-colorful";
@@ -17,7 +13,7 @@ export const ColorPicker = ({
   color: prevColor,
   onChange,
   faIcon = faFont,
-  borderStyle,
+  borderStyle = "border",
   showBar = true,
 }: {
   color: string;
@@ -57,16 +53,12 @@ export const ColorPicker = ({
 
   return (
     <Popover className="relative">
-      <PopoverButton
-        className={twMerge(
-          "bg-ui-controls flex size-10 flex-col items-center gap-1 rounded-md p-2",
-          borderStyle,
-        )}
-      >
+      <PopoverButton className={twMerge(
+        "flex size-10 flex-col items-center gap-1 rounded-md p-2",
+        borderStyle
+      )}>
         <FontAwesomeIcon icon={faIcon} color={prevColor} />
-        {showBar && (
-          <span className="h-1 w-full" style={{ backgroundColor: prevColor }} />
-        )}
+        {showBar && <span className="h-1 w-full" style={{ backgroundColor: prevColor }} />}
       </PopoverButton>
       <PopoverPanel
         anchor="bottom"

frontend/realtime/src/components/ui/TextEditor/TextEditor.tsx

@@ -151,11 +151,7 @@ export const TextEditor = ({
           onChangeTextDecoration={onChangeTextDecoration}
         />
       </div>
-      <Textarea
-        className="bg-white"
-        {...unionedTextAreaProps}
-        ref={textAreaRef}
-      />
+      <Textarea {...unionedTextAreaProps} ref={textAreaRef} />
     </div>
   );
 };

frontend/realtime/src/components/ui/TextEditor/styles.ts

@@ -1,4 +1,4 @@
 export const buttonGroupButton =
-  "size-10 rounded-none bg-ui-controls text-white/80 hover:bg-white/30 data-[selected=true]:text-white data-[selected=true]:bg-primary/30 disabled:hover:bg-transparent disabled:text-gray-400 border border-ui-border";
+  "size-10 rounded-none bg-white text-black hover:bg-gray-200 data-[selected=true]:text-primary-400 disabled:hover:bg-transparent disabled:text-gray-400";
 export const buttonGroupGrouper =
-  "flex divide-x overflow-hidden rounded-md border border-ui-border";
+  "flex divide-x overflow-hidden rounded-md border";

frontend/realtime/src/components/ui/Tooltip.tsx

@@ -7,18 +7,15 @@ export const Tooltip = ({
   tip,
   children,
   forceShow,
-  position = "top",
 }: {
   tip: string;
   children: ReactElement<any, string | JSXElementConstructor<any>>;
   forceShow?: boolean;
-  position?: "top" | "bottom" | "left" | "right";
 }) => {
   const clonedChildren = cloneElement(children, {
     "data-tooltip": tip,
     className: twMerge(
-      toolTipStyles.base,
-      toolTipStyles[position],
+      toolTipStyles,
       forceShow && "after:block before:block",
       children.props.className,
     ),

frontend/realtime/src/signals/uiAccess/toolbarMain/toolbarMain.ts

@@ -1,6 +1,5 @@
 import { signal } from "@preact/signals-react";
 import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum";
-import { isToolbarMode } from "~/signals/uiEvents/toolbarMain/toolbarMain";
 
 interface ToolbarMainSignalInterface {
   disabled: boolean;
@@ -43,17 +42,6 @@ export const toolbarMain = {
       },
     };
   },
-
-  setMode(mode: ToolbarMainButtonNames) {
-    if (!isToolbarMode(mode)) {
-      return;
-    }
-
-    this.changeButtonState(ToolbarMainButtonNames.SELECT, { active: false });
-    this.changeButtonState(ToolbarMainButtonNames.ERASER, { active: false });
-    this.changeButtonState(ToolbarMainButtonNames.PAINT, { active: false });
-    this.changeButtonState(mode, { active: true });
-  }
 };
 
 function initButtonStates() {

frontend/realtime/src/signals/uiEvents/toolbarMain/toolbarMain.ts

@@ -2,7 +2,6 @@ import { MouseEventHandler } from "react";
 import { signal, effect, Signal } from "@preact/signals-react";
 
 import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum";
-import { AppModes } from "~/KonvaApp/constants";
 
 const events = Object.values(ToolbarMainButtonNames).reduce(
   (acc, buttonName) => {
@@ -26,35 +25,22 @@ const effectsCleanups = Object.values(ToolbarMainButtonNames).reduce(
     [key in ToolbarMainButtonNames]: (() => void) | undefined;
   },
 );
-
-export const isToolbarMode = (buttonName: ToolbarMainButtonNames) => {
-  return appModeMap[buttonName] !== undefined;
-}
-
-export const appModeMap: Record<string, AppModes> = {
-  [ToolbarMainButtonNames.SELECT]: AppModes.SELECT,
-  [ToolbarMainButtonNames.ERASER]: AppModes.ERASER,
-  [ToolbarMainButtonNames.PAINT]: AppModes.PAINT,
-}
-
-const defaultHandler = (buttonName: ToolbarMainButtonNames, callback: MouseEventHandler<HTMLButtonElement>) => {
-  if (effectsCleanups[buttonName]) {
-    effectsCleanups[buttonName]();
-  }
-  effectsCleanups[buttonName] = effect(() => {
-    if (events[buttonName].value) {
-      callback(events[buttonName].value);
-      events[buttonName].value = undefined;
-    }
-  });
-}
-
 export const buttonEventsHandlers = Object.values(
   ToolbarMainButtonNames,
 ).reduce(
   (acc, buttonName) => {
     acc[buttonName] = {
-      onClick: (callback: MouseEventHandler<HTMLButtonElement>) => { defaultHandler(buttonName, callback) },
+      onClick: (callback: MouseEventHandler<HTMLButtonElement>) => {
+        if (effectsCleanups[buttonName]) {
+          effectsCleanups[buttonName]();
+        }
+        effectsCleanups[buttonName] = effect(() => {
+          if (events[buttonName].value) {
+            callback(events[buttonName].value);
+            events[buttonName].value = undefined;
+          }
+        });
+      },
     };
     return acc;
   },

frontend/realtime/tailwind.config.js

@@ -44,15 +44,14 @@ export default {
       info: colors.sky[500],
 
       // brand colors
-      primary: storytellerColors.azureRadiance,
+      primary: storytellerColors.sunglo,
       secondary: storytellerColors.gunpowder,
       tertiary: storytellerColors.aquamarineBlue,
 
       ui: {
-        background: "#101014",
-        panel: "#3E3E41",
-        border: "#FFFFFF0D",
-        controls: "#3E3E41",
+        background: colors.zinc[100],
+        panel: colors.white,
+        border: colors.zinc[200],
       },
     },
     backgroundImage: {

frontend/realtime/tailwind.stcolors.js

@@ -1,19 +1,19 @@
 export const storytellerColors = {
-  azureRadiance: {
-    DEFAULT: "#2d81ff", //=500
-    50: "#eef6ff",
-    100: "#d9ebff",
-    200: "#bcdcff",
-    300: "#8ec7ff",
-    400: "#59a7ff",
-    500: "#2d81ff",
-    600: "#1b63f5",
-    700: "#144ee1",
-    800: "#173fb6",
-    900: "#19398f",
-    950: "#142457",
+  sunglo: {
+    //sunglo
+    DEFAULT: "#e66462", //=500
+    50: "#fdf3f3",
+    100: "#fce4e4",
+    200: "#facfce",
+    300: "#f6acab",
+    400: "#ef7c7a",
+    500: "#e66462",
+    600: "#cf3533",
+    700: "#ae2927",
+    800: "#902624",
+    900: "#782524",
+    950: "#410f0e",
   },
-
   gunpowder: {
     DEFAULT: "#39394c", //=950
     50: "#f5f6f9",