super basic flux

2026-10-27 20:20:07 -04:00 · 2025-03-26 10:29:22 +05:30
22 changed files with 2096 additions and 1182 deletions
@@ -662,6 +662,7 @@ dependencies = [
 "base64 0.22.1",
 "bytes 1.10.0",
 "const_format",
+ "cudarc",
 "directories 6.0.0",
 "dirs 6.0.0",
 "filesys",
@@ -674,6 +675,7 @@ dependencies = [
 "ml_models",
 "ml_weights_registry",
 "once_cell",
+ "ort",
 "primitives",
 "rand 0.8.5",
 "reqwest 0.12.12",
@@ -7954,6 +7956,7 @@ dependencies = [
 "half",
 "lazy_static",
 "libc",
+ "libloading 0.7.4",
 "ndarray",
 "tar",
 "thiserror 1.0.63",
@@ -0,0 +1,288 @@
+use std::sync::{Arc, 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::prompt_cache::PromptCache;
+use candle_nn::VarBuilder;
+use crate::ml::model_config::ModelConfig;
+use anyhow::{anyhow, Error as E, Result};
+use candle_core::{DType, IndexOp, Tensor, D, Device, Module};
+use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
+use candle_transformers::models::stable_diffusion::lcm::LCMScheduler;
+use image::DynamicImage;
+use log::info;
+use rand::Rng;
+use tauri::{AppHandle, Emitter};
+use candle_transformers::models::{clip, t5, flux};
+use tokenizers::Tokenizer;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum FluxModel {
+  Schnell,
+  Dev,
+}
+
+pub struct FluxPipelineArgs<'a> {
+  pub image: &'a DynamicImage,
+  pub prompt: String,
+  pub model_cache: &'a ModelCache,
+  pub prompt_cache: &'a PromptCache,
+  pub app: &'a AppHandle,
+  pub flux_model: FluxModel,
+  pub active_blocks: Option<usize>,
+  pub use_cuda_stream: Option<bool>,
+  pub prefetch_next_batch: Option<bool>,
+  pub use_full_gpu: Option<bool>,
+  pub img2img_strength: Option<f64>,
+  pub seed: Option<u64>,
+  pub use_quantized_model: Option<bool>,
+  pub model_config: &'a ModelConfig,
+}
+
+pub fn flux_pipeline(args: FluxPipelineArgs<'_>) -> Result<RgbImage> {
+  let FluxPipelineArgs { prompt, model_cache, prompt_cache, app, image, flux_model, active_blocks, use_cuda_stream, prefetch_next_batch, use_full_gpu, img2img_strength, seed, use_quantized_model, model_config } = args;
+
+  if use_quantized_model.unwrap_or(false) {
+    #[cfg(feature = "cuda")]
+    candle_core::quantized::cuda::set_force_dmmv(false);
+  }
+
+  let height = 1024;
+  let width = 1024;
+
+  let quantized = use_quantized_model.unwrap_or(false);
+  let device = model_config.device.clone();
+  let dtype = device.bf16_default_to_f32();
+
+  println!("Starting image generation with FLUX model:");
+  println!("  Model: {:?}", flux_model);
+  println!("  Prompt: {}", prompt);
+  println!("  Dimensions: {}x{}", width, height);
+  println!("  Device: {:?}", model_config.device);
+
+  let seed_value = seed.unwrap_or_else(|| rand::thread_rng().gen());
+  device.set_seed(seed_value)?;
+  info!("Using seed: {}", seed_value);
+
+  // Setup FLUX model configuration
+  let flux_config = match flux_model {
+    FluxModel::Dev => flux::model::Config::dev(),
+    FluxModel::Schnell => flux::model::Config::schnell(),
+  };
+
+  let bf_repo = match flux_model {
+    FluxModel::Dev => model_config.hf_api.repo(hf_hub::Repo::model("black-forest-labs/FLUX.1-dev".to_string())),
+    FluxModel::Schnell => model_config.hf_api.repo(hf_hub::Repo::model("black-forest-labs/FLUX.1-schnell".to_string())),
+  };
+  // Get T5 embeddings
+  println!("Generating T5 embeddings for prompt...");
+  let t5_emb = {
+    let repo = model_config.hf_api.repo(hf_hub::Repo::with_revision("google/t5-v1_1-xxl".to_string(), hf_hub::RepoType::Model, "refs/pr/2".to_string()));
+    let model_file = repo.get("model.safetensors")?;
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype , &device)? };
+    let config_filename = repo.get("config.json")?;
+    let config = std::fs::read_to_string(config_filename)?;
+    let config: t5::Config = serde_json::from_str(&config)?;
+    let mut model = t5::T5EncoderModel::load(vb, &config)?;
+    let tokenizer_filename = model_config.hf_api.model("lmz/mt5-tokenizers".to_string()).get("t5-v1_1-xxl.tokenizer.json")?;
+    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+
+    let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
+    tokens.resize(256, 0);
+    let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+
+
+    model.forward(&input_token_ids)?
+  };
+  println!("T5 embeddings generated");
+
+  // Get CLIP embeddings
+  println!("Generating CLIP embeddings for prompt...");
+  let clip_emb = {
+    let repo = model_config.hf_api.repo(hf_hub::Repo::model(
+        "openai/clip-vit-large-patch14".to_string(),
+    ));
+    let model_file = repo.get("model.safetensors")?;
+    let vb =
+        unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+    // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
+    let config = clip::text_model::ClipTextConfig {
+        vocab_size: 49408,
+        projection_dim: 768,
+        activation: clip::text_model::Activation::QuickGelu,
+        intermediate_size: 3072,
+        embed_dim: 768,
+        max_position_embeddings: 77,
+        pad_with: None,
+        num_hidden_layers: 12,
+        num_attention_heads: 12,
+    };
+    let model =
+        clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
+    let tokenizer_filename = repo.get("tokenizer.json")?;
+    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+    let tokens = tokenizer
+        .encode(prompt.as_str(), true)
+        .map_err(E::msg)?
+        .get_ids()
+        .to_vec();
+    let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+
+
+    model.forward(&input_token_ids)?
+  };
+  println!("CLIP embeddings generated");
+
+  // Process input image to tensor
+  println!("Converting input image to tensor...");
+  let img_tensor = dynamic_image_to_tensor(image, &device, dtype)?;
+  println!("Image tensor shape: {:?}", img_tensor.shape());
+
+  // Create initial noise - this is important to match the expected format
+  println!("Generating initial noise for FLUX...");
+  let initial_tensor = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
+  println!("Initial tensor shape: {:?}", initial_tensor.shape());
+  
+  // Blend initial noise with input image if img2img_strength is specified
+  let initial_tensor = if let Some(strength) = img2img_strength {
+    if strength > 0.0 && strength < 1.0 {
+      println!("Using img2img with strength: {}", strength);
+      // Reshape the input tensor to match the noise tensor dimensions
+      let adjusted_img = if img_tensor.shape() != initial_tensor.shape() {
+        println!("Reshaping input image to match noise dimensions...");
+        // You may need more complex reshaping logic here depending on your use case
+        let (_, c, h, w) = initial_tensor.dims4()?;
+        img_tensor.upsample_nearest2d(h, w)?
+      } else {
+        img_tensor
+      };
+      
+      // Blend the noise and the image according to strength
+      let strength_tensor = Tensor::new(strength as f32, &device)?;
+      let inverse_strength = Tensor::new(1.0 - strength as f32, &device)?;
+      let weighted_noise = (&initial_tensor * &strength_tensor)?;
+      let weighted_img = (&adjusted_img * &inverse_strength)?;
+      (&weighted_noise + &weighted_img)?
+    } else {
+      initial_tensor
+    }
+  } else {
+    initial_tensor
+  };
+  println!("Final initial tensor shape: {:?}", initial_tensor.shape());
+
+  // Create FLUX state
+  println!("Creating FLUX state...");
+  let state = if quantized {
+    flux::sampling::State::new(
+        &t5_emb.to_dtype(DType::F32)?,
+        &clip_emb.to_dtype(DType::F32)?,
+        &initial_tensor.to_dtype(DType::F32)?,
+    )?
+  } else {
+    flux::sampling::State::new(&t5_emb, &clip_emb, &initial_tensor)?
+  };
+  println!("{state:?}");
+
+  // Set appropriate timesteps based on model
+  let timesteps = match flux_model {
+    FluxModel::Dev => flux::sampling::get_schedule(28, Some((state.img.dim(1)?, 0.5, 1.15))),
+    FluxModel::Schnell => flux::sampling::get_schedule(4, None),
+  };
+
+  println!("Starting FLUX denoising with timesteps: {:?}", timesteps);
+
+  // Actual model loading and inference
+  let img_latent = if quantized {
+    // Quantized model path
+    println!("Loading quantized FLUX model...");
+    let model_file = match flux_model {
+      FluxModel::Schnell => model_config.hf_api.repo(hf_hub::Repo::model("lmz/candle-flux".to_string())).get("flux1-schnell.gguf")?,
+      FluxModel::Dev => return Err(anyhow!("Quantized model not available for FLUX.1-dev")),
+    };
+
+    println!("Loading quantized model from {}", model_file.to_str().unwrap());
+
+    let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?;
+    let mut model = flux::quantized_model::Flux::new(&flux_config, vb)?;
+    println!("Quantized model loaded on {:?}", device);
+
+
+    // Run denoising
+    println!("Running denoising...");
+    let img_latent = flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?.to_dtype(dtype)?;
+    println!("Denoising completed");
+    img_latent
+  } else {
+    // Full precision model
+    println!("Loading FLUX model...");
+    let model_file = match flux_model {
+      FluxModel::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
+      FluxModel::Dev => bf_repo.get("flux1-dev.safetensors")?,
+    };
+
+
+    let mut model = if use_full_gpu.unwrap_or(false) {
+      // Load entire model on GPU
+      println!("Loading entire model directly to GPU...");
+      let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+      let mut model = flux::model::Flux::new(&flux_config, gpu_vb)?;
+      model.set_use_device_management(false);
+      model
+    } else {
+      // Use memory-efficient approach with blocks on CPU
+      println!("Using memory-efficient approach with blocks on CPU...");
+      let cpu_device = candle_core::Device::Cpu;
+      let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &cpu_device)? };
+      let mut model = flux::model::Flux::new_with_gpu_core(&flux_config, cpu_vb, &device)?;
+
+      // Set active blocks
+      let active_blocks_count = active_blocks.unwrap_or(1);
+      model.set_max_active_blocks(active_blocks_count);
+      println!("Using {} active blocks on GPU", active_blocks_count);
+
+      // Enable prefetching if requested
+      if prefetch_next_batch.unwrap_or(false) {
+        model.set_prefetch_next_batch(true);
+        println!("Prefetching enabled - will overlap transfers with computation");
+      }
+
+      model
+    };
+
+
+    // Run denoising
+    flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?
+  };
+
+  println!("FLUX denoising completed, unpacking latent image...");
+  let img_unpacked = flux::sampling::unpack(&img_latent, height, width)?;
+
+  // Load autoencoder for final decoding
+  println!("Loading autoencoder for final decoding...");
+  let ae_config = match flux_model {
+    FluxModel::Dev => flux::autoencoder::Config::dev(),
+    FluxModel::Schnell => flux::autoencoder::Config::schnell(),
+  };
+
+  let model_file = bf_repo.get("ae.safetensors")?;
+
+
+  let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+  let ae_model = flux::autoencoder::AutoEncoder::new(&ae_config, vb)?;
+
+
+  println!("Decoding final image...");
+  let img = ae_model.decode(&img_unpacked)?;
+
+  // Convert tensor to image
+  println!("Converting tensor to image...");
+  let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle_core::DType::U8)?;
+
+  // Convert to RGB image
+  let image = tensor_to_image_buffer(&img.i(0)?)?;
+
+  println!("FLUX pipeline completed successfully");
+  Ok(image)
+}
@@ -0,0 +1,3 @@
+pub mod flux_pipeline;
+
+pub use flux_pipeline::{flux_pipeline, FluxPipelineArgs, FluxModel};
@@ -4,4 +4,5 @@ pub mod model_cache;
 pub mod model_config;
 pub mod model_file;
 pub mod prompt_cache;
-pub mod stable_diffusion;
+pub mod stable_diffusion;
+pub mod flux;
@@ -29,6 +29,8 @@ primitives = { path = "../../lib/data/primitives" }
 # Internal ML
 # TODO: Only on Windows and Linux
 #ml_models = { path = "../ml_models", optional = true }
+# Bypass cudarc issues on Windows https://github.com/coreylowman/cudarc/issues/332
+cudarc = { version = "0.13.5", features = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc", "f16", "cuda-version-from-build-system", "dynamic-linking"], default-features=false }

 # External
 #openssl = { version = "0.10", features = ["vendored"] }
@@ -43,6 +45,7 @@ half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_di
 image = "0.25.5"
 log = "0.4"
 once_cell = "1.20.3"
+ort = { version = "1.16.2", features = ["cuda", "load-dynamic"] }
 rand = "0.8.5"
 # NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
 reqwest = { version = "0.12.12", features = ["stream"], default-features = false }
@@ -1,5 +1,5 @@
 #[cfg(not(target_os = "macos"))]
-use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
+use ml_models::ml::flux::flux_pipeline::{flux_pipeline, FluxPipelineArgs, FluxModel};

 use crate::state::app_config::AppConfig;
 use crate::state::app_dir::AppDataRoot;
@@ -84,31 +84,24 @@ async fn text_to_image_impl(
  app_data_root: &AppDataRoot,
 ) -> Result<RgbImage, String> {
  
-  let weights_dir = app_data_root.weights_dir();
-  let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
-  let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
-  let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
-  let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
-
-  let args = Args {
+  let args = FluxPipelineArgs {
    image: &image,
    prompt: prompt.to_string(),
-    uncond_prompt: "".to_string(),
    model_cache,
    prompt_cache: &prompt_cache,
+    app: &app,
+    flux_model: FluxModel::Schnell,
+    active_blocks: None,
+    use_cuda_stream: Some(true),
+    prefetch_next_batch: Some(false),
+    use_full_gpu: Some(true),
    img2img_strength: strength,
-    cfg_scale: config.cfg_scale,
-    use_flash_attn: true,
+    seed: None,
+    use_quantized_model: Some(true),
    model_config: &config.model_config,
-    maybe_seed: None,
-    scheduler_steps: config.scheduler_steps,
-    vae_path: &vae_path,
-    unet_path: &unet_path,
-    clip_json_path: &clip_json_path,
-    clip_weights_path: &clip_weights_path,
  };

-  let image = lcm_pipeline(args)
+  let image = flux_pipeline(args)
    .map_err(|err| format!("failure to encode image: {:?}", err))?;

  Ok(image)
@@ -8,6 +8,9 @@ pub enum NotificationModelType {
  TextTokenizer,
  Unet,
  Vae,
+  Json,
+  ClipEncoder,
+  FluxModel,
 }

 #[derive(Clone, Serialize)]
@@ -26,5 +29,17 @@ pub enum NotificationEvent<'a> {
  ModelDownloadComplete {
    model_name: &'a str,
    model_type: NotificationModelType,
-  }
+  },
+  ModelLoading {
+    model_name: String,
+    model_type: NotificationModelType,
+  },
+  ModelInference {
+    model_name: String,
+    model_type: NotificationModelType,
+  },
+  ModelInferenceComplete {
+    model_name: String,
+    model_type: NotificationModelType,
+  },
 }
@@ -75,6 +75,7 @@ pub fn run() {
      flip_image,
      infer_image,
      inpaint_image,
+      text_to_image,
      remove_background,
      text_to_image,
    ])
@@ -15,252 +15,281 @@ use tokenizers::Tokenizer;
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Args {
-    /// The prompt to be used for image generation.
-    #[arg(long, default_value = "A rusty robot walking on a beach")]
-    prompt: String,
+  /// The prompt to be used for image generation.
+  #[arg(long, default_value = "A rusty robot walking on a beach")]
+  prompt: String,

-    /// Run on CPU rather than on GPU.
-    #[arg(long)]
-    cpu: bool,
+  /// Run on CPU rather than on GPU.
+  #[arg(long)]
+  cpu: bool,

-    /// Use the quantized model.
-    #[arg(long)]
-    quantized: bool,
+  /// Use the quantized model.
+  #[arg(long)]
+  quantized: bool,

-    /// Enable tracing (generates a trace-timestamp.json file).
-    #[arg(long)]
-    tracing: bool,
+  /// Enable tracing (generates a trace-timestamp.json file).
+  #[arg(long)]
+  tracing: bool,

-    /// The height in pixels of the generated image.
-    #[arg(long)]
-    height: Option<usize>,
+  /// The height in pixels of the generated image.
+  #[arg(long)]
+  height: Option<usize>,

-    /// The width in pixels of the generated image.
-    #[arg(long)]
-    width: Option<usize>,
+  /// The width in pixels of the generated image.
+  #[arg(long)]
+  width: Option<usize>,

-    #[arg(long)]
-    decode_only: Option<String>,
+  #[arg(long)]
+  decode_only: Option<String>,

-    #[arg(long, value_enum, default_value = "schnell")]
-    model: Model,
+  #[arg(long, value_enum, default_value = "schnell")]
+  model: Model,

-    /// Use the slower kernels.
-    #[arg(long)]
-    use_dmmv: bool,
+  /// Use the slower kernels.
+  #[arg(long)]
+  use_dmmv: bool,

-    /// The seed to use when generating random samples.
-    #[arg(long)]
-    seed: Option<u64>,
+  /// The seed to use when generating random samples.
+  #[arg(long)]
+  seed: Option<u64>,
+
+  #[arg(long)]
+  cache_path: Option<String>,
+
+  /// Number of blocks to keep active on GPU (default: 1).
+  /// Only relevant when NOT using --use_full_gpu.
+  /// Higher values use more memory but may improve performance.
+  #[arg(long, default_value = "1")]
+  active_blocks: usize,
+
+  /// Use CUDA stream for faster device transfers (default: true).
+  /// Only relevant when running on GPU.
+  #[arg(long, default_value = "true")]
+  use_cuda_stream: bool,
+
+  /// Enable prefetching of the next batch of blocks while processing the current batch.
+  /// Only relevant when NOT using --use_full_gpu.
+  /// This can improve performance by overlapping memory transfers and computation.
+  #[arg(long, default_value = "false")]
+  prefetch_next_batch: bool,
+
+  /// Load entire model on GPU instead of using memory-efficient active blocks approach.
+  /// This uses more GPU memory but avoids CPU-GPU transfers completely.
+  /// When this flag is enabled, the --active_blocks and --prefetch_next_batch flags have no effect.
+  #[arg(long, default_value = "false")]
+  use_full_gpu: bool,
 }

 #[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
 enum Model {
-    Schnell,
-    Dev,
+  Schnell,
+  Dev,
 }

 fn run(args: Args) -> Result<()> {
-    use tracing_chrome::ChromeLayerBuilder;
-    use tracing_subscriber::prelude::*;
+  use tracing_chrome::ChromeLayerBuilder;
+  use tracing_subscriber::prelude::*;

-    let Args {
-        prompt,
-        cpu,
-        height,
-        width,
-        tracing,
-        decode_only,
-        model,
-        quantized,
-        ..
-    } = args;
-    let width = width.unwrap_or(1360);
-    let height = height.unwrap_or(768);
+  let Args { prompt, cpu, height, width, tracing, decode_only, model, quantized, active_blocks, use_cuda_stream, prefetch_next_batch, use_full_gpu, .. } = args;
+  let width = width.unwrap_or(1360);
+  let height = height.unwrap_or(768);

-    let _guard = if tracing {
-        let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
-        tracing_subscriber::registry().with(chrome_layer).init();
-        Some(guard)
-    } else {
-        None
+  let _guard = if tracing {
+    let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
+    tracing_subscriber::registry().with(chrome_layer).init();
+    Some(guard)
+  } else {
+    None
+  };
+
+  let api = match args.cache_path.as_ref() {
+    Some(path) => hf_hub::api::sync::ApiBuilder::from_cache(hf_hub::Cache::new(path.to_string().into())).build().map_err(anyhow::Error::msg)?,
+    None => {
+      let cache = hf_hub::Cache::from_env();
+      hf_hub::api::sync::ApiBuilder::from_cache(cache).build().map_err(anyhow::Error::msg)?
+    },
+  };
+  let bf_repo = {
+    let name = match model {
+      Model::Dev => "black-forest-labs/FLUX.1-dev",
+      Model::Schnell => "black-forest-labs/FLUX.1-schnell",
    };
+    api.repo(hf_hub::Repo::model(name.to_string()))
+  };

-    let api = hf_hub::api::sync::Api::new()?;
-    let bf_repo = {
-        let name = match model {
-            Model::Dev => "black-forest-labs/FLUX.1-dev",
-            Model::Schnell => "black-forest-labs/FLUX.1-schnell",
-        };
-        api.repo(hf_hub::Repo::model(name.to_string()))
-    };
-    let device = candle_examples::device(cpu)?;
-    if let Some(seed) = args.seed {
-        device.set_seed(seed)?;
+  // Use CUDA with stream for faster device transfers when not in CPU mode
+  let device = if cpu {
+    candle::Device::Cpu
+  } else {
+    #[cfg(feature = "cuda")]
+    {
+      if use_cuda_stream {
+        println!("Using CUDA with dedicated stream for faster device transfers");
+        // Use device ordinal 0 (first GPU)
+        candle::Device::new_cuda_with_stream(0)?
+      } else {
+        println!("Using standard CUDA device (no dedicated stream)");
+        candle::Device::new_cuda(0)?
+      }
    }
-    let dtype = device.bf16_default_to_f32();
-    let img = match decode_only {
-        None => {
-            let t5_emb = {
-                let repo = api.repo(hf_hub::Repo::with_revision(
-                    "google/t5-v1_1-xxl".to_string(),
-                    hf_hub::RepoType::Model,
-                    "refs/pr/2".to_string(),
-                ));
-                let model_file = repo.get("model.safetensors")?;
-                let vb =
-                    unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
-                let config_filename = repo.get("config.json")?;
-                let config = std::fs::read_to_string(config_filename)?;
-                let config: t5::Config = serde_json::from_str(&config)?;
-                let mut model = t5::T5EncoderModel::load(vb, &config)?;
-                let tokenizer_filename = api
-                    .model("lmz/mt5-tokenizers".to_string())
-                    .get("t5-v1_1-xxl.tokenizer.json")?;
-                let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
-                let mut tokens = tokenizer
-                    .encode(prompt.as_str(), true)
-                    .map_err(E::msg)?
-                    .get_ids()
-                    .to_vec();
-                tokens.resize(256, 0);
-                let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
-                println!("{input_token_ids}");
-                model.forward(&input_token_ids)?
-            };
-            println!("T5\n{t5_emb}");
-            let clip_emb = {
-                let repo = api.repo(hf_hub::Repo::model(
-                    "openai/clip-vit-large-patch14".to_string(),
-                ));
-                let model_file = repo.get("model.safetensors")?;
-                let vb =
-                    unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
-                // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
-                let config = clip::text_model::ClipTextConfig {
-                    vocab_size: 49408,
-                    projection_dim: 768,
-                    activation: clip::text_model::Activation::QuickGelu,
-                    intermediate_size: 3072,
-                    embed_dim: 768,
-                    max_position_embeddings: 77,
-                    pad_with: None,
-                    num_hidden_layers: 12,
-                    num_attention_heads: 12,
-                };
-                let model =
-                    clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
-                let tokenizer_filename = repo.get("tokenizer.json")?;
-                let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
-                let tokens = tokenizer
-                    .encode(prompt.as_str(), true)
-                    .map_err(E::msg)?
-                    .get_ids()
-                    .to_vec();
-                let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
-                println!("{input_token_ids}");
-                model.forward(&input_token_ids)?
-            };
-            println!("CLIP\n{clip_emb}");
-            let img = {
-                let cfg = match model {
-                    Model::Dev => flux::model::Config::dev(),
-                    Model::Schnell => flux::model::Config::schnell(),
-                };
-                let img = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
-                let state = if quantized {
-                    flux::sampling::State::new(
-                        &t5_emb.to_dtype(candle::DType::F32)?,
-                        &clip_emb.to_dtype(candle::DType::F32)?,
-                        &img.to_dtype(candle::DType::F32)?,
-                    )?
-                } else {
-                    flux::sampling::State::new(&t5_emb, &clip_emb, &img)?
-                };
-                let timesteps = match model {
-                    Model::Dev => {
-                        flux::sampling::get_schedule(50, Some((state.img.dim(1)?, 0.5, 1.15)))
-                    }
-                    Model::Schnell => flux::sampling::get_schedule(4, None),
-                };
-                println!("{state:?}");
-                println!("{timesteps:?}");
-                if quantized {
-                    let model_file = match model {
-                        Model::Schnell => api
-                            .repo(hf_hub::Repo::model("lmz/candle-flux".to_string()))
-                            .get("flux1-schnell.gguf")?,
-                        Model::Dev => todo!(),
-                    };
-                    let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(
-                        model_file, &device,
-                    )?;
+    #[cfg(not(feature = "cuda"))]
+    {
+      println!("CUDA not available, using CPU");
+      candle::Device::Cpu
+    }
+  };

-                    let model = flux::quantized_model::Flux::new(&cfg, vb)?;
-                    flux::sampling::denoise(
-                        &model,
-                        &state.img,
-                        &state.img_ids,
-                        &state.txt,
-                        &state.txt_ids,
-                        &state.vec,
-                        &timesteps,
-                        4.,
-                    )?
-                    .to_dtype(dtype)?
-                } else {
-                    let model_file = match model {
-                        Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
-                        Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
-                    };
-                    let vb = unsafe {
-                        VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)?
-                    };
-                    let model = flux::model::Flux::new(&cfg, vb)?;
-                    flux::sampling::denoise(
-                        &model,
-                        &state.img,
-                        &state.img_ids,
-                        &state.txt,
-                        &state.txt_ids,
-                        &state.vec,
-                        &timesteps,
-                        4.,
-                    )?
-                }
-            };
-            flux::sampling::unpack(&img, height, width)?
-        }
-        Some(file) => {
-            let mut st = candle::safetensors::load(file, &device)?;
-            st.remove("img").unwrap().to_dtype(dtype)?
-        }
-    };
-    println!("latent img\n{img}");
-
-    let img = {
-        let model_file = bf_repo.get("ae.safetensors")?;
+  if let Some(seed) = args.seed {
+    device.set_seed(seed)?;
+  }
+  let start_time = std::time::Instant::now();
+  let dtype = device.bf16_default_to_f32();
+  let img = match decode_only {
+    None => {
+      let t5_emb = {
+        let repo = api.repo(hf_hub::Repo::with_revision("google/t5-v1_1-xxl".to_string(), hf_hub::RepoType::Model, "refs/pr/2".to_string()));
+        let model_file = repo.get("model.safetensors")?;
        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+        let config_filename = repo.get("config.json")?;
+        let config = std::fs::read_to_string(config_filename)?;
+        let config: t5::Config = serde_json::from_str(&config)?;
+        let mut model = t5::T5EncoderModel::load(vb, &config)?;
+        let tokenizer_filename = api.model("lmz/mt5-tokenizers".to_string()).get("t5-v1_1-xxl.tokenizer.json")?;
+        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+        let mut tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
+        tokens.resize(256, 0);
+        let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+        println!("{input_token_ids}");
+        model.forward(&input_token_ids)?
+      };
+      println!("T5\n{t5_emb}");
+      let clip_emb = {
+        let repo = api.repo(hf_hub::Repo::model("openai/clip-vit-large-patch14".to_string()));
+        let model_file = repo.get("model.safetensors")?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+        // https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
+        let config = clip::text_model::ClipTextConfig { vocab_size: 49408, projection_dim: 768, activation: clip::text_model::Activation::QuickGelu, intermediate_size: 3072, embed_dim: 768, max_position_embeddings: 77, pad_with: None, num_hidden_layers: 12, num_attention_heads: 12 };
+        let model = clip::text_model::ClipTextTransformer::new(vb.pp("text_model"), &config)?;
+        let tokenizer_filename = repo.get("tokenizer.json")?;
+        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+        let tokens = tokenizer.encode(prompt.as_str(), true).map_err(E::msg)?.get_ids().to_vec();
+        let input_token_ids = Tensor::new(&tokens[..], &device)?.unsqueeze(0)?;
+        println!("{input_token_ids}");
+        model.forward(&input_token_ids)?
+      };
+      println!("CLIP\n{clip_emb}");
+      let img = {
        let cfg = match model {
-            Model::Dev => flux::autoencoder::Config::dev(),
-            Model::Schnell => flux::autoencoder::Config::schnell(),
+          Model::Dev => flux::model::Config::dev(),
+          Model::Schnell => flux::model::Config::schnell(),
        };
-        let model = flux::autoencoder::AutoEncoder::new(&cfg, vb)?;
-        model.decode(&img)?
+        let img = flux::sampling::get_noise(1, height, width, &device)?.to_dtype(dtype)?;
+        let state = if quantized { flux::sampling::State::new(&t5_emb.to_dtype(candle::DType::F32)?, &clip_emb.to_dtype(candle::DType::F32)?, &img.to_dtype(candle::DType::F32)?)? } else { flux::sampling::State::new(&t5_emb, &clip_emb, &img)? };
+        let timesteps = match model {
+          Model::Dev => flux::sampling::get_schedule(28, Some((state.img.dim(1)?, 0.5, 1.15))),
+          Model::Schnell => flux::sampling::get_schedule(4, None),
+        };
+        println!("{state:?}");
+        println!("{timesteps:?}");
+        if quantized {
+          // Note about quantized model memory management
+          let model_file = match model {
+            Model::Schnell => api.repo(hf_hub::Repo::model("lmz/candle-flux".to_string())).get("flux1-schnell.gguf")?,
+            Model::Dev => todo!(),
+          };
+          let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, &device)?;
+
+          let mut model = flux::quantized_model::Flux::new(&cfg, vb)?;
+          // Quantized models are always fully loaded to the selected device
+          if !cpu {
+            println!("Quantized model fully loaded to GPU with optimized memory layout");
+          }
+
+          flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?.to_dtype(dtype)?
+        } else {
+          let model_file = match model {
+            Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
+            Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
+          };
+
+          let mut model = if use_full_gpu {
+            // When use_full_gpu is true, load everything directly to GPU
+            println!("Loading entire model directly to GPU...");
+            let gpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+            let mut model = flux::model::Flux::new(&cfg, gpu_vb)?;
+            model.set_use_device_management(false);
+            println!("Model loaded entirely on GPU");
+            model
+          } else {
+            // Use memory-efficient approach with blocks on CPU
+            // Create a CPU device for initially loading all model weights
+            let cpu_device = candle::Device::Cpu;
+
+            // Load the entire model to CPU first (uses less memory)
+            println!("Loading model to CPU to minimize GPU memory usage...");
+            let cpu_vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &cpu_device)? };
+
+            // Create the model with CPU weights, then selectively move core components to GPU
+            let mut model = flux::model::Flux::new_with_gpu_core(&cfg, cpu_vb, &device)?;
+            println!("Model loaded (core components on GPU, blocks on CPU for reduced memory usage)");
+
+            // Set the number of active blocks to keep on GPU
+            model.set_max_active_blocks(active_blocks);
+            println!("Using {} active blocks on GPU", active_blocks);
+
+            // Enable prefetching if requested
+            if prefetch_next_batch {
+              model.set_prefetch_next_batch(true);
+              println!("Prefetching enabled - will overlap transfers with computation");
+            }
+
+            model
+          };
+
+          // Print final memory management approach being used
+          if use_full_gpu {
+            println!("Using full GPU memory - entire model on GPU");
+          } else {
+            println!("Using device management - blocks will be transferred between CPU and GPU as needed");
+          }
+
+          flux::sampling::denoise(&mut model, &state.img, &state.img_ids, &state.txt, &state.txt_ids, &state.vec, &timesteps, 4.)?
+        }
+      };
+      flux::sampling::unpack(&img, height, width)?
+    },
+    Some(file) => {
+      let mut st = candle::safetensors::load(file, &device)?;
+      st.remove("img").unwrap().to_dtype(dtype)?
+    },
+  };
+  println!("latent img\n{img}");
+
+  let end_time = std::time::Instant::now();
+  println!("Time taken to decode: {:?}", end_time.duration_since(start_time));
+  let img = {
+    let model_file = bf_repo.get("ae.safetensors")?;
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
+    let cfg = match model {
+      Model::Dev => flux::autoencoder::Config::dev(),
+      Model::Schnell => flux::autoencoder::Config::schnell(),
    };
-    println!("img\n{img}");
-    let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle::DType::U8)?;
-    let filename = match args.seed {
-        None => "out.jpg".to_string(),
-        Some(s) => format!("out-{s}.jpg"),
-    };
-    candle_examples::save_image(&img.i(0)?, filename)?;
-    Ok(())
+    let model = flux::autoencoder::AutoEncoder::new(&cfg, vb)?;
+    model.decode(&img)?
+  };
+  println!("img\n{img}");
+  let img = ((img.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(candle::DType::U8)?;
+  let filename = match args.seed {
+    None => "out.jpg".to_string(),
+    Some(s) => format!("out-{s}.jpg"),
+  };
+  candle_examples::save_image(&img.i(0)?, filename)?;
+  Ok(())
 }

 fn main() -> Result<()> {
-    let args = Args::parse();
-    #[cfg(feature = "cuda")]
-    candle::quantized::cuda::set_force_dmmv(args.use_dmmv);
-    run(args)
+  let args = Args::parse();
+  #[cfg(feature = "cuda")]
+  candle::quantized::cuda::set_force_dmmv(args.use_dmmv);
+  run(args)
 }
@@ -28,7 +28,8 @@
 //! ```
 //!
 //! [`Layer Normalization`]: https://arxiv.org/abs/1607.06450
-use candle::{DType, Module, Result, Tensor, D};
+use candle::{DType, Module, Result, Tensor, D, Device};
+use crate::offload::DeviceTransferable;

 #[derive(Debug, Clone, Copy, PartialEq)]
 pub struct LayerNormConfig {
@@ -201,3 +202,20 @@ pub fn rms_norm(size: usize, eps: f64, vb: crate::VarBuilder) -> Result<RmsNorm>
    };
    Ok(RmsNorm(layer_norm(size, config, vb)?))
 }
+
+impl DeviceTransferable for LayerNorm {
+    fn to_device(&self, device: &Device) -> Result<Self> {
+        Ok(Self {
+            weight: self.weight.to_device(device)?,
+            bias: self.bias.as_ref().map(|b| b.to_device(device)).transpose()?,
+            remove_mean: self.remove_mean,
+            eps: self.eps,
+        })
+    }
+}
+
+impl DeviceTransferable for RmsNorm {
+    fn to_device(&self, device: &Device) -> Result<Self> {
+        Ok(Self(self.0.to_device(device)?))
+    }
+}
@@ -27,6 +27,7 @@ pub mod kv_cache;
 pub mod layer_norm;
 pub mod linear;
 pub mod loss;
+pub mod offload;
 pub mod ops;
 pub mod optim;
 pub mod rnn;
@@ -58,3 +59,4 @@ pub use var_builder::VarBuilder;
 pub use var_map::VarMap;

 pub use candle::{Module, ModuleT};
+pub use offload::DeviceTransferable;
@@ -17,7 +17,9 @@
 //! assert_eq!(ys.to_vec2::<f32>()?, &[[210.0, 430.0, 650.0]]);
 //! # Ok(()) }
 //! ```
-use candle::{Result, Tensor};
+use candle::{Result, Tensor, Device};
+use crate::offload::DeviceTransferable;
+

 #[derive(Clone, Debug)]
 pub struct Linear {
@@ -88,3 +90,12 @@ pub fn linear_b(
        linear_no_bias(in_dim, out_dim, vb)
    }
 }
+
+impl DeviceTransferable for Linear {
+    fn to_device(&self, device: &Device) -> Result<Self> {
+        Ok(Self {
+            weight: self.weight.to_device(device)?,
+            bias: self.bias.as_ref().map(|b| b.to_device(device)).transpose()?,
+        })
+    }
+}
@@ -0,0 +1,18 @@
+use candle::{Device, Result};
+
+pub trait DeviceTransferable {
+    /// Move the module to the specified device
+    ///
+    /// # Arguments
+    ///
+    /// * `device` - The target device to move the module to
+    ///
+    /// # Returns
+    ///
+    /// A new instance of the module with all tensors moved to the target device,
+    /// or an error if the module cannot be moved.
+    fn to_device(&self, _device: &Device) -> Result<Self> where Self: Sized {
+        // Default implementation that rejects moving - components can override this
+        candle::bail!("to_device not implemented for this type")
+    }
+}
@@ -24,9 +24,39 @@
 use candle::{Result, Tensor};

 pub trait WithForward {
+    /// The forward pass for the Flux model.
+    ///
+    /// # Arguments
+    ///
+    /// * `img` - The input image tensor
+    /// * `img_ids` - The image position ids
+    /// * `txt` - The text tensor
+    /// * `txt_ids` - The text position ids
+    /// * `timesteps` - The timesteps tensor
+    /// * `y` - The conditioning vector
+    /// * `guidance` - An optional guidance vector
+    ///
+    /// # Returns
+    ///
+    /// The result of the forward pass
+    ///
+    /// # Note on Memory Optimization
+    ///
+    /// To optimize memory usage, we could implement a mechanism to move model blocks 
+    /// to CPU when not in use. This would require:
+    ///
+    /// 1. Adding `to_device` methods to all model blocks (DoubleStreamBlock, SingleStreamBlock, etc.)
+    /// 2. Implementing these methods by moving all tensors inside each block to the target device
+    /// 3. Modifying the forward method to:
+    ///    - Move each block to GPU when needed
+    ///    - Process the block
+    ///    - Move the block back to CPU when done
+    ///
+    /// This pattern could significantly reduce GPU memory usage for large models by
+    /// keeping only the currently active block in GPU memory.
    #[allow(clippy::too_many_arguments)]
    fn forward(
-        &self,
+        &mut self,
        img: &Tensor,
        img_ids: &Tensor,
        txt: &Tensor,
@@ -419,7 +419,7 @@ impl Flux {
 impl super::WithForward for Flux {
    #[allow(clippy::too_many_arguments)]
    fn forward(
-        &self,
+        &mut self,
        img: &Tensor,
        img_ids: &Tensor,
        txt: &Tensor,
@@ -93,7 +93,7 @@ pub fn unpack(xs: &Tensor, height: usize, width: usize) -> Result<Tensor> {

 #[allow(clippy::too_many_arguments)]
 pub fn denoise<M: super::WithForward>(
-    model: &M,
+    model: &mut M,
    img: &Tensor,
    img_ids: &Tensor,
    txt: &Tensor,
@@ -35,7 +35,7 @@
        "@fortawesome/pro-solid-svg-icons": "^6.6.0",
        "@fortawesome/react-fontawesome": "^0.2.2",
        "@preact/signals-react-transform": "^0.4.0",
-        "@tauri-apps/cli": "^2.2.7",
+        "@tauri-apps/cli": "^2.4.0",
        "@types/deep-equal": "^1.0.4",
        "@types/react": "^18.3.3",
        "@types/react-dom": "^18.3.0",
@@ -1259,11 +1259,10 @@
      }
    },
    "node_modules/@tauri-apps/cli": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli/-/cli-2.2.7.tgz",
-      "integrity": "sha512-ZnsS2B4BplwXP37celanNANiIy8TCYhvg5RT09n72uR/o+navFZtGpFSqljV8fy1Y4ixIPds8FrGSXJCN2BerA==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli/-/cli-2.4.0.tgz",
+      "integrity": "sha512-Esg7s20tuSULd2YF3lmtMa1vF7yr5eh/TlBHXjEyrC+XSD9aBxHVoXb6oz7oKybDY9Jf9OiBa0bf2PbybcmOLA==",
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "bin": {
        "tauri": "tauri.js"
      },
@@ -1275,27 +1274,27 @@
        "url": "https://opencollective.com/tauri"
      },
      "optionalDependencies": {
-        "@tauri-apps/cli-darwin-arm64": "2.2.7",
-        "@tauri-apps/cli-darwin-x64": "2.2.7",
-        "@tauri-apps/cli-linux-arm-gnueabihf": "2.2.7",
-        "@tauri-apps/cli-linux-arm64-gnu": "2.2.7",
-        "@tauri-apps/cli-linux-arm64-musl": "2.2.7",
-        "@tauri-apps/cli-linux-x64-gnu": "2.2.7",
-        "@tauri-apps/cli-linux-x64-musl": "2.2.7",
-        "@tauri-apps/cli-win32-arm64-msvc": "2.2.7",
-        "@tauri-apps/cli-win32-ia32-msvc": "2.2.7",
-        "@tauri-apps/cli-win32-x64-msvc": "2.2.7"
+        "@tauri-apps/cli-darwin-arm64": "2.4.0",
+        "@tauri-apps/cli-darwin-x64": "2.4.0",
+        "@tauri-apps/cli-linux-arm-gnueabihf": "2.4.0",
+        "@tauri-apps/cli-linux-arm64-gnu": "2.4.0",
+        "@tauri-apps/cli-linux-arm64-musl": "2.4.0",
+        "@tauri-apps/cli-linux-riscv64-gnu": "2.4.0",
+        "@tauri-apps/cli-linux-x64-gnu": "2.4.0",
+        "@tauri-apps/cli-linux-x64-musl": "2.4.0",
+        "@tauri-apps/cli-win32-arm64-msvc": "2.4.0",
+        "@tauri-apps/cli-win32-ia32-msvc": "2.4.0",
+        "@tauri-apps/cli-win32-x64-msvc": "2.4.0"
      }
    },
    "node_modules/@tauri-apps/cli-darwin-arm64": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.2.7.tgz",
-      "integrity": "sha512-54kcpxZ3X1Rq+pPTzk3iIcjEVY4yv493uRx/80rLoAA95vAC0c//31Whz75UVddDjJfZvXlXZ3uSZ+bnCOnt0A==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.4.0.tgz",
+      "integrity": "sha512-MVzYrahJBgDyzUJ2gNEU8H+0oCVEucN115+CVorFnidHcJ6DtDRMCaKLkpjOZNfJyag1WQ25fu7imvZSe0mz/g==",
      "cpu": [
        "arm64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "darwin"
@@ -1305,14 +1304,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-darwin-x64": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-x64/-/cli-darwin-x64-2.2.7.tgz",
-      "integrity": "sha512-Vgu2XtBWemLnarB+6LqQeLanDlRj7CeFN//H8bVVdjbNzxcSxsvbLYMBP8+3boa7eBnjDrqMImRySSgL6IrwTw==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-x64/-/cli-darwin-x64-2.4.0.tgz",
+      "integrity": "sha512-/4IdbWv6IWSuBn0WVe5JkiSIP1gZhXCZRcumSsYq3ZmOlq4BqXeXT36Oig5mlDnS/2/UpNS94kd8gOA1DNdIeQ==",
      "cpu": [
        "x64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "darwin"
@@ -1322,14 +1320,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-linux-arm-gnueabihf": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm-gnueabihf/-/cli-linux-arm-gnueabihf-2.2.7.tgz",
-      "integrity": "sha512-+Clha2iQAiK9zoY/KKW0KLHkR0k36O78YLx5Sl98tWkwI3OBZFg5H5WT1plH/4sbZIS2aLFN6dw58/JlY9Bu/g==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm-gnueabihf/-/cli-linux-arm-gnueabihf-2.4.0.tgz",
+      "integrity": "sha512-rOjlk3Vd6R847LXds4pOAFKUL5NVdSWlaiQvr4H9WDUwIWWoxnj7SQfpryTtElDb2wV7a0BNaOCXCpyFEAXjkw==",
      "cpu": [
        "arm"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "linux"
@@ -1339,14 +1336,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-linux-arm64-gnu": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-gnu/-/cli-linux-arm64-gnu-2.2.7.tgz",
-      "integrity": "sha512-Z/Lp4SQe6BUEOays9BQAEum2pvZF4w9igyXijP+WbkOejZx4cDvarFJ5qXrqSLmBh7vxrdZcLwoLk9U//+yQrg==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-gnu/-/cli-linux-arm64-gnu-2.4.0.tgz",
+      "integrity": "sha512-X/uCwao6R/weWT2y4f3JKJMeUsujo9H4nBMAv9RZhRsz93n9Amw9ETavAOP11pyhl57YkasvKTCRQN6FwsaoXg==",
      "cpu": [
        "arm64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "linux"
@@ -1356,14 +1352,29 @@
      }
    },
    "node_modules/@tauri-apps/cli-linux-arm64-musl": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-musl/-/cli-linux-arm64-musl-2.2.7.tgz",
-      "integrity": "sha512-+8HZ+txff/Y3YjAh80XcLXcX8kpGXVdr1P8AfjLHxHdS6QD4Md+acSxGTTNbplmHuBaSHJvuTvZf9tU1eDCTDg==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-musl/-/cli-linux-arm64-musl-2.4.0.tgz",
+      "integrity": "sha512-GhvQtrTjadW3eLSmfrSfybSqgJMZzUXC+0WqDzFovLug3a1a1go0m9QK9YGvYLkyEpTY5zRxLtwv+tbZXN4tZw==",
      "cpu": [
        "arm64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-riscv64-gnu": {
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-riscv64-gnu/-/cli-linux-riscv64-gnu-2.4.0.tgz",
+      "integrity": "sha512-NgeNihQ9uHS/ibMWLge5VA/BgsS/g8VPSVtCp8DSPyub3bBuCy79A8V+bzNKlMOiDVrqK4vQ//FS9kSxoJOtXw==",
+      "cpu": [
+        "riscv64"
+      ],
+      "dev": true,
      "optional": true,
      "os": [
        "linux"
@@ -1373,14 +1384,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-linux-x64-gnu": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.2.7.tgz",
-      "integrity": "sha512-ahlSnuCnUntblp9dG7/w5ZWZOdzRFi3zl0oScgt7GF4KNAOEa7duADsxPA4/FT2hLRa0SvpqtD4IYFvCxoVv3Q==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.4.0.tgz",
+      "integrity": "sha512-ebRmV2HLIVms1KlNNueQCp3OrXBv6cimU3mYEh5NbZ8dH88f2sF46dFCyPq8Qr/Zti4qPEaAArVG7RYFjfECPw==",
      "cpu": [
        "x64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "linux"
@@ -1390,14 +1400,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-linux-x64-musl": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.2.7.tgz",
-      "integrity": "sha512-+qKAWnJRSX+pjjRbKAQgTdFY8ecdcu8UdJ69i7wn3ZcRn2nMMzOO2LOMOTQV42B7/Q64D1pIpmZj9yblTMvadA==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.4.0.tgz",
+      "integrity": "sha512-FOp2cBFyq5LnUr3he95Z99PQm3nCSJv2GZNeH7UqmUpeHwdcYuhBERU7C+8VDJJPR98Q5fkcoV00Pc4nw0v5KQ==",
      "cpu": [
        "x64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "linux"
@@ -1407,14 +1416,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-win32-arm64-msvc": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-arm64-msvc/-/cli-win32-arm64-msvc-2.2.7.tgz",
-      "integrity": "sha512-aa86nRnrwT04u9D9fhf5JVssuAZlUCCc8AjqQjqODQjMd4BMA2+d4K9qBMpEG/1kVh95vZaNsLogjEaqSTTw4A==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-arm64-msvc/-/cli-win32-arm64-msvc-2.4.0.tgz",
+      "integrity": "sha512-SVf1wDagYsaFM+mpUYKmjNveKodcUSGPEM27WmMW4Enh6aXGzTJi4IYOE3GEFOJF1BpRNscslwE1Rd064kfpcg==",
      "cpu": [
        "arm64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "win32"
@@ -1424,14 +1432,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-win32-ia32-msvc": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-ia32-msvc/-/cli-win32-ia32-msvc-2.2.7.tgz",
-      "integrity": "sha512-EiJ5/25tLSQOSGvv+t6o3ZBfOTKB5S3vb+hHQuKbfmKdRF0XQu2YPdIi1CQw1DU97ZAE0Dq4frvnyYEKWgMzVQ==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-ia32-msvc/-/cli-win32-ia32-msvc-2.4.0.tgz",
+      "integrity": "sha512-j+fOFVeSSejk9hrUePY7bJuaYr+80xr+ftjXzxCj0CS0d2oSbq+lT8/zS514WemJk9e9yxUus+2ke/Ng17wkkQ==",
      "cpu": [
        "ia32"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "win32"
@@ -1441,14 +1448,13 @@
      }
    },
    "node_modules/@tauri-apps/cli-win32-x64-msvc": {
-      "version": "2.2.7",
-      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.2.7.tgz",
-      "integrity": "sha512-ZB8Kw90j8Ld+9tCWyD2fWCYfIrzbQohJ4DJSidNwbnehlZzP7wAz6Z3xjsvUdKtQ3ibtfoeTqVInzCCEpI+pWg==",
+      "version": "2.4.0",
+      "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.4.0.tgz",
+      "integrity": "sha512-nv84b3a8eI5Y7ksTLBKjjvtr9NOlAGGGo7OJbjKT+xZLdiPOZ0nJ2cT+4IdfnNAZ33pKJugAfuj1fBvQKeTy0w==",
      "cpu": [
        "x64"
      ],
      "dev": true,
-      "license": "Apache-2.0 OR MIT",
      "optional": true,
      "os": [
        "win32"
@@ -38,7 +38,7 @@
    "@fortawesome/pro-solid-svg-icons": "^6.6.0",
    "@fortawesome/react-fontawesome": "^0.2.2",
    "@preact/signals-react-transform": "^0.4.0",
-    "@tauri-apps/cli": "^2.2.7",
+    "@tauri-apps/cli": "^2.4.0",
    "@types/deep-equal": "^1.0.4",
    "@types/react": "^18.3.3",
    "@types/react-dom": "^18.3.0",
@@ -26,7 +26,7 @@ import {
  loadingProgress,
 } from "~/signals/uiEvents/loadingIndicator";
 import { ensureBase64Prefix } from "../EngineUtitlities/Base64Helpers";
-import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts";
+import { DecodeBase64ToImage } from "~/utilities/DecodeBase64ToImage.ts";

 interface ServerSettings {
  model_path: string;
@@ -65,10 +65,10 @@ export class RealTimeDrawEngine {
  private port: MessagePort | undefined;
  public captureCanvas: Konva.Rect;
  public backgroundRasterRect: Konva.Image;
-  
+

  public previewCanvas: Konva.Image;
-  
+


  public videoLoadingCanvas: VideoNode | undefined;
@@ -205,7 +205,7 @@ export class RealTimeDrawEngine {
      strokeWidth: 1,
      draggable: false,
    });
- 
+
    this.mediaLayerRef.add(this.captureCanvas);
    this.mediaLayerRef.add(this.backgroundRasterRect);
    this.mediaLayerRef.add(this.previewCanvas);
@@ -215,9 +215,9 @@ export class RealTimeDrawEngine {
    this.previewCanvas.setZIndex(1);
    // Add mouse events for preview canvas copying
    this.previewCopyListener();
-    
+
    this.listenToServerEvents();
-   
+
  }

  public async listenToServerEvents() {
@@ -1027,7 +1027,7 @@ export class RealTimeDrawEngine {
    const imageSource = new Image();
    imageSource.src = captureCanvasImage;

-    imageSource.onload = () => {      
+    imageSource.onload = () => {
      this.backgroundRasterRect.fill(color);
      this.backgroundRasterRect.image(imageSource)

@@ -824,31 +824,32 @@
  resolved "https://registry.npmjs.org/@tauri-apps/api/-/api-2.2.0.tgz"
  integrity sha512-R8epOeZl1eJEl603aUMIGb4RXlhPjpgxbGVEaqY+0G5JG9vzV/clNlzTeqc+NLYXVqXcn8mb4c5b9pJIUDEyAg==

-"@tauri-apps/cli-linux-x64-gnu@2.2.7":
-  version "2.2.7"
-  resolved "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.2.7.tgz"
-  integrity sha512-ahlSnuCnUntblp9dG7/w5ZWZOdzRFi3zl0oScgt7GF4KNAOEa7duADsxPA4/FT2hLRa0SvpqtD4IYFvCxoVv3Q==
+"@tauri-apps/cli-linux-x64-gnu@2.4.0":
+  version "2.4.0"
+  resolved "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.4.0.tgz"
+  integrity sha512-ebRmV2HLIVms1KlNNueQCp3OrXBv6cimU3mYEh5NbZ8dH88f2sF46dFCyPq8Qr/Zti4qPEaAArVG7RYFjfECPw==

-"@tauri-apps/cli-linux-x64-musl@2.2.7":
-  version "2.2.7"
-  resolved "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.2.7.tgz"
-  integrity sha512-+qKAWnJRSX+pjjRbKAQgTdFY8ecdcu8UdJ69i7wn3ZcRn2nMMzOO2LOMOTQV42B7/Q64D1pIpmZj9yblTMvadA==
+"@tauri-apps/cli-linux-x64-musl@2.4.0":
+  version "2.4.0"
+  resolved "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.4.0.tgz"
+  integrity sha512-FOp2cBFyq5LnUr3he95Z99PQm3nCSJv2GZNeH7UqmUpeHwdcYuhBERU7C+8VDJJPR98Q5fkcoV00Pc4nw0v5KQ==

-"@tauri-apps/cli@^2.2.7":
-  version "2.2.7"
-  resolved "https://registry.npmjs.org/@tauri-apps/cli/-/cli-2.2.7.tgz"
-  integrity sha512-ZnsS2B4BplwXP37celanNANiIy8TCYhvg5RT09n72uR/o+navFZtGpFSqljV8fy1Y4ixIPds8FrGSXJCN2BerA==
+"@tauri-apps/cli@^2.4.0":
+  version "2.4.0"
+  resolved "https://registry.npmjs.org/@tauri-apps/cli/-/cli-2.4.0.tgz"
+  integrity sha512-Esg7s20tuSULd2YF3lmtMa1vF7yr5eh/TlBHXjEyrC+XSD9aBxHVoXb6oz7oKybDY9Jf9OiBa0bf2PbybcmOLA==
  optionalDependencies:
-    "@tauri-apps/cli-darwin-arm64" "2.2.7"
-    "@tauri-apps/cli-darwin-x64" "2.2.7"
-    "@tauri-apps/cli-linux-arm-gnueabihf" "2.2.7"
-    "@tauri-apps/cli-linux-arm64-gnu" "2.2.7"
-    "@tauri-apps/cli-linux-arm64-musl" "2.2.7"
-    "@tauri-apps/cli-linux-x64-gnu" "2.2.7"
-    "@tauri-apps/cli-linux-x64-musl" "2.2.7"
-    "@tauri-apps/cli-win32-arm64-msvc" "2.2.7"
-    "@tauri-apps/cli-win32-ia32-msvc" "2.2.7"
-    "@tauri-apps/cli-win32-x64-msvc" "2.2.7"
+    "@tauri-apps/cli-darwin-arm64" "2.4.0"
+    "@tauri-apps/cli-darwin-x64" "2.4.0"
+    "@tauri-apps/cli-linux-arm-gnueabihf" "2.4.0"
+    "@tauri-apps/cli-linux-arm64-gnu" "2.4.0"
+    "@tauri-apps/cli-linux-arm64-musl" "2.4.0"
+    "@tauri-apps/cli-linux-riscv64-gnu" "2.4.0"
+    "@tauri-apps/cli-linux-x64-gnu" "2.4.0"
+    "@tauri-apps/cli-linux-x64-musl" "2.4.0"
+    "@tauri-apps/cli-win32-arm64-msvc" "2.4.0"
+    "@tauri-apps/cli-win32-ia32-msvc" "2.4.0"
+    "@tauri-apps/cli-win32-x64-msvc" "2.4.0"

 "@techstark/opencv-js@^4.10.0-release.1":
  version "4.10.0-release.1"