render bug fix

- Both modes' menus have been separated for better maintenance
- Slider is generic, can be reused
- New signals added for engine to hook into
- Sliders tested working on the UI side

crates/desktop/src-tauri/Cargo.toml

@@ -47,10 +47,4 @@ tokenizers = "0.15.2"
 
 [features]
 default = ["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",
-]
+cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]

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

@@ -1,150 +0,0 @@
-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

@@ -1,474 +0,0 @@
-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,6 +1,4 @@
-pub mod cumprod;
 pub mod image;
-pub mod lcm_scheduler;
 pub mod model_cache;
 pub mod model_file;
 pub mod models;

frontend/realtime/.gitignore

@@ -12,7 +12,7 @@ dist
 dist-ssr
 *.local
 backup
-
+.venv
 # Editor directories and files
 ~
 .vscode

frontend/realtime/paint_py_server/server.py

@@ -0,0 +1,527 @@
+import asyncio
+import base64
+import json
+import logging
+import pathlib
+import time
+from io import BytesIO
+from typing import Dict, Any, Optional, Tuple, Callable, List
+
+import torch
+import websockets
+from PIL import Image
+from diffusers import (
+    StableDiffusionXLPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+    AutoPipelineForImage2Image,
+    LCMScheduler,
+    UNet2DConditionModel,
+)
+from tqdm import tqdm
+from huggingface_hub import HfApi, hf_hub_download
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Global variables
+connected = False
+pipe = None
+model_initialized = False
+current_settings = {}
+
+# Custom TQDMProgressCallback for tracking download progress
+class ProgressCallback:
+    def __init__(self, websocket):
+        self.websocket = websocket
+        self.total = 0
+        self.current = 0
+        self.last_update_time = 0
+        self.last_progress = 0
+        self.stage = "downloading"
+        self.current_file = ""
+    
+    async def __call__(self, current, total, file_name=""):
+        self.total = total
+        self.current = current
+        self.current_file = file_name or self.current_file
+        
+        # Only send updates at most once every 500ms to avoid flooding the connection
+        current_time = time.time()
+        progress = int(100 * current / total) if total > 0 else 0
+        
+        if (current_time - self.last_update_time > 0.5 and progress != self.last_progress) or current >= total:
+            self.last_update_time = current_time
+            self.last_progress = progress
+            
+            await self.websocket.send(json.dumps({
+                'type': 'progress_update',
+                'stage': self.stage,
+                'progress': progress,
+                'file': self.current_file
+            }))
+
+# Pure functions
+def decode_base64_to_image(base64_string: str) -> Image.Image:
+    """Decode a base64 string to a PIL Image."""
+    try:
+        image_data = base64.b64decode(base64_string)
+        return Image.open(BytesIO(image_data))
+    except Exception as e:
+        logger.error(f"Error decoding base64 image: {e}")
+        raise ValueError(f"Invalid base64 image: {e}")
+
+def encode_image_to_base64(image: Image.Image) -> str:
+    """Encode a PIL Image to a base64 string."""
+    try:
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+    except Exception as e:
+        logger.error(f"Error encoding image to base64: {e}")
+        raise ValueError(f"Error encoding image: {e}")
+
+async def load_model(sdxl_checkpoint_path: pathlib.Path, progress_callback: Optional[ProgressCallback] = None) -> Tuple[AutoPipelineForImage2Image, bool]:
+    """Load the SDXL model with LCM scheduler and report progress."""
+    try:
+        logger.info(f"Loading model from {sdxl_checkpoint_path}")
+        
+        if progress_callback:
+            progress_callback.stage = "loading_unet"
+            progress_callback.current_file = "latent-consistency/lcm-sdxl"
+            await progress_callback(0, 100)
+        
+        # Load the UNet with LCM
+        unet = UNet2DConditionModel.from_pretrained(
+            "latent-consistency/lcm-sdxl",
+            torch_dtype=torch.float16,
+            variant="fp16",
+            use_safetensors=True,
+        )
+        
+        if progress_callback:
+            progress_callback.stage = "loading_base_model"
+            progress_callback.current_file = str(sdxl_checkpoint_path)
+            await progress_callback(25, 100)
+        
+        # Load the base model using the checkpoint path
+        pipe = AutoPipelineForImage2Image.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            unet=unet,
+            torch_dtype=torch.float16,
+            variant="fp16",
+           use_safetensors=True
+        ).to("cuda")
+        
+        if progress_callback:
+            progress_callback.stage = "configuring_scheduler"
+            await progress_callback(75, 100)
+        
+        # Replace the scheduler with LCM scheduler
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        
+        if progress_callback:
+            progress_callback.stage = "loading_lcm_adapter"
+            progress_callback.current_file = "latent-consistency/lcm-lora-sdxl"
+            await progress_callback(85, 100)
+        
+        # Load the default LCM adapter
+        pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm",from_safetensors=True)
+        
+        if progress_callback:
+            await progress_callback(100, 100)
+        
+        return pipe, True
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        if progress_callback:
+            progress_callback.stage = "error"
+            progress_callback.current_file = str(e)
+            await progress_callback(0, 100)
+        return None, False
+
+async def load_lora(pipe: AutoPipelineForImage2Image, lora_path: pathlib.Path, adapter_name: str = "lora", progress_callback: Optional[ProgressCallback] = None) -> bool:
+    """Load a LoRA adapter into the pipeline and report progress."""
+    try:
+        logger.info(f"Loading LoRA from {lora_path} as {adapter_name}")
+        
+        if progress_callback:
+            progress_callback.stage = "loading_lora"
+            progress_callback.current_file = str(lora_path)
+            await progress_callback(0, 100)
+        
+        # Load the LoRA weights
+        pipe.load_lora_weights(lora_path, adapter_name=adapter_name)
+        
+        if progress_callback:
+            await progress_callback(100, 100)
+            
+        return True
+    except Exception as e:
+        logger.error(f"Error loading LoRA: {e}")
+        if progress_callback:
+            progress_callback.stage = "error"
+            progress_callback.current_file = str(e)
+            await progress_callback(0, 100)
+        return False
+
+def generate_image(
+    pipe: AutoPipelineForImage2Image,
+    image: Image.Image,
+    prompt: str,
+    lora_strength: float = 1.0,
+    image_to_image_strength: float = 1.0,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 4,
+    guidance_scale: float = 1.0,
+    seed: Optional[int] = None
+) -> Optional[Image.Image]:
+    """Generate an image using the LCM model."""
+    try:
+        # Resize input image to match desired dimensions
+        image = image.resize((width, height))
+        
+        # Set up generator for reproducibility if seed is provided
+        generator = torch.manual_seed(seed) if seed is not None else None
+        
+        # Set lora adapter weights if available
+        if hasattr(pipe, 'set_adapters') and len(pipe.adapters) > 0:
+            # If we have both lcm and lora adapters
+            if "lcm" in pipe.adapters and "lora" in pipe.adapters:
+                pipe.set_adapters(["lcm", "lora"], adapter_weights=[1.0, lora_strength])
+            # If we only have lcm adapter
+            elif "lcm" in pipe.adapters:
+                pipe.set_adapters(["lcm"], adapter_weights=[1.0])
+        
+        # Generate the image
+        result = pipe(
+            prompt=prompt,
+            image=image,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            strength=image_to_image_strength,
+            width=width,
+            height=height,
+            generator=generator
+        )
+        
+        return result.images[0]
+    except Exception as e:
+        logger.error(f"Error generating image: {e}")
+        return None
+
+async def handle_client(websocket):
+    """Handle a client connection."""
+    global connected, pipe, model_initialized, current_settings
+    
+    if connected:
+        await websocket.send(json.dumps({
+            'type': 'connection_error',
+            'error': "Another client is already connected"
+        }))
+        await websocket.close(1008, "Another client is already connected")
+        return
+    
+    connected = True
+    logger.info("Client connected")
+    
+    # Create a progress callback for this connection
+    progress_callback = ProgressCallback(websocket)
+    
+    try:
+        async for message in websocket:
+            try:
+                data = json.loads(message)
+                logger.info(f"Received message type: {data.get('type', 'unknown')}")
+                logger.info(f"Full payload: {json.dumps(data, indent=2)}")
+                
+                # Check message type
+                message_type = data.get('type', '')
+                
+                # Handle setup request - initializes the model and settings
+                if message_type == 'setup':
+                    logger.info("Processing setup message")
+                    model_data = data.get('model', {})
+                    lora_data = data.get('lora')
+                    device = data.get('device', 'cuda')
+                    
+                    logger.info(f"Model data: {json.dumps(model_data, indent=2)}")
+                    logger.info(f"LoRA data: {json.dumps(lora_data, indent=2) if lora_data else 'None'}")
+                    logger.info(f"Device: {device}")
+                    
+                    # Extract setup parameters with detailed logging
+                    model_path = model_data.get('path', '')
+                    logger.info(f"Raw model path from message: '{model_path}'")
+                    
+                    # Check if the path is valid
+                    if not model_path:
+                        logger.error("Model path is empty")
+                        await websocket.send(json.dumps({
+                            'type': 'setup_response',
+                            'success': False,
+                            'error': 'Model path is required'
+                        }))
+                        continue
+                    
+                    # Convert to Path object
+                    sdxl_checkpoint_path = pathlib.Path(model_path)
+                    logger.info(f"Converted to pathlib.Path: '{sdxl_checkpoint_path}'")
+                    logger.info(f"Path is absolute: {sdxl_checkpoint_path.is_absolute()}")
+                    logger.info(f"Path exists check: {sdxl_checkpoint_path.exists()}")
+                    
+                    # If the path exists as a file
+                    if sdxl_checkpoint_path.is_file():
+                        logger.info(f"Path is a file: Yes")
+                        logger.info(f"File size: {sdxl_checkpoint_path.stat().st_size / (1024*1024):.2f} MB")
+                    else:
+                        logger.info(f"Path is a file: No")
+                    
+                    # Store current settings
+                    current_settings = {
+                        'sdxl_checkpoint_path': str(sdxl_checkpoint_path),
+                        'lora_path': str(lora_data['path']) if lora_data else None,
+                        'lora_alpha': float(lora_data['alpha']) if lora_data else 1.0,
+                        'device': device
+                    }
+                    logger.info(f"Current settings: {json.dumps(current_settings, indent=2)}")
+                    
+                    # Validate required parameters
+                    if not sdxl_checkpoint_path.exists():
+                        logger.error(f"Path does not exist: {sdxl_checkpoint_path}")
+                        # Try alternative path format
+                        alt_path = pathlib.Path(model_path.replace('\\', '/'))
+                        logger.info(f"Trying alternative path format: {alt_path}")
+                        logger.info(f"Alternative path exists check: {alt_path.exists()}")
+                        
+                        if alt_path.exists():
+                            logger.info("Using alternative path format")
+                            sdxl_checkpoint_path = alt_path
+                        else:
+                            await websocket.send(json.dumps({
+                                'type': 'setup_response',
+                                'success': False,
+                                'error': f'SDXL checkpoint path does not exist: {sdxl_checkpoint_path}'
+                            }))
+                            continue
+                    
+                    # Load the model with progress reporting
+                    logger.info(f"Loading model from: {sdxl_checkpoint_path}")
+                    pipe, success = await load_model(sdxl_checkpoint_path, progress_callback)
+                    
+                    if not success:
+                        await websocket.send(json.dumps({
+                            'type': 'setup_response',
+                            'success': False,
+                            'error': 'Failed to load model'
+                        }))
+                        continue
+                    
+                    # Load LoRA if specified
+                    lora_success = True
+                    if lora_data and lora_data['path'] and pathlib.Path(lora_data['path']).exists():
+                        lora_success = await load_lora(pipe, pathlib.Path(lora_data['path']), progress_callback=progress_callback)
+                        if not lora_success:
+                            await websocket.send(json.dumps({
+                                'type': 'setup_response',
+                                'success': False,
+                                'error': f'Failed to load LoRA from {pathlib.Path(lora_data["path"])}'
+                            }))
+                            continue
+                    
+                    model_initialized = True
+                    await websocket.send(json.dumps({
+                        'type': 'setup_response',
+                        'success': True,
+                        'message': 'Model initialized successfully',
+                        'settings': current_settings
+                    }))
+                
+                # Handle settings update - changes settings without full reinitialization
+                elif message_type == 'update_settings':
+                    if not model_initialized:
+                        await websocket.send(json.dumps({
+                            'type': 'update_settings_response',
+                            'success': False,
+                            'error': 'Model not initialized. Send a setup request first.'
+                        }))
+                        continue
+                    
+                    settings_data = data.get('settings', {})
+                    updates_applied = []
+                    errors = []
+
+                    # Check for LoRA path update
+                    if 'lora_path' in settings_data and settings_data['lora_path'] != current_settings.get('lora_path'):
+                        new_lora_path = pathlib.Path(settings_data['lora_path']) if settings_data['lora_path'] else None
+                        
+                        if new_lora_path and new_lora_path.exists():
+                            # Load the new LoRA with progress reporting
+                            lora_success = await load_lora(pipe, new_lora_path, progress_callback=progress_callback)
+                            if lora_success:
+                                current_settings['lora_path'] = str(new_lora_path)
+                                updates_applied.append('lora_path')
+                            else:
+                                errors.append(f'Failed to load LoRA from {new_lora_path}')
+                        elif new_lora_path:
+                            errors.append(f'LoRA path does not exist: {new_lora_path}')
+                        else:
+                            # If LoRA path is None, we'll just update the setting
+                            current_settings['lora_path'] = None
+                            updates_applied.append('lora_path')
+                    
+                    # Handle any other settings that can be updated without reloading the model
+                    # For example, default generation parameters
+                    for param in ['default_lora_strength', 'default_image_to_image_strength',
+                                 'default_width', 'default_height', 'default_steps',
+                                 'default_guidance_scale']:
+                        if param in settings_data:
+                            current_settings[param] = settings_data[param]
+                            updates_applied.append(param)
+                    
+                    # Send response
+                    if errors:
+                        await websocket.send(json.dumps({
+                            'type': 'update_settings_response',
+                            'success': False,
+                            'error': ', '.join(errors),
+                            'updates_applied': updates_applied,
+                            'current_settings': current_settings
+                        }))
+                    else:
+                        await websocket.send(json.dumps({
+                            'type': 'update_settings_response',
+                            'success': True,
+                            'message': 'Settings updated successfully',
+                            'updates_applied': updates_applied,
+                            'current_settings': current_settings
+                        }))
+                
+                # Handle generation request
+                elif message_type == 'generate':
+                    if not model_initialized:
+                        await websocket.send(json.dumps({
+                            'type': 'generate_response',
+                            'success': False,
+                            'error': 'Model not initialized. Send a setup request first.'
+                        }))
+                        continue
+                    
+                    generate_data = data.get('generate', {})
+                    
+                    # Extract generation parameters
+                    base64_image = generate_data.get('image')
+                    prompt = generate_data.get('prompt', '')
+                    
+                    # Use default settings if not provided
+                    lora_strength = float(generate_data.get('lora_strength', current_settings.get('default_lora_strength', 1.0)))
+                    image_to_image_strength = float(generate_data.get('image_to_image_strength', current_settings.get('default_image_to_image_strength', 1.0)))
+                    width = int(generate_data.get('generated_image_width', current_settings.get('default_width', 1024)))
+                    height = int(generate_data.get('generated_image_height', current_settings.get('default_height', 1024)))
+                    num_inference_steps = int(generate_data.get('num_inference_steps', current_settings.get('default_steps', 4)))
+                    guidance_scale = float(generate_data.get('guidance_scale', current_settings.get('default_guidance_scale', 1.0)))
+                    seed = int(generate_data.get('seed')) if 'seed' in generate_data else None
+                    
+                    # Validate required parameters
+                    if not base64_image:
+                        await websocket.send(json.dumps({
+                            'type': 'generate_response',
+                            'success': False,
+                            'error': 'Missing base64 image'
+                        }))
+                        continue
+                    
+                    # Decode the input image
+                    input_image = decode_base64_to_image(base64_image)
+                    
+                    # Update for generation in progress
+                    await websocket.send(json.dumps({
+                        'type': 'progress_update',
+                        'stage': 'generating',
+                        'progress': 0,
+                        'file': ""
+                    }))
+                    
+                    # Generate the image
+                    output_image = generate_image(
+                        pipe=pipe,
+                        image=input_image,
+                        prompt=prompt,
+                        lora_strength=lora_strength,
+                        image_to_image_strength=image_to_image_strength,
+                        width=width,
+                        height=height,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        seed=seed
+                    )
+                    
+                    # Update for generation complete
+                    await websocket.send(json.dumps({
+                        'type': 'progress_update',
+                        'stage': 'generating',
+                        'progress': 100,
+                        'file': ""
+                    }))
+                    
+                    if output_image is None:
+                        await websocket.send(json.dumps({
+                            'type': 'generate_response',
+                            'success': False,
+                            'error': 'Failed to generate image'
+                        }))
+                        continue
+                    
+                    # Encode the output image to base64
+                    output_base64 = encode_image_to_base64(output_image)
+                    
+                    # Send the result back to the client
+                    await websocket.send(json.dumps({
+                        'type': 'generate_response',
+                        'success': True,
+                        'image': output_base64
+                    }))
+                
+                # Handle status request
+                elif message_type == 'status':
+                    await websocket.send(json.dumps({
+                        'type': 'status_response',
+                        'success': True,
+                        'model_initialized': model_initialized,
+                        'current_settings': current_settings
+                    }))
+                
+                else:
+                    await websocket.send(json.dumps({
+                        'type': 'error',
+                        'error': f'Unknown message type: {message_type}'
+                    }))
+                    
+            except json.JSONDecodeError:
+                await websocket.send(json.dumps({
+                    'type': 'error',
+                    'error': 'Invalid JSON'
+                }))
+            except Exception as e:
+                logger.error(f"Error processing message: {e}")
+                await websocket.send(json.dumps({
+                    'type': 'error',
+                    'error': str(e)
+                }))
+    except websockets.exceptions.ConnectionClosed:
+        logger.info("Connection closed")
+    finally:
+        connected = False
+        logger.info("Client disconnected")
+
+async def main():
+    """Start the WebSocket server."""
+    server = await websockets.serve(handle_client, "localhost", 8765)
+    logger.info("Server started on ws://localhost:8765")
+    await server.wait_closed()
+
+if __name__ == "__main__":
+    asyncio.run(main()) 

frontend/realtime/src/KonvaApp/Engine.ts

@@ -1,9 +1,9 @@
 import Konva from "konva";
 
-
 import { uiAccess, uiEvents } from "~/signals";
 import { ShapeNode } from "./Nodes";
-
+import { PaintNode } from "./Nodes/PaintNode";
+import { PreviewCopyNode } from "./Nodes/PreviewCopy";
 import { UndoStackManager } from "./UndoRedo";
 import { CommandManager, MatteBox, SceneManager } from "./EngineUtitlities";
 import {
@@ -99,7 +99,6 @@ 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());
@@ -112,17 +111,8 @@ export class Engine {
     });
 
     // core layer for all the work done.
-
     this.offScreenCanvas = new OffscreenCanvas(0, 0);
 
-    this.realTimeDrawEngine = new RealTimeDrawEngine({
-      width: VideoResolutions.SQUARE_1024.width,
-      height: VideoResolutions.SQUARE_1024.height,
-      mediaLayerRef: this.mediaLayer,
-      bgLayerRef: this.bgLayer,
-      offScreenCanvas: this.offScreenCanvas
-    });
-
     // Collection of all Nodes
     this.nodesManager = new NodesManager();
     // Partial Collection of selected Nodes
@@ -130,6 +120,20 @@ export class Engine {
       nodeTransformerRef: this.nodeTransformer,
       mediaLayerRef: this.mediaLayer,
     });
+
+    this.realTimeDrawEngine = new RealTimeDrawEngine({
+      width: VideoResolutions.SQUARE_1024.width,
+      height: VideoResolutions.SQUARE_1024.height,
+      mediaLayerRef: this.mediaLayer,
+      offScreenCanvas: this.offScreenCanvas,
+      onDraw: async (canvas, lineBounds) => {
+        await this.addPaintNode(canvas, lineBounds);
+      },
+      onPreviewCopy: async (image) => {
+        await this.addPreviewCopy(image);
+      },
+    });
+
     // Selector Square to select Nodes
     this.selectorSquare = new SelectorSquare({
       captureCanvasRef: this.realTimeDrawEngine.captureCanvas,
@@ -141,7 +145,9 @@ export class Engine {
     this.uiLayer.add(this.selectorSquare.getKonvaNode());
 
     //Collection of commands for undo-redo
-    this.undoStackManager = new UndoStackManager();
+    this.undoStackManager = new UndoStackManager(() => {
+      this.realTimeDrawEngine.render();
+    });
     this.commandManager = new CommandManager({
       mediaLayerRef: this.mediaLayer,
       nodesManagerRef: this.nodesManager,
@@ -179,19 +185,31 @@ export class Engine {
     // some of the managers has events
     // hence, lastly, setup these events
     this.setupEventSystem();
-    this.setAppMode(AppModes.SELECT);
-
-
-
+    this.setAppMode(this.appMode);
   }
 
-
-
   private setAppMode(newAppMode: AppModes) {
     this.appMode = newAppMode;
     switch (this.appMode) {
+      case AppModes.PAINT: {
+        console.log("APPMODE: PAINT");
+        this.selectorSquare.disable(); // this breaks the paint if you use it in the wrong spot
+        this.selectionManager.disable(); // this prevents selection.
+        this.realTimeDrawEngine.enablePaintMode();
+        this.realTimeDrawEngine.disableDragging();
+
+        uiAccess.toolbarMain.enable();
+        uiAccess.toolbarMain.changeButtonState(ToolbarMainButtonNames.SELECT, {
+          active: false,
+        });
+        this.matteBox.disable();
+        return;
+      }
       case AppModes.SELECT: {
         console.log("APPMODE: SELECT");
+        this.realTimeDrawEngine.disablePaintMode();
+        this.realTimeDrawEngine.enableDragging();
+
         this.selectorSquare.enable();
         this.selectionManager.enable();
         uiAccess.toolbarMain.enable();
@@ -363,12 +381,30 @@ export class Engine {
     );
 
     // Listen to Toolbar Main
-    uiEvents.toolbarMain.UNDO.onClick(() => this.undoStackManager.undo());
-    uiEvents.toolbarMain.REDO.onClick(() => this.undoStackManager.redo());
+    uiEvents.toolbarMain.UNDO.onClick(() => {
+      this.undoStackManager.undo();
+    });
+    uiEvents.toolbarMain.REDO.onClick(() => {
+      this.undoStackManager.redo();
+      this.realTimeDrawEngine.render();
+    });
 
     uiEvents.toolbarMain.SAVE.onClick(async (/*event*/) => {
       await this.realTimeDrawEngine.saveOutput();
+    });
 
+    uiEvents.toolbarMain.SELECT.onClick(() => {
+      console.log("Toolbar Main >> Select");
+      this.setAppMode(AppModes.SELECT);
+    });
+
+    uiEvents.toolbarMain.PAINT.onClick(() => {
+      console.log("Toolbar Main >> Paint");
+      this.setAppMode(AppModes.PAINT);
+    });
+
+    uiEvents.toolbarMain.ERASER.onClick(() => {
+      console.log("Toolbar Main >> Eraser");
     });
 
     uiEvents.toolbarMain.PREVIEW.onClick(async () => {
@@ -399,8 +435,6 @@ export class Engine {
       this.addImage(image);
     });
 
-
-
     uiEvents.onAddTextToEngine((textdata) => {
       this.addText(textdata);
     });
@@ -465,7 +499,10 @@ export class Engine {
         return;
       }
       this.changeNodeColor(nodeColor);
-    })
+    });
+    uiEvents.toolbarMain.onPaintColorChanged((color) => {
+      this.realTimeDrawEngine.paintColor = color;
+    });
   }
 
   disableAllButtons() {
@@ -501,7 +538,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");
@@ -529,9 +566,9 @@ export class Engine {
       x: 10,
       y: 80,
       text: "",
-      fontSize: 24,
+      fontSize: 18,
       fontFamily: "Source Sans 3",
-      fill: "black",
+      fill: "white",
     });
     const anim = new Konva.Animation((frame) => {
       if (frame && import.meta.env.DEV) {
@@ -549,32 +586,37 @@ 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,
+          loaded: async () => {
+            this.realTimeDrawEngine.render();
+          },
         });
         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);
@@ -601,16 +643,30 @@ export class Engine {
       size: size,
       color: color,
       mediaLayerRef: this.mediaLayer,
-      selectionManagerRef: this.selectionManager
+      selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        await this.realTimeDrawEngine.render();
+      },
     });
     shapeNode.kNode.zIndex(1); // Replace desiredZIndex with a number
 
     this.commandManager.createNode(shapeNode);
+
     console.debug("Added shapenode:", shapeNode);
     console.debug("Added node's ID:", shapeNode.kNode._id);
-    //this.realTimeDrawEngine.addNodes(shapeNode);
   }
-
+  public addPreviewCopy(image: Image.Konva) {
+    // Start of Selection
+    const copyNode = new PreviewCopyNode({
+      image: image,
+      mediaLayerRef: this.mediaLayer,
+      selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        this.realTimeDrawEngine.render();
+      },
+    });
+    this.commandManager.createNode(copyNode);
+  }
   public addImage(imageFile: File) {
     const imageNode = new ImageNode({
       mediaLayerRef: this.mediaLayer,
@@ -618,10 +674,32 @@ export class Engine {
       canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
       imageFile: imageFile,
       selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        await this.realTimeDrawEngine.render();
+      },
     });
     imageNode.kNode.zIndex(1);
     this.commandManager.createNode(imageNode);
-    //this.realTimeDrawEngine.addNodes(imageNode);
+  }
+  public addPaintNode(
+    canvas: HTMLCanvasElement,
+    lineBounds: {
+      width: number;
+      height: number;
+      x: number;
+      y: number;
+    },
+  ) {
+    var node = new PaintNode({
+      canvasElement: canvas,
+      lineBounds: lineBounds,
+      mediaLayerRef: this.mediaLayer,
+      selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        await this.realTimeDrawEngine.render();
+      },
+    });
+    this.commandManager.createNode(node);
   }
 
   public addVideo(
@@ -665,7 +743,7 @@ export class Engine {
           if (
             longestVideoNode === undefined ||
             node.videoComponent.duration >
-            longestVideoNode.videoComponent.duration
+              longestVideoNode.videoComponent.duration
           ) {
             longestVideoNode = node;
           }

frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts

@@ -20,14 +20,15 @@ import {
   NodeTransformer,
   SelectionManager,
 } from "../NodesManagers";
-import { RenderEngine } from "../RenderingPrimitives/RenderEngine";
+import { RealTimeDrawEngine } from "../RenderingPrimitives/RealTimeDrawEngine";
+import { PaintNode } from "../Nodes/PaintNode";
 
 interface EngineReferences {
   mediaLayerRef: Konva.Layer;
   nodesManagerRef: NodesManager;
   nodeTransformerRef: NodeTransformer;
   selectionManagerRef: SelectionManager;
-  renderEngineRef: RenderEngine;
+  renderEngineRef: RealTimeDrawEngine;
   undoStackManagerRef: UndoStackManager;
 }
 
@@ -36,7 +37,7 @@ export class CommandManager {
   private nodesManagerRef: NodesManager;
   private nodeTransformerRef: NodeTransformer;
   private selectionManagerRef: SelectionManager;
-  private renderEngineRef: RenderEngine;
+  private renderEngineRef: RealTimeDrawEngine;
   private undoStackManagerRef: UndoStackManager;
 
   constructor(engineRefs: EngineReferences) {
@@ -48,7 +49,7 @@ export class CommandManager {
     this.undoStackManagerRef = engineRefs.undoStackManagerRef;
   }
 
-  createNode(node: VideoNode | ImageNode | TextNode | ShapeNode) {
+  createNode(node: MediaNode) {
     const command = new CreateCommand({
       nodes: new Set<MediaNode>([node]),
       mediaLayerRef: this.mediaLayerRef,
@@ -62,6 +63,7 @@ export class CommandManager {
     // Set the kNode's manual id to same as kNode's internal _id
     // This helps identify the kNode in the Konva layer with find
     node.kNode.id(node.kNode._id.toString());
+    this.renderEngineRef.render();
   }
   deleteNodes() {
     const nodes = this.selectionManagerRef.getSelectedNodes();
@@ -74,6 +76,7 @@ export class CommandManager {
       renderEngineRef: this.renderEngineRef,
     });
     this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
   }
   toggleLockNodes() {
     const nodes = this.selectionManagerRef.getSelectedNodes();
@@ -101,6 +104,7 @@ export class CommandManager {
       mediaLayerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
   }
   moveNodesDown() {
     const command = new MoveLayerDown({
@@ -109,6 +113,7 @@ export class CommandManager {
       mediaLayerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
   }
   translateNodes(props: {
     nodes: Set<MediaNode>;
@@ -120,6 +125,7 @@ export class CommandManager {
       layerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.pushCommand(command);
+    this.renderEngineRef.render();
   }
   transformNodes(props: {
     nodes: Set<MediaNode>;
@@ -131,6 +137,7 @@ export class CommandManager {
       layerRef: this.mediaLayerRef,
     });
     this.undoStackManagerRef.pushCommand(command);
+    this.renderEngineRef.render();
   }
   addChromaKey(props: { videoNode: VideoNode; newChromaColor: RGBColor }) {
     const command = new AddChromaKey(props);

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

@@ -68,7 +68,7 @@ export abstract class BaseNode {
     if (this.kNode instanceof Konva.Group) {
       const wrapperRect = this.kNode.findOne(".wrapper") as Konva.Rect;
       console.log("Wrapper rect:", wrapperRect);
-      if (wrapperRect !=null) {
+      if (wrapperRect != null) {
         wrapperRect.stroke(primaryOrange);
         wrapperRect.strokeWidth(highlightStrokeWidth);
       }
@@ -89,8 +89,8 @@ export abstract class BaseNode {
     }
     if (this.kNode instanceof Konva.Group) {
       const wrapperRect = this.kNode.findOne(".wrapper") as Konva.Rect;
-      console.log("Wrapper rect:", wrapperRect); 
-      if (wrapperRect !=null) {
+      console.log("Wrapper rect:", wrapperRect);
+      if (wrapperRect != null) {
         // Issue investigate
         wrapperRect.strokeWidth(0);
       }
@@ -138,7 +138,7 @@ export abstract class BaseNode {
       // console.log("handle select");
       if (!isMultiSelect) {
         // clear selection if not multislect
-        //console.log("No Shift >> no multiselect");
+        console.log("No Shift >> no multiselect");
         this.selectionManagerRef.clearSelection();
       }
       if (this.selectionManagerRef.isNodeSelected(this)) {
@@ -151,8 +151,25 @@ export abstract class BaseNode {
       }
     };
 
+    const handleRightClick = () => {
+      this.selectionManagerRef.clearSelection();
+      this.selectionManagerRef.selectNode(this, undefined, true);
+    }
+
+    this.kNode.on("contextmenu", (e) => {
+      console.log("CONTEXT MENU", e);
+      e.evt.preventDefault();
+      handleRightClick();
+    });
+
     this.kNode.on("mousedown", (e) => {
-      // console.log("MOUSE DOWN");
+      // https://developer.mozilla.org/en-US/docs/Web/API/UIEvent/which
+      // 3: Right mouse button
+      // If right click, let context event handle it
+      if (e.evt.button === 3) {
+        return;
+      }
+
       // Selection of Node
       if (!this.selectionManagerRef.isNodeSelected(this)) {
         this.isSelecting = true;
@@ -162,7 +179,13 @@ export abstract class BaseNode {
     });
 
     this.kNode.on("mouseup", (e) => {
-      // console.log("MOUSE UP");
+      // https://developer.mozilla.org/en-US/docs/Web/API/UIEvent/which
+      // 3: Right mouse button
+      // If right click, let context event handle it
+      if (e.evt.button === 3) {
+        return;
+      }
+
       // just coming out of dragging or selecting mode
       // no need to handle select
       if (this.selectionManagerRef.isDragging() || this.isSelecting) {

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

@@ -17,12 +17,13 @@ interface ImageNodeContructor {
   transform?: TransformationData;
   mediaLayerRef: Konva.Layer;
   selectionManagerRef: SelectionManager;
+  loaded: () => Promise<void>;
 }
 
 export class ImageNode extends NetworkedNode {
   public kNode: Konva.Image;
   public imageSize?: Size;
-
+  private loaded!: () => Promise<void>;
   constructor({
     canvasPosition,
     canvasSize,
@@ -32,6 +33,7 @@ export class ImageNode extends NetworkedNode {
     transform: existingTransform,
     mediaLayerRef,
     selectionManagerRef,
+    loaded,
   }: ImageNodeContructor) {
     // kNodes need to be created first to guaruntee
     // that it is not undefined in parent's context
@@ -53,6 +55,7 @@ export class ImageNode extends NetworkedNode {
       localFile: imageFile,
     });
     this.kNode = kNode;
+    this.loaded = loaded;
     // this.imageSize = minNodeSize;
     this.mediaLayerRef.add(this.kNode);
 
@@ -108,6 +111,7 @@ export class ImageNode extends NetworkedNode {
       this.listenToBaseKNode();
       this.mediaLayerRef.draw();
       //this.uploadImage(imageFile);
+      this.loaded();
     };
     imageComponent.onerror = () => {
       //this.setProgress({ progress: 0, status: UploadStatus.ERROR_ON_FILE });

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

@@ -0,0 +1,92 @@
+import Konva from "konva";
+import { SelectionManager } from "../NodesManagers";
+import { Position, NodeData } from "../types";
+import { NodeType } from "./constants";
+
+import { NetworkedNode } from "./NetworkedNode";
+
+interface PaintNodeConstructor {
+  canvasElement:HTMLCanvasElement;
+  lineBounds: {
+    width: number;
+    height: number;
+    x: number;
+    y: number;
+    };
+  mediaLayerRef: Konva.Layer;
+  selectionManagerRef: SelectionManager;
+  
+  loaded: () => Promise<void>;
+}
+
+export class PaintNode extends NetworkedNode {
+  retry(): void {
+    throw new Error("Method not implemented.");
+  }
+
+  constructor({
+    canvasElement,
+    lineBounds,
+    mediaLayerRef,
+    selectionManagerRef,
+    loaded,
+  }: PaintNodeConstructor) {
+
+    const cl = canvasElement
+
+    // Convert circle to image
+    const dataURL = cl.toDataURL();
+    const imageObj = new Image();
+    imageObj.src = dataURL;
+
+    const paintShape = new Konva.Image({
+        x: lineBounds.x,
+        y: lineBounds.y,
+        width: lineBounds.width,
+        height: lineBounds.height,
+        image: imageObj,
+        listening: false, // Disable event listening to allow captureCanvas to receive events
+        draggable: true,
+        // stroke: '#000000',
+        // strokeWidth: 2,
+        // dash: [5, 5], // Creates dotted outline THIS IS FOR DEBUG
+        globalCompositeOperation: 'source-over',
+        fill: 'transparent' 
+      });
+
+    imageObj.onload = () => {
+        this.mediaLayerRef.draw();
+        loaded();
+    };
+  
+    super({
+      selectionManagerRef: selectionManagerRef,
+      mediaLayerRef: mediaLayerRef,
+      kNode: paintShape,
+    });
+
+    this.mediaLayerRef.add(this.kNode);
+    this.listenToBaseKNode();
+    this.mediaLayerRef.draw();
+  }
+
+  public getNodeData(canvasPosition: Position) {
+    const data: NodeData = {
+      type: NodeType.SHAPE,
+      transform: {
+        position: {
+          x: this.kNode.position().x - canvasPosition.x,
+          y: this.kNode.position().y - canvasPosition.y,
+        },
+        size: this.kNode.size(),
+        rotation: this.kNode.rotation(),
+        scale: {
+          x: this.kNode.scaleX(),
+          y: this.kNode.scaleY(),
+        },
+        zIndex: this.kNode.getZIndex(),
+      },
+    };
+    return data;
+  }
+}

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

@@ -0,0 +1,58 @@
+import Konva from "konva";
+import { SelectionManager } from "../NodesManagers";
+import { Position, Size, NodeData, TransformationData } from "../types";
+import { NodeType } from "./constants";
+import { NodeUtilities } from "./NodeUtilities";
+import { NetworkedNode } from "./NetworkedNode";
+
+interface PreviewCopyConstructor {
+  image: Konva.Image;
+  mediaLayerRef: Konva.Layer;
+  selectionManagerRef: SelectionManager;
+  loaded: () => Promise<void>;
+}
+
+export class PreviewCopyNode extends NetworkedNode {
+  retry(): void {
+    throw new Error("Method not implemented.");
+  }
+
+  constructor({
+    image,
+    mediaLayerRef,
+    selectionManagerRef,
+    loaded,
+  }: PreviewCopyConstructor) {
+    // Create the actual shape inside the group
+
+    super({
+      selectionManagerRef: selectionManagerRef,
+      mediaLayerRef: mediaLayerRef,
+      kNode: image,
+    });
+
+    this.mediaLayerRef.add(this.kNode);
+    this.listenToBaseKNode();
+    this.mediaLayerRef.draw();
+  }
+
+  public getNodeData(canvasPosition: Position) {
+    const data: NodeData = {
+      type: NodeType.SHAPE,
+      transform: {
+        position: {
+          x: this.kNode.position().x - canvasPosition.x,
+          y: this.kNode.position().y - canvasPosition.y,
+        },
+        size: this.kNode.size(),
+        rotation: this.kNode.rotation(),
+        scale: {
+          x: this.kNode.scaleX(),
+          y: this.kNode.scaleY(),
+        },
+        zIndex: this.kNode.getZIndex(),
+      },
+    };
+    return data;
+  }
+}

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

@@ -20,10 +20,10 @@ interface ShapeNodeConstructor {
   transform?: TransformationData;
   mediaLayerRef: Konva.Layer;
   selectionManagerRef: SelectionManager;
+  loaded: () => Promise<void>;
 }
 
 export class ShapeNode extends NetworkedNode {
-  
   retry(): void {
     throw new Error("Method not implemented.");
   }
@@ -36,10 +36,11 @@ export class ShapeNode extends NetworkedNode {
     canvasSize,
     shapeType,
     size,
-    color = "#ff0000", // Default red if no color provided
+    color = "#2D81FF", // Default red if no color provided
     transform: existingTransform,
     mediaLayerRef,
     selectionManagerRef,
+    loaded,
   }: ShapeNodeConstructor) {
     // Create the actual shape inside the group
     const transform = NodeUtilities.getInitialTransform({
@@ -49,7 +50,7 @@ export class ShapeNode extends NetworkedNode {
     });
 
     let shape = null;
-    
+
     switch (shapeType) {
       case ShapeType.CIRCLE:
         const circleShape = new Konva.Circle({
@@ -72,13 +73,14 @@ export class ShapeNode extends NetworkedNode {
           width: circleShape.width(),
           height: circleShape.height(),
           ...transform,
-          fill: 'transparent',
+          fill: "transparent",
           strokeScaleEnabled: false,
           draggable: true,
         });
 
         imageObj.onload = () => {
           this.mediaLayerRef.draw();
+          loaded();
         };
 
         // Clean up the temporary circle
@@ -107,13 +109,14 @@ export class ShapeNode extends NetworkedNode {
           width: squareShape.width(),
           height: squareShape.height(),
           ...transform,
-          fill: 'transparent',
+          fill: "transparent",
           strokeScaleEnabled: false,
           draggable: true,
         });
 
         squareImageObj.onload = () => {
           this.mediaLayerRef.draw();
+          loaded()
         };
 
         // Clean up the temporary square
@@ -142,13 +145,14 @@ export class ShapeNode extends NetworkedNode {
           width: triangleShape.width(),
           height: triangleShape.height(),
           ...transform,
-          fill: 'transparent',
+          fill: "transparent",
           strokeScaleEnabled: false,
           draggable: true,
         });
 
         triangleImageObj.onload = () => {
           this.mediaLayerRef.draw();
+          loaded()
         };
 
         // Clean up the temporary triangle

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 = 10;
+export const highlightStrokeWidth = 2;
 export const minNodeSize = {
   width: 200,
   height: 200,

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

@@ -80,7 +80,7 @@ export class SelectionManager {
   public isDisabled() {
     return this.disable;
   }
-  public selectNode(node: MediaNode, doNotDraw?: boolean): boolean {
+  public selectNode(node: MediaNode, doNotDraw?: boolean, showContextMenu = false): boolean {
     if (
       this.disabled || // no-op when disabled
       this.selectedNodes.has(node) || // if the node is already selected
@@ -100,7 +100,11 @@ export class SelectionManager {
     this.selectedNodes.add(node);
     this.updateNodeTransformer();
     this.updateContextComponents();
-    this.showContextComponents();
+
+    console.log("SelectionManager > selectNode > showContextMenu", showContextMenu);
+    if (showContextMenu) {
+      this.showContextComponents();
+    }
 
     if (!doNotDraw) {
       this.mediaLayerRef.batchDraw();
@@ -189,8 +193,8 @@ export class SelectionManager {
     // done, clear data and show menu
     this.initialPositions.clear();
     this.updateContextComponents();
-    this.showContextComponents();
   }
+
   private getKNodeTransformation(kNode: Konva.Node) {
     return {
       kNodeId: kNode._id,

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

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

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

@@ -11,12 +11,21 @@ import { MediaNode } from "../types";
 
 import { RenderTask } from "./RenderTask";
 import { OffScreenSceneCanvas } from "./OffScreenSceneCanvas";
+import { Image } from "@tauri-apps/api/image";
+
+import { PaintNode } from "../Nodes/PaintNode";
+
+import {
+  ServerSetupPayload,
+  ServerSettingsPayload,
+  ServerResponse,
+} from "../types/ServerTypes";
 
 // https://www.aiseesoft.com/resource/phone-aspect-ratio-screen-resolution.html#:~:text=16%3A9%20Aspect%20Ratio
 
 export class RealTimeDrawEngine {
   private videoNodes: VideoNode[];
-  private imageNodes: (ImageNode | TextNode | ShapeNode)[];
+  private imageNodes: (ImageNode | TextNode | ShapeNode | PaintNode)[];
 
   private offScreenCanvas: OffscreenCanvas;
   private outputBitmap: ImageBitmap | undefined;
@@ -26,6 +35,7 @@ export class RealTimeDrawEngine {
   // capturing composite within window
 
   private mediaLayerRef: Konva.Layer;
+  private drawingsLayer: Konva.Layer; // New Layer for Drawings
 
   private height: number;
   private width: number;
@@ -46,20 +56,55 @@ export class RealTimeDrawEngine {
   public currentPrompt: string;
   public currentStrength: number;
 
+  // Paint Color
+  // paint Brush Size
+  // has to exit out of paint mode when shape or image are used.
+  public paintColor: string = "#000000";
+  public paintBrushSize: number = 5;
+  public isPaintMode: boolean = false;
+
+  private onDrawCallback?: (
+    canvas: HTMLCanvasElement,
+    lineBounds: {
+      width: number;
+      height: number;
+      x: number;
+      y: number;
+    },
+  ) => void;
+
+  private onPreviewCopyCallback?: (previewCopy: Konva.Image) => void; // New Callback
+
+  private serverSocket: WebSocket | null = null;
+
   constructor({
     width,
     height,
     mediaLayerRef,
     offScreenCanvas,
+    onDraw,
+    onPreviewCopy, // New Parameter
   }: {
     width: number;
     height: number;
     mediaLayerRef: Konva.Layer;
     offScreenCanvas: OffscreenCanvas;
+    onDraw?: (
+      canvas: HTMLCanvasElement,
+      lineBounds: {
+        width: number;
+        height: number;
+        x: number;
+        y: number;
+      },
+    ) => void;
+    onPreviewCopy?: (previewCopy: Konva.Image) => void; // New Parameter
   }) {
     this.videoLoadingCanvas = undefined;
     this.videoNodes = [];
     this.imageNodes = [];
+    this.onDrawCallback = onDraw;
+    this.onPreviewCopyCallback = onPreviewCopy; // Assign Callback
 
     // TODO: Make this dynamic and update this on change of canvas.
 
@@ -79,6 +124,12 @@ export class RealTimeDrawEngine {
     // this is the whole canvas
     this.mediaLayerRef = mediaLayerRef;
 
+    // Create a separate layer for drawings
+    this.drawingsLayer = new Konva.Layer({
+      clearBeforeDraw: true, // Ensures transparent background
+    });
+    this.mediaLayerRef.getStage()?.add(this.drawingsLayer); // to od pass in stage
+
     // Set background layer to red and media layer to green for visibility
 
     this.port = undefined;
@@ -121,24 +172,264 @@ export class RealTimeDrawEngine {
     // Add mouse events for preview canvas copying
     //this.previewCopyListener();
 
+    //this.startServer();
+  }
+
+  private isEnabled: boolean = false;
+  private cleanupFunction: (() => void) | null = null;
+
+  // this starts the python server
+  public startServer() {
+    // Create WebSocket connection
+    const socket = new WebSocket("ws://localhost:8765");
+
+    // Setup event handlers
+    socket.onopen = () => {
+      console.log("Connected to inference server");
+      // Send initial setup payload to load models
+      const setupPayload: ServerSetupPayload = {
+        type: "setup",
+        model: {
+          name: "stabilityai/sdxl-turbo",
+          path: "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/checkpoints/ponyDiffusionV6XL_v6StartWithThisOne.safetensors",
+          precision: "fp16",
+        },
+        lora: {
+          path: "F:/ComfyUI_windows_portable_nvidia/ComfyUI_windows_portable/ComfyUI/models/loras/LCM_LoRA_Weights_SDXL.safetensors",
+          alpha: 0.75,
+        },
+        device: "cuda",
+      };
+
+      socket.send(JSON.stringify(setupPayload));
+    };
+
+    // Handle incoming messages from server
+    socket.onmessage = (event) => {
+      try {
+        const response: ServerResponse = JSON.parse(event.data);
+        console.log("Server response:", response);
+        // Handle model loading progress updates
+        if (response.type === "progress") {
+          console.log(`Loading progress: ${response.percent}%`);
+          // You might want to update UI with loading progress here
+        }
+        // Handle successful model load
+        else if (
+          response.type === "success" &&
+          response.status === "model_loaded"
+        ) {
+          console.log("Model loaded successfully");
+          // Enable UI elements that depend on model being loaded
+        }
+        // Handle image generation results
+        else if (
+          response.type === "success" &&
+          response.status === "generation_complete"
+        ) {
+          if (response.image) {
+            // Convert base64 image to ImageBitmap for preview
+            this.base64ToImageBitmap(response.image).then((bitmap) => {
+              this.outputBitmap = bitmap;
+              this.previewCanvas.image(bitmap);
+              this.mediaLayerRef.draw();
+            });
+          }
+        }
+      } catch (error) {
+        console.error("Error parsing server response:", error);
+      }
+    };
+
+    socket.onerror = (error) => {
+      console.error("WebSocket error:", error);
+    };
+
+    socket.onclose = () => {
+      console.log("Disconnected from inference server");
+    };
+
+    // Store socket reference for later use
+    this.serverSocket = socket;
+  }
+
+  // Add method to update model settings
+  public updateServerSettings(settings: ServerSettingsPayload) {
+    if (this.serverSocket && this.serverSocket.readyState === WebSocket.OPEN) {
+      this.serverSocket.send(JSON.stringify(settings));
+    } else {
+      console.error(
+        "Cannot update settings: Server connection not established",
+      );
+    }
+  }
+
+  public paintMode() {
+    let isDrawing = false;
+    let currentLine: Konva.Line | null = null;
+
+    const startDrawing = (pos: { x: number; y: number }) => {
+      if (!this.isEnabled) return;
+      const stage = this.mediaLayerRef.getStage();
+      if (!stage) return;
+
+      // Convert pointer position to relative position within capture canvas
+      const captureBox = this.captureCanvas.getClientRect();
+      const relativeX = pos.x - captureBox.x;
+      const relativeY = pos.y - captureBox.y;
+
+      currentLine = new Konva.Line({
+        points: [relativeX, relativeY],
+        stroke: this.paintColor,
+        strokeWidth: 5,
+        lineCap: "round",
+        lineJoin: "round",
+        x: this.captureCanvas.x(),
+        y: this.captureCanvas.y(),
+        draggable: false,
+      });
+
+      this.drawingsLayer.add(currentLine); // Add to drawingsLayer
+      isDrawing = true;
+    };
+
+    const draw = (pos: { x: number; y: number }) => {
+      if (!this.isEnabled) return;
+      if (!isDrawing || !currentLine) return;
+
+      // Convert pointer position to relative position
+      const captureBox = this.captureCanvas.getClientRect();
+      const relativeX = pos.x - captureBox.x;
+      const relativeY = pos.y - captureBox.y;
+
+      const newPoints = currentLine.points().concat([relativeX, relativeY]);
+      currentLine.points(newPoints);
+      this.drawingsLayer.batchDraw();
+    };
+
+    const stopDrawing = () => {
+      if (!this.isEnabled) return;
+      if (!isDrawing || !currentLine) return;
+
+      // Store current line before resetting state
+      const lineToConvert = currentLine;
+
+      // Reset drawing state immediately so we can start a new stroke
+      isDrawing = false;
+      currentLine = null;
+
+      // Create a temporary layer to render the line
+      const tempLayer = new Konva.Layer({
+        clearBeforeDraw: true,
+      });
+      this.mediaLayerRef.getStage()?.add(tempLayer);
+      tempLayer.add(lineToConvert);
+
+      // Get the line canvas with transparent background
+      // Get the bounding box of the line
+      const lineBounds = lineToConvert.getClientRect();
+
+      // Create canvas with just enough size to contain the line
+      const lineCanvas = tempLayer.toCanvas({
+        x: lineBounds.x,
+        y: lineBounds.y,
+        width: lineBounds.width,
+        height: lineBounds.height,
+        pixelRatio: 1,
+      });
+
+      if (this.onDrawCallback) {
+        this.onDrawCallback(lineCanvas, lineBounds);
+      }
+
+      // Clean up
+      lineToConvert.destroy();
+      tempLayer.destroy();
+
+      this.drawingsLayer.batchDraw();
+    };
+
+    // Check if point is within capture canvas bounds
+    const isWithinCaptureCanvas = (pos: { x: number; y: number }) => {
+      const captureBox = this.captureCanvas.getClientRect();
+      return (
+        pos.x >= captureBox.x &&
+        pos.x <= captureBox.x + captureBox.width &&
+        pos.y >= captureBox.y &&
+        pos.y <= captureBox.y + captureBox.height
+      );
+    };
+
+    // Add event listeners
+    const stage = this.mediaLayerRef.getStage();
+    if (!stage) return;
+
+    stage.on("mousedown touchstart", (e) => {
+      const pos = stage.getPointerPosition();
+      if (pos && isWithinCaptureCanvas(pos)) {
+        startDrawing(pos);
+      }
+    });
+
+    stage.on("mousemove touchmove", (e) => {
+      const pos = stage.getPointerPosition();
+      if (pos && isWithinCaptureCanvas(pos)) {
+        draw(pos);
+      }
+    });
+
+    stage.on("mouseup touchend", async () => {
+      stopDrawing();
+      await this.render();
+    });
+
+    // Store cleanup function
+    this.cleanupFunction = () => {
+      stage.off("mousedown touchstart");
+      stage.off("mousemove touchmove");
+      stage.off("mouseup touchend");
+    };
+  }
+
+  public enablePaintMode() {
+    this.isEnabled = true;
+    if (!this.cleanupFunction) {
+      this.paintMode();
+    }
+  }
+
+  public enableDragging() {
+    // Enable dragging for all nodes in media layer
+    this.imageNodes?.forEach((node) => {
+      node.kNode.draggable(true);
+      node.kNode.listening(true);
+    });
+    this.mediaLayerRef.batchDraw();
+  }
+
+  public disableDragging() {
+    // Disable dragging for all nodes in media layer
+    this.imageNodes?.forEach((node) => {
+      node.kNode.draggable(false);
+      node.kNode.listening(false);
+    });
+    this.mediaLayerRef.batchDraw();
+  }
+
+  public disablePaintMode() {
+    this.isEnabled = false;
+    if (this.cleanupFunction) {
+      this.cleanupFunction();
+      this.cleanupFunction = null;
+    }
   }
 
   public previewCopyListener() {
-    this.previewCanvas.on('mousedown touchstart', () => {
+    this.previewCanvas.on("mousedown touchstart", () => {
       if (!this.outputBitmap) {
-        console.log("No preview image to copy"); 
+        console.log("No preview image to copy");
         return;
       }
-      
-      // TODO create it as a node.
-      
-      // const imageNode = new ImageNode({
-      //   mediaLayerRef: this.mediaLayer,
-      //   canvasPosition: this.renderEngine.captureCanvas.position(),
-      //   canvasSize: this.renderEngine.captureCanvas.size(),
-      //   imageFile: imageFile,
-      //   selectionManagerRef: this.selectionManager,
-      // });
 
       // Create draggable preview copy
       const previewCopy = new Konva.Image({
@@ -148,43 +439,49 @@ export class RealTimeDrawEngine {
         height: this.height,
         image: this.outputBitmap,
         draggable: true,
-        listening: true
+        listening: true,
       });
 
-      this.mediaLayerRef.add(previewCopy);
-      previewCopy.moveToTop();
-      this.mediaLayerRef.draw();
+      // this.drawingsLayer.add(previewCopy); // Add to drawingsLayer instead of mediaLayerRef
+      // this.drawingsLayer.batchDraw();
+
+      // Invoke the callback with the preview copy
+      if (this.onPreviewCopyCallback) {
+        this.onPreviewCopyCallback(previewCopy);
+      }
 
       // Start dragging immediately
       previewCopy.startDrag();
 
       // Handle drag events
-      previewCopy.on('dragmove', () => {
-        this.mediaLayerRef.draw();
+      previewCopy.on("dragmove", () => {
+        // this.drawingsLayer.draw();
       });
 
-      previewCopy.on('dragend', () => {
+      previewCopy.on("dragend", () => {
         const previewBox = previewCopy.getClientRect();
         const captureBox = this.captureCanvas.getClientRect();
-        
+
         if (Konva.Util.haveIntersection(previewBox, captureBox)) {
           // Snap to capture canvas position
           previewCopy.position({
             x: this.captureCanvas.x(),
-            y: this.captureCanvas.y()
+            y: this.captureCanvas.y(),
           });
-          this.mediaLayerRef.batchDraw();
+          //this.drawingsLayer.batchDraw();
         } else {
           // Remove if dropped outside capture area
-          previewCopy.destroy();
-          this.mediaLayerRef.batchDraw();
+          //previewCopy.destroy();
+          //this.drawingsLayer.batchDraw();
         }
       });
     });
   }
 
-  public findImageNodeById(id: string): (ImageNode | TextNode | ShapeNode | undefined) {
-    return this.imageNodes.find(node => {
+  public findImageNodeById(
+    id: string,
+  ): ImageNode | TextNode | ShapeNode | PaintNode | undefined {
+    return this.imageNodes.find((node) => {
       if (node.kNode) {
         return node.kNode.id() === id;
       }
@@ -327,21 +624,33 @@ export class RealTimeDrawEngine {
   private handleNodeDragEnd = async () => {
     // Clean up any existing state
     if (this.isProcessing) {
+      console.log("isProcessing Returning");
       return;
     }
     console.log("Node drag ended");
     this.isProcessing = true;
     await this.render();
   };
-  
-  public addNodes(node: MediaNode) {
-   
-    if (node instanceof ImageNode || node instanceof TextNode || node instanceof ShapeNode) {
+
+  public async addNodes(node: MediaNode) {
+    if (
+      node instanceof ImageNode ||
+      node instanceof TextNode ||
+      node instanceof ShapeNode ||
+      node instanceof PaintNode
+    ) {
       console.debug("Adding node:", node);
       this.imageNodes.push(node);
-      console.log(this.imageNodes)
-      node.kNode.on("dragend", this.handleNodeDragEnd);
+      console.log(this.imageNodes);
+      //node.kNode.on("dragend", this.handleNodeDragEnd);
     }
+
+    // ensure the layer doesn't move if added while painting.
+    if (this.isEnabled) {
+      this.disableDragging();
+    }
+
+    await this.render();
   }
 
   public removeNodes(node: MediaNode) {
@@ -351,7 +660,12 @@ export class RealTimeDrawEngine {
         node.kNode.off("dragend", this.handleNodeDragEnd);
         this.videoNodes.splice(index, 1);
       }
-    } else if (node instanceof ImageNode || node instanceof TextNode) {
+    } else if (
+      node instanceof ImageNode ||
+      node instanceof TextNode ||
+      node instanceof ShapeNode ||
+      node instanceof PaintNode
+    ) {
       const index = this.imageNodes.indexOf(node);
       if (index > -1) {
         node.kNode.off("dragend", this.handleNodeDragEnd);
@@ -375,7 +689,6 @@ export class RealTimeDrawEngine {
     try {
       const box = config.layerOfInterest.getClientRect();
       const stage = config.layerOfInterest.getStage();
-      
 
       const x = config.x !== undefined ? config.x : Math.floor(box.x);
       const y = config.y !== undefined ? config.y : Math.floor(box.y);
@@ -461,7 +774,7 @@ export class RealTimeDrawEngine {
     base64String: string,
   ): Promise<ImageBitmap> {
     // Create an image element
-    const img = new Image();
+    const img = document.createElement("img");
 
     // Convert base64 to data URL if it doesn't include the prefix
     const dataUrl = base64String.startsWith("data:")
@@ -512,6 +825,7 @@ export class RealTimeDrawEngine {
 
   public async render() {
     // only pick nodes that intersect wi th the canvas on screen bounds to freeze.
+
     this.mediaLayerRef.draw();
     // Output all nodes in mediaLayerRef
     const nodes = this.mediaLayerRef.getChildren();
@@ -533,7 +847,7 @@ export class RealTimeDrawEngine {
     })) as ImageBitmap;
 
     // Test code
-    if (false) {
+    if (true) {
       this.outputBitmap = bitmap;
       this.previewCanvas.image(bitmap);
       this.isProcessing = false;

frontend/realtime/src/KonvaApp/UndoRedo/CreateCommand.ts

@@ -6,7 +6,7 @@ import {
   NodeTransformer,
   SelectionManager,
 } from "../NodesManagers";
-import { RenderEngine } from "../RenderingPrimitives/RenderEngine";
+import { RealTimeDrawEngine } from "../RenderingPrimitives/RealTimeDrawEngine";
 
 import { toolbarNode } from "~/signals/uiAccess/toolbarNode";
 
@@ -19,7 +19,7 @@ export class CreateCommand implements ICommand {
   private nodeTransformerRef: NodeTransformer;
   private selectionManagerRef: SelectionManager;
 
-  private renderEngineRef: RenderEngine;
+  private renderEngineRef: RealTimeDrawEngine;
 
   constructor({
     nodes,
@@ -33,7 +33,7 @@ export class CreateCommand implements ICommand {
     mediaLayerRef: Konva.Layer;
     nodesManagerRef: NodesManager;
     nodeTransformerRef: NodeTransformer;
-    renderEngineRef: RenderEngine;
+    renderEngineRef: RealTimeDrawEngine;
     selectionManagerRef: SelectionManager;
   }) {
     this.nodes = nodes;
@@ -53,6 +53,7 @@ export class CreateCommand implements ICommand {
       this.renderEngineRef.addNodes(node);
     });
     this.mediaLayerRef.draw();
+
   }
 
   undo() {
@@ -71,5 +72,7 @@ export class CreateCommand implements ICommand {
     this.nodeTransformerRef.clear();
     toolbarNode.hide();
     this.mediaLayerRef.draw();
+    this.renderEngineRef.render();
+ 
   }
 }

frontend/realtime/src/KonvaApp/UndoRedo/UndoRedoManager.ts

@@ -5,6 +5,11 @@ export class UndoStackManager {
   private undoStack: ICommand[] = [];
   private redoStack: ICommand[] = [];
   private readonly STACK_LIMIT = 20;
+  private onCommandExecuted: () => void;
+
+  constructor(onCommandExecuted: () => void) {
+    this.onCommandExecuted = onCommandExecuted;
+  }
 
   setDisabled(newState: boolean) {
     this.disabled = newState;
@@ -50,6 +55,7 @@ export class UndoStackManager {
       }
       console.log("undo:", this.undoStack);
     }
+    this.onCommandExecuted();
   }
 
   redo() {
@@ -68,5 +74,6 @@ export class UndoStackManager {
 
       console.log("Redo:", this.redoStack);
     }
+    this.onCommandExecuted();
   }
 }

frontend/realtime/src/KonvaApp/constants.ts

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

frontend/realtime/src/KonvaApp/types/NodeData.ts

@@ -3,10 +3,11 @@ import { NodeType, BaseNode, ImageNode, TextNode, VideoNode, ShapeNode } from ".
 import { TransformationData } from "./Transformation";
 
 import { TextNodeData } from "./Text";
+import { PaintNode } from "../Nodes/PaintNode";
 
 export type MediaNode =
   // | NetworkedNode
-  BaseNode | VideoNode | ImageNode | TextNode | ShapeNode;
+  BaseNode | VideoNode | ImageNode | TextNode | ShapeNode | PaintNode;
 
 export type ImageNodeData = {
   mediaFileUrl: string;

frontend/realtime/src/KonvaRootComponent/SignaledPromptSlider.tsx

@@ -1,22 +1,28 @@
-
 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;
 
-    const sliderValue = promptStrength.value
-    const onSliderValueChanged = dispatchUiEvents.promptSettings.setPromptStrength
+  // This makes component re-render whenever the signal value changes
+  useSignalEffect(() => {
+    promptStrength.value;
+  });
 
-    // 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>
-    )
-}
+  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>
+  );
+};

frontend/realtime/src/KonvaRootComponent/SignaledPromptText.tsx

@@ -1,31 +1,17 @@
 import { useSignalEffect } from "@preact/signals-react";
-import { useCallback, useEffect, useRef } from "react";
+import { 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 = {
-    style: {
-      width: "530px",
-      height: "200px",
-    },
-    placeholder: "",
+    placeholder: "Type your prompt here...",
     label: "",
     resize: "none" as ResizeType,
+    rows: 3,
   };
 
   const textContent = promptText.value;
@@ -36,11 +22,18 @@ export const SignaledPromptText = ({
   // This makes component re-render whenever the signal value changes
   useSignalEffect(() => {
     promptText.value;
-  })
+  });
 
   return (
-    <div className="fixed bottom-24 left-1/2 -translate-x-1/2">
-      <Textarea {...TextAreaStates} ref={textAreaRef} onChange={onTextChanged} value={textContent} />
+    <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>
   );
 };

frontend/realtime/src/KonvaRootComponent/SignaledToolbarMain.tsx

@@ -50,6 +50,13 @@ 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];
     }
@@ -76,7 +83,7 @@ export const SignaledToolbarMain = ({
     dispatchers.loadingBarRetry(e);
   };
   return (
-    <div className="fixed bottom-4 left-1/2 -translate-x-1/2">
+    <div className="fixed left-4 top-1/2 -translate-y-1/2">
       <Transition
         as="div"
         className={twMerge(

frontend/realtime/src/PyServer/ServerClient.ts

@@ -0,0 +1,260 @@
+import {
+  ServerMessage,
+  WebSocketHandlers,
+  SetupRequest,
+  UpdateSettingsRequest,
+  GenerateRequest,
+  StatusRequest,
+  GenerationParameters,
+  ServerSettings,
+  createSetupRequest,
+  createUpdateSettingsRequest,
+  createGenerateRequest,
+  createStatusRequest,
+} from "./ServerTypes";
+
+/**
+ * ServerClient provides a typed interface for communicating with the Python WebSocket server
+ */
+export class ServerClient {
+  private ws: WebSocket | null = null;
+  private handlers: WebSocketHandlers = {};
+  private connectionPromise: Promise<void> | null = null;
+  private connectionResolve: (() => void) | null = null;
+  private connectionReject: ((error: Error) => void) | null = null;
+
+  /**
+   * Create a new ServerClient
+   * @param url The WebSocket URL to connect to
+   * @param handlers Optional event handlers for different message types
+   */
+  constructor(
+    private url: string,
+    handlers: Partial<WebSocketHandlers> = {},
+  ) {
+    this.handlers = { ...handlers };
+  }
+
+  /**
+   * Connect to the WebSocket server
+   * @returns A promise that resolves when connected
+   */
+  public connect(): Promise<void> {
+    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
+      return Promise.resolve();
+    }
+
+    if (this.connectionPromise) {
+      return this.connectionPromise;
+    }
+
+    this.connectionPromise = new Promise<void>((resolve, reject) => {
+      this.connectionResolve = resolve;
+      this.connectionReject = reject;
+
+      try {
+        this.ws = new WebSocket(this.url);
+
+        this.ws.onopen = () => {
+          if (this.connectionResolve) {
+            this.connectionResolve();
+          }
+          if (this.handlers.onOpen) {
+            this.handlers.onOpen();
+          }
+        };
+
+        this.ws.onclose = () => {
+          if (this.handlers.onClose) {
+            this.handlers.onClose();
+          }
+          this.ws = null;
+          this.connectionPromise = null;
+        };
+
+        this.ws.onerror = (event) => {
+          const error = new Error(`WebSocket error: ${event}`);
+          if (this.connectionReject) {
+            this.connectionReject(error);
+          }
+          if (this.handlers.onError) {
+            this.handlers.onError({
+              type: "error",
+              error: `Connection error: ${event}`,
+            });
+          }
+        };
+
+        this.ws.onmessage = (event) => {
+          try {
+            const message = JSON.parse(event.data) as ServerMessage;
+            this.handleMessage(message);
+          } catch (error) {
+            console.error("Error parsing message:", error);
+            if (this.handlers.onError) {
+              this.handlers.onError({
+                type: "error",
+                error: `Failed to parse message: ${error}`,
+              });
+            }
+          }
+        };
+      } catch (error) {
+        if (this.connectionReject) {
+          this.connectionReject(error as Error);
+        }
+        this.connectionPromise = null;
+      }
+    });
+
+    return this.connectionPromise;
+  }
+
+  /**
+   * Disconnect from the WebSocket server
+   */
+  public disconnect(): void {
+    if (this.ws) {
+      this.ws.close();
+      this.ws = null;
+    }
+    this.connectionPromise = null;
+  }
+
+  /**
+   * Send a message to the server
+   * @param message The message to send
+   * @returns A promise that resolves when the message is sent
+   */
+  private async send(message: ServerMessage): Promise<void> {
+    await this.connect();
+
+    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
+      throw new Error("WebSocket is not connected");
+    }
+
+    this.ws.send(JSON.stringify(message));
+  }
+
+  /**
+   * Handle an incoming message from the server
+   * @param message The message received from the server
+   */
+  private handleMessage(message: ServerMessage): void {
+    switch (message.type) {
+      case "setup_response":
+        if (this.handlers.onSetupResponse) {
+          this.handlers.onSetupResponse(message);
+        }
+        break;
+
+      case "update_settings_response":
+        if (this.handlers.onUpdateSettingsResponse) {
+          this.handlers.onUpdateSettingsResponse(message);
+        }
+        break;
+
+      case "generate_response":
+        if (this.handlers.onGenerateResponse) {
+          this.handlers.onGenerateResponse(message);
+        }
+        break;
+
+      case "status_response":
+        if (this.handlers.onStatusResponse) {
+          this.handlers.onStatusResponse(message);
+        }
+        break;
+
+      case "progress_update":
+        if (this.handlers.onProgressUpdate) {
+          this.handlers.onProgressUpdate(message);
+        }
+        break;
+
+      case "error":
+      case "connection_error":
+        if (this.handlers.onError) {
+          this.handlers.onError(message);
+        }
+        break;
+    }
+  }
+
+  /**
+   * Set event handlers
+   * @param handlers The event handlers to set
+   */
+  public setHandlers(handlers: Partial<WebSocketHandlers>): void {
+    this.handlers = { ...this.handlers, ...handlers };
+  }
+
+  /**
+   * Send a setup request to initialize the model
+   * @param sdxl_checkpoint_path Path to the SDXL checkpoint
+   * @param lora_path Optional path to a LoRA file
+   */
+  public async setup(
+    sdxl_checkpoint_path: string,
+    lora_path?: string,
+  ): Promise<void> {
+    const request = createSetupRequest(sdxl_checkpoint_path, lora_path);
+    await this.send(request);
+  }
+
+  /**
+   * Send an update settings request
+   * @param settings The settings to update
+   */
+  public async updateSettings(
+    settings: Partial<ServerSettings>,
+  ): Promise<void> {
+    const request = createUpdateSettingsRequest(settings);
+    await this.send(request);
+  }
+
+  /**
+   * Send a generate request
+   * @param params The generation parameters
+   */
+  public async generate(params: GenerationParameters): Promise<void> {
+    const request = createGenerateRequest(params);
+    await this.send(request);
+  }
+
+  /**
+   * Send a status request
+   */
+  public async requestStatus(): Promise<void> {
+    const request = createStatusRequest();
+    await this.send(request);
+  }
+}
+
+/**
+ * Create a base64 string from an image
+ * @param file The image file
+ * @returns A promise that resolves to the base64 string
+ */
+export function fileToBase64(file: File): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+    reader.readAsDataURL(file);
+    reader.onload = () => {
+      const result = reader.result as string;
+      // Remove the data URL prefix (e.g., "data:image/png;base64,")
+      const base64 = result.split(",")[1];
+      resolve(base64);
+    };
+    reader.onerror = reject;
+  });
+}
+
+/**
+ * Create an image from a base64 string
+ * @param base64 The base64 string
+ * @returns The image URL
+ */
+export function base64ToImageUrl(base64: string): string {
+  return `data:image/png;base64,${base64}`;
+}

frontend/realtime/src/PyServer/ServerTypes.ts

@@ -0,0 +1,216 @@
+// Server Message Types - Common base type for all messages
+export type MessageType =
+  | "setup"
+  | "setup_response"
+  | "update_settings"
+  | "update_settings_response"
+  | "generate"
+  | "generate_response"
+  | "status"
+  | "status_response"
+  | "progress_update"
+  | "error"
+  | "connection_error";
+
+// Base interface for all messages
+interface BaseMessage {
+  type: MessageType;
+}
+
+// ===== Setup Messages =====
+
+// Setup Request
+export interface SetupRequest extends BaseMessage {
+  type: "setup";
+  setup: {
+    sdxl_checkpoint_path: string;
+    lora_path?: string;
+  };
+}
+
+// Setup Response
+export interface SetupResponse extends BaseMessage {
+  type: "setup_response";
+  success: boolean;
+  message?: string;
+  error?: string;
+  settings?: ServerSettings;
+}
+
+// ===== Update Settings Messages =====
+
+// Default generation settings that can be configured
+export interface GenerationSettings {
+  default_lora_strength?: number;
+  default_image_to_image_strength?: number;
+  default_width?: number;
+  default_height?: number;
+  default_steps?: number;
+  default_guidance_scale?: number;
+}
+
+// Server settings - includes both paths and generation defaults
+export interface ServerSettings extends GenerationSettings {
+  sdxl_checkpoint_path: string;
+  lora_path?: string | null;
+}
+
+// Update Settings Request
+export interface UpdateSettingsRequest extends BaseMessage {
+  type: "update_settings";
+  settings: Partial<ServerSettings>;
+}
+
+// Update Settings Response
+export interface UpdateSettingsResponse extends BaseMessage {
+  type: "update_settings_response";
+  success: boolean;
+  message?: string;
+  error?: string;
+  updates_applied?: string[];
+  current_settings: ServerSettings;
+}
+
+// ===== Generate Messages =====
+
+// Generation parameters
+export interface GenerationParameters {
+  image: string; // base64 encoded image
+  prompt: string;
+  lora_strength?: number;
+  image_to_image_strength?: number;
+  generated_image_width?: number;
+  generated_image_height?: number;
+  num_inference_steps?: number;
+  guidance_scale?: number;
+  seed?: number;
+}
+
+// Generate Request
+export interface GenerateRequest extends BaseMessage {
+  type: "generate";
+  generate: GenerationParameters;
+}
+
+// Generate Response
+export interface GenerateResponse extends BaseMessage {
+  type: "generate_response";
+  success: boolean;
+  image?: string; // base64 encoded image
+  error?: string;
+}
+
+// ===== Status Messages =====
+
+// Status Request
+export interface StatusRequest extends BaseMessage {
+  type: "status";
+}
+
+// Status Response
+export interface StatusResponse extends BaseMessage {
+  type: "status_response";
+  success: boolean;
+  model_initialized: boolean;
+  current_settings: ServerSettings;
+}
+
+// ===== Progress Update =====
+
+// Progress stages
+export type ProgressStage =
+  | "downloading"
+  | "loading_unet"
+  | "loading_base_model"
+  | "configuring_scheduler"
+  | "loading_lcm_adapter"
+  | "loading_lora"
+  | "generating"
+  | "error";
+
+// Progress Update Message
+export interface ProgressUpdateMessage extends BaseMessage {
+  type: "progress_update";
+  stage: ProgressStage;
+  progress: number; // 0-100
+  file: string;
+}
+
+// ===== Error Messages =====
+
+// Error Message
+export interface ErrorMessage extends BaseMessage {
+  type: "error";
+  error: string;
+}
+
+// Connection Error Message
+export interface ConnectionErrorMessage extends BaseMessage {
+  type: "connection_error";
+  error: string;
+}
+
+// Union type for all possible server messages
+export type ServerMessage =
+  | SetupRequest
+  | SetupResponse
+  | UpdateSettingsRequest
+  | UpdateSettingsResponse
+  | GenerateRequest
+  | GenerateResponse
+  | StatusRequest
+  | StatusResponse
+  | ProgressUpdateMessage
+  | ErrorMessage
+  | ConnectionErrorMessage;
+
+// WebSocket event handlers for different message types
+export interface WebSocketHandlers {
+  onSetupResponse?: (response: SetupResponse) => void;
+  onUpdateSettingsResponse?: (response: UpdateSettingsResponse) => void;
+  onGenerateResponse?: (response: GenerateResponse) => void;
+  onStatusResponse?: (response: StatusResponse) => void;
+  onProgressUpdate?: (update: ProgressUpdateMessage) => void;
+  onError?: (error: ErrorMessage | ConnectionErrorMessage) => void;
+  onOpen?: () => void;
+  onClose?: () => void;
+}
+
+// Helper function to create a typed setup request
+export function createSetupRequest(
+  sdxl_checkpoint_path: string,
+  lora_path?: string,
+): SetupRequest {
+  return {
+    type: "setup",
+    setup: {
+      sdxl_checkpoint_path,
+      lora_path,
+    },
+  };
+}
+
+// Helper function to create a typed update settings request
+export function createUpdateSettingsRequest(
+  settings: Partial<ServerSettings>,
+): UpdateSettingsRequest {
+  return {
+    type: "update_settings",
+    settings,
+  };
+}
+
+// Helper function to create a typed generate request
+export function createGenerateRequest(
+  params: GenerationParameters,
+): GenerateRequest {
+  return {
+    type: "generate",
+    generate: params,
+  };
+}
+
+// Helper function to create a typed status request
+export function createStatusRequest(): StatusRequest {
+  return { type: "status" };
+}

frontend/realtime/src/app/global.css

@@ -4,10 +4,11 @@
 
 @layer base {
   :root {
-    @apply bg-ui-background font-sans text-[#222222];
+    @apply bg-ui-background font-sans text-white;
   }
   .pegboard {
-    background-image: radial-gradient(#3f3f466b 1px, transparent 0);
+    background-image: radial-gradient(#ffffff25 1px, transparent 0);
+    background-color: #1b1b1f;
     background-size: 36px 36px;
     @apply bg-center;
   }
@@ -41,7 +42,24 @@
   }
 
   input {
-    @apply border-0 px-2 py-1 outline-none -outline-offset-2 focus:outline-primary-500 focus:ring-0;
+    @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;
+    }
   }
 }
 
@@ -87,6 +105,11 @@
     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,6 +8,7 @@ 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
@@ -25,6 +26,7 @@ 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/ImageGenerator.tsx

@@ -0,0 +1,524 @@
+import React, { useEffect, useState, useRef } from "react";
+import {
+  ServerClient,
+  fileToBase64,
+  base64ToImageUrl,
+} from "../PyServer/ServerClient";
+import {
+  ProgressUpdateMessage,
+  SetupResponse,
+  GenerateResponse,
+  StatusResponse,
+  UpdateSettingsResponse,
+} from "../PyServer/ServerTypes";
+
+// Default server URL
+const DEFAULT_SERVER_URL = "ws://localhost:8765";
+
+export const ImageGenerator: React.FC = () => {
+  // Server connection state
+  const [isConnected, setIsConnected] = useState(false);
+  const [serverUrl, setServerUrl] = useState(DEFAULT_SERVER_URL);
+  const clientRef = useRef<ServerClient | null>(null);
+
+  // Model state
+  const [modelPath, setModelPath] = useState("");
+  const [loraPath, setLoraPath] = useState("");
+  const [isModelInitialized, setIsModelInitialized] = useState(false);
+
+  // Generation parameters
+  const [sourceImage, setSourceImage] = useState<File | null>(null);
+  const [sourceImageUrl, setSourceImageUrl] = useState<string | null>(null);
+  const [prompt, setPrompt] = useState("");
+  const [loraStrength, setLoraStrength] = useState(1.0);
+  const [imageToImageStrength, setImageToImageStrength] = useState(0.75);
+  const [steps, setSteps] = useState(4);
+  const [guidanceScale, setGuidanceScale] = useState(1.0);
+  const [width, setWidth] = useState(1024);
+  const [height, setHeight] = useState(1024);
+
+  // Results and UI state
+  const [generatedImageUrl, setGeneratedImageUrl] = useState<string | null>(
+    null,
+  );
+  const [isLoading, setIsLoading] = useState(false);
+  const [progress, setProgress] = useState<ProgressUpdateMessage | null>(null);
+  const [errorMessage, setErrorMessage] = useState<string | null>(null);
+
+  // Initialize the client
+  useEffect(() => {
+    const client = new ServerClient(serverUrl, {
+      onOpen: () => {
+        console.log("Connected to server");
+        setIsConnected(true);
+        setErrorMessage(null);
+
+        // Request status to see if model is already initialized
+        client.requestStatus().catch((error) => {
+          console.error("Error requesting status:", error);
+        });
+      },
+
+      onClose: () => {
+        console.log("Disconnected from server");
+        setIsConnected(false);
+        setIsModelInitialized(false);
+      },
+
+      onError: (error) => {
+        console.error("Error:", error);
+        setErrorMessage(error.error);
+        setIsLoading(false);
+      },
+
+      onProgressUpdate: (update) => {
+        console.log("Progress:", update);
+        setProgress(update);
+
+        if (update.stage !== "error") {
+          setErrorMessage(null);
+        }
+      },
+
+      onSetupResponse: (response) => {
+        console.log("Setup response:", response);
+        setIsLoading(false);
+
+        if (response.success) {
+          setIsModelInitialized(true);
+          setErrorMessage(null);
+        } else {
+          setErrorMessage(response.error || "Unknown error during setup");
+        }
+      },
+
+      onUpdateSettingsResponse: (response) => {
+        console.log("Update settings response:", response);
+        setIsLoading(false);
+
+        if (response.success) {
+          setErrorMessage(null);
+        } else {
+          setErrorMessage(response.error || "Unknown error updating settings");
+        }
+      },
+
+      onGenerateResponse: (response) => {
+        console.log("Generate response:", response);
+        setIsLoading(false);
+
+        if (response.success && response.image) {
+          setGeneratedImageUrl(base64ToImageUrl(response.image));
+          setErrorMessage(null);
+        } else {
+          setErrorMessage(response.error || "Unknown error generating image");
+        }
+      },
+
+      onStatusResponse: (response) => {
+        console.log("Status response:", response);
+
+        if (response.success) {
+          setIsModelInitialized(response.model_initialized);
+
+          if (response.current_settings) {
+            // Update UI with current settings
+            setModelPath(response.current_settings.sdxl_checkpoint_path || "");
+            setLoraPath(response.current_settings.lora_path || "");
+
+            // Update generation parameters if they exist
+            if (response.current_settings.default_lora_strength !== undefined) {
+              setLoraStrength(response.current_settings.default_lora_strength);
+            }
+
+            if (
+              response.current_settings.default_image_to_image_strength !==
+              undefined
+            ) {
+              setImageToImageStrength(
+                response.current_settings.default_image_to_image_strength,
+              );
+            }
+
+            if (response.current_settings.default_steps !== undefined) {
+              setSteps(response.current_settings.default_steps);
+            }
+
+            if (
+              response.current_settings.default_guidance_scale !== undefined
+            ) {
+              setGuidanceScale(
+                response.current_settings.default_guidance_scale,
+              );
+            }
+
+            if (response.current_settings.default_width !== undefined) {
+              setWidth(response.current_settings.default_width);
+            }
+
+            if (response.current_settings.default_height !== undefined) {
+              setHeight(response.current_settings.default_height);
+            }
+          }
+        }
+      },
+    });
+
+    clientRef.current = client;
+
+    // Connect to the server
+    client.connect().catch((error) => {
+      console.error("Error connecting to server:", error);
+      setErrorMessage(`Failed to connect to server: ${error.message}`);
+    });
+
+    return () => {
+      client.disconnect();
+      clientRef.current = null;
+    };
+  }, [serverUrl]);
+
+  // Handle source image selection
+  const handleImageChange = async (e: React.ChangeEvent<HTMLInputElement>) => {
+    const file = e.target.files?.[0];
+    if (file) {
+      setSourceImage(file);
+      setSourceImageUrl(URL.createObjectURL(file));
+    }
+  };
+
+  // Handle setup submission
+  const handleSetup = async () => {
+    if (!clientRef.current) return;
+
+    setIsLoading(true);
+    setErrorMessage(null);
+
+    try {
+      await clientRef.current.setup(modelPath, loraPath || undefined);
+    } catch (error) {
+      console.error("Error setting up model:", error);
+      setErrorMessage(`Error setting up model: ${(error as Error).message}`);
+      setIsLoading(false);
+    }
+  };
+
+  // Handle update settings
+  const handleUpdateSettings = async () => {
+    if (!clientRef.current) return;
+
+    setIsLoading(true);
+    setErrorMessage(null);
+
+    try {
+      await clientRef.current.updateSettings({
+        lora_path: loraPath || null,
+        default_lora_strength: loraStrength,
+        default_image_to_image_strength: imageToImageStrength,
+        default_steps: steps,
+        default_guidance_scale: guidanceScale,
+        default_width: width,
+        default_height: height,
+      });
+    } catch (error) {
+      console.error("Error updating settings:", error);
+      setErrorMessage(`Error updating settings: ${(error as Error).message}`);
+      setIsLoading(false);
+    }
+  };
+
+  // Handle generation
+  const handleGenerate = async () => {
+    if (!clientRef.current || !sourceImage) {
+      setErrorMessage("Please select a source image first");
+      return;
+    }
+
+    setIsLoading(true);
+    setErrorMessage(null);
+    setGeneratedImageUrl(null);
+
+    try {
+      const base64Image = await fileToBase64(sourceImage);
+
+      await clientRef.current.generate({
+        image: base64Image,
+        prompt,
+        lora_strength: loraStrength,
+        image_to_image_strength: imageToImageStrength,
+        num_inference_steps: steps,
+        guidance_scale: guidanceScale,
+        generated_image_width: width,
+        generated_image_height: height,
+      });
+    } catch (error) {
+      console.error("Error generating image:", error);
+      setErrorMessage(`Error generating image: ${(error as Error).message}`);
+      setIsLoading(false);
+    }
+  };
+
+  return (
+    <div className="mx-auto max-w-4xl p-6">
+      <h1 className="mb-6 text-2xl font-bold">SDXL Image Generator</h1>
+
+      {/* Connection and Setup */}
+      <div className="mb-6 rounded-lg bg-gray-100 p-4">
+        <h2 className="mb-2 text-xl font-semibold">Server Connection</h2>
+        <div className="mb-4 flex gap-2">
+          <input
+            type="text"
+            value={serverUrl}
+            onChange={(e) => setServerUrl(e.target.value)}
+            className="flex-grow rounded border p-2"
+            placeholder="WebSocket URL"
+          />
+          <button
+            onClick={() => clientRef.current?.connect()}
+            disabled={isConnected}
+            className="rounded bg-blue-500 px-4 py-2 text-white disabled:bg-gray-400"
+          >
+            Connect
+          </button>
+          <button
+            onClick={() => clientRef.current?.disconnect()}
+            disabled={!isConnected}
+            className="rounded bg-red-500 px-4 py-2 text-white disabled:bg-gray-400"
+          >
+            Disconnect
+          </button>
+        </div>
+
+        <div className="mb-2 flex items-center">
+          <span className="mr-2">Status:</span>
+          <span
+            className={`rounded px-2 py-1 text-white ${isConnected ? "bg-green-500" : "bg-red-500"}`}
+          >
+            {isConnected ? "Connected" : "Disconnected"}
+          </span>
+          <span className="ml-4 mr-2">Model:</span>
+          <span
+            className={`rounded px-2 py-1 text-white ${isModelInitialized ? "bg-green-500" : "bg-yellow-500"}`}
+          >
+            {isModelInitialized ? "Initialized" : "Not Initialized"}
+          </span>
+        </div>
+      </div>
+
+      {/* Model Setup */}
+      <div className="mb-6 rounded-lg bg-gray-100 p-4">
+        <h2 className="mb-2 text-xl font-semibold">Model Setup</h2>
+        <div className="space-y-4">
+          <div>
+            <label className="mb-1 block">SDXL Checkpoint Path:</label>
+            <input
+              type="text"
+              value={modelPath}
+              onChange={(e) => setModelPath(e.target.value)}
+              className="w-full rounded border p-2"
+              placeholder="/path/to/sdxl/model"
+            />
+          </div>
+          <div>
+            <label className="mb-1 block">LoRA Path (optional):</label>
+            <input
+              type="text"
+              value={loraPath}
+              onChange={(e) => setLoraPath(e.target.value)}
+              className="w-full rounded border p-2"
+              placeholder="/path/to/lora"
+            />
+          </div>
+          <div className="flex gap-2">
+            <button
+              onClick={handleSetup}
+              disabled={!isConnected || !modelPath || isLoading}
+              className="rounded bg-green-500 px-4 py-2 text-white disabled:bg-gray-400"
+            >
+              Initialize Model
+            </button>
+            <button
+              onClick={handleUpdateSettings}
+              disabled={!isConnected || !isModelInitialized || isLoading}
+              className="rounded bg-purple-500 px-4 py-2 text-white disabled:bg-gray-400"
+            >
+              Update Settings
+            </button>
+          </div>
+        </div>
+      </div>
+
+      {/* Generation Parameters */}
+      <div className="mb-6 rounded-lg bg-gray-100 p-4">
+        <h2 className="mb-2 text-xl font-semibold">Generation Parameters</h2>
+        <div className="grid grid-cols-1 gap-4 md:grid-cols-2">
+          <div>
+            <label className="mb-1 block">Prompt:</label>
+            <textarea
+              value={prompt}
+              onChange={(e) => setPrompt(e.target.value)}
+              className="h-20 w-full rounded border p-2"
+              placeholder="Enter your prompt..."
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">Source Image:</label>
+            <input
+              type="file"
+              accept="image/*"
+              onChange={handleImageChange}
+              className="mb-2"
+            />
+            {sourceImageUrl && (
+              <img
+                src={sourceImageUrl}
+                alt="Source"
+                className="max-h-40 rounded border"
+              />
+            )}
+          </div>
+
+          <div>
+            <label className="mb-1 block">
+              LoRA Strength: {loraStrength.toFixed(2)}
+            </label>
+            <input
+              type="range"
+              min="0"
+              max="1"
+              step="0.05"
+              value={loraStrength}
+              onChange={(e) => setLoraStrength(parseFloat(e.target.value))}
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">
+              Image-to-Image Strength: {imageToImageStrength.toFixed(2)}
+            </label>
+            <input
+              type="range"
+              min="0"
+              max="1"
+              step="0.05"
+              value={imageToImageStrength}
+              onChange={(e) =>
+                setImageToImageStrength(parseFloat(e.target.value))
+              }
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">Steps: {steps}</label>
+            <input
+              type="range"
+              min="1"
+              max="20"
+              step="1"
+              value={steps}
+              onChange={(e) => setSteps(parseInt(e.target.value))}
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">
+              Guidance Scale: {guidanceScale.toFixed(2)}
+            </label>
+            <input
+              type="range"
+              min="0"
+              max="10"
+              step="0.1"
+              value={guidanceScale}
+              onChange={(e) => setGuidanceScale(parseFloat(e.target.value))}
+              className="w-full"
+            />
+          </div>
+
+          <div>
+            <label className="mb-1 block">Width:</label>
+            <select
+              value={width}
+              onChange={(e) => setWidth(parseInt(e.target.value))}
+              className="w-full rounded border p-2"
+            >
+              <option value="512">512</option>
+              <option value="768">768</option>
+              <option value="1024">1024</option>
+              <option value="1280">1280</option>
+            </select>
+          </div>
+
+          <div>
+            <label className="mb-1 block">Height:</label>
+            <select
+              value={height}
+              onChange={(e) => setHeight(parseInt(e.target.value))}
+              className="w-full rounded border p-2"
+            >
+              <option value="512">512</option>
+              <option value="768">768</option>
+              <option value="1024">1024</option>
+              <option value="1280">1280</option>
+            </select>
+          </div>
+        </div>
+
+        <button
+          onClick={handleGenerate}
+          disabled={
+            !isConnected || !isModelInitialized || !sourceImage || isLoading
+          }
+          className="mt-4 rounded bg-blue-500 px-6 py-2 text-white disabled:bg-gray-400"
+        >
+          Generate Image
+        </button>
+      </div>
+
+      {/* Progress and Results */}
+      {isLoading && progress && (
+        <div className="mb-6 rounded-lg bg-gray-100 p-4">
+          <h2 className="mb-2 text-xl font-semibold">Progress</h2>
+          <div className="mb-2">
+            <div className="mb-1 flex justify-between">
+              <span>{progress.stage}</span>
+              <span>{progress.progress}%</span>
+            </div>
+            <div className="h-2.5 w-full rounded-full bg-gray-300">
+              <div
+                className="h-2.5 rounded-full bg-blue-600"
+                style={{ width: `${progress.progress}%` }}
+              ></div>
+            </div>
+            {progress.file && (
+              <div className="mt-1 text-sm text-gray-600">{progress.file}</div>
+            )}
+          </div>
+        </div>
+      )}
+
+      {errorMessage && (
+        <div className="mb-6 rounded border border-red-400 bg-red-100 px-4 py-3 text-red-700">
+          <p>{errorMessage}</p>
+        </div>
+      )}
+
+      {generatedImageUrl && (
+        <div className="mb-6 rounded-lg bg-gray-100 p-4">
+          <h2 className="mb-2 text-xl font-semibold">Generated Image</h2>
+          <div className="flex justify-center">
+            <img
+              src={generatedImageUrl}
+              alt="Generated"
+              className="max-w-full rounded border"
+            />
+          </div>
+        </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-3xl font-bold">Upload Image</DialogTitle>
+          <DialogTitle className="text-2xl 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-3xl font-bold">Edit Text</DialogTitle>
+          <DialogTitle className="text-2xl 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-gray-100 hover:bg-gray-200/60 transition-colors duration-150 ease-in-out",
+    "group cursor-pointer px-6 py-12 bg-ui-controls hover:bg-white/25 transition-colors duration-150 ease-in-out rounded-lg",
     !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 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 ",
+    "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",
     disabled && "pointer-events-none text-secondary-300",
     buttonProps.prominent &&
-      "highlight-button border border-primary-400/30 text-primary-500 hover:bg-primary-100/40 hover:border-primary-400/60",
+      "border border-primary-400/30 text-primary-500 hover:bg-primary-100/40 hover:border-primary-400/60",
     customButtonClassNames,
   );
 
@@ -57,6 +57,7 @@ export const ToolbarButton = ({
     <button
       className={mergedButtonClasses}
       disabled={disabled}
+      autoFocus={false}
       {...restButtonProps}
       onClick={(e) => {
         e.preventDefault();
@@ -73,7 +74,11 @@ export const ToolbarButton = ({
     </button>
   );
   if (tooltip) {
-    return <Tooltip tip={tooltip}>{Button}</Tooltip>;
+    return (
+      <Tooltip position="right" tip={tooltip}>
+        {Button}
+      </Tooltip>
+    );
   }
   return Button;
 };

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

@@ -1,29 +1,32 @@
-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,
-  faMagicWandSparkles,
+  faPaintbrush,
   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";
+import { dispatchUiEvents, uiEvents } from "~/signals/uiEvents";
+import { paintColor } from "~/signals/uiEvents/toolbarMain/paintMode";
+import { ColorPicker } from "~/components/ui/TextEditor/ColorPicker";
+import { PaintModeMenu } from "~/components/ui/ToolbarMain/PaintModeMenu";
+import { EraseModeMenu } from "~/components/ui/ToolbarMain/EraseModeMenu";
 
 export const ToolbarMain = ({
   disabled = false,
@@ -42,13 +45,14 @@ export const ToolbarMain = ({
     <>
       <div
         className={twMerge(
-          "flex w-fit items-center divide-x divide-ui-border",
+          "flex h-fit flex-col items-center gap-y-2",
           paperWrapperStyles,
+          "glass",
           disabled &&
           "pointer-events-none cursor-default bg-ui-border shadow-md",
         )}
       >
-        <div className="flex items-center gap-2 pr-2">
+        <div className="flex flex-col items-center gap-2">
           <ToolbarButton
             icon={faLocationArrow}
             iconProps={{ className: "fa-flip-horizontal" }}
@@ -62,62 +66,150 @@ export const ToolbarMain = ({
           />
 
           <Popover className="relative">
-            <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>
+            {({ 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>
+              </>
+            )}
           </Popover>
 
           <Popover className="relative">
             <PopoverButton as={Fragment}>
               <button
-                data-tooltip="Add..."
+                data-tooltip="Add Media"
                 className={twMerge(
-                  "size-10 rounded-lg p-2 hover:bg-gray-200/50",
-                  toolTipStyles,
+                  "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={faFilePlus} />
               </button>
             </PopoverButton>
             <PopoverPanel
-              anchor="bottom"
+              anchor="right"
               className={twMerge(
                 // "absolute bottom-full left-1/2 z-10 mb-2 -translate-x-1/2",
-                "flex flex-col [--anchor-gap:16px]",
+                "flex [--anchor-gap:12px]",
                 paperWrapperStyles,
               )}
             >
               <ToolbarButton icon={faImage} buttonProps={buttonProps.ADD_IMAGE}>
-                Add Image
+                <span className="text-[16px]">Add Image</span>
               </ToolbarButton>
               {/* <ToolbarButton icon={faFilm} buttonProps={buttonProps.ADD_VIDEO}>
                 Add Video
               </ToolbarButton> */}
             </PopoverPanel>
           </Popover>
-
         </div>
-        <div className="flex items-center gap-2 pl-2">
+
+        <hr className="w-full border-t border-white/15" />
+
+        <div className="flex flex-col items-center gap-2">
+          <div className="flex flex-col items-center gap-2">
+            { /* Conditionally show paint menu on mode selection */}
+            {buttonProps.PAINT.active ? (
+              <div className="relative">
+                <PaintModeMenu
+                  color={paintColor.value}
+                  onChange={dispatchUiEvents.toolbarMain.setPaintColor}
+                  faIcon={faPaintbrush}
+                  borderStyle="border-2  bg-primary/30 border-2 border-primary hover:bg-primary/30 text-white"
+                  showBar={false}
+                  staticIconColor="white"
+                  streamChanges
+                  defaultOpen={true}
+                  anchor="right"
+                  anchorGap={12}
+                  closeOnMouseLeave={true}
+                />
+                <div
+                  className="pointer-events-none absolute -bottom-1.5 -right-1.5 h-[18px] w-[18px] rounded-full border border-gray-400"
+                  style={{ backgroundColor: paintColor.value }}
+                />
+              </div>
+            ) : (
+              <ToolbarButton
+                icon={faPaintbrush}
+                buttonProps={buttonProps.PAINT}
+                tooltip="Paint Brush"
+              />
+            )}
+
+            { /* Conditionally show eraser menu on mode selection */}
+            {buttonProps.ERASER.active ? (
+              <EraseModeMenu
+                faIcon={faEraser}
+                borderStyle="border-2  bg-primary/30 border-2 border-primary hover:bg-primary/30 text-white"
+                defaultOpen={true}
+                anchor="right"
+                anchorGap={12}
+                closeOnMouseLeave={true}
+              />
+            )
+              : (
+                <ToolbarButton
+                  icon={faEraser}
+                  buttonProps={buttonProps.ERASER}
+                  tooltip="Eraser"
+                />
+              )}
+          </div>
+        </div>
+
+        <hr className="w-full border-t border-white/15" />
+
+        <div className="flex flex-col items-center gap-2">
           <ToolbarButton
             icon={faArrowRotateLeft}
             buttonProps={buttonProps.UNDO}
@@ -128,15 +220,11 @@ 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,6 +7,8 @@ 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/ToolbarNode/ToolbarNode.tsx

@@ -13,9 +13,6 @@ import {
 import { paperWrapperStyles, toolTipStyles } from "~/components/styles";
 import { ToolbarNodeButtonNames } from "./enums";
 import { ToolbarNodeButtonData } from "./data";
-import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
-import { Fragment } from "react/jsx-runtime";
-import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { ColorPicker } from "~/components/ui/TextEditor/ColorPicker";
 
 export interface ToolbarNodeProps {
@@ -54,7 +51,7 @@ export const ToolbarNode = ({
       className={twMerge(
         paperWrapperStyles,
         disabled && "pointer-events-none cursor-default bg-ui-border shadow-md",
-        "flex gap-2 transition",
+        "flex flex-col gap-2 transition",
       )}
     >
       <ToolbarButton

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

@@ -1,28 +1,27 @@
 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/Storyteller-Logo-Black.png"
+        src="/brand/mira-logo.png"
         alt="logo"
-        className="mr-4 h-8"
+        className="h-[30px] select-none pb-1"
       />
-      <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

@@ -0,0 +1,31 @@
+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,26 +1,35 @@
 import { twMerge } from "tailwind-merge";
 
 export const paperWrapperStyles =
-  "rounded-xl border border-ui-border bg-ui-panel p-2 shadow-lg";
+  "rounded-xl border border-ui-border p-1.5 shadow-lg bg-ui-panel";
 
 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 = 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 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 dialogPanelStyles = twMerge(
-  "flex max-w-5xl flex-col justify-between gap-4 p-6",
+  "flex max-w-5xl flex-col justify-between gap-4 p-6 bg-[#1C1C20]",
 );

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-gray-600 hover:bg-gray-500 text-white focus-visible:outline-secondary";
+        return "bg-ui-controls hover:bg-ui-controls/80 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-400 text-white focus-visible:outline-primary-600";
+        return "bg-primary hover:bg-primary-600 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">
+      <div className="relative flex h-10 rounded-md border border-ui-border">
         <ComboboxInput
-          className="rounded-md"
+          className="rounded-lg px-3"
           displayValue={(font: string) => font}
           onChange={(event) => setQuery(event.target.value)}
         />
-        <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 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>
       </div>
       <ComboboxOptions
         anchor="bottom"
-        className="mt-1 w-[var(--input-width)] rounded-lg border empty:invisible"
+        className="mt-1 w-[var(--input-width)] rounded-lg border border-ui-border empty:invisible"
       >
         {filteredOptions.map((option, idx) => (
           <ComboboxOption
             key={idx}
             value={option}
-            className="cursor-pointer bg-white p-2 data-[focus]:bg-primary-200 data-[selected]:font-medium"
+            className="bg-ui-controls cursor-pointer px-3 py-2 data-[focus]:bg-gray-600 data-[selected]:font-medium"
           >
             {option}
           </ComboboxOption>

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

@@ -1,6 +1,8 @@
 import { Transition } from "@headlessui/react";
 import React, { useState, useRef, useEffect, ReactNode } from "react";
 import { twMerge } from "tailwind-merge";
+import { Button } from "./Button";
+import { faMinus, faPlus } from "@fortawesome/pro-solid-svg-icons";
 
 interface SliderProps {
   min: number;
@@ -12,6 +14,9 @@ interface SliderProps {
   showTooltip?: boolean;
   tooltipContent?: ReactNode;
   suffix?: string;
+  innerLabel?: string;
+  showDecrement?: boolean,
+  showIncrement?: boolean,
 }
 
 export const Slider = ({
@@ -24,6 +29,9 @@ export const Slider = ({
   showTooltip,
   tooltipContent,
   suffix,
+  innerLabel,
+  showDecrement,
+  showIncrement
 }: SliderProps) => {
   const [localValue, setLocalValue] = useState(value);
   const sliderRef = useRef<HTMLDivElement>(null);
@@ -34,15 +42,19 @@ export const Slider = ({
     setLocalValue(value);
   }, [value]);
 
+  const setClampedValue = (value: number) => {
+    const roundedValue = Math.round(value / step) * step;
+    const clampedValue = Math.min(Math.max(roundedValue, min), max);
+    setLocalValue(clampedValue);
+    onChange(clampedValue);
+  }
+
   const updatePosition = (x: number) => {
     if (!sliderRef.current) return;
     const rect = sliderRef.current.getBoundingClientRect();
     const fraction = (x - rect.left) / rect.width;
     const newValue = fraction * (max - min) + min;
-    const roundedValue = Math.round(newValue / step) * step;
-    const clampedValue = Math.min(Math.max(roundedValue, min), max);
-    setLocalValue(clampedValue);
-    onChange(clampedValue);
+    setClampedValue(newValue);
   };
 
   const handleMouseDown = (e: React.MouseEvent<HTMLDivElement>) => {
@@ -79,34 +91,75 @@ export const Slider = ({
     }
   };
 
+  const handleDecrement = () => {
+    setClampedValue(localValue - step);
+  }
+
+  const handleIncrement = () => {
+    setClampedValue(localValue + step);
+  }
+
   const percentage = ((localValue - min) / (max - min)) * 100;
   const displayValue =
     localValue.toFixed(step.toString().split(".")[1]?.length || 0) +
     (suffix || "");
 
   return (
-    <div
-      ref={sliderRef}
-      className={twMerge(
-        "relative h-2.5 w-full cursor-pointer rounded-full bg-gray-500",
-        className,
+    <div className="flex w-full">
+      {showDecrement && (
+        <Button
+          icon={faMinus}
+          className="size-6 text-white/80 bg-transparent rounded-full hover:bg-white/10 focus-visible:outline-primary active:bg-primary/30 my-auto mr-1"
+          onClick={handleDecrement}
+        />
       )}
-      onMouseDown={handleMouseDown}
-    >
-      <div
-        className="absolute h-2.5 rounded-full bg-primary"
-        style={{ width: `${percentage}%` }}
-      ></div>
+
       <div
+        ref={sliderRef}
         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",
+          "glass group relative h-7 w-full cursor-pointer overflow-hidden rounded-lg border border-ui-border",
+          isDragging && "!bg-ui-controls/90",
+          className,
         )}
-        style={{ left: `${percentage}%` }}
         onMouseDown={handleMouseDown}
-        onMouseEnter={handleMouseEnter}
-        onMouseLeave={handleMouseLeave}
-      />
+      >
+        <div
+          className={twMerge(
+            "absolute h-7 bg-primary/30 transition-colors duration-300 group-hover:bg-primary/50",
+            isDragging && "!bg-primary/50",
+          )}
+          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
+            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>
+      </div>
+
+      {showIncrement && (
+        <Button
+          icon={faPlus}
+          className="size-6 text-white/80 bg-transparent rounded-full hover:bg-white/10 focus-visible:outline-primary active:bg-primary/30 my-auto ml-1"
+          onClick={handleIncrement}
+        />
+      )}
+
       {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,7 +27,6 @@ export const Textarea = forwardRef<HTMLTextAreaElement, TextareaInterface>(
     { className, label, resize = "vertical", id, ...rest }: TextareaInterface,
     ref,
   ) => {
-
     const resizeStyle = resizeStyles[resize];
 
     return (
@@ -38,17 +37,16 @@ export const Textarea = forwardRef<HTMLTextAreaElement, TextareaInterface>(
           ref={ref}
           id={id ? id : label ? kebabCase(label) : undefined}
           className={twMerge(
-            "border-ui-panel-border bg-ui-controls rounded-lg border px-3 py-2",
+            "rounded-lg border border-ui-border px-4 py-2.5 text-white outline-none focus:outline-none",
             className,
             resizeStyle,
           )}
           style={{
             outline: "2px solid transparent",
-            transition: "outline-color 0.15s ease-in-out",
           }}
           onFocus={(e) => {
             // disableHotkeyInput(DomLevels.INPUT);
-            e.currentTarget.style.outlineColor = "#e66462";
+            e.currentTarget.style.outlineColor = "#2d81ff";
           }}
           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,27 +1,46 @@
-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";
+import { ChangeEvent, useEffect, useState } from "react";
+import { HexAlphaColorPicker } from "react-colorful";
+import { twMerge } from "tailwind-merge";
 
+import { paperWrapperStyles } from "~/components/styles";
 import { Button } from "../Button";
 import { Input } from "../Input";
-import { paperWrapperStyles } from "~/components/styles";
 
 export const ColorPicker = ({
   color: prevColor,
   onChange,
   faIcon = faFont,
-  borderStyle = "border",
+  borderStyle,
   showBar = true,
+  anchor = "bottom",
+  anchorGap = 0,
+  fillBg = false,
+  staticIconColor,
+  streamChanges = false,
+  defaultOpen = false,
+  closeOnMouseLeave = false,
 }: {
   color: string;
   onChange: (newColor: string) => void;
   faIcon?: IconDefinition;
   borderStyle?: string;
   showBar?: boolean;
+  anchor?: "top" | "right" | "bottom" | "left";
+  anchorGap?: number;
+  fillBg?: boolean;
+  staticIconColor?: string;
+  streamChanges?: boolean;
+  defaultOpen?: boolean;
+  closeOnMouseLeave?: boolean;
 }) => {
+  const [isOpen, setIsOpen] = useState(defaultOpen);
   const [{ currColor, textInput }, setStates] = useState<{
     currColor: string;
     textInput: string;
@@ -30,13 +49,18 @@ export const ColorPicker = ({
     textInput: prevColor.substring(1),
   });
 
-  const handleSelectColor = (newColor: string) => {
-    setStates({ currColor: newColor, textInput: newColor.substring(1) });
-  };
-  const isHexColorCode = (color: string) => {
-    const hexColorRegex = /^(?:[0-9a-fA-F]{3}){1,2}$/;
-    return hexColorRegex.test(color);
+  useEffect(() => {
+    setIsOpen(defaultOpen);
+  }, [defaultOpen]);
+
+  const handleMouseLeave = (close: () => void) => {
+    if (closeOnMouseLeave) {
+      onChange(currColor);
+      setIsOpen(false);
+      close();
+    }
   };
+
   const handleTextInput = (e: ChangeEvent<HTMLInputElement>) => {
     if (isHexColorCode(e.target.value)) {
       setStates({
@@ -51,58 +75,100 @@ export const ColorPicker = ({
     }
   };
 
+  const isHexColorCode = (color: string) => {
+    const hexColorRegex = /^(?:[0-9a-fA-F]{3}){1,2}$/;
+    return hexColorRegex.test(color);
+  };
+
   return (
     <Popover className="relative">
-      <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 }} />}
-      </PopoverButton>
-      <PopoverPanel
-        anchor="bottom"
-        className={twMerge(
-          paperWrapperStyles,
-          "flex flex-col items-center gap-2",
-        )}
-      >
-        {({ close }) => {
-          return (
-            <>
-              <HexColorPicker color={currColor} onChange={handleSelectColor} />
-              <Input
-                style={{ width: "198px" }}
-                icon={faHashtag}
-                value={textInput}
-                onChange={handleTextInput}
+      {({ open, close }) => (
+        <>
+          <PopoverButton
+            className={twMerge(
+              "flex size-10 items-center gap-1 rounded-lg bg-ui-controls p-2",
+              borderStyle,
+              showBar ? "flex-col" : "justify-center",
+            )}
+            style={fillBg ? { backgroundColor: prevColor } : {}}
+            onMouseEnter={() => defaultOpen && setIsOpen(true)}
+          >
+            <FontAwesomeIcon
+              icon={faIcon}
+              color={staticIconColor ?? (fillBg ? undefined : prevColor)}
+            />
+            {showBar && (
+              <span
+                className="h-1 w-full"
+                style={{ backgroundColor: prevColor }}
               />
-              <div className="flex w-full justify-center gap-2">
-                <Button
-                  variant="secondary"
-                  onClick={() => {
+            )}
+          </PopoverButton>
+          {(open || isOpen) && (
+            <div onMouseLeave={() => handleMouseLeave(close)}>
+              <PopoverPanel
+                anchor={anchor}
+                className={twMerge(
+                  paperWrapperStyles,
+                  "flex flex-col items-center gap-2 overflow-hidden",
+                )}
+                style={
+                  {
+                    "--anchor-gap": `${anchorGap}px`,
+                  } as React.CSSProperties
+                }
+                static
+              >
+                <HexAlphaColorPicker
+                  className="overflow-hidden"
+                  color={currColor}
+                  onChange={(color) => {
                     setStates({
-                      currColor: prevColor,
-                      textInput: prevColor.substring(1),
+                      currColor: color,
+                      textInput: color.substring(1),
                     });
-                    close();
+                    if (streamChanges) {
+                      onChange(color);
+                    }
                   }}
-                >
-                  Cancel
-                </Button>
-                <Button
-                  onClick={() => {
-                    onChange(currColor);
-                    close();
-                  }}
-                >
-                  OK
-                </Button>
-              </div>
-            </>
-          );
-        }}
-      </PopoverPanel>
+                />
+                <Input
+                  style={{ width: "198px" }}
+                  icon={faHashtag}
+                  value={textInput}
+                  onChange={handleTextInput}
+                />
+                {!closeOnMouseLeave && (
+                  <div className="flex w-full justify-center gap-2">
+                    <Button
+                      variant="secondary"
+                      onClick={() => {
+                        setStates({
+                          currColor: prevColor,
+                          textInput: prevColor.substring(1),
+                        });
+                        setIsOpen(false);
+                        close();
+                      }}
+                    >
+                      Cancel
+                    </Button>
+                    <Button
+                      onClick={() => {
+                        onChange(currColor);
+                        setIsOpen(false);
+                        close();
+                      }}
+                    >
+                      OK
+                    </Button>
+                  </div>
+                )}
+              </PopoverPanel>
+            </div>
+          )}
+        </>
+      )}
     </Popover>
   );
 };

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

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

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

@@ -1,4 +1,4 @@
 export const buttonGroupButton =
-  "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";
+  "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";
 export const buttonGroupGrouper =
-  "flex divide-x overflow-hidden rounded-md border";
+  "flex divide-x overflow-hidden rounded-md border border-ui-border";

frontend/realtime/src/components/ui/ToolbarMain/EraseModeMenu.tsx

@@ -0,0 +1,102 @@
+import {
+  faFont,
+  faHashtag,
+  IconDefinition,
+} from "@fortawesome/pro-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
+import { ChangeEvent, useEffect, useState } from "react";
+import { HexColorPicker } from "react-colorful";
+import { twMerge } from "tailwind-merge";
+
+import { useSignals } from "@preact/signals-react/runtime";
+import { paperWrapperStyles } from "~/components/styles";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, paintBrushSize, setPaintBrushSize } from "~/signals/uiEvents/toolbarMain/paintMode";
+import { Button } from "../Button";
+import { Input } from "../Input";
+import { Slider } from "../Slider";
+import { eraseBrushSize, setEraseBrushSize } from "~/signals/uiEvents/toolbarMain/eraseMode";
+
+export const EraseModeMenu = ({
+  faIcon = faFont,
+  borderStyle,
+  anchor = "bottom",
+  anchorGap = 0,
+  defaultOpen = false,
+  closeOnMouseLeave = false,
+}: {
+  faIcon?: IconDefinition;
+  borderStyle?: string;
+  anchor?: "top" | "right" | "bottom" | "left";
+  anchorGap?: number;
+  defaultOpen?: boolean;
+  closeOnMouseLeave?: boolean;
+}) => {
+  const [isOpen, setIsOpen] = useState(defaultOpen);
+
+  useEffect(() => {
+    setIsOpen(defaultOpen);
+  }, [defaultOpen]);
+
+  const handleMouseLeave = (close: () => void) => {
+    if (closeOnMouseLeave) {
+      setIsOpen(false);
+      close();
+    }
+  };
+
+  // Rerender the component when these signals change:
+  useSignals();
+  const brushSize = eraseBrushSize.value;
+
+  return (
+    <Popover className="relative">
+      {({ open, close }) => (
+        <>
+          <PopoverButton
+            className={twMerge(
+              "flex size-10 items-center justify-center gap-1 rounded-lg bg-ui-controls p-2",
+              borderStyle,
+            )}
+            onMouseEnter={() => defaultOpen && setIsOpen(true)}
+          >
+            <FontAwesomeIcon
+              icon={faIcon}
+            />
+          </PopoverButton>
+          {(open || isOpen) && (
+            <div onMouseLeave={() => handleMouseLeave(close)}>
+              <PopoverPanel
+                anchor={anchor}
+                className={twMerge(
+                  paperWrapperStyles,
+                  "flex flex-col items-center gap-2 overflow-hidden w-[198px]",
+                )}
+                style={
+                  {
+                    "--anchor-gap": `${anchorGap}px`,
+                  } as React.CSSProperties
+                }
+                static
+              >
+                <div className="flex flex-col w-full gap-2">
+                  <p className="text-white font-medium text-sm w-full justify-start">Eraser Size:</p>
+                  <Slider
+                    min={BRUSH_MIN_SIZE}
+                    max={BRUSH_MAX_SIZE}
+                    value={brushSize}
+                    onChange={setEraseBrushSize}
+                    step={1}
+                    innerLabel={brushSize.toString()}
+                    showDecrement
+                    showIncrement
+                  />
+                </div>
+              </PopoverPanel>
+            </div>
+          )}
+        </>
+      )}
+    </Popover>
+  );
+};

frontend/realtime/src/components/ui/ToolbarMain/PaintModeMenu.tsx

@@ -0,0 +1,194 @@
+import {
+  faFont,
+  faHashtag,
+  IconDefinition,
+} from "@fortawesome/pro-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { Popover, PopoverButton, PopoverPanel } from "@headlessui/react";
+import { ChangeEvent, useEffect, useState } from "react";
+import { HexAlphaColorPicker } from "react-colorful";
+import { twMerge } from "tailwind-merge";
+
+import { useSignals } from "@preact/signals-react/runtime";
+import { paperWrapperStyles } from "~/components/styles";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, paintBrushSize, setPaintBrushSize } from "~/signals/uiEvents/toolbarMain/paintMode";
+import { Button } from "../Button";
+import { Input } from "../Input";
+import { Slider } from "../Slider";
+
+export const PaintModeMenu = ({
+  color: prevColor,
+  onChange,
+  faIcon = faFont,
+  borderStyle,
+  showBar = true,
+  anchor = "bottom",
+  anchorGap = 0,
+  fillBg = false,
+  staticIconColor,
+  streamChanges = false,
+  defaultOpen = false,
+  closeOnMouseLeave = false,
+}: {
+  color: string;
+  onChange: (newColor: string) => void;
+  faIcon?: IconDefinition;
+  borderStyle?: string;
+  showBar?: boolean;
+  anchor?: "top" | "right" | "bottom" | "left";
+  anchorGap?: number;
+  fillBg?: boolean;
+  staticIconColor?: string;
+  streamChanges?: boolean;
+  defaultOpen?: boolean;
+  closeOnMouseLeave?: boolean;
+}) => {
+  const [isOpen, setIsOpen] = useState(defaultOpen);
+  const [{ currColor, textInput }, setStates] = useState<{
+    currColor: string;
+    textInput: string;
+  }>({
+    currColor: prevColor,
+    textInput: prevColor.substring(1),
+  });
+
+  useEffect(() => {
+    setIsOpen(defaultOpen);
+  }, [defaultOpen]);
+
+  const handleMouseLeave = (close: () => void) => {
+    if (closeOnMouseLeave) {
+      onChange(currColor);
+      setIsOpen(false);
+      close();
+    }
+  };
+
+  const handleTextInput = (e: ChangeEvent<HTMLInputElement>) => {
+    if (isHexColorCode(e.target.value)) {
+      setStates({
+        currColor: `#${e.target.value}`,
+        textInput: e.target.value,
+      });
+    } else {
+      setStates((curr) => ({
+        ...curr,
+        textInput: e.target.value,
+      }));
+    }
+  };
+
+  const isHexColorCode = (color: string) => {
+    const hexColorRegex = /^(?:[0-9a-fA-F]{3}){1,2}$/;
+    return hexColorRegex.test(color);
+  };
+
+  // Rerender the component when these signals change:
+  useSignals();
+  const brushSize = paintBrushSize.value;
+
+  return (
+    <Popover className="relative">
+      {({ open, close }) => (
+        <>
+          <PopoverButton
+            className={twMerge(
+              "flex size-10 items-center gap-1 rounded-lg bg-ui-controls p-2",
+              borderStyle,
+              showBar ? "flex-col" : "justify-center",
+            )}
+            style={fillBg ? { backgroundColor: prevColor } : {}}
+            onMouseEnter={() => defaultOpen && setIsOpen(true)}
+          >
+            <FontAwesomeIcon
+              icon={faIcon}
+              color={staticIconColor ?? (fillBg ? undefined : prevColor)}
+            />
+            {showBar && (
+              <span
+                className="h-1 w-full"
+                style={{ backgroundColor: prevColor }}
+              />
+            )}
+          </PopoverButton>
+          {(open || isOpen) && (
+            <div onMouseLeave={() => handleMouseLeave(close)}>
+              <PopoverPanel
+                anchor={anchor}
+                className={twMerge(
+                  paperWrapperStyles,
+                  "flex flex-col items-center gap-2 overflow-hidden",
+                )}
+                style={
+                  {
+                    "--anchor-gap": `${anchorGap}px`,
+                  } as React.CSSProperties
+                }
+                static
+              >
+                <HexAlphaColorPicker
+                  className="overflow-hidden"
+                  color={currColor}
+                  onChange={(color) => {
+                    setStates({
+                      currColor: color,
+                      textInput: color.substring(1),
+                    });
+                    if (streamChanges) {
+                      onChange(color);
+                    }
+                  }}
+                />
+                <Input
+                  style={{ width: "198px" }}
+                  icon={faHashtag}
+                  value={textInput}
+                  onChange={handleTextInput}
+                />
+                <div className="flex flex-col w-full mt-2 gap-2">
+                  <p className="text-white font-medium text-sm w-full justify-start">Brush Size:</p>
+                  <Slider
+                    min={BRUSH_MIN_SIZE}
+                    max={BRUSH_MAX_SIZE}
+                    value={brushSize}
+                    onChange={setPaintBrushSize}
+                    step={1}
+                    innerLabel={brushSize.toString()}
+                    showDecrement
+                    showIncrement
+                  />
+                </div>
+                {!streamChanges && (
+                  <div className="flex w-full justify-center gap-2">
+                    <Button
+                      variant="secondary"
+                      onClick={() => {
+                        setStates({
+                          currColor: prevColor,
+                          textInput: prevColor.substring(1),
+                        });
+                        setIsOpen(false);
+                        close();
+                      }}
+                    >
+                      Cancel
+                    </Button>
+                    <Button
+                      onClick={() => {
+                        onChange(currColor);
+                        setIsOpen(false);
+                        close();
+                      }}
+                    >
+                      OK
+                    </Button>
+                  </div>
+                )}
+              </PopoverPanel>
+            </div>
+          )}
+        </>
+      )}
+    </Popover>
+  );
+};

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

@@ -7,15 +7,18 @@ 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,
+      toolTipStyles.base,
+      toolTipStyles[position],
       forceShow && "after:block before:block",
       children.props.className,
     ),

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

@@ -1,5 +1,6 @@
 import { signal } from "@preact/signals-react";
 import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum";
+import { isToolbarMode } from "~/signals/uiEvents/toolbarMain/toolbarMain";
 
 interface ToolbarMainSignalInterface {
   disabled: boolean;
@@ -42,6 +43,17 @@ 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/eraseMode.ts

@@ -0,0 +1,14 @@
+import { effect, signal } from "@preact/signals-react";
+import { clamp, normalize } from "~/utilities";
+import { BRUSH_MAX_SIZE, BRUSH_MIN_SIZE, DEFAULT_PAINT_BRUSH_SIZE, paintBrushSize } from "./paintMode";
+
+export const eraseBrushSize = signal<number>(DEFAULT_PAINT_BRUSH_SIZE);
+export const setEraseBrushSize = (brushSize: number) => {
+  eraseBrushSize.value = clamp(brushSize, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE);
+}
+
+export const onEraseBrushSizeChanged = (callback: (normalizedSize: number) => void) => {
+  effect(() => {
+    callback(normalize(paintBrushSize.value, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE));
+  });
+};

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

@@ -3,11 +3,18 @@ import {
   loadingBarRetryDispatch,
   loadingBarRetryEventHandler,
 } from "./loadingBar";
+import { onEraseBrushSizeChanged, setEraseBrushSize } from "./eraseMode";
+import { onPaintColorChanged, setPaintBrushSize, setPaintColor } from "./paintMode";
 export const dispatchers = {
   ...buttonDispatchers,
   loadingBarRetry: loadingBarRetryDispatch,
+  setPaintColor,
+  setPaintBrushSize,
+  setEraseBrushSize
 };
 export const eventsHandlers = {
   ...buttonEventsHandlers,
   loadingBarRetry: loadingBarRetryEventHandler,
+  onPaintColorChanged,
+  onEraseBrushSizeChanged
 };

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

@@ -0,0 +1,32 @@
+import { effect, signal } from "@preact/signals-react";
+import { clamp, normalize } from "~/utilities";
+
+export const DEFAULT_PAINT_COLOR: string = "#ff0000";
+export const paintColor = signal<string>(DEFAULT_PAINT_COLOR);
+export const setPaintColor = (data: string) => {
+  paintColor.value = data;
+};
+
+export const BRUSH_MIN_SIZE: number = 0;
+export const BRUSH_MAX_SIZE: number = 100;
+// NOTE: Must be between MIN and MAX sizes - there's no runtime checks
+export const DEFAULT_PAINT_BRUSH_SIZE: number = 50;
+export const paintBrushSize = signal<number>(DEFAULT_PAINT_BRUSH_SIZE);
+export const setPaintBrushSize = (brushSize: number) => {
+  paintBrushSize.value = clamp(brushSize, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE);
+  console.debug(`Paint brush size set to ${paintBrushSize.value}`);
+}
+
+export const onPaintColorChanged = (callback: (data: string) => void) => {
+  effect(() => {
+    if (paintColor.value) {
+      callback(paintColor.value);
+    }
+  });
+};
+
+export const onPaintBrushSizeChanged = (callback: (normalizedSize: number) => void) => {
+  effect(() => {
+    callback(normalize(paintBrushSize.value, BRUSH_MIN_SIZE, BRUSH_MAX_SIZE));
+  });
+};

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

@@ -2,6 +2,7 @@ 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) => {
@@ -25,22 +26,35 @@ 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>) => {
-        if (effectsCleanups[buttonName]) {
-          effectsCleanups[buttonName]();
-        }
-        effectsCleanups[buttonName] = effect(() => {
-          if (events[buttonName].value) {
-            callback(events[buttonName].value);
-            events[buttonName].value = undefined;
-          }
-        });
-      },
+      onClick: (callback: MouseEventHandler<HTMLButtonElement>) => { defaultHandler(buttonName, callback) },
     };
     return acc;
   },

frontend/realtime/src/types/ServerTypes.ts

@@ -0,0 +1,36 @@
+// Server communication types
+export interface ModelConfig {
+  name: string;
+  path: string;
+  precision: "fp16" | "fp32";
+}
+
+export interface LoraConfig {
+  path: string;
+  alpha: number;
+}
+
+export interface ServerSetupPayload {
+  type: "setup";
+  model: ModelConfig;
+  lora?: LoraConfig;
+  device: "cuda" | "cpu";
+}
+
+export interface ServerSettingsPayload {
+  type: "settings";
+  prompt?: string;
+  strength?: number;
+  steps?: number;
+  guidance_scale?: number;
+  seed?: number;
+  negative_prompt?: string;
+}
+
+export interface ServerResponse {
+  type: "progress" | "success" | "error";
+  status?: "model_loaded" | "generation_complete";
+  message?: string;
+  percent?: number;
+  image?: string; // base64 encoded image
+}

frontend/realtime/src/utilities/index.ts

@@ -10,3 +10,9 @@ export const getFileName = (file: File) => {
 export const getFileExtension = (file: File) => {
   return file.name.substring(file.name.lastIndexOf(".") + 1);
 };
+export const clamp = (num: number, min: number, max: number) => {
+  return Math.min(Math.max(num, min), max);
+}
+export const normalize = (num: number, min: number, max: number) => {
+  return (num - min) / (max - min);
+}

frontend/realtime/start_server.ps1

@@ -0,0 +1,5 @@
+# Activate virtual environment
+deactivate conda
+.\.venv\Scripts\activate.bat
+cd .\paint_py_server
+python server.py

frontend/realtime/tailwind.config.js

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

frontend/realtime/tailwind.stcolors.js

@@ -1,19 +1,19 @@
 export const storytellerColors = {
-  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",
+  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",
   },
+
   gunpowder: {
     DEFAULT: "#39394c", //=950
     50: "#f5f6f9",