implement cumprod

more work on timestamps (WIP)
cleanup
2026-11-01 07:15:54 -05:00 · 2025-02-25 00:25:19 -05:00 · 2025-02-24 19:36:31 -05:00 · 2025-02-24 17:48:24 -05:00 · 2025-02-24 17:45:08 -05:00 · 2025-02-24 17:22:25 -05:00
45 changed files with 1768 additions and 275 deletions
@@ -609,6 +609,7 @@ dependencies = [
 "candle-core",
 "candle-nn",
 "candle-transformers",
+ "cudarc",
 "filesys",
 "hf-hub",
 "image",
@@ -1424,14 +1425,14 @@ dependencies = [

 [[package]]
 name = "candle-core"
-version = "0.8.2"
+version = "0.8.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "855dfedff437d2681d68e1f34ae559d88b0dd84aa5a6b63f2c8e75ebdd875bbf"
+checksum = "bef79e00300acaaa6b0b33168e2ed5050f06125d1add87e7103c40a119083581"
 dependencies = [
 "byteorder",
 "candle-kernels",
 "cudarc",
- "gemm",
+ "gemm 0.17.1",
 "half",
 "memmap2",
 "num-traits",
@@ -1449,9 +1450,9 @@ dependencies = [

 [[package]]
 name = "candle-kernels"
-version = "0.8.2"
+version = "0.8.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "53343628fa470b7075c28c589b98735b4220b464e37ddbb8e117040e199f4787"
+checksum = "e3134cc19a080a326f95bdcbe057ecaa9658598e5146e37ba1a8cd550c342a9b"
 dependencies = [
 "bindgen_cuda",
 ]
@@ -2228,9 +2229,9 @@ dependencies = [

 [[package]]
 name = "cudarc"
-version = "0.12.1"
+version = "0.13.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38cd60a9a42ec83a2ed7effb0b1f073270264ea99da7acfc44f7e8d74dee0384"
+checksum = "4e29ce3bfa797c1183053ceb496316203ef561c183941c3c181500d9ade6daf4"
 dependencies = [
 "half",
 "libloading 0.8.6",
@@ -2891,6 +2892,15 @@ dependencies = [
 "reborrow",
 ]

+[[package]]
+name = "dyn-stack"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "490bd48eb68fffcfed519b4edbfd82c69cbe741d175b84f0e0cbe8c57cbe0bdd"
+dependencies = [
+ "bytemuck",
+]
+
 [[package]]
 name = "easyenv"
 version = "0.1.0"
@@ -3743,17 +3753,37 @@ version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6ab24cc62135b40090e31a76a9b2766a501979f3070fa27f689c27ec04377d32"
 dependencies = [
- "dyn-stack",
- "gemm-c32",
- "gemm-c64",
- "gemm-common",
- "gemm-f16",
- "gemm-f32",
- "gemm-f64",
+ "dyn-stack 0.10.0",
+ "gemm-c32 0.17.1",
+ "gemm-c64 0.17.1",
+ "gemm-common 0.17.1",
+ "gemm-f16 0.17.1",
+ "gemm-f32 0.17.1",
+ "gemm-f64 0.17.1",
 "num-complex",
 "num-traits",
 "paste",
- "raw-cpuid",
+ "raw-cpuid 10.7.0",
+ "seq-macro",
+]
+
+[[package]]
+name = "gemm"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab96b703d31950f1aeddded248bc95543c9efc7ac9c4a21fda8703a83ee35451"
+dependencies = [
+ "dyn-stack 0.13.0",
+ "gemm-c32 0.18.2",
+ "gemm-c64 0.18.2",
+ "gemm-common 0.18.2",
+ "gemm-f16 0.18.2",
+ "gemm-f32 0.18.2",
+ "gemm-f64 0.18.2",
+ "num-complex",
+ "num-traits",
+ "paste",
+ "raw-cpuid 11.4.0",
 "seq-macro",
 ]

@@ -3763,12 +3793,27 @@ version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b9c030d0b983d1e34a546b86e08f600c11696fde16199f971cd46c12e67512c0"
 dependencies = [
- "dyn-stack",
- "gemm-common",
+ "dyn-stack 0.10.0",
+ "gemm-common 0.17.1",
 "num-complex",
 "num-traits",
 "paste",
- "raw-cpuid",
+ "raw-cpuid 10.7.0",
+ "seq-macro",
+]
+
+[[package]]
+name = "gemm-c32"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6db9fd9f40421d00eea9dd0770045a5603b8d684654816637732463f4073847"
+dependencies = [
+ "dyn-stack 0.13.0",
+ "gemm-common 0.18.2",
+ "num-complex",
+ "num-traits",
+ "paste",
+ "raw-cpuid 11.4.0",
 "seq-macro",
 ]

@@ -3778,12 +3823,27 @@ version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fbb5f2e79fefb9693d18e1066a557b4546cd334b226beadc68b11a8f9431852a"
 dependencies = [
- "dyn-stack",
- "gemm-common",
+ "dyn-stack 0.10.0",
+ "gemm-common 0.17.1",
 "num-complex",
 "num-traits",
 "paste",
- "raw-cpuid",
+ "raw-cpuid 10.7.0",
+ "seq-macro",
+]
+
+[[package]]
+name = "gemm-c64"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dfcad8a3d35a43758330b635d02edad980c1e143dc2f21e6fd25f9e4eada8edf"
+dependencies = [
+ "dyn-stack 0.13.0",
+ "gemm-common 0.18.2",
+ "num-complex",
+ "num-traits",
+ "paste",
+ "raw-cpuid 11.4.0",
 "seq-macro",
 ]

@@ -3794,17 +3854,38 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a2e7ea062c987abcd8db95db917b4ffb4ecdfd0668471d8dc54734fdff2354e8"
 dependencies = [
 "bytemuck",
- "dyn-stack",
+ "dyn-stack 0.10.0",
 "half",
 "num-complex",
 "num-traits",
 "once_cell",
 "paste",
- "pulp",
- "raw-cpuid",
+ "pulp 0.18.22",
+ "raw-cpuid 10.7.0",
 "rayon",
 "seq-macro",
- "sysctl",
+ "sysctl 0.5.5",
+]
+
+[[package]]
+name = "gemm-common"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a352d4a69cbe938b9e2a9cb7a3a63b7e72f9349174a2752a558a8a563510d0f3"
+dependencies = [
+ "bytemuck",
+ "dyn-stack 0.13.0",
+ "half",
+ "libm",
+ "num-complex",
+ "num-traits",
+ "once_cell",
+ "paste",
+ "pulp 0.21.4",
+ "raw-cpuid 11.4.0",
+ "rayon",
+ "seq-macro",
+ "sysctl 0.6.0",
 ]

 [[package]]
@@ -3813,14 +3894,32 @@ version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7ca4c06b9b11952071d317604acb332e924e817bd891bec8dfb494168c7cedd4"
 dependencies = [
- "dyn-stack",
- "gemm-common",
- "gemm-f32",
+ "dyn-stack 0.10.0",
+ "gemm-common 0.17.1",
+ "gemm-f32 0.17.1",
 "half",
 "num-complex",
 "num-traits",
 "paste",
- "raw-cpuid",
+ "raw-cpuid 10.7.0",
+ "rayon",
+ "seq-macro",
+]
+
+[[package]]
+name = "gemm-f16"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cff95ae3259432f3c3410eaa919033cd03791d81cebd18018393dc147952e109"
+dependencies = [
+ "dyn-stack 0.13.0",
+ "gemm-common 0.18.2",
+ "gemm-f32 0.18.2",
+ "half",
+ "num-complex",
+ "num-traits",
+ "paste",
+ "raw-cpuid 11.4.0",
 "rayon",
 "seq-macro",
 ]
@@ -3831,12 +3930,27 @@ version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e9a69f51aaefbd9cf12d18faf273d3e982d9d711f60775645ed5c8047b4ae113"
 dependencies = [
- "dyn-stack",
- "gemm-common",
+ "dyn-stack 0.10.0",
+ "gemm-common 0.17.1",
 "num-complex",
 "num-traits",
 "paste",
- "raw-cpuid",
+ "raw-cpuid 10.7.0",
+ "seq-macro",
+]
+
+[[package]]
+name = "gemm-f32"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc8d3d4385393304f407392f754cd2dc4b315d05063f62cf09f47b58de276864"
+dependencies = [
+ "dyn-stack 0.13.0",
+ "gemm-common 0.18.2",
+ "num-complex",
+ "num-traits",
+ "paste",
+ "raw-cpuid 11.4.0",
 "seq-macro",
 ]

@@ -3846,12 +3960,27 @@ version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "aa397a48544fadf0b81ec8741e5c0fba0043008113f71f2034def1935645d2b0"
 dependencies = [
- "dyn-stack",
- "gemm-common",
+ "dyn-stack 0.10.0",
+ "gemm-common 0.17.1",
 "num-complex",
 "num-traits",
 "paste",
- "raw-cpuid",
+ "raw-cpuid 10.7.0",
+ "seq-macro",
+]
+
+[[package]]
+name = "gemm-f64"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35b2a4f76ce4b8b16eadc11ccf2e083252d8237c1b589558a49b0183545015bd"
+dependencies = [
+ "dyn-stack 0.13.0",
+ "gemm-common 0.18.2",
+ "num-complex",
+ "num-traits",
+ "paste",
+ "raw-cpuid 11.4.0",
 "seq-macro",
 ]

@@ -5237,9 +5366,9 @@ dependencies = [

 [[package]]
 name = "libm"
-version = "0.2.8"
+version = "0.2.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
+checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa"

 [[package]]
 name = "libredox"
@@ -6238,7 +6367,7 @@ version = "0.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "681030a937600a36906c185595136d26abfebb4aa9c65701cefcaf8578bb982b"
 dependencies = [
- "proc-macro-crate 3.1.0",
+ "proc-macro-crate 1.3.1",
 "proc-macro2",
 "quote",
 "syn 2.0.87",
@@ -7323,6 +7452,20 @@ dependencies = [
 "reborrow",
 ]

+[[package]]
+name = "pulp"
+version = "0.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95fb7a99b37aaef4c7dd2fd15a819eb8010bfc7a2c2155230d51f497316cad6d"
+dependencies = [
+ "bytemuck",
+ "cfg-if 1.0.0",
+ "libm",
+ "num-complex",
+ "reborrow",
+ "version_check",
+]
+
 [[package]]
 name = "qoi"
 version = "0.4.1"
@@ -7636,6 +7779,15 @@ dependencies = [
 "bitflags 1.3.2",
 ]

+[[package]]
+name = "raw-cpuid"
+version = "11.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "529468c1335c1c03919960dfefdb1b3648858c20d7ec2d0663e728e4a717efbc"
+dependencies = [
+ "bitflags 2.5.0",
+]
+
 [[package]]
 name = "raw-window-handle"
 version = "0.6.2"
@@ -9774,6 +9926,20 @@ dependencies = [
 "walkdir",
 ]

+[[package]]
+name = "sysctl"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc"
+dependencies = [
+ "bitflags 2.5.0",
+ "byteorder",
+ "enum-as-inner",
+ "libc",
+ "thiserror 1.0.63",
+ "walkdir",
+]
+
 [[package]]
 name = "system-configuration"
 version = "0.5.1"
@@ -10846,27 +11012,34 @@ checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"

 [[package]]
 name = "ug"
-version = "0.0.2"
+version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c4eef2ebfc18c67a6dbcacd9d8a4d85e0568cc58c82515552382312c2730ea13"
+checksum = "03719c61a91b51541f076dfdba45caacf750b230cefaa4b32d6f5411c3f7f437"
 dependencies = [
+ "gemm 0.18.2",
 "half",
+ "libloading 0.8.6",
+ "memmap2",
 "num",
+ "num-traits",
+ "num_cpus",
+ "rayon",
+ "safetensors",
 "serde",
- "serde_json",
 "thiserror 1.0.63",
+ "tracing",
+ "yoke",
 ]

 [[package]]
 name = "ug-cuda"
-version = "0.0.2"
+version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1c4dcab280ad0ef3957e153a82dcad608c954d02cf253b695322f502d1f8902e"
+checksum = "50758486d7941f8b0a636ba7e29455c07071f41590beac1fd307ec893e8db69a"
 dependencies = [
 "cudarc",
 "half",
 "serde",
- "serde_json",
 "thiserror 1.0.63",
 "ug",
 ]
@@ -11192,9 +11365,9 @@ dependencies = [

 [[package]]
 name = "version_check"
-version = "0.9.4"
+version = "0.9.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"

 [[package]]
 name = "videos"
@@ -4,4 +4,4 @@ image_height: 512
 # NB(bt): It looks like we can get down to 2 sampling steps, perhaps ?
 scheduler_steps: 2
 scheduler_samples: 1
-
+cfg_scale: 1.1
@@ -23,6 +23,9 @@ tauri-build = { version = "2.0.4", features = [] }
 filesys = { path = "../../lib/files/filesys" }
 memory_caching = { path = "../../lib/caching/memory_caching" }

+# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
+cudarc = "0.13.7"
+
 # External
 anyhow = "1.0"
 bytes = "1.10"
@@ -44,4 +47,10 @@ tokenizers = "0.15.2"

 [features]
 default = ["cuda"]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+cuda = [
+  "candle-core/cuda",
+  # TODO(bt,2025-02-22): cuDNN might speed up Stable Diffusion
+  # "candle-core/cudnn",
+  "candle-nn/cuda",
+  "candle-transformers/cuda",
+]
@@ -59,11 +59,11 @@ fn do_infer_image(
    prompt: prompt.to_string(),
    uncond_prompt: NEGATIVE_PROMPT.to_string(),
    //guidance_scale: Some(0.0),
-    guidance_scale: None,
    model_cache,
    configs: config,
    prompt_cache: &prompt_cache,
-    strength: strength,
+    i2i_strength: strength,
+    cfg_scale: config.cfg_scale,
  };

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

-
 pub struct Args<'a> {
    pub image: &'a DynamicImage,
    pub prompt: String,
    pub uncond_prompt: String,
-    pub guidance_scale: Option<f64>,
+    pub cfg_scale: Option<f64>,
+    pub i2i_strength: Option<u8>,
    pub configs: &'a AppConfig,
    pub model_cache: &'a ModelCache,
    pub prompt_cache: &'a PromptCache,
-    pub strength: Option<u8>,
 }

 pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
@@ -38,20 +38,33 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
    // TODO(bt,2025-02-18): The scheduler is `EulerAncestralDiscreteScheduler`, but we may want to port an LCM scheduler.
    //  This is a target for performance improvement
    //  See: https://github.com/huggingface/candle/issues/1331
-    println!("Building scheduler...");
    let mut scheduler = args.configs.sd_config.build_scheduler(
        args.configs.scheduler_steps)?;

    let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
+    args.configs.device.set_seed(seed)?;

    println!("Using seed: {}", seed);
-    args.configs.device.set_seed(seed)?;
+
+    let guidance_scale = match args.cfg_scale {
+        Some(guidance_scale) => guidance_scale,
+        None => match args.configs.sd_version {
+            StableDiffusionVersion::V1_5
+            | StableDiffusionVersion::V2_1
+            | StableDiffusionVersion::Xl => 7.5,
+            StableDiffusionVersion::Turbo => 0.,
+            _ => 0., // NB: Not sure what the other model families should use, so sticking with "0"
+        },
+    };
+
+    let use_guide_scale = guidance_scale > 1.0;
+
+    println!("Using guide scale: {}", use_guide_scale);

    info!("Checking if prompt is cached");
    let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?;
-    
-    let text_embeddings = if let Some(tensor) = maybe_cached {
-        info!("Prompt is cached!");
+
+    let mut text_embeddings = if let Some(tensor) = maybe_cached {
        tensor
    } else {
        info!("Prompt is NOT cached! Calculating embedding...");
@@ -66,12 +79,12 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
            false, // use_f16
            &args.configs.device,
            args.configs.dtype,
-            false, // use_guide_scale
+            use_guide_scale,
        )?;
        args.prompt_cache.store_copy(&args.prompt, &tensor)?;
        tensor
    };
-    
+
    println!("Text embeddings shape: {:?}", text_embeddings.shape());

    println!("Loading input image into tensor...");
@@ -97,8 +110,8 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
    println!("Initial latents shape: {:?}", latents.shape());

    println!("Calculating start step for diffusion process...");
-    
-    let img2img_strength = match args.strength {
+
+    let img2img_strength = match args.i2i_strength {
        None => 0.75f64,
        Some(strength) => (strength as f64) / 100.0f64,
    };
@@ -136,20 +149,29 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
        }

        println!("Processing step {}/{} @ timestamp = {}", timestep_index + 1, timesteps.len(), timestep);
-        
-        // TODO(bt,2025-08-18): Do we have to clone the tensor?
-        let latent_model_input = scheduler.scale_model_input(latents.clone(), timestep)?;
-        
-        let noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) {
-            Ok(pred) => {
-                pred
-            },
+
+        let latent_model_input = if use_guide_scale {
+            Tensor::cat(&[&latents, &latents], 0)?
+        } else {
+            latents.clone() // TODO(bt,2025-08-18): Do we have to clone the tensor? `scale_model_input` requires ownership
+        };
+
+        let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
+
+        let mut noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) {
+            Ok(pred) => pred,
            Err(e) => {
                println!("UNet inference failed with error: {}", e);
                return Err(anyhow::anyhow!("UNet inference failed: {}", e));
            },
        };

+        if use_guide_scale {
+            let chunks = noise_pred.chunk(2, 0)?;
+            let (noise_pred_uncond, noise_pred_text) = (&chunks[0], &chunks[1]);
+            noise_pred = (noise_pred_uncond + (guidance_scale * (noise_pred_text - noise_pred_uncond)?)?)?;
+        }
+
        latents = scheduler.step(&noise_pred, timestep, &latents)?;
    }

@@ -32,6 +32,8 @@ pub struct AppConfig {
  pub scheduler_samples: usize,
  
  pub seed: Option<u64>,
+
+  pub cfg_scale: Option<f64>,
  
  /// Probably shouldn't be here.
  pub hf_api: Api,
@@ -72,6 +74,7 @@ impl AppConfig {
      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(1),
      scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
      seed: yaml_configs.seed,
+      cfg_scale: yaml_configs.cfg_scale,
      hf_api,
    })
  }
@@ -12,6 +12,7 @@ pub struct YamlConfig {
  pub scheduler_steps: Option<usize>,
  pub scheduler_samples: Option<usize>,
  pub seed: Option<u64>,
+  pub cfg_scale: Option<f64>,
 }

 impl YamlConfig {
@@ -4525,6 +4525,8 @@
    },
    "node_modules/npm/node_modules/@isaacs/cliui/node_modules/string-width": {
      "version": "5.1.2",
+      "resolved": "https://registry.npmjs.org/string-width/-/string-width-5.1.2.tgz",
+      "integrity": "sha512-HnLOCR3vjcY8beoNLtcjZ5/nxn2afmME6lhrDrebokqMap+XbeW8n9TXpPDOqdGK5qcI3oT0GKTW6wC7EMiVqA==",
      "dev": true,
      "inBundle": true,
      "license": "MIT",
@@ -5265,6 +5267,8 @@
    },
    "node_modules/npm/node_modules/cross-spawn/node_modules/which": {
      "version": "2.0.2",
+      "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz",
+      "integrity": "sha512-BLI3Tl1TW3Pvl70l3yq3Y64i+awpwXqsGBYWkkqMtnbXgrMD+yj7rhW0kuEDxzJaYXGjEW5ogapKNMEKNMjibA==",
      "dev": true,
      "inBundle": true,
      "license": "ISC",
@@ -5313,6 +5317,8 @@
    },
    "node_modules/npm/node_modules/debug/node_modules/ms": {
      "version": "2.1.2",
+      "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.2.tgz",
+      "integrity": "sha512-sGkPx+VjMtmA6MX27oA4FBFELFCZZ4S4XqeGOXCv68tT+jb3vk/RyaKWP0PTKyWtmLSM0b+adUTEvbs1PEaH2w==",
      "dev": true,
      "inBundle": true,
      "license": "MIT"
@@ -7139,6 +7145,8 @@
    },
    "node_modules/npm/node_modules/wrap-ansi-cjs/node_modules/ansi-styles": {
      "version": "4.3.0",
+      "resolved": "https://registry.npmjs.org/ansi-styles/-/ansi-styles-4.3.0.tgz",
+      "integrity": "sha512-zbB9rCJAT1rbjiVDb2hqKFHNYLxgtk8NURxZ3IZwD3F6NtxbXZQCnnSi1Lkx+IDohdPlFp222wVALIheZJQSEg==",
      "dev": true,
      "inBundle": true,
      "license": "MIT",
@@ -7176,6 +7184,8 @@
    },
    "node_modules/npm/node_modules/wrap-ansi/node_modules/string-width": {
      "version": "5.1.2",
+      "resolved": "https://registry.npmjs.org/string-width/-/string-width-5.1.2.tgz",
+      "integrity": "sha512-HnLOCR3vjcY8beoNLtcjZ5/nxn2afmME6lhrDrebokqMap+XbeW8n9TXpPDOqdGK5qcI3oT0GKTW6wC7EMiVqA==",
      "dev": true,
      "inBundle": true,
      "license": "MIT",
@@ -2,6 +2,7 @@ import Konva from "konva";


 import { uiAccess, uiEvents } from "~/signals";
+import { ShapeNode } from "./Nodes";

 import { UndoStackManager } from "./UndoRedo";
 import { CommandManager, MatteBox, SceneManager } from "./EngineUtitlities";
@@ -27,7 +28,8 @@ import { LoadingVideosProvider } from "./EngineUtitlities/LoadingVideosProvider"

 import { VideoExtractionHandler } from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler";
 import { RealTimeDrawEngine } from "./RenderingPrimitives/RealTimeDrawEngine";
-
+import { NodeColor } from "~/signals/uiEvents/toolbarNode";
+import { ShapeType } from "./Nodes";
 // for testing loading files from system
 // import { FileUtilities } from "./FileUtilities/FileUtilities";

@@ -97,6 +99,7 @@ export class Engine {
    this.stage.add(this.nodeIsolationLayer);
    this.stage.add(this.uiLayer);

+
    // Konva Transformer
    this.nodeTransformer = new NodeTransformer();
    this.uiLayer.add(this.nodeTransformer.getKonvaNode());
@@ -177,6 +180,9 @@ export class Engine {
    // hence, lastly, setup these events
    this.setupEventSystem();
    this.setAppMode(AppModes.SELECT);
+
+
+
  }


@@ -359,7 +365,7 @@ export class Engine {
    // Listen to Toolbar Main
    uiEvents.toolbarMain.UNDO.onClick(() => this.undoStackManager.undo());
    uiEvents.toolbarMain.REDO.onClick(() => this.undoStackManager.redo());
-    
+
    uiEvents.toolbarMain.SAVE.onClick(async (/*event*/) => {
      await this.realTimeDrawEngine.saveOutput();

@@ -393,7 +399,8 @@ export class Engine {
      this.addImage(image);
    });

-  
+
+
    uiEvents.onAddTextToEngine((textdata) => {
      this.addText(textdata);
    });
@@ -424,7 +431,7 @@ export class Engine {
        }
      }
    });
-    
+
    uiEvents.promptEvents.onPromptStrengthChanged((strength) => {
      this.realTimeDrawEngine.currentStrength = strength;
    });
@@ -432,6 +439,33 @@ export class Engine {
    uiEvents.promptEvents.onPromptTextChanged((prompt) => {
      this.realTimeDrawEngine.currentPrompt = prompt;
    });
+
+    uiEvents.onAddShapeToEngine((shapeData) => {
+      switch (shapeData.shape) {
+        case "circle":
+          this.addShape(ShapeType.CIRCLE, 100);
+          break;
+        case "square":
+          this.addShape(ShapeType.SQUARE, 100);
+          break;
+        case "triangle":
+          this.addShape(ShapeType.TRIANGLE, 100);
+          break;
+      }
+    });
+
+    uiEvents.toolbarNode.color.onConfirmChanged((nodeColor) => {
+      console.log("Color change event triggered with:", nodeColor);
+      if (!nodeColor) {
+        console.warn("nodeColor is undefined or null");
+        return;
+      }
+      if (!nodeColor.kNodeId) {
+        console.warn("nodeColor.kNodeId is undefined or null");
+        return;
+      }
+      this.changeNodeColor(nodeColor);
+    })
  }

  disableAllButtons() {
@@ -467,7 +501,7 @@ export class Engine {
  }

  // Sandbox is quickly a way to test your idea.
-  public async sandbox() {}
+  public async sandbox() { }

  public onMessage(event: MessageEvent) {
    console.log("Message From Shared Worker");
@@ -514,6 +548,40 @@ export class Engine {
    this.addKeyboardShortcuts();
  }

+  private changeNodeColor(nodeColor: NodeColor) {
+    console.log('Changing node color:', nodeColor);
+    // Find the node that is selected with the id
+    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);
+    // 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);
+      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(), 
+          shapeType: shapeNode.shapeType,
+          size: shapeNode.size,
+          color: nodeColor.color,
+          mediaLayerRef: this.mediaLayer,
+          selectionManagerRef: this.selectionManager
+        });
+        newShapeNode.kNode.position(shapeNode.kNode.position());
+        newShapeNode.kNode.zIndex(shapeNode.kNode.zIndex());
+        
+        // Remove old node and add new one
+        shapeNode.kNode.destroy();
+        this.commandManager.createNode(newShapeNode);
+      }
+    }
+  }
+
  public addText(textNodeData: TextNodeData) {
    const textNode = new TextNode({
      textNodeData: textNodeData,
@@ -525,6 +593,24 @@ export class Engine {
    this.commandManager.createNode(textNode);
  }

+  public addShape(type: ShapeType, size: number, color?: string) {
+    const shapeNode = new ShapeNode({
+      canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
+      canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
+      shapeType: type,
+      size: size,
+      color: color,
+      mediaLayerRef: this.mediaLayer,
+      selectionManagerRef: this.selectionManager
+    });
+    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 addImage(imageFile: File) {
    const imageNode = new ImageNode({
      mediaLayerRef: this.mediaLayer,
@@ -533,9 +619,9 @@ export class Engine {
      imageFile: imageFile,
      selectionManagerRef: this.selectionManager,
    });
-
+    imageNode.kNode.zIndex(1);
    this.commandManager.createNode(imageNode);
-    this.realTimeDrawEngine.addNodes(imageNode);
+    //this.realTimeDrawEngine.addNodes(imageNode);
  }

  public addVideo(
@@ -579,7 +665,7 @@ export class Engine {
          if (
            longestVideoNode === undefined ||
            node.videoComponent.duration >
-              longestVideoNode.videoComponent.duration
+            longestVideoNode.videoComponent.duration
          ) {
            longestVideoNode = node;
          }
@@ -1,5 +1,5 @@
 import Konva from "konva";
-import { ImageNode, TextNode, VideoNode } from "../Nodes";
+import { ImageNode, ShapeNode, TextNode, VideoNode } from "../Nodes";
 import { MediaNode, Position, RGBColor, Transformation } from "../types";
 import {
  AddChromaKey,
@@ -48,7 +48,7 @@ export class CommandManager {
    this.undoStackManagerRef = engineRefs.undoStackManagerRef;
  }

-  createNode(node: VideoNode | ImageNode | TextNode) {
+  createNode(node: VideoNode | ImageNode | TextNode | ShapeNode) {
    const command = new CreateCommand({
      nodes: new Set<MediaNode>([node]),
      mediaLayerRef: this.mediaLayerRef,
@@ -58,6 +58,10 @@ export class CommandManager {
      renderEngineRef: this.renderEngineRef,
    });
    this.undoStackManagerRef.executeCommand(command);
+
+    // 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());
  }
  deleteNodes() {
    const nodes = this.selectionManagerRef.getSelectedNodes();
@@ -29,7 +29,7 @@ export abstract class BaseNode {
  }: {
    selectionManagerRef: SelectionManager;
    mediaLayerRef: Konva.Layer;
-    kNode: Konva.Image | Konva.Group;
+    kNode: Konva.Image | Konva.Group | Konva.Shape;
  }) {
    // console.log("Node constructed");
    // this.uuid = uuidv4();
@@ -67,8 +67,11 @@ export abstract class BaseNode {
    }
    if (this.kNode instanceof Konva.Group) {
      const wrapperRect = this.kNode.findOne(".wrapper") as Konva.Rect;
-      wrapperRect.stroke(primaryOrange);
-      wrapperRect.strokeWidth(highlightStrokeWidth);
+      console.log("Wrapper rect:", wrapperRect);
+      if (wrapperRect !=null) {
+        wrapperRect.stroke(primaryOrange);
+        wrapperRect.strokeWidth(highlightStrokeWidth);
+      }
      return;
    }
    if (import.meta.env.DEV) {
@@ -86,7 +89,11 @@ export abstract class BaseNode {
    }
    if (this.kNode instanceof Konva.Group) {
      const wrapperRect = this.kNode.findOne(".wrapper") as Konva.Rect;
-      wrapperRect.strokeWidth(0);
+      console.log("Wrapper rect:", wrapperRect); 
+      if (wrapperRect !=null) {
+        // Issue investigate
+        wrapperRect.strokeWidth(0);
+      }
      return;
    }
    if (import.meta.env.DEV) {
@@ -83,8 +83,10 @@ export class ImageNode extends NetworkedNode {
  }) {
    const imageComponent = new Image();
    imageComponent.crossOrigin = "anonymous";
+    // this.setProgress({ progress: 100, status: UploadStatus.SUCCESS });
+
    imageComponent.onload = () => {
-      this.setProgress({ progress: 0, status: UploadStatus.FILE_STAGED });
+      //this.setProgress({ progress: 0, status: UploadStatus.FILE_STAGED });
      this.imageSize = {
        width: imageComponent.width,
        height: imageComponent.height,
@@ -105,33 +107,33 @@ export class ImageNode extends NetworkedNode {
      this.kNode.fill(transparent);
      this.listenToBaseKNode();
      this.mediaLayerRef.draw();
-      this.uploadImage(imageFile);
+      //this.uploadImage(imageFile);
    };
    imageComponent.onerror = () => {
-      this.setProgress({ progress: 0, status: UploadStatus.ERROR_ON_FILE });
+      //this.setProgress({ progress: 0, status: UploadStatus.ERROR_ON_FILE });
    };
    imageComponent.src = URL.createObjectURL(imageFile);
  }
  private async loadImageFromUrl(mediaFileUrl: string) {
-    this.setProgress({ progress: 75, status: UploadStatus.LOADING });
+    //this.setProgress({ progress: 75, status: UploadStatus.LOADING });

    const newImage = new Image();
    newImage.crossOrigin = "anonymous";
    newImage.onerror = () => {
-      this.setProgress({ progress: 90, status: UploadStatus.ERROR_ON_LOAD });
+      //this.setProgress({ progress: 90, status: UploadStatus.ERROR_ON_LOAD });
    };
    newImage.onload = () => {
      // console.log("network image", newImage);
      this.kNode.image(newImage);
      this.kNode.draw();
-      this.setProgress({ progress: 100, status: UploadStatus.SUCCESS });
+      //this.setProgress({ progress: 100, status: UploadStatus.SUCCESS });
    };
    newImage.src = mediaFileUrl;
    this.listenToBaseKNode();
  }

  private async uploadImage(imageFile: File) {
-    this.setProgress({ progress: 10, status: UploadStatus.UPLOADING });
+    //this.setProgress({ progress: 10, status: UploadStatus.UPLOADING });

    const mediaUploadApi = new MediaUploadApi();
    const uploadResponse = await mediaUploadApi.UploadImage({
@@ -149,7 +151,7 @@ export class ImageNode extends NetworkedNode {
  }

  private async retreiveImage(mediaFileToken: string) {
-    this.setProgress({ progress: 50, status: UploadStatus.RETREIVING });
+    //this.setProgress({ progress: 50, status: UploadStatus.RETREIVING });
    const mediaFileApi = new MediaFilesApi();
    const mediaFileResponse = await mediaFileApi.GetMediaFileByToken({
      mediaFileToken: mediaFileToken,
@@ -0,0 +1,202 @@
+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";
+
+export enum ShapeType {
+  CIRCLE = "circle",
+  SQUARE = "square",
+  TRIANGLE = "triangle",
+}
+
+interface ShapeNodeConstructor {
+  canvasPosition: Position;
+  canvasSize: Size;
+  shapeType: ShapeType;
+  size: number;
+  color?: string; // Hex color string
+  transform?: TransformationData;
+  mediaLayerRef: Konva.Layer;
+  selectionManagerRef: SelectionManager;
+}
+
+export class ShapeNode extends NetworkedNode {
+  
+  retry(): void {
+    throw new Error("Method not implemented.");
+  }
+
+  public shapeType: ShapeType;
+  public size: number;
+
+  constructor({
+    canvasPosition,
+    canvasSize,
+    shapeType,
+    size,
+    color = "#ff0000", // Default red if no color provided
+    transform: existingTransform,
+    mediaLayerRef,
+    selectionManagerRef,
+  }: ShapeNodeConstructor) {
+    // Create the actual shape inside the group
+    const transform = NodeUtilities.getInitialTransform({
+      existingTransform,
+      canvasPosition,
+      canvasSize,
+    });
+
+    let shape = null;
+    
+    switch (shapeType) {
+      case ShapeType.CIRCLE:
+        const circleShape = new Konva.Circle({
+          radius: 50,
+          ...transform,
+          fill: color,
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        // Convert circle to image
+        const dataURL = circleShape.toDataURL();
+        const imageObj = new Image();
+        imageObj.src = dataURL;
+
+        shape = new Konva.Image({
+          image: imageObj,
+          x: circleShape.x(),
+          y: circleShape.y(),
+          width: circleShape.width(),
+          height: circleShape.height(),
+          ...transform,
+          fill: 'transparent',
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        imageObj.onload = () => {
+          this.mediaLayerRef.draw();
+        };
+
+        // Clean up the temporary circle
+        circleShape.destroy();
+        break;
+
+      case ShapeType.SQUARE:
+        const squareShape = new Konva.Rect({
+          width: 100,
+          height: 100,
+          ...transform,
+          fill: color,
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        // Convert square to image
+        const squareDataURL = squareShape.toDataURL();
+        const squareImageObj = new Image();
+        squareImageObj.src = squareDataURL;
+
+        shape = new Konva.Image({
+          image: squareImageObj,
+          x: squareShape.x(),
+          y: squareShape.y(),
+          width: squareShape.width(),
+          height: squareShape.height(),
+          ...transform,
+          fill: 'transparent',
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        squareImageObj.onload = () => {
+          this.mediaLayerRef.draw();
+        };
+
+        // Clean up the temporary square
+        squareShape.destroy();
+        break;
+
+      case ShapeType.TRIANGLE:
+        const triangleShape = new Konva.RegularPolygon({
+          sides: 3,
+          radius: 50,
+          ...transform,
+          fill: color,
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        // Convert triangle to image
+        const triangleDataURL = triangleShape.toDataURL();
+        const triangleImageObj = new Image();
+        triangleImageObj.src = triangleDataURL;
+
+        shape = new Konva.Image({
+          image: triangleImageObj,
+          x: triangleShape.x(),
+          y: triangleShape.y(),
+          width: triangleShape.width(),
+          height: triangleShape.height(),
+          ...transform,
+          fill: 'transparent',
+          strokeScaleEnabled: false,
+          draggable: true,
+        });
+
+        triangleImageObj.onload = () => {
+          this.mediaLayerRef.draw();
+        };
+
+        // Clean up the temporary triangle
+        triangleShape.destroy();
+        break;
+
+      default:
+        throw new Error("Invalid shape type");
+    }
+
+    super({
+      selectionManagerRef: selectionManagerRef,
+      mediaLayerRef: mediaLayerRef,
+      kNode: shape,
+    });
+
+    this.shapeType = shapeType;
+    this.size = size;
+
+    // Add shape to group
+    const centerPosition = NodeUtilities.positionNodeOnCanvasCenter({
+      canvasOffset: canvasPosition,
+      componentSize: canvasSize,
+      maxSize: canvasSize,
+    });
+    this.kNode.setPosition(centerPosition);
+    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;
+  }
+}
@@ -4,6 +4,7 @@ export enum NodeType {
  IMAGE = "image",
  TEXT = "text",
  VIDEO = "video",
+  SHAPE = "shape"
 }
 export const transparent = Colors.transparent;
 export const primaryOrange = Colors.primaryOrange;
@@ -3,3 +3,4 @@ export { BaseNode } from "./BaseNode";
 export { ImageNode } from "./ImageNode";
 export { VideoNode } from "./VideoNode";
 export { TextNode } from "./TextNode";
+export { ShapeNode , ShapeType} from "./ShapeNode";
@@ -4,11 +4,12 @@ import { MediaNode, Position, Transformation } from "../types";
 import { uiAccess, uiEvents } from "~/signals";
 import { NetworkedNode } from "../Nodes/NetworkedNode";
 import { LoadingBarStatus } from "~/components/ui";
-import { ImageNode, TextNode, VideoNode } from "../Nodes";
+import { ImageNode, ShapeNode, TextNode, VideoNode } from "../Nodes";
 import {
  calculateContextualsPosition,
  getImageNodeButtonStates,
  getMultiSelectButtonStates,
+  getShapeNodeButtonStates,
  getTextNodeButtonStates,
  getVideoNodeButtonStates,
 } from "./ToolbarNodeUtilities";
@@ -273,17 +274,29 @@ export class SelectionManager {
    if (this.selectedNodes.size === 1) {
      if (node instanceof ImageNode) {
        showOrUpdate({
+          knodeIds: [node.kNode.id()],
          locked: node.isLocked(),
          buttonStates: getImageNodeButtonStates({ locked: node.isLocked() }),
        });
      }
-      if (node instanceof TextNode) {
+      else if (node instanceof ShapeNode) {
+        console.debug("Selected node is of type ShapeNode", node.kNode.id())
        showOrUpdate({
+          knodeIds: [node.kNode.id()],
          locked: node.isLocked(),
-          buttonStates: getTextNodeButtonStates({ locked: node.isLocked() }),
+          buttonStates: getShapeNodeButtonStates({ locked: node.isLocked() }),
+          color: node.getNodeData({ x: 0, y: 0 })?.textNodeData?.color ?? "#000000",
        });
      }
-      if (node instanceof VideoNode) {
+      else if (node instanceof TextNode) {
+        showOrUpdate({
+          knodeIds: [node.kNode.id()],
+          locked: node.isLocked(),
+          buttonStates: getTextNodeButtonStates({ locked: node.isLocked() }),
+          color: node.getNodeData({ x: 0, y: 0 })?.textNodeData?.color ?? "#000000",
+        });
+      }
+      else if (node instanceof VideoNode) {
        showOrUpdate({
          locked: node.isLocked(),
          buttonStates: getVideoNodeButtonStates({ locked: node.isLocked() }),
@@ -51,7 +51,8 @@ export function getImageNodeButtonStates(
        buttonName === ButtonNames.AI_STYLIZE ||
        buttonName === ButtonNames.SEGMENTATION ||
        buttonName === ButtonNames.DOWNLOAD ||
-        buttonName === ButtonNames.CHROMA
+        buttonName === ButtonNames.CHROMA ||
+        buttonName === ButtonNames.COLOR
      ) {
        buttonStates[buttonName] = {
          disabled: true,
@@ -130,7 +131,9 @@ export function getVideoNodeButtonStates(
        return buttonStates;
      }
      // soon to come feature is disabled
-      if (buttonName === ButtonNames.AI_STYLIZE) {
+      if (buttonName === ButtonNames.AI_STYLIZE ||
+        buttonName === ButtonNames.COLOR
+      ) {
        buttonStates[buttonName] = {
          disabled: true,
          hidden: true,
@@ -149,7 +152,8 @@ export function getVideoNodeButtonStates(
    {} as ContextualToolbarProps["buttonStates"],
  );
 }
-export function getMultiSelectButtonStates(
+
+export function getShapeNodeButtonStates(
  props: { locked?: boolean | "unknown" } = {},
 ) {
  const ButtonNames = ToolbarNodeButtonNames;
@@ -157,7 +161,7 @@ export function getMultiSelectButtonStates(
    (buttonStates, buttonName) => {
      // transform button depends on locked state
      if (props.locked !== undefined && buttonName === ButtonNames.TRANSFORM) {
-        buttonStates[buttonName] = {
+        buttonStates.TRANSFORM = {
          disabled: props.locked === "unknown" || props.locked === true,
          hidden: false,
          active: true,
@@ -191,3 +195,47 @@ export function getMultiSelectButtonStates(
    {} as ContextualToolbarProps["buttonStates"],
  );
 }
+
+export function getMultiSelectButtonStates(
+  props: { locked?: boolean | "unknown" } = {},
+) {
+  const ButtonNames = ToolbarNodeButtonNames;
+  return Object.values(ButtonNames).reduce(
+    (buttonStates, buttonName) => {
+      // transform button depends on locked state
+      if (props.locked !== undefined && buttonName === ButtonNames.TRANSFORM) {
+        buttonStates[buttonName] = {
+          disabled: props.locked === "unknown" || props.locked === true,
+          hidden: false,
+          active: true,
+        };
+        return buttonStates;
+      }
+
+      // hidden buttons
+      if (
+        buttonName === ButtonNames.AI_STYLIZE ||
+        buttonName === ButtonNames.SEGMENTATION ||
+        buttonName === ButtonNames.DOWNLOAD ||
+        buttonName === ButtonNames.CHROMA ||
+        buttonName === ButtonNames.COLOR
+      ) {
+        buttonStates[buttonName] = {
+          disabled: true,
+          hidden: true,
+          active: false,
+        };
+        return buttonStates;
+      }
+
+      // all other buttons
+      buttonStates[buttonName] = {
+        disabled: false,
+        hidden: false,
+        active: false,
+      };
+      return buttonStates;
+    },
+    {} as ContextualToolbarProps["buttonStates"],
+  );
+}
@@ -3,12 +3,10 @@ import { Container } from "konva/lib/Container";
 import { Shape } from "konva/lib/Shape";
 import { Group } from "konva/lib/Group";

-
 import { invoke } from "@tauri-apps/api/core";

-
 import { FileUtilities } from "../FileUtilities/FileUtilities";
-import { ImageNode, VideoNode, TextNode } from "../Nodes";
+import { ImageNode, VideoNode, TextNode, ShapeNode, ShapeType } from "../Nodes";
 import { MediaNode } from "../types";

 import { RenderTask } from "./RenderTask";
@@ -18,7 +16,7 @@ import { OffScreenSceneCanvas } from "./OffScreenSceneCanvas";

 export class RealTimeDrawEngine {
  private videoNodes: VideoNode[];
-  private imageNodes: (ImageNode | TextNode)[];
+  private imageNodes: (ImageNode | TextNode | ShapeNode)[];

  private offScreenCanvas: OffscreenCanvas;
  private outputBitmap: ImageBitmap | undefined;
@@ -26,7 +24,7 @@ export class RealTimeDrawEngine {
  // private frames: ImageBitmap[];

  // capturing composite within window
-  private bgLayerRef: Konva.Layer;
+
  private mediaLayerRef: Konva.Layer;

  private height: number;
@@ -51,13 +49,11 @@ export class RealTimeDrawEngine {
  constructor({
    width,
    height,
-    bgLayerRef,
    mediaLayerRef,
    offScreenCanvas,
  }: {
    width: number;
    height: number;
-    bgLayerRef: Konva.Layer;
    mediaLayerRef: Konva.Layer;
    offScreenCanvas: OffscreenCanvas;
  }) {
@@ -80,8 +76,6 @@ export class RealTimeDrawEngine {
    this.offScreenCanvas.width = this.width;
    this.offScreenCanvas.height = this.height;

-    this.bgLayerRef = bgLayerRef;
-
    // this is the whole canvas
    this.mediaLayerRef = mediaLayerRef;

@@ -90,7 +84,7 @@ export class RealTimeDrawEngine {
    this.port = undefined;

    this.fps = 24;
-    
+
    this.currentPrompt = "";
    this.currentStrength = 100;
    // This is captures a subset of the medialayer ref
@@ -106,20 +100,6 @@ export class RealTimeDrawEngine {
      draggable: false,
    });

-    // Add preview canvas with same dimensions but different style
-    // this.previewCanvas = new Konva.Rect({
-    //   name: "PreviewCanvas",
-    //   x: this.positionX,
-    //   y: this.positionY,
-    //   width: this.width,
-    //   height: this.height,
-    //   fill: "white",
-    //   stroke: "blue",
-    //   strokeWidth: 1,
-    //   dash: [5, 5],
-    //   draggable: false,
-    // });
-
    this.previewCanvas = new Konva.Image({
      name: "PreviewCanvas",
      x: this.positionX,
@@ -127,21 +107,90 @@ export class RealTimeDrawEngine {
      width: this.width,
      height: this.height,
      image: undefined,
-      stroke: "blue",
+      stroke: "black",
      strokeWidth: 1,
-      dash: [5, 5],
      draggable: false,
+      fill: "white",
    });

-    this.bgLayerRef.add(this.captureCanvas);
-    this.bgLayerRef.add(this.previewCanvas);
+    this.mediaLayerRef.add(this.captureCanvas);
+    this.mediaLayerRef.add(this.previewCanvas);
    // send back
    this.captureCanvas.setZIndex(0);
-    this.previewCanvas.setZIndex(0);
+    this.previewCanvas.setZIndex(1);
+    // Add mouse events for preview canvas copying
+    //this.previewCopyListener();

-    //this.debug();
  }

+  public previewCopyListener() {
+    this.previewCanvas.on('mousedown touchstart', () => {
+      if (!this.outputBitmap) {
+        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({
+        x: this.previewCanvas.x(),
+        y: this.previewCanvas.y(),
+        width: this.width,
+        height: this.height,
+        image: this.outputBitmap,
+        draggable: true,
+        listening: true
+      });
+
+      this.mediaLayerRef.add(previewCopy);
+      previewCopy.moveToTop();
+      this.mediaLayerRef.draw();
+
+      // Start dragging immediately
+      previewCopy.startDrag();
+
+      // Handle drag events
+      previewCopy.on('dragmove', () => {
+        this.mediaLayerRef.draw();
+      });
+
+      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()
+          });
+          this.mediaLayerRef.batchDraw();
+        } else {
+          // Remove if dropped outside capture area
+          previewCopy.destroy();
+          this.mediaLayerRef.batchDraw();
+        }
+      });
+    });
+  }
+
+  public findImageNodeById(id: string): (ImageNode | TextNode | ShapeNode | undefined) {
+    return this.imageNodes.find(node => {
+      if (node.kNode) {
+        return node.kNode.id() === id;
+      }
+      return false;
+    });
+  }
  async updateCaptureCanvas(
    width: number | undefined,
    height: number | undefined,
@@ -284,12 +333,13 @@ export class RealTimeDrawEngine {
    this.isProcessing = true;
    await this.render();
  };
+  
  public addNodes(node: MediaNode) {
-    if (node instanceof VideoNode) {
-      this.videoNodes.push(node);
-      node.kNode.on("dragend", this.handleNodeDragEnd);
-    } else if (node instanceof ImageNode || node instanceof TextNode) {
+   
+    if (node instanceof ImageNode || node instanceof TextNode || node instanceof ShapeNode) {
+      console.debug("Adding node:", node);
      this.imageNodes.push(node);
+      console.log(this.imageNodes)
      node.kNode.on("dragend", this.handleNodeDragEnd);
    }
  }
@@ -325,6 +375,7 @@ 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);
@@ -384,8 +435,6 @@ export class RealTimeDrawEngine {
    }
  }

-
-
  private async imageBitmapToBase64(imageBitmap: ImageBitmap): Promise<string> {
    // Create a temporary canvas
    const canvas = document.createElement("canvas");
@@ -462,9 +511,11 @@ export class RealTimeDrawEngine {
  }

  public async render() {
-    // only pick nodes that intersect with the canvas on screen bounds to freeze.
+    // 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();
+    console.log("All nodes in mediaLayer:", nodes);
    console.log(
      `context: x:${this.positionX} y:${this.positionY} ${this.width} x ${this.height}`,
    );
@@ -486,16 +537,16 @@ export class RealTimeDrawEngine {
      this.outputBitmap = bitmap;
      this.previewCanvas.image(bitmap);
      this.isProcessing = false;
-      return
-    } 
+      return;
+    }

    try {
      const base64Bitmap = await this.imageBitmapToBase64(bitmap);

      const base64BitmapResponse = await invoke("infer_image", {
        image: base64Bitmap,
-        prompt:this.currentPrompt,
-        strength:this.currentStrength,
+        prompt: this.currentPrompt,
+        strength: this.currentStrength,
      });

      console.log(base64BitmapResponse);
@@ -510,6 +561,5 @@ export class RealTimeDrawEngine {
    } finally {
      this.isProcessing = false;
    }
-
  }
 }
@@ -1,4 +1,4 @@
-import { NodeType, BaseNode, ImageNode, TextNode, VideoNode } from "../Nodes";
+import { NodeType, BaseNode, ImageNode, TextNode, VideoNode, ShapeNode } from "../Nodes";

 import { TransformationData } from "./Transformation";

@@ -6,7 +6,7 @@ import { TextNodeData } from "./Text";

 export type MediaNode =
  // | NetworkedNode
-  BaseNode | VideoNode | ImageNode | TextNode;
+  BaseNode | VideoNode | ImageNode | TextNode | ShapeNode;

 export type ImageNodeData = {
  mediaFileUrl: string;
@@ -26,6 +26,10 @@ export type VideoNodeData = {
  extractionUrl?: string;
 };

+export type ShapeNodeData = {
+  shape: "triangle" | "square" | "circle";
+}
+
 export type NodeData = {
  type: NodeType;
  transform: TransformationData;
@@ -50,6 +50,20 @@ export const ContextualToolbarNode = () => {
    };
  }, []);

+  const handleColorChange = (color: string) => {
+    // Don't do anything if no node is known to be bound
+    if (uiAccess.toolbarNode.signal.value.knodeIds.length < 1) {
+      return;
+    }
+
+    // Update the UI component
+    uiAccess.toolbarNode.setColor(color);
+
+    // Signal the engine with the node ID (just grab the first -and expected only- node)
+    const nodeID = uiAccess.toolbarNode.signal.value.knodeIds[0];
+    dispatchUiEvents.toolbarNode.setSelectedColor({ kNodeId: nodeID, color });
+  }
+
  const boundedTopPosition = Math.min(position.y, winSize.height - 164);
  return (
    <Transition
@@ -71,6 +85,8 @@ export const ContextualToolbarNode = () => {
        onLockClicked={(e) => {
          dispatchUiEvents.toolbarNode.lock(e);
        }}
+        color={uiAccess.toolbarNode.signal.value.color}
+        onColorChange={handleColorChange}
      />
    </Transition>
  );
@@ -12,6 +12,7 @@ import { LayoutSignalType } from "./contextSignals/layout";
 import { AppUiContextInterface } from "./contextSignals/appUi";
 import { twMerge } from "tailwind-merge";
 import { paperWrapperStyles } from "~/components/styles";
+import { dispatchUiEvents } from "~/signals";

 export const SignaledToolbarMain = ({
  // layoutSignal,
@@ -43,6 +44,12 @@ export const SignaledToolbarMain = ({
        return () => appUiContext.openAddImage();
      case ToolbarMainButtonNames.ADD_VIDEO:
        return () => appUiContext.openAddVideo();
+      case ToolbarMainButtonNames.ADD_CIRCLE:
+        return () => dispatchUiEvents.addShapeToEngine({ shape: "circle" });
+      case ToolbarMainButtonNames.ADD_TRIANGLE:
+        return () => dispatchUiEvents.addShapeToEngine({ shape: "triangle" });
+      case ToolbarMainButtonNames.ADD_SQUARE:
+        return () => dispatchUiEvents.addShapeToEngine({ shape: "square" });
      default:
        return dispatchers[buttonName];
    }
@@ -54,10 +54,10 @@
  .dotted-pattern {
    background-image: radial-gradient(
      circle at 1px 1px,
-      rgb(0 0 0 / 0.25) 1px,
+      rgb(255 255 255 / 0.25) 1px,
      transparent 0
    );
-    background-size: 32px 32px;
+    background-size: 40px 40px;
  }

  .gradient-text {
@@ -1,107 +1,169 @@
 import Navbar from "../../../components/landing/Navbar";
-import Hero from "../../../components/landing/Hero";
-import FeatureCard from "../../../components/landing/FeatureCard";
 import Footer from "../../../components/landing/Footer";
+import { Button } from "~/components/ui";
+import { faApple } from "@fortawesome/free-brands-svg-icons";
+import { faDesktop } from "@fortawesome/pro-solid-svg-icons";
+import { useIsMobile } from "~/hooks/useIsMobile";

 const Landing = () => {
+  const isMobile = useIsMobile();
  const features = [
    {
-      title: "Step 1: Compose Your Scene",
+      title: "Lorem ipsum dolor sit",
      description:
-        "Begin your story on our intuitive canvas. Drop and arrange photos and videos to create layered scenes. Position, scale, and transform elements with precision to build your perfect composition. Our dynamic canvas makes it easy to bring your vision to life.",
-      image: "/images/feature_compositing.png",
-      gradient: "bg-gradient-to-br from-purple-500/20 to-blue-500/20",
+        "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt.",
+      image: "/images/features/ai-creation.jpg",
+      tag: "Core",
    },
    {
-      title: "Step 2: Smart Segmentation",
+      title: "Ut enim ad minim",
      description:
-        "Clean up your scenes with powerful AI segmentation. Remove green screens instantly, isolate subjects from any background, and perfect your layer masks with intelligent tools. Create professional-quality compositions with just a few clicks.",
-      image: "/images/feature_segmentation.png",
-      gradient: "bg-gradient-to-br from-emerald-500/20 to-cyan-500/20",
+        "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea.",
+      image: "/images/features/collaboration.jpg",
+      tag: "Workflow",
    },
    {
-      title: "Step 3: AI Style Transfer",
+      title: "Duis aute irure dolor",
      description:
-        "Transform your composition with AI style transfer magic. Choose from our curated collection of artistic styles. Apply consistent styling across your entire scene to create a cohesive visual masterpiece.",
-      image: "/images/feature_ai_stylize.png",
-      gradient: "bg-gradient-to-br from-rose-500/20 to-orange-500/20",
-    },
-    {
-      title: "Step 4: Render Your Video",
-      description:
-        "Render your styled composition into a stunning final movie. Export in high quality, ready to share with the world.",
-      image: "/images/feature_render.png",
-      gradient: "bg-gradient-to-br from-blue-500/20 to-indigo-500/20",
+        "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla.",
+      image: "/images/features/style-models.jpg",
+      tag: "Tools",
    },
  ];

-  const scrollToFeature = (index: number) => {
-    const element = document.getElementById(`feature-${index + 1}`);
-    if (element) {
-      element.scrollIntoView({ behavior: "smooth", block: "center" });
-    }
-  };
-
  return (
-    <div className="min-h-screen bg-white">
+    <div className="relative min-h-screen bg-[#101014] text-white">
      <Navbar />

-      <div>
-        <Hero onFeatureClick={scrollToFeature} />
-
-        <section className="relative overflow-hidden py-32">
-          {/* Decorative elements */}
-          <div className="dotted-pattern absolute inset-0 opacity-50" />
-
-          <div className="relative mx-auto max-w-7xl px-4 sm:px-6 lg:px-8">
-            {/* Section header with enhanced styling */}
-            <div className="mx-auto mb-24 max-w-3xl text-center">
-              <div className="mb-6 inline-flex items-center rounded-full bg-gradient-to-r from-purple-500/10 via-blue-500/10 to-purple-500/10 px-8 py-2.5">
-                <span className="whitespace-nowrap bg-gradient-to-r from-purple-600 via-blue-600 to-purple-600 bg-clip-text text-sm font-semibold text-transparent">
-                  AI-POWERED WORKFLOW
+      <div className="dotted-pattern absolute inset-0 z-[0] opacity-50" />
+      {/* Hero Section */}
+      <div className="relative flex overflow-hidden pt-24 xl:min-h-[1000px] xl:items-center xl:pt-0">
+        {/* Content Container */}
+        <div className="relative z-10 mx-auto grid h-full w-full max-w-[1920px] grid-cols-12 items-center px-6 md:px-16 xl:px-32">
+          <div className="z-10 col-span-12 w-full xl:absolute xl:left-0 xl:top-1/2 xl:max-w-4xl xl:-translate-y-1/2 xl:pl-32 before:xl:absolute before:xl:-inset-0 before:xl:-z-10 before:xl:mr-40 before:xl:rounded-full before:xl:bg-black/90 before:xl:to-transparent before:xl:blur-[180px]">
+            <div className="max-w-2xl">
+              <div className="mb-6">
+                <span className="text-lg font-semibold uppercase tracking-widest text-gray-400">
+                  MIRA
                </span>
              </div>
-              <h2 className="mb-6 text-4xl font-bold tracking-tight text-gray-900">
-                Create Stunning Videos in
-                <span className="relative">
-                  <span className="gradient-text relative z-10">
-                    {" "}
-                    Four Simple Steps
-                  </span>
-                  <span className="bg-rose-200/50 absolute -bottom-2 left-0 z-0 h-3 w-full" />
-                </span>
-              </h2>
-              <p className="text-xl leading-relaxed text-gray-600">
-                From composition to final render, transform your media into
-                cinematic stories with our powerful AI-enhanced workflow.
+
+              <h1 className="mb-6 font-bold leading-tight lg:text-6xl">
+                Realtime AI editor.
+                <br />
+                On your local machine.
+              </h1>
+
+              <p className="mb-10 max-w-xl text-lg leading-relaxed">
+                Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do
+                eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut
+                enim ad minim veniam, quis nostrud exercitation.
              </p>
-            </div>

-            {/* Features stack with enhanced visual separation */}
-            <div className="relative space-y-32">
-              {/* Vertical timeline line */}
-              <div className="to-rose-500/20 absolute left-1/2 top-0 h-full w-px bg-gradient-to-b from-purple-500/20 via-blue-500/20" />
-
-              {features.map((feature, index) => (
-                <div
-                  key={index}
-                  id={`feature-${index + 1}`}
-                  className="relative"
+              <div className="flex flex-col gap-4 sm:flex-row">
+                <Button
+                  className={`w-full transform rounded-lg px-8 py-4 font-semibold transition sm:w-auto ${
+                    isMobile
+                      ? "cursor-not-allowed bg-[#2D81FF]"
+                      : "bg-[#2D81FF] hover:scale-105 hover:bg-[#438AF6]"
+                  }`}
+                  disabled={isMobile}
+                  icon={isMobile ? faDesktop : faApple}
                >
-                  {/* Step number bubble */}
-                  <div className="absolute left-1/2 -translate-x-1/2 -translate-y-16">
-                    <div className="flex h-12 w-12 items-center justify-center rounded-full border bg-white shadow-lg">
-                      <span className="bg-gradient-to-br from-purple-600 to-purple-600 bg-clip-text text-xl font-bold text-transparent">
-                        {index + 1}
-                      </span>
-                    </div>
-                  </div>
-                  <FeatureCard {...feature} isReversed={index % 2 === 1} />
-                </div>
-              ))}
+                  {isMobile ? "Download on desktop" : "Download for MacOS"}
+                </Button>
+                <button className="w-full transform rounded-lg border border-white/30 px-8 py-4 font-semibold transition hover:scale-105 hover:bg-white/10 sm:w-auto">
+                  Learn more
+                </button>
+              </div>
            </div>
          </div>
-        </section>
+          <div className="col-span-12 mt-12 aspect-video w-full transform overflow-hidden rounded-xl border border-white/[16%] bg-transparent md:mt-16 xl:col-span-9 xl:col-start-4 xl:mt-0">
+            <img
+              src="/images/landing_hero.png"
+              alt="AI Video Creation"
+              className="h-full w-full object-cover object-top opacity-90"
+            />
+          </div>
+        </div>
+      </div>
+
+      {/* Features Section */}
+      <div className="relative z-10 mx-auto mb-32 max-w-[1920px] px-4 py-20 sm:px-24 lg:px-32">
+        <div className="mb-8 flex items-center justify-between">
+          <h1 className="text-4xl font-bold">Features</h1>
+          {/* <button className="text-sm text-[#0096CE] hover:text-[#0084B5]">
+            View all
+          </button> */}
+        </div>
+
+        <div className="grid gap-6 md:grid-cols-3">
+          {features.map((item, index) => (
+            <div
+              key={index}
+              className="group cursor-pointer overflow-hidden rounded-xl bg-[#1C1C20] transition hover:bg-[#27272B]"
+            >
+              <div className="aspect-video overflow-hidden">
+                <img
+                  src={item.image}
+                  alt={item.title}
+                  className="h-full w-full object-cover transition group-hover:scale-105"
+                />
+              </div>
+              <div className="p-6">
+                {/* <div className="mb-4 text-sm text-[#0096CE]">{item.tag}</div> */}
+                <h3 className="mb-2 text-xl font-semibold">{item.title}</h3>
+                <p className="text-gray-400">{item.description}</p>
+              </div>
+            </div>
+          ))}
+        </div>
+
+        {/* Installation Steps Section */}
+        <div className="mt-36">
+          <h1 className="mb-12 text-5xl font-bold">
+            How to install Mira AI Editor
+          </h1>
+          <div className="grid gap-6 md:grid-cols-3">
+            <div className="rounded-xl bg-[#1C1C20] p-8">
+              <div className="mb-4 w-fit rounded-full bg-white/15 px-3 py-1 text-sm font-medium text-white">
+                STEP 1
+              </div>
+              <h3 className="mb-4 text-2xl font-semibold">Download the app</h3>
+              <p className="mb-6 text-gray-400">
+                Download and install our application to start creating amazing
+                AI-powered content.
+              </p>
+              <Button className="rounded-lg bg-[#2D81FF] px-4 py-2 text-sm font-medium hover:bg-[#438AF6]">
+                Download App
+              </Button>
+            </div>
+
+            <div className="rounded-xl bg-[#1C1C20] p-8">
+              <div className="mb-4 w-fit rounded-full bg-white/15 px-3 py-1 text-sm font-medium text-white">
+                STEP 2
+              </div>
+              <h3 className="mb-4 text-2xl font-semibold">Create an account</h3>
+              <p className="mb-6 text-gray-400">
+                Sign up for an account to access all features and start your
+                journey.
+              </p>
+              <Button className="text-sm font-medium" variant="secondary">
+                Sign up
+              </Button>
+            </div>
+
+            <div className="rounded-xl bg-[#1C1C20] p-8">
+              <div className="mb-4 w-fit rounded-full bg-white/15 px-3 py-1 text-sm font-medium text-white">
+                STEP 3
+              </div>
+              <h3 className="mb-4 text-2xl font-semibold">Start creating</h3>
+              <p className="mb-6 text-gray-400">
+                Begin creating with our realtime AI editor!
+              </p>
+            </div>
+          </div>
+        </div>
      </div>

      <Footer />
@@ -42,12 +42,6 @@ export const DialogAddImage = ({
    closeCallback();
  }

-  function handleEnter() {
-    if (currFile) {
-      onAddImage(currFile);
-    }
-    handleClose();
-  }
  return (
    <Dialog open={isOpen} onClose={closeCallback} className="relative z-50">
      <div className={dialogBackgroundStyles}>
@@ -66,6 +60,10 @@ export const DialogAddImage = ({
              file={currFile}
              setFile={(file: File | null) => {
                setAssetFile(file);
+                if (file) {
+                  onAddImage(file);
+                  handleClose();
+                }
              }}
            />
            {currFile && (
@@ -84,9 +82,6 @@ export const DialogAddImage = ({
            <Button onClick={handleClose} variant="secondary">
              Cancel
            </Button>
-            <Button onClick={handleEnter} disabled={currFile === null}>
-              Enter
-            </Button>
          </div>
        </DialogPanel>
      </div>
@@ -5,13 +5,17 @@ import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import {
  faArrowRotateLeft,
  faArrowRotateRight,
+  faCircle,
  faFilePlus,
  faFilm,
  faFloppyDisk,
  faImage,
  faLocationArrow,
  faMagicWandSparkles,
+  faShapes,
+  faSquare,
  faText,
+  faTriangle,
 } from "@fortawesome/pro-solid-svg-icons";

 import { ToolbarButton } from "../ToolbarButton";
@@ -41,7 +45,7 @@ export const ToolbarMain = ({
          "flex w-fit items-center divide-x divide-ui-border",
          paperWrapperStyles,
          disabled &&
-            "pointer-events-none cursor-default bg-ui-border shadow-md",
+          "pointer-events-none cursor-default bg-ui-border shadow-md",
        )}
      >
        <div className="flex items-center gap-2 pr-2">
@@ -56,6 +60,33 @@ export const ToolbarMain = ({
            buttonProps={buttonProps.ADD_TEXT}
            tooltip="Add Text"
          />
+
+          <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>
+          </Popover>
+
          <Popover className="relative">
            <PopoverButton as={Fragment}>
              <button
@@ -79,9 +110,9 @@ export const ToolbarMain = ({
              <ToolbarButton icon={faImage} buttonProps={buttonProps.ADD_IMAGE}>
                Add Image
              </ToolbarButton>
-              <ToolbarButton icon={faFilm} buttonProps={buttonProps.ADD_VIDEO}>
+              {/* <ToolbarButton icon={faFilm} buttonProps={buttonProps.ADD_VIDEO}>
                Add Video
-              </ToolbarButton>
+              </ToolbarButton> */}
            </PopoverPanel>
          </Popover>

@@ -4,6 +4,9 @@ export enum ToolbarMainButtonNames {
  ADD_TEXT = "ADD_TEXT",
  ADD_IMAGE = "ADD_IMAGE",
  ADD_VIDEO = "ADD_VIDEO",
+  ADD_CIRCLE = "ADD_CIRCLE",
+  ADD_TRIANGLE = "ADD_TRIANGLE",
+  ADD_SQUARE = "ADD_SQUARE",
  AI_STYLIZE = "AI_STYLIZE",
  UNDO = "UNDO",
  REDO = "REDO",
@@ -4,14 +4,19 @@ import {
  faLockKeyholeOpen,
  faLockKeyhole,
  faUnlockKeyhole,
+  faPalette,
 } from "@fortawesome/pro-solid-svg-icons";
 import {
  ToolbarButton,
  ToolbarButtonProps,
 } from "~/components/features/ToolbarButton";
-import { paperWrapperStyles } from "~/components/styles";
+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 {
  disabled?: boolean;
@@ -24,6 +29,8 @@ export interface ToolbarNodeProps {
  buttonsProps?: {
    [key in ToolbarNodeButtonNames]: ToolbarButtonProps;
  };
+  color?: string;
+  onColorChange?: (color: string) => void;
 }
 export const ToolbarNode = ({
  disabled,
@@ -31,6 +38,8 @@ export const ToolbarNode = ({
  lockDisabled,
  onLockClicked,
  buttonsProps,
+  color = "#000000",
+  onColorChange = () => { },
 }: ToolbarNodeProps) => {
  const handleOnLockClicked: React.MouseEventHandler<HTMLButtonElement> = (
    e,
@@ -53,7 +62,7 @@ export const ToolbarNode = ({
          className: twMerge(
            locked && "text-primary hover:bg-primary hover:text-white",
            lockDisabledOrUnknown &&
-              "text-secondary-300 hover:text-secondary-300",
+            "text-secondary-300 hover:text-secondary-300",
          ),
          disabled: lockDisabledOrUnknown,
        }}
@@ -73,14 +82,24 @@ export const ToolbarNode = ({
      {ToolbarNodeButtonData.map((buttonDatum, idx) => {
        const buttonProps = buttonsProps?.[buttonDatum.name];

-        return (
-          <ToolbarButton
-            icon={buttonDatum.icon}
-            tooltip={buttonDatum.tooltip}
-            key={idx}
-            buttonProps={buttonProps}
-          />
-        );
+        if (!buttonProps || buttonProps.hidden) {
+          return;
+        }
+
+        if (buttonDatum.name !== ToolbarNodeButtonNames.COLOR) {
+          return (
+            <ToolbarButton
+              icon={buttonDatum.icon}
+              tooltip={buttonDatum.tooltip}
+              key={idx}
+              buttonProps={buttonProps}
+            />
+          );
+        } else {
+          return (
+            <ColorPicker color={color} onChange={onColorChange} faIcon={faPalette} borderStyle="" key={idx} />
+          )
+        }
      })}
    </div>
  );
@@ -8,6 +8,7 @@ import {
  faTrashCan,
  faVectorSquare,
  faDownload,
+  faPalette,
 } from "@fortawesome/pro-solid-svg-icons";

 export const ToolbarNodeButtonData = [
@@ -16,6 +17,11 @@ export const ToolbarNodeButtonData = [
    icon: faVectorSquare,
    tooltip: "Move",
  },
+  {
+    name: ButtonNames.COLOR,
+    icon: faPalette,
+    tooltip: "Color",
+  },
  {
    name: ButtonNames.CHROMA,
    icon: faTransporter,
@@ -7,4 +7,5 @@ export enum ToolbarNodeButtonNames {
  MOVE_LAYER_UP = "MOVE_LAYER_UP",
  MOVE_LAYER_DOWN = "MOVE_LAYER_DOWN",
  DOWNLOAD = "DOWNLOAD",
+  COLOR = "COLOR",
 }
@@ -9,18 +9,18 @@ const Footer = () => {
  const navigate = useNavigate();

  return (
-    <footer className="border-t border-gray-100 bg-gray-50">
+    <footer className="relative z-10 bg-[#18181C]">
      <div className="mx-auto max-w-7xl px-4 py-16 sm:px-6 lg:px-8">
        <div className="flex flex-col items-center space-y-10 text-center">
          <img
-            src="/brand/Storyteller-Logo-Black.png"
-            alt="Storyteller Board"
-            className="h-9"
+            src="/brand/mira-logo.png"
+            alt="Mira Realtime AI Editor"
+            className="h-[34px] pb-1"
          />

          <Button
            onClick={() => navigate("/signup")}
-            className="group rounded-xl px-8 py-4 text-lg"
+            className="text-md group rounded-xl bg-[#2D81FF] px-8 py-4 hover:bg-[#438AF6]"
          >
            <span className="relative flex items-center">
              Start Creating for Free
@@ -3,26 +3,31 @@ import { Button } from "../ui";

 const Navbar = () => {
  const navigate = useNavigate();
-
  return (
-    <nav className="fixed top-0 z-50 w-full border-b border-gray-100 bg-white/80 backdrop-blur-md">
-      <div className="mx-auto max-w-7xl px-4 sm:px-6 lg:px-8">
-        <div className="flex h-20 items-center justify-between">
+    <nav className="fixed top-0 z-50 w-full bg-[rgba(16,16,20,0.7)] backdrop-blur-3xl">
+      <div className="mx-auto max-w-[1920px] px-4 sm:px-24 lg:px-32">
+        <div className="flex h-[72px] items-center justify-between">
          <a href="/" className="flex items-center">
            <img
-              src="/brand/Storyteller-Logo-Black.png"
-              alt="Storyteller Board"
-              className="h-9"
+              src="/brand/mira-logo.png"
+              alt="Mira Realtime AI Editor"
+              className="h-[34px] pb-1"
            />
          </a>
-          <div className="flex items-center space-x-5">
-            <a
-              href="/login"
-              className="font-medium !text-gray-600 hover:!text-gray-900"
+          <div className="flex items-center space-x-3">
+            <Button
+              variant="secondary"
+              className=""
+              onClick={() => navigate("/signup")}
            >
-              Login
-            </a>
-            <Button onClick={() => navigate("/signup")}>Sign Up</Button>
+              Sign up
+            </Button>
+            <Button
+              className="hidden bg-[#2D81FF] hover:bg-[#438AF6] md:block"
+              onClick={() => []}
+            >
+              Download
+            </Button>
          </div>
        </div>
      </div>
@@ -1,6 +1,6 @@
 import { ChangeEvent, useState } from "react";
 import { twMerge } from "tailwind-merge";
-import { faFont, faHashtag } 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";
@@ -12,9 +12,15 @@ import { paperWrapperStyles } from "~/components/styles";
 export const ColorPicker = ({
  color: prevColor,
  onChange,
+  faIcon = faFont,
+  borderStyle = "border",
+  showBar = true,
 }: {
  color: string;
  onChange: (newColor: string) => void;
+  faIcon?: IconDefinition;
+  borderStyle?: string;
+  showBar?: boolean;
 }) => {
  const [{ currColor, textInput }, setStates] = useState<{
    currColor: string;
@@ -44,11 +50,15 @@ export const ColorPicker = ({
      }));
    }
  };
+
  return (
    <Popover className="relative">
-      <PopoverButton className="flex size-10 flex-col items-center gap-1 rounded-md border p-2">
-        <FontAwesomeIcon icon={faFont} />
-        <span className="h-1 w-full" style={{ backgroundColor: prevColor }} />
+      <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"
@@ -0,0 +1,17 @@
+import { useEffect, useState } from "react";
+
+export const useIsMobile = () => {
+  const [isMobile, setIsMobile] = useState(false);
+
+  useEffect(() => {
+    const checkMobile = () => {
+      return /Android|webOS|iPhone|iPad|iPod|BlackBerry|IEMobile|Opera Mini/i.test(
+        navigator.userAgent,
+      );
+    };
+
+    setIsMobile(checkMobile());
+  }, []);
+
+  return isMobile;
+};
@@ -10,12 +10,13 @@ type ButtonStates = {
  };
 };
 export interface ContextualToolbarProps extends ContextualUi {
-  knodeIds: number[];
+  knodeIds: string[];
  disabled: boolean;
  downloadUrl?: string;
  locked: boolean | "unknown";
  lockDisabled: boolean;
  buttonStates: ButtonStates;
+  color: string;
 }
 interface PartialContextualToolbarProps
  extends Partial<Omit<ContextualToolbarProps, "buttonStates">> {
@@ -32,6 +33,7 @@ const toolbarNodeSignal = signal<ContextualToolbarProps>({
  locked: false,
  lockDisabled: false,
  buttonStates: initButtonStates(),
+  color: "#000000",
 });

 export const toolbarNode = {
@@ -77,6 +79,11 @@ export const toolbarNode = {
          hidden: false,
          active: buttonStates[ButtonNames.TRANSFORM].active,
        },
+        [ButtonNames.COLOR]: {
+          disabled: locked,
+          hidden: false,
+          active: buttonStates[ButtonNames.COLOR].active,
+        },
      },
      locked: locked,
    };
@@ -135,6 +142,12 @@ export const toolbarNode = {
      },
    };
  },
+  setColor(color: string) {
+    toolbarNodeSignal.value = {
+      ...toolbarNodeSignal.value,
+      color,
+    }
+  },
  batchButtonStates({ buttonStates }: { buttonStates: Partial<ButtonStates> }) {
    const prevButtonStates = toolbarNodeSignal.value.buttonStates;
    toolbarNodeSignal.value = {
@@ -1,11 +1,12 @@
 import { signal, effect } from "@preact/signals-react";
-import { TextNodeData } from "~/KonvaApp/types";
+import { ShapeNodeData, TextNodeData } from "~/KonvaApp/types";

 // ADDING IMAGES
 const stagedImage = signal<File | null>(null);
 const addImageToEngine = (image: File) => {
  stagedImage.value = image;
 };
+
 const onGetStagedImage = (callback: (file: File) => void) => {
  effect(() => {
    if (stagedImage.value) {
@@ -48,15 +49,30 @@ const onAddTextToEngine = (callback: (data: TextNodeData) => void) => {
  });
 };

+// Add shapes
+const stagedShape = signal<ShapeNodeData | null>(null);
+const addShapeToEngine = (data: ShapeNodeData) => {
+  stagedShape.value = data;
+}
+const onAddShapeToEngine = (callback: (data: ShapeNodeData) => void) => {
+  effect(() => {
+    if (stagedShape.value) {
+      callback(stagedShape.value);
+    }
+  });
+};
+
 //EXPORTS
 export const dispatchers = {
  addImageToEngine,
  addVideoToEngine,
  addTextToEngine,
+  addShapeToEngine,
 };

 export const events = {
  onGetStagedImage,
  onGetStagedVideo,
  onAddTextToEngine,
+  onAddShapeToEngine,
 };
@@ -44,6 +44,11 @@ const buttonEventsRecords = Object.values(ButtonNames).reduce(

 const buttonEventsHandlers = Object.values(ButtonNames).reduce(
  (acc, buttonName) => {
+    // We define custom events/handlers for color changes
+    if (buttonName === ButtonNames.COLOR) {
+      return acc;
+    }
+
    acc[buttonName] = {
      onClick: (callback: MouseEventHandler<HTMLButtonElement>) => {
        const { eventSignal, effectCleanup, lastEventTimestamp } =
@@ -111,9 +116,27 @@ const retryEventHandler = (callback: MouseEventHandler<HTMLButtonElement>) => {
 const retryDispatcher = (e: React.MouseEvent<HTMLButtonElement>) => {
  retryEvent.value = e;
 };
+
+export type NodeColor = {
+  kNodeId: string;
+  color: string;
+}
+const selectedColorEvent = signal<NodeColor | null>();
+const setSelectedColor = (nodeColor: NodeColor) => {
+  selectedColorEvent.value = nodeColor;
+}
+const onSelectedColorChanged = (callback: (color: NodeColor) => void) => {
+  effect(() => {
+    if (selectedColorEvent.value) {
+      callback(selectedColorEvent.value);
+    }
+  });
+};
+
 export const dispatchers = {
  lock: lockDispatcher,
  retry: retryDispatcher,
+  setSelectedColor,
  ...buttonDispatchers,
 };
 export const eventsHandlers = {
@@ -123,5 +146,8 @@ export const eventsHandlers = {
  retry: {
    onClick: retryEventHandler,
  },
+  color: {
+    onConfirmChanged: onSelectedColorChanged,
+  },
  ...buttonEventsHandlers,
 };
@@ -54,6 +54,10 @@ export default {
        border: colors.zinc[200],
      },
    },
+    backgroundImage: {
+      "radial-gradient":
+        "radial-gradient(circle at left, var(--tw-gradient-stops))",
+    },
  },
  plugins: [],
 };
Author	SHA1	Message	Date
Brandon Thomas	b6a9d52e19	implement cumprod	2025-02-25 00:25:19 -05:00
Brandon Thomas	97fe9f90c9	more work on timestamps (WIP)	2025-02-24 19:36:31 -05:00
Brandon Thomas	5cec883b15	cleanup	2025-02-24 17:48:24 -05:00
Brandon Thomas	70fa084e30	step() mostly done	2025-02-24 17:45:08 -05:00
Brandon Thomas	8111c84d12	implementing step()	2025-02-24 17:22:25 -05:00
Brandon Thomas	a55e3f84dc	notes	2025-02-24 02:07:48 -05:00
Brandon Thomas	221861eb9f	clarify	2025-02-23 23:08:14 -05:00
Brandon Thomas	e9e6c66ba1	add_noise WIP	2025-02-23 22:38:31 -05:00
Brandon Thomas	d2da91a8bd	implement constructor, tests	2025-02-23 20:53:19 -05:00
Brandon Thomas	ee5efc378e	implement timestep initialization	2025-02-23 14:56:49 -05:00
Brandon Thomas	dbd06ee7b3	LCM implementation (wip)	2025-02-23 14:37:24 -05:00
Michael Chung	6077232cd3	Hide Preview Drag Layer as it is not done.	2025-02-21 17:21:43 -05:00
Michael Chung	f013373180	Copy Shift Still needs some UX work but it does work. Need to store it as a image node.	2025-02-21 00:39:59 -05:00
Michael Chung	f0c1fcf516	Got the color changes to work.	2025-02-20 22:34:17 -05:00
Brandon Thomas	9f392de9c5	bump packages to enable windows builds	2025-02-20 18:10:07 -05:00
BFlat	f5fb340de9	landing hero update (#419 )	2025-02-21 01:49:28 +07:00
Tanish Manku	0392bda03c	Forward color picker events to engine (#418 )	2025-02-20 11:38:29 -05:00
Michael Chung	1bca478621	Move color picker into the correct toolbar	2025-02-19 23:56:45 -05:00
Michael Chung	f9c966003a	Fixed issue with shapes not having the context menu component.	2025-02-19 23:18:28 -05:00
Michael Chung	2d8245cd55	Added Shapes. Removed Extra buttons.	2025-02-19 22:29:31 -05:00
Tanish Manku	9d1d605f5c	Add shapes to global toolbar and color picker to contextual (#417 ) Shapes scaffolding is fully set up, color picker isn't tied up to the nodes so it doesn't actually change the color yet	2025-02-19 21:55:06 -05:00
Michael Chung	1e38fbad5c	Turns out you need to convert these rasters into images.	2025-02-19 20:33:01 -05:00
Michael Chung	6e03df968d	Remove duplicate adds? Then fixed layering issue by removing useless background layer. Need to fix shape rendering.	2025-02-19 19:34:06 -05:00
Michael Chung	13e0e4bb02	Merge branch 'main' of github.com:storytold/storyteller-rust	2025-02-19 18:00:26 -05:00
Michael Chung	442c3579e1	Fixed Modal to have the image enter the scene right away. Added Shape Fixes resizes properly and drags. Removing Progress Fill.	2025-02-19 17:59:14 -05:00
bflatastic	b8a305e271	mobile download button disable landing	2025-02-19 22:58:22 +07:00
BFlat	a70061bfd3	realtime product landing (#415 )	2025-02-19 20:44:22 +07:00
Brandon Thomas	e7a392b8e2	add cfg (#416 )	2025-02-19 00:21:24 -05:00
Michael Chung	9147da22f2	ShapeNodes work need to fix transformations.	2025-02-19 00:09:47 -05:00