diff --git a/Cargo.lock b/Cargo.lock
index 604b7f7233..2f435255d3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -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",
diff --git a/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs b/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs
new file mode 100644
index 0000000000..b7e9c84807
--- /dev/null
+++ b/crates/desktop/ml_models/src/ml/flux/flux_pipeline.rs
@@ -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)
+}
diff --git a/crates/desktop/ml_models/src/ml/flux/mod.rs b/crates/desktop/ml_models/src/ml/flux/mod.rs
new file mode 100644
index 0000000000..ededd1ec2f
--- /dev/null
+++ b/crates/desktop/ml_models/src/ml/flux/mod.rs
@@ -0,0 +1,3 @@
+pub mod flux_pipeline;
+
+pub use flux_pipeline::{flux_pipeline, FluxPipelineArgs, FluxModel};
\ No newline at end of file
diff --git a/crates/desktop/ml_models/src/ml/mod.rs b/crates/desktop/ml_models/src/ml/mod.rs
index c889cdbae0..67908b8f65 100644
--- a/crates/desktop/ml_models/src/ml/mod.rs
+++ b/crates/desktop/ml_models/src/ml/mod.rs
@@ -4,4 +4,5 @@ pub mod model_cache;
 pub mod model_config;
 pub mod model_file;
 pub mod prompt_cache;
-pub mod stable_diffusion;
\ No newline at end of file
+pub mod stable_diffusion;
+pub mod flux;
\ No newline at end of file
diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml
index 9426e2d526..6ffaa00c3a 100644
--- a/crates/desktop/src-tauri/Cargo.toml
+++ b/crates/desktop/src-tauri/Cargo.toml
@@ -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 }
diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs
index 5ce4e18f24..fe798720b0 100644
--- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs
+++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs
@@ -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)
diff --git a/crates/desktop/src-tauri/src/events/notification_event.rs b/crates/desktop/src-tauri/src/events/notification_event.rs
index bacbcd73c2..5435c91850 100644
--- a/crates/desktop/src-tauri/src/events/notification_event.rs
+++ b/crates/desktop/src-tauri/src/events/notification_event.rs
@@ -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,
+  },
 }
diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs
index beb9a213ed..5252c940bd 100644
--- a/crates/desktop/src-tauri/src/lib.rs
+++ b/crates/desktop/src-tauri/src/lib.rs
@@ -75,6 +75,7 @@ pub fn run() {
       flip_image,
       infer_image,
       inpaint_image,
+      text_to_image,
       remove_background,
       text_to_image,
     ])
diff --git a/crates/vendor/candle/candle-examples/examples/flux/main.rs b/crates/vendor/candle/candle-examples/examples/flux/main.rs
index 12439892ce..a106012900 100644
--- a/crates/vendor/candle/candle-examples/examples/flux/main.rs
+++ b/crates/vendor/candle/candle-examples/examples/flux/main.rs
@@ -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)
 }
diff --git a/crates/vendor/candle/candle-nn/src/layer_norm.rs b/crates/vendor/candle/candle-nn/src/layer_norm.rs
index 468fe24d26..2385d5deab 100644
--- a/crates/vendor/candle/candle-nn/src/layer_norm.rs
+++ b/crates/vendor/candle/candle-nn/src/layer_norm.rs
@@ -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)?))
+    }
+}
\ No newline at end of file
diff --git a/crates/vendor/candle/candle-nn/src/lib.rs b/crates/vendor/candle/candle-nn/src/lib.rs
index 2113566d33..8a72b86017 100644
--- a/crates/vendor/candle/candle-nn/src/lib.rs
+++ b/crates/vendor/candle/candle-nn/src/lib.rs
@@ -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;
diff --git a/crates/vendor/candle/candle-nn/src/linear.rs b/crates/vendor/candle/candle-nn/src/linear.rs
index 96409042f4..a85beb04df 100644
--- a/crates/vendor/candle/candle-nn/src/linear.rs
+++ b/crates/vendor/candle/candle-nn/src/linear.rs
@@ -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()?,
+        })
+    }
+}
\ No newline at end of file
diff --git a/crates/vendor/candle/candle-nn/src/offload.rs b/crates/vendor/candle/candle-nn/src/offload.rs
new file mode 100644
index 0000000000..ea66adc3e6
--- /dev/null
+++ b/crates/vendor/candle/candle-nn/src/offload.rs
@@ -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")
+    }
+}
diff --git a/crates/vendor/candle/candle-transformers/src/models/flux/mod.rs b/crates/vendor/candle/candle-transformers/src/models/flux/mod.rs
index 1d2fa4ef33..cc8f45ffe8 100644
--- a/crates/vendor/candle/candle-transformers/src/models/flux/mod.rs
+++ b/crates/vendor/candle/candle-transformers/src/models/flux/mod.rs
@@ -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,
diff --git a/crates/vendor/candle/candle-transformers/src/models/flux/model.rs b/crates/vendor/candle/candle-transformers/src/models/flux/model.rs
index 17b4eb2532..ceef0824cd 100644
--- a/crates/vendor/candle/candle-transformers/src/models/flux/model.rs
+++ b/crates/vendor/candle/candle-transformers/src/models/flux/model.rs
@@ -1,626 +1,853 @@
-use candle::{DType, IndexOp, Result, Tensor, D};
-use candle_nn::{LayerNorm, Linear, RmsNorm, VarBuilder};
+use candle::{DType, IndexOp, Result, Tensor, D, Device};
+use candle_nn::{LayerNorm, Linear, RmsNorm, VarBuilder, DeviceTransferable};
+
+/// Trait for modules that can be moved between devices
+///
+/// This trait is intended to be implemented by model components that need
+/// to be transferred between CPU and GPU to save memory.
 
 // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/model.py#L12
 #[derive(Debug, Clone)]
 pub struct Config {
-    pub in_channels: usize,
-    pub vec_in_dim: usize,
-    pub context_in_dim: usize,
-    pub hidden_size: usize,
-    pub mlp_ratio: f64,
-    pub num_heads: usize,
-    pub depth: usize,
-    pub depth_single_blocks: usize,
-    pub axes_dim: Vec<usize>,
-    pub theta: usize,
-    pub qkv_bias: bool,
-    pub guidance_embed: bool,
+  pub in_channels: usize,
+  pub vec_in_dim: usize,
+  pub context_in_dim: usize,
+  pub hidden_size: usize,
+  pub mlp_ratio: f64,
+  pub num_heads: usize,
+  pub depth: usize,
+  pub depth_single_blocks: usize,
+  pub axes_dim: Vec<usize>,
+  pub theta: usize,
+  pub qkv_bias: bool,
+  pub guidance_embed: bool,
 }
 
 impl Config {
-    // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L32
-    pub fn dev() -> Self {
-        Self {
-            in_channels: 64,
-            vec_in_dim: 768,
-            context_in_dim: 4096,
-            hidden_size: 3072,
-            mlp_ratio: 4.0,
-            num_heads: 24,
-            depth: 19,
-            depth_single_blocks: 38,
-            axes_dim: vec![16, 56, 56],
-            theta: 10_000,
-            qkv_bias: true,
-            guidance_embed: true,
-        }
-    }
+  // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L32
+  pub fn dev() -> Self {
+    Self { in_channels: 64, vec_in_dim: 768, context_in_dim: 4096, hidden_size: 3072, mlp_ratio: 4.0, num_heads: 24, depth: 19, depth_single_blocks: 38, axes_dim: vec![16, 56, 56], theta: 10_000, qkv_bias: true, guidance_embed: true }
+  }
 
-    // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L64
-    pub fn schnell() -> Self {
-        Self {
-            in_channels: 64,
-            vec_in_dim: 768,
-            context_in_dim: 4096,
-            hidden_size: 3072,
-            mlp_ratio: 4.0,
-            num_heads: 24,
-            depth: 19,
-            depth_single_blocks: 38,
-            axes_dim: vec![16, 56, 56],
-            theta: 10_000,
-            qkv_bias: true,
-            guidance_embed: false,
-        }
-    }
+  // https://github.com/black-forest-labs/flux/blob/727e3a71faf37390f318cf9434f0939653302b60/src/flux/util.py#L64
+  pub fn schnell() -> Self {
+    Self { in_channels: 64, vec_in_dim: 768, context_in_dim: 4096, hidden_size: 3072, mlp_ratio: 4.0, num_heads: 24, depth: 19, depth_single_blocks: 38, axes_dim: vec![16, 56, 56], theta: 10_000, qkv_bias: true, guidance_embed: false }
+  }
 }
 
 fn layer_norm(dim: usize, vb: VarBuilder) -> Result<LayerNorm> {
-    let ws = Tensor::ones(dim, vb.dtype(), vb.device())?;
-    Ok(LayerNorm::new_no_bias(ws, 1e-6))
+  let ws = Tensor::ones(dim, vb.dtype(), vb.device())?;
+  Ok(LayerNorm::new_no_bias(ws, 1e-6))
 }
 
 fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
-    let dim = q.dim(D::Minus1)?;
-    let scale_factor = 1.0 / (dim as f64).sqrt();
-    let mut batch_dims = q.dims().to_vec();
-    batch_dims.pop();
-    batch_dims.pop();
-    let q = q.flatten_to(batch_dims.len() - 1)?;
-    let k = k.flatten_to(batch_dims.len() - 1)?;
-    let v = v.flatten_to(batch_dims.len() - 1)?;
-    let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?;
-    let attn_scores = candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(&v)?;
-    batch_dims.push(attn_scores.dim(D::Minus2)?);
-    batch_dims.push(attn_scores.dim(D::Minus1)?);
-    attn_scores.reshape(batch_dims)
+  let dim = q.dim(D::Minus1)?;
+  let scale_factor = 1.0 / (dim as f64).sqrt();
+  let mut batch_dims = q.dims().to_vec();
+  batch_dims.pop();
+  batch_dims.pop();
+  let q = q.flatten_to(batch_dims.len() - 1)?;
+  let k = k.flatten_to(batch_dims.len() - 1)?;
+  let v = v.flatten_to(batch_dims.len() - 1)?;
+  let attn_weights = (q.matmul(&k.t()?)? * scale_factor)?;
+  let attn_scores = candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(&v)?;
+  batch_dims.push(attn_scores.dim(D::Minus2)?);
+  batch_dims.push(attn_scores.dim(D::Minus1)?);
+  attn_scores.reshape(batch_dims)
 }
 
 fn rope(pos: &Tensor, dim: usize, theta: usize) -> Result<Tensor> {
-    if dim % 2 == 1 {
-        candle::bail!("dim {dim} is odd")
-    }
-    let dev = pos.device();
-    let theta = theta as f64;
-    let inv_freq: Vec<_> = (0..dim)
-        .step_by(2)
-        .map(|i| 1f32 / theta.powf(i as f64 / dim as f64) as f32)
-        .collect();
-    let inv_freq_len = inv_freq.len();
-    let inv_freq = Tensor::from_vec(inv_freq, (1, 1, inv_freq_len), dev)?;
-    let inv_freq = inv_freq.to_dtype(pos.dtype())?;
-    let freqs = pos.unsqueeze(2)?.broadcast_mul(&inv_freq)?;
-    let cos = freqs.cos()?;
-    let sin = freqs.sin()?;
-    let out = Tensor::stack(&[&cos, &sin.neg()?, &sin, &cos], 3)?;
-    let (b, n, d, _ij) = out.dims4()?;
-    out.reshape((b, n, d, 2, 2))
+  if dim % 2 == 1 {
+    candle::bail!("dim {dim} is odd")
+  }
+  let dev = pos.device();
+  let theta = theta as f64;
+  let inv_freq: Vec<_> = (0..dim).step_by(2).map(|i| 1f32 / theta.powf(i as f64 / dim as f64) as f32).collect();
+  let inv_freq_len = inv_freq.len();
+  let inv_freq = Tensor::from_vec(inv_freq, (1, 1, inv_freq_len), dev)?;
+  let inv_freq = inv_freq.to_dtype(pos.dtype())?;
+  let freqs = pos.unsqueeze(2)?.broadcast_mul(&inv_freq)?;
+  let cos = freqs.cos()?;
+  let sin = freqs.sin()?;
+  let out = Tensor::stack(&[&cos, &sin.neg()?, &sin, &cos], 3)?;
+  let (b, n, d, _ij) = out.dims4()?;
+  out.reshape((b, n, d, 2, 2))
 }
 
 fn apply_rope(x: &Tensor, freq_cis: &Tensor) -> Result<Tensor> {
-    let dims = x.dims();
-    let (b_sz, n_head, seq_len, n_embd) = x.dims4()?;
-    let x = x.reshape((b_sz, n_head, seq_len, n_embd / 2, 2))?;
-    let x0 = x.narrow(D::Minus1, 0, 1)?;
-    let x1 = x.narrow(D::Minus1, 1, 1)?;
-    let fr0 = freq_cis.get_on_dim(D::Minus1, 0)?;
-    let fr1 = freq_cis.get_on_dim(D::Minus1, 1)?;
-    (fr0.broadcast_mul(&x0)? + fr1.broadcast_mul(&x1)?)?.reshape(dims.to_vec())
+  let dims = x.dims();
+  let (b_sz, n_head, seq_len, n_embd) = x.dims4()?;
+  let x = x.reshape((b_sz, n_head, seq_len, n_embd / 2, 2))?;
+  let x0 = x.narrow(D::Minus1, 0, 1)?;
+  let x1 = x.narrow(D::Minus1, 1, 1)?;
+  let fr0 = freq_cis.get_on_dim(D::Minus1, 0)?;
+  let fr1 = freq_cis.get_on_dim(D::Minus1, 1)?;
+  (fr0.broadcast_mul(&x0)? + fr1.broadcast_mul(&x1)?)?.reshape(dims.to_vec())
 }
 
 pub(crate) fn attention(q: &Tensor, k: &Tensor, v: &Tensor, pe: &Tensor) -> Result<Tensor> {
-    let q = apply_rope(q, pe)?.contiguous()?;
-    let k = apply_rope(k, pe)?.contiguous()?;
-    let x = scaled_dot_product_attention(&q, &k, v)?;
-    x.transpose(1, 2)?.flatten_from(2)
+  let q = apply_rope(q, pe)?.contiguous()?;
+  let k = apply_rope(k, pe)?.contiguous()?;
+  let x = scaled_dot_product_attention(&q, &k, v)?;
+  x.transpose(1, 2)?.flatten_from(2)
 }
 
 pub(crate) fn timestep_embedding(t: &Tensor, dim: usize, dtype: DType) -> Result<Tensor> {
-    const TIME_FACTOR: f64 = 1000.;
-    const MAX_PERIOD: f64 = 10000.;
-    if dim % 2 == 1 {
-        candle::bail!("{dim} is odd")
-    }
-    let dev = t.device();
-    let half = dim / 2;
-    let t = (t * TIME_FACTOR)?;
-    let arange = Tensor::arange(0, half as u32, dev)?.to_dtype(candle::DType::F32)?;
-    let freqs = (arange * (-MAX_PERIOD.ln() / half as f64))?.exp()?;
-    let args = t
-        .unsqueeze(1)?
-        .to_dtype(candle::DType::F32)?
-        .broadcast_mul(&freqs.unsqueeze(0)?)?;
-    let emb = Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?.to_dtype(dtype)?;
-    Ok(emb)
+  const TIME_FACTOR: f64 = 1000.;
+  const MAX_PERIOD: f64 = 10000.;
+  if dim % 2 == 1 {
+    candle::bail!("{dim} is odd")
+  }
+  let dev = t.device();
+  let half = dim / 2;
+  let t = (t * TIME_FACTOR)?;
+  let arange = Tensor::arange(0, half as u32, dev)?.to_dtype(candle::DType::F32)?;
+  let freqs = (arange * (-MAX_PERIOD.ln() / half as f64))?.exp()?;
+  let args = t.unsqueeze(1)?.to_dtype(candle::DType::F32)?.broadcast_mul(&freqs.unsqueeze(0)?)?;
+  let emb = Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?.to_dtype(dtype)?;
+  Ok(emb)
 }
 
 #[derive(Debug, Clone)]
 pub struct EmbedNd {
-    #[allow(unused)]
-    dim: usize,
-    theta: usize,
-    axes_dim: Vec<usize>,
+  #[allow(unused)]
+  dim: usize,
+  theta: usize,
+  axes_dim: Vec<usize>,
 }
 
 impl EmbedNd {
-    pub fn new(dim: usize, theta: usize, axes_dim: Vec<usize>) -> Self {
-        Self {
-            dim,
-            theta,
-            axes_dim,
-        }
-    }
+  pub fn new(dim: usize, theta: usize, axes_dim: Vec<usize>) -> Self {
+    Self { dim, theta, axes_dim }
+  }
 }
 
 impl candle::Module for EmbedNd {
-    fn forward(&self, ids: &Tensor) -> Result<Tensor> {
-        let n_axes = ids.dim(D::Minus1)?;
-        let mut emb = Vec::with_capacity(n_axes);
-        for idx in 0..n_axes {
-            let r = rope(
-                &ids.get_on_dim(D::Minus1, idx)?,
-                self.axes_dim[idx],
-                self.theta,
-            )?;
-            emb.push(r)
-        }
-        let emb = Tensor::cat(&emb, 2)?;
-        emb.unsqueeze(1)
+  fn forward(&self, ids: &Tensor) -> Result<Tensor> {
+    let n_axes = ids.dim(D::Minus1)?;
+    let mut emb = Vec::with_capacity(n_axes);
+    for idx in 0..n_axes {
+      let r = rope(&ids.get_on_dim(D::Minus1, idx)?, self.axes_dim[idx], self.theta)?;
+      emb.push(r)
     }
+    let emb = Tensor::cat(&emb, 2)?;
+    emb.unsqueeze(1)
+  }
+}
+
+impl DeviceTransferable for EmbedNd {
+  fn to_device(&self, _device: &Device) -> Result<Self> {
+    // EmbedNd doesn't contain any tensors, just configuration
+    Ok(self.clone())
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct MlpEmbedder {
-    in_layer: Linear,
-    out_layer: Linear,
+  in_layer: Linear,
+  out_layer: Linear,
 }
 
 impl MlpEmbedder {
-    fn new(in_sz: usize, h_sz: usize, vb: VarBuilder) -> Result<Self> {
-        let in_layer = candle_nn::linear(in_sz, h_sz, vb.pp("in_layer"))?;
-        let out_layer = candle_nn::linear(h_sz, h_sz, vb.pp("out_layer"))?;
-        Ok(Self {
-            in_layer,
-            out_layer,
-        })
-    }
+  fn new(in_sz: usize, h_sz: usize, vb: VarBuilder) -> Result<Self> {
+    let in_layer = candle_nn::linear(in_sz, h_sz, vb.pp("in_layer"))?;
+    let out_layer = candle_nn::linear(h_sz, h_sz, vb.pp("out_layer"))?;
+    Ok(Self { in_layer, out_layer })
+  }
 }
 
 impl candle::Module for MlpEmbedder {
-    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        xs.apply(&self.in_layer)?.silu()?.apply(&self.out_layer)
-    }
+  fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+    xs.apply(&self.in_layer)?.silu()?.apply(&self.out_layer)
+  }
+}
+
+impl DeviceTransferable for MlpEmbedder {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { in_layer: self.in_layer.to_device(device)?, out_layer: self.out_layer.to_device(device)? })
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct QkNorm {
-    query_norm: RmsNorm,
-    key_norm: RmsNorm,
+  query_norm: RmsNorm,
+  key_norm: RmsNorm,
 }
 
 impl QkNorm {
-    fn new(dim: usize, vb: VarBuilder) -> Result<Self> {
-        let query_norm = vb.get(dim, "query_norm.scale")?;
-        let query_norm = RmsNorm::new(query_norm, 1e-6);
-        let key_norm = vb.get(dim, "key_norm.scale")?;
-        let key_norm = RmsNorm::new(key_norm, 1e-6);
-        Ok(Self {
-            query_norm,
-            key_norm,
-        })
-    }
+  fn new(dim: usize, vb: VarBuilder) -> Result<Self> {
+    let query_norm = vb.get(dim, "query_norm.scale")?;
+    let query_norm = RmsNorm::new(query_norm, 1e-6);
+    let key_norm = vb.get(dim, "key_norm.scale")?;
+    let key_norm = RmsNorm::new(key_norm, 1e-6);
+    Ok(Self { query_norm, key_norm })
+  }
 }
 
 struct ModulationOut {
-    shift: Tensor,
-    scale: Tensor,
-    gate: Tensor,
+  shift: Tensor,
+  scale: Tensor,
+  gate: Tensor,
 }
 
 impl ModulationOut {
-    fn scale_shift(&self, xs: &Tensor) -> Result<Tensor> {
-        xs.broadcast_mul(&(&self.scale + 1.)?)?
-            .broadcast_add(&self.shift)
-    }
+  fn scale_shift(&self, xs: &Tensor) -> Result<Tensor> {
+    xs.broadcast_mul(&(&self.scale + 1.)?)?.broadcast_add(&self.shift)
+  }
 
-    fn gate(&self, xs: &Tensor) -> Result<Tensor> {
-        self.gate.broadcast_mul(xs)
-    }
+  fn gate(&self, xs: &Tensor) -> Result<Tensor> {
+    self.gate.broadcast_mul(xs)
+  }
 }
 
 #[derive(Debug, Clone)]
 struct Modulation1 {
-    lin: Linear,
+  lin: Linear,
 }
 
 impl Modulation1 {
-    fn new(dim: usize, vb: VarBuilder) -> Result<Self> {
-        let lin = candle_nn::linear(dim, 3 * dim, vb.pp("lin"))?;
-        Ok(Self { lin })
-    }
+  fn new(dim: usize, vb: VarBuilder) -> Result<Self> {
+    let lin = candle_nn::linear(dim, 3 * dim, vb.pp("lin"))?;
+    Ok(Self { lin })
+  }
 
-    fn forward(&self, vec_: &Tensor) -> Result<ModulationOut> {
-        let ys = vec_
-            .silu()?
-            .apply(&self.lin)?
-            .unsqueeze(1)?
-            .chunk(3, D::Minus1)?;
-        if ys.len() != 3 {
-            candle::bail!("unexpected len from chunk {ys:?}")
-        }
-        Ok(ModulationOut {
-            shift: ys[0].clone(),
-            scale: ys[1].clone(),
-            gate: ys[2].clone(),
-        })
+  fn forward(&self, vec_: &Tensor) -> Result<ModulationOut> {
+    let ys = vec_.silu()?.apply(&self.lin)?.unsqueeze(1)?.chunk(3, D::Minus1)?;
+    if ys.len() != 3 {
+      candle::bail!("unexpected len from chunk {ys:?}")
     }
+    Ok(ModulationOut { shift: ys[0].clone(), scale: ys[1].clone(), gate: ys[2].clone() })
+  }
 }
 
 #[derive(Debug, Clone)]
 struct Modulation2 {
-    lin: Linear,
+  lin: Linear,
 }
 
 impl Modulation2 {
-    fn new(dim: usize, vb: VarBuilder) -> Result<Self> {
-        let lin = candle_nn::linear(dim, 6 * dim, vb.pp("lin"))?;
-        Ok(Self { lin })
-    }
+  fn new(dim: usize, vb: VarBuilder) -> Result<Self> {
+    let lin = candle_nn::linear(dim, 6 * dim, vb.pp("lin"))?;
+    Ok(Self { lin })
+  }
 
-    fn forward(&self, vec_: &Tensor) -> Result<(ModulationOut, ModulationOut)> {
-        let ys = vec_
-            .silu()?
-            .apply(&self.lin)?
-            .unsqueeze(1)?
-            .chunk(6, D::Minus1)?;
-        if ys.len() != 6 {
-            candle::bail!("unexpected len from chunk {ys:?}")
-        }
-        let mod1 = ModulationOut {
-            shift: ys[0].clone(),
-            scale: ys[1].clone(),
-            gate: ys[2].clone(),
-        };
-        let mod2 = ModulationOut {
-            shift: ys[3].clone(),
-            scale: ys[4].clone(),
-            gate: ys[5].clone(),
-        };
-        Ok((mod1, mod2))
+  fn forward(&self, vec_: &Tensor) -> Result<(ModulationOut, ModulationOut)> {
+    let ys = vec_.silu()?.apply(&self.lin)?.unsqueeze(1)?.chunk(6, D::Minus1)?;
+    if ys.len() != 6 {
+      candle::bail!("unexpected len from chunk {ys:?}")
     }
+    let mod1 = ModulationOut { shift: ys[0].clone(), scale: ys[1].clone(), gate: ys[2].clone() };
+    let mod2 = ModulationOut { shift: ys[3].clone(), scale: ys[4].clone(), gate: ys[5].clone() };
+    Ok((mod1, mod2))
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct SelfAttention {
-    qkv: Linear,
-    norm: QkNorm,
-    proj: Linear,
-    num_heads: usize,
+  qkv: Linear,
+  norm: QkNorm,
+  proj: Linear,
+  num_heads: usize,
 }
 
 impl SelfAttention {
-    fn new(dim: usize, num_heads: usize, qkv_bias: bool, vb: VarBuilder) -> Result<Self> {
-        let head_dim = dim / num_heads;
-        let qkv = candle_nn::linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?;
-        let norm = QkNorm::new(head_dim, vb.pp("norm"))?;
-        let proj = candle_nn::linear(dim, dim, vb.pp("proj"))?;
-        Ok(Self {
-            qkv,
-            norm,
-            proj,
-            num_heads,
-        })
-    }
+  fn new(dim: usize, num_heads: usize, qkv_bias: bool, vb: VarBuilder) -> Result<Self> {
+    let head_dim = dim / num_heads;
+    let qkv = candle_nn::linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?;
+    let norm = QkNorm::new(head_dim, vb.pp("norm"))?;
+    let proj = candle_nn::linear(dim, dim, vb.pp("proj"))?;
+    Ok(Self { qkv, norm, proj, num_heads })
+  }
 
-    fn qkv(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor)> {
-        let qkv = xs.apply(&self.qkv)?;
-        let (b, l, _khd) = qkv.dims3()?;
-        let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?;
-        let q = qkv.i((.., .., 0))?.transpose(1, 2)?;
-        let k = qkv.i((.., .., 1))?.transpose(1, 2)?;
-        let v = qkv.i((.., .., 2))?.transpose(1, 2)?;
-        let q = q.apply(&self.norm.query_norm)?;
-        let k = k.apply(&self.norm.key_norm)?;
-        Ok((q, k, v))
-    }
+  fn qkv(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor)> {
+    let qkv = xs.apply(&self.qkv)?;
+    let (b, l, _khd) = qkv.dims3()?;
+    let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?;
+    let q = qkv.i((.., .., 0))?.transpose(1, 2)?;
+    let k = qkv.i((.., .., 1))?.transpose(1, 2)?;
+    let v = qkv.i((.., .., 2))?.transpose(1, 2)?;
+    let q = q.apply(&self.norm.query_norm)?;
+    let k = k.apply(&self.norm.key_norm)?;
+    Ok((q, k, v))
+  }
 
-    #[allow(unused)]
-    fn forward(&self, xs: &Tensor, pe: &Tensor) -> Result<Tensor> {
-        let (q, k, v) = self.qkv(xs)?;
-        attention(&q, &k, &v, pe)?.apply(&self.proj)
-    }
+  #[allow(unused)]
+  fn forward(&self, xs: &Tensor, pe: &Tensor) -> Result<Tensor> {
+    let (q, k, v) = self.qkv(xs)?;
+    attention(&q, &k, &v, pe)?.apply(&self.proj)
+  }
 }
 
 #[derive(Debug, Clone)]
 struct Mlp {
-    lin1: Linear,
-    lin2: Linear,
+  lin1: Linear,
+  lin2: Linear,
 }
 
 impl Mlp {
-    fn new(in_sz: usize, mlp_sz: usize, vb: VarBuilder) -> Result<Self> {
-        let lin1 = candle_nn::linear(in_sz, mlp_sz, vb.pp("0"))?;
-        let lin2 = candle_nn::linear(mlp_sz, in_sz, vb.pp("2"))?;
-        Ok(Self { lin1, lin2 })
-    }
+  fn new(in_sz: usize, mlp_sz: usize, vb: VarBuilder) -> Result<Self> {
+    let lin1 = candle_nn::linear(in_sz, mlp_sz, vb.pp("0"))?;
+    let lin2 = candle_nn::linear(mlp_sz, in_sz, vb.pp("2"))?;
+    Ok(Self { lin1, lin2 })
+  }
 }
 
 impl candle::Module for Mlp {
-    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        xs.apply(&self.lin1)?.gelu()?.apply(&self.lin2)
-    }
+  fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+    xs.apply(&self.lin1)?.gelu()?.apply(&self.lin2)
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct DoubleStreamBlock {
-    img_mod: Modulation2,
-    img_norm1: LayerNorm,
-    img_attn: SelfAttention,
-    img_norm2: LayerNorm,
-    img_mlp: Mlp,
-    txt_mod: Modulation2,
-    txt_norm1: LayerNorm,
-    txt_attn: SelfAttention,
-    txt_norm2: LayerNorm,
-    txt_mlp: Mlp,
+  img_mod: Modulation2,
+  img_norm1: LayerNorm,
+  img_attn: SelfAttention,
+  img_norm2: LayerNorm,
+  img_mlp: Mlp,
+  txt_mod: Modulation2,
+  txt_norm1: LayerNorm,
+  txt_attn: SelfAttention,
+  txt_norm2: LayerNorm,
+  txt_mlp: Mlp,
+}
+
+impl DeviceTransferable for DoubleStreamBlock {
+  fn to_device(&self, _device: &Device) -> Result<Self> {
+    Ok(Self { img_mod: self.img_mod.to_device(_device)?, img_norm1: self.img_norm1.to_device(_device)?, img_attn: self.img_attn.to_device(_device)?, img_norm2: self.img_norm2.to_device(_device)?, img_mlp: self.img_mlp.to_device(_device)?, txt_mod: self.txt_mod.to_device(_device)?, txt_norm1: self.txt_norm1.to_device(_device)?, txt_attn: self.txt_attn.to_device(_device)?, txt_norm2: self.txt_norm2.to_device(_device)?, txt_mlp: self.txt_mlp.to_device(_device)? })
+  }
 }
 
 impl DoubleStreamBlock {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
-        let h_sz = cfg.hidden_size;
-        let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize;
-        let img_mod = Modulation2::new(h_sz, vb.pp("img_mod"))?;
-        let img_norm1 = layer_norm(h_sz, vb.pp("img_norm1"))?;
-        let img_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("img_attn"))?;
-        let img_norm2 = layer_norm(h_sz, vb.pp("img_norm2"))?;
-        let img_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("img_mlp"))?;
-        let txt_mod = Modulation2::new(h_sz, vb.pp("txt_mod"))?;
-        let txt_norm1 = layer_norm(h_sz, vb.pp("txt_norm1"))?;
-        let txt_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("txt_attn"))?;
-        let txt_norm2 = layer_norm(h_sz, vb.pp("txt_norm2"))?;
-        let txt_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("txt_mlp"))?;
-        Ok(Self {
-            img_mod,
-            img_norm1,
-            img_attn,
-            img_norm2,
-            img_mlp,
-            txt_mod,
-            txt_norm1,
-            txt_attn,
-            txt_norm2,
-            txt_mlp,
-        })
-    }
+  fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    let h_sz = cfg.hidden_size;
+    let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize;
+    let img_mod = Modulation2::new(h_sz, vb.pp("img_mod"))?;
+    let img_norm1 = layer_norm(h_sz, vb.pp("img_norm1"))?;
+    let img_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("img_attn"))?;
+    let img_norm2 = layer_norm(h_sz, vb.pp("img_norm2"))?;
+    let img_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("img_mlp"))?;
+    let txt_mod = Modulation2::new(h_sz, vb.pp("txt_mod"))?;
+    let txt_norm1 = layer_norm(h_sz, vb.pp("txt_norm1"))?;
+    let txt_attn = SelfAttention::new(h_sz, cfg.num_heads, cfg.qkv_bias, vb.pp("txt_attn"))?;
+    let txt_norm2 = layer_norm(h_sz, vb.pp("txt_norm2"))?;
+    let txt_mlp = Mlp::new(h_sz, mlp_sz, vb.pp("txt_mlp"))?;
+    Ok(Self { img_mod, img_norm1, img_attn, img_norm2, img_mlp, txt_mod, txt_norm1, txt_attn, txt_norm2, txt_mlp })
+  }
 
-    fn forward(
-        &self,
-        img: &Tensor,
-        txt: &Tensor,
-        vec_: &Tensor,
-        pe: &Tensor,
-    ) -> Result<(Tensor, Tensor)> {
-        let (img_mod1, img_mod2) = self.img_mod.forward(vec_)?; // shift, scale, gate
-        let (txt_mod1, txt_mod2) = self.txt_mod.forward(vec_)?; // shift, scale, gate
-        let img_modulated = img.apply(&self.img_norm1)?;
-        let img_modulated = img_mod1.scale_shift(&img_modulated)?;
-        let (img_q, img_k, img_v) = self.img_attn.qkv(&img_modulated)?;
+  fn forward(&self, img: &Tensor, txt: &Tensor, vec_: &Tensor, pe: &Tensor) -> Result<(Tensor, Tensor)> {
+    let (img_mod1, img_mod2) = self.img_mod.forward(vec_)?; // shift, scale, gate
+    let (txt_mod1, txt_mod2) = self.txt_mod.forward(vec_)?; // shift, scale, gate
+    let img_modulated = img.apply(&self.img_norm1)?;
+    let img_modulated = img_mod1.scale_shift(&img_modulated)?;
+    let (img_q, img_k, img_v) = self.img_attn.qkv(&img_modulated)?;
 
-        let txt_modulated = txt.apply(&self.txt_norm1)?;
-        let txt_modulated = txt_mod1.scale_shift(&txt_modulated)?;
-        let (txt_q, txt_k, txt_v) = self.txt_attn.qkv(&txt_modulated)?;
+    let txt_modulated = txt.apply(&self.txt_norm1)?;
+    let txt_modulated = txt_mod1.scale_shift(&txt_modulated)?;
+    let (txt_q, txt_k, txt_v) = self.txt_attn.qkv(&txt_modulated)?;
 
-        let q = Tensor::cat(&[txt_q, img_q], 2)?;
-        let k = Tensor::cat(&[txt_k, img_k], 2)?;
-        let v = Tensor::cat(&[txt_v, img_v], 2)?;
+    let q = Tensor::cat(&[txt_q, img_q], 2)?;
+    let k = Tensor::cat(&[txt_k, img_k], 2)?;
+    let v = Tensor::cat(&[txt_v, img_v], 2)?;
 
-        let attn = attention(&q, &k, &v, pe)?;
-        let txt_attn = attn.narrow(1, 0, txt.dim(1)?)?;
-        let img_attn = attn.narrow(1, txt.dim(1)?, attn.dim(1)? - txt.dim(1)?)?;
+    let attn = attention(&q, &k, &v, pe)?;
+    let txt_attn = attn.narrow(1, 0, txt.dim(1)?)?;
+    let img_attn = attn.narrow(1, txt.dim(1)?, attn.dim(1)? - txt.dim(1)?)?;
 
-        let img = (img + img_mod1.gate(&img_attn.apply(&self.img_attn.proj)?))?;
-        let img = (&img
-            + img_mod2.gate(
-                &img_mod2
-                    .scale_shift(&img.apply(&self.img_norm2)?)?
-                    .apply(&self.img_mlp)?,
-            )?)?;
+    let img = (img + img_mod1.gate(&img_attn.apply(&self.img_attn.proj)?))?;
+    let img = (&img + img_mod2.gate(&img_mod2.scale_shift(&img.apply(&self.img_norm2)?)?.apply(&self.img_mlp)?)?)?;
 
-        let txt = (txt + txt_mod1.gate(&txt_attn.apply(&self.txt_attn.proj)?))?;
-        let txt = (&txt
-            + txt_mod2.gate(
-                &txt_mod2
-                    .scale_shift(&txt.apply(&self.txt_norm2)?)?
-                    .apply(&self.txt_mlp)?,
-            )?)?;
+    let txt = (txt + txt_mod1.gate(&txt_attn.apply(&self.txt_attn.proj)?))?;
+    let txt = (&txt + txt_mod2.gate(&txt_mod2.scale_shift(&txt.apply(&self.txt_norm2)?)?.apply(&self.txt_mlp)?)?)?;
 
-        Ok((img, txt))
-    }
+    Ok((img, txt))
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct SingleStreamBlock {
-    linear1: Linear,
-    linear2: Linear,
-    norm: QkNorm,
-    pre_norm: LayerNorm,
-    modulation: Modulation1,
-    h_sz: usize,
-    mlp_sz: usize,
-    num_heads: usize,
+  linear1: Linear,
+  linear2: Linear,
+  norm: QkNorm,
+  pre_norm: LayerNorm,
+  modulation: Modulation1,
+  h_sz: usize,
+  mlp_sz: usize,
+  num_heads: usize,
+}
+
+impl DeviceTransferable for SingleStreamBlock {
+  fn to_device(&self, _device: &Device) -> Result<Self> {
+    Ok(Self { linear1: self.linear1.to_device(_device)?, linear2: self.linear2.to_device(_device)?, norm: self.norm.to_device(_device)?, pre_norm: self.pre_norm.to_device(_device)?, modulation: self.modulation.to_device(_device)?, h_sz: self.h_sz, mlp_sz: self.mlp_sz, num_heads: self.num_heads })
+  }
 }
 
 impl SingleStreamBlock {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
-        let h_sz = cfg.hidden_size;
-        let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize;
-        let head_dim = h_sz / cfg.num_heads;
-        let linear1 = candle_nn::linear(h_sz, h_sz * 3 + mlp_sz, vb.pp("linear1"))?;
-        let linear2 = candle_nn::linear(h_sz + mlp_sz, h_sz, vb.pp("linear2"))?;
-        let norm = QkNorm::new(head_dim, vb.pp("norm"))?;
-        let pre_norm = layer_norm(h_sz, vb.pp("pre_norm"))?;
-        let modulation = Modulation1::new(h_sz, vb.pp("modulation"))?;
-        Ok(Self {
-            linear1,
-            linear2,
-            norm,
-            pre_norm,
-            modulation,
-            h_sz,
-            mlp_sz,
-            num_heads: cfg.num_heads,
-        })
-    }
+  fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    let h_sz = cfg.hidden_size;
+    let mlp_sz = (h_sz as f64 * cfg.mlp_ratio) as usize;
+    let head_dim = h_sz / cfg.num_heads;
+    let linear1 = candle_nn::linear(h_sz, h_sz * 3 + mlp_sz, vb.pp("linear1"))?;
+    let linear2 = candle_nn::linear(h_sz + mlp_sz, h_sz, vb.pp("linear2"))?;
+    let norm = QkNorm::new(head_dim, vb.pp("norm"))?;
+    let pre_norm = layer_norm(h_sz, vb.pp("pre_norm"))?;
+    let modulation = Modulation1::new(h_sz, vb.pp("modulation"))?;
+    Ok(Self { linear1, linear2, norm, pre_norm, modulation, h_sz, mlp_sz, num_heads: cfg.num_heads })
+  }
 
-    fn forward(&self, xs: &Tensor, vec_: &Tensor, pe: &Tensor) -> Result<Tensor> {
-        let mod_ = self.modulation.forward(vec_)?;
-        let x_mod = mod_.scale_shift(&xs.apply(&self.pre_norm)?)?;
-        let x_mod = x_mod.apply(&self.linear1)?;
-        let qkv = x_mod.narrow(D::Minus1, 0, 3 * self.h_sz)?;
-        let (b, l, _khd) = qkv.dims3()?;
-        let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?;
-        let q = qkv.i((.., .., 0))?.transpose(1, 2)?;
-        let k = qkv.i((.., .., 1))?.transpose(1, 2)?;
-        let v = qkv.i((.., .., 2))?.transpose(1, 2)?;
-        let mlp = x_mod.narrow(D::Minus1, 3 * self.h_sz, self.mlp_sz)?;
-        let q = q.apply(&self.norm.query_norm)?;
-        let k = k.apply(&self.norm.key_norm)?;
-        let attn = attention(&q, &k, &v, pe)?;
-        let output = Tensor::cat(&[attn, mlp.gelu()?], 2)?.apply(&self.linear2)?;
-        xs + mod_.gate(&output)
-    }
+  fn forward(&self, xs: &Tensor, vec_: &Tensor, pe: &Tensor) -> Result<Tensor> {
+    let mod_ = self.modulation.forward(vec_)?;
+    let x_mod = mod_.scale_shift(&xs.apply(&self.pre_norm)?)?;
+    let x_mod = x_mod.apply(&self.linear1)?;
+    let qkv = x_mod.narrow(D::Minus1, 0, 3 * self.h_sz)?;
+    let (b, l, _khd) = qkv.dims3()?;
+    let qkv = qkv.reshape((b, l, 3, self.num_heads, ()))?;
+    let q = qkv.i((.., .., 0))?.transpose(1, 2)?;
+    let k = qkv.i((.., .., 1))?.transpose(1, 2)?;
+    let v = qkv.i((.., .., 2))?.transpose(1, 2)?;
+    let mlp = x_mod.narrow(D::Minus1, 3 * self.h_sz, self.mlp_sz)?;
+    let q = q.apply(&self.norm.query_norm)?;
+    let k = k.apply(&self.norm.key_norm)?;
+    let attn = attention(&q, &k, &v, pe)?;
+    let output = Tensor::cat(&[attn, mlp.gelu()?], 2)?.apply(&self.linear2)?;
+    xs + mod_.gate(&output)
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct LastLayer {
-    norm_final: LayerNorm,
-    linear: Linear,
-    ada_ln_modulation: Linear,
+  norm_final: LayerNorm,
+  linear: Linear,
+  ada_ln_modulation: Linear,
 }
 
 impl LastLayer {
-    fn new(h_sz: usize, p_sz: usize, out_c: usize, vb: VarBuilder) -> Result<Self> {
-        let norm_final = layer_norm(h_sz, vb.pp("norm_final"))?;
-        let linear = candle_nn::linear(h_sz, p_sz * p_sz * out_c, vb.pp("linear"))?;
-        let ada_ln_modulation = candle_nn::linear(h_sz, 2 * h_sz, vb.pp("adaLN_modulation.1"))?;
-        Ok(Self {
-            norm_final,
-            linear,
-            ada_ln_modulation,
-        })
-    }
+  fn new(h_sz: usize, p_sz: usize, out_c: usize, vb: VarBuilder) -> Result<Self> {
+    let norm_final = layer_norm(h_sz, vb.pp("norm_final"))?;
+    let linear = candle_nn::linear(h_sz, p_sz * p_sz * out_c, vb.pp("linear"))?;
+    let ada_ln_modulation = candle_nn::linear(h_sz, 2 * h_sz, vb.pp("adaLN_modulation.1"))?;
+    Ok(Self { norm_final, linear, ada_ln_modulation })
+  }
 
-    fn forward(&self, xs: &Tensor, vec: &Tensor) -> Result<Tensor> {
-        let chunks = vec.silu()?.apply(&self.ada_ln_modulation)?.chunk(2, 1)?;
-        let (shift, scale) = (&chunks[0], &chunks[1]);
-        let xs = xs
-            .apply(&self.norm_final)?
-            .broadcast_mul(&(scale.unsqueeze(1)? + 1.0)?)?
-            .broadcast_add(&shift.unsqueeze(1)?)?;
-        xs.apply(&self.linear)
-    }
+  fn forward(&self, xs: &Tensor, vec: &Tensor) -> Result<Tensor> {
+    let chunks = vec.silu()?.apply(&self.ada_ln_modulation)?.chunk(2, 1)?;
+    let (shift, scale) = (&chunks[0], &chunks[1]);
+    let xs = xs.apply(&self.norm_final)?.broadcast_mul(&(scale.unsqueeze(1)? + 1.0)?)?.broadcast_add(&shift.unsqueeze(1)?)?;
+    xs.apply(&self.linear)
+  }
+}
+
+impl DeviceTransferable for LastLayer {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { norm_final: self.norm_final.to_device(device)?, linear: self.linear.to_device(device)?, ada_ln_modulation: self.ada_ln_modulation.to_device(device)? })
+  }
 }
 
 #[derive(Debug, Clone)]
 pub struct Flux {
-    img_in: Linear,
-    txt_in: Linear,
-    time_in: MlpEmbedder,
-    vector_in: MlpEmbedder,
-    guidance_in: Option<MlpEmbedder>,
-    pe_embedder: EmbedNd,
-    double_blocks: Vec<DoubleStreamBlock>,
-    single_blocks: Vec<SingleStreamBlock>,
-    final_layer: LastLayer,
+  img_in: Linear,
+  txt_in: Linear,
+  time_in: MlpEmbedder,
+  vector_in: MlpEmbedder,
+  guidance_in: Option<MlpEmbedder>,
+  pe_embedder: EmbedNd,
+  double_blocks: Vec<DoubleStreamBlock>,
+  single_blocks: Vec<SingleStreamBlock>,
+  final_layer: LastLayer,
+  max_active_blocks: usize,    // Maximum number of blocks to keep on GPU simultaneously
+  prefetch_next_batch: bool,   // Whether to prefetch the next batch while processing the current one
+  use_device_management: bool, // Whether to use device management (moving blocks between CPU/GPU)
 }
 
 impl Flux {
-    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
-        let img_in = candle_nn::linear(cfg.in_channels, cfg.hidden_size, vb.pp("img_in"))?;
-        let txt_in = candle_nn::linear(cfg.context_in_dim, cfg.hidden_size, vb.pp("txt_in"))?;
-        let mut double_blocks = Vec::with_capacity(cfg.depth);
-        let vb_d = vb.pp("double_blocks");
-        for idx in 0..cfg.depth {
-            let db = DoubleStreamBlock::new(cfg, vb_d.pp(idx))?;
-            double_blocks.push(db)
-        }
-        let mut single_blocks = Vec::with_capacity(cfg.depth_single_blocks);
-        let vb_s = vb.pp("single_blocks");
-        for idx in 0..cfg.depth_single_blocks {
-            let sb = SingleStreamBlock::new(cfg, vb_s.pp(idx))?;
-            single_blocks.push(sb)
-        }
-        let time_in = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("time_in"))?;
-        let vector_in = MlpEmbedder::new(cfg.vec_in_dim, cfg.hidden_size, vb.pp("vector_in"))?;
-        let guidance_in = if cfg.guidance_embed {
-            let mlp = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("guidance_in"))?;
-            Some(mlp)
-        } else {
-            None
-        };
-        let final_layer =
-            LastLayer::new(cfg.hidden_size, 1, cfg.in_channels, vb.pp("final_layer"))?;
-        let pe_dim = cfg.hidden_size / cfg.num_heads;
-        let pe_embedder = EmbedNd::new(pe_dim, cfg.theta, cfg.axes_dim.to_vec());
-        Ok(Self {
-            img_in,
-            txt_in,
-            time_in,
-            vector_in,
-            guidance_in,
-            pe_embedder,
-            double_blocks,
-            single_blocks,
-            final_layer,
-        })
+  pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    let img_in = candle_nn::linear(cfg.in_channels, cfg.hidden_size, vb.pp("img_in"))?;
+    let txt_in = candle_nn::linear(cfg.context_in_dim, cfg.hidden_size, vb.pp("txt_in"))?;
+    let mut double_blocks = Vec::with_capacity(cfg.depth);
+    let vb_d = vb.pp("double_blocks");
+    for idx in 0..cfg.depth {
+      let db = DoubleStreamBlock::new(cfg, vb_d.pp(idx))?;
+      double_blocks.push(db)
     }
+    let mut single_blocks = Vec::with_capacity(cfg.depth_single_blocks);
+    let vb_s = vb.pp("single_blocks");
+    for idx in 0..cfg.depth_single_blocks {
+      let sb = SingleStreamBlock::new(cfg, vb_s.pp(idx))?;
+      single_blocks.push(sb)
+    }
+    let time_in = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("time_in"))?;
+    let vector_in = MlpEmbedder::new(cfg.vec_in_dim, cfg.hidden_size, vb.pp("vector_in"))?;
+    let guidance_in = if cfg.guidance_embed {
+      let mlp = MlpEmbedder::new(256, cfg.hidden_size, vb.pp("guidance_in"))?;
+      Some(mlp)
+    } else {
+      None
+    };
+    let final_layer = LastLayer::new(cfg.hidden_size, 1, cfg.in_channels, vb.pp("final_layer"))?;
+    let pe_dim = cfg.hidden_size / cfg.num_heads;
+    let pe_embedder = EmbedNd::new(pe_dim, cfg.theta, cfg.axes_dim.to_vec());
+    Ok(Self {
+      img_in,
+      txt_in,
+      time_in,
+      vector_in,
+      guidance_in,
+      pe_embedder,
+      double_blocks,
+      single_blocks,
+      final_layer,
+      max_active_blocks: 1,        // Default to keeping 1 block active
+      prefetch_next_batch: false,  // Default to no prefetching
+      use_device_management: true, // Default to using device management
+    })
+  }
+
+  /// Create a new Flux model with everything initially on CPU, then selectively
+  /// moving core components to GPU for better memory efficiency.
+  /// This approach avoids the initial GPU memory spike.
+  pub fn new_with_gpu_core(cfg: &Config, cpu_vb: VarBuilder, gpu_device: &Device) -> Result<Self> {
+    // Step 1: Create model with all components on CPU
+    // This is much more memory efficient as it never loads blocks to GPU
+    let cpu_model = Self::new(cfg, cpu_vb)?;
+
+    // Step 2: Move only the core components to GPU (keeping blocks on CPU)
+    // This selectively moves the lightweight components that need to be on GPU
+    // for efficient processing, while leaving the memory-heavy blocks on CPU
+    Ok(Self {
+      img_in: cpu_model.img_in.to_device(gpu_device)?,
+      txt_in: cpu_model.txt_in.to_device(gpu_device)?,
+      time_in: cpu_model.time_in.to_device(gpu_device)?,
+      vector_in: cpu_model.vector_in.to_device(gpu_device)?,
+      guidance_in: match cpu_model.guidance_in {
+        Some(g) => Some(g.to_device(gpu_device)?),
+        None => None,
+      },
+      pe_embedder: cpu_model.pe_embedder,
+      double_blocks: cpu_model.double_blocks,
+      single_blocks: cpu_model.single_blocks,
+      final_layer: cpu_model.final_layer.to_device(gpu_device)?,
+      max_active_blocks: cpu_model.max_active_blocks,
+      prefetch_next_batch: cpu_model.prefetch_next_batch,
+      use_device_management: cpu_model.use_device_management,
+    })
+  }
+
+  /// Set the maximum number of active blocks to keep on GPU simultaneously
+  pub fn set_max_active_blocks(&mut self, max_blocks: usize) {
+    self.max_active_blocks = max_blocks;
+  }
+
+  /// Enable or disable prefetching of the next batch
+  /// When enabled, this will attempt to load the next batch
+  /// of blocks to GPU while processing the current batch
+  pub fn set_prefetch_next_batch(&mut self, prefetch: bool) {
+    self.prefetch_next_batch = prefetch;
+  }
+
+  /// Enable or disable device management
+  /// When enabled (default), blocks are kept on CPU and only moved to GPU when needed
+  /// When disabled, all blocks stay on their current device (typically GPU)
+  pub fn set_use_device_management(&mut self, use_management: bool) {
+    self.use_device_management = use_management;
+  }
+}
+
+impl DeviceTransferable for Flux {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    let guidance_in = match &self.guidance_in {
+      Some(g) => Some(g.to_device(device)?),
+      None => None,
+    };
+
+    // When not using device management, move all blocks to the target device
+    let double_blocks = if !self.use_device_management {
+      let mut blocks = Vec::with_capacity(self.double_blocks.len());
+      for block in &self.double_blocks {
+        blocks.push(block.to_device(device)?);
+      }
+      blocks
+    } else {
+      // Otherwise keep them on their current device (CPU)
+      self.double_blocks.clone()
+    };
+
+    let single_blocks = if !self.use_device_management {
+      let mut blocks = Vec::with_capacity(self.single_blocks.len());
+      for block in &self.single_blocks {
+        blocks.push(block.to_device(device)?);
+      }
+      blocks
+    } else {
+      // Otherwise keep them on their current device (CPU)
+      self.single_blocks.clone()
+    };
+
+    Ok(Self { img_in: self.img_in.to_device(device)?, txt_in: self.txt_in.to_device(device)?, time_in: self.time_in.to_device(device)?, vector_in: self.vector_in.to_device(device)?, guidance_in, pe_embedder: self.pe_embedder.clone(), double_blocks, single_blocks, final_layer: self.final_layer.to_device(device)?, max_active_blocks: self.max_active_blocks, prefetch_next_batch: self.prefetch_next_batch, use_device_management: self.use_device_management })
+  }
 }
 
 impl super::WithForward for Flux {
-    #[allow(clippy::too_many_arguments)]
-    fn forward(
-        &self,
-        img: &Tensor,
-        img_ids: &Tensor,
-        txt: &Tensor,
-        txt_ids: &Tensor,
-        timesteps: &Tensor,
-        y: &Tensor,
-        guidance: Option<&Tensor>,
-    ) -> Result<Tensor> {
-        if txt.rank() != 3 {
-            candle::bail!("unexpected shape for txt {:?}", txt.shape())
-        }
-        if img.rank() != 3 {
-            candle::bail!("unexpected shape for img {:?}", img.shape())
-        }
-        let dtype = img.dtype();
-        let pe = {
-            let ids = Tensor::cat(&[txt_ids, img_ids], 1)?;
-            ids.apply(&self.pe_embedder)?
-        };
-        let mut txt = txt.apply(&self.txt_in)?;
-        let mut img = img.apply(&self.img_in)?;
-        let vec_ = timestep_embedding(timesteps, 256, dtype)?.apply(&self.time_in)?;
-        let vec_ = match (self.guidance_in.as_ref(), guidance) {
-            (Some(g_in), Some(guidance)) => {
-                (vec_ + timestep_embedding(guidance, 256, dtype)?.apply(g_in))?
-            }
-            _ => vec_,
-        };
-        let vec_ = (vec_ + y.apply(&self.vector_in))?;
-
-        // Double blocks
-        for block in self.double_blocks.iter() {
-            (img, txt) = block.forward(&img, &txt, &vec_, &pe)?
-        }
-        // Single blocks
-        let mut img = Tensor::cat(&[&txt, &img], 1)?;
-        for block in self.single_blocks.iter() {
-            img = block.forward(&img, &vec_, &pe)?;
-        }
-        let img = img.i((.., txt.dim(1)?..))?;
-        self.final_layer.forward(&img, &vec_)
+  #[allow(clippy::too_many_arguments)]
+  fn forward(&mut self, img: &Tensor, img_ids: &Tensor, txt: &Tensor, txt_ids: &Tensor, timesteps: &Tensor, y: &Tensor, guidance: Option<&Tensor>) -> Result<Tensor> {
+    if self.use_device_management {
+      self.forward_with_device_management(img, img_ids, txt, txt_ids, timesteps, y, guidance)
+    } else {
+      self.forward_without_device_management(img, img_ids, txt, txt_ids, timesteps, y, guidance)
     }
+  }
+}
+
+// This would be the implementation with full device management when available
+impl Flux {
+  /// Forward implementation with no device management - keeps all blocks on their current device.
+  /// This uses more GPU memory but may be faster since it avoids device transfers.
+  #[allow(clippy::too_many_arguments)]
+  fn forward_without_device_management(&mut self, img: &Tensor, img_ids: &Tensor, txt: &Tensor, txt_ids: &Tensor, timesteps: &Tensor, y: &Tensor, guidance: Option<&Tensor>) -> Result<Tensor> {
+    // Initial setup and preparation
+    if txt.rank() != 3 {
+      candle::bail!("unexpected shape for txt {:?}", txt.shape())
+    }
+    if img.rank() != 3 {
+      candle::bail!("unexpected shape for img {:?}", img.shape())
+    }
+    let dtype = img.dtype();
+    let pe = {
+      let ids = Tensor::cat(&[txt_ids, img_ids], 1)?;
+      ids.apply(&self.pe_embedder)?
+    };
+    let mut txt = txt.apply(&self.txt_in)?;
+    let mut img = img.apply(&self.img_in)?;
+    let vec_ = timestep_embedding(timesteps, 256, dtype)?.apply(&self.time_in)?;
+    let vec_ = match (self.guidance_in.as_ref(), guidance) {
+      (Some(g_in), Some(guidance)) => (vec_ + timestep_embedding(guidance, 256, dtype)?.apply(g_in))?,
+      _ => vec_,
+    };
+    let vec_ = (vec_ + y.apply(&self.vector_in))?;
+
+    // Process all double blocks at once (they're already on the device)
+    for block in &self.double_blocks {
+      let (new_img, new_txt) = block.forward(&img, &txt, &vec_, &pe)?;
+      img = new_img;
+      txt = new_txt;
+    }
+
+    // Combine tensors for single block processing
+    let mut img = Tensor::cat(&[&txt, &img], 1)?;
+
+    // Process all single blocks at once (they're already on the device)
+    for block in &self.single_blocks {
+      img = block.forward(&img, &vec_, &pe)?;
+    }
+
+    // Extract the relevant part of the output
+    let img = img.i((.., txt.dim(1)?..))?;
+
+    // Apply final layer
+    self.final_layer.forward(&img, &vec_)
+  }
+
+  /// Forward implementation with device management to reduce GPU memory usage.
+  /// Blocks are kept on CPU and only moved to GPU when needed.
+  #[allow(clippy::too_many_arguments)]
+  fn forward_with_device_management(&mut self, img: &Tensor, img_ids: &Tensor, txt: &Tensor, txt_ids: &Tensor, timesteps: &Tensor, y: &Tensor, guidance: Option<&Tensor>) -> Result<Tensor> {
+    // Initial setup and preparation
+    if txt.rank() != 3 {
+      candle::bail!("unexpected shape for txt {:?}", txt.shape())
+    }
+    if img.rank() != 3 {
+      candle::bail!("unexpected shape for img {:?}", img.shape())
+    }
+    let dtype = img.dtype();
+    let pe = {
+      let ids = Tensor::cat(&[txt_ids, img_ids], 1)?;
+      ids.apply(&self.pe_embedder)?
+    };
+    let mut txt = txt.apply(&self.txt_in)?;
+    let mut img = img.apply(&self.img_in)?;
+    let vec_ = timestep_embedding(timesteps, 256, dtype)?.apply(&self.time_in)?;
+    let vec_ = match (self.guidance_in.as_ref(), guidance) {
+      (Some(g_in), Some(guidance)) => (vec_ + timestep_embedding(guidance, 256, dtype)?.apply(g_in))?,
+      _ => vec_,
+    };
+    let vec_ = (vec_ + y.apply(&self.vector_in))?;
+
+    // Get original device for moving blocks back and forth
+    let original_device = img.device().clone();
+
+    // Calculate batch size (number of blocks to process at once)
+    let batch_size = self.max_active_blocks;
+    let num_double_batches = (self.double_blocks.len() + batch_size - 1) / batch_size;
+
+    if self.prefetch_next_batch {
+      println!("Processing {} double blocks in {} batches with prefetching enabled", self.double_blocks.len(), num_double_batches);
+    } else {
+      println!("Processing {} double blocks in {} batches", self.double_blocks.len(), num_double_batches);
+    }
+
+    // If prefetching is enabled, set up the first batch
+    let mut next_batch: Option<Vec<DoubleStreamBlock>> = None;
+    if self.prefetch_next_batch && num_double_batches > 1 {
+      // Prefetch the first batch
+      let mut prefetch_blocks = Vec::with_capacity(batch_size);
+      for i in 0..batch_size.min(self.double_blocks.len()) {
+        match self.double_blocks[i].to_device(&original_device) {
+          Ok(block) => prefetch_blocks.push(block),
+          Err(_) => prefetch_blocks.push(self.double_blocks[i].clone()),
+        }
+      }
+      next_batch = Some(prefetch_blocks);
+    }
+
+    // Process double blocks in batches
+    for batch_idx in 0..num_double_batches {
+      let start_idx = batch_idx * batch_size;
+      let end_idx = (start_idx + batch_size).min(self.double_blocks.len());
+
+      // Get the current GPU blocks (either prefetched or newly loaded)
+      let mut gpu_blocks = if let Some(prefetched) = next_batch.take() {
+        prefetched
+      } else {
+        // If not prefetched, load now
+        let mut blocks = Vec::with_capacity(end_idx - start_idx);
+        for i in start_idx..end_idx {
+          match self.double_blocks[i].to_device(&original_device) {
+            Ok(block) => blocks.push(block),
+            Err(_) => blocks.push(self.double_blocks[i].clone()),
+          }
+        }
+        blocks
+      };
+
+      // Start prefetching the next batch if enabled and not the last batch
+      if self.prefetch_next_batch && batch_idx < num_double_batches - 1 {
+        let next_start = (batch_idx + 1) * batch_size;
+        let next_end = (next_start + batch_size).min(self.double_blocks.len());
+
+        // Create a prefetch vector
+        let mut prefetch_blocks = Vec::with_capacity(next_end - next_start);
+        for i in next_start..next_end {
+          match self.double_blocks[i].to_device(&original_device) {
+            Ok(block) => prefetch_blocks.push(block),
+            Err(_) => prefetch_blocks.push(self.double_blocks[i].clone()),
+          }
+        }
+        next_batch = Some(prefetch_blocks);
+      }
+
+      // Process blocks
+      for block in &gpu_blocks {
+        let (new_img, new_txt) = block.forward(&img, &txt, &vec_, &pe)?;
+        img = new_img;
+        txt = new_txt;
+      }
+
+      // Let GPU blocks go out of scope
+    }
+
+    // Combine tensors for single block processing
+    let mut img = Tensor::cat(&[&txt, &img], 1)?;
+
+    // Process single blocks in batches
+    let num_single_batches = (self.single_blocks.len() + batch_size - 1) / batch_size;
+
+    // Reset next_batch for single blocks processing
+    let mut next_batch: Option<Vec<SingleStreamBlock>> = None;
+    if self.prefetch_next_batch && num_single_batches > 1 {
+      // Prefetch the first batch
+      let mut prefetch_blocks = Vec::with_capacity(batch_size);
+      for i in 0..batch_size.min(self.single_blocks.len()) {
+        match self.single_blocks[i].to_device(&original_device) {
+          Ok(block) => prefetch_blocks.push(block),
+          Err(_) => prefetch_blocks.push(self.single_blocks[i].clone()),
+        }
+      }
+      next_batch = Some(prefetch_blocks);
+    }
+
+    println!("Processing {} single blocks in {} batches", self.single_blocks.len(), num_single_batches);
+
+    for batch_idx in 0..num_single_batches {
+      let start_idx = batch_idx * batch_size;
+      let end_idx = (start_idx + batch_size).min(self.single_blocks.len());
+
+      // Get the current GPU blocks (either prefetched or newly loaded)
+      let mut gpu_blocks = if let Some(prefetched) = next_batch.take() {
+        prefetched
+      } else {
+        // If not prefetched, load now
+        let mut blocks = Vec::with_capacity(end_idx - start_idx);
+        for i in start_idx..end_idx {
+          match self.single_blocks[i].to_device(&original_device) {
+            Ok(block) => blocks.push(block),
+            Err(_) => blocks.push(self.single_blocks[i].clone()),
+          }
+        }
+        blocks
+      };
+
+      // Start prefetching the next batch if enabled and not the last batch
+      if self.prefetch_next_batch && batch_idx < num_single_batches - 1 {
+        let next_start = (batch_idx + 1) * batch_size;
+        let next_end = (next_start + batch_size).min(self.single_blocks.len());
+
+        // Create a prefetch vector
+        let mut prefetch_blocks = Vec::with_capacity(next_end - next_start);
+        for i in next_start..next_end {
+          match self.single_blocks[i].to_device(&original_device) {
+            Ok(block) => prefetch_blocks.push(block),
+            Err(_) => prefetch_blocks.push(self.single_blocks[i].clone()),
+          }
+        }
+        next_batch = Some(prefetch_blocks);
+      }
+
+      // Process blocks
+      for block in &gpu_blocks {
+        img = block.forward(&img, &vec_, &pe)?;
+      }
+
+      // Let GPU blocks go out of scope
+    }
+
+    // Extract the relevant part of the output
+    let img = img.i((.., txt.dim(1)?..))?;
+
+    // Apply final layer
+    self.final_layer.forward(&img, &vec_)
+  }
+}
+
+// Add DeviceTransferable for ModulationOut
+impl DeviceTransferable for ModulationOut {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { shift: self.shift.to_device(device)?, scale: self.scale.to_device(device)?, gate: self.gate.to_device(device)? })
+  }
+}
+
+// Replace conflicting implementations with placeholders
+// These will use the default trait implementation
+impl DeviceTransferable for Modulation1 {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { lin: self.lin.to_device(device)? })
+  }
+}
+impl DeviceTransferable for Modulation2 {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { lin: self.lin.to_device(device)? })
+  }
+}
+impl DeviceTransferable for QkNorm {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { query_norm: self.query_norm.to_device(device)?, key_norm: self.key_norm.to_device(device)? })
+  }
+}
+impl DeviceTransferable for SelfAttention {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { qkv: self.qkv.to_device(device)?, norm: self.norm.to_device(device)?, proj: self.proj.to_device(device)?, num_heads: self.num_heads })
+  }
+}
+impl DeviceTransferable for Mlp {
+  fn to_device(&self, device: &Device) -> Result<Self> {
+    Ok(Self { lin1: self.lin1.to_device(device)?, lin2: self.lin2.to_device(device)? })
+  }
 }
diff --git a/crates/vendor/candle/candle-transformers/src/models/flux/quantized_model.rs b/crates/vendor/candle/candle-transformers/src/models/flux/quantized_model.rs
index 0efeeab573..5d864fdf5f 100644
--- a/crates/vendor/candle/candle-transformers/src/models/flux/quantized_model.rs
+++ b/crates/vendor/candle/candle-transformers/src/models/flux/quantized_model.rs
@@ -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,
diff --git a/crates/vendor/candle/candle-transformers/src/models/flux/sampling.rs b/crates/vendor/candle/candle-transformers/src/models/flux/sampling.rs
index f3f0eafd4b..795f4c7183 100644
--- a/crates/vendor/candle/candle-transformers/src/models/flux/sampling.rs
+++ b/crates/vendor/candle/candle-transformers/src/models/flux/sampling.rs
@@ -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,
diff --git a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet_2d.rs b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet_2d.rs
index 421fbf269c..fcea8c7f7e 100644
--- a/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet_2d.rs
+++ b/crates/vendor/candle/candle-transformers/src/models/stable_diffusion/unet_2d.rs
@@ -5,7 +5,7 @@
 use super::embeddings::{TimestepEmbedding, Timesteps};
 use super::unet_2d_blocks::*;
 use crate::models::with_tracing::{conv2d, Conv2d};
-use candle::{Result, Tensor, Device, DType};
+use candle::{DType, Device, Result, Tensor};
 use candle_nn as nn;
 use candle_nn::Module;
 use std::collections::HashMap;
@@ -16,424 +16,689 @@ use std::collections::HashMap;
 // }
 #[derive(Debug, Clone, Copy)]
 pub struct BlockConfig {
-  pub out_channels: usize,
-  /// When `None` no cross-attn is used, when `Some(d)` then cross-attn is used and `d` is the
-  /// number of transformer blocks to be used.
-  pub use_cross_attn: Option<usize>,
-  pub attention_head_dim: usize,
+    pub out_channels: usize,
+    /// When `None` no cross-attn is used, when `Some(d)` then cross-attn is used and `d` is the
+    /// number of transformer blocks to be used.
+    pub use_cross_attn: Option<usize>,
+    pub attention_head_dim: usize,
 }
 
 #[derive(Debug, Clone)]
 pub struct UNet2DConditionModelConfig {
-  pub center_input_sample: bool,
-  pub flip_sin_to_cos: bool,
-  pub freq_shift: f64,
-  pub blocks: Vec<BlockConfig>,
-  pub layers_per_block: usize,
-  pub downsample_padding: usize,
-  pub mid_block_scale_factor: f64,
-  pub norm_num_groups: usize,
-  pub norm_eps: f64,
-  pub cross_attention_dim: usize,
-  pub sliced_attention_size: Option<usize>,
-  pub use_linear_projection: bool,
-  pub time_cond_proj_dim: Option<usize>,
-  pub time_embed_dim: Option<usize>,
+    pub center_input_sample: bool,
+    pub flip_sin_to_cos: bool,
+    pub freq_shift: f64,
+    pub blocks: Vec<BlockConfig>,
+    pub layers_per_block: usize,
+    pub downsample_padding: usize,
+    pub mid_block_scale_factor: f64,
+    pub norm_num_groups: usize,
+    pub norm_eps: f64,
+    pub cross_attention_dim: usize,
+    pub sliced_attention_size: Option<usize>,
+    pub use_linear_projection: bool,
+    pub time_cond_proj_dim: Option<usize>,
+    pub time_embed_dim: Option<usize>,
 }
 
 impl Default for UNet2DConditionModelConfig {
-  fn default() -> Self {
-    Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false, time_cond_proj_dim: Some(256), time_embed_dim: None }
-  }
+    fn default() -> Self {
+        Self {
+            center_input_sample: false,
+            flip_sin_to_cos: true,
+            freq_shift: 0.,
+            blocks: vec![
+                BlockConfig {
+                    out_channels: 320,
+                    use_cross_attn: Some(1),
+                    attention_head_dim: 8,
+                },
+                BlockConfig {
+                    out_channels: 640,
+                    use_cross_attn: Some(1),
+                    attention_head_dim: 8,
+                },
+                BlockConfig {
+                    out_channels: 1280,
+                    use_cross_attn: Some(1),
+                    attention_head_dim: 8,
+                },
+                BlockConfig {
+                    out_channels: 1280,
+                    use_cross_attn: None,
+                    attention_head_dim: 8,
+                },
+            ],
+            layers_per_block: 2,
+            downsample_padding: 1,
+            mid_block_scale_factor: 1.,
+            norm_num_groups: 32,
+            norm_eps: 1e-5,
+            cross_attention_dim: 1280,
+            sliced_attention_size: None,
+            use_linear_projection: false,
+            time_cond_proj_dim: Some(256),
+            time_embed_dim: None,
+        }
+    }
 }
 
 #[derive(Debug)]
 pub(crate) enum UNetDownBlock {
-  Basic(DownBlock2D),
-  CrossAttn(CrossAttnDownBlock2D),
+    Basic(DownBlock2D),
+    CrossAttn(CrossAttnDownBlock2D),
 }
 
 #[derive(Debug)]
 enum UNetUpBlock {
-  Basic(UpBlock2D),
-  CrossAttn(CrossAttnUpBlock2D),
+    Basic(UpBlock2D),
+    CrossAttn(CrossAttnUpBlock2D),
 }
 
 #[derive(Debug)]
 pub struct UNet2DConditionModel {
-  conv_in: Conv2d,
-  time_proj: Timesteps,
-  time_embedding: TimestepEmbedding,
-  down_blocks: Vec<UNetDownBlock>,
-  mid_block: UNetMidBlock2DCrossAttn,
-  up_blocks: Vec<UNetUpBlock>,
-  conv_norm_out: nn::GroupNorm,
-  conv_out: Conv2d,
-  span: tracing::Span,
-  config: UNet2DConditionModelConfig,
+    conv_in: Conv2d,
+    time_proj: Timesteps,
+    time_embedding: TimestepEmbedding,
+    down_blocks: Vec<UNetDownBlock>,
+    mid_block: UNetMidBlock2DCrossAttn,
+    up_blocks: Vec<UNetUpBlock>,
+    conv_norm_out: nn::GroupNorm,
+    conv_out: Conv2d,
+    span: tracing::Span,
+    config: UNet2DConditionModelConfig,
 }
 
 impl UNet2DConditionModel {
-  pub fn new(vs: nn::VarBuilder, in_channels: usize, out_channels: usize, use_flash_attn: bool, config: UNet2DConditionModelConfig) -> Result<Self> {
-    let n_blocks = config.blocks.len();
-    let b_channels = config.blocks[0].out_channels;
-    let bl_channels = config.blocks.last().unwrap().out_channels;
-    let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
-    let time_embed_dim = config.time_embed_dim.unwrap_or(b_channels * 4);
-    let conv_cfg = nn::Conv2dConfig { padding: 1, ..Default::default() };
-
-    // Log important configuration details
-    println!("Creating UNet2DConditionModel with {} blocks", n_blocks);
-    if let Some(proj_dim) = config.time_cond_proj_dim {
-      println!("UNet configured with time_cond_proj_dim: {}", proj_dim);
-      println!("This means the model expects guidance embeddings of dimension {}", proj_dim);
-    }
-
-    let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
-
-    let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
-    let time_embedding = TimestepEmbedding::new(vs.pp("time_embedding"), b_channels, time_embed_dim)?;
-
-    let vs_db = vs.pp("down_blocks");
-    let down_blocks = (0..n_blocks)
-      .map(|i| {
-        let BlockConfig { out_channels, use_cross_attn, attention_head_dim } = config.blocks[i];
-
-        // Enable automatic attention slicing if the config sliced_attention_size is set to 0.
-        let sliced_attention_size = match config.sliced_attention_size {
-          Some(0) => Some(attention_head_dim / 2),
-          _ => config.sliced_attention_size,
+    pub fn new(
+        vs: nn::VarBuilder,
+        in_channels: usize,
+        out_channels: usize,
+        use_flash_attn: bool,
+        config: UNet2DConditionModelConfig,
+    ) -> Result<Self> {
+        let n_blocks = config.blocks.len();
+        let b_channels = config.blocks[0].out_channels;
+        let bl_channels = config.blocks.last().unwrap().out_channels;
+        let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
+        let time_embed_dim = config.time_embed_dim.unwrap_or(b_channels * 4);
+        let conv_cfg = nn::Conv2dConfig {
+            padding: 1,
+            ..Default::default()
         };
 
-        let in_channels = if i > 0 { config.blocks[i - 1].out_channels } else { b_channels };
-        let db_cfg = DownBlock2DConfig { num_layers: config.layers_per_block, resnet_eps: config.norm_eps, resnet_groups: config.norm_num_groups, add_downsample: i < n_blocks - 1, downsample_padding: config.downsample_padding, ..Default::default() };
-        if let Some(transformer_layers_per_block) = use_cross_attn {
-          let config = CrossAttnDownBlock2DConfig { downblock: db_cfg, attn_num_head_channels: attention_head_dim, cross_attention_dim: config.cross_attention_dim, sliced_attention_size, use_linear_projection: config.use_linear_projection, transformer_layers_per_block };
-          let block = CrossAttnDownBlock2D::new(vs_db.pp(i.to_string()), in_channels, out_channels, Some(time_embed_dim), use_flash_attn, config)?;
-          Ok(UNetDownBlock::CrossAttn(block))
-        } else {
-          let block = DownBlock2D::new(vs_db.pp(i.to_string()), in_channels, out_channels, Some(time_embed_dim), db_cfg)?;
-          Ok(UNetDownBlock::Basic(block))
+        // Log important configuration details
+        println!("Creating UNet2DConditionModel with {} blocks", n_blocks);
+        if let Some(proj_dim) = config.time_cond_proj_dim {
+            println!("UNet configured with time_cond_proj_dim: {}", proj_dim);
+            println!(
+                "This means the model expects guidance embeddings of dimension {}",
+                proj_dim
+            );
         }
-      })
-      .collect::<Result<Vec<_>>>()?;
 
-    // https://github.com/huggingface/diffusers/blob/a76f2ad538e73b34d5fe7be08c8eb8ab38c7e90c/src/diffusers/models/unet_2d_condition.py#L462
-    let mid_transformer_layers_per_block = match config.blocks.last() {
-      None => 1,
-      Some(block) => block.use_cross_attn.unwrap_or(1),
-    };
-    let mid_cfg = UNetMidBlock2DCrossAttnConfig { resnet_eps: config.norm_eps, output_scale_factor: config.mid_block_scale_factor, cross_attn_dim: config.cross_attention_dim, attn_num_head_channels: bl_attention_head_dim, resnet_groups: Some(config.norm_num_groups), use_linear_projection: config.use_linear_projection, transformer_layers_per_block: mid_transformer_layers_per_block, ..Default::default() };
+        let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
 
-    let mid_block = UNetMidBlock2DCrossAttn::new(vs.pp("mid_block"), bl_channels, Some(time_embed_dim), use_flash_attn, mid_cfg)?;
+        let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
+        let time_embedding =
+            TimestepEmbedding::new(vs.pp("time_embedding"), b_channels, time_embed_dim)?;
 
-    let vs_ub = vs.pp("up_blocks");
-    let up_blocks = (0..n_blocks)
-      .map(|i| {
-        let BlockConfig { out_channels, use_cross_attn, attention_head_dim } = config.blocks[n_blocks - 1 - i];
+        let vs_db = vs.pp("down_blocks");
+        let down_blocks = (0..n_blocks)
+            .map(|i| {
+                let BlockConfig {
+                    out_channels,
+                    use_cross_attn,
+                    attention_head_dim,
+                } = config.blocks[i];
 
-        // Enable automatic attention slicing if the config sliced_attention_size is set to 0.
-        let sliced_attention_size = match config.sliced_attention_size {
-          Some(0) => Some(attention_head_dim / 2),
-          _ => config.sliced_attention_size,
+                // Enable automatic attention slicing if the config sliced_attention_size is set to 0.
+                let sliced_attention_size = match config.sliced_attention_size {
+                    Some(0) => Some(attention_head_dim / 2),
+                    _ => config.sliced_attention_size,
+                };
+
+                let in_channels = if i > 0 {
+                    config.blocks[i - 1].out_channels
+                } else {
+                    b_channels
+                };
+                let db_cfg = DownBlock2DConfig {
+                    num_layers: config.layers_per_block,
+                    resnet_eps: config.norm_eps,
+                    resnet_groups: config.norm_num_groups,
+                    add_downsample: i < n_blocks - 1,
+                    downsample_padding: config.downsample_padding,
+                    ..Default::default()
+                };
+                if let Some(transformer_layers_per_block) = use_cross_attn {
+                    let config = CrossAttnDownBlock2DConfig {
+                        downblock: db_cfg,
+                        attn_num_head_channels: attention_head_dim,
+                        cross_attention_dim: config.cross_attention_dim,
+                        sliced_attention_size,
+                        use_linear_projection: config.use_linear_projection,
+                        transformer_layers_per_block,
+                    };
+                    let block = CrossAttnDownBlock2D::new(
+                        vs_db.pp(i.to_string()),
+                        in_channels,
+                        out_channels,
+                        Some(time_embed_dim),
+                        use_flash_attn,
+                        config,
+                    )?;
+                    Ok(UNetDownBlock::CrossAttn(block))
+                } else {
+                    let block = DownBlock2D::new(
+                        vs_db.pp(i.to_string()),
+                        in_channels,
+                        out_channels,
+                        Some(time_embed_dim),
+                        db_cfg,
+                    )?;
+                    Ok(UNetDownBlock::Basic(block))
+                }
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        // https://github.com/huggingface/diffusers/blob/a76f2ad538e73b34d5fe7be08c8eb8ab38c7e90c/src/diffusers/models/unet_2d_condition.py#L462
+        let mid_transformer_layers_per_block = match config.blocks.last() {
+            None => 1,
+            Some(block) => block.use_cross_attn.unwrap_or(1),
+        };
+        let mid_cfg = UNetMidBlock2DCrossAttnConfig {
+            resnet_eps: config.norm_eps,
+            output_scale_factor: config.mid_block_scale_factor,
+            cross_attn_dim: config.cross_attention_dim,
+            attn_num_head_channels: bl_attention_head_dim,
+            resnet_groups: Some(config.norm_num_groups),
+            use_linear_projection: config.use_linear_projection,
+            transformer_layers_per_block: mid_transformer_layers_per_block,
+            ..Default::default()
         };
 
-        let prev_out_channels = if i > 0 { config.blocks[n_blocks - i].out_channels } else { bl_channels };
-        let in_channels = {
-          let index = if i == n_blocks - 1 { 0 } else { n_blocks - i - 2 };
-          config.blocks[index].out_channels
-        };
-        let ub_cfg = UpBlock2DConfig { num_layers: config.layers_per_block + 1, resnet_eps: config.norm_eps, resnet_groups: config.norm_num_groups, add_upsample: i < n_blocks - 1, ..Default::default() };
-        if let Some(transformer_layers_per_block) = use_cross_attn {
-          let config = CrossAttnUpBlock2DConfig { upblock: ub_cfg, attn_num_head_channels: attention_head_dim, cross_attention_dim: config.cross_attention_dim, sliced_attention_size, use_linear_projection: config.use_linear_projection, transformer_layers_per_block };
-          let block = CrossAttnUpBlock2D::new(vs_ub.pp(i.to_string()), in_channels, prev_out_channels, out_channels, Some(time_embed_dim), use_flash_attn, config)?;
-          Ok(UNetUpBlock::CrossAttn(block))
+        let mid_block = UNetMidBlock2DCrossAttn::new(
+            vs.pp("mid_block"),
+            bl_channels,
+            Some(time_embed_dim),
+            use_flash_attn,
+            mid_cfg,
+        )?;
+
+        let vs_ub = vs.pp("up_blocks");
+        let up_blocks = (0..n_blocks)
+            .map(|i| {
+                let BlockConfig {
+                    out_channels,
+                    use_cross_attn,
+                    attention_head_dim,
+                } = config.blocks[n_blocks - 1 - i];
+
+                // Enable automatic attention slicing if the config sliced_attention_size is set to 0.
+                let sliced_attention_size = match config.sliced_attention_size {
+                    Some(0) => Some(attention_head_dim / 2),
+                    _ => config.sliced_attention_size,
+                };
+
+                let prev_out_channels = if i > 0 {
+                    config.blocks[n_blocks - i].out_channels
+                } else {
+                    bl_channels
+                };
+                let in_channels = {
+                    let index = if i == n_blocks - 1 {
+                        0
+                    } else {
+                        n_blocks - i - 2
+                    };
+                    config.blocks[index].out_channels
+                };
+                let ub_cfg = UpBlock2DConfig {
+                    num_layers: config.layers_per_block + 1,
+                    resnet_eps: config.norm_eps,
+                    resnet_groups: config.norm_num_groups,
+                    add_upsample: i < n_blocks - 1,
+                    ..Default::default()
+                };
+                if let Some(transformer_layers_per_block) = use_cross_attn {
+                    let config = CrossAttnUpBlock2DConfig {
+                        upblock: ub_cfg,
+                        attn_num_head_channels: attention_head_dim,
+                        cross_attention_dim: config.cross_attention_dim,
+                        sliced_attention_size,
+                        use_linear_projection: config.use_linear_projection,
+                        transformer_layers_per_block,
+                    };
+                    let block = CrossAttnUpBlock2D::new(
+                        vs_ub.pp(i.to_string()),
+                        in_channels,
+                        prev_out_channels,
+                        out_channels,
+                        Some(time_embed_dim),
+                        use_flash_attn,
+                        config,
+                    )?;
+                    Ok(UNetUpBlock::CrossAttn(block))
+                } else {
+                    let block = UpBlock2D::new(
+                        vs_ub.pp(i.to_string()),
+                        in_channels,
+                        prev_out_channels,
+                        out_channels,
+                        Some(time_embed_dim),
+                        ub_cfg,
+                    )?;
+                    Ok(UNetUpBlock::Basic(block))
+                }
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        let conv_norm_out = nn::group_norm(
+            config.norm_num_groups,
+            b_channels,
+            config.norm_eps,
+            vs.pp("conv_norm_out"),
+        )?;
+        let conv_out = conv2d(b_channels, out_channels, 3, conv_cfg, vs.pp("conv_out"))?;
+        let span = tracing::span!(tracing::Level::TRACE, "unet2d");
+        Ok(Self {
+            conv_in,
+            time_proj,
+            time_embedding,
+            down_blocks,
+            mid_block,
+            up_blocks,
+            conv_norm_out,
+            conv_out,
+            span,
+            config,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        timestep: f64,
+        encoder_hidden_states: &Tensor,
+    ) -> Result<Tensor> {
+        let _enter = self.span.enter();
+        self.forward_with_additional_residuals(xs, timestep, encoder_hidden_states, None, None)
+    }
+
+    pub fn forward_with_additional_residuals(
+        &self,
+        xs: &Tensor,
+        timestep: f64,
+        encoder_hidden_states: &Tensor,
+        down_block_additional_residuals: Option<&[Tensor]>,
+        mid_block_additional_residual: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (bsize, _channels, height, width) = xs.dims4()?;
+        let device = xs.device();
+        let n_blocks = self.config.blocks.len();
+        let num_upsamplers = n_blocks - 1;
+        let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
+        let forward_upsample_size =
+            height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
+        // 0. center input if necessary
+        let xs = if self.config.center_input_sample {
+            ((xs * 2.0)? - 1.0)?
         } else {
-          let block = UpBlock2D::new(vs_ub.pp(i.to_string()), in_channels, prev_out_channels, out_channels, Some(time_embed_dim), ub_cfg)?;
-          Ok(UNetUpBlock::Basic(block))
+            xs.clone()
+        };
+        // 1. time
+        let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
+        let emb = self.time_proj.forward(&emb)?;
+        let time_emb = self.time_embedding.forward(&emb)?;
+        // 2. pre-process
+        let xs = self.conv_in.forward(&xs)?;
+        // 3. down
+        let mut down_block_res_xs = vec![xs.clone()];
+        let mut xs = xs;
+        for down_block in self.down_blocks.iter() {
+            let (_xs, res_xs) = match down_block {
+                UNetDownBlock::Basic(b) => b.forward(&xs, Some(&time_emb))?,
+                UNetDownBlock::CrossAttn(b) => {
+                    b.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?
+                }
+            };
+            down_block_res_xs.extend(res_xs);
+            xs = _xs;
         }
-      })
-      .collect::<Result<Vec<_>>>()?;
 
-    let conv_norm_out = nn::group_norm(config.norm_num_groups, b_channels, config.norm_eps, vs.pp("conv_norm_out"))?;
-    let conv_out = conv2d(b_channels, out_channels, 3, conv_cfg, vs.pp("conv_out"))?;
-    let span = tracing::span!(tracing::Level::TRACE, "unet2d");
-    Ok(Self { conv_in, time_proj, time_embedding, down_blocks, mid_block, up_blocks, conv_norm_out, conv_out, span, config })
-  }
-
-  pub fn forward(&self, xs: &Tensor, timestep: f64, encoder_hidden_states: &Tensor) -> Result<Tensor> {
-    let _enter = self.span.enter();
-    self.forward_with_additional_residuals(xs, timestep, encoder_hidden_states, None, None)
-  }
-
-  pub fn forward_with_additional_residuals(&self, xs: &Tensor, timestep: f64, encoder_hidden_states: &Tensor, down_block_additional_residuals: Option<&[Tensor]>, mid_block_additional_residual: Option<&Tensor>) -> Result<Tensor> {
-    let (bsize, _channels, height, width) = xs.dims4()?;
-    let device = xs.device();
-    let n_blocks = self.config.blocks.len();
-    let num_upsamplers = n_blocks - 1;
-    let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
-    let forward_upsample_size = height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
-    // 0. center input if necessary
-    let xs = if self.config.center_input_sample { ((xs * 2.0)? - 1.0)? } else { xs.clone() };
-    // 1. time
-    let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
-    let emb = self.time_proj.forward(&emb)?;
-    let time_emb = self.time_embedding.forward(&emb)?;
-    // 2. pre-process
-    let xs = self.conv_in.forward(&xs)?;
-    // 3. down
-    let mut down_block_res_xs = vec![xs.clone()];
-    let mut xs = xs;
-    for down_block in self.down_blocks.iter() {
-      let (_xs, res_xs) = match down_block {
-        UNetDownBlock::Basic(b) => b.forward(&xs, Some(&time_emb))?,
-        UNetDownBlock::CrossAttn(b) => b.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?,
-      };
-      down_block_res_xs.extend(res_xs);
-      xs = _xs;
-    }
-
-    let new_down_block_res_xs = if let Some(down_block_additional_residuals) = down_block_additional_residuals {
-      let mut v = vec![];
-      // A previous version of this code had a bug because of the addition being made
-      // in place via += hence modifying the input of the mid block.
-      for (i, residuals) in down_block_additional_residuals.iter().enumerate() {
-        v.push((&down_block_res_xs[i] + residuals)?)
-      }
-      v
-    } else {
-      down_block_res_xs
-    };
-    let mut down_block_res_xs = new_down_block_res_xs;
-
-    // 4. mid
-    let xs = self.mid_block.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?;
-    let xs = match mid_block_additional_residual {
-      None => xs,
-      Some(m) => (m + xs)?,
-    };
-    // 5. up
-    let mut xs = xs;
-    let mut upsample_size = None;
-    for (i, up_block) in self.up_blocks.iter().enumerate() {
-      let n_resnets = match up_block {
-        UNetUpBlock::Basic(b) => b.resnets.len(),
-        UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
-      };
-      let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
-      if i < n_blocks - 1 && forward_upsample_size {
-        let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
-        upsample_size = Some((h, w))
-      }
-      xs = match up_block {
-        UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size)?,
-        UNetUpBlock::CrossAttn(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size, Some(encoder_hidden_states))?,
-      };
-    }
-    // 6. post-process
-    let xs = self.conv_norm_out.forward(&xs)?;
-    let xs = nn::ops::silu(&xs)?;
-    self.conv_out.forward(&xs)
-  }
-
-  pub fn forward_with_guidance(&self, xs: &Tensor, timestep: f64, encoder_hidden_states: &Tensor, guidance_embedding: Option<&Tensor>) -> Result<Tensor> {
-    let _enter = self.span.enter();
-    let (bsize, _channels, height, width) = xs.dims4()?;
-    let device = xs.device();
-    let n_blocks = self.config.blocks.len();
-    let num_upsamplers = n_blocks - 1;
-    let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
-    let forward_upsample_size = height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
-
-    // 0. center input if necessary
-    let xs = if self.config.center_input_sample { ((xs * 2.0)? - 1.0)? } else { xs.clone() };
-
-    // 1. time
-    let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
-    let emb = self.time_proj.forward(&emb)?;
-    let mut time_emb = self.time_embedding.forward(&emb)?;
-
-    // Add guidance embedding if provided
-    if let Some(guidance_emb) = guidance_embedding {
-      // Get dimensions of both embeddings
-      let (batch_size, time_emb_dim) = time_emb.dims2()?;
-      let (guidance_batch, guidance_emb_dim) = guidance_emb.dims2()?;
-
-      println!("==== Guidance Embedding Integration ====");
-      println!("Time embedding dims: [{}, {}]", batch_size, time_emb_dim);
-      println!("Guidance embedding dims: [{}, {}]", guidance_batch, guidance_emb_dim);
-
-      // Use time_cond_proj_dim from config to determine expected guidance embedding dimension
-      let expected_guidance_dim = self.config.time_cond_proj_dim.unwrap_or(time_emb_dim);
-      println!("Expected guidance dimension from config: {}", expected_guidance_dim);
-
-      // Check if dimensions match the expected projection dimension
-      let processed_guidance_emb = if guidance_emb_dim != time_emb_dim {
-        // Convert guidance_emb to f32 for calculations
-        let guidance_emb_f32 = guidance_emb.to_dtype(DType::F32)?;
-        println!("Converting guidance embedding to F32 for processing");
-
-        // If the guidance embedding dimension matches the expected projection dimension,
-        // we need to project it to match the time embedding dimension
-        if guidance_emb_dim == expected_guidance_dim {
-          println!("Projecting guidance embedding from {} to {}", guidance_emb_dim, time_emb_dim);
-
-          // Use linear interpolation
-          let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
-          let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
-
-          // Get guidance embedding data (now f32)
-          let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
-          println!("Reshaping guidance data to vector of length {}", guidance_data.len());
-
-          // Perform linear interpolation
-          let mut interpolated_data = Vec::with_capacity(time_emb_dim);
-          for &t in &dst_positions {
-            // Find the two source positions that surround the target position
-            let mut i = 0;
-            while i < src_positions.len() - 1 && src_positions[i + 1] < t {
-              i += 1;
-            }
-
-            if i == src_positions.len() - 1 {
-              // At or beyond the last position
-              interpolated_data.push(guidance_data[i]);
-            } else if t <= src_positions[0] {
-              // At or before the first position
-              interpolated_data.push(guidance_data[0]);
+        let new_down_block_res_xs =
+            if let Some(down_block_additional_residuals) = down_block_additional_residuals {
+                let mut v = vec![];
+                // A previous version of this code had a bug because of the addition being made
+                // in place via += hence modifying the input of the mid block.
+                for (i, residuals) in down_block_additional_residuals.iter().enumerate() {
+                    v.push((&down_block_res_xs[i] + residuals)?)
+                }
+                v
             } else {
-              // Linear interpolation between two points
-              let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
-              let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
-              interpolated_data.push(value);
+                down_block_res_xs
+            };
+        let mut down_block_res_xs = new_down_block_res_xs;
+
+        // 4. mid
+        let xs = self
+            .mid_block
+            .forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?;
+        let xs = match mid_block_additional_residual {
+            None => xs,
+            Some(m) => (m + xs)?,
+        };
+        // 5. up
+        let mut xs = xs;
+        let mut upsample_size = None;
+        for (i, up_block) in self.up_blocks.iter().enumerate() {
+            let n_resnets = match up_block {
+                UNetUpBlock::Basic(b) => b.resnets.len(),
+                UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
+            };
+            let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
+            if i < n_blocks - 1 && forward_upsample_size {
+                let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
+                upsample_size = Some((h, w))
             }
-          }
+            xs = match up_block {
+                UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size)?,
+                UNetUpBlock::CrossAttn(b) => b.forward(
+                    &xs,
+                    &res_xs,
+                    Some(&time_emb),
+                    upsample_size,
+                    Some(encoder_hidden_states),
+                )?,
+            };
+        }
+        // 6. post-process
+        let xs = self.conv_norm_out.forward(&xs)?;
+        let xs = nn::ops::silu(&xs)?;
+        self.conv_out.forward(&xs)
+    }
 
-          println!("Created interpolated data of length {}", interpolated_data.len());
+    pub fn forward_with_guidance(
+        &self,
+        xs: &Tensor,
+        timestep: f64,
+        encoder_hidden_states: &Tensor,
+        guidance_embedding: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let _enter = self.span.enter();
+        let (bsize, _channels, height, width) = xs.dims4()?;
+        let device = xs.device();
+        let n_blocks = self.config.blocks.len();
+        let num_upsamplers = n_blocks - 1;
+        let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
+        let forward_upsample_size =
+            height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
 
-          // Create tensor and convert to time_emb's dtype
-          let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
-          let result = tensor.to_dtype(time_emb.dtype())?;
-          println!("Final projected guidance embedding shape: {:?}", result.shape());
-          result
+        // 0. center input if necessary
+        let xs = if self.config.center_input_sample {
+            ((xs * 2.0)? - 1.0)?
         } else {
-          // Unusual case - dimensions don't match expected pattern
-          println!("Warning: Guidance embedding dimension ({}) doesn't match expected projection dimension ({})", guidance_emb_dim, expected_guidance_dim);
+            xs.clone()
+        };
 
-          // Fallback to generic interpolation
-          let guidance_data = guidance_emb_f32.reshape((guidance_emb_dim,))?.to_vec1::<f32>()?;
+        // 1. time
+        let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
+        let emb = self.time_proj.forward(&emb)?;
+        let mut time_emb = self.time_embedding.forward(&emb)?;
 
-          let src_positions: Vec<f32> = (0..guidance_emb_dim).map(|i| i as f32 / (guidance_emb_dim - 1) as f32).collect();
-          let dst_positions: Vec<f32> = (0..time_emb_dim).map(|i| i as f32 / (time_emb_dim - 1) as f32).collect();
+        // Add guidance embedding if provided
+        if let Some(guidance_emb) = guidance_embedding {
+            // Get dimensions of both embeddings
+            let (batch_size, time_emb_dim) = time_emb.dims2()?;
+            let (guidance_batch, guidance_emb_dim) = guidance_emb.dims2()?;
 
-          let mut interpolated_data = Vec::with_capacity(time_emb_dim);
-          for &t in &dst_positions {
-            let mut i = 0;
-            while i < src_positions.len() - 1 && src_positions[i + 1] < t {
-              i += 1;
-            }
+            println!("==== Guidance Embedding Integration ====");
+            println!("Time embedding dims: [{}, {}]", batch_size, time_emb_dim);
+            println!(
+                "Guidance embedding dims: [{}, {}]",
+                guidance_batch, guidance_emb_dim
+            );
 
-            if i == src_positions.len() - 1 {
-              interpolated_data.push(guidance_data[i]);
-            } else if t <= src_positions[0] {
-              interpolated_data.push(guidance_data[0]);
+            // Use time_cond_proj_dim from config to determine expected guidance embedding dimension
+            let expected_guidance_dim = self.config.time_cond_proj_dim.unwrap_or(time_emb_dim);
+            println!(
+                "Expected guidance dimension from config: {}",
+                expected_guidance_dim
+            );
+
+            // Check if dimensions match the expected projection dimension
+            let processed_guidance_emb = if guidance_emb_dim != time_emb_dim {
+                // Convert guidance_emb to f32 for calculations
+                let guidance_emb_f32 = guidance_emb.to_dtype(DType::F32)?;
+                println!("Converting guidance embedding to F32 for processing");
+
+                // If the guidance embedding dimension matches the expected projection dimension,
+                // we need to project it to match the time embedding dimension
+                if guidance_emb_dim == expected_guidance_dim {
+                    println!(
+                        "Projecting guidance embedding from {} to {}",
+                        guidance_emb_dim, time_emb_dim
+                    );
+
+                    // Use linear interpolation
+                    let src_positions: Vec<f32> = (0..guidance_emb_dim)
+                        .map(|i| i as f32 / (guidance_emb_dim - 1) as f32)
+                        .collect();
+                    let dst_positions: Vec<f32> = (0..time_emb_dim)
+                        .map(|i| i as f32 / (time_emb_dim - 1) as f32)
+                        .collect();
+
+                    // Get guidance embedding data (now f32)
+                    let guidance_data = guidance_emb_f32
+                        .reshape((guidance_emb_dim,))?
+                        .to_vec1::<f32>()?;
+                    println!(
+                        "Reshaping guidance data to vector of length {}",
+                        guidance_data.len()
+                    );
+
+                    // Perform linear interpolation
+                    let mut interpolated_data = Vec::with_capacity(time_emb_dim);
+                    for &t in &dst_positions {
+                        // Find the two source positions that surround the target position
+                        let mut i = 0;
+                        while i < src_positions.len() - 1 && src_positions[i + 1] < t {
+                            i += 1;
+                        }
+
+                        if i == src_positions.len() - 1 {
+                            // At or beyond the last position
+                            interpolated_data.push(guidance_data[i]);
+                        } else if t <= src_positions[0] {
+                            // At or before the first position
+                            interpolated_data.push(guidance_data[0]);
+                        } else {
+                            // Linear interpolation between two points
+                            let t_norm =
+                                (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
+                            let value =
+                                guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
+                            interpolated_data.push(value);
+                        }
+                    }
+
+                    println!(
+                        "Created interpolated data of length {}",
+                        interpolated_data.len()
+                    );
+
+                    // Create tensor and convert to time_emb's dtype
+                    let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
+                    let result = tensor.to_dtype(time_emb.dtype())?;
+                    println!(
+                        "Final projected guidance embedding shape: {:?}",
+                        result.shape()
+                    );
+                    result
+                } else {
+                    // Unusual case - dimensions don't match expected pattern
+                    println!("Warning: Guidance embedding dimension ({}) doesn't match expected projection dimension ({})", 
+                             guidance_emb_dim, expected_guidance_dim);
+
+                    // Fallback to generic interpolation
+                    let guidance_data = guidance_emb_f32
+                        .reshape((guidance_emb_dim,))?
+                        .to_vec1::<f32>()?;
+
+                    let src_positions: Vec<f32> = (0..guidance_emb_dim)
+                        .map(|i| i as f32 / (guidance_emb_dim - 1) as f32)
+                        .collect();
+                    let dst_positions: Vec<f32> = (0..time_emb_dim)
+                        .map(|i| i as f32 / (time_emb_dim - 1) as f32)
+                        .collect();
+
+                    let mut interpolated_data = Vec::with_capacity(time_emb_dim);
+                    for &t in &dst_positions {
+                        let mut i = 0;
+                        while i < src_positions.len() - 1 && src_positions[i + 1] < t {
+                            i += 1;
+                        }
+
+                        if i == src_positions.len() - 1 {
+                            interpolated_data.push(guidance_data[i]);
+                        } else if t <= src_positions[0] {
+                            interpolated_data.push(guidance_data[0]);
+                        } else {
+                            let t_norm =
+                                (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
+                            let value =
+                                guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
+                            interpolated_data.push(value);
+                        }
+                    }
+
+                    println!(
+                        "Created interpolated data of length {}",
+                        interpolated_data.len()
+                    );
+                    let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
+                    let result = tensor.to_dtype(time_emb.dtype())?;
+                    println!(
+                        "Final fallback guidance embedding shape: {:?}",
+                        result.shape()
+                    );
+                    result
+                }
             } else {
-              let t_norm = (t - src_positions[i]) / (src_positions[i + 1] - src_positions[i]);
-              let value = guidance_data[i] * (1.0 - t_norm) + guidance_data[i + 1] * t_norm;
-              interpolated_data.push(value);
+                // Dimensions match, no processing needed
+                println!("Guidance dimension already matches time embedding dimension, no interpolation needed");
+                let result = guidance_emb
+                    .reshape((1, time_emb_dim))?
+                    .to_dtype(time_emb.dtype())?;
+                println!("Reshaped guidance embedding shape: {:?}", result.shape());
+                result
+            };
+
+            // Expand to match batch size if needed
+            let expanded_guidance = processed_guidance_emb.broadcast_as(time_emb.shape())?;
+            println!(
+                "Expanded guidance embedding shape: {:?}",
+                expanded_guidance.shape()
+            );
+
+            // Save the processed and expanded guidance embeddings to disk if requested
+            if let Ok(save_dir) = std::env::var("CANDLE_SAVE_TENSORS_DIR") {
+                // Save pre-expansion (original projection)
+                let proc_path = format!("{}/unet_processed_guidance_embedding.npy", save_dir);
+                println!("Saving processed guidance embedding to {}", proc_path);
+                if let Err(e) = processed_guidance_emb.write_npy(&proc_path) {
+                    println!("Error saving processed guidance embedding: {}", e);
+                } else {
+                    println!(
+                        "Successfully saved processed guidance embedding to {}",
+                        proc_path
+                    );
+                }
+
+                // Save post-expansion (final tensor used for addition to time_emb)
+                let exp_path = format!("{}/unet_expanded_guidance_embedding.npy", save_dir);
+                println!("Saving expanded guidance embedding to {}", exp_path);
+                if let Err(e) = expanded_guidance.write_npy(&exp_path) {
+                    println!("Error saving expanded guidance embedding: {}", e);
+                } else {
+                    println!(
+                        "Successfully saved expanded guidance embedding to {}",
+                        exp_path
+                    );
+                }
             }
-          }
 
-          println!("Created interpolated data of length {}", interpolated_data.len());
-          let tensor = Tensor::from_vec(interpolated_data, (1, time_emb_dim), device)?;
-          let result = tensor.to_dtype(time_emb.dtype())?;
-          println!("Final fallback guidance embedding shape: {:?}", result.shape());
-          result
-        }
-      } else {
-        // Dimensions match, no processing needed
-        println!("Guidance dimension already matches time embedding dimension, no interpolation needed");
-        let result = guidance_emb.reshape((1, time_emb_dim))?.to_dtype(time_emb.dtype())?;
-        println!("Reshaped guidance embedding shape: {:?}", result.shape());
-        result
-      };
-
-      // Expand to match batch size if needed
-      let expanded_guidance = processed_guidance_emb.broadcast_as(time_emb.shape())?;
-      println!("Expanded guidance embedding shape: {:?}", expanded_guidance.shape());
-
-      // Save the processed and expanded guidance embeddings to disk if requested
-      if let Ok(save_dir) = std::env::var("CANDLE_SAVE_TENSORS_DIR") {
-        // Save pre-expansion (original projection)
-        let proc_path = format!("{}/unet_processed_guidance_embedding.npy", save_dir);
-        println!("Saving processed guidance embedding to {}", proc_path);
-        if let Err(e) = processed_guidance_emb.write_npy(&proc_path) {
-          println!("Error saving processed guidance embedding: {}", e);
-        } else {
-          println!("Successfully saved processed guidance embedding to {}", proc_path);
+            // Add to time embeddings
+            time_emb = (time_emb + expanded_guidance)?;
+            println!(
+                "Final time embedding with guidance shape: {:?}",
+                time_emb.shape()
+            );
+            println!("==== End Guidance Integration ====");
         }
 
-        // Save post-expansion (final tensor used for addition to time_emb)
-        let exp_path = format!("{}/unet_expanded_guidance_embedding.npy", save_dir);
-        println!("Saving expanded guidance embedding to {}", exp_path);
-        if let Err(e) = expanded_guidance.write_npy(&exp_path) {
-          println!("Error saving expanded guidance embedding: {}", e);
-        } else {
-          println!("Successfully saved expanded guidance embedding to {}", exp_path);
+        // 2. pre-process
+        let xs = self.conv_in.forward(&xs)?;
+
+        // 3. down
+        let mut down_block_res_xs = vec![xs.clone()];
+        let mut xs = xs;
+        for down_block in self.down_blocks.iter() {
+            let (_xs, res_xs) = match down_block {
+                UNetDownBlock::Basic(b) => b.forward(&xs, Some(&time_emb))?,
+                UNetDownBlock::CrossAttn(b) => {
+                    b.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?
+                }
+            };
+            down_block_res_xs.extend(res_xs);
+            xs = _xs;
         }
-      }
 
-      // Add to time embeddings
-      time_emb = (time_emb + expanded_guidance)?;
-      println!("Final time embedding with guidance shape: {:?}", time_emb.shape());
-      println!("==== End Guidance Integration ====");
+        // 4. mid
+        let xs = self
+            .mid_block
+            .forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?;
+
+        // 5. up
+        let mut xs = xs;
+        let mut upsample_size = None;
+        for (i, up_block) in self.up_blocks.iter().enumerate() {
+            let n_resnets = match up_block {
+                UNetUpBlock::Basic(b) => b.resnets.len(),
+                UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
+            };
+            let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
+            if i < n_blocks - 1 && forward_upsample_size {
+                let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
+                upsample_size = Some((h, w))
+            }
+            xs = match up_block {
+                UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size)?,
+                UNetUpBlock::CrossAttn(b) => b.forward(
+                    &xs,
+                    &res_xs,
+                    Some(&time_emb),
+                    upsample_size,
+                    Some(encoder_hidden_states),
+                )?,
+            };
+        }
+
+        // 6. post-process
+        let xs = self.conv_norm_out.forward(&xs)?;
+        let xs = nn::ops::silu(&xs)?;
+        self.conv_out.forward(&xs)
     }
 
-    // 2. pre-process
-    let xs = self.conv_in.forward(&xs)?;
-
-    // 3. down
-    let mut down_block_res_xs = vec![xs.clone()];
-    let mut xs = xs;
-    for down_block in self.down_blocks.iter() {
-      let (_xs, res_xs) = match down_block {
-        UNetDownBlock::Basic(b) => b.forward(&xs, Some(&time_emb))?,
-        UNetDownBlock::CrossAttn(b) => b.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?,
-      };
-      down_block_res_xs.extend(res_xs);
-      xs = _xs;
+    pub fn load_with_weights(
+        config: UNet2DConditionModelConfig,
+        weights: HashMap<String, Tensor>,
+        device: &Device,
+        use_flash_attn: bool,
+        dtype: DType,
+    ) -> Result<Self> {
+        let vb = nn::VarBuilder::from_tensors(weights, dtype, device);
+        Self::new(vb, 4, 4, use_flash_attn, config)
     }
-
-    // 4. mid
-    let xs = self.mid_block.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?;
-
-    // 5. up
-    let mut xs = xs;
-    let mut upsample_size = None;
-    for (i, up_block) in self.up_blocks.iter().enumerate() {
-      let n_resnets = match up_block {
-        UNetUpBlock::Basic(b) => b.resnets.len(),
-        UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
-      };
-      let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
-      if i < n_blocks - 1 && forward_upsample_size {
-        let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
-        upsample_size = Some((h, w))
-      }
-      xs = match up_block {
-        UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size)?,
-        UNetUpBlock::CrossAttn(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size, Some(encoder_hidden_states))?,
-      };
-    }
-
-    // 6. post-process
-    let xs = self.conv_norm_out.forward(&xs)?;
-    let xs = nn::ops::silu(&xs)?;
-    self.conv_out.forward(&xs)
-  }
-
-  pub fn load_with_weights(config: UNet2DConditionModelConfig, weights: HashMap<String, Tensor>, device: &Device, use_flash_attn: bool, dtype: DType) -> Result<Self> {
-    let vb = nn::VarBuilder::from_tensors(weights, dtype, device);
-    Self::new(vb, 4, 4, use_flash_attn, config)
-  }
 }
diff --git a/frontend/realtime/package-lock.json b/frontend/realtime/package-lock.json
index 8f2b1bd6de..c574a6953d 100644
--- a/frontend/realtime/package-lock.json
+++ b/frontend/realtime/package-lock.json
@@ -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"
diff --git a/frontend/realtime/package.json b/frontend/realtime/package.json
index 74261bea47..66e08b8129 100644
--- a/frontend/realtime/package.json
+++ b/frontend/realtime/package.json
@@ -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",
diff --git a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts
index b5be251276..b050b5c023 100644
--- a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts
+++ b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts
@@ -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)
 
diff --git a/frontend/realtime/yarn.lock b/frontend/realtime/yarn.lock
index 637b09febf..6fb4000b26 100644
--- a/frontend/realtime/yarn.lock
+++ b/frontend/realtime/yarn.lock
@@ -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"